frost(sdpa): accept the cu_seq_len (prefix-sum) length form for THD

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]

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,

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")