Skip to content

Rank plans in one place: delete propose_plans, Router, and heuristics_sort - #528

Merged
Anerudhan merged 15 commits into
NVIDIA:developfrom
YangXu1990uiuc:dispatch-ranking
Aug 9, 2026
Merged

Rank plans in one place: delete propose_plans, Router, and heuristics_sort#528
Anerudhan merged 15 commits into
NVIDIA:developfrom
YangXu1990uiuc:dispatch-ranking

Conversation

@YangXu1990uiuc

@YangXu1990uiuc YangXu1990uiuc commented Aug 8, 2026

Copy link
Copy Markdown
Collaborator

Before submitting

  • I agree to license this contribution under the terms of LICENSE.txt.
  • I ran pre-commit run and committed any formatting changes. (black 26.3.1 -l 160)
  • I added GitHub labels: one cat-*, one or more mod-*, and one orig-*. (token lacks label permission — suggested: cat-cleanup, mod-python-fe, orig-nv-eng)

Affected area

Python graph dispatch (cudnn.pygraph plan ranking). No kernel changes.

Summary

An engine cannot rank plans. It sees neither its siblings nor the backend's entries, so propose_plans could only ever order its own knobs — and something downstream still had to merge the two sides, 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 never decided anything.

create_execution_plans([heur_mode.A, ...])
├─ validate() / _finalize_backend_layout() / _freeze() / _attach_facts()
└─ heuristics.rank(graph, _candidate_engines(), backend_plan_entries(), modes)
   └─ <family>.recommend(modes, facts, offered, backend_plans)
        [delegating backend entry, mode=None]   keeps the lead
        A | B      -> ours + theirs, order decided by measurement
        FALLBACK   -> our safe config + the backend's fallbacks
        OPENSOURCE -> our candidates only (no backend recommendation)

What comes back is graph.plans, position for position. An engine answers two questions only: can I serve this graph (check_support) and compile me this config (build_plan).

The manifest is the only way a python engine exists

register_backend / pygraph(backends=) / graph.backends / OUT_OF_TREE_ID_BASE are gone. 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 registered it, so one that could not serve the graph was still tried and failed at build instead of at classification.

An engine id now decodes straight from the manifest (manifest.engine_for_id), which is what lets create_execution_plan(engine_id, knobs) replay an autotune result on a fresh graph with nothing registered first.

The linear_attention suites had been using register_backend to pin cuTile over FROST. Their conftest verified the pin by inspecting the candidate list, which passes whether or not the pin took effect — so they ran for months against whichever engine the ranking picked, while the seam they pinned through was dead. The pin is now engine_utils.pin_engines() / apply_pin() by name after planning, and it raises when the pinned engine produced no plan.

The backend's plans arrive tagged by mode

Ranking the two sides against each other needs to know which backend entries are mode-A recommendations and which are fallbacks. This needed no C++ change: C++ appends each query to the same plan list and get_execution_plan_count() already exists, so asking one mode at a time gives the boundaries. Measured on a 512³ bf16 matmul (sm90, cuDNN 9.25):

mode plans shape
heur_mode.A [0:15] all knob-bearing
heur_mode.FALLBACK [15:17] bare eng0/eng7, no knobs

The two segments do not overlap.

heur_mode.OPENSOURCE is how FROST coverage gets measured

It is mode A without the backend's recommendation — these cells ARE the open-source implementation. Ask for [OPENSOURCE, A, FALLBACK] and every FROST config is tried first with the backend still behind it, so any graph that runs on a backend plan is one FROST does not cover. That replaces the old "python plans always outrank the backend" assumption, which was a placeholder, never a measurement.

One worked tuning rule, so the shape is copyable

sdpa/fwd/heuristics.py::_sm120_tiles is the only rule in this PR, and it is measured rather than invented: regret 1.009 geomean / 1.054 worst case against the best of the enumerated domain.

  • tile_m = 64 when the grid cannot fill the machine and each CTA has enough KV tiles to amortize the extra Q-tile loop; a causal mask counts as a halved effective grid because it halves the work per CTA.
  • tile_n = the largest that fits SMEM.

That second clause is the interesting one. The SM120 row declares tile_ms={64,128}, tile_ns={64,128}, so before this PR the choice fell through _sole() to None and landed on the adapter's _SM120_Q_TILES[0] default — and on an if self.tile_n is None branch that was quietly shrinking the KV tile to whatever fit. So None had been answering a capability question, not a tuning one; naming tile_n=128 skipped that branch and broke D=208–256 (106512 bytes wanted against the part's 101376). The fit arithmetic now lives once in config_sm120.smem_bytes(), beside the template it describes, and both the adapter's check and the ranking's choice call it.

To add another: write the function, list the cell in _TILE_RULE_CELLS, put the measurement in the commit. A cell absent from that set keeps its row's sole point per axis — the honest answer when nobody has timed it.

Deleted

why
BaseEngine.propose_plans + its 4 implementations all four were the default, verbatim
BaseEngine.default_knobs, BaseEngine.owns_id only fed propose_plans / zero callers
heuristics_sort merging is part of ranking, not a step after it
engines/router.py entirely — Router, default_router, set_router, pygraph(router=) policy has one home now; decline_types moved to base.py, where the engine contract already lives
register_backend, pygraph(backends=), graph.backends, OUT_OF_TREE_ID_BASE see above
engines.probe() (fwd + bwd) superseded by check_support; its only two callers were tests, which now ask analyze_for
graph.engine pure alias of selected_engine, zero callers
graph.from_serialized zero callers; serialize/deserialize are the pybind-era API and stay
engines.probe() callers (incl. one added by #485 mid-rebase) rewired to analyze_for
TorchMatmulEngine (test oracle) it reimplemented matmul/bias/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

Seven tests went with the concepts they tested (six checked registration-time id validation, which has no subject now that engines never declare their own ids). StubEngine replaces the torch oracle: 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, that every node arrives in build order with each input port resolved to the caller's storage and the virtual intermediate carrying none.

knobs=None no longer means "engine, pick for me". 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.

Testing

L40S / cuDNN 9.25 — full L0 sweep, this branch vs its base (develop @ cd5294f21), same command: 60 failures on each side, identical sets. No regressions and no new passes. (Those 60 are pre-existing on this box: rmsnorm x16, causal_conv1d x16, sdpa_fp32_rejected x12, layernorm x8, mhas_v2 x4, rubin_kernel_dispatch x2, native_backend_lowering x1, gdn_bprop x1.)

SM100 (Blackwell, 148 SM) / cuDNN 9.25 — the arch the L40S sweep skips entirely:

suite result
linear_attention (GDN / KDA / GDN-2) — where the pin rewrite lives 353 passed, 0 failed
dispatch guards (test_dispatch, test_graph_native, test_native_backend_lowering, test_import_boundaries) 131 passed, 1 pre-existing failure
sdpa + gemm frost, opt-in ON 4662 passed, 0 failed

SM120 (RTX PRO 6000 Blackwell Server Edition) / cuDNN 9.25 — on the final commit, after the rebase:

stage result
guards, opt-in off 117 passed
sdpa/ suites, opt-in on 139 passed, 0 failed
test_mhas_v2 routing 245 passed — frost:sdpa_fwd_prefill_sm120: 87, native:fp16-fwd: 158

The routing tally is unchanged from develop, which is the point: the tile rule changes which config the SM120 cell runs, not whether it is chosen.

Pre-existing, not from this change: test_a_replayed_plan_reports_its_own_notes fails here and also fails at the base — C++ on 9.25 no longer raises for an index one past the plan count.

Three codex review passes against develop. Round 1 found 6 (all fixed, including a serious one — see below), round 2 found 2, round 3 was clean.

What the review passes caught that I had missed

  • Four SDPA test files were carrying their pre-frost(sdpa): support ragged stats for SM100 frost sdpa forward engine #512 content while the production code they exercise is post-frost(sdpa): support ragged stats for SM100 frost sdpa forward engine #512. This branch is cherry-picked onto the GitHub develop, and the commit consolidating the SDPA test helpers was authored against a tree from before frost(sdpa): support ragged stats for SM100 frost sdpa forward engine #512 landed — so the cherry-pick took the whole file, not the helper edit, and reverted frost(sdpa): support ragged stats for SM100 frost sdpa forward engine #512's test additions with it. api_dsl.py/engines.py were untouched, which is why nothing looked wrong until an SM100 box ran the suite: 16 failures, all tests asserting the old contract against the new kernels. Restored from develop and re-applied only the three intended edits; verified mechanically that each file now differs from develop by exactly those.
  • The delegating backend entry led everything, including the family's OPENSOURCE block. It is not a pure OSS entry: Graph::build_plans tries the C++ OSS engine and falls through to the native configs already enqueued, so ahead of our OPENSOURCE plans it answers a coverage question with a native kernel.
  • Graph::create_execution_plans checks override_heuristics_query() first and returns before it reads the mode at all (deterministic SDPA backward, FP8 backward). Asking one mode at a time therefore appended the same engine-17 config once per mode. SDPA forward's recommend() deduped it; SDPA backward has no heuristics hook, so the duplicates reached graph.plans and an autotuner would time one config twice. Deduped at collection.
  • A plan's identity is (engine_id, knobs), never cpp_index — the index only says where one query put it.
  • test_dsl_sm100_band_right_uncovered_tail_rejected still called probe(). That test arrived with frost(sdpa): causal right-band widening + per-sequence THD bottom-right diagonal on SM100 #485, which this branch rebased onto after the deletion was written, so it would have taken out the whole Blackwell L0 suite.

Docs

docs/python_graph_and_execution_backends.md now carries the dispatch tree explicitly, corrects the three things this PR changed that it had stated as settled (delegate placement, plan identity, per-call mode success), and adds a table of what each machine actually covers — the suites SKIP on the wrong arch rather than fail, so a green sweep on one box says nothing about the others. python/cudnn/frost/README.md no longer describes propose_plans, heuristics_sort, Router, EngineRow/anchors/closed_under or register_backend.

Follow-ups (separate PRs)

  • More per-cell tuning rules, each with its measurements
  • FALLBACK is one config per cell today (the smallest tile the row admits); picking the handful that cover the plane needs measurements — TODO left in heuristics.py
  • A cost model that can compare a python config against a cuDNN engine on a common currency (predicted time), so _MEASURED_BEHIND becomes a number rather than a set

note to self: claude::304e9e55-1db7-4285-967f-001cb21032f3 — "审计 Frost Python DSL 引擎调用流程"
cwd /home/scratch.yanxu_libs/cudnn_frontend · workspace /home/scratch.yanxu_libs/fe_sm120

Summary by CodeRabbit

  • New Features

    • Added manifest-based engine discovery and unified ranked plan selection.
    • Added engine pinning for linear-attention operations.
    • Added SDPA recommendations, fallback plans, and improved tile selection.
    • Added shared-memory sizing for SM120 configurations.
  • Improvements

    • Improved plan deduplication, backend-mode tracking, fallback behavior, and engine-ID replay.
    • Updated dispatch and graph validation for more consistent engine selection.
  • Documentation

    • Expanded guidance on engine discovery, ranking, heuristics, fallback behavior, and configuration.

@coderabbitai

coderabbitai Bot commented Aug 8, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: c18f2f9c-2102-46e7-bd1d-ada444c39b80

📥 Commits

Reviewing files that changed from the base of the PR and between 717562d and bcfee39.

📒 Files selected for processing (3)
  • python/cudnn/_pygraph.py
  • python/cudnn/engines/base.py
  • python/cudnn/sdpa/fwd/heuristics.py
🚧 Files skipped from review as they are similar to previous changes (3)
  • python/cudnn/sdpa/fwd/heuristics.py
  • python/cudnn/engines/base.py
  • python/cudnn/_pygraph.py

📝 Walkthrough

Walkthrough

This PR replaces router-based planning with manifest-discovered engines and family heuristics. It unifies Python and backend plan ranking, removes registration APIs, adds SDPA recommendation logic, updates linear-attention pinning, and migrates tests and documentation.

Changes

Unified engine discovery and ranking

Layer / File(s) Summary
Engine contract and manifest model
docs/python_graph_and_execution_backends.md, python/cudnn/engines/*, python/cudnn/frost/README.md
The engine model now uses manifest-defined engines, family heuristics, decodable IDs, check_support(), and concrete PlanConfig values.
Unified ranking in pygraph
python/cudnn/_pygraph.py, python/cudnn/engines/heuristics.py
pygraph removes router and registration APIs, ranks Python and backend plans together, tracks heuristic modes, and deduplicates plan identities.
Family heuristics and consumer integration
python/cudnn/sdpa/*, python/cudnn/gemm/frost/engine.py, python/cudnn/linear_attention/*
SDPA forward adds recommendation and SM120 sizing logic. Linear-attention operations apply engine pins to ranked plans.
Test migration
test/python/test_dispatch.py, test/python/test_graph_native.py, test/python/test_native_backend_lowering.py, test/python/gemm/*, test/python/sdpa/*, test/python/linear_attention/*
Tests use manifest-backed fake engines, centralized ranking, shared plan selection, and analyzer-based eligibility checks.

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

Sequence Diagram(s)

sequenceDiagram
  participant Graph as pygraph
  participant Manifest as engines.manifest
  participant Backend as backend heuristics
  participant Rank as engines.heuristics.rank
  participant Engine as BaseEngine
  Graph->>Manifest: discover engines and decode engine_id
  Graph->>Backend: collect plans by heuristic mode
  Graph->>Rank: rank Python and backend plans
  Rank->>Engine: check_support(graph)
  Rank-->>Graph: return unified ranked plans
  Graph->>Engine: build selected plan
Loading

Possibly related PRs

Suggested labels: mod-frost

Suggested reviewers: anerudhan

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 60.51% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Title check ✅ Passed The title clearly summarizes the central change: unified plan ranking and removal of obsolete ranking APIs.
Description check ✅ Passed The description covers the affected area, changes, rationale, API impact, testing results, review findings, and follow-up work.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Comment @coderabbitai help to get the list of available commands.

@YangXu1990uiuc

Copy link
Copy Markdown
Collaborator Author

Verified on SM120 hardware (RTX 6000D, cuDNN 9.25):

stage result
dispatch guards (opt-in OFF) 69 passed
sdpa suites (opt-in ON) 137 passed, 0 failed
test_mhas_v2 routing 245 passed — frost:sdpa_fwd_prefill_sm120: 87, native:fp16-fwd: 158

The routing tally is unchanged from develop, which is the point: this PR moves where the ranking decision is made, not what it decides.

f8b1e64bb adds manifest.engine_for_id(). An engine id is fully decodable from the manifest — family block, then slot — so _owners_for_id no longer needs the engine to already be a candidate on that graph. That makes replaying a recorded (engine_id, knobs) work on a fresh graph, and it removes the reason register_backend() looked like a prerequisite for create_execution_plan(). Groundwork for dropping the out-of-tree engine concept, which is the next commit on this branch.

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.
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.
… 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.
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.
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.
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.
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.
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.
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.
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.
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.
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.
@YangXu1990uiuc
YangXu1990uiuc marked this pull request as ready for review August 8, 2026 19:33
@YangXu1990uiuc

Copy link
Copy Markdown
Collaborator Author

Rebased onto develop @ cd5294f21 (#485) and re-verified end to end.

target result
L40S full L0, this branch vs its base, same command 60 failures each, identical sets — no regressions, no new passes
SM100 linear_attention (where the pin rewrite lives) 353 passed, 0 failed
SM100 sdpa + gemm frost, opt-in ON 4662 passed, 0 failed
SM100 dispatch guards 131 passed, 1 pre-existing failure
codex review 3 passes; 6 findings + 2 findings fixed, round 3 clean

Two things the review passes caught that are worth calling out for anyone rebasing this branch again:

  1. Four SDPA test files had silently reverted frost(sdpa): support ragged stats for SM100 frost sdpa forward engine #512 — the branch is cherry-picked onto the GitHub develop, and a commit authored before frost(sdpa): support ragged stats for SM100 frost sdpa forward engine #512 landed brought the whole pre-frost(sdpa): support ragged stats for SM100 frost sdpa forward engine #512 file with it while the production code stayed post-frost(sdpa): support ragged stats for SM100 frost sdpa forward engine #512. Nothing looked wrong until an SM100 box ran the suite. Each of those files is now verified to differ from develop by exactly the intended edit.
  2. Graph::create_execution_plans checks override_heuristics_query() and returns before it reads the mode, so the per-mode query appends the same config once per mode. SDPA forward deduped it; SDPA backward has no heuristics hook and did not. Deduped at collection.

Not covered: SM120 hardware. The node's allocation lapsed mid-session and I could not re-run it after the rebase — the SM120 numbers in the description are from the pre-rebase tree. What is unverified there is #485's own SM120 kernel tweak combined with a one-line test-helper edit on this branch.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 4

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)
python/cudnn/_pygraph.py (1)

783-793: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win

Correct the build_operation_graph() lifecycle description.

validate() now lowers every backend-lowerable graph. This method no longer lowers only when no Python engine is registered. The current docstring gives callers an incorrect freeze and backend-layout lifecycle.

🤖 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/_pygraph.py` around lines 783 - 793, Update the
build_operation_graph() docstring to state that validate() lowers every
backend-lowerable graph, regardless of whether Python engines are registered.
Remove the outdated conditional-lowering and deferred-backend lifecycle
description, while retaining only accurate sequencing and plan-configuration
behavior.
🧹 Nitpick comments (4)
python/cudnn/sdpa/fwd/engines.py (1)

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

Sort __all__ to satisfy Ruff.

Line 772 violates RUF022. Sort the exported names in the order required by the configured linter.

Proposed fix
-__all__ = ["Capabilities", "EngineSpec", "ENGINE_SPECS", "SdpaFwdKnobs", "analyze_for", "build", "engine_name", "mismatch"]
+__all__ = ["Capabilities", "ENGINE_SPECS", "EngineSpec", "SdpaFwdKnobs", "analyze_for", "build", "engine_name", "mismatch"]
🤖 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/engines.py` at line 772, Sort the exported names in
__all__ alphabetically to satisfy Ruff RUF022, preserving the same set of
symbols: Capabilities, ENGINE_SPECS, EngineSpec, SdpaFwdKnobs, analyze_for,
build, engine_name, and mismatch.

Source: Linters/SAST tools

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

18-18: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Remove the stale # noqa: F401.

_select_engine is used at Line 84 and Line 202. The suppression is not needed and hides a real unused-import warning if the last usage is removed later.

♻️ Proposed change
-from frost_test_utils import select_engine as _select_engine  # noqa: F401
+from frost_test_utils import select_engine as _select_engine
🤖 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` at line 18, Remove the
stale `# noqa: F401` comment from the `_select_engine` import in
`test_sdpa_fwd_dsl_sm100.py`, leaving the import and its existing usages
unchanged.
test/python/test_dispatch.py (1)

703-703: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Rename the unused C binding to _C.

Ruff flags C as an unpacked-but-unused variable (RUF059) at these five g, C = _backend_first(monkeypatch) call sites. Rename it to _C at the sites where the tensor is not used.

Also applies to: 1156-1156, 1357-1357, 1390-1390, 1435-1435

🤖 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/test_dispatch.py` at line 703, Rename the unused second binding
from C to _C in each of the five _backend_first(monkeypatch) unpacking call
sites, while preserving g and any sites where the tensor binding is actually
used.

Source: Linters/SAST tools

test/python/test_graph_native.py (1)

525-536: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Extract the repeated Dummy engine boilerplate into one helper.

Four tests repeat the same import, class definition, and _offer call, differing only in the engine-id offset. None of them sets name, so every Dummy inherits BaseEngine.name == "base". _offer keys its slots mapping by e.name, so two such engines passed to one _offer call would collide on the key "base" and one slot would be dropped silently.

Add a module-level helper and give the engine an explicit name.

♻️ Proposed refactor
# module level, near the other test helpers
def _offer_dummy(monkeypatch, offset, node_type="MATMUL"):
    """Offer a no-op python engine so planning stays python-side (no C++ needed)."""
    from cudnn.engines import BaseEngine

    from test_dispatch import _FAKE, _offer

    class Dummy(BaseEngine):
        name = f"dummy_{offset}"
        engine_id = _FAKE + offset

        def execute(self, graph, tensor_data, ctx=None):
            pass

    return _offer(monkeypatch, Dummy(), node_type=node_type)

Each test then calls _offer_dummy(monkeypatch, 90) or _offer_dummy(monkeypatch, 93, node_type="GDN").

Also applies to: 555-566, 576-590, 626-636

🤖 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/test_graph_native.py` around lines 525 - 536, Extract the
repeated BaseEngine setup from the affected tests into a module-level
_offer_dummy helper near the existing test helpers. Have it accept monkeypatch,
offset, and optional node_type, define Dummy with a unique name such as
dummy_<offset>, set engine_id to _FAKE + offset, preserve the no-op execute
method, and return _offer’s result. Replace each repeated import/class/_offer
block with calls using its existing offset and node type.
🤖 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/engines/manifest.py`:
- Around line 305-320: Update engine_for_id() to resolve replay IDs using each
engine’s owned_id_range rather than requiring an exact match with the offered
base ID. After confirming the family owns the ID and the slot is offered, select
the instantiated engine whose declared range contains engine_id, while
preserving the existing gated-off behavior. Add a fresh-graph replay test
covering a non-base ID within an offered engine range.

In `@test/python/sdpa/frost/frost_test_utils.py`:
- Around line 86-88: Update the plan lookup in the test utility to match tile_m
and tile_n as exact values rather than using the substring condition `want in
n`. Prefer comparing the corresponding structured `PlanConfig.knobs` values;
otherwise enforce delimiters so values such as 128 cannot match 1280, while
preserving the existing missing-plan assertion.

In `@test/python/test_dispatch.py`:
- Around line 980-991: Update the test setup so _offer receives the probe
analyzer explicitly, matching the pattern used by
test_planning_attaches_facts_without_anyone_asking, instead of allowing _offer
to overwrite the probe_family manifest configuration. Remove the now-unused
manifest import and obsolete explanatory comment, while preserving the test’s
verification that ranking resolves the analyzer declared by
EngineFamily.analyzer.

In `@test/python/test_native_backend_lowering.py`:
- Line 424: Update the docstring of test_mixed_ranking_dispatch to reference
test_mixed_ranking_backend_slot_executes instead of the stale
test_mixed_router_backend_slot_executes name, preserving the existing
cross-reference format.

---

Outside diff comments:
In `@python/cudnn/_pygraph.py`:
- Around line 783-793: Update the build_operation_graph() docstring to state
that validate() lowers every backend-lowerable graph, regardless of whether
Python engines are registered. Remove the outdated conditional-lowering and
deferred-backend lifecycle description, while retaining only accurate sequencing
and plan-configuration behavior.

---

Nitpick comments:
In `@python/cudnn/sdpa/fwd/engines.py`:
- Line 772: Sort the exported names in __all__ alphabetically to satisfy Ruff
RUF022, preserving the same set of symbols: Capabilities, ENGINE_SPECS,
EngineSpec, SdpaFwdKnobs, analyze_for, build, engine_name, and mismatch.

In `@test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py`:
- Line 18: Remove the stale `# noqa: F401` comment from the `_select_engine`
import in `test_sdpa_fwd_dsl_sm100.py`, leaving the import and its existing
usages unchanged.

In `@test/python/test_dispatch.py`:
- Line 703: Rename the unused second binding from C to _C in each of the five
_backend_first(monkeypatch) unpacking call sites, while preserving g and any
sites where the tensor binding is actually used.

In `@test/python/test_graph_native.py`:
- Around line 525-536: Extract the repeated BaseEngine setup from the affected
tests into a module-level _offer_dummy helper near the existing test helpers.
Have it accept monkeypatch, offset, and optional node_type, define Dummy with a
unique name such as dummy_<offset>, set engine_id to _FAKE + offset, preserve
the no-op execute method, and return _offer’s result. Replace each repeated
import/class/_offer block with calls using its existing offset and node type.
🪄 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: b0de8de8-0a58-445b-a303-818ce9bd5b70

📥 Commits

Reviewing files that changed from the base of the PR and between cd5294f and 1a85c82.

📒 Files selected for processing (35)
  • docs/python_graph_and_execution_backends.md
  • python/cudnn/_pygraph.py
  • python/cudnn/engines/__init__.py
  • python/cudnn/engines/base.py
  • python/cudnn/engines/engine_ids.py
  • python/cudnn/engines/heuristics.py
  • python/cudnn/engines/manifest.py
  • python/cudnn/engines/router.py
  • python/cudnn/frost/README.md
  • python/cudnn/gemm/frost/engine.py
  • python/cudnn/linear_attention/engine_utils.py
  • python/cudnn/linear_attention/ops/gdn.py
  • python/cudnn/linear_attention/ops/gdn2.py
  • python/cudnn/linear_attention/ops/kda.py
  • python/cudnn/sdpa/bwd/engine.py
  • python/cudnn/sdpa/bwd/engines.py
  • python/cudnn/sdpa/fwd/api_dsl.py
  • python/cudnn/sdpa/fwd/config_sm120.py
  • python/cudnn/sdpa/fwd/engine.py
  • python/cudnn/sdpa/fwd/engines.py
  • python/cudnn/sdpa/fwd/heuristics.py
  • test/python/gemm/frost/test_frontend_integration.py
  • test/python/linear_attention/cutile/conftest.py
  • test/python/sdpa/frost/frost_test_utils.py
  • test/python/sdpa/frost/test_sdpa_bwd_dsl_sm120.py
  • test/python/sdpa/frost/test_sdpa_frontend_integration.py
  • test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py
  • test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py
  • test/python/sdpa/frost/test_sdpa_fwd_fp8_sm100.py
  • test/python/sdpa/frost/test_sdpa_fwd_mxfp8_sm100.py
  • test/python/sdpa/frost/test_sdpa_graph_analyzer.py
  • test/python/test_api_signature_parity.py
  • test/python/test_dispatch.py
  • test/python/test_graph_native.py
  • test/python/test_native_backend_lowering.py
💤 Files with no reviewable changes (4)
  • python/cudnn/sdpa/bwd/engines.py
  • python/cudnn/gemm/frost/engine.py
  • python/cudnn/engines/engine_ids.py
  • python/cudnn/engines/router.py

Comment thread python/cudnn/engines/manifest.py
Comment thread test/python/sdpa/frost/frost_test_utils.py Outdated
Comment thread test/python/test_dispatch.py
Comment thread test/python/test_native_backend_lowering.py
@YangXu1990uiuc YangXu1990uiuc self-assigned this Aug 8, 2026
@YangXu1990uiuc YangXu1990uiuc added cat-doc Documentation changes, examples, API references, tutorials, or wording fixes. mod-frontend cuDNN frontend APIs, operation graph construction, plans, and user-facing wrappers. orig-nv-eng Reported or requested by NVIDIA engineering. mod-cutedsl CuTeDSL kernels, generated kernels, examples, or related integration work. cat-cleanup labels Aug 8, 2026
@YangXu1990uiuc YangXu1990uiuc added this to the Frontend 1.28.0 milestone Aug 8, 2026
…uld 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 NVIDIA#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.
@YangXu1990uiuc

Copy link
Copy Markdown
Collaborator Author

All four CodeRabbit findings addressed in 9da8cee; replies inline on each thread.

Re-verified after the fixes:

target result
CPU dispatch suites 204 passed, 1 pre-existing failure
SM100 linear_attention 353 passed, 0 failed
SM100 sdpa + gemm frost 4662 passed, 0 failed
SM100 dispatch guards 131 passed, 1 pre-existing

Identical to the pre-fix run, so no regression from the changes.

One of the four I fixed in the opposite direction to what was suggested — id_end / owned_id_range are deleted rather than propagated, because nothing in the tree sets id_end and the range only ever existed to let a registered engine claim a block. Detail on that thread.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)
python/cudnn/engines/base.py (1)

56-67: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win

Document ImportError as a decline signal.

decline_types() now includes ImportError. The check_support() docstring at lines 211-213 still states that only NotImplementedError and cudnn.cudnnGraphNotSupportedError decline a graph. This gives engine authors an incorrect exception contract. Update the docstring to include ImportError.

🤖 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/engines/base.py` around lines 56 - 67, The check_support()
docstring must document ImportError as an additional exception type that
declines a graph. Update its exception-contract description to include
ImportError alongside NotImplementedError and cudnn.cudnnGraphNotSupportedError,
without changing the implementation.
🤖 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.

Outside diff comments:
In `@python/cudnn/engines/base.py`:
- Around line 56-67: The check_support() docstring must document ImportError as
an additional exception type that declines a graph. Update its
exception-contract description to include ImportError alongside
NotImplementedError and cudnn.cudnnGraphNotSupportedError, without changing the
implementation.

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 88d941bc-9723-4c52-b0c5-a3e59a9130c0

📥 Commits

Reviewing files that changed from the base of the PR and between 1a85c82 and 9da8cee.

📒 Files selected for processing (4)
  • python/cudnn/_pygraph.py
  • python/cudnn/engines/base.py
  • test/python/sdpa/frost/frost_test_utils.py
  • test/python/test_dispatch.py
🚧 Files skipped from review as they are similar to previous changes (3)
  • test/python/sdpa/frost/frost_test_utils.py
  • test/python/test_dispatch.py
  • python/cudnn/_pygraph.py

@YangXu1990uiuc

Copy link
Copy Markdown
Collaborator Author

SM120 gap closed — re-ran on the final commit 9da8cee9b (RTX PRO 6000 Blackwell Server Edition, 2u2g-spr-0105, cuDNN 9.25):

stage result
guards, opt-in off 117 passed
sdpa/ suites, opt-in on 139 passed, 0 failed
test_mhas_v2 routing 245 passed — frost:sdpa_fwd_prefill_sm120: 87, native:fp16-fwd: 158

The routing tally is byte-identical to develop, which is the point: the tile rule changes which config the SM120 cell runs, not whether it is chosen.

Full coverage on the final commit is now:

target result
L40S full L0, branch vs base, same command 60 failures each, identical sets
SM100 linear_attention 353 passed
SM100 sdpa + gemm frost 4662 passed
SM100 dispatch guards 131 passed, 1 pre-existing
SM120 (above) 501 passed
CPU dispatch suites 204 passed, 1 pre-existing

The only failure anywhere is test_a_replayed_plan_reports_its_own_notes, which fails identically at the base.

@Anerudhan
Anerudhan self-requested a review August 9, 2026 05:49
@Anerudhan

Copy link
Copy Markdown
Collaborator

@cudnn-ci-bot frost

@cudnn-ci-bot

Copy link
Copy Markdown

cuDNN CI bot commands

  • @cudnn-ci-bot status: confirm the bot is up.
  • @cudnn-ci-bot check: validate this PR without launching CI.
  • @cudnn-ci-bot run <targets>: mirror this PR's head SHA and launch a pipeline.
  • @cudnn-ci-bot run help: list the targets you can name.

Only allowlisted maintainers can use @cudnn-ci-bot check or @cudnn-ci-bot run.

@Anerudhan

Copy link
Copy Markdown
Collaborator

@cudnn-ci-bot run frost

@cudnn-ci-bot

Copy link
Copy Markdown

🚀 Running mirror pipeline

Branch: cudnn-gh/pr-528-9da8cee
Pipeline: 61785058
Targets: frost

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 NVIDIA#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.
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 NVIDIA#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.
@Anerudhan
Anerudhan merged commit 3fcee36 into NVIDIA:develop Aug 9, 2026
1 check passed
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.
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
egilliam-nv added a commit to egilliam-nv/cudnn-frontend that referenced this pull request Aug 10, 2026
 plan names carry knobs)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
egilliam-nv added a commit to egilliam-nv/cudnn-frontend that referenced this pull request Aug 12, 2026
 plan names carry knobs)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
YangXu1990uiuc pushed a commit that referenced this pull request Aug 13, 2026
* SM80 (A100) SDPA: FROST engines + cudnn.sdpa adapters

CuTe-DSL SDPA forward and backward for SM80 as FROST engines, plus
standalone cudnn.sdpa APIs.

* sdpa_fwd_prefill_sm80 joins the FrostSdpaFwdEngines family (the manifest
  row's sm_lo drops to 80); sdpa_bwd_sm80 introduces the first backward
  opset as a new FrostSdpaBwdEngines family on the reserved
  FROST_SDPA_BWD_ID_BASE block, with a frost_sdpa_bwd manifest row anchored
  on SDPA_BWD.
* The shared graph_analyzer learns sdpa_backward() graphs: backward facts,
  K/V transposed-input-view canonicalization, and a forward-direction gate
  in the shared mismatch().
* Kernels (fwd generic + d256; bwd generic + d64 fast path) build on the
  shared frost/tile_dsl library; torch-native host code, per-shape
  self-caching, stream-aware and CUDA-graph-capturable.
* Standalone SdpafwdSm80 / SdpabwdSm80 APIBase adapters + wrappers,
  including a packed-THD path the engines do not expose yet.
* Docs: FE OSS API pages + Attention.md sections.

Verified on A100: reference suites (103), engine/analyzer/stream suites
(123), and full test_mhas_v2 with FROST auto-selection: 1268 passed / 0
failed, 41.2% of graphs served on FROST.

The kernels originate from earlier internal work by Roman Anders.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* review: address CodeRabbit findings

* drop the stray fe_cuda_13.3.cfg (local build config, added by accident)
* THD forward: resolve the default softmax scale from the user's head dim
  BEFORE flavor padding (a d=96 call silently used 1/sqrt(128))
* THD forward/backward: reject dense-only features (bias / RoPE /
  block_mask / seq lens / scale_output / scheduler) instead of silently
  computing without them; drop the never-used backward max_s_q kwarg
* dense mask resolution: window_size=(-1, r) without is_causal is now
  rejected instead of silently selecting a 0-token SWA window
* bprop kernel: assert K's head dim matches Q's, and reject RoPE at
  d_qk > 128 (the sDQ SMEM staging exceeds the A100 budget beyond that)
* bwd reference in the fe_api suite anchors the causal diagonal top-left,
  matching the mask the wrapper actually requests
* docs: requirements sections no longer cite the retired ctm package;
  provenance notes restored in both kernel packages; THD feature gating
  documented
* stale comments: flavor envelope (dsv3->qwen), d64 'not yet routed',
  gptoss draft artifact, prefill header diagram / swizzle note /
  mainloop wait counts, d256 wait count

* review: actually drop the stray fe_cuda_13.3.cfg

The previous commit unstaged it, then a blanket git add -A swept the
untracked local file back in before committing. Now excluded locally.

* review: round 2 — backward twins of the round-1 fixes + test tiering

* bwd dense mask resolution: reject window_size=(-1, r) without is_causal
  instead of silently selecting a 0-token SWA window (the fwd adapter and
  both THD paths already resolve it this way)
* bwd THD: resolve the default softmax scale from the user's head dim
  BEFORE flavor padding (same silent-wrong-gradients bug as the fwd twin)
* fe_api sweeps (48 cases each) move from L0 to L2 per the repo guideline;
  a single representative smoke case stays at L0 for each direction
* underscore the unused unpacks Ruff flags in the fwd reference

* docs: purge the last ctm mentions — the retired DSL package is not a dependency of anything in this PR

* test: skip the SM80 fe_api suites when the CuTe DSL is missing or predates cutlass.experimental

The package imports are lazy (PEP 562) since the engine-family cleanup, so
a missing or old nvidia-cutlass-dsl no longer fails at wrapper import —
it erupts at kernel-load time mid-test. The oss:rel CI leg (older DSL)
showed 9 such errors; probe cutlass.experimental in the module skip so
those environments skip cleanly.

* review: round 3 — require the bound seq_len_q buffer + Ruff lints

* resolve_feature_operands: a graph that BINDS seq_len_q must get its
  buffer through _need() like every other feature operand — silently
  omitting it would execute a different query-length contract (missed
  BR-under-padding base and padded-row LSE trim)
* Ruff RUF059 (two unused unpacks) and E731 (lambda-to-def)

* fix: drop duplicated stale helpers the analyzer merge left behind

The rebase onto the engine-family cleanup appended our helper block
including copies of resolve_variant_pack (identical, harmless) and
tensor_desc_from_ir (STALE: still referencing the deleted
_DTYPE_FROM_CUDNN map). The stale copy shadowed upstream's fixed one and
took down every SM100/SM120 lowering with a NameError — invisible on
A100, where those paths skip; pipeline 61601513's Blackwell frost leg
caught it (716 failures).

* fix: never wrap the default stream handle in ExternalStream (CI determinism zeros)

Root cause of the frost_tests:sdpa[Ampere] determinism failures (104/104
is_determin backward configs, second execution returning all-zero grads):
the harness's cudnn handle carries raw stream 0, and _stream_ctx wrapped it
in torch.cuda.ExternalStream(0). On the CI image's NGC torch build
(2.12.0a0), every kernel launch inside that context after the compile run
silently no-ops. Reproduced and verified in the CI container itself
(gitlab/cudnn_frontend:cudnn_13.3.0 + the pipeline's build artifact):
before — run 0 correct, runs 1+ all-zero; after — bitwise-identical runs,
test_sdpa_random_bwd_L0 176 passed / 0 failed under -n 4.

_stream_ctx now maps a raw handle equal to torch's current/default stream
onto that torch stream object and reserves ExternalStream for genuine
foreign streams — the same guard fwd/api_dsl._torch_stream_context and
gemm/cutedsl/grouped/backend_utils.py already carry (these adapters were
the only unguarded spot in the tree).

Known residual, unrelated to CI: on NGC torch the FIRST execute after
re-pointing a handle to a brand-new stream still no-ops once on a cached
kernel (suspected DSL/tvm-ffi launch-state caching; to be reported
upstream — the kernels cannot compile without tvm-ffi, so it could not be
isolated further).

* test: pin SM80 plans via frost_test_utils.select_engine (post-#528 plan names carry knobs)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* review: fold SM80 adapters into api.py, configs into config_sm80.py, drop the docs pages

Per maintainer feedback on the PR:
- SdpafwdSm80/SdpabwdSm80 + wrappers move from api_sm80.py into each opset's
  api.py, following the SM100 classes there (one api.py per opset, no per-arch
  files). Lazy exports and the engine lowerings repoint; no signature changes.
- The per-flavor kernel configs leave kernels/ for parent-level config_sm80.py
  (one per direction), matching config_sm100/config_sm120.
- The FE OSS docs pages and their overview/Attention/llms.txt entries are
  dropped for now; the engines remain the PR's product.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* rebase: dedupe the bwd Capabilities.layouts field (#533 introduced its own copy)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

---------

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
vedaanta added a commit to vedaanta/cudnn-frontend that referenced this pull request Aug 13, 2026
Rebase adaptation: the feature-operand refactor (NVIDIA#493/NVIDIA#528) moved seq-lens
resolution into the presence-checked ga.resolve_feature_operands, which
raised "padding mask (seq_len_kv) requested but no buffer was provided"
for cu-form graphs (facts.seq_kv_t is None there). Either length form now
satisfies a side directly in the helper — the (B+1,) cu buffer travels
through the same operand slot — and the engines-side fallback becomes
dead code and is removed.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
vedaanta added a commit that referenced this pull request Aug 13, 2026
…522)

* frost(sdpa): accept the cu_seq_len (prefix-sum) length form for THD

cu_seq_len_q/kv ((B+1,) prefix sums, cuDNN 9.24+) is the length form
every major consumer natively holds — TE and PyTorch both launch a
device conversion kernel per call today purely to feed cuDNN's
per-batch SEQ_LEN form (TE: cu_seqlens_to_actual_seqlens each fwd+bwd;
PyT: 2x at::diff in MHA.cpp). The frost THD lowering already derives
both forms host-side from its inherent tolist round-trip, so accepting
cu is free:

- graph_analyzer: cu_seq_q_t/cu_seq_kv_t captured in facts (was only a
  bool); THD/padded graphs may carry either length form per side;
  both-forms-on-one-side stays analyzer-VALID (invalid means
  malformed-for-everyone) and is declined by the engine gate instead.
- engines: dedicated cu gate replacing the generic capability row —
  serving rows (SM100 f16, SM120) take THD cu graphs; dense cu graphs
  stay declined until the kernels grow a CU read mode
  (len = cu[b+1] - cu[b]); ambiguous both-forms graphs decline with a
  precise reason. Binding + lowering route the cu buffer through the
  same seq-lens execute argument.
- adapters: shared _thd_host_lens consumes either form in ONE D2H
  round-trip — per-batch lengths scan up to prefix sums, or cu
  differences down to lengths, with the prefix-sum invariants
  (starts at 0, non-decreasing) validated host-side where they are
  free to check. check_support rejects cu flags outside THD.
- tests: THD cu end-to-end (both ragged Stats layouts + zero-length /
  all-KV-zero degenerates) and analyzer probes (accept THD cu, decline
  ambiguous).

Testing (cc 10.0): cu tests 3 passed; sm100 THD/stats slice 116
passed; integration 10 passed; analyzer 71 passed.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* frost(sdpa): document the THD packed-only storage contract

The THD lowerings re-derive packed addressing as prefix(lens) x token
stride and never read the graph's bound ragged-offset values — TE-style
padded THD (offsets from cu_seqlens_padded != cu_seqlens, gaps between
sequences) is not served, and being runtime data it cannot be declined
at plan time. State it on Capabilities.thd where every serving row
inherits it.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* frost(sdpa): resolve cu_seq_len buffers through resolve_feature_operands

Rebase adaptation: the feature-operand refactor (#493/#528) moved seq-lens
resolution into the presence-checked ga.resolve_feature_operands, which
raised "padding mask (seq_len_kv) requested but no buffer was provided"
for cu-form graphs (facts.seq_kv_t is None there). Either length form now
satisfies a side directly in the helper — the (B+1,) cu buffer travels
through the same operand slot — and the engines-side fallback becomes
dead code and is removed.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* test(sdpa): SM120 cu_seq_len THD coverage (review)

Haobin: mirror the cu_seq_len tests into the SM120 suite. The THD
harness gains the same cu_lens binding as the SM100 one ((B+1,)
prefix-sum tensors through cu_seq_len_q/kv instead of per-batch
lengths), with cu variants of the stats test (both declared layouts)
and the degenerate-lengths sweep (zero-length sequence, all-KV-zero
dead rows, all-Q-zero no-op).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

---------

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cat-cleanup cat-doc Documentation changes, examples, API references, tutorials, or wording fixes. mod-cutedsl CuTeDSL kernels, generated kernels, examples, or related integration work. mod-frontend cuDNN frontend APIs, operation graph construction, plans, and user-facing wrappers. orig-nv-eng Reported or requested by NVIDIA engineering.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants