frost(sdpa): GQA/MQA, padding mask, right-band widening, and sink-token gradients for the SM120 f16 backward - #557
Merged
vedaanta merged 8 commits intoAug 13, 2026
Conversation
Adnios
added
cat-feature
Contributor
|
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 (4)
🚧 Files skipped from review as they are similar to previous changes (4)
📝 WalkthroughWalkthroughThe SM120 SDPA backward path adds GQA/MQA support, per-batch Q/KV sequence lengths, right-band masks, sink tokens, deterministic reduction, workspace handling, engine wiring, tests, and documentation. ChangesSM120 SDPA backward extensions
Estimated code review effort: 4 (Complex) | ~60 minutes Mergeability Score: 🟡 Moderate · up to This PR adds sink-token gradients, per-batch padding support, and expanded operand handling, but the current implementation can produce incorrect sink gradients for padded queries or fail with invalid/incompatible sink operands. The PR is not merge-ready until these bounded correctness and runtime issues are fixed or explicitly accepted. Sequence Diagram(s)sequenceDiagram
participant SdpaBinding
participant SdpaBwdDsl
participant SM120BackwardKernel
participant GQAReduction
SdpaBinding->>SdpaBwdDsl: Forward sequence lengths and sink tensors
SdpaBwdDsl->>SM120BackwardKernel: Launch specialized backward kernel
SM120BackwardKernel->>GQAReduction: Write grouped dK/dV partials
GQAReduction->>SdpaBwdDsl: Write final dK/dV outputs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
Collaborator
Author
|
@cudnn-ci-bot run frost |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-557-3048a00 |
![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: 4
🤖 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 `@docs/fe-oss-apis/attention/sdpa_bwd_sm120.md`:
- Around line 80-86: Add a suitable language identifier, such as text, to the
fenced code block containing the dot/main/reduce/cvt descriptions, without
changing its contents.
- Around line 233-234: Update the padding documentation in the attention SDPA
backward section to use the public wrapper parameter names seq_kv_lens and
seq_q_lens instead of seq_len_kv and seq_len_q, preserving the existing
required/optional and mask-composition descriptions.
- Around line 237-240: Update the GQA workspace description to size the dk_ws
and dv_ws partial buffers using padded head dimension D_padded
(head_dim_padded), not logical D. Alternatively, direct callers to
scratch_workspace_bytes() for the exact workspace requirement, while preserving
the existing element-count and layout details.
In `@python/cudnn/sdpa/bwd/api_dsl.py`:
- Around line 322-336: Update _checked_seq_lens to validate that seq_lens is on
the same CUDA device expected by the SDPA backward kernel before returning its
flattened view. Preserve the existing dtype, element-count, contiguity, and
reshape validations, and raise the established validation error with the tensor
name and device details when the check fails.
🪄 Autofix
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: 71640798-cba8-4ad3-8f02-631f06daf8ff
📒 Files selected for processing (7)
docs/fe-oss-apis/attention/sdpa_bwd_sm120.mdpython/cudnn/sdpa/bwd/api_dsl.pypython/cudnn/sdpa/bwd/config_sm120.pypython/cudnn/sdpa/bwd/engines.pypython/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.pytest/python/sdpa/frost/test_sdpa_bwd_dsl_sm120.pytest/python/sdpa/frost/test_sdpa_graph_analyzer.py
Collaborator
Author
|
@cudnn-ci-bot run frost |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-557-57ebcc6 |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 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/sdpa/bwd/api_dsl.py`:
- Around line 275-286: The sink descriptor validation near lines 275-286 must
require every non-None sink descriptor to use self.q_desc.device, while
preserving the existing shape, contiguity, and dtype checks. In
python/cudnn/sdpa/bwd/api_dsl.py lines 579-584, update the cache key
construction to include _tensor_signature(sink_token) whenever sink_token is
present, so cached APIs distinguish sink dtype, stride, and device; both sites
require direct changes.
In `@python/cudnn/sdpa/bwd/engines.py`:
- Around line 341-342: Update the lowering around sample_dsink and
check_support() to validate dSink_token’s dtype before kernel execution. Ensure
non-float32 facts.dsink_t descriptors are rejected, either by adding the
corresponding graph dtype check or by deriving sample_dsink’s dtype from
facts.dsink_t while preserving the fp32 kernel requirement.
In `@python/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py`:
- Around line 1655-1659: Update _dsink_kernel to accept optional seq_q_lens and
skip rows where q >= seq_q_lens[batch] before accumulating dSink. Propagate the
same optional specialization through the dSink compile and launch signatures,
preserving existing behavior when sequence lengths are absent.
In `@test/python/sdpa/frost/test_sdpa_graph_analyzer.py`:
- Around line 1029-1035: Add the pytest.mark.L0 decorator before
test_bwd_probe_accepts_right_band_widening and test_bwd_probe_accepts_sink in
test/python/sdpa/frost/test_sdpa_graph_analyzer.py at lines 1029-1035 and
1059-1064, respectively; no other changes are needed.
🪄 Autofix
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: 60dd6c33-1628-476b-ac02-6962ea19882a
📒 Files selected for processing (7)
docs/fe-oss-apis/attention/sdpa_bwd_sm120.mdpython/cudnn/sdpa/bwd/api_dsl.pypython/cudnn/sdpa/bwd/config_sm120.pypython/cudnn/sdpa/bwd/engines.pypython/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.pytest/python/sdpa/frost/test_sdpa_bwd_dsl_sm120.pytest/python/sdpa/frost/test_sdpa_graph_analyzer.py
🚧 Files skipped from review as they are similar to previous changes (1)
- docs/fe-oss-apis/attention/sdpa_bwd_sm120.md
Adnios
marked this pull request as draft
August 13, 2026 04:11
This was referenced Aug 13, 2026
Adnios
changed the title
Adnios
changed the title
Adnios
marked this pull request as ready for review
August 13, 2026 06:58
Collaborator
Author
|
@cudnn-ci-bot run frost |
1 similar comment
Collaborator
Author
|
@cudnn-ci-bot run frost |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-557-fa63cb2 |
vedaanta
approved these changes
Aug 13, 2026
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
3 participants
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.
Before submitting
pre-commit runand committed any formatting changes.cat-*, one or moremod-*, and oneorig-*(see label list).Affected area
Summary
GQA / MQA (H_q a multiple of H_kv, including H_kv == 1)
Padding mask (per-batch actual sequence lengths)
seq_len_q/kvRight-band widening (diagonal_band_right_bound > 0)
Sink-token backward (dSink_token output)
_dsink_kernelto summing over every batch b and query row q and returndSink_tokenoutputdsink[h] = -sum_{b,q} exp(sink[h] - LSE[b,h,q]) * delta[b,h,q]Why
Related issues
#381
API and compatibility impact
Testing
Summary by CodeRabbit
New Features
Documentation
Tests