[Bugfix][Spec Decode] Use group geometry for FlashAttention metadata (#53002)

Signed-off-by: mgoin <mgoin64@gmail.com>

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

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.

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}."

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