mirror of
https://github.com/deepseek-ai/DeepGEMM.git
synced 2026-11-05 00:47:26 +00:00
[PR #357] Draft: Add green-context split-kernel MegaMoE features #389
Labels
No labels
pull-request
No milestone
No project
No assignees
1 participant
Notifications
Due date
No due date set.
Dependencies
No dependencies set.
Reference
deepseek-ai/DeepGEMM#389
Loading…
Add table
Add a link
Reference in a new issue
No description provided.
Delete branch "%!s()"
Deleting a branch is permanent. Although the deleted branch may continue to exist for a short time before it actually gets removed, it CANNOT be undone in most cases. Continue?
📋 Pull Request Information
Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/357
Author: @RayWang96
Created: 6/12/2026
Status: 🔄 Open
Base:
nv_dev← Head:split_mega_moe📝 Commits (1)
bb83742Add green-context split-kernel MegaMoE (dispatch_l1_swiglu / l2_combine / combine_reduce)📊 Changes
13 files changed (+3803 additions, -4 deletions)
View changed files
📝
csrc/apis/mega.hpp(+153 -0)📝
csrc/jit/handle.hpp(+57 -0)➕
csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe_split.hpp(+1184 -0)📝
deep_gemm/__init__.py(+2 -0)📝
deep_gemm/include/deep_gemm/common/utils.cuh(+104 -0)➕
deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe_split/combine_reduce.cuh(+142 -0)➕
deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe_split/common.cuh(+27 -0)➕
deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe_split/dispatch_l1_swiglu.cuh(+1013 -0)➕
deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe_split/l2_combine.cuh(+660 -0)➕
deep_gemm/include/deep_gemm/layout/mega_moe_split.cuh(+109 -0)📝
deep_gemm/include/deep_gemm/ptx/ld_st.cuh(+6 -0)📝
deep_gemm/mega/__init__.py(+89 -4)➕
tests/test_mega_moe_split.py(+257 -0)📄 Description
Add green-context split-kernel MegaMoE (dispatch_l1_swiglu / l2_combine / combine_reduce)
Summary
This PR adds a split-kernel implementation of the FP8/FP4 MegaMoE alongside the existing fused
megakernel. Instead of one monolithic kernel that does everything, the MoE forward is decomposed
into three focused kernels that are wired into a single native CUDA graph:
dispatch_l1_swiglul2_combinecombine_reduceK1 and K2 run concurrently on disjoint SM partitions carved with CUDA green contexts
(
cudaGreenCtxCreate/cudaDevSmResourceSplit), coupled only through the HBM pool + arrivalmasks. K3 runs on all SMs after the barrier. Everything replays as one CUDA-graph launch, so the
external API matches the fused path.
The fused
fp8_fp4_mega_moekernel is left completely untouched — this is a second,parallel implementation. It is bitwise-identical to the fused output and up to ~11% faster
at large batch.
Why split it?
1. Maintainability
The fused megakernel interleaves dispatch, both GEMMs, activation, combine, and reduce inside one
persistent-CTA kernel that shares a single register / shared-memory / TMEM budget and a single
launch. That is fast but monolithic: tuning one stage (a GEMM tile, an occupancy target, the SM
count for L2) is entangled with every other stage.
2. Performance
Running K1 ‖ K2 concurrently overlaps the dispatch + Linear1 compute with Linear2 + combine,
which the monolithic kernel serializes inside its internal pipeline. At large batch this overlap
plus the per-kernel occupancy headroom yields a net speedup while remaining numerically identical.
3. As a possible reference
As a possible reference implementation for architectures such as Hopper
Correctness
tests/test_mega_moe_split.pybuilds the same inputs for both paths, runs the fused megakernel asthe reference, replays the split graph, and asserts
torch.equal(y_split, y_fused).All 32 swept configurations are bitwise-identical (
max_abs = 0).Benchmark results on B200
Sweep over token count × intermediate size × ranks. Fixed: hidden=7168, 6/384 experts,
SM split K1=96 / K2=52 / K3=148, fast_math=1, clamp=10. Timing: wall-clock best-of-20 (3 warmup),
2 GiB L2 flush + buffer reset each iter; fused & split measured by the same harness. Metrics
use the existing
test_mega_moe.pyformulas (TFLOPS = 2·recv·H·I·3; HBM = FP4 weights + acts + out;NVLink = recv·H·3).
ratio = split_us / fused_us(< 1 = split faster).procs=8, intermediate=3072 (48 experts/rank)
procs=8, intermediate=2048 (48 experts/rank)
procs=4, intermediate=3072 (96 experts/rank)
procs=4, intermediate=2048 (96 experts/rank)
Note: Without
CUDA_HOMEpointing at a 13.1+ toolkit, the split graph constructor throwsFailed to load CUDA runtime API.🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.