@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 7fd9cc036e

[KV Connector] Support MooncakeStore with hybrid DCP prefix caching (#53324)

Signed-off-by: wzhao18 <wzhao18.sz@gmail.com> Signed-off-by: wzhao18 <weizha@nvidia.com> Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai>

Wei Zhao committed Aug 28, 2026 at 22:10 UTC 7fd9cc036e3c0cbf0dfa4d1a9bbd755ed4922031

14 files changed +560 -85

tests/v1/core/test_kv_cache_utils.py

+35

@@ -2334,6 +2334,41 @@ def test_hidden_state_group_preserves_hybrid_prefix_cache_granularity():

2334 ) == (544, 136)

2335

2336

2337 +def test_resolve_dcp_kv_block_size_unwraps_uniform_type_specs():

2338 + attention = FullAttentionSpec(

2339 + block_size=16,

2340 + num_kv_heads=1,

2341 + head_size=64,

2342 + dtype=torch.float16,

2343 + )

2344 + mamba = MambaSpec(

2345 + block_size=16,

2346 + shapes=((1, 1),),

2347 + dtypes=(torch.float16,),

2348 + )

2349 + wrapped_attention = UniformTypeKVCacheSpecs(

2350 + block_size=16,

2351 + kv_cache_specs={"attention.0": attention, "attention.1": attention},

2352 + )

2353 + wrapped_mamba = UniformTypeKVCacheSpecs(

2354 + block_size=16,

2355 + kv_cache_specs={"mamba.0": mamba, "mamba.1": mamba},

2356 + )

2357 +

2358 + assert kv_cache_utils.resolve_dcp_kv_block_size(attention, 4) == 64

2359 + assert kv_cache_utils.resolve_dcp_kv_block_size(wrapped_attention, 4) == 64

2360 + assert kv_cache_utils.resolve_dcp_kv_block_size(mamba, 4) == 16

2361 + assert kv_cache_utils.resolve_dcp_kv_block_size(wrapped_mamba, 4) == 16

2362 +

2363 + scaled_attention = kv_cache_utils.resolve_dcp_kv_cache_spec(wrapped_attention, 4)

2364 + assert scaled_attention.block_size == 64

2365 + assert isinstance(scaled_attention, UniformTypeKVCacheSpecs)

2366 + assert all(

2367 + spec.block_size == 64 for spec in scaled_attention.kv_cache_specs.values()

2368 + )

2369 + assert kv_cache_utils.resolve_dcp_kv_cache_spec(wrapped_mamba, 4) is wrapped_mamba

2370 +

2371 +

2372 def test_multi_run_layer_compact_strides_place_hoisted_heads():

2373 """A layer-compact run region is its own dense allocation: under LHBNC the head

2374 groups sit between the layers and the blocks, so a run's block stride is one head

tests/v1/kv_connector/unit/test_mooncake_store_connector.py

+27 -3

@@ -19,7 +19,9 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store import (

19 worker,

20 )

21 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (

22 + MooncakeLookupResult,

23 MooncakeStoreConnectorMetadata,

24 + TailKeyBoundary,

25 )

26 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.metrics import (

27 MooncakeStoreConnectorStats,

@@ -121,10 +123,14 @@ def test_worker_methods_delegate_to_store_worker():

123 connector.bind_connector_metadata(metadata)

124

125 connector.register_kv_caches(kv_caches)

126 + connector.start_load_kv(MagicMock())

127 + connector.wait_for_save()

128 result = connector.get_finished(finished_req_ids)

129 invalid_block_ids = connector.get_block_ids_with_load_errors()

130

131 worker.register_kv_caches.assert_called_once_with(kv_caches)

132 + worker.start_load_kv.assert_called_once_with(metadata)

133 + worker.wait_for_save.assert_called_once_with(metadata)

134 worker.get_finished.assert_called_once_with(finished_req_ids, metadata)

135 assert result == ({"req-1"}, {"req-2"})

136 worker.get_block_ids_with_load_errors.assert_called_once_with()

@@ -359,7 +365,9 @@ def test_lookup_key_client_lookup_prepends_typed_tag():

365

366 # Blocking lookup (non_block defaults to False) runs on the executor and

367 # returns the resolved hit length.

362 - assert client.lookup("req0", num_tokens=128, block_hashes=[]) == 5

368 + result = client.lookup("req0", num_tokens=128, block_hashes=[])

369 + assert result is not None

370 + assert result.hit_length == 5

371

372 sent_frames = fake_socket.send_multipart.call_args[0][0]

373 assert sent_frames[0] == protocol.LOOKUP_MSG

@@ -394,7 +402,7 @@ def _poll_lookup(client, req_id, num_tokens=128, block_hashes=(), timeout=5.0):

402 while time.monotonic() < deadline:

403 result = client.lookup(req_id, num_tokens, list(block_hashes), non_block=True)

404 if result is not None:

397 - return result

405 + return result.hit_length

406 time.sleep(0.005)

407 return None

408

@@ -502,11 +510,13 @@ def test_get_num_new_matched_tokens_async_defers_then_reports():

510

511 # Lookup ready with a hit -> report need_to_allocate + async-load flag.

512 hit = 3 * block_size

505 - mock_client.lookup.return_value = hit

513 + boundary = TailKeyBoundary(group_id=0, num_tokens=4 * block_size)

514 + mock_client.lookup.return_value = MooncakeLookupResult(hit, (boundary,))

515 need, load_async = sched.get_num_new_matched_tokens(request, 0)

516 assert need == hit

517 assert load_async == sched.load_async

518 assert sched.load_specs["r1"].kvpool_cached_tokens == hit

519 + assert sched.load_specs["r1"].tail_key_boundaries == (boundary,)

520

521

522 def test_protocol_tags_are_distinct_and_non_empty():

@@ -519,6 +529,20 @@ def test_protocol_tags_are_distinct_and_non_empty():

529 assert protocol.RESP_OK != protocol.RESP_ERR

530

531

532 +def test_lookup_response_round_trip_preserves_tail_keys():

533 + result = MooncakeLookupResult(

534 + hit_length=20,

535 + tail_key_boundaries=(

536 + TailKeyBoundary(group_id=0, num_tokens=24),

537 + TailKeyBoundary(group_id=1, num_tokens=20),

538 + ),

539 + )

540 +

541 + assert protocol.decode_lookup_response(protocol.encode_lookup_response(result)) == (

542 + result

543 + )

544 +

545 +

546 def test_scheduler_reset_connector_cache_invokes_connector_reset():

547 """Cascade test: Scheduler.reset_prefix_cache(reset_connector=True)

548 cascades into MooncakeStoreConnector.reset_cache without dragging in

tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py

+4 -4

@@ -254,10 +254,10 @@ def test_e2e_swa_plus_full_save_then_lookup_hits():

254 worker.store = store

255

256 # Both groups stored all 4 blocks -> full hit.

257 - assert worker.lookup(num_tokens=65, block_hashes=hs) == 64

257 + assert worker.lookup(num_tokens=65, block_hashes=hs).hit_length == 64

258 # Exact-multiple prompt: the full hit is re-derived one block lower,

259 # where both groups' stored blocks still cover the SWA window.

260 - assert worker.lookup(num_tokens=64, block_hashes=hs) == 48

260 + assert worker.lookup(num_tokens=64, block_hashes=hs).hit_length == 48

261

262 # Evict SWA's first two blocks (outside its window of 32 tokens = 2 blocks).

263 swa_keys_outside_window = [

@@ -270,12 +270,12 @@ def test_e2e_swa_plus_full_save_then_lookup_hits():

270

271 # SWA window=32 -> only last 2 blocks must be present in SWA group.

272 # Full has all 4. Coordinator should still return 64.

273 - assert worker.lookup(num_tokens=65, block_hashes=hs) == 64

273 + assert worker.lookup(num_tokens=65, block_hashes=hs).hit_length == 64

274 # Exact-multiple prompt after eviction: the boundary one block lower

275 # needs SWA block 1, which is gone — no usable stored boundary remains

276 # (the pre-fix arithmetic clamp would have returned 48 and livelocked

277 # on load failure -> recompute -> same lookup).

278 - assert worker.lookup(num_tokens=64, block_hashes=hs) == 0

278 + assert worker.lookup(num_tokens=64, block_hashes=hs).hit_length == 0

279

280

281 def test_recv_skips_swa_blocks_before_window():

tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py

+3 -2

@@ -7,6 +7,7 @@ import pytest

8 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (

9 LoadSpec,

10 + MooncakeLookupResult,

11 MooncakeStoreWorkerMetadata,

12 ReqMeta,

13 RequestTracker,

@@ -702,9 +703,9 @@ class _StubLookupClient:

703 num_tokens: int,

704 block_hashes: list[bytes],

705 non_block: bool = False,

705 - ) -> int:

706 + ) -> MooncakeLookupResult:

707 self.num_tokens.append(num_tokens)

707 - return self._hit_tokens

708 + return MooncakeLookupResult(self._hit_tokens)

709

710

711 def test_full_external_hit_keeps_kvpool_cached_tokens_block_aligned():

tests/v1/kv_connector/unit/test_mooncake_store_worker.py

+249 -13

@@ -31,8 +31,10 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (

31 ChunkedTokenDatabase,

32 KeyMetadata,

33 LoadSpec,

34 + MooncakeLookupResult,

35 PoolKey,

36 ReqMeta,

37 + TailKeyBoundary,

38 )

39 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.metrics import (

40 MooncakeStoreConnectorStats,

@@ -150,6 +152,7 @@ def _make_load_req(

152 *,

153 token_len: int,

154 vllm_cached_tokens: int = 0,

155 + tail_key_boundaries: tuple[TailKeyBoundary, ...] = (),

156 ) -> ReqMeta:

157 return ReqMeta(

158 req_id=req_id,

@@ -161,6 +164,7 @@ def _make_load_req(

164 kvpool_cached_tokens=token_len,

165 can_load=True,

166 token_len=token_len,

167 + tail_key_boundaries=tail_key_boundaries,

168 ),

169 )

170

@@ -305,15 +309,16 @@ def _patch_worker_runtime(

309 local_ip: str = "10.0.0.7",

310 tp_rank: int = 0,

311 tp_size: int = 1,

312 + pcp_size: int = 1,

313 dcp_size: int = 1,

314 ) -> None:

310 - single_rank_group = SimpleNamespace(world_size=1, rank_in_group=0)

315 + pcp_group = SimpleNamespace(world_size=pcp_size, rank_in_group=0)

316 # DCP groups are contiguous splits of the TP group (see

317 # parallel_state.py), so dcp_rank == tp_rank % dcp_size.

318 dcp_group = SimpleNamespace(world_size=dcp_size, rank_in_group=tp_rank % dcp_size)

319 monkeypatch.setattr(worker, "get_tensor_model_parallel_rank", lambda: tp_rank)

320 monkeypatch.setattr(worker, "get_tensor_model_parallel_world_size", lambda: tp_size)

316 - monkeypatch.setattr(worker, "get_pcp_group", lambda: single_rank_group)

321 + monkeypatch.setattr(worker, "get_pcp_group", lambda: pcp_group)

322 monkeypatch.setattr(worker, "get_dcp_group", lambda: dcp_group)

323 monkeypatch.setattr(worker, "get_ip", lambda: local_ip)

324 monkeypatch.setattr(worker, "LookupKeyServer", MagicMock())

@@ -1433,6 +1438,29 @@ def test_recv_thread_uses_single_batch_when_no_disk_offload_budget(monkeypatch):

1438 store.batch_get_replica_desc.assert_not_called()

1439

1440

1441 +def test_recv_thread_keys_chunk_by_lookup_selected_boundary():

1442 + store = MagicMock()

1443 + store.batch_get_into_multi_buffers.return_value = [256, 256]

1444 + thread = _make_store_recving_thread(store)

1445 +

1446 + # 32-token hit over 16-token chunks: chunk 1 would default to the hash at

1447 + # the 32-token boundary (a1); the lookup matched the 48-token one (a2).

1448 + req = _make_load_req(

1449 + "req-a",

1450 + [b"a0", b"a1", b"a2"],

1451 + token_len=32,

1452 + tail_key_boundaries=(TailKeyBoundary(group_id=0, num_tokens=48),),

1453 + )

1454 +

1455 + thread._handle_request(req)

1456 +

1457 + keys = store.batch_get_into_multi_buffers.call_args.args[0]

1458 + assert keys == [

1459 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6130",

1460 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6132",

1461 + ]

1462 +

1463 +

1464 def test_recv_thread_logs_tier_summary_when_enabled(monkeypatch, caplog_vllm):

1465 monkeypatch.setenv("VLLM_MOONCAKE_STORE_TIER_LOG", "1")

1466 caplog_vllm.set_level(logging.INFO, logger=worker.logger.name)

@@ -1799,6 +1827,59 @@ def test_requester_worker_init_builds_replicate_config_for_preferred_segment(

1827 assert w.store_replicate_config.preferred_segment == "10.0.0.7:50053"

1828

1829

1830 +def test_worker_scales_uniform_attention_group_under_dcp(tmp_path, monkeypatch):

1831 + from vllm.v1.kv_cache_interface import (

1832 + KVCacheConfig,

1833 + UniformTypeKVCacheSpecs,

1834 + )

1835 +

1836 + store = MagicMock()

1837 + store.setup.return_value = 0

1838 + _install_fake_mooncake(monkeypatch, store)

1839 + _patch_worker_runtime(monkeypatch, tp_size=4, dcp_size=4)

1840 + monkeypatch.setenv(

1841 + "MOONCAKE_CONFIG_PATH",

1842 + _write_mooncake_config(

1843 + tmp_path,

1844 + {

1845 + "metadata_server": "http://metadata/endpoint",

1846 + "protocol": "tcp",

1847 + "device_name": "",

1848 + "master_server_address": "10.0.0.7:50051",

1849 + },

1850 + ),

1851 + )

1852 + inner_specs = {

1853 + "layer0": FullAttentionSpec(

1854 + block_size=16, num_kv_heads=8, head_size=64, dtype=None

1855 + ),

1856 + "layer1": FullAttentionSpec(

1857 + block_size=16, num_kv_heads=4, head_size=64, dtype=None

1858 + ),

1859 + }

1860 + uniform_spec = UniformTypeKVCacheSpecs(

1861 + block_size=16,

1862 + kv_cache_specs=inner_specs,

1863 + )

1864 + kv_cache_config = KVCacheConfig(

1865 + num_blocks=10,

1866 + kv_cache_tensors=[],

1867 + kv_cache_groups=[KVCacheGroupSpec(list(inner_specs), uniform_spec)],

1868 + )

1869 +

1870 + w = worker.MooncakeStoreWorker(

1871 + _make_vllm_config(decode_context_parallel_size=4),

1872 + kv_cache_config,

1873 + )

1874 +

1875 + scaled_spec = w._kv_cache_groups[0].kv_cache_spec

1876 + assert isinstance(scaled_spec, UniformTypeKVCacheSpecs)

1877 + assert scaled_spec.block_size == 64

1878 + assert all(spec.block_size == 64 for spec in scaled_spec.kv_cache_specs.values())

1879 + assert w.coord.attention_groups[0].spec.block_size == 64

1880 + assert w.token_dbs[0].block_size == 64

1881 +

1882 +

1883 @pytest.mark.parametrize("dcp_size", [1, 4])

1884 def test_worker_put_striding_covers_every_rank_get_namespace(

1885 tmp_path, monkeypatch, dcp_size

@@ -2668,12 +2749,12 @@ def test_lookup_rejects_boundary_missing_one_mamba_shard():

2749 # 33 tokens for two 16-token blocks: the hit stops below the request end,

2750 # so the full-hit re-derivation stays out of the shard accounting.

2751 worker.store.batch_is_exist.side_effect = lambda keys: [1] * len(keys)

2671 - assert worker.lookup(33, [b"h0", b"h1"]) == 32

2752 + assert worker.lookup(33, [b"h0", b"h1"]).hit_length == 32

2753

2754 worker.store.batch_is_exist.side_effect = lambda keys: [

2755 0 if "tp_rank:1" in k and "group:1" in k else 1 for k in keys

2756 ]

2676 - assert worker.lookup(33, [b"h0", b"h1"]) == 0

2757 + assert worker.lookup(33, [b"h0", b"h1"]).hit_length == 0

2758

2759

2760 def test_lookup_requires_all_dcp_rank_namespaces():

@@ -2684,7 +2765,7 @@ def test_lookup_requires_all_dcp_rank_namespaces():

2765 _refresh_group_tp_replication_factors(worker)

2766 worker.store.batch_is_exist.return_value = [1, 1, 0, 1]

2767

2687 - assert worker.lookup(16, [b"a0"]) == 0

2768 + assert worker.lookup(16, [b"a0"]).hit_length == 0

2769 assert worker.store.batch_is_exist.call_args.args[0] == [

2770 "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6130",

2771 "test-model@tp_rank:1@pcp0@dcp1@pp_rank:0@group:0@6130",

@@ -2696,7 +2777,7 @@ def test_lookup_requires_all_dcp_rank_namespaces():

2777 def test_lookup_partial_prefix_returns_first_hit_length():

2778 worker = _make_bare_worker()

2779 worker.store.batch_is_exist.return_value = [1, 1, 0]

2699 - assert worker.lookup(48, [b"a0", b"a1", b"a2"]) == 32

2780 + assert worker.lookup(48, [b"a0", b"a1", b"a2"]).hit_length == 32

2781

2782

2783 def test_lookup_partial_tail_uses_hash_alignment():

@@ -2736,7 +2817,7 @@ def test_lookup_partial_tail_uses_hash_alignment():

2817 _refresh_group_tp_replication_factors(worker)

2818 worker.store.batch_is_exist.return_value = [0, 0, 1, 0, 0, 1]

2819

2739 - assert worker.lookup(13, [b"h0", b"h1", b"h2"]) == 12

2820 + assert worker.lookup(13, [b"h0", b"h1", b"h2"]).hit_length == 12

2821

2822

2823 def test_lookup_full_hit_reuses_existing_boundary():

@@ -2744,7 +2825,10 @@ def test_lookup_full_hit_reuses_existing_boundary():

2825 worker = _make_bare_worker(block_size=16)

2826 worker.store.batch_is_exist.return_value = [1, 1]

2827

2747 - assert worker.lookup(32, [b"h0", b"h1"]) == 16

2828 + assert worker.lookup(32, [b"h0", b"h1"]) == MooncakeLookupResult(

2829 + hit_length=16,

2830 + tail_key_boundaries=(TailKeyBoundary(group_id=0, num_tokens=16),),

2831 + )

2832 assert worker.store.batch_is_exist.call_count == 1

2833

2834

@@ -2765,10 +2849,144 @@ def test_lookup_full_hit_with_eagle_pops_once_not_twice():

2849 # 64-token exact-multiple prompt, all 4 blocks stored: one eagle pop

2850 # gives 48; a spurious re-derivation (anchored at 48) would pop again

2851 # and return 32.

2768 - assert worker.lookup(64, [b"h0", b"h1", b"h2", b"h3"]) == 48

2852 + assert worker.lookup(64, [b"h0", b"h1", b"h2", b"h3"]).hit_length == 48

2853 assert worker.store.batch_is_exist.call_count == 1

2854

2855

2856 +def test_lookup_plan_resolves_group_tail_keys_from_existing_hashes():

2857 + from vllm.v1.kv_cache_interface import (

2858 + FullAttentionSpec,

2859 + KVCacheGroupSpec,

2860 + MambaSpec,

2861 + )

2862 +

2863 + worker = _make_bare_worker(block_size=16)

2864 + full = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None)

2865 + mamba = MambaSpec(

2866 + block_size=16,

2867 + shapes=((1, 1),),

2868 + dtypes=(torch.float32,),

2869 + mamba_cache_mode="align",

2870 + )

2871 + worker._kv_cache_groups = [

2872 + KVCacheGroupSpec(["full"], full),

2873 + KVCacheGroupSpec(["mamba"], mamba),

2874 + ]

2875 + worker.hash_block_size = 4

2876 + worker.token_dbs = [

2877 + ChunkedTokenDatabase(

2878 + KeyMetadata("test-model", 0, 0, 0, 0, group_id=group_id),

2879 + block_size=16,

2880 + hash_block_size=4,

2881 + )

2882 + for group_id in range(2)

2883 + ]

2884 + worker.coord = mooncake_store_worker.MooncakeStoreCoordinator(

2885 + worker._kv_cache_groups,

2886 + scheduler_block_size=16,

2887 + hash_block_size=4,

2888 + use_eagle=True,

2889 + )

2890 + _refresh_group_tp_replication_factors(worker)

2891 + hashes = [BlockHash(f"h{i}".encode()) for i in range(6)]

2892 + present = {

2893 + (0, bytes(hashes[3])),

2894 + (0, bytes(hashes[5])),

2895 + (1, bytes(hashes[4])),

2896 + # A later Mamba state also exists, but its exact hit-boundary hash

2897 + # must win and therefore needs no override.

2898 + (1, bytes(hashes[5])),

2899 + }

2900 +

2901 + def exists(keys):

2902 + return [

2903 + int(

2904 + any(

2905 + f"@group:{group_id}@{block_hash.hex()}" in key

2906 + for group_id, block_hash in present

2907 + )

2908 + )

2909 + for key in keys

2910 + ]

2911 +

2912 + worker.store.batch_is_exist.side_effect = exists

2913 +

2914 + result = worker.lookup(25, hashes)

2915 +

2916 + # The eagle drop trims the hit to 20 tokens, but the block that survives

2917 + # truncation is the one keyed at the 24-token boundary (hashes[5]), which

2918 + # the load path cannot derive from hit_length.

2919 + assert result.hit_length == 20

2920 + assert result.tail_key_boundaries == (

2921 + TailKeyBoundary(group_id=0, num_tokens=24),

2922 + TailKeyBoundary(group_id=1, num_tokens=20),

2923 + )

2924 +

2925 +

2926 +def test_lookup_plan_recovers_tail_key_after_multi_chunk_convergence():

2927 + from vllm.v1.kv_cache_interface import (

2928 + FullAttentionSpec,

2929 + KVCacheGroupSpec,

2930 + MambaSpec,

2931 + )

2932 +

2933 + worker = _make_bare_worker(block_size=16)

2934 + full = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None)

2935 + mamba = MambaSpec(

2936 + block_size=16,

2937 + shapes=((1, 1),),

2938 + dtypes=(torch.float32,),

2939 + mamba_cache_mode="align",

2940 + )

2941 + worker._kv_cache_groups = [

2942 + KVCacheGroupSpec(["full"], full),

2943 + KVCacheGroupSpec(["mamba"], mamba),

2944 + ]

2945 + worker.hash_block_size = 4

2946 + worker.token_dbs = [

2947 + ChunkedTokenDatabase(

2948 + KeyMetadata("test-model", 0, 0, 0, 0, group_id=group_id),

2949 + block_size=16,

2950 + hash_block_size=4,

2951 + )

2952 + for group_id in range(2)

2953 + ]

2954 + worker.coord = mooncake_store_worker.MooncakeStoreCoordinator(

2955 + worker._kv_cache_groups,

2956 + scheduler_block_size=16,

2957 + hash_block_size=4,

2958 + )

2959 + _refresh_group_tp_replication_factors(worker)

2960 + hashes = [BlockHash(f"h{i}".encode()) for i in range(12)]

2961 + present = {

2962 + (0, bytes(hashes[3])),

2963 + (0, bytes(hashes[7])),

2964 + (0, bytes(hashes[11])),

2965 + (1, bytes(hashes[4])),

2966 + }

2967 +

2968 + def exists(keys):

2969 + return [

2970 + int(

2971 + any(

2972 + f"@group:{group_id}@{block_hash.hex()}" in key

2973 + for group_id, block_hash in present

2974 + )

2975 + )

2976 + for key in keys

2977 + ]

2978 +

2979 + worker.store.batch_is_exist.side_effect = exists

2980 +

2981 + result = worker.lookup(49, hashes)

2982 +

2983 + assert result.hit_length == 20

2984 + assert result.tail_key_boundaries == (

2985 + TailKeyBoundary(group_id=0, num_tokens=32),

2986 + TailKeyBoundary(group_id=1, num_tokens=20),

2987 + )

2988 +

2989 +

2990 def test_lookup_full_hit_swa_degrades_when_no_stored_boundary_is_usable():

2991 """The motivating livelock: the producer of a 64-token prompt stored

2992 only its SWA tail window (blocks 2-3). The old arithmetic clamp turned

@@ -2790,7 +3008,7 @@ def test_lookup_full_hit_swa_degrades_when_no_stored_boundary_is_usable():

3008 )

3009 worker.store.batch_is_exist.return_value = [0, 0, 1, 1]

3010

2793 - assert worker.lookup(64, [b"h0", b"h1", b"h2", b"h3"]) == 0

3011 + assert worker.lookup(64, [b"h0", b"h1", b"h2", b"h3"]).hit_length == 0

3012 assert worker.store.batch_is_exist.call_count == 1

3013

3014

@@ -2811,7 +3029,7 @@ def test_lookup_swa_single_group_returns_full_when_tail_window_present():

3029 hash_block_size=worker.hash_block_size,

3030 )

3031 worker.store.batch_is_exist.return_value = [0, 0, 1, 1]

2814 - assert worker.lookup(65, [b"h0", b"h1", b"h2", b"h3"]) == 64

3032 + assert worker.lookup(65, [b"h0", b"h1", b"h2", b"h3"]).hit_length == 64

3033

3034

3035 def test_lookup_checks_all_potential_swa_hit_boundaries():

@@ -2865,7 +3083,7 @@ def test_lookup_checks_all_potential_swa_hit_boundaries():

3083 [f"h{i}".encode() for i in range(12)],

3084 )

3085

2868 - assert result == 32

3086 + assert result.hit_length == 32

3087 keys = worker.store.batch_is_exist.call_args.args[0]

3088 assert len(keys) == 6

3089 swa_keys = [key for key in keys if "@group:1@" in key]

@@ -2949,6 +3167,24 @@ def test_consumer_starts_send_thread_only_when_put_is_enabled(save_decode_cache)

3167 assert (worker.kv_send_thread is not None) == save_decode_cache

3168

3169

3170 +def test_start_load_kv_queues_async_load():

3171 + w = _make_bare_worker()

3172 + req = _make_load_req(

3173 + "load-req",

3174 + [b"h0", b"h1"],

3175 + token_len=32,

3176 + )

3177 + assert req.load_spec is not None

3178 + req.load_spec.token_len = 0

3179 + meta = mooncake_store_worker.MooncakeStoreConnectorMetadata(set(), set())

3180 + meta.add_request(req)

3181 +

3182 + w.start_load_kv(meta)

3183 +

3184 + assert req.load_spec.token_len == 32

3185 + assert w.recv_request_queue.get_nowait() is req

3186 +

3187 +

3188 def test_putting_consumer_queues_decode_save():

3189 w = _make_bare_worker(kv_role="kv_consumer", save_decode_cache=True)

3190 send_thread = MagicMock()

@@ -3385,7 +3621,7 @@ def test_lookup_records_mooncake_metrics():

3621 result = worker.lookup(33, [b"a0", b"a1"])

3622 stats = worker.get_kv_connector_stats()

3623

3388 - assert result == 32

3624 + assert result.hit_length == 32

3625 assert isinstance(stats, MooncakeStoreConnectorStats)

3626 assert len(stats.data["lookup_exists"]) == 1

3627 assert stats.data["lookup_exists"][0]["num_keys"] == 2

tests/v1/kv_connector/unit/utils.py

@@ -465,6 +465,7 @@ def make_kv_cache_config(

465 mamba_enabled: bool = False,

466 sw_size: int = 128,

467 num_blocks: int = 100,

468 + mamba_cache_mode: Literal["all", "align", "none"] = "none",

469 ) -> KVCacheConfig:

470 kv_cache_groups = [

471 KVCacheGroupSpec(

@@ -498,6 +499,7 @@ def make_kv_cache_config(

499 block_size=block_size,

500 shapes=((16,), (16,)),

501 dtypes=(torch.float16,),

502 + mamba_cache_mode=mamba_cache_mode,

503 ),

504 )

505 )

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py

+2 -5

@@ -109,11 +109,8 @@ class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA):

109 f"{cache_block_size} (mamba_cache_mode != 'align')"

110 )

111 pcp = vllm_config.parallel_config.prefill_context_parallel_size

112 - dcp = vllm_config.parallel_config.decode_context_parallel_size

113 - if len(kv_cache_config.transfer_groups) > 1 and pcp * dcp > 1:

114 - unsupported.append(

115 - f"PCP/DCP > 1 (pcp={pcp}, dcp={dcp}) with hybrid attention"

116 - )

112 + if len(kv_cache_config.transfer_groups) > 1 and pcp > 1:

113 + unsupported.append(f"PCP > 1 (pcp={pcp}) with hybrid attention")

114 if unsupported:

115 raise ValueError(

116 "MooncakeStoreConnector does not support: " + "; ".join(unsupported)

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py

+21 -11

@@ -55,6 +55,12 @@ class ExternalCachedBlockPool:

55 return [self._present_block] * len(group_ids)

56 return None

58 + def contains(self, group_id: int, block_hash: BlockHash) -> bool:

59 + """Return whether a group has a loadable key for this hash."""

60 + return (

61 + self._exists is not None and (group_id, bytes(block_hash)) in self._exists

62 + )

63 +

65 class MooncakeStoreCoordinator:

66 """Mirror of ``HybridKVCacheCoordinator.find_longest_cache_hit`` over an

@@ -67,6 +73,7 @@ class MooncakeStoreCoordinator:

73 hash_block_size: int,

74 use_eagle: bool = False,

75 retention_interval: int | None = None,

76 + dcp_world_size: int = 1,

77 ) -> None:

78 assert all(

79 g.kv_cache_spec.block_size % hash_block_size == 0 for g in kv_cache_groups

@@ -83,7 +90,7 @@ class MooncakeStoreCoordinator:

90 self.hash_block_size = hash_block_size

91 self.lcm_block_size = scheduler_block_size

92 self.enable_partial_hash_hits = partial_hash_hits_enabled(

86 - kv_cache_groups, hash_block_size

93 + kv_cache_groups, hash_block_size, dcp_world_size

94 )

95 self.use_eagle = use_eagle

96 # Mirror vLLM core's KVCacheCoordinator.retention_interval.

@@ -376,10 +383,11 @@ class MooncakeStoreCoordinator:

383 # Truncate full-attention hit_blocks to final converged length;

384 # other specs already trim themselves inside their hit logic. cdiv keeps

385 # the partial tail block when hit_length is not block-aligned.

379 - first_group = self.attention_groups[0]

380 - if isinstance(first_group.spec, FullAttentionSpec):

381 - num_blocks = cdiv(hit_length, first_group.spec.block_size)

382 - for group_id in first_group.group_ids:

386 + for group in self.attention_groups:

387 + if not isinstance(group.spec, FullAttentionSpec):

388 + continue

389 + num_blocks = cdiv(hit_length, group.spec.block_size)

390 + for group_id in group.group_ids:

391 full_blks = hit_blocks_by_group[group_id]

392 assert full_blks is not None

393 del full_blks[num_blocks:]

@@ -398,15 +406,17 @@ def _unwrap_spec(spec: KVCacheSpec) -> KVCacheSpec:

406

407

408 def partial_hash_hits_enabled(

401 - kv_cache_groups: list[KVCacheGroupSpec], hash_block_size: int

409 + kv_cache_groups: list[KVCacheGroupSpec],

410 + hash_block_size: int,

411 + dcp_world_size: int = 1,

412 ) -> bool:

403 - """Mirror of core's ``HybridKVCacheCoordinator.enable_partial_hash_hits``

404 - (its dcp == 1 clause holds: the connector rejects hybrid + DCP/PCP > 1).

405 - Single copy on purpose — scheduler and coordinator must not disagree.

406 - """

413 + """Match core's DCP-aware Mamba partial-hit condition."""

414 return any(

415 isinstance(spec := _unwrap_spec(g.kv_cache_spec), MambaSpec)

416 and spec.mamba_cache_mode == "align"

410 - and spec.block_size > hash_block_size

417 + and (

418 + (dcp_world_size == 1 and spec.block_size > hash_block_size)

419 + or (dcp_world_size > 1 and spec.block_size >= hash_block_size)

420 + )

421 for g in kv_cache_groups

422 )

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py

+34

@@ -300,6 +300,39 @@ class ChunkedTokenDatabase:

300 yield start_idx, end_idx, h

301

302

303 +@dataclass(frozen=True)

304 +class TailKeyBoundary:

305 + """Hash boundary used to key a group's tail block in the store.

306 +

307 + Attributes:

308 + group_id: KV-cache group containing the tail block.

309 + num_tokens: Token boundary whose prefix hash identifies the matched

310 + stored block. The loader uses

311 + ``block_hashes[num_tokens // hash_block_size - 1]`` instead of the

312 + hash implied by ``MooncakeLookupResult.hit_length``. This changes

313 + only the load key, not the reusable prefix.

314 + """

315 +

316 + group_id: int

317 + num_tokens: int

318 +

319 +

320 +@dataclass

321 +class MooncakeLookupResult:

322 + """Lookup result used to build the subsequent load request.

323 +

324 + Attributes:

325 + hit_length: Longest prefix that every KV-cache group can reuse after

326 + their individual cache hits converge.

327 + tail_key_boundaries: Hash boundary used to store each cache group's

328 + tail block when ``hit_length`` does not identify its store key.

329 + There is one entry per group for every nonzero hit.

330 + """

331 +

332 + hit_length: int

333 + tail_key_boundaries: tuple[TailKeyBoundary, ...] = ()

334 +

335 +

336 @dataclass

337 class LoadSpec:

338 """Specification for loading KV cache from external store."""

@@ -308,6 +341,7 @@ class LoadSpec:

341 kvpool_cached_tokens: int

342 can_load: bool

343 token_len: int = 0

344 + tail_key_boundaries: tuple[TailKeyBoundary, ...] = ()

345

346

347 @dataclass

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/protocol.py

+42 -1

@@ -16,7 +16,9 @@ Wire format (REQ/REP over IPC):

16 fixed-size block hash (0 when there are no hashes)

17 frame 3: raw block hashes concatenated back-to-back (each hash_len

18 bytes); the server splits on hash_len

19 - Response: [hit_count: u32 big-endian, 4 bytes]

19 + Response: hit_length (u32 big-endian, first 4 bytes), followed by zero

20 + or more 8-byte tail-key boundaries. Each entry is group_id

21 + (u32), then boundary_tokens (u32).

23 msg_type == RESET_MSG:

24 (no payload frames)

@@ -31,6 +33,11 @@ Mirrors the named-tag convention used by the NIXL connector (see

33 ``vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py``).

34 """

36 +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (

37 + MooncakeLookupResult,

38 + TailKeyBoundary,

39 +)

40 +

41 # Request message-type tags. Frame 0 of every request.

42 LOOKUP_MSG: bytes = b"lookup"

43 RESET_MSG: bytes = b"reset"

@@ -38,3 +45,37 @@ RESET_MSG: bytes = b"reset"

45 # Single-byte response status codes for admin commands.

46 RESP_OK: bytes = b"\x01"

47 RESP_ERR: bytes = b"\x00"

48 +

49 +# group_id (u32), boundary_tokens (u32).

50 +TAIL_KEY_BOUNDARY_ENTRY_SIZE: int = 8

51 +

52 +

53 +def encode_lookup_response(result: MooncakeLookupResult) -> bytes:

54 + hit_length = result.hit_length.to_bytes(4, "big")

55 + if not result.tail_key_boundaries:

56 + return hit_length

57 +

58 + payload = bytearray(hit_length)

59 + for boundary in result.tail_key_boundaries:

60 + payload.extend(boundary.group_id.to_bytes(4, "big"))

61 + payload.extend(boundary.num_tokens.to_bytes(4, "big"))

62 + return bytes(payload)

63 +

64 +

65 +def decode_lookup_response(payload: bytes) -> MooncakeLookupResult:

66 + if len(payload) < 4 or (len(payload) - 4) % TAIL_KEY_BOUNDARY_ENTRY_SIZE:

67 + raise ValueError("Invalid Mooncake lookup response")

68 +

69 + hit_length = int.from_bytes(payload[:4], "big")

70 + if len(payload) == 4:

71 + return MooncakeLookupResult(hit_length)

72 +

73 + boundaries = []

74 + for offset in range(4, len(payload), TAIL_KEY_BOUNDARY_ENTRY_SIZE):

75 + boundaries.append(

76 + TailKeyBoundary(

77 + group_id=int.from_bytes(payload[offset : offset + 4], "big"),

78 + num_tokens=int.from_bytes(payload[offset + 4 : offset + 8], "big"),

79 + )

80 + )

81 + return MooncakeLookupResult(hit_length, tuple(boundaries))

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py

+7 -3

@@ -76,7 +76,9 @@ class MooncakeStoreScheduler:

76 kv_cache_config, vllm_config

77 )

78 self.enable_partial_hash_hits = partial_hash_hits_enabled(

79 - kv_cache_config.kv_cache_groups, self._hash_block_size

79 + kv_cache_config.kv_cache_groups,

80 + self._hash_block_size,

81 + vllm_config.parallel_config.decode_context_parallel_size,

82 )

84 # Per-request state

@@ -115,15 +117,16 @@ class MooncakeStoreScheduler:

117 if request.num_tokens < align:

118 return 0, False

119

118 - num_external_hit_tokens = self.client.lookup(

120 + lookup_result = self.client.lookup(

121 request.request_id,

122 request.num_tokens,

123 request.block_hashes,

124 non_block=self.lookup_async,

125 )

124 - if num_external_hit_tokens is None:

126 + if lookup_result is None:

127 # Lookup not ready yet; scheduler will retry on a later step.

128 return None, False

129 + num_external_hit_tokens = lookup_result.hit_length

130

131 if num_external_hit_tokens < num_computed_tokens:

132 need_to_allocate = 0

@@ -145,6 +148,7 @@ class MooncakeStoreScheduler:

148 vllm_cached_tokens=num_computed_tokens,

149 kvpool_cached_tokens=num_external_hit_tokens,

150 can_load=False,

151 + tail_key_boundaries=lookup_result.tail_key_boundaries,

152 )

153

154 return need_to_allocate, self.load_async

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py

+98 -32

@@ -44,16 +44,20 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (

44 BlobBlockHashes,

45 ChunkedTokenDatabase,

46 KeyMetadata,

47 + MooncakeLookupResult,

48 MooncakeStoreConnectorMetadata,

49 MooncakeStoreWorkerMetadata,

50 PoolKey,

51 ReqMeta,

52 + TailKeyBoundary,

53 )

54 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.protocol import ( # noqa: E501

55 LOOKUP_MSG,

56 RESET_MSG,

57 RESP_ERR,

58 RESP_OK,

59 + decode_lookup_response,

60 + encode_lookup_response,

61 )

62 from vllm.logger import init_logger

63 from vllm.utils.math_utils import cdiv

@@ -63,11 +67,11 @@ from vllm.v1.attention.backends.utils import NULL_BLOCK_ID

67 from vllm.v1.core.kv_cache_utils import (

68 BlockHash,

69 maybe_convert_block_hash,

70 + resolve_dcp_kv_cache_spec,

71 resolve_kv_cache_block_sizes,

72 )

73 from vllm.v1.kv_cache_interface import (

74 KVCacheConfig,

70 - KVCacheGroupSpec,

75 KVCacheSpec,

76 MambaSpec,

77 MLAAttentionSpec,

@@ -1179,6 +1183,12 @@ class KVCacheStoreRecvingThread(KVTransferThread):

1183 # Skip chunks the consumer's per-group spec wouldn't populate

1184 # locally (e.g. SWA pre-window) even if the producer stored them.

1185 load_mask_per_group = self.coord.load_mask(req_meta.block_hashes, token_len)

1186 + tail_key_boundaries = {

1187 + boundary.group_id: boundary.num_tokens

1188 + for boundary in (

1189 + req_meta.load_spec.tail_key_boundaries # type: ignore[union-attr]

1190 + )

1191 + }

1192

1193 addr_list: list[list[int]] = []

1194 size_list: list[list[int]] = []

@@ -1193,6 +1203,13 @@ class KVCacheStoreRecvingThread(KVTransferThread):

1203 chunk_idx = start // db.block_size

1204 if chunk_idx >= len(mask) or not mask[chunk_idx]:

1205 continue

1206 + boundary_tokens = (

1207 + tail_key_boundaries.get(g_idx) if end == token_len else None

1208 + )

1209 + if boundary_tokens is not None:

1210 + block_hash = req_meta.block_hashes[

1211 + boundary_tokens // db.hash_block_size - 1

1212 + ]

1213 key_list.append(db.key_for(block_hash))

1214 chunks.append((start, end))

1215 g_addrs, g_sizes, g_block_ids = db.prepare_values(

@@ -1494,21 +1511,15 @@ class MooncakeStoreWorker:

1511 )

1512 return

1513

1497 - # Single-group + PCP/DCP > 1: scale the lone group's spec.block_size to

1498 - # self.block_size (= scheduler_block_size) so the coordinator's

1499 - # ``block_size % hash_block_size == 0`` invariant holds.

1500 - groups = list(kv_cache_config.transfer_groups)

1501 - if len(groups) == 1 and groups[0].kv_cache_spec.block_size != self.block_size:

1502 - g = groups[0]

1503 - groups = [

1504 - dataclasses.replace(

1505 - g,

1506 - kv_cache_spec=dataclasses.replace(

1507 - g.kv_cache_spec, block_size=self.block_size

1508 - ),

1509 - )

1510 - ]

1511 - self._kv_cache_groups: list[KVCacheGroupSpec] = groups

1514 + self._kv_cache_groups = [

1515 + dataclasses.replace(

1516 + group,

1517 + kv_cache_spec=resolve_dcp_kv_cache_spec(

1518 + group.kv_cache_spec, self.dcp_size

1519 + ),

1520 + )

1521 + for group in kv_cache_config.transfer_groups

1522 + ]

1523 spec_cfg = getattr(vllm_config, "speculative_config", None)

1524 use_eagle = bool(

1525 spec_cfg.use_eagle()

@@ -1521,6 +1532,7 @@ class MooncakeStoreWorker:

1532 hash_block_size=self.hash_block_size,

1533 use_eagle=use_eagle,

1534 retention_interval=kv_cache_config.prefix_cache_retention_interval,

1535 + dcp_world_size=self.dcp_size,

1536 )

1537 # One ChunkedTokenDatabase per group; addresses populated in

1538 # register_kv_caches once the kv-cache layout is known. Each group's

@@ -1881,7 +1893,9 @@ class MooncakeStoreWorker:

1893 return None

1894 return MooncakeStoreWorkerMetadata(completed_saves=completed_saves)

1895

1884 - def lookup(self, num_tokens: int, block_hashes: Sequence[BlockHash]) -> int:

1896 + def lookup(

1897 + self, num_tokens: int, block_hashes: Sequence[BlockHash]

1898 + ) -> MooncakeLookupResult:

1899 """Check how many prefix tokens exist in the store.

1900

1901 Checks across all rank-specific key namespaces that may be loaded. A

@@ -1889,11 +1903,11 @@ class MooncakeStoreWorker:

1903 the last token is recomputed for sampling.

1904 """

1905 if self._capacity_only:

1892 - return 0

1906 + return MooncakeLookupResult(0)

1907

1908 token_len = self.coord.align_lookup_length(num_tokens)

1909 if not block_hashes or token_len <= 0:

1896 - return 0

1910 + return MooncakeLookupResult(0)

1911

1912 # Build per-(group, hash) candidate keys expanded across rank namespaces.

1913 # candidate_meta stores the (group, hash_bytes) for key slice.

@@ -1934,7 +1948,7 @@ class MooncakeStoreWorker:

1948 candidate_meta.append((g_idx, bytes(h)))

1949

1950 if not candidate_keys:

1937 - return 0

1951 + return MooncakeLookupResult(0)

1952

1953 lookup_start = time.perf_counter()

1954 try:

@@ -1953,7 +1967,7 @@ class MooncakeStoreWorker:

1967 num_failed_keys=len(candidate_keys),

1968 )

1969 logger.error("Remote connection failed in lookup: %s", e)

1956 - return 0

1970 + return MooncakeLookupResult(0)

1971

1972 # A (group, hash) is "present" only when every namespace that will be

1973 # loaded has it (per-group count: sharded groups need every rank's

@@ -1970,7 +1984,7 @@ class MooncakeStoreWorker:

1984 self.hash_block_size,

1985 exists_set,

1986 )

1973 - _masks, hit_length = self.coord.find_longest_cache_hit(

1987 + _, hit_length = self.coord.find_longest_cache_hit(

1988 block_hashes,

1989 token_len,

1990 cached_block_pool,

@@ -1978,13 +1992,63 @@ class MooncakeStoreWorker:

1992 if hit_length >= num_tokens:

1993 usable_length = self.coord.align_lookup_length(num_tokens - 1)

1994 if usable_length <= 0:

1981 - return 0

1982 - _masks, hit_length = self.coord.find_longest_cache_hit(

1995 + return MooncakeLookupResult(0)

1996 + _, hit_length = self.coord.find_longest_cache_hit(

1997 block_hashes,

1998 usable_length,

1999 cached_block_pool,

2000 )

1987 - return hit_length

2001 + return MooncakeLookupResult(

2002 + hit_length,

2003 + self._tail_key_boundaries(

2004 + block_hashes,

2005 + hit_length,

2006 + cached_block_pool,

2007 + ),

2008 + )

2009 +

2010 + def _tail_key_boundaries(

2011 + self,

2012 + block_hashes: Sequence[BlockHash],

2013 + hit_length: int,

2014 + cached_block_pool: ExternalCachedBlockPool,

2015 + ) -> tuple[TailKeyBoundary, ...]:

2016 + """Return the hash boundary used to store each group's tail block.

2017 +

2018 + With fine-grained prefix matching, ``hit_length`` may fall within a

2019 + physical cache block and may not align with the hash boundary used to

2020 + store that block. For each KV-cache group, return the token boundary

2021 + whose hash was used as the store key.

2022 + """

2023 + if hit_length <= 0:

2024 + return ()

2025 +

2026 + boundaries = []

2027 + hit_boundary_hash_idx = hit_length // self.hash_block_size - 1

2028 + for group_id, db in enumerate(self.token_dbs):

2029 + chunk_id = cdiv(hit_length, db.block_size) - 1

2030 + boundary_tokens = hit_length

2031 + contains_hit_boundary = cached_block_pool.contains(

2032 + group_id, block_hashes[hit_boundary_hash_idx]

2033 + )

2034 + if not self.coord.enable_partial_hash_hits:

2035 + assert contains_hit_boundary

2036 + if not contains_hit_boundary:

2037 + next_chunk_hash_idx = min(

2038 + (chunk_id + 1) * db.block_size // self.hash_block_size,

2039 + len(block_hashes),

2040 + )

2041 + for hash_idx in range(hit_boundary_hash_idx + 1, next_chunk_hash_idx):

2042 + if cached_block_pool.contains(group_id, block_hashes[hash_idx]):

2043 + boundary_tokens = (hash_idx + 1) * self.hash_block_size

2044 + break

2045 + else:

2046 + raise AssertionError(

2047 + f"No tail key found for cache group {group_id} at "

2048 + f"hit length {hit_length}"

2049 + )

2050 + boundaries.append(TailKeyBoundary(group_id, boundary_tokens))

2051 + return tuple(boundaries)

2052

2053 def get_kv_events(self) -> list[BlockStored]:

2054 if self.enable_kv_events and self.kv_send_thread is not None:

@@ -2017,7 +2081,7 @@ class LookupKeyServer:

2081 """ZMQ server on worker rank 0 for the LookupKey admin channel.

2082

2083 Handles two request types, tagged at frame 0:

2020 - - ``LOOKUP_MSG``: prefix-cache hit query, returns hit count.

2084 + - ``LOOKUP_MSG``: prefix-cache hit query, returns its load plan.

2085 - ``RESET_MSG``: drains the send thread queue, then runs

2086 ``store.remove_all(force=True)``. Caller must have paused the

2087 scheduler first.

@@ -2054,7 +2118,7 @@ class LookupKeyServer:

2118 blob = all_frames[3].buffer

2119 block_hashes = BlobBlockHashes(blob, hash_len)

2120 result = self.store_worker.lookup(num_tokens, block_hashes)

2057 - self.socket.send(result.to_bytes(4, "big"))

2121 + self.socket.send(encode_lookup_response(result))

2122

2123 elif msg_type == RESET_MSG:

2124 try:

@@ -2114,9 +2178,11 @@ class LookupKeyClient:

2178 self.executor = ThreadPoolExecutor(

2179 max_workers=1, thread_name_prefix="MooncakeLookupClient"

2180 )

2117 - self.futures: dict[str, Future[int]] = {}

2181 + self.futures: dict[str, Future[MooncakeLookupResult]] = {}

2182

2119 - def _lookup(self, num_tokens: int, block_hashes: list[BlockHash]) -> int:

2183 + def _lookup(

2184 + self, num_tokens: int, block_hashes: list[BlockHash]

2185 + ) -> MooncakeLookupResult:

2186 hash_len = len(block_hashes[0]) if block_hashes else 0

2187 all_frames = (

2188 LOOKUP_MSG,

@@ -2126,7 +2192,7 @@ class LookupKeyClient:

2192 )

2193 self.socket.send_multipart(all_frames, copy=False)

2194 resp = self.socket.recv()

2129 - return int.from_bytes(resp, "big")

2195 + return decode_lookup_response(resp)

2196

2197 def lookup(

2198 self,

@@ -2134,7 +2200,7 @@ class LookupKeyClient:

2200 num_tokens: int,

2201 block_hashes: list[BlockHash],

2202 non_block: bool = False,

2137 - ) -> int | None:

2203 + ) -> MooncakeLookupResult | None:

2204 """If non_block is True, will return None until the result is ready,

2205 so the caller retries on a later step."""

2206 future = self.futures.get(req_id)

@@ -2147,7 +2213,7 @@ class LookupKeyClient:

2213 return future.result()

2214 except Exception as e:

2215 logger.error("Async Mooncake lookup failed for %s: %s", req_id, e)

2150 - return 0

2216 + return MooncakeLookupResult(0)

2217 finally:

2218 del self.futures[req_id]

2219

vllm/v1/core/kv_cache_coordinator.py

+7 -7

@@ -867,14 +867,14 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

867 if is_simple_hybrid:

868 break

869

870 - # Truncate full attention blocks to final hit_length (if present)

871 - first_group = self.attention_groups[0]

872 - if isinstance(first_group.spec, FullAttentionSpec):

873 - group_block_size = self.single_type_managers[

874 - first_group.group_ids[0]

875 - ].block_size

870 + # Truncate every full-attention group (target and draft) blocks

871 + # to final hit_length.

872 + for group in self.attention_groups:

873 + if not isinstance(group.spec, FullAttentionSpec):

874 + continue

875 + group_block_size = self.single_type_managers[group.group_ids[0]].block_size

876 num_blocks = cdiv(hit_length, group_block_size)

877 - for group_id in first_group.group_ids:

877 + for group_id in group.group_ids:

878 if (blks := hit_blocks_by_group[group_id]) is not None:

879 del blks[num_blocks:]

880 hit_length_by_group[group_id] = hit_length

vllm/v1/core/kv_cache_utils.py

+29 -4

@@ -35,6 +35,7 @@ from vllm.v1.kv_cache_interface import (

35 SlidingWindowSpec,

36 UniformTypeKVCacheSpecs,

37 compute_layout_strides,

38 + iter_layer_specs,

39 replace_as,

40 )

41 from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry

@@ -647,6 +648,33 @@ def hash_block_tokens(

648 )

649

650

651 +def resolve_dcp_kv_block_size(spec: KVCacheSpec, dcp_world_size: int) -> int:

652 + """Return the token span of a cache block under DCP."""

653 + layer_specs = iter_layer_specs(spec)

654 + if len(layer_specs) > 0 and all(

655 + isinstance(layer_spec, AttentionSpec) for layer_spec in layer_specs

656 + ):

657 + return spec.block_size * dcp_world_size

658 + return spec.block_size

659 +

660 +

661 +def resolve_dcp_kv_cache_spec(spec: KVCacheSpec, dcp_world_size: int) -> KVCacheSpec:

662 + """Return a KV cache spec with block sizes adjusted for DCP."""

663 + block_size = resolve_dcp_kv_block_size(spec, dcp_world_size)

664 + if block_size == spec.block_size:

665 + return spec

666 + if isinstance(spec, UniformTypeKVCacheSpecs):

667 + return replace(

668 + spec,

669 + block_size=block_size,

670 + kv_cache_specs={

671 + name: resolve_dcp_kv_cache_spec(layer_spec, dcp_world_size)

672 + for name, layer_spec in spec.kv_cache_specs.items()

673 + },

674 + )

675 + return replace(spec, block_size=block_size)

676 +

677 +

678 def resolve_kv_cache_block_sizes(

679 kv_cache_config: KVCacheConfig,

680 vllm_config: VllmConfig,

@@ -674,10 +702,7 @@ def resolve_kv_cache_block_sizes(

702 return bs, bs

703

704 group_block_sizes = [

677 - g.kv_cache_spec.block_size * dcp

678 - if isinstance(g.kv_cache_spec, AttentionSpec)

679 - else g.kv_cache_spec.block_size

680 - for g in groups

705 + resolve_dcp_kv_block_size(g.kv_cache_spec, dcp) for g in groups

706 ]

707 scheduler_block_size = math.lcm(*group_block_sizes)

708