From 182d3a732a57bdde5262e2140642f96c66ea380b Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Thu, 13 Aug 2026 11:41:35 -0700 Subject: [PATCH] sdpa fp8 sm100: fused LDTM row-max on cc10.3 (per-tensor FP8) cc10.3 has had the fused LDTM.STAT row-max (tcgen05.ld.red.f32.max) since bring-up, but only the MXFP8 kernel read the fused_ldtm_stat flag - the per-tensor FP8 kernel ran the manual tcgen05_ld + software reduction on every part. This wires the same specialization into the per-tensor kernel: the unmasked softmax path loads S_acc and reduces the row max in one op per 64-column chunk; masked iters keep the software path (the fused max reduces before a mask could apply); cc10.0 folds to the manual path unchanged (the flag is per-specialization state in TemplateParams, so each device class gets its own traced variant, as with MXFP8). The branch is byte-for-byte the one the SM107 sibling bakes (board-proven on cc10.7, where the LDTM lever carried most of a -7.5% step at 65k), and the helper + instruction are the same ones the MXFP8 kernel exercises on cc10.3 in production. SM100 (cc10.0) suite: 34 passed, behaviorally identical. Direct cc10.3 hardware validation is queued (B300); the fused path's only new surface here is this kernel's integration, identical in shape to both existing users. Co-Authored-By: Claude Fable 5 --- python/cudnn/sdpa/fwd/api_dsl.py | 6 +++- python/cudnn/sdpa/fwd/config_sm100.py | 2 +- .../fwd/kernels/prefill_d128_fp8_sm100.py | 29 +++++++++++++++---- 3 files changed, 30 insertions(+), 7 deletions(-) diff --git a/python/cudnn/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py index e8bc28633..dea9857a7 100644 --- a/python/cudnn/sdpa/fwd/api_dsl.py +++ b/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 diff --git a/python/cudnn/sdpa/fwd/config_sm100.py b/python/cudnn/sdpa/fwd/config_sm100.py index 1dedb6c5e..8cb5c8063 100644 --- a/python/cudnn/sdpa/fwd/config_sm100.py +++ b/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 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 cf49a9a54..6ddf2576f 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm100.py +++ b/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",