[RFC] Fused W8A8 MoE megakernel on H20

Background

The current blockwise FP8 MoE path runs the local compute in several CUDA kernels:

1. count and gather routed tokens
2. gate/up GroupGEMM
3. SiLU, gate multiplication and FP8 requantization
4. down GroupGEMM
5. top-k weighted reduction

There is no distributed communication involved here. The overhead comes from the kernel boundaries and from writing intermediate activations to global memory and reading them back in the next stage.

The roadmap already mentions Megakernel as a direction:

Fuse multiple consecutive operators into a single kernel to reduce inter-kernel launch overhead and intermediate memory traffic.

I have been experimenting with applying that idea to the W8A8 MoE path.

Design

The current prototype has two kernels:

• a cooperative planner that builds the expert histogram, prefix sums and row-to-token mapping
• a compute megakernel that performs the activation gather, both GEMMs, SiLU, amax/requantization and weighted output accumulation

The planner is separate because the histogram and prefix sums are grid-wide dependencies. Once the row mapping is available, the rest of the local MoE compute stays inside the megakernel and does not materialize the intermediate gate/up or down-GEMM outputs in global memory.

The implementation uses CuTe and SM90a primitives and currently targets H20.

Results so far

The measurements below were collected on H20. Both paths run back to back with the same normalized routing weights. The megakernel timing includes output zeroing, the planner and the compute kernel.

The megakernel is most useful for small-token workloads and TP-sharded MoE with a narrow local intermediate dimension. Native GroupGEMM is still better on some wide-intermediate or low-parallelism shapes, so I do not intend this as a universal replacement.

Question

Does the cooperative planner + compute megakernel split look like a reasonable implementation of the Megakernel direction for Fused MoE? If the structure makes sense, I can send the implementation and full benchmark details as a PR.

Start by reviewing the existing blockwise FP8 MoE path and the roadmap's Megakernel direction; the issue names no files or tests. Evaluate the proposed cooperative planner and compute megakernel against the listed H20 workloads; done is an agreed design followed by a PR with full benchmark details.