sdpa fp8 sm100: fused LDTM row-max on cc10.3 (per-tensor FP8)

python/cudnn/sdpa/fwd/api_dsl.py

@@ -943,7 +943,11 @@ def compile(self) -> None:
         # don't read this flag). Auto-set from the device capability so an SM103 run
         # picks the fused path with no user action.
         mxfp8 = self._fp8 and not self._pertensor
-        fused_ldtm_stat = mxfp8 and (self._device_cc == (10, 3))
+        # cc10.3 has the fused LDTM.STAT row-max: MXFP8 has used it since its
+        # bring-up; per-tensor FP8 now takes the same path (its kernel reads
+        # the flag identically — the SM107 sibling bakes it). cc10.0 lacks the
+        # instruction and keeps the manual reduction.
+        fused_ldtm_stat = self._fp8 and self._device_cc == (10, 3)
         sched_policy = self.sched_policy
         if mxfp8 and sched_policy == SCHED_NATURAL and self.window_right is not None:
             sched_policy = SCHED_LPT

python/cudnn/sdpa/fwd/config_sm100.py

@@ -89,7 +89,7 @@ class TemplateParams:
     thd_varlen: bool = False
     # cc10.3+ fuses the S_acc row-max into the LDTM (tcgen05.ld.red.f32.max); cc10.0
     # lacks it and uses the manual load + software reduction. Auto-set from the device
-    # capability at compile time (MXFP8 only; the f16/fp8 kernels do not read it).
+    # capability at compile time (MXFP8 and per-tensor FP8; the f16 kernels do not read it).
     fused_ldtm_stat: bool = False
 
 

python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm100.py

@@ -82,11 +82,13 @@
     SCHED_LPT_L2,
 )
 from cudnn.frost.tile_dsl.pointwise import (
-    # SM100: no LDTM.STAT — the MASK_NONE fast path uses manual tcgen05_ld +
-    # row_max_reduction (see _softmax_kv_body); tmem_load_max_reduction_tile
-    # is not imported.
+    # cc10.0: manual tcgen05_ld + row_max_reduction on the MASK_NONE path.
+    # cc10.3 (FUSED_LDTM_STAT) fuses load + row-max into one
+    # tcgen05.ld.red.f32.max via tmem_load_max_reduction_x64, exactly like
+    # the MXFP8 kernel (and the SM107 sibling, which bakes it).
     row_reduction_pair,
     row_max_reduction,
+    tmem_load_max_reduction_x64,
     vec_scale_pair,
     fp32_to_fp8_pack,
 )
@@ -145,6 +147,11 @@
     make_sdpa_helpers,
 )
 
+# cc10.3 fuses the S_acc row-max into the LDTM; cc10.0 lacks the instruction
+# and keeps the manual load + software reduction. Set by the adapter from the
+# device capability (see api_dsl.compile()).
+FUSED_LDTM_STAT = int(PARAMS.fused_ldtm_stat)
+
 CGA_SIZE = CFG.CGA_M * CFG.CGA_N
 
 CTA_GROUP_KIND = nvvm.CTAGroup.CTA_2 if CFG.CTA_MMA == 2 else nvvm.CTAGroup.CTA_1
@@ -1220,9 +1227,21 @@ def _softmax_kv_body(
         current_max_unscaled = chunks_max[0]
         for m in chunks_max[1:]:
             current_max_unscaled = cute.math.max(current_max_unscaled, m)
+    elif cutlass.const_expr(FUSED_LDTM_STAT != 0):
+        # cc10.3: one tcgen05.ld.red.f32.max per 64-col chunk does the S_acc
+        # load AND the row-max in one op (data regs + max at index CHUNK).
+        # Unmasked iters only — the fused max reduces before a mask could be
+        # applied (masked iters take the software path above).
+        res_chunks = [tmem_load_max_reduction_x64(s_addr_base + cutlass.Int32(c * CHUNK)) for c in range(N_CHUNKS)]
+        raw_chunks = [cutlass.Vector.from_elements(tuple(r[:CHUNK]), cutlass.Int32).bitcast(cutlass.Float32) for r in res_chunks]
+        chunks_max = [cutlass.Vector.from_elements((r[CHUNK],), cutlass.Int32).bitcast(cutlass.Float32)[0] for r in res_chunks]
+        reg_S_vec = vec_concat(raw_chunks)
+        current_max_unscaled = chunks_max[0]
+        for m in chunks_max[1:]:
+            current_max_unscaled = cute.math.max(current_max_unscaled, m)
     else:
-        # SM100: manual row-max (no LDTM.STAT / tmem_load_max_reduction_tile) —
-        # the masked path's pattern sans mask.  S_acc is FP32 regardless of dtype.
+        # cc10.0: manual row-max (no LDTM.STAT) — the masked path's pattern
+        # sans mask.  S_acc is FP32 regardless of dtype.
         raw_chunks = [
             nvvm.tcgen05_ld(
                 "32x32b",