Add SM120 per-tensor FP8 (e4m3) SDPA-forward engine - #509
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds an SM120 per-tensor E4M3 FP8 SDPA forward engine. The change includes a fused kernel, PTX helpers, API and graph plumbing, engine registration, tile heuristics, tests, and technical documentation. ChangesSM120 FP8 SDPA
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant GraphAnalyzer
participant SdpaFwdDslSm120
participant SM120FusedMultiHeadAttentionForward
participant CUDA
GraphAnalyzer->>SdpaFwdDslSm120: provide FP8 tensors and scale metadata
SdpaFwdDslSm120->>SM120FusedMultiHeadAttentionForward: dispatch validated FP8 execution
SM120FusedMultiHeadAttentionForward->>CUDA: load tiles and execute E4M3 attention
CUDA-->>SM120FusedMultiHeadAttentionForward: produce FP16 output, LSE, and Amax
SM120FusedMultiHeadAttentionForward-->>SdpaFwdDslSm120: return execution results
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 |
YangXu1990uiuc
added
cat-feature
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 6
🧹 Nitpick comments (9)
docs/fe-oss-apis/attention/sdpa-fp8-sm120.md (1)
23-25: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePending benchmark row.
The RTX PRO 6000 numbers are marked "pending". Track this so the document does not ship with a placeholder indefinitely.
Do you want me to open an issue to add the workstation measurements?
🤖 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 `@docs/fe-oss-apis/attention/sdpa-fp8-sm120.md` around lines 23 - 25, Track the pending RTX PRO 6000 Blackwell benchmark in the documentation workflow by opening or linking an issue for the workstation measurements, then replace the placeholder row with the measured numbers once available.python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py (4)
1119-1134: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDead sink branch.
__init__raises whenhas_sinkisTrue(line 225-226), soself.has_sinkis alwaysFalsein this kernel. The branch on lines 1124-1134 and thesinksparameter plumbing can never execute.validate_paramsalso does not rejecthas_sink=Truefor the FP8 template, so the only rejection is the__init__check.Keeping the branch is defensible if a sink variant is planned. If you keep it, add a short note that it is inherited from the f16 sibling and is currently unreachable.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py` around lines 1119 - 1134, Document the currently unreachable has_sink branch in the row normalization logic, noting that it is inherited from the f16 sibling and retained for a future sink variant. Keep the existing __init__ rejection and sinks parameter plumbing unchanged.
806-815: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUntested THD paths in an engine that rejects THD.
SdpaFwdDslSm120.check_supportrejects THD for FP8 (python/cudnn/sdpa/fwd/api_dsl.pyline 1585), and the engine spec does not declarethd. The THD branches in this kernel are therefore unreachable through the supported path and are not covered bytest/python/sdpa/frost/test_sdpa_fwd_fp8_sm120.py. The documentation states "THD: the packed-LSE plumbing is untested with the fp8 epilogue."
validate_params(PARAMS, allowed_dtypes=(DTYPE_E4M3,))does not rejectthd_varlen, so a direct template user can reach this untested code. Consider rejectingthd_varlenin__init__next to thehas_sinkcheck until the path is validated.🛡️ Proposed guard
if has_sink: raise ValueError("has_sink is not supported by the fp8 cell (Amax_S semantics)") + if thd_varlen: + raise NotImplementedError("thd_varlen is not validated on the fp8 cell yet") if thd_varlen and (thd_batch < 1 or thd_max_sq < 1):Also applies to: 837-842, 1243-1250, 1340-1344
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py` around lines 806 - 815, Reject thd_varlen during kernel initialization alongside the existing has_sink validation, before any THD-specific branches execute. Update the relevant constructor validation so direct template users receive an explicit unsupported-configuration error, while preserving non-THD FP8 behavior.
612-639: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRemove the unused
p_regsparameter and correct the docstring.
mma_pvdeclaresp_regsbut never reads it.load_p_fragsreloads P frommma_params.sP_warp(line 638). The docstring on lines 621-622 and 626 still describes register-resident P fragments, which contradicts the SMEM restage the module docstring describes on lines 24-28.online_softmaxreturnss_regs, andcompute_one_kv_tilepasses it asp_regs(line 746), so the return value is also dead.This keeps a reader looking for a register path that does not exist.
♻️ Proposed cleanup
`@cute.jit` def mma_pv( self, basic_params: SimpleNamespace, mma_params: SimpleNamespace, - p_regs: cutlass.Array, ) -> None: """Compute ``O += P @ V``. - P fragments are already packed in registers. V fragments are streamed - from the TMA-populated ``sV`` tile with ``ldmatrix``. + P fragments are reloaded from the warp-private ``sP`` restage tile + with ``ldmatrix``. V fragments are streamed from the TMA-populated + ``sV`` tile with ``ldmatrix``. :param basic_params: Per-CTA tensor metadata and lane mapping. - :param mma_params: Shared V tile and local O accumulator state. - :param p_regs: Register-resident packed P fragments from ``softmax``. + :param mma_params: Shared V tile, warp-private P tile, and local O + accumulator state. """Then update the call site at line 746:
- self.mma_pv(basic_params, mma_params, p_regs) + self.mma_pv(basic_params, mma_params)
online_softmaxcan drop itsreturn s_regsandcompute_one_kv_tileitsp_regs =binding at the same time.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py` around lines 612 - 639, Remove the unused p_regs parameter from mma_pv and update its docstring to describe P fragments being reloaded from mma_params.sP_warp. Remove the dead s_regs return from online_softmax, and change compute_one_kv_tile to call online_softmax without assigning p_regs; preserve the existing softmax and MMA behavior.
36-44: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAlign the constraint list with the enforced envelope.
Line 38 states head dimensions may be any multiple of 16 from 16 to 256. The kernel class accepts that range, but the FP8 engine and adapter gate
D_QK=D_V=128exactly (python/cudnn/sdpa/fwd/engines.pylines 515-516,python/cudnn/sdpa/fwd/api_dsl.pylines 1588-1591). The 8-bit fragment path is validated at d128 only, perdocs/fe-oss-apis/attention/sdpa-fp8-sm120.md. State that only d128 is validated so a direct template user does not read line 38 as a support claim.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py` around lines 36 - 44, Update the constraints documentation in the FP8 prefill kernel to state that D_QK and D_V must be 128, noting this is the only validated head dimension. Remove the broader 16–256 multiple-of-16 support claim while preserving the remaining constraints.python/cudnn/sdpa/fwd/config_sm120.py (1)
10-10: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueThe re-export comment states the FP8 template consumes
DTYPE_E4M3from this module. The FP8 kernel imports it fromcudnn.frost.tile_dsl.constantsdirectly (python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.pyline 56). Confirm which module is the intended source, then keep one path.🤖 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/config_sm120.py` at line 10, Resolve the intended source of DTYPE_E4M3 between config_sm120.py and the FP8 kernel, then keep a single import path. Update the import and its noqa/re-export comment in config_sm120.py or the corresponding import in prefill_fp8_sm120.py so the FP8 template consistently uses that source without a redundant re-export.python/cudnn/sdpa/fwd/api_dsl.py (2)
1869-1872: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winFour separate device-to-host syncs per execute.
_scalarcalls.item()once per descale tensor. That issues four independent device-to-host synchronizations on everyexecutecall. Batch them into one transfer.♻️ Proposed single-sync read
- def _scalar(t, default=1.0): - return float(t.reshape(-1)[0].item()) if t is not None else default - - dq, dk, dv, so = _scalar(descale_q), _scalar(descale_k), _scalar(descale_v), _scalar(scale_o) + _tensors = (descale_q, descale_k, descale_v, scale_o) + _present = [t.reshape(-1)[:1] for t in _tensors if t is not None] + _vals = torch.cat(_present).tolist() if _present else [] + _it = iter(_vals) + dq, dk, dv, so = (float(next(_it)) if t is not None else 1.0 for t in _tensors)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/sdpa/fwd/api_dsl.py` around lines 1869 - 1872, Update the scalar extraction in the execute path around `_scalar` to batch all non-None descale and output-scale tensors into a single device-to-host transfer, then unpack the returned values into dq, dk, dv, and so while preserving defaults for None inputs.
1877-1890: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDuplicated LSE and seq-length validation.
Lines 1877-1890 repeat the LSE presence checks from lines 1784-1791 and the
seq_kv_lensdummy construction from lines 1803-1811. Extract a small helper so the two paths cannot drift.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/sdpa/fwd/api_dsl.py` around lines 1877 - 1890, Extract the duplicated LSE presence/view validation and seq_kv_lens handling from the current path and the earlier path into a shared helper. Update both callers to use that helper, preserving the existing _checked_lse_view, _checked_seq_lens, and _dummy behavior and validation messages so the paths cannot diverge.test/python/sdpa/frost/test_sdpa_fwd_fp8_sm120.py (1)
89-89: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winThe
statsoutput is bound but never asserted.Line 89 allocates
lse, line 120 declaresstatsas a graph output, and line 130 binds it. No test reads it. The engine spec declaresstats=Trueandlse_optional=True, anddocs/fe-oss-apis/attention/sdpa-fp8-sm120.mdstates LSE agrees with the reference to ~1e-6. The kernel's LSE path also has a distinct-inftrim branch for padded rows (python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.pylines 1180-1186) that no test exercises.Add an LSE assertion against the reference row log-sum-exp, at least in the mask and padding tests.
💚 Proposed reference LSE
def _ref(qd, kd, vd, *, scale, is_causal=False, bottom_right=False, swa_window=None, seq_lens_kv=None): @@ scores = scores.masked_fill(masked, float("-inf")) probs = torch.softmax(scores, dim=-1) - return torch.matmul(probs, v_e), probs.max().item() + lse_ref = torch.logsumexp(scores, dim=-1) + return torch.matmul(probs, v_e), probs.max().item(), lse_refThen compare
lse.squeeze(-1)againstlse_refwith an fp32 tolerance in_check.As per path instructions: "Compare test results against a reference implementation using existing reference-module patterns and dtype-appropriate tolerances."
Also applies to: 130-130
🤖 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_fp8_sm120.py` at line 89, Update the affected mask and padding tests in the SDPA FP8 test flow to validate the bound stats output: compute the reference row log-sum-exp using the existing reference implementation, then pass lse.squeeze(-1) and lse_ref through _check with an fp32-appropriate tolerance. Preserve the existing output assertions while exercising padded-row -inf trimming.Source: Path instructions
🤖 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 `@docs/fe-oss-apis/attention/sdpa-fp8-sm120.md`:
- Around line 89-93: Fix the Markdown line wrapping in the fp8 O epilogue note
so “store-path” remains a single hyphenated word when rendered; rewrap the
sentence or place the hyphen at the end of the preceding line without changing
the text.
In `@python/cudnn/sdpa/fwd/api_dsl.py`:
- Around line 1869-1874: The FP8 execution path must not silently ignore graph
operands. In python/cudnn/sdpa/fwd/api_dsl.py:1869-1874, update _execute_fp8 to
honor descale_s and scale_s when computing softmax_scale_log2 and o_scale_fused,
or reject non-unit values through check_support. In
python/cudnn/sdpa/fwd/api_dsl.py:1746-1767, forward seq_q_lens into _execute_fp8
and honor shorter per-batch seq_len_q values, or reject such graphs in
check_support.
- Around line 1901-1907: Update the amax buffer setup around amax_s_buf and
amax_o_buf to preserve views into the caller’s tensors for non-contiguous
inputs, rather than using reshape paths that may allocate temporaries. Execute
both o_view.copy_(o_scratch) and amax_o_buf.div_ within
_torch_stream_context(current_stream, device), ensuring they are ordered after
the launch-stream kernel writes. Keep the existing zero-initialization ordering
unchanged.
In `@python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py`:
- Around line 1217-1223: Update the d_frag_pair processing to apply fmul2
pairwise across the eight-element slice: iterate over the four FP32 pairs, pass
each pair to fmul2, and process each resulting pair when computing lane_amax_o.
Preserve the existing row_sum_inv_vec half selection and row_valid masking, and
use cutlass.Array slicing with explicit offset and count semantics.
In `@test/python/sdpa/frost/test_sdpa_fwd_fp8_sm120.py`:
- Around line 11-13: The module docstring overstates negative-test coverage.
Either add negative tests covering FP8 output dtype, non-128 head dimension, and
attention-sink rejection in the existing SM120 test suite, or revise the
docstring to mention only the negative cases actually implemented, including
test_fp8_sm120_e5m2_not_offered.
- Around line 152-156: Rename the ambiguous O variable to out throughout _check
and every _run result assignment at the listed call sites, updating
corresponding references while leaving O_ref and all validation behavior
unchanged.
---
Nitpick comments:
In `@docs/fe-oss-apis/attention/sdpa-fp8-sm120.md`:
- Around line 23-25: Track the pending RTX PRO 6000 Blackwell benchmark in the
documentation workflow by opening or linking an issue for the workstation
measurements, then replace the placeholder row with the measured numbers once
available.
In `@python/cudnn/sdpa/fwd/api_dsl.py`:
- Around line 1869-1872: Update the scalar extraction in the execute path around
`_scalar` to batch all non-None descale and output-scale tensors into a single
device-to-host transfer, then unpack the returned values into dq, dk, dv, and so
while preserving defaults for None inputs.
- Around line 1877-1890: Extract the duplicated LSE presence/view validation and
seq_kv_lens handling from the current path and the earlier path into a shared
helper. Update both callers to use that helper, preserving the existing
_checked_lse_view, _checked_seq_lens, and _dummy behavior and validation
messages so the paths cannot diverge.
In `@python/cudnn/sdpa/fwd/config_sm120.py`:
- Line 10: Resolve the intended source of DTYPE_E4M3 between config_sm120.py and
the FP8 kernel, then keep a single import path. Update the import and its
noqa/re-export comment in config_sm120.py or the corresponding import in
prefill_fp8_sm120.py so the FP8 template consistently uses that source without a
redundant re-export.
In `@python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py`:
- Around line 1119-1134: Document the currently unreachable has_sink branch in
the row normalization logic, noting that it is inherited from the f16 sibling
and retained for a future sink variant. Keep the existing __init__ rejection and
sinks parameter plumbing unchanged.
- Around line 806-815: Reject thd_varlen during kernel initialization alongside
the existing has_sink validation, before any THD-specific branches execute.
Update the relevant constructor validation so direct template users receive an
explicit unsupported-configuration error, while preserving non-THD FP8 behavior.
- Around line 612-639: Remove the unused p_regs parameter from mma_pv and update
its docstring to describe P fragments being reloaded from mma_params.sP_warp.
Remove the dead s_regs return from online_softmax, and change
compute_one_kv_tile to call online_softmax without assigning p_regs; preserve
the existing softmax and MMA behavior.
- Around line 36-44: Update the constraints documentation in the FP8 prefill
kernel to state that D_QK and D_V must be 128, noting this is the only validated
head dimension. Remove the broader 16–256 multiple-of-16 support claim while
preserving the remaining constraints.
In `@test/python/sdpa/frost/test_sdpa_fwd_fp8_sm120.py`:
- Line 89: Update the affected mask and padding tests in the SDPA FP8 test flow
to validate the bound stats output: compute the reference row log-sum-exp using
the existing reference implementation, then pass lse.squeeze(-1) and lse_ref
through _check with an fp32-appropriate tolerance. Preserve the existing output
assertions while exercising padded-row -inf trimming.
🪄 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: 853c7542-99fb-4302-9f8b-5fe1f76ec58b
📒 Files selected for processing (8)
docs/fe-oss-apis/attention/sdpa-fp8-sm120.mdpython/cudnn/engines/manifest.pypython/cudnn/frost/tile_dsl/mma.pypython/cudnn/sdpa/fwd/api_dsl.pypython/cudnn/sdpa/fwd/config_sm120.pypython/cudnn/sdpa/fwd/engines.pypython/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.pytest/python/sdpa/frost/test_sdpa_fwd_fp8_sm120.py
5a26429 to
7a2d7f4
Compare
|
Rebased onto Offer the tile domain as plans, and pick by shape
FP8 only — the f16 cell wants MeasurementsRTX PRO 6000 Blackwell (sm120, 188 SMs), cuDNN 9.25.0.15, CUPTI kernel time,
One note on the baseline in the first commit messageThe 1.74–1.79x there is against the bf16 SM120 kernel, which is the right Also, the two paths turn out to overlap less than the "sibling" framing Unrelated and pre-existing, filed separately as #510: the torch glue around note to self: claude::304e9e55-1db7-4285-967f-001cb21032f3 — "审计 Frost Python DSL 引擎调用流程" |
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python/cudnn/sdpa/fwd/engines.py`:
- Line 558: Update the return annotation of knob_candidates to use a defined
List import or the built-in list[...] syntax, ensuring Optional and SdpaFwdKnobs
remain represented and typing.get_type_hints() resolves successfully.
In `@test/python/sdpa/frost/test_sdpa_fwd_fp8_sm120.py`:
- Around line 127-145: Extend the SDPA FP8 test helpers by making _ref return
torch.logsumexp(scores, dim=-1), then thread the produced lse and reference LSE
through _run and _check. In _check, compare the actual and reference LSE tensors
using the existing reference-comparison pattern with dtype-appropriate
tolerances, while preserving the current output and amax validations.
🪄 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: 94498916-8902-43a5-bfae-44ecf1495904
📒 Files selected for processing (6)
python/cudnn/sdpa/fwd/api_dsl.pypython/cudnn/sdpa/fwd/config_sm120.pypython/cudnn/sdpa/fwd/engine.pypython/cudnn/sdpa/fwd/engines.pytest/python/sdpa/frost/test_sdpa_fwd_fp8_sm120.pytest/python/sdpa/frost/test_sm120_fp8_tiles.py
🚧 Files skipped from review as they are similar to previous changes (1)
- python/cudnn/sdpa/fwd/api_dsl.py
|
Went through the review findings against the current code. Four fixed, one Fixed
Docstring promised negative tests that did not exist (Minor). Added Writing the first of those turned up something the review did not ask about: E741
Not a defect"Apply Real, but not fixed here
Closing it means either folding both scalars ( note to self: claude::304e9e55-1db7-4285-967f-001cb21032f3 — "审计 Frost Python DSL 引擎调用流程" |
…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.
…_sort (#528) * Take the backend's plans one heuristic mode at a time Ranking the two sides against each other needs to know which backend entries are mode-A recommendations and which are fallbacks -- "the backend's A ahead of ours, its fallbacks behind" cannot be said about one opaque list. Until now the whole thing arrived from a single create_execution_plans([A, FALLBACK]). No C++ change is needed. C++ appends each query to the same plan list, and get_execution_plan_count() already exists, so asking one mode at a time and reading the count after each gives the boundaries. Measured on a 512^3 bf16 matmul (sm90, cuDNN 9.25): A -> plans[0:15], all knob-bearing; FALLBACK -> plans[15:17], bare eng0/eng7 with no knobs; the two segments do not overlap. A mode with no configs raises, which is not a decline while another mode still has entries -- an OPENSOURCE-only query legitimately leaves the cuDNN modes empty. Only every mode failing means the backend has nothing, and then the last error is re-raised so the caller still reports why. * Move plan ranking out of the engines and into one heuristics function An engine cannot rank. It sees neither its siblings nor the backend's entries, so propose_plans could only ever order its own knobs -- and then something downstream had to merge the two sides anyway, which heuristics_sort did by concatenating and calling it ranking. All four in-tree propose_plans were the base class's default copied verbatim: the hook has never decided anything. create_execution_plans() now gathers the inputs (parsed facts, the family's offered ids, the backend's entries tagged by mode) and hands all of it to the graph's family in ONE call. What comes back IS graph.plans, position for position. An engine answers two questions: can I serve this graph (check_support), and compile me this config (build_plan). sdpa/fwd/heuristics.py is the first such hook, and it is deliberately a frame with no tuning in it: one entry per eligible cell at the config its capability row declares. Mode A and FALLBACK differ only in which backend entries they carry; OPENSOURCE is mode A without the backend's recommendation, since these cells ARE the open-source implementation. Real per-cell rules land on top. Deleted, all superseded or never used: BaseEngine.propose_plans + its 4 implementations BaseEngine.default_knobs only fed propose_plans heuristics_sort merging is part of ranking, not a step after engines/router.py entirely Router / default_router / set_router / pygraph(router=) -- policy has one home now, and decline_types moved to base.py where the engine contract already lives engines.probe() (fwd + bwd) superseded by check_support graph.engine pure alias of selected_engine, zero callers graph.from_serialized zero callers; serialize/deserialize are the pybind-era API and stay knobs=None no longer means "engine, pick for me" -- the heuristics name a concrete config. A None field survives only on an axis whose capability row declares no domain. That reading is what let one choice be made twice, once when ranking and once inside the adapter. * Update the dispatch tests to the ranking contract, and delete what it retired Ranking has one home, so a test that wants a specific order replaces heuristics.rank instead of subclassing Router. The _ranking() helper does that; it is the same monkeypatch idiom the rest of the suite already uses. Deleted rather than translated: test_set_router_frozen_after_planning the API it tested is gone test_a_claiming_engine_is_tried_before_the_backend asserted that python plans always outrank the backend, which is a per-cell measurement, not a rule. FROST coverage rides on heur_mode.OPENSOURCE instead: ask for it and any graph still landing on a backend plan is one FROST cannot serve Renamed for what they now test: test_mixed_ranking_dispatch, test_empty_ranking_output_rejected, test_mixed_ranking_backend_slot_executes, test_constructor_backends_validated_and_ranking_ids_checked. One assertion changed meaning: the backend is queried once PER MODE now, so _create_backend_plans records two create_execution_plans calls for [A, FALLBACK]. test_sdpa_graph_analyzer called engines.probe() twice; those two call analyze_for directly, so no production API exists only for tests. Six sdpa test files each carried a verbatim copy of _select_engine matching a bare engine name. Plans now read <engine>[<knobs>] because the heuristics name a concrete config for every entry, so they share frost_test_utils.select_engine, which matches on the engine. 208 passed. test_a_replayed_plan_reports_its_own_notes still fails and also fails on develop without this change -- C++ on 9.25 no longer raises for an index one past the plan count. * Update the design doc, and align the remaining sdpa test helpers The doc still described a Router with three pluggability levels, engines that propose their own plans, and heuristics_sort as the seam a cost model replaces. Rewritten to what dispatch now does: one call per graph into the family's heuristics hook, the backend's entries tagged by the mode that produced them, and heur_mode.OPENSOURCE as the way FROST coverage is measured rather than assumed. Also states plainly what register_backend is and is not. It installs an engine instance on one graph -- the hatch tests use to inject a fake. It does not make an engine rankable: an out-of-tree engine declares no Capabilities, so nothing can enumerate its configs or place it against the backend. The follow-up list now names removing that concept, since an engine id is decodable from the manifest alone. Six sdpa test files each carried a verbatim copy of _select_engine matching a bare engine name; the shared frost_test_utils.select_engine matches on the engine, which is what plan names now carry a config suffix for. 208 passed locally. The one failure is test_a_replayed_plan_reports_its_own_notes, which fails on develop without this change too. * Decode an engine id from the manifest, with nothing registered first An engine id is fully decodable from the manifest: the family owning the id block, then the slot within it. _owners_for_id only ever looked inside the graph's candidate set, so an id could be resolved only if something had already put that engine there -- which made register_backend look like a prerequisite for create_execution_plan() when it is really just one way to supply an instance. engine_for_id() closes that. _owners_for_id falls back to it, so replaying a recorded (engine_id, knobs) works on a fresh graph, including for an engine that is not a candidate for THAT graph -- there the replay is a deliberate pin, not a routing decision. A gated-off slot still resolves to None rather than being built. Groundwork for removing the out-of-tree engine concept entirely. * Remove the out-of-tree engine concept: the manifest is the only way in Every python engine now exists exactly one way. register_backend, pygraph(backends=), graph.backends and OUT_OF_TREE_ID_BASE are gone, and _candidate_engines() is the graph's family and nothing else. An out-of-tree engine could never be RANKED anyway: it declares no Capabilities, so nothing could enumerate its configs or place it against the backend. It was an entry point into the plan list, not into the decision. And being a candidate had nothing to do with fitness -- an engine was in the list because someone had registered it, so an engine that could not serve the graph was still tried, and failed at build instead of at classification. The linear_attention suites used register_backend to PIN an implementation -- cuTile rather than FROST. That is not what registration is for, and those engines are in the manifest already, so the pin is now by name and applied after planning through select_plan(): engine_utils.pin_engines() / apply_pin(). apply_pin raises when the pinned engine produced no plan, so a pin that stops working fails the first op call. The cutile conftest used to check the pin by inspecting the CANDIDATE list, which passes whether or not the pin took effect -- which is how it ran for months against whichever engine the ranking picked while the seam it pinned through was dead. That check is deleted; the pin enforces itself. heuristics: no engine sits outside a family now, so the "family-less engines go last" branch is gone and _without_a_family is _unranked -- the case it covers is a family that declares no heuristics hook, not an engine with no family. test_engine_router.py -> test_dispatch.py. It never tested a Router; it tested dispatch -- one plan list, the at-index APIs, select_plan's strict pin, one-shot planning, how a decline advances the walk, note filters reaching python plans, manifest classification, facts attachment. _offer(monkeypatch, *engines) replaces register_backend by putting the fakes in a manifest family, so the tests reach engines through the same path production does. Six tests deleted with the concept they tested -- all checked registration-time id validation, which has no subject now that engines never declare their own ids: test_register_backend_validation, test_engine_id_in_the_in_tree_region_is_rejected, test_a_registered_in_tree_engine_is_not_offered_twice, test_overlapping_declared_id_blocks_are_rejected, test_a_lying_owns_id_cannot_capture_another_engines_plans, test_constructor_backends_validated_and_ranking_ids_checked. What they protected is covered by test_family_id_blocks_are_disjoint and test_every_engine_spec_has_a_manifest_slot. BaseEngine.owns_id goes with them: zero callers, and its docstring already called it a convenience. Three tests needed real thought rather than a mechanical edit: - The "no family, no facts payload" test built a bare relu graph. relu names no family, so there is no python candidate, and the backend declines a 2-D pass-by-value tensor -- planning raised before the assertion. The claim under test is about the payload, not about the graph being servable. - The mutable-after-validate window was `not self._backends`: validate() lowers and freezes any graph the backend CAN lower, and registering an engine was the only way to skip that. With registration gone the window is exactly the ops with no backend lowering, which is what the property was always about. - test_api_signature_parity asserted {"backends", "router"} were keyword-only. Both are gone, so the assertion had no subject; what it protected is that nothing pygraph-only is POSITIONAL, which is now asserted directly. TorchMatmulEngine goes too. It reimplemented matmul, bias and relu in torch inside a dispatch test: the numeric assertions proved torch, not dispatch, and "torch_matmul" in plan names reads like something cuDNN ships. StubEngine replaces it -- same claim on the graph, no arithmetic, and it RECORDS what dispatch handed it, so the fusion test now asserts what was only implied before: every node arrives in build order, each input port resolved to the caller's storage, and the virtual intermediate carrying none. * Give the SM120 SDPA-forward cell a real tile rule, as the worked example The framework had no rule in it: every cell went to `_sole()` on each knob axis, which answers None the moment a row declares more than one value. The SM120 prefill row declares tile_ms={64,128}, tile_ns={64,128}, so its choice fell through to api_dsl's `_SM120_Q_TILES[0]` default -- the choice being made in the adapter is exactly what moving ranking out of the engines was meant to stop, and it left the frame with nothing showing how a rule is added. _sm120_tiles(facts) is that rule, and it is measured rather than invented: regret 1.009 geomean / 1.054 worst against the best of the enumerated domain. tile_n=128 always; tile_m=64 when the grid cannot fill the machine AND each CTA has enough KV tiles to amortize the extra Q-tile loop, with a causal mask counted as a halved effective grid because it halves the work per CTA. It reads facts and nothing else -- device_sm_count is already on the record. Shape a colleague can copy: write the function, list the cell in _TILE_RULE_CELLS, put the measurement in the commit. A cell absent from that set keeps the old behaviour (its row's sole point per axis), which is the honest answer when nobody has timed it. Mode A now emits the guess FIRST and the rest of the domain behind it, so a caller who autotunes has the runners-up and a caller who does not gets the best guess at index 0. FALLBACK takes the smallest tile the row admits -- the config that asks least of the device; picking real fallback configs per cell is a TODO left in the file. * Ask one function whether an SM120 tile fits, not two Naming tile_n=128 unconditionally broke D=208/224/240/256: the adapter's own `if self.tile_n is None` branch was quietly shrinking the KV tile to whatever fit SMEM, so leaving the knob None had been answering a CAPABILITY question, not a tuning one. Requesting a value skips that branch, and the request then fails the very check the branch existed to satisfy -- 106512 bytes wanted against the part's 101376. The fit arithmetic moves to config_sm120.smem_bytes(), beside the template it describes, and both callers use it: the adapter's check and the ranking's choice. The rule now reads "tile_n = the largest that fits, tile_m by occupancy", and the runners-up it offers are filtered the same way -- a config the kernel cannot fit is not a runner-up, it is an entry that sits in the list to decline at build. test_api_signature_parity asserted {"backends", "router"} were keyword-only. Both are gone, so the assertion had no subject; what it protected is that nothing pygraph-only is POSITIONAL, which is now asserted directly. * Query the backend for the modes the ranking will actually place The default mode list was written out twice -- once in _create_backend_plans, once in heuristics.default_modes. They agree today; a change to one alone would have the backend enumerate plans for a mode no family places, which reads as the family losing entries rather than as the query asking for the wrong thing. * Restore the #512 SDPA tests this branch had silently reverted Four test files were carrying their PRE-#512 content while the production code they exercise is post-#512. The branch is cherry-picked onto the github develop, and the commit that consolidated the sdpa test helpers was authored against a tree from before #512 landed -- so the cherry-pick took the whole file, not the helper edit, and reverted #512's test additions with it. api_dsl.py and engines.py were untouched by that, which is why nothing looked wrong until an SM100 box ran the suite: 16 failures, all of them tests asserting the old contract against the new kernels (a stats-less SM100 graph now carves a dummy LSE, so get_workspace_size() is b*h*s*4, not 0). Restored all four from gh/develop and re-applied only what this branch meant to change: - test_sdpa_fwd_dsl_sm100 / _sm120: the local verbatim copy of _select_engine -> frost_test_utils.select_engine. - test_sdpa_frontend_integration: plan-name lookups made suffix-aware. The heuristics now name a concrete config for every entry, so a plan reads "<engine>[<knobs>]" and names.index(_FROST) raises ValueError. - test_sdpa_graph_analyzer: engines.probe() is deleted, so _eligible asks analyze_for(...)[1] is None. The ragged-Stats coverage #512 added (token-major and head-major layouts, zero-length sequences, the analyzer acceptance test, the strict LSE presence contract in both directions) is back verbatim. * Bring the design doc to the architecture as it now stands The dispatch tree, written out: what create_execution_plans does in order, where the backend's per-mode entries come from, and where a family's rules sit. That tree was the first thing anyone asked for and the doc did not have it. Corrects three things the doc stated as settled that this PR changed: the delegating entry leads the BACKEND's block and not the family's (it falls through to native configs when the C++ OSS engine declines, so ahead of an OPENSOURCE block it answers a coverage question with a native kernel); a plan's identity is (engine_id, knobs) and never its cpp_index; whether a heuristic mode succeeded is tracked per call, not inferred from plan spans. Adds what each machine covers. The suites SKIP on the wrong arch rather than fail, so a green sweep on one box says nothing about the others -- defaulting to CUDA device 0 is how a whole SM100 run silently skips. Follow-ups now name what is actually left: one tuning rule exists, FALLBACK is a placeholder, _MEASURED_BEHIND is empty by design. * Answer the backend's plan query once per distinct config Two findings from the second review pass, both about APIs whose callers moved under the branch. Graph::create_execution_plans checks override_heuristics_query() FIRST and returns before it reads the mode at all -- deterministic SDPA backward and FP8 backward both override. Asking one mode at a time therefore appends the SAME engine-17 config once per mode, and backend_plan_entries() handed all of them back. SDPA forward's recommend() would have deduped them; SDPA BACKWARD declares no heuristics hook, so _unranked passed the duplicates straight into graph.plans and build_plans(ALL) or an autotuner would compile and time one config twice. Deduped at collection instead of in each family: a repeated (engine, knobs) in the backend's own list is never two different things, and the first index is the one whose mode span is real. test_dsl_sm100_band_right_uncovered_tail_rejected called fwd_engines.probe(). That test arrived with #485, which this branch rebased onto after probe() was already deleted here -- so it is a caller that did not exist when the deletion was written, and it would have taken out the whole Blackwell L0 suite with an AttributeError before reaching its assertion. * Address the CodeRabbit pass: an id names one engine, and two tests could not fail engine_for_id() matched an id exactly while _owners_for_id() matched a RANGE, so a replay could resolve one way for a candidate engine and another way on a fresh graph. Collapsed the other way from what was suggested: BaseEngine.id_end and owned_id_range are deleted and _owners_for_id is an equality test. The range existed so a REGISTERED engine could claim a block and registration could prove two blocks disjoint; nothing registers now, no shipped engine ever set id_end, and every range was [engine_id, engine_id + 1). Keeping it would have spread dead machinery to fix an asymmetry that only that machinery created. EngineFamily.id_end -- the family's block -- is a different thing and stays. test_ranking_and_engine_read_the_same_record declared a probe_family by hand and then called _offer(), whose own monkeypatch of MANIFEST won; the surviving family had no analyzer. It passed anyway because both sides call _facts_for(_probe_analyzer) directly, so the documented claim -- that the ranking resolves the analyzer from EngineFamily.analyzer -- went unexercised. Now declared through _offer, and it asserts the analyzer already ran BEFORE ranking, which is the part only planning can do. Verified by mutation: drop the analyzer declaration and the test fails. select_engine(tiles=) matched the rendered plan name by substring, so a request for tile_n=128 could select a tile_n=1280 plan and the test would pass having run something else. Matches PlanConfig.knobs structurally now. No caller on this branch -- the fp8 SM120 tile tests in #509 are the first, and they would have been the ones to hit it. Plus a cross-reference to a test renamed in this PR. 204 passed on the CPU suites; the one failure is the pre-existing test_a_replayed_plan_reports_its_own_notes. * License the one file this PR adds as Apache-2.0, not MIT sdpa/fwd/heuristics.py was created by copying the header from config_sm120.py, which is MIT -- so the new file inherited a tag that does not apply to it. Per LICENSING.md, the repo relicensed MIT -> Apache-2.0 in #408 and a file is kept under MIT for exactly two reasons: surviving lines from an external contributor who has not consented to relicensing, or derivation from third-party source. A file written from scratch at NVIDIA has neither, so Apache-2.0 is the correct tag -- as it already is on every other file this change adds content to (engines/heuristics.py, engines/manifest.py, sdpa/fwd/engine.py, and the tests). The MIT neighbours in sdpa/fwd are pre-existing files this PR only edits, and editing does not move a file between licenses. Also switches to the SPDX-FileCopyrightText form the Apache-2.0 files use. * Compress the comments this PR adds Six blocks broke the house rule that a call site explains only the non-obvious load-bearing fact and rationale/measurements go in the MR description -- which is where all of this already was, so it was duplicated, not lost. Cut: measurement detail from _sm120_tiles (1.5x at 64 CTAs, 240-vs-320 CTAs, 2-4%, 106 KB vs 99) down to a pointer at PR #528, keeping the two thresholds a reader needs and the warning that the rule is kernel-specific. _MEASURED_BEHIND lost two antitheses ("deliberate but NOT a measurement", "an experiment, not an edit") and a cross-reference the module docstring already makes. _owners_for_id, _create_backend_plans and the backend-dedup comment lost restated clauses. BaseEngine's note on the deleted id range stopped narrating the deletion -- that belongs to the commit that made it, where it is verbatim. 31 fewer added lines, 8 fewer comment lines; no claim, number or caveat dropped, only relocated to where it was already written. 199 passed; the one failure is the pre-existing test_a_replayed_plan_reports_its_own_notes.
d7dd85e to
5203c65
Compare
|
Rebuilt on the merged What changed structurallyRebuilt rather than rebased, because two things moved underneath this branch:
Sharing that rule is a measurement, not an assumption — 30 causal and non-causal shapes on a 188-SM part give 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: Testing
|
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
python/cudnn/sdpa/fwd/heuristics.py (1)
126-127: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
domaincomputes SMEM with the wrong element sizes for FP8.Line 81 passes
qkv_itemando_itemtosmem_bytes. Line 126 omits both, so the filter usesitemsize=2for the KV term on an FP8 graph. The KV tile is one byte per element there, so the filter overstates the footprint.Two consequences:
- The autotune candidate list can drop a tile combination the FP8 kernel fits.
- If the filter drops
best, the sort keymn != bestanchors on a value that is not inordered, and the rule's own choice disappears from the plan list.The FP8 row is d128-only today, so no combination is dropped yet. Reuse the same element sizes to keep the two call sites consistent.
♻️ Proposed fix
- best = _sm120_tiles(caps, facts) + best = _sm120_tiles(caps, facts) + qkv_item, o_item = (1, 2) if facts.is_fp8 else (2, 2) # The guess first, then the rest of the domain as autotune candidates: # the rule's regret is small but not zero, so the runners-up are worth # offering to a caller who measures. Configs the kernel cannot fit are # not runners-up -- they would sit in the list only to decline at build. - domain = [(m, n) for m in caps.tile_ms for n in caps.tile_ns if smem_bytes(facts.d_qk, facts.d_v, m, n) <= SMEM_CAPACITY_BYTES] + domain = [ + (m, n) + for m in caps.tile_ms + for n in caps.tile_ns + if smem_bytes(facts.d_qk, facts.d_v, m, n, qkv_item, o_item) <= SMEM_CAPACITY_BYTES + ]🤖 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/heuristics.py` around lines 126 - 127, Update the smem_bytes call inside the domain comprehension to pass the same qkv_item and o_item element-size arguments already used by the earlier call, preserving accurate FP8 KV footprint filtering and keeping best eligible for ordering.
🧹 Nitpick comments (4)
python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py (2)
669-682: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueThe prefetch comment does not describe this loop.
Lines 669-672 state that the next
v_frag's fragments are issued ahead of the current MMAs and that the index wraps so the last iteration reloadsv_frag0. The loop below issuesload_v_frags(v_frag, d_frag_pair)inside thed_frag_pairbody for the currentv_fragonly. There is no wrapped index and no software prefetch.Update the comment to match the code, or restore the prefetch if it was intended.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py` around lines 669 - 682, Update the comment immediately above the v_frag loop to describe the current execution: load_v_frags uses the current v_frag within each d_frag_pair iteration, with no ahead-of-time or wrapped prefetch. Do not claim an unconditional prefetch or final reload unless the loop is changed to implement that behavior.
1232-1237: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winReduce
Amax_Oto one atomic per warp.All compute lanes currently update the same
amax_oaddress. Use five butterfly shuffles, following the repository’s full-warp reduction pattern, then issue the atomic only whenlane == 0. This reduces atomics by 32×.threads_computeis 128 lanes forq_tile=64and 256 forq_tile=128;lane_amax_ois non-negative, so the bitcastInt32maximum remains exact.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py` around lines 1232 - 1237, Update the Amax_O reduction near amax_o_arr and the prims.AtomicOp.MAX call to perform a full-warp reduction of lane_amax_o using five butterfly shuffles, following the repository’s established pattern, then issue the atomic update only when lane == 0. Preserve the Int32 bitcast because lane_amax_o is non-negative, and ensure the logic supports both 128- and 256-lane threads_compute configurations.python/cudnn/sdpa/fwd/api_dsl.py (1)
1994-2010: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueThe LSE presence checks and the stream fallback repeat the caller's logic.
executealready runs the same twolse_desc/lse_tensorchecks at lines 1857-1864 before it dispatches here, and the samecurrent_stream is Nonefallback exists at lines 1926-1930. The duplication is harmless today, but the two copies can drift.Drop the repeated checks here and resolve
current_streamonce inexecutebefore the FP8 branch.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/sdpa/fwd/api_dsl.py` around lines 1994 - 2010, Remove the duplicated lse_desc/lse_tensor validation and current_stream fallback from the FP8 execution path around _checked_lse_view and seq_kv_t. Keep these checks in execute, and resolve current_stream there before dispatching to the FP8 branch so the downstream path reuses the single normalized value.python/cudnn/sdpa/fwd/engines.py (1)
699-709: 🚀 Performance & Scalability | 🔵 Trivial | ⚖️ Poor tradeoffAvoid the extra device-to-host read for quantized Q padding. Dense FP8/MXFP8 plans force
seq_q_lens_present=False, so this guard readsseq_q_buf.min().item()on every execute, including full-Q cases such astest_fp8_sm120_paddingthat only pad KV. Read the minimum once on the failure path. Treat CUDA-graph capture as unsupported for this lowering because the quantized adapter also calls.item()for scale tensors; this guard is not the sole capture blocker.🤖 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` around lines 699 - 709, The quantized Q-padding guard in the execution path currently performs seq_q_buf.min().item() on every run. Restructure the check around the existing FP8/MXFP8 and seq_q_lens_present conditions so the minimum is read only on the failure path when shorter per-batch Q lengths must be detected, while preserving the current error for unsupported padding. Mark this quantized lowering as unsupported during CUDA-graph capture, alongside the existing scale-tensor .item() limitation.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python/cudnn/sdpa/fwd/api_dsl.py`:
- Around line 341-355: Update _amax_slot to validate that a caller-provided
tensor has the required 4-byte element type before returning its view, matching
the dtype validation used by _checked_lse_view and _checked_seq_lens. Reject
incompatible dtypes with the established validation error behavior, while
preserving the existing dummy allocation and contiguous view handling.
In `@python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py`:
- Around line 195-226: Update the `__init__` docstring for `in_dtype` and
`out_dtype` to describe the actual FP8 contract enforced by the checks: Q/K/V
use `cutlass.Uint8` storage and O must be `cutlass.Float16`; remove the
contradictory Float16/BFloat16 and matching-dtype descriptions while leaving the
validation logic unchanged.
In `@test/python/sdpa/frost/test_sdpa_fwd_fp8_sm120.py`:
- Around line 11-14: Update the module docstring to reference the module-level
decline tests instead of the nonexistent TestNotOffered class, and add the
requested L0 test_fp8_sm120_sink_not_offered test alongside the other negative
tests. Have it assert that _fp8_graph_offers_sm120 rejects FP8_E4M3 input with
HALF output when sink=True, using the existing RNG decorator and imports.
---
Outside diff comments:
In `@python/cudnn/sdpa/fwd/heuristics.py`:
- Around line 126-127: Update the smem_bytes call inside the domain
comprehension to pass the same qkv_item and o_item element-size arguments
already used by the earlier call, preserving accurate FP8 KV footprint filtering
and keeping best eligible for ordering.
---
Nitpick comments:
In `@python/cudnn/sdpa/fwd/api_dsl.py`:
- Around line 1994-2010: Remove the duplicated lse_desc/lse_tensor validation
and current_stream fallback from the FP8 execution path around _checked_lse_view
and seq_kv_t. Keep these checks in execute, and resolve current_stream there
before dispatching to the FP8 branch so the downstream path reuses the single
normalized value.
In `@python/cudnn/sdpa/fwd/engines.py`:
- Around line 699-709: The quantized Q-padding guard in the execution path
currently performs seq_q_buf.min().item() on every run. Restructure the check
around the existing FP8/MXFP8 and seq_q_lens_present conditions so the minimum
is read only on the failure path when shorter per-batch Q lengths must be
detected, while preserving the current error for unsupported padding. Mark this
quantized lowering as unsupported during CUDA-graph capture, alongside the
existing scale-tensor .item() limitation.
In `@python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py`:
- Around line 669-682: Update the comment immediately above the v_frag loop to
describe the current execution: load_v_frags uses the current v_frag within each
d_frag_pair iteration, with no ahead-of-time or wrapped prefetch. Do not claim
an unconditional prefetch or final reload unless the loop is changed to
implement that behavior.
- Around line 1232-1237: Update the Amax_O reduction near amax_o_arr and the
prims.AtomicOp.MAX call to perform a full-warp reduction of lane_amax_o using
five butterfly shuffles, following the repository’s established pattern, then
issue the atomic update only when lane == 0. Preserve the Int32 bitcast because
lane_amax_o is non-negative, and ensure the logic supports both 128- and
256-lane threads_compute configurations.
🪄 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: e88b076e-930e-489d-a8c6-790d75a26663
📒 Files selected for processing (9)
python/cudnn/engines/manifest.pypython/cudnn/sdpa/fwd/api_dsl.pypython/cudnn/sdpa/fwd/config_sm120.pypython/cudnn/sdpa/fwd/engines.pypython/cudnn/sdpa/fwd/heuristics.pypython/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.pypython/cudnn/sdpa/graph_analyzer.pytest/python/sdpa/frost/test_sdpa_fwd_fp8_sm120.pytest/python/sdpa/frost/test_sm120_tile_rule.py
🚧 Files skipped from review as they are similar to previous changes (1)
- python/cudnn/engines/manifest.py
5203c65 to
360e48c
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@test/python/sdpa/frost/test_sdpa_fwd_fp8_sm120.py`:
- Around line 175-178: Update test_fp8_sm120_gqa to parameterize H_kv over 1 and
2, passing the parameter into _run so the test covers both MQA and GQA. Mark the
new parameterized Python test with the appropriate L0–L4 test level according to
its runtime and scope.
- Line 293: Replace the direct plan-name equality check with offers_engine for
both SM120 FP8 plan checks in
test/python/sdpa/frost/test_sdpa_fwd_fp8_sm120.py:247-247 and
test/python/sdpa/frost/test_sdpa_fwd_fp8_sm120.py:293-293, preserving the
existing engine_name(arch="sm120", fp8=True) target so names with knob suffixes
are recognized.
🪄 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: a5b77927-d11c-4884-969f-06218f37c8e1
📒 Files selected for processing (5)
python/cudnn/sdpa/fwd/api_dsl.pypython/cudnn/sdpa/fwd/engines.pypython/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.pypython/cudnn/sdpa/graph_analyzer.pytest/python/sdpa/frost/test_sdpa_fwd_fp8_sm120.py
🚧 Files skipped from review as they are similar to previous changes (4)
- python/cudnn/sdpa/graph_analyzer.py
- python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py
- python/cudnn/sdpa/fwd/engines.py
- python/cudnn/sdpa/fwd/api_dsl.py
360e48c to
e3af62c
Compare
P quantization implemented, per the FORT conventionFollowing up on the open scope question — implemented rather than shipped as v1. cuDNN’s The kernel now matches the backend’s FORT ordering — amax on the unscaled softmax result, then scale, then cast:
SM100 deliberately keeps the decline: its kernel has no The test that can actually fail
Cost~0.7–0.9% at large shapes — six of six 8192-cells positive in a 0.65–0.92% band. The 30-cell geomean (1.0089) sits at the ~1% run-to-run floor and is not separately informative. A free version exists and is deliberately left for its own change: fold Correction to my earlier numbersThe tile-rule regret I posted before ( Final verification
|
|
@cudnn-ci-bot run frost |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-509-e3af62c |
Scale_S scope: SM120 applies it, SM100 declines a non-reciprocal pairFollowing up on the P-quantization question. Implementing Why it fails. That kernel uses a lazy rescale: the running max is only refreshed when a tile exceeds it by This is the same constraint that makes the cuDNN backend's SM100 path ignore the pair. Why there is nothing to win. Sweeping the scale on B2xH8xS256 e4m3:
Flat to the digit from 1 to 64. e4m3 is floating point, so relative precision does not move with scale, and subtracting the row max is already a per-row placement — which strictly dominates a per-tensor scale. The accuracy gain predicted in the design doc does not exist; the doc has been corrected. So:
Validated at 48e70ca on both architectures: SM100 (parley Blackwell) 50 passed fp8+mxfp8 / 567 passed all sdpa; SM120 (RTX PRO 6000 Blackwell Server Edition) 18 passed fp8-sm120 / 179 passed all frost sdpa; 80 passed guards+tile-rule on each. note to self: claude::304e9e55-1db7-4285-967f-001cb21032f3 — "cudnn-FE #509 SM120 FP8 engine + Scale_S scope" |
|
@cudnn-ci-bot run frost |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-509-48e70ca |
| (SASS ``LDSM.8.MT1616``) — one issue covers a 32(kv) x 16(d-bytes) tile and | ||
| feeds two MMAs, register map (0, 2, 1, 3). | ||
| - P: fp32 softmax output packed to e4m3 with ``cvt.rn.satfinite.e4m3x2.f32`` | ||
| and staged through a per-warp SMEM tile (the k32 C->A fragment-column |
There was a problem hiding this comment.
I think P exists only in registers and is synchronized across threads using shfl_sync, rather than being staged in shared memory and read back through ldmatrix. This comment seems stale, and the other comments should also be updated to reflect the current kernel implementation.
There was a problem hiding this comment.
Confirmed, and the code says so in its own words — mma_pv notes that two shfl.sync and one prmt "replace the SMEM round trip (and the 16 KB it needed)". The header was describing a design that got removed. Rewritten to say P stays in registers and the k32 C->A fragment-column mismatch is a thread-quad exchange.
Thanks for the "other comments" nudge — P scale is fixed 1.0 on the next line went stale in this very PR, since Scale_S now multiplies P before the cast. Fixed in the same pass.
| src_format=prims.LoadSrcFormat.B8, | ||
| ) | ||
|
|
||
| # An ldmatrix is always in flight behind the tensor cores: the next |
There was a problem hiding this comment.
Seems stale comment. No prefetch or wrapping index exists in the loop; V loads are issued in-loop right before their MMAs.
There was a problem hiding this comment.
Confirmed. load_v_frags is called inside the inner loop immediately before the MMAs that consume it; there is no wrapping index.
I nearly replaced it with "the extra live registers cost more than the latency it hid", then checked the measurement: V prefetch landed within ±0.5% against a ~1% run-to-run floor, i.e. no effect — because keeping P in registers leaves little ldmatrix latency to hide in the first place. The comment now says that instead.
| scale_o_t: Any = None | ||
| # cuDNN's Scale_S/Descale_S: they quantize P, the softmax OUTPUT. The FROST | ||
| # kernels convert P unscaled, so execute declines a non-unit pair rather | ||
| # than dropping it (api_dsl._require_unit_s_scales). |
There was a problem hiding this comment.
_require_unit_s_scales -> _require_reciprocal_s_scales.
There was a problem hiding this comment.
Fixed. The rest of that comment was wrong in two further ways once I looked: it claimed "the FROST kernels convert P unscaled", but SM120 now applies the scales, and the criterion is non-reciprocal, not non-unit (a reciprocal pair is served exactly, since nothing is applied and nothing is owed back). Rewritten to name both rows.
| "SM120 fp8 serves the per-tensor SDPA_FP8 op only (no MXFP8 cell)", | ||
| ) | ||
| self._not_implemented_error_if(self.thd, "SM120 fp8: THD/varlen is not wired yet (dense only)") | ||
| self._value_error_if(self.has_sink, "SM120 fp8 does not support attention sinks (Amax_S semantics)") |
There was a problem hiding this comment.
Seems like we implement THD and sink related logics in the fp8 kernel, but not enable it?
There was a problem hiding this comment.
Good catch, and the two halves turned out to differ.
Sink is not implemented at all here — there is no sink math in the kernel, only a rejection (sinks is always None, and both the kernel ctor and the adapter refuse has_sink). Amax_S semantics with a sink column are undefined for this cell. Now documented as "no sink math exists here" rather than implying dormant support.
THD was real, and is now wired and tested. Four layers were blocking it, each hiding the next:
- the engine row declared
thd=False("deferred, dense execute only for v1"); sdpa_support_surface.hrejected(prop_major == 12 && is_ragged)outright — removed, and verified safe: with the FROST opt-in off, a ragged fp8 sm120 graph now declines at plan time via the backend's own engine-config check, so that guard was a redundant early-out;- the execute-time
seq_len_qguard rejected any per-batch length< S_q, which under THD is the definition of ragged; - the ragged LSE is head-major
(H, head_stride)— the kernel writeslse[head, q_row_base + row]— and its extent was pinned to the packed total, so a padded token capacity was inexpressible. The fake tensor is 2-D now, like the f16 cell's.
_thd_pack is shared with the f16 THD execute rather than duplicated; the f16 ragged suite is unchanged. New coverage: test_fp8_sm120_thd{,_cross,_stats,_gqa}.
| one `prmt.b32` produce each A register directly from the `cvt.rn.satfinite.e4m3x2` | ||
| results left in registers by the softmax. | ||
|
|
||
| v1 of this kernel staged P through a per-warp `16 × kv_tile` SMEM tile instead |
There was a problem hiding this comment.
Do we want to let the document include the design decisions behind the kernels, or only reflect their current state? It seems the other FE OSS API documents focus on kernel interfaces and high-level design.
There was a problem hiding this comment.
Agreed — dropped the file. It had become kernel design rationale (rejected alternatives, measured costs, tile-rule fits), which is not what the fe-oss-apis docs are for, and this change has no interface surface to document. The rationale worth keeping now lives at the code it explains or in the commit messages.
THD (ragged) now served on the SM120 FP8 enginePushed
On (2) — that guard also governs the native backend, so I did not simply delete it. With it removed and the FROST opt-in off, a ragged fp8 sm120 graph declines at plan time through the backend's own engine-config check ( On (3) — that guard was mine, added earlier in this PR for the dense path, where a short On (4) — the ragged LSE is head-major
Also parameterises the grouped-query test over
note to self: claude::304e9e55-1db7-4285-967f-001cb21032f3 — "cudnn-FE #509 SM120 FP8 engine + Scale_S scope" |
|
@cudnn-ci-bot run frost |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-509-a498899 |
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.
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.
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
a498899 to
c98a352
Compare
Rebased onto develop (#531) — conflicts resolved
Validated at
On the earlier
|
|
@cudnn-ci-bot run frost |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-509-c98a352 |
…nels; decline what TMA cannot express A THD tensor may declare a wider token stride than the packed h*d — e.g. a K/V view of a kv-interleaved [T, 2, H, D] buffer (token stride 2*h*d), the layout torch.nn.attention.varlen users produce by slicing a fused KV projection. The THD lowerings rebuilt packed (1, T, H, D) views with hardcoded strides, so such graphs were claimed and silently mis-addressed (100% of O wrong on both sdpa_fwd_prefill_sm120 and the sm100 flavors; caught by PR NVIDIA#516's fuzz coverage and PyTorch's own varlen suite). Native support, no fallback (AGENTS Hard Rule 2): - compile() on all five f16 kernels (sm120, sm100 d128/d192_d128/d256/ d512) takes optional caller-declared (batch, seq, head, elem) strides per tensor (lru cache-key); None keeps the compact specialization bit-for-bit. Strided fakes via make_fake_tensor, validated against the TMA 16-byte global-stride rule. - SM120: kv_tma_desc reads the tensor's strides instead of recomputing packed ones (Q/O offset math was already layout-driven); the entry validator accepts padded 16-byte-granular BSHD storage (compact = the equality special case). - SM100: the Q/K/V/O TMA descriptors are built from the tensor views, so declared strides flow in unchanged; the THD O-descriptor builder steps per-batch bases by O's declared seq-axis stride (o_tensor.stride[1]). - Adapters bind declared-stride (1, T, H, D) views directly. What TMA cannot express is REJECTED in check_support (NotImplementedError naming the offending strides), so the Router falls back to an engine that honors the declaration: non-innermost-contiguous head dim, or token/head strides that are not multiples of 8 elements (sub- granularity strides also violate the graph API's pointer-alignment contract for the backend, so declining is correct, not conservative). - The SM120 FP8 THD path (NVIDIA#509) keeps the packed contract for now: non-packed declarations are declined (_thd_check_strides_packed); extending native strides there is tracked as a follow-up. Verified (torch nightly cu132, ToT develop + PR NVIDIA#516's fuzz tests): gapped seeded repros pass with the frost engines serving natively on cc 10.0 (sm100) and RTX 5080 (sm120); 128-test fwd ragged L0 sweep slice green on cc 10.0 (all four sm100 flavors) and 84-test slice on sm120; ex-ops suite incl. kv-interleaved views 11/11 on both; dense fwd slice 182 passed (dense compile paths pass no strides -> unchanged); packed THD configs bit-for-bit unchanged. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…6 fwd kernels (#526) * frost(sdpa): native THD stride support in the SM100/SM120 f16 fwd kernels; decline what TMA cannot express A THD tensor may declare a wider token stride than the packed h*d — e.g. a K/V view of a kv-interleaved [T, 2, H, D] buffer (token stride 2*h*d), the layout torch.nn.attention.varlen users produce by slicing a fused KV projection. The THD lowerings rebuilt packed (1, T, H, D) views with hardcoded strides, so such graphs were claimed and silently mis-addressed (100% of O wrong on both sdpa_fwd_prefill_sm120 and the sm100 flavors; caught by PR #516's fuzz coverage and PyTorch's own varlen suite). Native support, no fallback (AGENTS Hard Rule 2): - compile() on all five f16 kernels (sm120, sm100 d128/d192_d128/d256/ d512) takes optional caller-declared (batch, seq, head, elem) strides per tensor (lru cache-key); None keeps the compact specialization bit-for-bit. Strided fakes via make_fake_tensor, validated against the TMA 16-byte global-stride rule. - SM120: kv_tma_desc reads the tensor's strides instead of recomputing packed ones (Q/O offset math was already layout-driven); the entry validator accepts padded 16-byte-granular BSHD storage (compact = the equality special case). - SM100: the Q/K/V/O TMA descriptors are built from the tensor views, so declared strides flow in unchanged; the THD O-descriptor builder steps per-batch bases by O's declared seq-axis stride (o_tensor.stride[1]). - Adapters bind declared-stride (1, T, H, D) views directly. What TMA cannot express is REJECTED in check_support (NotImplementedError naming the offending strides), so the Router falls back to an engine that honors the declaration: non-innermost-contiguous head dim, or token/head strides that are not multiples of 8 elements (sub- granularity strides also violate the graph API's pointer-alignment contract for the backend, so declining is correct, not conservative). - The SM120 FP8 THD path (#509) keeps the packed contract for now: non-packed declarations are declined (_thd_check_strides_packed); extending native strides there is tracked as a follow-up. Verified (torch nightly cu132, ToT develop + PR #516's fuzz tests): gapped seeded repros pass with the frost engines serving natively on cc 10.0 (sm100) and RTX 5080 (sm120); 128-test fwd ragged L0 sweep slice green on cc 10.0 (all four sm100 flavors) and 84-test slice on sm120; ex-ops suite incl. kv-interleaved views 11/11 on both; dense fwd slice 182 passed (dense compile paths pass no strides -> unchanged); packed THD configs bit-for-bit unchanged. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * docs(agents): Hard Rule 2 — serve the declared layout natively or decline, never adapt Closes the loophole Rule 1's letter leaves open: adapter-side normalization copies that make an unsupported layout runnable. Workspace carving does not legitimize a data-tensor copy (the carve exemption is for metadata and dead-slot dummies), the dense path's grandfathered normalization is not a license for new ones, and whatever check_support accepts the kernel must address natively. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * frost(sdpa): review hardening — validate runtime THD buffers, decline overlapping strides - _thd_view validates the runtime buffer against its declaration before reinterpreting storage: dtype/device must match and the base address must be 16-byte aligned (TMA global-address rule / assumed_align=16); as_strided already rejects views past the underlying allocation. - _thd_check_strides_native additionally requires covering (non-overlapping) strides — head >= d, token >= heads*head — matching the SM120 kernel's is_layout_supported, so sub-dense declarations are declined at check_support instead of failing at the per-execute compile (or racing on O writes on SM100). - Kernel _fake_bshd guards: the head dim must be innermost-contiguous; d256/d512 validate the O stride at BPE_O (the O storage dtype byte size). - Clearer SM120 layout-rejection message (the entry validator accepts padded storage now; the text still demanded compact). The THD host-prep stream binding flagged in the same review round is a pre-existing issue (#476) and is split into a separate PR. Addresses CodeRabbit review feedback on #526. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * frost(sdpa): make the THD native-stride gate dtype-aware (16 // itemsize) The gate hardcoded the TMA 16-byte global-stride rule as 8 elements, the f16/bf16 case. It lives in the shared base class, so express the quantum in the tensor's own element units — 8 at 2 B/elem, 16 at 1 B/elem (fp8), 4 at 4 B/elem — per descriptor, so mixed-precision declarations check each tensor at its own dtype. No behavior change for the f16 paths this PR enables; the fp8 native-stride follow-up (#537) inherits the correct quantum for free. Suggested by @Aneureka in review. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
3 participants
Before submitting
pre-commit runand committed any formatting changes. (black 26.3.1 -l 160 on all touched python files)cat-*, one or moremod-*, and oneorig-*(see label list). (token lacks label permission on this repo — suggested:cat-feature,mod-python-fe,mod-kernels,orig-community)Affected area
FE OSS kernels or CuTeDSL
Summary
New opt-in engine
sdpa_fwd_prefill_sm120_fp8(manifest slot 7): a per-tensor FP8 (e4m3) SDPA-forward prefill kernel for consumer Blackwell (SM120/SM121), reachable through the ordinarygraph.sdpa_fp8(...)op withCUDNN_FRONTEND_ENABLE_FROST_ENGINES=1.python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py— the e4m3 sibling ofprefill_f16_sm120.py, with the MMA lowered tomma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32. e4m3 travels as Uint8 storage end to end (TMA → ldmatrix → MMA consume bit patterns). K B-fragments via byte-preservingldmatrix.m8n8.x4.b16; V B-fragments via the hardware 8-bit transposedldmatrix.m16n16.x2.trans.b8(SASSLDSM.8.MT1616). FP16 O; fp32 softmax denominator.2*(t%4)+{0,1}while the k32 A-fragment wants four consecutive bytes; the two differ only by an exchange inside each thread quad, sopack_f8x2_pairs+ twoshfl.sync.idx+ oneprmt.b32build each A operand from thecvt.rn.satfinite.e4m3x2results the softmax leaves in registers. No SMEM round trip.descale_q*descale_kfolds into the softmax scale,descale_v*scale_ointo ano_scale_fusedepilogue scalar;Amax_SandAmax_Ovia bitcast-int32 atomic max into host-pre-zeroed buffers._sm120_fp8_speccapability row (including a newCapabilities.out_dtypes, sincefacts.dtype_opreviously never reachedmismatch()),SdpaFwdDslSm120fp8 compile/execute path, shared PTX helpers intile_dsl.mma.propose_plansnow offers the capability row's wholetile_ms x tile_nsdomain instead of one entry, so a caller can pin a point withcreate_execution_plan(engine_id, SdpaFwdKnobs(...)). Entry 0 carries no knobs and letsconfig_sm120.fp8_tile_choicedecide from shape and SM count.v1 envelope: E4M3 in / FP16 out, exact d128, causal / bottom-right (+SWA, +rect) / SWA / KV-padding masks, GQA/MQA, stats +
lse_optional; no sink, no THD, no E5M2 — all gated in the capability row and negative-tested.Why
SM120 had only the f16/bf16 FROST SDPA cell; the quantized-attention alternatives on this arch are third-party INT8/FP4 kernels with materially worse numerics.
All numbers below: RTX PRO 6000 Blackwell Server Edition (188 SMs, SM120), cuDNN 9.25.0.15, CUPTI kernel time, 6 s device warm-up before the first configuration, interleaved repeats. Three different baselines, which answer three different questions:
Accuracy is unchanged from v1: O max abs err 1e-3..3e-2 against an fp32 reference over dequantized inputs, LSE ~1e-6.
A note on the RTX 5080 table this PR previously carried (1.74–1.79x vs bf16): those were the pre-shfl kernel on a different part, and the ratio does not transfer. The same old kernel measured here gives only 1.27x–1.38x; the shfl path is what takes it to 1.65x–1.89x. The doc's claim that "the kernel is arch-identical so ratios are expected to transfer" was wrong and has been removed.
Measured and rejected
Recorded because they are the obvious next things to try, and each cost a measurement to rule out. The kernel is issue-bound after the shfl change (ncu: Compute 74%, DRAM 3.5%, L1/TEX 49%), which is why reordering work does not help — only removing instructions does.
softmax_op::rescale_thresholdand the SM100 fp8 cell do): 2.7–3.2x worse O error in the regime where 98% of P falls below e4m3's smallest normal. It trades accuracy for speed, and the speed half needs a branch this kernel's structure does not give for free.Measuring any of this needed a noise floor first: run-to-run spread is ~1% at s>=2048 but was 12% at s=512 until the probe warmed the whole device, which is larger than every effect above.
API and compatibility impact
Additive. The engine is opt-in behind
CUDNN_FRONTEND_ENABLE_FROST_ENGINES=1and declines anything outside its envelope, so graphs that do not match fall through to the backend as before.Capabilities.out_dtypesis a new field defaulting to empty, which only the quantized rows read.test_mhas_v2run serially withrandomseeded: identical routing tally and identical failure set againstdevelop.Related issues
descale_sandscale_sbefore the adapter is reached. Not SM120-specific and not fixed here.Summary by CodeRabbit
New Features
Documentation
Bug Fixes