[BugFix][Mooncake] Fix Mooncake saves from sparse Mamba block tables (#51362)
Signed-off-by: ZeldaHuang <zelda.huanghuang@gmail.com> Signed-off-by: Ziming Huang <zelda.huanghuang@gmail.com>
Ziming Huang committed
Aug 20, 2026 at 18:05 UTC
5b1e7a812b1358125b6d5291c51bbc65264f044d
3 files changed
+99
-21
tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py
+2
-1
@@ -225,7 +225,8 @@ def test_e2e_swa_plus_full_save_then_lookup_hits():
225
save_req = ReqMeta(
226
req_id="r0",
227
token_len_chunk=64,
228
- block_ids=([0, 1, 2, 3], [0, 1, 2, 3]),
228
+ # Block 0 is reserved as NULL_BLOCK_ID by the production block pool.
229
+ block_ids=([1, 2, 3, 4], [1, 2, 3, 4]),
230
block_hashes=hs,
231
can_save=True,
232
store_job_id=1,
tests/v1/kv_connector/unit/test_mooncake_store_worker.py
+84
-20
@@ -36,6 +36,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
36
from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.metrics import (
37
MooncakeStoreConnectorStats,
38
)
39
+from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
40
from vllm.v1.core.kv_cache_utils import BlockHash, maybe_convert_block_hash
41
42
@@ -173,7 +174,7 @@ def _make_store_req(req_id: str, block_hashes: list[bytes]) -> ReqMeta:
174
return ReqMeta(
175
req_id=req_id,
176
token_len_chunk=32,
176
- block_ids=([0, 1],),
177
+ block_ids=([1, 2],),
178
block_hashes=block_hashes,
179
can_save=True,
180
)
@@ -183,7 +184,7 @@ def _make_multi_group_store_req(req_id: str, block_hashes: list[bytes]) -> ReqMe
184
return ReqMeta(
185
req_id=req_id,
186
token_len_chunk=32,
186
- block_ids=([0, 1], [2, 3]),
187
+ block_ids=([1, 2], [3, 4]),
188
block_hashes=block_hashes,
189
can_save=True,
190
)
@@ -517,7 +518,7 @@ def test_store_sending_thread_delta_saves_only_new_full_attention_chunks():
518
ReqMeta(
519
req_id="req-a",
520
token_len_chunk=64,
520
- block_ids=([0, 1, 2, 3],),
521
+ block_ids=([1, 2, 3, 4],),
522
block_hashes=[b"a0", b"a1", b"a2", b"a3"],
523
can_save=True,
524
),
@@ -543,7 +544,7 @@ def test_store_sending_thread_delta_strides_with_local_phase():
544
ReqMeta(
545
req_id="req-a",
546
token_len_chunk=64,
546
- block_ids=([0, 1, 2, 3],),
547
+ block_ids=([1, 2, 3, 4],),
548
block_hashes=[b"a0", b"a1", b"a2", b"a3"],
549
can_save=True,
550
),
@@ -569,7 +570,7 @@ def test_tp_sharded_group_saves_every_block_on_every_rank():
570
ReqMeta(
571
req_id="req-a",
572
token_len_chunk=64,
572
- block_ids=([0, 1, 2, 3],),
573
+ block_ids=([1, 2, 3, 4],),
574
block_hashes=[b"a0", b"a1", b"a2", b"a3"],
575
can_save=True,
576
),
@@ -595,7 +596,7 @@ def test_store_sending_thread_retries_skipped_range_after_pressure():
596
ReqMeta(
597
req_id="req-a",
598
token_len_chunk=16,
598
- block_ids=([0],),
599
+ block_ids=([1],),
600
block_hashes=[b"a0"],
601
can_save=True,
602
),
@@ -614,7 +615,7 @@ def test_store_sending_thread_retries_skipped_range_after_pressure():
615
ReqMeta(
616
req_id="req-a",
617
token_len_chunk=64,
617
- block_ids=([0, 1, 2, 3],),
618
+ block_ids=([1, 2, 3, 4],),
619
block_hashes=[b"a0", b"a1", b"a2", b"a3"],
620
can_save=True,
621
),
@@ -686,6 +687,69 @@ def test_partial_tail_offload_skips_null_source_blocks():
687
assert addrs == [[0x1000 + 2 * 256], [0x1000 + 3 * 256]]
688
689
690
+def test_store_sending_thread_skips_null_sparse_group_blocks():
691
+ from vllm.v1.kv_cache_interface import (
692
+ FullAttentionSpec,
693
+ KVCacheGroupSpec,
694
+ MambaSpec,
695
+ )
696
+
697
+ store = MagicMock()
698
+ store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
699
+ store.batch_put_from_multi_buffers.side_effect = (
700
+ lambda keys, addrs, sizes, replicate_config: [256] * len(keys)
701
+ )
702
+ full = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None)
703
+ mamba = MambaSpec(
704
+ block_size=16,
705
+ shapes=((1, 1),),
706
+ dtypes=(torch.float32,),
707
+ mamba_cache_mode="align",
708
+ )
709
+ coord = mooncake_store_worker.MooncakeStoreCoordinator(
710
+ [
711
+ KVCacheGroupSpec(["full"], full),
712
+ KVCacheGroupSpec(["mamba"], mamba),
713
+ ],
714
+ scheduler_block_size=16,
715
+ hash_block_size=16,
716
+ )
717
+
718
+ db_full = ChunkedTokenDatabase(
719
+ KeyMetadata("test-model", 0, 0, 0, 0, group_id=0),
720
+ block_size=16,
721
+ )
722
+ db_full.set_kv_caches_base_addr([0x1000])
723
+ db_full.set_block_len([256])
724
+ db_mamba = ChunkedTokenDatabase(
725
+ KeyMetadata("test-model", 0, 0, 0, 0, group_id=1),
726
+ block_size=16,
727
+ )
728
+ db_mamba.set_kv_caches_base_addr([0x2000])
729
+ db_mamba.set_block_len([512])
730
+
731
+ thread = _make_store_sending_thread(
732
+ store,
733
+ coord=coord,
734
+ token_databases=[db_full, db_mamba],
735
+ )
736
+ _run_store_req(
737
+ thread,
738
+ ReqMeta(
739
+ req_id="req-a",
740
+ token_len_chunk=48,
741
+ block_ids=([1, 2, 3], [NULL_BLOCK_ID, 4, NULL_BLOCK_ID]),
742
+ block_hashes=[b"a0", b"a1", b"a2"],
743
+ can_save=True,
744
+ ),
745
+ )
746
+
747
+ keys, addrs, _, _ = store.batch_put_from_multi_buffers.call_args.args
748
+ mamba_keys = [key for key in keys if "@group:1" in key]
749
+ assert [key.rsplit("@", 1)[-1] for key in mamba_keys] == [b"a1".hex()]
750
+ assert all(addr[0] >= 0x2000 for key, addr in zip(keys, addrs) if "@group:1" in key)
751
+
752
+
753
def test_partial_tail_offload_replaces_stale_group_ids_after_filtering():
754
store = MagicMock()
755
store.batch_is_exist.return_value = [1, 0, 0]
@@ -752,7 +816,7 @@ def test_store_sending_thread_delta_start_rank_saves_second_local_chunk():
816
ReqMeta(
817
req_id="req-a",
818
token_len_chunk=64,
755
- block_ids=([0, 1, 2, 3],),
819
+ block_ids=([1, 2, 3, 4],),
820
block_hashes=[b"a0", b"a1", b"a2", b"a3"],
821
can_save=True,
822
),
@@ -859,7 +923,7 @@ def test_store_sending_thread_prepares_missing_chunks_once_per_group():
923
ReqMeta(
924
req_id="req-a",
925
token_len_chunk=48,
862
- block_ids=([0, 1, 2], [2, 1, 0]),
926
+ block_ids=([1, 2, 3], [3, 2, 1]),
927
block_hashes=[b"a0", b"a1", b"a2"],
928
can_save=True,
929
),
@@ -867,8 +931,8 @@ def test_store_sending_thread_prepares_missing_chunks_once_per_group():
931
932
db0.prepare_value.assert_not_called()
933
db1.prepare_value.assert_not_called()
870
- db0.prepare_values.assert_called_once_with([(0, 16), (32, 48)], [0, 1, 2])
871
- db1.prepare_values.assert_called_once_with([(16, 32), (32, 48)], [2, 1, 0])
934
+ db0.prepare_values.assert_called_once_with([(0, 16), (32, 48)], [1, 2, 3])
935
+ db1.prepare_values.assert_called_once_with([(16, 32), (32, 48)], [3, 2, 1])
936
937
keys, addrs, sizes, _ = store.batch_put_from_multi_buffers.call_args.args
938
assert [key.rsplit("@", 1)[-1] for key in keys] == [
@@ -877,7 +941,7 @@ def test_store_sending_thread_prepares_missing_chunks_once_per_group():
941
"6131",
942
"6132",
943
]
880
- assert addrs == [[0x1000], [0x1200], [0x2200], [0x2000]]
944
+ assert addrs == [[0x1100], [0x1300], [0x2400], [0x2200]]
945
assert sizes == [[256], [256], [512], [512]]
946
947
@@ -1250,7 +1314,7 @@ def test_store_sending_thread_multiple_segments_share_logical_group_id():
1314
"test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6130",
1315
"test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6131",
1316
]
1253
- assert addrs == [[0x1000, 0x2000], [0x1100, 0x2100]]
1317
+ assert addrs == [[0x1100, 0x2100], [0x1200, 0x2200]]
1318
assert sizes == [[256, 256], [256, 256]]
1319
assert config.group_ids == [
1320
"vllm-mooncake-store:test-model@6130",
@@ -1301,7 +1365,7 @@ def test_store_sending_thread_group_ids_share_across_kv_cache_groups():
1365
"test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:1@6130",
1366
"test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:1@6131",
1367
]
1304
- assert addrs == [[0x1000], [0x1100], [0x3200], [0x3300]]
1368
+ assert addrs == [[0x1100], [0x1200], [0x3300], [0x3400]]
1369
assert sizes == [[256], [256], [256], [256]]
1370
# Different vLLM KV cache groups for the same prefix chunk share the
1371
# same Mooncake lifecycle group id.
@@ -1935,7 +1999,7 @@ def test_store_sending_thread_clamps_token_len_to_lcm():
1999
ReqMeta(
2000
req_id="r0",
2001
token_len_chunk=33,
1938
- block_ids=([0, 1, 2],),
2002
+ block_ids=([1, 2, 3],),
2003
block_hashes=[b"a0", b"a1", b"a2"],
2004
can_save=True,
2005
),
@@ -1974,7 +2038,7 @@ def test_store_sending_thread_skips_when_token_len_below_lcm():
2038
ReqMeta(
2039
req_id="r0",
2040
token_len_chunk=32,
1977
- block_ids=([0, 1],),
2041
+ block_ids=([1, 2],),
2042
block_hashes=[b"a0", b"a1"],
2043
can_save=True,
2044
),
@@ -2050,7 +2114,7 @@ def test_store_sending_thread_only_stores_swa_blocks_in_window():
2114
ReqMeta(
2115
req_id="r0",
2116
token_len_chunk=64,
2053
- block_ids=([0, 1], list(range(8))),
2117
+ block_ids=([1, 2], list(range(1, 9))),
2118
block_hashes=hs,
2119
can_save=True,
2120
),
@@ -2126,7 +2190,7 @@ def test_store_sending_thread_delta_saves_only_new_swa_boundary_chunks():
2190
ReqMeta(
2191
req_id="r0",
2192
token_len_chunk=64,
2129
- block_ids=([0, 1], list(range(8))),
2193
+ block_ids=([1, 2], list(range(1, 9))),
2194
block_hashes=hs,
2195
can_save=True,
2196
),
@@ -2196,7 +2260,7 @@ def test_store_sending_thread_kv_events_use_group_chunk_metadata():
2260
ReqMeta(
2261
req_id="r0",
2262
token_len_chunk=32,
2199
- block_ids=([0], list(range(4))),
2263
+ block_ids=([1], list(range(1, 5))),
2264
block_hashes=hs,
2265
can_save=True,
2266
token_ids=list(range(32)),
@@ -2222,7 +2286,7 @@ def _make_event_store_req(token_len: int, token_ids_start: int = 0) -> ReqMeta:
2286
return ReqMeta(
2287
req_id="r0",
2288
token_len_chunk=token_len,
2225
- block_ids=(list(range(num_blocks)),),
2289
+ block_ids=(list(range(1, num_blocks + 1)),),
2290
block_hashes=[f"a{i}".encode() for i in range(num_blocks)],
2291
can_save=True,
2292
token_ids=list(range(token_ids_start, token_len)),
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py
+13
@@ -889,6 +889,19 @@ class KVCacheStoreSendingThread(KVTransferThread):
889
put_step=put_step,
890
put_step_rank=put_step_rank,
891
):
892
+ block_idx = start // db.block_size
893
+ group_blocks = block_ids_per_group[g_idx]
894
+ if block_idx >= len(group_blocks) or (
895
+ group_blocks[block_idx] == NULL_BLOCK_ID
896
+ ):
897
+ logger.debug(
898
+ "Skipping unavailable Mooncake store source block "
899
+ "(req=%s, group=%d, block=%d)",
900
+ req_id,
901
+ g_idx,
902
+ block_idx,
903
+ )
904
+ continue
905
starts.append(start)
906
ends.append(end)
907
keys.append(db.key_for(block_hash))