frost(sdpa): add more features like padded head dim, skv_tile=0, diagonal_band_right_bound to sm120 frost sdpa prefill kernel

[Aneureka]

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

Support three features for the SM120 FROST SDPA forward engine:

• Ragged S_kv tails served natively skv_tile=0 on the SM120 row: the sm120 prefill kernel's first (masked) KV step always covers the rightmost (and therefore any partial) tile, so no padding mask or synthesized lengths are needed.
• Head dims in multiples of 8 (TMA envelope, 8 for TMA 16B alignment requirement): the kernel compiles at head tiles rounded up to 16 (HEAD_TILE_GRANULE); envelope dims load K/V through per-chunk TMA copies over actual-extent descriptors, so pad columns hardware-zero-fill (exact zero terms in Q@K^T; O stores clip at the actual D_V). Exact head dims keep the original single-copy rank-5 load — no overhead on existing shapes.
• Causal right-band widening (diagonal_band_right_bound > 0): the causal machinery with the diagonal translated right by a compile-time R (window_right in the band model), for TOP_LEFT and BOTTOM_RIGHT alignments; composes with SWA / padding / stats / sink / THD / mixed head dims.

Why

Narrow the functional gap between the sm120 frost sdpa kernels and the native ones.

Related issues

Related to #381.

API and compatibility impact

• SdpaFwdDslSm120 now accepts: any d_qk/d_v multiple of 8 ≤ 256 (was: multiples of 16), window_size_right (with is_causal), and dense unmasked graphs with S_kv not a multiple of the KV tile
• Kernel compile() signature unchanged; window_size_right rides TemplateParams.window_right (band model)

Testing

• pytest sdpa/frost/test_sdpa_fwd_dsl_sm120.py -m "L0 or L1": 62 passed, 0 failed

Summary by CodeRabbit

• New Features

  • Expanded SM120 attention support to head dimensions divisible by 8, up to 256.
  • Added support for widening the causal attention window on the right.
  • Improved handling of non-aligned head dimensions with automatic padding.
  • Added support for ragged key/value sequence lengths and mixed head dimensions.

• Bug Fixes

  • Corrected masking, diagonal alignment, padding, statistics, and tile-boundary behavior for expanded attention configurations.
  • Improved validation for right-side window sizes and unsupported head dimensions.

[coderabbitai]

📝 Walkthrough

Walkthrough

SM120 SDPA now accepts head dimensions divisible by 8, rounds them to compile-time envelopes with TMA zero-padding, supports causal window_size_right, and handles ragged KV tails without alignment gating. Tests cover envelope dimensions, right-band masks, mixed dimensions, statistics, and dense or THD execution.

Changes

SM120 SDPA forward support

Layer / File(s)	Summary
Capability and validation contracts python/cudnn/sdpa/fwd/config_sm120.py, python/cudnn/sdpa/fwd/engines.py, python/cudnn/sdpa/fwd/api_dsl.py	SM120 head-tile constants now define the 16–256 envelope. Validation accepts positive multiples of 8 and non-negative causal right bands. Capability matching enables envelope dimensions, right-band widening, and ragged KV tails.
Head-dimension envelope and TMA loading python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py, test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py, test/python/sdpa/frost/test_sdpa_graph_analyzer.py	The kernel rounds Q/K and V/O dimensions independently. Exact dimensions use rank-5 TMA copies. Envelope dimensions use chunked rank-4 copies with zero-filled padding. Tests cover non-16-aligned dimensions and ragged KV tails.
Causal right-band execution python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py, test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py, test/python/sdpa/frost/test_sdpa_graph_analyzer.py	The kernel propagates window_size_right, shifts and widens causal diagonals, and expands KV tile coverage. Dense and THD tests cover both diagonal alignments and combined mask windows.

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

Sequence Diagram(s)

sequenceDiagram
  participant DenseDSLGraphRunner
  participant SM120FusedMultiHeadAttentionForward
  participant TMA
  participant CausalMask
  DenseDSLGraphRunner->>SM120FusedMultiHeadAttentionForward: Compile rounded head tiles and window_size_right
  SM120FusedMultiHeadAttentionForward->>TMA: Load exact or envelope-padded K/V tiles
  SM120FusedMultiHeadAttentionForward->>CausalMask: Apply shifted diagonal and right_slack
  CausalMask-->>SM120FusedMultiHeadAttentionForward: Valid column bounds
  SM120FusedMultiHeadAttentionForward-->>DenseDSLGraphRunner: Execute attention output

Loading

Possibly related PRs

• NVIDIA/cudnn-frontend#485: Both changes update SM120 causal right-band masking, validation, and related forward components.
• NVIDIA/cudnn-frontend#493: Both changes use shared SDPA engine infrastructure for head-dimension envelope capability handling.
• NVIDIA/cudnn-frontend#507: Both changes modify SM120 head-dimension handling and kernel compilation paths.

Suggested labels: mod-cutedsl

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 identifies the SM120 FROST SDPA kernel and names the main added features, including padded head dimensions, skv_tile=0, and right-band support.
Description check	✅ Passed	The description completes all required template sections and provides the feature scope, rationale, compatibility impact, related issue, and test command with results.
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.}

[Aneureka]

@cudnn-ci-bot run frost

[cudnn-ci-bot]

🚀 Running mirror pipeline

Branch: cudnn-gh/pr-531-c084afb
Pipeline: 61900149
Targets: frost

[coderabbitai]

🧹 Nitpick comments (1)

test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py (1)

129-134: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Make the reference reject the contradictory is_causal + window_size_right combination.

_apply_mask_kwargs treats a non-None window_size_right as a band and ignores is_causal. The reference applies both bounds, so j > lim wins over j > lim + R. If a future test passes both flags, the reference and the graph diverge silently and the test fails for the wrong reason.

No current test hits this. Add an assertion so the trap surfaces immediately.

♻️ Proposed guard

+    assert not (is_causal and window_size_right is not None), "a right bound makes the mask a band; do not also pass is_causal (see _apply_mask_kwargs)"
     if is_causal:
         masked = masked | (j > lim)
     if window_size_right is not None:
         masked = masked | (j > lim + window_size_right)

🤖 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/sdpa/frost/test_sdpa_fwd_dsl_sm120.py` around lines 129 - 134,
Update the reference mask logic around is_causal and window_size_right to assert
that both are not provided together, before applying either bound. Keep the
existing masking behavior unchanged for valid argument combinations.

🤖 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/sdpa/frost/test_sdpa_fwd_dsl_sm120.py`:
- Around line 129-134: Update the reference mask logic around is_causal and
window_size_right to assert that both are not provided together, before applying
either bound. Keep the existing masking behavior unchanged for valid argument
combinations.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 4be16fbe-155f-4160-9bfd-9e89e239df94

📥 Commits

Reviewing files that changed from the base of the PR and between 3fcee36 and c084afb.

📒 Files selected for processing (6)

• python/cudnn/sdpa/fwd/api_dsl.py
• python/cudnn/sdpa/fwd/config_sm120.py
• python/cudnn/sdpa/fwd/engines.py
• python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py
• test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py
• test/python/sdpa/frost/test_sdpa_graph_analyzer.py

[Adnios]

LGTM