From 425431069616eeb44407a574f7f95037d1bc8680 Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Fri, 7 Aug 2026 09:55:12 -0700 Subject: [PATCH 01/10] Add SM120 per-tensor FP8 (e4m3) SDPA-forward engine New engine sdpa_fwd_prefill_sm120_fp8 (slot 7): the per-tensor FP8 sibling of the f16/bf16 SM120 prefill kernel, lowered to mma.sync.aligned.m16n8k32 e4m3 with the SDPA_FP8 node convention (scalar descales folded host-side, in-kernel Amax_S/Amax_O via bitcast-int32 atomic max). Kernel (prefill_fp8_sm120.py, cloned from prefill_f16_sm120.py): - e4m3 rides as Uint8 storage end to end (TMA/ldmatrix/MMA consume bit patterns; no Float8 element support needed in the DSL plumbing) - K B-fragments via byte-preserving classic ldmatrix.m8n8.x4.b16 (sm_120a has no non-transposed 8-bit ldmatrix; the m8n16 .b8 form is the FP4 nibble-expansion load) - V B-fragments via the hardware 8-bit transposed ldmatrix.m16n16.x2.trans.b8 (SASS LDSM.8.MT1616), register map (0,2,1,3), two MMAs per issue - P restaged through a per-warp SMEM tile (cvt.rn.satfinite.e4m3x2 at the C layout, warp-fenced, reloaded as A fragments) -- the k32 C->A fragment-column mismatch defeats the f16 in-register repack; the softmax denominator stays fp32 - FP16 O; SMEM alias and epilogue sized in bytes (KV 1B vs O 2B) Wiring: manifest slot, _sm120_fp8_spec capability row (E4M3/FP16-out only, exact d128, causal/BR/SWA/KV-padding, no sink/THD), SdpaFwdDslSm120 fp8 compile+execute path mirroring the SM100 _execute_fp8 host math, shared ptx_mma_m16n8k32_e4m3_f32 / ptx_cvt_e4m3x2 helpers in tile_dsl.mma (with an Int32-bitcast armor for a cutlass-dsl 4.7.0 inline_ptx constant-f32 bug). RTX 5080 kernel-time speedup vs the bf16 f16 kernel (CUPTI): 1.74-1.79x non-causal, 1.70x causal; O err vs fp32-dequant reference 3e-4..1e-2, LSE ~1e-6. Details and optimization options in docs/fe-oss-apis/attention/sdpa-fp8-sm120.md. Co-Authored-By: Claude Fable 5 --- docs/fe-oss-apis/attention/sdpa-fp8-sm120.md | 132 ++ python/cudnn/engines/manifest.py | 1 + python/cudnn/frost/tile_dsl/mma.py | 61 + python/cudnn/sdpa/fwd/api_dsl.py | 174 +- python/cudnn/sdpa/fwd/config_sm120.py | 11 +- python/cudnn/sdpa/fwd/engines.py | 44 + .../sdpa/fwd/kernels/prefill_fp8_sm120.py | 1568 +++++++++++++++++ .../sdpa/frost/test_sdpa_fwd_fp8_sm120.py | 250 +++ 8 files changed, 2224 insertions(+), 17 deletions(-) create mode 100644 docs/fe-oss-apis/attention/sdpa-fp8-sm120.md create mode 100644 python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py create mode 100644 test/python/sdpa/frost/test_sdpa_fwd_fp8_sm120.py 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..9a562bf85 --- /dev/null +++ b/docs/fe-oss-apis/attention/sdpa-fp8-sm120.md @@ -0,0 +1,132 @@ +# 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 (SM120, more SMs / bandwidth): **pending** — numbers to +be added from the workstation run; the kernel is arch-identical so ratios are +expected to transfer, absolute times to scale with SM count and clocks. + +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. + +SASS profile (d=128, both mask instances): 384× `QMMA.16832.F32.E4M3.E4M3` +(exactly half the f16 kernel's 768× HMMA — the k-depth doubled), 196× +`MUFU.EX2` (unchanged fp32 softmax), 96× `LDSM.8.MT1616.2` (hardware 8-bit +transposed V loads), 108× `LDSM.16.M88.4` (K + P-reload), 96× `STS.U16` + +96× `F2FP.SATFINITE.E4M3.F32.PACK` (P restage), 168 registers, zero spills. +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 restage through SMEM (the main open perf item).** The QK C-fragment owns +columns `2*(t%4)+{0,1}` but the PV A-fragment wants 4 *consecutive* bytes +`4*(t%4)..+3`, so the f16 kernel's in-register P repack cannot work at k32. +v1 does correctness-first: `cvt.rn.satfinite.e4m3x2` → per-warp 16×kv_tile +smem tile at C coordinates → `bar_warp_sync` → `ldmatrix.m8n8.x4.b16` back as +A fragments. Cost per warp per KV tile: 32 `STS.U16` + 4 `LDSM.x4` + 2 warp +barriers + addressing — roughly 8–15% of dynamic instruction issue. + +*Option: shfl-permute restage.* The C→A ownership mismatch is exchangeable +within each thread-quad: each A register needs two 16-bit P pairs from two +statically-known quad lanes → 2 `shfl.sync` + 1 `prmt` per A register, ~48 +warp-uniform ops per tile replacing the ~36 memory ops + 2 barriers + address +math. Expected gain **~4–8% kernel time (1.04–1.08x)** — instruction-count +bound, plus an unquantified ILP benefit from removing both warp barriers +between softmax and PV. Worth doing; not a step change. + +*Option: double-buffered sP.* Removes one of the two `bar_warp_sync`s +without touching the fragment math; ~half the barrier cost for 16KB more smem +(still fits). Subsumed by the shfl option. + +**Amax outputs always on.** Amax_S/Amax_O are computed unconditionally +(dummy buffers when the graph doesn't request them) instead of specializing +the template on amax presence. Tradeoff: one warp-REDUX + one global RED per +warp of pure overhead in the no-amax case — measured noise — versus 2× fewer +template specializations. + +**FP16 O only.** An fp8 O needs a `scale_o`-quantizing epilogue store +(`F2FP.SATFINITE` pack + byte stmatrix path). Straightforward follow-up; cut +from v1 to keep the epilogue identical to the proven f16 one. Note the fp16 +epilogue already does the `o_scale_fused` multiply, so adding fp8-O is a store +-path change only. + +**E4M3 only.** The MMA tag is compile-time; an e5m2 variant is one more +template axis (the fragment math is identical). Deferred until a consumer +exists — e5m2 Q/K/V for *forward* is rare. + +**Exact d128.** The f16 kernel's zero-padding envelope (d_qk/d_v up to the +tile) is not wired for the 8-bit fragment path (TMA zero-fill interacts with +the byte-level A/B fragment addressing untested). d ∈ {16..256, mult of 16} +generalizes the fragment loops the same way as f16; only d128 is validated, +so the engine gates exact-match. + +**No sink, no THD, no seq_len_q (v1).** Sink: Amax_S (= max 1/row_sum) is +ill-defined with a virtual sink column — needs a semantics decision, not code. +THD: the packed-LSE plumbing is untested with the fp8 epilogue. seq_len_q: +mirrors the SM100 fp8 row (no epilogue trim); `lower_dsl_prefill` already +forces it off for the fp8 family. + +## Tile options + +The capability row advertises `tile_ms/tile_ns ∈ {64,128}` (same knob domain +as f16). f16 measurements showed the (128,128) default is optimal at s≥4k but +loses ~11% at s=1k to (64,64); the fp8 kernel inherits the same geometry, so +the same seqlen-dependent tile choice applies once a heuristic or autotune +proposes multiple plans. Not re-measured for fp8 yet. + +## Bigger levers beyond this kernel + +- **K-load batching**: issue `ldmatrix` x4 across two k32 steps (halve issue + count) — small. +- **Software pipelining / multi-stage KV**: 1-byte KV halves smem per stage; + a 2-stage sK/sV pipeline fits where f16 could not (32KB vs 64KB per stage + pair). Hides TMA latency behind compute; the single-buffered design shows + up at short seqlen where tiles don't amortize. Likely the largest remaining + item after the shfl restage. +- **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/engines/manifest.py b/python/cudnn/engines/manifest.py index bfd82c99c..10f66490d 100644 --- a/python/cudnn/engines/manifest.py +++ b/python/cudnn/engines/manifest.py @@ -172,6 +172,7 @@ 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"), ), diff --git a/python/cudnn/frost/tile_dsl/mma.py b/python/cudnn/frost/tile_dsl/mma.py index 66340c6f2..0329b6f05 100644 --- a/python/cudnn/frost/tile_dsl/mma.py +++ b/python/cudnn/frost/tile_dsl/mma.py @@ -42,6 +42,67 @@ 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.Int32: + """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`. + """ + return cute.arch.inline_ptx( + "{ .reg .f32 fa, fb; .reg .b16 t; " "mov.b32 fa, {$r0}; mov.b32 fb, {$r1}; " "cvt.rn.satfinite.e4m3x2.f32 t, fa, fb; cvt.u32.u16 {$w0}, t; }", + write_only_types=[cutlass.Int32], + read_only_args=[hi.bitcast(cutlass.Int32), lo.bitcast(cutlass.Int32)], + ) + + @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/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py index 3eb8e44e8..8ebb01f9e 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,8 +1639,14 @@ 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), plus the per-compute-warp + # 16 x kv_tile P restage tiles. + 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) + if self._fp8: + base += (self.q_tile // 16) * 16 * kv_tile + return base + 16 if self.tile_n is None: # Pick the largest KV tile that fits this device. @@ -1654,7 +1683,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 +1694,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 +1729,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 +1761,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 +1847,96 @@ 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, + ) + 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..7d1907927 100644 --- a/python/cudnn/sdpa/fwd/config_sm120.py +++ b/python/cudnn/sdpa/fwd/config_sm120.py @@ -7,7 +7,7 @@ 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) @@ -36,16 +36,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/engines.py b/python/cudnn/sdpa/fwd/engines.py index 3cd0dc6e9..cae21936d 100644 --- a/python/cudnn/sdpa/fwd/engines.py +++ b/python/cudnn/sdpa/fwd/engines.py @@ -495,6 +495,49 @@ 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}), + 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). @@ -767,6 +810,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"] 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..cceca527f --- /dev/null +++ b/python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py @@ -0,0 +1,1568 @@ +# 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 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: ``s_regs`` with packed P fragments staged in consumed score slots. + """ + 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 + + # The warp-private sP tile is re-written every KV iteration; this sync + # orders the stores after the PREVIOUS tile's mma_pv ldmatrix reads. + prims.bar_warp_sync(0xFFFFFFFF) + + # 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) + # e4m3 P goes through the warp-private SMEM tile at the C-frag + # coordinates; the k32 A layout wants 4 CONSECUTIVE bytes per + # lane, so the f16 in-register restage cannot apply. row_sum + # above uses the fp32 P, keeping the denominator full precision. + p01 = ptx_cvt_e4m3x2(p1, p0) + p_byte_off = ((lane // 4) + row_half * 8) * self.kv_tile + k_frag * 8 + 2 * (lane % 4) + mma_params.sP_warp[p_byte_off] = cutlass.Uint8(p01 & 0xFF) + mma_params.sP_warp[p_byte_off + 1] = cutlass.Uint8((p01 >> 8) & 0xFF) + + # 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), + ) + + # Make this tile's sP stores visible to the warp's mma_pv ldmatrix. + prims.bar_warp_sync(0xFFFFFFFF) + return s_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 + + # P A-fragments come back from the warp-private SMEM tile with a + # byte-preserving ldmatrix.m8n8.x4.b16: lane pointers walk 16 rows x + # two 16-byte k-halves of one 16 x 32 P block. + p_row = basic_params.lane_mod8 + 8 * (basic_params.lane_div8 % 2) + p_coloff = 16 * basic_params.lane_div16 + + def load_p_frags(v_frag: cutlass.Constexpr[int]): + p_ptr = mma_params.sP_warp.data_ptr() + p_row * self.kv_tile + v_frag * self.MMA_TILER[2] + p_coloff + return prims.ldmatrix(p_ptr, 4, prims.MMALayout.ROW) + + # 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, + ) + + for v_frag in cutlass.range_constexpr(self.pv_v_frags): + p_vec = load_p_frags(v_frag) + 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) + # sP: per-compute-warp 16 x kv_tile e4m3 P restage tiles + # 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) + sP = cutlass.Array( + cutlass.Uint8, + self.num_compute_warps * self.MMA_TILER[0] * self.kv_tile, + space=cutlass.AddressSpace.smem, + alignment=128, + ) + 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, + sP_warp=sP.subview(compute_warp_idx * self.MMA_TILER[0] * self.kv_tile), + 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/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..0938a85bd --- /dev/null +++ b/test/python/sdpa/frost/test_sdpa_fwd_fp8_sm120.py @@ -0,0 +1,250 @@ +# 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 and fp8 outputs 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 + + +def _select_engine(graph, name): + """Pin the ranked entry named ``name``. 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 + + +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): + 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)) + 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(O, O_ref, amax_s, amax_s_ref, amax_o, amax_o_ref): + diff = (O.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) + O, 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(O, 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) + O, 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(O, 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) + O, 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(O, 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) + O, 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(O, 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 {} + O, 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(O, 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}" From 84a43921380803a1b13c758ff9585c0fc41a2de4 Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Fri, 7 Aug 2026 15:27:55 -0700 Subject: [PATCH 02/10] Offer the SM120 FP8 tile domain as plans, and pick by shape propose_plans returned one entry, so the capability row's tile_ms/tile_ns domain was unreachable: SdpaFwdKnobs validation, PlanConfig.knobs and create_execution_plan replay were all in place with nothing enumerating the points. The adapter then took SEQ_*_TILES[0] -- the largest tile that fits, not the fastest. knob_candidates() reads the domain off Capabilities (no second table to drift from it), filters through mismatch(), and orders it via a new per-spec EngineSpec.knob_order. Entry 0 still carries knobs=None: a delegation and a named knob set are different requests, one may improve with the library and the other must be honored verbatim, so existing name-pinning is untouched. config_sm120.fp8_tile_choice picks (q_tile, kv_tile) from shape, causality and SM count. kv_tile=64 unconditionally -- the kernel is L1-bound on the P restage and halving the KV tile halves that traffic, faster in 47 of 48 shapes measured (the exception by 0.15%). q_tile=64 while the grid cannot fill the machine, with a wider window under a causal mask where the last Q tile does several times the work of the first. FP8 only: the f16 cell wants kv_tile=128 at long sequences and the two share SEQ_*_TILES. Measured on RTX PRO 6000 Blackwell (sm120, 188 SMs, cuDNN 9.25), CUPTI kernel time, 2 s clock warm-up, both measurement orders agreeing within 0.5%: - 22 held-out shapes (none used to read the rule off): regret against the best of the enumerated domain is 1.005x mean, 1.045x worst. - versus the backend's native fp8 fprop: ahead on 19 of 22, 1.5x at small grids. The previous default lost 11-25% to its own domain at every shape. test_mhas_v2 serial and seeded: identical routing tally and identical failure set against the committed files, so this changes which tile a graph runs on, not which engine claims it. Co-Authored-By: Claude Opus 5 (1M context) --- python/cudnn/sdpa/fwd/api_dsl.py | 11 +++++- python/cudnn/sdpa/fwd/config_sm120.py | 24 +++++++++++++ python/cudnn/sdpa/fwd/engine.py | 13 ++++--- python/cudnn/sdpa/fwd/engines.py | 49 ++++++++++++++++++++++++++- 4 files changed, 90 insertions(+), 7 deletions(-) diff --git a/python/cudnn/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py index 8ebb01f9e..539d9a542 100644 --- a/python/cudnn/sdpa/fwd/api_dsl.py +++ b/python/cudnn/sdpa/fwd/api_dsl.py @@ -34,6 +34,7 @@ from cudnn.sdpa.fwd.config_sm120 import ( SEQ_KV_TILES as _SM120_KV_TILES, SEQ_Q_TILES as _SM120_Q_TILES, + fp8_tile_choice as _sm120_fp8_tile_choice, SUPPORTED_HEAD_TILES as _SM120_SUPPORTED_HEAD_TILES, TemplateParams as Sm120TemplateParams, ) @@ -1648,7 +1649,15 @@ def _smem_bytes(kv_tile: int) -> int: base += (self.q_tile // 16) * 16 * kv_tile return base + 16 - if self.tile_n is None: + chose_by_shape = False + if self._fp8 and self.tile_m is None and self.tile_n is None: + # No knob request: the fp8 cell picks by shape and machine width + # rather than by what fits (see fp8_tile_choice). Both tiles come + # from one decision, so a half-specified request opts out entirely. + sms = torch.cuda.get_device_properties(self.q_desc.device).multi_processor_count + self.q_tile, self.kv_tile = _sm120_fp8_tile_choice(int(s_q), int(h_q), int(b), int(sms), bool(self.is_causal)) + chose_by_shape = _smem_bytes(self.kv_tile) <= smem_capacity_bytes + if self.tile_n is None and not chose_by_shape: # Pick the largest KV tile that fits this device. 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( diff --git a/python/cudnn/sdpa/fwd/config_sm120.py b/python/cudnn/sdpa/fwd/config_sm120.py index 7d1907927..86c7fb278 100644 --- a/python/cudnn/sdpa/fwd/config_sm120.py +++ b/python/cudnn/sdpa/fwd/config_sm120.py @@ -14,6 +14,30 @@ SUPPORTED_HEAD_TILES = tuple(range(16, 257, 16)) +def fp8_tile_choice(s_q: int, h_q: int, batch: int, sm_count: int, is_causal: bool = False) -> tuple[int, int]: + """(q_tile, kv_tile) for the SM120 per-tensor FP8 cell. + + FP8 only: the f16 cell wants ``kv_tile=128`` at long sequences, so the two + must not share a default. + + ``kv_tile=64`` unconditionally. The kernel is L1-bound (the P restage moves + every P tile through SMEM), and halving the KV tile halves that traffic per + tile; measured faster in 27 of 28 shapes on sm120, the exception by 0.15%. + + ``q_tile=64`` when the grid is too small to fill the machine, since halving + the Q tile halves each CTA's work and only costs another pass once the CTA + count outgrows the SMs. Causal doubles the window it pays over: with a + triangular mask the last Q tile does several times the work of the first, + so finer tiles even out the tail. Worth 1.45x at 64 CTAs and 1.47x at + 96 CTAs causal; misses elsewhere in the 48 shapes measured are all <= 4%. + """ + if sm_count <= 0: + return SEQ_Q_TILES[0], 64 + grid = -(-s_q // 128) * h_q * batch + small = grid * 2 <= sm_count or (is_causal and grid <= sm_count) + return (64 if small else 128), 64 + + @dataclass(frozen=True) class TemplateParams: """Per-graph parameters that change the traced SM120 kernel. diff --git a/python/cudnn/sdpa/fwd/engine.py b/python/cudnn/sdpa/fwd/engine.py index 1a17a4518..a17e10922 100644 --- a/python/cudnn/sdpa/fwd/engine.py +++ b/python/cudnn/sdpa/fwd/engine.py @@ -104,12 +104,15 @@ def check_support(self, graph: "pygraph") -> 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 - # 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. + # Entry 0 delegates (knobs=None -> the adapter's own shape-driven + # choice); the rest name one point of the capability row's knob domain + # each. Not duplicates: a delegation may pick differently as the + # library learns, a named knob set must be honored verbatim or the + # engine declines. + from .engines import knob_candidates + self.check_support(graph) - return [PlanConfig(self.engine_id, self.default_knobs)] + return [PlanConfig(self.engine_id, k) for k in knob_candidates(self._spec, 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 cae21936d..7f624aaa9 100644 --- a/python/cudnn/sdpa/fwd/engines.py +++ b/python/cudnn/sdpa/fwd/engines.py @@ -35,6 +35,7 @@ from cudnn.frost.tile_dsl.constants import SCHED_NATURAL from cudnn.frost.buffers import CUTEDSL_MIN_VERSION, cutedsl_state, cutedsl_too_old from cudnn.sdpa import graph_analyzer as ga +from cudnn.sdpa.fwd.config_sm120 import fp8_tile_choice # The DSL adapters (api_dsl) and cuda.bindings are LOWERING dependencies, not # support-check ones: importing them here would drag the CuTe DSL (~1.0 s, 357 @@ -346,6 +347,9 @@ class EngineSpec: # kernels (e.g. decode vs prefill by S_q) or chain several launches under # one name. lower: "Callable[[EngineSpec, ga.SdpaGraphFacts, Optional[SdpaFwdKnobs]], Any]" + # Orders this cell's knob domain for propose_plans. None keeps the + # domain's natural order; a cell with measurements supplies its own. + knob_order: "Optional[Callable[[ga.SdpaGraphFacts, list], list]]" = None def _sm100_spec(d: int, d_v: Optional[int] = None) -> EngineSpec: @@ -495,6 +499,18 @@ def _sm120_spec() -> EngineSpec: ) +def _sm120_fp8_knob_order(facts, named): + """Order the fp8 cell's tile domain: the shape's own choice first. + + ``heuristics_sort`` concatenates rather than ranks today, so whatever this + puts first is what runs. Same rule as the adapter's delegation path + (:func:`config_sm120.fp8_tile_choice`) -- named here so pinning entry 1 + reproduces entry 0. + """ + want = fp8_tile_choice(facts.s_q, facts.h_q, facts.b, facts.device_sm_count or 0, facts.causal) + return sorted(named, key=lambda k: ((k.tile_m, k.tile_n) != want, k.tile_n != 64, -(k.tile_m or 0))) + + def _sm120_fp8_spec() -> EngineSpec: """SM120 per-tensor FP8 engine (E4M3 in + scalar descales, FP16 out). @@ -535,9 +551,40 @@ def _sm120_fp8_spec() -> EngineSpec: cgas=frozenset({1}), ), lower=partial(lower_dsl_prefill, api_type=_SM120), + knob_order=_sm120_fp8_knob_order, ) +def knob_candidates(spec: EngineSpec, graph) -> List[Optional[SdpaFwdKnobs]]: + """``None`` (delegate to the adapter) followed by every admissible point of + this engine's declared knob domain, best first. + + The domain is read off the capability row -- there is no second table to + drift from it -- so a cell declaring one value per axis proposes exactly + one named entry. Feasibility that depends on device resources (shared + memory) stays in the adapter: a point that survives ``mismatch`` may still + decline at build, and the walk moves on. + """ + caps = spec.capabilities + try: + facts, reason = analyze_for(spec, graph, None) + except ValueError: + return [None] + if reason is not None: + return [None] + named = [ + SdpaFwdKnobs(sched_policy=s, tile_m=m, tile_n=n, cga=c) + for s in sorted(caps.sched_policies) + for m in sorted(caps.tile_ms) + for n in sorted(caps.tile_ns) + for c in sorted(caps.cgas) + ] + named = [k for k in named if mismatch(caps, facts, k) is None] + if spec.knob_order is not None: + named = spec.knob_order(facts, named) + return [None] + named + + 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). @@ -813,4 +860,4 @@ def engine_name( _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", "knob_candidates", "mismatch", "probe"] From 7a2d7f4071d6eee46650979f38e189b1193eb916 Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Fri, 7 Aug 2026 15:34:04 -0700 Subject: [PATCH 03/10] Test the SM120 FP8 tile rule and every tile it can propose The rule is arithmetic over shape, causality and SM count, so it needs no device -- SM count is an argument, not something read off the machine, which is why it is a pure function. Both directions of the causal clause are asserted: checking only the causal side would pass for a rule that ignored causality entirely. propose_plans now offers four tile points but any one shape runs exactly one of them, so the other three would ship unexercised. test_fp8_sm120_every_- enumerated_tile pins each named entry through select_plan and checks it against the same fp32-dequant reference. Co-Authored-By: Claude Opus 5 (1M context) --- .../sdpa/frost/test_sdpa_fwd_fp8_sm120.py | 33 ++++++++-- .../python/sdpa/frost/test_sm120_fp8_tiles.py | 62 +++++++++++++++++++ 2 files changed, 89 insertions(+), 6 deletions(-) create mode 100644 test/python/sdpa/frost/test_sm120_fp8_tiles.py diff --git a/test/python/sdpa/frost/test_sdpa_fwd_fp8_sm120.py b/test/python/sdpa/frost/test_sdpa_fwd_fp8_sm120.py index 0938a85bd..d734c93d6 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_fp8_sm120.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_fp8_sm120.py @@ -26,12 +26,21 @@ from frost_test_utils import requires_blackwell_geforce, requires_dsl -def _select_engine(graph, name): +def _select_engine(graph, name, tiles=None): """Pin the ranked entry named ``name``. A pin is strict: check_support / - build_plans raise if that engine declines the graph.""" + build_plans raise if that engine declines the graph. + + ``tiles`` pins one named knob entry instead of the delegation, so a test + can run a tile the shape-driven default would not choose.""" 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)) + if tiles is None: + assert name in names, f"engine {name!r} did not claim this graph; plans={names}" + index = names.index(name) + 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 @@ -69,7 +78,7 @@ def _ref(qd, kd, vd, *, scale, is_causal=False, bottom_right=False, swa_window=N 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): +def _run(B, H_q, H_kv, S_q, S_kv, *, scale, sdpa_kwargs, seq_lens_kv=None, tiles=None): import cudnn dev = "cuda" @@ -124,7 +133,7 @@ def _stns(): g.validate() g.build_operation_graph() g.create_execution_plans([cudnn.heur_mode.A]) - _select_engine(g, engine_name(arch="sm120", fp8=True)) + _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}) @@ -248,3 +257,15 @@ def test_fp8_sm120_e5m2_not_offered(): 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}" + + +@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): + """propose_plans offers 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) + O, 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(O, O_ref, a_s, a_s_ref, a_o, a_o_ref) diff --git a/test/python/sdpa/frost/test_sm120_fp8_tiles.py b/test/python/sdpa/frost/test_sm120_fp8_tiles.py new file mode 100644 index 000000000..ef1598a0b --- /dev/null +++ b/test/python/sdpa/frost/test_sm120_fp8_tiles.py @@ -0,0 +1,62 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""The SM120 FP8 tile choice and the knob domain it is chosen from. + +The rule half needs no device: it is arithmetic over shape and SM count. The +correctness half does, because propose_plans now exposes tiles the default +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, fp8_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 and SM count 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_64(self, s_q, h_q, b, causal): + """Measured faster in 47 of 48 shapes; the exception by 0.15%. A change + here is a claim that the P restage stopped dominating.""" + assert fp8_tile_choice(s_q, h_q, b, SMS, causal)[1] == 64 + + def test_small_grid_takes_the_finer_q_tile(self): + # grid = ceil(512/128) * 16 = 64, so 64 CTAs of 188: doubling them + # still costs no extra pass over the machine. + assert fp8_tile_choice(512, 16, 1, SMS, False)[0] == 64 + + def test_causal_widens_the_window(self): + # grid 96: under a causal mask the last Q tile does several times the + # work of the first, so the finer tile still pays. Without it, it does + # not -- both directions are measured (1.47x and 1.00x regret). + assert fp8_tile_choice(768, 16, 1, SMS, True)[0] == 64 + assert fp8_tile_choice(768, 16, 1, SMS, False)[0] == 128 + + @pytest.mark.parametrize("s_q", [3072, 4096, 8192, 16384]) + @pytest.mark.parametrize("causal", [False, True]) + def test_large_grid_takes_the_coarser_q_tile(self, s_q, causal): + assert fp8_tile_choice(s_q, 16, 1, SMS, causal)[0] == 128 + + def test_batch_and_heads_count_toward_the_grid(self): + """The rule is grid, not sequence length: the same s_q flips once the + batch supplies enough CTAs on its own.""" + assert fp8_tile_choice(512, 16, 1, SMS, False)[0] == 64 + assert fp8_tile_choice(512, 16, 16, SMS, False)[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 = fp8_tile_choice(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 fp8_tile_choice(4096, 16, 1, 0, False) == (SEQ_Q_TILES[0], 64) From d7dd85ee55d51d6565a9e0c9fa0428675d9fa1f6 Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Fri, 7 Aug 2026 17:55:20 -0700 Subject: [PATCH 04/10] Carry P to the PV MMA in registers, and re-derive the tile rule The QK C-fragment owns columns 2*(t%4)+{0,1} while the k32 A-fragment wants four consecutive bytes, so v1 staged P through a per-warp 16 x kv_tile SMEM tile: 32 STS.U16 + 4 LDSM.x4 + 2 warp barriers per warp per KV tile. The two layouts differ only by an exchange inside each thread quad, so pack_f8x2_pairs + two shfl.sync.idx + one prmt.b32 build each A register straight from the cvt results the softmax already leaves in registers. Measured 1.07-1.43x over 24 shape x tile combinations on an RTX PRO 6000 Blackwell (188 SMs, cuDNN 9.25), largest where the restage traffic was largest. The design notes predicted 1.04-1.08x from an instruction count; ncu says the kernel was L1-bound (72-77% L1/TEX against the backend kernel's 45-51%, DRAM 5-13% on both), so the binding resource was the L1 port and removing 36 memory ops per warp per tile was worth several times the issue-count estimate. It also frees 16 KB of SMEM (49 -> 33 KB at 128x128), which moves the tile optimum: kv_tile is now 128 in all 28 shapes measured, having been 64 in 47 of 48 on the old kernel. fp8_tile_choice is re-derived accordingly, and both it and its tests now say that a tile rule is a property of the kernel it was measured on. Held-out regret against the best of the enumerated domain is 1.009x mean / 1.089x worst over 22 shapes; against the backend's own fp8 fprop the engine is ahead on all 22, by 1.15x-1.93x. Three further optimizations from the sibling DKG kernel were measured and NOT taken. Prefetching K or V fragments one step ahead lands inside +-0.5%, and loading Q coalesced with a shfl transpose is reproducibly 0.5-2.3% slower at s=512 -- the shape it should help most, since Q is small enough to sit in L2. The run-to-run floor is ~1% at s>=2048 and was 12% at s=512 until the probe warmed the device first, which is what made these calls possible. Also from review: - Capabilities gains out_dtypes. facts.dtype_o never reached mismatch(), so a graph asking for fp8 O was claimed and then declined at build. The default is empty rather than unconstrained: a quantized row that forgets to declare serves nothing, which fails loudly. - The post-kernel amax_o.div_ and O copy-back move onto the launch stream, next to the resets that were already there (#510). - The negative tests the module docstring promised now exist (fp8 O, D != 128). Co-Authored-By: Claude Opus 5 (1M context) --- docs/fe-oss-apis/attention/sdpa-fp8-sm120.md | 144 +++++++++--------- python/cudnn/frost/tile_dsl/mma.py | 19 ++- python/cudnn/sdpa/fwd/api_dsl.py | 18 +-- python/cudnn/sdpa/fwd/config_sm120.py | 37 ++--- python/cudnn/sdpa/fwd/engines.py | 11 +- .../sdpa/fwd/kernels/prefill_fp8_sm120.py | 71 ++++----- .../sdpa/frost/test_sdpa_fwd_fp8_sm120.py | 95 ++++++++++-- .../python/sdpa/frost/test_sm120_fp8_tiles.py | 67 ++++---- 8 files changed, 283 insertions(+), 179 deletions(-) diff --git a/docs/fe-oss-apis/attention/sdpa-fp8-sm120.md b/docs/fe-oss-apis/attention/sdpa-fp8-sm120.md index 9a562bf85..544eec604 100644 --- a/docs/fe-oss-apis/attention/sdpa-fp8-sm120.md +++ b/docs/fe-oss-apis/attention/sdpa-fp8-sm120.md @@ -20,9 +20,18 @@ comparator, kernel time via CUPTI (min of 3 trials × 20 iters; host excluded): | 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 (SM120, more SMs / bandwidth): **pending** — numbers to -be added from the workstation run; the kernel is arch-identical so ratios are -expected to transfer, absolute times to scale with SM count and clocks. +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 @@ -30,13 +39,11 @@ O max abs err 3e-4..3e-3 at std `1/sqrt(d)` inputs, ~1e-2 at full-scale (std ~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. -SASS profile (d=128, both mask instances): 384× `QMMA.16832.F32.E4M3.E4M3` -(exactly half the f16 kernel's 768× HMMA — the k-depth doubled), 196× -`MUFU.EX2` (unchanged fp32 softmax), 96× `LDSM.8.MT1616.2` (hardware 8-bit -transposed V loads), 108× `LDSM.16.M88.4` (K + P-reload), 96× `STS.U16` + -96× `F2FP.SATFINITE.E4M3.F32.PACK` (P restage), 168 registers, zero spills. -The amax atomics compile to one warp `REDUX.MAX` + one global `RED` per warp — -negligible. +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 @@ -60,71 +67,70 @@ hardware 8-bit transpose added for exactly this MMA family. One issue covers a transposed smem layout (second TMA descriptor bank conflicts) or in-register `prmt` transposes — this instruction is why the port is clean. -**P restage through SMEM (the main open perf item).** The QK C-fragment owns -columns `2*(t%4)+{0,1}` but the PV A-fragment wants 4 *consecutive* bytes -`4*(t%4)..+3`, so the f16 kernel's in-register P repack cannot work at k32. -v1 does correctness-first: `cvt.rn.satfinite.e4m3x2` → per-warp 16×kv_tile -smem tile at C coordinates → `bar_warp_sync` → `ldmatrix.m8n8.x4.b16` back as -A fragments. Cost per warp per KV tile: 32 `STS.U16` + 4 `LDSM.x4` + 2 warp -barriers + addressing — roughly 8–15% of dynamic instruction issue. - -*Option: shfl-permute restage.* The C→A ownership mismatch is exchangeable -within each thread-quad: each A register needs two 16-bit P pairs from two -statically-known quad lanes → 2 `shfl.sync` + 1 `prmt` per A register, ~48 -warp-uniform ops per tile replacing the ~36 memory ops + 2 barriers + address -math. Expected gain **~4–8% kernel time (1.04–1.08x)** — instruction-count -bound, plus an unquantified ILP benefit from removing both warp barriers -between softmax and PV. Worth doing; not a step change. - -*Option: double-buffered sP.* Removes one of the two `bar_warp_sync`s -without touching the fragment math; ~half the barrier cost for 16KB more smem -(still fits). Subsumed by the shfl option. - -**Amax outputs always on.** Amax_S/Amax_O are computed unconditionally -(dummy buffers when the graph doesn't request them) instead of specializing -the template on amax presence. Tradeoff: one warp-REDUX + one global RED per -warp of pure overhead in the no-amax case — measured noise — versus 2× fewer -template specializations. - -**FP16 O only.** An fp8 O needs a `scale_o`-quantizing epilogue store -(`F2FP.SATFINITE` pack + byte stmatrix path). Straightforward follow-up; cut -from v1 to keep the epilogue identical to the proven f16 one. Note the fp16 -epilogue already does the `o_scale_fused` multiply, so adding fp8-O is a store --path change only. - -**E4M3 only.** The MMA tag is compile-time; an e5m2 variant is one more -template axis (the fragment math is identical). Deferred until a consumer -exists — e5m2 Q/K/V for *forward* is rare. - -**Exact d128.** The f16 kernel's zero-padding envelope (d_qk/d_v up to the -tile) is not wired for the 8-bit fragment path (TMA zero-fill interacts with -the byte-level A/B fragment addressing untested). d ∈ {16..256, mult of 16} -generalizes the fragment loops the same way as f16; only d128 is validated, -so the engine gates exact-match. - -**No sink, no THD, no seq_len_q (v1).** Sink: Amax_S (= max 1/row_sum) is -ill-defined with a virtual sink column — needs a semantics decision, not code. -THD: the packed-LSE plumbing is untested with the fp8 epilogue. seq_len_q: -mirrors the SM100 fp8 row (no epilogue trim); `lower_dsl_prefill` already -forces it off for the fp8 family. +**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}` (same knob domain -as f16). f16 measurements showed the (128,128) default is optimal at s≥4k but -loses ~11% at s=1k to (64,64); the fp8 kernel inherits the same geometry, so -the same seqlen-dependent tile choice applies once a heuristic or autotune -proposes multiple plans. Not re-measured for fp8 yet. +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.fp8_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. Held-out regret against the best +of the enumerated domain is 1.009x mean, 1.089x worst over 22 shapes; the +misses are causal, where the triangular mask changes the per-CTA balance in a +way grid alone does not capture. + +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. ## Bigger levers beyond this kernel -- **K-load batching**: issue `ldmatrix` x4 across two k32 steps (halve issue - count) — small. -- **Software pipelining / multi-stage KV**: 1-byte KV halves smem per stage; - a 2-stage sK/sV pipeline fits where f16 could not (32KB vs 64KB per stage - pair). Hides TMA latency behind compute; the single-buffered design shows - up at short seqlen where tiles don't amortize. Likely the largest remaining - item after the shfl restage. +- **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 diff --git a/python/cudnn/frost/tile_dsl/mma.py b/python/cudnn/frost/tile_dsl/mma.py index 0329b6f05..b2bf2da36 100644 --- a/python/cudnn/frost/tile_dsl/mma.py +++ b/python/cudnn/frost/tile_dsl/mma.py @@ -89,20 +89,33 @@ def ptx_mma_m16n8k32_e4m3_f32( @cute.jit -def ptx_cvt_e4m3x2(hi: cutlass.Float32, lo: cutlass.Float32) -> cutlass.Int32: +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; .reg .b16 t; " "mov.b32 fa, {$r0}; mov.b32 fb, {$r1}; " "cvt.rn.satfinite.e4m3x2.f32 t, fa, fb; cvt.u32.u16 {$w0}, t; }", - write_only_types=[cutlass.Int32], + "{ .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/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py index 539d9a542..f0b6698e7 100644 --- a/python/cudnn/sdpa/fwd/api_dsl.py +++ b/python/cudnn/sdpa/fwd/api_dsl.py @@ -1641,12 +1641,9 @@ 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. FP8: KV elements are 1 byte - # but O staging is FP16 (2 bytes), plus the per-compute-warp - # 16 x kv_tile P restage tiles. + # 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) - if self._fp8: - base += (self.q_tile // 16) * 16 * kv_tile return base + 16 chose_by_shape = False @@ -1655,7 +1652,7 @@ def _smem_bytes(kv_tile: int) -> int: # rather than by what fits (see fp8_tile_choice). Both tiles come # from one decision, so a half-specified request opts out entirely. sms = torch.cuda.get_device_properties(self.q_desc.device).multi_processor_count - self.q_tile, self.kv_tile = _sm120_fp8_tile_choice(int(s_q), int(h_q), int(b), int(sms), bool(self.is_causal)) + self.q_tile, self.kv_tile = _sm120_fp8_tile_choice(int(s_q), int(s_kv), int(h_q), int(b), int(sms)) chose_by_shape = _smem_bytes(self.kv_tile) <= smem_capacity_bytes if self.tile_n is None and not chose_by_shape: # Pick the largest KV tile that fits this device. @@ -1940,10 +1937,13 @@ def _scalar(t, default=1.0): cutlass.Float32(o_scale_fused), current_stream, ) - if o_needs_copy_back: - o_view.copy_(o_scratch) - if amax_o is not None: - amax_o_buf.div_(max(so, 1e-30)) + # 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( diff --git a/python/cudnn/sdpa/fwd/config_sm120.py b/python/cudnn/sdpa/fwd/config_sm120.py index 86c7fb278..1f4408a17 100644 --- a/python/cudnn/sdpa/fwd/config_sm120.py +++ b/python/cudnn/sdpa/fwd/config_sm120.py @@ -14,28 +14,31 @@ SUPPORTED_HEAD_TILES = tuple(range(16, 257, 16)) -def fp8_tile_choice(s_q: int, h_q: int, batch: int, sm_count: int, is_causal: bool = False) -> tuple[int, int]: +def fp8_tile_choice(s_q: int, s_kv: int, h_q: int, batch: int, sm_count: int) -> tuple[int, int]: """(q_tile, kv_tile) for the SM120 per-tensor FP8 cell. - FP8 only: the f16 cell wants ``kv_tile=128`` at long sequences, so the two - must not share a default. - - ``kv_tile=64`` unconditionally. The kernel is L1-bound (the P restage moves - every P tile through SMEM), and halving the KV tile halves that traffic per - tile; measured faster in 27 of 28 shapes on sm120, the exception by 0.15%. - - ``q_tile=64`` when the grid is too small to fill the machine, since halving - the Q tile halves each CTA's work and only costs another pass once the CTA - count outgrows the SMs. Causal doubles the window it pays over: with a - triangular mask the last Q tile does several times the work of the first, - so finer tiles even out the tail. Worth 1.45x at 64 CTAs and 1.47x at - 96 CTAs causal; misses elsewhere in the 48 shapes measured are all <= 4%. + FP8 only: this is read off the fp8 kernel's own measurements and the f16 + cell's optimum differs, so the two must not share a default. + + ``kv_tile=128`` unconditionally -- fastest in all 28 shapes measured. An + earlier revision of this 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. """ if sm_count <= 0: - return SEQ_Q_TILES[0], 64 + return 128, 128 grid = -(-s_q // 128) * h_q * batch - small = grid * 2 <= sm_count or (is_causal and grid <= sm_count) - return (64 if small else 128), 64 + 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) diff --git a/python/cudnn/sdpa/fwd/engines.py b/python/cudnn/sdpa/fwd/engines.py index 7f624aaa9..d62ed3810 100644 --- a/python/cudnn/sdpa/fwd/engines.py +++ b/python/cudnn/sdpa/fwd/engines.py @@ -114,6 +114,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) @@ -257,6 +261,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" @@ -403,6 +409,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, @@ -438,6 +445,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, @@ -507,7 +515,7 @@ def _sm120_fp8_knob_order(facts, named): (:func:`config_sm120.fp8_tile_choice`) -- named here so pinning entry 1 reproduces entry 0. """ - want = fp8_tile_choice(facts.s_q, facts.h_q, facts.b, facts.device_sm_count or 0, facts.causal) + want = fp8_tile_choice(facts.s_q, facts.s_kv, facts.h_q, facts.b, facts.device_sm_count or 0) return sorted(named, key=lambda k: ((k.tile_m, k.tile_n) != want, k.tile_n != 64, -(k.tile_m or 0))) @@ -532,6 +540,7 @@ def _sm120_fp8_spec() -> EngineSpec: 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, diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py b/python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py index cceca527f..a78fd23e2 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py @@ -54,7 +54,7 @@ from cutlass.experimental import primitives as prims from cudnn.frost.tile_dsl.constants import DTYPE_E4M3 -from cudnn.frost.tile_dsl.mma import ptx_cvt_e4m3x2, ptx_mma_m16n8k32_e4m3_f32 +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, @@ -481,17 +481,14 @@ def online_softmax( :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: ``s_regs`` with packed P fragments staged in consumed score slots. + :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 - - # The warp-private sP tile is re-written every KV iteration; this sync - # orders the stores after the PREVIOUS tile's mma_pv ldmatrix reads. - prims.bar_warp_sync(0xFFFFFFFF) + 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): @@ -580,14 +577,10 @@ def online_softmax( p0 = cute.math.exp2(in0, fastmath=True) p1 = cute.math.exp2(in1, fastmath=True) tile_sum = tile_sum + (p0 + p1) - # e4m3 P goes through the warp-private SMEM tile at the C-frag - # coordinates; the k32 A layout wants 4 CONSECUTIVE bytes per - # lane, so the f16 in-register restage cannot apply. row_sum - # above uses the fp32 P, keeping the denominator full precision. - p01 = ptx_cvt_e4m3x2(p1, p0) - p_byte_off = ((lane // 4) + row_half * 8) * self.kv_tile + k_frag * 8 + 2 * (lane % 4) - mma_params.sP_warp[p_byte_off] = cutlass.Uint8(p01 & 0xFF) - mma_params.sP_warp[p_byte_off + 1] = cutlass.Uint8((p01 >> 8) & 0xFF) + # 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) @@ -605,9 +598,7 @@ def online_softmax( (old_scale, old_scale), ) - # Make this tile's sP stores visible to the warp's mma_pv ldmatrix. - prims.bar_warp_sync(0xFFFFFFFF) - return s_regs + return p_regs @cute.jit def mma_pv( @@ -628,15 +619,23 @@ def mma_pv( o_regs = mma_params.o_regs lane = basic_params.lane - # P A-fragments come back from the warp-private SMEM tile with a - # byte-preserving ldmatrix.m8n8.x4.b16: lane pointers walk 16 rows x - # two 16-byte k-halves of one 16 x 32 P block. - p_row = basic_params.lane_mod8 + 8 * (basic_params.lane_div8 % 2) - p_coloff = 16 * basic_params.lane_div16 - - def load_p_frags(v_frag: cutlass.Constexpr[int]): - p_ptr = mma_params.sP_warp.data_ptr() + p_row * self.kv_tile + v_frag * self.MMA_TILER[2] + p_coloff - return prims.ldmatrix(p_ptr, 4, prims.MMALayout.ROW) + # 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 @@ -667,8 +666,18 @@ def load_v_frags(v_frag: cutlass.Constexpr[int], d_frag_pair: cutlass.Constexpr[ 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): - p_vec = load_p_frags(v_frag) + # 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 @@ -875,7 +884,6 @@ def kernel( # 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) - # sP: per-compute-warp 16 x kv_tile e4m3 P restage tiles # 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 @@ -889,12 +897,6 @@ def kernel( ) sK = sKV sV = sKV.subview(k_tile_bytes) - sP = cutlass.Array( - cutlass.Uint8, - self.num_compute_warps * self.MMA_TILER[0] * self.kv_tile, - space=cutlass.AddressSpace.smem, - alignment=128, - ) 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) @@ -1026,7 +1028,6 @@ def kernel( mma_params = SimpleNamespace( sK=sK, sV=sV, - sP_warp=sP.subview(compute_warp_idx * self.MMA_TILER[0] * self.kv_tile), o_regs=o_regs, ) softmax_params = SimpleNamespace( diff --git a/test/python/sdpa/frost/test_sdpa_fwd_fp8_sm120.py b/test/python/sdpa/frost/test_sdpa_fwd_fp8_sm120.py index d734c93d6..81b3d2c18 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_fp8_sm120.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_fp8_sm120.py @@ -10,7 +10,7 @@ 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 and fp8 outputs are covered by negative tests. +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. """ @@ -158,8 +158,8 @@ def _ref_kwargs(sdpa_kwargs): return out -def _check(O, O_ref, amax_s, amax_s_ref, amax_o, amax_o_ref): - diff = (O.float() - O_ref).abs().max().item() +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}" @@ -178,24 +178,24 @@ def _check(O, O_ref, amax_s, amax_s_ref, amax_o, amax_o_ref): @torch_fork_set_rng(seed=0) def test_fp8_sm120_masks(mask): scale = 1.0 / math.sqrt(128) - O, 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(O, O_ref, a_s, a_s_ref, a_o, a_o_ref) + 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) - O, 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(O, O_ref, a_s, a_s_ref, a_o, a_o_ref) + 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) - O, 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(O, O_ref, a_s, a_s_ref, a_o, a_o_ref) + 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 @@ -203,8 +203,8 @@ def test_fp8_sm120_bottom_right_rectangular(): 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) - O, 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(O, O_ref, a_s, a_s_ref, a_o, a_o_ref) + 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 @@ -214,8 +214,8 @@ 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 {} - O, 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(O, O_ref, a_s, a_s_ref, a_o, a_o_ref) + 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 @@ -259,6 +259,71 @@ def test_fp8_sm120_e5m2_not_offered(): 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) @@ -267,5 +332,5 @@ def test_fp8_sm120_every_enumerated_tile(tiles): 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) - O, 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(O, O_ref, a_s, a_s_ref, a_o, a_o_ref) + 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_sm120_fp8_tiles.py b/test/python/sdpa/frost/test_sm120_fp8_tiles.py index ef1598a0b..dc24a24a5 100644 --- a/test/python/sdpa/frost/test_sm120_fp8_tiles.py +++ b/test/python/sdpa/frost/test_sm120_fp8_tiles.py @@ -22,41 +22,48 @@ class TestTileChoice: """No device: shape and SM count 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_64(self, s_q, h_q, b, causal): - """Measured faster in 47 of 48 shapes; the exception by 0.15%. A change - here is a claim that the P restage stopped dominating.""" - assert fp8_tile_choice(s_q, h_q, b, SMS, causal)[1] == 64 - - def test_small_grid_takes_the_finer_q_tile(self): - # grid = ceil(512/128) * 16 = 64, so 64 CTAs of 188: doubling them - # still costs no extra pass over the machine. - assert fp8_tile_choice(512, 16, 1, SMS, False)[0] == 64 - - def test_causal_widens_the_window(self): - # grid 96: under a causal mask the last Q tile does several times the - # work of the first, so the finer tile still pays. Without it, it does - # not -- both directions are measured (1.47x and 1.00x regret). - assert fp8_tile_choice(768, 16, 1, SMS, True)[0] == 64 - assert fp8_tile_choice(768, 16, 1, SMS, False)[0] == 128 - - @pytest.mark.parametrize("s_q", [3072, 4096, 8192, 16384]) - @pytest.mark.parametrize("causal", [False, True]) - def test_large_grid_takes_the_coarser_q_tile(self, s_q, causal): - assert fp8_tile_choice(s_q, 16, 1, SMS, causal)[0] == 128 + def test_kv_tile_is_always_128(self, s_q, h_q, b): + """Fastest in all 28 shapes measured. 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 fp8_tile_choice(s_q, s_q, h_q, b, SMS)[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 fp8_tile_choice(512, 512, 16, 1, SMS)[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 fp8_tile_choice(1024, 1024, 16, 1, SMS)[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 fp8_tile_choice(2048, 2048, 16, 1, SMS)[0] == 64 + assert fp8_tile_choice(2048, 2048, 8, 2, SMS)[0] == 64 + + @pytest.mark.parametrize("s_q", [4096, 8192, 16384]) + def test_a_full_machine_takes_the_coarser_q_tile(self, s_q): + assert fp8_tile_choice(s_q, s_q, 16, 1, SMS)[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 fp8_tile_choice(2560, 2560, 12, 1, SMS)[0] == 64 # grid 240 + assert fp8_tile_choice(2560, 2560, 16, 1, SMS)[0] == 128 # grid 320 def test_batch_and_heads_count_toward_the_grid(self): - """The rule is grid, not sequence length: the same s_q flips once the - batch supplies enough CTAs on its own.""" - assert fp8_tile_choice(512, 16, 1, SMS, False)[0] == 64 - assert fp8_tile_choice(512, 16, 16, SMS, False)[0] == 128 + """The rule reads grid, not sequence length: the same s_q flips once + the batch supplies enough CTAs on its own.""" + assert fp8_tile_choice(2048, 2048, 16, 1, SMS)[0] == 64 + assert fp8_tile_choice(2048, 2048, 16, 4, SMS)[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 = fp8_tile_choice(s_q, 16, b, SMS, causal) - assert q in SEQ_Q_TILES and kv in SEQ_KV_TILES + q, kv = fp8_tile_choice(s_q, s_q, 16, b, SMS) + 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 fp8_tile_choice(4096, 16, 1, 0, False) == (SEQ_Q_TILES[0], 64) + assert fp8_tile_choice(4096, 4096, 16, 1, 0) == (128, 128) From ce8b3e10ee21f36a2e522a43be01c9059c679f87 Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Sat, 8 Aug 2026 00:37:34 -0700 Subject: [PATCH 05/10] Give the SM120 tile choice a causal term, and use one rule for both cells The f16/bf16 cell shared the fp8 cell's knob domain but had no rule of its own, so it took whatever fit: 128x128 everywhere on this part. Measuring it produced the same rule the fp8 cell uses, plus one term. A causal mask halves the work per CTA, so the machine empties sooner and the finer Q tile keeps paying further out. Folded in as a halved effective grid. Regret against the best of the enumerated {64,128}^2 domain: without with RTX PRO 6000, 24 bf16 cells 1.023 / 1.175 1.009 / 1.054 RTX PRO 6000, 28 fp8 cells 1.0078/ 1.107 1.0066/ 1.107 (2 fixed, 2 broken) RTX 5090 (170 SM), 14 fp8 1.0082/ 1.062 1.0007/ 1.006 22 held-out fp8 shapes 1.009 / 1.089 flips the 3 worst cells to optimal (mean / worst.) On the part the rule was trained on the term is a wash and the worst case is unchanged; it earns its place on the second part and the held-out set, where it fixes exactly the causal cells that were the previous misses. So fp8_tile_choice and f16_tile_choice collapse into one tile_choice, and the two knob_order callbacks into one. That also drops a leftover: the fp8 callback ranked its runner-up entries by tile_n != 64, from before the shfl change moved the KV optimum to 128. The tile-rule test file covers both cells now, so it loses the fp8 in its name. --- docs/fe-oss-apis/attention/sdpa-fp8-sm120.md | 37 ++++++-- python/cudnn/sdpa/fwd/api_dsl.py | 13 +-- python/cudnn/sdpa/fwd/config_sm120.py | 25 +++--- python/cudnn/sdpa/fwd/engines.py | 29 ++++--- .../python/sdpa/frost/test_sm120_fp8_tiles.py | 69 --------------- test/python/sdpa/frost/test_sm120_tiles.py | 85 +++++++++++++++++++ 6 files changed, 154 insertions(+), 104 deletions(-) delete mode 100644 test/python/sdpa/frost/test_sm120_fp8_tiles.py create mode 100644 test/python/sdpa/frost/test_sm120_tiles.py diff --git a/docs/fe-oss-apis/attention/sdpa-fp8-sm120.md b/docs/fe-oss-apis/attention/sdpa-fp8-sm120.md index 544eec604..43bfa7a9b 100644 --- a/docs/fe-oss-apis/attention/sdpa-fp8-sm120.md +++ b/docs/fe-oss-apis/attention/sdpa-fp8-sm120.md @@ -103,7 +103,7 @@ hide. 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.fp8_tile_choice` decide. +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 @@ -114,16 +114,43 @@ 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. Held-out regret against the best -of the enumerated domain is 1.009x mean, 1.089x worst over 22 shapes; the -misses are causal, where the triangular mask changes the per-CTA balance in a -way grid alone does not capture. +`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 diff --git a/python/cudnn/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py index f0b6698e7..1ea2b6832 100644 --- a/python/cudnn/sdpa/fwd/api_dsl.py +++ b/python/cudnn/sdpa/fwd/api_dsl.py @@ -34,7 +34,7 @@ from cudnn.sdpa.fwd.config_sm120 import ( SEQ_KV_TILES as _SM120_KV_TILES, SEQ_Q_TILES as _SM120_Q_TILES, - fp8_tile_choice as _sm120_fp8_tile_choice, + tile_choice as _sm120_tile_choice, SUPPORTED_HEAD_TILES as _SM120_SUPPORTED_HEAD_TILES, TemplateParams as Sm120TemplateParams, ) @@ -1647,12 +1647,13 @@ def _smem_bytes(kv_tile: int) -> int: return base + 16 chose_by_shape = False - if self._fp8 and self.tile_m is None and self.tile_n is None: - # No knob request: the fp8 cell picks by shape and machine width - # rather than by what fits (see fp8_tile_choice). Both tiles come - # from one decision, so a half-specified request opts out entirely. + if self.tile_m is None and self.tile_n is None: + # No knob request: pick by shape and machine width rather than by + # what fits. Each cell has its own rule -- they share a knob domain + # but not an optimum. Both tiles come from one decision, so a + # half-specified request opts out entirely. sms = torch.cuda.get_device_properties(self.q_desc.device).multi_processor_count - self.q_tile, self.kv_tile = _sm120_fp8_tile_choice(int(s_q), int(s_kv), int(h_q), int(b), int(sms)) + self.q_tile, self.kv_tile = _sm120_tile_choice(int(s_q), int(s_kv), int(h_q), int(b), int(sms), bool(self.is_causal)) chose_by_shape = _smem_bytes(self.kv_tile) <= smem_capacity_bytes if self.tile_n is None and not chose_by_shape: # Pick the largest KV tile that fits this device. diff --git a/python/cudnn/sdpa/fwd/config_sm120.py b/python/cudnn/sdpa/fwd/config_sm120.py index 1f4408a17..041b1cb63 100644 --- a/python/cudnn/sdpa/fwd/config_sm120.py +++ b/python/cudnn/sdpa/fwd/config_sm120.py @@ -14,17 +14,15 @@ SUPPORTED_HEAD_TILES = tuple(range(16, 257, 16)) -def fp8_tile_choice(s_q: int, s_kv: int, h_q: int, batch: int, sm_count: int) -> tuple[int, int]: - """(q_tile, kv_tile) for the SM120 per-tensor FP8 cell. +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. - FP8 only: this is read off the fp8 kernel's own measurements and the f16 - cell's optimum differs, so the two must not share a default. - - ``kv_tile=128`` unconditionally -- fastest in all 28 shapes measured. An - earlier revision of this 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. + ``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 @@ -32,10 +30,17 @@ def fp8_tile_choice(s_q: int, s_kv: int, h_q: int, batch: int, sm_count: int) -> 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 diff --git a/python/cudnn/sdpa/fwd/engines.py b/python/cudnn/sdpa/fwd/engines.py index d62ed3810..047de119c 100644 --- a/python/cudnn/sdpa/fwd/engines.py +++ b/python/cudnn/sdpa/fwd/engines.py @@ -35,7 +35,7 @@ from cudnn.frost.tile_dsl.constants import SCHED_NATURAL from cudnn.frost.buffers import CUTEDSL_MIN_VERSION, cutedsl_state, cutedsl_too_old from cudnn.sdpa import graph_analyzer as ga -from cudnn.sdpa.fwd.config_sm120 import fp8_tile_choice +from cudnn.sdpa.fwd.config_sm120 import tile_choice # The DSL adapters (api_dsl) and cuda.bindings are LOWERING dependencies, not # support-check ones: importing them here would drag the CuTe DSL (~1.0 s, 357 @@ -474,6 +474,18 @@ def _sm100_fp8_spec(d: int) -> EngineSpec: ) +def _sm120_knob_order(facts, named): + """Order an SM120 cell's tile domain: the shape's own choice first. + + ``heuristics_sort`` concatenates rather than ranks today, so whatever this + puts first is what runs. Same rule as the adapter's delegation path + (:func:`config_sm120.tile_choice`) -- named here so pinning entry 1 + reproduces entry 0. + """ + want = tile_choice(facts.s_q, facts.s_kv, facts.h_q, facts.b, facts.device_sm_count or 0, facts.causal) + return sorted(named, key=lambda k: ((k.tile_m, k.tile_n) != want, k.tile_n != 128, -(k.tile_m or 0))) + + def _sm120_spec() -> EngineSpec: return EngineSpec( name="sdpa_fwd_prefill_sm120", @@ -504,21 +516,10 @@ def _sm120_spec() -> EngineSpec: cgas=frozenset({1}), ), lower=partial(lower_dsl_prefill, api_type=_SM120), + knob_order=_sm120_knob_order, ) -def _sm120_fp8_knob_order(facts, named): - """Order the fp8 cell's tile domain: the shape's own choice first. - - ``heuristics_sort`` concatenates rather than ranks today, so whatever this - puts first is what runs. Same rule as the adapter's delegation path - (:func:`config_sm120.fp8_tile_choice`) -- named here so pinning entry 1 - reproduces entry 0. - """ - want = fp8_tile_choice(facts.s_q, facts.s_kv, facts.h_q, facts.b, facts.device_sm_count or 0) - return sorted(named, key=lambda k: ((k.tile_m, k.tile_n) != want, k.tile_n != 64, -(k.tile_m or 0))) - - def _sm120_fp8_spec() -> EngineSpec: """SM120 per-tensor FP8 engine (E4M3 in + scalar descales, FP16 out). @@ -560,7 +561,7 @@ def _sm120_fp8_spec() -> EngineSpec: cgas=frozenset({1}), ), lower=partial(lower_dsl_prefill, api_type=_SM120), - knob_order=_sm120_fp8_knob_order, + knob_order=_sm120_knob_order, ) diff --git a/test/python/sdpa/frost/test_sm120_fp8_tiles.py b/test/python/sdpa/frost/test_sm120_fp8_tiles.py deleted file mode 100644 index dc24a24a5..000000000 --- a/test/python/sdpa/frost/test_sm120_fp8_tiles.py +++ /dev/null @@ -1,69 +0,0 @@ -# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT - -"""The SM120 FP8 tile choice and the knob domain it is chosen from. - -The rule half needs no device: it is arithmetic over shape and SM count. The -correctness half does, because propose_plans now exposes tiles the default -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, fp8_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 and SM count 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)]) - def test_kv_tile_is_always_128(self, s_q, h_q, b): - """Fastest in all 28 shapes measured. 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 fp8_tile_choice(s_q, s_q, h_q, b, SMS)[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 fp8_tile_choice(512, 512, 16, 1, SMS)[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 fp8_tile_choice(1024, 1024, 16, 1, SMS)[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 fp8_tile_choice(2048, 2048, 16, 1, SMS)[0] == 64 - assert fp8_tile_choice(2048, 2048, 8, 2, SMS)[0] == 64 - - @pytest.mark.parametrize("s_q", [4096, 8192, 16384]) - def test_a_full_machine_takes_the_coarser_q_tile(self, s_q): - assert fp8_tile_choice(s_q, s_q, 16, 1, SMS)[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 fp8_tile_choice(2560, 2560, 12, 1, SMS)[0] == 64 # grid 240 - assert fp8_tile_choice(2560, 2560, 16, 1, SMS)[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 fp8_tile_choice(2048, 2048, 16, 1, SMS)[0] == 64 - assert fp8_tile_choice(2048, 2048, 16, 4, SMS)[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): - q, kv = fp8_tile_choice(s_q, s_q, 16, b, SMS) - 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 fp8_tile_choice(4096, 4096, 16, 1, 0) == (128, 128) 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..722ddfa31 --- /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 +propose_plans exposes tiles the default 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) From 65962a11a323ca5e13a41ade779628ed536d739f Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Sat, 8 Aug 2026 00:37:44 -0700 Subject: [PATCH 06/10] Ask the SDPA-forward eligibility question once per engine propose_plans called check_support(graph) and then knob_candidates(spec, graph), and each of those ran analyze_for -> mismatch(caps, facts, None). The facts parse itself is memoized per graph (_attach_facts -> _facts_for), so the waste was the capability match, once per engine per plan walk. knob_candidates now takes the facts the caller already has. _facts_or_decline is the single place the engine asks, returning the facts it just proved eligible; check_support is that call with the answer dropped. No behaviour change: knob_candidates could only be reached through a passing check_support, so its own decline branch was unreachable. --- python/cudnn/sdpa/fwd/engine.py | 24 ++++++++++++++---------- python/cudnn/sdpa/fwd/engines.py | 14 ++++++-------- 2 files changed, 20 insertions(+), 18 deletions(-) diff --git a/python/cudnn/sdpa/fwd/engine.py b/python/cudnn/sdpa/fwd/engine.py index a17e10922..8fb11380d 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,21 +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 check_support(self, graph: "pygraph") -> None: + self._facts_or_decline(graph) def propose_plans(self, graph: "pygraph") -> List[PlanConfig]: # Entry 0 delegates (knobs=None -> the adapter's own shape-driven @@ -111,8 +116,7 @@ def propose_plans(self, graph: "pygraph") -> List[PlanConfig]: # engine declines. from .engines import knob_candidates - self.check_support(graph) - return [PlanConfig(self.engine_id, k) for k in knob_candidates(self._spec, graph)] + return [PlanConfig(self.engine_id, k) for k in knob_candidates(self._spec, 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 047de119c..3cfc94050 100644 --- a/python/cudnn/sdpa/fwd/engines.py +++ b/python/cudnn/sdpa/fwd/engines.py @@ -565,10 +565,14 @@ def _sm120_fp8_spec() -> EngineSpec: ) -def knob_candidates(spec: EngineSpec, graph) -> List[Optional[SdpaFwdKnobs]]: +def knob_candidates(spec: EngineSpec, facts: "ga.SdpaGraphFacts") -> List[Optional[SdpaFwdKnobs]]: """``None`` (delegate to the adapter) followed by every admissible point of this engine's declared knob domain, best first. + Takes the facts the caller already established eligibility with, so the + knob-free ``mismatch`` is answered once per engine rather than here and in + ``check_support``. + The domain is read off the capability row -- there is no second table to drift from it -- so a cell declaring one value per axis proposes exactly one named entry. Feasibility that depends on device resources (shared @@ -576,12 +580,6 @@ def knob_candidates(spec: EngineSpec, graph) -> List[Optional[SdpaFwdKnobs]]: decline at build, and the walk moves on. """ caps = spec.capabilities - try: - facts, reason = analyze_for(spec, graph, None) - except ValueError: - return [None] - if reason is not None: - return [None] named = [ SdpaFwdKnobs(sched_policy=s, tile_m=m, tile_n=n, cga=c) for s in sorted(caps.sched_policies) @@ -600,7 +598,7 @@ def analyze_for(spec: EngineSpec, graph, knobs: Optional[SdpaFwdKnobs] = None): 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 From 33922bfb6b955f0b987f4692007e776c109e7e29 Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Sat, 8 Aug 2026 01:28:00 -0700 Subject: [PATCH 07/10] Take the backend's plans one heuristic mode at a time Ranking the two sides against each other needs to know which backend entries are mode-A recommendations and which are fallbacks -- "the backend's A ahead of ours, its fallbacks behind" cannot be said about one opaque list. Until now the whole thing arrived from a single create_execution_plans([A, FALLBACK]). No C++ change is needed. C++ appends each query to the same plan list, and get_execution_plan_count() already exists, so asking one mode at a time and reading the count after each gives the boundaries. Measured on a 512^3 bf16 matmul (sm90, cuDNN 9.25): A -> plans[0:15], all knob-bearing; FALLBACK -> plans[15:17], bare eng0/eng7 with no knobs; the two segments do not overlap. A mode with no configs raises, which is not a decline while another mode still has entries -- an OPENSOURCE-only query legitimately leaves the cuDNN modes empty. Only every mode failing means the backend has nothing, and then the last error is re-raised so the caller still reports why. --- python/cudnn/_pygraph.py | 50 +++++++++++++++++++++++++++++++----- python/cudnn/engines/base.py | 5 ++++ 2 files changed, 48 insertions(+), 7 deletions(-) diff --git a/python/cudnn/_pygraph.py b/python/cudnn/_pygraph.py index d41c4c3d4..9fcfcaac9 100644 --- a/python/cudnn/_pygraph.py +++ b/python/cudnn/_pygraph.py @@ -151,6 +151,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 @@ -990,6 +991,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: @@ -1108,6 +1110,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 +1132,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 +1322,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).""" diff --git a/python/cudnn/engines/base.py b/python/cudnn/engines/base.py index 73aa5545e..68ddb2830 100644 --- a/python/cudnn/engines/base.py +++ b/python/cudnn/engines/base.py @@ -63,11 +63,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) From 0750ff42f6ffa037c4cbce6cb0cafc9fa0776788 Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Sat, 8 Aug 2026 01:39:41 -0700 Subject: [PATCH 08/10] Move plan ranking out of the engines and into one heuristics function An engine cannot rank. It sees neither its siblings nor the backend's entries, so propose_plans could only ever order its own knobs -- and then something downstream had to merge the sides anyway, which heuristics_sort did by concatenating and calling it ranking. Now create_execution_plans() gathers the inputs (parsed facts, this family's offered ids, the backend's entries tagged by mode) and hands all of it to the graph's family in ONE call. What comes back IS graph.plans, position for position. sdpa/fwd/heuristics.py is the first such hook: mode A returns candidates worth timing with the tile rule's guess first, FALLBACK returns configs expected to build where A's choice may not, and both sides' fallbacks sort behind both sides' tuned entries. Which side leads is now a measurement rather than a default. sm120 fp8 is timed 1.20-1.83x the backend's native fp8 fprop, so it leads; a cell that has not been timed against the backend keeps the historical order and says so. Moving one is an experiment, not an edit. Deleted, all superseded or never used: BaseEngine.propose_plans + its 3 implementations (2 were verbatim copies of the default, the third was added last week -- nothing ever used the hook to decide anything) BaseEngine.default_knobs only fed propose_plans EngineSpec.knob_order ranking, in the wrong place knob_candidates ditto engines.probe() (fwd + bwd) zero callers, superseded by check_support heuristics_sort merging is part of ranking, not a step after Router.python_plans / Router class / default_router / set_router / pygraph(router=) policy has one home now; the router is plumbing graph.engine pure alias of selected_engine, zero callers graph.from_serialized zero callers; serialize/deserialize are the pybind-era API and stay api_dsl's shape-driven tile choice the tile rule ran twice, once here and once when ranking; resource feasibility stays knobs=None no longer means "engine, pick for me" -- heuristics always names a concrete config. A None field survives only on an axis whose capability row declares no domain. Tests are updated in the next commit; test_engine_router still builds Router subclasses to inject orderings, which is now done by replacing heuristics.rank. --- python/cudnn/_pygraph.py | 67 ++----------- python/cudnn/engines/__init__.py | 6 +- python/cudnn/engines/base.py | 23 ++--- python/cudnn/engines/heuristics.py | 99 +++++++++++++----- python/cudnn/engines/manifest.py | 45 ++++++--- python/cudnn/engines/router.py | 70 ++++--------- python/cudnn/gemm/frost/engine.py | 7 -- python/cudnn/sdpa/bwd/engine.py | 8 -- python/cudnn/sdpa/bwd/engines.py | 8 -- python/cudnn/sdpa/fwd/api_dsl.py | 18 ++-- python/cudnn/sdpa/fwd/engine.py | 10 -- python/cudnn/sdpa/fwd/engines.py | 56 +---------- python/cudnn/sdpa/fwd/heuristics.py | 150 ++++++++++++++++++++++++++++ 13 files changed, 301 insertions(+), 266 deletions(-) create mode 100644 python/cudnn/sdpa/fwd/heuristics.py diff --git a/python/cudnn/_pygraph.py b/python/cudnn/_pygraph.py index 9fcfcaac9..ef54c8096 100644 --- a/python/cudnn/_pygraph.py +++ b/python/cudnn/_pygraph.py @@ -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( @@ -130,12 +129,11 @@ def __init__( 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-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()) @@ -207,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). @@ -893,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. @@ -902,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 import router # One-shot planning (classic conformance: the C++ graph never supported # re-planning — a second call there appends plans by accident, and no @@ -928,16 +917,15 @@ def create_execution_plans(self, heuristics: Optional[List] = None) -> None: 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. + # 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. @@ -952,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 @@ -1854,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). @@ -1890,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..2e9be808e 100644 --- a/python/cudnn/engines/__init__.py +++ b/python/cudnn/engines/__init__.py @@ -24,9 +24,8 @@ is_backend_engine, is_python_engine, ) -from .heuristics import heuristics_sort from .manifest import MANIFEST, EngineFamily -from .router import Router, decline_types, default_router +from .router import decline_types def __getattr__(name: str): @@ -49,10 +48,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 68ddb2830..cb0c08e9e 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 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 @@ -166,7 +169,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 @@ -179,7 +181,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 = () @@ -220,16 +221,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..ade3a0ef4 100644 --- a/python/cudnn/engines/heuristics.py +++ b/python/cudnn/engines/heuristics.py @@ -1,41 +1,94 @@ # 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 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.""" + from .router import decline_types -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. + 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 - 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. - 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 _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) + + +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 10f66490d..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: @@ -175,6 +180,7 @@ def offered_ids(self) -> Dict[str, int]: "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, @@ -213,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 index d6639d0ac..d16c74c89 100644 --- a/python/cudnn/engines/router.py +++ b/python/cudnn/engines/router.py @@ -1,10 +1,16 @@ # 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. +"""Router: the plumbing that feeds the heuristics at create_execution_plans() time. - Python Graph API -> create_execution_plans() -> Router -> one ranked list - (graph.plans) + Python Graph API -> create_execution_plans() -> Router -> heuristics -> graph.plans + +Deliberately dumb. It gathers the inputs — the graph's python engine candidates +and the backend's own entries — calls ``heuristics.rank``, and checks that what +came back names engines this graph can dispatch to. It makes no ordering +decision of its own: policy lives in ``engines/heuristics.py`` and the family +hook it dispatches to, so there is exactly one place to change how plans are +ranked. 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 @@ -12,18 +18,9 @@ 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. +anything the caller added with ``register_backend()``. The backend's own entries +come from ``graph.backend_plan_entries()``, which is [] when the backend +declined the graph or is not installed. """ import logging @@ -50,38 +47,11 @@ def decline_types(): 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() +def plan(graph: "pygraph", engines: List[BaseEngine]) -> List[PlanConfig]: + """The ranked plan list for ``graph``. + + Through the module, not a bound name: ``heuristics.rank`` is the seam a + ranking policy replaces, and it must be swappable in a live process (tests, + experiments) without touching this file. + """ + return heuristics.rank(graph, engines, graph.backend_plan_entries(), graph._backend_heuristics) 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 1ea2b6832..1b293c693 100644 --- a/python/cudnn/sdpa/fwd/api_dsl.py +++ b/python/cudnn/sdpa/fwd/api_dsl.py @@ -34,7 +34,6 @@ from cudnn.sdpa.fwd.config_sm120 import ( SEQ_KV_TILES as _SM120_KV_TILES, SEQ_Q_TILES as _SM120_Q_TILES, - tile_choice as _sm120_tile_choice, SUPPORTED_HEAD_TILES as _SM120_SUPPORTED_HEAD_TILES, TemplateParams as Sm120TemplateParams, ) @@ -1646,17 +1645,12 @@ def _smem_bytes(kv_tile: int) -> int: base = max(kv_tile * (d_q + d_v) * self.dtype.itemsize, self.q_tile * d_v * o_item) return base + 16 - chose_by_shape = False - if self.tile_m is None and self.tile_n is None: - # No knob request: pick by shape and machine width rather than by - # what fits. Each cell has its own rule -- they share a knob domain - # but not an optimum. Both tiles come from one decision, so a - # half-specified request opts out entirely. - sms = torch.cuda.get_device_properties(self.q_desc.device).multi_processor_count - self.q_tile, self.kv_tile = _sm120_tile_choice(int(s_q), int(s_kv), int(h_q), int(b), int(sms), bool(self.is_causal)) - chose_by_shape = _smem_bytes(self.kv_tile) <= smem_capacity_bytes - if self.tile_n is None and not chose_by_shape: - # Pick the largest KV tile that fits this device. + if self.tile_n is None: + # 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, diff --git a/python/cudnn/sdpa/fwd/engine.py b/python/cudnn/sdpa/fwd/engine.py index 8fb11380d..bbb2722a9 100644 --- a/python/cudnn/sdpa/fwd/engine.py +++ b/python/cudnn/sdpa/fwd/engine.py @@ -108,16 +108,6 @@ def _facts_or_decline(self, graph: "pygraph"): def check_support(self, graph: "pygraph") -> None: self._facts_or_decline(graph) - def propose_plans(self, graph: "pygraph") -> List[PlanConfig]: - # Entry 0 delegates (knobs=None -> the adapter's own shape-driven - # choice); the rest name one point of the capability row's knob domain - # each. Not duplicates: a delegation may pick differently as the - # library learns, a named knob set must be honored verbatim or the - # engine declines. - from .engines import knob_candidates - - return [PlanConfig(self.engine_id, k) for k in knob_candidates(self._spec, 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 3cfc94050..d6b232b03 100644 --- a/python/cudnn/sdpa/fwd/engines.py +++ b/python/cudnn/sdpa/fwd/engines.py @@ -35,7 +35,6 @@ from cudnn.frost.tile_dsl.constants import SCHED_NATURAL from cudnn.frost.buffers import CUTEDSL_MIN_VERSION, cutedsl_state, cutedsl_too_old from cudnn.sdpa import graph_analyzer as ga -from cudnn.sdpa.fwd.config_sm120 import tile_choice # The DSL adapters (api_dsl) and cuda.bindings are LOWERING dependencies, not # support-check ones: importing them here would drag the CuTe DSL (~1.0 s, 357 @@ -353,9 +352,6 @@ class EngineSpec: # kernels (e.g. decode vs prefill by S_q) or chain several launches under # one name. lower: "Callable[[EngineSpec, ga.SdpaGraphFacts, Optional[SdpaFwdKnobs]], Any]" - # Orders this cell's knob domain for propose_plans. None keeps the - # domain's natural order; a cell with measurements supplies its own. - knob_order: "Optional[Callable[[ga.SdpaGraphFacts, list], list]]" = None def _sm100_spec(d: int, d_v: Optional[int] = None) -> EngineSpec: @@ -474,18 +470,6 @@ def _sm100_fp8_spec(d: int) -> EngineSpec: ) -def _sm120_knob_order(facts, named): - """Order an SM120 cell's tile domain: the shape's own choice first. - - ``heuristics_sort`` concatenates rather than ranks today, so whatever this - puts first is what runs. Same rule as the adapter's delegation path - (:func:`config_sm120.tile_choice`) -- named here so pinning entry 1 - reproduces entry 0. - """ - want = tile_choice(facts.s_q, facts.s_kv, facts.h_q, facts.b, facts.device_sm_count or 0, facts.causal) - return sorted(named, key=lambda k: ((k.tile_m, k.tile_n) != want, k.tile_n != 128, -(k.tile_m or 0))) - - def _sm120_spec() -> EngineSpec: return EngineSpec( name="sdpa_fwd_prefill_sm120", @@ -516,7 +500,6 @@ def _sm120_spec() -> EngineSpec: cgas=frozenset({1}), ), lower=partial(lower_dsl_prefill, api_type=_SM120), - knob_order=_sm120_knob_order, ) @@ -561,38 +544,9 @@ def _sm120_fp8_spec() -> EngineSpec: cgas=frozenset({1}), ), lower=partial(lower_dsl_prefill, api_type=_SM120), - knob_order=_sm120_knob_order, ) -def knob_candidates(spec: EngineSpec, facts: "ga.SdpaGraphFacts") -> List[Optional[SdpaFwdKnobs]]: - """``None`` (delegate to the adapter) followed by every admissible point of - this engine's declared knob domain, best first. - - Takes the facts the caller already established eligibility with, so the - knob-free ``mismatch`` is answered once per engine rather than here and in - ``check_support``. - - The domain is read off the capability row -- there is no second table to - drift from it -- so a cell declaring one value per axis proposes exactly - one named entry. Feasibility that depends on device resources (shared - memory) stays in the adapter: a point that survives ``mismatch`` may still - decline at build, and the walk moves on. - """ - caps = spec.capabilities - named = [ - SdpaFwdKnobs(sched_policy=s, tile_m=m, tile_n=n, cga=c) - for s in sorted(caps.sched_policies) - for m in sorted(caps.tile_ms) - for n in sorted(caps.tile_ns) - for c in sorted(caps.cgas) - ] - named = [k for k in named if mismatch(caps, facts, k) is None] - if spec.knob_order is not None: - named = spec.knob_order(facts, named) - return [None] + named - - 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). @@ -609,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).""" @@ -868,4 +814,4 @@ def engine_name( _sm120_fp8_spec(), ) -__all__ = ["Capabilities", "EngineSpec", "ENGINE_SPECS", "SdpaFwdKnobs", "analyze_for", "build", "engine_name", "knob_candidates", "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..d7f5f6d11 --- /dev/null +++ b/python/cudnn/sdpa/fwd/heuristics.py @@ -0,0 +1,150 @@ +# 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. + +Modes this family has nothing to say about (B, OPENSOURCE) contribute no python +plans; the backend's entries for them still rank normally. B is treated as A: +it asks for a wider search than A, and this family has no wider search 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"}) + +# Where a cell stands against the backend, from measurements only. +# +# MEASURED_AHEAD: timed against the backend's own kernel and faster, so its +# plans lead. sm120 fp8 is 1.20-1.83x the backend's native fp8 fprop across 28 +# shapes (docs/fe-oss-apis/attention/sdpa-fp8-sm120.md). +# +# MEASURED_BEHIND: timed and slower -- the backend's mode-A entries lead and +# this cell is the second choice. +# +# A cell in NEITHER set has not been timed against the backend. Those keep the +# historical order (python ahead of the backend), which was a placeholder from +# when heuristics_sort concatenated the two sides, never a measurement. Moving +# one is a measurement, not an opinion: time it, then add it to a set. +_MEASURED_AHEAD = frozenset({"sdpa_fwd_prefill_sm120_fp8"}) +_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. + + Only a cell measured SLOWER than the backend yields the lead; unmeasured + cells keep the historical order. See _MEASURED_AHEAD / _MEASURED_BEHIND. + """ + behind = {offered[name] for name in _MEASURED_BEHIND if name in offered} + return not all(cfg.engine_id in behind for cfg in plans) if plans else False + + +def recommend(modes: List[Any], facts, offered: Dict[str, int], backend_plans: List[PlanConfig]) -> List[PlanConfig]: + """The ranked plan list for this graph: FROST and backend entries interleaved.""" + wanted = {m for m in modes} + a_modes = {cudnn.heur_mode.A, cudnn.heur_mode.B} + ours_a = _mode_a(facts, offered) if wanted & a_modes else [] + ours_fb = _mode_fallback(facts, offered) if cudnn.heur_mode.FALLBACK in wanted else [] + + backend_a = [c for c in backend_plans if c.mode != cudnn.heur_mode.FALLBACK] + backend_fb = [c for c in backend_plans if c.mode == cudnn.heur_mode.FALLBACK] + + # Fallbacks last on both sides: a fallback outranking any tuned candidate is + # what makes a plan walk take a slow kernel while a fast one was available. + if _leads(offered, ours_a): + return ours_a + backend_a + ours_fb + backend_fb + return backend_a + ours_a + ours_fb + backend_fb From 85a90fd8de45553ca6c509af33d8517ee50081ae Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Sat, 8 Aug 2026 01:56:29 -0700 Subject: [PATCH 09/10] Update the dispatch tests to the ranking contract, and delete what it retired Ranking has one home, so a test that wants a specific order replaces heuristics.rank instead of subclassing Router. The _ranking() helper does that; it is the same monkeypatch idiom the rest of the suite already uses (326 calls across 20 files). Deleted rather than translated: test_set_router_frozen_after_planning the API it tested is gone test_a_claiming_engine_is_tried_before_the_backend asserted that python plans always outrank the backend, which is now a measurement per cell, not a rule FROST coverage no longer rides on that ordering. heur_mode.OPENSOURCE asks for an open-source implementation, which is what these cells are, so it puts everything FROST can build ahead of the backend whatever the measurements say. A test asks for it and whatever still lands on a backend plan is a graph FROST does not cover -- coverage measured rather than asserted, and the plan walk falls through to the backend so the test still checks numerics either way. Renamed for what they now test: test_mixed_ranking_dispatch, test_empty_ranking_output_rejected, test_mixed_ranking_backend_slot_executes, test_constructor_backends_validated_and_ranking_ids_checked. One assertion changed meaning: the backend is queried once PER MODE now, so _create_backend_plans records two create_execution_plans calls for [A, FALLBACK]. Known: test_a_replayed_plan_reports_its_own_notes fails, and it already failed at d7dd85ee5 (verified by running it against that tree) -- C++ no longer raises for an index one past the plan count on 9.25. Pre-existing, not this change. Out-of-tree engines lose the ability to offer several knob candidates: they declare no Capabilities, so nothing can enumerate their domain. register_backend remains what the tests use to install a fake engine. --- python/cudnn/_pygraph.py | 26 +-- python/cudnn/engines/__init__.py | 5 +- python/cudnn/engines/base.py | 16 +- python/cudnn/engines/heuristics.py | 4 +- python/cudnn/engines/router.py | 57 ----- python/cudnn/sdpa/fwd/heuristics.py | 49 ++-- .../gemm/frost/test_frontend_integration.py | 11 +- .../sdpa/frost/test_sdpa_fwd_fp8_sm120.py | 2 +- test/python/sdpa/frost/test_sm120_tiles.py | 2 +- test/python/test_engine_router.py | 210 +++++++++--------- test/python/test_native_backend_lowering.py | 18 +- 11 files changed, 169 insertions(+), 231 deletions(-) delete mode 100644 python/cudnn/engines/router.py diff --git a/python/cudnn/_pygraph.py b/python/cudnn/_pygraph.py index ef54c8096..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): @@ -128,7 +128,7 @@ 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 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 @@ -242,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``.""" @@ -830,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 @@ -891,7 +891,7 @@ def create_execution_plans(self, heuristics: Optional[List] = None) -> None: if not self._is_validated: self.validate() - from .engines import 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 @@ -910,14 +910,14 @@ 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() - plans = list(router.plan(self, self._candidate_engines())) + 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 @@ -964,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 @@ -1057,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 @@ -1367,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, @@ -1413,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: @@ -1449,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() diff --git a/python/cudnn/engines/__init__.py b/python/cudnn/engines/__init__.py index 2e9be808e..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, @@ -25,7 +25,6 @@ is_python_engine, ) from .manifest import MANIFEST, EngineFamily -from .router import decline_types def __getattr__(name: str): diff --git a/python/cudnn/engines/base.py b/python/cudnn/engines/base.py index cb0c08e9e..9a248d3e7 100644 --- a/python/cudnn/engines/base.py +++ b/python/cudnn/engines/base.py @@ -3,8 +3,8 @@ """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. ``check_support(graph)`` -> accept, or decline the whole graph by raising @@ -53,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. diff --git a/python/cudnn/engines/heuristics.py b/python/cudnn/engines/heuristics.py index ade3a0ef4..30a28a728 100644 --- a/python/cudnn/engines/heuristics.py +++ b/python/cudnn/engines/heuristics.py @@ -28,7 +28,7 @@ import logging from typing import TYPE_CHECKING, Any, List, Optional -from .base import BaseEngine, PlanConfig +from .base import BaseEngine, PlanConfig, decline_types if TYPE_CHECKING: from .._pygraph import pygraph @@ -45,8 +45,6 @@ def default_modes() -> List[Any]: def accepts(engine: BaseEngine, graph: "pygraph") -> bool: """Whether ``engine`` will serve ``graph``, declines being routing not error.""" - from .router import decline_types - try: engine.check_support(graph) except decline_types() as exc: diff --git a/python/cudnn/engines/router.py b/python/cudnn/engines/router.py deleted file mode 100644 index d16c74c89..000000000 --- a/python/cudnn/engines/router.py +++ /dev/null @@ -1,57 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Router: the plumbing that feeds the heuristics at create_execution_plans() time. - - Python Graph API -> create_execution_plans() -> Router -> heuristics -> graph.plans - -Deliberately dumb. It gathers the inputs — the graph's python engine candidates -and the backend's own entries — calls ``heuristics.rank``, and checks that what -came back names engines this graph can dispatch to. It makes no ordering -decision of its own: policy lives in ``engines/heuristics.py`` and the family -hook it dispatches to, so there is exactly one place to change how plans are -ranked. - -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 entries -come from ``graph.backend_plan_entries()``, which is [] when the backend -declined the graph or is not installed. -""" - -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) - - -def plan(graph: "pygraph", engines: List[BaseEngine]) -> List[PlanConfig]: - """The ranked plan list for ``graph``. - - Through the module, not a bound name: ``heuristics.rank`` is the seam a - ranking policy replaces, and it must be swappable in a live process (tests, - experiments) without touching this file. - """ - return heuristics.rank(graph, engines, graph.backend_plan_entries(), graph._backend_heuristics) diff --git a/python/cudnn/sdpa/fwd/heuristics.py b/python/cudnn/sdpa/fwd/heuristics.py index d7f5f6d11..7c866cb03 100644 --- a/python/cudnn/sdpa/fwd/heuristics.py +++ b/python/cudnn/sdpa/fwd/heuristics.py @@ -17,9 +17,13 @@ - **FALLBACK** — configs that are expected to build wherever mode A might not, ordered cheapest-resource first. Nothing here is chosen for speed. -Modes this family has nothing to say about (B, OPENSOURCE) contribute no python -plans; the backend's entries for them still rank normally. B is treated as A: -it asks for a wider search than A, and this family has no wider search to give. +``heur_mode.OPENSOURCE`` asks for an open-source implementation, which is what +these cells are, so it puts everything FROST can build ahead of the backend +whatever the measurements say. Tests use it to measure FROST coverage: ask for +it, and any graph that still lands on a backend plan is one FROST cannot serve. + +``heur_mode.B`` is treated as A -- it asks for a wider search than A, and this +family has no wider search to give. """ from __future__ import annotations @@ -39,20 +43,14 @@ # outlives its evidence. _TILE_RULE_CELLS = frozenset({"sdpa_fwd_prefill_sm120", "sdpa_fwd_prefill_sm120_fp8"}) -# Where a cell stands against the backend, from measurements only. -# -# MEASURED_AHEAD: timed against the backend's own kernel and faster, so its -# plans lead. sm120 fp8 is 1.20-1.83x the backend's native fp8 fprop across 28 +# 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). # -# MEASURED_BEHIND: timed and slower -- the backend's mode-A entries lead and -# this cell is the second choice. -# -# A cell in NEITHER set has not been timed against the backend. Those keep the -# historical order (python ahead of the backend), which was a placeholder from -# when heuristics_sort concatenated the two sides, never a measurement. Moving -# one is a measurement, not an opinion: time it, then add it to a set. -_MEASURED_AHEAD = frozenset({"sdpa_fwd_prefill_sm120_fp8"}) +# 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() @@ -124,25 +122,30 @@ def _mode_fallback(facts, offered: Dict[str, int]) -> List[PlanConfig]: def _leads(offered: Dict[str, int], plans: List[PlanConfig]) -> bool: - """Whether this family's mode-A plans outrank the backend's. - - Only a cell measured SLOWER than the backend yields the lead; unmeasured - cells keep the historical order. See _MEASURED_AHEAD / _MEASURED_BEHIND. - """ + """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 not all(cfg.engine_id in behind for cfg in plans) if plans else False + 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: FROST and backend entries interleaved.""" - wanted = {m for m in modes} - a_modes = {cudnn.heur_mode.A, cudnn.heur_mode.B} + wanted = set(modes) + # OPENSOURCE wants these cells tried, so it asks for candidates like A does. + a_modes = {cudnn.heur_mode.A, cudnn.heur_mode.B, cudnn.heur_mode.OPENSOURCE} ours_a = _mode_a(facts, offered) if wanted & a_modes else [] ours_fb = _mode_fallback(facts, offered) if cudnn.heur_mode.FALLBACK in wanted else [] backend_a = [c for c in backend_plans if c.mode != cudnn.heur_mode.FALLBACK] backend_fb = [c for c in backend_plans if c.mode == cudnn.heur_mode.FALLBACK] + if cudnn.heur_mode.OPENSOURCE in wanted: + # The caller asked for an open-source implementation, which is what these + # cells are. Everything FROST can build goes first, regardless of which + # side measures faster -- that is the mode's whole meaning, and it is how + # a test measures FROST coverage: ask for OPENSOURCE, and whatever runs + # on a backend plan instead is a graph FROST does not cover yet. + return ours_a + ours_fb + backend_a + backend_fb + # Fallbacks last on both sides: a fallback outranking any tuned candidate is # what makes a plan walk take a slow kernel while a fast one was available. if _leads(offered, ours_a): 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/test_sdpa_fwd_fp8_sm120.py b/test/python/sdpa/frost/test_sdpa_fwd_fp8_sm120.py index 81b3d2c18..e9e5f320f 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_fp8_sm120.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_fp8_sm120.py @@ -328,7 +328,7 @@ def test_fp8_sm120_non_128_head_dim_not_offered(D): @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): - """propose_plans offers the whole tile domain, but a shape only ever runs + """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) diff --git a/test/python/sdpa/frost/test_sm120_tiles.py b/test/python/sdpa/frost/test_sm120_tiles.py index 722ddfa31..a75bef6f0 100644 --- a/test/python/sdpa/frost/test_sm120_tiles.py +++ b/test/python/sdpa/frost/test_sm120_tiles.py @@ -6,7 +6,7 @@ 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 -propose_plans exposes tiles the default never picks, and nothing else in the +the heuristics expose tiles the best guess never picks, and nothing else in the suite runs a kernel at those. """ 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]) From a4fc449317cac472e8c861c6b96affa9cd0635c4 Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Sat, 8 Aug 2026 02:17:29 -0700 Subject: [PATCH 10/10] Answer each heuristic mode on its own terms, and share the test's plan lookup OPENSOURCE is mode A without the backend's recommendation: these cells ARE the open-source implementation and the backend's engines are not, so there is nothing of the backend's to place. Combine it to measure coverage -- [OPENSOURCE, A, FALLBACK] tries every FROST config first and still has the backend behind it, so a graph that runs on a backend plan is one FROST does not cover. recommend() is now a loop over the caller's modes rather than three blocks spliced together. Each mode contributes its own, so [A, FALLBACK] puts both sides' tuned candidates ahead of both sides' fallbacks by construction. An untagged backend entry is the delegating one C++ tries before its own engine_configs; it belongs to no mode and keeps the lead, or an OPENSOURCE caller would get a native kernel instead of the OSS one. Plan names now always read [], since the heuristics name a concrete config for every entry -- there is no knobs=None entry left to print as a bare engine name. Six test files each carried their own verbatim copy of _select_engine matching that bare name exactly; they now share frost_test_utils.select_engine, which matches on the engine. probe() is restored to its callers rather than deleted: I claimed it had none, which was wrong -- test_sdpa_graph_analyzer used it twice. Those two call analyze_for directly now, so no production API exists only for tests. sm120: sdpa suites 62 failed -> 4 (171 passed); test_mhas_v2 302 passed with the routing tally unchanged (frost sm120 x87, fp8 x4, native x158/x53). The 4 are head_dim 208-256, where the adapter used to drop to a KV tile that fits and the heuristics now name 128 -- traceback pending, not yet diagnosed. --- python/cudnn/sdpa/fwd/heuristics.py | 73 +++++++++++-------- test/python/sdpa/frost/frost_test_utils.py | 32 ++++++++ .../sdpa/frost/test_sdpa_bwd_dsl_sm120.py | 9 +-- .../frost/test_sdpa_frontend_integration.py | 25 +++++-- .../sdpa/frost/test_sdpa_fwd_dsl_sm100.py | 11 +-- .../sdpa/frost/test_sdpa_fwd_dsl_sm120.py | 9 +-- .../sdpa/frost/test_sdpa_fwd_fp8_sm100.py | 10 +-- .../sdpa/frost/test_sdpa_fwd_fp8_sm120.py | 18 +---- .../sdpa/frost/test_sdpa_fwd_mxfp8_sm100.py | 10 +-- .../sdpa/frost/test_sdpa_graph_analyzer.py | 8 +- 10 files changed, 102 insertions(+), 103 deletions(-) diff --git a/python/cudnn/sdpa/fwd/heuristics.py b/python/cudnn/sdpa/fwd/heuristics.py index 7c866cb03..9da7df115 100644 --- a/python/cudnn/sdpa/fwd/heuristics.py +++ b/python/cudnn/sdpa/fwd/heuristics.py @@ -17,13 +17,14 @@ - **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`` asks for an open-source implementation, which is what -these cells are, so it puts everything FROST can build ahead of the backend -whatever the measurements say. Tests use it to measure FROST coverage: ask for -it, and any graph that still lands on a backend plan is one FROST cannot serve. - -``heur_mode.B`` is treated as A -- it asks for a wider search than A, and this -family has no wider search to give. +``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 @@ -128,26 +129,38 @@ def _leads(offered: Dict[str, int], plans: List[PlanConfig]) -> bool: def recommend(modes: List[Any], facts, offered: Dict[str, int], backend_plans: List[PlanConfig]) -> List[PlanConfig]: - """The ranked plan list for this graph: FROST and backend entries interleaved.""" - wanted = set(modes) - # OPENSOURCE wants these cells tried, so it asks for candidates like A does. - a_modes = {cudnn.heur_mode.A, cudnn.heur_mode.B, cudnn.heur_mode.OPENSOURCE} - ours_a = _mode_a(facts, offered) if wanted & a_modes else [] - ours_fb = _mode_fallback(facts, offered) if cudnn.heur_mode.FALLBACK in wanted else [] - - backend_a = [c for c in backend_plans if c.mode != cudnn.heur_mode.FALLBACK] - backend_fb = [c for c in backend_plans if c.mode == cudnn.heur_mode.FALLBACK] - - if cudnn.heur_mode.OPENSOURCE in wanted: - # The caller asked for an open-source implementation, which is what these - # cells are. Everything FROST can build goes first, regardless of which - # side measures faster -- that is the mode's whole meaning, and it is how - # a test measures FROST coverage: ask for OPENSOURCE, and whatever runs - # on a backend plan instead is a graph FROST does not cover yet. - return ours_a + ours_fb + backend_a + backend_fb - - # Fallbacks last on both sides: a fallback outranking any tuned candidate is - # what makes a plan walk take a slow kernel while a fast one was available. - if _leads(offered, ours_a): - return ours_a + backend_a + ours_fb + backend_fb - return backend_a + ours_a + ours_fb + backend_fb + """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/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 index e9e5f320f..d95157980 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_fp8_sm120.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_fp8_sm120.py @@ -26,23 +26,7 @@ from frost_test_utils import requires_blackwell_geforce, requires_dsl -def _select_engine(graph, name, tiles=None): - """Pin the ranked entry named ``name``. A pin is strict: check_support / - build_plans raise if that engine declines the graph. - - ``tiles`` pins one named knob entry instead of the delegation, so a test - can run a tile the shape-driven default would not choose.""" - names = [graph.get_plan_name_at_index(i) for i in range(len(graph.plans))] - if tiles is None: - assert name in names, f"engine {name!r} did not claim this graph; plans={names}" - index = names.index(name) - 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 - +from frost_test_utils import select_engine as _select_engine # noqa: F401 pytestmark = [requires_blackwell_geforce, requires_dsl] 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]: