From bfb8cba168045e1f8b3ef0e6fec9e45c85464632 Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Thu, 13 Aug 2026 01:46:08 -0700 Subject: [PATCH] sdpa fp8 sm107: port the has_lse specialization; add a static SMEM guard Two follow-ups to #576: 1. #574 (has_lse specialization for the FP8/MXFP8 SM100 flavors) merged between #577 and #576, so the shared adapter now passes has_lse to the fp8 compile() while the SM107 sibling still had the pre-#574 signature - a TypeError on any Rubin fp8 compile at develop tip. This ports #574's fp8-kernel hunks onto the sibling verbatim (LSE None-specialization, compile(has_lse), the specialized epilogue), restoring signature parity with the shared call site. 2. Responding to review on #576: the 9-stage ring with BF16/FP16 O (~242 KiB) exceeds the STANDARD sm_10x 227 KiB per-CTA opt-in and is legal on GR100 only through the sm107 oversized-SMEM launch mode (function attribute 16), which the required internal cutlass-dsl toolchain enables for its sm_107a kernels - board-validated e4m3->bf16 across the full suite. Rather than rejecting output dtypes that demonstrably work on the target stack, a static import-time guard now accounts the geometry against the GR100 hardware budget so a future stage/width bump fails with a clear message instead of at launch. Validated: SM100 box 34 passed; SM107 board 34 passed (full fp8 e2e suite through the sibling, has_lse and no-lse populations both exercised via the suite's generate_stats matrix). Co-Authored-By: Claude Fable 5 --- .../fwd/kernels/prefill_d128_fp8_sm107.py | 67 ++++++++++++++----- 1 file changed, 49 insertions(+), 18 deletions(-) diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm107.py b/python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm107.py index d0374fac7..1a8852c59 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm107.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm107.py @@ -46,7 +46,7 @@ import os import sys from functools import lru_cache -from typing import Callable, Tuple +from typing import Callable, Optional, Tuple from cutlass.base_dsl.typing import Pointer # was the legacy DSL Pointer pre-DKG-bump @@ -76,6 +76,25 @@ CFG = _dc.replace(CFG, TILE_K_HW_BMM1=64, TILE_K_HW_BMM2=64, STAGES_KV=9) Cfg = type(CFG) + +# Static SMEM accounting for the post-override geometry. The 9-stage ring +# with BF16/FP16 O (~241 KiB) exceeds the STANDARD sm_10x 227 KiB per-CTA +# opt-in, and is legal on GR100 only through the sm107 oversized-SMEM +# launch mode (function attribute 16), which the required internal +# cutlass-dsl toolchain enables for its sm_107a kernels — board-validated +# e4m3->bf16 across the full suite. The guard below is against the GR100 +# HARDWARE budget so a future geometry bump fails at import with a clear +# message instead of at launch. +_GR100_SMEM_BUDGET = 320 * 1024 # usable per-CTA carveout (327 KiB capacity minus reserves) +_SMEM_BYTES = ( + CFG.TILES_Q * CFG.TILE_M * CFG.TILE_K * CFG.BPE # Q slabs + + CFG.STAGES_KV * (CFG.TILE_N * CFG.TILE_K // CFG.CTA_MMA) * CFG.BPE # K ring + + CFG.STAGES_KV * (CFG.TILE_O * CFG.TILE_N // CFG.CTA_MMA) * CFG.BPE # V ring + + CFG.TILES_Q * CFG.TILE_M * CFG.TILE_O * CFG.BPE_O # O slabs + + 2048 # barriers + scheduler + tmem-ptr slack (upper bound) +) +if _SMEM_BYTES > _GR100_SMEM_BUDGET: + raise ValueError(f"prefill_d128_fp8_sm107: SMEM {_SMEM_BYTES} B exceeds the GR100 budget ({_GR100_SMEM_BUDGET} B) — shrink STAGES_KV") TMA_QK_ITERS = _TMA.QK_ITERS TMA_VO_ITERS = _TMA.VO_ITERS TMA_QK_GRANU_ELEMS = _TMA.QK_GRANU_ELEMS @@ -246,7 +265,7 @@ def _kernel( tma_k_desc: cutlass.GridConstant[tmap.TensorMap], tma_v_desc: cutlass.GridConstant[tmap.TensorMap], tma_o_desc: cutlass.GridConstant[tmap.TensorMap], - lse_tensor: cute.Tensor, + lse_tensor: Optional[cute.Tensor], sinks_tensor: cute.Tensor, seq_kv_lens_tensor: cute.Tensor, o_desc_words: cute.Tensor, @@ -1548,7 +1567,7 @@ def _correction_warp_group( tidx, bars, sched, - lse_tensor: cute.Tensor, + lse_tensor: Optional[cute.Tensor], sinks_tensor: cute.Tensor, seq_kv_lens_tensor, n_q_supers, @@ -1736,16 +1755,20 @@ def _correction_warp_group( _s_q_b = cutlass.Int32(_cu[n_batch + batch_idx + cutlass.Int32(1)]) - _cu_q_b _row_valid = q_row_global < _s_q_b if _row_valid: - lse_arr = cutlass.make_array_view(lse_tensor) - lse_row = lse_arr[cutlass.Int32(0), head_idx, :] - lse_row[_cu_q_b + q_row_global] = lse_val + # has_lse=False: the Stats store is compiled out; the amax_s + # atomicMax is independent of it and always live. + if cutlass.const_expr(lse_tensor is not None): + lse_arr = cutlass.make_array_view(lse_tensor) + lse_row = lse_arr[cutlass.Int32(0), head_idx, :] + lse_row[_cu_q_b + q_row_global] = lse_val nvvm.atomicrmw(nvvm.AtomicOp.MAX, _amax_s_ptr, _beta_bits) else: _row_valid = q_row_global < seqlen_q if _row_valid: - lse_arr = cutlass.make_array_view(lse_tensor) - lse_row = lse_arr[batch_idx, head_idx, :] - lse_row[q_row_global] = lse_val + if cutlass.const_expr(lse_tensor is not None): + lse_arr = cutlass.make_array_view(lse_tensor) + lse_row = lse_arr[batch_idx, head_idx, :] + lse_row[q_row_global] = lse_val nvvm.atomicrmw(nvvm.AtomicOp.MAX, _amax_s_ptr, _beta_bits) # amax_o = max over valid rows of |o_scaled| (the fp32 pre-cast output). Divided @@ -1885,7 +1908,7 @@ def _host( k_tensor: cute.Tensor, v_tensor: cute.Tensor, o_tensor: cute.Tensor, - lse_tensor: cute.Tensor, + lse_tensor: Optional[cute.Tensor], sinks_tensor: cute.Tensor, seq_kv_lens_tensor: cute.Tensor, o_desc_words: cute.Tensor, @@ -1994,10 +2017,13 @@ def _tma_swz(byte_w: int): @lru_cache(maxsize=None) -def compile(b: int = 1, qh: int = 1, kh: int = 1, sq: int = 256, skv: int = 128) -> Callable: # noqa: A001 +def compile(b: int = 1, qh: int = 1, kh: int = 1, sq: int = 256, skv: int = 128, has_lse: bool = True) -> Callable: # noqa: A001 """Compile with ALL dims concrete — pins TMA strides at compile time. - THD/varlen: q/k/v/o/lse PACKED with batch dim 1; ``b`` = logical batch.""" + THD/varlen: q/k/v/o/lse PACKED with batch dim 1; ``b`` = logical batch. + ``has_lse=False`` compiles the LSE store out (the kernel specializes on a + ``None`` LSE argument) — callers without a Stats output pass no LSE buffer + at all; the amax_s/amax_o atomicMax writes are independent and unchanged.""" _fake_batch = 1 if CFG.THD_VARLEN else b fake_q = cute.runtime.make_fake_compact_tensor( STORAGE_DTYPE, @@ -2023,12 +2049,17 @@ def compile(b: int = 1, qh: int = 1, kh: int = 1, sq: int = 256, skv: int = 128) stride_order=(3, 2, 1, 0), assumed_align=16, ) - fake_lse = cute.runtime.make_fake_compact_tensor( - cutlass.Float32, - (_fake_batch, qh, sq), - stride_order=(2, 1, 0), - assumed_align=16, - ) + if not has_lse: + # No Stats output: the LSE argument is None-specialized and the store + # is compiled out entirely — no dummy buffer exists at any level. + fake_lse = None + else: + fake_lse = cute.runtime.make_fake_compact_tensor( + cutlass.Float32, + (_fake_batch, qh, sq), + stride_order=(2, 1, 0), + assumed_align=16, + ) # Always part of the ABI; unread when CFG.HAS_SINK == 0 (compile-time fold). fake_sinks = cute.runtime.make_fake_compact_tensor( cutlass.Float32,