diff --git a/python/cudnn/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py index 46b8dd449..bdb01c968 100644 --- a/python/cudnn/sdpa/fwd/api_dsl.py +++ b/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. 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 7504c7faf..23db76ed0 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py +++ b/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(): diff --git a/test/python/sdpa/random_config.py b/test/python/sdpa/random_config.py index 681fb8060..d4f638fb3 100644 --- a/test/python/sdpa/random_config.py +++ b/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_