Full causal-mask support for FROST SM120 SDPA backward engine

[Adnios]

Before submitting

• I agree to license this contribution under the terms of LICENSE.txt.
• I ran pre-commit run and committed any formatting changes.
• I added GitHub labels: one cat-*, one or more mod-*, and one orig-* (see label list).

Affected area

• FE OSS kernels or CuTeDSL

Summary

• Add top left causal mask
• Bottom-right with S_q > S_kv
• Reuse reference functions in sm120_sdpa_bwd/test/python/sdpa/fp16_ref.py

Why

Related issues

#381

API and compatibility impact

Testing

pytest test/python/sdpa/frost/test_sdpa_graph_analyzer.py -q -n 4
66 passed, 8 warnings in 5.23s

pytest test/python/sdpa/frost/test_sdpa_bwd_dsl_sm120.py -q -n 4
20 passed, 27 warnings in 12.76s

CUDNN_FRONTEND_ENABLE_FROST_ENGINES=1 pytest test/python/sdpa/frost -q -n 4
118 passed, 318 skipped, 100 warnings in 28.11s

CUDNN_FRONTEND_ENABLE_FROST_ENGINES=1 pytest "test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0" -q -n 4
=========================================== FROST routing =============================================
graphs on FROST engines: 134/554 (24.2%) -- transition goal is all-FROST
  frost:sdpa_fwd_prefill_sm120: 134
  native:fp16-bwd: 176
  native:fp16-fwd: 244
176 passed, 208 skipped, 309 warnings in 88.75s (0:01:28)

Summary by CodeRabbit

• New Features

  • Added support for top-left causal attention with differing query and key/value sequence lengths.
  • Expanded support for bottom-right causal attention when query sequences exceed key/value sequences.
  • Added configuration for selecting causal diagonal alignment.

• Bug Fixes

  • Prevented invalid gradient propagation for fully masked query rows.
  • Improved handling of sequence tails and causal masking across supported configurations.

[coderabbitai]

📝 Walkthrough

Walkthrough

Changes

SM120 causal alignment

Layer / File(s)	Summary
Causal configuration and eligibility python/cudnn/sdpa/bwd/config_sm120.py, python/cudnn/sdpa/bwd/api_dsl.py, python/cudnn/sdpa/bwd/engines.py	Adds causal_top_left, validates its use with causal mode, forwards it during compilation, and removes prior unequal-length restrictions while retaining bottom-right engine checks.
SM120 kernel masking python/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py	Selects top-left or bottom-right diagonal offsets, clamps query iterations, updates causal masking, and handles fully masked rows.
Causal alignment validation test/python/sdpa/frost/test_sdpa_bwd_dsl_sm120.py, test/python/sdpa/frost/test_sdpa_graph_analyzer.py	Uses canonical references and adds coverage for unequal lengths, mask variants, fully masked rows, and updated mismatch diagnostics.

Estimated code review effort: 4 (Complex) | ~45 minutes

Sequence Diagram(s)

sequenceDiagram
  participant SDPA_API
  participant EngineMismatch
  participant SM120Kernel
  SDPA_API->>EngineMismatch: submit causal sequence-length configuration
  EngineMismatch->>SM120Kernel: select eligible causal alignment
  SDPA_API->>SM120Kernel: forward causal_top_left
  SM120Kernel->>SM120Kernel: apply diagonal offset and masked-row handling

Loading

Possibly related PRs

• NVIDIA/cudnn-frontend#476: Introduced the SM120 SDPA backward DSL engine extended by this change.
• NVIDIA/cudnn-frontend#485: Adds related causal masking support for nonstandard sequence lengths in a different engine.
• NVIDIA/cudnn-frontend#486: Added the SM120 backward causal configuration extended by this change.

Suggested labels: cat-feature, orig-nv-eng, mod-cutedsl, mod-frost

Suggested reviewers: vedaanta

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% 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
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.
Title check	✅ Passed	The title clearly summarizes the main change: full causal-mask support for the FROST SM120 SDPA backward engine.
Description check	✅ Passed	The description covers the required sections, summarizes the changes, identifies the related issue, and lists exact test commands and results.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

🤖 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 `@test/python/sdpa/frost/test_sdpa_bwd_dsl_sm120.py`:
- Around line 280-282: Update the new test function
test_sdpa_bwd_dsl_sm120_causal_br_sq_gt_skv by adding the appropriate
pytest.mark.L0–L4 test-level marker alongside its existing decorators, following
the repository’s conventions.

🪄 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: 676f3149-0099-4bf0-abea-1c3a0c1e07f4

📥 Commits

Reviewing files that changed from the base of the PR and between f843ff8 and 597c260.

📒 Files selected for processing (6)

• python/cudnn/sdpa/bwd/api_dsl.py
• python/cudnn/sdpa/bwd/config_sm120.py
• python/cudnn/sdpa/bwd/engines.py
• python/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py
• test/python/sdpa/frost/test_sdpa_bwd_dsl_sm120.py
• test/python/sdpa/frost/test_sdpa_graph_analyzer.py

[vedaanta]

@cudnn-ci-bot run frost

[cudnn-ci-bot]

🚀 Running mirror pipeline

Branch: cudnn-gh/pr-498-597c260
Pipeline: 61417882
Targets: frost