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 #360] Add SM90 MegaMoE by cooperative #390

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

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/360
Author: @mpdfdfl
Created: 6/13/2026
Status: 🔄 Open

Base: main ← Head: sm90-mega-moe-pingpong-coop

📝 Commits (10+)

  • 006709a sm90 megamoe: add pingpong and cooperative FP8 kernels
  • 4dd2e94 sm90 megamoe: token-count dispatch and Python bindings
  • 1f84636 sm90 megamoe: correctness test and performance benchmark
  • 4cd92c4 chore: ignore profiling artifacts (*.ncu-rep, *.nsys-rep, profile_logs/)
  • 4146d0e fix mul2
  • 3f16c98 add sm90 megamoe bench
  • d7c89bd fix accuracy
  • 183fec5 fix accuracy
  • 6b5be1d add tma to HBM double buffer
  • 5fdbdda fix accuract

📊 Changes

14 files changed (+5003 additions, -8 deletions)

View changed files

📝 .gitignore (+6 -1)
📝 csrc/apis/mega.hpp (+141 -6)
➕ csrc/jit_kernels/heuristics/sm90_mega_moe.hpp (+235 -0)
➕ csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp (+229 -0)
📝 deep_gemm/__init__.py (+2 -0)
📝 deep_gemm/include/deep_gemm/comm/barrier.cuh (+6 -0)
📝 deep_gemm/include/deep_gemm/common/math.cuh (+9 -1)
➕ deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh (+1364 -0)
➕ deep_gemm/include/deep_gemm/scheduler/sm90_mega_moe.cuh (+241 -0)
📝 deep_gemm/mega/__init__.py (+57 -0)
➕ tests/bench_mega_moe_sm90.py (+1102 -0)
➕ tests/test_l1_single_token.py (+221 -0)
➕ tests/test_mega_moe_sm90.py (+711 -0)
➕ tests/test_mega_moe_sm90_mimo.py (+679 -0)

📄 Description

Add SM90 MegaMoE (cooperative)

Co-Authored-By: @OpenDarrenlu

First, thanks to #352 and #323 — this work builds on the ideas from those PRs.

We provide the SM90 fused MegaMoE through a single cooperative kernel
(BLOCK_M=64, BLOCK_N=256). The two math warpgroups cooperatively N-split one
tile — the first warpgroup owns the left 128 columns and the second owns the
right 128 columns (each an m64n128 WGMMA) — and share a single A-tile (weight)
load.

Why cooperative-only, and why the 64×256 tiling

We discovered an accuracy issue during end-to-end testing. We suspect the cause
is that the L1 output of the MoE layer is normally quantized per-128, whereas
our earlier implementation quantized it per-64. To fix this, we settled on a
single cooperative kernel in which the first warpgroup handles the left 128
columns and the other handles the right 128 columns, so the post-SwiGLU L1 output
is quantized per-128 (matching the standard DeepEP runner's
scale_block_size=128).

For the L2 GEMM, cooperative also beats a pingpong variant we tested: the
dequant FMAs inside the k-loop overlap the asynchronous WGMMAs, and the two
warpgroups share a single A-tile load.

The SM90 kernel / scheduler / heuristics are kept separate from the shared SM100
files, so SM100 behavior is unchanged.

Correctness test

tests/test_mega_moe_sm90.py checks the fused kernel against a PyTorch reference
(calc_diff < 0.01) over layered scenarios (smoke → shapes → edge cases →
random stress).

python tests/test_mega_moe_sm90.py --num-processes 8

Benchmark

tests/bench_mega_moe_sm90.py times the fused kernel end-to-end against the
DeepEP + DeepGEMM unfused baselines (V1 contiguous, V1 low-latency, V2
ElasticBuffer), same CUDA-event timing for all.

NVSHMEM_IBGDA_ENABLE=0 NVSHMEM_DISABLE_IBGDA=1 EP_DISABLE_GIN=1 \
python tests/bench_mega_moe_sm90.py --num-processes 8 \
  --hidden 4096 --intermediate-hidden 2048 --num-experts 256 --num-topk 8 \
  --baseline --baseline-version both \
  --batches 8 128 256 512 1024 2048 4096 8192

Here are our latest test results.

1. Single-operator benchmark (vs. deepep)

Shape — hidden=4096, moe_intermediate=2048, experts=256, topk=8 (MiMo-V2.5)

per-rank tokens ours (μs) V1-contig (μs) vs V1-contig V1-ll (μs) vs V1-ll V2 (μs) vs V2
8 262.3 1649.1 6.29x 272.6 1.04x 332.5 1.27x
128 282.5 1533.7 5.43x 327.1 1.16x 380.4 1.35x
256 348.4 1635.3 4.69x 389.3 1.12x 434.8 1.25x
512 497.4 1900.7 3.82x 557.0 1.12x 566.5 1.14x
1024 767.0 2336.9 3.05x 899.6 1.17x 888.0 1.16x
2048 1302.8 3341.1 2.56x 1597.4 1.23x 1543.5 1.18x
4096 2323.0 5320.2 2.29x - - 2878.6 1.24x
8192 4611.9 9408.6 2.04x - - 5516.0 1.20x

2. End-to-end prefill benchmark

Setup: input 16K, output 1, TTFT capped within 1s.

QPS Input token throughput (tok/s) TTFT (ms)
deepep 4.10 69635.40 1005.80
megamoe 5.15 86653.93 1037.03

📈 Compared with deepep, input token throughput improves by 25% (86653 vs 69635 tok/s).

3. Decode benchmark

Since this is a single-node deployment, TPOT is derived from the logs.

bs megamoe (median) deepep (median) Gap
36 18.69 ms 19.43 ms megamoe +4.0%
64 20.14 ms 20.94 ms megamoe +4.0%

TODO / future work

This is just a first version, and there is still a lot of room to optimize
inside the kernel. We would really love your help — let's optimize this kernel
together! Planned directions:

  • Block-tile scheduling. The L2 cache hit rate is currently low (~40%);
    tuning the block-tile / expert-wave scheduling strategy should reduce HBM
    traffic and help the large-M DRAM-bandwidth-bound cases.
  • Green-context split-kernel. Explore the approach from PR #357 — split the
    MoE into focused kernels (dispatch+L1+SwiGLU / L2+combine / combine-reduce)
    running concurrently on disjoint SM partitions via CUDA green contexts — to
    overlap stages the fused megakernel serializes internally.

🔄 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/360 **Author:** [@mpdfdfl](https://github.com/mpdfdfl) **Created:** 6/13/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `sm90-mega-moe-pingpong-coop` --- ### 📝 Commits (10+) - [`006709a`](https://github.com/deepseek-ai/DeepGEMM/commit/006709a45b2eb97b438f27a69c6264abb8dfbc8f) sm90 megamoe: add pingpong and cooperative FP8 kernels - [`4dd2e94`](https://github.com/deepseek-ai/DeepGEMM/commit/4dd2e94e9150b79aea5165c865a479ed72a93cb1) sm90 megamoe: token-count dispatch and Python bindings - [`1f84636`](https://github.com/deepseek-ai/DeepGEMM/commit/1f846368996038d8f4dbd97d428932aed4753c99) sm90 megamoe: correctness test and performance benchmark - [`4cd92c4`](https://github.com/deepseek-ai/DeepGEMM/commit/4cd92c4cca06d9b8bdbd94f86a464be913aaf81a) chore: ignore profiling artifacts (*.ncu-rep, *.nsys-rep, profile_logs/) - [`4146d0e`](https://github.com/deepseek-ai/DeepGEMM/commit/4146d0eee6e9268ef6222b992625e8072dfb61f8) fix mul2 - [`3f16c98`](https://github.com/deepseek-ai/DeepGEMM/commit/3f16c9850ce38f805c93ef7aa73e8f7c09c72ffe) add sm90 megamoe bench - [`d7c89bd`](https://github.com/deepseek-ai/DeepGEMM/commit/d7c89bdeecde60029b6c078c041881a8bbcbf975) fix accuracy - [`183fec5`](https://github.com/deepseek-ai/DeepGEMM/commit/183fec5fc197623fba1ea85ce1f4d1a53c111bb6) fix accuracy - [`6b5be1d`](https://github.com/deepseek-ai/DeepGEMM/commit/6b5be1d36c1d318999c795b9c3f6a67b4960a12a) add tma to HBM double buffer - [`5fdbdda`](https://github.com/deepseek-ai/DeepGEMM/commit/5fdbdda705505d0f1e83b3a530f40cdab0a9e280) fix accuract ### 📊 Changes **14 files changed** (+5003 additions, -8 deletions) <details> <summary>View changed files</summary> 📝 `.gitignore` (+6 -1) 📝 `csrc/apis/mega.hpp` (+141 -6) ➕ `csrc/jit_kernels/heuristics/sm90_mega_moe.hpp` (+235 -0) ➕ `csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp` (+229 -0) 📝 `deep_gemm/__init__.py` (+2 -0) 📝 `deep_gemm/include/deep_gemm/comm/barrier.cuh` (+6 -0) 📝 `deep_gemm/include/deep_gemm/common/math.cuh` (+9 -1) ➕ `deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh` (+1364 -0) ➕ `deep_gemm/include/deep_gemm/scheduler/sm90_mega_moe.cuh` (+241 -0) 📝 `deep_gemm/mega/__init__.py` (+57 -0) ➕ `tests/bench_mega_moe_sm90.py` (+1102 -0) ➕ `tests/test_l1_single_token.py` (+221 -0) ➕ `tests/test_mega_moe_sm90.py` (+711 -0) ➕ `tests/test_mega_moe_sm90_mimo.py` (+679 -0) </details> ### 📄 Description # Add SM90 MegaMoE (cooperative) Co-Authored-By: @OpenDarrenlu First, thanks to #352 and #323 — this work builds on the ideas from those PRs. We provide the SM90 fused MegaMoE through a single **cooperative** kernel (`BLOCK_M=64`, `BLOCK_N=256`). The two math warpgroups cooperatively N-split one tile — the first warpgroup owns the left 128 columns and the second owns the right 128 columns (each an `m64n128` WGMMA) — and share a single A-tile (weight) load. ### Why cooperative-only, and why the 64×256 tiling We discovered an accuracy issue during end-to-end testing. We suspect the cause is that the L1 output of the MoE layer is normally quantized **per-128**, whereas our earlier implementation quantized it **per-64**. To fix this, we settled on a single cooperative kernel in which the first warpgroup handles the left 128 columns and the other handles the right 128 columns, so the post-SwiGLU L1 output is quantized **per-128** (matching the standard DeepEP runner's `scale_block_size=128`). For the L2 GEMM, cooperative also beats a pingpong variant we tested: the dequant FMAs inside the k-loop overlap the asynchronous WGMMAs, and the two warpgroups share a single A-tile load. The SM90 kernel / scheduler / heuristics are kept separate from the shared SM100 files, so SM100 behavior is unchanged. ## Correctness test `tests/test_mega_moe_sm90.py` checks the fused kernel against a PyTorch reference (`calc_diff < 0.01`) over layered scenarios (smoke → shapes → edge cases → random stress). ```bash python tests/test_mega_moe_sm90.py --num-processes 8 ``` ## Benchmark `tests/bench_mega_moe_sm90.py` times the fused kernel end-to-end against the DeepEP + DeepGEMM unfused baselines (V1 contiguous, V1 low-latency, V2 ElasticBuffer), same CUDA-event timing for all. ```bash NVSHMEM_IBGDA_ENABLE=0 NVSHMEM_DISABLE_IBGDA=1 EP_DISABLE_GIN=1 \ python tests/bench_mega_moe_sm90.py --num-processes 8 \ --hidden 4096 --intermediate-hidden 2048 --num-experts 256 --num-topk 8 \ --baseline --baseline-version both \ --batches 8 128 256 512 1024 2048 4096 8192 ``` --- Here are our latest test results. ## 1. Single-operator benchmark (vs. deepep) **Shape — hidden=4096, moe_intermediate=2048, experts=256, topk=8 (MiMo-V2.5)** | per-rank tokens | ours (μs) | V1-contig (μs) | vs V1-contig | V1-ll (μs) | vs V1-ll | V2 (μs) | vs V2 | | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | | 8 | 262.3 | 1649.1 | 6.29x | 272.6 | 1.04x | 332.5 | 1.27x | | 128 | 282.5 | 1533.7 | 5.43x | 327.1 | 1.16x | 380.4 | 1.35x | | 256 | 348.4 | 1635.3 | 4.69x | 389.3 | 1.12x | 434.8 | 1.25x | | 512 | 497.4 | 1900.7 | 3.82x | 557.0 | 1.12x | 566.5 | 1.14x | | 1024 | 767.0 | 2336.9 | 3.05x | 899.6 | 1.17x | 888.0 | 1.16x | | 2048 | 1302.8 | 3341.1 | 2.56x | 1597.4 | 1.23x | 1543.5 | 1.18x | | 4096 | 2323.0 | 5320.2 | 2.29x | - | - | 2878.6 | 1.24x | | 8192 | 4611.9 | 9408.6 | 2.04x | - | - | 5516.0 | 1.20x | ## 2. End-to-end prefill benchmark **Setup: input 16K, output 1, TTFT capped within 1s.** | | QPS | Input token throughput (tok/s) | TTFT (ms) | | -------- | ---- | ------------------------------ | --------- | | deepep | 4.10 | 69635.40 | 1005.80 | | megamoe | 5.15 | 86653.93 | 1037.03 | 📈 Compared with deepep, **input token throughput improves by 25%** (86653 vs 69635 tok/s). ## 3. Decode benchmark Since this is a single-node deployment, TPOT is derived from the logs. | bs | megamoe (median) | deepep (median) | Gap | | -- | ---------------- | --------------- | ------------- | | 36 | 18.69 ms | 19.43 ms | megamoe +4.0% | | 64 | 20.14 ms | 20.94 ms | megamoe +4.0% | ## TODO / future work This is just a first version, and there is still a lot of room to optimize inside the kernel. **We would really love your help — let's optimize this kernel together!** Planned directions: - **Block-tile scheduling.** The L2 cache hit rate is currently low (~40%); tuning the block-tile / expert-wave scheduling strategy should reduce HBM traffic and help the large-M DRAM-bandwidth-bound cases. - **Green-context split-kernel.** Explore the approach from PR #357 — split the MoE into focused kernels (dispatch+L1+SwiGLU / L2+combine / combine-reduce) running concurrently on disjoint SM partitions via CUDA green contexts — to overlap stages the fused megakernel serializes internally. --- <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#390
No description provided.