add A4W4 and FP8 P2P transport support to MegaMoE

[Yaowu-Xiong]

Motivation

Extend MegaMoE V2 to support the DeepSeek V4 Pro A4W4 path alongside A8W4, while reducing Stage2 cross-rank traffic and latency with blockwise FP8 P2P scatter.

Technical Details

• Add end-to-end A4W4 support:
  • quantize BF16 activations to MXFP4 with per-1x32 E8M0 scales;
  • dispatch packed FP4 payloads and consume FP4 activations in GEMM1;
  • emit FP4 A2 activations/scales and run FP4×FP4 GEMM2;
  • preserve the interleaved gate/up weight contract used by MegaMoE V2.
• Add fp8_blockwise_1x32 Stage2 transport:
  • quantize weighted GEMM2 CShuffle rows to FP8 plus E8M0 scales before P2P stores;
  • dequantize in Combine and accumulate TopK contributions in FP32;
  • add cache-policy, padding-guard, persistent-CU, BM/BN, and strided scheduling profiles for small and large token buckets.
• Add deeper A-operand pipelining for selected BN256 Stage2 profiles, with explicit B two-stage and LDS-ring invariants.
• Centralize Stage1/Stage2 configuration selection, including activation-dtype overrides and the invariant that Stage2 BM must divide Stage1 SBM.

Test Plan

• Run 8-GPU DeepSeek V4 Pro accuracy tests for A4W4 and A8W4.
• Run isolated CUDAGraph Stage1, Stage2+Combine, prequant E2E, and BF16 E2E benchmarks across BS 1–16384.
• Compare BF16 P2P and fp8_blockwise_1x32 P2P transport for both A4W4 and A8W4.

Test Result

8× AMD Instinct MI355X DeepSeek V4 Pro
hidden/intermediate：7168/3072 , experts/top-k：384/6 .
Note: A4W4 and A8W4 performance data has been updated following the merge of PR #985 (2d65dea), which changed the benchmark results.

1、A4W4 vs A8W4

2、A4W4 FP8 P2P vs BF16 P2P

3、A8W4 FP8 P2P vs BF16 P2P

4、MegaMoE vs local mxfp_moe serial chain

Submission Checklist

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

[jhinpan]

A4W4 misses two committed goldens, and CI cannot see it

Both shapes _MEGA_BENCH_PARAMS gates for a4w4 miss _MEGA_PERF_BASELINE on an
idle 8×MI355X (prequant_e2e, medians): a4w4:512 0.5350 → 0.5602 (+4.7%),
a4w4:1024 0.7050 → 0.7451 (+5.7%). Accuracy is fine on both.

CI never catches this because prepare-mlir fails in Build LLVM once, so test
and the multi-GPU jobs are SKIPPED. Worth fixing on its own.

It entered with #985

state	bs=512	bs=1024
golden, committed at a62685de	0.5350	0.7050
ee9fb4fc — last A4W4 commit before the merges	0.5348	0.7272
0369a804 — first merge of main	0.5621	0.7414
9baf8cec / 972e1a94	0.5668 / 0.5671	0.7419 / 0.7416

ee9fb4fc reproduces the golden to 0.04%, so this node is comparable to the one
they were captured on. The jump is entirely at 0369a804, whose only commit
touching kernels/mega_moe/ is #985 (2d65dea1) — consistent with #985 having
been validated on A8W4 while A4W4 lived on this branch.

Ruled out by measurement

LLVM bump #945, two ways: running 0369a804 on the older wheel with buffer_ops.py
reverted still regresses (0.5627), and reverting kernels/mega_moe/ on the newer
wheel fully recovers (0.5340) · mega_moe_config.py: keeping #985's config and
reverting only the four kernel files recovers to 0.5353 / 0.7306 · Stage2:
176.62 → 177.16 µs · occupancy, VGPR/SGPR/LDS, grid: identical · inter-rank skew:
mean/max spread 0.0004 vs 0.0005 ms · grid_mult, persist_cu,
use_tile_resource, external_grouping/external_counting, and the
fz_safe_end_i32 unroll split: all within noise.

What is left

In the slower build the shader does strictly less work — SQ_BUSY_CYCLES -1.8%,
SQ_WAIT_ANY -2.6%, SQ_WAIT_INST_ANY -5.0%, SQ_INSTS_VALU -2.8%, same occupancy, no
rank skew — yet kernel min duration goes 315.28 → 339.64 µs (+7.7%).

Total work down, critical path up. That points at intra-kernel workgroup
scheduling, which is what #985 changed.

@GwilliamHu — does the payload-overlap schedule assume the A8W4 payload size? A FP4
row is 896 i32 against A8W4's 1792, so a chunking granularity derived from the A8W4
row length lands differently. Pinning this down needs per-wave (ATT) data or
knowledge of the intended schedule, which is why I am reporting rather than
patching.

Small patch, not worth a PR

On megamoe-a4w4-dispatch-perf
if useful: dp()'s dispatch-table pointer stays in VGPRs, so codegen wraps every
buffer access built from it in a scalarization waterfall — scalar loads take Stage1
from 4275 to 3351 static instructions (-21.6%). Plus num_dispatch_cu 160 → 64 for
the FP4 bucket-512 override, matching what _select_bounded_stage1 already uses for
that bucket.

Interleaved A/B, 5 rounds: bs=512 -1.3% but the distributions overlap (4/5, p≈0.19,
not separable from noise); bs=1024 -0.8% clean (5/5, disjoint). 0.8% against a 5%
regression is not worth review bandwidth.

[Yaowu-Xiong]

**jhinpan ** commented

Thanks for the detailed investigation.

The A4W4 performance numbers in this PR are indeed outdated after PR #985 (2d65dea) was merged. I'll update the A4W4/A8W4 benchmark results and corresponding baselines to reflect the latest measurements.

The regression does not appear to come from the A4W4 functionality itself, and accuracy remains unchanged. I'll also investigate the payload-overlap scheduling changes from #985 to better understand the A4W4 impact.

Agreed that the CI failure/skipped test issue should be addressed separately since it prevented this regression from being caught.

Note:
A4W4 and A8W4 performance data has been updated after PR #985 (2d65dea) was merged. The benchmark numbers below reflect the latest measurements.