# PR #51040 完整报告

- 仓库：`vllm-project/vllm`
- 标题：[ROCm][K3] Extend FP8 asm MLA prefill to non-divisor small head counts
- 合并时间：2026-08-27 01:47
- 原文链接：http://prhub.com.cn/vllm-project/vllm/pull/51040

---

# 执行摘要

- 一句话：K3 FP8 MLA prefill 支持非 16 整除小头数，长上下文不再 OOM
- 推荐动作：值得精读。这是 ROCm MLA FP8 推理栈中「kernel 对齐要求」与「真实模型头数」矛盾的典型案例，展示了如何用 replicate-pad/slice-back 以少量拷贝换取 asm kernel 的启用，以及如何通过统一 helper 避免 metadata 与 forward 宽度漂移。关注 `get_fp8_prefill_num_heads` 的设计、workspace 化输出 scratch 的做法，以及 review 中关于 `num_real_heads` in-kernel 支持的后续展望。

# 功能与动机

Kimi-K3 在 TP8 下每 rank 仅有 12 个 attention heads，而 AITER FP8 MLA asm prefill（`mla_prefill_ps_asm_fwd` + `mla_reduce_v1`）原门槛要求 `num_heads % 16 == 0`，导致 K3 回退到 BF16 FMHA 解压路径；该回退构建的 bf16 工作集不在 FP8 KV 池 accounting 范围内，长上下文（约 197k tokens）时激活 arena 耗尽 OOM，而 KV pool 使用率不足 4%。PR body 明确指出这是 #50578（asm decode pad 12→16）的 prefill 对应部分，利用 MLA 每个 query head 在共享 latent KV 上独立计算的性质，将 head 填充到 16 再切回是精确操作。

# 实现拆解

1. **放宽 `_fp8_prefill_enabled` 门槛 **（`vllm/v1/attention/backends/mla/rocm_aiter_mla.py`）：由 `_fp8_mla_prefill_supported() and self.num_heads % 16 == 0` 改为同时要求 `kv_cache_dtype_str == "fp8"`、`model_config.dtype == torch.bfloat16` 和 `AiterMLAHelper.is_valid_num_heads(self.num_heads)`；这样 bf16 KV 服务不会预留 PS workspace，fp16 模型也不会把 bf16 输出位模式误读为 fp16（review 中发现的真实 gap）。
2. **统一 head pad 目标到 helper**：新增 `AiterMLAHelper.get_fp8_prefill_num_heads`，并让 `_init_fp8_prefill_ps_buffers` 与 `_build_fp8_prefill_ps_metadata` 的 `num_head_k` 从 `max(16, num_heads)` 改为该 helper 结果，使 work/reduce maps 与 forward 实际 pad 宽度一致；该改动同时降低 partial tiles（gcd 从 4 变 16），实测 batch=1/qlen=512 时保留内存从 193.6 MiB 降到 68.6 MiB，为长上下文节省约 6 GiB。
3. **`_mla_fp8_prefill_attn` 内 replicate-pad 与 slice-back**：对非 16 倍数小头数（如 12）用 #50578 引入的 `get_mla_padded_q` 将 Q/K/V pad 到 16，运行 PS asm prefill 与 `mla_reduce_v1` 后把输出切回真实 `num_heads`。
4. **workspace 化输出 scratch**：`out_3d` 从每次 `torch.empty` 改为从全局 workspace manager 的 `get_simultaneous` 调用中获取，与 `logits`/`attn_lse`/`final_lse` 同一来源，避免 per-call 动态分配，并在 `_init_fp8_prefill_ps_buffers` 中预留最大形状；同时用 `ATTN_OUT_DTYPE = torch.bfloat16` 常量统一输出缓冲与 workspace 预留 dtype。
5. **测试覆盖**：`tests/kernels/attention/test_rocm_aiter_mla_fp8_prefill.py` 将单一 `NUM_HEADS = 16` 改为 `HEAD_COUNTS = [12, 16, 24]`，覆盖“非 divisor pad 到 16”“无需 pad 输出别名”“pad 到 32”三种 regime；`_make_impl` 与 `_build_prefill_metadata` 增加 `num_heads` 与 `attn_out_dtype` 参数，新增 `test_fp8_prefill_matches_reference` 参数化与 `test_fp8_prefill_metadata_width_matches_forward` 元数据宽度一致性用例。

关键文件：
- `vllm/v1/attention/backends/mla/rocm_aiter_mla.py`（模块 注意力后端；类别 source；类型 core-logic；符号 get_mla_padded_q, get_fp8_prefill_num_heads, _mla_fp8_prefill_attn, _init_fp8_prefill_ps_buffers）: 核心实现文件：放宽 FP8 prefill 头数门槛、新增统一 pad 头数 helper、将 PS metadata/workspace 预留与 forward pad 目标对齐、用 workspace manager 承载 padded-head 输出 scratch。
- `tests/kernels/attention/test_rocm_aiter_mla_fp8_prefill.py`（模块 kernel 测试；类别 test；类型 test-coverage；符号 _make_impl, _build_prefill_metadata, test_fp8_prefill_matches_reference, test_fp8_prefill_metadata_width_matches_forward）: 测试覆盖从单一 16 头扩展到 12/16/24 三种 regime，验证 pad/slice 正确性与 metadata 宽度与 forward 一致，并顺带暴露 fp16 模型误用 bf16 输出的缺陷。

关键符号：get_fp8_prefill_num_heads, _mla_fp8_prefill_attn, _init_fp8_prefill_ps_buffers, _build_fp8_prefill_ps_metadata, get_mla_padded_q

## 关键源码片段

### `vllm/v1/attention/backends/mla/rocm_aiter_mla.py`

核心实现文件：放宽 FP8 prefill 头数门槛、新增统一 pad 头数 helper、将 PS metadata/workspace 预留与 forward pad 目标对齐、用 workspace manager 承载 padded-head 输出 scratch。

```python
# vllm/v1/attention/backends/mla/rocm_aiter_mla.py（整理后核心片段）

def get_fp8_prefill_num_heads(num_heads: int) -> int:
    """FP8 asm prefill 与 mla_reduce_v1 要求 16 对齐的头数。

    对非 divisor 的小头数（如 K3 在 TP8 下每 rank 12 头），
    _mla_fp8_prefill_attn 会把 Q/K/V replicate-pad 到 16 头再 slice 回来；
    PS metadata 的 num_head_k 也必须用同一个目标宽度，否则 work/reduce
    maps 与 kernel 实际拿到的张量宽度不一致。统一走这里，同时避免
    max(16, num_heads) 在 24 头时产生更碎的 partial tiles。
    """
    return max(_AITER_MIN_MLA_HEADS, num_heads)  # 16 的倍数向上取整

# __init__ 中的启用门槛：FP8 KV + bf16 模型 + 有效头数，缺一不可
self._fp8_prefill_enabled = _fp8_mla_prefill_supported() and (
    kv_cache_dtype_str == "fp8"
    and vllm_config.model_config.dtype == torch.bfloat16
    and AiterMLAHelper.is_valid_num_heads(self.num_heads)
)

# _init_fp8_prefill_ps_buffers 中：metadata 按 pad 后的宽度构造
# 之前是 num_head_k = self.num_heads；对 24 头会保留 24 宽 work maps，
# 但 forward 实际 pad 到 32，导致 64 个 partial tiles（约 193.6 MiB）；
# 统一用 helper 后取 32，partial tiles 降到 16（约 68.6 MiB，gfx950 实测）。
num_head_k = AiterMLAHelper.get_fp8_prefill_num_heads(self.num_heads)

```

### `tests/kernels/attention/test_rocm_aiter_mla_fp8_prefill.py`

测试覆盖从单一 16 头扩展到 12/16/24 三种 regime，验证 pad/slice 正确性与 metadata 宽度与 forward 一致，并顺带暴露 fp16 模型误用 bf16 输出的缺陷。

```python
# tests/kernels/attention/test_rocm_aiter_mla_fp8_prefill.py（整理后核心片段）

# 覆盖三种 head 数 regime：
# 12 -> pad 到 16（非 divisor，tile+slice）——K3 在 TP8 下的真实场景
# 16 -> 无需 pad，输出直接别名调用方 buffer
# 24 -> pad 到 32，证明 16 以上也存在 pad 路径（当前无模型触发，防未来回归）
HEAD_COUNTS = [12, 16, 24]

# asm kernel 以原始指针写 bf16，输出缓冲与 workspace 预留必须同 dtype
ATTN_OUT_DTYPE = torch.bfloat16

@pytest.mark.parametrize("num_heads", HEAD_COUNTS)
@pytest.mark.parametrize("seq_len", [128, 512])
@torch.inference_mode()
def test_fp8_prefill_matches_reference(num_heads: int, seq_len: int) -> None:
    # 用生产 builder 构建 metadata，避免手搓导致与实现漂移
    metadata, total_q, builder = _build_prefill_metadata(
        [seq_len], device, num_heads, ATTN_OUT_DTYPE
    )
    # ... 构造 q/k/v，调用 _mla_fp8_prefill_attn，与 fp8 量化输入的
    # causal SDPA 参考实现对比，atol=1e-1 / rtol=5e-2 覆盖 fp8 量化噪声

def test_fp8_prefill_metadata_width_matches_forward(num_heads: int) -> None:
    # 断言 builder 中 fp8_prefill_* 的宽度 == get_fp8_prefill_num_heads(num_heads)，
    # 防止 metadata sizing 与 forward pad 目标脱钩（此前 max(16, n) 在 24 头时失配）
    ...

```

# 评论区精华

- **输出 copy 能否消除 **（Rohan138）：询问能否后续让 AITER 在 kernel 内支持 padding+slicing，并关注 `out_3d` 初始化对 CUDAGraph/ 持久 MLA/prefix caching 的影响。作者回应：`out` 是按 12 头打包的，kernel 写 16 头行，任何 view 都无法直接接收结果，必须 scratch+slice-back；输出 copy 仅约 6 KB/token，而 Q/K/V padding 约 28 KB/token，`num_real_heads` 参数可一次性消除四份拷贝；同时 `out_3d` 已改为从 workspace manager 获取。
- **测试中硬编码 dtype**（AndreasKaratzas）：质疑测试里为何硬编码 `attn_out_dtype=torch.bfloat16`。作者承认并修复，还顺带发现非测试路径的真实缺陷：fp16 模型 + fp8 KV cache 会走该路径导致 bf16 位被当作 fp16 解释，于是新增 `model_config.dtype == torch.bfloat16` 门禁。
- **TP<8（如 TP4 时 24 头）是否启用 **（shen-shanshan）：作者调研后表示 K3 因权重容量无法跑 TP4，并逐一核对现有 MLA 模型（K2 64、LongCat-Flash 64、Kimi-VL 16、DeepSeek-V3/R1 128）在各 TP 下均不会落入 16 以上非倍数区间；MiniMax-M2（48）与 GLM-4.6（96）虽会落入但均为 GQA 而非 MLA，因此当前无实际触发场景，但泛化改动仍值得做。
- **硬编码 16 的健壮性 **（shen-shanshan）：建议把所有 `16` 换成 `AiterMLAHelper._AITER_MIN_MLA_HEADS`。作者接受并进一步发现 PS metadata sizing 中的 `max(16, num_heads)` 也需要与 forward pad 目标联动，因此统一收敛到 `get_fp8_prefill_num_heads` 一个 helper。

- padded-head 输出 copy 能否消除 (design): 作者回应 copy 无法从 vLLM 侧消除（out 是 12 头打包，kernel 写 16 头行），但输出 copy 仅约 6 KB/token，Q/K/V padding 约 28 KB/token；AITER 加 num_real_heads 参数可一次消除四份拷贝，可开 follow-up。out_3d 已改为从 workspace manager 获取。
- 测试中硬编码 bf16 与 fp16 模型误用 (correctness): 测试使用 ATTN_OUT_DTYPE 常量；_fp8_prefill_enabled 新增 model_config.dtype == torch.bfloat16 门禁，fp16 模型回退到 BF16 prefill 路径。
- TP<8 时 24 头是否启用 asm prefill (question): 作者核实 K3 无法跑 TP4（权重容量），并逐一检查现有 MLA 模型（K2 64、LongCat-Flash 64、Kimi-VL 16、DeepSeek-V3/R1 128）在各 TP 下都不会落入 16 以上非倍数 band；MiniMax-M2/GLM-4.6 虽会但为 GQA。泛化改动仍值得。
- 硬编码 16 的健壮性 (design): 作者接受并把 PS metadata sizing 的 max(16, num_heads) 一并收敛到 get_fp8_prefill_num_heads，避免 forward pad 目标与 metadata 宽度脱钩。

# 风险与影响

- 风险：
 - **fp16 模型误用 bf16 输出**：asm prefill 以原始指针写 bf16，若 fp16 模型启用该路径会把 bf16 位模式重解释为 fp16 产生静默错误；已通过 `model_config.dtype == torch.bfloat16` 门禁修复，但依赖该门禁在后续重构中不被移除。
 - **pad/slice 的数值与形状一致性**：replicate-pad 依赖 MLA 每 head 独立计算的语义，若未来引入跨 head 的 kernel 或 MLA 结构变化，pad 后 slice 的精确性可能被破坏；测试覆盖 12/16/24 三种头数，但 24 头场景当前无真实模型。
 - **workspace 预留变化**：`num_head_k` 从 `max(16, num_heads)` 改为 `get_fp8_prefill_num_heads` 后，workspace 预留显著变小（约 6 GiB），若 PS metadata 与 forward 的 pad 目标不同步，会出现越界或预留不足；测试 `test_fp8_prefill_metadata_width_matches_forward` 正是针对此风险。
 - **CUDAGraph/prefix caching 交互**：`out_3d` 改为从 workspace manager 获取后，其地址稳定性依赖 workspace manager 的持续性；review 中 Rohan138 对此提出过疑虑，作者称已与 logits 等同等来源，但 CUDAGraph 捕获场景仍需回归验证。
 - **CUDA graph 与 breakable cudagraph**：PR body 的 serve 命令显式设置 `VLLM_USE_BREAKABLE_CUDAGRAPH=0`，说明该路径与 breakable cudagraph 的兼容性尚未完全验证，部署配置受限。
 - 影响：对 **Kimi-K3 on ROCm gfx950（MI355X）**用户是实质性修复：长上下文 prefill 从 OOM（约 197k tokens）提升到 470k/590k tokens 可正常完成，PS workspace 节省约 6 GiB，GSM8K 5-shot 准确率 97.04% 与基线 96.66% 相比无回退（在噪声范围内）。对 **其他 MLA 模型 **（K2、LongCat-Flash、Kimi-VL、DeepSeek-V3/R1）当前各 TP 下头数均为 16 的倍数或 16 以下，行为不变但门槛放宽后未来非 divisor 头数模型也可自动受益。对 **团队 **而言，该 PR 确立了“头数 pad 目标与 PS metadata sizing 必须共用一个 helper”的设计约束，并沉淀了 kernel 级数值对比测试与元数据宽度测试，降低后续回归风险。测试文件仅 gfx950 + AITER 构建可用，其他平台自动 skip。
 - 风险标记：核心路径变更 , 平台特定（gfx950/AITER）, 新增门禁依赖 , workspace 预留敏感 , CUDAGraph 交互待验证

# 关联脉络

- PR #50578 [ROCm][MLA] Support non-divisor small head counts in asm decode (pad 12 to 16): 本 PR 明确是 #50578 的 prefill 对应部分，复用了其 AiterMLAHelper.get_mla_padded_q 与头数 pad 语义；依赖其先合并。
- PR #48712 [ROCm][MLA] Gate PS workspace on FP8 KV cache: 本 PR 在 _fp8_prefill_enabled 中加入了与 #48712 相同的 FP8 KV 门禁（避免 bf16 KV 服务预留 PS workspace），PR body 说明与该 PR 组合使用。