[Bugfix] Restore part of bugfix #42650 after accidental deletion in #43241 (#47039)
Signed-off-by: zhanda <zhandazhu@gmail.com> Signed-off-by: Nikita Shapovalov <nikita@poolside.ai> Co-authored-by: Zhanda Zhu <49645678+zhandaz@users.noreply.github.com> Co-authored-by: Shang Wang <shangw@nvidia.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>
Nikita Shapovalov committed
Jun 30, 2026 at 20:07 UTC
a7732537f4304a36c8067fde0937999fff939b63
3 files changed
+36
-6
vllm/v1/attention/backends/flashinfer.py
+5
-3
@@ -62,6 +62,7 @@ from vllm.v1.attention.backends.utils import (
62
KVCacheLayoutType,
63
get_dcp_local_seq_lens,
64
get_kv_cache_layout,
65
+ get_num_attention_heads_from_layers,
66
get_per_layer_parameters,
67
infer_global_hyperparameters,
68
split_decodes_and_prefills,
@@ -635,9 +636,10 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
636
self.use_dcp and vllm_config.parallel_config.dcp_comm_backend == "a2a"
637
)
638
638
- self.num_qo_heads = self.model_config.get_num_attention_heads(
639
- self.vllm_config.parallel_config
640
- )
639
+ # Compatible with models with non-uniform per-layer head counts.
640
+ self.num_qo_heads = get_num_attention_heads_from_layers(
641
+ vllm_config, layer_names
642
+ ) or self.model_config.get_num_attention_heads(self.vllm_config.parallel_config)
643
644
self.num_kv_heads = self.kv_cache_spec.num_kv_heads
645
self.head_dim = self.kv_cache_spec.head_size
vllm/v1/attention/backends/triton_attn.py
+5
-3
@@ -33,6 +33,7 @@ from vllm.v1.attention.backend import (
33
from vllm.v1.attention.backends.utils import (
34
compute_mm_prefix_range_tensor,
35
get_kv_cache_layout,
36
+ get_num_attention_heads_from_layers,
37
)
38
from vllm.v1.attention.ops.triton_prefill_attention import context_attention_fwd
39
from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
@@ -110,9 +111,10 @@ class TritonAttentionMetadataBuilder(AttentionMetadataBuilder[TritonAttentionMet
111
self.block_size = kv_cache_spec.block_size
112
113
model_config = vllm_config.model_config
113
- self.num_heads_q = model_config.get_num_attention_heads(
114
- vllm_config.parallel_config
115
- )
114
+ # Compatible with models with non-uniform per-layer head counts.
115
+ self.num_heads_q = get_num_attention_heads_from_layers(
116
+ vllm_config, layer_names
117
+ ) or model_config.get_num_attention_heads(vllm_config.parallel_config)
118
self.num_heads_kv = model_config.get_num_kv_heads(vllm_config.parallel_config)
119
self.headdim = model_config.get_head_size()
120
vllm/v1/attention/backends/utils.py
+26
@@ -166,6 +166,32 @@ def get_per_layer_parameters(
166
return per_layer_params
167
168
169
+def get_num_attention_heads_from_layers(
170
+ vllm_config: VllmConfig, layer_names: list[str]
171
+) -> int | None:
172
+ """Per-TP-rank ``num_heads`` shared by the named Attention layers.
173
+
174
+ Use in metadata builders whose plan-time allocations depend on the
175
+ head count: the model-wide ``get_num_attention_heads()`` is wrong
176
+ for models with non-uniform per-layer head counts. All layers in
177
+ one attention group must agree on ``num_heads``; this is asserted.
178
+ Returns ``None`` when no matching Attention layer is found.
179
+ """
180
+ attn_layers = get_layers_from_vllm_config(
181
+ vllm_config,
182
+ AttentionLayerBase, # type: ignore[type-abstract]
183
+ layer_names,
184
+ )
185
+ if not attn_layers:
186
+ return None
187
+ heads = {layer.impl.num_heads for layer in attn_layers.values()}
188
+ assert len(heads) == 1, (
189
+ f"All layers in one attention group must share num_heads; "
190
+ f"got {heads} for {layer_names}."
191
+ )
192
+ return heads.pop()
193
+
194
+
195
def infer_global_hyperparameters(
196
per_layer_params: dict[str, PerLayerParameters],
197
) -> PerLayerParameters: