[Bugfix][KV Connector] MooncakeStore: exclude non-prefix-cacheable (QSA ring) groups; fix align-mode check

tests/v1/core/test_kv_cache_utils.py

@@ -35,7 +35,7 @@
 from vllm.sampling_params import SamplingParams
 from vllm.utils.hashing import sha256, sha256_cbor, xxhash, xxhash_cbor
 from vllm.utils.mem_constants import GiB_bytes
-from vllm.v1.core.kv_cache_manager import KVCacheManager
+from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager
 from vllm.v1.core.kv_cache_utils import (
     BlockHash,
     FreeKVCacheBlockQueue,
@@ -61,6 +61,7 @@
 )
 from vllm.v1.kv_cache_interface import (
     ChunkedLocalAttentionSpec,
+    CircularBufferSpec,
     FullAttentionSpec,
     HiddenStateCacheSpec,
     HiSparseHotSpec,
@@ -324,6 +325,45 @@ def test_kv_cache_config_selects_only_transferable_groups():
     )
 
 
+def test_kv_cache_config_selects_prefix_cacheable_groups():
+    """Prefix stores exclude scratch state without changing transfer groups."""
+    full_group = KVCacheGroupSpec(["full"], new_kv_cache_spec())
+    qsa_group = KVCacheGroupSpec(
+        ["qsa"],
+        CircularBufferSpec(
+            block_size=4,
+            num_kv_heads=1,
+            head_size=64,
+            head_size_v=0,
+            dtype=torch.float16,
+        ),
+    )
+    disabled_group = KVCacheGroupSpec(
+        ["disabled"], new_kv_cache_spec(), enable_kv_transfer=False
+    )
+    config = KVCacheConfig(
+        num_blocks=1,
+        kv_cache_tensors=[],
+        kv_cache_groups=[full_group, qsa_group, disabled_group],
+    )
+    assert config.transfer_group_ids == (0, 1)
+    assert config.select_transfer_block_ids(([1], [2], [3])) == ([1], [2])
+    assert config.prefix_cacheable_group_ids == (0,)
+    assert config.prefix_cacheable_groups == (full_group,)
+
+
+def test_kv_cache_blocks_selects_requested_groups():
+    blocks = KVCacheBlocks(
+        (
+            [KVCacheBlock(1)],
+            [KVCacheBlock(2)],
+            [KVCacheBlock(3)],
+        )
+    )
+
+    assert blocks.get_block_ids(group_ids=(0, 2)) == ([1], [3])
+
+
 def new_sliding_window_spec(
     block_size=16,
     num_kv_heads=2,

tests/v1/kv_connector/unit/test_mooncake_store_connector.py

@@ -30,6 +30,7 @@
     MooncakeStoreConnectorStats,
 )
 from vllm.v1.kv_cache_interface import (
+    CircularBufferSpec,
     FullAttentionSpec,
     KVCacheConfig,
     KVCacheGroupSpec,
@@ -91,6 +92,53 @@ def _make_kv_cache_config() -> KVCacheConfig:
     )
 
 
+def _make_qsa_hybrid_kv_cache_config(mamba_mode: str):
+    full_spec = FullAttentionSpec(
+        block_size=800, num_kv_heads=8, head_size=64, dtype=None
+    )
+    qsa_spec = CircularBufferSpec(
+        block_size=8,
+        num_kv_heads=1,
+        head_size=64,
+        head_size_v=0,
+        dtype=torch.float16,
+    )
+    mamba_spec = MambaSpec(
+        block_size=800,
+        shapes=((1, 1),),
+        dtypes=(torch.float32,),
+        mamba_cache_mode=mamba_mode,
+    )
+    return KVCacheConfig(
+        num_blocks=4,
+        kv_cache_tensors=[],
+        kv_cache_groups=[
+            KVCacheGroupSpec(["full"], full_spec),
+            KVCacheGroupSpec(["qsa"], qsa_spec),
+            KVCacheGroupSpec(["mamba"], mamba_spec),
+        ],
+    )
+
+
+def test_validation_accepts_aligned_mamba_after_block_size_rewrite():
+    vllm_config = _make_vllm_config()
+    vllm_config.cache_config.block_size = 8
+
+    mooncake_store_connector.MooncakeStoreConnector._validate_kv_cache_config(
+        vllm_config, _make_qsa_hybrid_kv_cache_config("align")
+    )
+
+
+def test_validation_rejects_mamba_mode_directly():
+    vllm_config = _make_vllm_config()
+    vllm_config.cache_config.block_size = 800
+
+    with pytest.raises(ValueError, match="mamba_cache_mode=\x27none\x27"):
+        mooncake_store_connector.MooncakeStoreConnector._validate_kv_cache_config(
+            vllm_config, _make_qsa_hybrid_kv_cache_config("none")
+        )
+
+
 def test_scheduler_requires_align_mode_for_mamba():
     vllm_config = _make_vllm_config()
     mamba_spec = MambaSpec(

tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py

@@ -759,8 +759,9 @@ def test_worker_setup_tolerates_finer_scratch_group():
     worker.pp_size = 1
     worker.num_kv_head = 8
     assert worker.coord.enable_partial_hash_hits
-    # The scratch DB is keyed at its own block size and never probed.
-    assert worker.token_dbs[2].hash_block_size == 4
+    # Only the participating full-attention and Mamba groups are registered.
+    assert len(worker.token_dbs) == 2
+    assert all(db.hash_block_size == 8 for db in worker.token_dbs)
 
     for g_idx, db in enumerate(worker.token_dbs):
         db.set_kv_caches_base_addr([g_idx * 10_000])
@@ -772,19 +773,19 @@ def test_worker_setup_tolerates_finer_scratch_group():
         block_size=worker.block_size,
         coord=worker.coord,
         tp_rank=0,
-        group_put_steps=[1, 1, 1],
+        group_put_steps=[1, 1],
         kv_role="kv_both",
         ready_event=threading.Event(),
         replicate_config=MagicMock(),
-        group_participates=[True, True, False],
+        group_participates=[True, True],
     )
 
     # Persist the sub-block partial tail at boundary 12 (keyed by hs[12//8-1]).
     hs = [BlockHash(bytes([i + 1]) * 8) for i in range(3)]
     req = ReqMeta(
         req_id="r0",
         token_len_chunk=0,
-        block_ids=([1], [2], [3]),
+        block_ids=([1], [2]),
         block_hashes=hs,
         can_save=True,
         num_prompt_tokens=20,
@@ -797,8 +798,7 @@ def test_worker_setup_tolerates_finer_scratch_group():
     # A 13-token prompt sharing the prefix must hit the first hash unit.
     assert worker.lookup(num_tokens=13, block_hashes=hs).hit_length == 8
     # The scratch group's namespace never enters the store.
-    scratch_prefix = worker.token_dbs[2].key_for(hs[0]).rsplit("@", 1)[0]
-    assert not any(key.startswith(scratch_prefix) for key in store._data)
+    assert not any("@group:2" in key for key in store._data)
 
 
 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():
     worker.tp_size = 1
     worker.pp_size = 1
     worker.num_kv_head = 8
-    assert worker.token_dbs[1].hash_block_size == 8
+    assert len(worker.token_dbs) == 1
+    assert worker.token_dbs[0].hash_block_size == 16
 
     raw_full = torch.zeros(4 * full.page_size_bytes, dtype=torch.int8)
     raw_ring = torch.zeros(4 * ring.page_size_bytes, dtype=torch.int8)
@@ -879,14 +880,14 @@ def _fake_thread_init(*args, **kwargs):
         kv_role=worker.kv_role,
         ready_event=threading.Event(),
         enable_kv_event=False,
-        group_participates=[True, False],
+        group_participates=[True],
     )
     hs = [BlockHash(bytes([i + 1]) * 4) for i in range(4)]
     save_req = ReqMeta(
         req_id="r0",
         token_len_chunk=64,
-        # The ring group holds one block for the request's lifetime.
-        block_ids=([1, 2, 3, 4], [1]),
+        # Scheduler metadata is projected to the participating store groups.
+        block_ids=([1, 2, 3, 4],),
         block_hashes=hs,
         can_save=True,
         store_job_id=1,

tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py

@@ -2,12 +2,15 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from types import SimpleNamespace
+from unittest.mock import patch
 
 import pytest
+import torch
 
 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
     LoadSpec,
     MooncakeLookupResult,
+    MooncakeStoreConnectorMetadata,
     MooncakeStoreWorkerMetadata,
     ReqMeta,
     RequestTracker,
@@ -16,7 +19,16 @@
     MooncakeStoreScheduler,
 )
 from vllm.v1.core.block_pool import BlockPool
+from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+from vllm.v1.core.kv_cache_utils import KVCacheBlock
 from vllm.v1.core.sched.output import KVConnectorBlockState
+from vllm.v1.kv_cache_interface import (
+    CircularBufferSpec,
+    FullAttentionSpec,
+    KVCacheConfig,
+    KVCacheGroupSpec,
+    MambaSpec,
+)
 
 
 def _make_bare_scheduler(
@@ -35,9 +47,9 @@ def _make_bare_scheduler(
     scheduler._block_size = 16
     scheduler._hash_block_size = hash_block_size
     scheduler.enable_partial_hash_hits = enable_partial_hash_hits
-    scheduler.kv_cache_config = SimpleNamespace(
-        select_transfer_block_ids=lambda block_ids: tuple(block_ids)
-    )
+    scheduler.kv_cache_config = SimpleNamespace()
+    scheduler._store_group_ids = (0,)
+    scheduler._store_group_id_by_kv_cache_group_id = {0: 0, 1: 1}
     scheduler.load_specs = {}
     scheduler._unfinished_request_ids = {"req-0"}
     scheduler._unfinished_requests = {}
@@ -228,14 +240,82 @@ def test_pending_load_for_non_chosen_connector_is_dropped():
     assert "req-0" not in scheduler._request_trackers
 
 
+def _make_qsa_hybrid_cache_config():
+    full = FullAttentionSpec(block_size=800, num_kv_heads=8, head_size=64, dtype=None)
+    circular = CircularBufferSpec(
+        block_size=8,
+        num_kv_heads=1,
+        head_size=64,
+        head_size_v=0,
+        dtype=torch.float16,
+    )
+    mamba = MambaSpec(
+        block_size=800,
+        shapes=((1, 1),),
+        dtypes=(torch.float32,),
+        mamba_cache_mode="align",
+    )
+    return KVCacheConfig(
+        num_blocks=4,
+        kv_cache_tensors=[],
+        kv_cache_groups=[
+            KVCacheGroupSpec(["full"], full),
+            KVCacheGroupSpec(["qsa"], circular),
+            KVCacheGroupSpec(["mamba"], mamba),
+        ],
+    )
+
+
+def test_scheduler_projects_nonprefix_groups_and_mamba_ids():
+    vllm_config = SimpleNamespace(
+        kv_transfer_config=SimpleNamespace(
+            kv_role="kv_both", kv_connector_extra_config={}
+        ),
+        kv_events_config=None,
+        cache_config=SimpleNamespace(
+            block_size=800, enable_prefix_caching=True, prefix_match_unit=None
+        ),
+        parallel_config=SimpleNamespace(decode_context_parallel_size=1, world_size=1),
+    )
+
+    with patch(
+        "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
+        "scheduler.LookupKeyClient"
+    ):
+        scheduler = MooncakeStoreScheduler(vllm_config, _make_qsa_hybrid_cache_config())
+
+    assert scheduler._store_group_ids == (0, 2)
+    assert scheduler._boundary_state_group_ids == frozenset({1})
+
+
+def test_current_save_block_ids_use_store_group_projection():
+    scheduler = _make_bare_scheduler()
+    scheduler.kv_cache_config = _make_qsa_hybrid_cache_config()
+    scheduler._store_group_ids = (0, 2)
+    meta = MooncakeStoreConnectorMetadata(set(), set())
+    req_meta = ReqMeta(
+        req_id="req-0",
+        token_len_chunk=800,
+        block_ids=(),
+        block_hashes=[b"h0"],
+        can_save=True,
+    )
+    meta.add_request(req_meta)
+    output = SimpleNamespace(
+        kv_connector_block_state=_make_connector_block_state(([10], [80], [30]))
+    )
+
+    scheduler._apply_current_save_block_ids(meta, output)
+
+    assert req_meta.block_ids == ([10], [30])
+
+
 def test_update_state_excludes_nontransfer_groups():
     """Store metadata must match the worker's registered cache groups."""
     scheduler = _make_bare_scheduler()
-    scheduler.kv_cache_config = SimpleNamespace(
-        select_transfer_block_ids=lambda block_ids: (block_ids[0],)
-    )
+    scheduler._store_group_ids = (0,)
     request = SimpleNamespace(request_id="req-1")
-    blocks = SimpleNamespace(get_block_ids=lambda: ([1, 2], [9]))
+    blocks = KVCacheBlocks(([KVCacheBlock(1), KVCacheBlock(2)], [KVCacheBlock(9)]))
 
     scheduler.update_state_after_alloc(request, blocks, num_external_tokens=32)
 
@@ -1003,6 +1083,7 @@ def test_pending_partial_tail_emits_offload_only_reqmeta():
 
 def test_finished_partial_tail_is_pre_pinned_as_store_job():
     scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True)
+    scheduler._store_group_ids = (0, 1)
     scheduler.client = SimpleNamespace(discard=lambda *_: None)
     request = SimpleNamespace(
         request_id="req-0",
@@ -1143,6 +1224,20 @@ def _make_offload_only_output(entries, block_ids=([0],)):
     )
 
 
+def test_boundary_state_group_ids_are_remapped_to_store_projection():
+    scheduler = _make_bare_scheduler(hash_block_size=800, enable_partial_hash_hits=True)
+    scheduler.kv_cache_config = _make_qsa_hybrid_cache_config()
+    scheduler._store_group_ids = (0, 2)
+    scheduler._store_group_id_by_kv_cache_group_id = {0: 0, 2: 1}
+    scheduler._boundary_state_group_ids = frozenset({1})
+    _register_offload_request(scheduler, prefill_end_tokens=800, num_prompt_tokens=800)
+    meta = MooncakeStoreConnectorMetadata(set(), set())
+
+    scheduler._handle_boundary_state_offloads({"req-0": [(2, 7, 800)]}, meta)
+
+    assert meta.requests[0].boundary_state_offloads == [(1, 7, 800)]
+
+
 def test_resumed_prefill_claims_boundaries_past_prompt_length():
     # A resumed request re-prefills its previously generated tokens, so its
     # save window (`prefill_end_tokens`) extends past `num_prompt_tokens`.
@@ -1195,6 +1290,7 @@ def test_boundary_state_job_pins_exact_blocks_once():
 
 def test_store_job_pins_current_non_null_non_mamba_blocks():
     scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True)
+    scheduler._store_group_ids = (0, 1)
     request = SimpleNamespace(
         all_token_ids=list(range(48)),
         block_hashes=[bytes([i]) for i in range(12)],