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 #323] feat: MegaMOE adaptation for SM90 #367

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

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/323
Author: @qiushixiaoyu
Created: 4/30/2026
Status: 🔄 Open

Base: main ← Head: main

📝 Commits (10+)

  • 88db27f Support tvm-ffi interfaces
  • 495fcfd Support Wheel Compilation
  • 3036815 Export the PDL utils of DeepGEMM (#34)
  • 607503a Expose BF16 grouped GEMM wrappers
  • 00c42b8 Setup tests for sgl-deep-gemm
  • a8459f1 Sm90 mega moe on sgl dev (#36)
  • eed753c Add Hopper mega moe test to runner (#46)
  • f3d0ca5 chore: bump apache-tvm-ffi 0.1.9 -> 0.1.11 (#47)
  • bb569e2 Add FP4 acts + MXF4 kind support and fused mega_moe_pre_dispatch kernel (#27)
  • b111331 Add DG_USE_FP8_COMBINE: FP8 + per-row UE8M0 SF on the second a2a (combine path) (#28)

📊 Changes

72 files changed (+18234 additions, -362 deletions)

View changed files

➕ build_sgl_deep_gemm.sh (+157 -0)
📝 csrc/apis/attention.hpp (+4 -0)
📝 csrc/apis/einsum.hpp (+4 -3)
📝 csrc/apis/gemm.hpp (+4 -0)
📝 csrc/apis/hyperconnection.hpp (+4 -0)
📝 csrc/apis/layout.hpp (+4 -0)
📝 csrc/apis/mega.hpp (+86 -11)
📝 csrc/apis/runtime.hpp (+4 -0)
➕ csrc/apis/sm90_mega.hpp (+420 -0)
📝 csrc/jit_kernels/heuristics/mega_moe.hpp (+35 -11)
➕ csrc/jit_kernels/heuristics/sm90_mega_moe.hpp (+454 -0)
📝 csrc/jit_kernels/impls/runtime_utils.hpp (+3 -1)
📝 csrc/jit_kernels/impls/sm100_bf16_gemm.hpp (+1 -1)
📝 csrc/jit_kernels/impls/sm100_bmk_bnk_mn.hpp (+1 -1)
📝 csrc/jit_kernels/impls/sm100_fp8_fp4_gemm_1d1d.hpp (+1 -1)
📝 csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp (+120 -17)
➕ csrc/jit_kernels/impls/sm100_mega_moe_pre_dispatch.hpp (+175 -0)
📝 csrc/jit_kernels/impls/sm100_tf32_hc_prenorm_gemm.hpp (+1 -1)
📝 csrc/jit_kernels/impls/sm90_bf16_gemm.hpp (+1 -1)
📝 csrc/jit_kernels/impls/sm90_bmk_bnk_mn.hpp (+1 -1)

...and 52 more files

📄 Description

Summary

Add SM90/Hopper FP8 MegaMoE decode support to DeepGEMM.

This change introduces a fused SM90 MegaMoE kernel path for FP8 weights and FP8 activations, including the C++/CUDA implementation, JIT entry point, Python API wrapper, Hopper-specific weight transform, scheduling heuristics, and a unified Hopper accuracy/performance test script.

Main Changes

  • Added the SM90 FP8 MegaMoE fused kernel implementation:

    • deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh
    • csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp

  • Extended MegaMoE API bindings with fp8_mega_moe for SM90.

  • Added Hopper-specific Python entry points:

    • deep_gemm.fp8_mega_moe
    • deep_gemm.transform_weights_for_mega_moe_sm90

  • Added SM90 MegaMoE scheduling/config heuristics.

  • Updated MegaMoE symmetric buffer handling for SM90 FP32 scale-factor layouts.

  • Added tests/test_mega_moe_hopper.py, covering:

    • Layered accuracy checks
    • Fused-kernel benchmark mode
    • Normal DeepEP + grouped-FP8 baseline comparison
    • Low-latency DeepEP baseline comparison

DeepSeekV4Flash（8 card H20）

batch MegaMoE us normal us normal/MegaMoE LL us LL/MegaMoE
1 188.0 505.0 2.686x 250.1 1.330x
2 239.0 615.8 2.576x 299.9 1.255x
4 361.1 828.1 2.293x 416.0 1.152x
8 432.2 950.4 2.199x 494.8 1.145x
16 482.2 1074.1 2.227x 536.5 1.112x
32 499.0 1087.6 2.180x 540.4 1.083x
64 507.6 1081.6 2.131x 550.9 1.085x
128 511.0 1079.8 2.113x 562.5 1.101x
256 528.4 1178.6 2.231x 625.8 1.184x
512 874.1 1238.4 1.417x - -
1024 1638.1 2101.2 1.283x - -
2048 2856.6 3564.9 1.248x - -
4096 5172.0 6378.5 1.233x - -
8192 9784.6 11959.2 1.222x - -

DeepSeekV4Pro（8 card H20）

batch MegaMoE us normal us normal/MegaMoE LL us LL/MegaMoE
1 410.2 910.2 2.219x 493.9 1.204x
2 560.8 1296.6 2.312x 669.5 1.194x
4 843.0 1861.5 2.208x 986.0 1.170x
8 1257.1 2736.1 2.176x 1383.0 1.100x
16 1507.6 3179.8 2.109x 1628.9 1.080x
32 1550.4 3298.0 2.127x 1708.5 1.102x
64 1554.4 3281.2 2.111x 1711.8 1.101x
128 1584.0 3334.6 2.105x 1736.9 1.097x
256 1613.5 3381.5 2.096x 1805.6 1.119x
512 3014.2 3466.1 1.150x - -
1024 4775.0 5373.1 1.125x - -
2048 7838.6 8723.8 1.113x - -
4096 13666.4 15115.5 1.106x - -
8192 25470.8 28148.6 1.105x - -

Benchmark：DeepSeekV4Flash CP8/EP8

Metric MegaMoE off MegaMoE on
GSM8K accuracy 0.975 0.970
GSM8K invalid 0.000 0.000

SLO-Compliant Total Throughput

Input/Output off on Improvement
3500/1500 4766.11 5143.19 +7.9%
32000/1500 11760.82 15176.28 +29.0%
128000/1500 13714.90 16044.73 +17.0%

Max Throughput

Input/Output off on Improvement
3500/1500 4708.56 5002.54 +6.2%
32000/1500 13177.93 16434.02 +24.7%
128000/1500 15888.99 19490.54 +22.7%

@LyricZhao Could you help review this PR when you have time? If the patch is too large for one PR, I’m happy to split it into smaller parts following your preference.

🔄 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/323 **Author:** [@qiushixiaoyu](https://github.com/qiushixiaoyu) **Created:** 4/30/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `main` --- ### 📝 Commits (10+) - [`88db27f`](https://github.com/deepseek-ai/DeepGEMM/commit/88db27f9a39d6e45346c0a370cd672fda501b62e) Support tvm-ffi interfaces - [`495fcfd`](https://github.com/deepseek-ai/DeepGEMM/commit/495fcfd5b9b9f08272ecb2b765e5d06af2c2bfe9) Support Wheel Compilation - [`3036815`](https://github.com/deepseek-ai/DeepGEMM/commit/303681509b98657c325060d70aee32f1e9f102ef) Export the PDL utils of DeepGEMM (#34) - [`607503a`](https://github.com/deepseek-ai/DeepGEMM/commit/607503ae5921b73bee12dd1278d992cdecd69581) Expose BF16 grouped GEMM wrappers - [`00c42b8`](https://github.com/deepseek-ai/DeepGEMM/commit/00c42b8d8bc988aebf8a4eeaee8006f88a5061b0) Setup tests for sgl-deep-gemm - [`a8459f1`](https://github.com/deepseek-ai/DeepGEMM/commit/a8459f16a0a05184f0cbb7d1090e54f02baea372) Sm90 mega moe on sgl dev (#36) - [`eed753c`](https://github.com/deepseek-ai/DeepGEMM/commit/eed753cf15a41438eef9d0d64163e02525115980) Add Hopper mega moe test to runner (#46) - [`f3d0ca5`](https://github.com/deepseek-ai/DeepGEMM/commit/f3d0ca5e8eb5d3f090397facc4e959b071965114) chore: bump apache-tvm-ffi 0.1.9 -> 0.1.11 (#47) - [`bb569e2`](https://github.com/deepseek-ai/DeepGEMM/commit/bb569e29d5c962ac7bddc882fd3c69c6b5da1f9f) Add FP4 acts + MXF4 kind support and fused mega_moe_pre_dispatch kernel (#27) - [`b111331`](https://github.com/deepseek-ai/DeepGEMM/commit/b111331a8adbc781ca394343db75d60ca07d86b7) Add DG_USE_FP8_COMBINE: FP8 + per-row UE8M0 SF on the second a2a (combine path) (#28) ### 📊 Changes **72 files changed** (+18234 additions, -362 deletions) <details> <summary>View changed files</summary> ➕ `build_sgl_deep_gemm.sh` (+157 -0) 📝 `csrc/apis/attention.hpp` (+4 -0) 📝 `csrc/apis/einsum.hpp` (+4 -3) 📝 `csrc/apis/gemm.hpp` (+4 -0) 📝 `csrc/apis/hyperconnection.hpp` (+4 -0) 📝 `csrc/apis/layout.hpp` (+4 -0) 📝 `csrc/apis/mega.hpp` (+86 -11) 📝 `csrc/apis/runtime.hpp` (+4 -0) ➕ `csrc/apis/sm90_mega.hpp` (+420 -0) 📝 `csrc/jit_kernels/heuristics/mega_moe.hpp` (+35 -11) ➕ `csrc/jit_kernels/heuristics/sm90_mega_moe.hpp` (+454 -0) 📝 `csrc/jit_kernels/impls/runtime_utils.hpp` (+3 -1) 📝 `csrc/jit_kernels/impls/sm100_bf16_gemm.hpp` (+1 -1) 📝 `csrc/jit_kernels/impls/sm100_bmk_bnk_mn.hpp` (+1 -1) 📝 `csrc/jit_kernels/impls/sm100_fp8_fp4_gemm_1d1d.hpp` (+1 -1) 📝 `csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp` (+120 -17) ➕ `csrc/jit_kernels/impls/sm100_mega_moe_pre_dispatch.hpp` (+175 -0) 📝 `csrc/jit_kernels/impls/sm100_tf32_hc_prenorm_gemm.hpp` (+1 -1) 📝 `csrc/jit_kernels/impls/sm90_bf16_gemm.hpp` (+1 -1) 📝 `csrc/jit_kernels/impls/sm90_bmk_bnk_mn.hpp` (+1 -1) _...and 52 more files_ </details> ### 📄 Description ## Summary Add SM90/Hopper FP8 MegaMoE decode support to DeepGEMM. This change introduces a fused SM90 MegaMoE kernel path for FP8 weights and FP8 activations, including the C++/CUDA implementation, JIT entry point, Python API wrapper, Hopper-specific weight transform, scheduling heuristics, and a unified Hopper accuracy/performance test script. ## Main Changes - Added the SM90 FP8 MegaMoE fused kernel implementation: - `deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh` - `csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp` - Extended MegaMoE API bindings with `fp8_mega_moe` for SM90. - Added Hopper-specific Python entry points: - `deep_gemm.fp8_mega_moe` - `deep_gemm.transform_weights_for_mega_moe_sm90` - Added SM90 MegaMoE scheduling/config heuristics. - Updated MegaMoE symmetric buffer handling for SM90 FP32 scale-factor layouts. - Added `tests/test_mega_moe_hopper.py`, covering: - Layered accuracy checks - Fused-kernel benchmark mode - Normal DeepEP + grouped-FP8 baseline comparison - Low-latency DeepEP baseline comparison ## DeepSeekV4Flash（8 card H20） batch | MegaMoE us | normal us | normal/MegaMoE | LL us | LL/MegaMoE -- | -- | -- | -- | -- | -- 1 | 188.0 | 505.0 | 2.686x | 250.1 | 1.330x 2 | 239.0 | 615.8 | 2.576x | 299.9 | 1.255x 4 | 361.1 | 828.1 | 2.293x | 416.0 | 1.152x 8 | 432.2 | 950.4 | 2.199x | 494.8 | 1.145x 16 | 482.2 | 1074.1 | 2.227x | 536.5 | 1.112x 32 | 499.0 | 1087.6 | 2.180x | 540.4 | 1.083x 64 | 507.6 | 1081.6 | 2.131x | 550.9 | 1.085x 128 | 511.0 | 1079.8 | 2.113x | 562.5 | 1.101x 256 | 528.4 | 1178.6 | 2.231x | 625.8 | 1.184x 512 | 874.1 | 1238.4 | 1.417x | - | - 1024 | 1638.1 | 2101.2 | 1.283x | - | - 2048 | 2856.6 | 3564.9 | 1.248x | - | - 4096 | 5172.0 | 6378.5 | 1.233x | - | - 8192 | 9784.6 | 11959.2 | 1.222x | - | - ## DeepSeekV4Pro（8 card H20） batch | MegaMoE us | normal us | normal/MegaMoE | LL us | LL/MegaMoE -- | -- | -- | -- | -- | -- 1 | 410.2 | 910.2 | 2.219x | 493.9 | 1.204x 2 | 560.8 | 1296.6 | 2.312x | 669.5 | 1.194x 4 | 843.0 | 1861.5 | 2.208x | 986.0 | 1.170x 8 | 1257.1 | 2736.1 | 2.176x | 1383.0 | 1.100x 16 | 1507.6 | 3179.8 | 2.109x | 1628.9 | 1.080x 32 | 1550.4 | 3298.0 | 2.127x | 1708.5 | 1.102x 64 | 1554.4 | 3281.2 | 2.111x | 1711.8 | 1.101x 128 | 1584.0 | 3334.6 | 2.105x | 1736.9 | 1.097x 256 | 1613.5 | 3381.5 | 2.096x | 1805.6 | 1.119x 512 | 3014.2 | 3466.1 | 1.150x | - | - 1024 | 4775.0 | 5373.1 | 1.125x | - | - 2048 | 7838.6 | 8723.8 | 1.113x | - | - 4096 | 13666.4 | 15115.5 | 1.106x | - | - 8192 | 25470.8 | 28148.6 | 1.105x | - | - ## Benchmark：DeepSeekV4Flash CP8/EP8 | Metric | MegaMoE off | MegaMoE on | |---|---:|---:| | GSM8K accuracy | 0.975 | 0.970 | | GSM8K invalid | 0.000 | 0.000 | ### SLO-Compliant Total Throughput | Input/Output | off | on | Improvement | |---|---:|---:|---:| | 3500/1500 | 4766.11 | 5143.19 | +7.9% | | 32000/1500 | 11760.82 | 15176.28 | +29.0% | | 128000/1500 | 13714.90 | 16044.73 | +17.0% | ### Max Throughput | Input/Output | off | on | Improvement | |---|---:|---:|---:| | 3500/1500 | 4708.56 | 5002.54 | +6.2% | | 32000/1500 | 13177.93 | 16434.02 | +24.7% | | 128000/1500 | 15888.99 | 19490.54 | +22.7% | @LyricZhao Could you help review this PR when you have time? If the patch is too large for one PR, I’m happy to split it into smaller parts following your preference. --- <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#367
No description provided.