diff --git a/python/cudnn/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py index 73785ea3c..3eb8e44e8 100644 --- a/python/cudnn/sdpa/fwd/api_dsl.py +++ b/python/cudnn/sdpa/fwd/api_dsl.py @@ -1470,6 +1470,8 @@ def _initialize_implementation(self) -> None: self.compute_capability: Optional[tuple[int, int]] = None self.head_dim_qk: Optional[int] = None self.head_dim_v: Optional[int] = None + self.thd_stats_head_major = False + self.thd_stats_head_stride = 0 self._k_mod = None def check_support(self) -> bool: @@ -1523,13 +1525,21 @@ def check_support(self) -> bool: self._check_tensor_shape(self.v_desc, (b, h_kv, s_kv, d_v), name="V") self._check_tensor_shape(self.o_desc, (b, h_q, s_q, d_v), 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") + if self.thd: + stride_h, stride_s = tuple(self.lse_desc.stride[1:]) + token_major = (stride_h, stride_s) == (1, h_q) + head_major = not token_major and stride_s == 1 and stride_h >= 1 + self._value_error_if( + not token_major and not head_major, + f"THD LSE must be packed token-major (stride_h == 1, stride_s == H) " + f"or head-major (stride_s == 1, stride_h == head_stride); got stride {self.lse_desc.stride}", + ) + self.thd_stats_head_major = head_major + self.thd_stats_head_stride = int(stride_h) if head_major else 0 + else: + self._value_error_if(not self.lse_desc.is_contiguous(), "LSE must be contiguous on SM120 DSL") for label, val in ( ("B", b), @@ -1704,10 +1714,17 @@ def execute( "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", + ) + 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", + ) 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: - self._not_implemented_error_if(lse_tensor is not None, "THD stats/LSE output is not plumbed yet") self._execute_thd( q_tensor, k_tensor, @@ -1717,18 +1734,11 @@ def execute( sinks, seq_kv_lens, seq_q_lens, + lse_tensor=lse_tensor, workspace=workspace, current_stream=current_stream, ) return - self._value_error_if( - self.lse_desc is not None and lse_tensor is None, - "lse_tensor is required by this compiled specialization", - ) - 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 = ( @@ -1777,7 +1787,9 @@ def execute( if o_needs_copy_back: o_view.copy_(o_scratch) - def _execute_thd(self, q_buf, k_buf, v_buf, o_buf, scale_softmax_log2, sinks, seq_kv_lens, seq_q_lens, workspace=None, current_stream=None): + def _execute_thd( + self, q_buf, k_buf, v_buf, o_buf, scale_softmax_log2, sinks, seq_kv_lens, seq_q_lens, lse_tensor=None, workspace=None, current_stream=None + ): """THD (ragged) execute: packed ``(1, T, H, D)`` views + cu_seqlens. Serves the same fully-packed contract as the SM100 THD path @@ -1785,7 +1797,10 @@ def _execute_thd(self, q_buf, k_buf, v_buf, o_buf, scale_softmax_log2, sinks, se 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. + per-execute compile and grid. ``lse_tensor``, when given, is the + caller's ragged Stats buffer, in its declared layout: token-major + packed ``(T, H)`` in the first ``T*H`` elements, or head-major + ``(H, head_stride)`` with tokens contiguous within each head row. """ b, qh, kh = self.batch_size, self.h_q, self.h_kv @@ -1822,15 +1837,36 @@ def _execute_thd(self, q_buf, k_buf, v_buf, o_buf, scale_softmax_log2, sinks, se if t_q == 0: return + lse = None + lse_valid = None # the valid-region view (first t_q tokens) the t_kv == 0 fill writes + if lse_tensor is not None: + if self.thd_stats_head_major: + head_stride = self.thd_stats_head_stride + self._value_error_if( + head_stride < t_q, + f"head-major THD LSE head_stride ({head_stride}) must cover the packed Q token total ({t_q})", + ) + lse = lse_tensor.as_strided((qh, head_stride), (head_stride, 1), lse_tensor.storage_offset()) + lse_valid = lse_tensor.as_strided((qh, t_q), (head_stride, 1), lse_tensor.storage_offset()) + else: + lse = lse_tensor.as_strided((t_q, qh), (qh, 1), lse_tensor.storage_offset()) + lse_valid = lse if t_kv == 0: - # Every row is dead: O := 0 (THD stats are not plumbed). A - # zero-token K/V view cannot back a TMA descriptor, so short-cut. + # Every row is dead: O := 0, LSE := -inf (or the sink alone — + # its column keeps the denominator alive). A zero-token K/V view + # cannot back a TMA descriptor, so short-cut both. o_buf.as_strided((t_q * qh * d_v,), (1,), o_buf.storage_offset()).zero_() + if lse_valid is not None: + if sinks is not None: + sinks_v = self._checked_sinks_1d(sinks) + sinks_v = sinks_v.reshape(qh, 1).expand(qh, t_q) if self.thd_stats_head_major else sinks_v.reshape(1, qh).expand(t_q, qh) + lse_valid.copy_(sinks_v) + else: + lse_valid.fill_(float("-inf")) return - # 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 are None-specialized like the LSE 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", @@ -1858,14 +1894,16 @@ def _packed(buf, tokens): d_qk=d_qk, d_v=d_v, max_sq=max_sq, - has_lse=False, + has_lse=self.lse_desc is not None, + lse_head_major=self.thd_stats_head_major, + lse_head_stride=self.thd_stats_head_stride, ) fn( _packed(q_buf, t_q), _packed(k_buf, t_kv), _packed(v_buf, t_kv), _packed(o_buf, t_q), - None, + lse, sinks_t, seq_q_dummy, meta, @@ -1876,11 +1914,14 @@ def _packed(buf, tokens): def scratch_workspace_bytes(self) -> int: if self.thd: # [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. + # No packed-LSE chunk: with a Stats output the kernel writes the + # caller's ragged Stats buffer directly (token-major (T, H) or + # head-major (H, head_stride)); without one it compiles with + # has_lse=False and 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 7275df746..3cd0dc6e9 100644 --- a/python/cudnn/sdpa/fwd/engines.py +++ b/python/cudnn/sdpa/fwd/engines.py @@ -484,6 +484,7 @@ def _sm120_spec() -> EngineSpec: padded_stats=True, thd=True, thd_bottom_right=True, + thd_stats=True, layouts=frozenset({"bshd", "dense_flex"}), sched_policies=frozenset({SCHED_NATURAL}), tile_ms=frozenset({64, 128}), @@ -691,8 +692,7 @@ def _execute(variant_pack, workspace=None, stream=None): k_tensor=k_buf, v_tensor=v_buf, o_tensor=o_buf, - # THD ignores lse_tensor (the packed LSE is api-level scratch). - lse_tensor=lse_buf.reshape(facts.b, facts.h_q, facts.s_q) if lse_buf is not None else None, + lse_tensor=lse_buf, scale_softmax=facts.scale, sinks=sinks_buf, seq_kv_lens=seq_kv_buf, diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py b/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py index 7a9d32f9d..165781d33 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_f16_sm120.py @@ -179,6 +179,7 @@ def __init__( 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, kv_tile: int = SEQ_KV_TILES[0], @@ -202,6 +203,10 @@ def __init__( 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 + token-major ``(T, H)``. :param head_tile_qk: Q/K head dimension (the QK^T contraction width). Must be a multiple of 16 between 16 and 256, inclusive. :param head_tile_v: V/O head dimension (the P@V output width). Same @@ -225,6 +230,7 @@ def __init__( 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 self.head_tile_v = head_tile_v @@ -729,7 +735,9 @@ def kernel( :param k: Key tensor. :param v: Value tensor. :param o: Output tensor. - :param lse: ``(B, H, Sq)`` fp32 log-sum-exp output, or ``None`` to + :param lse: fp32 log-sum-exp output — ``(B, H, Sq)`` dense; packed + token-major ``(T, H)`` or head-major ``(H, head_stride)`` (per + ``thd_lse_head_major``) under ``thd_varlen``; 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``. @@ -1093,12 +1101,19 @@ def kernel( 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. + # Packed ragged-Stats LSE, written directly in the + # caller's declared layout: token-major (T, H) or + # head-major (H, head_stride). + # 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 + if cutlass.const_expr(self.thd_lse_head_major): + lse_row = lse_arr[head_idx, :] + lse_row[q_row_base + lse_q_idx] = lse_out + else: + lse_row = lse_arr[q_row_base + lse_q_idx, :] + lse_row[head_idx] = lse_out else: # Rows at/past this batch's Q length trim to -inf. if lse_q_idx >= seqlen_q: @@ -1197,7 +1212,9 @@ 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, or ``None`` to + :param lse: fp32 log-sum-exp output — ``(B, H, Sq)`` dense; packed + token-major ``(T, H)`` or head-major ``(H, head_stride)`` (per + ``thd_lse_head_major``) under ``thd_varlen``; 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``. @@ -1225,10 +1242,22 @@ def __call__( 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 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.thd_varlen): + 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]): + 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])): + 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") + else: + 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],)): @@ -1339,6 +1368,8 @@ def compile( # noqa: A001 d_v: int = 128, max_sq: int = 0, has_lse: bool = True, + lse_head_major: bool = False, + lse_head_stride: int = 0, ) -> Callable: """Compile and cache one architecture-specific compact BSHD shape. @@ -1351,7 +1382,10 @@ def compile( # noqa: A001 ``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. + at all instead of a dummy. THD LSE is token-major ``(T, H)`` by default; + ``lse_head_major=True`` switches to head-major ``(H, lse_head_stride)`` + (FlashAttention's ``softmax_lse`` layout), where ``lse_head_stride`` is the + caller-declared head-row stride (``>= T``, a shape — part of the cache key). """ kernel = SM120FusedMultiHeadAttentionForward( @@ -1366,6 +1400,7 @@ def compile( # noqa: A001 thd_varlen=PARAMS.thd_varlen, thd_batch=b, thd_max_sq=max_sq, + thd_lse_head_major=lse_head_major, head_tile_qk=d_qk, head_tile_v=d_v, q_tile=PARAMS.q_tile, @@ -1396,11 +1431,15 @@ def compile( # noqa: A001 stride_order=(3, 2, 1, 0), assumed_align=16, ) + if PARAMS.thd_varlen: + fake_lse_shape = (qh, lse_head_stride) if lse_head_major else (sq, qh) + else: + fake_lse_shape = (fake_batch, qh, sq) fake_lse = ( cute.runtime.make_fake_compact_tensor( cutlass.Float32, - (fake_batch, qh, sq), - stride_order=(2, 1, 0), + fake_lse_shape, + stride_order=(1, 0) if PARAMS.thd_varlen else (2, 1, 0), assumed_align=4, ) if has_lse 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 a629452d6..4bbdda888 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py @@ -327,7 +327,7 @@ def _pack_thd(seqs: list[torch.Tensor], s_max: int) -> tuple[torch.Tensor, torch for s in seqs: cu.append(cu[-1] + s.shape[2]) storage = torch.zeros(b * s_max * h * d, dtype=seqs[0].dtype, device="cuda") - packed = storage[: cu[-1] * h * d].view(max(cu[-1], 1), h, d) + packed = storage[: max(cu[-1], 1) * h * d].view(max(cu[-1], 1), h, d) for i, s in enumerate(seqs): packed[cu[i] : cu[i + 1]].copy_(s[0].permute(1, 0, 2)) view = storage.as_strided((b, h, s_max, d), (s_max * h * d, d, h * d, 1)) @@ -348,8 +348,16 @@ def _run_thd_case( causal_bottom_right: bool = False, window_size_left: int | None = None, with_sink: bool = False, + check_stats: bool = False, + stats_layout: str = "token_major", ) -> None: - """Run a THD (ragged) graph on the SM120 engine vs per-sequence references.""" + """Run a THD (ragged) graph on the SM120 engine vs per-sequence references. + + ``stats_layout`` selects the ragged Stats declaration: ``token_major`` + (``[t, h]``, sequence stride ``h_q``) or ``head_major`` (``[h, t]``, + sequence stride 1 with a padded token-capacity head stride — FlashAttention's + ``softmax_lse`` layout, mirroring PR #462's harness convention). + """ _require_dsl() import cudnn @@ -391,7 +399,7 @@ def _run_thd_case( q=tq, k=tk, v=tv, - generate_stats=False, + generate_stats=check_stats, attn_scale=scale, use_padding_mask=True, seq_len_q=sq, @@ -408,10 +416,31 @@ def _run_thd_case( st = graph.tensor_like(sinks, name="sink") sdpa_kwargs["sink_token"] = st variant_pack[st] = sinks - o, _ = graph.sdpa(**sdpa_kwargs) + o, stats = graph.sdpa(**sdpa_kwargs) o.set_output(True).set_dim(list(o_view.shape)).set_stride(list(o_view.stride())) o.set_ragged_offset(ro) variant_pack[o] = o_view + stats_storage = None + t_cap = max(64, -(-sum(seq_q_lens) // 64) * 64) + if check_stats: + assert stats is not None + stats.set_output(True) + stats.set_data_type(cudnn.data_type.FLOAT) + if stats_layout == "head_major": + # [h, t]: tokens contiguous within a head, heads strided by the + # padded token capacity; offsets = cu_q * stride_s = cu_q. + stats_storage = torch.empty(h_q * t_cap, dtype=torch.float32, device="cuda") + stats.set_dim((batch, h_q, s_q_max, 1)).set_stride((h_q * t_cap, t_cap, 1, 1)) + stats_ro_t = (q_ro.flatten() // (head_dim * h_q)).view(batch + 1, 1, 1, 1).contiguous() + else: + # [t, h]: heads contiguous within a token; offsets = cu_q * h_q. + stats_storage = torch.empty(batch * s_q_max * h_q, dtype=torch.float32, device="cuda") + stats.set_dim((batch, h_q, s_q_max, 1)).set_stride((s_q_max * h_q, 1, h_q, 1)) + stats_ro_t = (q_ro.flatten() // head_dim).view(batch + 1, 1, 1, 1).contiguous() + stats_ro = graph.tensor_like(stats_ro_t, name="stats_ro") + stats.set_ragged_offset(stats_ro) + variant_pack[stats_ro] = stats_ro_t + variant_pack[stats] = stats_storage graph.validate() graph.build_operation_graph() @@ -427,11 +456,17 @@ def _run_thd_case( for n in seq_q_lens: cu.append(cu[-1] + n) packed_o = o_storage[: cu[-1] * h_q * d_v].view(max(cu[-1], 1), h_q, d_v) - for i, (nq, nkv) in enumerate(zip(seq_q_lens, seq_kv_lens)): + if check_stats and stats_layout == "head_major": + packed_stats = stats_storage.view(h_q, t_cap) # (H, head_stride); tokens at [:, cu[i]:cu[i+1]] + elif check_stats: + packed_stats = stats_storage[: cu[-1] * h_q].view(max(cu[-1], 1), h_q) # (T, H) + else: + packed_stats = None + for i, (nq, _nkv) in enumerate(zip(seq_q_lens, seq_kv_lens)): if nq == 0: continue got = packed_o[cu[i] : cu[i + 1]].permute(1, 0, 2).unsqueeze(0).float() - expected = _ref_sdpa_full( + ref = _ref_sdpa_full( q_seqs[i], k_seqs[i], v_seqs[i], @@ -440,8 +475,16 @@ def _run_thd_case( causal_bottom_right=causal_bottom_right, window_size_left=window_size_left, sinks=sinks, + return_stats=check_stats, ) + expected, expected_lse = ref if check_stats else (ref, None) torch.testing.assert_close(got, expected, atol=0.1, rtol=5e-2) + if check_stats: + if stats_layout == "head_major": + got_lse = packed_stats[:, cu[i] : cu[i + 1]].unsqueeze(0) # (H, T_i) -> (1, H, T_i) + else: + got_lse = packed_stats[cu[i] : cu[i + 1]].t().unsqueeze(0) # (T_i, H) -> (1, H, T_i) + torch.testing.assert_close(got_lse, expected_lse, atol=2e-2, rtol=2e-2) @pytest.mark.L0 @@ -748,11 +791,29 @@ def test_dsl_sm120_execute_contract_mismatches(): 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() + # THD LSE must be declared packed: token-major [t, h] or head-major + # [h, t]. A dense-contiguous declaration (stride (S*H, S, 1)) is valid + # head-major (head_stride S); a padded sequence stride matches NEITHER layout + # and is rejected up front instead of being silently mis-addressed. + lse_padded = torch.empty(4 * 128 * 2, dtype=torch.float32, device="cuda").as_strided((1, 4, 128), (4 * 128 * 2, 128 * 2, 2)) + with pytest.raises(ValueError, match="token-major"): + SdpaFwdDslSm120(sample_q=q, sample_k=k, sample_v=v, sample_o=o, sample_lse=lse_padded, thd=True).check_support() + api = SdpaFwdDslSm120(sample_q=q, sample_k=k, sample_v=v, sample_o=o, sample_lse=lse, thd=True) + assert api.check_support() and api.thd_stats_head_major and api.thd_stats_head_stride == 128 + + # THD execute keeps the same presence contract as dense: the raise fires + # before any packing or launch, so plain dense buffers suffice here. + lse_thd = torch.empty(1 * 4 * 128, dtype=torch.float32, device="cuda").as_strided((1, 4, 128), (128 * 4, 1, 4)) + api = SdpaFwdDslSm120(sample_q=q, sample_k=k, sample_v=v, sample_o=o, sample_lse=lse_thd, thd=True) + assert api.check_support() + api.compile() + with pytest.raises(ValueError, match="lse_tensor is required"): + api.execute(q_tensor=q, k_tensor=k, v_tensor=v, o_tensor=o, seq_q_lens=seq_kv, seq_kv_lens=seq_kv) + api = SdpaFwdDslSm120(sample_q=q, sample_k=k, sample_v=v, sample_o=o, thd=True) + assert api.check_support() + api.compile() + with pytest.raises(ValueError, match="without an LSE output"): + api.execute(q_tensor=q, k_tensor=k, v_tensor=v, o_tensor=o, seq_q_lens=seq_kv, seq_kv_lens=seq_kv, lse_tensor=lse_thd) @pytest.mark.L0 @@ -779,21 +840,59 @@ def test_dsl_sm120_thd_bottom_right(): _run_thd_case(seq_q_lens=[100, 60], seq_kv_lens=[180, 120], is_causal=True, causal_bottom_right=True) +@pytest.mark.L0 +@pytest.mark.parametrize("stats_layout", ["token_major", "head_major"]) +@torch_fork_set_rng(seed=30) +def test_dsl_sm120_thd_stats(stats_layout: str): + """THD + generate_stats: the ragged Stats output is written in the + caller's declared layout — token-major [t, h] or head-major [h, t] + (the SM100 rows reject this combination — thd_stats gap).""" + + _run_thd_case(seq_q_lens=[200, 150], seq_kv_lens=[200, 150], is_causal=True, check_stats=True, stats_layout=stats_layout) + + +@pytest.mark.L1 +@torch_fork_set_rng(seed=32) +def test_dsl_sm120_thd_swa_stats(): + """THD + causal left sliding window + ragged Stats: the window trims the + per-sequence LSE denominator, and a distinct compiled specialization + (window_size_left is a template parameter) carries the token-major store.""" + + _run_thd_case(seq_q_lens=[150, 90], seq_kv_lens=[150, 90], is_causal=True, window_size_left=32, check_stats=True) + + @pytest.mark.L1 +@pytest.mark.parametrize("stats_layout", ["token_major", "head_major"]) @torch_fork_set_rng(seed=25) -def test_dsl_sm120_thd_gqa_sink(): - """THD + GQA + attention sink through the packed epilogue fold.""" +def test_dsl_sm120_thd_gqa_sink(stats_layout: str): + """THD + GQA + attention sink through the packed epilogue fold, with the + sink entering the ragged Stats (both declared layouts).""" - _run_thd_case(seq_q_lens=[130, 70], seq_kv_lens=[130, 70], h_q=8, h_kv=2, is_causal=True, with_sink=True) + _run_thd_case(seq_q_lens=[130, 70], seq_kv_lens=[130, 70], h_q=8, h_kv=2, is_causal=True, with_sink=True, check_stats=True, stats_layout=stats_layout) @pytest.mark.L1 @torch_fork_set_rng(seed=26) def test_dsl_sm120_thd_zero_length_sequence(): """A zero-length sequence contributes no tokens and must not perturb its - packed neighbors.""" + packed neighbors (O and ragged Stats). The last sequence has Q tokens but + ZERO keys inside a live launch: its rows must come back O := 0 with + LSE := -inf through the kernel's row_sum <= 0 guard, not stale memory.""" + + _run_thd_case(seq_q_lens=[128, 0, 64], seq_kv_lens=[100, 0, 0], is_causal=True, check_stats=True) - _run_thd_case(seq_q_lens=[128, 0, 64], seq_kv_lens=[100, 0, 50], is_causal=True) + +@pytest.mark.L1 +@pytest.mark.parametrize("stats_layout", ["token_major", "head_major"]) +@pytest.mark.parametrize("with_sink", [False, True], ids=["no_sink", "sink"]) +@torch_fork_set_rng(seed=31) +def test_dsl_sm120_thd_all_kv_zero_stats(with_sink: bool, stats_layout: str): + """Every KV length zero: a zero-token K/V view cannot back a TMA + descriptor, so the adapter short-cut fills O := 0 and the ragged Stats + adapter-side — -inf, or the sink value alone (the sink column keeps the + softmax denominator alive) — in either declared layout.""" + + _run_thd_case(seq_q_lens=[64, 32], seq_kv_lens=[0, 0], with_sink=with_sink, check_stats=True, stats_layout=stats_layout) @pytest.mark.L0 @@ -872,7 +971,8 @@ def test_dsl_sm120_mixed_head_dims_causal_gqa_stats(): @pytest.mark.L0 @torch_fork_set_rng(seed=25) def test_dsl_sm120_thd_mixed_head_dims(): - """THD packed views carry per-tensor head dims (Q/K at 192, V/O at 128).""" + """THD packed views carry per-tensor head dims (Q/K at 192, V/O at 128), + with the ragged Stats checked under the bottom-right diagonal.""" _run_thd_case( seq_q_lens=[33, 128, 7], seq_kv_lens=[65, 128, 190], @@ -882,6 +982,7 @@ def test_dsl_sm120_thd_mixed_head_dims(): head_dim_v=128, is_causal=True, causal_bottom_right=True, + check_stats=True, ) @@ -940,11 +1041,14 @@ def test_dsl_sm120_mixed_head_dims_sink_bf16(): @pytest.mark.L0 +@pytest.mark.parametrize("stats_layout", ["token_major", "head_major"]) @torch_fork_set_rng(seed=29) -def test_dsl_sm120_thd_mixed_head_dims_sink(): - """THD + mixed head dims + sink: the packed per-tensor D views and the - sink denominator together, without a mask (the THD mixed-dim causal case - is covered above).""" +def test_dsl_sm120_thd_mixed_head_dims_sink(stats_layout: str): + """THD + mixed head dims + sink + ragged Stats: the packed per-tensor D + views and the sink denominator together, without a mask (the THD + mixed-dim causal case is covered above). The Stats ragged offsets derive + from Q's (d_qk-based) offsets, so both declared LSE layouts are pinned + under d_qk != d_v.""" _run_thd_case( seq_q_lens=[130, 70, 9], seq_kv_lens=[130, 70, 190], @@ -953,6 +1057,8 @@ def test_dsl_sm120_thd_mixed_head_dims_sink(): head_dim=192, head_dim_v=128, with_sink=True, + check_stats=True, + stats_layout=stats_layout, ) diff --git a/test/python/sdpa/frost/test_sdpa_graph_analyzer.py b/test/python/sdpa/frost/test_sdpa_graph_analyzer.py index 1e51f490c..bfd266915 100644 --- a/test/python/sdpa/frost/test_sdpa_graph_analyzer.py +++ b/test/python/sdpa/frost/test_sdpa_graph_analyzer.py @@ -647,6 +647,44 @@ def test_sm120_probe_accepts_thd_bottom_right(monkeypatch): assert engines.engine_name(arch="sm120") in _eligible(g) +def test_sm120_probe_accepts_thd_stats(monkeypatch): + """The SM120 epilogue writes cuDNN's token-major ragged Stats directly, + so THD + generate_stats is eligible.""" + monkeypatch.setattr(ga, "_device_cc", lambda: (12, 0)) + g = _mk_graph() + dims = (B, H, S, 128) + strides = (S * H * 128, 128, H * 128, 1) + q = g.tensor(dim=dims, stride=strides, data_type=DTYPE, name="q") + k = g.tensor(dim=dims, stride=strides, data_type=DTYPE, name="k") + v = g.tensor(dim=dims, stride=strides, data_type=DTYPE, name="v") + ro = g.tensor(dim=(B + 1, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT64, name="ro") + q.set_ragged_offset(ro) + k.set_ragged_offset(ro) + v.set_ragged_offset(ro) + seq_q = g.tensor(dim=(B, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT32, name="sq") + seq_kv = g.tensor(dim=(B, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT32, name="skv") + o, stats = g.sdpa( + name="s", + q=q, + k=k, + v=v, + attn_scale=0.1, + generate_stats=True, + use_causal_mask=True, + use_padding_mask=True, + seq_len_q=seq_q, + seq_len_kv=seq_kv, + ) + _finish_output(o, dims, strides) + o.set_ragged_offset(ro) + assert stats is not None + stats.set_output(True).set_dim((B, H, S, 1)).set_stride((S * H, 1, H, 1)) + stats.set_data_type(cudnn.data_type.FLOAT) + stats_ro = g.tensor(dim=(B + 1, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT64, name="stats_ro") + stats.set_ragged_offset(stats_ro) + assert engines.engine_name(arch="sm120") in _eligible(g) + + @pytest.mark.parametrize("side", ["cu_seq_len_q", "cu_seq_len_kv"]) def test_cu_seq_len_is_declined(side): """cu_seq_len_* (cuDNN 9.24+) are prefix sums — a different contract from