From 0a8eb07d79180fef5ded062dc00d3bcebcb329c4 Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Fri, 7 Aug 2026 13:31:23 -0700 Subject: [PATCH 1/8] frost(sdpa): support ragged stats for SM100 frost sdpa forward engine MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Port of the SM120 ragged-Stats support (#508) to the SM100 f16/bf16 flavors (d128, d192/d128, d256, d512). The SM100 THD kernels always write a packed LSE; it previously landed in workspace scratch in the kernels' native head-major (1, QH, T) packing. The epilogue store is now layout-aware and writes the caller's ragged Stats buffer directly in the graph's declared layout — no extra kernels on the execute path: declared layout | strides (dims (B, H, S, 1)) | kernel store ------------------|-----------------------------|--------------------------- token-major [t,h] | stride_h == 1, stride_s == H | lse[cu_q[b] + row, head] head-major [h,t] | stride_s == 1, stride_h >= T | lse[head, cu_q[b] + row] - config_sm100: new TemplateParams.thd_lse_token_major -> CFG.THD_LSE_TOKEN_MAJOR (THD-only, validated); kernels gain a lse_stride compile() shape for the head-major padded head stride (part of the per-shape cache key); the THD fake LSE drops to element alignment (user buffers only guarantee 4B). - SdpaFwdDslSm100: THD + sample_lse accepted with the same declared-layout validation as SM120; the packed-LSE workspace chunk is carved only for stats-less graphs; t_kv == 0 short-cut fills the Stats valid region with -inf (or the sink logit alone) in either layout; strict lse_tensor presence contract in both directions for THD. - engines: the SM100 f16 spec advertises thd_stats. Testing (cc 10.0): pytest test_sdpa_fwd_dsl_sm100.py -m "L0 or L1" (thd/graph_api slice: 102 passed; new stats/contract tests: 15 passed) pytest test_sdpa_graph_analyzer.py (69 passed) pytest test_sdpa_fwd_{fp8,mxfp8}_sm100.py (46 passed) Related to #381. Co-Authored-By: Claude Fable 5 --- python/cudnn/sdpa/fwd/api_dsl.py | 126 +++++++--- python/cudnn/sdpa/fwd/config_sm100.py | 14 ++ python/cudnn/sdpa/fwd/engines.py | 3 +- .../fwd/kernels/prefill_d128_f16_sm100.py | 43 +++- .../kernels/prefill_d192_d128_f16_sm100.py | 35 ++- .../fwd/kernels/prefill_d256_f16_sm100.py | 29 ++- .../fwd/kernels/prefill_d512_f16_sm100.py | 29 ++- .../sdpa/frost/test_sdpa_fwd_dsl_sm100.py | 228 ++++++++++++++++-- .../sdpa/frost/test_sdpa_graph_analyzer.py | 37 +++ 9 files changed, 475 insertions(+), 69 deletions(-) diff --git a/python/cudnn/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py index 3eb8e44e8..d05b34766 100644 --- a/python/cudnn/sdpa/fwd/api_dsl.py +++ b/python/cudnn/sdpa/fwd/api_dsl.py @@ -459,6 +459,8 @@ def _initialize_implementation(self) -> None: self.flavor: Optional[tuple[int, int]] = None self.mask_flags = 0 self.swa_window_runtime = 0 + self.thd_stats_token_major = False + self.thd_stats_head_stride = 0 self._k_mod = None def check_support(self) -> bool: @@ -545,13 +547,21 @@ def check_support(self) -> bool: ) self.dtype_o = self.dtype if self.lse_desc is not None: - # THD stats are not plumbed (the kernel's packed (1, H, T) LSE does - # not match cuDNN's ragged Stats contract) — reject the request - # instead of silently never writing the user's LSE. - self._not_implemented_error_if(self.thd, "THD stats/LSE output is not plumbed yet; construct without sample_lse") self._check_dtype(self.lse_desc, torch.float32, name="LSE") self._check_tensor_shape(self.lse_desc, (b, h_qo, s_qo), name="LSE") - self._value_error_if(not self.lse_desc.is_contiguous(), "LSE must be contiguous on SM100 DSL") + if self.thd: + stride_h, stride_s = tuple(self.lse_desc.stride[1:]) + token_major = (stride_h, stride_s) == (1, h_qo) + head_major = not token_major and stride_s == 1 and stride_h >= 1 + self._value_error_if( + not token_major and not head_major, + f"THD LSE must be packed token-major (stride_h == 1, stride_s == H) " + f"or head-major (stride_s == 1, stride_h == head_stride); got stride {self.lse_desc.stride}", + ) + self.thd_stats_token_major = token_major + self.thd_stats_head_stride = int(stride_h) if head_major else 0 + else: + self._value_error_if(not self.lse_desc.is_contiguous(), "LSE must be contiguous on SM100 DSL") self._value_error_if(not torch.cuda.is_available(), "CUDA must be available for SM100 DSL SDPA") device = self.q_desc.device @@ -713,6 +723,7 @@ def compile(self) -> None: seq_q_lens_present=self.seq_q_lens_present, sched_policy=sched_policy, thd_varlen=self.thd, + thd_lse_token_major=self.thd and self.thd_stats_token_major, fused_ldtm_stat=fused_ldtm_stat, ) self._k_mod = _load_sm100_kernel_module(self.flavor, params, fp8=self._fp8, pertensor=self._pertensor) @@ -760,15 +771,19 @@ def scratch_workspace_bytes(self) -> int: b, qh = self.batch_size, self.h_q if self.thd: # [slq32 | slk32 | meta(seq_kv, cu_q, cu_k) | o_desc | packed LSE | sinks dummy] - # The packed LSE is sized for the worst case t_q = B * S_q_max - # (per-execute t_q is a runtime value; every carve stays within - # this bound). o_desc: 16 int64 per sequence + 16 spare, the - # per-sequence O TMA descriptors the builder kernel fills. + # The packed-LSE scratch exists only when NO Stats output is + # declared (the kernel always writes an LSE; with a Stats output + # it writes the caller's ragged Stats buffer directly — token-major + # (T, H) or head-major (H, head_stride) — and no dummy is carved). + # It is sized for the worst case t_q = B * S_q_max (per-execute + # t_q is a runtime value; every carve stays within this bound). + # o_desc: 16 int64 per sequence + 16 spare, the per-sequence O TMA + # descriptors the builder kernel fills. return ( 2 * ws_align(b * 4) + ws_align((3 * b + 2) * 4) + ws_align((b * 16 + 16) * 8) - + ws_align(qh * b * self.s_q_max * 4) + + (0 if self.lse_desc is not None else ws_align(qh * b * self.s_q_max * 4)) + (0 if self.has_sink else ws_align(qh * 4)) ) if self._fp8: @@ -834,11 +849,15 @@ def execute( "lse_tensor is required by this compiled specialization", ) if self.thd: - # The kernel's packed-LSE scratch is api-level workspace; a - # user-facing THD LSE output is not plumbed, so reject rather than - # silently never writing the caller's buffer (check_support already - # rejects thd + sample_lse). - self._not_implemented_error_if(lse_tensor is not None, "THD stats/LSE output is not plumbed yet") + # A THD lse_tensor is bound in its DECLARED packed layout + # (token-major / head-major, recorded at check_support); without a + # sample_lse there is no layout to bind it under, so reject rather + # than guess. Dense keeps accepting an extra lse_tensor (the SM100 + # kernels always write an LSE — see the else-branch dummy below). + self._value_error_if( + self.lse_desc is None and lse_tensor is not None, + "this specialization was compiled without an LSE output; construct the API with sample_lse", + ) elif lse_tensor is not None: lse_tensor = self._checked_lse_view(lse_tensor) else: @@ -900,7 +919,17 @@ def execute( if self.thd: self._execute_thd( - q_tensor, k_tensor, v_tensor, o_tensor, scale_softmax_log2, sinks, seq_kv_lens, seq_q_lens, workspace=workspace, current_stream=current_stream + q_tensor, + k_tensor, + v_tensor, + o_tensor, + scale_softmax_log2, + sinks, + seq_kv_lens, + seq_q_lens, + lse_tensor=lse_tensor, + workspace=workspace, + current_stream=current_stream, ) return @@ -950,16 +979,21 @@ def execute( O_view.copy_(O_scratch) self._logger.debug("execute completed") - def _execute_thd(self, q_buf, k_buf, v_buf, o_buf, scale_softmax_log2, sinks, seq_len_kv, seq_q_lens, workspace=None, current_stream=None): + def _execute_thd(self, q_buf, k_buf, v_buf, o_buf, scale_softmax_log2, sinks, seq_len_kv, seq_q_lens, lse_tensor=None, workspace=None, current_stream=None): """THD / varlen execute: reconstruct the kernel's packed [1, T, H, D] views and metadata buffer from the cuDNN ragged buffers, then launch. With a ``workspace`` the metadata buffers (int32 length copies, the [seq_kv | cu_q | cu_k] buffer, the per-sequence O TMA descriptors, the - packed LSE, the sinks dummy) are carved from it — zero per-execute - allocations; without one they are torch-allocated (standalone use). - The host round-trip for the runtime totals (t_q / t_kv / unit count) - is inherent to the lowering — the packed extents are data-dependent — - and costs one D2H sync per length tensor, no device allocation.""" + packed-LSE scratch (stats-less graphs only), the sinks dummy) are + carved from it — zero per-execute allocations; without one they are + torch-allocated (standalone use). ``lse_tensor``, when given, is the + caller's ragged Stats buffer, written by the kernel directly in its + declared layout: token-major packed ``(T, H)`` in the first ``T*H`` + elements, or head-major ``(H, head_stride)`` with tokens contiguous + within each head row. The host round-trip for the runtime totals + (t_q / t_kv / unit count) is inherent to the lowering — the packed + extents are data-dependent — and costs one D2H sync per length + tensor, no device allocation.""" import cutlass dev = q_buf.device @@ -999,13 +1033,35 @@ def _execute_thd(self, q_buf, k_buf, v_buf, o_buf, scale_softmax_log2, sinks, se # * t_q == 0 — no query token exists anywhere, so the packed O/LSE have # zero rows: nothing to compute or write. # * t_kv == 0 — every query row is fully masked; cuDNN semantics for a - # dead row are O := 0 (LSE := -inf, but THD stats are not supported). + # dead row are O := 0, LSE := -inf (or the sink alone — its column + # keeps the softmax denominator alive). if t_q == 0: self._logger.debug("execute (THD): t_q == 0, nothing to do") return + lse = None + lse_valid = None # the valid-region view (first t_q tokens) the t_kv == 0 fill writes + if lse_tensor is not None: + if self.thd_stats_token_major: + lse = lse_tensor.as_strided((1, t_q, qh), (t_q * qh, qh, 1), lse_tensor.storage_offset()) + lse_valid = lse_tensor.as_strided((t_q, qh), (qh, 1), lse_tensor.storage_offset()) + else: + head_stride = self.thd_stats_head_stride + self._value_error_if( + head_stride < t_q, + f"head-major THD LSE head_stride ({head_stride}) must cover the packed Q token total ({t_q})", + ) + lse = lse_tensor.as_strided((1, qh, head_stride), (qh * head_stride, head_stride, 1), lse_tensor.storage_offset()) + lse_valid = lse_tensor.as_strided((qh, t_q), (head_stride, 1), lse_tensor.storage_offset()) if t_kv == 0: self._logger.debug("execute (THD): t_kv == 0, zeroing packed O") o_buf.as_strided((t_q * qh * d_v,), (1,), o_buf.storage_offset()).zero_() + if lse_valid is not None: + if sinks is not None: + sinks_v = self._checked_sinks_1d(sinks) + sinks_v = sinks_v.reshape(1, qh).expand(t_q, qh) if self.thd_stats_token_major else sinks_v.reshape(qh, 1).expand(qh, t_q) + lse_valid.copy_(sinks_v) + else: + lse_valid.fill_(float("-inf")) return # Per-sequence O TMA descriptors, filled by the kernel's builder pass. @@ -1023,9 +1079,13 @@ def _packed(buf, t, h, d): K = _packed(k_buf, t_kv, kh, d_qk) V = _packed(v_buf, t_kv, kh, d_v) O = _packed(o_buf, t_q, qh, d_v) - # Packed dummy LSE (THD stats are not plumbed): carved at the runtime - # t_q, always within the compile-time bound qh * B * S_q_max. - if carver is not None: + # LSE binding: the caller's ragged Stats buffer in its declared layout + # when a Stats output exists; otherwise a packed head-major scratch + # dummy carved at the runtime t_q (always within the compile-time + # bound qh * B * S_q_max) — the kernel always writes an LSE. + if lse is not None: + LSE = lse + elif carver is not None: LSE = carver.take(qh * t_q, torch.float32).reshape(1, qh, t_q) LSE.zero_() else: @@ -1038,7 +1098,19 @@ def _packed(buf, t, h, d): else: sinks_t = torch.zeros(qh, dtype=torch.float32, device=dev) - fn = self._k_mod.compile(b=b, qh=qh, kh=kh, sq=t_q, skv=t_kv, d_qk=d_qk, d_v=d_v) + fn = self._k_mod.compile( + b=b, + qh=qh, + kh=kh, + sq=t_q, + skv=t_kv, + d_qk=d_qk, + d_v=d_v, + # Head-major ragged Stats carry the caller-declared head-row + # stride (a shape, part of the compile cache key); token-major + # and the stats-less dummy are compact (0 -> sq). + lse_stride=(self.thd_stats_head_stride if (lse is not None and not self.thd_stats_token_major) else 0), + ) fn(Q, K, V, O, LSE, sinks_t, meta, o_desc, (b, qh, kh, t_q, t_kv, 0), cutlass.Float32(scale_softmax_log2), cutlass.Int32(units), stream=current_stream) self._logger.debug("execute (THD) completed") diff --git a/python/cudnn/sdpa/fwd/config_sm100.py b/python/cudnn/sdpa/fwd/config_sm100.py index 6814b4d27..f847ce540 100644 --- a/python/cudnn/sdpa/fwd/config_sm100.py +++ b/python/cudnn/sdpa/fwd/config_sm100.py @@ -77,6 +77,11 @@ class TemplateParams: seq_q_lens_present: bool = False sched_policy: int = SCHED_NATURAL thd_varlen: bool = False + # THD ragged-Stats layout: the packed LSE is written directly in the + # graph's declared layout. False (default) = head-major ``(1, H, T)`` — + # the kernels' native packing, also serving the no-stats scratch dummy; + # True = token-major ``(T, H)`` (cuDNN's TH1 ragged Stats contract). + thd_lse_token_major: 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). @@ -104,6 +109,8 @@ def _validate_params(flavor: str, k: TemplateParams) -> None: raise ValueError(f"{flavor}: THD/varlen implies per-sequence padded masking (MASK_PADDED)") if not k.seq_kv_lens_present: raise ValueError(f"{flavor}: THD/varlen requires SEQ_KV_LENS_PRESENT") + if k.thd_lse_token_major and not k.thd_varlen: + raise ValueError(f"{flavor}: THD_LSE_TOKEN_MAJOR only applies under THD_VARLEN (dense LSE is (B, H, S))") if k.seq_q_lens_present: if k.thd_varlen: raise ValueError(f"{flavor}: SEQ_Q_LENS_PRESENT is dense-only (THD carries per-sequence Q lengths via cu_seqlens)") @@ -264,6 +271,7 @@ class CfgD256: SEQ_Q_LENS_PRESENT: int = 0 THD_VARLEN: int = 0 + THD_LSE_TOKEN_MAJOR: int = 0 def _validate_cfg_d256(cfg: CfgD256) -> None: @@ -309,6 +317,7 @@ def make_cfg_d256(params: TemplateParams) -> Tuple[CfgD256, TmaIters]: SEQ_KV_LENS_PRESENT=1 if (params.thd_varlen or params.seq_kv_lens_present) else 0, SEQ_Q_LENS_PRESENT=int(params.seq_q_lens_present), THD_VARLEN=int(params.thd_varlen), + THD_LSE_TOKEN_MAJOR=int(params.thd_lse_token_major), ) _validate_cfg_d256(cfg) return cfg, _tma_iters(cfg) @@ -401,6 +410,7 @@ class CfgD512: SEQ_Q_LENS_PRESENT: int = 0 THD_VARLEN: int = 0 + THD_LSE_TOKEN_MAJOR: int = 0 def _validate_cfg_d512(cfg: CfgD512) -> None: @@ -450,6 +460,7 @@ def make_cfg_d512(params: TemplateParams) -> Tuple[CfgD512, TmaIters]: SEQ_KV_LENS_PRESENT=1 if (params.thd_varlen or params.seq_kv_lens_present) else 0, SEQ_Q_LENS_PRESENT=int(params.seq_q_lens_present), THD_VARLEN=int(params.thd_varlen), + THD_LSE_TOKEN_MAJOR=int(params.thd_lse_token_major), ) _validate_cfg_d512(cfg) return cfg, _tma_iters(cfg) @@ -545,6 +556,7 @@ class CfgD128: SEQ_Q_LENS_PRESENT: int = 0 THD_VARLEN: int = 0 + THD_LSE_TOKEN_MAJOR: int = 0 def _validate_cfg_d128(cfg: CfgD128) -> None: @@ -608,6 +620,7 @@ def make_cfg_d128(params: TemplateParams) -> Tuple[CfgD128, TmaIters]: SEQ_KV_LENS_PRESENT=1 if (params.thd_varlen or params.seq_kv_lens_present) else 0, SEQ_Q_LENS_PRESENT=int(params.seq_q_lens_present), THD_VARLEN=int(params.thd_varlen), + THD_LSE_TOKEN_MAJOR=int(params.thd_lse_token_major), ) _validate_cfg_d128(cfg) return cfg, _tma_iters(cfg) @@ -678,6 +691,7 @@ def make_cfg_d192(params: TemplateParams) -> Tuple[CfgD192, TmaIters]: SEQ_KV_LENS_PRESENT=1 if (params.thd_varlen or params.seq_kv_lens_present) else 0, SEQ_Q_LENS_PRESENT=int(params.seq_q_lens_present), THD_VARLEN=int(params.thd_varlen), + THD_LSE_TOKEN_MAJOR=int(params.thd_lse_token_major), ) _validate_cfg_d192(cfg) return cfg, _tma_iters(cfg) diff --git a/python/cudnn/sdpa/fwd/engines.py b/python/cudnn/sdpa/fwd/engines.py index 3cd0dc6e9..37d94bc40 100644 --- a/python/cudnn/sdpa/fwd/engines.py +++ b/python/cudnn/sdpa/fwd/engines.py @@ -157,7 +157,7 @@ class Capabilities: # compute it from the GLOBAL S_q — the THD variant of the # bottom_right_padded_seq_q gap above. thd_bottom_right: bool = False - thd_stats: bool = False # packed LSE output plumbing is a follow-up + thd_stats: bool = False # ragged Stats output (packed token-major / head-major LSE) # Dense padded + stats needs the per-batch seq_len_q LSE trim (padded # q-rows write LSE=-inf / O=0, cuDNN >= 9.14). Plumbed for the half # kernels via SEQ_Q_LENS_PRESENT; the FP8/MXFP8 kernels lack the epilogue @@ -368,6 +368,7 @@ def _sm100_spec(d: int, d_v: Optional[int] = None) -> EngineSpec: sink=True, stats=True, thd=True, + thd_stats=True, padded_stats=True, # The f16/bf16 lowering serves any dense B/H/S stride permutation # (padded strides included) with the head dim innermost; the diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py index 16b59d1a1..8615f8319 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py @@ -33,7 +33,9 @@ THD / varlen (``CFG.THD_VARLEN=1``): packed ``[1,T,H,D]`` Q/K/V + ``cu_seqlens`` coord offset (applied to BOTH Q slabs under TILES_Q=2), per-batch O -TMA-descriptor array (shared ``thd_sm100.py``), packed ``[1,QH,T]`` LSE — via +TMA-descriptor array (shared ``thd_sm100.py``), packed ragged-Stats LSE +(head-major ``[1,QH,head_stride]`` or token-major ``[1,T,QH]`` per +``CFG.THD_LSE_TOKEN_MAJOR``) — via the shared ``_common_sm100`` / ``thd_sm100`` mechanism (same as the SM100 qwen / dsv4 kernels). The dense ``[B,S,H,D]`` path is byte-identical (folds out at ``THD_VARLEN=0``). @@ -1812,15 +1814,22 @@ def _correction_warp_group( lse_val = cutlass.Float32(arith.select(row_trim.ir_value(), neg_inf_trim.ir_value(), lse_val.ir_value())) inv_sum = cutlass.Float32(arith.select(row_trim.ir_value(), cutlass.Float32(0.0).ir_value(), inv_sum.ir_value())) if cutlass.const_expr(CFG.THD_VARLEN): - # THD: q_row_global is sequence-local; LSE is packed [1,QH,T] → - # index [0, head, cu_q[b] + local], bound by per-sequence Q len S_q_b. + # THD: q_row_global is sequence-local; the packed ragged-Stats + # LSE is written in the caller's declared layout — head-major + # [1, QH, head_stride] (index [0, head, cu_q[b] + local]) or + # token-major [1, T, QH] (index [0, cu_q[b] + local, head]) — + # bound by per-sequence Q len S_q_b. _cu = cutlass.make_array_view(seq_kv_lens_tensor) _cu_q_b = cutlass.Int32(_cu[n_batch + batch_idx]) _s_q_b = cutlass.Int32(_cu[n_batch + batch_idx + cutlass.Int32(1)]) - _cu_q_b if q_row_global < _s_q_b: 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 + if cutlass.const_expr(CFG.THD_LSE_TOKEN_MAJOR): + lse_row = lse_arr[cutlass.Int32(0), _cu_q_b + q_row_global, :] + lse_row[head_idx] = lse_val + else: + lse_row = lse_arr[cutlass.Int32(0), head_idx, :] + lse_row[_cu_q_b + q_row_global] = lse_val else: if q_row_global < seqlen_q: lse_arr = cutlass.make_array_view(lse_tensor) @@ -2014,11 +2023,16 @@ 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, d_qk: int = CFG.TILE_K, d_v: int = CFG.TILE_O) -> Callable: # noqa: A001 +def compile( + b: int = 1, qh: int = 1, kh: int = 1, sq: int = 256, skv: int = 128, d_qk: int = CFG.TILE_K, d_v: int = CFG.TILE_O, lse_stride: int = 0 +) -> Callable: # noqa: A001 """Compile a kernel with ALL dims concrete to pin TMA descriptor strides at compile time. THD/varlen: q/k/v/o/lse are PACKED with batch dim 1 ([1,T,H,D]); ``b`` is the LOGICAL batch (sequence count) driving n_batch / metadata + O-desc sizes. + ``lse_stride`` (THD head-major LSE only) is the caller-declared head-row + stride of the packed [1, QH, head_stride] LSE (0 → compact, i.e. ``sq``); + it is a shape, so it is part of this cache key. ENVELOPE: ``d_qk`` / ``d_v`` are the ACTUAL head dims (defaults = the flavor's full TILE_K / TILE_O). The Q/K/V/O TMA descriptors are built from @@ -2057,11 +2071,24 @@ 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, ) + if CFG.THD_VARLEN: + # Packed ragged-Stats LSE in the caller's declared layout (align 4: + # the store is scalar f32 and the caller's Stats buffer only + # guarantees element alignment). Token-major [1, T, QH]; head-major + # [1, QH, head_stride] with head_stride >= T (compact when 0). + _lse_hs = lse_stride if lse_stride else sq + if _lse_hs < sq: + raise ValueError(f"THD head-major LSE head_stride ({_lse_hs}) must cover the packed Q token total ({sq})") + _fake_lse_shape = (1, sq, qh) if CFG.THD_LSE_TOKEN_MAJOR else (1, qh, _lse_hs) + else: + if lse_stride: + raise ValueError("lse_stride is THD-only (dense LSE is compact (B, H, Sq))") + _fake_lse_shape = (b, qh, sq) fake_lse = cute.runtime.make_fake_compact_tensor( cutlass.Float32, - (_fake_batch, qh, sq), + _fake_lse_shape, stride_order=(2, 1, 0), - assumed_align=16, + assumed_align=4 if CFG.THD_VARLEN else 16, ) # Sinks tensor always part of the ABI; read only when CFG.HAS_SINK == 1 (compile-time fold). fake_sinks = cute.runtime.make_fake_compact_tensor( diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py index 1e6534521..f8a549e82 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py @@ -1896,15 +1896,21 @@ def _correction_warp_group( lse_val = cutlass.Float32(arith.select(row_trim.ir_value(), neg_inf_trim.ir_value(), lse_val.ir_value())) inv_sum = cutlass.Float32(arith.select(row_trim.ir_value(), cutlass.Float32(0.0).ir_value(), inv_sum.ir_value())) if cutlass.const_expr(CFG.THD_VARLEN): - # THD: q_row_global is sequence-local; LSE is packed [1,QH,T] → - # index [0, head, cu_q[b] + local], bound by per-sequence Q len S_q_b. + # THD: q_row_global is sequence-local; the packed ragged-Stats + # LSE is written in the caller's declared layout — head-major + # [1, QH, head_stride] or token-major [1, T, QH] — bound by + # per-sequence Q len S_q_b. _cu = cutlass.make_array_view(seq_kv_lens_tensor) _cu_q_b = cutlass.Int32(_cu[n_batch + batch_idx]) _s_q_b = cutlass.Int32(_cu[n_batch + batch_idx + cutlass.Int32(1)]) - _cu_q_b if q_row_global < _s_q_b: 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 + if cutlass.const_expr(CFG.THD_LSE_TOKEN_MAJOR): + lse_row = lse_arr[cutlass.Int32(0), _cu_q_b + q_row_global, :] + lse_row[head_idx] = lse_val + else: + lse_row = lse_arr[cutlass.Int32(0), head_idx, :] + lse_row[_cu_q_b + q_row_global] = lse_val else: if q_row_global < seqlen_q: lse_arr = cutlass.make_array_view(lse_tensor) @@ -2106,7 +2112,9 @@ 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, d_qk: int = CFG.TILE_K, d_v: int = CFG.TILE_O) -> Callable: # noqa: A001 +def compile( + b: int = 1, qh: int = 1, kh: int = 1, sq: int = 256, skv: int = 128, d_qk: int = CFG.TILE_K, d_v: int = CFG.TILE_O, lse_stride: int = 0 +) -> Callable: # noqa: A001 """Compile a kernel with ALL dims concrete to pin TMA descriptor strides at compile time. THD/varlen: q/k/v/o/lse are PACKED with batch dim 1 ([1,T,H,D]); ``b`` is the @@ -2149,11 +2157,24 @@ 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, ) + if CFG.THD_VARLEN: + # Packed ragged-Stats LSE in the caller's declared layout (align 4: + # the store is scalar f32 and the caller's Stats buffer only + # guarantees element alignment). Token-major [1, T, QH]; head-major + # [1, QH, head_stride] with head_stride >= T (compact when 0). + _lse_hs = lse_stride if lse_stride else sq + if _lse_hs < sq: + raise ValueError(f"THD head-major LSE head_stride ({_lse_hs}) must cover the packed Q token total ({sq})") + _fake_lse_shape = (1, sq, qh) if CFG.THD_LSE_TOKEN_MAJOR else (1, qh, _lse_hs) + else: + if lse_stride: + raise ValueError("lse_stride is THD-only (dense LSE is compact (B, H, Sq))") + _fake_lse_shape = (b, qh, sq) fake_lse = cute.runtime.make_fake_compact_tensor( cutlass.Float32, - (_fake_batch, qh, sq), + _fake_lse_shape, stride_order=(2, 1, 0), - assumed_align=16, + assumed_align=4 if CFG.THD_VARLEN else 16, ) # Sinks tensor always part of the ABI; read only when CFG.HAS_SINK == 1 (compile-time fold). fake_sinks = cute.runtime.make_fake_compact_tensor( diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py index b2d5a967b..7bf1f556a 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py @@ -1491,8 +1491,12 @@ def _correction_warp_group( _s_q_b = cutlass.Int32(_cu[n_batch + batch_idx + cutlass.Int32(1)]) - _cu_q_b if q_row_global < _s_q_b: 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 + if cutlass.const_expr(CFG.THD_LSE_TOKEN_MAJOR): + lse_row = lse_arr[cutlass.Int32(0), _cu_q_b + q_row_global, :] + lse_row[head_idx] = lse_val + else: + lse_row = lse_arr[cutlass.Int32(0), head_idx, :] + lse_row[_cu_q_b + q_row_global] = lse_val else: if q_row_global < seqlen_q: lse_arr = cutlass.make_array_view(lse_tensor) @@ -1671,7 +1675,9 @@ 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, d_qk: int = CFG.TILE_K, d_v: int = CFG.TILE_O) -> Callable: # noqa: A001 +def compile( + b: int = 1, qh: int = 1, kh: int = 1, sq: int = 256, skv: int = 128, d_qk: int = CFG.TILE_K, d_v: int = CFG.TILE_O, lse_stride: int = 0 +) -> Callable: # noqa: A001 """ENVELOPE: ``d_qk`` / ``d_v`` are the ACTUAL head dims (defaults = full TILE_K / TILE_O). TMA descriptors carry these extents while the tile box stays the compile-time TILE geometry: loads past d_qk / d_v zero-fill @@ -1706,11 +1712,24 @@ 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, ) + if CFG.THD_VARLEN: + # Packed ragged-Stats LSE in the caller's declared layout (align 4: + # the store is scalar f32 and the caller's Stats buffer only + # guarantees element alignment). Token-major [1, T, QH]; head-major + # [1, QH, head_stride] with head_stride >= T (compact when 0). + _lse_hs = lse_stride if lse_stride else sq + if _lse_hs < sq: + raise ValueError(f"THD head-major LSE head_stride ({_lse_hs}) must cover the packed Q token total ({sq})") + _fake_lse_shape = (1, sq, qh) if CFG.THD_LSE_TOKEN_MAJOR else (1, qh, _lse_hs) + else: + if lse_stride: + raise ValueError("lse_stride is THD-only (dense LSE is compact (B, H, Sq))") + _fake_lse_shape = (b, qh, sq) fake_lse = cute.runtime.make_fake_compact_tensor( cutlass.Float32, - (_fake_batch, qh, sq), + _fake_lse_shape, stride_order=(2, 1, 0), - assumed_align=16, + assumed_align=4 if CFG.THD_VARLEN else 16, ) fake_sinks = cute.runtime.make_fake_compact_tensor( cutlass.Float32, diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py index f29e20000..4d7154c97 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py @@ -1151,8 +1151,12 @@ def _compute_warp_group( _s_q_b = cutlass.Int32(_cu[n_batch + batch_idx + cutlass.Int32(1)]) - _cu_q_b if q_row_global < _s_q_b: 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 + if cutlass.const_expr(CFG.THD_LSE_TOKEN_MAJOR): + lse_row = lse_arr[cutlass.Int32(0), _cu_q_b + q_row_global, :] + lse_row[head_idx] = lse + else: + lse_row = lse_arr[cutlass.Int32(0), head_idx, :] + lse_row[_cu_q_b + q_row_global] = lse else: if q_row_global < seqlen_q: lse_arr = cutlass.make_array_view(lse_tensor) @@ -1874,7 +1878,9 @@ 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, d_qk: int = CFG.TILE_K, d_v: int = CFG.TILE_O) -> Callable: +def compile( + b: int = 1, qh: int = 1, kh: int = 1, sq: int = 256, skv: int = 128, d_qk: int = CFG.TILE_K, d_v: int = CFG.TILE_O, lse_stride: int = 0 +) -> Callable: """ENVELOPE: ``d_qk`` / ``d_v`` are the ACTUAL head dims (defaults = full TILE_K / TILE_O). TMA descriptors carry these extents while the tile box stays the compile-time TILE geometry: loads past d_qk / d_v zero-fill @@ -1911,11 +1917,24 @@ 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, ) + if CFG.THD_VARLEN: + # Packed ragged-Stats LSE in the caller's declared layout (align 4: + # the store is scalar f32 and the caller's Stats buffer only + # guarantees element alignment). Token-major [1, T, QH]; head-major + # [1, QH, head_stride] with head_stride >= T (compact when 0). + _lse_hs = lse_stride if lse_stride else sq + if _lse_hs < sq: + raise ValueError(f"THD head-major LSE head_stride ({_lse_hs}) must cover the packed Q token total ({sq})") + _fake_lse_shape = (1, sq, qh) if CFG.THD_LSE_TOKEN_MAJOR else (1, qh, _lse_hs) + else: + if lse_stride: + raise ValueError("lse_stride is THD-only (dense LSE is compact (B, H, Sq))") + _fake_lse_shape = (b, qh, sq) fake_lse = cute.runtime.make_fake_compact_tensor( cutlass.Float32, - (_fake_batch, qh, sq), + _fake_lse_shape, stride_order=(2, 1, 0), - assumed_align=16, + assumed_align=4 if CFG.THD_VARLEN else 16, ) fake_sinks = cute.runtime.make_fake_compact_tensor( cutlass.Float32, diff --git a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py index 48751e8bc..51a6cf3af 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py @@ -113,9 +113,12 @@ def test_sdpa_fwd_dsl_sm100_graph_api(dtype, is_causal, d): _DTYPE_IDS = ["fp16", "bf16"] -def _ref_sdpa_full(q, k, v, *, scale, is_causal=False, bottom_right=False, swa_window=None, seq_kv_lens=None, sinks=None): +def _ref_sdpa_full(q, k, v, *, scale, is_causal=False, bottom_right=False, swa_window=None, seq_kv_lens=None, sinks=None, return_stats=False): """fp32 reference matching the SM100 DSL kernel's mask + sink semantics. - q/k/v are BHSD; GQA (h_q > h_kv) is handled by expanding K/V.""" + q/k/v are BHSD; GQA (h_q > h_kv) is handled by expanding K/V. With + ``return_stats`` also returns the (B, H_q, S_q) LSE — logsumexp over the + masked scores (the sink joins as one extra column; fully-masked rows are + -inf without one).""" b, h_q, s_q, _ = q.shape _, h_kv, s_kv, _ = v.shape dev = q.device @@ -139,11 +142,16 @@ def _ref_sdpa_full(q, k, v, *, scale, is_causal=False, bottom_right=False, swa_w if sinks is not None: sink_col = sinks.view(1, h_q, 1, 1).float().expand(b, h_q, s_q, 1).to(dev) - probs = torch.softmax(torch.cat([scores, sink_col], dim=-1), dim=-1) + full_scores = torch.cat([scores, sink_col], dim=-1) + probs = torch.softmax(full_scores, dim=-1) o = torch.matmul(probs[..., :s_kv], v_ref) else: - o = torch.matmul(torch.softmax(scores, dim=-1), v_ref) - return o.to(q.dtype) + full_scores = scores + o = torch.matmul(torch.softmax(scores, dim=-1).nan_to_num(0.0), v_ref) + if not return_stats: + return o.to(q.dtype) + lse = torch.logsumexp(full_scores, dim=-1) # fully-masked rows -> -inf (sink-less) + return o.to(q.dtype), lse def _require_dsl(): @@ -343,11 +351,32 @@ def test_dsl_sm100_execute_sink_lse_contract(): o_ref = _ref_sdpa_full(q, k, v, scale=scale, is_causal=True) torch.testing.assert_close(o, o_ref, atol=5e-2, rtol=3e-2) - # THD LSE output is not plumbed (packed (1, H, T) layout != cuDNN's ragged - # Stats contract): requesting it is rejected up front instead of being - # silently ignored. - with pytest.raises(NotImplementedError, match="THD stats/LSE"): - SdpaFwdDslSm100(sample_q=q, sample_k=k, sample_v=v, sample_o=o, sample_lse=lse, thd=True).check_support() + # THD LSE must be declared packed: token-major [t, h] or head-major + # [h, t]. A dense-contiguous declaration (stride (H*S, S, 1)) is valid + # head-major (head_stride S); a padded sequence stride matches NEITHER + # layout and is rejected up front instead of being silently mis-addressed. + lse_padded = torch.empty(h * s * 2, dtype=torch.float32, device="cuda").as_strided((b, h, s), (h * s * 2, s * 2, 2)) + with pytest.raises(ValueError, match="token-major"): + SdpaFwdDslSm100(sample_q=q, sample_k=k, sample_v=v, sample_o=o, sample_lse=lse_padded, thd=True).check_support() + api = SdpaFwdDslSm100(sample_q=q, sample_k=k, sample_v=v, sample_o=o, sample_lse=lse, thd=True) + assert api.check_support() and not api.thd_stats_token_major and api.thd_stats_head_stride == s + lse_tm = torch.empty(s * h, dtype=torch.float32, device="cuda").as_strided((b, h, s), (s * h, 1, h)) + api = SdpaFwdDslSm100(sample_q=q, sample_k=k, sample_v=v, sample_o=o, sample_lse=lse_tm, thd=True) + assert api.check_support() and api.thd_stats_token_major + + # THD execute keeps a strict presence contract in BOTH directions (dense + # keeps accepting an extra lse_tensor — the kernels always write an LSE): + # the raises fire before any packing or launch. + api = SdpaFwdDslSm100(sample_q=q, sample_k=k, sample_v=v, sample_o=o, sample_lse=lse_tm, thd=True) + assert api.check_support() + api.compile() + with pytest.raises(ValueError, match="lse_tensor is required"): + api.execute(q_tensor=q, k_tensor=k, v_tensor=v, o_tensor=o, seq_q_lens=seq_kv, seq_kv_lens=seq_kv) + api = SdpaFwdDslSm100(sample_q=q, sample_k=k, sample_v=v, sample_o=o, thd=True) + assert api.check_support() + api.compile() + with pytest.raises(ValueError, match="without an LSE output"): + api.execute(q_tensor=q, k_tensor=k, v_tensor=v, o_tensor=o, seq_q_lens=seq_kv, seq_kv_lens=seq_kv, lse_tensor=lse_tm) @pytest.mark.L0 @@ -573,14 +602,41 @@ def _mask_ref_kwargs(mask): }[mask] -def _run_dsl_thd_graph(q_pk, k_pk, v_pk, cu_q, cu_k, seq_lens_q, seq_lens_kv, *, scale, dtype, H_q, H_kv, d, sink=None, mask="causal"): - """Build + execute a packed THD/varlen graph; returns the flat packed O storage buffer.""" +def _run_dsl_thd_graph( + q_pk, + k_pk, + v_pk, + cu_q, + cu_k, + seq_lens_q, + seq_lens_kv, + *, + scale, + dtype, + H_q, + H_kv, + d, + sink=None, + mask="causal", + check_stats=False, + stats_layout="token_major", +): + """Build + execute a packed THD/varlen graph; returns the flat packed O + storage buffer — plus, with ``check_stats``, the flat Stats storage and + the padded token capacity of its head-major head stride. + + ``stats_layout`` selects the ragged Stats declaration: ``token_major`` + (``[t, h]``, sequence stride ``h_q``) or ``head_major`` (``[h, t]``, + sequence stride 1 with a padded token-capacity head stride — + FlashAttention's ``softmax_lse`` layout).""" import cudnn dev = "cuda" B = len(seq_lens_q) T_q, T_kv = cu_q[-1], cu_k[-1] - S_max_q, S_max_kv = max(seq_lens_q), max(seq_lens_kv) + # Clamp the declared KV extent: an all-zero seq_len_kv batch still needs a + # rank-legal (>0) graph dim; the padding mask carries the real lengths. + S_max_q, S_max_kv = max(seq_lens_q), max(max(seq_lens_kv), 1) def _dense_buf(packed, s_max, t, H): stride = (s_max * H * d, d, H * d, 1) @@ -613,16 +669,37 @@ def _dense_buf(packed, s_max, t, H): tq.set_ragged_offset(qro) tk.set_ragged_offset(kro) tv.set_ragged_offset(vro) - kw = dict(name="sdpa", q=tq, k=tk, v=tv, generate_stats=False, attn_scale=scale, use_padding_mask=True, seq_len_q=sq, seq_len_kv=skv) + kw = dict(name="sdpa", q=tq, k=tk, v=tv, generate_stats=check_stats, attn_scale=scale, use_padding_mask=True, seq_len_q=sq, seq_len_kv=skv) kw.update(_mask_graph_kwargs(mask)) vp = {tq: q_gpu, tk: k_gpu, tv: v_gpu, sq: slq, skv: slk, qro: ro_q, kro: ro_k, vro: ro_k, oro: ro_q} if sink is not None: st = g.tensor_like(sink) kw["sink_token"] = st vp[st] = sink - o, _ = g.sdpa(**kw) + o, stats = g.sdpa(**kw) o.set_output(True).set_dim([B, H_q, S_max_q, d]).set_stride(list(stride_q)) o.set_ragged_offset(oro) + stats_stor = None + t_cap = max(64, -(-T_q // 64) * 64) + if check_stats: + assert stats is not None + stats.set_output(True) + stats.set_data_type(cudnn.data_type.FLOAT) + if stats_layout == "head_major": + # [h, t]: tokens contiguous within a head, heads strided by the + # padded token capacity; offsets = cu_q * stride_s = cu_q. + stats_stor = torch.empty(H_q * t_cap, dtype=torch.float32, device=dev) + stats.set_dim((B, H_q, S_max_q, 1)).set_stride((H_q * t_cap, t_cap, 1, 1)) + stats_ro_t = (ro_q.flatten() // (H_q * d)).view(B + 1, 1, 1, 1).contiguous() + else: + # [t, h]: heads contiguous within a token; offsets = cu_q * h_q. + stats_stor = torch.empty(B * S_max_q * H_q, dtype=torch.float32, device=dev) + stats.set_dim((B, H_q, S_max_q, 1)).set_stride((S_max_q * H_q, 1, H_q, 1)) + stats_ro_t = (ro_q.flatten() // d).view(B + 1, 1, 1, 1).contiguous() + stats_ro = g.tensor_like(stats_ro_t, name="stats_ro") + stats.set_ragged_offset(stats_ro) + vp[stats_ro] = stats_ro_t + vp[stats] = stats_stor g.validate() g.build_operation_graph() @@ -633,7 +710,7 @@ def _dense_buf(packed, s_max, t, H): vp[o] = o_gpu g.execute(vp, torch.empty(max(g.get_workspace_size(), 1), device=dev, dtype=torch.uint8)) torch.cuda.synchronize() - return o_stor + return (o_stor, stats_stor, t_cap) if check_stats else o_stor def _combo_dense(d, dtype, H_q, H_kv, scale, sink_t, mask): @@ -693,6 +770,125 @@ def _cu(sl): torch.testing.assert_close(o_out, o_ref, atol=5e-2, rtol=3e-2) +def _run_thd_stats_case(*, seq_lens_q, seq_lens_kv, d=128, dtype=torch.float16, H_q=8, H_kv=8, mask="causal", with_sink=False, stats_layout="token_major"): + """Run a THD (ragged) graph with generate_stats and check O and the ragged + Stats against per-sequence references, in the declared Stats layout.""" + _require_dsl() + + dev = "cuda" + scale = 1.0 / math.sqrt(d) + B = len(seq_lens_q) + + def _cu(sl): + c = [0] + for s in sl: + c.append(c[-1] + s) + return c + + cu_q, cu_k = _cu(seq_lens_q), _cu(seq_lens_kv) + T_q, T_kv = cu_q[-1], cu_k[-1] + q_pk = torch.randn(T_q, H_q, d, device=dev, dtype=dtype) + k_pk = torch.randn(T_kv, H_kv, d, device=dev, dtype=dtype) + v_pk = torch.randn(T_kv, H_kv, d, device=dev, dtype=dtype) + sink_t = torch.randn(1, H_q, 1, 1, dtype=torch.float32, device=dev) if with_sink else None + + o_stor, stats_stor, t_cap = _run_dsl_thd_graph( + q_pk, + k_pk, + v_pk, + cu_q, + cu_k, + seq_lens_q, + seq_lens_kv, + scale=scale, + dtype=dtype, + H_q=H_q, + H_kv=H_kv, + d=d, + sink=sink_t, + mask=mask, + check_stats=True, + stats_layout=stats_layout, + ) + + if stats_layout == "head_major": + packed_stats = stats_stor.view(H_q, t_cap) # (H, head_stride); tokens at [:, cu[i]:cu[i+1]] + else: + packed_stats = stats_stor[: max(T_q, 1) * H_q].view(max(T_q, 1), H_q) # (T, H) + ref_kw = _mask_ref_kwargs(mask) + sinks = sink_t.flatten() if sink_t is not None else None + packed_o = o_stor[: max(T_q, 1) * H_q * d].view(max(T_q, 1), H_q, d) + for i, (nq, _nkv) in enumerate(zip(seq_lens_q, seq_lens_kv)): + if nq == 0: + continue + qb = q_pk[cu_q[i] : cu_q[i + 1]].permute(1, 0, 2).unsqueeze(0) + kb = k_pk[cu_k[i] : cu_k[i + 1]].permute(1, 0, 2).unsqueeze(0) + vb = v_pk[cu_k[i] : cu_k[i + 1]].permute(1, 0, 2).unsqueeze(0) + expected, expected_lse = _ref_sdpa_full(qb, kb, vb, scale=scale, sinks=sinks, return_stats=True, **ref_kw) + got_o = packed_o[cu_q[i] : cu_q[i + 1]].permute(1, 0, 2).unsqueeze(0) + torch.testing.assert_close(got_o, expected, atol=5e-2, rtol=3e-2) + if stats_layout == "head_major": + got_lse = packed_stats[:, cu_q[i] : cu_q[i + 1]].unsqueeze(0) # (H, T_i) -> (1, H, T_i) + else: + got_lse = packed_stats[cu_q[i] : cu_q[i + 1]].t().unsqueeze(0) # (T_i, H) -> (1, H, T_i) + torch.testing.assert_close(got_lse, expected_lse, atol=2e-2, rtol=2e-2) + + +@pytest.mark.L0 +@pytest.mark.parametrize("stats_layout", ["token_major", "head_major"]) +@pytest.mark.parametrize("d", _FLAVORS, ids=_FLAVOR_IDS) +@torch_fork_set_rng(seed=30) +def test_dsl_sm100_thd_stats(d, stats_layout): + """THD + generate_stats: the ragged Stats output is written in the + caller's declared layout — token-major [t, h] or head-major [h, t] — + across every f16 flavor.""" + + _run_thd_stats_case(seq_lens_q=[200, 150], seq_lens_kv=[200, 150], d=d, mask="causal", stats_layout=stats_layout) + + +@pytest.mark.L1 +@torch_fork_set_rng(seed=32) +def test_dsl_sm100_thd_swa_stats(): + """THD + causal left sliding window + ragged Stats: the window trims the + per-sequence LSE denominator.""" + + _run_thd_stats_case(seq_lens_q=[150, 90], seq_lens_kv=[150, 90], mask="swa", stats_layout="token_major") + + +@pytest.mark.L1 +@pytest.mark.parametrize("stats_layout", ["token_major", "head_major"]) +@torch_fork_set_rng(seed=25) +def test_dsl_sm100_thd_gqa_sink_stats(stats_layout): + """THD + GQA + attention sink, with the sink entering the ragged Stats + (both declared layouts).""" + + _run_thd_stats_case(seq_lens_q=[130, 70], seq_lens_kv=[130, 70], H_q=8, H_kv=2, mask="causal", with_sink=True, stats_layout=stats_layout) + + +@pytest.mark.L1 +@torch_fork_set_rng(seed=26) +def test_dsl_sm100_thd_zero_length_sequence_stats(): + """A zero-length sequence contributes no tokens and must not perturb its + packed neighbors (O and ragged Stats). The last sequence has Q tokens but + ZERO keys inside a live launch: its rows must come back O := 0 with + LSE := -inf through the kernel's row_dead guard, not stale memory.""" + + _run_thd_stats_case(seq_lens_q=[128, 0, 64], seq_lens_kv=[100, 0, 0], mask="causal", stats_layout="token_major") + + +@pytest.mark.L1 +@pytest.mark.parametrize("stats_layout", ["token_major", "head_major"]) +@pytest.mark.parametrize("with_sink", [False, True], ids=["no_sink", "sink"]) +@torch_fork_set_rng(seed=31) +def test_dsl_sm100_thd_all_kv_zero_stats(with_sink, stats_layout): + """Every KV length zero: a zero-token K/V view cannot back a TMA + descriptor, so the adapter short-cut fills O := 0 and the ragged Stats + adapter-side — -inf, or the sink value alone (the sink column keeps the + softmax denominator alive) — in either declared layout.""" + + _run_thd_stats_case(seq_lens_q=[64, 32], seq_lens_kv=[0, 0], mask="none", with_sink=with_sink, stats_layout=stats_layout) + + _COMBO_MASKS = { "dense": ["none", "causal", "causal_br", "swa", "padded"], # THD forces padding internally; bottom-right causal is a kernel gap (BR diff --git a/test/python/sdpa/frost/test_sdpa_graph_analyzer.py b/test/python/sdpa/frost/test_sdpa_graph_analyzer.py index bfd266915..a38798043 100644 --- a/test/python/sdpa/frost/test_sdpa_graph_analyzer.py +++ b/test/python/sdpa/frost/test_sdpa_graph_analyzer.py @@ -324,6 +324,43 @@ def test_probe_rejects_thd_bottom_right(): assert not _eligible(g) +def test_probe_accepts_thd_stats(): + """The SM100 epilogue writes cuDNN's ragged Stats directly (token-major + or head-major packed LSE), so THD + generate_stats is eligible.""" + g = _mk_graph() + dims = (B, H, S, D) + strides = (S * H * D, D, H * D, 1) + q = g.tensor(dim=dims, stride=strides, data_type=DTYPE, name="q") + k = g.tensor(dim=dims, stride=strides, data_type=DTYPE, name="k") + v = g.tensor(dim=dims, stride=strides, data_type=DTYPE, name="v") + ro = g.tensor(dim=(B + 1, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT64, name="ro") + q.set_ragged_offset(ro) + k.set_ragged_offset(ro) + v.set_ragged_offset(ro) + seq_q = g.tensor(dim=(B, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT32, name="sq") + seq_kv = g.tensor(dim=(B, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT32, name="skv") + o, stats = g.sdpa( + name="s", + q=q, + k=k, + v=v, + attn_scale=0.1, + generate_stats=True, + use_causal_mask=True, + use_padding_mask=True, + seq_len_q=seq_q, + seq_len_kv=seq_kv, + ) + _finish_output(o, dims, strides) + o.set_ragged_offset(ro) + assert stats is not None + stats.set_output(True).set_dim((B, H, S, 1)).set_stride((S * H, 1, H, 1)) + stats.set_data_type(cudnn.data_type.FLOAT) + stats_ro = g.tensor(dim=(B + 1, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT64, name="stats_ro") + stats.set_ragged_offset(stats_ro) + assert engines.engine_name(512) in _eligible(g) + + def test_probe_rejects_right_band_widening(): g = _mk_graph() q, k, v, dims, strides = _mk_qkv(g) From 93a96c1b140be3e6fc466ac898eb82220fe8c06f Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Fri, 7 Aug 2026 14:07:33 -0700 Subject: [PATCH 2/8] frost(sdpa): has_lse specialization + per-shape stats layout for SM100 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review follow-ups on the SM100 ragged-stats port, bringing the f16 kernels to full SM120 parity on the LSE contract: - has_lse specialization: the f16 kernels (d128, d192/d128, d256, d512) now None-specialize the LSE argument — a stats-less graph compiles the LSE store out. All dummy-LSE scratch disappears: dense inference graphs report get_workspace_size() == 0 (was b*h*s*4), THD keeps only its metadata chunks, and the SM100 f16 spec advertises lse_optional. The lse_tensor execute contract is strict in both directions (an unrequested dense lse_tensor is now rejected instead of silently written). FP8/MXFP8 kernels still write an LSE unconditionally and keep the engine-carved dummy. - No template parameter for the THD stats layout (mirrors SM120's per-compile keying): TemplateParams.thd_lse_token_major and CFG.THD_LSE_TOKEN_MAJOR are gone. The layout is a per-shape compile() specialization encoded in the LSE fake tensor's static layout — token-major binds its natural packed rank-2 (T, H) view, head-major keeps the native rank-3 (1, QH, head_stride) packing, and the epilogue branches on the static rank. - THD metadata built host-side (SM100 AND SM120 adapters): the two device-side torch.cumsum calls each allocated scan-temp storage and launched a kernel per execute. The [seq_kv | cu_q | cu_k] buffer is now built on the host from the (inherent) tolist round-trip and uploaded in one H2D copy; the slq/slk workspace copies go away too. test_workspace_carve_no_per_execute_allocs_and_guards is reworked to a THD graph (dense no longer needs a workspace) and asserts zero per-execute CUDA allocations. Testing (cc 10.0): sm100 suite -m "L0 or L1" -k "thd or graph_api or stats or contract" 111 passed; frontend integration 10 passed; fp8+mxfp8+analyzer 115 passed. The SM120 suite skips locally (no SM120 GPU); its metadata change is mechanically identical and CI-covered. Co-Authored-By: Claude Fable 5 --- python/cudnn/sdpa/fwd/api_dsl.py | 201 ++++++++---------- python/cudnn/sdpa/fwd/config_sm100.py | 14 -- python/cudnn/sdpa/fwd/engines.py | 14 +- .../fwd/kernels/prefill_d128_f16_sm100.py | 106 ++++++--- .../kernels/prefill_d192_d128_f16_sm100.py | 90 +++++--- .../fwd/kernels/prefill_d256_f16_sm100.py | 88 +++++--- .../fwd/kernels/prefill_d512_f16_sm100.py | 84 +++++--- .../frost/test_sdpa_frontend_integration.py | 84 ++++++-- .../sdpa/frost/test_sdpa_fwd_dsl_sm100.py | 24 ++- 9 files changed, 431 insertions(+), 274 deletions(-) diff --git a/python/cudnn/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py index d05b34766..2cbb9a998 100644 --- a/python/cudnn/sdpa/fwd/api_dsl.py +++ b/python/cudnn/sdpa/fwd/api_dsl.py @@ -723,7 +723,6 @@ def compile(self) -> None: seq_q_lens_present=self.seq_q_lens_present, sched_policy=sched_policy, thd_varlen=self.thd, - thd_lse_token_major=self.thd and self.thd_stats_token_major, fused_ldtm_stat=fused_ldtm_stat, ) self._k_mod = _load_sm100_kernel_module(self.flavor, params, fp8=self._fp8, pertensor=self._pertensor) @@ -744,7 +743,8 @@ def compile(self) -> None: # ENVELOPE: hand the f16/bf16 kernel the ACTUAL head dims so its # TMA descriptors carry the real extents (loads past them # zero-fill, O stores past d_v clip); the tile box stays the - # flavor's compile-time D. + # flavor's compile-time D. has_lse=False (no Stats output) + # compiles the LSE store out — no dummy buffer at any level. self._compiled_kernel = self._k_mod.compile( b=self.batch_size, qh=self.h_q, @@ -753,6 +753,7 @@ def compile(self) -> None: skv=self.s_k_max, d_qk=self.head_dim_qk, d_v=self.head_dim_v, + has_lse=self.lse_desc is not None, ) self._logger.debug("compile completed") @@ -770,22 +771,16 @@ def scratch_workspace_bytes(self) -> int: self._ensure_support_checked() b, qh = self.batch_size, self.h_q if self.thd: - # [slq32 | slk32 | meta(seq_kv, cu_q, cu_k) | o_desc | packed LSE | sinks dummy] - # The packed-LSE scratch exists only when NO Stats output is - # declared (the kernel always writes an LSE; with a Stats output - # it writes the caller's ragged Stats buffer directly — token-major - # (T, H) or head-major (H, head_stride) — and no dummy is carved). - # It is sized for the worst case t_q = B * S_q_max (per-execute - # t_q is a runtime value; every carve stays within this bound). - # o_desc: 16 int64 per sequence + 16 spare, the per-sequence O TMA - # descriptors the builder kernel fills. - return ( - 2 * ws_align(b * 4) - + ws_align((3 * b + 2) * 4) - + ws_align((b * 16 + 16) * 8) - + (0 if self.lse_desc is not None else ws_align(qh * b * self.s_q_max * 4)) - + (0 if self.has_sink else ws_align(qh * 4)) - ) + # [meta(seq_kv, cu_q, cu_k) | o_desc | sinks dummy] + # No packed-LSE chunk: with a Stats output the kernel writes the + # caller's ragged Stats buffer directly (token-major (T, H) or + # head-major (H, head_stride)); without one it compiles with + # has_lse=False and no LSE buffer exists at all. No slq/slk + # copies either: the metadata is built host-side from the tolist + # round-trip and uploaded in one H2D copy. o_desc: 16 int64 per + # sequence + 16 spare, the per-sequence O TMA descriptors the + # builder kernel fills. + return ws_align((3 * b + 2) * 4) + ws_align((b * 16 + 16) * 8) + (0 if self.has_sink else ws_align(qh * 4)) if self._fp8: return 0 # dense FP8/MXFP8: no per-execute scratch (dummies are cached one-time) # Dense padded-Q lens bind directly as their own kernel parameter @@ -819,10 +814,9 @@ def execute( ``workspace``: optional caller-provided scratch buffer (uint8, at least ``scratch_workspace_bytes()`` bytes). When given, every - per-execute scratch buffer ( - the THD metadata / O-descriptor / packed-LSE buffers) is carved from - it — zero per-execute allocations. When None (standalone use), those - buffers are torch-allocated as before. + per-execute scratch buffer (the THD metadata / O-descriptor buffers) + is carved from it — zero per-execute allocations. When None + (standalone use), those buffers are torch-allocated as before. """ self._logger.debug("Entering execute") if self._compiled_kernel is None: @@ -848,25 +842,25 @@ def execute( self.lse_desc is not None and lse_tensor is None, "lse_tensor is required by this compiled specialization", ) + # Strict presence contract, both directions: the f16 kernels are + # compiled with has_lse keyed on sample_lse (no Stats output -> the + # LSE store is compiled out and there is no LSE slot to bind), and a + # THD lse_tensor is bound in its DECLARED packed layout (recorded at + # check_support) — so an lse_tensor without a sample_lse cannot be + # honored and is rejected rather than silently dropped. The FP8/MXFP8 + # kernels (dense-only) still write an LSE unconditionally; their + # stats-less write lands in a cached write-only dummy (the FROST + # dispatch never reaches it: engines.lower_dsl_prefill carves the + # dummy from the caller's workspace instead). + self._value_error_if( + self.lse_desc is None and lse_tensor is not None and not self._fp8, + "this specialization was compiled without an LSE output; construct the API with sample_lse", + ) if self.thd: - # A THD lse_tensor is bound in its DECLARED packed layout - # (token-major / head-major, recorded at check_support); without a - # sample_lse there is no layout to bind it under, so reject rather - # than guess. Dense keeps accepting an extra lse_tensor (the SM100 - # kernels always write an LSE — see the else-branch dummy below). - self._value_error_if( - self.lse_desc is None and lse_tensor is not None, - "this specialization was compiled without an LSE output; construct the API with sample_lse", - ) + pass # bound in _execute_thd (declared packed layout) elif lse_tensor is not None: lse_tensor = self._checked_lse_view(lse_tensor) - else: - # The SM100 kernels always write an LSE (no has_lse specialization - # yet — follow-up): with no Stats output requested the write lands - # in a cached write-only dummy, allocated once per device rather - # than per execute. The FROST dispatch path never reaches this: - # engines.lower_dsl_prefill carves the dummy from the caller's - # workspace instead. + elif self._fp8: lse_tensor = self._dummy( "lse", q_tensor.device, @@ -965,7 +959,7 @@ def execute( K, V, O_scratch if o_needs_copy_back else O_view, - lse_tensor.reshape(self.batch_size, self.h_q, self.s_q_max), + lse_tensor.reshape(self.batch_size, self.h_q, self.s_q_max) if lse_tensor is not None else None, sinks_t, seq_kv_t, o_desc_dummy, @@ -984,16 +978,16 @@ def _execute_thd(self, q_buf, k_buf, v_buf, o_buf, scale_softmax_log2, sinks, se With a ``workspace`` the metadata buffers (int32 length copies, the [seq_kv | cu_q | cu_k] buffer, the per-sequence O TMA descriptors, the - packed-LSE scratch (stats-less graphs only), the sinks dummy) are - carved from it — zero per-execute allocations; without one they are - torch-allocated (standalone use). ``lse_tensor``, when given, is the - caller's ragged Stats buffer, written by the kernel directly in its - declared layout: token-major packed ``(T, H)`` in the first ``T*H`` - elements, or head-major ``(H, head_stride)`` with tokens contiguous - within each head row. The host round-trip for the runtime totals - (t_q / t_kv / unit count) is inherent to the lowering — the packed - extents are data-dependent — and costs one D2H sync per length - tensor, no device allocation.""" + sinks dummy) are carved from it — zero per-execute allocations; + without one they are torch-allocated (standalone use). ``lse_tensor``, + when given, is the caller's ragged Stats buffer, written by the + kernel directly in its declared layout: token-major packed ``(T, H)`` + in the first ``T*H`` elements, or head-major ``(H, head_stride)`` + with tokens contiguous within each head row; when ``None`` the kernel + compiles the LSE store out (has_lse=False) and no scratch exists. + The host round-trip for the runtime totals (t_q / t_kv / unit count) + is inherent to the lowering — the packed extents are data-dependent — + and costs one D2H sync per length tensor, no device allocation.""" import cutlass dev = q_buf.device @@ -1001,28 +995,22 @@ def _execute_thd(self, q_buf, k_buf, v_buf, o_buf, scale_softmax_log2, sinks, se slk_v = self._checked_seq_lens(seq_len_kv, "seq_kv_lens") b = slq_v.numel() carver = WorkspaceCarver(workspace, self.scratch_workspace_bytes(), "SdpaFwdDslSm100 (THD)") if workspace is not None else None - if carver is not None: - slq = carver.take(b, torch.int32) - slq.copy_(slq_v) - slk = carver.take(b, torch.int32) - slk.copy_(slk_v) - else: - slq = slq_v - slk = slk_v # Metadata buffer: [ seq_kv_lens(B) | cu_seqlens_q(B+1) | cu_seqlens_k(B+1) ], - # with the cumulative sums built in place (no torch.cat temporaries). + # built HOST-side from the (inherent) tolist round-trip and uploaded in + # ONE H2D copy: a device-side cumsum would allocate its scan-temp + # storage and launch kernels on the execute hot path. meta = carver.take(3 * b + 2, torch.int32) if carver is not None else torch.empty(3 * b + 2, dtype=torch.int32, device=dev) - cu_q = meta[b : 2 * b + 1] - cu_k = meta[2 * b + 1 :] - meta[0:b].copy_(slk) - cu_q[0:1].zero_() - torch.cumsum(slq, 0, dtype=torch.int32, out=cu_q[1:]) # dtype pinned: integer cumsum otherwise promotes to int64 - cu_k[0:1].zero_() - torch.cumsum(slk, 0, dtype=torch.int32, out=cu_k[1:]) - slq_host = slq.tolist() # one D2H sync; t_q/t_kv/units are runtime values - slk_host = slk.tolist() - t_q = int(sum(slq_host)) - t_kv = int(sum(slk_host)) + slq_host = slq_v.tolist() # one D2H sync; t_q/t_kv/units are runtime values + slk_host = slk_v.tolist() + cu_q_host = [0] + for n in slq_host: + cu_q_host.append(cu_q_host[-1] + int(n)) + cu_k_host = [0] + for n in slk_host: + cu_k_host.append(cu_k_host[-1] + int(n)) + meta.copy_(torch.tensor(slk_host + cu_q_host + cu_k_host, dtype=torch.int32)) + t_q = cu_q_host[-1] + t_kv = cu_k_host[-1] qh, kh = self.h_q, self.h_kv d_qk, d_v = self.head_dim_qk, self.head_dim_v @@ -1042,8 +1030,10 @@ def _execute_thd(self, q_buf, k_buf, v_buf, o_buf, scale_softmax_log2, sinks, se lse_valid = None # the valid-region view (first t_q tokens) the t_kv == 0 fill writes if lse_tensor is not None: if self.thd_stats_token_major: - lse = lse_tensor.as_strided((1, t_q, qh), (t_q * qh, qh, 1), lse_tensor.storage_offset()) - lse_valid = lse_tensor.as_strided((t_q, qh), (qh, 1), lse_tensor.storage_offset()) + # Natural packed rank-2 (T, H) view — the kernel's epilogue + # dispatches on this static rank. + lse = lse_tensor.as_strided((t_q, qh), (qh, 1), lse_tensor.storage_offset()) + lse_valid = lse else: head_stride = self.thd_stats_head_stride self._value_error_if( @@ -1080,16 +1070,9 @@ def _packed(buf, t, h, d): V = _packed(v_buf, t_kv, kh, d_v) O = _packed(o_buf, t_q, qh, d_v) # LSE binding: the caller's ragged Stats buffer in its declared layout - # when a Stats output exists; otherwise a packed head-major scratch - # dummy carved at the runtime t_q (always within the compile-time - # bound qh * B * S_q_max) — the kernel always writes an LSE. - if lse is not None: - LSE = lse - elif carver is not None: - LSE = carver.take(qh * t_q, torch.float32).reshape(1, qh, t_q) - LSE.zero_() - else: - LSE = torch.zeros(1, qh, t_q, dtype=torch.float32, device=dev) + # when a Stats output exists; None otherwise — the kernel compiles the + # LSE store out (has_lse=False), so no dummy buffer exists at all. + LSE = lse if sinks is not None: sinks_t = self._checked_sinks_1d(sinks) elif carver is not None: @@ -1106,9 +1089,12 @@ def _packed(buf, t, h, d): skv=t_kv, d_qk=d_qk, d_v=d_v, - # Head-major ragged Stats carry the caller-declared head-row - # stride (a shape, part of the compile cache key); token-major - # and the stats-less dummy are compact (0 -> sq). + # The Stats layout is a per-shape specialization (like d_qk/d_v): + # has_lse=False compiles the store out; token-major binds the + # packed rank-2 (T, H) view; head-major carries the caller-declared + # head-row stride (0 -> compact sq). + has_lse=lse is not None, + lse_token_major=lse is not None and self.thd_stats_token_major, lse_stride=(self.thd_stats_head_stride if (lse is not None and not self.thd_stats_token_major) else 0), ) fn(Q, K, V, O, LSE, sinks_t, meta, o_desc, (b, qh, kh, t_q, t_kv, 0), cutlass.Float32(scale_softmax_log2), cutlass.Int32(units), stream=current_stream) @@ -1882,29 +1868,24 @@ def _execute_thd( slq_v = self._checked_seq_lens(seq_q_lens, "seq_q_lens") slk_v = self._checked_seq_lens(seq_kv_lens, "seq_kv_lens") - if carver is not None: - slq = carver.take(b, torch.int32) - slq.copy_(slq_v) - slk = carver.take(b, torch.int32) - slk.copy_(slk_v) - else: - slq = slq_v - slk = slk_v # [seq_kv(B) | cu_q(B+1) | cu_k(B+1)] — bound as the kernel's # seq_kv_lens tensor; the leading B words alias the per-sequence KV # lengths so the kernel's existing padded-mask read works unchanged. + # Built HOST-side from the (inherent) tolist round-trip and uploaded + # in ONE H2D copy: a device-side cumsum would allocate its scan-temp + # storage and launch kernels on the execute hot path. meta = carver.take(3 * b + 2, torch.int32) if carver is not None else torch.empty(3 * b + 2, dtype=torch.int32, device=dev) - cu_q = meta[b : 2 * b + 1] - cu_k = meta[2 * b + 1 :] - meta[0:b].copy_(slk) - cu_q[0:1].zero_() - torch.cumsum(slq, 0, dtype=torch.int32, out=cu_q[1:]) - cu_k[0:1].zero_() - torch.cumsum(slk, 0, dtype=torch.int32, out=cu_k[1:]) - slq_host = slq.tolist() - slk_host = slk.tolist() - t_q = int(sum(slq_host)) - t_kv = int(sum(slk_host)) + slq_host = slq_v.tolist() + slk_host = slk_v.tolist() + cu_q_host = [0] + for n in slq_host: + cu_q_host.append(cu_q_host[-1] + int(n)) + cu_k_host = [0] + for n in slk_host: + cu_k_host.append(cu_k_host[-1] + int(n)) + meta.copy_(torch.tensor(slk_host + cu_q_host + cu_k_host, dtype=torch.int32)) + t_q = cu_q_host[-1] + t_kv = cu_k_host[-1] max_sq = max(slq_host) if slq_host else 0 if t_q == 0: @@ -1985,17 +1966,17 @@ def _packed(buf, tokens): def scratch_workspace_bytes(self) -> int: if self.thd: - # [slq32 | slk32 | meta(seq_kv, cu_q, cu_k)]. + # [meta(seq_kv, cu_q, cu_k)]. # No packed-LSE chunk: with a Stats output the kernel writes the # caller's ragged Stats buffer directly (token-major (T, H) or # head-major (H, head_stride)); without one it compiles with - # has_lse=False and no LSE buffer exists at all. No sinks-dummy - # chunk either: the kernel - # None-specializes on sinks. No O-descriptor chunk: SM120 stores O - # with plain guarded GMEM stores, so THD needs no per-sequence - # tensor maps. + # has_lse=False and no LSE buffer exists at all. No slq/slk + # copies either: the metadata is built host-side from the tolist + # round-trip. No sinks-dummy chunk: the kernel None-specializes + # on sinks. No O-descriptor chunk: SM120 stores O with plain + # guarded GMEM stores, so THD needs no per-sequence tensor maps. b = self.batch_size - return 2 * ws_align(b * 4) + ws_align((3 * b + 2) * 4) + return ws_align((3 * b + 2) * 4) return 0 diff --git a/python/cudnn/sdpa/fwd/config_sm100.py b/python/cudnn/sdpa/fwd/config_sm100.py index f847ce540..6814b4d27 100644 --- a/python/cudnn/sdpa/fwd/config_sm100.py +++ b/python/cudnn/sdpa/fwd/config_sm100.py @@ -77,11 +77,6 @@ class TemplateParams: seq_q_lens_present: bool = False sched_policy: int = SCHED_NATURAL thd_varlen: bool = False - # THD ragged-Stats layout: the packed LSE is written directly in the - # graph's declared layout. False (default) = head-major ``(1, H, T)`` — - # the kernels' native packing, also serving the no-stats scratch dummy; - # True = token-major ``(T, H)`` (cuDNN's TH1 ragged Stats contract). - thd_lse_token_major: 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). @@ -109,8 +104,6 @@ def _validate_params(flavor: str, k: TemplateParams) -> None: raise ValueError(f"{flavor}: THD/varlen implies per-sequence padded masking (MASK_PADDED)") if not k.seq_kv_lens_present: raise ValueError(f"{flavor}: THD/varlen requires SEQ_KV_LENS_PRESENT") - if k.thd_lse_token_major and not k.thd_varlen: - raise ValueError(f"{flavor}: THD_LSE_TOKEN_MAJOR only applies under THD_VARLEN (dense LSE is (B, H, S))") if k.seq_q_lens_present: if k.thd_varlen: raise ValueError(f"{flavor}: SEQ_Q_LENS_PRESENT is dense-only (THD carries per-sequence Q lengths via cu_seqlens)") @@ -271,7 +264,6 @@ class CfgD256: SEQ_Q_LENS_PRESENT: int = 0 THD_VARLEN: int = 0 - THD_LSE_TOKEN_MAJOR: int = 0 def _validate_cfg_d256(cfg: CfgD256) -> None: @@ -317,7 +309,6 @@ def make_cfg_d256(params: TemplateParams) -> Tuple[CfgD256, TmaIters]: SEQ_KV_LENS_PRESENT=1 if (params.thd_varlen or params.seq_kv_lens_present) else 0, SEQ_Q_LENS_PRESENT=int(params.seq_q_lens_present), THD_VARLEN=int(params.thd_varlen), - THD_LSE_TOKEN_MAJOR=int(params.thd_lse_token_major), ) _validate_cfg_d256(cfg) return cfg, _tma_iters(cfg) @@ -410,7 +401,6 @@ class CfgD512: SEQ_Q_LENS_PRESENT: int = 0 THD_VARLEN: int = 0 - THD_LSE_TOKEN_MAJOR: int = 0 def _validate_cfg_d512(cfg: CfgD512) -> None: @@ -460,7 +450,6 @@ def make_cfg_d512(params: TemplateParams) -> Tuple[CfgD512, TmaIters]: SEQ_KV_LENS_PRESENT=1 if (params.thd_varlen or params.seq_kv_lens_present) else 0, SEQ_Q_LENS_PRESENT=int(params.seq_q_lens_present), THD_VARLEN=int(params.thd_varlen), - THD_LSE_TOKEN_MAJOR=int(params.thd_lse_token_major), ) _validate_cfg_d512(cfg) return cfg, _tma_iters(cfg) @@ -556,7 +545,6 @@ class CfgD128: SEQ_Q_LENS_PRESENT: int = 0 THD_VARLEN: int = 0 - THD_LSE_TOKEN_MAJOR: int = 0 def _validate_cfg_d128(cfg: CfgD128) -> None: @@ -620,7 +608,6 @@ def make_cfg_d128(params: TemplateParams) -> Tuple[CfgD128, TmaIters]: SEQ_KV_LENS_PRESENT=1 if (params.thd_varlen or params.seq_kv_lens_present) else 0, SEQ_Q_LENS_PRESENT=int(params.seq_q_lens_present), THD_VARLEN=int(params.thd_varlen), - THD_LSE_TOKEN_MAJOR=int(params.thd_lse_token_major), ) _validate_cfg_d128(cfg) return cfg, _tma_iters(cfg) @@ -691,7 +678,6 @@ def make_cfg_d192(params: TemplateParams) -> Tuple[CfgD192, TmaIters]: SEQ_KV_LENS_PRESENT=1 if (params.thd_varlen or params.seq_kv_lens_present) else 0, SEQ_Q_LENS_PRESENT=int(params.seq_q_lens_present), THD_VARLEN=int(params.thd_varlen), - THD_LSE_TOKEN_MAJOR=int(params.thd_lse_token_major), ) _validate_cfg_d192(cfg) return cfg, _tma_iters(cfg) diff --git a/python/cudnn/sdpa/fwd/engines.py b/python/cudnn/sdpa/fwd/engines.py index 37d94bc40..831b7572a 100644 --- a/python/cudnn/sdpa/fwd/engines.py +++ b/python/cudnn/sdpa/fwd/engines.py @@ -147,8 +147,8 @@ class Capabilities: stats: bool = False # The adapter accepts lse_tensor=None (its kernel None-specializes the LSE # store), so a stats-less graph needs no dummy-LSE workspace chunk. Rows - # that keep False (the SM100 flavors) always write an LSE and get a carved - # dummy from lower_dsl_prefill when the graph has no Stats output. + # that keep False (the SM100 FP8/MXFP8 flavors) always write an LSE and get + # a carved dummy from lower_dsl_prefill when the graph has no Stats output. lse_optional: bool = False thd: bool = False cu_seq_len: bool = False # cu_seq_len_q / cu_seq_len_kv prefix sums (no row serves these yet) @@ -367,6 +367,7 @@ def _sm100_spec(d: int, d_v: Optional[int] = None) -> EngineSpec: padded=True, sink=True, stats=True, + lse_optional=True, thd=True, thd_stats=True, padded_stats=True, @@ -594,11 +595,10 @@ def lower_dsl_prefill( # build time and recorded on the executor as ``workspace_bytes`` — that # number is what the plan's CompiledPlan.get_workspace_size() reports. # - dummy LSE (dense, stats absent, non-lse_optional adapters): the - # SM100 kernels always write an LSE; without a Stats output it lands - # in b*h_q*s_q fp32 scratch. lse_optional adapters (SM120) compile the - # LSE store out instead and bind no buffer. (THD needs no engine-level - # LSE chunk — the packed THD LSE is part of the api-level scratch - # below.) + # SM100 FP8/MXFP8 kernels always write an LSE; without a Stats output + # it lands in b*h_q*s_q fp32 scratch. lse_optional adapters (the f16 + # flavors, SM120) compile the LSE store out instead and bind no + # buffer. (THD needs no engine-level LSE chunk either way.) # - synthesized seq_len_kv (skv_tail_via_padding rows): b int32. # - api-level scratch (api.scratch_workspace_bytes()): the dense padded # [seq_kv|seq_q] combine and the THD metadata/LSE buffers. diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py index 8615f8319..358732c5e 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py @@ -34,8 +34,8 @@ THD / varlen (``CFG.THD_VARLEN=1``): packed ``[1,T,H,D]`` Q/K/V + ``cu_seqlens`` coord offset (applied to BOTH Q slabs under TILES_Q=2), per-batch O TMA-descriptor array (shared ``thd_sm100.py``), packed ragged-Stats LSE -(head-major ``[1,QH,head_stride]`` or token-major ``[1,T,QH]`` per -``CFG.THD_LSE_TOKEN_MAJOR``) — via +(head-major ``[1,QH,head_stride]`` or token-major ``[T,QH]`` — the epilogue +branches on the LSE tensor's static rank) — via the shared ``_common_sm100`` / ``thd_sm100`` mechanism (same as the SM100 qwen / dsv4 kernels). The dense ``[B,S,H,D]`` path is byte-identical (folds out at ``THD_VARLEN=0``). @@ -250,7 +250,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, @@ -1629,7 +1629,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, seq_q_lens_tensor, @@ -1813,21 +1813,25 @@ def _correction_warp_group( neg_inf_trim = cutlass.Float32(float("-inf")) lse_val = cutlass.Float32(arith.select(row_trim.ir_value(), neg_inf_trim.ir_value(), lse_val.ir_value())) inv_sum = cutlass.Float32(arith.select(row_trim.ir_value(), cutlass.Float32(0.0).ir_value(), inv_sum.ir_value())) - if cutlass.const_expr(CFG.THD_VARLEN): + if cutlass.const_expr(lse_tensor is None): + pass # has_lse=False: the Stats store is compiled out + elif cutlass.const_expr(CFG.THD_VARLEN): # THD: q_row_global is sequence-local; the packed ragged-Stats # LSE is written in the caller's declared layout — head-major - # [1, QH, head_stride] (index [0, head, cu_q[b] + local]) or - # token-major [1, T, QH] (index [0, cu_q[b] + local, head]) — - # bound by per-sequence Q len S_q_b. + # rank-3 [1, QH, head_stride] (index [0, head, cu_q[b] + local]) + # or token-major rank-2 [T, QH] (index [cu_q[b] + local, head]) + # — bound by per-sequence Q len S_q_b. _cu = cutlass.make_array_view(seq_kv_lens_tensor) _cu_q_b = cutlass.Int32(_cu[n_batch + batch_idx]) _s_q_b = cutlass.Int32(_cu[n_batch + batch_idx + cutlass.Int32(1)]) - _cu_q_b if q_row_global < _s_q_b: lse_arr = cutlass.make_array_view(lse_tensor) - if cutlass.const_expr(CFG.THD_LSE_TOKEN_MAJOR): - lse_row = lse_arr[cutlass.Int32(0), _cu_q_b + q_row_global, :] + if cutlass.const_expr(len(lse_tensor.shape) == 2): + # token-major packed (T, H) + lse_row = lse_arr[_cu_q_b + q_row_global, :] lse_row[head_idx] = lse_val else: + # head-major packed (1, QH, head_stride) lse_row = lse_arr[cutlass.Int32(0), head_idx, :] lse_row[_cu_q_b + q_row_global] = lse_val else: @@ -1909,7 +1913,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, @@ -2024,15 +2028,28 @@ 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, d_qk: int = CFG.TILE_K, d_v: int = CFG.TILE_O, lse_stride: int = 0 + b: int = 1, + qh: int = 1, + kh: int = 1, + sq: int = 256, + skv: int = 128, + d_qk: int = CFG.TILE_K, + d_v: int = CFG.TILE_O, + has_lse: bool = True, + lse_token_major: bool = False, + lse_stride: int = 0, ) -> Callable: # noqa: A001 """Compile a kernel with ALL dims concrete to pin TMA descriptor strides at compile time. THD/varlen: q/k/v/o/lse are PACKED with batch dim 1 ([1,T,H,D]); ``b`` is the LOGICAL batch (sequence count) driving n_batch / metadata + O-desc sizes. - ``lse_stride`` (THD head-major LSE only) is the caller-declared head-row - stride of the packed [1, QH, head_stride] LSE (0 → compact, i.e. ``sq``); - it is a shape, so it is part of this cache key. + ``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. THD Stats layouts: ``lse_token_major=True`` = packed rank-2 (T, H); + default head-major = rank-3 [1, QH, head_stride], where ``lse_stride`` is + the caller-declared head-row stride (0 → compact, i.e. ``sq``). All three + are shapes/specializations of the traced code, so they are part of this + cache key. ENVELOPE: ``d_qk`` / ``d_v`` are the ACTUAL head dims (defaults = the flavor's full TILE_K / TILE_O). The Q/K/V/O TMA descriptors are built from @@ -2071,25 +2088,48 @@ def compile( stride_order=(3, 2, 1, 0), assumed_align=16, ) - if CFG.THD_VARLEN: - # Packed ragged-Stats LSE in the caller's declared layout (align 4: - # the store is scalar f32 and the caller's Stats buffer only - # guarantees element alignment). Token-major [1, T, QH]; head-major - # [1, QH, head_stride] with head_stride >= T (compact when 0). - _lse_hs = lse_stride if lse_stride else sq - if _lse_hs < sq: - raise ValueError(f"THD head-major LSE head_stride ({_lse_hs}) must cover the packed Q token total ({sq})") - _fake_lse_shape = (1, sq, qh) if CFG.THD_LSE_TOKEN_MAJOR else (1, qh, _lse_hs) + 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. + if lse_token_major or lse_stride: + raise ValueError("lse_token_major / lse_stride require has_lse=True") + fake_lse = None + elif CFG.THD_VARLEN: + # Packed ragged-Stats LSE in the caller's declared layout (align 4: the + # store is scalar f32 and the caller's Stats buffer only guarantees + # element alignment). Token-major = its natural packed rank-2 (T, H) + # view; head-major = the kernels' native rank-3 (1, QH, head_stride) + # packing with head_stride >= T (compact when 0). The epilogue store + # branches on the STATIC rank, so the layout is fully encoded in this + # fake tensor — no template parameter. + if lse_token_major: + if lse_stride: + raise ValueError("lse_stride is head-major-only (token-major (T, H) is compact)") + fake_lse = cute.runtime.make_fake_compact_tensor( + cutlass.Float32, + (sq, qh), + stride_order=(1, 0), + assumed_align=4, + ) + else: + _lse_hs = lse_stride if lse_stride else sq + if _lse_hs < sq: + raise ValueError(f"THD head-major LSE head_stride ({_lse_hs}) must cover the packed Q token total ({sq})") + fake_lse = cute.runtime.make_fake_compact_tensor( + cutlass.Float32, + (1, qh, _lse_hs), + stride_order=(2, 1, 0), + assumed_align=4, + ) else: - if lse_stride: - raise ValueError("lse_stride is THD-only (dense LSE is compact (B, H, Sq))") - _fake_lse_shape = (b, qh, sq) - fake_lse = cute.runtime.make_fake_compact_tensor( - cutlass.Float32, - _fake_lse_shape, - stride_order=(2, 1, 0), - assumed_align=4 if CFG.THD_VARLEN else 16, - ) + if lse_token_major or lse_stride: + raise ValueError("lse_token_major / lse_stride are THD-only (dense LSE is compact (B, H, Sq))") + fake_lse = cute.runtime.make_fake_compact_tensor( + cutlass.Float32, + (b, qh, sq), + stride_order=(2, 1, 0), + assumed_align=16, + ) # Sinks tensor always part of the ABI; read only when CFG.HAS_SINK == 1 (compile-time fold). fake_sinks = cute.runtime.make_fake_compact_tensor( cutlass.Float32, diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py index f8a549e82..2ed6c4755 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py @@ -295,7 +295,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, @@ -1708,7 +1708,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, seq_q_lens_tensor, @@ -1895,20 +1895,24 @@ def _correction_warp_group( neg_inf_trim = cutlass.Float32(float("-inf")) lse_val = cutlass.Float32(arith.select(row_trim.ir_value(), neg_inf_trim.ir_value(), lse_val.ir_value())) inv_sum = cutlass.Float32(arith.select(row_trim.ir_value(), cutlass.Float32(0.0).ir_value(), inv_sum.ir_value())) - if cutlass.const_expr(CFG.THD_VARLEN): + if cutlass.const_expr(lse_tensor is None): + pass # has_lse=False: the Stats store is compiled out + elif cutlass.const_expr(CFG.THD_VARLEN): # THD: q_row_global is sequence-local; the packed ragged-Stats # LSE is written in the caller's declared layout — head-major - # [1, QH, head_stride] or token-major [1, T, QH] — bound by - # per-sequence Q len S_q_b. + # rank-3 [1, QH, head_stride] or token-major rank-2 [T, QH] — + # bound by per-sequence Q len S_q_b. _cu = cutlass.make_array_view(seq_kv_lens_tensor) _cu_q_b = cutlass.Int32(_cu[n_batch + batch_idx]) _s_q_b = cutlass.Int32(_cu[n_batch + batch_idx + cutlass.Int32(1)]) - _cu_q_b if q_row_global < _s_q_b: lse_arr = cutlass.make_array_view(lse_tensor) - if cutlass.const_expr(CFG.THD_LSE_TOKEN_MAJOR): - lse_row = lse_arr[cutlass.Int32(0), _cu_q_b + q_row_global, :] + if cutlass.const_expr(len(lse_tensor.shape) == 2): + # token-major packed (T, H) + lse_row = lse_arr[_cu_q_b + q_row_global, :] lse_row[head_idx] = lse_val else: + # head-major packed (1, QH, head_stride) lse_row = lse_arr[cutlass.Int32(0), head_idx, :] lse_row[_cu_q_b + q_row_global] = lse_val else: @@ -1998,7 +2002,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, @@ -2113,7 +2117,16 @@ 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, d_qk: int = CFG.TILE_K, d_v: int = CFG.TILE_O, lse_stride: int = 0 + b: int = 1, + qh: int = 1, + kh: int = 1, + sq: int = 256, + skv: int = 128, + d_qk: int = CFG.TILE_K, + d_v: int = CFG.TILE_O, + has_lse: bool = True, + lse_token_major: bool = False, + lse_stride: int = 0, ) -> Callable: # noqa: A001 """Compile a kernel with ALL dims concrete to pin TMA descriptor strides at compile time. @@ -2157,25 +2170,48 @@ def compile( stride_order=(3, 2, 1, 0), assumed_align=16, ) - if CFG.THD_VARLEN: - # Packed ragged-Stats LSE in the caller's declared layout (align 4: - # the store is scalar f32 and the caller's Stats buffer only - # guarantees element alignment). Token-major [1, T, QH]; head-major - # [1, QH, head_stride] with head_stride >= T (compact when 0). - _lse_hs = lse_stride if lse_stride else sq - if _lse_hs < sq: - raise ValueError(f"THD head-major LSE head_stride ({_lse_hs}) must cover the packed Q token total ({sq})") - _fake_lse_shape = (1, sq, qh) if CFG.THD_LSE_TOKEN_MAJOR else (1, qh, _lse_hs) + 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. + if lse_token_major or lse_stride: + raise ValueError("lse_token_major / lse_stride require has_lse=True") + fake_lse = None + elif CFG.THD_VARLEN: + # Packed ragged-Stats LSE in the caller's declared layout (align 4: the + # store is scalar f32 and the caller's Stats buffer only guarantees + # element alignment). Token-major = its natural packed rank-2 (T, H) + # view; head-major = the kernels' native rank-3 (1, QH, head_stride) + # packing with head_stride >= T (compact when 0). The epilogue store + # branches on the STATIC rank, so the layout is fully encoded in this + # fake tensor — no template parameter. + if lse_token_major: + if lse_stride: + raise ValueError("lse_stride is head-major-only (token-major (T, H) is compact)") + fake_lse = cute.runtime.make_fake_compact_tensor( + cutlass.Float32, + (sq, qh), + stride_order=(1, 0), + assumed_align=4, + ) + else: + _lse_hs = lse_stride if lse_stride else sq + if _lse_hs < sq: + raise ValueError(f"THD head-major LSE head_stride ({_lse_hs}) must cover the packed Q token total ({sq})") + fake_lse = cute.runtime.make_fake_compact_tensor( + cutlass.Float32, + (1, qh, _lse_hs), + stride_order=(2, 1, 0), + assumed_align=4, + ) else: - if lse_stride: - raise ValueError("lse_stride is THD-only (dense LSE is compact (B, H, Sq))") - _fake_lse_shape = (b, qh, sq) - fake_lse = cute.runtime.make_fake_compact_tensor( - cutlass.Float32, - _fake_lse_shape, - stride_order=(2, 1, 0), - assumed_align=4 if CFG.THD_VARLEN else 16, - ) + if lse_token_major or lse_stride: + raise ValueError("lse_token_major / lse_stride are THD-only (dense LSE is compact (B, H, Sq))") + fake_lse = cute.runtime.make_fake_compact_tensor( + cutlass.Float32, + (b, qh, sq), + stride_order=(2, 1, 0), + assumed_align=16, + ) # Sinks tensor always part of the ABI; read only when CFG.HAS_SINK == 1 (compile-time fold). fake_sinks = cute.runtime.make_fake_compact_tensor( cutlass.Float32, diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py index 7bf1f556a..0820cbb13 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py @@ -175,7 +175,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, @@ -918,7 +918,7 @@ def _softmax_warp_group( tmem_ptr_i32, bars, sched, - lse_tensor: cute.Tensor, + lse_tensor: Optional[cute.Tensor], sinks_tensor: cute.Tensor, seq_kv_lens_tensor, seq_q_lens_tensor, @@ -1322,7 +1322,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, seq_q_lens_tensor, @@ -1485,16 +1485,20 @@ def _correction_warp_group( neg_inf_trim = cutlass.Float32(float("-inf")) lse_val = cutlass.Float32(arith.select(row_trim.ir_value(), neg_inf_trim.ir_value(), lse_val.ir_value())) inv_sum = cutlass.Float32(arith.select(row_trim.ir_value(), cutlass.Float32(0.0).ir_value(), inv_sum.ir_value())) - if cutlass.const_expr(CFG.THD_VARLEN): + if cutlass.const_expr(lse_tensor is None): + pass # has_lse=False: the Stats store is compiled out + elif cutlass.const_expr(CFG.THD_VARLEN): _cu = cutlass.make_array_view(seq_kv_lens_tensor) _cu_q_b = cutlass.Int32(_cu[n_batch + batch_idx]) _s_q_b = cutlass.Int32(_cu[n_batch + batch_idx + cutlass.Int32(1)]) - _cu_q_b if q_row_global < _s_q_b: lse_arr = cutlass.make_array_view(lse_tensor) - if cutlass.const_expr(CFG.THD_LSE_TOKEN_MAJOR): - lse_row = lse_arr[cutlass.Int32(0), _cu_q_b + q_row_global, :] + if cutlass.const_expr(len(lse_tensor.shape) == 2): + # token-major packed (T, H) + lse_row = lse_arr[_cu_q_b + q_row_global, :] lse_row[head_idx] = lse_val else: + # head-major packed (1, QH, head_stride) lse_row = lse_arr[cutlass.Int32(0), head_idx, :] lse_row[_cu_q_b + q_row_global] = lse_val else: @@ -1576,7 +1580,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, @@ -1676,7 +1680,16 @@ 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, d_qk: int = CFG.TILE_K, d_v: int = CFG.TILE_O, lse_stride: int = 0 + b: int = 1, + qh: int = 1, + kh: int = 1, + sq: int = 256, + skv: int = 128, + d_qk: int = CFG.TILE_K, + d_v: int = CFG.TILE_O, + has_lse: bool = True, + lse_token_major: bool = False, + lse_stride: int = 0, ) -> Callable: # noqa: A001 """ENVELOPE: ``d_qk`` / ``d_v`` are the ACTUAL head dims (defaults = full TILE_K / TILE_O). TMA descriptors carry these extents while the tile box @@ -1712,25 +1725,48 @@ def compile( stride_order=(3, 2, 1, 0), assumed_align=16, ) - if CFG.THD_VARLEN: - # Packed ragged-Stats LSE in the caller's declared layout (align 4: - # the store is scalar f32 and the caller's Stats buffer only - # guarantees element alignment). Token-major [1, T, QH]; head-major - # [1, QH, head_stride] with head_stride >= T (compact when 0). - _lse_hs = lse_stride if lse_stride else sq - if _lse_hs < sq: - raise ValueError(f"THD head-major LSE head_stride ({_lse_hs}) must cover the packed Q token total ({sq})") - _fake_lse_shape = (1, sq, qh) if CFG.THD_LSE_TOKEN_MAJOR else (1, qh, _lse_hs) + 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. + if lse_token_major or lse_stride: + raise ValueError("lse_token_major / lse_stride require has_lse=True") + fake_lse = None + elif CFG.THD_VARLEN: + # Packed ragged-Stats LSE in the caller's declared layout (align 4: the + # store is scalar f32 and the caller's Stats buffer only guarantees + # element alignment). Token-major = its natural packed rank-2 (T, H) + # view; head-major = the kernels' native rank-3 (1, QH, head_stride) + # packing with head_stride >= T (compact when 0). The epilogue store + # branches on the STATIC rank, so the layout is fully encoded in this + # fake tensor — no template parameter. + if lse_token_major: + if lse_stride: + raise ValueError("lse_stride is head-major-only (token-major (T, H) is compact)") + fake_lse = cute.runtime.make_fake_compact_tensor( + cutlass.Float32, + (sq, qh), + stride_order=(1, 0), + assumed_align=4, + ) + else: + _lse_hs = lse_stride if lse_stride else sq + if _lse_hs < sq: + raise ValueError(f"THD head-major LSE head_stride ({_lse_hs}) must cover the packed Q token total ({sq})") + fake_lse = cute.runtime.make_fake_compact_tensor( + cutlass.Float32, + (1, qh, _lse_hs), + stride_order=(2, 1, 0), + assumed_align=4, + ) else: - if lse_stride: - raise ValueError("lse_stride is THD-only (dense LSE is compact (B, H, Sq))") - _fake_lse_shape = (b, qh, sq) - fake_lse = cute.runtime.make_fake_compact_tensor( - cutlass.Float32, - _fake_lse_shape, - stride_order=(2, 1, 0), - assumed_align=4 if CFG.THD_VARLEN else 16, - ) + if lse_token_major or lse_stride: + raise ValueError("lse_token_major / lse_stride are THD-only (dense LSE is compact (B, H, Sq))") + fake_lse = cute.runtime.make_fake_compact_tensor( + cutlass.Float32, + (b, qh, sq), + stride_order=(2, 1, 0), + assumed_align=16, + ) fake_sinks = cute.runtime.make_fake_compact_tensor( cutlass.Float32, (qh,), diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py index 4d7154c97..a72e74e7f 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py @@ -324,7 +324,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, @@ -1145,16 +1145,20 @@ def _compute_warp_group( bars.mb_tma_o_full[chunk].arrive() q_row_global = q_super_idx * cutlass.Int32(CFG.TILES_Q * CFG.TILE_M) + tid_in_wg - if cutlass.const_expr(CFG.THD_VARLEN): + if cutlass.const_expr(lse_tensor is None): + pass # has_lse=False: the Stats store is compiled out + elif cutlass.const_expr(CFG.THD_VARLEN): _cu = cutlass.make_array_view(seq_kv_lens_tensor) _cu_q_b = cutlass.Int32(_cu[n_batch + batch_idx]) _s_q_b = cutlass.Int32(_cu[n_batch + batch_idx + cutlass.Int32(1)]) - _cu_q_b if q_row_global < _s_q_b: lse_arr = cutlass.make_array_view(lse_tensor) - if cutlass.const_expr(CFG.THD_LSE_TOKEN_MAJOR): - lse_row = lse_arr[cutlass.Int32(0), _cu_q_b + q_row_global, :] + if cutlass.const_expr(len(lse_tensor.shape) == 2): + # token-major packed (T, H) + lse_row = lse_arr[_cu_q_b + q_row_global, :] lse_row[head_idx] = lse else: + # head-major packed (1, QH, head_stride) lse_row = lse_arr[cutlass.Int32(0), head_idx, :] lse_row[_cu_q_b + q_row_global] = lse else: @@ -1779,7 +1783,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, @@ -1879,7 +1883,16 @@ 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, d_qk: int = CFG.TILE_K, d_v: int = CFG.TILE_O, lse_stride: int = 0 + b: int = 1, + qh: int = 1, + kh: int = 1, + sq: int = 256, + skv: int = 128, + d_qk: int = CFG.TILE_K, + d_v: int = CFG.TILE_O, + has_lse: bool = True, + lse_token_major: bool = False, + lse_stride: int = 0, ) -> Callable: """ENVELOPE: ``d_qk`` / ``d_v`` are the ACTUAL head dims (defaults = full TILE_K / TILE_O). TMA descriptors carry these extents while the tile box @@ -1917,25 +1930,48 @@ def compile( stride_order=(3, 2, 1, 0), assumed_align=16, ) - if CFG.THD_VARLEN: - # Packed ragged-Stats LSE in the caller's declared layout (align 4: - # the store is scalar f32 and the caller's Stats buffer only - # guarantees element alignment). Token-major [1, T, QH]; head-major - # [1, QH, head_stride] with head_stride >= T (compact when 0). - _lse_hs = lse_stride if lse_stride else sq - if _lse_hs < sq: - raise ValueError(f"THD head-major LSE head_stride ({_lse_hs}) must cover the packed Q token total ({sq})") - _fake_lse_shape = (1, sq, qh) if CFG.THD_LSE_TOKEN_MAJOR else (1, qh, _lse_hs) + 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. + if lse_token_major or lse_stride: + raise ValueError("lse_token_major / lse_stride require has_lse=True") + fake_lse = None + elif CFG.THD_VARLEN: + # Packed ragged-Stats LSE in the caller's declared layout (align 4: the + # store is scalar f32 and the caller's Stats buffer only guarantees + # element alignment). Token-major = its natural packed rank-2 (T, H) + # view; head-major = the kernels' native rank-3 (1, QH, head_stride) + # packing with head_stride >= T (compact when 0). The epilogue store + # branches on the STATIC rank, so the layout is fully encoded in this + # fake tensor — no template parameter. + if lse_token_major: + if lse_stride: + raise ValueError("lse_stride is head-major-only (token-major (T, H) is compact)") + fake_lse = cute.runtime.make_fake_compact_tensor( + cutlass.Float32, + (sq, qh), + stride_order=(1, 0), + assumed_align=4, + ) + else: + _lse_hs = lse_stride if lse_stride else sq + if _lse_hs < sq: + raise ValueError(f"THD head-major LSE head_stride ({_lse_hs}) must cover the packed Q token total ({sq})") + fake_lse = cute.runtime.make_fake_compact_tensor( + cutlass.Float32, + (1, qh, _lse_hs), + stride_order=(2, 1, 0), + assumed_align=4, + ) else: - if lse_stride: - raise ValueError("lse_stride is THD-only (dense LSE is compact (B, H, Sq))") - _fake_lse_shape = (b, qh, sq) - fake_lse = cute.runtime.make_fake_compact_tensor( - cutlass.Float32, - _fake_lse_shape, - stride_order=(2, 1, 0), - assumed_align=4 if CFG.THD_VARLEN else 16, - ) + if lse_token_major or lse_stride: + raise ValueError("lse_token_major / lse_stride are THD-only (dense LSE is compact (B, H, Sq))") + fake_lse = cute.runtime.make_fake_compact_tensor( + cutlass.Float32, + (b, qh, sq), + stride_order=(2, 1, 0), + assumed_align=16, + ) fake_sinks = cute.runtime.make_fake_compact_tensor( cutlass.Float32, (qh,), diff --git a/test/python/sdpa/frost/test_sdpa_frontend_integration.py b/test/python/sdpa/frost/test_sdpa_frontend_integration.py index 6e08da3f1..1534dcc20 100644 --- a/test/python/sdpa/frost/test_sdpa_frontend_integration.py +++ b/test/python/sdpa/frost/test_sdpa_frontend_integration.py @@ -126,11 +126,12 @@ def test_select_dsl_engine_runs_and_matches_torch(): g.check_support() g.build_plans() # Honest workspace: this graph is inference (no Stats output), so the - # engine carves a dummy LSE (B*H*S fp32) from the caller's buffer. + # kernel compiles the LSE store out (has_lse=False) — no dummy buffer, + # no workspace at all. ws_size = g.get_workspace_size() - assert ws_size >= B * H * S * 4 + assert ws_size == 0 - ws = torch.empty(ws_size, device="cuda", dtype=torch.uint8) + ws = torch.empty(max(ws_size, 1), device="cuda", dtype=torch.uint8) g.execute({q: q_gpu, k: k_gpu, v: v_gpu, o: o_gpu}, ws) torch.cuda.synchronize() @@ -282,16 +283,49 @@ def test_plan_name_contract(): @_SM100_DSL def test_workspace_carve_no_per_execute_allocs_and_guards(): - """The dummy-LSE (inference) path carves scratch from the caller's - workspace: steady-state executes make ZERO torch CUDA allocations, an - undersized buffer raises instead of corrupting, and get_workspace_size - reports the real requirement (not 0).""" - q_gpu = torch.randn(B, S, H, D, device="cuda", dtype=torch.float16).transpose(1, 2) - k_gpu = torch.randn(B, S, H, D, device="cuda", dtype=torch.float16).transpose(1, 2) - v_gpu = torch.randn(B, S, H, D, device="cuda", dtype=torch.float16).transpose(1, 2) - o_gpu = torch.empty(B, S, H, D, device="cuda", dtype=torch.float16).transpose(1, 2) - - g, q, k, v, o = _build_causal_sdpa() + """The THD metadata path carves scratch from the caller's workspace: + steady-state executes make ZERO torch CUDA allocations, an undersized + buffer raises instead of corrupting, and get_workspace_size reports the + real requirement (not 0). Dense f16 no longer needs a workspace at all — + has_lse=False compiles the stats-less LSE store out — so the carve + mechanics live on the THD lowering.""" + seq_lens = [200, 150] + t = sum(seq_lens) + s_max = max(seq_lens) + dims = (B, H, s_max, D) + strides = (s_max * H * D, D, H * D, 1) + stor = [torch.zeros(B * s_max * H * D, device="cuda", dtype=torch.float16) for _ in range(4)] + for buf in stor[:3]: + buf[: t * H * D].normal_() + q_gpu, k_gpu, v_gpu, o_gpu = (buf.as_strided(dims, strides) for buf in stor) + sl = torch.tensor(seq_lens, dtype=torch.int32, device="cuda").view(B, 1, 1, 1) + cu = torch.tensor([0, seq_lens[0], t], dtype=torch.int64, device="cuda") + ro_t = (cu * H * D).view(B + 1, 1, 1, 1) + + g = cudnn.pygraph(io_data_type=cudnn.data_type.HALF, intermediate_data_type=cudnn.data_type.FLOAT, compute_data_type=cudnn.data_type.FLOAT) + q = g.tensor(dim=dims, stride=strides, data_type=cudnn.data_type.HALF, name="q") + k = g.tensor(dim=dims, stride=strides, data_type=cudnn.data_type.HALF, name="k") + v = g.tensor(dim=dims, stride=strides, data_type=cudnn.data_type.HALF, name="v") + sq = g.tensor_like(sl) + skv = g.tensor_like(sl) + qro, kro, vro, oro = (g.tensor_like(ro_t) for _ in range(4)) + q.set_ragged_offset(qro) + k.set_ragged_offset(kro) + v.set_ragged_offset(vro) + o, _ = g.sdpa( + name="sdpa", + q=q, + k=k, + v=v, + attn_scale=1.0 / (D**0.5), + is_inference=True, + use_causal_mask=True, + use_padding_mask=True, + seq_len_q=sq, + seq_len_kv=skv, + ) + o.set_output(True).set_dim(dims).set_stride(strides) + o.set_ragged_offset(oro) _plan(g) _pin(g, _FROST) g.check_support() @@ -299,24 +333,30 @@ def test_workspace_carve_no_per_execute_allocs_and_guards(): assert g.selected_engine.name == _FROST ws_size = g.get_workspace_size() - assert ws_size >= B * H * S * 4 # dummy LSE: b*h_q*s_q fp32 + assert ws_size > 0 # THD metadata: [slq32 | slk32 | meta | o_desc | sinks dummy] + vp = {q: q_gpu, k: k_gpu, v: v_gpu, o: o_gpu, sq: sl, skv: sl, qro: ro_t, kro: ro_t, vro: ro_t, oro: ro_t} # Undersized / absent workspace: loud failure, no silent allocation. with pytest.raises(ValueError, match="workspace"): - g.execute({q: q_gpu, k: k_gpu, v: v_gpu, o: o_gpu}, torch.empty(1, device="cuda", dtype=torch.uint8)) + g.execute(vp, torch.empty(1, device="cuda", dtype=torch.uint8)) with pytest.raises((ValueError, TypeError), match="workspace"): - g.execute({q: q_gpu, k: k_gpu, v: v_gpu, o: o_gpu}, None) + g.execute(vp, None) ws = torch.empty(ws_size, device="cuda", dtype=torch.uint8) - g.execute({q: q_gpu, k: k_gpu, v: v_gpu, o: o_gpu}, ws) # warm-up: one-time dummy caches fill + g.execute(vp, ws) # warm-up: per-shape kernel compile + one-time caches torch.cuda.synchronize() stats_key = "allocation.all.allocated" before = torch.cuda.memory_stats().get(stats_key, 0) - g.execute({q: q_gpu, k: k_gpu, v: v_gpu, o: o_gpu}, ws) + g.execute(vp, ws) torch.cuda.synchronize() after = torch.cuda.memory_stats().get(stats_key, 0) - assert after == before, f"LSE path made {after - before} per-execute CUDA allocation(s); scratch must be carved from the workspace" - - ref = torch.nn.functional.scaled_dot_product_attention(q_gpu, k_gpu, v_gpu, is_causal=True, scale=1.0 / (D**0.5)) - torch.testing.assert_close(o_gpu, ref, atol=5e-2, rtol=3e-2) + assert after == before, f"THD path made {after - before} per-execute CUDA allocation(s); scratch must be carved from the workspace" + + packed_o = stor[3][: t * H * D].view(t, H, D) + for lo, hi in ((0, seq_lens[0]), (seq_lens[0], t)): + qb = stor[0][lo * H * D : hi * H * D].view(hi - lo, H, D).permute(1, 0, 2) + kb = stor[1][lo * H * D : hi * H * D].view(hi - lo, H, D).permute(1, 0, 2) + vb = stor[2][lo * H * D : hi * H * D].view(hi - lo, H, D).permute(1, 0, 2) + ref = torch.nn.functional.scaled_dot_product_attention(qb, kb, vb, is_causal=True, scale=1.0 / (D**0.5)) + torch.testing.assert_close(packed_o[lo:hi].permute(1, 0, 2), ref, atol=5e-2, rtol=3e-2) diff --git a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py index 51a6cf3af..22c171c6c 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py @@ -93,11 +93,11 @@ def test_sdpa_fwd_dsl_sm100_graph_api(dtype, is_causal, d): _select_engine(graph, engine_name(d)) graph.check_support() graph.build_plans() - # Honest workspace: no Stats output, so the engine carves a dummy LSE - # (b*h*s fp32) from the caller's buffer. - assert graph.get_workspace_size() == b * h * s * 4 + # Honest workspace: no Stats output, so the kernel compiles the LSE store + # out (has_lse=False) — no dummy buffer exists at any level. + assert graph.get_workspace_size() == 0 - workspace = torch.empty(graph.get_workspace_size(), device=device, dtype=torch.uint8) + workspace = torch.empty(max(graph.get_workspace_size(), 1), device=device, dtype=torch.uint8) graph.execute({q: q_gpu, k: k_gpu, v: v_gpu, o: o_gpu}, workspace) torch.cuda.synchronize() @@ -305,9 +305,9 @@ def test_dsl_sm100_execute_sink_lse_contract(): has_sink is a compile-time specialization: substituting a zeros dummy for missing sinks would silently change the softmax denominator (a zero sink logit still contributes exp(0) mass), and sinks passed to a sink-less - kernel would be silently dropped. lse_tensor stays accepted when no - sample_lse was given (the SM100 kernels always write an LSE; the FROST - dispatch hands in workspace scratch), but a requested LSE must be bound. + kernel would be silently dropped. Same for the LSE: has_lse is keyed on + sample_lse (no Stats output -> the store is compiled out), so a requested + LSE must be bound and an unrequested one is rejected, both directions. """ _require_dsl() from cudnn.sdpa.fwd.api_dsl import SdpaFwdDslSm100 @@ -344,8 +344,11 @@ def test_dsl_sm100_execute_sink_lse_contract(): api.execute(q_tensor=q, k_tensor=k, v_tensor=v, o_tensor=o, seq_kv_lens=seq_kv) with pytest.raises(ValueError, match="without per-batch Q lengths"): api.execute(q_tensor=q, k_tensor=k, v_tensor=v, o_tensor=o, seq_q_lens=seq_kv) - # No sample_lse and no lse_tensor: the kernel's mandatory LSE write lands - # in a cached dummy — no per-execute allocation, output still correct. + # No sample_lse: the kernel compiles with has_lse=False (LSE store folded + # out, no dummy buffer anywhere) and the output is still correct — while + # an unrequested lse_tensor is rejected (there is no LSE slot to bind). + with pytest.raises(ValueError, match="without an LSE output"): + api.execute(q_tensor=q, k_tensor=k, v_tensor=v, o_tensor=o, lse_tensor=lse) api.execute(q_tensor=q, k_tensor=k, v_tensor=v, o_tensor=o) torch.cuda.synchronize() o_ref = _ref_sdpa_full(q, k, v, scale=scale, is_causal=True) @@ -364,8 +367,7 @@ def test_dsl_sm100_execute_sink_lse_contract(): api = SdpaFwdDslSm100(sample_q=q, sample_k=k, sample_v=v, sample_o=o, sample_lse=lse_tm, thd=True) assert api.check_support() and api.thd_stats_token_major - # THD execute keeps a strict presence contract in BOTH directions (dense - # keeps accepting an extra lse_tensor — the kernels always write an LSE): + # THD execute keeps the same strict presence contract in both directions: # the raises fire before any packing or launch. api = SdpaFwdDslSm100(sample_q=q, sample_k=k, sample_v=v, sample_o=o, sample_lse=lse_tm, thd=True) assert api.check_support() From 7b5188b3835b378e3a07a10a8235fb26d4958235 Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Fri, 7 Aug 2026 14:32:01 -0700 Subject: [PATCH 3/8] frost(sdpa): drop the vestigial thd_stats capability axis With both the SM120 engine (#508) and the SM100 f16 flavors serving ragged Stats, every row with thd=True also had thd_stats=True, so the dedicated gate could never fire: THD + generate_stats eligibility now follows from thd AND stats alone (the FP8/MXFP8 rows keep thd=False). A future partial bring-up that lands THD before its stats plumbing re-adds the axis with its precise decline message. Co-Authored-By: Claude Fable 5 --- python/cudnn/sdpa/fwd/engines.py | 5 ----- test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py | 3 +-- 2 files changed, 1 insertion(+), 7 deletions(-) diff --git a/python/cudnn/sdpa/fwd/engines.py b/python/cudnn/sdpa/fwd/engines.py index 831b7572a..3a251d19d 100644 --- a/python/cudnn/sdpa/fwd/engines.py +++ b/python/cudnn/sdpa/fwd/engines.py @@ -157,7 +157,6 @@ class Capabilities: # compute it from the GLOBAL S_q — the THD variant of the # bottom_right_padded_seq_q gap above. thd_bottom_right: bool = False - thd_stats: bool = False # ragged Stats output (packed token-major / head-major LSE) # Dense padded + stats needs the per-batch seq_len_q LSE trim (padded # q-rows write LSE=-inf / O=0, cuDNN >= 9.14). Plumbed for the half # kernels via SEQ_Q_LENS_PRESENT; the FP8/MXFP8 kernels lack the epilogue @@ -314,8 +313,6 @@ def mismatch(capabilities: Capabilities, facts: "ga.SdpaGraphFacts", knobs: Opti "bottom-right causal with a dense padding mask carrying per-batch seq_len_q is not " "supported (kernel anchors the BR diagonal at the global S_q, not seq_len_q[b])" ) - if facts.thd and facts.wants_stats and not capabilities.thd_stats: - return "THD with generate_stats is not supported yet" if facts.padded and facts.wants_stats and not facts.thd and not capabilities.padded_stats: return "padding mask with generate_stats is not supported yet (per-batch seq_len_q LSE trim not plumbed)" @@ -369,7 +366,6 @@ def _sm100_spec(d: int, d_v: Optional[int] = None) -> EngineSpec: stats=True, lse_optional=True, thd=True, - thd_stats=True, padded_stats=True, # The f16/bf16 lowering serves any dense B/H/S stride permutation # (padded strides included) with the head dim innermost; the @@ -486,7 +482,6 @@ def _sm120_spec() -> EngineSpec: padded_stats=True, thd=True, thd_bottom_right=True, - thd_stats=True, layouts=frozenset({"bshd", "dense_flex"}), sched_policies=frozenset({SCHED_NATURAL}), tile_ms=frozenset({64, 128}), diff --git a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py index 4bbdda888..a05bef163 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py @@ -845,8 +845,7 @@ def test_dsl_sm120_thd_bottom_right(): @torch_fork_set_rng(seed=30) def test_dsl_sm120_thd_stats(stats_layout: str): """THD + generate_stats: the ragged Stats output is written in the - caller's declared layout — token-major [t, h] or head-major [h, t] - (the SM100 rows reject this combination — thd_stats gap).""" + caller's declared layout — token-major [t, h] or head-major [h, t].""" _run_thd_case(seq_q_lens=[200, 150], seq_kv_lens=[200, 150], is_causal=True, check_stats=True, stats_layout=stats_layout) From 2b8f8215c8f2d6ae10502af3ac53880149931c60 Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Fri, 7 Aug 2026 15:05:38 -0700 Subject: [PATCH 4/8] frost(sdpa): route the all-KV-zero THD case through the kernel dead-row path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review follow-up (AGENTS.md Rule 1 — execute is a zero-surprise hot path): the t_kv == 0 short-cut re-implemented the kernels' dead-row semantics adapter-side with zero_/fill_/copy_ writes — surprise kernel launches and a second copy of the same semantics that can drift. Both kernels already serve dead rows (row_sum <= 0 -> O := 0 and LSE := -inf, or the sink alone), pinned by the live-launch zero-KV sequence tests. The only launch blocker was the zero-token packed K/V view (a CuTe layout mode must be > 0), so the adapters now clamp the packed KV extent to ONE never-dereferenced token — every sequence's KV tile range is empty, so no K/V load is ever issued — bound over storage the contract already guarantees: Q backs K (kh*d_qk <= t_q*qh*d_qk), O backs V (kh*d_v <= t_q*qh*d_v). Views only; the short-cut, the adapter-side fills, and the O zero-fill are gone from both the SM100 and SM120 adapters, and AGENTS.md Rule 1 gains a bullet making the no-degenerate-path-fixups expectation explicit. Testing (cc 10.0): all_kv_zero + zero_length tests 5 passed (now exercising the kernel path); sm100 suite -k "thd or graph_api or stats or contract" 111 passed; integration + analyzer 79 passed. The SM120 suite skips locally (no SM120 GPU); its change is the same mechanical transformation, CI-covered. Co-Authored-By: Claude Fable 5 --- python/cudnn/AGENTS.md | 8 ++ python/cudnn/sdpa/fwd/api_dsl.py | 86 ++++++++++--------- .../sdpa/frost/test_sdpa_fwd_dsl_sm100.py | 9 +- .../sdpa/frost/test_sdpa_fwd_dsl_sm120.py | 9 +- 4 files changed, 62 insertions(+), 50 deletions(-) diff --git a/python/cudnn/AGENTS.md b/python/cudnn/AGENTS.md index 2de22ae4d..d28581a63 100644 --- a/python/cudnn/AGENTS.md +++ b/python/cudnn/AGENTS.md @@ -32,6 +32,14 @@ Numbered so reviews can cite them; the list grows — append, never renumber. raise, never fall back to a zeros dummy (zeros sinks change the softmax denominator; zeros seq lens mask every row — silently wrong output). A provided-but-uncompiled tensor must also raise, never be silently ignored. +- **No degenerate-path fixups either.** A runtime-degenerate input (e.g. an + all-zero THD ``seq_kv_lens``) routes through the kernel's own dead-row + path, the same code live launches use — if the packed extent would be zero + (a CuTe layout mode must be > 0), bind a never-dereferenced dummy view + over storage the contract already guarantees, not a fresh buffer. Never + re-implement the kernel's semantics adapter-side with `fill_`/`copy_`/ + `zero_` writes: those are surprise kernel launches on the execute path, + and a second implementation of the same semantics that can drift. ## Frontend-only kernel package layout diff --git a/python/cudnn/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py index 2cbb9a998..fee0fe36e 100644 --- a/python/cudnn/sdpa/fwd/api_dsl.py +++ b/python/cudnn/sdpa/fwd/api_dsl.py @@ -1015,25 +1015,20 @@ def _execute_thd(self, q_buf, k_buf, v_buf, o_buf, scale_softmax_log2, sinks, se qh, kh = self.h_q, self.h_kv d_qk, d_v = self.head_dim_qk, self.head_dim_v - # Degenerate totals (a CuTe layout mode must be > 0, so the kernel cannot - # be launched over an empty packing; these are runtime values, invisible - # to the plan-time probe): - # * t_q == 0 — no query token exists anywhere, so the packed O/LSE have - # zero rows: nothing to compute or write. - # * t_kv == 0 — every query row is fully masked; cuDNN semantics for a - # dead row are O := 0, LSE := -inf (or the sink alone — its column - # keeps the softmax denominator alive). + # Degenerate total (runtime value, invisible to the plan-time probe): + # t_q == 0 means no query token exists anywhere, so the packed O/LSE + # have zero rows — nothing to compute or write. (t_kv == 0 launches + # normally through the kernel's dead-row path; see the K/V binding + # below.) if t_q == 0: self._logger.debug("execute (THD): t_q == 0, nothing to do") return lse = None - lse_valid = None # the valid-region view (first t_q tokens) the t_kv == 0 fill writes if lse_tensor is not None: if self.thd_stats_token_major: # Natural packed rank-2 (T, H) view — the kernel's epilogue # dispatches on this static rank. lse = lse_tensor.as_strided((t_q, qh), (qh, 1), lse_tensor.storage_offset()) - lse_valid = lse else: head_stride = self.thd_stats_head_stride self._value_error_if( @@ -1041,18 +1036,6 @@ def _execute_thd(self, q_buf, k_buf, v_buf, o_buf, scale_softmax_log2, sinks, se f"head-major THD LSE head_stride ({head_stride}) must cover the packed Q token total ({t_q})", ) lse = lse_tensor.as_strided((1, qh, head_stride), (qh * head_stride, head_stride, 1), lse_tensor.storage_offset()) - lse_valid = lse_tensor.as_strided((qh, t_q), (head_stride, 1), lse_tensor.storage_offset()) - if t_kv == 0: - self._logger.debug("execute (THD): t_kv == 0, zeroing packed O") - o_buf.as_strided((t_q * qh * d_v,), (1,), o_buf.storage_offset()).zero_() - if lse_valid is not None: - if sinks is not None: - sinks_v = self._checked_sinks_1d(sinks) - sinks_v = sinks_v.reshape(1, qh).expand(t_q, qh) if self.thd_stats_token_major else sinks_v.reshape(qh, 1).expand(qh, t_q) - lse_valid.copy_(sinks_v) - else: - lse_valid.fill_(float("-inf")) - return # Per-sequence O TMA descriptors, filled by the kernel's builder pass. o_desc = carver.take(b * 16 + 16, torch.int64) if carver is not None else torch.zeros(b * 16 + 16, dtype=torch.int64, device=dev) @@ -1066,9 +1049,25 @@ def _packed(buf, t, h, d): return buf.as_strided((1, t, h, d), (t * h * d, h * d, d, 1), buf.storage_offset()) Q = _packed(q_buf, t_q, qh, d_qk) - K = _packed(k_buf, t_kv, kh, d_qk) - V = _packed(v_buf, t_kv, kh, d_v) O = _packed(o_buf, t_q, qh, d_v) + if t_kv == 0: + # Every query row is dead (all-zero seq_kv_lens): served by the + # KERNEL's own dead-row path (total_sum <= 0 -> O := 0 and + # LSE := -inf, or the sink alone — its column keeps the softmax + # denominator alive), exactly like a live launch's zero-KV + # sequences — no adapter-side fills on the execute hot path + # (AGENTS.md Rule 1). A zero-token packed K/V view cannot back a + # CuTe layout / TMA descriptor, so clamp the packed KV extent to + # ONE never-dereferenced token (every tile sees kv_left == + # kv_right == 0, so no K/V load is ever issued) bound over + # storage guaranteed large enough: Q backs K (kh*d_qk <= + # t_q*qh*d_qk) and O backs V (kh*d_v <= t_q*qh*d_v). + t_kv = 1 + K = q_buf.as_strided((1, 1, kh, d_qk), (kh * d_qk, kh * d_qk, d_qk, 1), q_buf.storage_offset()) + V = o_buf.as_strided((1, 1, kh, d_v), (kh * d_v, kh * d_v, d_v, 1), o_buf.storage_offset()) + else: + K = _packed(k_buf, t_kv, kh, d_qk) + V = _packed(v_buf, t_kv, kh, d_v) # LSE binding: the caller's ragged Stats buffer in its declared layout # when a Stats output exists; None otherwise — the kernel compiles the # LSE store out (has_lse=False), so no dummy buffer exists at all. @@ -1891,7 +1890,6 @@ def _execute_thd( if t_q == 0: return lse = None - lse_valid = None # the valid-region view (first t_q tokens) the t_kv == 0 fill writes if lse_tensor is not None: if self.thd_stats_head_major: head_stride = self.thd_stats_head_stride @@ -1900,23 +1898,8 @@ def _execute_thd( f"head-major THD LSE head_stride ({head_stride}) must cover the packed Q token total ({t_q})", ) lse = lse_tensor.as_strided((qh, head_stride), (head_stride, 1), lse_tensor.storage_offset()) - lse_valid = lse_tensor.as_strided((qh, t_q), (head_stride, 1), lse_tensor.storage_offset()) else: lse = lse_tensor.as_strided((t_q, qh), (qh, 1), lse_tensor.storage_offset()) - lse_valid = lse - if t_kv == 0: - # Every row is dead: O := 0, LSE := -inf (or the sink alone — - # its column keeps the denominator alive). A zero-token K/V view - # cannot back a TMA descriptor, so short-cut both. - o_buf.as_strided((t_q * qh * d_v,), (1,), o_buf.storage_offset()).zero_() - if lse_valid is not None: - if sinks is not None: - sinks_v = self._checked_sinks_1d(sinks) - sinks_v = sinks_v.reshape(qh, 1).expand(qh, t_q) if self.thd_stats_head_major else sinks_v.reshape(1, qh).expand(t_q, qh) - lse_valid.copy_(sinks_v) - else: - lse_valid.fill_(float("-inf")) - return # Sinks are None-specialized like the LSE when the graph has no sink # token. @@ -1932,6 +1915,25 @@ def _packed(buf, tokens): d = d_qk if buf is q_buf or buf is k_buf else d_v return buf.as_strided((1, tokens, heads, d), (tokens * heads * d, heads * d, d, 1), buf.storage_offset()) + if t_kv == 0: + # Every query row is dead (all-zero seq_kv_lens): served by the + # KERNEL's own dead-row path (row_sum <= 0 -> O := 0 and + # LSE := -inf, or the sink alone — its column keeps the softmax + # denominator alive), exactly like a live launch's zero-KV + # sequences — no adapter-side fills on the execute hot path + # (AGENTS.md Rule 1). A zero-token packed K/V view cannot back a + # CuTe layout, so clamp the packed KV extent to ONE + # never-dereferenced token (every sequence's KV tile range is + # empty, so no K/V load is ever issued) bound over storage + # guaranteed large enough: Q backs K (kh*d_qk <= t_q*qh*d_qk) and + # O backs V (kh*d_v <= t_q*qh*d_v). + t_kv = 1 + K = q_buf.as_strided((1, 1, kh, d_qk), (kh * d_qk, kh * d_qk, d_qk, 1), q_buf.storage_offset()) + V = o_buf.as_strided((1, 1, kh, d_v), (kh * d_v, kh * d_v, d_v, 1), o_buf.storage_offset()) + else: + K = _packed(k_buf, t_kv) + V = _packed(v_buf, t_kv) + if current_stream is None: current_stream = cuda.CUstream(torch.cuda.current_stream(dev).cuda_stream) @@ -1953,8 +1955,8 @@ def _packed(buf, tokens): ) fn( _packed(q_buf, t_q), - _packed(k_buf, t_kv), - _packed(v_buf, t_kv), + K, + V, _packed(o_buf, t_q), lse, sinks_t, diff --git a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py index 22c171c6c..ead66b815 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py @@ -883,10 +883,11 @@ def test_dsl_sm100_thd_zero_length_sequence_stats(): @pytest.mark.parametrize("with_sink", [False, True], ids=["no_sink", "sink"]) @torch_fork_set_rng(seed=31) def test_dsl_sm100_thd_all_kv_zero_stats(with_sink, stats_layout): - """Every KV length zero: a zero-token K/V view cannot back a TMA - descriptor, so the adapter short-cut fills O := 0 and the ragged Stats - adapter-side — -inf, or the sink value alone (the sink column keeps the - softmax denominator alive) — in either declared layout.""" + """Every KV length zero: the launch goes through the KERNEL's dead-row + path (O := 0, LSE := -inf, or the sink value alone — the sink column + keeps the softmax denominator alive) with the packed KV extent clamped + to one never-dereferenced token (a zero-token K/V view cannot back a + CuTe layout) — no adapter-side fills, in either declared layout.""" _run_thd_stats_case(seq_lens_q=[64, 32], seq_lens_kv=[0, 0], mask="none", with_sink=with_sink, stats_layout=stats_layout) diff --git a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py index a05bef163..62d863b8e 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py @@ -886,10 +886,11 @@ def test_dsl_sm120_thd_zero_length_sequence(): @pytest.mark.parametrize("with_sink", [False, True], ids=["no_sink", "sink"]) @torch_fork_set_rng(seed=31) def test_dsl_sm120_thd_all_kv_zero_stats(with_sink: bool, stats_layout: str): - """Every KV length zero: a zero-token K/V view cannot back a TMA - descriptor, so the adapter short-cut fills O := 0 and the ragged Stats - adapter-side — -inf, or the sink value alone (the sink column keeps the - softmax denominator alive) — in either declared layout.""" + """Every KV length zero: the launch goes through the KERNEL's dead-row + path (O := 0, LSE := -inf, or the sink value alone — the sink column + keeps the softmax denominator alive) with the packed KV extent clamped + to one never-dereferenced token (a zero-token K/V view cannot back a + CuTe layout) — no adapter-side fills, in either declared layout.""" _run_thd_case(seq_q_lens=[64, 32], seq_kv_lens=[0, 0], with_sink=with_sink, check_stats=True, stats_layout=stats_layout) From 580835486ca7824e84b0166f591103cdd45f3c3f Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Fri, 7 Aug 2026 16:45:05 -0700 Subject: [PATCH 5/8] frost(sdpa) tests: pin t_q == 0 as a complete no-op (all-zero seq_len_q) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The adapters early-return when no query token exists anywhere (t_q == 0: the packed O/Stats have zero rows, nothing to compute or write), but no test pinned it — the zero-length coverage always kept live Q tokens. The random backend sweeps can generate the case (a B=1 batch draws a zero seq_len_q with 10% probability), so frameworks do hit it. The THD harnesses now pre-fill the O and ragged Stats storages with a sentinel (2048.0, exact in fp16/bf16/fp32): live tests still compare the kernel-written packed region against the reference, and the new test_dsl_sm1xx_thd_all_q_zero_stats (both stats layouts, live KV and all-zero KV) asserts the buffers come back untouched end to end through the graph -> engine -> adapter stack. The SM100 harness gains the same declared-extent clamp for all-zero seq_len_q that it already had for seq_len_kv (SM120's harness had both). Testing (cc 10.0): new all_q_zero tests + sentinel-affected neighbors 11 passed; full sm100 THD/stats slice 100 passed. SM120 mirror is CI-covered (no SM120 GPU locally). Co-Authored-By: Claude Fable 5 --- .../sdpa/frost/test_sdpa_fwd_dsl_sm100.py | 36 ++++++++++++++++--- .../sdpa/frost/test_sdpa_fwd_dsl_sm120.py | 34 ++++++++++++++++-- 2 files changed, 63 insertions(+), 7 deletions(-) diff --git a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py index ead66b815..3645b02c4 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py @@ -111,6 +111,9 @@ def test_sdpa_fwd_dsl_sm100_graph_api(dtype, is_causal, d): _FLAVOR_IDS = ["dsv4_d512", "qwen_d256", "llama_d128"] _DTYPES = [torch.float16, torch.bfloat16] _DTYPE_IDS = ["fp16", "bf16"] +# Exact in fp16/bf16/fp32: pre-fills O/Stats storages in the THD harness so +# no-op paths (t_q == 0) can assert the buffers came back untouched. +_THD_SENTINEL = 2048.0 def _ref_sdpa_full(q, k, v, *, scale, is_causal=False, bottom_right=False, swa_window=None, seq_kv_lens=None, sinks=None, return_stats=False): @@ -636,9 +639,9 @@ def _run_dsl_thd_graph( dev = "cuda" B = len(seq_lens_q) T_q, T_kv = cu_q[-1], cu_k[-1] - # Clamp the declared KV extent: an all-zero seq_len_kv batch still needs a + # Clamp the declared extents: an all-zero seq_len batch still needs a # rank-legal (>0) graph dim; the padding mask carries the real lengths. - S_max_q, S_max_kv = max(seq_lens_q), max(max(seq_lens_kv), 1) + S_max_q, S_max_kv = max(max(seq_lens_q), 1), max(max(seq_lens_kv), 1) def _dense_buf(packed, s_max, t, H): stride = (s_max * H * d, d, H * d, 1) @@ -649,7 +652,10 @@ def _dense_buf(packed, s_max, t, H): _, q_gpu, stride_q = _dense_buf(q_pk, S_max_q, T_q, H_q) _, k_gpu, stride_kv = _dense_buf(k_pk, S_max_kv, T_kv, H_kv) _, v_gpu, _ = _dense_buf(v_pk, S_max_kv, T_kv, H_kv) - o_stor = torch.zeros(B * S_max_q * H_q * d, device=dev, dtype=dtype) + # Output/Stats storages carry a SENTINEL: the valid packed region is fully + # written by the kernel (compared against the reference), while everything + # else — the whole buffer when t_q == 0 — must come back untouched. + o_stor = torch.full((B * S_max_q * H_q * d,), _THD_SENTINEL, device=dev, dtype=dtype) o_gpu = o_stor.as_strided((B, H_q, S_max_q, d), stride_q) slq = torch.tensor(seq_lens_q, dtype=torch.int32, device=dev).view(B, 1, 1, 1) @@ -690,12 +696,12 @@ def _dense_buf(packed, s_max, t, H): if stats_layout == "head_major": # [h, t]: tokens contiguous within a head, heads strided by the # padded token capacity; offsets = cu_q * stride_s = cu_q. - stats_stor = torch.empty(H_q * t_cap, dtype=torch.float32, device=dev) + stats_stor = torch.full((H_q * t_cap,), _THD_SENTINEL, dtype=torch.float32, device=dev) stats.set_dim((B, H_q, S_max_q, 1)).set_stride((H_q * t_cap, t_cap, 1, 1)) stats_ro_t = (ro_q.flatten() // (H_q * d)).view(B + 1, 1, 1, 1).contiguous() else: # [t, h]: heads contiguous within a token; offsets = cu_q * h_q. - stats_stor = torch.empty(B * S_max_q * H_q, dtype=torch.float32, device=dev) + stats_stor = torch.full((B * S_max_q * H_q,), _THD_SENTINEL, dtype=torch.float32, device=dev) stats.set_dim((B, H_q, S_max_q, 1)).set_stride((S_max_q * H_q, 1, H_q, 1)) stats_ro_t = (ro_q.flatten() // d).view(B + 1, 1, 1, 1).contiguous() stats_ro = g.tensor_like(stats_ro_t, name="stats_ro") @@ -813,6 +819,13 @@ def _cu(sl): stats_layout=stats_layout, ) + if T_q == 0: + # No query token exists anywhere: execute must be a complete no-op — + # the sentinel-filled O and ragged Stats storages come back untouched. + assert (o_stor == _THD_SENTINEL).all(), "t_q == 0 wrote to O" + assert (stats_stor == _THD_SENTINEL).all(), "t_q == 0 wrote to the ragged Stats" + return + if stats_layout == "head_major": packed_stats = stats_stor.view(H_q, t_cap) # (H, head_stride); tokens at [:, cu[i]:cu[i+1]] else: @@ -892,6 +905,19 @@ def test_dsl_sm100_thd_all_kv_zero_stats(with_sink, stats_layout): _run_thd_stats_case(seq_lens_q=[64, 32], seq_lens_kv=[0, 0], mask="none", with_sink=with_sink, stats_layout=stats_layout) +@pytest.mark.L1 +@pytest.mark.parametrize("stats_layout", ["token_major", "head_major"]) +@torch_fork_set_rng(seed=34) +def test_dsl_sm100_thd_all_q_zero_stats(stats_layout): + """Every Q length zero (t_q == 0): no query token exists anywhere, so the + packed O/Stats have zero rows and execute must be a complete NO-OP — the + sentinel-filled buffers come back untouched, with live KV and with the + fully-degenerate all-zero KV as well.""" + + _run_thd_stats_case(seq_lens_q=[0, 0], seq_lens_kv=[50, 30], mask="none", stats_layout=stats_layout) + _run_thd_stats_case(seq_lens_q=[0, 0], seq_lens_kv=[0, 0], mask="none", stats_layout=stats_layout) + + _COMBO_MASKS = { "dense": ["none", "causal", "causal_br", "swa", "padded"], # THD forces padding internally; bottom-right causal is a kernel gap (BR diff --git a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py index 62d863b8e..ddcaf2869 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py @@ -47,6 +47,11 @@ def _select_engine(graph, name): return graph +# Exact in fp16/bf16/fp32: pre-fills O/Stats storages in the THD harness so +# no-op paths (t_q == 0) can assert the buffers came back untouched. +_THD_SENTINEL = 2048.0 + + def _bhsd( batch: int, heads: int, @@ -376,6 +381,10 @@ def _run_thd_case( k_view, _, k_ro = _pack_thd([s.contiguous() for s in k_seqs], s_kv_max) v_view, _, v_ro = _pack_thd([s.contiguous() for s in v_seqs], s_kv_max) o_view, o_storage, o_ro = _pack_thd([torch.zeros(1, h_q, max(n, 1), d_v, dtype=dtype, device="cuda")[:, :, :n] for n in seq_q_lens], s_q_max) + # SENTINEL fill: the kernel writes every valid packed O token (compared + # against the reference below); everything else — the whole buffer when + # t_q == 0 — must come back untouched. + o_storage.fill_(_THD_SENTINEL) sq_t = torch.tensor(seq_q_lens, dtype=torch.int32, device="cuda").view(batch, 1, 1, 1) skv_t = torch.tensor(seq_kv_lens, dtype=torch.int32, device="cuda").view(batch, 1, 1, 1) sinks = torch.randn(1, h_q, 1, 1, dtype=torch.float32, device="cuda") if with_sink else None @@ -429,12 +438,12 @@ def _run_thd_case( if stats_layout == "head_major": # [h, t]: tokens contiguous within a head, heads strided by the # padded token capacity; offsets = cu_q * stride_s = cu_q. - stats_storage = torch.empty(h_q * t_cap, dtype=torch.float32, device="cuda") + stats_storage = torch.full((h_q * t_cap,), _THD_SENTINEL, dtype=torch.float32, device="cuda") stats.set_dim((batch, h_q, s_q_max, 1)).set_stride((h_q * t_cap, t_cap, 1, 1)) stats_ro_t = (q_ro.flatten() // (head_dim * h_q)).view(batch + 1, 1, 1, 1).contiguous() else: # [t, h]: heads contiguous within a token; offsets = cu_q * h_q. - stats_storage = torch.empty(batch * s_q_max * h_q, dtype=torch.float32, device="cuda") + stats_storage = torch.full((batch * s_q_max * h_q,), _THD_SENTINEL, dtype=torch.float32, device="cuda") stats.set_dim((batch, h_q, s_q_max, 1)).set_stride((s_q_max * h_q, 1, h_q, 1)) stats_ro_t = (q_ro.flatten() // head_dim).view(batch + 1, 1, 1, 1).contiguous() stats_ro = graph.tensor_like(stats_ro_t, name="stats_ro") @@ -455,6 +464,14 @@ def _run_thd_case( cu = [0] for n in seq_q_lens: cu.append(cu[-1] + n) + if cu[-1] == 0: + # No query token exists anywhere (t_q == 0): execute must be a + # complete no-op — the sentinel-filled O and ragged Stats storages + # come back untouched. + assert (o_storage == _THD_SENTINEL).all(), "t_q == 0 wrote to O" + if check_stats: + assert (stats_storage == _THD_SENTINEL).all(), "t_q == 0 wrote to the ragged Stats" + return packed_o = o_storage[: cu[-1] * h_q * d_v].view(max(cu[-1], 1), h_q, d_v) if check_stats and stats_layout == "head_major": packed_stats = stats_storage.view(h_q, t_cap) # (H, head_stride); tokens at [:, cu[i]:cu[i+1]] @@ -895,6 +912,19 @@ def test_dsl_sm120_thd_all_kv_zero_stats(with_sink: bool, stats_layout: str): _run_thd_case(seq_q_lens=[64, 32], seq_kv_lens=[0, 0], with_sink=with_sink, check_stats=True, stats_layout=stats_layout) +@pytest.mark.L1 +@pytest.mark.parametrize("stats_layout", ["token_major", "head_major"]) +@torch_fork_set_rng(seed=34) +def test_dsl_sm120_thd_all_q_zero_stats(stats_layout: str): + """Every Q length zero (t_q == 0): no query token exists anywhere, so the + packed O/Stats have zero rows and execute must be a complete NO-OP — the + sentinel-filled buffers come back untouched, with live KV and with the + fully-degenerate all-zero KV as well.""" + + _run_thd_case(seq_q_lens=[0, 0], seq_kv_lens=[50, 30], check_stats=True, stats_layout=stats_layout) + _run_thd_case(seq_q_lens=[0, 0], seq_kv_lens=[0, 0], check_stats=True, stats_layout=stats_layout) + + @pytest.mark.L0 @torch_fork_set_rng(seed=9) def test_dsl_sm120_dense_flex_bhsd_contiguous(): From 9def37d3f923540b5b89bc826a61a9611c706e50 Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Fri, 7 Aug 2026 16:56:10 -0700 Subject: [PATCH 6/8] docs: tighten the AGENTS.md degenerate-path bullet Co-Authored-By: Claude Fable 5 --- python/cudnn/AGENTS.md | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/python/cudnn/AGENTS.md b/python/cudnn/AGENTS.md index d28581a63..156c3816e 100644 --- a/python/cudnn/AGENTS.md +++ b/python/cudnn/AGENTS.md @@ -32,14 +32,12 @@ Numbered so reviews can cite them; the list grows — append, never renumber. raise, never fall back to a zeros dummy (zeros sinks change the softmax denominator; zeros seq lens mask every row — silently wrong output). A provided-but-uncompiled tensor must also raise, never be silently ignored. -- **No degenerate-path fixups either.** A runtime-degenerate input (e.g. an - all-zero THD ``seq_kv_lens``) routes through the kernel's own dead-row - path, the same code live launches use — if the packed extent would be zero - (a CuTe layout mode must be > 0), bind a never-dereferenced dummy view - over storage the contract already guarantees, not a fresh buffer. Never - re-implement the kernel's semantics adapter-side with `fill_`/`copy_`/ - `zero_` writes: those are surprise kernel launches on the execute path, - and a second implementation of the same semantics that can drift. +- **No degenerate-path fixups.** Runtime-degenerate inputs (e.g. all-zero + THD ``seq_kv_lens``) go through the kernel's own dead-row path — never + re-implemented adapter-side with `fill_`/`copy_` writes (surprise kernel + launches, and a second copy of the semantics that can drift). If a packed + extent would be zero, bind a never-dereferenced dummy view over storage + the contract already guarantees. ## Frontend-only kernel package layout From 3962481f7aef2613860f85676a346afd52422247 Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Fri, 7 Aug 2026 18:52:29 -0700 Subject: [PATCH 7/8] frost(sdpa): put the compile() noqa on the def line (review) CodeRabbit: the A001 (builtin shadowing) suppression must sit on the `def compile(` line; the three kernels that had it on the closing-paren line were suppressing nothing. Co-Authored-By: Claude Fable 5 --- python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py | 4 ++-- python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py | 4 ++-- python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py | 4 ++-- python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py | 2 +- 4 files changed, 7 insertions(+), 7 deletions(-) diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py index 358732c5e..00a92bee4 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py @@ -2027,7 +2027,7 @@ def _tma_swz(byte_w: int): @lru_cache(maxsize=None) -def compile( +def compile( # noqa: A001 b: int = 1, qh: int = 1, kh: int = 1, @@ -2038,7 +2038,7 @@ def compile( has_lse: bool = True, lse_token_major: bool = False, lse_stride: int = 0, -) -> Callable: # noqa: A001 +) -> Callable: """Compile a kernel with ALL dims concrete to pin TMA descriptor strides at compile time. THD/varlen: q/k/v/o/lse are PACKED with batch dim 1 ([1,T,H,D]); ``b`` is the diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py index 2ed6c4755..673e32b67 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py @@ -2116,7 +2116,7 @@ def _tma_swz(byte_w: int): @lru_cache(maxsize=None) -def compile( +def compile( # noqa: A001 b: int = 1, qh: int = 1, kh: int = 1, @@ -2127,7 +2127,7 @@ def compile( has_lse: bool = True, lse_token_major: bool = False, lse_stride: int = 0, -) -> Callable: # noqa: A001 +) -> Callable: """Compile a kernel with ALL dims concrete to pin TMA descriptor strides at compile time. THD/varlen: q/k/v/o/lse are PACKED with batch dim 1 ([1,T,H,D]); ``b`` is the diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py index 0820cbb13..0a4e32809 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py @@ -1679,7 +1679,7 @@ def _tma_swz(byte_w: int): @lru_cache(maxsize=None) -def compile( +def compile( # noqa: A001 b: int = 1, qh: int = 1, kh: int = 1, @@ -1690,7 +1690,7 @@ def compile( has_lse: bool = True, lse_token_major: bool = False, lse_stride: int = 0, -) -> Callable: # noqa: A001 +) -> Callable: """ENVELOPE: ``d_qk`` / ``d_v`` are the ACTUAL head dims (defaults = full TILE_K / TILE_O). TMA descriptors carry these extents while the tile box stays the compile-time TILE geometry: loads past d_qk / d_v zero-fill diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py index a72e74e7f..d93c2a7fe 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py @@ -1882,7 +1882,7 @@ def _tma_swz(byte_w: int): @lru_cache(maxsize=None) -def compile( +def compile( # noqa: A001 b: int = 1, qh: int = 1, kh: int = 1, From fe4f15457ec1f0253b8b9d68c48253a78b0adda7 Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Fri, 7 Aug 2026 21:02:57 -0700 Subject: [PATCH 8/8] frost(sdpa): standardize on thd_stats_head_major / lse_head_major (review) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Haobin: SM100 stored thd_stats_token_major (False = its kernel's native head-major packing) while SM120 stores thd_stats_head_major (False = token-major) — each flag named after its kernel's NON-native layout. Standardize both adapters and every kernel compile() on the SM120 / contract-aligned vocabulary: thd_stats_head_major == False means token-major, matching cuDNN's TH1 ragged Stats recipe, and the SM100 kernels' compile() keywords become lse_head_major / lse_head_stride — identical signatures across all five kernels. Naming/polarity only; the layout is always derived explicitly from the graph's declared strides on every live path, so no behavior changes. Also fixes the stale [slq32 | slk32] workspace comment in the carve test. Testing (cc 10.0): sm100 stats/contract slice 17 passed. Co-Authored-By: Claude Fable 5 --- python/cudnn/sdpa/fwd/api_dsl.py | 18 +++---- .../fwd/kernels/prefill_d128_f16_sm100.py | 52 ++++++++++--------- .../kernels/prefill_d192_d128_f16_sm100.py | 43 +++++++-------- .../fwd/kernels/prefill_d256_f16_sm100.py | 43 +++++++-------- .../fwd/kernels/prefill_d512_f16_sm100.py | 43 +++++++-------- .../frost/test_sdpa_frontend_integration.py | 2 +- .../sdpa/frost/test_sdpa_fwd_dsl_sm100.py | 4 +- 7 files changed, 105 insertions(+), 100 deletions(-) diff --git a/python/cudnn/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py index fee0fe36e..cc4fc6954 100644 --- a/python/cudnn/sdpa/fwd/api_dsl.py +++ b/python/cudnn/sdpa/fwd/api_dsl.py @@ -459,7 +459,7 @@ def _initialize_implementation(self) -> None: self.flavor: Optional[tuple[int, int]] = None self.mask_flags = 0 self.swa_window_runtime = 0 - self.thd_stats_token_major = False + self.thd_stats_head_major = False self.thd_stats_head_stride = 0 self._k_mod = None @@ -558,7 +558,7 @@ def check_support(self) -> bool: f"THD LSE must be packed token-major (stride_h == 1, stride_s == H) " f"or head-major (stride_s == 1, stride_h == head_stride); got stride {self.lse_desc.stride}", ) - self.thd_stats_token_major = token_major + self.thd_stats_head_major = head_major self.thd_stats_head_stride = int(stride_h) if head_major else 0 else: self._value_error_if(not self.lse_desc.is_contiguous(), "LSE must be contiguous on SM100 DSL") @@ -1025,17 +1025,17 @@ def _execute_thd(self, q_buf, k_buf, v_buf, o_buf, scale_softmax_log2, sinks, se return lse = None if lse_tensor is not None: - if self.thd_stats_token_major: - # Natural packed rank-2 (T, H) view — the kernel's epilogue - # dispatches on this static rank. - lse = lse_tensor.as_strided((t_q, qh), (qh, 1), lse_tensor.storage_offset()) - else: + if self.thd_stats_head_major: head_stride = self.thd_stats_head_stride self._value_error_if( head_stride < t_q, f"head-major THD LSE head_stride ({head_stride}) must cover the packed Q token total ({t_q})", ) lse = lse_tensor.as_strided((1, qh, head_stride), (qh * head_stride, head_stride, 1), lse_tensor.storage_offset()) + else: + # Token-major (TH1, the default): natural packed rank-2 (T, H) + # view — the kernel's epilogue dispatches on this static rank. + lse = lse_tensor.as_strided((t_q, qh), (qh, 1), lse_tensor.storage_offset()) # Per-sequence O TMA descriptors, filled by the kernel's builder pass. o_desc = carver.take(b * 16 + 16, torch.int64) if carver is not None else torch.zeros(b * 16 + 16, dtype=torch.int64, device=dev) @@ -1093,8 +1093,8 @@ def _packed(buf, t, h, d): # packed rank-2 (T, H) view; head-major carries the caller-declared # head-row stride (0 -> compact sq). has_lse=lse is not None, - lse_token_major=lse is not None and self.thd_stats_token_major, - lse_stride=(self.thd_stats_head_stride if (lse is not None and not self.thd_stats_token_major) else 0), + lse_head_major=lse is not None and self.thd_stats_head_major, + lse_head_stride=(self.thd_stats_head_stride if (lse is not None and self.thd_stats_head_major) else 0), ) fn(Q, K, V, O, LSE, sinks_t, meta, o_desc, (b, qh, kh, t_q, t_kv, 0), cutlass.Float32(scale_softmax_log2), cutlass.Int32(units), stream=current_stream) self._logger.debug("execute (THD) completed") diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py index 00a92bee4..2f2c60bc9 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py @@ -2036,8 +2036,8 @@ def compile( # noqa: A001 d_qk: int = CFG.TILE_K, d_v: int = CFG.TILE_O, has_lse: bool = True, - lse_token_major: bool = False, - lse_stride: int = 0, + lse_head_major: bool = False, + lse_head_stride: int = 0, ) -> Callable: """Compile a kernel with ALL dims concrete to pin TMA descriptor strides at compile time. @@ -2045,10 +2045,11 @@ def compile( # noqa: A001 LOGICAL batch (sequence count) driving n_batch / metadata + O-desc sizes. ``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. THD Stats layouts: ``lse_token_major=True`` = packed rank-2 (T, H); - default head-major = rank-3 [1, QH, head_stride], where ``lse_stride`` is - the caller-declared head-row stride (0 → compact, i.e. ``sq``). All three - are shapes/specializations of the traced code, so they are part of this + at all. THD Stats layouts: token-major packed rank-2 (T, H) by default + (cuDNN's TH1 ragged Stats recipe); ``lse_head_major=True`` = rank-3 + [1, QH, head_stride], where ``lse_head_stride`` is the caller-declared + head-row stride (0 → compact, i.e. ``sq``). All three are + shapes/specializations of the traced code, so they are part of this cache key. ENVELOPE: ``d_qk`` / ``d_v`` are the ACTUAL head dims (defaults = the @@ -2091,39 +2092,40 @@ def compile( # noqa: A001 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. - if lse_token_major or lse_stride: - raise ValueError("lse_token_major / lse_stride require has_lse=True") + if lse_head_major or lse_head_stride: + raise ValueError("lse_head_major / lse_head_stride require has_lse=True") fake_lse = None elif CFG.THD_VARLEN: # Packed ragged-Stats LSE in the caller's declared layout (align 4: the # store is scalar f32 and the caller's Stats buffer only guarantees - # element alignment). Token-major = its natural packed rank-2 (T, H) - # view; head-major = the kernels' native rank-3 (1, QH, head_stride) - # packing with head_stride >= T (compact when 0). The epilogue store - # branches on the STATIC rank, so the layout is fully encoded in this - # fake tensor — no template parameter. - if lse_token_major: - if lse_stride: - raise ValueError("lse_stride is head-major-only (token-major (T, H) is compact)") + # element alignment). Token-major (the default — cuDNN's TH1 ragged + # Stats recipe) = its natural packed rank-2 (T, H) view; head-major = + # the kernels' native rank-3 (1, QH, head_stride) packing with + # head_stride >= T (compact when 0). The epilogue store branches on + # the STATIC rank, so the layout is fully encoded in this fake tensor + # — no template parameter. + if lse_head_major: + _lse_hs = lse_head_stride if lse_head_stride else sq + if _lse_hs < sq: + raise ValueError(f"THD head-major LSE head_stride ({_lse_hs}) must cover the packed Q token total ({sq})") fake_lse = cute.runtime.make_fake_compact_tensor( cutlass.Float32, - (sq, qh), - stride_order=(1, 0), + (1, qh, _lse_hs), + stride_order=(2, 1, 0), assumed_align=4, ) else: - _lse_hs = lse_stride if lse_stride else sq - if _lse_hs < sq: - raise ValueError(f"THD head-major LSE head_stride ({_lse_hs}) must cover the packed Q token total ({sq})") + if lse_head_stride: + raise ValueError("lse_head_stride is head-major-only (token-major (T, H) is compact)") fake_lse = cute.runtime.make_fake_compact_tensor( cutlass.Float32, - (1, qh, _lse_hs), - stride_order=(2, 1, 0), + (sq, qh), + stride_order=(1, 0), assumed_align=4, ) else: - if lse_token_major or lse_stride: - raise ValueError("lse_token_major / lse_stride are THD-only (dense LSE is compact (B, H, Sq))") + if lse_head_major or lse_head_stride: + raise ValueError("lse_head_major / lse_head_stride are THD-only (dense LSE is compact (B, H, Sq))") fake_lse = cute.runtime.make_fake_compact_tensor( cutlass.Float32, (b, qh, sq), diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py index 673e32b67..fe6773d21 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py @@ -2125,8 +2125,8 @@ def compile( # noqa: A001 d_qk: int = CFG.TILE_K, d_v: int = CFG.TILE_O, has_lse: bool = True, - lse_token_major: bool = False, - lse_stride: int = 0, + lse_head_major: bool = False, + lse_head_stride: int = 0, ) -> Callable: """Compile a kernel with ALL dims concrete to pin TMA descriptor strides at compile time. @@ -2173,39 +2173,40 @@ def compile( # noqa: A001 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. - if lse_token_major or lse_stride: - raise ValueError("lse_token_major / lse_stride require has_lse=True") + if lse_head_major or lse_head_stride: + raise ValueError("lse_head_major / lse_head_stride require has_lse=True") fake_lse = None elif CFG.THD_VARLEN: # Packed ragged-Stats LSE in the caller's declared layout (align 4: the # store is scalar f32 and the caller's Stats buffer only guarantees - # element alignment). Token-major = its natural packed rank-2 (T, H) - # view; head-major = the kernels' native rank-3 (1, QH, head_stride) - # packing with head_stride >= T (compact when 0). The epilogue store - # branches on the STATIC rank, so the layout is fully encoded in this - # fake tensor — no template parameter. - if lse_token_major: - if lse_stride: - raise ValueError("lse_stride is head-major-only (token-major (T, H) is compact)") + # element alignment). Token-major (the default — cuDNN's TH1 ragged + # Stats recipe) = its natural packed rank-2 (T, H) view; head-major = + # the kernels' native rank-3 (1, QH, head_stride) packing with + # head_stride >= T (compact when 0). The epilogue store branches on + # the STATIC rank, so the layout is fully encoded in this fake tensor + # — no template parameter. + if lse_head_major: + _lse_hs = lse_head_stride if lse_head_stride else sq + if _lse_hs < sq: + raise ValueError(f"THD head-major LSE head_stride ({_lse_hs}) must cover the packed Q token total ({sq})") fake_lse = cute.runtime.make_fake_compact_tensor( cutlass.Float32, - (sq, qh), - stride_order=(1, 0), + (1, qh, _lse_hs), + stride_order=(2, 1, 0), assumed_align=4, ) else: - _lse_hs = lse_stride if lse_stride else sq - if _lse_hs < sq: - raise ValueError(f"THD head-major LSE head_stride ({_lse_hs}) must cover the packed Q token total ({sq})") + if lse_head_stride: + raise ValueError("lse_head_stride is head-major-only (token-major (T, H) is compact)") fake_lse = cute.runtime.make_fake_compact_tensor( cutlass.Float32, - (1, qh, _lse_hs), - stride_order=(2, 1, 0), + (sq, qh), + stride_order=(1, 0), assumed_align=4, ) else: - if lse_token_major or lse_stride: - raise ValueError("lse_token_major / lse_stride are THD-only (dense LSE is compact (B, H, Sq))") + if lse_head_major or lse_head_stride: + raise ValueError("lse_head_major / lse_head_stride are THD-only (dense LSE is compact (B, H, Sq))") fake_lse = cute.runtime.make_fake_compact_tensor( cutlass.Float32, (b, qh, sq), diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py index 0a4e32809..c35fd04d7 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py @@ -1688,8 +1688,8 @@ def compile( # noqa: A001 d_qk: int = CFG.TILE_K, d_v: int = CFG.TILE_O, has_lse: bool = True, - lse_token_major: bool = False, - lse_stride: int = 0, + lse_head_major: bool = False, + lse_head_stride: int = 0, ) -> Callable: """ENVELOPE: ``d_qk`` / ``d_v`` are the ACTUAL head dims (defaults = full TILE_K / TILE_O). TMA descriptors carry these extents while the tile box @@ -1728,39 +1728,40 @@ def compile( # noqa: A001 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. - if lse_token_major or lse_stride: - raise ValueError("lse_token_major / lse_stride require has_lse=True") + if lse_head_major or lse_head_stride: + raise ValueError("lse_head_major / lse_head_stride require has_lse=True") fake_lse = None elif CFG.THD_VARLEN: # Packed ragged-Stats LSE in the caller's declared layout (align 4: the # store is scalar f32 and the caller's Stats buffer only guarantees - # element alignment). Token-major = its natural packed rank-2 (T, H) - # view; head-major = the kernels' native rank-3 (1, QH, head_stride) - # packing with head_stride >= T (compact when 0). The epilogue store - # branches on the STATIC rank, so the layout is fully encoded in this - # fake tensor — no template parameter. - if lse_token_major: - if lse_stride: - raise ValueError("lse_stride is head-major-only (token-major (T, H) is compact)") + # element alignment). Token-major (the default — cuDNN's TH1 ragged + # Stats recipe) = its natural packed rank-2 (T, H) view; head-major = + # the kernels' native rank-3 (1, QH, head_stride) packing with + # head_stride >= T (compact when 0). The epilogue store branches on + # the STATIC rank, so the layout is fully encoded in this fake tensor + # — no template parameter. + if lse_head_major: + _lse_hs = lse_head_stride if lse_head_stride else sq + if _lse_hs < sq: + raise ValueError(f"THD head-major LSE head_stride ({_lse_hs}) must cover the packed Q token total ({sq})") fake_lse = cute.runtime.make_fake_compact_tensor( cutlass.Float32, - (sq, qh), - stride_order=(1, 0), + (1, qh, _lse_hs), + stride_order=(2, 1, 0), assumed_align=4, ) else: - _lse_hs = lse_stride if lse_stride else sq - if _lse_hs < sq: - raise ValueError(f"THD head-major LSE head_stride ({_lse_hs}) must cover the packed Q token total ({sq})") + if lse_head_stride: + raise ValueError("lse_head_stride is head-major-only (token-major (T, H) is compact)") fake_lse = cute.runtime.make_fake_compact_tensor( cutlass.Float32, - (1, qh, _lse_hs), - stride_order=(2, 1, 0), + (sq, qh), + stride_order=(1, 0), assumed_align=4, ) else: - if lse_token_major or lse_stride: - raise ValueError("lse_token_major / lse_stride are THD-only (dense LSE is compact (B, H, Sq))") + if lse_head_major or lse_head_stride: + raise ValueError("lse_head_major / lse_head_stride are THD-only (dense LSE is compact (B, H, Sq))") fake_lse = cute.runtime.make_fake_compact_tensor( cutlass.Float32, (b, qh, sq), diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py index d93c2a7fe..7aa888d25 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py @@ -1891,8 +1891,8 @@ def compile( # noqa: A001 d_qk: int = CFG.TILE_K, d_v: int = CFG.TILE_O, has_lse: bool = True, - lse_token_major: bool = False, - lse_stride: int = 0, + lse_head_major: bool = False, + lse_head_stride: int = 0, ) -> Callable: """ENVELOPE: ``d_qk`` / ``d_v`` are the ACTUAL head dims (defaults = full TILE_K / TILE_O). TMA descriptors carry these extents while the tile box @@ -1933,39 +1933,40 @@ def compile( # noqa: A001 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. - if lse_token_major or lse_stride: - raise ValueError("lse_token_major / lse_stride require has_lse=True") + if lse_head_major or lse_head_stride: + raise ValueError("lse_head_major / lse_head_stride require has_lse=True") fake_lse = None elif CFG.THD_VARLEN: # Packed ragged-Stats LSE in the caller's declared layout (align 4: the # store is scalar f32 and the caller's Stats buffer only guarantees - # element alignment). Token-major = its natural packed rank-2 (T, H) - # view; head-major = the kernels' native rank-3 (1, QH, head_stride) - # packing with head_stride >= T (compact when 0). The epilogue store - # branches on the STATIC rank, so the layout is fully encoded in this - # fake tensor — no template parameter. - if lse_token_major: - if lse_stride: - raise ValueError("lse_stride is head-major-only (token-major (T, H) is compact)") + # element alignment). Token-major (the default — cuDNN's TH1 ragged + # Stats recipe) = its natural packed rank-2 (T, H) view; head-major = + # the kernels' native rank-3 (1, QH, head_stride) packing with + # head_stride >= T (compact when 0). The epilogue store branches on + # the STATIC rank, so the layout is fully encoded in this fake tensor + # — no template parameter. + if lse_head_major: + _lse_hs = lse_head_stride if lse_head_stride else sq + if _lse_hs < sq: + raise ValueError(f"THD head-major LSE head_stride ({_lse_hs}) must cover the packed Q token total ({sq})") fake_lse = cute.runtime.make_fake_compact_tensor( cutlass.Float32, - (sq, qh), - stride_order=(1, 0), + (1, qh, _lse_hs), + stride_order=(2, 1, 0), assumed_align=4, ) else: - _lse_hs = lse_stride if lse_stride else sq - if _lse_hs < sq: - raise ValueError(f"THD head-major LSE head_stride ({_lse_hs}) must cover the packed Q token total ({sq})") + if lse_head_stride: + raise ValueError("lse_head_stride is head-major-only (token-major (T, H) is compact)") fake_lse = cute.runtime.make_fake_compact_tensor( cutlass.Float32, - (1, qh, _lse_hs), - stride_order=(2, 1, 0), + (sq, qh), + stride_order=(1, 0), assumed_align=4, ) else: - if lse_token_major or lse_stride: - raise ValueError("lse_token_major / lse_stride are THD-only (dense LSE is compact (B, H, Sq))") + if lse_head_major or lse_head_stride: + raise ValueError("lse_head_major / lse_head_stride are THD-only (dense LSE is compact (B, H, Sq))") fake_lse = cute.runtime.make_fake_compact_tensor( cutlass.Float32, (b, qh, sq), diff --git a/test/python/sdpa/frost/test_sdpa_frontend_integration.py b/test/python/sdpa/frost/test_sdpa_frontend_integration.py index 1534dcc20..da964bf97 100644 --- a/test/python/sdpa/frost/test_sdpa_frontend_integration.py +++ b/test/python/sdpa/frost/test_sdpa_frontend_integration.py @@ -333,7 +333,7 @@ def test_workspace_carve_no_per_execute_allocs_and_guards(): assert g.selected_engine.name == _FROST ws_size = g.get_workspace_size() - assert ws_size > 0 # THD metadata: [slq32 | slk32 | meta | o_desc | sinks dummy] + assert ws_size > 0 # THD metadata: [meta(seq_kv, cu_q, cu_k) | o_desc | sinks dummy] vp = {q: q_gpu, k: k_gpu, v: v_gpu, o: o_gpu, sq: sl, skv: sl, qro: ro_t, kro: ro_t, vro: ro_t, oro: ro_t} # Undersized / absent workspace: loud failure, no silent allocation. diff --git a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py index 3645b02c4..d1e4d6212 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py @@ -365,10 +365,10 @@ def test_dsl_sm100_execute_sink_lse_contract(): with pytest.raises(ValueError, match="token-major"): SdpaFwdDslSm100(sample_q=q, sample_k=k, sample_v=v, sample_o=o, sample_lse=lse_padded, thd=True).check_support() api = SdpaFwdDslSm100(sample_q=q, sample_k=k, sample_v=v, sample_o=o, sample_lse=lse, thd=True) - assert api.check_support() and not api.thd_stats_token_major and api.thd_stats_head_stride == s + assert api.check_support() and api.thd_stats_head_major and api.thd_stats_head_stride == s lse_tm = torch.empty(s * h, dtype=torch.float32, device="cuda").as_strided((b, h, s), (s * h, 1, h)) api = SdpaFwdDslSm100(sample_q=q, sample_k=k, sample_v=v, sample_o=o, sample_lse=lse_tm, thd=True) - assert api.check_support() and api.thd_stats_token_major + assert api.check_support() and not api.thd_stats_head_major # THD execute keeps the same strict presence contract in both directions: # the raises fire before any packing or launch.