From 9ff7041b55555976359fe8a67eef59f3a055a822 Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Fri, 21 Aug 2026 20:07:19 -0400 Subject: [PATCH] [Bugfix][Spec Decode] Use group geometry for FlashAttention metadata (#53002) Signed-off-by: mgoin --- tests/v1/attention/test_group_head_counts.py | 59 +++++++++++++++++-- vllm/v1/attention/backends/flash_attn.py | 11 ++-- vllm/v1/attention/backends/utils.py | 6 +- .../gpu/spec_decode/dflash/speculator.py | 12 ++-- 4 files changed, 70 insertions(+), 18 deletions(-) diff --git a/tests/v1/attention/test_group_head_counts.py b/tests/v1/attention/test_group_head_counts.py index b799bccc29..9691926819 100644 --- a/tests/v1/attention/test_group_head_counts.py +++ b/tests/v1/attention/test_group_head_counts.py @@ -1,10 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Scheduler metadata sizes a scratchpad from the query head count, so it must -come from the builder's own group: the model-wide ``get_num_attention_heads()`` -is wrong for models that vary it per layer (e.g. Laguna), and too small a -scratchpad is indexed past its end. -""" +"""Attention metadata geometry must come from the builder's own group.""" from types import SimpleNamespace from unittest.mock import MagicMock, patch @@ -17,8 +13,9 @@ from vllm.v1.attention.backends.cpu_attn import ( CPUAttentionBackendImpl, CPUAttentionMetadataBuilder, ) +from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadataBuilder -pytestmark = pytest.mark.skipif( +requires_cpu = pytest.mark.skipif( not current_platform.is_cpu(), reason="CPU attention backend" ) @@ -69,6 +66,7 @@ def _build(layer_num_heads: list[int]) -> CPUAttentionMetadataBuilder: ) +@requires_cpu @pytest.mark.parametrize("group_num_heads", [MODEL_WIDE_NUM_HEADS, 64, 16]) def test_num_heads_comes_from_the_group(group_num_heads): """The group's own count wins, even when it is not the model-wide one.""" @@ -76,7 +74,56 @@ def test_num_heads_comes_from_the_group(group_num_heads): assert builder.num_heads == group_num_heads +@requires_cpu def test_mixed_head_counts_in_one_group_are_rejected(): """Grouping guarantees uniformity; a mixed group means that broke.""" with pytest.raises(AssertionError, match="share num_heads"): _build([MODEL_WIDE_NUM_HEADS, 64]) + + +def test_flash_attention_geometry_comes_from_the_group(): + """FA3 must use group geometry for layers without an ``impl`` wrapper.""" + layers = {f"layer_{i}": SimpleNamespace(num_heads=16) for i in range(2)} + vllm_config = MagicMock() + vllm_config.model_config.get_num_attention_heads.return_value = MODEL_WIDE_NUM_HEADS + vllm_config.model_config.get_num_kv_heads.return_value = NUM_KV_HEADS + vllm_config.model_config.get_head_size.return_value = 128 + vllm_config.model_config.rswa_window = None + vllm_config.model_config.is_mm_prefix_lm = False + vllm_config.parallel_config.cp_kv_cache_interleave_size = 1 + vllm_config.compilation_config.cudagraph_mode.has_full_cudagraphs.return_value = ( + False + ) + vllm_config.compilation_config.max_cudagraph_capture_size = None + kv_cache_spec = SimpleNamespace( + block_size=16, + num_kv_heads=2, + head_size=64, + dtype=torch.bfloat16, + ) + + with ( + patch( + "vllm.v1.attention.backends.utils.get_layers_from_vllm_config", + return_value=layers, + ), + patch( + "vllm.distributed.parallel_state.get_dcp_group", + side_effect=AssertionError, + ), + patch( + "vllm.v1.attention.backends.flash_attn.get_flash_attn_version", + return_value=3, + ), + ): + builder = FlashAttentionMetadataBuilder( + kv_cache_spec=kv_cache_spec, + layer_names=list(layers), + vllm_config=vllm_config, + device=torch.device("cpu"), + ) + + assert builder.aot_schedule + assert builder.num_heads_q == 16 + assert builder.num_heads_kv == 2 + assert builder.headdim == 64 diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index 7ccbddf124..76aaed4c5b 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -33,6 +33,7 @@ from vllm.v1.attention.backends.fa_utils import ( from vllm.v1.attention.backends.utils import ( fill_mm_prefix_query_ranges, get_dcp_local_seq_lens, + get_num_attention_heads_from_layers, ) from vllm.v1.attention.ops.dcp import ( cp_lse_ag_out_rs, @@ -435,12 +436,12 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad self.compilation_config = vllm_config.compilation_config self.attention_config = vllm_config.attention_config - self.num_heads_q = self.model_config.get_num_attention_heads( - self.parallel_config - ) - self.num_heads_kv = self.model_config.get_num_kv_heads(self.parallel_config) + self.num_heads_q = get_num_attention_heads_from_layers( + vllm_config, layer_names + ) or self.model_config.get_num_attention_heads(self.parallel_config) + self.num_heads_kv = kv_cache_spec.num_kv_heads self.kv_cache_dtype = kv_cache_spec.dtype - self.headdim = self.model_config.get_head_size() + self.headdim = kv_cache_spec.head_size self.block_size = kv_cache_spec.block_size self.max_num_splits = 0 # No upper bound on the number of splits. diff --git a/vllm/v1/attention/backends/utils.py b/vllm/v1/attention/backends/utils.py index f3378806c8..9a5224238a 100644 --- a/vllm/v1/attention/backends/utils.py +++ b/vllm/v1/attention/backends/utils.py @@ -9,6 +9,7 @@ from typing import ( Any, Literal, Protocol, + cast, get_args, ) @@ -263,7 +264,10 @@ def get_num_attention_heads_from_layers( ) if not attn_layers: return None - heads = {layer.impl.num_heads for layer in attn_layers.values()} + heads = { + cast(Any, getattr(layer, "impl", layer)).num_heads + for layer in attn_layers.values() + } assert len(heads) == 1, ( f"All layers in one attention group must share num_heads; " f"got {heads} for {layer_names}." diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index 92837afd68..684a3d3a44 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import copy from collections.abc import Mapping from typing import Any @@ -103,13 +104,12 @@ class DFlashSpeculator(DraftModelSpeculator): @property def attn_vllm_config(self) -> VllmConfig: # The draft's attention differs from the target's in causality. - return replace( - self.vllm_config, - attention_config=replace( - self.vllm_config.attention_config, - use_non_causal=self.requires_non_causal, - ), + config = copy.copy(super().attn_vllm_config) + config.attention_config = replace( + self.vllm_config.attention_config, + use_non_causal=self.requires_non_causal, ) + return config def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None: wants_full = cudagraph_mode.decode_mode() == CUDAGraphMode.FULL