feat(cco): gemm_ar MXFP8 support, and the GEMM usable without a collective

[yangyuhuiling]

What this adds

mori.ops.gemm_ar gains support for 32-wide ue8m0 (MXFP8), the quantisation
DeepSeek-V4.1-Flash actually ships (weight_block_size [32, 32],
scale_fmt "ue8m0"), and two ways to use the GEMM without a collective attached.

op	for	status
GemmAllReduceOp	a RowParallelLinear at prefill M	existing, now serves mxfp8 as well as blockscale
Mxfp8GemmOp (gemm.py)	any mxfp8 linear at prefill M	new
Mxfp8GemvOp (gemv.py)	any mxfp8 linear at decode M (<= 32)	new

The second and third exist because at V4.1-Flash's shapes most of the win is
the multiply, not the overlap -- fusing is worth 11-15% at best there, where
the GEMM alone is worth ~30% against the route SGLang would otherwise take. And
a ColumnParallelLinear has no all-reduce to fuse with at all.

CDNA4's v_mfma_scale_f32_16x16x128_f8f6f4 takes the ue8m0 scales as
instruction operands, so the 32-wide form needs no dequantisation arithmetic
at all -- unlike the existing _BlockScaleK promote/rescale chain, which is
there precisely because a 128-wide block cannot be expressed that way.

Headline numbers

MI355X (gfx950). Against mxfp8_native_blockscaled_linear, identical operands,
bf16 in and bf16 out with quantisation included on both sides, cold:

	wq_b 8192x1280	wo_b 5120x2048
M=2048	-24%	-34%
M=16384	-43%	-27%

Fused GEMM + all-reduce at V4.1-Flash's wo_b, TP4, M=16384: 1192us against
1592 for split/bf16, -25.1% on fused-sdma + fp8 gather over LSA.

Two layout results came out of making the GEMM fast, both about addresses
rather than bytes, and both are documented with the hardware counters that
settle them:

• Scales must be K-block major. Same instruction count, 3.2x the cache
  accesses row-major, +50-67% on the GEMM.
• Four M tiles pack into one scale dword, via opsel_b on the scaled MFMA.
  -29.1% SQ_INSTS_VMEM with VALU flat, worth -5 to -7.5%.

Correctness

• 66 single-GPU cases in tests/python/cco/test_gemm_ar_op.py, plus the
  existing 2-rank worker cases.
• Where the reference takes tl.dot_scaled, mori agrees with it bit for
  bit -- same instruction, same operands, and ue8m0 scales are exact powers of
  two. That is asserted, not just tolerated.
• The GEMV is bit-identical to SGLang's mxfp8_gemv at every M and both shapes.
• The 112-config correctness grid (check_gemv_grid.sh) passes on both real
  shapes and every M in {1, 2, 8, 16, 32}.

Notes for review

_gemm_a8w8_8wave.py is vendored from aiter and now diverges in two places,
not one. Mfma16x16x128 grew the scaled form of the atom: opsel_b_per_tile
on the constructor, scale_a/scale_b on call. Both default to off and the
unscaled path is byte-for-byte what it was, so a re-vendor stays a merge. The
file's header says so.

_PinnedLaunch reaches into FlyDSL internals (_call_state_cache, _sig)
to skip per-dispatch bookkeeping that measured 181.6us of Python per fused call
against 5.9us of actual hipModuleLaunchKernel. It verifies the argument tuple
element-by-element before pinning and falls back to the normal path on any
surprise, so a FlyDSL change degrades performance rather than correctness. It
is the difference between a layer being host-bound and not.

Two quantisation modes exist only to be measured against: mxfp8_unpacked
and mxfp8_row are the scale layouts mxfp8 was chosen over. They are ~5 lines
of plumbing and they keep the choice reproducible; happy to drop them if you
would rather not carry them.

bench_gemm_ar.py keeps its single-call graph timer deliberately. On this
box that carries a ~13.4us replay floor, which is 1-4% of the 300-1600us fused
all-reduce it measures and identical across the modes being compared; changing
it would move every number already published from it. New benchmarks of single
small kernels use benchmark/cco/flydsl/gemm_ar/timing.py instead, which
amortises the capture and reports cold as well as hot. The reasoning is in the
README under "Measurement traps", along with the two ways this branch got it
wrong before getting it right.

Docs moved. docs/MORI-GEMM-AR-BENCHMARK.md is now how-to-run plus a
summary (744 -> 177 lines); every measurement table lives in
python/mori/ops/gemm_ar/README.md, next to the code it is about.

Not in scope

• scatter_dtype='fp8' still raises NotImplementedError. Its regions are
  sized but nothing writes the payload; that leg is the largest remaining lever
  on the fused path.
• The shared expert's gate_up (N=1152) is not expressible -- N is 4.5 tiles of
  256. Measured on the narrow tile anyway and it never wins, so supports_gemm
  is left refusing it.

[jhchouuu]

Reviewed on 8x MI355X against the PR head. Four things to fix, all on the new public surface or in the benchmark contracts — none of them in the MXFP8 kernel numerics, which held up everywhere I pushed them. Details inline.

What checked out

• The GEMV masked-step path is correct well beyond the two shapes the grid test covers: 59 cases across NSTEPS 1–16, including waves > NSTEPS where 15 of 16 waves are entirely out of range. All within the fp8 floor, no NaN. Given the history in that code, this is the result that matters most.
• _WIDE_TILE_MIN_GRID = 140 picks the faster tile on 16/16 measured points, and both tables in gemm.py reproduce within 1–2% (M=16384: 165.9/240.0 and 146.9/212.0 against your 167.2/239.1 and 148.6/212.0). The grid-based threshold is doing what the M-based one could not.
• The GEMV tuned entries beat what _HEURISTIC would pick on every bucket, 6.7% to 28.8%.
• wo_a: a config sweep finds 13.9–17.5% over the heuristic at M=8/32 — consistent with, and slightly better than, the 12–16% the docstring admits to. Two _tune() entries would recover it.
• TP4 mxfp8, M=16384: fused+fp8 against split+bf16 is −24.0%, spread ≤0.15% over repeats. Fusion alone is 6.0% (bf16) and 9.3% (fp8); the rest is the wire, which also takes relL2 from 2.4e-3 to 2.3e-2. The PR is clear about that split, and it holds.

Not covered: anything comparing against SGLang (this environment is missing mori_mxfp8_common.py), and end-to-end serving — so the unreproduced hipErrorIllegalAddress in the docs is still open.

[jhchouuu]

The permlane guard is a lower bound where the store needs equality.

n_floor admits any BLOCK_N >= 256, but _LaneTransposeStoreC._emit asserts n_tiles_b == 2, i.e. BLOCK_N == 256 exactly. So block_n=512 passes here and dies a few frames deeper:

supports_gemm(n=8192, k=1280, block_n=512) -> True
  M=   64 narrow=True  rel_l2=1.659e-03  ok
  M= 1024 narrow=True  rel_l2=1.661e-03  ok
  M= 4096 narrow=False AssertionError: the permlane mapping pairs exactly two
                       N-tiles (BLOCK_N == 256); got n_tiles_b=4

An instance that serves small batches fine and fails once the batch grows is the worst shape for this to take.

Suggest BLOCK_N == 256 on the permlane branch rather than >= n_floor. Fixing it here covers GemmAllReduceOp too, rather than adding a third copy of the rule at each entry point.

[jhchouuu]

supports_gemm returns True for block_n=512 on shapes the wide path cannot actually store (see the note on kernels_fused.py:1140). If the guard there is tightened to == 256 this predicate follows automatically; if not, it needs its own check, since it is the documented way to ask.

(No issue with the narrow tile being checked here and not in __init__ — TILE_N_GRANULE=256 makes N % block_n == 0 imply N % 128 == 0, so that gap is unreachable.)

[jhchouuu]

The first implementation becomes its own reference, so it is never validated.

ref is seeded from the first got, making rel_l2 identically 0 for it and validated unconditionally true. The rel_l2 is None branch passes for the same reason. A MORI-only run — which the docs support — reports validated=true whatever it computed.

Later impls are still compared against the first, so a bug in one variant is caught; a bug shared by all of them, or a single-impl run, is not.

bench_gemm_ar.py:210 already has reference_partial() — an fp32 reference matching the same quantisation — used with a real tolerance at :486 and :506. Reusing it here seems better than writing a second one, and would make the new benchmark at least as strong as the existing one next to it. Failing that, validated should be None when there is nothing to compare against, not True.

[jhchouuu]

Same as bench_gemm.py:412: with --baseline none there is no reference, rel stays None, and that is treated as a pass.

[jhchouuu]

--scope kernel with --impl sglang compares two different inputs.

sglang is in the default --impl list and --scope kernel does not exclude it, but the two branches do not share operands: mori_kernel_call builds its own fp8 A and scales from seed 99 (:256) while the sglang branch runs sglang_linear_call(layer, x) on the outer bf16 x.

Since auto comes first and becomes ref, the effect is that sglang gets rel_l2'd against an unrelated result and reads as failing validation — a false accusation against the baseline rather than a false pass. The row also reports includes_quant=false (:381) while that branch does quantise, so the timing is not comparable either.

Either reject the combination, or have both sides consume the same pre-quantised input with quantisation outside the timed region.

[jhchouuu]

split-lsa is in this preset's mode list and the preset appends --gather-dtype fp8 to every mode, but bench_gemm_ar.py:281 rejects exactly that pair before the communicator is up:

--mode split-lsa has no fp8 gather leg (build_lsa_ar takes no gather_dtype);
use split-sdma or fused-sdma for --gather-dtype fp8

sweep.py fused-fp8 --dry-run emits three such commands (m=4096/8192/16384), so 3 of the 15 points cannot pass even on a healthy box. Suggest dropping the mode from this preset, or skipping it explicitly as unsupported.

[jhchouuu]

LGTM

great work!