[jhchouuu]

What

Adds a per-token scale row to EP v2 dispatch: an opaque, fixed-width row that travels with its token to the same recv slot. mori does not quantize or interpret it — it is a forwarding channel.

Cfg.scaleBytes        the row the CALLER has, in bytes; 0 = off
EpScaleStride(cfg)    the row MORI lays down: that, padded to 128 B
EpArgs.scalesBuf      source rows
EpArgs.offOutScales   destination region

Why

The MoE wants quantized activations. Today dispatch ships bf16 and the receiver quantizes once per received copy while it gathers — 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, but only if the e8m0 scale row can travel with the token.

At hidden 7168 that is a 224 B row, and it buys: bf16 14336 B/token/destination → 7424 with an fp8 payload, 3840 with an fp4 one.

The stride is mori's, not the caller's

A caller states the row it has; mori states the row it lays down, and the two differ by padding to 128 B.

This is not tidiness. A scale transfer is a run of consecutive destination slots, so the slot stride decides whether a run starts on a 128 B boundary — and TdmWholeOrSplit128 only yields a body for the part of a run that does; the rest degrades to the scalar global→global fallback. At the natural 224 B only every 4th run starts aligned. Measured at 512 tokens/rank, hidden 7168, EP4 on gfx1250: dispatch 157.7 us at 224 B, 94.8 us at 256 B.

Deriving it here rather than making the caller pad also deletes a kernel on the caller side — the padding lands where rows are already copied one at a time (the gfx1250 staging step, the portable body's per-token copy), so it costs nothing there. The pad bytes are written as zero rather than left as whatever the last launch put in the staging pool: nobody reads them, but they cross into a peer's memory, and "nobody reads it" is not something a buffer diff or a future consumer knows.

The padding applies on every arch, including the portable body that gains nothing from the alignment: a destination layout that changed with the arch would fork every consumer of it, and the pad is ~14% of a side channel that is itself ~3% of an fp8 payload.

Anything addressing the region itself must use the stride. recv_scales() returns a strided view of the meaningful dwords, so a caller that just wants the values never sees the padding; a caller handed the base pointer asks scale_stride_bytes().

What "off costs nothing" does and does not mean

Render omits default-valued fields, so a scaleBytes=0 Cfg renders byte-identical text to before. No new specialisation, no behaviour change for anyone not using the feature.

It does not mean nobody recompiles. The JIT cache key is name + toolchain signature + nic + flags + IncludeTreeHash(sourceDeps) + code, and EpSourceDeps covers include/mori, src/ops/dispatch_combine_v2 and src/cco — every header this PR touches. So the first use after upgrading recompiles once, for everyone.

That is correct behaviour rather than a cost. This PR inserts two fields into the middle of EpArgs (offOutScales, scalesBuf), which shifts every offset after them. If the cache key really were just the rendered text, a scaleBytes=0 user would get a kernel compiled against the old EpArgs paired with new host-side offsets, and read silently wrong arguments. IncludeTreeHash is what prevents that.

Notes for review

Both device bodies carry it, differently. The portable body copies straight to the peer, as it already does for the payload. The gfx1250 body stages the row and ships it through TDM as a fourth metadata field.

One number in two places is checked, not assumed. The stride reaches the device as a macro (the staging array is at file scope, before kCfg exists) and is recomputed from the Cfg inside the body; a static_assert requires them equal. The same shape of bug — a constant restated where the compiler cannot see both — has already bitten this file once.

EpDispatch1250xSlabBytes gives the per-warp LDS slab a floor when scales are on. The slab is sized off the payload dtype and shared with the metadata tile, so an fp8 payload halves it exactly when per-token metadata grows 52 → 276 B. Without the floor the metadata batch collapses ~11x — measured as dispatch 36.7 → 93.1 us at 512 tokens/rank. The floor is an empirical result from that one shape: the mechanism is general, the choice of the bf16 width as the target is not necessarily so for a very different hidden size. Note it is computed from Cfg.hiddenDim, which an fp4 caller has already halved, so an fp4 wire gets half the floor an fp8 one does — worth re-measuring if fp4 becomes the common case.

The scale staging pool is quadratic in world_size and is a __device__ global, so it is one copy per compiled variant: ~64 MiB at EP4 / maxTok 16384 / 256 B stride, ~256 MiB at EP8 / 16384, and ~672 MiB for a three-entry (block, warp) schedule. The per-peer stride has to be the peer's full recv capacity so the destination slot id indexes it directly, which is also what lets the existing ab + cc > recvCapM guard cover the new array — _stgCap does not bound it, and the two cross over as world_size grows. Making it world_size-independent needs the staging re-indexed block-locally, which would turn that guard from "skip this transfer" into "drop these tokens" and wants hardware validation. This PR therefore documents the footprint and adds a static_assert that fails the build at 1 GiB rather than letting a wider EP discover it at launch.

Combine never carries scales. It moves post-expert tokens, already in the combine dtype; MakeEpCfg zeroes scaleBytes for the combine kernel.

plan_api only lets an explicitly optional region be missing. A blanket except KeyError -> bind 0 would turn a typo in region_names, or a backend that forgot a required region, into a silent bind to offset 0 — which aliases the first region, so the kernel would scribble over it. An optional region also becomes mandatory again when the Request field that enables it is set (scaleBytes for out_scales), so "the backend must check its arena" is a mechanism rather than a comment.

Ep1250xLdsBytes is new. EpDispatch1250xSlabBytes used to be capped by EpCombine1250xLdsBudget: the same number, but for a shared physical reason rather than because dispatch follows combine, so retuning combine would have silently resized the dispatch slab.

The source scale row must be a whole number of dwords. The kernel walks it at a dword stride, so the Cfg validator cannot see a caller row of, say, 30 B; _unsupported checks scale_dim * scale_type_size directly.

Testing

Hardware, 4x gfx1250 (test_asym_dtype.py, MORI_V2_KERNEL_BACKEND=hip): fp8 and fp4 dispatch carrying a 224 B scale row, plus widths 32 (pads 4x), 256 (already aligned, must not grow) and 260 (not a power of two) — all byte-exact through the routing reverse map, at 8 / 64 / 512 tokens. bf16 + scales still passes.

End to end through aiter's MegaMoE (EP4, hidden 7168, topk 6): moving the quant to the sender is bit-identical — a8w4 scores 0.002174 against a torch fp32 reference on both the bf16 and the fp8 wire, a4w4 scores 0.042041 on both bf16 and fp4. Per-layer wall clock at 16384 tokens/rank, medians of 3 interleaved runs with disjoint ranges: fp8 −5.9% (4654 → 4378 us), fp4 −6.4% (3512 → 3286 us); 512 tokens is parity. Per-kernel attribution puts all of it in dispatch (449 → 197 us/layer for fp4) with the two GEMMs and combine unchanged.

CI gains two things. test_op.py with SCALE_DIM=32 on the HIP backend — without it the new channel had zero CI coverage. And test_asym_dtype.py with SCALE_DIM=224, fp8 and fp4, because test_op.py gives dispatch and combine a single dtype and the HIP backend refuses an fp8 combine — so it can only ever test scales on a bf16 wire, which is not the wire the row exists for. 224 is deliberate: it is hidden/32 at 7168 and it is not 128-aligned, so it exercises the padding.

Both CI steps run the portable body: the runners have no gfx1250, so the TDM body can only be exercised by hand.

🤖 Generated with Claude Code

[kawhil-amd]

LGTM