sdpa fp8: dense head-dim envelope + one engine per arch x dtype family

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"),

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<dqk>` and append `x<dv>` 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

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,24 @@ 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 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 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,
-            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, 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
         # actual tensor extents, and cuTensorMapEncodeTiled requires every
@@ -894,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),
@@ -906,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
@@ -1081,21 +1094,28 @@ 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,
                 sq=self.s_q_max,
                 skv=self.s_k_max,
                 has_lse=(self.lse_desc is not None) or self.split_kv > 1,
             )
+            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:
             # ENVELOPE: hand the f16/bf16 kernel the ACTUAL head dims so its
             # TMA descriptors carry the real extents (loads past them