diff --git a/python/cudnn/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py index 3c9e26f4d..aec5e345d 100644 --- a/python/cudnn/sdpa/fwd/api_dsl.py +++ b/python/cudnn/sdpa/fwd/api_dsl.py @@ -277,6 +277,8 @@ def __init__( scale_softmax: Optional[float] = None, seq_kv_lens_present: bool = False, seq_q_lens_present: bool = False, + cu_seq_q_lens: bool = False, + cu_seq_kv_lens: bool = False, has_sink: bool = False, thd: bool = False, dtype_o: Optional[torch.dtype] = None, @@ -327,6 +329,13 @@ def __init__( # FA's seqused_q) bound directly at execute — no packing, no # per-execute copies. Dense-only. self.seq_q_lens_present = bool(seq_q_lens_present) + # cu_seq_len form (cuDNN 9.24+): the corresponding seq-lens execute + # argument arrives as a (B+1,)-int32 PREFIX-SUM tensor instead of + # (B,) per-batch lengths. THD-only today: the ragged lowering derives + # both forms host-side from its inherent tolist round-trip; the dense + # kernels have no CU read mode yet (check_support rejects). + self.cu_seq_q_lens = bool(cu_seq_q_lens) + self.cu_seq_kv_lens = bool(cu_seq_kv_lens) self.has_sink = bool(has_sink) self.thd = bool(thd) # MXFP8: FP8 (E4M3/E5M2) Q/K/V in, half (BF16/FP16) O out. dtype_o overrides @@ -547,6 +556,49 @@ def _checked_seq_lens(self, seq_lens: torch.Tensor, name: str) -> torch.Tensor: ) return seq_lens.reshape(-1) + def _checked_cu_seq_lens(self, cu_seq_lens: torch.Tensor, name: str) -> torch.Tensor: + """Validate a caller-provided (B+1,)-int32 prefix-sum tensor (cu_seq_len form). + + Strictly a view, like :meth:`_checked_seq_lens`. The prefix-sum + INVARIANTS (starts at 0, non-decreasing) are runtime values — they are + validated host-side by the THD lowering's inherent tolist round-trip, + not here. + """ + self._value_error_if( + cu_seq_lens.dtype != torch.int32, + f"{name} must be int32; got {cu_seq_lens.dtype}", + ) + self._value_error_if( + cu_seq_lens.numel() != self.batch_size + 1, + f"{name} must have B + 1 = {self.batch_size + 1} elements (prefix sums); got {cu_seq_lens.numel()}", + ) + self._value_error_if( + not cu_seq_lens.is_contiguous(), + f"{name} must be contiguous (read as a flat (B+1,) view)", + ) + return cu_seq_lens.reshape(-1) + + def _thd_host_lens(self, seq_lens, name: str, cu_form: bool) -> tuple[list, list]: + """One inherent D2H round-trip -> (per-batch lens, prefix sums) host lists. + + Consumes EITHER length form: per-batch ``(B,)`` lengths (prefix sums + built by a Python scan) or the ``(B+1,)`` cu_seq_len prefix-sum form + (lengths are adjacent differences; the prefix-sum invariants are + validated here, where they are free to check). + """ + if cu_form: + cu_host = [int(x) for x in self._checked_cu_seq_lens(seq_lens, name).tolist()] + self._value_error_if( + cu_host[0] != 0 or any(cu_host[i] > cu_host[i + 1] for i in range(len(cu_host) - 1)), + f"{name} must be a non-decreasing prefix sum starting at 0; got {cu_host}", + ) + return [cu_host[i + 1] - cu_host[i] for i in range(len(cu_host) - 1)], cu_host + lens_host = [int(x) for x in self._checked_seq_lens(seq_lens, name).tolist()] + cu_host = [0] + for n in lens_host: + cu_host.append(cu_host[-1] + n) + return lens_host, cu_host + def _check_seq_lens_contract(self, seq_q_lens, seq_kv_lens) -> None: """Reject seq-length tensors inconsistent with the compiled specialization. @@ -804,6 +856,10 @@ def check_support(self) -> bool: ) if self.thd: self.seq_kv_lens_present = True + self._not_implemented_error_if( + (self.cu_seq_q_lens or self.cu_seq_kv_lens) and not self.thd, + "cu_seq_len_* is THD-only (the dense kernels have no CU read mode yet)", + ) # Dense padded-Q trim backstops (engines.lower_dsl_prefill never sets # these combinations; a direct caller could). self._value_error_if( @@ -1159,23 +1215,17 @@ def _execute_thd(self, q_buf, k_buf, v_buf, o_buf, scale_softmax_log2, sinks, se import cutlass dev = q_buf.device - 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() + b = self.batch_size carver = WorkspaceCarver(workspace, self.scratch_workspace_bytes(), "SdpaFwdDslSm100 (THD)") if workspace is not None else None # Metadata buffer: [ seq_kv_lens(B) | cu_seqlens_q(B+1) | cu_seqlens_k(B+1) ], # built HOST-side from the (inherent) tolist round-trip and uploaded in # ONE H2D copy: a device-side cumsum would allocate its scan-temp - # storage and launch kernels on the execute hot path. + # storage and launch kernels on the execute hot path. Either length + # form feeds it — per-batch (B,) lengths or the (B+1,) cu_seq_len + # prefix sums — at identical cost. meta = carver.take(3 * b + 2, torch.int32) if carver is not None else torch.empty(3 * b + 2, dtype=torch.int32, device=dev) - slq_host = slq_v.tolist() # one D2H sync; t_q/t_kv/units are runtime values - slk_host = slk_v.tolist() - cu_q_host = [0] - for n in slq_host: - cu_q_host.append(cu_q_host[-1] + int(n)) - cu_k_host = [0] - for n in slk_host: - cu_k_host.append(cu_k_host[-1] + int(n)) + slq_host, cu_q_host = self._thd_host_lens(seq_q_lens, "cu_seq_len_q" if self.cu_seq_q_lens else "seq_q_lens", self.cu_seq_q_lens) + slk_host, cu_k_host = self._thd_host_lens(seq_len_kv, "cu_seq_len_kv" if self.cu_seq_kv_lens else "seq_kv_lens", self.cu_seq_kv_lens) meta.copy_(torch.tensor(slk_host + cu_q_host + cu_k_host, dtype=torch.int32)) t_q = cu_q_host[-1] t_kv = cu_k_host[-1] @@ -1720,6 +1770,10 @@ def check_support(self) -> bool: if self.thd: self._value_error_if(self.seq_q_lens_present, "seq_q_lens_present is dense-only (THD carries per-sequence Q lengths via cu_seqlens)") self.seq_kv_lens_present = True + self._not_implemented_error_if( + (self.cu_seq_q_lens or self.cu_seq_kv_lens) and not self.thd, + "cu_seq_len_* is THD-only (the dense kernels have no CU read mode yet)", + ) self._value_error_if( self.sched_policy is not None and self.sched_policy != SCHED_NATURAL, f"SM120 DSL SDPA only supports sched_policy={SCHED_NATURAL}", @@ -2263,23 +2317,17 @@ def _thd_pack(self, q_buf, k_buf, v_buf, o_buf, seq_q_lens, seq_kv_lens, workspa dev = q_buf.device carver = WorkspaceCarver(workspace, self.scratch_workspace_bytes(), label) 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") # [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. # Built HOST-side from the (inherent) tolist round-trip and uploaded # in ONE H2D copy: a device-side cumsum would allocate its scan-temp - # storage and launch kernels on the execute hot path. + # storage and launch kernels on the execute hot path. Either length + # form feeds it — per-batch (B,) lengths or the (B+1,) cu_seq_len + # prefix sums — at identical cost. meta = carver.take(3 * b + 2, torch.int32) if carver is not None else torch.empty(3 * b + 2, dtype=torch.int32, device=dev) - slq_host = slq_v.tolist() - slk_host = slk_v.tolist() - cu_q_host = [0] - for n in slq_host: - cu_q_host.append(cu_q_host[-1] + int(n)) - cu_k_host = [0] - for n in slk_host: - cu_k_host.append(cu_k_host[-1] + int(n)) + slq_host, cu_q_host = self._thd_host_lens(seq_q_lens, "cu_seq_len_q" if self.cu_seq_q_lens else "seq_q_lens", self.cu_seq_q_lens) + slk_host, cu_k_host = self._thd_host_lens(seq_kv_lens, "cu_seq_len_kv" if self.cu_seq_kv_lens else "seq_kv_lens", self.cu_seq_kv_lens) meta.copy_(torch.tensor(slk_host + cu_q_host + cu_k_host, dtype=torch.int32)) t_q = cu_q_host[-1] t_kv = cu_k_host[-1] diff --git a/python/cudnn/sdpa/fwd/engines.py b/python/cudnn/sdpa/fwd/engines.py index 3ef503079..45b8698b5 100644 --- a/python/cudnn/sdpa/fwd/engines.py +++ b/python/cudnn/sdpa/fwd/engines.py @@ -157,8 +157,17 @@ class Capabilities: # that keep False (the SM100 FP8/MXFP8 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 lowerings assume FULLY-PACKED storage: the packed addressing is + # re-derived as prefix(lens) x token stride, and the graph's bound + # ragged-offset values are never read. TE-style padded THD (offsets + # from cu_seqlens_padded != cu_seqlens, gaps between sequences) is NOT + # served — and being runtime data, cannot be declined at plan time. thd: bool = False - cu_seq_len: bool = False # cu_seq_len_q / cu_seq_len_kv prefix sums (no row serves these yet) + # cu_seq_len_q / cu_seq_len_kv (B+1,) prefix sums (cuDNN 9.24+). Serving + # rows consume the form on THD host-side (lens = adjacent differences of + # the inherent tolist); dense cu graphs stay declined until the kernels + # grow a CU read mode (len = cu[b+1] - cu[b]) — see mismatch(). + cu_seq_len: bool = False # Dense padded + stats needs the per-batch seq_len_q LSE trim (padded # q-rows write LSE=-inf / O=0, cuDNN >= 9.14). Plumbed for the half # kernels via SEQ_Q_LENS_PRESENT; the FP8/MXFP8 kernels lack the epilogue @@ -317,13 +326,25 @@ def mismatch(capabilities: Capabilities, facts: "ga.SdpaGraphFacts", knobs: Opti (facts.has_sink, capabilities.sink, "sink token"), (facts.wants_stats, capabilities.stats, "stats output"), (facts.thd, capabilities.thd, "THD / ragged"), - (facts.has_cu_seq_len, capabilities.cu_seq_len, "cu_seq_len_q / cu_seq_len_kv"), ): if fact and not cap: return f"graph uses {label}, which this engine does not support" if facts.right_band_widening and facts.right_bound is not None and facts.right_bound < 0: return f"negative diagonal_band_right_bound ({facts.right_bound}) is not supported" + + if facts.has_cu_seq_len: + # cu_seq_len_* ((B+1,) prefix sums, cuDNN 9.24+). The THD lowering + # consumes either length form host-side; the dense kernels' CU read + # mode (len = cu[b+1] - cu[b]) is not plumbed yet, so dense cu graphs + # stay declined even on serving rows. + if not capabilities.cu_seq_len: + return "graph uses cu_seq_len_q / cu_seq_len_kv, which this engine does not support" + if not facts.thd: + return "cu_seq_len_* on dense graphs is not supported yet (kernel CU read mode not plumbed)" + if (facts.seq_q_t is not None and facts.cu_seq_q_t is not None) or (facts.seq_kv_t is not None and facts.cu_seq_kv_t is not None): + return "seq_len_* and cu_seq_len_* on the same side is ambiguous (backend precedence is not replicated here)" + if facts.bottom_right: if not (facts.causal or facts.right_band_widening): return "bottom-right alignment requires a causal upper bound (plain or right-widened)" @@ -389,6 +410,7 @@ def _sm100_spec(d: int, d_v: Optional[int] = None) -> EngineSpec: stats=True, lse_optional=True, thd=True, + cu_seq_len=True, padded_stats=True, # The f16/bf16 lowering serves any dense B/H/S stride permutation # (padded strides included) with the head dim innermost; the @@ -626,6 +648,7 @@ def _sm120_spec() -> EngineSpec: # of mask flags. Ragged S_kv is served natively with no synthesized # padding and no padded-path cost. skv_tile=0, + cu_seq_len=True, layouts=frozenset({"bshd", "dense_flex"}), sched_policies=frozenset({SCHED_NATURAL}), tile_ms=frozenset({64, 128}), @@ -713,6 +736,10 @@ def lower_dsl_prefill( # 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=seq_q_lens_present, + # cu_seq_len form (THD-only; the probe declined dense cu graphs): the + # adapter's seq-lens execute arguments carry (B+1,) prefix sums. + cu_seq_q_lens=facts.cu_seq_q_t is not None, + cu_seq_kv_lens=facts.cu_seq_kv_t is not None, has_sink=facts.has_sink, thd=facts.thd, dtype_o=facts.dtype_o if (facts.is_mxfp8 or facts.is_fp8) else None, @@ -755,6 +782,8 @@ def lower_dsl_prefill( sink_token=facts.sink_t, seq_len_kv=seq_kv_t, seq_len_q=seq_q_t, + cu_seq_len_q=facts.cu_seq_q_t, + cu_seq_len_kv=facts.cu_seq_kv_t, sf_q=facts.sf_q_t, sf_k=facts.sf_k_t, sf_v=facts.sf_v_t, diff --git a/python/cudnn/sdpa/graph_analyzer.py b/python/cudnn/sdpa/graph_analyzer.py index 3ba1f2d98..2fef9dfbf 100644 --- a/python/cudnn/sdpa/graph_analyzer.py +++ b/python/cudnn/sdpa/graph_analyzer.py @@ -232,10 +232,11 @@ class SdpaGraphFacts: padded: bool = False # per-batch KV lengths present (padding mask or THD) thd: bool = False # ragged (THD) Q/K/V - # cu_seq_len_q / cu_seq_len_kv (cuDNN 9.24+): prefix sums, a contract of - # their own — neither seq_len_* nor ragged_offset. A fact, not a verdict: - # no engine here implements it yet, but reading the graph as plain padded - # gave wrong output (14.9% of O on test_sdpa_mixed_seq_len_forms_L0[cu_q_brcm]). + # 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 + # the graph as plain padded gave wrong output — 14.9% of O on + # test_sdpa_mixed_seq_len_forms_L0[cu_q_brcm]). has_cu_seq_len: bool = False has_sink: bool = False wants_stats: bool = False @@ -253,6 +254,10 @@ class SdpaGraphFacts: sink_t: Any = None seq_kv_t: Any = None seq_q_t: Any = None + # (B+1,) prefix-sum IR refs (cuDNN 9.24+ cu_seq_len form); None when the + # graph carries the per-batch seq_len_* form on that side instead. + cu_seq_q_t: Any = None + cu_seq_kv_t: Any = None # Feature operands (bias / block-mask / score-stat outputs). bias_t: Any = None block_mask_t: Any = None @@ -474,30 +479,39 @@ def _square_transposed(dim: tuple, stride: tuple) -> bool: return _invalid(f"sliding-window length must be >= 1; got {left_bound}") window_left = (left_bound - 1) if left_bound is not None else None - has_cu_seq_len = any(rec.get(name) is not None for name in ("cu_seq_len_q", "cu_seq_len_kv")) + cu_seq_q = rec.get("cu_seq_len_q") + cu_seq_kv = rec.get("cu_seq_len_kv") + has_cu_seq_len = cu_seq_q is not None or cu_seq_kv is not None - # Padding / THD. + # Padding / THD. Per-batch lengths arrive as seq_len_* ((B,) lengths) or + # cu_seq_len_* ((B+1,) prefix sums, cuDNN 9.24+) per side; either form + # satisfies the length requirement. Both-on-one-side is NOT flagged + # invalid here (invalid means malformed-for-everyone; the backend accepts + # it with its own precedence) — an engine that serves the cu form must + # decline the ambiguous combination itself. thd = getattr(q, "ragged_offset", None) is not None use_padding_mask = bool(rec.get("use_padding_mask", False)) seq_len_kv = rec.get("seq_len_kv") seq_len_q = rec.get("seq_len_q") + q_lens_given = seq_len_q is not None or cu_seq_q is not None + kv_lens_given = seq_len_kv is not None or cu_seq_kv is not None seq_q_trim = False if thd: if getattr(k, "ragged_offset", None) is None or getattr(v, "ragged_offset", None) is None: return _invalid("THD (ragged) requires ragged Q, K, and V") - if seq_len_q is None or seq_len_kv is None: - return _invalid("THD (ragged) requires seq_len_q and seq_len_kv") + if not q_lens_given or not kv_lens_given: + return _invalid("THD (ragged) requires seq_len_q/cu_seq_len_q and seq_len_kv/cu_seq_len_kv") padded = True else: - if use_padding_mask and seq_len_kv is None: - return _invalid("use_padding_mask requires seq_len_kv") - 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 use_padding_mask and not kv_lens_given: + return _invalid("use_padding_mask requires seq_len_kv or cu_seq_len_kv") + seq_q_trim = q_lens_given and not use_padding_mask + padded = use_padding_mask and kv_lens_given + + # The kernels consume per-batch lengths / prefix sums 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), ("cu_seq_len_q", cu_seq_q), ("cu_seq_len_kv", cu_seq_kv)): if t is not None: if t.get_data_type() != cudnn.data_type.INT32: return _invalid(f"{name} must be int32; got {t.get_data_type()}") @@ -597,6 +611,8 @@ def _square_transposed(dim: tuple, stride: tuple) -> bool: sink_t=sink_token, seq_kv_t=seq_len_kv, seq_q_t=seq_len_q, + cu_seq_q_t=cu_seq_q, + cu_seq_kv_t=cu_seq_kv, sf_q_t=(dsc_q if is_mxfp8 else None), sf_k_t=(dsc_k if is_mxfp8 else None), sf_v_t=(dsc_v if is_mxfp8 else None), @@ -640,6 +656,9 @@ class SdpaBinding: sink_token: Any = None seq_len_kv: Any = None seq_len_q: Any = None + # (B+1,) prefix-sum form (cuDNN 9.24+); at most one form per side. + cu_seq_len_q: Any = None + cu_seq_len_kv: Any = None # MXFP8 block-scale (descale) tensors + Amax_O output. sf_q: Any = None sf_k: Any = None @@ -677,6 +696,8 @@ def bound_tensors(self) -> list: self.sink_token, self.seq_len_kv, self.seq_len_q, + self.cu_seq_len_q, + self.cu_seq_len_kv, self.sf_q, self.sf_k, self.sf_v, @@ -823,8 +844,17 @@ def _need(t_ref, label): ops = FeatureOperands(alibi=facts.has_alibi) if facts.padded: - ops.seq_kv_lens = _need(facts.seq_kv_t, "padding mask (seq_len_kv)") - if facts.seq_q_t is not None: + # Either length form satisfies a side: per-batch seq_len_* or the + # (B+1,) cu_seq_len_* prefix sums (cuDNN 9.24+) — the cu buffer + # travels through the same operand slot (the adapter was constructed + # knowing the form). + if facts.cu_seq_kv_t is not None: + ops.seq_kv_lens = _need(facts.cu_seq_kv_t, "padding mask (cu_seq_len_kv)") + else: + ops.seq_kv_lens = _need(facts.seq_kv_t, "padding mask (seq_len_kv)") + if facts.cu_seq_q_t is not None: + ops.seq_len_q = _need(facts.cu_seq_q_t, "per-batch query lengths (cu_seq_len_q)") + elif facts.seq_q_t is not None: ops.seq_len_q = _need(facts.seq_q_t, "per-batch query lengths (seq_len_q)") if facts.has_bias: ops.bias = _need(facts.bias_t, "bias") 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 0baea5fa7..657898709 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py @@ -691,6 +691,7 @@ def _run_dsl_thd_graph( mask="causal", check_stats=False, stats_layout="token_major", + cu_lens=False, ): """Build + execute a packed THD/varlen graph; returns the flat packed O storage buffer — plus, with ``check_stats``, the flat Stats storage and @@ -726,6 +727,10 @@ def _dense_buf(packed, s_max, t, H): slq = torch.tensor(seq_lens_q, dtype=torch.int32, device=dev).view(B, 1, 1, 1) slk = torch.tensor(seq_lens_kv, dtype=torch.int32, device=dev).view(B, 1, 1, 1) + # cu_lens: bind the (B+1,) prefix-sum form (cu_seq_len_q/kv, cuDNN 9.24+) + # instead of per-batch lengths. + cuq_t = torch.tensor(cu_q, dtype=torch.int32, device=dev).view(B + 1, 1, 1, 1) + cuk_t = torch.tensor(cu_k, dtype=torch.int32, device=dev).view(B + 1, 1, 1, 1) ro_q = (torch.tensor(cu_q, dtype=torch.int64, device=dev) * H_q * d).view(B + 1, 1, 1, 1) ro_k = (torch.tensor(cu_k, dtype=torch.int64, device=dev) * H_kv * d).view(B + 1, 1, 1, 1) @@ -734,8 +739,8 @@ def _dense_buf(packed, s_max, t, H): tq = g.tensor(dim=[B, H_q, S_max_q, d], stride=list(stride_q), data_type=io, name="q") tk = g.tensor(dim=[B, H_kv, S_max_kv, d], stride=list(stride_kv), data_type=io, name="k") tv = g.tensor(dim=[B, H_kv, S_max_kv, d], stride=list(stride_kv), data_type=io, name="v") - sq = g.tensor_like(slq) - skv = g.tensor_like(slk) + sq = g.tensor_like(cuq_t if cu_lens else slq) + skv = g.tensor_like(cuk_t if cu_lens else slk) qro = g.tensor_like(ro_q) kro = g.tensor_like(ro_k) vro = g.tensor_like(ro_k) @@ -743,9 +748,13 @@ def _dense_buf(packed, s_max, t, H): tq.set_ragged_offset(qro) tk.set_ragged_offset(kro) tv.set_ragged_offset(vro) - kw = dict(name="sdpa", q=tq, k=tk, v=tv, generate_stats=check_stats, attn_scale=scale, use_padding_mask=True, seq_len_q=sq, seq_len_kv=skv) + kw = dict(name="sdpa", q=tq, k=tk, v=tv, generate_stats=check_stats, attn_scale=scale, use_padding_mask=True) + if cu_lens: + kw.update(cu_seq_len_q=sq, cu_seq_len_kv=skv) + else: + kw.update(seq_len_q=sq, seq_len_kv=skv) kw.update(_mask_graph_kwargs(mask)) - vp = {tq: q_gpu, tk: k_gpu, tv: v_gpu, sq: slq, skv: slk, qro: ro_q, kro: ro_k, vro: ro_k, oro: ro_q} + vp = {tq: q_gpu, tk: k_gpu, tv: v_gpu, sq: (cuq_t if cu_lens else slq), skv: (cuk_t if cu_lens else slk), qro: ro_q, kro: ro_k, vro: ro_k, oro: ro_q} if sink is not None: st = g.tensor_like(sink) kw["sink_token"] = st @@ -849,7 +858,9 @@ def _cu(sl): torch.testing.assert_close(o_out, o_ref, atol=5e-2, rtol=3e-2) -def _run_thd_stats_case(*, seq_lens_q, seq_lens_kv, d=128, dtype=torch.float16, H_q=8, H_kv=8, mask="causal", with_sink=False, stats_layout="token_major"): +def _run_thd_stats_case( + *, seq_lens_q, seq_lens_kv, d=128, dtype=torch.float16, H_q=8, H_kv=8, mask="causal", with_sink=False, stats_layout="token_major", cu_lens=False +): """Run a THD (ragged) graph with generate_stats and check O and the ragged Stats against per-sequence references, in the declared Stats layout.""" _require_dsl() @@ -888,6 +899,7 @@ def _cu(sl): mask=mask, check_stats=True, stats_layout=stats_layout, + cu_lens=cu_lens, ) if T_q == 0: @@ -989,6 +1001,29 @@ def test_dsl_sm100_thd_all_q_zero_stats(stats_layout): _run_thd_stats_case(seq_lens_q=[0, 0], seq_lens_kv=[0, 0], mask="none", stats_layout=stats_layout) +@pytest.mark.L0 +@pytest.mark.parametrize("stats_layout", ["token_major", "head_major"]) +@torch_fork_set_rng(seed=35) +def test_dsl_sm100_thd_cu_seq_len_stats(stats_layout): + """THD with the cu_seq_len_q/kv length form ((B+1,) prefix sums, cuDNN + 9.24+ — the form TE/PyT/vLLM natively hold): the lowering derives the + per-batch lengths host-side from the same inherent tolist round-trip, so + results are identical to the seq_len form, ragged Stats included.""" + + _run_thd_stats_case(seq_lens_q=[200, 150], seq_lens_kv=[200, 150], mask="causal", stats_layout=stats_layout, cu_lens=True) + + +@pytest.mark.L1 +@torch_fork_set_rng(seed=36) +def test_dsl_sm100_thd_cu_seq_len_zero_lens(): + """cu_seq_len form with degenerate lengths: a zero-length sequence + (repeated prefix value) and an all-zero KV side keep the same dead-row / + no-op semantics as the seq_len form.""" + + _run_thd_stats_case(seq_lens_q=[128, 0, 64], seq_lens_kv=[100, 0, 0], mask="causal", cu_lens=True) + _run_thd_stats_case(seq_lens_q=[64, 32], seq_lens_kv=[0, 0], mask="none", cu_lens=True) + + _COMBO_MASKS = { "dense": ["none", "causal", "causal_br", "swa", "padded", "band", "band_br", "band_swa"], # THD forces padding internally, so its mask axis rides on top of that. 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 023527490..b6b38e0ca 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm120.py @@ -379,6 +379,7 @@ def _run_thd_case( with_sink: bool = False, check_stats: bool = False, stats_layout: str = "token_major", + cu_lens: bool = False, ) -> None: """Run a THD (ragged) graph on the SM120 engine vs per-sequence references. @@ -411,6 +412,16 @@ def _run_thd_case( o_storage.fill_(_THD_SENTINEL) sq_t = torch.tensor(seq_q_lens, dtype=torch.int32, device="cuda").view(batch, 1, 1, 1) skv_t = torch.tensor(seq_kv_lens, dtype=torch.int32, device="cuda").view(batch, 1, 1, 1) + + # cu_lens: bind the (B+1,) prefix-sum form (cu_seq_len_q/kv, cuDNN 9.24+) + # instead of per-batch lengths. + def _prefix(lens): + cu = [0] + for n in lens: + cu.append(cu[-1] + n) + return torch.tensor(cu, dtype=torch.int32, device="cuda").view(batch + 1, 1, 1, 1) + + cuq_t, cukv_t = (_prefix(seq_q_lens), _prefix(seq_kv_lens)) if cu_lens else (None, None) sinks = torch.randn(1, h_q, 1, 1, dtype=torch.float32, device="cuda") if with_sink else None io_dtype = cudnn.data_type.HALF if dtype == torch.float16 else cudnn.data_type.BFLOAT16 @@ -425,8 +436,8 @@ def _run_thd_case( tq.set_ragged_offset(rq) tk.set_ragged_offset(rk) tv.set_ragged_offset(rv) - sq = graph.tensor_like(sq_t, name="seq_q") - skv = graph.tensor_like(skv_t, name="seq_kv") + sq = graph.tensor_like(cuq_t, name="cu_seq_q") if cu_lens else graph.tensor_like(sq_t, name="seq_q") + skv = graph.tensor_like(cukv_t, name="cu_seq_kv") if cu_lens else graph.tensor_like(skv_t, name="seq_kv") sdpa_kwargs = dict( name="sdpa", q=tq, @@ -435,9 +446,11 @@ def _run_thd_case( generate_stats=check_stats, attn_scale=scale, use_padding_mask=True, - seq_len_q=sq, - seq_len_kv=skv, ) + if cu_lens: + sdpa_kwargs.update(cu_seq_len_q=sq, cu_seq_len_kv=skv) + else: + sdpa_kwargs.update(seq_len_q=sq, seq_len_kv=skv) _apply_mask_kwargs( sdpa_kwargs, cudnn, @@ -446,7 +459,17 @@ def _run_thd_case( window_size_left=window_size_left, window_size_right=window_size_right, ) - variant_pack = {tq: q_view, tk: k_view, tv: v_view, rq: q_ro, rk: k_ro, rv: v_ro, ro: o_ro, sq: sq_t, skv: skv_t} + variant_pack = { + tq: q_view, + tk: k_view, + tv: v_view, + rq: q_ro, + rk: k_ro, + rv: v_ro, + ro: o_ro, + sq: (cuq_t if cu_lens else sq_t), + skv: (cukv_t if cu_lens else skv_t), + } if sinks is not None: st = graph.tensor_like(sinks, name="sink") sdpa_kwargs["sink_token"] = st @@ -960,6 +983,30 @@ def test_dsl_sm120_thd_all_q_zero_stats(stats_layout: str): _run_thd_case(seq_q_lens=[0, 0], seq_kv_lens=[0, 0], check_stats=True, stats_layout=stats_layout) +@pytest.mark.L0 +@pytest.mark.parametrize("stats_layout", ["token_major", "head_major"]) +@torch_fork_set_rng(seed=35) +def test_dsl_sm120_thd_cu_seq_len_stats(stats_layout: str): + """THD with the cu_seq_len_q/kv length form ((B+1,) prefix sums, cuDNN + 9.24+ — the form TE/PyT/vLLM natively hold): the lowering derives the + per-batch lengths host-side from the same inherent tolist round-trip, so + results are identical to the seq_len form, ragged Stats included.""" + + _run_thd_case(seq_q_lens=[200, 150], seq_kv_lens=[200, 150], is_causal=True, check_stats=True, stats_layout=stats_layout, cu_lens=True) + + +@pytest.mark.L1 +@torch_fork_set_rng(seed=36) +def test_dsl_sm120_thd_cu_seq_len_zero_lens(): + """cu_seq_len form with degenerate lengths: a zero-length sequence + (repeated prefix value), an all-zero KV side (kernel dead-row path), and + the all-zero Q no-op keep the same semantics as the seq_len form.""" + + _run_thd_case(seq_q_lens=[128, 0, 64], seq_kv_lens=[100, 0, 0], is_causal=True, check_stats=True, cu_lens=True) + _run_thd_case(seq_q_lens=[64, 32], seq_kv_lens=[0, 0], check_stats=True, cu_lens=True) + _run_thd_case(seq_q_lens=[0, 0], seq_kv_lens=[50, 30], check_stats=True, cu_lens=True) + + @pytest.mark.L0 @torch_fork_set_rng(seed=9) def test_dsl_sm120_dense_flex_bhsd_contiguous(): diff --git a/test/python/sdpa/frost/test_sdpa_graph_analyzer.py b/test/python/sdpa/frost/test_sdpa_graph_analyzer.py index ab513da88..ac9b509ea 100644 --- a/test/python/sdpa/frost/test_sdpa_graph_analyzer.py +++ b/test/python/sdpa/frost/test_sdpa_graph_analyzer.py @@ -779,6 +779,43 @@ def test_sm120_probe_accepts_thd_stats(monkeypatch): assert engines.engine_name(arch="sm120") in _eligible(g) +def _mk_thd_cu_graph(*, extra_seq_len=False): + """Ragged (THD) graph carrying the cu_seq_len_q/kv (B+1,) prefix-sum form.""" + g = _mk_graph() + dims = (B, H, S, D) + strides = (S * H * D, D, H * D, 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) + cu_q = g.tensor(dim=(B + 1, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT32, name="cu_q") + cu_kv = g.tensor(dim=(B + 1, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT32, name="cu_kv") + kw = dict(cu_seq_len_q=cu_q, cu_seq_len_kv=cu_kv) + if extra_seq_len: + skv = g.tensor(dim=(B, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT32, name="skv") + kw["seq_len_kv"] = skv + o, _ = g.sdpa(name="s", q=q, k=k, v=v, attn_scale=0.1, is_inference=True, use_causal_mask=True, use_padding_mask=True, **kw) + _finish_output(o, dims, strides) + o.set_ragged_offset(ro) + return g + + +def test_probe_accepts_thd_cu_seq_len(): + """THD with the (B+1,) cu_seq_len prefix-sum form (cuDNN 9.24+) is served: + the lowering derives per-batch lengths host-side from its inherent tolist + round-trip.""" + assert engines.engine_name(512) in _eligible(_mk_thd_cu_graph()) + + +def test_probe_rejects_thd_cu_plus_seq_len(): + """Both forms on one side is ambiguous (the backend has its own + precedence, which the python engines do not replicate) — declined.""" + assert not _eligible(_mk_thd_cu_graph(extra_seq_len=True)) + + @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