[Bugfix][R3] Unwrap UniformTypeKVCacheSpecs when selecting the routed-experts KV group (#53240)

Signed-off-by: Hollow Man <hollowman@opensuse.org>

tests/model_executor/test_routed_experts_capture.py

@@ -25,6 +25,9 @@ from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
     KVCacheConfig,
     KVCacheGroupSpec,
+    MLAAttentionSpec,
+    SlidingWindowMLASpec,
+    UniformTypeKVCacheSpecs,
 )
 
 pytestmark = pytest.mark.cpu_test
@@ -321,15 +324,17 @@ def test_routed_experts_capturer_dp_unexpected_batch_raises():
     assert capturer.device_buffer[0, 0, 0].item() == -1
 
 
-def test_routed_experts_attention_group_is_shared_and_fail_closed(monkeypatch):
-    class FullAttentionSpec:
-        pass
-
-    monkeypatch.setattr(f"{_REC_MODULE}.FullAttentionSpec", FullAttentionSpec)
+def test_routed_experts_attention_group_is_shared_and_fail_closed():
+    """Both sides key routing data by this gid, so it must skip non-full-attention
+    groups rather than defaulting to 0, and fail closed when none exists."""
+    common = dict(num_kv_heads=1, head_size=1, dtype=torch.float32)
     config = SimpleNamespace(
         kv_cache_groups=[
-            SimpleNamespace(kv_cache_spec=object()),
-            SimpleNamespace(kv_cache_spec=FullAttentionSpec()),
+            KVCacheGroupSpec(
+                ["swa_layer"],
+                SlidingWindowMLASpec(block_size=4, sliding_window=8, **common),
+            ),
+            KVCacheGroupSpec(["full_layer"], FullAttentionSpec(block_size=4, **common)),
         ]
     )
     assert get_routed_experts_attn_gid(config) == 1
@@ -338,6 +343,40 @@ def test_routed_experts_attention_group_is_shared_and_fail_closed(monkeypatch):
         get_routed_experts_attn_gid(SimpleNamespace(kv_cache_groups=[]))
 
 
+def test_routed_experts_attention_group_unwraps_uniform_type_specs():
+    """DeepSeek-V4-shaped groups wrap their specs in ``UniformTypeKVCacheSpecs``.
+
+    The wrapper is not a ``FullAttentionSpec``, so a bare isinstance check finds
+    no group and fails closed on every worker. Unwrap semantics themselves are
+    covered by ``test_is_full_attention_spec_*`` in tests/v1/core.
+    """
+    common = dict(num_kv_heads=1, head_size=1, dtype=torch.float32)
+    swa_spec = SlidingWindowMLASpec(block_size=4, sliding_window=8, **common)
+    mla_spec = MLAAttentionSpec(block_size=4, **common)
+    config = SimpleNamespace(
+        kv_cache_groups=[
+            KVCacheGroupSpec(
+                ["swa_layer"],
+                UniformTypeKVCacheSpecs(
+                    block_size=4, kv_cache_specs={"swa_layer": swa_spec}
+                ),
+            ),
+            KVCacheGroupSpec(
+                ["mla_layer"],
+                UniformTypeKVCacheSpecs(
+                    block_size=4, kv_cache_specs={"mla_layer": mla_spec}
+                ),
+            ),
+        ]
+    )
+
+    assert get_routed_experts_attn_gid(config) == 1
+
+    swa_only = SimpleNamespace(kv_cache_groups=[config.kv_cache_groups[0]])
+    with pytest.raises(ValueError, match="requires a full-attention KV cache group"):
+        get_routed_experts_attn_gid(swa_only)
+
+
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
 def test_mrv2_async_output_returns_existing_routed_experts_field():
     from vllm.v1.outputs import ModelRunnerOutput, RoutedExpertsTensors

tests/v1/core/test_kv_cache_utils.py

@@ -61,6 +61,8 @@ from vllm.v1.kv_cache_interface import (
     UniformTypeKVCacheSpecs,
     get_kv_cache_spec_kind,
     get_kv_cache_spec_sliding_window,
+    is_full_attention_spec,
+    iter_layer_specs,
 )
 from vllm.v1.metrics.stats import CachingMetrics, PrefixCacheStats
 from vllm.v1.request import Request
@@ -3179,3 +3181,40 @@ def test_check_enough_kv_cache_memory_reserves_null_block():
     check_enough_kv_cache_memory(
         vllm_config, {"layer1": spec}, spec.page_size_bytes * 33
     )
+
+
+def test_is_full_attention_spec_unwraps_uniform_type_specs():
+    """``UniformTypeKVCacheSpecs`` is not a ``FullAttentionSpec``, so callers
+    scanning groups with a bare isinstance miss DeepSeek-V4-shaped configs."""
+    common = dict(num_kv_heads=1, head_size=1, dtype=torch.float32)
+    full = FullAttentionSpec(block_size=4, **common)
+    mla = MLAAttentionSpec(block_size=4, **common)
+    swa = SlidingWindowMLASpec(block_size=4, sliding_window=8, **common)
+
+    def wrap(**specs):
+        return UniformTypeKVCacheSpecs(block_size=4, kv_cache_specs=specs)
+
+    assert is_full_attention_spec(full)
+    assert is_full_attention_spec(mla)
+    assert not is_full_attention_spec(swa)
+
+    assert is_full_attention_spec(wrap(a=mla, b=mla))
+    assert not is_full_attention_spec(wrap(a=swa, b=swa))
+
+    # Every layer must qualify: a group holding a recycling sliding-window layer
+    # has no stable slot layout. Grouping forbids mixing today, but the wrapper
+    # is also built directly (e.g. projecting groups onto a PP worker).
+    assert not is_full_attention_spec(wrap(a=mla, b=swa))
+    assert not is_full_attention_spec(wrap())
+
+
+def test_iter_layer_specs_returns_group_members():
+    common = dict(num_kv_heads=1, head_size=1, dtype=torch.float32)
+    full = FullAttentionSpec(block_size=4, **common)
+    mla = MLAAttentionSpec(block_size=4, **common)
+
+    assert list(iter_layer_specs(full)) == [full]
+    wrapped = UniformTypeKVCacheSpecs(
+        block_size=4, kv_cache_specs={"a": full, "b": mla}
+    )
+    assert list(iter_layer_specs(wrapped)) == [full, mla]

vllm/model_executor/layers/fused_moe/routed_experts_capturer.py

@@ -17,7 +17,7 @@ from vllm.config import VllmConfig
 from vllm.distributed.parallel_state import get_tp_group
 from vllm.forward_context import get_forward_context
 from vllm.platforms import current_platform
-from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig
+from vllm.v1.kv_cache_interface import KVCacheConfig, is_full_attention_spec
 from vllm.v1.outputs import RoutedExpertsTensors
 
 logger = logging.getLogger(__name__)
@@ -307,7 +307,7 @@ def bind_routed_experts_capturer(
 def get_routed_experts_attn_gid(kv_cache_config: KVCacheConfig) -> int:
     """Return the full-attention KV cache group used for routed experts."""
     for gid, group in enumerate(kv_cache_config.kv_cache_groups):
-        if isinstance(group.kv_cache_spec, FullAttentionSpec):
+        if is_full_attention_spec(group.kv_cache_spec):
             return gid
     raise ValueError("Routed-experts capture requires a full-attention KV cache group.")
 

vllm/v1/attention/backends/flashinfer.py

@@ -85,7 +85,7 @@ from vllm.v1.kv_cache_interface import (
     KVCacheLayout,
     KVCacheSpec,
     KVQuantMode,
-    UniformTypeKVCacheSpecs,
+    iter_layer_specs,
 )
 from vllm.v1.utils import CpuGpuBuffer
 
@@ -981,12 +981,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
         if is_sm12x and vllm_config.parallel_config.decode_context_parallel_size > 1:
             return AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
 
-        # For UniformTypeKVCacheSpecs, check all contained specs
-        kv_specs = (
-            kv_cache_spec.kv_cache_specs.values()
-            if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs)
-            else [kv_cache_spec]
-        )
+        kv_specs = iter_layer_specs(kv_cache_spec)
         num_qo_heads = vllm_config.model_config.get_num_attention_heads(
             vllm_config.parallel_config
         )

vllm/v1/kv_cache_interface.py

@@ -1020,6 +1020,35 @@ class UniformTypeKVCacheSpecs(KVCacheSpec):
         )
 
 
+def iter_layer_specs(kv_cache_spec: KVCacheSpec) -> Collection[KVCacheSpec]:
+    """The per-layer specs a KV cache group spec covers.
+
+    ``UniformTypeKVCacheSpecs`` groups keep one spec per layer; every other
+    spec describes its group on its own. Returns the layer specs either way so
+    callers do not have to special-case the wrapper.
+    """
+    if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs):
+        return kv_cache_spec.kv_cache_specs.values()
+    return (kv_cache_spec,)
+
+
+def is_full_attention_spec(kv_cache_spec: KVCacheSpec) -> bool:
+    """Whether a KV cache group spec is (or wraps) full attention.
+
+    ``UniformTypeKVCacheSpecs`` is not itself a ``FullAttentionSpec``, so a bare
+    isinstance check misses groups that carry the wrapper -- DeepSeek-V4's MLA
+    layers, or any model taking the ``UniformTypeKVCacheSpecs.from_specs`` path.
+
+    Every layer must be full attention: a group holding a recycling
+    (sliding-window) layer has no stable slot layout, so callers that key data
+    by slot cannot use it.
+    """
+    layer_specs = iter_layer_specs(kv_cache_spec)
+    return len(layer_specs) > 0 and all(
+        isinstance(spec, FullAttentionSpec) for spec in layer_specs
+    )
+
+
 def get_kv_cache_spec_kind(kv_cache_spec: KVCacheSpec) -> KVCacheSpecKind:
     if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs):
         inner_kinds = {
@@ -1144,15 +1173,9 @@ class KVCacheConfig:
         """Whether attention groups store their KV cache at more than one precision."""
         kv_cache_precisions: set[tuple[torch.dtype, KVQuantMode]] = set()
         for group in self.kv_cache_groups:
-            group_spec = group.kv_cache_spec
-            group_specs = (
-                list(group_spec.kv_cache_specs.values())
-                if isinstance(group_spec, UniformTypeKVCacheSpecs)
-                else [group_spec]
-            )
             kv_cache_precisions.update(
                 (spec.dtype, spec.kv_quant_mode)
-                for spec in group_specs
+                for spec in iter_layer_specs(group.kv_cache_spec)
                 if isinstance(spec, AttentionSpec)
             )
         return len(kv_cache_precisions) > 1