[Core] Avoid KVCacheBlock.__eq__ invocations in FreeKVCacheBlockQueue #21005
+182
−32
Conversation
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
JialinOuyang-Meta
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 🚀 |
|
This pull request was exported from Phabricator. Differential Revision: D78292345 |
![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 introduces a significant performance optimization to the FreeKVCacheBlockQueue by implementing a doubly linked list with sentinel nodes. This change effectively removes expensive __eq__ comparisons on KVCacheBlock dataclasses, which should improve performance as demonstrated by the new benchmark. The implementation is a classic and well-executed approach.
My review focuses on ensuring the robustness of this new implementation. I've identified a couple of areas where adding validation checks could prevent potential crashes from state inconsistencies, making the system more resilient. These changes should have a negligible performance impact while significantly improving debuggability and correctness guarantees.
JialinOuyang-Meta
added a commit
to JialinOuyang-Meta/vllm-jialino
that referenced
this pull request
Jul 15, 2025
…project#21005) Summary: # Optimizations As a common trick for doubly linked list implementation, introducing fake head and tail nodes would significantly reduce the implementation overhead, and help us to get rid of dataclass.__eq__ comparison easily. - No dataclass.__eq__ invocation - Shorter code - Branchless All these combined should yield significant perf improvement for this piece of code. # Observations Per vLLM profiling, kv_cache_manager.allocate_slots consumed non-negligible cost for each prefill. |{F1980260529}|{F1980260481}|{F1980260497}| By zooming in, we could see the stack of FreeKVCacheBlockQueue.popleft is non-trivial. popleft -> remove -> string.__eq__ which is mainly coming from dataclasses (i.e. KVCacheBlock) equal comparison. Per [dataclasses python library doc](https://docs.python.org/3/library/dataclasses.html#dataclasses.dataclass) ``` dataclasses.dataclass(*, init=True, repr=True, eq=True, order=False, unsafe_hash=False, frozen=False, match_args=True, kw_only=False, slots=False, weakref_slot=False) eq: If true (the default), an __eq__() method will be generated. This method compares the class as if it were a tuple of its fields, in order. Both instances in the comparison must be of the identical type. If the class already defines __eq__(), this parameter is ignored. ``` Test Plan: # Result Typically, block_size is set to 16, so in production usage, we might likely allocate 10-1000 blocks. In this range, the optimization gave us up to ~1ms TTFT savings (the improvements are more significant on long inputs). |After|Before| |{F1980286936}|{F1980286941}| Rollback Plan: Reviewed By: CuiCoco Differential Revision: D78292345
JialinOuyang-Meta
force-pushed
the
export-D78292345
branch
from
574454a to
53e5f05
Compare
JialinOuyang-Meta
added a commit
to JialinOuyang-Meta/vllm-jialino
that referenced
this pull request
Jul 15, 2025
…project#21005) Summary: # Optimizations As a common trick for doubly linked list implementation, introducing fake head and tail nodes would significantly reduce the implementation overhead, and help us to get rid of dataclass.__eq__ comparison easily. - No dataclass.__eq__ invocation - Shorter code - Branchless All these combined should yield significant perf improvement for this piece of code. # Observations Per vLLM profiling, kv_cache_manager.allocate_slots consumed non-negligible cost for each prefill. |{F1980260529}|{F1980260481}|{F1980260497}| By zooming in, we could see the stack of FreeKVCacheBlockQueue.popleft is non-trivial. popleft -> remove -> string.__eq__ which is mainly coming from dataclasses (i.e. KVCacheBlock) equal comparison. Per [dataclasses python library doc](https://docs.python.org/3/library/dataclasses.html#dataclasses.dataclass) ``` dataclasses.dataclass(*, init=True, repr=True, eq=True, order=False, unsafe_hash=False, frozen=False, match_args=True, kw_only=False, slots=False, weakref_slot=False) eq: If true (the default), an __eq__() method will be generated. This method compares the class as if it were a tuple of its fields, in order. Both instances in the comparison must be of the identical type. If the class already defines __eq__(), this parameter is ignored. ``` Test Plan: # Result Typically, block_size is set to 16, so in production usage, we might likely allocate 10-1000 blocks. In this range, the optimization gave us up to ~1ms TTFT savings (the improvements are more significant on long inputs). |After|Before| |{F1980286936}|{F1980286941}| Rollback Plan: Reviewed By: CuiCoco Differential Revision: D78292345
JialinOuyang-Meta
force-pushed
the
export-D78292345
branch
from
53e5f05 to
f0a6c84
Compare
|
This pull request was exported from Phabricator. Differential Revision: D78292345 |
1 similar comment
|
This pull request was exported from Phabricator. Differential Revision: D78292345 |
JialinOuyang-Meta
added a commit
to JialinOuyang-Meta/vllm-jialino
that referenced
this pull request
Jul 15, 2025
…project#21005) Summary: Pull Request resolved: vllm-project#21005 # Optimizations As a common trick for doubly linked list implementation, introducing fake head and tail nodes would significantly reduce the implementation overhead, and help us to get rid of dataclass.__eq__ comparison easily. - No dataclass.__eq__ invocation - Shorter code - Branchless All these combined should yield significant perf improvement for this piece of code. # Observations Per vLLM profiling, kv_cache_manager.allocate_slots consumed non-negligible cost for each prefill. |{F1980260529}|{F1980260481}|{F1980260497}| By zooming in, we could see the stack of FreeKVCacheBlockQueue.popleft is non-trivial. popleft -> remove -> string.__eq__ which is mainly coming from dataclasses (i.e. KVCacheBlock) equal comparison. Per [dataclasses python library doc](https://docs.python.org/3/library/dataclasses.html#dataclasses.dataclass) ``` dataclasses.dataclass(*, init=True, repr=True, eq=True, order=False, unsafe_hash=False, frozen=False, match_args=True, kw_only=False, slots=False, weakref_slot=False) eq: If true (the default), an __eq__() method will be generated. This method compares the class as if it were a tuple of its fields, in order. Both instances in the comparison must be of the identical type. If the class already defines __eq__(), this parameter is ignored. ``` Test Plan: # Result Typically, block_size is set to 16, so in production usage, we might likely allocate 10-1000 blocks. In this range, the optimization gave us up to ~1ms TTFT savings (the improvements are more significant on long inputs). |After|Before| |{F1980286936}|{F1980286941}| Rollback Plan: Reviewed By: CuiCoco Differential Revision: D78292345
JialinOuyang-Meta
force-pushed
the
export-D78292345
branch
from
f0a6c84 to
36411c3
Compare
JialinOuyang-Meta
changed the title
…project#21005) Summary: Pull Request resolved: vllm-project#21005 # Optimizations As a common trick for doubly linked list implementation, introducing fake head and tail nodes would significantly reduce the implementation overhead, and help us to get rid of dataclass.__eq__ comparison easily. - No dataclass.__eq__ invocation - Shorter code - Branchless All these combined should yield significant perf improvement for this piece of code. # Observations Per vLLM profiling, kv_cache_manager.allocate_slots consumed non-negligible cost for each prefill. |{F1980260529}|{F1980260481}|{F1980260497}| By zooming in, we could see the stack of FreeKVCacheBlockQueue.popleft is non-trivial. popleft -> remove -> string.__eq__ which is mainly coming from dataclasses (i.e. KVCacheBlock) equal comparison. Per [dataclasses python library doc](https://docs.python.org/3/library/dataclasses.html#dataclasses.dataclass) ``` dataclasses.dataclass(*, init=True, repr=True, eq=True, order=False, unsafe_hash=False, frozen=False, match_args=True, kw_only=False, slots=False, weakref_slot=False) eq: If true (the default), an __eq__() method will be generated. This method compares the class as if it were a tuple of its fields, in order. Both instances in the comparison must be of the identical type. If the class already defines __eq__(), this parameter is ignored. ``` Test Plan: # Result Typically, block_size is set to 16, so in production usage, we might likely allocate 10-1000 blocks. In this range, the optimization gave us up to ~1ms TTFT savings (the improvements are more significant on long inputs). |After|Before| |{F1980286936}|{F1980286941}| Rollback Plan: Reviewed By: CuiCoco Differential Revision: D78292345 Signed-off-by: Jialin Ouyang <jialino@meta.com>
JialinOuyang-Meta
force-pushed
the
export-D78292345
branch
from
36411c3 to
608f1f5
Compare
Signed-off-by: Jialin Ouyang <jialino@meta.com>
Signed-off-by: Jialin Ouyang <jialino@meta.com>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
Summary:
Optimizations
As a common trick for doubly linked list implementation, introducing fake head and tail nodes would significantly reduce the implementation overhead, and help us to get rid of dataclass.eq comparison easily.
All these combined should yield significant perf improvement for this piece of code.
Observations
Per vLLM profiling, kv_cache_manager.allocate_slots consumed non-negligible cost for each prefill.
|{F1980260529}|{F1980260481}|{F1980260497}|
By zooming in, we could see the stack of FreeKVCacheBlockQueue.popleft is non-trivial. popleft -> remove -> string.eq which is mainly coming from dataclasses (i.e. KVCacheBlock) equal comparison.
Per dataclasses python library doc
Test Plan:
Result
Typically, block_size is set to 16, so in production usage, we might likely allocate 10-1000 blocks. In this range, the optimization gave us up to ~1ms TTFT savings (the improvements are more significant on long inputs).
Benchmark
After
|
Before
Stack
After
Before
Rollback Plan:
Reviewed By: CuiCoco
Differential Revision: D78292345