frost(sdpa): dense-padded cu_seq_len support (kernel CU read mode)

python/cudnn/sdpa/fwd/api_dsl.py

@@ -331,9 +331,11 @@ def __init__(
         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).
+        # (B,) per-batch lengths. THD derives both forms host-side from its
+        # inherent tolist round-trip; dense graphs bind the (B+1,) tensor
+        # directly and the f16 kernels read len = cu[b+1] - cu[b] on device
+        # (the sync-free dense hot path stays sync-free). The FP8/MXFP8
+        # kernels are not plumbed (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)
@@ -856,10 +858,24 @@ def check_support(self) -> bool:
         )
         if self.thd:
             self.seq_kv_lens_present = True
+        # cu_seq_len form: THD consumes either form host-side; dense graphs
+        # use the f16 kernels' CU read mode (len = cu[b+1] - cu[b]). The
+        # FP8/MXFP8 execute paths are not plumbed for the (B+1,) form.
         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.cu_seq_q_lens or self.cu_seq_kv_lens) and self._fp8 and not self.thd,
+            "dense cu_seq_len_* is not plumbed for the FP8/MXFP8 kernels",
         )
+        if not self.thd:
+            # The cu flags declare the FORM of the corresponding lengths
+            # argument; the *_present flags still declare its PRESENCE.
+            self._value_error_if(
+                self.cu_seq_kv_lens and not self.seq_kv_lens_present,
+                "cu_seq_kv_lens declares the form of the KV lengths and requires seq_kv_lens_present=True",
+            )
+            self._value_error_if(
+                self.cu_seq_q_lens and not self.seq_q_lens_present,
+                "cu_seq_q_lens declares the form of the Q lengths and requires seq_q_lens_present=True",
+            )
         # Dense padded-Q trim backstops (engines.lower_dsl_prefill never sets
         # these combinations; a direct caller could).
         self._value_error_if(
@@ -941,6 +957,10 @@ def compile(self) -> None:
             has_sink=self.has_sink,
             seq_kv_lens_present=self.seq_kv_lens_present,
             seq_q_lens_present=self.seq_q_lens_present,
+            # Dense-only kernel CU read mode; the THD cu form is consumed
+            # host-side and compiles to the same THD specialization.
+            seq_kv_lens_cu=self.cu_seq_kv_lens and not self.thd,
+            seq_q_lens_cu=self.cu_seq_q_lens and not self.thd,
             sched_policy=sched_policy,
             thd_varlen=self.thd,
             fused_ldtm_stat=fused_ldtm_stat,
@@ -1163,17 +1183,22 @@ def execute(
             else self._dummy("sinks", device, lambda: torch.zeros(self.h_q, dtype=torch.float32, device=device))
         )
         seq_kv_t = (
-            self._checked_seq_lens(seq_kv_lens, "seq_kv_lens")
+            (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"))
             if seq_kv_lens is not None
             else self._dummy("seq_kv", device, lambda: torch.zeros(self.batch_size, dtype=torch.int32, device=device))
         )
         # Dense padded-Q trim: per-batch Q lengths are their OWN kernel
         # parameter (compiled in only when seq_q_lens_present — the kernel
         # signature is specialized on `None`, so the flag-off ABI is
-        # unchanged). The caller's (B,)-int32 device tensor is bound directly
+        # unchanged). The caller's device tensor — (B,)-int32 lengths or the
+        # (B+1,)-int32 cu prefix sums under cu_seq_q_lens — is bound directly
         # as a validated view — zero allocations/copies on the execute hot
         # path, stable pointer (CUDA-graph-capture friendly).
-        seq_q_t = self._checked_seq_lens(seq_q_lens, "seq_q_lens") if self.seq_q_lens_present else None
+        seq_q_t = (
+            (self._checked_cu_seq_lens(seq_q_lens, "cu_seq_len_q") if self.cu_seq_q_lens else self._checked_seq_lens(seq_q_lens, "seq_q_lens"))
+            if self.seq_q_lens_present
+            else None
+        )
         o_desc_dummy = self._dummy("o_desc", device, lambda: torch.zeros(1, dtype=torch.int64, device=device))
 
         import cutlass
@@ -1770,10 +1795,19 @@ 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)",
-        )
+        # cu_seq_len form: THD consumes either form host-side; dense graphs
+        # use the f16 kernel's CU read mode (len = cu[b+1] - cu[b]). The cu
+        # flags declare the FORM of the corresponding lengths argument; the
+        # *_present flags still declare its PRESENCE.
+        if not self.thd:
+            self._value_error_if(
+                self.cu_seq_kv_lens and not self.seq_kv_lens_present,
+                "cu_seq_kv_lens declares the form of the KV lengths and requires seq_kv_lens_present=True",
+            )
+            self._value_error_if(
+                self.cu_seq_q_lens and not self.seq_q_lens_present,
+                "cu_seq_q_lens declares the form of the Q lengths and requires seq_q_lens_present=True",
+            )
         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}",
@@ -1896,6 +1930,12 @@ def check_support(self) -> bool:
             )
             self._value_error_if(self.has_sink, "SM120 fp8 does not support attention sinks (Amax_S semantics)")
             self._value_error_if(self.seq_q_lens_present and not self.thd, "SM120 fp8 does not support per-batch seq_len_q")
+            # The fp8 kernel has no dense CU read mode (its template rejects
+            # the cu flags); THD cu is consumed host-side and stays served.
+            self._not_implemented_error_if(
+                (self.cu_seq_q_lens or self.cu_seq_kv_lens) and not self.thd,
+                "cu_seq_len_* is not plumbed for the SM120 fp8 kernel",
+            )
             self._value_error_if(
                 any(d not in _SM120_FP8_HEAD_TILES for d in (d_q, d_v)),
                 f"SM120 fp8 requires D_QK and D_V to be multiples of 32 within 32..256 (k32 contraction and 1-byte "
@@ -1998,6 +2038,10 @@ def compile(self) -> None:
             bottom_right=self.causal_bottom_right,
             seq_q_lens_present=self.seq_q_lens_present,
             seq_kv_lens_present=self.seq_kv_lens_present,
+            # Dense-only kernel CU read mode; the THD cu form is consumed
+            # host-side and compiles to the same THD specialization.
+            seq_q_lens_cu=self.cu_seq_q_lens and not self.thd,
+            seq_kv_lens_cu=self.cu_seq_kv_lens and not self.thd,
             has_sink=self.has_sink,
             thd_varlen=self.thd,
             q_tile=self.q_tile,
@@ -2117,7 +2161,7 @@ def execute(
         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 = (
-            self._checked_seq_lens(seq_q_lens, "seq_q_lens")
+            (self._checked_cu_seq_lens(seq_q_lens, "cu_seq_len_q") if self.cu_seq_q_lens else self._checked_seq_lens(seq_q_lens, "seq_q_lens"))
             if seq_q_lens is not None
             else self._dummy(
                 "seq_q_lens",
@@ -2126,7 +2170,7 @@ def execute(
             )
         )
         seq_kv_lens = (
-            self._checked_seq_lens(seq_kv_lens, "seq_kv_lens")
+            (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"))
             if seq_kv_lens is not None
             else self._dummy(
                 "seq_kv_lens",

python/cudnn/sdpa/fwd/config_sm100.py

@@ -85,6 +85,13 @@ class TemplateParams:
     # convention). Dense-only — THD carries per-sequence Q lengths via
     # cu_seqlens instead.
     seq_q_lens_present: bool = False
+    # cu_seq_len form (cuDNN 9.24+): the corresponding seq-lens kernel
+    # parameter is the (B+1,)-int32 PREFIX-SUM tensor instead of (B,)
+    # per-batch lengths; the kernels read len = cu[b+1] - cu[b] on device (the
+    # sync-free dense hot path stays sync-free). Dense-only — THD already
+    # carries prefix sums in its packed [kv_lens | cu_q | cu_kv] metadata.
+    seq_kv_lens_cu: bool = False
+    seq_q_lens_cu: bool = False
     sched_policy: int = SCHED_NATURAL
     thd_varlen: bool = False
     # cc10.3+ fuses the S_acc row-max into the LDTM (tcgen05.ld.red.f32.max); cc10.0
@@ -120,6 +127,13 @@ def _validate_params(flavor: str, k: TemplateParams) -> None:
             raise ValueError(f"{flavor}: SEQ_Q_LENS_PRESENT is dense-only (THD carries per-sequence Q lengths via cu_seqlens)")
         if not k.seq_kv_lens_present:
             raise ValueError(f"{flavor}: SEQ_Q_LENS_PRESENT requires SEQ_KV_LENS_PRESENT (padding mask)")
+    if k.seq_kv_lens_cu or k.seq_q_lens_cu:
+        if k.thd_varlen:
+            raise ValueError(f"{flavor}: seq_*_lens_cu is dense-only (the THD metadata buffer already carries cu_seqlens)")
+        if k.seq_kv_lens_cu and not k.seq_kv_lens_present:
+            raise ValueError(f"{flavor}: SEQ_KV_LENS_CU declares the FORM of the KV lengths and requires SEQ_KV_LENS_PRESENT")
+        if k.seq_q_lens_cu and not k.seq_q_lens_present:
+            raise ValueError(f"{flavor}: SEQ_Q_LENS_CU declares the FORM of the Q lengths and requires SEQ_Q_LENS_PRESENT")
     if k.sched_policy not in (SCHED_NATURAL, SCHED_LPT):
         raise ValueError(f"{flavor}: only SCHED_NATURAL (0) / SCHED_LPT (1) are wired up; got {k.sched_policy}")
 
@@ -288,6 +302,10 @@ class CfgD256:
 
     SEQ_KV_LENS_PRESENT: int = 0
     SEQ_Q_LENS_PRESENT: int = 0
+    # cu_seq_len form: the corresponding lens parameter is the (B+1,)
+    # prefix-sum tensor; per-batch length = cu[b+1] - cu[b]. Dense-only.
+    SEQ_KV_LENS_CU: int = 0
+    SEQ_Q_LENS_CU: int = 0
 
     THD_VARLEN: int = 0
 
@@ -335,6 +353,8 @@ def make_cfg_d256(params: TemplateParams) -> Tuple[CfgD256, TmaIters]:
         SCHEDULER_POLICY=params.sched_policy,
         SEQ_KV_LENS_PRESENT=1 if (params.thd_varlen or params.seq_kv_lens_present) else 0,
         SEQ_Q_LENS_PRESENT=int(params.seq_q_lens_present),
+        SEQ_KV_LENS_CU=int(params.seq_kv_lens_cu),
+        SEQ_Q_LENS_CU=int(params.seq_q_lens_cu),
         THD_VARLEN=int(params.thd_varlen),
     )
     _validate_cfg_d256(cfg)
@@ -427,6 +447,10 @@ class CfgD512:
 
     SEQ_KV_LENS_PRESENT: int = 0
     SEQ_Q_LENS_PRESENT: int = 0
+    # cu_seq_len form: the corresponding lens parameter is the (B+1,)
+    # prefix-sum tensor; per-batch length = cu[b+1] - cu[b]. Dense-only.
+    SEQ_KV_LENS_CU: int = 0
+    SEQ_Q_LENS_CU: int = 0
 
     THD_VARLEN: int = 0
 
@@ -478,6 +502,8 @@ def make_cfg_d512(params: TemplateParams) -> Tuple[CfgD512, TmaIters]:
         SCHEDULER_POLICY=params.sched_policy,
         SEQ_KV_LENS_PRESENT=1 if (params.thd_varlen or params.seq_kv_lens_present) else 0,
         SEQ_Q_LENS_PRESENT=int(params.seq_q_lens_present),
+        SEQ_KV_LENS_CU=int(params.seq_kv_lens_cu),
+        SEQ_Q_LENS_CU=int(params.seq_q_lens_cu),
         THD_VARLEN=int(params.thd_varlen),
     )
     _validate_cfg_d512(cfg)
@@ -573,6 +599,10 @@ class CfgD128:
 
     SEQ_KV_LENS_PRESENT: int = 0
     SEQ_Q_LENS_PRESENT: int = 0
+    # cu_seq_len form: the corresponding lens parameter is the (B+1,)
+    # prefix-sum tensor; per-batch length = cu[b+1] - cu[b]. Dense-only.
+    SEQ_KV_LENS_CU: int = 0
+    SEQ_Q_LENS_CU: int = 0
 
     THD_VARLEN: int = 0
 
@@ -638,6 +668,8 @@ def make_cfg_d128(params: TemplateParams) -> Tuple[CfgD128, TmaIters]:
         SCHEDULER_POLICY=params.sched_policy,
         SEQ_KV_LENS_PRESENT=1 if (params.thd_varlen or params.seq_kv_lens_present) else 0,
         SEQ_Q_LENS_PRESENT=int(params.seq_q_lens_present),
+        SEQ_KV_LENS_CU=int(params.seq_kv_lens_cu),
+        SEQ_Q_LENS_CU=int(params.seq_q_lens_cu),
         THD_VARLEN=int(params.thd_varlen),
     )
     _validate_cfg_d128(cfg)
@@ -709,6 +741,8 @@ def make_cfg_d192(params: TemplateParams) -> Tuple[CfgD192, TmaIters]:
         CORRECTION_REGS=40 if _mask_flags_from(params) == MASK_NONE else 88,
         SEQ_KV_LENS_PRESENT=1 if (params.thd_varlen or params.seq_kv_lens_present) else 0,
         SEQ_Q_LENS_PRESENT=int(params.seq_q_lens_present),
+        SEQ_KV_LENS_CU=int(params.seq_kv_lens_cu),
+        SEQ_Q_LENS_CU=int(params.seq_q_lens_cu),
         THD_VARLEN=int(params.thd_varlen),
     )
     _validate_cfg_d192(cfg)

python/cudnn/sdpa/fwd/config_sm120.py

@@ -61,6 +61,12 @@ class TemplateParams:
     bottom_right: bool = False
     seq_q_lens_present: bool = False
     seq_kv_lens_present: bool = False
+    # cu_seq_len form (cuDNN 9.24+): the corresponding seq-lens kernel
+    # argument is the (B+1,)-int32 PREFIX-SUM tensor instead of (B,)
+    # per-batch lengths; the kernel reads len = cu[b+1] - cu[b] on device.
+    # Dense-only — THD already carries prefix sums in its metadata tensor.
+    seq_q_lens_cu: bool = False
+    seq_kv_lens_cu: bool = False
     has_sink: bool = False
     thd_varlen: bool = False
     q_tile: int = SEQ_Q_TILES[0]
@@ -71,14 +77,16 @@ def validate_params(
     params: TemplateParams,
     allowed_dtypes: tuple[int, ...] = (DTYPE_BF16, DTYPE_FP16),
     allow_right_band: bool = True,
+    allow_cu: bool = True,
 ) -> None:
     """Validate the SM120 template specialization.
 
     Reachable failures should already have been rejected by the engine
     capabilities or adapter support checks; this validation is a backstop for
     direct template use. ``allowed_dtypes`` defaults to the FP16/BF16 template's
     set; the FP8 template passes its own. ``allow_right_band=False`` also
-    rejects a widened right band, which the FP8 template does not plumb.
+    rejects a widened right band, which the FP8 template does not plumb;
+    ``allow_cu=False`` likewise rejects the dense cu_seq_len read mode.
     """
 
     if params.dtype_qkv not in allowed_dtypes:
@@ -100,3 +108,12 @@ def validate_params(
             raise ValueError("SM120 SDPA: thd_varlen requires seq_kv_lens_present (the THD metadata tensor)")
         if params.seq_q_lens_present:
             raise ValueError("SM120 SDPA: seq_q_lens_present is dense-only (THD carries per-sequence Q lengths via cu_seqlens)")
+    if params.seq_q_lens_cu or params.seq_kv_lens_cu:
+        if not allow_cu:
+            raise ValueError("SM120 SDPA: the dense cu_seq_len read mode is not plumbed for this template")
+        if params.thd_varlen:
+            raise ValueError("SM120 SDPA: seq_*_lens_cu is dense-only (the THD metadata tensor already carries cu_seqlens)")
+        if params.seq_q_lens_cu and not params.seq_q_lens_present:
+            raise ValueError("SM120 SDPA: seq_q_lens_cu declares the FORM of the Q lengths and requires seq_q_lens_present")
+        if params.seq_kv_lens_cu and not params.seq_kv_lens_present:
+            raise ValueError("SM120 SDPA: seq_kv_lens_cu declares the FORM of the KV lengths and requires seq_kv_lens_present")