Repository navigation
feat(cco): Fuse DSV4-Pro wo_b's GEMM with its all-reduce over cco SDMA - #662
Conversation
Adds a FlyDSL 1-stage/2-stage all-reduce over cco symmetric windows, plus the
harness and tests to compare it with aiter's custom all-reduce. Target shape is
DeepSeek-V4-Pro wo_b on 8x MI355X: [M, 7168] bf16, TP8.
Structure follows benchmark/cco/triton/gemm_a2a: pure-Python layout math that is
unit-testable without a GPU, a torchrun benchmark emitting one RESULT_JSON line
with critical-rank decomposition, and a pytest that runs it on 2/4 ranks.
Protocol parity with aiter's custom_all_reduce.cuh: per-(block, peer) signal
slots with a monotonic _flag[block] that is never reset (graph-replay safe),
peer rotation in the reduce-scatter, stage 2 reusing stage 1's thread->index
mapping (cross-device visibility only holds between threads with the same tid),
and matching accumulation order so results are bitwise comparable.
Status: correctness done -- rel_l2 == 0.0 against a float32 reference on 2 and 8
ranks at M=64/512/4096, n=7168. Performance is not yet at parity (M=4096: 798us
vs aiter 265.7us). The likely cause is the one deliberate divergence from aiter,
noted in kernels_lsa.py: each thread reads all peers straight into registers
instead of staging through LDS and reducing on warp 0. A/B of the two is next,
along with the SDMA backend and folding the aiter baseline into this harness.
Two bugs in examples/cco/python/05_flydsl_lsa_allreduce surfaced while porting
it, both of which make its cross-GPU barrier a no-op. The example still passes
because its inputs are host-written before the kernel, so nothing in it actually
depends on the barrier; the 2-stage all-reduce here does, which is how they
showed up. Both are worked around in _compat.py rather than in the example:
1. Its cache-modifier constants are gfx942-specific. On gfx950 the aux
encoding is bit0=SC0, bit1=NT, bit2=SC1, so SC0|SC1=3 assembles to
"sc0 nt" -- a non-temporal hint that does not bypass L2. Confirmed by
reading the emitted ISA (aux=3 -> sc0 nt, aux=5 -> sc0 sc1).
2. rocdl.raw.ptr.buffer.load is side-effect-free, so the scf.While spin loop
is dead-code-eliminated outright: the barrier compiles down to a signal
store and nothing else.
The signal path here is therefore atomic load/store with an explicit syncscope
(system to publish, agent to poll, mirroring aiter's __scoped_atomic_*), which
is neither removable nor dependent on hand-encoded cache bits.
Also carried over from aiter: _wave_uniform_i64 on peer base addresses. Window
.lsa_ptr is an opaque extern call, so without it the backend cannot prove the
buffer descriptor is wave-uniform and wraps every 16-byte access in a
readfirstlane waterfall loop.
Two fixes, found by A/B rather than by inspection. At [4096, 7168] bf16 on 8x MI355X the kernel goes from 798us to 286us; aiter is 271us. 1. All-gather loop nesting (the 2.2x). Stage 2 had the peer loop *outside* the index loop, so it drained one peer completely before starting the next -- pinning the all-gather to a single xGMI link for 7 consecutive passes. aiter unrolls the peer loop inside the index loop (custom_all_reduce.cuh:566) so all 8 links stream at once. The arithmetic matches the measurement: one 7.34MB slice over one 55GB/s link is ~133us, and 133 + 7*95 is the 798us we were seeing. The old nesting is kept as `--gather sequential`. It is the more natural way to write the loop and the penalty is invisible in the output, so it is worth being able to demonstrate rather than just warn about. 2. Grid size (a further 1.25x). We were using aiter's 80 blocks, but the two kernels put a different amount of work on each thread: aiter's thread takes one pack from one peer (its warp_id picks the peer), ours takes one index across all 8. Matching aiter's grid therefore over-subscribes the xGMI request queues by 8x, and the time rises monotonically past ~32 blocks (286us @28 -> 361 @80 -> 583 @256). Cap is now 24, and the per-block stride is `threads` rather than `threads / world`. K_MAX_BLOCKS stays 80: it sizes the signal array, which is a correctness bound, not a tuning knob. The two are now separate. Best-of-3, 101 iterations, CUDA-graph replay, vs aiter custom AR: M aiter LSA ratio 64 25.6us 30.7us 1.20 128 27.2 41.8 1.54 512 53.6 66.5 1.24 1024 85.6 100.6 1.17 2048 146.1 160.3 1.10 4096 271.1 285.9 1.05 8192 558.8 531.7 0.95 Parity (+-15%) holds for M >= 1024. Below that both implementations sit on a payload-independent floor -- ours is ~28us and does not move between M=8 and M=64 -- so those ratios are two constants divided, not throughput. Best-of-3 is not defensive: single runs of *either* implementation vary by up to 1.7x (aiter's M=64 spanned 25.6-44.0us), and an earlier single-run sweep reported 204us at M=2048 where the stable figure is 160us. Two things checked and left alone: - aiter's 1-stage/2-stage threshold is right for us too. Forcing 1-stage is worse at every small size (M=512: 150us vs 68us), so the extra barrier costs less than the 4x traffic. - This box is affected by known-issue 3 (ROCm 7.2 puts uncached VMM allocations in the coarse-grained pool): CCO_UNCACHED_WINDOW=0 changes nothing, which is the documented tell. aiter uses IPC handles and does not pay it. Not quantified here, and it does not account for the whole small-M gap. The earlier note blaming the gap on our register-direct reduce vs aiter's LDS-staged one was wrong -- that divergence is real but costs nothing, and the register version is what makes the smaller grid necessary. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Same algorithm, window and wire traffic as the LSA 2-stage kernel; the copy
engines move the bytes instead of the CUs. Three kernels (scatter push /
local reduce / gather push) because both transfers need a device-wide
ordering point and a kernel boundary is the only one that is free and
graph-capturable.
Arrival is signalled over LSA, not SDMA. This is forced, not a preference:
FlyDSL's Sdma exposes no wait_signal, and its put wrapper hard-codes
remoteSignal=false (cco_device_wrapper.cpp:140), so a peer cannot observe an
SDMA signal at all. That leaves gcnasm's opus_sdma_a2a_quiet_notify shape --
no per-PUT signal, quiet_queue on the sender, then an LSA atomic into the
peer's flag. The flag array, its monotonic epoch and the spin loop are reused
from kernels_lsa; only row 0 is needed, since one warp issues the pushes.
Resolves the open question the plan flagged. FlyDSL's Sdma._xfer docstring
said a no-signal put "cannot be drained by quiet"; the gcnasm README said the
opposite. The README is right -- cco.hpp:1707 states quiet()/quietQueue() are
"independent of the signals", and quietQueue drains the queue read pointer
(cco.hpp:1790). Verified rather than assumed: --sdma-signal A/Bs the trailing
ATOMIC, and the no-signal default validates rel_l2=0 on 2/4/8 ranks while
being slightly faster (M=512: 79.2us vs 84.0). The misleading docstring is
corrected here.
Measured on 8x MI355X, [M, 7168] bf16, best of 3 x 101 iterations:
M aiter LSA (/aiter) SDMA (/LSA)
64 25.6us 30.7us 1.20 47.4us 1.54
128 27.2 41.8 1.54 51.7 1.24
512 53.6 66.5 1.24 79.2 1.19
1024 85.6 100.6 1.17 116.4 1.16
2048 146.1 160.3 1.10 194.2 1.21
4096 271.1 285.9 1.05 342.4 1.20
8192 558.8 531.7 0.95 637.8 1.20
SDMA does not reach the plan's +-15%-of-aiter bar (it lands at 1.26-1.48x for
M >= 512) and on this evidence it will not: the gap decomposes into a ~12us
fixed floor -- SDMA's payload-independent time is ~40us against LSA's ~28us,
which is two transfer phases times the engine's ~6us dispatch cost -- plus
~18% lower sustained bandwidth than CU-issued loads, still present at M=8192.
Neither is a tuning knob. cco's own chunking is not implicated:
CCO_SDMA_MAX_COPY_BYTES is 1GB, so each 7.3MB slice is a single packet.
This backend is not here to be faster. It is here because the transfers occupy
one warp of one block rather than the whole grid, which is the property step 3
needs to hide an all-reduce inside a GEMM epilogue. Reporting it as a
standalone collective is the wrong scoreboard, so the numbers above are
recorded as a baseline for that comparison rather than as a result.
layout.py gains an optional `recv` region (one landing slice per peer, slot
`rank` unused so the index is the peer rank). LSA passes recv_slots=0 and its
window is unchanged.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Step 3 of the wo_b plan. The GEMM's C lands directly in a registered cco
window, and each destination's slice is pushed by the copy engine as soon as
its last tile is written, so the reduce-scatter overlaps the rest of the GEMM.
Reduce and all-gather are unchanged, reused from ar/kernels_sdma via the new
build_sdma_phases with scatter swapped for a drain-only twin.
The result is negative, and cleanly attributed.
8x MI355X, [4096, 7168] out, K=1024, graph replay, median of 51:
gemm-only 49.6us
split-lsa 327.0
split-sdma 382.0
fused-sdma 450.2 (chunks=2)
Fused is 66-68us worse, and by the same 66-68us at K=4096 and K=8192 -- a
fixed cost, not a scaling one. It is the release fence:
* --fence all (every lane) costs 70us over the 448 blocks: 452.0 vs 382.4.
* --fence leader is *wrong*, measured not inferred: relL2 goes 2.35e-3 ->
8.9e-3 on 8 ranks. s_waitcnt vmcnt(0) plus s_barrier retires every lane's
stores and is still not enough, which is exactly what
cco_device_wrapper.cpp:105-111 says -- leaderOnly orders thread 0 only and
does not substitute for a release from every producer lane. I had assumed
the opposite (that the engine reads through our own L2 so no fence was
needed at all); the first 8-rank run disagreed at 3.76e-3 vs split's
2.35e-3, which is why the modes now cross-check each other in the tests
rather than each being judged "small enough" alone.
* --fence none is incorrect too and exists only to price the ceiling: 373.0
vs split's 382.0. So even with a free fence the overlap is worth ~9us,
against a GEMM of 49.6us.
Why ~9us and not ~50: 448 blocks over 256 CUs is 1.75 waves, so no
destination's first chunk completes before ~57% of the GEMM has elapsed, and
each push then pays the engine's ~6us dispatch. Only the tail is available to
hide anything in. Finer chunking helps as far as it can (chunks=2 is 450.2 vs
chunks=1's 505.6) but the slice is only 2 BLOCK_M row-bands here, so 2 is the
cap.
The obstacle underneath is the ratio the plan flagged as risk 6: the transfer
is ~2.7x the GEMM (7.3MB per link at ~50GB/s is ~147us, twice, against 49.6us
of compute), so the links are the critical path with or without overlap and
perfect fusion could remove at most the GEMM. Raising K does not rescue it --
K and the collective are independent axes.
Two things that did work and are worth keeping:
* Destination-rotated, chunk-major tile order (gcnasm's
opus_direct_stripe_tile). aiter's linear order finishes destination 0's
whole slice first, so the last destination's link only starts at the end of
the GEMM. Rotating costs nothing measurable: 49.72us vs 49.80.
* The completion-counter protocol itself. Monotonic per (dest, chunk)
counters with a modulo test elect exactly one CTA per chunk, so no
per-destination spin lock is needed (gcnasm needed one because several CTAs
could submit), and never resetting them keeps the kernel graph-safe.
The LSA transport is deliberately not implemented; LSA_NOTE in kernels_fused
records why. It needs StoreC coalesced first (its _store_bf16 is a per-lane
scatter that gcnasm measured at 0.26x when pushed to a peer), and the ceiling
that rewrite competes for is the ~50us just shown to be unreachable.
kernels_fused.py pins a copy of aiter's kernel_gemm rather than adding a hook,
because the original's main loop has a runtime `if wave_m == 1` and FlyDSL
only rewrites `if` inside the @flyc.kernel function's own AST -- factoring the
body into a shared helper makes that line raise at trace time. Everything
reusable is imported; the copy is the ~90-line pipeline only, and
test_gemm_ar.py asserts it is bit-identical to aiter's kernel so it cannot rot
silently.
Not a production claim: aiter's tuned CSV picks ck/asm/cktile for
(N=7168, K=1024), not this host.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
… fence
rocprofv3 --kernel-trace contradicts the attribution in the previous commit.
I had blamed a flat ~70us of fence spread across the pipeline, inferred from
end-to-end totals. Per-kernel it is nothing like that (8x MI355X,
[4096, 7168] out, K=1024, median of 11, us):
config gemm scatter/drain reduce gather sum correct
split 39.1 136.8 59.0 137.3 372.3 yes
fused --fence none 64.5 103.4 59.8 136.6 364.2 NO
fused --fence leader 83.5 96.9 59.9 136.2 376.5 NO
fused --fence agent 156.1 67.3 59.8 137.0 420.1 yes
fused --fence all 175.9 71.1 59.3 137.3 443.6 yes
Two corrections to what I said before. The overlap is not "worth ~9us" -- it
is worth 69.5us, the scatter dropping from 136.8 to 67.3, more than the whole
39.1us GEMM. And the cost is not distributed: all of it is inside the GEMM
kernel, +117us. Registers are unchanged (VGPR=128, SGPR=112, scratch=0), so
it is not occupancy.
The +117us splits three ways:
* +25us for the tail with no fence at all (barrier, atomic, and the copy
engine reading C while the GEMM still writes it).
* +92us for `buffer_wbl2 sc1` and the `s_waitcnt vmcnt(0)` on it. This is
quadratic: buffer_wbl2 flushes the *entire* L2, so block k writes back every
tile blocks 1..k-1 already wrote, 448 blocks over. The split path gets one
bulk release at kernel end for free.
* +20us more for system scope over agent scope. The ISA shows why:
__threadfence_system is bidirectional and emits
`buffer_wbl2 sc0 sc1 ; s_waitcnt ; buffer_inv sc0 sc1`, and that buffer_inv
discards the B weights the GEMM is streaming, 3584 times.
Three attempts to make the release cheap, all failed, all now selectable so
the measurements can be reproduced:
* agent scope (new default): correct, drops the buffer_inv, buys 20us of 112.
cco exposes no agent-scope release -- only cco_system_fence -- so this is
emitted as a raw llvm.fence via a new _compat.release_fence. That gap is
worth closing in the device API regardless of this benchmark.
* one fence per block instead of per wave (--fence leader): incorrect, relL2
9.6e-3, indistinguishable from no fence, even though every lane has passed
s_waitcnt vmcnt(0) and a block barrier. So the release is genuinely
per-wave, which is the factor of 8 in 3584.
* write-through C stores (--fence writethrough, a one-argument cache modifier
on StoreC's copy atom): not a release on its own (relL2 1.72e-2), and
combined with the agent fence (--fence wt-agent, 432.5us) it is
indistinguishable from the agent fence alone (432.4). buffer_wbl2 costs the
same whether or not the lines are dirty.
So the conclusion is narrower and more useful than "the compute:comm ratio is
wrong". What is missing is a **per-tile release**: every primitive available
publishes the whole cache, while the GEMM needs to publish 448 tiles. The
counter protocol and the destination-rotated tile order are sound -- they are
what produced the 69.5us of real overlap -- and would pay immediately on a
transport whose producer-side release is per-tile.
The bottom line is unchanged: fused 432.4us loses to split-sdma 382.0 and to
gemm+LSA 327.0, which is still the number to beat.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The trace decoder *is* installed -- /opt/rocm-7.2.4/lib/librocprof-trace-
decoder.so. My earlier "not available" was a bad search: /opt/rocm is a
symlink and find(1) does not follow those without -L. With
--att-library-path the decode works and attributes the epilogue exactly.
The tail is 26% of the kernel's latency, and almost all of it is waiting:
vaddr cycles % kernel instruction
19396 91336 5.1% s_waitcnt vmcnt(0) retire this block's C
19400 22984 1.3% s_barrier
19404 292 0.0% buffer_wbl2 sc1 <- the op itself is free
19412 246804 13.9% s_waitcnt vmcnt(0) <- waiting on the writeback
19608+ 91612 5.1% waits bracketing the counter atomic
The writeback wait is the largest single line item in the whole kernel --
larger than any main-loop stall -- and it is quadratic, since buffer_wbl2
flushes the entire L2 and every one of the 448 blocks waits for every earlier
block's tile.
The last row was my own mistake and is now fixed. I had given the counter
atomic acq_rel, so it emitted its own buffer_wbl2/buffer_inv pair costing more
per wave than the explicit fence. Both halves were redundant: the caller has
already issued the release, and the winning block never reads what it counts
(the copy engine does), so it needs no acquire. Relaxed:
fused --fence agent gemm 156.1 -> 136.9us, total 432.4 -> 413.4us
still relL2 2.35e-3, matching split exactly.
Cumulative from profiling this path: 450.2 -> 413.4us, of which 20us was
agent-vs-system scope (dropping a buffer_inv that discarded the B weights the
GEMM streams) and 19us this atomic.
Conclusion is unchanged and now has instruction-level evidence: what is
missing is a per-tile release. Every primitive available publishes the whole
cache while the GEMM has 448 tiles to publish. Fused 413.4us still loses to
split-sdma 372.3 and to gemm+LSA 327.0.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Profiling the fusion turned up a much larger miss in the plain SDMA
all-reduce. Its reduce kernel was inheriting cfg.blocks, i.e. LSA_BLOCK_CAP,
which is 24. That cap exists to throttle outstanding *xGMI* requests, and it
is right for the LSA kernel. The SDMA reduce reads local HBM. Same constant,
opposite regime.
At [4096, 7168] it left the kernel with 24 x 512 threads x 8 packs = 1.6MB in
flight, about a fifth of what covers HBM latency: 66.1MB moved in 59.3us, or
1.11 TB/s on an ~8 TB/s part. Sized on its own (SDMA_REDUCE_BLOCK_CAP = 256,
flat from 256 to 512 in a sweep) it is 13.5us, 4.9 TB/s.
standalone SDMA all-reduce, [M, 7168] bf16, best of 2 x 101 iters:
M aiter LSA (/aiter) SDMA before SDMA after (/LSA)
64 25.6us 30.7us 1.20 47.4us 45.6us 1.49
512 53.6 66.5 1.24 79.2 73.6 1.11
2048 146.1 160.3 1.10 194.2 167.6 1.05
4096 271.1 285.9 1.05 342.4 294.5 1.03
8192 558.8 531.7 0.95 637.8 544.6 1.02
So the Step 2 conclusion needs correcting: SDMA is not "~1.20x slower than LSA
and not meant to win". Past M=2048 it is within 5% of LSA, and at M=8192 it
beats aiter. What remains is almost entirely the ~40us dispatch floor, which
is why decode sizes still lose ~1.5x.
Same fix flows into the fused GEMM path:
split-sdma 382.0 -> 334.3us
fused-sdma 413.4 -> 365.4us
split-lsa 327.0us (unchanged)
Also settles, negatively, whether the fused epilogue's release can be
per-block instead of per-wave. `--fence agent-leader` emits the fence inline
in our own kernel rather than through cco's leaderOnly wrapper, ruling out the
compiler having sunk it past the atomic. Still incorrect: relL2 1.36e-2. Two
independent implementations agree, so one wave's buffer_wbl2 genuinely does
not publish its sibling waves' stores even after s_waitcnt vmcnt(0) and
s_barrier. Its *time* is the useful part though -- 322.5us, which would beat
split-lsa -- so the 8x reduction in fence count is worth 41us and the prize
for a cheaper release is now priced.
Retracting one claim from the previous commit: "buffer_wbl2 costs the same
whether or not the lines are dirty" was based on a broken experiment.
--fence writethrough asks for sc0|sc1 but the ISA shows the store emitted with
`sc0` only, so the data stayed dirty in L2 and nothing was tested. The copy
atom's cache_modifier is a two-value enum (0=cached, 2=nt), not the aux
bitmask that raw_ptr_buffer_store takes, so sc0|sc1 is not expressible through
StoreC at all. Whether write-through would help is still open.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…t stores)
Adds two probes that separate the three things the fused epilogue costs, which
previously could only be inferred from differences of totals.
--no-put keeps the counter protocol and the fence but drops the transfer
nt-agent non-temporal C stores + agent fence
At [4096, 7168], K=1024, 8 ranks:
none --no-put 206.7us epilogue without a fence, no transfer
agent --no-put 284.9us +78.2us <- the release fence, alone
agent 364.3us +79.4us <- the transfers
nt-agent 411.2us
Three results:
1. The barrier and atomic are nearly free. A GEMM with no epilogue at all is
~197us in this pipeline, so the counter protocol costs ~9us, not the ~25us
I had inferred. Most of what I attributed to "the tail itself" was actually
the copy engine reading C while the GEMM writes it.
2. The release fence is 78.2us, now measured with no SDMA in the picture at
all rather than derived from a difference of totals.
3. The overlap is 57.4us: the transfers add 79.4us on top of the fenced
no-put run, against 136.8us for the standalone scatter kernel.
Which prices the whole thing: 334.3 (split) - 57.4 (overlap) + 9 (protocol)
= ~284us if the release were free, against split-lsa's 327. That is a 13% win,
materially better than the "~10-20us" I estimated from totals, and it makes
the StoreC coalescing rewrite worth costing out rather than dismissing.
nt-agent is a negative result and stays in as one: non-temporal C stores make
it *worse* (411.2 vs 364.3). They do not make buffer_wbl2 cheaper, and they
cost the GEMM's stores. The only untried way to avoid the writeback is real
write-through (sc0|sc1), which the copy atom cannot express -- it needs a raw
raw_ptr_buffer_store, and probably needs StoreC coalesced first, since its
per-lane 2-byte scatter bypassing L2 would write 32B partial lines to HBM.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
… default
Running the X1 precursor experiment turned up something more important than
the experiment: the fused path is intermittently wrong, and has been. Every
"validated=True" I reported for it was a single sample. Repeating the runs:
config passes note
chunks=1, --fence agent, acq_rel atomic 6/6 the only stable one
chunks=2, --fence agent, acq_rel atomic 5/6 RACY
chunks=2, --fence agent, monotonic atomic 3/5 RACY, worse
chunks=1, --fence raw-wt-agent 2/4 RACY
standalone SDMA all-reduce (control) 5/5 stable
Two races, both mine:
1. chunks > 1 shares an SDMA queue. The modulo test elects one block per
*chunk*, not one per *queue*, so two winners post to queue `dest` at once.
The plan called for a per-destination submit lock and I skipped it, arguing
the counter already elects a single CTA -- true, and about the wrong thing.
Default is now chunks=1; >1 stays settable to reproduce it.
2. The counter atomic needed its acquire half. I made it relaxed two commits
ago for a 19us win, reasoning the winner never reads what it counts. Wrong:
the winner needs the acquire so the other blocks' releases are ordered
before *its put*. acq_rel is the default again, and it measurably lowers the
failure rate (5/6 vs 3/5) though it does not fix race 1.
The X1 answer is no. Write-through C stores looked like the fix -- they make
buffer_wbl2 cheap because there is nothing dirty to drain, GEMM 139.1 -> 97.6us
and total 333.8 against split-lsa's 327 -- and they are racy at chunks=1 too,
so that number describes a kernel that does not reliably compute the right
answer. Not worth building the StoreC coalescing rewrite on.
Getting there needed two ISA checks, because both earlier write-through
attempts silently did nothing:
- fx.rocdl.BufferCopy*(cache_modifier) is a two-value enum (0=cached, 2=nt),
not an aux bitmask, so sc0|sc1 cannot be asked for through the copy atom.
- _compat's CM_SC1 was 4. Assembling each policy and diffing the encoding
(llvm-mc, gfx950) puts SC0 at insn bit 14, SC1 at bit 15, NT at bit 17 =
CPol 1, 16, 2. So SC1 is 16; 4 is DLC, RDNA-only, silently ignored on CDNA.
Latent -- the barrier uses atomics and payload stores use CM_CACHED -- but
it invalidated two experiments before showing up in a disassembly.
One genuinely useful sub-result survives: the un-coalesced per-lane 2-byte C
store does *not* collapse when it bypasses L2. That was the open question
about whether StoreC must be coalesced first, and the answer is no.
Corrected standings, correct configurations only:
split-lsa 324.4us
split-sdma 338.0
fused-sdma chunks=1, agent 440.0
Adds test_fused_is_stable_across_repeats, which runs the fused path three
times and requires the results to be *identical*, not merely each small. It
skips below 8 GPUs, since the queue race needs more than one chunk per
destination and the wo_b shape only has one at 2 ranks. A one-shot check
would not have caught either race.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…348us
From reading gcnasm's opus_gemm_a2a_lsa, which fuses a GEMM with an a2a and
gets 16-27% out of it. Its ChunkFused epilogue
(gemm_a16w16_quad_subtile_kernel_template.hpp:701-741) uses **no cache fence
at all** -- s_waitcnt vmcnt(0), s_barrier, an __ATOMIC_ACQ_REL counter, and a
per-destination submit spin lock. No __threadfence_system, no buffer_wbl2.
That is correct and I had missed it: the acq_rel counter *is* the release. Its
ISA already emits a buffer_wbl2/buffer_inv pair, on thread 0 only -- once per
block, not once per wave. The explicit per-wave fence I had added on top was
redundant, and it was 90us:
fused chunks=1, agent fence 440.0us gemm 152.4
fused chunks=1, no fence 348.5us gemm 62.7 stable 6/6
The earlier conclusion that "the release fence is mandatory and costs 78us"
was wrong. --fence none had only ever been measured at chunks=2, where a
separate bug corrupts the result, so the fence was being credited for fixing
something it had nothing to do with.
Also ports gcnasm's per-destination submit lock (atomic test-and-set with
s_sleep backoff, ISA-verified) and its post-submit barrier. Neither fixes
chunks>1, which stays racy in roughly 1 run in 3 and is now documented as
unexplained after ruling out, by measurement: the lock, the barrier, the fence
at both scopes, and the atomic ordering. The only factor that separates working
from broken is whether the put is issued mid-GEMM or at the end. Default stays
chunks=1; note that makes the overlap ~zero (scatter 129.0 vs split's 136.8),
which is why fused still loses.
Corrected standings:
split-lsa 324.4us
split-sdma 338.0
fused-sdma chunks=1, no fence 348.5
With the fence gone the entire remaining deficit is +23.6us inside the GEMM --
barrier, counter, and the copy engine reading C while the GEMM writes it --
against a 7.8us overlap gain.
Four other things from that README worth having on record here: its collective
is one phase where an all-reduce is three, so fusion can only ever touch a
third of the traffic; its compute:comm is 2.6:1 where wo_b is 0.29:1; its best
mode (Direct LSA) has the epilogue store straight into the peer, so it has no
mid-kernel publication problem at all; and it also found chunked PUTs gave no
end-to-end benefit on the non-fused path.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…ut is racy
--mode fused-lsa has the GEMM epilogue store straight into the destination
rank's window, gcnasm's best a2a mode. No staging, no copy engine, no
completion counter, no submit lock, no release fence -- the only sync left is
the cross-rank barrier after the kernel.
Structurally it does what fused-sdma never could (rocprofv3, us):
config gemm scatter/barrier reduce gather sum
split-sdma 39.1 136.8 13.5 137.3 326.8
fused-sdma c=1 62.7 129.0 13.0 136.2 341.0
fused-lsa (Direct) 166.6 13.3 11.1 135.5 326.6
The scatter collapses 136.8 -> 13.3us; the transfer is entirely absorbed. And
the un-coalesced store is far better than feared: 7.34MB per link in 166.6us is
~44 GB/s, 82% of the SDMA scatter's 53.7, against the 0.26x gcnasm measured for
a lane-scatter pushed to a peer.
It is not correct, and finding out why turns the StoreC coalescing rewrite from
an optimisation into a prerequisite:
- cached peer stores validate 2-3 runs in 6. The lines are homed in the peer
but sit dirty in whichever XCD's L2 the block ran on, and nothing
downstream can clean that: the barrier kernel is one block and buffer_wbl2
writes back only its own XCD's L2, one of eight.
- uncached sc0|sc1 peer stores are *consistently* wrong (relL2 ~3e-2) and
slower (426 vs 341us). _store_bf16 writes one bf16, so bypassing L2 makes
every store a 2-byte partial-line write and concurrent sub-32B writes into
a sector across the fabric lose updates. An uncached recv load in the
reduce does not change it, so it is the store, not a stale read.
L2 merging is what makes the narrow store work at all; L2 residency is what
makes it invisible. Only a wider store separates them -- which is exactly why
gcnasm needed its ds_bpermute pair-coalesced C store.
Even coalesced it likely ties rather than wins: 326.6 against split-lsa's ~324.
LSA's 2-stage does read+reduce+write in one pass over the wire so its reduce is
free, while Direct LSA writes to the peer and then pays a separate local reduce,
with a store rate ~18% under LSA's read rate.
Also plumbs two switches the mode needs: build_sdma_phases gains
reduce_self_from_recv (the epilogue writes my own slice into my recv slot, so
j=0 comes from there) and recv_uncached.
fused-lsa is deliberately left out of test_gemm_ar_modes_agree while it is racy.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…worth it
StoreC emits 128 buffer_store_short per block because the MFMA accumulator
layout gives each lane four consecutive *rows* (stride c_cols = 14336B), not
columns. The eight bf16 that would make a 16-byte store live in eight
different lanes. Ports gcnasm's mfma_adaptor_swap_ab (opus.hpp:2064, literally
base::operator()(b, a, c) with dim_c() redefined) as --swap-ab: computing
(A B)^T in the accumulator's own layout puts four consecutive *columns* in a
lane. A/B scale loads swap with it, scalar for A and vec4 for B.
Bit-identical to aiter's kernel at three shapes including wo_b, guarded by
test_swap_ab_is_bitwise_identical. And slower:
variant per lane lanes/row bytes/row rows/instr gemm-only
unswapped (128 short) 2B 16 32B 4 50.4us
--swap-ab (32 dwordx2) 8B 4 32B 16 55.6us
+ 16B/lane (probe) 16B 4 64B 16 48.1us
Store instructions drop 4x and total kernel instructions 30%, and it does not
help, because the transaction count is unchanged -- both layouts put 32 bytes
in a row and need 16 of them. The swap just moves those 16 from four
instructions onto one, so per-instruction address fan-out goes 4 -> 16. ATT:
buffer_store_dwordx2 costs 1022 cycles each against buffer_store_short's 114.
--store-probe answers whether the permlane stage is worth building before
building it: it emits the 16B-per-lane / 64B-per-row pattern with deliberately
wrong data, so it prices the access pattern with the shuffle's cost removed.
48.1us against 50.4 is a -4.6% ceiling locally, and on fused-lsa -- where the
store is 14% of the kernel with 60x the stall -- it is a wash
(340.8/338.7/333.1 vs 336.4/335.3/337.7), because that phase is xGMI
bandwidth-limited, not transaction-limited. So the permlane stage is not built.
The floor is the 16-column MFMA tile: 16 bf16 is 32 bytes, a quarter line,
whatever the operand order. 64B needs two N-tiles merged, which needs a shuffle
in *both* layouts; 128B needs merging across waves through LDS, which gcnasm
measured at a 1-3% regression.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
I dismissed this last commit on a guess that permlane16_swap + shfl would eat
the 4.6% the --store-probe ceiling showed. The guess was wrong: it needs two
permlane16_swap per M-tile, no ds_bpermute and no LDS, and they cost 4,260
cycles to save about 33,000.
The mapping falls out once permlane16_swap's semantics are measured rather
than assumed -- vdst' = [X.r0, Y.r0, X.r2, Y.r2], vsrc' = [X.r1, Y.r1, X.r3,
Y.r3]:
(A, B) = permlane16_swap(tile0.d0, tile1.d0)
(C, D) = permlane16_swap(tile0.d1, tile1.d1)
lane group g stores (A, C, B, D)
which lands g = 0,1,2,3 on columns 0-7, 16-23, 8-15, 24-31 -- together columns
0..31 contiguously. 16 bytes per lane, 64 per row, and the column permutation
is absorbed into the address.
Bit-identical to aiter's kernel at three shapes including wo_b. GEMM kernel
median of 42 dispatches, 8 ranks:
variant bytes/row store instrs store cycles gemm kernel
unswapped 32B 128 short 14,608 38.30us
--swap-ab 32B 32 dwordx2 32,720 slower
--swap-ab --permlane 64B 16 dwordx4 14,604 37.04us
The store is not where the time goes, which was the part I had wrong in both
directions. Its total latency is unchanged, 14,608 vs 14,604, and --swap-ab
alone is slower because at 32B/row the same 16 transactions get squeezed onto a
quarter as many instructions (1022 cycles each vs 114). permlane's 64B rows
just put the store back where it started.
What pays is what a 2-byte store drags along. Thread trace, unswapped ->
permlane:
v_mul_f32_e32 257 -> 1 17,664 -> 68 cyc
v_pk_mul_f32 0 -> 128 0 -> 9,276
v_lshlrev_b32_e32 159 -> 45 11,304 -> 3,828
v_add_u32_e32 133 -> 21 8,944 -> 1,604
v_cndmask_b32_e64 96 -> 8 6,368 -> 656
v_cvt_pk_bf16_f32 128 -> 64 8,384 -> 4,244
v_permlane16_swap_b32_e64 0 -> 32 0 -> 4,260
128 stores need 128 addresses, 128 bounds predicates and 128 scalar multiplies;
16 need 16. And the swapped layout makes B's scale a vec4, so the scaling packs
into v_pk_mul_f32. Kernel instructions 2062 -> 1421, registers unchanged
(VGPR 128 / SGPR 112 / 0 scratch).
Two limits: it does nothing on fused-lsa (335-338us either way), where the
store phase is xGMI bandwidth-limited rather than issue-limited; and 64B is one
wave's ceiling, since a wave owns N_TILES_B*16 = 32 columns.
Worth pushing back into aiter's StoreC: bit-exact, no register cost, and the
3.3% is on the GEMM itself, independent of any all-reduce.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
… GEMM
Third stage, from the ds_bpermute block at
opus_gemm_a2a_lsa/gemm_a16w16_quad_subtile_kernel_template.hpp:491-510. After
the permlane stage a row's four 8-column chunks sit in lanes 16 apart, so a
16-lane group touches 16 rows at 16 bytes each. Transposing the lane index --
lane l' = 4r'+q' pulls from the lane holding (row r', chunk q'), with g being
q's two bits swapped, 0,1,2,3 -> 0,2,1,3 -- puts adjacent lanes on one row, so
lanes 0-3 write 64 contiguous bytes and a group covers 4 rows instead of 16.
GEMM kernel, median of 42 dispatches, 8 ranks, [4096,7168] K=1024:
stage gemm store instrs store cycles
aiter as-is 38.30us 128 short 14,608
--swap-ab slower 32 dwordx2 32,720
--swap-ab --permlane 37.04us 16 dwordx4 14,604
--swap-ab --permlane --lane-transpose 35.74us 16 dwordx4 3,852
Bit-identical to aiter at all three shapes including wo_b, registers unchanged
(VGPR 128 / SGPR 112 / 0 scratch) at every stage.
The striking part is that the third stage changes neither the instructions nor
the addresses -- same 16 buffer_store_dwordx4 over the same 64 addresses, only
which lane holds which -- and the store's own latency falls 14,604 -> 3,852.
So the coalescer is sensitive to lane adjacency and not just to the address
set, which is what gcnasm's "pair-coalesced" comment is really about. The 64
ds_bpermute cost 9,604 cycles.
Priced before building, as with the previous stage: a --store-probe variant
emits the adjacent-lane access pattern with deliberately wrong data, and
measured 45.4 against 48.1us, so the real version at 46.3 lands close to its
ceiling.
Note the three stages pay for different things. --swap-ab alone is slower.
--permlane does not speed the store up at all (14,608 -> 14,604); it pays by
deleting what a 2-byte store drags along -- 128 addresses, 128 bounds
predicates and 128 scalar multiplies become 16, and B's scale becoming a vec4
packs the scaling into v_pk_mul_f32 (v_mul_f32_e32 257 -> 1). Only
--lane-transpose actually makes the store faster.
Still nothing on fused-lsa (335-338us either way): that store phase is xGMI
bandwidth-limited. And 64 bytes is one wave's ceiling here, since a wave owns
N_TILES_B*16 = 32 columns; gcnasm reaches a 128-byte line only by tiling four
wave_id_n waves across adjacent 16-column runs.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…bad combos
The C-store variant is picked by a chain of const_expr branches, and there was
no `direct_lsa and ... lane_transpose` branch. So `--mode fused-lsa
--lane-transpose` did not error -- it fell through and quietly ran stages 1+2.
That produced a wrong conclusion. I reported "the store optimization does
nothing on fused-lsa, that phase is xGMI bandwidth-limited" on the strength of
a measurement that only ever exercised the two stages which do *not* speed the
store up (permlane leaves it at 14,608 -> 14,604); stage 3, the one that takes
it to 3,852, was never in the binary.
With the branch actually wired, the conclusion survives, now for a reason that
was measured rather than assumed:
gemm barrier reduce gather sum
fused-lsa 177.4us 5.4 11.1 135.3 329.1
+ all three 175.6us 7.8 11.1 135.0 329.4
against 39.0 -> 35.7us for the same three stages on the split path's local
store. Stage 3 buys the local memory pipeline's sensitivity to lane adjacency;
a peer store crosses the fabric, and that phase is bounded by ~44 GB/s of link
bandwidth, which transaction shaping does not move.
Also adds argument validation for the store flags -- permlane requires swap_ab,
lane_transpose requires permlane, store_probe requires swap_ab, and the two
combinations that have no branch are rejected by name. Silent fallthrough is
what made the original mistake possible and it should not be possible again.
Current standings at [4096,7168] K=1024, 8 ranks, graph replay, median of 3,
default C-store / with all three stages:
gemm-only 49.6 / 46.8us
split-lsa 324.8 / 322.8 correct
split-sdma 333.2 / 331.4 correct
fused-sdma 348.6 / 340.2 correct
fused-lsa 334.2 / -- RACY, 1 run in 3
The C-store work barely moves the collectives -- the GEMM is ~15% of the
pipeline, so 6.7% of it is ~3us. Its value is on the GEMM itself and is
independent of any all-reduce.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…gh scf.for
Implemented gcnasm's persistent scheduler (one workgroup per CU striding over
tiles), got it bit-exact, measured it, and reverted it. Recording the result
and the two non-obvious bugs it surfaced, because the next person to read
gcnasm's PERSISTENT=1 will want to try the same thing.
The tail it targets is real and large. Sweeping N at M=4096, K=1024:
N tiles remainder time
4096 256 0 32.64us
4352 272 16 45.88us <- 6% more work, 40% more time
6144 384 128 47.40
7168 448 192 49.80 <- wo_b
But wo_b sits at remainder 192, the benign end of gcnasm's own sweep (+0.77%
there, against +9.97% at remainder 8). And the implementation is a ~2x
regression for a reason unrelated to scheduling:
variant gemm scratch
one workgroup per tile 38.54us 0 B
body wrapped in an scf.for 74.48us 156 B/thread
fly-promote-regmem-to-vectorssa does not promote make_rmem_tensor values
declared inside an scf.for region, so the MFMA accumulators spill. Persistent
does beat the *looped* baseline by ~3% (82.3 vs 84.8us at remainder 192, 76.8
vs 79.8 at remainder 16), consistent with gcnasm -- but that is a gain over an
already 94%-regressed kernel.
Two findings kept in the docstring:
- The prologue's `if wave_m == 1: s_barrier()` runs the half-waves one
barrier out of phase on purpose. In a loop that offset accumulates by one
per tile, and from the second tile on the halves rendezvous at mismatched
program points and the LDS double-buffering races -- silently, as partial
corruption inside otherwise-correct tiles. A compensating
`if wave_m == 0: s_barrier()` per tile fixes it exactly; that is how the
persistent version was made bit-exact.
- The LDS handles must be rebound per tile. Hoisting them makes eight
shared-address-space pointers loop-carried and fails to legalize.
Also worth knowing for anyone wrapping a FlyDSL kernel in a runtime loop:
`range(...)` has to appear literally in the `for` statement or the AST rewriter
does not see it and Python evaluates it eagerly. Assigning it to a variable
first raises "dynamic 'ArithValue' has no Python integer representation", and
the workaround of always emitting the loop (step = tile count when not
persistent) makes the default path pay the spill too.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…the promoter
The previous commit blamed fly-promote-regmem-to-vectorssa for not promoting
through scf.for. That is wrong, and checking it properly says so:
- the pass handles scf::ForOp (rewriteForOp)
- it promoted all 451 register allocas in the persistent kernel, zero left
afterwards and zero llvm.alloca in the final IR
- the emitted loop carries no iter_args at all
- the 156 bytes are ordinary register-allocator spill:
.vgpr_spill_count = 38, with 38 scratch_store_dword / 38 scratch_load_dword
The real wall is the register budget, from the ISA kernel metadata:
variant VGPR AGPR SGPR V-spill scratch instrs gemm
as committed 256 0 50 0 0B 2052 38.54us
+ 3-stage C store 254 0 51 0 0B 1573 35.74us
body inside an scf.for 256 0 68 38 156B 2260 74.72us
persistent, grid 256 256 0 67 38 156B 2264 72.22us
The baseline already uses all 256 VGPRs with zero spill. A runtime loop asks for
38 more -- the K pipeline's accumulators and operand fragments die at the end of
a tile in the flat version, but a loop makes them potentially live across the
back-edge -- and there is nowhere to put them. --amdgpu-num-vgpr 256/512 changes
nothing because the cap was never the constraint. Persistent and non-persistent
spill identically, which is what pins the cost on scf.for rather than on the
scheduling idea.
The pass does have one real inefficiency, just not one we hit: an alloca
declared *outside* a loop becomes an iter_arg unconditionally even when fully
overwritten each iteration, because collectTouchedRegAllocaInRegion records on
any load or store with no liveness test. Allocas declared inside the loop are
correctly materialised as ub.poison; ours are all inside. Minimal repro and the
IR diff are described in the docstring.
Also notes that the 3-stage C store is the only change measured to reduce
register pressure (256 -> 254 VGPR, 2052 -> 1573 instructions, because 128
two-byte stores need 128 addresses and 128 predicates live at once) -- relevant
if anyone revisits this, though 2 VGPRs against 38 is not close.
Working tree reverted; the persistent path is not kept.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Found by review, not by me: the prologue's `if wave_m == 1: s_barrier()` gives
waves 4-7 one extra barrier, and s_barrier is a counting rendezvous, so from
then on waves 0-3 run one phase ahead for the rest of the kernel. A one-shot
GEMM does not care -- the trailing unmatched barrier is released when the other
half exits, and the offset is deliberate, it staggers the two halves of M.
The fused epilogue does care. `wait_barrier(0)` after store_c is supposed to
mean "every wave's C tile has retired". Under the offset it rendezvouses waves
0-3 -- which contain thread 0, hence the counter and the transfer -- with waves
4-7 sitting at the *previous* barrier, before their stores. Thread 0 then counts
the tile complete and pushes it while half of it is unwritten.
One `if wave_m == 0: s_barrier()` before store_c re-balances the counts.
--chunks 2 goes from failing about one run in three to 12 of 12, and it is the
configuration that actually overlaps, so it is the default again.
That changes the Step 3 conclusion. 8 ranks, [4096,7168], K=1024, median of 3-4:
default C-store + 3-stage C-store
split-lsa 326.8us 321.8us
split-sdma 334.6 330.9
fused-sdma chunks=2 326.2 323.0
Fusing is no longer a loss: it beats split-sdma by ~2.5% and ties split-lsa. It
still does not win, because the overlap (scatter 136.8 -> 126.9us) is roughly
cancelled by what the epilogue costs the GEMM (39.1 -> 60.9), and reduce plus
all-gather -- 46% of the pipeline -- are untouched by construction.
Retracting two earlier diagnoses of this race, both wrong:
- "chunks > 1 shares an SDMA queue". gcnasm's per-destination submit lock was
ported and ISA-verified and did not fix anything. It is kept because cco's
one-issuing-warp-per-queue rule still applies, but it was not the bug.
- "the counter atomic needs its acquire half". acq_rel did appear to lower the
failure rate against monotonic, and with a one-in-three intermittent failure
that comparison was noise. acq_rel stays because release/acquire is right
for a producer handing tiles to a consumer, not because it was measured.
fused-lsa is still racy and is a separate problem -- its peer stores sit dirty
in eight XCDs' L2s, which no barrier fixes.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…ir headers The code fix landed in the previous commit but its docstring did not -- the edit's anchor had already been mangled by an earlier replacement that left duplicated section headers behind. Both are fixed here. Replaces "The chunking race, unexplained" and "Result: the fusion does not pay" with the actual cause and the new standings, and retracts the two wrong diagnoses (shared SDMA queue, missing acquire on the counter) in place rather than deleting them, since both are things a reader would otherwise try again. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
… does The compensating `if wave_m == 0: s_barrier()` added last commit was gated on `fuse`. It should not be. gcnasm closes the pair at opus_gemm_a2a_lsa/gemm_a16w16_quad_subtile_kernel_template.hpp:693, before store_c and *outside* its `if constexpr (ChunkFused)` -- unconditional. aiter's 8-wave kernel opens the stagger at gemm_a8w8_8wave.py:449 and never closes it, which is where the pinned copy inherited the imbalance. Gating it on `fuse` was reasoning from the symptom rather than the invariant. The split path has nothing after store_c so it cannot observe the imbalance today -- but leaving the barrier counts unbalanced is what makes the next thing added after the epilogue silently wrong, which is exactly how this bug arrived, and it is also what the persistent-tile experiment had to work around independently. Costs nothing measurable on the split path: +8 instructions, VGPR 256 -> 256, SGPR 50 -> 42, zero spill either way, gemm-only unchanged within noise. fused-sdma chunks=2 remains 6/6 stable at ~326us, and the full suite is 28/28. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…ot a win
Review spotted that the epilogue's four store_c calls each reload the scales,
and they do: the calls share base rows pairwise and base columns pairwise, so
every scale is fetched twice. The source-attributed thread trace counts it
exactly -- 16 buffer_load_dwordx4 at gemm_a8w8_8wave.py:191 where 8 addresses
are distinct, 8 buffer_load_dword at :199 where 4 are. 12 of 24 loads redundant
plus 12 redundant address computations, 4,276 cycles or 0.8% of the kernel. The
compiler cannot merge them because all four calls write the same reg_f32_*
register buffer, so they are a chain of overwrites rather than pure loads.
store_all() loads each once. It removes exactly the predicted loads and is
slightly slower:
variant VGPR scratch instrs dwordx4 dword gemm
--permlane --lane-transpose 254 0B 1581 72 16 33.32us
+ --hoist-scales 256 0B 1595 68 8 33.96us
Keeping both scale sets live across all four stores costs more in register
moves and re-materialised addresses than the twelve loads were worth, and it
spends the last two VGPRs of headroom. Bit-exact either way, and the bitwise
test now covers it.
Kept as a default-off switch rather than deleted: the redundant fraction grows
with N_TILES_B, so a larger BLOCK_N changes the arithmetic, and it flips sign if
register pressure ever loosens. Re-measuring a flag is cheaper than re-deriving
why it was rejected.
One honest note on the test runs: a single unidentified failure appeared once in
the combined suite and did not reproduce in seven subsequent runs (full suite
x3, fused stability and modes-agree x3, AR suite x2, all green). It is recorded
here rather than dismissed, since an intermittent failure is exactly what this
epilogue has produced twice before.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…lish
fused-lsa was wrong roughly half the time. The cause: in the LSA 2-stage
all-reduce the producer fences (ar/kernels_lsa.py, every lane of every block,
right after the tmp stores). Here the producer is the GEMM and the only fence
was in the separate barrier kernel -- one block, so one XCD's L2 out of eight.
The other seven kept the peer-homed lines dirty while the LSA flag, a
system-scope atomic, overtook them. Fencing at the end of the GEMM fixes it:
10 of 10 at the fp8 floor.
The fence can be leaderOnly, worth 80us (434 -> 355), and that is only legal
because the half-wave barrier pair is now closed: wait_barrier(0) really does
mean every wave's stores have retired into this CU's L2. It measured as
incorrect before that fix, which is what made a per-wave release look mandatory.
Two things that did not work and are recorded rather than dropped:
- uncached sc0|sc1 peer stores, so there would be nothing to flush.
Consistently wrong (~1.7e-2) at 2 bytes per lane and still at 16 after the
permlane stage, so the "2-byte partial-line writes lose updates" theory was
wrong and the cause is unknown.
- the 3-stage C store, worth 6.7% on a local store, is worth nothing here
(189.4us of GEMM either way) -- the peer store is xGMI bandwidth-bound.
Standings, 8 ranks, [4096,7168] K=1024, 3-stage C store, failures out of 10:
split-lsa 318.8us 0/10
fused-sdma chunks=1 342 0/10
fused-sdma chunks=2 325 3/10 fastest, still racy
split-sdma 329.7 0/10
fused-lsa 3-stage, leader 359.0 0/10
split-lsa still wins. Both fusions pay a separate local reduce that LSA's
stage 1 gets for free, and reduce plus all-gather are 46% of the pipeline and
untouched by construction.
Two corrections to the previous commit, from wider sampling:
- chunks=2 is NOT fixed. The barrier-phase fix took it from about 1 in 3 to
3 in 10; a second, rarer bug remains. Default goes back to 1. Adding a
fence there makes it worse, not better (leader 3/6 wrong, agent 5/6),
which is unexplained.
- the bench's validation gate was 5e-3, loose enough to pass a corrupt run at
3.97e-3 while the fused path was racing. Tightened to 3e-3: the fp8 floor
is 2.35e-3 and the corruption lands at 4-9e-3.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Direct LSA had been wrong 2-3 runs in 6 since it was written. The fix is one
fence in the right place. The LSA 2-stage all-reduce publishes from the kernel
that produced the data (ar/kernels_lsa.py fences right after its tmp stores, in
every block). Direct LSA's producer is the GEMM, and the only fence lived in the
separate barrier kernel -- one block, so one XCD's L2 out of eight. The other
seven kept the peer-homed lines dirty and the LSA flag, a system-scope atomic,
overtook them. Fencing in the GEMM gives 10/10 bit-correct runs.
--direct-fence leader (thread 0 only) rather than every lane is worth 80us,
434 -> 355. It is legal only because the half-wave barrier pair is now closed:
wait_barrier(0) really does mean every wave's stores have retired into this CU's
L2, so one wave writing it back covers all eight. It measured as incorrect
before that fix, which is what made a per-wave release look mandatory.
Two things that did not work, recorded so they are not retried: uncached
(sc0|sc1) peer stores are consistently wrong at *every* store width, including
16B/lane with --permlane --lane-transpose, which kills the partial-line-write
explanation and leaves the real one unknown; and an uncached recv load in the
reduce changes nothing, which rules out a stale read on the consumer side.
It still does not beat split-lsa. 8 ranks, [4096,7168], K=1024, all three
C-store stages, median of 4, correctness over 10 runs:
mode time correct
split-lsa 318.8us 10/10 <- best
fused-sdma chunks=2 325.0 7/10 racy
split-sdma 329.7 10/10
fused-sdma chunks=1 342.0 10/10
fused-lsa 359.0 10/10
The 41us gap for fused-lsa is structural: LSA's 2-stage does read+reduce+write
in one pass over the wire so its reduce is free, while Direct LSA writes to the
peer and then pays a separate 10.9us reduce, runs its stores ~14% under LSA's
read rate, and adds the publishing fence LSA gets for free by being local.
Absorbing the scatter completely (136.8 -> 10.1us) still does not cover that.
Also corrects two things about fused-sdma. --chunks 2 is *not* fixed: the
barrier-phase fix took it from 1-in-3 to 3-in-10, so the default goes back to 1
(10/10, ~17us slower). And the bench's validation gate was too loose -- the
corruption lands at 4-9e-3 against an fp8 floor of 2.35e-3, so 5e-3 reported a
corrupt run as validated, which happened at 3.97e-3. Gate is 3e-3 now.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…ge grids
At the real DSV4-Pro prefill shape ([16384, 7168] K=2048) none of the GEMM+AR
modes validated -- relL2 4e-3 to 2e-2, different on every run. The all-reduce
was innocent: bench_ar.py is bit-exact at the same 224 MiB payload for both
LSA and SDMA. The GEMM is what is wrong, and aiter's own
flydsl_8wave_gemm_a8 reproduces it, so this is upstream and not a
transcription error in the pinned copy.
The main loop closes with
wait_barrier(2 * N_LDS_STEPS_A + N_LDS_STEPS_B) # s_waitcnt vmcnt(n); s_barrier
which leaves too many global->LDS prefetches in flight across the barrier, so a
wave passes it and refills an LDS buffer another wave has not finished
consuming. The corruption is always a whole (wave_m sub-tile x all four wave_n)
region -- the signature of a shared A-side LDS buffer -- and the values are
wrong rather than zero, so it does not look like missing output.
The trigger is grid size, not kernel duration: clean at 896 blocks even with
K=16384, corrupt at 1792 blocks even with K=512.
Measured boundary at [16384, 7168] K=512 on 8x MI355X, 4-5 runs per count:
256x256 corrupts at 6 (aiter's), 5 and 4, and is clean at 3, 2 and 0; 128x256
corrupts at 4 (aiter's) and 3, and is clean at 2, 1 and 0. Both thresholds are
N_LDS_STEPS_A + N_LDS_STEPS_B - 1, which is what this commit uses. It costs
nothing -- 241/240/245 us against 240/246/237 us for aiter's count at
[16384, 7168] K=2048 -- and takes all four modes to 5/5 correct at relL2
2.350e-3, the fp8 floor.
Not a complete theory, and the comment at the site says so: 128x512 never
reproduces even with aiter's count, and forcing *any* of the 14 sync points in
the loop to vmcnt(0) also fixes it, so the same under-wait may exist elsewhere.
This is the smallest change that is measurably correct here. Still to do: push
it upstream to aiter.
Also add validation to --mode gemm-only. Without that control a GEMM bug reads
as an all-reduce bug, because every other mode validates the *sum* -- which is
why this took as long as it did to find.
--swap-ab, --permlane and --lane-transpose were opt-in because they were added
one rung at a time and each needed its own A/B. Measured together at the shape
the model actually runs -- [16384, 7168] K=2048, 8x MI355X, median of 31 -- they
help every mode and change no output (all still relL2 2.350e-3):
gemm-only 243.0 -> 230.8us -5.0%
fused-sdma 1342.6 -> 1299.9 -3.2%
fused-lsa 1600.9 -> 1563.2 -2.4%
split-lsa 1275.2 -> 1263.8 -0.9%
split-sdma 1269.8 -> 1260.8 -0.7%
They are now BooleanOptionalAction with default True, so a rung can still be
taken out with its --no- form. The dependency errors name that form, since
hitting them now means "--no-swap-ab without --no-permlane" rather than the
old "--permlane without --swap-ab".
Where the win is, from the per-kernel trace, is worth writing down because it
is not where the ATT profile first suggests. The stages coalesce the store
*issue*, so they pay off when C lands in local memory -- fused-sdma's GEMM goes
297.3 -> 256.6us -- and do nothing when it goes straight to a peer:
fused-lsa's GEMM is 954.5 -> 952.1, i.e. noise, even though 56% of that kernel
is in the C-store path. That path is xGMI-bound, and ATT already said so: 99%
of its store time is stall, not issue. Fewer, wider stores do not shorten a
wait for the link.
compile_fused_gemm_scatter keeps all three defaulting to False. Its default has
to stay "aiter's kernel verbatim" -- that is the premise of
test_pinned_copy_matches_aiter_kernel_bitwise, and the ladder is enumerated
explicitly by test_swap_ab_is_bitwise_identical.
Also refresh the module headline, which was still quoting [4096, 7168] K=1024 --
a shape that does not occur in the model.
--chunks > 1 is the whole overlap mechanism -- with one chunk per destination
the tile counter only fires when the destination's entire slice is done, which
under the rotated tile order is the end of the GEMM, so nothing overlaps. It
was pinned to 1 because it produced wrong output about 3 runs in 10.
That was never the chunk protocol. It was aiter's 8-wave GEMM under-counting
one s_waitcnt in its main loop (fixed two commits ago), which corrupted output
non-deterministically on large grids whatever the epilogue did. At
[16384, 7168] K=2048 on 8 ranks, chunks 1/2/4/8 are now 40/40 correct with no
hang, and chunks=8 alone is 25/25.
Default is now 8, halved until it divides the M-tiles per destination so small
M still runs (M=2048 -> 1, M=4096 -> 2, M=16384 -> 8).
Per-kernel medians, [16384, 7168] K=2048:
chunks GEMM drain reduce gather total end-to-end
1 257.5 493.8 44.0 496.4 1291.7 1299.6
2 266.5 379.4 44.2 496.4 1186.6 1194.9
8 265.0 303.7 43.8 496.6 1109.1 1115.7
The drain falls 38% while the GEMM grows 7.5us for the extra counter atomics
and the submit lock. Finer is worse: at 8 each PUT is 3.5 MiB, and 16 (via
--block-m 128) halves that to 1.75 MiB, under the knee in the SDMA bandwidth
curve, for 1130.8us.
Which flips the headline. At the real prefill shape:
fused-sdma chunks=8 1114.1us <- best
split-sdma 1261.7
split-lsa 1262.4
fused-lsa 1571.4
gemm-only 228.9
11.7% under split, and 25.3% under the 1491.3us the same layer costs in the
running model (GEMM 348.0 + NCCL 1143.3, medians over its 61 layers).
fused-lsa is unchanged and still loses: it spends 730us of its GEMM pushing C
over xGMI where the copy engines move the same bytes in 499, and ATT shows 99%
of that store time is stall, so the C-store coalescing buys it nothing.
Two hangs were seen at 8 ranks while the sweep was being set up and never
reproduced in the 100+ runs after. The submit lock is a plain test-and-set
spin, so this is not proof it cannot happen; the notes say to use timeout when
sweeping and to treat a hang as a finding.
Also rewrite the two stale sections this invalidates: the "chunks race" note in
kernels_fused, which concluded the race was unsolved and the default had to
stay 1, and the headline in bench_gemm_ar, which concluded fusing loses.
The destination-rotated chunk-major order we already had is gcnasm's
ChunkFused ordering, transposed -- its destination is on N and ours is on M, so
its `m_in_chunk -> peer -> inner` is our `n -> dest -> m_tile`. Worth
confirming what that is worth now that chunks actually overlap: at
[16384, 7168] K=2048, rotated 1114.7us against linear 1323.0, i.e. 18.6%. At
chunks=1 it was worth almost nothing, which is the same statement as "without
rotation the chunks buy nothing" -- linear at chunks=8 (1323.0) is no better
than chunks=1 (1299.6).
What we did *not* have is gcnasm's other variant, opus_direct_stripe_tile: it
rotates the destination every `m_stripe` tiles of the other axis (2) rather
than after a whole chunk. gcnasm applies it only to Direct LSA and keeps
chunk-major for its chunked-SDMA path, and that split reproduces exactly:
mode n_stripe=1 n_stripe=2
fused-lsa 1560.7us 1366.3 -12.5%
fused-sdma 1116.0 1113.7 noise
Per kernel, all of fused-lsa's gain is in the GEMM -- 937.9 -> 788.8us, -15.9%
-- plus 36us of rank skew off the barrier that follows it. reduce and gather do
not move. The mechanism is what the comment says: with chunk-major, the 256
resident blocks retire their peer stores in bursts of 28 to the same
destination; striping by 2 spreads them across all seven links. The cost is
that a chunk completes later, which is exactly why it is wrong for the chunked
path and why gcnasm does not use it there either. 4 is worse than 2 (1392.1).
So `n_stripe` defaults to 2 for fused-lsa and 1 elsewhere, mirroring gcnasm.
`--n-stripe N` overrides; 0 means the per-mode default.
fused-lsa is still last (1365.6 against split-lsa's 1262.4 and fused-sdma's
1114.5), but the gap to split is down from +24% to +8%. 10/10 correct at the
new default, and the test suite passes.
…major path
n_stripe was two branches -- 1 meant chunk-major and >1 meant the striped
order -- which left the finest setting untestable and hid that the two are the
same family. Chunk-major *is* a stripe of N // BLOCK_N, and measuring it both
ways agrees (1569.0 against 1573.9). So the branch now keys on
`n_stripe >= n_blocks_const` and the parameter spans 1..n_blocks.
Full sweep at [16384, 7168] K=2048, 8 ranks, the divisors of 28:
n_stripe 1 2 4 7 14 28
fused-lsa 1370.9 1362.5 1404.9 1402.9 1437.0 1569.0
fused-sdma 1117.9 1117.1 1114.8 1124.0 1117.5 1115.7
fused-lsa is monotone from 28 down to 2 and flat below that; 2 stays the
default. fused-sdma is flat across the whole range, so its default is now
spelled "chunk-major" rather than 1, which is what it always was.
The sweep kills two plausible explanations, both worth recording so they are
not proposed again:
* Not XCD locality. dest = (block / n_stripe) % 8 and blocks go round-robin
over the 8 XCDs, so a destination is fed by exactly min(n_stripe, 8) XCDs.
14 and 28 spread a destination's traffic over all eight XCDs and 2 confines
it to two -- and 2 is the fast one. More XCD coverage is monotonically
worse.
* Not chunk-completion timing. If it were, fused-sdma would prefer
chunk-major; it is indifferent. With 256 resident blocks out of 1792 the
first chunk completes inside the first wave of blocks whatever the order.
What is left is the simple reading: for peer stores, rotating the destination
often keeps all seven links busy, and 2 is already often enough.
Folding chunk-major into the n_stripe range made `n_stripe == 1` stop meaning "unstriped", but the guard still read `n_stripe != 1 and not rotated`. With `fuse=False` the default resolves to chunk-major (N // BLOCK_N, e.g. 2 at N=512), so every non-rotated build -- the split baseline, and test_pinned_copy_matches_aiter_kernel_bitwise -- now raised. Committed in 253c71b with the test failure visible; this is the fix. The guard is only about a *real* stripe, so it is now `n_stripe < N // BLOCK_N and not rotated`. Also bound n_stripe to [1, N // BLOCK_N] and require it to divide the N-tile count, which the striped index math assumed silently. Test suite passes (12/12) and the five modes are unchanged at [16384, 7168] K=2048: gemm-only 233.9, fused-sdma 1114.6, split-sdma 1264.2, split-lsa 1264.4, fused-lsa 1359.0, all validating.
… doubles
The benchmark's GEMM was a8w8 per-token/per-channel while the model dispatches
a block-scale kernel whose scales vary along K, so every fused-vs-split number
was being drawn at a compute:comm of 0.44 where the real one is 0.76 -- our
GEMM was 1.7x cheaper than the one the fusion actually has to hide behind.
`--quant blockscale` adds the model's form: A quantised 1x128, B 128x128, fp32
scales. `--quant ptpc` keeps the old path and stays the default.
The scale layouts are the ones aiter's gemm_a8w8_blockscale_bpreshuffle
consumes: sa logically [M, K/128] but **column-major** (sglang's
materialize_bpreshuffle_fp8_scale is t().contiguous().t()), sb [N/128, K/128]
row-major. Established by sweeping the four combinations against that kernel --
only this one lands at the fp8 floor, the others give 0.17 to 0.32.
Applying the scales needs work in the mainloop, since
`sum_k s_a[m,k] * s_b[nb,k] * d_k` does not factor. The obvious form -- promote
each K-block into a second accumulator and zero the MFMA one -- does not work,
and the reason is worth recording: zeroing removes the only dependency between
K-blocks, the K loop is `range_constexpr` so fully unrolled, and the scheduler
then hoists every block's MFMAs above every promotion. Measured at K=512, all
64 MFMAs in the first 400 instructions and all the arithmetic after, keeping
K_ITERS accumulator sets live at once: 256 VGPR with 81 spilled and the GEMM 10x
slow (4052us). `rocdl.sched_barrier(0)` between them does not stop it.
So the accumulator instead carries the running sum rescaled, which keeps a real
data dependency and needs no second accumulator:
t_0 = d_0; t_k = t_{k-1} * (s_{k-1}/s_k) + d_k; out = t_{K-1} * s_{K-1}
170 VGPR, 0 spill, against 160 for ptpc at the same tile. BLOCK_M defaults to
128 for blockscale, which is why the tile is 128x256 there.
Results at [16384, 7168] K=2048 on 8x MI355X, median of 31:
--quant ptpc --quant blockscale
fused-sdma 1114.5us 1146.2us
split-sdma 1261.7 1465.5
split-lsa 1262.4 1469.9
fused-lsa 1365.6 1699.6
gemm-only 228.9 388.5
The fusion margin over split goes 11.7% -> 21.8%, which is the point: the GEMM
grows 228.9 -> 388.5us while the 489us of link time does not, so there is more
compute to hide the transfer behind. Against the model's own 1491.3us per layer
that is 23.1% under. And our blockscale GEMM at 388.5 is within 4% of the
373.9us aiter's gemm_a8w8_blockscale_bpreshuffle measures standalone here, so
the operating point is now the model's.
Verification: a new test compares our kernel against the kernel the model
dispatches and against an fp32 reference (not bitwise -- different accumulation
order); 14/14 tests pass; fused-sdma, fused-lsa and split-sdma each validate
8/8 at the production shape.
One FlyDSL trap hit on the way, noted at the site: a `def` nested inside a
`@flyc.kernel` function gets its AST rewritten, and the rewrite of an `if`
inside it turned a captured object into a local, so reading it raised
UnboundLocalError and a `nonlocal` on another raised "no binding". Module-level
methods are left alone, so the rescale lives on _BlockScaleK.
jhchouuu
left a comment
jhchouuu
left a comment
There was a problem hiding this comment.
Reviewed commit 693fd5451280bc0a8664b67847fb981284e07839. Requesting changes for the reproduced public-API correctness issues: column-major A-scale handling, same-instance dynamic-M window reuse, and the unsupported K=128 pipeline being accepted. Additional inline comments cover API validation, lifecycle, usability, and clearly marked non-blocking maintainability suggestions.
Validation used an independent Podman container on 8x MI355X, ROCm 7.14.60850, PyTorch 2.12.0+rocm7.14.0, FlyDSL 0.2.4, and a wheel built with BUILD_CCO_SDMA=ON. Imports were checked outside the source checkout. Kernel P2P/VMM prerequisites were enabled. The existing test files passed: 48 passed, 1 skipped (the optional aiter comparison).
The main performance claim reproduced at M=16384/N=7168/K=2048, blockscale, median of 11 graph replays after 3 warmups, maximum over ranks:
| Mode | Time | relL2 |
|---|---|---|
| split-sdma | 1471.4 us | 2.35e-3 |
| fused-sdma | 1137.8 us | 2.35e-3 |
| fused-sdma + FP8/LSA | 979.8 us | 2.49e-2 |
That is a 22.7% latency reduction for BF16 fused versus split. The performance result does not establish the correctness of the public wrapper, since the benchmark calls the lower-level builders. Full SGLang/model-level evaluation and the NCCL baseline were not rerun. Host-only/stub probes are identified as such in the comments, separately from real GPU reproductions.
| a_fp8.contiguous().view(torch.int8).view(-1), | ||
| b_preshuffled.contiguous().view(torch.int8).view(-1), | ||
| c.view(-1), | ||
| a_scale.reshape(-1), |
There was a problem hiding this comment.
[P1] Preserve the documented column-major A-scale order.
For a logical [M, K/128] column-major scale tensor, reshape(-1) creates a copy in logical row order. The kernel instead indexes scales as kb * M + row, so this changes which scale belongs to each K block. The benchmark and single-GPU test avoid this by passing sa.t().reshape(-1).contiguous().
Reproduced through the installed public GemmAllReduceOp on 2 MI355X ranks with M=512, N=1024, K=512, using make_operands(..., quant="blockscale") from this PR's benchmark. sa.shape == (512, 4) and sa.stride() == (1, 512). Passing sa gives relL2=0.26081455 on both ranks; passing sa.t().reshape(-1) to the same op with the same operands gives 0.00244654.
Please preserve the physical order for the documented 2-D layout and explicitly distinguish it from an already flattened physical-order input. This needs a public-op test, since the existing kernel-level tests bypass this conversion.
| m=m, | ||
| n=self.n, | ||
| recv_slots=self.world_size, | ||
| counter_chunks=counter_chunks(m, self.world_size, self.block_m), |
There was a problem hiding this comment.
[P1] Keep control-region offsets stable when reusing the op for another M.
counter_chunks changes the counter-region size, which moves lock_off and input_off, but this instance reuses one window that was zeroed only during construction. On 8 ranks:
| M | chunks | lock_off | input_off |
|---|---|---|---|
| 4096 | 4 | 5632 | 5760 |
| 8192 | 8 | 5760 | 5888 |
Thus the second configuration interprets the previous GEMM payload as its submit locks. With one GemmAllReduceOp(comm, n=7168, k=512, m_max=8192), the sequence M=4096 -> M=8192 produced a correct first result (relL2=0.00235044) and then hung until a 90-second timeout (exit 124). A separate inspection read nonzero payload words such as 964770028 and 958839021 from the new lock region before launching the second call. Scales were flattened into the correct physical order to isolate this from the scale issue.
As a diagnostic control, synchronizing ranks and zeroing the window before each call made both calls complete with relL2 approximately 0.00235. Please allocate a fixed control layout for the capacity and define how counters are initialized/reused across shapes; the completion modulus also depends on the shape. The solution should remain valid when graphs for different M values are replayed alternately, rather than relying only on a Python-side last-M check.
| 2 <= world_size <= 8 | ||
| and n % block_n == 0 | ||
| and n % SCALE_BLOCK_K == 0 | ||
| and k % SCALE_BLOCK_K == 0 |
There was a problem hiding this comment.
[P2] Reject K=128 unless a single-K-block pipeline is implemented.
The divisibility check admits K=128, but the GEMM unconditionally prefetches block 1 and executes both K_ITERS - 2 and K_ITERS - 1 tail steps. Reproduced with supports(512, 1024, 128, 2) == True: construction and execution succeed, yet both ranks return finite output with relL2=1.41160023. The otherwise equivalent K=256 control gives 0.00244162.
Please require K >= 256 in both the support predicate and execution validation, or implement the one-block case. There is a related mismatch for world sizes 3/5/6/7: supports() returns True, while ArConfig.validate() rejects them because 512 threads is not divisible by that world size. These checks should share the actual execution constraints.
| return (p - 1) * self.nbytes | ||
| return 2 * (p - 1) * self.nbytes // p | ||
|
|
||
| def validate(self) -> None: |
There was a problem hiding this comment.
[P2] Validate mode strings and configuration constraints before allocating resources.
validate() does not check the dtype/transport strings: gather_dtype="FP8" is accepted but selects BF16, and gather_dtype="fp8", gather_transport="sla" is accepted but selects SDMA. A configuration typo silently changes the precision or performance path.
Other constraints are checked only after construction. A host-side probe with a recording FakeComm confirmed that sdma_queues=0, block_m=64, and n=512, gather_dtype="fp8" all reach allocation/window/dev-comm creation before _compiled() rejects them; block_m=64 yields an empty AssertionError. This probe tests Python validation order, not whether the real CCO API accepts invalid queues.
Please centralize legal mode values, positive sizes/queue counts, tile restrictions, and the FP8 N-alignment constraint before acquiring resources. In particular, the public op always enables permlane, which requires block_n=256. Literal annotations would help callers discover values, with runtime ValueErrors reporting the supplied and allowed values.
| gemm, parts, c, out = self._compiled(m) | ||
| stream = fx.Stream(torch.cuda.current_stream()) | ||
| gemm( | ||
| a_fp8.contiguous().view(torch.int8).view(-1), |
There was a problem hiding this comment.
API usability: validate tensor metadata before byte reinterpretation.
The public call currently checks M only. BF16 A/B tensors and a BF16 A-scale tensor can reach this launch boundary: A/B are reinterpreted as int8 bytes and the scale dtype is forwarded unchanged. I confirmed this with a kernel stub, without launching a GPU kernel on invalid inputs.
Please validate the cheap metadata here: operand dimensions and N/K, the supported FP8 encoding (or an explicitly documented raw-byte representation), FP32 scale dtype/layout, and the expected GPU/device. Also make the contiguous-input policy explicit. This would turn common integration mistakes into actionable Python errors before JIT/launch, rather than letting the kernel interpret incompatible storage. Type annotations on the public boundary can complement these checks.
|
|
||
|
|
||
| LSA_NOTE = """\ | ||
| The LSA transport is not implemented, and the reason is a measurement rather than |
There was a problem hiding this comment.
Non-blocking cleanup: remove stale exported notes and consolidate the experiment narrative.
LSA_NOTE still says the LSA transport is not implemented, although this file implements direct LSA and the benchmark exposes fused-lsa. It is also exported in __all__.
There are related stale descriptions elsewhere: kernels_sdma.py says fused quantize is the default while the signature uses False; the README switches between K=2048 and K=1024 and retains “Nothing here is a production kernel”; comments still describe a second blockscale accumulator although _BlockScaleK now rescales one running accumulator.
Please remove the obsolete exported string, keep current invariants and their reasons beside the code, and maintain measurements/history in one document. This should make the synchronization decisions easier to review without losing the useful negative results.
| quant: str = "ptpc", | ||
| sdma_queues: int = 8, | ||
| swap_ab: bool = False, | ||
| store_probe: bool = False, |
There was a problem hiding this comment.
Non-blocking simplification: move deliberately incorrect probes out of the shipped kernel path.
This entry point has 22 keyword-only arguments and 11 fence modes, including _AdjacentLaneProbeC/_WideStoreProbeC paths whose output is wrong by construction and raw-wt* experiments documented as racy. The public op uses a much smaller, fixed subset of these choices.
Please consider moving the incorrect probes and historical tuning variants under benchmark/experimental, leaving the supported kernel path with a small set of validated choices. Negative results and their reproduction commits are useful to retain in documentation. The completion counters, submit locking, and barrier invariants themselves have a clear purpose and should not be removed merely to reduce line count.
| return first_pid_m + intra_group_m, pid_n | ||
|
|
||
|
|
||
| def compile_fp8_gemm_8w( |
There was a problem hiding this comment.
Non-blocking simplification: vendor the components actually used, or make this full pipeline an explicit benchmark reference.
The module docstring says kernels_fused.py reuses the loaders/MFMA/store helpers but does not call this pipeline. Repository search found no caller of compile_fp8_gemm_8w in this PR, yet the remainder of this file carries another full GEMM compile/launch implementation and documents that its upstream under-wait remains unfixed.
Keeping roughly 250 unused lines of a second pipeline in the installed package adds another implementation for maintainers to reason about. Please retain the required components with upstream revision/license attribution; if the original full pipeline is needed as a reference, give it an explicit benchmark location and caller.
| #: sums that are then added across every rank, so its error compounds. They | ||
| #: also quantise in different places, hence the two scale granularities | ||
| #: below. ``ArConfig(comm_dtype=...)`` sets both at once. | ||
| scatter_dtype: str = "bf16" |
There was a problem hiding this comment.
Non-blocking simplification: defer the unimplemented FP8-scatter layout surface.
The scatter dtype adds payload/scale/recv layout branches and dedicated tests, but build_sdma_phases unconditionally raises NotImplementedError for FP8 scatter. It also means the public comm_dtype="fp8" shorthand selects a configuration the SDMA implementation cannot execute.
The working FP8-gather feature does not require this unimplemented scatter surface. Consider landing the scatter layout/options together with its implementation and numerical tests, keeping this PR focused on BF16 scatter plus the supported gather choices.
|
|
||
|
|
||
| @pytest.mark.parametrize("mode", ["split-sdma", "fused-sdma", "split-lsa"]) | ||
| def test_gemm_ar_modes_agree(mode): |
There was a problem hiding this comment.
Add public-op and changing-input coverage alongside the kernel benchmarks.
The current tests construct kernels/phases directly; the CI check for GemmAllReduceOp only imports the class. I reproduced 48 passed, 1 skipped across the three PR test files while the public op still had the column-major scale conversion error and the 4096 -> 8192 window-reuse hang described in the inline comments.
Please exercise GemmAllReduceOp itself with the documented scale layout, a same-instance 4096 -> 8192 -> 4096 sequence, repeated changing inputs, and graph replay after warmup. Explicitly test both FP8 gather transports: the existing standalone FP8 test uses SDMA, whereas the public op defaults to LSA. This would cover caller-visible behavior and state transitions that additional layout-getter assertions or constant-input benchmark repetitions cannot check.
046391f threaded a `geom` tuple through every helper in kernels_sdma -- 21 call sites and three changed signatures -- to hoist the two descriptor loads out of `lsa_ptr`. Same optimisation, wrong seam. `CachedWindow` reads the geometry in its constructor and computes `lsa_ptr` from it, so the kernels change by one line each: w = cco.Window(win) -> w = cco.CachedWindow(win) and nothing else. kernels_sdma.py goes back to the pre-046391f6 text; the 14 lines that remain are the 7 swaps. `Window` keeps `lsa_ptr` unchanged, and `lsa_geometry`/`lsa_ptr_at` are gone -- the class is the API now. The C++ side from 046391f stands: `cco_lsa_win_base` / `cco_lsa_stride` take an `address_space(1)` pointer, so each is one `global_load` rather than the `flat_load` that `cco_lsa_ptr`'s generic cast forces, and a `flat_load` counts against `lgkmcnt` as well as `vmcnt`. It still does not measure. A/B/A' in one session, m=16384 n=7168 k=2048, 8x MI355X, max_rank_time_us: wire Window CachedWindow Window again bf16/sdma 1114.5 1111.4 1115.4 fp8/sdma 983.3 983.6 983.6 fp8/lsa 926.6 924.3 926.0 B-A is smaller than A'-A. This is the fourth independent null for the idea, after the band-loop hoist, the Python geometry API and the global_load conversion. Note for anyone re-measuring: this box drifts ~4% between mornings, across every configuration at once. Today's absolute numbers are all ~4% under 046391f's on identical code, which is why the A/B is the only line worth reading. 44 passed, 1 skipped. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…eal win
Four more sites: both kernels in kernels_lsa.py, and kernels_fused.py's w_pre and
its one inline lsa_ptr. kernels_fused.py:1488 stays on Window, see below.
split-lsa is the first configuration where this measures. A/B, three alternating
repeats, m=16384 n=7168 k=2048, 8x MI355X, max_rank_time_us:
mode Window CachedWindow
split-lsa 1264.69 1264.23 1264.85 1250.25 1256.85 1254.53 -10.7us
fused-sdma 1110.45 1109.77 1112.85 1110.53 1113.69 1112.61 null
-0.85% against a Window spread of 0.6us across its three runs.
The reason is the shape of the kernel, not the number of call sites.
ar_1stage/ar_2stage build nine peer addresses in *every block* of a short kernel,
so two flat_loads apiece is a real fraction of it. Every other site converted so
far -- 21 in kernels_sdma, both here in kernels_fused -- runs once per launch
against a body lasting a millisecond, and all of them measured nothing. Count
address constructions per launch, not `grep -c lsa_ptr`.
kernels_fused.py:1488 cannot be converted. That `w` is read inside a dynamic
`if`, and FlyDSL captures every variable an if body reads as state, requiring
each to be a single MLIR value:
TypeError: state variable 'w' is CachedWindow, not an MLIR Value
Making CachedWindow a cco_struct does not help -- three fields extract to three
values:
TypeError: state variable 'w' is list, not an MLIR Value
Window passes only because it has exactly one field. Reverted that experiment;
CachedWindow stays a plain object, and the limitation is now documented with the
error text so nobody retries it.
44 passed, 1 skipped.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
kernel_gemm_scatter built the window three times: w_pre at the top, then a throwaway `cco.Window(win)` inline for c_rsrc, and a third one for the scatter epilogue. All three are in the same function with no nested def between them, so w_pre was in scope for all of them -- the rebuilds were left over, not needed. The epilogue one also needed the two addresses moved out of the dynamic `if`, which is what let it use w_pre at all: FlyDSL captures every variable an if body reads as state and requires single MLIR values, so a CachedWindow read inside fails. Both bases are rank-local with constant offsets, identical across threads, so hoisting them is free and leaves only plain Int64 in the captured state. That retires the "kernels_fused.py:1488 stays on Window" exception from 536eba6 -- the file now builds the window exactly once. No timing change, which is what the per-launch rule predicts: these addresses are built once per launch, not per block. fused-sdma bf16 1111.3us, fp8/lsa 930.0, fused-lsa 1360.3, all inside the spreads in 536eba6. relL2 unchanged at 2.35e-3 bf16 / 2.49e-2 fp8, validated on all three. 44 passed, 1 skipped. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
… changed d5e26fa removed the kernels_fused.py:1488 Window the docstring cited as the worked example. Hoisting the addresses out of the branch is the general answer anyway, and it does not rot when a line number moves. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
81fe2fe converted kernels_sdma's 21 lsa_ptr sites and I measured three wire configurations, all green. --fuse-reduce-push is off by default, so neither the test suite nor any sweep ever traced sdma_reduce_push, and it had been dead since: TypeError: state variable 'w' is CachedWindow, not an MLIR Value Same cause and same fix as kernels_fused: the two bases are rank-local at a constant offset, so they are invariant across bands and identical on every thread. Hoisting them above the band loop keeps the window out of the if's captured state and computes them once instead of once per band. Smoke-tested the other off-default flags for the same fault -- publish=fence, fuse-quantize both ways, direct-fence=all, gather-bands=4, sdma-queues=1 -- all validate. Answering the original question: fusing the push into the reduce still does not pay. M=16384, three alternating repeats: off 1111.4 1111.5 1112.9 ON 1127.2 1129.4 1128.8 +16.4us, +1.5% Worse than the +8.5us recorded before CachedWindow, and the spreads are 1.5 and 2.2us. Staying off by default. Also: --fuse-quantize defaulted to True while build_sdma_phases defaults it to False, and its help claimed "same numbers either way". It costs 23.4us on fp8/lsa (926.6 vs 903.2, three repeats each) and 18.5us on fp8/sdma. Flipped the bench default to match the library, so the bench stops reporting the op ~23us slower than it runs. 44 passed, 1 skipped. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Adds the lsa_ptr / CachedWindow line to the negative results: four ways of hoisting the window geometry, all null on kernels_sdma and kernels_fused, one win of -10.7us on split-lsa, and the rule that separates them -- address constructions per launch, not call sites. Also re-measures the two existing negative results with repeats: fuse_reduce_push +16.4us, fuse_quantize +23.4us. And retracts something I got wrong and committed. I read today's numbers against this page and found everything ~4% faster, and wrote that up -- in 81fe2fe's message and in CachedWindow's docstring -- as the box drifting between mornings. It does not. The benchmark defaults to --quant ptpc; every number on this page is --quant blockscale. Measured back to back: mode doc today blockscale today ptpc gemm-only 369.4 367.9 221.0 split-sdma 1472.9 1462.8 1259.0 fused bf16 1148.8 1144.9 1112.1 The page reproduces within ~1%. There is no drift; there was a flag I did not pin. The A/B pairs are unaffected -- both arms always ran the same flags -- but the explanation attached to them was wrong, so it is corrected at both sites. Also confirms the PR comparison in blockscale, the mode the headline is quoted in: split-lsa 1463.6 -> 1455.6us, everything else flat. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…hat layer
The op layer had no tests at all -- test_gemm_ar.py goes through bench_gemm_ar.py,
which calls the low-level builders -- so three correctness bugs sat behind a
green suite. Adds tests/python/cco/test_gemm_ar_op.py (27 cases) and a worker
that drives GemmAllReduceOp itself.
P1, all reproduced before fixing:
* The A scale was flattened with reshape(-1). The kernel reads (row, kb) at
kb*M + row, so for the column-major [M, K/128] tensor the quantiser returns
that pairs every scale with the wrong K block: relL2 0.26 instead of 2.4e-3.
_flatten_a_scale now distinguishes [M, K/128], [K/128, M] and an
already-flat physical-order buffer, and the three agree bitwise.
* counter_chunks changes with M, which moved lock_off and input_off, so a second
M read the previous shape's payload as control state. The counter region is
now sized by capacity (MAX_CHUNKS) and carries one counter set per shape --
the election test is `seq % tiles_per_chunk == 0` on a counter that is never
reset, and tiles_per_chunk depends on M, so sharing a set means electing on
the other shape's residue.
* supports() admitted K=128, which returns finite garbage (relL2 1.41): the
mainloop prefetches block 1 and runs two tail steps, so K_ITERS must be >= 2.
It also disagreed with ArConfig.validate() on world sizes 3/5/6/7. It now
builds the config and validates it rather than keeping a second copy of the
rules.
P2:
* ArConfig.validate() rejects misspelled mode strings ("FP8" selected bf16,
"sla" selected SDMA), non-positive sizes, and the fp8 N multiple -- before
anything is allocated.
* __call__ validates operand dtype/shape/device. A/B go in as int8, so a bf16
tensor used to reach the kernel and come back finite and wrong.
* close() / context manager. It releases the window and memory, not the
dev-comm: destroying that and then letting the communicator tear down aborts
the process. The window is the point -- ~700 MiB at the model shape.
* pad_rows keys its buffer on (device, dtype, K), not just capacity, and checks
m_pad >= m. A bf16 buffer reused for an fp32 input silently rounded it.
* RESULT_JSON records the effective wire, transport, quantize fusion, queues and
resolved tile order. An fp8/lsa and an fp8/sdma run both said only
mode="fused-sdma".
* The cached plan is a NamedTuple with the phase tail already resolved.
Also: removed LSA_NOTE, which said the LSA transport is not implemented in the
file that implements it; corrected the --fence help, which named 'agent' as the
default when the default is 'none' and called 'none' incorrect. Measured, twice
each at m=16384 blockscale on 8 ranks: none 1143.5/1144.8us, writethrough
1148.5/1145.4, agent 1797.0/1798.5, all 1936.1/1937.5 -- all four at relL2
2.350e-3, so the explicit fences cost 650-790us and did not move accuracy.
README gains a runnable end-to-end example; op.padded_m(m, world_size) in the
old one raised TypeError.
One defect found by the new tests and NOT fixed: after a call at some M, the
first call at a *larger* M is wrong about half the time, always with the same
value, with the repeat correct. 8 runs each on 2 ranks: 512->1024 fails 4/8,
1024->512 fails 0/8, and 2048->4096 fails too, so it is the growth and not the
chunk count. Five calls at one M with different operands each time are all
correct, which rules out a stale read at a fixed shape. The op now raises
instead of returning the wrong answer, and the test pins that.
73 passed, 1 skipped.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…ch other Root-causes and fixes the defect e3514e2 could only guard against: after a call at some M, the first call at a larger M came back wrong about half the time. The row map was the same on every failing run -- at world=2 n=1024, m=1024 had rows 512..767 wrong, contiguous, and nothing else. Those rows are bytes 3152128..3676416, which is exactly where m=512's recv slot for peer 1 sat, and their contents matched peer 1's *unreduced* GEMM output at m=512 to 9.5e-07, i.e. that peer's scatter payload from the previous call, still in place. Every offset past the control area is a running sum of m-sized regions, so two shapes sharing a window put different regions at the same bytes. The overlaps at that shape: m=1024 input x m=512 output 1048576 B m=1024 output x m=512 recv0/1 524288 B each ArConfig.capacity_m lays the payload out for the largest m an instance will serve, while sizes stay m-based -- the map moves, not the bytes moved. The op passes m_max for every shape, so a region only ever aliases itself. The 15 offset properties now read their sizes from `self._cap`, a copy with m replaced; with capacity_m unset it is `self`, so nothing else changes and the benchmark, which builds its own configs, is untouched. Verified: 512 -> 1024 is 8/8 correct where it was 4/8 wrong; 2048 -> 4096 6/6. The growth guard and its test are gone -- growing M is supported now, and test_one_instance_serves_two_m_values_in_any_order covers both directions. A host-only test pins the map being M-independent, and asserts it is *not* without capacity_m, which is the property that broke. No performance change, as expected since the map is identical when m == m_max: blockscale at the model shape, split-sdma 1460.0us, fused bf16 1149.4, fp8/lsa 949.0, split-lsa 1453.4, all within noise of the same run before. 74 passed, 1 skipped. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
End-to-end regression after the op-layer fixes, same protocol as the page
describes, two or three runs each. GPU busy:
unfused 1095.5ms (1100.2 1095.1 1091.2)
fused, bf16 1021.0 (1018.7 1025.1 1019.2)
fused, fp8/sdma 1026.5 (1025.9 1027.1)
fused, fp8/lsa 1049.3 (1048.3 1050.2)
This page's layer table has fp8/lsa at 949us against bf16's 1149 at the same
shape, so fp8/lsa should have been the best of the three fused rows and is the
worst. The likely reason is rank skew -- a pull needs the peer's slice finished
at the moment it reads, a push lets each producer send as soon as its own band is
done -- and only the benchmark has every rank in lockstep on an idle box. SGLang
now defaults the transport to sdma; GemmAllReduceOp keeps lsa, which is right for
the benchmark it was measured in.
Fusing itself is worth more than the old capture said: -6.8% against -2.2%. The
unfused row reproduces, the bf16 fused row does not and is 56ms better. That
measurement predates the move to mori.ops.gemm_ar and several fixes here; the
cause was not isolated, and the page says so rather than picking one.
Pointers added at the headline and at the transport table, because a reader who
stops at either would take away the opposite conclusion.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…e doc's e2e claims
Two corrections, one of them to a bug I introduced yesterday.
**_flatten_a_scale must not infer the layout from the shape.** It treated a 2-D
[M, K/128] tensor as logical and transposed it. That is right for the
column-major one the review reported, and wrong for the one
aiter_per1x128_quant(transpose_scale=True) returns, which writes K/128-major
data and keeps the (M, K/128) shape -- shape and stride are identical in both
cases and only the producer knows which. Guessing took SGLang from relL2 3.9e-3
to 0.36. It now accepts 1-D, [K/128, M], and [M, K/128] column-major, and
rejects [M, K/128] row-major with a message naming both readings; SGLang passes
.reshape(-1) explicitly.
Verified on the server's own dumped operands at [11264, 7168] K=2048 on 8 ranks:
flat 3.868e-3, ten repeats identical; the transposed spelling 0.365.
**The end-to-end section is restored.** Yesterday's rewrite -- "fusing is worth
-6.8%", "the fp8 gather does not pay", "the LSA pull costs 28ms" -- was measured
against a mori built with BUILD_CCO_SDMA=OFF, where every put silently does
nothing. The all-reduce then returns mostly the local slice, the model still
answers fluently, every mori kernel still shows up in the profile, and the fused
path measures faster than it is because it moves no data. fp8/lsa looked worst
of the three for the same reason: the pull is the one leg that does not use
SDMA, so it was the only one actually working.
Re-measured on a build that has it:
unfused 1096.0ms wall 1.1969s
fused, bf16 1070.3 1.1728 -2.3%
fused, fp8/sdma 1052.1 1.1455 -4.0%
fused, fp8/lsa 1041.2 1.1385 -5.0%
which is the page's original 1101.9 / 1077.2 / 1050.4 / 1048.2 to about half a
percent, in the same order. The page now leads with how to notice the failure:
perplexity, 862511 against 3.26, because generated text and the kernel trace
both look fine, and a short prompt never enters the path at all (fusing needs
M >= 4096).
75 passed, 1 skipped.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
… op tests cco_lsa_win_base and cco_lsa_stride went into CCO_DEVICE_FUNCTIONS without the two tests that pin that surface: the count tripwire (68 -> 70) and the flydsl-vs-triton ABI comparison, whose extern list had to name them. Both are doing their job -- the surface changed and they said so; I had only run the gemm_ar files locally. Also wires tests/python/cco/test_gemm_ar_op.py into the numerics step. It was added in e3514e2 and never run by CI, which is the same gap it exists to close: the other two files drive the benchmark, which calls the low-level builders, so nothing covered the public op. Raises that step's timeout to 1800s -- it spawns several 2-rank workers with an 20s settle between them. Ran the CCO steps that work here: triton bindings/lsa/sdma 11 passed, rocm import order 2 passed, sdma api 12 passed. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
|
I think there remain two things to try, first is to tune the sdma queue number to max the xgmi bw (like the multi-qp mode in ibgda); second is the reduce kernel impl way, investigate the one-shot all-reduce or two-shot all-reduce which will be better. Anyway, that is a nice work! @yangyuhuiling |
|
And for local reduce + gather path, is possible to combine them to one kernel, use direct vector store rather than using sdma? |
max_shapes defaulted to 8. Each distinct M needs its own counter set -- the
election modulus depends on M, so two shapes sharing a set elect on each other's
residue -- and the ninth M raised. A server's M changes with the batch, so under
real traffic it raised within seconds:
mori fused wo_b failed and is disabled for this process; falling back to the
split path: this op already serves 8 distinct M values [4096, 5120, 6144, ...
and SGLang turns any exception from this path into a permanent fallback, so the
whole optimisation switched itself off. Two thirds of a 128k-prefix benchmark
matrix measured the unfused path without saying so.
The default is now every M the window can legally take: a padded M is a multiple
of world_size*block_m and at most m_max, so m_max/(world_size*block_m) slots can
never be exhausted by a legal call -- 16 at the model shape, covering the 13
values from 4096 to 16384. A slot is world_size*MAX_CHUNKS*4 bytes, 256 B at
tp8, so 16 of them cost 4 KiB of a 700 MiB window. Rationing them bought nothing
and cost the feature.
75 passed, 1 skipped.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…measures A mori built without BUILD_CCO_SDMA=ON -- the default, since the flag only flips with BUILD_BENCHMARK -- compiles the SDMA puts out. Every symbol is still there, every kernel still launches, every put still returns, and nothing moves. The all-reduce then yields mostly the local slice: the model keeps answering fluently, the profile keeps showing all the kernels, and the fused path measures *faster* than it is because it is not moving data. An end-to-end campaign read -6.8% instead of -2.3% that way, and the only thing that caught it was perplexity, 862511 against 3.26. GemmAllReduceOp.self_test() writes rank+1 over the input region, runs the split order (scatter -> reduce -> gather; fused_order opens with drain, which assumes a GEMM epilogue already pushed), and checks every element came back as sum(1..world_size). With the puts inert each rank sees only its own share and the check fails. It runs at m_max, compiling the shape the first real call needs anyway, and reads mori.cco.device._build_flags.BUILD_CCO_SDMA only to make the error message name the cause. The test asserts both halves, because a check that cannot fail is what let this happen. Zeroing the landing slots before the negative half is not a detail: the op has served other cases by then, so the slots still hold a previous call's peer data, reducing that produces a plausible sum, and the teeth test passes while testing nothing -- the same stale-read effect that made the dead transport look correct. Also in this commit, because the tests needed it: * All six GPU cases now come from one 2-rank spawn behind a module fixture. Six back-to-back launches outran the SDMA queue reclaim (hsaKmtCreateQueueExt, anvil.cpp:237) and the casualties were as often the neighbouring file's tests as this file's, since the cost falls on whoever launches next. The file also drops from 305s to 177s. * A spawn that fails with that signature is retried once, with a longer settle. Nothing else is retried. * counter_chunks() raises with the shape it was given instead of "max() iterable argument is empty" when m_pad is below one row band. 75 passed, 1 skipped. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
jhchouuu
left a comment
jhchouuu
left a comment
There was a problem hiding this comment.
Re-reviewed 1531cfe8ddd23abed2897182b7cc73ccf73b11c3. The four inline issues below are the remaining merge blockers from this review. The earlier style, vendoring-attribution, and module-organization suggestions are non-blocking.
The original column-major scale, dynamic-M window, and K=128 findings have targeted fixes, and those regression cases now pass. Validation on 8x MI355X with ROCm 7.14.60850, PyTorch 2.12.0+rocm7.14.0, FlyDSL 0.2.4, and BUILD_CCO_SDMA=ON:
- Existing benchmark + numerical/public-op/ABI tests: 88 passed, 1 skipped (optional aiter comparison).
- Additional 8-rank changing-input tests alternating captured graphs for M=4096/8192/16384 passed for BF16, FP8/LSA, and FP8/SDMA.
- After waiting for another CI workload to exit, the isolated model-shape performance test passed: 1130.0 us fused vs 1466.9 us split, approximately 23% lower latency. No external GPU processes were observed in that measurement interval.
- The dependency audit found no aiter runtime dependency or external prebuilt-GEMM import in this path; a GEMM numerical test also passed with those imports blocked.
Performance and dependency structure are therefore not the reasons for this request-changes review. The accepted FNUZ input produces incorrect values, and the remaining exception/validation/quickstart paths need the concrete fixes described inline. Host fault-injection results are identified separately from GPU numerical reproductions.
| #: fp8 encodings the kernel's byte reinterpretation is valid for. Both are 1 | ||
| #: byte and the GEMM only ever sees the bytes, but forwarding a bf16 tensor here | ||
| #: reinterprets two rows as one and produces finite nonsense. | ||
| FP8_DTYPES = (torch.float8_e4m3fnuz, torch.float8_e4m3fn) |
There was a problem hiding this comment.
[P1] Reject FNUZ inputs until the kernel supports their encoding.
These two dtypes are both one byte, but their encodings are not interchangeable. The kernel's MMA atom is fixed to fx.Float8E4M3FN (_gemm_a8w8_8wave.py:286), while __call__ only reinterprets the input as int8 bytes. Admitting float8_e4m3fnuz here therefore produces incorrect results for an input the API explicitly accepts.
Reproduced on both MI355X machines with an initialized 2-rank communicator, M=512, N=1024, K=512, A/B filled with numeric ones, FP32 A/B scales filled with ones, and B passed through preshuffle_b:
| Input dtype | Byte encoding of 1 | Expected output | Actual output on both ranks |
|---|---|---|---|
float8_e4m3fn |
0x38 |
1024 | 1024 |
float8_e4m3fnuz |
0x40 |
1024 | 4096 |
Every output element is wrong by the same factor of four; relL2 is approximately 3. op.self_test() passed before this comparison, so the collective was moving data correctly.
Please restrict this gfx950 path to the encoding it actually implements, or add an explicit FNUZ conversion/kernel path. Add a numerical test for each accepted dtype; checking element size is insufficient.
| # reference to each handle, so a half-built op leaks the window for | ||
| # as long as the communicator lives. | ||
| self._closed = False | ||
| self.close() |
There was a problem hiding this comment.
[P2] Initialize cleanup state before entering the resource-acquisition try block.
This rollback calls close(), which immediately executes self._cache.clear(), but _cache is not initialized until line 381, after the entire try block succeeds. Any allocation/window/dev-comm failure therefore reaches cleanup with _cache absent, masking the original exception and preventing the handle-release loop from running.
A host fault-injection probe let alloc_mem and register_window succeed and made create_dev_comm raise RuntimeError. The observed result was AttributeError: 'GemmAllReduceOp' object has no attribute '_cache', with 2 resources acquired and 0 closed. The injected error survived only in __context__. These were recording fake resources, not a deliberately induced hardware OOM.
Please initialize _cache, _pad_in, and all other cleanup state before resource acquisition, and test failures at each acquisition step. Rollback should release the handles acquired by this construction and propagate the original error. This does not require destroying communicator-scoped shared SDMA queues.
| max_shapes = default_max_shapes( | ||
| padded_m(m_max, comm.nranks, block_m), comm.nranks, block_m | ||
| ) | ||
| if block_m < 1 or m_max < 1 or sdma_queues < 1 or max_shapes < 1: |
There was a problem hiding this comment.
[P2] Share the actual tile constraints between supports(), construction, and execution.
The K=128/world-size cases are fixed, but block_m is only checked for positivity here, whereas the kernel requires BLOCK_M >= 128 and BLOCK_M % 128 == 0.
Reproduced at M=512, N=1024, K=512, world_size=2:
supports(..., block_m=64)returnsTrue.- Construction with
block_m=64reaches all three resource acquisitions;_compiled(512)then raises an emptyAssertionError. supports(..., block_m=0)raisesZeroDivisionError; construction also calculates padding/default_max_shapes before reaching this positive-value check.
The constructor/preparation ordering was checked with recording fake resources, without launching an invalid tile on a GPU. Please validate tile parameters before any padding arithmetic or allocation and use the same execution constraints in supports(). Unsupported shapes should be rejected at that boundary instead of failing during the first real call.
| uid = [bytes(Communicator.get_unique_id()) if rank == 0 else None] | ||
| dist.broadcast_object_list(uid, src=0) | ||
|
|
||
| with Communicator.init(world, rank, uid[0], per_rank_vmm=vmm) as comm: |
There was a problem hiding this comment.
[P2] Reconstruct UniqueId before passing the broadcast bytes to Communicator.init.
uid[0] is a bytes object, but Communicator.init accepts the Python UniqueId wrapper or its underlying Cython object and does not convert bytes automatically. Executing this call raises:
TypeError: Argument 'uid' has incorrect type (expected mori.cco.cco.UniqueId, got bytes)
Please import UniqueId and call Communicator.init(world, rank, UniqueId.from_bytes(uid[0]), per_rank_vmm=vmm), as the benchmark and new worker already do.
Please also bring the adjacent contract in line with the final code: _flatten_a_scale now rejects logical row-major [M,K/128] rather than accepting either memory order; the default shape capacity is no longer 8; and close() deliberately retains the dev-comm until communicator teardown. The example's returned view also needs to be consumed or cloned, with the necessary synchronization, before the op context frees its backing window.
…ollback Four findings from the second review, all reproduced. **FNUZ is not the encoding this kernel implements.** FP8_DTYPES admitted both e4m3 variants on the theory that the launch only sees bytes and both are one byte wide. The width is the same; the exponent bias is not -- 8 against OCP's 7 -- and the MMA atom is fixed to fx.Float8E4M3FN. Every A and every B element therefore reads at twice its value and the product comes back 4x too large: finite, plausible, and wrong. The reviewer measured exactly that, 4096 where 1024 was expected, relL2 ~3. sglang reaches this on gfx942, where normalize_e4m3fn_to_e4m3fnuz converts weights to FNUZ, so the accepting-and-corrupting version was reachable from a real caller. Accept only float8_e4m3fn. The comment justifying the old tuple was the actual defect -- it argued from element size, which is evidence of nothing. **The rollback could not run.** __init__'s except block calls close(), whose first statement is self._cache.clear(), but _cache was assigned after the try block. Any acquisition failure died in cleanup with AttributeError, demoting the real error to __context__ and skipping the handle-release loop the rollback exists for. Fault injection at create_dev_comm showed 2 resources acquired and 0 closed. Initialise all cleanup state before the first thing that can raise. **Tile constraints were not shared.** The kernel asserts BLOCK_M >= 128 and BLOCK_M % 128 == 0 (bare, so an empty AssertionError), while the op only checked block_m >= 1. block_m=64 passed supports(), constructed, acquired all three resources, and failed on the first real call. block_m=0 was worse: default_max_shapes and padded_m both divide by world_size * block_m, and they ran before the positivity check, so it raised ZeroDivisionError from inside the padding arithmetic. _tile_constraints() now states the granules once and runs first in both supports() and __init__, before anything divides. supports() also answers False for world_size < 1 for the same reason. **README.** Communicator.init takes the UniqueId wrapper, not the bytes that survive a broadcast, so the quickstart raised TypeError as written. Three other places had drifted from the code: the a_scale contract (row- major [M,K/128] is rejected now, and the reason is that the two 2-D spellings are indistinguishable by shape yet need opposite treatment), the max_shapes default (every legal M, not 8), and close() (the dev-comm is deliberately left to the communicator). The example now clones the result before the context frees the window behind it. Tests: the tile boundary at 64/1/192/0/-128 and block_n 128/384/0, that 256x512 is still accepted, world_size 0 and -1, that FNUZ is not in FP8_DTYPES, and a fault injection at each of the three acquisition steps asserting the original error propagates and everything acquired before it was released. The injection stubs from_gpu_ptr: the fake handles carry a null pointer and the window-zeroing write faults the context otherwise, which then surfaces in whatever runs next. 44 pass in test_gemm_ar_op.py, including the 2-rank GPU cases. test_flydsl_ar.py::test_flydsl_ar_matches_aiter_bitwise fails here, but it fails identically on the unmodified tree. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
|
@kawhil-amd Thanks — both of these have measurements behind them already, so let me put the data here rather than just an opinion. 1. Tuning the SDMA queue countThe multi-QP analogy does not carry over, because the hardware queues are already per The lever that does move xGMI bandwidth here is chunk size, and it is tuned. Per-kernel at the model shape (M=16384, N=7168, K=2048, 8 ranks):
The drain falls 38% while the GEMM grows 7.5 µs for the extra counter atomics and the lock. Going finer loses: at chunks=8 each PUT is 3.5 MiB; 16 halves it to 1.75 MiB, under the knee in the SDMA bandwidth curve, for 1130.8 µs. Bandwidth on this path is a function of bytes per op and essentially nothing else, so the tuning is "do not split the transfer more than you need for overlap" rather than "add queues". 2. Fusing local reduce + gather into one kernelTwo related things have been measured, and your suggestion is a third that neither quite covers. Being precise about which is which: (a) Fusing them with SDMA puts fired from the reduce — commit The cause is not the traversal or the puts. Handing a range to a copy engine needs a system-scope release from every block that wrote it: 256 blocks × bands fences at ~0.24 µs. The whole reduce is only 42 µs, so the overlap cannot repay it. Isolated at 8 bands: fence without the vmcnt drain 1649.8 µs, drain without the fence 1164.5 µs — the drain is free, the fence is the entire cost. Dropping the fence is not a fix: 1157.9 µs at 4 bands (still behind) and racy, relL2 1.8e-2 at 32 bands against a 2.35e-3 floor, differing per rank. (b) Direct vector store instead of SDMA, on the scatter leg — that is (c) Your exact suggestion — one kernel that reduces and then vector-stores the result to peers — is not the same as either, and I want to be clear that I have not run it. It sidesteps (a)'s objection entirely: with no external engine there is no range to publish, so the release fence disappears. What it inherits is (b)'s ratio, 730/499 ≈ 1.46× for CU stores against copy engines on the push direction. Applied to the 496.6 µs gather that is ~725 µs, against 540 µs for reduce+gather today — so the expectation is that it loses by roughly that margin, but that is an extrapolation across two different legs, not a measurement. Worth noting the pull direction already does what you are describing and wins: If you think (c) is worth settling rather than extrapolating, I am happy to build it — the reduce already has the values in registers at the right moment, so it is a contained change. Say the word and I will post numbers. The transferable result from all of this, for whatever the next fusion candidate is: hiding a transfer behind compute pays only when what is hidden is much larger than the cost of publishing it. That held for the GEMM's fused scatter (437 µs of transfer behind a compute-bound GEMM) and failed for the reduce (42 µs, bandwidth-saturated). |
|
@jhchouuu All four are fixed in [P1] FNUZ. Rejected — The comment justifying the old tuple was the real defect. It argued from element size, which is evidence of nothing, and I have replaced it with the encoding argument plus a note that widening this needs a kernel path and a numerical test per dtype. Worth recording that this was reachable from a real caller, not just from the API surface: sglang runs [P2] Cleanup state. [P2] Tile constraints. New [P2] README. Verification. 44 pass in |
pre-commit's only complaint on f4caced. The repo pins black 25.1.0, which still wraps the condition rather than the message on a multi-line assert. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
jhchouuu
left a comment
jhchouuu
left a comment
There was a problem hiding this comment.
Approving after re-reviewing dc8ba7fbb6dd4f8d2dea523de76237c49a51e533 and rerunning the experiments. The FNUZ rejection and acquisition-failure rollback fixes are verified, and the original UID conversion issue is corrected.
The three inline points are non-blocking follow-ups for your assessment, with 1 (supports vs construction) and 2 (README scale flattening) especially worth another look. Their concrete observations are still relevant; this approval should not be read as confirmation that those details are already resolved.
Validation on MI355X with ROCm 7.14.60850, PyTorch 2.12.0+rocm7.14.0, FlyDSL 0.2.4, and BUILD_CCO_SDMA=ON:
- Existing public-op, numerical, ABI, and benchmark tests: 101 passed, 1 skipped.
- FNUZ was rejected independently for A and B; the supported FN path computed the expected result.
- Failure injection at allocation, window registration, and dev-comm creation propagated the original errors and released the resources already acquired.
- At M=16384/N=7168/K=2048, blockscale: 1130.5 us fused vs 1462.7 us split, approximately 22.7% lower latency, with both numerical checks passing.
The separate presence-gate probe for follow-up 3 intentionally exposed an assertion failure and is not included in the passing-test count. Thanks for addressing the earlier correctness issues.
| ) | ||
| except ValueError: | ||
| return False | ||
| return True |
There was a problem hiding this comment.
Follow-up 1 — non-blocking, but please prioritize checking the support contract.
Could you recheck the accepted 256x512 case? At M=4096, N=7168, K=2048, world_size=8, supports(..., block_m=256, block_n=512) returns True, while constructing GemmAllReduceOp with the same configuration still raises ValueError: block_n must be 256 because the op always enables permlane. I reproduced this through the constructor; it rejects before allocating resources.
The new test_supports_accepts_the_legal_tile_multiples checks only the predicate, so it currently enshrines that disagreement. If supports() describes the public op, it should apply the same restriction; if it intentionally describes a broader kernel family, that scope should be explicit, since its current docstring promises agreement with construction. Please consider validating both sides in the test.
| one must be read flat, the row-major one must be transposed. Guessing gets one | ||
| of them silently wrong (relL2 0.26 either way). Pass `.reshape(-1)` or `.t()` | ||
| yourself to say which you have. `aiter_per1x128_quant(transpose_scale=True)` | ||
| returns the column-major one, so `.reshape(-1)` is right for it. |
There was a problem hiding this comment.
Follow-up 2 — non-blocking, but please prioritize rechecking this scale guidance.
For a real PyTorch column-major [M, K/128] tensor with stride (1, M), reshape(-1) walks logical order and can reorder the physical scale buffer. In a real 2-rank run at M=512, N=1024, K=512, the same operands gave:
| A scale passed to the op | relL2 |
|---|---|
column-major sa directly |
0.00244654 |
sa.reshape(-1) as suggested here |
0.26081455 |
sa.t().reshape(-1) |
0.00244654 |
The wrapper handles the direct column-major case correctly now. Please consider changing this sentence so readers do not recreate the original ordering error before calling it. If a particular quantizer writes kb-major data into a tensor whose metadata remains row-major, describe that specific contract separately from a tensor whose stride is actually (1, M); .reshape(-1) has different implications in those two cases. Passing sa directly is the simplest documented option for the latter.
| pytest.importorskip("aiter", reason="aiter baseline not installed in this venv") | ||
| record = _run_bench(world_size, "lsa", 256, 1024, {"AR_CHECK_VS_AITER": "1"}) | ||
| assert record["validated"] is True | ||
| assert record.get("rel_l2_vs_aiter") == 0.0 |
There was a problem hiding this comment.
Follow-up 3 — non-blocking: please check whether the optional aiter comparison is fully wired.
The worker benchmark currently does not read AR_CHECK_VS_AITER, call an aiter reference, or emit rel_l2_vs_aiter. Once the dependency-presence gate is passed, this therefore fails as None == 0.0 even when the actual AR result validates.
I reproduced that by supplying an empty module only to bypass importorskip and then running the real 2-rank benchmark. No aiter numerical implementation was supplied or executed; this was a diagnostic of the test's missing comparison path. To clarify my earlier validation totals, this test was SKIPPED, not passed, in our normal environments because aiter was absent.
This predates the latest two commits, but please consider implementing the optional reference and result field, or adjusting/removing the test until that comparison exists, so an environment with aiter installed does not fail on a missing record field.
yangyuhuiling
changed the title
#691) `compile_fused_gemm_scatter` took a `fence` argument with eleven values. Every caller in the tree -- `op.py`, the tests, and `bench_gemm_ar.py`'s own default -- left it at `none`, so the two `const_expr` chains it drove never emitted anything and three `StoreC` subclasses were unreachable. They are unreachable for a reason that is worth stating where the code is rather than in an option nobody picks: **the counter is already the release.** The `acq_rel` atomic_add emits its own `buffer_wbl2`/`buffer_inv` pair, on thread 0 only, so an added `__threadfence_system` redoes that work per lane. gcnasm's `opus_chunk_sdma_submit` has the same shape and no fence, which is what pointed at it. Measured at m=16384 blockscale on 8 ranks: no fence 1143.5/1144.8us, `agent` 1797.0/1798.5, `all` 1936.1/1937.5 -- all three at relL2 2.350e-3. The fences cost 650-790us and bought no accuracy. That paragraph now sits in the epilogue instead of in `--help`. Removed: the parameter, its validation list, the `{fence[0]}` kernel-name tag, both dispatch chains, `_RawWriteThroughStoreC` / `_WriteThroughStoreC` / `_NonTemporalStoreC`, the now-dead `release_fence` import, and the benchmark's `--fence` with its result field. `direct_fence` is untouched -- it is Direct LSA's, it is live, and it defaults to `leader`. Two stale pieces of text went with it, both of which contradicted the `none` default they shipped alongside in #662: the module docstring's protocol sketch still listed `__threadfence_system()` as a step, and the epilogue comment still read "This fence is not optional". Read literally, either would send someone looking for a release that is not there. `kernels_sdma`'s own `publish` knob is a different parameter and is still in use; only its two comments pointing at the deleted `_WriteThroughStoreC` were trimmed. Verified rather than assumed, 3 runs a side at m=16384 blockscale on 8 ranks: median 1145.8 -> 1146.2us, +0.03%, against a 1.9-2.8us spread within each side, and relL2 bit-identical across all six (2.350369e-03). Co-authored-by: xiangch <xiangch@amd.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Fuses DSV4-Pro's
wo_bGEMM with its all-reduce over cco's SDMA copy engines,shipped as
mori.ops.gemm_ar. The fused kernel writes each destination's rowband straight into a symmetric window and hands it to the copy engines as soon
as its tiles are done, so the transfer overlaps the GEMM.
[16384, 7168]K=2048, 8x MI355X,blockscale): 1419.5 us as the model runs it today (GEMM + NCCL) against
1148.8 fused. An optional fp8 all-gather leg takes it to 979.3 us, -31%.
gather. Smaller than the layer win because
wo_bis ~12% of the profile.GemmAllReduceOp, mirroringEpDispatchCombineOp: it owns thesymmetric window, the per-M kernel cache and the row padding.
supports()answers whether a shape is expressible.
BUILD_CCO_SDMA=ON. With it off every put silently does nothing,the model still answers fluently, and the fused path measures faster than it
is because it moves no data.
GemmAllReduceOp.self_test()now catches this atconstruction; the CI image and install script set the flag.
Full measurements, how to reproduce each one, the model-level evaluation of the
fp8 option, and the negative results kept deliberately are in
docs/MORI-GEMM-AR-BENCHMARK.md.🤖 Generated with Claude Code