Fix latent TMEM WAR race in the SM100 DSA backward dKV drain (head_dim 576) - #396
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughThe backward kernel adds dKV2/dKV3 T2R completion barriers, updates synchronization around MMA and loop transitions, and changes dKV reductions to use registers before global atomic reduction. ChangesSparse attention backward TMEM synchronization
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant MMA
participant TMEM
participant T2R
participant reduce_dKV
MMA->>TMEM: overwrite dKV columns
reduce_dKV->>T2R: load dKV values
T2R->>reduce_dKV: signal completion barrier
reduce_dKV->>reduce_dKV: reduce register values
reduce_dKV->>MMA: allow subsequent overwrite
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
Anerudhan
added
cat-bug
|
@cudnn-ci-bot run |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-396-ef50ddc |
|
Thanks for helping fix the TMEM WAR race when head_dim=576. During my review, I realized that head_dim=512 also has such issue. My comments are all about the head_dim=512, could you help fix it in your PR? Your fix for head_dim=576 is good for me! |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In
`@python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm100.py`:
- Around line 2201-2209: Remove the duplicate unconditional dKV0/dKV1 reduction
block following the preceding branch, including its barrier arrival. Refactor
that branch into one unified register-based path that performs t2r_dKV, the TMEM
fence, barrier synchronization, and reduce_dKV_from_reg exactly once for both
values.
- Around line 1651-1655: Remove the stray outer if cutlass.const_expr(not
self.same_hdim_kv) preceding the “Gemm dKV = dO @ P part2” comments, leaving
only the intended conditional that follows them. Ensure the surrounding block
has valid indentation and parsing.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: c50a772f-a26b-4ba2-93dc-624a5f908cf7
📒 Files selected for processing (1)
python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm100.py
|
@Jie-Fang Thanks for the thorough refactor — the split T2R + signal protocol is much cleaner, and pulling the same_hdim path onto it is a good catch (it removes the I've applied all three suggestion blocks. Batching them hit a bit of line-drift: since the first suggestion inserts lines, the anchors for the part-2 dKV2 (else branch) and the dKV0/dKV1 reduce suggestions landed slightly off, so the originals weren't fully replaced — that left a dangling The branch now matches the refactor exactly as you laid it out: every reduce site loads dKV into registers, fences, signals the MMA warp that the TMEM columns are free, then does the global atomic_add, with same_hdim and not-same_hdim on the same protocol. (Minor note: |
|
LGTM, thanks! @Anerudhan this PR can be merged after the #395 |
|
@cudnn-ci-bot run |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-396-8578906 |
…6 path) In the not-same_hdim_kv path (head_dim 576 / head_dim_v 512), the reduce warps read dKV2/dKV3 from TMEM inside store_dKV, after their t2r_dKV4_done arrive. The MMA warp's next-iteration dKV0/dKV1 gemms overwrite the same TMEM columns (tmem_dKV2_offset == tmem_dKV0_offset, tmem_dKV3_offset == tmem_dKV1_offset) with no happens-before edge to those reads: mma_reduce_dKV_pipeline has 2 stages, so the part1 producer_acquire of generation 3i+3 only orders against the consumer_release of generation 3i+1 (the dKV4 generation), and no named barrier separates part2's TMEM reads (generation 3i+2) from the issue of the next part1's overwrites. Correctness currently depends on the tensor core still draining queued work when the reduce warps issue their loads. The window reproduces deterministically under adversarial delay: a 1 us spin inserted in the reduce warps between the two part2 store_dKV calls (no other change) corrupts dkv columns [384:512) on every run at S=2048 H=64 D=576 topk=512 bf16 (rel L2 vs fp32 autograd 1.0e0 vs 5.6e-3 baseline); a 10 us spin before both calls corrupts [256:384) as well. Delays at already-synchronized points (after part1's register-staged T2Rs, or before the MMA warp's dKV0 issue) never corrupt. Fix, mirroring part1's existing register-staging pattern: T2R dKV2/dKV3 into registers, fence, arrive on a new t2r_dKV23_done named barrier (id 9), and only then run the global-memory atomic reduction; the MMA warp waits on that barrier before issuing the next iteration's dKV0 (skipped on the first iteration, balanced by a final arrive after the loop, like t2r_dKV4_done). Register peak is unchanged (part1 already holds two fragments). The same_hdim_kv (head_dim 512) path is untouched and compiles identically. Measured on B200 (S=8192 H=64 topk=1024 bf16, ABAB paired): head_dim 576: 5.398 -> 5.840 ms median (+8.0%), the cost of no longer overlapping the next iteration's dKV issue with the dKV2/dKV3 readback; head_dim 512: +0.03% (noise). dq stays bitwise-identical to develop on both head dims; dkv/d_sink match the fp32 autograd reference at the baseline rel L2 on both, and the adversarial-delay matrix that trips develop is clean on the fixed kernel. Signed-off-by: zky <51477259+zkyue@users.noreply.github.com>
Co-authored-by: Jie Fang <jief@nvidia.com> Signed-off-by: zky <51477259+zkyue@users.noreply.github.com>
Completes the remove/relocate edits from the code review that the inline suggestion blocks could not express on their own: - NVIDIA#2: remove the now-duplicate same_hdim t2r_dKV4 wait before dKV1 - NVIDIA#4: remove the now-duplicate same_hdim t2r_dKV01 wait before dKV3 - NVIDIA#6: switch same_hdim dKV2/dKV3 drain to split T2R + reduce_dKV_from_reg Also repairs line-drift from batch-applying the earlier suggestions (their anchors shifted after the first suggestion inserted lines): - NVIDIA#3: remove a leftover dangling `if cutlass.const_expr(not self.same_hdim_kv):` with no body (IndentationError) left because the original lines were not replaced - NVIDIA#5: remove the leftover old dKV0/dKV1 if/else block that sat before the new unified block, which would otherwise double the reduce/atomic_add and the t2r_dKV01 barrier arrival Final state is the intended split-T2R + signal drain refactor: every reduce site loads dKV into registers, fences, signals the MMA warp that the TMEM columns are free, then does the global atomic_add; same_hdim and not-same_hdim now follow the same protocol. dq is bitwise-identical to ef50ddc; head_dim 576 and 512 both verified race-free (compute-sanitizer racecheck 0). Signed-off-by: zky <51477259+zkyue@users.noreply.github.com>
8578906 to
b980ba7
Compare
|
Hi @Anerudhan — done. I've rebased The conflict was a Re-validated on SM100 for both head_dim 512 and 576: dKV/d_sink match both post-#395 |
|
Thanks @zkyue merging. |
4 participants
Before submitting
pre-commit runand committed any formatting changes. (clang-format skipped — no C/CUDA files; black passes.)Affected area
FE OSS kernels or CuTeDSL
Summary
In
FlashAttentionDSABackwardSm100.reduce_dKV, thepart2reduction readsdKV2/dKV3 out of TMEM inside
store_dKV. On thehead_dim == 576(
not same_hdim_kv) layout those columns alias dKV0/dKV1, and the nextiteration's
mma_reduce_dKVproducer_acquireonly orders against theconsumer release two generations back (the dKV4 generation) — not against these
dKV2/dKV3 reads. The MMA warp can therefore overwrite the columns before the
reduce warp has read them: a write-after-read race on TMEM that silently
corrupts dKV.
This PR reads dKV2/dKV3 into registers and arrives on a dedicated named barrier
(id 9) before the (slow) global atomic reduction, so the TMEM columns are free
for the next iteration's dKV0/dKV1 overwrite. This mirrors the T2R-then-arrive
pattern the kernel already uses in
part1. Thesame_hdim_kv(head_dim 512)path is unchanged — there dKV2/dKV3 map to distinct free TMEM columns and are
already safe (verified below).
Why
The race is timing-hidden: today's scheduling keeps the tensor-core queue
backed up so the reduce warp wins, but any perturbation that lets the MMA warp
catch up before the reads complete corrupts dKV, with no code change required.
Minimal reproduction: a 1 µs delay inserted between the two
part2store_dKVcalls on the 576 path corrupts dKV 3/3 runs (dkv relL2 5.6e-3 → 1.0), and the
corrupted TMEM column blocks match the missing-edge analysis exactly.
Related issues
Companion to #395 (two other latent synchronization bugs in the same
kernel); independent fix, different code path.
API and compatibility impact
No functional change to any current result; dq stays bitwise-identical.
This fix only touches the
not same_hdim_kv(head_dim 576) path.of closing the window — an extra commit→wake→2×T2R→barrier handshake per
iteration on the MMA warp's critical path (~0.5 µs/iter, matching the measured
delta). The stock kernel's lower 576 time is precisely the un-synchronized race.
This is not the default/production configuration; we consider correctness-first
the right call for a silent-gradient-corruption bug and are happy to discuss the
tradeoff, but wanted to flag the regression explicitly for maintainer judgment.
Testing
B200 (SM100), CUDA 13.3, nvidia-cutlass-dsl[cu13] 4.5.2, torch 2.13.
Adversarial-delay matrix, S=2048 H=64 D=576 topk=512 bf16 sink-on; control dkv
relL2 5.64e-3 (5 TMEM column blocks uniform):
Exactly and only the two read points the static edge table predicts have no
ordering flip under delay; protected blocks (b0/b1/b4) clean in all cases; dq
bitwise throughout. On the fixed kernel every previously-triggering delay is
clean (×3), and dq sha is byte-identical to stock on both 576 and 512. The 512
path is verified safe unmodified: its adversarial matrix (including the case a
naive alias reading predicts should fire) is all clean.
Adversarial-delay repro harness
Method: a `__nanosleep(DSA_PROBE_NS)` is inserted at a chosen reduce-warp read site (selected by an env key) in a copy of the kernel, and the result is compared against the fp32 autograd reference; per-128-column-block dkv relL2 isolates which TMEM columns are corrupted. dq (TMA single-writer) is the determinism control. The full probe harness + delay-injected tree are available on request.Summary by CodeRabbit