[Kernel] Pipe attn_logits_soft_cap through paged attention TPU kernels#12482
[Kernel] Pipe attn_logits_soft_cap through paged attention TPU kernels#12482alexm-redhat merged 1 commit intovllm-project:mainfrom
Conversation
|
👋 Hi! Thank you for contributing to the vLLM project. Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can do one of these:
🚀 |
Signed-off-by: Fenghui Zhang <fhzhang@google.com>
fenghuizhang
force-pushed
the
fork/main
branch
from
30a6814 to
ef40b6f
Compare
|
cc @vanbasten23 |
alexm-redhat
left a comment
alexm-redhat
left a comment
There was a problem hiding this comment.
@fenghuizhang thanks for doing this! LGTM
| context_lens, | ||
| block_tables, | ||
| pages_per_compute_block, | ||
| megacore_mode=megacore_mode, |
There was a problem hiding this comment.
megacore_mode is checked internally for None?
There was a problem hiding this comment.
Yep, we fixed it on the ptxla side. No workaround needed here.
mgoin
added
tpu
vllm-project#12482) Signed-off-by: Fenghui Zhang <fhzhang@google.com>
vllm-project#12482) Signed-off-by: Fenghui Zhang <fhzhang@google.com> Signed-off-by: Isotr0py <2037008807@qq.com>
vllm-project#12482) Signed-off-by: Fenghui Zhang <fhzhang@google.com>
vllm-project#12482) Signed-off-by: Fenghui Zhang <fhzhang@google.com>
4 participants
Pipe attn_logits_soft_cap through paged_attention, this will unblock some of our models' adoption.
Note that the changed code currently doesn't have unit tests. I will add one later.
This is the same as #12294. We created a new PR as I messed up the sign off and previous commit chain.
Thanks,