[FlyDSL] [CI] feat(mega_moe/gfx1250): quantize before dispatch, on an fp8 or fp4 wire by jhchouuu · Pull Request #4984 · ROCm/aiter

jhchouuu · 2026-08-25T08:11:28Z

What

MegaMoE's EP dispatch can carry an already-quantized payload — fp8 or fp4 — instead of bf16.

MegaMoEGfx1250(dispatch_wire=...)        bf16 | fp8 | fp4     default bf16
$MEGA_DISPATCH_WIRE                      same, when the kwarg is not given
--dispatch_wire                          the harness flag

The default is bf16, which is what every existing caller gets today.

Why

Dispatch ships bf16 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:

wire	payload	scale row	total
bf16	14336 B	—	14336
fp8	7168	256	7424
fp4	3584	256	3840

MX quant is destination-independent, so moving it changes no bytes — the numbers below are bit-identical, not "close enough".

Results

Module (test_mega_moe_gfx1250.py, EP4, hidden 7168, topk 6, medians of 3 interleaved, ranges disjoint), per-layer wall clock at 16384 tokens/rank:

quant	bf16 wire	quantized wire
a8w4	4654 us	4378 us (fp8)	−5.9%
a4w4	3512 us	3286 us (fp4)	−6.4%

512 tokens/rank is parity. Per-kernel attribution puts all of it in dispatch — 449 → 197 us/layer for fp4 — with both GEMMs and combine unchanged.

End to end, DSV4-Pro on 4x gfx1250, DP4+EP4, 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%

How it works

Sender. per_1x32_mx_quant_hip once per local token, into the wire dtype the GEMM wants. The wire must MATCH that dtype (a8w4 → fp8, a4w4 → fp4): a mismatch is a row-width error, not a slow path, so it is checked rather than inferred.

Receiver. The route-gather kernel keeps its gather and its scale preshuffle and drops only the quant — the preshuffle cannot move to the sender, because its destination is the grouped row this rank assigns, which no sender knows. A branch inside _emit_quant_block_loop loads what the quant pass would have computed; the store pass, which owns the preshuffled-scale addressing, stays shared.

Combine is untouched and still carries bf16. combine_token_nbytes is split from the old token_nbytes on purpose: deriving the combine slot stride from a now-wire-dependent token size would silently halve it while the rows staged for combine are still bf16.

Naming. Everything that means the dispatch wire says so — is_quant_dispatch_wire, dispatch_scale_nbytes, dispatch_scale_dst_nbytes, dispatch_wire_elem_count, _DispatchWire, and the arena region disp_out_scales. _DispatchWire is deliberately not generalised: three of its four fields are structurally dispatch-only (mori carries no scales on a combine, a combine quant would happen inside the gemm2 epilogue on an LDS tile rather than host-side on a whole tensor, and recv_dtype exists only to build a torch view combine has no equivalent of), so sharing the table would hand a future combine wire three dead fields.

Also fixed

a4w4 could not reach the grouped GEMM. is_grouped_a4w4 required q_dtype_w == fp4x2 while the a8w4 arm also accepted w1.dtype == uint8. mxfp4 weights arrive in both spellings (ATOM's loader keeps them uint8), so a4w4 with uint8 weights failed the eligibility test and fell through to the 2-stage mxfp4 kernels — a silent detour to a different kernel family, which on this shape has no tuned config and aborts.
a1_scale named two different tensors in grouped_gemm_gfx1250_a8w4: the parameter is the caller's per-token e8m0 rows, and the quant pass rebound the same name to the preshuffled grouped scale. Both are uint8 with a plausible shape, so a wrong read was silent.
The gfx1250 MegaMoE CI job was not running. It passes --combine scatter_fused, a value [UT] Support a4w4 in test_mega_moe #5052 removed, so it exited on argparse.

Dependency

The quantizing wire needs a mori whose EP dispatch forwards a per-token scale row — ROCm/mori#593 (merged). The import is inside the property that needs it, so a bf16 wire works with any mori; asking for fp8/fp4 without one raises a message naming the PR.

Testing

4x gfx1250, test_mega_moe_gfx1250.py against an fp32 reference:

a8w4 + bf16 wire   logits_diff 0.002174
a8w4 + fp8  wire   logits_diff 0.002174     <- identical
a4w4 + bf16 wire   logits_diff 0.042041
a4w4 + fp4  wire   logits_diff 0.042041     <- identical

The harness gains --dispatch_wire {auto,bf16,fp8,fp4} and drives AITER_FORCE_A8W4 off the quant key, so -q a4w4_mxfp4 cannot silently measure a8w4.

Note for anyone with a launch script

MEGA_WIRE was the name during development and is now MEGA_DISPATCH_WIRE, because combine will want a wire of its own. The old name raises rather than falling back: an env var that is silently ignored would send a run that asked for fp4 down the bf16 path and report nothing. Setting both names to the same value is allowed, so a fleet can be rolled over one script at a time.

🤖 Generated with Claude Code

Dispatch sends bf16 and every receiver quantizes each copy it gets, 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 halves (fp8) or quarters (fp4) what crosses the fabric; only the WMMA scale preshuffle has to stay on the receiver, because its destination is the grouped row that rank assigns. The wire must match what the expert GEMM wants for its A operand -- a8w4 -> fp8, a4w4 -> fp4 -- so it is checked, not inferred: a mismatch is a row-width error that would read into the next token's bytes rather than a slow path. MegaMoEStage2Config.dispatch_wire: bf16 | fp8 | fp4 (mori backend only, which is the one with a channel for the scale row). combine_token_nbytes stays bf16 and separate, so the combine slot stride does not follow the wire down. The e8m0 row is padded to 128 B. Dword is all mori's validator asks for, but TdmWholeOrSplit128 only yields a body for the part of a run that starts 128 B aligned, and at the natural 224 B stride only every 4th token does. mori is told hidden_dim in ELEMENTS at its own element size: fp8 and fp4 both transport as byte8, so an fp4 wire halves the count itself. Receiver side, the route gather learns a prequantized producer: it loads the payload dwords and the e8m0 byte the quant pass would have computed, and the store pass -- the only copy of the preshuffled-scale addressing -- is shared verbatim. No software pipelining needed here: this tree's quant/store split already lets the loads cluster. test_mega_moe_gfx1250.py gains --mega_wire {auto,bf16,fp8,fp4}, allows scatter_fused for a4w4_mxfp4, and drives AITER_FORCE_A8W4 off the quant key -- otherwise -q a4w4_mxfp4 silently measured a8w4. Compile-checked only (no run): 8/8 route-gather variants, 6/6 mori dispatch plans (bf16/fp8/fp4 x 8/16 warps) emit .hsaco, and the four host paths produce the payload/scale widths the GEMM expects.

Three things kept an fp4 wire from ever being measurable, none of them in the wire itself. grouped_moe_gfx1250: mxfp4 weights arrive either as fp4x2 or as the uint8 view of the same bytes (ATOM's loader keeps them uint8; MegaMoE accepts both). Only the a8w4 arm said so, so a4w4-with-uint8-weights failed the eligibility test and fell through to the 2-stage mxfp4 kernels -- a silent detour to a different kernel family, which on this shape has no tuned config and aborts. The next statement already normalized both spellings for the CSV key, so the asymmetry was an oversight, not a rule. test harness: MegaMoE rejects anything but g1u1 interleave, while resolve_spec gave a4w4 the SEPARATED default, so -q a4w4_mxfp4 --combine scatter_fused could not start. Forcing INTERLEAVE alone then produced uncorrelated output, because shuffle_group still laid the weights out for the 2-stage family (e8m0_shuffle) rather than the grouped one (n32k4 + gguu->gugu rows) -- the MX keys differ only in ACTIVATION dtype, so under MegaMoE both take the grouped prep. run_matrix.sh: never check accuracy in a timed run. The fp32 reference is a per-expert torch loop; at 16384 tokens/rank it pins all four ranks for tens of minutes and reads as a hang. Correctness now runs once per wire at 256 tokens.

github-actions · 2026-08-25T08:12:02Z

🏷️ 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 4984 --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.

The wire's e8m0 rows were being copied into a 128 B-strided buffer before dispatch, because mori laid a row down at exactly the width it was given and the alignment is what makes the transfer fast. mori now derives that stride itself, so the copy -- one extra kernel per dispatch, over every token's scales -- goes away and the quant op's output goes straight onto the wire. scale_nbytes is now what we SEND (hidden/32, packed); scale_dst_nbytes is what ARRIVES (mori's stride), and it is asked of mori rather than recomputed, so a change to the alignment cannot leave the two out of step.

yanbo's 36dac59 was cherry-picked here and then landed upstream separately, in a revised form, as ROCm#5052. Keeping both makes the two versions collide; the mega_moe.py half of that commit stays, since it is not on main.

Conflict was confined to the multi-GPU test: yanbo's a4w4 support landed upstream as ROCm#5052 in a revised form (gate_mode simplified, shuffle_group unified via an fp4x2 view, scatter_fused renamed fused, Communicator.init given per_rank_vmm). Resolved by taking main's file wholesale and re-applying only this PR's wire plumbing on top, so the renames stay consistent. AITER_FORCE_A8W4 is now driven by -q rather than setdefault, because the wire has to match the GEMM's A operand.

…ger exists ROCm#5052 narrowed --combine to base|fused; the CI line still says scatter_fused, so the job has been exiting on argparse rather than running the test.

Flipping it to mori was an out-of-scope behaviour change: a quantizing wire already requires the mori backend, and __post_init__ says so with a message naming the kwarg, so the default never had to move. The dataclass field default is unreachable anyway -- MegaMoEGfx1250.__init__ always supplies a value -- but the $MEGA_DISPATCH fallback beside it is live, and it decided the backend for every caller that sets neither.

…orked --dispatch_commu_dtype has had no effect since the file was created in ROCm#4785: resolve_spec wrote transport/prequant/is_fp8/fp8_dtype into the spec and nothing ever read them, quant_tokens_fp8() had no callers, and the base combine hardcodes transport_dtype = bf16. So this is not two overlapping knobs -- it is one live knob and one corpse. --mega_wire takes over the name (--dispatch_wire, matching the kwarg it feeds) and the dead half goes with the flag it belonged to. _FP8_DTYPE stays: the per_Token / per_128x128 weight quant still uses it.

…e dispatch wire Half the names this PR added already carried the direction (dispatch_token_nbytes, combine_token_nbytes); the other half did not, and a combine wire is coming. is_quant_wire is the one that mattered: it is a pure dispatch predicate with nine call sites, so 'if config.is_quant_wire' written inside _combine() would take the wrong branch AND RUN -- a bf16 reduce over fp8 bytes, wrong numbers, no error. _DispatchWire is deliberately not generalised: three of its four fields are structurally dispatch-only, so sharing the table would hand combine three dead ones.

… name is fatal Combine will want a wire of its own, so the unqualified name had to go while it is still unreleased. The stale name raises rather than falling back: an env var that is silently ignored sends a run that asked for fp4 down the bf16 path and reports nothing, which is exactly the failure a wire benchmark cannot survive. Setting both names to the same value is allowed, so a fleet can be rolled over one script at a time.

The parameter is the caller's per-token e8m0 rows; forty lines later the quant pass rebinds the same name to the preshuffled grouped scale. Both are uint8 with a plausible shape, so a wrong read is silent. The incoming one now has its own name and the rebinding introduces the grouped meaning exactly once.

Three renamed lines went past 88 columns, which the pre-checks job enforces. The env guard was also true of one caller only: the library evaluates it just when dispatch_wire is None, and the test always passes the kwarg, so the harness -- the launch path most likely to carry a stale MEGA_WIRE -- would have measured bf16 in silence. The helper is now public and the test's argparse default calls it, imported lazily so the mega package is not pulled in before FLYDSL_GPU_ARCH is set. _WIRE_FOR_QUANT picks up the dispatch prefix too; it was the same ambiguity the commit before this one exists to remove.

It carries geometry and the dispatch knobs and not one stage2 parameter. Stage2 in this package is the gemm2 epilogue fused into combine, which Stage2ScatterContext already names correctly and which takes nothing from here -- so the config was the only place the word was wrong. Name only: the fields, the flat kwargs and the public surface are unchanged, so no caller moves.

yadaish · 2026-09-01T05:41:23Z

+    for fp4 down the bf16 path and reports nothing, which is the one failure
+    mode a wire benchmark cannot survive.
+    """
+    stale, current = os.environ.get("MEGA_WIRE"), os.environ.get("MEGA_DISPATCH_WIRE")


what about using one more readable name? These two looks quite similar.

MEGA_WIRE is the old name, not a second knob, the live name is MEGA_DISPATCH_WIRE, it appears only in this guard. A launch script that still exports it fails loudly instead of silently serving bf16, and the reference goes away once those scripts are updated.

And these will also be refactor after dispatch + gemm1 were implemented.
cc @yanboshao

yadaish · 2026-09-01T05:42:49Z

LGTM

jhchouuu added 2 commits August 25, 2026 08:06

jhchouuu requested a review from a team August 25, 2026 08:11

github-actions Bot changed the title ~~feat(mega_moe/gfx1250): quantize before dispatch, on an fp8 or fp4 wire~~ [FlyDSL] feat(mega_moe/gfx1250): quantize before dispatch, on an fp8 or fp4 wire Aug 25, 2026

jhchouuu force-pushed the jhchouuu/mega-moe-quantized-wire branch from 50d147c to d265dc6 Compare August 25, 2026 08:34

jhchouuu changed the title ~~[FlyDSL] feat(mega_moe/gfx1250): quantize before dispatch, on an fp8 or fp4 wire~~ feat(mega_moe/gfx1250): quantize before dispatch, on an fp8 or fp4 wire Aug 25, 2026

jhchouuu mentioned this pull request Aug 25, 2026

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

Open

jhchouuu requested review from XingerZhu and yanboshao August 25, 2026 09:18

support a4w4 test

36dac59

github-actions Bot changed the title ~~feat(mega_moe/gfx1250): quantize before dispatch, on an fp8 or fp4 wire~~ [FlyDSL] feat(mega_moe/gfx1250): quantize before dispatch, on an fp8 or fp4 wire Aug 26, 2026

jhchouuu changed the title ~~[FlyDSL] feat(mega_moe/gfx1250): quantize before dispatch, on an fp8 or fp4 wire~~ feat(mega_moe/gfx1250): quantize before dispatch, on an fp8 or fp4 wire Aug 28, 2026

jhchouuu added 2 commits August 31, 2026 02:09

revert(test): drop the duplicated a4w4 test rewrite

e7560fe

yanbo's 36dac59 was cherry-picked here and then landed upstream separately, in a revised form, as ROCm#5052. Keeping both makes the two versions collide; the mega_moe.py half of that commit stays, since it is not on main.

github-actions Bot changed the title ~~feat(mega_moe/gfx1250): quantize before dispatch, on an fp8 or fp4 wire~~ [FlyDSL] feat(mega_moe/gfx1250): quantize before dispatch, on an fp8 or fp4 wire Aug 31, 2026

junhaha666 previously approved these changes Aug 31, 2026

View reviewed changes

yanboshao reviewed Aug 31, 2026

View reviewed changes

Comment thread op_tests/multigpu_tests/test_mega_moe_gfx1250.py Outdated

yanboshao reviewed Aug 31, 2026

View reviewed changes

Comment thread aiter/ops/flydsl/kernels/mega_moe_gfx1250/mega_moe.py Outdated

jhchouuu added 7 commits August 31, 2026 14:45

fix(ci): the gfx1250 MegaMoE job passes a --combine value that no lon…

edd47d6

…ger exists ROCm#5052 narrowed --combine to base|fused; the CI line still says scatter_fused, so the job has been exiting on argparse rather than running the test.

jhchouuu dismissed junhaha666’s stale review via ebc9506 August 31, 2026 15:09

github-actions Bot changed the title ~~[FlyDSL] feat(mega_moe/gfx1250): quantize before dispatch, on an fp8 or fp4 wire~~ [FlyDSL] [CI] feat(mega_moe/gfx1250): quantize before dispatch, on an fp8 or fp4 wire Aug 31, 2026

github-actions Bot added the CI label Aug 31, 2026

github-actions Bot added the FlyDSL label Aug 31, 2026

jhchouuu requested a review from yadaish September 1, 2026 03:19

yadaish reviewed Sep 1, 2026

View reviewed changes

Navigation Menu

Uh oh!

[FlyDSL] [CI] feat(mega_moe/gfx1250): quantize before dispatch, on an fp8 or fp4 wire - #4984

[FlyDSL] [CI] feat(mega_moe/gfx1250): quantize before dispatch, on an fp8 or fp4 wire#4984
jhchouuu wants to merge 14 commits into
ROCm:mainfrom
jhchouuu:jhchouuu/mega-moe-quantized-wire

jhchouuu commented Aug 25, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Aug 25, 2026

Uh oh!

Uh oh!

Uh oh!

yadaish Sep 1, 2026

Uh oh!

jhchouuu Sep 1, 2026

Uh oh!

yadaish commented Sep 1, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Conversation

jhchouuu commented Aug 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What

Why

Results

How it works

Also fixed

Dependency

Testing

Note for anyone with a launch script

Uh oh!

github-actions Bot commented Aug 25, 2026

🏷️ CI Guide

Uh oh!

Uh oh!

Uh oh!

yadaish Sep 1, 2026

Choose a reason for hiding this comment

Uh oh!

jhchouuu Sep 1, 2026

Choose a reason for hiding this comment

Uh oh!

yadaish commented Sep 1, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

jhchouuu commented Aug 25, 2026 •

edited

Loading