Fix two latent synchronization bugs in the SM100 DSA backward kernel

[zkyue]

Before submitting

• I agree to license this contribution under the terms of LICENSE.txt.
• I ran pre-commit run and committed any formatting changes. (clang-format skipped — no C/CUDA files; black + black-jupyter pass.)

Affected area

FE OSS kernels or CuTeDSL

Summary

dsa_bwd_sm100.py (FlashAttentionDSABackwardSm100) carries two latent
synchronization bugs. Neither changes any current numeric result, but each
corrupts gradients under a configuration the code already anticipates:

1. Staged compute→MMA store views for P and dS.
   P_smem_layout_store_staged / dS_smem_layout_store_staged are built with
   self.load_mma_K_stage instead of self.compute_mma_P_stage /
   self.compute_mma_dS_stage, and in compute() the store-side partitions
   tRS_sP / tRS_sdS are sliced at the producer's initial index once before
   the tile loop, so stsm always writes slot 0 while producer_state.advance()
   cycles the mbarrier phases and the MMA consumer reads slots 0..N-1. Dormant
   only because every compute→MMA stage count is 1 today; deepening either
   pipeline yields silent NaNs.

2. TMEM dealloc lifetime at the kernel tail.
   Compute warp 0 calls dealloc_tmem as soon as its own compute() finishes;
   nothing orders that after the reduce warps' final dKV T2R. The existing
   t2r_dKV01/4_done barriers sync reduce with the MMA warp only — compute
   warp 0 is in neither. A successor CTA can re-allocate and overwrite the
   columns while a T2R is still in flight.

Why

1. With e.g. compute_mma_dS_stage = 2 the MMA warp consumes odd tiles from a
   smem slot the compute warps never wrote; dQ/dKV pick up NaNs within one tile
   pair. Deepening these pipelines is a natural next step (overlapping
   softmax-scale work with the dKV/dQ GEMMs) and hits silent corruption with no
   hint the store paths are the cause. The optimize dsa bwd sm100 kernel #318 rewrite additionally hoisted the
   store partition out of the loop (item 3 of this bug).
2. tcgen05.dealloc frees the TMEM columns for a successor CTA on the same SM;
   if that CTA allocates and its MMA starts writing while the predecessor's T2R
   is still in flight, the drained dKV values are silently corrupted. The window
   is the kernel tail (compute finishing its dQ stores while reduce is still
   draining the last dKV tile) — timing-dependent and silent, i.e. incorrect by
   construction even though today's scheduling happens to win the race.

Fixes: (1) build both store-view layouts with their pipeline's stage count and
select the store slot from the producer state, keeping the stage == 1 slot-0
path behind cutlass.const_expr (byte-identical default); (2) a dedicated named
barrier (id 9) — reduce warps arrive after reduce_dKV() returns (all their T2R
already fenced: store_dKV runs fence_view_async_tmem_load internally, the
split t2r_dKV paths fence at their call sites), compute warp 0 arrive-and-waits
right before dealloc_tmem.

Related issues

None filed; found by code inspection + staged-configuration experiments while
porting downstream optimizations onto the DSA backward.

API and compatibility impact

No functional change to any current result.

• Bug 1: at the committed (all-1) stage settings the emitted SM100 cubin is
  byte-identical (same sha256) — zero change for every current config; only
  unblocks future work that deepens the compute→MMA pipelines (within the 227 KB
  smem budget; the 576/512 variant's existing SharedStorage assert still guards
  the over-budget case).
• Bug 2: no measurable cost at the production-like shape (S=8192, H=64, D=512,
  topk=1024, bf16, sink on) on B200 — paired interleaved A/B, 200+ samples/side,
  sides swapped; medians within noise, minimums differ 0.006%. dq bitwise;
  dkv/d_sink unchanged within fp32-atomic run-to-run jitter.

Testing

B200 (SM100), CUDA 13.3, nvidia-cutlass-dsl[cu13] 4.5.2, torch 2.13. pre-commit run
on both files: black + black-jupyter pass (clang-format n/a).

Bug 1. New L0 regression test test_DSA_sparse_attention_backward_staged_store
(deepens both compute→MMA stages to 2, requires dq bitwise == stage-1 baseline):
1 passed; fails with NaN on the unpatched parent. Standalone repro on unpatched
fe7c8b0 (S=2048 topk=512 H=64 D=512 bf16):

variant	dq NaNs	dkv NaNs	dq bitwise vs stock
stock (all stages 1)	0	0	—
compute_mma_dS_stage=2	7168	28672	no
compute_mma_P_stage=2	0	28672	yes (P feeds only dKV)
both = 2	28672	110592	no

Patched tree: every staged variant dq-bitwise vs stock, dkv/d_sink within
fp32-atomic run-to-run jitter (control in script). Default-config cubin sha256
identical before/after.

Bug 2. dq bitwise; paired A/B no measurable cost (above).

Pre-existing test_DSA_sparse_attention_backward_wrapper[...512...] cases OOM on
our heavily shared B200 (fp32 reference ~1 GiB vs <1 GiB free) identically on the
unpatched base — unrelated to this change; maintainer CI with a dedicated GPU
should pass them (stage-1 cubin is identical).

Repro script (self-contained; flips stage constants via monkeypatch, no tree edit)

```python #!/usr/bin/env python3 """Repro for the latent staged-store bugs in the SM100 DSA backward kernel (cudnn-frontend, python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm100.py).

Self-contained: needs only torch + an importable cudnn python package with
the DSA CuTe-DSL path (no other files). It does NOT modify the installed
tree: the stage constants are flipped by monkeypatching _setup_attributes,
which is exactly what future in-tree staging work would do by editing them.

variant 'stock' : stages as committed (all 1) -> reference
variant 'ds2' : compute_mma_dS_stage = 2
variant 'p2' : compute_mma_P_stage = 2
variant 'ds2p2' : both = 2

On an UNFIXED tree: ds2/p2/ds2p2 produce NaNs (dq and/or dkv) -> exit 2
On a FIXED tree : dq bitwise == stock for every variant, dkv /
d_sink within fp32-atomic run-to-run jitter -> exit 0

Judging criteria: dq is TMA-stored with a single writer per element, hence
deterministic -> bitwise gate. dkv and d_sink are fp32 atomic reductions and
are NOT bitwise-stable run to run even at fixed config -> gated on NaN count
and relative L2 against a measured stock-rerun jitter control.

Usage: CUDA_VISIBLE_DEVICES= python3 repro_staged_store_standalone.py
"""

import sys

import torch

from cudnn import DSA
import cudnn.deepseek_sparse_attention.sparse_attention_backward.dsa_bwd_sm100 as kmod
import cudnn.deepseek_sparse_attention.sparse_attention_backward._interface_sm100 as imod
import cudnn.deepseek_sparse_attention.sparse_attention_backward.api as amod

S, TOPK, H, D = 2048, 512, 64, 512 # topk/64 = 8 tiles: pipelines really cycle
DTYPE = torch.bfloat16

def make_case(seed=0):
g = torch.Generator(device="cuda").manual_seed(seed)
q = torch.randn(S, H, D, device="cuda", dtype=torch.float32, generator=g).to(DTYPE) / 10
kv = torch.randn(S, D, device="cuda", dtype=torch.float32, generator=g).to(DTYPE) / 10
dout = torch.randn(S, H, D, device="cuda", dtype=torch.float32, generator=g).to(DTYPE) / 10
sink = torch.linspace(-2.0, 2.0, H, device="cuda", dtype=torch.float32)
idx = torch.empty(S, TOPK, device="cuda", dtype=torch.int32)
for i in range(0, S, 1024):
r = torch.rand(min(1024, S - i), S, device="cuda", generator=g)
idx[i : i + r.shape[0]] = r.argsort(dim=-1)[:, :TOPK].to(torch.int32)
tl = torch.full((S,), TOPK, dtype=torch.int32, device="cuda")
scale = D ** -0.5
# chunked reference forward (produces out + KV-only natural-log lse)
out = torch.empty(S, H, D, dtype=DTYPE, device="cuda")
lse = torch.empty(S, H, dtype=torch.float32, device="cuda")
for i in range(0, S, 128):
kv_g = kv[idx[i : i + 128].long()].float()
sc = torch.einsum("chd,ckd->chk", q[i : i + 128].float(), kv_g) * scale
l = torch.logsumexp(sc, dim=-1)
p = torch.exp(sc - torch.logaddexp(l, sink.view(1, H)).unsqueeze(-1))
out[i : i + 128] = torch.einsum("chk,ckv->chv", p, kv_g).to(DTYPE)
lse[i : i + 128] = l
return q, kv, sink, dout, idx, tl, out, lse, scale

def main():
q, kv, sink, dout, idx, tl, out, lse, scale = make_case()
orig = kmod.FlashAttentionDSABackwardSm100._setup_attributes

def run(overrides):
    def patched(self):
        orig(self)
        for k, v in overrides.items():
            setattr(self, k, v)

    kmod.FlashAttentionDSABackwardSm100._setup_attributes = patched
    imod.flash_attn_bwd_sm100.compile_cache.clear()
    amod._cache_of_SparseAttentionBackwardObjects.clear()
    try:
        dq = torch.full_like(q, float("nan"))
        dkv = torch.zeros_like(kv)
        res = DSA.sparse_attention_backward_wrapper(
            q, kv, out, dout, lse, sink, idx,
            softmax_scale=scale, topk_length=tl, dq=dq, dkv=dkv)
        torch.cuda.synchronize()
        return dq, dkv, res["d_sink"].clone()
    finally:
        kmod.FlashAttentionDSABackwardSm100._setup_attributes = orig
        imod.flash_attn_bwd_sm100.compile_cache.clear()
        amod._cache_of_SparseAttentionBackwardObjects.clear()

def rel_l2(a, b):
    return ((a.float() - b.float()).norm() / b.float().norm().clamp_min(1e-30)).item()

variants = {"stock": {}, "ds2": {"compute_mma_dS_stage": 2},
            "p2": {"compute_mma_P_stage": 2},
            "ds2p2": {"compute_mma_dS_stage": 2, "compute_mma_P_stage": 2}}
results = {name: run(ov) for name, ov in variants.items()}
for name, (dq, dkv, ds) in results.items():
    print(f"[{name:>6}] dq NaNs={int(torch.isnan(dq).sum())}  dkv NaNs={int(torch.isnan(dkv).sum())}")

dq0, dkv0, ds0 = results["stock"]
assert not torch.isnan(dq0).any(), "baseline NaN: environment problem"
dq_r, dkv_r, ds_r = run({})  # jitter control
assert torch.equal(dq_r, dq0), "stock dq not deterministic run-to-run"
dkv_tol = max(10 * rel_l2(dkv_r, dkv0), 1e-5)
ds_tol = max(10 * rel_l2(ds_r, ds0), 1e-5)
print(f"[control] stock rerun: dkv relL2={rel_l2(dkv_r, dkv0):.3e} (atomic jitter baseline)")

bug = False
for name, (dq, dkv, ds) in results.items():
    if name == "stock":
        continue
    nan = int(torch.isnan(dq).sum() + torch.isnan(dkv).sum() + torch.isnan(ds).sum())
    ok = (torch.equal(dq, dq0) and nan == 0
          and rel_l2(dkv, dkv0) <= dkv_tol and rel_l2(ds, ds0) <= ds_tol)
    print(f"{name:>6} vs stock: {'OK' if ok else 'CORRUPTED'} "
          f"(dq bitwise={torch.equal(dq, dq0)}, NaNs={nan}, dkv relL2={rel_l2(dkv, dkv0):.3e})")
    bug |= not ok

print("VERDICT:", "BUG REPRODUCED" if bug else "CLEAN")
sys.exit(2 if bug else 0)

if name == "main":
main()

</details>

<!-- This is an auto-generated comment: release notes by coderabbit.ai -->
## Summary by CodeRabbit

* **Bug Fixes**
  * Improved sparse-attention backward execution to ensure shared-memory deallocation happens only after all required asynchronous reduce steps complete.
  * Fixed staged storage selection for intermediate attention data to behave consistently across different pipeline stage settings.
* **Tests**
  * Added a CUDA-only SM100 regression test covering the staged-store backward path.
  * Verifies no NaNs in gradients and checks staged results match the stage-1 baseline with tight numerical tolerances.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: b61ae3c9-b74f-4db2-bcef-fc1f87e7b5cd

📥 Commits

Reviewing files that changed from the base of the PR and between 9ae03e5 and 0f23df4.

📒 Files selected for processing (2)

• python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm100.py
• test/python/fe_api/dsa/test_DSA_sparse_attention_backward.py

🚧 Files skipped from review as they are similar to previous changes (1)

• test/python/fe_api/dsa/test_DSA_sparse_attention_backward.py

---

📝 Walkthrough

Walkthrough

dsa_bwd_sm100.py updates SM100 staged-store slot selection and coordinates TMEM deallocation with reduce-warps’ final operations. A CUDA regression test compares baseline and two-stage configurations for gradient validity and numerical parity.

Changes

SM100 sparse-attention backward

Layer / File(s)	Summary
Stage-specific store layouts python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm100.py	P and dS staged shared-memory layouts now use their corresponding compute-stage parameters.
Producer-state store slots python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm100.py	Compute-stage tensor partitions and copies select slot zero for single-stage operation and producer-state indices for multi-stage operation.
TMEM deallocation ordering and regression coverage python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm100.py, test/python/fe_api/dsa/test_DSA_sparse_attention_backward.py	A dedicated barrier delays compute warp 0 TMEM deallocation until reduce-warps finish final dKV T2R operations; the CUDA regression test compares baseline and staged configurations.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Sequence Diagram(s)

sequenceDiagram
  participant ReduceWarps
  participant tmem_dealloc_barrier
  participant ComputeWarp0
  participant TMEM
  ReduceWarps->>ReduceWarps: Complete final dKV T2R operations
  ReduceWarps->>tmem_dealloc_barrier: arrive()
  ComputeWarp0->>tmem_dealloc_barrier: arrive_and_wait()
  tmem_dealloc_barrier-->>ComputeWarp0: All reduce warps arrived
  ComputeWarp0->>TMEM: dealloc_tmem()

Loading

Possibly related PRs

• NVIDIA/cudnn-frontend#396: Both changes update TMEM synchronization around dKV T2R completion in the SM100 sparse-attention backward path.

Suggested reviewers: liujane-dev

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 22.22% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is concise and accurately summarizes the main SM100 DSA backward synchronization fixes.
Description check	✅ Passed	The description covers the required sections with clear summary, rationale, compatibility, and testing details.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

<sub>Comment @coderabbitai help to get the list of available commands.</sub>

[coderabbitai]

🧹 Nitpick comments (1)

test/python/fe_api/dsa/test_DSA_sparse_attention_backward.py (1)

209-218: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win

Exercise the P and dS stage controls independently.

Setting both stages to 2 can still pass if one path accidentally uses the other path’s stage count or producer state. Add {P: 2, dS: 1} and {P: 1, dS: 2} cases to verify the stage-specific contract.

Proposed coverage

-    dq, dkv, d_sink = run({"compute_mma_P_stage": 2, "compute_mma_dS_stage": 2})
-
-    assert not torch.isnan(dq).any() and not torch.isnan(dkv).any() and not torch.isnan(d_sink).any(), "staged store paths produced NaN gradients"
-    assert torch.equal(dq, dq_ref), "dq must be bitwise-identical between stage 1 and stage 2"
+    for overrides in (
+        {"compute_mma_P_stage": 2},
+        {"compute_mma_dS_stage": 2},
+        {"compute_mma_P_stage": 2, "compute_mma_dS_stage": 2},
+    ):
+        dq, dkv, d_sink = run(overrides)
+        assert not torch.isnan(dq).any()
+        assert not torch.isnan(dkv).any()
+        assert not torch.isnan(d_sink).any()
+        assert torch.equal(dq, dq_ref)
+        assert rel_l2(dkv, dkv_ref) < 1e-4
+        assert rel_l2(d_sink, d_sink_ref) < 1e-4

🤖 Prompt for 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.

In `@test/python/fe_api/dsa/test_DSA_sparse_attention_backward.py` around lines
209 - 218, Extend the staged backward coverage around run and the existing
stage-2 assertions by adding independent cases for compute_mma_P_stage 2 with
compute_mma_dS_stage 1, and compute_mma_P_stage 1 with compute_mma_dS_stage 2.
Validate each case against the stage-1 baseline using the same NaN, dq bitwise,
and dkv/d_sink relative-parity checks, ensuring each stage control is exercised
independently.

🤖 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.

Nitpick comments:
In `@test/python/fe_api/dsa/test_DSA_sparse_attention_backward.py`:
- Around line 209-218: Extend the staged backward coverage around run and the
existing stage-2 assertions by adding independent cases for compute_mma_P_stage
2 with compute_mma_dS_stage 1, and compute_mma_P_stage 1 with
compute_mma_dS_stage 2. Validate each case against the stage-1 baseline using
the same NaN, dq bitwise, and dkv/d_sink relative-parity checks, ensuring each
stage control is exercised independently.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 8328ecbb-7f61-4f01-8a2a-6094a816f3cc

📥 Commits

Reviewing files that changed from the base of the PR and between f3ee97b and 9ae03e5.

📒 Files selected for processing (2)

• python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm100.py
• test/python/fe_api/dsa/test_DSA_sparse_attention_backward.py

[Anerudhan]

Thanks for the contribution @zkyue . @Jie-Fang will be reviewing this.

Aside, there is a merge conflict in test/python/fe_api/dsa/test_DSA_sparse_attention_backward.py . Can you please take a look?

[zkyue]

Thanks @Anerudhan! I've rebased onto the latest develop and resolved the conflict in test/python/fe_api/dsa/test_DSA_sparse_attention_backward.py.

It was a clean union — kept the existing test_DSA_sparse_attention_backward_qh32_uses_per_query_topk_without_padding test intact and re-appended our test_DSA_sparse_attention_backward_staged_store regression test after it. The kernel changes are unchanged, and the PR is mergeable again.

[Jie-Fang]

I think there's no successor CTA in our kernel. 1 CTA/SM is in this kernel so that there's no issue even if deallocating tmem early. Could you modify the comments to not make readers confused?

But I agree that we need to add a namedbarrier to ensure the dealloc tmem happens after the store dKV.

[zkyue]

Thanks for the review! You're right that this kernel is 1 CTA/SM, so there is no successor CTA and the old wording was misleading. I've reworded the two comments (the tmem_dealloc_barrier definition and the dealloc_tmem site) to attribute the barrier to the intra-CTA ordering instead: compute warp 0's dealloc_tmem must happen after the reduce warps have drained their final dKV T2R reads (the reads that feed store_dKV), otherwise the dealloc would race those in-flight T2R loads within the same CTA. The successor-CTA framing is removed. The barrier and all logic are unchanged — comment-only in 18c0277.

[Jie-Fang]

looks good, thank you!

[Jie-Fang]

@zkyue, @Anerudhan

only a small comment, LGTM, thanks a lot. I will keep reviewing the next PR.

[Jie-Fang]

@Anerudhan I think we can merge this PR

[Anerudhan]

Thanks @Jie-Fang

[Anerudhan]

@cudnn-ci-bot run

[Anerudhan]

@cudnn-ci-bot run

[cudnn-ci-bot]

🚀 Running mirror pipeline

Branch: cudnn-gh/pr-395-18c0277
Pipeline: 58722513