From c6526fd271f2426d599bb373756d9a3a232d3930 Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Fri, 21 Aug 2026 20:55:45 -0700 Subject: [PATCH 1/3] sdpa fp8: serve the dense d<=128 envelope for per-tensor FP8 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per-tensor FP8 now runs head dims below the d128 tile through the same zero-padding ENVELOPE the f16/bf16 flavors use, rebuilt on the merged knob/capability framework (#692) and the d192-era native-shapes gate: - engines: the d128/d128 fp8 row declares d_envelope with d_pad_multiple=16 (TMA 16-byte global-stride rule at 1 byte/elem). New notch Capabilities.thd_d_envelope (default True); the fp8 row sets False — its packed THD compile key carries no head-dim entries, so the envelope is dense-only there and mismatch() says so at probe time. - adapter: check_support admits dense per-tensor D_QK/D_V <= 128, multiples of 16, alongside the native shapes; compile hands the kernel the ACTUAL head dims on the per-tensor d128 flavor. - kernels (both d128 fp8 siblings, SM100/SM107, in lockstep): compile() grows d_qk/d_v. TMA descriptors carry the real extents while the tile box stays the compile-time D: OOB loads zero-fill (exact in FP8) and O stores clip at d_v. The stride guard checks d_v against the fp8 input BPE (which subsumes the O side at BPE_O in {1, 2}). d192/d128 and MXFP8 stay exact-native (SF plumbing not audited for zero-padding). This is the landing zone for the ViT d=72-in-80 contract (e.g. Qwen3-VL vision encoders) without caller-side re-padding to 128; the descales are scalars, so the envelope is arch-independent. Co-Authored-By: Claude Fable 5 --- python/cudnn/sdpa/fwd/api_dsl.py | 50 ++++++++++----- python/cudnn/sdpa/fwd/engines.py | 25 +++++++- .../fwd/kernels/prefill_d128_fp8_sm100.py | 26 ++++++-- .../fwd/kernels/prefill_d128_fp8_sm107.py | 26 ++++++-- test/python/sdpa/frost/test_sdpa_fp8_sm107.py | 59 ++++++++++++++++++ .../sdpa/frost/test_sdpa_fwd_fp8_sm100.py | 62 ++++++++++++++++++- 6 files changed, 220 insertions(+), 28 deletions(-) diff --git a/python/cudnn/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py index db0bc38f6..7fde0677c 100644 --- a/python/cudnn/sdpa/fwd/api_dsl.py +++ b/python/cudnn/sdpa/fwd/api_dsl.py @@ -225,9 +225,11 @@ def _pick_flavor(d_qk: int, d_v: int) -> tuple[int, int]: the tile box stays the compile-time D, so loads past d_qk / d_v hardware zero-fill (adding exact zero terms to every QK^T dot product — S, softmax and P·V are bit-identical to the unpadded problem) and O stores past d_v - are OOB-clipped. FP8/MXFP8 uses exact native shapes (gated in - check_support); alignment (d % 8, the TMA 16-byte global-stride rule at - 2 bytes/elem) is also gated in check_support / engines.mismatch. + are OOB-clipped. Per-tensor FP8 serves the same envelope under its d128 + tile (dense only, d % 16 at 1 byte/elem); d192 FP8 and MXFP8 use exact + native shapes. All of it is gated in check_support / engines.mismatch, + including the f16 alignment rule (d % 8, the TMA 16-byte global-stride + rule at 2 bytes/elem). """ for flavor in _SM100_FLAVORS: fdqk, fdv = flavor @@ -855,13 +857,22 @@ def check_support(self) -> bool: f"SdpaFwdDslSm100 requires {_allowed_msg}; found SM{major}{minor} on {device}", ) - # FP8 paths use exact native shapes. SM100 supports d128/d128 and - # d192/d128; Rubin currently supports only per-tensor FP8 d128. + # FP8 native shapes: SM100 serves d128/d128 and d192/d128; Rubin + # currently serves only per-tensor FP8 d128. Per-tensor FP8 ALSO + # serves the dense d<=128 ENVELOPE under the d128 tile (TMA + # zero-padding, like the f16 flavors — exact in FP8, and the descales + # are scalars, so the envelope is arch-independent): head dims must be + # multiples of 16 (TMA 16-byte global-stride rule at 1 byte/elem). + # THD stays native-shape (the packed THD compile key carries no + # head-dim entries — engines.thd_d_envelope) and MXFP8 stays exact + # (SF plumbing not audited for zero-padding). fp8_shapes = _sm100_fp8_shapes(self._pertensor, self._device_cc) + _fp8_envelope_ok = self._pertensor and not self.thd and max(int(d_qk), int(d_v)) <= 128 and int(d_qk) % 16 == 0 and int(d_v) % 16 == 0 self._value_error_if( - self._fp8 and (int(d_qk), int(d_v)) not in fp8_shapes, - f"{'FP8' if self._pertensor else 'MXFP8'} (E4M3/E5M2 inputs) requires an exact native shape in {sorted(fp8_shapes)} " - f"(no envelope padding); got (D_QK={d_qk}, D_V={d_v})", + self._fp8 and (int(d_qk), int(d_v)) not in fp8_shapes and not _fp8_envelope_ok, + f"{'FP8' if self._pertensor else 'MXFP8'} (E4M3/E5M2 inputs) requires a native shape in {sorted(fp8_shapes)}" + + (" or the dense d128 envelope (D_QK/D_V <= 128, multiples of 16)" if self._pertensor else " (no envelope padding)") + + f"; got (D_QK={d_qk}, D_V={d_v})", ) # Envelope alignment gate: the TMA descriptors are built from the # actual tensor extents, and cuTensorMapEncodeTiled requires every @@ -1081,14 +1092,13 @@ def compile(self) -> None: # entries (see _thd_compile_kwargs). self._compiled_kernel = self._k_mod.compile(**self._thd_compile_kwargs()) elif self._fp8: - # FP8/MXFP8 kernels use exact native shapes (gated in check_support); - # their compile() has no envelope head-dim parameters. has_lse=False - # (no Stats output) compiles the LSE store out — no dummy buffer at - # any level (the amax_o atomicMax write is independent). A split - # REQUIRES the in-kernel LSE: the per-split LSE is the combine - # weight (the kernel skips its own amax write under a split; the - # combine reports the amax of the RECOMBINED O instead). - self._compiled_kernel = self._k_mod.compile( + # has_lse=False (no Stats output) compiles the LSE store out — no + # dummy buffer at any level (the amax_o atomicMax write is + # independent). A split REQUIRES the in-kernel LSE: the per-split + # LSE is the combine weight (the kernel skips its own amax write + # under a split; the combine reports the amax of the RECOMBINED O + # instead). + fp8_kwargs = dict( b=self.batch_size, qh=self.h_q, kh=self.h_kv, @@ -1096,6 +1106,14 @@ def compile(self) -> None: skv=self.s_k_max, has_lse=(self.lse_desc is not None) or self.split_kv > 1, ) + if self._pertensor and self.flavor == (128, 128): + # ENVELOPE (per-tensor d128 only): hand the kernel the ACTUAL + # head dims so its TMA descriptors carry the real extents + # (loads past them zero-fill — exact in FP8; O stores past d_v + # clip). The d192/d128 and MXFP8 kernels are exact-native + # (gated in check_support) and take no head-dim parameters. + fp8_kwargs.update(d_qk=self.head_dim_qk, d_v=self.head_dim_v) + self._compiled_kernel = self._k_mod.compile(**fp8_kwargs) else: # ENVELOPE: hand the f16/bf16 kernel the ACTUAL head dims so its # TMA descriptors carry the real extents (loads past them diff --git a/python/cudnn/sdpa/fwd/engines.py b/python/cudnn/sdpa/fwd/engines.py index d02ca81f1..71fa7190f 100644 --- a/python/cudnn/sdpa/fwd/engines.py +++ b/python/cudnn/sdpa/fwd/engines.py @@ -125,8 +125,14 @@ class Capabilities: d_envelope: bool = False # Alignment rule for envelope rows: graph head dims must be multiples of # this. 8 = the TMA 16-byte global-stride rule at 2 bytes/elem (SM100 DSL - # rows). 1 = no constraint (the SM80 lowering pads host-side). + # rows). 16 = the same rule at 1 byte/elem (per-tensor FP8). 1 = no + # constraint (the SM80 lowering pads host-side). d_pad_multiple: int = 8 + # Whether the THD (ragged) leg serves the same head-dim envelope. The + # per-tensor FP8 row sets False: its packed THD compile key carries no + # head-dim entries (native-tile contract), so the envelope is dense-only + # there and THD graphs must match the native dims exactly. + thd_d_envelope: bool = True dtypes: frozenset = frozenset({cudnn.data_type.HALF, cudnn.data_type.BFLOAT16}) # cudnn.data_type, see graph_analyzer 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) @@ -330,6 +336,11 @@ def mismatch(capabilities: Capabilities, facts: "ga.SdpaGraphFacts", knobs: Opti f"envelope zero-padding requires D_QK/D_V multiples of {m} (TMA 16-byte " f"global-stride constraint); graph has D_QK={facts.d_qk}/D_V={facts.d_v}" ) + if facts.thd and not capabilities.thd_d_envelope and (facts.d_qk not in capabilities.d_qk or facts.d_v not in capabilities.d_v): + return ( + f"THD (ragged) rides the packed native-tile leg on this engine (D_QK in " + f"{sorted(capabilities.d_qk)} / D_V in {sorted(capabilities.d_v)}); the head-dim envelope is dense-only" + ) elif facts.d_qk not in capabilities.d_qk or facts.d_v not in capabilities.d_v: 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.s_q == 1 and not capabilities.decode: @@ -553,7 +564,14 @@ def _sm100_fp8_spec( sink_dtypes: Optional[frozenset] = None, arch: str = "sm100", ) -> EngineSpec: - """Exact-shape per-tensor FP8 engine with scalar descales. + """Per-tensor FP8 engine with scalar descales. + + The d128/d128 cell serves the dense ``d <= 128`` ENVELOPE (TMA + zero-padding, like the f16 flavors — exact in FP8, and the descales are + scalars so no per-column plumbing is affected); head dims must be + multiples of 16 (the TMA 16-byte global-stride rule at 1 byte/elem). + d192/d128 stays exact-native, and THD keeps native dims on both cells + (the packed THD compile key carries no head-dim entries). One row per ARCH LINE (``arch``: "sm100" = pre-Rubin Blackwell 100-106, "sm107" = Rubin line 107-119), because the two lowerings genuinely @@ -600,6 +618,9 @@ def _sm100_fp8_spec( phase="prefill", d_qk=frozenset({d}), d_v=frozenset({d_v}), + d_envelope=(d, d_v) == (128, 128), + d_pad_multiple=16, + thd_d_envelope=False, dtypes=dtypes, out_dtypes=frozenset({cudnn.data_type.HALF, cudnn.data_type.BFLOAT16, cudnn.data_type.FP8_E4M3, cudnn.data_type.FP8_E5M2}), is_fp8=True, diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm100.py index ad3551db2..dccee4a34 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm100.py @@ -2204,9 +2204,19 @@ def compile( # noqa: A001 has_lse: bool = True, lse_head_major: bool = False, lse_head_stride: int = 0, + d_qk: int = CFG.TILE_K, + d_v: int = CFG.TILE_O, ) -> Callable: """Compile with ALL dims concrete — pins TMA strides at compile time. + ``d_qk``/``d_v`` <= the d128 tile serve the dense ENVELOPE: the TMA + descriptors carry the ACTUAL extents while the tile box stays the + compile-time D, so loads past them hardware zero-fill (exact in FP8 — + S/softmax/P·V are bit-identical to the unpadded problem) and O stores + past ``d_v`` are OOB-clipped. Head dims like the ViT d=72-in-80 contract + run without caller-side re-padding. THD serves native tile dims only + (the packed compile key carries no head-dim entries). + THD/varlen: q/k/v/o/lse PACKED with batch dim 1 ([1,T,H,D]); ``b`` is the LOGICAL batch (sequence count) driving n_batch / metadata + O-desc sizes. ``sq``/``skv`` are IGNORED under THD — the packed token totals are runtime @@ -2221,6 +2231,14 @@ def compile( # noqa: A001 # Each split's LSE is not optional under KV split — it IS the weight # the combine reduces with. Without it the partials cannot be recombined. raise ValueError("split_kv > 1 requires has_lse=True (the per-split LSE drives the combine)") + if not (0 < d_qk <= CFG.TILE_K and 0 < d_v <= CFG.TILE_O): + raise ValueError(f"fp8 d128 envelope: need 0 < d_qk <= {CFG.TILE_K} and 0 < d_v <= {CFG.TILE_O}; got ({d_qk}, {d_v})") + if (d_qk * CFG.BPE) % 16 != 0 or (d_v * CFG.BPE) % 16 != 0: + # d_v strides BOTH V (BPE) and O (BPE_O >= BPE); the fp8 input side is + # the binding TMA 16-byte global-stride constraint. + raise ValueError(f"fp8 d128 envelope: d_qk/d_v global strides must be 16-byte multiples (TMA rule at BPE={CFG.BPE}); got ({d_qk}, {d_v})") + if CFG.THD_VARLEN and (d_qk != CFG.TILE_K or d_v != CFG.TILE_O): + raise ValueError("THD/varlen serves native tile dims only (the packed compile key carries no head-dim entries); leave d_qk/d_v at the defaults") _fake_batch = 1 if CFG.THD_VARLEN else b # KV split: O and LSE are the PARTIAL workspaces, stacked split-major on # the batch axis (B*SPLIT_KV). Q/K/V keep the real batch. THD packs the @@ -2236,25 +2254,25 @@ def compile( # noqa: A001 skv = cute.sym_int(divisibility=1) fake_q = cute.runtime.make_fake_compact_tensor( STORAGE_DTYPE, - (_fake_batch, sq, qh, CFG.TILE_K), + (_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, CFG.TILE_K), + (_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, CFG.TILE_O), + (_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, - (_o_batch, sq, qh, CFG.TILE_O), + (_o_batch, sq, qh, d_v), stride_order=(3, 2, 1, 0), assumed_align=16, ) diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm107.py b/python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm107.py index bc469732f..82a096521 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm107.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm107.py @@ -2289,9 +2289,19 @@ def compile( # noqa: A001 has_lse: bool = True, lse_head_major: bool = False, lse_head_stride: int = 0, + d_qk: int = CFG.TILE_K, + d_v: int = CFG.TILE_O, ) -> Callable: """Compile with ALL dims concrete — pins TMA strides at compile time. + ``d_qk``/``d_v`` <= the d128 tile serve the dense ENVELOPE: the TMA + descriptors carry the ACTUAL extents while the tile box stays the + compile-time D, so loads past them hardware zero-fill (exact in FP8 — + S/softmax/P·V are bit-identical to the unpadded problem) and O stores + past ``d_v`` are OOB-clipped. Head dims like the ViT d=72-in-80 contract + run without caller-side re-padding. THD serves native tile dims only + (the packed compile key carries no head-dim entries). + THD/varlen: q/k/v/o/lse PACKED with batch dim 1 ([1,T,H,D]); ``b`` is the LOGICAL batch (sequence count) driving n_batch / metadata + O-desc sizes. ``sq``/``skv`` are IGNORED under THD — the packed token totals are runtime @@ -2302,6 +2312,14 @@ def compile( # noqa: A001 ``has_lse=False`` compiles the LSE store out (the kernel specializes on a ``None`` LSE argument) — callers without a Stats output pass no LSE buffer at all; the amax_o atomicMax write is independent and unchanged.""" + if not (0 < d_qk <= CFG.TILE_K and 0 < d_v <= CFG.TILE_O): + raise ValueError(f"fp8 d128 envelope: need 0 < d_qk <= {CFG.TILE_K} and 0 < d_v <= {CFG.TILE_O}; got ({d_qk}, {d_v})") + if (d_qk * CFG.BPE) % 16 != 0 or (d_v * CFG.BPE) % 16 != 0: + # d_v strides BOTH V (BPE) and O (BPE_O >= BPE); the fp8 input side is + # the binding TMA 16-byte global-stride constraint. + raise ValueError(f"fp8 d128 envelope: d_qk/d_v global strides must be 16-byte multiples (TMA rule at BPE={CFG.BPE}); got ({d_qk}, {d_v})") + if CFG.THD_VARLEN and (d_qk != CFG.TILE_K or d_v != CFG.TILE_O): + raise ValueError("THD/varlen serves native tile dims only (the packed compile key carries no head-dim entries); leave d_qk/d_v at the defaults") _fake_batch = 1 if CFG.THD_VARLEN else b if CFG.THD_VARLEN: # Dynamic packed token totals: one symbol per ragged group (Q/O and a @@ -2311,25 +2329,25 @@ def compile( # noqa: A001 skv = cute.sym_int(divisibility=1) fake_q = cute.runtime.make_fake_compact_tensor( STORAGE_DTYPE, - (_fake_batch, sq, qh, CFG.TILE_K), + (_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, CFG.TILE_K), + (_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, CFG.TILE_O), + (_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, CFG.TILE_O), + (_fake_batch, sq, qh, d_v), stride_order=(3, 2, 1, 0), assumed_align=16, ) diff --git a/test/python/sdpa/frost/test_sdpa_fp8_sm107.py b/test/python/sdpa/frost/test_sdpa_fp8_sm107.py index fb8b42f84..6378da083 100644 --- a/test/python/sdpa/frost/test_sdpa_fp8_sm107.py +++ b/test/python/sdpa/frost/test_sdpa_fp8_sm107.py @@ -270,3 +270,62 @@ def test_fp8_softmax_f16_e2e(): assert err <= 0.1 * ref.abs().max().item(), f"{precision}: max err {err} vs fp32 reference" xerr = (outs[_c.data_type.HALF] - outs[_c.data_type.FLOAT]).abs().max().item() assert xerr <= 0.05 * ref.abs().max().item(), f"HALF-vs-FLOAT softmax divergence {xerr}" + + +def test_fp8_d128_rows_serve_dense_envelope(): + """BOTH d128/d128 per-tensor FP8 rows (sm100 and sm107) serve the dense + d<=128 head-dim ENVELOPE (TMA zero-padding — exact in FP8; d % 16 at + 1 byte/elem, and arch-independent since the descales are scalars). THD + stays native-tile (the packed THD compile key carries no head-dim + entries) and the d192/d128 and MXFP8 rows stay exact-native.""" + from cudnn.sdpa.fwd import engines + + caps = {s.name: s.capabilities for s in engines.ENGINE_SPECS} + for arch in ("sm100", "sm107"): + row = caps[engines.engine_name(128, arch=arch, fp8=True)] + assert row.d_envelope, arch + assert row.d_pad_multiple == 16, arch + assert not row.thd_d_envelope, arch + assert not caps[engines.engine_name(192, d_v=128, fp8=True)].d_envelope + assert not caps[engines.engine_name(128, mxfp8=True)].d_envelope + + +def _fp8_facts(**kw): + import cudnn + from cudnn.sdpa import graph_analyzer as ga + + base = dict( + b=2, + h_q=4, + h_kv=4, + s_q=384, + s_kv=384, + d_qk=80, + d_v=80, + dtype=cudnn.data_type.FP8_E4M3, + dtype_o=cudnn.data_type.BFLOAT16, + is_fp8=True, + device_cc=(10, 0), + ) + base.update(kw) + return ga.SdpaGraphFacts(**base) + + +def test_fp8_envelope_mismatch_rules(): + """Honest eligibility for the fp8 envelope: mismatch() admits dense d80 on + the d128 rows of both arch lines, enforces d % 16, and keeps THD + native-tile — no plan may enter the ranked list only to die at build in + the adapter.""" + from cudnn.sdpa.fwd import engines + + caps = {s.name: s.capabilities for s in engines.ENGINE_SPECS} + sm100 = caps[engines.engine_name(128, fp8=True)] + assert engines.mismatch(sm100, _fp8_facts()) is None + assert engines.mismatch(sm100, _fp8_facts(d_qk=96, d_v=64)) is None + assert "multiples of 16" in engines.mismatch(sm100, _fp8_facts(d_qk=88, d_v=88)) + assert "dense-only" in engines.mismatch(sm100, _fp8_facts(thd=True, padded=True)) + assert engines.mismatch(sm100, _fp8_facts(d_qk=128, d_v=128, thd=True, padded=True)) is None + # The Rubin row serves the same dense envelope (the ViT d=72-in-80 case). + sm107 = caps[engines.engine_name(128, arch="sm107", fp8=True)] + assert engines.mismatch(sm107, _fp8_facts(device_cc=(10, 7))) is None + assert "dense-only" in engines.mismatch(sm107, _fp8_facts(device_cc=(10, 7), thd=True, padded=True)) 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 a8622dd96..a640e0aa9 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_fp8_sm100.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_fp8_sm100.py @@ -4,7 +4,8 @@ """End-to-end tests for the FROST SM100 DSL per-tensor FP8 SDPA-forward engine. Drives ``graph.sdpa_fp8`` (FP8 E4M3/E5M2 Q/K/V + scalar per-tensor descales) routed to -the exact d128/d128 or d192/d128 engine, and validates O against an fp32-dequant +the d128/d128 engine (which also serves the dense d<=128 head-dim ENVELOPE via TMA +zero-padding) or the exact d192/d128 engine, and validates O against an fp32-dequant reference. ``Amax_O`` is produced in-kernel (atomicMax over the pre-cast fp32 values) and checked. @@ -101,6 +102,7 @@ def _run( sync_debug=False, d_qk=128, d_v=128, + engine_shape=None, ): import cudnn @@ -163,7 +165,12 @@ def _stns(): g.validate() g.build_operation_graph() g.create_execution_plans([cudnn.heur_mode.A]) - _select_engine(g, engine_name(d_qk, arch=_D128_ARCH if (d_qk, d_v) == (128, 128) else "sm100", d_v=d_v, fp8=True)) + # engine_shape pins a row other than the graph's own dims — the d128 rows + # serve smaller head dims through their dense ENVELOPE (zero-padding). + # The arch pick rides the ENGINE dims: an enveloped d80 graph pins the + # d128 row of the device's arch line (sm100 or sm107). + e_qk, e_v = engine_shape if engine_shape is not None else (d_qk, d_v) + _select_engine(g, engine_name(e_qk, arch=_D128_ARCH if (e_qk, e_v) == (128, 128) else "sm100", d_v=e_v, fp8=True)) g.check_support() g.build_plans() if not stats: @@ -319,6 +326,57 @@ def test_fp8_d192_d128_zero_length_kv(): _check(out, o_ref, torch.float16, "e4m3", a_o, a_o_ref) +@pytest.mark.L0 +@pytest.mark.parametrize("mask", ["none", "causal"]) +@pytest.mark.parametrize("dims", [(80, 80), (96, 64)], ids=["d80", "d96_d64"]) +@torch_fork_set_rng(seed=0) +def test_fp8_head_dim_envelope(dims, mask): + """Per-tensor FP8 serves the dense d<=128 head-dim ENVELOPE on the d128 + row (TMA zero-padding — exact in FP8, arch-independent since the descales + are scalars): the ViT d=72-in-80 contract's landing zone. Head dims must + be multiples of 16 (TMA 16-byte global-stride rule at 1 byte/elem).""" + d_qk, d_v = dims + scale = 1.0 / math.sqrt(d_qk) + out, o_ref, a_o, a_o_ref = _run( + 2, + 4, + 4, + 384, + 384, + "e4m3", + torch.bfloat16, + scale=scale, + sdpa_kwargs=_MASKS[mask], + d_qk=d_qk, + d_v=d_v, + engine_shape=(128, 128), + ) + _check(out, o_ref, torch.bfloat16, "e4m3", a_o, a_o_ref) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +def test_fp8_head_dim_envelope_padded(): + """The envelope composed with the KV padding mask — the ViT production + shape (non-causal, seqlen padded up to a tile multiple, d=80).""" + out, o_ref, a_o, a_o_ref = _run( + 2, + 4, + 4, + 384, + 384, + "e4m3", + torch.bfloat16, + scale=1.0 / math.sqrt(72), + sdpa_kwargs={}, + seq_lens_kv=[384, 250], + d_qk=80, + d_v=80, + engine_shape=(128, 128), + ) + _check(out, o_ref, torch.bfloat16, "e4m3", a_o, a_o_ref) + + @pytest.mark.L0 @pytest.mark.parametrize("in_key", _INS) @torch_fork_set_rng(seed=0) From 02241042b477454d45025181f36fddba5ed86314 Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Sat, 22 Aug 2026 13:24:37 -0700 Subject: [PATCH 2/3] =?UTF-8?q?sdpa=20engines:=20one=20row=20per=20arch=20?= =?UTF-8?q?x=20dtype=20family=20=E2=80=94=20head=20dim=20is=20a=20lowering?= =?UTF-8?q?=20concern?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The 12-row table (per-head-dim cells) becomes 7 family rows: sdpa_fwd_prefill_{sm100, sm100_mxfp8, sm100_fp8, sm107_fp8, sm120, sm120_fp8, sm80}. Kernel-flavor choice (which head-dim tile) happens inside the lowering (api_dsl._pick_flavor, smallest covering flavor), not in the ranked list. Capabilities head-dim vocabulary shrinks to two fields: d_shapes (the native flavor shapes) and d_pad_multiple (envelope alignment; 0 = exact native shapes only — MXFP8, whose SF plumbing is not audited for zero-padding). d_envelope and the thd_d_envelope notch are deleted; thd_d_shapes / split_d_shapes express the flavors that carry the THD leg / SplitHelpers (the quantized families wire both in d128 only). The d192x128 fp8 kernel gains the same d_qk/d_v envelope params as the d128 siblings, so the per-tensor family is uniformly enveloped ((160, 96) e2e rides the d192 flavor). Honesty fixes that fall out as row data: - f16 and MXFP8 rows shrink to SM 100-106 (no Rubin lowering): Rubin graphs are ineligible at probe time instead of erroring at build, and the f16/MXFP8/stream/async/split suites gain a requires_pre_rubin_blackwell gate so the Rubin CI lane can widen its FROST_TEST_PATHS to the whole frost directory (the lane's config note asks exactly for this). - Fixes test_dispatch.py::test_every_engine_spec_has_a_manifest_slot, red on develop since #651: the sm107 row had no manifest slot, so FrostSdpaFwdEngines silently never built it — Rubin per-tensor FP8 was unreachable on the graph route. The family rows get fresh append-only slots (11-14); the per-head-dim slots are retired. - The adapter honors an explicit softmax_precision=FLOAT on every per-tensor flavor (the pipeline each already runs; place() hands it out from the row domain). HALF stays d128 + cc10.7. Cross-flavor plan ranking (running a small-d graph on a larger flavor via select_plan) is retired with the per-d rows; if flavor A/B testing is wanted it should return as a knob, not as engine identity. Co-Authored-By: Claude Fable 5 --- python/cudnn/engines/manifest.py | 17 +- python/cudnn/frost/README.md | 15 +- python/cudnn/sdpa/fwd/api_dsl.py | 44 +-- python/cudnn/sdpa/fwd/engines.py | 273 ++++++++---------- .../kernels/prefill_d192_d128_fp8_sm100.py | 32 +- test/python/sdpa/frost/frost_test_utils.py | 8 + .../sdpa/frost/test_sdpa_execute_is_async.py | 6 +- test/python/sdpa/frost/test_sdpa_fp8_sm107.py | 55 ++-- .../frost/test_sdpa_frontend_integration.py | 25 +- .../sdpa/frost/test_sdpa_fwd_dsl_sm100.py | 14 +- .../sdpa/frost/test_sdpa_fwd_fp8_sm100.py | 44 ++- .../sdpa/frost/test_sdpa_fwd_heuristics.py | 34 ++- .../sdpa/frost/test_sdpa_fwd_mxfp8_sm100.py | 10 +- .../frost/test_sdpa_fwd_split_kv_sm100.py | 4 +- .../sdpa/frost/test_sdpa_graph_analyzer.py | 48 +-- .../sdpa/frost/test_sdpa_stream_ordering.py | 4 +- .../sdpa/frost/test_sdpa_stream_respect.py | 4 +- .../sdpa/frost/test_split_kv_heuristic.py | 16 +- 18 files changed, 341 insertions(+), 312 deletions(-) diff --git a/python/cudnn/engines/manifest.py b/python/cudnn/engines/manifest.py index be040eb90..8ad0be2ce 100644 --- a/python/cudnn/engines/manifest.py +++ b/python/cudnn/engines/manifest.py @@ -174,18 +174,19 @@ def offered_ids(self) -> Dict[str, int]: "cudnn.sdpa.fwd.engine", "FrostSdpaFwdEngines", # Slots are FIXED FOREVER; append the next free one, never reorder. + # RETIRED (one engine per arch x dtype family absorbed the per-head-dim + # rows; kernel-flavor choice moved into the lowering — never reuse): + # 0 sm100_d128, 1 sm100_d256, 2 sm100_d512, 3 sm100_d128_mxfp8, + # 4 sm100_d128_fp8, 6 sm100_d192_d128, 9 sm100_d192_d128_fp8, + # 10 sm100_d192_d128_mxfp8 slots={ - "sdpa_fwd_prefill_sm100_d128": EngineSlot(0, opt_in=True), - "sdpa_fwd_prefill_sm100_d256": EngineSlot(1, opt_in=True), - "sdpa_fwd_prefill_sm100_d512": EngineSlot(2, opt_in=True), - "sdpa_fwd_prefill_sm100_d128_mxfp8": EngineSlot(3, opt_in=True), - "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), "sdpa_fwd_prefill_sm80": EngineSlot(8, opt_in=True), - "sdpa_fwd_prefill_sm100_d192_d128_fp8": EngineSlot(9, opt_in=True), - "sdpa_fwd_prefill_sm100_d192_d128_mxfp8": EngineSlot(10, opt_in=True), + "sdpa_fwd_prefill_sm100": EngineSlot(11, opt_in=True), + "sdpa_fwd_prefill_sm100_mxfp8": EngineSlot(12, opt_in=True), + "sdpa_fwd_prefill_sm100_fp8": EngineSlot(13, opt_in=True), + "sdpa_fwd_prefill_sm107_fp8": EngineSlot(14, opt_in=True), }, analyzer=("cudnn.sdpa.graph_analyzer", "analyze"), heuristics=("cudnn.sdpa.fwd.heuristics", "recommend"), diff --git a/python/cudnn/frost/README.md b/python/cudnn/frost/README.md index 193193f2e..c27fbe2c1 100644 --- a/python/cudnn/frost/README.md +++ b/python/cudnn/frost/README.md @@ -687,17 +687,18 @@ sdpa_bwd_sm100_d128 (future) one row serves several compute capabilities. The row's `Capabilities.arches` set is the source of truth for exactly which; the name never enumerates minors. -- Head dimensions are omitted when one engine accepts a domain of dimensions, - as the SM120 prefill engine does. `Capabilities.d_qk` and `d_v` are the - source of truth for that domain. -- Geometry-specific engines use `d` and append `x` only when the two - head dimensions differ. +- Head dimensions never appear in engine names: one engine per + arch x dtype family accepts a DOMAIN of dimensions and its lowering picks + the kernel flavor (the smallest native shape covering the graph). + `Capabilities.d_shapes` (native flavor shapes) plus `d_pad_multiple` + (envelope alignment; 0 = exact shapes only) are the source of truth for + that domain. - No version counters. If a genuinely distinct second engine ever serves the same cell, give it a descriptive variant suffix (e.g. `_cga4`), not a number. - Names are for humans; `engine_id` is for machines. Pin by index (`select_plan`) or replay by id -- never by parsing a name. -- `cudnn.sdpa.fwd.engines.engine_name(d)` computes geometry-specific names; - omit `d` for dimension-agnostic engines. +- `cudnn.sdpa.fwd.engines.engine_name(arch=..., fp8=..., mxfp8=...)` computes + the family names (test/user convenience). ## Kernel templates and TemplateParams diff --git a/python/cudnn/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py index 7fde0677c..46b8dd449 100644 --- a/python/cudnn/sdpa/fwd/api_dsl.py +++ b/python/cudnn/sdpa/fwd/api_dsl.py @@ -857,21 +857,23 @@ def check_support(self) -> bool: f"SdpaFwdDslSm100 requires {_allowed_msg}; found SM{major}{minor} on {device}", ) - # FP8 native shapes: SM100 serves d128/d128 and d192/d128; Rubin - # currently serves only per-tensor FP8 d128. Per-tensor FP8 ALSO - # serves the dense d<=128 ENVELOPE under the d128 tile (TMA - # zero-padding, like the f16 flavors — exact in FP8, and the descales - # are scalars, so the envelope is arch-independent): head dims must be - # multiples of 16 (TMA 16-byte global-stride rule at 1 byte/elem). - # THD stays native-shape (the packed THD compile key carries no - # head-dim entries — engines.thd_d_envelope) and MXFP8 stays exact - # (SF plumbing not audited for zero-padding). + # FP8 flavor shapes: SM100 serves d128/d128 and d192/d128; Rubin + # currently serves only per-tensor FP8 d128. Per-tensor FP8 serves + # the dense ENVELOPE of every flavor it has (TMA zero-padding, like + # the f16 flavors — exact in FP8, and the descales are scalars, so + # the envelope is arch-independent): head dims componentwise <= a + # flavor shape and multiples of 16 (TMA 16-byte global-stride rule + # at 1 byte/elem). THD stays native-shape (the packed THD compile + # key carries no head-dim entries — engines.thd_d_shapes) and MXFP8 + # stays exact (SF plumbing not audited for zero-padding). fp8_shapes = _sm100_fp8_shapes(self._pertensor, self._device_cc) - _fp8_envelope_ok = self._pertensor and not self.thd and max(int(d_qk), int(d_v)) <= 128 and int(d_qk) % 16 == 0 and int(d_v) % 16 == 0 + _fp8_envelope_ok = ( + self._pertensor and not self.thd and any(int(d_qk) <= sq and int(d_v) <= sv for sq, sv in fp8_shapes) and int(d_qk) % 16 == 0 and int(d_v) % 16 == 0 + ) self._value_error_if( self._fp8 and (int(d_qk), int(d_v)) not in fp8_shapes and not _fp8_envelope_ok, f"{'FP8' if self._pertensor else 'MXFP8'} (E4M3/E5M2 inputs) requires a native shape in {sorted(fp8_shapes)}" - + (" or the dense d128 envelope (D_QK/D_V <= 128, multiples of 16)" if self._pertensor else " (no envelope padding)") + + (" — or, dense only, its envelope (head dims <= a flavor shape, multiples of 16)" if self._pertensor else " (no envelope padding)") + f"; got (D_QK={d_qk}, D_V={d_v})", ) # Envelope alignment gate: the TMA descriptors are built from the @@ -905,8 +907,8 @@ def check_support(self) -> bool: from cudnn import data_type as _cudnn_dtype self._value_error_if( - self.softmax_precision is not None and not (self._fp8 and self._pertensor and self.flavor == (128, 128)), - "softmax_precision is served on the d128 per-tensor FP8 path only (other flavors run the f32 pipeline)", + self.softmax_precision is not None and not (self._fp8 and self._pertensor), + "softmax_precision is served on the per-tensor FP8 path only (other families run the f32 pipeline)", ) self._value_error_if( self.softmax_precision is not None and self.softmax_precision not in (_cudnn_dtype.FLOAT, _cudnn_dtype.HALF), @@ -917,8 +919,8 @@ def check_support(self) -> bool: # split engine rows: only sdpa_fwd_prefill_sm107_d128_fp8 declares # HALF in its softmax_precisions domain). self._value_error_if( - self.softmax_precision == _cudnn_dtype.HALF and self._device_cc != (10, 7), - "softmax_precision=HALF is served for per-tensor FP8 on cc10.7 only (FLOAT is the default everywhere)", + self.softmax_precision == _cudnn_dtype.HALF and (self._device_cc != (10, 7) or self.flavor != (128, 128)), + "softmax_precision=HALF is served for per-tensor FP8 d128 on cc10.7 only (FLOAT is the default everywhere)", ) if self.split_kv > 1: # Split-KV: partials weighted by the per-split LSE, recombined by @@ -1106,12 +1108,12 @@ def compile(self) -> None: skv=self.s_k_max, has_lse=(self.lse_desc is not None) or self.split_kv > 1, ) - if self._pertensor and self.flavor == (128, 128): - # ENVELOPE (per-tensor d128 only): hand the kernel the ACTUAL - # head dims so its TMA descriptors carry the real extents - # (loads past them zero-fill — exact in FP8; O stores past d_v - # clip). The d192/d128 and MXFP8 kernels are exact-native - # (gated in check_support) and take no head-dim parameters. + if self._pertensor: + # ENVELOPE (per-tensor only): hand the kernel the ACTUAL head + # dims so its TMA descriptors carry the real extents (loads + # past them zero-fill — exact in FP8; O stores past d_v clip). + # Both fp8 flavor kernels take these; MXFP8 is exact-native + # (gated in check_support) and takes no head-dim parameters. fp8_kwargs.update(d_qk=self.head_dim_qk, d_v=self.head_dim_v) self._compiled_kernel = self._k_mod.compile(**fp8_kwargs) else: diff --git a/python/cudnn/sdpa/fwd/engines.py b/python/cudnn/sdpa/fwd/engines.py index 71fa7190f..22d549b55 100644 --- a/python/cudnn/sdpa/fwd/engines.py +++ b/python/cudnn/sdpa/fwd/engines.py @@ -113,26 +113,29 @@ class Capabilities: sm_lo: int sm_hi: int phase: str - d_qk: frozenset[int] - d_v: frozenset[int] - # Head-dim ENVELOPE: when True the lowering ALSO serves any graph with - # d_qk/d_v <= the native caps (both multiples of 8, the TMA 16-byte - # global-stride rule at 2 bytes/elem) via TMA zero-padding — the kernel's - # descriptors carry the ACTUAL extents, so padded contraction columns load - # as exact zeros (S/softmax unchanged) and O stores past d_v are - # OOB-clipped. False = only the native dims above are eligible (MXFP8, - # whose SF plumbing is not audited for zero-padding). - d_envelope: bool = False - # Alignment rule for envelope rows: graph head dims must be multiples of - # this. 8 = the TMA 16-byte global-stride rule at 2 bytes/elem (SM100 DSL - # rows). 16 = the same rule at 1 byte/elem (per-tensor FP8). 1 = no - # constraint (the SM80 lowering pads host-side). + # Head-dim DOMAIN. One engine per arch x dtype family: the head dim is a + # LOWERING concern — the adapter picks the smallest kernel flavor whose + # native shape covers the graph (api_dsl._pick_flavor) — not an engine + # identity. ``d_shapes`` is the set of NATIVE flavor shapes (d_qk, d_v) + # that lowering picks among. + d_shapes: frozenset + # Envelope + alignment rule. When > 0: any graph (d_qk, d_v) componentwise + # <= some native shape AND a multiple of this is served via TMA + # zero-padding — the kernel's descriptors carry the ACTUAL extents, so + # padded contraction columns load as exact zeros (S/softmax unchanged) and + # O stores past d_v are OOB-clipped. 8 = the TMA 16-byte global-stride + # rule at 2 bytes/elem (f16/bf16), 16 = the same rule at 1 byte/elem + # (per-tensor FP8), 1 = no constraint (the SM80 lowering pads host-side). + # 0 = NO envelope: exact native shapes only (MXFP8, whose SF plumbing is + # not audited for zero-padding). d_pad_multiple: int = 8 - # Whether the THD (ragged) leg serves the same head-dim envelope. The - # per-tensor FP8 row sets False: its packed THD compile key carries no - # head-dim entries (native-tile contract), so the envelope is dense-only - # there and THD graphs must match the native dims exactly. - thd_d_envelope: bool = True + # Shapes whose kernels carry the THD leg. None = THD (when the ``thd`` + # capability is set) serves the same head-dim domain as dense — the f16 + # THD compile key carries the head dims, so THD rides the envelope. A set + # = THD graphs must match one of these NATIVE shapes exactly: the + # quantized rows' packed THD compile key carries no head-dim entries + # (native-tile contract), so their envelope is dense-only. + thd_d_shapes: Optional[frozenset] = None dtypes: frozenset = frozenset({cudnn.data_type.HALF, cudnn.data_type.BFLOAT16}) # cudnn.data_type, see graph_analyzer 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) @@ -240,6 +243,11 @@ class Capabilities: # whose kernels wire the split path AND whose adapter launches the combine # widen this (the SM100 f16 rows today). split_kvs: frozenset[int] = frozenset({1}) + # Shapes whose kernel flavors wire SplitHelpers. None = every flavor in + # d_shapes does (f16/SM120). A set = split_kv > 1 is honored only when + # the graph's dims are covered by a member (the quantized families wire + # the split path in the d128 flavor only). + split_d_shapes: Optional[frozenset] = None # Softmax-precision domain (cudnn.data_type values). Empty = unserved. # Arch-dependent membership (the f16x2 exponent arm exists only in the # SM107 sibling kernel) is expressed by SPLITTING the row per arch line — @@ -311,6 +319,10 @@ def mismatch(capabilities: Capabilities, facts: "ga.SdpaGraphFacts", knobs: Opti # QUANTIZED partials would lose what the split must be # numerically neutral about. return "split_kv > 1 on a quantized graph requires a bf16/fp16 O" + if capabilities.split_d_shapes is not None and not any(facts.d_qk <= sq and facts.d_v <= sv for sq, sv in capabilities.split_d_shapes): + return ( + f"split_kv > 1 is wired only in the {sorted(capabilities.split_d_shapes)} kernel " f"flavors; graph has D_QK={facts.d_qk}/D_V={facts.d_v}" + ) cc = facts.device_cc sm = None if cc is None else cc[0] * 10 + cc[1] if sm is None or not (capabilities.sm_lo <= sm <= capabilities.sm_hi): @@ -323,26 +335,26 @@ def mismatch(capabilities: Capabilities, facts: "ga.SdpaGraphFacts", knobs: Opti if cutedsl_too_old(version): want = ".".join(str(v) for v in CUTEDSL_MIN_VERSION) return f"requires nvidia-cutlass-dsl >= {want}; found {version[1]}" - if capabilities.d_envelope: - # Envelope row: native caps are upper bounds (TMA zero-padding semantics - # — see Capabilities.d_envelope). Alignment: TMA global strides must be - # 16-byte multiples; compact BSHD H-stride is D * 2 bytes -> D % 8. - cap_qk, cap_v = max(capabilities.d_qk), max(capabilities.d_v) - if facts.d_qk > cap_qk or facts.d_v > cap_v: - return f"serves D_QK<={cap_qk}/D_V<={cap_v} (envelope); graph has D_QK={facts.d_qk}/D_V={facts.d_v}" + shapes = sorted(capabilities.d_shapes) + if capabilities.d_pad_multiple: + # Envelope family: native flavor shapes are upper bounds (TMA + # zero-padding semantics — see Capabilities.d_shapes/d_pad_multiple); + # the lowering picks the smallest covering flavor. + if not any(facts.d_qk <= sq and facts.d_v <= sv for sq, sv in capabilities.d_shapes): + return f"no kernel-flavor envelope covers (D_QK={facts.d_qk}, D_V={facts.d_v}); native shapes: {shapes}" m = capabilities.d_pad_multiple if m > 1 and (facts.d_qk % m != 0 or facts.d_v % m != 0): return ( f"envelope zero-padding requires D_QK/D_V multiples of {m} (TMA 16-byte " f"global-stride constraint); graph has D_QK={facts.d_qk}/D_V={facts.d_v}" ) - if facts.thd and not capabilities.thd_d_envelope and (facts.d_qk not in capabilities.d_qk or facts.d_v not in capabilities.d_v): - return ( - f"THD (ragged) rides the packed native-tile leg on this engine (D_QK in " - f"{sorted(capabilities.d_qk)} / D_V in {sorted(capabilities.d_v)}); the head-dim envelope is dense-only" - ) - elif facts.d_qk not in capabilities.d_qk or facts.d_v not in capabilities.d_v: - 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}" + elif (facts.d_qk, facts.d_v) not in capabilities.d_shapes: + return f"serves exact native shapes {shapes} (no envelope padding); graph has D_QK={facts.d_qk}/D_V={facts.d_v}" + if facts.thd and capabilities.thd_d_shapes is not None and (facts.d_qk, facts.d_v) not in capabilities.thd_d_shapes: + return ( + f"THD (ragged) rides the packed native-tile leg on this engine " + f"(shapes {sorted(capabilities.thd_d_shapes)}); the head-dim envelope is dense-only" + ) if facts.s_q == 1 and not capabilities.decode: return "s_q == 1 (decode) is out of scope for the SM80 prefill kernels" if facts.dtype not in capabilities.dtypes: @@ -460,18 +472,19 @@ class EngineSpec: lower: "Callable[[EngineSpec, ga.SdpaGraphFacts, Optional[SdpaFwdKnobs]], Any]" -def _sm100_spec(d: int, d_v: Optional[int] = None) -> EngineSpec: - d_v = d if d_v is None else d_v - suffix = f"d{d}" if d_v == d else f"d{d}_d{d_v}" +def _sm100_spec() -> EngineSpec: + """f16/bf16 SM100-family engine: ONE row; the adapter picks the smallest + kernel flavor (d128 / d192xd128 / d256 / d512) covering the graph's head + dims (api_dsl._pick_flavor), and every flavor serves its envelope via TMA + zero-padding. sm_hi=106: no f16 lowering exists on the Rubin line — when + one lands it gets its own row (the per-arch-line row doctrine).""" return EngineSpec( - name=f"sdpa_fwd_prefill_sm100_{suffix}", + name="sdpa_fwd_prefill_sm100", capabilities=Capabilities( sm_lo=_BLACKWELL[0], - sm_hi=_BLACKWELL[1], + sm_hi=106, phase="prefill", - d_qk=frozenset({d}), - d_v=frozenset({d_v}), - d_envelope=True, # native tile box d; smaller dims via TMA zero-padding + d_shapes=frozenset({(128, 128), (192, 128), (256, 256), (512, 512)}), dtypes=frozenset({cudnn.data_type.HALF, cudnn.data_type.BFLOAT16}), causal=True, bottom_right=True, @@ -509,7 +522,7 @@ def _sm100_spec(d: int, d_v: Optional[int] = None) -> EngineSpec: ) -def _sm100_mxfp8_spec(d: int, d_v: Optional[int] = None) -> EngineSpec: +def _sm100_mxfp8_spec() -> EngineSpec: """Block-scale MXFP8 engine (E4M3/E5M2 + per-32-block E8M0 SF). THD/varlen (d128/d128 only — the d192/d128 kernel is dense-only) rides the @@ -520,17 +533,20 @@ def _sm100_mxfp8_spec(d: int, d_v: Optional[int] = None) -> EngineSpec: SF dims stay the dense capacity, like the ragged Q/K/V storage. """ - d_v = d if d_v is None else d_v - suffix = f"d{d}" if d_v == d else f"d{d}_d{d_v}" - thd = (d, d_v) == (128, 128) return EngineSpec( - name=f"sdpa_fwd_prefill_sm100_{suffix}_mxfp8", + name="sdpa_fwd_prefill_sm100_mxfp8", capabilities=Capabilities( sm_lo=_BLACKWELL[0], - sm_hi=_BLACKWELL[1], + sm_hi=106, # no Rubin MXFP8 lowering phase="prefill", - d_qk=frozenset({d}), - d_v=frozenset({d_v}), + # Exact native shapes only (d_pad_multiple=0): the SF plumbing is + # not audited for envelope zero-padding. + d_shapes=frozenset({(128, 128), (192, 128)}), + d_pad_multiple=0, + # Only the d128 kernel carries the write_thd_meta THD leg and + # wires SplitHelpers; the d192x128 file is dense-only. + thd_d_shapes=frozenset({(128, 128)}), + split_d_shapes=frozenset({(128, 128)}), 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, @@ -542,53 +558,48 @@ def _sm100_mxfp8_spec(d: int, d_v: Optional[int] = None) -> EngineSpec: sink=True, stats=True, lse_optional=True, - thd=thd, - cu_seq_len=thd, + thd=True, + cu_seq_len=True, sched_policies=frozenset({SCHED_NATURAL, SCHED_LPT, SCHED_LPT_L2}), tile_ms=frozenset({128}), tile_ns=frozenset({128}), cgas=frozenset({2}), - # Only the d128 mxfp8 kernel wires SplitHelpers; the split path - # also needs a half-precision O (mismatch's facts x knobs gate). - split_kvs=frozenset({1, 2, 4}) if d == 128 else frozenset({1}), + # The split path also needs a half-precision O (mismatch's + # facts x knobs gate) and rides the d128 flavor (split_d_shapes). + split_kvs=frozenset({1, 2, 4}), ), lower=partial(lower_dsl_prefill, api_type=_SM100), ) -def _sm100_fp8_spec( - d: int, - d_v: Optional[int] = None, - *, - dtypes: Optional[frozenset] = None, - sink_dtypes: Optional[frozenset] = None, - arch: str = "sm100", -) -> EngineSpec: +def _sm100_fp8_spec(*, arch: str = "sm100") -> EngineSpec: """Per-tensor FP8 engine with scalar descales. - The d128/d128 cell serves the dense ``d <= 128`` ENVELOPE (TMA - zero-padding, like the f16 flavors — exact in FP8, and the descales are - scalars so no per-column plumbing is affected); head dims must be - multiples of 16 (the TMA 16-byte global-stride rule at 1 byte/elem). - d192/d128 stays exact-native, and THD keeps native dims on both cells - (the packed THD compile key carries no head-dim entries). - - One row per ARCH LINE (``arch``: "sm100" = pre-Rubin Blackwell 100-106, - "sm107" = Rubin line 107-119), because the two lowerings genuinely - diverge and a shared row could only describe their union with knob x arch - notches. Each row declares exactly what its own kernel carries: - + ONE row per ARCH LINE (``arch``: "sm100" = pre-Rubin Blackwell 100-106, + "sm107" = Rubin line 107-119): the two lowerings genuinely diverge and a + shared row could only describe their union with knob x arch notches. + Within a row the head dim is a LOWERING concern — the adapter picks the + kernel flavor (d128 or d192xd128) covering the graph. Each row declares + exactly what its own kernels carry: + + - d_shapes: the sm100 row picks between the d128 and d192xd128 flavors; + Rubin has only the d128 sibling, so a Rubin d192 graph is ineligible + at probe time instead of a late build error. + - The ENVELOPE (d_pad_multiple=16, the TMA 16-byte global-stride rule at + 1 byte/elem): smaller head dims ride TMA zero-padding — exact in FP8, + and the descales are scalars so no per-column plumbing is affected. + THD keeps native dims (thd_d_shapes: the packed THD compile key + carries no head-dim entries). - softmax_precisions: the f16x2 exponent arm lives only in the SM107 - sibling kernel, so only that row admits HALF. - - split_kvs: only the SM100 d128 kernel wires SplitHelpers; the SM107 - sibling has no split path yet. + sibling kernel, so only that row admits HALF. FLOAT is the pipeline + every flavor already runs. + - split_kvs / split_d_shapes: only the SM100 d128 kernel wires + SplitHelpers; the SM107 sibling has no split path yet, and the + d192x128 file forks its own scheduler and has none either. - sched_policies: the LPT/LPT_L2 remap is not yet ported to the SM107 sibling (issue #653); {NATURAL} keeps requests honest AND routes the graph path around the un-ported derivation (place() hands the adapter an explicit policy from this domain). - - d192/d128 exists only on the sm100 row (no Rubin d192 kernel), so a - Rubin d192 graph is now ineligible at probe time instead of a late - build error. Padding mask (per-batch ``seq_len_kv`` → KV-side masking) is supported: KV-only padding leaves every query row real, so each row's total_sum > 0 and the @@ -600,28 +611,20 @@ def _sm100_fp8_spec( the same THD leg. """ - d_v = d if d_v is None else d_v - suffix = f"d{d}" if d_v == d else f"d{d}_d{d_v}" - thd = (d, d_v) == (128, 128) rubin_row = arch == "sm107" - assert not (rubin_row and (d, d_v) != (128, 128)), "Rubin serves only per-tensor FP8 d128" - if dtypes is None: - dtypes = frozenset({cudnn.data_type.FP8_E4M3, cudnn.data_type.FP8_E5M2}) return EngineSpec( - name=f"sdpa_fwd_prefill_{arch}_{suffix}_fp8", + name=f"sdpa_fwd_prefill_{arch}_fp8", capabilities=Capabilities( # Ranges, not the parts that exist today (see sm_lo above): the - # split point is the Rubin line — 100-106 runs the SM100 module, + # split point is the Rubin line — 100-106 runs the SM100 modules, # 107-119 the SM107 sibling. sm_lo=107 if rubin_row else _BLACKWELL[0], sm_hi=_BLACKWELL[1] if rubin_row else 106, phase="prefill", - d_qk=frozenset({d}), - d_v=frozenset({d_v}), - d_envelope=(d, d_v) == (128, 128), + d_shapes=frozenset({(128, 128)}) if rubin_row else frozenset({(128, 128), (192, 128)}), d_pad_multiple=16, - thd_d_envelope=False, - dtypes=dtypes, + thd_d_shapes=frozenset({(128, 128)}), + 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, @@ -630,11 +633,10 @@ def _sm100_fp8_spec( swa=True, padded=True, sink=True, - sink_dtypes=sink_dtypes, stats=True, lse_optional=True, - thd=thd, - cu_seq_len=thd, + thd=True, + cu_seq_len=True, # The fp8 kernel lacks the SEQ_Q_LENS_PRESENT epilogue trim, but its # only reachable padded+stats population is KV-only padding with # full-length seq_len_q (the fp8 suite; test_mhas_v2 fp8 padding @@ -655,16 +657,13 @@ def _sm100_fp8_spec( cgas=frozenset({2}), # f16x2-softmax arm: only the SM107 sibling kernel carries the # path (MUFU EX2.F16x2 exists below cc10.7 but no other file wires - # it). FLOAT is the f32 pipeline every serving row already runs. - softmax_precisions=( - frozenset({cudnn.data_type.FLOAT, cudnn.data_type.HALF}) if rubin_row else (frozenset({cudnn.data_type.FLOAT}) if d == 128 else frozenset()) - ), - # Only the SM100 d128 fp8 kernel wires SplitHelpers (the SM107 - # sibling has no split path yet; the d192x128 file forks its own - # scheduler and has none either). Split partials reduce in half - # precision, so mismatch()'s facts x knobs gate additionally - # requires a bf16/fp16 O on the quantized rows. - split_kvs=frozenset({1, 2, 4}) if (d == 128 and not rubin_row) else frozenset({1}), + # it). FLOAT is the f32 pipeline every flavor already runs. + softmax_precisions=(frozenset({cudnn.data_type.FLOAT, cudnn.data_type.HALF}) if rubin_row else frozenset({cudnn.data_type.FLOAT})), + # Split partials reduce in half precision, so mismatch()'s + # facts x knobs gate additionally requires a bf16/fp16 O on the + # quantized rows; split_d_shapes pins it to the d128 flavor. + split_kvs=frozenset({1}) if rubin_row else frozenset({1, 2, 4}), + split_d_shapes=frozenset({(128, 128)}), ), lower=partial(lower_dsl_prefill, api_type=_SM100), ) @@ -684,9 +683,7 @@ def _sm80_spec() -> EngineSpec: sm_lo=80, sm_hi=80, # A100 exactly: the kernels assume its 164 KiB opt-in SMEM phase="prefill", - d_qk=frozenset({256}), - d_v=frozenset({256}), - d_envelope=True, # flavor envelopes; host-side zero-padding + d_shapes=frozenset({(256, 256)}), # flavor envelopes; host-side zero-padding d_pad_multiple=1, dtypes=frozenset({cudnn.data_type.HALF, cudnn.data_type.BFLOAT16}), bias=True, @@ -724,9 +721,9 @@ def _sm120_spec() -> EngineSpec: sm_lo=_BLACKWELL_GEFORCE[0], sm_hi=_BLACKWELL_GEFORCE[1], phase="prefill", - d_qk=frozenset(SUPPORTED_HEAD_TILES), - d_v=frozenset(SUPPORTED_HEAD_TILES), - d_envelope=True, + # The kernel picks its Q/K and V head tiles independently, so the + # native shapes are the cross product of the supported tiles. + d_shapes=frozenset((tq, tv) for tq in SUPPORTED_HEAD_TILES for tv in SUPPORTED_HEAD_TILES), dtypes=frozenset({cudnn.data_type.HALF, cudnn.data_type.BFLOAT16}), causal=True, bottom_right=True, @@ -1029,36 +1026,29 @@ def _execute(variant_pack, workspace=None, stream=None): def engine_name( - d: Optional[int] = None, phase: str = "prefill", arch: str = "sm100", mxfp8: bool = False, fp8: bool = False, - d_v: Optional[int] = None, ) -> str: - """The registered engine name for a coverage cell (test/user convenience).""" - - name = f"sdpa_fwd_{phase}_{arch}" - if d is not None: - if d_v is None or d_v == d: - name += f"_d{d}" - else: - name += f"_d{d}_d{d_v}" + """The registered engine name for a coverage cell (test/user convenience). + + One engine per arch x dtype family — head dims are a lowering concern + (kernel-flavor pick), not part of the engine identity.""" + suffix = "_mxfp8" if mxfp8 else "_fp8" if fp8 else "" - return name + suffix + return f"sdpa_fwd_{phase}_{arch}" + suffix # ORDER MATTERS: this is the PREFERENCE order — ``engine.FrostSdpaFwdEngines()`` # wraps the specs in it, so the plans they propose reach graph.plans in this -# order and the build walk tries them top-down. With head-dim envelopes every -# small-d graph is eligible for ALL covering f16 flavors, so the f16 rows are -# listed smallest-first to make first-eligible == smallest-covering flavor -# (the tightest tile, least padded work). We deliberately keep mismatch() a -# pure upper-bound check instead of rejecting non-tightest flavors: an -# explicit select_plan() of a larger covering flavor stays legal (useful for -# flavor A/B testing), only the ranking prefers the tightest. +# order and the build walk tries them top-down. One engine per arch x dtype +# family: kernel-FLAVOR choice (which head-dim tile) happens inside the +# lowering (api_dsl._pick_flavor, smallest covering flavor — the tightest +# tile, least padded work), so at most one row of a family is eligible per +# device and the order only breaks ties across families. # Engine IDS do NOT follow this order — they are pinned per name in -# engine._ID_OFFSETS and never move. +# engines/manifest.py and never move. def _sm120_fp8_spec() -> EngineSpec: """SM120 per-tensor FP8 engine (E4M3/E5M2 in + scalar descales, FP16/BF16/FP8 out). @@ -1093,9 +1083,7 @@ def _sm120_fp8_spec() -> EngineSpec: sm_lo=_BLACKWELL_GEFORCE[0], sm_hi=_BLACKWELL_GEFORCE[1], phase="prefill", - d_qk=frozenset(SUPPORTED_HEAD_TILES_FP8), - d_v=frozenset(SUPPORTED_HEAD_TILES_FP8), - d_envelope=True, # native tile box d; smaller dims via TMA zero-padding + d_shapes=frozenset((tq, tv) for tq in SUPPORTED_HEAD_TILES_FP8 for tv in SUPPORTED_HEAD_TILES_FP8), d_pad_multiple=16, # TMA 16-byte global-stride rule at 1 byte/elem 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}), @@ -1127,19 +1115,10 @@ def _sm120_fp8_spec() -> EngineSpec: ENGINE_SPECS = ( - _sm100_spec(128), - _sm100_spec(192, d_v=128), - _sm100_spec(256), - _sm100_spec(512), - _sm100_mxfp8_spec(128), - _sm100_mxfp8_spec(192, d_v=128), - _sm100_fp8_spec(128), - _sm100_fp8_spec(128, arch="sm107"), - _sm100_fp8_spec( - 192, - d_v=128, - dtypes=frozenset({cudnn.data_type.FP8_E4M3, cudnn.data_type.FP8_E5M2}), - ), + _sm100_spec(), + _sm100_mxfp8_spec(), + _sm100_fp8_spec(), + _sm100_fp8_spec(arch="sm107"), _sm120_spec(), _sm120_fp8_spec(), _sm80_spec(), diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_fp8_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_fp8_sm100.py index 3979b4767..a63680b37 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_fp8_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_fp8_sm100.py @@ -2257,33 +2257,55 @@ def _tma_swz(byte_w: int): @lru_cache(maxsize=None) -def compile(b: int = 1, qh: int = 1, kh: int = 1, sq: int = 256, skv: int = 128, has_lse: bool = True) -> Callable: # noqa: A001 +def compile( # noqa: A001 + b: int = 1, + qh: int = 1, + kh: int = 1, + sq: int = 256, + skv: int = 128, + has_lse: bool = True, + d_qk: int = CFG.TILE_K, + d_v: int = CFG.TILE_O, +) -> Callable: """Compile with ALL dims concrete — pins TMA strides at compile time. + ``d_qk``/``d_v`` <= the native (192, 128) tile serve the dense ENVELOPE: + the TMA descriptors carry the ACTUAL extents while the tile box stays the + compile-time D, so loads past them hardware zero-fill (exact in FP8 — + S/softmax/P·V are bit-identical to the unpadded problem, including the + statically-offset second Q/K chunk) and O stores past ``d_v`` are + OOB-clipped. + ``has_lse=False`` specializes the LSE argument to ``None`` and removes the Stats store while retaining the independent amax writes.""" + if not (0 < d_qk <= CFG.TILE_K and 0 < d_v <= CFG.TILE_O): + raise ValueError(f"fp8 d192 envelope: need 0 < d_qk <= {CFG.TILE_K} and 0 < d_v <= {CFG.TILE_O}; got ({d_qk}, {d_v})") + if (d_qk * CFG.BPE) % 16 != 0 or (d_v * CFG.BPE) % 16 != 0: + # d_v strides BOTH V (BPE) and O (BPE_O >= BPE); the fp8 input side is + # the binding TMA 16-byte global-stride constraint. + raise ValueError(f"fp8 d192 envelope: d_qk/d_v global strides must be 16-byte multiples (TMA rule at BPE={CFG.BPE}); got ({d_qk}, {d_v})") _fake_batch = b fake_q = cute.runtime.make_fake_compact_tensor( STORAGE_DTYPE, - (_fake_batch, sq, qh, CFG.TILE_K), + (_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, CFG.TILE_K), + (_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, CFG.TILE_O), + (_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, CFG.TILE_O), + (_fake_batch, sq, qh, d_v), stride_order=(3, 2, 1, 0), assumed_align=16, ) diff --git a/test/python/sdpa/frost/frost_test_utils.py b/test/python/sdpa/frost/frost_test_utils.py index 26472c42e..bb1c1915f 100644 --- a/test/python/sdpa/frost/frost_test_utils.py +++ b/test/python/sdpa/frost/frost_test_utils.py @@ -28,6 +28,14 @@ def _active_sm(): _SM is None or not (100 <= _SM <= 119), reason="needs an SM100-line GPU (100 <= SM <= 119), have " + ("none" if _SM is None else f"sm_{_SM}"), ) +# Pre-Rubin gate for the suites whose lowerings do not exist on the Rubin +# line (f16/bf16 and MXFP8 SM100 paths; Rubin serves per-tensor FP8 only) — +# these must SKIP on cc10.7 so the Rubin CI lane can run the whole frost +# directory (the lane's FROST_TEST_PATHS note asks exactly for this). +requires_pre_rubin_blackwell = pytest.mark.skipif( + _SM is None or not (100 <= _SM <= 106), + reason="needs a pre-Rubin SM100-line GPU (100 <= SM <= 106; no f16/MXFP8 Rubin lowering), have " + ("none" if _SM is None else f"sm_{_SM}"), +) requires_blackwell_geforce = pytest.mark.skipif( _SM is None or not (120 <= _SM <= 129), reason="needs an SM120-line GPU, have " + ("none" if _SM is None else f"sm_{_SM}"), diff --git a/test/python/sdpa/frost/test_sdpa_execute_is_async.py b/test/python/sdpa/frost/test_sdpa_execute_is_async.py index 3fa0d27a1..affe89f0b 100644 --- a/test/python/sdpa/frost/test_sdpa_execute_is_async.py +++ b/test/python/sdpa/frost/test_sdpa_execute_is_async.py @@ -30,7 +30,7 @@ import cudnn from cudnn.engines import is_python_engine -from frost_test_utils import requires_blackwell, requires_dsl +from frost_test_utils import requires_pre_rubin_blackwell, requires_dsl pytestmark = [pytest.mark.L0] @@ -88,7 +88,7 @@ def _build(d): return g, vp, ws, o_buf -@requires_blackwell +@requires_pre_rubin_blackwell @requires_dsl @pytest.mark.parametrize("d", [256, 512]) def test_execute_reads_no_device_memory_to_the_host(monkeypatch, d): @@ -126,7 +126,7 @@ def guard(*a, **kw): assert not caught -@requires_blackwell +@requires_pre_rubin_blackwell @requires_dsl def test_execute_without_a_handle_is_cuda_graph_capturable(): """The no-handle path, which is the one that resolves the stream itself. diff --git a/test/python/sdpa/frost/test_sdpa_fp8_sm107.py b/test/python/sdpa/frost/test_sdpa_fp8_sm107.py index 6378da083..1121a93dc 100644 --- a/test/python/sdpa/frost/test_sdpa_fp8_sm107.py +++ b/test/python/sdpa/frost/test_sdpa_fp8_sm107.py @@ -67,14 +67,15 @@ def test_per_tensor_fp8_rows_split_per_arch_line(): from cudnn.sdpa.fwd import engines caps = {s.name: s.capabilities for s in engines.ENGINE_SPECS} - sm100 = caps[engines.engine_name(128, fp8=True)] - sm107 = caps[engines.engine_name(128, arch="sm107", fp8=True)] - d192 = caps[engines.engine_name(192, d_v=128, fp8=True)] + sm100 = caps[engines.engine_name(fp8=True)] + sm107 = caps[engines.engine_name(arch="sm107", fp8=True)] # Arch ranges tile the SM100 family at the Rubin boundary, no overlap. assert (sm100.sm_lo, sm100.sm_hi) == (100, 106) assert (sm107.sm_lo, sm107.sm_hi) == (107, 119) - assert (d192.sm_lo, d192.sm_hi) == (100, 106) # no Rubin d192 kernel + # Kernel flavors are row DATA: Rubin has no d192 sibling. + assert sm100.d_shapes == frozenset({(128, 128), (192, 128)}) + assert sm107.d_shapes == frozenset({(128, 128)}) # The f16x2 exponent arm is Rubin-row data, not a notch. assert sm100.softmax_precisions == frozenset({_c.data_type.FLOAT}) @@ -113,8 +114,8 @@ def facts(cc): ) caps = {s.name: s.capabilities for s in engines.ENGINE_SPECS} - sm100 = caps[engines.engine_name(128, fp8=True)] - sm107 = caps[engines.engine_name(128, arch="sm107", fp8=True)] + sm100 = caps[engines.engine_name(fp8=True)] + sm107 = caps[engines.engine_name(arch="sm107", fp8=True)] half = engines.SdpaFwdKnobs(softmax_precision=_c.data_type.HALF) assert "domain" in engines.mismatch(sm100, facts((10, 0)), half) @@ -184,17 +185,14 @@ def test_softmax_points_never_propose_half(): sm_lo=100, sm_hi=119, phase="prefill", - d_qk=frozenset({128}), - d_v=frozenset({128}), + d_shapes=frozenset({(128, 128)}), softmax_precisions=frozenset({_c.data_type.FLOAT, _c.data_type.HALF}), ) assert _softmax_points(lit) == [_c.data_type.FLOAT] - dark = Capabilities(sm_lo=100, sm_hi=119, phase="prefill", d_qk=frozenset({128}), d_v=frozenset({128})) + dark = Capabilities(sm_lo=100, sm_hi=119, phase="prefill", d_shapes=frozenset({(128, 128)})) assert _softmax_points(dark) == [None] # A HALF-only row must still not get HALF auto-proposed (numerics-changing). - half_only = Capabilities( - sm_lo=100, sm_hi=119, phase="prefill", d_qk=frozenset({128}), d_v=frozenset({128}), softmax_precisions=frozenset({_c.data_type.HALF}) - ) + half_only = Capabilities(sm_lo=100, sm_hi=119, phase="prefill", d_shapes=frozenset({(128, 128)}), softmax_precisions=frozenset({_c.data_type.HALF})) assert _softmax_points(half_only) == [None] @@ -272,22 +270,20 @@ def test_fp8_softmax_f16_e2e(): assert xerr <= 0.05 * ref.abs().max().item(), f"HALF-vs-FLOAT softmax divergence {xerr}" -def test_fp8_d128_rows_serve_dense_envelope(): - """BOTH d128/d128 per-tensor FP8 rows (sm100 and sm107) serve the dense - d<=128 head-dim ENVELOPE (TMA zero-padding — exact in FP8; d % 16 at - 1 byte/elem, and arch-independent since the descales are scalars). THD - stays native-tile (the packed THD compile key carries no head-dim - entries) and the d192/d128 and MXFP8 rows stay exact-native.""" +def test_fp8_rows_serve_dense_envelope(): + """BOTH per-tensor FP8 rows (sm100 and sm107) serve the dense head-dim + ENVELOPE of their kernel flavors (TMA zero-padding — exact in FP8; + d % 16 at 1 byte/elem, and arch-independent since the descales are + scalars). THD stays native-tile (the packed THD compile key carries no + head-dim entries) and MXFP8 stays exact-native (d_pad_multiple=0).""" from cudnn.sdpa.fwd import engines caps = {s.name: s.capabilities for s in engines.ENGINE_SPECS} for arch in ("sm100", "sm107"): - row = caps[engines.engine_name(128, arch=arch, fp8=True)] - assert row.d_envelope, arch + row = caps[engines.engine_name(arch=arch, fp8=True)] assert row.d_pad_multiple == 16, arch - assert not row.thd_d_envelope, arch - assert not caps[engines.engine_name(192, d_v=128, fp8=True)].d_envelope - assert not caps[engines.engine_name(128, mxfp8=True)].d_envelope + assert row.thd_d_shapes == frozenset({(128, 128)}), arch + assert caps[engines.engine_name(mxfp8=True)].d_pad_multiple == 0 def _fp8_facts(**kw): @@ -319,13 +315,20 @@ def test_fp8_envelope_mismatch_rules(): from cudnn.sdpa.fwd import engines caps = {s.name: s.capabilities for s in engines.ENGINE_SPECS} - sm100 = caps[engines.engine_name(128, fp8=True)] + sm100 = caps[engines.engine_name(fp8=True)] assert engines.mismatch(sm100, _fp8_facts()) is None assert engines.mismatch(sm100, _fp8_facts(d_qk=96, d_v=64)) is None + # The d192xd128 flavor serves its envelope too (kernel takes d_qk/d_v). + assert engines.mismatch(sm100, _fp8_facts(d_qk=160, d_v=96)) is None + assert "no kernel-flavor envelope" in engines.mismatch(sm100, _fp8_facts(d_qk=160, d_v=160)) assert "multiples of 16" in engines.mismatch(sm100, _fp8_facts(d_qk=88, d_v=88)) assert "dense-only" in engines.mismatch(sm100, _fp8_facts(thd=True, padded=True)) assert engines.mismatch(sm100, _fp8_facts(d_qk=128, d_v=128, thd=True, padded=True)) is None - # The Rubin row serves the same dense envelope (the ViT d=72-in-80 case). - sm107 = caps[engines.engine_name(128, arch="sm107", fp8=True)] + # THD at the d192 native shape is dense-only (thd_d_shapes = {(128, 128)}). + assert "dense-only" in engines.mismatch(sm100, _fp8_facts(d_qk=192, d_v=128, thd=True, padded=True)) + # The Rubin row serves the same dense envelope (the ViT d=72-in-80 case) + # but has no d192 flavor at all. + sm107 = caps[engines.engine_name(arch="sm107", fp8=True)] assert engines.mismatch(sm107, _fp8_facts(device_cc=(10, 7))) is None + assert "no kernel-flavor envelope" in engines.mismatch(sm107, _fp8_facts(device_cc=(10, 7), d_qk=192, d_v=128)) assert "dense-only" in engines.mismatch(sm107, _fp8_facts(device_cc=(10, 7), thd=True, padded=True)) diff --git a/test/python/sdpa/frost/test_sdpa_frontend_integration.py b/test/python/sdpa/frost/test_sdpa_frontend_integration.py index 4da7cffcf..75dd84a53 100644 --- a/test/python/sdpa/frost/test_sdpa_frontend_integration.py +++ b/test/python/sdpa/frost/test_sdpa_frontend_integration.py @@ -15,14 +15,14 @@ from cudnn.engines import MANIFEST, is_backend_engine, is_python_engine from cudnn.sdpa.fwd.engines import engine_name -from frost_test_utils import requires_blackwell, requires_dsl, _dsl_installed, _is_plan_for +from frost_test_utils import requires_pre_rubin_blackwell, requires_dsl, _dsl_installed, _is_plan_for -_FROST = engine_name(512) # matches the D=512 graphs below +_FROST = engine_name() # matches the D=512 graphs below _GPU = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs GPU") -_SM100 = requires_blackwell -_SM100_DSL = pytest.mark.skipif(requires_blackwell.args[0] or requires_dsl.args[0], reason="needs an SM100-line GPU with the cutedsl extra") +_SM100 = requires_pre_rubin_blackwell +_SM100_DSL = pytest.mark.skipif(requires_pre_rubin_blackwell.args[0] or requires_dsl.args[0], reason="needs an SM100-line GPU with the cutedsl extra") # The default pytest.ini addopts is `-m L0`; mark the whole module so it runs. pytestmark = pytest.mark.L0 @@ -176,10 +176,11 @@ def test_default_plan_runs_and_matches_torch(): @_SM100_DSL -def test_envelope_lists_every_covering_flavor_smallest_first(): - """Head-dim ENVELOPE: a (64,64) graph is eligible for every f16 flavor, and - the family proposes them smallest-covering-first, so the tightest tile is - the first python entry of the ranked list; the padded run matches torch.""" +def test_envelope_serves_small_dims_through_one_family_engine(): + """Head-dim ENVELOPE: a (64,64) graph is served by the ONE f16 family + engine — kernel-flavor choice (which head-dim tile) happens inside the + lowering (api_dsl._pick_flavor, smallest covering flavor), not in the + ranked list; the padded run matches torch.""" d = 64 q_gpu = torch.randn(B, S, H, d, device="cuda", dtype=torch.float16).transpose(1, 2) k_gpu = torch.randn(B, S, H, d, device="cuda", dtype=torch.float16).transpose(1, 2) @@ -189,14 +190,12 @@ def test_envelope_lists_every_covering_flavor_smallest_first(): g, q, k, v, o = _build_causal_sdpa(d=d) _plan(g) names = _plan_names(g) - for flavor in (128, 256, 512): - assert any(_is_plan_for(n, engine_name(flavor)) for n in names) # every covering flavor + assert any(_is_plan_for(n, engine_name()) for n in names) python = [i for i, p in enumerate(g.plans) if is_python_engine(p.engine_id)] - assert _is_plan_for(names[python[0]], engine_name(128)) # tightest flavor first g.select_plan(python[0]) g.check_support() g.build_plans() - assert g.selected_engine.name == engine_name(128) + assert g.selected_engine.name == engine_name() ws = torch.empty(max(g.get_workspace_size(), 1), device="cuda", dtype=torch.uint8) g.execute({q: q_gpu, k: k_gpu, v: v_gpu, o: o_gpu}, ws) @@ -245,7 +244,7 @@ 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" - "i = next((i for i, n in enumerate(names) if n.split('[')[0] == 'sdpa_fwd_prefill_sm100_d512'), None)\n" + "i = next((i for i, n in enumerate(names) if n.split('[')[0] == 'sdpa_fwd_prefill_sm100'), None)\n" "assert i is not None, names\n" "g.select_plan(i)\n" "g.check_support()\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 92996e21b..7504c7faf 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py @@ -12,12 +12,12 @@ from test_utils import torch_fork_set_rng from cudnn.sdpa.fwd.engines import engine_name -from frost_test_utils import requires_blackwell, requires_dsl, _dsl_installed +from frost_test_utils import requires_pre_rubin_blackwell, requires_dsl, _dsl_installed from frost_test_utils import select_engine as _select_engine # noqa: F401 -pytestmark = requires_blackwell +pytestmark = requires_pre_rubin_blackwell def _ref_sdpa(q, k, v, *, is_causal, scale): @@ -81,7 +81,7 @@ def test_sdpa_fwd_dsl_sm100_graph_api(dtype, is_causal, d): graph.validate() graph.build_operation_graph() graph.create_execution_plans([cudnn.heur_mode.A]) - _select_engine(graph, engine_name(d)) + _select_engine(graph, engine_name()) graph.check_support() graph.build_plans() # Honest workspace: no Stats output, so the kernel compiles the LSE store @@ -206,7 +206,7 @@ def _run_dsl_graph(q_gpu, k_gpu, v_gpu, *, scale, dtype, sdpa_kwargs, seq_len_kv g.validate() g.build_operation_graph() g.create_execution_plans([cudnn.heur_mode.A]) - _select_engine(g, engine_name(q_gpu.shape[-1], d_v=d_v)) + _select_engine(g, engine_name()) g.check_support() g.build_plans() vp[o] = o_gpu @@ -638,7 +638,7 @@ def _dense_buf(packed): g.validate() g.build_operation_graph() g.create_execution_plans([cudnn.heur_mode.A]) - _select_engine(g, engine_name(d)) + _select_engine(g, engine_name()) g.check_support() g.build_plans() vp = {tq: q_gpu, tk: k_gpu, tv: v_gpu, o: o_gpu, sq: slq, skv: slk, qro: ro, kro: ro, vro: ro, oro: ro} @@ -728,7 +728,7 @@ def _dense_buf(packed, s_max, t): g.validate() g.build_operation_graph() g.create_execution_plans([cudnn.heur_mode.A]) - _select_engine(g, engine_name(d)) + _select_engine(g, engine_name()) g.check_support() g.build_plans() vp = {tq: q_gpu, tk: k_gpu, tv: v_gpu, o: o_gpu, sq: slq, skv: slk, qro: ro_q, kro: ro_k, vro: ro_k, oro: ro_q} @@ -898,7 +898,7 @@ def _dense_buf(packed, s_max, t, H): g.validate() g.build_operation_graph() g.create_execution_plans([cudnn.heur_mode.A]) - _select_engine(g, engine_name(d)) + _select_engine(g, engine_name()) g.check_support() g.build_plans() vp[o] = o_gpu 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 a640e0aa9..d884647ce 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_fp8_sm100.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_fp8_sm100.py @@ -33,11 +33,11 @@ pytestmark = [requires_blackwell, requires_dsl] -# The per-tensor FP8 rows are split per arch line (sm100 = pre-Rubin -# Blackwell 100-106, sm107 = the Rubin line): pin the row that serves the -# device under test. d192/d128 exists on the sm100 row only. +# ONE per-tensor FP8 engine per arch line (sm100 = pre-Rubin Blackwell +# 100-106, sm107 = the Rubin line): pin the engine that serves the device +# under test. The d192xd128 kernel flavor exists on the sm100 engine only. _D128_ARCH = "sm107" if _SM == 107 else "sm100" -_skip_on_rubin = pytest.mark.skipif(_SM == 107, reason="d192/d128 per-tensor FP8 has no Rubin kernel (the sm107 row serves d128 only)") +_skip_on_rubin = pytest.mark.skipif(_SM == 107, reason="the d192xd128 per-tensor FP8 flavor has no Rubin kernel (sm107 serves d128 only)") _FP8 = {"e4m3": torch.float8_e4m3fn, "e5m2": torch.float8_e5m2} _FP8_MAX = {"e4m3": 448.0, "e5m2": 57344.0} @@ -102,7 +102,6 @@ def _run( sync_debug=False, d_qk=128, d_v=128, - engine_shape=None, ): import cudnn @@ -165,12 +164,9 @@ def _stns(): g.validate() g.build_operation_graph() g.create_execution_plans([cudnn.heur_mode.A]) - # engine_shape pins a row other than the graph's own dims — the d128 rows - # serve smaller head dims through their dense ENVELOPE (zero-padding). - # The arch pick rides the ENGINE dims: an enveloped d80 graph pins the - # d128 row of the device's arch line (sm100 or sm107). - e_qk, e_v = engine_shape if engine_shape is not None else (d_qk, d_v) - _select_engine(g, engine_name(e_qk, arch=_D128_ARCH if (e_qk, e_v) == (128, 128) else "sm100", d_v=e_v, fp8=True)) + # ONE fp8 engine per arch line — kernel-flavor choice (d128 vs d192xd128, + # and the head-dim envelope) happens inside the lowering. + _select_engine(g, engine_name(arch=_D128_ARCH, fp8=True)) g.check_support() g.build_plans() if not stats: @@ -349,7 +345,28 @@ def test_fp8_head_dim_envelope(dims, mask): sdpa_kwargs=_MASKS[mask], d_qk=d_qk, d_v=d_v, - engine_shape=(128, 128), + ) + _check(out, o_ref, torch.bfloat16, "e4m3", a_o, a_o_ref) + + +@pytest.mark.L0 +@_skip_on_rubin +@torch_fork_set_rng(seed=0) +def test_fp8_head_dim_envelope_d192_flavor(): + """The d192xd128 flavor serves its envelope too: (160, 96) rides the + (192, 128) kernel with TMA zero-padding on both sides.""" + out, o_ref, a_o, a_o_ref = _run( + 2, + 4, + 4, + 384, + 384, + "e4m3", + torch.bfloat16, + scale=1.0 / math.sqrt(160), + sdpa_kwargs={}, + d_qk=160, + d_v=96, ) _check(out, o_ref, torch.bfloat16, "e4m3", a_o, a_o_ref) @@ -372,7 +389,6 @@ def test_fp8_head_dim_envelope_padded(): seq_lens_kv=[384, 250], d_qk=80, d_v=80, - engine_shape=(128, 128), ) _check(out, o_ref, torch.bfloat16, "e4m3", a_o, a_o_ref) @@ -569,7 +585,7 @@ def _stns(): g.validate() g.build_operation_graph() g.create_execution_plans([cudnn.heur_mode.A]) - _select_engine(g, engine_name(128, arch=_D128_ARCH, fp8=True)) + _select_engine(g, engine_name(arch=_D128_ARCH, fp8=True)) g.check_support() g.build_plans() vp.update({o: o_gpu, amx_o: amax_o}) diff --git a/test/python/sdpa/frost/test_sdpa_fwd_heuristics.py b/test/python/sdpa/frost/test_sdpa_fwd_heuristics.py index 1e3940d6d..7f34147ed 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_heuristics.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_heuristics.py @@ -25,8 +25,8 @@ from cudnn.sdpa.fwd.heuristics import _MAX_SETS_PER_ENGINE, recommend from cudnn.sdpa.graph_analyzer import SdpaGraphFacts -_D128 = "sdpa_fwd_prefill_sm100_d128" -_OFFERED = {_D128: 20500, "sdpa_fwd_prefill_sm100_d256": 20501} +_F16 = "sdpa_fwd_prefill_sm100" +_OFFERED = {_F16: 20500, "sdpa_fwd_prefill_sm100_fp8": 20501} def _facts(**over): @@ -50,15 +50,15 @@ def _facts(**over): @pytest.mark.L0 def test_recommend_emits_multiple_complete_sets_per_engine(): plans = recommend("A", _facts(), _OFFERED) - d128 = [p for p in plans if p.engine_id == 20500] - assert len(d128) >= 3, "expected sched + split runners behind the primary" - for p in d128: + f16 = [p for p in plans if p.engine_id == 20500] + assert len(f16) >= 3, "expected sched + split runners behind the primary" + for p in f16: k = p.knobs # Complete assignment: every axis the row declares carries a value. assert None not in (k.sched_policy, k.tile_m, k.tile_n, k.cga, k.split_kv) assert p.mode is None and p.cpp_index is None - assert len({p.knobs for p in d128}) == len(d128), "duplicate knob sets emitted" - assert len(d128) <= _MAX_SETS_PER_ENGINE + assert len({p.knobs for p in f16}) == len(f16), "duplicate knob sets emitted" + assert len(f16) <= _MAX_SETS_PER_ENGINE @pytest.mark.L0 @@ -69,9 +69,9 @@ def test_recommend_primary_reproduces_the_derived_scheduler(): causal = recommend("A", _facts(), _OFFERED) assert causal[0].knobs.sched_policy == 2 # SCHED_LPT_L2 dense = recommend("A", _facts(causal=False), _OFFERED) - dense_d128 = [p for p in dense if p.engine_id == 20500] - assert dense_d128[0].knobs.sched_policy == 0 # SCHED_NATURAL - assert all(p.knobs.sched_policy == 0 for p in dense_d128), "mask-free graphs gain nothing from LPT runners" + dense_f16 = [p for p in dense if p.engine_id == 20500] + assert dense_f16[0].knobs.sched_policy == 0 # SCHED_NATURAL + assert all(p.knobs.sched_policy == 0 for p in dense_f16), "mask-free graphs gain nothing from LPT runners" @pytest.mark.L0 @@ -129,10 +129,12 @@ def test_fallback_kind_is_least_demanding(): def _is_sm100() -> bool: + # Pre-Rubin only: the executable tier drives the f16 family, which has no + # Rubin lowering (Rubin serves per-tensor FP8 only). if not torch.cuda.is_available(): return False - major, _ = torch.cuda.get_device_capability(0) - return major == 10 + major, minor = torch.cuda.get_device_capability(0) + return major == 10 and minor <= 6 def _dsl_available() -> bool: @@ -185,9 +187,9 @@ def test_split_kv_plan_pinned_by_name_matches_reference(): if want == 1: pytest.skip("this part is small enough that the shape already fills it") names = [g.get_plan_name_at_index(i) for i in range(len(g.plans))] - d128 = [n for n in names if "sm100_d128" in n] - assert len(d128) >= 3, f"expected knob-suffixed duplicates of the d128 cell: {d128}" - split_idx = next(i for i, n in enumerate(names) if "sm100_d128" in n and f"split_kv={want}" in n) + f16 = [n for n in names if n.split("[")[0] == "sdpa_fwd_prefill_sm100"] + assert len(f16) >= 3, f"expected knob-suffixed duplicates of the f16 family engine: {f16}" + split_idx = next(i for i, n in enumerate(names) if n.split("[")[0] == "sdpa_fwd_prefill_sm100" and f"split_kv={want}" in n) g.select_plan(split_idx) g.check_support() g.build_plans() @@ -225,7 +227,7 @@ def test_runner_up_sched_plan_builds_and_matches_the_winner(): that set and executes correctly — honored, not silently degraded.""" g, (q, k, v, o, st), (B, H, SQ, SKV, D) = _build_decodeish_graph() names = [g.get_plan_name_at_index(i) for i in range(len(g.plans))] - nat_idx = next(i for i, n in enumerate(names) if "sm100_d128" in n and "sched_policy=0" in n and "split_kv=1" in n) + nat_idx = next(i for i, n in enumerate(names) if n.split("[")[0] == "sdpa_fwd_prefill_sm100" and "sched_policy=0" in n and "split_kv=1" in n) g.select_plan(nat_idx) g.check_support() g.build_plans() 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 e63644e82..57f5c3acc 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_mxfp8_sm100.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_mxfp8_sm100.py @@ -32,12 +32,12 @@ from test_utils import torch_fork_set_rng from cudnn.sdpa.fwd.engines import engine_name -from frost_test_utils import requires_blackwell, requires_dsl +from frost_test_utils import requires_pre_rubin_blackwell, requires_dsl from frost_test_utils import select_engine as _select_engine # noqa: F401 -pytestmark = [requires_blackwell, requires_dsl] +pytestmark = [requires_pre_rubin_blackwell, requires_dsl] _FP8 = {"e4m3": torch.float8_e4m3fn, "e5m2": torch.float8_e5m2} @@ -167,7 +167,7 @@ def _sf(dims): g.validate() g.build_operation_graph() g.create_execution_plans([cudnn.heur_mode.A]) - _select_engine(g, engine_name(d_qk, d_v=d_v, mxfp8=True)) + _select_engine(g, engine_name(mxfp8=True)) g.check_support() g.build_plans() if not stats: @@ -449,7 +449,7 @@ def _sf(dims): g.validate() g.build_operation_graph() g.create_execution_plans([cudnn.heur_mode.A]) - _select_engine(g, engine_name(d_qk, d_v=d_v, mxfp8=True)) + _select_engine(g, engine_name(mxfp8=True)) g.check_support() g.build_plans() g.execute( @@ -670,7 +670,7 @@ def _sf(dims): g.validate() g.build_operation_graph() g.create_execution_plans([cudnn.heur_mode.A]) - _select_engine(g, engine_name(128, mxfp8=True)) + _select_engine(g, engine_name(mxfp8=True)) g.check_support() g.build_plans() vp.update({o: o_gpu, amax_o: amax}) diff --git a/test/python/sdpa/frost/test_sdpa_fwd_split_kv_sm100.py b/test/python/sdpa/frost/test_sdpa_fwd_split_kv_sm100.py index 4934d222e..4502e21ff 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_split_kv_sm100.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_split_kv_sm100.py @@ -18,9 +18,9 @@ import pytest import torch -from frost_test_utils import requires_blackwell, requires_dsl +from frost_test_utils import requires_pre_rubin_blackwell, requires_dsl -pytestmark = [requires_blackwell, requires_dsl] +pytestmark = [requires_pre_rubin_blackwell, requires_dsl] D = 128 TILE_N = 128 diff --git a/test/python/sdpa/frost/test_sdpa_graph_analyzer.py b/test/python/sdpa/frost/test_sdpa_graph_analyzer.py index 18469a90a..79bd0e659 100644 --- a/test/python/sdpa/frost/test_sdpa_graph_analyzer.py +++ b/test/python/sdpa/frost/test_sdpa_graph_analyzer.py @@ -77,7 +77,8 @@ def test_engines_registered(): (row,) = [r for r in MANIFEST if r.factory == "FrostSdpaFwdEngines"] assert is_python_engine(row.engine_id) assert row.id_end - row.engine_id >= len(engines.ENGINE_SPECS) - assert engines.engine_name(512) == "sdpa_fwd_prefill_sm100_d512" + assert engines.engine_name() == "sdpa_fwd_prefill_sm100" + assert engines.engine_name(arch="sm107", fp8=True) == "sdpa_fwd_prefill_sm107_fp8" def test_single_sdpa_node_found(): @@ -99,7 +100,7 @@ def test_probe_accepts_dsv4_causal(): q, k, v, dims, strides = _mk_qkv(g) o, _ = g.sdpa(name="s", q=q, k=k, v=v, attn_scale=0.1, is_inference=True, use_causal_mask=True) _finish_output(o, dims, strides) - assert engines.engine_name(512) in _eligible(g) + assert engines.engine_name() in _eligible(g) def test_probe_accepts_bf16(): @@ -115,7 +116,7 @@ def test_probe_accepts_bf16(): v = g.tensor(dim=dims, stride=strides, data_type=cudnn.data_type.BFLOAT16, name="v") o, _ = g.sdpa(name="s", q=q, k=k, v=v, attn_scale=0.1, is_inference=True) _finish_output(o, dims, strides, dtype=cudnn.data_type.BFLOAT16) - assert engines.engine_name(512) in _eligible(g) + assert engines.engine_name() in _eligible(g) def test_probe_rejects_uncoverable_head_dim(): @@ -144,9 +145,9 @@ def test_probe_envelope_covers_small_head_dim(): o, _ = g.sdpa(name="s", q=q, k=k, v=v, attn_scale=0.1, is_inference=True) _finish_output(o, dims, strides) elig = _eligible(g) - assert {engines.engine_name(128), engines.engine_name(256), engines.engine_name(512)} <= elig + assert engines.engine_name() in elig ordered = [s.name for s in engines.ENGINE_SPECS if s.name in elig] - assert ordered[0] == engines.engine_name(128) + assert ordered[0] == engines.engine_name() def test_probe_envelope_mixed_dims_pick_covering_flavor(): @@ -160,14 +161,13 @@ def test_probe_envelope_mixed_dims_pick_covering_flavor(): o, _ = g.sdpa(name="s", q=q, k=k, v=v, attn_scale=0.1, is_inference=True) _finish_output(o, (B, H, S, d_v), (S * H * d_v, d_v, H * d_v, 1)) elig = _eligible(g) - assert engines.engine_name(128) not in elig - assert {engines.engine_name(192, d_v=128), engines.engine_name(256), engines.engine_name(512)} <= elig + assert engines.engine_name() in elig ordered = [s.name for s in engines.ENGINE_SPECS if s.name in elig] - assert ordered[0] == engines.engine_name(192, d_v=128) + assert ordered[0] == engines.engine_name() def test_d192_fp8_sink_dtype_support(): - spec = next(s for s in engines.ENGINE_SPECS if s.name == engines.engine_name(192, d_v=128, fp8=True)) + spec = next(s for s in engines.ENGINE_SPECS if s.name == engines.engine_name(fp8=True)) def facts(dtype, *, sink): return ga.SdpaGraphFacts( @@ -252,7 +252,7 @@ def test_probe_accepts_seq_len_q_with_padding_mask(): seq_len_q=seq_q, ) _finish_output(o, dims, strides) - assert engines.engine_name(512) in _eligible(g) + assert engines.engine_name() in _eligible(g) def test_probe_rejects_non_int32_seq_len(): @@ -297,7 +297,7 @@ def test_probe_accepts_bottom_right_with_padded_seq_len_q(): seq_len_q=seq_q, ) _finish_output(o, dims, strides) - assert engines.engine_name(512) in _eligible(g) + assert engines.engine_name() in _eligible(g) def test_probe_rejects_seq_len_q_without_padding_mask(): @@ -340,7 +340,7 @@ def _mk_thd_qkvo(g, *, mask_kwargs, d: int = D): def test_probe_accepts_thd_top_left_causal(): g = _mk_graph() _mk_thd_qkvo(g, mask_kwargs=dict(use_causal_mask=True)) - assert engines.engine_name(512) in _eligible(g) + assert engines.engine_name() in _eligible(g) def test_probe_accepts_thd_bottom_right(): @@ -348,7 +348,7 @@ def test_probe_accepts_thd_bottom_right(): # own (seq_len_q[b], seq_len_kv[b]) via the cu_seqlen metadata. g = _mk_graph() _mk_thd_qkvo(g, mask_kwargs=dict(use_causal_mask_bottom_right=True)) - assert engines.engine_name(512) in _eligible(g) + assert engines.engine_name() in _eligible(g) def test_probe_accepts_thd_stats(): @@ -385,7 +385,7 @@ def test_probe_accepts_thd_stats(): stats.set_data_type(cudnn.data_type.FLOAT) stats_ro = g.tensor(dim=(B + 1, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT64, name="stats_ro") stats.set_ragged_offset(stats_ro) - assert engines.engine_name(512) in _eligible(g) + assert engines.engine_name() in _eligible(g) def test_probe_accepts_right_band_widening(): @@ -395,7 +395,7 @@ def test_probe_accepts_right_band_widening(): q, k, v, dims, strides = _mk_qkv(g) o, _ = g.sdpa(name="s", q=q, k=k, v=v, attn_scale=0.1, is_inference=True, diagonal_band_right_bound=16) _finish_output(o, dims, strides) - assert engines.engine_name(512) in _eligible(g) + assert engines.engine_name() in _eligible(g) facts = ga.analyze(g) assert facts.right_band_widening and facts.right_bound == 16 and not facts.causal @@ -487,7 +487,7 @@ def test_resolve_padding_mask_with_seq_len_kv(): _finish_output(o, dims, strides) cfg = _facts(g) assert cfg.padded is True - assert engines.engine_name(512) in _eligible(g) + assert engines.engine_name() in _eligible(g) def test_resolve_generate_stats(): @@ -533,7 +533,7 @@ def test_probe_accepts_ragged_skv_via_synth_padding(): v = g.tensor(dim=(B, H, s_kv, D), stride=(s_kv * H * D, D, H * D, 1), data_type=DTYPE, name="v") o, _ = g.sdpa(name="s", q=q, k=k, v=v, attn_scale=0.1, is_inference=True) _finish_output(o, (B, H, S, D), (S * H * D, D, H * D, 1)) - assert engines.engine_name(512) in _eligible(g) + assert engines.engine_name() in _eligible(g) def test_probe_accepts_ragged_skv_with_top_left_causal(): @@ -545,7 +545,7 @@ def test_probe_accepts_ragged_skv_with_top_left_causal(): v = g.tensor(dim=(B, H, s_kv, D), stride=(s_kv * H * D, D, H * D, 1), data_type=DTYPE, name="v") o, _ = g.sdpa(name="s", q=q, k=k, v=v, attn_scale=0.1, is_inference=True, use_causal_mask=True) _finish_output(o, (B, H, S, D), (S * H * D, D, H * D, 1)) - assert engines.engine_name(512) in _eligible(g) + assert engines.engine_name() in _eligible(g) def _mk_eligible_graph(): @@ -558,7 +558,7 @@ def _mk_eligible_graph(): def test_knob_request_within_domain_keeps_engine_eligible(): g = _mk_eligible_graph() - assert engines.engine_name(512) in _eligible(g, engines.SdpaFwdKnobs(sched_policy=0, tile_m=128, tile_n=128, cga=2)) + assert engines.engine_name() in _eligible(g, engines.SdpaFwdKnobs(sched_policy=0, tile_m=128, tile_n=128, cga=2)) def test_knob_request_outside_domain_rejects_engine(): @@ -575,7 +575,7 @@ def test_knob_request_lpt_sched_is_in_domain(): # The SM100 rows advertise all three scheduler policies (the static/CLC # remap serves them); an explicit LPT request stays eligible. g = _mk_eligible_graph() - assert engines.engine_name(512) in _eligible(g, engines.SdpaFwdKnobs(sched_policy=1)) + assert engines.engine_name() in _eligible(g, engines.SdpaFwdKnobs(sched_policy=1)) def test_knob_request_unsupported_tile_rejects_engine(): @@ -596,7 +596,7 @@ def test_knob_request_wrong_vocabulary_rejects_engine(): def test_knob_request_none_fields_are_no_preference(): g = _mk_eligible_graph() - assert engines.engine_name(512) in _eligible(g, engines.SdpaFwdKnobs()) + assert engines.engine_name() in _eligible(g, engines.SdpaFwdKnobs()) # --------------------------------------------------------------------------- @@ -677,7 +677,7 @@ def build(request_amax_s): amx_s.set_output(True).set_dim((1, 1, 1, 1)).set_stride((1, 1, 1, 1)).set_data_type(cudnn.data_type.FLOAT) return gg - fp8_name = engines.engine_name(128, fp8=True) + fp8_name = engines.engine_name(fp8=True) assert fp8_name in _eligible(build(request_amax_s=False)) assert fp8_name not in _eligible(build(request_amax_s=True)) @@ -689,7 +689,7 @@ def test_probe_accepts_bottom_right_with_swa(): q, k, v, dims, strides = _mk_qkv(g) o, _ = g.sdpa(name="s", q=q, k=k, v=v, attn_scale=0.1, is_inference=True, use_causal_mask_bottom_right=True, sliding_window_length=64) _finish_output(o, dims, strides) - assert engines.engine_name(512) in _eligible(g) + assert engines.engine_name() in _eligible(g) def test_sm120_probe_accepts_bottom_right_with_swa(monkeypatch): @@ -893,7 +893,7 @@ def test_probe_accepts_thd_cu_seq_len(): """THD with the (B+1,) cu_seq_len prefix-sum form (cuDNN 9.24+) is served: the lowering derives per-batch lengths host-side from its inherent tolist round-trip.""" - assert engines.engine_name(512) in _eligible(_mk_thd_cu_graph()) + assert engines.engine_name() in _eligible(_mk_thd_cu_graph()) def test_probe_rejects_thd_cu_plus_seq_len(): diff --git a/test/python/sdpa/frost/test_sdpa_stream_ordering.py b/test/python/sdpa/frost/test_sdpa_stream_ordering.py index b76f0fbed..a381a2757 100644 --- a/test/python/sdpa/frost/test_sdpa_stream_ordering.py +++ b/test/python/sdpa/frost/test_sdpa_stream_ordering.py @@ -23,9 +23,9 @@ import torch from test_utils import torch_fork_set_rng -from frost_test_utils import requires_blackwell, requires_dsl +from frost_test_utils import requires_pre_rubin_blackwell, requires_dsl -pytestmark = requires_blackwell +pytestmark = requires_pre_rubin_blackwell # ~0.5-1 s of spin: long enough that an unordered kernel launch on the # default stream reliably overtakes the side-stream mutation. diff --git a/test/python/sdpa/frost/test_sdpa_stream_respect.py b/test/python/sdpa/frost/test_sdpa_stream_respect.py index 1adaa5bf9..223d41afe 100644 --- a/test/python/sdpa/frost/test_sdpa_stream_respect.py +++ b/test/python/sdpa/frost/test_sdpa_stream_respect.py @@ -18,9 +18,9 @@ import cudnn from cudnn.engines import is_python_engine -from frost_test_utils import requires_blackwell, requires_dsl, _dsl_installed +from frost_test_utils import requires_pre_rubin_blackwell, requires_dsl, _dsl_installed -pytestmark = [pytest.mark.L0, requires_blackwell, requires_dsl] +pytestmark = [pytest.mark.L0, requires_pre_rubin_blackwell, requires_dsl] _B, _H, _S = 2, 8, 256 _HALF, _F32 = cudnn.data_type.HALF, cudnn.data_type.FLOAT diff --git a/test/python/sdpa/frost/test_split_kv_heuristic.py b/test/python/sdpa/frost/test_split_kv_heuristic.py index 84cebf106..5a5304878 100644 --- a/test/python/sdpa/frost/test_split_kv_heuristic.py +++ b/test/python/sdpa/frost/test_split_kv_heuristic.py @@ -200,7 +200,7 @@ def test_split_request_outside_the_domain_makes_the_engine_ineligible(requested) """The default row serves only split_kv=1 ("off"): a split request on a row whose lowering has no split path is honored-or-ineligible, never silently dropped.""" - caps = Capabilities(sm_lo=100, sm_hi=100, phase="prefill", d_qk=frozenset({128}), d_v=frozenset({128})) + caps = Capabilities(sm_lo=100, sm_hi=100, phase="prefill", d_shapes=frozenset({(128, 128)})) assert caps.split_kvs == frozenset({1}) why = mismatch(caps, _facts(), SdpaFwdKnobs(split_kv=requested)) assert why is not None and "split_kv" in why @@ -209,7 +209,7 @@ def test_split_request_outside_the_domain_makes_the_engine_ineligible(requested) def test_no_split_and_explicit_one_leave_the_engine_eligible(): """No preference passes; so does an EXPLICIT split_kv=1 — "do not split" is a real point on the axis, not an ignored request.""" - caps = Capabilities(sm_lo=100, sm_hi=100, phase="prefill", d_qk=frozenset({128}), d_v=frozenset({128})) + caps = Capabilities(sm_lo=100, sm_hi=100, phase="prefill", d_shapes=frozenset({(128, 128)})) for knobs in (SdpaFwdKnobs(split_kv=None), SdpaFwdKnobs(split_kv=1)): why = mismatch(caps, _facts(), knobs) or "" assert "split_kv" not in why @@ -232,8 +232,7 @@ def test_split_declines_when_the_kv_tail_needs_synthesized_padding(): sm_lo=100, sm_hi=100, phase="prefill", - d_qk=frozenset({128}), - d_v=frozenset({128}), + d_shapes=frozenset({(128, 128)}), skv_tail_via_padding=True, split_kvs=frozenset({1, 2, 4}), ) @@ -255,11 +254,8 @@ def test_split_domains_match_the_wired_lowerings(): advertising = {sp.name for sp in ENGINE_SPECS if sp.capabilities.split_kvs != frozenset({1})} assert advertising == { - "sdpa_fwd_prefill_sm100_d128", - "sdpa_fwd_prefill_sm100_d256", - "sdpa_fwd_prefill_sm100_d512", - "sdpa_fwd_prefill_sm100_d192_d128", - "sdpa_fwd_prefill_sm100_d128_mxfp8", - "sdpa_fwd_prefill_sm100_d128_fp8", + "sdpa_fwd_prefill_sm100", + "sdpa_fwd_prefill_sm100_mxfp8", + "sdpa_fwd_prefill_sm100_fp8", "sdpa_fwd_prefill_sm120", }, f"split domains drifted from the wired lowerings: {sorted(advertising)}" From 1ecac50710b9219c319b4408abb5cc1644614135 Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Sat, 22 Aug 2026 14:08:11 -0700 Subject: [PATCH 3/3] test: exercise the d192-flavor fp8 envelope through the direct adapter API MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The pygraph sdpa_fp8 node validator (C++ frontend) still bounds the graph route at d_qk <= 128 (%16) / exact (192, 128) — a pre-FE-OSS shape whitelist — so the (160, 96) region of the d192xd128 flavor's envelope is reachable through the standalone API only. Relaxing that validation to describe-not-judge is a separate C++ question for the maintainers; the kernel capability and the engine row are validated here regardless. Co-Authored-By: Claude Fable 5 --- .../sdpa/frost/test_sdpa_fwd_fp8_sm100.py | 43 ++++++++++++------- 1 file changed, 28 insertions(+), 15 deletions(-) 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 d884647ce..fae5d1e29 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_fp8_sm100.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_fp8_sm100.py @@ -354,21 +354,34 @@ def test_fp8_head_dim_envelope(dims, mask): @torch_fork_set_rng(seed=0) def test_fp8_head_dim_envelope_d192_flavor(): """The d192xd128 flavor serves its envelope too: (160, 96) rides the - (192, 128) kernel with TMA zero-padding on both sides.""" - out, o_ref, a_o, a_o_ref = _run( - 2, - 4, - 4, - 384, - 384, - "e4m3", - torch.bfloat16, - scale=1.0 / math.sqrt(160), - sdpa_kwargs={}, - d_qk=160, - d_v=96, - ) - _check(out, o_ref, torch.bfloat16, "e4m3", a_o, a_o_ref) + (192, 128) kernel with TMA zero-padding on both sides. + + Direct adapter API: the pygraph ``sdpa_fp8`` node validator still bounds + the GRAPH route at d_qk <= 128 (%16) / exact (192, 128) — a pre-FE-OSS + shape whitelist in the C++ frontend — so this region of the envelope is + reachable through the standalone API only until that validation is + relaxed to describe rather than judge.""" + from cudnn.sdpa.fwd.api_dsl import SdpaFwdDslSm100 + + B, H, S, d_qk, d_v = 2, 4, 384, 160, 96 + dev = "cuda" + Q8 = (torch.randn(B, H, S, d_qk, device=dev) * 0.5).to(torch.float8_e4m3fn) + K8 = (torch.randn(B, H, S, d_qk, device=dev) * 0.5).to(torch.float8_e4m3fn) + V8 = (torch.randn(B, H, S, d_v, device=dev) * 0.5).to(torch.float8_e4m3fn) + out = torch.empty(B, H, S, d_v, device=dev, dtype=torch.bfloat16) + lse = torch.empty(B, H, S, device=dev, dtype=torch.float32) + scale = 1.0 / math.sqrt(d_qk) + + api = SdpaFwdDslSm100(sample_q=Q8, sample_k=K8, sample_v=V8, sample_o=out, sample_lse=lse, scale_softmax=scale, pertensor_fp8=True) + assert api.check_support() + api.compile() + api.execute(q_tensor=Q8, k_tensor=K8, v_tensor=V8, o_tensor=out, lse_tensor=lse) + torch.cuda.synchronize() + + o_ref = torch.softmax(Q8.float() @ K8.float().transpose(-1, -2) * scale, dim=-1) @ V8.float() + err = (out.float() - o_ref).abs().max().item() + assert err <= 5e-2 + 0.05 * o_ref.abs().max().item(), f"(160,96) envelope mismatch: max err {err}" + assert not torch.isnan(out.float()).any() @pytest.mark.L0