@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 6bc6f2d86d

[1/N][Core] add partial prefix cache primitives (#45939)

Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai>

Jiangyun Zhu committed Jun 22, 2026 at 14:43 UTC 6bc6f2d86d7800f878a4a13495b51a6e3b728c37
5 files changed +719 -41
tests/v1/core/prefix_cache/test_partial_prefix_cache_primitives.py new
+460
@@ -0,0 +1,460 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from collections.abc import Callable
5 +
6 +import pytest
7 +
8 +import vllm.v1.core.kv_cache_utils as kv_cache_utils
9 +from vllm.distributed.kv_events import BlockRemoved, BlockStored
10 +from vllm.sampling_params import SamplingParams
11 +from vllm.utils.hashing import sha256
12 +from vllm.v1.core.block_pool import BlockPool
13 +from vllm.v1.core.kv_cache_utils import (
14 + BlockHash,
15 + BlockHashListWithBlockSize,
16 + KVCacheBlock,
17 + get_request_block_hasher,
18 + hash_block_tokens,
19 + init_none_hash,
20 +)
21 +from vllm.v1.request import Request
22 +
23 +pytestmark = pytest.mark.cpu_test
24 +
25 +
26 +@pytest.fixture(autouse=True)
27 +def _auto_init_hash_fn():
28 + init_none_hash(sha256)
29 +
30 +
31 +def make_request(
32 + request_id: str,
33 + prompt_token_ids: list[int],
34 + hash_block_size: int,
35 + hash_fn: Callable,
36 +) -> Request:
37 + sampling_params = SamplingParams(max_tokens=17)
38 + sampling_params.update_from_generation_config({}, eos_token_id=100)
39 + return Request(
40 + request_id=request_id,
41 + prompt_token_ids=prompt_token_ids,
42 + sampling_params=sampling_params,
43 + pooling_params=None,
44 + block_hasher=get_request_block_hasher(hash_block_size, hash_fn),
45 + )
46 +
47 +
48 +def boundary_hash(req: Request, hash_block_size: int, num_tokens: int) -> BlockHash:
49 + # Every boundary at a hash_block_size multiple is just the fine-grained
50 + # chain hash ending there.
51 + return req.block_hashes[num_tokens // hash_block_size - 1]
52 +
53 +
54 +def cache_full_block_and_partial_tail(
55 + token_ids: list[int],
56 + *,
57 + enable_kv_cache_events: bool = False,
58 +) -> tuple[BlockPool, Request, list[KVCacheBlock], BlockHash]:
59 + hash_block_size = 2
60 + block_size = 6
61 + kv_cache_group_id = 0
62 + req = make_request("0", token_ids, hash_block_size, sha256)
63 + pool = BlockPool(
64 + num_gpu_blocks=3,
65 + enable_caching=True,
66 + hash_block_size=hash_block_size,
67 + enable_kv_cache_events=enable_kv_cache_events,
68 + )
69 + blocks = pool.get_new_blocks(2)
70 +
71 + pool.cache_full_blocks(
72 + request=req,
73 + blocks=blocks,
74 + num_cached_blocks=0,
75 + num_full_blocks=1,
76 + block_size=block_size,
77 + kv_cache_group_id=kv_cache_group_id,
78 + )
79 + partial_hash = boundary_hash(req, hash_block_size, len(token_ids))
80 + assert pool.cache_partial_block(
81 + request=req,
82 + block=blocks[1],
83 + num_tokens=len(token_ids),
84 + kv_cache_group_id=kv_cache_group_id,
85 + block_size=block_size,
86 + )
87 + return pool, req, blocks, partial_hash
88 +
89 +
90 +def test_boundary_hashes_reuse_fine_grained_chain():
91 + hash_block_size = 2
92 + block_size = 6
93 + token_ids = [0, 0, 1, 1, 2, 2, 3, 3, 4, 4]
94 + req = make_request("0", token_ids, hash_block_size, sha256)
95 +
96 + coarse = BlockHashListWithBlockSize(req.block_hashes, hash_block_size, block_size)
97 + # The block_size=6 full-block hash is the fine hash at the 6-token boundary,
98 + # not a concatenation of the three fine hashes inside the block.
99 + assert coarse[0] == req.block_hashes[6 // hash_block_size - 1]
100 + assert coarse[0] != BlockHash(
101 + req.block_hashes[0] + req.block_hashes[1] + req.block_hashes[2]
102 + )
103 + # A partial tail at 10 tokens is the fine hash at the 10-token boundary,
104 + # which chains over the entire prefix.
105 + tail_hash = boundary_hash(req, hash_block_size, 10)
106 + assert tail_hash == req.block_hashes[4]
107 + assert tail_hash == hash_block_tokens(sha256, req.block_hashes[3], token_ids[8:10])
108 +
109 +
110 +def test_cache_partial_block_kv_cache_events():
111 + hash_block_size = 4
112 + block_size = 12
113 + kv_cache_group_id = 2
114 +
115 + pool = BlockPool(
116 + num_gpu_blocks=2,
117 + enable_caching=True,
118 + hash_block_size=hash_block_size,
119 + enable_kv_cache_events=True,
120 + )
121 + req = make_request(
122 + "req_partial_events",
123 + prompt_token_ids=list(range(hash_block_size * 2)),
124 + hash_block_size=hash_block_size,
125 + hash_fn=sha256,
126 + )
127 +
128 + block = pool.get_new_blocks(1)[0]
129 + partial_entry_hash = pool.cache_partial_block(
130 + request=req,
131 + block=block,
132 + num_tokens=hash_block_size * 2,
133 + kv_cache_group_id=kv_cache_group_id,
134 + block_size=block_size,
135 + )
136 +
137 + events = pool.take_events()
138 + assert len(events) == 1
139 + stored_event = events[0]
140 + assert isinstance(stored_event, BlockStored)
141 + assert partial_entry_hash is not None
142 + assert stored_event.block_hashes == [
143 + kv_cache_utils.maybe_convert_block_hash(req.block_hashes[1])
144 + ]
145 + assert stored_event.parent_block_hash == kv_cache_utils.maybe_convert_block_hash(
146 + req.block_hashes[0]
147 + )
148 + assert stored_event.token_ids == req.all_token_ids[hash_block_size:]
149 + assert stored_event.block_size == 4
150 + assert stored_event.group_idx == kv_cache_group_id
151 +
152 + duplicate_entry_hash = pool.cache_partial_block(
153 + request=req,
154 + block=block,
155 + num_tokens=hash_block_size * 2,
156 + kv_cache_group_id=kv_cache_group_id,
157 + block_size=block_size,
158 + )
159 + assert duplicate_entry_hash == partial_entry_hash
160 + assert pool.take_events() == []
161 +
162 + pool.free_blocks([block])
163 + pool.get_new_blocks(1)
164 + events = pool.take_events()
165 + assert len(events) == 1
166 + removed_event = events[0]
167 + assert isinstance(removed_event, BlockRemoved)
168 + assert removed_event.block_hashes == stored_event.block_hashes
169 + assert removed_event.group_idx == kv_cache_group_id
170 +
171 +
172 +def test_partial_block_replacement_emits_remove_then_store_events():
173 + hash_block_size = 2
174 + block_size = 6
175 + kv_cache_group_id = 0
176 + req = make_request("0", [0, 0, 1, 1, 2, 2, 3, 3], hash_block_size, sha256)
177 + pool = BlockPool(
178 + num_gpu_blocks=3,
179 + enable_caching=True,
180 + hash_block_size=hash_block_size,
181 + enable_kv_cache_events=True,
182 + )
183 + blocks = pool.get_new_blocks(2)
184 +
185 + pool.cache_full_blocks(
186 + request=req,
187 + blocks=blocks,
188 + num_cached_blocks=0,
189 + num_full_blocks=1,
190 + block_size=block_size,
191 + kv_cache_group_id=kv_cache_group_id,
192 + )
193 + partial_hash_8 = boundary_hash(req, hash_block_size, 8)
194 + assert pool.cache_partial_block(
195 + request=req,
196 + block=blocks[1],
197 + num_tokens=8,
198 + kv_cache_group_id=kv_cache_group_id,
199 + block_size=block_size,
200 + )
201 + assert pool.get_cached_block(partial_hash_8, [kv_cache_group_id]) == [blocks[1]]
202 + pool.take_events()
203 +
204 + req.append_output_token_ids([4, 4])
205 + partial_hash_10 = boundary_hash(req, hash_block_size, 10)
206 + assert pool.cache_partial_block(
207 + request=req,
208 + block=blocks[1],
209 + num_tokens=10,
210 + kv_cache_group_id=kv_cache_group_id,
211 + block_size=block_size,
212 + )
213 + events = pool.take_events()
214 +
215 + assert len(events) == 2
216 + removed_event, stored_event = events
217 + assert isinstance(removed_event, BlockRemoved)
218 + assert removed_event.block_hashes == [
219 + kv_cache_utils.maybe_convert_block_hash(partial_hash_8)
220 + ]
221 + assert removed_event.group_idx == kv_cache_group_id
222 + assert isinstance(stored_event, BlockStored)
223 + assert stored_event.block_hashes == [
224 + kv_cache_utils.maybe_convert_block_hash(partial_hash_10)
225 + ]
226 + assert stored_event.parent_block_hash == kv_cache_utils.maybe_convert_block_hash(
227 + boundary_hash(req, hash_block_size, 8)
228 + )
229 + assert stored_event.token_ids == req.all_token_ids[8:10]
230 + assert stored_event.block_size == hash_block_size
231 + assert stored_event.group_idx == kv_cache_group_id
232 + assert pool.get_cached_block(partial_hash_8, [kv_cache_group_id]) is None
233 + assert pool.get_cached_block(partial_hash_10, [kv_cache_group_id]) == [blocks[1]]
234 +
235 +
236 +def test_later_request_hits_cached_partial_tail():
237 + hash_block_size = 2
238 + block_size = 6
239 + kv_cache_group_id = 0
240 + cached_token_ids = [0, 0, 1, 1, 2, 2, 3, 3, 4, 4]
241 + req = make_request("0", cached_token_ids, hash_block_size, sha256)
242 + pool = BlockPool(
243 + num_gpu_blocks=3,
244 + enable_caching=True,
245 + hash_block_size=hash_block_size,
246 + )
247 + blocks = pool.get_new_blocks(2)
248 +
249 + pool.cache_full_blocks(
250 + request=req,
251 + blocks=blocks,
252 + num_cached_blocks=0,
253 + num_full_blocks=1,
254 + block_size=block_size,
255 + kv_cache_group_id=kv_cache_group_id,
256 + )
257 + partial_hash_10 = boundary_hash(req, hash_block_size, 10)
258 + assert pool.cache_partial_block(
259 + request=req,
260 + block=blocks[1],
261 + num_tokens=10,
262 + kv_cache_group_id=kv_cache_group_id,
263 + block_size=block_size,
264 + )
265 +
266 + replay = make_request("1", cached_token_ids, hash_block_size, sha256)
267 + replay_hash_10 = boundary_hash(replay, hash_block_size, 10)
268 + assert replay_hash_10 == partial_hash_10
269 + assert pool.get_cached_block(replay_hash_10, [kv_cache_group_id]) == [blocks[1]]
270 +
271 + extended = make_request("2", cached_token_ids + [10], hash_block_size, sha256)
272 + extended_hash_10 = boundary_hash(extended, hash_block_size, 10)
273 + assert extended_hash_10 == partial_hash_10
274 + assert pool.get_cached_block(extended_hash_10, [kv_cache_group_id]) == [blocks[1]]
275 +
276 +
277 +def test_cache_partial_block_uses_fine_grained_boundary_hash():
278 + hash_block_size = 2
279 + block_size = 6
280 + kv_cache_group_id = 0
281 + token_ids = [0, 0, 1, 1, 2, 2, 3, 3, 4, 4]
282 + req = make_request("0", token_ids, hash_block_size, sha256)
283 + pool = BlockPool(
284 + num_gpu_blocks=3,
285 + enable_caching=True,
286 + hash_block_size=hash_block_size,
287 + )
288 + blocks = pool.get_new_blocks(2)
289 +
290 + pool.cache_full_blocks(
291 + request=req,
292 + blocks=blocks,
293 + num_cached_blocks=0,
294 + num_full_blocks=1,
295 + block_size=block_size,
296 + kv_cache_group_id=kv_cache_group_id,
297 + )
298 +
299 + partial_entry_hash = pool.cache_partial_block(
300 + request=req,
301 + block=blocks[1],
302 + num_tokens=10,
303 + kv_cache_group_id=kv_cache_group_id,
304 + block_size=block_size,
305 + )
306 + # The partial entry is keyed by the fine-grained hash at the 10-token
307 + # boundary, regardless of the owning group's block_size.
308 + expected = boundary_hash(req, hash_block_size, 10)
309 + assert partial_entry_hash == kv_cache_utils.make_block_hash_with_group_id(
310 + expected, kv_cache_group_id
311 + )
312 + assert pool.get_cached_block(expected, [kv_cache_group_id]) == [blocks[1]]
313 +
314 +
315 +def test_cache_partial_block_requires_hash_boundary():
316 + hash_block_size = 2
317 + block_size = 4
318 + req = make_request("0", [0, 0, 1, 1], hash_block_size, sha256)
319 + pool = BlockPool(
320 + num_gpu_blocks=2,
321 + enable_caching=True,
322 + hash_block_size=hash_block_size,
323 + )
324 + block = pool.get_new_blocks(1)[0]
325 +
326 + with pytest.raises(AssertionError):
327 + pool.cache_partial_block(
328 + request=req,
329 + block=block,
330 + num_tokens=3,
331 + kv_cache_group_id=0,
332 + block_size=block_size,
333 + )
334 +
335 +
336 +def test_cache_partial_block_duplicate_checks_all_blocks_for_hash():
337 + hash_block_size = 2
338 + block_size = 4
339 + kv_cache_group_id = 0
340 + req = make_request("0", [0, 0, 1, 1], hash_block_size, sha256)
341 + pool = BlockPool(
342 + num_gpu_blocks=4,
343 + enable_caching=True,
344 + hash_block_size=hash_block_size,
345 + )
346 + blocks = pool.get_new_blocks(2)
347 +
348 + first_entry_hash = pool.cache_partial_block(
349 + request=req,
350 + block=blocks[0],
351 + num_tokens=2,
352 + kv_cache_group_id=kv_cache_group_id,
353 + block_size=block_size,
354 + )
355 + second_entry_hash = pool.cache_partial_block(
356 + request=req,
357 + block=blocks[1],
358 + num_tokens=2,
359 + kv_cache_group_id=kv_cache_group_id,
360 + block_size=block_size,
361 + )
362 + assert first_entry_hash == second_entry_hash
363 +
364 + duplicate_entry_hash = pool.cache_partial_block(
365 + request=req,
366 + block=blocks[1],
367 + num_tokens=2,
368 + kv_cache_group_id=kv_cache_group_id,
369 + block_size=block_size,
370 + )
371 + assert duplicate_entry_hash == second_entry_hash
372 + assert pool.cached_block_hashes_by_block == {}
373 +
374 +
375 +def test_reset_prefix_cache_clears_partial_entry_metadata():
376 + pool, req, blocks, partial_hash_10 = cache_full_block_and_partial_tail(
377 + [0, 0, 1, 1, 2, 2, 3, 3, 4, 4]
378 + )
379 + full_hash = BlockHashListWithBlockSize(req.block_hashes, 2, 6)[0]
380 +
381 + assert pool.get_cached_block(full_hash, [0]) == [blocks[0]]
382 + assert pool.get_cached_block(partial_hash_10, [0]) == [blocks[1]]
383 +
384 + pool.free_blocks(blocks)
385 + assert pool.reset_prefix_cache()
386 +
387 + assert pool.get_cached_block(full_hash, [0]) is None
388 + assert pool.get_cached_block(partial_hash_10, [0]) is None
389 + assert pool.cached_block_hashes_by_block == {}
390 +
391 +
392 +def test_evict_cached_block_removes_full_hash_and_partial_entry():
393 + pool, req, blocks, partial_hash_10 = cache_full_block_and_partial_tail(
394 + [0, 0, 1, 1, 2, 2, 3, 3, 4, 4]
395 + )
396 + full_hash = BlockHashListWithBlockSize(req.block_hashes, 2, 6)[0]
397 +
398 + assert pool.get_cached_block(full_hash, [0]) == [blocks[0]]
399 + assert pool.get_cached_block(partial_hash_10, [0]) == [blocks[1]]
400 +
401 + pool.evict_blocks({blocks[0].block_id, blocks[1].block_id})
402 +
403 + assert pool.get_cached_block(full_hash, [0]) is None
404 + assert pool.get_cached_block(partial_hash_10, [0]) is None
405 + assert pool.cached_block_hashes_by_block == {}
406 +
407 +
408 +def test_partial_block_promotes_to_direct_full_block_hash():
409 + hash_block_size = 2
410 + block_size = 6
411 + kv_cache_group_id = 0
412 + token_ids = [0, 0, 1, 1, 2, 2, 3, 3, 4, 4]
413 + req = make_request("0", token_ids, hash_block_size, sha256)
414 + pool = BlockPool(
415 + num_gpu_blocks=3,
416 + enable_caching=True,
417 + hash_block_size=hash_block_size,
418 + )
419 + blocks = pool.get_new_blocks(2)
420 +
421 + pool.cache_full_blocks(
422 + request=req,
423 + blocks=blocks,
424 + num_cached_blocks=0,
425 + num_full_blocks=1,
426 + block_size=block_size,
427 + kv_cache_group_id=kv_cache_group_id,
428 + )
429 + partial_hash_10 = boundary_hash(req, hash_block_size, 10)
430 + assert pool.cache_partial_block(
431 + request=req,
432 + block=blocks[1],
433 + num_tokens=10,
434 + kv_cache_group_id=kv_cache_group_id,
435 + block_size=block_size,
436 + )
437 + assert pool.get_cached_block(partial_hash_10, [kv_cache_group_id]) == [blocks[1]]
438 +
439 + req.append_output_token_ids([5, 5])
440 + full_hashes = BlockHashListWithBlockSize(
441 + req.block_hashes, hash_block_size, block_size
442 + )
443 + promoted_full_hash = full_hashes[1]
444 + # The promoted full-block hash is the fine hash at the 12-token boundary,
445 + # not a concatenation of the fine hashes inside the block.
446 + assert promoted_full_hash == req.block_hashes[12 // hash_block_size - 1]
447 + assert promoted_full_hash != BlockHash(
448 + req.block_hashes[3] + req.block_hashes[4] + req.block_hashes[5]
449 + )
450 +
451 + pool.cache_full_blocks(
452 + request=req,
453 + blocks=blocks,
454 + num_cached_blocks=1,
455 + num_full_blocks=2,
456 + block_size=block_size,
457 + kv_cache_group_id=kv_cache_group_id,
458 + )
459 + assert pool.get_cached_block(promoted_full_hash, [kv_cache_group_id]) == [blocks[1]]
460 + assert pool.get_cached_block(partial_hash_10, [kv_cache_group_id]) is None
tests/v1/core/test_kv_cache_utils.py
+1 -1
@@ -225,7 +225,7 @@ def test_kv_cache_block():
225
226 # Test block hash setting and resetting
227 block_hash = make_block_hash_with_group_id(BlockHash(b"abc"), 0)
228 - block.block_hash = block_hash
228 + block.set_block_hash(block_hash)
229 assert block.block_hash == block_hash
230
231 block.reset_hash()
tests/v1/core/test_prefix_caching.py
+1 -1
@@ -2003,7 +2003,7 @@ def test_maybe_evict_cached_block():
2003 assert len(pool.blocks) == len(block_hashes)
2004 # Manually add all blocks to cached_blocks
2005 for block, block_hash in zip(pool.blocks, block_hashes):
2006 - block.block_hash = block_hash
2006 + block.set_block_hash(block_hash)
2007 pool.cached_block_hash_to_block.insert(block_hash, block)
2008
2009 block0, block1, block2, block3 = pool.blocks
vllm/v1/core/block_pool.py
+219 -23
@@ -72,6 +72,20 @@ class BlockHashToBlockMap:
72 self._unexpected_blocks_type(blocks)
73 return None
74
75 + def contain(self, key: BlockHashWithGroupId, block_id: int) -> bool:
76 + """
77 + Checks whether the key maps to the given block ID.
78 + """
79 + blocks = self._cache.get(key)
80 + if blocks is None:
81 + return False
82 + if isinstance(blocks, KVCacheBlock):
83 + return blocks.block_id == block_id
84 + if isinstance(blocks, dict):
85 + return block_id in blocks
86 + self._unexpected_blocks_type(blocks)
87 + return False
88 +
89 def insert(self, key: BlockHashWithGroupId, block: KVCacheBlock) -> None:
90 """
91 Inserts the KVCacheBlock to the cache
@@ -169,6 +183,7 @@ class BlockPool:
183
184 # Cache for block lookup
185 self.cached_block_hash_to_block: BlockHashToBlockMap = BlockHashToBlockMap()
186 + self.cached_block_hashes_by_block: dict[int, set[BlockHashWithGroupId]] = {}
187
188 # To represent a placeholder block with block_id=0.
189 # The ref_cnt of null_block is not maintained, needs special care to
@@ -245,7 +260,6 @@ class BlockPool:
260 if num_cached_blocks >= num_full_blocks:
261 return
262 new_full_blocks = blocks[num_cached_blocks:num_full_blocks]
248 - assert len(request.block_hashes) >= num_full_blocks
263 assert block_mask is None or len(block_mask) == len(new_full_blocks)
264 if block_size == self.hash_block_size:
265 # Common case.
@@ -254,11 +268,10 @@ class BlockPool:
268 # block_size is a multiple of hash_block_size. This happens when
269 # different KV cache groups have different block sizes.
270 assert block_size % self.hash_block_size == 0
257 - # Recalculate block_hashes at the granularity of block_size, using
258 - # the original block_hashes (at the granularity of hash_block_size).
271 block_hashes = BlockHashListWithBlockSize(
272 request.block_hashes, self.hash_block_size, block_size
273 )
274 + assert len(block_hashes) >= num_full_blocks
275
276 new_block_hashes = block_hashes[num_cached_blocks:]
277 new_hashes: list[ExternalBlockHash] | None = (
@@ -270,15 +283,27 @@ class BlockPool:
283 # in align mode. We skip null blocks here.
284 if blk.is_null or (block_mask is not None and not block_mask[i]):
285 continue
273 - assert blk.block_hash is None
286 block_hash = new_block_hashes[i]
287 + num_hash_tokens = (num_cached_blocks + i + 1) * block_size
288
289 # Update and added the full block to the cache.
290 block_hash_with_group_id = make_block_hash_with_group_id(
291 block_hash, kv_cache_group_id
292 )
280 - blk.block_hash = block_hash_with_group_id
281 - self.cached_block_hash_to_block.insert(block_hash_with_group_id, blk)
293 + if blk.block_hash is not None:
294 + # The only valid case where a "new full block" already has a
295 + # hash is partial->full promotion of the same cache block.
296 + assert (
297 + blk.block_hash_num_tokens is not None
298 + and blk.block_hash_num_tokens < num_hash_tokens
299 + )
300 + removed_hashes = self._remove_cached_block_hashes(blk)
301 + self._emit_block_removed_events(removed_hashes)
302 + self._insert_block_hash(
303 + block_hash_with_group_id,
304 + blk,
305 + num_tokens=num_hash_tokens,
306 + )
307 if new_hashes is not None:
308 new_hashes.append(maybe_convert_block_hash(block_hash))
309
@@ -330,6 +355,190 @@ class BlockPool:
355 )
356 )
357
358 + def cache_partial_block(
359 + self,
360 + request: Request,
361 + block: KVCacheBlock,
362 + num_tokens: int,
363 + kv_cache_group_id: int,
364 + block_size: int,
365 + ) -> BlockHashWithGroupId | None:
366 + """Register a partial prefix-cache entry for an existing block.
367 +
368 + Prefix-cache keys normally identify full cache blocks. A partial entry
369 + makes an existing cache block reachable from a fine-grained prefix
370 + boundary inside that block without allocating or copying a new
371 + ``KVCacheBlock``.
372 +
373 + The partial entry is lookup metadata owned by ``block``. If ``block``
374 + has no primary hash, the key becomes its primary hash. If the block
375 + already has a primary hash, the partial entry is tracked in
376 + ``cached_block_hashes_by_block`` so eviction, reset, and promotion can
377 + remove every hash key that points to the block.
378 +
379 + Args:
380 + request: Request whose token IDs and block hashes define the
381 + partial entry.
382 + block: Existing cache block to make reachable from the partial
383 + prefix boundary.
384 + num_tokens: Prefix length represented by the partial entry. It
385 + must be a positive multiple of ``self.hash_block_size`` and
386 + cannot exceed the request's computed block hashes.
387 + kv_cache_group_id: KV cache group that owns the partial entry.
388 + block_size: Cache block size for the owning group. The partial
389 + entry hash itself is always the prefix-chain hash at
390 + ``num_tokens``; ``block_size`` is used to assert that the
391 + entry is partial within the owning cache block.
392 +
393 + Returns:
394 + The hash key with group ID if a partial entry can be registered;
395 + otherwise ``None`` for null blocks.
396 + """
397 + if block.is_null:
398 + return None
399 +
400 + assert block_size > self.hash_block_size
401 + assert block_size % self.hash_block_size == 0
402 + assert num_tokens % block_size != 0
403 + block_hash = self._get_partial_block_hash(request, num_tokens)
404 + num_hash_blocks = num_tokens // self.hash_block_size
405 + block_hash_with_group_id = make_block_hash_with_group_id(
406 + block_hash, kv_cache_group_id
407 + )
408 + already_cached = block.block_hash == block_hash_with_group_id or (
409 + self.cached_block_hash_to_block.contain(
410 + block_hash_with_group_id, block.block_id
411 + )
412 + )
413 + if (
414 + not already_cached
415 + and block.block_hash is not None
416 + and block.block_hash_num_tokens is not None
417 + and block.block_hash_num_tokens < num_hash_blocks * self.hash_block_size
418 + ):
419 + removed_hashes = self._remove_cached_block_hashes(block)
420 + self._emit_block_removed_events(removed_hashes)
421 + self._insert_block_hash(
422 + block_hash_with_group_id,
423 + block,
424 + num_tokens=num_hash_blocks * self.hash_block_size,
425 + )
426 + if self.enable_kv_cache_events and not already_cached:
427 + parent_hash, block_start = self._get_partial_block_parent_hash_and_start(
428 + request, num_tokens
429 + )
430 + parent_block_hash = (
431 + maybe_convert_block_hash(parent_hash)
432 + if parent_hash is not None
433 + else None
434 + )
435 + block_end = num_tokens
436 + curr_mm_idx = -1 if block_start > 0 else 0
437 + extra_keys, _ = generate_block_hash_extra_keys(
438 + request, block_start, block_end, curr_mm_idx
439 + )
440 + self.kv_event_queue.append(
441 + BlockStored(
442 + block_hashes=[maybe_convert_block_hash(block_hash)],
443 + parent_block_hash=parent_block_hash,
444 + token_ids=request.all_token_ids[block_start:block_end],
445 + block_size=block_end - block_start,
446 + lora_id=request.lora_request.adapter_id
447 + if request.lora_request
448 + else None,
449 + medium=MEDIUM_GPU,
450 + lora_name=request.lora_request.name
451 + if request.lora_request
452 + else None,
453 + extra_keys=[extra_keys],
454 + group_idx=kv_cache_group_id,
455 + )
456 + )
457 + return block_hash_with_group_id
458 +
459 + def _get_partial_block_hash(
460 + self,
461 + request: Request,
462 + num_tokens: int,
463 + ) -> BlockHash:
464 + assert num_tokens % self.hash_block_size == 0
465 + num_hash_blocks = num_tokens // self.hash_block_size
466 + assert 0 < num_hash_blocks <= len(request.block_hashes)
467 +
468 + # Each hash_block_size hash chains over its full prefix, so the partial
469 + # entry for any group block size is the hash at that prefix boundary.
470 + return request.block_hashes[num_hash_blocks - 1]
471 +
472 + def _get_partial_block_parent_hash_and_start(
473 + self,
474 + request: Request,
475 + num_tokens: int,
476 + ) -> tuple[BlockHash | None, int]:
477 + num_hash_blocks = num_tokens // self.hash_block_size
478 + parent_hash = (
479 + request.block_hashes[num_hash_blocks - 2] if num_hash_blocks > 1 else None
480 + )
481 + block_start = (num_hash_blocks - 1) * self.hash_block_size
482 + return parent_hash, block_start
483 +
484 + def _remove_cached_block_hashes(
485 + self,
486 + block: KVCacheBlock,
487 + ) -> list[BlockHashWithGroupId]:
488 + block_hashes: list[BlockHashWithGroupId] = []
489 + if block.block_hash is not None:
490 + block_hashes.append(block.block_hash)
491 + block_hashes.extend(self.cached_block_hashes_by_block.pop(block.block_id, ()))
492 + if not block_hashes:
493 + return []
494 +
495 + removed_hashes: list[BlockHashWithGroupId] = []
496 + for block_hash in block_hashes:
497 + if (
498 + self.cached_block_hash_to_block.pop(block_hash, block.block_id)
499 + is not None
500 + ):
501 + removed_hashes.append(block_hash)
502 + block.reset_hash()
503 + return removed_hashes
504 +
505 + def _emit_block_removed_events(
506 + self,
507 + block_hashes: list[BlockHashWithGroupId],
508 + ) -> None:
509 + if not self.enable_kv_cache_events:
510 + return
511 + for block_hash in block_hashes:
512 + self.kv_event_queue.append(
513 + BlockRemoved(
514 + block_hashes=[maybe_convert_block_hash(get_block_hash(block_hash))],
515 + medium=MEDIUM_GPU,
516 + group_idx=get_group_id(block_hash),
517 + )
518 + )
519 +
520 + def _insert_block_hash(
521 + self,
522 + block_hash_with_group_id: BlockHashWithGroupId,
523 + block: KVCacheBlock,
524 + num_tokens: int | None,
525 + ) -> None:
526 + if block.block_hash == block_hash_with_group_id:
527 + return
528 +
529 + if self.cached_block_hash_to_block.contain(
530 + block_hash_with_group_id, block.block_id
531 + ):
532 + return
533 +
534 + if block.block_hash is None:
535 + block.set_block_hash(block_hash_with_group_id, num_tokens=num_tokens)
536 + else:
537 + self.cached_block_hashes_by_block.setdefault(block.block_id, set()).add(
538 + block_hash_with_group_id
539 + )
540 + self.cached_block_hash_to_block.insert(block_hash_with_group_id, block)
541 +
542 def get_new_blocks(self, num_blocks: int) -> list[KVCacheBlock]:
543 """Get new blocks from the free block pool.
544
@@ -377,26 +586,12 @@ class BlockPool:
586 if self.metrics_collector:
587 self.metrics_collector.on_block_evicted(block)
588
380 - block_hash = block.block_hash
381 - if block_hash is None:
589 + evicted_hashes = self._remove_cached_block_hashes(block)
590 + if not evicted_hashes:
591 # The block doesn't have hash, eviction is not needed
592 return False
593
385 - if self.cached_block_hash_to_block.pop(block_hash, block.block_id) is None:
386 - # block not found in cached_block_hash_to_block,
387 - # eviction is not needed
388 - return False
389 -
390 - block.reset_hash()
391 -
392 - if self.enable_kv_cache_events:
393 - self.kv_event_queue.append(
394 - BlockRemoved(
395 - block_hashes=[maybe_convert_block_hash(get_block_hash(block_hash))],
396 - medium=MEDIUM_GPU,
397 - group_idx=get_group_id(block_hash),
398 - )
399 - )
594 + self._emit_block_removed_events(evicted_hashes)
595 return True
596
597 def touch(self, blocks: Sequence[KVCacheBlock]) -> None:
@@ -478,6 +673,7 @@ class BlockPool:
673
674 # Remove all hashes so that no new blocks will hit.
675 self.cached_block_hash_to_block = BlockHashToBlockMap()
676 + self.cached_block_hashes_by_block.clear()
677
678 # Remove all hashes from all blocks.
679 for block in self.blocks:
vllm/v1/core/kv_cache_utils.py
+38 -16
@@ -125,6 +125,9 @@ class KVCacheBlock:
125 # The hash key (block hash + group id) of the block, only available
126 # when the block is full and cached.
127 _block_hash: BlockHashWithGroupId | None = None
128 + # Number of prefix tokens covered by _block_hash. For full blocks this is
129 + # the full block boundary; partial aliases can end inside a cache block.
130 + _block_hash_num_tokens: int | None = None
131
132 # Used to construct a doubly linked list for free blocks.
133 # These two attributes should only be manipulated by FreeKVCacheBlockQueue.
@@ -138,16 +141,25 @@ class KVCacheBlock:
141 def block_hash(self) -> BlockHashWithGroupId | None:
142 return self._block_hash
143
141 - @block_hash.setter
142 - def block_hash(self, block_hash: BlockHashWithGroupId):
143 - assert self.block_hash is None, (
144 + @property
145 + def block_hash_num_tokens(self) -> int | None:
146 + return self._block_hash_num_tokens
147 +
148 + def set_block_hash(
149 + self,
150 + block_hash: BlockHashWithGroupId,
151 + num_tokens: int | None = None,
152 + ) -> None:
153 + assert self.block_hash is None and self._block_hash_num_tokens is None, (
154 "The block already has a hash. This should not happen."
155 )
156 self._block_hash = block_hash
157 + self._block_hash_num_tokens = num_tokens
158
159 def reset_hash(self):
160 """Reset the block hash when the block is evicted."""
161 self._block_hash = None
162 + self._block_hash_num_tokens = None
163
164 def __repr__(self) -> str:
165 # Use block_id instead of KVCacheBlock object to avoid calling __repr__
@@ -158,6 +170,7 @@ class KVCacheBlock:
170 f"KVCacheBlock(block_id={self.block_id}, "
171 f"ref_cnt={self.ref_cnt}, "
172 f"_block_hash={self._block_hash!r}, "
173 + f"_block_hash_num_tokens={self._block_hash_num_tokens}, "
174 f"prev_free_block={prev_block_id}, "
175 f"next_free_block={next_block_id})"
176 )
@@ -658,18 +671,24 @@ def resolve_kv_cache_block_sizes(
671
672
673 def get_request_block_hasher(
661 - block_size: int,
674 + hash_block_size: int,
675 caching_hash_fn: Callable[[Any], bytes],
676 ) -> Callable[[Request], list[BlockHash]]:
677 """
678 Returns a function which computes the list of un-computed block hashes
666 - of a request."""
679 + of a request.
680 +
681 + Hashes are computed at ``hash_block_size`` granularity and chained over the
682 + full prefix, so each hash uniquely fingerprints the prefix ending at its
683 + boundary. Coarser group block sizes and partial-cache boundaries reuse
684 + these hashes directly (see ``BlockHashListWithBlockSize``).
685 + """
686
687 def request_block_hasher(request: Request) -> list[BlockHash]:
669 - start_token_idx = len(request.block_hashes) * block_size
688 + start_token_idx = len(request.block_hashes) * hash_block_size
689 num_tokens = request.num_tokens
690
672 - if start_token_idx + block_size > num_tokens:
691 + if start_token_idx + hash_block_size > num_tokens:
692 # Early stop when there no new full blocks created.
693 return []
694
@@ -686,7 +705,7 @@ def get_request_block_hasher(
705 )
706 new_block_hashes: list[BlockHash] = []
707 while True:
689 - end_token_idx = start_token_idx + block_size
708 + end_token_idx = start_token_idx + hash_block_size
709 if end_token_idx > num_tokens:
710 # We only hash full blocks
711 break
@@ -703,7 +722,7 @@ def get_request_block_hasher(
722 )
723
724 new_block_hashes.append(block_hash)
706 - start_token_idx += block_size
725 + start_token_idx += hash_block_size
726 prev_block_hash_value = block_hash
727
728 return new_block_hashes
@@ -2132,11 +2151,14 @@ class BlockHashListWithBlockSize:
2151
2152 Currently, only scaling up by an integer factor is supported (i.e.,
2153 `target_block_size` is a multiple of `hash_block_size`). Conversion is
2135 - performed lazily on access for efficiency, by concatenating consecutive
2136 - hashes at `hash_block_size` to form each hash at `target_block_size`.
2154 + performed lazily on access for efficiency. Each `hash_block_size` hash is
2155 + already chained over its entire prefix, so the hash at the last
2156 + `hash_block_size` boundary of a `target_block_size` block uniquely
2157 + fingerprints that block's prefix; we use it directly.
2158
2159 Example (`hash_block_size` = 16, `target_block_size` = 32):
2139 - concatenating two 16-size hashes yields one 32-size hash:
2160 + the second 16-size hash already covers tokens 0-31, so it is the 32-size
2161 + hash:
2162
2163 Block hashes with block_size 16:
2164 | Token Range | 0-15 | 16-31 | 32-47 | 48-63 |
@@ -2146,7 +2168,7 @@ class BlockHashListWithBlockSize:
2168 Block hashes with block_size 32:
2169 | Token Range | 0-31 | 32-63 |
2170 |-------------|------|-------|
2149 - | Hash | AB | CD |
2171 + | Hash | B | D |
2172
2173 Args:
2174 block_hashes: Block hashes to convert, computed at `hash_block_size`.
@@ -2188,9 +2210,9 @@ class BlockHashListWithBlockSize:
2210 yield self._get_value_at(i)
2211
2212 def _get_value_at(self, idx: int) -> BlockHash:
2191 - base = idx * self.scale_factor
2192 - end = base + self.scale_factor
2193 - return BlockHash(b"".join(self.block_hashes[base:end]))
2213 + # The last hash_block_size hash within the target block already chains
2214 + # over the whole prefix, so it is the target block's hash.
2215 + return self.block_hashes[(idx + 1) * self.scale_factor - 1]
2216
2217
2218 BlockHashList = list[BlockHash] | BlockHashListWithBlockSize