diff --git a/docs/fe-oss-apis/attention/sdpa-fp8-sm120.md b/docs/fe-oss-apis/attention/sdpa-fp8-sm120.md new file mode 100644 index 000000000..43bfa7a9b --- /dev/null +++ b/docs/fe-oss-apis/attention/sdpa-fp8-sm120.md @@ -0,0 +1,165 @@ +# SM120 per-tensor FP8 SDPA forward: performance, options, and tradeoffs + +Engine `sdpa_fwd_prefill_sm120_fp8` (kernel +`python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py`) is the per-tensor FP8 +(e4m3) sibling of the f16/bf16 SM120 prefill kernel, reachable through the +ordinary `graph.sdpa_fp8(...)` op with +`CUDNN_FRONTEND_ENABLE_FROST_ENGINES=1`. This note records the measured +baseline, every design decision that trades performance against something +else, and the optimization options on the table — so the next change here +starts from data instead of archaeology. + +## Measured baseline + +RTX 5080 (84 SMs, SM120), CUDA 13.3, cutlass-dsl 4.7.0, bf16 f16-kernel as the +comparator, kernel time via CUPTI (min of 3 trials × 20 iters; host excluded): + +| shape (d=128) | f16 kernel | fp8 kernel | speedup | +|---|---|---|---| +| b2 h24 s4096 non-causal | 3699 µs | 2129 µs | **1.74x** | +| b2 h24 s8192 non-causal | 14832 µs | 8266 µs | **1.79x** | +| b1 h32 s4096 causal | 1364 µs | 804 µs | **1.70x** | + +RTX PRO 6000 Blackwell Server Edition (188 SMs, SM120, cuDNN 9.25.0.15), +against the **backend's own fp8 fprop** — the path these graphs take today — +at the tiles the shape rule picks, over 22 shapes it was not fitted on: +**1.15x–1.93x, ahead on all of them**, widest at small grids. Against the +sibling DKG SM120 fp8 kernel at matched tiles, best-vs-best: **1.05x–1.08x +behind** (it was 1.21x–1.34x before P moved out of SMEM). + +Two caveats on reading any of these. The two paths overlap less than "fp8 +sibling" suggests: this engine emits FP16 O and the backend's fp8 fprop +declines FP16 O, so most fp8 graphs cannot move either way. And the 1.74–1.79x +above is against the **bf16 kernel** — it answers "what does fp8 buy", not +"is this faster than what already runs". + +Numerics (vs fp32 reference on dequantized inputs, scale-folded descales): +O max abs err 3e-4..3e-3 at std `1/sqrt(d)` inputs, ~1e-2 at full-scale (std +1.0) inputs — dominated by the in-kernel P→e4m3 quantization. LSE agrees to +~1e-6 (the softmax path is fp32 end to end). Amax_S matches the fp32 +reference to float precision; Amax_O to P-quantization tolerance. + +The mainloop is `QMMA.16832.F32.E4M3.E4M3` (half the f16 kernel's HMMA count +— the k-depth doubled), `MUFU.EX2` for the unchanged fp32 softmax, +`LDSM.8.MT1616` for the hardware 8-bit transposed V loads, and `LDSM.16.M88.4` +for K. P never reaches SMEM (see below). The amax atomics compile to one warp +`REDUX.MAX` + one global `RED` per warp — negligible. + +## Design decisions and their tradeoffs + +**e4m3 as Uint8 storage.** The kernel never does elementwise math on Q/K/V, so +bytes flow TMA → ldmatrix → MMA as bit patterns and no Float8 element support +is needed anywhere in the DSL plumbing (Array, TensorMap, pointer paths). +Tradeoff: the ABI is a `uint8` view (`torch` fp8 tensors pass as +`.view(torch.uint8)`), and dtype identity lives in the adapter, not the type +system. + +**K loads via classic `ldmatrix.m8n8.x4.b16`.** sm_120a has no non-transposed +8-bit ldmatrix (the `m8n16 .b8` form in the ISA is the FP4 nibble-expansion +load — using it on e4m3 corrupts data silently). The b16 form is +byte-preserving and gathers the m16n8k32 B fragment exactly when each lane +points at one 16-byte K-segment. No tradeoff vs a hypothetical native +instruction; one extra issue per k32 step vs f16's per-k16 amortization. + +**V loads via `ldmatrix.m16n16.x2.trans.b8`** (SASS `LDSM.8.MT1616`): the +hardware 8-bit transpose added for exactly this MMA family. One issue covers a +32(kv)×16(d-bytes) tile and feeds two MMAs. Without it, V would need a +transposed smem layout (second TMA descriptor bank conflicts) or in-register +`prmt` transposes — this instruction is why the port is clean. + +**P reaches the PV MMA through shfl, not SMEM.** The QK C-fragment owns columns +`2*(t%4)+{0,1}` while the k32 A-fragment wants 4 consecutive bytes, so the f16 +kernel's in-register repack does not apply. The two layouts differ only by an +exchange inside each thread quad, so `pack_f8x2_pairs` + two `shfl.sync.idx` + +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 +(32 `STS.U16` + 4 `LDSM.x4` + 2 warp barriers per warp per KV tile). Replacing +it measured **1.07–1.43x** across 24 shape × tile combinations on an RTX PRO +6000 Blackwell, largest at `kv_tile=128` where the restage traffic was largest. +Two things about that are worth recording, because the v1 notes predicted +neither: + +- The kernel was **L1-bound, not issue-bound**. ncu put L1/TEX throughput at + 72–77% against the backend fp8 kernel's 45–51%, with DRAM at 5–13% on both. + An estimate that counted instructions put the shfl route at 1.04–1.08x; the + binding resource was the L1 port, and the measured gain was several times + that. +- It freed 16 KB of SMEM (49 → 33 KB at 128×128), which **moved the tile + optimum** — see below. + +*Rejected: prefetching K or V fragments one step ahead, and loading Q coalesced +with a shfl transpose.* All three are in the sibling DKG kernel. Measured here +against the shfl-P baseline, K and V prefetch land inside ±0.5% (the run-to-run +floor at s≥2048 is ~1%), and the coalesced Q load is **reproducibly 0.5–2.3% +slower at s=512**, the shape it was supposed to help most — Q is small enough +to sit in L2, so the eight shuffles per d-fragment cost more than the saved +transactions. Absent the SMEM restage there is little ldmatrix latency left to +hide. + +## Tile options + +The capability row advertises `tile_ms/tile_ns ∈ {64,128}` and +`propose_plans` offers every point, so a caller can pin one with +`create_execution_plan(engine_id, SdpaFwdKnobs(tile_m=..., tile_n=...))`. +Entry 0 carries no knobs and lets `config_sm120.tile_choice` decide. + +`kv_tile=128` in all 28 shapes measured. **This reversed when P moved out of +SMEM**: the restage tile was `(q_tile/16) x 16 x kv_tile`, so its traffic grew +with the KV tile and made 64 the better choice, by 11–25% at every size. The +f16 measurements this file used to extrapolate from said 128 was optimal only +at s>=4k; neither that nor the fp8 v1 rule survives a change to how P is +carried. A tile rule is a property of the kernel it was measured on. + +`q_tile=64` while the grid cannot fill the machine *and* the sequence is long +enough to amortize the extra Q-tile loop — `grid*2 <= SMs`, or +`grid*2 <= 3*SMs` with at least 12 KV tiles. A causal mask halves the work per +CTA, so the machine empties sooner and the finer Q tile keeps paying further +out; that enters as a halved effective grid. + +The 1.5x-SM bound was moved in after a held-out shape at 320 CTAs missed by +1.19x while 240 CTAs was correct, so that pair is no longer independent +evidence — `test_the_grid_bound_sits_between_240_and_320_ctas` pins both +points, and a genuine re-validation needs fresh shapes. + +### One rule, both cells + +The f16/bf16 cell used the same knob domain but took whatever fit, which on +this part meant 128x128 everywhere. Measuring it produced the same rule, so +there is one `tile_choice` rather than one per cell. + +The two sweeps disagree on how much each tile wins by — `kv_tile=128` leads by +2–4% in bf16 against a uniform margin in fp8, since bf16 KV is two bytes and +there is no P quantization — but never on which tile wins. The causal term was +found on the f16 sweep and then measured on the fp8 cell, where it also helps: + +| dataset | rule without the causal term | with it | +|---|---|---| +| RTX PRO 6000, 24 bf16 cells | 1.023x mean regret, 1.175x worst | **1.009x, 1.054x** | +| RTX PRO 6000, 28 fp8 cells | 1.0078x, 1.107x | 1.0066x, 1.107x (2 fixed, 2 broken) | +| RTX 5090 (170 SMs), 14 fp8 cells | 1.0082x, 1.062x | **1.0007x, 1.006x** | +| 22 held-out fp8 shapes | 1.009x, 1.089x | flips the 3 worst cells to optimal | + +Regret is against the best of the enumerated `{64,128}²` domain, so 1.00x means +the rule picked the tile an exhaustive sweep would have. On the training part +the term is a wash in the mean and leaves the worst case untouched; the case +for it is that it is a clear win on the second part and on the held-out set, +where it fixes exactly the causal cells that were the previous rule's misses. + +Evidence scope: the fp8 side is validated on 22 held-out shapes and on a second +part. **The bf16 side has been measured on one part only**, and contributes no +held-out set. + +## Bigger levers beyond this kernel + +- **Multi-stage KV**: 1-byte KV halves smem per stage, and the 16 KB the P + restage used to hold is now free, so a 2-stage sK/sV pipeline fits at + 128x128 (33 KB single-buffered today). This hides TMA latency behind + compute, which is a different thing from the *register*-level prefetch of K + and V fragments measured and rejected above. Largest remaining item. +- **NVFP4 / mixed-precision PV**: `MmaMXF8F6F4Op` supports mixed + (e4m3, e2m1) operand pairs on SM120, and `LDSM.U4` (incl. the transposed + form) hardware-unpacks fp4 from smem — a P=e4m3 × V=e2m1 PV MMA halves V + bandwidth without quantizing P below 8 bits. That is a new kernel (SF + channel, sub-byte addressing), not an upgrade of this one. diff --git a/python/cudnn/_pygraph.py b/python/cudnn/_pygraph.py index d41c4c3d4..eb95ea417 100644 --- a/python/cudnn/_pygraph.py +++ b/python/cudnn/_pygraph.py @@ -5,11 +5,11 @@ All graph structure and attributes are kept in Python. Graph construction is backend-agnostic; a backend is chosen at create_execution_plans() time by the -Router, and the backend-specific representation (e.g. the C++ cuDNN graph) is +heuristics, and the backend-specific representation (e.g. the C++ cuDNN graph) is generated lazily only then. Execution flow (unification proposal): - build ops -> create_execution_plans() -> Router -> selected backend + build ops -> create_execution_plans() -> heuristics -> selected backend (a registered native engine, or the cuDNN Graph backend by lazy lowering) Example with a native backend (pass torch tensors directly): @@ -94,7 +94,6 @@ def __init__( *, # ---- new (keyword-only: never shifts the classic positional order) -- backends: Optional[List["BaseEngine"]] = None, - router: Any = None, **kwargs, ): self._context = GraphContext( @@ -129,13 +128,12 @@ def __init__( self._is_built: bool = False self._data_bindings: Dict[int, Any] = {} # uid -> tensor data for auto-bound inputs - # Backend routing (see engines/router.py). Graph construction is - # backend-agnostic. At create_execution_plans() the Router builds a flat - # ranked plan list (python engines + cuDNN) in one shared engine-id + # Backend routing (see engines/heuristics.py). Graph construction is + # backend-agnostic. At create_execution_plans() the heuristics build a + # flat ranked plan list (python engines + cuDNN) in one shared engine-id # space; each plan is dispatched by its id (is_python_engine -> python # registry, else lower to cuDNN). ``_plan_index`` selects the plan to run. self._backends: List["BaseEngine"] = [] - self._router = router # None => engines.router.default_router at route time self._plans: List[Any] = [] # list[PlanConfig], populated by create_execution_plans() self._planning_done: bool = False # create_execution_plans() ran (one-shot) self._frozen: bool = False # whole-surface freeze (set by _freeze()) @@ -151,6 +149,7 @@ def __init__( self._facts: Dict[Any, Any] = {} # analyzer callable -> its record; see _facts_for() self._backend_declined: Optional[Exception] = None # why the backend has no entries self._backend_entries: Optional[List[Any]] = None # backend_plan_entries(), once + self._backend_mode_spans: List[Any] = [] # (mode, lo, hi) over the C++ plan list self._barred_names: set = set() # deselect_engines() self._workspace_limit: Optional[int] = None # deselect_workspace_greater_than() self._note_filters: List[Any] = [] # (kind, note, keep) from the classic note filters @@ -206,15 +205,6 @@ def register_backend(self, engine: "BaseEngine") -> "pygraph": self._backends.append(engine) return self - def set_router(self, router: Any) -> "pygraph": - """Override the plan-list / ranking policy for this graph. Must be set - before create_execution_plans() (a later router cannot affect the - already-planned list).""" - if self._planning_done: - raise RuntimeError("cannot set a router after create_execution_plans(); planning is one-shot — build a new graph") - self._router = router - return self - @property def backends(self) -> List["BaseEngine"]: """Engines added with ``register_backend()`` (the out-of-tree hatch). @@ -252,7 +242,7 @@ def _engine_for(self, cfg) -> Optional["BaseEngine"]: def _owners_for_id(self, engine_id: int) -> List["BaseEngine"]: """Candidate engines whose DECLARED range contains ``engine_id``. - The single owner lookup: dispatch, the Router's output validation and + The single owner lookup: dispatch, the ranking's output validation and replay all go through it, so "who runs this id" has one answer computed one way. Never a subclass predicate — registration can prove intervals disjoint, and cannot prove anything about an arbitrary ``owns_id``.""" @@ -840,7 +830,7 @@ def validate(self) -> None: def build_operation_graph(self) -> None: """Validate the graph; lower to C++ when no python engines are registered. - Backend selection is deferred to create_execution_plans() (the Router + Backend selection is deferred to create_execution_plans() (the heuristics stage). With python engines registered, nothing is lowered here (a graph routed to a python engine never touches C++). Without them — the classic sequencing — lowering happens now, so plan-configuration and query @@ -892,8 +882,8 @@ def create_execution_plans(self, heuristics: Optional[List] = None) -> None: graph (discovered from ``engines.manifest`` — no registration call, no environment variable) and the backend's own ranked recommendation (``backend_plan_entries()``, [] when the backend declined the graph or - is not installed). ``engines.heuristics.heuristics_sort`` decides the - order; ``build_plans()`` walks it. + is not installed). ``engines.heuristics.rank`` decides the order; + ``build_plans()`` walks it. Args: heuristics: cuDNN heuristic modes for the backend's recommendation. @@ -901,7 +891,7 @@ def create_execution_plans(self, heuristics: Optional[List] = None) -> None: if not self._is_validated: self.validate() - from .engines.router import default_router + from .engines.heuristics import rank # One-shot planning (classic conformance: the C++ graph never supported # re-planning — a second call there appends plans by accident, and no @@ -920,23 +910,22 @@ def create_execution_plans(self, heuristics: Optional[List] = None) -> None: # lands (a no-op for a graph that was never lowered), and it writes # through object.__setattr__ precisely to bypass the freeze — so the # snapshot has to come after it, not merely after validate(). Doing it - # here rather than inside Router keeps the epoch boundary out of + # here rather than inside the heuristics keeps the epoch boundary out of # overridable policy: every engine probe and the ranking then read one # set of facts describing a graph that can no longer change. self._finalize_backend_layout() self._freeze() self._attach_facts() - router = self._router or default_router - plans = list(router.plan(self, self._candidate_engines())) - # Validate the FINAL router output: every entry must name an engine this - # graph can actually dispatch to. + plans = list(rank(self, self._candidate_engines(), self.backend_plan_entries(), self._backend_heuristics)) + # Validate the FINAL ranking: every entry must name an engine this graph + # can actually dispatch to. from .engines.engine_ids import is_python_engine known = {e.engine_id for e in self._candidate_engines()} for cfg in plans: if is_python_engine(cfg.engine_id) and not self._owners_for_id(cfg.engine_id): - raise ValueError(f"router produced a plan for unknown python engine_id {cfg.engine_id} (known: {sorted(known)})") + raise ValueError(f"heuristics produced a plan for unknown python engine_id {cfg.engine_id} (known: {sorted(known)})") if not plans: # Say WHY, or the user is left guessing which side had nothing: the # backend's own rejection is the usual answer. @@ -951,7 +940,7 @@ def _candidate_engines(self) -> List["BaseEngine"]: ``register_backend()`` added, then the in-tree manifest families whose coarse key matches. Registered engines lead because bringing your own engine is an explicit act; the library's own table is the default. - (Candidate ORDER is not rank — ``heuristics_sort`` ranks.) Cached: the + (Candidate ORDER is not rank — ``heuristics.rank`` ranks.) Cached: the graph is frozen at planning time anyway.""" if self._candidates is None: from .engines import manifest @@ -975,7 +964,7 @@ def _finalize_backend_layout(self) -> None: A failure here is the backend DECLINING (a python engine may still serve the graph), recorded as backend_plan_entries() records one. Not routed through that, which also runs the ~178 ms C++ plan query a - Router may never ask for. + the heuristics may never ask for. """ import cudnn @@ -990,6 +979,7 @@ def _finalize_backend_layout(self) -> None: self._cpp_tensors.clear() self._cpp_bog_done = False self._cpp_plans_created = False + self._backend_mode_spans.clear() self._backend_entries = [] def _attach_facts(self) -> None: @@ -1067,7 +1057,7 @@ def backend_plan_entries(self) -> List[Any]: Answered ONCE per graph: a second C++ create_execution_plans() appends to the same plan list (``enqueue_engine_configs`` -> ``back_inserter``), - so re-querying would report every backend plan twice. A Router may + so re-querying would report every backend plan twice. The heuristics may therefore call it freely to place the entries where it wants. """ import cudnn @@ -1108,6 +1098,7 @@ def backend_plan_entries(self) -> List[Any]: self._cpp_tensors.clear() self._cpp_bog_done = False self._cpp_plans_created = False + self._backend_mode_spans.clear() self._backend_entries = [] return self._backend_entries try: @@ -1129,7 +1120,7 @@ def backend_plan_entries(self) -> List[Any]: entries = [] for i in range(self._lowered_graph.get_execution_plan_count()): engine_id, knobs = self._lowered_graph.get_engine_and_knobs_at_index(i) - entries.append(PlanConfig(engine_id, knobs, cpp_index=i)) + entries.append(PlanConfig(engine_id, knobs, cpp_index=i, mode=self._mode_of_backend_plan(i))) import cudnn asked_oss = any(h == cudnn.heur_mode.OPENSOURCE for h in (self._backend_heuristics or [])) @@ -1319,14 +1310,47 @@ def _lower_backend_graph(self) -> None: self._sync_ir_shapes_from_backend() def _create_backend_plans(self) -> None: - """Run the backend's heuristics for the lowered graph (once).""" + """Run the backend's heuristics for the lowered graph (once), ONE MODE AT A TIME. + + A call per mode rather than one call listing them all. C++ appends each + query to the same plan list, so asking separately and reading + get_execution_plan_count() after each is what says which entries came + from which mode — and ranking needs that, since "the backend's mode-A + entries ahead of ours, its fallbacks behind" is not expressible against + one opaque list. + + A mode with no configs raises; that is not a decline, since another mode + may still have entries (an OPENSOURCE-only query legitimately leaves the + cuDNN modes empty). Only every mode failing means the backend has + nothing, and the last error is re-raised so the caller reports why. + """ import cudnn - if not self._cpp_plans_created: - heur = self._backend_heuristics or [cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK] - self._lowered_graph.create_execution_plans(heur) - self._cpp_plans_created = True - self._forward_note_filters() # deferred from _filter_notes(): the plans exist now + if self._cpp_plans_created: + return + modes = self._backend_heuristics or [cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK] + at, failure = self._lowered_graph.get_execution_plan_count(), None + for mode in modes: + try: + self._lowered_graph.create_execution_plans([mode]) + except cudnn.cudnnGraphNotSupportedError as exc: + failure = exc + continue + now = self._lowered_graph.get_execution_plan_count() + if now > at: + self._backend_mode_spans.append((mode, at, now)) + at = now + if not self._backend_mode_spans and failure is not None: + raise failure + self._cpp_plans_created = True + self._forward_note_filters() # deferred from _filter_notes(): the plans exist now + + def _mode_of_backend_plan(self, cpp_index: int): + """The heuristic mode whose query produced the backend plan at ``cpp_index``.""" + for mode, lo, hi in self._backend_mode_spans: + if lo <= cpp_index < hi: + return mode + return None def _lower_backend_plan(self) -> None: """Lower to C++ (if not already) and create the backend plans (once).""" @@ -1343,7 +1367,7 @@ def _check_plan_index(self, index: int) -> int: def _materialize_backend_plan(self, index: int): """Ensure the backend entry at ``index`` has a real C++ plan index. - A replayed entry (``cpp_index=None``, from a custom Router or an + A replayed entry (``cpp_index=None``, from custom heuristics or an autotune result) does not exist in the backend's list until it is appended. Every entry point that needs a concrete plan — check_support, the build walk, workspace, behaviour notes, execute — goes through here, @@ -1389,7 +1413,7 @@ def check_support(self) -> None: what makes "backend first, python next" actually reachable; if nothing builds, the walk raises with every failure listed. """ - from .engines.router import decline_types + from .engines.base import decline_types eng = self.selected_engine if eng is not None: @@ -1425,7 +1449,7 @@ def build_plans(self, *args, ctx: Any = None, **kwargs) -> None: """ import cudnn - from .engines.router import decline_types + from .engines.base import decline_types if not self._planning_done: self.create_execution_plans() @@ -1818,12 +1842,6 @@ def __repr__(self) -> str: return json.dumps(self.inspect(), default=str, indent=2) - @property - def engine(self) -> Optional["BaseEngine"]: - """The python engine for the selected plan, or None for the backend path. - Populated after create_execution_plans().""" - return self.selected_engine - def serialize(self): """Serialize the graph (classic passthrough). @@ -1854,37 +1872,6 @@ def deserialize(self, *args, **kwargs) -> None: self._lowered_graph.deserialize(*args, **kwargs) self._is_built = True - @classmethod - def from_serialized(cls, data, handle: Optional[int] = None, **kwargs) -> "pygraph": - """Create a pygraph from serialized data. - - This is a convenience method that creates a minimal graph and deserializes into it. - - Args: - data: Serialized graph data (from serialize()). - handle: Optional cuDNN handle for AoT compilation. - **kwargs: Additional arguments passed to the constructor. - - Returns: - pygraph: Deserialized graph ready for execution. - """ - import cudnn - - # Create a new graph with a fresh C++ graph - graph = cls(**kwargs) - graph._lowered_graph = cudnn._pybind_module.backend_graph( - io_data_type=graph._context.io_data_type, - intermediate_data_type=graph._context.intermediate_data_type, - compute_data_type=graph._context.compute_data_type, - ) - - if handle is not None: - graph._lowered_graph.deserialize(handle, data) - else: - graph._lowered_graph.deserialize(data) - graph._is_built = True - return graph - def _lower_to_cpp(self) -> Any: """Lower Python graph to C++ (the internal ``_pybind_module.backend_graph``).""" import cudnn diff --git a/python/cudnn/engines/__init__.py b/python/cudnn/engines/__init__.py index 2d1449bfd..700ca1b39 100644 --- a/python/cudnn/engines/__init__.py +++ b/python/cudnn/engines/__init__.py @@ -4,7 +4,7 @@ """Execution engines for pygraph. Pluggable engines in one flat engine-id space with the cuDNN backend. At -``create_execution_plans()`` the Router ranks the engines that claim the graph +``create_execution_plans()`` the heuristics rank the engines that claim the graph against the backend's own recommendation into ONE list (``graph.plans``); ``build_plans()`` walks it. Graph construction stays engine-agnostic. @@ -14,7 +14,7 @@ the out-of-tree escape hatch. """ -from .base import BaseEngine, CompiledPlan, ExecutionContext, PlanConfig +from .base import BaseEngine, decline_types, CompiledPlan, ExecutionContext, PlanConfig from .engine_ids import ( BACKEND_ENGINE_ID_BASE, BACKEND_HEURISTIC_ENGINE_ID, @@ -24,9 +24,7 @@ is_backend_engine, is_python_engine, ) -from .heuristics import heuristics_sort from .manifest import MANIFEST, EngineFamily -from .router import Router, decline_types, default_router def __getattr__(name: str): @@ -49,10 +47,7 @@ def __getattr__(name: str): "CompiledPlan", "ExecutionContext", "PlanConfig", - "Router", - "default_router", "decline_types", - "heuristics_sort", "MANIFEST", "EngineFamily", "BACKEND_ENGINE_ID_BASE", diff --git a/python/cudnn/engines/base.py b/python/cudnn/engines/base.py index 73aa5545e..9a248d3e7 100644 --- a/python/cudnn/engines/base.py +++ b/python/cudnn/engines/base.py @@ -1,15 +1,18 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Backend (engine) contract for the Python graph: plan -> compile -> execute. +"""Backend (engine) contract for the Python graph: claim -> compile -> execute. -A backend is one of the interchangeable implementations the Router dispatches -to (Python DSLs, a naive reference, the cuDNN Graph backend, ...). The +A backend is one of the interchangeable implementations the ranked plan list +dispatches to (Python DSLs, a naive reference, the cuDNN Graph backend, ...). The lifecycle mirrors a real JIT/DSL engine: - 1. ``propose_plans(graph)`` -> candidate ``PlanConfig`` entries (one per - configuration the engine wants ranked; decline the whole graph by raising - ``NotImplementedError`` / ``cudnn.cudnnGraphNotSupportedError``). + 1. ``check_support(graph)`` -> accept, or decline the whole graph by raising + ``NotImplementedError`` / ``cudnn.cudnnGraphNotSupportedError``. An engine + does NOT propose its own plans: which configs to try, in what order, and + where the backend's own entries belong is one comparison across every + candidate, which no engine can make from the inside. That lives in + ``engines/heuristics.py`` and the family hook it dispatches to. 2. ``build_plan(graph, plan)`` -> a ``CompiledPlan`` — the expensive JIT step, run ONCE per (graph, selected plan) at ``graph.build_plans()`` time; the compiled artifact lives on the graph, so one engine instance is safely @@ -50,6 +53,18 @@ def execute(self, graph, uid_to_data, ctx): from ..pygraph import pygraph +def decline_types(): + """The exception types that mean "this engine does not serve this graph". + + ImportError counts: an engine whose optional dependency is absent cannot + serve the graph, and since lowering imports are deferred past check_support + that only becomes visible at build time. + """ + import cudnn + + return (NotImplementedError, cudnn.cudnnGraphNotSupportedError, ImportError) + + @dataclass(frozen=True) class PlanConfig: """One candidate execution plan: an engine id + its knobs. @@ -63,11 +78,16 @@ class PlanConfig: ``cpp_index`` is set only on backend entries: the position this plan holds in the lowered graph's own plan list, so building it is one ``build_plan_at_index`` instead of a rebuild from (engine_id, knobs). + + ``mode`` is the heuristic mode that produced the entry. Ranking needs it: + "the backend's mode-A entries ahead of ours, its fallbacks behind" cannot + be said about a list whose entries do not remember where they came from. """ engine_id: int knobs: Any = None cpp_index: Any = None + mode: Any = None @dataclass(frozen=True) @@ -161,7 +181,6 @@ class BaseEngine(ABC): engine_id: Stable id in the shared flat engine-id space, in the reserved Python region (>= PYTHON_ENGINE_ID_BASE). Subclasses MUST declare it; the base default (None) is rejected at register_backend(). - default_knobs: Optional default tuning knobs for this engine's plan. behavior_notes / numerical_notes: what this engine's plans are, in the same vocabulary the backend's plans answer in, so deselect_behavior_notes(...) and friends mean one thing across the @@ -174,7 +193,6 @@ class BaseEngine(ABC): # base intentionally has none so a forgotten override fails at registration # instead of silently colliding with another engine. engine_id: Any = None - default_knobs: Any = None behavior_notes: tuple = () numerical_notes: tuple = () @@ -215,16 +233,6 @@ def check_support(self, graph: "pygraph") -> None: """ _ = graph - def propose_plans(self, graph: "pygraph") -> List[PlanConfig]: - """Candidate plans for ``graph``, in this engine's preference order. - - Default: one plan with ``default_knobs`` when ``check_support`` accepts. - Engines with several viable configurations override this to expose them - to ranking/autotune (each entry's knobs reach ``build_plan`` verbatim). - """ - self.check_support(graph) - return [PlanConfig(self.engine_id, self.default_knobs)] - def build_plan(self, graph: "pygraph", plan: PlanConfig, ctx: "ExecutionContext" = None) -> CompiledPlan: """Compile ``graph`` for ``plan`` (the expensive step; run once per graph/plan at build_plans() time). ``ctx`` carries the build context — diff --git a/python/cudnn/engines/heuristics.py b/python/cudnn/engines/heuristics.py index f472bfffb..30a28a728 100644 --- a/python/cudnn/engines/heuristics.py +++ b/python/cudnn/engines/heuristics.py @@ -1,41 +1,92 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Frontend heuristics: rank python and backend plans into ONE list. +"""Frontend heuristics: produce the ranked plan list for a graph. -``create_execution_plans()`` collects both sides and hands them here. The -returned order IS the order the build walk tries. +``create_execution_plans([heur_mode.A, ...])`` collects the inputs — the parsed +facts, the engine ids on offer, and the backend's own entries tagged with the +mode that produced each — and hands all of it to the graph's family. What the +family returns IS ``graph.plans``, position for position. -PLACEHOLDER: claiming engines first, then the backend. Real ranking — per -operation, over a cost model that can compare a CuTe tile config against a cuDNN -engine — replaces the body of :func:`heuristics_sort` and nothing else. +One function, everything in view. Ranking is a comparison, so whoever ranks has +to see both sides: an engine cannot see its siblings, and a family that only +saw its own engines could not decide whether the backend belongs in front of +them. That is why there is no per-engine ``propose_plans`` and no second +merge step after this one — splitting the decision is what forced the previous +design to concatenate and call it ranking. + +An engine answers two questions only: can I serve this graph +(``check_support``) and compile me this config (``build_plan``). + +A family that declares no ``heuristics`` hook, and any engine outside a family +(``register_backend()``), falls back to one default plan per accepting engine, +ahead of the backend's entries. """ from __future__ import annotations -from typing import TYPE_CHECKING, List +import logging +from typing import TYPE_CHECKING, Any, List, Optional -from .base import PlanConfig +from .base import BaseEngine, PlanConfig, decline_types if TYPE_CHECKING: from .._pygraph import pygraph +_LOG = logging.getLogger("cudnn.engines.heuristics") + + +def default_modes() -> List[Any]: + """The modes assumed when the caller named none — the backend's own default.""" + import cudnn + + return [cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK] + + +def accepts(engine: BaseEngine, graph: "pygraph") -> bool: + """Whether ``engine`` will serve ``graph``, declines being routing not error.""" + try: + engine.check_support(graph) + except decline_types() as exc: + _LOG.debug("engine %s declined the graph: %s", engine.name, exc) + return False + return True -def heuristics_sort(graph: "pygraph", python_plans: List[PlanConfig], backend_plans: List[PlanConfig]) -> List[PlanConfig]: - """Rank the two sides into one list. Either side may be empty. - Which python engines are in the list at all is the manifest's decision - (``EngineFamily.offered``, where ``CUDNN_FRONTEND_ENABLE_FROST_ENGINES`` is - read); a plan that reaches here has already been admitted. +def _without_a_family(graph: "pygraph", engines: List[BaseEngine], backend_plans: List[PlanConfig]) -> List[PlanConfig]: + """The ranking for a graph no family speaks for: accepting engines, then the backend.""" + return [PlanConfig(e.engine_id, None) for e in engines if accepts(e, graph)] + list(backend_plans) - Placeholder ranking, deliberately. The facts plumbing this reads FROM is - already in place and unused: ``graph._facts_for(analyzer)`` holds the record - the family's analyzer produced (planning attached it before this ran), and - ``manifest.family_for(graph)`` names the family whose vocabulary it is - in. A real policy has two layers — order the family's own engines using - those facts, then merge that against the backend's entries on a common - currency (predicted time), which is the only comparison that has to work - across families. Neither is written yet; the seam is here so that writing - them does not mean re-plumbing the graph. + +def rank(graph: "pygraph", engines: List[BaseEngine], backend_plans: List[PlanConfig], modes: Optional[List[Any]] = None) -> List[PlanConfig]: + """The ranked plan list for ``graph`` — what ``create_execution_plans`` stores. + + ``engines`` are this graph's python candidates and ``backend_plans`` the + backend's own entries, each already tagged with its ``mode``. """ - return python_plans + backend_plans + from . import manifest + + modes = list(modes) if modes else default_modes() + family = manifest.family_for(graph) + recommend = manifest.resolve_heuristics(family) if family is not None else None + if recommend is None: + return _without_a_family(graph, engines, backend_plans) + + analyzer = manifest.resolve_analyzer(family) + facts = graph._facts_for(analyzer) if analyzer is not None else None + if facts is None: + # The family claims the graph by node type but its analyzer cannot + # express it. Nothing to rank its engines on; the backend serves it. + return list(backend_plans) + + offered = {e.name: e.engine_id for e in engines if family.owns(e.engine_id)} + plans = list(recommend(modes, facts, offered, list(backend_plans))) + own = set(offered.values()) + for cfg in plans: + from .engine_ids import is_python_engine + + if is_python_engine(cfg.engine_id) and cfg.engine_id not in own: + raise ValueError(f"heuristics for {family.name} returned python engine_id {cfg.engine_id}, which the family does not own or offer") + # register_backend() engines sit outside every family, so no family + # heuristic speaks for them; they go last, ahead of nothing. + return plans + [PlanConfig(e.engine_id, None) for e in engines if not family.owns(e.engine_id) and accepts(e, graph)] diff --git a/python/cudnn/engines/manifest.py b/python/cudnn/engines/manifest.py index bfd82c99c..72473cb96 100644 --- a/python/cudnn/engines/manifest.py +++ b/python/cudnn/engines/manifest.py @@ -85,6 +85,11 @@ class EngineFamily: # ("module", "callable") producing this family's facts from a graph, or None # while a family still reads the graph inside its engines. analyzer: Optional[Tuple[str, str]] = None + # ("module", "callable") recommending (engine_id, knobs) for one heuristic + # mode, given this family's facts. The family is the smallest scope that can + # rank -- an engine cannot see its siblings. None falls back to one + # default-knob plan per eligible engine. + heuristics: Optional[Tuple[str, str]] = None @property def id_end(self) -> int: @@ -172,8 +177,10 @@ def offered_ids(self) -> Dict[str, int]: "sdpa_fwd_prefill_sm100_d128_fp8": EngineSlot(4, opt_in=True), "sdpa_fwd_prefill_sm120": EngineSlot(5, opt_in=True), "sdpa_fwd_prefill_sm100_d192_d128": EngineSlot(6, opt_in=True), + "sdpa_fwd_prefill_sm120_fp8": EngineSlot(7, opt_in=True), }, analyzer=("cudnn.sdpa.graph_analyzer", "analyze"), + heuristics=("cudnn.sdpa.fwd.heuristics", "recommend"), ), EngineFamily( FROST_SDPA_BWD_ID_BASE, @@ -212,31 +219,42 @@ def family_for(graph) -> Optional[EngineFamily]: return next(f for f in MANIFEST if f.name == name) -def resolve_analyzer(family: EngineFamily): - """The family's facts callable, or None when it declares no analyzer. +def _resolve(family: EngineFamily, ref: Optional[Tuple[str, str]], what: str): + """Import a ("module", "callable") declaration, or None when absent. - Importing it is the caller's decision, not this module's: keeping - ``analyzer`` a pair of strings is what lets the coarse key stay - import-free. Planning resolves it and attaches the record to the frozen - graph; the family's engines then read that same record back. + Importing is the caller's decision, not this module's: keeping these + declarations pairs of strings is what lets the coarse key stay import-free. + A missing optional dependency makes the hook absent, not the graph + unplannable -- importing one pulls in its package (cudnn.sdpa.__init__ -> + cuda.bindings, cutlass), so without this a planning call raises instead of + falling back to the backend. """ - if family.analyzer is None: + if ref is None: return None import importlib - module, attr = family.analyzer + module, attr = ref try: return getattr(importlib.import_module(module), attr) except ImportError as exc: - # Same contract as instantiate(): a missing optional dependency makes - # the family absent, not the graph unplannable. Importing an analyzer - # pulls in its package (cudnn.sdpa.__init__ -> cuda.bindings, cutlass), - # so without this a planning call raises instead of falling back to the - # backend. - _LOG.info("analyzer for %s is unavailable in this environment: %s", family.name, exc) + _LOG.info("%s for %s is unavailable in this environment: %s", what, family.name, exc) return None +def resolve_analyzer(family: EngineFamily): + """The family's facts callable, or None when it declares no analyzer. + + Planning resolves it and attaches the record to the frozen graph; the + family's heuristics and engines then read that same record back. + """ + return _resolve(family, family.analyzer, "analyzer") + + +def resolve_heuristics(family: EngineFamily): + """The family's plan-recommending callable, or None when it declares none.""" + return _resolve(family, family.heuristics, "heuristics") + + def instantiate(family: EngineFamily, ids: Dict[str, int]): """Import the family's module and build the engines named in ``ids``. diff --git a/python/cudnn/engines/router.py b/python/cudnn/engines/router.py deleted file mode 100644 index d6639d0ac..000000000 --- a/python/cudnn/engines/router.py +++ /dev/null @@ -1,87 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Router: builds the ranked execution-plan list at create_execution_plans() time. - - Python Graph API -> create_execution_plans() -> Router -> one ranked list - (graph.plans) - -The list is flat ``PlanConfig(engine_id, knobs)`` entries mixing python engines -and backend engines in any order. Position is rank; ``engine_id`` is identity -and the key the build walk dispatches on — the two are independent, so an -engine from anywhere can sit anywhere in the list. - -Candidates come from ``engines.manifest`` (the library's own static table) plus -anything the caller added with ``register_backend()``. The backend's own ranked -entries come from ``graph.backend_plan_entries()``, which is [] when the backend -declined the graph or is not installed. Merging the two sides is -``engines.heuristics.heuristics_sort``'s job — that is the seam a real cost -model replaces. - -Policy is pluggable at three levels: subclass ``Router`` and override ``plan()``; -pass one per graph (``pygraph(router=...)`` / ``graph.set_router()``); or swap -the process-wide ``default_router``. A Router places the backend's entries by -calling ``graph.backend_plan_entries()`` (answered once per graph) and putting -the result where it wants — there is no placeholder to expand, so the list a -Router returns is the list, position for position. -""" - -import logging -from typing import TYPE_CHECKING, List - -from . import heuristics -from .base import BaseEngine, PlanConfig - -if TYPE_CHECKING: - from .._pygraph import pygraph - -_LOG = logging.getLogger("cudnn.engines.router") - - -def decline_types(): - """The exception types that mean "this engine does not serve this graph". - - ImportError counts: an engine whose optional dependency is absent cannot - serve the graph, and since lowering imports are deferred past check_support - that only becomes visible at build time. - """ - import cudnn - - return (NotImplementedError, cudnn.cudnnGraphNotSupportedError, ImportError) - - -class Router: - """Default policy: every claiming python engine, ranked against the backend.""" - - def python_plans(self, graph: "pygraph", engines: List[BaseEngine]) -> List[PlanConfig]: - """Proposals from the python engines that claim ``graph``, in candidate order. - - An engine declines with ``NotImplementedError`` / - ``cudnnGraphNotSupportedError`` only; any other exception is a bug in the - engine and propagates instead of silently costing the user a kernel. - """ - decline = decline_types() - out: List[PlanConfig] = [] - for engine in engines: - try: - proposals = engine.propose_plans(graph) - except decline as exc: - _LOG.debug("engine %s declined the graph: %s", engine.name, exc) - continue - for cfg in proposals: - lo, hi = engine.owned_id_range - if not lo <= cfg.engine_id < hi: # no identity injection - raise ValueError(f"engine {engine.name!r} proposed a plan with foreign engine_id {cfg.engine_id}") - out.extend(proposals) - return out - - def plan(self, graph: "pygraph", engines: List[BaseEngine]) -> List[PlanConfig]: - # Through the module, not a bound name: heuristics.heuristics_sort is the - # seam a real ranking policy replaces, and it must be swappable in a - # live process (tests, experiments) without touching this file. - return heuristics.heuristics_sort(graph, self.python_plans(graph, engines), graph.backend_plan_entries()) - - -# Process-wide default. Assign a Router subclass to change global policy, or pass -# one to pygraph(router=...) / graph.set_router(...) per graph. -default_router = Router() diff --git a/python/cudnn/frost/tile_dsl/mma.py b/python/cudnn/frost/tile_dsl/mma.py index 66340c6f2..b2bf2da36 100644 --- a/python/cudnn/frost/tile_dsl/mma.py +++ b/python/cudnn/frost/tile_dsl/mma.py @@ -42,6 +42,80 @@ def ptx_mma_m16n8k16_f32( ) +@cute.jit +def ptx_mma_m16n8k32_e4m3_f32( + a0: cutlass.Int32, + a1: cutlass.Int32, + a2: cutlass.Int32, + a3: cutlass.Int32, + b0: cutlass.Int32, + b1: cutlass.Int32, + c0: cutlass.Float32, + c1: cutlass.Float32, + c2: cutlass.Float32, + c3: cutlass.Float32, +) -> tuple[cutlass.Float32, cutlass.Float32, cutlass.Float32, cutlass.Float32]: + """``mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32`` (SM89+/SM120). + + The C operands travel as ``Int32`` bit patterns and are ``mov.b32``'d into + ``.f32`` temps inside the asm block: cutlass-dsl 4.7.0's ``inline_ptx`` + fails in libNVVM when a compile-time-constant ``Float32`` reaches + ``read_only_args``, and an accumulator's zero-init can fold to a constant. + """ + return cute.arch.inline_ptx( + "{ .reg .f32 fc<4>; " + "mov.b32 fc0, {$r6}; mov.b32 fc1, {$r7}; mov.b32 fc2, {$r8}; mov.b32 fc3, {$r9}; " + "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " + "{{$w0},{$w1},{$w2},{$w3}}, {{$r0},{$r1},{$r2},{$r3}}, {{$r4},{$r5}}, {fc0,fc1,fc2,fc3}; }", + write_only_types=[ + cutlass.Float32, + cutlass.Float32, + cutlass.Float32, + cutlass.Float32, + ], + read_only_args=[ + a0, + a1, + a2, + a3, + b0, + b1, + c0.bitcast(cutlass.Int32), + c1.bitcast(cutlass.Int32), + c2.bitcast(cutlass.Int32), + c3.bitcast(cutlass.Int32), + ], + ) + + +@cute.jit +def ptx_cvt_e4m3x2(hi: cutlass.Float32, lo: cutlass.Float32) -> cutlass.Uint16: + """Pack two fp32 into e4m3 bytes: low byte = e4m3(lo), byte 1 = e4m3(hi). + + ``cvt.rn.satfinite.e4m3x2.f32`` matches torch's ``.to(float8_e4m3fn)`` + bit-exactly. Operands ride as Int32 bit patterns for the same + constant-operand ``inline_ptx`` reason as :func:`ptx_mma_m16n8k32_e4m3_f32`. + + Stays 16-bit so two results pair into one MMA operand register with + :func:`pack_f8x2_pairs`. + """ + return cute.arch.inline_ptx( + "{ .reg .f32 fa, fb; " "mov.b32 fa, {$r0}; mov.b32 fb, {$r1}; " "cvt.rn.satfinite.e4m3x2.f32 {$w0}, fa, fb; }", + write_only_types=[cutlass.Uint16], + read_only_args=[hi.bitcast(cutlass.Int32), lo.bitcast(cutlass.Int32)], + ) + + +@cute.jit +def pack_f8x2_pairs(pair0: cutlass.Uint16, pair1: cutlass.Uint16) -> cutlass.Int32: + """Two e4m3x2 halves into one 32-bit MMA A/B operand (pair0 = low half).""" + return cute.arch.inline_ptx( + "mov.b32 $0, {$1, $2};", + write_only_types=[cutlass.Int32], + read_only_args=[pair0, pair1], + ) + + @cute.jit def mma_ss(desc, desc_a_base, desc_b_base, tmem_c, tmem_sf_a=None, tmem_sf_b=None, accumulate: bool = False, k_start: int = 0, k_count=None): if cutlass.const_expr(desc.cta_group == 1): diff --git a/python/cudnn/gemm/frost/engine.py b/python/cudnn/gemm/frost/engine.py index 897e9e26e..eee374a8c 100644 --- a/python/cudnn/gemm/frost/engine.py +++ b/python/cudnn/gemm/frost/engine.py @@ -79,13 +79,6 @@ def check_support(self, graph: "pygraph") -> None: # at the engine boundary that is a decline, not a user error. raise NotImplementedError(f"frost_gemm: {exc}") from exc - def propose_plans(self, graph: "pygraph") -> List[PlanConfig]: - # One plan today: tile selection happens inside build. Exposing - # kernel_registry.candidates() as several knob-bearing plans is what - # makes routed autotune work over FROST configs — a follow-up. - self.check_support(graph) - return [PlanConfig(self.engine_id, None)] - def build_plan(self, graph: "pygraph", plan: PlanConfig, ctx: ExecutionContext = None) -> CompiledPlan: from .graph_analyzer import build_gemm_plan diff --git a/python/cudnn/sdpa/bwd/engine.py b/python/cudnn/sdpa/bwd/engine.py index 5eecbf177..1ed5a5b44 100644 --- a/python/cudnn/sdpa/bwd/engine.py +++ b/python/cudnn/sdpa/bwd/engine.py @@ -100,14 +100,6 @@ def check_support(self, graph: "pygraph") -> None: if reason is not None: raise NotImplementedError(f"{self.name}: {reason}") - def propose_plans(self, graph: "pygraph") -> List[PlanConfig]: - # One plan, no knobs: nothing proposes a tuning request today, so the - # engine runs at its capability row's default tiles. A knob search - # (SdpaBwdKnobs over Capabilities.tile_ms/tile_ns) becomes several - # entries here; each one's knobs reach build_plan verbatim. - self.check_support(graph) - return [PlanConfig(self.engine_id, self.default_knobs)] - def build_plan(self, graph: "pygraph", plan: PlanConfig, ctx: ExecutionContext = None) -> CompiledPlan: from .engines import build diff --git a/python/cudnn/sdpa/bwd/engines.py b/python/cudnn/sdpa/bwd/engines.py index 7fbae51d2..cd03ef2ea 100644 --- a/python/cudnn/sdpa/bwd/engines.py +++ b/python/cudnn/sdpa/bwd/engines.py @@ -251,14 +251,6 @@ def analyze_for(spec: EngineSpec, graph, knobs: Optional[SdpaBwdKnobs] = None): return facts, mismatch(spec.capabilities, facts, knobs) -def probe(spec: EngineSpec, graph, knobs: Optional[SdpaBwdKnobs] = None) -> bool: - _, reason = analyze_for(spec, graph, knobs) - if reason is not None: - _LOG.debug("cudnn.sdpa: %s ineligible: %s", spec.name, reason) - return False - return True - - def build(spec: EngineSpec, graph, knobs: Optional[SdpaBwdKnobs] = None): """Lower ``spec`` for ``graph``, or raise the bare ineligibility reason (the caller — the engine — names itself in the message).""" diff --git a/python/cudnn/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py index 3eb8e44e8..1b293c693 100644 --- a/python/cudnn/sdpa/fwd/api_dsl.py +++ b/python/cudnn/sdpa/fwd/api_dsl.py @@ -69,6 +69,9 @@ _SM100_TILE_N = 128 _SM120_KERNEL_FILE = "prefill_f16_sm120.py" +# Per-tensor FP8 sibling (E4M3 in as Uint8 storage, FP16 out, mma.sync +# m16n8k32); selected by the graph op (sdpa_fp8) via check_support's dtype. +_SM120_FP8_KERNEL_FILE = "prefill_fp8_sm120.py" _SM120_DTYPE_QKV_CODE = { torch.bfloat16: DTYPE_BF16, torch.float16: DTYPE_FP16, @@ -204,7 +207,9 @@ def _load_sm100_kernel_module(flavor: tuple[int, int], params: Sm100TemplatePara return _load_kernel_template(filename, params, tag) -def _load_sm120_kernel_module(params: Sm120TemplateParams): +def _load_sm120_kernel_module(params: Sm120TemplateParams, fp8: bool = False): + if fp8: + return _load_kernel_template(_SM120_FP8_KERNEL_FILE, params, tag="sdpa_fwd_sm120_fp8") return _load_kernel_template(_SM120_KERNEL_FILE, params, tag="sdpa_fwd_sm120") @@ -1564,18 +1569,36 @@ def check_support(self) -> bool: f"D_V ({d_v}) must be one of {_SM120_SUPPORTED_HEAD_TILES}", ) - self.dtype = self._check_dtype(self.q_desc, [torch.float16, torch.bfloat16], name="Q") + self.dtype = self._check_dtype(self.q_desc, [torch.float16, torch.bfloat16, torch.float8_e4m3fn], name="Q") + self._fp8 = self.dtype == torch.float8_e4m3fn for desc in (self.k_desc, self.v_desc, self.o_desc): - self._check_dtype( - desc, - self.dtype, - name=desc.name, - extra_error_msg=f"{desc.name} must match Q", - ) + if self._fp8 and desc is self.o_desc: + # The fp8 kernel's epilogue emits FP16 only (see the kernel + # docstring); fp8 O would need a scale_o quantizing store. + self._check_dtype(desc, torch.float16, name="O", extra_error_msg="SM120 fp8 emits FP16 O only") + else: + self._check_dtype( + desc, + self.dtype, + name=desc.name, + extra_error_msg=f"{desc.name} must match Q", + ) self._value_error_if( desc.device != self.q_desc.device, f"{desc.name} must be on device {self.q_desc.device}, got {desc.device}", ) + if self._fp8: + self._value_error_if( + not self._pertensor, + "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)") + self._value_error_if(self.seq_q_lens_present, "SM120 fp8 does not support per-batch seq_len_q") + self._value_error_if( + (d_q, d_v) != (128, 128), + f"SM120 fp8 requires D_QK=D_V=128 (no zero-padding envelope on the 8-bit fragment path); got ({d_q}, {d_v})", + ) self._value_error_if( self.q_desc.device.type != "cuda", @@ -1616,11 +1639,18 @@ def check_support(self) -> bool: def _smem_bytes(kv_tile: int) -> int: # One K tile (D_QK wide) + one V tile (D_V wide), aliased with the - # q_tile x D_V output staging tile. - return max(kv_tile * (d_q + d_v) * self.dtype.itemsize, self.q_tile * d_v * self.dtype.itemsize) + 16 + # q_tile x D_V output staging tile. FP8: KV elements are 1 byte + # but O staging is FP16 (2 bytes). + o_item = 2 if self._fp8 else self.dtype.itemsize + base = max(kv_tile * (d_q + d_v) * self.dtype.itemsize, self.q_tile * d_v * o_item) + return base + 16 if self.tile_n is None: - # Pick the largest KV tile that fits this device. + # Largest KV tile this device fits. RESOURCE feasibility, not a + # performance choice: which tile is fastest is a measurement, and it + # is made once, in sdpa/fwd/heuristics.py. A graph reaching here + # already carries the tiles that decided on; only a direct caller of + # this adapter leaves them unset. self.kv_tile = next((t for t in _SM120_KV_TILES if _smem_bytes(t) <= smem_capacity_bytes), self.kv_tile) self._not_implemented_error_if( _smem_bytes(self.kv_tile) > smem_capacity_bytes, @@ -1654,7 +1684,7 @@ def compile(self) -> None: return params = Sm120TemplateParams( - dtype_qkv=_SM120_DTYPE_QKV_CODE[self.dtype], + dtype_qkv=DTYPE_E4M3 if self._fp8 else _SM120_DTYPE_QKV_CODE[self.dtype], is_causal=self.is_causal, causal_bottom_right=self.causal_bottom_right, window_size_left=self.window_size_left, @@ -1665,7 +1695,7 @@ def compile(self) -> None: q_tile=self.q_tile, kv_tile=self.kv_tile, ) - self._k_mod = _load_sm120_kernel_module(params) + self._k_mod = _load_sm120_kernel_module(params, fp8=self._fp8) if self.thd: # The packed token totals (and max sequence length) are runtime # values, so the per-shape compile is deferred to execute(). @@ -1700,6 +1730,15 @@ def execute( scale_softmax: Optional[float] = None, workspace: Optional[torch.Tensor] = None, current_stream: Optional[cuda.CUstream] = None, + descale_q: Optional[torch.Tensor] = None, + descale_k: Optional[torch.Tensor] = None, + descale_v: Optional[torch.Tensor] = None, + scale_o: Optional[torch.Tensor] = None, + amax_s: Optional[torch.Tensor] = None, + amax_o: Optional[torch.Tensor] = None, + sf_q: Optional[torch.Tensor] = None, + sf_k: Optional[torch.Tensor] = None, + sf_v: Optional[torch.Tensor] = None, ) -> None: """Execute tensors matching the compiled specialization.""" @@ -1723,6 +1762,28 @@ def execute( "this specialization was compiled without an LSE output; construct the API with sample_lse", ) scale_val = self.scale_softmax if scale_softmax is None or scale_softmax == 0.0 else float(scale_softmax) + if self._fp8: + self._value_error_if( + any(t is not None for t in (sf_q, sf_k, sf_v)), + "SM120 fp8 is per-tensor (scalar descales); block-scale SF tensors are MXFP8-only", + ) + self._execute_fp8( + q_tensor, + k_tensor, + v_tensor, + o_tensor, + lse_tensor, + scale_val, + seq_kv_lens, + descale_q, + descale_k, + descale_v, + scale_o, + amax_s, + amax_o, + current_stream=current_stream, + ) + return scale_softmax_log2 = scale_val * math.log2(math.e) if self.thd: self._execute_thd( @@ -1787,6 +1848,99 @@ def execute( if o_needs_copy_back: o_view.copy_(o_scratch) + def _execute_fp8( + self, + q_tensor, + k_tensor, + v_tensor, + o_tensor, + lse_tensor, + scale_val, + seq_kv_lens, + descale_q, + descale_k, + descale_v, + scale_o, + amax_s, + amax_o, + current_stream=None, + ): + """Per-tensor FP8 execute (dense): SM100 convention on the SM120 kernel. + + ``descale_q*descale_k`` folds into ``scale_softmax_log2`` and + ``descale_v*scale_o`` into the kernel's ``o_scale_fused`` scalar. + ``Amax_S`` is produced in-kernel (bitcast-int32 atomicMax of the + per-row ``1/row_sum``) into the caller's pre-zeroed buffer; ``Amax_O`` + is ``max|o_scaled|/scale_o`` post-kernel (exact for the FP16 output). + E4M3 tensors travel as ``uint8`` views — the kernel consumes bit + patterns (see the kernel docstring). + """ + import cutlass + + 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) + scale_softmax_log2 = scale_val * dq * dk * math.log2(math.e) + o_scale_fused = dv * so + device = q_tensor.device + + self._value_error_if( + self.lse_desc is not None and lse_tensor is None, + "lse_tensor is required by this compiled specialization", + ) + self._value_error_if( + self.lse_desc is None and lse_tensor is not None, + "this specialization was compiled without an LSE output; construct the API with sample_lse", + ) + lse = self._checked_lse_view(lse_tensor) if lse_tensor is not None else None + seq_kv_t = ( + self._checked_seq_lens(seq_kv_lens, "seq_kv_lens") + if seq_kv_lens is not None + else self._dummy("seq_kv_lens", device, lambda: torch.zeros(self.batch_size, dtype=torch.int32, device=device)) + ) + seq_q_dummy = self._dummy("seq_q_lens", device, lambda: torch.zeros(self.batch_size, dtype=torch.int32, device=device)) + if current_stream is None: + current_stream = cuda.CUstream(torch.cuda.current_stream(device).cuda_stream) + + q = self._to_bshd(q_tensor).view(torch.uint8) + k = self._to_bshd(k_tensor).view(torch.uint8) + v = self._to_bshd(v_tensor).view(torch.uint8) + o_view, o_needs_copy_back, o_scratch = self._to_bshd_writable(o_tensor) + o = o_scratch if o_needs_copy_back else o_view + + # amax_s / amax_o: the kernel atomicMax'es into these buffers, so they + # MUST start at 0, reset on the LAUNCH stream (ordering vs the kernel). + amax_s_buf = amax_s.reshape(-1)[:1] if amax_s is not None else self._dummy("amax_s", device, lambda: torch.zeros(1, dtype=torch.float32, device=device)) + amax_o_buf = amax_o.reshape(-1)[:1] if amax_o is not None else self._dummy("amax_o", device, lambda: torch.zeros(1, dtype=torch.float32, device=device)) + with _torch_stream_context(current_stream, device): + amax_s_buf.zero_() + amax_o_buf.zero_() + + self._compiled_kernel( + q, + k, + v, + o, + lse, + None, # sinks: fp8 cell rejects has_sink + seq_q_dummy, + seq_kv_t, + amax_s_buf.view(torch.int32), + amax_o_buf.view(torch.int32), + cutlass.Float32(scale_softmax_log2), + cutlass.Float32(o_scale_fused), + current_stream, + ) + # Both of these consume what the kernel just wrote, so they belong on + # the launch stream for the same reason the resets above do. + with _torch_stream_context(current_stream, device): + if o_needs_copy_back: + o_view.copy_(o_scratch) + if amax_o is not None: + amax_o_buf.div_(max(so, 1e-30)) + self._logger.debug("execute (SM120 FP8 per-tensor) completed") + def _execute_thd( self, q_buf, k_buf, v_buf, o_buf, scale_softmax_log2, sinks, seq_kv_lens, seq_q_lens, lse_tensor=None, workspace=None, current_stream=None ): diff --git a/python/cudnn/sdpa/fwd/config_sm120.py b/python/cudnn/sdpa/fwd/config_sm120.py index 8e2d2c07f..041b1cb63 100644 --- a/python/cudnn/sdpa/fwd/config_sm120.py +++ b/python/cudnn/sdpa/fwd/config_sm120.py @@ -7,13 +7,45 @@ from dataclasses import dataclass -from cudnn.frost.tile_dsl.constants import DTYPE_BF16, DTYPE_FP16 +from cudnn.frost.tile_dsl.constants import DTYPE_BF16, DTYPE_E4M3, DTYPE_FP16 # noqa: F401 (DTYPE_E4M3 re-exported for the FP8 template) SEQ_Q_TILES = (128, 64) SEQ_KV_TILES = (128, 64) SUPPORTED_HEAD_TILES = tuple(range(16, 257, 16)) +def tile_choice(s_q: int, s_kv: int, h_q: int, batch: int, sm_count: int, is_causal: bool) -> tuple[int, int]: + """(q_tile, kv_tile) for the SM120 SDPA-forward cells, fp8 and f16 alike. + + ``kv_tile=128`` unconditionally -- fastest in all 28 fp8 shapes measured, + and in the f16 sweep too (by 2-4% rather than a uniform margin). An earlier + revision of the fp8 kernel staged P through SMEM, and that traffic grew + with the KV tile, which made 64 the better choice; the shfl path removed it + and the optimum moved. A tile rule is a property of the kernel it was + measured on, not of the hardware. + + ``q_tile=64`` when the grid cannot fill the machine AND each CTA has enough + KV tiles for the extra Q-tile loop to amortize: halving the Q tile doubles + the CTA count, which only pays while SMs sit idle, and adds loop overhead + that only long sequences absorb. Worth 1.5x at 64 CTAs. The 1.5x SM-count + bound is where the two stop trading evenly: 240 CTAs still want the finer + tile on this part, 320 want the coarser one by 1.19x. + + 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. Both cells want that term; see the doc for the four + datasets it was measured against. + """ + if sm_count <= 0: + return 128, 128 + grid = -(-s_q // 128) * h_q * batch + if is_causal: + grid //= 2 + kv_tiles = -(-s_kv // 128) + fine = grid * 2 <= sm_count or (grid * 2 <= 3 * sm_count and kv_tiles >= 12) + return (64 if fine else 128), 128 + + @dataclass(frozen=True) class TemplateParams: """Per-graph parameters that change the traced SM120 kernel. @@ -36,16 +68,17 @@ class TemplateParams: kv_tile: int = SEQ_KV_TILES[0] -def validate_params(params: TemplateParams) -> None: +def validate_params(params: TemplateParams, allowed_dtypes: tuple[int, ...] = (DTYPE_BF16, DTYPE_FP16)) -> None: """Validate the SM120 template specialization. Reachable failures should already have been rejected by the engine capabilities or adapter support checks; this validation is a backstop for - direct template use. + direct template use. ``allowed_dtypes`` defaults to the FP16/BF16 template's + set; the FP8 template passes its own. """ - if params.dtype_qkv not in (DTYPE_BF16, DTYPE_FP16): - raise ValueError(f"SM120 SDPA: dtype_qkv must be DTYPE_BF16 ({DTYPE_BF16}) or DTYPE_FP16 ({DTYPE_FP16}); got {params.dtype_qkv}") + if params.dtype_qkv not in allowed_dtypes: + raise ValueError(f"SM120 SDPA: dtype_qkv must be one of {allowed_dtypes}; got {params.dtype_qkv}") if params.causal_bottom_right and not params.is_causal: raise ValueError("SM120 SDPA: causal_bottom_right requires is_causal=True") if params.window_size_left is not None and params.window_size_left < 0: diff --git a/python/cudnn/sdpa/fwd/engine.py b/python/cudnn/sdpa/fwd/engine.py index 1a17a4518..bbb2722a9 100644 --- a/python/cudnn/sdpa/fwd/engine.py +++ b/python/cudnn/sdpa/fwd/engine.py @@ -15,7 +15,7 @@ only the engine contract around them. """ -from typing import TYPE_CHECKING, Any, List, Optional +from typing import TYPE_CHECKING, Any, List from cudnn import behavior_note from cudnn.engines.base import BaseEngine, CompiledPlan, ExecutionContext, PlanConfig @@ -87,29 +87,26 @@ def __init__(self, spec: "EngineSpec", engine_id: int): self.name = spec.name self.engine_id = engine_id - def _decline_reason(self, graph: "pygraph", knobs) -> Optional[str]: + def _facts_or_decline(self, graph: "pygraph"): + """The parsed graph, or NotImplementedError naming why this cell declines. + + The one eligibility question this engine asks, so callers that also need + the facts do not ask it twice. + """ from .engines import analyze_for try: - _, reason = analyze_for(self._spec, graph, knobs) + facts, reason = analyze_for(self._spec, graph, None) except ValueError as exc: # ValueError is the analyzer's internal "cannot express this graph"; # at the engine boundary that is a decline, not a user error. - return str(exc) - return reason - - def check_support(self, graph: "pygraph") -> None: - reason = self._decline_reason(graph, None) + facts, reason = None, str(exc) if reason is not None: raise NotImplementedError(f"{self.name}: {reason}") + return facts - def propose_plans(self, graph: "pygraph") -> List[PlanConfig]: - # One plan, no knobs: nothing proposes a tuning request today, so the - # engines run at their capability row's default tile/schedule. A knob - # search (SdpaFwdKnobs over Capabilities.tile_ms/tile_ns/cgas) becomes - # several entries here; each one's knobs reach build_plan verbatim. - self.check_support(graph) - return [PlanConfig(self.engine_id, self.default_knobs)] + def check_support(self, graph: "pygraph") -> None: + self._facts_or_decline(graph) def build_plan(self, graph: "pygraph", plan: PlanConfig, ctx: ExecutionContext = None) -> CompiledPlan: from .engines import build diff --git a/python/cudnn/sdpa/fwd/engines.py b/python/cudnn/sdpa/fwd/engines.py index 3cd0dc6e9..d6b232b03 100644 --- a/python/cudnn/sdpa/fwd/engines.py +++ b/python/cudnn/sdpa/fwd/engines.py @@ -113,6 +113,10 @@ class Capabilities: # SM120, whose lowering has no zero-padding path wired yet). d_envelope: bool = False dtypes: frozenset = frozenset({cudnn.data_type.HALF, cudnn.data_type.BFLOAT16}) # cudnn.data_type, see graph_analyzer + # O dtype domain. Only the quantized rows declare it: elsewhere O must + # equal Q, which facts.uniform_dtype already enforces. A quantized row + # that leaves it empty serves nothing, which is the loud failure. + out_dtypes: frozenset = frozenset() is_mxfp8: bool = False # block-scale MXFP8 engine (FP8 in + per-32-block E8M0 SF) is_fp8: bool = False # per-tensor FP8 engine (FP8 in + scalar descales) @@ -256,6 +260,8 @@ def mismatch(capabilities: Capabilities, facts: "ga.SdpaGraphFacts", knobs: Opti return f"serves D_QK in {sorted(capabilities.d_qk)}/D_V in {sorted(capabilities.d_v)}; graph has D_QK={facts.d_qk}/D_V={facts.d_v}" if facts.dtype not in capabilities.dtypes: return f"dtype {facts.dtype} not in {sorted(str(d) for d in capabilities.dtypes)}" + if (capabilities.is_fp8 or capabilities.is_mxfp8) and facts.dtype_o not in capabilities.out_dtypes: + return f"O dtype {facts.dtype_o} not in {sorted(str(d) for d in capabilities.out_dtypes)}" if (facts.is_mxfp8, facts.is_fp8) != (capabilities.is_mxfp8, capabilities.is_fp8): quant = "block-scale MXFP8 (sdpa_mxfp8)" if capabilities.is_mxfp8 else "per-tensor FP8 (sdpa_fp8)" if capabilities.is_fp8 else "half (sdpa)" return f"this engine serves only {quant} graphs" @@ -399,6 +405,7 @@ def _sm100_mxfp8_spec(d: int) -> EngineSpec: d_qk=frozenset({d}), d_v=frozenset({d}), dtypes=frozenset({cudnn.data_type.FP8_E4M3, cudnn.data_type.FP8_E5M2}), + out_dtypes=frozenset({cudnn.data_type.HALF, cudnn.data_type.BFLOAT16, cudnn.data_type.FP8_E4M3, cudnn.data_type.FP8_E5M2}), is_mxfp8=True, causal=True, bottom_right=True, @@ -434,6 +441,7 @@ def _sm100_fp8_spec(d: int) -> EngineSpec: d_qk=frozenset({d}), d_v=frozenset({d}), dtypes=frozenset({cudnn.data_type.FP8_E4M3, cudnn.data_type.FP8_E5M2}), + out_dtypes=frozenset({cudnn.data_type.HALF, cudnn.data_type.BFLOAT16, cudnn.data_type.FP8_E4M3, cudnn.data_type.FP8_E5M2}), is_fp8=True, causal=True, bottom_right=True, @@ -495,12 +503,56 @@ def _sm120_spec() -> EngineSpec: ) +def _sm120_fp8_spec() -> EngineSpec: + """SM120 per-tensor FP8 engine (E4M3 in + scalar descales, FP16 out). + + Same mma.sync architecture as the f16 SM120 cell with the MMA lowered to + m16n8k32 e4m3; ``descale_q*descale_k`` folds into the softmax scale and + ``descale_v*scale_o`` into an epilogue scalar, so the kernel adds only the + Amax_S/Amax_O atomics over the f16 sibling. E4M3 only (no E5M2 tag in the + kernel yet), FP16 O only, exact d128 (no zero-padding envelope on the + 8-bit fragment path), no sink (Amax_S semantics with a sink column are + undefined here), and THD deferred like the SM100 fp8 v1. + """ + + return EngineSpec( + name="sdpa_fwd_prefill_sm120_fp8", + capabilities=Capabilities( + sm_lo=_BLACKWELL_GEFORCE[0], + sm_hi=_BLACKWELL_GEFORCE[1], + phase="prefill", + d_qk=frozenset({128}), + d_v=frozenset({128}), + dtypes=frozenset({cudnn.data_type.FP8_E4M3}), + out_dtypes=frozenset({cudnn.data_type.HALF}), + is_fp8=True, + causal=True, + bottom_right=True, + bottom_right_with_swa=True, + bottom_right_padded_seq_q=True, + swa=True, + padded=True, + stats=True, + lse_optional=True, + # Same caveat as the SM100 fp8 row: no SEQ_Q_LENS epilogue trim, + # but fp8 graphs cannot carry seq_len_q here (lower_dsl_prefill + # forces seq_q_lens_present=False for the fp8 family). + padded_stats=True, + sched_policies=frozenset({SCHED_NATURAL}), + tile_ms=frozenset({64, 128}), + tile_ns=frozenset({64, 128}), + cgas=frozenset({1}), + ), + lower=partial(lower_dsl_prefill, api_type=_SM120), + ) + + def analyze_for(spec: EngineSpec, graph, knobs: Optional[SdpaFwdKnobs] = None): """``(facts, reason)``: the parsed graph and the first reason ``spec`` cannot serve it under ``knobs`` (``None`` when it can). The single eligibility entry point, shared by :func:`probe`, :func:`build` - and ``engine.FrostSdpaFwdEngine.check_support``. ``knobs`` is the plan's + and ``engine.FrostSdpaFwdEngine._facts_or_decline``. ``knobs`` is the plan's tuning request (``PlanConfig.knobs``), ``None`` for no preference. """ # The record validate() attached, not a fresh parse: one per graph, shared @@ -511,14 +563,6 @@ def analyze_for(spec: EngineSpec, graph, knobs: Optional[SdpaFwdKnobs] = None): return facts, mismatch(spec.capabilities, facts, knobs) -def probe(spec: EngineSpec, graph, knobs: Optional[SdpaFwdKnobs] = None) -> bool: - _, reason = analyze_for(spec, graph, knobs) - if reason is not None: - _LOG.debug("cudnn.sdpa: %s ineligible: %s", spec.name, reason) - return False - return True - - def build(spec: EngineSpec, graph, knobs: Optional[SdpaFwdKnobs] = None): """Lower ``spec`` for ``graph``, or raise the bare ineligibility reason (the caller — the engine — names itself in the message).""" @@ -767,6 +811,7 @@ def engine_name( _sm100_mxfp8_spec(128), _sm100_fp8_spec(128), _sm120_spec(), + _sm120_fp8_spec(), ) -__all__ = ["Capabilities", "EngineSpec", "ENGINE_SPECS", "SdpaFwdKnobs", "analyze_for", "build", "engine_name", "mismatch", "probe"] +__all__ = ["Capabilities", "EngineSpec", "ENGINE_SPECS", "SdpaFwdKnobs", "analyze_for", "build", "engine_name", "mismatch"] diff --git a/python/cudnn/sdpa/fwd/heuristics.py b/python/cudnn/sdpa/fwd/heuristics.py new file mode 100644 index 000000000..9da7df115 --- /dev/null +++ b/python/cudnn/sdpa/fwd/heuristics.py @@ -0,0 +1,166 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""How the SDPA-forward family ranks plans for a graph. + +``engines/heuristics.rank`` hands over the parsed facts, this family's engine +ids, and the backend's entries tagged by mode; what :func:`recommend` returns is +``graph.plans``, position for position. The whole comparison is here because +that is the only place it can be made — a cell cannot see its siblings, and +neither side of the FROST/backend split can place the other. + +Per mode: + +- **A** — candidates that are all worth running, best guess first. A caller who + does not autotune takes entry 0; one who does builds the first few and times + them. The tile rule (``config_sm120.tile_choice``) picks the guess. +- **FALLBACK** — configs that are expected to build wherever mode A might not, + ordered cheapest-resource first. Nothing here is chosen for speed. + +``heur_mode.OPENSOURCE`` is mode A without the backend's recommendation: these +cells ARE the open-source implementation, and the backend's engines are not. +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. + +``heur_mode.B`` is answered as A: it asks for a wider search than A, and this +family has none to give. +""" + +from __future__ import annotations + +from typing import Any, Dict, List + +import cudnn + +from cudnn.engines.base import PlanConfig +from cudnn.sdpa.fwd.config_sm120 import tile_choice +from cudnn.sdpa.fwd.engines import ENGINE_SPECS, Capabilities, SdpaFwdKnobs, mismatch + +# Cells whose tile choice has measurements behind it (see +# docs/fe-oss-apis/attention/sdpa-fp8-sm120.md). Named rather than derived from +# the capability row: a rule belongs to the kernel it was measured on, and a new +# cell inheriting one by accident of its knob domain is how a tile rule +# outlives its evidence. +_TILE_RULE_CELLS = frozenset({"sdpa_fwd_prefill_sm120", "sdpa_fwd_prefill_sm120_fp8"}) + +# Cells timed against the backend's own kernel and found SLOWER: the backend's +# mode-A entries lead and the cell is the second choice. sm120 fp8 is not here +# because it measures 1.20-1.83x the backend's native fp8 fprop across 28 +# shapes (docs/fe-oss-apis/attention/sdpa-fp8-sm120.md). +# +# A cell absent from this set has either measured faster or not been timed at +# all; both lead, which is deliberate. Coverage does not depend on it -- a +# caller who wants FROST tried first asks for heur_mode.OPENSOURCE. +_MEASURED_BEHIND: frozenset = frozenset() + + +def _sole(values): + """The only value on an axis, or None where the row offers a choice.""" + return next(iter(values)) if len(values) == 1 else None + + +def _knobs(caps: Capabilities, tile_m, tile_n) -> SdpaFwdKnobs: + """A knob request for one point. A field stays None only where the row + declares no domain for that axis — the engine then has no say to honour.""" + return SdpaFwdKnobs(sched_policy=_sole(caps.sched_policies), tile_m=tile_m, tile_n=tile_n, cga=_sole(caps.cgas)) + + +def _admissible(caps: Capabilities, facts, knobs: SdpaFwdKnobs) -> bool: + return mismatch(caps, facts, knobs) is None + + +def _eligible(facts, offered: Dict[str, int]): + """(engine_id, spec) for each offered cell whose capability row admits ``facts``.""" + for spec in ENGINE_SPECS: + engine_id = offered.get(spec.name) + if engine_id is not None and mismatch(spec.capabilities, facts, None) is None: + yield engine_id, spec + + +def _mode_a(facts, offered: Dict[str, int]) -> List[PlanConfig]: + """Candidates worth timing, best guess first.""" + out = [] + for engine_id, spec in _eligible(facts, offered): + caps = spec.capabilities + if spec.name not in _TILE_RULE_CELLS: + # No rule measured for this cell: its capability row has one point + # per axis, so there is nothing to choose between anyway. + knobs = _knobs(caps, _sole(caps.tile_ms), _sole(caps.tile_ns)) + if _admissible(caps, facts, knobs): + out.append(PlanConfig(engine_id, knobs, mode=cudnn.heur_mode.A)) + continue + best = tile_choice(facts.s_q, facts.s_kv, facts.h_q, facts.b, facts.device_sm_count or 0, facts.causal) + # The guess first, then the rest of the domain as autotune candidates: + # the rule's regret is small but not zero, so the runner-up is worth + # offering to a caller who measures. + ordered = sorted( + ((m, n) for m in caps.tile_ms for n in caps.tile_ns), + key=lambda mn: (mn != best, mn[1] != best[1], -mn[0]), + ) + for tile_m, tile_n in ordered: + knobs = _knobs(caps, tile_m, tile_n) + if _admissible(caps, facts, knobs): + out.append(PlanConfig(engine_id, knobs, mode=cudnn.heur_mode.A)) + return out + + +def _mode_fallback(facts, offered: Dict[str, int]) -> List[PlanConfig]: + """Configs expected to build where mode A's choice may not. + + TODO: today this is the smallest tile the row admits — the config that asks + least of the device, which is the one thing a fallback must be. Once a cell + has features its largest tiles cannot serve, this becomes the handful of + configs that between them cover the whole plane, chosen from measurements. + """ + out = [] + for engine_id, spec in _eligible(facts, offered): + caps = spec.capabilities + knobs = _knobs(caps, min(caps.tile_ms, default=None), min(caps.tile_ns, default=None)) + if _admissible(caps, facts, knobs): + out.append(PlanConfig(engine_id, knobs, mode=cudnn.heur_mode.FALLBACK)) + return out + + +def _leads(offered: Dict[str, int], plans: List[PlanConfig]) -> bool: + """Whether this family's mode-A plans outrank the backend's. See _MEASURED_BEHIND.""" + behind = {offered[name] for name in _MEASURED_BEHIND if name in offered} + return bool(plans) and not all(cfg.engine_id in behind for cfg in plans) + + +def recommend(modes: List[Any], facts, offered: Dict[str, int], backend_plans: List[PlanConfig]) -> List[PlanConfig]: + """The ranked plan list for this graph, mode by mode in the caller's order. + + Each mode contributes a block and the blocks concatenate, so asking for + [A, FALLBACK] puts every tuned candidate — both sides' — ahead of every + fallback. A plan repeated across modes keeps its first position: building + the same config twice only costs the caller a JIT compile. + """ + a_modes = (cudnn.heur_mode.A, cudnn.heur_mode.B) + # An untagged backend entry is the delegating one: candidates C++ holds but + # never exposes as plans, which Graph::build_plans tries BEFORE its own + # engine_configs. It belongs to no mode and must keep the lead, or an + # OPENSOURCE caller gets a native kernel instead of the OSS one. + out: List[PlanConfig] = [c for c in backend_plans if c.mode is None] + for mode in modes: + if mode == cudnn.heur_mode.OPENSOURCE: + # Mode A without the backend's recommendation: the caller asked for + # an open-source implementation and the backend's engines are not + # one. Nothing to place, so the measurements do not come into it. + out += _mode_a(facts, offered) + elif mode in a_modes: + # B asks for a wider search than A and this family has none to give, + # so it answers as A does. The backend answered B on its own terms. + ours = _mode_a(facts, offered) + theirs = [c for c in backend_plans if c.mode == mode] + out += (ours + theirs) if _leads(offered, ours) else (theirs + ours) + elif mode == cudnn.heur_mode.FALLBACK: + out += _mode_fallback(facts, offered) + [c for c in backend_plans if c.mode == mode] + + seen, ranked = set(), [] + for cfg in out: + key = (cfg.engine_id, repr(cfg.knobs), cfg.cpp_index) + if key not in seen: + seen.add(key) + ranked.append(cfg) + return ranked diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py b/python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py new file mode 100644 index 000000000..a78fd23e2 --- /dev/null +++ b/python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py @@ -0,0 +1,1569 @@ +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +""" +A fused multi-head attention (FMHA) per-tensor FP8 (e4m3) kernel for the NVIDIA +Blackwell SM120 family (SM120 and SM121), sibling of ``prefill_f16_sm120.py``. + +Same architecture as the f16 kernel — dedicated TMA load warp, GMEM-direct Q, +fp32 online softmax in registers, right-to-left masked KV walk — with the MMA +lowered to ``mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32``: + +- Q/K/V arrive PRE-QUANTIZED e4m3 and travel as **Uint8 storage**: the kernel + never does elementwise math on them, so bytes flow TMA -> ldmatrix -> MMA as + bit patterns and no Float8 element support is needed in the DSL plumbing. + ``descale_q * descale_k`` folds host-side into ``softmax_scale_log2`` and + ``descale_v * scale_o`` into the ``o_scale_fused`` scalar (cuDNN SDPA_FP8 + node convention; see the SM100 fp8 adapter). +- K B-fragments: classic byte-preserving ``ldmatrix.m8n8.x4.b16`` (sm_120a has + no non-transposed 8-bit ldmatrix; the b16 form gathers the right bytes when + each lane points at one 16-byte K-segment). +- V B-fragments: hardware 8-bit transposed ``ldmatrix.m16n16.x2.trans.b8`` + (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 + mismatch defeats the f16 kernel's in-register restage), reloaded as A + fragments with ``ldmatrix.m8n8.x4.b16``. P scale is fixed 1.0 (values below + 2^-9 flush to zero). The softmax denominator uses the fp32 P. +- O is Float16; the epilogue and the sKV/sO SMEM alias are sized in BYTES + because KV (1B) and O (2B) element sizes differ. +- ``Amax_S`` (max over valid rows of ``1/row_sum``) and ``Amax_O`` + (max ``|o_scaled|`` pre-cast) are produced via bitcast-int32 atomic max on + 1-element Int32 buffers the host pre-zeros on the launch stream; the host + divides Amax_O by ``scale_o`` afterwards (SM100 fp8 convention). + +Constraints: +* Input dtype: e4m3 only (as Uint8 storage); output dtype Float16 +* Head dimensions must be multiples of 16 between 16 and 256, inclusive +* Q heads must be divisible by the number of K/V heads +* Q/K/V/O use compact BSHD storage +* Supported CTA Q/KV tiles are 128 or 64 +* ``has_sink`` is not supported (Amax_S semantics with a sink column are + undefined in this cell; the adapter declines such graphs) +""" + +from functools import lru_cache, partial +from types import SimpleNamespace +from typing import Callable, Optional, Type + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.experimental.cuda as cuda +import cutlass.cute as cute + +from cutlass.experimental import primitives as prims +from cudnn.frost.tile_dsl.constants import DTYPE_E4M3 +from cudnn.frost.tile_dsl.mma import pack_f8x2_pairs, ptx_cvt_e4m3x2, ptx_mma_m16n8k32_e4m3_f32 +from cudnn.frost.tile_dsl.swizzle import swizzle_xor +from cudnn.sdpa.fwd.config_sm120 import ( + SEQ_KV_TILES as _SEQ_KV_TILES, + SEQ_Q_TILES as _SEQ_Q_TILES, + SUPPORTED_HEAD_TILES as _SUPPORTED_HEAD_TILES, + TemplateParams, + validate_params, +) + +# The FROST loader injects one immutable specialization before executing this +# module. A direct import uses dense e4m3 defaults. +PARAMS: TemplateParams = globals().get("FROST_TEMPLATE_PARAMS", TemplateParams(dtype_qkv=DTYPE_E4M3)) +validate_params(PARAMS, allowed_dtypes=(DTYPE_E4M3,)) + +# e4m3 travels as raw bytes: TMA, ldmatrix, and the MMA consume bit patterns, +# so Uint8 storage sidesteps Float8 element support in the DSL plumbing. +STORAGE_DTYPE = cutlass.Uint8 +OUT_STORAGE_DTYPE = cutlass.Float16 + +# --------------------------------------------------------------------------- +# PTX and layout helpers. +# --------------------------------------------------------------------------- + + +@cute.jit +def nvvm_threadquad_reduction_max(val: cutlass.Float32) -> cutlass.Float32: + """Butterfly thread-quad (4 lanes) reduction max via shfl.sync.bfly.""" + val = cute.arch.fmax( + val, + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=val, + offset=2, + mask_and_clamp=0x1F, + kind=prims.Shfl.BFLY, + ), + ) + val = cute.arch.fmax( + val, + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=val, + offset=1, + mask_and_clamp=0x1F, + kind=prims.Shfl.BFLY, + ), + ) + return val + + +@cute.jit +def nvvm_threadquad_reduction_sum(val: cutlass.Float32) -> cutlass.Float32: + """Butterfly thread-quad (4 lanes) reduction sum via shfl.sync.bfly.""" + val = val + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=val, + offset=2, + mask_and_clamp=0x1F, + kind=prims.Shfl.BFLY, + ) + val = val + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=val, + offset=1, + mask_and_clamp=0x1F, + kind=prims.Shfl.BFLY, + ) + return val + + +@cute.jit +def pack_to_i32( + src: tuple, + dtype: cutlass.Constexpr[Type[cutlass.Numeric]], +) -> cutlass.Int32: + """Pack four 8-bit or two 16-bit values into one 32-bit register.""" + vals = cutlass.Vector.from_elements(src, dtype) + return vals.bitcast(cutlass.Int32)[0] + + +def ceil_div(a: int, b: int) -> int: + """Return the ceiling division of a by b.""" + return (a + b - 1) // b + + +fmul2 = partial(prims.mul_packed_f32x2, ftz=False, rnd=prims.FPRoundingMode.RN) +fma2 = partial(prims.fma_packed_f32x2, ftz=False, rnd=prims.FPRoundingMode.RN) + + +# --------------------------------------------------------------------------- +# Main kernel class +# --------------------------------------------------------------------------- + + +class SM120FusedMultiHeadAttentionForward: + """Configure and launch the SM120/SM121 per-tensor FP8 FMHA prefill kernel.""" + + SEQ_Q_TILES = _SEQ_Q_TILES + SEQ_KV_TILES = _SEQ_KV_TILES + SUPPORTED_HEAD_TILES = _SUPPORTED_HEAD_TILES + MMA_TILER = (16, 8, 32) # mma.sync.aligned.m16n8k32 (e4m3) + + @staticmethod + def is_layout_supported( + shape: tuple[int, ...], + stride: tuple[int, ...], + ) -> bool: + """Return whether a BSHD tensor uses compact storage.""" + + if len(shape) != 4 or len(stride) != 4: + return False + _, sequence, heads, head_dim = shape + return stride == ( + sequence * heads * head_dim, + heads * head_dim, + head_dim, + 1, + ) + + def __init__( + self, + in_dtype: Type[cutlass.Numeric] = cutlass.Uint8, + out_dtype: Type[cutlass.Numeric] = cutlass.Float16, + is_causal: bool = False, + causal_bottom_right: bool = False, + window_size_left: int | None = None, + seq_q_lens_present: bool = False, + seq_kv_lens_present: bool = False, + has_sink: bool = False, + thd_varlen: bool = False, + thd_batch: int = 1, + thd_max_sq: int = 0, + head_tile_qk: int = 128, + head_tile_v: int = 128, + kv_tile: int = SEQ_KV_TILES[0], + q_tile: int = SEQ_Q_TILES[0], + ): + """Initialize the FMHA prefill kernel configuration. + + :param in_dtype: Q/K/V element type (Float16 or BFloat16). + :param out_dtype: O element type. Must match ``in_dtype``. + :param is_causal: Apply an upper causal bound to QK. + :param causal_bottom_right: Shift the causal diagonal by ``Skv - Sq``. + :param window_size_left: Inclusive left-window offset, or ``None``. + :param seq_q_lens_present: Read per-batch query lengths at runtime. + :param seq_kv_lens_present: Read per-batch key/value lengths at runtime. + :param has_sink: Fold the per-Q-head sink logit from the ``sinks`` + tensor into the softmax denominator; when ``False`` the ``sinks`` + argument is an unused dummy. + :param thd_varlen: THD (ragged) mode — Q/K/V/O and LSE are fully + packed batch-1 views, ``seq_kv_lens`` is the + ``[seq_kv(B) | cu_q(B+1) | cu_k(B+1)]`` metadata tensor, and the + grid covers ``ceil(thd_max_sq / q_tile)`` tiles per sequence. + :param thd_batch: THD only: the real sequence count B. + :param thd_max_sq: THD only: the longest sequence's Q length. + :param head_tile_qk: Q/K head dimension (the QK^T contraction width). + Must be a multiple of 16 between 16 and 256, inclusive. + :param head_tile_v: V/O head dimension (the P@V output width). Same + constraint as ``head_tile_qk``. + :param q_tile: Query sequence tile size. + :param kv_tile: Key/value sequence tile size. + """ + + if in_dtype != cutlass.Uint8: + raise ValueError("fp8 kernel takes e4m3 Q/K/V as Uint8 storage") + if out_dtype != cutlass.Float16: + raise ValueError("fp8 kernel emits Float16 O only") + if has_sink: + raise ValueError("has_sink is not supported by the fp8 cell (Amax_S semantics)") + if thd_varlen and (thd_batch < 1 or thd_max_sq < 1): + raise ValueError("thd_varlen requires thd_batch >= 1 and thd_max_sq >= 1") + self.in_dtype = in_dtype + self.out_dtype = out_dtype + self.is_causal = is_causal + self.causal_bottom_right = causal_bottom_right + self.window_size_left = window_size_left + self.seq_q_lens_present = seq_q_lens_present + self.seq_kv_lens_present = seq_kv_lens_present + self.has_sink = has_sink + self.thd_varlen = thd_varlen + self.thd_batch = thd_batch + self.thd_max_sq = thd_max_sq + + self.head_tile_qk = head_tile_qk + self.head_tile_v = head_tile_v + self.q_tile = q_tile + self.kv_tile = kv_tile + + # Warp roles + if self.q_tile == 128: + self.compute_warp_ids = (0, 1, 2, 3, 4, 5, 6, 7) + self.load_warp_id = 8 + self.empty_warp_ids = (9, 10, 11) + self.num_warps = 12 + else: + self.compute_warp_ids = (0, 1, 2, 3) + self.load_warp_id = 4 + self.empty_warp_ids = (5, 6, 7) + self.num_warps = 8 + self.num_compute_warps = len(self.compute_warp_ids) + self.num_load_warps = 1 + + self.bar_compute_sync = 1 + self.bar_k_consumed = 2 + self.bar_v_consumed = 3 + + self.threads_per_cta = cute.arch.WARP_SIZE * self.num_warps + self.threads_load = cute.arch.WARP_SIZE * self.num_load_warps + self.threads_compute = cute.arch.WARP_SIZE * self.num_compute_warps + self.threads_kv_pipeline = self.threads_compute + self.threads_load + + self._setup_attributes() + + def _setup_attributes(self): + """Compute derived tile, MMA, and TMA constants from the configuration.""" + + # Tiling + self.k_tile_elems = self.kv_tile * self.head_tile_qk + self.v_tile_elems = self.kv_tile * self.head_tile_v + self.o_tile_elems = self.q_tile * self.head_tile_v + + # MMA + self.qk_k_frags = self.kv_tile // self.MMA_TILER[1] + self.qk_d_frags = self.head_tile_qk // self.MMA_TILER[2] + self.pv_v_frags = self.kv_tile // self.MMA_TILER[2] + self.pv_d_frags = self.head_tile_v // self.MMA_TILER[1] + + # TMA + def get_swizzle(head_tile: int): + head_bytes = head_tile * self.in_dtype.bytes + for swizzle, span in ( + (cuda.TensorMapSwizzle.s128b, 128), + (cuda.TensorMapSwizzle.s64b, 64), + (cuda.TensorMapSwizzle.s32b, 32), + ): + if head_bytes % span == 0: + return swizzle, head_bytes // span, head_tile // (head_bytes // span) + raise ValueError(f"Unsupported TMA inner dimension: {head_bytes} B") + + self.k_tma_swizzle, self.k_tma_swizzle_chunks, self.k_swizzle_chunk_elems = get_swizzle(self.head_tile_qk) + self.v_tma_swizzle, self.v_tma_swizzle_chunks, self.v_swizzle_chunk_elems = get_swizzle(self.head_tile_v) + + @cute.jit + def load_one_kv_tile( + self, + s_dst: cutlass.Array, + tma_desc: cutlass.GridConstant[cuda.TensorMap], + mbar: cutlass.Array, + batch_idx: cutlass.Int32, + head_idx: cutlass.Int32, + seq_coord: cutlass.Int32, + ) -> None: + """Launch one TMA load for a complete K/V tile into swizzled SMEM. + + The tensor map exposes compact ``(B, S, H, D)`` storage through a + logical ``(B, H, I, S, C)`` view, where ``D = I * C``. Its TMA-order + dimensions are ``(C, S, I, H, B)``, so one rank-5 copy covers every + head chunk and uses coordinates ``(c, seq, i, head, batch)``. + + :param s_dst: Swizzled SMEM destination tile. + :param tma_desc: K or V tensor map descriptor. + :param mbar: TMA completion mbarrier for this stream. + :param batch_idx: Batch index. + :param head_idx: Attention head index. + :param seq_coord: Starting sequence row for the K/V tile. + """ + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx(mbar, tma_desc.global_tx_bytes()) + prims.cp_async_bulk_tensor_shared_cta_global( + s_dst, + tma_desc.get_ptr(), + (0, seq_coord, 0, head_idx, batch_idx), + mbar, + ) + + @cute.jit + def load_q_tile( + self, + basic_params: SimpleNamespace, + ) -> cutlass.Array: + """Load the warp-owned Q tile directly from GMEM into MMA A registers. + + :param basic_params: Per-CTA tensor metadata, lane mapping, and Q base + offsets. + :return: Packed Q fragments arranged for ``mma.sync`` A operands. + """ + q_regs = cutlass.Array( + cutlass.Int32, + self.qk_d_frags * 4, + alignment=16, + ) + + # First row and column owned by this lane in each MMA A fragment. + # m16n8k32 e4m3 A layout: a0 = A[r0, 4q..4q+3], a1 = A[r0+8, same], + # a2/a3 = +16 in k; each reg is 4 bytes packed little-endian, loaded + # as one aligned 4-byte GMEM access. + row0 = basic_params.lane // 4 + col0 = (basic_params.lane % 4) * 4 + + row0_in_cta = basic_params.q_warp_row0 + row0 + col0_in_cta = col0 + q_regs_offset = 0 + for _ in cutlass.range_constexpr(self.qk_d_frags): + mma_offsets_in_cta = ( + (row0_in_cta, col0_in_cta), + (row0_in_cta + 8, col0_in_cta), + (row0_in_cta, col0_in_cta + 16), + (row0_in_cta + 8, col0_in_cta + 16), + ) + for i in cutlass.range_constexpr(4): + row_in_cta, col_in_cta = mma_offsets_in_cta[i] + cur_q_seq_idx = basic_params.q_seq_idx + row_in_cta + q_packed = cutlass.Int32(0) + if cur_q_seq_idx < basic_params.seqlen_q and col_in_cta < basic_params.head_dim_qk: + q_quad = (basic_params.q_ptr + basic_params.q_head_off + cur_q_seq_idx * basic_params.q_seq_stride + col_in_cta).load(count=4, alignment=4) + q_packed = q_quad.bitcast(cutlass.Int32)[0] + q_regs[q_regs_offset + i] = q_packed + + col0_in_cta += self.MMA_TILER[2] + q_regs_offset += 4 + + return q_regs + + @cute.jit + def mma_qk( + self, + basic_params: SimpleNamespace, + mma_params: SimpleNamespace, + q_regs: cutlass.Array, + ): + """Compute ``S = Q @ K.T``. + + Q fragments are supplied in registers by ``load_q_tile``. K fragments + are read from the TMA-populated ``sK`` tile with ``ldmatrix``. + + :param basic_params: Per-CTA tensor metadata and lane mapping. + :param mma_params: Shared K tile and local O accumulator state. + :param q_regs: Register-resident packed Q fragments. + :return: Register-resident QK score fragments. + """ + s_regs = cutlass.Array( + cutlass.Float32, + self.qk_k_frags * 4, + alignment=16, + ) + for i in cutlass.range_constexpr(self.qk_k_frags * 4): + s_regs[i] = cutlass.Float32(0.0) + + # 8-bit K path: byte-preserving ldmatrix.m8n8.x4.b16 — each lane points + # at one 16-byte K-segment; tile pairs cover (n8 kv rows) x (k16-half) + # of one k32 d_frag, so k_vec[0],[1] form the m16n8k32 B fragment of + # the first n8 block and k_vec[2],[3] the second. + k_row_in_frag_pair = basic_params.lane_div16 * 8 + basic_params.lane_mod8 # which half of k-frag pair # which row in half + k_col_in_frag_pair = (basic_params.lane_div8 % 2) * 16 # which k16-half of the k32 d-frag + + def load_k_frag_pair(k_frag_pair: cutlass.Constexpr[int], d_frag: cutlass.Constexpr[int]): + k_row_in_cta = k_frag_pair * 16 + k_row_in_frag_pair + k_col_in_cta = d_frag * self.MMA_TILER[2] + k_col_in_frag_pair + k_chunk = k_col_in_cta // self.k_swizzle_chunk_elems + k_col_in_chunk = k_col_in_cta % self.k_swizzle_chunk_elems + k_physical_row = k_chunk * self.kv_tile + k_row_in_cta + k_smem_ptr = ( + mma_params.sK.data_ptr() + + k_physical_row * self.k_swizzle_chunk_elems + + swizzle_xor( + k_physical_row, + k_col_in_chunk, + self.k_swizzle_chunk_elems, + self.in_dtype.bytes, + ) + ) + return prims.ldmatrix(k_smem_ptr, 4, prims.MMALayout.ROW) + + for k_frag_pair in cutlass.range_constexpr(self.qk_k_frags // 2): + for d_frag in cutlass.range_constexpr(self.qk_d_frags): + k_vec = load_k_frag_pair(k_frag_pair, d_frag) + q_off = d_frag * 4 + s_off = (k_frag_pair * 2) * 4 + s_regs[s_off:4] = ptx_mma_m16n8k32_e4m3_f32( + q_regs[q_off + 0], + q_regs[q_off + 1], + q_regs[q_off + 2], + q_regs[q_off + 3], + k_vec[0], + k_vec[1], + s_regs[s_off + 0], + s_regs[s_off + 1], + s_regs[s_off + 2], + s_regs[s_off + 3], + ) + s_regs[s_off + 4 : 4] = ptx_mma_m16n8k32_e4m3_f32( + q_regs[q_off + 0], + q_regs[q_off + 1], + q_regs[q_off + 2], + q_regs[q_off + 3], + k_vec[2], + k_vec[3], + s_regs[s_off + 4], + s_regs[s_off + 5], + s_regs[s_off + 6], + s_regs[s_off + 7], + ) + + return s_regs + + @cute.jit + def online_softmax( + self, + basic_params: SimpleNamespace, + mma_params: SimpleNamespace, + softmax_params: SimpleNamespace, + s_regs: cutlass.Array, + kv_seq_idx: cutlass.Int32, + in_mask_steps: cutlass.Constexpr[bool], + is_first_kv_tile: cutlass.Constexpr[bool], + ): + """Online softmax and stage packed P in registers for the PV MMA. + + :param basic_params: Per-CTA tensor metadata and lane mapping. + :param mma_params: Local output accumulator state to rescale. + :param softmax_params: Row max/sum state and log2 softmax scale. + :param s_regs: Register-resident QK score fragments from ``mma_qk``. + :param kv_seq_idx: Absolute K/V row offset for this tile. + :param in_mask_steps: Whether this iteration needs causal or tail predicates. + :param is_first_kv_tile: Whether this is the first processed K/V tile + for the current Q tile. + :return: packed e4m3 P fragments, indexed ``[k_frag * 2 + row_half]``. + """ + lane = basic_params.lane + o_regs = mma_params.o_regs + row_max = softmax_params.row_max + row_sum = softmax_params.row_sum + softmax_scale_log2 = softmax_params.softmax_scale_log2 + p_regs = cutlass.Array(cutlass.Uint16, self.qk_k_frags * 2) + + # Each lane owns four S registers split across two Q rows after Q@K^T. + for row_half in cutlass.range_constexpr(2): + s_reg_idx_lo = row_half * 2 + s_reg_idx_hi = row_half * 2 + 1 + + q_row_in_cta = basic_params.q_warp_row0 + (lane // 4) + row_half * 8 + + # Resolve mask bounds for this query row. ``valid_cols`` is an + # exclusive upper bound; ``first_valid_col`` is inclusive. + q_position = basic_params.q_seq_idx + q_row_in_cta + diagonal_offset = cutlass.Int32(0) + if cutlass.const_expr(self.causal_bottom_right): + diagonal_offset = basic_params.seqlen_k - basic_params.seqlen_q + diagonal_position = q_position + diagonal_offset + + valid_cols = basic_params.seqlen_k + if cutlass.const_expr(self.is_causal): + valid_cols = cute.math.max( + cutlass.Int32(0), + cute.math.min(diagonal_position + 1, basic_params.seqlen_k), + ) + + first_valid_col = cutlass.Int32(0) + if cutlass.const_expr(self.window_size_left is not None): + first_valid_col = cute.math.max( + cutlass.Int32(0), + diagonal_position - self.window_size_left, + ) + + # Reduce max across this lane's S values for the current Q row. + cur_max = cutlass.Float32(-cutlass.Float32.inf) + for k_frag in cutlass.range_constexpr(self.qk_k_frags): + s_off = k_frag * 4 + s0 = s_regs[s_off + s_reg_idx_lo] + s1 = s_regs[s_off + s_reg_idx_hi] + if cutlass.const_expr(in_mask_steps): + k_col0 = kv_seq_idx + k_frag * 8 + 2 * (lane % 4) + k_col1 = k_col0 + 1 + valid0 = k_col0 >= first_valid_col and k_col0 < valid_cols + valid1 = k_col1 >= first_valid_col and k_col1 < valid_cols + if not valid0: + s0 = -cutlass.Float32.inf + if not valid1: + s1 = -cutlass.Float32.inf + s_regs[s_off + s_reg_idx_lo] = s0 + s_regs[s_off + s_reg_idx_hi] = s1 + cur_max = cute.arch.fmax(cur_max, cute.arch.fmax(s0, s1)) + + # The four lanes that share one Q row reduce to the tile row max. + cur_max = nvvm_threadquad_reduction_max(cur_max) + + # Update row_max and compute the old-output correction factor. + old_scale = cutlass.Float32(1.0) + if cutlass.const_expr(is_first_kv_tile): + new_max = cur_max + else: + row_max_prev = row_max[row_half] + new_max = cute.arch.fmax(row_max_prev, cur_max) + need_correct = True + if cutlass.const_expr(in_mask_steps): + need_correct = new_max > -cutlass.Float32.inf + if need_correct: + old_scale = cute.math.exp2( + (row_max_prev - new_max) * softmax_scale_log2, + fastmath=True, + ) + row_max[row_half] = new_max + + # Compute P, accumulate the per-lane partial sum, and stage P. + exp_max = new_max + if cutlass.const_expr(in_mask_steps): + if exp_max == -cutlass.Float32.inf: + exp_max = cutlass.Float32(0.0) + neg_exp_max_scaled = -(exp_max * softmax_scale_log2) + tile_sum = cutlass.Float32(0.0) + for k_frag in cutlass.range_constexpr(self.qk_k_frags): + s_off = k_frag * 4 + s0 = s_regs[s_off + s_reg_idx_lo] + s1 = s_regs[s_off + s_reg_idx_hi] + in0, in1 = fma2( + (s0, s1), + (softmax_scale_log2, softmax_scale_log2), + (neg_exp_max_scaled, neg_exp_max_scaled), + ) + p0 = cute.math.exp2(in0, fastmath=True) + p1 = cute.math.exp2(in1, fastmath=True) + tile_sum = tile_sum + (p0 + p1) + # P stays in registers at the C-fragment coordinates; mma_pv + # redistributes it to the k32 A layout with shfl. row_sum above + # uses the fp32 P, keeping the denominator full precision. + p_regs[k_frag * 2 + row_half] = ptx_cvt_e4m3x2(p1, p0) + + # Reduce tile_sum across the four lanes that own one Q row. + tile_sum = nvvm_threadquad_reduction_sum(tile_sum) + + # Correct row_sum and rescale O when row_max changes. + if cutlass.const_expr(is_first_kv_tile): + row_sum[row_half] = tile_sum + else: + row_sum[row_half] = row_sum[row_half] * old_scale + tile_sum + + for d_frag in cutlass.range_constexpr(self.pv_d_frags): + o_off = d_frag * 4 + row_half * 2 + o_regs[o_off + 0], o_regs[o_off + 1] = fmul2( + (o_regs[o_off + 0], o_regs[o_off + 1]), + (old_scale, old_scale), + ) + + return p_regs + + @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``. + + :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``. + """ + o_regs = mma_params.o_regs + lane = basic_params.lane + + # The QK C-fragment gives each lane columns 2*(t%4)+{0,1} while the k32 + # A-fragment wants 4 consecutive bytes; the two differ only by an + # exchange inside each thread quad, so two shfl and one prmt replace + # the SMEM round trip (and the 16 KB it needed). + lane_mod4 = lane % 4 + src0 = (lane // 4) * 4 + (lane_mod4 % 2) * 2 + selector = cutlass.Int32(0x5410) if lane_mod4 < 2 else cutlass.Int32(0x7632) + + def pack_p_cols(k_frag0: cutlass.Constexpr[int], row_half: cutlass.Constexpr[int]) -> cutlass.Int32: + pairs = pack_f8x2_pairs(p_regs[k_frag0 * 2 + row_half], p_regs[(k_frag0 + 1) * 2 + row_half]) + lo = prims.shfl_sync(thread_mask=0xFFFFFFFF, val=pairs, offset=src0, mask_and_clamp=0x1F, kind=prims.Shfl.IDX) + hi = prims.shfl_sync(thread_mask=0xFFFFFFFF, val=pairs, offset=src0 + 1, mask_and_clamp=0x1F, kind=prims.Shfl.IDX) + return cute.arch.inline_ptx( + "prmt.b32 $0, $1, $2, $3;", + write_only_types=[cutlass.Int32], + read_only_args=[lo, hi, selector], + ) + + # V B-fragments use the hardware 8-bit transposed load + # ``ldmatrix.m16n16.x2.trans.b8`` (SASS LDSM.8.MT1616): every lane + # supplies the start of smem kv-row ``v_frag*32 + lane`` at one 16-byte + # d-chunk; one issue covers 32(kv) x 16(d) and feeds TWO MMAs with + # register map (0, 2, 1, 3). + def load_v_frags(v_frag: cutlass.Constexpr[int], d_frag_pair: cutlass.Constexpr[int]): + v_row_in_cta = v_frag * self.MMA_TILER[2] + lane + v_col_in_cta = d_frag_pair * 16 + v_chunk = v_col_in_cta // self.v_swizzle_chunk_elems + v_col_in_chunk = v_col_in_cta % self.v_swizzle_chunk_elems + v_physical_row = v_chunk * self.kv_tile + v_row_in_cta + sV_ptr = ( + mma_params.sV.data_ptr() + + v_physical_row * self.v_swizzle_chunk_elems + + swizzle_xor( + v_physical_row, + v_col_in_chunk, + self.v_swizzle_chunk_elems, + self.in_dtype.bytes, + ) + ) + return prims.ldmatrix( + sV_ptr, + 4, + prims.MMALayout.COL, + shape=prims.LoadShape.M16N16, + src_format=prims.LoadSrcFormat.B8, + ) + + # An ldmatrix is always in flight behind the tensor cores: the next + # v_frag's fragments are issued before this one's MMAs run. The index + # wraps so the prefetch is unconditional -- the last iteration reloads + # v_frag 0 and drops it, which is cheaper than a branch in the loop. + for v_frag in cutlass.range_constexpr(self.pv_v_frags): + # One k32 PV step consumes four QK k-fragments, paired (0,1) and (2,3). + p_vec = ( + pack_p_cols(v_frag * 4 + 0, 0), + pack_p_cols(v_frag * 4 + 0, 1), + pack_p_cols(v_frag * 4 + 2, 0), + pack_p_cols(v_frag * 4 + 2, 1), + ) + for d_frag_pair in cutlass.range_constexpr(self.pv_d_frags // 2): + v_vec = load_v_frags(v_frag, d_frag_pair) + o_off = (d_frag_pair * 2) * 4 + o_regs[o_off:4] = ptx_mma_m16n8k32_e4m3_f32( + p_vec[0], + p_vec[1], + p_vec[2], + p_vec[3], + v_vec[0], + v_vec[2], + o_regs[o_off + 0], + o_regs[o_off + 1], + o_regs[o_off + 2], + o_regs[o_off + 3], + ) + o_regs[o_off + 4 : 4] = ptx_mma_m16n8k32_e4m3_f32( + p_vec[0], + p_vec[1], + p_vec[2], + p_vec[3], + v_vec[1], + v_vec[3], + o_regs[o_off + 4], + o_regs[o_off + 5], + o_regs[o_off + 6], + o_regs[o_off + 7], + ) + + @cute.jit + def compute_one_kv_tile( + self, + basic_params: SimpleNamespace, + mma_params: SimpleNamespace, + softmax_params: SimpleNamespace, + q_regs: cutlass.Array, + num_kv_tiles: cutlass.Int32, + kv_tile_idx: cutlass.Int32, + in_mask_steps: cutlass.Constexpr[bool], + is_first_kv_tile: cutlass.Constexpr[bool], + ) -> None: + """One compute-side iteration of the right-to-left FMHA prefill loop. + + :param basic_params: Per-CTA tensor metadata, lane mapping, and TMA mbarriers. + :param mma_params: Shared-memory tiles and local MMA state. + :param softmax_params: Online softmax row state. + :param q_regs: Register-resident packed Q fragments. + :param num_kv_tiles: Number of K/V tiles processed by this CTA. + :param kv_tile_idx: K/V tile index processed by this iteration. + :param in_mask_steps: Whether this tile needs causal or K-tail masking. + :param is_first_kv_tile: Whether this tile initializes the online softmax state. + """ + + # The K/V loop walks tile indices in reverse order. The mbarrier parity + # still follows the load iteration count: 0, 1, 0, 1, ... + tma_phase = (num_kv_tiles - 1 - kv_tile_idx) & cutlass.Int32(1) + while not prims.mbarrier_try_wait_parity(basic_params.k_tma_mbar, tma_phase): + pass + + s_regs = self.mma_qk(basic_params, mma_params, q_regs) + prims.barrier_cta_arrive(self.bar_k_consumed, self.threads_kv_pipeline) + + while not prims.mbarrier_try_wait_parity(basic_params.v_tma_mbar, tma_phase): + pass + + p_regs = self.online_softmax( + basic_params, + mma_params, + softmax_params, + s_regs, + kv_tile_idx * self.kv_tile, + in_mask_steps, + is_first_kv_tile, + ) + + self.mma_pv(basic_params, mma_params, p_regs) + prims.barrier_cta_arrive(self.bar_v_consumed, self.threads_kv_pipeline) + + @cute.kernel + def kernel( + self, + q: cute.Tensor, + k: cute.Tensor, + v: cute.Tensor, + o: cute.Tensor, + lse: Optional[cute.Tensor], + sinks: Optional[cute.Tensor], + seq_q_lens: cute.Tensor, + seq_kv_lens: cute.Tensor, + amax_s: cute.Tensor, + amax_o: cute.Tensor, + tma_k_desc: cutlass.GridConstant[cuda.TensorMap], + tma_v_desc: cutlass.GridConstant[cuda.TensorMap], + softmax_scale_log2: cutlass.Float32, + o_scale_fused: cutlass.Float32, + ) -> None: + """SM120 per-tensor FP8 FMHA prefill kernel. + + :param q: Query tensor (e4m3 as Uint8 storage). + :param k: Key tensor (e4m3 as Uint8 storage). + :param v: Value tensor (e4m3 as Uint8 storage). + :param o: Output tensor (Float16). + :param lse: ``(B, H, Sq)`` fp32 log-sum-exp output, or ``None`` to + compile the LSE store out (the DSL specializes on ``None``). + :param sinks: Unused; must be ``None`` (fp8 cell rejects has_sink). + :param seq_q_lens: Per-batch query lengths, or an unused dummy tensor. + :param seq_kv_lens: Per-batch key/value lengths, or an unused dummy tensor. + :param amax_s: 1-element Int32 buffer, host pre-zeroed on the launch + stream; receives bitcast-fp32 atomic max of ``1/row_sum`` over + valid rows (cuDNN Amax_S convention). Pass a dummy when unused. + :param amax_o: 1-element Int32 buffer, host pre-zeroed; receives + bitcast-fp32 atomic max of ``|o_scaled|`` pre-cast. The host + divides by ``scale_o`` afterwards. Pass a dummy when unused. + :param tma_k_desc: Tensor map descriptor for K. + :param tma_v_desc: Tensor map descriptor for V. + :param softmax_scale_log2: ``softmax_scale * descale_q * descale_k * + log2(e)``, pre-folded host-side. + :param o_scale_fused: ``descale_v * scale_o``, pre-folded host-side. + """ + tidx, _, _ = cute.arch.thread_idx() + q_tile_idx, batch_idx, head_idx = cute.arch.block_idx() + if cutlass.const_expr(self.is_causal): + # Causal work grows with the Q tile. Launch long tiles first to + # avoid leaving a few expensive CTAs in the final scheduler waves. + grid_q, _, _ = cute.arch.grid_dim() + q_tile_idx = grid_q - q_tile_idx - 1 + q_seq_idx = q_tile_idx * self.q_tile + + lane = tidx % cute.arch.WARP_SIZE + warp = cute.arch.warp_idx() + + seqlen_q = cutlass.Int32(q.shape[1]) + seqlen_k = cutlass.Int32(k.shape[1]) + q_row_base = cutlass.Int32(0) + kv_row_base = cutlass.Int32(0) + if cutlass.const_expr(self.thd_varlen): + # THD: seq_kv_lens is the metadata tensor [seq_kv(B) | cu_q(B+1) | cu_k(B+1)]. + # Per-sequence lengths come from the prefix sums; the bases offset + # every packed (1, T, H, D) access below. + n_batch = (seq_kv_lens.shape[0] - 2) // 3 + meta = cutlass.make_array_view(seq_kv_lens) + q_row_base = cutlass.Int32(meta[n_batch + batch_idx]) + seqlen_q = cutlass.Int32(meta[n_batch + batch_idx + 1]) - q_row_base + kv_row_base = cutlass.Int32(meta[2 * n_batch + 1 + batch_idx]) + seqlen_k = cutlass.Int32(meta[2 * n_batch + 1 + batch_idx + 1]) - kv_row_base + else: + if cutlass.const_expr(self.seq_q_lens_present): + seqlen_q = cute.math.max( + cutlass.Int32(0), + cute.math.min(seq_q_lens[batch_idx], cutlass.Int32(q.shape[1])), + ) + if cutlass.const_expr(self.seq_kv_lens_present): + seqlen_k = cute.math.max( + cutlass.Int32(0), + cute.math.min(seq_kv_lens[batch_idx], cutlass.Int32(k.shape[1])), + ) + + num_heads_q = q.shape[2] + num_heads_kv = k.shape[2] + head_dim_qk = q.shape[3] + head_dim_v = v.shape[3] + q_ptr = q.iterator.raw_ptr() + o_ptr = o.iterator.raw_ptr() + + q_batch_stride, q_seq_stride, q_head_stride, _ = q.stride + o_batch_stride, o_seq_stride, o_head_stride, _ = o.stride + if cutlass.const_expr(self.thd_varlen): + # Packed view has batch 1: the sequence's token base replaces the + # batch stride term, and every Q/O row index below stays + # sequence-local. + q_head_off = q_row_base * q_seq_stride + head_idx * q_head_stride + o_head_off = q_row_base * o_seq_stride + head_idx * o_head_stride + else: + q_head_off = batch_idx * q_batch_stride + head_idx * q_head_stride + o_head_off = batch_idx * o_batch_stride + head_idx * o_head_stride + kv_head_idx = head_idx // (num_heads_q // num_heads_kv) + + num_kv_tiles = ceil_div(seqlen_k, self.kv_tile) + if cutlass.const_expr(self.thd_varlen): + # The grid covers ceil(max_seq_q / q_tile) tiles per sequence; a + # tile past this sequence's Q length has no rows to produce. + # Zeroing its KV work makes the whole CTA drain through the + # barriers without loads, compute, or stores. + if q_seq_idx >= seqlen_q: + num_kv_tiles = cutlass.Int32(0) + if cutlass.const_expr(self.is_causal): + causal_k_end = q_seq_idx + self.q_tile + if cutlass.const_expr(self.causal_bottom_right): + causal_k_end += seqlen_k - seqlen_q + causal_k_end = cute.math.max(cutlass.Int32(0), cute.math.min(causal_k_end, seqlen_k)) + num_kv_tiles_causal = ceil_div(causal_k_end, self.kv_tile) + num_kv_tiles = cute.math.min(num_kv_tiles, num_kv_tiles_causal) + + min_kv_tile = cutlass.Int32(0) + if cutlass.const_expr(self.window_size_left is not None): + first_q_position = q_seq_idx + if cutlass.const_expr(self.causal_bottom_right): + first_q_position += seqlen_k - seqlen_q + first_valid_col = cute.math.max(cutlass.Int32(0), first_q_position - self.window_size_left) + min_kv_tile = first_valid_col // self.kv_tile + has_kv_work = num_kv_tiles > 0 and (num_kv_tiles - 1) >= min_kv_tile + + # Shared-memory layout (BYTE-sized: KV elements are 1 byte, O staging + # elements are 2 — an element-counted max() would under-allocate the + # sO alias): + # sK: one kv_tile x head_tile_qk K tile (e4m3 bytes) + # sV: one kv_tile x head_tile_v V tile (e4m3 bytes) + # The epilogue later aliases sKV as the q_tile x head_tile_v sO + # staging tile after compute warps finish consuming the final K/V tile. + k_tile_bytes = self.k_tile_elems * self.in_dtype.bytes + v_tile_bytes = self.v_tile_elems * self.in_dtype.bytes + o_stage_bytes = self.o_tile_elems * self.out_dtype.bytes + sKV = cutlass.Array( + k.dtype, + max(k_tile_bytes + v_tile_bytes, o_stage_bytes), + space=cutlass.AddressSpace.smem, + alignment=128, + ) + sK = sKV + sV = sKV.subview(k_tile_bytes) + tma_mbar = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + k_tma_mbar = tma_mbar + v_tma_mbar = tma_mbar.subview(1) + + # Initialize the TMA completion barriers before any load or compute warp + # can touch the K/V pipeline. + if warp == self.load_warp_id: + if prims.elect_sync(): + prims.prefetch_tensormap(tma_k_desc.get_ptr()) + prims.prefetch_tensormap(tma_v_desc.get_ptr()) + prims.mbarrier_init(k_tma_mbar, 1) + prims.mbarrier_init(v_tma_mbar, 1) + prims.fence_mbarrier_init() + prims.barrier_cta_sync(0) + + # ///////////////////////////////////////////////////////////////////////////// + # LOAD K/V + # ///////////////////////////////////////////////////////////////////////////// + if warp == self.load_warp_id: + prims.setmaxregister(40, prims.SetMaxRegisterAction.DECREASE) + + # THD collapses the packed view's batch coordinate to 0; the + # per-sequence token base rides the seq coordinate instead. Every + # K/V load (including the first) must apply both, or batch >= 1 + # reads the wrong packed rows. + tma_batch_idx = batch_idx + if cutlass.const_expr(self.thd_varlen): + tma_batch_idx = cutlass.Int32(0) + + # The attention loop walks K/V tiles right-to-left so causal and + # tail-masked tiles are processed before fully unmasked tiles. + kv_seq_idx = (num_kv_tiles - 1) * self.kv_tile + if has_kv_work: + self.load_one_kv_tile( + sK, + tma_k_desc, + k_tma_mbar, + tma_batch_idx, + kv_head_idx, + kv_row_base + kv_seq_idx, + ) + self.load_one_kv_tile( + sV, + tma_v_desc, + v_tma_mbar, + tma_batch_idx, + kv_head_idx, + kv_row_base + kv_seq_idx, + ) + + kv_seq_idx -= self.kv_tile + while kv_seq_idx >= min_kv_tile * self.kv_tile: + prims.barrier_cta_sync( + self.bar_k_consumed, + thread_count=self.threads_kv_pipeline, + ) + self.load_one_kv_tile( + sK, + tma_k_desc, + k_tma_mbar, + tma_batch_idx, + kv_head_idx, + kv_row_base + kv_seq_idx, + ) + + prims.barrier_cta_sync( + self.bar_v_consumed, + thread_count=self.threads_kv_pipeline, + ) + self.load_one_kv_tile( + sV, + tma_v_desc, + v_tma_mbar, + tma_batch_idx, + kv_head_idx, + kv_row_base + kv_seq_idx, + ) + kv_seq_idx -= self.kv_tile + # ///////////////////////////////////////////////////////////////////////////// + # COMPUTE + # ///////////////////////////////////////////////////////////////////////////// + elif warp < self.load_warp_id: + prims.setmaxregister(232, prims.SetMaxRegisterAction.INCREASE) + + compute_warp_idx = warp + q_warp_row0 = compute_warp_idx * self.MMA_TILER[0] + + lane_div8 = lane // 8 + lane_mod8 = lane % 8 + lane_div16 = lane // 16 + + # Per-lane row_max and row_sum for online softmax. Each lane owns + # two Q rows: lane//4 and lane//4 + 8 within this compute warp. + row_max = cutlass.Array(cutlass.Float32, 2, alignment=16) + row_sum = cutlass.Array(cutlass.Float32, 2, alignment=16) + for i in cutlass.range_constexpr(2): + row_max[i] = -cutlass.Float32.inf + row_sum[i] = 0.0 + + # Per-lane fp32 accumulator for O = P @ V. + o_regs = cutlass.Array( + cutlass.Float32, + self.pv_d_frags * 4, + alignment=16, + ) + for i in cutlass.range_constexpr(self.pv_d_frags * 4): + o_regs[i] = 0.0 + + basic_params = SimpleNamespace( + seqlen_q=seqlen_q, + seqlen_k=seqlen_k, + head_dim_qk=head_dim_qk, + q_ptr=q_ptr, + batch_idx=batch_idx, + head_idx=head_idx, + q_seq_idx=q_seq_idx, + q_head_off=q_head_off, + q_seq_stride=q_seq_stride, + q_warp_row0=q_warp_row0, + lane=lane, + lane_div8=lane_div8, + lane_mod8=lane_mod8, + lane_div16=lane_div16, + tma_k_desc=tma_k_desc, + tma_v_desc=tma_v_desc, + k_tma_mbar=k_tma_mbar, + v_tma_mbar=v_tma_mbar, + ) + mma_params = SimpleNamespace( + sK=sK, + sV=sV, + o_regs=o_regs, + ) + softmax_params = SimpleNamespace( + row_max=row_max, + row_sum=row_sum, + softmax_scale_log2=softmax_scale_log2, + ) + + # Load Q into registers. + q_regs = self.load_q_tile(basic_params) + + # Main attention loop. + mask_steps = 1 + if cutlass.const_expr(self.is_causal): + mask_steps = ceil_div(self.q_tile, self.kv_tile) + if cutlass.const_expr(self.causal_bottom_right): + # The shifted diagonal can straddle one additional KV tile. + mask_steps = ceil_div(self.q_tile + self.kv_tile - 1, self.kv_tile) + left_mask_steps = 1 + if cutlass.const_expr(self.window_size_left is not None): + left_mask_steps = ceil_div(self.q_tile + self.kv_tile - 1, self.kv_tile) + + kv_tile_idx = num_kv_tiles - 1 + # Phase 1: potentially masked iterations. + for step in cutlass.range_constexpr(mask_steps): + if kv_tile_idx >= min_kv_tile: + self.compute_one_kv_tile( + basic_params, + mma_params, + softmax_params, + q_regs, + num_kv_tiles, + kv_tile_idx, + in_mask_steps=True, + is_first_kv_tile=(step == 0), + ) + kv_tile_idx -= 1 + + # Phase 2: remaining fully unmasked iterations. + while kv_tile_idx > min_kv_tile + (left_mask_steps - 1): + self.compute_one_kv_tile( + basic_params, + mma_params, + softmax_params, + q_regs, + num_kv_tiles, + kv_tile_idx, + in_mask_steps=False, + is_first_kv_tile=False, + ) + kv_tile_idx -= 1 + + # The sliding-window left edge sweeps across the Q tile and can + # therefore cut through more than one K/V tile. + if cutlass.const_expr(self.window_size_left is not None): + for _ in cutlass.range_constexpr(left_mask_steps): + if kv_tile_idx >= min_kv_tile: + self.compute_one_kv_tile( + basic_params, + mma_params, + softmax_params, + q_regs, + num_kv_tiles, + kv_tile_idx, + in_mask_steps=True, + is_first_kv_tile=False, + ) + kv_tile_idx -= 1 + else: + if kv_tile_idx >= 0: + self.compute_one_kv_tile( + basic_params, + mma_params, + softmax_params, + q_regs, + num_kv_tiles, + kv_tile_idx, + in_mask_steps=False, + is_first_kv_tile=False, + ) + + # Per-row O normalization factor and natural-log LSE. + # The thread-quad reductions left row_max / row_sum replicated across + # the four lanes that share a Q row, so every lane finalizes the two + # rows it owns without further exchange. row_max holds the raw (unscaled) + # score max; the scale is applied in log2 domain and converted with ln(2). + # With has_sink, the per-head sink logit joins the softmax denominator + # as a virtual column with no V row: it rescales O, enters the LSE, + # and gives a row with no visible key a finite LSE (the sink alone). + LN2 = cutlass.Float32(0.6931471805599453) + row_sum_inv = cutlass.Array(cutlass.Float32, 2, alignment=8) + row_lse = cutlass.Array(cutlass.Float32, 2, alignment=8) + for row_half in cutlass.range_constexpr(2): + row_max_nat = row_max[row_half] * softmax_scale_log2 * LN2 + if cutlass.const_expr(self.has_sink): + sinks_arr = cutlass.make_array_view(sinks) + sink_logit = cutlass.Float32(sinks_arr[head_idx]) + new_max = cute.arch.fmax(row_max_nat, sink_logit) + # alpha re-normalizes the loop's accumulator and sum from + # row_max_nat to the sink-extended max; it is 0 for a row + # with no visible key, so O := 0 falls out. + alpha = cute.math.exp(row_max_nat - new_max, fastmath=True) + new_sum = row_sum[row_half] * alpha + cute.math.exp(sink_logit - new_max, fastmath=True) + row_sum_inv[row_half] = alpha / new_sum + row_lse[row_half] = new_max + cute.math.log(new_sum, fastmath=True) + else: + inv = cutlass.Float32(0.0) + if row_sum[row_half] > 0.0: + inv = cute.math.rcp(row_sum[row_half], approx=True, ftz=True) + row_sum_inv[row_half] = inv + lse_val = row_max_nat + cute.math.log( + cute.math.max(row_sum[row_half], cutlass.Float32(1e-30)), + fastmath=True, + ) + # Rows with no visible key write -inf / O := 0. + if row_sum[row_half] <= 0.0: + lse_val = -cutlass.Float32.inf + row_lse[row_half] = lse_val + + # Amax_S = max over valid rows of the raw 1/row_sum (cuDNN + # convention: the softmax-probability amax proxy), captured BEFORE + # o_scale_fused folds into the normalization factor. One lane per + # thread-quad carries the replicated row state; bitcast-int32 + # atomic max is exact for non-negative fp32. + amax_s_arr = cutlass.make_array_view(amax_s) + if lane % 4 == 0: + for row_half in cutlass.range_constexpr(2): + amax_q_idx = q_seq_idx + q_warp_row0 + (lane // 4) + row_half * 8 + if amax_q_idx < seqlen_q: + prims.atomicrmw( + prims.AtomicOp.MAX, + amax_s_arr, + row_sum_inv[row_half].bitcast(cutlass.Int32), + ) + for row_half in cutlass.range_constexpr(2): + row_sum_inv[row_half] = row_sum_inv[row_half] * o_scale_fused + + if cutlass.const_expr(lse is not None): + if lane % 4 == 0: + lse_arr = cutlass.make_array_view(lse) + for row_half in cutlass.range_constexpr(2): + lse_q_idx = q_seq_idx + q_warp_row0 + (lane // 4) + row_half * 8 + lse_out = cutlass.Float32(row_lse[row_half]) + if cutlass.const_expr(self.thd_varlen): + # Packed (1, H, T) LSE: rows past this sequence's Q + # length belong to the NEXT sequence — never written, + # and there is no padded region to trim. + if lse_q_idx < seqlen_q: + lse_row = lse_arr[0, head_idx, :] + lse_row[q_row_base + lse_q_idx] = lse_out + else: + # Rows at/past this batch's Q length trim to -inf. + if lse_q_idx >= seqlen_q: + lse_out = -cutlass.Float32.inf + if lse_q_idx < q.shape[1]: + lse_row = lse_arr[batch_idx, head_idx, :] + lse_row[lse_q_idx] = lse_out + + prims.barrier_cta_sync(self.bar_compute_sync, thread_count=self.threads_compute) + + # Epilogue: normalize O, stage it through an stmatrix-friendly SMEM + # layout, then store one contiguous 8-element vector per lane to GMEM. + sO = sKV + row_sum_inv_vec = cutlass.Vector.from_elements( + ( + row_sum_inv[0], + row_sum_inv[0], + row_sum_inv[1], + row_sum_inv[1], + row_sum_inv[0], + row_sum_inv[0], + row_sum_inv[1], + row_sum_inv[1], + ), + cutlass.Float32, + ) + # sO is a BYTE array (it aliases 1-byte KV storage); each + # (warp, d_frag_pair) block stages 16 x 16 Float16 = 512 bytes, + # one 16-byte stmatrix row per lane. Row validity for Amax_O + # mirrors the store-time trim: rows at/past seqlen_q store zeros + # and must not contribute. + o_block_bytes = 16 * 16 * self.out_dtype.bytes + row_valid = cutlass.Array(cutlass.Float32, 2, alignment=8) + for row_half in cutlass.range_constexpr(2): + amax_q_idx = q_seq_idx + q_warp_row0 + (lane // 4) + row_half * 8 + row_valid[row_half] = cutlass.Float32(1.0) if amax_q_idx < seqlen_q else cutlass.Float32(0.0) + lane_amax_o = cutlass.Float32(0.0) + for d_frag_pair in cutlass.range_constexpr(self.pv_d_frags // 2): + o_off = (d_frag_pair * 2) * 4 + o_scaled = fmul2(o_regs[o_off:8], row_sum_inv_vec) + for i in cutlass.range_constexpr(8): + # row_sum_inv_vec order: halves alternate 0,0,1,1,0,0,1,1. + half = (i // 2) % 2 + lane_amax_o = cute.arch.fmax(lane_amax_o, cute.math.abs(o_scaled[i]) * row_valid[half]) + o_packed = o_scaled.to(self.out_dtype).bitcast(cutlass.Int32) + sO_ptr = sO.data_ptr() + (compute_warp_idx * (self.pv_d_frags // 2) + d_frag_pair) * o_block_bytes + lane * 16 + prims.stmatrix( + sO_ptr, + o_packed, + prims.MMALayout.ROW, + ) + amax_o_arr = cutlass.make_array_view(amax_o) + prims.atomicrmw( + prims.AtomicOp.MAX, + amax_o_arr, + lane_amax_o.bitcast(cutlass.Int32), + ) + + store_row = lane_mod8 + ((lane_div8) % 2) * 8 + store_col = lane_div16 * 8 + store_q_seq_idx = q_seq_idx + q_warp_row0 + store_row + for d_frag_pair in cutlass.range_constexpr(self.pv_d_frags // 2): + store_col_in_cta = d_frag_pair * 16 + store_col + if cutlass.const_expr(self.thd_varlen): + # Packed storage: rows past this sequence's Q length are + # the NEXT sequence's tokens — no store, and never the + # dense path's zero-fill. + if store_q_seq_idx < seqlen_q and store_col_in_cta < head_dim_v: + gO_ptr = o_ptr + o_head_off + store_q_seq_idx * o_seq_stride + store_col_in_cta + sO_ptr = sO.data_ptr() + (compute_warp_idx * (self.pv_d_frags // 2) + d_frag_pair) * o_block_bytes + lane * 16 + gO_ptr.store(sO_ptr.load(count=16, alignment=16).bitcast(self.out_dtype), alignment=16) + else: + if store_q_seq_idx < q.shape[1] and store_col_in_cta < head_dim_v: + gO_ptr = o_ptr + o_head_off + store_q_seq_idx * o_seq_stride + store_col_in_cta + if store_q_seq_idx < seqlen_q: + sO_ptr = sO.data_ptr() + (compute_warp_idx * (self.pv_d_frags // 2) + d_frag_pair) * o_block_bytes + lane * 16 + gO_ptr.store(sO_ptr.load(count=16, alignment=16).bitcast(self.out_dtype), alignment=16) + else: + zero_vec = cutlass.Vector.from_elements( + ( + o.dtype(0.0), + o.dtype(0.0), + o.dtype(0.0), + o.dtype(0.0), + o.dtype(0.0), + o.dtype(0.0), + o.dtype(0.0), + o.dtype(0.0), + ), + o.dtype, + ) + gO_ptr.store(zero_vec, alignment=16) + + # ///////////////////////////////////////////////////////////////////////////// + # EMPTY + # ///////////////////////////////////////////////////////////////////////////// + else: + prims.setmaxregister(40, prims.SetMaxRegisterAction.DECREASE) + + @cute.jit + def __call__( + self, + q: cute.Tensor, + k: cute.Tensor, + v: cute.Tensor, + o: cute.Tensor, + lse: Optional[cute.Tensor], + sinks: Optional[cute.Tensor], + seq_q_lens: cute.Tensor, + seq_kv_lens: cute.Tensor, + amax_s: cute.Tensor, + amax_o: cute.Tensor, + softmax_scale_log2: cutlass.Float32, + o_scale_fused: cutlass.Float32, + stream: cuda_driver.CUstream, + ) -> None: + """Launch the SM120 per-tensor FP8 FMHA kernel. + + :param q: Query tensor with shape ``(B, Sq, H, D)`` (e4m3 as Uint8). + :param k: Key tensor with shape ``(B, Sk, H, D)`` (e4m3 as Uint8). + :param v: Value tensor with shape ``(B, Sk, H, D)`` (e4m3 as Uint8). + :param o: Output tensor with shape ``(B, Sq, H, D)`` (Float16). + :param lse: ``(B, H, Sq)`` fp32 log-sum-exp output, or ``None`` to + compile the LSE store out entirely (no dummy buffer needed). + :param sinks: Must be ``None`` (fp8 cell rejects has_sink). + :param seq_q_lens: Per-batch query lengths, or an unused dummy tensor. + :param seq_kv_lens: Per-batch key/value lengths, or an unused dummy tensor. + :param amax_s: 1-element Int32 amax buffer (pre-zeroed; dummy ok). + :param amax_o: 1-element Int32 amax buffer (pre-zeroed; dummy ok). + :param softmax_scale_log2: ``softmax_scale * descale_q * descale_k * log2(e)``. + :param o_scale_fused: ``descale_v * scale_o``. + :param stream: CUDA stream used for the launch. + """ + head_dim_qk = q.shape[3] + head_dim_v = v.shape[3] + if cutlass.const_expr(head_dim_qk != k.shape[3] or head_dim_qk != self.head_tile_qk): + raise ValueError("runtime Q/K head dimensions must match the kernel head_tile_qk") + if cutlass.const_expr(head_dim_v != o.shape[3] or head_dim_v != self.head_tile_v): + raise ValueError("runtime V/O head dimensions must match the kernel head_tile_v") + if cutlass.const_expr( + q.shape[0] != k.shape[0] + or k.shape[:3] != v.shape[:3] + or q.shape[0] != o.shape[0] + or q.shape[1] != o.shape[1] + or q.shape[2] != o.shape[2] + or q.shape[2] % k.shape[2] != 0 + ): + raise ValueError("runtime Q/K/V/O batch, sequence, or head geometry mismatch") + for name, tensor in (("Q", q), ("K", k), ("V", v), ("O", o)): + if cutlass.const_expr(not self.is_layout_supported(tensor.shape, tensor.stride)): + raise ValueError(f"{name} must use compact BSHD storage") + if cutlass.const_expr(lse is not None): + if cutlass.const_expr(lse.shape != (q.shape[0], q.shape[2], q.shape[1])): + raise ValueError("LSE must have shape (B, H, Sq)") + if cutlass.const_expr(lse.stride != (q.shape[2] * q.shape[1], q.shape[1], 1)): + raise ValueError("LSE must be compact row-major") + if cutlass.const_expr(self.has_sink != (sinks is not None)): + raise ValueError("sinks must be provided exactly when the kernel is configured with has_sink") + if cutlass.const_expr(sinks is not None and sinks.shape != (q.shape[2],)): + raise ValueError("sinks must have shape (H,)") + if cutlass.const_expr(self.thd_varlen): + if cutlass.const_expr(q.shape[0] != 1): + raise ValueError("THD Q/K/V/O must be packed batch-1 views") + if cutlass.const_expr(seq_kv_lens.shape != (3 * self.thd_batch + 2,)): + raise ValueError("THD seq_kv_lens must be the (3*B+2,) metadata tensor") + + # Split D into I contiguous C-element chunks while preserving the + # compact (B, S, H, D) global-memory address calculation. TMA order + # (C, S, I, H, B) linearizes the SMEM destination as [I][kv_tile][C]. + k_tma_layout = cute.make_layout( + ( + k.shape[0], + k.shape[2], + self.k_tma_swizzle_chunks, + k.shape[1], + self.k_swizzle_chunk_elems, + ), + stride=( + k.shape[1] * k.shape[2] * head_dim_qk, + head_dim_qk, + self.k_swizzle_chunk_elems, + k.shape[2] * head_dim_qk, + 1, + ), + ) + k_tma_box = ( + 1, + 1, + self.k_tma_swizzle_chunks, + self.kv_tile, + self.k_swizzle_chunk_elems, + ) + tma_k_desc = cuda.create_tensor_map_tiled_from_view( + cute.make_tensor(k.iterator, k_tma_layout), + box_dims=k_tma_box, + stride_order=(4, 3, 2, 1, 0), + swizzle=self.k_tma_swizzle, + ) + v_tma_layout = cute.make_layout( + ( + v.shape[0], + v.shape[2], + self.v_tma_swizzle_chunks, + v.shape[1], + self.v_swizzle_chunk_elems, + ), + stride=( + v.shape[1] * v.shape[2] * head_dim_v, + head_dim_v, + self.v_swizzle_chunk_elems, + v.shape[2] * head_dim_v, + 1, + ), + ) + v_tma_box = ( + 1, + 1, + self.v_tma_swizzle_chunks, + self.kv_tile, + self.v_swizzle_chunk_elems, + ) + tma_v_desc = cuda.create_tensor_map_tiled_from_view( + cute.make_tensor(v.iterator, v_tma_layout), + box_dims=v_tma_box, + stride_order=(4, 3, 2, 1, 0), + swizzle=self.v_tma_swizzle, + ) + self.kernel( + q, + k, + v, + o, + lse, + sinks, + seq_q_lens, + seq_kv_lens, + amax_s, + amax_o, + tma_k_desc, + tma_v_desc, + softmax_scale_log2, + o_scale_fused, + ).launch( + # THD: ceil(max_seq_q / q_tile) tiles per sequence over the real + # batch count (the packed view's batch mode is 1); tiles past a + # shorter sequence's length drain without work. + grid=( + ceil_div(self.thd_max_sq, self.q_tile) if cutlass.const_expr(self.thd_varlen) else ceil_div(q.shape[1], self.q_tile), + self.thd_batch if cutlass.const_expr(self.thd_varlen) else q.shape[0], + q.shape[2], + ), + block=(self.threads_per_cta, 1, 1), + stream=stream, + min_blocks_per_mp=1, + ) + + +@lru_cache(maxsize=None) +def compile( # noqa: A001 + compute_capability: tuple[int, int], + b: int = 1, + qh: int = 1, + kh: int = 1, + sq: int = 128, + skv: int = 128, + d_qk: int = 128, + d_v: int = 128, + max_sq: int = 0, + has_lse: bool = True, +) -> Callable: + """Compile and cache one architecture-specific compact BSHD shape. + + ``d_qk`` is the Q/K head dim (QK^T contraction width) and ``d_v`` the V/O + head dim (P@V output width); they are independent, e.g. (192, 128). + + THD specializations pack the batch: ``b`` is the real sequence count, + ``sq``/``skv`` are the packed token totals, and ``max_sq`` (the longest + sequence's Q length) sizes the per-sequence grid. + + ``has_lse=False`` compiles the LSE store out (the kernel specializes on a + ``None`` LSE argument) — callers that don't want stats pass no LSE buffer + at all instead of a dummy. + """ + + kernel = SM120FusedMultiHeadAttentionForward( + in_dtype=STORAGE_DTYPE, + out_dtype=OUT_STORAGE_DTYPE, + is_causal=PARAMS.is_causal, + causal_bottom_right=PARAMS.causal_bottom_right, + window_size_left=PARAMS.window_size_left, + seq_q_lens_present=PARAMS.seq_q_lens_present, + seq_kv_lens_present=PARAMS.seq_kv_lens_present, + has_sink=PARAMS.has_sink, + thd_varlen=PARAMS.thd_varlen, + thd_batch=b, + thd_max_sq=max_sq, + head_tile_qk=d_qk, + head_tile_v=d_v, + q_tile=PARAMS.q_tile, + kv_tile=PARAMS.kv_tile, + ) + fake_batch = 1 if PARAMS.thd_varlen else b + fake_q = cute.runtime.make_fake_compact_tensor( + STORAGE_DTYPE, + (fake_batch, sq, qh, d_qk), + stride_order=(3, 2, 1, 0), + assumed_align=16, + ) + fake_k = cute.runtime.make_fake_compact_tensor( + STORAGE_DTYPE, + (fake_batch, skv, kh, d_qk), + stride_order=(3, 2, 1, 0), + assumed_align=16, + ) + fake_v = cute.runtime.make_fake_compact_tensor( + STORAGE_DTYPE, + (fake_batch, skv, kh, d_v), + stride_order=(3, 2, 1, 0), + assumed_align=16, + ) + fake_o = cute.runtime.make_fake_compact_tensor( + OUT_STORAGE_DTYPE, + (fake_batch, sq, qh, d_v), + stride_order=(3, 2, 1, 0), + assumed_align=16, + ) + fake_lse = ( + cute.runtime.make_fake_compact_tensor( + cutlass.Float32, + (fake_batch, qh, sq), + stride_order=(2, 1, 0), + assumed_align=4, + ) + if has_lse + else None + ) + fake_sinks = ( + cute.runtime.make_fake_compact_tensor( + cutlass.Float32, + (qh,), + stride_order=(0,), + assumed_align=4, + ) + if PARAMS.has_sink + else None + ) + fake_seq_q_lens = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, + (b,), + stride_order=(0,), + assumed_align=4, + ) + fake_seq_kv_lens = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, + (3 * b + 2,) if PARAMS.thd_varlen else (b,), # THD: [ seq_kv(B) | cu_q(B+1) | cu_k(B+1) ] + stride_order=(0,), + assumed_align=4, + ) + # Amax buffers are Int32 at the ABI (bitcast-fp32 atomic max targets); + # the adapter passes torch fp32 buffers as .view(torch.int32). + fake_amax_s = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, + (1,), + stride_order=(0,), + assumed_align=4, + ) + fake_amax_o = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, + (1,), + stride_order=(0,), + assumed_align=4, + ) + return cute.compile( + kernel, + fake_q, + fake_k, + fake_v, + fake_o, + fake_lse, + fake_sinks, + fake_seq_q_lens, + fake_seq_kv_lens, + fake_amax_s, + fake_amax_o, + cutlass.Float32(1.0), + cutlass.Float32(1.0), + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=False), + options="--enable-tvm-ffi", + ) diff --git a/test/python/gemm/frost/test_frontend_integration.py b/test/python/gemm/frost/test_frontend_integration.py index ac1240eea..948603b0c 100644 --- a/test/python/gemm/frost/test_frontend_integration.py +++ b/test/python/gemm/frost/test_frontend_integration.py @@ -16,7 +16,7 @@ from gemm_test_utils import requires_sm100 import cudnn -from cudnn.engines import MANIFEST, OUT_OF_TREE_ID_BASE, PlanConfig, Router, is_backend_engine, is_python_engine +from cudnn.engines import MANIFEST, OUT_OF_TREE_ID_BASE, PlanConfig, heuristics, is_backend_engine, is_python_engine pytestmark = pytest.mark.L0 @@ -262,7 +262,7 @@ def test_ranked_list_routes_and_all_routes_agree(route): @_GPU -def test_build_walk_falls_through_a_declining_plan(caplog): +def test_build_walk_falls_through_a_declining_plan(caplog, monkeypatch): """A plan that declines at build time is logged and the walk moves to the next entry — here a python engine ranked ahead of the backend, so the graph still builds and executes natively with no exception reaching the user @@ -283,13 +283,10 @@ def execute(self, graph, tensor_data, ctx=None): boom = Boom() - class BoomFirst(Router): - def plan(self, graph, engines): - return [PlanConfig(boom.engine_id)] + graph.backend_plan_entries() - a, b, bias_t, ref = _operands() g, A, B, bias, Y = _build_matmul_bias_relu() - g.set_router(BoomFirst()).register_backend(boom) + monkeypatch.setattr(heuristics, "rank", lambda graph, engines, backend_plans, modes=None: [PlanConfig(boom.engine_id)] + list(backend_plans)) + g.register_backend(boom) _plan(g) assert g.get_plan_name_at_index(0) == "frost_fake_always_fails" g.check_support() diff --git a/test/python/sdpa/frost/frost_test_utils.py b/test/python/sdpa/frost/frost_test_utils.py index f540b05ec..53838d6c0 100644 --- a/test/python/sdpa/frost/frost_test_utils.py +++ b/test/python/sdpa/frost/frost_test_utils.py @@ -61,3 +61,35 @@ def _dsl_usable(): def _dsl_installed() -> bool: """For the few call sites that gate inside a test body rather than on it.""" return _DSL_OK + + +def _is_plan_for(plan_name, engine) -> bool: + """A plan reads ``[]``: the heuristics name a concrete config + for every entry, so match on the engine, not on the whole plan name.""" + return plan_name == engine or plan_name.startswith(engine + "[") + + +def select_engine(graph, name, tiles=None): + """Pin the ranked entry for engine ``name`` (graph.plans holds the backend's + plans and the python engines' in one list). A pin is strict: check_support / + build_plans raise if that engine declines the graph. + + The FIRST entry for that engine is the heuristics' own best guess for this + shape. ``tiles`` pins a different one, so a test can run a tile the best + guess would not choose. + """ + names = [graph.get_plan_name_at_index(i) for i in range(len(graph.plans))] + if tiles is None: + index = next((i for i, n in enumerate(names) if _is_plan_for(n, name)), None) + assert index is not None, f"engine {name!r} did not claim this graph; plans={names}" + else: + want = f"tile_m={tiles[0]}, tile_n={tiles[1]}" + index = next((i for i, n in enumerate(names) if n.startswith(name + "[") and want in n), None) + assert index is not None, f"no plan for tiles {tiles}; plans={names}" + graph.select_plan(index) + return graph + + +def offers_engine(graph, name) -> bool: + """Whether any ranked entry is a plan for engine ``name``.""" + return any(_is_plan_for(graph.get_plan_name_at_index(i), name) for i in range(len(graph.plans))) diff --git a/test/python/sdpa/frost/test_sdpa_bwd_dsl_sm120.py b/test/python/sdpa/frost/test_sdpa_bwd_dsl_sm120.py index e25c5d1e6..582dc9dd1 100644 --- a/test/python/sdpa/frost/test_sdpa_bwd_dsl_sm120.py +++ b/test/python/sdpa/frost/test_sdpa_bwd_dsl_sm120.py @@ -46,14 +46,7 @@ def _require_dsl() -> None: pytest.skip("cutlass/dsl not installed") -def _select_engine(graph, name): - """Pin the ranked entry named ``name`` (graph.plans holds the backend's - plans and the python engines' in one list). A pin is strict: check_support / - build_plans raise if that engine declines the graph.""" - names = [graph.get_plan_name_at_index(i) for i in range(len(graph.plans))] - assert name in names, f"engine {name!r} did not claim this graph; plans={names}" - graph.select_plan(names.index(name)) - return graph +from frost_test_utils import select_engine as _select_engine # noqa: F401 def _bhsd(batch: int, heads: int, sequence: int, head_dim: int, dtype: torch.dtype, empty: bool = False) -> torch.Tensor: diff --git a/test/python/sdpa/frost/test_sdpa_frontend_integration.py b/test/python/sdpa/frost/test_sdpa_frontend_integration.py index 6e08da3f1..0ac4c2663 100644 --- a/test/python/sdpa/frost/test_sdpa_frontend_integration.py +++ b/test/python/sdpa/frost/test_sdpa_frontend_integration.py @@ -53,10 +53,17 @@ def _plan_names(g): return [g.get_plan_name_at_index(i) for i in range(len(g.plans))] +def _is_plan_for(plan_name, engine): + """A plan reads ``[]``: the heuristics name a concrete config + for every entry, so match on the engine, not on the whole plan name.""" + return plan_name == engine or plan_name.startswith(engine + "[") + + def _index_of(g, name): names = _plan_names(g) - assert name in names, f"no plan named {name!r} in {names}" - return names.index(name) + index = next((i for i, n in enumerate(names) if _is_plan_for(n, name)), None) + assert index is not None, f"no plan for engine {name!r} in {names}" + return index def _pin(g, name): @@ -94,9 +101,9 @@ def test_eligible_graph_lists_matching_dsl_engine(): g, q, k, v, o = _build_causal_sdpa() _plan(g) names = _plan_names(g) - assert _FROST in names + assert any(_is_plan_for(n, _FROST) for n in names) assert g.get_execution_plan_count() == len(names) == len(g.plans) - assert is_python_engine(g.plans[names.index(_FROST)].engine_id) + assert is_python_engine(g.plans[_index_of(g, _FROST)].engine_id) @_SM100 @@ -188,9 +195,9 @@ def test_envelope_lists_every_covering_flavor_smallest_first(): _plan(g) names = _plan_names(g) for flavor in (128, 256, 512): - assert engine_name(flavor) in names # every covering flavor + assert any(_is_plan_for(n, engine_name(flavor)) for n in names) # every covering flavor python = [i for i, p in enumerate(g.plans) if is_python_engine(p.engine_id)] - assert names[python[0]] == engine_name(128) # tightest flavor first + assert _is_plan_for(names[python[0]], engine_name(128)) # tightest flavor first g.select_plan(python[0]) g.check_support() g.build_plans() @@ -243,8 +250,10 @@ def test_no_magic_import_required(): "o.set_output(True).set_dim(q_gpu.shape).set_stride(q_gpu.stride())\n" "g.validate(); g.build_operation_graph(); g.create_execution_plans([cudnn.heur_mode.A])\n" "names = [g.get_plan_name_at_index(i) for i in range(len(g.plans))]\n" - "assert 'sdpa_fwd_prefill_sm100_d512' in names, names\n" - "g.select_plan(names.index('sdpa_fwd_prefill_sm100_d512'))\n" + "want = 'sdpa_fwd_prefill_sm100_d512'\n" + "i = next((i for i, n in enumerate(names) if n == want or n.startswith(want + '[')), None)\n" + "assert i is not None, names\n" + "g.select_plan(i)\n" "g.check_support()\n" "print('ELIGIBLE-WITHOUT-IMPORT')\n" ) diff --git a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py index 48751e8bc..823b07d3d 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py @@ -15,16 +15,7 @@ from frost_test_utils import requires_blackwell, requires_dsl, _dsl_installed -def _select_engine(graph, name): - """Pin the ranked entry named ``name`` (graph.plans holds the backend's - plans and the python engines' in one list). A pin is strict: check_support / - build_plans raise if that engine declines, so an ineligible config cannot - silently fall back to native cuDNN.""" - names = [graph.get_plan_name_at_index(i) for i in range(len(graph.plans))] - assert name in names, f"engine {name!r} did not claim this graph; plans={names}" - graph.select_plan(names.index(name)) - return graph - +from frost_test_utils import select_engine as _select_engine # noqa: F401 pytestmark = requires_blackwell diff --git a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py index 4bbdda888..036f3e52e 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py @@ -37,14 +37,7 @@ def _require_dsl() -> None: pytest.skip("cutlass/dsl not installed") -def _select_engine(graph, name): - """Pin the ranked entry named ``name`` (graph.plans holds the backend's - plans and the python engines' in one list). A pin is strict: check_support / - build_plans raise if that engine declines the graph.""" - names = [graph.get_plan_name_at_index(i) for i in range(len(graph.plans))] - assert name in names, f"engine {name!r} did not claim this graph; plans={names}" - graph.select_plan(names.index(name)) - return graph +from frost_test_utils import select_engine as _select_engine # noqa: F401 def _bhsd( diff --git a/test/python/sdpa/frost/test_sdpa_fwd_fp8_sm100.py b/test/python/sdpa/frost/test_sdpa_fwd_fp8_sm100.py index 5444cd7de..bb2765b15 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_fp8_sm100.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_fp8_sm100.py @@ -26,15 +26,7 @@ from frost_test_utils import requires_blackwell, requires_dsl -def _select_engine(graph, name): - """Pin the ranked entry named ``name`` (graph.plans holds the backend's - plans and the python engines' in one list). A pin is strict: check_support / - build_plans raise if that engine declines the graph.""" - names = [graph.get_plan_name_at_index(i) for i in range(len(graph.plans))] - assert name in names, f"engine {name!r} did not claim this graph; plans={names}" - graph.select_plan(names.index(name)) - return graph - +from frost_test_utils import select_engine as _select_engine # noqa: F401 pytestmark = [requires_blackwell, requires_dsl] diff --git a/test/python/sdpa/frost/test_sdpa_fwd_fp8_sm120.py b/test/python/sdpa/frost/test_sdpa_fwd_fp8_sm120.py new file mode 100644 index 000000000..d95157980 --- /dev/null +++ b/test/python/sdpa/frost/test_sdpa_fwd_fp8_sm120.py @@ -0,0 +1,320 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""End-to-end tests for the FROST SM120 DSL per-tensor FP8 SDPA-forward engine. + +Drives ``graph.sdpa_fp8`` (FP8 E4M3 Q/K/V + scalar per-tensor descales) routed to +the ``sdpa_fwd_prefill_sm120_fp8`` engine, and validates O against an fp32-dequant +reference. ``Amax_S`` and ``Amax_O`` are both produced in-kernel (bitcast-int32 +atomicMax over the pre-cast fp32 values); both are checked. + +SM120 v1 envelope (see engines._sm120_fp8_spec): E4M3 in / FP16 out only, exact +d128, causal / bottom-right / SWA / KV-padding masks; no sink (Amax_S semantics), +no THD. E5M2 inputs, FP8 outputs and non-128 head dims are covered by negative tests. + +Requires: SM120/SM121 (consumer Blackwell), cutlass-dsl. Skips otherwise. +""" + +import math + +import pytest +import torch + +from test_utils import torch_fork_set_rng + +from cudnn.sdpa.fwd.engines import engine_name +from frost_test_utils import requires_blackwell_geforce, requires_dsl + + +from frost_test_utils import select_engine as _select_engine # noqa: F401 + +pytestmark = [requires_blackwell_geforce, requires_dsl] + +_E4M3_MAX = 448.0 + + +def _quant(x): + dq = (x.abs().amax().clamp_min(1e-8) / _E4M3_MAX).item() + return (x / dq).clamp(-_E4M3_MAX, _E4M3_MAX).to(torch.float8_e4m3fn), dq + + +def _ref(qd, kd, vd, *, scale, is_causal=False, bottom_right=False, swa_window=None, seq_lens_kv=None): + b, h_q, s_q, _ = qd.shape + _, h_kv, s_kv, _ = vd.shape + dev = qd.device + g = h_q // h_kv + k_e = kd.repeat_interleave(g, dim=1) + v_e = vd.repeat_interleave(g, dim=1) + scores = torch.matmul(qd, k_e.transpose(-1, -2)) * scale + i = torch.arange(s_q, device=dev).view(1, 1, s_q, 1) + j = torch.arange(s_kv, device=dev).view(1, 1, 1, s_kv) + masked = torch.zeros(1, 1, s_q, s_kv, dtype=torch.bool, device=dev) + if is_causal: + lim = i + (s_kv - s_q) if bottom_right else i + masked = masked | (j > lim) + if swa_window is not None: + masked = masked | (j < i - swa_window) + if seq_lens_kv is not None: + slk = torch.as_tensor(seq_lens_kv, device=dev, dtype=torch.long).view(b, 1, 1, 1) + masked = masked | (j >= slk) + scores = scores.masked_fill(masked, float("-inf")) + probs = torch.softmax(scores, dim=-1) + return torch.matmul(probs, v_e), probs.max().item() + + +def _run(B, H_q, H_kv, S_q, S_kv, *, scale, sdpa_kwargs, seq_lens_kv=None, tiles=None): + import cudnn + + dev = "cuda" + D = 128 + Qf = torch.randn(B, H_q, S_q, D, device=dev) * 0.5 + Kf = torch.randn(B, H_kv, S_kv, D, device=dev) * 0.5 + Vf = torch.randn(B, H_kv, S_kv, D, device=dev) * 0.5 + Q8, dq = _quant(Qf) + K8, dk = _quant(Kf) + V8, dv = _quant(Vf) + + def bshd(x8): + return x8.permute(0, 2, 1, 3).contiguous().transpose(1, 2) + + Qb, Kb, Vb = bshd(Q8), bshd(K8), bshd(V8) + Ob = torch.empty(B, S_q, H_q, D, device=dev, dtype=torch.float16).transpose(1, 2) + lse = torch.empty(B, H_q, S_q, 1, device=dev, dtype=torch.float32) + amax_s = torch.zeros(1, 1, 1, 1, device=dev, dtype=torch.float32) + amax_o = torch.zeros(1, 1, 1, 1, device=dev, dtype=torch.float32) + + def sc(val): + return torch.tensor([[[[val]]]], dtype=torch.float32, device=dev) + + dqt, dkt, dvt, dst, sst, sot = sc(dq), sc(dk), sc(dv), sc(1.0), sc(1.0), sc(1.0) + + g = cudnn.pygraph(io_data_type=cudnn.data_type.FP8_E4M3, intermediate_data_type=cudnn.data_type.FLOAT, compute_data_type=cudnn.data_type.FLOAT) + q = g.tensor_like(Qb) + k = g.tensor_like(Kb) + v = g.tensor_like(Vb) + + def _stns(): + return g.tensor(dim=[1, 1, 1, 1], stride=[1, 1, 1, 1], data_type=cudnn.data_type.FLOAT) + + dqn, dkn, dvn, dsn, ssn, son = (_stns() for _ in range(6)) + kw = dict(q=q, k=k, v=v, descale_q=dqn, descale_k=dkn, descale_v=dvn, descale_s=dsn, scale_s=ssn, scale_o=son, attn_scale=scale, generate_stats=True) + vp = {q: Qb, k: Kb, v: Vb, dqn: dqt, dkn: dkt, dvn: dvt, dsn: dst, ssn: sst, son: sot} + if seq_lens_kv is not None: + slq = torch.full((B, 1, 1, 1), S_q, dtype=torch.int32, device=dev) + slk = torch.tensor(seq_lens_kv, dtype=torch.int32, device=dev).reshape(B, 1, 1, 1) + sq_h = g.tensor(dim=[B, 1, 1, 1], stride=[1, 1, 1, 1], data_type=cudnn.data_type.INT32) + skv_h = g.tensor(dim=[B, 1, 1, 1], stride=[1, 1, 1, 1], data_type=cudnn.data_type.INT32) + kw.update(use_padding_mask=True, seq_len_q=sq_h, seq_len_kv=skv_h) + vp[sq_h] = slq + vp[skv_h] = slk + kw.update(sdpa_kwargs) + o, stats, amx_s, amx_o = g.sdpa_fp8(**kw) + o.set_output(True).set_dim(list(Ob.shape)).set_stride(list(Ob.stride())).set_data_type(cudnn.data_type.HALF) + stats.set_output(True).set_dim([B, H_q, S_q, 1]).set_stride([H_q * S_q, S_q, 1, 1]).set_data_type(cudnn.data_type.FLOAT) + amx_s.set_output(True).set_dim([1, 1, 1, 1]).set_stride([1, 1, 1, 1]).set_data_type(cudnn.data_type.FLOAT) + amx_o.set_output(True).set_dim([1, 1, 1, 1]).set_stride([1, 1, 1, 1]).set_data_type(cudnn.data_type.FLOAT) + + g.validate() + g.build_operation_graph() + g.create_execution_plans([cudnn.heur_mode.A]) + _select_engine(g, engine_name(arch="sm120", fp8=True), tiles=tiles) + g.check_support() + g.build_plans() + vp.update({o: Ob, stats: lse, amx_s: amax_s, amx_o: amax_o}) + g.execute(vp, torch.empty(max(g.get_workspace_size(), 1), device=dev, dtype=torch.uint8)) + torch.cuda.synchronize() + + ref_kw = _ref_kwargs(sdpa_kwargs) + o_ref, amax_s_ref = _ref(Q8.float() * dq, K8.float() * dk, V8.float() * dv, scale=scale, seq_lens_kv=seq_lens_kv, **ref_kw) + return Ob, o_ref, amax_s.item(), amax_s_ref, amax_o.item(), o_ref.abs().max().item() + + +def _ref_kwargs(sdpa_kwargs): + out = {} + if sdpa_kwargs.get("use_causal_mask"): + out["is_causal"] = True + if sdpa_kwargs.get("use_causal_mask_bottom_right"): + out["is_causal"] = True + out["bottom_right"] = True + lb = sdpa_kwargs.get("left_bound") + if lb is not None: + out["swa_window"] = lb - 1 + return out + + +def _check(out, o_ref, amax_s, amax_s_ref, amax_o, amax_o_ref): + diff = (out.float() - o_ref).abs().max().item() + assert diff <= 5e-2, f"max|O-ref|={diff:.4f} > 0.05" + assert abs(amax_s - amax_s_ref) <= 0.03, f"amax_s {amax_s:.4f} vs ref {amax_s_ref:.4f}" + assert abs(amax_o - amax_o_ref) <= 0.03, f"amax_o {amax_o:.4f} vs ref {amax_o_ref:.4f}" + + +_MASKS = { + "none": {}, + "causal": dict(use_causal_mask=True), + "causal_br": dict(use_causal_mask_bottom_right=True), + "swa": dict(use_causal_mask=True, left_bound=65), +} + + +@pytest.mark.L0 +@pytest.mark.parametrize("mask", list(_MASKS)) +@torch_fork_set_rng(seed=0) +def test_fp8_sm120_masks(mask): + scale = 1.0 / math.sqrt(128) + out, o_ref, a_s, a_s_ref, a_o, a_o_ref = _run(2, 8, 8, 256, 256, scale=scale, sdpa_kwargs=_MASKS[mask]) + _check(out, o_ref, a_s, a_s_ref, a_o, a_o_ref) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +def test_fp8_sm120_gqa(): + scale = 1.0 / math.sqrt(128) + out, o_ref, a_s, a_s_ref, a_o, a_o_ref = _run(2, 8, 2, 256, 256, scale=scale, sdpa_kwargs=dict(use_causal_mask=True)) + _check(out, o_ref, a_s, a_s_ref, a_o, a_o_ref) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +def test_fp8_sm120_bottom_right_rectangular(): + scale = 1.0 / math.sqrt(128) + out, o_ref, a_s, a_s_ref, a_o, a_o_ref = _run(2, 8, 8, 128, 256, scale=scale, sdpa_kwargs=dict(use_causal_mask_bottom_right=True)) + _check(out, o_ref, a_s, a_s_ref, a_o, a_o_ref) + + +@pytest.mark.L1 +@torch_fork_set_rng(seed=0) +def test_fp8_sm120_multi_tile_long_seq(): + # 1k x 1k exercises the multi-KV-tile online-softmax rescale path. + scale = 1.0 / math.sqrt(128) + out, o_ref, a_s, a_s_ref, a_o, a_o_ref = _run(1, 4, 4, 1024, 1024, scale=scale, sdpa_kwargs=dict(use_causal_mask=True)) + _check(out, o_ref, a_s, a_s_ref, a_o, a_o_ref) + + +@pytest.mark.L0 +@pytest.mark.parametrize("causal", [False, True]) +@torch_fork_set_rng(seed=0) +def test_fp8_sm120_padding(causal): + # KV padding: batch 0 uses all 256 KV cols, batch 1 only 192 (partial tile). + scale = 1.0 / math.sqrt(128) + sk = dict(use_causal_mask=True) if causal else {} + out, o_ref, a_s, a_s_ref, a_o, a_o_ref = _run(2, 8, 8, 256, 256, scale=scale, sdpa_kwargs=sk, seq_lens_kv=[256, 192]) + _check(out, o_ref, a_s, a_s_ref, a_o, a_o_ref) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +def test_fp8_sm120_e5m2_not_offered(): + """The v1 kernel hardcodes the e4m3 MMA tag; E5M2 graphs must not route here.""" + import cudnn + + dev = "cuda" + B, H, S, D = 1, 4, 256, 128 + X = torch.randn(B, S, H, D, device=dev).to(torch.float8_e5m2).transpose(1, 2) + g = cudnn.pygraph(io_data_type=cudnn.data_type.FP8_E5M2, intermediate_data_type=cudnn.data_type.FLOAT, compute_data_type=cudnn.data_type.FLOAT) + q, k, v = g.tensor_like(X), g.tensor_like(X), g.tensor_like(X) + scalars = [g.tensor(dim=[1, 1, 1, 1], stride=[1, 1, 1, 1], data_type=cudnn.data_type.FLOAT) for _ in range(6)] + o, stats, amx_s, amx_o = g.sdpa_fp8( + q=q, + k=k, + v=v, + descale_q=scalars[0], + descale_k=scalars[1], + descale_v=scalars[2], + descale_s=scalars[3], + scale_s=scalars[4], + scale_o=scalars[5], + attn_scale=1.0 / math.sqrt(D), + generate_stats=True, + ) + o.set_output(True).set_dim([B, H, S, D]).set_stride([S * H * D, D, H * D, 1]).set_data_type(cudnn.data_type.HALF) + stats.set_output(True).set_dim([B, H, S, 1]).set_stride([H * S, S, 1, 1]).set_data_type(cudnn.data_type.FLOAT) + amx_s.set_output(True).set_dim([1, 1, 1, 1]).set_stride([1, 1, 1, 1]).set_data_type(cudnn.data_type.FLOAT) + amx_o.set_output(True).set_dim([1, 1, 1, 1]).set_stride([1, 1, 1, 1]).set_data_type(cudnn.data_type.FLOAT) + g.validate() + g.build_operation_graph() + try: + g.create_execution_plans([cudnn.heur_mode.A]) + except cudnn.cudnnGraphNotSupportedError: + # Nothing — python engine or backend — serves E5M2 here: also a pass + # (the point is only that the e4m3-tagged sm120 fp8 cell declined). + return + names = [g.get_plan_name_at_index(i) for i in range(len(g.plans))] + assert engine_name(arch="sm120", fp8=True) not in names, f"E5M2 graph must not offer the sm120 fp8 engine; plans={names}" + + +def _fp8_graph_offers_sm120(io_dtype, o_dtype, D=128, sink=False): + """Build one sdpa_fp8 graph and report whether the sm120 fp8 cell claims it. + + A capability rejection is the point, so nothing is executed; a graph that + no engine at all serves counts as declined too. + """ + import cudnn + + dev = "cuda" + B, H, S = 1, 4, 256 + torch_in = torch.float8_e5m2 if io_dtype == cudnn.data_type.FP8_E5M2 else torch.float8_e4m3fn + X = torch.randn(B, S, H, D, device=dev).to(torch_in).transpose(1, 2) + g = cudnn.pygraph(io_data_type=io_dtype, intermediate_data_type=cudnn.data_type.FLOAT, compute_data_type=cudnn.data_type.FLOAT) + q, k, v = g.tensor_like(X), g.tensor_like(X), g.tensor_like(X) + scalars = [g.tensor(dim=[1, 1, 1, 1], stride=[1, 1, 1, 1], data_type=cudnn.data_type.FLOAT) for _ in range(6)] + kw = dict( + q=q, + k=k, + v=v, + descale_q=scalars[0], + descale_k=scalars[1], + descale_v=scalars[2], + descale_s=scalars[3], + scale_s=scalars[4], + scale_o=scalars[5], + attn_scale=1.0 / math.sqrt(D), + generate_stats=True, + ) + if sink: + kw["sink_token"] = g.tensor(dim=[1, H, 1, 1], stride=[H, 1, 1, 1], data_type=cudnn.data_type.FLOAT) + o, stats, amx_s, amx_o = g.sdpa_fp8(**kw) + o.set_output(True).set_dim([B, H, S, D]).set_stride([S * H * D, D, H * D, 1]).set_data_type(o_dtype) + stats.set_output(True).set_dim([B, H, S, 1]).set_stride([H * S, S, 1, 1]).set_data_type(cudnn.data_type.FLOAT) + for t in (amx_s, amx_o): + t.set_output(True).set_dim([1, 1, 1, 1]).set_stride([1, 1, 1, 1]).set_data_type(cudnn.data_type.FLOAT) + try: + g.validate() + g.build_operation_graph() + g.create_execution_plans([cudnn.heur_mode.A]) + except (cudnn.cudnnGraphNotSupportedError, RuntimeError, ValueError): + # The op itself may refuse the shape before any engine is consulted; + # for "this cell must not claim it" that is the same answer. + return False + return engine_name(arch="sm120", fp8=True) in [g.get_plan_name_at_index(i) for i in range(len(g.plans))] + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +def test_fp8_sm120_fp8_output_not_offered(): + """The epilogue stores FP16; an fp8 O would need a quantizing store.""" + import cudnn + + assert not _fp8_graph_offers_sm120(cudnn.data_type.FP8_E4M3, cudnn.data_type.FP8_E4M3) + + +@pytest.mark.L0 +@pytest.mark.parametrize("D", [64, 256]) +@torch_fork_set_rng(seed=0) +def test_fp8_sm120_non_128_head_dim_not_offered(D): + """The 8-bit fragment path has no zero-padding envelope, so d is exact.""" + import cudnn + + assert not _fp8_graph_offers_sm120(cudnn.data_type.FP8_E4M3, cudnn.data_type.HALF, D=D) + + +@pytest.mark.L0 +@pytest.mark.parametrize("tiles", [(64, 64), (64, 128), (128, 64), (128, 128)]) +@torch_fork_set_rng(seed=0) +def test_fp8_sm120_every_enumerated_tile(tiles): + """The heuristics offer the whole tile domain, but a shape only ever runs + one point of it, so the rest would ship untested. S_q=256 keeps both q_tile + values meaningful (one full tile at 128, two at 64).""" + scale = 1.0 / math.sqrt(128) + out, o_ref, a_s, a_s_ref, a_o, a_o_ref = _run(2, 8, 8, 256, 256, scale=scale, sdpa_kwargs=dict(use_causal_mask=True), tiles=tiles) + _check(out, o_ref, a_s, a_s_ref, a_o, a_o_ref) diff --git a/test/python/sdpa/frost/test_sdpa_fwd_mxfp8_sm100.py b/test/python/sdpa/frost/test_sdpa_fwd_mxfp8_sm100.py index 4f86bd180..445bf2bcd 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_mxfp8_sm100.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_mxfp8_sm100.py @@ -30,15 +30,7 @@ from frost_test_utils import requires_blackwell, requires_dsl -def _select_engine(graph, name): - """Pin the ranked entry named ``name`` (graph.plans holds the backend's - plans and the python engines' in one list). A pin is strict: check_support / - build_plans raise if that engine declines the graph.""" - names = [graph.get_plan_name_at_index(i) for i in range(len(graph.plans))] - assert name in names, f"engine {name!r} did not claim this graph; plans={names}" - graph.select_plan(names.index(name)) - return graph - +from frost_test_utils import select_engine as _select_engine # noqa: F401 pytestmark = [requires_blackwell, requires_dsl] diff --git a/test/python/sdpa/frost/test_sdpa_graph_analyzer.py b/test/python/sdpa/frost/test_sdpa_graph_analyzer.py index bfd266915..29a81f1fd 100644 --- a/test/python/sdpa/frost/test_sdpa_graph_analyzer.py +++ b/test/python/sdpa/frost/test_sdpa_graph_analyzer.py @@ -17,11 +17,11 @@ def _eligible(graph, knobs=None): """Names of the FROST SDPA engines whose caps match this graph. - ``knobs`` is passed straight to the probe: graph.set_engine_knobs() was - removed with the monkey-patch dispatch layer, and a knob request is a + ``knobs`` is passed straight to the capability match: graph.set_engine_knobs() + was removed with the monkey-patch dispatch layer, and a knob request is a property of a PLAN (engines.base.PlanConfig.knobs), not of the graph. """ - return {s.name for s in engines.ENGINE_SPECS if engines.probe(s, graph, knobs)} + return {s.name for s in engines.ENGINE_SPECS if engines.analyze_for(s, graph, knobs)[1] is None} # The default pytest.ini addopts is `-m L0`; mark the whole module so it runs. @@ -782,7 +782,7 @@ def test_sm120_knob_domains(monkeypatch): def _bwd_eligible(graph, knobs=None): """Names of the FROST SDPA-backward engines whose caps match this graph.""" - return {s.name for s in bwd_engines.ENGINE_SPECS if bwd_engines.probe(s, graph, knobs)} + return {s.name for s in bwd_engines.ENGINE_SPECS if bwd_engines.analyze_for(s, graph, knobs)[1] is None} def _bshd_strides(h: int, s: int, d: int) -> tuple[int, int, int, int]: diff --git a/test/python/sdpa/frost/test_sm120_tiles.py b/test/python/sdpa/frost/test_sm120_tiles.py new file mode 100644 index 000000000..a75bef6f0 --- /dev/null +++ b/test/python/sdpa/frost/test_sm120_tiles.py @@ -0,0 +1,85 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""The SM120 tile choice and the knob domain it is chosen from. + +One rule for both cells (fp8 and f16): the fp8 and bf16 sweeps disagree on how +much each tile wins by, never on which one wins. The rule half needs no device +-- it is arithmetic over shape and SM count. The correctness half does, because +the heuristics expose tiles the best guess never picks, and nothing else in the +suite runs a kernel at those. +""" + +import pytest + +from cudnn.sdpa.fwd.config_sm120 import SEQ_KV_TILES, SEQ_Q_TILES, tile_choice + +SMS = 188 # RTX PRO 6000 Blackwell, the machine the rule was measured on + +# Arithmetic over shape and SM count: no GPU, no DSL, so no arch gate either. +pytestmark = pytest.mark.L0 + + +class TestTileChoice: + """No device: shape, SM count and causality in, (q_tile, kv_tile) out.""" + + @pytest.mark.parametrize("s_q,h_q,b", [(512, 16, 1), (4096, 16, 1), (1024, 16, 1), (16384, 16, 1), (2048, 32, 4)]) + @pytest.mark.parametrize("causal", [False, True]) + def test_kv_tile_is_always_128(self, s_q, h_q, b, causal): + """Fastest in all 28 fp8 shapes measured and in the bf16 sweep. It was + 64 while P was staged through SMEM, because that traffic scaled with + the KV tile -- a change here is a claim about the kernel, not about the + hardware.""" + assert tile_choice(s_q, s_q, h_q, b, SMS, causal)[1] == 128 + + def test_grid_too_small_to_fill_the_machine_takes_the_finer_q_tile(self): + # grid = ceil(512/128) * 16 = 64 CTAs of 188: doubling them still fits, + # and it is worth 1.5x. + assert tile_choice(512, 512, 16, 1, SMS, False)[0] == 64 + + def test_a_short_sequence_does_not_amortize_the_finer_q_tile(self): + # Same underfilled machine, but 8 KV tiles cannot absorb the extra + # Q-tile loop, so the coarse tile wins by 1.22x. + assert tile_choice(1024, 1024, 16, 1, SMS, False)[0] == 128 + + def test_underfilled_and_long_takes_the_finer_q_tile(self): + # grid 256 with 16 KV tiles: both conditions hold, worth 1.06-1.11x. + assert tile_choice(2048, 2048, 16, 1, SMS, False)[0] == 64 + assert tile_choice(2048, 2048, 8, 2, SMS, False)[0] == 64 + + @pytest.mark.parametrize("s_q", [4096, 8192, 16384]) + def test_a_full_machine_takes_the_coarser_q_tile(self, s_q): + assert tile_choice(s_q, s_q, 16, 1, SMS, False)[0] == 128 + + def test_the_grid_bound_sits_between_240_and_320_ctas(self): + """Where the two tiles stop trading evenly on a 188-SM part: 240 CTAs + still want the finer tile, 320 want the coarser one by 1.19x.""" + assert tile_choice(2560, 2560, 12, 1, SMS, False)[0] == 64 # grid 240 + assert tile_choice(2560, 2560, 16, 1, SMS, False)[0] == 128 # grid 320 + + def test_batch_and_heads_count_toward_the_grid(self): + """The rule reads grid, not sequence length: the same s_q flips once + the batch supplies enough CTAs on its own.""" + assert tile_choice(2048, 2048, 16, 1, SMS, False)[0] == 64 + assert tile_choice(2048, 2048, 16, 4, SMS, False)[0] == 128 + + def test_a_causal_mask_widens_the_window_for_the_finer_q_tile(self): + """A triangular mask halves the work per CTA, so the machine empties + sooner. Both directions are asserted -- checking only the causal side + would pass for a rule that ignored causality.""" + assert tile_choice(1024, 1024, 16, 1, SMS, True)[0] == 64 + assert tile_choice(1024, 1024, 16, 1, SMS, False)[0] == 128 + + def test_causality_does_not_override_a_full_machine(self): + assert tile_choice(8192, 8192, 16, 1, SMS, True)[0] == 128 + + def test_choice_is_always_in_the_domain(self): + for s_q in (128, 512, 1024, 4096, 32768): + for b in (1, 8): + for causal in (False, True): + q, kv = tile_choice(s_q, s_q, 16, b, SMS, causal) + assert q in SEQ_Q_TILES and kv in SEQ_KV_TILES + + def test_unknown_sm_count_falls_back_without_dividing_by_zero(self): + assert tile_choice(4096, 4096, 16, 1, 0, False) == (128, 128) + assert tile_choice(4096, 4096, 16, 1, 0, True) == (128, 128) diff --git a/test/python/test_engine_router.py b/test/python/test_engine_router.py index 2bf7ef83b..29239bad9 100644 --- a/test/python/test_engine_router.py +++ b/test/python/test_engine_router.py @@ -1,9 +1,9 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""CPU tests for the Router + BaseEngine contract of the unified dispatch. +"""CPU tests for the ranking + BaseEngine contract of the unified dispatch. -These run without a GPU: they exercise pygraph -> Router -> ONE ranked plan +These run without a GPU: they exercise pygraph -> heuristics -> ONE ranked plan list -> engine-id dispatch, using throwaway engines defined in this file (``TorchMatmulEngine`` is the pure-torch matmul(+bias/relu) oracle the execute tests run against). This is the CI-safe proof that the unification contract @@ -11,7 +11,7 @@ There is ONE plan list: the python engines that claim the graph (registered here plus whatever ``engines/manifest.py`` matches) and the backend's own -ranked entries, merged by ``engines.heuristics.heuristics_sort``. Which of +ranked entries, ordered by ``engines.heuristics.rank``. Which of those participate depends on the machine, so every assertion below addresses a plan by NAME or engine id and pins it with ``select_plan()`` — never by a hard-coded absolute index. @@ -22,7 +22,7 @@ torch = pytest.importorskip("torch") from cudnn._pygraph import pygraph -from cudnn.engines import BaseEngine, Router, OUT_OF_TREE_ID_BASE, is_backend_engine, is_python_engine +from cudnn.engines import BaseEngine, OUT_OF_TREE_ID_BASE, is_backend_engine, is_python_engine from cudnn.engines.base import resolve_node_buffers from cudnn.engines.engine_ids import BACKEND_HEURISTIC_ENGINE_ID from cudnn.graph_types import NodeType @@ -136,7 +136,9 @@ def execute(self, graph, tensor_data, ctx=None): g.matmul(a, b, name="mm") g.register_backend(Accepts()).register_backend(Declines()) - plans = Router().plan(g, g.backends) + from cudnn.engines import heuristics + + plans = heuristics.rank(g, g.backends, g.backend_plan_entries()) ids = [p.engine_id for p in plans] assert _OOT + 10 in ids # the supporting python engine assert _OOT + 50 not in ids # the declining one never joins @@ -283,16 +285,16 @@ def test_no_python_engine_plan_list_is_backend_only(): b = g.tensor(dim=[8, 4], name="B") g.matmul(a, b, name="mm") - from cudnn.engines.router import default_router + from cudnn.engines import heuristics - plans = default_router.plan(g, []) + plans = heuristics.rank(g, [], g.backend_plan_entries()) assert all(is_backend_engine(p.engine_id) for p in plans) g._plans = plans assert g.selected_engine is None # backend path -def test_compiled_plan_lifecycle_knobs_and_reuse(): - """Review item 1 acceptance: multiple knob proposals from one engine; the +def test_compiled_plan_lifecycle_knobs_and_reuse(monkeypatch): + """Review item 1 acceptance: several knob candidates for one engine; the selected plan's knobs reach build_plan; compilation runs once per plan and the artifact is reused; caller workspace + stream context reach execute.""" from cudnn.engines import CompiledPlan, PlanConfig @@ -314,15 +316,18 @@ class Tunable(BaseEngine): name = "tunable" engine_id = _OOT + 40 - def propose_plans(self, graph): - return [PlanConfig(self.engine_id, {"tile": 128}), PlanConfig(self.engine_id, {"tile": 256})] - def build_plan(self, graph, plan, ctx=None): compiled_log.append(plan.knobs) return TunablePlan(plan.knobs) + tunable = Tunable() + _ranking( + monkeypatch, + lambda graph, engines, backend_plans, modes=None: [PlanConfig(tunable.engine_id, {"tile": 128}), PlanConfig(tunable.engine_id, {"tile": 256})], + ) + g = pygraph() - g.register_backend(Tunable()) + g.register_backend(tunable) C = g.matmul(torch.randn(2, 2), torch.randn(2, 2)) g.create_execution_plans() tuned = _python_indices(g) @@ -342,7 +347,7 @@ def build_plan(self, graph, plan, ctx=None): # same engine instance on a second graph: no state collision g2 = pygraph() - g2.register_backend(Tunable()) + g2.register_backend(tunable) C2 = g2.matmul(torch.randn(2, 2), torch.randn(2, 2)) g2.create_execution_plans() g2.select_plan(_python_indices(g2)[0]) @@ -350,6 +355,17 @@ def build_plan(self, graph, plan, ctx=None): assert compiled_log == [{"tile": 256}, {"tile": 128}] # g2 compiled its own plan +def _ranking(monkeypatch, fn): + """Replace the ranking policy for one graph's planning. + + What a Router subclass used to do. Ranking has one home now, so a test + that wants a specific order says so here rather than by subclassing. + """ + from cudnn.engines import heuristics + + monkeypatch.setattr(heuristics, "rank", fn) + + def _mk_engine(id_off, knobs=None, log=None): from cudnn.engines import CompiledPlan @@ -363,10 +379,11 @@ def execute(self, graph, tensor_data, ctx): class _E(BaseEngine): name = f"e{id_off}" engine_id = _OOT + id_off - default_knobs = knobs def build_plan(self, graph, plan, ctx=None): - return _Plan(plan.knobs) + # The ranking names the knobs now; this fake keeps its own marker + # for the entries a test did not give any. + return _Plan(plan.knobs if plan.knobs is not None else knobs) def execute(self, graph, tensor_data, ctx=None): pass @@ -392,9 +409,10 @@ def test_planning_is_one_shot(): assert log[-1] == "old" # the planned artifact, unchanged -def test_mixed_router_ordering_dispatch(monkeypatch): - """Follow-up item 2: dispatch honors arbitrary Router ordering (backend - entry in the middle), never a python-prefix assumption. +def test_mixed_ranking_dispatch(monkeypatch): + """Follow-up item 2: dispatch honors arbitrary ranking (a backend entry in + the MIDDLE), never a python-prefix assumption. Not hypothetical any more -- + the sdpa-fwd heuristics place backend entries between their own. The backend's own entries are stubbed so the ordering is the same with or without a cuDNN that accepts this toy graph; real execution THROUGH the @@ -408,11 +426,11 @@ def test_mixed_router_ordering_dispatch(monkeypatch): stub = PlanConfig(0, {}, cpp_index=0) monkeypatch.setattr(pygraph, "backend_plan_entries", lambda self: [stub]) - class Interleaved(Router): - def plan(self, graph, backends): - return [PlanConfig(ea.engine_id, "A")] + graph.backend_plan_entries() + [PlanConfig(eb.engine_id, "B")] + _ranking( + monkeypatch, lambda graph, engines, backend_plans, modes=None: [PlanConfig(ea.engine_id, "A")] + list(backend_plans) + [PlanConfig(eb.engine_id, "B")] + ) - g = pygraph(router=Interleaved()) + g = pygraph() g.register_backend(ea).register_backend(eb) C = g.matmul(torch.randn(2, 2), torch.randn(2, 2)) g.create_execution_plans() @@ -455,50 +473,35 @@ def execute(self, graph, tensor_data, ctx=None): g.build_plans() -def test_empty_router_output_rejected(): - """A Router returning [] is an error — there is no legal empty planning +def test_empty_ranking_output_rejected(monkeypatch): + """Ranking that returns [] is an error — there is no legal empty planning state (it would defeat the one-shot flag and every needs-planning check).""" import cudnn - class Empty(Router): - def plan(self, graph, backends): - return [] + _ranking(monkeypatch, lambda graph, engines, backend_plans, modes=None: []) - g = pygraph(router=Empty()) + g = pygraph() g.matmul(torch.randn(2, 2), torch.randn(2, 2)) with pytest.raises(cudnn.cudnnGraphNotSupportedError, match="no engine"): g.create_execution_plans() - # the failed call did NOT consume the one-shot: fixing the router by + # the failed call did NOT consume the one-shot: fixing the ranking by # rebuilding the graph is the documented path, but the graph must not be # left half-planned either assert not g._planning_done -def test_set_router_frozen_after_planning(): - """set_router() after planning raises (it could not affect the already - planned list; accepting it silently would lie).""" - g = pygraph() - g.register_backend(_mk_engine(70)) - g.matmul(torch.randn(2, 2), torch.randn(2, 2)) - g.create_execution_plans() - with pytest.raises(RuntimeError, match="one-shot"): - g.set_router(Router()) - - -def test_plan_count_is_the_whole_ranked_list(): +def test_plan_count_is_the_whole_ranked_list(monkeypatch): """get_execution_plan_count() counts ONE list — python engines and backend engines alike (``graph.plans``), so a python-only graph reports its python plans instead of raising.""" eng = _mk_engine(71) - class PythonOnly(Router): - def plan(self, graph, backends): - from cudnn.engines import PlanConfig + from cudnn.engines import PlanConfig - return [PlanConfig(eng.engine_id)] # ``backends`` also holds the in-tree candidates + _ranking(monkeypatch, lambda graph, engines, backend_plans, modes=None: [PlanConfig(eng.engine_id)]) - g = pygraph(router=PythonOnly()) + g = pygraph() g.register_backend(eng) C = g.matmul(torch.randn(2, 2), torch.randn(2, 2)) g.create_execution_plans() @@ -508,9 +511,9 @@ def plan(self, graph, backends): g.execute({C: torch.empty(2, 2)}) # the routed python plan still runs -def test_constructor_backends_validated_and_proposals_checked(): - """Follow-up item 6: constructor path uses registration validation; foreign - engine ids in proposals are rejected.""" +def test_constructor_backends_validated_and_ranking_ids_checked(monkeypatch): + """Follow-up item 6: constructor path uses registration validation; a + ranking naming an id no engine owns is rejected.""" from cudnn.engines import PlanConfig class NoId(BaseEngine): @@ -520,19 +523,17 @@ def execute(self, graph, tensor_data, ctx=None): with pytest.raises(ValueError, match="engine_id"): pygraph(backends=[NoId()]) - class Impostor(BaseEngine): - name = "impostor" + class Ordinary(BaseEngine): + name = "ordinary" engine_id = _OOT + 63 - def propose_plans(self, graph): - return [PlanConfig(_OOT + 99, None)] # foreign id - def execute(self, graph, tensor_data, ctx=None): pass - g = pygraph(backends=[Impostor()]) + _ranking(monkeypatch, lambda graph, engines, backend_plans, modes=None: [PlanConfig(_OOT + 99, None)]) + g = pygraph(backends=[Ordinary()]) g.matmul(torch.randn(2, 2), torch.randn(2, 2)) - with pytest.raises(ValueError, match="foreign engine_id"): + with pytest.raises(ValueError, match="unknown python engine_id"): g.create_execution_plans() @@ -584,17 +585,6 @@ def _family(name): return next(f for f in MANIFEST if f.name == name) -def test_a_claiming_engine_is_tried_before_the_backend(): - """Ranking the python side second is how the frost job reported 0/3201 - graphs on FROST while passing. The opt-in is not a ranking concept: it - decides which engines are offered, in manifest.EngineFamily.matches.""" - from cudnn.engines import heuristics - from cudnn.engines.base import PlanConfig - - py, be = [PlanConfig(_OOT + 0, None)], [PlanConfig(0, {}, cpp_index=0)] - assert heuristics.heuristics_sort(None, py, be) == py + be - - def test_note_filters_reach_python_plans(monkeypatch): """The four classic note filters used to fall through to C++, so they filtered backend plans and silently skipped every python one. A python @@ -633,7 +623,7 @@ def _fresh(): assert idx not in g._barred_indices(), "engine declares no numerical notes" -def test_a_barred_note_advances_the_walk(): +def test_a_barred_note_advances_the_walk(monkeypatch): """Barring by note must change which plan RUNS, not just _barred_indices().""" import cudnn @@ -644,12 +634,10 @@ class Jitted(TorchMatmulEngine): engine_id = _OOT + 81 behavior_notes = (cudnn.behavior_note.RUNTIME_COMPILATION,) - class PythonFirst(Router): - def plan(self, graph, engines): - return self.python_plans(graph, engines) + [PlanConfig(0, {}, cpp_index=0)] + _ranking(monkeypatch, lambda graph, engines, backend_plans, modes=None: [PlanConfig(e.engine_id, None) for e in engines] + [PlanConfig(0, {}, cpp_index=0)]) def _fresh(): - g = pygraph(router=PythonFirst()) + g = pygraph() g.register_backend(Jitted()) C = g.matmul(torch.randn(2, 3), torch.randn(3, 2)) g._lowered_graph = _FakeBackend() @@ -719,16 +707,18 @@ def test_a_note_filter_set_before_planning_reaches_the_backend(): g.deselect_behavior_notes([cudnn.behavior_note.RUNTIME_COMPILATION]) assert fake.calls == [], "nothing to filter yet — the backend has no plans" g._create_backend_plans() - assert fake.calls == ["create_execution_plans", "deselect_behavior_notes"], fake.calls + # One create_execution_plans per heuristic mode (A, FALLBACK): that is how + # the backend's entries get tagged with the mode that produced them. + assert fake.calls == ["create_execution_plans", "create_execution_plans", "deselect_behavior_notes"], fake.calls -def test_cuda_graph_capture_declines_a_python_plan(): +def test_cuda_graph_capture_declines_a_python_plan(monkeypatch): """populate/update_cuda_graph record the BACKEND's plan. On a python plan they used to reach __getattr__ and report "graph not lowered yet", which points at the wrong thing.""" import cudnn - g, C = _backend_first() + g, C = _backend_first(monkeypatch) _pin(g, "torch_matmul") g.build_plans() for name in ("populate_cuda_graph", "update_cuda_graph"): @@ -983,10 +973,13 @@ def test_ranking_and_engine_read_the_same_record(monkeypatch): seen = {} - class Recording(Router): - def plan(self, graph, engines): - seen["ranking"] = graph._facts_for(_probe_analyzer) - return super().plan(graph, engines) + from cudnn.engines import heuristics + + real_rank = heuristics.rank + + def recording_rank(graph, engines, backend_plans, modes=None): + seen["ranking"] = graph._facts_for(_probe_analyzer) + return real_rank(graph, engines, backend_plans, modes) class Reader(TorchMatmulEngine): name = "reader" @@ -1002,7 +995,8 @@ def check_support(self, graph): monkeypatch.setattr(manifest, "MANIFEST", (family,)) monkeypatch.setattr(manifest, "_ANCHOR_NODE_TO_FAMILY", {"MATMUL": "probe_family"}) - g = pygraph(router=Recording(), backends=[Reader()]) + monkeypatch.setattr(heuristics, "rank", recording_rank) + g = pygraph(backends=[Reader()]) g.matmul(torch.randn(2, 3), torch.randn(3, 2)) g.create_execution_plans() @@ -1010,7 +1004,7 @@ def check_support(self, graph): assert len(_PROBE_CALLS) == 1, "one graph, one parse" -def test_facts_are_recomputed_when_the_graph_grows(): +def test_facts_are_recomputed_when_the_graph_grows(monkeypatch): """Facts describe the graph AS READ; a graph that gained a node since is a different graph.""" @@ -1076,15 +1070,13 @@ def select_numeric_notes(self, notes): return self -def _backend_first(*, check=None, build=None): +def _backend_first(monkeypatch, *, check=None, build=None): """A graph whose ranked list is [backend, python], with a scripted backend.""" from cudnn.engines.base import PlanConfig - class BackendFirst(Router): - def plan(self, graph, engines): - return [PlanConfig(0, {}, cpp_index=0)] + self.python_plans(graph, engines) + _ranking(monkeypatch, lambda graph, engines, backend_plans, modes=None: [PlanConfig(0, {}, cpp_index=0)] + [PlanConfig(e.engine_id, None) for e in engines]) - g = pygraph(router=BackendFirst()) + g = pygraph() g.register_backend(TorchMatmulEngine()) C = g.matmul(torch.randn(2, 3), torch.randn(3, 2)) g._lowered_graph = _FakeBackend(check=check, build=build) @@ -1092,14 +1084,16 @@ def plan(self, graph, engines): return g, C -def test_backend_check_support_decline_does_not_abort_the_walk(): +def test_backend_check_support_decline_does_not_abort_the_walk(monkeypatch): """An aggregate backend check_support() answers for the BACKEND, not for one plan. Letting it raise from build() aborted the walk before it reached a python entry that can serve the graph.""" import cudnn g, C = _backend_first( - check=cudnn.cudnnGraphNotSupportedError("backend cannot serve this graph"), build=cudnn.cudnnGraphNotSupportedError("backend cannot serve this graph") + monkeypatch, + check=cudnn.cudnnGraphNotSupportedError("backend cannot serve this graph"), + build=cudnn.cudnnGraphNotSupportedError("backend cannot serve this graph"), ) g.build() assert g.selected_engine is not None and g.selected_engine.name == "torch_matmul" @@ -1135,11 +1129,9 @@ def build_plan(self, graph, plan, ctx=None): from cudnn.engines.base import PlanConfig - class BackendFirst(Router): - def plan(self, graph, engines): - return [PlanConfig(0, {}, cpp_index=0)] + self.python_plans(graph, engines) + _ranking(monkeypatch, lambda graph, engines, backend_plans, modes=None: [PlanConfig(0, {}, cpp_index=0)] + [PlanConfig(e.engine_id, None) for e in engines]) - g = pygraph(router=BackendFirst()) + g = pygraph() g.register_backend(Recording()) C = g.matmul(torch.randn(2, 3), torch.randn(3, 2)) g._lowered_graph = _FakeBackend(build=cudnn.cudnnGraphNotSupportedError("backend build declined")) @@ -1168,10 +1160,10 @@ def boom(): g.backend_plan_entries() -def test_pinned_plan_that_was_deselected_raises(): +def test_pinned_plan_that_was_deselected_raises(monkeypatch): """select_plan() and deselect_engines() contradicting each other is a caller error, not a licence to run a third plan.""" - g, C = _backend_first() + g, C = _backend_first(monkeypatch) g.create_execution_plans() idx = _index_of(g, "torch_matmul") g.select_plan(idx) @@ -1275,7 +1267,7 @@ def owns_id(self, engine_id): g._engine_for(PlanConfig(family.engine_id, None)) -def test_replayed_backend_entry_addresses_the_plan_it_replayed(): +def test_replayed_backend_entry_addresses_the_plan_it_replayed(monkeypatch): """``create_execution_plan`` APPENDS in C++ and ``build_plans`` short-circuits once a candidate exists, so a replayed entry must be addressed by the index it landed at — the plain calls would run whichever plan the backend already had.""" @@ -1306,11 +1298,9 @@ def get_workspace_size_plan_at_index(self, i, *a, **k): from cudnn.engines.base import PlanConfig - class ReplayRouter(Router): - def plan(self, graph, engines): - return [PlanConfig(0, {}, cpp_index=0), PlanConfig(7, {})] # the 2nd is a replay + _ranking(monkeypatch, lambda graph, engines, backend_plans, modes=None: [PlanConfig(0, {}, cpp_index=0), PlanConfig(7, {})]) # the 2nd is a replay - g = pygraph(router=ReplayRouter()) + g = pygraph() g.matmul(torch.randn(2, 3), torch.randn(3, 2)) be = FakeAppendingBackend() g._lowered_graph = be @@ -1380,7 +1370,7 @@ def declines(): def test_the_router_output_is_the_plan_list_position_for_position(monkeypatch): - """Nothing rewrites what a Router returned. The delegating entry that + """Nothing rewrites what the ranking returned. The delegating entry that backend_plan_entries() appends under heur_mode.OPENSOURCE used to share its id with a 'the backend goes here' placeholder the frontend expanded, so planning spliced the backend's whole list in a second time — wrong count, @@ -1431,7 +1421,7 @@ def get_engine_count(self): def test_backend_entries_are_queried_once_per_graph(): """A second C++ create_execution_plans() APPENDS to the same plan list, so - asking the backend twice reports every plan twice. The Router and the marker + asking the backend twice reports every plan twice. The ranking and the marker expansion both ask.""" g = pygraph() g.register_backend(TorchMatmulEngine()) @@ -1440,12 +1430,12 @@ def test_backend_entries_are_queried_once_per_graph(): assert g.backend_plan_entries() is first -def test_at_index_build_honours_the_exclusions(): +def test_at_index_build_honours_the_exclusions(monkeypatch): """deselect_engines() / deselect_workspace_greater_than() are properties of the plan, not of the walk: the list is never filtered (indices stay stable), so build_plan_at_index() has to apply them too or it compiles and selects a plan the caller excluded.""" - g, C = _backend_first() + g, C = _backend_first(monkeypatch) g.create_execution_plans() idx = _index_of(g, "torch_matmul") g.deselect_engines(["torch_matmul"]) @@ -1475,10 +1465,10 @@ def build_plan(self, graph, plan, ctx=None): g.build_plan_at_index(idx) -def test_behaviour_notes_reject_a_negative_index(): +def test_behaviour_notes_reject_a_negative_index(monkeypatch): """Python indexing would quietly answer for the LAST plan; every other at-index API rejects it.""" - g, C = _backend_first() + g, C = _backend_first(monkeypatch) g.create_execution_plans() with pytest.raises(IndexError, match="out of range"): g.get_behavior_notes_for_plan_at_index(-1) @@ -1519,11 +1509,11 @@ def test_frost_opt_in_does_not_leak_out_of_the_frost_suites(): assert not manifest.opt_in_engines_enabled(), "CUDNN_FRONTEND_ENABLE_FROST_ENGINES leaked into the default-path tests" -def test_at_index_queries_answer_from_the_unified_list(): +def test_at_index_queries_answer_from_the_unified_list(monkeypatch): """get_engine_and_knobs_at_index() must describe the plan the caller just saw at that index — forwarding a unified index to C++ reported the backend's entry for a python plan's slot.""" - g, C = _backend_first() + g, C = _backend_first(monkeypatch) g.create_execution_plans() idx = _index_of(g, "torch_matmul") eid, knobs = g.get_engine_and_knobs_at_index(idx) @@ -1532,11 +1522,11 @@ def test_at_index_queries_answer_from_the_unified_list(): assert g.get_behavior_notes_for_plan_at_index(idx) == [] # engine declares none -def test_create_execution_plan_appends_a_python_plan(): +def test_create_execution_plan_appends_a_python_plan(monkeypatch): """The deterministic-replay idiom: record (engine_id, knobs), rebuild it later, address it with count-1. It has to work for a python engine id too, or 'one id space' is only true for the backend.""" - g, C = _backend_first() + g, C = _backend_first(monkeypatch) g.create_execution_plans() before = g.get_execution_plan_count() g.create_execution_plan(TorchMatmulEngine.engine_id, None) diff --git a/test/python/test_native_backend_lowering.py b/test/python/test_native_backend_lowering.py index 5663e79d8..94ab0a20f 100644 --- a/test/python/test_native_backend_lowering.py +++ b/test/python/test_native_backend_lowering.py @@ -421,12 +421,12 @@ def test_native_rmsnorm_lowers_to_backend(): torch.testing.assert_close(ivb, ivref, atol=5e-3, rtol=5e-3) -def test_mixed_router_backend_slot_executes(): - """Review round 4: the backend entry of a MIXED router is selectable and +def test_mixed_ranking_backend_slot_executes(monkeypatch): + """Review round 4: the backend entry of a MIXED ranking is selectable and actually executes through the backend (lowering triggered), with routed indices stable across that lowering; the pinned python plan still runs afterwards with its own knobs.""" - from cudnn.engines import BaseEngine, PlanConfig, Router, is_backend_engine + from cudnn.engines import BaseEngine, PlanConfig, heuristics, is_backend_engine from cudnn.engines.engine_ids import OUT_OF_TREE_ID_BASE ran = [] @@ -448,11 +448,9 @@ def execute(self, graph, uid_to_data, ctx=None): py_engine = PyMatmul() - class CudnnFirst(Router): - def plan(self, graph, backends): - # ``backends`` also carries the in-tree manifest candidates, so name - # the engine this test means instead of taking the first one. - return graph.backend_plan_entries() + [PlanConfig(py_engine.engine_id)] + # ``engines`` also carries the in-tree manifest candidates, so name the + # engine this test means instead of taking the first one. + monkeypatch.setattr(heuristics, "rank", lambda graph, engines, backend_plans, modes=None: list(backend_plans) + [PlanConfig(py_engine.engine_id)]) h = _handle() a = torch.randn(1, M, K, device="cuda", dtype=torch.float16) @@ -460,9 +458,7 @@ def plan(self, graph, backends): c = torch.empty(1, M, N, device="cuda", dtype=torch.float16) ref = (a.float() @ b.float()).half() - g = pygraph( - handle=h, io_data_type=cudnn.data_type.HALF, intermediate_data_type=cudnn.data_type.FLOAT, compute_data_type=cudnn.data_type.FLOAT, router=CudnnFirst() - ) + g = pygraph(handle=h, io_data_type=cudnn.data_type.HALF, intermediate_data_type=cudnn.data_type.FLOAT, compute_data_type=cudnn.data_type.FLOAT) g.register_backend(py_engine) A = g.tensor(dim=[1, M, K], stride=[M * K, K, 1]) B = g.tensor(dim=[1, K, N], stride=[K * N, N, 1])