From 0743c455ebbaa4d62cd834663754b96f2e5a1abf Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Fri, 14 Aug 2026 18:10:04 -0700 Subject: [PATCH 1/7] frost(sdpa): make over-launched THD grid units dead by contract (issue #552, SM100 d128) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A grid unit past the live total (sum of ceil(S_q_b/CGA_TILE_M)*QH) used to fall through _thd_decode with batch 0 — aliasing live tile (0,0,0): duplicated compute, racy (identical-byte) O/LSE writes, and full wasted KV loops when batch 0 is zero-length. Now it keeps the batch == n_batch sentinel, which neutralizes every consumer through in-bounds metadata reads: eff_seqlen_kv reads cu_q[0] == 0 (empty KV range in every role, in every mask mode), the epilogue's per-sequence Q length goes negative (the LSE predicate never fires), and the O-store role skips the TMA store explicitly (descriptor slot n_batch is never built). This is the enabler for the issue #552 endgame: an envelope grid sized from plan-time declarations needs over-launch to be harmless before the exact host-computed unit count (and its .tolist() D2H sync) can go. The exact unit count moves to _thd_unit_count() — the seam the envelope replaces, and the seam the new test pads (+7 units) to pin the contract across a zero-length leading sequence, a zero-length middle sequence, and all-KV-zero, in both Stats layouts. d128 f16 only: the shared decode change is inert for the other SM100 families (their grids stay exact) until their O-store roles get the same skip. Co-Authored-By: Claude Fable 5 --- python/cudnn/sdpa/fwd/api_dsl.py | 16 +++++++++++--- .../cudnn/sdpa/fwd/kernels/_common_sm100.py | 9 +++++++- .../fwd/kernels/prefill_d128_f16_sm100.py | 10 ++++++--- .../sdpa/frost/test_sdpa_fwd_dsl_sm100.py | 22 +++++++++++++++++++ 4 files changed, 50 insertions(+), 7 deletions(-) diff --git a/python/cudnn/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py index 6b672d79e..f19a2bd33 100644 --- a/python/cudnn/sdpa/fwd/api_dsl.py +++ b/python/cudnn/sdpa/fwd/api_dsl.py @@ -1236,6 +1236,18 @@ def _key(desc): o_stride=_key(self.o_desc), ) + def _thd_unit_count(self, slq_host) -> int: + """One THD unit per CGA-height slice of each sequence's Q rows. + + The kernel tolerates OVER-counting (issue #552): units past the live + total decode to the batch == n_batch dead sentinel — every role takes + the empty-KV path and neither O nor LSE is written — so a launch grid + sized above the exact total changes nothing but occupancy. The + envelope-grid redesign relies on this; the over-launch test pads it + to pin the contract.""" + cga_tile_m = int(self._k_mod.CGA_TILE_M) + return self.h_q * sum((l + cga_tile_m - 1) // cga_tile_m for l in slq_host) + 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. @@ -1311,9 +1323,7 @@ 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) + units = self._thd_unit_count(slq_host) # 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 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..65940921e 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py @@ -834,9 +834,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], 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..1f6e0fa36 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,28 @@ 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 + + exact = SdpaFwdDslSm100._thd_unit_count + monkeypatch.setattr(SdpaFwdDslSm100, "_thd_unit_count", lambda self, slq: exact(self, slq) + 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") + + _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. From 861dfc238a759bea2d41dbbf98305347bd2be69b Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Fri, 14 Aug 2026 19:09:03 -0700 Subject: [PATCH 2/7] frost(sdpa): build the THD metadata device-side; KV lengths never reach the host (issue #552, SM100 d128) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The KV-side .tolist() D2H sync, the host cumsum, and the pageable H2D metadata upload are gone on THD_DEVICE_META modules (SM100 d128): - The per-execute setup kernel (the single-thread O-descriptor builder) grows a first phase that builds the [kv_lens | cu_q | cu_k] metadata buffer on device from the CALLER's length tensors — (B,) per-batch lengths (serial cumsum, B is small) or the (B+1,) cu prefix-sum form — then builds the O descriptors from the cu values it just wrote (same thread, program order). The form rides a runtime bitmask and the fake lens tensors compile with dynamic extents, so both forms bind one artifact and no compile key grows (Rule 4). - The K/V ragged views bind their buffers' CAPACITY (numel // token stride; shared floor — K and V bind one dynamic token symbol) instead of the device-only packed total. Loads never step past the real per-sequence lengths the kernel reads from the device metadata, so the over-claim only widens the TMA descriptors' bound — the test harness's envelope storages already exercised exactly this. The zero-KV clamp is re-keyed from the packed total to capacity: all-zero lengths over live storage now launch normally through the kernel's per-sequence dead-row path; only zero-numel K/V buffers take the one-dummy-token clamp. - Knowingly given up: the KV-side cu prefix-sum invariant check — a validation that needs a device read is not a validation (Rule 3); it is caller contract now. The Q lengths still take ONE tolist: they size the exact launch grid. That sync dies with the plan-time envelope grid (next), whose dead-unit kernel contract is already in. Tests: a guard that FAILS if KV lengths ever reach the host (full numerics in both length forms), and an all-zero-KV two-shapes test (live capacity vs zero-numel clamp). Legacy modules (d192/d256/d512, SM120) keep the host-built-meta path until their ports. Co-Authored-By: Claude Fable 5 --- python/cudnn/AGENTS.md | 15 ++- python/cudnn/sdpa/fwd/api_dsl.py | 121 ++++++++++++++---- .../fwd/kernels/prefill_d128_f16_sm100.py | 48 ++++++- python/cudnn/sdpa/fwd/kernels/thd_sm100.py | 83 ++++++++++++ .../sdpa/frost/test_sdpa_fwd_dsl_sm100.py | 58 +++++++++ 5 files changed, 285 insertions(+), 40 deletions(-) diff --git a/python/cudnn/AGENTS.md b/python/cudnn/AGENTS.md index d5f125e22..222de2126 100644 --- a/python/cudnn/AGENTS.md +++ b/python/cudnn/AGENTS.md @@ -99,12 +99,15 @@ 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. + constant and no compile is keyed on it. The KV half is DONE on + `THD_DEVICE_META` modules (SM100 d128): the setup kernel builds the + metadata buffer device-side from the caller's length tensors and the K/V + views bind their buffers' capacity, so the KV lengths never reach the + host. Only the Q lengths still do — they size the exact launch grid; + removing that needs the plan-time-max (`b * s_q_max`) grid, whose + in-kernel dead-unit exit is already in (`_thd_decode`'s `batch == n_batch` + sentinel) (issue #552). `sdpa_fwd_wrapper_sm80` shows the alternative — it + requires `max_s_q` from the caller 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 f19a2bd33..3377d8063 100644 --- a/python/cudnn/sdpa/fwd/api_dsl.py +++ b/python/cudnn/sdpa/fwd/api_dsl.py @@ -1260,34 +1260,46 @@ 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): with a THD_DEVICE_META module the KV + lengths NEVER reach the host — the setup kernel builds the metadata + buffer device-side and the K/V views bind their buffers' capacity — + leaving ONE D2H sync, the Q lengths, which size the exact launch + grid (dies with the envelope grid). Legacy modules still take one + sync per length tensor and upload host-built metadata. Neither path + keys any compile: 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. + # Metadata buffer: [ seq_kv_lens(B) | cu_seqlens_q(B+1) | cu_seqlens_k(B+1) ]. + # Modules with THD_DEVICE_META build it DEVICE-side in the setup + # kernel from the caller's length tensors (either form) — no host + # cumsum, no H2D, and the KV lengths NEVER reach the host (issue + # #552); the KV-side prefix-sum invariants become caller contract + # (a validation that needs a device read is not a validation — + # AGENTS.md Rule 3). The Q lengths still take one tolist round-trip: + # they size the exact launch grid (dies with the envelope grid). + # Legacy modules build meta host-side from both tolists and upload it + # in ONE H2D copy. The torch work (allocation, any D2H length reads, + # any 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. + dev_meta = getattr(self._k_mod, "THD_DEVICE_META", False) + q_lens_dev = kv_lens_dev = None 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)) + if dev_meta: + 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") + ) + else: + 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] qh, kh = self.h_q, self.h_kv d_qk, d_v = self.head_dim_qk, self.head_dim_v @@ -1331,15 +1343,30 @@ 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) + if dev_meta: + # 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) + else: + # Legacy modules: the packed KV total from the host round-trip. + t_kv = cu_k_host[-1] 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 (legacy modules: all-zero seq_kv_lens): + # 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). @@ -1374,7 +1401,45 @@ 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) + if dev_meta: + # 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, + ) + else: + 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") @staticmethod 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 65940921e..156288f91 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,14 @@ # 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 (issue #552) — the adapter passes them through and +# never reads the KV lengths to the host. Modules without this flag still +# take the adapter's legacy host-built-meta path. +THD_DEVICE_META = True _dispatch_decode_initial = _sdpa_h.dispatch_decode_initial _dispatch_decode_payload = _sdpa_h.dispatch_decode_payload _thd_tma_offsets = _sdpa_h.thd_tma_offsets @@ -1939,6 +1944,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 @@ -2002,18 +2016,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]), @@ -2205,6 +2225,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, @@ -2221,6 +2254,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..6ec950aa4 100644 --- a/python/cudnn/sdpa/fwd/kernels/thd_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/thd_sm100.py @@ -13,6 +13,89 @@ 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 + (copied as-is, per-batch KV lengths derived by diff), 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: + for b in cutlass.range(0, n_batch + cutlass.Int32(1), 1, unroll=1): + meta[cuq0 + b] = cutlass.Int32(ql[b]) + 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: + meta[cuk0] = cutlass.Int32(kl[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)]) + 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 1f6e0fa36..2329e2492 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py @@ -1115,6 +1115,64 @@ def test_dsl_sm100_thd_over_launched_units_are_dead(monkeypatch): _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_kv_lens_never_reach_host(monkeypatch): + """Issue #552 (D2H removal, KV leg): with a THD_DEVICE_META module the KV + lengths are consumed ONLY by the setup kernel's device-side metadata + build — any host read of them (the old tolist round-trip) is a + regression. The guard rejects a KV-side _thd_host_lens call while full + numerics run in both length forms; the Q side still legitimately syncs + once for the exact launch grid (dies with the envelope grid).""" + _require_dsl() + from cudnn.sdpa.fwd.api_dsl import SdpaFwdDslSm100 + + orig = SdpaFwdDslSm100._thd_host_lens + + def guard(self, seq_lens, name, cu_form): + assert "kv" not in name, f"KV lengths reached the host via {name} (device-meta modules must not read them)" + return orig(self, seq_lens, name, cu_form) + + 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.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. From 6d860eff4bfd7cd1b5ae4c0339197887102a012b Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Fri, 14 Aug 2026 19:43:07 -0700 Subject: [PATCH 3/7] =?UTF-8?q?frost(sdpa):=20plan-time=20envelope=20THD?= =?UTF-8?q?=20grid=20=E2=80=94=20zero=20host=20reads,=20CUDA-graph=20captu?= =?UTF-8?q?rable=20(issue=20#552,=20SM100=20d128)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The last THD execute D2H sync is gone on THD_DEVICE_META modules: the launch grid is the PLAN-TIME envelope b * ceil(S_q_declared/CGA_TILE_M) * qh (every length is bounded by the declared S_q), and units past the live total exit through the dead-unit kernel contract — no loads, no O/LSE writes, one empty-mainloop barrier dance each. Q/O and the shared Stats token extent bind their buffers' CAPACITY (they share one dynamic token symbol; writes are bounded on device by the per-batch O-descriptor extents and the LSE row predicate). The head-major head_stride-covers-t_q check joins the caller contract on these modules (Rule 3), and the degenerate early-return keys on zero CAPACITY, not the packed total — all-zero lengths over live storage launch normally and touch nothing. Execute now reads NOTHING from device memory: fully async and CUDA-graph capturable, pinned by two new tests — execute under torch.cuda.set_sync_debug_mode('error'), and a capture/replay test that REWRITES the lengths in-place between replays (the replay honors them: they are read on device, nothing host-side is baked into the graph). Dead-tile tax, measured (B200 d128, B=8 QH=8 S_decl=16k, CGA_TILE_M=512): zero at full/near-full declarations, +0.2% at half-length, and ~145 ns per exposed dead unit when live work is tiny (64 of 2048 units: 33us -> 319us). Realistic THD prefill declarations (S_decl ~ max length in batch) pay <0.2%; far-oversized declarations pay the same tax the C++ backend's s_max grid does — a capped persistent grid reading a device-side live-unit count would bound it by resident clusters (follow-up, noted in AGENTS.md). The AGENTS Rule 3 THD known-violation entry is RESOLVED for these modules and now records the remaining ports (other SM100 families, SM120) and the dead-tile-tax follow-up. Co-Authored-By: Claude Fable 5 --- python/cudnn/AGENTS.md | 25 ++-- python/cudnn/sdpa/fwd/api_dsl.py | 105 +++++++++++----- .../sdpa/frost/test_sdpa_fwd_dsl_sm100.py | 112 ++++++++++++++++-- 3 files changed, 189 insertions(+), 53 deletions(-) diff --git a/python/cudnn/AGENTS.md b/python/cudnn/AGENTS.md index 222de2126..102a45107 100644 --- a/python/cudnn/AGENTS.md +++ b/python/cudnn/AGENTS.md @@ -97,17 +97,20 @@ 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. The KV half is DONE on - `THD_DEVICE_META` modules (SM100 d128): the setup kernel builds the - metadata buffer device-side from the caller's length tensors and the K/V - views bind their buffers' capacity, so the KV lengths never reach the - host. Only the Q lengths still do — they size the exact launch grid; - removing that needs the plan-time-max (`b * s_q_max`) grid, whose - in-kernel dead-unit exit is already in (`_thd_decode`'s `batch == n_batch` - sentinel) (issue #552). `sdpa_fwd_wrapper_sm80` shows the alternative — it - requires `max_s_q` from the caller rather than deriving it. + SM100 and SM120). RESOLVED on `THD_DEVICE_META` modules (SM100 d128, + issue #552) — the reference for the remaining ports: 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: the OTHER SM100 + families and SM120 (exact host-computed grid + host-built metadata), 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 3377d8063..aacbc3429 100644 --- a/python/cudnn/sdpa/fwd/api_dsl.py +++ b/python/cudnn/sdpa/fwd/api_dsl.py @@ -1237,17 +1237,35 @@ def _key(desc): ) def _thd_unit_count(self, slq_host) -> int: - """One THD unit per CGA-height slice of each sequence's Q rows. + """One THD unit per CGA-height slice of each sequence's Q rows — + the EXACT total, host-computed from the length round-trip (legacy + modules only). The kernel tolerates OVER-counting (issue #552): units past the live total decode to the batch == n_batch dead sentinel — every role takes the empty-KV path and neither O nor LSE is written — so a launch grid - sized above the exact total changes nothing but occupancy. The - envelope-grid redesign relies on this; the over-launch test pads it - to pin the contract.""" + sized above the exact total changes nothing but occupancy.""" cga_tile_m = int(self._k_mod.CGA_TILE_M) return self.h_q * sum((l + cga_tile_m - 1) // cga_tile_m for l in slq_host) + 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. @@ -1260,14 +1278,16 @@ 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. - Host round-trips (issue #552): with a THD_DEVICE_META module the KV - lengths NEVER reach the host — the setup kernel builds the metadata - buffer device-side and the K/V views bind their buffers' capacity — - leaving ONE D2H sync, the Q lengths, which size the exact launch - grid (dies with the envelope grid). Legacy modules still take one - sync per length tensor and upload host-built metadata. Neither path - keys any compile: the kernels compile with DYNAMIC token extents, so - a new packed total re-binds the same artifact.""" + Host round-trips (issue #552): with a THD_DEVICE_META module there + are 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. Legacy modules still take one sync per + length tensor, upload host-built metadata, and launch the exact + grid. Neither path keys any compile: the kernels compile with + DYNAMIC token extents, so a new packed total re-binds the same + artifact.""" import cutlass dev = q_buf.device @@ -1277,10 +1297,10 @@ def _execute_thd(self, q_buf, k_buf, v_buf, o_buf, scale_softmax_log2, sinks, se # Modules with THD_DEVICE_META build it DEVICE-side in the setup # kernel from the caller's length tensors (either form) — no host # cumsum, no H2D, and the KV lengths NEVER reach the host (issue - # #552); the KV-side prefix-sum invariants become caller contract - # (a validation that needs a device read is not a validation — - # AGENTS.md Rule 3). The Q lengths still take one tolist round-trip: - # they size the exact launch grid (dies with the envelope grid). + # #552); the prefix-sum invariants become caller contract (a + # validation that needs a device read is not a validation — + # AGENTS.md Rule 3), and NOTHING is read back — the grid is the + # plan-time envelope. # Legacy modules build meta host-side from both tolists and upload it # in ONE H2D copy. The torch work (allocation, any D2H length reads, # any H2D upload) runs on the LAUNCH stream so it is ordered against @@ -1290,36 +1310,60 @@ def _execute_thd(self, q_buf, k_buf, v_buf, o_buf, scale_softmax_log2, sinks, se q_lens_dev = kv_lens_dev = None 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) if dev_meta: 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") ) else: + 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] 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.) + if dev_meta: + # 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) + else: + t_q = cu_q_host[-1] + + # Degenerate Q side: zero CAPACITY (device-meta modules — no token is + # addressable, and a zero-token view cannot back a TMA descriptor) or + # a zero packed total (legacy modules): nothing to compute or write. + # All-zero LENGTHS over live storage launch normally on device-meta + # modules — 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 = 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 not dev_meta: + # Device-meta modules cannot know t_q host-side; the + # cover invariant is caller contract there (a validation + # that needs a device read is not a validation — Rule 3). + 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) @@ -1335,7 +1379,10 @@ 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) - units = self._thd_unit_count(slq_host) + # Device-meta modules launch the PLAN-TIME envelope grid (dead units + # exit by kernel contract); legacy modules the exact host-computed + # total. + units = self._thd_unit_envelope() if dev_meta else self._thd_unit_count(slq_host) # 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 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 2329e2492..e30cca41e 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py @@ -1108,8 +1108,8 @@ def test_dsl_sm100_thd_over_launched_units_are_dead(monkeypatch): _require_dsl() from cudnn.sdpa.fwd.api_dsl import SdpaFwdDslSm100 - exact = SdpaFwdDslSm100._thd_unit_count - monkeypatch.setattr(SdpaFwdDslSm100, "_thd_unit_count", lambda self, slq: exact(self, slq) + 7) + 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") @@ -1117,27 +1117,113 @@ def test_dsl_sm100_thd_over_launched_units_are_dead(monkeypatch): @pytest.mark.L0 @torch_fork_set_rng(seed=39) -def test_dsl_sm100_thd_kv_lens_never_reach_host(monkeypatch): - """Issue #552 (D2H removal, KV leg): with a THD_DEVICE_META module the KV - lengths are consumed ONLY by the setup kernel's device-side metadata - build — any host read of them (the old tolist round-trip) is a - regression. The guard rejects a KV-side _thd_host_lens call while full - numerics run in both length forms; the Q side still legitimately syncs - once for the exact launch grid (dies with the envelope grid).""" +def test_dsl_sm100_thd_lens_never_reach_host(monkeypatch): + """Issue #552 (D2H removal): with a THD_DEVICE_META module 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 - orig = SdpaFwdDslSm100._thd_host_lens - def guard(self, seq_lens, name, cu_form): - assert "kv" not in name, f"KV lengths reached the host via {name} (device-meta modules must not read them)" - return orig(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_() + 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(0) + torch.cuda.synchronize() + assert torch.equal(o, o_ref) + + +@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(): From baca6be9893d0f064185832a4b1811ca7fb5d31d Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Fri, 14 Aug 2026 19:57:06 -0700 Subject: [PATCH 4/7] frost(sdpa): port the THD device-meta + envelope path to d192_d128 / d256 / d512 (issue #552) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Mechanical replication of the d128 template — the decode dead-sentinel and the setup kernel were already shared, so each family needed only: the O-store role's dead-unit skip (batch == n_batch never stores through the unbuilt pad descriptor slot), the THD_DEVICE_META flag + setup-kernel import swap, the _host lens-tensor params + launch swap, and the compile() fakes (dynamic extents; no compile key grows). All four SM100 f16 THD families now execute with zero host reads on the plan-time envelope grid. New test: d192/d128 (native MLA head dims) THD numerics via the direct API — the graph THD harness assumes d_qk == d_v, so this flavor's THD leg had no coverage. Still on the legacy host-meta path: SM120 f16/fp8 (different engine class and grid mechanism — needs its own dead-unit contract and setup kernel, and SM120 CI to validate). Co-Authored-By: Claude Fable 5 --- .../kernels/prefill_d192_d128_f16_sm100.py | 46 +++++++++++++++++-- .../fwd/kernels/prefill_d256_f16_sm100.py | 46 +++++++++++++++++-- .../fwd/kernels/prefill_d512_f16_sm100.py | 46 +++++++++++++++++-- .../sdpa/frost/test_sdpa_fwd_dsl_sm100.py | 38 +++++++++++++++ 4 files changed, 161 insertions(+), 15 deletions(-) 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..4d05840a5 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,13 @@ 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. +THD_DEVICE_META = True _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 +911,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 +2064,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 +2143,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 +2334,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 +2363,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..5a4bc9bfb 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,13 @@ _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. +THD_DEVICE_META = True _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 +638,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 +1611,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 +1676,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 +1852,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 +1881,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..c188f5900 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,13 @@ 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. +THD_DEVICE_META = True _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 +1758,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 +1811,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 +1876,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 +2054,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 +2083,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/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py index e30cca41e..79723cff0 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py @@ -1171,6 +1171,44 @@ def test_dsl_sm100_thd_execute_never_syncs(): assert torch.equal(o, o_ref) +@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(): From 69eb53b48c9c90cfbe4c8d4c5d86a2260fc0a749 Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Fri, 14 Aug 2026 23:21:43 -0700 Subject: [PATCH 5/7] =?UTF-8?q?frost(sdpa):=20drop=20the=20THD=5FDEVICE=5F?= =?UTF-8?q?META=20migration=20seam=20=E2=80=94=20one=20THD=20execute=20pat?= =?UTF-8?q?h=20on=20SM100?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The flag existed so each migration commit could land atomically: d128 flipped to the device-meta path while the other families stayed correct (and tested) on the host-meta path. With all four f16 families ported it is dead code — SdpaFwdDslSm100._execute_thd is reachable only for f16 modules (the per-tensor-FP8 and MXFP8 executes return before the THD dispatch, and SM120 has its own _execute_thd) and every f16 module has the device-meta ABI. Remove the flag, the legacy branches (host tolists, cumsum, H2D meta upload, exact host-computed grid), and the now-unused _thd_unit_count. A future SM100 module without the device-meta ABI fails loudly at the launch call rather than silently taking a host path. Co-Authored-By: Claude Fable 5 --- python/cudnn/AGENTS.md | 9 +- python/cudnn/sdpa/fwd/api_dsl.py | 207 ++++++------------ .../fwd/kernels/prefill_d128_f16_sm100.py | 6 +- .../kernels/prefill_d192_d128_f16_sm100.py | 1 - .../fwd/kernels/prefill_d256_f16_sm100.py | 1 - .../fwd/kernels/prefill_d512_f16_sm100.py | 1 - .../sdpa/frost/test_sdpa_fwd_dsl_sm100.py | 4 +- 7 files changed, 82 insertions(+), 147 deletions(-) diff --git a/python/cudnn/AGENTS.md b/python/cudnn/AGENTS.md index 102a45107..82573f1a1 100644 --- a/python/cudnn/AGENTS.md +++ b/python/cudnn/AGENTS.md @@ -97,15 +97,16 @@ 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). RESOLVED on `THD_DEVICE_META` modules (SM100 d128, - issue #552) — the reference for the remaining ports: the kernels compile + 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: the OTHER SM100 - families and SM120 (exact host-computed grid + host-built metadata), and + 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 diff --git a/python/cudnn/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py index aacbc3429..304d5f12d 100644 --- a/python/cudnn/sdpa/fwd/api_dsl.py +++ b/python/cudnn/sdpa/fwd/api_dsl.py @@ -1236,18 +1236,6 @@ def _key(desc): o_stride=_key(self.o_desc), ) - def _thd_unit_count(self, slq_host) -> int: - """One THD unit per CGA-height slice of each sequence's Q rows — - the EXACT total, host-computed from the length round-trip (legacy - modules only). - - The kernel tolerates OVER-counting (issue #552): units past the live - total decode to the batch == n_batch dead sentinel — every role takes - the empty-KV path and neither O nor LSE is written — so a launch grid - sized above the exact total changes nothing but occupancy.""" - cga_tile_m = int(self._k_mod.CGA_TILE_M) - return self.h_q * sum((l + cga_tile_m - 1) // cga_tile_m for l in slq_host) - def _thd_unit_envelope(self) -> int: """PLAN-TIME upper bound on live THD units: ``B * ceil(S_q_declared / CGA_TILE_M) * QH``. @@ -1278,73 +1266,52 @@ 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. - Host round-trips (issue #552): with a THD_DEVICE_META module there - are 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. Legacy modules still take one sync per - length tensor, upload host-built metadata, and launch the exact - grid. Neither path keys any compile: the kernels compile with - DYNAMIC token extents, so a new packed total re-binds the same - artifact.""" + 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) ]. - # Modules with THD_DEVICE_META build it DEVICE-side in the setup - # kernel from the caller's length tensors (either form) — no host - # cumsum, no H2D, and the KV lengths NEVER reach the host (issue - # #552); the prefix-sum invariants become caller contract (a - # validation that needs a device read is not a validation — - # AGENTS.md Rule 3), and NOTHING is read back — the grid is the - # plan-time envelope. - # Legacy modules build meta host-side from both tolists and upload it - # in ONE H2D copy. The torch work (allocation, any D2H length reads, - # any 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. - dev_meta = getattr(self._k_mod, "THD_DEVICE_META", False) - q_lens_dev = kv_lens_dev = None + # Metadata buffer: [ seq_kv_lens(B) | cu_seqlens_q(B+1) | cu_seqlens_k(B+1) ], + # 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) - if dev_meta: - 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") - ) - else: - 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)) + 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 - if dev_meta: - # 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) - else: - t_q = cu_q_host[-1] - - # Degenerate Q side: zero CAPACITY (device-meta modules — no token is - # addressable, and a zero-token view cannot back a TMA descriptor) or - # a zero packed total (legacy modules): nothing to compute or write. - # All-zero LENGTHS over live storage launch normally on device-meta - # modules — every unit decodes dead and neither O nor LSE is touched. + # 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: @@ -1353,15 +1320,9 @@ def _execute_thd(self, q_buf, k_buf, v_buf, o_buf, scale_softmax_log2, sinks, se 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 - if not dev_meta: - # Device-meta modules cannot know t_q host-side; the - # cover invariant is caller contract there (a validation - # that needs a device read is not a validation — Rule 3). - 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()) @@ -1379,10 +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) - # Device-meta modules launch the PLAN-TIME envelope grid (dead units - # exit by kernel contract); legacy modules the exact host-computed - # total. - units = self._thd_unit_envelope() if dev_meta else self._thd_unit_count(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 @@ -1390,22 +1349,18 @@ 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) - if dev_meta: - # 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) - else: - # Legacy modules: the packed KV total from the host round-trip. - t_kv = cu_k_host[-1] + # 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: - # No KV storage at all (legacy modules: all-zero seq_kv_lens): + # 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), @@ -1448,45 +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) - if dev_meta: - # 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, - ) - else: - 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 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 156288f91..6c8976a42 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py @@ -193,10 +193,8 @@ _TENSOR_MAP_QWORDS = TENSOR_MAP_QWORDS # The setup kernel builds the THD metadata buffer DEVICE-side from the -# caller's length tensors (issue #552) — the adapter passes them through and -# never reads the KV lengths to the host. Modules without this flag still -# take the adapter's legacy host-built-meta path. -THD_DEVICE_META = True +# 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 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 4d05840a5..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 @@ -206,7 +206,6 @@ def _bounds_for_tile( # 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. -THD_DEVICE_META = True _dispatch_decode_initial = _sdpa_h.dispatch_decode_initial _dispatch_decode_payload = _sdpa_h.dispatch_decode_payload _thd_tma_offsets = _sdpa_h.thd_tma_offsets 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 5a4bc9bfb..4365ac423 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py @@ -126,7 +126,6 @@ # 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. -THD_DEVICE_META = True _dispatch_decode_initial = _sdpa_h.dispatch_decode_initial _dispatch_decode_payload = _sdpa_h.dispatch_decode_payload _thd_tma_offsets = _sdpa_h.thd_tma_offsets 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 c188f5900..322449e05 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py @@ -317,7 +317,6 @@ class KernelTmemLayout: # 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. -THD_DEVICE_META = True _dispatch_decode_initial = _sdpa_h.dispatch_decode_initial _dispatch_decode_payload = _sdpa_h.dispatch_decode_payload _thd_tma_offsets = _sdpa_h.thd_tma_offsets 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 79723cff0..351b7de7b 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py @@ -1118,8 +1118,8 @@ def test_dsl_sm100_thd_over_launched_units_are_dead(monkeypatch): @pytest.mark.L0 @torch_fork_set_rng(seed=39) def test_dsl_sm100_thd_lens_never_reach_host(monkeypatch): - """Issue #552 (D2H removal): with a THD_DEVICE_META module the length - tensors are consumed ONLY on device — the setup kernel builds the + """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 From 43617148b6609577cbee780980830c70bab1ee5a Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Fri, 14 Aug 2026 23:51:54 -0700 Subject: [PATCH 6/7] frost(sdpa): normalize THD cu prefixes device-side; restore sync-debug mode in the async test (PR #606 review) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The device-side metadata build now subtracts element 0 from a cu prefix-sum tensor before writing cu_q / cu_k: 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 packed offsets the tiles read or the dead-unit sentinel's cu_q[0] == 0 empty-KV guarantee. The old host path raised on cu[0] != 0; a device build cannot raise, it normalizes. Regression test: base-0 and base-1000 cu tensors over the same buffers produce bitwise-identical O. The zero-sync test now restores the caller's sync-debug mode instead of resetting it to 0. Co-Authored-By: Claude Fable 5 --- python/cudnn/sdpa/fwd/kernels/thd_sm100.py | 15 ++++++--- .../sdpa/frost/test_sdpa_fwd_dsl_sm100.py | 33 ++++++++++++++++++- 2 files changed, 43 insertions(+), 5 deletions(-) diff --git a/python/cudnn/sdpa/fwd/kernels/thd_sm100.py b/python/cudnn/sdpa/fwd/kernels/thd_sm100.py index 6ec950aa4..517d47305 100644 --- a/python/cudnn/sdpa/fwd/kernels/thd_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/thd_sm100.py @@ -30,7 +30,12 @@ def build_thd_meta_o_descs_kernel( 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 - (copied as-is, per-batch KV lengths derived by diff), per side via + (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 @@ -44,8 +49,9 @@ def build_thd_meta_o_descs_kernel( 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]) + meta[cuq0 + b] = cutlass.Int32(ql[b]) - base_q else: acc = cutlass.Int32(0) meta[cuq0] = cutlass.Int32(0) @@ -53,9 +59,10 @@ def build_thd_meta_o_descs_kernel( acc = acc + cutlass.Int32(ql[b]) meta[cuq0 + b + cutlass.Int32(1)] = acc if kv_is_cu: - meta[cuk0] = cutlass.Int32(kl[0]) + 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)]) + 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) 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 351b7de7b..dae32bd8b 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py @@ -1162,15 +1162,46 @@ def test_dsl_sm100_thd_execute_never_syncs(): 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(0) + 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(): From 690f8bf128fbeee843a14d6f08a96aac34cc2d32 Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Sat, 15 Aug 2026 01:49:04 -0700 Subject: [PATCH 7/7] frost(sdpa): read the MXFP8 SF batch from the graph facts, not the bound tensor's shape MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit An F8_128x4-reordered scale-factor tensor is an opaque byte layout: the graph declares logical [B, H, s_padded, d/32] dims, but callers legally bind any right-sized buffer — the TE-style swizzle producers (and test_mhas_v2's torch-only fallback) hand over flat [B*H*s_padded, 4] tensors. _reshape_sf took B from sf.shape[0], so any frost-served graph with a flat-bound SF blew up with 'MXFP8 SF size mismatch' (sf.shape[0]=B*H*s_padded misread as B). Latent since the mxfp8 engine landed; surfaced on develop CI (test_sdpa_mxfp8_fwd_L0[test193] failing on every Blackwell frost_tests:sdpa run) once the random sweep produced a padded-stride config the backend's own plans decline — frost is the fallback there (1/162 of the mxfp8 sweep), so the sweep only now reached this path. Fix: B comes from self.batch_size; only the total byte count is validated; the reshape is unchanged. Full mxfp8 fwd L0 sweep: 162 passed. --- python/cudnn/sdpa/fwd/api_dsl.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/python/cudnn/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py index 304d5f12d..fd9bbc1f5 100644 --- a/python/cudnn/sdpa/fwd/api_dsl.py +++ b/python/cudnn/sdpa/fwd/api_dsl.py @@ -1433,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)