[Perf] SM8x sparse-MLA prefill fallback re-reads the shared MLA latent once per head (~64x redundant KV traffic)

快速结论：当你看到 [Perf] SM8x sparse-MLA prefill fallback re-reads the shared MLA latent once per head (~64x redundant KV traffic) 这类性能描述时，通常发生在 Ampere（sm_80）显卡上用 vLLM 跑 DSA / sparse-MLA 长 prompt prefill 的场景，且性能瓶颈被定位到 _sparse_mla_fwd_with_sink_kernel。优先排查你的代码树里是否真的存在这个 kernel 文件，以及它来自哪里（上游还是本地补丁）。

适用环境：硬件 2x A100-PCIE-40GB（compute capability 8.0），tensor parallel=2；模型 zai-org/GLM-5.3-Flash（DSA / is_v32 sparse-MLA 路径，index_topk=2048，num_attention_heads=64，kv_lora_rank=512，qk_rope_head_dim=0，NoPE 布局）；KV cache 为 bf16；工作负载为单请求 16384-token 随机 prompt，并发 1；vLLM 为某个包含 vllm/v1/attention/backends/mla/sparse_mla_kernels.py 的本地 checkout。

最快修复方案：暂无确认的一步修复方案。Issue 最终被作者撤回，结论是这些 kernel 文件并非上游 vLLM 代码，而是本地 SM80 port 补丁集（pr3-sm80-port.patch）新增的文件，因此上游无法直接修复。

注意事项：Issue 中的性能数字、64x 冗余结论、以及“上游为 SM8x 提供 portable Triton fallback”的前提均已被作者本人撤回或修正。若你遇到同类现象，先确认 kernel 来源，再考虑是否值得向对应补丁维护者反馈。

问题场景

用户在 Ampere（sm_80）显卡上用 vLLM 跑 GLM-5.3-Flash（DSA / is_v32 sparse-MLA 路径）的长 prompt prefill，单请求 16384 token，开启 tensor parallel=2。profiling 发现 _sparse_mla_fwd_with_sink_kernel 占了 58.8% 的 CUDA 时间（单次 prefill，9 次调用），TTFT 约 56 秒。用户最初认为该 kernel 对共享的 MLA latent KV 每个 head 重复读取一次，造成约 64x 冗余 KV 流量，并判断这是主要性能瓶颈。

报错原文

[Perf] SM8x sparse-MLA prefill fallback re-reads the shared MLA latent once per head (~64x redundant KV traffic)

原因分析

最初的判断是：kernel 的 grid 为 (num_tokens, active_heads)，head_idx 只出现在 q 和 output 的地址计算里，不出现在 KV 地址计算里，因此每个 head 的 program 都会重新加载同一段共享的 MLA latent KV，理论上造成 num_heads 倍的冗余流量。

但作者随后自己做了 head-blocked 变体测试，结论发生反转：

冗余确实存在，但不是瓶颈。BLOCK_H=1（完全不做 KV 复用，行为与原始一致）仍然比原始快 2.78x；继续加大到 BLOCK_H=8 只再多约 3%。说明该 kernel 不是 KV-bandwidth bound。
真正的性能驱动可能是 tile 形状。原始 kernel 把 q 加载为 1-D (BLOCK_D,) 向量、running_acc 也是 1-D；改成 2-D 形式本身就能带来约 2.8x 提升，具体原因作者也不清楚，属于经验可复现但缺乏机理解释。
中间值非单调。BLOCK_H=4 比 1、2、8 都慢，疑似 occupancy 效应，不宜据此读趋势；BLOCK_H=16 明显劣化（0.58x），疑似 register spilling。
数值差异。各变体相对原始 kernel 均非 bit-exact，max|diff| 在小 shape 为 1.965e-05，真实 shape 为 2.819e-05，来自 reduction 顺序不同（fp32 累加 bf16 输入，作者认为可接受）；BLOCK_H=16 为 1.170e-03，约差 60 倍，疑似 spilling 所致。

最后，维护者指出 sparse_mla_kernels.py 在 main 上返回 404，代码搜索 _sparse_mla_fwd_with_sink_kernel 无结果；main 上的 sparse MLA 文件是 flashmla_sparse.py、flashinfer_mla_sparse.py（含 _sm90 / _sm120 变体）、flashattn_mla_sparse.py、rocm_aiter_mla_sparse.py。作者确认这些文件来自本地 SM80 port 补丁集 pr3-sm80-port.patch（新增 sparse_mla_kernels.py，3517 行，另有 sparse_mla_env.py），并非上游代码，因此撤回了整个 Issue。

环境排查

确认你的 vLLM checkout 中是否存在 vllm/v1/attention/backends/mla/sparse_mla_kernels.py；在 upstream main 上该文件不存在（返回 404）。
用代码搜索确认 _sparse_mla_fwd_with_sink_kernel 是否命中；upstream main 上无命中。
确认当前 tree 是否应用了本地/第三方的 SM80 port 补丁集（例如 pr3-sm80-port.patch），并核对补丁引入的文件清单。
确认硬件是否为 compute capability 8.0（Ampere），因为该 Triton fallback 路径与 SM8x 相关。
确认 upstream main 上现有 sparse MLA 后端文件（flashmla_sparse.py、flashinfer_mla_sparse.py 的 _sm90 / _sm120 变体、flashattn_mla_sparse.py、rocm_aiter_mla_sparse.py），判断你的问题是否只存在于本地补丁路径。
若使用 tensor parallel，注意 active_heads 是 pre-shard 还是 post-shard 计数，会影响冗余倍数的估算（TP=2 时作者按 post-shard 应为 32 而非 64，但该点未被最终确认）。

解决步骤

先在本地确认 kernel 来源：检查 sparse_mla_kernels.py 与 flashmla_sparse_sm8x.py 是否为补丁新增文件（例如查看补丁头部是否有 new file mode 100644）。
如果这些文件来自本地 SM80 port 补丁集，则不要向 vLLM 上游提 Issue 或 PR，改为向补丁集维护者反馈。
可优先尝试在本地 kernel 上做 shape 变换实验：把 q 与 running_acc 改为二维 (BLOCK_H, BLOCK_D) tile，配合 BLOCK_H 取值 1、2、8 做对比测试；作者实测该改动可将 16K prefill 上 kernel 时间从 22.3 s 降到约 8 s。
避免直接采用 BLOCK_H=4（非单调劣化）与 BLOCK_H=16（疑似 spilling、数值偏差约 60 倍）这两档。
若无法解释 2-D 形式为何更快，建议先以 stopwatch 数据记录，作为后续进一步分析或向 Triton 相关方请教的依据，而不是直接合入无法解释的性能补丁。

验证方法

在相同 profiler window、相同 prompt 长度与参数下对比 _sparse_mla_fwd_with_sink_kernel 的 CUDA 耗时。作者在 standalone kernel harness 上的对照结果可作为参考：原始 311.208 ms（1.00x）、BLOCK_H=1 111.853 ms（2.78x）、BLOCK_H=2 108.929 ms（2.86x）、BLOCK_H=4 135.885 ms（2.29x）、BLOCK_H=8 108.185 ms（2.88x）、BLOCK_H=16 532.517 ms（0.58x）。同时需检查 max|diff| 是否落在可接受范围（小 shape 1.965e-05、真实 shape 2.819e-05）。注意以上均为 kernel harness 结果，作者本人也未确认 end-to-end 表现，也未验证是否能迁移到 SM90 的 flash_mla_sparse_fwd 路径。