[Mooncake] Only check and store new KV cache range (#46412)
Signed-off-by: wzhao18 <wzhao18.sz@gmail.com> Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai>
Wei Zhao committed
Jun 24, 2026 at 06:10 UTC
d20dbf921b9074c68cd36181cbd961fb46c62caf
6 files changed
+487
-60
tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py
+41
@@ -246,6 +246,20 @@ def test_store_mask_swa_wider_window_covers_more_blocks_per_lcm():
246
assert masks[1] == [False, False, True, True, False, False, True, True]
247
248
249
+def test_store_mask_swa_prefix_stable_as_aligned_length_grows():
250
+ full = _full(32)
251
+ swa = _swa(block_size=8, sliding_window=8)
252
+ groups = [KVCacheGroupSpec(["L0"], full), KVCacheGroupSpec(["L1"], swa)]
253
+ coord = _make_coord(groups, hash_block_size=8)
254
+
255
+ shorter = coord.store_mask(64)[1]
256
+ longer = coord.store_mask(128)[1]
257
+
258
+ assert shorter is not None
259
+ assert longer is not None
260
+ assert longer[: len(shorter)] == shorter
261
+
262
+
263
def test_store_mask_dsv4_5_groups_full_mla_plus_4_swa():
264
"""DSV4-shaped: full-MLA(B=256) + 4 SWA groups with B in {64, 64, 4, 8}
265
and varied sliding windows. lcm=256, hash_block_size=4. Two lcm segments
@@ -355,6 +369,33 @@ def test_store_mask_retention_interval_keeps_segment_and_replay_tails():
369
assert masks[1] == [i in (7, 11, 15) for i in range(16)]
370
371
372
+def test_store_mask_suffix_matches_full_mask_tail():
373
+ coord = _make_coord(_retention_groups(), hash_block_size=8, retention_interval=64)
374
+ full = coord.store_mask(128, num_prompt_tokens=100)
375
+ suffix = coord.store_mask(128, start_token=64, num_prompt_tokens=100)
376
+
377
+ for g_idx, cache_group in enumerate(coord.kv_cache_groups):
378
+ block_size = cache_group.kv_cache_spec.block_size
379
+ start_chunk = 64 // block_size
380
+ end_chunk = 128 // block_size
381
+ full_mask = full[g_idx]
382
+ if full_mask is None:
383
+ assert suffix[g_idx] is None
384
+ else:
385
+ assert suffix[g_idx] == full_mask[start_chunk:end_chunk]
386
+
387
+
388
+def test_store_mask_retention_prefix_stable_as_aligned_length_grows():
389
+ coord = _make_coord(_retention_groups(), hash_block_size=8, retention_interval=0)
390
+
391
+ shorter = coord.store_mask(64, num_prompt_tokens=100)[1]
392
+ longer = coord.store_mask(128, num_prompt_tokens=100)[1]
393
+
394
+ assert shorter is not None
395
+ assert longer is not None
396
+ assert longer[: len(shorter)] == shorter
397
+
398
+
399
# ----- Eagle / MTP interaction with load_mask -----
400
401
tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py
+2
-2
@@ -337,5 +337,5 @@ def test_chunked_token_database_hash_block_size_smaller_than_block_size():
337
assert out[1][0] == 16 and out[1][1] == 32
338
# Each chunk's hash is its last (4th) fine hash, which already chains the
339
# prior three.
340
- assert out[0][2].chunk_hash == fine_hashes[3].hex()
341
- assert out[1][2].chunk_hash == fine_hashes[7].hex()
340
+ assert out[0][2].hex() == fine_hashes[3].hex()
341
+ assert out[1][2].hex() == fine_hashes[7].hex()
tests/v1/kv_connector/unit/test_mooncake_store_worker.py
+340
-5
@@ -36,6 +36,25 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.metrics import
36
from vllm.v1.core.kv_cache_utils import BlockHash
37
38
39
+class _RecordingBlockHashes:
40
+ def __init__(self, values: list[bytes]):
41
+ self.values = values
42
+ self.accessed: list[int] = []
43
+
44
+ def __len__(self):
45
+ return len(self.values)
46
+
47
+ def __getitem__(self, idx):
48
+ if isinstance(idx, slice):
49
+ return [self[i] for i in range(*idx.indices(len(self)))]
50
+ if idx < 0:
51
+ idx += len(self.values)
52
+ if not 0 <= idx < len(self.values):
53
+ raise IndexError(idx)
54
+ self.accessed.append(idx)
55
+ return self.values[idx]
56
+
57
+
58
def _default_send_coord() -> mooncake_store_worker.MooncakeStoreCoordinator:
59
from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheGroupSpec
60
@@ -53,6 +72,8 @@ def _make_store_sending_thread(
72
coord: mooncake_store_worker.MooncakeStoreCoordinator | None = None,
73
token_databases: list[ChunkedTokenDatabase] | None = None,
74
block_size: int = 16,
75
+ tp_rank: int = 0,
76
+ put_step: int = 1,
77
replicate_config: object | None = None,
78
) -> mooncake_store_worker.KVCacheStoreSendingThread:
79
if coord is None:
@@ -67,8 +88,8 @@ def _make_store_sending_thread(
88
token_databases=token_databases,
89
block_size=block_size,
90
coord=coord,
70
- tp_rank=0,
71
- put_step=1,
91
+ tp_rank=tp_rank,
92
+ put_step=put_step,
93
kv_role="kv_producer",
94
ready_event=threading.Event(),
95
replicate_config=replicate_config,
@@ -391,6 +412,248 @@ def test_store_sending_thread_records_mooncake_metrics():
412
assert stats.data["save_put"][0]["status"] == "ok"
413
414
415
+def test_process_tokens_uses_mask_num_as_start_chunk():
416
+ db = ChunkedTokenDatabase(
417
+ KeyMetadata("test-model", 0, 0, 0, 0),
418
+ block_size=32,
419
+ hash_block_size=8,
420
+ )
421
+ block_hashes = _RecordingBlockHashes([bytes([i]) for i in range(16)])
422
+
423
+ results = list(
424
+ db.process_tokens(
425
+ token_len=96,
426
+ block_hashes=block_hashes,
427
+ mask_num=64,
428
+ )
429
+ )
430
+
431
+ assert block_hashes.accessed == [11]
432
+ assert results == [(64, 96, bytes([11]))]
433
+
434
+
435
+def test_process_tokens_applies_chunk_mask_before_hash_access():
436
+ db = ChunkedTokenDatabase(
437
+ KeyMetadata("test-model", 0, 0, 0, 0),
438
+ block_size=32,
439
+ hash_block_size=8,
440
+ )
441
+ block_hashes = _RecordingBlockHashes([bytes([i]) for i in range(16)])
442
+
443
+ results = list(
444
+ db.process_tokens(
445
+ token_len=128,
446
+ block_hashes=block_hashes,
447
+ mask_num=64,
448
+ chunk_mask=[False, True],
449
+ )
450
+ )
451
+
452
+ assert block_hashes.accessed == [15]
453
+ assert results == [(96, 128, bytes([15]))]
454
+
455
+
456
+def test_process_tokens_applies_stride_before_hash_access():
457
+ db = ChunkedTokenDatabase(
458
+ KeyMetadata("test-model", 0, 0, 0, 0),
459
+ block_size=32,
460
+ hash_block_size=8,
461
+ )
462
+ block_hashes = _RecordingBlockHashes([bytes([i]) for i in range(16)])
463
+
464
+ results = list(
465
+ db.process_tokens(
466
+ token_len=128,
467
+ block_hashes=block_hashes,
468
+ mask_num=64,
469
+ put_step=2,
470
+ put_step_rank=1,
471
+ )
472
+ )
473
+
474
+ assert block_hashes.accessed == [15]
475
+ assert results == [(96, 128, bytes([15]))]
476
+
477
+
478
+def test_store_sending_thread_delta_saves_only_new_full_attention_chunks():
479
+ store = MagicMock()
480
+ store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
481
+ store.batch_put_from_multi_buffers.return_value = [256, 256]
482
+ thread = _make_store_sending_thread(store)
483
+
484
+ thread.add_stored_request("req-a")
485
+ thread._saved_offset["req-a"] = 32
486
+ thread._handle_request(
487
+ ReqMeta(
488
+ req_id="req-a",
489
+ token_len_chunk=64,
490
+ block_ids=([0, 1, 2, 3],),
491
+ block_hashes=[b"a0", b"a1", b"a2", b"a3"],
492
+ can_save=True,
493
+ )
494
+ )
495
+
496
+ keys = store.batch_is_exist.call_args.args[0]
497
+ assert keys == [
498
+ "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6132",
499
+ "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6133",
500
+ ]
501
+ assert store.batch_put_from_multi_buffers.call_args.args[0] == keys
502
+
503
+
504
+def test_store_sending_thread_delta_strides_with_local_phase():
505
+ store = MagicMock()
506
+ store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
507
+ store.batch_put_from_multi_buffers.return_value = [256]
508
+ thread = _make_store_sending_thread(store, tp_rank=0, put_step=2)
509
+
510
+ thread.add_stored_request("req-a")
511
+ thread._saved_offset["req-a"] = 16
512
+ thread._handle_request(
513
+ ReqMeta(
514
+ req_id="req-a",
515
+ token_len_chunk=64,
516
+ block_ids=([0, 1, 2, 3],),
517
+ block_hashes=[b"a0", b"a1", b"a2", b"a3"],
518
+ can_save=True,
519
+ )
520
+ )
521
+
522
+ keys = store.batch_is_exist.call_args.args[0]
523
+ assert keys == [
524
+ "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6132",
525
+ ]
526
+ assert store.batch_put_from_multi_buffers.call_args.args[0] == keys
527
+
528
+
529
+def test_store_sending_thread_retries_skipped_range_after_pressure():
530
+ store = MagicMock()
531
+ store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
532
+ store.batch_put_from_multi_buffers.side_effect = lambda keys, *a: [256] * len(keys)
533
+ thread = _make_store_sending_thread(store, tp_rank=0, put_step=2)
534
+
535
+ # Under pressure the request is skipped without persisting anything, so the
536
+ # saved offset stays at 0.
537
+ thread._store_pressure_active = True
538
+ thread._skip_store_requests.add("req-a")
539
+
540
+ thread.add_stored_request("req-a")
541
+ thread._handle_request(
542
+ ReqMeta(
543
+ req_id="req-a",
544
+ token_len_chunk=16,
545
+ block_ids=([0],),
546
+ block_hashes=[b"a0"],
547
+ can_save=True,
548
+ )
549
+ )
550
+
551
+ store.batch_is_exist.assert_not_called()
552
+ assert thread._saved_offset.get("req-a", 0) == 0
553
+
554
+ thread._store_pressure_active = False
555
+ thread._skip_store_requests.clear()
556
+
557
+ # The next batch resumes from offset 0, re-covering the chunk skipped under
558
+ # pressure (chunk 0) rather than losing it.
559
+ thread.add_stored_request("req-a")
560
+ thread._handle_request(
561
+ ReqMeta(
562
+ req_id="req-a",
563
+ token_len_chunk=64,
564
+ block_ids=([0, 1, 2, 3],),
565
+ block_hashes=[b"a0", b"a1", b"a2", b"a3"],
566
+ can_save=True,
567
+ )
568
+ )
569
+
570
+ keys = store.batch_is_exist.call_args.args[0]
571
+ assert keys == [
572
+ "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6130",
573
+ "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6132",
574
+ ]
575
+ assert store.batch_put_from_multi_buffers.call_args.args[0] == keys
576
+
577
+
578
+def test_store_sending_thread_delta_start_rank_saves_second_local_chunk():
579
+ store = MagicMock()
580
+ store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
581
+ store.batch_put_from_multi_buffers.return_value = [256, 256]
582
+ thread = _make_store_sending_thread(store, tp_rank=1, put_step=2)
583
+
584
+ thread.add_stored_request("req-a")
585
+ thread._saved_offset["req-a"] = 16
586
+ thread._handle_request(
587
+ ReqMeta(
588
+ req_id="req-a",
589
+ token_len_chunk=64,
590
+ block_ids=([0, 1, 2, 3],),
591
+ block_hashes=[b"a0", b"a1", b"a2", b"a3"],
592
+ can_save=True,
593
+ )
594
+ )
595
+
596
+ keys = store.batch_is_exist.call_args.args[0]
597
+ assert keys == [
598
+ "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6131",
599
+ "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6133",
600
+ ]
601
+ assert store.batch_put_from_multi_buffers.call_args.args[0] == keys
602
+
603
+
604
+def test_store_sending_thread_delta_saves_only_new_masked_chunks():
605
+ store = MagicMock()
606
+ store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
607
+ store.batch_put_from_multi_buffers.side_effect = (
608
+ lambda keys, addrs, sizes, replicate_config: [256] * len(keys)
609
+ )
610
+ coord = SimpleNamespace(
611
+ lcm_block_size=16,
612
+ store_mask=lambda token_len, start_token, num_prompt_tokens=None: (
613
+ None,
614
+ [True, False],
615
+ ),
616
+ )
617
+
618
+ db_full = ChunkedTokenDatabase(
619
+ KeyMetadata("test-model", 0, 0, 0, 0, group_id=0),
620
+ block_size=16,
621
+ )
622
+ db_full.set_kv_caches_base_addr([0x1000])
623
+ db_full.set_block_len([256])
624
+ db_masked = ChunkedTokenDatabase(
625
+ KeyMetadata("test-model", 0, 0, 0, 0, group_id=1),
626
+ block_size=16,
627
+ )
628
+ db_masked.set_kv_caches_base_addr([0x2000])
629
+ db_masked.set_block_len([256])
630
+
631
+ thread = _make_store_sending_thread(
632
+ store,
633
+ coord=coord,
634
+ token_databases=[db_full, db_masked],
635
+ )
636
+
637
+ thread.add_stored_request("req-a")
638
+ thread._saved_offset["req-a"] = 32
639
+ thread._handle_request(
640
+ ReqMeta(
641
+ req_id="req-a",
642
+ token_len_chunk=64,
643
+ block_ids=([0, 1, 2, 3], [0, 1, 2, 3]),
644
+ block_hashes=[b"a0", b"a1", b"a2", b"a3"],
645
+ can_save=True,
646
+ )
647
+ )
648
+
649
+ keys = store.batch_is_exist.call_args.args[0]
650
+ full_hashes = [k.rsplit("@", 1)[-1] for k in keys if "@group:0" in k]
651
+ masked_hashes = [k.rsplit("@", 1)[-1] for k in keys if "@group:1" in k]
652
+
653
+ assert full_hashes == [b"a2".hex(), b"a3".hex()]
654
+ assert masked_hashes == [b"a2".hex()]
655
+
656
+
657
def test_store_sending_thread_only_skips_on_no_available_handle():
658
store = MagicMock()
659
store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
@@ -943,7 +1206,8 @@ def test_worker_put_striding_covers_every_rank_get_namespace(
1206
db = w.token_dbs[0]
1207
token_len = len(block_hashes) * db.block_size
1208
keys = [
946
- key.to_string() for _, _, key in db.process_tokens(token_len, block_hashes)
1209
+ PoolKey(db.metadata, block_hash.hex()).to_string()
1210
+ for _, _, block_hash in db.process_tokens(token_len, block_hashes)
1211
]
1212
assert len(keys) == len(block_hashes)
1213
# PUT side: mirrors KVCacheStoreSendingThread's striding slice.
@@ -1047,8 +1311,8 @@ def test_store_sending_thread_only_stores_swa_blocks_in_window():
1311
1312
store = MagicMock()
1313
store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
1050
- store.batch_put_from_multi_buffers.side_effect = lambda keys, addrs, sizes: (
1051
- [256] * len(keys)
1314
+ store.batch_put_from_multi_buffers.side_effect = (
1315
+ lambda keys, addrs, sizes, replicate_config: [256] * len(keys)
1316
)
1317
1318
full_spec = FullAttentionSpec(
@@ -1113,6 +1377,77 @@ def test_store_sending_thread_only_stores_swa_blocks_in_window():
1377
assert swa_hashes == {hs[3].hex(), hs[7].hex()}
1378
1379
1380
+def test_store_sending_thread_delta_saves_only_new_swa_boundary_chunks():
1381
+ from vllm.v1.kv_cache_interface import (
1382
+ FullAttentionSpec,
1383
+ KVCacheGroupSpec,
1384
+ SlidingWindowSpec,
1385
+ )
1386
+
1387
+ store = MagicMock()
1388
+ store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
1389
+ store.batch_put_from_multi_buffers.side_effect = (
1390
+ lambda keys, addrs, sizes, replicate_config: [256] * len(keys)
1391
+ )
1392
+
1393
+ full_spec = FullAttentionSpec(
1394
+ block_size=32, num_kv_heads=8, head_size=64, dtype=None
1395
+ )
1396
+ swa_spec = SlidingWindowSpec(
1397
+ block_size=8,
1398
+ num_kv_heads=8,
1399
+ head_size=64,
1400
+ dtype=None,
1401
+ sliding_window=8,
1402
+ )
1403
+ coord = mooncake_store_worker.MooncakeStoreCoordinator(
1404
+ [KVCacheGroupSpec(["L0"], full_spec), KVCacheGroupSpec(["L1"], swa_spec)],
1405
+ scheduler_block_size=32,
1406
+ hash_block_size=8,
1407
+ )
1408
+
1409
+ db_full = ChunkedTokenDatabase(
1410
+ KeyMetadata("test-model", 0, 0, 0, 0, group_id=0),
1411
+ block_size=32,
1412
+ hash_block_size=8,
1413
+ )
1414
+ db_full.set_kv_caches_base_addr([0x1000])
1415
+ db_full.set_block_len([512])
1416
+ db_swa = ChunkedTokenDatabase(
1417
+ KeyMetadata("test-model", 0, 0, 0, 0, group_id=1),
1418
+ block_size=8,
1419
+ hash_block_size=8,
1420
+ )
1421
+ db_swa.set_kv_caches_base_addr([0x2000])
1422
+ db_swa.set_block_len([128])
1423
+
1424
+ thread = _make_store_sending_thread(
1425
+ store,
1426
+ coord=coord,
1427
+ token_databases=[db_full, db_swa],
1428
+ block_size=32,
1429
+ )
1430
+
1431
+ hs = [bytes([i + 1]) * 4 for i in range(8)]
1432
+ thread.add_stored_request("r0")
1433
+ thread._saved_offset["r0"] = 32
1434
+ thread._handle_request(
1435
+ ReqMeta(
1436
+ req_id="r0",
1437
+ token_len_chunk=64,
1438
+ block_ids=([0, 1], list(range(8))),
1439
+ block_hashes=hs,
1440
+ can_save=True,
1441
+ )
1442
+ )
1443
+
1444
+ keys = store.batch_put_from_multi_buffers.call_args.args[0]
1445
+ full_hashes = [k.rsplit("@", 1)[-1] for k in keys if "@group:0" in k]
1446
+ swa_hashes = [k.rsplit("@", 1)[-1] for k in keys if "@group:1" in k]
1447
+ assert full_hashes == [hs[7].hex()]
1448
+ assert swa_hashes == [hs[7].hex()]
1449
+
1450
+
1451
def test_store_sending_thread_kv_events_use_group_chunk_metadata():
1452
from vllm.v1.core.kv_cache_utils import BlockHash, maybe_convert_block_hash
1453
from vllm.v1.kv_cache_interface import (
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py
+16
-9
@@ -8,6 +8,7 @@ from typing import cast
8
from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
9
chunk_hashes_for_block_size,
10
)
11
+from vllm.utils.math_utils import cdiv
12
from vllm.v1.core.block_pool import BlockPool
13
from vllm.v1.core.kv_cache_utils import (
14
BlockHash,
@@ -173,19 +174,21 @@ class MooncakeStoreCoordinator:
174
def store_mask(
175
self,
176
aligned_token_len: int,
177
+ start_token: int = 0,
178
num_prompt_tokens: int | None = None,
179
) -> tuple[list[bool] | None, ...]:
178
- """Per-group store masks.
180
+ """Per-group store masks for the suffix starting at ``start_token``.
181
180
- ``mask[g][i]`` is True iff chunk ``i`` of group ``g`` should be
181
- written to the store so a future cache hit can consume it. ``None`` is
182
- the all-True sentinel.
182
+ ``mask[g][i]`` is True iff the i-th chunk of group ``g`` *after*
183
+ ``start_token`` should be written to the store so a future cache hit
184
+ can consume it. ``None`` is the all-True sentinel for the suffix.
185
186
Reuses the engine's ``SingleTypeKVCacheManager.reachable_block_mask``
187
so the store retains exactly the blocks the local prefix cache would.
188
"""
189
return self._reachable_masks(
190
aligned_token_len,
191
+ start_token,
192
retention_interval=self.retention_interval,
193
num_prompt_tokens=num_prompt_tokens,
194
)
@@ -202,6 +205,7 @@ class MooncakeStoreCoordinator:
205
"""
206
return self._reachable_masks(
207
aligned_token_len,
208
+ 0,
209
retention_interval=None,
210
num_prompt_tokens=None,
211
)
@@ -209,6 +213,7 @@ class MooncakeStoreCoordinator:
213
def _reachable_masks(
214
self,
215
aligned_token_len: int,
216
+ start_token: int,
217
*,
218
retention_interval: int | None,
219
num_prompt_tokens: int | None,
@@ -220,20 +225,22 @@ class MooncakeStoreCoordinator:
225
masks: list[list[bool] | None] = []
226
for g_idx, g in enumerate(self.kv_cache_groups):
227
spec = _unwrap_spec(g.kv_cache_spec)
223
- num_chunks = aligned_token_len // spec.block_size
228
+ end_chunk = aligned_token_len // spec.block_size
229
+ start_chunk = min(end_chunk, max(0, cdiv(start_token, spec.block_size)))
230
manager_cls = KVCacheSpecRegistry.get_manager_class(spec)
231
assert manager_cls is not None
232
+ use_eagle = g_idx in self.eagle_group_ids
233
mask = manager_cls.reachable_block_mask(
227
- start_block=0,
228
- end_block=num_chunks,
234
+ start_block=start_chunk,
235
+ end_block=end_chunk,
236
alignment_tokens=self.lcm_block_size,
237
kv_cache_spec=spec,
231
- use_eagle=g_idx in self.eagle_group_ids,
238
+ use_eagle=use_eagle,
239
retention_interval=retention_interval,
240
num_prompt_tokens=num_prompt_tokens,
241
)
242
if mask is not None:
236
- assert len(mask) == num_chunks
243
+ assert len(mask) == end_chunk - start_chunk
244
masks.append(mask)
245
return tuple(masks)
246
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py
+31
-11
@@ -180,9 +180,10 @@ class ChunkedTokenDatabase:
180
)
181
self.kv_caches_base_addr: list[int] = []
182
self.block_len: list[int] = []
183
+ self._key_prefix = PoolKey.build_prefix(metadata)
184
184
- def _make_key_by_hash(self, chunk_hash: str) -> PoolKey:
185
- return PoolKey(self.metadata, chunk_hash)
185
+ def key_for(self, chunk_hash: BlockHash) -> str:
186
+ return PoolKey.build_key_string(self._key_prefix, chunk_hash.hex())
187
188
def set_kv_caches_base_addr(self, kv_caches_base_addr: list[int]):
189
self.kv_caches_base_addr = kv_caches_base_addr
@@ -215,8 +216,17 @@ class ChunkedTokenDatabase:
216
token_len: int,
217
block_hashes: list[BlockHash],
218
mask_num: int = 0,
218
- ) -> Iterable[tuple[int, int, PoolKey]]:
219
- """Process tokens and yield (start_idx, end_idx, pool_key) tuples.
219
+ *,
220
+ chunk_mask: list[bool] | None = None,
221
+ put_step: int = 1,
222
+ put_step_rank: int = 0,
223
+ ) -> Iterable[tuple[int, int, BlockHash]]:
224
+ """Process tokens and yield (start_idx, end_idx, block_hash) tuples.
225
+
226
+ When there are fewer KV heads than TP ranks, chunks are distributed
227
+ across TP ranks to avoid duplicate load/store. The assignment keys off
228
+ the absolute ``chunk_id`` so a given chunk always lands on the same
229
+ rank regardless of where the processed suffix begins.
230
231
Args:
232
token_len: Total number of tokens.
@@ -224,20 +234,30 @@ class ChunkedTokenDatabase:
234
When ``block_size > hash_block_size`` each group's ``block_size`` chunk
235
is keyed by its last sub-hash via ``chunk_hashes_for_block_size``.
236
mask_num: Number of tokens to skip from the beginning.
237
+ chunk_mask: Optional mask relative to the first chunk after
238
+ ``mask_num``. False entries are skipped before hash access.
239
+ put_step: Stride for distributing chunks across ranks.
240
+ put_step_rank: ``chunk_id % put_step`` value this rank stores.
241
"""
242
+ assert put_step > 0
243
if not block_hashes:
244
return
230
- chunk_hashes: Iterable[BlockHash] = chunk_hashes_for_block_size(
245
+ chunk_hashes: Sequence[BlockHash] = chunk_hashes_for_block_size(
246
block_hashes, self.hash_block_size, self.block_size
247
)
233
- for chunk_id, h in enumerate(chunk_hashes):
248
+ start_chunk = max(0, cdiv(mask_num, self.block_size))
249
+ max_chunks = min(len(chunk_hashes), cdiv(token_len, self.block_size))
250
+ if chunk_mask is not None:
251
+ max_chunks = min(max_chunks, start_chunk + len(chunk_mask))
252
+ for chunk_id in range(start_chunk, max_chunks):
253
+ if chunk_mask is not None and not chunk_mask[chunk_id - start_chunk]:
254
+ continue
255
+ if chunk_id % put_step != put_step_rank:
256
+ continue
257
+ h = chunk_hashes[chunk_id]
258
start_idx = chunk_id * self.block_size
235
- if start_idx >= token_len:
236
- break
259
end_idx = min(start_idx + self.block_size, token_len)
238
- if start_idx < mask_num:
239
- continue
240
- yield start_idx, end_idx, self._make_key_by_hash(h.hex())
260
+ yield start_idx, end_idx, h
261
262
263
@dataclass
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py
+57
-33
@@ -473,6 +473,10 @@ class KVCacheStoreSendingThread(KVTransferThread):
473
self._store_pressure_active = False
474
self._skip_store_requests: set[str] = set()
475
476
+ # Per-request high-water mark of tokens actually persisted; the next
477
+ # batch resumes here, so pressure-skipped or failed ranges are retried.
478
+ self._saved_offset: dict[str, int] = {}
479
+
480
def add_stored_request(self, req_id: str):
481
with self.done_task_lock:
482
self.stored_requests[req_id] += 1
@@ -487,6 +491,13 @@ class KVCacheStoreSendingThread(KVTransferThread):
491
if req_id in self.stored_requests:
492
del self.stored_requests[req_id]
493
self._skip_store_requests.discard(req_id)
494
+ self._saved_offset.pop(req_id, None)
495
+
496
+ def _record_saved(self, req_id: str, token_len: int) -> None:
497
+ # Guard on liveness so a concurrent finish/preempt pop isn't recreated.
498
+ with self.done_task_lock:
499
+ if req_id in self.stored_requests:
500
+ self._saved_offset[req_id] = token_len
501
502
def _should_skip_request(self, req_id: str) -> bool:
503
with self.done_task_lock:
@@ -526,6 +537,7 @@ class KVCacheStoreSendingThread(KVTransferThread):
537
try:
538
if token_len == 0:
539
return
540
+
541
if self._should_skip_request(req_id):
542
logger.debug(
543
"Skipping Mooncake store for request %s while CPU/disk "
@@ -534,40 +546,42 @@ class KVCacheStoreSendingThread(KVTransferThread):
546
)
547
return
548
549
+ # Resume from where this rank left off; only the new suffix is saved.
550
+ save_start = self._saved_offset.get(req_id, 0)
551
+
552
# Within each lcm region only per-spec relevant chunks are loaded
553
# (e.g., SWA or linear attn), so mask out irrelevant chunks
554
store_masks = self.coord.store_mask(
540
- token_len, num_prompt_tokens=req_meta.num_prompt_tokens
555
+ token_len,
556
+ save_start,
557
+ num_prompt_tokens=req_meta.num_prompt_tokens,
558
)
559
+
560
starts: list[int] = []
561
ends: list[int] = []
562
keys: list[str] = []
545
- block_hashes: list[BlockHash] = []
563
+ kv_event_block_hashes: list[BlockHash] = []
564
group_indices: list[int] = []
565
for g_idx, db in enumerate(self.token_databases):
548
- mask = store_masks[g_idx]
549
- for chunk_idx, (start, end, key) in enumerate(
550
- db.process_tokens(token_len, req_meta.block_hashes)
566
+ # Rotate the stride phase per group to balance load across ranks.
567
+ put_step_rank = (self.tp_rank + g_idx) % self.put_step
568
+ for start, end, block_hash in db.process_tokens(
569
+ token_len,
570
+ req_meta.block_hashes,
571
+ mask_num=save_start,
572
+ chunk_mask=store_masks[g_idx],
573
+ put_step=self.put_step,
574
+ put_step_rank=put_step_rank,
575
):
552
- if mask is not None and (
553
- chunk_idx >= len(mask) or not mask[chunk_idx]
554
- ):
555
- continue
576
starts.append(start)
577
ends.append(end)
558
- keys.append(key.to_string())
559
- block_hashes.append(BlockHash(bytes.fromhex(key.chunk_hash)))
578
+ keys.append(db.key_for(block_hash))
579
+ if self.enable_kv_event:
580
+ kv_event_block_hashes.append(block_hash)
581
group_indices.append(g_idx)
582
562
- # Apply put_step striding for TP
563
- sl = slice(self.tp_rank % self.put_step, None, self.put_step)
564
- starts = starts[sl]
565
- ends = ends[sl]
566
- keys = keys[sl]
567
- block_hashes = block_hashes[sl]
568
- group_indices = group_indices[sl]
569
-
583
if not keys:
584
+ self._record_saved(req_id, token_len)
585
return
586
587
# Check which blocks already exist (dedup)
@@ -593,13 +607,18 @@ class KVCacheStoreSendingThread(KVTransferThread):
607
]
608
609
if not missing_indices:
610
+ self._record_saved(req_id, token_len)
611
return
612
598
- starts = [starts[i] for i in missing_indices]
599
- ends = [ends[i] for i in missing_indices]
600
- keys = [keys[i] for i in missing_indices]
601
- block_hashes = [block_hashes[i] for i in missing_indices]
602
- group_indices = [group_indices[i] for i in missing_indices]
613
+ if len(missing_indices) != len(keys):
614
+ starts = [starts[i] for i in missing_indices]
615
+ ends = [ends[i] for i in missing_indices]
616
+ keys = [keys[i] for i in missing_indices]
617
+ if self.enable_kv_event:
618
+ kv_event_block_hashes = [
619
+ kv_event_block_hashes[i] for i in missing_indices
620
+ ]
621
+ group_indices = [group_indices[i] for i in missing_indices]
622
623
logger.debug(
624
"Storing KV cache for %d blocks (groups=%s) for request %s",
@@ -612,8 +631,11 @@ class KVCacheStoreSendingThread(KVTransferThread):
631
sizes: list[list[int]] = []
632
stored_events: list[BlockStored] = []
633
# parent_block_hash chains live within a group, not across.
615
- prev_key_per_group: dict[int, Any] = {}
616
- new_block_hashes = [maybe_convert_block_hash(bh) for bh in block_hashes]
634
+ if self.enable_kv_event:
635
+ prev_key_per_group: dict[int, Any] = {}
636
+ new_block_hashes = [
637
+ maybe_convert_block_hash(bh) for bh in kv_event_block_hashes
638
+ ]
639
640
for idx, (s, e, g_idx) in enumerate(
641
zip(starts, ends, group_indices, strict=True)
@@ -687,11 +709,13 @@ class KVCacheStoreSendingThread(KVTransferThread):
709
"batch succeeds",
710
req_id,
711
)
690
- elif self._clear_store_pressure():
691
- logger.info(
692
- "Mooncake CPU/disk offloading pressure cleared after a "
693
- "successful store batch"
694
- )
712
+ else:
713
+ self._record_saved(req_id, token_len)
714
+ if self._clear_store_pressure():
715
+ logger.info(
716
+ "Mooncake CPU/disk offloading pressure cleared "
717
+ "after a successful store batch"
718
+ )
719
except Exception as e:
720
self._record_operation(
721
"save_put",
@@ -775,7 +799,7 @@ class KVCacheStoreRecvingThread(KVTransferThread):
799
block_id_list: list[int] = []
800
for g_idx, db in enumerate(self.token_databases):
801
mask = load_mask_per_group[g_idx]
778
- for start, end, key in db.process_tokens(
802
+ for start, end, block_hash in db.process_tokens(
803
token_len, req_meta.block_hashes, mask_num
804
):
805
chunk_idx = start // db.block_size
@@ -784,7 +808,7 @@ class KVCacheStoreRecvingThread(KVTransferThread):
808
addr, size, block_id = db.prepare_value(
809
start, end, req_meta.block_ids[g_idx]
810
)
787
- key_list.append(key.to_string())
811
+ key_list.append(db.key_for(block_hash))
812
addr_list.append(addr)
813
size_list.append(size)
814
block_id_list.append(block_id)