frost(sdpa): derive THD token capacity from the view's element span (fixes #613)

python/cudnn/sdpa/fwd/api_dsl.py

@@ -463,6 +463,38 @@ 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_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 +1536,7 @@ 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))
         if lse_tokens_cap is not None:
             t_q = min(t_q, lse_tokens_cap)
         if t_q == 0:
@@ -1527,9 +1557,7 @@ 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))
         if t_kv == 0:
             # No KV storage at all:
             # every query row is dead — served by the KERNEL's own dead-row
@@ -3002,8 +3030,7 @@ 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.

test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py

@@ -1393,6 +1393,74 @@ def _check(seq_lens):
     _check([64, 33])
 
 
+@pytest.mark.L0
+@torch_fork_set_rng(seed=0)
+def test_dsl_sm100_thd_interleaved_kv_views():
+    """THD K/V bound as strided VIEWS of one fused [T, 2, H, D] record — the
+    layout torch.nn.attention.varlen users produce by slicing a fused KV
+    projection (declared token stride 2*H*D, V at element offset H*D).
+
+    Regression for issue #613: the execute-time token capacity must come from
+    the view's element SPAN. numel()-derived extents HALVE for these views
+    (the TMA descriptors then cut off half the tokens — silently wrong O),
+    and storage-derived extents over-claim into allocator slack, whose NaN
+    bit patterns poison masked rows through P(=0) * V."""
+    _require_dsl()
+    import cudnn
+
+    dev = "cuda"
+    H, d = 8, 128
+    dtype = torch.bfloat16
+    seq_lens = [200, 150, 47]
+    B, S_max, T = len(seq_lens), max(seq_lens), sum(seq_lens)
+    cu = [0]
+    for s_i in seq_lens:
+        cu.append(cu[-1] + s_i)
+    scale = 1.0 / math.sqrt(d)
+
+    q_pk = torch.randn(T, H, d, device=dev, dtype=dtype)
+    kv_rec = torch.randn(T, 2, H, d, device=dev, dtype=dtype)  # fused record
+    k_view, v_view = kv_rec[:, 0], kv_rec[:, 1]  # token stride 2*H*d, V offset H*d
+
+    def _graph_run(k_buf, v_buf, kv_token_stride):
+        q_stride = (S_max * H * d, d, H * d, 1)
+        kv_stride = (S_max * kv_token_stride, d, kv_token_stride, 1)
+        io = cudnn.data_type.BFLOAT16
+        g = cudnn.pygraph(io_data_type=io, intermediate_data_type=cudnn.data_type.FLOAT, compute_data_type=cudnn.data_type.FLOAT)
+        tq = g.tensor(dim=[B, H, S_max, d], stride=list(q_stride), data_type=io, name="q")
+        tk = g.tensor(dim=[B, H, S_max, d], stride=list(kv_stride), data_type=io, name="k")
+        tv = g.tensor(dim=[B, H, S_max, d], stride=list(kv_stride), data_type=io, name="v")
+        slq = torch.tensor(seq_lens, dtype=torch.int32, device=dev).view(B, 1, 1, 1)
+        cu_t = torch.tensor(cu, dtype=torch.int64, device=dev)
+        ro_q = (cu_t * H * d).view(B + 1, 1, 1, 1)
+        ro_kv = (cu_t * kv_token_stride).view(B + 1, 1, 1, 1)
+        sq, skv = g.tensor_like(slq), g.tensor_like(slq)
+        qro, kro, vro, oro = (g.tensor_like(ro_q) for _ in range(4))
+        tq.set_ragged_offset(qro)
+        tk.set_ragged_offset(kro)
+        tv.set_ragged_offset(vro)
+        o, _ = g.sdpa(
+            name="sdpa", q=tq, k=tk, v=tv, generate_stats=False, attn_scale=scale, use_causal_mask=True, use_padding_mask=True, seq_len_q=sq, seq_len_kv=skv
+        )
+        o.set_output(True).set_dim([B, H, S_max, d]).set_stride(list(q_stride))
+        o.set_ragged_offset(oro)
+        g.validate()
+        g.build_operation_graph()
+        g.create_execution_plans([cudnn.heur_mode.A])
+        _select_engine(g, engine_name())
+        g.check_support()
+        g.build_plans()
+        o_buf = torch.zeros(T, H, d, device=dev, dtype=dtype)
+        vp = {tq: q_pk, tk: k_buf, tv: v_buf, o: o_buf, sq: slq, skv: slq, qro: ro_q, kro: ro_kv, vro: ro_kv, oro: ro_q}
+        g.execute(vp, torch.empty(max(g.get_workspace_size(), 1), device=dev, dtype=torch.uint8))
+        torch.cuda.synchronize()
+        return o_buf
+
+    o_views = _graph_run(k_view, v_view, 2 * H * d)
+    o_packed = _graph_run(k_view.contiguous(), v_view.contiguous(), H * d)
+    assert torch.equal(o_views, o_packed), f"interleaved K/V views diverge from packed binding: max|diff|={(o_views - o_packed).abs().max().item()}"
+
+
 @pytest.mark.L1
 @torch_fork_set_rng(seed=40)
 def test_dsl_sm100_thd_kv_zero_capacity_clamp():

test/python/sdpa/random_config.py

@@ -382,8 +382,13 @@ def __call__(self, rng, rng_data_seed, rng_geom_seed=None):
         randoms_.shape_stats = (randoms_.batches, randoms_.h_q, randoms_.s_q, 1)
 
         if randoms_.is_ragged:  # Ideally Q, O, and Stats are all ragged
-            randoms_.stride_q = get_strides_from_layout(randoms_.shape_q, "bshd")
-            randoms_.stride_o = get_strides_from_layout(randoms_.shape_o, "bshd")
+            # Q/O strides stay None: fill_derived_fields (called before
+            # return) draws the seeded per-tensor token gaps there, with the
+            # auto-packed fallbacks for the cu / offset-multiplier forms
+            # (#538) and 1-byte data types (#537). Assigning packed strides
+            # here would bypass that knob for the whole randomized fleet —
+            # exactly how the numel()//token_stride capacity bug (issue #613)
+            # stayed invisible to these sweeps.
             if randoms_.ragged_stats_layout == "head_major":
                 # [h, t] stats: tokens contiguous within a head, heads strided by the whole packed
                 # buffer. This is FlashAttention's / PyTorch varlen's softmax_lse layout; unlike the
@@ -423,8 +428,7 @@ def __call__(self, rng, rng_data_seed, rng_geom_seed=None):
         randoms_.shape_v = (randoms_.batches, randoms_.h_v, randoms_.s_kv, randoms_.d_v)
 
         if randoms_.is_ragged:  # Ideally K ragged and V ragged
-            randoms_.stride_k = get_strides_from_layout(randoms_.shape_k, "bshd")
-            randoms_.stride_v = get_strides_from_layout(randoms_.shape_v, "bshd")
+            pass  # K/V strides stay None -> seeded token gaps in fill_derived_fields
 
         else:
             indices = [0, 1, 2]
@@ -442,6 +446,10 @@ def __call__(self, rng, rng_data_seed, rng_geom_seed=None):
             randoms_.stride_k = get_strides_from_indices(randoms_.shape_k, indices, gaps_k, rng)
             randoms_.stride_v = get_strides_from_indices(randoms_.shape_v, indices, gaps_v, rng)
 
+        # Fill whatever was left None (ragged Q/K/V/O strides) through the
+        # same derivation every ExecConfig takes — one source of truth for
+        # the ragged token-gap draw and its fallbacks.
+        randoms_.fill_derived_fields()
         return randoms_