feat(sdpa): accept max_total_seq_len_q/kv on the forward SDPA node

docs/operations/Attention.md

@@ -72,6 +72,7 @@ The support matrix is based on the latest cudnn backend version 9.18.1
 
    **Ragged Offset Multiplier (cuDNN 9.24+, UNIFIED forward only):**
 - `tensor.set_ragged_offset_multiplier(value)` lets the ragged offsets be stored in coarser units; the engine multiplies each offset by `value` to recover element offsets.
+- `max_total_seq_len_q` / `max_total_seq_len_kv` declare the **packed token totals** of the ragged Q and K/V. A ragged tensor's dims stay `(B, H, S_max, D)` and the per-sequence starts live in a device-side offset tensor, so the packed total is not otherwise expressible in the graph. Supplying it lets the implementation bound the token axis exactly rather than inferring an upper bound from the bound buffers' extents — which matters when a buffer is allocated larger than the tokens it holds, since rows past the real total are masked but still take part in `P @ V` and so must be finite. The values only ever tighten the inferred bound, never widen it, and are accepted only on a ragged layout. `sdpa_backward` has taken the same two arguments since cuDNN 9.6.
 - Example: with a multiplier of $H \times D$, a token-unit cumulative-sequence-length tensor (e.g. `cu_seq_len_q`) can be bound directly as the ragged offset, avoiding a conversion pass.
 
    **Memory Layout visualization:**
@@ -305,6 +306,8 @@ graph.sdpa(
     paged_attention_k_table=None,         # Page table for K container
     paged_attention_v_table=None,         # Page table for V container
     paged_attention_max_seq_len_kv=None,  # Max KV sequence length for paged attention
+    max_total_seq_len_q=None,             # Packed token total for Q (ragged tensors)
+    max_total_seq_len_kv=None,            # Packed token total for KV (ragged tensors)
     generate_stats=None,                  # Output softmax stats for training (True/False)
     implementation=AUTO,                  # SDPA implementation: AUTO, COMPOSITE, UNIFIED
     unfuse_fma=False,                     # Use unfused mul/add in the softmax computation

include/cudnn_frontend/graph_properties.h

@@ -2004,6 +2004,16 @@ class SDPA_attributes : public Attributes<SDPA_attributes> {
     std::optional<float> dropout_probability;
     std::optional<float> attn_scale_value;
     std::optional<int> max_seq_len_kv;
+
+    // Packed (ragged) token totals, mirroring SDPA_backward_attributes. A
+    // frontend-side hint only: never lowered to a backend attribute. Ragged
+    // layouts describe extents as (B, H, S_max, D) plus a device ragged-offset
+    // tensor, so the packed total is not otherwise expressible -- consumers
+    // that must bound the token axis have to infer it from buffer geometry
+    // instead. See docs/operations/Attention.md.
+    std::optional<int64_t> max_total_seq_len_q;
+    std::optional<int64_t> max_total_seq_len_kv;
+
     AttentionScoreModifier_t attention_score_modifier = nullptr;
     DataType_t mma_core_mode                          = DataType_t::NOT_SET;
 
@@ -2083,6 +2093,8 @@ class SDPA_attributes : public Attributes<SDPA_attributes> {
                                    dropout_probability,
                                    attn_scale_value,
                                    max_seq_len_kv,
+                                   max_total_seq_len_q,
+                                   max_total_seq_len_kv,
                                    mma_core_mode,
                                    left_bound,
                                    right_bound,
@@ -2124,6 +2136,21 @@ class SDPA_attributes : public Attributes<SDPA_attributes> {
         return *this;
     }
 
+    // Packed token total of the ragged Q (and O/Stats, which share its token
+    // axis). Only meaningful on a ragged/packed layout; ignored otherwise.
+    SDPA_attributes&
+    set_max_total_seq_len_q(int64_t const value) {
+        max_total_seq_len_q = value;
+        return *this;
+    }
+
+    // Packed token total of the ragged K/V.
+    SDPA_attributes&
+    set_max_total_seq_len_kv(int64_t const value) {
+        max_total_seq_len_kv = value;
+        return *this;
+    }
+
     SDPA_attributes&
     set_bias(std::shared_ptr<Tensor_attributes> value) {
         inputs[SDPA_attributes::input_names::Bias] = std::move(value);

include/cudnn_frontend/node/scaled_dot_product_flash_attention.h

@@ -535,6 +535,19 @@ class SDPANodeBase : public NodeCRTP<DerivedT> {
                 "tensor.");
         }
 
+        // validate options for max_total_seq_len (mirrors SDPA_backward_attributes)
+        {
+            bool const is_ragged = attributes.inputs.at(input_names::Q)->get_ragged_offset() ||
+                                   attributes.inputs.at(input_names::K)->get_ragged_offset() ||
+                                   attributes.inputs.at(input_names::V)->get_ragged_offset() ||
+                                   attributes.outputs.at(output_names::O)->get_ragged_offset();
+            RETURN_CUDNN_FRONTEND_ERROR_IF(
+                (attributes.max_total_seq_len_q.has_value() || attributes.max_total_seq_len_kv.has_value()) &&
+                    !is_ragged,
+                error_code_t::GRAPH_NOT_SUPPORTED,
+                "max_total_seq_len_q/kv is only supported with packed (ragged) layout");
+        }
+
 #undef CUDNN_FE_VALIDATE_STRIDE
 
         return {error_code_t::OK, ""};

python/cudnn/sdpa/fwd/api_dsl.py

@@ -293,6 +293,8 @@ def __init__(
         cu_seq_kv_lens: bool = False,
         has_sink: bool = False,
         thd: bool = False,
+        max_total_seq_len_q: Optional[int] = None,
+        max_total_seq_len_kv: Optional[int] = None,
         dtype_o: Optional[torch.dtype] = None,
         pertensor_fp8: bool = False,
         sched_policy: Optional[int] = None,
@@ -352,6 +354,14 @@ def __init__(
         self.cu_seq_kv_lens = bool(cu_seq_kv_lens)
         self.has_sink = bool(has_sink)
         self.thd = bool(thd)
+        # Caller-declared packed token totals. These only ever TIGHTEN the
+        # execute-time token extents (they are min'd against the capacity the
+        # bound buffers can address), so a wrong or stale value cannot make a
+        # launch address memory the caller does not own -- it can only make it
+        # address less. None = not declared; the extent falls back to the
+        # buffer-derived capacity.
+        self.max_total_seq_len_q = None if max_total_seq_len_q is None else int(max_total_seq_len_q)
+        self.max_total_seq_len_kv = None if max_total_seq_len_kv is None else int(max_total_seq_len_kv)
         # MXFP8: FP8 (E4M3/E5M2) Q/K/V in, half (BF16/FP16) O out. dtype_o overrides
         # the output dtype; None inherits Q's dtype. _fp8 is set in check_support once
         # Q's dtype is known.
@@ -463,6 +473,56 @@ def _thd_check_strides_packed(self) -> None:
                 f"{desc.name}: non-packed THD strides {tuple(desc.stride)} are not supported by the FP8 path yet",
             )
 
+    def _thd_capacity(self, buf: torch.Tensor, desc: TensorDesc, packed: bool = False) -> int:
+        """Token CAPACITY of a THD buffer under the strides the view will
+        bind: the largest T whose final token's ROW still fits inside the
+        buffer's own element SPAN (``1 + sum((size_i - 1) * stride_i)``).
+
+        Why the span, and not numel or the untyped storage (issue #613):
+
+        - ``numel() // token_stride`` halves non-packed VIEWS — a K/V slice
+          of a kv-interleaved ``[T, 2, H, D]`` record holds T tokens but only
+          ``T*H*D`` of the record's elements — silently truncating the TMA
+          extent (half the tokens never load).
+        - The untyped storage over-claims into ALLOCATOR SLACK. That is not
+          benign: rows between the real packed total and the extent are
+          masked but still multiplied (``P == 0`` times V), so they must be
+          FINITE — TMA zero-fill only covers rows at or beyond the extent.
+          A slack row carrying NaN bit patterns poisons whole sequences
+          through ``0 * NaN``.
+
+        The span is exact on both edges: every row below the returned
+        capacity lies fully inside caller-provided (finite) elements, and
+        every row at or beyond it is TMA-clipped to zeros."""
+        h, d = desc.shape[1], desc.shape[3]
+        if packed:
+            ts, hs, es = h * d, d, 1
+        else:
+            (ts, hs, es), _ = self._thd_declared(desc)
+        if buf.numel() == 0:
+            return 0
+        span = 1 + sum((size - 1) * stride for size, stride in zip(buf.shape, buf.stride()))
+        row = (h - 1) * hs + (d - 1) * es + 1
+        return 0 if span < row else (span - row) // ts + 1
+
+    def _thd_declared_total(self, cap: int, declared: Optional[int]) -> int:
+        """Tighten a buffer-derived token capacity with the caller's declared
+        packed total (``sdpa(max_total_seq_len_q/kv=...)``).
+
+        A ragged graph declares ``(B, H, S_max, D)`` plus device ragged
+        offsets, so the packed total is not expressible as a dim and
+        ``_thd_capacity`` has to infer an upper bound from buffer geometry.
+        That bound is safe but loose: rows between the real total and the
+        capacity are masked, yet still multiplied (``P == 0`` times V), so
+        they must be finite -- an over-allocated buffer whose tail was never
+        written is a hazard (issue #624). Declaring the total makes the TMA
+        extent exact, which puts that tail out of reach entirely.
+
+        Always a MIN: the declaration can only tighten the capacity, never
+        exceed it, so a stale or wrong value cannot push an access outside the
+        caller's own allocation."""
+        return cap if declared is None else min(cap, max(int(declared), 0))
+
     def _thd_view(self, buf: torch.Tensor, desc: TensorDesc, tokens: int) -> torch.Tensor:
         """The declared-stride ``(1, T, H, D)`` view over a THD buffer's storage.
 
@@ -1504,9 +1564,8 @@ def _thd_pack(self, q_buf, k_buf, v_buf, o_buf, sinks, seq_kv_lens, seq_q_lens,
         kv_lens_dev = self._checked_cu_seq_lens(seq_kv_lens, "cu_seq_len_kv") if self.cu_seq_kv_lens else self._checked_seq_lens(seq_kv_lens, "seq_kv_lens")
         lens_form = (1 if self.cu_seq_q_lens else 0) | (2 if self.cu_seq_kv_lens else 0)
 
-        (q_ts, _, _), _ = self._thd_declared(self.q_desc)
-        (o_ts, _, _), _ = self._thd_declared(self.o_desc)
-        t_q = min(q_buf.numel() // q_ts, o_buf.numel() // o_ts)
+        t_q = min(self._thd_capacity(q_buf, self.q_desc), self._thd_capacity(o_buf, self.o_desc))
+        t_q = self._thd_declared_total(t_q, self.max_total_seq_len_q)
         if lse_tokens_cap is not None:
             t_q = min(t_q, lse_tokens_cap)
         if t_q == 0:
@@ -1527,9 +1586,8 @@ def _thd_pack(self, q_buf, k_buf, v_buf, o_buf, sinks, seq_kv_lens, seq_q_lens,
 
         Q = self._thd_view(q_buf, self.q_desc, t_q)
         O = self._thd_view(o_buf, self.o_desc, t_q)
-        (k_ts, _, _), _ = self._thd_declared(self.k_desc)
-        (v_ts, _, _), _ = self._thd_declared(self.v_desc)
-        t_kv = min(k_buf.numel() // k_ts, v_buf.numel() // v_ts)
+        t_kv = min(self._thd_capacity(k_buf, self.k_desc), self._thd_capacity(v_buf, self.v_desc))
+        t_kv = self._thd_declared_total(t_kv, self.max_total_seq_len_kv)
         if t_kv == 0:
             # No KV storage at all:
             # every query row is dead — served by the KERNEL's own dead-row
@@ -3002,15 +3060,16 @@ def _thd_pack(self, q_buf, k_buf, v_buf, o_buf, seq_q_lens, seq_kv_lens, workspa
         def _cap(buf, desc, heads, d):
             # Token CAPACITY under the strides the view will bind: declared
             # (f16, TMA-expressible by check_support) or packed (FP8).
-            ts = self._thd_declared(desc)[0][0] if declared_views else heads * d
-            return buf.numel() // ts
+            return self._thd_capacity(buf, desc, packed=not declared_views)
 
         # Q/O (and a token-major LSE) bind ONE dynamic token symbol; K/V the
         # other — shared floors.
         t_q = min(_cap(q_buf, self.q_desc, qh, d_qk), _cap(o_buf, self.o_desc, qh, d_v))
+        t_q = self._thd_declared_total(t_q, self.max_total_seq_len_q)
         if lse_tokens_cap is not None:
             t_q = min(t_q, lse_tokens_cap)
         t_kv = min(_cap(k_buf, self.k_desc, kh, d_qk), _cap(v_buf, self.v_desc, kh, d_v))
+        t_kv = self._thd_declared_total(t_kv, self.max_total_seq_len_kv)
 
         if t_q == 0:
             return None

python/cudnn/sdpa/fwd/engines.py

@@ -841,6 +841,12 @@ def lower_dsl_prefill(
         cu_seq_kv_lens=facts.cu_seq_kv_t is not None,
         has_sink=facts.has_sink,
         thd=facts.thd,
+        # Caller-declared packed token totals (issue #624): when present the
+        # adapter binds EXACT token extents instead of the buffer-derived
+        # capacity, putting an over-allocated buffer's uninitialized tail out
+        # of TMA reach. Only ever tightens (see _thd_declared_total).
+        max_total_seq_len_q=facts.max_total_seq_len_q,
+        max_total_seq_len_kv=facts.max_total_seq_len_kv,
         dtype_o=facts.dtype_o if (facts.is_mxfp8 or facts.is_fp8) else None,
         pertensor_fp8=facts.is_fp8,
         sched_policy=knobs.sched_policy if knobs is not None else None,

python/cudnn/sdpa/graph_analyzer.py

@@ -232,6 +232,13 @@ class SdpaGraphFacts:
 
     padded: bool = False  # per-batch KV lengths present (padding mask or THD)
     thd: bool = False  # ragged (THD) Q/K/V
+    # Caller-declared packed token totals (sdpa(max_total_seq_len_q/kv=...)).
+    # Ragged graphs describe extents as (B, H, S_max, D) plus device ragged
+    # offsets, so the packed total is not otherwise expressible; when supplied
+    # it bounds the token axis EXACTLY instead of being inferred from buffer
+    # geometry. None = not declared (infer).
+    max_total_seq_len_q: Optional[int] = None
+    max_total_seq_len_kv: Optional[int] = None
     # cu_seq_len_q / cu_seq_len_kv (cuDNN 9.24+): (B+1,) prefix sums, a
     # contract of their own — neither seq_len_* nor ragged_offset. A fact,
     # not a verdict: engines that don't consume the form must decline (reading
@@ -587,6 +594,8 @@ def _square_transposed(dim: tuple, stride: tuple) -> bool:
         seq_q_trim=seq_q_trim,
         padded=padded,
         thd=thd,
+        max_total_seq_len_q=rec.get("max_total_seq_len_q"),
+        max_total_seq_len_kv=rec.get("max_total_seq_len_kv"),
         has_cu_seq_len=has_cu_seq_len,
         has_sink=sink_token is not None,
         wants_stats=wants_stats,

python/pygraph/pygraph.h

@@ -452,6 +452,8 @@ class PyGraph {
          std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& paged_attention_k_table,
          std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& paged_attention_v_table,
          py::object const& paged_attention_max_seq_len_kv,
+         py::object const& max_total_seq_len_q,
+         py::object const& max_total_seq_len_kv,
          cudnn_frontend::DataType_t const& compute_data_type,
          std::string const& name,
          std::optional<PyCallback> fn,
@@ -526,6 +528,8 @@ class PyGraph {
              std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& paged_attention_k_table,
              std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& paged_attention_v_table,
              py::object const& paged_attention_max_seq_len_kv,
+             py::object const& max_total_seq_len_q,
+             py::object const& max_total_seq_len_kv,
              cudnn_frontend::DataType_t const& compute_data_type,
              std::string const& name,
              std::optional<PyCallback> fn,
@@ -562,6 +566,8 @@ class PyGraph {
                bool const use_padding_mask,
                std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& seq_len_q,
                std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& seq_len_kv,
+               py::object const& max_total_seq_len_q,
+               py::object const& max_total_seq_len_kv,
                std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& cu_seq_len_q,
                std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& cu_seq_len_kv);
 
@@ -850,6 +856,8 @@ class PyGraph {
                   std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& paged_attention_k_table,
                   std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& paged_attention_v_table,
                   py::object const& paged_attention_max_seq_len_kv,
+                  py::object const& max_total_seq_len_q,
+                  py::object const& max_total_seq_len_kv,
                   cudnn_frontend::DataType_t const& compute_data_type,
                   std::string const& name,
                   std::optional<PyCallback> fn,