feat(mori-v2): let the MoE dispatch wire be chosen, for the fp4 GEMM it feeds

[jhchouuu]

What

Lets the MoE dispatch wire be chosen for the mori-v2 (MegaMoE) path, so it can ship an already-quantized payload instead of bf16:

MEGA_WIRE=bf16 | fp8 | fp4      default bf16 -- unchanged behaviour

Why

Dispatch ships bf16 today and every receiver quantizes the copy it got, so a token routed to topk peers is quantized topk times on the same values. Quantizing on the sender does it once per local token and shrinks the wire; at hidden 7168 a token crosses as 14336 B on bf16, 7424 on fp8, 3840 on fp4.

The wire has to match the expert GEMM's A operand — on gfx1250 that is fp4 unless AITER_FORCE_A8W4=1. A mismatch is a row-width error, not a slow path, so wire= and force_a8w4= are logged on the same line: a failed start is then readable instead of a width assert with no context.

Measured

DSV4-Pro on 4x gfx1250, DP4 + EP4 + fake-eplb, ISL/OSL 1024/1024, concurrency 2048, 8192 requests:

	baseline (bf16 wire)	MEGA_WIRE=fp4
total token throughput	57851.6 tok/s	60226.5 tok/s	+4.1%
mean TTFT	11967.9 ms	11248.0 ms	−6.0%
mean TPOT	58.93 ms	56.79 ms	−3.6%
benchmark duration	290.0 s	278.6 s	−3.9%

Two containers off the same image on the same box, only the env var differing. At the MoE-layer level the same change is −6.4% per layer at 16384 tokens/rank, all of it in dispatch.

Output is bit-identical either way — MX quant is destination-independent, so moving it to the sender changes no bytes (test_mega_moe_gfx1250.py scores 0.042041 against an fp32 reference on both wires).

Depends on

ROCm/aiter#4984, which adds MegaMoEGfx1250(dispatch_wire=...) and the quantizing wire behind it. Without it the dispatch_wire= keyword is unknown and construction raises — including on the bf16 default — so this should land after it.

The quantizing wire additionally needs a mori that forwards a per-token scale row (ROCm/mori#593); aiter raises a message naming it if you ask for fp8/fp4 without one. The bf16 default needs neither.

Note on the launch script

MEGA_WIRE is read from the environment, and the DSV4 serve scripts under /app/scripts are not part of this repo, so turning the wire on is a launch-time change:

MEGA_WIRE=fp4 python -m atom.entrypoints.openai_server -tp 4 --model ... \
  --enable-expert-parallel --enable-dp-attention ...

Confirm it took effect from the startup log:

[MORI-V2] Created MegaMoE: ... dispatch=mori wire=fp4 force_a8w4=0

🤖 Generated with Claude Code

[github-actions]

🏷️ CI Guide

Runs automatically on every eligible PR before approval:

• ✅ Pre Checkin: Black, Ruff, catalog schema validation, non-GPU unit tests

Heavy model tests:

• ✅ Run after the PR is approved and Pre Checkin passes
• ✅ Run immediately when an approval review is submitted
• ✅ Can be requested before approval with labels

Label	Tests
ci:full	Run all heavy PR model tests: native ATOM, vLLM, and SGLang
ci:atom	Run native ATOM model accuracy tests
ci:vllm	Run ATOM vLLM OOT model accuracy tests
ci:sglang	Run ATOM SGLang model accuracy tests

Heavy jobs are skipped when the PR is not approved and no matching ci:* label is present.
Add labels via the sidebar or gh pr edit 2018 --add-label <label>