MORI GEMM + All-Reduce Benchmark

Measurements for mori.ops.gemm_ar, the fused fp8 GEMM + all-reduce. The design and the API live next to the code in python/mori/ops/gemm_ar/README.md; this file is how to reproduce the numbers and what they were.

Every number below was taken on 8x MI355X (gfx950), one node, at the shape [M, 7168] with K=2048 — DeepSeek-V4-Pro’s wo_b under TP8 with --chunked-prefill-size 16384. Kernel timings are the median over 11 graph-replayed iterations, maximum over ranks, on an otherwise idle box. Run-to-run spread at this shape is about 2%, so differences below that are not differences.

Table of Contents

Running the benchmark

Needs a mori built with BUILD_CCO_SDMA=ON. Setting MORI_ENABLE_SDMA in the environment only rebuilds the device bitcode — a host library built without the flag has no SDMA queues, every put silently does nothing, and the all-reduce quietly produces zeros.

cd /path/to/mori
BUILD_CCO_SDMA=ON pip install .

MORI_ENABLE_SDMA=1 MORI_SOCKET_IFNAME=lo \
  torchrun --standalone --nproc_per_node=8 \
  benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py \
  --mode fused-sdma --quant blockscale -m 16384 -n 7168 -k 2048

--mode selects what is measured, all reaching the same end state:

mode

what runs

gemm-only

the GEMM alone, to size the ceiling

split-sdma

gemm then a 4-kernel SDMA all-reduce

fused-sdma

GEMM with the scatter fused into its epilogue

split-lsa

gemm then the 2-kernel LSA all-reduce

fused-lsa

GEMM storing straight into peers

--quant blockscale is the model’s own quantisation (A 1x128, B 128x128, fp32 scales) and is the column to read. --gather-dtype fp8 and --gather-transport {sdma,lsa} select the wire; see the fp8 wire.

Each run prints a RESULT_JSON line with max_rank_time_us, rel_l2 and validated, so a sweep can be parsed rather than eyeballed.

Between runs, give the SDMA queues time to drain. Back-to-back 8-rank runs hit hsaKmtCreateQueueExt failures (anvil.cpp:237) if a previous run’s ranks have not exited. ~10s is enough; a leftover server holding queues is not.

Headline

--quant blockscale, median of 11:

M

split-sdma

fused-sdma

fused-sdma + fp8 gather

4096

398.1 us

351.0 us

329.5 us

8192

722.0

621.1

539.3

16384

1472.9

1148.8

979.3

Fusing is worth -22% at M=16384; the fp8 gather a further -15%.

Both survive to the server, in the same order. See the end-to-end numbers.

For scale, the same layer as the model runs it today (a separate GEMM then an NCCL all-reduce) measures 1419.5 us at M=16384, and the GEMM alone is 369.4 us.

Where the time goes

Per layer, captured in SGLang over one 20000-token prefill (not the standalone benchmark — this is the pipeline as the model drives it):

phase

bf16

fp8 / sdma

fp8 / lsa

gemm

441.9 us

438.2 us

441.3 us

drain

180.2

170.9

204.8

reduce

41.7

42.5

43.1

quantize

—

11.8

11.9

gather

437.8

256.4

7.5 (barrier only)

dequantize

—

61.0

—

pull

—

—

230.5

wo_b layer

1101.6

980.8

939.1

-11.0%

-14.7%

Two things to read out of the bf16 column:

  • gather is the bottleneck, not drain. It moves 196 MiB at 470 GB/s, which is 7 xGMI links flat out, so halving its bytes halves its time.

  • drain’s apparent 1140 GB/s is not a bandwidth. Seven links cannot do that. It is the tell that the scatter’s pushes already went out from the GEMM epilogue and the drain is only waiting for the tail — which is why the scatter leg has far less to give than its byte count suggests, and why it is still bf16.

The fp8 wire

--gather-dtype fp8 sends the all-gather leg as e4m3 with one fp32 scale per row. The reduce still accumulates in fp32 and output is still bf16; only the wire changes. The scatter leg is unchanged — it carries partial sums that are then added across every rank, so its error would compound rather than being a single rounding.

Who moves the gather

--gather-transport

how

sdma

copy engines push, a second kernel widens

lsa (default)

CUs pull over xGMI and widen on the way in

A copy engine has no ALU, so for SDMA the widening cannot be the same step: it is a second kernel that reads the landed fp8 back out of local HBM, 98 MiB a layer. A CU pull has those bytes in registers already.

gather

M=16384

bf16 / sdma

1150.7 us

fp8 / sdma

1018.9

fp8 / lsa

957.3

The ordering holds in the server too: 1070.3ms of GPU busy for bf16 against 1052.1 for fp8/sdma and 1041.2 for fp8/lsa.

The pull also removes fp8’s small-M penalty. With the SDMA gather the two conversion kernels were a fixed cost against a transfer that shrinks with M, so fp8 measured +2.3% (slower) at M=4096. With the pull there is no fixed conversion cost left and fp8 wins wherever fusing does.

The pull grid

The single most important tuning parameter, and it is not obvious from the source. These are xGMI reads, so the grid throttles outstanding remote requests rather than covering HBM latency — it wants roughly a tenth of what the local conversion kernels want.

blocks

16

24

32

48

64

80

128

256

512

us

1261

1091

1006

962

959

963

1018

1138

1184

Flat from 48 to 80, steep either side. The first implementation launched 512 — the grid the local quantize kernel uses — and lost to SDMA by 17%, which looked like “LSA is the wrong transport” rather than “the grid is wrong”.

Note this is also not LSA_BLOCK_CAP’s 24: that cap is for a kernel moving bf16 with no arithmetic, while this one moves half the bytes and dequantises them, so it needs more waves in flight to keep the links fed.

Sweep it with MORI_GEMM_AR_PULL_BLOCKS:

for b in 16 24 32 48 64 80 128 256; do
  MORI_GEMM_AR_PULL_BLOCKS=$b MORI_ENABLE_SDMA=1 MORI_SOCKET_IFNAME=lo \
    torchrun --standalone --nproc_per_node=8 \
    benchmark/cco/flydsl/gemm_ar/bench_gemm_ar.py --mode fused-sdma \
    --quant blockscale --gather-dtype fp8 --gather-transport lsa \
    -m 16384 -n 7168 -k 2048
  sleep 10
done

What fp8 costs, numerically

relL2 against an fp32 host reference goes from 2.35e-3 (bf16 wire, which is bitwise exact through the collective) to 2.49e-2.

That is a floor, not a tuning problem. e4m3 carries 3 mantissa bits, and scale granularity barely moves it — measured in torch on a [2048, 7168] standard normal payload:

scale granularity

relL2

scale bytes

per row (7168)

2.646e-2

0.06%

per 512

2.631e-2

0.78%

per 256

2.609e-2

1.56%

per 128

2.572e-2

3.12%

per 32

2.399e-2

12.5%

200x the scale bytes buys 9%. Per-row is therefore the right choice, and ~2.5e-2 is what fp8 costs.

Model-level evaluation

The kernel-level cost above is large — 6.7x the bf16 wire. Whether it matters is a different question, and it needs the model, so this section was measured in SGLang on DeepSeek-V4-Pro at TP8. The fp8 path was verified live throughout: SGLANG_DEBUG_FUSED_WO_B_AR=1 logs relL2 per layer call during the very requests being scored.

At the layer, relL2 against the unfused path, 488 layer calls:

wire

min

median

max

bf16

3.706e-3

—

4.046e-3

fp8 / sdma

1.435e-2

2.505e-2

2.654e-2

fp8 / lsa

2.153e-2

2.496e-2

2.683e-2

At the model output, it is not detectable. Scoring 10941 tokens of real source text in a single prefill (mean logprob; lower is a worse model):

run

mean logprob

ppl

bf16

-1.184169

3.2680

bf16, rerun

-1.182018

3.2609

bf16, again

-1.180106

3.2547

fp8

-1.183895

3.2671

fp8, with debug

-1.183358

3.2653

fp8 lands inside the bf16 run-to-run band. Paired per token against the same bf16 run:

pair

mean d

sd d

max abs d

CONTROL bf16 rerun

+0.002151

0.229

3.42

CONTROL bf16 again

+0.004063

0.230

4.05

TEST fp8

+0.000274

0.219

2.93

On every statistic, fp8 is closer to bf16 than bf16 is to itself.

That band is wide because this model is already strongly non-deterministic: two bf16 runs disagree on ~48% of tokens by more than 0.01 logprob, and greedy decode diverges within 10-20 tokens. The likely cause is the MoE stage-2 epilogue, which accumulates with atomic_fadd. This is also why greedy token agreement is useless as a metric here — the bf16-vs-bf16 control is as divergent as bf16-vs-fp8:

pair

~12k tok

~24k tok

CONTROL bf16 vs bf16

35.9%

15.6%

TEST bf16 vs fp8

28.1%

23.4%

Needle-in-a-haystack retrieval at 15140 tokens is 24/24 on both wires — saturated, so it bounds gross damage without resolving anything finer.

What this does and does not say. It says fp8 causes no gross degradation and no measurable shift in next-token distribution on one scoring task. It does not say quality is unaffected on long-chain reasoning, code or maths — that needs a task benchmark, which has not been run. Note also that short prompts and decode never reach this path at all (it engages only at M >= 4096), so only long-prefill workloads are affected.

Negative results

Kept because each reads as obviously right and the reason it is not cannot be seen from the source.

Folding the narrowing into the reduce (fuse_quantize=True) saves a 28 MiB re-read and a kernel launch, and costs 20 us:

us

split reduce + quantize

957.4

fused, row stashed in registers

977.1

fused, row re-read

982.0

Not register pressure — the re-reading variant keeps no stash and is no better. It is the thread map: a per-row amax cannot be taken by a block holding only part of a row, so fusing forces one-wave-per-row, where sdma_reduce walks packs with a flat grid stride and streams a block through all 8 source slices at once.

Re-measured with three repeats each, the cost holds and is if anything larger: 926.6 against 903.2 on fp8/lsa, 980.4 against 961.9 on fp8/sdma. build_sdma_phases has always defaulted it off, so the op and SGLang take the fast path; the benchmark defaulted it on until 96bd9522, and any number produced by an older driver without an explicit --no-fuse-quantize reads about 23 us slow.

Firing the gather’s puts from inside the reduce (fuse_reduce_push=True) looked like the safest of the three: the push sends this rank’s own slice, so unlike the pull it has no cross-rank dependency, and SDMA is a copy engine so it costs no CU time. It reaches parity and not a win — against 1148.7us unfused:

bands

1

4

8

16

32

publish="writethrough"

1162.1

1157.2

1160.3

1190.1

1323.5

publish="fence"

1227.3

1380.1

1630.3

2137.2

3026.7

The gap between those rows is the useful result, and it is a lesson about how to pay for a release rather than whether to.

Handing a range to a copy engine does need one: the engine reads over the fabric, not through a CU’s cache, so s_waitcnt vmcnt(0) alone is not enough — it only retires the stores as far as this XCD’s L2. But the release can be paid two ways. Releasing to system scope after the stores is L2-writeback work charged once per block per band (256 x bands of it, ~61us per band, against a reduce that is only 42us in total — unrepayable). Storing with sc0+sc1 so the bytes never stop in L1 or L2 makes the waitcnt itself the release, and that is free here: applying the same store policy to the plain unfused reduce moves it 1153.4 -> 1148.7us, i.e. nothing. This output is written once and nothing local reads it again before the gather, so holding it in L2 bought nothing.

What remains after that fix is small on both sides and nearly cancels. At bands=1 the publish carries the mechanism’s cost with none of its benefit — 1162.1 vs 1148.7, about 13us for the per-band 256-block wait_barrier, the counter atomic and the elected block’s locked puts. Four bands buy back about 5us of overlap before the sync cost takes over again.

Dropping the release entirely is not an option even though it briefly looks like one: with cached stores and no fence the kernel reaches 1157.9us, but at 32 bands it produced relL2 1.8e-2 against the 2.35e-3 floor, differing per rank. The same 32 bands are exact under either correct publish mode, which rules out an indexing bug.

The contrast with the GEMM’s fused scatter is the transferable part: there the publish is amortised against a 437us transfer hidden behind a compute-bound GEMM; here against 42us of bandwidth-saturated reduce. The mechanism pays when what is hidden is much larger than the cost of publishing it.

Re-measured after the window-geometry work below, with three alternating repeats rather than a sweep, it is a clearer loss than the table suggests: 1128.4 us on against 1112.0 off, +16.4 us, against spreads of 2.2 and 1.5 us.

Hoisting the window geometry out of lsa_ptr. cco_lsa_ptr is winBase + peer*stride + offset and loads both fields on every call, through a generic pointer – which has to be a flat_load, since the compiler cannot rule out LDS, so it counts against lgkmcnt as well as vmcnt. FlyDSL emits it as an opaque extern call, and a kernel storing through addresses derived from that base gives LLVM no way to prove the loads are not clobbered.

Reading the geometry once and doing the arithmetic in the DSL removes all of that. It was tried four ways – hoisting out the band loop, a lsa_geometry() API, global_load accessors in C++ (cco_lsa_win_base / cco_lsa_stride, which take an address_space(1) pointer so each is a single global_load), and finally cco.CachedWindow, which reads both in its constructor so a kernel changes by one line. All four measured nothing on kernels_sdma (21 call sites) and kernels_fused, in every wire configuration.

It pays in exactly one place (ptpc, three alternating repeats):

Window

CachedWindow

split-lsa

1264.69 1264.23 1264.85

1250.25 1256.85 1254.53

fused-sdma

1110.45 1109.77 1112.85

1110.53 1113.69 1112.61

-10.7 us on split-lsa, against a 0.6 us spread; nothing on fused-sdma. ar_1stage/ar_2stage build nine peer addresses in every block of a short kernel; everywhere else the addresses are built once per launch against a body that runs for a millisecond. Count address constructions per launch, not grep -c lsa_ptr.

The same holds against PR #662’s branch in blockscale, two alternating repeats: split-lsa 1463.6 -> 1455.6 us, while split-sdma (1461.2 -> 1459.9), fused-sdma bf16 (1144.9 -> 1145.8) and fp8/lsa (949.2 -> 949.8) do not move.

Two things worth carrying. A CachedWindow cannot cross an scf.if – FlyDSL captures every variable an if body reads as state and requires single MLIR values, which Window satisfies only by having exactly one field – and the way out is to compute the addresses before the branch, which is what the offsets usually allow. And pin --quant when comparing against anything: the benchmark defaults to ptpc, ~3% faster than the blockscale every number on this page is quoted in, and reading one against the other looks exactly like a machine that drifts overnight.

A CK-shaped 4-wave GEMM, chasing a 22% gap against CK’s block-scale kernel at the same shape, reached CK’s instruction mix and not its speed. Ten hypotheses were falsified by measurement; hardware counters show identical SQ_INSTS_MFMA (7,340,032) and SQ_VALU_MFMA_BUSY_CYCLES (234,881,024), VALU within 1%, MemUnitStalled at approximately zero — but SQ_WAIT_ANY 156.0M against 120.4M. A K-sweep puts the whole difference per-iteration: our fixed cost is lower (51.6 us against 66.0), while each K-block costs 20.9 us against 14.5. The gap is wait, not work. See commits cc696762, 54fef960, 9f990637, 1dead3fc.

Reproducing the end-to-end numbers

The SGLang numbers need the integration branch and a mori built with BUILD_CCO_SDMA=ON on PYTHONPATH.

export MORI_ENABLE_SDMA=1 MORI_SOCKET_IFNAME=lo
export PYTHONPATH=/path/to/mori-with-sdma
export SGLANG_OPT_FUSED_WO_B_AR=1
export SGLANG_OPT_FUSED_WO_B_AR_FP8_GATHER=1     # optional, the fp8 wire
export SGLANG_DEBUG_FUSED_WO_B_AR=1              # optional, logs per-layer relL2

sglang serve --model-path <DeepSeek-V4-Pro> --tp 8 \
  --attention-backend dsv4 --page-size 256 --chunked-prefill-size 16384 \
  --mem-fraction-static 0.88 --kv-cache-dtype fp8_e4m3 \
  --enforce-shared-experts-fusion

--mem-fraction-static has to leave room for the symmetric window, which is VMM memory outside torch’s allocator: 700 MiB on the bf16 wire and 812 MiB on fp8, which needs the extra staging region.

The capture protocol matters more than it looks. Profile a different prompt than the one used to warm up, flush_cache between them, and compare the same request on both sides — an earlier A/B without those controls reported a +8.3% regression that did not exist.

warm = "Pack my box with five dozen liquor jugs. " * 2800
main = "The quick brown fox jumps over the lazy dog. " * 2800
gen(warm); post("/flush_cache"); post("/start_profile")
gen(main); post("/stop_profile")

GPU busy time over that capture:

GPU busy over that capture:

GPU busy

wall

vs unfused

unfused (GEMM + NCCL)

1096.0 ms

1.1969 s

—

fused, bf16 wire

1070.3

1.1728

-2.3%

fused, fp8 / sdma

1052.1

1.1455

-4.0%

fused, fp8 / lsa

1041.2

1.1385

-5.0%

An earlier capture of the same four read 1101.9 / 1077.2 / 1050.4 / 1048.2, so this reproduces to about half a percent.

Check BUILD_CCO_SDMA=ON before believing any end-to-end number. With it off every put silently does nothing: the all-reduce returns mostly the local slice, the model still answers fluently, every mori kernel still appears in the profile, and the fused path measures faster than it is because it is not moving data – -6.8% instead of -2.3%, with fp8/lsa appearing worst of the three rather than best, since the pull is the one leg that does not go through SDMA. Perplexity catches it and nothing cheaper does: 862511 against 3.26 on the same text. A short prompt cannot catch it either, because fusing needs M >= 4096.

The layer-level win is larger than the end-to-end one because wo_b is about 12% of the profile.