From 127822e2dda71de4ba5abbbc7d5686b5247d189b Mon Sep 17 00:00:00 2001 From: Emil Gilliam Date: Wed, 8 Jul 2026 15:08:26 -0700 Subject: [PATCH 1/2] Expose cu_seq_len_q/kv on the sdpa_fp8 python binding The unified-engine FP8/MXFP8 forward (cuDNN 9.25+) accepts cumulative sequence lengths, and the C++ API has supported them on the fp8 node since 1.25 (SDPA_fp8_attributes aliases SDPA_attributes), but the python sdpa_fp8 binding hardcoded cu_seq_len_q/kv to nullptr. Expose them as kwargs (appended last to preserve positional backward compatibility) so python callers can use fp8 + cu_seq_len; the python-native pygraph capture/replay layer forwards them without changes. - python/pygraph/{pygraph.h,sdpa.cpp}: add cu_seq_len_q/kv parameters to PyGraph::sdpa_fp8 and its m.def, with docstring entries (requires cuDNN 9.25+ and the UNIFIED implementation). Remove a stale "Deprecated, use sdpa_unified instead" comment: implementation selection is automatic (or explicit via the implementation attribute), and python fp8 users are expected to call sdpa_fp8. - docs/operations/Attention.md: document cu_seq_len_q/kv on the fp16/bf16 C++ and python APIs (missed in #266), the ragged offset multiplier (missed in #290), and the fp8 varlen surface incl. the new kwargs. - test/python/sdpa/fp8.py: support is_cu_seq_len and with_ragged_offset_multiplier configs (mirroring fp16.py): cu_seq_len graph tensors, token-coarse offsets with per-tensor multipliers on Q/K/V/O, version gating at 9.25. - test/python/test_mhas_v2.py: test_sdpa_fp8_fwd_ragged_L0 now draws ragged / cu_ragged / cu_ragged_mult with equal weight. sdpa_mxfp8 is intentionally untouched: it has no varlen surface at all (no padding mask or seq_len kwargs), so cu_seq_len support there is a separate feature. Validated against cuDNN 9.25 (test_sdpa_fp8_fwd_ragged_L0): H100 10 passed / 22 skipped (pre-existing Hopper config limits), Blackwell 24 passed / 8 skipped (head-dim limits); the passing draws include 23 is_cu_seq_len=True and 9 multiplier configs, zero failures. Co-Authored-By: Claude Fable 5 --- docs/operations/Attention.md | 23 +++++++++++++- python/pygraph/pygraph.h | 4 ++- python/pygraph/sdpa.cpp | 12 ++++---- test/python/sdpa/fp8.py | 58 +++++++++++++++++++++++++++++------- test/python/test_mhas_v2.py | 2 +- 5 files changed, 81 insertions(+), 18 deletions(-) diff --git a/docs/operations/Attention.md b/docs/operations/Attention.md index bceaba78b..a2c7b15a4 100644 --- a/docs/operations/Attention.md +++ b/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 value); SDPA_attributes& set_seq_len_kv(std::shared_ptr 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 value); +SDPA_attributes& set_cu_seq_len_kv(std::shared_ptr 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 diff --git a/python/pygraph/pygraph.h b/python/pygraph/pygraph.h index 0eee32987..4019f3560 100644 --- a/python/pygraph/pygraph.h +++ b/python/pygraph/pygraph.h @@ -524,7 +524,9 @@ class PyGraph { std::shared_ptr score_sum_exp, std::shared_ptr sink_token, bool const unfuse_fma, - cudnn_frontend::AttentionImplementation_t const& implementation); + cudnn_frontend::AttentionImplementation_t const& implementation, + std::shared_ptr& cu_seq_len_q, + std::shared_ptr& cu_seq_len_kv); // MXFP8 SDPA forward - uses block-wise scale factors (E8M0 with F8_128x4 reordering) // return [o, stats, amax_o] diff --git a/python/pygraph/sdpa.cpp b/python/pygraph/sdpa.cpp index dabad1782..738c989cb 100644 --- a/python/pygraph/sdpa.cpp +++ b/python/pygraph/sdpa.cpp @@ -505,7 +505,6 @@ PyGraph::sdpa_backward(std::shared_ptr return {dQ, dK, dV}; } -// Deprecated, use sdpa_unified instead std::array, 4> PyGraph::sdpa_fp8(std::shared_ptr& q, std::shared_ptr& k, @@ -542,12 +541,11 @@ PyGraph::sdpa_fp8(std::shared_ptr& q, std::shared_ptr score_sum_exp, std::shared_ptr sink_token, bool const unfuse_fma, - cudnn_frontend::AttentionImplementation_t const& implementation) { + cudnn_frontend::AttentionImplementation_t const& implementation, + std::shared_ptr& cu_seq_len_q, + std::shared_ptr& cu_seq_len_kv) { cudnn_frontend::DataType_t mma_core_mode = cudnn_frontend::DataType_t::FP8_E4M3; std::shared_ptr block_mask = nullptr; - // cu_seq_len_q/cu_seq_len_kv are not exposed via the fp8 path. - std::shared_ptr cu_seq_len_q = nullptr; - std::shared_ptr 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_& 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_& 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. + 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. 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. diff --git a/test/python/sdpa/fp8.py b/test/python/sdpa/fp8.py index 0da4cc596..c9802845d 100644 --- a/test/python/sdpa/fp8.py +++ b/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 diff --git a/test/python/test_mhas_v2.py b/test/python/test_mhas_v2.py index 84b6dd1f8..746cf2d1d 100644 --- a/test/python/test_mhas_v2.py +++ b/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) From 148e08e2a8ec3efa4d81eacf83f67de5797ba885 Mon Sep 17 00:00:00 2001 From: Emil Gilliam Date: Wed, 8 Jul 2026 17:34:08 -0700 Subject: [PATCH 2/2] Complete cu_seq_len docstring constraints on sdpa_fp8 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address review: the runtime-visible docstring now carries the same constraints as the sdpa() docstring and Attention.md — set together, use_padding_mask=True, cuDNN 9.25+ and the UNIFIED implementation. Co-Authored-By: Claude Fable 5 --- python/pygraph/sdpa.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/pygraph/sdpa.cpp b/python/pygraph/sdpa.cpp index 738c989cb..57cc22c03 100644 --- a/python/pygraph/sdpa.cpp +++ b/python/pygraph/sdpa.cpp @@ -1304,8 +1304,8 @@ init_pygraph_sdpa_submodule(py::class_& 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. - 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. + 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.