From da67de4fd3343147bdc7ee4261fdbabe11e723b6 Mon Sep 17 00:00:00 2001 From: Emil Gilliam Date: Wed, 22 Jul 2026 14:25:05 -0700 Subject: [PATCH] Support mixed-form sequence lengths in SDPA forward (cuDNN 9.26+) Allow the Q and KV sides of SDPA forward to independently choose their sequence length representation: per-batch (seq_len_*) or cumulative (cu_seq_len_*). Previously the two forms were mutually exclusive across the whole operation, forcing paged-attention integrations -- whose natural state is a cumulative Q prefix sum alongside per-batch KV lengths -- to materialize a KV-side prefix sum purely to satisfy the pairing rule. - sdpa_support_surface: validation is now per-side (exactly one form per side, sides provided together); the paged requirement accepts either form per side; mixed forms are gated on cuDNN 9.26+ in the UNIFIED surface (the initial public 9.25 release predates the backend-side support, and a 9.25 build that has it cannot be told apart from one that does not). Hoist a has_input helper for the presence checks. - diagonal_band_mask: relax the node's metadata seqlen check to per-side exclusivity, matching the SDPA rules. - tests: per-side form control in the SDPA test harness (cu_seq_len_sides) and a deterministic mixed-form test with non-uniform lengths, including bottom-right causal cases that guard the backend's band-mask alignment derivation. Validated on H100 against a cuDNN 9.27 dev build: mixed forms in both directions match the both-per-batch reference for non-paged, ragged, and paged (ragged Q + page tables, cu_seq_len_q + seq_len_kv) graphs; the deterministic tests skip below the version gate. Co-Authored-By: Claude Fable 5 --- .../cudnn_frontend/node/diagonal_band_mask.h | 10 +-- .../node/sdpa_support_surface.h | 69 ++++++++++++------- test/python/sdpa/fp16.py | 26 ++++--- test/python/sdpa/random_config.py | 10 +++ test/python/test_mhas_v2.py | 59 ++++++++++++++++ 5 files changed, 135 insertions(+), 39 deletions(-) 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..0a43d2ec7 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 < 92600) { + 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.26.0 or above"); + } + if (effective_cudnn_ver >= 92500) { allowed_input_names.insert({input_names::Descale_Q, input_names::Descale_K, diff --git a/test/python/sdpa/fp16.py b/test/python/sdpa/fp16.py index 2ef97058f..e9d0b865a 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.26.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.26.0 or higher.") + pytest.skip("mixed-form sequence lengths (cumulative on one side only) require cuDNN 9.26.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..6944857d3 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.26+. + 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 746cf2d1d..2cea541b9 100644 --- a/test/python/test_mhas_v2.py +++ b/test/python/test_mhas_v2.py @@ -1024,6 +1024,65 @@ 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.26+ (skips below via exec_sdpa). + """ + test = SDPATestConfig(**env_info, implementation=cudnn.attention_implementation.AUTO) + 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.AUTO, + ) + 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