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 #409] [MERGED] Add SM100 NVFP4 MegaMoE with fused BF16 shared experts #425

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

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/409
Author: @zhou9402
Created: 8/17/2026
Status: ✅ Merged
Merged: 8/27/2026
Merged by: @RayWang96

Base: nv_dev ← Head: codex/nvfp4-fused-shared-oss-pr

📝 Commits (10+)

  • 9d96576 Add NVFP4 MegaMoE with fused BF16 shared experts
  • d0fdd21 Add accurate fused shared experts to NVFP4 MegaMoE
  • a76eede Trim NVFP4 MegaMoE integration surface
  • cb1f99d Address NVFP4 MegaMoE review feedback
  • 4d73079 Fix NVFP4 MegaMoE ring reuse and edge cases
  • 2deb26a Tighten NVFP4 MegaMoE validation
  • 67f0168 Guard byte-based GEMM heuristics from sub-byte types
  • 1807590 Keep legacy MegaMoE metadata compact
  • 814d9bc Validate NVFP4 MegaMoE buffer layout
  • 0c067ed Clarify NVFP4 heuristic and combine-slot invariants

📊 Changes

16 files changed (+2909 additions, -51 deletions)

View changed files

📝 csrc/apis/mega.hpp (+232 -15)
📝 csrc/jit_kernels/heuristics/common.hpp (+7 -0)
📝 csrc/jit_kernels/heuristics/mega_moe.hpp (+28 -13)
📝 csrc/jit_kernels/heuristics/sm100.hpp (+3 -2)
📝 csrc/jit_kernels/heuristics/sm120.hpp (+3 -3)
📝 csrc/jit_kernels/heuristics/sm90.hpp (+1 -1)
📝 csrc/jit_kernels/impls/runtime_utils.hpp (+2 -0)
➕ csrc/jit_kernels/impls/sm100_fp4_fp4_mega_moe.hpp (+306 -0)
📝 deep_gemm/__init__.py (+1 -0)
📝 deep_gemm/include/deep_gemm/common/types.cuh (+7 -1)
➕ deep_gemm/include/deep_gemm/impls/sm100_fp4_fp4_mega_moe.cuh (+1778 -0)
📝 deep_gemm/include/deep_gemm/layout/mega_moe.cuh (+37 -13)
📝 deep_gemm/include/deep_gemm/ptx/tcgen05.cuh (+24 -0)
📝 deep_gemm/include/deep_gemm/scheduler/mega_moe.cuh (+3 -2)
📝 deep_gemm/mega/__init__.py (+62 -1)
➕ tests/test_nvfp4_mega_moe.py (+415 -0)

📄 Description

Summary

  • Add a dedicated SM100 NVFP4 x NVFP4 MegaMoE kernel.
  • Fuse the optional BF16 shared expert into the same persistent MegaMoE kernel.
  • Apply routed scaling after the routed top-k reduction and before adding the shared-expert output.
  • Preserve ModelOpt NVFP4 scales without affecting the existing BF16 or FP8 x FP4 MegaMoE paths.

Design highlights

  • Native NVFP4 routed path. Both routed activations and weights use packed E2M1 data with per-16-element E4M3 scale factors. This requires a bit-aware symmetric-buffer layout, NVFP4-specific TMA/UTCCP handling, and an in-kernel L1 requantization epilogue.

  • Post-L2 top-k weighting. Top-k routing weights are stored in a new NVFP4-specific buffer and applied to the BF16 routed output after L2, immediately before the combine write. Applying top-k weights directly in the L1 epilogue can reduce the intermediate magnitude before NVFP4 quantization and introduce additional precision loss.

  • ModelOpt scale handling. Per-expert gate/up scales are applied before SwiGLU. The down-projection input scale normalizes the in-kernel NVFP4 intermediate quantization and is folded back into the L2 output scale: l2_scale = l2_alpha * a2_scale.

  • Correct routed/shared ordering. The kernel accumulates the routed top-k outputs, rounds the routed sum to BF16, applies routed_scaling_factor, and then adds the BF16 shared-expert output.

  • BF16 shared expert inside MegaMoE. Optional BF16 shared L1 and L2 work is inserted into the existing scheduler as SharedLinear1 and SharedLinear2 phases. These tasks reuse the existing TMA-load, MMA, and epilogue warp roles; shared computation uses ordinary BF16 UMMA without scale factors or UTCCP.

  • Separate shared K tiling. Routed NVFP4 and shared BF16 tasks reuse the same shared-memory byte tiles but contain different numbers of elements, so SHARED_BLOCK_K describes the BF16 K extent independently from the routed NVFP4 K extent.

  • Fused shared epilogues. Shared L1 applies SwiGLU and stores a BF16 intermediate. Shared L2 writes to an additional local combine slot, allowing the final combine phase to perform routed scaling and the shared add without separate shared-expert launches.

  • Compatibility. The new NVFP4 path does not affect the existing BF16 or FP8 x FP4 MegaMoE paths.

Validation

Test Result
Boundary/adversarial accuracy, 42 cases PASS; max rel-L2 0.005421
Random fuzz, 22 cases PASS; max rel-L2 0.002655
Shared-only accuracy, 10 cases PASS; max rel-L2 0.001134
GLM-5.2 layer replay, 75 layers PASS; fused closer to FP32 in 69/75
GLM-5.2 deterministic E2E, 0/8K context PASS; serial 32/32, fused 32/32
Clean GB300 build and 4-rank correctness PASS; max rel-L2 0.001315

Performance

Pure CUDA-kernel time on 8x GB300 with EP8. All three arms have the same BF16-input to BF16-output model boundary and include routed computation, dispatch/combine, routed scaling, the BF16 shared expert, and the final add.

  • FlashInfer modular baseline: FlashInfer one-sided All-to-All dispatch -> TRT-LLM NVFP4 routed experts -> FlashInfer one-sided All-to-All combine -> routed scaling -> serial DeepGEMM BF16 shared expert -> add. This arm does not call MegaMoE. Optimized FlashInfer configurations are used, and the NVFP4 expert tactic is autotuned independently for each workload size. For each M, the table reports the lowest audited latency among runs with the correct tactic and identical timing boundary.
  • DG unfused: routed-only sm100_fp4_fp4_mega_moe -> routed scaling -> serial DeepGEMM BF16 shared expert -> add.
  • DG fused: one sm100_fp4_fp4_mega_moe kernel containing the routed path, routed scaling, BF16 shared expert, and final add.
M/rank FlashInfer modular, no MegaMoE (us) DG routed-only + serial shared (us) DG fused shared (us) Fused vs FI Fusion gain
1 132.064 131.584 105.504 20.11% 19.82%
2 154.096 146.608 124.784 19.02% 14.89%
4 168.720 164.816 140.464 16.75% 14.78%
8 189.232 183.104 160.512 15.18% 12.34%
16 199.360 190.528 166.672 16.40% 12.52%
32 203.792 192.400 168.464 17.34% 12.44%
64 209.184 195.232 171.200 18.16% 12.31%
128 257.328 203.536 179.216 30.36% 11.95%
256 268.633 221.920 190.960 28.91% 13.95%
512 346.718 275.192 234.177 32.46% 14.90%
1024 539.041 417.849 365.346 32.22% 12.57%
2048 918.737 690.976 604.784 34.17% 12.47%
4096 1721.712 1400.384 1248.720 27.47% 10.83%
8192 3345.690 2539.136 2218.608 33.69% 12.62%
16384 6559.996 5400.112 4871.792 25.73% 9.78%

Fusion gain compares the two DeepGEMM arms. Fused vs FI compares against the modular, non-MegaMoE FlashInfer baseline.

Complete test recipe

Revisions and environment

  • Hardware: 8x GB300, two hosts x four GPUs, EP8
  • CUDA 13, PyTorch 2.13.0+cu130, vLLM 0.17.2rc1.dev4473+g3ada0d7f1, FlashInfer 0.6.16rc5

Build and correctness


CUDA_VISIBLE_DEVICES=0,1,2,3 \
python tests/test_nvfp4_mega_moe.py \
  --num-processes 4 \
  --num-max-tokens-per-rank 128 \
  --num-tokens 128 \
  --hidden 6144 \
  --intermediate-hidden 2048 \
  --num-experts 256 \
  --num-topk 8 \
  --num-shared-experts 1 \
  --per-expert-alphas 1 \
  --per-expert-a2 1 \
  --routed-scaling-factor 0.75 \
  --num-correctness-tests 2

The boundary, zero-token, masked-routing, and shared-only cases use the same command with the corresponding --num-tokens, --zero-token-rank, --masked-ratio, --routing-hot-rank, and --num-shared-experts settings.

Performance shape and inputs

  • Hidden size: 6144
  • Routed/shared intermediate size: 2048
  • Routed experts: 256; top-k: 8; BF16 shared experts: 1
  • Routed scaling factor: 0.75
  • M/rank: 1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384
  • Symmetric-buffer token capacity per rank: 16512
  • Deterministic seed: 20260811
  • All arms use identical BF16 inputs, top-k metadata, ModelOpt NVFP4 weights/scales, and BF16 shared weights.
  • Data generation, weight conversion, JIT compilation, and backend autotuning are completed before measurement.

Timed call graphs

FI modular:
  BF16->NVFP4 quantization
  -> FlashInfer one-sided dispatch prepare/sanitize/dispatch
  -> TRT-LLM routing + NVFP4 BMM1 + NVFP4 BMM2 + finalize
  -> FlashInfer combine prepare/combine
  -> routed BF16 scale
  -> DeepGEMM BF16 shared L1 + SwiGLU + BF16 shared L2
  -> BF16 add

DG unfused:
  routed-only sm100_fp4_fp4_mega_moe
  -> routed BF16 scale
  -> DeepGEMM BF16 shared L1 + SwiGLU + BF16 shared L2
  -> BF16 add

DG fused:
  sm100_fp4_fp4_mega_moe with routed scaling and BF16 shared expert enabled

Timing

  • Timing follows deep_gemm.testing.bench_kineto and reports selected CUDA-kernel time only.
  • Perform five untimed warmup calls before profiling.
  • Use Kineto schedule wait=0, warmup=1, active=1 with 30 calls in the measured phase.
  • Before each measured call, evict L2 with the repository benchmark's 8 GB eviction buffer. Eviction, CUDA sleep, distributed barriers, CPU launch gaps, router/top-k selection, data generation, weight conversion, JIT compilation, and autotuning are excluded.
  • FI backend-internal quantization, routing, dispatch, expert, finalize, and combine CUDA kernels are included.
  • Sum the selected CUDA rows per rank and report the maximum across all eight ranks.

🔄 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/409 **Author:** [@zhou9402](https://github.com/zhou9402) **Created:** 8/17/2026 **Status:** ✅ Merged **Merged:** 8/27/2026 **Merged by:** [@RayWang96](https://github.com/RayWang96) **Base:** `nv_dev` ← **Head:** `codex/nvfp4-fused-shared-oss-pr` --- ### 📝 Commits (10+) - [`9d96576`](https://github.com/deepseek-ai/DeepGEMM/commit/9d96576ecbd148872f20c2f90ae073137ef1d58c) Add NVFP4 MegaMoE with fused BF16 shared experts - [`d0fdd21`](https://github.com/deepseek-ai/DeepGEMM/commit/d0fdd21540bdcb70b2c376512aa6cc89b61551df) Add accurate fused shared experts to NVFP4 MegaMoE - [`a76eede`](https://github.com/deepseek-ai/DeepGEMM/commit/a76eedef576da1091dd1cfdf0330001d2073119f) Trim NVFP4 MegaMoE integration surface - [`cb1f99d`](https://github.com/deepseek-ai/DeepGEMM/commit/cb1f99df066de9f32ae50aacc80ff6ee57fa65e4) Address NVFP4 MegaMoE review feedback - [`4d73079`](https://github.com/deepseek-ai/DeepGEMM/commit/4d730794143fe4562132cc7974b5262cf62716ad) Fix NVFP4 MegaMoE ring reuse and edge cases - [`2deb26a`](https://github.com/deepseek-ai/DeepGEMM/commit/2deb26abcfb66afce8cca9f9a3fc7127697cb03e) Tighten NVFP4 MegaMoE validation - [`67f0168`](https://github.com/deepseek-ai/DeepGEMM/commit/67f0168c5a3152c973378dc94466f414181ff5d4) Guard byte-based GEMM heuristics from sub-byte types - [`1807590`](https://github.com/deepseek-ai/DeepGEMM/commit/1807590cbf9903c525397b157fc463b915c90c89) Keep legacy MegaMoE metadata compact - [`814d9bc`](https://github.com/deepseek-ai/DeepGEMM/commit/814d9bc8547864ae2183c33b78caecb49d91d767) Validate NVFP4 MegaMoE buffer layout - [`0c067ed`](https://github.com/deepseek-ai/DeepGEMM/commit/0c067ed4c845c20564ee230a5105529bafbaa959) Clarify NVFP4 heuristic and combine-slot invariants ### 📊 Changes **16 files changed** (+2909 additions, -51 deletions) <details> <summary>View changed files</summary> 📝 `csrc/apis/mega.hpp` (+232 -15) 📝 `csrc/jit_kernels/heuristics/common.hpp` (+7 -0) 📝 `csrc/jit_kernels/heuristics/mega_moe.hpp` (+28 -13) 📝 `csrc/jit_kernels/heuristics/sm100.hpp` (+3 -2) 📝 `csrc/jit_kernels/heuristics/sm120.hpp` (+3 -3) 📝 `csrc/jit_kernels/heuristics/sm90.hpp` (+1 -1) 📝 `csrc/jit_kernels/impls/runtime_utils.hpp` (+2 -0) ➕ `csrc/jit_kernels/impls/sm100_fp4_fp4_mega_moe.hpp` (+306 -0) 📝 `deep_gemm/__init__.py` (+1 -0) 📝 `deep_gemm/include/deep_gemm/common/types.cuh` (+7 -1) ➕ `deep_gemm/include/deep_gemm/impls/sm100_fp4_fp4_mega_moe.cuh` (+1778 -0) 📝 `deep_gemm/include/deep_gemm/layout/mega_moe.cuh` (+37 -13) 📝 `deep_gemm/include/deep_gemm/ptx/tcgen05.cuh` (+24 -0) 📝 `deep_gemm/include/deep_gemm/scheduler/mega_moe.cuh` (+3 -2) 📝 `deep_gemm/mega/__init__.py` (+62 -1) ➕ `tests/test_nvfp4_mega_moe.py` (+415 -0) </details> ### 📄 Description ## Summary - Add a dedicated SM100 NVFP4 x NVFP4 MegaMoE kernel. - Fuse the optional BF16 shared expert into the same persistent MegaMoE kernel. - Apply routed scaling after the routed top-k reduction and before adding the shared-expert output. - Preserve ModelOpt NVFP4 scales without affecting the existing BF16 or FP8 x FP4 MegaMoE paths. ## Design highlights - **Native NVFP4 routed path.** Both routed activations and weights use packed E2M1 data with per-16-element E4M3 scale factors. This requires a bit-aware symmetric-buffer layout, NVFP4-specific TMA/UTCCP handling, and an in-kernel L1 requantization epilogue. - **Post-L2 top-k weighting.** Top-k routing weights are stored in a new NVFP4-specific buffer and applied to the BF16 routed output after L2, immediately before the combine write. Applying top-k weights directly in the L1 epilogue can reduce the intermediate magnitude before NVFP4 quantization and introduce additional precision loss. - **ModelOpt scale handling.** Per-expert gate/up scales are applied before SwiGLU. The down-projection input scale normalizes the in-kernel NVFP4 intermediate quantization and is folded back into the L2 output scale: `l2_scale = l2_alpha * a2_scale`. - **Correct routed/shared ordering.** The kernel accumulates the routed top-k outputs, rounds the routed sum to BF16, applies `routed_scaling_factor`, and then adds the BF16 shared-expert output. - **BF16 shared expert inside MegaMoE.** Optional BF16 shared L1 and L2 work is inserted into the existing scheduler as `SharedLinear1` and `SharedLinear2` phases. These tasks reuse the existing TMA-load, MMA, and epilogue warp roles; shared computation uses ordinary BF16 UMMA without scale factors or UTCCP. - **Separate shared K tiling.** Routed NVFP4 and shared BF16 tasks reuse the same shared-memory byte tiles but contain different numbers of elements, so `SHARED_BLOCK_K` describes the BF16 K extent independently from the routed NVFP4 K extent. - **Fused shared epilogues.** Shared L1 applies SwiGLU and stores a BF16 intermediate. Shared L2 writes to an additional local combine slot, allowing the final combine phase to perform routed scaling and the shared add without separate shared-expert launches. - **Compatibility.** The new NVFP4 path does not affect the existing BF16 or FP8 x FP4 MegaMoE paths. ## Validation | Test | Result | |---|---:| | Boundary/adversarial accuracy, 42 cases | PASS; max rel-L2 0.005421 | | Random fuzz, 22 cases | PASS; max rel-L2 0.002655 | | Shared-only accuracy, 10 cases | PASS; max rel-L2 0.001134 | | GLM-5.2 layer replay, 75 layers | PASS; fused closer to FP32 in 69/75 | | GLM-5.2 deterministic E2E, 0/8K context | PASS; serial 32/32, fused 32/32 | | Clean GB300 build and 4-rank correctness | PASS; max rel-L2 0.001315 | ## Performance Pure CUDA-kernel time on 8x GB300 with EP8. All three arms have the same BF16-input to BF16-output model boundary and include routed computation, dispatch/combine, routed scaling, the BF16 shared expert, and the final add. - **FlashInfer modular baseline:** FlashInfer one-sided All-to-All dispatch -> TRT-LLM NVFP4 routed experts -> FlashInfer one-sided All-to-All combine -> routed scaling -> serial DeepGEMM BF16 shared expert -> add. This arm does not call MegaMoE. Optimized FlashInfer configurations are used, and the NVFP4 expert tactic is autotuned independently for each workload size. For each M, the table reports the lowest audited latency among runs with the correct tactic and identical timing boundary. - **DG unfused:** routed-only `sm100_fp4_fp4_mega_moe` -> routed scaling -> serial DeepGEMM BF16 shared expert -> add. - **DG fused:** one `sm100_fp4_fp4_mega_moe` kernel containing the routed path, routed scaling, BF16 shared expert, and final add. | M/rank | FlashInfer modular, no MegaMoE (us) | DG routed-only + serial shared (us) | DG fused shared (us) | Fused vs FI | Fusion gain | |---:|---:|---:|---:|---:|---:| | 1 | 132.064 | 131.584 | 105.504 | 20.11% | 19.82% | | 2 | 154.096 | 146.608 | 124.784 | 19.02% | 14.89% | | 4 | 168.720 | 164.816 | 140.464 | 16.75% | 14.78% | | 8 | 189.232 | 183.104 | 160.512 | 15.18% | 12.34% | | 16 | 199.360 | 190.528 | 166.672 | 16.40% | 12.52% | | 32 | 203.792 | 192.400 | 168.464 | 17.34% | 12.44% | | 64 | 209.184 | 195.232 | 171.200 | 18.16% | 12.31% | | 128 | 257.328 | 203.536 | 179.216 | 30.36% | 11.95% | | 256 | 268.633 | 221.920 | 190.960 | 28.91% | 13.95% | | 512 | 346.718 | 275.192 | 234.177 | 32.46% | 14.90% | | 1024 | 539.041 | 417.849 | 365.346 | 32.22% | 12.57% | | 2048 | 918.737 | 690.976 | 604.784 | 34.17% | 12.47% | | 4096 | 1721.712 | 1400.384 | 1248.720 | 27.47% | 10.83% | | 8192 | 3345.690 | 2539.136 | 2218.608 | 33.69% | 12.62% | | 16384 | 6559.996 | 5400.112 | 4871.792 | 25.73% | 9.78% | `Fusion gain` compares the two DeepGEMM arms. `Fused vs FI` compares against the modular, non-MegaMoE FlashInfer baseline. <details> <summary>Complete test recipe</summary> ### Revisions and environment - Hardware: 8x GB300, two hosts x four GPUs, EP8 - CUDA 13, PyTorch `2.13.0+cu130`, vLLM `0.17.2rc1.dev4473+g3ada0d7f1`, FlashInfer `0.6.16rc5` ### Build and correctness ```bash CUDA_VISIBLE_DEVICES=0,1,2,3 \ python tests/test_nvfp4_mega_moe.py \ --num-processes 4 \ --num-max-tokens-per-rank 128 \ --num-tokens 128 \ --hidden 6144 \ --intermediate-hidden 2048 \ --num-experts 256 \ --num-topk 8 \ --num-shared-experts 1 \ --per-expert-alphas 1 \ --per-expert-a2 1 \ --routed-scaling-factor 0.75 \ --num-correctness-tests 2 ``` The boundary, zero-token, masked-routing, and shared-only cases use the same command with the corresponding `--num-tokens`, `--zero-token-rank`, `--masked-ratio`, `--routing-hot-rank`, and `--num-shared-experts` settings. ### Performance shape and inputs - Hidden size: `6144` - Routed/shared intermediate size: `2048` - Routed experts: `256`; top-k: `8`; BF16 shared experts: `1` - Routed scaling factor: `0.75` - M/rank: `1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384` - Symmetric-buffer token capacity per rank: `16512` - Deterministic seed: `20260811` - All arms use identical BF16 inputs, top-k metadata, ModelOpt NVFP4 weights/scales, and BF16 shared weights. - Data generation, weight conversion, JIT compilation, and backend autotuning are completed before measurement. ### Timed call graphs ```text FI modular: BF16->NVFP4 quantization -> FlashInfer one-sided dispatch prepare/sanitize/dispatch -> TRT-LLM routing + NVFP4 BMM1 + NVFP4 BMM2 + finalize -> FlashInfer combine prepare/combine -> routed BF16 scale -> DeepGEMM BF16 shared L1 + SwiGLU + BF16 shared L2 -> BF16 add DG unfused: routed-only sm100_fp4_fp4_mega_moe -> routed BF16 scale -> DeepGEMM BF16 shared L1 + SwiGLU + BF16 shared L2 -> BF16 add DG fused: sm100_fp4_fp4_mega_moe with routed scaling and BF16 shared expert enabled ``` ### Timing - Timing follows `deep_gemm.testing.bench_kineto` and reports selected CUDA-kernel time only. - Perform five untimed warmup calls before profiling. - Use Kineto schedule `wait=0, warmup=1, active=1` with 30 calls in the measured phase. - Before each measured call, evict L2 with the repository benchmark's 8 GB eviction buffer. Eviction, CUDA sleep, distributed barriers, CPU launch gaps, router/top-k selection, data generation, weight conversion, JIT compilation, and autotuning are excluded. - FI backend-internal quantization, routing, dispatch, expert, finalize, and combine CUDA kernels are included. - Sum the selected CUDA rows per rank and report the maximum across all eight ranks. </details> --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>
bsakaguchi 2026-10-26 21:17:38 +00:00
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#425
No description provided.