diff --git a/tests/model_executor/test_routed_experts_capture.py b/tests/model_executor/test_routed_experts_capture.py index 21b7ba4239..b33b62cb4a 100644 --- a/tests/model_executor/test_routed_experts_capture.py +++ b/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 diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 870c139ae0..39f18c6d4f 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/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] diff --git a/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py b/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py index fa57ea383f..edd3e3beaa 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py +++ b/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.") diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index 792056f61d..c9aea6cb83 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/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 ) diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index efc284bcbb..b660f49317 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/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