[PR #352] [CLOSED] Add SM90 FP8 MegaMoE split kernel #384

New issue

Closed

opened 2026-10-26 21:17:06 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 21:17:06 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/352
Author: @AichenF
Created: 6/3/2026
Status: ❌ Closed

Base: main ← Head: megamoe_sm90

📝 Commits (10+)

5e012d4 Optimize SM90 MegaMoE kernels
fb73273 Layer local SM90 split tuning on PR1
0a544b3 Optimize SM90 MegaMoE split path
615ba0a Split SM90 MegaMoE into L1 and L2 kernels
eece053 Clean SM90 MegaMoE split path
c6f98a3 Merge pull request #1 from AichenF/sm90-megamoe-split-optimize
03a5c9a Add SM90 MegaMoE config selector
60927a1 Merge remote-tracking branch 'aichenf/megamoe_sm90' into megamoe_sm90
1a5bd09 Add SM90 MegaMoE H20 H200 profile selector
b033430 Refactor SM90 MegaMoE split kernels

📊 Changes

15 files changed (+4404 additions, -37 deletions)

View changed files

📝 csrc/apis/mega.hpp (+142 -7)
📝 csrc/jit/compiler.hpp (+6 -0)
📝 csrc/jit/handle.hpp (+5 -2)
📝 csrc/jit_kernels/heuristics/mega_moe.hpp (+741 -0)
➕ csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp (+271 -0)
📝 deep_gemm/__init__.py (+2 -0)
📝 deep_gemm/include/deep_gemm/comm/barrier.cuh (+7 -1)
📝 deep_gemm/include/deep_gemm/common/math.cuh (+9 -1)
➕ deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh (+2270 -0)
📝 deep_gemm/include/deep_gemm/scheduler/mega_moe.cuh (+71 -19)
📝 deep_gemm/mega/__init__.py (+55 -0)
📝 deep_gemm/testing/bench.py (+36 -6)
➕ tests/bench_mega_moe_sm90.py (+259 -0)
📝 tests/test_mega_moe.py (+1 -1)
➕ tests/test_mega_moe_sm90.py (+529 -0)

📄 Description

co-author： @jychen21

Summary

Adds an SM90 (Hopper) FP8 MegaMoE implementation alongside the existing SM100 fp8_fp4_mega_moe. The kernel is split into two cooperating launches(_l1_impl + _l2_impl), each tuned independently per phase. Activation SF layout matches the DeepEP convention (per-128 K float), so SF tensors can be physically shared between DeepEP and this kernel.
Speedup over DeepEP baseline on H20: 1.15× – 2.04× (see table below).
This branch is built on top of the initial SM90 MegaMoE port in https://github.com/deepseek-ai/DeepGEMM/pull/323 — thanks to the authors of that PR for the groundwork.

What's new

API: deep_gemm.fp8_mega_moe(...) and transform_weights_for_mega_moe_sm90(...). Same SymmBuffer as SM100; slice dtype is arch-aware (float SF on SM90, packed UE8M0 int on SM100).
Kernel: deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh — one templated _core instantiated twice via MegaMoEPhasePolicy<Linear1|Linear2>, no cross-phase if-constexpr branching at use-sites.
Scheduler: mega_moe.cuh gains kClusterSize / kL{1,2}NMajorSchedule template params for_each_phase_block; SM100 fused path unchanged (cluster_size=2 default preserved).
Selector: csrc/jit_kernels/heuristics/mega_moe.hpp — shape-keyed candidate enumeration with H20 / H200 empirical hints.
Tests / bench: tests/test_mega_moe_sm90.py (smoke / heuristic / shape / edge / stress layers, calc_diff < 0.01 against BF16/FP32 reference); tests/bench_mega_moe_sm90.py.

Generic touches

deep_gemm/include/deep_gemm/common/math.cuh: added mul2() portable shim (__fmul2_rn is sm_100+ only).
deep_gemm/include/deep_gemm/scheduler/mega_moe.cuh: kClusterSize template parameter; existing 2-CTA path preserved as the default.
deep_gemm/include/deep_gemm/comm/barrier.cuh: gate NVLink-barrier timeout printf behind DG_NVLINK_BARRIER_TIMEOUT_PRINTF; use DG_TRAP_ONLY_DEVICE_ASSERT.
csrc/jit/{compiler,handle}.hpp: DG_JIT_WITH_DEVICE_DEBUG (cuda-gdb DWARF) and DG_JIT_FORCE_LEGACY_LOAD (old-driver cuModuleLoad fallback) escape hatches.
deep_gemm/testing/bench.py: configurable L2-flush size (DG_BENCH_FLUSH_L2_BYTES); optional split-kernel timing breakdown (DG_SM90_MOE_REPORT_SPLIT_KERNELS).

Benchmark — H20-3e, 8 ranks, hidden=7168, ih=2048, num_experts=256, topk=8

M	This branch (μs)	DeepEP baseline (μs)	Speedup
8	783.5	~1520	1.94×
16	942.5	~1624	1.72×
32	848.8	~1650	1.94×
64	881.3	~1647	1.87×
128	819.7	~1670	2.04×
256	1141.6	~1740	1.52×
260	1384.3	~1730	1.25×
512	1958.4	~2545	1.30×
819	2819.4	~3356	1.19×
1024	3286.0	~4221	1.29×
2048	6024.0	~7322	1.22×
4096	11347.0	~13254	1.17×
8192	22039.0	~25396	1.15×

M ≤ 128 are 10-run medians; M ≥ 256 are from the full sweep. Numbers for M ≤ 128 are 10-run medians because at that scale per-call time is only 700–950 μs, and the result is dominated by routing-draw noise, idle-SM under-utilisation, and roughly M-independent launch/barrier overhead that doesn't amortise away.

_{🔄 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/352 **Author:** [@AichenF](https://github.com/AichenF) **Created:** 6/3/2026 **Status:** ❌ Closed **Base:** `main` ← **Head:** `megamoe_sm90` --- ### 📝 Commits (10+) - [`5e012d4`](https://github.com/deepseek-ai/DeepGEMM/commit/5e012d431084f13011b4f86e4b534024d6b54dd4) Optimize SM90 MegaMoE kernels - [`fb73273`](https://github.com/deepseek-ai/DeepGEMM/commit/fb732739a958fc1aabf7cbbe72e7880c7715b453) Layer local SM90 split tuning on PR1 - [`0a544b3`](https://github.com/deepseek-ai/DeepGEMM/commit/0a544b3fdf8d4c811d70a9d10dae10cca038578d) Optimize SM90 MegaMoE split path - [`615ba0a`](https://github.com/deepseek-ai/DeepGEMM/commit/615ba0ab7dd600fb4d2b9e02113c96edfe5f7332) Split SM90 MegaMoE into L1 and L2 kernels - [`eece053`](https://github.com/deepseek-ai/DeepGEMM/commit/eece053b2770a8eb22943ce844acbaf9dad18d46) Clean SM90 MegaMoE split path - [`c6f98a3`](https://github.com/deepseek-ai/DeepGEMM/commit/c6f98a3d03e6ee1dd0386389bd0c3db4951e49cc) Merge pull request #1 from AichenF/sm90-megamoe-split-optimize - [`03a5c9a`](https://github.com/deepseek-ai/DeepGEMM/commit/03a5c9a3b83074cd7a5b0cc4a9a261b157572b77) Add SM90 MegaMoE config selector - [`60927a1`](https://github.com/deepseek-ai/DeepGEMM/commit/60927a1f5d27703de7d6c5d328e3ae3ffd6aa5ba) Merge remote-tracking branch 'aichenf/megamoe_sm90' into megamoe_sm90 - [`1a5bd09`](https://github.com/deepseek-ai/DeepGEMM/commit/1a5bd09a42f7670d7a072acb3124fc808c48f794) Add SM90 MegaMoE H20 H200 profile selector - [`b033430`](https://github.com/deepseek-ai/DeepGEMM/commit/b03343004d8fc0edc20829911e18c4b31af03726) Refactor SM90 MegaMoE split kernels ### 📊 Changes **15 files changed** (+4404 additions, -37 deletions) <details> <summary>View changed files</summary> 📝 `csrc/apis/mega.hpp` (+142 -7) 📝 `csrc/jit/compiler.hpp` (+6 -0) 📝 `csrc/jit/handle.hpp` (+5 -2) 📝 `csrc/jit_kernels/heuristics/mega_moe.hpp` (+741 -0) ➕ `csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp` (+271 -0) 📝 `deep_gemm/__init__.py` (+2 -0) 📝 `deep_gemm/include/deep_gemm/comm/barrier.cuh` (+7 -1) 📝 `deep_gemm/include/deep_gemm/common/math.cuh` (+9 -1) ➕ `deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh` (+2270 -0) 📝 `deep_gemm/include/deep_gemm/scheduler/mega_moe.cuh` (+71 -19) 📝 `deep_gemm/mega/__init__.py` (+55 -0) 📝 `deep_gemm/testing/bench.py` (+36 -6) ➕ `tests/bench_mega_moe_sm90.py` (+259 -0) 📝 `tests/test_mega_moe.py` (+1 -1) ➕ `tests/test_mega_moe_sm90.py` (+529 -0) </details> ### 📄 Description co-author： @jychen21 Summary - Adds an SM90 (Hopper) FP8 MegaMoE implementation alongside the existing SM100 fp8_fp4_mega_moe. The kernel is split into two cooperating launches(_l1_impl + _l2_impl), each tuned independently per phase. Activation SF layout matches the DeepEP convention (per-128 K float), so SF tensors can be physically shared between DeepEP and this kernel. - Speedup over DeepEP baseline on H20: 1.15× – 2.04× (see table below). - This branch is built on top of the initial SM90 MegaMoE port in https://github.com/deepseek-ai/DeepGEMM/pull/323 — thanks to the authors of that PR for the groundwork. What's new - API: deep_gemm.fp8_mega_moe(...) and transform_weights_for_mega_moe_sm90(...). Same SymmBuffer as SM100; slice dtype is arch-aware (float SF on SM90, packed UE8M0 int on SM100). - Kernel: deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh — one templated _core instantiated twice via MegaMoEPhasePolicy<Linear1|Linear2>, no cross-phase if-constexpr branching at use-sites. - Scheduler: mega_moe.cuh gains kClusterSize / kL{1,2}NMajorSchedule template params for_each_phase_block<BlockPhase>; SM100 fused path unchanged (cluster_size=2 default preserved). - Selector: csrc/jit_kernels/heuristics/mega_moe.hpp — shape-keyed candidate enumeration with H20 / H200 empirical hints. - Tests / bench: tests/test_mega_moe_sm90.py (smoke / heuristic / shape / edge / stress layers, calc_diff < 0.01 against BF16/FP32 reference); tests/bench_mega_moe_sm90.py. Generic touches - deep_gemm/include/deep_gemm/common/math.cuh: added mul2() portable shim (__fmul2_rn is sm_100+ only). - deep_gemm/include/deep_gemm/scheduler/mega_moe.cuh: kClusterSize template parameter; existing 2-CTA path preserved as the default. - deep_gemm/include/deep_gemm/comm/barrier.cuh: gate NVLink-barrier timeout printf behind DG_NVLINK_BARRIER_TIMEOUT_PRINTF; use DG_TRAP_ONLY_DEVICE_ASSERT. - csrc/jit/{compiler,handle}.hpp: DG_JIT_WITH_DEVICE_DEBUG (cuda-gdb DWARF) and DG_JIT_FORCE_LEGACY_LOAD (old-driver cuModuleLoad fallback) escape hatches. - deep_gemm/testing/bench.py: configurable L2-flush size (DG_BENCH_FLUSH_L2_BYTES); optional split-kernel timing breakdown (DG_SM90_MOE_REPORT_SPLIT_KERNELS). Benchmark — H20-3e, 8 ranks, hidden=7168, ih=2048, num_experts=256, topk=8 | M | This branch (μs) | DeepEP baseline (μs) | Speedup | |------|------------------|----------------------|---------| | 8 | 783.5 | ~1520 | **1.94×** | | 16 | 942.5 | ~1624 | 1.72× | | 32 | 848.8 | ~1650 | **1.94×** | | 64 | 881.3 | ~1647 | 1.87× | | 128 | 819.7 | ~1670 | **2.04×** | | 256 | 1141.6 | ~1740 | 1.52× | | 260 | 1384.3 | ~1730 | 1.25× | | 512 | 1958.4 | ~2545 | 1.30× | | 819 | 2819.4 | ~3356 | 1.19× | | 1024 | 3286.0 | ~4221 | 1.29× | | 2048 | 6024.0 | ~7322 | 1.22× | | 4096 | 11347.0 | ~13254 | 1.17× | | 8192 | 22039.0 | ~25396 | 1.15× | M ≤ 128 are 10-run medians; M ≥ 256 are from the full sweep. Numbers for M ≤ 128 are 10-run medians because at that scale per-call time is only 700–950 μs, and the result is dominated by routing-draw noise, idle-SM under-utilisation, and roughly M-independent launch/barrier overhead that doesn't amortise away. --- <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:06 +00:00