Allocate kv_cache with stride order#16605
Conversation
|
👋 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 🚀 |
LucasWilkinson
left a comment
There was a problem hiding this comment.
Thanks for doing this! this is looking much better! Left a few nits, please also fix the pre-commit failures
7cd508d to
8a6c7fe
Compare
LucasWilkinson
left a comment
There was a problem hiding this comment.
LGTM, thanks for the contribution!
c13e5d6 to
343378f
Compare
|
This pull request has merge conflicts that must be resolved before it can be |
Head branch was pushed to by a user without write access
18d6e6e to
aba3c45
Compare
|
Can you merge from main to fix docker build? |
Signed-off-by: shuw <shuw@nvidia.com>
aba3c45 to
33d5969
Compare
Done. |
|
Please check whether the failing kernels test is related to this PR |
Signed-off-by: shuw <shuw@nvidia.com>
7dbc7ef to
a48be76
Compare
Only test_cache.py failure is related. Fixed by reducing tensor size to avoid OOM. The spec-decoding tests should be irrelevant. |
Signed-off-by: shuw <shuw@nvidia.com>
Signed-off-by: shuw <shuw@nvidia.com>
Signed-off-by: shuw <shuw@nvidia.com> Signed-off-by: Agata Dobrzyniewicz <adobrzyniewicz@habana.ai>
Signed-off-by: shuw <shuw@nvidia.com> Signed-off-by: Mu Huai <tianbowen.tbw@antgroup.com>
Signed-off-by: shuw <shuw@nvidia.com> Signed-off-by: Yuqi Zhang <yuqizhang@google.com>
Allow KV cache manager to support an stride order to the allocation which the attention backend could provide. Mainly affect Flashinfer backend. Ref. #8200
@tlrmchlsmth @LucasWilkinson