DSA: fix 576-wide d_sink, test/document top-k index semantics, B300 benchmarks#405
Open
vedaanta wants to merge 3 commits into
Open
DSA: fix 576-wide d_sink, test/document top-k index semantics, B300 benchmarks#405vedaanta wants to merge 3 commits into
vedaanta wants to merge 3 commits into
Conversation
Contributor
📝 WalkthroughWalkthroughThe changes document B300 sparse-attention benchmarks, clarify backward top-k index semantics, compute ChangesDSA backward behavior
B300 benchmark documentation
Estimated code review effort: 3 (Moderate) | ~20 minutes 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
The kernel's sum_dSink postprocess only runs when D_qk == D_v, so the
576-wide MLA latent (512 value dims + 64 RoPE dims) silently returned a
zero sink gradient. Compute it with a closed-form epilogue from tensors
already at hand:
d_sink_h = -sum_t sigmoid(sink_h - lse[t,h]) * (out . dout)[t,h]
which is the derivative of the sink-normalized softmax through the
KV-only LSE. Precision matches the kernel's own D_qk == D_v sum_dSink
path (both consume 16-bit out/dout).
Fixes the failing
test_DSA_sparse_attention_backward_wrapper[False-512-576-512-32].
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The backward kernel supports two top-k index conventions that were undocumented and untested: without topk_length, invalid topk_idxs entries are -1 sentinels (the FlashMLA convention, zero-filled rows); with topk_length, valid indices must be a compact prefix and per-row sentinel checks are skipped. Add a parity test that the same valid index set produces the same gradients through both compiled variants, and document the semantics — including that topk_length values must be >= 1 and the inert-row (padding) recipe: topk_length=1, any valid index, and a zeroed dout row. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…alysis Wall-clock sweeps for d_qk=512 (topk 512/2048) and d_qk=576 (topk 2048) across seqlens 1k-32k, plus Nsight Compute speed-of-light numbers for the main backward kernel. Key observations recorded in the README: MMA utilization is set by topk (per-CTA inner-loop depth), not sequence length; the kernel is memory-pipe-bound (top-k gather + fp32 dKV accumulation); the seqlen-32k dip is the KV working set exceeding L2; the 576-wide latent runs ~7 SOL points below 512 due to the non-128-aligned RoPE tail tiles. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
vedaanta
force-pushed
the
vagarwalla/dsa-api-layout
branch
from
July 17, 2026 21:45
acb10b1 to
336417e
Compare
vedaanta
changed the title
cudnn.dsa API layout — fwd/bwd wrappers, torch/JAX ops, benchmarks
vedaanta
marked this pull request as ready for review
July 17, 2026 22:54
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (1)
test/python/fe_api/dsa/test_DSA_sparse_attention_backward.py (1)
236-249: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd coverage for the documented inert-row contract.
The test only compares valid-prefix and sentinel-padded inputs. Add a zeroed-
doutrow withtopk_length = 1, then assert that itsdqis zero and changing its valid index does not alterdkvord_sink.🤖 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 236 - 249, Extend the test around ref_sparse_attention_forward to include a row with topk_length equal to 1 and an all-zero dout row. Assert that this row produces zero dq, then change its single valid top-k index and verify dkv and d_sink remain unchanged; preserve the existing valid-prefix and sentinel comparisons.
🤖 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 `@benchmark/dsa/README.md`:
- Around line 154-155: Update the comparison text in the README to qualify the
performance gap by sequence length: state that the 576-wide latent trails the
512-wide case by approximately 7 SOL points through 16k, widening to
approximately 10 points at 32k.
In `@docs/fe-oss-apis/dsa.md`:
- Around line 122-130: Update the “Top-k index semantics” documentation to
define every valid topk_idxs entry as satisfying 0 <= topk_idxs < total_S_kv.
Clarify that negative entries and entries greater than or equal to total_S_kv
are invalid, while preserving the existing sentinel and topk_length rules.
In `@test/python/fe_api/dsa/test_DSA_sparse_attention_backward.py`:
- Around line 251-257: Restrict the exception handling around the two
DSA.sparse_attention_backward_wrapper calls to the specific known unsupported
condition, such as ImportError, ValueError, or NotImplementedError, and remove
RuntimeError from the skip list. Unexpected runtime failures from compilation,
launch, validation, or implementation must propagate and fail the test.
---
Nitpick comments:
In `@test/python/fe_api/dsa/test_DSA_sparse_attention_backward.py`:
- Around line 236-249: Extend the test around ref_sparse_attention_forward to
include a row with topk_length equal to 1 and an all-zero dout row. Assert that
this row produces zero dq, then change its single valid top-k index and verify
dkv and d_sink remain unchanged; preserve the existing valid-prefix and sentinel
comparisons.
🪄 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: 4f4b205d-9d8d-4126-a7ca-a901b73ca57d
📒 Files selected for processing (4)
benchmark/dsa/README.mddocs/fe-oss-apis/dsa.mdpython/cudnn/deepseek_sparse_attention/sparse_attention_backward/_interface_sm100.pytest/python/fe_api/dsa/test_DSA_sparse_attention_backward.py
Comment on lines
+154
to
+155
| - The 576-wide latent runs ~7 SOL points below 512 at the same top-k: the | ||
| non-128-aligned 64-dim RoPE tail takes separate tail-tile MMAs. |
Contributor
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟡 Minor | ⚡ Quick win
Qualify the “~7 SOL points” comparison.
The 32k row is 44.3% for d_qk=512 versus 34.2% for d_qk=576, a 10.1-point gap. State that the gap is ~7 points through 16k and widens to ~10 points at 32k.
🤖 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 `@benchmark/dsa/README.md` around lines 154 - 155, Update the comparison text
in the README to qualify the performance gap by sequence length: state that the
576-wide latent trails the 512-wide case by approximately 7 SOL points through
16k, widening to approximately 10 points at 32k.
Comment on lines
+122
to
+130
| - **Top-k index semantics** — without `topk_length`, invalid `topk_idxs` | ||
| entries must be `-1` sentinels (the FlashMLA convention). With | ||
| `topk_length`, the first `topk_length[i]` entries of each row must all be | ||
| valid indices — this compact-prefix fast path skips per-row sentinel | ||
| checks, so `-1` entries inside the prefix are not allowed. `topk_length` | ||
| values must be `>= 1`; to make a query row inert (e.g. padding rows for | ||
| fixed-shape capture), use `topk_length = 1` with any valid index and a | ||
| zeroed `dout` row, which yields an exactly-zero `dq` row and no `dkv` / | ||
| `d_sink` contribution. |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Map the relevant files first
git ls-files docs/fe-oss-apis/dsa.md test/python/fe_api/dsa/dsa_reference.py
# Show the relevant section of the docs with line numbers
sed -n '100,150p' docs/fe-oss-apis/dsa.md | cat -n
# Inspect the reference helper for top-k index handling
rg -n "topk|total_S_kv|invalid|sentinel|topk_length" test/python/fe_api/dsa/dsa_reference.py -n -A 4 -B 4
# If needed, show the surrounding implementation
python3 - <<'PY'
from pathlib import Path
p = Path('test/python/fe_api/dsa/dsa_reference.py')
text = p.read_text()
for needle in ['topk_length', 'total_S_kv', 'topk_idxs']:
idx = text.find(needle)
print(f"\n== {needle} @ {idx} ==")
if idx != -1:
start = max(0, idx - 500)
end = min(len(text), idx + 1800)
print(text[start:end])
PYRepository: NVIDIA/cudnn-frontend
Length of output: 18739
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Search the docs for any other definition of top-k bounds or topk_length semantics
rg -n "topk_length|topk_max|total_S_kv|valid indices|sentinel|FlashMLA" docs/fe-oss-apis/dsa.md -n -A 3 -B 3
# Search the reference helper for any explicit assertion on topk_length bounds
rg -n "assert|clamp|min=1|max=|topk_length" test/python/fe_api/dsa/dsa_reference.py -n -A 3 -B 3Repository: NVIDIA/cudnn-frontend
Length of output: 13212
Document the top-k index range. “Valid index” should explicitly mean 0 <= topk_idxs < total_S_kv; the current wording leaves the accepted range implicit, even though the helper treats negative and >= total_S_kv entries as invalid.
🤖 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 `@docs/fe-oss-apis/dsa.md` around lines 122 - 130, Update the “Top-k index
semantics” documentation to define every valid topk_idxs entry as satisfying 0
<= topk_idxs < total_S_kv. Clarify that negative entries and entries greater
than or equal to total_S_kv are invalid, while preserving the existing sentinel
and topk_length rules.
Comment on lines
+251
to
+257
| try: | ||
| from cudnn import DSA | ||
|
|
||
| r_len = DSA.sparse_attention_backward_wrapper(q, kv, out, dout, lse, attn_sink, topk_idxs, topk_length=topk_length) | ||
| r_sen = DSA.sparse_attention_backward_wrapper(q, kv, out, dout, lse, attn_sink, idx_sentinel, topk_length=None) | ||
| except (ImportError, ValueError, NotImplementedError, RuntimeError) as e: | ||
| pytest.skip(f"Unsupported testcase: {e}") |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Do not skip unexpected RuntimeErrors.
This catches kernel compilation, launch, validation, and implementation regressions as “unsupported testcase” skips. Restrict the skip to a known unsupported condition and let unexpected runtime failures fail the test; otherwise this parity test cannot reliably guard the compiled variants.
🤖 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
251 - 257, Restrict the exception handling around the two
DSA.sparse_attention_backward_wrapper calls to the specific known unsupported
condition, such as ImportError, ValueError, or NotImplementedError, and remove
RuntimeError from the skip list. Unexpected runtime failures from compilation,
launch, validation, or implementation must propagate and fail the test.
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
1 participant
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.
Three small, independent improvements to the DeepSeek Sparse Attention backward — no API changes; the existing
cudnn.DSA.sparse_attention_backward_wrappersurface is untouched.1. Bug fix: zero
d_sinkfor the 576-wide latentThe kernel's
sum_dSinkpostprocess only runs whenD_qk == D_v, so the 576-wide MLA latent (512 value dims + 64 RoPE dims) silently returned a zero sink gradient. Fixed with a closed-form epilogue in_interface_sm100.py(d_sink_h = -sum_t sigmoid(sink_h - lse[t,h]) * (out . dout)[t,h]); the failingtest_DSA_sparse_attention_backward_wrapper[False-512-576-512-32-...]now passes.Known pre-existing issue (reproducible on clean develop, unchanged by this PR):
[True-512-576-512-32-...]fails on dq accuracy whentopk_lengthis combined withD_qk=576; kernel-level fix tracked separately.2. Test + docs: top-k index semantics
The kernel supports two index conventions that were undocumented and untested:
-1sentinels withouttopk_length(the FlashMLA convention) vs the compact-prefix fast path with it. Added a parity test that the same valid index set produces the same gradients through both compiled variants, and documented the semantics indocs/fe-oss-apis/dsa.md— including thattopk_lengthvalues must be>= 1and the inert-row (padding) recipe:topk_length=1, any valid index, and a zeroeddoutrow (verified to produce exactly-zerodqrows and nodkv/d_sinkcontribution).3. Benchmarks: B300 results with MMA-utilization analysis
Wall-clock sweeps for
d_qk=512(topk 512/2048) andd_qk=576(topk 2048) across seqlens 1k–32k, plus Nsight Compute tensor-pipe speed-of-light numbers for the main backward kernel. Highlights: MMA utilization is set bytopk(44% at topk=2048 / d512, ~37% at d576), is sequence-length invariant, the kernel is memory-pipe-bound (top-k gather + fp32 dKV accumulation), and the seqlen-32k dip is the KV working set exceeding L2.Testing (SM100-class GPU)
test_DSA_sparse_attention_backward.py: 4 passed, 1 skipped, 1 pre-existing failure (the dq issue above, identical on clean develop).Summary by CodeRabbit
Documentation
Bug Fixes
Tests