[FlyDSL] Fix cross-wave LDS race on the last K tile in 8-wave GEMM and gather_kv_b_proj by yuychang · Pull Request #6539 · ROCm/aiter

yuychang · 2026-10-14T14:09:40Z

Motivation

The FlyDSL 8-wave a8w8 bpreshuffle GEMM (flydsl_bpreshuffle_8w_*, gfx950) intermittently returns wrong results: whole 16–64 row (or column) bands of an output tile are off by one K tile (max relative error ~0.15). The same input gives different results on repeated launches. It hits every tile/swizzle/wave variant and is not tied to M: about 1% of calls fail at mid M, 50–90% at M=32768. 291 tuned rows on main use these kernels (Kimi-K3, GLM-5.2, GLM-5.3-flash, Qwen3.5).

Technical Details

All 8 waves fill each LDS tile, and waves 4–7 run one barrier behind waves 0–3. In the peeled second-to-last K step, the last K tile's LDS buffers were read before any wait-plus-barrier had drained the lagging waves' loads, so waves 0–3 could read LDS that waves 4–7 had not finished filling.
Fix: replace the barrier after the final a_next1 load with wait_barrier(N_LDS_STEPS_A) (one extra s_waitcnt in the ISA). The mxfp8 8-wave kernel already does the equivalent.
gather_kv_b_proj 8-wave (gather_gemm_8wave.py): its tighter main-loop wait (SAFE_VMCNT) can still leave the last K_ITERS-1 a0 load in flight when N_LDS_STEPS_A >= N_LDS_STEPS_B (e.g. BLOCK_M = 256), and waves 0–3 read it before the loading waves' next wait. It did not reproduce on gfx950, but it gets the same fix. Its b1 tile was always read after vmcnt(0).
Barrier pairing counts the s_barrier at the end of every mfma.call. With it, b_next0 is read only after both wave groups have executed the new wait.
New op_tests/test_flydsl_gemm_a8w8_8wave.py checks every one of --repeat launches per shape (into a NaN-filled output, against an fp32 reference) for both kernels, with one summary table each.

Test Plan

python op_tests/test_flydsl_gemm_a8w8_8wave.py (both kernels, 20 launches per shape).
Repeated-launch repro against an fp32 reference on fixed vs unfixed trees, 30–100 launches per config.
python -m pytest op_tests/test_flydsl_gather_kv_b_proj.py.
Perf: run_perftest, fixed and unfixed interleaved, 3 rounds.

Test Result

MI355X (gfx950), ROCm 7.2.

Kernel	Unfixed	Fixed
a8w8 GEMM, failing shapes (e.g. N=6400 K=7168 M=32768, 10 variants)	4–18 of 30 launches wrong	0 of 30
a8w8 GEMM, N=7168 K=3584 small/mid M	up to 1 of 100 wrong	0 of 100
gather_kv_b_proj, 36 configs + stress	0 of 30–100 (not reproducible)	0 of 30–100

The new test fails on unfixed main (up to 19 of 20 launches wrong) and passes with the fix. test_flydsl_gather_kv_b_proj.py passes (100 tests).
Perf is unchanged within noise: a8w8 GEMM within ±1% from M=1024 to 32768; gather_kv_b_proj within noise, about +0.5% at 32 heads.

Submission Checklist

Look over the contributing guidelines at https://github.com/ROCm/TheRock/blob/main/GOVERNANCE.md#pull-requests.

🤖 Generated with Claude Code

…8w8 GEMM - The peeled K_ITERS - 2 step read the K_ITERS - 1 b0/a0/b1 LDS tiles with no vmcnt wait + barrier pairing. - Waves 4..7 lag one barrier, so waves 0..3 could read rows/cols waves 4..7 had not finished DMA-ing into LDS. - Symptom: intermittent 16-64 row (or column) bands per tile off by one K tile, max rel err ~0.15, any M. - Replace the barrier after the final a_next1 load with wait_barrier(N_LDS_STEPS_A); perf is unchanged. - Add op_tests/test_flydsl_gemm_a8w8_8wave.py, which checks every one of --repeat launches per shape/kernel.

…ather_kv_b_proj - gather_gemm_8wave.py has the 8-wave GEMM's tail: a plain barrier after the final a_next1 load. - Its SAFE_VMCNT loop wait can leave the K_ITERS - 1 b1 loads (and at BLOCK_M >= 256 the last a0 step) in flight. - Waves then read those b1 / a0 rows before any vmcnt wait of the wave that loaded them covers the LDS DMA. - Not reproduced on gfx950: 0 bad of 30-100 launches per config, also with an HBM hog and with fault injection. - Same fix as the GEMM: wait_barrier(N_LDS_STEPS_A) replaces that barrier; one s_waitcnt in the ISA, perf unchanged. - Extend op_tests/test_flydsl_gemm_a8w8_8wave.py with a repeated-launch gather_kv_b_proj check and its own table.

github-actions · 2026-10-14T14:10:24Z

🏷️ CI Guide

Runs automatically on every PR:

✅ Pre-checks (submodule verification, code formatting)
✅ Aiter op tests (gfx942 + gfx950)
✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
`ci:gfx1250-ffm-triton`	Run the five-shard gfx1250 FFM Triton test suite
`ci:triton-300x`	Run an additional Triton test job on MI300X in PRs (added automatically when gfx942 configs change); main branch always runs both MI35X and MI300X
`ci:triton-355`	Run the full Triton test suite on MI35X, not only the tests the change affects
`ci:windows`	Windows JIT op smoke test on the gfx1201 ROCm runner
`ci:multi-gpu`	Aiter multi-GPU tests on the 8-GPU runner
`ci:sglang`	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
`ci:atom`	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
`ci:atom_full`	ATOM accuracy suite for PR and main models from ATOM `models_accuracy.json`
`ci:vllm`	vLLM benchmark: DeepSeek-V4-Pro, MiniMax-M3-MXFP4
`ci:all`	All standard extended tests (excludes `ci:atom_full`)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 6539 --add-label <label>

One backend per PR:
A PR changes one kernel backend: [Triton/Gluon] (Triton and Gluon count as one), [HIP], [ASM], [CK], [OPUS] or [FlyDSL]. If the title ends up with two backend tags, split the PR -- as stacked pull requests when one part cannot merge without the other.

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to stop the title rewrites; labels stay in sync either way.

Copilot

🔵 Needs a closer look

The remaining B0 synchronization gap needs correction and human review of gfx950 correctness and performance evidence.

1 open finding

Final B0 read remains unprotected, allowing stale data · New

What changed in this PR

This PR targets intermittent last-K-tile corruption in FlyDSL’s gfx950 8-wave GEMM and fused gather projection kernels.

Changes:

Adds wait-before-barrier synchronization to both pipeline tails.
Adds repeated-launch correctness checks and performance summaries.

File	Description
op_tests/test_flydsl_gemm_a8w8_8wave.py	Adds repeated-launch checks for both kernels.
aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py	Adds a final-prefetch wait and ordering comments.
aiter/ops/flydsl/kernels/gather_gemm_8wave.py	Adds the corresponding gather pipeline wait.

🧠 Review effort: Balanced

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

- mfma.call ends in s_barrier: waves 0..3 pass c10's only after waves 4..7 ran the tail vmcnt(N_LDS_STEPS_A). - So b_next0 is read after both wave groups drained the K_ITERS - 1 b0/a0/b1 loads; the drain needs no move. - gather_kv_b_proj: only the last a0 load (N_LDS_STEPS_A >= N_LDS_STEPS_B) is at risk; b1 is read after vmcnt(0). - Comment-only: final ISA is byte-identical for the 256x256/128x256 GEMM and two gather_kv_b_proj variants.

Copilot

🔵 Needs a closer look

Cross-wave synchronization warrants human confirmation of gfx950 instruction ordering and the reported correctness and performance results.

0 open findings

1 resolved since last review

Final B0 read remains unprotected, allowing stale data

🧠 Review effort: Balanced

AndreasKaratzas · 2026-10-23T16:09:18Z

This also explains the intermittent test_gemm_a8w8 failures on MI35X CI ("ck bpreshuffle: catastrophic error" at n=7168, k=768/1024), e.g. #6691, the main nightly and this one from 10-10. I ended up at the same one-line change. With it, back-to-back launches over all 300 gfx950 tuned rows that use the 8-wave kernel (flydsl 0.4.0) gave 0 bad outputs in ~1.6M calls vs ~44k on main, and kernel time is unchanged (median 1.0006x). Main actually fails on almost every 8-wave row, not just the Kimi shapes. Would be great to get this in so CI stops flaking.

rasmith

Looks good to me for the most part, but someone with more AITER experience should probably approve, or maybe a second set of eyes should also approve.

rasmith · 2026-10-23T16:42:04Z

        # K_ITERS - 1 tile here or c10 / c11 read stale A1 data.
        a1_g2s.load(a_next1, (K_ITERS - 1) * BLOCK_K)
-        rocdl.s_barrier()
+        # Every wave loads a slice of each LDS tile, and waves 4..7 run one


Can we simplify these comments (this one and the one below) so they are more straightforward to read?

Agreed. This is moot now: I'm closing this PR because #6368 already fixes the race on main (details in my PR comment).

rasmith · 2026-10-23T16:43:32Z

+    run_perftest,
+)
+
+torch.set_default_device("cuda")


Do we also need torch.accelerator.set_device_index("cuda")? Not sure, just asking.

Not needed: torch.set_default_device("cuda") already uses the current device. This is moot now anyway, because I'm closing this PR (details in my PR comment).

rasmith · 2026-10-23T16:47:46Z

The barrier fix does seem to make sense to me.

rasmith · 2026-10-23T16:52:13Z

+]
+
+
+def run_torch(xq, wq, x_scale, w_scale, dtype=dtypes.bf16):


Could we call this ref_gemm and the one below ref_gather? It seems thats what these function are doing, the run_torch providing the reference gemm impl and the other one providing the reference gather impl.

Agreed, those are better names. This is moot now: I'm closing this PR because #6368 already fixes the race on main (details in my PR comment).

solinzby1 · 2026-10-26T03:10:23Z

fix already in main which pr 6368.

yuychang · 2026-10-27T16:30:54Z

Thanks @solinzby1, @rasmith and @AndreasKaratzas.

@solinzby1 is right: #6368 already fixes this on main. Its wait_barrier(0) at the start of the last K step drains every load, so this PR's GEMM change does nothing after merging main.

I checked with back-to-back launches on all 300 gfx950 tuned 8-wave rows (about 1.2M calls each):

Before [FlyDSL] [gfx950] Fix 8-wave a8w8 GEMM reading LDS before the last tile lands #6368: 23,060 bad results
main: 0
main + this PR: 0

The gather_kv_b_proj race does not reproduce on main either (0 bad results in 840k calls).

So I'm closing this PR. Thanks for the reviews.

yuychang added 2 commits October 14, 2026 13:31

yuychang requested review from a team and a balanced review from Copilot October 14, 2026 14:09

github-actions Bot added the FlyDSL label Oct 14, 2026

Copilot started reviewing on behalf of yuychang October 14, 2026 14:10

View session

Copilot AI reviewed Oct 14, 2026

View reviewed changes

Comment thread aiter/ops/flydsl/kernels/gemm_a8w8_8wave.py Resolved

Copilot AI balanced review requested due to automatic review settings October 14, 2026 14:54

Copilot started reviewing on behalf of yuychang October 14, 2026 14:54

View session

Copilot AI reviewed Oct 14, 2026

View reviewed changes

zufayu requested a review from yadaish October 15, 2026 02:10

zufayu assigned yadaish Oct 15, 2026

zufayu requested review from coderfeli and solinzby1 October 15, 2026 06:36

zufayu assigned coderfeli and unassigned yadaish Oct 15, 2026

yuychang added the ci:sglang label Oct 21, 2026

rasmith reviewed Oct 23, 2026

View reviewed changes

yuychang closed this Oct 27, 2026

		]


		def run_torch(xq, wq, x_scale, w_scale, dtype=dtypes.bf16):

Conversation

yuychang commented Oct 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Technical Details

Test Plan

Test Result

Submission Checklist

Uh oh!

github-actions Bot commented Oct 14, 2026

🏷️ CI Guide

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🔵 Needs a closer look

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🔵 Needs a closer look

Uh oh!

AndreasKaratzas commented Oct 23, 2026

Uh oh!

rasmith left a comment

Choose a reason for hiding this comment

Uh oh!

rasmith Oct 23, 2026

Choose a reason for hiding this comment

Uh oh!

yuychang Oct 27, 2026

Choose a reason for hiding this comment

Uh oh!

rasmith Oct 23, 2026

Choose a reason for hiding this comment

Uh oh!

yuychang Oct 27, 2026

Choose a reason for hiding this comment

Uh oh!

rasmith commented Oct 23, 2026

Uh oh!

rasmith Oct 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yuychang Oct 27, 2026

Choose a reason for hiding this comment

Uh oh!

solinzby1 commented Oct 26, 2026

Uh oh!

yuychang commented Oct 27, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

yuychang commented Oct 14, 2026 •

edited

Loading

rasmith Oct 23, 2026 •

edited

Loading