frost(sdpa): strict lse/sink/seq-lens execute contract, no torch.empty in execute - #484
Merged
Aneureka merged 5 commits intoAug 6, 2026
Merged
Conversation
Init-time flags are compile-time kernel specializations, so execute must match them exactly — silently substituting a zeros sink changes the softmax denominator (a zero logit still adds exp(0) mass), and sinks passed to a sink-less kernel were silently dropped. Both directions now raise, on SM100 and SM12x alike. The SM12x kernel None-specializes lse and sinks (cutlass.const_expr): with no sample_lse the LSE store is compiled out and no buffer — dummy or otherwise — is bound. The new Capabilities.lse_optional flag lets lower_dsl_prefill skip the dummy-LSE workspace carve for such adapters (dense stats-less SM12x workspace is now 0), and the SM12x THD scratch drops its packed-LSE and sinks-dummy chunks. The SM100 kernels still write an LSE unconditionally (has_lse specialization there is a follow-up), so the standalone stats-less path keeps a dummy — but cached once per device instead of torch.empty per execute; the FROST dispatch path continues to carve it from the caller's workspace. Caller-provided lse/sinks are now shape/dtype-validated at execute; lse_tensor must additionally be contiguous, because the kernel writes through the bound view and a silent reshape copy would swallow the output. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The implicit sinks.to(torch.float32) allocated and launched a cast kernel on the execute hot path (and broke CUDA-graph pointer stability) whenever a caller passed non-fp32 sinks. The kernels only consume fp32 sink logits, so make that a contract: _checked_sinks_1d validates dtype (+ contiguity, since the flat (H_q,) rebind must stay a view) and the graph analyzer rejects non-fp32 sink tokens up front, keeping dispatch failures at check_support instead of execute. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…never cast Same treatment as sinks: the implicit seq_lens.reshape(-1).to(torch.int32) allocated and launched a cast kernel on the execute hot path for any non-int32 or non-contiguous input, and presence was never checked against the compiled specialization — a seq_kv_lens_present kernel executed without lengths ran on a zeros dummy, masking every row (silently wrong output), while lengths passed to a specialization compiled without them were silently ignored. _checked_seq_lens validates dtype/numel/contiguity and binds a true view; _check_seq_lens_contract enforces presence both ways (THD requires both tensors outright — they source the packed cu_seqlens metadata); the graph analyzer rejects non-int32 seq_len tensors at probe time. The dispatch lowering forwards seq_q only when the compiled specialization consumes it, keeping the FP8/MXFP8 padded-Q-trim gap (not plumbed in those kernels) a documented drop instead of a new execute-time error. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ule 1: execute() validates, never converts or allocates First entry of a growing, citable rule set distilled from review findings. Rule 1 captures the lse/sink/seq-lens execute-contract work: strict init/execute specialization matching, checked views instead of implicit .to()/reshape copies, and workspace carving instead of per-execute torch.empty. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
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 (3)
🚧 Files skipped from review as they are similar to previous changes (2)
📝 WalkthroughWalkthroughSDPA forward execution now validates optional tensor contracts, rejects implicit conversions, and supports LSE-free SM120 specialization. SM100 retains dummy-LSE handling where required. Tests cover dtype, specialization, workspace, and execution behavior. ChangesSDPA forward execution contracts
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant SdpaFwdDslSm120
participant SM120FusedMultiHeadAttentionForward
participant SM120Kernel
Caller->>SdpaFwdDslSm120: execute optional tensors
SdpaFwdDslSm120->>SdpaFwdDslSm120: validate specialization and tensor contracts
SdpaFwdDslSm120->>SM120FusedMultiHeadAttentionForward: bind validated views
SM120FusedMultiHeadAttentionForward->>SM120Kernel: launch with optional LSE
SM120Kernel-->>Caller: attention output and optional LSE
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
vedaanta
marked this pull request as ready for review
August 5, 2026 05:30
![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: 2
🧹 Nitpick comments (1)
test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py (1)
312-342: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd successful optional-binding coverage.
The failure assertions do not test valid sink binding. A regression that rejects every sink or misbinds a valid FP32 sink will pass this test.
Execute the sink-enabled API with
lse_tensor=lseandsinks=sink. Compareowith_ref_sdpa_full(..., sinks=sink). Also execute the API created withoutsample_lsewith a compatiblelse_tensor, as documented in this test.Proposed test additions
with pytest.raises(ValueError, match="sinks must be float32"): api.execute(q_tensor=q, k_tensor=k, v_tensor=v, o_tensor=o, lse_tensor=lse, sinks=sink.to(torch.bfloat16)) + api.execute(q_tensor=q, k_tensor=k, v_tensor=v, o_tensor=o, lse_tensor=lse, sinks=sink) + torch.cuda.synchronize() + torch.testing.assert_close( + o, + _ref_sdpa_full(q, k, v, scale=scale, is_causal=True, sinks=sink), + atol=5e-2, + rtol=3e-2, + ) + api = SdpaFwdDslSm100(sample_q=q, sample_k=k, sample_v=v, sample_o=o, is_causal=True) assert api.check_support() api.compile() + api.execute(q_tensor=q, k_tensor=k, v_tensor=v, o_tensor=o, lse_tensor=lse)🤖 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_sm100.py` around lines 312 - 342, Extend the successful execution coverage for SdpaFwdDslSm100: after the existing sink validation, execute the sink-enabled api with lse_tensor=lse and sinks=sink, then compare o against _ref_sdpa_full using sinks=sink. Also execute the specialization created without sample_lse with a compatible lse_tensor, preserving its existing no-LSE path coverage.
🤖 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/fwd/api_dsl.py`:
- Around line 783-812: Reject non-null lse_tensor for THD in the SM100 execute
path before any LSE view or execution setup, since THD stats output is
unsupported. Apply the same validation to non-null sample_lse in the SM120
adapter’s execute path, alongside its existing execute-time checks. Ensure THD
compile-time has_lse=False is enforced consistently in both directions: THD
cannot receive stats output, and non-THD behavior remains unchanged.
In `@test/python/sdpa/frost/test_sdpa_graph_analyzer.py`:
- Around line 233-252: Add an appropriate L0 test-level marker to both newly
added graph-analysis tests, including test_probe_rejects_non_int32_seq_len and
the test at the referenced additional location, since neither requires GPU
execution.
---
Nitpick comments:
In `@test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py`:
- Around line 312-342: Extend the successful execution coverage for
SdpaFwdDslSm100: after the existing sink validation, execute the sink-enabled
api with lse_tensor=lse and sinks=sink, then compare o against _ref_sdpa_full
using sinks=sink. Also execute the specialization created without sample_lse
with a compatible lse_tensor, preserving its existing no-LSE path coverage.
🪄 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: 92b0446a-3a4e-4e3d-afd6-ded4535cc9d4
📒 Files selected for processing (9)
python/cudnn/AGENTS.mdpython/cudnn/frost/README.mdpython/cudnn/sdpa/fwd/api_dsl.pypython/cudnn/sdpa/fwd/engines.pypython/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.pypython/cudnn/sdpa/graph_analyzer.pytest/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.pytest/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.pytest/python/sdpa/frost/test_sdpa_graph_analyzer.py
Comment on lines
+233
to
+252
| def test_probe_rejects_non_int32_seq_len(): | ||
| """The kernels consume per-batch lengths as int32 directly — no implicit | ||
| cast anywhere on the execute path — so an int64 seq_len is ineligible.""" | ||
| g = _mk_graph() | ||
| q, k, v, dims, strides = _mk_qkv(g) | ||
| seq_kv = g.tensor(dim=(B, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT64, name="seq_kv64") | ||
| o, _ = g.sdpa( | ||
| name="s", | ||
| q=q, | ||
| k=k, | ||
| v=v, | ||
| attn_scale=0.1, | ||
| is_inference=True, | ||
| use_padding_mask=True, | ||
| seq_len_kv=seq_kv, | ||
| ) | ||
| _finish_output(o, dims, strides) | ||
| assert not _eligible(g) | ||
|
|
||
|
|
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Add a test level marker to both new tests.
Mark each test with an L0-L4 level. These graph-analysis tests should use L0 if they do not require GPU execution.
As per coding guidelines: “Mark every new Python test with a level from L0 through L4.”
Also applies to: 346-359
🤖 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_graph_analyzer.py` around lines 233 - 252,
Add an appropriate L0 test-level marker to both newly added graph-analysis
tests, including test_probe_rejects_non_int32_seq_len and the test at the
referenced additional location, since neither requires GPU execution.
Source: Coding guidelines
2 tasks
Aneureka
approved these changes
Aug 5, 2026
Aneureka
left a comment
Aneureka
left a comment
Member
There was a problem hiding this comment.
LGTM overall; left a minor comment regarding the behavior for user-provided LSE tensor when using THD.
…ng them THD stats are not plumbed: the kernels emit a head-major packed (1, H, T) LSE into api-level scratch, which does not match cuDNN's ragged Stats contract, so a caller's sample_lse/lse_tensor was never written. SM100 was doubly wrong — it REQUIRED lse_tensor when sample_lse was given, then dropped it. check_support now rejects thd + sample_lse up front on both adapters, and the SM100 execute THD branch raises on a provided lse_tensor (matching SM120). Dispatch is unaffected: Capabilities.thd_stats=False already keeps THD+generate_stats graphs from reaching the adapters. Addresses Haobin's review comment on PR NVIDIA#484. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Collaborator
Author
|
@cudnn-ci-bot run |
Member
|
@cudnn-ci-bot run frost |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-484-4fa6afb |
Member
|
No new failures found; will merge the PR. Thanks! |
3 tasks
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
4 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.
What
Replaces GitLab MR !2305 (closed in favor of this PR).
execute()must match the init-time specialization exactly, and must never allocate or convert on the hot path.ValueError.cutlass.const_expr): nosample_lse→ the LSE store is compiled out, no buffer bound at all. NewCapabilities.lse_optionallets the lowering skip the dummy-LSE workspace carve (dense stats-less SM120 workspace: b·h·s·4 → 0); SM120 THD scratch drops its packed-LSE and sinks-dummy chunks.has_lsespecialization is a follow-up), so the standalone stats-less path keeps a dummy — cached once per device instead oftorch.emptyper execute. Dispatch path unchanged (workspace carve)._checked_lse_view/_checked_sinks_1d/_checked_seq_lensenforce dtype (fp32 sinks/LSE, int32 seq lens), shape, and contiguity, and bind true views — the implicit.to(...)/reshape-copy kernel launches are gone. The graph analyzer rejects non-fp32 sink tokens and non-int32 seq lens at probe time.python/cudnn/AGENTS.md: starts a numbered, growing hard-rules list for agents; Rule 1 codifies this execute-contract discipline.Testing
SM100 box (cc 10.0): full frost SDPA fwd suites (sm100 + sm120 + fp8 + mxfp8 + integration + analyzer) — 354 passed, 37 skipped, 0 failed (one pre-existing environment-only ranking failure deselected; it fails identically on the unmodified base). New contract tests:
test_dsl_sm100_execute_sink_lse_contract(passes on SM100),test_dsl_sm120_execute_contract_mismatches(needs SM120 HW), analyzer rejection tests for non-fp32 sinks / non-int32 seq lens. All SM120 sink×lse specializations + THD JIT-compile clean.🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Bug Fixes
Tests