frost(sdpa): causal right-band widening + per-sequence THD bottom-right diagonal on SM100 - #485
Merged
Merged
Conversation
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:
📝 WalkthroughWalkthroughChangesRight-banded causal masking now uses normalized window bounds across the SDPA API, engine selection, SM100 and SM120 configurations, tile bounds, and prefill kernels. THD bottom-right paths resolve per-batch query lengths. Tests cover dense, THD, bottom-right, sliding-window, eligibility, and invalid-bound cases. Right-banded causal SDPA
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant SdpaFwdDsl
participant EngineLowering
participant SM100Kernel
participant apply_mask_chunk
SdpaFwdDsl->>EngineLowering: pass window_size_right
EngineLowering->>SM100Kernel: compile normalized window bounds
SM100Kernel->>SM100Kernel: resolve effective query length
SM100Kernel->>apply_mask_chunk: apply causal mask with window_right
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 06:04
![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
🧹 Nitpick comments (1)
test/python/sdpa/frost/test_sdpa_graph_analyzer.py (1)
315-316: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winReuse
_facts(g)in the positive analyzer probe.The existing helper already checks that analysis returns a fact object and that
facts.invalidisNone. This test callsga.analyze(g)directly and skips those checks.Proposed change
- facts = ga.analyze(g) + facts = _facts(g)🤖 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 315 - 316, Update the positive analyzer probe to obtain its result through the existing _facts(g) helper instead of calling ga.analyze(g) directly, then retain the current assertions for right_band_widening, window_right, and causal.
🤖 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 552-556: Reformat the causal_covers_tail assignment in the
surrounding SDPA logic using Black’s 160-character line limit, wrapping the
boolean condition without changing its behavior.
---
Nitpick comments:
In `@test/python/sdpa/frost/test_sdpa_graph_analyzer.py`:
- Around line 315-316: Update the positive analyzer probe to obtain its result
through the existing _facts(g) helper instead of calling ga.analyze(g) directly,
then retain the current assertions for right_band_widening, window_right, and
causal.
🪄 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: 16871c1a-8e11-43d6-adfe-c528fbf96514
📒 Files selected for processing (13)
python/cudnn/frost/tile_dsl/mask.pypython/cudnn/sdpa/fwd/api_dsl.pypython/cudnn/sdpa/fwd/config_sm100.pypython/cudnn/sdpa/fwd/engines.pypython/cudnn/sdpa/fwd/kernels/_common_sm100.pypython/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.pypython/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm100.pypython/cudnn/sdpa/fwd/kernels/prefill_d128_mxfp8_sm100.pypython/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.pypython/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.pypython/cudnn/sdpa/graph_analyzer.pytest/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.pytest/python/sdpa/frost/test_sdpa_graph_analyzer.py
2 tasks
vedaanta
added
mod-frontend
This was referenced Aug 6, 2026
vedaanta
force-pushed
the
vagarwalla/frost-rbw-thd-br
branch
from
August 6, 2026 20:50
6b54d0f to
53b209e
Compare
Contributor
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
Collaborator
Author
|
@cudnn-ci-bot run frost |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-485-53b209e |
vedaanta
force-pushed
the
vagarwalla/frost-rbw-thd-br
branch
from
August 6, 2026 23:10
53b209e to
0171127
Compare
Contributor
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
Collaborator
Author
|
Pushed Root cause: develop's new Fix: same plumbing as the other five SM100 kernels — Validation on the updated revision:
|
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 `@python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py`:
- Around line 215-218: Update both pure causal helper call sites in the prefill
attention masking flow to accept and use eff_seqlen_kv, then clamp last_live
after applying CFG.BAND_RIGHT to eff_seqlen_kv - kv_col_base - 1. Preserve
existing masking behavior for aligned lengths and non-positive right bands, and
add regression coverage for a partial final KV tile with positive
CFG.BAND_RIGHT.
🪄 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: 77544783-2902-4e5c-a289-b27bf2c28bef
📒 Files selected for processing (14)
python/cudnn/frost/tile_dsl/mask.pypython/cudnn/sdpa/fwd/api_dsl.pypython/cudnn/sdpa/fwd/config_sm100.pypython/cudnn/sdpa/fwd/engines.pypython/cudnn/sdpa/fwd/kernels/_common_sm100.pypython/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.pypython/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm100.pypython/cudnn/sdpa/fwd/kernels/prefill_d128_mxfp8_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.pypython/cudnn/sdpa/graph_analyzer.pytest/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.pytest/python/sdpa/frost/test_sdpa_graph_analyzer.py
🚧 Files skipped from review as they are similar to previous changes (13)
- python/cudnn/frost/tile_dsl/mask.py
- python/cudnn/sdpa/graph_analyzer.py
- python/cudnn/sdpa/fwd/api_dsl.py
- python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py
- test/python/sdpa/frost/test_sdpa_graph_analyzer.py
- test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py
- python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py
- python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py
- python/cudnn/sdpa/fwd/kernels/_common_sm100.py
- python/cudnn/sdpa/fwd/engines.py
- python/cudnn/sdpa/fwd/config_sm100.py
- python/cudnn/sdpa/fwd/kernels/prefill_d128_mxfp8_sm100.py
- python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm100.py
Comment on lines
+215
to
+218
| if cutlass.const_expr(CFG.BAND_RIGHT != 0): | ||
| # Right-band widening: the causal upper limit sits BAND_RIGHT columns | ||
| # right of the diagonal (cuDNN diagonal_band_right_bound). | ||
| last_live = last_live + cutlass.Int32(CFG.BAND_RIGHT) |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Clamp BAND_RIGHT to the effective KV length.
If eff_seqlen_kv is not tile-aligned and CFG.BAND_RIGHT > 0, these helpers can mark columns >= eff_seqlen_kv as live. The pure causal branches call these helpers directly, so the eff_seqlen_kv handling in the generic apply_mask_chunk branch does not protect them.
Pass eff_seqlen_kv to both helpers and clamp last_live to eff_seqlen_kv - kv_col_base - 1. Add regression coverage for a partial final KV tile with a positive CFG.BAND_RIGHT.
Suggested fix
-def _apply_top_left_causal_mask_chunk(reg_S, q_abs, kv_col_base, N: int = 64):
+def _apply_top_left_causal_mask_chunk(reg_S, q_abs, kv_col_base, eff_seqlen_kv, N: int = 64):
last_live = q_abs - kv_col_base
if cutlass.const_expr(CFG.BAND_RIGHT != 0):
last_live = last_live + cutlass.Int32(CFG.BAND_RIGHT)
+ last_live = cute.math.min(last_live, eff_seqlen_kv - kv_col_base - cutlass.Int32(1))
-def _apply_bottom_right_causal_mask_chunk(reg_S, q_abs, kv_col_base, causal_diag, N: int = 64):
+def _apply_bottom_right_causal_mask_chunk(
+ reg_S, q_abs, kv_col_base, causal_diag, eff_seqlen_kv, N: int = 64
+):
last_live = q_abs + causal_diag - kv_col_base
if cutlass.const_expr(CFG.BAND_RIGHT != 0):
last_live = last_live + cutlass.Int32(CFG.BAND_RIGHT)
+ last_live = cute.math.min(last_live, eff_seqlen_kv - kv_col_base - cutlass.Int32(1))Also applies to: 236-237
🤖 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
215 - 218, Update both pure causal helper call sites in the prefill attention
masking flow to accept and use eff_seqlen_kv, then clamp last_live after
applying CFG.BAND_RIGHT to eff_seqlen_kv - kv_col_base - 1. Preserve existing
masking behavior for aligned lengths and non-positive right bands, and add
regression coverage for a partial final KV tile with positive CFG.BAND_RIGHT.
Collaborator
Author
|
@cudnn-ci-bot run frost |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-485-0171127 |
vedaanta
force-pushed
the
vagarwalla/frost-rbw-thd-br
branch
from
August 7, 2026 23:50
0171127 to
c3721ad
Compare
Collaborator
Author
|
Rebased onto develop head
window_left: Optional[int] # None = no lower bound; W keeps kv >= q-W
window_right: Optional[int] # None = no upper bound; 0 = causal; R widens
bottom_right: bool # anchors the band's diagonal
Validation on the rebased + refactored branch: analyzer + sm100 + sm120 fwd suites 464 pass / 0 fail (49 sm120-only skips on the SM100 box), black-clean. |
3 tasks
Collaborator
Author
|
@cudnn-ci-bot run frost |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-485-c3721ad |
vedaanta
force-pushed
the
vagarwalla/frost-rbw-thd-br
branch
from
August 8, 2026 01:43
c3721ad to
c779cb8
Compare
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py (1)
272-304: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winTest the exact tail-coverage boundary.
Line 272 tests
s_q + R < s_kv. Line 292 testss_q + R > s_kv. Neither case tests the required accepted boundary,s_q + R == s_kv.Use
s_kv = 232in an accepted partial-tile control case. Assert that at least one matching engine probes successfully before asserting rejection fors_kv = 200. This makes Line 304 prove that tail coverage causes the rejection.🤖 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 272 - 304, Update the SM100 band-right tests to add an accepted partial-tile control case with s_kv = 232, where s_q + R == s_kv, and assert that at least one matching engine probes successfully. Keep the existing s_kv = 200 rejection case and ensure its failed probe assertion follows the accepted control, proving rejection is caused by uncovered tail columns.
🤖 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_sm100.py`:
- Around line 272-304: Update the SM100 band-right tests to add an accepted
partial-tile control case with s_kv = 232, where s_q + R == s_kv, and assert
that at least one matching engine probes successfully. Keep the existing s_kv =
200 rejection case and ensure its failed probe assertion follows the accepted
control, proving rejection is caused by uncovered tail columns.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 0e2b460a-4414-42be-903c-13bba514dacf
📒 Files selected for processing (1)
test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py
YangXu1990uiuc
approved these changes
Aug 8, 2026
vedaanta
force-pushed
the
vagarwalla/frost-rbw-thd-br
branch
from
August 8, 2026 03:19
c779cb8 to
908c5ad
Compare
Collaborator
Author
|
Root-caused and fixed the Root cause — a rename the band-model refactor could not see: the KV-loop bounds feed read the right-band bound via a string attribute name, Fix (
Validation: local fwd_ragged back to 128/128; frost sm100 + analyzer suites green. The full local mhas_v2 run also surfaces |
Aneureka
approved these changes
Aug 8, 2026
| is_causal: bool = False, | ||
| causal_bottom_right: bool = False, | ||
| window_size_left: Optional[int] = None, | ||
| window_size_right: Optional[int] = None, |
Member
There was a problem hiding this comment.
nit: do we want to unify the naming by changing window_size_* to window_*?
Collaborator
Author
|
@cudnn-ci-bot run frost |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-485-908c5ad |
…ht diagonal on SM100 Right-band widening (diagonal_band_right_bound > 0) lowers as MASK_CAUSAL with a compile-time BAND_RIGHT offset on the diagonal, applied in both the KV-loop bounds and the per-element mask, so all five SM100 rows (f16 d128/d256/d512, fp8, mxfp8) serve it; the analyzer records the bound as facts.window_right and the KV-tail coverage rule accounts for the widened band. THD bottom-right now anchors the diagonal at each sequence's own (seq_len_q[b], seq_len_kv[b]) via a resolve_seqlen_q helper over the cu_seqlen metadata (folds out unless THD+BR), and the thd_bottom_right capability gap flag is removed. test_mhas_v2 routing: 526/2874 graphs on FROST (was 244/2753), 1951 pass / 0 fail; frost sm100 suites 447 pass (one pre-existing frontend-integration ranking failure reproduces on pristine frost_devel). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The mask parameterization had four spellings of the same thing across the
stack (sm100 config: mask_flags bitmask + swa_window + band_right +
causal_bottom_right; sm120 config: is_causal + causal_bottom_right +
window_size_left; the adapter: mask_flags + swa_window_runtime; the analyzer
facts: window_left + right_bound + bottom_right). Unify on the band model
the analyzer, FlashAttention, and CUTLASS FMHA already share — per-side
OFFSETS from the (aligned) diagonal, None = unbounded:
window_left: Optional[int] # None = no lower bound; W keeps kv >= q-W
window_right: Optional[int] # None = no upper bound; 0 = causal; R widens
bottom_right: bool # anchors the band's diagonal
TemplateParams (sm100 + sm120) and the adapter carry only these three; the
kernels' MASK_FLAGS bits are DERIVED in make_cfg_* (trace-time branch folding
is unchanged), and the CFG value fields rename SWA_WINDOW -> WINDOW_LEFT,
BAND_RIGHT -> WINDOW_RIGHT, CAUSAL_BOTTOM_RIGHT -> BOTTOM_RIGHT along with
the shared bounds/mask helper vocabulary. The canonical band is resolved
once, in the shared adapter base __init__.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
vedaanta
force-pushed
the
vagarwalla/frost-rbw-thd-br
branch
from
August 8, 2026 05:32
908c5ad to
13f200e
Compare
Collaborator
Author
|
@cudnn-ci-bot run frost |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-485-13f200e |
3 tasks
YangXu1990uiuc
added a commit
to YangXu1990uiuc/cudnn-frontend
that referenced
this pull request
Aug 8, 2026
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 NVIDIA#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.
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.
YangXu1990uiuc
added a commit
to YangXu1990uiuc/cudnn-frontend
that referenced
this pull request
Aug 9, 2026
E4M3 in with scalar descales, FP16 out, d128 dense. Same mma.sync architecture as the f16 SM120 cell with the MMA lowered to m16n8k32.e4m3; descale_q*descale_k folds into the softmax scale and descale_v*scale_o into an epilogue scalar, so the kernel adds only the Amax_S/Amax_O atomics over its f16 sibling. Rebuilt on develop rather than rebased, because two things moved underneath it: - NVIDIA#485 unified the mask parameterization onto one band model (window_left / window_right / bottom_right). The kernel is ported the way NVIDIA#485 ported the f16 sibling: causal_bottom_right -> bottom_right internally, and the translation at the make_cfg call site. The adapter needed nothing -- NVIDIA#485 kept the public is_causal/window_size_left arguments and resolves the band once in the base. - NVIDIA#528 moved plan ranking out of the engines. The tile choice is no longer an engine-side propose_plans/knob_order/fp8_tile_choice trio (~140 lines); the cell joins _TILE_RULE_CELLS and _sm120_tiles ranks it. Sharing that rule is a measurement, not an assumption: sweeping 30 causal and non-causal shapes on a 188-SM part gives regret 1.0058 geomean / 1.155 worst, against the f16 rule's own 1.009 / 1.054. The worst cell is a limit of the FEATURES -- B1xH16xS2048 and B2xH8xS2048 causal arrive with identical inputs (grid 128, 16 KV tiles) and have opposite optima, so no threshold separates them. Recorded in the docstring rather than papered over with a second rule. smem_bytes() now sizes its two terms independently. FP8 stages a byte per KV element but still writes O in half, so one itemsize cannot describe both; the f16 path is unchanged (out_itemsize defaults to itemsize). Without this the shared rule would drop tile_n to 64 for wide FP8 heads that fit 128 -- latent today at d128, wrong from d208 up. test_sm120_tile_rule.py covers it, replacing the fp8-only tile test with one over the shared rule. Every FP8 operand is now honoured or rejected, never dropped (AGENTS.md Rule 1): - descale_s / scale_s reached no code at all -- the analyzer never recorded them. Threaded through facts, SdpaBinding, bound_tensors and execute. S IS converted to e4m3 for the PV MMA but UNSCALED, so a reciprocal pair is not equivalent to unity (it asks for a different quantization range); anything but the exact unit pair is declined. - Per-batch seq_len_q is dropped by the quantized lowerings. Harmless while it equals S_q, wrong below it -- O and a finite LSE are written past the valid length. Checked at execute, where the device value is readable; a plan-time decline would reject the equal-length case that works today. - amax_s / amax_o used reshape(), which silently COPIES a non-contiguous input: the kernel would write the copy and the caller read back zeros. view() now, raising instead. Applies to the SM100 fp8 path, which had the same bug. Capabilities gains out_dtypes, declared only by the quantized rows, so an unservable O dtype is a decline rather than a build failure. Verified: SM120 (RTX PRO 6000 Blackwell, 188 SM) fp8 16, all sdpa 177, guards + tile rule 80, test_mhas_v2 245 with routing unchanged at frost 87 / native 158. SM100 (Blackwell) sdpa+gemm 4684, linear_attention 353. CPU dispatch suites 150. All pass.
YangXu1990uiuc
added a commit
to YangXu1990uiuc/cudnn-frontend
that referenced
this pull request
Aug 9, 2026
E4M3 in with scalar descales, FP16 out, d128 dense. Same mma.sync architecture as the f16 SM120 cell with the MMA lowered to m16n8k32.e4m3; descale_q*descale_k folds into the softmax scale and descale_v*scale_o into an epilogue scalar, so the kernel adds only the Amax_S/Amax_O atomics over its f16 sibling. Rebuilt on develop rather than rebased, because two things moved underneath it: - NVIDIA#485 unified the mask parameterization onto one band model (window_left / window_right / bottom_right). The kernel is ported the way NVIDIA#485 ported the f16 sibling: causal_bottom_right -> bottom_right internally, and the translation at the make_cfg call site. The adapter needed nothing -- NVIDIA#485 kept the public is_causal/window_size_left arguments and resolves the band once in the base. - NVIDIA#528 moved plan ranking out of the engines. The tile choice is no longer an engine-side propose_plans/knob_order/fp8_tile_choice trio (~140 lines); the cell joins _TILE_RULE_CELLS and _sm120_tiles ranks it. Sharing that rule is a measurement, not an assumption: sweeping 30 causal and non-causal shapes on a 188-SM part gives regret 1.0058 geomean / 1.155 worst, against the f16 rule's own 1.009 / 1.054. The worst cell is a limit of the FEATURES -- B1xH16xS2048 and B2xH8xS2048 causal arrive with identical inputs (grid 128, 16 KV tiles) and have opposite optima, so no threshold separates them. Recorded in the docstring rather than papered over with a second rule. smem_bytes() now sizes its two terms independently. FP8 stages a byte per KV element but still writes O in half, so one itemsize cannot describe both; the f16 path is unchanged (out_itemsize defaults to itemsize). Without this the shared rule would drop tile_n to 64 for wide FP8 heads that fit 128 -- latent today at d128, wrong from d208 up. test_sm120_tile_rule.py covers it, replacing the fp8-only tile test with one over the shared rule. Every FP8 operand is now honoured or rejected, never dropped (AGENTS.md Rule 1): - descale_s / scale_s reached no code at all -- the analyzer never recorded them. Threaded through facts, SdpaBinding, bound_tensors and execute. S IS converted to e4m3 for the PV MMA but UNSCALED, so a reciprocal pair is not equivalent to unity (it asks for a different quantization range); anything but the exact unit pair is declined. - Per-batch seq_len_q is dropped by the quantized lowerings. Harmless while it equals S_q, wrong below it -- O and a finite LSE are written past the valid length. Checked at execute, where the device value is readable; a plan-time decline would reject the equal-length case that works today. - amax_s / amax_o used reshape(), which silently COPIES a non-contiguous input: the kernel would write the copy and the caller read back zeros. view() now, raising instead. Applies to the SM100 fp8 path, which had the same bug. Capabilities gains out_dtypes, declared only by the quantized rows, so an unservable O dtype is a decline rather than a build failure. Verified: SM120 (RTX PRO 6000 Blackwell, 188 SM) fp8 16, all sdpa 177, guards + tile rule 80, test_mhas_v2 245 with routing unchanged at frost 87 / native 158. SM100 (Blackwell) sdpa+gemm 4684, linear_attention 353. CPU dispatch suites 150. All pass.
YangXu1990uiuc
added a commit
to YangXu1990uiuc/cudnn-frontend
that referenced
this pull request
Aug 9, 2026
E4M3 in with scalar descales, FP16 out, d128 dense. Same mma.sync architecture as the f16 SM120 cell with the MMA lowered to m16n8k32.e4m3; descale_q*descale_k folds into the softmax scale and descale_s*descale_v*scale_o into an epilogue scalar, so the kernel adds only the Amax_S/Amax_O atomics over its f16 sibling. Rebuilt on develop rather than rebased, because two things moved underneath it: - NVIDIA#485 unified the mask parameterization onto one band model (window_left / window_right / bottom_right). The kernel is ported the way NVIDIA#485 ported the f16 sibling: causal_bottom_right -> bottom_right internally, and the translation at the make_cfg call site. The adapter needed nothing -- NVIDIA#485 kept the public is_causal/window_size_left arguments and resolves the band once in the base. - NVIDIA#528 moved plan ranking out of the engines. The tile choice is no longer an engine-side propose_plans/knob_order/fp8_tile_choice trio (~140 lines); the cell joins _TILE_RULE_CELLS and _sm120_tiles ranks it. Sharing that rule is a measurement: 30 seeded causal and non-causal shapes on a 188-SM part give regret 1.0046 geomean / 1.039 worst, against the f16 rule's own 1.009 / 1.054. Most cells sit within the ~1% run-to-run floor, so a single sweep's worst cell is often noise -- an unseeded run of the same code reported 1.155 at one shape that the seeded repeat shows as a tie. What survives repetition is that the misses cluster on causal shapes. P quantization is implemented, following the backend's FORT ordering. cuDNN's Scale_S/Descale_S quantize P -- the softmax OUTPUT, not the scores: the graph applies Scale_S after softmax and after Amax_S, and hands Descale_S to bmm2. This kernel previously converted P to e4m3 unscaled, so both operands reached no math and any graph supplying real S scales -- which the standard contract does -- got a wrong answer silently. P is now scaled before the cast and descale_s folded into o_scale_fused, while tile_sum keeps consuming the unscaled P so the softmax denominator and Amax_S are unaffected. Cost ~0.7-0.9% at large shapes. test_fp8_sm120_s_scales_are_actually_applied is the falsifying test: the two scales are reciprocal in normal use, so applying both and ignoring both give the same O -- it breaks the reciprocity and requires O to track the gain. smem_bytes() sizes its two terms independently: FP8 stages a byte per KV element but still writes O in half, so one itemsize cannot describe both. Without it the shared rule drops tile_n to 64 for wide FP8 heads that fit 128 -- latent at d128, wrong from d208 up. test_sm120_tile_rule.py covers it, replacing the fp8-only tile test with one over the shared rule. Every other FP8 operand is honoured or rejected, never dropped (AGENTS.md Rule 1): descale_s/scale_s reached no code at all (the analyzer never recorded them, so bound_tensors never resolved them); per-batch seq_len_q is dropped by the quantized lowerings, harmless while it equals S_q and wrong below it, now checked at execute where the device value is readable; amax_s/amax_o used reshape(), which silently COPIES a non-contiguous input so the kernel wrote the copy and the caller read back zeros -- view() now, which also fixes the SM100 path. Capabilities gains out_dtypes so an unservable O dtype declines rather than failing at build. Verified on the final commit: SM120 (RTX PRO 6000 Blackwell, 188 SM) guards + tile rule 139, all sdpa 179, fp8 18 twice with identical results, test_mhas_v2 245 with routing unchanged at frost 87 / native 158. SM100 (Blackwell) sdpa + gemm 4684, linear_attention 353. CPU dispatch suites 150.
3 tasks
YangXu1990uiuc
added a commit
to YangXu1990uiuc/cudnn-frontend
that referenced
this pull request
Aug 10, 2026
E4M3 in with scalar descales, FP16 out, d128 dense. Same mma.sync architecture as the f16 SM120 cell with the MMA lowered to m16n8k32.e4m3; descale_q*descale_k folds into the softmax scale and descale_s*descale_v*scale_o into an epilogue scalar, so the kernel adds only the Amax_S/Amax_O atomics over its f16 sibling. Rebuilt on develop rather than rebased, because two things moved underneath it: - NVIDIA#485 unified the mask parameterization onto one band model (window_left / window_right / bottom_right). The kernel is ported the way NVIDIA#485 ported the f16 sibling: causal_bottom_right -> bottom_right internally, and the translation at the make_cfg call site. The adapter needed nothing -- NVIDIA#485 kept the public is_causal/window_size_left arguments and resolves the band once in the base. - NVIDIA#528 moved plan ranking out of the engines. The tile choice is no longer an engine-side propose_plans/knob_order/fp8_tile_choice trio (~140 lines); the cell joins _TILE_RULE_CELLS and _sm120_tiles ranks it. Sharing that rule is a measurement: 30 seeded causal and non-causal shapes on a 188-SM part give regret 1.0046 geomean / 1.039 worst, against the f16 rule's own 1.009 / 1.054. Most cells sit within the ~1% run-to-run floor, so a single sweep's worst cell is often noise -- an unseeded run of the same code reported 1.155 at one shape that the seeded repeat shows as a tie. What survives repetition is that the misses cluster on causal shapes. P quantization is implemented, following the backend's FORT ordering. cuDNN's Scale_S/Descale_S quantize P -- the softmax OUTPUT, not the scores: the graph applies Scale_S after softmax and after Amax_S, and hands Descale_S to bmm2. This kernel previously converted P to e4m3 unscaled, so both operands reached no math and any graph supplying real S scales -- which the standard contract does -- got a wrong answer silently. P is now scaled before the cast and descale_s folded into o_scale_fused, while tile_sum keeps consuming the unscaled P so the softmax denominator and Amax_S are unaffected. Cost ~0.7-0.9% at large shapes. test_fp8_sm120_s_scales_are_actually_applied is the falsifying test: the two scales are reciprocal in normal use, so applying both and ignoring both give the same O -- it breaks the reciprocity and requires O to track the gain. smem_bytes() sizes its two terms independently: FP8 stages a byte per KV element but still writes O in half, so one itemsize cannot describe both. Without it the shared rule drops tile_n to 64 for wide FP8 heads that fit 128 -- latent at d128, wrong from d208 up. test_sm120_tile_rule.py covers it, replacing the fp8-only tile test with one over the shared rule. Every other FP8 operand is honoured or rejected, never dropped (AGENTS.md Rule 1): descale_s/scale_s reached no code at all (the analyzer never recorded them, so bound_tensors never resolved them); per-batch seq_len_q is dropped by the quantized lowerings, harmless while it equals S_q and wrong below it, now checked at execute where the device value is readable; amax_s/amax_o used reshape(), which silently COPIES a non-contiguous input so the kernel wrote the copy and the caller read back zeros -- view() now, which also fixes the SM100 path. Capabilities gains out_dtypes so an unservable O dtype declines rather than failing at build. Verified on the final commit: SM120 (RTX PRO 6000 Blackwell, 188 SM) guards + tile rule 139, all sdpa 179, fp8 18 twice with identical results, test_mhas_v2 245 with routing unchanged at frost 87 / native 158. SM100 (Blackwell) sdpa + gemm 4684, linear_attention 353. CPU dispatch suites 150.
YangXu1990uiuc
added a commit
that referenced
this pull request
Aug 10, 2026
* Add the SM120 per-tensor FP8 (e4m3) SDPA-forward engine E4M3 in with scalar descales, FP16 out, d128 dense. Same mma.sync architecture as the f16 SM120 cell with the MMA lowered to m16n8k32.e4m3; descale_q*descale_k folds into the softmax scale and descale_s*descale_v*scale_o into an epilogue scalar, so the kernel adds only the Amax_S/Amax_O atomics over its f16 sibling. Rebuilt on develop rather than rebased, because two things moved underneath it: - #485 unified the mask parameterization onto one band model (window_left / window_right / bottom_right). The kernel is ported the way #485 ported the f16 sibling: causal_bottom_right -> bottom_right internally, and the translation at the make_cfg call site. The adapter needed nothing -- #485 kept the public is_causal/window_size_left arguments and resolves the band once in the base. - #528 moved plan ranking out of the engines. The tile choice is no longer an engine-side propose_plans/knob_order/fp8_tile_choice trio (~140 lines); the cell joins _TILE_RULE_CELLS and _sm120_tiles ranks it. Sharing that rule is a measurement: 30 seeded causal and non-causal shapes on a 188-SM part give regret 1.0046 geomean / 1.039 worst, against the f16 rule's own 1.009 / 1.054. Most cells sit within the ~1% run-to-run floor, so a single sweep's worst cell is often noise -- an unseeded run of the same code reported 1.155 at one shape that the seeded repeat shows as a tie. What survives repetition is that the misses cluster on causal shapes. P quantization is implemented, following the backend's FORT ordering. cuDNN's Scale_S/Descale_S quantize P -- the softmax OUTPUT, not the scores: the graph applies Scale_S after softmax and after Amax_S, and hands Descale_S to bmm2. This kernel previously converted P to e4m3 unscaled, so both operands reached no math and any graph supplying real S scales -- which the standard contract does -- got a wrong answer silently. P is now scaled before the cast and descale_s folded into o_scale_fused, while tile_sum keeps consuming the unscaled P so the softmax denominator and Amax_S are unaffected. Cost ~0.7-0.9% at large shapes. test_fp8_sm120_s_scales_are_actually_applied is the falsifying test: the two scales are reciprocal in normal use, so applying both and ignoring both give the same O -- it breaks the reciprocity and requires O to track the gain. smem_bytes() sizes its two terms independently: FP8 stages a byte per KV element but still writes O in half, so one itemsize cannot describe both. Without it the shared rule drops tile_n to 64 for wide FP8 heads that fit 128 -- latent at d128, wrong from d208 up. test_sm120_tile_rule.py covers it, replacing the fp8-only tile test with one over the shared rule. Every other FP8 operand is honoured or rejected, never dropped (AGENTS.md Rule 1): descale_s/scale_s reached no code at all (the analyzer never recorded them, so bound_tensors never resolved them); per-batch seq_len_q is dropped by the quantized lowerings, harmless while it equals S_q and wrong below it, now checked at execute where the device value is readable; amax_s/amax_o used reshape(), which silently COPIES a non-contiguous input so the kernel wrote the copy and the caller read back zeros -- view() now, which also fixes the SM100 path. Capabilities gains out_dtypes so an unservable O dtype declines rather than failing at build. Verified on the final commit: SM120 (RTX PRO 6000 Blackwell, 188 SM) guards + tile rule 139, all sdpa 179, fp8 18 twice with identical results, test_mhas_v2 245 with routing unchanged at frost 87 / native 158. SM100 (Blackwell) sdpa + gemm 4684, linear_attention 353. CPU dispatch suites 150. * SM100 FP8: accept reciprocal S scales, decline only non-reciprocal The unit-only guard was too strict and regressed a path that worked. A kernel that converts P unscaled still returns the RIGHT O for a reciprocal pair -- no scale was applied, so none is owed back -- and that is the normal case, descale_s = 1/scale_s. A NON-reciprocal pair is a different request, O scaled by descale_s*scale_s, and ignoring it is silently wrong. That is what the guard should catch, and now all it catches. SM120 implements the scaling (FORT ordering). This row does not, and should not: - No headroom. The lazy-rescale skip (RESCALE_THRESHOLD=8) refreshes the running max only when a tile exceeds it by 2^8, so P is bounded by 256, not 1. e4m3 tops out at 448, so the range above 1.0 is already spent on that skip and only scale_s <= 448/256 = 1.75 is provably safe. This is why the cuDNN backend's SM100 path ignores the pair too -- same kernel structure, same constraint. - Nothing to gain. Measured on B2xH8xS256 e4m3, max|O-ref| is flat to the digit across scale_s 1 -> 64 (swa .0239 throughout) and degrades only once tiles start saturating (swa .0807 at 448). e4m3 is floating point, so relative precision does not move with scale, and subtracting the row max already places P per ROW -- strictly better than a per-tensor scale. Implementing it anyway was tried and reverted: it passed 45/46 fp8 cases and failed swa-e4m3 at .0686 > .05, which is the saturation above. * SM120 FP8: serve THD (ragged), and address review Review (@Aneureka): - The kernel header described P as staged through a per-warp SMEM tile and reloaded with ldmatrix. It is not: mma_pv keeps P in registers and two shfl.sync + one prmt do the k32 C->A exchange, which is what removed the SMEM round trip and its 16 KB. Header rewritten; "P scale is fixed 1.0" went stale in this same PR and is corrected too. - The V-fragment loop claimed a wrapping one-step-ahead prefetch. Loads are in-loop, immediately before their MMAs; prefetch was measured at within +/-0.5% against a ~1% noise floor and dropped, because keeping P in registers leaves little ldmatrix latency to hide. Comment now says that. - graph_analyzer named api_dsl._require_unit_s_scales. Renamed, and the rest of that comment was wrong in two more ways: SM120 does apply the S scales now, and the criterion is non-RECIPROCAL, not non-unit. - THD/sink were implemented-but-unreachable. They differ: sink has no math here at all (only a rejection, `sinks` is always None) and is documented as such; THD was real and is now wired. THD wiring, in the order the layers had to be corrected: - The engine row declared thd=False ("deferred, dense execute only for v1"), so ragged graphs never reached the adapter. - sdpa_support_surface.h rejected (prop_major == 12 && is_ragged) outright. Removing it is safe and verified: with the FROST opt-in off, a ragged fp8 sm120 graph now declines at plan time through the backend's own engine-config check ("No valid engine configs"), so the guard was a redundant early-out, not a capability statement. - The execute-time seq_len_q guard rejected any per-batch length < S_q. Under THD that is the definition of ragged, and the packed layout gives each sequence its own extent, so nothing is written past a valid length. Exempted. - The ragged LSE is head-major (H, head_stride), not token-major -- the kernel writes lse[head, q_row_base + row]. The fake tensor is now 2-D like the f16 cell's, so the caller's padded token capacity is expressible instead of being pinned to the packed total, and the host rank check knows about it. _thd_pack is shared with the f16 THD execute rather than copied; the f16 ragged suite is unchanged (8 passed). Also parameterizes the grouped-query test over H_kv in {1, 2} so MQA is covered (CodeRabbit), and drops docs/fe-oss-apis/attention/sdpa-fp8-sm120.md: it carried kernel design rationale, which is not what the fe-oss-apis docs are for, and this change has no interface surface to document. Validated at this commit: sm120 (RTX PRO 6000 Blackwell) fp8 incl. THD+MQA 24 passed; all frost sdpa 183 passed; native ragged fp8 declines cleanly sm100 (parley Blackwell) fp8+mxfp8 50 passed; all sdpa 567 passed no GPU dispatch + tile rule 80 passed
This was referenced Aug 13, 2026
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
5 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.
Two mask-machinery features for the SM100 FROST SDPA fwd engines, closing the two biggest single-reason routing gaps after bwd (per a test_mhas_v2 mismatch census on SM100: right-band widening solely blocked 220 graphs, THD bottom-right 31+). Supersedes GitLab MR !2308 (closed).
Causal right-band widening (
diagonal_band_right_bound > 0)MASK_CAUSALwith a compile-timeBAND_RIGHToffset on the diagonal, applied in bothcompute_kv_loop_bounds(hi/lo tile limits) andapply_mask_chunk(per-element), so all five SM100 rows — f16 d128/d256/d512, fp8, mxfp8 — serve it (caps flipped).facts.window_right; a new_band_covers_kv_tailhelper accounts for the widened band in the S_kv%128 tail rule (top-left: covered iffs_q + R <= s_kv; bottom-right: iffR == 0).window_size_right(validated>= 0, requires the causal band); SM120 rejects it explicitly (not plumbed there).THD bottom-right per-sequence diagonal
resolve_seqlen_qhelper reads each sequence's actual Q length from the packed cu_seqlen metadata (the same buffer_thd_decodereads) and feeds both the KV-loop bounds and the mask diagonal; it folds out entirely unlessTHD_VARLEN && CAUSAL_BOTTOM_RIGHT.thd_bottom_rightcapability gap flag and its gate are deleted (SM120 already anchored per sequence).Validation (B200, cuDNN 9.23, cutlass-dsl internal)
test_mhas_v2full run: 1951 pass / 0 fail, routing 526/2874 graphs on FROST (18.3%), up from 244/2753 (8.9%) — d128 103→221, d256 86→224, fp8 55→81.band,band_br,band_swa(dense) andcausal_br,band,band_br(THD); analyzer probe tests flipped to accept, plus a negative-bound rejection test.🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Bug Fixes
Tests