add A4W4 and FP8 P2P transport support to MegaMoE v2

[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 .

A4W4：V1 与 V2 BF16-P2P transport:

V2 A4W4：FP8-P2P 与 BF16-P2P

V2 A8W4：FP8-P2P 与 BF16-P2P

Submission Checklist

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