执行摘要

• 一句话：修复 FlashInfer/Triton metadata 构建器的 num_qo_heads 回归
• 推荐动作：该 PR 值得快速合并，因为它恢复了一个已知正确的修复，解决了特定模型的严重错误。建议合并后补充针对非均匀 head 数模型的测试用例，防止今后类似回归。

功能与动机

修复 #42650 修复被 #43241 意外删除的问题。#41651 和 #47037 报告了同一类错误：在 num_attention_heads_per_layer 不统一的模型上，FlashInfer 和 Triton attention backend 因使用了模型全局的 head 数而非实际 Attention 层的值，导致 kv-cache 组分配不足，运行时产生非法内存访问错误。

实现拆解

1. 在 utils.py 中新增 get_num_attention_heads_from_layers 辅助函数：该函数通过 get_layers_from_vllm_config 获取指定名称的 AttentionLayerBase 实例，然后收集每个 layer.impl.num_heads 组成集合并断言所有层一致（一个 attention group 内必须统一），返回该值；若找不到 attention 层则返回 None。
2. 更新 flashinfer.py 的 FlashInferAttentionMetadataBuilder.__init__：修改 self.num_qo_heads 的赋值逻辑，优先调用 get_num_attention_heads_from_layers(vllm_config, layer_names)，如果返回 None 则回退到 model_config.get_num_attention_heads(parallel_config)，并更新 import 列表。
3. 更新 triton_attn.py 的 TritonAttentionMetadataBuilder.__init__：同样修改 self.num_heads_q 的赋值逻辑，优先调用 get_num_attention_heads_from_layers 再回退，并更新 import 列表。
4. 回退逻辑兼容所有模型：对于未设置 num_attention_heads_per_layer 的普通模型，get_num_attention_heads_from_layers 返回 None，此时仍使用原全局值，行为完全不变。

关键文件：

• vllm/v1/attention/backends/utils.py（模块 注意力模块；类别 source；类型 core-logic；符号 get_num_attention_heads_from_layers）: 新增 get_num_attention_heads_from_layers 函数，这是修复的核心逻辑。函数读取实际 Attention 层的 num_heads 并断言同组一致性，返回 per-layer 值或 None。
• vllm/v1/attention/backends/flashinfer.py（模块 注意力模块；类别 source；类型 core-logic）: 修改 FlashInferAttentionMetadataBuilder.__init__，将 self.num_qo_heads 从全局值改为优先使用 get_num_attention_heads_from_layers，并添加回退。同时更新 import。
• vllm/v1/attention/backends/triton_attn.py（模块 注意力模块；类别 source；类型 core-logic）: 修改 TritonAttentionMetadataBuilder.__init__，将 self.num_heads_q 从全局值改为优先使用 get_num_attention_heads_from_layers，并添加回退。同时更新 import。

关键符号：get_num_attention_heads_from_layers

关键源码片段

vllm/v1/attention/backends/utils.py

新增 get_num_attention_heads_from_layers 函数，这是修复的核心逻辑。函数读取实际 Attention 层的 num_heads 并断言同组一致性，返回 per-layer 值或 None。

# vllm/v1/attention/backends/utils.py
def get_num_attention_heads_from_layers(
    vllm_config: VllmConfig, layer_names: list[str]
) -> int | None:
    """Per-TP-rank ``num_heads`` shared by the named Attention layers.
​
    Use in metadata builders whose plan-time allocations depend on the
    head count: the model-wide ``get_num_attention_heads()`` is wrong
    for models with non-uniform per-layer head counts. All layers in
    one attention group must agree on ``num_heads``; this is asserted.
    Returns ``None`` when no matching Attention layer is found.
    """
    attn_layers = get_layers_from_vllm_config(
        vllm_config,
        AttentionLayerBase, # type: ignore[type-abstract]
        layer_names,
    )
    # 如果找不到 attention 层（如 profile 阶段尚未构建），返回 None
    if not attn_layers:
        return None
    # 收集所有层的 num_heads，并断言它们必须一致
    heads = {layer.impl.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}."
    )
    return heads.pop()

评论区精华

该 PR 没有 review 评论讨论，但有两名 reviewer 批准。

• TheEpicDolphin 承认他在 #43241 的 rebase 中意外删除了 #42650 的变更，并为此道歉。
• mgoin 简单表态 +1 并感谢。

• claude[bot] 自动评论但由于是从 fork 发起的 PR 而无法执行自动审查。

• 暂无高价值评论线程

风险与影响

• 风险：
  1. 回归风险低：变更仅涉及三个源文件，改动逻辑与已合并的 #42650 一致，只是恢复已审阅过的代码。新的 get_num_attention_heads_from_layers 函数有回退逻辑（找不到层时返回 None，然后使用全局值），对现有模型行为无影响。
  2. 缺少测试：本次恢复性变更没有附带对应的测试文件。虽然 #42650 当时可能有测试，但在此 PR 中未引入，建议后续补充针对非均匀 head 数模型的集成测试。
  3. 仅影响 FlashInfer 和 Triton attention backend：flashinfer.py 和 triton_attn.py 是两个主要后端，其他后端未受影响。
• 影响：
  1. 对用户的影响：使用具有非均匀 num_attention_heads_per_layer 的模型（如 poolside/Laguna-XS.2-FP8）的用户将不再遭遇 FlashInfer 或 Triton 后端的非法内存访问错误。其他模型用户不受影响。
  2. 对系统的影响：仅在 attention metadata 构建阶段增加一次轻量级层查询开销，不影响推理性能。
  3. 对团队的影响：修复了一个回归问题，恢复了对这部分模型功能的支持。
     - 风险标记：缺少测试覆盖

关联脉络

• PR #42650 [Bugfix] Source num_qo_heads from Attention layers in Flashinfer/Triton metadata builders: 本 PR 恢复的是 #42650 的变更，完全相同的修复逻辑。
• PR #43241 [Model Runner V2][Spec Decode] Add Gemma4 MTP support: #43241 在 rebase 过程中意外删除了 #42650 的代码，导致本 PR 修复的回归。
• PR #41651 [Bug] Original bug report: #41651 是 #42650 修复的原始 bug 报告，本 PR 重新修复该问题。