deepseek-ai/DeepGEMM
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepGEMM.git synced 2026-11-05 00:47:26 +00:00
Code Issues 162 Projects Packages Wiki Activity

[PR #411] feat: add SM90 FP8 x MXFP4 MegaMoE support #427

New issue
Open
opened 2026-10-26 21:17:39 +00:00 by bsakaguchi · 0 comments
bsakaguchi commented 2026-10-26 21:17:39 +00:00
Owner
Copy link

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/411
Author: @huangzhilin-hzl
Created: 8/20/2026
Status: 🔄 Open

Base: nv_dev ← Head: molou/sm90_humming_mxfp4afp8_megamoe_dev

📝 Commits (3)

  • ec72c60 feat(sm90): add FP8 x MXFP4 MegaMoE support
  • f026a32 fix(sm90): harden cross-rank MegaMoE paths
  • 5d32216 perf(sm90): integrate latest MXFP4 MegaMoE optimizations

📊 Changes

20 files changed (+8479 additions, -60 deletions)

View changed files

➕ csrc/apis/sm90_mega.hpp (+429 -0)
📝 csrc/jit/handle.hpp (+58 -6)
📝 csrc/jit/kernel_runtime.hpp (+22 -12)
➕ csrc/jit_kernels/heuristics/sm90_mega_moe.hpp (+138 -0)
📝 csrc/jit_kernels/impls/runtime_utils.hpp (+6 -3)
➕ csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp (+472 -0)
📝 csrc/python_api.cpp (+2 -0)
📝 deep_gemm/__init__.py (+5 -0)
➕ deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh (+4033 -0)
📝 deep_gemm/include/deep_gemm/layout/mega_moe.cuh (+89 -5)
➕ deep_gemm/include/deep_gemm/layout/sm90_mega_moe.cuh (+20 -0)
📝 deep_gemm/include/deep_gemm/mma/sm90.cuh (+47 -0)
📝 deep_gemm/include/deep_gemm/ptx/ld_st.cuh (+14 -0)
📝 deep_gemm/include/deep_gemm/ptx/wgmma.cuh (+4 -0)
📝 deep_gemm/include/deep_gemm/scheduler/mega_moe.cuh (+188 -34)
📝 deep_gemm/mega/__init__.py (+252 -0)
➕ deep_gemm/mega/mxfp4.py (+427 -0)
➕ tests/bench_mega_moe_sm90.py (+672 -0)
➕ tests/test_mega_moe_mxfp4_preprocess.py (+485 -0)
➕ tests/test_mega_moe_sm90.py (+1116 -0)

📄 Description

co-author： @guzekai01

Inspired by the SM90 FP8×FP8 MegaMoE implementation introduced in deepseek-ai/DeepGEMM#383 and the SM90 FP8×NVFP4 work in AichenF/DeepGEMM@megamoe_nvfp4_dev_m, this PR adds an optimized FP8×MXFP4 MegaMoE backend for SM90 GPUs. Keeping routed-expert weights in MXFP4 enables native reuse of Hugging Face MXFP4 MoE checkpoints, while substantially reducing GPU memory pressure by avoiding a persistent FP8 weight expansion.

Summary

  • Unified SM90 persistent execution.
    Unlike #383, this implementation fuses dispatch, routed/shared L1 and L2, and combine into one persistent kernel. It adapts the SM100 ring-buffer and wave-interleaved scheduler to a cluster-free Hopper design using register-resident WGMMA, one CTA per task, and two resident CTAs per SM.

  • Humming-inspired MXFP4×FP8 compute path.
    Routed weights remain packed MXFP4 in global memory. Model-load preprocessing rebases the K32 UE8M0 exponents into bounded relative scales plus one FP32 secondary scale per expert, following the fused MXFP4×FP8 approach used by Humming. During execution, packed E2M1 tiles are loaded into compact shared-memory staging, decoded and expanded on the fly into E4M3 FP8 tiles, and then consumed by Hopper WGMMA. This avoids materializing a persistent FP8 copy of the routed weights.

Benchmark Results

The benchmark follows the H20 workload shapes and timing methodology from #383. Results were collected on 8× NVIDIA H20 with fast_math=1 and cold-L2 measurements.

vs PR #383 is calculated as PR #383 / FP8×MXFP4 - 1; vs DeepEP is calculated as DeepEP / FP8×MXFP4 - 1. Positive values indicate that this implementation is faster.

H20 High Throughput

Model M FP8×MXFP4 (us) PR #383 FP8xFP8 (us) vs PR #383 DeepEP HT (us) vs DeepEP
Flash 8 272.2 273.1 +0.3% 566.7 +108.2%
Flash 16 311.1 304.4 -2.1% 619.2 +99.0%
Flash 32 323.1 302.0 -6.5% 618.6 +91.4%
Flash 64 333.9 340.7 +2.0% 622.6 +86.5%
Flash 128 408.5 414.4 +1.4% 638.5 +56.3%
Flash 256 471.7 569.5 +20.7% 688.7 +46.0%
Flash 512 881.3 922.0 +4.6% 1057.3 +20.0%
Flash 1024 1449.0 1516.6 +4.7% 1983.4 +36.9%
Flash 2048 2665.0 2735.1 +2.6% 3419.5 +28.3%
Flash 4096 5005.0 5116.0 +2.2% 6087.7 +21.6%
Flash 8192 9741.0 9749.0 +0.1% 11779.2 +20.9%
Pro 8 700.2 768.0 +9.7% 1349.2 +92.7%
Pro 16 913.4 950.3 +4.0% 1585.8 +73.6%
Pro 32 985.2 1026.3 +4.2% 1770.2 +79.7%
Pro 64 1013.0 1059.9 +4.6% 1787.6 +76.5%
Pro 128 1178.5 1201.0 +1.9% 1803.9 +53.1%
Pro 256 1593.0 1639.9 +2.9% 1857.7 +16.6%
Pro 512 2474.0 2599.0 +5.1% 2898.8 +17.2%
Pro 1024 3852.0 4036.0 +4.8% 5412.7 +40.5%
Pro 2048 6865.0 6986.0 +1.8% 8067.9 +17.5%
Pro 4096 12832.0 12932.0 +0.8% 14614.0 +13.9%
Pro 8192 24912.0 24777.0 -0.5% 28184.2 +13.1%

H20 Low Latency

Model M FP8×MXFP4 (us) PR #383 FP8xFP8 (us) vs PR #383 DeepEP LL (us) vs DeepEP
Flash 8 272.2 273.1 +0.3% 479.7 +76.2%
Flash 16 311.1 304.4 -2.1% 528.4 +69.9%
Flash 32 323.1 302.0 -6.5% 531.8 +64.6%
Flash 64 333.9 340.7 +2.0% 552.2 +65.4%
Flash 128 408.5 414.4 +1.4% 552.8 +35.3%
Pro 8 700.2 768.0 +9.7% 1265.6 +80.8%
Pro 16 913.4 950.3 +4.0% 1493.3 +63.5%
Pro 32 985.2 1026.3 +4.2% 1667.2 +69.2%
Pro 64 1013.0 1059.9 +4.6% 1683.8 +66.2%
Pro 128 1178.5 1201.0 +1.9% 1714.1 +45.4%

TODO

  • Support additional MoE architectures, including K3-style latent MoE.
  • Continue fine-grained kernel performance tuning and comprehensive numerical-accuracy validation.
  • Integrate with inference frameworks such as SGLang to evaluate end-to-end performance and model accuracy.

🔄 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/411 **Author:** [@huangzhilin-hzl](https://github.com/huangzhilin-hzl) **Created:** 8/20/2026 **Status:** 🔄 Open **Base:** `nv_dev` ← **Head:** `molou/sm90_humming_mxfp4afp8_megamoe_dev` --- ### 📝 Commits (3) - [`ec72c60`](https://github.com/deepseek-ai/DeepGEMM/commit/ec72c60ca81f4ef0ea8531926e7a61bcd2a77ae5) feat(sm90): add FP8 x MXFP4 MegaMoE support - [`f026a32`](https://github.com/deepseek-ai/DeepGEMM/commit/f026a32667a88aacef2ce4982c15ec17fec34b19) fix(sm90): harden cross-rank MegaMoE paths - [`5d32216`](https://github.com/deepseek-ai/DeepGEMM/commit/5d32216686b982a39eabccca9419af430a60cfc2) perf(sm90): integrate latest MXFP4 MegaMoE optimizations ### 📊 Changes **20 files changed** (+8479 additions, -60 deletions) <details> <summary>View changed files</summary> ➕ `csrc/apis/sm90_mega.hpp` (+429 -0) 📝 `csrc/jit/handle.hpp` (+58 -6) 📝 `csrc/jit/kernel_runtime.hpp` (+22 -12) ➕ `csrc/jit_kernels/heuristics/sm90_mega_moe.hpp` (+138 -0) 📝 `csrc/jit_kernels/impls/runtime_utils.hpp` (+6 -3) ➕ `csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp` (+472 -0) 📝 `csrc/python_api.cpp` (+2 -0) 📝 `deep_gemm/__init__.py` (+5 -0) ➕ `deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh` (+4033 -0) 📝 `deep_gemm/include/deep_gemm/layout/mega_moe.cuh` (+89 -5) ➕ `deep_gemm/include/deep_gemm/layout/sm90_mega_moe.cuh` (+20 -0) 📝 `deep_gemm/include/deep_gemm/mma/sm90.cuh` (+47 -0) 📝 `deep_gemm/include/deep_gemm/ptx/ld_st.cuh` (+14 -0) 📝 `deep_gemm/include/deep_gemm/ptx/wgmma.cuh` (+4 -0) 📝 `deep_gemm/include/deep_gemm/scheduler/mega_moe.cuh` (+188 -34) 📝 `deep_gemm/mega/__init__.py` (+252 -0) ➕ `deep_gemm/mega/mxfp4.py` (+427 -0) ➕ `tests/bench_mega_moe_sm90.py` (+672 -0) ➕ `tests/test_mega_moe_mxfp4_preprocess.py` (+485 -0) ➕ `tests/test_mega_moe_sm90.py` (+1116 -0) </details> ### 📄 Description co-author： @guzekai01 Inspired by the SM90 FP8×FP8 MegaMoE implementation introduced in [deepseek-ai/DeepGEMM#383](https://github.com/deepseek-ai/DeepGEMM/pull/383) and the SM90 FP8×NVFP4 work in [AichenF/DeepGEMM@megamoe_nvfp4_dev_m](https://github.com/AichenF/DeepGEMM/tree/megamoe_nvfp4_dev_m), this PR adds an optimized FP8×MXFP4 MegaMoE backend for SM90 GPUs. Keeping routed-expert weights in MXFP4 enables native reuse of Hugging Face MXFP4 MoE checkpoints, while substantially reducing GPU memory pressure by avoiding a persistent FP8 weight expansion. ## Summary - **Unified SM90 persistent execution.** Unlike #383, this implementation fuses dispatch, routed/shared L1 and L2, and combine into one persistent kernel. It adapts the SM100 ring-buffer and wave-interleaved scheduler to a cluster-free Hopper design using register-resident WGMMA, one CTA per task, and two resident CTAs per SM. - **Humming-inspired MXFP4×FP8 compute path.** Routed weights remain packed MXFP4 in global memory. Model-load preprocessing rebases the K32 UE8M0 exponents into bounded relative scales plus one FP32 secondary scale per expert, following the fused MXFP4×FP8 approach used by Humming. During execution, packed E2M1 tiles are loaded into compact shared-memory staging, decoded and expanded on the fly into E4M3 FP8 tiles, and then consumed by Hopper WGMMA. This avoids materializing a persistent FP8 copy of the routed weights. ## Benchmark Results The benchmark follows the H20 workload shapes and timing methodology from #383. Results were collected on **8× NVIDIA H20** with `fast_math=1` and cold-L2 measurements. `vs PR #383` is calculated as `PR #383 / FP8×MXFP4 - 1`; `vs DeepEP` is calculated as `DeepEP / FP8×MXFP4 - 1`. Positive values indicate that this implementation is faster. ### H20 High Throughput | Model | M | FP8×MXFP4 (us) | PR #383 FP8xFP8 (us) | vs PR #383 | DeepEP HT (us) | vs DeepEP | |---|---:|---:|---:|---:|---:|---:| | Flash | 8 | 272.2 | 273.1 | +0.3% | 566.7 | +108.2% | | Flash | 16 | 311.1 | 304.4 | -2.1% | 619.2 | +99.0% | | Flash | 32 | 323.1 | 302.0 | -6.5% | 618.6 | +91.4% | | Flash | 64 | 333.9 | 340.7 | +2.0% | 622.6 | +86.5% | | Flash | 128 | 408.5 | 414.4 | +1.4% | 638.5 | +56.3% | | Flash | 256 | 471.7 | 569.5 | +20.7% | 688.7 | +46.0% | | Flash | 512 | 881.3 | 922.0 | +4.6% | 1057.3 | +20.0% | | Flash | 1024 | 1449.0 | 1516.6 | +4.7% | 1983.4 | +36.9% | | Flash | 2048 | 2665.0 | 2735.1 | +2.6% | 3419.5 | +28.3% | | Flash | 4096 | 5005.0 | 5116.0 | +2.2% | 6087.7 | +21.6% | | Flash | 8192 | 9741.0 | 9749.0 | +0.1% | 11779.2 | +20.9% | | Pro | 8 | 700.2 | 768.0 | +9.7% | 1349.2 | +92.7% | | Pro | 16 | 913.4 | 950.3 | +4.0% | 1585.8 | +73.6% | | Pro | 32 | 985.2 | 1026.3 | +4.2% | 1770.2 | +79.7% | | Pro | 64 | 1013.0 | 1059.9 | +4.6% | 1787.6 | +76.5% | | Pro | 128 | 1178.5 | 1201.0 | +1.9% | 1803.9 | +53.1% | | Pro | 256 | 1593.0 | 1639.9 | +2.9% | 1857.7 | +16.6% | | Pro | 512 | 2474.0 | 2599.0 | +5.1% | 2898.8 | +17.2% | | Pro | 1024 | 3852.0 | 4036.0 | +4.8% | 5412.7 | +40.5% | | Pro | 2048 | 6865.0 | 6986.0 | +1.8% | 8067.9 | +17.5% | | Pro | 4096 | 12832.0 | 12932.0 | +0.8% | 14614.0 | +13.9% | | Pro | 8192 | 24912.0 | 24777.0 | -0.5% | 28184.2 | +13.1% | ### H20 Low Latency | Model | M | FP8×MXFP4 (us) | PR #383 FP8xFP8 (us) | vs PR #383 | DeepEP LL (us) | vs DeepEP | |---|---:|---:|---:|---:|---:|---:| | Flash | 8 | 272.2 | 273.1 | +0.3% | 479.7 | +76.2% | | Flash | 16 | 311.1 | 304.4 | -2.1% | 528.4 | +69.9% | | Flash | 32 | 323.1 | 302.0 | -6.5% | 531.8 | +64.6% | | Flash | 64 | 333.9 | 340.7 | +2.0% | 552.2 | +65.4% | | Flash | 128 | 408.5 | 414.4 | +1.4% | 552.8 | +35.3% | | Pro | 8 | 700.2 | 768.0 | +9.7% | 1265.6 | +80.8% | | Pro | 16 | 913.4 | 950.3 | +4.0% | 1493.3 | +63.5% | | Pro | 32 | 985.2 | 1026.3 | +4.2% | 1667.2 | +69.2% | | Pro | 64 | 1013.0 | 1059.9 | +4.6% | 1683.8 | +66.2% | | Pro | 128 | 1178.5 | 1201.0 | +1.9% | 1714.1 | +45.4% | ## TODO - Support additional MoE architectures, including K3-style latent MoE. - Continue fine-grained kernel performance tuning and comprehensive numerical-accuracy validation. - Integrate with inference frameworks such as SGLang to evaluate end-to-end performance and model accuracy. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>
Sign in to join this conversation.
No Branch/Tag specified
Branches Tags
main
public-sync-20261014
0930-release
nv_dev
fix-mega-moe-fence
public-release-2609
public-release-2607
public-release-2606
mega-update
public-release-260416
ko3n1g/ci/build-wheels
nv_dev_f8e8fb5
nv_dev_950d31a
nv_dev_67fc648
nv_dev_c491439
nv_dev_4ff3f54
v2.1.1.post3
v2.1.1.post2
v2.1.1.post1
v2.1.1
v2.1.0
Labels
Clear labels   
pull-request
Mirrored from GitHub Pull Request

No labels
pull-request
Milestone
Clear milestone
No items
No milestone
Projects
Clear projects
No items
No project
Assignees
Clear assignees
No assignees
1 participant
Notifications
Due date
The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference
deepseek-ai/DeepGEMM#427
No description provided.