[Bugfix][Core] Isolate hidden-state cache from DeepSeek-V4 MLA groups

tests/v1/core/test_kv_cache_utils.py

@@ -2241,6 +2241,84 @@ def test_group_and_unify_kv_cache_specs_mixed_page_size_groups():
     assert layer_names == {"mla.0", "mla.1", "swa.0"}
 
 
+def new_hidden_state_spec(block_size=16):
+    return HiddenStateCacheSpec(
+        block_size=block_size,
+        num_kv_heads=4,
+        head_size=4096,
+        dtype=torch.bfloat16,
+    )
+
+
+def _dsv4_specs_with_hidden():
+    return {
+        "mla.0": new_mla_spec(),
+        "mla.1": new_mla_spec(),
+        "swa.0": new_swa_mla_spec(head_size=1024),
+        "cache_only_layers.0": new_hidden_state_spec(),
+    }
+
+
+def test_group_and_unify_kv_cache_specs_excludes_hidden_state():
+    specs = _dsv4_specs_with_hidden()
+    grouped = group_and_unify_kv_cache_specs(specs)
+    grouped_without_hidden = group_and_unify_kv_cache_specs(
+        {name: spec for name, spec in specs.items() if "cache_only" not in name}
+    )
+
+    assert grouped is not None
+    assert grouped_without_hidden is not None
+    assert {
+        name for group in grouped for name in group.kv_cache_specs
+    } == {"mla.0", "mla.1", "swa.0"}
+    assert [sorted(group.get_page_sizes()) for group in grouped] == [
+        sorted(group.get_page_sizes()) for group in grouped_without_hidden
+    ]
+    assert set(specs) == {
+        "mla.0", "mla.1", "swa.0", "cache_only_layers.0"
+    }
+
+
+def test_hidden_state_page_does_not_force_dsv4_packing():
+    # The hidden-state page is intentionally much larger than the attention
+    # pages. It must not make an otherwise uniform MLA/SWA layout enter DSV4
+    # tuple packing.
+    specs = {
+        "mla.0": new_mla_spec(),
+        "swa.0": new_swa_mla_spec(),
+        "cache_only_layers.0": new_hidden_state_spec(),
+    }
+    assert group_and_unify_kv_cache_specs(specs) is None
+
+
+def test_get_kv_cache_groups_isolates_hidden_state():
+    specs = _dsv4_specs_with_hidden()
+    groups = get_kv_cache_groups(_grouping_config(), specs)
+
+    assert {
+        name for group in groups for name in group.layer_names
+    } == set(specs)
+    assert sum(
+        "cache_only_layers.0" in group.layer_names for group in groups
+    ) == 1
+
+    hidden_group = next(
+        group for group in groups if "cache_only_layers.0" in group.layer_names
+    )
+    assert isinstance(hidden_group.kv_cache_spec, UniformTypeKVCacheSpecs)
+    assert hidden_group.layer_names == ["cache_only_layers.0"]
+    assert list(hidden_group.kv_cache_spec.kv_cache_specs) == [
+        "cache_only_layers.0"
+    ]
+    assert all(
+        isinstance(spec, HiddenStateCacheSpec)
+        for spec in hidden_group.kv_cache_spec.kv_cache_specs.values()
+    )
+    assert hidden_group.kv_cache_spec.block_size == specs[
+        "cache_only_layers.0"
+    ].block_size
+
+
 def new_indexer_mla_spec(block_size=16):
     # Sparse-attention indexer k_cache: an MLAAttentionSpec with a much smaller
     # page size than the main MLA attention (uint8, small head), so their pages

tests/v1/kv_connector/unit/test_hidden_states_connector.py

@@ -17,6 +17,7 @@
     KVCacheGroupSpec,
     MLAAttentionSpec,
     SlidingWindowMLASpec,
+    UniformTypeKVCacheSpecs,
 )
 
 
@@ -68,13 +69,17 @@ def test_find_group_id_locates_hidden_group_last():
 
 def test_find_group_id_raises_when_no_hidden_group_and_multiple_groups():
     cfg = _config(_full(16), _full(16))
-    with pytest.raises(ValueError, match="Could not uniquely identify"):
+    with pytest.raises(
+        ValueError, match="isolated in exactly one KV cache group"
+    ):
         ExampleHiddenStatesConnector._find_cache_kv_group_id(cfg)
 
 
 def test_find_group_id_raises_when_multiple_hidden_groups():
     cfg = _config(_hidden(22), _hidden(22))
-    with pytest.raises(ValueError, match="Could not uniquely identify"):
+    with pytest.raises(
+        ValueError, match="isolated in exactly one KV cache group"
+    ):
         ExampleHiddenStatesConnector._find_cache_kv_group_id(cfg)
 
 
@@ -99,19 +104,37 @@ def test_get_block_size_falls_back_to_cache_config_when_no_kv_cache_config():
     assert block_size == 16
 
 
-# ---- MLA-verifier absorption ------------------------------------------------
+# ---- DSV4 isolation and connector lookup -----------------------------------
 
 
-def test_find_group_id_errors_clearly_when_absorbed_by_mla_swa_verifier():
-    # HiddenStateCacheSpec subclasses MLAAttentionSpec, so an MLA + sliding-
-    # window MLA verifier absorbs it into the MLA group instead of isolating it.
+def test_find_group_id_locates_wrapped_hidden_group_and_block_size():
+    hidden = _hidden(22)
+    hidden_group = UniformTypeKVCacheSpecs(
+        block_size=22, kv_cache_specs={"cache_only_layers.61": hidden}
+    )
+    cfg = _config(_full(528), hidden_group)
+    assert ExampleHiddenStatesConnector._find_cache_kv_group_id(cfg) == 1
+
+    vllm_config = SimpleNamespace(cache_config=SimpleNamespace(block_size=528))
+    assert ExampleHiddenStatesConnector._get_cache_block_size(
+        vllm_config, cfg, cache_kv_group_id=1
+    ) == 22
+
+
+def test_dsv4_groups_isolate_hidden_states_for_connector():
+    # HiddenStateCacheSpec is an MLAAttentionSpec subclass. The DSV4 grouping
+    # path must nevertheless leave it in a hidden-only UniformType group.
     dt = torch.bfloat16
     spec = {
         "layers.0.mla": MLAAttentionSpec(
             block_size=64, num_kv_heads=1, head_size=576, dtype=dt
         ),
         "layers.1.swa": SlidingWindowMLASpec(
-            block_size=64, num_kv_heads=1, head_size=576, dtype=dt, sliding_window=512
+            block_size=64,
+            num_kv_heads=1,
+            head_size=1024,
+            dtype=dt,
+            sliding_window=512,
         ),
         "cache_only_layers.61": _hidden(64),
     }
@@ -120,7 +143,22 @@ def test_find_group_id_errors_clearly_when_absorbed_by_mla_swa_verifier():
         speculative_config=None,
     )
     groups = get_kv_cache_groups(vllm_config, spec)
-    assert not any(isinstance(g.kv_cache_spec, HiddenStateCacheSpec) for g in groups)
+    hidden_groups = [
+        group
+        for group in groups
+        if "cache_only_layers.61" in group.layer_names
+    ]
+    assert len(hidden_groups) == 1
+    assert isinstance(hidden_groups[0].kv_cache_spec, UniformTypeKVCacheSpecs)
+    assert all(
+        isinstance(inner, HiddenStateCacheSpec)
+        for inner in hidden_groups[0].kv_cache_spec.kv_cache_specs.values()
+    )
     cfg = SimpleNamespace(kv_cache_groups=groups)
-    with pytest.raises(ValueError, match="MLA verifiers are unsupported"):
+    hidden_group_id = next(
+        gid for gid, group in enumerate(groups) if group is hidden_groups[0]
+    )
+    assert (
         ExampleHiddenStatesConnector._find_cache_kv_group_id(cfg)
+        == hidden_group_id
+    )

vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py

@@ -118,13 +118,26 @@ def _find_cache_kv_group_id(cls, kv_cache_config: "KVCacheConfig | None") -> int
         if kv_cache_config is None:
             return 0
 
-        from vllm.v1.kv_cache_interface import HiddenStateCacheSpec
+        from vllm.v1.kv_cache_interface import (
+            HiddenStateCacheSpec,
+            UniformTypeKVCacheSpecs,
+        )
+
+        def _is_hidden_spec(spec) -> bool:
+            if isinstance(spec, HiddenStateCacheSpec):
+                return True
+            if isinstance(spec, UniformTypeKVCacheSpecs) and spec.kv_cache_specs:
+                return all(
+                    isinstance(inner, HiddenStateCacheSpec)
+                    for inner in spec.kv_cache_specs.values()
+                )
+            return False
 
         groups = kv_cache_config.kv_cache_groups
         group_ids = [
             gid
             for gid, group in enumerate(groups)
-            if isinstance(group.kv_cache_spec, HiddenStateCacheSpec)
+            if _is_hidden_spec(group.kv_cache_spec)
         ]
         if len(group_ids) == 1:
             return group_ids[0]
@@ -133,7 +146,7 @@ def _find_cache_kv_group_id(cls, kv_cache_config: "KVCacheConfig | None") -> int
         raise ValueError(
             "Could not uniquely identify the extract-hidden-states KV cache "
             f"group among {len(groups)} groups; the hidden-states layer must be "
-            "isolated in its own group (MLA verifiers are unsupported)."
+            "isolated in exactly one KV cache group."
         )
 
     @staticmethod

vllm/v1/core/kv_cache_utils.py

@@ -1605,13 +1605,18 @@ def group_and_unify_kv_cache_specs(
     Group the KV cache specs and unify each group into one UniformTypeKVCacheSpecs.
     Currently, this is only used for DeepseekV4.
     """
+    non_hidden_specs = {
+        name: spec
+        for name, spec in kv_cache_spec.items()
+        if not isinstance(spec, HiddenStateCacheSpec)
+    }
     if not any(
-        isinstance(spec, SlidingWindowMLASpec) for spec in kv_cache_spec.values()
+        isinstance(spec, SlidingWindowMLASpec) for spec in non_hidden_specs.values()
     ):
         return None
 
     # SlidingWindowMLASpec models with uniform page sizes don't need tuple packing.
-    page_sizes = {spec.page_size_bytes for spec in kv_cache_spec.values()}
+    page_sizes = {spec.page_size_bytes for spec in non_hidden_specs.values()}
     if len(page_sizes) <= 1:
         return None
 
@@ -1622,7 +1627,7 @@ def group_and_unify_kv_cache_specs(
     # NOTE: Here we group SWA layers by (block_size, sliding_window), which separates
     # SWA layers, C4I+C4A layers, and C128A layers into three different groups. It can
     # be fragile with only block_size and sliding_window as keys, but fine for now.
-    for name, spec in kv_cache_spec.items():
+    for name, spec in non_hidden_specs.items():
         if isinstance(spec, SlidingWindowMLASpec):
             grouped_swa_mla_specs[(spec.block_size, spec.sliding_window)][name] = spec
         elif isinstance(spec, MLAAttentionSpec):
@@ -1832,6 +1837,21 @@ def get_kv_cache_groups(
         # UniformTypeKVCacheSpecs.
         kv_cache_groups = _get_kv_cache_groups_uniform_groups(grouped_specs)
         _annotate_eagle_groups_deepseek_v4(vllm_config, kv_cache_spec, kv_cache_groups)
+        hidden_specs = {
+            k: v
+            for k, v in kv_cache_spec.items()
+            if isinstance(v, HiddenStateCacheSpec)
+        }
+        if hidden_specs:
+            # Keep UniformTypeKVCacheSpecs so DSV4 packed-layout detection
+            # still sees an all-uniform group list.
+            hidden_uniform = UniformTypeKVCacheSpecs.from_specs(hidden_specs)
+            assert hidden_uniform is not None, (
+                "HiddenStateCacheSpec layers must form a uniform KV group"
+            )
+            kv_cache_groups.append(
+                KVCacheGroupSpec(list(hidden_specs), hidden_uniform)
+            )
         return kv_cache_groups
 
     # Pull HiddenStateCacheSpec layers out before the general multi-group