[Bugfix][KV Connector] MooncakeStore: exclude non-prefix-cacheable (QSA ring) groups; fix align-mode check (#55027)
Signed-off-by: Zhewen Li <zhewenli@inferact.ai> Co-authored-by: Zhewen Li <zhewenli@inferact.ai> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
Zhewen Li committed
Sep 13, 2026 at 19:18 UTC
58d45fd767ff81229afd92543400f5c8c93f4000
11 files changed
+381
-52
tests/v1/core/test_kv_cache_utils.py
+41
-1
@@ -35,7 +35,7 @@ from vllm.multimodal.inputs import (
35
from vllm.sampling_params import SamplingParams
36
from vllm.utils.hashing import sha256, sha256_cbor, xxhash, xxhash_cbor
37
from vllm.utils.mem_constants import GiB_bytes
38
-from vllm.v1.core.kv_cache_manager import KVCacheManager
38
+from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager
39
from vllm.v1.core.kv_cache_utils import (
40
BlockHash,
41
FreeKVCacheBlockQueue,
@@ -61,6 +61,7 @@ from vllm.v1.hisparse.layout import (
61
)
62
from vllm.v1.kv_cache_interface import (
63
ChunkedLocalAttentionSpec,
64
+ CircularBufferSpec,
65
FullAttentionSpec,
66
HiddenStateCacheSpec,
67
HiSparseHotSpec,
@@ -324,6 +325,45 @@ def test_kv_cache_config_selects_only_transferable_groups():
325
)
326
327
328
+def test_kv_cache_config_selects_prefix_cacheable_groups():
329
+ """Prefix stores exclude scratch state without changing transfer groups."""
330
+ full_group = KVCacheGroupSpec(["full"], new_kv_cache_spec())
331
+ qsa_group = KVCacheGroupSpec(
332
+ ["qsa"],
333
+ CircularBufferSpec(
334
+ block_size=4,
335
+ num_kv_heads=1,
336
+ head_size=64,
337
+ head_size_v=0,
338
+ dtype=torch.float16,
339
+ ),
340
+ )
341
+ disabled_group = KVCacheGroupSpec(
342
+ ["disabled"], new_kv_cache_spec(), enable_kv_transfer=False
343
+ )
344
+ config = KVCacheConfig(
345
+ num_blocks=1,
346
+ kv_cache_tensors=[],
347
+ kv_cache_groups=[full_group, qsa_group, disabled_group],
348
+ )
349
+ assert config.transfer_group_ids == (0, 1)
350
+ assert config.select_transfer_block_ids(([1], [2], [3])) == ([1], [2])
351
+ assert config.prefix_cacheable_group_ids == (0,)
352
+ assert config.prefix_cacheable_groups == (full_group,)
353
+
354
+
355
+def test_kv_cache_blocks_selects_requested_groups():
356
+ blocks = KVCacheBlocks(
357
+ (
358
+ [KVCacheBlock(1)],
359
+ [KVCacheBlock(2)],
360
+ [KVCacheBlock(3)],
361
+ )
362
+ )
363
+
364
+ assert blocks.get_block_ids(group_ids=(0, 2)) == ([1], [3])
365
+
366
+
367
def new_sliding_window_spec(
368
block_size=16,
369
num_kv_heads=2,
tests/v1/kv_connector/unit/test_mooncake_store_connector.py
+48
@@ -30,6 +30,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.metrics import
30
MooncakeStoreConnectorStats,
31
)
32
from vllm.v1.kv_cache_interface import (
33
+ CircularBufferSpec,
34
FullAttentionSpec,
35
KVCacheConfig,
36
KVCacheGroupSpec,
@@ -91,6 +92,53 @@ def _make_kv_cache_config() -> KVCacheConfig:
92
)
93
94
95
+def _make_qsa_hybrid_kv_cache_config(mamba_mode: str):
96
+ full_spec = FullAttentionSpec(
97
+ block_size=800, num_kv_heads=8, head_size=64, dtype=None
98
+ )
99
+ qsa_spec = CircularBufferSpec(
100
+ block_size=8,
101
+ num_kv_heads=1,
102
+ head_size=64,
103
+ head_size_v=0,
104
+ dtype=torch.float16,
105
+ )
106
+ mamba_spec = MambaSpec(
107
+ block_size=800,
108
+ shapes=((1, 1),),
109
+ dtypes=(torch.float32,),
110
+ mamba_cache_mode=mamba_mode,
111
+ )
112
+ return KVCacheConfig(
113
+ num_blocks=4,
114
+ kv_cache_tensors=[],
115
+ kv_cache_groups=[
116
+ KVCacheGroupSpec(["full"], full_spec),
117
+ KVCacheGroupSpec(["qsa"], qsa_spec),
118
+ KVCacheGroupSpec(["mamba"], mamba_spec),
119
+ ],
120
+ )
121
+
122
+
123
+def test_validation_accepts_aligned_mamba_after_block_size_rewrite():
124
+ vllm_config = _make_vllm_config()
125
+ vllm_config.cache_config.block_size = 8
126
+
127
+ mooncake_store_connector.MooncakeStoreConnector._validate_kv_cache_config(
128
+ vllm_config, _make_qsa_hybrid_kv_cache_config("align")
129
+ )
130
+
131
+
132
+def test_validation_rejects_mamba_mode_directly():
133
+ vllm_config = _make_vllm_config()
134
+ vllm_config.cache_config.block_size = 800
135
+
136
+ with pytest.raises(ValueError, match="mamba_cache_mode=\x27none\x27"):
137
+ mooncake_store_connector.MooncakeStoreConnector._validate_kv_cache_config(
138
+ vllm_config, _make_qsa_hybrid_kv_cache_config("none")
139
+ )
140
+
141
+
142
def test_scheduler_requires_align_mode_for_mamba():
143
vllm_config = _make_vllm_config()
144
mamba_spec = MambaSpec(
tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py
+12
-11
@@ -759,8 +759,9 @@ def test_worker_setup_tolerates_finer_scratch_group():
759
worker.pp_size = 1
760
worker.num_kv_head = 8
761
assert worker.coord.enable_partial_hash_hits
762
- # The scratch DB is keyed at its own block size and never probed.
763
- assert worker.token_dbs[2].hash_block_size == 4
762
+ # Only the participating full-attention and Mamba groups are registered.
763
+ assert len(worker.token_dbs) == 2
764
+ assert all(db.hash_block_size == 8 for db in worker.token_dbs)
765
766
for g_idx, db in enumerate(worker.token_dbs):
767
db.set_kv_caches_base_addr([g_idx * 10_000])
@@ -772,11 +773,11 @@ def test_worker_setup_tolerates_finer_scratch_group():
773
block_size=worker.block_size,
774
coord=worker.coord,
775
tp_rank=0,
775
- group_put_steps=[1, 1, 1],
776
+ group_put_steps=[1, 1],
777
kv_role="kv_both",
778
ready_event=threading.Event(),
779
replicate_config=MagicMock(),
779
- group_participates=[True, True, False],
780
+ group_participates=[True, True],
781
)
782
783
# Persist the sub-block partial tail at boundary 12 (keyed by hs[12//8-1]).
@@ -784,7 +785,7 @@ def test_worker_setup_tolerates_finer_scratch_group():
785
req = ReqMeta(
786
req_id="r0",
787
token_len_chunk=0,
787
- block_ids=([1], [2], [3]),
788
+ block_ids=([1], [2]),
789
block_hashes=hs,
790
can_save=True,
791
num_prompt_tokens=20,
@@ -797,8 +798,7 @@ def test_worker_setup_tolerates_finer_scratch_group():
798
# A 13-token prompt sharing the prefix must hit the first hash unit.
799
assert worker.lookup(num_tokens=13, block_hashes=hs).hit_length == 8
800
# The scratch group's namespace never enters the store.
800
- scratch_prefix = worker.token_dbs[2].key_for(hs[0]).rsplit("@", 1)[0]
801
- assert not any(key.startswith(scratch_prefix) for key in store._data)
801
+ assert not any("@group:2" in key for key in store._data)
802
803
804
def test_ring_scratch_group_is_never_stored_and_does_not_block_hits():
@@ -838,7 +838,8 @@ def test_ring_scratch_group_is_never_stored_and_does_not_block_hits():
838
worker.tp_size = 1
839
worker.pp_size = 1
840
worker.num_kv_head = 8
841
- assert worker.token_dbs[1].hash_block_size == 8
841
+ assert len(worker.token_dbs) == 1
842
+ assert worker.token_dbs[0].hash_block_size == 16
843
844
raw_full = torch.zeros(4 * full.page_size_bytes, dtype=torch.int8)
845
raw_ring = torch.zeros(4 * ring.page_size_bytes, dtype=torch.int8)
@@ -879,14 +880,14 @@ def test_ring_scratch_group_is_never_stored_and_does_not_block_hits():
880
kv_role=worker.kv_role,
881
ready_event=threading.Event(),
882
enable_kv_event=False,
882
- group_participates=[True, False],
883
+ group_participates=[True],
884
)
885
hs = [BlockHash(bytes([i + 1]) * 4) for i in range(4)]
886
save_req = ReqMeta(
887
req_id="r0",
888
token_len_chunk=64,
888
- # The ring group holds one block for the request's lifetime.
889
- block_ids=([1, 2, 3, 4], [1]),
889
+ # Scheduler metadata is projected to the participating store groups.
890
+ block_ids=([1, 2, 3, 4],),
891
block_hashes=hs,
892
can_save=True,
893
store_job_id=1,
tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py
+103
-7
@@ -2,12 +2,15 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from types import SimpleNamespace
5
+from unittest.mock import patch
6
7
import pytest
8
+import torch
9
10
from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
11
LoadSpec,
12
MooncakeLookupResult,
13
+ MooncakeStoreConnectorMetadata,
14
MooncakeStoreWorkerMetadata,
15
ReqMeta,
16
RequestTracker,
@@ -16,7 +19,16 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.scheduler impor
19
MooncakeStoreScheduler,
20
)
21
from vllm.v1.core.block_pool import BlockPool
22
+from vllm.v1.core.kv_cache_manager import KVCacheBlocks
23
+from vllm.v1.core.kv_cache_utils import KVCacheBlock
24
from vllm.v1.core.sched.output import KVConnectorBlockState
25
+from vllm.v1.kv_cache_interface import (
26
+ CircularBufferSpec,
27
+ FullAttentionSpec,
28
+ KVCacheConfig,
29
+ KVCacheGroupSpec,
30
+ MambaSpec,
31
+)
32
33
34
def _make_bare_scheduler(
@@ -35,9 +47,9 @@ def _make_bare_scheduler(
47
scheduler._block_size = 16
48
scheduler._hash_block_size = hash_block_size
49
scheduler.enable_partial_hash_hits = enable_partial_hash_hits
38
- scheduler.kv_cache_config = SimpleNamespace(
39
- select_transfer_block_ids=lambda block_ids: tuple(block_ids)
40
- )
50
+ scheduler.kv_cache_config = SimpleNamespace()
51
+ scheduler._store_group_ids = (0,)
52
+ scheduler._store_group_id_by_kv_cache_group_id = {0: 0, 1: 1}
53
scheduler.load_specs = {}
54
scheduler._unfinished_request_ids = {"req-0"}
55
scheduler._unfinished_requests = {}
@@ -228,14 +240,82 @@ def test_pending_load_for_non_chosen_connector_is_dropped():
240
assert "req-0" not in scheduler._request_trackers
241
242
243
+def _make_qsa_hybrid_cache_config():
244
+ full = FullAttentionSpec(block_size=800, num_kv_heads=8, head_size=64, dtype=None)
245
+ circular = CircularBufferSpec(
246
+ block_size=8,
247
+ num_kv_heads=1,
248
+ head_size=64,
249
+ head_size_v=0,
250
+ dtype=torch.float16,
251
+ )
252
+ mamba = MambaSpec(
253
+ block_size=800,
254
+ shapes=((1, 1),),
255
+ dtypes=(torch.float32,),
256
+ mamba_cache_mode="align",
257
+ )
258
+ return KVCacheConfig(
259
+ num_blocks=4,
260
+ kv_cache_tensors=[],
261
+ kv_cache_groups=[
262
+ KVCacheGroupSpec(["full"], full),
263
+ KVCacheGroupSpec(["qsa"], circular),
264
+ KVCacheGroupSpec(["mamba"], mamba),
265
+ ],
266
+ )
267
+
268
+
269
+def test_scheduler_projects_nonprefix_groups_and_mamba_ids():
270
+ vllm_config = SimpleNamespace(
271
+ kv_transfer_config=SimpleNamespace(
272
+ kv_role="kv_both", kv_connector_extra_config={}
273
+ ),
274
+ kv_events_config=None,
275
+ cache_config=SimpleNamespace(
276
+ block_size=800, enable_prefix_caching=True, prefix_match_unit=None
277
+ ),
278
+ parallel_config=SimpleNamespace(decode_context_parallel_size=1, world_size=1),
279
+ )
280
+
281
+ with patch(
282
+ "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
283
+ "scheduler.LookupKeyClient"
284
+ ):
285
+ scheduler = MooncakeStoreScheduler(vllm_config, _make_qsa_hybrid_cache_config())
286
+
287
+ assert scheduler._store_group_ids == (0, 2)
288
+ assert scheduler._boundary_state_group_ids == frozenset({1})
289
+
290
+
291
+def test_current_save_block_ids_use_store_group_projection():
292
+ scheduler = _make_bare_scheduler()
293
+ scheduler.kv_cache_config = _make_qsa_hybrid_cache_config()
294
+ scheduler._store_group_ids = (0, 2)
295
+ meta = MooncakeStoreConnectorMetadata(set(), set())
296
+ req_meta = ReqMeta(
297
+ req_id="req-0",
298
+ token_len_chunk=800,
299
+ block_ids=(),
300
+ block_hashes=[b"h0"],
301
+ can_save=True,
302
+ )
303
+ meta.add_request(req_meta)
304
+ output = SimpleNamespace(
305
+ kv_connector_block_state=_make_connector_block_state(([10], [80], [30]))
306
+ )
307
+
308
+ scheduler._apply_current_save_block_ids(meta, output)
309
+
310
+ assert req_meta.block_ids == ([10], [30])
311
+
312
+
313
def test_update_state_excludes_nontransfer_groups():
314
"""Store metadata must match the worker's registered cache groups."""
315
scheduler = _make_bare_scheduler()
234
- scheduler.kv_cache_config = SimpleNamespace(
235
- select_transfer_block_ids=lambda block_ids: (block_ids[0],)
236
- )
316
+ scheduler._store_group_ids = (0,)
317
request = SimpleNamespace(request_id="req-1")
238
- blocks = SimpleNamespace(get_block_ids=lambda: ([1, 2], [9]))
318
+ blocks = KVCacheBlocks(([KVCacheBlock(1), KVCacheBlock(2)], [KVCacheBlock(9)]))
319
320
scheduler.update_state_after_alloc(request, blocks, num_external_tokens=32)
321
@@ -1003,6 +1083,7 @@ def test_pending_partial_tail_emits_offload_only_reqmeta():
1083
1084
def test_finished_partial_tail_is_pre_pinned_as_store_job():
1085
scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True)
1086
+ scheduler._store_group_ids = (0, 1)
1087
scheduler.client = SimpleNamespace(discard=lambda *_: None)
1088
request = SimpleNamespace(
1089
request_id="req-0",
@@ -1143,6 +1224,20 @@ def _make_offload_only_output(entries, block_ids=([0],)):
1224
)
1225
1226
1227
+def test_boundary_state_group_ids_are_remapped_to_store_projection():
1228
+ scheduler = _make_bare_scheduler(hash_block_size=800, enable_partial_hash_hits=True)
1229
+ scheduler.kv_cache_config = _make_qsa_hybrid_cache_config()
1230
+ scheduler._store_group_ids = (0, 2)
1231
+ scheduler._store_group_id_by_kv_cache_group_id = {0: 0, 2: 1}
1232
+ scheduler._boundary_state_group_ids = frozenset({1})
1233
+ _register_offload_request(scheduler, prefill_end_tokens=800, num_prompt_tokens=800)
1234
+ meta = MooncakeStoreConnectorMetadata(set(), set())
1235
+
1236
+ scheduler._handle_boundary_state_offloads({"req-0": [(2, 7, 800)]}, meta)
1237
+
1238
+ assert meta.requests[0].boundary_state_offloads == [(1, 7, 800)]
1239
+
1240
+
1241
def test_resumed_prefill_claims_boundaries_past_prompt_length():
1242
# A resumed request re-prefills its previously generated tokens, so its
1243
# save window (`prefill_end_tokens`) extends past `num_prompt_tokens`.
@@ -1195,6 +1290,7 @@ def test_boundary_state_job_pins_exact_blocks_once():
1290
1291
def test_store_job_pins_current_non_null_non_mamba_blocks():
1292
scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True)
1293
+ scheduler._store_group_ids = (0, 1)
1294
request = SimpleNamespace(
1295
all_token_ids=list(range(48)),
1296
block_hashes=[bytes([i]) for i in range(12)],
tests/v1/kv_connector/unit/test_mooncake_store_worker.py
+90
@@ -46,8 +46,11 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.metrics import
46
from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
47
from vllm.v1.core.kv_cache_utils import BlockHash, maybe_convert_block_hash
48
from vllm.v1.kv_cache_interface import (
49
+ CircularBufferSpec,
50
FullAttentionSpec,
51
+ KVCacheConfig,
52
KVCacheGroupSpec,
53
+ MambaSpec,
54
RSWASpec,
55
)
56
from vllm.v1.kv_cache_layout import KVCacheLayout
@@ -342,6 +345,43 @@ def _make_kv_cache_config(
345
)
346
347
348
+def _make_qsa_hybrid_kv_cache_config():
349
+ return KVCacheConfig(
350
+ num_blocks=10,
351
+ kv_cache_tensors=[],
352
+ kv_cache_groups=[
353
+ KVCacheGroupSpec(
354
+ ["full"],
355
+ FullAttentionSpec(
356
+ block_size=800,
357
+ num_kv_heads=8,
358
+ head_size=64,
359
+ dtype=None,
360
+ ),
361
+ ),
362
+ KVCacheGroupSpec(
363
+ ["qsa"],
364
+ CircularBufferSpec(
365
+ block_size=8,
366
+ num_kv_heads=1,
367
+ head_size=64,
368
+ head_size_v=0,
369
+ dtype=torch.float16,
370
+ ),
371
+ ),
372
+ KVCacheGroupSpec(
373
+ ["mamba"],
374
+ MambaSpec(
375
+ block_size=800,
376
+ shapes=((1, 1),),
377
+ dtypes=(torch.float32,),
378
+ mamba_cache_mode="align",
379
+ ),
380
+ ),
381
+ ],
382
+ )
383
+
384
+
385
def _write_mooncake_config(tmp_path, config: dict[str, object]) -> str:
386
config_path = tmp_path / "mooncake_config.json"
387
config_path.write_text(json.dumps(config), encoding="utf-8")
@@ -2107,6 +2147,31 @@ def test_recv_thread_reports_unsplittable_key_larger_than_budget():
2147
assert thread.get_and_clear_block_ids_with_load_errors() == {0, 1, 2}
2148
2149
2150
+def test_worker_init_excludes_nonprefix_cache_groups(monkeypatch):
2151
+ store = MagicMock()
2152
+ store.setup.return_value = 0
2153
+ _install_fake_mooncake(monkeypatch, store)
2154
+ _patch_worker_runtime(monkeypatch)
2155
+ monkeypatch.setattr(
2156
+ worker.MooncakeStoreConfig,
2157
+ "load_from_config",
2158
+ staticmethod(lambda: _make_config()),
2159
+ )
2160
+ vllm_config = _make_vllm_config()
2161
+ vllm_config.cache_config.block_size = 800
2162
+ vllm_config.cache_config.enable_prefix_caching = True
2163
+ vllm_config.cache_config.prefix_match_unit = None
2164
+
2165
+ store_worker = worker.MooncakeStoreWorker(
2166
+ vllm_config, _make_qsa_hybrid_kv_cache_config()
2167
+ )
2168
+
2169
+ assert [
2170
+ group.kv_cache_spec.block_size for group in store_worker._kv_cache_groups
2171
+ ] == [800, 800]
2172
+ assert [db.block_size for db in store_worker.token_dbs] == [800, 800]
2173
+
2174
+
2175
def test_requester_worker_init_uses_positional_setup(tmp_path, monkeypatch):
2176
store = MagicMock()
2177
store.setup.return_value = 0
@@ -4144,6 +4209,31 @@ def test_register_kv_caches_uses_transfer_group_memory_domain():
4209
)
4210
4211
4212
+def test_register_kv_caches_excludes_nonprefix_cacheable_layers():
4213
+ num_blocks = 2
4214
+ worker = _make_bare_worker(num_gpu_blocks=num_blocks)
4215
+ store_layout = MagicMock()
4216
+ worker.token_dbs[0].store_layout = store_layout
4217
+ full_cache = torch.zeros(num_blocks, 4, dtype=torch.float16)
4218
+ qsa_cache = torch.zeros(num_blocks, 4, dtype=torch.float16)
4219
+
4220
+ _register_with_mocked_threads(
4221
+ worker,
4222
+ {"layer0": full_cache, "qsa": qsa_cache},
4223
+ )
4224
+
4225
+ store_layout.register_kv_caches.assert_called_once()
4226
+ registered_caches, registered_num_blocks = (
4227
+ store_layout.register_kv_caches.call_args.args
4228
+ )
4229
+ assert registered_num_blocks == num_blocks
4230
+ assert len(registered_caches) == 1
4231
+ assert registered_caches[0] is full_cache
4232
+ worker.store.register_buffer.assert_called_once_with(
4233
+ full_cache.data_ptr(), full_cache.untyped_storage().nbytes()
4234
+ )
4235
+
4236
+
4237
def test_register_kv_caches_separate_head_groups():
4238
# LHBNC gives each head group its own region; the K/V split doubles heads.
4239
layout = KVCacheLayout.LHBNC
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py
+12
-10
@@ -132,20 +132,22 @@ class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA):
132
from vllm.v1.kv_cache_interface import CrossAttentionSpec, MambaSpec
133
134
unsupported: list[str] = []
135
- cache_block_size = vllm_config.cache_config.block_size
136
- for g_idx, g in enumerate(kv_cache_config.transfer_groups):
137
- spec = g.kv_cache_spec
135
+ store_group_ids = kv_cache_config.prefix_cacheable_group_ids
136
+ if not store_group_ids:
137
+ raise ValueError(
138
+ "MooncakeStore requires at least one prefix-cacheable KV cache group"
139
+ )
140
+ for group_id in store_group_ids:
141
+ spec = kv_cache_config.kv_cache_groups[group_id].kv_cache_spec
142
if isinstance(spec, CrossAttentionSpec):
139
- unsupported.append(f"group {g_idx}: CrossAttentionSpec")
140
- # Enforce Mamba align mode
141
- if isinstance(spec, MambaSpec) and spec.block_size != cache_block_size:
143
+ unsupported.append(f"group {group_id}: CrossAttentionSpec")
144
+ if isinstance(spec, MambaSpec) and spec.mamba_cache_mode != "align":
145
unsupported.append(
143
- f"group {g_idx}: MambaSpec with block_size="
144
- f"{spec.block_size} != cache_config.block_size="
145
- f"{cache_block_size} (mamba_cache_mode != 'align')"
146
+ f"group {group_id}: mamba_cache_mode="
147
+ f"{spec.mamba_cache_mode!r} != 'align'"
148
)
149
pcp = vllm_config.parallel_config.prefill_context_parallel_size
148
- if len(kv_cache_config.transfer_groups) > 1 and pcp > 1:
150
+ if len(store_group_ids) > 1 and pcp > 1:
151
unsupported.append(f"PCP > 1 (pcp={pcp}) with hybrid attention")
152
if unsupported:
153
raise ValueError(
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py
+1
-1
@@ -452,7 +452,7 @@ def _unwrap_spec(spec: KVCacheSpec) -> KVCacheSpec:
452
453
454
def partial_hash_hits_enabled(
455
- kv_cache_groups: list[KVCacheGroupSpec],
455
+ kv_cache_groups: Sequence[KVCacheGroupSpec],
456
hash_block_size: int,
457
dcp_world_size: int = 1,
458
) -> bool:
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py
+30
-17
@@ -71,19 +71,27 @@ class MooncakeStoreScheduler:
71
)
72
self.client = LookupKeyClient(vllm_config)
73
self.kv_cache_config = kv_cache_config
74
+ self._store_group_ids = kv_cache_config.prefix_cacheable_group_ids
75
+ # Map scheduler group IDs to store group indices. Groups outside the
76
+ # store projection, such as the QSA ring, map to None and are skipped.
77
+ self._store_group_id_by_kv_cache_group_id = {
78
+ group_id: store_group_id
79
+ for store_group_id, group_id in enumerate(self._store_group_ids)
80
+ }
81
+ store_groups = kv_cache_config.prefix_cacheable_groups
82
83
# Align with the engine's own scheduler_block_size and hash_block_size.
84
self._block_size, self._hash_block_size = resolve_kv_cache_block_sizes(
85
kv_cache_config, vllm_config
86
)
87
self.enable_partial_hash_hits = partial_hash_hits_enabled(
80
- kv_cache_config.kv_cache_groups,
88
+ store_groups,
89
self._hash_block_size,
90
vllm_config.parallel_config.decode_context_parallel_size,
91
)
92
mamba_groups = {
93
group_id: group.kv_cache_spec
86
- for group_id, group in enumerate(kv_cache_config.kv_cache_groups)
94
+ for group_id, group in enumerate(store_groups)
95
if isinstance(group.kv_cache_spec, MambaSpec)
96
}
97
assert all(
@@ -173,8 +181,8 @@ class MooncakeStoreScheduler:
181
"""Update state after block allocation."""
182
local_block_ids: tuple[list[int], ...] = ()
183
if num_external_tokens > 0:
176
- local_block_ids = self.kv_cache_config.select_transfer_block_ids(
177
- blocks.get_block_ids()
184
+ local_block_ids = blocks.get_block_ids(
185
+ group_ids=self._store_group_ids,
186
)
187
188
self._unfinished_requests[request.request_id] = (request, local_block_ids)
@@ -239,10 +247,7 @@ class MooncakeStoreScheduler:
247
request_real = request_tuple[0] # type: ignore[index]
248
249
unfolded_block_ids = tuple(
242
- blocks.copy()
243
- for blocks in self.kv_cache_config.select_transfer_block_ids(
244
- request.block_ids
245
- )
250
+ request.block_ids[group_id].copy() for group_id in self._store_group_ids
251
)
252
253
prefill_tokens = _new_req_prefill_tokens(request)
@@ -278,8 +283,8 @@ class MooncakeStoreScheduler:
283
for i, req_id in enumerate(cached_reqs.req_ids):
284
new_block_ids = cached_reqs.new_block_ids[i]
285
if new_block_ids:
281
- new_block_ids = self.kv_cache_config.select_transfer_block_ids(
282
- new_block_ids
286
+ new_block_ids = tuple(
287
+ new_block_ids[group_id] for group_id in self._store_group_ids
288
)
289
290
req_meta = None
@@ -438,7 +443,9 @@ class MooncakeStoreScheduler:
443
assert block_ids is not None, (
444
f"Missing current block table for store request {req_meta.req_id}"
445
)
441
- req_meta.block_ids = block_ids
446
+ req_meta.block_ids = tuple(
447
+ block_ids[group_id] for group_id in self._store_group_ids
448
+ )
449
450
def _reference_save_blocks(self, meta: MooncakeStoreConnectorMetadata) -> None:
451
"""Take a GPU block reference for every store job this step emits.
@@ -509,12 +516,15 @@ class MooncakeStoreScheduler:
516
return False
517
518
pinned_block_ids: list[int] = []
512
- for group_id, block_id, _ in partial_tail_offloads:
513
- if group_id not in self._boundary_state_group_ids:
519
+ remapped_offloads: list[tuple[int, int, int]] = []
520
+ for group_id, block_id, boundary in partial_tail_offloads:
521
+ store_group_id = self._store_group_id_by_kv_cache_group_id.get(group_id)
522
+ if store_group_id not in self._boundary_state_group_ids:
523
return False
524
if block_id == NULL_BLOCK_ID:
525
return False
526
pinned_block_ids.append(block_id)
527
+ remapped_offloads.append((store_group_id, block_id, boundary))
528
pinned_block_ids = list(dict.fromkeys(pinned_block_ids))
529
530
pool = self._gpu_block_pool
@@ -530,12 +540,14 @@ class MooncakeStoreScheduler:
540
self._finished_partial_tail_metas[request.request_id] = ReqMeta(
541
req_id=request.request_id,
542
token_len_chunk=0,
533
- block_ids=tuple(group.copy() for group in block_ids),
543
+ block_ids=tuple(
544
+ block_ids[group_id].copy() for group_id in self._store_group_ids
545
+ ),
546
block_hashes=list(request.block_hashes),
547
can_save=True,
548
num_prompt_tokens=tracker.prefill_end_tokens,
549
store_job_id=store_job_id,
538
- boundary_state_offloads=partial_tail_offloads,
550
+ boundary_state_offloads=remapped_offloads,
551
)
552
tracker.has_pending_offload = True
553
# The store job owns exact block refs, so request cleanup need not wait.
@@ -574,9 +586,10 @@ class MooncakeStoreScheduler:
586
continue
587
if block_id == NULL_BLOCK_ID:
588
continue
577
- if group_id not in self._boundary_state_group_ids:
589
+ store_group_id = self._store_group_id_by_kv_cache_group_id.get(group_id)
590
+ if store_group_id not in self._boundary_state_group_ids:
591
continue
579
- accepted.append((group_id, block_id, boundary_tokens))
592
+ accepted.append((store_group_id, block_id, boundary_tokens))
593
if not accepted:
594
continue
595
tracker.has_pending_offload = True
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py
+10
-3
@@ -1672,10 +1672,11 @@ class MooncakeStoreWorker:
1672
dataclasses.replace(
1673
group,
1674
kv_cache_spec=resolve_dcp_kv_cache_spec(
1675
- group.kv_cache_spec, self.dcp_size
1675
+ group.kv_cache_spec,
1676
+ self.dcp_size,
1677
),
1678
)
1678
- for group in kv_cache_config.transfer_groups
1679
+ for group in kv_cache_config.prefix_cacheable_groups
1680
]
1681
spec_cfg = getattr(vllm_config, "speculative_config", None)
1682
use_eagle_block_drop = bool(
@@ -1923,10 +1924,16 @@ class MooncakeStoreWorker:
1924
use_group_regions = self._kv_cache_config.hisparse_host_num_blocks is not None
1925
1926
if not use_group_regions:
1927
+ store_layer_names = {
1928
+ layer_name
1929
+ for group in self._kv_cache_groups
1930
+ for layer_name in group.layer_names
1931
+ }
1932
seen_storage_ptrs: set[int] = set()
1933
cache_tensors = [
1934
group_kernel_blocks(_repr_tensor(cache), self.num_blocks)
1929
- for cache in kv_caches.values()
1935
+ for layer_name, cache in kv_caches.items()
1936
+ if layer_name in store_layer_names
1937
]
1938
for cache in cache_tensors:
1939
cache_storage = cache.untyped_storage()
vllm/v1/core/kv_cache_manager.py
+17
-2
@@ -67,30 +67,45 @@ class KVCacheBlocks:
67
def get_block_ids(
68
self,
69
allow_none: Literal[False] = False,
70
+ *,
71
+ group_ids: Sequence[int] | None = None,
72
) -> tuple[list[int], ...]: ...
73
74
@overload
75
def get_block_ids(
76
self,
77
allow_none: Literal[True] = True,
78
+ *,
79
+ group_ids: Sequence[int] | None = None,
80
) -> tuple[list[int], ...] | None: ...
81
82
def get_block_ids(
83
self,
84
allow_none: bool = False,
85
+ *,
86
+ group_ids: Sequence[int] | None = None,
87
) -> tuple[list[int], ...] | None:
88
"""
89
Converts the KVCacheBlocks instance to block_ids.
90
91
+ Args:
92
+ allow_none: Return None when every selected group is empty.
93
+ group_ids: KV cache groups to include. Includes all groups by default.
94
+
95
Returns:
96
tuple[list[int], ...]: A tuple of lists where:
97
- the outer tuple corresponds to KV cache groups
98
- each inner list contains the block_ids of the blocks in that
99
group
100
"""
91
- if allow_none and all(len(group) == 0 for group in self.blocks):
101
+ groups = (
102
+ self.blocks
103
+ if group_ids is None
104
+ else tuple(self.blocks[group_id] for group_id in group_ids)
105
+ )
106
+ if allow_none and all(len(group) == 0 for group in groups):
107
return None
93
- return tuple([blk.block_id for blk in group] for group in self.blocks)
108
+ return tuple([blk.block_id for blk in group] for group in groups)
109
110
def get_unhashed_block_ids(self) -> list[int]:
111
"""Get block_ids of unhashed blocks from KVCacheBlocks instance."""
vllm/v1/kv_cache_interface.py
+17
@@ -1450,6 +1450,23 @@ class KVCacheConfig:
1450
if group.enable_kv_transfer
1451
)
1452
1453
+ @cached_property
1454
+ def prefix_cacheable_group_ids(self) -> tuple[int, ...]:
1455
+ """IDs of transferable groups eligible for hash-addressed stores."""
1456
+ return tuple(
1457
+ group_id
1458
+ for group_id in self.transfer_group_ids
1459
+ if self.kv_cache_groups[group_id].kv_cache_spec.prefix_cacheable
1460
+ )
1461
+
1462
+ @cached_property
1463
+ def prefix_cacheable_groups(self) -> tuple[KVCacheGroupSpec, ...]:
1464
+ """Transferable groups eligible for hash-addressed stores."""
1465
+ return tuple(
1466
+ self.kv_cache_groups[group_id]
1467
+ for group_id in self.prefix_cacheable_group_ids
1468
+ )
1469
+
1470
@cached_property
1471
def transfer_groups(self) -> tuple[KVCacheGroupSpec, ...]:
1472
"""Cache groups that participate in external KV transfer."""