Support mixed-form sequence lengths in SDPA forward (cuDNN 9.25+)

docs/operations/Attention.md

@@ -321,10 +321,10 @@ graph.sdpa(
 - `bias` (Optional[cudnn_tensor]): Additive bias mask for attention scores. Supports broadcasting.
 - `block_mask` (Optional[cudnn_tensor]): Block-level mask for 128x128 tiles. Only supported with UNIFIED implementation.
 - `use_alibi_mask` (Optional[bool]): Enable ALiBi (Attention with Linear Biases) positional encoding. Requires `diagonal_band_right_bound=0`.
-- `use_padding_mask` (Optional[bool]): Enable variable sequence length masking. Must also provide `seq_len_q` and `seq_len_kv`.
+- `use_padding_mask` (Optional[bool]): Enable variable sequence length masking. Must also provide a Q-side and a KV-side length tensor, each in per-batch (`seq_len_q`/`seq_len_kv`) or cumulative (`cu_seq_len_q`/`cu_seq_len_kv`) form.
 - `seq_len_q` (Optional[cudnn_tensor]): Per-batch query sequence lengths with shape $(B, 1, 1, 1)$.
 - `seq_len_kv` (Optional[cudnn_tensor]): Per-batch key/value sequence lengths with shape $(B, 1, 1, 1)$.
-- `cu_seq_len_q` (Optional[cudnn_tensor]): Cumulative query sequence lengths (prefix sums with a leading 0) with shape $(B+1, 1, 1, 1)$ or 1-D $(B+1,)$ (promoted automatically), int32 or int64. Mutually exclusive with `seq_len_q`/`seq_len_kv`; must be set together with `cu_seq_len_kv` and requires `use_padding_mask=True`. Requires cuDNN 9.24+ and the UNIFIED implementation.
+- `cu_seq_len_q` (Optional[cudnn_tensor]): Cumulative query sequence lengths (prefix sums with a leading 0) with shape $(B+1, 1, 1, 1)$ or 1-D $(B+1,)$ (promoted automatically), int32 or int64. Mutually exclusive with `seq_len_q` (a side cannot use both forms); a KV-side length tensor (`seq_len_kv` or `cu_seq_len_kv`) must also be provided, and `use_padding_mask=True` is required. The two sides may use different forms (e.g. `cu_seq_len_q` with `seq_len_kv`), which requires cuDNN 9.25+. Supplying `cu_seq_len_q` requires cuDNN 9.24+ and the UNIFIED implementation.
 - `cu_seq_len_kv` (Optional[cudnn_tensor]): Cumulative key/value sequence lengths; same shape, type, and constraints as `cu_seq_len_q`.
 - `diagonal_alignment` (Optional[cudnn.diagonal_alignment]): Alignment for diagonal masking. `TOP_LEFT` for standard causal, `BOTTOM_RIGHT` for prefix-LM style.
 - `diagonal_band_left_bound` (Optional[int]): Left bound for sliding window attention. Masks columns at or before `row_idx - left_bound`.
@@ -911,7 +911,7 @@ Args:
     use_padding_mask (Optional[bool]): Enable variable sequence length masking. Requires seq_len_q/seq_len_kv or cu_seq_len_q/cu_seq_len_kv. Default is False.
     seq_len_q (Optional[cudnn_tensor]): Per-batch query sequence lengths with shape (B, 1, 1, 1).
     seq_len_kv (Optional[cudnn_tensor]): Per-batch key/value sequence lengths with shape (B, 1, 1, 1).
-    cu_seq_len_q (Optional[cudnn_tensor]): Cumulative query sequence lengths (prefix sums with a leading 0) with shape (B+1, 1, 1, 1) or 1-D (B+1,), int32 or int64. Mutually exclusive with seq_len_q/seq_len_kv; must be set together with cu_seq_len_kv. Requires cuDNN 9.25+ and the UNIFIED implementation.
+    cu_seq_len_q (Optional[cudnn_tensor]): Cumulative query sequence lengths (prefix sums with a leading 0) with shape (B+1, 1, 1, 1) or 1-D (B+1,), int32 or int64. Mutually exclusive with seq_len_q (a side cannot use both forms); a KV-side length tensor (seq_len_kv or cu_seq_len_kv) must also be provided. The two sides may use different forms (e.g. cu_seq_len_q with seq_len_kv). Requires cuDNN 9.25+ and the UNIFIED implementation (the FP8 path requires 9.25+ for cumulative sequence lengths in any form).
     cu_seq_len_kv (Optional[cudnn_tensor]): Cumulative key/value sequence lengths; same shape, type, and constraints as cu_seq_len_q.
     compute_data_type (Optional[cudnn.data_type]): The data type for computation. Default is NOT_SET.
     name (Optional[str]): The name of the operation.

include/cudnn_frontend/node/diagonal_band_mask.h

@@ -29,10 +29,12 @@ class DiagonalBandMaskNodeBase : public NodeCRTP<DerivedT> {
         RETURN_CUDNN_FRONTEND_ERROR_IF(has_left_bound() && has_shift_right_bound(),
                                        error_code_t::INVALID_VALUE,
                                        "DiagonalBandMaskNode cannot have both left_bound and shift_right_bound");
-        RETURN_CUDNN_FRONTEND_ERROR_IF(
-            (has_seq_len_q() || has_seq_len_kv()) && (has_cu_seq_len_q() || has_cu_seq_len_kv()),
-            error_code_t::INVALID_VALUE,
-            "SEQ_LEN_Q / SEQ_LEN_KV and CU_SEQ_LEN_Q / CU_SEQ_LEN_KV are mutually exclusive");
+        RETURN_CUDNN_FRONTEND_ERROR_IF(has_seq_len_q() && has_cu_seq_len_q(),
+                                       error_code_t::INVALID_VALUE,
+                                       "SEQ_LEN_Q and CU_SEQ_LEN_Q are mutually exclusive");
+        RETURN_CUDNN_FRONTEND_ERROR_IF(has_seq_len_kv() && has_cu_seq_len_kv(),
+                                       error_code_t::INVALID_VALUE,
+                                       "SEQ_LEN_KV and CU_SEQ_LEN_KV are mutually exclusive");
 
         return pre_validate_node_extra();
     }

include/cudnn_frontend/node/sdpa_support_surface.h

@@ -84,26 +84,32 @@ SDPA_attributes::validate_sdpa_support_surface(const detail::Context& context,
                                    error_code_t::GRAPH_NOT_SUPPORTED,
                                    "Bias mask data type cannot be boolean");
 
-    // validate options for padding mask: padding requires per-sequence length tensors,
-    // either as (SEQ_LEN_Q + SEQ_LEN_KV) or as (CU_SEQ_LEN_Q + CU_SEQ_LEN_KV).
-    RETURN_CUDNN_FRONTEND_ERROR_IF(
-        padding_mask && (!has_seq_len_q || !has_seq_len_kv) && (!has_cu_seq_len_q || !has_cu_seq_len_kv),
-        error_code_t::ATTRIBUTE_NOT_SET,
-        "Padding mask requires seq_len_q/seq_len_kv (or cu_seq_len_q/cu_seq_len_kv) to be set.");
+    // validate options for padding mask: padding requires a per-sequence length
+    // tensor on each side. Each side independently uses exactly one
+    // representation -- per-batch (seq_len_*) or cumulative (cu_seq_len_*) --
+    // and the two sides may use different forms, e.g. cumulative Q lengths
+    // with per-batch KV lengths.
+    bool const has_any_seq_len_q  = has_seq_len_q || has_cu_seq_len_q;
+    bool const has_any_seq_len_kv = has_seq_len_kv || has_cu_seq_len_kv;
+    RETURN_CUDNN_FRONTEND_ERROR_IF(padding_mask && !(has_any_seq_len_q && has_any_seq_len_kv),
+                                   error_code_t::ATTRIBUTE_NOT_SET,
+                                   "Padding mask requires seq_len_q/seq_len_kv (or cu_seq_len_q/cu_seq_len_kv) "
+                                   "to be set.");
     RETURN_CUDNN_FRONTEND_ERROR_IF(
-        (!padding_mask && !attention_score_modifier) &&
-            (has_seq_len_q || has_seq_len_kv || has_cu_seq_len_q || has_cu_seq_len_kv),
+        (!padding_mask && !attention_score_modifier) && (has_any_seq_len_q || has_any_seq_len_kv),
         error_code_t::ATTRIBUTE_NOT_SET,
         "seq_len_q/seq_len_kv (or cu_seq_len_q/cu_seq_len_kv) needs to be set only if padding mask is enabled.");
 
-    // Cumulative sequence length tensors must be set together, and are mutually
-    // exclusive with the (per-batch) seq_len tensors.
-    RETURN_CUDNN_FRONTEND_ERROR_IF(has_cu_seq_len_q != has_cu_seq_len_kv,
-                                   error_code_t::ATTRIBUTE_NOT_SET,
-                                   "cu_seq_len_q and cu_seq_len_kv must both be set or both unset.");
-    RETURN_CUDNN_FRONTEND_ERROR_IF((has_cu_seq_len_q || has_cu_seq_len_kv) && (has_seq_len_q || has_seq_len_kv),
+    RETURN_CUDNN_FRONTEND_ERROR_IF(has_seq_len_q && has_cu_seq_len_q,
+                                   error_code_t::INVALID_VALUE,
+                                   "seq_len_q and cu_seq_len_q cannot both be set.");
+    RETURN_CUDNN_FRONTEND_ERROR_IF(has_seq_len_kv && has_cu_seq_len_kv,
                                    error_code_t::INVALID_VALUE,
-                                   "Cannot specify both seq_len tensors and cu_seq_len tensors.");
+                                   "seq_len_kv and cu_seq_len_kv cannot both be set.");
+    RETURN_CUDNN_FRONTEND_ERROR_IF(has_any_seq_len_q != has_any_seq_len_kv,
+                                   error_code_t::ATTRIBUTE_NOT_SET,
+                                   "A Q-side and a KV-side sequence length tensor must be provided together "
+                                   "(each side may independently use seq_len_* or cu_seq_len_*).");
 
     RETURN_CUDNN_FRONTEND_ERROR_IF(is_ragged && ((padding_mask == false) && (attention_score_modifier == nullptr)),
                                    error_code_t::GRAPH_NOT_SUPPORTED,
@@ -301,10 +307,10 @@ SDPA_attributes::validate_sdpa_support_surface(const detail::Context& context,
                                        "Paged caches are not supported in combination with ragged offsets.");
 
         RETURN_CUDNN_FRONTEND_ERROR_IF(
-            is_paged && !((has_seq_len_q && has_seq_len_kv) || (has_cu_seq_len_q && has_cu_seq_len_kv)),
+            is_paged && !(has_any_seq_len_q && has_any_seq_len_kv),
             error_code_t::GRAPH_NOT_SUPPORTED,
             "Paged caches can only be used in combination with padding mask and variable sequence lengths "
-            "for both Q and KV (via seq_len_q/seq_len_kv or cu_seq_len_q/cu_seq_len_kv).");
+            "for both Q and KV (each side independently via seq_len_* or cu_seq_len_*).");
 
         RETURN_CUDNN_FRONTEND_ERROR_IF(
             !is_paged && max_seq_kv_explicit,
@@ -389,21 +395,24 @@ SDPA_attributes::validate_sdpa_support_surface(const detail::Context& context,
 inline error_t
 SDPA_attributes::verify_sdpa_support_surface_for_implementation(const detail::Context& context,
                                                                 AttentionImplementation_t impl) const {
+    auto const has_input = [this](input_names name) {
+        auto const it = inputs.find(name);
+        return it != inputs.end() && it->second != nullptr;
+    };
+
     switch (impl) {
         case AttentionImplementation_t::AUTO:
             // This function should not be called with AUTO.
             return {error_code_t::INVALID_VALUE,
                     "Can't call verify_sdpa_support_surface_for_implementation with impl=AUTO"};
         case AttentionImplementation_t::COMPOSITE:
-            for (const auto& [key, value] : inputs) {
-                RETURN_CUDNN_FRONTEND_ERROR_IF(key == input_names::Block_mask && value != nullptr,
-                                               error_code_t::GRAPH_NOT_SUPPORTED,
-                                               "Composite SDPA node doesn't support Block_mask input");
-                RETURN_CUDNN_FRONTEND_ERROR_IF(
-                    (key == input_names::CU_SEQ_LEN_Q || key == input_names::CU_SEQ_LEN_KV) && value != nullptr,
-                    error_code_t::GRAPH_NOT_SUPPORTED,
-                    "Composite SDPA node doesn't support CU_SEQ_LEN_Q / CU_SEQ_LEN_KV inputs");
-            }
+            RETURN_CUDNN_FRONTEND_ERROR_IF(has_input(input_names::Block_mask),
+                                           error_code_t::GRAPH_NOT_SUPPORTED,
+                                           "Composite SDPA node doesn't support Block_mask input");
+            RETURN_CUDNN_FRONTEND_ERROR_IF(
+                has_input(input_names::CU_SEQ_LEN_Q) || has_input(input_names::CU_SEQ_LEN_KV),
+                error_code_t::GRAPH_NOT_SUPPORTED,
+                "Composite SDPA node doesn't support CU_SEQ_LEN_Q / CU_SEQ_LEN_KV inputs");
             // The ragged offset multiplier is only supported by the unified forward engine.
             // Reject it here so auto-select routes such graphs to the unified implementation.
             for (const auto& [key, value] : inputs) {
@@ -459,6 +468,14 @@ SDPA_attributes::verify_sdpa_support_surface_for_implementation(const detail::Co
                 allowed_input_msg += ", CU_SEQ_LEN_Q, CU_SEQ_LEN_KV";
             }
 
+            if (effective_cudnn_ver < 92500) {
+                RETURN_CUDNN_FRONTEND_ERROR_IF(
+                    has_input(input_names::CU_SEQ_LEN_Q) != has_input(input_names::CU_SEQ_LEN_KV),
+                    error_code_t::GRAPH_NOT_SUPPORTED,
+                    "Mixed-form sequence lengths (cu_seq_len on one side with seq_len on the other) "
+                    "require cuDNN 9.25.0 or above");
+            }
+
             if (effective_cudnn_ver >= 92500) {
                 allowed_input_names.insert({input_names::Descale_Q,
                                             input_names::Descale_K,

python/pygraph/sdpa.cpp

@@ -1138,8 +1138,8 @@ init_pygraph_sdpa_submodule(py::class_<PyGraph>& m) {
                     score_sum_exp (Optional[cudnn_tensor]): The numerically stable sum of exponents using normalized values wrt max score.
                     sink_token (Optional[cudnn_tensor]): The sink attention token tensor. Shape is (1, h_q, 1, 1), type is float32.
                     unfuse_fma (Optional[bool]): For SM100: use unfused __fmul_rn + __fadd_rn instead of ffma2 in softmax. Default is False.
-                    cu_seq_len_q (Optional[cudnn_tensor]): Cumulative sequence length of the query, shape (b+1, 1, 1, 1) or 1-D (b+1,) (promoted automatically), int32 or int64. Mutually exclusive with seq_len_q. Requires cuDNN 9.24.0 or newer and the UNIFIED implementation.
-                    cu_seq_len_kv (Optional[cudnn_tensor]): Cumulative sequence length of the key, shape (b+1, 1, 1, 1) or 1-D (b+1,) (promoted automatically), int32 or int64. Mutually exclusive with seq_len_kv. Requires cuDNN 9.24.0 or newer and the UNIFIED implementation.
+                    cu_seq_len_q (Optional[cudnn_tensor]): Cumulative sequence length of the query, shape (b+1, 1, 1, 1) or 1-D (b+1,) (promoted automatically), int32 or int64. Mutually exclusive with seq_len_q; pair with a KV-side tensor (seq_len_kv or cu_seq_len_kv) and set use_padding_mask=True. Requires cuDNN 9.24+ and UNIFIED (9.25+ if the two sides use different forms).
+                    cu_seq_len_kv (Optional[cudnn_tensor]): Cumulative sequence length of the key, shape (b+1, 1, 1, 1) or 1-D (b+1,) (promoted automatically), int32 or int64. Mutually exclusive with seq_len_kv; pair with a Q-side tensor (seq_len_q or cu_seq_len_q) and set use_padding_mask=True. Requires cuDNN 9.24+ and UNIFIED (9.25+ if the two sides use different forms).
                 Preferred masking Args:
                     diagonal_alignment (Optional[cudnn.diagonal_alignment]): One of {"TOP_LEFT", "BOTTOM_RIGHT"}. E.g., causal masking can be performed by setting diagonal_alignment=TOP_LEFT, and diagonal_band_right_bound=0. Default is TOP_LEFT.
                     diagonal_band_left_bound (Optional[int]): An integer >= 1 specifying the offset to the left of the main diagonal to attend to. Default is None, implying +Inf.
@@ -1304,8 +1304,8 @@ init_pygraph_sdpa_submodule(py::class_<PyGraph>& m) {
                     sink_token (Optional[cudnn_tensor]): Sink token bias for streaming attention. Default is None.
                     unfuse_fma (Optional[bool]): For SM100: use unfused __fmul_rn + __fadd_rn instead of ffma2 in softmax. Default is False.
                     implementation (Optional[cudnn.attention_implementation]): Which underlying implementation to use in the cuDNN backend. Default is AUTO (recommended).
-                    cu_seq_len_q (Optional[cudnn_tensor]): Cumulative sequence length of the query, shape (b+1, 1, 1, 1) or 1-D (b+1,) (promoted automatically), int32 or int64. Mutually exclusive with seq_len_q; must be set together with cu_seq_len_kv and requires use_padding_mask=True. Requires cuDNN 9.25.0 or newer and the UNIFIED implementation.
-                    cu_seq_len_kv (Optional[cudnn_tensor]): Cumulative sequence length of the key, shape (b+1, 1, 1, 1) or 1-D (b+1,) (promoted automatically), int32 or int64. Mutually exclusive with seq_len_kv; must be set together with cu_seq_len_q and requires use_padding_mask=True. Requires cuDNN 9.25.0 or newer and the UNIFIED implementation.
+                    cu_seq_len_q (Optional[cudnn_tensor]): Cumulative sequence length of the query, shape (b+1, 1, 1, 1) or 1-D (b+1,) (promoted automatically), int32 or int64. Mutually exclusive with seq_len_q; pair with a KV-side tensor (seq_len_kv or cu_seq_len_kv) and set use_padding_mask=True. The two sides may use different forms. Requires cuDNN 9.25+ and UNIFIED.
+                    cu_seq_len_kv (Optional[cudnn_tensor]): Cumulative sequence length of the key, shape (b+1, 1, 1, 1) or 1-D (b+1,) (promoted automatically), int32 or int64. Mutually exclusive with seq_len_kv; pair with a Q-side tensor (seq_len_q or cu_seq_len_q) and set use_padding_mask=True. The two sides may use different forms. Requires cuDNN 9.25+ and UNIFIED.
                 Preferred masking Args:
                     diagonal_alignment (Optional[cudnn.diagonal_alignment]): One of {"TOP_LEFT", "BOTTOM_RIGHT"}. E.g., causal masking can be performed by setting diagonal_alignment=TOP_LEFT, and right_bound=0. Default is TOP_LEFT.
                     left_bound (Optional[int]): An integer >= 1 specifying the offset to the left of the main diagonal to attend to. Default is None, implying +Inf.

samples/cpp/sdpa/fp16_fwd_with_cu_seq_len.cpp

@@ -43,12 +43,15 @@ For example, with b=3 and actual_seq_len = {12, 20, 8}, cu_seq_len = {0, 12, 32,
 A 1-D (b+1,) tensor is also accepted and promoted automatically to (b+1, 1, 1, 1).
 
 Constraints (enforced by the frontend; see SDPA_attributes::validate_sdpa_support_surface):
-    - cu_seq_len_q and cu_seq_len_kv must both be set or both unset.
-    - cu_seq_len_* are mutually exclusive with seq_len_q / seq_len_kv.
-    - padding_mask must be true when cu_seq_len_* are set.
+    - A Q-side and a KV-side sequence length tensor must both be provided (or both
+      omitted). Each side independently uses at most one form: per-batch (seq_len_*)
+      or cumulative (cu_seq_len_*); the two forms cannot both be set on the same side.
+    - The two sides may use different forms (e.g. cu_seq_len_q with seq_len_kv). Such
+      mixed forms require cuDNN >= 9.25.0; this sample uses cu_seq_len on both sides.
+    - padding_mask must be true when sequence length tensors are set.
     - Only the UNIFIED SDPA implementation supports cu_seq_len_*; the COMPOSITE path
       will reject the inputs explicitly.
-    - Requires cuDNN >= 9.24.0.
+    - Supplying cu_seq_len_* requires cuDNN >= 9.24.0.
 */
 
 // Tensors in forward pass