Rank plans in one place: delete propose_plans, Router, and heuristics_sort - #528
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (3)
🚧 Files skipped from review as they are similar to previous changes (3)
📝 WalkthroughWalkthroughThis 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. ChangesUnified engine discovery and ranking
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
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
|
Verified on SM120 hardware (RTX 6000D, cuDNN 9.25):
The routing tally is unchanged from develop, which is the point: this PR moves where the ranking decision is made, not what it decides.
|
b75a915 to
103bf3f
Compare
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.
103bf3f to
71a6af7
Compare
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.
|
Rebased onto
Two things the review passes caught that are worth calling out for anyone rebasing this branch again:
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. |
There was a problem hiding this comment.
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 winCorrect 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 winSort
__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 valueRemove the stale
# noqa: F401.
_select_engineis 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 valueRename the unused
Cbinding to_C.Ruff flags
Cas an unpacked-but-unused variable (RUF059) at these fiveg, C = _backend_first(monkeypatch)call sites. Rename it to_Cat 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 valueExtract the repeated
Dummyengine boilerplate into one helper.Four tests repeat the same import, class definition, and
_offercall, differing only in the engine-id offset. None of them setsname, so everyDummyinheritsBaseEngine.name == "base"._offerkeys itsslotsmapping bye.name, so two such engines passed to one_offercall 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
📒 Files selected for processing (35)
docs/python_graph_and_execution_backends.mdpython/cudnn/_pygraph.pypython/cudnn/engines/__init__.pypython/cudnn/engines/base.pypython/cudnn/engines/engine_ids.pypython/cudnn/engines/heuristics.pypython/cudnn/engines/manifest.pypython/cudnn/engines/router.pypython/cudnn/frost/README.mdpython/cudnn/gemm/frost/engine.pypython/cudnn/linear_attention/engine_utils.pypython/cudnn/linear_attention/ops/gdn.pypython/cudnn/linear_attention/ops/gdn2.pypython/cudnn/linear_attention/ops/kda.pypython/cudnn/sdpa/bwd/engine.pypython/cudnn/sdpa/bwd/engines.pypython/cudnn/sdpa/fwd/api_dsl.pypython/cudnn/sdpa/fwd/config_sm120.pypython/cudnn/sdpa/fwd/engine.pypython/cudnn/sdpa/fwd/engines.pypython/cudnn/sdpa/fwd/heuristics.pytest/python/gemm/frost/test_frontend_integration.pytest/python/linear_attention/cutile/conftest.pytest/python/sdpa/frost/frost_test_utils.pytest/python/sdpa/frost/test_sdpa_bwd_dsl_sm120.pytest/python/sdpa/frost/test_sdpa_frontend_integration.pytest/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.pytest/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.pytest/python/sdpa/frost/test_sdpa_fwd_fp8_sm100.pytest/python/sdpa/frost/test_sdpa_fwd_mxfp8_sm100.pytest/python/sdpa/frost/test_sdpa_graph_analyzer.pytest/python/test_api_signature_parity.pytest/python/test_dispatch.pytest/python/test_graph_native.pytest/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
…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.
|
All four CodeRabbit findings addressed in 9da8cee; replies inline on each thread. Re-verified after the fixes:
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 — |
There was a problem hiding this comment.
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 winDocument
ImportErroras a decline signal.
decline_types()now includesImportError. Thecheck_support()docstring at lines 211-213 still states that onlyNotImplementedErrorandcudnn.cudnnGraphNotSupportedErrordecline a graph. This gives engine authors an incorrect exception contract. Update the docstring to includeImportError.🤖 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
📒 Files selected for processing (4)
python/cudnn/_pygraph.pypython/cudnn/engines/base.pytest/python/sdpa/frost/frost_test_utils.pytest/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
|
SM120 gap closed — re-ran on the final commit
The routing tally is byte-identical to Full coverage on the final commit is now:
The only failure anywhere is |
|
@cudnn-ci-bot frost |
|
cuDNN CI bot commands
Only allowlisted maintainers can use |
|
@cudnn-ci-bot run frost |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-528-9da8cee |
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.
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.
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.
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.
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.
* 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
* 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>
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>
…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>
Before submitting
pre-commit runand committed any formatting changes. (black 26.3.1 -l 160)cat-*, one or moremod-*, and oneorig-*. (token lacks label permission — suggested:cat-cleanup,mod-python-fe,orig-nv-eng)Affected area
Python graph dispatch (
cudnn.pygraphplan ranking). No kernel changes.Summary
An engine cannot rank plans. It sees neither its siblings nor the backend's entries, so
propose_planscould only ever order its own knobs — and something downstream still had to merge the two sides, whichheuristics_sortdid by concatenating and calling it ranking.All four in-tree
propose_planswere the base class's default copied verbatim. The hook never decided anything.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_BASEare gone. An out-of-tree engine could never be ranked anyway: it declares noCapabilities, 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 letscreate_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_backendto 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 nowengine_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):heur_mode.A[0:15]heur_mode.FALLBACK[15:17]eng0/eng7, no knobsThe two segments do not overlap.
heur_mode.OPENSOURCEis how FROST coverage gets measuredIt 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_tilesis 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 = 64when 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()toNoneand landed on the adapter's_SM120_Q_TILES[0]default — and on anif self.tile_n is Nonebranch that was quietly shrinking the KV tile to whatever fit. SoNonehad been answering a capability question, not a tuning one; namingtile_n=128skipped that branch and broke D=208–256 (106512 bytes wanted against the part's 101376). The fit arithmetic now lives once inconfig_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
BaseEngine.propose_plans+ its 4 implementationsBaseEngine.default_knobs,BaseEngine.owns_idpropose_plans/ zero callersheuristics_sortengines/router.pyentirely —Router,default_router,set_router,pygraph(router=)decline_typesmoved tobase.py, where the engine contract already livesregister_backend,pygraph(backends=),graph.backends,OUT_OF_TREE_ID_BASEengines.probe()(fwd + bwd)check_support; its only two callers were tests, which now askanalyze_forgraph.engineselected_engine, zero callersgraph.from_serializedserialize/deserializeare the pybind-era API and stayengines.probe()callers (incl. one added by #485 mid-rebase)analyze_forTorchMatmulEngine(test oracle)torch_matmulin plan names reads like something cuDNN shipsSeven 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).
StubEnginereplaces 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=Noneno 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:
linear_attention(GDN / KDA / GDN-2) — where the pin rewrite livestest_dispatch,test_graph_native,test_native_backend_lowering,test_import_boundaries)sdpa+gemmfrost, opt-in ONSM120 (RTX PRO 6000 Blackwell Server Edition) / cuDNN 9.25 — on the final commit, after the rebase:
sdpa/suites, opt-in ontest_mhas_v2routingfrost:sdpa_fwd_prefill_sm120: 87,native:fp16-fwd: 158The 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_notesfails here and also fails at the base — C++ on 9.25 no longer raises for an index one past the plan count.Three
codex reviewpasses againstdevelop. 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
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.pywere 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 fromdevelopand re-applied only the three intended edits; verified mechanically that each file now differs fromdevelopby exactly those.Graph::build_planstries 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_planschecksoverride_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'srecommend()deduped it; SDPA backward has no heuristics hook, so the duplicates reachedgraph.plansand an autotuner would time one config twice. Deduped at collection.(engine_id, knobs), nevercpp_index— the index only says where one query put it.test_dsl_sm100_band_right_uncovered_tail_rejectedstill calledprobe(). 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.mdnow 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.mdno longer describespropose_plans,heuristics_sort,Router,EngineRow/anchors/closed_underorregister_backend.Follow-ups (separate PRs)
heuristics.py_MEASURED_BEHINDbecomes a number rather than a setnote 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
Improvements
Documentation