diff --git a/python/cudnn/AGENTS.md b/python/cudnn/AGENTS.md index 0c8e1b856..d5f125e22 100644 --- a/python/cudnn/AGENTS.md +++ b/python/cudnn/AGENTS.md @@ -97,11 +97,14 @@ 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). `t_q`/`t_kv` reach the host only because `T` is a - compile-time constant. Compile on the `b * s_q_max` envelope and pass `T` as a - runtime argument, as the f16 kernels already do for head dims; - `sdpa_fwd_wrapper_sm80` shows the other half — it requires `max_s_q` from the - caller rather than deriving it. + 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. - 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 @@ -122,6 +125,65 @@ When auditing this list, grep for the ARGUMENT, not the call shape: `device="cpu"` finds `to(dtype=..., device="cpu")`, which `to(device="cpu")` misses. +**Rule 4 — compile keys are PLAN-TIME-ONLY: never key a kernel compile on +runtime data values.** + +`cute.compile` takes seconds. Anything an execute path feeds into a +compile-cache key (an `lru_cache`d `compile()` wrapper, a template parameter, +a fake-tensor extent) must be derivable from the graph declaration alone — +tensor dtypes, declared strides, head counts, head dims, flags. Values read +out of runtime tensors (THD packed token totals, max sequence lengths, batch +contents) change every step under continuous batching, so a key that includes +them degenerates into a fresh multi-second compile per `execute()` — a +pathology that no correctness test catches (issue #552 is the case study: +`sq=t_q, skv=t_kv` in the THD compile key). Rule 3 bans the read that feeds +such a key; this rule bans the key itself — a runtime value that arrives +legally (a caller-passed host scalar, an `int(tensor.shape[...])`) still must +not become a compile key. + +- **Runtime extents compile DYNAMIC.** Use `cute.sym_int()` in the fake + tensors (one symbol per ragged group) so one compiled artifact re-binds any + total; runtime scalars the launch needs (grid extents like THD `max_sq`) + are `cutlass.Int32` call arguments, never compile parameters. +- **Derived values count.** A stride tuple whose batch stride is + `t_q * token_stride` smuggles the runtime total into the key just as + surely as `sq=t_q` — normalize it out (zero the never-stepped batch + stride, rebuild it symbolically kernel-side). +- **Compile at plan time, re-bind at execute.** With a plan-time-only key + there is no reason to defer: `compile()` builds the artifact once and the + execute path's cached call must be a guaranteed hit. Guard it with a + cache-miss regression test (see + `test_dsl_sm100_thd_compile_key_plan_time_only`), not by inspection. +- **Known open cleanup (issue #604)**: the SM80 engines' `_compile_cached` + (#493) still keys `SQ`/`SKV` under `THD_VARLEN` — migrate it to dynamic + token extents like the SM100/SM120 THD compiles rather than copying its + pattern. + +**Rule 5 — every torch operation on the execute path is ordered on the +LAUNCH stream, never implicitly on torch's current stream.** + +The kernel launches on the stream carried by the execute-time handle +(`ExecutionContext.stream`), but torch enqueues work — H2D metadata uploads, +buffer resets (`zero_()`), post-kernel reductions (`div_()`, `copy_()`), +and the caching allocator's stream-tagging of fresh blocks — on +`torch.cuda.current_stream()`. When the two differ, the prep and the kernel +race (PR #543 is the case study: the THD `[seq_kv | cu_q | cu_k]` upload vs +the kernel that reads it). + +- **Resolve the launch stream FIRST**, before any torch work in the execute + path, and run every torch op (including allocator calls: workspace-less + fallback allocations, cached-dummy first use) inside + `_torch_stream_context(current_stream, device)` — see the fp8/mxfp8 amax + resets and both `_execute_thd` paths in `sdpa/fwd/api_dsl.py`. +- **Consumers too, not just producers**: anything reading what the kernel + wrote (`amax_o.div_()`, an O scratch copy-back) belongs on the launch + stream for the same reason. +- The PyTorch-integration path launches on torch's current stream, where the + context is a no-op — the race only bites direct graph-API users with an + explicit handle stream, which is exactly why tests miss it. Order the work + by construction rather than relying on the common case. + + ## Frontend-only kernel package layout ``` diff --git a/python/cudnn/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py index 614e3f8b5..6b672d79e 100644 --- a/python/cudnn/sdpa/fwd/api_dsl.py +++ b/python/cudnn/sdpa/fwd/api_dsl.py @@ -954,8 +954,13 @@ def compile(self) -> None: ) self._k_mod = _load_sm100_kernel_module(self.flavor, params, fp8=self._fp8, pertensor=self._pertensor, rubin=(self._device_cc == (10, 7))) if self.thd: - # T (total tokens) is a runtime value, so the per-shape compile is deferred to execute(). - self._compiled_kernel = "thd-deferred" + # The THD compile key is PLAN-TIME-ONLY (the packed token totals + # compile as dynamic extents — issue #552), so compile HERE like + # every dense specialization; execute()'s lru-cached call re-binds + # this artifact. (The all-KV-zero clamp swaps the K/V strides and + # mints its own entry on first hit.) FP8 THD is not wired + # (_execute_thd is f16-only); keep the deferred sentinel there. + self._compiled_kernel = "thd-deferred" if self._fp8 else self._k_mod.compile(**self._thd_compile_kwargs()) elif self._fp8: # FP8/MXFP8 kernels are exact-match d128 (gated in check_support); # their compile() has no envelope head-dim parameters. has_lse=False @@ -1196,6 +1201,41 @@ def execute( O_view.copy_(O_scratch) self._logger.debug("execute completed") + def _thd_compile_kwargs(self) -> dict: + """The THD compile key — PLAN-TIME-ONLY by contract (issue #552). + + The packed token totals are runtime values and compile as DYNAMIC + extents; everything here (logical batch, heads, head dims, the Stats + specialization, the declared strides with the batch stride zeroed — + ``_thd_view``'s batch stride is ``t * token_stride``, a runtime value + that never steps at batch extent 1) is known when the graph is built, + so ``compile()`` compiles eagerly and ``_execute_thd``'s lru-cached + call re-binds the same artifact for every packed total.""" + + def _key(desc): + (ts, hs, es), _ = self._thd_declared(desc) + return (0, ts, hs, es) + + has_lse = self.lse_desc is not None + return dict( + b=self.batch_size, + qh=self.h_q, + kh=self.h_kv, + d_qk=self.head_dim_qk, + d_v=self.head_dim_v, + # The Stats layout is a per-shape specialization (like d_qk/d_v): + # has_lse=False compiles the store out; token-major binds the + # packed rank-2 (T, H) view; head-major carries the caller-declared + # head-row stride (0 -> compact t_q). + has_lse=has_lse, + lse_head_major=has_lse and self.thd_stats_head_major, + lse_head_stride=(self.thd_stats_head_stride if (has_lse and self.thd_stats_head_major) else 0), + q_stride=_key(self.q_desc), + k_stride=_key(self.k_desc), + v_stride=_key(self.v_desc), + o_stride=_key(self.o_desc), + ) + 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. @@ -1209,8 +1249,12 @@ def _execute_thd(self, q_buf, k_buf, v_buf, o_buf, scale_softmax_log2, sinks, se with tokens contiguous within each head row; when ``None`` the kernel compiles the LSE store out (has_lse=False) and no scratch exists. The host round-trip for the runtime totals (t_q / t_kv / unit count) - is inherent to the lowering — the packed extents are data-dependent — - and costs one D2H sync per length tensor, no device allocation.""" + 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.""" import cutlass dev = q_buf.device @@ -1221,11 +1265,15 @@ def _execute_thd(self, q_buf, k_buf, v_buf, o_buf, scale_softmax_log2, sinks, se # 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. - 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)) + # 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. + with _torch_stream_context(current_stream, dev): + meta = carver.take(3 * b + 2, torch.int32) if carver is not None else torch.empty(3 * b + 2, dtype=torch.int32, device=dev) + slq_host, cu_q_host = self._thd_host_lens(seq_q_lens, "cu_seq_len_q" if self.cu_seq_q_lens else "seq_q_lens", self.cu_seq_q_lens) + slk_host, cu_k_host = self._thd_host_lens(seq_len_kv, "cu_seq_len_kv" if self.cu_seq_kv_lens else "seq_kv_lens", self.cu_seq_kv_lens) + meta.copy_(torch.tensor(slk_host + cu_q_host + cu_k_host, dtype=torch.int32)) t_q = cu_q_host[-1] t_kv = cu_k_host[-1] @@ -1254,10 +1302,15 @@ def _execute_thd(self, q_buf, k_buf, v_buf, o_buf, scale_softmax_log2, sinks, se # view — the kernel's epilogue dispatches on this static rank. lse = lse_tensor.as_strided((t_q, qh), (qh, 1), lse_tensor.storage_offset()) - # Per-sequence O TMA descriptors, filled by the kernel's builder pass. - o_desc = carver.take(b * 16 + 16, torch.int64) if carver is not None else torch.zeros(b * 16 + 16, dtype=torch.int64, device=dev) - if carver is not None: - o_desc.zero_() + # Per-sequence O TMA descriptors. No zero-init: the kernel's builder + # pass copies every qword of each sequence's slot from the base + # descriptor (then patches address/extent) before the fence and + # before any consumer read, so stale bytes never survive — a fill + # here is a wasted kernel launch on the execute hot path (Rule 1). + # Allocated on the launch stream (allocator stream-tagging + ordering + # 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) @@ -1292,35 +1345,25 @@ def _execute_thd(self, q_buf, k_buf, v_buf, o_buf, scale_softmax_log2, sinks, se LSE = lse if sinks is not None: sinks_t = self._checked_sinks_1d(sinks) - elif carver is not None: - sinks_t = carver.take(qh, torch.float32) - sinks_t.zero_() else: - sinks_t = torch.zeros(qh, dtype=torch.float32, device=dev) - - fn = self._k_mod.compile( - b=b, - qh=qh, - kh=kh, - sq=t_q, - skv=t_kv, - d_qk=d_qk, - d_v=d_v, - # The Stats layout is a per-shape specialization (like d_qk/d_v): - # has_lse=False compiles the store out; token-major binds the - # packed rank-2 (T, H) view; head-major carries the caller-declared - # head-row stride (0 -> compact sq). - has_lse=lse is not None, - lse_head_major=lse is not None and self.thd_stats_head_major, - lse_head_stride=(self.thd_stats_head_stride if (lse is not None and self.thd_stats_head_major) else 0), - # Declared strides of the bound views (cache-key): compact views - # reproduce the packed specialization; native non-packed views - # compile their strides into the kernel's addressing. - q_stride=tuple(Q.stride()), - k_stride=tuple(K.stride()), - v_stride=tuple(V.stride()), - o_stride=tuple(O.stride()), - ) + # Dummy sinks bound on the launch stream (ordering vs the kernel). + with _torch_stream_context(current_stream, dev): + if carver is not None: + sinks_t = carver.take(qh, torch.float32) + sinks_t.zero_() + else: + sinks_t = torch.zeros(qh, dtype=torch.float32, device=dev) + + # PLAN-TIME-ONLY compile key (issue #552: keying on the packed totals + # degenerated into a per-step recompile under continuous batching): + # this lru-cached call re-binds the artifact compile() already built. + # The K/V strides are taken from the BOUND views because the + # all-KV-zero clamp above swaps in packed batch-1 views (that rare + # shape mints its own cache entry); the batch stride is zeroed out of + # the key (a runtime value the kernel rebuilds symbolically). + 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) self._logger.debug("execute (THD) completed") @@ -2002,10 +2045,14 @@ def compile(self) -> None: ) self._k_mod = _load_sm120_kernel_module(params, fp8=self._fp8) if self.thd: - # The packed token totals (and max sequence length) are runtime - # values, so the per-shape compile is deferred to execute(). - self._compiled_kernel = "thd-deferred" - self._logger.debug("compile completed (THD per-shape compile deferred)") + # The THD compile key is PLAN-TIME-ONLY (the packed token totals + # compile as dynamic extents and max_sq is a runtime launch + # argument — issue #552), so compile HERE like every dense + # specialization; execute()'s lru-cached call re-binds this + # artifact. (The all-KV-zero clamp swaps the K/V strides and + # mints its own entry on first hit.) + self._compiled_kernel = self._k_mod.compile(**self._thd_compile_kwargs()) + self._logger.debug("compile completed (THD, dynamic token extents)") return self._compiled_kernel = self._k_mod.compile( compute_capability=self.compute_capability, @@ -2152,6 +2199,7 @@ def execute( seq_q_lens, seq_kv_lens, cutlass.Float32(scale_softmax_log2), + cutlass.Int32(0), # thd_max_sq: THD-only runtime grid extent current_stream, ) if o_needs_copy_back: @@ -2229,7 +2277,7 @@ def _scalar(t, default=1.0): # scalars, the amax protocol) is identical. pack = None if self.thd: - pack = self._thd_pack(q, k, v, o, seq_q_lens, seq_kv_lens, workspace, "SdpaFwdDslSm120 (FP8 THD)") + pack = self._thd_pack(q, k, v, o, seq_q_lens, seq_kv_lens, workspace, "SdpaFwdDslSm120 (FP8 THD)", current_stream=current_stream) if pack is None: return # This kernel's ragged LSE store is head-major (H, head_stride): @@ -2254,19 +2302,10 @@ def _scalar(t, default=1.0): fn = self._compiled_kernel if pack is not None: - fn = self._k_mod.compile( - compute_capability=self.compute_capability, - b=self.batch_size, - qh=self.h_q, - kh=self.h_kv, - sq=pack.t_q, - skv=pack.t_kv, - d_qk=self.head_dim_qk, - d_v=self.head_dim_v, - max_sq=pack.max_sq, - has_lse=self.lse_desc is not None, - lse_head_stride=self.thd_stats_head_stride, - ) + # PLAN-TIME-ONLY compile key (issue #552): this lru-cached call + # re-binds the artifact compile() already built — the packed + # totals are dynamic extents and max_sq is a launch argument. + fn = self._k_mod.compile(**self._thd_compile_kwargs()) fn( pack.Q if pack is not None else q, pack.K if pack is not None else k, @@ -2280,6 +2319,7 @@ def _scalar(t, default=1.0): cutlass.Float32(scale_softmax_log2), cutlass.Float32(o_scale_fused), cutlass.Float32(ss), + cutlass.Int32(pack.max_sq if pack is not None else 0), current_stream, ) # Both of these consume what the kernel just wrote, so they belong on @@ -2291,15 +2331,55 @@ def _scalar(t, default=1.0): amax_o_buf.div_(max(so, 1e-30)) self._logger.debug("execute (SM120 FP8 per-tensor) completed") - def _thd_pack(self, q_buf, k_buf, v_buf, o_buf, seq_q_lens, seq_kv_lens, workspace, label, declared_views=False): + def _thd_compile_kwargs(self) -> dict: + """The THD compile key — PLAN-TIME-ONLY by contract (issue #552). + + The packed token totals compile as DYNAMIC extents and ``max_sq`` is + a runtime launch argument, so everything here is known when the graph + is built: ``compile()`` compiles eagerly and the execute paths' + lru-cached calls re-bind the same artifact for every packed total.""" + has_lse = self.lse_desc is not None + kwargs = dict( + compute_capability=self.compute_capability, + b=self.batch_size, + qh=self.h_q, + kh=self.h_kv, + d_qk=self.head_dim_qk, + d_v=self.head_dim_v, + has_lse=has_lse, + lse_head_stride=self.thd_stats_head_stride, + ) + if self._fp8: + # The FP8 cell serves only the packed contract (no stride keys) + # and its ragged LSE store is head-major-only. + return kwargs + + def _key(desc): + (ts, hs, es), _ = self._thd_declared(desc) + return (0, ts, hs, es) + + kwargs.update( + lse_head_major=self.thd_stats_head_major, + q_stride=_key(self.q_desc), + k_stride=_key(self.k_desc), + v_stride=_key(self.v_desc), + o_stride=_key(self.o_desc), + ) + return kwargs + + def _thd_pack(self, q_buf, k_buf, v_buf, o_buf, seq_q_lens, seq_kv_lens, workspace, label, declared_views=False, current_stream=None): """Shared THD (ragged) packing: cu_seqlens metadata + ``(1, T, H, D)`` views. Serves the same fully-packed contract as the SM100 THD path (``ragged_offset == cumsum(seq_len) * H * D`` from 0, multiplier 1); the offsets are re-derived from ``seq_len_q``/``seq_len_kv``. The two - ``.tolist()`` D2H syncs are inherent — the packed totals and the - longest sequence's Q length are runtime values that size the - per-execute compile and grid. + ``.tolist()`` D2H syncs feed the metadata upload, the ragged views' + extents and the launch grid (they no longer key any compile — the + kernels compile with DYNAMIC token extents, issue #552; removing the + sync itself needs the plan-time-max grid redesign tracked there). + The torch work (allocation, length + reads, the H2D upload) runs on ``current_stream`` — the LAUNCH + stream — so it is ordered against the kernel that consumes it. Returns ``None`` when the packed Q total is zero (nothing to launch). """ @@ -2316,10 +2396,11 @@ def _thd_pack(self, q_buf, k_buf, v_buf, o_buf, seq_q_lens, seq_kv_lens, workspa # 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. - 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_kv_lens, "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)) + 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_kv_lens, "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] max_sq = max(slq_host) if slq_host else 0 @@ -2356,6 +2437,10 @@ def _view(buf, desc, tokens, heads, d): K = _view(k_buf, self.k_desc, t_kv, kh, d_qk) V = _view(v_buf, self.v_desc, t_kv, kh, d_v) + # The cached seq_q dummy is allocated/zeroed once, on the launch + # stream (first-use ordering vs the kernel that reads it). + with _torch_stream_context(current_stream, dev): + seq_q_dummy = self._dummy("seq_q_lens", dev, lambda: torch.zeros(b, dtype=torch.int32, device=dev)) return SimpleNamespace( meta=meta, t_q=t_q, @@ -2365,7 +2450,7 @@ def _view(buf, desc, tokens, heads, d): K=K, V=V, O=_view(o_buf, self.o_desc, t_q, qh, d_v), - seq_q_dummy=self._dummy("seq_q_lens", dev, lambda: torch.zeros(b, dtype=torch.int32, device=dev)), + seq_q_dummy=seq_q_dummy, ) def _execute_thd( @@ -2379,7 +2464,14 @@ def _execute_thd( within each head row. """ - pack = self._thd_pack(q_buf, k_buf, v_buf, o_buf, seq_q_lens, seq_kv_lens, workspace, "SdpaFwdDslSm120 (THD)", declared_views=True) + # Resolve the launch stream BEFORE packing: the metadata upload inside + # _thd_pack must be ordered against the kernel launch below. + if current_stream is None: + current_stream = cuda.CUstream(torch.cuda.current_stream(q_buf.device).cuda_stream) + + pack = self._thd_pack( + q_buf, k_buf, v_buf, o_buf, seq_q_lens, seq_kv_lens, workspace, "SdpaFwdDslSm120 (THD)", declared_views=True, current_stream=current_stream + ) if pack is None: return @@ -2399,33 +2491,17 @@ def _execute_thd( # token. sinks_t = self._checked_sinks_1d(sinks) if sinks is not None else None - if current_stream is None: - current_stream = cuda.CUstream(torch.cuda.current_stream(q_buf.device).cuda_stream) - import cutlass - fn = self._k_mod.compile( - compute_capability=self.compute_capability, - b=self.batch_size, - qh=self.h_q, - kh=self.h_kv, - sq=pack.t_q, - skv=pack.t_kv, - d_qk=self.head_dim_qk, - d_v=self.head_dim_v, - max_sq=pack.max_sq, - has_lse=self.lse_desc is not None, - lse_head_major=self.thd_stats_head_major, - lse_head_stride=self.thd_stats_head_stride, - # Declared strides of the bound views (cache-key): compact views - # reproduce the packed specialization; native non-packed views - # compile their strides into the Q/O offset math and K/V TMA - # descriptors. - q_stride=tuple(pack.Q.stride()), - k_stride=tuple(pack.K.stride()), - v_stride=tuple(pack.V.stride()), - o_stride=tuple(pack.O.stride()), - ) + # PLAN-TIME-ONLY compile key (issue #552): this lru-cached call + # re-binds the artifact compile() already built. The K/V strides are + # taken from the BOUND views because the all-KV-zero clamp swaps in + # packed batch-1 views (that rare shape mints its own cache entry); + # the batch stride is zeroed out of the key (a runtime value the + # kernel rebuilds symbolically). + kwargs = self._thd_compile_kwargs() + kwargs.update(k_stride=(0, *pack.K.stride()[1:]), v_stride=(0, *pack.V.stride()[1:])) + fn = self._k_mod.compile(**kwargs) fn( pack.Q, pack.K, @@ -2436,6 +2512,7 @@ def _execute_thd( pack.seq_q_dummy, pack.meta, cutlass.Float32(scale_softmax_log2), + cutlass.Int32(pack.max_sq), current_stream, ) 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 fd70b3cf3..3e7d31baa 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py @@ -1938,6 +1938,11 @@ def _host( stream: _cuda_driver.CUstream = None, ) -> None: B, QH, KH, SQ, SKV, _ = problem_size + if cutlass.const_expr(CFG.THD_VARLEN): + # Packed token totals are runtime values (dynamic extents); the + # problem_size slots are 0 by contract. + SQ = q_tensor.shape[1] + SKV = k_tensor.shape[1] # Tensors are [B, S, H, D] with stride_order=(3, 2, 1, 0); D is fastest. # K is split along seq under cga2 — box rows are per-CTA. V is split along d_v @@ -2059,6 +2064,13 @@ def compile( # noqa: A001 THD/varlen: q/k/v/o/lse are PACKED with batch dim 1 ([1,T,H,D]); ``b`` is the LOGICAL batch (sequence count) driving n_batch / metadata + O-desc sizes. + ``sq``/``skv`` are IGNORED under THD — the packed token totals are runtime + values (they change every step under continuous batching), so the token + extents compile DYNAMIC (``cute.sym_int``) and the cache key stays + plan-time-only; callers must not pass them (a stray value would only mint + a redundant cache entry). THD ``q_stride``/... carry a ZERO batch stride + (the real view's batch stride is ``t_q * token_stride``, a runtime value; + the fake rebuilds it symbolically — batch extent is 1, it never steps). ``has_lse=False`` compiles the LSE store out (the kernel specializes on a ``None`` LSE argument) — callers without a Stats output pass no LSE buffer at all. THD Stats layouts: token-major packed rank-2 (T, H) by default @@ -2081,6 +2093,12 @@ def compile( # noqa: A001 if (d_qk * CFG.BPE) % 16 != 0 or (d_v * CFG.BPE_O) % 16 != 0: raise ValueError(f"d128 envelope: d_qk*BPE and d_v*BPE must be 16-byte multiples (TMA global-stride rule); got ({d_qk}, {d_v}) at BPE={CFG.BPE}") _fake_batch = 1 if CFG.THD_VARLEN else b + if CFG.THD_VARLEN: + # Dynamic packed token totals: one symbol per ragged group (Q/O and + # the LSE share t_q; K/V share t_kv), so a new total re-binds the same + # compiled artifact instead of minting a new one (issue #552). + sq = cute.sym_int(divisibility=1) + skv = cute.sym_int(divisibility=1) def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): if stride is None: @@ -2090,6 +2108,10 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): for axis in (1, 2): # seq/head global strides feed TMA: 16-byte rule if (stride[axis] * bpe) % 16 != 0: raise ValueError(f"declared stride {stride} axis {axis} must be a 16-byte multiple at BPE={bpe} (TMA global-stride rule)") + if CFG.THD_VARLEN: + # Batch stride = tokens * token_stride (`_thd_view`'s envelope), + # a runtime value: rebuild it from the dynamic token extent. + return cute.runtime.make_fake_tensor(dtype, shape, (shape[1] * stride[1], stride[1], stride[2], stride[3]), assumed_align=16) return cute.runtime.make_fake_tensor(dtype, shape, tuple(stride), assumed_align=16) fake_q = _fake_bshd((_fake_batch, sq, qh, d_qk), q_stride) @@ -2112,9 +2134,9 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): # the STATIC rank, so the layout is fully encoded in this fake tensor # — no template parameter. if lse_head_major: + # head_stride covering t_q is validated at execute (t_q is a + # runtime value); 0 = compact = the dynamic token total itself. _lse_hs = lse_head_stride if lse_head_stride else sq - if _lse_hs < sq: - raise ValueError(f"THD head-major LSE head_stride ({_lse_hs}) must cover the packed Q token total ({sq})") fake_lse = cute.runtime.make_fake_compact_tensor( cutlass.Float32, (1, qh, _lse_hs), @@ -2189,7 +2211,9 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): fake_sinks, fake_seq_kv_lens, fake_o_desc, - (b, qh, kh, sq, skv, 0), + # THD: the packed totals are runtime values carried by the (dynamic) + # tensor extents — _host reads them from the views' shapes. + (b, qh, kh, 0, 0, 0) if CFG.THD_VARLEN else (b, qh, kh, sq, skv, 0), cutlass.Float32(0.0), cutlass.Int32(0), fake_seq_q_lens, 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 8f2face2b..81693eab8 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 @@ -2033,6 +2033,11 @@ def _host( stream: _cuda_driver.CUstream = None, ) -> None: B, QH, KH, SQ, SKV, _ = problem_size + if cutlass.const_expr(CFG.THD_VARLEN): + # Packed token totals are runtime values (dynamic extents); the + # problem_size slots are 0 by contract. + SQ = q_tensor.shape[1] + SKV = k_tensor.shape[1] # Tensors are [B, S, H, D] with stride_order=(3, 2, 1, 0); D is fastest. # K is split along seq under cga2 — box rows are per-CTA. V is split along d_v @@ -2154,6 +2159,13 @@ def compile( # noqa: A001 THD/varlen: q/k/v/o/lse are PACKED with batch dim 1 ([1,T,H,D]); ``b`` is the LOGICAL batch (sequence count) driving n_batch / metadata + O-desc sizes. + ``sq``/``skv`` are IGNORED under THD — the packed token totals are runtime + values (they change every step under continuous batching), so the token + extents compile DYNAMIC (``cute.sym_int``) and the cache key stays + plan-time-only; callers must not pass them (a stray value would only mint + a redundant cache entry). THD ``q_stride``/... carry a ZERO batch stride + (the real view's batch stride is ``t_q * token_stride``, a runtime value; + the fake rebuilds it symbolically — batch extent is 1, it never steps). ENVELOPE: ``d_qk`` / ``d_v`` are the ACTUAL head dims (defaults = the flavor's full TILE_K / TILE_O). The Q/K/V/O TMA descriptors are built from @@ -2168,6 +2180,12 @@ def compile( # noqa: A001 if (d_qk * CFG.BPE) % 16 != 0 or (d_v * CFG.BPE_O) % 16 != 0: raise ValueError(f"d192 envelope: d_qk*BPE and d_v*BPE must be 16-byte multiples (TMA global-stride rule); got ({d_qk}, {d_v}) at BPE={CFG.BPE}") _fake_batch = 1 if CFG.THD_VARLEN else b + if CFG.THD_VARLEN: + # Dynamic packed token totals: one symbol per ragged group (Q/O and + # the LSE share t_q; K/V share t_kv), so a new total re-binds the same + # compiled artifact instead of minting a new one (issue #552). + sq = cute.sym_int(divisibility=1) + skv = cute.sym_int(divisibility=1) def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): if stride is None: @@ -2177,6 +2195,10 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): for axis in (1, 2): # seq/head global strides feed TMA: 16-byte rule if (stride[axis] * bpe) % 16 != 0: raise ValueError(f"declared stride {stride} axis {axis} must be a 16-byte multiple at BPE={bpe} (TMA global-stride rule)") + if CFG.THD_VARLEN: + # Batch stride = tokens * token_stride (`_thd_view`'s envelope), + # a runtime value: rebuild it from the dynamic token extent. + return cute.runtime.make_fake_tensor(dtype, shape, (shape[1] * stride[1], stride[1], stride[2], stride[3]), assumed_align=16) return cute.runtime.make_fake_tensor(dtype, shape, tuple(stride), assumed_align=16) fake_q = _fake_bshd((_fake_batch, sq, qh, d_qk), q_stride) @@ -2199,9 +2221,9 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): # the STATIC rank, so the layout is fully encoded in this fake tensor # — no template parameter. if lse_head_major: + # head_stride covering t_q is validated at execute (t_q is a + # runtime value); 0 = compact = the dynamic token total itself. _lse_hs = lse_head_stride if lse_head_stride else sq - if _lse_hs < sq: - raise ValueError(f"THD head-major LSE head_stride ({_lse_hs}) must cover the packed Q token total ({sq})") fake_lse = cute.runtime.make_fake_compact_tensor( cutlass.Float32, (1, qh, _lse_hs), @@ -2276,7 +2298,9 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): fake_sinks, fake_seq_kv_lens, fake_o_desc, - (b, qh, kh, sq, skv, 0), + # THD: the packed totals are runtime values carried by the (dynamic) + # tensor extents — _host reads them from the views' shapes. + (b, qh, kh, 0, 0, 0) if CFG.THD_VARLEN else (b, qh, kh, sq, skv, 0), cutlass.Float32(0.0), cutlass.Int32(0), fake_seq_q_lens, 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 36a4295e5..7eef3b540 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py @@ -1606,6 +1606,11 @@ def _host( stream: _cuda_driver.CUstream = None, ) -> None: B, QH, KH, SQ, SKV, _ = problem_size + if cutlass.const_expr(CFG.THD_VARLEN): + # Packed token totals are runtime values (dynamic extents); the + # problem_size slots are 0 by contract. + SQ = q_tensor.shape[1] + SKV = k_tensor.shape[1] _O_GRANU_ELEMS = CFG.O_SWZ_BYTES // CFG.BPE_O qk_box_q = (1, CFG.TILE_M, 1, TMA_QK_GRANU_ELEMS) @@ -1712,12 +1717,25 @@ def compile( # noqa: A001 TILE_K / TILE_O). TMA descriptors carry these extents while the tile box stays the compile-time TILE geometry: loads past d_qk / d_v zero-fill (exact zeros in the QK^T / P·V contractions), O stores past d_v clip. - d * BPE must be a 16-byte multiple (TMA global-stride rule -> d % 8).""" + d * BPE must be a 16-byte multiple (TMA global-stride rule -> d % 8). + + THD/varlen: ``sq``/``skv`` are IGNORED — the packed token totals are + runtime values (they change every step under continuous batching), so the + token extents compile DYNAMIC (``cute.sym_int``) and the cache key stays + plan-time-only; callers must not pass them. THD strides carry a ZERO batch + stride (the real view's batch stride is ``t_q * token_stride``, a runtime + value; the fake rebuilds it symbolically — batch extent 1 never steps).""" if not (0 < d_qk <= CFG.TILE_K and 0 < d_v <= CFG.TILE_O): raise ValueError(f"d256 envelope: need 0 < d_qk <= {CFG.TILE_K} and 0 < d_v <= {CFG.TILE_O}; got ({d_qk}, {d_v})") if (d_qk * CFG.BPE) % 16 != 0 or (d_v * CFG.BPE_O) % 16 != 0: raise ValueError(f"d256 envelope: d_qk*BPE and d_v*BPE must be 16-byte multiples (TMA global-stride rule); got ({d_qk}, {d_v}) at BPE={CFG.BPE}") _fake_batch = 1 if CFG.THD_VARLEN else b + if CFG.THD_VARLEN: + # Dynamic packed token totals: one symbol per ragged group (Q/O and + # the LSE share t_q; K/V share t_kv), so a new total re-binds the same + # compiled artifact instead of minting a new one (issue #552). + sq = cute.sym_int(divisibility=1) + skv = cute.sym_int(divisibility=1) def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): if stride is None: @@ -1727,6 +1745,10 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): for axis in (1, 2): # seq/head global strides feed TMA: 16-byte rule if (stride[axis] * bpe) % 16 != 0: raise ValueError(f"declared stride {stride} axis {axis} must be a 16-byte multiple at BPE={bpe} (TMA global-stride rule)") + if CFG.THD_VARLEN: + # Batch stride = tokens * token_stride (`_thd_view`'s envelope), + # a runtime value: rebuild it from the dynamic token extent. + return cute.runtime.make_fake_tensor(dtype, shape, (shape[1] * stride[1], stride[1], stride[2], stride[3]), assumed_align=16) return cute.runtime.make_fake_tensor(dtype, shape, tuple(stride), assumed_align=16) fake_q = _fake_bshd((_fake_batch, sq, qh, d_qk), q_stride) @@ -1749,9 +1771,9 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): # the STATIC rank, so the layout is fully encoded in this fake tensor # — no template parameter. if lse_head_major: + # head_stride covering t_q is validated at execute (t_q is a + # runtime value); 0 = compact = the dynamic token total itself. _lse_hs = lse_head_stride if lse_head_stride else sq - if _lse_hs < sq: - raise ValueError(f"THD head-major LSE head_stride ({_lse_hs}) must cover the packed Q token total ({sq})") fake_lse = cute.runtime.make_fake_compact_tensor( cutlass.Float32, (1, qh, _lse_hs), @@ -1820,7 +1842,9 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): fake_sinks, fake_seq_kv_lens, fake_o_desc, - (b, qh, kh, sq, skv, 0), + # THD: the packed totals are runtime values carried by the (dynamic) + # tensor extents — _host reads them from the views' shapes. + (b, qh, kh, 0, 0, 0) if CFG.THD_VARLEN else (b, qh, kh, sq, skv, 0), cutlass.Float32(0.0), cutlass.Int32(0), fake_seq_q_lens, 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 374eb9bf2..906d9cc9d 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py @@ -1806,6 +1806,11 @@ def _host( stream: _cuda_driver.CUstream = None, ) -> None: B, QH, KH, SQ, SKV, _ = problem_size + if cutlass.const_expr(CFG.THD_VARLEN): + # Packed token totals are runtime values (dynamic extents); the + # problem_size slots are 0 by contract. + SQ = q_tensor.shape[1] + SKV = k_tensor.shape[1] _O_GRANU_ELEMS = CFG.O_SWZ_BYTES // CFG.BPE_O qk_box_q = (1, CFG.TILE_M, 1, TMA_QK_GRANU_ELEMS) @@ -1914,12 +1919,25 @@ def compile( # noqa: A001 (exact zeros in the QK^T / P·V contractions), O stores past d_v clip. The Q∪V∪O SMEM alias slabs keep their full compile-time extents — only the GMEM descriptor extents change. d * BPE must be a 16-byte multiple - (TMA global-stride rule -> d % 8).""" + (TMA global-stride rule -> d % 8). + + THD/varlen: ``sq``/``skv`` are IGNORED — the packed token totals are + runtime values (they change every step under continuous batching), so the + token extents compile DYNAMIC (``cute.sym_int``) and the cache key stays + plan-time-only; callers must not pass them. THD strides carry a ZERO batch + stride (the real view's batch stride is ``t_q * token_stride``, a runtime + value; the fake rebuilds it symbolically — batch extent 1 never steps).""" if not (0 < d_qk <= CFG.TILE_K and 0 < d_v <= CFG.TILE_O): raise ValueError(f"d512 envelope: need 0 < d_qk <= {CFG.TILE_K} and 0 < d_v <= {CFG.TILE_O}; got ({d_qk}, {d_v})") if (d_qk * CFG.BPE) % 16 != 0 or (d_v * CFG.BPE_O) % 16 != 0: raise ValueError(f"d512 envelope: d_qk*BPE and d_v*BPE must be 16-byte multiples (TMA global-stride rule); got ({d_qk}, {d_v}) at BPE={CFG.BPE}") _fake_batch = 1 if CFG.THD_VARLEN else b + if CFG.THD_VARLEN: + # Dynamic packed token totals: one symbol per ragged group (Q/O and + # the LSE share t_q; K/V share t_kv), so a new total re-binds the same + # compiled artifact instead of minting a new one (issue #552). + sq = cute.sym_int(divisibility=1) + skv = cute.sym_int(divisibility=1) def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): if stride is None: @@ -1929,6 +1947,10 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): for axis in (1, 2): # seq/head global strides feed TMA: 16-byte rule if (stride[axis] * bpe) % 16 != 0: raise ValueError(f"declared stride {stride} axis {axis} must be a 16-byte multiple at BPE={bpe} (TMA global-stride rule)") + if CFG.THD_VARLEN: + # Batch stride = tokens * token_stride (`_thd_view`'s envelope), + # a runtime value: rebuild it from the dynamic token extent. + return cute.runtime.make_fake_tensor(dtype, shape, (shape[1] * stride[1], stride[1], stride[2], stride[3]), assumed_align=16) return cute.runtime.make_fake_tensor(dtype, shape, tuple(stride), assumed_align=16) fake_q = _fake_bshd((_fake_batch, sq, qh, d_qk), q_stride) @@ -1951,9 +1973,9 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): # the STATIC rank, so the layout is fully encoded in this fake tensor # — no template parameter. if lse_head_major: + # head_stride covering t_q is validated at execute (t_q is a + # runtime value); 0 = compact = the dynamic token total itself. _lse_hs = lse_head_stride if lse_head_stride else sq - if _lse_hs < sq: - raise ValueError(f"THD head-major LSE head_stride ({_lse_hs}) must cover the packed Q token total ({sq})") fake_lse = cute.runtime.make_fake_compact_tensor( cutlass.Float32, (1, qh, _lse_hs), @@ -2022,7 +2044,9 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): fake_sinks, fake_seq_kv_lens, fake_o_desc, - (b, qh, kh, sq, skv, 0), + # THD: the packed totals are runtime values carried by the (dynamic) + # tensor extents — _host reads them from the views' shapes. + (b, qh, kh, 0, 0, 0) if CFG.THD_VARLEN else (b, qh, kh, sq, skv, 0), cutlass.Float32(0.0), cutlass.Int32(0), fake_seq_q_lens, diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py b/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py index 72d9f3d34..e3d4f75cb 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py @@ -199,7 +199,6 @@ def __init__( has_sink: bool = False, thd_varlen: bool = False, thd_batch: int = 1, - thd_max_sq: int = 0, thd_lse_head_major: bool = False, head_tile_qk: int = 128, head_tile_v: int = 128, @@ -228,7 +227,6 @@ def __init__( ``[seq_kv(B) | cu_q(B+1) | cu_k(B+1)]`` metadata tensor, and the grid covers ``ceil(thd_max_sq / q_tile)`` tiles per sequence. :param thd_batch: THD only: the real sequence count B. - :param thd_max_sq: THD only: the longest sequence's Q length. :param thd_lse_head_major: THD only: the packed LSE is head-major ``(H, head_stride)`` (FlashAttention's ``softmax_lse`` layout; tokens contiguous within a head, ``head_stride >= T``) instead of the default @@ -246,8 +244,8 @@ def __init__( if out_dtype != in_dtype: raise ValueError("out_dtype must match in_dtype") - if thd_varlen and (thd_batch < 1 or thd_max_sq < 1): - raise ValueError("thd_varlen requires thd_batch >= 1 and thd_max_sq >= 1") + if thd_varlen and thd_batch < 1: + raise ValueError("thd_varlen requires thd_batch >= 1") self.in_dtype = in_dtype self.out_dtype = in_dtype self.is_causal = is_causal @@ -266,7 +264,6 @@ def __init__( self.has_sink = has_sink self.thd_varlen = thd_varlen self.thd_batch = thd_batch - self.thd_max_sq = thd_max_sq self.thd_lse_head_major = thd_lse_head_major self.head_tile_qk = head_tile_qk @@ -1292,6 +1289,7 @@ def __call__( seq_q_lens: cute.Tensor, seq_kv_lens: cute.Tensor, softmax_scale_log2: cutlass.Float32, + thd_max_sq: cutlass.Int32, stream: cuda_driver.CUstream, ) -> None: """Launch the SM120 cutlass FMHA kernel. @@ -1309,6 +1307,9 @@ def __call__( :param seq_q_lens: Per-batch query lengths, or an unused dummy tensor. :param seq_kv_lens: Per-batch key/value lengths, or an unused dummy tensor. :param softmax_scale_log2: ``softmax_scale * log2(e)``. + :param thd_max_sq: THD only: the longest sequence's Q length (a + RUNTIME value — it sizes the per-sequence grid without entering + the compile cache key); 0 / ignored when dense. :param stream: CUDA stream used for the launch. """ head_dim_qk = q.shape[3] @@ -1319,12 +1320,20 @@ def __call__( raise ValueError("runtime V/O head dimensions must round up (by the head-tile granule) to the kernel head_tile_v") if cutlass.const_expr(head_dim_qk % 8 != 0 or head_dim_v % 8 != 0): raise ValueError("head dimensions must be multiples of 8 (TMA 16-byte global-stride rule at 2 B/elem)") + + # THD compiles the token extents DYNAMIC (mode 1 is a symbol, not an + # int), so only statically-known modes can be compared at trace time; + # the adapter builds the ragged views from shared totals, so the + # dynamic seq extents match by construction. + def _static_neq(a, b): + return isinstance(a, int) and isinstance(b, int) and a != b + if cutlass.const_expr( - q.shape[0] != k.shape[0] - or k.shape[:3] != v.shape[:3] - or q.shape[0] != o.shape[0] - or q.shape[1] != o.shape[1] - or q.shape[2] != o.shape[2] + _static_neq(q.shape[0], k.shape[0]) + or any(_static_neq(a, b) for a, b in zip(k.shape[:3], v.shape[:3])) + or _static_neq(q.shape[0], o.shape[0]) + or _static_neq(q.shape[1], o.shape[1]) + or _static_neq(q.shape[2], o.shape[2]) or q.shape[2] % k.shape[2] != 0 ): raise ValueError("runtime Q/K/V/O batch, sequence, or head geometry mismatch") @@ -1337,13 +1346,17 @@ def __call__( ) if cutlass.const_expr(lse is not None): if cutlass.const_expr(self.thd_varlen): + # The packed token total (q.shape[1]) is DYNAMIC under THD, so + # only the static modes are trace-checkable; head_stride >= T + # and the (T, H) extent are validated by the adapter at + # execute, which builds both views from the same total. if cutlass.const_expr(self.thd_lse_head_major): - if cutlass.const_expr(lse.shape[0] != q.shape[2] or lse.shape[1] < q.shape[1]): + if cutlass.const_expr(lse.shape[0] != q.shape[2]): raise ValueError("head-major THD LSE must have shape (H, head_stride) with head_stride >= T") if cutlass.const_expr(lse.stride != (lse.shape[1], 1)): raise ValueError("head-major THD LSE must be compact row-major") else: - if cutlass.const_expr(lse.shape != (q.shape[1], q.shape[2])): + if cutlass.const_expr(lse.shape[1] != q.shape[2]): raise ValueError("THD LSE must have shape (T, H)") if cutlass.const_expr(lse.stride != (q.shape[2], 1)): raise ValueError("THD LSE must be compact token-major") @@ -1421,7 +1434,7 @@ def kv_tma_desc(t, head_dim, head_tile, swizzle, swizzle_chunks, swizzle_chunk_e # batch count (the packed view's batch mode is 1); tiles past a # shorter sequence's length drain without work. grid=( - ceil_div(self.thd_max_sq, self.q_tile) if cutlass.const_expr(self.thd_varlen) else ceil_div(q.shape[1], self.q_tile), + ceil_div(thd_max_sq, cutlass.Int32(self.q_tile)) if cutlass.const_expr(self.thd_varlen) else ceil_div(q.shape[1], self.q_tile), self.thd_batch if cutlass.const_expr(self.thd_varlen) else q.shape[0], q.shape[2], ), @@ -1441,7 +1454,6 @@ def compile( # noqa: A001 skv: int = 128, d_qk: int = 128, d_v: int = 128, - max_sq: int = 0, has_lse: bool = True, lse_head_major: bool = False, lse_head_stride: int = 0, @@ -1455,9 +1467,15 @@ def compile( # noqa: A001 ``d_qk`` is the Q/K head dim (QK^T contraction width) and ``d_v`` the V/O head dim (P@V output width); they are independent, e.g. (192, 128). - THD specializations pack the batch: ``b`` is the real sequence count, - ``sq``/``skv`` are the packed token totals, and ``max_sq`` (the longest - sequence's Q length) sizes the per-sequence grid. + THD specializations pack the batch: ``b`` is the real sequence count and + ``sq``/``skv`` are IGNORED — the packed token totals are runtime values + (they change every step under continuous batching), so the token extents + compile DYNAMIC (``cute.sym_int``) and the cache key stays plan-time-only; + callers must not pass them. ``max_sq`` (the longest sequence's Q length, + which sizes the per-sequence grid) is likewise a RUNTIME ``__call__`` + argument, not a compile parameter. THD strides carry a ZERO batch stride + (the real view's batch stride is ``t * token_stride``, a runtime value; + the fake rebuilds it symbolically — batch extent 1 never steps). ``has_lse=False`` compiles the LSE store out (the kernel specializes on a ``None`` LSE argument) — callers that don't want stats pass no LSE buffer @@ -1479,7 +1497,6 @@ def compile( # noqa: A001 has_sink=PARAMS.has_sink, thd_varlen=PARAMS.thd_varlen, thd_batch=b, - thd_max_sq=max_sq, thd_lse_head_major=lse_head_major, head_tile_qk=round_up_head_tile(d_qk), head_tile_v=round_up_head_tile(d_v), @@ -1487,10 +1504,20 @@ def compile( # noqa: A001 kv_tile=PARAMS.kv_tile, ) fake_batch = 1 if PARAMS.thd_varlen else b + if PARAMS.thd_varlen: + # Dynamic packed token totals: one symbol per ragged group (Q/O and + # the LSE share t_q; K/V share t_kv), so a new total re-binds the same + # compiled artifact instead of minting a new one (issue #552). + sq = cute.sym_int(divisibility=1) + skv = cute.sym_int(divisibility=1) def _fake_bshd(shape, stride): if stride is None: return cute.runtime.make_fake_compact_tensor(STORAGE_DTYPE, shape, stride_order=(3, 2, 1, 0), assumed_align=16) + if PARAMS.thd_varlen: + # Batch stride = tokens * token_stride (`_thd_view`'s envelope), + # a runtime value: rebuild it from the dynamic token extent. + return cute.runtime.make_fake_tensor(STORAGE_DTYPE, shape, (shape[1] * stride[1], stride[1], stride[2], stride[3]), assumed_align=16) return cute.runtime.make_fake_tensor(STORAGE_DTYPE, shape, tuple(stride), assumed_align=16) fake_q = _fake_bshd((fake_batch, sq, qh, d_qk), q_stride) @@ -1544,6 +1571,7 @@ def _fake_bshd(shape, stride): fake_seq_q_lens, fake_seq_kv_lens, cutlass.Float32(1.0), + cutlass.Int32(0), # thd_max_sq: runtime grid extent (THD) cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=False), options="--enable-tvm-ffi", ) diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py b/python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py index 1241db9ae..2723cb625 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_fp8_sm120.py @@ -171,6 +171,11 @@ def is_layout_supported( if len(shape) != 4 or len(stride) != 4: return False _, sequence, heads, head_dim = shape + if not isinstance(sequence, int): + # THD: the packed token total is DYNAMIC; the batch stride is the + # matching symbol (batch extent 1 never steps), so only the + # static inner strides are trace-checkable. + return stride[1:] == (heads * head_dim, head_dim, 1) return stride == ( sequence * heads * head_dim, heads * head_dim, @@ -190,7 +195,6 @@ def __init__( has_sink: bool = False, thd_varlen: bool = False, thd_batch: int = 1, - thd_max_sq: int = 0, head_tile_qk: int = 128, head_tile_v: int = 128, kv_tile: int = SEQ_KV_TILES[0], @@ -215,7 +219,6 @@ def __init__( ``[seq_kv(B) | cu_q(B+1) | cu_k(B+1)]`` metadata tensor, and the grid covers ``ceil(thd_max_sq / q_tile)`` tiles per sequence. :param thd_batch: THD only: the real sequence count B. - :param thd_max_sq: THD only: the longest sequence's Q length. :param head_tile_qk: Q/K head dimension (the QK^T contraction width). Must be a multiple of 32 between 32 and 256, inclusive. :param head_tile_v: V/O head dimension (the P@V output width). Same @@ -230,8 +233,8 @@ def __init__( raise ValueError("fp8 kernel emits Float16 O only") if has_sink: raise ValueError("has_sink is not supported by the fp8 cell (Amax_S semantics)") - if thd_varlen and (thd_batch < 1 or thd_max_sq < 1): - raise ValueError("thd_varlen requires thd_batch >= 1 and thd_max_sq >= 1") + if thd_varlen and thd_batch < 1: + raise ValueError("thd_varlen requires thd_batch >= 1") self.in_dtype = in_dtype self.out_dtype = out_dtype self.is_causal = is_causal @@ -242,7 +245,6 @@ def __init__( self.has_sink = has_sink self.thd_varlen = thd_varlen self.thd_batch = thd_batch - self.thd_max_sq = thd_max_sq self.head_tile_qk = head_tile_qk self.head_tile_v = head_tile_v @@ -1301,6 +1303,7 @@ def __call__( softmax_scale_log2: cutlass.Float32, o_scale_fused: cutlass.Float32, scale_s: cutlass.Float32, + thd_max_sq: cutlass.Int32, stream: cuda_driver.CUstream, ) -> None: """Launch the SM120 per-tensor FP8 FMHA kernel. @@ -1325,12 +1328,20 @@ def __call__( raise ValueError("runtime Q/K head dimensions must match the kernel head_tile_qk") if cutlass.const_expr(head_dim_v != o.shape[3] or head_dim_v != self.head_tile_v): raise ValueError("runtime V/O head dimensions must match the kernel head_tile_v") + + # THD compiles the token extents DYNAMIC (mode 1 is a symbol, not an + # int), so only statically-known modes can be compared at trace time; + # the adapter builds the ragged views from shared totals, so the + # dynamic seq extents match by construction. + def _static_neq(a, b): + return isinstance(a, int) and isinstance(b, int) and a != b + if cutlass.const_expr( - q.shape[0] != k.shape[0] - or k.shape[:3] != v.shape[:3] - or q.shape[0] != o.shape[0] - or q.shape[1] != o.shape[1] - or q.shape[2] != o.shape[2] + _static_neq(q.shape[0], k.shape[0]) + or any(_static_neq(a, b) for a, b in zip(k.shape[:3], v.shape[:3])) + or _static_neq(q.shape[0], o.shape[0]) + or _static_neq(q.shape[1], o.shape[1]) + or _static_neq(q.shape[2], o.shape[2]) or q.shape[2] % k.shape[2] != 0 ): raise ValueError("runtime Q/K/V/O batch, sequence, or head geometry mismatch") @@ -1344,8 +1355,8 @@ def __call__( # only the head extent is pinned. if cutlass.const_expr(len(lse.shape) != 2 or lse.shape[0] != q.shape[2]): raise ValueError("THD LSE must have shape (H, head_stride)") - if cutlass.const_expr(lse.shape[1] < q.shape[1]): - raise ValueError("THD LSE head_stride must cover the packed Q token total") + # head_stride >= T is validated by the adapter at execute: + # the packed total (q.shape[1]) is DYNAMIC under THD. if cutlass.const_expr(lse.stride != (lse.shape[1], 1)): raise ValueError("THD LSE must be head-major with unit token stride") else: @@ -1442,7 +1453,7 @@ def __call__( # batch count (the packed view's batch mode is 1); tiles past a # shorter sequence's length drain without work. grid=( - ceil_div(self.thd_max_sq, self.q_tile) if cutlass.const_expr(self.thd_varlen) else ceil_div(q.shape[1], self.q_tile), + ceil_div(thd_max_sq, cutlass.Int32(self.q_tile)) if cutlass.const_expr(self.thd_varlen) else ceil_div(q.shape[1], self.q_tile), self.thd_batch if cutlass.const_expr(self.thd_varlen) else q.shape[0], q.shape[2], ), @@ -1462,7 +1473,6 @@ def compile( # noqa: A001 skv: int = 128, d_qk: int = 128, d_v: int = 128, - max_sq: int = 0, has_lse: bool = True, lse_head_stride: int = 0, ) -> Callable: @@ -1471,9 +1481,13 @@ def compile( # noqa: A001 ``d_qk`` is the Q/K head dim (QK^T contraction width) and ``d_v`` the V/O head dim (P@V output width); they are independent, e.g. (192, 128). - THD specializations pack the batch: ``b`` is the real sequence count, - ``sq``/``skv`` are the packed token totals, and ``max_sq`` (the longest - sequence's Q length) sizes the per-sequence grid. + THD specializations pack the batch: ``b`` is the real sequence count and + ``sq``/``skv`` are IGNORED — the packed token totals are runtime values + (they change every step under continuous batching), so the token extents + compile DYNAMIC (``cute.sym_int``) and the cache key stays plan-time-only; + callers must not pass them. ``max_sq`` (the longest sequence's Q length, + which sizes the per-sequence grid) is likewise a RUNTIME ``__call__`` + argument, not a compile parameter. ``has_lse=False`` compiles the LSE store out (the kernel specializes on a ``None`` LSE argument) — callers that don't want stats pass no LSE buffer @@ -1495,13 +1509,18 @@ def compile( # noqa: A001 has_sink=PARAMS.has_sink, thd_varlen=PARAMS.thd_varlen, thd_batch=b, - thd_max_sq=max_sq, head_tile_qk=d_qk, head_tile_v=d_v, q_tile=PARAMS.q_tile, kv_tile=PARAMS.kv_tile, ) fake_batch = 1 if PARAMS.thd_varlen else b + if PARAMS.thd_varlen: + # Dynamic packed token totals: one symbol per ragged group (Q/O share + # t_q; K/V share t_kv), so a new total re-binds the same compiled + # artifact instead of minting a new one (issue #552). + sq = cute.sym_int(divisibility=1) + skv = cute.sym_int(divisibility=1) fake_q = cute.runtime.make_fake_compact_tensor( STORAGE_DTYPE, (fake_batch, sq, qh, d_qk), @@ -1580,6 +1599,7 @@ def compile( # noqa: A001 cutlass.Float32(1.0), cutlass.Float32(1.0), cutlass.Float32(1.0), + cutlass.Int32(0), # thd_max_sq: runtime grid extent (THD) 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 a779f72dc..ef2957a9b 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py @@ -1041,6 +1041,58 @@ def test_dsl_sm100_thd_cu_seq_len_zero_lens(): _run_thd_stats_case(seq_lens_q=[64, 32], seq_lens_kv=[0, 0], mask="none", cu_lens=True) +@pytest.mark.L0 +@torch_fork_set_rng(seed=37) +def test_dsl_sm100_thd_compile_key_plan_time_only(): + """Issue #552: the THD compile key carries NO packed totals. + + ``compile()`` builds the one artifact at plan time (the token extents + compile dynamic), and executes with DIFFERENT packed totals re-bind it — + zero ``cute.compile`` calls on the execute path. Keying the compile on + the totals degenerated into a fresh multi-second compile per step under + continuous batching (the totals change every step), and correctness is + checked per total to prove one artifact serves them all. + """ + _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() + # Plan-time compile: no deferred sentinel, the artifact already exists. + assert api._compiled_kernel != "thd-deferred" + info_plan = api._k_mod.compile.cache_info() + + def _run_and_check(seq_lens): + 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) + 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 + + _run_and_check([200, 150]) + _run_and_check([64, 33]) + info_exec = api._k_mod.compile.cache_info() + assert info_exec.misses == info_plan.misses, "a THD execute minted a new kernel compile (runtime data leaked into the compile key)" + assert info_exec.hits >= info_plan.hits + 2 + + _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. diff --git a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py index b6b38e0ca..ef10fbafa 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py @@ -1007,6 +1007,59 @@ def test_dsl_sm120_thd_cu_seq_len_zero_lens(): _run_thd_case(seq_q_lens=[0, 0], seq_kv_lens=[50, 30], check_stats=True, cu_lens=True) +@pytest.mark.L0 +@torch_fork_set_rng(seed=37) +def test_dsl_sm120_thd_compile_key_plan_time_only(): + """Issue #552: the THD compile key carries NO packed totals. + + ``compile()`` builds the one artifact at plan time (the token extents + compile dynamic, ``max_sq`` is a runtime launch argument), and executes + with DIFFERENT packed totals re-bind it — zero ``cute.compile`` calls on + the execute path. Keying the compile on the totals degenerated into a + fresh multi-second compile per step under continuous batching (the + totals change every step), and correctness is checked per total to prove + one artifact serves them all. + """ + _require_dsl() + from cudnn.sdpa.fwd.api_dsl import SdpaFwdDslSm120 + + 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 = SdpaFwdDslSm120(sample_q=q, sample_k=k, sample_v=v, sample_o=o, thd=True) + assert api.check_support() + api.compile() + # Plan-time compile: no deferred sentinel, the artifact already exists. + assert api._compiled_kernel != "thd-deferred" + info_plan = api._k_mod.compile.cache_info() + + def _run_and_check(seq_lens): + 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) + 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 + + _run_and_check([200, 150]) + _run_and_check([64, 33]) + info_exec = api._k_mod.compile.cache_info() + assert info_exec.misses == info_plan.misses, "a THD execute minted a new kernel compile (runtime data leaked into the compile key)" + assert info_exec.hits >= info_plan.hits + 2 + + @pytest.mark.L0 @torch_fork_set_rng(seed=9) def test_dsl_sm120_dense_flex_bhsd_contiguous():