[Bugfix][Spec Decode] Reapply group geometry for FlashAttention metadata (#53336)
Signed-off-by: mgoin <mgoin64@gmail.com>
Misha Goin committed
Aug 24, 2026 at 11:08 UTC
f620499ee3fe18131d71b02e1e8e5f1cf984cf1c
4 files changed
+70
-18
tests/v1/attention/test_group_head_counts.py
+53
-6
@@ -1,10 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Scheduler metadata sizes a scratchpad from the query head count, so it must
4
-come from the builder's own group: the model-wide ``get_num_attention_heads()``
5
-is wrong for models that vary it per layer (e.g. Laguna), and too small a
6
-scratchpad is indexed past its end.
7
-"""
3
+"""Attention metadata geometry must come from the builder's own group."""
4
5
from types import SimpleNamespace
6
from unittest.mock import MagicMock, patch
@@ -17,8 +13,9 @@ from vllm.v1.attention.backends.cpu_attn import (
13
CPUAttentionBackendImpl,
14
CPUAttentionMetadataBuilder,
15
)
16
+from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadataBuilder
17
21
-pytestmark = pytest.mark.skipif(
18
+requires_cpu = pytest.mark.skipif(
19
not current_platform.is_cpu(), reason="CPU attention backend"
20
)
21
@@ -69,6 +66,7 @@ def _build(layer_num_heads: list[int]) -> CPUAttentionMetadataBuilder:
66
)
67
68
69
+@requires_cpu
70
@pytest.mark.parametrize("group_num_heads", [MODEL_WIDE_NUM_HEADS, 64, 16])
71
def test_num_heads_comes_from_the_group(group_num_heads):
72
"""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):
74
assert builder.num_heads == group_num_heads
75
76
77
+@requires_cpu
78
def test_mixed_head_counts_in_one_group_are_rejected():
79
"""Grouping guarantees uniformity; a mixed group means that broke."""
80
with pytest.raises(AssertionError, match="share num_heads"):
81
_build([MODEL_WIDE_NUM_HEADS, 64])
82
+
83
+
84
+def test_flash_attention_geometry_comes_from_the_group():
85
+ """FA3 must use group geometry for layers without an ``impl`` wrapper."""
86
+ layers = {f"layer_{i}": SimpleNamespace(num_heads=16) for i in range(2)}
87
+ vllm_config = MagicMock()
88
+ vllm_config.model_config.get_num_attention_heads.return_value = MODEL_WIDE_NUM_HEADS
89
+ vllm_config.model_config.get_num_kv_heads.return_value = NUM_KV_HEADS
90
+ vllm_config.model_config.get_head_size.return_value = 128
91
+ vllm_config.model_config.rswa_window = None
92
+ vllm_config.model_config.is_mm_prefix_lm = False
93
+ vllm_config.parallel_config.cp_kv_cache_interleave_size = 1
94
+ vllm_config.compilation_config.cudagraph_mode.has_full_cudagraphs.return_value = (
95
+ False
96
+ )
97
+ vllm_config.compilation_config.max_cudagraph_capture_size = None
98
+ kv_cache_spec = SimpleNamespace(
99
+ block_size=16,
100
+ num_kv_heads=2,
101
+ head_size=64,
102
+ dtype=torch.bfloat16,
103
+ )
104
+
105
+ with (
106
+ patch(
107
+ "vllm.v1.attention.backends.utils.get_layers_from_vllm_config",
108
+ return_value=layers,
109
+ ),
110
+ patch(
111
+ "vllm.distributed.parallel_state.get_dcp_group",
112
+ side_effect=AssertionError,
113
+ ),
114
+ patch(
115
+ "vllm.v1.attention.backends.flash_attn.get_flash_attn_version",
116
+ return_value=3,
117
+ ),
118
+ ):
119
+ builder = FlashAttentionMetadataBuilder(
120
+ kv_cache_spec=kv_cache_spec,
121
+ layer_names=list(layers),
122
+ vllm_config=vllm_config,
123
+ device=torch.device("cpu"),
124
+ )
125
+
126
+ assert builder.aot_schedule
127
+ assert builder.num_heads_q == 16
128
+ assert builder.num_heads_kv == 2
129
+ assert builder.headdim == 64
vllm/v1/attention/backends/flash_attn.py
+6
-5
@@ -33,6 +33,7 @@ from vllm.v1.attention.backends.fa_utils import (
33
from vllm.v1.attention.backends.utils import (
34
fill_mm_prefix_query_ranges,
35
get_dcp_local_seq_lens,
36
+ get_num_attention_heads_from_layers,
37
)
38
from vllm.v1.attention.ops.dcp import (
39
cp_lse_ag_out_rs,
@@ -397,12 +398,12 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
398
self.compilation_config = vllm_config.compilation_config
399
self.attention_config = vllm_config.attention_config
400
400
- self.num_heads_q = self.model_config.get_num_attention_heads(
401
- self.parallel_config
402
- )
403
- self.num_heads_kv = self.model_config.get_num_kv_heads(self.parallel_config)
401
+ self.num_heads_q = get_num_attention_heads_from_layers(
402
+ vllm_config, layer_names
403
+ ) or self.model_config.get_num_attention_heads(self.parallel_config)
404
+ self.num_heads_kv = kv_cache_spec.num_kv_heads
405
self.kv_cache_dtype = kv_cache_spec.dtype
405
- self.headdim = self.model_config.get_head_size()
406
+ self.headdim = kv_cache_spec.head_size
407
self.block_size = kv_cache_spec.block_size
408
409
self.max_num_splits = 0 # No upper bound on the number of splits.
vllm/v1/attention/backends/utils.py
+5
-1
@@ -8,6 +8,7 @@ from typing import (
8
TYPE_CHECKING,
9
Any,
10
Protocol,
11
+ cast,
12
)
13
14
import numpy as np
@@ -375,7 +376,10 @@ def get_num_attention_heads_from_layers(
376
)
377
if not attn_layers:
378
return None
378
- heads = {layer.impl.num_heads for layer in attn_layers.values()}
379
+ heads = {
380
+ cast(Any, getattr(layer, "impl", layer)).num_heads
381
+ for layer in attn_layers.values()
382
+ }
383
assert len(heads) == 1, (
384
f"All layers in one attention group must share num_heads; "
385
f"got {heads} for {layer_names}."
vllm/v1/worker/gpu/spec_decode/dflash/speculator.py
+6
-6
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import copy
4
from collections.abc import Mapping
5
from typing import Any
6
@@ -103,13 +104,12 @@ class DFlashSpeculator(DraftModelSpeculator):
104
@property
105
def attn_vllm_config(self) -> VllmConfig:
106
# The draft's attention differs from the target's in causality.
106
- return replace(
107
- self.vllm_config,
108
- attention_config=replace(
109
- self.vllm_config.attention_config,
110
- use_non_causal=self.requires_non_causal,
111
- ),
107
+ config = copy.copy(super().attn_vllm_config)
108
+ config.attention_config = replace(
109
+ self.vllm_config.attention_config,
110
+ use_non_causal=self.requires_non_causal,
111
)
112
+ return config
113
114
def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None:
115
wants_full = cudagraph_mode.decode_mode() == CUDAGraphMode.FULL