A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 23:15 UTC
[Bugfix] Source num_qo_heads from Attention layers in Flashinfer/Triton metadata builders (#42650)
Signed-off-by: zhanda <zhandazhu@gmail.com> Co-authored-by: Shang Wang <shangw@nvidia.com>
Zhanda Zhu committed
May 22, 2026 at 12:10 UTC
c7624bea5ebba1c688eb4c216bd4ede7a94f2a82
3 files changed
+39
-7
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,
@@ -607,9 +608,10 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
608
self.use_dcp and vllm_config.parallel_config.dcp_comm_backend == "a2a"
609
)
610
610
- self.num_qo_heads = self.model_config.get_num_attention_heads(
611
- self.vllm_config.parallel_config
612
- )
611
+ # Compatible with models with non-uniform per-layer head counts.
612
+ self.num_qo_heads = get_num_attention_heads_from_layers(
613
+ vllm_config, layer_names
614
+ ) or self.model_config.get_num_attention_heads(self.vllm_config.parallel_config)
615
616
self.num_kv_heads = self.kv_cache_spec.num_kv_heads
617
self.head_dim = self.kv_cache_spec.head_size
vllm/v1/attention/backends/triton_attn.py
+8
-4
@@ -30,7 +30,10 @@ from vllm.v1.attention.backend import (
30
CommonAttentionMetadata,
31
MultipleOf,
32
)
33
-from vllm.v1.attention.backends.utils import get_kv_cache_layout
33
+from vllm.v1.attention.backends.utils import (
34
+ get_kv_cache_layout,
35
+ get_num_attention_heads_from_layers,
36
+)
37
from vllm.v1.attention.ops.triton_prefill_attention import context_attention_fwd
38
from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
39
triton_reshape_and_cache_flash,
@@ -139,9 +142,10 @@ class TritonAttentionMetadataBuilder(AttentionMetadataBuilder[TritonAttentionMet
142
self.block_size = kv_cache_spec.block_size
143
144
model_config = vllm_config.model_config
142
- self.num_heads_q = model_config.get_num_attention_heads(
143
- vllm_config.parallel_config
144
- )
145
+ # Compatible with models with non-uniform per-layer head counts.
146
+ self.num_heads_q = get_num_attention_heads_from_layers(
147
+ vllm_config, layer_names
148
+ ) or model_config.get_num_attention_heads(vllm_config.parallel_config)
149
self.num_heads_kv = model_config.get_num_kv_heads(vllm_config.parallel_config)
150
self.headdim = model_config.get_head_size()
151
vllm/v1/attention/backends/utils.py
+26
@@ -136,6 +136,32 @@ def get_per_layer_parameters(
136
return per_layer_params
137
138
139
+def get_num_attention_heads_from_layers(
140
+ vllm_config: VllmConfig, layer_names: list[str]
141
+) -> int | None:
142
+ """Per-TP-rank ``num_heads`` shared by the named Attention layers.
143
+
144
+ Use in metadata builders whose plan-time allocations depend on the
145
+ head count: the model-wide ``get_num_attention_heads()`` is wrong
146
+ for models with non-uniform per-layer head counts. All layers in
147
+ one attention group must agree on ``num_heads``; this is asserted.
148
+ Returns ``None`` when no matching Attention layer is found.
149
+ """
150
+ attn_layers = get_layers_from_vllm_config(
151
+ vllm_config,
152
+ AttentionLayerBase, # type: ignore[type-abstract]
153
+ layer_names,
154
+ )
155
+ if not attn_layers:
156
+ return None
157
+ heads = {layer.impl.num_heads for layer in attn_layers.values()}
158
+ assert len(heads) == 1, (
159
+ f"All layers in one attention group must share num_heads; "
160
+ f"got {heads} for {layer_names}."
161
+ )
162
+ return heads.pop()
163
+
164
+
165
def infer_global_hyperparameters(
166
per_layer_params: dict[str, PerLayerParameters],
167
) -> PerLayerParameters: