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). At the same time, we were pleasantly surprised to find
that with this layout the cooperative implementation also delivers better
performance at small shapes.

As a result, the former pingpong kernel and the token-count routing threshold
were removed — fp8_mega_moe now always runs the cooperative kernel.
fp8_mega_moe_cooperative is kept as an explicit named entry for A/B testing.

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	265.9	1649.1	6.20x	272.6	1.03x	332.5	1.25x
128	283.8	1533.7	5.40x	327.1	1.15x	380.4	1.34x
256	354.4	1635.3	4.61x	389.3	1.10x	434.8	1.23x
512	501.1	1900.7	3.79x	557.0	1.11x	566.5	1.13x
1024	767.9	2336.9	3.04x	899.6	1.17x	888.0	1.16x
2048	1312.4	3341.1	2.55x	1597.4	1.22x	1543.5	1.18x
4096	2377.1	5320.2	2.24x	-	-	2878.6	1.21x
8192	4695.2	9408.6	2.00x	-	-	5516.0	1.17x

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