@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 056bc2e166
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 3, 08:00 UTC

[KVConnector][DSV4] HMA support for Mooncake store connector (#42828)

Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>

Yifan Qiao committed May 19, 2026 at 01:07 UTC 056bc2e16646599a96ac94e761c953e680e6fba9
10 files changed +1835 -446
tests/v1/kv_connector/unit/test_mooncake_store_connector.py
+72 -94
@@ -9,12 +9,17 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
9 KVConnectorRole,
10 )
11 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store import (
12 - connector,
13 - worker,
12 + connector as mooncake_store_connector,
13 )
15 -from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( # noqa: E501
14 +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
15 MooncakeStoreConnectorMetadata,
16 )
17 +from vllm.v1.kv_cache_interface import (
18 + FullAttentionSpec,
19 + KVCacheConfig,
20 + KVCacheGroupSpec,
21 + KVCacheTensor,
22 +)
23 from vllm.v1.outputs import KVConnectorOutput
24
25 from .utils import create_vllm_config
@@ -27,6 +32,17 @@ def _make_vllm_config():
32 )
33
34
35 +def _make_kv_cache_config() -> KVCacheConfig:
36 + """Single-group full-attention KVCacheConfig — enough for the connector
37 + constructor's validate() pass."""
38 + spec = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None)
39 + return KVCacheConfig(
40 + num_blocks=4,
41 + kv_cache_tensors=[KVCacheTensor(size=8192, shared_by=["layer0"])],
42 + kv_cache_groups=[KVCacheGroupSpec(["layer0"], spec)],
43 + )
44 +
45 +
46 def _make_block_stored() -> BlockStored:
47 return BlockStored(
48 block_hashes=[b"hash"],
@@ -41,6 +57,7 @@ def _make_block_stored() -> BlockStored:
57
58 def test_scheduler_role_initializes_store_scheduler_only():
59 vllm_config = _make_vllm_config()
60 + kv_cache_config = _make_kv_cache_config()
61
62 with (
63 set_current_vllm_config(vllm_config),
@@ -53,75 +70,19 @@ def test_scheduler_role_initializes_store_scheduler_only():
70 "connector.MooncakeStoreWorker"
71 ) as mock_worker,
72 ):
56 - conn = connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.SCHEDULER)
73 + connector = mooncake_store_connector.MooncakeStoreConnector(
74 + vllm_config, KVConnectorRole.SCHEDULER, kv_cache_config
75 + )
76
58 - mock_scheduler.assert_called_once_with(vllm_config)
77 + mock_scheduler.assert_called_once_with(vllm_config, kv_cache_config)
78 mock_worker.assert_not_called()
60 - assert conn.connector_scheduler is mock_scheduler.return_value
61 - assert conn.connector_worker is None
62 -
63 -
64 -def test_worker_role_initializes_store_worker_on_rank0():
65 - vllm_config = _make_vllm_config()
66 -
67 - with (
68 - set_current_vllm_config(vllm_config),
69 - patch(
70 - "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
71 - "connector.MooncakeStoreScheduler"
72 - ) as mock_scheduler,
73 - patch(
74 - "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
75 - "connector.MooncakeStoreWorker"
76 - ) as mock_worker,
77 - ):
78 - conn = connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.WORKER)
79 -
80 - mock_scheduler.assert_not_called()
81 - mock_worker.assert_called_once_with(vllm_config)
82 - assert conn.connector_scheduler is None
83 - assert conn.connector_worker is mock_worker.return_value
84 -
85 -
86 -def test_worker_role_initializes_on_nonzero_rank():
87 - vllm_config = _make_vllm_config()
88 - vllm_config.parallel_config.rank = 1
89 -
90 - with (
91 - set_current_vllm_config(vllm_config),
92 - patch(
93 - "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
94 - "connector.MooncakeStoreWorker"
95 - ) as mock_worker,
96 - ):
97 - connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.WORKER)
98 -
99 - mock_worker.assert_called_once_with(vllm_config)
100 -
101 -
102 -def test_lookup_rpc_path_uses_data_parallel_index_in_dense_dp():
103 - vllm_config = _make_vllm_config()
104 - vllm_config.parallel_config.data_parallel_rank = 0
105 - vllm_config.parallel_config.data_parallel_index = 3
106 -
107 - path = worker.get_zmq_rpc_path_lookup(vllm_config)
108 -
109 - assert path.endswith("_dp_rank3")
110 -
111 -
112 -def test_lookup_rpc_path_uses_local_rank_when_local_engines_only():
113 - vllm_config = _make_vllm_config()
114 - vllm_config.parallel_config.data_parallel_index = 7
115 - vllm_config.parallel_config.data_parallel_rank_local = 1
116 - vllm_config.parallel_config.data_parallel_hybrid_lb = True
117 -
118 - path = worker.get_zmq_rpc_path_lookup(vllm_config)
119 -
120 - assert path.endswith("_dp_rank1")
79 + assert connector.connector_scheduler is mock_scheduler.return_value
80 + assert connector.connector_worker is None
81
82
83 def test_worker_methods_delegate_to_store_worker():
84 vllm_config = _make_vllm_config()
85 + kv_cache_config = _make_kv_cache_config()
86 kv_caches = {"layer0": MagicMock()}
87 metadata = MooncakeStoreConnectorMetadata(set(), set())
88 finished_req_ids = {"req-1"}
@@ -133,22 +94,25 @@ def test_worker_methods_delegate_to_store_worker():
94 "connector.MooncakeStoreWorker"
95 ) as mock_worker_cls,
96 ):
136 - conn = connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.WORKER)
97 + connector = mooncake_store_connector.MooncakeStoreConnector(
98 + vllm_config, KVConnectorRole.WORKER, kv_cache_config
99 + )
100
138 - worker_inst = mock_worker_cls.return_value
139 - worker_inst.get_finished.return_value = ({"req-1"}, {"req-2"})
140 - conn.bind_connector_metadata(metadata)
101 + worker = mock_worker_cls.return_value
102 + worker.get_finished.return_value = ({"req-1"}, {"req-2"})
103 + connector.bind_connector_metadata(metadata)
104
142 - conn.register_kv_caches(kv_caches)
143 - result = conn.get_finished(finished_req_ids)
105 + connector.register_kv_caches(kv_caches)
106 + result = connector.get_finished(finished_req_ids)
107
145 - worker_inst.register_kv_caches.assert_called_once_with(kv_caches)
146 - worker_inst.get_finished.assert_called_once_with(finished_req_ids, metadata)
108 + worker.register_kv_caches.assert_called_once_with(kv_caches)
109 + worker.get_finished.assert_called_once_with(finished_req_ids, metadata)
110 assert result == ({"req-1"}, {"req-2"})
111
112
113 def test_get_kv_connector_kv_cache_events_returns_none_when_empty():
114 vllm_config = _make_vllm_config()
115 + kv_cache_config = _make_kv_cache_config()
116
117 with (
118 set_current_vllm_config(vllm_config),
@@ -157,14 +121,17 @@ def test_get_kv_connector_kv_cache_events_returns_none_when_empty():
121 "connector.MooncakeStoreWorker"
122 ) as mock_worker_cls,
123 ):
160 - conn = connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.WORKER)
124 + connector = mooncake_store_connector.MooncakeStoreConnector(
125 + vllm_config, KVConnectorRole.WORKER, kv_cache_config
126 + )
127
128 mock_worker_cls.return_value.get_kv_events.return_value = []
163 - assert conn.get_kv_connector_kv_cache_events() is None
129 + assert connector.get_kv_connector_kv_cache_events() is None
130
131
132 def test_get_kv_connector_kv_cache_events_wraps_worker_events():
133 vllm_config = _make_vllm_config()
134 + kv_cache_config = _make_kv_cache_config()
135 event = _make_block_stored()
136
137 with (
@@ -174,12 +141,14 @@ def test_get_kv_connector_kv_cache_events_wraps_worker_events():
141 "connector.MooncakeStoreWorker"
142 ) as mock_worker_cls,
143 ):
177 - conn = connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.WORKER)
144 + connector = mooncake_store_connector.MooncakeStoreConnector(
145 + vllm_config, KVConnectorRole.WORKER, kv_cache_config
146 + )
147
148 mock_worker_cls.return_value.get_kv_events.return_value = [event]
180 - kv_events = conn.get_kv_connector_kv_cache_events()
149 + kv_events = connector.get_kv_connector_kv_cache_events()
150
182 - assert isinstance(kv_events, connector.MooncakeStoreKVEvents)
151 + assert isinstance(kv_events, mooncake_store_connector.MooncakeStoreKVEvents)
152 assert kv_events.get_number_of_workers() == 1
153 assert kv_events.get_all_events() == [event]
154
@@ -187,6 +156,7 @@ def test_get_kv_connector_kv_cache_events_wraps_worker_events():
156 def test_prefer_cross_layer_blocks_from_config():
157 # Default: disabled
158 vllm_config = _make_vllm_config()
159 + kv_cache_config = _make_kv_cache_config()
160 with (
161 set_current_vllm_config(vllm_config),
162 patch(
@@ -194,8 +164,10 @@ def test_prefer_cross_layer_blocks_from_config():
164 "connector.MooncakeStoreScheduler"
165 ),
166 ):
197 - conn = connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.SCHEDULER)
198 - assert conn.prefer_cross_layer_blocks is False
167 + connector = mooncake_store_connector.MooncakeStoreConnector(
168 + vllm_config, KVConnectorRole.SCHEDULER, kv_cache_config
169 + )
170 + assert connector.prefer_cross_layer_blocks is False
171
172 # Enabled via config
173 vllm_config_enabled = create_vllm_config(
@@ -210,14 +182,15 @@ def test_prefer_cross_layer_blocks_from_config():
182 "connector.MooncakeStoreScheduler"
183 ),
184 ):
213 - conn_enabled = connector.MooncakeStoreConnector(
214 - vllm_config_enabled, KVConnectorRole.SCHEDULER
185 + connector_enabled = mooncake_store_connector.MooncakeStoreConnector(
186 + vllm_config_enabled, KVConnectorRole.SCHEDULER, kv_cache_config
187 )
216 - assert conn_enabled.prefer_cross_layer_blocks is True
188 + assert connector_enabled.prefer_cross_layer_blocks is True
189
190
191 def test_register_cross_layers_kv_cache_delegates_to_worker():
192 vllm_config = _make_vllm_config()
193 + kv_cache_config = _make_kv_cache_config()
194
195 with (
196 set_current_vllm_config(vllm_config),
@@ -226,18 +199,21 @@ def test_register_cross_layers_kv_cache_delegates_to_worker():
199 "connector.MooncakeStoreWorker"
200 ) as mock_worker_cls,
201 ):
229 - conn = connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.WORKER)
202 + connector = mooncake_store_connector.MooncakeStoreConnector(
203 + vllm_config, KVConnectorRole.WORKER, kv_cache_config
204 + )
205
206 fake_tensor = MagicMock()
207 fake_backend = MagicMock()
233 - conn.register_cross_layers_kv_cache(fake_tensor, fake_backend)
208 + connector.register_cross_layers_kv_cache(fake_tensor, fake_backend)
209
235 - worker_inst = mock_worker_cls.return_value
236 - worker_inst.register_cross_layers_kv_caches.assert_called_once_with(fake_tensor)
210 + worker = mock_worker_cls.return_value
211 + worker.register_cross_layers_kv_caches.assert_called_once_with(fake_tensor)
212
213
214 def test_update_connector_output_and_take_events():
215 vllm_config = _make_vllm_config()
216 + kv_cache_config = _make_kv_cache_config()
217 event = _make_block_stored()
218
219 with (
@@ -247,12 +223,14 @@ def test_update_connector_output_and_take_events():
223 "connector.MooncakeStoreScheduler"
224 ),
225 ):
250 - conn = connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.SCHEDULER)
226 + connector = mooncake_store_connector.MooncakeStoreConnector(
227 + vllm_config, KVConnectorRole.SCHEDULER, kv_cache_config
228 + )
229
252 - kv_events = connector.MooncakeStoreKVEvents(num_workers=1)
230 + kv_events = mooncake_store_connector.MooncakeStoreKVEvents(num_workers=1)
231 kv_events.add_events([event])
254 - conn.update_connector_output(KVConnectorOutput(kv_cache_events=kv_events))
232 + connector.update_connector_output(KVConnectorOutput(kv_cache_events=kv_events))
233
256 - assert conn._kv_cache_events is kv_events
257 - assert list(conn.take_events()) == [event]
258 - assert conn._kv_cache_events is None
234 + assert connector._kv_cache_events is kv_events
235 + assert list(connector.take_events()) == [event]
236 + assert connector._kv_cache_events is None
tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py new
+302
@@ -0,0 +1,302 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from math import lcm
5 +
6 +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.coordinator import ( # noqa: E501
7 + ExternalCachedBlockPool,
8 + MooncakeStoreCoordinator,
9 +)
10 +from vllm.v1.core.kv_cache_utils import BlockHash, BlockHashListWithBlockSize
11 +from vllm.v1.kv_cache_interface import (
12 + FullAttentionSpec,
13 + KVCacheGroupSpec,
14 + SlidingWindowSpec,
15 +)
16 +
17 +
18 +def _make_coord(groups, hash_block_size, use_eagle=False):
19 + """Construct a coordinator using the natural LCM of group block sizes as
20 + the scheduler block size — mirrors ``resolve_kv_cache_block_sizes`` for
21 + the test fixtures."""
22 + block_sizes = [g.kv_cache_spec.block_size for g in groups]
23 + scheduler_block_size = lcm(*block_sizes)
24 + return MooncakeStoreCoordinator(
25 + groups,
26 + scheduler_block_size=scheduler_block_size,
27 + hash_block_size=hash_block_size,
28 + use_eagle=use_eagle,
29 + )
30 +
31 +
32 +# ----- ExternalCachedBlockPool -----
33 +
34 +
35 +def test_external_cached_block_pool_tautological_returns_present_for_any_hash():
36 + cmap = ExternalCachedBlockPool()
37 + h = BlockHash(b"\xaa" * 4)
38 + res = cmap.get_cached_block(h, [0, 1])
39 + assert res is not None
40 + assert len(res) == 2
41 + assert res[0] is not cmap.null_block
42 + assert res[1] is not cmap.null_block
43 +
44 +
45 +def test_external_cached_block_pool_hit_all_groups():
46 + h = BlockHash(b"\x11\x22\x33\x44")
47 + cmap = ExternalCachedBlockPool({(0, bytes(h)), (1, bytes(h))})
48 + res = cmap.get_cached_block(h, [0, 1])
49 + assert res is not None
50 + assert len(res) == 2
51 + assert res[0] is not cmap.null_block
52 + assert res[1] is not cmap.null_block
53 +
54 +
55 +def test_external_cached_block_pool_miss_one_group():
56 + h = BlockHash(b"\x11\x22\x33\x44")
57 + cmap = ExternalCachedBlockPool({(0, bytes(h))})
58 + assert cmap.get_cached_block(h, [0, 1]) is None
59 +
60 +
61 +def test_external_cached_block_pool_unknown_hash():
62 + h_known = BlockHash(b"\x01" * 4)
63 + h_unknown = BlockHash(b"\x02" * 4)
64 + cmap = ExternalCachedBlockPool({(0, bytes(h_known))})
65 + assert cmap.get_cached_block(h_unknown, [0]) is None
66 +
67 +
68 +# ----- Helpers -----
69 +
70 +
71 +def _full(block_size=16, sliding_window=None):
72 + return FullAttentionSpec(
73 + block_size=block_size,
74 + num_kv_heads=8,
75 + head_size=64,
76 + dtype=None,
77 + sliding_window=sliding_window,
78 + )
79 +
80 +
81 +def _swa(block_size=16, sliding_window=32):
82 + return SlidingWindowSpec(
83 + block_size=block_size,
84 + num_kv_heads=8,
85 + head_size=64,
86 + dtype=None,
87 + sliding_window=sliding_window,
88 + )
89 +
90 +
91 +def _hashes(n: int) -> list[BlockHash]:
92 + return [BlockHash(bytes([i + 1]) * 4) for i in range(n)]
93 +
94 +
95 +# ----- Single-group coordinator -----
96 +
97 +
98 +def test_coordinator_single_full_attention_all_hits():
99 + groups = [KVCacheGroupSpec(["L0"], _full(16))]
100 + coord = _make_coord(groups, hash_block_size=16)
101 + hs = _hashes(4)
102 + cmap = ExternalCachedBlockPool({(0, bytes(h)) for h in hs})
103 + masks, hit = coord.find_longest_cache_hit(hs, max_length=64, cached_block_pool=cmap)
104 + assert hit == 64
105 + assert masks[0] == [True, True, True, True]
106 +
107 +
108 +def test_coordinator_single_full_attention_partial_prefix():
109 + groups = [KVCacheGroupSpec(["L0"], _full(16))]
110 + coord = _make_coord(groups, hash_block_size=16)
111 + hs = _hashes(4)
112 + cmap = ExternalCachedBlockPool({(0, bytes(hs[0])), (0, bytes(hs[1]))})
113 + masks, hit = coord.find_longest_cache_hit(hs, max_length=64, cached_block_pool=cmap)
114 + assert hit == 32
115 + assert masks[0] == [True, True]
116 +
117 +
118 +def test_coordinator_single_full_attention_no_hits():
119 + groups = [KVCacheGroupSpec(["L0"], _full(16))]
120 + coord = _make_coord(groups, hash_block_size=16)
121 + hs = _hashes(4)
122 + cmap = ExternalCachedBlockPool(set())
123 + masks, hit = coord.find_longest_cache_hit(hs, max_length=64, cached_block_pool=cmap)
124 + assert hit == 0
125 + assert masks[0] == []
126 +
127 +
128 +def test_coordinator_single_swa_tautological_pool_masks_pre_window():
129 + """SWA tautological-pool: hit_length spans full prefix, mask is
130 + tail-window only."""
131 + groups = [KVCacheGroupSpec(["L0"], _swa(block_size=16, sliding_window=32))]
132 + coord = _make_coord(groups, hash_block_size=16)
133 + hs = _hashes(4) # 4 chunks * 16 tokens
134 + cmap = ExternalCachedBlockPool()
135 + masks, hit = coord.find_longest_cache_hit(hs, max_length=64, cached_block_pool=cmap)
136 + assert hit == 64
137 + # ceil((sw-1)/block_size) = ceil(31/16) = 2 tail blocks.
138 + assert masks[0][-2:] == [True, True]
139 + assert all(not m for m in masks[0][:-2])
140 +
141 +
142 +# ----- Hybrid coordinator (single-group worker, multi-group coordinator) -----
143 +
144 +
145 +def test_coordinator_hybrid_full_plus_swa_all_hit():
146 + groups = [
147 + KVCacheGroupSpec(["L0"], _full(16)),
148 + KVCacheGroupSpec(["L1"], _swa(16, 32)),
149 + ]
150 + coord = _make_coord(groups, hash_block_size=16)
151 + hs = _hashes(4)
152 + cmap = ExternalCachedBlockPool({(g, bytes(h)) for g in (0, 1) for h in hs})
153 + _masks, hit = coord.find_longest_cache_hit(
154 + hs, max_length=64, cached_block_pool=cmap
155 + )
156 + assert hit == 64
157 +
158 +
159 +def test_coordinator_hybrid_hole_in_full_clips_both():
160 + groups = [
161 + KVCacheGroupSpec(["L0"], _full(16)),
162 + KVCacheGroupSpec(["L1"], _swa(16, 32)),
163 + ]
164 + coord = _make_coord(groups, hash_block_size=16)
165 + hs = _hashes(4)
166 + exists = {(0, bytes(hs[0])), (0, bytes(hs[2])), (0, bytes(hs[3]))}
167 + exists |= {(1, bytes(h)) for h in hs}
168 + cmap = ExternalCachedBlockPool(exists)
169 + _masks, hit = coord.find_longest_cache_hit(
170 + hs, max_length=64, cached_block_pool=cmap
171 + )
172 + assert hit == 16
173 +
174 +
175 +def test_coordinator_group_block_size_double_hash():
176 + """Group block_size=32 over hash_block_size=16 hashes: adjacent
177 + hashes merge before pool lookup."""
178 + groups = [
179 + KVCacheGroupSpec(["L0"], _full(16)),
180 + KVCacheGroupSpec(["L1"], _full(32)),
181 + ]
182 + coord = _make_coord(groups, hash_block_size=16)
183 + hs = _hashes(4)
184 + big_hashes = list(BlockHashListWithBlockSize(hs, 16, 32))
185 + exists = {(0, bytes(h)) for h in hs}
186 + exists |= {(1, bytes(bh)) for bh in big_hashes}
187 + cmap = ExternalCachedBlockPool(exists)
188 + _masks, hit = coord.find_longest_cache_hit(
189 + hs, max_length=64, cached_block_pool=cmap
190 + )
191 + assert hit == 64
192 + assert hit % 32 == 0
193 +
194 +
195 +# ----- store_mask -----
196 +
197 +
198 +def test_store_mask_full_attention_all_true():
199 + groups = [KVCacheGroupSpec(["L0"], _full(16))]
200 + coord = _make_coord(groups, hash_block_size=16)
201 + masks = coord.store_mask(64)
202 + assert masks == ([True, True, True, True],)
203 +
204 +
205 +def test_store_mask_zero_aligned_returns_empty_per_group():
206 + groups = [
207 + KVCacheGroupSpec(["L0"], _full(16)),
208 + KVCacheGroupSpec(["L1"], _swa(16, 32)),
209 + ]
210 + coord = _make_coord(groups, hash_block_size=16)
211 + masks = coord.store_mask(0)
212 + assert masks == ([], [])
213 +
214 +
215 +def test_store_mask_swa_only_window_around_each_lcm_boundary():
216 + """Hybrid full-attn(block=32) + SWA(block=8, sw=8). lcm=32. With
217 + aligned=64 the SWA group should mark exactly the blocks ending at 32
218 + and 64 (i.e. blocks 3 and 7 at block_size=8); the rest can never
219 + participate in any future hit."""
220 + full = _full(32)
221 + swa = _swa(block_size=8, sliding_window=8)
222 + groups = [KVCacheGroupSpec(["L0"], full), KVCacheGroupSpec(["L1"], swa)]
223 + coord = _make_coord(groups, hash_block_size=8)
224 + masks = coord.store_mask(64)
225 + # Full-attn: 2 chunks * 32 tokens.
226 + assert masks[0] == [True, True]
227 + # SWA: 8 chunks * 8 tokens. Only chunks ending at 32 and 64 are stored.
228 + assert masks[1] == [False, False, False, True, False, False, False, True]
229 +
230 +
231 +def test_store_mask_swa_wider_window_covers_more_blocks_per_lcm():
232 + """Same hybrid layout but sliding_window=16 (= 2 SWA blocks). Each lcm
233 + boundary should now span two SWA tail blocks."""
234 + full = _full(32)
235 + swa = _swa(block_size=8, sliding_window=16)
236 + groups = [KVCacheGroupSpec(["L0"], full), KVCacheGroupSpec(["L1"], swa)]
237 + coord = _make_coord(groups, hash_block_size=8)
238 + masks = coord.store_mask(64)
239 + assert masks[0] == [True, True]
240 + # Boundary at 32: blocks ending in [16, 32) — chunks 2 and 3.
241 + # Boundary at 64: chunks 6 and 7. Others stay False.
242 + assert masks[1] == [False, False, True, True, False, False, True, True]
243 +
244 +
245 +def test_store_mask_dsv4_5_groups_full_mla_plus_4_swa():
246 + """DSV4-shaped: full-MLA(B=256) + 4 SWA groups with B in {64, 64, 4, 8}
247 + and varied sliding windows. lcm=256, hash_block_size=4. Two lcm segments
248 + (aligned_len=512). Validates that the tile-once strategy produces the
249 + expected per-segment tail-window pattern, repeated."""
250 + full_mla = _full(block_size=256)
251 + swa_64_sw128 = _swa(block_size=64, sliding_window=128)
252 + swa_64_sw512 = _swa(block_size=64, sliding_window=512)
253 + swa_4_sw16 = _swa(block_size=4, sliding_window=16)
254 + swa_8_sw64 = _swa(block_size=8, sliding_window=64)
255 + groups = [
256 + KVCacheGroupSpec(["L0"], full_mla),
257 + KVCacheGroupSpec(["L1"], swa_64_sw128),
258 + KVCacheGroupSpec(["L2"], swa_64_sw512),
259 + KVCacheGroupSpec(["L3"], swa_4_sw16),
260 + KVCacheGroupSpec(["L4"], swa_8_sw64),
261 + ]
262 + coord = _make_coord(groups, hash_block_size=4)
263 + assert coord.lcm_block_size == 256
264 + masks = coord.store_mask(512)
265 +
266 + # Full-MLA: 2 chunks of 256, both stored.
267 + assert masks[0] == [True, True]
268 + # SWA(64, sw=128): tail = ceil(127/64) = 2; C = 256/64 = 4.
269 + # Per-segment template = [F,F,T,T]; tiled twice.
270 + assert masks[1] == [False, False, True, True] * 2
271 + # SWA(64, sw=512): tail = 8 >= C = 4 → entire segment True.
272 + assert masks[2] == [True] * 8
273 + # SWA(4, sw=16): tail = ceil(15/4) = 4; C = 256/4 = 64.
274 + # Last 4 of each 64-chunk segment True.
275 + assert masks[3] == ([False] * 60 + [True] * 4) * 2
276 + # SWA(8, sw=64): tail = ceil(63/8) = 8; C = 256/8 = 32.
277 + # Last 8 of each 32-chunk segment True.
278 + assert masks[4] == ([False] * 24 + [True] * 8) * 2
279 +
280 +
281 +def test_store_mask_fast_path_all_block_sizes_equal_lcm():
282 + """When every non-full-attn group already aligns to lcm_block_size, the
283 + fast path returns all-True without invoking find_longest_cache_hit."""
284 + full = _full(block_size=64)
285 + swa = _swa(block_size=64, sliding_window=128)
286 + groups = [KVCacheGroupSpec(["L0"], full), KVCacheGroupSpec(["L1"], swa)]
287 + coord = _make_coord(groups, hash_block_size=64)
288 + assert coord.lcm_block_size == 64
289 + masks = coord.store_mask(256)
290 + # Every block in every group is True — no sub-lcm filtering possible.
291 + assert masks == ([True] * 4, [True] * 4)
292 +
293 +
294 +def test_store_mask_fast_path_single_attention_group():
295 + """Two groups sharing the same SWA spec collapse to one attention group;
296 + no lcm filter applies, every chunk is True."""
297 + swa = _swa(block_size=16, sliding_window=32)
298 + groups = [KVCacheGroupSpec(["L0"], swa), KVCacheGroupSpec(["L1"], swa)]
299 + coord = _make_coord(groups, hash_block_size=16)
300 + assert len(coord.attention_groups) == 1
301 + masks = coord.store_mask(64)
302 + assert masks == ([True] * 4, [True] * 4)
tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py new
+342
@@ -0,0 +1,342 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""End-to-end save->lookup test for MooncakeStoreConnector on a hybrid
4 +(SWA + Full) attention config, using a dict-backed mock store."""
5 +
6 +import sys
7 +import threading
8 +import types
9 +from unittest.mock import MagicMock, patch
10 +
11 +import torch
12 +
13 +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store import (
14 + worker as mooncake_store_worker,
15 +)
16 +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.coordinator import ( # noqa: E501
17 + MooncakeStoreCoordinator,
18 +)
19 +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
20 + ChunkedTokenDatabase,
21 + KeyMetadata,
22 + LoadSpec,
23 + ReqMeta,
24 +)
25 +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.worker import ( # noqa: E501
26 + KVCacheStoreRecvingThread,
27 + KVCacheStoreSendingThread,
28 +)
29 +from vllm.v1.core.kv_cache_utils import BlockHash
30 +from vllm.v1.kv_cache_interface import (
31 + FullAttentionSpec,
32 + KVCacheConfig,
33 + KVCacheGroupSpec,
34 + KVCacheTensor,
35 + SlidingWindowSpec,
36 +)
37 +
38 +
39 +class _DictStore:
40 + """In-memory MooncakeDistributedStore stand-in."""
41 +
42 + def __init__(self):
43 + self._data: dict[str, bytes] = {}
44 +
45 + def setup(self, *_args, **_kwargs):
46 + return 0
47 +
48 + def register_buffer(self, addr, length):
49 + return 0
50 +
51 + def batch_is_exist(self, keys):
52 + return [1 if k in self._data else 0 for k in keys]
53 +
54 + def batch_put_from_multi_buffers(self, keys, addrs, sizes, *_args, **_kwargs):
55 + for k in keys:
56 + self._data[k] = b"x"
57 + return [0] * len(keys)
58 +
59 + def batch_get_into_multi_buffers(self, keys, addrs, sizes, *_args, **_kwargs):
60 + return [0 if k in self._data else -1 for k in keys]
61 +
62 +
63 +def _minimal_vllm_config(cache_block_size=16):
64 + cfg = MagicMock()
65 + cfg.cache_config.block_size = cache_block_size
66 + cfg.cache_config.num_gpu_blocks = 4
67 + cfg.cache_config.hash_block_size = None
68 + cfg.cache_config.enable_prefix_caching = True
69 + cfg.parallel_config.prefill_context_parallel_size = 1
70 + cfg.parallel_config.decode_context_parallel_size = 1
71 + cfg.parallel_config.pipeline_parallel_size = 1
72 + cfg.parallel_config.world_size = 1
73 + cfg.parallel_config.rank = 0
74 + cfg.parallel_config.data_parallel_rank_local = 0
75 + cfg.parallel_config.data_parallel_size_local = 1
76 + cfg.kv_transfer_config.kv_role = "kv_both"
77 + cfg.kv_transfer_config.kv_connector_extra_config = {}
78 + cfg.kv_events_config = None
79 + cfg.model_config.model = "/tmp/m"
80 + cfg.model_config.use_mla = False
81 + cfg.model_config.get_num_layers.return_value = 2
82 + cfg.model_config.get_total_num_kv_heads.return_value = 8
83 + cfg.model_config.max_model_len = 4096
84 + cfg.scheduler_config.max_num_batched_tokens = 8192
85 + # Without this, MagicMock's truthy use_eagle() triggers the coordinator's
86 + # "use_eagle && nothing annotated → flag all groups" fallback.
87 + cfg.speculative_config = None
88 + return cfg
89 +
90 +
91 +def _build_worker_with_dict_store(vllm_config, kv_cache_config, store):
92 + """Build a MooncakeStoreWorker patching all distributed dependencies."""
93 + fake_mooncake_store = types.ModuleType("mooncake.store")
94 + fake_mooncake_store.MooncakeDistributedStore = lambda: store # type: ignore[attr-defined]
95 + fake_mooncake_store.ReplicateConfig = MagicMock # type: ignore[attr-defined]
96 + with (
97 + patch.dict(sys.modules, {"mooncake.store": fake_mooncake_store}),
98 + patch.object(mooncake_store_worker, "MooncakeStoreConfig") as MCfg,
99 + ):
100 + sc = MCfg.load_from_env.return_value
101 + sc.metadata_server = ""
102 + sc.global_segment_size = 1 << 20
103 + sc.local_buffer_size = 1 << 20
104 + sc.protocol = "tcp"
105 + sc.device_name = ""
106 + sc.master_server_address = ""
107 + sc.enable_offload = False
108 + with (
109 + patch(
110 + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake"
111 + ".store.worker.get_tensor_model_parallel_rank",
112 + return_value=0,
113 + ),
114 + patch(
115 + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake"
116 + ".store.worker.get_tensor_model_parallel_world_size",
117 + return_value=1,
118 + ),
119 + patch(
120 + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake"
121 + ".store.worker.get_pcp_group"
122 + ) as mock_pcp,
123 + patch(
124 + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake"
125 + ".store.worker.get_dcp_group"
126 + ) as mock_dcp,
127 + patch(
128 + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake"
129 + ".store.worker.get_ip",
130 + return_value="127.0.0.1",
131 + ),
132 + ):
133 + mock_pcp.return_value.world_size = 1
134 + mock_dcp.return_value.world_size = 1
135 + worker = mooncake_store_worker.MooncakeStoreWorker(
136 + vllm_config, kv_cache_config=kv_cache_config
137 + )
138 + return worker
139 +
140 +
141 +def test_e2e_swa_plus_full_save_then_lookup_hits():
142 + """
143 + E2E: build a SWA+Full hybrid worker, save all blocks via the sending
144 + thread (synchronously), then verify lookup returns the full hit length.
145 + Also verify that evicting SWA's early blocks (outside its window) still
146 + allows a full hit because the window covers the tail.
147 + """
148 + full = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None)
149 + swa = SlidingWindowSpec(
150 + block_size=16,
151 + num_kv_heads=8,
152 + head_size=64,
153 + dtype=None,
154 + sliding_window=32,
155 + )
156 + cfg = KVCacheConfig(
157 + num_blocks=4,
158 + kv_cache_tensors=[
159 + KVCacheTensor(size=8192, shared_by=["L0"]),
160 + KVCacheTensor(size=8192, shared_by=["L1"]),
161 + ],
162 + kv_cache_groups=[
163 + KVCacheGroupSpec(["L0"], full),
164 + KVCacheGroupSpec(["L1"], swa),
165 + ],
166 + )
167 + vllm_config = _minimal_vllm_config(cache_block_size=16)
168 + store = _DictStore()
169 +
170 + worker = _build_worker_with_dict_store(vllm_config, cfg, store)
171 + worker.tp_size = 1
172 + worker.pp_size = 1
173 + worker.put_step = 1
174 + worker.num_kv_head = 8
175 +
176 + # Register kv_caches using mocked thread classes so register_kv_caches
177 + # doesn't try to start real background threads (which set ready_event).
178 + kv_caches = {
179 + "L0": torch.zeros(2, 4, 8, 8, 64),
180 + "L1": torch.zeros(2, 4, 8, 8, 64),
181 + }
182 +
183 + def _fake_thread_init(*args, **kwargs):
184 + """Mock thread that sets all threading.Event args so waits don't block."""
185 + for v in list(args) + list(kwargs.values()):
186 + if isinstance(v, threading.Event):
187 + v.set()
188 + m = MagicMock()
189 + m.start = lambda: None
190 + return m
191 +
192 + with (
193 + patch.object(
194 + mooncake_store_worker,
195 + "KVCacheStoreSendingThread",
196 + side_effect=_fake_thread_init,
197 + ),
198 + patch.object(
199 + mooncake_store_worker,
200 + "KVCacheStoreRecvingThread",
201 + side_effect=_fake_thread_init,
202 + ),
203 + ):
204 + worker.register_kv_caches(kv_caches)
205 +
206 + # Now build a real sending thread (no .start()) over the worker's token_dbs
207 + # and the dict-backed store, so _handle_request runs synchronously.
208 + ready = threading.Event()
209 + send_thread = KVCacheStoreSendingThread(
210 + store=store,
211 + token_databases=worker.token_dbs,
212 + block_size=worker.block_size,
213 + coord=worker.coord,
214 + tp_rank=worker.tp_rank,
215 + put_step=worker.put_step,
216 + kv_role=worker.kv_role,
217 + ready_event=ready,
218 + enable_kv_event=False,
219 + )
220 +
221 + hs = [BlockHash(bytes([i + 1]) * 4) for i in range(4)]
222 + save_req = ReqMeta(
223 + req_id="r0",
224 + token_len_chunk=64,
225 + block_ids=([0, 1, 2, 3], [0, 1, 2, 3]),
226 + block_hashes=hs,
227 + can_save=True,
228 + )
229 + send_thread.add_stored_request("r0")
230 + # Put the request in the queue so task_done() doesn't underflow.
231 + send_thread.request_queue.put(save_req)
232 + req = send_thread.request_queue.get()
233 + send_thread._handle_request(req)
234 +
235 + # Point worker.store at the dict store (the worker constructor captured
236 + # the MagicMock; replace with the real dict store for lookup).
237 + worker.store = store
238 +
239 + # Both groups stored all 4 blocks -> full hit.
240 + assert worker.lookup(token_len=64, block_hashes=hs) == 64
241 +
242 + # Evict SWA's first two blocks (outside its window of 32 tokens = 2 blocks).
243 + swa_keys_outside_window = [
244 + k
245 + for k in list(store._data.keys())
246 + if "@group:1" in k and (("@" + hs[0].hex()) in k or ("@" + hs[1].hex()) in k)
247 + ]
248 + for k in swa_keys_outside_window:
249 + del store._data[k]
250 +
251 + # SWA window=32 -> only last 2 blocks must be present in SWA group.
252 + # Full has all 4. Coordinator should still return 64.
253 + assert worker.lookup(token_len=64, block_hashes=hs) == 64
254 +
255 +
256 +def test_recv_skips_swa_blocks_before_window():
257 + """Producer stored every block for both groups; consumer must only fetch
258 + SWA blocks within the sliding window, not the head."""
259 + full = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None)
260 + # sliding_window=32, block_size=16 → 2 contiguous blocks within window.
261 + swa = SlidingWindowSpec(
262 + block_size=16,
263 + num_kv_heads=8,
264 + head_size=64,
265 + dtype=None,
266 + sliding_window=32,
267 + )
268 + groups = [
269 + KVCacheGroupSpec(["L0"], full),
270 + KVCacheGroupSpec(["L1"], swa),
271 + ]
272 + md0 = KeyMetadata("m", 0, 0, 0, 0, group_id=0)
273 + md1 = KeyMetadata("m", 0, 0, 0, 0, group_id=1)
274 + db_full = ChunkedTokenDatabase(md0, block_size=16, hash_block_size=16)
275 + db_swa = ChunkedTokenDatabase(md1, block_size=16, hash_block_size=16)
276 + db_full.set_kv_caches_base_addr([0])
277 + db_full.set_block_len([1024])
278 + db_swa.set_kv_caches_base_addr([1 << 20])
279 + db_swa.set_block_len([1024])
280 +
281 + requested_keys: list[str] = []
282 +
283 + class _CapturingStore:
284 + def batch_get_into_multi_buffers(self, keys, addrs, sizes):
285 + requested_keys.extend(keys)
286 + return [0] * len(keys)
287 +
288 + ready = threading.Event()
289 + coord = MooncakeStoreCoordinator(
290 + groups, scheduler_block_size=16, hash_block_size=16
291 + )
292 + recv = KVCacheStoreRecvingThread(
293 + store=_CapturingStore(),
294 + token_databases=[db_full, db_swa],
295 + block_size=16,
296 + tp_rank=0,
297 + ready_event=ready,
298 + coord=coord,
299 + )
300 +
301 + hs = [BlockHash(bytes([i + 1]) * 4) for i in range(4)]
302 + req = ReqMeta(
303 + req_id="r0",
304 + token_len_chunk=64,
305 + block_ids=([0, 1, 2, 3], [0, 1, 2, 3]),
306 + block_hashes=hs,
307 + load_spec=LoadSpec(
308 + vllm_cached_tokens=0, kvpool_cached_tokens=64, can_load=True, token_len=64
309 + ),
310 + )
311 + recv.request_queue.put(req)
312 + recv._handle_request(recv.request_queue.get())
313 +
314 + full_keys = [k for k in requested_keys if "@group:0" in k]
315 + swa_keys = [k for k in requested_keys if "@group:1" in k]
316 + # Full attention: load every block (4).
317 + assert len(full_keys) == 4
318 + # SWA: only the 2 tail-window blocks (hashes hs[2], hs[3]).
319 + assert len(swa_keys) == 2
320 + swa_hashes = {k.rsplit("@", 1)[-1] for k in swa_keys}
321 + assert swa_hashes == {hs[2].hex(), hs[3].hex()}
322 +
323 +
324 +def test_chunked_token_database_hash_block_size_smaller_than_block_size():
325 + """DSv4-style: hash_block_size=4, group block_size=16 — process_tokens
326 + must merge every 4 fine hashes into one chunk hash via
327 + BlockHashListWithBlockSize."""
328 + md = KeyMetadata("m", 0, 0, 0, 0, group_id=3)
329 + db = ChunkedTokenDatabase(md, block_size=16, hash_block_size=4)
330 + db.set_kv_caches_base_addr([0])
331 + db.set_block_len([512])
332 + # 8 fine-grained hashes (32 tokens at hash_block_size=4) → 2 group chunks.
333 + fine_hashes = [BlockHash(bytes([i + 1]) * 4) for i in range(8)]
334 + out = list(db.process_tokens(token_len=32, block_hashes=fine_hashes))
335 + assert len(out) == 2
336 + assert out[0][0] == 0 and out[0][1] == 16
337 + assert out[1][0] == 16 and out[1][1] == 32
338 + # Each chunk's hash is the concatenation of 4 fine hashes.
339 + expected0 = b"".join(fine_hashes[0:4]).hex()
340 + expected1 = b"".join(fine_hashes[4:8]).hex()
341 + assert out[0][2].chunk_hash == expected0
342 + assert out[1][2].chunk_hash == expected1
tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py new
+111
@@ -0,0 +1,111 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from types import SimpleNamespace
5 +
6 +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
7 + RequestTracker,
8 +)
9 +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.scheduler import (
10 + MooncakeStoreScheduler,
11 +)
12 +
13 +
14 +def _make_bare_scheduler() -> MooncakeStoreScheduler:
15 + scheduler = object.__new__(MooncakeStoreScheduler)
16 + scheduler.kv_role = "kv_both"
17 + scheduler.original_block_size = 16
18 + scheduler._block_size = 16
19 + scheduler._discard_partial_chunks = True
20 + scheduler.load_specs = {}
21 + scheduler._preempted_req_ids = set()
22 + scheduler._unfinished_request_ids = {"req-0"}
23 + scheduler._unfinished_requests = {}
24 + scheduler._request_trackers = {}
25 + return scheduler
26 +
27 +
28 +def _make_scheduler_output(*, scheduled_spec_tokens: list[int] | None):
29 + return SimpleNamespace(
30 + finished_req_ids=set(),
31 + preempted_req_ids=set(),
32 + scheduled_new_reqs=[],
33 + scheduled_cached_reqs=SimpleNamespace(
34 + req_ids=["req-0"],
35 + new_block_ids=[([2],)],
36 + num_computed_tokens=[44],
37 + ),
38 + num_scheduled_tokens={"req-0": 4},
39 + scheduled_spec_decode_tokens=(
40 + {"req-0": scheduled_spec_tokens} if scheduled_spec_tokens else {}
41 + ),
42 + )
43 +
44 +
45 +def _add_unfinished_request(
46 + scheduler: MooncakeStoreScheduler,
47 + *,
48 + token_ids: list[int],
49 + block_hashes: list[bytes],
50 + prefill_end_tokens: int,
51 +) -> None:
52 + request = SimpleNamespace(
53 + all_token_ids=token_ids,
54 + block_hashes=block_hashes,
55 + num_output_placeholders=0,
56 + )
57 + scheduler._unfinished_requests["req-0"] = (request, ([0, 1],))
58 + scheduler._request_trackers["req-0"] = RequestTracker(
59 + req_id="req-0",
60 + token_len=44,
61 + allocated_block_ids=([0, 1],),
62 + num_saved_tokens=32,
63 + token_ids=token_ids[:44],
64 + prefill_end_tokens=prefill_end_tokens,
65 + )
66 +
67 +
68 +def test_cached_request_with_spec_decode_does_not_save_scheduled_drafts():
69 + # Drafts in scheduled_spec_decode_tokens are not appended to all_token_ids
70 + # yet, so the tracker's token_len does not advance and num_tokens_to_save
71 + # stays below chunk_boundary — the save is naturally skipped.
72 + scheduler = _make_bare_scheduler()
73 + _add_unfinished_request(
74 + scheduler,
75 + token_ids=list(range(44)),
76 + block_hashes=[b"h0", b"h1"],
77 + prefill_end_tokens=48,
78 + )
79 +
80 + meta = scheduler.build_connector_meta(
81 + _make_scheduler_output(scheduled_spec_tokens=[101, 102, 103])
82 + )
83 +
84 + assert meta.requests == []
85 + tracker = scheduler._request_trackers["req-0"]
86 + assert tracker.token_len == 44
87 + assert tracker.num_saved_tokens == 32
88 + assert tracker.allocated_block_ids == ([0, 1, 2],)
89 +
90 +
91 +def test_cached_request_without_spec_decode_keeps_current_step_save_overlap():
92 + scheduler = _make_bare_scheduler()
93 + _add_unfinished_request(
94 + scheduler,
95 + token_ids=list(range(48)),
96 + block_hashes=[b"h0", b"h1", b"h2"],
97 + prefill_end_tokens=48,
98 + )
99 +
100 + meta = scheduler.build_connector_meta(
101 + _make_scheduler_output(scheduled_spec_tokens=None)
102 + )
103 +
104 + assert len(meta.requests) == 1
105 + req_meta = meta.requests[0]
106 + assert req_meta.req_id == "req-0"
107 + assert req_meta.can_save is True
108 + assert req_meta.token_len_chunk == 48
109 + tracker = scheduler._request_trackers["req-0"]
110 + assert tracker.token_len == 48
111 + assert tracker.num_saved_tokens == 48
tests/v1/kv_connector/unit/test_mooncake_store_worker.py
+357 -117
@@ -19,7 +19,10 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake import (
19 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store import (
20 worker,
21 )
22 -from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( # noqa: E501
22 +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store import (
23 + worker as mooncake_store_worker,
24 +)
25 +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
26 ChunkedTokenDatabase,
27 KeyMetadata,
28 LoadSpec,
@@ -27,20 +30,37 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
30 )
31
32
33 +def _default_send_coord() -> mooncake_store_worker.MooncakeStoreCoordinator:
34 + from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheGroupSpec
35 +
36 + spec = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None)
37 + return mooncake_store_worker.MooncakeStoreCoordinator(
38 + [KVCacheGroupSpec(["layer0"], spec)],
39 + scheduler_block_size=16,
40 + hash_block_size=16,
41 + )
42 +
43 +
44 def _make_store_sending_thread(
45 store: MagicMock,
46 *,
47 + coord: mooncake_store_worker.MooncakeStoreCoordinator | None = None,
48 + token_databases: list[ChunkedTokenDatabase] | None = None,
49 + block_size: int = 16,
50 replicate_config: object | None = None,
34 -) -> worker.KVCacheStoreSendingThread:
35 - token_database = ChunkedTokenDatabase(
36 - KeyMetadata("test-model", 0, 0, 0, 0), block_size=16
37 - )
38 - token_database.set_kv_caches_base_addr([0x1000])
39 - token_database.set_block_len([256])
40 - thread = worker.KVCacheStoreSendingThread(
51 +) -> mooncake_store_worker.KVCacheStoreSendingThread:
52 + if coord is None:
53 + coord = _default_send_coord()
54 + if token_databases is None:
55 + db = ChunkedTokenDatabase(KeyMetadata("test-model", 0, 0, 0, 0), block_size=16)
56 + db.set_kv_caches_base_addr([0x1000])
57 + db.set_block_len([256])
58 + token_databases = [db]
59 + thread = mooncake_store_worker.KVCacheStoreSendingThread(
60 store=store,
42 - token_database=token_database,
43 - block_size=16,
61 + token_databases=token_databases,
62 + block_size=block_size,
63 + coord=coord,
64 tp_rank=0,
65 put_step=1,
66 kv_role="kv_producer",
@@ -55,18 +75,27 @@ def _make_store_recving_thread(
75 store: MagicMock,
76 *,
77 disk_offload_buffer_budget_bytes: int | None = None,
58 -) -> worker.KVCacheStoreRecvingThread:
78 +) -> mooncake_store_worker.KVCacheStoreRecvingThread:
79 + from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheGroupSpec
80 +
81 token_database = ChunkedTokenDatabase(
82 KeyMetadata("test-model", 0, 0, 0, 0), block_size=16
83 )
84 token_database.set_kv_caches_base_addr([0x1000])
85 token_database.set_block_len([256])
64 - thread = worker.KVCacheStoreRecvingThread(
86 + spec = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None)
87 + coord = mooncake_store_worker.MooncakeStoreCoordinator(
88 + [KVCacheGroupSpec(["layer0"], spec)],
89 + scheduler_block_size=16,
90 + hash_block_size=16,
91 + )
92 + thread = mooncake_store_worker.KVCacheStoreRecvingThread(
93 store=store,
66 - token_database=token_database,
94 + token_databases=[token_database],
95 block_size=16,
96 tp_rank=0,
97 ready_event=threading.Event(),
98 + coord=coord,
99 disk_offload_buffer_budget_bytes=disk_offload_buffer_budget_bytes,
100 )
101 thread.request_queue.task_done = MagicMock()
@@ -83,7 +112,7 @@ def _make_load_req(
112 return ReqMeta(
113 req_id=req_id,
114 token_len_chunk=token_len,
86 - block_ids=list(range(len(block_hashes))),
115 + block_ids=(list(range(len(block_hashes))),),
116 block_hashes=block_hashes,
117 load_spec=LoadSpec(
118 vllm_cached_tokens=vllm_cached_tokens,
@@ -98,7 +127,7 @@ def _make_store_req(req_id: str, block_hashes: list[bytes]) -> ReqMeta:
127 return ReqMeta(
128 req_id=req_id,
129 token_len_chunk=32,
101 - block_ids=[0, 1],
130 + block_ids=([0, 1],),
131 block_hashes=block_hashes,
132 can_save=True,
133 original_block_size=16,
@@ -149,10 +178,31 @@ def _make_vllm_config(
178 parallel_config=SimpleNamespace(
179 pipeline_parallel_size=1,
180 rank=0,
181 + decode_context_parallel_size=1,
182 + prefill_context_parallel_size=1,
183 ),
184 kv_transfer_config=_FakeKVTransferConfig(extra_config=extra_config),
185 cache_config=SimpleNamespace(block_size=16, num_gpu_blocks=10),
186 kv_events_config=SimpleNamespace(enable_kv_cache_events=False),
187 + speculative_config=None,
188 + )
189 +
190 +
191 +def _make_kv_cache_config(*, block_size: int = 16) -> object:
192 + """Minimal single-group KVCacheConfig for topology tests."""
193 + from vllm.v1.kv_cache_interface import (
194 + FullAttentionSpec,
195 + KVCacheConfig,
196 + KVCacheGroupSpec,
197 + )
198 +
199 + spec = FullAttentionSpec(
200 + block_size=block_size, num_kv_heads=8, head_size=64, dtype=None
201 + )
202 + return KVCacheConfig(
203 + num_blocks=10,
204 + kv_cache_tensors=[],
205 + kv_cache_groups=[KVCacheGroupSpec(["layer0"], spec)],
206 )
207
208
@@ -445,9 +495,9 @@ def test_recv_thread_uses_single_batch_when_no_disk_offload_budget(monkeypatch):
495 assert store.batch_get_into_multi_buffers.call_count == 1
496 keys, addrs, sizes = store.batch_get_into_multi_buffers.call_args.args
497 assert keys == [
448 - "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6130",
449 - "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6131",
450 - "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6132",
498 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6130",
499 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6131",
500 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6132",
501 ]
502 assert sizes == [[256], [256], [256]]
503 store.batch_get_replica_desc.assert_not_called()
@@ -467,9 +517,9 @@ def test_recv_thread_logs_tier_summary_when_enabled(monkeypatch, caplog_vllm):
517 token_len=48,
518 )
519 expected_keys = [
470 - "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6130",
471 - "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6131",
472 - "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6132",
520 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6130",
521 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6131",
522 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6132",
523 ]
524 store.batch_get_replica_desc.return_value = {
525 expected_keys[0]: [_ReplicaDesc("memory")],
@@ -521,10 +571,10 @@ def test_recv_thread_uses_ratio_scaled_budget_for_first_pass_split():
571 first_keys = store.batch_get_into_multi_buffers.call_args_list[0].args[0]
572 second_keys = store.batch_get_into_multi_buffers.call_args_list[1].args[0]
573 assert first_keys == [
524 - "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6130",
574 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6130",
575 ]
576 assert second_keys == [
527 - "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6131",
577 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6131",
578 ]
579
580
@@ -556,14 +606,14 @@ def test_recv_thread_splits_disk_offload_loads_by_budget():
606 first_sizes = store.batch_get_into_multi_buffers.call_args_list[0].args[2]
607 second_sizes = store.batch_get_into_multi_buffers.call_args_list[1].args[2]
608 assert first_keys == [
559 - "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6130",
560 - "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6131",
609 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6130",
610 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6131",
611 ]
612 assert second_keys == [
563 - "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6132",
613 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6132",
614 ]
565 - base_addr = thread.token_database.kv_caches_base_addr[0]
566 - block_len = thread.token_database.block_len[0]
615 + base_addr = thread.token_databases[0].kv_caches_base_addr[0]
616 + block_len = thread.token_databases[0].block_len[0]
617 assert first_addrs == [[base_addr], [base_addr + block_len]]
618 assert second_addrs == [[base_addr + 2 * block_len]]
619 expected_size = block_len
@@ -610,9 +660,9 @@ def test_recv_thread_skips_split_when_budget_holds_all_keys():
660
661 assert store.batch_get_into_multi_buffers.call_count == 1
662 assert store.batch_get_into_multi_buffers.call_args_list[0].args[0] == [
613 - "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6130",
614 - "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6131",
615 - "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@6132",
663 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6130",
664 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6131",
665 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6132",
666 ]
667
668
@@ -654,7 +704,7 @@ def test_requester_worker_init_uses_positional_setup(tmp_path, monkeypatch):
704 },
705 ),
706 )
657 - w = worker.MooncakeStoreWorker(_make_vllm_config())
707 + w = worker.MooncakeStoreWorker(_make_vllm_config(), _make_kv_cache_config())
708
709 assert not hasattr(w, "_isolate_offload_resources")
710 assert store.setup.call_args.args == (
@@ -690,7 +740,7 @@ def test_requester_worker_init_prefers_local_hostname_override(
740 },
741 ),
742 )
693 - worker.MooncakeStoreWorker(_make_vllm_config())
743 + worker.MooncakeStoreWorker(_make_vllm_config(), _make_kv_cache_config())
744
745 assert store.setup.call_args.args[0] == "worker-a:50053"
746
@@ -718,7 +768,7 @@ def test_requester_worker_init_skips_disk_budget_when_offload_disabled(
768 },
769 ),
770 )
721 - w = worker.MooncakeStoreWorker(_make_vllm_config())
771 + w = worker.MooncakeStoreWorker(_make_vllm_config(), _make_kv_cache_config())
772
773 assert w.disk_offload_buffer_budget_bytes is None
774
@@ -748,7 +798,8 @@ def test_requester_worker_init_builds_replicate_config_for_preferred_segment(
798 extra_config={
799 "preferred_segment": "10.0.0.7:50053",
800 }
751 - )
801 + ),
802 + _make_kv_cache_config(),
803 )
804
805 assert isinstance(w.store_replicate_config, fake_replicate_config_cls)
@@ -760,6 +811,158 @@ def test_requester_worker_init_builds_replicate_config_for_preferred_segment(
811 # ---------------------------------------------------------------------------
812
813
814 +def test_store_sending_thread_clamps_token_len_to_lcm():
815 + """Partial chunks past the last lcm boundary aren't stored — cache hits
816 + are always lcm-aligned (mirrors HybridKVCacheCoordinator)."""
817 + store = MagicMock()
818 + store.batch_is_exist.return_value = [0, 0]
819 + store.batch_put_from_multi_buffers.return_value = [256, 256]
820 + # Default coord: single full-attn block_size=16, lcm=16.
821 + # token_len_chunk=33 clamps to 32 → 2 chunks (not 3 with a partial 1-token chunk).
822 + thread = _make_store_sending_thread(store)
823 +
824 + thread.add_stored_request("r0")
825 + thread._handle_request(
826 + ReqMeta(
827 + req_id="r0",
828 + token_len_chunk=33,
829 + block_ids=([0, 1, 2],),
830 + block_hashes=[b"a0", b"a1", b"a2"],
831 + can_save=True,
832 + original_block_size=16,
833 + )
834 + )
835 +
836 + keys = store.batch_put_from_multi_buffers.call_args.args[0]
837 + assert len(keys) == 2
838 +
839 +
840 +def test_store_sending_thread_skips_when_token_len_below_lcm():
841 + """Requests shorter than lcm_block_size cannot produce any aligned chunk,
842 + so neither the existence check nor the put should be issued."""
843 + from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheGroupSpec
844 +
845 + store = MagicMock()
846 + # lcm=64 via single full-attn block_size=64.
847 + spec = FullAttentionSpec(block_size=64, num_kv_heads=8, head_size=64, dtype=None)
848 + coord = mooncake_store_worker.MooncakeStoreCoordinator(
849 + [KVCacheGroupSpec(["L"], spec)],
850 + scheduler_block_size=64,
851 + hash_block_size=64,
852 + )
853 + db = ChunkedTokenDatabase(
854 + KeyMetadata("test-model", 0, 0, 0, 0, group_id=0),
855 + block_size=64,
856 + hash_block_size=64,
857 + )
858 + db.set_kv_caches_base_addr([0x1000])
859 + db.set_block_len([1024])
860 + thread = _make_store_sending_thread(
861 + store, coord=coord, token_databases=[db], block_size=64
862 + )
863 +
864 + thread.add_stored_request("r0")
865 + thread._handle_request(
866 + ReqMeta(
867 + req_id="r0",
868 + token_len_chunk=32,
869 + block_ids=([0, 1],),
870 + block_hashes=[b"a0", b"a1"],
871 + can_save=True,
872 + original_block_size=64,
873 + )
874 + )
875 +
876 + store.batch_is_exist.assert_not_called()
877 + store.batch_put_from_multi_buffers.assert_not_called()
878 + assert thread.stored_requests["r0"] == 0
879 +
880 +
881 +def test_store_sending_thread_only_stores_swa_blocks_in_window():
882 + """For SWA groups, only blocks within ``sliding_window`` of an
883 + lcm-aligned boundary are stored — a block outside every such window
884 + can never participate in any future hit (mirrors recv-side masking).
885 +
886 + Setup: full-attn (block_size=32) + SWA (block_size=8, sliding_window=8)
887 + → lcm=32. With token_len=64 there are two LCM boundaries (32, 64);
888 + the SWA group should store exactly one block per boundary: the block
889 + ending at 32 and the block ending at 64.
890 + """
891 + from vllm.v1.kv_cache_interface import (
892 + FullAttentionSpec,
893 + KVCacheGroupSpec,
894 + SlidingWindowSpec,
895 + )
896 +
897 + store = MagicMock()
898 + store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
899 + store.batch_put_from_multi_buffers.side_effect = lambda keys, addrs, sizes: (
900 + [256] * len(keys)
901 + )
902 +
903 + full_spec = FullAttentionSpec(
904 + block_size=32, num_kv_heads=8, head_size=64, dtype=None
905 + )
906 + swa_spec = SlidingWindowSpec(
907 + block_size=8,
908 + num_kv_heads=8,
909 + head_size=64,
910 + dtype=None,
911 + sliding_window=8,
912 + )
913 + coord = mooncake_store_worker.MooncakeStoreCoordinator(
914 + [KVCacheGroupSpec(["L0"], full_spec), KVCacheGroupSpec(["L1"], swa_spec)],
915 + scheduler_block_size=32,
916 + hash_block_size=8,
917 + )
918 +
919 + db_full = ChunkedTokenDatabase(
920 + KeyMetadata("test-model", 0, 0, 0, 0, group_id=0),
921 + block_size=32,
922 + hash_block_size=8,
923 + )
924 + db_full.set_kv_caches_base_addr([0x1000])
925 + db_full.set_block_len([512])
926 + db_swa = ChunkedTokenDatabase(
927 + KeyMetadata("test-model", 0, 0, 0, 0, group_id=1),
928 + block_size=8,
929 + hash_block_size=8,
930 + )
931 + db_swa.set_kv_caches_base_addr([0x2000])
932 + db_swa.set_block_len([128])
933 +
934 + thread = _make_store_sending_thread(
935 + store,
936 + coord=coord,
937 + token_databases=[db_full, db_swa],
938 + block_size=32,
939 + )
940 +
941 + hs = [bytes([i + 1]) * 4 for i in range(8)]
942 + thread.add_stored_request("r0")
943 + thread._handle_request(
944 + ReqMeta(
945 + req_id="r0",
946 + token_len_chunk=64,
947 + block_ids=([0, 1], list(range(8))),
948 + block_hashes=hs,
949 + can_save=True,
950 + original_block_size=32,
951 + )
952 + )
953 +
954 + keys = store.batch_put_from_multi_buffers.call_args.args[0]
955 + full_keys = [k for k in keys if "@group:0" in k]
956 + swa_keys = [k for k in keys if "@group:1" in k]
957 + # Full-attn: 2 blocks (chunks ending at 32 and 64).
958 + assert len(full_keys) == 2
959 + # SWA: only the two blocks ending at lcm boundaries (32 and 64), i.e.
960 + # blocks covering tokens [24,32) and [56,64) — hashes hs[3] and hs[7].
961 + assert len(swa_keys) == 2
962 + swa_hashes = {k.rsplit("@", 1)[-1] for k in swa_keys}
963 + assert swa_hashes == {hs[3].hex(), hs[7].hex()}
964 +
965 +
966 def _auto_set_ready_event(*args, **kwargs):
967 """Side effect for mocked thread constructors that auto-sets ready_event."""
968 for arg in args:
@@ -771,37 +974,108 @@ def _auto_set_ready_event(*args, **kwargs):
974 return MagicMock()
975
976
977 +def _register_with_mocked_threads(
978 + worker: mooncake_store_worker.MooncakeStoreWorker,
979 + kv_caches: dict[str, torch.Tensor],
980 +) -> None:
981 + """Call register_kv_caches with the I/O transfer threads mocked out."""
982 + prefix = "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.worker."
983 + with (
984 + patch(prefix + "KVCacheStoreSendingThread", side_effect=_auto_set_ready_event),
985 + patch(prefix + "KVCacheStoreRecvingThread", side_effect=_auto_set_ready_event),
986 + ):
987 + worker.register_kv_caches(kv_caches)
988 +
989 +
990 def _make_bare_worker(
991 *,
992 num_gpu_blocks: int = 10,
993 block_size: int = 16,
994 kv_role: str = "kv_both",
779 -) -> worker.MooncakeStoreWorker:
995 +) -> mooncake_store_worker.MooncakeStoreWorker:
996 """Construct a MooncakeStoreWorker via __new__, bypassing __init__.
997
998 Sets only the attributes that register_kv_caches() reads so we can
999 test the stride-based layout detection without a real
1000 MooncakeDistributedStore.
1001 """
786 - w = object.__new__(worker.MooncakeStoreWorker)
787 - w.cache_config = MagicMock()
788 - w.cache_config.num_gpu_blocks = num_gpu_blocks
789 - w.store = MagicMock()
790 - w.store.register_buffer.return_value = 0
791 - w.use_mla = False
792 - w.token_database = ChunkedTokenDatabase(
793 - KeyMetadata("test-model", 0, 0, 0, 0), block_size=block_size
794 - )
795 - w.kv_role = kv_role
796 - w.block_size = block_size
797 - w.tp_rank = 0
798 - w.put_step = 1
799 - w.enable_kv_events = False
800 - w.disk_offload_buffer_budget_bytes = None
801 - w.kv_send_thread = None
802 - w.kv_recv_thread = None
803 - w.store_replicate_config = SimpleNamespace()
804 - return w
1002 + worker = object.__new__(mooncake_store_worker.MooncakeStoreWorker)
1003 + worker.cache_config = MagicMock()
1004 + worker.cache_config.num_gpu_blocks = num_gpu_blocks
1005 + worker.store = MagicMock()
1006 + worker.store.register_buffer.return_value = 0
1007 + worker.use_mla = False
1008 + worker.kv_role = kv_role
1009 + worker.block_size = block_size
1010 + worker.tp_rank = 0
1011 + worker.put_step = 1
1012 + worker.enable_kv_events = False
1013 + worker.kv_send_thread = None
1014 + worker.kv_recv_thread = None
1015 + worker.tp_size = 1
1016 + worker.num_kv_head = 1
1017 + worker.pp_size = 1
1018 + # Minimal single-full-attention-group config so the coordinator-based
1019 + # lookup path works (the connector no longer carries a legacy single-group
1020 + # path; everything flows through the coordinator).
1021 + from vllm.v1.kv_cache_interface import (
1022 + FullAttentionSpec,
1023 + KVCacheGroupSpec,
1024 + )
1025 +
1026 + worker.disk_offload_buffer_budget_bytes = None
1027 + worker.store_replicate_config = SimpleNamespace()
1028 +
1029 + spec = FullAttentionSpec(
1030 + block_size=block_size, num_kv_heads=8, head_size=64, dtype=None
1031 + )
1032 + group = KVCacheGroupSpec(["layer0", "__cross_layer__"], spec)
1033 + worker._kv_cache_groups = [group]
1034 + worker.pcp_size = 1
1035 + worker.dcp_size = 1
1036 + worker.hash_block_size = block_size
1037 + worker.metadata = KeyMetadata("test-model", 0, 0, 0, 0)
1038 + # Pre-build a single-group token_dbs so lookup-only tests don't have to
1039 + # call register_kv_caches.
1040 + worker.token_dbs = [
1041 + ChunkedTokenDatabase(
1042 + KeyMetadata("test-model", 0, 0, 0, 0, group_id=0),
1043 + block_size=block_size,
1044 + hash_block_size=block_size,
1045 + )
1046 + ]
1047 + worker.coord = mooncake_store_worker.MooncakeStoreCoordinator(
1048 + worker._kv_cache_groups,
1049 + scheduler_block_size=block_size,
1050 + hash_block_size=block_size,
1051 + )
1052 + return worker
1053 +
1054 +
1055 +def test_lookup_partial_prefix_returns_first_hit_length():
1056 + worker = _make_bare_worker()
1057 + worker.store.batch_is_exist.return_value = [1, 1, 0]
1058 + assert worker.lookup(48, [b"a0", b"a1", b"a2"]) == 32
1059 +
1060 +
1061 +def test_lookup_swa_single_group_returns_full_when_tail_window_present():
1062 + """Single-SWA, sliding_window=32 (= 2 blocks): producer stored only the
1063 + tail. Coordinator-driven lookup returns full prefix even though the
1064 + pre-window blocks are absent."""
1065 + from vllm.v1.kv_cache_interface import KVCacheGroupSpec, SlidingWindowSpec
1066 +
1067 + worker = _make_bare_worker(block_size=16)
1068 + swa = SlidingWindowSpec(
1069 + block_size=16, num_kv_heads=8, head_size=64, dtype=None, sliding_window=32
1070 + )
1071 + worker._kv_cache_groups = [KVCacheGroupSpec(["layer0"], swa)]
1072 + worker.coord = mooncake_store_worker.MooncakeStoreCoordinator(
1073 + worker._kv_cache_groups,
1074 + scheduler_block_size=worker.hash_block_size,
1075 + hash_block_size=worker.hash_block_size,
1076 + )
1077 + worker.store.batch_is_exist.return_value = [0, 0, 1, 1]
1078 + assert worker.lookup(64, [b"h0", b"h1", b"h2", b"h3"]) == 64
1079
1080
1081 # ---------------------------------------------------------------------------
@@ -812,34 +1086,17 @@ def _make_bare_worker(
1086 def test_register_kv_caches_blocks_first_single_segment():
1087 """Blocks-first layout (FlashInfer/MLA): one segment per layer."""
1088 num_blocks = 10
815 - page_size_elements = 64 # elements per block
816 - w = _make_bare_worker(num_gpu_blocks=num_blocks)
1089 + page_size_elements = 64
1090 + worker = _make_bare_worker(num_gpu_blocks=num_blocks)
1091
1092 # Shape: (num_blocks, page_size_elements) — blocks outermost, no outer_dims
1093 tensor = torch.zeros(num_blocks, page_size_elements, dtype=torch.float16)
1094 + _register_with_mocked_threads(worker, {"layer0": tensor})
1095
821 - with (
822 - patch(
823 - "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
824 - "worker.KVCacheStoreSendingThread",
825 - side_effect=_auto_set_ready_event,
826 - ),
827 - patch(
828 - "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
829 - "worker.KVCacheStoreRecvingThread",
830 - side_effect=_auto_set_ready_event,
831 - ),
832 - ):
833 - w.register_kv_caches({"layer0": tensor})
834 -
835 - assert len(w.kv_caches_base_addr) == 1
836 - assert w.kv_caches_base_addr[0] == tensor.untyped_storage().data_ptr()
837 -
838 - expected_block_len = tensor.untyped_storage().nbytes() // num_blocks
839 - assert len(w.block_len) == 1
840 - assert w.block_len[0] == expected_block_len
841 -
842 - w.store.register_buffer.assert_called_once_with(
1096 + db = worker.token_dbs[0]
1097 + assert db.kv_caches_base_addr == [tensor.untyped_storage().data_ptr()]
1098 + assert db.block_len == [tensor.untyped_storage().nbytes() // num_blocks]
1099 + worker.store.register_buffer.assert_called_once_with(
1100 tensor.untyped_storage().data_ptr(),
1101 tensor.untyped_storage().nbytes(),
1102 )
@@ -851,8 +1108,7 @@ def test_register_kv_caches_kv_first_two_segments():
1108 block_size_tokens = 16
1109 num_kv_heads = 4
1110 head_size = 8
854 -
855 - w = _make_bare_worker(num_gpu_blocks=num_blocks)
1111 + worker = _make_bare_worker(num_gpu_blocks=num_blocks)
1112
1113 # Shape: (2, num_blocks, block_size, num_kv_heads, head_size) — K/V outermost
1114 tensor = torch.zeros(
@@ -863,32 +1119,13 @@ def test_register_kv_caches_kv_first_two_segments():
1119 head_size,
1120 dtype=torch.float16,
1121 )
1122 + _register_with_mocked_threads(worker, {"layer0": tensor})
1123
867 - with (
868 - patch(
869 - "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
870 - "worker.KVCacheStoreSendingThread",
871 - side_effect=_auto_set_ready_event,
872 - ),
873 - patch(
874 - "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
875 - "worker.KVCacheStoreRecvingThread",
876 - side_effect=_auto_set_ready_event,
877 - ),
878 - ):
879 - w.register_kv_caches({"layer0": tensor})
880 -
881 - # K/V-first: dim 0 has stride > page_size, so 2 segments
882 - assert len(w.kv_caches_base_addr) == 2
883 - assert len(w.block_len) == 2
884 -
885 - el = tensor.element_size()
886 - seg_stride = tensor.stride(0) * el # stride of the K/V dim in bytes
1124 + db = worker.token_dbs[0]
1125 + seg_stride = tensor.stride(0) * tensor.element_size()
1126 base = tensor.untyped_storage().data_ptr()
888 - assert w.kv_caches_base_addr[0] == base
889 - assert w.kv_caches_base_addr[1] == base + seg_stride
890 - assert w.block_len[0] == seg_stride // num_blocks
891 - assert w.block_len[1] == seg_stride // num_blocks
1127 + assert db.kv_caches_base_addr == [base, base + seg_stride]
1128 + assert db.block_len == [seg_stride // num_blocks] * 2
1129
1130
1131 def test_register_kv_caches_cross_layer_single_segment():
@@ -897,7 +1134,7 @@ def test_register_kv_caches_cross_layer_single_segment():
1134 num_layers = 4
1135 per_layer_page_elements = 64 # elements per layer per block
1136
900 - w = _make_bare_worker(num_gpu_blocks=num_blocks)
1137 + worker = _make_bare_worker(num_gpu_blocks=num_blocks)
1138
1139 # Cross-layer blocks-first tensor: all layers packed into a single
1140 # contiguous block. Shape (num_blocks, num_layers * per_layer_page)
@@ -918,10 +1155,11 @@ def test_register_kv_caches_cross_layer_single_segment():
1155 ),
1156 ):
1157 # Use the cross-layer wrapper key, same as register_cross_layers_kv_caches
921 - w.register_kv_caches({"__cross_layer__": tensor})
1158 + worker.register_kv_caches({"__cross_layer__": tensor})
1159
923 - assert len(w.kv_caches_base_addr) == 1
924 - assert w.kv_caches_base_addr[0] == tensor.untyped_storage().data_ptr()
1160 + db = worker.token_dbs[0]
1161 + assert len(db.kv_caches_base_addr) == 1
1162 + assert db.kv_caches_base_addr[0] == tensor.untyped_storage().data_ptr()
1163
1164 expected_block_len = tensor.untyped_storage().nbytes() // num_blocks
1165 # block_len should be per_layer_page_size * num_layers
@@ -929,11 +1167,11 @@ def test_register_kv_caches_cross_layer_single_segment():
1167 expected_block_len
1168 == num_layers * per_layer_page_elements * tensor.element_size()
1169 )
932 - assert len(w.block_len) == 1
933 - assert w.block_len[0] == expected_block_len
1170 + assert len(db.block_len) == 1
1171 + assert db.block_len[0] == expected_block_len
1172
1173 # Also verify via register_cross_layers_kv_caches wrapper
936 - w2 = _make_bare_worker(num_gpu_blocks=num_blocks)
1174 + worker2 = _make_bare_worker(num_gpu_blocks=num_blocks)
1175 with (
1176 patch(
1177 "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
@@ -946,10 +1184,11 @@ def test_register_kv_caches_cross_layer_single_segment():
1184 side_effect=_auto_set_ready_event,
1185 ),
1186 ):
949 - w2.register_cross_layers_kv_caches(tensor)
1187 + worker2.register_cross_layers_kv_caches(tensor)
1188
951 - assert w2.kv_caches_base_addr == w.kv_caches_base_addr
952 - assert w2.block_len == w.block_len
1189 + db2 = worker2.token_dbs[0]
1190 + assert db2.kv_caches_base_addr == db.kv_caches_base_addr
1191 + assert db2.block_len == db.block_len
1192
1193
1194 # ---------------------------------------------------------------------------
@@ -1070,7 +1309,8 @@ def test_topology_standalone_store_with_disk_offload(tmp_path, monkeypatch):
1309 )
1310
1311 w = worker.MooncakeStoreWorker(
1073 - _make_vllm_config(extra_config={"preferred_segment": "10.0.0.7:50053"})
1312 + _make_vllm_config(extra_config={"preferred_segment": "10.0.0.7:50053"}),
1313 + _make_kv_cache_config(),
1314 )
1315
1316 # setup() receives global_segment_size=0 and the configured local buffer.
@@ -1114,7 +1354,7 @@ def test_topology_embedded_cpu_only(tmp_path, monkeypatch):
1354 ),
1355 )
1356
1117 - w = worker.MooncakeStoreWorker(_make_vllm_config())
1357 + w = worker.MooncakeStoreWorker(_make_vllm_config(), _make_kv_cache_config())
1358
1359 # setup() receives global_segment_size=4 GiB (rank contributes a segment).
1360 assert store.setup.call_args.args == (
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py
+50 -3
@@ -26,6 +26,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
26 KVConnectorBase_V1,
27 KVConnectorMetadata,
28 KVConnectorRole,
29 + SupportsHMA,
30 )
31 from vllm.forward_context import ForwardContext
32 from vllm.logger import init_logger
@@ -76,7 +77,7 @@ class MooncakeStoreKVEvents(KVConnectorKVEvents):
77 return f"<MooncakeStoreKVEvents events={self.get_all_events()}>"
78
79
79 -class MooncakeStoreConnector(KVConnectorBase_V1):
80 +class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA):
81 """KV connector using MooncakeDistributedStore as shared KV pool."""
82
83 @property
@@ -87,6 +88,36 @@ class MooncakeStoreConnector(KVConnectorBase_V1):
88 == "true"
89 )
90
91 + @staticmethod
92 + def _validate_kv_cache_config(
93 + vllm_config: VllmConfig, kv_cache_config: KVCacheConfig
94 + ) -> None:
95 + from vllm.v1.kv_cache_interface import CrossAttentionSpec, MambaSpec
96 +
97 + unsupported: list[str] = []
98 + cache_block_size = vllm_config.cache_config.block_size
99 + for g_idx, g in enumerate(kv_cache_config.kv_cache_groups):
100 + spec = g.kv_cache_spec
101 + if isinstance(spec, CrossAttentionSpec):
102 + unsupported.append(f"group {g_idx}: CrossAttentionSpec")
103 + # Enforce Mamba align mode
104 + if isinstance(spec, MambaSpec) and spec.block_size != cache_block_size:
105 + unsupported.append(
106 + f"group {g_idx}: MambaSpec with block_size="
107 + f"{spec.block_size} != cache_config.block_size="
108 + f"{cache_block_size} (mamba_cache_mode != 'align')"
109 + )
110 + pcp = vllm_config.parallel_config.prefill_context_parallel_size
111 + dcp = vllm_config.parallel_config.decode_context_parallel_size
112 + if len(kv_cache_config.kv_cache_groups) > 1 and pcp * dcp > 1:
113 + unsupported.append(
114 + f"PCP/DCP > 1 (pcp={pcp}, dcp={dcp}) with hybrid attention"
115 + )
116 + if unsupported:
117 + raise ValueError(
118 + "MooncakeStoreConnector does not support: " + "; ".join(unsupported)
119 + )
120 +
121 def __init__(
122 self,
123 vllm_config: VllmConfig,
@@ -99,6 +130,9 @@ class MooncakeStoreConnector(KVConnectorBase_V1):
130 kv_cache_config=kv_cache_config, # type: ignore[arg-type]
131 )
132 assert vllm_config.kv_transfer_config is not None
133 + assert kv_cache_config is not None, "kv_cache_config is required"
134 + self._validate_kv_cache_config(vllm_config, kv_cache_config)
135 + self._kv_cache_config = kv_cache_config
136 self.kv_role = vllm_config.kv_transfer_config.kv_role
137 self._kv_cache_events: MooncakeStoreKVEvents | None = None
138
@@ -106,9 +140,11 @@ class MooncakeStoreConnector(KVConnectorBase_V1):
140 self.connector_worker: MooncakeStoreWorker | None = None
141
142 if role == KVConnectorRole.SCHEDULER:
109 - self.connector_scheduler = MooncakeStoreScheduler(vllm_config)
143 + self.connector_scheduler = MooncakeStoreScheduler(
144 + vllm_config, kv_cache_config
145 + )
146 else:
111 - self.connector_worker = MooncakeStoreWorker(vllm_config)
147 + self.connector_worker = MooncakeStoreWorker(vllm_config, kv_cache_config)
148
149 # ============================================================
150 # Scheduler-side methods
@@ -146,6 +182,13 @@ class MooncakeStoreConnector(KVConnectorBase_V1):
182 self,
183 request: Request,
184 block_ids: list[int],
185 + ) -> tuple[bool, dict[str, Any] | None]:
186 + return self.request_finished_all_groups(request, (block_ids,))
187 +
188 + def request_finished_all_groups(
189 + self,
190 + request: Request,
191 + block_ids: tuple[list[int], ...],
192 ) -> tuple[bool, dict[str, Any] | None]:
193 assert self.connector_scheduler is not None
194 return self.connector_scheduler.request_finished(request, block_ids)
@@ -184,6 +227,10 @@ class MooncakeStoreConnector(KVConnectorBase_V1):
227 self, kv_cache: torch.Tensor, attn_backend: type
228 ):
229 assert self.connector_worker is not None
230 + assert (
231 + self._kv_cache_config is not None
232 + and len(self._kv_cache_config.kv_cache_groups) == 1
233 + ), "Cross-layer KV cache does not supported with hybrid models"
234 self.connector_worker.register_cross_layers_kv_caches(kv_cache)
235
236 def start_load_kv(self, forward_context: ForwardContext, **kwargs: Any) -> None:
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py new
+290
@@ -0,0 +1,290 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""External-store cache-hit coordinator for MooncakeStoreConnector."""
4 +
5 +from typing import cast
6 +
7 +from vllm.v1.core.block_pool import BlockPool
8 +from vllm.v1.core.kv_cache_utils import (
9 + BlockHash,
10 + BlockHashList,
11 + BlockHashListWithBlockSize,
12 + KVCacheBlock,
13 +)
14 +from vllm.v1.core.single_type_kv_cache_manager import (
15 + SingleTypeKVCacheManager,
16 + spec_manager_map,
17 +)
18 +from vllm.v1.kv_cache_interface import (
19 + FullAttentionSpec,
20 + KVCacheGroupSpec,
21 + KVCacheSpec,
22 + UniformTypeKVCacheSpecs,
23 +)
24 +
25 +# Dummy placeholder hash for store_mask's template computation.
26 +_DUMMY_BLOCK_HASH = BlockHash(b"\x00" * 32)
27 +
28 +
29 +class ExternalCachedBlockPool:
30 + """Duck-typed BlockPool backed by a ``(group_id, hash)`` exists set."""
31 +
32 + def __init__(self, exists: set[tuple[int, bytes]] | None = None) -> None:
33 + # ``exists=None`` is used on the recv side where hit_length is already
34 + # determined and we just want each spec's manager to apply its own mask.
35 + self._exists = exists
36 + self.null_block = KVCacheBlock(block_id=0)
37 + # Dummy ID 1 for present block for duck-typing.
38 + self._present_block = KVCacheBlock(block_id=1)
39 +
40 + def get_cached_block(
41 + self,
42 + block_hash: BlockHash,
43 + group_ids: list[int],
44 + ) -> list[KVCacheBlock] | None:
45 + # Mirrors BlockPool.get_cached_block: hit only when every group_id
46 + # (groups sharing a spec) has the hash cached.
47 + if self._exists is None:
48 + return [self._present_block] * len(group_ids)
49 + h = bytes(block_hash)
50 + if all((g, h) in self._exists for g in group_ids):
51 + return [self._present_block] * len(group_ids)
52 + return None
53 +
54 +
55 +class MooncakeStoreCoordinator:
56 + """Mirror of ``HybridKVCacheCoordinator.find_longest_cache_hit`` over an
57 + ``ExternalCachedBlockPool``."""
58 +
59 + def __init__(
60 + self,
61 + kv_cache_groups: list[KVCacheGroupSpec],
62 + scheduler_block_size: int,
63 + hash_block_size: int,
64 + use_eagle: bool = False,
65 + ) -> None:
66 + assert all(
67 + g.kv_cache_spec.block_size % hash_block_size == 0 for g in kv_cache_groups
68 + ), "block_size must be divisible by hash_block_size"
69 + assert scheduler_block_size % hash_block_size == 0, (
70 + f"scheduler_block_size ({scheduler_block_size}) must be a multiple of "
71 + f"hash_block_size ({hash_block_size})"
72 + )
73 + assert all(
74 + scheduler_block_size % g.kv_cache_spec.block_size == 0
75 + for g in kv_cache_groups
76 + ), "scheduler_block_size must be a multiple of each group's block_size"
77 + self.kv_cache_groups = kv_cache_groups
78 + self.hash_block_size = hash_block_size
79 + self.lcm_block_size = scheduler_block_size
80 + self.use_eagle = use_eagle
81 + self._verify_and_split_kv_cache_groups()
82 +
83 + def _verify_and_split_kv_cache_groups(self) -> None:
84 + """Mirrors KVCacheCoordinator.verify_and_split_kv_cache_groups but
85 + dispatches via spec_manager_map (we don't allocate managers).
86 + """
87 + attention_groups: list[
88 + tuple[KVCacheSpec, list[int], type[SingleTypeKVCacheManager]]
89 + ] = []
90 + for i, g in enumerate(self.kv_cache_groups):
91 + spec = _unwrap_spec(g.kv_cache_spec)
92 + manager_cls = spec_manager_map[type(spec)]
93 + for existing_spec, group_ids, existing_cls in attention_groups:
94 + if existing_spec == spec:
95 + assert manager_cls is existing_cls
96 + group_ids.append(i)
97 + break
98 + else:
99 + attention_groups.append((spec, [i], manager_cls))
100 + # Full attention first (matches upstream convergence ordering).
101 + self.attention_groups = sorted(
102 + attention_groups,
103 + key=lambda x: not isinstance(x[0], FullAttentionSpec),
104 + )
105 + self.eagle_attn_group_indices: set[int] = {
106 + i
107 + for i, (_, group_ids, _) in enumerate(self.attention_groups)
108 + if any(self.kv_cache_groups[gid].is_eagle_group for gid in group_ids)
109 + }
110 + if self.use_eagle and not self.eagle_attn_group_indices:
111 + self.eagle_attn_group_indices = set(range(len(self.attention_groups)))
112 +
113 + def find_longest_cache_hit(
114 + self,
115 + block_hashes: list[BlockHash],
116 + max_length: int,
117 + cached_block_pool: ExternalCachedBlockPool,
118 + ) -> tuple[tuple[list[bool], ...], int]:
119 + """Returns ``(load_mask_per_group, hit_length)``. ``mask[g][i]`` is True iff
120 + group ``g`` populates chunk ``i`` locally (e.g. SWA and Mamba tail-only);
121 + recv-side callers skip False slots."""
122 + blocks_per_group, hit_length = self._find_hit_blocks(
123 + block_hashes, max_length, cached_block_pool
124 + )
125 + masks = tuple(
126 + [blk is not cached_block_pool.null_block for blk in blocks]
127 + for blocks in blocks_per_group
128 + )
129 + return masks, hit_length
130 +
131 + def load_mask(
132 + self,
133 + block_hashes: list[BlockHash],
134 + token_len: int,
135 + ) -> tuple[list[bool], ...]:
136 + """Per-group load masks: ``mask[g][i]`` is True iff group ``g``'s
137 + spec would populate chunk ``i`` locally at length ``token_len``
138 + (e.g. SWA / Mamba tail-only).
139 + """
140 + masks, _ = self.find_longest_cache_hit(
141 + block_hashes, token_len, ExternalCachedBlockPool()
142 + )
143 + return masks
144 +
145 + def store_mask(self, aligned_token_len: int) -> tuple[list[bool], ...]:
146 + """Per-group store masks: ``mask[g][i]`` is True iff chunk ``i`` of
147 + group ``g`` would be populated by some future cache hit at length
148 + ``L = N * lcm_block_size <= aligned_token_len``.
149 + """
150 + assert aligned_token_len % self.lcm_block_size == 0, (
151 + f"aligned_token_len ({aligned_token_len}) must be a multiple of "
152 + f"lcm_block_size ({self.lcm_block_size})"
153 + )
154 + if aligned_token_len == 0:
155 + return tuple([] for _ in self.kv_cache_groups)
156 +
157 + num_chunks_per_group = [
158 + aligned_token_len // g.kv_cache_spec.block_size
159 + for g in self.kv_cache_groups
160 + ]
161 +
162 + # Fast path: single group or full attn groups or uniform block_sizes
163 + if all(
164 + isinstance(spec, FullAttentionSpec)
165 + or spec.block_size == self.lcm_block_size
166 + for spec, _, _ in self.attention_groups
167 + ):
168 + return tuple([True] * n for n in num_chunks_per_group)
169 +
170 + n_segments = aligned_token_len // self.lcm_block_size
171 + dummy_hashes: list[BlockHash] = [_DUMMY_BLOCK_HASH] * (
172 + self.lcm_block_size // self.hash_block_size
173 + )
174 + template_masks, _ = self.find_longest_cache_hit(
175 + dummy_hashes,
176 + max_length=self.lcm_block_size,
177 + cached_block_pool=ExternalCachedBlockPool(),
178 + )
179 + return tuple(
180 + list(template_masks[g]) * n_segments
181 + for g in range(len(self.kv_cache_groups))
182 + )
183 +
184 + def block_hashes_for_spec(
185 + self, block_hashes: list[BlockHash], spec: KVCacheSpec
186 + ) -> BlockHashList:
187 + if spec.block_size == self.hash_block_size:
188 + return block_hashes
189 + return BlockHashListWithBlockSize(
190 + block_hashes, self.hash_block_size, spec.block_size
191 + )
192 +
193 + def _find_hit_blocks(
194 + self,
195 + block_hashes: list[BlockHash],
196 + max_length: int,
197 + cached_block_pool: ExternalCachedBlockPool,
198 + ) -> tuple[tuple[list[KVCacheBlock], ...], int]:
199 + """Mirrors HybridKVCacheCoordinator.find_longest_cache_hit but
200 + dispatches via spec_manager_map (we don't allocate managers).
201 + """
202 + if len(self.attention_groups) == 1:
203 + spec, group_ids, manager_cls = self.attention_groups[0]
204 + hashes = self.block_hashes_for_spec(block_hashes, spec)
205 + hit_blocks = manager_cls.find_longest_cache_hit(
206 + block_hashes=hashes,
207 + max_length=max_length,
208 + kv_cache_group_ids=group_ids,
209 + block_pool=cast(BlockPool, cached_block_pool),
210 + kv_cache_spec=spec,
211 + use_eagle=(0 in self.eagle_attn_group_indices),
212 + alignment_tokens=spec.block_size,
213 + )
214 + num_groups = len(self.kv_cache_groups)
215 + blocks_by_group: list[list[KVCacheBlock]] = [[] for _ in range(num_groups)]
216 + for gid, blks in zip(group_ids, hit_blocks, strict=True):
217 + blocks_by_group[gid] = blks
218 + return tuple(blocks_by_group), len(hit_blocks[0]) * spec.block_size
219 +
220 + num_groups = len(self.kv_cache_groups)
221 + hit_length = max_length
222 + hit_blocks_by_group: list[list[KVCacheBlock] | None] = [None] * num_groups
223 +
224 + is_simple_hybrid = len(self.attention_groups) == 2 and isinstance(
225 + self.attention_groups[0][0], FullAttentionSpec
226 + )
227 + eagle_verified: set[int] = set()
228 +
229 + while True:
230 + curr_hit_length = hit_length
231 +
232 + for idx, (spec, group_ids, manager_cls) in enumerate(self.attention_groups):
233 + cached = hit_blocks_by_group[group_ids[0]]
234 + if isinstance(spec, FullAttentionSpec) and cached is not None:
235 + curr_hit_length = (
236 + curr_hit_length // spec.block_size * spec.block_size
237 + )
238 + continue
239 +
240 + use_eagle = (
241 + idx in self.eagle_attn_group_indices and idx not in eagle_verified
242 + )
243 + _max_length = curr_hit_length
244 + if use_eagle:
245 + _max_length = min(curr_hit_length + spec.block_size, max_length)
246 + hashes = self.block_hashes_for_spec(block_hashes, spec)
247 + hit_blocks = manager_cls.find_longest_cache_hit(
248 + block_hashes=hashes,
249 + max_length=_max_length,
250 + kv_cache_group_ids=group_ids,
251 + block_pool=cast(BlockPool, cached_block_pool),
252 + kv_cache_spec=spec,
253 + use_eagle=use_eagle,
254 + alignment_tokens=self.lcm_block_size,
255 + )
256 + _new_hit_length = len(hit_blocks[0]) * spec.block_size
257 + if use_eagle:
258 + eagle_verified.add(idx)
259 + elif _new_hit_length < curr_hit_length:
260 + eagle_verified.clear()
261 + curr_hit_length = _new_hit_length
262 + for gid, blocks in zip(group_ids, hit_blocks, strict=True):
263 + hit_blocks_by_group[gid] = blocks
264 +
265 + if curr_hit_length >= hit_length:
266 + break
267 + hit_length = curr_hit_length
268 + if is_simple_hybrid:
269 + break
270 +
271 + # Truncate full-attention hit_blocks to final converged length;
272 + # other specs already trim themselves inside their hit logic.
273 + spec0, group_ids0, _ = self.attention_groups[0]
274 + if isinstance(spec0, FullAttentionSpec):
275 + num_blocks = hit_length // spec0.block_size
276 + for gid in group_ids0:
277 + full_blks = hit_blocks_by_group[gid]
278 + assert full_blks is not None
279 + del full_blks[num_blocks:]
280 +
281 + return (
282 + tuple(blks if blks is not None else [] for blks in hit_blocks_by_group),
283 + hit_length,
284 + )
285 +
286 +
287 +def _unwrap_spec(spec: KVCacheSpec) -> KVCacheSpec:
288 + if isinstance(spec, UniformTypeKVCacheSpecs):
289 + return next(iter(spec.kv_cache_specs.values()))
290 + return spec
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py
+47 -33
@@ -15,7 +15,10 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
15 )
16 from vllm.logger import init_logger
17 from vllm.utils.math_utils import cdiv
18 -from vllm.v1.core.kv_cache_utils import BlockHash
18 +from vllm.v1.core.kv_cache_utils import (
19 + BlockHash,
20 + BlockHashListWithBlockSize,
21 +)
22
23 logger = init_logger(__name__)
24
@@ -29,6 +32,7 @@ class KeyMetadata:
32 pcp_rank: int
33 dcp_rank: int
34 pp_rank: int
35 + group_id: int = 0
36
37
38 @dataclass(order=True)
@@ -46,6 +50,7 @@ class PoolKey:
50 self.key_metadata.pcp_rank,
51 self.key_metadata.dcp_rank,
52 self.key_metadata.pp_rank,
53 + self.key_metadata.group_id,
54 self.chunk_hash,
55 )
56 )
@@ -57,6 +62,7 @@ class PoolKey:
62 f"@pcp{self.key_metadata.pcp_rank}"
63 f"@dcp{self.key_metadata.dcp_rank}"
64 f"@pp_rank:{self.key_metadata.pp_rank}"
65 + f"@group:{self.key_metadata.group_id}"
66 f"@{self.chunk_hash}"
67 )
68
@@ -64,9 +70,20 @@ class PoolKey:
70 class ChunkedTokenDatabase:
71 """Maps token positions to store keys and GPU memory addresses."""
72
67 - def __init__(self, metadata: KeyMetadata, block_size: int):
73 + def __init__(
74 + self,
75 + metadata: KeyMetadata,
76 + block_size: int,
77 + hash_block_size: int | None = None,
78 + ):
79 self.metadata = metadata
80 self.block_size = block_size
81 + self.hash_block_size = hash_block_size or block_size
82 + if self.block_size % self.hash_block_size != 0:
83 + raise ValueError(
84 + f"block_size ({self.block_size}) must be a multiple of "
85 + f"hash_block_size ({self.hash_block_size})"
86 + )
87 self.kv_caches_base_addr: list[int] = []
88 self.block_len: list[int] = []
89
@@ -77,9 +94,6 @@ class ChunkedTokenDatabase:
94 self.kv_caches_base_addr = kv_caches_base_addr
95
96 def set_block_len(self, block_len: list[int]):
80 - for length in block_len:
81 - if length % self.block_size != 0:
82 - raise ValueError(f"block_len {length} % {self.block_size} != 0")
97 self.block_len = block_len
98
99 def prepare_value(
@@ -96,7 +110,8 @@ class ChunkedTokenDatabase:
110 length = len(self.block_len)
111 for index, base_addr in enumerate(self.kv_caches_base_addr):
112 addr = base_addr + block_id * self.block_len[index % length]
99 - size = self.block_len[index % length] // self.block_size * (end - start)
113 + assert (end - start) % self.block_size == 0
114 + size = self.block_len[index % length] * cdiv(end - start, self.block_size)
115 addr_list.append(addr)
116 size_list.append(size)
117 return addr_list, size_list, block_id
@@ -104,38 +119,34 @@ class ChunkedTokenDatabase:
119 def process_tokens(
120 self,
121 token_len: int,
107 - block_hashes: list[BlockHash] | list[str],
122 + block_hashes: list[BlockHash],
123 mask_num: int = 0,
124 ) -> Iterable[tuple[int, int, PoolKey]]:
125 """Process tokens and yield (start_idx, end_idx, pool_key) tuples.
126
127 Args:
128 token_len: Total number of tokens.
114 - block_hashes: Block hashes for each block.
129 + block_hashes: Block hashes computed at ``hash_block_size`` granularity.
130 + When ``block_size > hash_block_size`` consecutive hashes are merged
131 + up to the group's ``block_size`` via ``BlockHashListWithBlockSize``.
132 mask_num: Number of tokens to skip from the beginning.
133 """
134 if not block_hashes:
135 return
119 - if not isinstance(block_hashes[0], str):
120 - block_hashes = [
121 - h.hex() # type: ignore[union-attr]
122 - for h in block_hashes
123 - ]
124 - for chunk_id, hash_val in enumerate(block_hashes):
136 + if self.block_size == self.hash_block_size:
137 + chunk_hashes: Iterable[BlockHash] = block_hashes
138 + else:
139 + chunk_hashes = BlockHashListWithBlockSize(
140 + block_hashes, self.hash_block_size, self.block_size
141 + )
142 + for chunk_id, h in enumerate(chunk_hashes):
143 start_idx = chunk_id * self.block_size
144 if start_idx >= token_len:
145 break
146 end_idx = min(start_idx + self.block_size, token_len)
147 if start_idx < mask_num:
148 continue
131 - else:
132 - yield (
133 - start_idx,
134 - end_idx,
135 - self._make_key_by_hash(
136 - hash_val # type: ignore[arg-type]
137 - ),
138 - )
149 + yield start_idx, end_idx, self._make_key_by_hash(h.hex())
150
151
152 @dataclass
@@ -154,7 +165,7 @@ class RequestTracker:
165
166 req_id: str
167 token_len: int
157 - allocated_block_ids: list[int]
168 + allocated_block_ids: tuple[list[int], ...]
169 num_saved_tokens: int = 0
170 token_ids: list[int] | None = None
171 # Snapshot of the prefill range length at tracker creation time.
@@ -166,15 +177,18 @@ class RequestTracker:
177 self,
178 new_block_ids: tuple[list[int], ...] | list[int],
179 ) -> None:
169 - if len(new_block_ids) == 0:
170 - new_block_ids = []
171 - elif isinstance(new_block_ids, tuple):
172 - new_block_ids = new_block_ids[0]
173 - elif isinstance(new_block_ids, list):
174 - pass
175 - else:
176 - raise ValueError(f"Unsupported new_block_ids type {type(new_block_ids)}")
177 - self.allocated_block_ids.extend(new_block_ids)
180 + # Backward-compat: accept a single list (broadcast to single group).
181 + if isinstance(new_block_ids, list):
182 + new_block_ids = (new_block_ids,)
183 + if len(new_block_ids) != len(self.allocated_block_ids):
184 + raise ValueError(
185 + f"Group count mismatch: tracker has "
186 + f"{len(self.allocated_block_ids)} groups, update has "
187 + f"{len(new_block_ids)}"
188 + )
189 + for existing, new in zip(self.allocated_block_ids, new_block_ids, strict=True):
190 + if new:
191 + existing.extend(new)
192
193
194 @dataclass
@@ -183,7 +197,7 @@ class ReqMeta:
197
198 req_id: str
199 token_len_chunk: int
186 - block_ids: list[int]
200 + block_ids: tuple[list[int], ...]
201 block_hashes: list[BlockHash]
202
203 can_save: bool | None = None
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py
+27 -19
@@ -22,7 +22,9 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.worker import (
22 )
23 from vllm.logger import init_logger
24 from vllm.v1.core.kv_cache_manager import KVCacheBlocks
25 +from vllm.v1.core.kv_cache_utils import resolve_kv_cache_block_sizes
26 from vllm.v1.core.sched.output import NewRequestData, SchedulerOutput
27 +from vllm.v1.kv_cache_interface import KVCacheConfig
28 from vllm.v1.request import Request
29
30 logger = init_logger(__name__)
@@ -45,7 +47,11 @@ def _new_req_prefill_tokens(request: NewRequestData) -> list[int]:
47 class MooncakeStoreScheduler:
48 """Scheduler-side component for MooncakeStoreConnector."""
49
48 - def __init__(self, vllm_config: VllmConfig):
50 + def __init__(
51 + self,
52 + vllm_config: VllmConfig,
53 + kv_cache_config: KVCacheConfig,
54 + ):
55 assert vllm_config.kv_transfer_config is not None
56 self.kv_role = vllm_config.kv_transfer_config.kv_role
57 self.load_async = vllm_config.kv_transfer_config.kv_connector_extra_config.get(
@@ -56,11 +62,11 @@ class MooncakeStoreScheduler:
62 self.pcp_size = vllm_config.parallel_config.prefill_context_parallel_size
63 self.dcp_size = vllm_config.parallel_config.decode_context_parallel_size
64 self.original_block_size = vllm_config.cache_config.block_size
59 - self._block_size = vllm_config.cache_config.block_size
60 - if self.pcp_size > 1:
61 - self._block_size *= self.pcp_size
62 - if self.dcp_size > 1:
63 - self._block_size *= self.dcp_size
65 + # LCM for multi-group HMA; bs * pcp * dcp for single-group. Matches
66 + # the engine's own scheduler block size by construction.
67 + self._block_size, self._hash_block_size = resolve_kv_cache_block_sizes(
68 + kv_cache_config, vllm_config
69 + )
70
71 self._discard_partial_chunks = (
72 vllm_config.kv_transfer_config.get_from_extra_config(
@@ -72,7 +78,7 @@ class MooncakeStoreScheduler:
78 self.load_specs: dict[str, LoadSpec] = {} # to be loaded
79 self._request_trackers: dict[str, RequestTracker] = {} # scheduled new requests
80 self._preempted_req_ids: set[str] = set() # preempted requests
75 - self._unfinished_requests: dict[str, tuple[Request, list[int]]] = {}
81 + self._unfinished_requests: dict[str, tuple[Request, tuple[list[int], ...]]] = {}
82 self._unfinished_request_ids: set[str] = set()
83
84 def get_num_new_matched_tokens(
@@ -126,9 +132,9 @@ class MooncakeStoreScheduler:
132 num_external_tokens: int,
133 ):
134 """Update state after block allocation."""
129 - local_block_ids: list[int] = []
135 + local_block_ids: tuple[list[int], ...] = ()
136 if num_external_tokens > 0:
131 - local_block_ids = blocks.get_block_ids()[0]
137 + local_block_ids = blocks.get_block_ids()
138
139 self._unfinished_requests[request.request_id] = (request, local_block_ids)
140 self._unfinished_request_ids.add(request.request_id)
@@ -189,11 +195,12 @@ class MooncakeStoreScheduler:
195 request_tuple = self._unfinished_requests.get(request.req_id)
196 request_real = request_tuple[0] # type: ignore[index]
197
192 - if not isinstance(request.block_ids[0], list):
193 - unfolded_block_ids = request.block_ids.copy()
198 + if isinstance(request.block_ids, tuple):
199 + # Multi-group: preserve per-group structure.
200 + unfolded_block_ids = tuple(b.copy() for b in request.block_ids)
201 else:
195 - # TODO: support HMA
196 - unfolded_block_ids = request.block_ids[0].copy()
202 + # Single-group legacy: list[int] -> 1-tuple.
203 + unfolded_block_ids = (request.block_ids.copy(),)
204
205 prefill_tokens = _new_req_prefill_tokens(request)
206 request_tracker = RequestTracker(
@@ -237,9 +244,9 @@ class MooncakeStoreScheduler:
244 if req_id in self._preempted_req_ids:
245 # Resumed after preemption
246 if isinstance(new_block_ids, tuple):
240 - block_ids_list = new_block_ids[0].copy()
247 + new_block_ids = tuple(b.copy() for b in new_block_ids)
248 else:
242 - block_ids_list = new_block_ids.copy()
249 + new_block_ids = (new_block_ids.copy(),)
250 self._preempted_req_ids.discard(req_id)
251 load_spec = self.load_specs.pop(req_id, None)
252 request_tuple = self._unfinished_requests.get(req_id)
@@ -254,7 +261,7 @@ class MooncakeStoreScheduler:
261 request_tracker = RequestTracker(
262 req_id=req_id,
263 token_len=num_tokens_to_compute,
257 - allocated_block_ids=block_ids_list,
264 + allocated_block_ids=new_block_ids,
265 num_saved_tokens=0,
266 token_ids=prefill_tokens[:num_tokens_to_compute].copy(),
267 prefill_end_tokens=len(prefill_tokens),
@@ -361,7 +368,7 @@ class MooncakeStoreScheduler:
368 def request_finished(
369 self,
370 request: Request,
364 - block_ids: list[int],
371 + block_ids: tuple[list[int], ...],
372 ) -> tuple[bool, dict[str, Any] | None]:
373 """Determine whether to delay freeing blocks for async save."""
374 if self.kv_role == "kv_consumer":
@@ -370,11 +377,12 @@ class MooncakeStoreScheduler:
377 assert tracker is not None
378 if tracker.num_saved_tokens <= 0:
379 return False, None
373 - delay_free_blocks = len(block_ids) > 0
380 + total_blocks = sum(len(g) for g in block_ids)
381 + delay_free_blocks = total_blocks > 0
382 if delay_free_blocks:
383 logger.debug(
384 "Delaying free of %d blocks for request %s",
377 - len(block_ids),
385 + total_blocks,
386 request.request_id,
387 )
388 return delay_free_blocks, None
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py
+237 -180
@@ -10,6 +10,7 @@ Includes the store worker, transfer threads, lookup server,
10 and MooncakeDistributedStore integration.
11 """
12
13 +import dataclasses
14 import json
15 import os
16 import queue
@@ -36,15 +37,25 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake import rdma_utils
37 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_utils import (
38 get_mooncake_dp_engine_index,
39 )
39 -from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
40 +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.coordinator import ( # noqa: E501
41 + ExternalCachedBlockPool,
42 + MooncakeStoreCoordinator,
43 +)
44 +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( # noqa: E501
45 ChunkedTokenDatabase,
46 KeyMetadata,
47 MooncakeStoreConnectorMetadata,
48 + PoolKey,
49 ReqMeta,
50 )
51 from vllm.logger import init_logger
52 from vllm.utils.network_utils import get_ip, make_zmq_socket
47 -from vllm.v1.core.kv_cache_utils import BlockHash, maybe_convert_block_hash
53 +from vllm.v1.core.kv_cache_utils import (
54 + BlockHash,
55 + maybe_convert_block_hash,
56 + resolve_kv_cache_block_sizes,
57 +)
58 +from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheGroupSpec
59 from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder
60
61 logger = init_logger(__name__)
@@ -321,7 +332,7 @@ class KVTransferThread(threading.Thread):
332 def __init__(
333 self,
334 store: Any,
324 - token_database: ChunkedTokenDatabase,
335 + token_databases: list[ChunkedTokenDatabase],
336 block_size: int,
337 tp_rank: int,
338 ready_event: threading.Event,
@@ -332,7 +343,7 @@ class KVTransferThread(threading.Thread):
343 self.ready_event = ready_event
344 self.block_size = block_size
345 self.tp_rank = tp_rank
335 - self.token_database = token_database
346 + self.token_databases = token_databases
347 self.done_task_lock = threading.Lock()
348 self.request_queue: queue.Queue[Any] = queue.Queue()
349 self.finished_requests: set[str] = set()
@@ -385,7 +396,8 @@ class KVCacheStoreSendingThread(KVTransferThread):
396 def __init__(
397 self,
398 store: Any,
388 - token_database: ChunkedTokenDatabase,
399 + coord: MooncakeStoreCoordinator,
400 + token_databases: list[ChunkedTokenDatabase],
401 block_size: int,
402 tp_rank: int,
403 put_step: int,
@@ -396,13 +408,14 @@ class KVCacheStoreSendingThread(KVTransferThread):
408 ):
409 super().__init__(
410 store,
399 - token_database,
411 + token_databases,
412 block_size,
413 tp_rank,
414 ready_event,
415 name="KVCacheStoreSendingThread",
416 )
417 self.put_step = put_step
418 + self.coord = coord
419 self.kv_role = kv_role
420 self.stored_requests: defaultdict[str, int] = defaultdict(int)
421 self.enable_kv_event = enable_kv_event
@@ -449,14 +462,20 @@ class KVCacheStoreSendingThread(KVTransferThread):
462 return True
463
464 def _handle_request(self, req_meta: ReqMeta):
452 - token_len = req_meta.token_len_chunk
453 - block_ids = req_meta.block_ids
465 + # Cache hits are always a multiple of ``lcm_block_size`` tokens
466 + lcm_block_size = self.coord.lcm_block_size
467 + token_len = req_meta.token_len_chunk // lcm_block_size * lcm_block_size
468 + block_ids_per_group = req_meta.block_ids
469 req_id = req_meta.req_id
470 current_event = req_meta.current_event
471
472 if req_id not in self.stored_requests:
473 self.request_queue.task_done()
474 return
475 + if token_len == 0:
476 + self.dec_stored_request(req_id)
477 + self.request_queue.task_done()
478 + return
479 if self._should_skip_request(req_id):
480 logger.debug(
481 "Skipping Mooncake store for request %s while CPU/disk offloading "
@@ -467,23 +486,34 @@ class KVCacheStoreSendingThread(KVTransferThread):
486 self.request_queue.task_done()
487 return
488
470 - starts = []
471 - ends = []
472 - keys = []
489 + # Within each lcm region only per-spec relevant chunks are loaded
490 + # (e.g., SWA or linear attn), so mask out irrelevant chunks
491 + store_masks = self.coord.store_mask(token_len)
492 + starts: list[int] = []
493 + ends: list[int] = []
494 + keys: list[str] = []
495 block_hashes: list[BlockHash] = []
474 - for index, (start, end, key) in enumerate(
475 - self.token_database.process_tokens(token_len, req_meta.block_hashes)
476 - ):
477 - starts.append(start)
478 - ends.append(end)
479 - keys.append(key.to_string())
480 - block_hashes.append(req_meta.block_hashes[index])
496 + group_indices: list[int] = []
497 + for g_idx, db in enumerate(self.token_databases):
498 + mask = store_masks[g_idx]
499 + for chunk_idx, (start, end, key) in enumerate(
500 + db.process_tokens(token_len, req_meta.block_hashes)
501 + ):
502 + if chunk_idx >= len(mask) or not mask[chunk_idx]:
503 + continue
504 + starts.append(start)
505 + ends.append(end)
506 + keys.append(key.to_string())
507 + block_hashes.append(req_meta.block_hashes[chunk_idx])
508 + group_indices.append(g_idx)
509
510 # Apply put_step striding for TP
483 - starts = starts[self.tp_rank % self.put_step :: self.put_step]
484 - ends = ends[self.tp_rank % self.put_step :: self.put_step]
485 - keys = keys[self.tp_rank % self.put_step :: self.put_step]
486 - block_hashes = block_hashes[self.tp_rank % self.put_step :: self.put_step]
511 + sl = slice(self.tp_rank % self.put_step, None, self.put_step)
512 + starts = starts[sl]
513 + ends = ends[sl]
514 + keys = keys[sl]
515 + block_hashes = block_hashes[sl]
516 + group_indices = group_indices[sl]
517
518 if not keys:
519 self.dec_stored_request(req_id)
@@ -501,38 +531,37 @@ class KVCacheStoreSendingThread(KVTransferThread):
531 ends = [ends[i] for i in missing_indices]
532 keys = [keys[i] for i in missing_indices]
533 block_hashes = [block_hashes[i] for i in missing_indices]
534 + group_indices = [group_indices[i] for i in missing_indices]
535
536 logger.debug(
506 - "Storing KV cache for %d out of %d blocks "
507 - "(missing_count=%d) for request %s",
537 + "Storing KV cache for %d blocks (groups=%s) for request %s",
538 len(keys),
509 - token_len // self.block_size,
510 - len(missing_indices),
539 + set(group_indices),
540 req_id,
541 )
542
514 - addrs = []
515 - sizes = []
543 + addrs: list[list[int]] = []
544 + sizes: list[list[int]] = []
545 stored_events: list[BlockStored] = []
517 - prev_key = None
546 + # parent_block_hash chains live within a group, not across.
547 + prev_key_per_group: dict[int, Any] = {}
548 new_block_hashes = [maybe_convert_block_hash(bh) for bh in block_hashes]
549
520 - for index, start in enumerate(starts):
521 - addr, size, _ = self.token_database.prepare_value(
522 - start, ends[index], block_ids
523 - )
550 + for idx, (s, e, g_idx) in enumerate(
551 + zip(starts, ends, group_indices, strict=True)
552 + ):
553 + db = self.token_databases[g_idx]
554 + addr, size, _ = db.prepare_value(s, e, block_ids_per_group[g_idx])
555 addrs.append(addr)
556 sizes.append(size)
557
558 if self.enable_kv_event:
559 token_ids = (
529 - req_meta.token_ids[start : ends[index]]
530 - if req_meta.token_ids is not None
531 - else None
560 + req_meta.token_ids[s:e] if req_meta.token_ids is not None else None
561 )
562 stored_event = BlockStored(
534 - block_hashes=[new_block_hashes[index]],
535 - parent_block_hash=prev_key,
563 + block_hashes=[new_block_hashes[idx]],
564 + parent_block_hash=prev_key_per_group.get(g_idx),
565 token_ids=token_ids,
566 block_size=req_meta.original_block_size,
567 lora_id=None,
@@ -540,7 +569,7 @@ class KVCacheStoreSendingThread(KVTransferThread):
569 lora_name=None,
570 )
571 stored_events.append(stored_event)
543 - prev_key = new_block_hashes[index]
572 + prev_key_per_group[g_idx] = new_block_hashes[idx]
573
574 if current_event is not None:
575 current_event.synchronize()
@@ -600,7 +629,8 @@ class KVCacheStoreRecvingThread(KVTransferThread):
629 def __init__(
630 self,
631 store: Any,
603 - token_database: ChunkedTokenDatabase,
632 + coord: MooncakeStoreCoordinator,
633 + token_databases: list[ChunkedTokenDatabase],
634 block_size: int,
635 tp_rank: int,
636 ready_event: threading.Event,
@@ -608,7 +638,7 @@ class KVCacheStoreRecvingThread(KVTransferThread):
638 ):
639 super().__init__(
640 store,
611 - token_database,
641 + token_databases,
642 block_size,
643 tp_rank,
644 ready_event,
@@ -622,6 +652,7 @@ class KVCacheStoreRecvingThread(KVTransferThread):
652 disk_offload_buffer_budget_bytes
653 )
654 )
655 + self.coord = coord
656
657 def _handle_request(self, req_meta: ReqMeta):
658 token_len = req_meta.load_spec.token_len # type: ignore[union-attr]
@@ -632,32 +663,31 @@ class KVCacheStoreRecvingThread(KVTransferThread):
663 * self.block_size
664 )
665
635 - addr_list = []
636 - size_list = []
637 - key_list = []
638 - for start, end, key in self.token_database.process_tokens(
639 - token_len, req_meta.block_hashes, mask_num
640 - ):
641 - addr, size, _ = self.token_database.prepare_value(
642 - start, end, req_meta.block_ids
643 - )
644 - key_list.append(key.to_string())
645 - addr_list.append(addr)
646 - size_list.append(size)
666 + # Skip chunks the consumer's per-group spec wouldn't populate
667 + # locally (e.g. SWA pre-window) even if the producer stored them.
668 + load_mask_per_group = self.coord.load_mask(req_meta.block_hashes, token_len)
669 +
670 + addr_list: list[list[int]] = []
671 + size_list: list[list[int]] = []
672 + key_list: list[str] = []
673 + for g_idx, db in enumerate(self.token_databases):
674 + mask = load_mask_per_group[g_idx]
675 + for start, end, key in db.process_tokens(
676 + token_len, req_meta.block_hashes, mask_num
677 + ):
678 + chunk_idx = start // db.block_size
679 + if chunk_idx >= len(mask) or not mask[chunk_idx]:
680 + continue
681 + addr, size, _ = db.prepare_value(start, end, req_meta.block_ids[g_idx])
682 + key_list.append(key.to_string())
683 + addr_list.append(addr)
684 + size_list.append(size)
685
686 # Rotate lists by tp_rank for load balancing
649 - key_list_c = (
650 - key_list[self.tp_rank % len(key_list) :]
651 - + key_list[: self.tp_rank % len(key_list)]
652 - )
653 - addr_list_c = (
654 - addr_list[self.tp_rank % len(addr_list) :]
655 - + addr_list[: self.tp_rank % len(addr_list)]
656 - )
657 - size_list_c = (
658 - size_list[self.tp_rank % len(size_list) :]
659 - + size_list[: self.tp_rank % len(size_list)]
660 - )
687 + rotation = self.tp_rank % len(key_list)
688 + key_list_c = key_list[rotation:] + key_list[:rotation]
689 + addr_list_c = addr_list[rotation:] + addr_list[:rotation]
690 + size_list_c = size_list[rotation:] + size_list[:rotation]
691
692 load_batches = [(key_list_c, addr_list_c, size_list_c)]
693 if self.usable_disk_offload_buffer_budget_bytes is not None:
@@ -737,7 +767,11 @@ class KVCacheStoreRecvingThread(KVTransferThread):
767 class MooncakeStoreWorker:
768 """Worker-side component for MooncakeStoreConnector."""
769
740 - def __init__(self, vllm_config: VllmConfig):
770 + def __init__(
771 + self,
772 + vllm_config: VllmConfig,
773 + kv_cache_config: KVCacheConfig,
774 + ):
775 try:
776 from mooncake.store import ( # type: ignore
777 MooncakeDistributedStore,
@@ -771,11 +805,9 @@ class MooncakeStoreWorker:
805 )
806 self.cache_config = vllm_config.cache_config
807 self.original_block_size = self.cache_config.block_size
774 - self.block_size = self.cache_config.block_size
775 - if self.pcp_size > 1:
776 - self.block_size *= self.pcp_size
777 - if self.dcp_size > 1:
778 - self.block_size *= self.dcp_size
808 + self.block_size, self.hash_block_size = resolve_kv_cache_block_sizes(
809 + kv_cache_config, vllm_config
810 + )
811 self.num_layers = model_config.get_num_layers(parallel_config)
812
813 self.use_mla = False
@@ -806,8 +838,6 @@ class MooncakeStoreWorker:
838 pp_rank=self.pp_rank,
839 )
840
809 - self.token_database = ChunkedTokenDatabase(self.metadata, self.block_size)
810 -
841 # Initialize MooncakeDistributedStore with its own TransferEngine
842 store_config = MooncakeStoreConfig.load_from_env()
843 extra_config = (
@@ -892,6 +922,45 @@ class MooncakeStoreWorker:
922 self.kv_recv_thread: KVCacheStoreRecvingThread | None = None
923 self.finished_store_req: set[str] = set()
924
925 + self._kv_cache_config = kv_cache_config
926 + # Single-group + PCP/DCP > 1: scale the lone group's spec.block_size to
927 + # self.block_size (= scheduler_block_size) so the coordinator's
928 + # ``block_size % hash_block_size == 0`` invariant holds.
929 + groups = list(kv_cache_config.kv_cache_groups)
930 + if len(groups) == 1 and groups[0].kv_cache_spec.block_size != self.block_size:
931 + g = groups[0]
932 + groups = [
933 + dataclasses.replace(
934 + g,
935 + kv_cache_spec=dataclasses.replace(
936 + g.kv_cache_spec, block_size=self.block_size
937 + ),
938 + )
939 + ]
940 + self._kv_cache_groups: list[KVCacheGroupSpec] = groups
941 + spec_cfg = getattr(vllm_config, "speculative_config", None)
942 + use_eagle = bool(
943 + spec_cfg.use_eagle()
944 + if spec_cfg is not None and callable(getattr(spec_cfg, "use_eagle", None))
945 + else False
946 + )
947 + self.coord = MooncakeStoreCoordinator(
948 + self._kv_cache_groups,
949 + scheduler_block_size=self.block_size,
950 + hash_block_size=self.hash_block_size,
951 + use_eagle=use_eagle,
952 + )
953 + # One ChunkedTokenDatabase per group; addresses populated in
954 + # register_kv_caches once the kv-cache layout is known.
955 + self.token_dbs: list[ChunkedTokenDatabase] = [
956 + ChunkedTokenDatabase(
957 + dataclasses.replace(self.metadata, group_id=g_idx),
958 + g.kv_cache_spec.block_size,
959 + hash_block_size=self.hash_block_size,
960 + )
961 + for g_idx, g in enumerate(self._kv_cache_groups)
962 + ]
963 +
964 def register_cross_layers_kv_caches(self, kv_cache: torch.Tensor) -> None:
965 """Register a cross-layers KV cache tensor.
966
@@ -901,93 +970,86 @@ class MooncakeStoreWorker:
970 """
971 self.register_kv_caches({"__cross_layer__": kv_cache})
972
904 - def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
973 + def register_kv_caches(
974 + self,
975 + kv_caches: dict[str, torch.Tensor | list[torch.Tensor]],
976 + ) -> None:
977 """Register KV cache tensors and start transfer threads."""
906 - # TODO(yifan): we haven't supported HMA yet.
907 - first_kv_cache = next(iter(kv_caches.values()))
978 + if not kv_caches:
979 + logger.warning("No KV caches to offload.")
980 + return
981 +
982 + # Resolve each entry to a representative tensor for storage
983 + # deduplication. For attention layers the value is already a tensor;
984 + # for Mamba layers it is a list of tensors that all share the same
985 + # underlying raw storage, so we take the first one.
986 + def _repr_tensor(v: torch.Tensor | list[torch.Tensor]) -> torch.Tensor:
987 + assert isinstance(v, torch.Tensor | list)
988 + return v if isinstance(v, torch.Tensor) else v[0]
989
909 - # num_blocks from cache_config is authoritative (set after
910 - # profiling, before KV cache allocation).
990 assert self.cache_config.num_gpu_blocks is not None
991 self.num_blocks = self.cache_config.num_gpu_blocks
992
914 - # Detect the KV cache memory layout using the stride-based
915 - # approach from simple_kv_offload/worker.py.
916 - #
917 - # The physical layout varies across attention backends:
918 - # FlashAttn/ROCm : (2, num_blocks, ...) → K/V outermost
919 - # FlashInfer/MLA : (num_blocks, ...) → blocks outermost
920 - #
921 - # We derive page_size_bytes = storage.nbytes() // num_blocks,
922 - # then classify dims: any dim whose byte-stride exceeds
923 - # page_size_bytes must be an outer segment dim (e.g. the K/V
924 - # dim of size 2). For those backends we register each segment
925 - # (K, V) as a separate base-address so that the per-block
926 - # offset arithmetic in prepare_value() stays correct.
927 - storage = first_kv_cache.untyped_storage()
928 - el = first_kv_cache.element_size()
929 - page_size_bytes = storage.nbytes() // self.num_blocks
930 - outer_dims = [
931 - d
932 - for d in range(first_kv_cache.ndim)
933 - if first_kv_cache.stride(d) * el > page_size_bytes
934 - ]
935 -
936 - # Register buffers with the store (deduplicate shared storages)
937 - # and record per-segment base addresses for every layer.
993 seen_ptrs: set[int] = set()
939 - self.kv_caches_base_addr: list[int] = []
940 - self.block_len: list[int] = []
994 + addrs: list[int] = []
995 + block_lens: list[int] = []
996
942 - for cache in kv_caches.values():
997 + for value in kv_caches.values():
998 + cache = _repr_tensor(value)
999 cache_storage = cache.untyped_storage()
1000 base_addr = cache_storage.data_ptr()
1001 + if base_addr in seen_ptrs:
1002 + continue
1003 + seen_ptrs.add(base_addr)
1004 region_len = cache_storage.nbytes()
1005
947 - if base_addr not in seen_ptrs:
948 - seen_ptrs.add(base_addr)
949 - ret = self.store.register_buffer(base_addr, region_len)
950 - if ret != 0:
951 - logger.error(
952 - "register_buffer failed for addr %#x len %d: %d",
953 - base_addr,
954 - region_len,
955 - ret,
956 - )
1006 + ret = self.store.register_buffer(base_addr, region_len)
1007 + if ret != 0:
1008 + logger.error(
1009 + "register_buffer failed for addr %#x len %d: %d",
1010 + base_addr,
1011 + region_len,
1012 + ret,
1013 + )
1014
1015 + # Detect layout via stride: a dim whose byte-stride exceeds
1016 + # page_size_bytes is an outer segment dim (e.g. the K/V dim of
1017 + # FlashAttn's (2, num_blocks, ...)). FlashInfer/MLA's blocks-
1018 + # outermost layout has no such dim and yields a single segment.
1019 + el = cache.element_size()
1020 + page_size_bytes = region_len // self.num_blocks
1021 + outer_dims = [
1022 + d for d in range(cache.ndim) if cache.stride(d) * el > page_size_bytes
1023 + ]
1024 if not outer_dims:
1025 # Blocks-first layout (FlashInfer / MLA): one segment.
960 - self.kv_caches_base_addr.append(base_addr)
961 - self.block_len.append(page_size_bytes)
1026 + addrs.append(base_addr)
1027 + block_lens.append(page_size_bytes)
1028 else:
1029 # K/V-first layout (FlashAttn / ROCm): split segments.
1030 seg_stride = cache.stride(outer_dims[0]) * el
1031 for idx in range(cache.shape[outer_dims[0]]):
966 - self.kv_caches_base_addr.append(base_addr + idx * seg_stride)
967 - self.block_len.append(seg_stride // self.num_blocks)
1032 + addrs.append(base_addr + idx * seg_stride)
1033 + block_lens.append(seg_stride // self.num_blocks)
1034
1035 logger.info(
970 - "Registering KV_Caches. use_mla: %s, shape %s, "
971 - "num_blocks: %d, block_len: %s, "
972 - "per_key_bytes: %d, "
973 - "num_segments: %d",
974 - self.use_mla,
975 - first_kv_cache.shape,
1036 + "Registered KV caches: num_groups=%d, num_segments=%d, num_blocks=%d",
1037 + len(self.token_dbs),
1038 + len(addrs),
1039 self.num_blocks,
977 - list(set(self.block_len)),
978 - sum(self.block_len),
979 - len(self.kv_caches_base_addr),
1040 )
1041
982 - self.token_database.set_kv_caches_base_addr(self.kv_caches_base_addr)
983 - self.token_database.set_block_len(self.block_len)
1042 + for db in self.token_dbs:
1043 + db.set_kv_caches_base_addr(addrs)
1044 + db.set_block_len(block_lens)
1045
1046 # Start transfer threads
1047 if self.kv_role in ["kv_producer", "kv_both"]:
1048 ready_event_sending = threading.Event()
1049 self.kv_send_thread = KVCacheStoreSendingThread(
1050 self.store,
990 - self.token_database,
1051 + self.coord,
1052 + self.token_dbs,
1053 self.block_size,
1054 self.tp_rank,
1055 self.put_step,
@@ -1001,7 +1063,8 @@ class MooncakeStoreWorker:
1063 ready_event_recving = threading.Event()
1064 self.kv_recv_thread = KVCacheStoreRecvingThread(
1065 self.store,
1004 - self.token_database,
1066 + self.coord,
1067 + self.token_dbs,
1068 self.block_size,
1069 self.tp_rank,
1070 ready_event_recving,
@@ -1122,62 +1185,55 @@ class MooncakeStoreWorker:
1185
1186 return finished_sending
1187
1125 - def lookup(
1126 - self,
1127 - token_len: int,
1128 - block_hashes: list[BlockHash],
1129 - ) -> int:
1188 + def lookup(self, token_len: int, block_hashes: list[BlockHash]) -> int:
1189 """Check how many prefix tokens exist in the store.
1190
1191 Checks across all TP ranks and PP ranks.
1192 """
1134 - end = 0
1135 - keys: list[str] = []
1136 - try:
1137 - starts: list[int] = []
1138 - for start, end, key in self.token_database.process_tokens(
1139 - token_len, block_hashes
1140 - ):
1141 - keys.append(key.to_string())
1142 - starts.append(start)
1193 + if not block_hashes or token_len <= 0:
1194 + return 0
1195
1144 - # Expand keys for all TP ranks
1145 - multi_tp_keys = keys[:]
1146 - for i in range(1, min(self.tp_size, self.num_kv_head)):
1147 - for item in keys:
1148 - new_str = item.replace("@tp_rank:0", f"@tp_rank:{i}", 1)
1149 - multi_tp_keys.append(new_str)
1150 -
1151 - # Expand keys for all PP ranks
1152 - pp_base_keys = multi_tp_keys.copy()
1153 - for i in range(1, self.pp_size):
1154 - for item in pp_base_keys:
1155 - new_str = item.replace("@pp_rank:0", f"@pp_rank:{i}", 1)
1156 - multi_tp_keys.append(new_str)
1157 -
1158 - res = self.store.batch_is_exist(multi_tp_keys)
1159 -
1160 - num_block = len(keys)
1161 - multi_tp_values = [
1162 - res[i * num_block : (i + 1) * num_block]
1163 - for i in range(min(self.tp_size, self.num_kv_head) * self.pp_size)
1164 - ]
1165 - index = self._find_min_first_non_one_index(multi_tp_values)
1166 - if index != -1:
1167 - return starts[index]
1196 + # Build per-(group, hash) candidate keys expanded across TP/PP.
1197 + # candidate_meta[i] is the (group_id, hash_bytes) for candidate_keys[i].
1198 + candidate_keys: list[str] = []
1199 + candidate_meta: list[tuple[int, bytes]] = []
1200 + tp_count = min(self.tp_size, self.num_kv_head)
1201 + for g_idx, db in enumerate(self.token_dbs):
1202 + spec_block_size = db.block_size
1203 + group_hashes = self.coord.block_hashes_for_spec(
1204 + block_hashes, self._kv_cache_groups[g_idx].kv_cache_spec
1205 + )
1206 + for chunk_id, h in enumerate(group_hashes):
1207 + start_idx = chunk_id * spec_block_size
1208 + if start_idx >= token_len:
1209 + break
1210 + for tp in range(tp_count):
1211 + for pp in range(self.pp_size):
1212 + md = dataclasses.replace(db.metadata, tp_rank=tp, pp_rank=pp)
1213 + candidate_keys.append(PoolKey(md, h.hex()).to_string())
1214 + candidate_meta.append((g_idx, bytes(h)))
1215 +
1216 + if not candidate_keys:
1217 + return 0
1218 +
1219 + try:
1220 + res = self.store.batch_is_exist(candidate_keys)
1221 except Exception as e:
1222 logger.error("Remote connection failed in lookup: %s", e)
1223 return 0
1171 - return end
1224
1173 - @staticmethod
1174 - def _find_min_first_non_one_index(
1175 - arr: list[list[int]],
1176 - ) -> int:
1177 - try:
1178 - return min(idx for row in arr for idx, val in enumerate(row) if val != 1)
1179 - except ValueError:
1180 - return -1
1225 + # A (group, hash) is "present" only when every TP*PP rank has it.
1226 + expected_per_key = max(1, tp_count * self.pp_size)
1227 + present_count: dict[tuple[int, bytes], int] = {}
1228 + for gh, exists in zip(candidate_meta, res, strict=True):
1229 + if exists == 1:
1230 + present_count[gh] = present_count.get(gh, 0) + 1
1231 + exists_set = {gh for gh, c in present_count.items() if c >= expected_per_key}
1232 +
1233 + _masks, hit_length = self.coord.find_longest_cache_hit(
1234 + block_hashes, token_len, ExternalCachedBlockPool(exists_set)
1235 + )
1236 + return hit_length
1237
1238 def get_kv_events(self) -> list[BlockStored]:
1239 if self.enable_kv_events and self.kv_send_thread is not None:
@@ -1220,7 +1276,8 @@ class LookupKeyServer:
1276 token_len = int.from_bytes(all_frames[0], byteorder="big")
1277 hash_frames = all_frames[1:]
1278 hashes_str = self.decoder.decode(hash_frames)
1223 - result = self.store_worker.lookup(token_len, hashes_str)
1279 + block_hashes = [BlockHash(bytes.fromhex(s)) for s in hashes_str]
1280 + result = self.store_worker.lookup(token_len, block_hashes)
1281 response = result.to_bytes(4, "big")
1282 self.socket.send(response)
1283