Add SM90 MegaMoE by cooperative

[mpdfdfl]

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 Draft: Add green-context split-kernel MegaMoE features #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.

[Rachmanino]

hi, which machine are you working on, h20, h100 or h200?

[mpdfdfl]

hi, which machine are you working on, h20, h100 or h200?

we working on H200