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 ``` diff --git a/python/cudnn/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py index 3d1763404..7712dcb6b 100644 --- a/python/cudnn/sdpa/fwd/api_dsl.py +++ b/python/cudnn/sdpa/fwd/api_dsl.py @@ -360,6 +360,82 @@ 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 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 % 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"{quantum}-element — multiples, and non-overlapping: head stride >= {d}, " + f"token stride >= heads * head stride)", + ) + + 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. + + 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: """The caller's 1-element amax storage, or a cached dummy. @@ -570,6 +646,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 +1196,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 +1218,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 +1247,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 +1740,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 +2230,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 +2275,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 +2298,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 +2324,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 +2362,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..fd70b3cf3 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,21 @@ 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, 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] * 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=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..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 @@ -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,21 @@ 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, 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] * 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=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..36a4295e5 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,21 @@ 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, 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] * 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, 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 5d712f91d..374eb9bf2 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,21 @@ 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, 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] * 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, 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 cf3e1d3df..72d9f3d34 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, @@ -1318,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): @@ -1359,18 +1375,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 +1445,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 +1487,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):