diff --git a/python/cudnn/AGENTS.md b/python/cudnn/AGENTS.md index a4e84c608..2de22ae4d 100644 --- a/python/cudnn/AGENTS.md +++ b/python/cudnn/AGENTS.md @@ -8,6 +8,31 @@ The `cudnn` Python package: pybind11-backed graph API plus pure-Python **fronten - Never add an eager `import torch` / `import cutlass` to `__init__.py` or anything it imports transitively. `api_base.py` itself imports them at top level, which is why kernel classes must only be reachable through the lazy table. - Reuse the existing `[cutedsl]` extra (`pyproject.toml` optional-dependencies) unless a kernel truly needs a new package. +## Hard rules + +Numbered so reviews can cite them; the list grows — append, never renumber. + +**Rule 1 — `execute()` is a zero-surprise hot path: validate, never convert, never allocate.** + +- **No implicit conversions.** Never `.to(dtype)`, and never a `reshape()` that can + copy, on an execute argument: both silently allocate and launch a kernel per + call, and the fresh pointer breaks CUDA-graph capture. Worse, for an *output* + tensor a reshape copy swallows the kernel's write. Validate dtype / shape / + contiguity and bind a true view (`.view()` or a checked `reshape`), raising + `ValueError` otherwise — see `_checked_lse_view` / `_checked_sinks_1d` / + `_checked_seq_lens` in `sdpa/fwd/api_dsl.py`. +- **No per-execute allocations.** No `torch.empty`/`torch.zeros` inside + `execute()`: scratch is carved from the caller's workspace + (`scratch_workspace_bytes()` contract), and a dead ABI slot may use a + one-time cached dummy (`_dummy`) at most. Prefer compiling the unused + operand out entirely (CuTeDSL specializes on `None` via + `cutlass.const_expr` — see the SM120 SDPA kernel's optional lse/sinks). +- **Init-time flags are compile-time specializations; `execute()` must match + them exactly, in both directions.** A required-but-missing tensor must + raise, never fall back to a zeros dummy (zeros sinks change the softmax + denominator; zeros seq lens mask every row — silently wrong output). A + provided-but-uncompiled tensor must also raise, never be silently ignored. + ## Frontend-only kernel package layout ``` diff --git a/python/cudnn/frost/README.md b/python/cudnn/frost/README.md index e737872ec..ab60ddb66 100644 --- a/python/cudnn/frost/README.md +++ b/python/cudnn/frost/README.md @@ -100,8 +100,10 @@ Reading that sequence line by line: still forwarded to the lowered C++ graph. - **`get_workspace_size()` is honest.** For a python plan it returns `CompiledPlan.get_workspace_size()`, the executor's real requirement (for - the graph above: the dummy-LSE scratch `b*h*s*4 = 16384` bytes, because an - inference graph has no Stats output). `execute()` forwards the caller's + the graph above on an SM100 engine: the dummy-LSE scratch `b*h*s*4 = 16384` + bytes, because an inference graph has no Stats output and the SM100 kernels + always write an LSE; `lse_optional` adapters like SM120 compile the LSE + store out instead and report 0). `execute()` forwards the caller's buffer through `ExecutionContext.workspace` and the executor carves its scratch out of it in 128-byte-aligned chunks, never touching bytes at or beyond the reported size -- no hidden per-execute allocation, stable diff --git a/python/cudnn/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py index 27cc4bb93..c38441dc3 100644 --- a/python/cudnn/sdpa/fwd/api_dsl.py +++ b/python/cudnn/sdpa/fwd/api_dsl.py @@ -289,6 +289,102 @@ def _dummy(self, key: str, device: torch.device, factory: Callable[[], torch.Ten self._dummy_cache[cache_key] = tensor return tensor + def _checked_lse_view(self, lse_tensor: torch.Tensor) -> torch.Tensor: + """Validate a caller-provided LSE buffer and return the kernel's (B, H_q, S_q) view. + + The kernel WRITES through the returned view, so this must be a true + view: a silent ``reshape`` copy of a non-contiguous buffer would + receive the output and be dropped, leaving the caller's LSE unwritten. + """ + self._value_error_if( + lse_tensor.dtype != torch.float32, + f"lse_tensor must be float32; got {lse_tensor.dtype}", + ) + expected = self.batch_size * self.h_q * self.s_q_max + self._value_error_if( + lse_tensor.numel() != expected, + f"lse_tensor must have B*H_q*S_q = {expected} elements; got {lse_tensor.numel()}", + ) + self._value_error_if( + not lse_tensor.is_contiguous(), + "lse_tensor must be contiguous (the kernel writes through this buffer)", + ) + return lse_tensor.view(self.batch_size, self.h_q, self.s_q_max) + + def _checked_sinks_1d(self, sinks: torch.Tensor) -> torch.Tensor: + """Validate caller-provided sink logits and return the kernel's (H_q,) fp32 view. + + Strictly a view: the kernels consume fp32 sinks directly, and an + implicit ``.to(float32)`` here would allocate and launch a cast kernel + on the execute hot path (and break CUDA-graph pointer stability). + """ + self._value_error_if( + sinks.dtype != torch.float32, + f"sinks must be float32; got {sinks.dtype}", + ) + self._value_error_if( + sinks.numel() != self.h_q, + f"sinks must have H_q = {self.h_q} elements; got {sinks.numel()}", + ) + self._value_error_if( + not sinks.is_contiguous(), + "sinks must be contiguous (bound to the kernel as a flat (H_q,) view)", + ) + return sinks.reshape(-1) + + def _checked_seq_lens(self, seq_lens: torch.Tensor, name: str) -> torch.Tensor: + """Validate caller-provided per-batch lengths and return the kernel's (B,) int32 view. + + Strictly a view: an implicit ``.to(torch.int32)`` here would allocate + and launch a cast kernel on the execute hot path (and break CUDA-graph + pointer stability). + """ + self._value_error_if( + seq_lens.dtype != torch.int32, + f"{name} must be int32; got {seq_lens.dtype}", + ) + self._value_error_if( + seq_lens.numel() != self.batch_size, + f"{name} must have B = {self.batch_size} elements; got {seq_lens.numel()}", + ) + self._value_error_if( + not seq_lens.is_contiguous(), + f"{name} must be contiguous (bound to the kernel as a flat (B,) view)", + ) + return seq_lens.reshape(-1) + + def _check_seq_lens_contract(self, seq_q_lens, seq_kv_lens) -> None: + """Reject seq-length tensors inconsistent with the compiled specialization. + + Like sinks, presence is a compile-time specialization: substituting a + zeros dummy for a required tensor masks every row (silently wrong + output), and lengths passed to a specialization compiled without them + are silently ignored. THD is exempt — it always requires both (they + source the packed cu_seqlens metadata). + """ + if self.thd: + self._value_error_if( + seq_q_lens is None or seq_kv_lens is None, + "THD execute requires seq_q_lens and seq_kv_lens", + ) + return + self._value_error_if( + self.seq_kv_lens_present and seq_kv_lens is None, + "seq_kv_lens is required by this compiled specialization", + ) + self._value_error_if( + not self.seq_kv_lens_present and seq_kv_lens is not None, + "this specialization was compiled without per-batch KV lengths; construct the API with seq_kv_lens_present=True", + ) + self._value_error_if( + self.seq_q_lens_present and seq_q_lens is None, + "seq_q_lens is required by this compiled specialization", + ) + self._value_error_if( + not self.seq_q_lens_present and seq_q_lens is not None, + "this specialization was compiled without per-batch Q lengths; construct the API with seq_q_lens_present=True", + ) + @abstractmethod def scratch_workspace_bytes(self) -> int: """Return the per-execution scratch requirement for this implementation.""" @@ -416,6 +512,10 @@ def check_support(self) -> bool: ) self.dtype_o = self.dtype if self.lse_desc is not None: + # THD stats are not plumbed (the kernel's packed (1, H, T) LSE does + # not match cuDNN's ragged Stats contract) — reject the request + # instead of silently never writing the user's LSE. + self._not_implemented_error_if(self.thd, "THD stats/LSE output is not plumbed yet; construct without sample_lse") self._check_dtype(self.lse_desc, torch.float32, name="LSE") self._check_tensor_shape(self.lse_desc, (b, h_qo, s_qo), name="LSE") self._value_error_if(not self.lse_desc.is_contiguous(), "LSE must be contiguous on SM100 DSL") @@ -684,12 +784,39 @@ def execute( scale_val = self.scale_softmax if scale_softmax is None or scale_softmax == 0.0 else float(scale_softmax) scale_softmax_log2 = scale_val * math.log2(math.e) + self._value_error_if( + self.has_sink and sinks is None, + "sinks is required by this compiled specialization", + ) + self._value_error_if( + not self.has_sink and sinks is not None, + "this specialization was compiled without sink support; construct the API with has_sink=True", + ) + self._check_seq_lens_contract(seq_q_lens, seq_kv_lens) self._value_error_if( self.lse_desc is not None and lse_tensor is None, "lse_tensor is required by this compiled specialization", ) - if lse_tensor is None: - lse_tensor = torch.empty((self.batch_size, self.h_q, self.s_q_max), dtype=torch.float32, device=q_tensor.device) + if self.thd: + # The kernel's packed-LSE scratch is api-level workspace; a + # user-facing THD LSE output is not plumbed, so reject rather than + # silently never writing the caller's buffer (check_support already + # rejects thd + sample_lse). + self._not_implemented_error_if(lse_tensor is not None, "THD stats/LSE output is not plumbed yet") + elif lse_tensor is not None: + lse_tensor = self._checked_lse_view(lse_tensor) + else: + # The SM100 kernels always write an LSE (no has_lse specialization + # yet — follow-up): with no Stats output requested the write lands + # in a cached write-only dummy, allocated once per device rather + # than per execute. The FROST dispatch path never reaches this: + # engines.lower_dsl_prefill carves the dummy from the caller's + # workspace instead. + lse_tensor = self._dummy( + "lse", + q_tensor.device, + lambda: torch.empty((self.batch_size, self.h_q, self.s_q_max), dtype=torch.float32, device=q_tensor.device), + ) if self._fp8 and self._pertensor: # Per-tensor FP8 (sdpa_fp8): scalar descales fold into the softmax scale @@ -748,31 +875,22 @@ def execute( device = q_tensor.device sinks_t = ( - sinks.reshape(-1).to(torch.float32) + self._checked_sinks_1d(sinks) if sinks is not None else self._dummy("sinks", device, lambda: torch.zeros(self.h_q, dtype=torch.float32, device=device)) ) seq_kv_t = ( - seq_kv_lens.reshape(-1).to(torch.int32) + self._checked_seq_lens(seq_kv_lens, "seq_kv_lens") if seq_kv_lens is not None else self._dummy("seq_kv", device, lambda: torch.zeros(self.batch_size, dtype=torch.int32, device=device)) ) # Dense padded-Q trim: per-batch Q lengths are their OWN kernel # parameter (compiled in only when seq_q_lens_present — the kernel # signature is specialized on `None`, so the flag-off ABI is - # unchanged). The caller's (B,)-int32 device tensor is bound - # directly: reshape(-1) is a view and .to(torch.int32) a no-op for - # the canonical contiguous int32 input, so the execute hot path - # performs zero allocations/copies and the kernel sees a stable - # pointer (CUDA-graph-capture friendly). - seq_q_t = None - if self.seq_q_lens_present: - if seq_q_lens is None: - raise ValueError("SdpaFwdDsl.execute: seq_q_lens_present requires a seq_len_q tensor") - # Direct bind (no cat, no carve): reshape(-1)/.to(int32) are - # no-op views for the canonical contiguous int32 input, so this - # path needs no workspace scratch at all. - seq_q_t = seq_q_lens.reshape(-1).to(torch.int32) + # unchanged). The caller's (B,)-int32 device tensor is bound directly + # as a validated view — zero allocations/copies on the execute hot + # path, stable pointer (CUDA-graph-capture friendly). + seq_q_t = self._checked_seq_lens(seq_q_lens, "seq_q_lens") if self.seq_q_lens_present else None o_desc_dummy = self._dummy("o_desc", device, lambda: torch.zeros(1, dtype=torch.int64, device=device)) import cutlass @@ -809,18 +927,18 @@ def _execute_thd(self, q_buf, k_buf, v_buf, o_buf, scale_softmax_log2, sinks, se import cutlass dev = q_buf.device - if seq_q_lens is None or seq_len_kv is None: - raise ValueError("THD execute requires seq_len_q and seq_len_kv") - b = seq_q_lens.numel() + slq_v = self._checked_seq_lens(seq_q_lens, "seq_q_lens") + slk_v = self._checked_seq_lens(seq_len_kv, "seq_kv_lens") + b = slq_v.numel() carver = WorkspaceCarver(workspace, self.scratch_workspace_bytes(), "SdpaFwdDslSm100 (THD)") if workspace is not None else None if carver is not None: slq = carver.take(b, torch.int32) - slq.copy_(seq_q_lens.reshape(-1)) + slq.copy_(slq_v) slk = carver.take(b, torch.int32) - slk.copy_(seq_len_kv.reshape(-1)) + slk.copy_(slk_v) else: - slq = seq_q_lens.reshape(-1).to(torch.int32) - slk = seq_len_kv.reshape(-1).to(torch.int32) + slq = slq_v + slk = slk_v # Metadata buffer: [ seq_kv_lens(B) | cu_seqlens_q(B+1) | cu_seqlens_k(B+1) ], # with the cumulative sums built in place (no torch.cat temporaries). meta = carver.take(3 * b + 2, torch.int32) if carver is not None else torch.empty(3 * b + 2, dtype=torch.int32, device=dev) @@ -877,7 +995,7 @@ def _packed(buf, t, h, d): else: LSE = torch.zeros(1, qh, t_q, dtype=torch.float32, device=dev) if sinks is not None: - sinks_t = sinks.reshape(-1).to(torch.float32) + sinks_t = self._checked_sinks_1d(sinks) elif carver is not None: sinks_t = carver.take(qh, torch.float32) sinks_t.zero_() @@ -960,12 +1078,10 @@ def _execute_mxfp8( lse = lse_tensor.reshape(b, h_q, sq) sinks_t = ( - sinks.reshape(-1).to(torch.float32) - if sinks is not None - else self._dummy("sinks", device, lambda: torch.zeros(h_q, dtype=torch.float32, device=device)) + self._checked_sinks_1d(sinks) if sinks is not None else self._dummy("sinks", device, lambda: torch.zeros(h_q, dtype=torch.float32, device=device)) ) seq_kv_t = ( - seq_kv_lens.reshape(-1).to(torch.int32) + self._checked_seq_lens(seq_kv_lens, "seq_kv_lens") if seq_kv_lens is not None else self._dummy("seq_kv", device, lambda: torch.zeros(b, dtype=torch.int32, device=device)) ) @@ -1050,12 +1166,10 @@ def _scalar(t, default=1.0): lse = lse_tensor.reshape(b, h_q, sq) sinks_t = ( - sinks.reshape(-1).to(torch.float32) - if sinks is not None - else self._dummy("sinks", device, lambda: torch.zeros(h_q, dtype=torch.float32, device=device)) + self._checked_sinks_1d(sinks) if sinks is not None else self._dummy("sinks", device, lambda: torch.zeros(h_q, dtype=torch.float32, device=device)) ) seq_kv_t = ( - seq_kv_lens.reshape(-1).to(torch.int32) + self._checked_seq_lens(seq_kv_lens, "seq_kv_lens") if seq_kv_lens is not None else self._dummy("seq_kv", device, lambda: torch.zeros(b, dtype=torch.int32, device=device)) ) @@ -1366,6 +1480,10 @@ def check_support(self) -> bool: self._check_tensor_shape(self.v_desc, (b, h_kv, s_kv, d_q), name="V") self._check_tensor_shape(self.o_desc, (b, h_q, s_q, d_q), name="O") if self.lse_desc is not None: + # THD stats are not plumbed (the kernel's packed (1, H, T) LSE does + # not match cuDNN's ragged Stats contract) — reject the request + # instead of silently never writing the user's LSE. + self._not_implemented_error_if(self.thd, "THD stats/LSE output is not plumbed yet; construct without sample_lse") self._check_dtype(self.lse_desc, torch.float32, name="LSE") self._check_tensor_shape(self.lse_desc, (b, h_q, s_q), name="LSE") self._value_error_if(not self.lse_desc.is_contiguous(), "LSE must be contiguous on SM120 DSL") @@ -1504,6 +1622,9 @@ def compile(self) -> None: sq=self.s_q_max, skv=self.s_k_max, d=self.head_dim, + # No sample_lse -> the LSE store is compiled out; execute() then + # binds no LSE buffer at all (no dummy, no allocation). + has_lse=self.lse_desc is not None, ) self._logger.debug("compile completed") @@ -1533,6 +1654,7 @@ def execute( not self.has_sink and sinks is not None, "this specialization was compiled without sink support; construct the API with has_sink=True", ) + self._check_seq_lens_contract(seq_q_lens, seq_kv_lens) scale_val = self.scale_softmax if scale_softmax is None or scale_softmax == 0.0 else float(scale_softmax) scale_softmax_log2 = scale_val * math.log2(math.e) if self.thd: @@ -1554,20 +1676,14 @@ def execute( self.lse_desc is not None and lse_tensor is None, "lse_tensor is required by this compiled specialization", ) - if lse_tensor is None: - lse_tensor = torch.empty((self.batch_size, self.h_q, self.s_q_max), dtype=torch.float32, device=q_tensor.device) - lse = lse_tensor.reshape(self.batch_size, self.h_q, self.s_q_max) - sinks_t = ( - sinks.reshape(-1).to(torch.float32) - if sinks is not None - else self._dummy( - "sinks", - q_tensor.device, - lambda: torch.zeros(self.h_q, dtype=torch.float32, device=q_tensor.device), - ) + self._value_error_if( + self.lse_desc is None and lse_tensor is not None, + "this specialization was compiled without an LSE output; construct the API with sample_lse", ) + lse = self._checked_lse_view(lse_tensor) if lse_tensor is not None else None + sinks_t = self._checked_sinks_1d(sinks) if sinks is not None else None seq_q_lens = ( - seq_q_lens.reshape(-1).to(torch.int32) + self._checked_seq_lens(seq_q_lens, "seq_q_lens") if seq_q_lens is not None else self._dummy( "seq_q_lens", @@ -1576,7 +1692,7 @@ def execute( ) ) seq_kv_lens = ( - seq_kv_lens.reshape(-1).to(torch.int32) + self._checked_seq_lens(seq_kv_lens, "seq_kv_lens") if seq_kv_lens is not None else self._dummy( "seq_kv_lens", @@ -1627,14 +1743,16 @@ def _execute_thd(self, q_buf, k_buf, v_buf, o_buf, scale_softmax_log2, sinks, se dev = q_buf.device carver = WorkspaceCarver(workspace, self.scratch_workspace_bytes(), "SdpaFwdDslSm120 (THD)") if workspace is not None else None + slq_v = self._checked_seq_lens(seq_q_lens, "seq_q_lens") + slk_v = self._checked_seq_lens(seq_kv_lens, "seq_kv_lens") if carver is not None: slq = carver.take(b, torch.int32) - slq.copy_(seq_q_lens.reshape(-1)) + slq.copy_(slq_v) slk = carver.take(b, torch.int32) - slk.copy_(seq_kv_lens.reshape(-1)) + slk.copy_(slk_v) else: - slq = seq_q_lens.reshape(-1).to(torch.int32) - slk = seq_kv_lens.reshape(-1).to(torch.int32) + slq = slq_v + slk = slk_v # [seq_kv(B) | cu_q(B+1) | cu_k(B+1)] — bound as the kernel's # seq_kv_lens tensor; the leading B words alias the per-sequence KV # lengths so the kernel's existing padded-mask read works unchanged. @@ -1660,15 +1778,10 @@ def _execute_thd(self, q_buf, k_buf, v_buf, o_buf, scale_softmax_log2, sinks, se o_buf.as_strided((t_q * qh * d,), (1,), o_buf.storage_offset()).zero_() return - # Packed dummy LSE (THD stats are not plumbed): carved at the runtime - # t_q, always within the compile-time bound qh * B * S_q_max. - lse = carver.take(qh * t_q, torch.float32) if carver is not None else torch.empty(qh * t_q, dtype=torch.float32, device=dev) - lse = lse.reshape(1, qh, t_q) - if sinks is not None: - sinks_t = sinks.reshape(-1).to(torch.float32) - else: - sinks_t = carver.take(qh, torch.float32) if carver is not None else torch.empty(qh, dtype=torch.float32, device=dev) - sinks_t.zero_() + # THD stats are not plumbed: the kernel compiles with has_lse=False, so + # there is no packed-LSE buffer (dummy or otherwise) to bind. Sinks are + # None-specialized the same way when the graph has no sink token. + sinks_t = self._checked_sinks_1d(sinks) if sinks is not None else None seq_q_dummy = self._dummy( "seq_q_lens", dev, @@ -1693,13 +1806,14 @@ def _packed(buf, tokens): skv=t_kv, d=d, max_sq=max_sq, + has_lse=False, ) fn( _packed(q_buf, t_q), _packed(k_buf, t_kv), _packed(v_buf, t_kv), _packed(o_buf, t_q), - lse, + None, sinks_t, seq_q_dummy, meta, @@ -1709,13 +1823,14 @@ def _packed(buf, tokens): def scratch_workspace_bytes(self) -> int: if self.thd: - # [slq32 | slk32 | meta(seq_kv, cu_q, cu_k) | packed LSE | sinks dummy]. - # The packed LSE is sized for the worst case t_q = B * S_q_max; every - # per-execute carve stays within this bound. - # No O-descriptor chunk: SM120 stores O with plain guarded GMEM stores, - # so THD needs no per-sequence tensor maps. - b, qh = self.batch_size, self.h_q - return 2 * ws_align(b * 4) + ws_align((3 * b + 2) * 4) + ws_align(qh * b * self.s_q_max * 4) + (0 if self.has_sink else ws_align(qh * 4)) + # [slq32 | slk32 | meta(seq_kv, cu_q, cu_k)]. + # No packed-LSE chunk: THD stats are not plumbed and the kernel is + # compiled with has_lse=False, so no LSE buffer exists at all. No + # sinks-dummy chunk either: the kernel None-specializes on sinks. + # No O-descriptor chunk: SM120 stores O with plain guarded GMEM + # stores, so THD needs no per-sequence tensor maps. + b = self.batch_size + return 2 * ws_align(b * 4) + ws_align((3 * b + 2) * 4) return 0 diff --git a/python/cudnn/sdpa/fwd/engines.py b/python/cudnn/sdpa/fwd/engines.py index c0cdbcf35..cff041ed8 100644 --- a/python/cudnn/sdpa/fwd/engines.py +++ b/python/cudnn/sdpa/fwd/engines.py @@ -115,6 +115,11 @@ class Capabilities: padded: bool = False sink: bool = False stats: bool = False + # The adapter accepts lse_tensor=None (its kernel None-specializes the LSE + # store), so a stats-less graph needs no dummy-LSE workspace chunk. Rows + # that keep False (the SM100 flavors) always write an LSE and get a carved + # dummy from lower_dsl_prefill when the graph has no Stats output. + lse_optional: bool = False thd: bool = False # True = the kernel anchors the THD bottom-right diagonal at each sequence's # own (seq_len_q[b], seq_len_kv[b]). Rows that keep False (the SM100 flavors) @@ -426,6 +431,7 @@ def _sm120_spec() -> EngineSpec: padded=True, sink=True, stats=True, + lse_optional=True, padded_stats=True, thd=True, thd_bottom_right=True, @@ -495,6 +501,13 @@ def lower_dsl_prefill( seq_q_t = facts.seq_q_t if facts.padded else None seq_kv_t = facts.seq_kv_t if facts.padded else None + # Mirrors the seq_q_lens_present constructor argument below. Execute + # forwards seq_q only when the compiled specialization consumes it (or THD, + # which sources cu_seqlens from it) — the adapter rejects mismatches, so a + # buffer the FP8/MXFP8 kernels can't honor (dense padded-Q trim is not + # plumbed there — known gap) is dropped here rather than erroring at + # execute. + seq_q_lens_present = facts.padded and not facts.thd and facts.seq_q_t is not None and not (facts.is_mxfp8 or facts.is_fp8) api = api_type( sample_q=ga.tensor_desc_from_ir(facts.q_t, name="q"), sample_k=ga.tensor_desc_from_ir(facts.k_t, name="k"), @@ -510,7 +523,7 @@ def lower_dsl_prefill( # enabled whenever a dense padded graph carries per-batch Q lengths. # THD carries Q lengths via cu_seqlens; the FP8/MXFP8 kernels are not # plumbed (their specs also keep padded_stats=False). - seq_q_lens_present=(facts.padded and not facts.thd and facts.seq_q_t is not None and not (facts.is_mxfp8 or facts.is_fp8)), + seq_q_lens_present=seq_q_lens_present, has_sink=facts.has_sink, thd=facts.thd, dtype_o=facts.dtype_o if (facts.is_mxfp8 or facts.is_fp8) else None, @@ -527,14 +540,20 @@ def lower_dsl_prefill( # buffer is carved from the CALLER's workspace, so its size is fixed here at # build time and recorded on the executor as ``workspace_bytes`` — that # number is what the plan's CompiledPlan.get_workspace_size() reports. - # - dummy LSE (dense, stats absent): the kernel always writes an LSE; - # without a Stats output it lands in b*h_q*s_q fp32 scratch. - # (THD needs no engine-level LSE chunk — the packed THD LSE is part - # of the api-level scratch below.) + # - dummy LSE (dense, stats absent, non-lse_optional adapters): the + # SM100 kernels always write an LSE; without a Stats output it lands + # in b*h_q*s_q fp32 scratch. lse_optional adapters (SM120) compile the + # LSE store out instead and bind no buffer. (THD needs no engine-level + # LSE chunk — the packed THD LSE is part of the api-level scratch + # below.) # - synthesized seq_len_kv (skv_tail_via_padding rows): b int32. # - api-level scratch (api.scratch_workspace_bytes()): the dense padded # [seq_kv|seq_q] combine and the THD metadata/LSE buffers. - dummy_lse_bytes = 0 if (not spec.capabilities.stats or facts.stats_t is not None or facts.thd) else ws_align(facts.b * facts.h_q * facts.s_q * 4) + dummy_lse_bytes = ( + 0 + if (not spec.capabilities.stats or spec.capabilities.lse_optional or facts.stats_t is not None or facts.thd) + else ws_align(facts.b * facts.h_q * facts.s_q * 4) + ) synth_kv_bytes = ws_align(facts.b * 4) if synth_kv_padding else 0 api_scratch_bytes = api.scratch_workspace_bytes() total_workspace_bytes = dummy_lse_bytes + synth_kv_bytes + api_scratch_bytes @@ -593,9 +612,10 @@ def _execute(variant_pack, workspace=None, stream=None): # re-validates so a direct call cannot silently corrupt memory. carver = WorkspaceCarver(workspace, total_workspace_bytes, spec.name) if total_workspace_bytes else None lse_buf = resolved.get(id(binding.stats)) if binding.stats is not None else None - if lse_buf is None and spec.capabilities.stats and not facts.thd: + if lse_buf is None and spec.capabilities.stats and not spec.capabilities.lse_optional and not facts.thd: # Dummy LSE for stats-less dense graphs — carved, not allocated - # (uninitialized is fine: the kernel writes every row). + # (uninitialized is fine: the kernel writes every row). lse_optional + # adapters take lse_tensor=None instead. lse_buf = carver.take(facts.b * facts.h_q * facts.s_q, torch.float32) sinks_buf = resolved.get(id(binding.sink_token)) if binding.sink_token is not None else None seq_kv_buf = resolved.get(id(binding.seq_len_kv)) if binding.seq_len_kv is not None else None @@ -623,7 +643,7 @@ def _execute(variant_pack, workspace=None, stream=None): scale_softmax=facts.scale, sinks=sinks_buf, seq_kv_lens=seq_kv_buf, - seq_q_lens=seq_q_buf, + seq_q_lens=seq_q_buf if (seq_q_lens_present or facts.thd) else None, # Stream from the execute-time handle (raw CUstream int, the # ExecutionContext's stream); None keeps the default stream. current_stream=_cuda_driver.CUstream(stream) if stream is not None else None, diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py b/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py index 3c1664065..33763b4e6 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py @@ -40,7 +40,7 @@ from functools import lru_cache, partial from types import SimpleNamespace -from typing import Callable, Type +from typing import Callable, Optional, Type import cuda.bindings.driver as cuda_driver import cutlass @@ -772,8 +772,8 @@ def kernel( k: cute.Tensor, v: cute.Tensor, o: cute.Tensor, - lse: cute.Tensor, - sinks: cute.Tensor, + lse: Optional[cute.Tensor], + sinks: Optional[cute.Tensor], seq_q_lens: cute.Tensor, seq_kv_lens: cute.Tensor, tma_k_desc: cutlass.GridConstant[cuda.TensorMap], @@ -786,9 +786,10 @@ def kernel( :param k: Key tensor. :param v: Value tensor. :param o: Output tensor. - :param lse: ``(B, H, Sq)`` fp32 log-sum-exp output. - :param sinks: ``(H,)`` fp32 per-Q-head sink logits, or an unused - dummy tensor when the kernel is configured without ``has_sink``. + :param lse: ``(B, H, Sq)`` fp32 log-sum-exp output, or ``None`` to + compile the LSE store out (the DSL specializes on ``None``). + :param sinks: ``(H,)`` fp32 per-Q-head sink logits; ``None`` iff the + kernel is configured without ``has_sink``. :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 tma_k_desc: Tensor map descriptor for K. @@ -1140,25 +1141,26 @@ def kernel( lse_val = -cutlass.Float32.inf row_lse[row_half] = lse_val - if lane % 4 == 0: - lse_arr = cutlass.make_array_view(lse) - for row_half in cutlass.range_constexpr(2): - lse_q_idx = q_seq_idx + q_warp_row0 + (lane // 4) + row_half * 8 - lse_out = cutlass.Float32(row_lse[row_half]) - if cutlass.const_expr(self.thd_varlen): - # Packed (1, H, T) LSE: rows past this sequence's Q - # length belong to the NEXT sequence — never written, - # and there is no padded region to trim. - if lse_q_idx < seqlen_q: - lse_row = lse_arr[0, head_idx, :] - lse_row[q_row_base + lse_q_idx] = lse_out - else: - # Rows at/past this batch's Q length trim to -inf. - if lse_q_idx >= seqlen_q: - lse_out = -cutlass.Float32.inf - if lse_q_idx < q.shape[1]: - lse_row = lse_arr[batch_idx, head_idx, :] - lse_row[lse_q_idx] = lse_out + if cutlass.const_expr(lse is not None): + if lane % 4 == 0: + lse_arr = cutlass.make_array_view(lse) + for row_half in cutlass.range_constexpr(2): + lse_q_idx = q_seq_idx + q_warp_row0 + (lane // 4) + row_half * 8 + lse_out = cutlass.Float32(row_lse[row_half]) + if cutlass.const_expr(self.thd_varlen): + # Packed (1, H, T) LSE: rows past this sequence's Q + # length belong to the NEXT sequence — never written, + # and there is no padded region to trim. + if lse_q_idx < seqlen_q: + lse_row = lse_arr[0, head_idx, :] + lse_row[q_row_base + lse_q_idx] = lse_out + else: + # Rows at/past this batch's Q length trim to -inf. + if lse_q_idx >= seqlen_q: + lse_out = -cutlass.Float32.inf + if lse_q_idx < q.shape[1]: + lse_row = lse_arr[batch_idx, head_idx, :] + lse_row[lse_q_idx] = lse_out prims.barrier_cta_sync(self.bar_compute_sync, thread_count=self.threads_compute) @@ -1237,8 +1239,8 @@ def __call__( k: cute.Tensor, v: cute.Tensor, o: cute.Tensor, - lse: cute.Tensor, - sinks: cute.Tensor, + lse: Optional[cute.Tensor], + sinks: Optional[cute.Tensor], seq_q_lens: cute.Tensor, seq_kv_lens: cute.Tensor, softmax_scale_log2: cutlass.Float32, @@ -1250,9 +1252,10 @@ def __call__( :param k: Key tensor with shape ``(B, Sk, H, D)``. :param v: Value tensor with shape ``(B, Sk, H, D)``. :param o: Output tensor with shape ``(B, Sq, H, D)``. - :param lse: ``(B, H, Sq)`` fp32 log-sum-exp output. - :param sinks: ``(H,)`` fp32 per-Q-head sink logits, or an unused - dummy tensor when the kernel is configured without ``has_sink``. + :param lse: ``(B, H, Sq)`` fp32 log-sum-exp output, or ``None`` to + compile the LSE store out entirely (no dummy buffer needed). + :param sinks: ``(H,)`` fp32 per-Q-head sink logits; must be ``None`` + exactly when the kernel is configured without ``has_sink``. :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)``. @@ -1273,11 +1276,14 @@ def __call__( for name, tensor in (("Q", q), ("K", k), ("V", v), ("O", o)): if cutlass.const_expr(not self.is_layout_supported(tensor.shape, tensor.stride)): raise ValueError(f"{name} must use compact BSHD storage") - if cutlass.const_expr(lse.shape != (q.shape[0], q.shape[2], q.shape[1])): - raise ValueError("LSE must have shape (B, H, Sq)") - if cutlass.const_expr(lse.stride != (q.shape[2] * q.shape[1], q.shape[1], 1)): - raise ValueError("LSE must be compact row-major") - if cutlass.const_expr(self.has_sink and sinks.shape != (q.shape[2],)): + if cutlass.const_expr(lse is not None): + if cutlass.const_expr(lse.shape != (q.shape[0], q.shape[2], q.shape[1])): + raise ValueError("LSE must have shape (B, H, Sq)") + if cutlass.const_expr(lse.stride != (q.shape[2] * q.shape[1], q.shape[1], 1)): + raise ValueError("LSE must be compact row-major") + if cutlass.const_expr(self.has_sink != (sinks is not None)): + raise ValueError("sinks must be provided exactly when the kernel is configured with has_sink") + if cutlass.const_expr(sinks is not None and sinks.shape != (q.shape[2],)): raise ValueError("sinks must have shape (H,)") if cutlass.const_expr(self.thd_varlen): if cutlass.const_expr(q.shape[0] != 1): @@ -1360,12 +1366,17 @@ def compile( # noqa: A001 skv: int = 128, d: int = 128, max_sq: int = 0, + has_lse: bool = True, ) -> Callable: """Compile and cache one architecture-specific compact BSHD shape. 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. + + ``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 + at all instead of a dummy. """ kernel = SM120FusedMultiHeadAttentionForward( @@ -1409,17 +1420,25 @@ def compile( # noqa: A001 stride_order=(3, 2, 1, 0), assumed_align=16, ) - fake_lse = cute.runtime.make_fake_compact_tensor( - cutlass.Float32, - (fake_batch, qh, sq), - stride_order=(2, 1, 0), - assumed_align=4, + fake_lse = ( + cute.runtime.make_fake_compact_tensor( + cutlass.Float32, + (fake_batch, qh, sq), + stride_order=(2, 1, 0), + assumed_align=4, + ) + if has_lse + else None ) - fake_sinks = cute.runtime.make_fake_compact_tensor( - cutlass.Float32, - (qh,), - stride_order=(0,), - assumed_align=4, + fake_sinks = ( + cute.runtime.make_fake_compact_tensor( + cutlass.Float32, + (qh,), + stride_order=(0,), + assumed_align=4, + ) + if PARAMS.has_sink + else None ) fake_seq_q_lens = cute.runtime.make_fake_compact_tensor( cutlass.Int32, diff --git a/python/cudnn/sdpa/graph_analyzer.py b/python/cudnn/sdpa/graph_analyzer.py index d81547011..5b1ce64e2 100644 --- a/python/cudnn/sdpa/graph_analyzer.py +++ b/python/cudnn/sdpa/graph_analyzer.py @@ -350,11 +350,26 @@ def _extract_facts(rec: dict) -> SdpaGraphFacts: seq_q_trim = seq_len_q is not None and not use_padding_mask padded = use_padding_mask and seq_len_kv is not None + # The kernels consume per-batch lengths as int32 directly; there is no + # implicit conversion anywhere on the execute path (it would allocate and + # launch a cast kernel). + for name, t in (("seq_len_q", seq_len_q), ("seq_len_kv", seq_len_kv)): + if t is not None: + t_dtype = _DTYPE_FROM_CUDNN.get(t.get_data_type()) + if t_dtype != torch.int32: + return _invalid(f"{name} must be int32; got {t_dtype}") + sink_token = rec.get("sink_token") if sink_token is not None: sink_dim = tuple(sink_token.get_dim()) if sink_dim != (1, h_q, 1, 1): return _invalid(f"sink_token must be (1, H_q, 1, 1); got {sink_dim}") + # The kernels consume fp32 sink logits directly; there is no implicit + # conversion anywhere on the execute path (it would allocate and + # launch a cast kernel). + sink_dtype = _DTYPE_FROM_CUDNN.get(sink_token.get_data_type()) + if sink_dtype != torch.float32: + return _invalid(f"sink_token must be float32; got {sink_dtype}") generate_stats = rec.get("generate_stats") is_inference = rec.get("is_inference") 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 e328675da..eb0c8cacd 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py @@ -287,6 +287,67 @@ def test_dsl_sm100_sink(dtype, d): torch.testing.assert_close(o, o_ref, atol=5e-2, rtol=3e-2) +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +def test_dsl_sm100_execute_sink_lse_contract(): + """execute() rejects sinks inconsistent with the compiled specialization. + + has_sink is a compile-time specialization: substituting a zeros dummy for + missing sinks would silently change the softmax denominator (a zero sink + logit still contributes exp(0) mass), and sinks passed to a sink-less + kernel would be silently dropped. lse_tensor stays accepted when no + sample_lse was given (the SM100 kernels always write an LSE; the FROST + dispatch hands in workspace scratch), but a requested LSE must be bound. + """ + _require_dsl() + from cudnn.sdpa.fwd.api_dsl import SdpaFwdDslSm100 + + b, h, s, d = 1, 4, 256, 128 + q, k, v = (_bhsd(b, h, s, d, torch.float16) for _ in range(3)) + o = torch.empty_like(q) + lse = torch.empty(b, h, s, dtype=torch.float32, device="cuda") + sink = torch.randn(1, h, 1, 1, dtype=torch.float32, device="cuda") + scale = 1.0 / math.sqrt(d) + + api = SdpaFwdDslSm100(sample_q=q, sample_k=k, sample_v=v, sample_o=o, sample_lse=lse, is_causal=True, has_sink=True) + assert api.check_support() + api.compile() + with pytest.raises(ValueError, match="sinks is required"): + api.execute(q_tensor=q, k_tensor=k, v_tensor=v, o_tensor=o, lse_tensor=lse) + with pytest.raises(ValueError, match="lse_tensor is required"): + api.execute(q_tensor=q, k_tensor=k, v_tensor=v, o_tensor=o, sinks=sink) + # Sinks are consumed as fp32 directly — no implicit cast (which would + # allocate and launch a kernel on the execute hot path). + with pytest.raises(ValueError, match="sinks must be float32"): + api.execute(q_tensor=q, k_tensor=k, v_tensor=v, o_tensor=o, lse_tensor=lse, sinks=sink.to(torch.bfloat16)) + + api = SdpaFwdDslSm100(sample_q=q, sample_k=k, sample_v=v, sample_o=o, is_causal=True) + assert api.check_support() + api.compile() + with pytest.raises(ValueError, match="without sink support"): + api.execute(q_tensor=q, k_tensor=k, v_tensor=v, o_tensor=o, sinks=sink) + # Same contract for per-batch lengths: a specialization compiled without + # them must not silently ignore a provided tensor (nor, the other way, + # substitute a zeros dummy that would mask every row). + seq_kv = torch.full((b,), s, dtype=torch.int32, device="cuda") + with pytest.raises(ValueError, match="without per-batch KV lengths"): + api.execute(q_tensor=q, k_tensor=k, v_tensor=v, o_tensor=o, seq_kv_lens=seq_kv) + with pytest.raises(ValueError, match="without per-batch Q lengths"): + api.execute(q_tensor=q, k_tensor=k, v_tensor=v, o_tensor=o, seq_q_lens=seq_kv) + # No sample_lse and no lse_tensor: the kernel's mandatory LSE write lands + # in a cached dummy — no per-execute allocation, output still correct. + api.execute(q_tensor=q, k_tensor=k, v_tensor=v, o_tensor=o) + torch.cuda.synchronize() + o_ref = _ref_sdpa_full(q, k, v, scale=scale, is_causal=True) + torch.testing.assert_close(o, o_ref, atol=5e-2, rtol=3e-2) + + # THD LSE output is not plumbed (packed (1, H, T) layout != cuDNN's ragged + # Stats contract): requesting it is rejected up front instead of being + # silently ignored. + with pytest.raises(NotImplementedError, match="THD stats/LSE"): + SdpaFwdDslSm100(sample_q=q, sample_k=k, sample_v=v, sample_o=o, sample_lse=lse, thd=True).check_support() + + @pytest.mark.L0 @pytest.mark.parametrize("d", _FLAVORS, ids=_FLAVOR_IDS) @pytest.mark.parametrize("dtype", _DTYPES, ids=_DTYPE_IDS) 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 bac480886..e393ff99a 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py @@ -302,8 +302,9 @@ def _run_dsl_graph( _select_engine(graph, engine_name(arch="sm120")) graph.check_support() graph.build_plans() - lse_bytes = batch * heads * sequence * 4 - expected_workspace = 0 if return_stats else (lse_bytes + 127) // 128 * 128 + # Honest workspace: the SM120 kernel None-specializes the LSE store, so a + # stats-less graph needs no dummy-LSE chunk — dense workspace is always 0. + expected_workspace = 0 assert graph.get_workspace_size() == expected_workspace variant_pack[o] = o_gpu @@ -694,6 +695,69 @@ def test_dsl_sm120_wrapper_lse(): torch.testing.assert_close(lse_tensor, expected_lse, atol=2e-2, rtol=2e-2) +@pytest.mark.L0 +@torch_fork_set_rng(seed=17) +def test_dsl_sm120_execute_contract_mismatches(): + """execute() rejects lse/sinks that contradict the compiled specialization. + + The kernel specializes on sink and LSE presence at compile time, so a + mismatch at execute is a hard error: substituting a zeros sink would + silently change the softmax denominator, and a provided-but-uncompiled + LSE would be silently left unwritten. + """ + + _require_dsl() + from cudnn.sdpa.fwd.api_dsl import SdpaFwdDslSm120 + + q = _bhsd(1, 4, 128, 128, torch.float16) + k = _bhsd(1, 4, 128, 128, torch.float16) + v = _bhsd(1, 4, 128, 128, torch.float16) + o = torch.empty_like(q) + lse = torch.empty(1, 4, 128, dtype=torch.float32, device="cuda") + sinks = torch.randn(1, 4, 1, 1, dtype=torch.float32, device="cuda") + + # Compiled WITH sink + LSE: both must be provided at execute. + api = SdpaFwdDslSm120(sample_q=q, sample_k=k, sample_v=v, sample_o=o, sample_lse=lse, has_sink=True) + assert api.check_support() + api.compile() + with pytest.raises(ValueError, match="sinks is required"): + api.execute(q_tensor=q, k_tensor=k, v_tensor=v, o_tensor=o, lse_tensor=lse) + with pytest.raises(ValueError, match="lse_tensor is required"): + api.execute(q_tensor=q, k_tensor=k, v_tensor=v, o_tensor=o, sinks=sinks) + # Sinks are consumed as fp32 directly — no implicit cast (which would + # allocate and launch a kernel on the execute hot path). + with pytest.raises(ValueError, match="sinks must be float32"): + api.execute(q_tensor=q, k_tensor=k, v_tensor=v, o_tensor=o, lse_tensor=lse, sinks=sinks.to(torch.bfloat16)) + + # Compiled WITHOUT sink or LSE: providing either is rejected, and the + # matching call runs with no LSE buffer anywhere (store compiled out). + api = SdpaFwdDslSm120(sample_q=q, sample_k=k, sample_v=v, sample_o=o) + assert api.check_support() + api.compile() + with pytest.raises(ValueError, match="without sink support"): + api.execute(q_tensor=q, k_tensor=k, v_tensor=v, o_tensor=o, sinks=sinks) + with pytest.raises(ValueError, match="without an LSE output"): + api.execute(q_tensor=q, k_tensor=k, v_tensor=v, o_tensor=o, lse_tensor=lse) + # Same contract for per-batch lengths: a specialization compiled without + # them must not silently ignore a provided tensor (nor, the other way, + # substitute a zeros dummy that would mask every row). + seq_kv = torch.full((1,), 128, dtype=torch.int32, device="cuda") + with pytest.raises(ValueError, match="without per-batch KV lengths"): + api.execute(q_tensor=q, k_tensor=k, v_tensor=v, o_tensor=o, seq_kv_lens=seq_kv) + with pytest.raises(ValueError, match="without per-batch Q lengths"): + api.execute(q_tensor=q, k_tensor=k, v_tensor=v, o_tensor=o, seq_q_lens=seq_kv) + api.execute(q_tensor=q, k_tensor=k, v_tensor=v, o_tensor=o) + torch.cuda.synchronize() + expected = _ref_sdpa_full(q, k, v, scale=1.0 / math.sqrt(128)) + torch.testing.assert_close(o.float(), expected, atol=0.1, rtol=5e-2) + + # THD LSE output is not plumbed (packed (1, H, T) layout != cuDNN's ragged + # Stats contract): requesting it is rejected up front instead of being + # silently ignored. + with pytest.raises(NotImplementedError, match="THD stats/LSE"): + SdpaFwdDslSm120(sample_q=q, sample_k=k, sample_v=v, sample_o=o, sample_lse=lse, thd=True).check_support() + + @pytest.mark.L0 @torch_fork_set_rng(seed=22) def test_dsl_sm120_thd(): diff --git a/test/python/sdpa/frost/test_sdpa_graph_analyzer.py b/test/python/sdpa/frost/test_sdpa_graph_analyzer.py index aa58a51b3..bdab44ac3 100644 --- a/test/python/sdpa/frost/test_sdpa_graph_analyzer.py +++ b/test/python/sdpa/frost/test_sdpa_graph_analyzer.py @@ -230,6 +230,26 @@ def test_probe_accepts_seq_len_q_with_padding_mask(): assert engines.engine_name(512) in _eligible(g) +def test_probe_rejects_non_int32_seq_len(): + """The kernels consume per-batch lengths as int32 directly — no implicit + cast anywhere on the execute path — so an int64 seq_len is ineligible.""" + g = _mk_graph() + q, k, v, dims, strides = _mk_qkv(g) + seq_kv = g.tensor(dim=(B, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT64, name="seq_kv64") + o, _ = g.sdpa( + name="s", + q=q, + k=k, + v=v, + attn_scale=0.1, + is_inference=True, + use_padding_mask=True, + seq_len_kv=seq_kv, + ) + _finish_output(o, dims, strides) + assert not _eligible(g) + + def test_probe_rejects_bottom_right_with_padded_seq_len_q(): # Kernel gap (pre-existing): the BR diagonal is anchored at the global S_q, # so dense padding with per-batch seq_len_q would shift it wrongly. @@ -323,6 +343,20 @@ def test_probe_rejects_bad_sink_shape(): assert not _eligible(g) +def test_probe_rejects_non_fp32_sink(): + """The kernels consume fp32 sink logits directly — no implicit cast anywhere + on the execute path — so a non-fp32 sink token is ineligible up front.""" + g = _mk_graph() + q, k, v, dims, strides = _mk_qkv(g) + bad_sink = g.tensor(dim=(1, H, 1, 1), stride=(H, 1, 1, 1), data_type=cudnn.data_type.BFLOAT16, name="bf16sink") + try: + o, _ = g.sdpa(name="s", q=q, k=k, v=v, attn_scale=0.1, is_inference=True, sink_token=bad_sink) + except TypeError: + pytest.skip("this cuDNN wheel's sdpa() binding predates sink_token") + _finish_output(o, dims, strides) + assert not _eligible(g) + + def test_resolve_causal_plus_swa(): g = _mk_graph() q, k, v, dims, strides = _mk_qkv(g)