MegaMoEV1 on Gfx950

[GwilliamHu]

Motivation

This PR adds MegaMoEV1, a single end-to-end intra-node EP-MoE operator (MegaMoE.forward(x_bf16, wts,
topk_ids) -> bf16) that folds the entire layer — quant → fused stage-1 (dispatch ⊕ group-GEMM1) →
fused stage-2 (GEMM2 ⊕ combine-scatter) → bf16.

Technical Details

New package kernels/mega_moe/. mega_moe.py::MegaMoE drives both stages, each with a non-fused
fallback (enable_fused_stage1/2); stage-1→stage-2 handoff is the stable Stage1Output buffer set;

Fused stage-1 —
P2P writes straight into the GEMM's expert-major layout), fixed-slot
dispatch removes the whole moe_sorting kernel; raw_a_scale folded into the K-loop removes the
scale-sort kernel; expert-major layout removes GEMM gather; A compact
scheme auto-switches from fixedslot by buffer size.

Fused stage-2 —
GEMM2 and the EP combine are fused: Rewrites the GEMM2 epilogue's
local store into a remote P2P buffer_store, inlining the combine Stage-1 per-token scatter
directly into GEMM2's CShuffle epilogue (4 extra i64 peer-buffer addrs on the launcher).

Test Plan

• tests/kernels/test_mega_moe.py — 8-rank torchrun --standalone --nproc_per_node=8 on 8×MI355X
  (gfx950): full end-to-end MegaMoE.forward (stage-1 + stage-2) accuracy + CUDAGraph device-time
  perf, plus stage-1-only mode, across 3 networks (v4_flash 4096/2048/256, r1_v3 7168/2048/256,
  v4_pro 7168/3072/384) × 15 batch sizes (1…32768) × a8w4 / a4w4, vs the unfused ops(dispatch → moe_sorting → scale_sort → GEMM1 → GEMM2->combine).
• Accuracy: relL2(mega, unfused ops) and relL2(mega, f16 reference) over multiple seeds/bs.
• CI: registered in tests/arch_compat.py (CDNA4-gated); E2E perf-regression gate via
  tests/kernels/mega_moe_perf_baseline.json (golden forward() CUDAGraph latency; passes if measured
  ≤ golden × 1.05).

Test Result

• Coverage: all 3 networks × 15 bs complete — no skip / OOM / FAIL. (v4_pro bs32768 omitted:
  fused-stage-2 a2 buffer exceeds the 32-bit num_records cap.)
• Accuracy: relL2(mega, unfused) ≈ 1e-5 end-to-end (bit-level identical, pure fp8 quant jitter)
  across seeds/bs; relL2(mega, f16) = 0.027 (fp8) / 0.143 (fp4).
• Performance (8×GFX950, CUDAGraph device-time, 8-rank mean):
  • end-to-end (stage-1 + fused stage-2) generally 1.08–1.40×, strongest at small/mid batch;

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.