frost(spda): support ragged stats for SM120 frost sdpa forward engine - #508
Merged
vedaanta merged 4 commits intoAug 7, 2026
Merged
Conversation
Aneureka
added
cat-feature
Member
Author
|
@cudnn-ci-bot run frost |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-508-b21391f |
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 (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughSM120 THD execution now supports caller-provided ragged LSE output in token-major and head-major layouts. Runtime validation, kernel specialization, zero-KV handling, workspace behavior, and test coverage were updated. ChangesSM120 THD LSE support
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant SM120Execute
participant THDHelper
participant SM120Kernel
participant LSEBuffer
SM120Execute->>THDHelper: pass optional LSE tensor
THDHelper->>LSEBuffer: construct token-major or head-major view
THDHelper->>SM120Kernel: compile and dispatch with LSE layout
SM120Kernel->>LSEBuffer: write ragged LSE values
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![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: 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_fwd_dsl_sm120.py`:
- Line 465: In the loop over seq_q_lens and seq_kv_lens, rename the unused
unpacked variable nkv to _nkv so Ruff B007 recognizes it as intentionally
unused.
🪄 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: c5e5fd89-b614-41d4-be1c-2a9ba0b88bca
📒 Files selected for processing (5)
python/cudnn/sdpa/fwd/api_dsl.pypython/cudnn/sdpa/fwd/engines.pypython/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.pytest/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.pytest/python/sdpa/frost/test_sdpa_graph_analyzer.py
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
vedaanta
approved these changes
Aug 7, 2026
Member
Author
|
@cudnn-ci-bot run frost |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-508-05f1e2b |
3 tasks
vedaanta
added a commit
to vedaanta/cudnn-frontend
that referenced
this pull request
Aug 7, 2026
With both the SM120 engine (NVIDIA#508) and the SM100 f16 flavors serving ragged Stats, every row with thd=True also had thd_stats=True, so the dedicated gate could never fire: THD + generate_stats eligibility now follows from thd AND stats alone (the FP8/MXFP8 rows keep thd=False). A future partial bring-up that lands THD before its stats plumbing re-adds the axis with its precise decline message. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
3 tasks
vedaanta
added a commit
that referenced
this pull request
Aug 8, 2026
…#512) * frost(sdpa): support ragged stats for SM100 frost sdpa forward engine Port of the SM120 ragged-Stats support (#508) to the SM100 f16/bf16 flavors (d128, d192/d128, d256, d512). The SM100 THD kernels always write a packed LSE; it previously landed in workspace scratch in the kernels' native head-major (1, QH, T) packing. The epilogue store is now layout-aware and writes the caller's ragged Stats buffer directly in the graph's declared layout — no extra kernels on the execute path: declared layout | strides (dims (B, H, S, 1)) | kernel store ------------------|-----------------------------|--------------------------- token-major [t,h] | stride_h == 1, stride_s == H | lse[cu_q[b] + row, head] head-major [h,t] | stride_s == 1, stride_h >= T | lse[head, cu_q[b] + row] - config_sm100: new TemplateParams.thd_lse_token_major -> CFG.THD_LSE_TOKEN_MAJOR (THD-only, validated); kernels gain a lse_stride compile() shape for the head-major padded head stride (part of the per-shape cache key); the THD fake LSE drops to element alignment (user buffers only guarantee 4B). - SdpaFwdDslSm100: THD + sample_lse accepted with the same declared-layout validation as SM120; the packed-LSE workspace chunk is carved only for stats-less graphs; t_kv == 0 short-cut fills the Stats valid region with -inf (or the sink logit alone) in either layout; strict lse_tensor presence contract in both directions for THD. - engines: the SM100 f16 spec advertises thd_stats. Testing (cc 10.0): pytest test_sdpa_fwd_dsl_sm100.py -m "L0 or L1" (thd/graph_api slice: 102 passed; new stats/contract tests: 15 passed) pytest test_sdpa_graph_analyzer.py (69 passed) pytest test_sdpa_fwd_{fp8,mxfp8}_sm100.py (46 passed) Related to #381. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * frost(sdpa): has_lse specialization + per-shape stats layout for SM100 Review follow-ups on the SM100 ragged-stats port, bringing the f16 kernels to full SM120 parity on the LSE contract: - has_lse specialization: the f16 kernels (d128, d192/d128, d256, d512) now None-specialize the LSE argument — a stats-less graph compiles the LSE store out. All dummy-LSE scratch disappears: dense inference graphs report get_workspace_size() == 0 (was b*h*s*4), THD keeps only its metadata chunks, and the SM100 f16 spec advertises lse_optional. The lse_tensor execute contract is strict in both directions (an unrequested dense lse_tensor is now rejected instead of silently written). FP8/MXFP8 kernels still write an LSE unconditionally and keep the engine-carved dummy. - No template parameter for the THD stats layout (mirrors SM120's per-compile keying): TemplateParams.thd_lse_token_major and CFG.THD_LSE_TOKEN_MAJOR are gone. The layout is a per-shape compile() specialization encoded in the LSE fake tensor's static layout — token-major binds its natural packed rank-2 (T, H) view, head-major keeps the native rank-3 (1, QH, head_stride) packing, and the epilogue branches on the static rank. - THD metadata built host-side (SM100 AND SM120 adapters): the two device-side torch.cumsum calls each allocated scan-temp storage and launched a kernel per execute. The [seq_kv | cu_q | cu_k] buffer is now built on the host from the (inherent) tolist round-trip and uploaded in one H2D copy; the slq/slk workspace copies go away too. test_workspace_carve_no_per_execute_allocs_and_guards is reworked to a THD graph (dense no longer needs a workspace) and asserts zero per-execute CUDA allocations. Testing (cc 10.0): sm100 suite -m "L0 or L1" -k "thd or graph_api or stats or contract" 111 passed; frontend integration 10 passed; fp8+mxfp8+analyzer 115 passed. The SM120 suite skips locally (no SM120 GPU); its metadata change is mechanically identical and CI-covered. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * frost(sdpa): drop the vestigial thd_stats capability axis With both the SM120 engine (#508) and the SM100 f16 flavors serving ragged Stats, every row with thd=True also had thd_stats=True, so the dedicated gate could never fire: THD + generate_stats eligibility now follows from thd AND stats alone (the FP8/MXFP8 rows keep thd=False). A future partial bring-up that lands THD before its stats plumbing re-adds the axis with its precise decline message. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * frost(sdpa): route the all-KV-zero THD case through the kernel dead-row path Review follow-up (AGENTS.md Rule 1 — execute is a zero-surprise hot path): the t_kv == 0 short-cut re-implemented the kernels' dead-row semantics adapter-side with zero_/fill_/copy_ writes — surprise kernel launches and a second copy of the same semantics that can drift. Both kernels already serve dead rows (row_sum <= 0 -> O := 0 and LSE := -inf, or the sink alone), pinned by the live-launch zero-KV sequence tests. The only launch blocker was the zero-token packed K/V view (a CuTe layout mode must be > 0), so the adapters now clamp the packed KV extent to ONE never-dereferenced token — every sequence's KV tile range is empty, so no K/V load is ever issued — bound over storage the contract already guarantees: Q backs K (kh*d_qk <= t_q*qh*d_qk), O backs V (kh*d_v <= t_q*qh*d_v). Views only; the short-cut, the adapter-side fills, and the O zero-fill are gone from both the SM100 and SM120 adapters, and AGENTS.md Rule 1 gains a bullet making the no-degenerate-path-fixups expectation explicit. Testing (cc 10.0): all_kv_zero + zero_length tests 5 passed (now exercising the kernel path); sm100 suite -k "thd or graph_api or stats or contract" 111 passed; integration + analyzer 79 passed. The SM120 suite skips locally (no SM120 GPU); its change is the same mechanical transformation, CI-covered. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * frost(sdpa) tests: pin t_q == 0 as a complete no-op (all-zero seq_len_q) The adapters early-return when no query token exists anywhere (t_q == 0: the packed O/Stats have zero rows, nothing to compute or write), but no test pinned it — the zero-length coverage always kept live Q tokens. The random backend sweeps can generate the case (a B=1 batch draws a zero seq_len_q with 10% probability), so frameworks do hit it. The THD harnesses now pre-fill the O and ragged Stats storages with a sentinel (2048.0, exact in fp16/bf16/fp32): live tests still compare the kernel-written packed region against the reference, and the new test_dsl_sm1xx_thd_all_q_zero_stats (both stats layouts, live KV and all-zero KV) asserts the buffers come back untouched end to end through the graph -> engine -> adapter stack. The SM100 harness gains the same declared-extent clamp for all-zero seq_len_q that it already had for seq_len_kv (SM120's harness had both). Testing (cc 10.0): new all_q_zero tests + sentinel-affected neighbors 11 passed; full sm100 THD/stats slice 100 passed. SM120 mirror is CI-covered (no SM120 GPU locally). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * docs: tighten the AGENTS.md degenerate-path bullet Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * frost(sdpa): put the compile() noqa on the def line (review) CodeRabbit: the A001 (builtin shadowing) suppression must sit on the `def compile(` line; the three kernels that had it on the closing-paren line were suppressing nothing. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * frost(sdpa): standardize on thd_stats_head_major / lse_head_major (review) Haobin: SM100 stored thd_stats_token_major (False = its kernel's native head-major packing) while SM120 stores thd_stats_head_major (False = token-major) — each flag named after its kernel's NON-native layout. Standardize both adapters and every kernel compile() on the SM120 / contract-aligned vocabulary: thd_stats_head_major == False means token-major, matching cuDNN's TH1 ragged Stats recipe, and the SM100 kernels' compile() keywords become lse_head_major / lse_head_stride — identical signatures across all five kernels. Naming/polarity only; the layout is always derived explicitly from the graph's declared strides on every live path, so no behavior changes. Also fixes the stale [slq32 | slk32] workspace comment in the carve test. Testing (cc 10.0): sm100 stats/contract slice 17 passed. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
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
FE OSS kernels or CuTeDSL
Summary
Enable
thd_stats=Truefor the SM120 FROST SDPA forward engine. The ragged stats output is written directly in the layout the graph declares — both packed layouts in use are supported:Why
To enhance the SM120 FROST SDPA forward engine.
Related issues
Related to #381.
API and compatibility impact
SdpaFwdDslSm120now accepts THD graphs with a stats output (both declared layouts); previously rejected withNotImplementedError. Dense stats behavior unchanged.Testing
pytest sdpa/frost/test_sdpa_fwd_dsl_sm120.py -m "L0 or L1"Summary by CodeRabbit
New Features
Bug Fixes