[FlyDSL] Replace the split-K atomic combine with a workspace + reduce kernel - #4622
Open
JohnQinAMD wants to merge 3 commits into
Open
[FlyDSL] Replace the split-K atomic combine with a workspace + reduce kernel#4622JohnQinAMD wants to merge 3 commits into
JohnQinAMD wants to merge 3 commits into
Conversation
Contributor
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
|
… kernel
The FlyDSL split-K GEMM combined its splits with a bf16 atomic fadd straight
into C. That accumulate needs an initialised target, so the kernels carried a
handshake to get one: the first-arriving block zeroed the C tile, raised a
signal, every other block spin-waited on it, and the last one to depart reset
the counter for the next launch.
That handshake requires the counter to be zero when a launch starts, and
nothing can guarantee that under CUDA graph capture. The counter is zeroed
only where it is allocated, which happens outside the captured region, so a
recorded graph contains no memset and a dirty counter can never be cleared.
The small_m family, whose reduction is gated on a block observing prev == 0,
then spins forever: no first arrival, so no signal, so no exit from the wait,
so no reset. The one path that could clean up is blocked by the state it
would clean.
Rather than protect the counter, remove the reason it exists. Each split now
writes its fp32 accumulator with plain stores into its own disjoint slot of a
[SLOTS, m, N] workspace, and a new reduce kernel sums the slots in fixed
order, folds bias, casts and writes C once. Ordering between the two launches
comes from the stream -- inside a graph, from the dependency edge between two
nodes. Nothing is shared, so there is no coordination state a replay can
poison. This mirrors csrc/opus_gemm, which reduces the same way ("no
atomic_add, no self-clear, no semaphore").
Measured on MI355X (gfx950) across 71 shape x split_k x tiling configs,
against the atomic combine:
accuracy better on 71/71; max abs error 3.476 -> 0.999, i.e. bounded
by half a bf16 ULP -- correctly rounded. Partials are no
longer rounded to bf16 before being accumulated.
determinism bit-identical run to run on all configs; the atomic path
differed in 27-75% of elements on 5 of 7 spot checks.
performance graph replay median 0.951x, mean 0.911x, faster on 45/71.
Concentrated where split-K is actually selected: m=1 is
0.805x. Eager is ~1.10x from one extra host dispatch.
capture replays match eager bit-for-bit, and still do after the
workspace is filled with garbage and NaN.
The workspace is unpadded and stores are masked to row < m: the reduce reads
only real rows, so writing a grid-aligned M_PAD tile was pure write
amplification (32x at m=1 with BLOCK_M=32).
The legacy combine is deleted rather than kept behind a flag. There is no
measured shape where it wins, and it is outright wrong in one: with
PERSISTENT_N_TILES > 1 the small_m combine produces garbage (max abs error
~106, 60-85% of elements, differing between two runs of the same binary)
because its per-tile arrival counter is shared with blocks at a different
tile in their sequence. That bug predates this change and reproduces on a
pristine checkout; the redesign makes it moot.
Both kernel files end up smaller than before (-153 and -181 lines).
test_flydsl_splitk_hgemm.py's thresholds were loosened to 4 bf16 ULP with 1%
of elements allowed to be arbitrarily wrong, to tolerate the atomic path.
Tightened to one ULP at 100% on all four cases; all now measure 0.0000.
op_tests/test_flydsl_splitk_workspace.py adds correctness, bit-reproducibility
and graph-replay-matches-eager coverage, plus a conformant @benchmark sweep.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
JohnQinAMD
force-pushed
the
flydsl-splitk-workspace-reduce
branch
from
August 7, 2026 06:55
a485f58 to
4a5e2ff
Compare
Contributor
There was a problem hiding this comment.
Pull request overview
This PR replaces FlyDSL split‑K HGEMM’s in-place bf16 atomic accumulation (and its semaphore/signal coordination) with a two-stage approach: each split writes fp32 partials into a per-device fp32 workspace, and a dedicated reduce kernel deterministically sums the slots into C.
Changes:
- Main split‑K kernels now write fp32 partials to a
[slots, m, n]workspace (no atomics, no cross-block barrier state). - New
splitk_reducekernel reduces workspace slots in fixed fp32 order (optionally folds bias once) and writesConce. - Host launcher replaces semaphore/signal cache with a growable per-device fp32 workspace plus a prewarm helper for CUDA graph capture.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| op_tests/test_flydsl_splitk_workspace.py | Adds correctness/determinism/graph-replay tests and a perf sweep for the new workspace+reduce combine. |
| aiter/ops/flydsl/test_flydsl_splitk_hgemm.py | Tightens split‑K precision thresholds now that accumulation is fp32 and deterministic. |
| aiter/ops/flydsl/kernels/splitk_reduce.py | Introduces the new reduce kernel that sums workspace slots and writes C. |
| aiter/ops/flydsl/kernels/splitk_hgemm.py | Replaces split‑K atomic combine path with fp32 workspace stores and updates kernel ABI. |
| aiter/ops/flydsl/kernels/small_m_hgemm.py | Removes semaphore/signal barrier logic and writes split‑K partials into workspace. |
| aiter/ops/flydsl/gemm_kernels.py | Implements workspace allocation/growth/graph-capture prewarm; launches main+reduce for split‑K. |
| aiter/aot/flydsl/gemm.py | Updates AOT compilation inputs for the new kernel ABI (workspace pointer). |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Four review comments from the Copilot reviewer on ROCm#4622. 1. AOT precompiled only the main split-K kernel, so every split-K config still JIT-compiled the reduce kernel on first use -- exactly the cost GEMM AOT precompilation exists to remove. Compile the reduce kernel too when split_k > 1. The slot count is part of the layout contract between the main kernel, the reduce kernel and the AOT precompiler, so it moves into `_split_k_workspace_slots` and all three read it from there instead of restating `split_k * (block_k_warps if hgemm else 1)`. 2. `_get_split_k_workspace` tested `is_current_stream_capturing()` on the *current* stream while allocating on the passed-in `stream`. When a caller supplies an explicit launch stream those can disagree, and the guard would then miss a capturing stream and allocate into the graph's private pool -- the one thing it exists to prevent. Run the check and the allocation under the same `torch.cuda.device` + `torch.cuda.stream` context. Only the grow path reaches it, so this stays off the hot path. `flydsl_splitk_prewarm_capture_workspace` had the same split, plus `torch.cuda.Stream()` binds to whichever device is current: prewarming for a non-current device tested, and registered the workspace on, another device's stream. Resolve both under `device`. 3. The comment on SPLIT_K_WORKSPACE_MAX_BYTES claimed the fp32 workspace was kept "well below" 2GiB while the constant was exactly 2**31, and attributed the bound to i32 byte offsets. The real bound is the buffer descriptor: `num_records` is a 32-bit BYTE count (clamped to 0xFFFFFFFF in buffer_ops) and voffset is a 32-bit element offset scaled to bytes, so addressing wraps at 4GiB. The constant is right -- half the wrap point -- and the comment now says so. Same fix to the error message. 4. `Case.overrides` used a mutable `{}` as a NamedTuple default, shared by every instance that omits it. Use MappingProxyType so it is actually immutable rather than suppressed with a noqa; same for the shared SMALL_M dict. Verified on MI355X (gfx950) with the patched tree mounted over the installed package in the image pinned in the PR description: pytest op_tests/test_flydsl_splitk_workspace.py 21 passed test_flydsl_splitk_hgemm.py 4/4, max_delta 0.0000 AOT split_k=2 2 kernels compiled AOT split_k=1 1 kernel compiled Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
`flydsl_splitk_prewarm_capture_workspace` was added with no production
caller: only the new test used it, while the real path --
`tuned_gemm.flydsl_gemm` -> `flydsl_hgemm` -- never called it. The
workspace cannot grow inside a CUDA graph capture, so capturing a shape
that had not already run eagerly at that `m` raised RuntimeError instead
of degrading. Reproduced: run m=1 eagerly, then capture m=64 with the same
config and the capture aborts.
Call it from `flydsl_gemm`, the same placement and for the same reason as
`_opus_prewarm_capture_workspace` in `gemm()`. Dispatch is the right layer
because only the tuner knows the `split_k` that decides the size. Soft-fails
like the opus one: a prewarm that cannot run must not break eager callers,
and capture re-surfaces the problem as a clear error anyway.
Calling it per GEMM required making it cheap when warm. It ended every call
with `capture_stream.synchronize()`, which was free when only a test called
it once but is a host-side sync per eager split-K GEMM from dispatch. Sync
only when the call actually allocated. Resolve `device.index` for the same
reason: an index-less `cuda` device missed the cache lookup and would have
synced every time.
Verified on MI355X (gfx950), patched tree mounted over the installed package:
capture of an unwarmed shape, before RuntimeError (the repro above)
same shape, after captures; replay bit-identical to
eager with the workspace NaN-filled
warm path no realloc, so no per-GEMM sync
split_k=1 / index-less device no-op / no realloc
pytest test_flydsl_splitk_workspace.py 21 passed
test_flydsl_splitk_hgemm.py 4/4, max_delta 0.0000
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
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.
Purpose
What breaks
A FlyDSL split-K GEMM recorded into a CUDA graph can spin forever on replay. The device wedges — it does not return a wrong number and it does not recover. The same design also makes the op non-deterministic run to run, and costs accuracy, because partials are rounded to bf16 before they are accumulated.
Root cause
It is a five-step chain. Each step is forced by the one above it.
1. The splits are combined by a bf16 atomic add straight into C.
kernels/splitk_hgemm.py:1059onmain:2. An atomic add needs an initialised target, so the kernel needs a barrier to produce one. Somebody must zero the C tile before anyone adds to it, and everybody else must wait for that. That is the entire purpose of the semaphore/signal pair:
zero_c()signal[idx] = 1split_k_barrier()3. That barrier only works if the counter is 0 when a launch begins. Step 1 elects "first arrival" by
atomicAdd(...) == 0, and step 4 elects "last departure" by comparing against a fixed arrival count. Both are false if the counter starts anywhere other than 0.4. The counter is zeroed exactly once — where it is allocated.
gemm_kernels.py:700onmain:After that, the invariant is maintained solely by the kernel resetting the counter on its way out.
5. Under CUDA graph capture, neither holds. The
lru_cacheentry was allocated the first time the shape ran — during warmup, before capture began — so the recorded graph contains the kernel launch and no memset. Nothing re-establishes the entry state on replay. A counter left non-zero therefore stays non-zero forever: clearing it was never part of the graph.For
small_mthat is fatal, because its reduction is gated on the same state it would have to clean (kernels/small_m_hgemm.py:712-800, condensed — the source is FlyDSL MLIR builder calls):No first arrival → no signal → the wait never exits → the reset never runs. The only path that could clean up is blocked by the state it would clean.
The generic
hgemmfamily survives a dirty counter because there its counter only elects who does the cleanup, with nothing conditioned on a zero start — it completes, with the counter drifting. That is why the currently dispatched family shows no symptom.What changed
The counter exists only to make in-place atomic accumulation safe. So rather than protect the counter, this removes the accumulation that needs it.
In the main kernel: atomic add into C → plain store into a private slot.
New:
kernels/splitk_reduce.py(145 lines) sums the slots. One launch, no coordination:Ordering between the two launches comes from the stream — inside a graph, from the dependency edge between the two nodes. Nothing is shared between blocks, so there is no coordination state a replay can poison.
On the host: a counter cache → a growable workspace.
_get_split_k_tensorsand_check_split_k_semaphore_capacityare gone, replaced by_get_split_k_workspacesizingslots * m * nfp32 elements, plusflydsl_splitk_prewarm_capture_workspace()so the buffer can be sized before capture starts — growth inside a captured region is not allowed. This followsaiter/tuned_gemm.py::_opus_prewarm_capture_workspace, which does the same for opus.Deleted, not kept behind a flag:
zero_c(), the signal store,split_k_barrier(), the counter reset, the bf16 atomic accumulate,SPLIT_K_SEMAPHORE_MAX_LEN, and the semaphore/signal kernel ABI arguments. Both kernel files end up smaller than before (splitk_hgemm.py1121 → 983,small_m_hgemm.py1412 → 1242).This is not a new pattern in this repo.
csrc/opus_gemmalready reduces exactly this way and says so in its ABI header — "no atomic_add, no self-clear, no semaphore". That option is not open to the ASM a16w16 split-K GEMM, which ships as precompiled.cobinaries and therefore has to fix its semaphore instead (#4494); the FlyDSL kernels are generated from the Python DSL, so the counter can simply go.Test plan
Three things need proving, and they use different tools.
1. Correctness, reproducibility and capture safety —
op_tests/test_flydsl_splitk_workspace.pyNew file. Seven
(shape, split_k, family)cases, three assertions each:test_workspace_matches_referencetest_workspace_is_deterministictest_graph_replay_matches_eagerThe third one deliberately corrupts the workspace between replays (
1e30,NaN,1e30). Without that it would be worthless: a clean capture-and-replay passes on the atomic combine as well, so it has to seed the adversarial state that the old design could not recover from.2. Roofline sweep — same file, run directly
@benchmarksweep over 6 shapes x 3 split_k x 5 tilings, filtered to launchable configs. ItsTB/scolumn counts the workspace round trip (2 * slots * m * n * 4bytes), so the design's real cost appears in the table instead of being hidden:3. Before/after A/B — two trees, one script
The legacy combine is deleted, so it cannot be selected with a flag. Instead the same script runs against two checked-out trees, mounted over the installed package:
ab_full.pywalks the same config space, times each config twice — once eager, once as a captured graph replayed 30 times — records max abs error against an fp32 torch reference, and writes JSON. Configs an unsupported tiling rejects are recorded as such and excluded from both sides, so the comparison is over the identical set. 71 configs measured on both trees.Full environment for all three:
Test results
MI355X (gfx950). Ratios below are this PR ÷
main, so < 1.0 is this PR being faster or more accurate.Suites
pytest op_tests/test_flydsl_splitk_workspace.pyerr0 on every onetest_flydsl_splitk_hgemm.py(pre-existing)max_delta = 0.0000That last file's thresholds had been loosened to tolerate the atomic path —
max_delta_limitof 32.0 and 8.0 with 1% of elements allowed to be arbitrarily wrong, i.e. 4 bf16 ULP. Tightened here to one ULP at 100% on all four cases, since the workspace path measures exactly 0.Accuracy — better on 71 of 71 configs, none worse
main(bf16 atomic accumulate)0.9994 is half a bf16 ULP at that magnitude — correctly rounded, the floor for a bf16 output. The atomic path cannot reach it because it rounds each partial to bf16 before adding. Its error also varies run to run.
Determinism
Bit-identical across two identical runs on all 71 configs. The atomic path differed in 27-75% of elements on 5 of 7 spot checks.
Performance
Split-K is selected for skinny-M decode shapes, and that is where the win lands:
The remaining regressions are the second launch, not the kernel body. Timed in isolation inside a graph, the reduce node costs 1.0-1.3us regardless of workload — 1.06us to move 640KB, 1.01us to move 57KB. It is launch-bound, and that floor matches the residual almost exactly:
m16_n7168_k512_spk2regresses by 1.55us against a 1.27us reduce. Removing that cost means removing the second launch, which is the design.Eager pays one extra host dispatch, a roughly constant 2-5us. That is implementation cost, not design cost.
Capture safety
A clean capture-and-replay proves nothing here: it passes on the atomic combine too, because the counter only wedges the kernel once it has been left dirty. So both sides are measured with adversarial state seeded before capture. The object differs, necessarily — after this change there is no counter to seed, so the workspace is seeded instead:
main— counter seeded dirtyThe
mainrow is cited from earlier work on this defect, not produced by this PR's suite: it used a seeded-semaphore test against the unpatched tree, reproduced in the public image pinned above (1 failed in 124.56s, childTimeoutExpired). That test cannot have an "after" column here, because the semaphore it seeds no longer exists.The
this PRrow istest_graph_replay_matches_eager, which replays four times and refills the workspace with1e30,NaN,1e30between them. Every replay stays bit-identical to eager, because the combine reads only slots it wrote in the same launch.A pre-existing bug this removes
Legacy split-K with
PERSISTENT_N_TILES > 1insmall_mdoes not merely lose precision, it produces wrong output. Reproduced on a pristine checkout ofmain, so it predates this change:main, run 1main, run 2mainTwo runs of the same binary disagree, and
split_k=1— which never enters the combine — is exact.prepare_split_k_tile/split_k_barrierare re-entered per tile against a counter shared with blocks that are at a different tile in their own sequence, so the "first arrival zeroes / last departure resets" invariants do not hold. Not fixed separately: the redesign removes the code that carries the bug.Overlap and limits
#4494 fixes the same class of defect on the ASM a16w16 split-K path, which cannot take this approach because it ships as precompiled binaries. The two are independent and touch different files.
small_mis not currently dispatched —iter_small_m_registry_configsis commented out atgemm_kernels.py:24andaiter/configs/holds 0small_mentries. The hang this removes is therefore latent rather than something users hit today. The accuracy, determinism and capture-safety results apply to the dispatchedhgemmfamily immediately.Tool assistance
Claude Opus 5 (1M context) assisted with the root-cause analysis, the redesign, the measurement harness and drafting this description.