deepseek-ai/DeepGEMM

Fork 0

mirror of https://github.com/deepseek-ai/DeepGEMM.git synced 2026-11-05 00:47:26 +00:00

[PR #445] fix(sm90): support parameterized SwiGLU in FP8 MegaMoE #451

New issue

Open

opened 2026-10-26 21:18:00 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 21:18:00 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/445
Author: @shiyang814-cpu
Created: 9/14/2026
Status: 🔄 Open

Base: nv_dev ← Head: bugfix/sm90-megamoe-parameterized-swiglu-deepseek

📝 Commits (2)

45cef20 fix(sm90): support parameterized SwiGLU in FP8 MegaMoE
8645ed6 fix(sm90): align JIT template arguments

📊 Changes

5 files changed (+75 additions, -21 deletions)

View changed files

📝 csrc/apis/sm90_mega.hpp (+13 -5)
📝 csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp (+10 -0)
📝 deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh (+14 -9)
📝 deep_gemm/mega/__init__.py (+4 -1)
📝 tests/test_mega_moe_sm90.py (+34 -6)

📄 Description

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

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/deepseek-ai/DeepGEMM/pull/445 **Author:** [@shiyang814-cpu](https://github.com/shiyang814-cpu) **Created:** 9/14/2026 **Status:** 🔄 Open **Base:** `nv_dev` ← **Head:** `bugfix/sm90-megamoe-parameterized-swiglu-deepseek` --- ### 📝 Commits (2) - [`45cef20`](https://github.com/deepseek-ai/DeepGEMM/commit/45cef20018c4a58600361af8ca30555df7fe9084) fix(sm90): support parameterized SwiGLU in FP8 MegaMoE - [`8645ed6`](https://github.com/deepseek-ai/DeepGEMM/commit/8645ed67eead469d585d72e34b9a8ec93b218397) fix(sm90): align JIT template arguments ### 📊 Changes **5 files changed** (+75 additions, -21 deletions) <details> <summary>View changed files</summary> 📝 `csrc/apis/sm90_mega.hpp` (+13 -5) 📝 `csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp` (+10 -0) 📝 `deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh` (+14 -9) 📝 `deep_gemm/mega/__init__.py` (+4 -1) 📝 `tests/test_mega_moe_sm90.py` (+34 -6) </details> ### 📄 Description ## Summary This PR adds parameterized SwiGLU support to the SM90 FP8 MegaMoE kernel. The existing implementation assumes standard SwiGLU: ```text SiLU(gate) * up ``` Some MoE models use parameterized OAI-SwiGLU instead: ```text gate = clamp(gate, max=limit) up = clamp(up, min=-limit, max=limit) gate * sigmoid(alpha * gate) * (up + up_bias) ``` For example, MiniMax-M3 uses: ```text 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: ```text 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: ```text activation_alpha = 1.0 activation_up_bias = 0.0 ``` Therefore, existing callers that do not provide the new parameters retain the original behavior: ```text 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 ```text 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: ```text 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. ```text 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 ```bash python3 -m py_compile \ deep_gemm/mega/__init__.py \ tests/test_mega_moe_sm90.py git diff --check ``` --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi added the

pull-request

label

2026-10-26 21:18:00 +00:00

bsakaguchi referenced this issue

2026-10-26 21:18:03 +00:00

[PR #451] [Bugfix] Reset MegaMoE combine readiness for CUDA graph replay #456

bsakaguchi referenced this issue

2026-10-26 21:18:21 +00:00

[PR #478] [MERGED] Fix mega_mhc split barriers under CUDA graph replay #478