diff --git a/include/cudnn_frontend/graph_properties.h b/include/cudnn_frontend/graph_properties.h index aa0091e45..bbb61afad 100644 --- a/include/cudnn_frontend/graph_properties.h +++ b/include/cudnn_frontend/graph_properties.h @@ -1984,6 +1984,8 @@ class SDPA_attributes : public Attributes { Bias, SEQ_LEN_Q, SEQ_LEN_KV, + CU_SEQ_LEN_Q, + CU_SEQ_LEN_KV, Seed, Offset, Dropout_mask, @@ -2109,6 +2111,18 @@ class SDPA_attributes : public Attributes { return *this; } + SDPA_attributes& + set_cu_seq_len_q(std::shared_ptr value) { + inputs[SDPA_attributes::input_names::CU_SEQ_LEN_Q] = std::move(value); + return *this; + } + + SDPA_attributes& + set_cu_seq_len_kv(std::shared_ptr value) { + inputs[SDPA_attributes::input_names::CU_SEQ_LEN_KV] = std::move(value); + return *this; + } + SDPA_attributes& set_diagonal_alignment(DiagonalAlignment_t const alignment) { diagonal_alignment = alignment; @@ -2741,7 +2755,7 @@ class DiagonalBandMask_attributes : public Attributes> inputs; enum class output_names { Y }; std::unordered_map> outputs; diff --git a/include/cudnn_frontend/node/diagonal_band_mask.h b/include/cudnn_frontend/node/diagonal_band_mask.h index aa301f975..243d1e349 100644 --- a/include/cudnn_frontend/node/diagonal_band_mask.h +++ b/include/cudnn_frontend/node/diagonal_band_mask.h @@ -29,7 +29,16 @@ 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 pre_validate_node_extra(); + } + + virtual error_t + pre_validate_node_extra() const { return {error_code_t::OK, ""}; } @@ -67,6 +76,18 @@ class DiagonalBandMaskNodeBase : public NodeCRTP { return ((seq_len_KV_it) != attributes.inputs.end() && seq_len_KV_it->second != nullptr); } + bool + has_cu_seq_len_q() const { + auto cu_seq_len_Q_it = attributes.inputs.find(DiagonalBandMask_attributes::input_names::CU_SEQ_LEN_Q); + return ((cu_seq_len_Q_it) != attributes.inputs.end() && cu_seq_len_Q_it->second != nullptr); + } + + bool + has_cu_seq_len_kv() const { + auto cu_seq_len_KV_it = attributes.inputs.find(DiagonalBandMask_attributes::input_names::CU_SEQ_LEN_KV); + return ((cu_seq_len_KV_it) != attributes.inputs.end() && cu_seq_len_KV_it->second != nullptr); + } + bool has_left_bound() const { auto left_bound_it = attributes.inputs.find(DiagonalBandMask_attributes::input_names::LeftBound); @@ -107,6 +128,17 @@ class CompositeDiagonalBandMaskNode : public DiagonalBandMaskNodeBase= 92400) + NV_CUDNN_FE_DYNAMIC_CHECK_CUDNN_BACKEND_VERSION(92400, cu_seq_len_cudnn_ver_error); + + if (has_cu_seq_len_q()) { + auto cu_seq_len_Q = + attributes.inputs.find(DiagonalBandMask_attributes::input_names::CU_SEQ_LEN_Q)->second; + auto backend_cu_seq_len_Q = tensors[cu_seq_len_Q->get_uid()]->get_desc()->get_backend_descriptor(); + _CUDNN_CHECK_CUDNN_ERROR(detail::set_attribute(diagonal_band_mask_operation->get_backend_descriptor(), + CUDNN_ATTR_OPERATION_DIAGONAL_BAND_MASK_CU_SEQ_LEN_QDESC, + CUDNN_TYPE_BACKEND_DESCRIPTOR, + 1, + &backend_cu_seq_len_Q)); + } + + if (has_cu_seq_len_kv()) { + auto cu_seq_len_KV = + attributes.inputs.find(DiagonalBandMask_attributes::input_names::CU_SEQ_LEN_KV)->second; + auto backend_cu_seq_len_KV = tensors[cu_seq_len_KV->get_uid()]->get_desc()->get_backend_descriptor(); + _CUDNN_CHECK_CUDNN_ERROR( + detail::set_attribute(diagonal_band_mask_operation->get_backend_descriptor(), + CUDNN_ATTR_OPERATION_DIAGONAL_BAND_MASK_CU_SEQ_LEN_KVDESC, + CUDNN_TYPE_BACKEND_DESCRIPTOR, + 1, + &backend_cu_seq_len_KV)); + } +#else + return cu_seq_len_cudnn_ver_error; +#endif + } + if (has_left_bound()) { auto left_bound = attributes.inputs.find(DiagonalBandMask_attributes::input_names::LeftBound)->second; auto backend_left_bound = tensors[left_bound->get_uid()]->get_desc()->get_backend_descriptor(); @@ -369,6 +435,8 @@ INode::diagonal_band_mask(std::shared_ptr x, std::shared_ptr b, std::shared_ptr seq_len_q, std::shared_ptr seq_len_kv, + std::shared_ptr cu_seq_len_q, + std::shared_ptr cu_seq_len_kv, std::shared_ptr left_bound, std::shared_ptr shift_right_bound, DiagonalBandMask_attributes attributes) { @@ -380,6 +448,12 @@ INode::diagonal_band_mask(std::shared_ptr x, if (seq_len_kv) { attributes.inputs[DiagonalBandMask_attributes::input_names::SEQ_LEN_KV] = seq_len_kv; } + if (cu_seq_len_q) { + attributes.inputs[DiagonalBandMask_attributes::input_names::CU_SEQ_LEN_Q] = cu_seq_len_q; + } + if (cu_seq_len_kv) { + attributes.inputs[DiagonalBandMask_attributes::input_names::CU_SEQ_LEN_KV] = cu_seq_len_kv; + } if (left_bound) { attributes.inputs[DiagonalBandMask_attributes::input_names::LeftBound] = left_bound; } diff --git a/include/cudnn_frontend/node/scaled_dot_product_flash_attention.h b/include/cudnn_frontend/node/scaled_dot_product_flash_attention.h index e773a75f7..6e14aab0d 100644 --- a/include/cudnn_frontend/node/scaled_dot_product_flash_attention.h +++ b/include/cudnn_frontend/node/scaled_dot_product_flash_attention.h @@ -22,6 +22,23 @@ namespace cudnn_frontend::graph { +// If a 1-D length/offset index tensor [n] is supplied, promote it to the 4-D +// [n, 1, 1, 1] form that the cuDNN backend requires for seq_len / cu_seq_len and +// ragged-offset tensors. No-op if the tensor is unset or is not 1-D. +inline void +promote_1d_index_tensor_to_4d(std::shared_ptr const& tensor) { + if (tensor == nullptr) { + return; + } + auto const& dim = tensor->get_dim(); + auto const& stride = tensor->get_stride(); + if (dim.size() == 1) { + int64_t const n = dim[0]; + int64_t const leading_stride = (stride.size() == 1) ? stride[0] : 1; + tensor->set_dim({n, 1, 1, 1}).set_stride({leading_stride, 1, 1, 1}); + } +} + namespace attn::score_modifiers { // clang-format off @@ -61,7 +78,9 @@ inline std::shared_ptr sliding_window_mask( int64_t s_q, int64_t s_kv, std::shared_ptr s_q_ptr, - std::shared_ptr s_kv_ptr + std::shared_ptr s_kv_ptr, + std::shared_ptr cu_s_q_ptr = nullptr, + std::shared_ptr cu_s_kv_ptr = nullptr ); inline std::shared_ptr alibi_mask( @@ -87,6 +106,31 @@ class SDPANodeBase : public NodeCRTP { std::shared_ptr alibi_slopes; int64_t alibi_slopes_size = 0; + // Promote any 1-D seq_len / cu_seq_len / ragged-offset index tensors to the 4-D + // [n, 1, 1, 1] form the cuDNN backend requires (see promote_1d_index_tensor_to_4d). + void + promote_index_tensors_to_4d() { + for (auto& key : + {input_names::SEQ_LEN_Q, input_names::SEQ_LEN_KV, input_names::CU_SEQ_LEN_Q, input_names::CU_SEQ_LEN_KV}) { + auto const it = attributes.inputs.find(key); + if (it != attributes.inputs.end()) { + promote_1d_index_tensor_to_4d(it->second); + } + } + for (auto& [key, value] : attributes.inputs) { + CUDNN_FRONTEND_UNUSED(key); + if (value != nullptr) { + promote_1d_index_tensor_to_4d(value->get_ragged_offset()); + } + } + for (auto& [key, value] : attributes.outputs) { + CUDNN_FRONTEND_UNUSED(key); + if (value != nullptr) { + promote_1d_index_tensor_to_4d(value->get_ragged_offset()); + } + } + } + public: SDPA_attributes attributes; @@ -117,6 +161,18 @@ class SDPANodeBase : public NodeCRTP { return ((seq_len_KV_it) != attributes.inputs.end() && seq_len_KV_it->second != nullptr); } + bool + has_cu_seq_len_q() const { + auto cu_seq_len_Q_it = attributes.inputs.find(SDPA_attributes::input_names::CU_SEQ_LEN_Q); + return ((cu_seq_len_Q_it) != attributes.inputs.end() && cu_seq_len_Q_it->second != nullptr); + } + + bool + has_cu_seq_len_kv() const { + auto cu_seq_len_KV_it = attributes.inputs.find(SDPA_attributes::input_names::CU_SEQ_LEN_KV); + return ((cu_seq_len_KV_it) != attributes.inputs.end() && cu_seq_len_KV_it->second != nullptr); + } + // Helper function to detect MXFP8 (microscaling FP8) mode // MXFP8 uses block-wise scale factors with E8M0 data type and F8_128x4 reordering // When detected, we use block_scale_dequantize before matmuls instead of pointwise descale after @@ -373,6 +429,10 @@ class SDPANodeBase : public NodeCRTP { error_t infer_properties_node() override final { + // If a 1-D length/offset index tensor [n] is supplied, promote it to the 4-D + // [n, 1, 1, 1] (contiguous) form the cuDNN backend requires. + promote_index_tensors_to_4d(); + if (attributes.generate_stats.value_or(false)) { auto stats = attributes.outputs.at(output_names::Stats); auto stats_dim = stats->get_dim(); @@ -483,6 +543,16 @@ class SDPANodeBase : public NodeCRTP { add_tensor(seq_len_kv_it->second); } + auto const cu_seq_len_q_it = attributes.inputs.find(input_names::CU_SEQ_LEN_Q); + if (cu_seq_len_q_it != attributes.inputs.end()) { + add_tensor(cu_seq_len_q_it->second); + } + + auto const cu_seq_len_kv_it = attributes.inputs.find(input_names::CU_SEQ_LEN_KV); + if (cu_seq_len_kv_it != attributes.inputs.end()) { + add_tensor(cu_seq_len_kv_it->second); + } + for (auto const& tensor : {attributes.inputs.at(input_names::Q), attributes.inputs.at(input_names::K), attributes.inputs.at(input_names::V), @@ -808,7 +878,9 @@ class CompositeSDPANode : public SDPANodeBase { s_q, s_kv, s_q_ptr, - s_kv_ptr); + s_kv_ptr, + /*cu_s_q_ptr=*/nullptr, + /*cu_s_kv_ptr=*/nullptr); sub_nodes.emplace_back(node_); } @@ -1028,6 +1100,31 @@ class CompositeSDPABackwardNode : public NodeCRTP { mutable bool is_deterministic_algorithm_supported_on_blackwell = false; // Will be edited in pre_validate_node() mutable bool is_d256_on_blackwell = false; // Will be edited in pre_validate_node() + // Promote any 1-D seq_len / ragged-offset index tensors to the 4-D + // [n, 1, 1, 1] form the cuDNN backend requires (see promote_1d_index_tensor_to_4d). + void + promote_index_tensors_to_4d() { + // TODO: Handle CU_SEQ_LEN_Q and CU_SEQ_LEN_KV once bprop supports these. + for (auto& key : {input_names::SEQ_LEN_Q, input_names::SEQ_LEN_KV}) { + auto const it = attributes.inputs.find(key); + if (it != attributes.inputs.end()) { + promote_1d_index_tensor_to_4d(it->second); + } + } + for (auto& [key, value] : attributes.inputs) { + CUDNN_FRONTEND_UNUSED(key); + if (value != nullptr) { + promote_1d_index_tensor_to_4d(value->get_ragged_offset()); + } + } + for (auto& [key, value] : attributes.outputs) { + CUDNN_FRONTEND_UNUSED(key); + if (value != nullptr) { + promote_1d_index_tensor_to_4d(value->get_ragged_offset()); + } + } + } + public: mutable SDPA_backward_attributes attributes; // Will be edited in pre_validate_node() for workaround padding mask @@ -1329,6 +1426,10 @@ class CompositeSDPABackwardNode : public NodeCRTP { error_t infer_properties_node() override final { + // If a 1-D length/offset index tensor [n] is supplied, promote it to the 4-D + // [n, 1, 1, 1] (contiguous) form the cuDNN backend requires. + promote_index_tensors_to_4d(); + // clang-format off if (detail::get_backend_version() < 90600 && (attributes.max_total_seq_len_q.has_value() || attributes.max_total_seq_len_kv.has_value())) { CUDNN_FE_LOG_LABEL_ENDL("WARNING: sdpa_backward.attributes.max_total_seq_len has been set, but cuDNN version is below 9.6.0 does not support max_total_seq_len_q. The workspace memory size required to execute this graph may be unexpectedly large"); @@ -1711,7 +1812,9 @@ class CompositeSDPABackwardNode : public NodeCRTP { s_q, s_kv, s_q_ptr, - s_kv_ptr); + s_kv_ptr, + /*cu_s_q_ptr=*/nullptr, + /*cu_s_kv_ptr=*/nullptr); sub_nodes.emplace_back(std::move(node_)); } @@ -2141,14 +2244,20 @@ class UnifiedSDPANode : public SDPANodeBase { if (attributes.left_bound.has_value() || attributes.right_bound.has_value()) { if (!subgraph) init_subgraph(); - auto s_q = attributes.inputs[input_names::Q]->get_dim()[2]; - auto s_kv = infer_s_kv(); - auto s_kv_ptr = attributes.inputs.find(input_names::SEQ_LEN_KV) != attributes.inputs.end() - ? attributes.inputs[input_names::SEQ_LEN_KV] - : nullptr; - auto s_q_ptr = attributes.inputs.find(input_names::SEQ_LEN_Q) != attributes.inputs.end() - ? attributes.inputs[input_names::SEQ_LEN_Q] - : nullptr; + auto s_q = attributes.inputs[input_names::Q]->get_dim()[2]; + auto s_kv = infer_s_kv(); + auto s_kv_ptr = attributes.inputs.find(input_names::SEQ_LEN_KV) != attributes.inputs.end() + ? attributes.inputs[input_names::SEQ_LEN_KV] + : nullptr; + auto s_q_ptr = attributes.inputs.find(input_names::SEQ_LEN_Q) != attributes.inputs.end() + ? attributes.inputs[input_names::SEQ_LEN_Q] + : nullptr; + auto cu_s_kv_ptr = attributes.inputs.find(input_names::CU_SEQ_LEN_KV) != attributes.inputs.end() + ? attributes.inputs[input_names::CU_SEQ_LEN_KV] + : nullptr; + auto cu_s_q_ptr = attributes.inputs.find(input_names::CU_SEQ_LEN_Q) != attributes.inputs.end() + ? attributes.inputs[input_names::CU_SEQ_LEN_Q] + : nullptr; subgraph_output = attn::score_modifiers::sliding_window_mask(subgraph, subgraph_output, @@ -2158,7 +2267,9 @@ class UnifiedSDPANode : public SDPANodeBase { s_q, s_kv, s_q_ptr, - s_kv_ptr); + s_kv_ptr, + cu_s_q_ptr, + cu_s_kv_ptr); } if (subgraph) { @@ -2230,7 +2341,7 @@ class UnifiedSDPANode : public SDPANodeBase { managed_backend_descriptor_t& raw_operations, std::unordered_map>& tensors) const override final { CUDNN_FRONTEND_UNUSED(operations); - CUDNN_FE_LOG_LABEL("INFO: " << "Building UnifiedSDPANode operations " << attributes.name << " "); + CUDNN_FE_LOG_LABEL("INFO: Building UnifiedSDPANode operations " << attributes.name << " "); auto cudnn_ver_error = error_t{error_code_t::GRAPH_NOT_SUPPORTED, "Unified SDPA node requires cuDNN 9.13.1"}; #if (CUDNN_VERSION >= 91301) @@ -2332,10 +2443,15 @@ class UnifiedSDPANode : public SDPANodeBase { #endif } - // Paged attention attributes - if (is_paged_k() || is_paged_v() || has_seq_len_q() || has_seq_len_kv()) { - auto paged_cudnn_ver_error = error_t{error_code_t::GRAPH_NOT_SUPPORTED, - "Paged attention in unified SDPA node requires cuDNN 9.15.0"}; + // Variable-length and paged-attention attributes: + // - paged attention K/V page tables (needs cuDNN >= 9.15.0) + // - per-batch sequence-length tensors (SEQ_LEN_Q / SEQ_LEN_KV; needs cuDNN >= 9.15.0) + // - cumulative sequence-length tensors (CU_SEQ_LEN_Q / CU_SEQ_LEN_KV; needs cuDNN >= 9.24.0). + if (is_paged_k() || is_paged_v() || has_seq_len_q() || has_seq_len_kv() || has_cu_seq_len_q() || + has_cu_seq_len_kv()) { + auto paged_cudnn_ver_error = + error_t{error_code_t::GRAPH_NOT_SUPPORTED, + "Paged attention or per-batch sequence lengths in unified SDPA node require cuDNN 9.15.0"}; #if (CUDNN_VERSION >= 91500) NV_CUDNN_FE_DYNAMIC_CHECK_CUDNN_BACKEND_VERSION(91500, paged_cudnn_ver_error); @@ -2379,12 +2495,44 @@ class UnifiedSDPANode : public SDPANodeBase { &backend_seq_len_KV)); } + if (has_cu_seq_len_q() || has_cu_seq_len_kv()) { + auto cu_seq_len_cudnn_ver_error = + error_t{error_code_t::GRAPH_NOT_SUPPORTED, + "Cumulative sequence length tensors in unified SDPA node require cuDNN 9.24.0"}; +#if (CUDNN_VERSION >= 92400) + NV_CUDNN_FE_DYNAMIC_CHECK_CUDNN_BACKEND_VERSION(92400, cu_seq_len_cudnn_ver_error); + + if (has_cu_seq_len_q()) { + auto cu_seq_len_Q = attributes.inputs.find(SDPA_attributes::input_names::CU_SEQ_LEN_Q)->second; + auto backend_cu_seq_len_Q = tensors[cu_seq_len_Q->get_uid()]->get_desc()->get_backend_descriptor(); + _CUDNN_CHECK_CUDNN_ERROR(detail::set_attribute(unified_sdpa_operation->get_backend_descriptor(), + CUDNN_ATTR_OPERATION_SDPA_FWD_CU_SEQ_LEN_QDESC, + CUDNN_TYPE_BACKEND_DESCRIPTOR, + 1, + &backend_cu_seq_len_Q)); + } + + if (has_cu_seq_len_kv()) { + auto cu_seq_len_KV = attributes.inputs.find(SDPA_attributes::input_names::CU_SEQ_LEN_KV)->second; + auto backend_cu_seq_len_KV = + tensors[cu_seq_len_KV->get_uid()]->get_desc()->get_backend_descriptor(); + _CUDNN_CHECK_CUDNN_ERROR(detail::set_attribute(unified_sdpa_operation->get_backend_descriptor(), + CUDNN_ATTR_OPERATION_SDPA_FWD_CU_SEQ_LEN_KVDESC, + CUDNN_TYPE_BACKEND_DESCRIPTOR, + 1, + &backend_cu_seq_len_KV)); + } +#else + return cu_seq_len_cudnn_ver_error; +#endif + } + // Ignore attributes.max_seq_len_kv, because unified engine doesn't need it (it's harmless if set). // Ignore attributes.padding_mask, because unified engine already applies an implicit padding mask - // if seq_len_Q and seq_len_KV are both provided. We already checked in - // `SDPA_attributes::validate_sdpa_support_surface()` that padding_mask must be true if and - // only if seq_len_Q and seq_len_KV are both set, so we don't need to check it here. + // if either (seq_len_Q and seq_len_KV) or (cu_seq_len_Q and cu_seq_len_KV) are provided. + // We already checked in `SDPA_attributes::validate_sdpa_support_surface()` that padding_mask + // must be true if and only if one of those pairs is set, so we don't need to check it here. #else return paged_cudnn_ver_error; #endif diff --git a/include/cudnn_frontend/node/sdpa_fp8_bwd.h b/include/cudnn_frontend/node/sdpa_fp8_bwd.h index 29b8f2005..0cec5a85f 100644 --- a/include/cudnn_frontend/node/sdpa_fp8_bwd.h +++ b/include/cudnn_frontend/node/sdpa_fp8_bwd.h @@ -12,6 +12,7 @@ #include "reduction.h" #include "softmax.h" #include "block_scale_dequantize.h" +#include "scaled_dot_product_flash_attention.h" // For promote_1d_index_tensor_to_4d() namespace cudnn_frontend::graph { @@ -22,6 +23,31 @@ class SDPAFP8BackwardNode : public NodeCRTP { private: mutable bool is_deterministic_algorithm_supported_on_blackwell = false; // Will be edited in pre_validate_node() + // Promote any 1-D seq_len / ragged-offset index tensors to the 4-D + // [n, 1, 1, 1] form the cuDNN backend requires (see promote_1d_index_tensor_to_4d). + void + promote_index_tensors_to_4d() { + // TODO: Handle CU_SEQ_LEN_Q and CU_SEQ_LEN_KV once fp8 bprop supports these. + for (auto& key : {input_names::SEQ_LEN_Q, input_names::SEQ_LEN_KV}) { + auto const it = attributes.inputs.find(key); + if (it != attributes.inputs.end()) { + promote_1d_index_tensor_to_4d(it->second); + } + } + for (auto& [key, value] : attributes.inputs) { + CUDNN_FRONTEND_UNUSED(key); + if (value != nullptr) { + promote_1d_index_tensor_to_4d(value->get_ragged_offset()); + } + } + for (auto& [key, value] : attributes.outputs) { + CUDNN_FRONTEND_UNUSED(key); + if (value != nullptr) { + promote_1d_index_tensor_to_4d(value->get_ragged_offset()); + } + } + } + public: mutable SDPA_fp8_backward_attributes attributes; // mutable to allow auto-routing to deterministic in pre_validate_node() @@ -454,6 +480,10 @@ class SDPAFP8BackwardNode : public NodeCRTP { error_t infer_properties_node() override final { + // If a 1-D length/offset index tensor [n] is supplied, promote it to the 4-D + // [n, 1, 1, 1] form the cuDNN backend requires. + promote_index_tensors_to_4d(); + return {error_code_t::OK, ""}; } @@ -819,7 +849,9 @@ class SDPAFP8BackwardNode : public NodeCRTP { s_q, s_kv, s_q_ptr, - s_kv_ptr); + s_kv_ptr, + /*cu_s_q_ptr=*/nullptr, + /*cu_s_kv_ptr=*/nullptr); sub_nodes.emplace_back(std::move(node_)); } diff --git a/include/cudnn_frontend/node/sdpa_support_surface.h b/include/cudnn_frontend/node/sdpa_support_surface.h index 223539b34..af2bc6ad3 100644 --- a/include/cudnn_frontend/node/sdpa_support_surface.h +++ b/include/cudnn_frontend/node/sdpa_support_surface.h @@ -54,6 +54,11 @@ SDPA_attributes::validate_sdpa_support_surface(const detail::Context& context, auto const& seq_len_kv = inputs.find(SDPA_attributes::input_names::SEQ_LEN_KV); bool const has_seq_len_kv = (seq_len_kv != inputs.end()) && (seq_len_kv->second != nullptr); + auto const& cu_seq_len_q = inputs.find(SDPA_attributes::input_names::CU_SEQ_LEN_Q); + bool const has_cu_seq_len_q = (cu_seq_len_q != inputs.end()) && (cu_seq_len_q->second != nullptr); + auto const& cu_seq_len_kv = inputs.find(SDPA_attributes::input_names::CU_SEQ_LEN_KV); + bool const has_cu_seq_len_kv = (cu_seq_len_kv != inputs.end()) && (cu_seq_len_kv->second != nullptr); + // validation TODO: // - validate stats has valid dims @@ -79,13 +84,26 @@ 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 - RETURN_CUDNN_FRONTEND_ERROR_IF(padding_mask && (!has_seq_len_q || !has_seq_len_kv), - error_code_t::ATTRIBUTE_NOT_SET, - "Padding mask requires seq_len_q and seq_len_kv to be set."); - RETURN_CUDNN_FRONTEND_ERROR_IF((!padding_mask && !attention_score_modifier) && (has_seq_len_q || has_seq_len_kv), + // 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."); + 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), + 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, - "seq_len_q and seq_len_kv needs to be set only if padding mask is enabled."); + "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), + error_code_t::INVALID_VALUE, + "Cannot specify both seq_len tensors and cu_seq_len tensors."); RETURN_CUDNN_FRONTEND_ERROR_IF(is_ragged && ((padding_mask == false) && (attention_score_modifier == nullptr)), error_code_t::GRAPH_NOT_SUPPORTED, @@ -282,10 +300,11 @@ SDPA_attributes::validate_sdpa_support_surface(const detail::Context& context, error_code_t::GRAPH_NOT_SUPPORTED, "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), - 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."); + 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)), + 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)."); RETURN_CUDNN_FRONTEND_ERROR_IF( !is_paged && max_seq_kv_explicit, @@ -380,6 +399,10 @@ SDPA_attributes::verify_sdpa_support_surface_for_implementation(const detail::Co 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"); } break; case AttentionImplementation_t::UNIFIED: { @@ -419,6 +442,11 @@ SDPA_attributes::verify_sdpa_support_surface_for_implementation(const detail::Co allowed_input_msg += ", Bias, Seed, Offset, SINK_TOKEN"; } + if (effective_cudnn_ver >= 92400) { + allowed_input_names.insert({input_names::CU_SEQ_LEN_Q, input_names::CU_SEQ_LEN_KV}); + allowed_input_msg += ", CU_SEQ_LEN_Q, CU_SEQ_LEN_KV"; + } + for (const auto& [key, value] : inputs) { if (allowed_input_names.find(key) == allowed_input_names.end() && value != nullptr) { return {error_code_t::GRAPH_NOT_SUPPORTED, allowed_input_msg}; diff --git a/include/cudnn_frontend/node/softmax.h b/include/cudnn_frontend/node/softmax.h index 9c44468f4..1aa56894a 100644 --- a/include/cudnn_frontend/node/softmax.h +++ b/include/cudnn_frontend/node/softmax.h @@ -101,7 +101,7 @@ class CompositeSoftmaxNode : public SoftmaxNodeBase { !((!has_stats() && !has_max() && !has_sum_exp()) || // (has_stats() && !has_max() && !has_sum_exp()) || // (!has_stats() && has_max() && has_sum_exp())), - error_code_t::INVALID_VALUE, + error_code_t::GRAPH_NOT_SUPPORTED, "CompositeSoftmaxNode can only output certain combinations of stats, max and sum_exp: " "stats only, max and sum_exp only, or none of the above."); diff --git a/include/cudnn_frontend/node_interface.h b/include/cudnn_frontend/node_interface.h index bf7ebe949..8d46ece33 100644 --- a/include/cudnn_frontend/node_interface.h +++ b/include/cudnn_frontend/node_interface.h @@ -427,6 +427,8 @@ class INode { std::shared_ptr b, std::shared_ptr seq_len_q, std::shared_ptr seq_len_kv, + std::shared_ptr cu_seq_len_q, + std::shared_ptr cu_seq_len_kv, std::shared_ptr left_bound, std::shared_ptr shift_right_bound, DiagonalBandMask_attributes attributes); diff --git a/include/cudnn_frontend/utils/attn_score_modifiers.h b/include/cudnn_frontend/utils/attn_score_modifiers.h index e0597907b..a09e53840 100644 --- a/include/cudnn_frontend/utils/attn_score_modifiers.h +++ b/include/cudnn_frontend/utils/attn_score_modifiers.h @@ -47,6 +47,8 @@ causal_mask_bottom_right(std::shared_ptr graph, ); } +// Padding mask for the legacy CompositeSDPANode, not used in UnifiedSDPANode. +// cu_seq_len_kv and cu_seq_len_q are explicitly not supported. [[maybe_unused]] inline std::shared_ptr padding_mask(std::shared_ptr graph, std::shared_ptr attention_score, @@ -168,7 +170,9 @@ sliding_window_mask(std::shared_ptr graph, int64_t s_q, int64_t s_kv, std::shared_ptr s_q_ptr, - std::shared_ptr s_kv_ptr) { + std::shared_ptr s_kv_ptr, + std::shared_ptr cu_s_q_ptr, + std::shared_ptr cu_s_kv_ptr) { std::shared_ptr return_mask = attention_score; // Note: the right and left bound subtrees can be constructed in different ways as well that yield functionally @@ -179,11 +183,19 @@ sliding_window_mask(std::shared_ptr graph, if (right_bound.has_value()) { std::shared_ptr effective_s_q; std::shared_ptr effective_s_kv; + std::shared_ptr effective_cu_s_q; + std::shared_ptr effective_cu_s_kv; if (diagonal_alignment == DiagonalAlignment_t::BOTTOM_RIGHT) { - effective_s_q = - s_q_ptr != nullptr ? s_q_ptr : std::make_shared(static_cast(s_q)); - effective_s_kv = - s_kv_ptr != nullptr ? s_kv_ptr : std::make_shared(static_cast(s_kv)); + effective_s_q = s_q_ptr; + effective_s_kv = s_kv_ptr; + effective_cu_s_q = cu_s_q_ptr; + effective_cu_s_kv = cu_s_kv_ptr; + if (effective_s_q == nullptr && effective_cu_s_q == nullptr) { + effective_s_q = std::make_shared(static_cast(s_q)); + } + if (effective_s_kv == nullptr && effective_cu_s_kv == nullptr) { + effective_s_kv = std::make_shared(static_cast(s_kv)); + } } std::shared_ptr effective_shift_right_bound; @@ -196,6 +208,8 @@ sliding_window_mask(std::shared_ptr graph, /*b=*/std::make_shared(get_negative_inf_value()), /*seq_len_q=*/effective_s_q, /*seq_len_kv=*/effective_s_kv, + /*cu_seq_len_q=*/effective_cu_s_q, + /*cu_seq_len_kv=*/effective_cu_s_kv, /*left_bound=*/nullptr, /*shift_right_bound=*/effective_shift_right_bound, DiagonalBandMask_attributes() @@ -208,6 +222,8 @@ sliding_window_mask(std::shared_ptr graph, std::shared_ptr effective_left_bound; std::shared_ptr effective_s_q; std::shared_ptr effective_s_kv; + std::shared_ptr effective_cu_s_q; + std::shared_ptr effective_cu_s_kv; // When the diagonal is top left aligned: setup a graph so we can compare column + window_size > row // All elements for which column + window_size > row, will be retained, all others will be masked out @@ -221,7 +237,7 @@ sliding_window_mask(std::shared_ptr graph, // Setup a graph so we can compare column + window_size - (s_kv - s_q) > row // Optimization with fixed sequence lengths: single pointwise addition for the left-hand of the comparison // Again, all elements satisfying the comparison will be retained. - else if (s_kv_ptr == nullptr && s_q_ptr == nullptr) { + else if (s_kv_ptr == nullptr && s_q_ptr == nullptr && cu_s_kv_ptr == nullptr && cu_s_q_ptr == nullptr) { effective_left_bound = std::make_shared((float)(left_bound.value() - s_kv + s_q)); } // With bottom right diagonal alignment: general case when at least one of Q and KV have variable sequence @@ -230,16 +246,24 @@ sliding_window_mask(std::shared_ptr graph, // Also here, all elements satisfying the comparison will be retained. else { effective_left_bound = std::make_shared((int32_t)left_bound.value()); - effective_s_kv = - s_kv_ptr != nullptr ? s_kv_ptr : std::make_shared(static_cast(s_kv)); - effective_s_q = - s_q_ptr != nullptr ? s_q_ptr : std::make_shared(static_cast(s_q)); + effective_s_q = s_q_ptr; + effective_s_kv = s_kv_ptr; + effective_cu_s_q = cu_s_q_ptr; + effective_cu_s_kv = cu_s_kv_ptr; + if (effective_s_q == nullptr && effective_cu_s_q == nullptr) { + effective_s_q = std::make_shared(static_cast(s_q)); + } + if (effective_s_kv == nullptr && effective_cu_s_kv == nullptr) { + effective_s_kv = std::make_shared(static_cast(s_kv)); + } } return_mask = graph->diagonal_band_mask(/*x=*/return_mask, /*b=*/std::make_shared(get_negative_inf_value()), /*seq_len_q=*/effective_s_q, /*seq_len_kv=*/effective_s_kv, + /*cu_seq_len_q=*/effective_cu_s_q, + /*cu_seq_len_kv=*/effective_cu_s_kv, /*left_bound=*/effective_left_bound, /*shift_right_bound=*/nullptr, DiagonalBandMask_attributes() diff --git a/include/cudnn_frontend/utils/serialize.h b/include/cudnn_frontend/utils/serialize.h index 9fb36bef2..7cdd969ef 100644 --- a/include/cudnn_frontend/utils/serialize.h +++ b/include/cudnn_frontend/utils/serialize.h @@ -365,6 +365,8 @@ NLOHMANN_JSON_SERIALIZE_ENUM(SDPA_attributes::input_names, {SDPA_attributes::input_names::Bias, "Bias"}, {SDPA_attributes::input_names::SEQ_LEN_Q, "SEQ_LEN_Q"}, {SDPA_attributes::input_names::SEQ_LEN_KV, "SEQ_LEN_KV"}, + {SDPA_attributes::input_names::CU_SEQ_LEN_Q, "CU_SEQ_LEN_Q"}, + {SDPA_attributes::input_names::CU_SEQ_LEN_KV, "CU_SEQ_LEN_KV"}, {SDPA_attributes::input_names::Seed, "Seed"}, {SDPA_attributes::input_names::Offset, "Offset"}, {SDPA_attributes::input_names::Dropout_mask, "Dropout_mask"}, @@ -523,6 +525,8 @@ NLOHMANN_JSON_SERIALIZE_ENUM(DiagonalBandMask_attributes::input_names, {DiagonalBandMask_attributes::input_names::X, "X"}, {DiagonalBandMask_attributes::input_names::SEQ_LEN_Q, "SEQ_LEN_Q"}, {DiagonalBandMask_attributes::input_names::SEQ_LEN_KV, "SEQ_LEN_KV"}, + {DiagonalBandMask_attributes::input_names::CU_SEQ_LEN_Q, "CU_SEQ_LEN_Q"}, + {DiagonalBandMask_attributes::input_names::CU_SEQ_LEN_KV, "CU_SEQ_LEN_KV"}, {DiagonalBandMask_attributes::input_names::LeftBound, "LeftBound"}, {DiagonalBandMask_attributes::input_names::ShiftRightBound, "ShiftRightBound"}, {DiagonalBandMask_attributes::input_names::B, "B"}, diff --git a/python/pygraph/pygraph.h b/python/pygraph/pygraph.h index 9c8dd6295..9891c6d74 100644 --- a/python/pygraph/pygraph.h +++ b/python/pygraph/pygraph.h @@ -449,7 +449,9 @@ class PyGraph { std::shared_ptr score_max, std::shared_ptr score_sum_exp, std::shared_ptr sink_token, - bool const unfuse_fma); + bool const unfuse_fma, + std::shared_ptr& cu_seq_len_q, + std::shared_ptr& cu_seq_len_kv); // return [dQ, dK, dV] std::array, 3> @@ -816,6 +818,8 @@ class PyGraph { bool const use_padding_mask, std::shared_ptr& seq_len_q, std::shared_ptr& seq_len_kv, + std::shared_ptr& cu_seq_len_q, + std::shared_ptr& cu_seq_len_kv, cudnn_frontend::DiagonalAlignment_t const& diagonal_alignment, py::object const& left_bound, py::object const& right_bound, diff --git a/python/pygraph/sdpa.cpp b/python/pygraph/sdpa.cpp index 13ddaa3d4..e6278f575 100644 --- a/python/pygraph/sdpa.cpp +++ b/python/pygraph/sdpa.cpp @@ -24,6 +24,8 @@ PyGraph::sdpa_internal(std::shared_ptr bool const use_padding_mask, std::shared_ptr& seq_len_q, std::shared_ptr& seq_len_kv, + std::shared_ptr& cu_seq_len_q, + std::shared_ptr& cu_seq_len_kv, cudnn_frontend::DiagonalAlignment_t const& diagonal_alignment, py::object const& left_bound, py::object const& right_bound, @@ -54,6 +56,8 @@ PyGraph::sdpa_internal(std::shared_ptr .set_padding_mask(use_padding_mask) .set_seq_len_q(seq_len_q) .set_seq_len_kv(seq_len_kv) + .set_cu_seq_len_q(cu_seq_len_q) + .set_cu_seq_len_kv(cu_seq_len_kv) .set_diagonal_alignment(diagonal_alignment) .set_compute_data_type(compute_data_type) ._set_mma_core_mode(mma_core_mode) @@ -246,7 +250,9 @@ PyGraph::sdpa(std::shared_ptr& q, std::shared_ptr score_max, std::shared_ptr score_sum_exp, std::shared_ptr sink_token, - bool const unfuse_fma) { + bool const unfuse_fma, + 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::HALF; std::shared_ptr descale_q = nullptr; std::shared_ptr descale_k = nullptr; @@ -314,6 +320,8 @@ PyGraph::sdpa(std::shared_ptr& q, use_padding_mask, seq_len_q, seq_len_kv, + cu_seq_len_q, + cu_seq_len_kv, actual_diagonal_alignment, actual_left_bound, actual_right_bound, @@ -535,6 +543,9 @@ PyGraph::sdpa_fp8(std::shared_ptr& q, std::shared_ptr sink_token) { 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; @@ -595,6 +606,8 @@ PyGraph::sdpa_fp8(std::shared_ptr& q, use_padding_mask, seq_len_q, seq_len_kv, + cu_seq_len_q, + cu_seq_len_kv, actual_diagonal_alignment, actual_left_bound, actual_right_bound, @@ -1093,6 +1106,8 @@ init_pygraph_sdpa_submodule(py::class_& m) { py::arg_v("score_sum_exp", nullptr), py::arg_v("sink_token", nullptr), py::arg_v("unfuse_fma", false), + py::arg_v("cu_seq_len_q", nullptr), + py::arg_v("cu_seq_len_kv", nullptr), R"pbdoc( Perform scaled dot product attention. @@ -1104,8 +1119,8 @@ init_pygraph_sdpa_submodule(py::class_& m) { bias (Optional[cudnn_tensor]): The bias data for attention. Default is None. use_alibi_mask (Optional[bool]): Whether to use alibi mask. Default is False. use_padding_mask (Optional[bool]): Whether to use padding mask. Default is False. - seq_len_q (Optional[cudnn_tensor]): The sequence length of the query. - seq_len_kv (Optional[cudnn_tensor]): The sequence length of the key. + seq_len_q (Optional[cudnn_tensor]): The sequence length of the query. Accepts shape (b, 1, 1, 1) or 1-D (b,) (promoted automatically). + seq_len_kv (Optional[cudnn_tensor]): The sequence length of the key. Accepts shape (b, 1, 1, 1) or 1-D (b,) (promoted automatically). dropout (Optional[Union[Tuple[(probability: float, seed: cudnn_tensor, offset: cudnn_tensor)], Tuple[mask: cudnn_tensor, scale: cudnn_tensor]]]): Whether to do dropout. Default is None. rng_dump (Optional[cudnn_tensor]): Debug tensor to dump the Philox RNG dropout mask. Default is None. paged_attention_k_table (Optional[cudnn_tensor]): The page table to look up offsets into 'k' @@ -1119,6 +1134,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. 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. @@ -1183,8 +1200,8 @@ init_pygraph_sdpa_submodule(py::class_& m) { dBias (Optional[cudnn_tensor]): The dBias data for attention. Default is None. use_alibi_mask (Optional[bool]): Whether to use alibi mask. Default is False. use_padding_mask (Optional[bool]): Whether to use padding mask. Default is False. - seq_len_q (Optional[cudnn_tensor]): The sequence length of the query. - seq_len_kv (Optional[cudnn_tensor]): The sequence length of the key. + seq_len_q (Optional[cudnn_tensor]): The sequence length of the query. Accepts shape (b, 1, 1, 1) or 1-D (b,) (promoted automatically). + seq_len_kv (Optional[cudnn_tensor]): The sequence length of the key. Accepts shape (b, 1, 1, 1) or 1-D (b,) (promoted automatically). max_total_seq_len_q (Optional[int]): The maximum number of query sequence tokens for all batches, used for workspace allocation, max_total_seq_len_kv (Optional[int]): The maximum number of key/value sequence tokens for all batches, used for workspace allocation, dropout (Optional[Union[Tuple[(probability: float, seed: cudnn_tensor, offset: cudnn_tensor)], Tuple[mask: cudnn_tensor, scale: cudnn_tensor]]]): Whether to do dropout. Default is None. @@ -1263,8 +1280,8 @@ init_pygraph_sdpa_submodule(py::class_& m) { bias (Optional[cudnn_tensor]): The bias data for attention. Default is None. use_alibi_mask (Optional[bool]): Whether to use alibi mask. Default is False. use_padding_mask (Optional[bool]): Whether to use padding mask. Default is False. - seq_len_q (Optional[cudnn_tensor]): The sequence length of the query. - seq_len_kv (Optional[cudnn_tensor]): The sequence length of the key. + seq_len_q (Optional[cudnn_tensor]): The sequence length of the query. Accepts shape (b, 1, 1, 1) or 1-D (b,) (promoted automatically). + seq_len_kv (Optional[cudnn_tensor]): The sequence length of the key. Accepts shape (b, 1, 1, 1) or 1-D (b,) (promoted automatically). dropout (Optional[Union[Tuple[(probability: float, seed: cudnn_tensor, offset: cudnn_tensor)], Tuple[mask: cudnn_tensor, scale: cudnn_tensor]]]): Whether to do dropout. Default is None. rng_dump (Optional[cudnn_tensor]): Debug tensor to dump the Philox RNG dropout mask. Default is None. paged_attention_k_table (Optional[cudnn_tensor]): The page table to look up offsets into 'k'. Default is None. @@ -1417,8 +1434,8 @@ init_pygraph_sdpa_submodule(py::class_& m) { scale_dP (cudnn_tensor): Scale factor for dP gradient. attn_scale (Optional[Union[float, cudnn_tensor]]): The scale factor for attention. Default is None. use_padding_mask (bool): Whether it is an inference step or training step. - seq_len_q (Optional[cudnn_tensor]): The sequence length of the query. - seq_len_kv (Optional[cudnn_tensor]): The sequence length of the key. + seq_len_q (Optional[cudnn_tensor]): The sequence length of the query. Accepts shape (b, 1, 1, 1) or 1-D (b,) (promoted automatically). + seq_len_kv (Optional[cudnn_tensor]): The sequence length of the key. Accepts shape (b, 1, 1, 1) or 1-D (b,) (promoted automatically). use_causal_mask (Optional[bool]): Whether to use causal mask. Default is False. use_causal_mask_bottom_right (Optional[bool]): Whether to use bottom right aligned causal mask. Default is False. diagonal_alignment (Optional[cudnn.diagonal_alignment]): One of {"TOP_LEFT", "BOTTOM_RIGHT"}. Default is TOP_LEFT. @@ -1498,8 +1515,8 @@ init_pygraph_sdpa_submodule(py::class_& m) { descale_dO_T (cudnn_tensor): Descale factor for transposed output gradient. attn_scale (Optional[Union[float, cudnn_tensor]]): The scale factor for attention. Default is None. use_padding_mask (bool): Whether it is an inference step or training step. - seq_len_q (Optional[cudnn_tensor]): The sequence length of the query. - seq_len_kv (Optional[cudnn_tensor]): The sequence length of the key. + seq_len_q (Optional[cudnn_tensor]): The sequence length of the query. Accepts shape (b, 1, 1, 1) or 1-D (b,) (promoted automatically). + seq_len_kv (Optional[cudnn_tensor]): The sequence length of the key. Accepts shape (b, 1, 1, 1) or 1-D (b,) (promoted automatically). use_causal_mask (Optional[bool]): Whether to use causal mask. Default is False. use_causal_mask_bottom_right (Optional[bool]): Whether to use bottom right aligned causal mask. Default is False. diagonal_alignment (Optional[cudnn.diagonal_alignment]): Alignment of the diagonal band. Default is TOP_LEFT. diff --git a/samples/cpp/CMakeLists.txt b/samples/cpp/CMakeLists.txt index c64c013e9..6cb727149 100644 --- a/samples/cpp/CMakeLists.txt +++ b/samples/cpp/CMakeLists.txt @@ -10,6 +10,7 @@ add_executable( sdpa/fp16_bwd_with_flexible_graphs.cpp sdpa/fp16_fwd_with_custom_dropout.cpp sdpa/fp16_fwd_with_paged_caches.cpp + sdpa/fp16_fwd_with_cu_seq_len.cpp sdpa/fp16_dynamic_shapes.cpp sdpa/fp16_fwd_paged_decode_and_prefill.cpp sdpa/fp16_fwd_with_cudagraphs.cpp diff --git a/samples/cpp/sdpa/fp16_fwd_with_cu_seq_len.cpp b/samples/cpp/sdpa/fp16_fwd_with_cu_seq_len.cpp new file mode 100644 index 000000000..96e2c713c --- /dev/null +++ b/samples/cpp/sdpa/fp16_fwd_with_cu_seq_len.cpp @@ -0,0 +1,211 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a + * copy of this software and associated documentation files (the "Software"), + * to deal in the Software without restriction, including without limitation + * the rights to use, copy, modify, merge, publish, distribute, sublicense, + * and/or sell copies of the Software, and to permit persons to whom the + * Software is furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL + * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING + * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + * DEALINGS IN THE SOFTWARE. + */ + +#include +#include "../utils/helpers.h" + +#include + +#include +namespace fe = cudnn_frontend; + +/* +Run this example by using command: +bin/samples "Toy sdpa forward with cu_seq_len" + +This example shows how to construct a sdpa forward graph that supplies per-batch +sequence lengths via cumulative-sequence-length tensors (cu_seq_len_q / cu_seq_len_kv) +instead of the more common seq_len_q / seq_len_kv form. + +cu_seq_len_* tensors have shape (b+1, 1, 1, 1) and store the prefix-sum of the +per-batch actual sequence lengths, with a leading 0: + cu_seq_len[i] = sum(actual_seq_len[0..i-1]) for i in [0, b] +For example, with b=3 and actual_seq_len = {12, 20, 8}, cu_seq_len = {0, 12, 32, 40}. +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. + - Only the UNIFIED SDPA implementation supports cu_seq_len_*; the COMPOSITE path + will reject the inputs explicitly. + - Requires cuDNN >= 9.24.0. +*/ + +// Tensors in forward pass +#define Q_UID 1 +#define K_UID 2 +#define V_UID 3 +#define O_UID 4 +#define STATS_UID 5 +#define CU_SEQ_LEN_Q_UID 6 +#define CU_SEQ_LEN_KV_UID 7 + +std::shared_ptr +create_sdpa_forward_graph_with_cu_seq_len(int64_t const b, + int64_t const h_q, + int64_t const h_k, + int64_t const h_v, + int64_t const s_q, + int64_t const s_kv, + int64_t const d_qk, + int64_t const d_v, + float const attn_scale = 1.0f, + bool const generate_stats = true) { + auto graph = std::make_shared(); + graph->set_io_data_type(fe::DataType_t::BFLOAT16) + .set_intermediate_data_type(fe::DataType_t::FLOAT) + .set_compute_data_type(fe::DataType_t::FLOAT); + + auto Q = graph->tensor(fe::graph::Tensor_attributes() + .set_name("Q") + .set_uid(Q_UID) + .set_dim({b, h_q, s_q, d_qk}) + .set_stride({h_q * s_q * d_qk, s_q * d_qk, d_qk, 1})); + + auto K = graph->tensor(fe::graph::Tensor_attributes() + .set_name("K") + .set_uid(K_UID) + .set_dim({b, h_k, s_kv, d_qk}) + .set_stride({h_k * s_kv * d_qk, s_kv * d_qk, d_qk, 1})); + + auto V = graph->tensor(fe::graph::Tensor_attributes() + .set_name("V") + .set_uid(V_UID) + .set_dim({b, h_v, s_kv, d_v}) + .set_stride({h_v * s_kv * d_v, s_kv * d_v, d_v, 1})); + + // Cumulative sequence-length tensors are int32. The frontend accepts either the 4-D + // (b+1, 1, 1, 1) form the backend requires, or a 1-D (b+1,) form that is promoted to + // 4-D automatically. Here we demonstrate the 1-D form. + auto cu_seq_q = graph->tensor(fe::graph::Tensor_attributes() + .set_name("cu_seq_q") + .set_uid(CU_SEQ_LEN_Q_UID) + .set_dim({b + 1}) + .set_stride({1}) + .set_data_type(fe::DataType_t::INT32)); + + auto cu_seq_kv = graph->tensor(fe::graph::Tensor_attributes() + .set_name("cu_seq_kv") + .set_uid(CU_SEQ_LEN_KV_UID) + .set_dim({b + 1}) + .set_stride({1}) + .set_data_type(fe::DataType_t::INT32)); + + auto sdpa_options = fe::graph::SDPA_attributes() + .set_name("flash_attention_cu_seq_len") + .set_generate_stats(generate_stats) + .set_attn_scale(attn_scale) + .set_padding_mask(true) + .set_cu_seq_len_q(cu_seq_q) + .set_cu_seq_len_kv(cu_seq_kv) + // cu_seq_len_* is unified-only; force the implementation so the + // sample fails loudly instead of silently falling back to composite. + .set_implementation(fe::AttentionImplementation_t::UNIFIED); + + auto [O, Stats] = graph->sdpa(Q, K, V, sdpa_options); + + O->set_output(true).set_dim({b, h_q, s_q, d_v}).set_stride({h_q * d_v, d_v, b * h_q * d_v, 1}).set_uid(O_UID); + + if (generate_stats) { + Stats->set_output(true).set_data_type(fe::DataType_t::FLOAT).set_uid(STATS_UID); + } else { + assert(Stats == nullptr); + } + + return graph; +} + +TEST_CASE("Toy sdpa forward with cu_seq_len", "[graph][sdpa][flash][forward][cu_seq_len]") { + if (cudnnGetVersion() < 92400) { + SKIP("cu_seq_len_q/cu_seq_len_kv require cuDNN 9.24.0 or above"); + return; + } + + int64_t const b = 3; // batch size + int64_t const h_q = 4; // head count for Q + int64_t const h_k = 4; // head count for K + int64_t const h_v = 4; // head count for V + int64_t const s_q = 1024; // q tensor is padded to this seq length + int64_t const s_kv = 1024; // k and v tensor is padded to this seq length + int64_t const d_qk = 128; // hidden dim + int64_t const d_v = 128; // hidden dim + bool const generate_stats = true; + float const attn_scale = 0.123f; + + auto handle_ptr = create_cudnn_handle(); + auto handle = *handle_ptr; + + auto graph = + create_sdpa_forward_graph_with_cu_seq_len(b, h_q, h_k, h_v, s_q, s_kv, d_qk, d_v, attn_scale, generate_stats); + + REQUIRE(graph->build(handle, {fe::HeurMode_t::A}).is_good()); + + Surface q_tensor(b * h_q * s_q * d_qk); + Surface k_tensor(b * h_k * d_qk * s_kv); + Surface v_tensor(b * h_v * d_v * s_kv); + Surface o_tensor(b * s_q * h_q * d_v); + + // Compute prefix-sum (cumulative) sequence lengths: shape (b+1,), with cu[0] == 0. + // For this toy example all batches use the same actual length (20), but cu_seq_len + // supports arbitrary per-batch lengths. + std::vector hostActualSeqlenQ(b, 20); + std::vector hostActualSeqlenKV(b, 20); + + std::vector hostCuSeqlenQ(b + 1, 0); + std::vector hostCuSeqlenKV(b + 1, 0); + for (int64_t i = 0; i < b; ++i) { + hostCuSeqlenQ[i + 1] = hostCuSeqlenQ[i] + hostActualSeqlenQ[i]; + hostCuSeqlenKV[i + 1] = hostCuSeqlenKV[i] + hostActualSeqlenKV[i]; + } + + Surface devCuSeqlenQ(b + 1); + Surface devCuSeqlenKV(b + 1); + + CUDA_CHECK(cudaMemcpy( + devCuSeqlenQ.devPtr, hostCuSeqlenQ.data(), sizeof(hostCuSeqlenQ[0]) * (b + 1), cudaMemcpyHostToDevice)); + CUDA_CHECK(cudaMemcpy( + devCuSeqlenKV.devPtr, hostCuSeqlenKV.data(), sizeof(hostCuSeqlenKV[0]) * (b + 1), cudaMemcpyHostToDevice)); + CUDA_CHECK(cudaDeviceSynchronize()); + + std::unordered_map variant_pack = { + {Q_UID, q_tensor.devPtr}, + {K_UID, k_tensor.devPtr}, + {V_UID, v_tensor.devPtr}, + {O_UID, o_tensor.devPtr}, + {CU_SEQ_LEN_Q_UID, devCuSeqlenQ.devPtr}, + {CU_SEQ_LEN_KV_UID, devCuSeqlenKV.devPtr}, + }; + + Surface statsTensor(b * h_q * s_q * 1); + if (generate_stats) { + variant_pack[STATS_UID] = statsTensor.devPtr; + } + + int64_t workspace_size = 0; + REQUIRE(graph->get_workspace_size(workspace_size).is_good()); + Surface workspace(workspace_size); + + REQUIRE(graph->execute(handle, variant_pack, workspace.devPtr).is_good()); + + CUDA_CHECK(cudaDeviceSynchronize()); +} diff --git a/test/python/sdpa/fp16.py b/test/python/sdpa/fp16.py index 5719c9a6c..736a15619 100644 --- a/test/python/sdpa/fp16.py +++ b/test/python/sdpa/fp16.py @@ -59,6 +59,8 @@ class TensorUid(IntEnum): k_rot = 33 dQ_rot = 34 dK_rot = 35 + cu_seq_len_q = 36 + cu_seq_len_kv = 37 def validate_config(cfg): if not all((x > 0 and type(x) == int) for x in (cfg.batches, cfg.d_qk, cfg.d_v, cfg.s_q, cfg.s_kv, cfg.h_q, cfg.h_k, cfg.h_v)): @@ -76,6 +78,10 @@ def validate_config(cfg): if cfg.is_ragged: assert cfg.is_padding == True, "is_ragged=True and is_padding=False not allowed" + 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 isinstance(cfg.seq_len_q, (list, tuple)), "input 'seq_len_q' must be list or tuple" if cfg.is_padding: assert len(cfg.seq_len_q) == cfg.batches, f"wrong 'seq_len_q' length" @@ -109,6 +115,10 @@ def validate_config(cfg): print("@@@@ Overall result: WAIVED, certain combinations of softmax outputs require cuDNN 9.20.0 or higher.") pytest.skip("certain combinations of softmax outputs require cuDNN 9.20.0 or higher") + if cudnn_version < "9.24.0" and cfg.is_cu_seq_len: + 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") + def allocate_tensors(cfg, rng_data_gen, perf=False): allocs = {} @@ -146,10 +156,18 @@ def allocate_tensors(cfg, rng_data_gen, perf=False): allocs[TensorUid.dV] = alloc_tensor(cfg.shape_v, cfg.data_type, strides=cfg.stride_v) allocs[TensorUid.dO] = alloc_tensor(cfg.shape_o, cfg.data_type, strides=cfg.stride_o, rng=rng_data_gen, mean=0.0, std=0.1, sparse_int=si) - seq_len_q_gpu = torch.tensor(cfg.seq_len_q, dtype=torch.int32, device="cuda").view(-1, 1, 1, 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, 1, 1, 1) if len(cfg.seq_len_kv) > 0 else None - allocs[TensorUid.seq_len_q] = (seq_len_q_gpu, None, None) - allocs[TensorUid.seq_len_kv] = (seq_len_kv_gpu, None, None) + 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). + 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) + allocs[TensorUid.seq_len_kv] = (seq_len_kv_gpu, None, None) if cfg.is_ragged: allocs[TensorUid.q_ragged_offset] = ((prefix_sum(seq_len_q_gpu) * cfg.h_q * cfg.d_qk).to(torch.int64), None, None) @@ -241,8 +259,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, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT32) if cfg.is_padding else None - seq_len_kv = graph.tensor(uid=int(TensorUid.seq_len_kv), dim=(cfg.batches, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT32) if cfg.is_padding 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 + + 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 seed = offset = dropout_tuple = rng_dump = None if cfg.is_dropout: @@ -251,11 +272,11 @@ def create_forward_graph(cfg, tensors, cudnn_handle): dropout_tuple = (cfg.dropout_prob, seed, offset) rng_dump = graph.tensor(uid=int(TensorUid.rng_dump), dim=(cfg.batches, cfg.h_q, cfg.s_q, cfg.s_kv), stride=(cfg.h_q * cfg.s_q * cfg.s_kv, cfg.s_q * cfg.s_kv, cfg.s_kv, 1), data_type=cudnn.data_type.FLOAT) - q_ragged_offset = graph.tensor(uid=int(TensorUid.q_ragged_offset), dim=(cfg.batches + 1, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT64) if cfg.is_ragged else None - k_ragged_offset = graph.tensor(uid=int(TensorUid.k_ragged_offset), dim=(cfg.batches + 1, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT64) if cfg.is_ragged else None - v_ragged_offset = graph.tensor(uid=int(TensorUid.v_ragged_offset), dim=(cfg.batches + 1, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT64) if cfg.is_ragged else None - o_ragged_offset = graph.tensor(uid=int(TensorUid.o_ragged_offset), dim=(cfg.batches + 1, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT64) if cfg.is_ragged else None - stats_ragged_offset = graph.tensor(uid=int(TensorUid.stats_ragged_offset), dim=(cfg.batches + 1, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT64) if cfg.is_ragged else None + q_ragged_offset = graph.tensor(uid=int(TensorUid.q_ragged_offset), dim=(cfg.batches + 1,), stride=(1,), data_type=cudnn.data_type.INT64) if cfg.is_ragged else None + k_ragged_offset = graph.tensor(uid=int(TensorUid.k_ragged_offset), dim=(cfg.batches + 1,), stride=(1,), data_type=cudnn.data_type.INT64) if cfg.is_ragged else None + v_ragged_offset = graph.tensor(uid=int(TensorUid.v_ragged_offset), dim=(cfg.batches + 1,), stride=(1,), data_type=cudnn.data_type.INT64) if cfg.is_ragged else None + o_ragged_offset = graph.tensor(uid=int(TensorUid.o_ragged_offset), dim=(cfg.batches + 1,), stride=(1,), data_type=cudnn.data_type.INT64) if cfg.is_ragged else None + stats_ragged_offset = graph.tensor(uid=int(TensorUid.stats_ragged_offset), dim=(cfg.batches + 1,), stride=(1,), data_type=cudnn.data_type.INT64) if cfg.is_ragged else None if cfg.is_ragged: q.set_ragged_offset(q_ragged_offset) @@ -299,6 +320,8 @@ def create_forward_graph(cfg, tensors, cudnn_handle): use_padding_mask=cfg.is_padding, seq_len_q=seq_len_q, seq_len_kv=seq_len_kv, + cu_seq_len_q=cu_seq_len_q, + cu_seq_len_kv=cu_seq_len_kv, diagonal_band_left_bound=cfg.left_bound, diagonal_band_right_bound=cfg.right_bound, diagonal_alignment=cfg.diag_align, @@ -349,6 +372,8 @@ def create_forward_graph(cfg, tensors, cudnn_handle): int(TensorUid.block_mask): tensors.get(TensorUid.block_mask), int(TensorUid.seq_len_q): tensors.get(TensorUid.seq_len_q), int(TensorUid.seq_len_kv): tensors.get(TensorUid.seq_len_kv), + int(TensorUid.cu_seq_len_q): tensors.get(TensorUid.cu_seq_len_q), + int(TensorUid.cu_seq_len_kv): tensors.get(TensorUid.cu_seq_len_kv), int(TensorUid.q_ragged_offset): tensors.get(TensorUid.q_ragged_offset), int(TensorUid.k_ragged_offset): tensors.get(TensorUid.k_ragged_offset), int(TensorUid.v_ragged_offset): tensors.get(TensorUid.v_ragged_offset), @@ -408,8 +433,8 @@ def create_backward_graph(cfg, tensors, cudnn_handle, max_t_q, max_t_kv): bias = graph.tensor(uid=int(TensorUid.bias), dim=bias_dim, stride=bias_stride, data_type=cudnn_dtype) if cfg.is_bias else None dBias = graph.tensor(uid=int(TensorUid.dBias), dim=bias_dim, stride=bias_stride, data_type=cudnn_dtype) if cfg.is_bias and not(cfg.d_qk == 256 and cfg.d_v == 256) else None - seq_len_q = graph.tensor(uid=int(TensorUid.seq_len_q), dim=(cfg.batches, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT32) if cfg.is_padding else None - seq_len_kv = graph.tensor(uid=int(TensorUid.seq_len_kv), dim=(cfg.batches, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT32) if cfg.is_padding 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 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 else None seed = offset = dropout_tuple = None if cfg.is_dropout: @@ -459,11 +484,11 @@ def create_backward_graph(cfg, tensors, cudnn_handle, max_t_q, max_t_kv): dK_orig.set_uid(int(TensorUid.dK)).set_output(True).set_data_type(cudnn_dtype).set_dim(cfg.shape_k).set_stride(cfg.stride_k) if cfg.is_ragged: - q_ragged_offset = graph.tensor(uid=int(TensorUid.q_ragged_offset), dim=(cfg.batches + 1, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT64) - k_ragged_offset = graph.tensor(uid=int(TensorUid.k_ragged_offset), dim=(cfg.batches + 1, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT64) - v_ragged_offset = graph.tensor(uid=int(TensorUid.v_ragged_offset), dim=(cfg.batches + 1, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT64) - o_ragged_offset = graph.tensor(uid=int(TensorUid.o_ragged_offset), dim=(cfg.batches + 1, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT64) - stats_ragged_offset = graph.tensor(uid=int(TensorUid.stats_ragged_offset), dim=(cfg.batches + 1, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT64) + q_ragged_offset = graph.tensor(uid=int(TensorUid.q_ragged_offset), dim=(cfg.batches + 1,), stride=(1,), data_type=cudnn.data_type.INT64) + k_ragged_offset = graph.tensor(uid=int(TensorUid.k_ragged_offset), dim=(cfg.batches + 1,), stride=(1,), data_type=cudnn.data_type.INT64) + v_ragged_offset = graph.tensor(uid=int(TensorUid.v_ragged_offset), dim=(cfg.batches + 1,), stride=(1,), data_type=cudnn.data_type.INT64) + o_ragged_offset = graph.tensor(uid=int(TensorUid.o_ragged_offset), dim=(cfg.batches + 1,), stride=(1,), data_type=cudnn.data_type.INT64) + stats_ragged_offset = graph.tensor(uid=int(TensorUid.stats_ragged_offset), dim=(cfg.batches + 1,), stride=(1,), data_type=cudnn.data_type.INT64) q.set_ragged_offset(q_ragged_offset) k.set_ragged_offset(k_ragged_offset) v.set_ragged_offset(v_ragged_offset) @@ -594,6 +619,8 @@ def compute_and_compare_reference(cfg, allocs, tensors, diffs): dO_gpu = tensors.get(TensorUid.dO) seq_len_q_gpu = tensors.get(TensorUid.seq_len_q) seq_len_kv_gpu = tensors.get(TensorUid.seq_len_kv) + cu_seq_len_q_gpu = tensors.get(TensorUid.cu_seq_len_q) + cu_seq_len_kv_gpu = tensors.get(TensorUid.cu_seq_len_kv) block_mask_gpu = tensors.get(TensorUid.block_mask) bias_gpu = tensors.get(TensorUid.bias) rng_dump_gpu = tensors.get(TensorUid.rng_dump) @@ -628,8 +655,18 @@ def compute_and_compare_reference(cfg, allocs, tensors, diffs): k_ref = k_ref.to(cfg.data_type).float() dO_ref = dO_gpu.detach().float() if dO_gpu is not None else None - seq_len_q_ref = seq_len_q_gpu.flatten().detach() if seq_len_q_gpu is not None else None - seq_len_kv_ref = seq_len_kv_gpu.flatten().detach() if seq_len_kv_gpu is not None else None + # The reference accepts the per-batch seq_lens in either form: explicit seq_len_* + # tensors or the cumulative cu_seq_len_* tensors (recovered via consecutive diffs). + def _to_seq_len_ref(seq_len_gpu, cu_seq_len_gpu): + if seq_len_gpu is not None: + return seq_len_gpu.flatten().detach() + if cu_seq_len_gpu is not None: + cu = cu_seq_len_gpu.flatten().detach() + return (cu[1:] - cu[:-1]).to(torch.int32) + return None + + seq_len_q_ref = _to_seq_len_ref(seq_len_q_gpu, cu_seq_len_q_gpu) + seq_len_kv_ref = _to_seq_len_ref(seq_len_kv_gpu, cu_seq_len_kv_gpu) block_mask_ref = block_mask_gpu.detach() if block_mask_gpu is not None else None bias_ref = bias_gpu.detach().float() if bias_gpu is not None else None rng_dump_ref = rng_dump_gpu.detach().float() if rng_dump_gpu is not None else None diff --git a/test/python/sdpa/fp8.py b/test/python/sdpa/fp8.py index ab4f2e624..92e954d44 100644 --- a/test/python/sdpa/fp8.py +++ b/test/python/sdpa/fp8.py @@ -114,21 +114,21 @@ def generate_graph_fwd(cudnn_itype, cudnn_otype, b, h_q, h_k, h_v, s_qo, s_kv, d v = graph_fwd.tensor(uid=GraphFwdUid.v, dim=(num_blocks, h_v, block_size, d_vo), stride=(block_size * h_v * d_vo, block_size * d_vo, d_vo, 1), data_type=cudnn_itype) use_padding_mask = True - kv_seq_len = graph_fwd.tensor(uid=GraphFwdUid.kv_seq_len, dim=(b, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT32) - q_seq_len = graph_fwd.tensor(uid=GraphFwdUid.q_seq_len, dim=(b, 1, 1, 1), stride=(1, 1, 1, 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_seq_len = graph_fwd.tensor(uid=GraphFwdUid.q_seq_len, dim=(b,), stride=(1,), data_type=cudnn.data_type.INT32) k_block_table = graph_fwd.tensor(uid=GraphFwdUid.k_block_table, dim=(b, 1, table_size, 1), stride=(table_size, table_size, 1, 1), data_type=cudnn.data_type.INT32) v_block_table = graph_fwd.tensor(uid=GraphFwdUid.v_block_table, dim=(b, 1, table_size, 1), stride=(table_size, table_size, 1, 1), data_type=cudnn.data_type.INT32) if is_ragged: use_padding_mask = True - q_seq_len = graph_fwd.tensor(uid=GraphFwdUid.q_seq_len, dim=(b, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT32) - kv_seq_len = graph_fwd.tensor(uid=GraphFwdUid.kv_seq_len, dim=(b, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT32) - - q_ragged_offset = graph_fwd.tensor(uid=int(GraphFwdUid.q_ragged_offset), dim=(b + 1, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT64) - k_ragged_offset = graph_fwd.tensor(uid=int(GraphFwdUid.k_ragged_offset), dim=(b + 1, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT64) - v_ragged_offset = graph_fwd.tensor(uid=int(GraphFwdUid.v_ragged_offset), dim=(b + 1, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT64) - o_ragged_offset = graph_fwd.tensor(uid=int(GraphFwdUid.o_ragged_offset), dim=(b + 1, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT64) - stats_ragged_offset = graph_fwd.tensor(uid=int(GraphFwdUid.stats_ragged_offset), dim=(b + 1, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT64) if generate_stats else None + 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) + v_ragged_offset = graph_fwd.tensor(uid=int(GraphFwdUid.v_ragged_offset), dim=(b + 1,), stride=(1,), data_type=cudnn.data_type.INT64) + o_ragged_offset = graph_fwd.tensor(uid=int(GraphFwdUid.o_ragged_offset), dim=(b + 1,), stride=(1,), data_type=cudnn.data_type.INT64) + stats_ragged_offset = graph_fwd.tensor(uid=int(GraphFwdUid.stats_ragged_offset), dim=(b + 1,), stride=(1,), data_type=cudnn.data_type.INT64) if generate_stats else None q.set_ragged_offset(q_ragged_offset) k.set_ragged_offset(k_ragged_offset) v.set_ragged_offset(v_ragged_offset) @@ -198,15 +198,15 @@ def generate_graph_bwd(cudnn_itype, cudnn_otype, b, h_q, h_k, h_v, s_qo, s_kv, d if is_ragged: use_padding_mask = True - seq_len_q = graph_bwd.tensor(uid=int(GraphBwdUid.q_seq_len), dim=(b, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT32) - seq_len_kv = graph_bwd.tensor(uid=int(GraphBwdUid.kv_seq_len), dim=(b, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT32) - - q_ragged_offset = graph_bwd.tensor(uid=int(GraphBwdUid.q_ragged_offset), dim=(b + 1, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT64) - k_ragged_offset = graph_bwd.tensor(uid=int(GraphBwdUid.k_ragged_offset), dim=(b + 1, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT64) - v_ragged_offset = graph_bwd.tensor(uid=int(GraphBwdUid.v_ragged_offset), dim=(b + 1, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT64) - o_ragged_offset = graph_bwd.tensor(uid=int(GraphBwdUid.o_ragged_offset), dim=(b + 1, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT64) - stats_ragged_offset = graph_bwd.tensor(uid=int(GraphBwdUid.stats_ragged_offset), dim=(b + 1, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT64) - dO_ragged_offset = graph_bwd.tensor(uid=int(GraphBwdUid.dO_ragged_offset), dim=(b + 1, 1, 1, 1), stride=(1, 1, 1, 1), data_type=cudnn.data_type.INT64) + seq_len_q = graph_bwd.tensor(uid=int(GraphBwdUid.q_seq_len), dim=(b,), stride=(1,), data_type=cudnn.data_type.INT32) + seq_len_kv = graph_bwd.tensor(uid=int(GraphBwdUid.kv_seq_len), dim=(b,), stride=(1,), data_type=cudnn.data_type.INT32) + + q_ragged_offset = graph_bwd.tensor(uid=int(GraphBwdUid.q_ragged_offset), dim=(b + 1,), stride=(1,), data_type=cudnn.data_type.INT64) + k_ragged_offset = graph_bwd.tensor(uid=int(GraphBwdUid.k_ragged_offset), dim=(b + 1,), stride=(1,), data_type=cudnn.data_type.INT64) + v_ragged_offset = graph_bwd.tensor(uid=int(GraphBwdUid.v_ragged_offset), dim=(b + 1,), stride=(1,), data_type=cudnn.data_type.INT64) + o_ragged_offset = graph_bwd.tensor(uid=int(GraphBwdUid.o_ragged_offset), dim=(b + 1,), stride=(1,), data_type=cudnn.data_type.INT64) + stats_ragged_offset = graph_bwd.tensor(uid=int(GraphBwdUid.stats_ragged_offset), dim=(b + 1,), stride=(1,), data_type=cudnn.data_type.INT64) + dO_ragged_offset = graph_bwd.tensor(uid=int(GraphBwdUid.dO_ragged_offset), dim=(b + 1,), stride=(1,), data_type=cudnn.data_type.INT64) q.set_ragged_offset(q_ragged_offset) k.set_ragged_offset(k_ragged_offset) v.set_ragged_offset(v_ragged_offset) @@ -326,8 +326,8 @@ def exec_sdpa_fp8(cfg, request, cudnn_handle): seq_len_kv_list = cfg.seq_len_kv if hasattr(cfg, 'seq_len_kv') and cfg.seq_len_kv else [] if is_ragged: - seq_len_q_gpu = torch.tensor(seq_len_q_list, dtype=torch.int32, device="cuda").view(-1, 1, 1, 1) - seq_len_kv_gpu = torch.tensor(seq_len_kv_list, dtype=torch.int32, device="cuda").view(-1, 1, 1, 1) + seq_len_q_gpu = torch.tensor(seq_len_q_list, dtype=torch.int32, device="cuda").view(-1) + seq_len_kv_gpu = torch.tensor(seq_len_kv_list, dtype=torch.int32, device="cuda").view(-1) 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) @@ -440,14 +440,14 @@ def exec_sdpa_fp8(cfg, request, cudnn_handle): if is_paged: variant_pack[int(GraphFwdUid.k)] = container_k_gpu variant_pack[int(GraphFwdUid.v)] = container_v_gpu - variant_pack[int(GraphFwdUid.kv_seq_len)] = torch.full((b, 1, 1, 1), s_kv, device="cuda", dtype=torch.int32) - variant_pack[int(GraphFwdUid.q_seq_len)] = torch.full((b, 1, 1, 1), s_qo, device="cuda", dtype=torch.int32) + variant_pack[int(GraphFwdUid.kv_seq_len)] = torch.full((b,), s_kv, device="cuda", dtype=torch.int32) + variant_pack[int(GraphFwdUid.q_seq_len)] = torch.full((b,), s_qo, device="cuda", dtype=torch.int32) variant_pack[int(GraphFwdUid.k_block_table)] = k_block_table_gpu 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, 1, 1, 1) - variant_pack[int(GraphFwdUid.kv_seq_len)] = torch.tensor(seq_len_kv_list, dtype=torch.int32, device="cuda").view(-1, 1, 1, 1) + 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/sdpa/random_config.py b/test/python/sdpa/random_config.py index fbe53315e..02ff20790 100644 --- a/test/python/sdpa/random_config.py +++ b/test/python/sdpa/random_config.py @@ -78,6 +78,10 @@ class ExecConfig: is_bias: bool = None is_block_mask: bool = None is_padding: bool = None + # When True, supply per-batch sequence lengths via cu_seq_len_q/cu_seq_len_kv + # (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 is_ragged: bool = None is_dropout: bool = None is_determin: bool = None @@ -252,8 +256,9 @@ def __call__(self, rng, rng_data_seed, rng_geom_seed=None): randoms_.d_qk, randoms_.d_v = randoms["d_qk_d_v"] randoms_.h_q, randoms_.h_k, randoms_.h_v = randoms["head_count"] - randoms_.is_ragged = randoms["is_ragged_or_padded_or_full"] == "ragged" - randoms_.is_padding = randoms["is_ragged_or_padded_or_full"] == "padded" or randoms["is_ragged_or_padded_or_full"] == "ragged" + randoms_.is_ragged = randoms["is_ragged_or_padded_or_full"] in ("ragged", "cu_ragged") + randoms_.is_padding = randoms["is_ragged_or_padded_or_full"] in ("padded", "ragged", "cu_padded", "cu_ragged") + randoms_.is_cu_seq_len = randoms["is_ragged_or_padded_or_full"] in ("cu_padded", "cu_ragged") if randoms["is_ragged_or_padded_or_full"] != "full": # ~10% chance of 0-length sequence for each batch diff --git a/test/python/test_mhas_v2.py b/test/python/test_mhas_v2.py index 93ee1a4da..d83d2ef67 100644 --- a/test/python/test_mhas_v2.py +++ b/test/python/test_mhas_v2.py @@ -149,7 +149,7 @@ def test_sdpa_random_fwd_unified_L0(env_info, test_no, request, cudnn_handle): diag_align=RandomChoice({cudnn.diagonal_alignment.TOP_LEFT : 1, cudnn.diagonal_alignment.BOTTOM_RIGHT : 1}), is_bias=RandomChoice({True : 1, False : 3}), is_alibi=RandomChoice({True : 1, False : 3}), - is_ragged_or_padded_or_full=RandomChoice({"ragged" : 0, "padded" : 1, "full" : 1}), + is_ragged_or_padded_or_full=RandomChoice({"ragged" : 0, "padded" : 1, "cu_padded" : 1, "full" : 1}), with_unfuse_fma=RandomChoice({True : 1, False : 1}), # Randomly enable unfuse_fma for SM100 with_score_max=RandomChoice({True : 1, False : 3}), with_score_sum_exp=RandomChoice({True : 1, False : 3}), @@ -407,7 +407,7 @@ def test_sdpa_random_fwd_ragged_unified_L0(env_info, test_no, request, cudnn_han data_type=RandomChoice({torch.float16 : 1, torch.bfloat16 : 2}), with_sliding_mask=SlidingWindowMaskGenerator(no_mask=10), # Modified from non-unified test diag_align=RandomChoice({cudnn.diagonal_alignment.TOP_LEFT : 1, cudnn.diagonal_alignment.BOTTOM_RIGHT : 0}), # Modified from non-unified test - 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, "padded" : 0, "full" : 0}), with_score_max=RandomChoice({True : 1, False : 3}), with_score_sum_exp=RandomChoice({True : 1, False : 3}), with_sink_token=RandomChoice({True : 1, False : 3}), @@ -513,7 +513,7 @@ def test_sdpa_fwd_paged_unified_L0(env_info, test_no, request, cudnn_handle): data_type=RandomChoice({torch.float16 : 1, torch.bfloat16 : 2}), with_sliding_mask=SlidingWindowMaskGenerator(no_mask=10), # Modified from non-unified test diag_align=RandomChoice({cudnn.diagonal_alignment.TOP_LEFT : 1, cudnn.diagonal_alignment.BOTTOM_RIGHT : 0}), # Modified from non-unified test - is_ragged_or_padded_or_full=RandomChoice({"ragged" : 0, "padded" : 1, "full" : 0}), + is_ragged_or_padded_or_full=RandomChoice({"ragged" : 0, "padded" : 1, "cu_padded" : 1, "full" : 0}), block_size=RandomBlockSize(min=1, max=1024, with_high_probability=[1,32,128]), ) as randomization_ctx: test.cfg = randomization_ctx(rng, data_seed, geom_seed)