From 154161e5ea03d59e82c772f1cbec918bde9fc63d Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Mon, 10 Aug 2026 11:21:01 -0700 Subject: [PATCH 1/4] frost(sdpa): native THD stride support in the SM100/SM120 f16 fwd kernels; decline what TMA cannot express MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A THD tensor may declare a wider token stride than the packed h*d — e.g. a K/V view of a kv-interleaved [T, 2, H, D] buffer (token stride 2*h*d), the layout torch.nn.attention.varlen users produce by slicing a fused KV projection. The THD lowerings rebuilt packed (1, T, H, D) views with hardcoded strides, so such graphs were claimed and silently mis-addressed (100% of O wrong on both sdpa_fwd_prefill_sm120 and the sm100 flavors; caught by PR #516's fuzz coverage and PyTorch's own varlen suite). Native support, no fallback (AGENTS Hard Rule 2): - compile() on all five f16 kernels (sm120, sm100 d128/d192_d128/d256/ d512) takes optional caller-declared (batch, seq, head, elem) strides per tensor (lru cache-key); None keeps the compact specialization bit-for-bit. Strided fakes via make_fake_tensor, validated against the TMA 16-byte global-stride rule. - SM120: kv_tma_desc reads the tensor's strides instead of recomputing packed ones (Q/O offset math was already layout-driven); the entry validator accepts padded 16-byte-granular BSHD storage (compact = the equality special case). - SM100: the Q/K/V/O TMA descriptors are built from the tensor views, so declared strides flow in unchanged; the THD O-descriptor builder steps per-batch bases by O's declared seq-axis stride (o_tensor.stride[1]). - Adapters bind declared-stride (1, T, H, D) views directly. What TMA cannot express is REJECTED in check_support (NotImplementedError naming the offending strides), so the Router falls back to an engine that honors the declaration: non-innermost-contiguous head dim, or token/head strides that are not multiples of 8 elements (sub- granularity strides also violate the graph API's pointer-alignment contract for the backend, so declining is correct, not conservative). - The SM120 FP8 THD path (#509) keeps the packed contract for now: non-packed declarations are declined (_thd_check_strides_packed); extending native strides there is tracked as a follow-up. Verified (torch nightly cu132, ToT develop + PR #516's fuzz tests): gapped seeded repros pass with the frost engines serving natively on cc 10.0 (sm100) and RTX 5080 (sm120); 128-test fwd ragged L0 sweep slice green on cc 10.0 (all four sm100 flavors) and 84-test slice on sm120; ex-ops suite incl. kv-interleaved views 11/11 on both; dense fwd slice 182 passed (dense compile paths pass no strides -> unchanged); packed THD configs bit-for-bit unchanged. Co-Authored-By: Claude Fable 5 --- python/cudnn/sdpa/fwd/api_dsl.py | 107 +++++++++++++++--- .../fwd/kernels/prefill_d128_f16_sm100.py | 43 +++---- .../kernels/prefill_d192_d128_f16_sm100.py | 43 +++---- .../fwd/kernels/prefill_d256_f16_sm100.py | 43 +++---- .../fwd/kernels/prefill_d512_f16_sm100.py | 43 +++---- .../sdpa/fwd/kernels/prefill_f16_sm120.py | 75 ++++++------ python/cudnn/sdpa/fwd/kernels/thd_sm100.py | 7 +- 7 files changed, 212 insertions(+), 149 deletions(-) diff --git a/python/cudnn/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py index 3d1763404..b58e4e458 100644 --- a/python/cudnn/sdpa/fwd/api_dsl.py +++ b/python/cudnn/sdpa/fwd/api_dsl.py @@ -360,6 +360,56 @@ def _to_bshd_writable(tensor: torch.Tensor): scratch = torch.empty_like(view, memory_format=torch.contiguous_format) return view, True, scratch + # -- THD declared-stride binding ------------------------------------------ + # A THD tensor may DECLARE a wider token stride than the packed h*d — e.g. + # a K/V view of a kv-interleaved [T, 2, H, D] buffer (token stride 2*h*d), + # the layout torch.nn.attention.varlen users produce by slicing a fused KV + # projection. The f16 kernels address declared strides NATIVELY + # (layout-driven offset math + TMA-encoded strides); declarations the + # hardware cannot express are REJECTED in check_support — no + # normalization-copy fallback (AGENTS.md Hard Rule 2) — so the router + # picks an engine that honors them instead. + + @staticmethod + def _thd_declared(desc: TensorDesc): + """(token, head, elem) strides a THD tensor declares, and whether they + are the packed contract (h*d, d, 1).""" + h, d = desc.shape[1], desc.shape[3] + st = (int(desc.stride[2]), int(desc.stride[1]), int(desc.stride[3])) + return st, st == (h * d, d, 1) + + def _thd_check_strides_native(self) -> None: + """Reject THD stride declarations the f16 kernels cannot address + natively: TMA's 16-byte global-stride rule at 2 bytes/elem — the head + dim must be innermost-contiguous (elem stride 1) and the token/head + strides multiples of 8 elements (which also keeps every per-sequence + ragged base 16-byte aligned). Whole-token gaps always qualify for + supported head dims; sub-token gaps only in multiples of 8 elements.""" + for desc in (self.q_desc, self.k_desc, self.v_desc, self.o_desc): + (ts, hs, es), _ = self._thd_declared(desc) + self._not_implemented_error_if( + es != 1 or ts % 8 != 0 or hs % 8 != 0, + f"{desc.name} THD strides {tuple(desc.stride)} are not TMA-expressible " + f"(head dim must be innermost-contiguous and token/head strides 16-byte multiples)", + ) + + def _thd_check_strides_packed(self) -> None: + """FP8 THD serves only the packed contract for now (its kernel and + harness are not audited for declared strides) — decline anything else + rather than adapt (AGENTS.md Hard Rule 2).""" + for desc in (self.q_desc, self.k_desc, self.v_desc, self.o_desc): + _, packed = self._thd_declared(desc) + self._not_implemented_error_if( + not packed, + f"{desc.name}: non-packed THD strides {tuple(desc.stride)} are not supported by the FP8 path yet", + ) + + def _thd_view(self, buf: torch.Tensor, desc: TensorDesc, tokens: int) -> torch.Tensor: + """The declared-stride ``(1, T, H, D)`` view over a THD buffer's storage.""" + h, d = desc.shape[1], desc.shape[3] + (ts, hs, es), _ = self._thd_declared(desc) + return buf.as_strided((1, tokens, h, d), (max(tokens, 1) * ts, ts, hs, es), buf.storage_offset()) + def _amax_slot(self, tensor, name: str, device: torch.device) -> torch.Tensor: """The caller's 1-element amax storage, or a cached dummy. @@ -570,6 +620,9 @@ def check_support(self) -> bool: f"strides allowed); got stride {_stride} shape {_shape}", ) + if self.thd: + self._thd_check_strides_native() + b, h_qo, s_qo, d_qk = self.q_desc.shape _, h_kv, s_kv, _ = self.k_desc.shape _, _, _, d_v = self.v_desc.shape @@ -1117,11 +1170,12 @@ def _execute_thd(self, q_buf, k_buf, v_buf, o_buf, scale_softmax_log2, sinks, se cga_tile_m = int(self._k_mod.CGA_TILE_M) units = qh * sum((l + cga_tile_m - 1) // cga_tile_m for l in slq_host) - 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) - O = _packed(o_buf, t_q, qh, d_v) + # Declared-stride (1, T, H, D) views, addressed NATIVELY by the kernel + # (the Q/K/V/O TMA descriptors are built from the tensor views, and + # the THD O-descriptor builder steps by O's declared seq stride); + # check_support rejected any declaration TMA cannot express. + Q = self._thd_view(q_buf, self.q_desc, t_q) + O = self._thd_view(o_buf, self.o_desc, t_q) 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 @@ -1138,8 +1192,8 @@ def _packed(buf, t, h, d): 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) + K = self._thd_view(k_buf, self.k_desc, t_kv) + V = self._thd_view(v_buf, self.v_desc, t_kv) # 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. @@ -1167,6 +1221,13 @@ def _packed(buf, t, h, d): 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), + # Declared strides of the bound views (cache-key): compact views + # reproduce the packed specialization; native non-packed views + # compile their strides into the kernel's addressing. + q_stride=tuple(Q.stride()), + k_stride=tuple(K.stride()), + v_stride=tuple(V.stride()), + o_stride=tuple(O.stride()), ) 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") @@ -1653,6 +1714,11 @@ def check_support(self) -> bool: f"non-broadcast, non-overlapping strides (any B/H/S order, padded " f"strides allowed); got stride {desc.stride} shape {desc.shape}", ) + if self.thd: + if self._pertensor: + self._thd_check_strides_packed() + else: + self._thd_check_strides_native() b, h_q, s_q, d_q = self.q_desc.shape _, h_kv, s_kv, _ = self.k_desc.shape @@ -2138,7 +2204,7 @@ def _scalar(t, default=1.0): amax_o_buf.div_(max(so, 1e-30)) self._logger.debug("execute (SM120 FP8 per-tensor) completed") - def _thd_pack(self, q_buf, k_buf, v_buf, o_buf, seq_q_lens, seq_kv_lens, workspace, label): + def _thd_pack(self, q_buf, k_buf, v_buf, o_buf, seq_q_lens, seq_kv_lens, workspace, label, declared_views=False): """Shared THD (ragged) packing: cu_seqlens metadata + ``(1, T, H, D)`` views. Serves the same fully-packed contract as the SM100 THD path @@ -2183,6 +2249,13 @@ def _thd_pack(self, q_buf, k_buf, v_buf, o_buf, seq_q_lens, seq_kv_lens, workspa def _packed(buf, tokens, heads, d): return buf.as_strided((1, tokens, heads, d), (tokens * heads * d, heads * d, d, 1), buf.storage_offset()) + def _view(buf, desc, tokens, heads, d): + # declared_views: the f16 kernel addresses declared strides + # natively (check_support rejected inexpressible ones); the FP8 + # path keeps the packed contract (check_support declined + # anything else). + return self._thd_view(buf, desc, tokens) if declared_views else _packed(buf, tokens, heads, d) + 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 @@ -2199,18 +2272,18 @@ def _packed(buf, tokens, heads, d): 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) + K = _view(k_buf, self.k_desc, t_kv, kh, d_qk) + V = _view(v_buf, self.v_desc, t_kv, kh, d_v) return SimpleNamespace( meta=meta, t_q=t_q, t_kv=t_kv, max_sq=max_sq, - Q=_packed(q_buf, t_q, qh, d_qk), + Q=_view(q_buf, self.q_desc, t_q, qh, d_qk), K=K, V=V, - O=_packed(o_buf, t_q, qh, d_v), + O=_view(o_buf, self.o_desc, t_q, qh, d_v), seq_q_dummy=self._dummy("seq_q_lens", dev, lambda: torch.zeros(b, dtype=torch.int32, device=dev)), ) @@ -2225,7 +2298,7 @@ def _execute_thd( within each head row. """ - pack = self._thd_pack(q_buf, k_buf, v_buf, o_buf, seq_q_lens, seq_kv_lens, workspace, "SdpaFwdDslSm120 (THD)") + pack = self._thd_pack(q_buf, k_buf, v_buf, o_buf, seq_q_lens, seq_kv_lens, workspace, "SdpaFwdDslSm120 (THD)", declared_views=True) if pack is None: return @@ -2263,6 +2336,14 @@ def _execute_thd( has_lse=self.lse_desc is not None, lse_head_major=self.thd_stats_head_major, lse_head_stride=self.thd_stats_head_stride, + # Declared strides of the bound views (cache-key): compact views + # reproduce the packed specialization; native non-packed views + # compile their strides into the Q/O offset math and K/V TMA + # descriptors. + q_stride=tuple(pack.Q.stride()), + k_stride=tuple(pack.K.stride()), + v_stride=tuple(pack.V.stride()), + o_stride=tuple(pack.O.stride()), ) fn( pack.Q, 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 a478a3447..033cbed07 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py @@ -2007,7 +2007,7 @@ def _tma_swz(byte_w: int): seq_kv_lens_tensor, cutlass.Int32(QH), cutlass.Int32(B), - cutlass.Int32(o_tensor.shape[3]), + cutlass.Int32(o_tensor.stride[1]), ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream) grid_shape = (n_thd_units * cutlass.Int32(CFG.CGA_M), cutlass.Int32(1), cutlass.Int32(1)) else: @@ -2050,6 +2050,10 @@ def compile( # noqa: A001 has_lse: bool = True, lse_head_major: bool = False, lse_head_stride: int = 0, + q_stride: Optional[tuple] = None, + k_stride: Optional[tuple] = None, + v_stride: Optional[tuple] = None, + o_stride: Optional[tuple] = None, ) -> Callable: """Compile a kernel with ALL dims concrete to pin TMA descriptor strides at compile time. @@ -2077,30 +2081,19 @@ def compile( # noqa: A001 if (d_qk * CFG.BPE) % 16 != 0 or (d_v * CFG.BPE_O) % 16 != 0: raise ValueError(f"d128 envelope: d_qk*BPE and d_v*BPE must be 16-byte multiples (TMA global-stride rule); got ({d_qk}, {d_v}) at BPE={CFG.BPE}") _fake_batch = 1 if CFG.THD_VARLEN else b - fake_q = cute.runtime.make_fake_compact_tensor( - STORAGE_DTYPE, - (_fake_batch, sq, qh, d_qk), - stride_order=(3, 2, 1, 0), - assumed_align=16, - ) - fake_k = cute.runtime.make_fake_compact_tensor( - STORAGE_DTYPE, - (_fake_batch, skv, kh, d_qk), - stride_order=(3, 2, 1, 0), - assumed_align=16, - ) - fake_v = cute.runtime.make_fake_compact_tensor( - STORAGE_DTYPE, - (_fake_batch, skv, kh, d_v), - stride_order=(3, 2, 1, 0), - assumed_align=16, - ) - fake_o = cute.runtime.make_fake_compact_tensor( - STORAGE_DTYPE, - (_fake_batch, sq, qh, d_v), - stride_order=(3, 2, 1, 0), - assumed_align=16, - ) + + def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE): + if stride is None: + return cute.runtime.make_fake_compact_tensor(dtype, shape, stride_order=(3, 2, 1, 0), assumed_align=16) + for axis in (1, 2): # seq/head global strides feed TMA: 16-byte rule + if (stride[axis] * CFG.BPE) % 16 != 0: + raise ValueError(f"declared stride {stride} axis {axis} must be a 16-byte multiple at BPE={CFG.BPE} (TMA global-stride rule)") + return cute.runtime.make_fake_tensor(dtype, shape, tuple(stride), assumed_align=16) + + fake_q = _fake_bshd((_fake_batch, sq, qh, d_qk), q_stride) + fake_k = _fake_bshd((_fake_batch, skv, kh, d_qk), k_stride) + fake_v = _fake_bshd((_fake_batch, skv, kh, d_v), v_stride) + fake_o = _fake_bshd((_fake_batch, sq, qh, d_v), o_stride, dtype=STORAGE_DTYPE) 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. 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 3cb4f2158..2dbe49d82 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 @@ -2102,7 +2102,7 @@ def _tma_swz(byte_w: int): seq_kv_lens_tensor, cutlass.Int32(QH), cutlass.Int32(B), - cutlass.Int32(o_tensor.shape[3]), + cutlass.Int32(o_tensor.stride[1]), ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream) grid_shape = (n_thd_units * cutlass.Int32(CFG.CGA_M), cutlass.Int32(1), cutlass.Int32(1)) else: @@ -2145,6 +2145,10 @@ def compile( # noqa: A001 has_lse: bool = True, lse_head_major: bool = False, lse_head_stride: int = 0, + q_stride: Optional[tuple] = None, + k_stride: Optional[tuple] = None, + v_stride: Optional[tuple] = None, + o_stride: Optional[tuple] = None, ) -> Callable: """Compile a kernel with ALL dims concrete to pin TMA descriptor strides at compile time. @@ -2164,30 +2168,19 @@ def compile( # noqa: A001 if (d_qk * CFG.BPE) % 16 != 0 or (d_v * CFG.BPE_O) % 16 != 0: raise ValueError(f"d192 envelope: d_qk*BPE and d_v*BPE must be 16-byte multiples (TMA global-stride rule); got ({d_qk}, {d_v}) at BPE={CFG.BPE}") _fake_batch = 1 if CFG.THD_VARLEN else b - fake_q = cute.runtime.make_fake_compact_tensor( - STORAGE_DTYPE, - (_fake_batch, sq, qh, d_qk), - stride_order=(3, 2, 1, 0), - assumed_align=16, - ) - fake_k = cute.runtime.make_fake_compact_tensor( - STORAGE_DTYPE, - (_fake_batch, skv, kh, d_qk), - stride_order=(3, 2, 1, 0), - assumed_align=16, - ) - fake_v = cute.runtime.make_fake_compact_tensor( - STORAGE_DTYPE, - (_fake_batch, skv, kh, d_v), - stride_order=(3, 2, 1, 0), - assumed_align=16, - ) - fake_o = cute.runtime.make_fake_compact_tensor( - STORAGE_DTYPE, - (_fake_batch, sq, qh, d_v), - stride_order=(3, 2, 1, 0), - assumed_align=16, - ) + + def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE): + if stride is None: + return cute.runtime.make_fake_compact_tensor(dtype, shape, stride_order=(3, 2, 1, 0), assumed_align=16) + for axis in (1, 2): # seq/head global strides feed TMA: 16-byte rule + if (stride[axis] * CFG.BPE) % 16 != 0: + raise ValueError(f"declared stride {stride} axis {axis} must be a 16-byte multiple at BPE={CFG.BPE} (TMA global-stride rule)") + return cute.runtime.make_fake_tensor(dtype, shape, tuple(stride), assumed_align=16) + + fake_q = _fake_bshd((_fake_batch, sq, qh, d_qk), q_stride) + fake_k = _fake_bshd((_fake_batch, skv, kh, d_qk), k_stride) + fake_v = _fake_bshd((_fake_batch, skv, kh, d_v), v_stride) + fake_o = _fake_bshd((_fake_batch, sq, qh, d_v), o_stride, dtype=STORAGE_DTYPE) 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. 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 34b43aed8..c808d9958 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py @@ -1661,7 +1661,7 @@ def _tma_swz(byte_w: int): seq_kv_lens_tensor, cutlass.Int32(QH), cutlass.Int32(B), - cutlass.Int32(o_tensor.shape[3]), + cutlass.Int32(o_tensor.stride[1]), ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream) grid_shape = (n_thd_units * cutlass.Int32(CFG.CGA_M), cutlass.Int32(1), cutlass.Int32(1)) else: @@ -1703,6 +1703,10 @@ def compile( # noqa: A001 has_lse: bool = True, lse_head_major: bool = False, lse_head_stride: int = 0, + q_stride: Optional[tuple] = None, + k_stride: Optional[tuple] = None, + v_stride: Optional[tuple] = None, + o_stride: Optional[tuple] = None, ) -> 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 @@ -1714,30 +1718,19 @@ def compile( # noqa: A001 if (d_qk * CFG.BPE) % 16 != 0 or (d_v * CFG.BPE_O) % 16 != 0: raise ValueError(f"d256 envelope: d_qk*BPE and d_v*BPE must be 16-byte multiples (TMA global-stride rule); got ({d_qk}, {d_v}) at BPE={CFG.BPE}") _fake_batch = 1 if CFG.THD_VARLEN else b - fake_q = cute.runtime.make_fake_compact_tensor( - STORAGE_DTYPE, - (_fake_batch, sq, qh, d_qk), - stride_order=(3, 2, 1, 0), - assumed_align=16, - ) - fake_k = cute.runtime.make_fake_compact_tensor( - STORAGE_DTYPE, - (_fake_batch, skv, kh, d_qk), - stride_order=(3, 2, 1, 0), - assumed_align=16, - ) - fake_v = cute.runtime.make_fake_compact_tensor( - STORAGE_DTYPE, - (_fake_batch, skv, kh, d_v), - stride_order=(3, 2, 1, 0), - assumed_align=16, - ) - fake_o = cute.runtime.make_fake_compact_tensor( - OUT_STORAGE_DTYPE, - (_fake_batch, sq, qh, d_v), - stride_order=(3, 2, 1, 0), - assumed_align=16, - ) + + def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE): + if stride is None: + return cute.runtime.make_fake_compact_tensor(dtype, shape, stride_order=(3, 2, 1, 0), assumed_align=16) + for axis in (1, 2): # seq/head global strides feed TMA: 16-byte rule + if (stride[axis] * CFG.BPE) % 16 != 0: + raise ValueError(f"declared stride {stride} axis {axis} must be a 16-byte multiple at BPE={CFG.BPE} (TMA global-stride rule)") + return cute.runtime.make_fake_tensor(dtype, shape, tuple(stride), assumed_align=16) + + fake_q = _fake_bshd((_fake_batch, sq, qh, d_qk), q_stride) + fake_k = _fake_bshd((_fake_batch, skv, kh, d_qk), k_stride) + fake_v = _fake_bshd((_fake_batch, skv, kh, d_v), v_stride) + fake_o = _fake_bshd((_fake_batch, sq, qh, d_v), o_stride, dtype=OUT_STORAGE_DTYPE) 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. 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 5d712f91d..e66b99003 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py @@ -1861,7 +1861,7 @@ def _tma_swz(byte_w: int): seq_kv_lens_tensor, cutlass.Int32(QH), cutlass.Int32(B), - cutlass.Int32(o_tensor.shape[3]), + cutlass.Int32(o_tensor.stride[1]), ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream) grid_shape = (n_thd_units * cutlass.Int32(CFG.CGA_M), cutlass.Int32(1), cutlass.Int32(1)) else: @@ -1903,6 +1903,10 @@ def compile( # noqa: A001 has_lse: bool = True, lse_head_major: bool = False, lse_head_stride: int = 0, + q_stride: Optional[tuple] = None, + k_stride: Optional[tuple] = None, + v_stride: Optional[tuple] = None, + o_stride: Optional[tuple] = None, ) -> 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 @@ -1916,30 +1920,19 @@ def compile( # noqa: A001 if (d_qk * CFG.BPE) % 16 != 0 or (d_v * CFG.BPE_O) % 16 != 0: raise ValueError(f"d512 envelope: d_qk*BPE and d_v*BPE must be 16-byte multiples (TMA global-stride rule); got ({d_qk}, {d_v}) at BPE={CFG.BPE}") _fake_batch = 1 if CFG.THD_VARLEN else b - fake_q = cute.runtime.make_fake_compact_tensor( - STORAGE_DTYPE, - (_fake_batch, sq, qh, d_qk), - stride_order=(3, 2, 1, 0), - assumed_align=16, - ) - fake_k = cute.runtime.make_fake_compact_tensor( - STORAGE_DTYPE, - (_fake_batch, skv, kh, d_qk), - stride_order=(3, 2, 1, 0), - assumed_align=16, - ) - fake_v = cute.runtime.make_fake_compact_tensor( - STORAGE_DTYPE, - (_fake_batch, skv, kh, d_v), - stride_order=(3, 2, 1, 0), - assumed_align=16, - ) - fake_o = cute.runtime.make_fake_compact_tensor( - OUT_STORAGE_DTYPE, - (_fake_batch, sq, qh, d_v), - stride_order=(3, 2, 1, 0), - assumed_align=16, - ) + + def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE): + if stride is None: + return cute.runtime.make_fake_compact_tensor(dtype, shape, stride_order=(3, 2, 1, 0), assumed_align=16) + for axis in (1, 2): # seq/head global strides feed TMA: 16-byte rule + if (stride[axis] * CFG.BPE) % 16 != 0: + raise ValueError(f"declared stride {stride} axis {axis} must be a 16-byte multiple at BPE={CFG.BPE} (TMA global-stride rule)") + return cute.runtime.make_fake_tensor(dtype, shape, tuple(stride), assumed_align=16) + + fake_q = _fake_bshd((_fake_batch, sq, qh, d_qk), q_stride) + fake_k = _fake_bshd((_fake_batch, skv, kh, d_qk), k_stride) + fake_v = _fake_bshd((_fake_batch, skv, kh, d_v), v_stride) + fake_o = _fake_bshd((_fake_batch, sq, qh, d_v), o_stride, dtype=OUT_STORAGE_DTYPE) 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. diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py b/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py index cf3e1d3df..9868f15d4 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py @@ -162,17 +162,29 @@ def is_layout_supported( shape: tuple[int, ...], stride: tuple[int, ...], ) -> bool: - """Return whether a BSHD tensor uses compact storage.""" + """Return whether a BSHD tensor uses storage the kernel can address. + + The head dim must be innermost-contiguous, and the head/seq strides + must be 16-byte multiples (x8 elements at 2 bytes/elem — the TMA + interior-stride and vectorized-access granularity) covering the dims + below them (compact or padded; sub-dense would alias). The compact + layout is the special case with equality everywhere; a padded seq + stride is e.g. a THD K/V view of a kv-interleaved [T, 2, H, D] + buffer (token stride 2*h*d). + """ if len(shape) != 4 or len(stride) != 4: return False - _, sequence, heads, head_dim = shape - return stride == ( - sequence * heads * head_dim, - heads * head_dim, - head_dim, - 1, - ) + batch, sequence, heads, head_dim = shape + if stride[3] != 1: + return False + if stride[2] % 8 != 0 or stride[2] < head_dim: + return False + if stride[1] % 8 != 0 or stride[1] < heads * stride[2]: + return False + if batch != 1 and stride[0] < sequence * stride[1]: + return False + return True def __init__( self, @@ -1359,18 +1371,23 @@ def __call__( # check to zero-fill past it, so a rank-4 view keeps the ACTUAL extent # innermost — TMA order (D, S, H, B) — and load_one_kv_tile steps the # head coordinate per swizzle-span chunk. + # K/V TMA layouts read the TENSORS' strides (batch/seq/head), not + # packed recomputations: a THD view may declare a wider token stride + # (e.g. a K/V slice of a kv-interleaved [T, 2, H, D] buffer), and TMA + # encodes it directly (interior strides must be 16-byte multiples; + # check_support declines declarations TMA cannot express). def kv_tma_desc(t, head_dim, head_tile, swizzle, swizzle_chunks, swizzle_chunk_elems): if cutlass.const_expr(head_dim == head_tile): layout = cute.make_layout( (t.shape[0], t.shape[2], swizzle_chunks, t.shape[1], swizzle_chunk_elems), - stride=(t.shape[1] * t.shape[2] * head_dim, head_dim, swizzle_chunk_elems, t.shape[2] * head_dim, 1), + stride=(t.stride[0], t.stride[2], swizzle_chunk_elems, t.stride[1], 1), ) box = (1, 1, swizzle_chunks, self.kv_tile, swizzle_chunk_elems) stride_order = (4, 3, 2, 1, 0) else: layout = cute.make_layout( (t.shape[0], t.shape[2], t.shape[1], head_dim), - stride=(t.shape[1] * t.shape[2] * head_dim, head_dim, t.shape[2] * head_dim, 1), + stride=(t.stride[0], t.stride[2], t.stride[1], 1), ) box = (1, 1, self.kv_tile, swizzle_chunk_elems) stride_order = (3, 2, 1, 0) @@ -1424,6 +1441,10 @@ def compile( # noqa: A001 has_lse: bool = True, lse_head_major: bool = False, lse_head_stride: int = 0, + q_stride: Optional[tuple[int, int, int, int]] = None, + k_stride: Optional[tuple[int, int, int, int]] = None, + v_stride: Optional[tuple[int, int, int, int]] = None, + o_stride: Optional[tuple[int, int, int, int]] = None, ) -> Callable: """Compile and cache one architecture-specific compact BSHD shape. @@ -1462,30 +1483,16 @@ def compile( # noqa: A001 kv_tile=PARAMS.kv_tile, ) fake_batch = 1 if PARAMS.thd_varlen else b - fake_q = cute.runtime.make_fake_compact_tensor( - STORAGE_DTYPE, - (fake_batch, sq, qh, d_qk), - stride_order=(3, 2, 1, 0), - assumed_align=16, - ) - fake_k = cute.runtime.make_fake_compact_tensor( - STORAGE_DTYPE, - (fake_batch, skv, kh, d_qk), - stride_order=(3, 2, 1, 0), - assumed_align=16, - ) - fake_v = cute.runtime.make_fake_compact_tensor( - STORAGE_DTYPE, - (fake_batch, skv, kh, d_v), - stride_order=(3, 2, 1, 0), - assumed_align=16, - ) - fake_o = cute.runtime.make_fake_compact_tensor( - STORAGE_DTYPE, - (fake_batch, sq, qh, d_v), - stride_order=(3, 2, 1, 0), - assumed_align=16, - ) + + def _fake_bshd(shape, stride): + if stride is None: + return cute.runtime.make_fake_compact_tensor(STORAGE_DTYPE, shape, stride_order=(3, 2, 1, 0), assumed_align=16) + return cute.runtime.make_fake_tensor(STORAGE_DTYPE, shape, tuple(stride), assumed_align=16) + + fake_q = _fake_bshd((fake_batch, sq, qh, d_qk), q_stride) + fake_k = _fake_bshd((fake_batch, skv, kh, d_qk), k_stride) + fake_v = _fake_bshd((fake_batch, skv, kh, d_v), v_stride) + fake_o = _fake_bshd((fake_batch, sq, qh, d_v), o_stride) if PARAMS.thd_varlen: fake_lse_shape = (qh, lse_head_stride) if lse_head_major else (sq, qh) else: diff --git a/python/cudnn/sdpa/fwd/kernels/thd_sm100.py b/python/cudnn/sdpa/fwd/kernels/thd_sm100.py index 7b7baa643..6fe9077c0 100644 --- a/python/cudnn/sdpa/fwd/kernels/thd_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/thd_sm100.py @@ -21,7 +21,7 @@ def build_o_descs_kernel( seq_kv_lens_t: cute.Tensor, n_qh: cutlass.Int32, n_batch: cutlass.Int32, - d_v: cutlass.Int32, + o_row_stride: cutlass.Int32, ) -> None: if nvvm.elect_sync(): o_ptr = o_tensor.iterator.raw_ptr() @@ -29,7 +29,10 @@ def build_o_descs_kernel( src_words = Pointer(base_o_desc.get_ptr(), dtype=cutlass.Int64) cu = cutlass.make_array_view(seq_kv_lens_t) cuq0 = n_batch - row_elems = n_qh * d_v + # The DECLARED per-token element stride of packed O (n_qh * d_v when + # compact; wider for e.g. an interleaved-buffer view) — the per-batch + # descriptor bases must step in real rows. + row_elems = o_row_stride for b in cutlass.range(0, n_batch, 1, unroll=1): dptr = desc_base + b * cutlass.Int32(TENSOR_MAP_QWORDS) for i in cutlass.range_constexpr(TENSOR_MAP_QWORDS): From 046f1ea752f7f5113a95f23ae1ab48b770eac385 Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Mon, 10 Aug 2026 11:21:01 -0700 Subject: [PATCH 2/4] =?UTF-8?q?docs(agents):=20Hard=20Rule=202=20=E2=80=94?= =?UTF-8?q?=20serve=20the=20declared=20layout=20natively=20or=20decline,?= =?UTF-8?q?=20never=20adapt?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes the loophole Rule 1's letter leaves open: adapter-side normalization copies that make an unsupported layout runnable. Workspace carving does not legitimize a data-tensor copy (the carve exemption is for metadata and dead-slot dummies), the dense path's grandfathered normalization is not a license for new ones, and whatever check_support accepts the kernel must address natively. Co-Authored-By: Claude Fable 5 --- python/cudnn/AGENTS.md | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/python/cudnn/AGENTS.md b/python/cudnn/AGENTS.md index 156c3816e..1a93961f3 100644 --- a/python/cudnn/AGENTS.md +++ b/python/cudnn/AGENTS.md @@ -39,6 +39,34 @@ Numbered so reviews can cite them; the list grows — append, never renumber. extent would be zero, bind a never-dereferenced dummy view over storage the contract already guarantees. +**Rule 2 — `execute()` launches exactly the kernels the plan promised: +serve the declared layout natively, or decline — never adapt.** + +Rule 1 bans implicit conversions and allocations; this rule bans the loophole +that survives its letter: "helpful" adapter-side work that makes an +unsupported input runnable. + +- **No hidden kernel launches.** A gather/scatter "normalization" copy, a + `.contiguous()`, a layout repack, a scatter-back after the launch — each is + an extra kernel that silently changes the measured perf profile per + configuration. **Carving the copy's scratch from the caller's workspace + does NOT make it acceptable**: Rule 1's workspace-carve exemption covers + metadata buffers and dead-slot dummies, never data-tensor copies. +- **Can't address the declared layout natively? Decline in + `check_support()`** (`NotImplementedError` naming the offending tensor and + its strides) so the Router picks an engine that honors the declaration. + Silent wrong results are the worst failure mode; a silent slow path is the + second worst — both hide behind a green test. See + `_thd_check_strides_native` in `sdpa/fwd/api_dsl.py`. +- **Precedent is not a license.** The SM100 dense path's compact-BSHD + normalization (`dense_layout_ok`: "one gather/scatter copy otherwise") + predates this rule and is grandfathered — do not cite it to justify a new + copy path, and treat migrating it to serve-or-decline as open cleanup. +- The flip side of declining: whatever `check_support()` ACCEPTS, the kernel + must address natively (layout-driven offset math, strides encoded in TMA + descriptors) — acceptance is a promise about the execute path, not about + what the adapter can patch up. + ## Frontend-only kernel package layout ``` From f13de3ca041c90cf9f41240cf894cbd341968baa Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Mon, 10 Aug 2026 14:30:34 -0700 Subject: [PATCH 3/4] =?UTF-8?q?frost(sdpa):=20review=20hardening=20?= =?UTF-8?q?=E2=80=94=20validate=20runtime=20THD=20buffers,=20decline=20ove?= =?UTF-8?q?rlapping=20strides?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - _thd_view validates the runtime buffer against its declaration before reinterpreting storage: dtype/device must match and the base address must be 16-byte aligned (TMA global-address rule / assumed_align=16); as_strided already rejects views past the underlying allocation. - _thd_check_strides_native additionally requires covering (non-overlapping) strides — head >= d, token >= heads*head — matching the SM120 kernel's is_layout_supported, so sub-dense declarations are declined at check_support instead of failing at the per-execute compile (or racing on O writes on SM100). - Kernel _fake_bshd guards: the head dim must be innermost-contiguous; d256/d512 validate the O stride at BPE_O (the O storage dtype byte size). - Clearer SM120 layout-rejection message (the entry validator accepts padded storage now; the text still demanded compact). The THD host-prep stream binding flagged in the same review round is a pre-existing issue (#476) and is split into a separate PR. Addresses CodeRabbit review feedback on #526. Co-Authored-By: Claude Fable 5 --- python/cudnn/sdpa/fwd/api_dsl.py | 30 ++++++++++++++++--- .../fwd/kernels/prefill_d128_f16_sm100.py | 8 +++-- .../kernels/prefill_d192_d128_f16_sm100.py | 8 +++-- .../fwd/kernels/prefill_d256_f16_sm100.py | 10 ++++--- .../fwd/kernels/prefill_d512_f16_sm100.py | 10 ++++--- .../sdpa/fwd/kernels/prefill_f16_sm120.py | 6 +++- 6 files changed, 53 insertions(+), 19 deletions(-) diff --git a/python/cudnn/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py index b58e4e458..2b0cc5efe 100644 --- a/python/cudnn/sdpa/fwd/api_dsl.py +++ b/python/cudnn/sdpa/fwd/api_dsl.py @@ -384,13 +384,19 @@ def _thd_check_strides_native(self) -> None: dim must be innermost-contiguous (elem stride 1) and the token/head strides multiples of 8 elements (which also keeps every per-sequence ragged base 16-byte aligned). Whole-token gaps always qualify for - supported head dims; sub-token gaps only in multiples of 8 elements.""" + supported head dims; sub-token gaps only in multiples of 8 elements. + The strides must also COVER the tensor (head >= d, token >= h*head): + an overlapping declaration would alias distinct O rows onto the same + storage (a write race) and is outside the kernels' addressing + contract.""" for desc in (self.q_desc, self.k_desc, self.v_desc, self.o_desc): (ts, hs, es), _ = self._thd_declared(desc) + h, d = desc.shape[1], desc.shape[3] self._not_implemented_error_if( - es != 1 or ts % 8 != 0 or hs % 8 != 0, + es != 1 or ts % 8 != 0 or hs % 8 != 0 or hs < d or ts < h * hs, f"{desc.name} THD strides {tuple(desc.stride)} are not TMA-expressible " - f"(head dim must be innermost-contiguous and token/head strides 16-byte multiples)", + f"(head dim must be innermost-contiguous, token/head strides 16-byte " + f"multiples, and non-overlapping: head stride >= {d}, token stride >= heads * head stride)", ) def _thd_check_strides_packed(self) -> None: @@ -405,9 +411,25 @@ def _thd_check_strides_packed(self) -> None: ) def _thd_view(self, buf: torch.Tensor, desc: TensorDesc, tokens: int) -> torch.Tensor: - """The declared-stride ``(1, T, H, D)`` view over a THD buffer's storage.""" + """The declared-stride ``(1, T, H, D)`` view over a THD buffer's storage. + + Validates the RUNTIME buffer against the declaration before + reinterpreting its storage: the dtype/device must match what was + declared, and the base address must be 16-byte aligned — the kernels + are compiled with ``assumed_align=16`` and TMA requires it of the + descriptor's global address, so a misaligned slice would fault (or + worse) instead of erroring here. ``as_strided`` itself rejects views + that extend past the underlying storage.""" h, d = desc.shape[1], desc.shape[3] (ts, hs, es), _ = self._thd_declared(desc) + self._value_error_if( + buf.dtype != desc.dtype or buf.device != desc.device, + f"{desc.name}: runtime buffer ({buf.dtype}, {buf.device}) does not match its declaration ({desc.dtype}, {desc.device})", + ) + self._value_error_if( + buf.data_ptr() % 16 != 0, + f"{desc.name}: runtime buffer base address must be 16-byte aligned (TMA global-address rule); got data_ptr() % 16 == {buf.data_ptr() % 16}", + ) return buf.as_strided((1, tokens, h, d), (max(tokens, 1) * ts, ts, hs, es), buf.storage_offset()) def _amax_slot(self, tensor, name: str, device: torch.device) -> torch.Tensor: 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 033cbed07..fd70b3cf3 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py @@ -2082,12 +2082,14 @@ def compile( # noqa: A001 raise ValueError(f"d128 envelope: d_qk*BPE and d_v*BPE must be 16-byte multiples (TMA global-stride rule); got ({d_qk}, {d_v}) at BPE={CFG.BPE}") _fake_batch = 1 if CFG.THD_VARLEN else b - def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE): + def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): if stride is None: return cute.runtime.make_fake_compact_tensor(dtype, shape, stride_order=(3, 2, 1, 0), assumed_align=16) + if stride[3] != 1: + raise ValueError(f"declared stride {stride}: the head dim must be innermost-contiguous (stride[3] == 1)") for axis in (1, 2): # seq/head global strides feed TMA: 16-byte rule - if (stride[axis] * CFG.BPE) % 16 != 0: - raise ValueError(f"declared stride {stride} axis {axis} must be a 16-byte multiple at BPE={CFG.BPE} (TMA global-stride rule)") + if (stride[axis] * bpe) % 16 != 0: + raise ValueError(f"declared stride {stride} axis {axis} must be a 16-byte multiple at BPE={bpe} (TMA global-stride rule)") return cute.runtime.make_fake_tensor(dtype, shape, tuple(stride), assumed_align=16) fake_q = _fake_bshd((_fake_batch, sq, qh, d_qk), q_stride) 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 2dbe49d82..8f2face2b 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 @@ -2169,12 +2169,14 @@ def compile( # noqa: A001 raise ValueError(f"d192 envelope: d_qk*BPE and d_v*BPE must be 16-byte multiples (TMA global-stride rule); got ({d_qk}, {d_v}) at BPE={CFG.BPE}") _fake_batch = 1 if CFG.THD_VARLEN else b - def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE): + def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): if stride is None: return cute.runtime.make_fake_compact_tensor(dtype, shape, stride_order=(3, 2, 1, 0), assumed_align=16) + if stride[3] != 1: + raise ValueError(f"declared stride {stride}: the head dim must be innermost-contiguous (stride[3] == 1)") for axis in (1, 2): # seq/head global strides feed TMA: 16-byte rule - if (stride[axis] * CFG.BPE) % 16 != 0: - raise ValueError(f"declared stride {stride} axis {axis} must be a 16-byte multiple at BPE={CFG.BPE} (TMA global-stride rule)") + if (stride[axis] * bpe) % 16 != 0: + raise ValueError(f"declared stride {stride} axis {axis} must be a 16-byte multiple at BPE={bpe} (TMA global-stride rule)") return cute.runtime.make_fake_tensor(dtype, shape, tuple(stride), assumed_align=16) fake_q = _fake_bshd((_fake_batch, sq, qh, d_qk), q_stride) 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 c808d9958..36a4295e5 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py @@ -1719,18 +1719,20 @@ def compile( # noqa: A001 raise ValueError(f"d256 envelope: d_qk*BPE and d_v*BPE must be 16-byte multiples (TMA global-stride rule); got ({d_qk}, {d_v}) at BPE={CFG.BPE}") _fake_batch = 1 if CFG.THD_VARLEN else b - def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE): + def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): if stride is None: return cute.runtime.make_fake_compact_tensor(dtype, shape, stride_order=(3, 2, 1, 0), assumed_align=16) + if stride[3] != 1: + raise ValueError(f"declared stride {stride}: the head dim must be innermost-contiguous (stride[3] == 1)") for axis in (1, 2): # seq/head global strides feed TMA: 16-byte rule - if (stride[axis] * CFG.BPE) % 16 != 0: - raise ValueError(f"declared stride {stride} axis {axis} must be a 16-byte multiple at BPE={CFG.BPE} (TMA global-stride rule)") + if (stride[axis] * bpe) % 16 != 0: + raise ValueError(f"declared stride {stride} axis {axis} must be a 16-byte multiple at BPE={bpe} (TMA global-stride rule)") return cute.runtime.make_fake_tensor(dtype, shape, tuple(stride), assumed_align=16) fake_q = _fake_bshd((_fake_batch, sq, qh, d_qk), q_stride) fake_k = _fake_bshd((_fake_batch, skv, kh, d_qk), k_stride) fake_v = _fake_bshd((_fake_batch, skv, kh, d_v), v_stride) - fake_o = _fake_bshd((_fake_batch, sq, qh, d_v), o_stride, dtype=OUT_STORAGE_DTYPE) + fake_o = _fake_bshd((_fake_batch, sq, qh, d_v), o_stride, dtype=OUT_STORAGE_DTYPE, bpe=CFG.BPE_O) 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. 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 e66b99003..374eb9bf2 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py @@ -1921,18 +1921,20 @@ def compile( # noqa: A001 raise ValueError(f"d512 envelope: d_qk*BPE and d_v*BPE must be 16-byte multiples (TMA global-stride rule); got ({d_qk}, {d_v}) at BPE={CFG.BPE}") _fake_batch = 1 if CFG.THD_VARLEN else b - def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE): + def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): if stride is None: return cute.runtime.make_fake_compact_tensor(dtype, shape, stride_order=(3, 2, 1, 0), assumed_align=16) + if stride[3] != 1: + raise ValueError(f"declared stride {stride}: the head dim must be innermost-contiguous (stride[3] == 1)") for axis in (1, 2): # seq/head global strides feed TMA: 16-byte rule - if (stride[axis] * CFG.BPE) % 16 != 0: - raise ValueError(f"declared stride {stride} axis {axis} must be a 16-byte multiple at BPE={CFG.BPE} (TMA global-stride rule)") + if (stride[axis] * bpe) % 16 != 0: + raise ValueError(f"declared stride {stride} axis {axis} must be a 16-byte multiple at BPE={bpe} (TMA global-stride rule)") return cute.runtime.make_fake_tensor(dtype, shape, tuple(stride), assumed_align=16) fake_q = _fake_bshd((_fake_batch, sq, qh, d_qk), q_stride) fake_k = _fake_bshd((_fake_batch, skv, kh, d_qk), k_stride) fake_v = _fake_bshd((_fake_batch, skv, kh, d_v), v_stride) - fake_o = _fake_bshd((_fake_batch, sq, qh, d_v), o_stride, dtype=OUT_STORAGE_DTYPE) + fake_o = _fake_bshd((_fake_batch, sq, qh, d_v), o_stride, dtype=OUT_STORAGE_DTYPE, bpe=CFG.BPE_O) 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. diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py b/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py index 9868f15d4..72d9f3d34 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py @@ -1330,7 +1330,11 @@ def __call__( raise ValueError("runtime Q/K/V/O batch, sequence, or head geometry mismatch") for name, tensor in (("Q", q), ("K", k), ("V", v), ("O", o)): if cutlass.const_expr(not self.is_layout_supported(tensor.shape, tensor.stride)): - raise ValueError(f"{name} must use compact BSHD storage") + raise ValueError( + f"{name} layout is not supported: BSHD with the head dim innermost-contiguous " + f"and non-overlapping seq/head strides that are multiples of 8 elements " + f"(compact or padded); got shape {tuple(tensor.shape)} stride {tuple(tensor.stride)}" + ) if cutlass.const_expr(lse is not None): if cutlass.const_expr(self.thd_varlen): if cutlass.const_expr(self.thd_lse_head_major): From 1c563ca475fd1f908a5c7e63356363b73c51fa19 Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Mon, 10 Aug 2026 22:58:45 -0700 Subject: [PATCH 4/4] frost(sdpa): make the THD native-stride gate dtype-aware (16 // itemsize) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The gate hardcoded the TMA 16-byte global-stride rule as 8 elements, the f16/bf16 case. It lives in the shared base class, so express the quantum in the tensor's own element units — 8 at 2 B/elem, 16 at 1 B/elem (fp8), 4 at 4 B/elem — per descriptor, so mixed-precision declarations check each tensor at its own dtype. No behavior change for the f16 paths this PR enables; the fp8 native-stride follow-up (#537) inherits the correct quantum for free. Suggested by @Aneureka in review. Co-Authored-By: Claude Fable 5 --- python/cudnn/sdpa/fwd/api_dsl.py | 30 +++++++++++++++++------------- 1 file changed, 17 insertions(+), 13 deletions(-) diff --git a/python/cudnn/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py index 2b0cc5efe..7712dcb6b 100644 --- a/python/cudnn/sdpa/fwd/api_dsl.py +++ b/python/cudnn/sdpa/fwd/api_dsl.py @@ -379,24 +379,28 @@ def _thd_declared(desc: TensorDesc): return st, st == (h * d, d, 1) def _thd_check_strides_native(self) -> None: - """Reject THD stride declarations the f16 kernels cannot address - natively: TMA's 16-byte global-stride rule at 2 bytes/elem — the head - dim must be innermost-contiguous (elem stride 1) and the token/head - strides multiples of 8 elements (which also keeps every per-sequence - ragged base 16-byte aligned). Whole-token gaps always qualify for - supported head dims; sub-token gaps only in multiples of 8 elements. - The strides must also COVER the tensor (head >= d, token >= h*head): - an overlapping declaration would alias distinct O rows onto the same - storage (a write race) and is outside the kernels' addressing - contract.""" + """Reject THD stride declarations the kernels cannot address + natively: TMA's 16-byte global-stride rule — the head dim must be + innermost-contiguous (elem stride 1) and the token/head strides + multiples of ``16 // itemsize`` elements (which also keeps every + per-sequence ragged base 16-byte aligned). Whole-token gaps always + qualify for supported head dims; sub-token gaps only in 16-byte + multiples. The strides must also COVER the tensor (head >= d, + token >= h*head): an overlapping declaration would alias distinct O + rows onto the same storage (a write race) and is outside the + kernels' addressing contract.""" for desc in (self.q_desc, self.k_desc, self.v_desc, self.o_desc): (ts, hs, es), _ = self._thd_declared(desc) h, d = desc.shape[1], desc.shape[3] + # The 16-byte TMA rule in this tensor's OWN element units: 8 at + # 2 B/elem (f16/bf16), 16 at 1 B/elem (fp8), 4 at 4 B/elem. + quantum = 16 // desc.dtype.itemsize self._not_implemented_error_if( - es != 1 or ts % 8 != 0 or hs % 8 != 0 or hs < d or ts < h * hs, + es != 1 or ts % quantum != 0 or hs % quantum != 0 or hs < d or ts < h * hs, f"{desc.name} THD strides {tuple(desc.stride)} are not TMA-expressible " - f"(head dim must be innermost-contiguous, token/head strides 16-byte " - f"multiples, and non-overlapping: head stride >= {d}, token stride >= heads * head stride)", + f"(head dim must be innermost-contiguous, token/head strides 16-byte — " + f"{quantum}-element — multiples, and non-overlapping: head stride >= {d}, " + f"token stride >= heads * head stride)", ) def _thd_check_strides_packed(self) -> None: