frost(sdpa): support ragged stats for SM100 frost sdpa forward engine - #512
Merged
Conversation
Port of the SM120 ragged-Stats support (NVIDIA#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 NVIDIA#381. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughSM100 and SM120 SDPA forward paths now support optional THD Stats/LSE output in token-major and head-major layouts. Metadata uses one host-built upload. Workspace allocation is reduced. Zero-KV inputs use guarded kernel views. ChangesTHD Stats and LSE support
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant SDPA_API
participant HostMetadata
participant SM100_SM120_Executor
participant CallerBuffers
SDPA_API->>HostMetadata: Build cumulative THD metadata
HostMetadata->>SM100_SM120_Executor: Upload combined metadata
SDPA_API->>SM100_SM120_Executor: Pass output and optional LSE views
SM100_SM120_Executor->>CallerBuffers: Write output and packed Stats/LSE
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
vedaanta
added
cat-feature
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>
vedaanta
marked this pull request as ready for review
August 7, 2026 21:07
![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 (4)
test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py (1)
839-848: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider moving the layout axis of this L0 test to L1.
test_dsl_sm100_thd_statsis markedL0and sweeps_FLAVORS×stats_layout. Each case triggers a per-shape THD JIT compile, so the case count multiplies the L0 runtime. Keep one layout atL0and move the second layout toL1, or mark the whole testL1. The other new stats tests already cover both layouts atL1.As per coding guidelines: "Mark every new Python test with a level from
L0throughL4; keepL0tests fast and place large parameter sweeps at higher levels."🤖 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 839 - 848, Reduce the L0 runtime of test_dsl_sm100_thd_stats by moving the layout sweep to L1: retain one stats_layout case at L0 and run the other at L1, or mark the entire test L1. Preserve coverage of both token_major and head_major layouts across the existing _FLAVORS parameterization.Source: Coding guidelines
python/cudnn/sdpa/fwd/api_dsl.py (1)
976-987: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueTwo prose descriptions still list the removed sequence-length workspace chunks. This PR moved THD metadata construction host-side and dropped the
slq/slkint32 device copies, but both descriptions of the workspace layout still name them.SdpaFwdDslSm100.scratch_workspace_bytes()now reserves the metadata buffer, the O-descriptor array, and the optional sinks dummy only.
python/cudnn/sdpa/fwd/api_dsl.py#L976-L987: remove "int32 length copies" from the_execute_thddocstring list of carved buffers.test/python/sdpa/frost/test_sdpa_frontend_integration.py#L336-L336: change the inline comment from[slq32 | slk32 | meta | o_desc | sinks dummy]to[meta | o_desc | sinks dummy].🤖 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 `@python/cudnn/sdpa/fwd/api_dsl.py` around lines 976 - 987, The workspace layout documentation still references removed sequence-length copies. In python/cudnn/sdpa/fwd/api_dsl.py:976-987, update _execute_thd’s docstring to list only the metadata buffer, per-sequence O TMA descriptors, and optional sinks dummy; in test/python/sdpa/frost/test_sdpa_frontend_integration.py:336, update the inline workspace layout comment to [meta | o_desc | sinks dummy].test/python/sdpa/frost/test_sdpa_graph_analyzer.py (1)
331-332: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDerive the expected engine name from
D.The graph is built with the module-level
Dhead dim, but the assertion hardcodesengines.engine_name(512). The d512 engine declaresd_envelope=True, so it stays eligible for any smallerdthat is a multiple of 8; the assertion therefore passes without checking the flavor that actually matchesD. Useengines.engine_name(D)so the test tracksD, or pin the head dim to 512 in this test.♻️ Proposed change
- assert engines.engine_name(512) in _eligible(g) + assert engines.engine_name(D) in _eligible(g)Also applies to: 361-361
🤖 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 331 - 332, Update the expected engine-name assertion in the SDPA graph analyzer test to derive the name from the module-level head dimension D, using engines.engine_name(D) instead of hardcoding 512. Ensure the related assertion at the other occurrence also tracks D so the test validates the engine flavor matching the configured head dimension.python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py (1)
2119-2130: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
compile()docstrings omit the new LSE parameters in three of the four flavors.has_lse,lse_token_major, andlse_strideare new public compile parameters and part of the cache key. Onlyprefill_d128_f16_sm100.pydocuments them (lines 2044-2052). Copy that paragraph into each remaining flavor so the four kernels state one contract.
python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py#L2119-L2130: add the LSE-parameter paragraph to thecompile()docstring at lines 2131-2143.python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py#L1682-L1693: add the same paragraph to thecompile()docstring at lines 1694-1698.python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py#L1885-L1896: add the same paragraph to thecompile()docstring at lines 1897-1903.🤖 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 `@python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py` around lines 2119 - 2130, Update the compile() docstrings in python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py (2119-2130), prefill_d256_f16_sm100.py (1682-1693), and prefill_d512_f16_sm100.py (1885-1896) by copying the existing LSE-parameter paragraph from prefill_d128_f16_sm100.py, documenting has_lse, lse_token_major, and lse_stride consistently in all three flavors.
🤖 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 1045-1055: Update the t_kv == 0 handling in the relevant execution
paths, including the SM120 _execute_thd method, so zeroing o_buf and updating
lse_valid via copy_ or fill_ are enclosed in
_torch_stream_context(current_stream, device). Ensure every zero-KV write is
enqueued on current_stream before returning, matching the existing
_execute_mxfp8 and _execute_fp8 pattern.
In `@python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py`:
- Around line 1885-1896: Add the `# noqa: A001` suppression directly to the `def
compile(` line in the `compile` function definition; do not place it on the
return annotation line.
---
Nitpick comments:
In `@python/cudnn/sdpa/fwd/api_dsl.py`:
- Around line 976-987: The workspace layout documentation still references
removed sequence-length copies. In python/cudnn/sdpa/fwd/api_dsl.py:976-987,
update _execute_thd’s docstring to list only the metadata buffer, per-sequence O
TMA descriptors, and optional sinks dummy; in
test/python/sdpa/frost/test_sdpa_frontend_integration.py:336, update the inline
workspace layout comment to [meta | o_desc | sinks dummy].
In `@python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py`:
- Around line 2119-2130: Update the compile() docstrings in
python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py (2119-2130),
prefill_d256_f16_sm100.py (1682-1693), and prefill_d512_f16_sm100.py (1885-1896)
by copying the existing LSE-parameter paragraph from prefill_d128_f16_sm100.py,
documenting has_lse, lse_token_major, and lse_stride consistently in all three
flavors.
In `@test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py`:
- Around line 839-848: Reduce the L0 runtime of test_dsl_sm100_thd_stats by
moving the layout sweep to L1: retain one stats_layout case at L0 and run the
other at L1, or mark the entire test L1. Preserve coverage of both token_major
and head_major layouts across the existing _FLAVORS parameterization.
In `@test/python/sdpa/frost/test_sdpa_graph_analyzer.py`:
- Around line 331-332: Update the expected engine-name assertion in the SDPA
graph analyzer test to derive the name from the module-level head dimension D,
using engines.engine_name(D) instead of hardcoding 512. Ensure the related
assertion at the other occurrence also tracks D so the test validates the engine
flavor matching the configured head dimension.
🪄 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: 769fc96e-55a4-480c-a185-a681d10cef7e
📒 Files selected for processing (9)
python/cudnn/sdpa/fwd/api_dsl.pypython/cudnn/sdpa/fwd/engines.pypython/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.pypython/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.pypython/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.pypython/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.pytest/python/sdpa/frost/test_sdpa_frontend_integration.pytest/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.pytest/python/sdpa/frost/test_sdpa_graph_analyzer.py
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>
…ow 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>
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>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
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>
Collaborator
Author
|
@cudnn-ci-bot run frost |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-512-3962481 |
Anerudhan
self-requested a review
August 8, 2026 03:02
Anerudhan
approved these changes
Aug 8, 2026
Aneureka
reviewed
Aug 8, 2026
Aneureka
left a comment
Aneureka
left a comment
Member
There was a problem hiding this comment.
LGTM overall with minor comments.
…view) 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>
Collaborator
Author
|
@cudnn-ci-bot run frost |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-512-fe4f154 |
Aneureka
approved these changes
Aug 8, 2026
This was referenced Aug 8, 2026
Closed
YangXu1990uiuc
added a commit
to YangXu1990uiuc/cudnn-frontend
that referenced
this pull request
Aug 8, 2026
Four test files were carrying their PRE-NVIDIA#512 content while the production code they exercise is post-NVIDIA#512. The branch is cherry-picked onto the github develop, and the commit that consolidated the sdpa test helpers was authored against a tree from before NVIDIA#512 landed -- so the cherry-pick took the whole file, not the helper edit, and reverted NVIDIA#512's test additions with it. api_dsl.py and engines.py were untouched by that, which is why nothing looked wrong until an SM100 box ran the suite: 16 failures, all of them tests asserting the old contract against the new kernels (a stats-less SM100 graph now carves a dummy LSE, so get_workspace_size() is b*h*s*4, not 0). Restored all four from gh/develop and re-applied only what this branch meant to change: - test_sdpa_fwd_dsl_sm100 / _sm120: the local verbatim copy of _select_engine -> frost_test_utils.select_engine. - test_sdpa_frontend_integration: plan-name lookups made suffix-aware. The heuristics now name a concrete config for every entry, so a plan reads "<engine>[<knobs>]" and names.index(_FROST) raises ValueError. - test_sdpa_graph_analyzer: engines.probe() is deleted, so _eligible asks analyze_for(...)[1] is None. The ragged-Stats coverage NVIDIA#512 added (token-major and head-major layouts, zero-length sequences, the analyzer acceptance test, the strict LSE presence contract in both directions) is back verbatim.
3 tasks
Anerudhan
pushed a commit
that referenced
this pull request
Aug 9, 2026
…_sort (#528) * Take the backend's plans one heuristic mode at a time Ranking the two sides against each other needs to know which backend entries are mode-A recommendations and which are fallbacks -- "the backend's A ahead of ours, its fallbacks behind" cannot be said about one opaque list. Until now the whole thing arrived from a single create_execution_plans([A, FALLBACK]). No C++ change is needed. C++ appends each query to the same plan list, and get_execution_plan_count() already exists, so asking one mode at a time and reading the count after each gives the boundaries. Measured on a 512^3 bf16 matmul (sm90, cuDNN 9.25): A -> plans[0:15], all knob-bearing; FALLBACK -> plans[15:17], bare eng0/eng7 with no knobs; the two segments do not overlap. A mode with no configs raises, which is not a decline while another mode still has entries -- an OPENSOURCE-only query legitimately leaves the cuDNN modes empty. Only every mode failing means the backend has nothing, and then the last error is re-raised so the caller still reports why. * Move plan ranking out of the engines and into one heuristics function An engine cannot rank. It sees neither its siblings nor the backend's entries, so propose_plans could only ever order its own knobs -- and then something downstream had to merge the two sides anyway, which heuristics_sort did by concatenating and calling it ranking. All four in-tree propose_plans were the base class's default copied verbatim: the hook has never decided anything. create_execution_plans() now gathers the inputs (parsed facts, the family's offered ids, the backend's entries tagged by mode) and hands all of it to the graph's family in ONE call. What comes back IS graph.plans, position for position. An engine answers two questions: can I serve this graph (check_support), and compile me this config (build_plan). sdpa/fwd/heuristics.py is the first such hook, and it is deliberately a frame with no tuning in it: one entry per eligible cell at the config its capability row declares. Mode A and FALLBACK differ only in which backend entries they carry; OPENSOURCE is mode A without the backend's recommendation, since these cells ARE the open-source implementation. Real per-cell rules land on top. Deleted, all superseded or never used: BaseEngine.propose_plans + its 4 implementations BaseEngine.default_knobs only fed propose_plans heuristics_sort merging is part of ranking, not a step after engines/router.py entirely Router / default_router / set_router / pygraph(router=) -- policy has one home now, and decline_types moved to base.py where the engine contract already lives engines.probe() (fwd + bwd) superseded by check_support graph.engine pure alias of selected_engine, zero callers graph.from_serialized zero callers; serialize/deserialize are the pybind-era API and stay knobs=None no longer means "engine, pick for me" -- the heuristics name a concrete config. A None field survives only on an axis whose capability row declares no domain. That reading is what let one choice be made twice, once when ranking and once inside the adapter. * Update the dispatch tests to the ranking contract, and delete what it retired Ranking has one home, so a test that wants a specific order replaces heuristics.rank instead of subclassing Router. The _ranking() helper does that; it is the same monkeypatch idiom the rest of the suite already uses. Deleted rather than translated: test_set_router_frozen_after_planning the API it tested is gone test_a_claiming_engine_is_tried_before_the_backend asserted that python plans always outrank the backend, which is a per-cell measurement, not a rule. FROST coverage rides on heur_mode.OPENSOURCE instead: ask for it and any graph still landing on a backend plan is one FROST cannot serve Renamed for what they now test: test_mixed_ranking_dispatch, test_empty_ranking_output_rejected, test_mixed_ranking_backend_slot_executes, test_constructor_backends_validated_and_ranking_ids_checked. One assertion changed meaning: the backend is queried once PER MODE now, so _create_backend_plans records two create_execution_plans calls for [A, FALLBACK]. test_sdpa_graph_analyzer called engines.probe() twice; those two call analyze_for directly, so no production API exists only for tests. Six sdpa test files each carried a verbatim copy of _select_engine matching a bare engine name. Plans now read <engine>[<knobs>] because the heuristics name a concrete config for every entry, so they share frost_test_utils.select_engine, which matches on the engine. 208 passed. test_a_replayed_plan_reports_its_own_notes still fails and also fails on develop without this change -- C++ on 9.25 no longer raises for an index one past the plan count. * Update the design doc, and align the remaining sdpa test helpers The doc still described a Router with three pluggability levels, engines that propose their own plans, and heuristics_sort as the seam a cost model replaces. Rewritten to what dispatch now does: one call per graph into the family's heuristics hook, the backend's entries tagged by the mode that produced them, and heur_mode.OPENSOURCE as the way FROST coverage is measured rather than assumed. Also states plainly what register_backend is and is not. It installs an engine instance on one graph -- the hatch tests use to inject a fake. It does not make an engine rankable: an out-of-tree engine declares no Capabilities, so nothing can enumerate its configs or place it against the backend. The follow-up list now names removing that concept, since an engine id is decodable from the manifest alone. Six sdpa test files each carried a verbatim copy of _select_engine matching a bare engine name; the shared frost_test_utils.select_engine matches on the engine, which is what plan names now carry a config suffix for. 208 passed locally. The one failure is test_a_replayed_plan_reports_its_own_notes, which fails on develop without this change too. * Decode an engine id from the manifest, with nothing registered first An engine id is fully decodable from the manifest: the family owning the id block, then the slot within it. _owners_for_id only ever looked inside the graph's candidate set, so an id could be resolved only if something had already put that engine there -- which made register_backend look like a prerequisite for create_execution_plan() when it is really just one way to supply an instance. engine_for_id() closes that. _owners_for_id falls back to it, so replaying a recorded (engine_id, knobs) works on a fresh graph, including for an engine that is not a candidate for THAT graph -- there the replay is a deliberate pin, not a routing decision. A gated-off slot still resolves to None rather than being built. Groundwork for removing the out-of-tree engine concept entirely. * Remove the out-of-tree engine concept: the manifest is the only way in Every python engine now exists exactly one way. register_backend, pygraph(backends=), graph.backends and OUT_OF_TREE_ID_BASE are gone, and _candidate_engines() is the graph's family and nothing else. An out-of-tree engine could never be RANKED anyway: it declares no Capabilities, so nothing could enumerate its configs or place it against the backend. It was an entry point into the plan list, not into the decision. And being a candidate had nothing to do with fitness -- an engine was in the list because someone had registered it, so an engine that could not serve the graph was still tried, and failed at build instead of at classification. The linear_attention suites used register_backend to PIN an implementation -- cuTile rather than FROST. That is not what registration is for, and those engines are in the manifest already, so the pin is now by name and applied after planning through select_plan(): engine_utils.pin_engines() / apply_pin(). apply_pin raises when the pinned engine produced no plan, so a pin that stops working fails the first op call. The cutile conftest used to check the pin by inspecting the CANDIDATE list, which passes whether or not the pin took effect -- which is how it ran for months against whichever engine the ranking picked while the seam it pinned through was dead. That check is deleted; the pin enforces itself. heuristics: no engine sits outside a family now, so the "family-less engines go last" branch is gone and _without_a_family is _unranked -- the case it covers is a family that declares no heuristics hook, not an engine with no family. test_engine_router.py -> test_dispatch.py. It never tested a Router; it tested dispatch -- one plan list, the at-index APIs, select_plan's strict pin, one-shot planning, how a decline advances the walk, note filters reaching python plans, manifest classification, facts attachment. _offer(monkeypatch, *engines) replaces register_backend by putting the fakes in a manifest family, so the tests reach engines through the same path production does. Six tests deleted with the concept they tested -- all checked registration-time id validation, which has no subject now that engines never declare their own ids: test_register_backend_validation, test_engine_id_in_the_in_tree_region_is_rejected, test_a_registered_in_tree_engine_is_not_offered_twice, test_overlapping_declared_id_blocks_are_rejected, test_a_lying_owns_id_cannot_capture_another_engines_plans, test_constructor_backends_validated_and_ranking_ids_checked. What they protected is covered by test_family_id_blocks_are_disjoint and test_every_engine_spec_has_a_manifest_slot. BaseEngine.owns_id goes with them: zero callers, and its docstring already called it a convenience. Three tests needed real thought rather than a mechanical edit: - The "no family, no facts payload" test built a bare relu graph. relu names no family, so there is no python candidate, and the backend declines a 2-D pass-by-value tensor -- planning raised before the assertion. The claim under test is about the payload, not about the graph being servable. - The mutable-after-validate window was `not self._backends`: validate() lowers and freezes any graph the backend CAN lower, and registering an engine was the only way to skip that. With registration gone the window is exactly the ops with no backend lowering, which is what the property was always about. - test_api_signature_parity asserted {"backends", "router"} were keyword-only. Both are gone, so the assertion had no subject; what it protected is that nothing pygraph-only is POSITIONAL, which is now asserted directly. TorchMatmulEngine goes too. It reimplemented matmul, bias and relu in torch inside a dispatch test: the numeric assertions proved torch, not dispatch, and "torch_matmul" in plan names reads like something cuDNN ships. StubEngine replaces it -- same claim on the graph, no arithmetic, and it RECORDS what dispatch handed it, so the fusion test now asserts what was only implied before: every node arrives in build order, each input port resolved to the caller's storage, and the virtual intermediate carrying none. * Give the SM120 SDPA-forward cell a real tile rule, as the worked example The framework had no rule in it: every cell went to `_sole()` on each knob axis, which answers None the moment a row declares more than one value. The SM120 prefill row declares tile_ms={64,128}, tile_ns={64,128}, so its choice fell through to api_dsl's `_SM120_Q_TILES[0]` default -- the choice being made in the adapter is exactly what moving ranking out of the engines was meant to stop, and it left the frame with nothing showing how a rule is added. _sm120_tiles(facts) is that rule, and it is measured rather than invented: regret 1.009 geomean / 1.054 worst against the best of the enumerated domain. tile_n=128 always; tile_m=64 when the grid cannot fill the machine AND each CTA has enough KV tiles to amortize the extra Q-tile loop, with a causal mask counted as a halved effective grid because it halves the work per CTA. It reads facts and nothing else -- device_sm_count is already on the record. Shape a colleague can copy: write the function, list the cell in _TILE_RULE_CELLS, put the measurement in the commit. A cell absent from that set keeps the old behaviour (its row's sole point per axis), which is the honest answer when nobody has timed it. Mode A now emits the guess FIRST and the rest of the domain behind it, so a caller who autotunes has the runners-up and a caller who does not gets the best guess at index 0. FALLBACK takes the smallest tile the row admits -- the config that asks least of the device; picking real fallback configs per cell is a TODO left in the file. * Ask one function whether an SM120 tile fits, not two Naming tile_n=128 unconditionally broke D=208/224/240/256: the adapter's own `if self.tile_n is None` branch was quietly shrinking the KV tile to whatever fit SMEM, so leaving the knob None had been answering a CAPABILITY question, not a tuning one. Requesting a value skips that branch, and the request then fails the very check the branch existed to satisfy -- 106512 bytes wanted against the part's 101376. The fit arithmetic moves to config_sm120.smem_bytes(), beside the template it describes, and both callers use it: the adapter's check and the ranking's choice. The rule now reads "tile_n = the largest that fits, tile_m by occupancy", and the runners-up it offers are filtered the same way -- a config the kernel cannot fit is not a runner-up, it is an entry that sits in the list to decline at build. test_api_signature_parity asserted {"backends", "router"} were keyword-only. Both are gone, so the assertion had no subject; what it protected is that nothing pygraph-only is POSITIONAL, which is now asserted directly. * Query the backend for the modes the ranking will actually place The default mode list was written out twice -- once in _create_backend_plans, once in heuristics.default_modes. They agree today; a change to one alone would have the backend enumerate plans for a mode no family places, which reads as the family losing entries rather than as the query asking for the wrong thing. * Restore the #512 SDPA tests this branch had silently reverted Four test files were carrying their PRE-#512 content while the production code they exercise is post-#512. The branch is cherry-picked onto the github develop, and the commit that consolidated the sdpa test helpers was authored against a tree from before #512 landed -- so the cherry-pick took the whole file, not the helper edit, and reverted #512's test additions with it. api_dsl.py and engines.py were untouched by that, which is why nothing looked wrong until an SM100 box ran the suite: 16 failures, all of them tests asserting the old contract against the new kernels (a stats-less SM100 graph now carves a dummy LSE, so get_workspace_size() is b*h*s*4, not 0). Restored all four from gh/develop and re-applied only what this branch meant to change: - test_sdpa_fwd_dsl_sm100 / _sm120: the local verbatim copy of _select_engine -> frost_test_utils.select_engine. - test_sdpa_frontend_integration: plan-name lookups made suffix-aware. The heuristics now name a concrete config for every entry, so a plan reads "<engine>[<knobs>]" and names.index(_FROST) raises ValueError. - test_sdpa_graph_analyzer: engines.probe() is deleted, so _eligible asks analyze_for(...)[1] is None. The ragged-Stats coverage #512 added (token-major and head-major layouts, zero-length sequences, the analyzer acceptance test, the strict LSE presence contract in both directions) is back verbatim. * Bring the design doc to the architecture as it now stands The dispatch tree, written out: what create_execution_plans does in order, where the backend's per-mode entries come from, and where a family's rules sit. That tree was the first thing anyone asked for and the doc did not have it. Corrects three things the doc stated as settled that this PR changed: the delegating entry leads the BACKEND's block and not the family's (it falls through to native configs when the C++ OSS engine declines, so ahead of an OPENSOURCE block it answers a coverage question with a native kernel); a plan's identity is (engine_id, knobs) and never its cpp_index; whether a heuristic mode succeeded is tracked per call, not inferred from plan spans. Adds what each machine covers. The suites SKIP on the wrong arch rather than fail, so a green sweep on one box says nothing about the others -- defaulting to CUDA device 0 is how a whole SM100 run silently skips. Follow-ups now name what is actually left: one tuning rule exists, FALLBACK is a placeholder, _MEASURED_BEHIND is empty by design. * Answer the backend's plan query once per distinct config Two findings from the second review pass, both about APIs whose callers moved under the branch. Graph::create_execution_plans checks override_heuristics_query() FIRST and returns before it reads the mode at all -- deterministic SDPA backward and FP8 backward both override. Asking one mode at a time therefore appends the SAME engine-17 config once per mode, and backend_plan_entries() handed all of them back. SDPA forward's recommend() would have deduped them; SDPA BACKWARD declares no heuristics hook, so _unranked passed the duplicates straight into graph.plans and build_plans(ALL) or an autotuner would compile and time one config twice. Deduped at collection instead of in each family: a repeated (engine, knobs) in the backend's own list is never two different things, and the first index is the one whose mode span is real. test_dsl_sm100_band_right_uncovered_tail_rejected called fwd_engines.probe(). That test arrived with #485, which this branch rebased onto after probe() was already deleted here -- so it is a caller that did not exist when the deletion was written, and it would have taken out the whole Blackwell L0 suite with an AttributeError before reaching its assertion. * Address the CodeRabbit pass: an id names one engine, and two tests could not fail engine_for_id() matched an id exactly while _owners_for_id() matched a RANGE, so a replay could resolve one way for a candidate engine and another way on a fresh graph. Collapsed the other way from what was suggested: BaseEngine.id_end and owned_id_range are deleted and _owners_for_id is an equality test. The range existed so a REGISTERED engine could claim a block and registration could prove two blocks disjoint; nothing registers now, no shipped engine ever set id_end, and every range was [engine_id, engine_id + 1). Keeping it would have spread dead machinery to fix an asymmetry that only that machinery created. EngineFamily.id_end -- the family's block -- is a different thing and stays. test_ranking_and_engine_read_the_same_record declared a probe_family by hand and then called _offer(), whose own monkeypatch of MANIFEST won; the surviving family had no analyzer. It passed anyway because both sides call _facts_for(_probe_analyzer) directly, so the documented claim -- that the ranking resolves the analyzer from EngineFamily.analyzer -- went unexercised. Now declared through _offer, and it asserts the analyzer already ran BEFORE ranking, which is the part only planning can do. Verified by mutation: drop the analyzer declaration and the test fails. select_engine(tiles=) matched the rendered plan name by substring, so a request for tile_n=128 could select a tile_n=1280 plan and the test would pass having run something else. Matches PlanConfig.knobs structurally now. No caller on this branch -- the fp8 SM120 tile tests in #509 are the first, and they would have been the ones to hit it. Plus a cross-reference to a test renamed in this PR. 204 passed on the CPU suites; the one failure is the pre-existing test_a_replayed_plan_reports_its_own_notes. * License the one file this PR adds as Apache-2.0, not MIT sdpa/fwd/heuristics.py was created by copying the header from config_sm120.py, which is MIT -- so the new file inherited a tag that does not apply to it. Per LICENSING.md, the repo relicensed MIT -> Apache-2.0 in #408 and a file is kept under MIT for exactly two reasons: surviving lines from an external contributor who has not consented to relicensing, or derivation from third-party source. A file written from scratch at NVIDIA has neither, so Apache-2.0 is the correct tag -- as it already is on every other file this change adds content to (engines/heuristics.py, engines/manifest.py, sdpa/fwd/engine.py, and the tests). The MIT neighbours in sdpa/fwd are pre-existing files this PR only edits, and editing does not move a file between licenses. Also switches to the SPDX-FileCopyrightText form the Apache-2.0 files use. * Compress the comments this PR adds Six blocks broke the house rule that a call site explains only the non-obvious load-bearing fact and rationale/measurements go in the MR description -- which is where all of this already was, so it was duplicated, not lost. Cut: measurement detail from _sm120_tiles (1.5x at 64 CTAs, 240-vs-320 CTAs, 2-4%, 106 KB vs 99) down to a pointer at PR #528, keeping the two thresholds a reader needs and the warning that the rule is kernel-specific. _MEASURED_BEHIND lost two antitheses ("deliberate but NOT a measurement", "an experiment, not an edit") and a cross-reference the module docstring already makes. _owners_for_id, _create_backend_plans and the backend-dedup comment lost restated clauses. BaseEngine's note on the deleted id range stopped narrating the deletion -- that belongs to the commit that made it, where it is verbatim. 31 fewer added lines, 8 fewer comment lines; no claim, number or caveat dropped, only relocated to where it was already written. 199 passed; the one failure is the pre-existing test_a_replayed_plan_reports_its_own_notes.
This was referenced Aug 13, 2026
vedaanta
added a commit
that referenced
this pull request
Aug 13, 2026
…574) Follow-up to #512 (f16 flavors): the SM100 FP8/MXFP8 forward kernels now None-specialize their LSE argument. compile() grows has_lse (True default); has_lse=False builds a None fake-LSE and the epilogue guards the Stats store with cutlass.const_expr(lse_tensor is not None), compiling it out entirely — the amax_s/amax_o atomicMax writes are independent and unchanged. Adapter + engine plumbing: - SdpaFwdDslSm100.compile() keys has_lse on sample_lse for the fp8/mxfp8 branch; execute() drops the cached dense dummy-LSE fallback and applies the strict lse_tensor presence contract (both directions) to fp8 too; _execute_fp8/_execute_mxfp8 pass lse=None through to the kernel. - engines.py: _sm100_fp8_spec/_sm100_mxfp8_spec flip lse_optional=True. Every lower_dsl_prefill row is now lse_optional, so the engine-level dummy-LSE carve (dummy_lse_bytes + the carver.take in _execute) is dead and removed. The SM80 row keeps the False default but lowers through lower_sm80_prefill, which never read the flag. Payoff: stats-less FP8/MXFP8 inference graphs report get_workspace_size()==0 like f16, one fewer GMEM write per tile, and the strict lse execute contract is uniform across all flavors. Tests (SM100, cc 10.0): - test_sdpa_fwd_{fp8,mxfp8}_sm100.py: new stats-less zero-workspace tests (generate_stats=False, workspace==0, O/amax vs ref) and an fp8 execute lse-contract test mirroring the f16 one; 55 passed (-m "L0 or L1"). - test_sdpa_graph_analyzer.py: 78 passed. - test_sdpa_fwd_dsl_sm100.py -k graph_api: 12 passed. Fixes #523 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
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.
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 ragged THD Stats for the SM100 FROST SDPA forward engine (all four f16/bf16 flavors: d128, d192/d128, d256, d512) — the SM100 counterpart of #508 — and bring the SM100 f16 kernels to SM120 parity on
has_lse:Ragged Stats, written in the declared layout. The THD epilogue stores the caller's ragged Stats buffer directly — no extra kernels on the execute path:
No template parameter is added (matching how SM120 keys this per compile): the layout is a per-shape
compile()specialization, fully encoded in the LSE fake tensor's static layout — token-major binds its natural packed rank-2(T, H)view, head-major keeps the kernels' native rank-3(1, QH, head_stride)packing, and the epilogue branches on the static rank.has_lsespecialization (SM120 parity). The f16 kernels now None-specialize the LSE argument: a stats-less graph compiles the LSE store out, so the workspace dummy-LSE chunks disappear at every level — plain dense inference graphs reportget_workspace_size() == 0, and THD keeps only its metadata/o-desc chunks. Thelse_tensorexecute contract becomes strict in both directions (required when compiled with a Stats output, rejected when compiled without — previously an unrequested denselse_tensorwas silently written). The SM100 f16 spec advertiseslse_optional; the FP8/MXFP8 kernels still write an LSE unconditionally and keep the engine-carved dummy.THD metadata built host-side. Surfaced by the stricter carve test: the two device-side
torch.cumsumcalls each allocate scan-temp storage and launch a kernel on the execute hot path. 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 — zero per-execute CUDA allocations, asserted bytest_workspace_carve_no_per_execute_allocs_and_guards(reworked to a THD graph, since dense no longer needs a workspace at all). Applied to the SM120 adapter too, which had the same pattern.Why
Brings the SM100 engines to parity with the SM120 engine for THD +
generate_statsgraphs, and closes the known SM100has_lsefollow-up. THD/packed graphs require the Stats to be packed as well; both packed layouts in use are served (cuDNN/TE's token-major TH1 and FlashAttention's head-majorsoftmax_lse). With both engines now serving ragged Stats, the vestigialthd_statscapability axis is dropped — THD +generate_statseligibility follows fromthdandstatsalone.Related issues
Related to #381. Follow-up to #508.
API and compatibility impact
SdpaFwdDslSm100now accepts THD graphs with a stats output (both declared packed layouts); previously rejected withNotImplementedError.get_workspace_size()shrinks accordingly.lse_tensorexecute contract on the f16 flavors: passing anlse_tensorto a specialization compiled withoutsample_lsenow raisesValueError(was silently accepted). Dense stats behavior with a Stats output is unchanged.Testing
On a cc 10.0 (SM100) GPU:
pytest sdpa/frost/test_sdpa_fwd_dsl_sm100.py -m "L0 or L1" -k "thd or graph_api or stats or contract"— 111 passed (pre-existing THD/dense + the new stats/contract tests: both layouts × d128/d256/d512, SWA, GQA+sink, zero-length sequence, all-KV-zero ±sink, all-Q-zero no-op (sentinel-checked), strict execute contract).pytest sdpa/frost/test_sdpa_frontend_integration.py— 10 passed, 1 skipped (workspace-carve test reworked to THD; asserts zero per-execute CUDA allocations).pytest sdpa/frost/test_sdpa_graph_analyzer.py— 69 passed.pytest sdpa/frost/test_sdpa_fwd_{fp8,mxfp8}_sm100.py— 46 passed (fp8/mxfp8 THD stays deferred; their dummy-LSE path is unchanged).test_sdpa_fwd_dsl_sm120.pyskips on this box (no SM120 GPU); the SM120 side of the host-built-metadata change is mechanically identical to the SM100 one and is covered by CI.🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Improvements