Repository navigation
build(cco): compile the SDMA path by default - #710
Merged
Merged
Conversation
`BUILD_CCO_SDMA` was OFF unless `BUILD_BENCHMARK=ON` turned it on. What OFF produces is a build whose SDMA puts compile away: every symbol is still there, every kernel still launches, every put still returns, and no bytes move. The all-reduce then returns mostly the local slice. `GemmAllReduceOp.self_test` exists because of one such run -- it measured *faster* than the real thing (-6.8% against -2.3%) precisely because it was not moving data, and the only signal that caught it was perplexity, 862511 against 3.26. Measured, same tree, clean build each side: OFF wall 27s libmori_cco.so 1,701,544 B ON wall 26s libmori_cco.so 1,707,632 B (+6,088 B, +0.36%) It adds no dependency: `anvil.cpp` is in MORI_APP_SOURCES unconditionally and `find_package(hsakmt REQUIRED)` is too, so the host transport and its link deps were already there either way. The flag reaches only cco's own layer -- the queue setup in `cco_init.cpp`, the `cco_sdma_*` wrapper symbols (58 lines of 316), the device API in `cco.hpp`, and one extra JIT bitcode variant that is cached per config. And it creates nothing on its own. `ccoSdmaSetupCommQueues` returns before any allocation unless `MORI_ENABLE_SDMA` is set, which is an independent runtime gate that stays off by default -- verified by building ON and creating a two-rank communicator with and without the variable: both succeed, and the one without allocates no queues. So for anyone not opting in, this changes nothing; for anyone who does, it is the difference between working and silently inert. CI already built with `BUILD_CCO_SDMA=ON` in both `ci_cco.yml` and `ci.yml`, so ON is the configuration that was being tested and OFF was the less-covered one. Also add the diagnostic that was missing on the other side of the gate: when a caller explicitly asks for queues (`sdmaQueueCount > 0`) and none will be created, say which of the two reasons applies. It stays silent for the default `0`, so a comm that never wanted SDMA logs nothing. Not fatal -- a comm without SDMA is legal and every other path still works. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Seven fixes from reviewing the first commit. Two of them mattered. **The diagnostic was silent in the one build that fails silently.** It sat inside `#if BUILD_CCO_SDMA`, so a BUILD_CCO_SDMA=OFF build -- the escape hatch this PR keeps, and what every older image ships -- still said nothing while the identical configuration on an ON build was diagnosed. That is backwards: OFF is the build whose puts move no bytes. Moved out of the guard, with its own reason string. **`BUILD_CCO_SDMA=1` built a host library with SDMA and a device wrapper without it.** setup.py passed the raw value to CMake, which takes 1/TRUE/yes as true, but baked the flag with `== "ON"`, which does not. So `BUILD_CCO_SDMA=1 pip install .` produced host queues and device puts that compile away -- the same silent-zeros failure this default exists to prevent, entered from the other side. Normalised once, before both uses. Verified: `BUILD_CCO_SDMA=1` now bakes True. The rest: - Widen the gate. `sdmaQueueCount` defaults to 0 and almost every caller leaves it there, so gating on `> 0` kept the common configuration silent. MORI_ENABLE_SDMA is also a way of asking. - Guard `getSdmaQueue`, which returns null for a pair it never connected or a channel past what that pair got, and was dereferenced unchecked. Pre-existing, but this PR is what compiles that loop into every default build. - Keep one CI job on BUILD_CCO_SDMA=OFF. It used to be covered for free by being the default; without it the `#else`, the cco.hpp fallback and the cpp-test skip rule break only for whoever sets the escape hatch. - sweep.py read a missing `_build_flags` as OFF and aborted. With ON the default that guess is wrong; None means unknown, as op.py already had it. - "MORI_ENABLE_SDMA is not enabled", since IsEnvVarEnabled is also false for =0. - Trim the rationale comment and fix the docstring in test_gemm_ar_op.py that still called OFF the default. Two findings checked and not taken. `anySdmaCapable` scans worldSize while the fill loop covers only the lsa range, but `canSDMA == sameHost` and the fill includes self, so the all-null table it implies cannot happen. And cco.hpp's `#define BUILD_CCO_SDMA 0` fallback is left alone: flipping it would make an out-of-tree TU emit SDMA code against a library that may be OFF, which is the same mismatch pointing the other way. Verified on MI355X, both builds, three configurations each: silent when nobody asked, diagnosed when asked and unavailable, working when enabled. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
yangyuhuiling
pushed a commit
that referenced
this pull request
Oct 19, 2026
Rebased off the gemm_a2a branch (#695) onto main: nothing in gemm_ag imports gemm_a2a, and the 14 commits applied without conflict. * tests: _gemm_ag_pull_worker runs bench_gemm_ag.py through runpy.run_path, which does not put the script's directory on sys.path, so the `_bench_utils` split (25d419b) left test_pull_initializes_recycled_window failing on every run with ModuleNotFoundError. * mori.ops: register gemm_ag with the lazy FlyDSL submodules. * ci_cco.yml: run the four gemm_ag test files, after the GEMM+AR steps. Their 8-rank SDMA process groups leave queues that take seconds to come back; gemm_ag's tests settle and retry around that, gemm_ar's do not, and run right after this step six of them failed in hsaKmtCreateQueueExt (anvil.cpp:210) while passing on their own. As written the step passes here: 213 tests in 22m50s. * .gitignore: drop the local experiments entry; that archive lives outside the repository. * docs: BUILD_CCO_SDMA is a build flag and on by default since #710, so stop exporting it at run time; example paths no longer point at a private checkout. * README tables re-measured on main 8ae8156 (with #603 and #711), 8x MI355X, ROCm 7.14 / PyTorch 2.12 / FlyDSL 0.3.4.1: 54 configurations, all validated, relL2 <= 9.95e-7. Conclusions hold: 128x128 + split SDMA leads the Pro shape (355.9us, 12.6% under Torch + RCCL), fusing pays only for the combined Flash N=1024 at M=16384 (4.3%). Host and stack changed with MORI, so the deltas against September are not attributed. Co-Authored-By: Claude <noreply@anthropic.com>
yangyuhuiling
added a commit
that referenced
this pull request
Nov 4, 2026
* gemm_ag: fp8 GEMM fused with an all-gather
A third fusion operator on gemm_ar's 8-wave fp8 GEMM, alongside gemm_ar's
all-reduce and gemm_a2a's all-to-all. Every rank computes its own [M, N] and
every rank ends up with all `world` of them concatenated along rows -- the
DeepSeek V4-Pro prefill-CP wkv_gate shape, where each context-parallel rank
scores its own tokens and then everyone needs every token's score.
Eight modes: gemm-only, gemm-to-window, split-rccl, split-lsa-{push,pull},
fused-lsa, split-sdma, fused-sdma. ptpc/blockscale/mxfp8 throughout.
All-gather shards nothing, so three things gemm_a2a needs are absent: the
staging region (the payload is already contiguous and already the same for
every destination), the per-(destination, chunk) counter (one chunk arms every
push at once), and the destination rotation (a broadcast tile belongs to
everyone). The layout's N rule reduces to `n % block_n`, which is what lets
this op run the ratio-128 wkv_gate shape at N=1024 -- a shape a2a_config
rejects at world_size=8.
Because the collective is a broadcast it can be driven from either end, so
both directions are implemented and compared. That is also the first thing the
measurements say: pull needs the producer to publish C with system scope,
since it reads a peer's HBM while the peer's GEMM left the line dirty in its
own L2. Without it split-lsa-pull validated at 1.66e-3 under ptpc and 2.05e-1
under blockscale -- a correctness bug visible in exactly one quantisation,
which is why the mode matrix is parametrised over all three.
gemm_ar's _SwapABStoreC gains a `peer_rsrcs` list alongside `peer_rsrc`. A
broadcast stores one tile to every peer, and doing that with `world` separate
store objects would repeat the scale loads, the convert and the permlane
shuffle with it. A single resource yields a one-element tuple and identical IR,
so gemm_ar and gemm_a2a are unchanged.
Measured on 8x MI355X: the split paths beat RCCL by 8-9%, and the fused paths
lose to their own split baselines. The fusion itself works -- at M=512
fused-sdma cuts exposed comm 50.9us -> 21.8us -- but the epilogue's per-block
publication roughly doubles the GEMM. See docs/MORI-GEMM-AG-BENCHMARK.md.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* docs(gemm_ag): the measured tables, and where the fused time goes
Two wkv_gate shapes x two sequence lengths x every mode, plus the full
quantisation matrix and the fused-sdma chunk ladder, on 8x MI355X.
Every mori transport beats RCCL by 10-25%, and push beats pull everywhere --
the opposite of gemm_ar's gather leg, which found pull worth 5.6-10.5% over a
*copy-engine* push. The GEMM phase is equal to within 0.5us across the pair, so
pull's mandatory sc0|sc1 C store is not the cause; a push is a fire-and-forget
store where a pull is a load that must round-trip. That is worth knowing beyond
the ranking, because pull is where a low-precision wire would want its
dequantize.
The fused paths lose, and --no-put and --fence none between them say why. The
overlap itself works: fused-sdma cuts exposed comm by 25-59%, a larger fraction
than gemm_ar's all-reduce absorbs. The release fence is free at these grid
sizes (32-128 blocks, against the 9216 where gemm_a2a calls it the most
expensive line in its epilogue) and the bookkeeping is 2-6us. What costs is
*issuing* the puts: 39-55us, nearly constant across four cells whose payloads
differ by 8x, which rules out bandwidth contention and lands on gemm_ar's own
constant of ~2us per SDMA packet regardless of size.
At --chunks 4 and world=8 that is 28 puts posted serially by one thread under
one lock, where gemm_a2a posts the same 28 from 28 different blocks -- its
counter elects per (destination, chunk), so each elected block owns one
destination. A broadcast chunk arms every destination at once, which is what
made the counter cheaper and the posting expensive. Fixing it is a change to
the epilogue, not to the design, and the doc says what it is.
Separately, the bound: all-gather's G0/C0 goes as K/(world-1) where both
siblings go as K, because they shard the payload by world and it does not. It
is the one operator in the family where scaling out makes fusion less
attractive, and the one where a narrower wire is worth more than any epilogue
tuning.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* gemm_ag: issue a chunk's packets from world-1 lanes, not one thread
The fused-SDMA epilogue was spending 39-55us of GEMM time issuing a chunk's
world-1 SDMA packets back to back from thread 0 -- nearly constant across four
cells whose payloads differ 8x, which is SDMA's ~2us per packet regardless of
size, the constant gemm_ar's layout already states.
--post lanes gives each of lanes 0..world-2 one packet. What made that awkward
is that a broadcast chunk arms every destination at once, so the natural counter
is one per chunk -- and the thread that wins a single counter then has to tell
world-2 other lanes, which costs a barrier and an LDS round trip for one dword.
So the counter goes back to one per (chunk, lane), gemm_a2a's shape for a
different reason: every block increments all world-1 in one instruction, they
cross the threshold together, and each lane learns from its own atomic. No
broadcast.
The submit lock had to go with it. world-1 lanes retire in lockstep, so two
waves each holding part of a lock set and spinning for the rest could never
reach their release. The same collision -- two chunks of one destination posting
at once -- is separated by queue instead (chunk % sdma_queues), hence the new
sdma_queues >= chunks requirement, which the benchmark derives rather than asks
for.
Measured, ptpc r4: the fused GEMM phase goes flat at 57.4-60.4us across c1..c8
at M=512 and 66.6-71.9 at M=2048, where serial climbs to 105.2 and 168.8. At
M=2048 c8 that is -96.9us of GEMM time. Best total 113.1 -> 105.5 (-6.7%) and
224.1 -> 217.2 (-3.1%). All twelve quant x chunk cells validate at the same
relL2 as every other mode.
It still does not beat split-sdma (105.5 against 104.1; 217.2 against 210.5),
and fixing it exposed that the first round's headline was wrong: serial's drain
shrank with chunk count not because the fusion absorbed communication but
because the GEMM was stalled posting and the copy engines got a head start,
which the phase split booked as absorbed comm. With the posting fixed the drain
grows with chunk count instead -- the ~2us per packet moves from the issue side
to the wire, and chunking an all-gather multiplies packets by world-1 per chunk
because every peer gets the same slab. --chunks 1 is now the best fused setting
everywhere, and the remaining gap is ~10us of bookkeeping against ~7.5us of real
early-start.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* gemm_ag: a bf16-in / fp32-or-bf16-out path
gemm_ag was built on gemm_ar's fp8 GEMM, and the layer it exists for is not
fp8: DeepSeek V4-Pro's wkv_gate goes through linear_bf16_fp32, and the planning
doc says in as many words that the first implementation must keep that input
precision and FP32 output convention rather than substituting the FP8 GEMM. So
the operator had the right shapes and the wrong precision, and its numbers could
not be put next to the FP32 baseline they are meant to replace.
_gemm_a16w16_8wave.py ports gcnasm's gemm_a16w16_quad_subtile_kernel_template
at its own instantiation <512, 256, 256, 64, bf16, bf16, bf16, float>. The port
is small because gemm_ar's _gemm_a8w8_8wave.py is already a FlyDSL rendering of
that same template with fp8 operands -- same 8 waves, same 2x2 half-tiles, same
swapped-AB MFMA, same permlane store. Four things change: the MFMA becomes
16x16x32 bf16, E_K goes 1 -> 2, VEC goes 16 -> 8, and there are no scales. The
LDS keeps mori's XOR swizzle restated in bytes (swizzle_row128b), asserted in
the tests to be the fp8 function value-for-value at elem_bytes=1.
The FP32 store turns out to be the cheap one. The A/B swap puts four
consecutive columns in each lane; at two bytes that is 8 B and needs a
permlane16_swap to pair two N-tiles into 16 B, and at four bytes it is already
a buffer_store_dwordx4 with no shuffle and no convert.
elem_bytes becomes load-bearing rather than decorative: layout validates it,
kernels_lsa derives its pack width from it, and the benchmark gains
--in-dtype {fp8,bf16} and --out-dtype {bf16,fp32}, both defaulting so every
existing number stays reproducible.
Measured, 8x MI355X, K=7168 N=2048. Accuracy first: relL2 9.9e-7 for fp32 out
against 1.66e-3 for bf16 out -- three orders of magnitude, and it shows the
1.66e-3 was always the bf16 *output* rounding rather than the fp8 inputs.
All 32 cells validate. Fusing now breaks even (fused-sdma within 0.5% of
split-sdma at every cell, with 4-9% of the transfer genuinely absorbed) where
under fp8 it lost by 3-7% -- the bf16 GEMM is 2.2x larger, which is exactly the
G0/C0 lever the bound section describes. fused-lsa gets worse, and fp32 makes it
much worse (616 against split-sdma's 428), because its epilogue pushes world
copies of every element through the CUs.
The honest headline: against the planning doc's FP32 baseline the collective leg
agrees to 3.4% (327.8 against 339.2) but our GEMM is 28% slower than
torch.matmul (137.6 against 107.2), so at the target shape this is not yet a win
and the whole deficit is in the multiply. BLOCK_K is pinned to the template's
64, xcd_swizzle is off and no tile sweep has been run; tuning that is worth more
than any further epilogue work.
Also recorded: FlyDSL caches compiled kernels keyed on the @flyc.jit wrapper's
name, not the kernel body, so a fixed bug can look unfixed. Here it produced a
thoroughly convincing false signal -- odd K_ITERS bit-exact, even K_ITERS broken
-- which looks exactly like a double-buffer parity bug and was not one.
One cell (fused-sdma bf16 M=2048) failed validation once and passed the next 13
runs. Left in the doc as unresolved rather than dismissed.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* gemm_ag: sweep the bf16 GEMM's tile, and find a bug instead of a speedup
The bf16 port landed 28% behind torch.matmul, called out as the largest number
on the table. The cause is not the mainloop: the template's 256x256 tile
launches ceil(M/256)*ceil(N/256) workgroups, which at the wkv_gate shape
(M=N=2048) is 64 against an MI355X's 256 CUs. Three quarters of the GPU idles
and nothing inside the loop can reach that.
Measured at M=2048 N=2048 K=7168, fp32 out, one GPU, torch.matmul at 84.8us:
256/256 (template) 64 wg 176.0us 342 TF/s
256/128 128 wg 139.6us 431 TF/s
128/256 128 wg 134.2us 448 TF/s
128/128 256 wg 116.5us 516 TF/s
waves_per_eu (1/2/4) and xcd_swizzle (0/1/4/8) are inert -- every cell within
1%. Across shapes the rule is not "maximise the grid" either: at M=4096,
128/128 gives 512 workgroups and 162.8us against 128/256's 256 workgroups at
155.2, so once there is one workgroup per CU the larger tile amortises better.
pick_tile takes the largest tile whose grid still covers the CUs, which
reproduces the measured best at all four shapes.
**Neither faster tile ships.** 128/128 races: one wrong fp32 C in 48 single-GPU
runs, and about half the ranks per 8-rank launch. 128/256 is worse than it first
looked -- gemm-to-window and split-rccl validated 4/4 and 3/3, then came back
1/2 on a later pass, while fused-lsa is 0/4 throughout. BLOCK_M=256 has not
failed once across every repeat run. So the failure rate at 128 is around a
third of launches, and the tidy "split paths are sound, fused paths are not"
story the first pass suggested did not survive repetition.
The obvious suspect is the release fence, since halving BLOCK_M doubles the
blocks that must publish. It is wrong: --fence all and --direct-fence all are
0/3 each.
Buying 24% for a wrong answer a third of the time is not a trade, so pick_tile
returns the template's tile and the speedup does not reach the default. What
ships is the knowledge: the tile is worth 24-34% and is the only knob that is,
the mechanism is grid occupancy rather than scheduling (so the next attempt is
split-K), and there is a latent correctness bug that halving BLOCK_M exposes
which was invisible at the template's tile. Finding it is worth 33% on the GEMM
and is now reproducible.
Also here: layout accepts block_n=128 under an fp32 C (a bf16 store pairs two N
tiles through permlane16_swap and cannot); the fused-SDMA C store moves to a
buffer resource at this rank's recv slot, matching gemm_ar, which fixed a wrong
answer at M=2048 fp32 (relL2 6.9e-3 -> 9.9e-7); and compile_bf16_gemm_ag now
rejects a tile that disagrees with its config, which would hang the completion
counter rather than corrupt it.
All eight modes validate at the shipped default.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* gemm_ag: the BLOCK_M=128 corruption is the mainloop's vmcnt, confirmed
The tile sweep left a bug it could not explain: BLOCK_M=128 is 24-33% faster
and corrupts about a third of 8-rank launches, including gemm-only, which has
no collective in it at all. This pins the cause.
`wait_barrier(N)` is `s_waitcnt vmcnt(N)` plus `s_barrier`, so it is the only
thing making a G2S load's LDS write visible before the next iteration's S2R
reads it. Sweeping that one number, BM=128 fp32 M=2048, gemm-only, six
launches each:
vmcnt(2A+B) = 4 (inherited) 5/6 91us
vmcnt(A+2B) = 5 (derived) 4/6 90us
vmcnt(0) (everything) 6/6 150us
So it is the waits, and `vmcnt(0)` is correct. A `--wait-policy` knob is added
to make that reproducible rather than a one-off.
It also shows the *derivation* is wrong, which matters more than the knob.
Walking the mainloop -- entering with A+2B outstanding, issuing A1@k+1, B0@k+2,
A0@k+2, B1@k+2, needing the 4th item retired -- gives a requirement of
wait <= A+2B, under which:
256/256 A=2 B=2 in-flight 14 wait 6 -> retires 8, needs 8, margin 0
128/256 A=1 B=2 in-flight 11 wait 4 -> retires 7, needs 6, margin 1
The tile with *zero* margin is the one that is clean and the tile with margin
is the one that fails, so the model does not describe the hardware. Something
is outstanding on vmcnt that this counting does not include -- the S2R reads
are the obvious suspect (16 per iteration at BM=128, 24 at BM=256) if they
lower to flat loads rather than ds_read, since flat counts against vmcnt.
Settling that needs the generated ISA, which the FlyDSL cache does not keep
(pickles, not objects).
Nothing about the shipped configuration changes: the default is still the
template's BLOCK_M=256 at wait_policy="tuned", which is what every mode
validates on. vmcnt(0) is not a fix worth taking either -- at 150us it is
slower than BLOCK_M=256's 137.6, so it spends more than the tile saves.
What this buys is a bounded next step. The question is no longer "why does a
smaller tile corrupt" but "how many vmem ops are actually in flight at that
one instruction", which the ISA answers directly.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* gemm_ag: the ISA clears the vmcnt counts, and the kernel with them
The wait-count hypothesis was confirmed by a sweep (vmcnt(0) is 6/6 where
vmcnt(4) is 5/6) and is refuted by the generated code. Both are worth having.
Dumping with GPU_DUMP_CODE_OBJECT=1 and llvm-objdump --mcpu=gfx950, at
M=2048 N=2048 K=7168 fp32:
* the S2R loads are `ds_read_b128` -- lgkmcnt, not vmcnt. They were never part
of this accounting, which kills the leading theory outright: 16 per iteration
at BM=128, 24 at BM=256;
* the G2S loads are `buffer_load_dwordx4`, exactly A+B per load call, 6 per
iteration at BM=128 and 8 at BM=256;
* `vmcnt(4)` and `vmcnt(6)` each retire precisely one iteration's worth,
leaving in flight exactly the k+2 prefetches. Both are correct;
* the two skeletons are structurally identical, 8 barriers per iteration in
both, with `s_waitcnt lgkmcnt(0)` correctly between the barrier pair that
guards the write-after-read on each LDS buffer;
* neither spills, and both get one workgroup per CU. The only asymmetry is
152 VGPRs / 96 KiB LDS against 248 / 128 KiB, which is just the tile.
One more candidate came out of reading it and died the same way.
gemm_ar.wait_barrier puts s_waitcnt and s_barrier in a single inline-asm string
with constraints="" -- no memory clobber -- so the compiler is never told there
is a barrier in there and could legally move an LDS access across it, which
would be a write-after-read on the double-buffered LDS and would plausibly
depend on register pressure. Emitting the wait as asm and the barrier as
rocdl.s_barrier() produces byte-identical code. That option is therefore not
kept: an option that provably cannot differ is not an option.
So vmcnt(0) fixing the corruption is timing perturbation, not a corrected
count. Put beside the other end of the evidence -- 0 wrong in 64 single-GPU
runs at every tile, no collective in the loop -- the weight has moved off this
kernel. What reproduces it is eight concurrent processes; what never
reproduces it is one idle GPU running the same kernel on the same data.
Next place to look is the harness or the environment: whether the benchmark's
own fp32 reference can be corrupted under memory pressure, and whether the cco
window's VMM allocation interacts with a plain torch.empty C (ROCm 7.2.x is
known to route uncached VMM allocations to the coarse-grained pool). The
--wait-policy knob stays so the sweep is reproducible.
Nothing about the shipped configuration changes.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
* gemm_ag: fix bf16 pipeline ordering and split-pull publication
Order LDS reads before buffer reuse and drain the final prefetches. Honor peer_uncached for unfused C stores and publish from every producer lane. Initialize the full benchmark window so stale arrival flags cannot bypass the first pull barrier.
Add pipeline dependency, exact numerical, graph replay, and dirty-window regressions. Validation: 151 related tests passed; BM=128/256 with BF16/FP32 output each passed five independent eight-rank benchmark launches.
* gemm_ag: drain all SDMA queues before publishing arrival
Chunked fused producers can submit through every queue to a peer. Wait for all of those queues before signalling completion, while retaining the existing split and single-queue waits.
Validate the publication contract with 24 queue-state cases. The corrected path also passed eight-rank BF16-to-FP32 graph replay checks with one, two, and four chunks.
* gemm_ag: keep direct benchmarks and reusable pipeline kernels
Expose split-K FP32 partials, serialized SDMA chunk posting and release-only
producer publication through the kernel API. Retain direct GEMM, transport
and pipeline benchmarks with shared graph timing and changed-input checks.
Preserve the FP8 wait-policy argument fix and reject nonfinite peer output.
Keep source generators, historical experiment plans and raw datasets out of
the branch. Document the supported commands and retain the measured findings.
Validation: 129 CPU checks; 54 gfx950 numerical/graph cases; 11 eight-rank
BF16/FP8 transport and MORI/split-K/Torch pipeline checks; eight-rank native
GEMM comparison; Black 25.1.0 and Ruff 0.12.0.
* docs(gemm_ag): keep benchmark reports alongside the operator
* docs(gemm_ag): consolidate benchmarks and trim experiment history
* gemm_ag: identify model shapes and benchmark V4.1 Flash
Label the Pro ratio-4 projection measurements and document Flash's separate
512-wide projections and combined 1024-wide candidate at K=5120. Record
8-GPU transport, M-sweep, pipeline and complete projection-order comparisons,
including interleave restoration and KV/gate materialization costs.
Use thread-local capture for RCCL to avoid HIP watchdog event queries
invalidating global graph capture, and destroy process groups on CLI exit.
Validation: 52 reported configurations/repeats on v2-015 passed initial and
changed-input checks; maximum relative L2 error 8.04e-7. Two benchmark CPU
checks, Black, Ruff, documentation links and numeric checks passed. Raw
plans, scripts, logs and excluded attempts are archived outside the repo.
* docs(gemm_ag): record smaller-tile fusion experiments
* gemm_ag: rebase onto main, fix the pull test, add CI, re-measure
Rebased off the gemm_a2a branch (#695) onto main: nothing in
gemm_ag imports gemm_a2a, and the 14 commits applied without conflict.
* tests: _gemm_ag_pull_worker runs bench_gemm_ag.py through runpy.run_path,
which does not put the script's directory on sys.path, so the
`_bench_utils` split (25d419b) left test_pull_initializes_recycled_window
failing on every run with ModuleNotFoundError.
* mori.ops: register gemm_ag with the lazy FlyDSL submodules.
* ci_cco.yml: run the four gemm_ag test files, after the GEMM+AR steps. Their
8-rank SDMA process groups leave queues that take seconds to come back;
gemm_ag's tests settle and retry around that, gemm_ar's do not, and run
right after this step six of them failed in hsaKmtCreateQueueExt
(anvil.cpp:210) while passing on their own. As written the step passes
here: 213 tests in 22m50s.
* .gitignore: drop the local experiments entry; that archive lives outside
the repository.
* docs: BUILD_CCO_SDMA is a build flag and on by default since #710, so stop
exporting it at run time; example paths no longer point at a private
checkout.
* README tables re-measured on main 8ae8156 (with #603 and #711), 8x MI355X,
ROCm 7.14 / PyTorch 2.12 / FlyDSL 0.3.4.1: 54 configurations, all
validated, relL2 <= 9.95e-7. Conclusions hold: 128x128 + split SDMA leads
the Pro shape (355.9us, 12.6% under Torch + RCCL), fusing pays only for
the combined Flash N=1024 at M=16384 (4.3%). Host and stack changed with
MORI, so the deltas against September are not attributed.
Co-Authored-By: Claude <noreply@anthropic.com>
* docs(gemm_ag): re-measure the larger-M Pro study and refresh the summary
The README's opening summary still quoted the September figures, and its
strongest fusion claim (11% at M=16384) rested on a study the previous
re-measurement did not cover. Re-taken on the same stack: split SDMA against
sc_release c4 at N=2048, K=7168 gains 5.3% at M=8192 and 11.9% at M=16384,
and loses at M=4096 (c4 -4.8%, c2 -1.6%) -- September's v2-015 pairs gave
4.9-5.2%, 11.1% and -3.2%.
Co-Authored-By: Claude <noreply@anthropic.com>
* gemm_ag: take #785's pipeline fix after merging main
kernels_fused.py carries its own copy of gemm_ar's 8-wave fp8 mainloop, and
the merge of main brought #785 to gemm_ar only. Apply the same three changes:
drain the peeled tail's producers with wait_barrier(0) before it reads the
last step's B0/A0, restore the loop's end-of-step count to 2A+B, and make the
loop's S0 a read_done_barrier(). The comment telling readers not to restore
2A+B goes, since #785 reversed it. The BF16 quad-subtile GEMM is a different
pipeline and is untouched.
The four gemm_ag test files pass as the CI step runs them (213), and at K=256
-- where the main loop runs zero times -- fp8 fused-sdma validates 3/3 under
ptpc and 1/1 under mxfp8, split-sdma 1/1, all at relL2 1.66e-3 on 8x MI355X.
Co-Authored-By: Claude <noreply@anthropic.com>
* gemm_ag: bind the LSA kernel's receive-slot offsets outside its body
The copy kernel read cfg.recv_slot_off() inside @flyc.kernel. FlyDSL keys
its compile cache on scalar closure values, and those offsets depend on the
window's capacity, which nothing else in the key carries, so two configs
differing only there could share a cached kernel with the wrong offsets.
gemm_ar fixed the same thing in #748. Bind them to a tuple of ints instead.
Co-Authored-By: Claude <noreply@anthropic.com>
* docs(gemm_ag): re-measure after merging main, and record #748's short tail
crsuse2 v3-054, ROCm 7.2 / FlyDSL 0.4.0: every conclusion holds -- split
SDMA 10.4% under Torch + RCCL at Pro M=2048, fusion paying from M=8192
(4.8%, 11.2% at M=16384), the combined Flash projection 4.8% at M=16384.
gemm_ar's short final chunk was ported and measured; no variant beat the
best even split, so it is recorded rather than kept.
Co-Authored-By: Claude <noreply@anthropic.com>
* docs(gemm_ag): Torch + RCCL at every shape, and a GEMM / AG phase split
Against Torch + RCCL the best MORI path is 10-16% faster on V4-Pro and
6-11% on V4.1-Flash. A same-run phase split over all four GEMM x transport
combinations attributes it: SDMA's all-gather is 12.5-13.9% under RCCL's at
every shape, while the 128x128 BF16 GEMM is level or ahead only up to
M=4096 and 22-52% behind Torch's at larger M -- so Torch's GEMM followed by
SDMA is the strongest non-fused path, and fusion beats it only on Pro at
M >= 8192.
Co-Authored-By: Claude <noreply@anthropic.com>
* docs(gemm_ag): retune the BF16 tile per shape, and why fusion does not follow
Sweeping all four tiles: the right one matches Torch's GEMM from M=8192 and
beats it by 23% at Pro M=4096, and pick_tile's rule names it at every
shape, but TILE_CANDIDATES holds only 256x256. The faster tile speeds the
split path, not the fused one: 256x256 launches a single round of
workgroups at Pro M=8192 and Flash M=16384, so every chunk completes at
once and there is nothing left to overlap. Best path per shape is now
11.7-16.1% under Torch + RCCL at M >= 4096.
Co-Authored-By: Claude <noreply@anthropic.com>
* gemm_ag: SDMA only -- remove the LSA transports
split-lsa-push, split-lsa-pull and fused-lsa go, with kernels_lsa.py, the
pull worker and its tests, as gemm_a2a (#695) did with its own. At the Pro
shape the best LSA path was 4.6% behind split SDMA and fused LSA 16-19%
behind; LSA led only on the smallest payload measured (Flash N=512, M=2048,
by 3%) and trailed by up to 2% on the other M=2048 shapes, which does not
pay for a second transport in the test matrix. The dated tables keep their
LSA rows, marked removed.
The SDMA barrier keeps its spin loop, now local to kernels_sdma. The fused
epilogues store into this rank's own slot only, so gemm_ar's _SwapABStoreC
no longer needs the peer_rsrcs list this PR added: it is back to main's, and
the PR no longer touches gemm_ar. The unfused bf16 GEMM's sc0|sc1 publication
existed for the pull baseline and is rejected now; peer_uncached remains for
the fused path's sc_release configuration.
Co-Authored-By: Claude <noreply@anthropic.com>
* gemm_ag: pick the BF16 tile per shape and per path, and the fused chunk count
TILE_CANDIDATES now holds all four tiles. pick_tile's unfused rule (largest
tile whose grid covers the CUs) names the fastest GEMM at all six measured
shapes. The fused path gets its own rule: the largest tile other than
256x256 whose grid runs at least two workgroup rounds, with one chunk per
round (pick_chunks) -- the best fused configuration at Pro M=8192 / 16384
and Flash M=16384, 3.9% / 7.4% / 2.2% under split SDMA and 1.9-3.4% under
the 128x128 fused runs. The benchmark uses both when --block-m, --block-n
and --chunks are left at 0 (--chunks now defaults to 0).
The smaller tiles were held back until the LDS ordering fix had broader
coverage; test_every_tile_validates_split_and_fused adds 8-rank coverage
for all four, split and fused.
Co-Authored-By: Claude <noreply@anthropic.com>
* gemm_ag: benchmarks use the installed mori, not the checkout's source tree
_bench_utils.py put the checkout's python/ at the front of sys.path, which
shadows the installed package. CI installs mori non-editable, so the source
tree there has no compiled mori.cco.cco and every benchmark process the tests
spawn died on import: the GEMM+AG CI step has failed (24 tests) since it was
added. An editable install, where the extension sits in the tree, hid it.
Drop the insertion, and stop the benchmark README from exporting PYTHONPATH at
the checkout for the same reason.
Co-Authored-By: Claude <noreply@anthropic.com>
* gemm_ag: add GemmAgOp, a recv-reuse credit, and fold the tests and docs
GemmAgOp is the public wrapper: it owns the window, picks tile and chunk
counts per M (with overrides), keeps one counter slot per shape, pads rows,
and has a self_test with teeth. supports()/shape_constraint() state the
shape rules up front.
The fused FP8 and BF16 epilogues and the split gather now wait for a
per-peer credit before pushing into that peer's recv: every rank posts its
call epoch at entry. Without it a rank a call ahead could overwrite a result
a slower rank was still reading; the benchmark replays identical operands,
so a stale read there was the right answer and nothing caught it. The
back-to-back case in the new op test fails with the credit off.
Tests: test_gemm_ag.py now holds everything that needs no process group
(layout, the LDS and SDMA models, the benchmark contract, single-GPU
numerics) plus the benchmark's two baselines; test_gemm_ag_op.py runs every
multi-rank case through GemmAgOp from one 8-rank spawn. EXPERIMENTS.md is
merged into README.md.
Co-Authored-By: Claude <noreply@anthropic.com>
* gemm_ag: apply black
Co-Authored-By: Claude <noreply@anthropic.com>
* gemm_ag: reset the election counters, add mode="auto", and wrap-safe waits
Review follow-up.
- The fused epilogue's election counters ran on and were tested modulo the
chunk's tile count, which needed a counter set per M (GemmAgOp refused a
ninth distinct M) and elects the wrong tile once a counter crosses 2^31,
or 2^32 unless the count is a power of two. The tile that completes a
chunk now puts its counter back to zero, as gemm_a2a does, in both the
fp8 and the bf16 kernels; one counter set serves every M and chunk count,
and max_shapes / the layout's shape slots are gone.
- _spin_until compares (int32)(cur - flag) < 0, like a2a_gemm's
_spin_reach: the unsigned compare passed at once on the call whose epoch
wraps to 0.
- GemmAgOp: mode="auto", now the default, fuses an M once even a 256x256
tile's grid covers the CUs and splits below it. That fits every pair in
the README's table, where the fused tile's round count (2 at both Pro
M=4096 and 8192) does not; through the op at P=4 and P=8 it picks the
faster path at M=2048, 4096 and 8192. fence and peer_uncached are
exposed, defaulting to the README's sc_release. The CU count comes from
the device. M pads to 128 (256 under mxfp8, or a pinned tile's BLOCK_M)
instead of always 256; pick_tile already only picks tiles dividing M.
- Tests: one fused instance through ten distinct M with every counter back
at zero afterwards; alternating M back to back with no host sync and a
held reader on an auto instance that switches between fused and split,
with the credit off as the negative control; the auto rule on the
README's shapes; chunk-count and alternating cases pin mode="fused".
Co-Authored-By: Claude <noreply@anthropic.com>
---------
Co-authored-by: xiangch <xiangch@amd.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-authored-by: Xiang.Chen <x@x>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
BUILD_CCO_SDMAwas OFF unlessBUILD_BENCHMARK=ONturned it on. What OFF produces is a build whose SDMA puts compile away: every symbol is still there, every kernel still launches, every put still returns, and no bytes move. The all-reduce then returns mostly the local slice.GemmAllReduceOp.self_testexists because of exactly one such run. It measured faster than the real thing (-6.8% against -2.3%) precisely because it was not moving data, and the only signal that caught it was perplexity: 862511 against 3.26.What it costs
Same tree, clean build each side:
libmori_cco.soOFFON+6,088 B (+0.36%), no measurable build time.
It adds no dependency.
anvil.cppis inMORI_APP_SOURCESunconditionally andfind_package(hsakmt REQUIRED)is too, so the host transport and its link deps were already there either way. The flag reaches only cco's own layer: the queue setup incco_init.cpp, thecco_sdma_*wrapper symbols (58 lines of 316), the device API incco.hpp, and one extra JIT bitcode variant that is cached per config.Why it is safe
It creates nothing on its own.
ccoSdmaSetupCommQueuesreturns before any allocation unlessMORI_ENABLE_SDMAis set, and that runtime gate is independent and still defaults off.Verified on the new default build, two ranks:
MORI_ENABLE_SDMAsdmaQueueCount1So for anyone not opting in this changes nothing; for anyone who does, it is the difference between working and silently inert.
CI already built with
BUILD_CCO_SDMA=ONin bothci_cco.ymlandci.yml, so ON was already the configuration under test and OFF was the less-covered one.The diagnostic
Case B above used to be silent too. When a caller explicitly asks for queues (
sdmaQueueCount > 0;0means "whatever the env wants") and none will be created,cco_init.cppnow says which of the two reasons applies:At ERROR level rather than WARN because the default global level is ERROR — a warning here would be invisible to exactly the person it is for. Not fatal: a comm without SDMA is legal and every other path still works. Silent for the default
sdmaQueueCount=0, so a comm that never wanted SDMA logs nothing.Verification
On MI355X, from this branch:
pip install .bakesBUILD_CCO_SDMA=True;BUILD_CCO_SDMA=OFFstill builds and bakesFalse— the escape hatch workstests/python/cco/test_sdma_api.py— 12 passedtest_rocm_bootstrap.py+test_rocm_import_order.pywithMORI_ENABLE_SDMAunset — 9 passedOne note on the box used: an unrelated container was holding SDMA engines part of the time, and
sdma_queue_count=8failed withERROR code: 6atanvil.cpp:237while 1 and 2 succeeded. That is the contentionkernels_sdma.pyalready documents ("56 queues to use 7"), and6isHSAKMT_STATUS_NO_MEMORY— a process-memory condition, not queue exhaustion (see #685). It cleared on its own and case C passes above; this change does not touch queue creation.Docs
Updated the two places that called OFF "the default" and the build instructions that told users to pass the flag by hand. The
BUILD_CCO_SDMA=ONrequirement is unchanged everywhere — only the need to type it.🤖 Generated with Claude Code