-
-
Notifications
You must be signed in to change notification settings - Fork 10.5k
[Core] Replace empty list with None in KVCacheBlocks for GC optimization #24964
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Conversation
200514f
to
aa2e98d
Compare
Resolve #24321 |
3967910
to
01c383f
Compare
Discussed with @heheda12345 offline. The main discussion point is to come up with solutions to minimize the code signature changes and avoid if else:
Change the PR to draft before addressing the concern. |
98c812b
to
ab0ecb6
Compare
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Nice catch. Can't imagine it can be implemented in such a clean way.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks @Jialin
vllm/v1/core/kv_cache_manager.py
Outdated
if any(blocks): | ||
# Only create new KVCacheBlocks for non-empty blocks | ||
return KVCacheBlocks( | ||
tuple(tuple(block_per_group) for block_per_group in blocks)) | ||
else: | ||
return self.empty_block_list |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
nit: simplify
if any(blocks): | |
# Only create new KVCacheBlocks for non-empty blocks | |
return KVCacheBlocks( | |
tuple(tuple(block_per_group) for block_per_group in blocks)) | |
else: | |
return self.empty_block_list | |
if not any(blocks): | |
return self.empty_block_list | |
# Only create new KVCacheBlocks for non-empty blocks | |
return KVCacheBlocks( | |
tuple(tuple(block_per_group) for block_per_group in blocks)) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Also per other comment, in the non-empty case we should maybe just return KVCacheBlocks(blocks)
? And allow the sub block lists to be tuple or list ...
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I had a quick discussion with @heheda12345 2 weeks ago about the design. As we're creating new precomputed empty blocks, we would love to claim immutability of KVCacheBlocks in order to avoid mis-operations
(e.g. append items on the empty KVCacheBlocks or additional if else checks on types or existence (our first approach was replacing tuple as optional[tuple]).
So if possible, we would love to prefer keeping it as tuple.
This pull request has merge conflicts that must be resolved before it can be |
Signed-off-by: Jialin Ouyang <Jialin.Ouyang@gmail.com>
Signed-off-by: Jialin Ouyang <Jialin.Ouyang@gmail.com>
Signed-off-by: Jialin Ouyang <Jialin.Ouyang@gmail.com>
Signed-off-by: Jialin Ouyang <Jialin.Ouyang@gmail.com>
Signed-off-by: Jialin Ouyang <Jialin.Ouyang@gmail.com>
Signed-off-by: Jialin Ouyang <Jialin.Ouyang@gmail.com>
@Jialin can you rebase the PR? |
Purpose
Majority of the the Generation 0 collect objects are related to KVCacheBlocks, and we found most of them are actually referring to empty blocks.

In this PR, we mainly made 2 changes
Test Plan & Test Result
Model: facebook/opt-125m
Prefill-Heavy workload: Input 2000, Output 48
Decode-Heavy workload: Input 48, Output 2000
We could see, GC costs are significantly smaller with the PR (especially for decode-heavy workload). It could roughly convert to 3-4% throughput improvements.

Essential Elements of an Effective PR Description Checklist
supported_models.md
andexamples
for a new model.