A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 22:35 UTC
[Misc] Improve code readability of KVCacheManager (#21673)
Signed-off-by: tanruixiang <tanruixiang0104@gmail.com> Signed-off-by: Ruixiang Tan <819464715@qq.com> Signed-off-by: GitHub <noreply@github.com>
Ruixiang Tan committed
Jul 30, 2025 at 22:20 UTC
8f4a1c9a04b36cb7527e67f1fea96c4f05ed0e03
6 files changed
+18
-22
tests/v1/core/test_kv_cache_utils.py
+2
-2
@@ -112,9 +112,9 @@ def test_kv_cache_block():
112
assert block.block_hash is None
113
114
# Test reference count manipulation
115
- block.incr_ref()
115
+ block.ref_cnt += 1
116
assert block.ref_cnt == 1
117
- block.decr_ref()
117
+ block.ref_cnt -= 1
118
assert block.ref_cnt == 0
119
120
# Test block hash setting and resetting
vllm/v1/core/block_pool.py
+1
-1
@@ -276,7 +276,7 @@ class BlockPool:
276
# candidate), so remove it.
277
if block.ref_cnt == 0 and not block.is_null:
278
self.free_block_queue.remove(block)
279
- block.incr_ref()
279
+ block.ref_cnt += 1
280
281
def free_blocks(self, ordered_blocks: Iterable[KVCacheBlock]) -> None:
282
"""Free a list of blocks. The blocks should be ordered by their
vllm/v1/core/kv_cache_coordinator.py
+6
-3
@@ -126,14 +126,17 @@ class KVCacheCoordinator(ABC):
126
def get_num_common_prefix_blocks(self, request_id: str,
127
num_running_requests: int) -> list[int]:
128
"""
129
- Get the number of common prefix blocks for a request.
129
+ Get the number of common prefix blocks for all requests in the RUNNING
130
+ state for each kv cache group.
131
132
Args:
133
request_id: The request ID.
133
- num_running_requests: The number of requests in the RUNNING state.
134
+ num_running_requests: The total number of requests in the RUNNING
135
+ state.
136
137
Returns:
136
- list[int]: The number of common prefix blocks.
138
+ list[int]: The number of common prefix blocks for all requests in
139
+ the RUNNING state for each kv cache group.
140
"""
141
num_blocks_per_group = [
142
manager.get_num_common_prefix_blocks(request_id,
vllm/v1/core/kv_cache_manager.py
+1
-4
@@ -170,10 +170,6 @@ class KVCacheManager:
170
self.block_size, request)
171
self.req_to_block_hashes[request.request_id] = block_hashes
172
173
- if self.log_stats:
174
- assert self.prefix_cache_stats is not None
175
- self.prefix_cache_stats.requests += 1
176
-
173
# NOTE: When all tokens hit the cache, we must recompute the last token
174
# to obtain logits. Thus, set max_cache_hit_length to prompt_length - 1.
175
# This can trigger recomputation of an entire block, rather than just
@@ -187,6 +183,7 @@ class KVCacheManager:
183
184
if self.log_stats:
185
assert self.prefix_cache_stats is not None
186
+ self.prefix_cache_stats.requests += 1
187
self.prefix_cache_stats.queries += request.num_tokens
188
self.prefix_cache_stats.hits += num_new_computed_tokens
189
vllm/v1/core/kv_cache_utils.py
-8
@@ -154,14 +154,6 @@ class KVCacheBlock:
154
# Whether the block is a null block that should never be cached.
155
is_null: bool = False
156
157
- # TODO(Jialin): For performance, let callers handle ref_cnt bumps to
158
- # avoid function calls.
159
- def incr_ref(self):
160
- self.ref_cnt += 1
161
-
162
- def decr_ref(self):
163
- self.ref_cnt -= 1
164
-
157
@property
158
def block_hash(self) -> Optional[BlockHashWithGroupId]:
159
return self._block_hash
vllm/v1/core/single_type_kv_cache_manager.py
+8
-4
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import itertools
4
from abc import ABC, abstractmethod
5
from collections import defaultdict
6
from typing import Callable
@@ -177,14 +178,17 @@ class SingleTypeKVCacheManager(ABC):
178
def get_num_common_prefix_blocks(self, request_id: str,
179
num_running_requests: int) -> int:
180
"""
180
- Get the number of common prefix blocks for a request.
181
+ Get the number of common prefix blocks for all requests in the RUNNING
182
+ state.
183
184
Args:
185
request_id: The request ID.
184
- num_running_requests: The number of requests in the RUNNING state.
186
+ num_running_requests: The total number of requests in the RUNNING
187
+ state.
188
189
Returns:
187
- The number of common prefix blocks.
190
+ The number of common prefix blocks for all requests in the RUNNING
191
+ state.
192
"""
193
194
raise NotImplementedError
@@ -264,7 +268,7 @@ class FullAttentionManager(SingleTypeKVCacheManager):
268
computed_blocks: tuple[list[KVCacheBlock], ...] = tuple(
269
[] for _ in range(len(kv_cache_group_ids)))
270
max_num_blocks = max_length // kv_cache_spec.block_size
267
- for i, block_hash in zip(range(max_num_blocks), block_hashes):
271
+ for block_hash in itertools.islice(block_hashes, max_num_blocks):
272
# block_hashes is a chain of block hashes. If a block hash is not
273
# in the cached_block_hash_to_id, the following block hashes are
274
# not computed yet for sure.