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 15, 07:15 UTC
[Bugfix] get_num_blocks_to_allocate with null_block (#19031)
Signed-off-by: Chen Zhang <zhangch99@outlook.com>
Chen Zhang committed
Jun 4, 2025 at 06:30 UTC
b5fd9506c14bed640210a7f3d0adb03a024afdbe
4 files changed
+32
-4
tests/v1/core/test_specialized_manager.py
+23
@@ -144,3 +144,26 @@ def test_sliding_window_remove_skipped_blocks():
144
# of removed blocks should be [1003, 1002].
145
manager.remove_skipped_blocks("test", 11)
146
assert_block_id(block_table, [null_block_id] * 4 + original_block_ids[4:])
147
+
148
+
149
+def test_get_num_blocks_to_allocate():
150
+ block_size = 2
151
+ sliding_window_spec = SlidingWindowSpec(
152
+ block_size=block_size,
153
+ num_kv_heads=1,
154
+ head_size=1,
155
+ dtype=torch.float32,
156
+ sliding_window=4, # Placeholder value, not related to test result
157
+ use_mla=False,
158
+ )
159
+
160
+ block_pool = BlockPool(num_gpu_blocks=100, enable_caching=True)
161
+ manager = get_sliding_window_manager(sliding_window_spec, block_pool)
162
+ cached_blocks_1 = [KVCacheBlock(i + 1) for i in range(10)]
163
+ cached_blocks_2 = [block_pool.null_block for _ in range(5)
164
+ ] + [KVCacheBlock(i + 1) for i in range(5)]
165
+
166
+ assert manager.get_num_blocks_to_allocate("1", 20 * block_size,
167
+ cached_blocks_1) == 20
168
+ assert manager.get_num_blocks_to_allocate("2", 20 * block_size,
169
+ cached_blocks_2) == 15
vllm/v1/core/block_pool.py
+3
-2
@@ -63,6 +63,7 @@ class BlockPool:
63
# The ref_cnt of null_block is not maintained, needs special care to
64
# avoid freeing it.
65
self.null_block = self.free_block_queue.popleft()
66
+ self.null_block.is_null = True
67
68
self.enable_kv_cache_events = enable_kv_cache_events
69
self.kv_event_queue: list[KVCacheEvent] = []
@@ -252,7 +253,7 @@ class BlockPool:
253
for block in blocks:
254
# ref_cnt=0 means this block is in the free list (i.e. eviction
255
# candidate), so remove it.
255
- if block.ref_cnt == 0 and block != self.null_block:
256
+ if block.ref_cnt == 0 and not block.is_null:
257
self.free_block_queue.remove(block)
258
block.incr_ref()
259
@@ -267,7 +268,7 @@ class BlockPool:
268
for block in ordered_blocks:
269
block.decr_ref()
270
# null_block should not be added to the free list.
270
- if block.ref_cnt == 0 and block != self.null_block:
271
+ if block.ref_cnt == 0 and not block.is_null:
272
self.free_block_queue.append(block)
273
274
def reset_prefix_cache(self) -> bool:
vllm/v1/core/kv_cache_utils.py
+3
@@ -125,6 +125,9 @@ class KVCacheBlock:
125
prev_free_block: Optional["KVCacheBlock"] = None
126
next_free_block: Optional["KVCacheBlock"] = None
127
128
+ # Whether the block is a null block that should never be cached.
129
+ is_null: bool = False
130
+
131
def incr_ref(self):
132
self.ref_cnt += 1
133
vllm/v1/core/single_type_kv_cache_manager.py
+3
-2
@@ -83,8 +83,9 @@ class SingleTypeKVCacheManager(ABC):
83
# free queue and ref_cnt == 0), it will be changed from a free block
84
# to a computed block when the request is allocated, so we also count
85
# it as needed to be allocated.
86
- num_evictable_computed_blocks = sum(blk.ref_cnt == 0
87
- for blk in new_computed_blocks)
86
+ num_evictable_computed_blocks = sum(
87
+ blk.ref_cnt == 0 and not blk.is_null
88
+ for blk in new_computed_blocks)
89
return ((num_new_blocks + num_evictable_computed_blocks) *
90
self.num_kv_cache_groups)
91