fix(sm90): support parameterized SwiGLU in FP8 MegaMoE

[shiyang814-cpu]

Summary

This PR adds parameterized SwiGLU support to the SM90 FP8 MegaMoE kernel.

The existing implementation assumes standard SwiGLU:

SiLU(gate) * up

Some MoE models use parameterized OAI-SwiGLU instead:

gate = clamp(gate, max=limit)
up = clamp(up, min=-limit, max=limit)

gate * sigmoid(alpha * gate) * (up + up_bias)

For example, MiniMax-M3 uses:

alpha = 1.702
up_bias = 1.0
clamp = 7.0

Without these parameters, the SM90 MegaMoE kernel applies standard SwiGLU semantics to models that require OAI-SwiGLU, which produces incorrect model outputs.

Root Cause

The SM90 FP8 MegaMoE epilogues currently hardcode:

gate * sigmoid(gate) * up

Although the activation clamp is configurable, the sigmoid alpha and up-projection bias are not represented by the Python API, host API, or JIT kernel template.

As a result, the generated kernel cannot reproduce parameterized SwiGLU semantics. This affects both the optimized swapAB epilogue and the non-swapAB epilogue.

Changes

• Add activation_alpha and activation_up_bias to the Python SM90 FP8 MegaMoE API.
• Propagate both parameters through the SM90 host API.
• Include both parameters in the generated JIT kernel specialization and cache identity.
• Apply parameterized SwiGLU semantics in both:
  • swapAB epilogue
  • non-swapAB epilogue
• Accept swigluoai in the SM90 buffer allocation and execution APIs.
• Extend the existing Hopper accuracy suite with an OAI-SwiGLU scenario using:
  • activation_alpha=1.702
  • activation_up_bias=1.0
  • activation_clamp=7.0

Compatibility

The new parameters default to standard SwiGLU behavior:

activation_alpha = 1.0  activation_up_bias = 0.0

Therefore, existing callers that do not provide the new parameters retain the original behavior:

gate * sigmoid(gate) * up

The new activation name is additive, and no behavior change is expected for existing SM90 MegaMoE workloads.

Validation

The change was validated on Hopper GPUs in a real two-node prefill/decode disaggregated deployment.

Environment

Model: MiniMax-M3-FP8
GPU: NVIDIA H20 (SM90)
Prefill: TP=8, EP=8, DeepEP normal
Decode: TP=8, EP=8, SM90 FP8 MegaMoE
KV transfer: Mooncake
Weight format: FP8 E4M3

MiniMax-M3 uses the following OAI-SwiGLU parameters:

activation_alpha = 1.702
activation_up_bias = 1.0
activation_clamp = 7.0

Kernel Accuracy

The SM90 MegaMoE kernel output was compared against the reference implementation on 8 Hopper GPUs.

Default swapAB + OAI-SwiGLU: calc_diff = 0.0000

Non-swapAB + OAI-SwiGLU: calc_diff = 0.0000

Standard SwiGLU regression: calc_diff = 0.0000

Both the swapAB and non-swapAB epilogues therefore preserve the expected parameterized SwiGLU semantics.

End-to-End P/D Validation

The patched kernel was also validated in a two-node P/D deployment:

• DeepEP normal was active on the prefill node.

• SM90 FP8 MegaMoE was active on the decode node.

• No fallback to Triton or another MoE backend occurred.

• Deterministic English generation passed.

• Deterministic Chinese generation passed.

• MMLU smoke evaluation passed: 8/8.

• No corrupted output, scheduler exception, or kernel traceback was observed.

The DeepGEMM JIT cache was cleared before validation to ensure that the generated SM90 kernels were rebuilt from the patched headers.

Static Checks

python3 -m py_compile \
  deep_gemm/mega/__init__.py \
  tests/test_mega_moe_sm90.py

git diff --check

[ds-review-bot]

🤖 ds-review-bot Code Review

v6

JIT 格式串缺失占位符会阻断新增功能及现有标准 SwiGLU 路径。其余参数接口和激活公式与关联 MR 一致；当前环境无 CUDA 工具链，未复跑 GPU 数值测试。

v5

本 MR 将 DeepGEMM #87 的参数化 SwiGLU（swigluoai）支持移植到 nv_dev 的 SM90 FP8 MegaMoE。跨仓契约核对结果：Python 入口 fp8_mega_moe(..., activation, activation_alpha=1.0, activation_up_bias=0.0, activation_clamp=None, fast_math)、C++ fp8_mega_moe 形参顺序、内核模板参数顺序（kActivationClamp, kActivationAlpha, kActivationUpBias, kFastMath, kFP8SwapAB, kBF16ScaledAccumRequested）、swapAB/非 swapAB 两条 epilogue 的数学实现、以及测试参考 gate * sigmoid(alpha*gate) * (up + up_bias) 均与对端 PR 一致；本仓使用 pybind11 直接绑定，因此对端 csrc/tvm_ffi_api.cpp 的改动在本仓无需对应，属预期差异。to_string(float) 使用十六进制浮点格式，1.702 等参数可精确传入模板。

但存在一个阻塞性问题：csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp 的 generate_impl 中，fmt 模板只新增了 1 行 {}, 占位符，而实参新增了 2 个（alpha、up_bias），导致 25 个占位符对应 26 个实参。fmt 不会诊断多余实参，结果是 {}{} 被 swap_ab 与 bf16_scaled_accum 填充（生成 truefalse 等非法标识符），而 phase_template_args（nmajor_schedule / direct_l2_scatter / one_warp_cleanup）被整体丢弃。这会使所有 SM90 FP8 MegaMoE 内核（包括普通 swiglu）在 JIT 编译阶段失败。对端 DeepGEMM #87 正确地新增了 2 行 {},，本 MR 移植时遗漏了一行；描述中的 calc_diff=0 验证来自对端仓库，本仓当前代码无法通过该验证。修复为在 fmt 模板中再补一行 {},。

另有两处非阻塞建议：activation 字符串与 alpha/up_bias 未做关联校验（swiglu 传非默认 alpha 会静默生效，与对端行为一致）；测试中 _swiglu_fp32 的 docstring 仍描述旧公式。

v4

Cross-repo contract with sgl-project#87 is consistent at the public API level. deep_gemm.mega.fp8_mega_moe (re-exported as deep_gemm.fp8_mega_moe) exposes activation_alpha (default 1.0) and activation_up_bias (default 0.0) in exactly the same position/order as the upstream sgl_deep_gemm.fp8_mega_moe change (inserted after 'activation' and before 'activation_clamp'), the pybind binding in csrc/apis/sm90_mega.hpp forwards them in the same order, and both get_symm_buffer_size_for_sm90_mega_moe and fp8_mega_moe accept activation in {'swiglu','swigluoai'}. The kernel math matches the upstream contract: clamp_gate(g)=min(g,c), clamp_up(u)=clamp(u,-c,c), result = clamp_gate(g)sigmoid(alphaclamp_gate(g))*(clamp_up(u)+up_bias), with alpha=1.0/up_bias=0.0 reducing exactly to the old swiglu path. The reference/test wiring is also consistent. However, the port introduces one actionable code-generation defect that violates the host->kernel template calling convention in csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp.

Files reviewed: 5
配对 MR checkout: 1
Issues found: 🔴 3 critical | 🔵 2 suggestion
Inline comments posted: 5