diff --git a/docs/operations/Attention.md b/docs/operations/Attention.md index a2c7b15a4..e7430fa74 100644 --- a/docs/operations/Attention.md +++ b/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. diff --git a/include/cudnn_frontend/node/diagonal_band_mask.h b/include/cudnn_frontend/node/diagonal_band_mask.h index 243d1e349..62dba8c56 100644 --- a/include/cudnn_frontend/node/diagonal_band_mask.h +++ b/include/cudnn_frontend/node/diagonal_band_mask.h @@ -29,10 +29,12 @@ class DiagonalBandMaskNodeBase : public NodeCRTP { 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(); } diff --git a/include/cudnn_frontend/node/sdpa_support_surface.h b/include/cudnn_frontend/node/sdpa_support_surface.h index 8c59a5937..9d7cddd5d 100644 --- a/include/cudnn_frontend/node/sdpa_support_surface.h +++ b/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, diff --git a/python/pygraph/sdpa.cpp b/python/pygraph/sdpa.cpp index 57cc22c03..5b4e5191d 100644 --- a/python/pygraph/sdpa.cpp +++ b/python/pygraph/sdpa.cpp @@ -1138,8 +1138,8 @@ init_pygraph_sdpa_submodule(py::class_& 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_& 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. diff --git a/samples/cpp/sdpa/fp16_fwd_with_cu_seq_len.cpp b/samples/cpp/sdpa/fp16_fwd_with_cu_seq_len.cpp index 96e2c713c..1674e32f3 100644 --- a/samples/cpp/sdpa/fp16_fwd_with_cu_seq_len.cpp +++ b/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 diff --git a/test/python/sdpa/fp16.py b/test/python/sdpa/fp16.py index 2ef97058f..34f7e4918 100644 --- a/test/python/sdpa/fp16.py +++ b/test/python/sdpa/fp16.py @@ -81,6 +81,7 @@ def validate_config(cfg): if cfg.is_cu_seq_len: assert cfg.is_padding == True, "is_cu_seq_len=True requires is_padding=True" assert cfg.is_train == False, "is_cu_seq_len=True is forward-only (cu_seq_len is not plumbed for backward)" + assert cfg.cu_seq_len_sides in ("both", "q", "kv"), f"invalid cu_seq_len_sides={cfg.cu_seq_len_sides}" assert isinstance(cfg.seq_len_q, (list, tuple)), "input 'seq_len_q' must be list or tuple" if cfg.is_padding: @@ -119,6 +120,10 @@ def validate_config(cfg): print("@@@@ Overall result: WAIVED, cu_seq_len_q/cu_seq_len_kv require cuDNN 9.24.0 or higher.") pytest.skip("cu_seq_len_q/cu_seq_len_kv require cuDNN 9.24.0 or higher") + if cudnn_version < "9.25.0" and cfg.is_cu_seq_len and cfg.cu_seq_len_sides != "both": + print("@@@@ Overall result: WAIVED, mixed-form sequence lengths require cuDNN 9.25.0 or higher.") + pytest.skip("mixed-form sequence lengths (cumulative on one side only) require cuDNN 9.25.0 or higher") + def allocate_tensors(cfg, rng_data_gen, perf=False): allocs = {} @@ -159,14 +164,17 @@ def allocate_tensors(cfg, rng_data_gen, perf=False): seq_len_q_gpu = torch.tensor(cfg.seq_len_q, dtype=torch.int32, device="cuda").view(-1) if len(cfg.seq_len_q) > 0 else None seq_len_kv_gpu = torch.tensor(cfg.seq_len_kv, dtype=torch.int32, device="cuda").view(-1) if len(cfg.seq_len_kv) > 0 else None - if cfg.is_cu_seq_len: - # When using cu_seq_len, the seq_len_q/seq_len_kv tensors are not part of the - # graph; instead, supply 1-D (b+1,) int32 prefix-sums of the per-batch seq_lens - # (the frontend promotes these to the 4-D form the backend requires). + # A side using the cumulative form gets a 1-D (b+1,) int32 prefix-sum of its + # per-batch seq_lens (the frontend promotes these to the 4-D form the backend + # requires) instead of the regular per-batch tensor. The two sides choose + # independently (cu_seq_len_sides), so mixed forms are exercised too. + if cfg.is_cu_seq_len_q(): allocs[TensorUid.cu_seq_len_q] = (prefix_sum(seq_len_q_gpu).to(torch.int32).view(-1), None, None) - allocs[TensorUid.cu_seq_len_kv] = (prefix_sum(seq_len_kv_gpu).to(torch.int32).view(-1), None, None) else: allocs[TensorUid.seq_len_q] = (seq_len_q_gpu, None, None) + if cfg.is_cu_seq_len_kv(): + allocs[TensorUid.cu_seq_len_kv] = (prefix_sum(seq_len_kv_gpu).to(torch.int32).view(-1), None, None) + else: allocs[TensorUid.seq_len_kv] = (seq_len_kv_gpu, None, None) if cfg.is_ragged: @@ -267,11 +275,11 @@ def create_forward_graph(cfg, tensors, cudnn_handle): block_mask_dim = (cfg.batches, cfg.h_q, (cfg.s_q + TILE_M - 1) // TILE_M, ((cfg.s_kv + TILE_N - 1) // TILE_N + 7) // 8) block_mask = graph.tensor(uid=int(TensorUid.block_mask), dim=block_mask_dim, stride=(block_mask_dim[1]*block_mask_dim[2]*block_mask_dim[3], block_mask_dim[2]*block_mask_dim[3], block_mask_dim[3], 1), data_type=cudnn.data_type.UINT8) if cfg.is_block_mask else None - seq_len_q = graph.tensor(uid=int(TensorUid.seq_len_q), dim=(cfg.batches,), stride=(1,), data_type=cudnn.data_type.INT32) if (cfg.is_padding and not cfg.is_cu_seq_len) else None - seq_len_kv = graph.tensor(uid=int(TensorUid.seq_len_kv), dim=(cfg.batches,), stride=(1,), data_type=cudnn.data_type.INT32) if (cfg.is_padding and not cfg.is_cu_seq_len) else None + seq_len_q = graph.tensor(uid=int(TensorUid.seq_len_q), dim=(cfg.batches,), stride=(1,), data_type=cudnn.data_type.INT32) if (cfg.is_padding and not cfg.is_cu_seq_len_q()) else None + seq_len_kv = graph.tensor(uid=int(TensorUid.seq_len_kv), dim=(cfg.batches,), stride=(1,), data_type=cudnn.data_type.INT32) if (cfg.is_padding and not cfg.is_cu_seq_len_kv()) else None - cu_seq_len_q = graph.tensor(uid=int(TensorUid.cu_seq_len_q), dim=(cfg.batches + 1,), stride=(1,), data_type=cudnn.data_type.INT32) if cfg.is_cu_seq_len else None - cu_seq_len_kv = graph.tensor(uid=int(TensorUid.cu_seq_len_kv), dim=(cfg.batches + 1,), stride=(1,), data_type=cudnn.data_type.INT32) if cfg.is_cu_seq_len else None + cu_seq_len_q = graph.tensor(uid=int(TensorUid.cu_seq_len_q), dim=(cfg.batches + 1,), stride=(1,), data_type=cudnn.data_type.INT32) if cfg.is_cu_seq_len_q() else None + cu_seq_len_kv = graph.tensor(uid=int(TensorUid.cu_seq_len_kv), dim=(cfg.batches + 1,), stride=(1,), data_type=cudnn.data_type.INT32) if cfg.is_cu_seq_len_kv() else None seed = offset = dropout_tuple = rng_dump = None if cfg.is_dropout: diff --git a/test/python/sdpa/random_config.py b/test/python/sdpa/random_config.py index adc661518..a8758b08d 100644 --- a/test/python/sdpa/random_config.py +++ b/test/python/sdpa/random_config.py @@ -88,6 +88,10 @@ class ExecConfig: # (cumulative sequence-length tensors of shape (b+1, 1, 1, 1)) instead of # the regular per-batch seq_len_q/seq_len_kv tensors. Implies is_padding=True. is_cu_seq_len: bool = None + # Which sides use the cumulative form when is_cu_seq_len is True: "both" + # (default), "q" (cu_seq_len_q with seq_len_kv), or "kv" (seq_len_q with + # cu_seq_len_kv). The mixed forms require cuDNN 9.25+. + cu_seq_len_sides: str = "both" is_ragged: bool = None is_dropout: bool = None is_determin: bool = None @@ -144,6 +148,12 @@ class ExecConfig: def is_train(self): return not self.is_infer + def is_cu_seq_len_q(self): + return bool(self.is_cu_seq_len) and self.cu_seq_len_sides in ("both", "q") + + def is_cu_seq_len_kv(self): + return bool(self.is_cu_seq_len) and self.cu_seq_len_sides in ("both", "kv") + def fill_derived_fields(self): """ Fill in derived fields (shapes, strides) from basic dims. diff --git a/test/python/test_mhas_v2.py b/test/python/test_mhas_v2.py index 3cb932b89..7fc355d3d 100644 --- a/test/python/test_mhas_v2.py +++ b/test/python/test_mhas_v2.py @@ -1024,6 +1024,67 @@ def test_sdpa_mxfp8_bwd_L0(env_info, test_no, request, cudnn_handle): # # Single repro test # # =================== +MIXED_SEQ_LEN_FORM_CASES = [ + ("q", cudnn.diagonal_alignment.TOP_LEFT, None), + ("kv", cudnn.diagonal_alignment.TOP_LEFT, None), + ("q", cudnn.diagonal_alignment.BOTTOM_RIGHT, 0), + ("kv", cudnn.diagonal_alignment.BOTTOM_RIGHT, 0), +] + + +@pytest.mark.parametrize( + "cu_sides,diag_align,right_bound", + MIXED_SEQ_LEN_FORM_CASES, + ids=["cu_q", "cu_kv", "cu_q_brcm", "cu_kv_brcm"], +) +@pytest.mark.L0 +def test_sdpa_mixed_seq_len_forms_L0(env_info, cu_sides, diag_align, right_bound, request, cudnn_handle): + """Mixed-form sequence lengths: cumulative on one side, per-batch on the other. + + Deterministic configs with non-uniform per-batch lengths, so misreading one + side's form cannot produce a passing result. The bottom-right causal cases + guard the DiagonalBandMask alignment derivation, which must treat the two + sides' forms independently. Requires cuDNN 9.25+ (skips below via exec_sdpa). + """ + # Mixed forms are gated in the UNIFIED surface; request it explicitly rather + # than relying on AUTO resolving to it. + test = SDPATestConfig(**env_info, implementation=cudnn.attention_implementation.UNIFIED) + test.cfg = ExecConfig( + data_type=torch.bfloat16, + rng_data_seed=1234, + rng_geom_seed=5678, + is_alibi=False, + is_infer=True, + is_paged=False, + is_bias=False, + is_block_mask=False, + is_padding=True, + is_cu_seq_len=True, + cu_seq_len_sides=cu_sides, + is_ragged=False, + is_dropout=False, + is_determin=False, + batches=4, + d_qk=64, + d_v=64, + s_q=256, + s_kv=512, + h_q=3, + h_k=3, + h_v=3, + diag_align=diag_align, + left_bound=None, + right_bound=right_bound, + seq_len_q=[128, 100, 256, 37], + seq_len_kv=[96, 64, 512, 200], + implementation=cudnn.attention_implementation.UNIFIED, + ) + test.cfg.fill_derived_fields() + test.showConfig((request.node.name, len(MIXED_SEQ_LEN_FORM_CASES)), request) + + exec_sdpa(test.cfg, request, cudnn_handle) + + @pytest.mark.skipif("not config.getoption('--repro')", reason="used with '--repro' only") @pytest.mark.L0 @pytest.mark.L1