[Bugfix][R3] Unwrap UniformTypeKVCacheSpecs when selecting the routed-experts KV group (#53240)
Signed-off-by: Hollow Man <hollowman@opensuse.org>
ℍ𝕠𝕝𝕝𝕠𝕨 𝕄𝕒𝕟 committed
Aug 22, 2026 at 05:13 UTC
b2db227a7c4c5e55f85524a094b607e5d27408b4
5 files changed
+119
-23
tests/model_executor/test_routed_experts_capture.py
+46
-7
@@ -25,6 +25,9 @@ from vllm.v1.kv_cache_interface import (
25
FullAttentionSpec,
26
KVCacheConfig,
27
KVCacheGroupSpec,
28
+ MLAAttentionSpec,
29
+ SlidingWindowMLASpec,
30
+ UniformTypeKVCacheSpecs,
31
)
32
33
pytestmark = pytest.mark.cpu_test
@@ -321,15 +324,17 @@ def test_routed_experts_capturer_dp_unexpected_batch_raises():
324
assert capturer.device_buffer[0, 0, 0].item() == -1
325
326
324
-def test_routed_experts_attention_group_is_shared_and_fail_closed(monkeypatch):
325
- class FullAttentionSpec:
326
- pass
327
-
328
- monkeypatch.setattr(f"{_REC_MODULE}.FullAttentionSpec", FullAttentionSpec)
327
+def test_routed_experts_attention_group_is_shared_and_fail_closed():
328
+ """Both sides key routing data by this gid, so it must skip non-full-attention
329
+ groups rather than defaulting to 0, and fail closed when none exists."""
330
+ common = dict(num_kv_heads=1, head_size=1, dtype=torch.float32)
331
config = SimpleNamespace(
332
kv_cache_groups=[
331
- SimpleNamespace(kv_cache_spec=object()),
332
- SimpleNamespace(kv_cache_spec=FullAttentionSpec()),
333
+ KVCacheGroupSpec(
334
+ ["swa_layer"],
335
+ SlidingWindowMLASpec(block_size=4, sliding_window=8, **common),
336
+ ),
337
+ KVCacheGroupSpec(["full_layer"], FullAttentionSpec(block_size=4, **common)),
338
]
339
)
340
assert get_routed_experts_attn_gid(config) == 1
@@ -338,6 +343,40 @@ def test_routed_experts_attention_group_is_shared_and_fail_closed(monkeypatch):
343
get_routed_experts_attn_gid(SimpleNamespace(kv_cache_groups=[]))
344
345
346
+def test_routed_experts_attention_group_unwraps_uniform_type_specs():
347
+ """DeepSeek-V4-shaped groups wrap their specs in ``UniformTypeKVCacheSpecs``.
348
+
349
+ The wrapper is not a ``FullAttentionSpec``, so a bare isinstance check finds
350
+ no group and fails closed on every worker. Unwrap semantics themselves are
351
+ covered by ``test_is_full_attention_spec_*`` in tests/v1/core.
352
+ """
353
+ common = dict(num_kv_heads=1, head_size=1, dtype=torch.float32)
354
+ swa_spec = SlidingWindowMLASpec(block_size=4, sliding_window=8, **common)
355
+ mla_spec = MLAAttentionSpec(block_size=4, **common)
356
+ config = SimpleNamespace(
357
+ kv_cache_groups=[
358
+ KVCacheGroupSpec(
359
+ ["swa_layer"],
360
+ UniformTypeKVCacheSpecs(
361
+ block_size=4, kv_cache_specs={"swa_layer": swa_spec}
362
+ ),
363
+ ),
364
+ KVCacheGroupSpec(
365
+ ["mla_layer"],
366
+ UniformTypeKVCacheSpecs(
367
+ block_size=4, kv_cache_specs={"mla_layer": mla_spec}
368
+ ),
369
+ ),
370
+ ]
371
+ )
372
+
373
+ assert get_routed_experts_attn_gid(config) == 1
374
+
375
+ swa_only = SimpleNamespace(kv_cache_groups=[config.kv_cache_groups[0]])
376
+ with pytest.raises(ValueError, match="requires a full-attention KV cache group"):
377
+ get_routed_experts_attn_gid(swa_only)
378
+
379
+
380
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
381
def test_mrv2_async_output_returns_existing_routed_experts_field():
382
from vllm.v1.outputs import ModelRunnerOutput, RoutedExpertsTensors
tests/v1/core/test_kv_cache_utils.py
+39
@@ -61,6 +61,8 @@ from vllm.v1.kv_cache_interface import (
61
UniformTypeKVCacheSpecs,
62
get_kv_cache_spec_kind,
63
get_kv_cache_spec_sliding_window,
64
+ is_full_attention_spec,
65
+ iter_layer_specs,
66
)
67
from vllm.v1.metrics.stats import CachingMetrics, PrefixCacheStats
68
from vllm.v1.request import Request
@@ -3179,3 +3181,40 @@ def test_check_enough_kv_cache_memory_reserves_null_block():
3181
check_enough_kv_cache_memory(
3182
vllm_config, {"layer1": spec}, spec.page_size_bytes * 33
3183
)
3184
+
3185
+
3186
+def test_is_full_attention_spec_unwraps_uniform_type_specs():
3187
+ """``UniformTypeKVCacheSpecs`` is not a ``FullAttentionSpec``, so callers
3188
+ scanning groups with a bare isinstance miss DeepSeek-V4-shaped configs."""
3189
+ common = dict(num_kv_heads=1, head_size=1, dtype=torch.float32)
3190
+ full = FullAttentionSpec(block_size=4, **common)
3191
+ mla = MLAAttentionSpec(block_size=4, **common)
3192
+ swa = SlidingWindowMLASpec(block_size=4, sliding_window=8, **common)
3193
+
3194
+ def wrap(**specs):
3195
+ return UniformTypeKVCacheSpecs(block_size=4, kv_cache_specs=specs)
3196
+
3197
+ assert is_full_attention_spec(full)
3198
+ assert is_full_attention_spec(mla)
3199
+ assert not is_full_attention_spec(swa)
3200
+
3201
+ assert is_full_attention_spec(wrap(a=mla, b=mla))
3202
+ assert not is_full_attention_spec(wrap(a=swa, b=swa))
3203
+
3204
+ # Every layer must qualify: a group holding a recycling sliding-window layer
3205
+ # has no stable slot layout. Grouping forbids mixing today, but the wrapper
3206
+ # is also built directly (e.g. projecting groups onto a PP worker).
3207
+ assert not is_full_attention_spec(wrap(a=mla, b=swa))
3208
+ assert not is_full_attention_spec(wrap())
3209
+
3210
+
3211
+def test_iter_layer_specs_returns_group_members():
3212
+ common = dict(num_kv_heads=1, head_size=1, dtype=torch.float32)
3213
+ full = FullAttentionSpec(block_size=4, **common)
3214
+ mla = MLAAttentionSpec(block_size=4, **common)
3215
+
3216
+ assert list(iter_layer_specs(full)) == [full]
3217
+ wrapped = UniformTypeKVCacheSpecs(
3218
+ block_size=4, kv_cache_specs={"a": full, "b": mla}
3219
+ )
3220
+ assert list(iter_layer_specs(wrapped)) == [full, mla]
vllm/model_executor/layers/fused_moe/routed_experts_capturer.py
+2
-2
@@ -17,7 +17,7 @@ from vllm.config import VllmConfig
17
from vllm.distributed.parallel_state import get_tp_group
18
from vllm.forward_context import get_forward_context
19
from vllm.platforms import current_platform
20
-from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig
20
+from vllm.v1.kv_cache_interface import KVCacheConfig, is_full_attention_spec
21
from vllm.v1.outputs import RoutedExpertsTensors
22
23
logger = logging.getLogger(__name__)
@@ -307,7 +307,7 @@ def bind_routed_experts_capturer(
307
def get_routed_experts_attn_gid(kv_cache_config: KVCacheConfig) -> int:
308
"""Return the full-attention KV cache group used for routed experts."""
309
for gid, group in enumerate(kv_cache_config.kv_cache_groups):
310
- if isinstance(group.kv_cache_spec, FullAttentionSpec):
310
+ if is_full_attention_spec(group.kv_cache_spec):
311
return gid
312
raise ValueError("Routed-experts capture requires a full-attention KV cache group.")
313
vllm/v1/attention/backends/flashinfer.py
+2
-7
@@ -85,7 +85,7 @@ from vllm.v1.kv_cache_interface import (
85
KVCacheLayout,
86
KVCacheSpec,
87
KVQuantMode,
88
- UniformTypeKVCacheSpecs,
88
+ iter_layer_specs,
89
)
90
from vllm.v1.utils import CpuGpuBuffer
91
@@ -981,12 +981,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
981
if is_sm12x and vllm_config.parallel_config.decode_context_parallel_size > 1:
982
return AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
983
984
- # For UniformTypeKVCacheSpecs, check all contained specs
985
- kv_specs = (
986
- kv_cache_spec.kv_cache_specs.values()
987
- if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs)
988
- else [kv_cache_spec]
989
- )
984
+ kv_specs = iter_layer_specs(kv_cache_spec)
985
num_qo_heads = vllm_config.model_config.get_num_attention_heads(
986
vllm_config.parallel_config
987
)
vllm/v1/kv_cache_interface.py
+30
-7
@@ -1020,6 +1020,35 @@ class UniformTypeKVCacheSpecs(KVCacheSpec):
1020
)
1021
1022
1023
+def iter_layer_specs(kv_cache_spec: KVCacheSpec) -> Collection[KVCacheSpec]:
1024
+ """The per-layer specs a KV cache group spec covers.
1025
+
1026
+ ``UniformTypeKVCacheSpecs`` groups keep one spec per layer; every other
1027
+ spec describes its group on its own. Returns the layer specs either way so
1028
+ callers do not have to special-case the wrapper.
1029
+ """
1030
+ if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs):
1031
+ return kv_cache_spec.kv_cache_specs.values()
1032
+ return (kv_cache_spec,)
1033
+
1034
+
1035
+def is_full_attention_spec(kv_cache_spec: KVCacheSpec) -> bool:
1036
+ """Whether a KV cache group spec is (or wraps) full attention.
1037
+
1038
+ ``UniformTypeKVCacheSpecs`` is not itself a ``FullAttentionSpec``, so a bare
1039
+ isinstance check misses groups that carry the wrapper -- DeepSeek-V4's MLA
1040
+ layers, or any model taking the ``UniformTypeKVCacheSpecs.from_specs`` path.
1041
+
1042
+ Every layer must be full attention: a group holding a recycling
1043
+ (sliding-window) layer has no stable slot layout, so callers that key data
1044
+ by slot cannot use it.
1045
+ """
1046
+ layer_specs = iter_layer_specs(kv_cache_spec)
1047
+ return len(layer_specs) > 0 and all(
1048
+ isinstance(spec, FullAttentionSpec) for spec in layer_specs
1049
+ )
1050
+
1051
+
1052
def get_kv_cache_spec_kind(kv_cache_spec: KVCacheSpec) -> KVCacheSpecKind:
1053
if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs):
1054
inner_kinds = {
@@ -1144,15 +1173,9 @@ class KVCacheConfig:
1173
"""Whether attention groups store their KV cache at more than one precision."""
1174
kv_cache_precisions: set[tuple[torch.dtype, KVQuantMode]] = set()
1175
for group in self.kv_cache_groups:
1147
- group_spec = group.kv_cache_spec
1148
- group_specs = (
1149
- list(group_spec.kv_cache_specs.values())
1150
- if isinstance(group_spec, UniformTypeKVCacheSpecs)
1151
- else [group_spec]
1152
- )
1176
kv_cache_precisions.update(
1177
(spec.dtype, spec.kv_quant_mode)
1155
- for spec in group_specs
1178
+ for spec in iter_layer_specs(group.kv_cache_spec)
1179
if isinstance(spec, AttentionSpec)
1180
)
1181
return len(kv_cache_precisions) > 1