Prhub

#51772 [Attention][MLA] Fuse Kimi-K3 chunked-context K/V packing

原始 PR 作者 zyongye 合并时间 2026-08-13 13:42 文件变更 15 提交数 4 评论 6 代码增减 +996 / -35

performance kernel kimi attention nvidia

执行摘要

K3 上下文 K/V 打包融合为单 kernel，修复 fp8 cache 启动崩溃

PR body 指出原实现每 chunk 要 cast kv_nope 到 fp8、cast k_pe、拼接 [k_nope | k_pe]，并重复量化 new-token epilogue 已量化过的 query，目标是“让尾部折叠成一个 kernel per chunk”以削减 launch 与拷贝。第二 commit 的动机是修复确定性崩溃：cache dtype 与 kv_b_proj dtype 相互独立——plain fp8 cache 的 gather 结果是 fp8，而 stock Kimi-K3 checkpoint 的 kv_b_proj 是 bf16，首发版的加载期断言导致 --kv-cache-dtype fp8 下每 rank 在 weight load 阶段死亡。

值得精读。核心看点：① supports_out() 契约在 chunked-context 上的复用与 data_ptr 兜底断言；② cache dtype 与 kv_b_proj dtype 独立性这个容易踩坑的边界条件（第二 commit 的修复过程是很好的负样本教材）；③ “融合实现必须与通用实现逐 chunk 位级对照”的测试策略，以及 kernel 测试中“fp8 必须位级一致而非近似”的严谨断言。

讨论亮点

PR 的 review 评论为空：jeejeelee 直接 APPROVED（无文字），claude[bot] 自动回复表示 fork PR 关闭自动 review。最有价值的讨论在 commit 记录与 PR body 中：第一 commit 的加载期断言导致 stock checkpoint 在 --kv-cache-dtype fp8 下全 rank 崩溃；作者在发 PR 前自测发现并修复，根因是“cache dtype 与 kv_b_proj dtype 相互独立”，取舍是恢复条件 cast、删除断言，并强调“对 fp8 kv_b_proj 而言 .to 是 no-op，融合本身 untouched”。另一个设计决策是 supports_out() 的复用：run_chunk 用 data_ptr() 断言兜底“后端声称支持 out 却未写入”，把跨后端契约违约转成快速失败。

实现拆解

  1. 新增 fused K/V pack 算子链：vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py 新增 Python 入口 fused_mla_kv_concat / fused_mla_kv_concat_quant_fp8（含 _empty_full_key 辅助）；csrc/libtorch_stable/ 同步新增声明（ops.h）、torch 绑定（torch_bindings.cpp）与 CUDA kernel fusedKimiK3MLAKVConcatPackKernel。kernel 直接消费 strided 的 kv_b_proj 输出与仍处于 fp8 cache 布局的 k_pe，输出连续 key（fp8 路径还有连续 fp8 V）；fp8 cast 走 __nv_cvt_*_to_fp8x2 原生 pairwise 转换器，保证与 .to(torch.float8_e4m3fn) 位级一致；grid 被 cap，长 context 走 grid-stride loop。launchPdl 泛化为 launchPdlSlots 支持 rows-per-warp 与 max_blocks 配置。

  2. K3 层自持 context 循环：vllm/models/kimi_k3/nvidia/mla.py 新增 _compute_prefill_context 覆盖 impl 版本，内嵌 run_chunk（_gather_context_latent → 按 _get_kv_b_proj_input_dtype 条件 cast → kv_b_proj → fused pack → run_prefill_context_chunk）；_gather_context_latent 按 cache dtype（fp8_ds_mla / bf16 / plain fp8）分发 gather 路径，与通用 impl 保持一致。单 chunk 快路径直接返回其 partial，decode context parallelism 仍走 impl._context_parallel_compute_prefill_context 不加融合。

  3. 打通 supports_out() 契约：vllm/v1/attention/backends/mla/prefill/base.py 抽象方法 run_prefill_context_chunk 增加可选 out 参数，supports_out() 文档同步扩展；trtllm_ragged.py 实现 out 直写（out 为 None 时才自行分配），flash_attn.py、flashinfer.py、tokenspeed_mla.py、aiter_flash_attn.py 同步适配签名。后端声明支持 out 时，accumulator 可在任何 chunk 运行前按固定形状分配，并用新增的 neutralize_empty_context_partials 初始化无 context 的 prefill；非 continuation chunk 直写 output[token_slice]，累计时仅折叠 lse；accumulate_mla_context_chunk 新增 output_written 参数并对 continuation chunk 加断言。

  4. 测试配套：新增 tests/models/kimi_k3/test_mla_prefill_context.py，用 _RecordingPrefillBackend（模拟 honors-out / 不 honors-out 两类后端）与 _KVBProj（模拟 fp8/bf16 权重输入契约）把 fused 循环与 MLACommonBaseImpl._compute_prefill_context 逐 chunk 对照（q/k/v 与合并结果 atol=0/rtol=0），覆盖 2 种 cache dtype × 2 种 kv_b_proj dtype × 2 种 out 模式；test_kimi_k3_mla_fused_epilogue.py 扩展 strided 输入与 fp8 位级一致性用例；test_mla_prefill_registry.py 适配新签名。

  5. bugfix（第二 commit）：首发版把 impl 的逐 chunk latent cast 换成加载期断言 kv_b_proj 必须直接消费 fp8，stock checkpoint（bf16 kv_b_proj + fp8 cache）全 rank 崩溃。修复为在 run_chunk 内恢复 _get_kv_b_proj_input_dtype 驱动的条件 cast 并删除断言——fp8 kv_b_proj 场景下 .to 是 no-op，融合收益不受影响；新增测试专门覆盖 bf16_kv_b_proj + fp8 cache 两种 out 模式（第一 commit 上 2 个 FAILED，第二 commit 后转绿）。

文件 模块 状态 重要度
vllm/models/kimi_k3/nvidia/mla.py 注意力层 modified 8.75
tests/models/kimi_k3/test_mla_prefill_context.py 测试 added 8.3
vllm/model_executor/layers/attention/mla_attention.py 公共注意力层 modified 7.24
csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu 内核 modified 5.85
vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py 算子封装 modified 6.55
vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py 预填充后端 modified 6.06
vllm/v1/attention/backends/mla/prefill/base.py 预填充后端 modified 5.65

关键符号

_compute_prefill_context run_chunk _gather_context_latent fused_mla_kv_concat fused_mla_kv_concat_quant_fp8 neutralize_empty_context_partials accumulate_mla_context_chunk run_prefill_context_chunk fusedKimiK3MLAKVConcatPackKernel writeFullKeyPack

关键源码片段

vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py infrastructure

Python 侧算子封装：提供 fused_mla_kv_concat / fused_mla_kv_concat_quant_fp8 入口，处理 k_pe reshape 与空 token 边界

def fused_mla_kv_concat(
    k_nope: torch.Tensor, # [T, H, qk_nope_head_dim]，允许 strided
    k_pe: torch.Tensor, # [T, rope] 或 [T, 1, rope]，与 k_nope 同 dtype
) -> torch.Tensor:
    """一次 launch 内把 key 拼成连续布局的 [k_nope | k_pe]。
​
    这是 fused_mla_key_concat_kv_cache_insert 的 chunked-context 版本：
    没有 query、没有 cache insert、没有 RoPE（gathered 的 k_pe 已旋转）。
    k_nope 是 kv_b_proj 输出的 strided 半段，k_pe 是 gather workspace 的
    strided 视图，因此两者都无需先变 contiguous。
    """
    k_pe = k_pe.reshape(k_pe.shape[0], k_pe.shape[-1])
    k = _empty_full_key(k_nope, k_pe, k_nope.dtype)
    if k.shape[0]:
        torch.ops._C.fused_kimi_k3_mla_kv_concat(k_nope, k_pe, k)
    return k
​
​
def fused_mla_kv_concat_quant_fp8(
    k_nope: torch.Tensor, # [T, H, qk_nope_head_dim]，允许 strided
    k_pe: torch.Tensor, # [T, rope] 或 [T, 1, rope]，k_nope 的 dtype 或 fp8
    v: torch.Tensor, # [T, H, v_head_dim]，允许 strided
) -> tuple[torch.Tensor, torch.Tensor]:
    """fused_mla_kv_concat 之外再对 key 与 v 做 fp8 cast。
​
    k_pe 可能已经是 fp8：plain fp8 cache 的 gather 不做反量化，这些字节
    直接原样拷入（由 kernel 内 KPE_FP8 分支处理）。返回连续的 (k_fp8, v_fp8)，
    cast 使用原生 pairwise 转换器，与 torch 的 .to(fp8) 位级一致。
    """
    k_pe = k_pe.reshape(k_pe.shape[0], k_pe.shape[-1])
    fp8 = torch.float8_e4m3fn
    k_fp8 = _empty_full_key(k_nope, k_pe, fp8)
    v_fp8 = torch.empty(v.shape, dtype=fp8, device=v.device)
    if k_fp8.shape[0]:
        torch.ops._C.fused_kimi_k3_mla_kv_concat_quant_fp8(
            k_nope, k_pe, v, k_fp8, v_fp8
        )
    return k_fp8, v_fp8

评论区精华

bf16 kv_b_proj + fp8 cache 加载期崩溃 正确性

第一 commit 把 impl 的逐 chunk latent cast 替换成加载期断言 `kv_b_proj` 必须直接消费 fp8，导致 stock checkpoint（bf16 `kv_b_proj` + fp8 cache）在 `--kv-cache-dtype fp8` 下每 rank 权重加载即崩：`AssertionError: ... this checkpoint's kv_b_proj wants torch.bfloat16`。作者在发 PR 前自测发现，commit ebd55a1 修复。

结论：在 run_chunk 内恢复 `_get_kv_b_proj_input_dtype` 驱动的条件 cast，删除加载期断言；对 fp8 `kv_b_proj` 该 `.to` 是 no-op，融合收益保留，fusion 本体未动。新增测试覆盖该组合（第一 commit 上 2 个 FAILED，修复后转绿）。 · 已解决

supports_out() 契约复用与 continuation chunk 处理 设计

借助现有 supports_out() 能力把 out 贯穿到 run_prefill_context_chunk：后端承诺写 out 时 accumulator 可先按固定形状分配，非 continuation chunk 直写 output[token_slice]，每 chunk 只折叠 lse；continuation chunk 因前缀 token 需与已有 partial merge 而保留自身 buffer。run_chunk 用 data_ptr() 断言兜底后端违约。

结论：已实现并测试覆盖 honors_out 两种模式；wrote_in_place 断言验证非 continuation chunk 全部原地写入，continuation chunk 不写。 · 已解决

fp8 位级一致性而非近似比对 测试

kernel 测试断言 fused pack 与 torch `.to(fp8)` 位级一致（atol=0 / rtol=0），理由是 pack 使用原生 pairwise 转换器；测试还覆盖 k_pe 已是 fp8 布局（plain fp8 cache gather 不反量化）的字节直拷路径。

结论：已覆盖 bf16/fp16 输入、num_tokens 含 0、num_heads 3/4、k_pe dtype 两种；位级断言保证 cast 语义不被 kernel 实现悄悄改变。 · 已解决

风险与影响

  1. 路径核心但范围受限：融合仅作用于 Kimi-K3 的 NVIDIA chunked-context prefill（测试 skipif 非 CUDA），但该路径是 K3 长上下文推理主干；q.dtype == prefill.q_data_type 是硬断言，若 future new-token epilogue 行为变化会直接 fail-fast，不会静默错算。
  2. 跨后端契约变更：run_prefill_context_chunk 签名在 base + 5 个后端同步修改，默认 out=None 保持旧行为；任何后端误报 supports_out() 会被 data_ptr 断言击中，后续新增后端必须遵守契约。
  3. 位级 fp8 一致性：依赖 __nv_cvt_*_to_fp8x2 原生转换器在具体 arch 上的行为，测试仅在 B300 / CUDA 13 验证；ROCm 走 copyChunk8 老路径，本次 CI 未覆盖。
  4. 公共 helper 回归面：mla_attention.py 的 neutralize_empty_context_partials 与 accumulate_mla_context_chunk 被通用 MLA 路径复用，output_written 默认 False 保持原语义，但通用 impl 的任何后续改动都会传导到 K3 融合循环。

对用户/系统：Kimi-K3 在 NVIDIA 上的 chunked-context prefill 每 chunk 减少 1-2 次 kernel launch 与输出拷贝，fp8 cache 场景下同时消除重复 query 量化；stock checkpoint 此前无法以 fp8 cache 加载的阻塞性 bug 被修复。对通用 MLA：所有 MLA prefill 后端接口新增可选 out 参数（向后兼容），未来任何模型的自持 fused prefill 循环都可复用该直写模式。对团队：新增 1 个 CUDA kernel 与一套“融合路径 vs 通用路径逐 chunk 对照”的高价值测试方法论，可作为 fused 实现等价性的回归基线。影响程度中等偏上，局限于 K3/NVIDIA 路径。

K3 专属路径（仅 NVIDIA） prefill 后端契约扩展 位级 fp8 依赖原生转换器 启动崩溃修复需跨硬件回归 公共 MLA helper 被通用路径复用

关联 Issue

#48051 [ROCm] Don't re-quantize bf16 MLA kv_b_proj to fp4/fp8 for GLM MoE DSA

完整报告

Markdown PDF

参与讨论