frost(sdpa): optimize and enhance sm120 sdpa forward fp8 kernel

python/cudnn/sdpa/fwd/api_dsl.py

@@ -33,6 +33,7 @@
     SEQ_KV_TILES as _SM120_KV_TILES,
     SEQ_Q_TILES as _SM120_Q_TILES,
     SUPPORTED_HEAD_TILE_MAX as _SM120_HEAD_TILE_MAX,
+    SUPPORTED_HEAD_TILES_FP8 as _SM120_FP8_HEAD_TILES,
     TemplateParams as Sm120TemplateParams,
     smem_bytes as _sm120_smem_bytes,
 )
@@ -61,11 +62,14 @@ def _require_reciprocal_s_scales(descale_s: float, scale_s: float) -> None:
 
     This row cannot apply Scale_S, and would gain nothing if it could:
 
-    - No headroom. The lazy-rescale skip (RESCALE_THRESHOLD=8) refreshes the
-      running max only when a tile exceeds it by 2^8, so P is bounded by 256,
-      not 1. e4m3 tops out at 448, so any scale_s > 448/256 = 1.75 can saturate
-      a lazily-skipped tile. Measured on B2xH8xS256 e4m3: max|O-ref| is flat
-      from scale_s 1 to 64 and degrades at 448 (swa .0239 -> .0807).
+    - No headroom. The lazy-rescale skip refreshes the running max only when
+      a tile exceeds it by RESCALE_THRESHOLD -- 4.0 for the fp8 dtypes
+      (config_sm100.rescale_threshold; 8.0 is the dataclass default the fp8
+      path overrides) -- so P is bounded by 2^4 = 16, not 1. e4m3 tops out at
+      448, so any scale_s > 448/16 = 28 can saturate a lazily-skipped tile.
+      Measured on B2xH8xS256 e4m3: max|O-ref| is flat from scale_s 1 to 64
+      (the analytical bound is conservative) and degrades at 448
+      (swa .0239 -> .0807).
     - Nothing to gain. e4m3 is floating point, so relative precision does not
       move with scale, and subtracting the row max already places P per ROW —
       strictly better than a per-tensor scale. Hence the flat error above.
@@ -1839,8 +1843,9 @@ def check_support(self) -> bool:
             self._value_error_if(self.has_sink, "SM120 fp8 does not support attention sinks (Amax_S semantics)")
             self._value_error_if(self.seq_q_lens_present and not self.thd, "SM120 fp8 does not support per-batch seq_len_q")
             self._value_error_if(
-                (d_q, d_v) != (128, 128),
-                f"SM120 fp8 requires D_QK=D_V=128 (no zero-padding envelope on the 8-bit fragment path); got ({d_q}, {d_v})",
+                any(d not in _SM120_FP8_HEAD_TILES for d in (d_q, d_v)),
+                f"SM120 fp8 requires D_QK and D_V to be multiples of 32 within 32..256 (k32 contraction and 1-byte "
+                f"TMA swizzle span; no zero-padding envelope on the 8-bit fragment path); got ({d_q}, {d_v})",
             )
 
         self._value_error_if(
@@ -1896,8 +1901,6 @@ def check_support(self) -> bool:
 
         def _smem_bytes(kv_tile: int) -> int:
             # FP8 stages a byte per KV element but still writes O in half.
-            # (The FP8 row requires exact d128, so its padded dims are the
-            # actual ones; the envelope padding is the f16 cell's.)
             return _sm120_smem_bytes(d_qp, d_vp, self.q_tile, kv_tile, self.dtype.itemsize, 2 if self._fp8 else self.dtype.itemsize)
 
         if self.tile_n is None:

python/cudnn/sdpa/fwd/config_sm120.py

@@ -16,6 +16,8 @@
 SUPPORTED_HEAD_TILE_MIN = 16
 SUPPORTED_HEAD_TILE_MAX = 256
 SUPPORTED_HEAD_TILES = tuple(range(SUPPORTED_HEAD_TILE_MIN, SUPPORTED_HEAD_TILE_MAX + 1, HEAD_TILE_GRANULE))
+FP8_HEAD_TILE_GRANULE = 32
+SUPPORTED_HEAD_TILES_FP8 = tuple(range(FP8_HEAD_TILE_GRANULE, SUPPORTED_HEAD_TILE_MAX + 1, FP8_HEAD_TILE_GRANULE))
 
 # SMEM the SM120 parts expose to a kernel. The adapter asks cutlass for the
 # authoritative number at build time; this constant lets the ranking answer

python/cudnn/sdpa/fwd/engines.py

@@ -267,11 +267,11 @@ def mismatch(capabilities: Capabilities, facts: "ga.SdpaGraphFacts", knobs: Opti
         return f"serves D_QK in {sorted(capabilities.d_qk)}/D_V in {sorted(capabilities.d_v)}; graph has D_QK={facts.d_qk}/D_V={facts.d_v}"
     if facts.dtype not in capabilities.dtypes:
         return f"dtype {facts.dtype} not in {sorted(str(d) for d in capabilities.dtypes)}"
-    if (capabilities.is_fp8 or capabilities.is_mxfp8) and facts.dtype_o not in capabilities.out_dtypes:
-        return f"O dtype {facts.dtype_o} not in {sorted(str(d) for d in capabilities.out_dtypes)}"
     if (facts.is_mxfp8, facts.is_fp8) != (capabilities.is_mxfp8, capabilities.is_fp8):
         quant = "block-scale MXFP8 (sdpa_mxfp8)" if capabilities.is_mxfp8 else "per-tensor FP8 (sdpa_fp8)" if capabilities.is_fp8 else "half (sdpa)"
         return f"this engine serves only {quant} graphs"
+    if (capabilities.is_fp8 or capabilities.is_mxfp8) and facts.dtype_o not in capabilities.out_dtypes:
+        return f"O dtype {facts.dtype_o} not in {sorted(str(d) for d in capabilities.out_dtypes)}"
     if not facts.uniform_dtype:
         return "K/V dtypes must match Q" if (facts.is_mxfp8 or facts.is_fp8) else "K/V/O dtypes must match Q"
     if facts.thd:
@@ -718,10 +718,9 @@ def _execute(variant_pack, workspace=None, stream=None):
         # construction, and the packed layout gives each sequence its own
         # extent, so nothing is written past a valid length.
         if (facts.is_mxfp8 or facts.is_fp8) and not facts.thd and seq_q_buf is not None and not seq_q_lens_present:
-            if int(seq_q_buf.min().item()) < int(facts.s_q):
-                raise NotImplementedError(
-                    f"per-tensor FP8/MXFP8: per-batch seq_len_q shorter than S_q={facts.s_q} is not plumbed; got min {int(seq_q_buf.min().item())}"
-                )
+            min_seq_q = int(seq_q_buf.min().item())
+            if min_seq_q < int(facts.s_q):
+                raise NotImplementedError(f"per-tensor FP8/MXFP8: per-batch seq_len_q shorter than S_q={facts.s_q} is not plumbed; got min {min_seq_q}")
         execute_kwargs = dict(
             q_tensor=q_buf,
             k_tensor=k_buf,
@@ -811,20 +810,22 @@ def _sm120_fp8_spec() -> EngineSpec:
     m16n8k32 e4m3; ``descale_q*descale_k`` folds into the softmax scale and
     ``descale_v*scale_o`` into an epilogue scalar, so the kernel adds only the
     Amax_S/Amax_O atomics over the f16 sibling. E4M3 only (no E5M2 tag in the
-    kernel yet), FP16 O only, exact d128 (no zero-padding envelope on the
+    kernel yet), FP16 O only, head dims any multiple of 32 up to 256 with the
+    QK^T and P@V sides independent (exact — no zero-padding envelope on the
     8-bit fragment path), and no sink (Amax_S semantics with a sink column are
     undefined here). THD (ragged) is served with token-major Stats; head-major
     ragged Stats stays f16-only (the fp8 kernel carries no such specialization).
     """
+    from cudnn.sdpa.fwd.config_sm120 import SUPPORTED_HEAD_TILES_FP8
 
     return EngineSpec(
         name="sdpa_fwd_prefill_sm120_fp8",
         capabilities=Capabilities(
             sm_lo=_BLACKWELL_GEFORCE[0],
             sm_hi=_BLACKWELL_GEFORCE[1],
             phase="prefill",
-            d_qk=frozenset({128}),
-            d_v=frozenset({128}),
+            d_qk=frozenset(SUPPORTED_HEAD_TILES_FP8),
+            d_v=frozenset(SUPPORTED_HEAD_TILES_FP8),
             dtypes=frozenset({cudnn.data_type.FP8_E4M3}),
             out_dtypes=frozenset({cudnn.data_type.HALF}),
             is_fp8=True,

python/cudnn/sdpa/fwd/heuristics.py

@@ -78,8 +78,8 @@ def _sm120_tiles(caps: Capabilities, facts) -> Tuple[int, int]:
     tile_m = 64 if fine else 128
     # FP8 stages a byte per KV element but still writes O in half, so the two
     # SMEM terms size differently -- see config_sm120.smem_bytes.
-    qkv_item, o_item = (1, 2) if facts.is_fp8 else (2, 2)
-    fits = [n for n in sorted(caps.tile_ns, reverse=True) if smem_bytes(facts.d_qk, facts.d_v, tile_m, n, qkv_item, o_item) <= SMEM_CAPACITY_BYTES]
+    qkv_itemsize, o_itemsize = (1, 2) if facts.is_fp8 else (2, 2)
+    fits = [n for n in sorted(caps.tile_ns, reverse=True) if smem_bytes(facts.d_qk, facts.d_v, tile_m, n, qkv_itemsize, o_itemsize) <= SMEM_CAPACITY_BYTES]
     return tile_m, (fits[0] if fits else min(caps.tile_ns))
 
 
@@ -124,7 +124,10 @@ def _mode_a(facts, offered: Dict[str, int], mode) -> List[PlanConfig]:
         # the rule's regret is small but not zero, so the runners-up are worth
         # offering to a caller who measures. Configs the kernel cannot fit are
         # not runners-up -- they would sit in the list only to decline at build.
-        domain = [(m, n) for m in caps.tile_ms for n in caps.tile_ns if smem_bytes(facts.d_qk, facts.d_v, m, n) <= SMEM_CAPACITY_BYTES]
+        qkv_itemsize, o_itemsize = (1, 2) if facts.is_fp8 else (2, 2)
+        domain = [
+            (m, n) for m in caps.tile_ms for n in caps.tile_ns if smem_bytes(facts.d_qk, facts.d_v, m, n, qkv_itemsize, o_itemsize) <= SMEM_CAPACITY_BYTES
+        ]
         ordered = sorted(domain or [best], key=lambda mn: (mn != best, mn[1] != best[1], -mn[0]))
         for tile_m, tile_n in ordered:
             knobs = _knobs(caps, tile_m, tile_n)