diff --git a/python/cudnn/AGENTS.md b/python/cudnn/AGENTS.md index d5f125e22..82573f1a1 100644 --- a/python/cudnn/AGENTS.md +++ b/python/cudnn/AGENTS.md @@ -97,14 +97,21 @@ Known violations, all pre-existing and each needing a kernel-side change, so none is precedent: - THD `cu_seqlens` host cumsum (`sdpa/fwd/api_dsl.py`, `_execute_thd` on both - SM100 and SM120). The compile-side half is DONE — the THD kernels compile - with dynamic token extents (Rule 4), so `T` is no longer a compile-time - constant and no compile is keyed on it. `t_q`/`t_kv` still reach the host - for the metadata upload, the ragged views' extents, and the launch - grid; removing that needs the plan-time-max (`b * s_q_max`) grid with - in-kernel dead-tile exit and the device-side metadata read (issue #552). - `sdpa_fwd_wrapper_sm80` shows the other half — it requires `max_s_q` from - the caller rather than deriving it. + SM100 and SM120). RESOLVED on all SM100 f16 families + (issue #552) — the reference for the SM120 port: the kernels compile + with dynamic token extents (Rule 4), the setup kernel builds the metadata + buffer device-side from the caller's length tensors, every ragged view + binds its buffer's capacity, and the grid is the plan-time envelope + (`b * ceil(s_q_declared / CGA_TILE_M) * qh`) with dead units exiting via + `_thd_decode`'s `batch == n_batch` sentinel — zero host reads, pinned by + the sync-debug and CUDA-graph capture tests. Still open: SM120 (exact + host-computed grid + host-built metadata; its own engine class and grid + mechanism, so its own port), and + the envelope's dead-tile tax under far-oversized declarations (measured + ~145 ns/dead unit exposed, <0.2% at realistic declarations; a capped + persistent grid reading a device-side live-unit count would bound it by + resident clusters). `sdpa_fwd_wrapper_sm80` shows the caller-provided + alternative — it requires `max_s_q` rather than deriving it. - Per-tensor FP8 descale readback (`_scalar` in the same file): fold on device, passing the pointers, as the backend FP8 sdpa does. - The FP8/MXFP8 `seq_len_q` guard in `sdpa/fwd/engines.py`. This one cannot be diff --git a/python/cudnn/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py index 6b672d79e..fd9bbc1f5 100644 --- a/python/cudnn/sdpa/fwd/api_dsl.py +++ b/python/cudnn/sdpa/fwd/api_dsl.py @@ -1236,6 +1236,24 @@ def _key(desc): o_stride=_key(self.o_desc), ) + def _thd_unit_envelope(self) -> int: + """PLAN-TIME upper bound on live THD units: + ``B * ceil(S_q_declared / CGA_TILE_M) * QH``. + + Every sequence's length is bounded by the declared S_q (the padding + contract), so this covers ``Σ_b ceil(s_b / tile) * QH``. Units past + the live total are DEAD by kernel contract (the decode's + ``batch == n_batch`` sentinel): no loads, no O/LSE writes, one + empty-mainloop barrier dance each. The grid is host-known at PLAN + time — execute reads nothing from the lengths — which removes the + last THD D2H sync and unblocks CUDA-graph capture (issue #552). The + dead-tile tax mirrors the C++ backend's THD grid strategy; callers + declaring S_q far above their live totals pay it. The over-launch + test pads this to pin the dead-unit contract.""" + cga_tile_m = int(self._k_mod.CGA_TILE_M) + s_q_decl = int(self.q_desc.shape[2]) + return self.batch_size * ((s_q_decl + cga_tile_m - 1) // cga_tile_m) * self.h_q + 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. @@ -1248,54 +1266,65 @@ def _execute_thd(self, q_buf, k_buf, v_buf, o_buf, scale_softmax_log2, sinks, se 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) - feeds the metadata upload, the ragged views' extents and the exact - launch grid, and costs one D2H sync per length tensor, no device - allocation. It no longer keys any compile: the kernels compile with - DYNAMIC token extents (issue #552), so a new packed total re-binds - the same artifact; removing the sync itself needs the plan-time-max - grid redesign tracked there.""" + Host round-trips (issue #552): NONE — the lengths never reach the + host (the setup kernel builds the metadata buffer device-side), + every ragged view binds its buffer's capacity, and the launch grid + is the plan-time envelope (dead units exit by kernel contract) — + the execute is fully async and CUDA-graph capturable. No compile is + keyed on runtime data: the kernels compile with DYNAMIC token + extents, so a new packed total re-binds the same artifact.""" import cutlass dev = q_buf.device b = self.batch_size carver = WorkspaceCarver(workspace, self.scratch_workspace_bytes(), "SdpaFwdDslSm100 (THD)") if workspace is not None else None # Metadata buffer: [ seq_kv_lens(B) | cu_seqlens_q(B+1) | cu_seqlens_k(B+1) ], - # 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. Either length - # form feeds it — per-batch (B,) lengths or the (B+1,) cu_seq_len - # prefix sums — at identical cost. The torch work (allocation, D2H - # length reads, the H2D upload) runs on the LAUNCH stream so it is - # ordered against the kernel that consumes it — the execute-time - # handle may carry a stream that is not torch's current. + # built DEVICE-side by the setup kernel from the caller's length + # tensors (either form) — no host cumsum, no H2D, and the lengths + # NEVER reach the host (issue #552); the prefix-sum invariants are + # caller contract (a validation that needs a device read is not a + # validation — AGENTS.md Rule 3). The torch allocations run on the + # LAUNCH stream so they are ordered against the kernels that consume + # them — the execute-time handle may carry a stream that is not + # torch's current. with _torch_stream_context(current_stream, dev): meta = carver.take(3 * b + 2, torch.int32) if carver is not None else torch.empty(3 * b + 2, dtype=torch.int32, device=dev) - slq_host, cu_q_host = self._thd_host_lens(seq_q_lens, "cu_seq_len_q" if self.cu_seq_q_lens else "seq_q_lens", self.cu_seq_q_lens) - slk_host, cu_k_host = self._thd_host_lens(seq_len_kv, "cu_seq_len_kv" if self.cu_seq_kv_lens else "seq_kv_lens", self.cu_seq_kv_lens) - 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] + q_lens_dev = self._checked_cu_seq_lens(seq_q_lens, "cu_seq_len_q") if self.cu_seq_q_lens else self._checked_seq_lens(seq_q_lens, "seq_q_lens") + kv_lens_dev = self._checked_cu_seq_lens(seq_len_kv, "cu_seq_len_kv") if self.cu_seq_kv_lens else self._checked_seq_lens(seq_len_kv, "seq_kv_lens") qh, kh = self.h_q, self.h_kv d_qk, d_v = self.head_dim_qk, self.head_dim_v - # 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.) + # Q/O token extent = buffer CAPACITY (issue #552: the packed Q total + # lives on device only). Q, O and a token-major (or compact + # head-major) LSE bind ONE dynamic token symbol, so take the shared + # floor. Writes never step past the real per-sequence lengths the + # kernel reads from the device metadata (O through the per-batch + # descriptors' extents, LSE via the row predicate), so the + # over-claim only widens the TMA descriptors' bound. + (q_ts, _, _), _ = self._thd_declared(self.q_desc) + (o_ts, _, _), _ = self._thd_declared(self.o_desc) + t_q = min(q_buf.numel() // q_ts, o_buf.numel() // o_ts) + if lse_tensor is not None and not (self.thd_stats_head_major and self.thd_stats_head_stride): + t_q = min(t_q, lse_tensor.numel() // qh) + + # Degenerate Q side: zero CAPACITY — no token is addressable, and a + # zero-token view cannot back a TMA descriptor — nothing to compute + # or write. All-zero LENGTHS over live storage launch normally: + # every unit decodes dead and neither O nor LSE is touched. + # (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") + self._logger.debug("execute (THD): no addressable Q token, nothing to do") return lse = None if lse_tensor is not None: if self.thd_stats_head_major: + # head_stride covering the packed total is caller contract + # (t_q is a device value — Rule 3). 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})", - ) + if head_stride == 0: + head_stride = t_q # compact: the token extent itself 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) @@ -1311,9 +1340,8 @@ def _execute_thd(self, q_buf, k_buf, v_buf, o_buf, scale_softmax_log2, sinks, se # vs the builder pass). with _torch_stream_context(current_stream, dev): o_desc = carver.take(b * 16 + 16, torch.int64) if carver is not None else torch.empty(b * 16 + 16, dtype=torch.int64, device=dev) - # One THD unit per CGA-height slice of each sequence's Q rows. - 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) + # The PLAN-TIME envelope grid — dead units exit by kernel contract. + units = self._thd_unit_envelope() # 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 @@ -1321,15 +1349,26 @@ def _execute_thd(self, q_buf, k_buf, v_buf, o_buf, scale_softmax_log2, sinks, se # 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) + # KV token extent = buffer CAPACITY (issue #552: the packed KV total + # lives on device only). K and V bind the same dynamic token symbol, + # so take the shared floor. Loads never step past the REAL + # per-sequence lengths the kernel reads from the device metadata, so + # an over-claimed extent only widens the TMA descriptors' bound. + # All-zero lengths with a live buffer launch normally through the + # kernel's per-sequence dead-row path. + (k_ts, _, _), _ = self._thd_declared(self.k_desc) + (v_ts, _, _), _ = self._thd_declared(self.v_desc) + t_kv = min(k_buf.numel() // k_ts, v_buf.numel() // v_ts) 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 == + # No KV storage at all: + # every query row is dead — 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). @@ -1364,7 +1403,29 @@ def _execute_thd(self, q_buf, k_buf, v_buf, o_buf, scale_softmax_log2, sinks, se kwargs = self._thd_compile_kwargs() kwargs.update(k_stride=(0, *K.stride()[1:]), v_stride=(0, *V.stride()[1:])) fn = self._k_mod.compile(**kwargs) - 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) + # The caller's length tensors ride to the setup kernel, which builds + # the metadata buffer device-side; the form bitmask is a runtime + # value (no compile key grows). problem_size sq/skv slots are 0 by + # the THD contract (_host reads the dynamic extents). + lens_form = (1 if self.cu_seq_q_lens else 0) | (2 if self.cu_seq_kv_lens else 0) + fn( + Q, + K, + V, + O, + LSE, + sinks_t, + meta, + o_desc, + (b, qh, kh, 0, 0, 0), + cutlass.Float32(scale_softmax_log2), + cutlass.Int32(units), + None, + q_lens_dev, + kv_lens_dev, + cutlass.Int32(lens_form), + stream=current_stream, + ) self._logger.debug("execute (THD) completed") @staticmethod @@ -1372,14 +1433,20 @@ def _ceil_div(x: int, a: int) -> int: return (x + a - 1) // a def _reshape_sf(self, sf: torch.Tensor, h: int, n_tiles: int, sf_smem_size: int) -> torch.Tensor: - """cuDNN F8_128x4 scale-factor tensor (FP8_E8M0, ``[B, H, *, *]``) → the - kernel's per-tile int8 view ``[B, H, n_tiles, sf_smem_size]``. + """cuDNN F8_128x4 scale-factor tensor (FP8_E8M0) → the kernel's per-tile + int8 view ``[B, H, n_tiles, sf_smem_size]``. cuDNN packs the 128×4 SF atom contiguously (``F8_128x4`` reordering); a Q/K tile is 128 rows × d/32 d-blocks and a V tile is 128 rows × 4 s-blocks, so each tile is exactly ``sf_smem_size`` E8M0 bytes and this is a pure reshape. + + A reordered tensor is an opaque byte layout: callers legally bind it under + any shape with the right byte count (the graph declares logical + ``[B, H, s_padded, d/32]`` dims; TE-style producers hand over flat + ``[B·H·s_padded, 4]`` swizzle output). So B comes from the graph facts, + never from ``sf.shape[0]``, and only the total size is validated. """ - b = sf.shape[0] + b = self.batch_size flat = sf.contiguous() if flat.dtype != torch.int8: flat = flat.view(torch.int8) diff --git a/python/cudnn/sdpa/fwd/kernels/_common_sm100.py b/python/cudnn/sdpa/fwd/kernels/_common_sm100.py index c04f2bb6e..41f38dcf9 100644 --- a/python/cudnn/sdpa/fwd/kernels/_common_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/_common_sm100.py @@ -449,7 +449,14 @@ def _thd_decode(linear_cta, seq_kv_lens_t, n_batch, n_qh, cta_in_pair): cu = cutlass.make_array_view(seq_kv_lens_t) cuq0 = n_batch acc = cutlass.Int32(0) - f_batch = cutlass.Int32(0) + # DEAD-unit sentinel (issue #552 over-launch): a unit no sequence + # claims (u >= sum of live units) keeps batch == n_batch. That index + # makes every downstream consumer a no-op through IN-BOUNDS metadata + # reads: _resolve_seqlen_kv reads meta[n_batch] = cu_q[0] = 0 (empty + # KV range in every role), the epilogue's per-sequence Q length + # cu[2n+1]-cu[2n] goes negative (LSE predicate never fires), and the + # O-store role skips the TMA store explicitly (batch >= n_batch). + f_batch = n_batch f_head = cutlass.Int32(0) f_qc = cutlass.Int32(0) done = cutlass.Int32(0) 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 3e7d31baa..6c8976a42 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py @@ -189,9 +189,12 @@ # Supported at cga1 and cga2 (TILES_Q=2 → two Q slabs / O stores per tile). # seq_kv_lens overloaded as the THD metadata buffer (int32 len 3B+2): # [0..B-1]=seq_kv_lens [B..2B]=cu_q(B+1) [2B+1..3B+1]=cu_k(B+1) -from cudnn.sdpa.fwd.kernels.thd_sm100 import build_o_descs_kernel as _build_o_descs_kernel, TENSOR_MAP_QWORDS +from cudnn.sdpa.fwd.kernels.thd_sm100 import build_thd_meta_o_descs_kernel as _build_thd_meta_o_descs_kernel, TENSOR_MAP_QWORDS _TENSOR_MAP_QWORDS = TENSOR_MAP_QWORDS +# The setup kernel builds the THD metadata buffer DEVICE-side from the +# caller's length tensors and the adapter launches the plan-time envelope +# grid (issue #552) — no length ever reaches the host. _dispatch_decode_initial = _sdpa_h.dispatch_decode_initial _dispatch_decode_payload = _sdpa_h.dispatch_decode_payload _thd_tma_offsets = _sdpa_h.thd_tma_offsets @@ -834,9 +837,13 @@ def _tmastg_warp_group( # (base at the sequence's packed row, seq extent = S_q_b → a box # past S_q_b is OOB-clipped). q_row coord is sequence-local; the # batch coord collapses to 0. Both slabs share one descriptor. - o_desc_ptr = (o_desc_words.iterator.raw_ptr() + batch_idx * cutlass.Int32(_TENSOR_MAP_QWORDS)).tospace(cutlass.AddressSpace.generic) - o_slice = tma_slice_runtime_desc(o_desc_ptr, cutlass.Int32(0), head_idx, q_row_base + cutlass.Int32(qs * CFG.TILE_M), cutlass.Int32(0)) - tma_store_tile(sO[qs], o_slice) + # DEAD unit (batch == n_batch, over-launched grid — issue #552): + # no O rows exist and descriptor slot n_batch is never built, so + # skip the store; the barrier protocol below still runs. + if batch_idx < n_batch: + o_desc_ptr = (o_desc_words.iterator.raw_ptr() + batch_idx * cutlass.Int32(_TENSOR_MAP_QWORDS)).tospace(cutlass.AddressSpace.generic) + o_slice = tma_slice_runtime_desc(o_desc_ptr, cutlass.Int32(0), head_idx, q_row_base + cutlass.Int32(qs * CFG.TILE_M), cutlass.Int32(0)) + tma_store_tile(sO[qs], o_slice) else: tma_store_tile( sO[qs], @@ -1935,6 +1942,15 @@ def _host( # Dense padded-Q trim: separate (B,)-int32 per-batch Q lengths; None # (and absent from the compiled ABI) unless CFG.SEQ_Q_LENS_PRESENT. seq_q_lens_tensor: Optional[cute.Tensor] = None, + # THD device metadata build (issue #552): the CALLER's Q/KV length + # tensors — (B,) per-batch lengths or (B+1,) cu prefix sums, per side via + # thd_lens_form (bit 0: Q is cu, bit 1: KV is cu) — consumed only by the + # setup kernel, which writes the [kv|cu_q|cu_k] metadata buffer + # (seq_kv_lens_tensor) device-side. None (folded out of the ABI) for + # dense graphs. + thd_q_lens_tensor: Optional[cute.Tensor] = None, + thd_kv_lens_tensor: Optional[cute.Tensor] = None, + thd_lens_form: Optional[cutlass.Int32] = None, stream: _cuda_driver.CUstream = None, ) -> None: B, QH, KH, SQ, SKV, _ = problem_size @@ -1998,18 +2014,24 @@ def _tma_swz(byte_w: int): grid_q_supers = q_clusters * CFG.CTA_MMA q_supers = grid_q_supers if cutlass.const_expr(CFG.THD_VARLEN): - # THD: build the per-batch O descriptor array (reuse tma_o_desc over the - # packed [1,T,QH,D_v] O as base), then launch the exact flat - # batch-outermost grid (n_thd_units = Σ_b ceil(S_q_b/CGA_TILE_M)*QH, - # host-computed); grid_x = n_thd_units * CGA_M. Works at cga1 (CGA_M=1). + # THD setup launch: build the [kv|cu_q|cu_k] metadata buffer + # DEVICE-side from the caller's length tensors (no host cumsum, no + # H2D — issue #552), then the per-batch O descriptor array (reuse + # tma_o_desc over the packed [1,T,QH,D_v] O as base). Main grid: the + # exact flat batch-outermost grid (n_thd_units = + # Σ_b ceil(S_q_b/CGA_TILE_M)*QH, host-computed); grid_x = + # n_thd_units * CGA_M. Works at cga1 (CGA_M=1). # ENVELOPE: the packed-O row stride is QH * ACTUAL d_v (o_tensor's # static inner extent), not QH * TILE_O — the per-batch descriptor # bases must step in real rows or every batch >= 1 lands OOB. - _build_o_descs_kernel( + _build_thd_meta_o_descs_kernel( o_tensor, tma_o_desc, o_desc_words, seq_kv_lens_tensor, + thd_q_lens_tensor, + thd_kv_lens_tensor, + thd_lens_form, cutlass.Int32(QH), cutlass.Int32(B), cutlass.Int32(o_tensor.stride[1]), @@ -2201,6 +2223,19 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): stride_order=(0,), assumed_align=16, ) + # THD: the caller's Q/KV length tensors, consumed by the setup kernel's + # device-side metadata build. DYNAMIC extents — (B,) per-batch lengths and + # (B+1,) cu prefix sums bind the same artifact; the form rides the runtime + # thd_lens_form bitmask, so no compile key grows (Rule 4). align 4: bound + # directly, only natural int32 alignment is guaranteed. + if CFG.THD_VARLEN: + fake_thd_q_lens = cute.runtime.make_fake_compact_tensor(cutlass.Int32, (cute.sym_int(divisibility=1),), stride_order=(0,), assumed_align=4) + fake_thd_kv_lens = cute.runtime.make_fake_compact_tensor(cutlass.Int32, (cute.sym_int(divisibility=1),), stride_order=(0,), assumed_align=4) + fake_thd_lens_form = cutlass.Int32(0) + else: + fake_thd_q_lens = None + fake_thd_kv_lens = None + fake_thd_lens_form = None return cute.compile( _host, fake_q, @@ -2217,6 +2252,9 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): cutlass.Float32(0.0), cutlass.Int32(0), fake_seq_q_lens, + fake_thd_q_lens, + fake_thd_kv_lens, + fake_thd_lens_form, stream=cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=False), options="--enable-tvm-ffi", ) 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 c5e5cb659..7040ceaf8 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 @@ -200,9 +200,12 @@ def _bounds_for_tile( # Supported at cga1 and cga2 (TILES_Q=2 → two Q slabs / O stores per tile). # seq_kv_lens overloaded as the THD metadata buffer (int32 len 3B+2): # [0..B-1]=seq_kv_lens [B..2B]=cu_q(B+1) [2B+1..3B+1]=cu_k(B+1) -from cudnn.sdpa.fwd.kernels.thd_sm100 import build_o_descs_kernel as _build_o_descs_kernel, TENSOR_MAP_QWORDS +from cudnn.sdpa.fwd.kernels.thd_sm100 import build_thd_meta_o_descs_kernel as _build_thd_meta_o_descs_kernel, TENSOR_MAP_QWORDS _TENSOR_MAP_QWORDS = TENSOR_MAP_QWORDS +# The setup kernel builds the THD metadata buffer DEVICE-side from the +# caller's length tensors and the adapter launches the plan-time envelope +# grid (issue #552) — no length ever reaches the host. _dispatch_decode_initial = _sdpa_h.dispatch_decode_initial _dispatch_decode_payload = _sdpa_h.dispatch_decode_payload _thd_tma_offsets = _sdpa_h.thd_tma_offsets @@ -907,9 +910,13 @@ def _tmastg_warp_group( # (base at the sequence's packed row, seq extent = S_q_b → a box # past S_q_b is OOB-clipped). q_row coord is sequence-local; the # batch coord collapses to 0. Both slabs share one descriptor. - o_desc_ptr = (o_desc_words.iterator.raw_ptr() + batch_idx * cutlass.Int32(_TENSOR_MAP_QWORDS)).tospace(cutlass.AddressSpace.generic) - o_slice = tma_slice_runtime_desc(o_desc_ptr, cutlass.Int32(0), head_idx, q_row_base + cutlass.Int32(qs * CFG.TILE_M), cutlass.Int32(0)) - tma_store_tile(sO[qs], o_slice) + # DEAD unit (batch == n_batch, envelope grid — issue #552): no O + # rows exist and descriptor slot n_batch is never built, so skip + # the store; the barrier protocol below still runs. + if batch_idx < n_batch: + o_desc_ptr = (o_desc_words.iterator.raw_ptr() + batch_idx * cutlass.Int32(_TENSOR_MAP_QWORDS)).tospace(cutlass.AddressSpace.generic) + o_slice = tma_slice_runtime_desc(o_desc_ptr, cutlass.Int32(0), head_idx, q_row_base + cutlass.Int32(qs * CFG.TILE_M), cutlass.Int32(0)) + tma_store_tile(sO[qs], o_slice) else: tma_store_tile( sO[qs], @@ -2056,6 +2063,15 @@ def _host( # Dense padded-Q trim: separate (B,)-int32 per-batch Q lengths; None # (and absent from the compiled ABI) unless CFG.SEQ_Q_LENS_PRESENT. seq_q_lens_tensor: Optional[cute.Tensor] = None, + # THD device metadata build (issue #552): the CALLER's Q/KV length + # tensors — (B,) per-batch lengths or (B+1,) cu prefix sums, per side via + # thd_lens_form (bit 0: Q is cu, bit 1: KV is cu) — consumed only by the + # setup kernel, which writes the [kv|cu_q|cu_k] metadata buffer + # (seq_kv_lens_tensor) device-side. None (folded out of the ABI) for + # dense graphs. + thd_q_lens_tensor: Optional[cute.Tensor] = None, + thd_kv_lens_tensor: Optional[cute.Tensor] = None, + thd_lens_form: Optional[cutlass.Int32] = None, stream: _cuda_driver.CUstream = None, ) -> None: B, QH, KH, SQ, SKV, _ = problem_size @@ -2126,11 +2142,14 @@ def _tma_swz(byte_w: int): # ENVELOPE: the packed-O row stride is QH * ACTUAL d_v (o_tensor's # static inner extent), not QH * TILE_O — the per-batch descriptor # bases must step in real rows or every batch >= 1 lands OOB. - _build_o_descs_kernel( + _build_thd_meta_o_descs_kernel( o_tensor, tma_o_desc, o_desc_words, seq_kv_lens_tensor, + thd_q_lens_tensor, + thd_kv_lens_tensor, + thd_lens_form, cutlass.Int32(QH), cutlass.Int32(B), cutlass.Int32(o_tensor.stride[1]), @@ -2314,6 +2333,19 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): stride_order=(0,), assumed_align=16, ) + # THD: the caller's Q/KV length tensors, consumed by the setup kernel's + # device-side metadata build. DYNAMIC extents — (B,) per-batch lengths and + # (B+1,) cu prefix sums bind the same artifact; the form rides the runtime + # thd_lens_form bitmask, so no compile key grows (Rule 4). align 4: bound + # directly, only natural int32 alignment is guaranteed. + if CFG.THD_VARLEN: + fake_thd_q_lens = cute.runtime.make_fake_compact_tensor(cutlass.Int32, (cute.sym_int(divisibility=1),), stride_order=(0,), assumed_align=4) + fake_thd_kv_lens = cute.runtime.make_fake_compact_tensor(cutlass.Int32, (cute.sym_int(divisibility=1),), stride_order=(0,), assumed_align=4) + fake_thd_lens_form = cutlass.Int32(0) + else: + fake_thd_q_lens = None + fake_thd_kv_lens = None + fake_thd_lens_form = None return cute.compile( _host, fake_q, @@ -2330,6 +2362,9 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): cutlass.Float32(0.0), cutlass.Int32(0), fake_seq_q_lens, + fake_thd_q_lens, + fake_thd_kv_lens, + fake_thd_lens_form, stream=cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=False), # Keep -uumn so the P exp2/store sched_res_busy_xu64 region affects # backend scheduling. 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 7eef3b540..4365ac423 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py @@ -120,9 +120,12 @@ _resolve_seqlen_q = _sdpa_h.resolve_seqlen_q -from cudnn.sdpa.fwd.kernels.thd_sm100 import build_o_descs_kernel as _build_o_descs_kernel, TENSOR_MAP_QWORDS +from cudnn.sdpa.fwd.kernels.thd_sm100 import build_thd_meta_o_descs_kernel as _build_thd_meta_o_descs_kernel, TENSOR_MAP_QWORDS _TENSOR_MAP_QWORDS = TENSOR_MAP_QWORDS +# The setup kernel builds the THD metadata buffer DEVICE-side from the +# caller's length tensors and the adapter launches the plan-time envelope +# grid (issue #552) — no length ever reaches the host. _dispatch_decode_initial = _sdpa_h.dispatch_decode_initial _dispatch_decode_payload = _sdpa_h.dispatch_decode_payload _thd_tma_offsets = _sdpa_h.thd_tma_offsets @@ -634,9 +637,13 @@ def _tmastg_warp_group( q_row_coord = q_super_idx * cutlass.Int32(CFG.TILES_Q * CFG.TILE_M) if cutlass.const_expr(CFG.THD_VARLEN): - o_desc_ptr = (o_desc_words.iterator.raw_ptr() + batch_idx * cutlass.Int32(_TENSOR_MAP_QWORDS)).tospace(cutlass.AddressSpace.generic) - o_slice = tma_slice_runtime_desc(o_desc_ptr, cutlass.Int32(0), head_idx, q_row_coord, cutlass.Int32(0)) - tma_store_tile(sO[0], o_slice) + # DEAD unit (batch == n_batch, envelope grid — issue #552): no O + # rows exist and descriptor slot n_batch is never built, so skip + # the store; the barrier protocol below still runs. + if batch_idx < n_batch: + o_desc_ptr = (o_desc_words.iterator.raw_ptr() + batch_idx * cutlass.Int32(_TENSOR_MAP_QWORDS)).tospace(cutlass.AddressSpace.generic) + o_slice = tma_slice_runtime_desc(o_desc_ptr, cutlass.Int32(0), head_idx, q_row_coord, cutlass.Int32(0)) + tma_store_tile(sO[0], o_slice) else: tma_store_tile( sO[0], @@ -1603,6 +1610,15 @@ def _host( # Dense padded-Q trim: separate (B,)-int32 per-batch Q lengths; None # (and absent from the compiled ABI) unless CFG.SEQ_Q_LENS_PRESENT. seq_q_lens_tensor: Optional[cute.Tensor] = None, + # THD device metadata build (issue #552): the CALLER's Q/KV length + # tensors — (B,) per-batch lengths or (B+1,) cu prefix sums, per side via + # thd_lens_form (bit 0: Q is cu, bit 1: KV is cu) — consumed only by the + # setup kernel, which writes the [kv|cu_q|cu_k] metadata buffer + # (seq_kv_lens_tensor) device-side. None (folded out of the ABI) for + # dense graphs. + thd_q_lens_tensor: Optional[cute.Tensor] = None, + thd_kv_lens_tensor: Optional[cute.Tensor] = None, + thd_lens_form: Optional[cutlass.Int32] = None, stream: _cuda_driver.CUstream = None, ) -> None: B, QH, KH, SQ, SKV, _ = problem_size @@ -1659,11 +1675,14 @@ def _tma_swz(byte_w: int): # ENVELOPE: the packed-O row stride is QH * ACTUAL d_v (o_tensor's # static inner extent), not QH * TILE_O — the per-batch descriptor # bases must step in real rows or every batch >= 1 lands OOB. - _build_o_descs_kernel( + _build_thd_meta_o_descs_kernel( o_tensor, tma_o_desc, o_desc_words, seq_kv_lens_tensor, + thd_q_lens_tensor, + thd_kv_lens_tensor, + thd_lens_form, cutlass.Int32(QH), cutlass.Int32(B), cutlass.Int32(o_tensor.stride[1]), @@ -1832,6 +1851,19 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): stride_order=(0,), assumed_align=16, ) + # THD: the caller's Q/KV length tensors, consumed by the setup kernel's + # device-side metadata build. DYNAMIC extents — (B,) per-batch lengths and + # (B+1,) cu prefix sums bind the same artifact; the form rides the runtime + # thd_lens_form bitmask, so no compile key grows (Rule 4). align 4: bound + # directly, only natural int32 alignment is guaranteed. + if CFG.THD_VARLEN: + fake_thd_q_lens = cute.runtime.make_fake_compact_tensor(cutlass.Int32, (cute.sym_int(divisibility=1),), stride_order=(0,), assumed_align=4) + fake_thd_kv_lens = cute.runtime.make_fake_compact_tensor(cutlass.Int32, (cute.sym_int(divisibility=1),), stride_order=(0,), assumed_align=4) + fake_thd_lens_form = cutlass.Int32(0) + else: + fake_thd_q_lens = None + fake_thd_kv_lens = None + fake_thd_lens_form = None return cute.compile( _host, fake_q, @@ -1848,6 +1880,9 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): cutlass.Float32(0.0), cutlass.Int32(0), fake_seq_q_lens, + fake_thd_q_lens, + fake_thd_kv_lens, + fake_thd_lens_form, stream=cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=False), options="--enable-tvm-ffi", ) 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 906d9cc9d..322449e05 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py @@ -311,9 +311,12 @@ class KernelTmemLayout: _decode_payload = _sdpa_h.decode_payload -from cudnn.sdpa.fwd.kernels.thd_sm100 import build_o_descs_kernel as _build_o_descs_kernel, TENSOR_MAP_QWORDS +from cudnn.sdpa.fwd.kernels.thd_sm100 import build_thd_meta_o_descs_kernel as _build_thd_meta_o_descs_kernel, TENSOR_MAP_QWORDS _TENSOR_MAP_QWORDS = TENSOR_MAP_QWORDS +# The setup kernel builds the THD metadata buffer DEVICE-side from the +# caller's length tensors and the adapter launches the plan-time envelope +# grid (issue #552) — no length ever reaches the host. _dispatch_decode_initial = _sdpa_h.dispatch_decode_initial _dispatch_decode_payload = _sdpa_h.dispatch_decode_payload _thd_tma_offsets = _sdpa_h.thd_tma_offsets @@ -1754,9 +1757,13 @@ def _tmastg_warp_group( q_row_coord = q_super_idx * cutlass.Int32(CFG.TILES_Q * CFG.TILE_M) if cutlass.const_expr(CFG.THD_VARLEN): - o_desc_ptr = (o_desc_words.iterator.raw_ptr() + batch_idx * cutlass.Int32(_TENSOR_MAP_QWORDS)).tospace(cutlass.AddressSpace.generic) - o_slice = tma_slice_runtime_desc(o_desc_ptr, cutlass.Int32(0), head_idx, q_row_coord, cutlass.Int32(0)) - tma_store_tile(sO[0], o_slice) + # DEAD unit (batch == n_batch, envelope grid — issue #552): no O + # rows exist and descriptor slot n_batch is never built, so skip + # the store; the barrier protocol below still runs. + if batch_idx < n_batch: + o_desc_ptr = (o_desc_words.iterator.raw_ptr() + batch_idx * cutlass.Int32(_TENSOR_MAP_QWORDS)).tospace(cutlass.AddressSpace.generic) + o_slice = tma_slice_runtime_desc(o_desc_ptr, cutlass.Int32(0), head_idx, q_row_coord, cutlass.Int32(0)) + tma_store_tile(sO[0], o_slice) else: tma_store_tile( sO[0], @@ -1803,6 +1810,15 @@ def _host( # Dense padded-Q trim: separate (B,)-int32 per-batch Q lengths; None # (and absent from the compiled ABI) unless CFG.SEQ_Q_LENS_PRESENT. seq_q_lens_tensor: Optional[cute.Tensor] = None, + # THD device metadata build (issue #552): the CALLER's Q/KV length + # tensors — (B,) per-batch lengths or (B+1,) cu prefix sums, per side via + # thd_lens_form (bit 0: Q is cu, bit 1: KV is cu) — consumed only by the + # setup kernel, which writes the [kv|cu_q|cu_k] metadata buffer + # (seq_kv_lens_tensor) device-side. None (folded out of the ABI) for + # dense graphs. + thd_q_lens_tensor: Optional[cute.Tensor] = None, + thd_kv_lens_tensor: Optional[cute.Tensor] = None, + thd_lens_form: Optional[cutlass.Int32] = None, stream: _cuda_driver.CUstream = None, ) -> None: B, QH, KH, SQ, SKV, _ = problem_size @@ -1859,11 +1875,14 @@ def _tma_swz(byte_w: int): # ENVELOPE: the packed-O row stride is QH * ACTUAL d_v (o_tensor's # static inner extent), not QH * TILE_O — the per-batch descriptor # bases must step in real rows or every batch >= 1 lands OOB. - _build_o_descs_kernel( + _build_thd_meta_o_descs_kernel( o_tensor, tma_o_desc, o_desc_words, seq_kv_lens_tensor, + thd_q_lens_tensor, + thd_kv_lens_tensor, + thd_lens_form, cutlass.Int32(QH), cutlass.Int32(B), cutlass.Int32(o_tensor.stride[1]), @@ -2034,6 +2053,19 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): stride_order=(0,), assumed_align=16, ) + # THD: the caller's Q/KV length tensors, consumed by the setup kernel's + # device-side metadata build. DYNAMIC extents — (B,) per-batch lengths and + # (B+1,) cu prefix sums bind the same artifact; the form rides the runtime + # thd_lens_form bitmask, so no compile key grows (Rule 4). align 4: bound + # directly, only natural int32 alignment is guaranteed. + if CFG.THD_VARLEN: + fake_thd_q_lens = cute.runtime.make_fake_compact_tensor(cutlass.Int32, (cute.sym_int(divisibility=1),), stride_order=(0,), assumed_align=4) + fake_thd_kv_lens = cute.runtime.make_fake_compact_tensor(cutlass.Int32, (cute.sym_int(divisibility=1),), stride_order=(0,), assumed_align=4) + fake_thd_lens_form = cutlass.Int32(0) + else: + fake_thd_q_lens = None + fake_thd_kv_lens = None + fake_thd_lens_form = None return cute.compile( _host, fake_q, @@ -2050,6 +2082,9 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): cutlass.Float32(0.0), cutlass.Int32(0), fake_seq_q_lens, + fake_thd_q_lens, + fake_thd_kv_lens, + fake_thd_lens_form, stream=cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=False), options="--enable-tvm-ffi", ) diff --git a/python/cudnn/sdpa/fwd/kernels/thd_sm100.py b/python/cudnn/sdpa/fwd/kernels/thd_sm100.py index 6fe9077c0..517d47305 100644 --- a/python/cudnn/sdpa/fwd/kernels/thd_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/thd_sm100.py @@ -13,6 +13,96 @@ TENSOR_MAP_QWORDS = 128 // 8 +@cute.kernel +def build_thd_meta_o_descs_kernel( + o_tensor: cute.Tensor, + base_o_desc: cutlass.GridConstant[tmap.TensorMap], + o_desc_words: cute.Tensor, + meta_t: cute.Tensor, + q_lens_t: cute.Tensor, + kv_lens_t: cute.Tensor, + lens_form: cutlass.Int32, + n_qh: cutlass.Int32, + n_batch: cutlass.Int32, + o_row_stride: cutlass.Int32, +) -> None: + """Per-execute THD setup, one elected thread (issue #552, D2H removal): + build the [seq_kv_lens(B) | cu_seqlens_q(B+1) | cu_seqlens_k(B+1)] metadata + buffer DEVICE-side from the caller's length tensors — ``(B,)`` per-batch + lengths (serial cumsum; B is small) or the ``(B+1,)`` cu prefix-sum form + (NORMALIZED by subtracting element 0 — the packed buffers are addressed + from token 0, so a cu tensor sliced from a larger prefix means the same + lengths, and the host can no longer validate ``cu[0] == 0`` (Rule 3), so + an unnormalized base must not leak into the offsets the tiles and the + dead-unit sentinel read; per-batch KV lengths are adjacent diffs either + way), per side via + ``lens_form`` (bit 0: Q is cu, bit 1: KV is cu) — then build the per-batch + O TMA descriptors from the cu_q values just written (same thread, program + order). Replaces the host tolist → cumsum → H2D round-trip with work + inside the setup launch that already existed for the descriptors.""" + if nvvm.elect_sync(): + meta = cutlass.make_array_view(meta_t) + ql = cutlass.make_array_view(q_lens_t) + kl = cutlass.make_array_view(kv_lens_t) + cuq0 = n_batch + cuk0 = cutlass.Int32(2) * n_batch + cutlass.Int32(1) + q_is_cu = (lens_form & cutlass.Int32(1)) != cutlass.Int32(0) + kv_is_cu = (lens_form & cutlass.Int32(2)) != cutlass.Int32(0) + if q_is_cu: + base_q = cutlass.Int32(ql[0]) + for b in cutlass.range(0, n_batch + cutlass.Int32(1), 1, unroll=1): + meta[cuq0 + b] = cutlass.Int32(ql[b]) - base_q + else: + acc = cutlass.Int32(0) + meta[cuq0] = cutlass.Int32(0) + for b in cutlass.range(0, n_batch, 1, unroll=1): + acc = acc + cutlass.Int32(ql[b]) + meta[cuq0 + b + cutlass.Int32(1)] = acc + if kv_is_cu: + base_k = cutlass.Int32(kl[0]) + meta[cuk0] = cutlass.Int32(0) + for b in cutlass.range(0, n_batch, 1, unroll=1): + meta[cuk0 + b + cutlass.Int32(1)] = cutlass.Int32(kl[b + cutlass.Int32(1)]) - base_k + meta[b] = cutlass.Int32(kl[b + cutlass.Int32(1)]) - cutlass.Int32(kl[b]) + else: + acc_k = cutlass.Int32(0) + meta[cuk0] = cutlass.Int32(0) + for b in cutlass.range(0, n_batch, 1, unroll=1): + lkv = cutlass.Int32(kl[b]) + meta[b] = lkv + acc_k = acc_k + lkv + meta[cuk0 + b + cutlass.Int32(1)] = acc_k + # Per-batch O descriptors, from the cu_q values written above (same + # thread — plain program order, no fence needed for the meta reads). + o_ptr = o_tensor.iterator.raw_ptr() + desc_base = o_desc_words.iterator.raw_ptr() + src_words = Pointer(base_o_desc.get_ptr(), dtype=cutlass.Int64) + 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): + (dptr + i).store((src_words + i).load()) + cu_q_b = cutlass.Int32(meta[cuq0 + b]) + s_i = cutlass.Int32(meta[cuq0 + b + cutlass.Int32(1)]) - cu_q_b + row_base = o_ptr + cu_q_b * row_elems + nvvm.tensormap_replace( + nvvm.TensormapField.GLOBAL_ADDRESS, + dptr, + new_value=row_base.toint(cutlass.Int64), + ) + nvvm.tensormap_replace( + nvvm.TensormapField.GLOBAL_DIM, + dptr, + new_value=s_i, + ord=2, + ) + nvvm.fence_proxy_release( + nvvm.MemScope.GPU, + from_proxy=nvvm.Proxy.GENERIC, + to_proxy=nvvm.Proxy.TENSORMAP, + ) + + @cute.kernel def build_o_descs_kernel( o_tensor: cute.Tensor, 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 ef2957a9b..dae32bd8b 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py @@ -1093,6 +1093,241 @@ def _run_and_check(seq_lens): assert info_exec.hits >= info_plan.hits + 2 +@pytest.mark.L0 +@torch_fork_set_rng(seed=38) +def test_dsl_sm100_thd_over_launched_units_are_dead(monkeypatch): + """Issue #552 (envelope-grid enabler): grid units past the live total are + DEAD by kernel contract — the decode maps them to the batch == n_batch + sentinel, every role takes the empty-KV path (eff_seqlen_kv reads + cu_q[0] == 0), and neither O nor the ragged Stats is written (LSE + predicate goes negative; the O-store role skips the TMA store). Padding + the launch grid must therefore change nothing. Covers a zero-length + LEADING sequence (pre-fix, dead units decoded to batch 0 and ran full KV + loops there), a zero-length middle sequence, all-KV-zero (dead units on + top of the packed-KV clamp), and both Stats layouts.""" + _require_dsl() + from cudnn.sdpa.fwd.api_dsl import SdpaFwdDslSm100 + + envelope = SdpaFwdDslSm100._thd_unit_envelope + monkeypatch.setattr(SdpaFwdDslSm100, "_thd_unit_envelope", lambda self: envelope(self) + 7) + _run_thd_stats_case(seq_lens_q=[128, 0, 64], seq_lens_kv=[100, 0, 30], mask="causal", stats_layout="token_major") + _run_thd_stats_case(seq_lens_q=[0, 64, 32], seq_lens_kv=[50, 40, 0], mask="none", stats_layout="head_major") + _run_thd_stats_case(seq_lens_q=[64, 32], seq_lens_kv=[0, 0], mask="none", stats_layout="token_major") + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=39) +def test_dsl_sm100_thd_lens_never_reach_host(monkeypatch): + """Issue #552 (D2H removal): the length tensors are consumed ONLY on + device — the setup kernel builds the + metadata, the views bind buffer capacities, and the grid is the + plan-time envelope — so ANY host read of them (the old tolist + round-trip) is a regression. The guard rejects every _thd_host_lens + call while full numerics run in both length forms.""" + _require_dsl() + from cudnn.sdpa.fwd.api_dsl import SdpaFwdDslSm100 + + def guard(self, seq_lens, name, cu_form): + raise AssertionError(f"lengths reached the host via {name} (device-meta modules must not read them)") + + monkeypatch.setattr(SdpaFwdDslSm100, "_thd_host_lens", guard) + _run_thd_stats_case(seq_lens_q=[200, 150], seq_lens_kv=[180, 120], mask="causal", stats_layout="token_major") + _run_thd_stats_case(seq_lens_q=[200, 150], seq_lens_kv=[180, 120], mask="causal", stats_layout="head_major", cu_lens=True) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=41) +def test_dsl_sm100_thd_execute_never_syncs(): + """Issue #552 endgame (d128): the THD execute performs NO synchronizing + CUDA call — no length D2H, no pageable H2D, no device/stream sync — + pinned by torch's sync debug mode ("error"), which raises on any. The + grid is the plan-time envelope, the ragged views bind buffer + capacities, and the metadata is built device-side; results are bitwise + identical to an unguarded execute.""" + _require_dsl() + from cudnn.sdpa.fwd.api_dsl import SdpaFwdDslSm100 + + b, h, s, d = 2, 4, 256, 128 + dtype = torch.float16 + q, k, v = (_bhsd(b, h, s, d, dtype) for _ in range(3)) + o = torch.zeros_like(q) + api = SdpaFwdDslSm100(sample_q=q, sample_k=k, sample_v=v, sample_o=o, thd=True) + assert api.check_support() + api.compile() + lens = torch.tensor([200, 150], dtype=torch.int32, device="cuda") + + # Warm-up outside the guarded region: allocator pools and lazy launcher + # state populate here, so the guarded execute reuses cached blocks. + api.execute(q_tensor=q, k_tensor=k, v_tensor=v, o_tensor=o, seq_q_lens=lens, seq_kv_lens=lens) + torch.cuda.synchronize() + o_ref = o.clone() + o.zero_() + prev_sync_mode = torch.cuda.get_sync_debug_mode() + torch.cuda.set_sync_debug_mode(2) + try: + api.execute(q_tensor=q, k_tensor=k, v_tensor=v, o_tensor=o, seq_q_lens=lens, seq_kv_lens=lens) + finally: + torch.cuda.set_sync_debug_mode(prev_sync_mode) + torch.cuda.synchronize() + assert torch.equal(o, o_ref) + + +@pytest.mark.L1 +@torch_fork_set_rng(seed=44) +def test_dsl_sm100_thd_cu_nonzero_base_normalized(): + """The device-side metadata build NORMALIZES cu prefix sums (subtracts + element 0): the packed buffers are addressed from token 0, so a cu + tensor sliced from a larger prefix (cu[0] != 0) means the same lengths + and must produce bitwise-identical results. The old host path errored + on cu[0] != 0; a device-side build cannot raise (Rule 3), it + normalizes.""" + _require_dsl() + from cudnn.sdpa.fwd.api_dsl import SdpaFwdDslSm100 + + b, h, s, d = 2, 4, 256, 128 + dtype = torch.float16 + q, k, v = (_bhsd(b, h, s, d, dtype) for _ in range(3)) + o = torch.zeros_like(q) + api = SdpaFwdDslSm100(sample_q=q, sample_k=k, sample_v=v, sample_o=o, thd=True, cu_seq_q_lens=True, cu_seq_kv_lens=True) + assert api.check_support() + api.compile() + + def _run(base): + cu = torch.tensor([base, base + 200, base + 350], dtype=torch.int32, device="cuda") + o.zero_() + api.execute(q_tensor=q, k_tensor=k, v_tensor=v, o_tensor=o, seq_q_lens=cu, seq_kv_lens=cu) + torch.cuda.synchronize() + return o.clone() + + assert torch.equal(_run(0), _run(1000)) + + +@pytest.mark.L1 +@torch_fork_set_rng(seed=43) +def test_dsl_sm100_thd_d192_d128_device_meta(): + """d192/d128 (native MLA head dims) THD through the device-meta + + envelope path: per-sequence numerics via the direct API — the graph THD + harness assumes d_qk == d_v, so this flavor's THD leg is pinned here.""" + _require_dsl() + from cudnn.sdpa.fwd.api_dsl import SdpaFwdDslSm100 + + b, h, s = 2, 4, 256 + d_qk, d_v = 192, 128 + dtype = torch.float16 + scale = 1.0 / math.sqrt(d_qk) + q, k = (_bhsd(b, h, s, d_qk, dtype) for _ in range(2)) + v = _bhsd(b, h, s, d_v, dtype) + o = torch.zeros_like(v) + api = SdpaFwdDslSm100(sample_q=q, sample_k=k, sample_v=v, sample_o=o, thd=True) + assert api.check_support() + api.compile() + seq_lens = [200, 150] + lens = torch.tensor(seq_lens, dtype=torch.int32, device="cuda") + api.execute(q_tensor=q, k_tensor=k, v_tensor=v, o_tensor=o, seq_q_lens=lens, seq_kv_lens=lens) + torch.cuda.synchronize() + base_q = q.transpose(1, 2).reshape(b * s, h, d_qk) + base_k = k.transpose(1, 2).reshape(b * s, h, d_qk) + base_v = v.transpose(1, 2).reshape(b * s, h, d_v) + base_o = o.transpose(1, 2).reshape(b * s, h, d_v) + off = 0 + for length in seq_lens: + qs = base_q[off : off + length].float() + ks = base_k[off : off + length].float() + vs = base_v[off : off + length].float() + scores = torch.einsum("lhd,mhd->hlm", qs, ks) * scale + ref = torch.einsum("hlm,mhd->lhd", torch.softmax(scores, dim=-1), vs) + torch.testing.assert_close(base_o[off : off + length].float(), ref, atol=5e-2, rtol=3e-2) + off += length + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=42) +def test_dsl_sm100_thd_execute_cuda_graph_capture(): + """Issue #552 endgame (d128): THD execute is CUDA-GRAPH CAPTURABLE — no + D2H, no pageable H2D, plan-time envelope grid. Capture once, then + replay with DIFFERENT lengths written into the same device tensors: the + replay must honor them (per-sequence lengths are read on device by the + setup and main kernels), proving no host value was baked into the + graph.""" + _require_dsl() + from cudnn.sdpa.fwd.api_dsl import SdpaFwdDslSm100 + + b, h, s, d = 2, 4, 256, 128 + dtype = torch.float16 + scale = 1.0 / math.sqrt(d) + q, k, v = (_bhsd(b, h, s, d, dtype) for _ in range(3)) + o = torch.zeros_like(q) + api = SdpaFwdDslSm100(sample_q=q, sample_k=k, sample_v=v, sample_o=o, thd=True) + assert api.check_support() + api.compile() + lens = torch.tensor([200, 150], dtype=torch.int32, device="cuda") + + def _check(seq_lens): + base_q = q.transpose(1, 2).reshape(b * s, h, d) + base_k = k.transpose(1, 2).reshape(b * s, h, d) + base_v = v.transpose(1, 2).reshape(b * s, h, d) + base_o = o.transpose(1, 2).reshape(b * s, h, d) + off = 0 + for length in seq_lens: + qs = base_q[off : off + length].float() + ks = base_k[off : off + length].float() + vs = base_v[off : off + length].float() + scores = torch.einsum("lhd,mhd->hlm", qs, ks) * scale + ref = torch.einsum("hlm,mhd->lhd", torch.softmax(scores, dim=-1), vs) + torch.testing.assert_close(base_o[off : off + length].float(), ref, atol=5e-2, rtol=3e-2) + off += length + + # Warm-up (compile artifacts, allocator, launcher state), then capture. + api.execute(q_tensor=q, k_tensor=k, v_tensor=v, o_tensor=o, seq_q_lens=lens, seq_kv_lens=lens) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + api.execute(q_tensor=q, k_tensor=k, v_tensor=v, o_tensor=o, seq_q_lens=lens, seq_kv_lens=lens) + graph.replay() + torch.cuda.synchronize() + _check([200, 150]) + # New lengths into the SAME device tensor — replay must honor them. + lens.copy_(torch.tensor([64, 33], dtype=torch.int32, device="cuda")) + graph.replay() + torch.cuda.synchronize() + _check([64, 33]) + + +@pytest.mark.L1 +@torch_fork_set_rng(seed=40) +def test_dsl_sm100_thd_kv_zero_capacity_clamp(): + """All-zero KV lengths, both storage shapes (issue #552): with LIVE K/V + storage the launch binds the buffers' capacity and the kernel's + per-sequence dead-row path zeroes O; with ZERO-numel K/V buffers the + packed-KV clamp binds the one never-dereferenced dummy token instead (a + zero-token view cannot back a TMA descriptor). Both must produce O == 0 + on every live row.""" + _require_dsl() + from cudnn.sdpa.fwd.api_dsl import SdpaFwdDslSm100 + + b, h, s, d = 2, 4, 128, 128 + dtype = torch.float16 + q = _bhsd(b, h, s, d, dtype) + k, v = _bhsd(b, h, s, d, dtype), _bhsd(b, h, s, d, dtype) + o = torch.full_like(q, 7.0) + api = SdpaFwdDslSm100(sample_q=q, sample_k=k, sample_v=v, sample_o=o, thd=True) + assert api.check_support() + api.compile() + lens_q = torch.tensor([100, 80], dtype=torch.int32, device="cuda") + lens_kv = torch.zeros(2, dtype=torch.int32, device="cuda") + + def _run(k_buf, v_buf): + o.fill_(7.0) + api.execute(q_tensor=q, k_tensor=k_buf, v_tensor=v_buf, o_tensor=o, seq_q_lens=lens_q, seq_kv_lens=lens_kv) + torch.cuda.synchronize() + return o.transpose(1, 2).reshape(b * s, h, d)[: 100 + 80].clone() + + live_capacity = _run(k, v) + zero_capacity = _run(torch.empty(0, dtype=dtype, device="cuda"), torch.empty(0, dtype=dtype, device="cuda")) + assert (live_capacity == 0).all(), "dead rows must be zeroed by the kernel's dead-row path" + assert torch.equal(live_capacity, zero_capacity) + + _COMBO_MASKS = { "dense": ["none", "causal", "causal_br", "swa", "padded", "band", "band_br", "band_swa", "swa_br"], # THD forces padding internally, so its mask axis rides on top of that.