[V0 Deprecation] Refactor kv cache from list to element#37487
Merged
Conversation
Signed-off-by: yewentao256 <zhyanwentao@126.com>
yewentao256
requested review from
ApostaC,
LucasWilkinson,
MatthewBonanni,
NickLucche,
njhill,
orozery,
sighingnow,
tdoublep and
tjtanaa
as code owners
mergify
bot
added
deepseek
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request refactors the kv_cache by removing the outer list wrapper, simplifying its structure from a list of one element to just the element itself (a tensor or a tuple of tensors). This change is consistently applied across various components, including attention layers, mamba-based layers, and their corresponding test files. The modifications simplify code by removing unnecessary [0] indexing when accessing the kv_cache. The change in _cleanup_profiling_kv_cache is a good addition that makes the cleanup logic more robust to the different types of kv_cache. The refactoring appears to be correct and improves code clarity.
Member
|
The NIXL failure seems like it might be relevant? |
hmellor
reviewed
Mar 19, 2026
Comment on lines
+184
to
+194
| @@ -185,15 +185,13 @@ def inject_kv_into_layer( | |||
| if kv_cache_attr is None: | |||
| continue | |||
|
|
|||
| kv_cache_layer = kv_cache_attr[0] | |||
|
|
|||
| filename = self._generate_filename_debug( | |||
| layer_name, request.token_ids, request.mm_hashes | |||
| ) | |||
| kv_cache = safetensors.torch.load_file(filename)["kv_cache"].cuda() | |||
| if isinstance(attn_metadata, dict): | |||
| inject_kv_into_layer( | |||
| kv_cache_layer, | |||
| kv_cache_attr, | |||
Member
There was a problem hiding this comment.
Could we rename this to kv_cache_layer?
Member
Author
There was a problem hiding this comment.
Done, thanks! And also fix the previous CI issue
hmellor
approved these changes
Mar 19, 2026
RhizoNymph
pushed a commit
to RhizoNymph/vllm
that referenced
this pull request
Mar 26, 2026
…#37487) Signed-off-by: yewentao256 <zhyanwentao@126.com>
HenryTangDev
pushed a commit
to HenryTangMain/vllm
that referenced
this pull request
Mar 27, 2026
…#37487) Signed-off-by: yewentao256 <zhyanwentao@126.com>
khairulkabir1661
pushed a commit
to khairulkabir1661/vllm
that referenced
this pull request
Mar 27, 2026
…#37487) Signed-off-by: yewentao256 <zhyanwentao@126.com>
Monishver11
pushed a commit
to Monishver11/vllm
that referenced
this pull request
Mar 27, 2026
…#37487) Signed-off-by: yewentao256 <zhyanwentao@126.com> Signed-off-by: Monishver Chandrasekaran <monishverchandrasekaran@gmail.com>
nithinvc
pushed a commit
to nithinvc/vllm
that referenced
this pull request
Mar 27, 2026
…#37487) Signed-off-by: yewentao256 <zhyanwentao@126.com> Signed-off-by: Nithin Chalapathi <nithin.ch10@gmail.com>
JiantaoXu
pushed a commit
to JiantaoXu/vllm
that referenced
this pull request
Mar 28, 2026
…#37487) Signed-off-by: yewentao256 <zhyanwentao@126.com>
vrdn-23
pushed a commit
to vrdn-23/vllm
that referenced
this pull request
Mar 30, 2026
…#37487) Signed-off-by: yewentao256 <zhyanwentao@126.com> Signed-off-by: Vinay Damodaran <vrdn@hey.com>
EricccYang
pushed a commit
to EricccYang/vllm
that referenced
this pull request
Apr 1, 2026
…#37487) Signed-off-by: yewentao256 <zhyanwentao@126.com> Signed-off-by: EricccYang <yangyang4991@gmail.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Purpose
A follow up for #37195 of removing the virtual engine, this PR further refactor the kv cache from list to element to clean the code
Tests in CI