执行摘要

• 一句话：CPU 注意力分组 head count 改从层读取，修复 Laguna 解码崩溃
• 推荐动作：值得精读。重点关注三点：一是 commit 演进从"按不同 head count 构建多份元数据字典"收敛为"利用 attention group 以 num_heads_q 为 key 的不变量直接取组内 head count"，是典型的用既有不变量简化设计的好案例；二是 get_num_attention_heads_from_layers() 的跨后端复用，体现 vLLM 正在标准化分组元数据的读取方式；三是"用断言拒绝不可能状态"（混合 head count 的 group）而非静默容忍的防御风格。

功能与动机

PR body 明确指出：CPUAttentionMetadataBuilder 用模型全局 query head count 给 split-KV scratchpad 定大小，但 kernel 运行在每个层自己的 count 上，因此像 Laguna 这样逐层变化 head count 的模型会索引越过 scratchpad 边界，decode 阶段 segfault 或 hang。作者用仪器化数据证实了越界：每 KV head 只有 6336 bytes，而 split-KV 写路径需要 8384 bytes（sizing 假设 6 个 query head 对应 1 个 KV head，实际 64-head 层用 stride 8 寻址）。GSM8K 全量评测在修复前直接以 execute_model 超时告终。

实现拆解

实现分 4 步完成：

1. 变更入口（核心逻辑）：修改 vllm/v1/attention/backends/cpu_attn.py 中 CPUAttentionMetadataBuilder.__init__。原先 self.num_heads = vllm_config.model_config.get_num_attention_heads(parallel_config) 取模型全局 count；改为 self.num_heads = get_num_attention_heads_from_layers(vllm_config, layer_names) or vllm_config.model_config.get_num_attention_heads(parallel_config)。get_num_attention_heads_from_layers 从 vllm.v1.attention.backends.utils 导入，是 flashinfer 与 triton 后端已经为同样原因使用的既有 helper；or 回退保证 helper 取不到层时行为不变。
2. 解析时机调整（设计简化）：window_size 从"首次 build() 时惰性解析"改为 __init__ 中直接 self.window_size = self._group_sliding_window()，删除 self.window_size: int | None = None 声明与 build() 内的 if self.window_size is None 分支。原理是 layer_names 在 builder 构造时已确定、层属性是常量；同时解决 mypy 在嵌套调度元数据 helper 中无法收窄 Optional 的问题。初始方案中的 per-count 元数据字典与 extra_num_heads 属性在最终提交中被移除。
3. 测试配套：新增 tests/v1/attention/test_group_head_counts.py，仅在 current_platform.is_cpu() 时运行。用 MagicMock + SimpleNamespace 构造 attention group 的层集合，参数化验证 builder.num_heads 优先取 group 自身 count（覆盖 48/64/16 三种取值），并断言混合 head count 的 group 会触发 AssertionError（match "share num_heads"）。
4. CI 配套：.buildkite/hardware_tests/cpu.yaml 在 CPU-Kernel Tests 的 source_file_dependencies 与 pytest 命令中增加 tests/v1/attention/test_group_head_counts.py，确保该回归测试在 CPU 硬件 CI 上执行。

关键文件：

• vllm/v1/attention/backends/cpu_attn.py（模块 CPU 后端；类别 source；类型 core-logic；符号 CPUAttentionMetadataBuilder.init, CPUAttentionMetadataBuilder._group_sliding_window, CPUAttentionMetadataBuilder.build）: 核心修复所在。CPUAttentionMetadataBuilder 的 head count 来源从模型全局值改为 group 内层联合取值，window_size 解析提前到构造期，直接决定 split-KV scratchpad 大小与 kernel 写入边界。
• tests/v1/attention/test_group_head_counts.py（模块 注意力测试；类别 test；类型 test-coverage；符号 _layers, _build, test_num_heads_comes_from_the_group, test_mixed_head_counts_in_one_group_are_rejected）: 新增单元测试，镜像 test_group_sliding_window.py 的结构，直接锁定回归：验证 group 自身 head count 优先于模型全局值，且混合 head count 的 group 被断言拒绝。
• .buildkite/hardware_tests/cpu.yaml（模块 CI 配置；类别 test；类型 test-coverage）: CPU 硬件 CI 配置，将新单测加入 source_file_dependencies 与 pytest 命令，保证 CPU 回归测试在真实硬件上执行。

关键符号：CPUAttentionMetadataBuilder.init, CPUAttentionMetadataBuilder._group_sliding_window, CPUAttentionMetadataBuilder.build, get_num_attention_heads_from_layers, test_num_heads_comes_from_the_group, test_mixed_head_counts_in_one_group_are_rejected

关键源码片段

vllm/v1/attention/backends/cpu_attn.py

核心修复所在。CPUAttentionMetadataBuilder 的 head count 来源从模型全局值改为 group 内层联合取值，window_size 解析提前到构造期，直接决定 split-KV scratchpad 大小与 kernel 写入边界。

class CPUAttentionMetadataBuilder(AttentionMetadataBuilder[CPUAttentionMetadata]):
    def __init__(
        self,
        kv_cache_spec: AttentionSpec,
        layer_names: list[str],
        vllm_config: VllmConfig,
        device: torch.device,
    ) -> None:
        super().__init__(kv_cache_spec, layer_names, vllm_config, device)
​
        self.kv_cache_spec = kv_cache_spec
        self.vllm_config = vllm_config
​
        parallel_config = vllm_config.parallel_config
        self.num_kv_heads = kv_cache_spec.num_kv_heads
        # 调度元数据会按 query head count 分配 split-KV scratchpad，
        # 因此该 count 必须取自本 group 覆盖的层：模型全局 count 对
        # 每层 head count 不同的模型（如 Laguna）是错误的，
        # 过小的 scratchpad 会被越界索引并导致 decode 段崩溃。
        self.num_heads = get_num_attention_heads_from_layers(
            vllm_config, layer_names
        ) or vllm_config.model_config.get_num_attention_heads(parallel_config)
        self.head_dim = kv_cache_spec.head_size
        self.dtype = vllm_config.model_config.dtype
        # group 的滑窗同样在构造期解析：layer_names 在 builder 构建时已确定，
        # 层属性是常量，无需再等到首次 build() 时惰性求值，
        # 也因此去掉了 window_size 的 Optional 类型。
        self.window_size = self._group_sliding_window()
        self.block_size = vllm_config.cache_config.block_size
        self.kv_cache_dtype = vllm_config.cache_config.cache_dtype
        self.isa = _get_attn_isa(
            self.dtype,
            self.block_size,
            self.head_dim,
            self.kv_cache_dtype,
        )
        self.is_cross_attention = isinstance(kv_cache_spec, CrossAttentionSpec)
        self.is_encoder_only_attention = isinstance(
            kv_cache_spec, EncoderOnlyAttentionSpec
        )
​
    def _group_sliding_window(self) -> int:
        """返回本 group 所有层共享的滑窗大小，无则返回 -1。
​
        窗口从层而非 group spec 读取：一个 KV cache group 可能同时包含
        带窗口与不带窗口的层（例如 Gemma-3 关闭 hybrid KV cache manager
        的场景），而这里构建的调度元数据由整个 group 共享，因此只能假设
        所有层一致认可的窗口。
        """
        layers = get_layers_from_vllm_config(
            self.vllm_config, Attention, self.layer_names
        )
        windows = {
            layer.impl.sliding_window
            for layer in layers.values()
            if isinstance(layer.impl, CPUAttentionBackendImpl)
        }
        if len(windows) != 1:
            return -1
        window = windows.pop()
        return -1 if window is None else window

tests/v1/attention/test_group_head_counts.py

新增单元测试，镜像 test_group_sliding_window.py 的结构，直接锁定回归：验证 group 自身 head count 优先于模型全局值，且混合 head count 的 group 被断言拒绝。

def _build(layer_num_heads: list[int]) -> CPUAttentionMetadataBuilder:
    # 用 mock 层集合模拟一个 attention group：每层只暴露 num_heads 与
    # sliding_window 两个属性，模型全局 head count 固定为 48，
    # 用于验证 builder 优先采用 group 自己的 count 而非模型全局值。
    layers = _layers(layer_num_heads)
    vllm_config = MagicMock()
    vllm_config.model_config.dtype = torch.bfloat16
    vllm_config.model_config.get_num_attention_heads.return_value = MODEL_WIDE_NUM_HEADS
    vllm_config.cache_config.block_size = 16
    vllm_config.cache_config.cache_dtype = "auto"
    kv_cache_spec = SimpleNamespace(num_kv_heads=NUM_KV_HEADS, head_size=64)
​
    with (
        patch(
            "vllm.v1.attention.backends.utils.get_layers_from_vllm_config",
            return_value=layers,
        ),
        patch(
            "vllm.v1.attention.backends.cpu_attn.get_layers_from_vllm_config",
            return_value=layers,
        ),
    ):
        return CPUAttentionMetadataBuilder(
            kv_cache_spec=kv_cache_spec,
            layer_names=list(layers),
            vllm_config=vllm_config,
            device=torch.device("cpu"),
        )
​
​
@pytest.mark.parametrize("group_num_heads", [MODEL_WIDE_NUM_HEADS, 64, 16])
def test_num_heads_comes_from_the_group(group_num_heads):
    # group 自身的 count 优先，即便它不同于模型全局值（48）；
    # 64 对应 Laguna 的滑动层，16 对应更小的注意力分组。
    builder = _build([group_num_heads, group_num_heads])
    assert builder.num_heads == group_num_heads

评论区精华

review 中三条实质讨论均已解决：

1. 测试平台限制：bigPYJ1151 建议新测试加 current_platform.is_cpu() 检查、只跑 CPU 后端，并加入 .buildkite/hardware_tests/cpu.yaml，作者回复 "Done"，最终测试文件带 pytestmark = pytest.mark.skipif(...) 且 CPU CI 已纳入。
2. 常量属性提前解析：bigPYJ1151 问能否把 build() 中惰性解析的 window_size（及当时的 extra_num_heads）移到 __init__，因为这些属性是常量；作者回复 "Done"，最终把 window_size 改为构造期解析，顺带去掉 Optional 类型。
3. head count 来源：bigPYJ1151 指出 num_heads 应直接从 layer_names[0] 读取，因为同一 attention group 的层必然共享同一 num_heads；作者回复采用既有 helper get_num_attention_heads_from_layers()，理由是 triton_attn 和 flashinfer 已经为此使用同一 helper，比手取 layer_names[0] 更稳健。

• 新测试应限 CPU 平台并纳入 CPU CI (testing): 已采纳：测试文件带 pytestmark = pytest.mark.skipif(not current_platform.is_cpu())，且 cpu.yaml 的依赖与命令均增加该测试。
• window_size 与 extra_num_heads 惰性解析提前到 init (design): 已采纳：window_size 改为构造期解析，删除 Optional 类型与 build() 内惰性分支；后续方案演进中还移除了 extra_num_heads 相关逻辑。
• head count 应从 group 的层读取而非模型全局值 (correctness): 采用既有 helper get_num_attention_heads_from_layers()，与 flashinfer、triton 后端保持一致，并保留模型全局值作为回退。

风险与影响

• 风险：
  1. CPU 后端核心路径：所有 CPU 推理请求都会经过 CPUAttentionMetadataBuilder 的构造与 build()，scratchpad 大小直接决定 split-KV kernel 的写入边界，属推理正确性关键路径；不过对 head count 均匀的模型，helper 返回值与模型全局值一致，保持原 single-blob 路径，无额外分配或查找。
  2. window_size 构造期解析：从惰性解析改为 __init__ 即解析，若未来出现"层尚未注册就构建 builder"的调用路径会提前报错。作者声称层在 builder 构建时已存在，单测覆盖了构造路径，但 fast_build 等未来路径需留意。
  3. 依赖 group 一致性断言：get_num_attention_heads_from_layers 对混合 head count 会触发 AssertionError（"share num_heads"），这依赖"attention group 以 num_heads_q 为 key"这一不变量；若未来分组规则变化，会以显式断言暴露而非静默越界，总体更安全但仍是隐性契约。
  4. 测试覆盖局限：新增单测是 mock 级，Laguna 真机验证（reproducer、GSM8K）未固化进 CI，后续回归只能靠 CPU CI 中的单测拦截。
     - 影响：影响范围：本 PR 只影响 CPU 后端 v1 注意力路径。对 Laguna 这类每层 head count 变化的模型，从"必然 segfault/hang"变为可正常推理（TP=2 下 300-token 解码 47 秒完成，GSM8K 全量 0.90 exact match）；对 head count 均匀的模型无行为变化。系统层面，元数据构建逻辑更简洁（去掉 Optional 与惰性初始化分支），并且让 CPU 后端与 flashinfer、triton 后端在"分组元数据从层读取属性"这一模式上对齐。该 PR 被归入 v0.28.0 cherry picks 里程碑，说明需要回传稳定分支。
     - 风险标记：CPU 后端核心路径变更, 窗口解析提前到构造期, 依赖 group 一致性断言, e2e 验证未固化 CI

关联脉络

• PR #52161 [Bugfix] Detect all attention-spelling variants in ModelConfig.is_hybrid: 同为注意力配置/元数据与实际 kernel 行为不一致导致的推理崩溃修复，且都针对 v1 注意力路径，属于同一条正确性修复线。
• PR #52512 [Bugfix][MLA] Do not use Dense MHA for GLM-5.2: 同为注意力分组/后端选择与模型配置不一致导致的输出退化或崩溃修复，验证了"配置必须贴合层实际行为"这一主题。