[TPU] fix kv_cache_update kernel block size choosing logic #21007
Conversation
Signed-off-by: Chengji Yao <chengjiyao@google.com>
yaochengji
requested review from
WoosukKwon,
alexm-redhat,
comaniac,
njhill,
robertgshaw2-redhat and
ywang96
as code owners
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels. Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run 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 either: Add 🚀 |
|
@bythew3i @tengyifei could you take a look? |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request fixes a TPU CI test by adjusting the KV cache update logic. A critical issue was identified where the application could crash with a ZeroDivisionError when running attention-free models on TPU, due to page_size_bytes being zero. A fix has been suggested to handle this edge case.
Signed-off-by: Chengji Yao <chengjiyao@google.com>
Signed-off-by: Chengji Yao <chengjiyao@google.com>
yaochengji
force-pushed
the
chengji/fix-ci-test
branch
from
1c81513 to
4de8313
Compare
|
Could you summarize what problem this PR is addressing? |
|
|
||
| kv_cache_dtype_bits = dtype_bits(kv_cache_dtype) | ||
| return (block_size * num_combined_kv_heads * padded_head_size * | ||
| kv_cache_dtype_bits // 8) |
There was a problem hiding this comment.
what is 8 in the kv_cache_dtype_bits // 8?
There was a problem hiding this comment.
1 bytes = 8 bits
| vmem_limit = 32 * 1024 * 1024 | ||
| # The default vmem_limit_bytes of a pallas kernel is 32MB. Here we | ||
| # calculate num_slices_per_block based on 16MB in case any register spills. | ||
| vmem_limit = 16 * 1024 * 1024 |
There was a problem hiding this comment.
I wonder why we reduce the vmem_limit from 32mb to 16mb?
There was a problem hiding this comment.
As noted in the comment, there might be register spills.
@vanbasten23 It is used fix the potential VMEM OOM due to incorrect vmem estimation. |
…ect#21007) Signed-off-by: Chengji Yao <chengjiyao@google.com> Signed-off-by: x22x22 <wadeking@qq.com>
…ect#21007) Signed-off-by: Chengji Yao <chengjiyao@google.com>
…ect#21007) Signed-off-by: Chengji Yao <chengjiyao@google.com>
…ect#21007) Signed-off-by: Chengji Yao <chengjiyao@google.com> Signed-off-by: Jinzhen Lin <linjinzhen@hotmail.com>
…ect#21007) Signed-off-by: Chengji Yao <chengjiyao@google.com> Signed-off-by: Paul Pak <paulpak58@gmail.com>
…ect#21007) Signed-off-by: Chengji Yao <chengjiyao@google.com> Signed-off-by: Diego-Castan <diego.castan@ibm.com>
…ect#21007) Signed-off-by: Chengji Yao <chengjiyao@google.com> Signed-off-by: taehyun <taehyun.kim@motiftech.io>
…ect#21007) Signed-off-by: Chengji Yao <chengjiyao@google.com>
4 participants
Purpose
Fix TPU CI test.
Test Plan
pytest -s -v tests/v1/tpu/test_basic.py
Test Result
Passed.
(Optional) Documentation Update