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