#37018 [Kernel] Fix SM90 FP8 decode regression with benchmarked M/K/N routing

原始 PR 作者 RunFMe 合并时间 2026-08-30 07:59 文件变更 1 提交数 4 评论 3 代码增减 +4 / -5

执行摘要

限流 SM90 FP8 路由至大 M 包络，修复 decode 回归

PR body 明确指出 “The SM90 row/column-scaled FP8 selector introduced by #34318 can route decode GEMMs to torch._scaled_mm based on K/N alone”，带来 H100 W8A8 回归：模型 Llama-3.1-8B-Instruct-FP8-dynamic、batch size 1、吞吐约 215 → 约 190 tok/s，受影响下投影为 M=1, K=14336, N=4096。直接剖析显示该形状下 AOT 远快于 Torch（H200 上 26.50 µs 对 51.12 µs，AOT 快 1.93 倍；H100 上 32.19 µs 对 53.55 µs，AOT 快 1.66 倍），说明 decode 的 M 过小，Torch 后端的启动与张量化开销无法被大 K/N 掩盖，路由必须以 M 为先决条件。

值得精读。虽然仅 9 行改动，但它是“基准驱动路由选择”的教科书案例：先用宽网格确定优势面，再用随机化边界确认锁定阈值，最后对噪声区域采取保守策略。内核维护者和 FP8 性能工程师可以复刻这套验证协议；普通使用者了解结论即可——decode 一律走 AOT，大 prefill 宽投影走 Torch。建议后续为路由谓词补充形状注释表与轻量单测，防止回归。

讨论亮点

本 PR 的讨论非常收敛：唯一正式 review 来自 BBuf（APPROVED）：“It seems resonable, approved!”，评审未提出技术质疑，改动以 PR 内大量基准数据自证。另一条线程是 CI 复验：通过 /rerun-test test/registered/quant/test_w8a8_quantization.py 在 1-gpu-h100 上通过，但 pr-test-extra 与 AMD ROCm 网格状态为失败且 PR 内未解释原因。

实现拆解

变更入口：唯一改动文件为 python/sglang/kernels/ops/gemm/__init__.py，修改 Fp8ScaledMMOp 依赖的 rowwise FP8 Torch/AOT 路由选择器 _prefer_torch_rowwise_fp8。
谓词重构：旧谓词 (k >= 5376 and n >= 3584) or (k >= 3584 and m >= 8192) 的第一项完全不看 M，第二项也只是把 m >= 8192 与 k >= 3584 简单叠加；新谓词改为 m >= 8192 and ((k >= 4096 and n >= 6144) or (k >= 7168 and n >= 5376))，M 成为硬性前置门槛，K/N 包络按“宽输出投影”（QKV/gate/up 族）与“大 K down 投影”两类实测优势形状收窄。
效果边界：M < 8192 的 decode 与小 prefill 全部留在 AOT；报告的 K=14336, N=4096 down 投影、窄 TP8 投影 5376x3584 / 3584x5376 以及未测量形状留在 AOT；M >= 8192 的宽 QKV/gate/up 与宽 down 投影继续走 Torch，保住大 prefill 收益。
验证与配套：未新增测试文件，但 PR 给出了三阶段基准协议——132 案例宽网格（11 个投影族 × 12 个 M 值）、两次随机化边界确认（44/33 案例与 44 案例）、精确工作负载热身与 30~40 样本时序统计；H100 上新谓词选中的 32 个形状全部偏向 Torch（中位数 +19.48%），H200 同样全部偏向 Torch（中位数 +2.50%）。CI 侧重跑了受影响回归测试 test/registered/quant/test_w8a8_quantization.py，在 1-gpu-h100 上通过。
提交演进：4 个 commit 显示从“Gate SM90 rowwise FP8 routing on M”到“Refine SM90 FP8 routing envelope”的两轮收敛，说明先加 M 门槛、再细化 K/N 包络的迭代过程，最后由 BBuf 修复 lint 并合入。

文件	模块	状态	重要度
`python/sglang/kernels/ops/gemm/__init__.py`	路由选择	modified	4.08

关键符号

_prefer_torch_rowwise_fp8

关键源码片段

python/sglang/kernels/ops/gemm/__init__.py core-logic

唯一改动文件，包含 Fp8ScaledMMOp 的 rowwise FP8 路由选择器 _prefer_torch_rowwise_fp8 ；本次修改的谓词直接决定 decode 与大 prefill GEMM 走 AOT 还是 torch._scaled_mm ，是修复回归的核心。

# python/sglang/kernels/ops/gemm/__init__.py
# Fp8ScaledMMOp 依赖此函数决定 rowwise FP8 GEMM 走 AOT 内核还是 Torch
def _prefer_torch_rowwise_fp8(m, k, n, ...):
    # 既有早期条件分支保持原样：不满足前置条件的形状在这里返回 False，
    # 由 AOT 内核执行，规避 Torch 后端在对应场景下的劣势

    # 本次修复的核心改动（替换了旧的 K/N-only 谓词）：
    # 1) m >= 8192 前置门槛：decode 与 M < 8192 的小 prefill 全部留在 AOT。
    # 实测 M=1、K=14336、N=4096 的 Llama-3.1-8B-Instruct-FP8-dynamic down
    # 投影在 H200 上 AOT 比 torch._scaled_mm 快 1.93 倍（26.50 µs 对 51.12 µs），
    # H100 上快 1.66 倍（32.19 µs 对 53.55 µs）；
    # 2) 只有在 M 足够大时，才把宽输出投影（gate/up/QKV，如 4096x6144、
    # 5376x21504）和宽 down 投影（如 14336x5376、28672x5376）路由到 Torch，
    # 这两类形状在 H100/H200 上均被 132 案例网格与随机边界确认证明更优；
    # 3) 窄投影（如 TP8 的 5376x3584、3584x5376）与未测量的形状留在 AOT，
    # 避免基于噪声调优造成二次回归。
    return (
        m >= 8192
        and ((k >= 4096 and n >= 6144) or (k >= 7168 and n >= 5376))
    )

评论区精华

路由改动合理性确认 other

唯一正式 review 来自 BBuf：“It seems resonable, approved!”，未提出技术质疑。改动以 PR 内详尽的 H100/H200 基准自证，评审无需追加讨论。

结论：已批准，无修改要求。 · 已解决

W8A8 回归测试复验 测试

issue 评论触发 /rerun-test test/registered/quant/test_w8a8_quantization.py ，在 1-gpu-h100 网格通过；与此同时 pr-test-extra 与 AMD ROCm 网格显示失败状态，PR 内未解释原因。

结论：受回归影响的 W8A8 测试复验通过；extra/AMD CI 失败列入风险，未在合并前澄清。 · 已解决

风险与影响

路由谓词是未经单元测试锁定的魔数，后续重构 _prefer_torch_rowwise_fp8 可能无意破坏该包络而不被感知；建议补充覆盖“decode 必走 AOT、大 M 宽投影走 Torch”的轻量单测。
实测形状仅覆盖 11 个投影族 × 12 个 M 值，其他模型族或 TP 配置的未测量形状可能偏离最优。作者明确采用“不在噪声上路由”的保守策略，最坏情况是牺牲少量潜在收益，而非引入新的明显回归。
m >= 8192 与 CUDA graph prefill 的 batch 上限耦合：未来若支持更大 prefill M，需要重做边界基准（M=8192 与 8320 的测量已显示收益在边界处规律变化）。
pr-test-extra 与 AMD ROCm 网格状态失败且未在 PR 内澄清；虽然改动只影响 SM90 rowwise FP8 路由，理论上与 AMD 无关，但合并前未完全闭环。

用户侧：H100 上 W8A8 模型小 batch 解码吞吐从约 190 tok/s 恢复到约 215 tok/s（约 12%）；大 prefill 用户仍享受 Torch 后端收益（H100 中位 +19.48%）。系统侧：影响范围仅限 rowwise FP8 GEMM 的后端选择，不触碰 AOT kernel 实现、非 FP8 路径或 CUDA graph 机制，属于低风险、高针对性的修复。团队侧：提供了一个可复用的“基准驱动路由调优”协议（宽网格 + 随机边界确认 + 精确工作负载热身），后续新架构或新形状的路由调整可以直接复用该验证方式。

路由谓词无单元测试锁定 M 门槛依赖实测基准样本未覆盖形状可能次优 CI extra/AMD 未闭环

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接，后续同步到相关引用后会出现在这里。

完整报告

执行摘要

一句话：限流 SM90 FP8 路由至大 M 包络，修复 decode 回归
推荐动作：值得精读。虽然仅 9 行改动，但它是“基准驱动路由选择”的教科书案例：先用宽网格确定优势面，再用随机化边界确认锁定阈值，最后对噪声区域采取保守策略。内核维护者和 FP8 性能工程师可以复刻这套验证协议；普通使用者了解结论即可——decode 一律走 AOT，大 prefill 宽投影走 Torch。建议后续为路由谓词补充形状注释表与轻量单测，防止回归。

功能与动机

实现拆解

变更入口：唯一改动文件为 python/sglang/kernels/ops/gemm/__init__.py，修改 Fp8ScaledMMOp 依赖的 rowwise FP8 Torch/AOT 路由选择器 _prefer_torch_rowwise_fp8。
谓词重构：旧谓词 (k >= 5376 and n >= 3584) or (k >= 3584 and m >= 8192) 的第一项完全不看 M，第二项也只是把 m >= 8192 与 k >= 3584 简单叠加；新谓词改为 m >= 8192 and ((k >= 4096 and n >= 6144) or (k >= 7168 and n >= 5376))，M 成为硬性前置门槛，K/N 包络按“宽输出投影”（QKV/gate/up 族）与“大 K down 投影”两类实测优势形状收窄。
效果边界：M < 8192 的 decode 与小 prefill 全部留在 AOT；报告的 K=14336, N=4096 down 投影、窄 TP8 投影 5376x3584 / 3584x5376 以及未测量形状留在 AOT；M >= 8192 的宽 QKV/gate/up 与宽 down 投影继续走 Torch，保住大 prefill 收益。
验证与配套：未新增测试文件，但 PR 给出了三阶段基准协议——132 案例宽网格（11 个投影族 × 12 个 M 值）、两次随机化边界确认（44/33 案例与 44 案例）、精确工作负载热身与 30~40 样本时序统计；H100 上新谓词选中的 32 个形状全部偏向 Torch（中位数 +19.48%），H200 同样全部偏向 Torch（中位数 +2.50%）。CI 侧重跑了受影响回归测试 test/registered/quant/test_w8a8_quantization.py，在 1-gpu-h100 上通过。
提交演进：4 个 commit 显示从“Gate SM90 rowwise FP8 routing on M”到“Refine SM90 FP8 routing envelope”的两轮收敛，说明先加 M 门槛、再细化 K/N 包络的迭代过程，最后由 BBuf 修复 lint 并合入。

关键文件：

python/sglang/kernels/ops/gemm/__init__.py（模块路由选择；类别 source；类型 core-logic；符号 _prefer_torch_rowwise_fp8）: 唯一改动文件，包含 Fp8ScaledMMOp 的 rowwise FP8 路由选择器 _prefer_torch_rowwise_fp8 ；本次修改的谓词直接决定 decode 与大 prefill GEMM 走 AOT 还是 torch._scaled_mm ，是修复回归的核心。

关键符号：_prefer_torch_rowwise_fp8

关键源码片段

`python/sglang/kernels/ops/gemm/init.py`

# python/sglang/kernels/ops/gemm/__init__.py
# Fp8ScaledMMOp 依赖此函数决定 rowwise FP8 GEMM 走 AOT 内核还是 Torch
def _prefer_torch_rowwise_fp8(m, k, n, ...):
    # 既有早期条件分支保持原样：不满足前置条件的形状在这里返回 False，
    # 由 AOT 内核执行，规避 Torch 后端在对应场景下的劣势

    # 本次修复的核心改动（替换了旧的 K/N-only 谓词）：
    # 1) m >= 8192 前置门槛：decode 与 M < 8192 的小 prefill 全部留在 AOT。
    # 实测 M=1、K=14336、N=4096 的 Llama-3.1-8B-Instruct-FP8-dynamic down
    # 投影在 H200 上 AOT 比 torch._scaled_mm 快 1.93 倍（26.50 µs 对 51.12 µs），
    # H100 上快 1.66 倍（32.19 µs 对 53.55 µs）；
    # 2) 只有在 M 足够大时，才把宽输出投影（gate/up/QKV，如 4096x6144、
    # 5376x21504）和宽 down 投影（如 14336x5376、28672x5376）路由到 Torch，
    # 这两类形状在 H100/H200 上均被 132 案例网格与随机边界确认证明更优；
    # 3) 窄投影（如 TP8 的 5376x3584、3584x5376）与未测量的形状留在 AOT，
    # 避免基于噪声调优造成二次回归。
    return (
        m >= 8192
        and ((k >= 4096 and n >= 6144) or (k >= 7168 and n >= 5376))
    )

评论区精华

路由改动合理性确认 (other): 已批准，无修改要求。
W8A8 回归测试复验 (testing): 受回归影响的 W8A8 测试复验通过；extra/AMD CI 失败列入风险，未在合并前澄清。

风险与影响

风险：
1. 路由谓词是未经单元测试锁定的魔数，后续重构 _prefer_torch_rowwise_fp8 可能无意破坏该包络而不被感知；建议补充覆盖“decode 必走 AOT、大 M 宽投影走 Torch”的轻量单测。
2. 实测形状仅覆盖 11 个投影族 × 12 个 M 值，其他模型族或 TP 配置的未测量形状可能偏离最优。作者明确采用“不在噪声上路由”的保守策略，最坏情况是牺牲少量潜在收益，而非引入新的明显回归。
3. m >= 8192 与 CUDA graph prefill 的 batch 上限耦合：未来若支持更大 prefill M，需要重做边界基准（M=8192 与 8320 的测量已显示收益在边界处规律变化）。
4. pr-test-extra 与 AMD ROCm 网格状态失败且未在 PR 内澄清；虽然改动只影响 SM90 rowwise FP8 路由，理论上与 AMD 无关，但合并前未完全闭环。
  - 影响：用户侧：H100 上 W8A8 模型小 batch 解码吞吐从约 190 tok/s 恢复到约 215 tok/s（约 12%）；大 prefill 用户仍享受 Torch 后端收益（H100 中位 +19.48%）。系统侧：影响范围仅限 rowwise FP8 GEMM 的后端选择，不触碰 AOT kernel 实现、非 FP8 路径或 CUDA graph 机制，属于低风险、高针对性的修复。团队侧：提供了一个可复用的“基准驱动路由调优”协议（宽网格 + 随机边界确认 + 精确工作负载热身），后续新架构或新形状的路由调整可以直接复用该验证方式。
  - 风险标记：路由谓词无单元测试锁定, M 门槛依赖实测基准样本, 未覆盖形状可能次优, CI extra/AMD 未闭环

关联脉络

PR #34318 引入 SM90 FP8 路由选择器（上游 PR，标题未在材料中提供）: 本 PR 的直接修复对象：#34318 引入的 K/N-only 路由导致 H100 decode 回归，PR body 引用了该 PR 下的 H100 W8A8 回归报告链接。

#37018 [Kernel] Fix SM90 FP8 decode regression with benchmarked M/K/N routing

执行摘要

限流 SM90 FP8 路由至大 M 包络，修复 decode 回归

实现拆解

评论区精华

风险与影响

关联 Issue

未识别关联 Issue

完整报告

参与讨论