Repository navigation
gemm_ag: GEMM fused with an all-gather over cco - #807
Open
yangyuhuiling wants to merge 18 commits into
Open
yangyuhuiling wants to merge 18 commits into
yangyuhuiling wants to merge 18 commits into
Conversation
A third fusion operator on gemm_ar's 8-wave fp8 GEMM, alongside gemm_ar's
all-reduce and gemm_a2a's all-to-all. Every rank computes its own [M, N] and
every rank ends up with all `world` of them concatenated along rows -- the
DeepSeek V4-Pro prefill-CP wkv_gate shape, where each context-parallel rank
scores its own tokens and then everyone needs every token's score.
Eight modes: gemm-only, gemm-to-window, split-rccl, split-lsa-{push,pull},
fused-lsa, split-sdma, fused-sdma. ptpc/blockscale/mxfp8 throughout.
All-gather shards nothing, so three things gemm_a2a needs are absent: the
staging region (the payload is already contiguous and already the same for
every destination), the per-(destination, chunk) counter (one chunk arms every
push at once), and the destination rotation (a broadcast tile belongs to
everyone). The layout's N rule reduces to `n % block_n`, which is what lets
this op run the ratio-128 wkv_gate shape at N=1024 -- a shape a2a_config
rejects at world_size=8.
Because the collective is a broadcast it can be driven from either end, so
both directions are implemented and compared. That is also the first thing the
measurements say: pull needs the producer to publish C with system scope,
since it reads a peer's HBM while the peer's GEMM left the line dirty in its
own L2. Without it split-lsa-pull validated at 1.66e-3 under ptpc and 2.05e-1
under blockscale -- a correctness bug visible in exactly one quantisation,
which is why the mode matrix is parametrised over all three.
gemm_ar's _SwapABStoreC gains a `peer_rsrcs` list alongside `peer_rsrc`. A
broadcast stores one tile to every peer, and doing that with `world` separate
store objects would repeat the scale loads, the convert and the permlane
shuffle with it. A single resource yields a one-element tuple and identical IR,
so gemm_ar and gemm_a2a are unchanged.
Measured on 8x MI355X: the split paths beat RCCL by 8-9%, and the fused paths
lose to their own split baselines. The fusion itself works -- at M=512
fused-sdma cuts exposed comm 50.9us -> 21.8us -- but the epilogue's per-block
publication roughly doubles the GEMM. See docs/MORI-GEMM-AG-BENCHMARK.md.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Two wkv_gate shapes x two sequence lengths x every mode, plus the full quantisation matrix and the fused-sdma chunk ladder, on 8x MI355X. Every mori transport beats RCCL by 10-25%, and push beats pull everywhere -- the opposite of gemm_ar's gather leg, which found pull worth 5.6-10.5% over a *copy-engine* push. The GEMM phase is equal to within 0.5us across the pair, so pull's mandatory sc0|sc1 C store is not the cause; a push is a fire-and-forget store where a pull is a load that must round-trip. That is worth knowing beyond the ranking, because pull is where a low-precision wire would want its dequantize. The fused paths lose, and --no-put and --fence none between them say why. The overlap itself works: fused-sdma cuts exposed comm by 25-59%, a larger fraction than gemm_ar's all-reduce absorbs. The release fence is free at these grid sizes (32-128 blocks, against the 9216 where gemm_a2a calls it the most expensive line in its epilogue) and the bookkeeping is 2-6us. What costs is *issuing* the puts: 39-55us, nearly constant across four cells whose payloads differ by 8x, which rules out bandwidth contention and lands on gemm_ar's own constant of ~2us per SDMA packet regardless of size. At --chunks 4 and world=8 that is 28 puts posted serially by one thread under one lock, where gemm_a2a posts the same 28 from 28 different blocks -- its counter elects per (destination, chunk), so each elected block owns one destination. A broadcast chunk arms every destination at once, which is what made the counter cheaper and the posting expensive. Fixing it is a change to the epilogue, not to the design, and the doc says what it is. Separately, the bound: all-gather's G0/C0 goes as K/(world-1) where both siblings go as K, because they shard the payload by world and it does not. It is the one operator in the family where scaling out makes fusion less attractive, and the one where a narrower wire is worth more than any epilogue tuning. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The fused-SDMA epilogue was spending 39-55us of GEMM time issuing a chunk's world-1 SDMA packets back to back from thread 0 -- nearly constant across four cells whose payloads differ 8x, which is SDMA's ~2us per packet regardless of size, the constant gemm_ar's layout already states. --post lanes gives each of lanes 0..world-2 one packet. What made that awkward is that a broadcast chunk arms every destination at once, so the natural counter is one per chunk -- and the thread that wins a single counter then has to tell world-2 other lanes, which costs a barrier and an LDS round trip for one dword. So the counter goes back to one per (chunk, lane), gemm_a2a's shape for a different reason: every block increments all world-1 in one instruction, they cross the threshold together, and each lane learns from its own atomic. No broadcast. The submit lock had to go with it. world-1 lanes retire in lockstep, so two waves each holding part of a lock set and spinning for the rest could never reach their release. The same collision -- two chunks of one destination posting at once -- is separated by queue instead (chunk % sdma_queues), hence the new sdma_queues >= chunks requirement, which the benchmark derives rather than asks for. Measured, ptpc r4: the fused GEMM phase goes flat at 57.4-60.4us across c1..c8 at M=512 and 66.6-71.9 at M=2048, where serial climbs to 105.2 and 168.8. At M=2048 c8 that is -96.9us of GEMM time. Best total 113.1 -> 105.5 (-6.7%) and 224.1 -> 217.2 (-3.1%). All twelve quant x chunk cells validate at the same relL2 as every other mode. It still does not beat split-sdma (105.5 against 104.1; 217.2 against 210.5), and fixing it exposed that the first round's headline was wrong: serial's drain shrank with chunk count not because the fusion absorbed communication but because the GEMM was stalled posting and the copy engines got a head start, which the phase split booked as absorbed comm. With the posting fixed the drain grows with chunk count instead -- the ~2us per packet moves from the issue side to the wire, and chunking an all-gather multiplies packets by world-1 per chunk because every peer gets the same slab. --chunks 1 is now the best fused setting everywhere, and the remaining gap is ~10us of bookkeeping against ~7.5us of real early-start. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
gemm_ag was built on gemm_ar's fp8 GEMM, and the layer it exists for is not
fp8: DeepSeek V4-Pro's wkv_gate goes through linear_bf16_fp32, and the planning
doc says in as many words that the first implementation must keep that input
precision and FP32 output convention rather than substituting the FP8 GEMM. So
the operator had the right shapes and the wrong precision, and its numbers could
not be put next to the FP32 baseline they are meant to replace.
_gemm_a16w16_8wave.py ports gcnasm's gemm_a16w16_quad_subtile_kernel_template
at its own instantiation <512, 256, 256, 64, bf16, bf16, bf16, float>. The port
is small because gemm_ar's _gemm_a8w8_8wave.py is already a FlyDSL rendering of
that same template with fp8 operands -- same 8 waves, same 2x2 half-tiles, same
swapped-AB MFMA, same permlane store. Four things change: the MFMA becomes
16x16x32 bf16, E_K goes 1 -> 2, VEC goes 16 -> 8, and there are no scales. The
LDS keeps mori's XOR swizzle restated in bytes (swizzle_row128b), asserted in
the tests to be the fp8 function value-for-value at elem_bytes=1.
The FP32 store turns out to be the cheap one. The A/B swap puts four
consecutive columns in each lane; at two bytes that is 8 B and needs a
permlane16_swap to pair two N-tiles into 16 B, and at four bytes it is already
a buffer_store_dwordx4 with no shuffle and no convert.
elem_bytes becomes load-bearing rather than decorative: layout validates it,
kernels_lsa derives its pack width from it, and the benchmark gains
--in-dtype {fp8,bf16} and --out-dtype {bf16,fp32}, both defaulting so every
existing number stays reproducible.
Measured, 8x MI355X, K=7168 N=2048. Accuracy first: relL2 9.9e-7 for fp32 out
against 1.66e-3 for bf16 out -- three orders of magnitude, and it shows the
1.66e-3 was always the bf16 *output* rounding rather than the fp8 inputs.
All 32 cells validate. Fusing now breaks even (fused-sdma within 0.5% of
split-sdma at every cell, with 4-9% of the transfer genuinely absorbed) where
under fp8 it lost by 3-7% -- the bf16 GEMM is 2.2x larger, which is exactly the
G0/C0 lever the bound section describes. fused-lsa gets worse, and fp32 makes it
much worse (616 against split-sdma's 428), because its epilogue pushes world
copies of every element through the CUs.
The honest headline: against the planning doc's FP32 baseline the collective leg
agrees to 3.4% (327.8 against 339.2) but our GEMM is 28% slower than
torch.matmul (137.6 against 107.2), so at the target shape this is not yet a win
and the whole deficit is in the multiply. BLOCK_K is pinned to the template's
64, xcd_swizzle is off and no tile sweep has been run; tuning that is worth more
than any further epilogue work.
Also recorded: FlyDSL caches compiled kernels keyed on the @flyc.jit wrapper's
name, not the kernel body, so a fixed bug can look unfixed. Here it produced a
thoroughly convincing false signal -- odd K_ITERS bit-exact, even K_ITERS broken
-- which looks exactly like a double-buffer parity bug and was not one.
One cell (fused-sdma bf16 M=2048) failed validation once and passed the next 13
runs. Left in the doc as unresolved rather than dismissed.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The bf16 port landed 28% behind torch.matmul, called out as the largest number
on the table. The cause is not the mainloop: the template's 256x256 tile
launches ceil(M/256)*ceil(N/256) workgroups, which at the wkv_gate shape
(M=N=2048) is 64 against an MI355X's 256 CUs. Three quarters of the GPU idles
and nothing inside the loop can reach that.
Measured at M=2048 N=2048 K=7168, fp32 out, one GPU, torch.matmul at 84.8us:
256/256 (template) 64 wg 176.0us 342 TF/s
256/128 128 wg 139.6us 431 TF/s
128/256 128 wg 134.2us 448 TF/s
128/128 256 wg 116.5us 516 TF/s
waves_per_eu (1/2/4) and xcd_swizzle (0/1/4/8) are inert -- every cell within
1%. Across shapes the rule is not "maximise the grid" either: at M=4096,
128/128 gives 512 workgroups and 162.8us against 128/256's 256 workgroups at
155.2, so once there is one workgroup per CU the larger tile amortises better.
pick_tile takes the largest tile whose grid still covers the CUs, which
reproduces the measured best at all four shapes.
**Neither faster tile ships.** 128/128 races: one wrong fp32 C in 48 single-GPU
runs, and about half the ranks per 8-rank launch. 128/256 is worse than it first
looked -- gemm-to-window and split-rccl validated 4/4 and 3/3, then came back
1/2 on a later pass, while fused-lsa is 0/4 throughout. BLOCK_M=256 has not
failed once across every repeat run. So the failure rate at 128 is around a
third of launches, and the tidy "split paths are sound, fused paths are not"
story the first pass suggested did not survive repetition.
The obvious suspect is the release fence, since halving BLOCK_M doubles the
blocks that must publish. It is wrong: --fence all and --direct-fence all are
0/3 each.
Buying 24% for a wrong answer a third of the time is not a trade, so pick_tile
returns the template's tile and the speedup does not reach the default. What
ships is the knowledge: the tile is worth 24-34% and is the only knob that is,
the mechanism is grid occupancy rather than scheduling (so the next attempt is
split-K), and there is a latent correctness bug that halving BLOCK_M exposes
which was invisible at the template's tile. Finding it is worth 33% on the GEMM
and is now reproducible.
Also here: layout accepts block_n=128 under an fp32 C (a bf16 store pairs two N
tiles through permlane16_swap and cannot); the fused-SDMA C store moves to a
buffer resource at this rank's recv slot, matching gemm_ar, which fixed a wrong
answer at M=2048 fp32 (relL2 6.9e-3 -> 9.9e-7); and compile_bf16_gemm_ag now
rejects a tile that disagrees with its config, which would hang the completion
counter rather than corrupt it.
All eight modes validate at the shipped default.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The tile sweep left a bug it could not explain: BLOCK_M=128 is 24-33% faster
and corrupts about a third of 8-rank launches, including gemm-only, which has
no collective in it at all. This pins the cause.
`wait_barrier(N)` is `s_waitcnt vmcnt(N)` plus `s_barrier`, so it is the only
thing making a G2S load's LDS write visible before the next iteration's S2R
reads it. Sweeping that one number, BM=128 fp32 M=2048, gemm-only, six
launches each:
vmcnt(2A+B) = 4 (inherited) 5/6 91us
vmcnt(A+2B) = 5 (derived) 4/6 90us
vmcnt(0) (everything) 6/6 150us
So it is the waits, and `vmcnt(0)` is correct. A `--wait-policy` knob is added
to make that reproducible rather than a one-off.
It also shows the *derivation* is wrong, which matters more than the knob.
Walking the mainloop -- entering with A+2B outstanding, issuing A1@k+1, B0@k+2,
A0@k+2, B1@k+2, needing the 4th item retired -- gives a requirement of
wait <= A+2B, under which:
256/256 A=2 B=2 in-flight 14 wait 6 -> retires 8, needs 8, margin 0
128/256 A=1 B=2 in-flight 11 wait 4 -> retires 7, needs 6, margin 1
The tile with *zero* margin is the one that is clean and the tile with margin
is the one that fails, so the model does not describe the hardware. Something
is outstanding on vmcnt that this counting does not include -- the S2R reads
are the obvious suspect (16 per iteration at BM=128, 24 at BM=256) if they
lower to flat loads rather than ds_read, since flat counts against vmcnt.
Settling that needs the generated ISA, which the FlyDSL cache does not keep
(pickles, not objects).
Nothing about the shipped configuration changes: the default is still the
template's BLOCK_M=256 at wait_policy="tuned", which is what every mode
validates on. vmcnt(0) is not a fix worth taking either -- at 150us it is
slower than BLOCK_M=256's 137.6, so it spends more than the tile saves.
What this buys is a bounded next step. The question is no longer "why does a
smaller tile corrupt" but "how many vmem ops are actually in flight at that
one instruction", which the ISA answers directly.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The wait-count hypothesis was confirmed by a sweep (vmcnt(0) is 6/6 where vmcnt(4) is 5/6) and is refuted by the generated code. Both are worth having. Dumping with GPU_DUMP_CODE_OBJECT=1 and llvm-objdump --mcpu=gfx950, at M=2048 N=2048 K=7168 fp32: * the S2R loads are `ds_read_b128` -- lgkmcnt, not vmcnt. They were never part of this accounting, which kills the leading theory outright: 16 per iteration at BM=128, 24 at BM=256; * the G2S loads are `buffer_load_dwordx4`, exactly A+B per load call, 6 per iteration at BM=128 and 8 at BM=256; * `vmcnt(4)` and `vmcnt(6)` each retire precisely one iteration's worth, leaving in flight exactly the k+2 prefetches. Both are correct; * the two skeletons are structurally identical, 8 barriers per iteration in both, with `s_waitcnt lgkmcnt(0)` correctly between the barrier pair that guards the write-after-read on each LDS buffer; * neither spills, and both get one workgroup per CU. The only asymmetry is 152 VGPRs / 96 KiB LDS against 248 / 128 KiB, which is just the tile. One more candidate came out of reading it and died the same way. gemm_ar.wait_barrier puts s_waitcnt and s_barrier in a single inline-asm string with constraints="" -- no memory clobber -- so the compiler is never told there is a barrier in there and could legally move an LDS access across it, which would be a write-after-read on the double-buffered LDS and would plausibly depend on register pressure. Emitting the wait as asm and the barrier as rocdl.s_barrier() produces byte-identical code. That option is therefore not kept: an option that provably cannot differ is not an option. So vmcnt(0) fixing the corruption is timing perturbation, not a corrected count. Put beside the other end of the evidence -- 0 wrong in 64 single-GPU runs at every tile, no collective in the loop -- the weight has moved off this kernel. What reproduces it is eight concurrent processes; what never reproduces it is one idle GPU running the same kernel on the same data. Next place to look is the harness or the environment: whether the benchmark's own fp32 reference can be corrupted under memory pressure, and whether the cco window's VMM allocation interacts with a plain torch.empty C (ROCm 7.2.x is known to route uncached VMM allocations to the coarse-grained pool). The --wait-policy knob stays so the sweep is reproducible. Nothing about the shipped configuration changes. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Order LDS reads before buffer reuse and drain the final prefetches. Honor peer_uncached for unfused C stores and publish from every producer lane. Initialize the full benchmark window so stale arrival flags cannot bypass the first pull barrier. Add pipeline dependency, exact numerical, graph replay, and dirty-window regressions. Validation: 151 related tests passed; BM=128/256 with BF16/FP32 output each passed five independent eight-rank benchmark launches.
Chunked fused producers can submit through every queue to a peer. Wait for all of those queues before signalling completion, while retaining the existing split and single-queue waits. Validate the publication contract with 24 queue-state cases. The corrected path also passed eight-rank BF16-to-FP32 graph replay checks with one, two, and four chunks.
Expose split-K FP32 partials, serialized SDMA chunk posting and release-only producer publication through the kernel API. Retain direct GEMM, transport and pipeline benchmarks with shared graph timing and changed-input checks. Preserve the FP8 wait-policy argument fix and reject nonfinite peer output. Keep source generators, historical experiment plans and raw datasets out of the branch. Document the supported commands and retain the measured findings. Validation: 129 CPU checks; 54 gfx950 numerical/graph cases; 11 eight-rank BF16/FP8 transport and MORI/split-K/Torch pipeline checks; eight-rank native GEMM comparison; Black 25.1.0 and Ruff 0.12.0.
Label the Pro ratio-4 projection measurements and document Flash's separate 512-wide projections and combined 1024-wide candidate at K=5120. Record 8-GPU transport, M-sweep, pipeline and complete projection-order comparisons, including interleave restoration and KV/gate materialization costs. Use thread-local capture for RCCL to avoid HIP watchdog event queries invalidating global graph capture, and destroy process groups on CLI exit. Validation: 52 reported configurations/repeats on v2-015 passed initial and changed-input checks; maximum relative L2 error 8.04e-7. Two benchmark CPU checks, Black, Ruff, documentation links and numeric checks passed. Raw plans, scripts, logs and excluded attempts are archived outside the repo.
Rebased off the gemm_a2a branch (#695) onto main: nothing in gemm_ag imports gemm_a2a, and the 14 commits applied without conflict. * tests: _gemm_ag_pull_worker runs bench_gemm_ag.py through runpy.run_path, which does not put the script's directory on sys.path, so the `_bench_utils` split (25d419b) left test_pull_initializes_recycled_window failing on every run with ModuleNotFoundError. * mori.ops: register gemm_ag with the lazy FlyDSL submodules. * ci_cco.yml: run the four gemm_ag test files, after the GEMM+AR steps. Their 8-rank SDMA process groups leave queues that take seconds to come back; gemm_ag's tests settle and retry around that, gemm_ar's do not, and run right after this step six of them failed in hsaKmtCreateQueueExt (anvil.cpp:210) while passing on their own. As written the step passes here: 213 tests in 22m50s. * .gitignore: drop the local experiments entry; that archive lives outside the repository. * docs: BUILD_CCO_SDMA is a build flag and on by default since #710, so stop exporting it at run time; example paths no longer point at a private checkout. * README tables re-measured on main 8ae8156 (with #603 and #711), 8x MI355X, ROCm 7.14 / PyTorch 2.12 / FlyDSL 0.3.4.1: 54 configurations, all validated, relL2 <= 9.95e-7. Conclusions hold: 128x128 + split SDMA leads the Pro shape (355.9us, 12.6% under Torch + RCCL), fusing pays only for the combined Flash N=1024 at M=16384 (4.3%). Host and stack changed with MORI, so the deltas against September are not attributed. Co-Authored-By: Claude <noreply@anthropic.com>
The README's opening summary still quoted the September figures, and its strongest fusion claim (11% at M=16384) rested on a study the previous re-measurement did not cover. Re-taken on the same stack: split SDMA against sc_release c4 at N=2048, K=7168 gains 5.3% at M=8192 and 11.9% at M=16384, and loses at M=4096 (c4 -4.8%, c2 -1.6%) -- September's v2-015 pairs gave 4.9-5.2%, 11.1% and -3.2%. Co-Authored-By: Claude <noreply@anthropic.com>
kernels_fused.py carries its own copy of gemm_ar's 8-wave fp8 mainloop, and the merge of main brought #785 to gemm_ar only. Apply the same three changes: drain the peeled tail's producers with wait_barrier(0) before it reads the last step's B0/A0, restore the loop's end-of-step count to 2A+B, and make the loop's S0 a read_done_barrier(). The comment telling readers not to restore 2A+B goes, since #785 reversed it. The BF16 quad-subtile GEMM is a different pipeline and is untouched. The four gemm_ag test files pass as the CI step runs them (213), and at K=256 -- where the main loop runs zero times -- fp8 fused-sdma validates 3/3 under ptpc and 1/1 under mxfp8, split-sdma 1/1, all at relL2 1.66e-3 on 8x MI355X. Co-Authored-By: Claude <noreply@anthropic.com>
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds
mori.ops.gemm_ag: a GEMM followed by an all-gather along rows, built oncco. Each rank computes
C = A @ B.Tfrom its ownA [M, K]and a replicatedB [N, K]; every rank ends with[world*M, N], source rankrat rows[r*M, (r+1)*M). FP8 inputs give BF16 output (gemm_ar's 8-wave GEMM);BF16 inputs give BF16 or FP32 output through a new quad-subtile BF16 GEMM,
which is what DeepSeek-V4's compressor projection (
wkv_gate) runs.The shape it is aimed at is a context-parallel projection whose output every
rank needs: DeepSeek-V4-Pro ratio-4
compressor.wkv_gate(K=7168, N=2048)and V4.1-Flash ratio-2 (K=5120, N=512 per projection, 1024 combined).
Eight modes, from
gemm-onlythrough RCCL, LSA push/pull, fused LSA, splitand fused SDMA; the window paths write C straight into the rank's own receive
slot, so every peer receives the same contiguous slab and no staging layout
is needed. The only change outside
gemm_agispeer_rsrcsongemm_ar's_SwapABStoreC, so a broadcast epilogue does the scale/convert/permlane onceand repeats only the store; with one resource it emits the same IR as before.
Results
8x MI355X, ROCm 7.14, PyTorch 2.12, FlyDSL 0.3.4.1, this branch on main
8ae8156; BF16 in / FP32 out, 128x128 tile, 5 rounds x 101 graph samples,max over ranks per round, median of rounds, every configuration in a fresh
process group on an idle box. All 61 configurations validated (relL2 <=
9.95e-7, initial output plus two changed-input replays).
DeepSeek-V4-Pro
wkv_gate, N=2048 K=7168, per-rank M:At M=2048 the win is the tile and the transport, not fusion: split SDMA is
12.6% under Torch + RCCL. Fusion starts paying from M=8192. V4.1-Flash at
M=2048 is within a few percent across the leading paths (LSA push 126.4us at
N=512, split SDMA / fused c1 203.2 / 202.8 at N=1024); its combined N=1024
projection gains 4.3% from fusion at M=16384. Full tables, the pure-GEMM
comparison and the two-stream pipeline are in
python/mori/ops/gemm_ag/README.md.Worth a reviewer's attention
gemm_a2a). Nothing ingemm_agimports it, so it is rebased onto main and independent. The docs and
comments still compare against
gemm_a2ain places ("unlike gemm_a2a's..."); they read as references to gemm_a2a: fuse a column-sharded all-to-all into the fp8 GEMM epilogue #695 until that lands.
_gemm_ag_pull_worker.pyruns the benchmark throughrunpy.run_path, whichdoes not put the script's directory on
sys.path, so once its helpers movedinto
_bench_utils.pyevery run raisedModuleNotFoundError. Fixed in thelast commit.
on purpose. Its 8-rank SDMA process groups leave queues that take seconds to
reclaim; the
gemm_agtests settle and retry around that, thegemm_aronesdo not, and run right after it six
gemm_artests failed inhsaKmtCreateQueueExt(anvil.cpp:210) while passing on their own. Givinggemm_ar's spawns the same settle/retry is left for a separate change.were re-taken on a different host and stack (ROCm 7.14 / PyTorch 2.12 /
FlyDSL 0.3.4.1, against ROCm 7.2.x / PyTorch 2.10-2.11 / FlyDSL 0.2.4), after
CCO SDMA micro optimizations #603 and fix(cco): two SDMA device-post deadlocks (multi-queue publish, commit snapshot) #711. Every conclusion held; the deltas are not attributed.
EXPERIMENTS.mdkeeps the September study, marked as such.Test plan
213 passed in 22m50s on 8x MI355X, run exactly as the new CI step does. The
gemm_arsuites (92 tests) pass on this branch run on their own. blackclean; ruff adds nothing over
mainin the touched files.Not in this PR
A public op wrapper with
self_test, settle/retry for thegemm_artestspawns, and model integration (SGLang's Flash low-ratio CP path currently
gathers hidden states before projection; this operator is the reordered
candidate).
🤖 Generated with Claude Code