Expose cu_seq_len_q/kv on the sdpa_fp8 python binding

docs/operations/Attention.md

@@ -70,6 +70,10 @@ The support matrix is based on the latest cudnn backend version 9.18.1
 - Contains cumulative token offsets in **elements** (not bytes)
 - Last element is the total number of tokens
 
+   **Ragged Offset Multiplier (cuDNN 9.24+, UNIFIED forward only):**
+- `tensor.set_ragged_offset_multiplier(value)` lets the ragged offsets be stored in coarser units; the engine multiplies each offset by `value` to recover element offsets.
+- Example: with a multiplier of $H \times D$, a token-unit cumulative-sequence-length tensor (e.g. `cu_seq_len_q`) can be bound directly as the ragged offset, avoiding a conversion pass.
+
    **Memory Layout visualization:**
 
        *Example:*
@@ -216,6 +220,13 @@ SDPA_attributes& set_padding_mask(bool const value);
 // integer tensor that specifies the sequence length of each batch
 SDPA_attributes& set_seq_len_q(std::shared_ptr<Tensor_attributes> value);
 SDPA_attributes& set_seq_len_kv(std::shared_ptr<Tensor_attributes> value);
+
+// integer tensor of shape (B+1, 1, 1, 1) that specifies the cumulative sequence
+// lengths (prefix sums, leading 0) of each batch. Mutually exclusive with
+// set_seq_len_q/set_seq_len_kv; both tensors must be set together.
+// Requires cuDNN 9.24+ and the UNIFIED implementation.
+SDPA_attributes& set_cu_seq_len_q(std::shared_ptr<Tensor_attributes> value);
+SDPA_attributes& set_cu_seq_len_kv(std::shared_ptr<Tensor_attributes> value);
 // ==========================  END     var len options =====================
 
 // ========================== BEGIN score mod options =====================
@@ -284,6 +295,8 @@ graph.sdpa(
     use_padding_mask=False,               # Enable variable sequence length masking
     seq_len_q=None,                       # Per-batch query sequence lengths
     seq_len_kv=None,                      # Per-batch key/value sequence lengths
+    cu_seq_len_q=None,                    # Cumulative query sequence lengths (UNIFIED only)
+    cu_seq_len_kv=None,                   # Cumulative key/value sequence lengths (UNIFIED only)
     diagonal_alignment=TOP_LEFT,          # Diagonal alignment: TOP_LEFT or BOTTOM_RIGHT
     diagonal_band_left_bound=None,        # Left bound for sliding window (None = no bound)
     diagonal_band_right_bound=None,       # Right bound for causal mask (0 = causal, None = no bound)
@@ -311,6 +324,8 @@ graph.sdpa(
 - `use_padding_mask` (Optional[bool]): Enable variable sequence length masking. Must also provide `seq_len_q` and `seq_len_kv`.
 - `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_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`.
 - `diagonal_band_right_bound` (Optional[int]): Right bound for causal masking. Set to 0 for causal mask. Masks columns beyond `row_idx + right_bound`.
@@ -893,6 +908,11 @@ Args:
     scale_o (cudnn_tensor): Scale factor for output.
     attn_scale (Optional[Union[float, cudnn_tensor]]): The scale factor for attention. Default is None.
     use_causal_mask (Optional[bool]): Whether to use causal mask. Default is False.
+    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_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.
     generate_stats (Optional[bool]): If true, compute and output softmax stats (useful at training time). Default is None, but one of {generate_stats, is_inference} must be set.
@@ -911,12 +931,13 @@ Returns:
 The current FP8 support is a subset of the options supported in FP16 and BF16 support.
 - Attention scale (`attn_scale`): Applies a scaling factor to attention scores before the softmax, such as $\frac{1}{\sqrt{\text{d}}}$. Set to 1.0 by default.
 - Causal mask: Fills the upper triangular matrix of attention scores with negative infinity.
+- Padding mask (`use_padding_mask`): Variable sequence lengths, provided either as per-batch lengths (`seq_len_q`/`seq_len_kv`) or as cumulative sequence lengths (`cu_seq_len_q`/`cu_seq_len_kv`; cuDNN 9.25+, UNIFIED implementation only).
 
 #### Limitations
 
 - Requires Hopper (SM90) or newer architecture.
 - Head dimension must be a multiple of 16.
-- Limited masking options compared to FP16/BF16 (causal mask only).
+- Limited masking options compared to FP16/BF16 (causal and padding masks only).
 - Requires explicit scale/descale tensors for all FP8 inputs and outputs.
 
 #### Tensors

python/pygraph/pygraph.h

@@ -524,7 +524,9 @@ class PyGraph {
              std::shared_ptr<cudnn_frontend::graph::Tensor_attributes> score_sum_exp,
              std::shared_ptr<cudnn_frontend::graph::Tensor_attributes> sink_token,
              bool const unfuse_fma,
-             cudnn_frontend::AttentionImplementation_t const& implementation);
+             cudnn_frontend::AttentionImplementation_t const& implementation,
+             std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& cu_seq_len_q,
+             std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& cu_seq_len_kv);
 
     // MXFP8 SDPA forward - uses block-wise scale factors (E8M0 with F8_128x4 reordering)
     // return [o, stats, amax_o]

python/pygraph/sdpa.cpp

@@ -505,7 +505,6 @@ PyGraph::sdpa_backward(std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>
     return {dQ, dK, dV};
 }
 
-// Deprecated, use sdpa_unified instead
 std::array<std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>, 4>
 PyGraph::sdpa_fp8(std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& q,
                   std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& k,
@@ -542,12 +541,11 @@ PyGraph::sdpa_fp8(std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& q,
                   std::shared_ptr<cudnn_frontend::graph::Tensor_attributes> score_sum_exp,
                   std::shared_ptr<cudnn_frontend::graph::Tensor_attributes> sink_token,
                   bool const unfuse_fma,
-                  cudnn_frontend::AttentionImplementation_t const& implementation) {
+                  cudnn_frontend::AttentionImplementation_t const& implementation,
+                  std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& cu_seq_len_q,
+                  std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& cu_seq_len_kv) {
     cudnn_frontend::DataType_t mma_core_mode                             = cudnn_frontend::DataType_t::FP8_E4M3;
     std::shared_ptr<cudnn_frontend::graph::Tensor_attributes> block_mask = nullptr;
-    // cu_seq_len_q/cu_seq_len_kv are not exposed via the fp8 path.
-    std::shared_ptr<cudnn_frontend::graph::Tensor_attributes> cu_seq_len_q  = nullptr;
-    std::shared_ptr<cudnn_frontend::graph::Tensor_attributes> cu_seq_len_kv = nullptr;
 
     // Handle sliding_window to left_bound mapping for backward compatibility
     py::object actual_left_bound = left_bound;
@@ -1271,6 +1269,8 @@ init_pygraph_sdpa_submodule(py::class_<PyGraph>& m) {
           py::arg_v("sink_token", nullptr),
           py::arg_v("unfuse_fma", false),
           py::arg_v("implementation", cudnn_frontend::AttentionImplementation_t::AUTO),
+          py::arg_v("cu_seq_len_q", nullptr),
+          py::arg_v("cu_seq_len_kv", nullptr),
           R"pbdoc(
                 Perform scaled dot product attention with fp8 datatype inputs and outputs.
 
@@ -1304,6 +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.
                 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.

test/python/sdpa/fp8.py

@@ -49,6 +49,8 @@ class GraphFwdUid(IntEnum):
     o_ragged_offset = 20
     stats_ragged_offset = 21
     sink_token = 22
+    cu_seq_len_q = 23
+    cu_seq_len_kv = 24
 
 class GraphBwdUid(IntEnum):
     q = 100
@@ -87,12 +89,14 @@ class GraphBwdUid(IntEnum):
     sink_token = 133
     dSink_token = 134
 
-def generate_graph_fwd(cudnn_itype, cudnn_otype, b, h_q, h_k, h_v, s_qo, s_kv, d_qk, d_vo, attn_scale, block_size, is_ragged=False, generate_stats=True, left_bound=None, right_bound=None, diag_align=None, with_sink_token=False, implementation=cudnn.attention_implementation.AUTO):
+def generate_graph_fwd(cudnn_itype, cudnn_otype, b, h_q, h_k, h_v, s_qo, s_kv, d_qk, d_vo, attn_scale, block_size, is_ragged=False, generate_stats=True, left_bound=None, right_bound=None, diag_align=None, with_sink_token=False, is_cu_seq_len=False, with_ragged_offset_multiplier=False, implementation=cudnn.attention_implementation.AUTO):
     graph_fwd = cudnn.pygraph(io_data_type=cudnn_itype, intermediate_data_type=cudnn.data_type.FLOAT, compute_data_type=cudnn.data_type.FLOAT)
 
     use_padding_mask = None
     kv_seq_len = None
     q_seq_len = None
+    cu_seq_len_q = None
+    cu_seq_len_kv = None
     k_block_table = None
     v_block_table = None
 
@@ -121,8 +125,12 @@ def generate_graph_fwd(cudnn_itype, cudnn_otype, b, h_q, h_k, h_v, s_qo, s_kv, d
 
     if is_ragged:
         use_padding_mask = True
-        q_seq_len = graph_fwd.tensor(uid=GraphFwdUid.q_seq_len, dim=(b,), stride=(1,), data_type=cudnn.data_type.INT32)
-        kv_seq_len = graph_fwd.tensor(uid=GraphFwdUid.kv_seq_len, dim=(b,), stride=(1,), data_type=cudnn.data_type.INT32)
+        if is_cu_seq_len:
+            cu_seq_len_q = graph_fwd.tensor(uid=GraphFwdUid.cu_seq_len_q, dim=(b + 1,), stride=(1,), data_type=cudnn.data_type.INT32)
+            cu_seq_len_kv = graph_fwd.tensor(uid=GraphFwdUid.cu_seq_len_kv, dim=(b + 1,), stride=(1,), data_type=cudnn.data_type.INT32)
+        else:
+            q_seq_len = graph_fwd.tensor(uid=GraphFwdUid.q_seq_len, dim=(b,), stride=(1,), data_type=cudnn.data_type.INT32)
+            kv_seq_len = graph_fwd.tensor(uid=GraphFwdUid.kv_seq_len, dim=(b,), stride=(1,), data_type=cudnn.data_type.INT32)
 
         q_ragged_offset = graph_fwd.tensor(uid=int(GraphFwdUid.q_ragged_offset), dim=(b + 1,), stride=(1,), data_type=cudnn.data_type.INT64)
         k_ragged_offset = graph_fwd.tensor(uid=int(GraphFwdUid.k_ragged_offset), dim=(b + 1,), stride=(1,), data_type=cudnn.data_type.INT64)
@@ -132,6 +140,12 @@ def generate_graph_fwd(cudnn_itype, cudnn_otype, b, h_q, h_k, h_v, s_qo, s_kv, d
         q.set_ragged_offset(q_ragged_offset)
         k.set_ragged_offset(k_ragged_offset)
         v.set_ragged_offset(v_ragged_offset)
+        if with_ragged_offset_multiplier:
+            # Offsets are stored in coarser units (divided out in the allocation);
+            # the engine multiplies back to element offsets.
+            q.set_ragged_offset_multiplier(d_qk)
+            k.set_ragged_offset_multiplier(d_qk)
+            v.set_ragged_offset_multiplier(d_vo)
 
     q_descale = graph_fwd.tensor(uid=GraphFwdUid.q_descale, dim=(1, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.FLOAT)
     k_descale = graph_fwd.tensor(uid=GraphFwdUid.k_descale, dim=(1, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.FLOAT)
@@ -150,6 +164,7 @@ def generate_graph_fwd(cudnn_itype, cudnn_otype, b, h_q, h_k, h_v, s_qo, s_kv, d
         scale_s=s_scale, descale_s=s_descale, scale_o=o_scale,
         generate_stats=generate_stats, attn_scale=attn_scale, use_causal_mask=False,
         use_padding_mask=use_padding_mask, seq_len_kv=kv_seq_len, seq_len_q=q_seq_len,
+        cu_seq_len_q=cu_seq_len_q, cu_seq_len_kv=cu_seq_len_kv,
         paged_attention_k_table=k_block_table, paged_attention_v_table=v_block_table,
         paged_attention_max_seq_len_kv=s_kv,
         left_bound=left_bound, right_bound=right_bound,
@@ -165,6 +180,8 @@ def generate_graph_fwd(cudnn_itype, cudnn_otype, b, h_q, h_k, h_v, s_qo, s_kv, d
     o.set_uid(GraphFwdUid.o).set_output(True).set_dim((b, h_q, s_qo, d_vo)).set_stride(stride_o).set_data_type(cudnn_otype)
     if is_ragged:
         o.set_ragged_offset(o_ragged_offset)
+        if with_ragged_offset_multiplier:
+            o.set_ragged_offset_multiplier(d_vo)
 
     if generate_stats:
         stats_stride = (s_qo * h_q, 1, h_q, 1) if is_ragged else (s_qo * h_q, s_qo, 1, 1)
@@ -301,6 +318,13 @@ def exec_sdpa_fp8(cfg, request, cudnn_handle):
     if torch.cuda.get_device_capability()[0] < 9:
         pytest.skip("SDPA FP8 requires Hopper or higher")
 
+    is_cu_seq_len = bool(getattr(cfg, 'is_cu_seq_len', False))
+    with_ragged_offset_multiplier = bool(getattr(cfg, 'with_ragged_offset_multiplier', False))
+    if (is_cu_seq_len or with_ragged_offset_multiplier) and cudnn_version < "9.25.0":
+        pytest.skip("cu_seq_len / ragged offset multiplier for FP8 requires cuDNN 9.25.0 or higher (unified engine)")
+    if is_cu_seq_len:
+        assert cfg.is_infer, "is_cu_seq_len=True is forward-only (cu_seq_len is not plumbed for backward)"
+
     torch_itype = cfg.data_type
     torch_otype = cfg.output_type if hasattr(cfg, 'output_type') and cfg.output_type else cfg.data_type
     cudnn_itype = convert_to_cudnn_type(torch_itype)
@@ -332,15 +356,25 @@ def exec_sdpa_fp8(cfg, request, cudnn_handle):
         max_t_q = max(64, ((seq_len_q_gpu.sum().item() + 63) // 64) * 64)
         max_t_kv = max(64, ((seq_len_kv_gpu.sum().item() + 63) // 64) * 64)
 
-        q_ragged_offset_gpu = (prefix_sum(seq_len_q_gpu) * h_q * d_qk).to(torch.int64)
-        k_ragged_offset_gpu = (prefix_sum(seq_len_kv_gpu) * h_k * d_qk).to(torch.int64)
-        v_ragged_offset_gpu = (prefix_sum(seq_len_kv_gpu) * h_v * d_vo).to(torch.int64)
-        o_ragged_offset_gpu = (prefix_sum(seq_len_q_gpu) * h_q * d_vo).to(torch.int64)
+        # With the ragged offset multiplier, offsets are stored in coarser units
+        # (divided by the per-tensor multiplier; always divides evenly) and the
+        # engine scales them back to element offsets.
+        q_off_mult = d_qk if with_ragged_offset_multiplier else 1
+        k_off_mult = d_qk if with_ragged_offset_multiplier else 1
+        v_off_mult = d_vo if with_ragged_offset_multiplier else 1
+        o_off_mult = d_vo if with_ragged_offset_multiplier else 1
+        q_ragged_offset_gpu = (prefix_sum(seq_len_q_gpu) * h_q * d_qk // q_off_mult).to(torch.int64)
+        k_ragged_offset_gpu = (prefix_sum(seq_len_kv_gpu) * h_k * d_qk // k_off_mult).to(torch.int64)
+        v_ragged_offset_gpu = (prefix_sum(seq_len_kv_gpu) * h_v * d_vo // v_off_mult).to(torch.int64)
+        o_ragged_offset_gpu = (prefix_sum(seq_len_q_gpu) * h_q * d_vo // o_off_mult).to(torch.int64)
         stats_ragged_offset_gpu = (prefix_sum(seq_len_q_gpu) * h_q * 1).to(torch.int64)
+        if is_cu_seq_len:
+            cu_seq_len_q_gpu = prefix_sum(seq_len_q_gpu).to(torch.int32).view(-1)
+            cu_seq_len_kv_gpu = prefix_sum(seq_len_kv_gpu).to(torch.int32).view(-1)
 
     # Build forward graph (always needed)
     try:
-        graph_fwd = generate_graph_fwd(cudnn_itype, cudnn_otype, b, h_q, h_k, h_v, s_qo, s_kv, d_qk, d_vo, attn_scale, block_size, is_ragged=is_ragged, left_bound=left_bound, right_bound=right_bound, diag_align=diag_align, with_sink_token=with_sink_token, implementation=cfg.implementation)
+        graph_fwd = generate_graph_fwd(cudnn_itype, cudnn_otype, b, h_q, h_k, h_v, s_qo, s_kv, d_qk, d_vo, attn_scale, block_size, is_ragged=is_ragged, left_bound=left_bound, right_bound=right_bound, diag_align=diag_align, with_sink_token=with_sink_token, is_cu_seq_len=is_cu_seq_len, with_ragged_offset_multiplier=with_ragged_offset_multiplier, implementation=cfg.implementation)
         graph_fwd.validate()
         graph_fwd.build_operation_graph()
         graph_fwd.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK])
@@ -447,8 +481,12 @@ def exec_sdpa_fp8(cfg, request, cudnn_handle):
         variant_pack[int(GraphFwdUid.v_block_table)] = v_block_table_gpu
 
     if is_ragged:
-        variant_pack[int(GraphFwdUid.q_seq_len)] = torch.tensor(seq_len_q_list, dtype=torch.int32, device="cuda").view(-1)
-        variant_pack[int(GraphFwdUid.kv_seq_len)] = torch.tensor(seq_len_kv_list, dtype=torch.int32, device="cuda").view(-1)
+        if is_cu_seq_len:
+            variant_pack[int(GraphFwdUid.cu_seq_len_q)] = cu_seq_len_q_gpu
+            variant_pack[int(GraphFwdUid.cu_seq_len_kv)] = cu_seq_len_kv_gpu
+        else:
+            variant_pack[int(GraphFwdUid.q_seq_len)] = torch.tensor(seq_len_q_list, dtype=torch.int32, device="cuda").view(-1)
+            variant_pack[int(GraphFwdUid.kv_seq_len)] = torch.tensor(seq_len_kv_list, dtype=torch.int32, device="cuda").view(-1)
         variant_pack[int(GraphFwdUid.q_ragged_offset)] = q_ragged_offset_gpu
         variant_pack[int(GraphFwdUid.k_ragged_offset)] = k_ragged_offset_gpu
         variant_pack[int(GraphFwdUid.v_ragged_offset)] = v_ragged_offset_gpu

test/python/test_mhas_v2.py

@@ -843,7 +843,7 @@ def test_sdpa_fp8_fwd_ragged_L0(env_info, test_no, request, cudnn_handle):
         output_type=RandomChoice({torch.float8_e4m3fn: 1, torch.float8_e5m2: 1, torch.float16: 2}),
         with_sliding_mask=SlidingWindowMaskGenerator(no_mask=10),
         diag_align=RandomChoice({cudnn.diagonal_alignment.TOP_LEFT: 1}),
-        is_ragged_or_padded_or_full=RandomChoice({"ragged": 1, "padded": 0, "full": 0}),
+        is_ragged_or_padded_or_full=RandomChoice({"ragged": 1, "cu_ragged": 1, "cu_ragged_mult": 1, "padded": 0, "full": 0}),
     ) as randomization_ctx:
         test.cfg = randomization_ctx(rng, data_seed, geom_seed)
     test.showConfig(test_no, request)