[V1] Allocate kv_cache with stride order for V1 #18775
Conversation
Signed-off-by: nicklucche <nlucches@redhat.com>
NickLucche
requested review from
WoosukKwon,
robertgshaw2-redhat,
njhill,
ywang96,
comaniac and
alexm-redhat
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 🚀 |
| try: | ||
| kv_cache_stride_order = self.attn_backends[ | ||
| i].get_kv_cache_stride_order() | ||
| assert len(kv_cache_stride_order) == len( | ||
| kv_cache_shape) | ||
| except (AttributeError, NotImplementedError): | ||
| kv_cache_stride_order = tuple( | ||
| range(len(kv_cache_shape))) |
There was a problem hiding this comment.
The try/catch seems a little sketchy. but I don't see a better way.
I thought about handling it in the abstract AttentionBackend, but there we don't know how many axes the kv_cache_shape is.
There was a problem hiding this comment.
yep same thought
| if default_stride == rnd_stride: | ||
| assert all(kv.is_contiguous() for kv in model_runner.kv_caches) | ||
| else: | ||
| assert all(not kv.is_contiguous() for kv in model_runner.kv_caches) |
There was a problem hiding this comment.
Might be a little better to check the exact stride order -- although this is ambiguous if one of the dims is length 1
def stride_order(tensor: torch.Tensor):
# returns a tuple of axis indices sorted from largest stride to smallest
# i.e. slowest-varying dim first, fastest-varying dim last
strides = tensor.stride()
return tuple(sorted(range(tensor.dim()), key=lambda i: strides[i], reverse=True))
|
Pre-commit is actually fine in this PR, issue is with CI itself when running the step Might be worth another spin. |
Signed-off-by: nicklucche <nlucches@redhat.com> Signed-off-by: amit <amit.man@gmail.com>
Signed-off-by: nicklucche <nlucches@redhat.com> Signed-off-by: amit <amit.man@gmail.com>
Signed-off-by: nicklucche <nlucches@redhat.com> Signed-off-by: minpeter <kali2005611@gmail.com>
Basically a port of #16605 for V1. The same motivations still stand, plus we will need to have a different kv cache layout to optimize nixl xfer for the P/D case.