Rank SDPA plans in one place, and give SM120 one measured tile rule - #527
Closed
YangXu1990uiuc wants to merge 10 commits into
Closed
Rank SDPA plans in one place, and give SM120 one measured tile rule#527YangXu1990uiuc wants to merge 10 commits into
YangXu1990uiuc wants to merge 10 commits into
Conversation
New engine sdpa_fwd_prefill_sm120_fp8 (slot 7): the per-tensor FP8 sibling of the f16/bf16 SM120 prefill kernel, lowered to mma.sync.aligned.m16n8k32 e4m3 with the SDPA_FP8 node convention (scalar descales folded host-side, in-kernel Amax_S/Amax_O via bitcast-int32 atomic max). Kernel (prefill_fp8_sm120.py, cloned from prefill_f16_sm120.py): - e4m3 rides as Uint8 storage end to end (TMA/ldmatrix/MMA consume bit patterns; no Float8 element support needed in the DSL plumbing) - K B-fragments via byte-preserving classic ldmatrix.m8n8.x4.b16 (sm_120a has no non-transposed 8-bit ldmatrix; the m8n16 .b8 form is the FP4 nibble-expansion load) - V B-fragments via the hardware 8-bit transposed ldmatrix.m16n16.x2.trans.b8 (SASS LDSM.8.MT1616), register map (0,2,1,3), two MMAs per issue - P restaged through a per-warp SMEM tile (cvt.rn.satfinite.e4m3x2 at the C layout, warp-fenced, reloaded as A fragments) -- the k32 C->A fragment-column mismatch defeats the f16 in-register repack; the softmax denominator stays fp32 - FP16 O; SMEM alias and epilogue sized in bytes (KV 1B vs O 2B) Wiring: manifest slot, _sm120_fp8_spec capability row (E4M3/FP16-out only, exact d128, causal/BR/SWA/KV-padding, no sink/THD), SdpaFwdDslSm120 fp8 compile+execute path mirroring the SM100 _execute_fp8 host math, shared ptx_mma_m16n8k32_e4m3_f32 / ptx_cvt_e4m3x2 helpers in tile_dsl.mma (with an Int32-bitcast armor for a cutlass-dsl 4.7.0 inline_ptx constant-f32 bug). RTX 5080 kernel-time speedup vs the bf16 f16 kernel (CUPTI): 1.74-1.79x non-causal, 1.70x causal; O err vs fp32-dequant reference 3e-4..1e-2, LSE ~1e-6. Details and optimization options in docs/fe-oss-apis/attention/sdpa-fp8-sm120.md. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
propose_plans returned one entry, so the capability row's tile_ms/tile_ns domain was unreachable: SdpaFwdKnobs validation, PlanConfig.knobs and create_execution_plan replay were all in place with nothing enumerating the points. The adapter then took SEQ_*_TILES[0] -- the largest tile that fits, not the fastest. knob_candidates() reads the domain off Capabilities (no second table to drift from it), filters through mismatch(), and orders it via a new per-spec EngineSpec.knob_order. Entry 0 still carries knobs=None: a delegation and a named knob set are different requests, one may improve with the library and the other must be honored verbatim, so existing name-pinning is untouched. config_sm120.fp8_tile_choice picks (q_tile, kv_tile) from shape, causality and SM count. kv_tile=64 unconditionally -- the kernel is L1-bound on the P restage and halving the KV tile halves that traffic, faster in 47 of 48 shapes measured (the exception by 0.15%). q_tile=64 while the grid cannot fill the machine, with a wider window under a causal mask where the last Q tile does several times the work of the first. FP8 only: the f16 cell wants kv_tile=128 at long sequences and the two share SEQ_*_TILES. Measured on RTX PRO 6000 Blackwell (sm120, 188 SMs, cuDNN 9.25), CUPTI kernel time, 2 s clock warm-up, both measurement orders agreeing within 0.5%: - 22 held-out shapes (none used to read the rule off): regret against the best of the enumerated domain is 1.005x mean, 1.045x worst. - versus the backend's native fp8 fprop: ahead on 19 of 22, 1.5x at small grids. The previous default lost 11-25% to its own domain at every shape. test_mhas_v2 serial and seeded: identical routing tally and identical failure set against the committed files, so this changes which tile a graph runs on, not which engine claims it. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The rule is arithmetic over shape, causality and SM count, so it needs no device -- SM count is an argument, not something read off the machine, which is why it is a pure function. Both directions of the causal clause are asserted: checking only the causal side would pass for a rule that ignored causality entirely. propose_plans now offers four tile points but any one shape runs exactly one of them, so the other three would ship unexercised. test_fp8_sm120_every_- enumerated_tile pins each named entry through select_plan and checks it against the same fp32-dequant reference. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The QK C-fragment owns columns 2*(t%4)+{0,1} while the k32 A-fragment wants
four consecutive bytes, so v1 staged P through a per-warp 16 x kv_tile SMEM
tile: 32 STS.U16 + 4 LDSM.x4 + 2 warp barriers per warp per KV tile. The two
layouts differ only by an exchange inside each thread quad, so pack_f8x2_pairs
+ two shfl.sync.idx + one prmt.b32 build each A register straight from the
cvt results the softmax already leaves in registers.
Measured 1.07-1.43x over 24 shape x tile combinations on an RTX PRO 6000
Blackwell (188 SMs, cuDNN 9.25), largest where the restage traffic was
largest. The design notes predicted 1.04-1.08x from an instruction count; ncu
says the kernel was L1-bound (72-77% L1/TEX against the backend kernel's
45-51%, DRAM 5-13% on both), so the binding resource was the L1 port and
removing 36 memory ops per warp per tile was worth several times the
issue-count estimate.
It also frees 16 KB of SMEM (49 -> 33 KB at 128x128), which moves the tile
optimum: kv_tile is now 128 in all 28 shapes measured, having been 64 in 47 of
48 on the old kernel. fp8_tile_choice is re-derived accordingly, and both it
and its tests now say that a tile rule is a property of the kernel it was
measured on. Held-out regret against the best of the enumerated domain is
1.009x mean / 1.089x worst over 22 shapes; against the backend's own fp8
fprop the engine is ahead on all 22, by 1.15x-1.93x.
Three further optimizations from the sibling DKG kernel were measured and
NOT taken. Prefetching K or V fragments one step ahead lands inside +-0.5%,
and loading Q coalesced with a shfl transpose is reproducibly 0.5-2.3% slower
at s=512 -- the shape it should help most, since Q is small enough to sit in
L2. The run-to-run floor is ~1% at s>=2048 and was 12% at s=512 until the
probe warmed the device first, which is what made these calls possible.
Also from review:
- Capabilities gains out_dtypes. facts.dtype_o never reached mismatch(), so a
graph asking for fp8 O was claimed and then declined at build. The default
is empty rather than unconstrained: a quantized row that forgets to declare
serves nothing, which fails loudly.
- The post-kernel amax_o.div_ and O copy-back move onto the launch stream,
next to the resets that were already there (NVIDIA#510).
- The negative tests the module docstring promised now exist (fp8 O, D != 128).
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…ells
The f16/bf16 cell shared the fp8 cell's knob domain but had no rule of its own,
so it took whatever fit: 128x128 everywhere on this part. Measuring it produced
the same rule the fp8 cell uses, plus one term.
A causal mask halves the work per CTA, so the machine empties sooner and the
finer Q tile keeps paying further out. Folded in as a halved effective grid.
Regret against the best of the enumerated {64,128}^2 domain:
without with
RTX PRO 6000, 24 bf16 cells 1.023 / 1.175 1.009 / 1.054
RTX PRO 6000, 28 fp8 cells 1.0078/ 1.107 1.0066/ 1.107 (2 fixed, 2 broken)
RTX 5090 (170 SM), 14 fp8 1.0082/ 1.062 1.0007/ 1.006
22 held-out fp8 shapes 1.009 / 1.089 flips the 3 worst cells to optimal
(mean / worst.) On the part the rule was trained on the term is a wash and the
worst case is unchanged; it earns its place on the second part and the held-out
set, where it fixes exactly the causal cells that were the previous misses.
So fp8_tile_choice and f16_tile_choice collapse into one tile_choice, and the
two knob_order callbacks into one. That also drops a leftover: the fp8 callback
ranked its runner-up entries by tile_n != 64, from before the shfl change moved
the KV optimum to 128.
The tile-rule test file covers both cells now, so it loses the fp8 in its name.
propose_plans called check_support(graph) and then knob_candidates(spec, graph), and each of those ran analyze_for -> mismatch(caps, facts, None). The facts parse itself is memoized per graph (_attach_facts -> _facts_for), so the waste was the capability match, once per engine per plan walk. knob_candidates now takes the facts the caller already has. _facts_or_decline is the single place the engine asks, returning the facts it just proved eligible; check_support is that call with the answer dropped. No behaviour change: knob_candidates could only be reached through a passing check_support, so its own decline branch was unreachable.
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 sides anyway, which heuristics_sort did by
concatenating and calling it ranking.
Now create_execution_plans() gathers the inputs (parsed facts, this 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. sdpa/fwd/heuristics.py is the first such hook: mode A returns
candidates worth timing with the tile rule's guess first, FALLBACK returns
configs expected to build where A's choice may not, and both sides' fallbacks
sort behind both sides' tuned entries.
Which side leads is now a measurement rather than a default. sm120 fp8 is timed
1.20-1.83x the backend's native fp8 fprop, so it leads; a cell that has not been
timed against the backend keeps the historical order and says so. Moving one is
an experiment, not an edit.
Deleted, all superseded or never used:
BaseEngine.propose_plans + its 3 implementations (2 were verbatim copies of
the default, the third was added last week -- nothing ever used the hook to
decide anything)
BaseEngine.default_knobs only fed propose_plans
EngineSpec.knob_order ranking, in the wrong place
knob_candidates ditto
engines.probe() (fwd + bwd) zero callers, superseded by check_support
heuristics_sort merging is part of ranking, not a step after
Router.python_plans / Router class / default_router / set_router /
pygraph(router=) policy has one home now; the router is plumbing
graph.engine pure alias of selected_engine, zero callers
graph.from_serialized zero callers; serialize/deserialize are the
pybind-era API and stay
api_dsl's shape-driven tile choice the tile rule ran twice, once here and
once when ranking; resource feasibility stays
knobs=None no longer means "engine, pick for me" -- heuristics always names a
concrete config. A None field survives only on an axis whose capability row
declares no domain.
Tests are updated in the next commit; test_engine_router still builds Router
subclasses to inject orderings, which is now done by replacing heuristics.rank.
… 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 (326 calls
across 20 files).
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 now a
measurement per cell, not a rule
FROST coverage no longer rides on that ordering. heur_mode.OPENSOURCE asks for
an open-source implementation, which is what these cells are, so it puts
everything FROST can build ahead of the backend whatever the measurements say.
A test asks for it and whatever still lands on a backend plan is a graph FROST
does not cover -- coverage measured rather than asserted, and the plan walk
falls through to the backend so the test still checks numerics either way.
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].
Known: test_a_replayed_plan_reports_its_own_notes fails, and it already failed
at d7dd85e (verified by running it against that tree) -- C++ no longer raises
for an index one past the plan count on 9.25. Pre-existing, not this change.
Out-of-tree engines lose the ability to offer several knob candidates: they
declare no Capabilities, so nothing can enumerate their domain. register_backend
remains what the tests use to install a fake engine.
…n lookup OPENSOURCE is mode A without the backend's recommendation: these cells ARE the open-source implementation and the backend's engines are not, so there is nothing of the backend's to place. Combine it to measure coverage -- [OPENSOURCE, A, FALLBACK] tries every FROST config first and still has the backend behind it, so a graph that runs on a backend plan is one FROST does not cover. recommend() is now a loop over the caller's modes rather than three blocks spliced together. Each mode contributes its own, so [A, FALLBACK] puts both sides' tuned candidates ahead of both sides' fallbacks by construction. An untagged backend entry is the delegating one C++ tries before its own engine_configs; it belongs to no mode and keeps the lead, or an OPENSOURCE caller would get a native kernel instead of the OSS one. Plan names now always read <engine>[<knobs>], since the heuristics name a concrete config for every entry -- there is no knobs=None entry left to print as a bare engine name. Six test files each carried their own verbatim copy of _select_engine matching that bare name exactly; they now share frost_test_utils.select_engine, which matches on the engine. probe() is restored to its callers rather than deleted: I claimed it had none, which was wrong -- test_sdpa_graph_analyzer used it twice. Those two call analyze_for directly now, so no production API exists only for tests. sm120: sdpa suites 62 failed -> 4 (171 passed); test_mhas_v2 302 passed with the routing tally unchanged (frost sm120 x87, fp8 x4, native x158/x53). The 4 are head_dim 208-256, where the adapter used to drop to a KV tile that fits and the heuristics now name 128 -- traceback pending, not yet diagnosed.
Contributor
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Comment |
Collaborator
Author
|
For a diff without #509's commits (GitHub will not take a fork branch as a base, so this PR targets |
Closed
3 tasks
Collaborator
Author
|
Superseded by #528, which is the same framework change based on |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Before submitting
pre-commit runand committed any formatting changes. (black 26.3.1 -l 160 on all touched python files)cat-*, one or moremod-*, and oneorig-*. (token lacks label permission — suggested:cat-feature,mod-python-fe,orig-nv-eng)Affected area
FE OSS kernels or CuTeDSL — python dispatch (
cudnn.pygraphplan ranking) and the SM120 SDPA-forward tile rule.Draft, stacked on #509. GitHub will not take a fork branch as a base, so this targets
developand the diff carries #509's six commits too. Review only the six commits fromce8b3e10eonward — everything at or belowd7dd85ee5belongs to #509 and disappears from this diff once it merges.Summary
An engine cannot rank plans. It sees neither its siblings nor the backend's entries, so
propose_planscould only order its own knobs — and something downstream still had to merge the two sides, whichheuristics_sortdid by concatenating and calling it ranking. This moves the whole decision to one function per engine family and deletes what that makes redundant.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 backend's plans now 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. That 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.
One SM120 tile rule instead of two, with a causal term
The f16/bf16 cell shared the fp8 cell's knob domain but had no rule, so it took whatever fit — 128x128 everywhere on this part. Measuring it produced the same rule plus one term: a causal mask halves the work per CTA, so the machine empties sooner and the finer Q tile keeps paying further out. Regret against the best of the enumerated
{64,128}²domain:On the part the rule was trained on the term is a wash; it earns its place on the second part and the held-out set, where it fixes exactly the causal cells that were the previous misses. So
fp8_tile_choiceandf16_tile_choicecollapse into onetile_choice.Which side leads is a measurement
_MEASURED_BEHINDnames cells timed slower than the backend; those follow it. sm120 fp8 is not there because it measures 1.20–1.83x the backend's native fp8 fprop across 28 shapes. A cell that has not been timed keeps the historical order and says so in a comment — moving one is an experiment, not an edit.Coverage no longer rides on that ordering:
heur_mode.OPENSOURCEputs everything FROST can build first whatever the measurements say, so a test asks for it and any graph still landing on a backend plan is one FROST cannot serve.Deleted
BaseEngine.propose_plans+ its 3 implementationsBaseEngine.default_knobspropose_plansEngineSpec.knob_order,knob_candidatesheuristics_sortRouterclass,default_router,set_router,pygraph(router=),engines/router.pydecline_typesmoved tobase.pywhere the contract isgraph.engineselected_engine, zero callersgraph.from_serializedserialize/deserializeare the pybind-era API and stayapi_dsl's shape-driven tile choiceknobs=Noneno longer means "engine, pick for me". A None field survives only on an axis whose capability row declares no domain.Testing
test_mhas_v2routingKnown, not yet fixed:
test_dsl_sm120_representative_head_dimensions[208..256]. At those head dims the adapter used to drop to a KV tile that fits; the heuristics now name 128 and the test pins the first entry. Traceback pending.Pre-existing, not from this change:
test_a_replayed_plan_reports_its_own_notesfails here and also fails at the base commitd7dd85ee5— C++ on 9.25 no longer raises for an index one past the plan count.Still to come on this branch
register_backend/pygraph(backends=)/OUT_OF_TREE_ID_BASEgo,_owners_for_iddecodes ids straight fromMANIFEST, and test fakes inject a fake family instead_MEASURED_BEHIND(today it is empty, i.e. "unmeasured keeps the old order")note to self: claude::304e9e55-1db7-4285-967f-001cb21032f3 — "审计 Frost Python DSL 引擎调用流程"
cwd /home/scratch.yanxu_libs/cudnn_frontend · workspace /home/scratch.yanxu_libs/fe_sm120