Add sliding-window attention support to the SM120 FROST 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

Adds sliding-window attention (SWA) support to the sdpa_bwd_sm120 FROST engine

Perf data

==== bwd+swa summary (17 rows; us medians [delayed]; ratios = other/frost, >1 means frost wins) ====
shape                                  frost     cudnn       fa2  vs_cudnn    vs_fa2
B1 S4096 H16 D64 causal                400.1     558.2     434.3     1.395     1.086
B1 S4096 H16 D64 causal W511           119.3     171.0     129.4     1.432     1.084
B1 S4096 H16 D64 causal W1023          193.9     276.6     209.2     1.426     1.079
B1 S16384 H16 D64 causal              4623.0    6447.7    5011.0     1.395     1.084
B1 S16384 H16 D64 causal W511          562.6     791.8     586.3     1.407     1.042
B1 S16384 H16 D64 causal W1023         817.1    1169.2     867.5     1.431     1.062
B1 S4096 H16 D128 causal               675.5     920.7     735.8     1.363     1.089
B1 S4096 H16 D128 causal W511          216.0     312.1     241.3     1.445     1.117
B1 S4096 H16 D128 causal W1023         327.8     462.7     362.1     1.412     1.105
B1 S16384 H16 D128 causal             8664.3   11865.5    9409.3     1.369     1.086
B1 S16384 H16 D128 causal W511        1032.2    1400.4    1094.5     1.357     1.060
B1 S16384 H16 D128 causal W1023       1521.1    2083.0    1641.7     1.369     1.079
B4 S2048 H16 D64 causal                396.5     561.0     422.7     1.415     1.066
B4 S2048 H16 D64 causal W255           172.1     256.3     188.4     1.489     1.095
B4 S2048 H16 D128 causal               752.6    1061.6     804.2     1.411     1.069
B4 S2048 H16 D128 causal W255          388.3     530.0     401.2     1.365     1.033
B1 S4096 H16 D128 dense W511           815.3    1129.5     913.4     1.385     1.120
---- vs cudnn: wins 17/17, geomean 1.404x
---- vs fa2: wins 17/17, geomean 1.080x

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
24 passed, 27 warnings in 12.76s

CUDNN_FRONTEND_ENABLE_FROST_ENGINES=1 pytest test/python/sdpa/frost -q -n 4
122 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 optional left sliding-window attention support for SM120 SDPA backward operations.
  • Supports causal, non-causal, bottom-right causal, unequal-length, and tail-window configurations.
  • Rejects invalid negative window sizes with a clear validation error.

• Tests

  • Added coverage for sliding-window execution and graph eligibility across supported configurations.

[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: 78446660-fc1e-48c2-ba15-0b9902f4757f

📥 Commits

Reviewing files that changed from the base of the PR and between ab9efe1 and 21f5af3.

📒 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

---

📝 Walkthrough

Walkthrough

SM120 SDPA backward now supports optional left sliding-window masking. The change updates API propagation, validation, engine lowering, kernel masking and LSE handling, cache specialization, and coverage for causal and non-causal configurations.

Changes

SM120 sliding-window backward support

Layer / File(s)	Summary
API contracts and lowering python/cudnn/sdpa/bwd/api_dsl.py, python/cudnn/sdpa/bwd/config_sm120.py, python/cudnn/sdpa/bwd/engines.py	The API and template parameters accept window_size_left. Negative values raise ValueError. The engine advertises sliding-window support and forwards the graph window size.
Kernel masking and reconstruction python/cudnn/sdpa/bwd/kernels/bprop_f16_sm120.py	The backward kernel applies sliding-window bounds to query blocks and scores. It handles fully masked rows during LSE loading and probability reconstruction.
Reference, graph, and analyzer tests test/python/sdpa/frost/test_sdpa_bwd_dsl_sm120.py, test/python/sdpa/frost/test_sdpa_graph_analyzer.py	Reference and graph helpers pass the converted window length. Tests cover causal, non-causal, bottom-right causal, unequal-length, and tail configurations. Analyzer expectations accept supported sliding-window cases.

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

Sequence Diagram(s)

sequenceDiagram
  participant GraphAnalyzer
  participant SdpaBwdDsl
  participant SM120BackwardKernel
  GraphAnalyzer->>SdpaBwdDsl: provide facts.window_left
  SdpaBwdDsl->>SM120BackwardKernel: compile window_size_left
  SM120BackwardKernel->>SM120BackwardKernel: apply sliding-window score masking

Loading

Possibly related PRs

• NVIDIA/cudnn-frontend#498: Extends related SM120 SDPA backward masking paths through the API, engine, kernel, and tests.

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
Title check	✅ Passed	The title clearly and concisely describes the main change: adding sliding-window attention support to the SM120 FROST SDPA backward engine.
Description check	✅ Passed	The description covers the affected area, summary, issue, performance data, and exact test results, but the Why and API compatibility sections remain incomplete.
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

---

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

[Adnios]

@cudnn-ci-bot run frost

[cudnn-ci-bot]

🚀 Running mirror pipeline

Branch: cudnn-gh/pr-505-21f5af3
Pipeline: 61519634
Targets: frost

[Aneureka]

LGTM!

[Anerudhan]

What is the percentage of test_mhas_v2.py it supports?

[Adnios]

What is the percentage of test_mhas_v2.py it supports?

From the pipeline:

================================ FROST routing =================================
graphs on FROST engines: 281/2040 (13.8%) -- transition goal is all-FROST
  frost:sdpa_fwd_prefill_sm120: 281
  native:fp16-bwd: 192
  native:fp16-fwd: 1433
  native:fp8-fwd: 134