MegaMoE on Gfx950

[GwilliamHu]

Motivation

Reduce A8W4 expert-parallel MoE latency by overlapping communication with computation and eliminating intermediate kernel launches.

Technical Details

• Add MegaMoEV2 with per-1x32 quantization, fused dispatch/sort/GEMM1, and fused GEMM2/weighted P2P combine.
• Support fixed-slot and compact dispatch with static MI355X configurations selected from the runtime token count and max_tok_per_rank (MTPR).
• No online collective autotuning, tuning artifacts, or FLYDSL_AUTOTUNE environment variable is required.

Usage Notes

• MTPR must be a positive power of two and cover the largest per-rank token count for the operator instance.
• For decode/CUDA graphs, use the smallest power of two covering the graph batch, normally the graph batch itself:
  • MTPR <= 128 selects the low-latency fixed-slot path.
  • MTPR >= 256 selects the compact path.
• Setting MTPR to 8192 for a small decode batch is valid, but forces compact mode and loses the fixed-slot latency advantage.
• For prefill or chunked prefill, use the smallest power of two covering the maximum per-rank chunk.
• Runtime token counts do not need to be powers of two; the nearest static token bucket is selected automatically.
• All ranks in an EP group must use the same MTPR and collective call order. Multiple graph capacities should use corresponding operator instances
  consistently on every rank.

Test Plan

• Run 8-GPU MI355X v4_pro A8W4 EP8 accuracy and performance sweeps for 1–32768 tokens per rank.
• Run static-configuration unit tests, standalone GEMM checks, and Python style checks.