[yangyuhuiling]

Adds mori.ops.ag_gemm: the all-gather that feeds a ColumnParallelLinear after
sequence parallelism, fused into the fp8 GEMM. Each rank holds its own tokens
A_r [M, K] and its weight shard B_r [N, K], and ends with
concat_s(A_s) @ B_r.T, [world*M, N]. This is the AMD counterpart of
vllm-project/vllm#54151 (Kimi-K3 AG-GEMM).

Draft: ptpc only. blockscale and mxfp8 -- what DeepSeek and Kimi actually
run -- need a push of their own; see the end.

Using it

with AgGemmOp(comm, n=N, k=K, m_max=M_MAX) as op:
    op.self_test()                                   # collective
    m = op.padded_m(x.shape[0])                      # same m on every rank
    a_fp8, a_scale = quantize_per_token(op.pad_rows(x, m))
    c = op(a_fp8, preshuffle_b(b_fp8), a_scale, b_scale)   # [world*m, N] bf16

One instance serves one (n, k) and any M up to m_max, compiling each M once.
Modes fused (default) and split. op.local_rows(m) is this rank's slot in
the window, to quantise into and skip the copy; out= takes a caller's buffer.
Operands the launch would read as raw bytes raise ValueError; the constructor
refuses a communicator without SDMA queues or spanning more than a node. The
guards follow GemmA2AOp.

N is any multiple of 16: the last 256-column tile may be partial, so Kimi-K3's
TP8 shards (6288, 3648) go in unpadded and come back contiguous.

How it works

The window holds the gathered operand (recv [world][M][K] fp8 and its per-row
scale); this rank's rows go into recv[rank]. The GEMM is gemm_ar's 8-wave
fp8 kernel reading that one tensor, with three differences:

• Block 0 is the poster. One lane per peer pushes this rank's rows over SDMA
  in row bands: put, quiet, then a system-scope store of the band's arrival
  flag into the peer. All peers at once, unlike #54151's ring -- on MI355X each
  peer is its own xGMI link, so a ring would use one of seven at a time.
• Tile order. Own rows first, never waiting; then (chunk, source, row tile,
  column tile), the source rotated by rank.
• The wait. A remote tile's thread 0 spins on its band's flag and fences
  before the prologue issues an A load.

A barrier-only kernel closes every call and advances the epoch the flags are
compared against (nothing is reset, so graph replay works).

Results

Kimi-K3 TP8 column projections, K=7168. fp8 ptpc, 8x MI355X, ROCm 7.14 /
PyTorch 2.12 / FlyDSL 0.3.4.1, graph replay, median of 2 interleaved rounds on an
idle box. All 92 runs validated at the fp8 floor (relL2 1.66e-3) against every
row of C for two operand sets. Times in us; split-rccl is
all_gather_into_tensor + the same GEMM. Taken before N was relaxed, so at the
padded 6400 / 3840:

global M	KDA split-sdma	KDA fused (best c)	vs RCCL	MLA split-sdma	MLA fused (best c)	vs RCCL
4096	258.7	237.9 (c2)	1.08x	178.8	176.8 (c1)	1.06x
8192	462.6	377.9 (c4)	1.27x	324.4	276.6 (c2)	1.20x
16384	864.4	702.9 (c4)	1.27x	618.0	486.0 (c4)	1.33x
32768	1693.1	1322.3 (c4)	1.32x	1212.5	896.9 (c4)	1.47x

The chunks are the mechanism: one chunk barely beats the split path, because
each remote tile then waits for a peer's whole slab. #54151's ratios on GB300
(BF16, ring) are 1.26-1.36x KDA and 1.07-1.54x MLA at global M >= 3072 -- the
same order, on different hardware, precision and GEMM, so only the ratios
compare.

The real shards against the padded widths, fused c4, FlyDSL 0.4.0:

global M	KDA 6288	padded 6400	MLA 3648	padded 3840
8192	374.8	373.3	271.6	283.3
32768	1300.3	1295.5	881.7	884.9

No faster, as expected: the partial tile still computes all 256 columns, the
same work as the padded one. What unpadded N buys is no padded weight copy and a
contiguous result instead of a column slice the next layer would copy. Full
tables in python/mori/ops/ag_gemm/README.md.

Worth a reviewer's attention

• The poster is in the GEMM, not on a second stream. Remote tiles spin; a
  grid whose resident blocks all spin on data a not-yet-dispatched push kernel
  would send never finishes. Block 0 is dispatched first. Its own tile pays the
  push time, and it is a local tile nothing waits on.
• The closing barrier is load-bearing. It advances the epoch, and it keeps a
  rank a call ahead from refilling a slot a slower rank still reads.
  --back-to-back 8 (no host sync, a different rank held each call) fails
  deterministically without it (relL2 0.8-1.3); it does not separate those two
  roles.
• One AgGemmOp switching M moves the receive slots. They are laid out for
  the current M so the GEMM reads one contiguous operand. That is safe for the
  same reason consecutive calls are: every rank's GEMM has read the slots before
  any rank passes the closing barrier. Covered by an alternating-M test.
• The partial column tile is masked by existing bounds, not new code. B and
  its scale go through buffer descriptors sized to the tensors (out of range
  reads 0), and the lane-transpose store drops any 8-column group at or past N.
  N % 16 keeps every store group and 4-column scale load on one side of N.
• The pinned mainloop carries fix(gemm-ar): drain the peeled tail's producers instead of tightening the loop #785. The kernel copies gemm_ar's 8-wave
  pipeline; its barrier sequence matches gemm_ar on main, tail drain
  included.
• CachedWindow cannot be captured by a runtime if. The kernels take its
  base/stride out before branching and build peer addresses from them, as its
  own docstring recommends; using the object inside a branch fails at trace time.
• bench_ag_gemm.py asks for a lazily-created cuda:nccl group alongside gloo
  for the RCCL baseline.

Test plan

cd /tmp && MORI_ENABLE_SDMA=1 MORI_SOCKET_IFNAME=lo python -m pytest -v \
  tests/python/cco/test_ag_gemm.py tests/python/cco/test_ag_gemm_op.py

This is the new CI step, Run CCO AG+GEMM numerics, back to back, placed after
the GEMM+AR steps because only these tests settle and retry around SDMA queue
reclaim. 41 passed on 8x MI355X from a non-editable install (ROCm 7.14, FlyDSL
0.4.0), run exactly as the step does, and the two CI steps after it pass right
behind it:

• test_ag_gemm.py -- layout on the CPU; every mode through the benchmark
  against an fp32 reference for two operand sets plus --back-to-back 8; a
  partial last column tile (N=1168) back to back, fused and split.
• test_ag_gemm_op.py -- shape rules on the CPU; from one 8-rank spawn:
  self_test and its teeth, a padded token count, alternating M, the zero-copy
  path and out=, split equal to fused byte for byte, back-to-back calls with a
  held rank, operand checks, close, and a partial column tile in both modes.

black and ruff clean.

Not done yet

• blockscale / mxfp8 -- their A scale is K-block major, so a row band is strided
  and needs its own push.
• Chunk counts above 4 (not measured), and a model integration.

🤖 Generated with Claude Code

[jhchouuu]

Tested d3509047 on MI355X at P=4 and P=8 (ROCm 7.14, FlyDSL 0.4.0), merged onto current main (the only conflict is python/mori/ops/__init__.py's _LAZY_SUBMODULES, which needs ag_gemm added next to the new a2a_gemm).

What passes:

• test_ag_gemm.py and test_ag_gemm_op.py (41 passed on 4 and on 8 GPUs), and bench_ag_gemm.py --back-to-back 8.
• At P=4, m_local=4096, N=6400, K=7168: fused c4 runs 720.9 us, against 1118.2 us for split-sdma and 1159.3 us for split-rccl.
• I also checked switching M while a rank is held. With no host sync, M = 512, 1024, 768, 1024, 256, 1024, and a different rank held before each call, every call lands at the fp8 floor in both fused and split. The claim in the code holds. The op test does not exercise it, though, because rel() and worst() sync the host between calls; this would be worth adding as a test.

1. self_test can fail on a healthy box when kernels are already compiled

self_test checks the receive slots after its own synchronize(), but nothing stops a faster rank from returning and starting its next op(...). That call's poster writes into the same recv slots of the slower rank. I warmed the op with one call, then ran self_test on every rank with rank 1's host 0.5 s slow after its sync, then made one call. Rank 1 raised:

ag_gemm self-test failed: the slot from rank 0 holds 31, expected 1. Check that mori was built with BUILD_CCO_SDMA=ON ...

On a cold op the first call's JIT compile hides the window, which is why the worker never sees it; a health check on a warm op hits it. Collecting the failures, then comm.barrier(), then raising fixes it. Raising before the barrier would leave the healthy ranks hanging.

2. The epoch wait breaks when the epoch wraps

The band and barrier waits use gemm_ar's _spin_until, an unsigned cur < ep, with ep = epoch + 1. I preset the window's control area to 0xFFFFFFFE and made four validated calls with fresh operands:

• call 1, where ep wraps to 0, comes back at relL2 1.23;
• the others are at the floor;
• without the preset all four are at the floor.

It takes 2^32 calls on one instance to get there, but one instance serves every layer with the same (n, k). a2a_gemm on main now has _spin_reach ((int32)(cur - target) < 0) for exactly this; using it in both spins fixes it.

Smaller points

• Full system fence per remote tile. Each remote tile runs cco_system_fence after its spin (kernels.py:297), whose release half writes back L2. An acquire-only fence, as a2a_gemm does, validates (including --back-to-back 8 at both world sizes). Two rounds each, us, max over ranks:

  	system fence	acquire only	change
  P=4, KDA N=6400, m_local=4096, c4	720.7 / 725.4	712.1 / 711.8	-1.5%
  P=4, MLA N=3648, m_local=1024, c2	221.8 / 225.2	222.9 / 222.5	~0
  P=8, KDA N=6400, m_local=4096, c4	1292.1 / 1295.5	1237.9 / 1237.2	-4.3%
  P=8, MLA N=3840, m_local=1024, c2	260.7 / 258.8	253.4 / 253.6	-2.4%

• DEFAULT_CHUNKS=4 against the MLA table. The comment says 4 wins from 1024 rows per rank up and 2 at 512, but the README's own MLA table has c1 ahead at global M 4096 and c2 ahead at 8192. I re-measured it at P=8, N=3840 (two rounds, max over ranks, us; chunks_for gives c2 and c4 here):

  global M	split-rccl	split-sdma	c1	c2	c4	default vs best
  4096	200.9	180.2	172.2	182.9	-	+6.2%
  8192	360.7	332.0	320.5	260.3	282.4	+8.5%

  At global M=4096 the default is also slower than split-sdma. A shape-aware default, or at least a corrected comment, would help.

• validate() accepts block_m 384/512. The mainloop cannot serve those: 512 does not fit in LDS, and 384 is the A > B shape behind #857's wait_barrier(2A+B) gap. Restricting block_m to the served tiles would close that.

[yangyuhuiling]

Thanks. Every point is addressed in bde45d7, on top of a merge with main (7adb171), where ag_gemm now sits next to a2a_gemm in _LAZY_SUBMODULES.

1. self_test. It now collects the failing slots, calls comm.barrier(), and only then raises. So no rank can leave and start its next call while a peer is still reading its slots, and a failing rank does not leave the healthy ones waiting in the barrier.

2. Epoch wrap. Both the band wait and the barrier wait now use a local _spin_reach, (int32)(cur - target) < 0, as in a2a_gemm. New op test test_calls_across_the_epoch_wrap_stay_correct: it presets the whole control area to the state 0xFFFFFFFE calls leave, then makes four validated calls with fresh operands, fused and split.

• All eight calls are at 1.66e-3.
• With the old unsigned compare patched back in, the wrapping call fails: relL2 1.32 fused, 1.29 split.

Smaller points:

• Fence. A remote tile's wait now takes an acquire-only fence (llvm.fence acquire), as a2a_gemm does. Your table is the measurement behind it.
• Chunk default. It is now the most bands, up to 4, that still hold MIN_BAND_TILES = 20 output tiles each. That reproduces the best count at all eight cells of the README's tables at P=8: KDA c2/c4/c4/c4; MLA c1 at 512 rows per rank, c2 at 1024, and c4 above. That includes the two MLA cells you re-measured. 20 sits between the 15-tile band that lost (MLA c2 at 512 rows) and the 25-tile ones that won. A pinned chunks= still takes the largest divisor up to it. The op worker pins 4, because at its N=1024 the default would cut a single band and most of its cases exist to exercise the band waits.
• validate() now accepts only block_m=256.
• Switching M under a held rank is now tested. case_alternating_m runs M = 512, 1024, 768, 1024, 256, 1024 back to back, with no host sync and a different rank held before each call, following your sequence.

Testing. 8x MI355X, ci-cco image, FlyDSL 0.4.0: test_ag_gemm.py and test_ag_gemm_op.py on 8 GPUs, 49 passed. Pre-commit passes.

[jhchouuu]

Re-tested bde45d7b, merged onto current main, on MI355X (ROCm 7.14, FlyDSL 0.4.0). Every point is addressed. LGTM.

• Tests. test_ag_gemm.py and test_ag_gemm_op.py pass (49).
• self_test. My reproducer now passes on every rank, and the call after it lands at the floor. The reproducer warms the op, makes rank 1's host 0.5 s late after its sync, then makes one call.
• Epoch wrap. I preset the control area to 0xFFFFFFFE. All four calls land at the floor in both fused and split; before, the wrapping call came back at 1.23.
• validate(), the acquire fence and the held-rank alternating-M test are as you describe.

One note on the chunk default, not blocking. MIN_BAND_TILES reproduces the P=8 tables, but it counts a band's tiles without the world size. At P=4 it picks more bands than the best in three of the four cells I measured. Kimi-K3 shapes, P=4, us, median of three rounds (spread about 2% at most):

N	m_local	c1	c2	c4	split-sdma	default
3648 (MLA)	512	148.1	156.1	-	148.4	c1
3648 (MLA)	1024	213.6	222.4	240.6	227.9	c2, +4.1%
6400 (KDA)	512	156.2	168.4	-	165.1	c2, +7.8%, no faster than split
6400 (KDA)	1024	298.5	233.3	253.4	305.9	c4, +8.6%

A band waits on the rows of P-1 peers. Counting its tiles over those peers fits all twelve cells, the eight in your P=8 tables and these four: (tiles // d) * n_tiles * (world - 1) >= MIN_BAND_TILES * 7. That is exactly today's rule at P=8, and at P=4 it gives c1 / c1 / c1 / c2 here. If you would rather keep the rule as it is, a note in the comment that it was fitted at P=8 would do.