[RFC]: DeepSeek-V4.1-Flash performance on ROCm

快速结论：该 RFC 描述的是 DeepSeek-V4.1-Flash 在 AMD MI355X（gfx950）上跑 vLLM 时的性能问题，而非硬报错——单条解码步会发起约 14,020 次 kernel launch，其中 85% 来自 mHC seam 的 4×4 矩阵 Sinkhorn 记账。优先排查方向是把 mHC pre/post 从逐 op 的 PyTorch 实现切到 AITER/融合实现，并用带 CUDA graph 的正常服务跑分，而不是用 `–enforce-eager` 的 trace 数字下结论。

适用环境：Issue 已确认的环境为 vLLM、AMD MI355X（gfx950）单节点 8 卡、TP4、MXFP4 MoE + DSpark MTP、ROCm（roctracer 介入 HSA 队列时与 hipGraph replay 死锁的某 ROCm build）、TileLang 0.1.10、`main` commit `c1b69aa0d4`；评测使用 lm-eval gsm8k 5-shot，harness gate 0.91。

最快修复方案：暂无确认的一步修复方案。Issue 中已验证并合入的是 #56503（mHC pre 走 AITER，每个 seam 从 141 次 launch 降到 4 次，最差相对误差 1.2e-5，gsm8k 不变）；#56513（融合 mHC post+pre，conc 1 下平均 ITL 降低 1.28%，95% CI [0.07%, 2.48%]）在 Issue 关闭时为 draft 状态。

注意事项：启用 CUDA graph 时不要在该 ROCm build 上做 roctracer 式 profiling，`/start_profile` 后的第一步不会完成，worker 会卡在 `queue_interposition.cpp` 等待 HSA signal；trace 中只有 launch 数量和 kernel 组成是 graph-invariant 可用的，墙钟时间、占用率、idle gap 都因 eager dispatch 加 profiler 开销而偏高，不能用来做结论。吞吐和 ITL 数字来自启用 graph 的正常服务运行。另外该节点两半 GPU 对相同工作有 0.96% 差异，小于约 3% 的收益必须用 crossover 或交错 A/B 设计测量，顺序 A/B 会得出错误结论。

问题场景

用户在 8x MI355X（gfx950）节点上以 TP4 运行 DeepSeek-V4.1-Flash，MoE 使用 MXFP4 量化并带 DSpark MTP，在并发从 1 升到 32 时，输出吞吐从 35.89 tok/s 涨到 469.41 tok/s，但 ITL p50 只从 0.0256 恶化到 0.0477，即并发提升 32 倍只付出 1.9x 的 ITL 代价，说明设备远未饱和，每个 decode step 存在一笔与并发无关的固定开销。抓取 decode trace 后定位到：单流 ITL 25.6 ms 的一步里发出 14,020 次 kernel launch，其中每层 350 次里有 85% 属于 mHC block，而 mHC 只是在处理一个 4×4 矩阵。

报错原文

[RFC]: DeepSeek-V4.1-Flash performance on ROCm

At a single-stream ITL of 25.6 ms, a decode step issues 14,020 kernel launches,
and 85% of each layer's 350 launches are attributable to the mHC block -
bookkeeping on a 4x4 matrix.

launches/layer | share | source
162 | 48.6% | mHC sinkhorn loop
122 | 36.6% | mHC pre/post mix + eager RMSNorm
24  | 7.2%  | MoE / FFN (already fused)
12  | 3.6%  | dtype casts/copies
6   | 1.8%  | attention + fused norm/quant
4   | 1.2%  | TP collective
3   | 0.9%  | other

注意：Issue 本身是 RFC/性能记录，上面这段是 Issue 正文的核心度量描述，不是抛出的异常栈。

原因分析

最可能的原因是 DeepSeek-V4.1 的 AMD 实现没有跟随 V4 的算子迁移路径。`vllm/models/deepseek_v4_1/amd/model.py` 直接 import `mhc_pre_delayed_torch` / `mhc_post_torch`，因此每个 sublayer seam 都要花约 141 次 launch 去对 4×4 矩阵做 Sinkhorn 归一化。V4 已经通过 `MHCPreOp` / `MHCPostOp` 分发并走到 AITER，而 V4.1 从未迁移，原因是 V4.1 用上一个 sublayer 的 pre-mix 来折叠当前 sublayer 的输入，而 AITER 的 `mhc_pre` 既不返回自己的 pre-mix，也不接受传入的 pre-mix。

其次是 eager RMSNorm 和未融合的 mHC seam 本身：Issue 中提到，未融合的 mHC seam 从不请求 `store_nt` 或 `large_m_splitk`，这在 decode 阶段没有效果，但在 prefill 阶段是实打实的损失（列为跟踪项 5）。

此外，Issue 后段的 nightly profiling 显示，mHC 之外的第二大块是 sparse indexer（DSA），并发 1/4/16 下占 decode step 的 4.1%/6.2%/11.4%，而它并不在最初的 RFC 覆盖范围内——如果只优化 mHC seam 后仍然觉得慢，这可能原因就在 indexer 上。

环境排查

确认 GPU 型号与架构：AMD MI355X，gfx950；Issue 为单节点 8 卡。
确认并行配置：TP4。
确认模型与量化：DeepSeek-V4.1-Flash，MXFP4 MoE + DSpark MTP。
确认 vLLM 分支/commit：Issue 中在未打补丁的 `main` 上用 `c1b69aa0d4` 复现过 TileLang 相关失败。
确认 TileLang 版本：pinned tilelang 0.1.10。
确认 ROCm build 是否与 roctracer 的 HSA 队列介入、hipGraph replay 存在死锁；如存在，profiling 时不能开 CUDA graph。
确认基准方法：lm-eval gsm8k 5-shot，harness gate 0.91；Issue 基线 strict-match 0.9704 ± 0.0047，flexible-extract 0.9697 ± 0.0047。
确认 A/B 测量设计：该 8 卡节点两半对相同工作有 0.96% 差异，小于约 3% 的收益需 crossover 或交错设计。

解决步骤

先用启用 CUDA graph 的正常服务方式复现吞吐与 ITL，记录 concurrency 1/2/4/8/16/32 下的 out tok/s、out tok/s per GPU、TTFT p50、ITL p50、E2EL p50，作为后续对比基线。
确认当前分支是否已包含 #56503（mHC pre -> AITER）。若未包含，按该 PR 将 V4.1 的 seam 迁移到 `MHCPreOp` / `MHCPostOp` 并走 AITER，同时正确传入上一 sublayer 的 pre-mix。合入后每个 seam 的 launch 从 141 降到 4，最差相对误差 1.2e-5。
在 #56503 之上评估 #56513（fused mHC post+pre，draft，stacked on #56503）。Issue 中在 rebase 到 `main` 后为单 commit，实测 conc 1 下平均 ITL 降低 1.28%，95% CI [0.07%, 2.48%]，gfx950 上 decode 尺寸 batch 每个 seam 从 5 次 launch 降到 4 次。
不要重复做 1.3（fused all-reduce + mHC post）。Issue 作者已 unclaim，原因是它与 1.2 是二选一而非叠加：两者消耗同一个 post step，各自只值六次 launch 里的一次，且 AITER 没有同时覆盖 AR+post+pre 的 kernel；all-reduce 变体还需要把未归约的 attention 输出传出来，而它今天在 `_o_proj` 内部被消费。
测量任何改动时使用 crossover 或交错 A/B。Issue 中一次单方向 A/B 把 #56513 的效果测成 +0.32%，交换两半 GPU 后为 +2.20%，消除 GPU 组项后的 crossover 估计才是 1.28%。
如果 mHC seam 优化后仍慢，按 Issue 后段的 nightly profiling 拆分 indexer 各子项（mqa logits、candidate mask、candidate select、topk per row），确认是否落在 DSA sparse indexer 上；Issue 显示 conc 16 时 indexer total 为 2.954 ms，占 decode step 11.4%。
若需要处理 prefill，检查未融合 mHC seam 是否请求了 `store_nt` 和 `large_m_splitk`（跟踪项 5）；decode 下该项无效，prefill 下是缺失的优化。
注意独立于本 RFC 的一个已有缺陷：`mhc_pre_delayed_tilelang` 带 `norm_weight` 时，RMS 平方和在 thread-divergent branch 内归约并只收回一半，导致 `layer_input` 高出 √2 倍。该问题在未打补丁的 `main`（`c1b69aa0d4`）、gfx950、pinned tilelang 0.1.10 上，从 hidden size 512 到 7168 全部复现，mix 输出不受影响。DSV4.1 从不请求该 seam 的 fused norm，因此本 RFC 不依赖它，但 `test_deepseek_v41_mhc_pre_delayed[fused_norm=True]` 会在 `main` 上独立失败。

验证方法

在启用 CUDA graph 的正常服务运行下，对比改动前后的 out tok/s、out tok/s per GPU、TTFT p50、ITL p50、E2EL p50；同时在 lm-eval gsm8k 5-shot 上对比 strict-match 与 flexible-extract，不能只看是否过了 0.91 的 harness gate，而要与 Issue 中给出的基线数值（0.9704 / 0.9697，以及 #56503 后的 0.9727 / 0.9719、#56503+#56513 后的 0.9735 / 0.9727）比对，因为一个改动可能过了 0.91 但已经破坏了别的东西。若用 trace 辅助定位，只读取 launch 数量和 kernel 组成，并确认 profiling 未开启 CUDA graph。