[Yaowu-Xiong]

Motivation

Extend MegaMoE 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.
• 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.

Test Plan

Hardware: Single-node 8× AMD MI355X (gfx950).
Software: PyTorch 2.9.1 + ROCm 7.2, FlyDSL 0.3.1, Mori 1.1.2.dev32.
Model: DeepSeek-V4-Pro (Quark MXFP4, H=7168, I=3072, 384 routed experts, topk=6, 61 layers).
Serving: TP8, DP-attention, Expert Parallel, max-model-len=9472, max-num-batched-tokens=8192, max-num-seqs=512, FP8 KV cache, block 16, prefix cache off. EPLB interval/window 200, balancedness threshold 2.0.
Workload: ISL/OSL 8192/1024, range ratio 1.0, ignore EOS, seed 0, inf rate, 64 warmup requests.
EPLB: r0 = 0 redundant (naive, util 0.85); r64 = 64 redundant (biased, util 0.90).
MegaMoE: MTPR 8192, Stage2 P2P fp8_blockwise_1x32.
Prompts: c512 = 5120, c4096 = 8192.

Test Result

After fixing the hang issue, MegaMoE exhibits some performance regression in small batch size scenarios. The performance data has been updated with the measurements obtained after the fix.

r0, Concurrency 512

Metric	MegaMoE A8W4	MegaMoE A4W4	A4W4 vs A8W4
Completed requests	5120/5120	5120/5120	—
Total throughput	43,227.67 tok/s	45,583.18 tok/s	+5.45%
Output throughput	4,803.07 tok/s	5,064.80 tok/s	+5.45%
Request throughput	4.691 req/s	4.946 req/s	+5.45%
Mean TTFT	29.29 s	26.91 s	-8.13%
Mean TPOT	77.66 ms	74.38 ms	-4.22%
P99 ITL	78.99 ms	77.65 ms	-1.70%
Mean E2EL	108.73 s	102.99 s	-5.28%

r0, Concurrency 4096

Metric	MegaMoE A8W4	MegaMoE A4W4	A4W4 vs A8W4
Completed requests	8192/8192	8192/8192	—
Total throughput	59,008.72 tok/s	61,487.22 tok/s	+4.20%
Output throughput	6,556.52 tok/s	6,831.91 tok/s	+4.20%
Request throughput	6.403 req/s	6.670 req/s	+4.17%
Mean TTFT	233.81 s	228.44 s	-2.30%
Mean TPOT	395.66 ms	375.71 ms	-5.04%
Mean E2EL	638.57 s	612.79 s	-4.04%

r64, Concurrency 512

Metric	MegaMoE A8W4	MegaMoE A4W4	A4W4 vs A8W4
Completed requests	5120/5120	5120/5120	—
Total throughput	45,537.88 tok/s	46,641.24 tok/s	+2.42%
Output throughput	5,059.76 tok/s	5,182.36 tok/s	+2.42%
Request throughput	4.941 req/s	5.061 req/s	+2.42%
Mean TTFT	27.03 s	24.89 s	-7.92%
Mean TPOT	74.60 ms	74.26 ms	-0.46%
P99 ITL	78.18 ms	80.69 ms	+3.21%
Mean E2EL	103.34 s	100.85 s	-2.41%

r64, Concurrency 4096

Metric	MegaMoE A8W4	MegaMoE A4W4	A4W4 vs A8W4
Completed requests	8192/8192	8192/8192	—
Total throughput	63,773.40 tok/s	67,195.18 tok/s	+5.36%
Output throughput	7,085.93 tok/s	7,466.13 tok/s	+5.37%
Request throughput	6.919 req/s	7.291 req/s	+5.38%
Mean TTFT	211.59 s	194.29 s	-8.17%
Mean TPOT	373.15 ms	358.09 ms	-4.04%
Mean E2EL	590.91 s	560.61 s	-5.13%

Summary: MegaMoE A4W4 vs Standard Mori A8W4 (FP8 dispatch)

Scene	Mori A8W4 tok/s	MegaMoE A4W4 tok/s	Throughput	TTFT	TPOT	E2EL
r0/c512	35,673.16	45,583.18	+27.78%	-39.87%	-11.36%	-21.14%
r0/c4096	42,969.03	61,487.22	+43.10%	-36.38%	-25.83%	-30.15%
r64/c512	37,887.71	46,641.24	+23.10%	-39.35%	-8.38%	-18.64%
r64/c4096	46,054.78	67,195.18	+45.90%	-39.94%	-25.86%	-31.43%

Summary: MegaMoE A4W4 vs Mori A4W4 — same quant, DeepSeek-R1 (sglang)

Isolates the backend (both sides A4W4, so the delta is purely MegaMoE fused
megakernel vs Standard Mori dispatch — unlike the V4-Pro table above, which mixes
quant and backend). Same sglang (PR#35619 branch) and same aiter for both runs;
DeepSeek-R1-MXFP4, EPLB r0, concurrency 512, inf rate. Both: 2048/2048 completed,
0 errors. Reported on two routing workloads: random (ISL/OSL 8192/1024,
prefill-heavy) and sharegpt (real-traffic prompt/routing distribution,
decode-heavy).

Note (data correction): an earlier revision of this table reported Mori A4W4
at 36,654 tok/s (→ +33.6%). That figure came from a non-quant-aligned Mori run;
re-running Mori under the same A4W4 dispatch, same machine, and same
aiter/sglang build yields 45,723.99 tok/s, so the true same-quant backend gap on
random is +3.5% (not +33.6%). All numbers above are the corrected, quant-aligned measurements.

random (ISL/OSL 8192/1024, prefill-heavy)

Metric	Mori A4W4	MegaMoE A4W4	MegaMoE vs Mori
Total throughput	45,723.99 tok/s	47,311.56 tok/s	+3.5%
Input throughput	40,642.99 tok/s	42,054.15 tok/s	+3.5%
Output throughput	5,080.99 tok/s	5,257.41 tok/s	+3.5%
Request throughput	4.96 req/s	5.13 req/s	+3.4%
Benchmark duration	412.74 s	398.89 s	-3.4%
Median TTFT	21,927.66 ms	23,133.00 ms	+5.5%
Median TPOT	79.23 ms	74.52 ms	-5.9%
Median ITL	55.05 ms	48.55 ms	-11.8%

sharegpt (real routing, decode-heavy)

Metric	Mori A4W4	MegaMoE A4W4	MegaMoE vs Mori
Total throughput	7,524.07 tok/s	8,460.70 tok/s	+12.4%
Input throughput	4,811.10 tok/s	5,410.01 tok/s	+12.4%
Output throughput	2,712.97 tok/s	3,050.69 tok/s	+12.4%
Request throughput	14.06 req/s	15.81 req/s	+12.4%
Benchmark duration	145.67 s	129.55 s	-11.1%
Median TTFT	448.76 ms	398.53 ms	-11.2%
Median TPOT	166.41 ms	148.99 ms	-10.5%
Median ITL	131.98 ms	112.98 ms	-14.4%

At the same A4W4 quant (so the delta is purely backend), MegaMoE's fused megakernel
with P2P all-to-all overlap wins on both workloads. On the real-traffic sharegpt
routing it delivers +12.4% total throughput with ~11% lower TTFT and TPOT.
On the prefill-heavy random workload (ISL=8192, dispatch-bandwidth bound) the gap
narrows to +3.5% total throughput with ~6% lower TPOT; here Mori's TTFT is
marginally lower, since the long-prefill phase is dominated by dispatch bandwidth
rather than kernel fusion. The fusion/overlap benefit is largest in the
decode-bound regime that dominates real serving traffic.

End-to-end accuracy (GSM8K, DeepSeek-R1-MXFP4, same A4W4)

Full end-to-end serving accuracy on GSM8K (not just the operator relL2 gate),
measured on the same aligned A4W4 runs, Invalid 0.000 for all:

Workload	Mori A4W4	MegaMoE A4W4
random	0.965	0.965
sharegpt	0.950	0.975

MegaMoE A4W4 matches Mori on random (0.965 = 0.965) and is within noise on
sharegpt (0.975 vs 0.950), confirming the fused A4W4 backend is accuracy-neutral
end-to-end — consistent with the operator-level gate (A4W4 relL2 < 0.28).

Accuracy (GSM8K, DeepSeek-V4-Pro)

All backends pass the operator-level accuracy gate (A4W4 relL2 < 0.25, A8W4 relL2 < 0.10)

Test Result

MegaMoE A4W4 delivers +2.4% to +5.5% total throughput over MegaMoE A8W4 across the four DeepSeek-V4-Pro serving configurations (r0/c512 +5.5%, r0/c4096 +4.2%, r64/c512 +2.4%, r64/c4096 +5.4%), by halving the weight and activation bandwidth. Accuracy is preserved (GSM8K full 0.9530 vs 0.9522, both PASS).

Submission Checklist

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

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 4980 --add-label <label>

PR title tags:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf] and op tags like [MLA] are left untouched. Add the no-auto-title label to opt this PR out of title tagging.

[github-actions]

No activity for 15 days, so this is now labelled stale. A push or a comment clears it; keep-open exempts it. Nothing is closed automatically.