@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 481e481be7

[2/N][Core] support partial prefix cache hit for hybrid model (#46384)

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

Jiangyun Zhu committed Jul 12, 2026 at 13:37 UTC 481e481be786c1ca3229e26aa34c15ffd22375af
21 files changed +1651 -175
tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py new
+816
@@ -0,0 +1,816 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Fine-grained partial prefix-cache hits for hybrid (full attention + mamba
4 +"align") models: scheduler chunk splitting, partial tail registration, CoW
5 +on partial hits, and same-step deferral."""
6 +
7 +from types import SimpleNamespace
8 +
9 +import pytest
10 +import torch
11 +
12 +from tests.v1.core.test_prefix_caching import make_kv_cache_manager, make_request
13 +from vllm.utils.hashing import sha256
14 +from vllm.v1.core.kv_cache_utils import (
15 + KVCacheBlockCopy,
16 + get_block_hash,
17 + get_group_id,
18 + init_none_hash,
19 +)
20 +from vllm.v1.core.sched.scheduler import Scheduler
21 +from vllm.v1.kv_cache_interface import (
22 + FullAttentionSpec,
23 + KVCacheConfig,
24 + KVCacheGroupSpec,
25 + MambaSpec,
26 +)
27 +
28 +
29 +@pytest.fixture(autouse=True)
30 +def _auto_init_hash_fn():
31 + init_none_hash(sha256)
32 +
33 +
34 +def test_mamba_align_split_partial_tail_schedule():
35 + """Chunk ends with partial hits on: block-aligned chunks, one extra stop
36 + at the prompt's last hash boundary (registering the partial tail), then
37 + the remaining tokens. block=512, hash=32, prompt=10000, budget=8192:
38 + 0 -> 8192 -> 9728 -> 9984 -> 10000."""
39 + block_size = 512
40 + hash_block_size = 32
41 + mock = SimpleNamespace(
42 + cache_config=SimpleNamespace(block_size=block_size),
43 + use_eagle=False,
44 + hash_block_size=hash_block_size,
45 + mamba_partial_cache_hit=True,
46 + )
47 + split = Scheduler._mamba_block_aligned_split
48 +
49 + req = make_request("0", [0] * 10000, hash_block_size, sha256)
50 + req.num_computed_tokens = 0
51 + assert split(self=mock, request=req, num_new_tokens=8192) == 8192
52 + req.num_computed_tokens = 8192
53 + # Stop at the last block boundary (9728).
54 + assert split(self=mock, request=req, num_new_tokens=1808) == 1536
55 + req.num_computed_tokens = 9728
56 + # Extra stop at the prompt's last hash boundary (9984).
57 + assert split(self=mock, request=req, num_new_tokens=272) == 256
58 + req.num_computed_tokens = 9984
59 + # Final 16 tokens run unchanged (no mid-block-resume stop: the next
60 + # block boundary is past the last block boundary).
61 + assert split(self=mock, request=req, num_new_tokens=16) == 16
62 +
63 + # Partial hits off: no extra stop, the tail runs in one chunk.
64 + mock.mamba_partial_cache_hit = False
65 + req.num_computed_tokens = 9728
66 + assert split(self=mock, request=req, num_new_tokens=272) == 272
67 + mock.mamba_partial_cache_hit = True
68 +
69 + # A request resumed mid-block (partial hash hit at 9984): the first chunk
70 + # stops at the next block boundary (10240), later chunk ends re-align.
71 + req2 = make_request("1", [0] * 12000, hash_block_size, sha256)
72 + req2.num_computed_tokens = 9984
73 + assert split(self=mock, request=req2, num_new_tokens=2016) == 256
74 + req2.num_computed_tokens = 10240
75 + assert split(self=mock, request=req2, num_new_tokens=1000) == 512
76 +
77 +
78 +def test_hybrid_mamba_align_partial_hash_hit():
79 + hash_block_size = 2
80 + mamba_block_size = 2 * hash_block_size
81 + kv_cache_config = KVCacheConfig(
82 + num_blocks=20,
83 + kv_cache_tensors=[],
84 + kv_cache_groups=[
85 + KVCacheGroupSpec(
86 + ["full"],
87 + FullAttentionSpec(
88 + block_size=hash_block_size,
89 + num_kv_heads=1,
90 + head_size=1,
91 + dtype=torch.float32,
92 + ),
93 + ),
94 + KVCacheGroupSpec(
95 + ["mamba"],
96 + MambaSpec(
97 + block_size=mamba_block_size,
98 + shapes=(1, 1),
99 + dtypes=(torch.float32,),
100 + mamba_cache_mode="align",
101 + ),
102 + ),
103 + ],
104 + )
105 + manager = make_kv_cache_manager(
106 + kv_cache_config=kv_cache_config,
107 + max_model_len=8192,
108 + enable_caching=True,
109 + hash_block_size=hash_block_size,
110 + )
111 +
112 + req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)
113 + computed_blocks, num_computed = manager.get_computed_blocks(req0)
114 + assert num_computed == 0
115 + blocks = manager.allocate_slots(req0, 6, num_computed, computed_blocks)
116 + assert blocks is not None
117 + manager.free(req0)
118 + manager.new_step_starts()
119 +
120 + partial_mamba_hash = req0.block_hashes[6 // hash_block_size - 1]
121 + partial_mamba_block = manager.block_pool.get_cached_block(
122 + partial_mamba_hash, kv_cache_group_ids=[1]
123 + )
124 + assert partial_mamba_block is not None
125 + assert partial_mamba_block[0].block_hash_num_tokens == 6
126 +
127 + req1 = make_request("1", [0, 0, 1, 1, 2, 2, 3, 3], hash_block_size, sha256)
128 + computed_blocks, num_computed = manager.get_computed_blocks(req1)
129 + assert num_computed == 6
130 + assert [len(group) for group in computed_blocks.blocks] == [3, 2]
131 +
132 + new_blocks = manager.allocate_slots(req1, 2, num_computed, computed_blocks)
133 + assert new_blocks is not None
134 + mamba_new_block_ids = new_blocks.get_block_ids()[1]
135 + assert len(mamba_new_block_ids) == 1
136 + assert mamba_new_block_ids[0] != partial_mamba_block[0].block_id
137 + assert manager.get_blocks("1").get_block_ids()[1][1] == mamba_new_block_ids[0]
138 + assert partial_mamba_block[0].block_hash is not None
139 + assert get_block_hash(partial_mamba_block[0].block_hash) == partial_mamba_hash
140 + assert get_group_id(partial_mamba_block[0].block_hash) == 1
141 + assert partial_mamba_block[0].block_hash_num_tokens == 6
142 + copies, _ = manager.take_kv_cache_block_copies()
143 + assert (
144 + KVCacheBlockCopy(
145 + src_block_id=partial_mamba_block[0].block_id,
146 + dst_block_id=mamba_new_block_ids[0],
147 + )
148 + in copies
149 + )
150 + assert manager.get_blocks("1").blocks[1][1].block_hash_num_tokens == 8
151 +
152 +
153 +def test_hybrid_mamba_partial_tail_owner_uses_cow_on_continue():
154 + hash_block_size = 2
155 + block_size = 2 * hash_block_size
156 + kv_cache_config = KVCacheConfig(
157 + num_blocks=24,
158 + kv_cache_tensors=[],
159 + kv_cache_groups=[
160 + KVCacheGroupSpec(
161 + ["full"],
162 + FullAttentionSpec(
163 + block_size=hash_block_size,
164 + num_kv_heads=1,
165 + head_size=1,
166 + dtype=torch.float32,
167 + ),
168 + ),
169 + KVCacheGroupSpec(
170 + ["mamba"],
171 + MambaSpec(
172 + block_size=block_size,
173 + shapes=(1, 1),
174 + dtypes=(torch.float32,),
175 + mamba_cache_mode="align",
176 + ),
177 + ),
178 + ],
179 + )
180 + manager = make_kv_cache_manager(
181 + kv_cache_config=kv_cache_config,
182 + max_model_len=8192,
183 + enable_caching=True,
184 + hash_block_size=hash_block_size,
185 + )
186 +
187 + req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)
188 + computed_blocks, num_computed = manager.get_computed_blocks(req0)
189 + assert num_computed == 0
190 + assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None
191 +
192 + partial_mamba_hash = req0.block_hashes[6 // hash_block_size - 1]
193 + partial_mamba_block = manager.block_pool.get_cached_block(
194 + partial_mamba_hash, kv_cache_group_ids=[1]
195 + )
196 + assert partial_mamba_block is not None
197 + partial_mamba_block_id = partial_mamba_block[0].block_id
198 + assert manager.get_blocks("0").get_block_ids()[1][1] == partial_mamba_block_id
199 +
200 + req0.num_computed_tokens = 6
201 + req0.append_output_token_ids([3])
202 + new_blocks = manager.allocate_slots(req0, 1)
203 + assert new_blocks is not None
204 +
205 + # Reversed CoW for the owning request: it keeps its own block (the
206 + # worker's block table is append-only), and no new mamba block is handed
207 + # to the worker. The prefix-cache entry is moved to a private copy that
208 + # the queued block copy fills before the next forward.
209 + assert new_blocks.get_block_ids()[1] == []
210 + assert manager.get_blocks("0").get_block_ids()[1][1] == partial_mamba_block_id
211 + copies, _ = manager.take_kv_cache_block_copies()
212 + cow_copy = next(c for c in copies if c.src_block_id == partial_mamba_block_id)
213 + assert cow_copy.dst_block_id != partial_mamba_block_id
214 + # The source block gave up the hash; the copy target now owns the entry.
215 + assert partial_mamba_block[0].block_hash is None
216 + moved = manager.block_pool.get_cached_block(
217 + partial_mamba_hash, kv_cache_group_ids=[1]
218 + )
219 + assert moved is not None
220 + assert moved[0].block_id == cow_copy.dst_block_id
221 + assert get_block_hash(moved[0].block_hash) == partial_mamba_hash
222 + assert get_group_id(moved[0].block_hash) == 1
223 + assert moved[0].block_hash_num_tokens == 6
224 +
225 +
226 +def test_hybrid_mamba_partial_tail_owner_continue_preserves_later_hit():
227 + hash_block_size = 2
228 + block_size = 2 * hash_block_size
229 + kv_cache_config = KVCacheConfig(
230 + num_blocks=32,
231 + kv_cache_tensors=[],
232 + kv_cache_groups=[
233 + KVCacheGroupSpec(
234 + ["full"],
235 + FullAttentionSpec(
236 + block_size=hash_block_size,
237 + num_kv_heads=1,
238 + head_size=1,
239 + dtype=torch.float32,
240 + ),
241 + ),
242 + KVCacheGroupSpec(
243 + ["mamba"],
244 + MambaSpec(
245 + block_size=block_size,
246 + shapes=(1, 1),
247 + dtypes=(torch.float32,),
248 + mamba_cache_mode="align",
249 + ),
250 + ),
251 + ],
252 + )
253 + manager = make_kv_cache_manager(
254 + kv_cache_config=kv_cache_config,
255 + max_model_len=8192,
256 + enable_caching=True,
257 + hash_block_size=hash_block_size,
258 + )
259 +
260 + req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)
261 + computed_blocks, num_computed = manager.get_computed_blocks(req0)
262 + assert num_computed == 0
263 + assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None
264 +
265 + partial_mamba_hash = req0.block_hashes[6 // hash_block_size - 1]
266 + partial_mamba_block = manager.block_pool.get_cached_block(
267 + partial_mamba_hash, kv_cache_group_ids=[1]
268 + )
269 + assert partial_mamba_block is not None
270 + partial_mamba_block_id = partial_mamba_block[0].block_id
271 +
272 + req0.num_computed_tokens = 6
273 + req0.append_output_token_ids([3])
274 + assert manager.allocate_slots(req0, 1) is not None
275 + # The owner moved the prefix-cache entry to a private copy; capture its id.
276 + owner_copies, _ = manager.take_kv_cache_block_copies()
277 + cow_copy = next(c for c in owner_copies if c.src_block_id == partial_mamba_block_id)
278 + moved_block_id = cow_copy.dst_block_id
279 + manager.new_step_starts()
280 +
281 + req1 = make_request("1", [0, 0, 1, 1, 2, 2, 4, 4], hash_block_size, sha256)
282 + computed_blocks, num_computed = manager.get_computed_blocks(req1)
283 + assert num_computed == 6
284 + # The later request hits the moved (private-copy) entry, not the source.
285 + assert computed_blocks.get_block_ids()[1][1] == moved_block_id
286 +
287 + new_blocks = manager.allocate_slots(req1, 2, num_computed, computed_blocks)
288 + assert new_blocks is not None
289 + mamba_new_block_ids = new_blocks.get_block_ids()[1]
290 + assert len(mamba_new_block_ids) == 1
291 + assert mamba_new_block_ids[0] != moved_block_id
292 + # The hitting request CoWs from the moved entry into its own private block.
293 + copies, _ = manager.take_kv_cache_block_copies()
294 + assert (
295 + KVCacheBlockCopy(
296 + src_block_id=moved_block_id,
297 + dst_block_id=mamba_new_block_ids[0],
298 + )
299 + in copies
300 + )
301 +
302 +
303 +def test_hybrid_mamba_moved_partial_entry_defers_same_step_hit():
304 + """The owner's move re-arms the same-step guard: the moved entry is
305 + filled by this step's copy, and chained same-step copies read stale
306 + sources, so a request hitting it in the move step must be deferred."""
307 + hash_block_size = 2
308 + block_size = 2 * hash_block_size
309 + kv_cache_config = KVCacheConfig(
310 + num_blocks=32,
311 + kv_cache_tensors=[],
312 + kv_cache_groups=[
313 + KVCacheGroupSpec(
314 + ["full"],
315 + FullAttentionSpec(
316 + block_size=hash_block_size,
317 + num_kv_heads=1,
318 + head_size=1,
319 + dtype=torch.float32,
320 + ),
321 + ),
322 + KVCacheGroupSpec(
323 + ["mamba"],
324 + MambaSpec(
325 + block_size=block_size,
326 + shapes=(1, 1),
327 + dtypes=(torch.float32,),
328 + mamba_cache_mode="align",
329 + ),
330 + ),
331 + ],
332 + )
333 + manager = make_kv_cache_manager(
334 + kv_cache_config=kv_cache_config,
335 + max_model_len=8192,
336 + enable_caching=True,
337 + hash_block_size=hash_block_size,
338 + )
339 +
340 + req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)
341 + computed_blocks, num_computed = manager.get_computed_blocks(req0)
342 + assert num_computed == 0
343 + assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None
344 + manager.new_step_starts()
345 +
346 + # The owning request continues decoding: the partial entry moves to a
347 + # private copy in this step.
348 + req0.num_computed_tokens = 6
349 + req0.append_output_token_ids([3])
350 + assert manager.allocate_slots(req0, 1) is not None
351 +
352 + # A request hitting the moved entry in the SAME step must be deferred.
353 + req1 = make_request("1", [0, 0, 1, 1, 2, 2, 4, 4], hash_block_size, sha256)
354 + computed_blocks, num_computed = manager.get_computed_blocks(req1)
355 + assert num_computed == 6
356 + assert manager.allocate_slots(req1, 2, num_computed, computed_blocks) is None
357 +
358 + # Next step the moved entry is consumable.
359 + manager.new_step_starts()
360 + computed_blocks, num_computed = manager.get_computed_blocks(req1)
361 + assert num_computed == 6
362 + assert manager.allocate_slots(req1, 2, num_computed, computed_blocks) is not None
363 +
364 +
365 +def test_hybrid_full_attention_partial_hash_hit_uses_cow():
366 + hash_block_size = 2
367 + block_size = 2 * hash_block_size
368 + kv_cache_config = KVCacheConfig(
369 + num_blocks=24,
370 + kv_cache_tensors=[],
371 + kv_cache_groups=[
372 + KVCacheGroupSpec(
373 + ["full"],
374 + FullAttentionSpec(
375 + block_size=block_size,
376 + num_kv_heads=1,
377 + head_size=1,
378 + dtype=torch.float32,
379 + ),
380 + ),
381 + KVCacheGroupSpec(
382 + ["mamba"],
383 + MambaSpec(
384 + block_size=block_size,
385 + shapes=(1, 1),
386 + dtypes=(torch.float32,),
387 + mamba_cache_mode="align",
388 + ),
389 + ),
390 + ],
391 + )
392 + manager = make_kv_cache_manager(
393 + kv_cache_config=kv_cache_config,
394 + max_model_len=8192,
395 + enable_caching=True,
396 + hash_block_size=hash_block_size,
397 + )
398 +
399 + req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)
400 + computed_blocks, num_computed = manager.get_computed_blocks(req0)
401 + assert num_computed == 0
402 + assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None
403 + manager.free(req0)
404 + manager.new_step_starts()
405 +
406 + partial_full_hash = req0.block_hashes[6 // hash_block_size - 1]
407 + partial_full_block = manager.block_pool.get_cached_block(
408 + partial_full_hash, kv_cache_group_ids=[0]
409 + )
410 + assert partial_full_block is not None
411 +
412 + req1 = make_request("1", [0, 0, 1, 1, 2, 2, 3, 3], hash_block_size, sha256)
413 + computed_blocks, num_computed = manager.get_computed_blocks(req1)
414 + assert num_computed == 6
415 + assert [len(group) for group in computed_blocks.blocks] == [2, 2]
416 +
417 + new_blocks = manager.allocate_slots(req1, 2, num_computed, computed_blocks)
418 + assert new_blocks is not None
419 + full_new_block_ids = new_blocks.get_block_ids()[0]
420 + assert len(full_new_block_ids) == 1
421 + assert full_new_block_ids[0] != partial_full_block[0].block_id
422 + assert partial_full_block[0].block_hash is not None
423 + assert get_block_hash(partial_full_block[0].block_hash) == partial_full_hash
424 + assert get_group_id(partial_full_block[0].block_hash) == 0
425 + assert partial_full_block[0].block_hash_num_tokens == 6
426 + copies, retained = manager.take_kv_cache_block_copies()
427 + assert (
428 + KVCacheBlockCopy(
429 + src_block_id=partial_full_block[0].block_id,
430 + dst_block_id=full_new_block_ids[0],
431 + )
432 + in copies
433 + )
434 + assert partial_full_block[0].ref_cnt == 1
435 + manager.block_pool.free_blocks(retained)
436 + assert partial_full_block[0].ref_cnt == 0
437 +
438 +
439 +def test_hybrid_partial_hit_cow_target_starts_uncached():
440 + hash_block_size = 2
441 + block_size = 2 * hash_block_size
442 + kv_cache_config = KVCacheConfig(
443 + num_blocks=32,
444 + kv_cache_tensors=[],
445 + kv_cache_groups=[
446 + KVCacheGroupSpec(
447 + ["full"],
448 + FullAttentionSpec(
449 + block_size=block_size,
450 + num_kv_heads=1,
451 + head_size=1,
452 + dtype=torch.float32,
453 + ),
454 + ),
455 + KVCacheGroupSpec(
456 + ["mamba"],
457 + MambaSpec(
458 + block_size=block_size,
459 + shapes=(1, 1),
460 + dtypes=(torch.float32,),
461 + mamba_cache_mode="align",
462 + ),
463 + ),
464 + ],
465 + )
466 + manager = make_kv_cache_manager(
467 + kv_cache_config=kv_cache_config,
468 + max_model_len=8192,
469 + enable_caching=True,
470 + hash_block_size=hash_block_size,
471 + )
472 +
473 + req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)
474 + computed_blocks, num_computed = manager.get_computed_blocks(req0)
475 + assert num_computed == 0
476 + assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None
477 + manager.free(req0)
478 + manager.new_step_starts()
479 +
480 + partial_hash = req0.block_hashes[6 // hash_block_size - 1]
481 + partial_full_block = manager.block_pool.get_cached_block(
482 + partial_hash, kv_cache_group_ids=[0]
483 + )
484 + partial_mamba_block = manager.block_pool.get_cached_block(
485 + partial_hash, kv_cache_group_ids=[1]
486 + )
487 + assert partial_full_block is not None
488 + assert partial_mamba_block is not None
489 +
490 + req1 = make_request("1", [0, 0, 1, 1, 2, 2, 3, 3], hash_block_size, sha256)
491 + computed_blocks, num_computed = manager.get_computed_blocks(req1)
492 + assert num_computed == 6
493 +
494 + new_blocks = manager.allocate_slots(
495 + req1,
496 + 2,
497 + num_computed,
498 + computed_blocks,
499 + delay_cache_blocks=True,
500 + )
501 + assert new_blocks is not None
502 +
503 + full_cow_block = manager.get_blocks("1").blocks[0][1]
504 + mamba_cow_block = manager.get_blocks("1").blocks[1][1]
505 + assert full_cow_block.block_id != partial_full_block[0].block_id
506 + assert mamba_cow_block.block_id != partial_mamba_block[0].block_id
507 + assert full_cow_block.block_hash is None
508 + assert full_cow_block.block_hash_num_tokens is None
509 + assert mamba_cow_block.block_hash is None
510 + assert mamba_cow_block.block_hash_num_tokens is None
511 +
512 + assert partial_full_block[0].block_hash is not None
513 + assert get_block_hash(partial_full_block[0].block_hash) == partial_hash
514 + assert get_group_id(partial_full_block[0].block_hash) == 0
515 + assert partial_full_block[0].block_hash_num_tokens == 6
516 + assert partial_mamba_block[0].block_hash is not None
517 + assert get_block_hash(partial_mamba_block[0].block_hash) == partial_hash
518 + assert get_group_id(partial_mamba_block[0].block_hash) == 1
519 + assert partial_mamba_block[0].block_hash_num_tokens == 6
520 +
521 +
522 +def test_hybrid_partial_hash_truncates_full_attention_hit_length():
523 + hash_block_size = 2
524 + block_size = 2 * hash_block_size
525 + kv_cache_config = KVCacheConfig(
526 + num_blocks=24,
527 + kv_cache_tensors=[],
528 + kv_cache_groups=[
529 + KVCacheGroupSpec(
530 + ["full"],
531 + FullAttentionSpec(
532 + block_size=block_size,
533 + num_kv_heads=1,
534 + head_size=1,
535 + dtype=torch.float32,
536 + ),
537 + ),
538 + KVCacheGroupSpec(
539 + ["mamba"],
540 + MambaSpec(
541 + block_size=block_size,
542 + shapes=(1, 1),
543 + dtypes=(torch.float32,),
544 + mamba_cache_mode="align",
545 + ),
546 + ),
547 + ],
548 + )
549 + manager = make_kv_cache_manager(
550 + kv_cache_config=kv_cache_config,
551 + max_model_len=8192,
552 + enable_caching=True,
553 + hash_block_size=hash_block_size,
554 + )
555 + pool = manager.block_pool
556 + req = make_request(
557 + "0",
558 + [0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5],
559 + hash_block_size,
560 + sha256,
561 + )
562 +
563 + full_blocks = pool.get_new_blocks(3)
564 + pool.cache_full_blocks(
565 + request=req,
566 + blocks=full_blocks,
567 + num_cached_blocks=0,
568 + num_full_blocks=2,
569 + block_size=block_size,
570 + kv_cache_group_id=0,
571 + )
572 + pool.cache_partial_block(
573 + request=req,
574 + block=full_blocks[2],
575 + num_tokens=10,
576 + kv_cache_group_id=0,
577 + block_size=block_size,
578 + )
579 +
580 + mamba_block = pool.get_new_blocks(1)[0]
581 + pool.cache_partial_block(
582 + request=req,
583 + block=mamba_block,
584 + num_tokens=6,
585 + kv_cache_group_id=1,
586 + block_size=block_size,
587 + )
588 +
589 + computed_blocks, num_computed = manager.get_computed_blocks(req)
590 + assert num_computed == 6
591 + assert [len(group) for group in computed_blocks.blocks] == [2, 2]
592 +
593 +
594 +def test_cow_retained_blocks_returned_for_release():
595 + """new_step_starts returns the CoW copy retentions instead of freeing
596 + them; the scheduler owns releasing them once the copy has run."""
597 + hash_block_size = 2
598 + block_size = 2 * hash_block_size
599 + kv_cache_config = KVCacheConfig(
600 + num_blocks=24,
601 + kv_cache_tensors=[],
602 + kv_cache_groups=[
603 + KVCacheGroupSpec(
604 + ["full"],
605 + FullAttentionSpec(
606 + block_size=hash_block_size,
607 + num_kv_heads=1,
608 + head_size=1,
609 + dtype=torch.float32,
610 + ),
611 + ),
612 + KVCacheGroupSpec(
613 + ["mamba"],
614 + MambaSpec(
615 + block_size=block_size,
616 + shapes=(1, 1),
617 + dtypes=(torch.float32,),
618 + mamba_cache_mode="align",
619 + ),
620 + ),
621 + ],
622 + )
623 + manager = make_kv_cache_manager(
624 + kv_cache_config=kv_cache_config,
625 + max_model_len=8192,
626 + enable_caching=True,
627 + hash_block_size=hash_block_size,
628 + )
629 + req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)
630 + computed_blocks, num_computed = manager.get_computed_blocks(req0)
631 + assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None
632 +
633 + # The owner's move queues a copy and retains both endpoints.
634 + req0.num_computed_tokens = 6
635 + req0.append_output_token_ids([3])
636 + assert manager.allocate_slots(req0, 1) is not None
637 + (cow_copy,), retained = manager.take_kv_cache_block_copies()
638 + assert {b.block_id for b in retained} == {
639 + cow_copy.src_block_id,
640 + cow_copy.dst_block_id,
641 + }
642 + # Not freed yet: the retention refs are still held.
643 + assert all(b.ref_cnt > 0 for b in retained)
644 + manager.block_pool.free_blocks(retained)
645 +
646 +
647 +def test_free_cow_retained_blocks_defers_until_copy_step_processed():
648 + """Scheduler releases CoW retentions immediately when the copy's step has
649 + been processed (or deferral is off), and defers them otherwise."""
650 + from collections import deque
651 +
652 + freed: list = []
653 + blocks = [SimpleNamespace(block_id=7), SimpleNamespace(block_id=9)]
654 + mock = SimpleNamespace(
655 + kv_cache_manager=SimpleNamespace(
656 + block_pool=SimpleNamespace(free_blocks=freed.extend)
657 + ),
658 + deferred_frees=deque(),
659 + defer_block_free=True,
660 + processed_step_seq=2,
661 + )
662 + free = Scheduler._free_cow_retained_blocks
663 +
664 + # Copy step still in flight: deferred with its fence.
665 + free(mock, list(blocks), fence_seq=3)
666 + assert not freed
667 + assert mock.deferred_frees == deque([(3, blocks[::-1])])
668 +
669 + # Copy step processed: freed immediately.
670 + mock.processed_step_seq = 3
671 + free(mock, list(blocks), fence_seq=3)
672 + assert freed == blocks
673 +
674 + # Deferral disabled: freed immediately regardless of the fence.
675 + freed.clear()
676 + mock.deferred_frees.clear()
677 + mock.defer_block_free = False
678 + mock.processed_step_seq = 0
679 + free(mock, list(blocks), fence_seq=3)
680 + assert freed == blocks
681 +
682 +
683 +def test_full_attention_eagle_drops_one_hash_unit():
684 + """With fine-grained partial hits, eagle rewinds the hit by one hash unit
685 + instead of a whole cache block: the tail block's KV is append-only, so it
686 + still covers the reduced length and stays in the hit as a partial block."""
687 + from vllm.v1.core.block_pool import BlockPool
688 + from vllm.v1.core.single_type_kv_cache_manager import FullAttentionManager
689 +
690 + hash_block_size = 2
691 + block_size = 4
692 + pool = BlockPool(
693 + num_gpu_blocks=10, enable_caching=True, hash_block_size=hash_block_size
694 + )
695 + spec = FullAttentionSpec(
696 + block_size=block_size, num_kv_heads=1, head_size=1, dtype=torch.float32
697 + )
698 + req = make_request("0", [0, 0, 1, 1, 2, 2, 3, 3], hash_block_size, sha256)
699 +
700 + def find(drop_eagle_block):
701 + return FullAttentionManager.find_longest_cache_hit(
702 + block_hashes=req.block_hashes,
703 + max_length=8,
704 + kv_cache_group_ids=[0],
705 + block_pool=pool,
706 + kv_cache_spec=spec,
707 + drop_eagle_block=drop_eagle_block,
708 + alignment_tokens=hash_block_size,
709 + )
710 +
711 + # Two full cached blocks (hit 8): eagle rewinds to 6, keeping the last
712 + # block as a partial hit instead of dropping it to 4.
713 + blocks = pool.get_new_blocks(2)
714 + pool.cache_full_blocks(
715 + request=req,
716 + blocks=blocks,
717 + num_cached_blocks=0,
718 + num_full_blocks=2,
719 + block_size=block_size,
720 + kv_cache_group_id=0,
721 + )
722 + hit_blocks, hit_length = find(drop_eagle_block=False)
723 + assert (hit_length, len(hit_blocks[0])) == (8, 2)
724 + hit_blocks, hit_length = find(drop_eagle_block=True)
725 + assert (hit_length, len(hit_blocks[0])) == (6, 2)
726 +
727 + # A partial tail at 6 (block 1 not fully cached): eagle rewinds to the
728 + # block boundary and trims the tail block.
729 + pool2 = BlockPool(
730 + num_gpu_blocks=10, enable_caching=True, hash_block_size=hash_block_size
731 + )
732 + pool = pool2
733 + blocks = pool.get_new_blocks(2)
734 + pool.cache_full_blocks(
735 + request=req,
736 + blocks=blocks[:1],
737 + num_cached_blocks=0,
738 + num_full_blocks=1,
739 + block_size=block_size,
740 + kv_cache_group_id=0,
741 + )
742 + assert (
743 + pool.cache_partial_block(
744 + request=req,
745 + block=blocks[1],
746 + num_tokens=6,
747 + kv_cache_group_id=0,
748 + block_size=block_size,
749 + )
750 + is not None
751 + )
752 + hit_blocks, hit_length = find(drop_eagle_block=False)
753 + assert (hit_length, len(hit_blocks[0])) == (6, 2)
754 + hit_blocks, hit_length = find(drop_eagle_block=True)
755 + assert (hit_length, len(hit_blocks[0])) == (4, 1)
756 +
757 +
758 +def test_hybrid_partial_hit_with_eagle_stays_within_group_blocks():
759 + """Regression: with eagle, the mamba group must not receive the eagle
760 + lookup margin — its finder never applies the drop, so it could return a
761 + hit past the blocks the (dropped) full-attention group covers, crashing
762 + the consumer's CoW with block_idx >= len(req_blocks)."""
763 + hash_block_size = 2
764 + block_size = 2 * hash_block_size
765 + kv_cache_config = KVCacheConfig(
766 + num_blocks=32,
767 + kv_cache_tensors=[],
768 + kv_cache_groups=[
769 + KVCacheGroupSpec(
770 + ["full"],
771 + FullAttentionSpec(
772 + block_size=block_size,
773 + num_kv_heads=1,
774 + head_size=1,
775 + dtype=torch.float32,
776 + ),
777 + ),
778 + KVCacheGroupSpec(
779 + ["mamba"],
780 + MambaSpec(
781 + block_size=block_size,
782 + shapes=(1, 1),
783 + dtypes=(torch.float32,),
784 + mamba_cache_mode="align",
785 + ),
786 + ),
787 + ],
788 + )
789 + manager = make_kv_cache_manager(
790 + kv_cache_config=kv_cache_config,
791 + max_model_len=8192,
792 + enable_caching=True,
793 + hash_block_size=hash_block_size,
794 + use_eagle=True,
795 + )
796 +
797 + # The owner prefills in scheduler-split style: stop at the block boundary
798 + # (4), then at the prompt's last hash boundary (6, partial entries).
799 + req0 = make_request("0", [7] * 6, hash_block_size, sha256)
800 + computed_blocks, num_computed = manager.get_computed_blocks(req0)
801 + assert manager.allocate_slots(req0, 4, num_computed, computed_blocks) is not None
802 + req0.num_computed_tokens = 4
803 + manager.new_step_starts()
804 + assert manager.allocate_slots(req0, 2) is not None
805 + req0.num_computed_tokens = 6
806 + manager.new_step_starts()
807 +
808 + # A longer request with eagle: full attention drops the partial tail, so
809 + # the joint hit must fall back to the block boundary the FA blocks cover.
810 + req1 = make_request("1", [7] * 6 + [9] * 2, hash_block_size, sha256)
811 + computed_blocks, num_computed = manager.get_computed_blocks(req1)
812 + assert num_computed == 4
813 + assert all(
814 + len(group) * block_size >= num_computed for group in computed_blocks.blocks
815 + )
816 + assert manager.allocate_slots(req1, 4, num_computed, computed_blocks) is not None
tests/v1/core/test_deferred_block_free.py
+62
@@ -18,6 +18,7 @@ from unittest.mock import PropertyMock, patch
18 import pytest
19
20 from vllm.config import VllmConfig
21 +from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy
22 from vllm.v1.core.sched.output import SchedulerOutput
23 from vllm.v1.outputs import ModelRunnerOutput
24 from vllm.v1.request import RequestStatus
@@ -412,3 +413,64 @@ def test_non_async_abort_defers_via_last_sched_seq():
413 scheduler.update_from_output(out0, _make_model_runner_output(out0))
414 assert not scheduler.deferred_frees
415 assert pool.get_num_free_blocks() == num_free_initially
416 +
417 +
418 +def test_cow_retentions_deferred_until_copy_step_processed():
419 + """The endpoints of a queued KV block copy must stay out of the free
420 + pool until the step that runs the copy has been processed. Freed
421 + earlier, an endpoint can be reallocated (e.g. as a PD KV-load
422 + destination) and overwritten by a transfer that is not ordered against
423 + the copy still pending in the in-flight step.
424 + """
425 + scheduler = _create_deferring_scheduler()
426 + pool = scheduler.kv_cache_manager.block_pool
427 + manager = scheduler.kv_cache_manager.coordinator.single_type_managers[0]
428 +
429 + request = create_requests(
430 + num_requests=1,
431 + num_tokens=NUM_PROMPT_TOKENS,
432 + max_tokens=5,
433 + stop_token_ids=[STOP_TOKEN_ID],
434 + )[0]
435 + scheduler.add_request(request)
436 +
437 + # Simulate a partial-hit CoW performed while scheduling step 1, whose
438 + # hitting request was freed within the same step: the copy rides out0
439 + # and each endpoint stays alive only through its copy retention.
440 + src_block, dst_block = pool.get_new_blocks(2)
441 + block_copy = KVCacheBlockCopy(
442 + src_block_id=src_block.block_id, dst_block_id=dst_block.block_id
443 + )
444 + manager._pending_cow_copies.append((src_block, dst_block))
445 + out0 = scheduler.schedule()
446 + assert out0.kv_cache_block_copies == [block_copy]
447 +
448 + # Exhaust the rest of the pool so the copy endpoints are the only blocks
449 + # a new request could receive, then add one that fits exactly in them.
450 + pool.get_new_blocks(pool.get_num_free_blocks())
451 + late_request = create_requests(
452 + num_requests=1,
453 + num_tokens=2 * scheduler.block_size,
454 + max_tokens=5,
455 + req_ids=["late"],
456 + )[0]
457 + scheduler.add_request(late_request)
458 +
459 + # Step 2 is scheduled while step 1 (which runs the copy) is still in
460 + # flight: the retentions are released against the copy's fence, so the
461 + # endpoints must not reach the free pool -- the late request must not be
462 + # scheduled onto them.
463 + out1 = scheduler.schedule()
464 + assert src_block.ref_cnt == 1
465 + assert dst_block.ref_cnt == 1
466 + assert scheduler.deferred_frees
467 + assert not out1.scheduled_new_reqs
468 +
469 + # Step 1's output is processed: the copy has run, endpoints return to
470 + # the pool and the late request can be scheduled onto them safely.
471 + scheduler.update_from_output(out0, _make_model_runner_output(out0))
472 + assert src_block.ref_cnt == 0
473 + assert dst_block.ref_cnt == 0
474 + assert not scheduler.deferred_frees
475 + out2 = scheduler.schedule()
476 + assert [r.req_id for r in out2.scheduled_new_reqs] == ["late"]
tests/v1/core/test_kv_cache_utils.py
+60
@@ -2597,3 +2597,63 @@ def test_hma_not_disabled_when_kv_events_enabled():
2597 assert vllm_config.scheduler_config.disable_hybrid_kv_cache_manager is False, (
2598 "kv_events_config must not force-disable the hybrid KV cache manager."
2599 )
2600 +
2601 +
2602 +def test_resolve_block_hashes_gate():
2603 + # Resolve symbols through the module so they stay consistent with each other
2604 + # even after other tests reload ``kv_cache_utils``.
2605 + resolve_block_hashes = kv_cache_utils.resolve_block_hashes
2606 + BlockHashListWithBlockSize = kv_cache_utils.BlockHashListWithBlockSize
2607 + # Raw, hash_block_size-granularity hashes (contents are opaque here).
2608 + raw = [BlockHash(bytes([i])) for i in range(8)]
2609 +
2610 + # block_size == hash_block_size: always reuse the raw hashes.
2611 + assert resolve_block_hashes(raw, 2, 2, alignment_tokens=2) is raw
2612 + assert (
2613 + resolve_block_hashes(
2614 + raw, 2, 2, supports_fine_grained_hash_lookup=True, alignment_tokens=2
2615 + )
2616 + is raw
2617 + )
2618 +
2619 + # Fine-grained manager, partial hits ON (alignment_tokens == hash_block_size
2620 + # < block_size): keep raw hashes so the manager can scan at hash granularity.
2621 + assert (
2622 + resolve_block_hashes(
2623 + raw, 2, 4, supports_fine_grained_hash_lookup=True, alignment_tokens=2
2624 + )
2625 + is raw
2626 + )
2627 +
2628 + # Fine-grained manager, partial hits OFF (alignment_tokens ==
2629 + # scheduler_block_size >= block_size): must fall back to a block-size view,
2630 + # exactly like the pre-refactor coordinator did.
2631 + off = resolve_block_hashes(
2632 + raw, 2, 4, supports_fine_grained_hash_lookup=True, alignment_tokens=4
2633 + )
2634 + assert isinstance(off, BlockHashListWithBlockSize)
2635 + assert off.scale_factor == 2
2636 +
2637 + # Non-fine-grained manager (e.g. sliding window): always a block-size view
2638 + # when block_size != hash_block_size, regardless of alignment_tokens.
2639 + swa = resolve_block_hashes(
2640 + raw, 2, 4, supports_fine_grained_hash_lookup=False, alignment_tokens=2
2641 + )
2642 + assert isinstance(swa, BlockHashListWithBlockSize)
2643 + assert swa.scale_factor == 2
2644 +
2645 +
2646 +def test_resolve_block_hashes_rejects_mismatched_view():
2647 + resolve_block_hashes = kv_cache_utils.resolve_block_hashes
2648 + BlockHashListWithBlockSize = kv_cache_utils.BlockHashListWithBlockSize
2649 + raw = [BlockHash(bytes([i])) for i in range(8)]
2650 +
2651 + # A view built at this block_size is returned as-is (idempotent).
2652 + view = BlockHashListWithBlockSize(raw, 2, 4)
2653 + assert resolve_block_hashes(view, 2, 4) is view
2654 +
2655 + # A view built at a different block_size must fail loudly rather than be
2656 + # silently reinterpreted at the wrong granularity.
2657 + mismatched = BlockHashListWithBlockSize(raw, 2, 8)
2658 + with pytest.raises(AssertionError):
2659 + resolve_block_hashes(mismatched, 2, 4)
tests/v1/core/test_prefix_caching.py
+4 -1
@@ -1072,7 +1072,10 @@ def test_hybrid_cache_mamba_align_shared_prefix_detection():
1072 # Next, validate scheduler logic for num_uncached_common_prefix_tokens > 0
1073 # Create minimal mock with just the needed attributes
1074 mock = SimpleNamespace(
1075 - cache_config=SimpleNamespace(block_size=block_size), use_eagle=False
1075 + cache_config=SimpleNamespace(block_size=block_size),
1076 + use_eagle=False,
1077 + hash_block_size=block_size,
1078 + mamba_partial_cache_hit=False,
1079 )
1080 num_new_tokens_adjusted = Scheduler._mamba_block_aligned_split(
1081 self=mock,
tests/v1/core/test_single_type_kv_cache_manager.py
+8 -6
@@ -93,7 +93,7 @@ def test_chunked_local_attention_possible_cached_prefix():
93 kv_cache_spec=chunked_local_attention_spec,
94 drop_eagle_block=False,
95 alignment_tokens=block_size,
96 - )[0]
96 + )[0][0]
97 assert len(computed_blocks) == expect_length
98
99 assert all(
@@ -164,7 +164,7 @@ def test_sliding_window_possible_cached_prefix():
164 kv_cache_spec=sliding_window_spec,
165 drop_eagle_block=False,
166 alignment_tokens=block_size,
167 - )[0]
167 + )[0][0]
168 assert len(computed_blocks) == expect_length
169
170 assert all(
@@ -398,13 +398,13 @@ def test_get_num_blocks_to_allocate():
398
399 assert (
400 manager.get_num_blocks_to_allocate(
401 - "1", 20 * block_size, cached_blocks_1, 0, 20 * block_size
401 + "1", 20 * block_size, cached_blocks_1, 0, 0, 20 * block_size
402 )
403 == 20
404 )
405 assert (
406 manager.get_num_blocks_to_allocate(
407 - "2", 20 * block_size, cached_blocks_2, 0, 20 * block_size
407 + "2", 20 * block_size, cached_blocks_2, 0, 0, 20 * block_size
408 )
409 == 15
410 )
@@ -434,6 +434,7 @@ def test_evictable_cached_blocks_not_double_allocated():
434 num_tokens=2 * block_size,
435 new_computed_blocks=[evictable_block],
436 total_computed_tokens=block_size,
437 + num_local_computed_tokens=block_size,
438 num_tokens_main_model=2 * block_size,
439 )
440 # Free capacity check should count evictable cached blocks, but allocation
@@ -474,13 +475,13 @@ def test_chunked_local_attention_get_num_blocks_to_allocate():
475
476 assert (
477 manager.get_num_blocks_to_allocate(
477 - "1", 20 * block_size, cached_blocks_1, 0, 20 * block_size
478 + "1", 20 * block_size, cached_blocks_1, 0, 0, 20 * block_size
479 )
480 == 20
481 )
482 assert (
483 manager.get_num_blocks_to_allocate(
483 - "2", 20 * block_size, cached_blocks_2, 0, 20 * block_size
484 + "2", 20 * block_size, cached_blocks_2, 0, 0, 20 * block_size
485 )
486 == 15
487 )
@@ -524,6 +525,7 @@ def test_predictor_matches_allocator_blocks_calculation_with_admission_cap():
525 num_tokens=num_tokens,
526 new_computed_blocks=[],
527 total_computed_tokens=total_computed,
528 + num_local_computed_tokens=0,
529 num_tokens_main_model=num_tokens,
530 )
531 new_blocks = manager.allocate_new_blocks(
tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py
+15 -15
@@ -37,7 +37,7 @@ def _make_coord(groups, hash_block_size, use_eagle=False, retention_interval=Non
37
38
39 def test_external_cached_block_pool_tautological_returns_present_for_any_hash():
40 - cmap = ExternalCachedBlockPool()
40 + cmap = ExternalCachedBlockPool(16)
41 h = BlockHash(b"\xaa" * 4)
42 res = cmap.get_cached_block(h, [0, 1])
43 assert res is not None
@@ -48,7 +48,7 @@ def test_external_cached_block_pool_tautological_returns_present_for_any_hash():
48
49 def test_external_cached_block_pool_hit_all_groups():
50 h = BlockHash(b"\x11\x22\x33\x44")
51 - cmap = ExternalCachedBlockPool({(0, bytes(h)), (1, bytes(h))})
51 + cmap = ExternalCachedBlockPool(16, {(0, bytes(h)), (1, bytes(h))})
52 res = cmap.get_cached_block(h, [0, 1])
53 assert res is not None
54 assert len(res) == 2
@@ -58,14 +58,14 @@ def test_external_cached_block_pool_hit_all_groups():
58
59 def test_external_cached_block_pool_miss_one_group():
60 h = BlockHash(b"\x11\x22\x33\x44")
61 - cmap = ExternalCachedBlockPool({(0, bytes(h))})
61 + cmap = ExternalCachedBlockPool(16, {(0, bytes(h))})
62 assert cmap.get_cached_block(h, [0, 1]) is None
63
64
65 def test_external_cached_block_pool_unknown_hash():
66 h_known = BlockHash(b"\x01" * 4)
67 h_unknown = BlockHash(b"\x02" * 4)
68 - cmap = ExternalCachedBlockPool({(0, bytes(h_known))})
68 + cmap = ExternalCachedBlockPool(16, {(0, bytes(h_known))})
69 assert cmap.get_cached_block(h_unknown, [0]) is None
70
71
@@ -103,7 +103,7 @@ def test_coordinator_single_full_attention_all_hits():
103 groups = [KVCacheGroupSpec(["L0"], _full(16))]
104 coord = _make_coord(groups, hash_block_size=16)
105 hs = _hashes(4)
106 - cmap = ExternalCachedBlockPool({(0, bytes(h)) for h in hs})
106 + cmap = ExternalCachedBlockPool(16, {(0, bytes(h)) for h in hs})
107 masks, hit = coord.find_longest_cache_hit(hs, max_length=64, cached_block_pool=cmap)
108 assert hit == 64
109 assert masks[0] == [True, True, True, True]
@@ -113,7 +113,7 @@ def test_coordinator_single_full_attention_partial_prefix():
113 groups = [KVCacheGroupSpec(["L0"], _full(16))]
114 coord = _make_coord(groups, hash_block_size=16)
115 hs = _hashes(4)
116 - cmap = ExternalCachedBlockPool({(0, bytes(hs[0])), (0, bytes(hs[1]))})
116 + cmap = ExternalCachedBlockPool(16, {(0, bytes(hs[0])), (0, bytes(hs[1]))})
117 masks, hit = coord.find_longest_cache_hit(hs, max_length=64, cached_block_pool=cmap)
118 assert hit == 32
119 assert masks[0] == [True, True]
@@ -123,7 +123,7 @@ def test_coordinator_single_full_attention_no_hits():
123 groups = [KVCacheGroupSpec(["L0"], _full(16))]
124 coord = _make_coord(groups, hash_block_size=16)
125 hs = _hashes(4)
126 - cmap = ExternalCachedBlockPool(set())
126 + cmap = ExternalCachedBlockPool(16, set())
127 masks, hit = coord.find_longest_cache_hit(hs, max_length=64, cached_block_pool=cmap)
128 assert hit == 0
129 assert masks[0] == []
@@ -135,7 +135,7 @@ def test_coordinator_single_swa_tautological_pool_masks_pre_window():
135 groups = [KVCacheGroupSpec(["L0"], _swa(block_size=16, sliding_window=32))]
136 coord = _make_coord(groups, hash_block_size=16)
137 hs = _hashes(4) # 4 chunks * 16 tokens
138 - cmap = ExternalCachedBlockPool()
138 + cmap = ExternalCachedBlockPool(16)
139 masks, hit = coord.find_longest_cache_hit(hs, max_length=64, cached_block_pool=cmap)
140 assert hit == 64
141 # ceil((sw-1)/block_size) = ceil(31/16) = 2 tail blocks.
@@ -153,7 +153,7 @@ def test_coordinator_hybrid_full_plus_swa_all_hit():
153 ]
154 coord = _make_coord(groups, hash_block_size=16)
155 hs = _hashes(4)
156 - cmap = ExternalCachedBlockPool({(g, bytes(h)) for g in (0, 1) for h in hs})
156 + cmap = ExternalCachedBlockPool(16, {(g, bytes(h)) for g in (0, 1) for h in hs})
157 _masks, hit = coord.find_longest_cache_hit(
158 hs, max_length=64, cached_block_pool=cmap
159 )
@@ -169,7 +169,7 @@ def test_coordinator_hybrid_hole_in_full_clips_both():
169 hs = _hashes(4)
170 exists = {(0, bytes(hs[0])), (0, bytes(hs[2])), (0, bytes(hs[3]))}
171 exists |= {(1, bytes(h)) for h in hs}
172 - cmap = ExternalCachedBlockPool(exists)
172 + cmap = ExternalCachedBlockPool(16, exists)
173 _masks, hit = coord.find_longest_cache_hit(
174 hs, max_length=64, cached_block_pool=cmap
175 )
@@ -188,7 +188,7 @@ def test_coordinator_group_block_size_double_hash():
188 big_hashes = list(chunk_hashes_for_block_size(hs, 16, 32))
189 exists = {(0, bytes(h)) for h in hs}
190 exists |= {(1, bytes(bh)) for bh in big_hashes}
191 - cmap = ExternalCachedBlockPool(exists)
191 + cmap = ExternalCachedBlockPool(16, exists)
192 _masks, hit = coord.find_longest_cache_hit(
193 hs, max_length=64, cached_block_pool=cmap
194 )
@@ -405,7 +405,7 @@ def test_lookup_with_eagle_pops_last_full_attention_block():
405 groups = [KVCacheGroupSpec(["L0"], _full(16))]
406 coord = _make_coord(groups, hash_block_size=16, use_eagle=True)
407 hs = _hashes(4)
408 - cmap = ExternalCachedBlockPool({(0, bytes(h)) for h in hs})
408 + cmap = ExternalCachedBlockPool(16, {(0, bytes(h)) for h in hs})
409 _masks, hit = coord.find_longest_cache_hit(
410 hs, max_length=64, cached_block_pool=cmap
411 )
@@ -426,7 +426,7 @@ def test_load_mask_with_eagle_does_not_double_prune_full_attention():
426 groups = [KVCacheGroupSpec(["L0"], _full(16))]
427 coord = _make_coord(groups, hash_block_size=16, use_eagle=True)
428 hs = _hashes(4)
429 - cmap = ExternalCachedBlockPool({(0, bytes(h)) for h in hs})
429 + cmap = ExternalCachedBlockPool(16, {(0, bytes(h)) for h in hs})
430 _masks, hit = coord.find_longest_cache_hit(
431 hs, max_length=64, cached_block_pool=cmap
432 )
@@ -450,7 +450,7 @@ def test_load_mask_with_eagle_hybrid_full_plus_swa():
450 coord = _make_coord(groups, hash_block_size=16, use_eagle=True)
451 hs = _hashes(4)
452 exists = {(g, bytes(h)) for g in (0, 1) for h in hs}
453 - cmap = ExternalCachedBlockPool(exists)
453 + cmap = ExternalCachedBlockPool(16, exists)
454 _masks, hit = coord.find_longest_cache_hit(
455 hs, max_length=64, cached_block_pool=cmap
456 )
@@ -469,7 +469,7 @@ def test_load_mask_without_eagle_unchanged():
469 groups = [KVCacheGroupSpec(["L0"], _full(16))]
470 coord = _make_coord(groups, hash_block_size=16, use_eagle=False)
471 hs = _hashes(4)
472 - cmap = ExternalCachedBlockPool({(0, bytes(h)) for h in hs})
472 + cmap = ExternalCachedBlockPool(16, {(0, bytes(h)) for h in hs})
473 _masks, hit = coord.find_longest_cache_hit(
474 hs, max_length=64, cached_block_pool=cmap
475 )
tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py
+1
@@ -65,6 +65,7 @@ def _minimal_vllm_config(cache_block_size=16):
65 cfg.cache_config.block_size = cache_block_size
66 cfg.cache_config.num_gpu_blocks = 4
67 cfg.cache_config.hash_block_size = None
68 + cfg.cache_config.prefix_match_unit = None
69 cfg.cache_config.enable_prefix_caching = True
70 cfg.parallel_config.prefill_context_parallel_size = 1
71 cfg.parallel_config.decode_context_parallel_size = 1
vllm/config/cache.py
+10 -10
@@ -53,17 +53,17 @@ class CacheConfig:
53 """Whether block_size was explicitly provided. Derived automatically."""
54 user_specified_mamba_block_size: bool = field(default=False, init=False)
55 """Whether mamba_block_size was explicitly provided. Derived automatically."""
56 - hash_block_size: int | None = Field(default=None, gt=0)
57 - """Block size (in tokens) used for computing Request's block_hashes.
56 + prefix_match_unit: int | None = Field(default=None, gt=0)
57 + """The finest token boundary (in tokens) a prefix-cache hit can land on.
58
59 - This can be set to a finer granularity than the physical KV cache block
60 - sizes (e.g. 8) as long as every KV cache group's `block_size` is divisible
61 - by it. This enables prefix-caching keys to be computed at the finest common
62 - granularity and then merged for larger physical block sizes.
59 + Prefix-cache keys are computed every `prefix_match_unit` tokens. It can
60 + be set finer than the physical KV cache block sizes (e.g. 32 vs a
61 + 1024-token hybrid-model block) as long as every KV cache group's
62 + `block_size` is divisible by it, enabling cache hits at boundaries
63 + inside a physical block. It controls matching granularity only, not how
64 + often states are stored.
65
64 - This config is not static default. If left unspecified, vLLM will choose a
65 - default based on the resolved KV cache groups (typically the smallest KV
66 - cache block size when there are multiple groups).
66 + This equals to the `hash_block_size` used throughout the KV cache code.
67 """
68 gpu_memory_utilization: float = Field(default=0.92, gt=0, le=1)
69 """The fraction of GPU memory to be used for the model executor, which can
@@ -209,7 +209,7 @@ class CacheConfig:
209 "enable_prefix_caching",
210 "prefix_caching_hash_algo",
211 # Prefix-caching implementation detail (doesn't affect compiled graph).
212 - "hash_block_size",
212 + "prefix_match_unit",
213 "mamba_page_size_padded",
214 "skip_page_size_padded",
215 "user_specified_block_size",
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py
+17 -9
@@ -29,10 +29,15 @@ from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
29 class ExternalCachedBlockPool:
30 """Duck-typed BlockPool backed by a ``(group_id, hash)`` exists set."""
31
32 - def __init__(self, exists: set[tuple[int, bytes]] | None = None) -> None:
32 + def __init__(
33 + self,
34 + hash_block_size: int,
35 + exists: set[tuple[int, bytes]] | None = None,
36 + ) -> None:
37 # ``exists=None`` is used on the recv side where hit_length is already
38 # determined and we just want each spec's manager to apply its own mask.
39 self._exists = exists
40 + self.hash_block_size = hash_block_size
41 self.null_block = KVCacheBlock(block_id=0)
42 # Dummy ID 1 for present block for duck-typing.
43 self._present_block = KVCacheBlock(block_id=1)
@@ -166,7 +171,7 @@ class MooncakeStoreCoordinator:
171 masks, _ = self.find_longest_cache_hit(
172 block_hashes,
173 token_len,
169 - ExternalCachedBlockPool(),
174 + ExternalCachedBlockPool(self.hash_block_size),
175 apply_eagle=False,
176 )
177 return masks
@@ -270,7 +275,7 @@ class MooncakeStoreCoordinator:
275 if len(self.attention_groups) == 1:
276 spec, group_ids, manager_cls = self.attention_groups[0]
277 hashes = self.block_hashes_for_spec(block_hashes, spec)
273 - hit_blocks = manager_cls.find_longest_cache_hit(
278 + hit_blocks, hit_length = manager_cls.find_longest_cache_hit(
279 block_hashes=hashes, # type: ignore[arg-type]
280 max_length=max_length,
281 kv_cache_group_ids=group_ids,
@@ -283,11 +288,12 @@ class MooncakeStoreCoordinator:
288 blocks_by_group: list[list[KVCacheBlock]] = [[] for _ in range(num_groups)]
289 for gid, blks in zip(group_ids, hit_blocks, strict=True):
290 blocks_by_group[gid] = blks
286 - return tuple(blocks_by_group), len(hit_blocks[0]) * spec.block_size
291 + return tuple(blocks_by_group), hit_length
292
293 num_groups = len(self.kv_cache_groups)
294 hit_length = max_length
295 hit_blocks_by_group: list[list[KVCacheBlock] | None] = [None] * num_groups
296 + hit_length_by_group: list[int] = [0] * num_groups
297
298 is_simple_hybrid = len(self.attention_groups) == 2 and isinstance(
299 self.attention_groups[0][0], FullAttentionSpec
@@ -298,10 +304,11 @@ class MooncakeStoreCoordinator:
304 curr_hit_length = hit_length
305
306 for idx, (spec, group_ids, manager_cls) in enumerate(self.attention_groups):
301 - cached = hit_blocks_by_group[group_ids[0]]
307 + first_group_id = group_ids[0]
308 + cached = hit_blocks_by_group[first_group_id]
309 if isinstance(spec, FullAttentionSpec) and cached is not None:
303 - curr_hit_length = (
304 - curr_hit_length // spec.block_size * spec.block_size
310 + curr_hit_length = min(
311 + curr_hit_length, hit_length_by_group[first_group_id]
312 )
313 continue
314
@@ -310,7 +317,7 @@ class MooncakeStoreCoordinator:
317 if drop_eagle_block:
318 _max_length = min(curr_hit_length + spec.block_size, max_length)
319 hashes = self.block_hashes_for_spec(block_hashes, spec)
313 - hit_blocks = manager_cls.find_longest_cache_hit(
320 + hit_blocks, _new_hit_length = manager_cls.find_longest_cache_hit(
321 block_hashes=hashes, # type: ignore[arg-type]
322 max_length=_max_length,
323 kv_cache_group_ids=group_ids,
@@ -319,7 +326,6 @@ class MooncakeStoreCoordinator:
326 drop_eagle_block=drop_eagle_block,
327 alignment_tokens=self.lcm_block_size,
328 )
322 - _new_hit_length = len(hit_blocks[0]) * spec.block_size
329 if drop_eagle_block:
330 eagle_verified.add(idx)
331 elif _new_hit_length < curr_hit_length:
@@ -327,6 +333,7 @@ class MooncakeStoreCoordinator:
333 curr_hit_length = _new_hit_length
334 for gid, blocks in zip(group_ids, hit_blocks, strict=True):
335 hit_blocks_by_group[gid] = blocks
336 + hit_length_by_group[gid] = _new_hit_length
337
338 if curr_hit_length >= hit_length:
339 break
@@ -343,6 +350,7 @@ class MooncakeStoreCoordinator:
350 full_blks = hit_blocks_by_group[gid]
351 assert full_blks is not None
352 del full_blks[num_blocks:]
353 + hit_length_by_group[gid] = hit_length
354
355 return (
356 tuple(blks if blks is not None else [] for blks in hit_blocks_by_group),
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py
+3 -1
@@ -1523,7 +1523,9 @@ class MooncakeStoreWorker:
1523 }
1524
1525 _masks, hit_length = self.coord.find_longest_cache_hit(
1526 - block_hashes, token_len, ExternalCachedBlockPool(exists_set)
1526 + block_hashes,
1527 + token_len,
1528 + ExternalCachedBlockPool(self.hash_block_size, exists_set),
1529 )
1530 return hit_length
1531
vllm/engine/arg_utils.py
+5
@@ -693,6 +693,7 @@ class EngineArgs:
693 mamba_cache_dtype: MambaDType = CacheConfig.mamba_cache_dtype
694 mamba_ssm_cache_dtype: MambaDType = CacheConfig.mamba_ssm_cache_dtype
695 mamba_block_size: int | None = get_field(CacheConfig, "mamba_block_size")
696 + prefix_match_unit: int | None = get_field(CacheConfig, "prefix_match_unit")
697 mamba_cache_mode: MambaCacheMode = CacheConfig.mamba_cache_mode
698
699 mamba_backend: MambaBackendEnum = MambaBackendEnum.TRITON
@@ -1194,6 +1195,9 @@ class EngineArgs:
1195 cache_group.add_argument(
1196 "--mamba-block-size", **cache_kwargs["mamba_block_size"]
1197 )
1198 + cache_group.add_argument(
1199 + "--prefix-match-unit", **cache_kwargs["prefix_match_unit"]
1200 + )
1201 cache_group.add_argument(
1202 "--mamba-cache-mode", **cache_kwargs["mamba_cache_mode"]
1203 )
@@ -1899,6 +1903,7 @@ class EngineArgs:
1903 mamba_cache_dtype=self.mamba_cache_dtype,
1904 mamba_ssm_cache_dtype=self.mamba_ssm_cache_dtype,
1905 mamba_block_size=self.mamba_block_size,
1906 + prefix_match_unit=self.prefix_match_unit,
1907 mamba_cache_mode=self.mamba_cache_mode,
1908 kv_offloading_size=self.kv_offloading_size,
1909 kv_offloading_backend=self.kv_offloading_backend,
vllm/v1/core/block_pool.py
+24 -2
@@ -260,7 +260,9 @@ class BlockPool:
260 return
261 new_full_blocks = blocks[num_cached_blocks:num_full_blocks]
262 assert block_mask is None or len(block_mask) == len(new_full_blocks)
263 - block_hashes = resolve_block_hashes(request, self.hash_block_size, block_size)
263 + block_hashes = resolve_block_hashes(
264 + request.block_hashes, self.hash_block_size, block_size
265 + )
266
267 new_block_hashes = block_hashes[num_cached_blocks:]
268 new_hashes: list[ExternalBlockHash] | None = (
@@ -390,7 +392,9 @@ class BlockPool:
392 if not self.enable_kv_cache_events or num_cached_blocks == 0:
393 return
394
393 - block_hashes = resolve_block_hashes(request, self.hash_block_size, block_size)
395 + block_hashes = resolve_block_hashes(
396 + request.block_hashes, self.hash_block_size, block_size
397 + )
398
399 # Collect external hashes and extra_keys for cached blocks.
400 cached_hashes: list[ExternalBlockHash] = []
@@ -622,6 +626,24 @@ class BlockPool:
626 )
627 self.cached_block_hash_to_block.insert(block_hash_with_group_id, block)
628
629 + def move_block_hashes(
630 + self,
631 + src_block: KVCacheBlock,
632 + dst_block: KVCacheBlock,
633 + ) -> None:
634 + """Re-point ``src_block``'s prefix-cache entries to ``dst_block``.
635 +
636 + Used when the request owning ``src_block`` keeps writing into it
637 + : the prefix cache holds a private copy (``dst_block``)
638 + under the same hashes instead. Entries stay live; no events emitted.
639 + """
640 + assert dst_block.block_hash is None
641 + assert dst_block.block_id not in self.cached_block_hashes_by_block
642 + num_tokens = src_block.block_hash_num_tokens
643 + for block_hash in self._remove_cached_block_hashes(src_block):
644 + # `num_tokens` only applies to the first (primary) insertion.
645 + self._insert_block_hash(block_hash, dst_block, num_tokens=num_tokens)
646 +
647 def get_new_blocks(self, num_blocks: int) -> list[KVCacheBlock]:
648 """Get new blocks from the free block pool.
649
vllm/v1/core/kv_cache_coordinator.py
+71 -43
@@ -5,12 +5,11 @@ from collections.abc import Sequence
5 from typing import NamedTuple
6
7 from vllm import envs
8 +from vllm.utils.math_utils import cdiv
9 from vllm.v1.core.block_pool import BlockPool
10 from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector
11 from vllm.v1.core.kv_cache_utils import (
12 BlockHash,
12 - BlockHashList,
13 - BlockHashListWithBlockSize,
13 KVCacheBlock,
14 )
15 from vllm.v1.core.single_type_kv_cache_manager import (
@@ -135,6 +134,7 @@ class KVCacheCoordinator(ABC):
134 new_computed_blocks: tuple[Sequence[KVCacheBlock], ...],
135 num_encoder_tokens: int,
136 total_computed_tokens: int,
137 + num_local_computed_tokens: int,
138 num_tokens_main_model: int,
139 apply_admission_cap: bool = False,
140 ) -> int:
@@ -150,6 +150,8 @@ class KVCacheCoordinator(ABC):
150 num_encoder_tokens: The number of encoder tokens for allocating
151 blocks for cross-attention.
152 total_computed_tokens: Include both local and external tokens.
153 + num_local_computed_tokens: The number of local prefix-cache computed
154 + tokens.
155 num_tokens_main_model: The number of tokens for the main model (aka target
156 model in spec decode). w/o spec decode, it is num_tokens;
157 with spec decode, it is num_tokens - num_lookahead_tokens.
@@ -171,6 +173,7 @@ class KVCacheCoordinator(ABC):
173 num_encoder_tokens,
174 [],
175 0,
176 + 0,
177 num_encoder_tokens,
178 apply_admission_cap=apply_admission_cap,
179 )
@@ -180,6 +183,7 @@ class KVCacheCoordinator(ABC):
183 num_tokens,
184 new_computed_blocks[i],
185 total_computed_tokens,
186 + num_local_computed_tokens,
187 num_tokens_main_model,
188 apply_admission_cap=apply_admission_cap,
189 )
@@ -370,7 +374,7 @@ class KVCacheCoordinator(ABC):
374 pass
375
376 def new_step_starts(self) -> None:
373 - """Called when a new step is started."""
377 + """Notify each manager that a new step is starting."""
378 for manager in self.single_type_managers:
379 manager.new_step_starts()
380
@@ -483,7 +487,7 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):
487 block_hashes: list[BlockHash],
488 max_cache_hit_length: int,
489 ) -> tuple[tuple[list[KVCacheBlock], ...], int]:
486 - hit_blocks = self.single_type_managers[0].find_longest_cache_hit(
490 + hit_blocks, hit_length = self.single_type_managers[0].find_longest_cache_hit(
491 block_hashes=block_hashes,
492 max_length=max_cache_hit_length,
493 kv_cache_group_ids=[0],
@@ -494,7 +498,7 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):
498 dcp_world_size=self.dcp_world_size,
499 pcp_world_size=self.pcp_world_size,
500 )
497 - return hit_blocks, len(hit_blocks[0]) * self.block_size
501 + return hit_blocks, hit_length
502
503
504 class SpecGroup(NamedTuple):
@@ -572,8 +576,26 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
576 "full-attention and Mamba groups, got: "
577 f"{type(g.kv_cache_spec).__name__}."
578 )
579 + # Partial hash hits are limited to full-attention + mamba ("align")
580 + # without context parallelism.
581 + self.enable_partial_hash_hits = dcp_world_size == 1 and any(
582 + isinstance(g.kv_cache_spec, MambaSpec)
583 + and g.kv_cache_spec.mamba_cache_mode == "align"
584 + and g.kv_cache_spec.block_size > hash_block_size
585 + for g in kv_cache_config.kv_cache_groups
586 + )
587 self.verify_and_split_kv_cache_groups()
588
589 + @property
590 + def _cache_hit_alignment_tokens(self) -> int:
591 + # Fine-grained partial hits may return hash-block-aligned lengths;
592 + # otherwise it must stay scheduler-block-aligned.
593 + return (
594 + self.hash_block_size
595 + if self.enable_partial_hash_hits
596 + else self.scheduler_block_size
597 + )
598 +
599 def verify_and_split_kv_cache_groups(self) -> None:
600 """
601 Groups KV cache groups by their spec type for efficient batch processing
@@ -617,14 +639,19 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
639 self.single_type_managers[gid].use_eagle = True
640
641 def cache_blocks(self, request: Request, num_computed_tokens: int) -> None:
620 - # Cache hits in this coordinator are always a multiple of
621 - # ``scheduler_block_size`` tokens (see ``find_longest_cache_hit``).
622 - # Within an aligned region, SWA groups may only consult a subset of blocks
623 - # per ``scheduler_block_size``-segment so the unused blocks also stay
624 - # out of the prefix-cache hash map.
625 - aligned_num_computed_tokens = (
626 - num_computed_tokens // self.scheduler_block_size * self.scheduler_block_size
627 - )
642 + if self.enable_partial_hash_hits:
643 + aligned_num_computed_tokens = num_computed_tokens
644 + else:
645 + # Cache hits in this coordinator are always a multiple of
646 + # ``scheduler_block_size`` tokens (see ``find_longest_cache_hit``).
647 + # Within an aligned region, SWA groups may only consult a subset of
648 + # blocks per ``scheduler_block_size``-segment so the unused blocks
649 + # also stay out of the prefix-cache hash map.
650 + aligned_num_computed_tokens = (
651 + num_computed_tokens
652 + // self.scheduler_block_size
653 + * self.scheduler_block_size
654 + )
655 for manager in self.single_type_managers:
656 num_tokens_to_cache = aligned_num_computed_tokens
657 # EAGLE groups match one block past each aligned boundary and drop
@@ -667,17 +694,11 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
694 - The number of tokens of the longest cache hit.
695 """
696
670 - def _get_block_hashes(block_size: int) -> BlockHashList:
671 - if block_size == self.hash_block_size:
672 - return block_hashes
673 - return BlockHashListWithBlockSize(
674 - block_hashes, self.hash_block_size, block_size
675 - )
676 -
697 num_groups = len(self.kv_cache_config.kv_cache_groups)
698 hit_length = max_cache_hit_length
699 longest_hit_length = 0
700 hit_blocks_by_group: list[list[KVCacheBlock] | None] = [None] * num_groups
701 + hit_length_by_group: list[int] = [0] * num_groups
702
703 # Simple hybrid (1 full attn + 1 other): one iteration suffices.
704 # Full attn is always first if it exists.
@@ -696,40 +717,53 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
717 for idx, (spec, group_ids, manager_cls, use_eagle) in enumerate(
718 self.attention_groups
719 ):
699 - group_block_size = self.single_type_managers[group_ids[0]].block_size
700 - cached_blocks = hit_blocks_by_group[group_ids[0]]
720 + first_group_id = group_ids[0]
721 + # DCP/PCP shard each block's KV across ranks, so the manager's
722 + # effective block size may exceed the spec's.
723 + group_block_size = self.single_type_managers[first_group_id].block_size
724 + cached_blocks = hit_blocks_by_group[first_group_id]
725 if isinstance(spec, FullAttentionSpec) and cached_blocks is not None:
726 # Full attention is downward-closed: we only need to look
727 # up cached blocks once; on subsequent iterations just trim
728 # to the (reduced) current hit length.
705 - curr_hit_length = (
706 - curr_hit_length // group_block_size * group_block_size
729 + curr_hit_length = min(
730 + curr_hit_length, hit_length_by_group[first_group_id]
731 )
732 continue
733
734 drop_eagle_block = use_eagle and idx not in eagle_verified
735
736 _max_length = curr_hit_length
713 - if drop_eagle_block:
714 - # Eagle needs to match one more block and then pop the last.
737 + # Eagle matches one extra drop unit (one hash unit for
738 + # fine-grained managers, else one cache block) and then drops
739 + # it, landing back at the candidate length. No margin for
740 + # mamba: its finder never drops (draft models have no mamba
741 + # layers), so the hit would grow past the candidate.
742 + if drop_eagle_block and not isinstance(spec, MambaSpec):
743 + eagle_margin = (
744 + self.hash_block_size
745 + if self.enable_partial_hash_hits
746 + and manager_cls.supports_fine_grained_hash_lookup
747 + and group_block_size > self.hash_block_size
748 + else group_block_size
749 + )
750 _max_length = min(
716 - curr_hit_length + group_block_size, max_cache_hit_length
751 + curr_hit_length + eagle_margin, max_cache_hit_length
752 )
718 - hit_blocks = manager_cls.find_longest_cache_hit(
719 - block_hashes=_get_block_hashes(group_block_size),
753 + hit_blocks, _new_hit_length = manager_cls.find_longest_cache_hit(
754 + block_hashes=block_hashes,
755 max_length=_max_length,
756 kv_cache_group_ids=group_ids,
757 block_pool=self.block_pool,
758 kv_cache_spec=spec,
759 drop_eagle_block=drop_eagle_block,
725 - alignment_tokens=self.scheduler_block_size,
760 + alignment_tokens=self._cache_hit_alignment_tokens,
761 dcp_world_size=(
762 self.dcp_world_size
763 if isinstance(spec, FullAttentionSpec)
764 else 1
765 ),
766 )
732 - _new_hit_length = len(hit_blocks[0]) * group_block_size
767 if drop_eagle_block:
768 eagle_verified.add(idx)
769 elif _new_hit_length < curr_hit_length:
@@ -738,6 +772,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
772 curr_hit_length = _new_hit_length
773 for group_id, blocks in zip(group_ids, hit_blocks):
774 hit_blocks_by_group[group_id] = blocks
775 + hit_length_by_group[group_id] = _new_hit_length
776
777 longest_hit_length = max(longest_hit_length, curr_hit_length)
778
@@ -753,10 +788,11 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
788 group_block_size = self.single_type_managers[
789 first_group.group_ids[0]
790 ].block_size
756 - num_blocks = hit_length // group_block_size
791 + num_blocks = cdiv(hit_length, group_block_size)
792 for group_id in first_group.group_ids:
793 if (blks := hit_blocks_by_group[group_id]) is not None:
794 del blks[num_blocks:]
795 + hit_length_by_group[group_id] = hit_length
796
797 # Uncached shared prefix detection: If any attn. group cached a longer prefix
798 # than the current prefix, it is an uncached common prefix across requests:
@@ -776,28 +812,20 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
812 (blocks_per_group, hit_lengths_per_group)
813 """
814
779 - def _get_block_hashes(kv_cache_spec: KVCacheSpec) -> BlockHashList:
780 - if kv_cache_spec.block_size == self.hash_block_size:
781 - return block_hashes
782 - return BlockHashListWithBlockSize(
783 - block_hashes, self.hash_block_size, kv_cache_spec.block_size
784 - )
785 -
815 num_groups = len(self.kv_cache_config.kv_cache_groups)
816 hit_blocks: list[list[KVCacheBlock]] = [[] for _ in range(num_groups)]
817 hit_lengths: list[int] = [0] * num_groups
818
819 for spec, group_ids, manager_cls, use_eagle in self.attention_groups:
791 - blocks = manager_cls.find_longest_cache_hit(
792 - block_hashes=_get_block_hashes(spec),
820 + blocks, group_hit = manager_cls.find_longest_cache_hit(
821 + block_hashes=block_hashes,
822 max_length=max_cache_hit_length,
823 kv_cache_group_ids=group_ids,
824 block_pool=self.block_pool,
825 kv_cache_spec=spec,
826 drop_eagle_block=use_eagle,
798 - alignment_tokens=self.scheduler_block_size,
827 + alignment_tokens=self._cache_hit_alignment_tokens,
828 )
800 - group_hit = len(blocks[0]) * spec.block_size
829 for gid, blks in zip(group_ids, blocks):
830 hit_blocks[gid] = blks
831 hit_lengths[gid] = group_hit
vllm/v1/core/kv_cache_manager.py
+21 -2
@@ -11,7 +11,7 @@ from vllm.logger import init_logger
11 from vllm.utils.math_utils import cdiv
12 from vllm.v1.core.kv_cache_coordinator import get_kv_cache_coordinator
13 from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector
14 -from vllm.v1.core.kv_cache_utils import KVCacheBlock
14 +from vllm.v1.core.kv_cache_utils import KVCacheBlock, KVCacheBlockCopy
15 from vllm.v1.kv_cache_interface import (
16 AttentionSpec,
17 CrossAttentionSpec,
@@ -404,6 +404,7 @@ class KVCacheManager:
404 new_computed_blocks=new_computed_block_list,
405 num_encoder_tokens=num_encoder_tokens,
406 total_computed_tokens=total_computed_tokens,
407 + num_local_computed_tokens=num_local_computed_tokens,
408 num_tokens_main_model=full_num_tokens,
409 apply_admission_cap=True,
410 )
@@ -438,6 +439,7 @@ class KVCacheManager:
439 num_encoder_tokens=num_encoder_tokens,
440 total_computed_tokens=num_local_computed_tokens
441 + num_external_computed_tokens,
442 + num_local_computed_tokens=num_local_computed_tokens,
443 num_tokens_main_model=num_tokens_main_model,
444 )
445
@@ -665,6 +667,23 @@ class KVCacheManager:
667 ids.extend(mgr.take_new_block_ids())
668 return ids
669
670 + def take_kv_cache_block_copies(
671 + self,
672 + ) -> tuple[list[KVCacheBlockCopy], list[KVCacheBlock]]:
673 + """Drain pending copies and return their retained endpoints."""
674 + pending_copies: list[tuple[KVCacheBlock, KVCacheBlock]] = []
675 + for mgr in self.coordinator.single_type_managers:
676 + pending_copies.extend(mgr.take_pending_cow_copies())
677 + copies = [
678 + KVCacheBlockCopy(
679 + src_block_id=source_block.block_id,
680 + dst_block_id=cow_block.block_id,
681 + )
682 + for source_block, cow_block in pending_copies
683 + ]
684 + retained_blocks = [block for pair in pending_copies for block in pair]
685 + return copies, retained_blocks
686 +
687 def new_step_starts(self) -> None:
669 - """Called when a new step is started."""
688 + """Notify the coordinator that a new step is starting."""
689 self.coordinator.new_step_starts()
vllm/v1/core/kv_cache_utils.py
+38 -19
@@ -10,7 +10,7 @@ from collections import defaultdict
10 from collections.abc import Callable, Iterable, Iterator, Sequence
11 from dataclasses import dataclass, replace
12 from functools import partial
13 -from typing import Any, NewType, TypeAlias, cast, overload
13 +from typing import Any, NamedTuple, NewType, TypeAlias, cast, overload
14
15 from vllm import envs
16 from vllm.config import VllmConfig
@@ -126,7 +126,7 @@ class KVCacheBlock:
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.
129 + # the full block boundary; partial entries 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.
@@ -176,6 +176,11 @@ class KVCacheBlock:
176 )
177
178
179 +class KVCacheBlockCopy(NamedTuple):
180 + src_block_id: int
181 + dst_block_id: int
182 +
183 +
184 class FreeKVCacheBlockQueue:
185 """This class organizes a list of KVCacheBlock objects to a doubly linked
186 list of free blocks. We implement this class instead of using Python
@@ -631,11 +636,11 @@ def resolve_kv_cache_block_sizes(
636 Mamba groups keep their full per-rank state and are not scaled.
637 - ``hash_block_size`` is the granularity at which ``Request.block_hashes``
638 is computed. Single group: equals scheduler block size. Multiple groups:
634 - ``cache_config.hash_block_size`` override if set, else the GCD of group
635 - block sizes; every group's block size must be divisible by it. Returns
636 - the scheduler block size (i.e. disables finer hashing) if block hashing
637 - is inactive or a mamba group's block size diverges from the cache
638 - block size (mamba_cache_mode != "align").
639 + ``cache_config.prefix_match_unit`` override if set, else the GCD of
640 + group block sizes; every group's block size must be divisible by it.
641 + Returns the scheduler block size (i.e. disables finer hashing) if block
642 + hashing is inactive or a mamba group's block size diverges from the
643 + cache block size (mamba_cache_mode != "align").
644 """
645 cache_config = vllm_config.cache_config
646 dcp = vllm_config.parallel_config.decode_context_parallel_size
@@ -671,14 +676,14 @@ def resolve_kv_cache_block_sizes(
676 ):
677 return scheduler_block_size, scheduler_block_size
678
674 - requested = cache_config.hash_block_size
679 + requested = cache_config.prefix_match_unit
680 hash_block_size = (
681 requested if requested is not None else math.gcd(*group_block_sizes)
682 )
683 if any(bs % hash_block_size != 0 for bs in group_block_sizes):
684 raise ValueError(
680 - f"Invalid hash_block_size={hash_block_size}; all KV cache group "
681 - f"block sizes must be divisible by hash_block_size. "
685 + f"Invalid prefix_match_unit={hash_block_size}; all KV cache group "
686 + f"block sizes must be divisible by prefix_match_unit. "
687 f"Got group block sizes={group_block_sizes}."
688 )
689 return scheduler_block_size, hash_block_size
@@ -2252,19 +2257,33 @@ BlockHashList = list[BlockHash] | BlockHashListWithBlockSize
2257
2258
2259 def resolve_block_hashes(
2255 - request: Request,
2260 + block_hashes: BlockHashList,
2261 hash_block_size: int,
2262 block_size: int,
2263 + *,
2264 + supports_fine_grained_hash_lookup: bool = False,
2265 + alignment_tokens: int | None = None,
2266 ) -> BlockHashList:
2259 - """Resolve the block-hash view for ``request`` at ``block_size``.
2267 + """Resolve the block-hash view at ``block_size``.
2268
2261 - When ``block_size`` equals ``hash_block_size``, reuse the request's
2262 - precomputed ``block_hashes`` directly; otherwise recalculate at
2263 - ``block_size`` granularity (``block_size`` must be a multiple of
2264 - ``hash_block_size``, which happens when KV cache groups differ in
2265 - block size).
2269 + When ``block_size`` equals ``hash_block_size``, reuse the precomputed block
2270 + hashes directly; otherwise view them at ``block_size`` granularity.
2271 + Fine-grained lookup keeps the original hashes for partial cache hits.
2272 """
2273 if block_size == hash_block_size:
2268 - return request.block_hashes
2274 + return block_hashes
2275 + if isinstance(block_hashes, BlockHashListWithBlockSize):
2276 + # Already a block-size view
2277 + assert block_hashes.scale_factor == block_size // hash_block_size
2278 + return block_hashes
2279 + # Fine-grained partial hits keep the raw hashes. The caller passes
2280 + # alignment_tokens = hash_block_size to enable them, else >= block_size.
2281 + if (
2282 + supports_fine_grained_hash_lookup
2283 + and alignment_tokens is not None
2284 + and alignment_tokens < block_size
2285 + and block_size % alignment_tokens == 0
2286 + ):
2287 + return block_hashes
2288 assert block_size % hash_block_size == 0
2270 - return BlockHashListWithBlockSize(request.block_hashes, hash_block_size, block_size)
2289 + return BlockHashListWithBlockSize(block_hashes, hash_block_size, block_size)
vllm/v1/core/sched/output.py
+5
@@ -16,10 +16,12 @@ if TYPE_CHECKING:
16 from vllm.multimodal.inputs import MultiModalFeatureSpec
17 from vllm.pooling_params import PoolingParams
18 from vllm.sampling_params import SamplingParams
19 + from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy
20 from vllm.v1.request import Request
21 else:
22 ECConnectorMetadata = object
23 KVConnectorMetadata = object
24 + KVCacheBlockCopy = object
25 LoRARequest = object
26 MultiModalFeatureSpec = object
27 PoolingParams = object
@@ -240,6 +242,9 @@ class SchedulerOutput:
242 # preventing stale NaN/data from corrupting attention or SSM computation.
243 new_block_ids_to_zero: list[int] | None = None
244
245 + # CoW copies to apply after zeroing new blocks and before forward.
246 + kv_cache_block_copies: list[KVCacheBlockCopy] | None = None
247 +
248 # Dynamic speculative decoding: optimal K chosen by scheduler.
249 # Number of spec tokens to schedule for the next step.
250 num_spec_tokens_to_schedule: int = 0
vllm/v1/core/sched/scheduler.py
+68 -8
@@ -259,6 +259,7 @@ class Scheduler(SchedulerInterface):
259 # Create the KV cache manager.
260 if hash_block_size is None:
261 hash_block_size = block_size
262 + self.hash_block_size = hash_block_size
263 self.kv_cache_manager = KVCacheManager(
264 kv_cache_config=kv_cache_config,
265 max_model_len=self.max_model_len,
@@ -297,6 +298,13 @@ class Scheduler(SchedulerInterface):
298 self.need_mamba_block_aligned_split = (
299 self.has_mamba_layers and self.cache_config.mamba_cache_mode == "align"
300 )
301 + # A finer prefix_match_unit is configured: a mamba partial tail entry
302 + # can only be registered by a step ending exactly at the prompt's last
303 + # hash boundary, so the split adds that stop.
304 + self.mamba_partial_cache_hit = (
305 + self.need_mamba_block_aligned_split
306 + and self.hash_block_size < self.block_size
307 + )
308
309 # Counts of non-empty steps scheduled / processed. update_from_output
310 # is called once per scheduled step in FIFO order, so these stay in sync.
@@ -368,9 +376,23 @@ class Scheduler(SchedulerInterface):
376 if self.use_eagle:
377 last_cache_position = max(last_cache_position - block_size, 0)
378 num_computed_tokens_after_sched = num_computed_tokens + num_new_tokens
371 - if num_computed_tokens_after_sched < last_cache_position:
372 - # align to block_size
373 - num_new_tokens = num_new_tokens // block_size * block_size
379 + next_boundary = (num_computed_tokens // block_size + 1) * block_size
380 + if (
381 + num_computed_tokens % block_size != 0
382 + and next_boundary <= last_cache_position
383 + and num_computed_tokens_after_sched > next_boundary
384 + ):
385 + # Resumed mid-block (partial hash hit): stop the first chunk
386 + # at the next block boundary so chunk ends re-align to blocks
387 + # and the boundary state is materialized before running past.
388 + num_new_tokens = next_boundary - num_computed_tokens
389 + elif num_computed_tokens_after_sched < last_cache_position:
390 + # Align the chunk END (not its length) to block_size;
391 + # identical to flooring the length when the start is
392 + # block-aligned. May yield 0 (insufficient budget to reach
393 + # the next boundary); the caller then skips the request.
394 + aligned_end = num_computed_tokens_after_sched // block_size * block_size
395 + num_new_tokens = max(aligned_end - num_computed_tokens, 0)
396 elif (
397 num_computed_tokens
398 < last_cache_position
@@ -378,9 +400,23 @@ class Scheduler(SchedulerInterface):
400 ):
401 # force to cache the last chunk
402 num_new_tokens = last_cache_position - num_computed_tokens
381 - else:
382 - # prefill the last few tokens
383 - pass
403 + elif self.mamba_partial_cache_hit:
404 + # Prefill of the final partial block: stop once at the
405 + # prompt's last hash boundary so the mamba partial tail entry
406 + # can be registered.
407 + tail_boundary = (
408 + request.num_prompt_tokens
409 + // self.hash_block_size
410 + * self.hash_block_size
411 + )
412 + if (
413 + num_computed_tokens
414 + < tail_boundary
415 + < num_computed_tokens_after_sched
416 + and tail_boundary < request.num_prompt_tokens
417 + and tail_boundary > last_cache_position
418 + ):
419 + num_new_tokens = tail_boundary - num_computed_tokens
420
421 # Marconi cache admission optimization:
422 # cache common prefixes by scheduling num_new_tokens = common prefix length
@@ -1084,6 +1120,16 @@ class Scheduler(SchedulerInterface):
1120 new_block_ids_to_zero = (
1121 (new_attn_block_ids or None) if self.needs_kv_cache_zeroing else None
1122 )
1123 + kv_cache_block_copies, cow_retained_blocks = (
1124 + self.kv_cache_manager.take_kv_cache_block_copies()
1125 + )
1126 + if kv_cache_block_copies:
1127 + # The copies run with this step's execution; the first non-empty
1128 + # step at or after it gets seq `sched_step_seq + 1` (0-token steps
1129 + # do not advance the seq), and its completion implies the copies
1130 + # have run.
1131 + self._free_cow_retained_blocks(cow_retained_blocks, self.sched_step_seq + 1)
1132 + pending_kv_cache_block_copies = kv_cache_block_copies or None
1133
1134 # Dynamic speculative decoding: compute optimal K
1135 num_spec_tokens_to_schedule = self.num_spec_tokens
@@ -1108,6 +1154,7 @@ class Scheduler(SchedulerInterface):
1154 finished_req_ids=self.finished_req_ids,
1155 free_encoder_mm_hashes=self.encoder_cache_manager.get_freed_mm_hashes(),
1156 new_block_ids_to_zero=new_block_ids_to_zero,
1157 + kv_cache_block_copies=pending_kv_cache_block_copies,
1158 num_spec_tokens_to_schedule=num_spec_tokens_to_schedule,
1159 )
1160
@@ -2148,11 +2195,23 @@ class Scheduler(SchedulerInterface):
2195 if blocks:
2196 self.deferred_frees.append((self.sched_step_seq, blocks))
2197
2198 + def _free_cow_retained_blocks(
2199 + self, blocks: list[KVCacheBlock], fence_seq: int
2200 + ) -> None:
2201 + """Release CoW copy retentions, deferring their return to the block
2202 + pool while the step that runs the copy may still be in flight.
2203 + """
2204 + if not self.defer_block_free or fence_seq <= self.processed_step_seq:
2205 + self.kv_cache_manager.block_pool.free_blocks(blocks)
2206 + return
2207 + self.deferred_frees.append((fence_seq, blocks[::-1]))
2208 +
2209 def _drain_deferred_frees(self):
2210 """Return deferred blocks whose fence step has completed.
2211
2154 - Entries are appended with monotonically non-decreasing fences, so
2155 - stop at the first one that is still pending.
2212 + Fences are appended in near-monotonic order (a CoW retention fence
2213 + can lead request-free fences by one step), so stop at the first
2214 + pending one; any satisfied entry behind it is merely freed later.
2215 """
2216 while self.deferred_frees:
2217 fence, _ = self.deferred_frees[0]
@@ -2401,6 +2460,7 @@ class Scheduler(SchedulerInterface):
2460 new_computed_blocks=self.kv_cache_manager.empty_kv_cache_blocks.blocks,
2461 num_encoder_tokens=0,
2462 total_computed_tokens=request.num_computed_tokens,
2463 + num_local_computed_tokens=request.num_computed_tokens,
2464 num_tokens_main_model=full_num_tokens,
2465 apply_admission_cap=True,
2466 )
vllm/v1/core/single_type_kv_cache_manager.py
+363 -57
@@ -4,13 +4,16 @@ import itertools
4 from abc import ABC, abstractmethod
5 from collections import defaultdict
6 from collections.abc import Sequence
7 +from typing import ClassVar
8
9 from vllm.utils.math_utils import cdiv
10 from vllm.v1.core.block_pool import BlockPool
11 from vllm.v1.core.kv_cache_utils import (
12 BlockHashList,
13 + BlockHashListWithBlockSize,
14 BlockHashWithGroupId,
15 KVCacheBlock,
16 + resolve_block_hashes,
17 )
18 from vllm.v1.kv_cache_interface import (
19 ChunkedLocalAttentionSpec,
@@ -36,6 +39,8 @@ class SingleTypeKVCacheManager(ABC):
39 logic of one specific type of attention layer.
40 """
41
42 + supports_fine_grained_hash_lookup: ClassVar[bool] = False
43 +
44 def __init__(
45 self,
46 kv_cache_spec: KVCacheSpec,
@@ -106,16 +111,34 @@ class SingleTypeKVCacheManager(ABC):
111 # determining the attention groups.
112 self.use_eagle = False
113
114 + # Partial-hit copy-on-write bookkeeping. Populated only by fine-grained
115 + # managers (full attention, mamba "align"); harmlessly empty elsewhere.
116 + self._partial_hit_reqs: dict[str, tuple[int, KVCacheBlock]] = {}
117 + self._pending_cow_copies: list[tuple[KVCacheBlock, KVCacheBlock]] = []
118 +
119 @classmethod
120 def _get_num_evictable_blocks(cls, blocks: Sequence[KVCacheBlock]):
121 return sum(blk.ref_cnt == 0 and not blk.is_null for blk in blocks)
122
123 + def _has_partial_local_hit(
124 + self,
125 + new_computed_blocks: Sequence[KVCacheBlock],
126 + num_local_computed_tokens: int,
127 + ) -> bool:
128 + # The local prefix-cache hit ends inside one of this manager's
129 + # blocks: the shared tail block needs CoW.
130 + return (
131 + len(new_computed_blocks) > 0
132 + and num_local_computed_tokens % self.block_size != 0
133 + )
134 +
135 def get_num_blocks_to_allocate(
136 self,
137 request_id: str,
138 num_tokens: int,
139 new_computed_blocks: Sequence[KVCacheBlock],
140 total_computed_tokens: int,
141 + num_local_computed_tokens: int,
142 num_tokens_main_model: int,
143 apply_admission_cap: bool = False,
144 ) -> int:
@@ -130,6 +153,8 @@ class SingleTypeKVCacheManager(ABC):
153 prefix caching.
154 total_computed_tokens: Include both local and external computed
155 tokens.
156 + num_local_computed_tokens: The number of local prefix-cache computed
157 + tokens.
158 num_tokens_main_model: The number of tokens for the main model (aka target
159 model in spec decode). w/o spec decode, it is num_tokens;
160 with spec decode, it is num_tokens - num_lookahead_tokens.
@@ -189,6 +214,10 @@ class SingleTypeKVCacheManager(ABC):
214 num_evictable_blocks = self._get_num_evictable_blocks(
215 new_computed_blocks[num_skipped_new_computed_blocks:]
216 )
217 + if self._has_partial_local_hit(new_computed_blocks, num_local_computed_tokens):
218 + # Reserve the extra block that allocate_new_blocks pulls for the
219 + # partial-hit CoW redirect.
220 + num_new_blocks += 1
221 return num_new_blocks + num_evictable_blocks
222
223 def add_local_computed_blocks(
@@ -242,6 +271,13 @@ class SingleTypeKVCacheManager(ABC):
271 # them so cache_blocks() will not try to re-cache blocks that already
272 # have a block_hash set.
273 self.num_cached_block[request_id] = len(req_blocks)
274 + if self._has_partial_local_hit(new_computed_blocks, num_local_computed_tokens):
275 + # Record the partial tail for the CoW redirect in
276 + # allocate_new_blocks; cap the cached count at the full blocks so
277 + # cache_blocks() re-caches the private copy once full.
278 + block_idx = num_local_computed_tokens // self.block_size
279 + self._partial_hit_reqs[request_id] = (block_idx, new_computed_blocks[-1])
280 + self.num_cached_block[request_id] = block_idx
281
282 def allocate_external_computed_blocks(
283 self,
@@ -299,17 +335,29 @@ class SingleTypeKVCacheManager(ABC):
335 Returns:
336 The new allocated blocks.
337 """
338 + cow_blocks: list[KVCacheBlock] = []
339 + if request_id in self._partial_hit_reqs:
340 + # Partial hit: redirect the shared tail to a private CoW block.
341 + # Replacing in place keeps the length-based allocation below
342 + # correct; the extra block was reserved by
343 + # get_num_blocks_to_allocate.
344 + block_idx, source_block = self._partial_hit_reqs.pop(request_id)
345 + cow_block = self.block_pool.get_new_blocks(1)[0]
346 + self._apply_cow(request_id, block_idx, source_block, cow_block)
347 + self.new_block_ids.append(cow_block.block_id)
348 + cow_blocks.append(cow_block)
349 +
350 req_blocks = self.req_to_blocks[request_id]
351 num_required_blocks = cdiv(num_tokens, self.block_size)
352 num_new_blocks = num_required_blocks - len(req_blocks)
353 if num_new_blocks <= 0:
306 - return []
354 + return cow_blocks
355 else:
356 new_blocks = self.block_pool.get_new_blocks(num_new_blocks)
357 req_blocks.extend(new_blocks)
358 if self._record_new_block_ids:
359 self.new_block_ids.extend(b.block_id for b in new_blocks)
312 - return new_blocks
360 + return cow_blocks + new_blocks
361
362 def take_new_block_ids(self) -> list[int]:
363 """Drain and return block IDs allocated since the last call."""
@@ -317,6 +365,36 @@ class SingleTypeKVCacheManager(ABC):
365 self.new_block_ids = []
366 return ids
367
368 + def take_pending_cow_copies(
369 + self,
370 + ) -> list[tuple[KVCacheBlock, KVCacheBlock]]:
371 + """Drain pending CoW source and destination block pairs."""
372 + pending_copies = self._pending_cow_copies
373 + self._pending_cow_copies = []
374 + return pending_copies
375 +
376 + def _apply_cow(
377 + self,
378 + request_id: str,
379 + block_idx: int,
380 + source_block: KVCacheBlock,
381 + cow_block: KVCacheBlock,
382 + ) -> None:
383 + """Redirect a partial prefix-cache hit to a private CoW block.
384 +
385 + Both copy endpoints stay retained until the copy has run on the worker,
386 + so a same-step free cannot recycle them: ``source_block`` keeps its
387 + hit-ref, ``cow_block`` takes an extra ref beyond the one handed to the
388 + request.
389 + """
390 + req_blocks = self.req_to_blocks[request_id]
391 + assert block_idx < len(req_blocks)
392 + assert req_blocks[block_idx] is source_block
393 + assert not source_block.is_null and source_block.ref_cnt > 0
394 + req_blocks[block_idx] = cow_block
395 + self._pending_cow_copies.append((source_block, cow_block))
396 + cow_block.ref_cnt += 1
397 +
398 def cache_blocks(
399 self,
400 request: Request,
@@ -398,6 +476,7 @@ class SingleTypeKVCacheManager(ABC):
476 # Default to [] in case a request is freed (aborted) before alloc.
477 req_blocks = self.req_to_blocks.pop(request_id, [])
478 self.num_cached_block.pop(request_id, None)
479 + self._partial_hit_reqs.pop(request_id, None)
480 return req_blocks
481
482 def free(self, request_id: str) -> None:
@@ -439,7 +518,7 @@ class SingleTypeKVCacheManager(ABC):
518 alignment_tokens: int,
519 dcp_world_size: int = 1,
520 pcp_world_size: int = 1,
442 - ) -> tuple[list[KVCacheBlock], ...]:
521 + ) -> tuple[tuple[list[KVCacheBlock], ...], int]:
522 """
523 Get the longest cache hit prefix of the blocks that is not longer than
524 `max_length`. The prefix should be a common prefix hit for all the
@@ -466,11 +545,9 @@ class SingleTypeKVCacheManager(ABC):
545 pcp_world_size: The world size of prefill context parallelism.
546
547 Returns:
469 - A list of cached blocks with skipped blocks replaced by null block
470 - for each kv cache group in `kv_cache_group_ids`.
471 - Return a list of length `len(kv_cache_group_ids)`, where the i-th
472 - element is a list of cached blocks for the i-th kv cache group
473 - in `kv_cache_group_ids`.
548 + A tuple containing cached blocks and the exact cache-hit length in
549 + tokens. The cached block tuple has skipped blocks replaced by null
550 + blocks for each kv cache group in `kv_cache_group_ids`.
551 For example, sliding window manager should return a list like
552 ([NULL, NULL, KVCacheBlock(7), KVCacheBlock(8)]) for block size 4
553 and sliding window 8 and len(kv_cache_group_ids) = 1.
@@ -558,11 +635,12 @@ class SingleTypeKVCacheManager(ABC):
635 return 0
636
637 def new_step_starts(self) -> None:
561 - # do nothing by default
638 return None
639
640
641 class FullAttentionManager(SingleTypeKVCacheManager):
642 + supports_fine_grained_hash_lookup: ClassVar[bool] = True
643 +
644 @classmethod
645 def find_longest_cache_hit(
646 cls,
@@ -575,42 +653,131 @@ class FullAttentionManager(SingleTypeKVCacheManager):
653 alignment_tokens: int,
654 dcp_world_size: int = 1,
655 pcp_world_size: int = 1,
578 - ) -> tuple[list[KVCacheBlock], ...]:
656 + ) -> tuple[tuple[list[KVCacheBlock], ...], int]:
657 assert isinstance(
658 kv_cache_spec, FullAttentionSpec | ChunkedLocalAttentionSpec
659 ), (
660 "FullAttentionManager can only be used for full attention "
661 "and chunked local attention groups"
662 )
585 - computed_blocks: tuple[list[KVCacheBlock], ...] = tuple(
586 - [] for _ in range(len(kv_cache_group_ids))
587 - )
663 block_size = kv_cache_spec.block_size
664 if dcp_world_size * pcp_world_size > 1:
665 + # DCP/PCP shard each block's KV across ranks; hashes must be
666 + # viewed at the sharded (scaled) block size.
667 block_size *= dcp_world_size * pcp_world_size
591 - max_num_blocks = max_length // block_size
592 - for block_hash in itertools.islice(block_hashes, max_num_blocks):
593 - # block_hashes is a chain of block hashes. If a block hash is not
594 - # in the cached_block_hash_to_id, the following block hashes are
595 - # not computed yet for sure.
596 - if cached_block := block_pool.get_cached_block(
597 - block_hash, kv_cache_group_ids
598 - ):
599 - for computed, cached in zip(computed_blocks, cached_block):
668 + block_hashes = resolve_block_hashes(
669 + block_hashes,
670 + block_pool.hash_block_size,
671 + block_size,
672 + supports_fine_grained_hash_lookup=cls.supports_fine_grained_hash_lookup,
673 + alignment_tokens=alignment_tokens,
674 + )
675 +
676 + # Fine-grained mode (alignment_tokens == hash_block_size <
677 + # block_size): resolve_block_hashes kept the raw hash-granularity
678 + # list so interior boundaries can be probed.
679 + fine_grained = (
680 + alignment_tokens < block_size and block_size % alignment_tokens == 0
681 + )
682 + if fine_grained:
683 + assert isinstance(block_hashes, list)
684 + full_block_hashes: BlockHashList = BlockHashListWithBlockSize(
685 + block_hashes, alignment_tokens, block_size
686 + )
687 + else:
688 + full_block_hashes = block_hashes
689 +
690 + computed_blocks: tuple[list[KVCacheBlock], ...] = tuple(
691 + [] for _ in range(len(kv_cache_group_ids))
692 + )
693 + # Phase 1: longest run of cached full blocks from the start. A missing
694 + # block implies every later block misses too (chained hashes).
695 + for block_hash in itertools.islice(full_block_hashes, max_length // block_size):
696 + cached_block = block_pool.get_cached_block(block_hash, kv_cache_group_ids)
697 + if not cached_block:
698 + break
699 + for computed, cached in zip(computed_blocks, cached_block):
700 + computed.append(cached)
701 + hit_length = len(computed_blocks[0]) * block_size
702 +
703 + # Phase 2 (fine-grained only): extend into the first non-full block by
704 + # probing its interior hash boundaries high-to-low (longest hit first).
705 + if fine_grained:
706 + assert isinstance(block_hashes, list)
707 + scale_factor = block_size // alignment_tokens
708 + first_partial_idx = len(computed_blocks[0]) * scale_factor
709 + max_partial_idx = min(
710 + first_partial_idx + scale_factor - 1,
711 + max_length // alignment_tokens,
712 + len(block_hashes),
713 + )
714 + for fine_idx in range(max_partial_idx - 1, first_partial_idx - 1, -1):
715 + cached_tail = block_pool.get_cached_block(
716 + block_hashes[fine_idx], kv_cache_group_ids
717 + )
718 + if not cached_tail:
719 + continue
720 + for computed, cached in zip(computed_blocks, cached_tail):
721 computed.append(cached)
601 - else:
722 + hit_length = (fine_idx + 1) * alignment_tokens
723 break
603 - if drop_eagle_block and computed_blocks[0]:
604 - # Need to drop the last matched block if eagle is enabled.
605 - for computed in computed_blocks:
606 - computed.pop()
607 - while (
608 - block_size != alignment_tokens # Faster for common case.
609 - and len(computed_blocks[0]) * block_size % alignment_tokens != 0
610 - ):
611 - for computed in computed_blocks:
612 - computed.pop()
613 - return computed_blocks
724 +
725 + # Eagle needs the tokens right before the generation point recomputed:
726 + # drop one hash unit when fine-grained (the tail block's KV is
727 + # append-only, so it still covers the reduced length), else one cache
728 + # block.
729 + if drop_eagle_block and hit_length > 0:
730 + hit_length -= min(alignment_tokens, block_size)
731 + # Round down to the alignment; a no-op when fine-grained (hits land on
732 + # hash boundaries by construction) and when alignment_tokens ==
733 + # block_size. Then trim blocks past the new tail.
734 + hit_length -= hit_length % alignment_tokens
735 + num_blocks = cdiv(hit_length, block_size)
736 + for computed in computed_blocks:
737 + del computed[num_blocks:]
738 + return computed_blocks, hit_length
739 +
740 + def cache_blocks(
741 + self,
742 + request: Request,
743 + num_tokens: int,
744 + retention_interval: int | None = None,
745 + ) -> None:
746 + super().cache_blocks(request, num_tokens, retention_interval=retention_interval)
747 + hash_block_size = self.block_pool.hash_block_size
748 + if self.block_size == hash_block_size:
749 + return
750 + self._cache_partial_tail_block(request, num_tokens)
751 +
752 + def _cache_partial_tail_block(
753 + self,
754 + request: Request,
755 + num_tokens: int,
756 + ) -> None:
757 + """Cache the prompt tail when it ends inside a cache block.
758 +
759 + Only the final prompt hash boundary is registered as a partial
760 + prefix-cache entry; intermediate hash boundaries inside the same cache
761 + block are intentionally skipped.
762 + """
763 + hash_block_size = self.block_pool.hash_block_size
764 + boundary_tokens = request.num_prompt_tokens // hash_block_size * hash_block_size
765 + if boundary_tokens == 0 or boundary_tokens > num_tokens:
766 + return
767 + if boundary_tokens % self.block_size == 0:
768 + return
769 +
770 + blocks = self.req_to_blocks[request.request_id]
771 + block_idx = boundary_tokens // self.block_size
772 + if block_idx >= len(blocks):
773 + return
774 + self.block_pool.cache_partial_block(
775 + request=request,
776 + block=blocks[block_idx],
777 + num_tokens=boundary_tokens,
778 + kv_cache_group_id=self.kv_cache_group_id,
779 + block_size=self.block_size,
780 + )
781
782 def get_num_common_prefix_blocks(self, running_request_id: str) -> int:
783 blocks = self.req_to_blocks[running_request_id]
@@ -699,12 +866,23 @@ class SlidingWindowManager(SingleTypeKVCacheManager):
866 alignment_tokens: int,
867 dcp_world_size: int = 1,
868 pcp_world_size: int = 1,
702 - ) -> tuple[list[KVCacheBlock], ...]:
869 + ) -> tuple[tuple[list[KVCacheBlock], ...], int]:
870 assert isinstance(kv_cache_spec, SlidingWindowSpec), (
871 "SlidingWindowManager can only be used for sliding window groups"
872 )
873 assert dcp_world_size == 1, "DCP not support sliding window attn now."
874 assert pcp_world_size == 1, "PCP not support sliding window attn now."
875 + # Fine-grained partial hits are not supported for sliding window now
876 + assert alignment_tokens % kv_cache_spec.block_size == 0, (
877 + "SlidingWindowManager does not support fine-grained (partial) cache hits"
878 + )
879 + block_hashes = resolve_block_hashes(
880 + block_hashes,
881 + block_pool.hash_block_size,
882 + kv_cache_spec.block_size,
883 + supports_fine_grained_hash_lookup=cls.supports_fine_grained_hash_lookup,
884 + alignment_tokens=alignment_tokens,
885 + )
886
887 # The number of contiguous blocks needed for a prefix cache hit.
888 sliding_window_contiguous_blocks = cls._contiguous_blocks_for_hit(
@@ -717,7 +895,7 @@ class SlidingWindowManager(SingleTypeKVCacheManager):
895 # sliding_window_contiguous_blocks),
896 # which is good for low cache hit rate scenarios.
897 max_num_blocks = max_length // kv_cache_spec.block_size
720 - computed_blocks = tuple(
898 + computed_blocks: tuple[list[KVCacheBlock], ...] = tuple(
899 [block_pool.null_block] * max_num_blocks
900 for _ in range(len(kv_cache_group_ids))
901 )
@@ -772,7 +950,8 @@ class SlidingWindowManager(SingleTypeKVCacheManager):
950 ):
951 for computed in computed_blocks:
952 computed.pop()
775 - return computed_blocks
953 + hit_length = len(computed_blocks[0]) * block_size
954 + return computed_blocks, hit_length
955
956 @classmethod
957 def reachable_block_mask(
@@ -893,7 +1072,7 @@ class ChunkedLocalAttentionManager(SingleTypeKVCacheManager):
1072 alignment_tokens: int,
1073 dcp_world_size: int = 1,
1074 pcp_world_size: int = 1,
896 - ) -> tuple[list[KVCacheBlock], ...]:
1075 + ) -> tuple[tuple[list[KVCacheBlock], ...], int]:
1076 """
1077 For chunked local attention, we need to find the longest cache hit
1078 prefix of the blocks that is not longer than `max_length`. The prefix
@@ -942,6 +1121,13 @@ class ChunkedLocalAttentionManager(SingleTypeKVCacheManager):
1121 "KV cache groups with different block sizes are not compatible with "
1122 "chunked local attention now"
1123 )
1124 + block_hashes = resolve_block_hashes(
1125 + block_hashes,
1126 + block_pool.hash_block_size,
1127 + kv_cache_spec.block_size,
1128 + supports_fine_grained_hash_lookup=cls.supports_fine_grained_hash_lookup,
1129 + alignment_tokens=alignment_tokens,
1130 + )
1131 max_num_blocks = max_length // kv_cache_spec.block_size
1132 if max_length > 0:
1133 local_attention_start_idx = (
@@ -971,7 +1157,8 @@ class ChunkedLocalAttentionManager(SingleTypeKVCacheManager):
1157 computed.append(cached)
1158 else:
1159 break
974 - return computed_blocks
1160 + hit_length = len(computed_blocks[0]) * kv_cache_spec.block_size
1161 + return computed_blocks, hit_length
1162
1163 def get_num_skipped_tokens(self, num_computed_tokens: int) -> int:
1164 """
@@ -1027,6 +1214,8 @@ class ChunkedLocalAttentionManager(SingleTypeKVCacheManager):
1214
1215
1216 class MambaManager(SingleTypeKVCacheManager):
1217 + supports_fine_grained_hash_lookup: ClassVar[bool] = True
1218 +
1219 def __init__(
1220 self, kv_cache_spec: MambaSpec, block_pool: BlockPool, **kwargs
1221 ) -> None:
@@ -1035,9 +1224,9 @@ class MambaManager(SingleTypeKVCacheManager):
1224 # recurrent state. Undo the DCP/PCP block_size scaling that the base
1225 # class applies for attention groups whose KV cache is partitioned.
1226 self.block_size = kv_cache_spec.block_size
1038 - self.cached_blocks_this_step: set[BlockHashWithGroupId] = set()
1227 self.mamba_cache_mode = kv_cache_spec.mamba_cache_mode
1228 self.num_speculative_blocks: int = kv_cache_spec.num_speculative_blocks
1229 + self.cached_blocks_this_step: set[BlockHashWithGroupId] = set()
1230 if self.mamba_cache_mode == "align":
1231 # Mapping from request ID to the index of the block
1232 # allocated in the previous step
@@ -1057,17 +1246,46 @@ class MambaManager(SingleTypeKVCacheManager):
1246 alignment_tokens: int,
1247 dcp_world_size: int = 1,
1248 pcp_world_size: int = 1,
1060 - ) -> tuple[list[KVCacheBlock], ...]:
1249 + ) -> tuple[tuple[list[KVCacheBlock], ...], int]:
1250 assert isinstance(kv_cache_spec, MambaSpec), (
1251 "MambaManager can only be used for mamba groups"
1252 )
1253 assert dcp_world_size == 1, "DCP not support mamba now."
1254 assert pcp_world_size == 1, "PCP not support mamba now."
1255 + block_hashes = resolve_block_hashes(
1256 + block_hashes,
1257 + block_pool.hash_block_size,
1258 + kv_cache_spec.block_size,
1259 + supports_fine_grained_hash_lookup=cls.supports_fine_grained_hash_lookup,
1260 + alignment_tokens=alignment_tokens,
1261 + )
1262 computed_blocks: tuple[list[KVCacheBlock], ...] = tuple(
1263 [] for _ in range(len(kv_cache_group_ids))
1264 )
1265 + hit_length = 0
1266
1267 block_size = kv_cache_spec.block_size
1268 + if alignment_tokens < block_size and block_size % alignment_tokens == 0:
1269 + assert isinstance(block_hashes, list)
1270 + hash_block_size = alignment_tokens
1271 + scale_factor = block_size // hash_block_size
1272 + max_num_partial_units = min(
1273 + max_length // hash_block_size, len(block_hashes)
1274 + )
1275 + for fine_idx in range(max_num_partial_units - 1, -1, -1):
1276 + num_tokens = (fine_idx + 1) * hash_block_size
1277 + block_hash = block_hashes[fine_idx]
1278 + if cached_block := block_pool.get_cached_block(
1279 + block_hash, kv_cache_group_ids
1280 + ):
1281 + block_idx = fine_idx // scale_factor
1282 + for computed, cached in zip(computed_blocks, cached_block):
1283 + computed.extend([block_pool.null_block] * block_idx)
1284 + computed.append(cached)
1285 + hit_length = num_tokens
1286 + break
1287 + return computed_blocks, hit_length
1288 +
1289 max_num_blocks = max_length // block_size
1290 # Search from right to left and early stop when a match is found.
1291 for i in range(max_num_blocks - 1, -1, -1):
@@ -1089,9 +1307,10 @@ class MambaManager(SingleTypeKVCacheManager):
1307 # so we insert dummy blocks at the beginning:
1308 computed.extend([block_pool.null_block] * i)
1309 computed.append(cached)
1310 + hit_length = (i + 1) * block_size
1311 break # we just need the last match - early stopping
1312
1094 - return computed_blocks
1313 + return computed_blocks, hit_length
1314
1315 @classmethod
1316 def reachable_block_mask(
@@ -1189,6 +1408,7 @@ class MambaManager(SingleTypeKVCacheManager):
1408 num_tokens: int,
1409 new_computed_blocks: Sequence[KVCacheBlock],
1410 total_computed_tokens: int,
1411 + num_local_computed_tokens: int,
1412 num_tokens_main_model: int,
1413 apply_admission_cap: bool = False,
1414 ) -> int:
@@ -1214,6 +1434,7 @@ class MambaManager(SingleTypeKVCacheManager):
1434 num_tokens,
1435 new_computed_blocks,
1436 total_computed_tokens,
1437 + num_local_computed_tokens,
1438 num_tokens_main_model,
1439 apply_admission_cap=apply_admission_cap,
1440 )
@@ -1235,15 +1456,26 @@ class MambaManager(SingleTypeKVCacheManager):
1456 - len(new_computed_blocks)
1457 - len(self.req_to_blocks[request_id])
1458 )
1459 + has_partial_hit = (
1460 + self._has_partial_local_hit(
1461 + new_computed_blocks, num_local_computed_tokens
1462 + )
1463 + or request_id in self._partial_hit_reqs
1464 + )
1465 + if has_partial_hit:
1466 + num_new_blocks = max(num_new_blocks, 0) + 1
1467 if num_new_blocks > 0:
1468 if request_id in self._allocated_block_reqs:
1469 # Old request. Needs at most 1 more blocks as we can reuse the
1470 # speculative blocks in previous step.
1242 - num_new_blocks = 1
1471 + num_new_blocks = 1 + int(has_partial_hit)
1472 else:
1244 - # First prefill. Allocate 1 block for running state and the
1245 - # speculative blocks.
1246 - num_new_blocks = 1 + self.num_speculative_blocks
1473 + # First prefill. Allocate 1 block for running state, the
1474 + # speculative blocks, and one extra block if a partial cache
1475 + # hit must be copy-on-written before the new tokens run.
1476 + num_new_blocks = (
1477 + 1 + self.num_speculative_blocks + int(has_partial_hit)
1478 + )
1479
1480 num_evictable_computed_blocks = self._get_num_evictable_blocks(
1481 new_computed_blocks
@@ -1275,9 +1507,11 @@ class MambaManager(SingleTypeKVCacheManager):
1507 num_required_blocks = (
1508 cdiv(num_tokens, self.block_size) + self.num_speculative_blocks
1509 )
1510 + partial_hit = self._partial_hit_reqs.get(request_id)
1511 + has_partial_hit = partial_hit is not None
1512 # `num_required_blocks` might be less than `len(req_blocks)` if blocks are
1513 # over-allocated at last round.
1280 - if num_required_blocks <= len(req_blocks):
1514 + if num_required_blocks <= len(req_blocks) and not has_partial_hit:
1515 return []
1516 else:
1517 prev_block_len = len(req_blocks)
@@ -1317,14 +1551,42 @@ class MambaManager(SingleTypeKVCacheManager):
1551 else:
1552 break
1553 num_new_blocks = num_required_blocks - len(req_blocks)
1554 + if has_partial_hit:
1555 + num_new_blocks = max(num_new_blocks, 0) + 1
1556 if blocks_allocated:
1321 - assert num_new_blocks <= 1
1557 + assert num_new_blocks <= 1 + int(has_partial_hit)
1558 else:
1323 - assert num_new_blocks <= self.num_speculative_blocks + 1
1559 + assert num_new_blocks <= self.num_speculative_blocks + 1 + int(
1560 + has_partial_hit
1561 + )
1562 new_blocks = self.block_pool.get_new_blocks(num_new_blocks)
1563 + returned_blocks = req_blocks[prev_block_len:]
1564 + if partial_hit is not None:
1565 + block_idx, source_block = partial_hit
1566 + cow_block = new_blocks[0]
1567 + new_blocks = new_blocks[1:]
1568 + if blocks_allocated:
1569 + # The worker block table of a running request is
1570 + # append-only, so the request must stay on
1571 + # source_block. Move the cache entry to cow_block
1572 + # instead; the queued copy fills it before forward
1573 + # overwrites source_block.
1574 + assert req_blocks[block_idx] is source_block
1575 + self.block_pool.move_block_hashes(source_block, cow_block)
1576 + self._pending_cow_copies.append((source_block, cow_block))
1577 + source_block.ref_cnt += 1
1578 + if cow_block.block_hash is not None:
1579 + # The moved entry is only filled by this step's
1580 + # copy, so defer same-step hits on it.
1581 + self.cached_blocks_this_step.add(cow_block.block_hash)
1582 + else:
1583 + self._apply_cow(request_id, block_idx, source_block, cow_block)
1584 + returned_blocks = [cow_block] + returned_blocks
1585 req_blocks.extend(new_blocks)
1586 self._allocated_block_reqs.add(request_id)
1327 - return req_blocks[prev_block_len:]
1587 + self._partial_hit_reqs.pop(request_id, None)
1588 + returned_blocks.extend(new_blocks)
1589 + return returned_blocks
1590
1591 def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]:
1592 if self.mamba_cache_mode == "align":
@@ -1349,6 +1611,10 @@ class MambaManager(SingleTypeKVCacheManager):
1611 num_cached_blocks_before = self.num_cached_block.get(request.request_id, 0)
1612 super().cache_blocks(request, num_tokens, retention_interval=retention_interval)
1613 num_cached_blocks_after = self.num_cached_block.get(request.request_id, 0)
1614 + if self.mamba_cache_mode == "align":
1615 + partial_hash = self._cache_partial_tail_block(request, num_tokens)
1616 + if partial_hash is not None:
1617 + self.cached_blocks_this_step.add(partial_hash)
1618 if num_cached_blocks_after > num_cached_blocks_before:
1619 for block in self.req_to_blocks[request.request_id][
1620 num_cached_blocks_before:num_cached_blocks_after
@@ -1364,6 +1630,44 @@ class MambaManager(SingleTypeKVCacheManager):
1630 def new_step_starts(self) -> None:
1631 self.cached_blocks_this_step.clear()
1632
1633 + def _cache_partial_tail_block(
1634 + self,
1635 + request: Request,
1636 + num_tokens: int,
1637 + ) -> BlockHashWithGroupId | None:
1638 + hash_block_size = self.block_pool.hash_block_size
1639 + if self.block_size == hash_block_size:
1640 + return None
1641 + if num_tokens % self.block_size == 0:
1642 + return None
1643 + if num_tokens % hash_block_size != 0:
1644 + return None
1645 + latest_prompt_hash_boundary = (
1646 + request.num_prompt_tokens // hash_block_size
1647 + ) * hash_block_size
1648 + if num_tokens != latest_prompt_hash_boundary:
1649 + return None
1650 +
1651 + block_idx = num_tokens // self.block_size
1652 + blocks = self.req_to_blocks[request.request_id]
1653 + if block_idx >= len(blocks):
1654 + return None
1655 + source_block = blocks[block_idx]
1656 + if source_block.is_null:
1657 + return None
1658 +
1659 + partial_hash = self.block_pool.cache_partial_block(
1660 + request=request,
1661 + block=source_block,
1662 + num_tokens=num_tokens,
1663 + kv_cache_group_id=self.kv_cache_group_id,
1664 + block_size=self.block_size,
1665 + )
1666 + if partial_hash is not None:
1667 + self._partial_hit_reqs[request.request_id] = (block_idx, source_block)
1668 + self.num_cached_block[request.request_id] = block_idx
1669 + return partial_hash
1670 +
1671
1672 class CrossAttentionManager(SingleTypeKVCacheManager):
1673 """Manager for cross-attention KV cache in encoder-decoder models."""
@@ -1415,7 +1719,7 @@ class CrossAttentionManager(SingleTypeKVCacheManager):
1719 alignment_tokens: int,
1720 dcp_world_size: int = 1,
1721 pcp_world_size: int = 1,
1418 - ) -> tuple[list[KVCacheBlock], ...]:
1722 + ) -> tuple[tuple[list[KVCacheBlock], ...], int]:
1723 assert isinstance(kv_cache_spec, CrossAttentionSpec), (
1724 "CrossAttentionManager can only be used for cross-attention groups"
1725 )
@@ -1435,16 +1739,18 @@ class SinkFullAttentionManager(FullAttentionManager):
1739 block_pool: BlockPool,
1740 enable_caching: bool,
1741 kv_cache_group_id: int,
1742 + scheduler_block_size: int,
1743 dcp_world_size: int = 1,
1744 pcp_world_size: int = 1,
1745 ):
1746 super().__init__(
1442 - kv_cache_spec,
1443 - block_pool,
1444 - enable_caching,
1445 - kv_cache_group_id,
1446 - dcp_world_size,
1447 - pcp_world_size,
1747 + kv_cache_spec=kv_cache_spec,
1748 + block_pool=block_pool,
1749 + enable_caching=enable_caching,
1750 + kv_cache_group_id=kv_cache_group_id,
1751 + scheduler_block_size=scheduler_block_size,
1752 + dcp_world_size=dcp_world_size,
1753 + pcp_world_size=pcp_world_size,
1754 )
1755 sink_len = kv_cache_spec.sink_len
1756 assert sink_len is not None and sink_len > 0 and sink_len % self.block_size == 0
vllm/v1/worker/gpu/model_runner.py
+10 -1
@@ -112,7 +112,7 @@ from vllm.v1.worker.gpu.spec_decode.utils import DraftTokensHandler
112 from vllm.v1.worker.gpu.states import RequestState
113 from vllm.v1.worker.gpu.structured_outputs import StructuredOutputsWorker
114 from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin
115 -from vllm.v1.worker.utils import KVBlockZeroer
115 +from vllm.v1.worker.utils import KVBlockZeroer, copy_kv_cache_blocks_inplace
116
117 logger = init_logger(__name__)
118
@@ -838,6 +838,15 @@ class GPUModelRunner(LoRAModelRunnerMixin):
838 assert self.kv_block_zeroer is not None
839 self.kv_block_zeroer.zero_block_ids(scheduler_output.new_block_ids_to_zero)
840
841 + # Apply copy-on-write block copies for partial prefix-cache hits, after
842 + # zeroing new blocks and before the forward pass reads them.
843 + if scheduler_output.kv_cache_block_copies:
844 + copy_kv_cache_blocks_inplace(
845 + self.kv_caches,
846 + self.kv_cache_config.num_blocks,
847 + scheduler_output.kv_cache_block_copies,
848 + )
849 +
850 def prepare_inputs(
851 self, scheduler_output: SchedulerOutput, batch_desc: BatchExecutionDescriptor
852 ) -> InputBatch:
vllm/v1/worker/gpu_model_runner.py
+7
@@ -230,6 +230,7 @@ from .utils import (
230 KVBlockZeroer,
231 add_kv_sharing_layers_to_kv_cache_groups,
232 bind_kv_cache,
233 + copy_kv_cache_blocks_inplace,
234 prepare_kernel_block_sizes,
235 sanity_check_mm_encoder_outputs,
236 )
@@ -1188,6 +1189,12 @@ class GPUModelRunner(
1189 # stale NaN/data from corrupting attention or SSM computation.
1190 if scheduler_output.new_block_ids_to_zero:
1191 self._zero_block_ids(scheduler_output.new_block_ids_to_zero)
1192 + if scheduler_output.kv_cache_block_copies:
1193 + copy_kv_cache_blocks_inplace(
1194 + self.kv_caches,
1195 + self.kv_cache_config.num_blocks,
1196 + scheduler_output.kv_cache_block_copies,
1197 + )
1198
1199 # Free the cached encoder outputs.
1200 for mm_hash in scheduler_output.free_encoder_mm_hashes:
vllm/v1/worker/utils.py
+43 -1
@@ -2,11 +2,12 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import math
4 from collections import defaultdict
5 -from collections.abc import Iterable
5 +from collections.abc import Iterable, Sequence
6 from dataclasses import dataclass, field
7 from itertools import product as iprod
8 from typing import Any
9
10 +import numpy as np
11 import torch
12
13 from vllm.config import CacheConfig, VllmConfig
@@ -18,11 +19,13 @@ from vllm.platforms import current_platform
19 from vllm.triton_utils import tl, triton
20 from vllm.utils.math_utils import largest_power_of_2_divisor
21 from vllm.utils.mem_utils import MemorySnapshot, format_gib
22 +from vllm.utils.torch_utils import async_tensor_h2d
23 from vllm.v1.attention.backend import (
24 AttentionBackend,
25 AttentionMetadataBuilder,
26 MultipleOf,
27 )
28 +from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy
29 from vllm.v1.kv_cache_interface import (
30 AttentionSpec,
31 EncoderOnlyAttentionSpec,
@@ -535,6 +538,45 @@ def bind_kv_cache(
538 forward_context[layer_name].kv_cache = kv_cache
539
540
541 +def copy_kv_cache_blocks_inplace(
542 + kv_caches: Iterable[torch.Tensor | list[torch.Tensor]],
543 + num_blocks: int,
544 + kv_cache_block_copies: Sequence[KVCacheBlockCopy],
545 +) -> None:
546 + if not kv_cache_block_copies:
547 + return
548 +
549 + storage_tensors: list[torch.Tensor] = []
550 + seen_storage: set[int] = set()
551 + for entry in kv_caches:
552 + # Mamba layers hold a list of state tensors; attention layers a single
553 + # tensor. Both alias the shared block-major backing storage.
554 + tensors = entry if isinstance(entry, (list, tuple)) else (entry,)
555 + for tensor in tensors:
556 + ptr = tensor.untyped_storage().data_ptr()
557 + if ptr in seen_storage:
558 + continue
559 + seen_storage.add(ptr)
560 + storage_tensors.append(tensor)
561 +
562 + if not storage_tensors:
563 + return
564 + device = storage_tensors[0].device
565 + indices_np = np.array(kv_cache_block_copies, dtype=np.int64)
566 + indices = async_tensor_h2d(indices_np, device=device)
567 + src_indices, dst_indices = indices.unbind(dim=1)
568 +
569 + for tensor in storage_tensors:
570 + assert tensor.device == device
571 + blocks = torch.empty(0, dtype=torch.uint8, device=device)
572 + blocks.set_(tensor.untyped_storage())
573 + # Block-major backing storage: block i owns the contiguous byte range
574 + # [i * page_size, (i + 1) * page_size).
575 + assert blocks.numel() % num_blocks == 0
576 + blocks = blocks.view(num_blocks, -1)
577 + blocks[dst_indices] = blocks[src_indices]
578 +
579 +
580 def is_residual_scattered_for_sp(
581 vllm_config: VllmConfig, num_input_tokens: int
582 ) -> bool: