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

A quantizing wire also runs only on mori's dispatch, which is the only one that carries the scale row, so this layer names that backend rather than leaving it to $MEGA_DISPATCH (whose default is flydsl).

Measured

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

	bf16 wire	fp4 wire
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).

Usage

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

MEGA_DISPATCH_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

MEGA_WIRE was the name during development. It now raises rather than falling back, because an env var that is silently ignored would serve bf16 while the operator believed they were on fp4. Setting both names to the same value is allowed, so a fleet can be rolled over one script at a time.

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, merged); aiter raises a message naming it if you ask for fp8/fp4 without one. The bf16 default needs neither.

🤖 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>