@tduncan / Nvidia-Optimization-Stack-VLLM / commits / d742856610

[3/N][Core][KV Connector] Support reliable partial-tail KV offload for sub-block prompts (#49502)

Signed-off-by: Dao Le <Dao007forever@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Dao007forever committed Jul 26, 2026 at 23:22 UTC d74285661064541f12be14e49a16defa27de67c6
13 files changed +1448 -49
tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py
+293
@@ -223,6 +223,299 @@ def test_hybrid_mamba_partial_tail_owner_uses_cow_on_continue():
223 assert moved[0].block_hash_num_tokens == 6
224
225
226 +def test_take_partial_tail_offloads_returns_cow_target():
227 + """The connector offload hand-off exposes the mamba CoW *target* block Y
228 + (the durable boundary state), not the overwritten source X, and only at
229 + the CoW step."""
230 + hash_block_size = 2
231 + block_size = 2 * hash_block_size
232 + kv_cache_config = KVCacheConfig(
233 + num_blocks=24,
234 + kv_cache_tensors=[],
235 + kv_cache_groups=[
236 + KVCacheGroupSpec(
237 + ["full"],
238 + FullAttentionSpec(
239 + block_size=hash_block_size,
240 + num_kv_heads=1,
241 + head_size=1,
242 + dtype=torch.float32,
243 + ),
244 + ),
245 + KVCacheGroupSpec(
246 + ["mamba"],
247 + MambaSpec(
248 + block_size=block_size,
249 + shapes=(1, 1),
250 + dtypes=(torch.float32,),
251 + mamba_cache_mode="align",
252 + ),
253 + ),
254 + ],
255 + )
256 + manager = make_kv_cache_manager(
257 + kv_cache_config=kv_cache_config,
258 + max_model_len=8192,
259 + enable_caching=True,
260 + hash_block_size=hash_block_size,
261 + )
262 +
263 + req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)
264 + computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)
265 + assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None
266 +
267 + # Step A registered the partial tail but has not CoW'd yet: no offload.
268 + assert manager.take_partial_tail_offloads() == {}
269 +
270 + partial_mamba_hash = req0.block_hashes[6 // hash_block_size - 1]
271 + source_block = manager.block_pool.get_cached_block(
272 + partial_mamba_hash, kv_cache_group_ids=[1]
273 + )
274 + assert source_block is not None
275 + source_block_id = source_block[0].block_id
276 +
277 + # Step B: the producer continues, triggering the CoW X->Y.
278 + req0.num_computed_tokens = 6
279 + req0.append_output_token_ids([3])
280 + assert manager.allocate_slots(req0, 1) is not None
281 +
282 + offloads = manager.take_partial_tail_offloads()
283 + assert list(offloads.keys()) == ["0"]
284 + assert len(offloads["0"]) == 1
285 + group_id, block_id, boundary_tokens = offloads["0"][0]
286 + assert group_id == 1 # the mamba group
287 + assert boundary_tokens == 6
288 + copies, _ = manager.take_kv_cache_block_copies()
289 + cow_copy = next(c for c in copies if c.src_block_id == source_block_id)
290 + # The offload points at the durable CoW target Y, not the overwritten X.
291 + assert block_id == cow_copy.dst_block_id
292 + assert block_id != source_block_id
293 + # Draining clears it.
294 + assert manager.take_partial_tail_offloads() == {}
295 +
296 + # The hand-off pinned Y (its CoW retention is released after this step,
297 + # and Y is off the request block table); freeing the request unpins it.
298 + cow_block = manager.block_pool.blocks[block_id]
299 + pinned_ref = cow_block.ref_cnt
300 + assert pinned_ref >= 1
301 + manager.free(req0)
302 + assert cow_block.ref_cnt == pinned_ref - 1
303 +
304 +
305 +def test_partial_tail_pin_survives_released_cow_retention():
306 + """If the CoW retention is released before the hand-off is drained
307 + (immediate-free mode), the drain must rescue the cow block from the free
308 + queue: a raw ref increment would leave a ref>0 block allocatable, and the
309 + next allocation would pop it and assert."""
310 + hash_block_size = 2
311 + block_size = 2 * hash_block_size
312 + kv_cache_config = KVCacheConfig(
313 + num_blocks=24,
314 + kv_cache_tensors=[],
315 + kv_cache_groups=[
316 + KVCacheGroupSpec(
317 + ["full"],
318 + FullAttentionSpec(
319 + block_size=hash_block_size,
320 + num_kv_heads=1,
321 + head_size=1,
322 + dtype=torch.float32,
323 + ),
324 + ),
325 + KVCacheGroupSpec(
326 + ["mamba"],
327 + MambaSpec(
328 + block_size=block_size,
329 + shapes=(1, 1),
330 + dtypes=(torch.float32,),
331 + mamba_cache_mode="align",
332 + ),
333 + ),
334 + ],
335 + )
336 + manager = make_kv_cache_manager(
337 + kv_cache_config=kv_cache_config,
338 + max_model_len=8192,
339 + enable_caching=True,
340 + hash_block_size=hash_block_size,
341 + )
342 + req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)
343 + computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)
344 + assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None
345 + req0.num_computed_tokens = 6
346 + req0.append_output_token_ids([3])
347 + assert manager.allocate_slots(req0, 1) is not None
348 +
349 + # Retention released before the drain (defer_block_free=False ordering).
350 + _copies, retained = manager.take_kv_cache_block_copies()
351 + manager.block_pool.free_blocks(retained)
352 +
353 + offloads = manager.take_partial_tail_offloads()
354 + ((_group_id, block_id, boundary_tokens),) = offloads["0"]
355 + assert boundary_tokens == 6
356 + cow_block = manager.block_pool.blocks[block_id]
357 + assert cow_block.ref_cnt == 1
358 +
359 + # The pinned block is out of the free queue: draining every free block
360 + # neither trips the allocator's ref_cnt assert nor hands it out.
361 + new_blocks = manager.block_pool.get_new_blocks(
362 + manager.block_pool.get_num_free_blocks()
363 + )
364 + assert block_id not in {b.block_id for b in new_blocks}
365 +
366 +
367 +def test_partial_tail_offload_dropped_when_request_freed_before_drain():
368 + """A hand-off recorded in the same scheduling pass as the request's death
369 + must not be drained: its release hook has already run, so draining would
370 + leak a pinned block."""
371 + hash_block_size = 2
372 + block_size = 2 * hash_block_size
373 + kv_cache_config = KVCacheConfig(
374 + num_blocks=24,
375 + kv_cache_tensors=[],
376 + kv_cache_groups=[
377 + KVCacheGroupSpec(
378 + ["full"],
379 + FullAttentionSpec(
380 + block_size=hash_block_size,
381 + num_kv_heads=1,
382 + head_size=1,
383 + dtype=torch.float32,
384 + ),
385 + ),
386 + KVCacheGroupSpec(
387 + ["mamba"],
388 + MambaSpec(
389 + block_size=block_size,
390 + shapes=(1, 1),
391 + dtypes=(torch.float32,),
392 + mamba_cache_mode="align",
393 + ),
394 + ),
395 + ],
396 + )
397 + manager = make_kv_cache_manager(
398 + kv_cache_config=kv_cache_config,
399 + max_model_len=8192,
400 + enable_caching=True,
401 + hash_block_size=hash_block_size,
402 + )
403 + req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)
404 + computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)
405 + assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None
406 + req0.num_computed_tokens = 6
407 + req0.append_output_token_ids([3])
408 + assert manager.allocate_slots(req0, 1) is not None
409 +
410 + # The request dies (preempt/abort) before the scheduler drains.
411 + manager.block_pool.free_blocks(manager.pop_blocks_for_free(req0))
412 + assert manager.take_partial_tail_offloads() == {}
413 +
414 +
415 +def test_take_partial_tail_offloads_empty_without_partial_tail():
416 + """A prompt ending on a block boundary registers no partial tail, so there
417 + is nothing to offload."""
418 + hash_block_size = 2
419 + block_size = 2 * hash_block_size
420 + kv_cache_config = KVCacheConfig(
421 + num_blocks=24,
422 + kv_cache_tensors=[],
423 + kv_cache_groups=[
424 + KVCacheGroupSpec(
425 + ["full"],
426 + FullAttentionSpec(
427 + block_size=hash_block_size,
428 + num_kv_heads=1,
429 + head_size=1,
430 + dtype=torch.float32,
431 + ),
432 + ),
433 + KVCacheGroupSpec(
434 + ["mamba"],
435 + MambaSpec(
436 + block_size=block_size,
437 + shapes=(1, 1),
438 + dtypes=(torch.float32,),
439 + mamba_cache_mode="align",
440 + ),
441 + ),
442 + ],
443 + )
444 + manager = make_kv_cache_manager(
445 + kv_cache_config=kv_cache_config,
446 + max_model_len=8192,
447 + enable_caching=True,
448 + hash_block_size=hash_block_size,
449 + )
450 +
451 + # 4-token prompt ends exactly on the mamba block boundary (block_size=4).
452 + req0 = make_request("0", [0, 0, 1, 1], hash_block_size, sha256)
453 + computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)
454 + assert manager.allocate_slots(req0, 4, num_computed, computed_blocks) is not None
455 + assert manager.take_partial_tail_offloads() == {}
456 +
457 + req0.num_computed_tokens = 4
458 + req0.append_output_token_ids([2])
459 + assert manager.allocate_slots(req0, 1) is not None
460 + assert manager.take_partial_tail_offloads() == {}
461 +
462 +
463 +def test_truncate_computed_blocks_preserves_sparse_prefix_positions():
464 + """truncate_computed_blocks slices each group by its own block size,
465 + keeps null placeholders in the retained prefix, and leaves the original
466 + lookup result untouched (pure view, no refcount changes)."""
467 + hash_block_size = 2
468 + kv_cache_config = KVCacheConfig(
469 + num_blocks=24,
470 + kv_cache_tensors=[],
471 + kv_cache_groups=[
472 + KVCacheGroupSpec(
473 + ["full"],
474 + FullAttentionSpec(
475 + block_size=hash_block_size,
476 + num_kv_heads=1,
477 + head_size=1,
478 + dtype=torch.float32,
479 + ),
480 + ),
481 + KVCacheGroupSpec(
482 + ["mamba"],
483 + MambaSpec(
484 + block_size=2 * hash_block_size,
485 + shapes=(1, 1),
486 + dtypes=(torch.float32,),
487 + mamba_cache_mode="align",
488 + ),
489 + ),
490 + ],
491 + )
492 + manager = make_kv_cache_manager(
493 + kv_cache_config=kv_cache_config,
494 + max_model_len=8192,
495 + enable_caching=True,
496 + hash_block_size=hash_block_size,
497 + )
498 + producer = make_request("producer", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)
499 + blocks, num_computed, _ = manager.get_computed_blocks(producer)
500 + assert manager.allocate_slots(producer, 6, num_computed, blocks) is not None
501 + manager.free(producer)
502 + manager.new_step_starts()
503 +
504 + consumer = make_request(
505 + "consumer", [0, 0, 1, 1, 2, 2, 3, 3], hash_block_size, sha256
506 + )
507 + blocks, num_computed, _ = manager.get_computed_blocks(consumer)
508 + assert num_computed == 6
509 + assert [len(group) for group in blocks.blocks] == [3, 2]
510 + assert blocks.blocks[1][0].is_null
511 +
512 + truncated = manager.truncate_computed_blocks(blocks, 4)
513 +
514 + assert [len(group) for group in truncated.blocks] == [2, 1]
515 + assert truncated.blocks[1][0].is_null
516 + assert [len(group) for group in blocks.blocks] == [3, 2]
517 +
518 +
519 def test_hybrid_mamba_partial_tail_owner_continue_preserves_later_hit():
520 hash_block_size = 2
521 block_size = 2 * hash_block_size
tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py
+55
@@ -3,6 +3,8 @@
3
4 from math import lcm
5
6 +import torch
7 +
8 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.coordinator import ( # noqa: E501
9 ExternalCachedBlockPool,
10 MooncakeStoreCoordinator,
@@ -14,10 +16,20 @@ from vllm.v1.core.kv_cache_utils import BlockHash
16 from vllm.v1.kv_cache_interface import (
17 FullAttentionSpec,
18 KVCacheGroupSpec,
19 + MambaSpec,
20 SlidingWindowSpec,
21 )
22
23
24 +def _mamba_align(block_size=32):
25 + return MambaSpec(
26 + block_size=block_size,
27 + shapes=((1, 1),),
28 + dtypes=(torch.float32,),
29 + mamba_cache_mode="align",
30 + )
31 +
32 +
33 def _make_coord(groups, hash_block_size, use_eagle=False, retention_interval=None):
34 """Construct a coordinator using the natural LCM of group block sizes as
35 the scheduler block size — mirrors ``resolve_kv_cache_block_sizes`` for
@@ -196,6 +208,49 @@ def test_coordinator_group_block_size_double_hash():
208 assert hit % 32 == 0
209
210
211 +# ----- Fine-grained partial hits (full attention + mamba "align") -----
212 +
213 +
214 +def test_coordinator_fine_grained_partial_tail_hit():
215 + """K3 shape: FA + mamba-align, block_size=32 over hash_block_size=16. When
216 + both groups have the sub-block boundary hash, the reconciled hit lands on
217 + the hash boundary (48), not the block boundary (32)."""
218 + groups = [
219 + KVCacheGroupSpec(["L0"], _full(32)),
220 + KVCacheGroupSpec(["L1"], _mamba_align(32)),
221 + ]
222 + coord = _make_coord(groups, hash_block_size=16)
223 + assert coord.enable_partial_hash_hits
224 + hs = _hashes(4) # 4 hash units of 16 = 64 tokens; block 0 = [0,32), etc.
225 + # Both groups: full block 0 (key = last sub-hash hs[1]) + partial boundary
226 + # at token 48 (key = hs[2]). No hs[3] -> block 1 is not full.
227 + exists = {(g, bytes(h)) for g in (0, 1) for h in (hs[1], hs[2])}
228 + cmap = ExternalCachedBlockPool(16, exists)
229 + _masks, hit = coord.find_longest_cache_hit(
230 + hs, max_length=64, cached_block_pool=cmap
231 + )
232 + assert hit == 48
233 +
234 +
235 +def test_coordinator_fine_grained_clips_when_one_group_missing_tail():
236 + """If only one group has the sub-block boundary, min-convergence clips the
237 + reconciled hit back to the block boundary (32)."""
238 + groups = [
239 + KVCacheGroupSpec(["L0"], _full(32)),
240 + KVCacheGroupSpec(["L1"], _mamba_align(32)),
241 + ]
242 + coord = _make_coord(groups, hash_block_size=16)
243 + hs = _hashes(4)
244 + # Full block 0 for both; partial boundary hs[2] only for FA (group 0).
245 + exists = {(g, bytes(hs[1])) for g in (0, 1)}
246 + exists |= {(0, bytes(hs[2]))}
247 + cmap = ExternalCachedBlockPool(16, exists)
248 + _masks, hit = coord.find_longest_cache_hit(
249 + hs, max_length=64, cached_block_pool=cmap
250 + )
251 + assert hit == 32
252 +
253 +
254 # ----- store_mask -----
255
256
tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py
+244 -3
@@ -32,6 +32,7 @@ from vllm.v1.kv_cache_interface import (
32 KVCacheConfig,
33 KVCacheGroupSpec,
34 KVCacheTensor,
35 + MambaSpec,
36 SlidingWindowSpec,
37 )
38
@@ -333,14 +334,14 @@ def test_recv_skips_swa_blocks_before_window():
334
335 def test_chunked_token_database_hash_block_size_smaller_than_block_size():
336 """DSv4-style: hash_block_size=4, group block_size=16 — process_tokens
336 - keys each 16-token chunk by its last fine hash, keeping the Mooncake key
337 - at one digest instead of concatenating all 4 fine hashes."""
337 + keys each chunk by its ending fine hash, including a partial tail."""
338 md = KeyMetadata("m", 0, 0, 0, 0, group_id=3)
339 db = ChunkedTokenDatabase(md, block_size=16, hash_block_size=4)
340 db.set_kv_caches_base_addr([0])
341 db.set_block_len([512])
342 - # 8 fine-grained hashes (32 tokens at hash_block_size=4) → 2 group chunks.
342 fine_hashes = [BlockHash(bytes([i + 1]) * 4) for i in range(8)]
343 +
344 + # 8 fine-grained hashes (32 tokens at hash_block_size=4) → 2 group chunks.
345 out = list(db.process_tokens(token_len=32, block_hashes=fine_hashes))
346 assert len(out) == 2
347 assert out[0][0] == 0 and out[0][1] == 16
@@ -349,3 +350,243 @@ def test_chunked_token_database_hash_block_size_smaller_than_block_size():
350 # prior three.
351 assert out[0][2].hex() == fine_hashes[3].hex()
352 assert out[1][2].hex() == fine_hashes[7].hex()
353 +
354 + # Sub-block hit: emit the partial chunk under its ending fine hash.
355 + out = list(db.process_tokens(token_len=12, block_hashes=fine_hashes[:3]))
356 + assert [(s, e) for s, e, _ in out] == [(0, 12)]
357 + assert out[0][2].hex() == fine_hashes[2].hex()
358 +
359 + # Cross-block hit: emit both the full chunk and its partial tail.
360 + out = list(db.process_tokens(token_len=28, block_hashes=fine_hashes[:7]))
361 + assert [(s, e) for s, e, _ in out] == [(0, 16), (16, 28)]
362 + assert out[0][2].hex() == fine_hashes[3].hex()
363 + assert out[1][2].hex() == fine_hashes[6].hex()
364 +
365 +
366 +def test_sub_block_partial_tail_offload_reads_cow_block():
367 + """Sub-block prompt (the 900/128/1536 shape, scaled to 12/4/16): the
368 + partial tail is offloaded for both groups under the boundary sub-hash. The
369 + full-attention block is read from the request block table; the mamba block
370 + is the core-provided CoW target, not block_ids."""
371 + full = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None)
372 + mamba = MambaSpec(
373 + block_size=16,
374 + shapes=((1, 1),),
375 + dtypes=(torch.float32,),
376 + mamba_cache_mode="align",
377 + )
378 + groups = [
379 + KVCacheGroupSpec(["L0"], full),
380 + KVCacheGroupSpec(["L1"], mamba),
381 + ]
382 + coord = MooncakeStoreCoordinator(groups, scheduler_block_size=16, hash_block_size=4)
383 + assert coord.enable_partial_hash_hits
384 +
385 + class _RecordingStore(_DictStore):
386 + def __init__(self):
387 + super().__init__()
388 + self.puts: dict[str, list[int]] = {}
389 +
390 + def batch_put_from_multi_buffers(self, keys, addrs, sizes, *a, **k):
391 + for key, addr in zip(keys, addrs):
392 + self.puts[key] = addr
393 + return super().batch_put_from_multi_buffers(keys, addrs, sizes, *a, **k)
394 +
395 + store = _RecordingStore()
396 + token_dbs = []
397 + for g_idx in range(2):
398 + db = ChunkedTokenDatabase(
399 + KeyMetadata("m", 0, 0, 0, 0, group_id=g_idx),
400 + block_size=16,
401 + hash_block_size=4,
402 + )
403 + db.set_kv_caches_base_addr([g_idx * 10_000])
404 + db.set_block_len([512])
405 + token_dbs.append(db)
406 +
407 + send = KVCacheStoreSendingThread(
408 + store=store,
409 + coord=coord,
410 + token_databases=token_dbs,
411 + block_size=16,
412 + tp_rank=0,
413 + group_put_steps=[1, 1],
414 + kv_role="kv_both",
415 + ready_event=threading.Event(),
416 + replicate_config=MagicMock(),
417 + )
418 +
419 + # The surrounding metadata may describe a longer resumed replay, but the
420 + # handoff identifies the exact state boundary to persist.
421 + hs = [BlockHash(bytes([i + 1]) * 4) for i in range(5)]
422 + mamba_cow_block = 7
423 + req = ReqMeta(
424 + req_id="r0",
425 + token_len_chunk=0,
426 + block_ids=([1], [2]),
427 + block_hashes=hs,
428 + can_save=True,
429 + num_prompt_tokens=20,
430 + partial_tail_offloads=[(1, mamba_cow_block, 12)],
431 + )
432 +
433 + send._maybe_offload_partial_tail(req)
434 +
435 + # boundary = 12 // 4 * 4 = 12 -> keyed by hs[12 // 4 - 1] = hs[2].
436 + partial_hash = hs[2]
437 + fa_key = token_dbs[0].key_for(partial_hash)
438 + mamba_key = token_dbs[1].key_for(partial_hash)
439 + assert set(store.puts) == {fa_key, mamba_key}
440 + # FA reads block_ids[0][0] = block 1: addr = base(0) + 1 * 512.
441 + assert store.puts[fa_key] == [512]
442 + # Mamba reads the CoW block 7, not block_ids[1][0]=2.
443 + assert store.puts[mamba_key] == [10_000 + mamba_cow_block * 512]
444 +
445 +
446 +def test_offload_syncs_event_before_put():
447 + """An offload-carrying meta synchronizes its CoW-fence event before the
448 + store put reads the blocks, then completes in one pass and drains the
449 + completion counter."""
450 + full = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None)
451 + mamba = MambaSpec(
452 + block_size=16,
453 + shapes=((1, 1),),
454 + dtypes=(torch.float32,),
455 + mamba_cache_mode="align",
456 + )
457 + groups = [
458 + KVCacheGroupSpec(["L0"], full),
459 + KVCacheGroupSpec(["L1"], mamba),
460 + ]
461 + coord = MooncakeStoreCoordinator(groups, scheduler_block_size=16, hash_block_size=4)
462 + event = MagicMock()
463 +
464 + class _FencedStore(_DictStore):
465 + def batch_put_from_multi_buffers(self, keys, addrs, sizes, *a, **k):
466 + assert event.synchronize.called, "put must run after the event sync"
467 + return super().batch_put_from_multi_buffers(keys, addrs, sizes, *a, **k)
468 +
469 + store = _FencedStore()
470 + token_dbs = []
471 + for g_idx in range(2):
472 + db = ChunkedTokenDatabase(
473 + KeyMetadata("m", 0, 0, 0, 0, group_id=g_idx),
474 + block_size=16,
475 + hash_block_size=4,
476 + )
477 + db.set_kv_caches_base_addr([g_idx * 10_000])
478 + db.set_block_len([512])
479 + token_dbs.append(db)
480 +
481 + send = KVCacheStoreSendingThread(
482 + store=store,
483 + coord=coord,
484 + token_databases=token_dbs,
485 + block_size=16,
486 + tp_rank=0,
487 + group_put_steps=[1, 1],
488 + kv_role="kv_both",
489 + ready_event=threading.Event(),
490 + replicate_config=MagicMock(),
491 + )
492 +
493 + hs = [BlockHash(bytes([i + 1]) * 4) for i in range(3)]
494 + req = ReqMeta(
495 + req_id="r1",
496 + token_len_chunk=0,
497 + block_ids=([1], [2]),
498 + block_hashes=hs,
499 + can_save=True,
500 + num_prompt_tokens=12,
501 + partial_tail_offloads=[(1, 7, 12)],
502 + )
503 + req.current_event = event
504 + send.add_stored_request("r1")
505 +
506 + send.request_queue.put(req)
507 + send._handle_request(send.request_queue.get())
508 + assert send.request_queue.qsize() == 0
509 + assert store._data
510 + assert send.stored_requests["r1"] == 0
511 + event.synchronize.assert_called_once()
512 +
513 +
514 +def test_sub_block_partial_tail_offload_covers_smaller_group_blocks():
515 + """The K3-shaped 900/128/1536 scenario scaled to 12/4/16, with a
516 + full-attention group whose block (4) is smaller than the lcm (16): the
517 + offload must persist every FA block up to the boundary — the normal save
518 + floors to the lcm, so those blocks are otherwise never written and the
519 + consumer's per-group lookup would miss. The mamba boundary block still
520 + reads the core-provided CoW target."""
521 + full = FullAttentionSpec(block_size=4, num_kv_heads=8, head_size=64, dtype=None)
522 + mamba = MambaSpec(
523 + block_size=16,
524 + shapes=((1, 1),),
525 + dtypes=(torch.float32,),
526 + mamba_cache_mode="align",
527 + )
528 + groups = [
529 + KVCacheGroupSpec(["L0"], full),
530 + KVCacheGroupSpec(["L1"], mamba),
531 + ]
532 + coord = MooncakeStoreCoordinator(groups, scheduler_block_size=16, hash_block_size=4)
533 + assert coord.enable_partial_hash_hits
534 +
535 + class _RecordingStore(_DictStore):
536 + def __init__(self):
537 + super().__init__()
538 + self.puts: dict[str, list[int]] = {}
539 +
540 + def batch_put_from_multi_buffers(self, keys, addrs, sizes, *a, **k):
541 + for key, addr in zip(keys, addrs):
542 + self.puts[key] = addr
543 + return super().batch_put_from_multi_buffers(keys, addrs, sizes, *a, **k)
544 +
545 + store = _RecordingStore()
546 + token_dbs = []
547 + for g_idx, block_size in enumerate([4, 16]):
548 + db = ChunkedTokenDatabase(
549 + KeyMetadata("m", 0, 0, 0, 0, group_id=g_idx),
550 + block_size=block_size,
551 + hash_block_size=4,
552 + )
553 + db.set_kv_caches_base_addr([g_idx * 10_000])
554 + db.set_block_len([512])
555 + token_dbs.append(db)
556 +
557 + send = KVCacheStoreSendingThread(
558 + store=store,
559 + coord=coord,
560 + token_databases=token_dbs,
561 + block_size=16,
562 + tp_rank=0,
563 + group_put_steps=[1, 1],
564 + kv_role="kv_both",
565 + ready_event=threading.Event(),
566 + replicate_config=MagicMock(),
567 + )
568 +
569 + hs = [BlockHash(bytes([i + 1]) * 4) for i in range(3)] # 3 hash units = 12 tok
570 + mamba_cow_block = 7
571 + req = ReqMeta(
572 + req_id="r2",
573 + token_len_chunk=0,
574 + block_ids=([1, 2, 3], [4]),
575 + block_hashes=hs,
576 + can_save=True,
577 + num_prompt_tokens=12,
578 + partial_tail_offloads=[(1, mamba_cow_block, 12)],
579 + )
580 +
581 + send._maybe_offload_partial_tail(req)
582 +
583 + # FA (block 4): full blocks ending at 4, 8 and 12, keyed by their normal
584 + # block-end hashes; mamba (block 16): the partial boundary block under
585 + # the boundary sub-hash, read from the CoW target.
586 + expected = {
587 + token_dbs[0].key_for(hs[0]): [1 * 512],
588 + token_dbs[0].key_for(hs[1]): [2 * 512],
589 + token_dbs[0].key_for(hs[2]): [3 * 512],
590 + token_dbs[1].key_for(hs[2]): [10_000 + mamba_cow_block * 512],
591 + }
592 + assert store.puts == expected
tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py
+231 -1
@@ -13,11 +13,15 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.scheduler impor
13 )
14
15
16 -def _make_bare_scheduler() -> MooncakeStoreScheduler:
16 +def _make_bare_scheduler(
17 + *, hash_block_size: int = 16, enable_partial_hash_hits: bool = False
18 +) -> MooncakeStoreScheduler:
19 scheduler = object.__new__(MooncakeStoreScheduler)
20 scheduler.kv_role = "kv_both"
21 scheduler.lookup_async = False
22 scheduler._block_size = 16
23 + scheduler._hash_block_size = hash_block_size
24 + scheduler.enable_partial_hash_hits = enable_partial_hash_hits
25 scheduler.load_specs = {}
26 scheduler._unfinished_request_ids = {"req-0"}
27 scheduler._unfinished_requests = {}
@@ -136,6 +140,7 @@ def test_preemption_resets_tracker_before_request_finished():
140 block_hashes=[b"h0", b"h1"],
141 prefill_end_tokens=48,
142 )
143 + scheduler._request_trackers["req-0"].has_pending_offload = True
144
145 scheduler.build_connector_meta(_make_preemption_scheduler_output())
146
@@ -144,6 +149,7 @@ def test_preemption_resets_tracker_before_request_finished():
149 assert tracker.allocated_block_ids == ()
150 assert tracker.num_saved_tokens == 0
151 assert tracker.token_ids is None
152 + assert tracker.has_pending_offload is False
153 assert tracker.prefill_end_tokens == 0
154 request = SimpleNamespace(request_id="req-0")
155 assert scheduler.request_finished(request, ([0, 1],)) == (False, None)
@@ -536,3 +542,227 @@ def test_full_external_hit_with_full_local_hit_skips_load():
542 assert need_to_allocate == 0
543 assert load_async is False
544 assert "req-0" not in scheduler.load_specs
545 +
546 +
547 +def test_partial_hash_hit_block_aligned_local_loads_partial_tail():
548 + # Fine-grained on (hash=4, block=16): a block-aligned local hit can pull a
549 + # sub-block remote hit (24 = a hash boundary inside block 1). Loads [16, 24).
550 + scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True)
551 + scheduler.load_async = True
552 + scheduler.client = _StubLookupClient(hit_tokens=24)
553 +
554 + request = SimpleNamespace(
555 + request_id="req-0",
556 + num_tokens=32,
557 + block_hashes=[b"h0", b"h1", b"h2", b"h3", b"h4", b"h5", b"h6", b"h7"],
558 + )
559 +
560 + need_to_allocate, load_async = scheduler.get_num_new_matched_tokens(
561 + request, num_computed_tokens=16
562 + )
563 +
564 + assert need_to_allocate == 8
565 + assert load_async is True
566 + load_spec = scheduler.load_specs["req-0"]
567 + assert load_spec.vllm_cached_tokens == 16
568 + assert load_spec.kvpool_cached_tokens == 24
569 +
570 +
571 +def test_partial_hash_hit_no_remote_gain_skips_load():
572 + # Core always presents a block-aligned local hit (it floors a sub-block
573 + # tail before calling the connector). When the remote hit does not exceed
574 + # that block-aligned local hit, nothing is loaded.
575 + scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True)
576 + scheduler.load_async = True
577 + scheduler.client = _StubLookupClient(hit_tokens=16)
578 +
579 + request = SimpleNamespace(
580 + request_id="req-0",
581 + num_tokens=32,
582 + block_hashes=[b"h0", b"h1", b"h2", b"h3", b"h4", b"h5", b"h6", b"h7"],
583 + )
584 +
585 + need_to_allocate, load_async = scheduler.get_num_new_matched_tokens(
586 + request, num_computed_tokens=16
587 + )
588 +
589 + assert need_to_allocate == 0
590 + assert load_async is False
591 + assert "req-0" not in scheduler.load_specs
592 +
593 +
594 +def test_sub_block_prompt_looks_up_with_fine_grained():
595 + # A prompt smaller than one block (12 < block 16). With fine-grained partial
596 + # hits the sub-block prefix is worth looking up (floor is the hash unit 4,
597 + # not a full block), so a remote partial hit is loaded. Pre-change the
598 + # block-size floor returned (0, False) for such prompts.
599 + scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True)
600 + scheduler.load_async = True
601 + scheduler.client = _StubLookupClient(hit_tokens=8)
602 +
603 + request = SimpleNamespace(
604 + request_id="req-0",
605 + num_tokens=12,
606 + block_hashes=[b"h0", b"h1", b"h2"],
607 + )
608 +
609 + need_to_allocate, load_async = scheduler.get_num_new_matched_tokens(
610 + request, num_computed_tokens=0
611 + )
612 +
613 + assert need_to_allocate == 8
614 + assert load_async is True
615 + assert scheduler.load_specs["req-0"].kvpool_cached_tokens == 8
616 +
617 +
618 +def test_sub_block_prompt_not_looked_up_without_fine_grained():
619 + # Without fine-grained partial hits, sub-block prompts still skip the lookup
620 + # (there is no full block, and no sub-block key granularity).
621 + scheduler = _make_bare_scheduler()
622 + scheduler.client = _StubLookupClient(hit_tokens=8)
623 +
624 + request = SimpleNamespace(
625 + request_id="req-0",
626 + num_tokens=12,
627 + block_hashes=[b"h0", b"h1", b"h2"],
628 + )
629 +
630 + need_to_allocate, load_async = scheduler.get_num_new_matched_tokens(
631 + request, num_computed_tokens=0
632 + )
633 +
634 + assert need_to_allocate == 0
635 + assert load_async is False
636 + assert "req-0" not in scheduler.load_specs
637 +
638 +
639 +def test_pending_partial_tail_emits_offload_only_reqmeta():
640 + # A sub-block prompt never produces a block-aligned save, so the partial-
641 + # tail offload arriving this step is emitted as an offload-only ReqMeta
642 + # (can_save=True so it takes the normal enqueue path, token_len_chunk=0 so
643 + # the worker skips the normal save). Pending-offload state delays the free
644 + # without advancing the normal-save watermark before the put succeeds.
645 + scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True)
646 + request = SimpleNamespace(
647 + all_token_ids=list(range(12)),
648 + block_hashes=[b"h0", b"h1", b"h2"],
649 + num_output_placeholders=0,
650 + num_prompt_tokens=12,
651 + )
652 + scheduler._unfinished_requests["req-0"] = (request, ([0],))
653 + scheduler._request_trackers["req-0"] = RequestTracker(
654 + req_id="req-0",
655 + token_len=12,
656 + allocated_block_ids=([0],),
657 + num_saved_tokens=0,
658 + token_ids=list(range(12)),
659 + prefill_end_tokens=12,
660 + )
661 +
662 + out = SimpleNamespace(
663 + finished_req_ids=set(),
664 + preempted_req_ids=set(),
665 + scheduled_new_reqs=[],
666 + scheduled_cached_reqs=SimpleNamespace(
667 + req_ids=[],
668 + new_block_ids=[],
669 + num_computed_tokens=[],
670 + resumed_req_ids=set(),
671 + ),
672 + num_scheduled_tokens={},
673 + scheduled_spec_decode_tokens={},
674 + partial_tail_offloads={"req-0": [(1, 7, 12)]},
675 + )
676 +
677 + meta = scheduler.build_connector_meta(out)
678 +
679 + assert len(meta.requests) == 1
680 + req_meta = meta.requests[0]
681 + assert req_meta.req_id == "req-0"
682 + assert req_meta.can_save is True
683 + assert req_meta.token_len_chunk == 0
684 + assert req_meta.partial_tail_offloads == [(1, 7, 12)]
685 + assert req_meta.num_prompt_tokens == 12
686 + assert req_meta.block_ids == ([0],)
687 + tracker = scheduler._request_trackers["req-0"]
688 + assert tracker.num_saved_tokens == 0
689 + assert tracker.has_pending_offload is True
690 + request = SimpleNamespace(request_id="req-0")
691 + assert scheduler.request_finished(request, ([0],)) == (True, None)
692 +
693 +
694 +def test_resumed_partial_tail_uses_handoff_boundary():
695 + scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True)
696 + request = SimpleNamespace(
697 + all_token_ids=list(range(20)),
698 + block_hashes=[b"h0", b"h1", b"h2", b"h3", b"h4"],
699 + num_output_placeholders=0,
700 + num_prompt_tokens=12,
701 + )
702 + scheduler._unfinished_requests["req-0"] = (request, ([0, 1],))
703 + scheduler._request_trackers["req-0"] = RequestTracker(
704 + req_id="req-0",
705 + token_len=20,
706 + allocated_block_ids=([0, 1],),
707 + num_saved_tokens=0,
708 + token_ids=list(range(20)),
709 + # Resumption replays prompt + previously generated tokens.
710 + prefill_end_tokens=20,
711 + )
712 +
713 + out = SimpleNamespace(
714 + finished_req_ids=set(),
715 + preempted_req_ids=set(),
716 + scheduled_new_reqs=[],
717 + scheduled_cached_reqs=SimpleNamespace(
718 + req_ids=[],
719 + new_block_ids=[],
720 + num_computed_tokens=[],
721 + resumed_req_ids=set(),
722 + ),
723 + num_scheduled_tokens={},
724 + scheduled_spec_decode_tokens={},
725 + partial_tail_offloads={"req-0": [(1, 7, 12)]},
726 + )
727 +
728 + meta = scheduler.build_connector_meta(out)
729 +
730 + assert len(meta.requests) == 1
731 + assert meta.requests[0].partial_tail_offloads == [(1, 7, 12)]
732 + # Ordinary metadata retains the full resumed prefill range.
733 + assert meta.requests[0].num_prompt_tokens == 20
734 + tracker = scheduler._request_trackers["req-0"]
735 + assert tracker.num_saved_tokens == 0
736 + assert tracker.has_pending_offload is True
737 +
738 +
739 +def test_resumed_partial_tail_attached_to_save_keeps_handoff_boundary():
740 + scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True)
741 + request = SimpleNamespace(
742 + all_token_ids=list(range(48)),
743 + block_hashes=[b"h0", b"h1", b"h2"],
744 + num_output_placeholders=0,
745 + num_prompt_tokens=36,
746 + )
747 + scheduler._unfinished_requests["req-0"] = (request, ([0, 1],))
748 + scheduler._request_trackers["req-0"] = RequestTracker(
749 + req_id="req-0",
750 + token_len=44,
751 + allocated_block_ids=([0, 1],),
752 + num_saved_tokens=32,
753 + token_ids=list(range(44)),
754 + prefill_end_tokens=48,
755 + )
756 + out = _make_scheduler_output(scheduled_spec_tokens=None)
757 + out.partial_tail_offloads = {"req-0": [(0, 7, 36)]}
758 +
759 + meta = scheduler.build_connector_meta(out)
760 +
761 + assert len(meta.requests) == 1
762 + assert meta.requests[0].can_save is True
763 + assert meta.requests[0].partial_tail_offloads == [(0, 7, 36)]
764 + assert meta.requests[0].num_prompt_tokens == 48
765 + # Ordinary saving still covers the full resumed prefill range.
766 + tracker = scheduler._request_trackers["req-0"]
767 + assert tracker.num_saved_tokens == 48
768 + assert tracker.has_pending_offload is True
tests/v1/kv_connector/unit/test_mooncake_store_worker.py
+122
@@ -599,6 +599,88 @@ def test_store_sending_thread_retries_skipped_range_after_pressure():
599 assert store.batch_put_from_multi_buffers.call_args.args[0] == keys
600
601
602 +def _make_partial_tail_send_thread(store):
603 + coord = SimpleNamespace(
604 + enable_partial_hash_hits=True,
605 + hash_block_size=4,
606 + lcm_block_size=16,
607 + )
608 + db = ChunkedTokenDatabase(
609 + KeyMetadata("test-model", 0, 0, 0, 0),
610 + block_size=4,
611 + hash_block_size=4,
612 + )
613 + db.set_kv_caches_base_addr([0x1000])
614 + db.set_block_len([256])
615 + return _make_store_sending_thread(
616 + store,
617 + coord=coord,
618 + token_databases=[db],
619 + )
620 +
621 +
622 +def _make_partial_tail_req(block_ids: list[int]) -> ReqMeta:
623 + return ReqMeta(
624 + req_id="req-a",
625 + token_len_chunk=0,
626 + block_ids=(block_ids,),
627 + block_hashes=[b"a0", b"a1", b"a2"],
628 + can_save=True,
629 + partial_tail_offloads=[(1, 7, 12)],
630 + )
631 +
632 +
633 +def test_partial_tail_offload_skips_null_source_blocks():
634 + store = MagicMock()
635 + store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
636 + store.batch_put_from_multi_buffers.return_value = [256, 256]
637 + thread = _make_partial_tail_send_thread(store)
638 +
639 + assert thread._maybe_offload_partial_tail(_make_partial_tail_req([0, 2, 3]))
640 +
641 + keys, addrs, _sizes, _replicate_config = (
642 + store.batch_put_from_multi_buffers.call_args.args
643 + )
644 + assert keys == [
645 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6131",
646 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6132",
647 + ]
648 + assert addrs == [[0x1000 + 2 * 256], [0x1000 + 3 * 256]]
649 +
650 +
651 +def test_partial_tail_offload_honors_active_pressure_gate():
652 + store = MagicMock()
653 + thread = _make_partial_tail_send_thread(store)
654 + thread._store_pressure_active = True
655 + thread._skip_store_requests.add("req-a")
656 + thread.add_stored_request("req-a")
657 +
658 + thread._handle_request(_make_partial_tail_req([1, 2, 3]))
659 +
660 + store.batch_is_exist.assert_not_called()
661 + store.batch_put_from_multi_buffers.assert_not_called()
662 + assert thread.stored_requests["req-a"] == 0
663 +
664 +
665 +def test_partial_tail_put_failure_activates_pressure_gate():
666 + store = MagicMock()
667 + store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
668 + store.batch_put_from_multi_buffers.return_value = [256, -200, 256]
669 + thread = _make_partial_tail_send_thread(store)
670 + thread.add_stored_request("req-a")
671 +
672 + thread._handle_request(_make_partial_tail_req([1, 2, 3]))
673 +
674 + assert thread._store_pressure_active is True
675 + assert thread._skip_store_requests == {"req-a"}
676 + assert thread._saved_offset.get("req-a", 0) == 0
677 + assert thread.stored_requests["req-a"] == 0
678 +
679 + thread.add_stored_request("req-a")
680 + thread._handle_request(_make_partial_tail_req([1, 2, 3]))
681 + assert store.batch_put_from_multi_buffers.call_count == 1
682 +
683 +
684 def test_store_sending_thread_delta_start_rank_saves_second_local_chunk():
685 store = MagicMock()
686 store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
@@ -1939,6 +2021,46 @@ def test_lookup_partial_prefix_returns_first_hit_length():
2021 assert worker.lookup(48, [b"a0", b"a1", b"a2"]) == 32
2022
2023
2024 +def test_lookup_partial_tail_uses_hash_alignment():
2025 + """A stored sub-block tail can serve a request extending past it."""
2026 + from vllm.v1.kv_cache_interface import (
2027 + FullAttentionSpec,
2028 + KVCacheGroupSpec,
2029 + MambaSpec,
2030 + )
2031 +
2032 + worker = _make_bare_worker(block_size=16)
2033 + full = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None)
2034 + mamba = MambaSpec(
2035 + block_size=16,
2036 + shapes=((1, 1),),
2037 + dtypes=(torch.float32,),
2038 + mamba_cache_mode="align",
2039 + )
2040 + worker._kv_cache_groups = [
2041 + KVCacheGroupSpec(["full"], full),
2042 + KVCacheGroupSpec(["mamba"], mamba),
2043 + ]
2044 + worker.hash_block_size = 4
2045 + worker.token_dbs = [
2046 + ChunkedTokenDatabase(
2047 + KeyMetadata("test-model", 0, 0, 0, 0, group_id=group_id),
2048 + block_size=16,
2049 + hash_block_size=4,
2050 + )
2051 + for group_id in range(2)
2052 + ]
2053 + worker.coord = mooncake_store_worker.MooncakeStoreCoordinator(
2054 + worker._kv_cache_groups,
2055 + scheduler_block_size=16,
2056 + hash_block_size=4,
2057 + )
2058 + _refresh_group_tp_replication_factors(worker)
2059 + worker.store.batch_is_exist.return_value = [0, 0, 1, 0, 0, 1]
2060 +
2061 + assert worker.lookup(13, [b"h0", b"h1", b"h2"]) == 12
2062 +
2063 +
2064 def test_lookup_full_hit_reuses_existing_boundary():
2065 """A full hit is re-derived below the request end without another RPC."""
2066 worker = _make_bare_worker(block_size=16)
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py
+53 -13
@@ -21,6 +21,7 @@ from vllm.v1.kv_cache_interface import (
21 FullAttentionSpec,
22 KVCacheGroupSpec,
23 KVCacheSpec,
24 + MambaSpec,
25 UniformTypeKVCacheSpecs,
26 )
27 from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
@@ -83,6 +84,9 @@ class MooncakeStoreCoordinator:
84 self.kv_cache_groups = kv_cache_groups
85 self.hash_block_size = hash_block_size
86 self.lcm_block_size = scheduler_block_size
87 + self.enable_partial_hash_hits = partial_hash_hits_enabled(
88 + kv_cache_groups, hash_block_size
89 + )
90 self.use_eagle = use_eagle
91 # Mirror vLLM core's KVCacheCoordinator.retention_interval.
92 self.retention_interval = retention_interval
@@ -94,7 +98,12 @@ class MooncakeStoreCoordinator:
98 self._verify_and_split_kv_cache_groups()
99
100 def align_lookup_length(self, length: int) -> int:
97 - return length // self.lcm_block_size * self.lcm_block_size
101 + alignment = (
102 + self.hash_block_size
103 + if self.enable_partial_hash_hits
104 + else self.lcm_block_size
105 + )
106 + return length // alignment * alignment
107
108 def _verify_and_split_kv_cache_groups(self) -> None:
109 """Mirrors KVCacheCoordinator.verify_and_split_kv_cache_groups but
@@ -226,9 +235,14 @@ class MooncakeStoreCoordinator:
235 retention_interval: int | None,
236 num_prompt_tokens: int | None,
237 ) -> tuple[list[bool] | None, ...]:
229 - assert aligned_token_len % self.lcm_block_size == 0, (
238 + mask_alignment = (
239 + self.hash_block_size
240 + if self.enable_partial_hash_hits
241 + else self.lcm_block_size
242 + )
243 + assert aligned_token_len % mask_alignment == 0, (
244 f"aligned_token_len ({aligned_token_len}) must be a multiple of "
231 - f"lcm_block_size ({self.lcm_block_size})"
245 + f"{mask_alignment}"
246 )
247 masks: list[list[bool] | None] = []
248 for g_idx, g in enumerate(self.kv_cache_groups):
@@ -278,17 +292,21 @@ class MooncakeStoreCoordinator:
292 one already removed by the lookup.
293 """
294 eagle_indices = self.eagle_attn_group_indices if apply_eagle else set()
295 + alignment_tokens = (
296 + self.hash_block_size
297 + if self.enable_partial_hash_hits
298 + else self.lcm_block_size
299 + )
300 if len(self.attention_groups) == 1:
301 spec, group_ids, manager_cls = self.attention_groups[0]
283 - hashes = self.block_hashes_for_spec(block_hashes, spec)
302 hit_blocks, hit_length = manager_cls.find_longest_cache_hit(
285 - block_hashes=hashes, # type: ignore[arg-type]
303 + block_hashes=block_hashes, # type: ignore[arg-type]
304 max_length=max_length,
305 kv_cache_group_ids=group_ids,
306 block_pool=cast(BlockPool, cached_block_pool),
307 kv_cache_spec=spec,
308 drop_eagle_block=(0 in eagle_indices),
291 - alignment_tokens=spec.block_size,
309 + alignment_tokens=alignment_tokens,
310 )
311 num_groups = len(self.kv_cache_groups)
312 blocks_by_group: list[list[KVCacheBlock]] = [[] for _ in range(num_groups)]
@@ -320,17 +338,23 @@ class MooncakeStoreCoordinator:
338
339 drop_eagle_block = idx in eagle_indices and idx not in eagle_verified
340 _max_length = curr_hit_length
323 - if drop_eagle_block:
324 - _max_length = min(curr_hit_length + spec.block_size, max_length)
325 - hashes = self.block_hashes_for_spec(block_hashes, spec)
341 + if drop_eagle_block and not isinstance(spec, MambaSpec):
342 + eagle_margin = (
343 + self.hash_block_size
344 + if self.enable_partial_hash_hits
345 + and manager_cls.supports_fine_grained_hash_lookup
346 + and spec.block_size > self.hash_block_size
347 + else spec.block_size
348 + )
349 + _max_length = min(curr_hit_length + eagle_margin, max_length)
350 hit_blocks, _new_hit_length = manager_cls.find_longest_cache_hit(
327 - block_hashes=hashes, # type: ignore[arg-type]
351 + block_hashes=block_hashes, # type: ignore[arg-type]
352 max_length=_max_length,
353 kv_cache_group_ids=group_ids,
354 block_pool=cast(BlockPool, cached_block_pool),
355 kv_cache_spec=spec,
356 drop_eagle_block=drop_eagle_block,
333 - alignment_tokens=self.lcm_block_size,
357 + alignment_tokens=alignment_tokens,
358 )
359 if drop_eagle_block:
360 eagle_verified.add(idx)
@@ -348,10 +372,11 @@ class MooncakeStoreCoordinator:
372 break
373
374 # Truncate full-attention hit_blocks to final converged length;
351 - # other specs already trim themselves inside their hit logic.
375 + # other specs already trim themselves inside their hit logic. cdiv keeps
376 + # the partial tail block when hit_length is not block-aligned.
377 spec0, group_ids0, _ = self.attention_groups[0]
378 if isinstance(spec0, FullAttentionSpec):
354 - num_blocks = hit_length // spec0.block_size
379 + num_blocks = cdiv(hit_length, spec0.block_size)
380 for gid in group_ids0:
381 full_blks = hit_blocks_by_group[gid]
382 assert full_blks is not None
@@ -368,3 +393,18 @@ def _unwrap_spec(spec: KVCacheSpec) -> KVCacheSpec:
393 if isinstance(spec, UniformTypeKVCacheSpecs):
394 return next(iter(spec.kv_cache_specs.values()))
395 return spec
396 +
397 +
398 +def partial_hash_hits_enabled(
399 + kv_cache_groups: list[KVCacheGroupSpec], hash_block_size: int
400 +) -> bool:
401 + """Mirror of core's ``HybridKVCacheCoordinator.enable_partial_hash_hits``
402 + (its dcp == 1 clause holds: the connector rejects hybrid + DCP/PCP > 1).
403 + Single copy on purpose — scheduler and coordinator must not disagree.
404 + """
405 + return any(
406 + isinstance(spec := _unwrap_spec(g.kv_cache_spec), MambaSpec)
407 + and spec.mamba_cache_mode == "align"
408 + and spec.block_size > hash_block_size
409 + for g in kv_cache_groups
410 + )
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py
+27 -8
@@ -228,16 +228,28 @@ class ChunkedTokenDatabase:
228 n = len(chunks)
229 starts = np.fromiter((c[0] for c in chunks), dtype=np.int64, count=n)
230 spans = np.fromiter((c[1] for c in chunks), dtype=np.int64, count=n) - starts
231 - assert not (spans % self.block_size).any()
231 + assert not (spans % self.hash_block_size).any()
232 bids = np.fromiter(
233 (block_ids[i] for i in (starts // self.block_size).tolist()),
234 dtype=np.int64,
235 count=n,
236 )
237 addrs = base[None, :] + bids[:, None] * blen[None, :]
238 - sizes = blen[None, :] * (spans // self.block_size)[:, None]
238 + block_counts = (spans + self.block_size - 1) // self.block_size
239 + sizes = blen[None, :] * block_counts[:, None]
240 return addrs.tolist(), sizes.tolist(), bids.tolist()
241
242 + def prepare_value_for_block(self, block_id: int) -> tuple[list[int], list[int]]:
243 + """Return addresses and sizes for one physical block slot."""
244 + addr_list = []
245 + size_list = []
246 + length = len(self.block_len)
247 + for index, base_addr in enumerate(self.kv_caches_base_addr):
248 + addr = base_addr + block_id * self.block_len[index % length]
249 + addr_list.append(addr)
250 + size_list.append(self.block_len[index % length])
251 + return addr_list, size_list
252 +
253 def process_tokens(
254 self,
255 token_len: int,
@@ -256,7 +268,8 @@ class ChunkedTokenDatabase:
268 rank regardless of where the processed suffix begins.
269
270 Args:
259 - token_len: Total number of tokens.
271 + token_len: Total number of tokens. Must be hash-block aligned and
272 + covered by ``block_hashes`` when hashes are present.
273 block_hashes: Block hashes computed at ``hash_block_size`` granularity.
274 When ``block_size > hash_block_size`` each group's ``block_size`` chunk
275 is keyed by its last sub-hash via ``chunk_hashes_for_block_size``.
@@ -269,11 +282,10 @@ class ChunkedTokenDatabase:
282 assert put_step > 0
283 if not block_hashes:
284 return
272 - chunk_hashes: Sequence[BlockHash] = chunk_hashes_for_block_size(
273 - block_hashes, self.hash_block_size, self.block_size
274 - )
285 + assert token_len % self.hash_block_size == 0
286 + assert token_len // self.hash_block_size <= len(block_hashes)
287 start_chunk = max(0, cdiv(mask_num, self.block_size))
276 - max_chunks = min(len(chunk_hashes), cdiv(token_len, self.block_size))
288 + max_chunks = cdiv(token_len, self.block_size)
289 if chunk_mask is not None:
290 max_chunks = min(max_chunks, start_chunk + len(chunk_mask))
291 for chunk_id in range(start_chunk, max_chunks):
@@ -281,9 +293,9 @@ class ChunkedTokenDatabase:
293 continue
294 if chunk_id % put_step != put_step_rank:
295 continue
284 - h = chunk_hashes[chunk_id]
296 start_idx = chunk_id * self.block_size
297 end_idx = min(start_idx + self.block_size, token_len)
298 + h = block_hashes[end_idx // self.hash_block_size - 1]
299 yield start_idx, end_idx, h
300
301
@@ -306,6 +318,7 @@ class RequestTracker:
318 allocated_block_ids: tuple[list[int], ...]
319 num_saved_tokens: int = 0
320 token_ids: list[int] | None = None
321 + has_pending_offload: bool = False
322 # Snapshot of the prefill range length at tracker creation time.
323 # For a fresh request this is len(prompt). For a resumed-from-preemption
324 # request it includes previously-generated tokens, which are re-prefilled.
@@ -316,6 +329,7 @@ class RequestTracker:
329 self.allocated_block_ids = ()
330 self.num_saved_tokens = 0
331 self.token_ids = None
332 + self.has_pending_offload = False
333 self.prefill_end_tokens = 0
334
335 def update(
@@ -352,6 +366,11 @@ class ReqMeta:
366
367 token_ids: list[int] | None = None
368 num_prompt_tokens: int | None = None
369 + # Core-provided per-mamba-group
370 + # (group_id, cow_block_id, boundary_tokens) for this request's partial tail.
371 + # Present only on the producer's CoW step; drives the connector's offload
372 + # (the FA group's block is derived from block_ids and boundary_tokens).
373 + partial_tail_offloads: list[tuple[int, int, int]] | None = None
374
375 @staticmethod
376 def from_request_tracker(
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py
+53 -2
@@ -11,6 +11,9 @@ from vllm.config import VllmConfig
11 from vllm.distributed.kv_transfer.kv_connector.v1.base import (
12 KVConnectorMetadata,
13 )
14 +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.coordinator import ( # noqa: E501
15 + partial_hash_hits_enabled,
16 +)
17 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( # noqa: E501
18 LoadSpec,
19 MooncakeStoreConnectorMetadata,
@@ -63,6 +66,9 @@ class MooncakeStoreScheduler:
66 self._block_size, self._hash_block_size = resolve_kv_cache_block_sizes(
67 kv_cache_config, vllm_config
68 )
69 + self.enable_partial_hash_hits = partial_hash_hits_enabled(
70 + kv_cache_config.kv_cache_groups, self._hash_block_size
71 + )
72
73 # Per-request state
74 self.load_specs: dict[str, LoadSpec] = {} # to be loaded
@@ -80,7 +86,12 @@ class MooncakeStoreScheduler:
86 Returns ``(None, False)`` when an async lookup is still in flight,
87 signaling the scheduler to retry this request on a later step.
88 """
83 - if request.num_tokens < self._block_size:
89 + # Fine-grained hits may land on a hash boundary inside a block; without
90 + # partial hits, prefixes shorter than one physical block are skipped.
91 + align = (
92 + self._hash_block_size if self.enable_partial_hash_hits else self._block_size
93 + )
94 + if request.num_tokens < align:
95 return 0, False
96
97 num_external_hit_tokens = self.client.lookup(
@@ -338,6 +349,44 @@ class MooncakeStoreScheduler:
349 if req_meta is not None:
350 meta.add_request(req_meta)
351
352 + # Flush partial-tail offloads in the step they arrive: the CoW copy is
353 + # enqueued before the connector event records, so this step's event
354 + # fences the cow block. Ride the request's save meta when present, else
355 + # emit an offload-only ReqMeta (token_len_chunk=0 skips the normal
356 + # save; can_save=True takes the normal enqueue path).
357 + step_partial_tails = getattr(scheduler_output, "partial_tail_offloads", None)
358 + if step_partial_tails and not force_skip_save:
359 + pending = dict(step_partial_tails)
360 + for req_meta in meta.requests:
361 + if req_meta.can_save:
362 + groups = pending.pop(req_meta.req_id, None)
363 + if groups:
364 + req_meta.partial_tail_offloads = groups
365 + tracker = self._request_trackers.get(req_meta.req_id)
366 + if tracker is not None:
367 + tracker.has_pending_offload = True
368 + for req_id, groups in pending.items():
369 + tracker = self._request_trackers.get(req_id)
370 + req_tuple = self._unfinished_requests.get(req_id)
371 + if tracker is None or req_tuple is None:
372 + # Request finished/preempted within this step; its blocks
373 + # are going away, so the offload is conservatively dropped.
374 + logger.debug("Dropping partial-tail offload for request %s", req_id)
375 + continue
376 + assert len({boundary for _, _, boundary in groups}) == 1
377 + tracker.has_pending_offload = True
378 + meta.add_request(
379 + ReqMeta(
380 + req_id=req_id,
381 + token_len_chunk=0,
382 + block_ids=tracker.allocated_block_ids,
383 + block_hashes=req_tuple[0].block_hashes,
384 + can_save=True,
385 + num_prompt_tokens=tracker.prefill_end_tokens,
386 + partial_tail_offloads=groups,
387 + )
388 + )
389 +
390 return meta
391
392 def request_finished(
@@ -352,7 +401,9 @@ class MooncakeStoreScheduler:
401 # Missing tracker can happen when the request is aborted before the
402 # connector observes the normal finished lifecycle or is preempted
403 # before finishing.
355 - if tracker is None or tracker.num_saved_tokens <= 0:
404 + if tracker is None or (
405 + tracker.num_saved_tokens <= 0 and not tracker.has_pending_offload
406 + ):
407 return False, None
408 total_blocks = sum(len(g) for g in block_ids)
409 delay_free_blocks = total_blocks > 0
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py
+193 -16
@@ -61,6 +61,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.protocol import
61 from vllm.logger import init_logger
62 from vllm.utils.math_utils import cdiv
63 from vllm.utils.network_utils import get_ip, make_zmq_socket
64 +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
65 from vllm.v1.core.kv_cache_utils import (
66 BlockHash,
67 maybe_convert_block_hash,
@@ -531,6 +532,164 @@ class KVCacheStoreSendingThread(KVTransferThread):
532 self._skip_store_requests.clear()
533 return True
534
535 + def _maybe_offload_partial_tail(self, req_meta: ReqMeta) -> bool:
536 + """Offload the request's sub-block partial tail (its last prompt hash
537 + boundary) so a later request can hit the sub-block prefix.
538 +
539 + Covers every block from the normal save's lcm floor to the boundary:
540 + the normal save floors to ``lcm_block_size``, so a smaller-block
541 + group's full blocks in that gap are never persisted elsewhere, and
542 + the consumer's lookup needs every group at every probed boundary.
543 + Full blocks are keyed by their block-end hash, the partial boundary
544 + block by the boundary sub-hash; the mamba "align" boundary block is
545 + the core-provided CoW block. All keys are deduped against the store.
546 +
547 + Returns:
548 + True when no put is needed or every put succeeds, False otherwise.
549 + """
550 + if not self.coord.enable_partial_hash_hits or not req_meta.block_hashes:
551 + return True
552 + partial_tail_offloads = req_meta.partial_tail_offloads
553 + if not partial_tail_offloads:
554 + return True
555 + hash_block_size = self.coord.hash_block_size
556 + boundaries = {boundary for _, _, boundary in partial_tail_offloads}
557 + if len(boundaries) != 1:
558 + raise ValueError(
559 + "Partial-tail offloads for one request must share a boundary"
560 + )
561 + boundary = boundaries.pop()
562 + if boundary == 0:
563 + return True
564 + if boundary // hash_block_size - 1 >= len(req_meta.block_hashes):
565 + return True
566 + mamba_offloads = {
567 + group_id: block_id for group_id, block_id, _ in partial_tail_offloads
568 + }
569 +
570 + keys: list[str] = []
571 + addrs: list[list[int]] = []
572 + sizes: list[list[int]] = []
573 + saved = self._saved_offset.get(req_meta.req_id, 0)
574 + for g_idx, db in enumerate(self.token_databases):
575 + group_blocks = req_meta.block_ids[g_idx]
576 + # Distribute across ranks by the same rule as normal chunks.
577 + put_step = self.group_put_steps[g_idx]
578 + put_step_rank = (self.tp_rank + g_idx) % put_step
579 + # Always include the boundary block: its sub-hash key is written
580 + # only here, even if normal saves already advanced past it.
581 + last_block = cdiv(boundary, db.block_size) - 1
582 + for block_idx in range(
583 + min(saved // db.block_size, last_block), last_block + 1
584 + ):
585 + if block_idx % put_step != put_step_rank:
586 + continue
587 + valid_end = min((block_idx + 1) * db.block_size, boundary)
588 + key_hash = req_meta.block_hashes[valid_end // hash_block_size - 1]
589 + if (
590 + g_idx in mamba_offloads
591 + and valid_end == boundary
592 + and boundary % db.block_size != 0
593 + ):
594 + block_id = mamba_offloads[g_idx]
595 + else:
596 + if block_idx >= len(group_blocks):
597 + continue
598 + block_id = group_blocks[block_idx]
599 + if block_id == NULL_BLOCK_ID:
600 + logger.debug(
601 + "Skipping unavailable partial-tail source block "
602 + "(req=%s, group=%d, block=%d)",
603 + req_meta.req_id,
604 + g_idx,
605 + block_idx,
606 + )
607 + continue
608 + addr, size = db.prepare_value_for_block(block_id)
609 + keys.append(db.key_for(key_hash))
610 + addrs.append(addr)
611 + sizes.append(size)
612 +
613 + if not keys:
614 + return True
615 + exists_start = time.perf_counter()
616 + try:
617 + exists = self.store.batch_is_exist(keys)
618 + except Exception as e:
619 + self._record_operation(
620 + "save_exists",
621 + exists_start,
622 + len(keys),
623 + status="error",
624 + num_failed_keys=len(keys),
625 + )
626 + logger.error(
627 + "Failed to check partial-tail keys for request %s: %s",
628 + req_meta.req_id,
629 + e,
630 + )
631 + return False
632 + self._record_operation("save_exists", exists_start, len(keys))
633 + missing = [i for i, e in enumerate(exists) if e != 1]
634 + if not missing:
635 + return True
636 + keys = [keys[i] for i in missing]
637 + addrs = [addrs[i] for i in missing]
638 + sizes = [sizes[i] for i in missing]
639 + if req_meta.current_event is not None:
640 + # Fence the CoW block copy enqueued earlier this step.
641 + req_meta.current_event.synchronize()
642 + batch_bytes = _sum_batch_bytes(sizes)
643 + put_start = time.perf_counter()
644 + try:
645 + res = self.store.batch_put_from_multi_buffers(
646 + keys, addrs, sizes, self.replicate_config
647 + )
648 + except Exception as e:
649 + self._record_operation(
650 + "save_put",
651 + put_start,
652 + len(keys),
653 + num_bytes=batch_bytes,
654 + status="error",
655 + num_failed_keys=len(keys),
656 + )
657 + logger.error(
658 + "Failed to put partial-tail keys for request %s: %s",
659 + req_meta.req_id,
660 + e,
661 + )
662 + return False
663 +
664 + failed = [i for i, value in enumerate(res) if value < 0]
665 + self._record_operation(
666 + "save_put",
667 + put_start,
668 + len(keys),
669 + num_bytes=batch_bytes,
670 + status="partial_failure" if failed else "ok",
671 + num_failed_keys=len(failed),
672 + )
673 + if failed:
674 + failed_codes = {res[i] for i in failed}
675 + logger.warning(
676 + "Partial-tail put failed for request %s: %d/%d keys failed (codes=%s)",
677 + req_meta.req_id,
678 + len(failed),
679 + len(keys),
680 + failed_codes,
681 + )
682 + if MOONCAKE_NO_AVAILABLE_HANDLE in failed_codes:
683 + self._mark_request_skipped_for_pressure(req_meta.req_id)
684 + return False
685 +
686 + if self._clear_store_pressure():
687 + logger.info(
688 + "Mooncake CPU/disk offloading pressure cleared after a "
689 + "successful partial-tail batch"
690 + )
691 + return True
692 +
693 def _handle_request(self, req_meta: ReqMeta):
694 # Cache hits are always a multiple of ``lcm_block_size`` tokens, which
695 # is also ``store_mask``'s precondition.
@@ -548,9 +707,6 @@ class KVCacheStoreSendingThread(KVTransferThread):
707 # so the scheduler can release the GPU blocks it pinned for this
708 # request (via `delay_free_blocks`) even when the store path raises.
709 try:
551 - if token_len == 0:
552 - return
553 -
710 if self._should_skip_request(req_id):
711 logger.debug(
712 "Skipping Mooncake store for request %s while CPU/disk "
@@ -559,6 +715,16 @@ class KVCacheStoreSendingThread(KVTransferThread):
715 )
716 return
717
718 + # Offload the sub-block partial tail (independent of the normal
719 + # block-aligned save, which may be skipped this step).
720 + if req_meta.partial_tail_offloads is not None and not (
721 + self._maybe_offload_partial_tail(req_meta)
722 + ):
723 + return
724 +
725 + if token_len == 0:
726 + return
727 +
728 # Resume from where this rank left off; only the new suffix is saved.
729 save_start = self._saved_offset.get(req_id, 0)
730
@@ -1385,7 +1551,7 @@ class MooncakeStoreWorker:
1551 self.recv_request_queue.put(request)
1552
1553 assert self.load_async, "load_async must be True for better performance."
1388 - # Issue stores with CUDA event synchronization
1554 + # Issue stores with CUDA event synchronization.
1555 if self.kv_role in ["kv_producer", "kv_both"]:
1556 current_event = None
1557 for request in meta.requests:
@@ -1501,21 +1667,32 @@ class MooncakeStoreWorker:
1667 # candidate_meta stores the (group, hash_bytes) for key slice.
1668 candidate_keys: list[str] = []
1669 candidate_meta: list[tuple[int, bytes]] = []
1504 - lookup_masks = self.coord.lookup_mask(token_len)
1670 + fine_grained = self.coord.enable_partial_hash_hits
1671 + lookup_masks = None if fine_grained else self.coord.lookup_mask(token_len)
1672 for g_idx, db in enumerate(self.token_dbs):
1673 spec_block_size = db.block_size
1507 - lookup_mask = lookup_masks[g_idx]
1674 key_prefixes = self._lookup_key_prefixes[g_idx]
1509 - group_hashes = self.coord.block_hashes_for_spec(
1510 - block_hashes, self._kv_cache_groups[g_idx].kv_cache_spec
1511 - )
1512 - max_chunks = min(len(group_hashes), cdiv(token_len, spec_block_size))
1513 - mask_limit = (
1514 - max_chunks if lookup_mask is None else min(max_chunks, len(lookup_mask))
1515 - )
1516 - for chunk_id in range(mask_limit):
1517 - if lookup_mask is not None and not lookup_mask[chunk_id]:
1518 - continue
1675 + if fine_grained:
1676 + max_units = min(len(block_hashes), token_len // self.hash_block_size)
1677 + unit_ids: range | list[int] = range(max_units)
1678 + group_hashes: Sequence[BlockHash] = block_hashes
1679 + else:
1680 + lookup_mask = lookup_masks[g_idx] # type: ignore[index]
1681 + group_hashes = self.coord.block_hashes_for_spec(
1682 + block_hashes, self._kv_cache_groups[g_idx].kv_cache_spec
1683 + )
1684 + max_chunks = min(len(group_hashes), cdiv(token_len, spec_block_size))
1685 + mask_limit = (
1686 + max_chunks
1687 + if lookup_mask is None
1688 + else min(max_chunks, len(lookup_mask))
1689 + )
1690 + unit_ids = [
1691 + chunk_id
1692 + for chunk_id in range(mask_limit)
1693 + if lookup_mask is None or lookup_mask[chunk_id]
1694 + ]
1695 + for chunk_id in unit_ids:
1696 h = group_hashes[chunk_id]
1697 hash_hex = h.hex()
1698 for key_prefix in key_prefixes:
vllm/v1/core/kv_cache_manager.py
+62 -1
@@ -186,6 +186,10 @@ class KVCacheManager:
186 tuple(() for _ in range(self.num_kv_cache_groups))
187 )
188
189 + # Off-table cow blocks handed to a KV connector for partial-tail
190 + # offload; pinned until the request's blocks are freed.
191 + self._partial_tail_pins: dict[str, list[KVCacheBlock]] = {}
192 +
193 @property
194 def usage(self) -> float:
195 """Get the KV cache usage.
@@ -568,6 +572,9 @@ class KVCacheManager:
572 Args:
573 request: The request to free the blocks.
574 """
575 + pins = self._partial_tail_pins.pop(request.request_id, None)
576 + if pins:
577 + self.block_pool.free_blocks(pins)
578 self.coordinator.free(request.request_id)
579
580 def remove_skipped_blocks(
@@ -600,7 +607,14 @@ class KVCacheManager:
607 Returns:
608 The request's blocks in allocation order.
609 """
603 - return self.coordinator.pop_blocks_for_free(request.request_id)
610 + blocks = self.coordinator.pop_blocks_for_free(request.request_id)
611 + # Pins ride the same (possibly deferred) free as the request blocks.
612 + # Preemption may release a pin under a still-queued offload — the same
613 + # exposure normal saves of table blocks already have.
614 + pins = self._partial_tail_pins.pop(request.request_id, None)
615 + if pins:
616 + blocks = pins + blocks
617 + return blocks
618
619 def evict_blocks(self, block_ids: set[int]) -> None:
620 """evict blocks from the prefix cache by their block IDs.
@@ -760,6 +774,25 @@ class KVCacheManager:
774 # Only create new KVCacheBlocks for non-empty blocks
775 return KVCacheBlocks(blocks) if any(blocks) else self.empty_kv_cache_blocks
776
777 + def truncate_computed_blocks(
778 + self, blocks: KVCacheBlocks, num_computed_tokens: int
779 + ) -> KVCacheBlocks:
780 + """Return a lookup-result view truncated at an aligned token endpoint.
781 +
782 + Pure slicing: refcounts are untouched and ``blocks`` is not mutated.
783 + """
784 + truncated: list[list[KVCacheBlock]] = []
785 + for group_blocks, manager in zip(
786 + blocks.blocks,
787 + self.coordinator.single_type_managers,
788 + strict=True,
789 + ):
790 + assert num_computed_tokens % manager.block_size == 0
791 + num_blocks = num_computed_tokens // manager.block_size
792 + assert num_blocks <= len(group_blocks)
793 + truncated.append(list(group_blocks[:num_blocks]))
794 + return self.create_kv_cache_blocks(tuple(truncated))
795 +
796 def take_new_block_ids(self) -> list[int]:
797 """Drain and return new attention block IDs for zeroing."""
798 ids: list[int] = []
@@ -812,6 +845,34 @@ class KVCacheManager:
845 retained_blocks = [block for pair in pending_copies for block in pair]
846 return copies, retained_blocks
847
848 + def take_partial_tail_offloads(self) -> dict[str, list[tuple[int, int, int]]]:
849 + """Drain producer partial-tail offload hand-offs per request.
850 +
851 + Returns ``{request_id: [(group_id, block_id, boundary_tokens), ...]}``
852 + for the durable boundary blocks of producers' last-prompt-boundary
853 + partial tails. Only mamba "align" groups contribute; empty otherwise.
854 + A KV connector reads the referenced blocks and offloads them so a later
855 + request can hit the sub-block prefix.
856 +
857 + Each handed-off block lives off the request block table, so it is
858 + pinned here and unpinned when the request's blocks are freed — for a
859 + producer with saved tokens, after the connector reports sends done.
860 + """
861 + offloads: dict[str, list[tuple[int, int, int]]] = {}
862 + for mgr in self.coordinator.single_type_managers:
863 + for (
864 + req_id,
865 + group_id,
866 + block,
867 + boundary_tokens,
868 + ) in mgr.take_pending_partial_tail_offloads():
869 + self.block_pool.touch((block,))
870 + self._partial_tail_pins.setdefault(req_id, []).append(block)
871 + offloads.setdefault(req_id, []).append(
872 + (group_id, block.block_id, boundary_tokens)
873 + )
874 + return offloads
875 +
876 def new_step_starts(self) -> None:
877 """Notify the coordinator that a new step is starting."""
878 self.coordinator.new_step_starts()
vllm/v1/core/sched/output.py
+6
@@ -258,6 +258,12 @@ class SchedulerOutput:
258 # CoW copies to apply after zeroing new blocks and before forward.
259 kv_cache_block_copies: list[KVCacheBlockCopy] | None = None
260
261 + # Producer partial-tail offload hand-off for external KV connectors:
262 + # {request_id: [(group_id, block_id, boundary_tokens), ...]} pointing at
263 + # the durable boundary block of a producer's last-prompt-boundary partial
264 + # tail (mamba "align" CoW target). None unless partial hash hits are active.
265 + partial_tail_offloads: dict[str, list[tuple[int, int, int]]] | None = None
266 +
267 # Dynamic speculative decoding: optimal K chosen by scheduler.
268 # Number of spec tokens to schedule for the next step.
269 num_spec_tokens_to_schedule: int = 0
vllm/v1/core/sched/scheduler.py
+46 -2
@@ -740,9 +740,16 @@ class Scheduler(SchedulerInterface):
740
741 # Get externally-cached tokens if using a KVConnector.
742 if self.connector is not None:
743 + # Present a block-aligned local hit to the connector so
744 + # a strictly longer remote hit can supersede a local
745 + # sub-block tail without racing its copy-on-write.
746 + partial_tail = num_new_local_computed_tokens % self.block_size
747 + block_aligned_local = (
748 + num_new_local_computed_tokens - partial_tail
749 + )
750 ext_tokens, load_kv_async = (
751 self.connector.get_num_new_matched_tokens(
745 - request, num_new_local_computed_tokens
752 + request, block_aligned_local
753 )
754 )
755
@@ -754,7 +761,27 @@ class Scheduler(SchedulerInterface):
761 step_skipped_waiting.prepend_request(request)
762 continue
763
757 - num_external_computed_tokens = ext_tokens
764 + if partial_tail and ext_tokens > partial_tail:
765 + # Remote strictly exceeds the full local hit: drop the
766 + # sub-block tail so no CoW is needed, and let the load
767 + # cover it. Trim the partial block out of the local
768 + # computed blocks so it is not adopted from the cache.
769 + new_computed_blocks = (
770 + self.kv_cache_manager.truncate_computed_blocks(
771 + new_computed_blocks, block_aligned_local
772 + )
773 + )
774 + num_new_local_computed_tokens = block_aligned_local
775 + num_external_computed_tokens = ext_tokens
776 + elif partial_tail:
777 + # Remote does not exceed the full local hit: keep the
778 + # local sub-block tail and load nothing external.
779 + num_external_computed_tokens = 0
780 + # Nothing to load remotely -> not an async-load step;
781 + # clearing avoids the `load_kv_async` assert below.
782 + load_kv_async = False
783 + else:
784 + num_external_computed_tokens = ext_tokens
785
786 if hit_diverged and num_external_computed_tokens == 0:
787 # No external tokens back the deeper local hit, so its
@@ -1108,6 +1135,22 @@ class Scheduler(SchedulerInterface):
1135 self.prev_step_scheduled_req_ids.clear()
1136 self.prev_step_scheduled_req_ids.update(num_scheduled_tokens.keys())
1137
1138 + # Producer partial-tail hand-off for external KV connectors. Drained
1139 + # before the CoW retentions are released below, so the pin lands while
1140 + # the cow block still holds a retention ref. Without a producer-side
1141 + # connector nothing consumes the hand-off, so skip the drain (and its
1142 + # pin); the manager drops stale entries when the request's blocks are
1143 + # popped for free.
1144 + pending_partial_tail_offloads = None
1145 + if (
1146 + self.connector is not None
1147 + and self.vllm_config.kv_transfer_config is not None
1148 + and self.vllm_config.kv_transfer_config.is_kv_producer
1149 + ):
1150 + pending_partial_tail_offloads = (
1151 + self.kv_cache_manager.take_partial_tail_offloads() or None
1152 + )
1153 +
1154 kv_cache_block_copies, cow_retained_blocks = (
1155 self.kv_cache_manager.take_kv_cache_block_copies()
1156 )
@@ -1153,6 +1196,7 @@ class Scheduler(SchedulerInterface):
1196 free_encoder_mm_hashes=self.encoder_cache_manager.get_freed_mm_hashes(),
1197 new_block_ids_to_zero=self._get_new_block_ids_to_zero(),
1198 kv_cache_block_copies=pending_kv_cache_block_copies,
1199 + partial_tail_offloads=pending_partial_tail_offloads,
1200 num_spec_tokens_to_schedule=num_spec_tokens_to_schedule,
1201 ec_manager_metadata=self.encoder_cache_manager.get_manager_metadata(),
1202 )
vllm/v1/core/single_type_kv_cache_manager.py
+63 -3
@@ -115,6 +115,15 @@ class SingleTypeKVCacheManager(ABC):
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 + # Partial-tail offload hand-off for external KV connectors: when a
119 + # producer registers its last-prompt-boundary partial tail and the
120 + # durable boundary block is not on the append-only request block table
121 + # (mamba "align" CoW target), record the request, group, block, and
122 + # exact token boundary so a connector can offload it under the right
123 + # hash. Populated only by mamba "align".
124 + self._pending_partial_tail_offloads: list[
125 + tuple[str, int, KVCacheBlock, int]
126 + ] = []
127
128 @classmethod
129 def _get_num_evictable_blocks(cls, blocks: Sequence[KVCacheBlock]):
@@ -378,6 +387,21 @@ class SingleTypeKVCacheManager(ABC):
387 self._pending_cow_copies = []
388 return pending_copies
389
390 + def take_pending_partial_tail_offloads(
391 + self,
392 + ) -> list[tuple[str, int, KVCacheBlock, int]]:
393 + """Drain producer partial-tail hand-offs.
394 +
395 + Entries are ``(req_id, group_id, block, boundary_tokens)``.
396 +
397 + Only mamba "align" populates this. The block lives off the request
398 + block table, so the caller must pin it until the connector has read
399 + it — nothing else keeps it alive once the CoW retention is released.
400 + """
401 + pending = self._pending_partial_tail_offloads
402 + self._pending_partial_tail_offloads = []
403 + return pending
404 +
405 def _apply_cow(
406 self,
407 request_id: str,
@@ -693,7 +717,8 @@ class FullAttentionManager(SingleTypeKVCacheManager):
717 alignment_tokens < block_size and block_size % alignment_tokens == 0
718 )
719 if fine_grained:
696 - assert isinstance(block_hashes, list)
720 + # list or lazy BlobBlockHashes view
721 + assert isinstance(block_hashes, Sequence)
722 full_block_hashes: BlockHashList = BlockHashListWithBlockSize(
723 block_hashes, alignment_tokens, block_size
724 )
@@ -716,7 +741,8 @@ class FullAttentionManager(SingleTypeKVCacheManager):
741 # Phase 2 (fine-grained only): extend into the first non-full block by
742 # probing its interior hash boundaries high-to-low (longest hit first).
743 if fine_grained:
719 - assert isinstance(block_hashes, list)
744 + # list or lazy BlobBlockHashes view
745 + assert isinstance(block_hashes, Sequence)
746 scale_factor = block_size // alignment_tokens
747 first_partial_idx = len(computed_blocks[0]) * scale_factor
748 max_partial_idx = min(
@@ -1244,6 +1270,11 @@ class MambaManager(SingleTypeKVCacheManager):
1270 self.last_state_block_idx: dict[str, int] = {}
1271 # The set of the requests that have been allocated blocks
1272 self._allocated_block_reqs: set[str] = set()
1273 + # Requests that registered their own last-prompt-boundary partial
1274 + # tail (producers). On the next step's CoW the boundary state moves
1275 + # into a private cow_block; we record that block for connector
1276 + # offload (see _pending_partial_tail_offloads).
1277 + self._producer_partial_tail_reqs: dict[str, int] = {}
1278
1279 @classmethod
1280 def find_longest_cache_hit(
@@ -1277,7 +1308,8 @@ class MambaManager(SingleTypeKVCacheManager):
1308
1309 block_size = kv_cache_spec.block_size
1310 if alignment_tokens < block_size and block_size % alignment_tokens == 0:
1280 - assert isinstance(block_hashes, list)
1311 + # list or lazy BlobBlockHashes view
1312 + assert isinstance(block_hashes, Sequence)
1313 hash_block_size = alignment_tokens
1314 scale_factor = block_size // hash_block_size
1315 max_num_partial_units = min(
@@ -1590,6 +1622,21 @@ class MambaManager(SingleTypeKVCacheManager):
1622 self.block_pool.move_block_hashes(source_block, cow_block)
1623 self._pending_cow_copies.append((source_block, cow_block))
1624 source_block.ref_cnt += 1
1625 + boundary_tokens = self._producer_partial_tail_reqs.pop(
1626 + request_id, None
1627 + )
1628 + if boundary_tokens is not None:
1629 + # This CoW preserved a producer's own boundary
1630 + # state in cow_block; hand it to the connector for
1631 + # partial-tail offload once the copy has run.
1632 + self._pending_partial_tail_offloads.append(
1633 + (
1634 + request_id,
1635 + self.kv_cache_group_id,
1636 + cow_block,
1637 + boundary_tokens,
1638 + )
1639 + )
1640 if cow_block.block_hash is not None:
1641 # The moved entry is only filled by this step's
1642 # copy, so defer same-step hits on it.
@@ -1607,6 +1654,14 @@ class MambaManager(SingleTypeKVCacheManager):
1654 if self.mamba_cache_mode == "align":
1655 self._allocated_block_reqs.discard(request_id)
1656 self.last_state_block_idx.pop(request_id, None)
1657 + self._producer_partial_tail_reqs.pop(request_id, None)
1658 + # A hand-off whose request died in this same scheduling pass must
1659 + # not reach the connector: its unpin hook (free) has already run.
1660 + self._pending_partial_tail_offloads = [
1661 + entry
1662 + for entry in self._pending_partial_tail_offloads
1663 + if entry[0] != request_id
1664 + ]
1665 return super().pop_blocks_for_free(request_id)
1666
1667 def get_num_skipped_tokens(self, num_computed_tokens: int) -> int:
@@ -1681,6 +1736,11 @@ class MambaManager(SingleTypeKVCacheManager):
1736 if partial_hash is not None:
1737 self._partial_hit_reqs[request.request_id] = (block_idx, source_block)
1738 self.num_cached_block[request.request_id] = block_idx
1739 + # Producer of this partial tail: the boundary state currently lives
1740 + # in ``source_block`` but the next step's forward overwrites it. The
1741 + # upcoming CoW copies it into a durable cow_block; record the req so
1742 + # allocate_new_blocks hands that block to the connector for offload.
1743 + self._producer_partial_tail_reqs[request.request_id] = num_tokens
1744 return partial_hash
1745
1746