[SM90] Optimize FP8 MegaMoE: fused shared expert, interleaved scheduler, L2 epilogue, 4WG heuristic and decode swap-AB by Helicoplacoidea · Pull Request #69 · sgl-project/DeepGEMM

Helicoplacoidea · 2026-08-02T13:29:13Z

Summary

Relative to the dev branch, this PR optimizes the SM90 FP8 MegaMoE kernel and syncs its
scheduler with DeepGEMM's official July release. Main changes:

Interleaved L1/L2 scheduler + producer–consumer model. Ported SM100's interleaved scheduler
to SM90 (cluster = 1); a dedicated producer warp drives a 2-stage SMEM task-ring with global
atomic L1/L2 counters, replacing wave-based coarse scheduling.
Fused shared expert. Folded into the kernel as SharedLinear1 / SharedLinear2 phases — one
launch produces routed + shared, replacing the env-gated two-stream fallback; the symmetric buffer
grows to a 10-tuple and transform_shared_weights_for_mega_moe_sm90 is exposed.
L2 epilogue bank swizzle. An 8-column-granularity XOR swizzle removes SMEM bank conflicts,
applied symmetrically on the STS write and LDS read sides (self-cancelling).
4-WG launch heuristic. The split-MN threshold now depends on intermediate_hidden and GPU SM
count: H20 pro/flash prefer 2-WG, and the 2-WG split-N BLOCK_N=256 selection is un-gated so
flash picks n=256.
Decode swap-AB → BLOCK_N=256. The swap-AB path supports 2-WG N-split (two m64n64 slabs),
with the L1-output TMA descriptor covering the full tile.
Thread/register fix + cleanup. Dispatch + non-epilogue fixed at 32 + 96; dead scheduler and
Option-A/chunk paths removed; accuracy layer 6 (fused shared-expert) added.

Accuracy

All changes are performance-only and do not change numerical results:

L2 epilogue bank swizzle: symmetric STS/LDS XOR permutation, self-cancelling.
4-WG heuristic + BLOCK_N=256 selection: only changes kernel shape/tile granularity; each
shape accumulates the same full-tile GEMM.
swap-AB → BLOCK_N=256: same GEMM with FP32 accumulation, m64n128 split into two m64n64
slabs.
Interleaved scheduler + fused shared expert: single-operator and model-level accuracy match
the dev branch, no regression.

Benchmark results

Environment: 8× NVIDIA H20 (SM90), CUDA / nvcc 13.2. Every token point runs the fused kernel and
the DeepEP baseline on the same config. Timing is robust to per-run jitter: each implementation is
launched 10 times in a fresh 8-rank process group, the 8 per-rank kernel times are averaged per
run, and the median across the 10 runs is reported. pr avg = fused-kernel time (kineto region
sm90_fp8_mega_moe_impl, fused single launch including the shared expert when SE is on); sgl avg
= the sgl-deep-gemm counterpart time; baseline avg = DeepEP dispatch + 2× grouped FP8 GEMM +
SwiGLU + combine (CUDA-event wall-clock median). All are averaged over 8 ranks.
vs baseline speedup = baseline_avg ÷ pr_avg, vs sgl speedup = sgl_avg ÷ pr_avg.

sgl (sgl avg) SE-on runs the shared expert serially, out-of-kernel. sgl's SM90 MegaMoE has no
fused shared expert, so with num_shared_experts=1 the shared expert is evaluated as a dense FP8
w8a8 MLP (L1 GEMM → SwiGLU + FP8 quant → L2 GEMM) on the same stream, serialized with the routed
mega kernel, and the two outputs are summed (matching sglang's forward_mega_moe single-stream
path). Its sgl avg is the CUDA-event median of that combined sequence;
vs sgl speedup = sgl_avg ÷ pr_avg (fused single launch vs serial shared + routed).

Shapes:

pro (default): E=384, H=7168, IH=3072, topk=6
flash (modified): E=256, H=4096, IH=2048, topk=6
token counts are per-rank, sweep [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192]

SE off · pro

pro (default: E=384 / H=7168 / IH=3072 / topk=6)

tokens	pr avg	baseline avg	vs baseline speedup	sgl avg	vs sgl speedup	timing ranks
1	335.1	536.6	1.60×	333.9	1.00×	8
2	439.8	694.4	1.58×	481.9	1.10×	8
4	609.3	1014.9	1.67×	652.9	1.07×	8
8	890.2	1410.4	1.58×	949.8	1.07×	8
16	1016.4	1670.8	1.64×	1097.7	1.08×	8
32	1110.2	1786.2	1.61×	1184.4	1.07×	8
64	1164.4	1768.9	1.52×	1246.5	1.07×	8
128	1282.8	1849.5	1.44×	1365.1	1.06×	8
256	1634.6	1924.9	1.18×	1643.2	1.01×	8
512	2493.6	3198.1	1.28×	3090.9	1.24×	8
1024	4050.4	5666.0	1.40×	4850.8	1.20×	8
2048	7193.8	8792.9	1.22×	7898.3	1.10×	8
4096	13281.4	15975.0	1.20×	13725.2	1.03×	8
8192	25463.3	30564.4	1.20×	25606.5	1.01×	8

SE off · flash

flash (modified: E=256 / H=4096 / IH=2048 / topk=6)

tokens	pr avg	baseline avg	vs baseline speedup	sgl avg	vs sgl speedup	timing ranks
1	169.8	262.4	1.55×	167.4	0.99×	8
2	203.7	334.2	1.64×	195.1	0.96×	8
4	268.0	439.8	1.64×	275.2	1.03×	8
8	342.2	493.2	1.44×	344.1	1.01×	8
16	358.9	582.2	1.62×	365.3	1.02×	8
32	373.4	618.5	1.66×	379.6	1.02×	8
64	390.8	607.9	1.56×	393.8	1.01×	8
128	475.9	630.8	1.33×	507.0	1.07×	8
256	524.8	739.0	1.41×	560.2	1.07×	8
512	935.6	1215.9	1.30×	927.6	0.99×	8
1024	1550.2	2212.4	1.43×	1703.1	1.10×	8
2048	2781.7	3862.0	1.39×	2923.5	1.05×	8
4096	5227.3	6967.6	1.33×	5237.4	1.00×	8
8192	9816.4	13424.4	1.37×	9841.2	1.00×	8

SE on (num_shared_experts=1) · pro

pro (default: E=384 / H=7168 / IH=3072 / topk=6)

tokens	pr avg	baseline avg	vs baseline speedup	sgl avg	vs sgl speedup	timing ranks
1	345.7	607.5	1.76×	385.9	1.12×	8
2	471.6	770.4	1.63×	508.1	1.08×	8
4	642.3	1100.5	1.71×	707.6	1.10×	8
8	914.4	1466.9	1.60×	996.9	1.09×	8
16	1048.1	1721.9	1.64×	1171.2	1.12×	8
32	1135.9	1803.9	1.59×	1240.3	1.09×	8
64	1256.1	1831.4	1.46×	1306.2	1.04×	8
128	1357.6	1925.0	1.42×	1448.6	1.07×	8
256	1767.7	2089.1	1.18×	1809.9	1.02×	8
512	2736.4	3399.4	1.24×	3357.8	1.23×	8
1024	4550.2	6321.8	1.39×	5467.2	1.20×	8
2048	8182.9	9942.4	1.22×	9059.2	1.11×	8
4096	15269.0	18304.4	1.20×	16102.2	1.05×	8
8192	29610.4	35205.8	1.19×	30451.8	1.03×	8

SE on (num_shared_experts=1) · flash

flash (modified: E=256 / H=4096 / IH=2048 / topk=6)

tokens	pr avg	baseline avg	vs baseline speedup	sgl avg	vs sgl speedup	timing ranks
1	187.3	298.0	1.59×	201.4	1.08×	8
2	228.1	373.9	1.64×	240.7	1.06×	8
4	277.4	488.2	1.76×	303.1	1.09×	8
8	331.4	535.8	1.62×	348.7	1.05×	8
16	358.2	669.9	1.87×	384.8	1.07×	8
32	384.2	633.9	1.65×	403.9	1.05×	8
64	419.6	672.2	1.60×	440.4	1.05×	8
128	504.1	685.5	1.36×	547.9	1.09×	8
256	592.1	816.6	1.38×	641.7	1.08×	8
512	995.7	1375.6	1.38×	1060.3	1.06×	8
1024	1750.0	2494.0	1.43×	1955.4	1.12×	8
2048	3168.4	4403.9	1.39×	3420.9	1.08×	8
4096	5990.9	7985.6	1.33×	6233.0	1.04×	8
8192	11373.2	15469.5	1.36×	11837.6	1.04×	8

DeepSeekV4Flash-Base Benchmark setup

Model: DeepSeek-V4-Flash-Base
Backend: sglang (in-process serving, --backend sglang)
Workload: random-ids, fixed shape 32,000 input × 1,500 output tokens
256 prompts, 32 max concurrency, 64 warmup requests, --seed 42

python3 -m sglang.benchmark.serving \
  --backend sglang --host localhost --port 31000 \
  --model /home/admin/DeepSeek-V4-Flash-Base/ \
  --dataset-name random-ids \
  --random-input-len 32000 --random-output-len 1500 --random-range-ratio 1.0 \
  --num-prompts 256 --max-concurrency 32 --warmup-requests 64 --seed 42

Results

Metric	sgl	pr	Speedup
Output token throughput (tok/s)	906.17	944.05	+4.2%
Total token throughput (tok/s)	20237.84	21083.80	+4.2%
Mean TTFT (ms)	27195	26303	-3.3%
Mean TPOT (ms)	17.18	16.35	-4.8%

Fusing SE improves output token throughput by ~4.2% while also reducing mean
TTFT and mean TPOT by ~3–5% on a long-context (32k) decode workload.

DeepSeekV4Flash-FP8 sgl-eval End-to-End Evaluation

End-to-end accuracy and throughput measured with sgl-eval on the same 8× H20
machine.

GSM8K: 1319 examples, single-shot, --num-threads 32.
GPQA-Diamond: 198 examples × 16 repeats, --thinking, --temperature 1.0 --top-p 1.0.

GSM8K

Metric	sgl	pr	Delta
Examples	1319 (single-shot)	1319 (single-shot)	—
Wall time	383.2s	359.1s	-6.3%
Throughput	1630 tok/s	1696 tok/s	+4.0%
Output tokens	625K	609K	-2.6%
score	97.27%	96.59%	-0.68pp
stop_rate	100.00%	100.00%	—
truncated_rate	0.00%	0.00%	—
error_rate	0.00%	0.00%	—

GPQA (Diamond)

Metric	sgl	pr	Delta
Examples	198 × 16	198 × 16	—
Wall time	10354.5s	10020.5s	-3.2%
Throughput	3357 tok/s	3454 tok/s	+2.9%
Output tokens	34.8M	34.6M	-0.6%
pass@1 (avg-of-16)	87.25% ± 1.60%	87.50% ± 1.27%	+0.25pp
pass@16	96.46%	95.96%	-0.50pp
majority@16	87.88%	90.40%	+2.52pp
no_answer	0.09%	0.00%	-0.09pp
stop_rate	100.00%	100.00%	—
truncated_rate	0.00%	0.00%	—
error_rate	0.00%	0.00%	—

Helicoplacoidea · 2026-08-20T07:34:56Z

@Fridge003 Could you please take a look when you get a chance?

fix and add cache hint

Replay of PR sgl-project#69 onto the re-based `dev` branch (merge-base reset from the old deepseek-ai#364 lineage to current upstream `dev`). The upstream rewrite made SM90 keep a `LegacyMegaMoEScheduler` while adding the SM100 producer-warp `TaskInfo` scheduler; this commit re-integrates the SM90 work on that base. What this brings (the PR's original scope, unchanged): * SM90 MegaMoE moves from the wave-based scheduler to an SM100-style producer-warp interleaved scheduler. A dedicated producer warp claims L1/L2 tasks from two global atomic counters and interleaves them 1:1 after a warmup phase, so an early pool block's L2 tiles overlap the L1 compute of later blocks instead of waiting for a whole L1 wave to drain. * The `kNumExpertsPerWave` / ring-token heuristic is dropped from the SM90 config and host, since the interleaved scheduler needs no wave bookkeeping. * Fused shared experts (`kNumSharedExperts > 0`) on SM90: two extra phases (`SharedLinear1` / `SharedLinear2`) with their own task counters, wired through config, host descriptors, the TMA loader, the L1/L2 epilogues, and the python / FF buffers. Merge notes: * `scheduler/mega_moe.cuh`: kept upstream's `TaskInfo<>` / producer `MegaMoEScheduler` / `LegacyMegaMoEScheduler` (needed by the SM100 kernels) and appended the SM90 `SM90MegaMoETaskInfo` + `MegaMoEInterleavedScheduler`. The duplicated `get_num_l1_warmup_waves` helper resolves to the single upstream definition; `is_l1_phase` / `is_shared_phase` are added. * `sm90_fp8_mega_moe.cuh` + `heuristics/sm90_mega_moe.hpp`: taken from the PR side. The only upstream cuh delta vs the old base was a rename of the wave scheduler in a path the SM90 kernel no longer uses, so no behaviour is lost. Co-Authored-By: Claude <noreply@anthropic.com>

…ffer_for_mega_moe The three-way replay kept both copies of the kwarg that the rebase added in different positions (upstream placed it mid-signature; the PR placed it last), leaving a duplicated parameter and a SyntaxError at import time. Callers pass it by keyword, so keep only the last-position copy to match the SM90 variant (get_symm_buffer_for_sm90_mega_moe). Co-Authored-By: Claude <noreply@anthropic.com>

Both the rebase base and the PR branch carried the free function `ptx::ld_volatile(const uint32_t*)`; applying one on top of the other left two identical definitions (lines 190 and 202), which nvcc rejects ("already been defined"). Remove the first copy; the uint64_t and uint32_t overloads remain. Co-Authored-By: Claude <noreply@anthropic.com>

SM90 previously carried its own port of the interleaved scheduler (`SM90MegaMoETaskInfo` + `MegaMoEInterleavedScheduler`, ~350 lines) that duplicated the SM100 `MegaMoEScheduler` mechanism. This unifies them so both archs share one scheduler, while guaranteeing the SM100 path is byte-identical. Approach: * Add `uint32_t kClusterSize = 2` and `typename WorkspaceT = layout::Workspace` template parameters (inserted after the defaulted `kNumSharedExperts`, so the existing SM100 12-positional-arg instantiation is untouched). * Replace the hard-coded `... / 2` cluster arithmetic with `... / kClusterSize` (`kNumL1Clusters`, `kNumL2Clusters`, warmup `kNumSMs / kClusterSize`, `shared_mainloop` `kShapeN / BLOCK_N / kClusterSize`). With the SM100 default `kClusterSize = 2` every formula reduces exactly to the pre-unification code. * Guard the two cluster-dependent behaviours with `if constexpr`: - `publish_task`: single-CTA writes the slot to local SMEM + arrive; the 2-CTA path keeps `st_async_cluster` (SM100, unchanged). - consumer `get_next_task`: single-CTA releases the slot inline (SM90 contract, empty barrier init `2 + kNumEpilogueWarps`); the 2-CTA path keeps deferred release via `release_task_info()` (SM100, unchanged). * Delete `SM90MegaMoETaskInfo` (byte-identical to `TaskInfo<kHasShared>`) and `MegaMoEInterleavedScheduler`; the SM90 kernel now instantiates the unified `sched::MegaMoEScheduler<..., kClusterSize=1, layout::SM90Workspace>` and reads `task_info.n_cluster_idx` (equal to a block N index when the cluster is a single CTA). `kNumRingBlocks` is assert-only on the SM90 path and is fed the pool block count. Net effect: the SM90 scheduler code shrinks by ~350 lines and future scheduler improvements land on one struct instead of two. Co-Authored-By: Claude <noreply@anthropic.com>

…helpers Align the SM90 kernel with the SM100 pattern of asking the task-info struct about its phase instead of decoding a raw `BlockPhase` enum through namespace-level helper functions. * Add `TaskInfo::is_l1()` (true for both routed and shared L1 phases), next to the existing `is_shared()`. Both are simple `block_phase` comparisons. * Delete the now-redundant namespace free functions `is_l1_phase()` / `is_shared_phase()`; their semantics are exactly `is_l1()` / the gated `is_shared()` and they had no other callers. * The three SM90 epilogue loader/math lambdas now take `const auto& task_info` instead of seven unpacked fields, re-deriving only the members they use (`num_k_blocks` from `task_info.shape_k / BLOCK_K`, `n_block_idx` from `task_info.n_cluster_idx`, etc.) and `is_shared`/`is_l1` via the new members. Call sites collapse to `process_*(task_info)`. * Remove the now-dead `get_num_k_blocks` phase lambda and the four `kNum*BlockKs` constants it was the sole user of. Semantics are unchanged: `task_info.is_shared()` expands to `kHasSharedExperts ? (block_phase > Linear2) : false`, identical to the old `kHasSharedExperts and is_shared_phase(...)`. No SM100 code is touched. Co-Authored-By: Claude <noreply@anthropic.com>

After the SM90 kernel moved onto the unified producer-warp MegaMoEScheduler (cluster=1) and SM100 already used it, the wave-based LegacyMegaMoEScheduler had no remaining instantiations anywhere. Delete the struct and its comment to shed ~170 lines of dead code. Co-Authored-By: Claude <noreply@anthropic.com>

The four top-of-file device templates `sm90_fp8_mega_moe_clamp_gate`, `sm90_fp8_mega_moe_clamp_up`, `sm90_fp8_mega_moe_silu` and `sm90_fp8_mega_moe_swiglu` were never called anywhere. The kernel computes SwiGLU inline in the L1 epilogue with byte-identical logic (`__expf` + `fast_rcp` silu, `min`/`min(max(...))` clamps), so these were dead leftovers of an earlier implementation. Remove the unused cluster; `kFastMath` / `kActivationClamp` template params remain used by the kernel. Co-Authored-By: Claude <noreply@anthropic.com>

Helicoplacoidea · 2026-09-01T06:20:15Z

pure cleanup / build fixes on top of the rebase — no numeric or scheduling-semantics changes:

Scheduler unification: SM90 now uses the same producer–consumer MegaMoEScheduler as SM100 (kClusterSize=1, WorkspaceT templated). The wave-based LegacyMegaMoEScheduler had no remaining instantiations and was removed; the get_num_k_blocks lambda and kNum*BlockKs constants were dropped in favor of task_info.shape_k / BLOCK_K.
PR Fix SM90 MegaMoE .template disambiguation on nvcc13 #74 (nvcc13 fix) included: chained member-template calls now use .template get_base_ptr() to fix nvcc ≥13 "type name is not allowed" (5 call sites). No behavior change.
Dead-code removal: dropped the never-called sm90_fp8_mega_moe_swiglu/clamp_gate/clamp_up/silu templates (SwiGLU is computed inline in the L1 epilogue with identical logic), plus the now-redundant is_l1_phase/is_shared_phase helpers.

Helicoplacoidea changed the title ~~[SM90] Optimize FP8 MegaMoE across L2 epilogue, 4WG launch heuristic and decode swap-AB~~ [SM90] Optimize FP8 MegaMoE: fused shared expert, interleaved scheduler, L2 epilogue, 4WG heuristic and decode swap-AB Aug 20, 2026

Fridge003 force-pushed the dev branch from 80b2c44 to 5be14ba Compare August 29, 2026 07:15

zcnrex and others added 4 commits August 30, 2026 14:19

Add NVFP4 mega MoE support (sgl-project#80)

b2aece8

fix and add cache hint

1693c4e

Merge pull request sgl-project#81 from zcnrex/megamoe-nvfp4-fix

ef19220

fix and add cache hint

Helicoplacoidea force-pushed the sgl branch from 514c649 to ea1e1f4 Compare August 31, 2026 08:12

Helicoplacoidea and others added 3 commits August 31, 2026 16:36

Fridge003 force-pushed the dev branch 3 times, most recently from e4d6b08 to b7fb616 Compare August 31, 2026 10:24

Helicoplacoidea and others added 3 commits August 31, 2026 18:53

Navigation Menu

Uh oh!

[SM90] Optimize FP8 MegaMoE: fused shared expert, interleaved scheduler, L2 epilogue, 4WG heuristic and decode swap-AB - #69

[SM90] Optimize FP8 MegaMoE: fused shared expert, interleaved scheduler, L2 epilogue, 4WG heuristic and decode swap-AB#69
Helicoplacoidea wants to merge 10 commits into
sgl-project:devfrom
Helicoplacoidea:sgl

Helicoplacoidea commented Aug 2, 2026 •

edited

Loading

Uh oh!

Helicoplacoidea commented Aug 20, 2026

Uh oh!

Helicoplacoidea commented Sep 1, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

Helicoplacoidea commented Aug 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Accuracy

Benchmark results

SE off · pro

SE off · flash

SE on (num_shared_experts=1) · pro

SE on (num_shared_experts=1) · flash

DeepSeekV4Flash-Base Benchmark setup

Results

DeepSeekV4Flash-FP8 sgl-eval End-to-End Evaluation

GSM8K

GPQA (Diamond)

Uh oh!

Helicoplacoidea commented Aug 20, 2026

Uh oh!

Helicoplacoidea commented Sep 1, 2026

pure cleanup / build fixes on top of the rebase — no numeric or scheduling-semantics changes:

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Helicoplacoidea commented Aug 2, 2026 •

edited

Loading