Prhub

#37159 [Kernel] Add GB300 Triton MoE configs for GLM-4.5 FP8

原始 PR 作者 BBuf 合并时间 2026-08-31 21:31 文件变更 1 提交数 3 评论 1 代码增减 +58 / -0

执行摘要

GB300 GLM-4.5 FP8 新增 Triton MoE 配置,吞吐提升约 22%

GLM-4.5-FP8 的 per-rank TP8 形态是 E=161、H=5120、I=192、topk=9 的 per-channel FP8 Triton MoE,SGLang 的 Triton 3.7.1 GB300 配置表没有这条 lane,只能回落到通用启发式调度。该缺失是在验证 KDA-Pilot PR #197 时定位到的;作者还评估了更大的自定义 kernel,但 tuned generic kernel 更快,因此最终只提交配置变更,避免引入更大 kernel 的维护成本。

值得性能工程师精读。它展示了 SGLang Triton MoE 配置系统的设计要点:通过文件名编码 device、Triton 版本、dtype、量化模式与形状选择器,配置与代码解耦。同时也呈现了完整的调优验证流程:microbenchmark 定位、端到端服务压测、精度对比与 CUDA Graph 捕获验证。需要注意该 PR 最终没有携带自动化测试,建议关注合入后该 lane 是否会被后续重构影响。

讨论亮点

本 PR 没有 review 评论,唯一的 PR 评论是 /tag-and-rerun-ci 重跑 CI 指令。核心设计权衡记录在 PR body 中:作者评估了更大的自定义 kernel,但 tuned generic kernel 更快,因此只提交更小的配置变更。commit 历史显示第二个 commit 明确移除了 config availability 测试,说明作者最终选择不把配置可用性测试纳入合入集,依赖配置目录结构和 fallback 机制保证隔离。

实现拆解

  1. 定位配置入口:SGLang 的 Triton MoE runner 按 device_name、Triton 版本、dtype、量化 scale 模式、E/N 形状从 triton_utils/configs// 目录读取 JSON 配置,缺少匹配文件时回落到启发式调度,因此新增配置文件即可精准覆盖目标 lane。
  2. 新增配置表:文件 python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/triton_3_7_1/E=161,N=192,device_name=NVIDIA_GB300,dtype=fp8_w8a8,per_channel_quant=True.json 提供 M=1/16/256/512/1024/4096/16384 七档配置。小 batch 用 BLOCK_SIZE_M=16 与较少 stages 控制延迟;M=16 将 BLOCK_SIZE_N/K 提到 128 并用 4 级流水提升吞吐;M>=256 切换为 BLOCK_SIZE_M=64、GROUP_SIZE_M=32,充分使用 SM103 的并行度并改善 L2 复用。
  3. 精度验证:默认与 tuned 调度对比使用真实 [E,N] per-channel scale 布局,M=1/16/256/2048 的最大绝对差为 2.4e-7、最大平均相对差为 1.67e-8;真实 TP8 服务完成全部 74 次 prefill 与 52 次 decode CUDA Graph 捕获,数值行为稳定。
  4. 性能验证:fused-experts 全调用(up GEMM + SiLU-mul + down GEMM + 路由加权 + combine)在 M=1/16/256/512 分别提速 1.067x、1.440x、1.727x、1.522x;两节点 TP8 GLM-4.5-FP8 服务端到端 decode +21.88%、混合 +21.48%,三个随机种子下结果一致。
  5. 测试配套:第一版曾包含一个 config availability CPU 测试,在第二个 commit 中被移除,最终合入集只有配置文件本身,未引入自动化测试。
文件 模块 状态 重要度
python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/triton_3_7_1/E=161,N=192,device_name=NVIDIA_GB300,dtype=fp8_w8a8,per_channel_quant=True.json MoE 配置 added 5.14

关键源码片段

python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/triton_3_7_1/E=161,N=192,device_name=NVIDIA_GB300,dtype=fp8_w8a8,per_channel_quant=True.json configuration

唯一变更文件,为 GLM-4.5-FP8 在 GB300 + Triton 3.7.1 的 per-channel FP8 Triton MoE lane 提供 7 档 tuned 配置;文件名选择器确保只影响该精确组合,其余硬件与形状保持原 fallback。

// SGLang Triton MoE tuned config: E=161, N=192, GB300 (SM103), Triton 3.7.1
// dtype=fp8_w8a8, per_channel_quant=True. Key is token count M.
// 只有文件名精确匹配此组合才命中,其他硬件 / 版本 / 量化 / 形状走原 fallback。
{
    "1": {   // decode 单 token:小 N 切块 32,3 级流水压低延迟
        "BLOCK_SIZE_M": 16,
        "BLOCK_SIZE_N": 32,
        "BLOCK_SIZE_K": 64,
        "GROUP_SIZE_M": 1,
        "num_warps": 4,
        "num_stages": 3
    },
    "16": {  // 小 batch:N 与 K 放大到 128,4 级流水提升吞吐
        "BLOCK_SIZE_M": 16,
        "BLOCK_SIZE_N": 128,
        "BLOCK_SIZE_K": 128,
        "GROUP_SIZE_M": 1,
        "num_warps": 4,
        "num_stages": 4
    },
    "256": { // 中等 batch:GROUP_SIZE_M=32 改善 L2 复用
        "BLOCK_SIZE_M": 16,
        "BLOCK_SIZE_N": 128,
        "BLOCK_SIZE_K": 64,
        "GROUP_SIZE_M": 32,
        "num_warps": 4,
        "num_stages": 3
    },
    "512": { // 大 batch:BLOCK_SIZE_M 升到 64,充分利用 SM103 并行度
        "BLOCK_SIZE_M": 64,
        "BLOCK_SIZE_N": 128,
        "BLOCK_SIZE_K": 64,
        "GROUP_SIZE_M": 32,
        "num_warps": 4,
        "num_stages": 3
    },
    // 1024/4096/16384 沿用大 batch 最优组合,覆盖 chunked prefill 至 16K
    "1024": {
        "BLOCK_SIZE_M": 64,
        "BLOCK_SIZE_N": 128,
        "BLOCK_SIZE_K": 64,
        "GROUP_SIZE_M": 32,
        "num_warps": 4,
        "num_stages": 3
    },
    "4096": {
        "BLOCK_SIZE_M": 64,
        "BLOCK_SIZE_N": 128,
        "BLOCK_SIZE_K": 64,
        "GROUP_SIZE_M": 32,
        "num_warps": 4,
        "num_stages": 3
    },
    "16384": {
        "BLOCK_SIZE_M": 64,
        "BLOCK_SIZE_N": 128,
        "BLOCK_SIZE_K": 64,
        "GROUP_SIZE_M": 32,
        "num_warps": 4,
        "num_stages": 3
    }
}

评论区精华

没有提炼出高价值讨论线程

当前评论区没有形成足够清晰的争议点或结论,后续有更多讨论时会体现在这里。

风险与影响

回归面极小:配置文件名强绑定 NVIDIA_GB300 + Triton 3.7.1 + fp8_w8a8 + per_channel_quant=True + E=161/N=192,其它硬件、Triton 版本、量化模式与 MoE 形状均不会命中这条 lane,会继续走原有 fallback。主要风险是测试缺失:配置没有自动化测试保护,未来配置查找逻辑重构或目录整理可能使该 lane 静默丢失。其次是 Triton 版本耦合:用户使用 Triton 3.8+ 时文件名不匹配会走 fallback,行为安全但性能退回启发式。FP8 调度顺序变化带来 2.4e-7 级输出差异,对 bitexact 场景需要留意。另外 tuned 配置只覆盖到 16384 token,更大的 prefill 会依赖 fallback。

影响范围仅限 GLM-4.5-FP8 在 GB300 上的 Triton MoE 路径:用户侧端到端吞吐提升约 21-22%、TPOT 下降约 19%,无需改动代码或启动参数。系统侧只是配置表新增一条 lane,无运行时行为变更风险。团队侧这是一次典型的纯配置调优实践,验证了 SGLang Triton MoE 配置架构(文件名即选择器)的扩展性,为后续其它模型在特定硬件上的调优提供了可复制的流程模板。

测试被移除 配置强绑定硬件与 Triton 版本 FP8 数值微小差异

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论