#39232 [Perf] trtllm_mla: reuse the fused fp8 KV/Q prepare on target verify

原始 PR 作者 zhendonghua 合并时间 2026-09-14 11:59 文件变更 1 提交数 1 评论 5 代码增减 +25 / -7

执行摘要

验证路径复用解码融合准备，MLA 每层 preparation 由 5 次 launch 降为 1 次

动机来自 issue #39107：作者在 SGLang v0.5.18 基线中发现，TRTLLMMLABackend.forward_decode 已经把「bf16 -> fp8 量化 + KV scatter + [q_nope | q_rope] 拼接」融合进一次 set_mla_kv_concat_q_fp8 launch，但 target verification 走的是 forward_extend，在同一输入上仍无条件执行非融合链路——「That is five kernel launches per MLA layer instead of one」。该路径正是 cos_sin_cache is None 的 case（Kimi K3 的 skip_rope=True，absorbed 路径不做 RoPE），带 RoPE 的模型走 flashinfer 的 mla_rope_quantize_fp8 路径不受影响。issue 中的诊断数据显示 B300 TP8 dummy Kimi K3 + MTP1 + FP8 KV 下输出吞吐 106.45 -> 108.40 tok/s（+1.83%），TPOT 9.1597 -> 8.9923 ms（-1.83%），并明确提出「Would sharing the existing decode preparation path with eligible target-verify calls be an appropriate upstream change? The dispatch fix can be considered separately from selecting CUDA versus Triton for the fused implementation」——本 PR 采纳的正是其中的分发层修复，而不是新的 Triton kernel。

推荐精读，尤其是正在做 MLA / speculative verify 性能优化或统一内存池相关改动的同学。值得关注的设计决策有三点：(1) 用 fused_fp8_query 单哨兵同时控制「跳过重复 KV 写入」和「跳过 query concat」，比在多处重复判断 gate 更难写错；(2) 复用 forward_decode 的守卫序列（save_kv_cache + 静态 gate + _resolve_fused_write_loc）而不是新建路径，使 decode 与 verify 的融合语义天然一致；(3) 明确把 dispatch 修复与 kernel 选择（CUDA vs Triton）切开，压缩了评审面。

阅读时建议同步确认两件事：一是 PR body 承诺的 test/registered/kernel/attention/test_trtllm_mla_fp8_prepare.py 是否已在其他 PR 落地，否则需补上 verify 路径的回归测试；二是真实权重（Kimi K3 + trtllm_mla + fp8 KV + MTP）的精度与吞吐结论是否已补齐，因为当前所有数字都来自 acceptance length=1.0 的 dummy 诊断，不能直接当作上线依据。

讨论亮点

本 PR 的 review 通道非常干净：review_comments_count 为 0，唯一一条 review 来自 Qiaolin-Yu，状态 APPROVED、正文为空，说明维护者对「复用 decode 融合准备 + 完整回退」这一改动形态没有提出异议。实质讨论全部发生在关联 issue #39107 与其评论区：

设计取舍被显式拆开：issue 里作者提出「The dispatch fix can be considered separately from selecting CUDA versus Triton for the fused implementation」，本 PR 只做分发层修复，沿用既有 CUDA JIT kernel set_mla_kv_concat_q_fp8，不引入新的 stride-aware Triton 实现，降低了评审与回归面。
收益口径由作者在评论区补强：「Fused kernel saves 4us per layer. A verify forward has 24 MLA layers, which saves ~0.1ms in total per verify forward.」并贴了 before/after profiler 截图。这与 issue 表格中 rank-0 GPU target-verify median 8.6397 -> 8.4645 ms（-2.03%）、preparation launches/MLA layer/verify step 5 -> 1 相互印证。
数据可信边界被作者自己标注：issue 中明确「These are short, non-interleaved dummy-model diagnostics; acceptance length was 1.0 in both arms. They do not establish real-weight throughput, task quality, or PDL's independent contribution.」PR body 也留有「real-weight numbers to follow」。
CI 层面出现过噪声：PR 状态区显示 Extra 与 AMD ROCm 10 lane 为失败，作者连续两次 /rerun-failed-ci 并执行 /tag-and-rerun-ci，base lane 最终绿灯后合入。

实现拆解

触发入口：唯一改动文件 python/sglang/srt/layers/attention/trtllm_mla_backend.py，改动集中在 TRTLLMMLABackend.forward_extend 的 fp8 分支。触发条件是既有的三元组：self.data_type == torch.float8_e4m3fn、forward_batch.forward_mode.is_target_verify()、cos_sin_cache is None。前置的 flashinfer fallback 判断保持不变，仍显式排除 target-verify / draft-extend，避免 verify 被绑到 prefill 规划的 wrapper buffer。
新增哨兵变量与守卫式融合尝试：在 merge_query = q_rope is not None 之后新增 fused_fp8_query = None，作为「融合是否命中」的唯一判据。在 cos_sin_cache is None 分支内先判断 save_kv_cache and self._fused_set_kv_concat_q_fp8（既有的静态开关），再用新增的 _resolve_fused_write_loc(forward_batch) 解析写地址；只有 loc is not None 时才调用 _set_kv_and_concat_q_fp8_fused(layer, loc, q, q_rope, k, k_rope)。这与 forward_decode 的守卫完全同构，复用了 unified pool 下已为 TARGET_VERIFY 填好的 capture-stable 写位置 _decode_kernel_loc，因此不需要新 kernel，也不改动静态 gate。
回退链路完整保留：融合未命中时（写关闭、gate 关闭、解析不到 loc，或 covered_fp8 不接受布局导致 helper 返回 None），仍执行原来的 mla_quantize_without_rope_for_fp8(q, q_rope, k.squeeze(1), k_rope.squeeze(1))。KV 写入条件由 if save_kv_cache: 收紧为 if save_kv_cache and fused_fp8_query is None:，避免融合分支已经写过 set_mla_kv_concat_q_fp8 后再用 set_mla_kv_buffer 重复散写一次；query 准备由 if merge_query: 改为 if fused_fp8_query is not None: q = fused_fp8_query / elif merge_query:，即融合成功时直接采用 kernel 产出的 fp8 query，绕过 concat_mla_absorb_q_general。
影响面外溢到子类后端：CuteDslMLABackend 与 TokenspeedMLABackend 继承 forward_extend，因此它们的 verify 路径同样受益，包括 Kimi K3 默认 DCP 配置（通过 speculative_attention_mode="decode" 把 verify 路由到 cutedsl_mla）。draft-extend 与普通 prefill 不受影响，因为融合尝试只挂在 is_target_verify() 条件下；RoPE 模型继续走 mla_quantize_and_rope_for_fp8。
测试与验证配套：PR body 声称新增 test/registered/kernel/attention/test_trtllm_mla_fp8_prepare.py（基于 TestFusedFp8WriteGate 的 bare-backend fixture，断言 verify 走融合准备、在 translating pool 下写到 capture-stable loc、gate 关闭时保留非融合链路并对齐 aten fp8 参考），但最终 diff 只包含 1 个源码文件（+25/-7），该测试文件未体现在合并内容中，需要确认是否拆到其他 PR 或遗漏。已有契约由 test_kimi_k3_prerequisite_ops（batch 64）保证融合 kernel 与 aten 链路 byte-identical；PR body 明确标注 GPU 测试未本地运行、Kimi K3 + trtllm_mla + fp8 KV + MTP 精度跑尚未完成，accuracy/speed benchmark 清单项仍为 pending。

文件	模块	状态	重要度
`python/sglang/srt/layers/attention/trtllm_mla_backend.py`	注意力后端	modified	6.87

关键符号

TRTLLMMLABackend.forward_extend _resolve_fused_write_loc _set_kv_and_concat_q_fp8_fused mla_quantize_without_rope_for_fp8

关键源码片段

python/sglang/srt/layers/attention/trtllm_mla_backend.py core-logic

唯一改动文件，承载全部逻辑：TRTLLMMLABackend.forward_extend 的 fp8 + target verify 分支新增融合准备尝试（`_resolve_fused_write_loc` + `_set_kv_and_concat_q_fp8_fused`），并以 `fused_fp8_query` 哨兵控制跳过 KV 写入与 query concat，未命中时完整回退到 `mla_quantize_without_rope_for_fp8` + `set_mla_kv_buffer`。改动同时通过继承影响 CuteDslMLABackend / TokenspeedMLABackend 的 verify 路径。

# python/sglang/srt/layers/attention/trtllm_mla_backend.py
# TRTLLMMLABackend.forward_extend()：fp8 + target verify 分支（PR#39232 改造后）
def forward_extend(
    self,
    q, k, v,
    layer: RadixAttention,
    forward_batch: ForwardBatch,
    save_kv_cache: bool = True,
    q_rope=None,
    k_rope=None,
    cos_sin_cache=None, # None 表示 absorbed 路径不需要应用 RoPE（Kimi K3 skip_rope=True）
    is_neox: bool = False,
    llama_4_scaling=None,
):
    # 真正的 extend 才允许回退 flashinfer prefill 实现；verify / draft-extend
    # 不能命中，否则 CUDA Graph 会被绑到 prefill 规划的 wrapper buffer 上
    if (
        forward_batch.forward_mode.is_extend()
        and not forward_batch.forward_mode.is_target_verify()
        and not forward_batch.forward_mode.is_draft_extend_v2()
        and self.forward_prefill_metadata is not None
        and self.forward_prefill_metadata.fallback_to_flashinfer_impl
    ):
        return super().forward_extend(q, k, v, layer, forward_batch, save_kv_cache, q_rope, k_rope)

    merge_query = q_rope is not None
    # 本次 PR 新增的哨兵：非 None 表示融合路径已写出 KV cache 并产出 fp8 query
    fused_fp8_query = None
    if (self.data_type == torch.float8_e4m3fn) and forward_batch.forward_mode.is_target_verify():
        assert q_rope is not None and k_rope is not None
        if cos_sin_cache is None:
            # 与 forward_decode 完全同构的守卫：需要写 KV、静态 gate 打开、
            # 且能解析出写地址时才尝试融合
            if save_kv_cache and self._fused_set_kv_concat_q_fp8:
                # unified pool 下 TARGET_VERIFY 的 _decode_kernel_loc 已填充，
                # 复用它可让 capture / replay 的写地址保持稳定
                loc = self._resolve_fused_write_loc(forward_batch)
                if loc is not None:
                    # 一次 launch 完成 bf16->fp8 量化 + KV scatter + [q_nope | q_rope] 拼接
                    fused_fp8_query = self._set_kv_and_concat_q_fp8_fused(
                        layer=layer,
                        loc=loc,
                        q=q,
                        q_rope=q_rope,
                        k=k,
                        k_rope=k_rope,
                    ) # 布局不被 covered_fp8 接受时返回 None，走下面的回退
            if fused_fp8_query is None:
                # 回退：原有 5 次 launch 链路（concat + 3 次 aten fp8 cast + Triton scatter）
                q, k, k_rope = mla_quantize_without_rope_for_fp8(
                    q, q_rope, k.squeeze(1), k_rope.squeeze(1)
                )
        else:
            # 需要应用 RoPE 的模型继续走 flashinfer 的 mla_rope_quantize_fp8 路径
            q, k, k_rope = mla_quantize_and_rope_for_fp8(
                q, q_rope, k.squeeze(1), k_rope.squeeze(1),
                forward_batch.positions, cos_sin_cache, is_neox,
                self.kv_lora_rank, self.qk_rope_head_dim,
            )
        merge_query = False

    # 融合路径已经在 kernel 内散写过 KV，这里必须跳过，否则会重复写入
    if save_kv_cache and fused_fp8_query is None:
        assert k is not None and k_rope is not None, (
            "For populating trtllm_mla kv cache, both k_nope and k_rope should be not None."
        )
        if self._decode_kernel_loc is not None:
            self.token_to_kv_pool.set_mla_kv_buffer(layer, self._decode_kernel_loc, k, k_rope)
        else:
            self.token_to_kv_pool.set_mla_kv_buffer(layer, forward_batch.out_cache_loc, k, k_rope)

    # query 准备：融合命中时直接采用 kernel 产出的 fp8 query，绕过 concat
    if fused_fp8_query is not None:
        q = fused_fp8_query
    elif merge_query:
        # fp16 路径仍把 q_nope 与 q_rope 拼成完整 head_dim
        q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim)
        q_rope_reshaped = q_rope.view(-1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim)
        q = concat_mla_absorb_q_general(q_nope, q_rope_reshaped)

    q = q.view(-1, layer.tp_q_head_num, layer.head_dim)

评论区精华

是否把 decode 的融合准备直接共享给 target verify 设计

issue #39107 中作者提问「Would sharing the existing decode preparation path with eligible target-verify calls be an appropriate upstream change? The dispatch fix can be considered separately from selecting CUDA versus Triton for the fused implementation.」；PR body 进一步论证融合 kernel 本身与 decode 无关（per-token / per-(token, head)、q/k 行 stride 是运行时参数、capture-stable 写地址在 unified pool 下已为 TARGET_VERIFY 填好）。

结论：采纳：verify 复用同一条守卫式融合路径，且把 dispatch 修复与 kernel 选择（CUDA JIT vs stride-aware Triton）拆成两件事，本 PR 只做前者。 · 已解决

收益数字的口径与可信边界 性能

issue 表格给出 output throughput 106.45 -> 108.40 tok/s（+1.83%）、TPOT 9.1597 -> 8.9923 ms（-1.83%）、rank-0 GPU target-verify median 8.6397 -> 8.4645 ms（-2.03%）、preparation launches/MLA layer/verify step 5 -> 1；作者在评论区补充「Fused kernel saves 4us per layer. A verify forward has 24 MLA layers, which saves ~0.1ms in total per verify forward.」并附 before/after profiler 截图。

结论：方向一致且多口径互相印证，但全部来自 B300 TP8 dummy Kimi K3、acceptance length 固定 1.0 的诊断场景；作者明确声明不构成真实权重吞吐与任务质量结论，real-weight numbers 待补。 · partially-resolved

精度与 GPU 测试尚未闭环 测试

PR body 的 Checklist 把「Provide accuracy and speed benchmark results」标为 pending，并写明「the GPU tests have not been run locally yet, and a Kimi K3 + trtllm_mla + fp8 KV + MTP accuracy run is still pending」；同时 body 描述的 `test/registered/kernel/attention/test_trtllm_mla_fp8_prepare.py` 未出现在最终 diff（changed_files_count=1，仅源码文件 +25/-7）。

结论：唯一一条 review 由 Qiaolin-Yu 以 APPROVED 通过且无正文，改动在精度跑未完成、自带测试未见入库的情况下合入，后续需要补齐验证。 · 待处理

CI 噪声处理 other

PR 状态区显示 Latest PR Test (Extra) 与 AMD ROCm 10 两条 lane 失败，作者先 `/tag-and-rerun-ci`，随后两次 `/rerun-failed-ci` 推动重跑。

结论：base lane 最终通过并完成合并，失败判定为基础设施噪声而非改动本身；但 AMD lane 失败与本次 blackwell/fp8 相关改动之间的关系未被明确说明。 · 已解决

风险与影响

跳过 KV 写入的正确性风险（最高优先）：融合成功后代码不再执行 set_mla_kv_buffer，KV 行的落点完全由 _resolve_fused_write_loc(forward_batch) 决定。若在某个 verify 配置下该 loc 与 forward_batch.out_cache_loc 语义不一致（例如非 unified pool 路径、或 loc 解析到 _decode_kernel_loc 而后续 consumer 仍按 out_cache_loc 读取），会出现 KV 写错位置而静默产生精度问题，而非崩溃。守卫 loc is not None 只保证地址存在，不保证语义等价。
CUDA Graph capture/replay 一致性：verify 在 decode 图捕获路径上运行，融合是否命中是一个运行时判定；如果 _fused_set_kv_concat_q_fp8、_resolve_fused_write_loc 的返回值在 capture 与 replay 之间发生翻转，图内会执行到与捕获时不同的分支。代码注释已经提示过同类隐患（forward_prefill_metadata 是 prefill 的 stale leftover，捕获期不清空），_decode_kernel_loc 在 unified pool 下为 TARGET_VERIFY 填充的状态同样依赖这一约定。
测试覆盖与 diff 不一致：PR body 描述的 test/registered/kernel/attention/test_trtllm_mla_fp8_prepare.py 未出现在最终 diff（changed_files_count=1），若该测试没有随其他 PR 落地，则「verify 走融合准备」「gate 关闭时走回退」两条断言在主线缺少直接回归保护，只能依赖 test_kimi_k3_prerequisite_ops 的 batch 64 契约。
精度验证未闭环：作者自述 GPU 测试未本地运行、真实权重精度跑 pending，checklist 的 accuracy/speed benchmark 项未勾选；MTP 场景下 acceptance length 固定为 1.0 的 dummy 诊断不能代表真实接受长度下的端到端表现。
跨后端继承带来的隐性影响面：CuteDslMLABackend 与 TokenspeedMLABackend 继承同一 forward_extend，改动会同时作用于它们的 verify 路径（含 Kimi K3 默认 DCP 配置），影响面比「只改 trtllm_mla」的字面范围更大，但同样没有针对这些后端的专门验证。
收益依赖特定配置：收益只在 blackwell + fp8 KV + MLA + 无 RoPE（skip_rope=True）+ MTP 的组合上兑现，其他 MLA 模型走 flashinfer RoPE 路径完全不感知，不要按全局 1.8% 外推。

用户侧：使用 Kimi K3 等 MLA 模型、开启 fp8 KV 与 MTP 推测解码、且运行在 Blackwell（B300 等）上的用户，target verify 阶段每个 MLA 层的 preparation 从 5 次 kernel launch 降到 1 次；按作者评论区的口径，24 层 MLA 的一个 verify forward 约省 0.1 ms，端到端 TPOT 观测到约 1.8% 改善（dummy 模型口径）。带 RoPE 的 MLA 模型与普通 prefill 用户无感知。

系统侧：改动落在 single-file 的注意力后端分发逻辑上，不引入新 kernel、不动静态 gate、不改 schema/配置/CLI，因此部署与回滚成本低。但由于 forward_extend 被 CuteDsl / Tokenspeed 子类继承，实际生效范围覆盖多条 MLA verify 通路。

团队侧：这是 speculative-decoding + MLA 后端维护者需要关注的典型模式——「融合内核的可用性判定」与「普通 extend 语义」解耦，靠哨兵变量把三条后续动作（KV 写入、query 准备、回退量化）串成一致的开关。对后续做类似分发修复的人有直接参考价值；同时也把 _decode_kernel_loc 在 TARGET_VERIFY 下的可用性变成一条隐性契约，需要在统一内存池相关改动中持续维护。

核心路径变更跳过重复 KV 写入依赖写地址语义一致新增测试未体现在 diff 精度与真实权重验证未完成跨后端继承影响面扩大 CUDA Graph capture/replay 分支一致性

关联 Issue

#39107 [Feature] TRTLLM MLA target verification misses fused FP8 KV/Q preparation in forward_extend

完整报告

执行摘要

一句话：验证路径复用解码融合准备，MLA 每层 preparation 由 5 次 launch 降为 1 次
推荐动作：推荐精读，尤其是正在做 MLA / speculative verify 性能优化或统一内存池相关改动的同学。值得关注的设计决策有三点：(1) 用 fused_fp8_query 单哨兵同时控制「跳过重复 KV 写入」和「跳过 query concat」，比在多处重复判断 gate 更难写错；(2) 复用 forward_decode 的守卫序列（save_kv_cache + 静态 gate + _resolve_fused_write_loc）而不是新建路径，使 decode 与 verify 的融合语义天然一致；(3) 明确把 dispatch 修复与 kernel 选择（CUDA vs Triton）切开，压缩了评审面。

功能与动机

实现拆解

触发入口：唯一改动文件 python/sglang/srt/layers/attention/trtllm_mla_backend.py，改动集中在 TRTLLMMLABackend.forward_extend 的 fp8 分支。触发条件是既有的三元组：self.data_type == torch.float8_e4m3fn、forward_batch.forward_mode.is_target_verify()、cos_sin_cache is None。前置的 flashinfer fallback 判断保持不变，仍显式排除 target-verify / draft-extend，避免 verify 被绑到 prefill 规划的 wrapper buffer。
新增哨兵变量与守卫式融合尝试：在 merge_query = q_rope is not None 之后新增 fused_fp8_query = None，作为「融合是否命中」的唯一判据。在 cos_sin_cache is None 分支内先判断 save_kv_cache and self._fused_set_kv_concat_q_fp8（既有的静态开关），再用新增的 _resolve_fused_write_loc(forward_batch) 解析写地址；只有 loc is not None 时才调用 _set_kv_and_concat_q_fp8_fused(layer, loc, q, q_rope, k, k_rope)。这与 forward_decode 的守卫完全同构，复用了 unified pool 下已为 TARGET_VERIFY 填好的 capture-stable 写位置 _decode_kernel_loc，因此不需要新 kernel，也不改动静态 gate。
回退链路完整保留：融合未命中时（写关闭、gate 关闭、解析不到 loc，或 covered_fp8 不接受布局导致 helper 返回 None），仍执行原来的 mla_quantize_without_rope_for_fp8(q, q_rope, k.squeeze(1), k_rope.squeeze(1))。KV 写入条件由 if save_kv_cache: 收紧为 if save_kv_cache and fused_fp8_query is None:，避免融合分支已经写过 set_mla_kv_concat_q_fp8 后再用 set_mla_kv_buffer 重复散写一次；query 准备由 if merge_query: 改为 if fused_fp8_query is not None: q = fused_fp8_query / elif merge_query:，即融合成功时直接采用 kernel 产出的 fp8 query，绕过 concat_mla_absorb_q_general。
影响面外溢到子类后端：CuteDslMLABackend 与 TokenspeedMLABackend 继承 forward_extend，因此它们的 verify 路径同样受益，包括 Kimi K3 默认 DCP 配置（通过 speculative_attention_mode="decode" 把 verify 路由到 cutedsl_mla）。draft-extend 与普通 prefill 不受影响，因为融合尝试只挂在 is_target_verify() 条件下；RoPE 模型继续走 mla_quantize_and_rope_for_fp8。
测试与验证配套：PR body 声称新增 test/registered/kernel/attention/test_trtllm_mla_fp8_prepare.py（基于 TestFusedFp8WriteGate 的 bare-backend fixture，断言 verify 走融合准备、在 translating pool 下写到 capture-stable loc、gate 关闭时保留非融合链路并对齐 aten fp8 参考），但最终 diff 只包含 1 个源码文件（+25/-7），该测试文件未体现在合并内容中，需要确认是否拆到其他 PR 或遗漏。已有契约由 test_kimi_k3_prerequisite_ops（batch 64）保证融合 kernel 与 aten 链路 byte-identical；PR body 明确标注 GPU 测试未本地运行、Kimi K3 + trtllm_mla + fp8 KV + MTP 精度跑尚未完成，accuracy/speed benchmark 清单项仍为 pending。

关键文件：

python/sglang/srt/layers/attention/trtllm_mla_backend.py（模块注意力后端；类别 source；类型 core-logic；符号 TRTLLMMLABackend.forward_extend, _resolve_fused_write_loc, _set_kv_and_concat_q_fp8_fused, mla_quantize_without_rope_for_fp8）: 唯一改动文件，承载全部逻辑：TRTLLMMLABackend.forward_extend 的 fp8 + target verify 分支新增融合准备尝试（_resolve_fused_write_loc + _set_kv_and_concat_q_fp8_fused），并以 fused_fp8_query 哨兵控制跳过 KV 写入与 query concat，未命中时完整回退到 mla_quantize_without_rope_for_fp8 + set_mla_kv_buffer。改动同时通过继承影响 CuteDslMLABackend / TokenspeedMLABackend 的 verify 路径。

关键符号：TRTLLMMLABackend.forward_extend, _resolve_fused_write_loc, _set_kv_and_concat_q_fp8_fused, mla_quantize_without_rope_for_fp8

关键源码片段

`python/sglang/srt/layers/attention/trtllm_mla_backend.py`

唯一改动文件，承载全部逻辑：TRTLLMMLABackend.forward_extend 的 fp8 + target verify 分支新增融合准备尝试（_resolve_fused_write_loc + _set_kv_and_concat_q_fp8_fused），并以 fused_fp8_query 哨兵控制跳过 KV 写入与 query concat，未命中时完整回退到 mla_quantize_without_rope_for_fp8 + set_mla_kv_buffer。改动同时通过继承影响 CuteDslMLABackend / TokenspeedMLABackend 的 verify 路径。

# python/sglang/srt/layers/attention/trtllm_mla_backend.py
# TRTLLMMLABackend.forward_extend()：fp8 + target verify 分支（PR#39232 改造后）
def forward_extend(
    self,
    q, k, v,
    layer: RadixAttention,
    forward_batch: ForwardBatch,
    save_kv_cache: bool = True,
    q_rope=None,
    k_rope=None,
    cos_sin_cache=None, # None 表示 absorbed 路径不需要应用 RoPE（Kimi K3 skip_rope=True）
    is_neox: bool = False,
    llama_4_scaling=None,
):
    # 真正的 extend 才允许回退 flashinfer prefill 实现；verify / draft-extend
    # 不能命中，否则 CUDA Graph 会被绑到 prefill 规划的 wrapper buffer 上
    if (
        forward_batch.forward_mode.is_extend()
        and not forward_batch.forward_mode.is_target_verify()
        and not forward_batch.forward_mode.is_draft_extend_v2()
        and self.forward_prefill_metadata is not None
        and self.forward_prefill_metadata.fallback_to_flashinfer_impl
    ):
        return super().forward_extend(q, k, v, layer, forward_batch, save_kv_cache, q_rope, k_rope)

    merge_query = q_rope is not None
    # 本次 PR 新增的哨兵：非 None 表示融合路径已写出 KV cache 并产出 fp8 query
    fused_fp8_query = None
    if (self.data_type == torch.float8_e4m3fn) and forward_batch.forward_mode.is_target_verify():
        assert q_rope is not None and k_rope is not None
        if cos_sin_cache is None:
            # 与 forward_decode 完全同构的守卫：需要写 KV、静态 gate 打开、
            # 且能解析出写地址时才尝试融合
            if save_kv_cache and self._fused_set_kv_concat_q_fp8:
                # unified pool 下 TARGET_VERIFY 的 _decode_kernel_loc 已填充，
                # 复用它可让 capture / replay 的写地址保持稳定
                loc = self._resolve_fused_write_loc(forward_batch)
                if loc is not None:
                    # 一次 launch 完成 bf16->fp8 量化 + KV scatter + [q_nope | q_rope] 拼接
                    fused_fp8_query = self._set_kv_and_concat_q_fp8_fused(
                        layer=layer,
                        loc=loc,
                        q=q,
                        q_rope=q_rope,
                        k=k,
                        k_rope=k_rope,
                    ) # 布局不被 covered_fp8 接受时返回 None，走下面的回退
            if fused_fp8_query is None:
                # 回退：原有 5 次 launch 链路（concat + 3 次 aten fp8 cast + Triton scatter）
                q, k, k_rope = mla_quantize_without_rope_for_fp8(
                    q, q_rope, k.squeeze(1), k_rope.squeeze(1)
                )
        else:
            # 需要应用 RoPE 的模型继续走 flashinfer 的 mla_rope_quantize_fp8 路径
            q, k, k_rope = mla_quantize_and_rope_for_fp8(
                q, q_rope, k.squeeze(1), k_rope.squeeze(1),
                forward_batch.positions, cos_sin_cache, is_neox,
                self.kv_lora_rank, self.qk_rope_head_dim,
            )
        merge_query = False

    # 融合路径已经在 kernel 内散写过 KV，这里必须跳过，否则会重复写入
    if save_kv_cache and fused_fp8_query is None:
        assert k is not None and k_rope is not None, (
            "For populating trtllm_mla kv cache, both k_nope and k_rope should be not None."
        )
        if self._decode_kernel_loc is not None:
            self.token_to_kv_pool.set_mla_kv_buffer(layer, self._decode_kernel_loc, k, k_rope)
        else:
            self.token_to_kv_pool.set_mla_kv_buffer(layer, forward_batch.out_cache_loc, k, k_rope)

    # query 准备：融合命中时直接采用 kernel 产出的 fp8 query，绕过 concat
    if fused_fp8_query is not None:
        q = fused_fp8_query
    elif merge_query:
        # fp16 路径仍把 q_nope 与 q_rope 拼成完整 head_dim
        q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim)
        q_rope_reshaped = q_rope.view(-1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim)
        q = concat_mla_absorb_q_general(q_nope, q_rope_reshaped)

    q = q.view(-1, layer.tp_q_head_num, layer.head_dim)

评论区精华

设计取舍被显式拆开：issue 里作者提出「The dispatch fix can be considered separately from selecting CUDA versus Triton for the fused implementation」，本 PR 只做分发层修复，沿用既有 CUDA JIT kernel set_mla_kv_concat_q_fp8，不引入新的 stride-aware Triton 实现，降低了评审与回归面。
收益口径由作者在评论区补强：「Fused kernel saves 4us per layer. A verify forward has 24 MLA layers, which saves ~0.1ms in total per verify forward.」并贴了 before/after profiler 截图。这与 issue 表格中 rank-0 GPU target-verify median 8.6397 -> 8.4645 ms（-2.03%）、preparation launches/MLA layer/verify step 5 -> 1 相互印证。
数据可信边界被作者自己标注：issue 中明确「These are short, non-interleaved dummy-model diagnostics; acceptance length was 1.0 in both arms. They do not establish real-weight throughput, task quality, or PDL's independent contribution.」PR body 也留有「real-weight numbers to follow」。
CI 层面出现过噪声：PR 状态区显示 Extra 与 AMD ROCm 10 lane 为失败，作者连续两次 /rerun-failed-ci 并执行 /tag-and-rerun-ci，base lane 最终绿灯后合入。
是否把 decode 的融合准备直接共享给 target verify (design): 采纳：verify 复用同一条守卫式融合路径，且把 dispatch 修复与 kernel 选择（CUDA JIT vs stride-aware Triton）拆成两件事，本 PR 只做前者。
收益数字的口径与可信边界 (performance): 方向一致且多口径互相印证，但全部来自 B300 TP8 dummy Kimi K3、acceptance length 固定 1.0 的诊断场景；作者明确声明不构成真实权重吞吐与任务质量结论，real-weight numbers 待补。
精度与 GPU 测试尚未闭环 (testing): 唯一一条 review 由 Qiaolin-Yu 以 APPROVED 通过且无正文，改动在精度跑未完成、自带测试未见入库的情况下合入，后续需要补齐验证。
CI 噪声处理 (other): base lane 最终通过并完成合并，失败判定为基础设施噪声而非改动本身；但 AMD lane 失败与本次 blackwell/fp8 相关改动之间的关系未被明确说明。

风险与影响

风险：
1. 跳过 KV 写入的正确性风险（最高优先）：融合成功后代码不再执行 set_mla_kv_buffer，KV 行的落点完全由 _resolve_fused_write_loc(forward_batch) 决定。若在某个 verify 配置下该 loc 与 forward_batch.out_cache_loc 语义不一致（例如非 unified pool 路径、或 loc 解析到 _decode_kernel_loc 而后续 consumer 仍按 out_cache_loc 读取），会出现 KV 写错位置而静默产生精度问题，而非崩溃。守卫 loc is not None 只保证地址存在，不保证语义等价。
2. CUDA Graph capture/replay 一致性：verify 在 decode 图捕获路径上运行，融合是否命中是一个运行时判定；如果 _fused_set_kv_concat_q_fp8、_resolve_fused_write_loc 的返回值在 capture 与 replay 之间发生翻转，图内会执行到与捕获时不同的分支。代码注释已经提示过同类隐患（forward_prefill_metadata 是 prefill 的 stale leftover，捕获期不清空），_decode_kernel_loc 在 unified pool 下为 TARGET_VERIFY 填充的状态同样依赖这一约定。
3. 测试覆盖与 diff 不一致：PR body 描述的 test/registered/kernel/attention/test_trtllm_mla_fp8_prepare.py 未出现在最终 diff（changed_files_count=1），若该测试没有随其他 PR 落地，则「verify 走融合准备」「gate 关闭时走回退」两条断言在主线缺少直接回归保护，只能依赖 test_kimi_k3_prerequisite_ops 的 batch 64 契约。
4. 精度验证未闭环：作者自述 GPU 测试未本地运行、真实权重精度跑 pending，checklist 的 accuracy/speed benchmark 项未勾选；MTP 场景下 acceptance length 固定为 1.0 的 dummy 诊断不能代表真实接受长度下的端到端表现。
5. 跨后端继承带来的隐性影响面：CuteDslMLABackend 与 TokenspeedMLABackend 继承同一 forward_extend，改动会同时作用于它们的 verify 路径（含 Kimi K3 默认 DCP 配置），影响面比「只改 trtllm_mla」的字面范围更大，但同样没有针对这些后端的专门验证。
6. 收益依赖特定配置：收益只在 blackwell + fp8 KV + MLA + 无 RoPE（skip_rope=True）+ MTP 的组合上兑现，其他 MLA 模型走 flashinfer RoPE 路径完全不感知，不要按全局 1.8% 外推。
  - 影响：用户侧：使用 Kimi K3 等 MLA 模型、开启 fp8 KV 与 MTP 推测解码、且运行在 Blackwell（B300 等）上的用户，target verify 阶段每个 MLA 层的 preparation 从 5 次 kernel launch 降到 1 次；按作者评论区的口径，24 层 MLA 的一个 verify forward 约省 0.1 ms，端到端 TPOT 观测到约 1.8% 改善（dummy 模型口径）。带 RoPE 的 MLA 模型与普通 prefill 用户无感知。

风险标记：核心路径变更, 跳过重复 KV 写入依赖写地址语义一致, 新增测试未体现在 diff, 精度与真实权重验证未完成, 跨后端继承影响面扩大, CUDA Graph capture/replay 分支一致性

关联脉络

PR #39171 [DeepSeek-V4.1] Bump FlashMLA to the fork's rebase head (v4.1 kernels): 同属 Blackwell 上的 MLA 注意力内核演进线：该 PR 升级 FlashMLA 内核并为 B300 补原生编译，本 PR 则在 trtllm_mla 后端复用融合 fp8 preparation，两者共同指向 MLA decode/verify 路径的 kernel 效率优化。
PR #39219 [Fix] Don't write conv state from the fused KDA verify kernel: 同一类「融合内核在 verify 路径改变写行为」的改动，主题高度对称：本 PR 让融合路径接管 KV 写入、并靠哨兵跳过重复写；该 PR 则移除 verify 融合内核的死回写以修复 CTA 调度竞态，可作为写路径变更的风险参照。
PR #37914 [Unified Cache][7/N] Support MTP, EAGLE, and DSpark draft KV caches in the external linker: 涉及 MTP/EAGLE 等推测解码的 KV 打包与校验，与本 PR 的 target verify + MTP1 场景（24 个 MLA 层的 verify preparation）共享同一套 speculative-decoding KV 处理语境。
PR #37418 [unified-memory] Enable prefill cuda-graph capture: 统一内存池与 CUDA Graph 抓取相关改动，本 PR 依赖的关键前提之一正是「unified pool 下 TARGET_VERIFY 已填好 capture-stable 的 _decode_kernel_loc」，两者的内存池/图捕获契约需要协同维护。

#39232 [Perf] trtllm_mla: reuse the fused fp8 KV/Q prepare on target verify

执行摘要

验证路径复用解码融合准备，MLA 每层 preparation 由 5 次 launch 降为 1 次

实现拆解

评论区精华

风险与影响

关联 Issue

完整报告

参与讨论