diff --git a/python/cudnn/AGENTS.md b/python/cudnn/AGENTS.md index 2de22ae4d..156c3816e 100644 --- a/python/cudnn/AGENTS.md +++ b/python/cudnn/AGENTS.md @@ -32,6 +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.** 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 diff --git a/python/cudnn/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py index 3eb8e44e8..cc4fc6954 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_head_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_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") self._value_error_if(not torch.cuda.is_available(), "CUDA must be available for SM100 DSL SDPA") device = self.q_desc.device @@ -733,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, @@ -742,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") @@ -759,18 +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 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.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 @@ -804,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: @@ -833,21 +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: - # 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") + 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, @@ -900,7 +913,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 @@ -936,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, @@ -950,13 +973,18 @@ 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). + 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.""" @@ -967,46 +995,47 @@ 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 - # 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, but THD stats are not supported). + # 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 - 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_() - return + lse = None + if lse_tensor is not None: + 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) @@ -1020,16 +1049,29 @@ 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) - # 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 = carver.take(qh * t_q, torch.float32).reshape(1, qh, t_q) - LSE.zero_() + 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: - LSE = torch.zeros(1, qh, t_q, dtype=torch.float32, device=dev) + 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. + LSE = lse if sinks is not None: sinks_t = self._checked_sinks_1d(sinks) elif carver is not None: @@ -1038,7 +1080,22 @@ 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, + # 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_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") @@ -1810,35 +1867,29 @@ 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: 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 @@ -1847,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. @@ -1879,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) @@ -1900,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, @@ -1913,17 +1968,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/engines.py b/python/cudnn/sdpa/fwd/engines.py index 3cd0dc6e9..3a251d19d 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) @@ -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 # packed LSE output plumbing is a follow-up # 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)" @@ -367,6 +364,7 @@ def _sm100_spec(d: int, d_v: Optional[int] = None) -> EngineSpec: padded=True, sink=True, stats=True, + lse_optional=True, thd=True, padded_stats=True, # The f16/bf16 lowering serves any dense B/H/S stride permutation @@ -484,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}), @@ -593,11 +590,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 16b59d1a1..2f2c60bc9 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 ``[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``). @@ -248,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, @@ -1627,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, @@ -1811,16 +1813,27 @@ 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): - # 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. + 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 + # 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) - lse_row = lse_arr[cutlass.Int32(0), head_idx, :] - lse_row[_cu_q_b + q_row_global] = lse_val + 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: if q_row_global < seqlen_q: lse_arr = cutlass.make_array_view(lse_tensor) @@ -1900,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, @@ -2014,11 +2027,30 @@ 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( # noqa: A001 + 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_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. 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. + ``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: 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 flavor's full TILE_K / TILE_O). The Q/K/V/O TMA descriptors are built from @@ -2057,12 +2089,49 @@ def compile(b: int = 1, qh: int = 1, kh: int = 1, sq: int = 256, skv: int = 128, stride_order=(3, 2, 1, 0), assumed_align=16, ) - fake_lse = cute.runtime.make_fake_compact_tensor( - cutlass.Float32, - (_fake_batch, qh, sq), - stride_order=(2, 1, 0), - assumed_align=16, - ) + if not has_lse: + # No Stats output: the LSE argument is None-specialized and the store + # is compiled out entirely — no dummy buffer exists at any level. + 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 (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, + (1, qh, _lse_hs), + stride_order=(2, 1, 0), + assumed_align=4, + ) + else: + 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, + (sq, qh), + stride_order=(1, 0), + assumed_align=4, + ) + else: + 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), + 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 1e6534521..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 @@ -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,16 +1895,26 @@ 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): - # 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. + 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 + # 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) - lse_row = lse_arr[cutlass.Int32(0), head_idx, :] - lse_row[_cu_q_b + q_row_global] = lse_val + 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: if q_row_global < seqlen_q: lse_arr = cutlass.make_array_view(lse_tensor) @@ -1992,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, @@ -2106,7 +2116,18 @@ 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( # noqa: A001 + 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_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. THD/varlen: q/k/v/o/lse are PACKED with batch dim 1 ([1,T,H,D]); ``b`` is the @@ -2149,12 +2170,49 @@ def compile(b: int = 1, qh: int = 1, kh: int = 1, sq: int = 256, skv: int = 128, stride_order=(3, 2, 1, 0), assumed_align=16, ) - fake_lse = cute.runtime.make_fake_compact_tensor( - cutlass.Float32, - (_fake_batch, qh, sq), - stride_order=(2, 1, 0), - assumed_align=16, - ) + if not has_lse: + # No Stats output: the LSE argument is None-specialized and the store + # is compiled out entirely — no dummy buffer exists at any level. + 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 (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, + (1, qh, _lse_hs), + stride_order=(2, 1, 0), + assumed_align=4, + ) + else: + 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, + (sq, qh), + stride_order=(1, 0), + assumed_align=4, + ) + else: + 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), + 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 b2d5a967b..c35fd04d7 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,14 +1485,22 @@ 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) - lse_row = lse_arr[cutlass.Int32(0), head_idx, :] - lse_row[_cu_q_b + q_row_global] = lse_val + 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: if q_row_global < seqlen_q: lse_arr = cutlass.make_array_view(lse_tensor) @@ -1572,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, @@ -1671,7 +1679,18 @@ 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( # noqa: A001 + 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_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 stays the compile-time TILE geometry: loads past d_qk / d_v zero-fill @@ -1706,12 +1725,49 @@ def compile(b: int = 1, qh: int = 1, kh: int = 1, sq: int = 256, skv: int = 128, stride_order=(3, 2, 1, 0), assumed_align=16, ) - fake_lse = cute.runtime.make_fake_compact_tensor( - cutlass.Float32, - (_fake_batch, qh, sq), - stride_order=(2, 1, 0), - assumed_align=16, - ) + if not has_lse: + # No Stats output: the LSE argument is None-specialized and the store + # is compiled out entirely — no dummy buffer exists at any level. + 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 (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, + (1, qh, _lse_hs), + stride_order=(2, 1, 0), + assumed_align=4, + ) + else: + 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, + (sq, qh), + stride_order=(1, 0), + assumed_align=4, + ) + else: + 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), + 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 f29e20000..7aa888d25 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,14 +1145,22 @@ 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) - lse_row = lse_arr[cutlass.Int32(0), head_idx, :] - lse_row[_cu_q_b + q_row_global] = lse + 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: if q_row_global < seqlen_q: lse_arr = cutlass.make_array_view(lse_tensor) @@ -1775,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, @@ -1874,7 +1882,18 @@ 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( # noqa: A001 + 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_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 stays the compile-time TILE geometry: loads past d_qk / d_v zero-fill @@ -1911,12 +1930,49 @@ def compile(b: int = 1, qh: int = 1, kh: int = 1, sq: int = 256, skv: int = 128, stride_order=(3, 2, 1, 0), assumed_align=16, ) - fake_lse = cute.runtime.make_fake_compact_tensor( - cutlass.Float32, - (_fake_batch, qh, sq), - stride_order=(2, 1, 0), - assumed_align=16, - ) + if not has_lse: + # No Stats output: the LSE argument is None-specialized and the store + # is compiled out entirely — no dummy buffer exists at any level. + 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 (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, + (1, qh, _lse_hs), + stride_order=(2, 1, 0), + assumed_align=4, + ) + else: + 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, + (sq, qh), + stride_order=(1, 0), + assumed_align=4, + ) + else: + 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), + 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..da964bf97 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: [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. 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 48751e8bc..d1e4d6212 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() @@ -111,11 +111,17 @@ 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): +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 +145,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(): @@ -297,9 +308,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 @@ -336,18 +347,41 @@ 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) 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 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 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. + 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 +607,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 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(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) @@ -591,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) @@ -613,16 +677,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.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.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") + stats.set_ragged_offset(stats_ro) + vp[stats_ro] = stats_ro_t + vp[stats] = stats_stor g.validate() g.build_operation_graph() @@ -633,7 +718,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 +778,146 @@ 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 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: + 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: 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) + + +@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 4bbdda888..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]] @@ -845,8 +862,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) @@ -887,14 +903,28 @@ 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) +@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(): 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)