From 0db957875ef7e06f9e94abe5b0285857327084c2 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Fri, 24 Apr 2026 17:37:30 -0700 Subject: [PATCH 01/32] Update GQA to support Gemma4 --- .../cpu/bert/attention_parameters.h | 4 + .../contrib_ops/cpu/bert/gqa_attention_base.h | 39 ++++-- .../cpu/bert/group_query_attention.cc | 28 ++++- .../cpu/bert/group_query_attention_helper.h | 114 +++++++++++++++++- .../contrib_ops/cuda/bert/attention_data.h | 4 + .../cuda/bert/group_query_attention.cc | 49 ++++++-- .../cuda/bert/group_query_attention_impl.cu | 17 +++ .../core/graph/contrib_ops/bert_defs.cc | 15 +++ 8 files changed, 240 insertions(+), 30 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/bert/attention_parameters.h b/onnxruntime/contrib_ops/cpu/bert/attention_parameters.h index f316a0dfdf91c..101ec88df375c 100644 --- a/onnxruntime/contrib_ops/cpu/bert/attention_parameters.h +++ b/onnxruntime/contrib_ops/cpu/bert/attention_parameters.h @@ -101,6 +101,10 @@ struct GroupQueryAttentionParameters : AttentionParameters { KVQuantizationType k_quant_type = KVQuantizationType::NONE; KVQuantizationType v_quant_type = KVQuantizationType::NONE; int kv_cache_bit_width = 0; + + // External KV parameters for KV-shared layers (e.g., Gemma4) + bool use_external_kv = false; // When true, use external K,V tensors instead of internal KV cache + int external_kv_sequence_length = 0; // Sequence length of external KV tensors }; // Parameters deduced from node attributes and inputs/outputs. diff --git a/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h b/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h index adc7b623ec8c4..8d2468bb7b21e 100644 --- a/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h +++ b/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h @@ -104,6 +104,10 @@ class GQAAttentionBase { bool past_present_share_buffer = past_key_data == present_key_data && past_value_data == present_value_data; + // External KV mode: K and V are nullptr, past_key/past_value contain the external KV data. + // Skip KV cache concatenation and use external KV directly. + const bool use_external_kv = parameters.use_external_kv; + const T* k = packed_qkv ? Q + num_heads_ * sequence_length * head_size : K; T* output_qk_buffer = output_qk != nullptr ? output_qk->MutableData() : nullptr; @@ -112,7 +116,7 @@ class GQAAttentionBase { ComputeAttentionProbs(static_cast(attention_probs), Q, k, head_sink, seqlens_k->Data(), attention_bias_data, batch_size, sequence_length, total_sequence_length, attention_bias_shape, seqlen_past_kv_cache, seqlen_present_kv_cache, head_size, past_key_data, present_key_data, output_qk_buffer, - past_present_share_buffer, packed_qkv, is_prompt, tp, allocator); + past_present_share_buffer, packed_qkv, is_prompt, use_external_kv, tp, allocator); // Compute the attentionScore * Value: out(B, N, S, H_v) = attention_probs(B, N, S, T) x V(B, N, T, H_v) const T* v = packed_qkv ? Q + (num_heads_ + kv_num_heads_) * sequence_length * head_size : V; @@ -120,12 +124,12 @@ class GQAAttentionBase { seqlens_k->Data(), batch_size, sequence_length, seqlen_past_kv_cache, seqlen_present_kv_cache, head_size, hidden_size, past_value_data, present_value_data, past_present_share_buffer, packed_qkv, - is_prompt, tp, allocator); + is_prompt, use_external_kv, tp, allocator); } else { ComputeAttentionProbs(static_cast(attention_probs), Q, k, head_sink, seqlens_k->Data(), attention_bias_data, batch_size, sequence_length, total_sequence_length, attention_bias_shape, seqlen_past_kv_cache, seqlen_present_kv_cache, head_size, past_key_data, present_key_data, output_qk_buffer, - past_present_share_buffer, packed_qkv, is_prompt, tp, allocator); + past_present_share_buffer, packed_qkv, is_prompt, use_external_kv, tp, allocator); // Compute the attentionScore * Value: out(B, N, S, H_v) = attention_probs(B, N, S, T) x V(B, N, T, H_v) const T* v = packed_qkv ? Q + (num_heads_ + kv_num_heads_) * sequence_length * head_size : V; @@ -133,7 +137,7 @@ class GQAAttentionBase { seqlens_k->Data(), batch_size, sequence_length, seqlen_past_kv_cache, seqlen_present_kv_cache, head_size, hidden_size, past_value_data, present_value_data, past_present_share_buffer, packed_qkv, - is_prompt, tp, allocator); + is_prompt, use_external_kv, tp, allocator); } return Status::OK(); @@ -164,6 +168,7 @@ class GQAAttentionBase { const bool past_present_share_buffer, // whether present key and value share the same buffer const bool packed_qkv, // whether Q, K, V are packed const bool is_prompt, // whether it is prompt + const bool use_external_kv, // whether using external KV (skip KV concat) ThreadPool* tp, // thread pool AllocatorPtr allocator) const { // allocator for temporary buffer const ptrdiff_t packed_batch_stride = @@ -237,12 +242,21 @@ class GQAAttentionBase { } const T* k; - if (packed_qkv) { + if (use_external_kv) { + // External KV mode: use past_key directly (it holds the external KV data in BNSH format). + // No new K to concatenate — the external KV is the complete key sequence. + k = past_key + (i / kv_num_heads_factor) * present_buff_chunk_length; + // Also copy to present for output pass-through + if (present_key != nullptr && !past_present_share_buffer) { + memcpy(present_key + (i / kv_num_heads_factor) * present_buff_chunk_length, + k, SafeInt(total_seqlen) * head_size * sizeof(T)); + } + } else if (packed_qkv) { k = K + packed_batch_stride * batch_index + kv_input_chunk_length * (head_index / kv_num_heads_factor); } else { k = K + kv_input_chunk_length * (i / kv_num_heads_factor); } - if (nullptr != present_key) { + if (!use_external_kv && nullptr != present_key) { k = ConcatStateChunkGQA(past_key, k, present_key, present_buff_chunk_length, past_buff_chunk_length, past_chunk_length, kv_input_chunk_length, past_present_share_buffer, i / kv_num_heads_factor); @@ -392,6 +406,7 @@ class GQAAttentionBase { const bool past_present_share_buffer, // whether present key and value share the same buffer const bool packed_qkv, // whether Q, K, V are packed const bool is_prompt, // whether it is prompt + const bool use_external_kv, // whether using external KV (skip KV concat) ThreadPool* tp, AllocatorPtr allocator) const { const ptrdiff_t packed_batch_stride = @@ -445,12 +460,20 @@ class GQAAttentionBase { const size_t past_chunk_length = SafeInt(past_seqlen) * head_size; const T* v; - if (packed_qkv) { + if (use_external_kv) { + // External KV mode: use past_value directly (it holds the external KV data in BNSH format). + v = past_value + (i / kv_num_heads_factor) * present_buff_chunk_length; + // Copy to present for output pass-through + if (present_value != nullptr && !past_present_share_buffer) { + memcpy(present_value + (i / kv_num_heads_factor) * present_buff_chunk_length, + v, SafeInt(total_seqlen) * head_size * sizeof(T)); + } + } else if (packed_qkv) { v = V + packed_batch_stride * batch_index + kv_input_chunk_length * (head_index / kv_num_heads_factor); } else { v = V + kv_input_chunk_length * (i / kv_num_heads_factor); } - if (nullptr != present_value) { + if (!use_external_kv && nullptr != present_value) { v = ConcatStateChunkGQA(past_value, v, present_value, present_buff_chunk_length, past_buff_chunk_length, past_chunk_length, kv_input_chunk_length, past_present_share_buffer, i / kv_num_heads_factor); diff --git a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc index 5698bcb659f20..7513c64aeeac0 100644 --- a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc @@ -55,6 +55,8 @@ Status GroupQueryAttention::Compute(OpKernelContext* context) const { const Tensor* position_ids = context->Input(9); const Tensor* attention_bias = context->Input(10); const Tensor* head_sink = context->Input(11); + const Tensor* external_key = context->Input(14); + const Tensor* external_value = context->Input(15); GroupQueryAttentionParameters parameters = {}; ORT_RETURN_IF_ERROR(group_query_attention_helper::CheckInputs(query, @@ -71,13 +73,16 @@ Status GroupQueryAttention::Compute(OpKernelContext* context) const { total_seqlen_tensor, scale_, softcap_, - 0)); + 0, + external_key != nullptr)); ORT_RETURN_IF_ERROR(group_query_attention_helper::CheckCustomAttentionInputs(position_ids, attention_bias, head_sink, parameters)); + ORT_RETURN_IF_ERROR(group_query_attention_helper::CheckAndSetExternalKV(external_key, external_value, parameters)); + const int batch_size = parameters.batch_size; const int sequence_length = parameters.sequence_length; const int present_kv_seqlen = parameters.seqlen_present_kv_cache; @@ -125,7 +130,11 @@ Status GroupQueryAttention::Compute(OpKernelContext* context) const { OrtValue Q; OrtValue K; OrtValue V; - if (packed_qkv) { + if (parameters.use_external_kv) { + // External KV mode: only Q needs transposing. K,V come from external tensors. + ORT_RETURN_IF_ERROR(MaybeTransposeToBNSH( + allocator, batch_size, num_heads_, sequence_length, head_size, query, Q)); + } else if (packed_qkv) { ORT_RETURN_IF_ERROR(MaybeTransposeToBNSH( allocator, batch_size, num_heads_ + 2 * kv_num_heads_, sequence_length, head_size, query, Q)); } else { @@ -141,8 +150,8 @@ Status GroupQueryAttention::Compute(OpKernelContext* context) const { OrtValue RotaryQ; OrtValue RotaryK; T* q_rotary = Q.GetMutable()->MutableData(); - T* k_rotary = packed_qkv ? nullptr : K.GetMutable()->MutableData(); - if (do_rotary_) { + T* k_rotary = (packed_qkv || parameters.use_external_kv) ? nullptr : K.GetMutable()->MutableData(); + if (do_rotary_ && !parameters.use_external_kv) { ORT_ENFORCE(cos_cache != nullptr && sin_cache != nullptr, "cos_cache and sin_cache must be provided when do_rotary is true"); // Initialize rotary parameters rotary_embedding_helper::RotaryParameters rotary_params = {}; @@ -232,9 +241,16 @@ Status GroupQueryAttention::Compute(OpKernelContext* context) const { const T* head_sink_data = (head_sink != nullptr) ? head_sink->Data() : nullptr; + // When external KV is provided, use it in place of past_key/past_value for attention computation. + // External KV is pre-computed from another layer (KV-shared layers, e.g., Gemma4). + const Tensor* effective_past_key = parameters.use_external_kv ? external_key : past_key; + const Tensor* effective_past_value = parameters.use_external_kv ? external_value : past_value; + // Compute the attention score and apply the score to V - return ApplyAttention(q_rotary, packed_qkv ? nullptr : k_rotary, packed_qkv ? nullptr : V.Get().Data(), - head_sink_data, attention_bias, past_key, past_value, output, present_k, present_v, + const T* k_data = (packed_qkv || parameters.use_external_kv) ? nullptr : k_rotary; + const T* v_data = (packed_qkv || parameters.use_external_kv) ? nullptr : V.Get().Data(); + return ApplyAttention(q_rotary, k_data, v_data, + head_sink_data, attention_bias, effective_past_key, effective_past_value, output, present_k, present_v, output_qk, seqlens_k, parameters, allocator, context); } } // namespace contrib diff --git a/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h b/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h index f5399e307fbca..f730f5dbdd82d 100644 --- a/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h +++ b/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h @@ -96,6 +96,27 @@ Status Check_QKV(const T* packed_qkv, const T* value, const int num_heads, const return Status::OK(); } +template +Status Check_Q_Only(const T* query, const int num_heads, const int kv_num_heads, + int& batch_size, int& sequence_length, int& q_hidden_size, int& kv_hidden_size, int& head_size) { + const auto& query_dims = query->Shape().GetDims(); + if (query_dims.size() != 3) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'query' is expected to have 3 dimensions, got ", + query_dims.size()); + } + batch_size = static_cast(query_dims[0]); + sequence_length = static_cast(query_dims[1]); + q_hidden_size = static_cast(query_dims[2]); + head_size = static_cast(q_hidden_size) / num_heads; + if (head_size % 8 != 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "head_size must be a multiple of 8. Got head_size % 8 == ", + head_size % 8); + } + kv_hidden_size = head_size * kv_num_heads; + return Status::OK(); +} + template Status CheckPast(const T* past_key, const T* past_value, int batch_size, int kv_num_heads, int head_size, int kv_cache_bit_width, int& past_sequence_length) { @@ -157,6 +178,56 @@ Status CheckPast(const T* past_key, const T* past_value, int batch_size, int kv_ return Status::OK(); } +template +Status CheckExternalKV(const T* external_key, const T* external_value, int batch_size, int kv_num_heads, + int& external_sequence_length) { + if (external_key == nullptr || external_value == nullptr) { + if (external_key != nullptr || external_value != nullptr) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'external_key' and 'external_value' shall be both present or both absent."); + } + return Status::OK(); + } + + const auto& ext_key_dims = external_key->Shape().GetDims(); + const auto& ext_value_dims = external_value->Shape().GetDims(); + + if (ext_key_dims.size() != 4) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'external_key' is expected to have 4 dimensions (BNSH), got ", + ext_key_dims.size()); + } + if (ext_value_dims.size() != 4) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'external_value' is expected to have 4 dimensions (BNSH), got ", + ext_value_dims.size()); + } + if (ext_key_dims[0] != batch_size) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'external_key' dimension 0 should be batch_size, got ", ext_key_dims[0]); + } + if (ext_value_dims[0] != batch_size) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'external_value' dimension 0 should be batch_size, got ", ext_value_dims[0]); + } + if (ext_key_dims[1] != kv_num_heads) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'external_key' shall have kv_num_heads, got ", ext_key_dims[1]); + } + if (ext_value_dims[1] != kv_num_heads) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'external_value' shall have kv_num_heads, got ", ext_value_dims[1]); + } + if (ext_key_dims[2] != ext_value_dims[2]) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'external_key' and 'external_value' should have same sequence length dimension."); + } + // Note: head_size validation is relaxed here — external KV may have different head_size + // than the query (e.g., Gemma4 global layers with head_dim=512 vs local head_dim=256). + external_sequence_length = static_cast(ext_key_dims[2]); + return Status::OK(); +} + template Status CheckRotaryCaches(const T* cos_cache, const T* sin_cache, int head_size, int total_sequence_length, int& rotary_dim) { @@ -207,7 +278,8 @@ Status CheckInputs(const T* query, const T* total_seqlen, float scale, float softcap, - int kv_cache_bit_width) { + int kv_cache_bit_width, + bool has_external_kv = false) { // Note: Here S* is seqlen_past_kv_cache, S+ is seqlen_present_kv_cache // past_key : (B, N_k, S*, H) or (B, N_k, S+, H) or nullptr // past_value : (B, N_k, S*, H) or (B, N_k, S+, H) or nullptr @@ -242,8 +314,14 @@ Status CheckInputs(const T* query, int q_hidden_size = 0; int kv_hidden_size = 0; int head_size = 0; - const bool is_packed_qkv = key == nullptr; - if (!is_packed_qkv) { + // When external KV is provided, key/value can be nullptr without implying packed QKV. + // In this mode, query contains only Q (not packed QKV). + const bool is_packed_qkv = (key == nullptr) && !has_external_kv; + if (has_external_kv && key == nullptr) { + // Q-only mode: query is just Q, K and V come from external tensors + ORT_RETURN_IF_ERROR(Check_Q_Only(query, num_heads, kv_num_heads, batch_size, sequence_length, + q_hidden_size, kv_hidden_size, head_size)); + } else if (!is_packed_qkv) { ORT_RETURN_IF_ERROR(Check_Q_K_V(query, key, value, num_heads, kv_num_heads, batch_size, sequence_length, q_hidden_size, kv_hidden_size, head_size)); } else { @@ -350,12 +428,13 @@ Status CheckInputs(const T* query, float scale, float softcap, int kv_cache_bit_width, - int max_threads_per_block) { + int max_threads_per_block, + bool has_external_kv = false) { if (max_threads_per_block > 0 && num_heads > max_threads_per_block) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "num_heads should be no larger than ", max_threads_per_block); } - return CheckInputs(query, key, value, past_key, past_value, cos_cache, sin_cache, parameters, num_heads, kv_num_heads, seqlens_k, total_seqlen, scale, softcap, kv_cache_bit_width); + return CheckInputs(query, key, value, past_key, past_value, cos_cache, sin_cache, parameters, num_heads, kv_num_heads, seqlens_k, total_seqlen, scale, softcap, kv_cache_bit_width, has_external_kv); } template @@ -445,6 +524,31 @@ inline Status CheckNoQKOutput(int num_outputs, int qk_output) { return Status::OK(); } +// Validate and configure external KV inputs for KV-shared layers. +// Call this after CheckInputs to set up external KV parameters. +template +Status CheckAndSetExternalKV(const T* external_key, const T* external_value, + GroupQueryAttentionParameters& parameters) { + if (external_key == nullptr && external_value == nullptr) { + return Status::OK(); + } + + int external_sequence_length = 0; + ORT_RETURN_IF_ERROR(CheckExternalKV(external_key, external_value, + parameters.batch_size, parameters.kv_num_heads, + external_sequence_length)); + + parameters.use_external_kv = true; + parameters.external_kv_sequence_length = external_sequence_length; + + // When using external KV, the total sequence length for attention is determined + // by the external KV tensor, not the internal KV cache. + parameters.total_sequence_length = external_sequence_length; + parameters.seqlen_present_kv_cache = external_sequence_length; + + return Status::OK(); +} + } // namespace group_query_attention_helper } // namespace contrib } // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/attention_data.h b/onnxruntime/contrib_ops/cuda/bert/attention_data.h index 486bf05bd86d5..a95e9b818ca81 100644 --- a/onnxruntime/contrib_ops/cuda/bert/attention_data.h +++ b/onnxruntime/contrib_ops/cuda/bert/attention_data.h @@ -192,6 +192,10 @@ struct GroupQueryAttentionData { U* present_key = nullptr; U* present_value = nullptr; + // External KV for KV-shared layers (e.g., Gemma4) + const U* external_key = nullptr; + const U* external_value = nullptr; + // Kernel Flags bool use_flash_attention = false; bool use_memory_efficient_attention = false; diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc index 3b6b5f9079ebe..f8de4d44261f5 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc @@ -164,6 +164,8 @@ Status GroupQueryAttention::ComputeInternal(OpKernelContext* context) cons const Tensor* head_sink = context->Input(11); const Tensor* k_scale = context->Input(12); const Tensor* v_scale = context->Input(13); + const Tensor* external_key = context->Input(14); + const Tensor* external_value = context->Input(15); if (k_quant_type_ != KVQuantizationType::NONE) { if (k_scale == nullptr) { @@ -219,12 +221,16 @@ Status GroupQueryAttention::ComputeInternal(OpKernelContext* context) cons scale_, softcap_, kv_cache_bit_width_, - device_prop.maxThreadsPerBlock)); + device_prop.maxThreadsPerBlock, + external_key != nullptr)); ORT_RETURN_IF_ERROR(group_query_attention_helper::CheckCustomAttentionInputs(position_ids, attention_bias, head_sink, parameters)); + + ORT_RETURN_IF_ERROR(group_query_attention_helper::CheckAndSetExternalKV(external_key, external_value, parameters)); + parameters.local_window_size = local_window_size_; parameters.is_unidirectional = is_unidirectional_; parameters.use_smooth_softmax = use_smooth_softmax_ || head_sink != nullptr; @@ -236,6 +242,14 @@ Status GroupQueryAttention::ComputeInternal(OpKernelContext* context) cons parameters.do_rotary = do_rotary_; parameters.rotary_interleaved = rotary_interleaved_; + // When using external KV, disable rotary embedding — the external KV already has RoPE applied + // from the source layer, and the caller is expected to pre-apply RoPE to Q. + if (parameters.use_external_kv && parameters.do_rotary) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "do_rotary must be 0 when using external_key/external_value. " + "Pre-apply RoPE to Q and use already-rotated K from the source layer."); + } + // The current GQA CUDA implementation will never be able to have a QK output. // GQA CUDA uses either flash attention or memory efficient attention. Neither kernel supports returning the QK output. ORT_RETURN_IF_ERROR(group_query_attention_helper::CheckNoQKOutput( @@ -287,15 +301,26 @@ Status GroupQueryAttention::ComputeInternal(OpKernelContext* context) cons data.k_scale = k_scale == nullptr ? nullptr : reinterpret_cast(k_scale->DataRaw()); data.v_scale = v_scale == nullptr ? nullptr : reinterpret_cast(v_scale->DataRaw()); - data.past_key = (past_key == nullptr) ? nullptr : reinterpret_cast(past_key->Data()); - data.past_value = (past_value == nullptr) ? nullptr : reinterpret_cast(past_value->Data()); - - data.present_key = reinterpret_cast(present_key_output->MutableData()); - data.present_value = reinterpret_cast(present_value_output->MutableData()); - - // Compute past_present_share_buffer early since it's needed for flash attention path selection. - // This compares the final pointer values after quantization handling. - parameters.past_present_share_buffer = (data.past_key == data.present_key); + if (parameters.use_external_kv) { + // External KV mode: use external tensors as the KV source for attention. + // The external KV is treated as "past" KV since it's already computed. + // No KV cache update is performed — the present outputs are copies/views of external KV. + data.external_key = reinterpret_cast(external_key->Data()); + data.external_value = reinterpret_cast(external_value->Data()); + data.past_key = data.external_key; + data.past_value = data.external_value; + data.present_key = reinterpret_cast(present_key_output->MutableData()); + data.present_value = reinterpret_cast(present_value_output->MutableData()); + // Mark as shared buffer so the kernel treats external KV as already-populated cache + parameters.past_present_share_buffer = false; + } else { + data.past_key = (past_key == nullptr) ? nullptr : reinterpret_cast(past_key->Data()); + data.past_value = (past_value == nullptr) ? nullptr : reinterpret_cast(past_value->Data()); + data.present_key = reinterpret_cast(present_key_output->MutableData()); + data.present_value = reinterpret_cast(present_value_output->MutableData()); + // Compute past_present_share_buffer early since it's needed for flash attention path selection. + parameters.past_present_share_buffer = (data.past_key == data.present_key); + } bool is_inputs_quantized = (k_quant_type_ != KVQuantizationType::NONE) || (v_quant_type_ != KVQuantizationType::NONE); constexpr bool is_int8 = std::is_same::value; @@ -519,7 +544,9 @@ Status GroupQueryAttention::ComputeInternal(OpKernelContext* context) cons } // Validate past_value pointer consistency (past_present_share_buffer was computed early after pointer setup) - if (parameters.past_present_share_buffer) { + if (parameters.use_external_kv) { + // External KV mode: past and present are separate (external source -> present output) + } else if (parameters.past_present_share_buffer) { ORT_ENFORCE(data.past_value == data.present_value, "past_value and present_value must be the same tensor when past_present_share_buffer is true"); } else { ORT_ENFORCE(data.past_value != data.present_value, "past_value and present_value must be different tensors when past_present_share_buffer is false"); diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu index c617de747ccf7..d1a4bde9fe20d 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu @@ -97,6 +97,23 @@ Status PrepareQKV( q_out = nullptr; } + // External KV mode: the external KV data is already set as past_key/past_value. + // Copy it into the present buffers and skip the KV append/RoPE logic. + if (parameters.use_external_kv) { + U* k = reinterpret_cast(data.present_key); + U* v = reinterpret_cast(data.present_value); + int external_seq_len = parameters.external_kv_sequence_length; + + // Copy external KV into present buffers + size_t kv_copy_size = (size_t)batch_size * kv_num_heads * external_seq_len * head_size * sizeof(U); + CUDA_CALL_THROW(cudaMemcpyAsync(k, data.past_key, kv_copy_size, cudaMemcpyDeviceToDevice, stream)); + CUDA_CALL_THROW(cudaMemcpyAsync(v, data.past_value, kv_copy_size, cudaMemcpyDeviceToDevice, stream)); + + // Q is used directly from the input + q = reinterpret_cast(data.query); + return Status::OK(); + } + U* k = reinterpret_cast(data.present_key); U* v = reinterpret_cast(data.present_value); int max_cache_length = parameters.seqlen_present_kv_cache; diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index 1209446c6a367..9b13cc170a27d 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -1314,6 +1314,21 @@ ONNX_MS_OPERATOR_SET_SCHEMA( OpSchema::Optional) .Input(12, "k_scale", "Scale tensor for past_key.", "T_KV_SCALE", OpSchema::Optional) .Input(13, "v_scale", "Scale tensor for past_value.", "T_KV_SCALE", OpSchema::Optional) + .Input(14, + "external_key", + "External pre-computed key tensor in BNSH format (batch_size, kv_num_heads, external_seq_len, head_size). " + "Used for KV-shared layers that borrow K,V from another layer's present KV output. " + "When provided, the operator skips its internal KV cache update and uses this tensor directly " + "for attention computation. RoPE is not applied to external keys (assumed already applied).", + "T_CACHE", + OpSchema::Optional) + .Input(15, + "external_value", + "External pre-computed value tensor in BNSH format (batch_size, kv_num_heads, external_seq_len, head_size). " + "Must be provided together with external_key. When provided, the operator uses this tensor " + "for attention computation instead of the internal KV cache.", + "T_CACHE", + OpSchema::Optional) .Output(0, "output", "3D output tensor with shape (batch_size, sequence_length, hidden_size)", From ec041dbb8294f095374871b22262d5489a4442bd Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 27 Apr 2026 10:44:38 -0700 Subject: [PATCH 02/32] Fix the op --- onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h | 8 ++++---- onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc | 7 +++++++ .../contrib_ops/cpu/bert/group_query_attention_helper.h | 7 +++++++ 3 files changed, 18 insertions(+), 4 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h b/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h index 8d2468bb7b21e..f2ea775dd78ff 100644 --- a/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h +++ b/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h @@ -246,8 +246,8 @@ class GQAAttentionBase { // External KV mode: use past_key directly (it holds the external KV data in BNSH format). // No new K to concatenate — the external KV is the complete key sequence. k = past_key + (i / kv_num_heads_factor) * present_buff_chunk_length; - // Also copy to present for output pass-through - if (present_key != nullptr && !past_present_share_buffer) { + // Copy to present for output pass-through (once per KV head, not per Q head) + if (present_key != nullptr && !past_present_share_buffer && head_index % kv_num_heads_factor == 0) { memcpy(present_key + (i / kv_num_heads_factor) * present_buff_chunk_length, k, SafeInt(total_seqlen) * head_size * sizeof(T)); } @@ -463,8 +463,8 @@ class GQAAttentionBase { if (use_external_kv) { // External KV mode: use past_value directly (it holds the external KV data in BNSH format). v = past_value + (i / kv_num_heads_factor) * present_buff_chunk_length; - // Copy to present for output pass-through - if (present_value != nullptr && !past_present_share_buffer) { + // Copy to present for output pass-through (once per KV head, not per Q head) + if (present_value != nullptr && !past_present_share_buffer && head_index % kv_num_heads_factor == 0) { memcpy(present_value + (i / kv_num_heads_factor) * present_buff_chunk_length, v, SafeInt(total_seqlen) * head_size * sizeof(T)); } diff --git a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc index 7513c64aeeac0..6cba0c32757b5 100644 --- a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc @@ -83,6 +83,13 @@ Status GroupQueryAttention::Compute(OpKernelContext* context) const { ORT_RETURN_IF_ERROR(group_query_attention_helper::CheckAndSetExternalKV(external_key, external_value, parameters)); + // External KV mode requires do_rotary=0 — K already has RoPE from the source layer + if (parameters.use_external_kv && do_rotary_) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "do_rotary must be 0 when using external_key/external_value. " + "Pre-apply RoPE to Q and use already-rotated K from the source layer."); + } + const int batch_size = parameters.batch_size; const int sequence_length = parameters.sequence_length; const int present_kv_seqlen = parameters.seqlen_present_kv_cache; diff --git a/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h b/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h index f730f5dbdd82d..7b72cbe3b9804 100644 --- a/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h +++ b/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h @@ -543,6 +543,13 @@ Status CheckAndSetExternalKV(const T* external_key, const T* external_value, // When using external KV, the total sequence length for attention is determined // by the external KV tensor, not the internal KV cache. + // Validate that the original total_sequence_length doesn't exceed external KV length. + if (parameters.total_sequence_length > external_sequence_length) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "total_sequence_length (", parameters.total_sequence_length, + ") exceeds external KV sequence length (", external_sequence_length, + "). Ensure seqlens_k is consistent with the external KV tensor size."); + } parameters.total_sequence_length = external_sequence_length; parameters.seqlen_present_kv_cache = external_sequence_length; From 7139053b63617de2d4c5716de604c9f0223f3130 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 27 Apr 2026 11:02:58 -0700 Subject: [PATCH 03/32] Fix lint error --- onnxruntime/contrib_ops/cpu/bert/attention_parameters.h | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/bert/attention_parameters.h b/onnxruntime/contrib_ops/cpu/bert/attention_parameters.h index 101ec88df375c..08345727677a6 100644 --- a/onnxruntime/contrib_ops/cpu/bert/attention_parameters.h +++ b/onnxruntime/contrib_ops/cpu/bert/attention_parameters.h @@ -103,8 +103,8 @@ struct GroupQueryAttentionParameters : AttentionParameters { int kv_cache_bit_width = 0; // External KV parameters for KV-shared layers (e.g., Gemma4) - bool use_external_kv = false; // When true, use external K,V tensors instead of internal KV cache - int external_kv_sequence_length = 0; // Sequence length of external KV tensors + bool use_external_kv = false; // When true, use external K,V tensors instead of internal KV cache + int external_kv_sequence_length = 0; // Sequence length of external KV tensors }; // Parameters deduced from node attributes and inputs/outputs. From 0bf639413032c0af0b4ea0b5da8b2e605f4146f1 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 27 Apr 2026 14:03:45 -0700 Subject: [PATCH 04/32] Fix webgpu build --- onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h b/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h index 7b72cbe3b9804..b94db4c9774e6 100644 --- a/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h +++ b/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h @@ -279,7 +279,7 @@ Status CheckInputs(const T* query, float scale, float softcap, int kv_cache_bit_width, - bool has_external_kv = false) { + bool has_external_kv) { // Note: Here S* is seqlen_past_kv_cache, S+ is seqlen_present_kv_cache // past_key : (B, N_k, S*, H) or (B, N_k, S+, H) or nullptr // past_value : (B, N_k, S*, H) or (B, N_k, S+, H) or nullptr From 3c07e1dd8d4ea4bc678b5faf00348db5b4c93220 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 27 Apr 2026 14:32:09 -0700 Subject: [PATCH 05/32] fix webgpu build --- onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc b/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc index fd72f751ee810..98f09a496f579 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc @@ -212,7 +212,8 @@ Status GroupQueryAttention::ComputeInternal(onnxruntime::webgpu::ComputeContext& scale_, softcap_, 0, - context.DeviceLimits().maxComputeInvocationsPerWorkgroup)); + static_cast(context.DeviceLimits().maxComputeInvocationsPerWorkgroup), + /*has_external_kv=*/false)); params.use_smooth_softmax = use_smooth_softmax_; params.rotary_interleaved = rotary_interleaved_; From f9051967780107f35ef999b5104762936b35f6df Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 27 Apr 2026 14:58:36 -0700 Subject: [PATCH 06/32] Address copilot comments --- .../cpu/bert/group_query_attention.cc | 7 + .../cpu/bert/group_query_attention_helper.h | 44 +++++- .../cuda/bert/group_query_attention.cc | 7 + .../cuda/bert/group_query_attention_impl.cu | 6 +- .../core/graph/contrib_ops/bert_defs.cc | 9 +- .../group_query_attention_op_test.cc | 127 ++++++++++++++++++ 6 files changed, 193 insertions(+), 7 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc index 6cba0c32757b5..d126e68f69edd 100644 --- a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc @@ -83,6 +83,13 @@ Status GroupQueryAttention::Compute(OpKernelContext* context) const { ORT_RETURN_IF_ERROR(group_query_attention_helper::CheckAndSetExternalKV(external_key, external_value, parameters)); + // External KV is mutually exclusive with provided key/value inputs + if (parameters.use_external_kv && (key != nullptr || value != nullptr)) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "key and value (inputs 1/2) must not be provided when using external_key/external_value. " + "External KV replaces the K,V projections entirely."); + } + // External KV mode requires do_rotary=0 — K already has RoPE from the source layer if (parameters.use_external_kv && do_rotary_) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, diff --git a/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h b/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h index b94db4c9774e6..0a458c3a8e124 100644 --- a/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h +++ b/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h @@ -107,7 +107,18 @@ Status Check_Q_Only(const T* query, const int num_heads, const int kv_num_heads, batch_size = static_cast(query_dims[0]); sequence_length = static_cast(query_dims[1]); q_hidden_size = static_cast(query_dims[2]); - head_size = static_cast(q_hidden_size) / num_heads; + if (q_hidden_size % num_heads != 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "q_hidden_size (", q_hidden_size, ") must be divisible by num_heads (", num_heads, + ") in Q-only mode (external KV). Got q_hidden_size % num_heads == ", + q_hidden_size % num_heads); + } + head_size = q_hidden_size / num_heads; + if (head_size == 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "head_size must be > 0. Got q_hidden_size=", q_hidden_size, + ", num_heads=", num_heads); + } if (head_size % 8 != 0) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "head_size must be a multiple of 8. Got head_size % 8 == ", @@ -180,7 +191,7 @@ Status CheckPast(const T* past_key, const T* past_value, int batch_size, int kv_ template Status CheckExternalKV(const T* external_key, const T* external_value, int batch_size, int kv_num_heads, - int& external_sequence_length) { + int head_size, int kv_cache_bit_width, int& external_sequence_length) { if (external_key == nullptr || external_value == nullptr) { if (external_key != nullptr || external_value != nullptr) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, @@ -222,8 +233,18 @@ Status CheckExternalKV(const T* external_key, const T* external_value, int batch return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'external_key' and 'external_value' should have same sequence length dimension."); } - // Note: head_size validation is relaxed here — external KV may have different head_size - // than the query (e.g., Gemma4 global layers with head_dim=512 vs local head_dim=256). + // Validate head dimension (dim 3). For 4-bit quantized KV cache, the stored dimension is head_size / 2. + int expected_head_dim = (kv_cache_bit_width == 4) ? (head_size / 2) : head_size; + if (ext_key_dims[3] != expected_head_dim) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'external_key' dimension 3 should be head_size (", expected_head_dim, + "), got ", ext_key_dims[3]); + } + if (ext_value_dims[3] != expected_head_dim) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'external_value' dimension 3 should be head_size (", expected_head_dim, + "), got ", ext_value_dims[3]); + } external_sequence_length = static_cast(ext_key_dims[2]); return Status::OK(); } @@ -536,11 +557,26 @@ Status CheckAndSetExternalKV(const T* external_key, const T* external_value, int external_sequence_length = 0; ORT_RETURN_IF_ERROR(CheckExternalKV(external_key, external_value, parameters.batch_size, parameters.kv_num_heads, + parameters.head_size, parameters.kv_cache_bit_width, external_sequence_length)); parameters.use_external_kv = true; parameters.external_kv_sequence_length = external_sequence_length; + // External KV mode is incompatible with packed QKV — query must contain only Q. + if (parameters.is_packed_qkv) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "external_key/external_value cannot be used with packed QKV input. " + "Provide query as Q-only (without K,V) when using external KV."); + } + + // External KV replaces the internal KV cache — past_key/past_value should not be provided. + if (parameters.seqlen_past_kv_cache > 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "past_key/past_value should not be provided when using external_key/external_value. " + "External KV replaces the internal KV cache entirely."); + } + // When using external KV, the total sequence length for attention is determined // by the external KV tensor, not the internal KV cache. // Validate that the original total_sequence_length doesn't exceed external KV length. diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc index f3fdb27e09b10..2562e72451bd8 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc @@ -233,6 +233,13 @@ Status GroupQueryAttention::ComputeInternal(OpKernelContext* context) cons ORT_RETURN_IF_ERROR(group_query_attention_helper::CheckAndSetExternalKV(external_key, external_value, parameters)); + // External KV is mutually exclusive with provided key/value inputs + if (parameters.use_external_kv && (key != nullptr || value != nullptr)) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "key and value (inputs 1/2) must not be provided when using external_key/external_value. " + "External KV replaces the K,V projections entirely."); + } + parameters.local_window_size = local_window_size_; parameters.is_unidirectional = is_unidirectional_; parameters.use_smooth_softmax = use_smooth_softmax_ || head_sink != nullptr; diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu index e8a222dcc4396..4f0da8fda4d08 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu @@ -105,8 +105,10 @@ Status PrepareQKV( U* v = reinterpret_cast(data.present_value); int external_seq_len = parameters.external_kv_sequence_length; - // Copy external KV into present buffers - size_t kv_copy_size = (size_t)batch_size * kv_num_heads * external_seq_len * head_size * sizeof(U); + // For 4-bit quantized KV cache, the stored head dimension is head_size/2 (two nibbles per byte). + // Use the packed dimension to compute the correct copy size. + int cache_head_dim = (parameters.kv_cache_bit_width == 4) ? (head_size + 1) / 2 : head_size; + size_t kv_copy_size = (size_t)batch_size * kv_num_heads * external_seq_len * cache_head_dim * sizeof(U); CUDA_CALL_THROW(cudaMemcpyAsync(k, data.past_key, kv_copy_size, cudaMemcpyDeviceToDevice, stream)); CUDA_CALL_THROW(cudaMemcpyAsync(v, data.past_value, kv_copy_size, cudaMemcpyDeviceToDevice, stream)); diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index 9b13cc170a27d..ccecc28ec6afc 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -240,13 +240,20 @@ void BaseGroupQueryAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceConte if (ctx.getNumOutputs() >= 3) { // has present output const auto* past_key_type = ctx.getInputType(past_key_index); + // external_key is at input index 14 for GroupQueryAttention + const auto* external_key_type = (ctx.getNumInputs() > 14) ? ctx.getInputType(14) : nullptr; if (past_key_type != nullptr) { // present_key and present_value have the same type as past_key/past_value. // This allows them to be int8 or packed uint8 when quantization is enabled. ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, past_key_index, 1); // present_key ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, past_key_index + 1, 2); // present_value + } else if (external_key_type != nullptr) { + // When external KV is provided (inputs 14/15), present outputs should match + // the external KV type (T_CACHE), not the query type (T). + ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 14, 1); // present_key from external_key + ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 15, 2); // present_value from external_value } else { - // If no past state, present is the same type as query. + // If no past state and no external KV, present is the same type as query. ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 0, 1); ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 0, 2); } diff --git a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc index 0690094031bb8..593fc55c01e99 100644 --- a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc @@ -307,5 +307,132 @@ TEST(GroupQueryAttentionTest, SeqlensKWrongLength) { {}, nullptr, &execution_providers); } +// ============================================================================ +// External KV tests (inputs 14/15: external_key, external_value) +// ============================================================================ + +// Helper for external KV tests +static void RunGQAExternalKVTest( + int external_seq_len, + OpTester::ExpectResult expect, + const std::string& expected_message, + bool provide_key_value = false, + bool provide_past = false, + bool do_rotary = false) { + constexpr int batch_size = 1; + constexpr int sequence_length = 1; + constexpr int num_heads = 2; + constexpr int kv_num_heads = 1; + constexpr int head_size = 8; + constexpr int hidden_size = num_heads * head_size; + constexpr int kv_hidden_size = kv_num_heads * head_size; + + OpTester tester("GroupQueryAttention", 1, onnxruntime::kMSDomain); + tester.AddAttribute("num_heads", static_cast(num_heads)); + tester.AddAttribute("kv_num_heads", static_cast(kv_num_heads)); + if (do_rotary) { + tester.AddAttribute("do_rotary", 1); + } + + // Query (Q-only when using external KV) + std::vector query_data(batch_size * sequence_length * hidden_size, 1.0f); + tester.AddInput("query", {batch_size, sequence_length, hidden_size}, query_data); + + // Key/Value inputs (should be absent for external KV) + if (provide_key_value) { + std::vector key_data(batch_size * sequence_length * kv_hidden_size, 1.0f); + std::vector value_data(batch_size * sequence_length * kv_hidden_size, 1.0f); + tester.AddInput("key", {batch_size, sequence_length, kv_hidden_size}, key_data); + tester.AddInput("value", {batch_size, sequence_length, kv_hidden_size}, value_data); + } else { + tester.AddOptionalInputEdge(); // key + tester.AddOptionalInputEdge(); // value + } + + // Past key/value (should be absent for external KV) + if (provide_past) { + std::vector past_k(batch_size * kv_num_heads * 4 * head_size, 0.5f); + std::vector past_v(batch_size * kv_num_heads * 4 * head_size, 0.5f); + tester.AddInput("past_key", {batch_size, kv_num_heads, 4, head_size}, past_k); + tester.AddInput("past_value", {batch_size, kv_num_heads, 4, head_size}, past_v); + } else { + tester.AddOptionalInputEdge(); // past_key + tester.AddOptionalInputEdge(); // past_value + } + + // seqlens_k = external_seq_len - 1 (historical convention) + tester.AddInput("seqlens_k", {batch_size}, {static_cast(external_seq_len - 1)}); + tester.AddInput("total_sequence_length", {1}, {static_cast(external_seq_len)}); + + tester.AddOptionalInputEdge(); // cos_cache (7) + tester.AddOptionalInputEdge(); // sin_cache (8) + tester.AddOptionalInputEdge(); // position_ids (9) + tester.AddOptionalInputEdge(); // attention_bias (10) + tester.AddOptionalInputEdge(); // head_sink (11) + tester.AddOptionalInputEdge(); // k_scale (12) + tester.AddOptionalInputEdge(); // v_scale (13) + + // External key/value (inputs 14/15) — BNSH format + std::vector ext_key(batch_size * kv_num_heads * external_seq_len * head_size, 0.5f); + std::vector ext_value(batch_size * kv_num_heads * external_seq_len * head_size, 0.5f); + tester.AddInput("external_key", {batch_size, kv_num_heads, external_seq_len, head_size}, ext_key); + tester.AddInput("external_value", {batch_size, kv_num_heads, external_seq_len, head_size}, ext_value); + + // Outputs + int present_seq_len = std::max(1, external_seq_len); + tester.AddOutput("output", {batch_size, sequence_length, hidden_size}, + std::vector(batch_size * sequence_length * hidden_size, 0.0f)); + tester.AddOutput("present_key", {batch_size, kv_num_heads, present_seq_len, head_size}, + std::vector(batch_size * kv_num_heads * present_seq_len * head_size, 0.0f)); + tester.AddOutput("present_value", {batch_size, kv_num_heads, present_seq_len, head_size}, + std::vector(batch_size * kv_num_heads * present_seq_len * head_size, 0.0f)); + + if (expect == OpTester::ExpectResult::kExpectSuccess) { + tester.SetOutputTolerance(1e6f); + } + + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + tester.Run(expect, expected_message, {}, nullptr, &execution_providers); +} + +// Basic: external KV with Q-only query should succeed +TEST(GroupQueryAttentionTest, ExternalKV_BasicSuccess) { + RunGQAExternalKVTest( + /*external_seq_len=*/8, + OpTester::ExpectResult::kExpectSuccess, + ""); +} + +// Reject: external KV with key/value inputs provided (mutual exclusivity) +TEST(GroupQueryAttentionTest, ExternalKV_RejectsProvidedKeyValue) { + RunGQAExternalKVTest( + /*external_seq_len=*/8, + OpTester::ExpectResult::kExpectFailure, + "key and value (inputs 1/2) must not be provided", + /*provide_key_value=*/true); +} + +// Reject: external KV with past_key/past_value provided +TEST(GroupQueryAttentionTest, ExternalKV_RejectsPastKV) { + RunGQAExternalKVTest( + /*external_seq_len=*/8, + OpTester::ExpectResult::kExpectFailure, + "past_key/past_value should not be provided", + /*provide_key_value=*/false, + /*provide_past=*/true); +} + +// Reject: external KV with do_rotary=1 +TEST(GroupQueryAttentionTest, ExternalKV_RejectsDoRotary) { + RunGQAExternalKVTest( + /*external_seq_len=*/8, + OpTester::ExpectResult::kExpectFailure, + "do_rotary must be 0", + /*provide_key_value=*/false, + /*provide_past=*/false, + /*do_rotary=*/true); +} + } // namespace test } // namespace onnxruntime From d2ead38cfe01669faa157b78338f4ee58d9181c9 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Tue, 28 Apr 2026 00:59:51 -0700 Subject: [PATCH 07/32] Make GQA present_key/present_value outputs optional for KV-shared layers --- .../cpu/bert/attention_parameters.h | 4 - .../contrib_ops/cpu/bert/gqa_attention_base.h | 47 ++---- .../cpu/bert/group_query_attention.cc | 42 +---- .../cpu/bert/group_query_attention_helper.h | 157 +----------------- .../contrib_ops/cuda/bert/attention_data.h | 4 - .../cuda/bert/group_query_attention.cc | 60 ++----- .../cuda/bert/group_query_attention_impl.cu | 23 +-- .../webgpu/bert/group_query_attention.cc | 3 +- .../core/graph/contrib_ops/bert_defs.cc | 30 +--- .../group_query_attention_op_test.cc | 151 +++++++---------- 10 files changed, 113 insertions(+), 408 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/bert/attention_parameters.h b/onnxruntime/contrib_ops/cpu/bert/attention_parameters.h index 08345727677a6..f316a0dfdf91c 100644 --- a/onnxruntime/contrib_ops/cpu/bert/attention_parameters.h +++ b/onnxruntime/contrib_ops/cpu/bert/attention_parameters.h @@ -101,10 +101,6 @@ struct GroupQueryAttentionParameters : AttentionParameters { KVQuantizationType k_quant_type = KVQuantizationType::NONE; KVQuantizationType v_quant_type = KVQuantizationType::NONE; int kv_cache_bit_width = 0; - - // External KV parameters for KV-shared layers (e.g., Gemma4) - bool use_external_kv = false; // When true, use external K,V tensors instead of internal KV cache - int external_kv_sequence_length = 0; // Sequence length of external KV tensors }; // Parameters deduced from node attributes and inputs/outputs. diff --git a/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h b/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h index f2ea775dd78ff..1f03cf9f105a2 100644 --- a/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h +++ b/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h @@ -85,7 +85,9 @@ class GQAAttentionBase { if (past_key != nullptr && past_value != nullptr) { seqlen_past_kv_cache = static_cast(past_key->Shape().GetDims()[2]); } - int seqlen_present_kv_cache = static_cast(present_key->Shape().GetDims()[2]); + int seqlen_present_kv_cache = present_key != nullptr + ? static_cast(present_key->Shape().GetDims()[2]) + : parameters.seqlen_present_kv_cache; // Compute the attention score. bool gqa_mlas_supported = MlasGQASupported(CblasNoTrans, CblasTrans) && @@ -104,10 +106,6 @@ class GQAAttentionBase { bool past_present_share_buffer = past_key_data == present_key_data && past_value_data == present_value_data; - // External KV mode: K and V are nullptr, past_key/past_value contain the external KV data. - // Skip KV cache concatenation and use external KV directly. - const bool use_external_kv = parameters.use_external_kv; - const T* k = packed_qkv ? Q + num_heads_ * sequence_length * head_size : K; T* output_qk_buffer = output_qk != nullptr ? output_qk->MutableData() : nullptr; @@ -116,7 +114,7 @@ class GQAAttentionBase { ComputeAttentionProbs(static_cast(attention_probs), Q, k, head_sink, seqlens_k->Data(), attention_bias_data, batch_size, sequence_length, total_sequence_length, attention_bias_shape, seqlen_past_kv_cache, seqlen_present_kv_cache, head_size, past_key_data, present_key_data, output_qk_buffer, - past_present_share_buffer, packed_qkv, is_prompt, use_external_kv, tp, allocator); + past_present_share_buffer, packed_qkv, is_prompt, tp, allocator); // Compute the attentionScore * Value: out(B, N, S, H_v) = attention_probs(B, N, S, T) x V(B, N, T, H_v) const T* v = packed_qkv ? Q + (num_heads_ + kv_num_heads_) * sequence_length * head_size : V; @@ -124,12 +122,12 @@ class GQAAttentionBase { seqlens_k->Data(), batch_size, sequence_length, seqlen_past_kv_cache, seqlen_present_kv_cache, head_size, hidden_size, past_value_data, present_value_data, past_present_share_buffer, packed_qkv, - is_prompt, use_external_kv, tp, allocator); + is_prompt, tp, allocator); } else { ComputeAttentionProbs(static_cast(attention_probs), Q, k, head_sink, seqlens_k->Data(), attention_bias_data, batch_size, sequence_length, total_sequence_length, attention_bias_shape, seqlen_past_kv_cache, seqlen_present_kv_cache, head_size, past_key_data, present_key_data, output_qk_buffer, - past_present_share_buffer, packed_qkv, is_prompt, use_external_kv, tp, allocator); + past_present_share_buffer, packed_qkv, is_prompt, tp, allocator); // Compute the attentionScore * Value: out(B, N, S, H_v) = attention_probs(B, N, S, T) x V(B, N, T, H_v) const T* v = packed_qkv ? Q + (num_heads_ + kv_num_heads_) * sequence_length * head_size : V; @@ -137,7 +135,7 @@ class GQAAttentionBase { seqlens_k->Data(), batch_size, sequence_length, seqlen_past_kv_cache, seqlen_present_kv_cache, head_size, hidden_size, past_value_data, present_value_data, past_present_share_buffer, packed_qkv, - is_prompt, use_external_kv, tp, allocator); + is_prompt, tp, allocator); } return Status::OK(); @@ -168,7 +166,6 @@ class GQAAttentionBase { const bool past_present_share_buffer, // whether present key and value share the same buffer const bool packed_qkv, // whether Q, K, V are packed const bool is_prompt, // whether it is prompt - const bool use_external_kv, // whether using external KV (skip KV concat) ThreadPool* tp, // thread pool AllocatorPtr allocator) const { // allocator for temporary buffer const ptrdiff_t packed_batch_stride = @@ -180,7 +177,7 @@ class GQAAttentionBase { const size_t past_buff_chunk_length = past_buffer_sequence_length * head_size; // L x H const size_t present_buff_chunk_length = present_buffer_sequence_length * head_size; // T x H - if (!past_present_share_buffer) { + if (present_key && !past_present_share_buffer) { memset((void*)present_key, 0, batch_size * kv_num_heads_ * present_buffer_sequence_length * head_size * sizeof(T)); @@ -242,21 +239,12 @@ class GQAAttentionBase { } const T* k; - if (use_external_kv) { - // External KV mode: use past_key directly (it holds the external KV data in BNSH format). - // No new K to concatenate — the external KV is the complete key sequence. - k = past_key + (i / kv_num_heads_factor) * present_buff_chunk_length; - // Copy to present for output pass-through (once per KV head, not per Q head) - if (present_key != nullptr && !past_present_share_buffer && head_index % kv_num_heads_factor == 0) { - memcpy(present_key + (i / kv_num_heads_factor) * present_buff_chunk_length, - k, SafeInt(total_seqlen) * head_size * sizeof(T)); - } - } else if (packed_qkv) { + if (packed_qkv) { k = K + packed_batch_stride * batch_index + kv_input_chunk_length * (head_index / kv_num_heads_factor); } else { k = K + kv_input_chunk_length * (i / kv_num_heads_factor); } - if (!use_external_kv && nullptr != present_key) { + if (nullptr != present_key) { k = ConcatStateChunkGQA(past_key, k, present_key, present_buff_chunk_length, past_buff_chunk_length, past_chunk_length, kv_input_chunk_length, past_present_share_buffer, i / kv_num_heads_factor); @@ -406,7 +394,6 @@ class GQAAttentionBase { const bool past_present_share_buffer, // whether present key and value share the same buffer const bool packed_qkv, // whether Q, K, V are packed const bool is_prompt, // whether it is prompt - const bool use_external_kv, // whether using external KV (skip KV concat) ThreadPool* tp, AllocatorPtr allocator) const { const ptrdiff_t packed_batch_stride = @@ -417,7 +404,7 @@ class GQAAttentionBase { const size_t past_buff_chunk_length = past_buffer_sequence_length * head_size; // L x H const size_t present_buff_chunk_length = present_buffer_sequence_length * head_size; // T x H - if (!past_present_share_buffer) { + if (present_value && !past_present_share_buffer) { memset((void*)present_value, 0, batch_size * kv_num_heads_ * present_buffer_sequence_length * head_size * sizeof(T)); @@ -460,20 +447,12 @@ class GQAAttentionBase { const size_t past_chunk_length = SafeInt(past_seqlen) * head_size; const T* v; - if (use_external_kv) { - // External KV mode: use past_value directly (it holds the external KV data in BNSH format). - v = past_value + (i / kv_num_heads_factor) * present_buff_chunk_length; - // Copy to present for output pass-through (once per KV head, not per Q head) - if (present_value != nullptr && !past_present_share_buffer && head_index % kv_num_heads_factor == 0) { - memcpy(present_value + (i / kv_num_heads_factor) * present_buff_chunk_length, - v, SafeInt(total_seqlen) * head_size * sizeof(T)); - } - } else if (packed_qkv) { + if (packed_qkv) { v = V + packed_batch_stride * batch_index + kv_input_chunk_length * (head_index / kv_num_heads_factor); } else { v = V + kv_input_chunk_length * (i / kv_num_heads_factor); } - if (!use_external_kv && nullptr != present_value) { + if (nullptr != present_value) { v = ConcatStateChunkGQA(past_value, v, present_value, present_buff_chunk_length, past_buff_chunk_length, past_chunk_length, kv_input_chunk_length, past_present_share_buffer, i / kv_num_heads_factor); diff --git a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc index d126e68f69edd..5ee2f31539bae 100644 --- a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc @@ -55,8 +55,6 @@ Status GroupQueryAttention::Compute(OpKernelContext* context) const { const Tensor* position_ids = context->Input(9); const Tensor* attention_bias = context->Input(10); const Tensor* head_sink = context->Input(11); - const Tensor* external_key = context->Input(14); - const Tensor* external_value = context->Input(15); GroupQueryAttentionParameters parameters = {}; ORT_RETURN_IF_ERROR(group_query_attention_helper::CheckInputs(query, @@ -73,30 +71,13 @@ Status GroupQueryAttention::Compute(OpKernelContext* context) const { total_seqlen_tensor, scale_, softcap_, - 0, - external_key != nullptr)); + 0)); ORT_RETURN_IF_ERROR(group_query_attention_helper::CheckCustomAttentionInputs(position_ids, attention_bias, head_sink, parameters)); - ORT_RETURN_IF_ERROR(group_query_attention_helper::CheckAndSetExternalKV(external_key, external_value, parameters)); - - // External KV is mutually exclusive with provided key/value inputs - if (parameters.use_external_kv && (key != nullptr || value != nullptr)) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "key and value (inputs 1/2) must not be provided when using external_key/external_value. " - "External KV replaces the K,V projections entirely."); - } - - // External KV mode requires do_rotary=0 — K already has RoPE from the source layer - if (parameters.use_external_kv && do_rotary_) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "do_rotary must be 0 when using external_key/external_value. " - "Pre-apply RoPE to Q and use already-rotated K from the source layer."); - } - const int batch_size = parameters.batch_size; const int sequence_length = parameters.sequence_length; const int present_kv_seqlen = parameters.seqlen_present_kv_cache; @@ -144,11 +125,7 @@ Status GroupQueryAttention::Compute(OpKernelContext* context) const { OrtValue Q; OrtValue K; OrtValue V; - if (parameters.use_external_kv) { - // External KV mode: only Q needs transposing. K,V come from external tensors. - ORT_RETURN_IF_ERROR(MaybeTransposeToBNSH( - allocator, batch_size, num_heads_, sequence_length, head_size, query, Q)); - } else if (packed_qkv) { + if (packed_qkv) { ORT_RETURN_IF_ERROR(MaybeTransposeToBNSH( allocator, batch_size, num_heads_ + 2 * kv_num_heads_, sequence_length, head_size, query, Q)); } else { @@ -164,8 +141,8 @@ Status GroupQueryAttention::Compute(OpKernelContext* context) const { OrtValue RotaryQ; OrtValue RotaryK; T* q_rotary = Q.GetMutable()->MutableData(); - T* k_rotary = (packed_qkv || parameters.use_external_kv) ? nullptr : K.GetMutable()->MutableData(); - if (do_rotary_ && !parameters.use_external_kv) { + T* k_rotary = packed_qkv ? nullptr : K.GetMutable()->MutableData(); + if (do_rotary_) { ORT_ENFORCE(cos_cache != nullptr && sin_cache != nullptr, "cos_cache and sin_cache must be provided when do_rotary is true"); // Initialize rotary parameters rotary_embedding_helper::RotaryParameters rotary_params = {}; @@ -255,16 +232,11 @@ Status GroupQueryAttention::Compute(OpKernelContext* context) const { const T* head_sink_data = (head_sink != nullptr) ? head_sink->Data() : nullptr; - // When external KV is provided, use it in place of past_key/past_value for attention computation. - // External KV is pre-computed from another layer (KV-shared layers, e.g., Gemma4). - const Tensor* effective_past_key = parameters.use_external_kv ? external_key : past_key; - const Tensor* effective_past_value = parameters.use_external_kv ? external_value : past_value; - // Compute the attention score and apply the score to V - const T* k_data = (packed_qkv || parameters.use_external_kv) ? nullptr : k_rotary; - const T* v_data = (packed_qkv || parameters.use_external_kv) ? nullptr : V.Get().Data(); + const T* k_data = packed_qkv ? nullptr : k_rotary; + const T* v_data = packed_qkv ? nullptr : V.Get().Data(); return ApplyAttention(q_rotary, k_data, v_data, - head_sink_data, attention_bias, effective_past_key, effective_past_value, output, present_k, present_v, + head_sink_data, attention_bias, past_key, past_value, output, present_k, present_v, output_qk, seqlens_k, parameters, allocator, context); } } // namespace contrib diff --git a/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h b/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h index 0a458c3a8e124..f65568700c0c9 100644 --- a/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h +++ b/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h @@ -96,38 +96,6 @@ Status Check_QKV(const T* packed_qkv, const T* value, const int num_heads, const return Status::OK(); } -template -Status Check_Q_Only(const T* query, const int num_heads, const int kv_num_heads, - int& batch_size, int& sequence_length, int& q_hidden_size, int& kv_hidden_size, int& head_size) { - const auto& query_dims = query->Shape().GetDims(); - if (query_dims.size() != 3) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'query' is expected to have 3 dimensions, got ", - query_dims.size()); - } - batch_size = static_cast(query_dims[0]); - sequence_length = static_cast(query_dims[1]); - q_hidden_size = static_cast(query_dims[2]); - if (q_hidden_size % num_heads != 0) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "q_hidden_size (", q_hidden_size, ") must be divisible by num_heads (", num_heads, - ") in Q-only mode (external KV). Got q_hidden_size % num_heads == ", - q_hidden_size % num_heads); - } - head_size = q_hidden_size / num_heads; - if (head_size == 0) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "head_size must be > 0. Got q_hidden_size=", q_hidden_size, - ", num_heads=", num_heads); - } - if (head_size % 8 != 0) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "head_size must be a multiple of 8. Got head_size % 8 == ", - head_size % 8); - } - kv_hidden_size = head_size * kv_num_heads; - return Status::OK(); -} - template Status CheckPast(const T* past_key, const T* past_value, int batch_size, int kv_num_heads, int head_size, int kv_cache_bit_width, int& past_sequence_length) { @@ -189,66 +157,6 @@ Status CheckPast(const T* past_key, const T* past_value, int batch_size, int kv_ return Status::OK(); } -template -Status CheckExternalKV(const T* external_key, const T* external_value, int batch_size, int kv_num_heads, - int head_size, int kv_cache_bit_width, int& external_sequence_length) { - if (external_key == nullptr || external_value == nullptr) { - if (external_key != nullptr || external_value != nullptr) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "Input 'external_key' and 'external_value' shall be both present or both absent."); - } - return Status::OK(); - } - - const auto& ext_key_dims = external_key->Shape().GetDims(); - const auto& ext_value_dims = external_value->Shape().GetDims(); - - if (ext_key_dims.size() != 4) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "Input 'external_key' is expected to have 4 dimensions (BNSH), got ", - ext_key_dims.size()); - } - if (ext_value_dims.size() != 4) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "Input 'external_value' is expected to have 4 dimensions (BNSH), got ", - ext_value_dims.size()); - } - if (ext_key_dims[0] != batch_size) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "Input 'external_key' dimension 0 should be batch_size, got ", ext_key_dims[0]); - } - if (ext_value_dims[0] != batch_size) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "Input 'external_value' dimension 0 should be batch_size, got ", ext_value_dims[0]); - } - if (ext_key_dims[1] != kv_num_heads) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "Input 'external_key' shall have kv_num_heads, got ", ext_key_dims[1]); - } - if (ext_value_dims[1] != kv_num_heads) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "Input 'external_value' shall have kv_num_heads, got ", ext_value_dims[1]); - } - if (ext_key_dims[2] != ext_value_dims[2]) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "Input 'external_key' and 'external_value' should have same sequence length dimension."); - } - // Validate head dimension (dim 3). For 4-bit quantized KV cache, the stored dimension is head_size / 2. - int expected_head_dim = (kv_cache_bit_width == 4) ? (head_size / 2) : head_size; - if (ext_key_dims[3] != expected_head_dim) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "Input 'external_key' dimension 3 should be head_size (", expected_head_dim, - "), got ", ext_key_dims[3]); - } - if (ext_value_dims[3] != expected_head_dim) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "Input 'external_value' dimension 3 should be head_size (", expected_head_dim, - "), got ", ext_value_dims[3]); - } - external_sequence_length = static_cast(ext_key_dims[2]); - return Status::OK(); -} - template Status CheckRotaryCaches(const T* cos_cache, const T* sin_cache, int head_size, int total_sequence_length, int& rotary_dim) { @@ -299,8 +207,7 @@ Status CheckInputs(const T* query, const T* total_seqlen, float scale, float softcap, - int kv_cache_bit_width, - bool has_external_kv) { + int kv_cache_bit_width) { // Note: Here S* is seqlen_past_kv_cache, S+ is seqlen_present_kv_cache // past_key : (B, N_k, S*, H) or (B, N_k, S+, H) or nullptr // past_value : (B, N_k, S*, H) or (B, N_k, S+, H) or nullptr @@ -335,14 +242,8 @@ Status CheckInputs(const T* query, int q_hidden_size = 0; int kv_hidden_size = 0; int head_size = 0; - // When external KV is provided, key/value can be nullptr without implying packed QKV. - // In this mode, query contains only Q (not packed QKV). - const bool is_packed_qkv = (key == nullptr) && !has_external_kv; - if (has_external_kv && key == nullptr) { - // Q-only mode: query is just Q, K and V come from external tensors - ORT_RETURN_IF_ERROR(Check_Q_Only(query, num_heads, kv_num_heads, batch_size, sequence_length, - q_hidden_size, kv_hidden_size, head_size)); - } else if (!is_packed_qkv) { + const bool is_packed_qkv = (key == nullptr); + if (!is_packed_qkv) { ORT_RETURN_IF_ERROR(Check_Q_K_V(query, key, value, num_heads, kv_num_heads, batch_size, sequence_length, q_hidden_size, kv_hidden_size, head_size)); } else { @@ -449,13 +350,12 @@ Status CheckInputs(const T* query, float scale, float softcap, int kv_cache_bit_width, - int max_threads_per_block, - bool has_external_kv = false) { + int max_threads_per_block) { if (max_threads_per_block > 0 && num_heads > max_threads_per_block) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "num_heads should be no larger than ", max_threads_per_block); } - return CheckInputs(query, key, value, past_key, past_value, cos_cache, sin_cache, parameters, num_heads, kv_num_heads, seqlens_k, total_seqlen, scale, softcap, kv_cache_bit_width, has_external_kv); + return CheckInputs(query, key, value, past_key, past_value, cos_cache, sin_cache, parameters, num_heads, kv_num_heads, seqlens_k, total_seqlen, scale, softcap, kv_cache_bit_width); } template @@ -545,53 +445,6 @@ inline Status CheckNoQKOutput(int num_outputs, int qk_output) { return Status::OK(); } -// Validate and configure external KV inputs for KV-shared layers. -// Call this after CheckInputs to set up external KV parameters. -template -Status CheckAndSetExternalKV(const T* external_key, const T* external_value, - GroupQueryAttentionParameters& parameters) { - if (external_key == nullptr && external_value == nullptr) { - return Status::OK(); - } - - int external_sequence_length = 0; - ORT_RETURN_IF_ERROR(CheckExternalKV(external_key, external_value, - parameters.batch_size, parameters.kv_num_heads, - parameters.head_size, parameters.kv_cache_bit_width, - external_sequence_length)); - - parameters.use_external_kv = true; - parameters.external_kv_sequence_length = external_sequence_length; - - // External KV mode is incompatible with packed QKV — query must contain only Q. - if (parameters.is_packed_qkv) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "external_key/external_value cannot be used with packed QKV input. " - "Provide query as Q-only (without K,V) when using external KV."); - } - - // External KV replaces the internal KV cache — past_key/past_value should not be provided. - if (parameters.seqlen_past_kv_cache > 0) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "past_key/past_value should not be provided when using external_key/external_value. " - "External KV replaces the internal KV cache entirely."); - } - - // When using external KV, the total sequence length for attention is determined - // by the external KV tensor, not the internal KV cache. - // Validate that the original total_sequence_length doesn't exceed external KV length. - if (parameters.total_sequence_length > external_sequence_length) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "total_sequence_length (", parameters.total_sequence_length, - ") exceeds external KV sequence length (", external_sequence_length, - "). Ensure seqlens_k is consistent with the external KV tensor size."); - } - parameters.total_sequence_length = external_sequence_length; - parameters.seqlen_present_kv_cache = external_sequence_length; - - return Status::OK(); -} - } // namespace group_query_attention_helper } // namespace contrib } // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/attention_data.h b/onnxruntime/contrib_ops/cuda/bert/attention_data.h index 7c2f805b5292b..c54a1fea9ad3a 100644 --- a/onnxruntime/contrib_ops/cuda/bert/attention_data.h +++ b/onnxruntime/contrib_ops/cuda/bert/attention_data.h @@ -192,10 +192,6 @@ struct GroupQueryAttentionData { U* present_key = nullptr; U* present_value = nullptr; - // External KV for KV-shared layers (e.g., Gemma4) - const U* external_key = nullptr; - const U* external_value = nullptr; - // Kernel Flags bool use_flash_attention = false; bool use_memory_efficient_attention = false; diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc index 2562e72451bd8..9563292f9187c 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc @@ -166,8 +166,6 @@ Status GroupQueryAttention::ComputeInternal(OpKernelContext* context) cons const Tensor* head_sink = context->Input(11); const Tensor* k_scale = context->Input(12); const Tensor* v_scale = context->Input(13); - const Tensor* external_key = context->Input(14); - const Tensor* external_value = context->Input(15); if (k_quant_type_ != KVQuantizationType::NONE) { if (k_scale == nullptr) { @@ -223,23 +221,13 @@ Status GroupQueryAttention::ComputeInternal(OpKernelContext* context) cons scale_, softcap_, kv_cache_bit_width_, - device_prop.maxThreadsPerBlock, - external_key != nullptr)); + device_prop.maxThreadsPerBlock)); ORT_RETURN_IF_ERROR(group_query_attention_helper::CheckCustomAttentionInputs(position_ids, attention_bias, head_sink, parameters)); - ORT_RETURN_IF_ERROR(group_query_attention_helper::CheckAndSetExternalKV(external_key, external_value, parameters)); - - // External KV is mutually exclusive with provided key/value inputs - if (parameters.use_external_kv && (key != nullptr || value != nullptr)) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "key and value (inputs 1/2) must not be provided when using external_key/external_value. " - "External KV replaces the K,V projections entirely."); - } - parameters.local_window_size = local_window_size_; parameters.is_unidirectional = is_unidirectional_; parameters.use_smooth_softmax = use_smooth_softmax_ || head_sink != nullptr; @@ -251,14 +239,6 @@ Status GroupQueryAttention::ComputeInternal(OpKernelContext* context) cons parameters.do_rotary = do_rotary_; parameters.rotary_interleaved = rotary_interleaved_; - // When using external KV, disable rotary embedding — the external KV already has RoPE applied - // from the source layer, and the caller is expected to pre-apply RoPE to Q. - if (parameters.use_external_kv && parameters.do_rotary) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "do_rotary must be 0 when using external_key/external_value. " - "Pre-apply RoPE to Q and use already-rotated K from the source layer."); - } - // The current GQA CUDA implementation will never be able to have a QK output. // GQA CUDA uses either flash attention or memory efficient attention. Neither kernel supports returning the QK output. ORT_RETURN_IF_ERROR(group_query_attention_helper::CheckNoQKOutput( @@ -310,26 +290,12 @@ Status GroupQueryAttention::ComputeInternal(OpKernelContext* context) cons data.k_scale = k_scale == nullptr ? nullptr : reinterpret_cast(k_scale->DataRaw()); data.v_scale = v_scale == nullptr ? nullptr : reinterpret_cast(v_scale->DataRaw()); - if (parameters.use_external_kv) { - // External KV mode: use external tensors as the KV source for attention. - // The external KV is treated as "past" KV since it's already computed. - // No KV cache update is performed — the present outputs are copies/views of external KV. - data.external_key = reinterpret_cast(external_key->Data()); - data.external_value = reinterpret_cast(external_value->Data()); - data.past_key = data.external_key; - data.past_value = data.external_value; - data.present_key = reinterpret_cast(present_key_output->MutableData()); - data.present_value = reinterpret_cast(present_value_output->MutableData()); - // Mark as shared buffer so the kernel treats external KV as already-populated cache - parameters.past_present_share_buffer = false; - } else { - data.past_key = (past_key == nullptr) ? nullptr : reinterpret_cast(past_key->Data()); - data.past_value = (past_value == nullptr) ? nullptr : reinterpret_cast(past_value->Data()); - data.present_key = reinterpret_cast(present_key_output->MutableData()); - data.present_value = reinterpret_cast(present_value_output->MutableData()); - // Compute past_present_share_buffer early since it's needed for flash attention path selection. - parameters.past_present_share_buffer = (data.past_key == data.present_key); - } + data.past_key = (past_key == nullptr) ? nullptr : reinterpret_cast(past_key->Data()); + data.past_value = (past_value == nullptr) ? nullptr : reinterpret_cast(past_value->Data()); + data.present_key = (present_key_output != nullptr) ? reinterpret_cast(present_key_output->MutableData()) : nullptr; + data.present_value = (present_value_output != nullptr) ? reinterpret_cast(present_value_output->MutableData()) : nullptr; + // Compute past_present_share_buffer early since it's needed for flash attention path selection. + parameters.past_present_share_buffer = (data.past_key != nullptr && data.past_key == data.present_key); bool is_inputs_quantized = (k_quant_type_ != KVQuantizationType::NONE) || (v_quant_type_ != KVQuantizationType::NONE); constexpr bool is_int8 = std::is_same::value; @@ -594,12 +560,12 @@ Status GroupQueryAttention::ComputeInternal(OpKernelContext* context) cons } // Validate past_value pointer consistency (past_present_share_buffer was computed early after pointer setup) - if (parameters.use_external_kv) { - // External KV mode: past and present are separate (external source -> present output) - } else if (parameters.past_present_share_buffer) { - ORT_ENFORCE(data.past_value == data.present_value, "past_value and present_value must be the same tensor when past_present_share_buffer is true"); - } else { - ORT_ENFORCE(data.past_value != data.present_value, "past_value and present_value must be different tensors when past_present_share_buffer is false"); + if (data.present_value != nullptr) { + if (parameters.past_present_share_buffer) { + ORT_ENFORCE(data.past_value == data.present_value, "past_value and present_value must be the same tensor when past_present_share_buffer is true"); + } else { + ORT_ENFORCE(data.past_value != data.present_value, "past_value and present_value must be different tensors when past_present_share_buffer is false"); + } } data.output = reinterpret_cast(output->MutableData()); diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu index 4f0da8fda4d08..3ce396989b181 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu @@ -98,23 +98,12 @@ Status PrepareQKV( q_out = nullptr; } - // External KV mode: the external KV data is already set as past_key/past_value. - // Copy it into the present buffers and skip the KV append/RoPE logic. - if (parameters.use_external_kv) { - U* k = reinterpret_cast(data.present_key); - U* v = reinterpret_cast(data.present_value); - int external_seq_len = parameters.external_kv_sequence_length; - - // For 4-bit quantized KV cache, the stored head dimension is head_size/2 (two nibbles per byte). - // Use the packed dimension to compute the correct copy size. - int cache_head_dim = (parameters.kv_cache_bit_width == 4) ? (head_size + 1) / 2 : head_size; - size_t kv_copy_size = (size_t)batch_size * kv_num_heads * external_seq_len * cache_head_dim * sizeof(U); - CUDA_CALL_THROW(cudaMemcpyAsync(k, data.past_key, kv_copy_size, cudaMemcpyDeviceToDevice, stream)); - CUDA_CALL_THROW(cudaMemcpyAsync(v, data.past_value, kv_copy_size, cudaMemcpyDeviceToDevice, stream)); - - // Q is used directly from the input - q = reinterpret_cast(data.query); - return Status::OK(); + // present_key/present_value are required for the CUDA path since flash attention + // and memory-efficient attention read directly from the present KV buffers. + // The CPU path supports optional present outputs for KV-shared layers. + if (data.present_key == nullptr || data.present_value == nullptr) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "present_key and present_value outputs are required for the CUDA GroupQueryAttention kernel."); } U* k = reinterpret_cast(data.present_key); diff --git a/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc b/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc index 98f09a496f579..5fff0516c7ce3 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc @@ -212,8 +212,7 @@ Status GroupQueryAttention::ComputeInternal(onnxruntime::webgpu::ComputeContext& scale_, softcap_, 0, - static_cast(context.DeviceLimits().maxComputeInvocationsPerWorkgroup), - /*has_external_kv=*/false)); + static_cast(context.DeviceLimits().maxComputeInvocationsPerWorkgroup))); params.use_smooth_softmax = use_smooth_softmax_; params.rotary_interleaved = rotary_interleaved_; diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index ccecc28ec6afc..e8ec04586a9d6 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -240,20 +240,13 @@ void BaseGroupQueryAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceConte if (ctx.getNumOutputs() >= 3) { // has present output const auto* past_key_type = ctx.getInputType(past_key_index); - // external_key is at input index 14 for GroupQueryAttention - const auto* external_key_type = (ctx.getNumInputs() > 14) ? ctx.getInputType(14) : nullptr; if (past_key_type != nullptr) { // present_key and present_value have the same type as past_key/past_value. // This allows them to be int8 or packed uint8 when quantization is enabled. ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, past_key_index, 1); // present_key ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, past_key_index + 1, 2); // present_value - } else if (external_key_type != nullptr) { - // When external KV is provided (inputs 14/15), present outputs should match - // the external KV type (T_CACHE), not the query type (T). - ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 14, 1); // present_key from external_key - ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 15, 2); // present_value from external_value } else { - // If no past state and no external KV, present is the same type as query. + // If no past state, present is the same type as query. ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 0, 1); ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 0, 2); } @@ -1321,21 +1314,6 @@ ONNX_MS_OPERATOR_SET_SCHEMA( OpSchema::Optional) .Input(12, "k_scale", "Scale tensor for past_key.", "T_KV_SCALE", OpSchema::Optional) .Input(13, "v_scale", "Scale tensor for past_value.", "T_KV_SCALE", OpSchema::Optional) - .Input(14, - "external_key", - "External pre-computed key tensor in BNSH format (batch_size, kv_num_heads, external_seq_len, head_size). " - "Used for KV-shared layers that borrow K,V from another layer's present KV output. " - "When provided, the operator skips its internal KV cache update and uses this tensor directly " - "for attention computation. RoPE is not applied to external keys (assumed already applied).", - "T_CACHE", - OpSchema::Optional) - .Input(15, - "external_value", - "External pre-computed value tensor in BNSH format (batch_size, kv_num_heads, external_seq_len, head_size). " - "Must be provided together with external_key. When provided, the operator uses this tensor " - "for attention computation instead of the internal KV cache.", - "T_CACHE", - OpSchema::Optional) .Output(0, "output", "3D output tensor with shape (batch_size, sequence_length, hidden_size)", @@ -1345,13 +1323,15 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "present state key with support for format BNSH. When past_key uses same tensor as present_key" "(k-v buffer), it is of length max_sequence_length... otherwise of length past_sequence_length +" "kv_sequence_length.", - "T_CACHE") + "T_CACHE", + OpSchema::Optional) .Output(2, "present_value", "present state value with support for format BNSH. When past_value uses same tensor as present_value" "(k-v buffer), it is of length max_sequence_length... otherwise of length past_sequence_length +" "kv_sequence_length.", - "T_CACHE") + "T_CACHE", + OpSchema::Optional) .Output(3, "output_qk", "Values of QK matrix multiplication, either before or after softmax normalization", diff --git a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc index 593fc55c01e99..cc00d1bf61fbb 100644 --- a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc @@ -308,19 +308,18 @@ TEST(GroupQueryAttentionTest, SeqlensKWrongLength) { } // ============================================================================ -// External KV tests (inputs 14/15: external_key, external_value) +// Optional present_key/present_value output tests // ============================================================================ -// Helper for external KV tests -static void RunGQAExternalKVTest( - int external_seq_len, +// Helper for tests with optional present outputs. +// When omit_present=true, present_key and present_value outputs are not connected. +static void RunGQAOptionalPresentTest( + int batch_size, + int sequence_length, + int total_seq_len, + bool omit_present, OpTester::ExpectResult expect, - const std::string& expected_message, - bool provide_key_value = false, - bool provide_past = false, - bool do_rotary = false) { - constexpr int batch_size = 1; - constexpr int sequence_length = 1; + const std::string& expected_message) { constexpr int num_heads = 2; constexpr int kv_num_heads = 1; constexpr int head_size = 8; @@ -330,62 +329,43 @@ static void RunGQAExternalKVTest( OpTester tester("GroupQueryAttention", 1, onnxruntime::kMSDomain); tester.AddAttribute("num_heads", static_cast(num_heads)); tester.AddAttribute("kv_num_heads", static_cast(kv_num_heads)); - if (do_rotary) { - tester.AddAttribute("do_rotary", 1); - } - // Query (Q-only when using external KV) std::vector query_data(batch_size * sequence_length * hidden_size, 1.0f); tester.AddInput("query", {batch_size, sequence_length, hidden_size}, query_data); - // Key/Value inputs (should be absent for external KV) - if (provide_key_value) { - std::vector key_data(batch_size * sequence_length * kv_hidden_size, 1.0f); - std::vector value_data(batch_size * sequence_length * kv_hidden_size, 1.0f); - tester.AddInput("key", {batch_size, sequence_length, kv_hidden_size}, key_data); - tester.AddInput("value", {batch_size, sequence_length, kv_hidden_size}, value_data); - } else { - tester.AddOptionalInputEdge(); // key - tester.AddOptionalInputEdge(); // value - } + std::vector key_data(batch_size * sequence_length * kv_hidden_size, 0.5f); + tester.AddInput("key", {batch_size, sequence_length, kv_hidden_size}, key_data); - // Past key/value (should be absent for external KV) - if (provide_past) { - std::vector past_k(batch_size * kv_num_heads * 4 * head_size, 0.5f); - std::vector past_v(batch_size * kv_num_heads * 4 * head_size, 0.5f); - tester.AddInput("past_key", {batch_size, kv_num_heads, 4, head_size}, past_k); - tester.AddInput("past_value", {batch_size, kv_num_heads, 4, head_size}, past_v); - } else { - tester.AddOptionalInputEdge(); // past_key - tester.AddOptionalInputEdge(); // past_value - } + std::vector value_data(batch_size * sequence_length * kv_hidden_size, 0.5f); + tester.AddInput("value", {batch_size, sequence_length, kv_hidden_size}, value_data); + + tester.AddOptionalInputEdge(); // past_key + tester.AddOptionalInputEdge(); // past_value + + tester.AddInput("seqlens_k", {batch_size}, {static_cast(total_seq_len - 1)}); + tester.AddInput("total_sequence_length", {1}, {static_cast(total_seq_len)}); + + tester.AddOptionalInputEdge(); // cos_cache + tester.AddOptionalInputEdge(); // sin_cache + tester.AddOptionalInputEdge(); // position_ids + tester.AddOptionalInputEdge(); // attention_bias + tester.AddOptionalInputEdge(); // head_sink - // seqlens_k = external_seq_len - 1 (historical convention) - tester.AddInput("seqlens_k", {batch_size}, {static_cast(external_seq_len - 1)}); - tester.AddInput("total_sequence_length", {1}, {static_cast(external_seq_len)}); - - tester.AddOptionalInputEdge(); // cos_cache (7) - tester.AddOptionalInputEdge(); // sin_cache (8) - tester.AddOptionalInputEdge(); // position_ids (9) - tester.AddOptionalInputEdge(); // attention_bias (10) - tester.AddOptionalInputEdge(); // head_sink (11) - tester.AddOptionalInputEdge(); // k_scale (12) - tester.AddOptionalInputEdge(); // v_scale (13) - - // External key/value (inputs 14/15) — BNSH format - std::vector ext_key(batch_size * kv_num_heads * external_seq_len * head_size, 0.5f); - std::vector ext_value(batch_size * kv_num_heads * external_seq_len * head_size, 0.5f); - tester.AddInput("external_key", {batch_size, kv_num_heads, external_seq_len, head_size}, ext_key); - tester.AddInput("external_value", {batch_size, kv_num_heads, external_seq_len, head_size}, ext_value); - - // Outputs - int present_seq_len = std::max(1, external_seq_len); + // Output 0: output (always required) tester.AddOutput("output", {batch_size, sequence_length, hidden_size}, std::vector(batch_size * sequence_length * hidden_size, 0.0f)); - tester.AddOutput("present_key", {batch_size, kv_num_heads, present_seq_len, head_size}, - std::vector(batch_size * kv_num_heads * present_seq_len * head_size, 0.0f)); - tester.AddOutput("present_value", {batch_size, kv_num_heads, present_seq_len, head_size}, - std::vector(batch_size * kv_num_heads * present_seq_len * head_size, 0.0f)); + + if (omit_present) { + // Omit present_key and present_value — they are optional + tester.AddOptionalOutputEdge(); // present_key + tester.AddOptionalOutputEdge(); // present_value + } else { + int present_seq_len = total_seq_len; + tester.AddOutput("present_key", {batch_size, kv_num_heads, present_seq_len, head_size}, + std::vector(batch_size * kv_num_heads * present_seq_len * head_size, 0.0f)); + tester.AddOutput("present_value", {batch_size, kv_num_heads, present_seq_len, head_size}, + std::vector(batch_size * kv_num_heads * present_seq_len * head_size, 0.0f)); + } if (expect == OpTester::ExpectResult::kExpectSuccess) { tester.SetOutputTolerance(1e6f); @@ -396,42 +376,37 @@ static void RunGQAExternalKVTest( tester.Run(expect, expected_message, {}, nullptr, &execution_providers); } -// Basic: external KV with Q-only query should succeed -TEST(GroupQueryAttentionTest, ExternalKV_BasicSuccess) { - RunGQAExternalKVTest( - /*external_seq_len=*/8, +// Baseline: GQA with present outputs connected works as before +TEST(GroupQueryAttentionTest, OptionalPresent_WithPresent) { + RunGQAOptionalPresentTest( + /*batch_size=*/1, + /*sequence_length=*/4, + /*total_seq_len=*/4, + /*omit_present=*/false, OpTester::ExpectResult::kExpectSuccess, ""); } -// Reject: external KV with key/value inputs provided (mutual exclusivity) -TEST(GroupQueryAttentionTest, ExternalKV_RejectsProvidedKeyValue) { - RunGQAExternalKVTest( - /*external_seq_len=*/8, - OpTester::ExpectResult::kExpectFailure, - "key and value (inputs 1/2) must not be provided", - /*provide_key_value=*/true); -} - -// Reject: external KV with past_key/past_value provided -TEST(GroupQueryAttentionTest, ExternalKV_RejectsPastKV) { - RunGQAExternalKVTest( - /*external_seq_len=*/8, - OpTester::ExpectResult::kExpectFailure, - "past_key/past_value should not be provided", - /*provide_key_value=*/false, - /*provide_past=*/true); +// KV-shared layer scenario: present outputs omitted, attention uses K,V directly +TEST(GroupQueryAttentionTest, OptionalPresent_WithoutPresent) { + RunGQAOptionalPresentTest( + /*batch_size=*/1, + /*sequence_length=*/4, + /*total_seq_len=*/4, + /*omit_present=*/true, + OpTester::ExpectResult::kExpectSuccess, + ""); } -// Reject: external KV with do_rotary=1 -TEST(GroupQueryAttentionTest, ExternalKV_RejectsDoRotary) { - RunGQAExternalKVTest( - /*external_seq_len=*/8, - OpTester::ExpectResult::kExpectFailure, - "do_rotary must be 0", - /*provide_key_value=*/false, - /*provide_past=*/false, - /*do_rotary=*/true); +// Batched: present outputs omitted with batch_size > 1 +TEST(GroupQueryAttentionTest, OptionalPresent_Batched) { + RunGQAOptionalPresentTest( + /*batch_size=*/2, + /*sequence_length=*/3, + /*total_seq_len=*/3, + /*omit_present=*/true, + OpTester::ExpectResult::kExpectSuccess, + ""); } } // namespace test From aad74efae671df82a813b519b3fcc0e3fb493c83 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Tue, 28 Apr 2026 11:17:09 -0700 Subject: [PATCH 08/32] Fix tests --- onnxruntime/test/contrib_ops/group_query_attention_op_test.cc | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc index cc00d1bf61fbb..1d57488d51363 100644 --- a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc @@ -342,7 +342,8 @@ static void RunGQAOptionalPresentTest( tester.AddOptionalInputEdge(); // past_key tester.AddOptionalInputEdge(); // past_value - tester.AddInput("seqlens_k", {batch_size}, {static_cast(total_seq_len - 1)}); + std::vector seqlens_k_data(batch_size, static_cast(total_seq_len - 1)); + tester.AddInput("seqlens_k", {batch_size}, seqlens_k_data); tester.AddInput("total_sequence_length", {1}, {static_cast(total_seq_len)}); tester.AddOptionalInputEdge(); // cos_cache From 64005dd59d91c1fa65b4aa0aa0f8ac752167879c Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Tue, 28 Apr 2026 17:32:23 -0700 Subject: [PATCH 09/32] Update the docs --- docs/ContribOperators.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 9aa44a1600ae6..45e85fcd9c9d5 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -2671,14 +2671,14 @@ This version of the operator has been available since version 1 of the 'com.micr
Scale tensor for past_value.
-#### Outputs (3 - 4) +#### Outputs (1 - 4)
output : T
3D output tensor with shape (batch_size, sequence_length, hidden_size)
-
present_key : T_CACHE
+
present_key (optional) : T_CACHE
present state key with support for format BNSH. When past_key uses same tensor as present_key(k-v buffer), it is of length max_sequence_length... otherwise of length past_sequence_length +kv_sequence_length.
-
present_value : T_CACHE
+
present_value (optional) : T_CACHE
present state value with support for format BNSH. When past_value uses same tensor as present_value(k-v buffer), it is of length max_sequence_length... otherwise of length past_sequence_length +kv_sequence_length.
output_qk (optional) : T
Values of QK matrix multiplication, either before or after softmax normalization
From 2b3a2ce374f9324f37e02f0197f8d12b312a160f Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Fri, 1 May 2026 11:10:03 -0700 Subject: [PATCH 10/32] Address comments --- .../contrib_ops/cpu/bert/group_query_attention.cc | 11 +++++++++++ .../contrib_ops/cuda/bert/group_query_attention.cc | 7 +++++++ .../test/contrib_ops/group_query_attention_op_test.cc | 11 +++++++++++ 3 files changed, 29 insertions(+) diff --git a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc index 5ee2f31539bae..44d3895abf672 100644 --- a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc @@ -113,6 +113,17 @@ Status GroupQueryAttention::Compute(OpKernelContext* context) const { Tensor* present_k = context->Output(1, present_k_shape); Tensor* present_v = context->Output(2, present_v_shape); + // Optional present outputs are only safe for first-prompt with no past KV. + // When past exists or total_sequence_length > sequence_length, the attention + // GEMMs use total_seqlen which requires a concatenated past+current KV buffer + // that only ConcatStateChunkGQA builds into present_key/present_value. + if ((present_k == nullptr || present_v == nullptr) && !parameters.is_first_prompt) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "present_key and present_value outputs are required when past state exists " + "(total_sequence_length > sequence_length). Omitting present outputs is only " + "supported for first-prompt inference with no past KV cache."); + } + std::vector output_qk_shape{static_cast(batch_size), static_cast(num_heads_), static_cast(parameters.sequence_length), static_cast(parameters.total_sequence_length)}; Tensor* output_qk = context->Output(3, output_qk_shape); diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc index 9563292f9187c..0ad62f62daa4c 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc @@ -267,6 +267,13 @@ Status GroupQueryAttention::ComputeInternal(OpKernelContext* context) cons Tensor* present_key_output = context->Output(1, present_shape); // present_key Tensor* present_value_output = context->Output(2, present_shape); // present_value + // Optional present outputs are only safe for first-prompt with no past KV. + if ((present_key_output == nullptr || present_value_output == nullptr) && !parameters.is_first_prompt) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "present_key and present_value outputs are required when past state exists. " + "Omitting present outputs is only supported for first-prompt inference."); + } + IAllocatorUniquePtr k_buffer; IAllocatorUniquePtr v_buffer; IAllocatorUniquePtr rotary_buffer; diff --git a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc index 1d57488d51363..fe32ab9e95329 100644 --- a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc @@ -410,5 +410,16 @@ TEST(GroupQueryAttentionTest, OptionalPresent_Batched) { ""); } +// Reject: omitting present outputs when total_seq_len > sequence_length (decode with past) +TEST(GroupQueryAttentionTest, OptionalPresent_RejectWithPast) { + RunGQAOptionalPresentTest( + /*batch_size=*/1, + /*sequence_length=*/1, + /*total_seq_len=*/5, + /*omit_present=*/true, + OpTester::ExpectResult::kExpectFailure, + "present_key and present_value outputs are required when past state exists"); +} + } // namespace test } // namespace onnxruntime From 9a14803a5d80ede166cda5688a57cdc75934ebb7 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Fri, 1 May 2026 11:50:20 -0700 Subject: [PATCH 11/32] Address copilot comments --- .../cuda/bert/group_query_attention.cc | 19 +- .../cuda/bert/group_query_attention_impl.cu | 8 - .../group_query_attention_op_test.cc | 207 ++++++++++++------ 3 files changed, 159 insertions(+), 75 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc index 0ad62f62daa4c..97df349656c07 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc @@ -274,6 +274,17 @@ Status GroupQueryAttention::ComputeInternal(OpKernelContext* context) cons "Omitting present outputs is only supported for first-prompt inference."); } + // When present outputs are omitted, allocate internal scratch buffers so the + // CUDA kernels (flash attention, MEA, unfused) have a valid KV workspace. + // This keeps behavior consistent with the CPU EP. + IAllocatorUniquePtr present_key_scratch; + IAllocatorUniquePtr present_value_scratch; + if (present_key_output == nullptr || present_value_output == nullptr) { + size_t present_kv_bytes = present_shape.Size() * sizeof(U); + present_key_scratch = GetScratchBuffer(present_kv_bytes, context->GetComputeStream()); + present_value_scratch = GetScratchBuffer(present_kv_bytes, context->GetComputeStream()); + } + IAllocatorUniquePtr k_buffer; IAllocatorUniquePtr v_buffer; IAllocatorUniquePtr rotary_buffer; @@ -299,8 +310,12 @@ Status GroupQueryAttention::ComputeInternal(OpKernelContext* context) cons data.past_key = (past_key == nullptr) ? nullptr : reinterpret_cast(past_key->Data()); data.past_value = (past_value == nullptr) ? nullptr : reinterpret_cast(past_value->Data()); - data.present_key = (present_key_output != nullptr) ? reinterpret_cast(present_key_output->MutableData()) : nullptr; - data.present_value = (present_value_output != nullptr) ? reinterpret_cast(present_value_output->MutableData()) : nullptr; + data.present_key = (present_key_output != nullptr) + ? reinterpret_cast(present_key_output->MutableData()) + : reinterpret_cast(present_key_scratch.get()); + data.present_value = (present_value_output != nullptr) + ? reinterpret_cast(present_value_output->MutableData()) + : reinterpret_cast(present_value_scratch.get()); // Compute past_present_share_buffer early since it's needed for flash attention path selection. parameters.past_present_share_buffer = (data.past_key != nullptr && data.past_key == data.present_key); diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu index 3ce396989b181..ebb6a0b0da215 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu @@ -98,14 +98,6 @@ Status PrepareQKV( q_out = nullptr; } - // present_key/present_value are required for the CUDA path since flash attention - // and memory-efficient attention read directly from the present KV buffers. - // The CPU path supports optional present outputs for KV-shared layers. - if (data.present_key == nullptr || data.present_value == nullptr) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "present_key and present_value outputs are required for the CUDA GroupQueryAttention kernel."); - } - U* k = reinterpret_cast(data.present_key); U* v = reinterpret_cast(data.present_value); int max_cache_length = parameters.seqlen_present_kv_cache; diff --git a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc index fe32ab9e95329..9b0940a96d6e4 100644 --- a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc @@ -311,37 +311,34 @@ TEST(GroupQueryAttentionTest, SeqlensKWrongLength) { // Optional present_key/present_value output tests // ============================================================================ -// Helper for tests with optional present outputs. -// When omit_present=true, present_key and present_value outputs are not connected. -static void RunGQAOptionalPresentTest( +// Run GQA with the given inputs and return the output tensor as a vector. +// This lets us compare outputs between present-connected and present-omitted runs. +static std::vector RunGQAAndGetOutput( int batch_size, int sequence_length, - int total_seq_len, - bool omit_present, - OpTester::ExpectResult expect, - const std::string& expected_message) { - constexpr int num_heads = 2; - constexpr int kv_num_heads = 1; - constexpr int head_size = 8; - constexpr int hidden_size = num_heads * head_size; - constexpr int kv_hidden_size = kv_num_heads * head_size; + const std::vector& query_data, + const std::vector& key_data, + const std::vector& value_data, + int num_heads, + int kv_num_heads, + int head_size, + bool omit_present) { + const int hidden_size = num_heads * head_size; + const int kv_hidden_size = kv_num_heads * head_size; + const int total_seq_len = sequence_length; // first-prompt: no past OpTester tester("GroupQueryAttention", 1, onnxruntime::kMSDomain); tester.AddAttribute("num_heads", static_cast(num_heads)); tester.AddAttribute("kv_num_heads", static_cast(kv_num_heads)); - std::vector query_data(batch_size * sequence_length * hidden_size, 1.0f); tester.AddInput("query", {batch_size, sequence_length, hidden_size}, query_data); - - std::vector key_data(batch_size * sequence_length * kv_hidden_size, 0.5f); tester.AddInput("key", {batch_size, sequence_length, kv_hidden_size}, key_data); - - std::vector value_data(batch_size * sequence_length * kv_hidden_size, 0.5f); tester.AddInput("value", {batch_size, sequence_length, kv_hidden_size}, value_data); tester.AddOptionalInputEdge(); // past_key tester.AddOptionalInputEdge(); // past_value + // First-prompt: seqlens_k = total_seq_len - 1 per GQA convention std::vector seqlens_k_data(batch_size, static_cast(total_seq_len - 1)); tester.AddInput("seqlens_k", {batch_size}, seqlens_k_data); tester.AddInput("total_sequence_length", {1}, {static_cast(total_seq_len)}); @@ -352,73 +349,153 @@ static void RunGQAOptionalPresentTest( tester.AddOptionalInputEdge(); // attention_bias tester.AddOptionalInputEdge(); // head_sink - // Output 0: output (always required) + // Use a placeholder output with large tolerance — we extract the actual values below. + const int output_size = batch_size * sequence_length * hidden_size; tester.AddOutput("output", {batch_size, sequence_length, hidden_size}, - std::vector(batch_size * sequence_length * hidden_size, 0.0f)); + std::vector(output_size, 0.0f)); if (omit_present) { - // Omit present_key and present_value — they are optional tester.AddOptionalOutputEdge(); // present_key tester.AddOptionalOutputEdge(); // present_value } else { - int present_seq_len = total_seq_len; - tester.AddOutput("present_key", {batch_size, kv_num_heads, present_seq_len, head_size}, - std::vector(batch_size * kv_num_heads * present_seq_len * head_size, 0.0f)); - tester.AddOutput("present_value", {batch_size, kv_num_heads, present_seq_len, head_size}, - std::vector(batch_size * kv_num_heads * present_seq_len * head_size, 0.0f)); - } - - if (expect == OpTester::ExpectResult::kExpectSuccess) { - tester.SetOutputTolerance(1e6f); + const int present_size = batch_size * kv_num_heads * total_seq_len * head_size; + tester.AddOutput("present_key", {batch_size, kv_num_heads, total_seq_len, head_size}, + std::vector(present_size, 0.0f)); + tester.AddOutput("present_value", {batch_size, kv_num_heads, total_seq_len, head_size}, + std::vector(present_size, 0.0f)); } + tester.SetOutputTolerance(1e6f); // We compare fetched outputs ourselves std::vector> execution_providers; execution_providers.push_back(DefaultCpuExecutionProvider()); - tester.Run(expect, expected_message, {}, nullptr, &execution_providers); -} + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); -// Baseline: GQA with present outputs connected works as before -TEST(GroupQueryAttentionTest, OptionalPresent_WithPresent) { - RunGQAOptionalPresentTest( - /*batch_size=*/1, - /*sequence_length=*/4, - /*total_seq_len=*/4, - /*omit_present=*/false, - OpTester::ExpectResult::kExpectSuccess, - ""); + // Extract the output tensor values + auto fetches = tester.GetFetches(); + const float* out_data = fetches[0].Get().Data(); + return std::vector(out_data, out_data + output_size); } -// KV-shared layer scenario: present outputs omitted, attention uses K,V directly -TEST(GroupQueryAttentionTest, OptionalPresent_WithoutPresent) { - RunGQAOptionalPresentTest( - /*batch_size=*/1, - /*sequence_length=*/4, - /*total_seq_len=*/4, - /*omit_present=*/true, - OpTester::ExpectResult::kExpectSuccess, - ""); +// Core correctness test: output must be identical whether present outputs +// are connected or omitted (first-prompt, no past KV). +TEST(GroupQueryAttentionTest, OptionalPresent_OutputMatchesWithAndWithout) { + constexpr int batch_size = 1; + constexpr int sequence_length = 4; + constexpr int num_heads = 2; + constexpr int kv_num_heads = 1; + constexpr int head_size = 8; + constexpr int hidden_size = num_heads * head_size; + constexpr int kv_hidden_size = kv_num_heads * head_size; + + // Deterministic non-trivial inputs + std::vector query_data(batch_size * sequence_length * hidden_size); + std::vector key_data(batch_size * sequence_length * kv_hidden_size); + std::vector value_data(batch_size * sequence_length * kv_hidden_size); + for (size_t i = 0; i < query_data.size(); i++) query_data[i] = 0.1f * static_cast(i % 7 + 1); + for (size_t i = 0; i < key_data.size(); i++) key_data[i] = 0.2f * static_cast(i % 5 + 1); + for (size_t i = 0; i < value_data.size(); i++) value_data[i] = 0.3f * static_cast(i % 3 + 1); + + auto output_with_present = RunGQAAndGetOutput( + batch_size, sequence_length, query_data, key_data, value_data, + num_heads, kv_num_heads, head_size, /*omit_present=*/false); + + auto output_without_present = RunGQAAndGetOutput( + batch_size, sequence_length, query_data, key_data, value_data, + num_heads, kv_num_heads, head_size, /*omit_present=*/true); + + ASSERT_EQ(output_with_present.size(), output_without_present.size()); + for (size_t i = 0; i < output_with_present.size(); i++) { + EXPECT_NEAR(output_with_present[i], output_without_present[i], 1e-5f) + << "Output mismatch at index " << i; + } + + // Sanity: output should not be all zeros (proves the kernel actually computed something) + bool all_zero = true; + for (float v : output_with_present) { + if (v != 0.0f) { + all_zero = false; + break; + } + } + EXPECT_FALSE(all_zero) << "Output should not be all zeros"; } -// Batched: present outputs omitted with batch_size > 1 -TEST(GroupQueryAttentionTest, OptionalPresent_Batched) { - RunGQAOptionalPresentTest( - /*batch_size=*/2, - /*sequence_length=*/3, - /*total_seq_len=*/3, - /*omit_present=*/true, - OpTester::ExpectResult::kExpectSuccess, - ""); +// Batched: same correctness check with batch_size > 1 +TEST(GroupQueryAttentionTest, OptionalPresent_BatchedOutputMatch) { + constexpr int batch_size = 2; + constexpr int sequence_length = 3; + constexpr int num_heads = 2; + constexpr int kv_num_heads = 1; + constexpr int head_size = 8; + constexpr int hidden_size = num_heads * head_size; + constexpr int kv_hidden_size = kv_num_heads * head_size; + + std::vector query_data(batch_size * sequence_length * hidden_size); + std::vector key_data(batch_size * sequence_length * kv_hidden_size); + std::vector value_data(batch_size * sequence_length * kv_hidden_size); + for (size_t i = 0; i < query_data.size(); i++) query_data[i] = 0.15f * static_cast(i % 11 + 1); + for (size_t i = 0; i < key_data.size(); i++) key_data[i] = 0.25f * static_cast(i % 7 + 1); + for (size_t i = 0; i < value_data.size(); i++) value_data[i] = 0.35f * static_cast(i % 5 + 1); + + auto output_with = RunGQAAndGetOutput( + batch_size, sequence_length, query_data, key_data, value_data, + num_heads, kv_num_heads, head_size, /*omit_present=*/false); + + auto output_without = RunGQAAndGetOutput( + batch_size, sequence_length, query_data, key_data, value_data, + num_heads, kv_num_heads, head_size, /*omit_present=*/true); + + ASSERT_EQ(output_with.size(), output_without.size()); + for (size_t i = 0; i < output_with.size(); i++) { + EXPECT_NEAR(output_with[i], output_without[i], 1e-5f) + << "Batched output mismatch at index " << i; + } } // Reject: omitting present outputs when total_seq_len > sequence_length (decode with past) TEST(GroupQueryAttentionTest, OptionalPresent_RejectWithPast) { - RunGQAOptionalPresentTest( - /*batch_size=*/1, - /*sequence_length=*/1, - /*total_seq_len=*/5, - /*omit_present=*/true, - OpTester::ExpectResult::kExpectFailure, - "present_key and present_value outputs are required when past state exists"); + constexpr int batch_size = 1; + constexpr int sequence_length = 1; + constexpr int total_seq_len = 5; + constexpr int num_heads = 2; + constexpr int kv_num_heads = 1; + constexpr int head_size = 8; + constexpr int hidden_size = num_heads * head_size; + constexpr int kv_hidden_size = kv_num_heads * head_size; + + OpTester tester("GroupQueryAttention", 1, onnxruntime::kMSDomain); + tester.AddAttribute("num_heads", static_cast(num_heads)); + tester.AddAttribute("kv_num_heads", static_cast(kv_num_heads)); + + tester.AddInput("query", {batch_size, sequence_length, hidden_size}, + std::vector(batch_size * sequence_length * hidden_size, 1.0f)); + tester.AddInput("key", {batch_size, sequence_length, kv_hidden_size}, + std::vector(batch_size * sequence_length * kv_hidden_size, 0.5f)); + tester.AddInput("value", {batch_size, sequence_length, kv_hidden_size}, + std::vector(batch_size * sequence_length * kv_hidden_size, 0.5f)); + + tester.AddOptionalInputEdge(); // past_key + tester.AddOptionalInputEdge(); // past_value + + tester.AddInput("seqlens_k", {batch_size}, {static_cast(total_seq_len - 1)}); + tester.AddInput("total_sequence_length", {1}, {static_cast(total_seq_len)}); + + tester.AddOptionalInputEdge(); // cos_cache + tester.AddOptionalInputEdge(); // sin_cache + tester.AddOptionalInputEdge(); // position_ids + tester.AddOptionalInputEdge(); // attention_bias + tester.AddOptionalInputEdge(); // head_sink + + tester.AddOutput("output", {batch_size, sequence_length, hidden_size}, + std::vector(batch_size * sequence_length * hidden_size, 0.0f)); + tester.AddOptionalOutputEdge(); // present_key — omitted + tester.AddOptionalOutputEdge(); // present_value — omitted + + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectFailure, + "present_key and present_value outputs are required when past state exists", + {}, nullptr, &execution_providers); } } // namespace test From 2ef269cd42f783c885408089bd01ba391ce89527 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Fri, 1 May 2026 12:39:34 -0700 Subject: [PATCH 12/32] Address comments --- .../contrib_ops/cpu/bert/gqa_attention_base.h | 2 +- .../cpu/bert/group_query_attention.cc | 16 +++-- .../cuda/bert/group_query_attention.cc | 12 +++- .../webgpu/bert/group_query_attention.cc | 3 +- .../group_query_attention_op_test.cc | 63 ++++++++++++++++--- 5 files changed, 80 insertions(+), 16 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h b/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h index 1f03cf9f105a2..60676f11acd3f 100644 --- a/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h +++ b/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h @@ -87,7 +87,7 @@ class GQAAttentionBase { } int seqlen_present_kv_cache = present_key != nullptr ? static_cast(present_key->Shape().GetDims()[2]) - : parameters.seqlen_present_kv_cache; + : parameters.total_sequence_length; // Compute the attention score. bool gqa_mlas_supported = MlasGQASupported(CblasNoTrans, CblasTrans) && diff --git a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc index 44d3895abf672..8de2088a37bd8 100644 --- a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc @@ -113,14 +113,20 @@ Status GroupQueryAttention::Compute(OpKernelContext* context) const { Tensor* present_k = context->Output(1, present_k_shape); Tensor* present_v = context->Output(2, present_v_shape); - // Optional present outputs are only safe for first-prompt with no past KV. - // When past exists or total_sequence_length > sequence_length, the attention - // GEMMs use total_seqlen which requires a concatenated past+current KV buffer - // that only ConcatStateChunkGQA builds into present_key/present_value. + // present_key and present_value must be both present or both absent. + if ((present_k == nullptr) != (present_v == nullptr)) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "present_key and present_value must be both provided or both omitted."); + } + + // Optional present outputs are only safe when is_first_prompt + // (sequence_length == total_sequence_length, i.e., no past KV to concatenate). + // When past exists, the attention GEMMs use total_seqlen which requires a + // concatenated past+current KV buffer built by ConcatStateChunkGQA into present. if ((present_k == nullptr || present_v == nullptr) && !parameters.is_first_prompt) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "present_key and present_value outputs are required when past state exists " - "(total_sequence_length > sequence_length). Omitting present outputs is only " + "(sequence_length != total_sequence_length). Omitting present outputs is only " "supported for first-prompt inference with no past KV cache."); } diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc index 97df349656c07..7e306428546fa 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc @@ -267,10 +267,18 @@ Status GroupQueryAttention::ComputeInternal(OpKernelContext* context) cons Tensor* present_key_output = context->Output(1, present_shape); // present_key Tensor* present_value_output = context->Output(2, present_shape); // present_value - // Optional present outputs are only safe for first-prompt with no past KV. + // present_key and present_value must be both present or both absent. + if ((present_key_output == nullptr) != (present_value_output == nullptr)) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "present_key and present_value must be both provided or both omitted."); + } + + // Optional present outputs are only safe when is_first_prompt + // (sequence_length == total_sequence_length, i.e., no past KV to concatenate). if ((present_key_output == nullptr || present_value_output == nullptr) && !parameters.is_first_prompt) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "present_key and present_value outputs are required when past state exists. " + "present_key and present_value outputs are required when past state exists " + "(sequence_length != total_sequence_length). " "Omitting present outputs is only supported for first-prompt inference."); } diff --git a/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc b/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc index 5fff0516c7ce3..b8a618857adb1 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc @@ -8,6 +8,7 @@ #include "contrib_ops/webgpu/bert/rotary_embedding.h" #include "contrib_ops/webgpu/bert/flash_attention.h" +#include "core/common/narrow.h" #include "core/providers/webgpu/webgpu_supported_types.h" #include "core/providers/webgpu/shader_helper.h" @@ -212,7 +213,7 @@ Status GroupQueryAttention::ComputeInternal(onnxruntime::webgpu::ComputeContext& scale_, softcap_, 0, - static_cast(context.DeviceLimits().maxComputeInvocationsPerWorkgroup))); + onnxruntime::narrow(context.DeviceLimits().maxComputeInvocationsPerWorkgroup))); params.use_smooth_softmax = use_smooth_softmax_; params.rotary_interleaved = rotary_interleaved_; diff --git a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc index 9b0940a96d6e4..69f72680ca345 100644 --- a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc @@ -313,6 +313,7 @@ TEST(GroupQueryAttentionTest, SeqlensKWrongLength) { // Run GQA with the given inputs and return the output tensor as a vector. // This lets us compare outputs between present-connected and present-omitted runs. +// When use_cuda=true, runs on CUDA EP instead of CPU EP. static std::vector RunGQAAndGetOutput( int batch_size, int sequence_length, @@ -322,7 +323,8 @@ static std::vector RunGQAAndGetOutput( int num_heads, int kv_num_heads, int head_size, - bool omit_present) { + bool omit_present, + bool use_cuda = false) { const int hidden_size = num_heads * head_size; const int kv_hidden_size = kv_num_heads * head_size; const int total_seq_len = sequence_length; // first-prompt: no past @@ -367,7 +369,11 @@ static std::vector RunGQAAndGetOutput( tester.SetOutputTolerance(1e6f); // We compare fetched outputs ourselves std::vector> execution_providers; - execution_providers.push_back(DefaultCpuExecutionProvider()); + if (use_cuda) { + execution_providers.push_back(DefaultCudaExecutionProvider()); + } else { + execution_providers.push_back(DefaultCpuExecutionProvider()); + } tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); // Extract the output tensor values @@ -376,9 +382,9 @@ static std::vector RunGQAAndGetOutput( return std::vector(out_data, out_data + output_size); } -// Core correctness test: output must be identical whether present outputs -// are connected or omitted (first-prompt, no past KV). -TEST(GroupQueryAttentionTest, OptionalPresent_OutputMatchesWithAndWithout) { +// Regression: omitting optional present outputs must not change the attention output +// compared to when present outputs are connected (first-prompt, no past KV). +TEST(GroupQueryAttentionTest, OptionalPresent_OmittingDoesNotChangeOutput) { constexpr int batch_size = 1; constexpr int sequence_length = 4; constexpr int num_heads = 2; @@ -420,8 +426,8 @@ TEST(GroupQueryAttentionTest, OptionalPresent_OutputMatchesWithAndWithout) { EXPECT_FALSE(all_zero) << "Output should not be all zeros"; } -// Batched: same correctness check with batch_size > 1 -TEST(GroupQueryAttentionTest, OptionalPresent_BatchedOutputMatch) { +// Regression (batched): same equivalence check with batch_size > 1 +TEST(GroupQueryAttentionTest, OptionalPresent_BatchedOmitMatchesConnected) { constexpr int batch_size = 2; constexpr int sequence_length = 3; constexpr int num_heads = 2; @@ -498,5 +504,48 @@ TEST(GroupQueryAttentionTest, OptionalPresent_RejectWithPast) { {}, nullptr, &execution_providers); } +// Regression (CUDA): omitting present outputs on CUDA EP must produce the same +// attention output as when present outputs are connected. The CUDA path allocates +// internal scratch buffers to serve as KV workspace for flash/MEA/unfused kernels. +TEST(GroupQueryAttentionTest, OptionalPresent_CudaOmitMatchesConnected) { + constexpr int batch_size = 1; + constexpr int sequence_length = 4; + constexpr int num_heads = 2; + constexpr int kv_num_heads = 1; + constexpr int head_size = 8; + constexpr int hidden_size = num_heads * head_size; + constexpr int kv_hidden_size = kv_num_heads * head_size; + + std::vector query_data(batch_size * sequence_length * hidden_size); + std::vector key_data(batch_size * sequence_length * kv_hidden_size); + std::vector value_data(batch_size * sequence_length * kv_hidden_size); + for (size_t i = 0; i < query_data.size(); i++) query_data[i] = 0.1f * static_cast(i % 7 + 1); + for (size_t i = 0; i < key_data.size(); i++) key_data[i] = 0.2f * static_cast(i % 5 + 1); + for (size_t i = 0; i < value_data.size(); i++) value_data[i] = 0.3f * static_cast(i % 3 + 1); + + auto output_with = RunGQAAndGetOutput( + batch_size, sequence_length, query_data, key_data, value_data, + num_heads, kv_num_heads, head_size, /*omit_present=*/false, /*use_cuda=*/true); + + auto output_without = RunGQAAndGetOutput( + batch_size, sequence_length, query_data, key_data, value_data, + num_heads, kv_num_heads, head_size, /*omit_present=*/true, /*use_cuda=*/true); + + ASSERT_EQ(output_with.size(), output_without.size()); + for (size_t i = 0; i < output_with.size(); i++) { + EXPECT_NEAR(output_with[i], output_without[i], 1e-5f) + << "CUDA output mismatch at index " << i; + } + + bool all_zero = true; + for (float v : output_with) { + if (v != 0.0f) { + all_zero = false; + break; + } + } + EXPECT_FALSE(all_zero) << "CUDA output should not be all zeros"; +} + } // namespace test } // namespace onnxruntime From 6cbe62cbc10762e6ef53d05278607c9400f16ac8 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Fri, 1 May 2026 13:19:48 -0700 Subject: [PATCH 13/32] Fix unit tests --- .../test/contrib_ops/group_query_attention_op_test.cc | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc index 69f72680ca345..8c158366c47e7 100644 --- a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc @@ -508,6 +508,11 @@ TEST(GroupQueryAttentionTest, OptionalPresent_RejectWithPast) { // attention output as when present outputs are connected. The CUDA path allocates // internal scratch buffers to serve as KV workspace for flash/MEA/unfused kernels. TEST(GroupQueryAttentionTest, OptionalPresent_CudaOmitMatchesConnected) { + auto cuda_ep = DefaultCudaExecutionProvider(); + if (!cuda_ep) { + GTEST_SKIP() << "CUDA EP not available"; + } + constexpr int batch_size = 1; constexpr int sequence_length = 4; constexpr int num_heads = 2; From 3db82c6b3f3709630abd9f427ffd022c32b31660 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Fri, 1 May 2026 13:52:27 -0700 Subject: [PATCH 14/32] Fix comments --- .../contrib_ops/webgpu/bert/group_query_attention.cc | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc b/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc index b8a618857adb1..1735c6e1ec3da 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc @@ -240,7 +240,15 @@ Status GroupQueryAttention::ComputeInternal(onnxruntime::webgpu::ComputeContext& std::vector present_kv_shape(present_dims); Tensor* present_key = context.Output(1, present_kv_shape); Tensor* present_value = context.Output(2, present_kv_shape); - parameters.past_present_share_buffer_ = present_key != nullptr && present_value != nullptr && past_key != nullptr && past_value != nullptr && past_key->DataRaw() == present_key->DataRaw() && past_value->DataRaw() == present_value->DataRaw(); + + // WebGPU flash attention requires present_key/present_value as working KV buffers. + if (present_key == nullptr || present_value == nullptr) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "WebGPU GroupQueryAttention requires present_key and present_value outputs. " + "Optional present outputs are supported on CPU and CUDA EPs only."); + } + + parameters.past_present_share_buffer_ = past_key != nullptr && past_value != nullptr && past_key->DataRaw() == present_key->DataRaw() && past_value->DataRaw() == present_value->DataRaw(); ORT_ENFORCE(parameters.total_sequence_length_ <= parameters.seqlen_present_kv_cache_, "Total sequence length cannot be greater than the existing KV cache length."); From bd023f1339f07a87925bf140297ce76bef8ae3be Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 4 May 2026 08:39:48 -0700 Subject: [PATCH 15/32] Address comments --- .../cpu/bert/group_query_attention.cc | 16 +-- .../cuda/bert/group_query_attention.cc | 12 +-- .../group_query_attention_op_test.cc | 100 ++++++++++++++++-- 3 files changed, 108 insertions(+), 20 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc index 8de2088a37bd8..b07ccd371a8ae 100644 --- a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc @@ -119,15 +119,15 @@ Status GroupQueryAttention::Compute(OpKernelContext* context) const { "present_key and present_value must be both provided or both omitted."); } - // Optional present outputs are only safe when is_first_prompt - // (sequence_length == total_sequence_length, i.e., no past KV to concatenate). - // When past exists, the attention GEMMs use total_seqlen which requires a - // concatenated past+current KV buffer built by ConcatStateChunkGQA into present. - if ((present_k == nullptr || present_v == nullptr) && !parameters.is_first_prompt) { + // Omitting present outputs is only safe when past_key is not provided. + // When past_key exists, ConcatStateChunkGQA must build a concatenated + // past+current KV buffer in present_key/present_value for the attention GEMMs. + // KV-shared layers (e.g., Gemma 4) legitimately omit present during decode: + // they receive borrowed KV via key/value inputs with no past of their own. + if ((present_k == nullptr || present_v == nullptr) && past_key != nullptr) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "present_key and present_value outputs are required when past state exists " - "(sequence_length != total_sequence_length). Omitting present outputs is only " - "supported for first-prompt inference with no past KV cache."); + "present_key and present_value outputs are required when past_key is provided. " + "Omitting present outputs is only supported when there is no past KV cache."); } std::vector output_qk_shape{static_cast(batch_size), static_cast(num_heads_), static_cast(parameters.sequence_length), static_cast(parameters.total_sequence_length)}; diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc index 7e306428546fa..1cf98895f6e1b 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc @@ -273,13 +273,13 @@ Status GroupQueryAttention::ComputeInternal(OpKernelContext* context) cons "present_key and present_value must be both provided or both omitted."); } - // Optional present outputs are only safe when is_first_prompt - // (sequence_length == total_sequence_length, i.e., no past KV to concatenate). - if ((present_key_output == nullptr || present_value_output == nullptr) && !parameters.is_first_prompt) { + // Omitting present outputs is only safe when past_key is not provided. + // KV-shared layers (e.g., Gemma 4) omit present during decode: they receive + // borrowed KV via key/value inputs with no past of their own. + if ((present_key_output == nullptr || present_value_output == nullptr) && past_key != nullptr) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "present_key and present_value outputs are required when past state exists " - "(sequence_length != total_sequence_length). " - "Omitting present outputs is only supported for first-prompt inference."); + "present_key and present_value outputs are required when past_key is provided. " + "Omitting present outputs is only supported when there is no past KV cache."); } // When present outputs are omitted, allocate internal scratch buffers so the diff --git a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc index 8c158366c47e7..a4123d733257d 100644 --- a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc @@ -458,11 +458,96 @@ TEST(GroupQueryAttentionTest, OptionalPresent_BatchedOmitMatchesConnected) { } } -// Reject: omitting present outputs when total_seq_len > sequence_length (decode with past) -TEST(GroupQueryAttentionTest, OptionalPresent_RejectWithPast) { +// KV-shared decode: Q has length 1, K/V have full context length, no past, +// present omitted. This is the Gemma 4 KV-shared layer decode scenario. +TEST(GroupQueryAttentionTest, OptionalPresent_KVSharedDecode) { + constexpr int batch_size = 1; + constexpr int sequence_length = 1; // decode: single token query + constexpr int kv_seq_len = 8; // borrowed KV from source layer + constexpr int num_heads = 2; + constexpr int kv_num_heads = 1; + constexpr int head_size = 8; + constexpr int hidden_size = num_heads * head_size; + constexpr int kv_hidden_size = kv_num_heads * head_size; + + std::vector query_data(batch_size * sequence_length * hidden_size); + std::vector key_data(batch_size * kv_seq_len * kv_hidden_size); + std::vector value_data(batch_size * kv_seq_len * kv_hidden_size); + for (size_t i = 0; i < query_data.size(); i++) query_data[i] = 0.1f * static_cast(i % 7 + 1); + for (size_t i = 0; i < key_data.size(); i++) key_data[i] = 0.2f * static_cast(i % 5 + 1); + for (size_t i = 0; i < value_data.size(); i++) value_data[i] = 0.3f * static_cast(i % 3 + 1); + + // Run with present connected + auto run_test = [&](bool omit_present) { + OpTester tester("GroupQueryAttention", 1, onnxruntime::kMSDomain); + tester.AddAttribute("num_heads", static_cast(num_heads)); + tester.AddAttribute("kv_num_heads", static_cast(kv_num_heads)); + + tester.AddInput("query", {batch_size, sequence_length, hidden_size}, query_data); + tester.AddInput("key", {batch_size, kv_seq_len, kv_hidden_size}, key_data); + tester.AddInput("value", {batch_size, kv_seq_len, kv_hidden_size}, value_data); + + tester.AddOptionalInputEdge(); // past_key — none (KV-shared layer) + tester.AddOptionalInputEdge(); // past_value + + tester.AddInput("seqlens_k", {batch_size}, {static_cast(kv_seq_len - 1)}); + tester.AddInput("total_sequence_length", {1}, {static_cast(kv_seq_len)}); + + tester.AddOptionalInputEdge(); // cos_cache + tester.AddOptionalInputEdge(); // sin_cache + tester.AddOptionalInputEdge(); // position_ids + tester.AddOptionalInputEdge(); // attention_bias + tester.AddOptionalInputEdge(); // head_sink + + const int output_size = batch_size * sequence_length * hidden_size; + tester.AddOutput("output", {batch_size, sequence_length, hidden_size}, + std::vector(output_size, 0.0f)); + + if (omit_present) { + tester.AddOptionalOutputEdge(); + tester.AddOptionalOutputEdge(); + } else { + const int present_size = batch_size * kv_num_heads * kv_seq_len * head_size; + tester.AddOutput("present_key", {batch_size, kv_num_heads, kv_seq_len, head_size}, + std::vector(present_size, 0.0f)); + tester.AddOutput("present_value", {batch_size, kv_num_heads, kv_seq_len, head_size}, + std::vector(present_size, 0.0f)); + } + tester.SetOutputTolerance(1e6f); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); + + auto fetches = tester.GetFetches(); + const float* out = fetches[0].Get().Data(); + return std::vector(out, out + output_size); + }; + + auto output_with = run_test(/*omit_present=*/false); + auto output_without = run_test(/*omit_present=*/true); + + ASSERT_EQ(output_with.size(), output_without.size()); + for (size_t i = 0; i < output_with.size(); i++) { + EXPECT_NEAR(output_with[i], output_without[i], 1e-5f) + << "KV-shared decode output mismatch at index " << i; + } + bool all_zero = true; + for (float v : output_with) { + if (v != 0.0f) { + all_zero = false; + break; + } + } + EXPECT_FALSE(all_zero) << "Output should not be all zeros"; +} + +// Reject: omitting present outputs when past_key is provided (KV cache concatenation needed) +TEST(GroupQueryAttentionTest, OptionalPresent_RejectWithPastKey) { constexpr int batch_size = 1; constexpr int sequence_length = 1; - constexpr int total_seq_len = 5; + constexpr int past_seq_len = 4; + constexpr int total_seq_len = past_seq_len + sequence_length; constexpr int num_heads = 2; constexpr int kv_num_heads = 1; constexpr int head_size = 8; @@ -480,8 +565,11 @@ TEST(GroupQueryAttentionTest, OptionalPresent_RejectWithPast) { tester.AddInput("value", {batch_size, sequence_length, kv_hidden_size}, std::vector(batch_size * sequence_length * kv_hidden_size, 0.5f)); - tester.AddOptionalInputEdge(); // past_key - tester.AddOptionalInputEdge(); // past_value + // Provide past_key/past_value — this triggers the rejection when present is omitted + tester.AddInput("past_key", {batch_size, kv_num_heads, past_seq_len, head_size}, + std::vector(batch_size * kv_num_heads * past_seq_len * head_size, 0.3f)); + tester.AddInput("past_value", {batch_size, kv_num_heads, past_seq_len, head_size}, + std::vector(batch_size * kv_num_heads * past_seq_len * head_size, 0.3f)); tester.AddInput("seqlens_k", {batch_size}, {static_cast(total_seq_len - 1)}); tester.AddInput("total_sequence_length", {1}, {static_cast(total_seq_len)}); @@ -500,7 +588,7 @@ TEST(GroupQueryAttentionTest, OptionalPresent_RejectWithPast) { std::vector> execution_providers; execution_providers.push_back(DefaultCpuExecutionProvider()); tester.Run(OpTester::ExpectResult::kExpectFailure, - "present_key and present_value outputs are required when past state exists", + "present_key and present_value outputs are required when past_key is provided", {}, nullptr, &execution_providers); } From f0035aab052e2b63b1d47e667968f444009d4105 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 4 May 2026 09:56:43 -0700 Subject: [PATCH 16/32] Fix unit tests --- .../group_query_attention_op_test.cc | 33 +++++++++---------- 1 file changed, 16 insertions(+), 17 deletions(-) diff --git a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc index a4123d733257d..03dcd2e6ed013 100644 --- a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc @@ -458,12 +458,12 @@ TEST(GroupQueryAttentionTest, OptionalPresent_BatchedOmitMatchesConnected) { } } -// KV-shared decode: Q has length 1, K/V have full context length, no past, -// present omitted. This is the Gemma 4 KV-shared layer decode scenario. -TEST(GroupQueryAttentionTest, OptionalPresent_KVSharedDecode) { +// KV-shared first-prompt: longer sequence with no past, present omitted. +// This simulates a KV-shared layer during prefill where Q, K, V all have +// the full prompt length and no KV cache is maintained. +TEST(GroupQueryAttentionTest, OptionalPresent_KVSharedFirstPrompt) { constexpr int batch_size = 1; - constexpr int sequence_length = 1; // decode: single token query - constexpr int kv_seq_len = 8; // borrowed KV from source layer + constexpr int sequence_length = 8; // full prompt length constexpr int num_heads = 2; constexpr int kv_num_heads = 1; constexpr int head_size = 8; @@ -471,27 +471,26 @@ TEST(GroupQueryAttentionTest, OptionalPresent_KVSharedDecode) { constexpr int kv_hidden_size = kv_num_heads * head_size; std::vector query_data(batch_size * sequence_length * hidden_size); - std::vector key_data(batch_size * kv_seq_len * kv_hidden_size); - std::vector value_data(batch_size * kv_seq_len * kv_hidden_size); + std::vector key_data(batch_size * sequence_length * kv_hidden_size); + std::vector value_data(batch_size * sequence_length * kv_hidden_size); for (size_t i = 0; i < query_data.size(); i++) query_data[i] = 0.1f * static_cast(i % 7 + 1); for (size_t i = 0; i < key_data.size(); i++) key_data[i] = 0.2f * static_cast(i % 5 + 1); for (size_t i = 0; i < value_data.size(); i++) value_data[i] = 0.3f * static_cast(i % 3 + 1); - // Run with present connected auto run_test = [&](bool omit_present) { OpTester tester("GroupQueryAttention", 1, onnxruntime::kMSDomain); tester.AddAttribute("num_heads", static_cast(num_heads)); tester.AddAttribute("kv_num_heads", static_cast(kv_num_heads)); tester.AddInput("query", {batch_size, sequence_length, hidden_size}, query_data); - tester.AddInput("key", {batch_size, kv_seq_len, kv_hidden_size}, key_data); - tester.AddInput("value", {batch_size, kv_seq_len, kv_hidden_size}, value_data); + tester.AddInput("key", {batch_size, sequence_length, kv_hidden_size}, key_data); + tester.AddInput("value", {batch_size, sequence_length, kv_hidden_size}, value_data); - tester.AddOptionalInputEdge(); // past_key — none (KV-shared layer) + tester.AddOptionalInputEdge(); // past_key tester.AddOptionalInputEdge(); // past_value - tester.AddInput("seqlens_k", {batch_size}, {static_cast(kv_seq_len - 1)}); - tester.AddInput("total_sequence_length", {1}, {static_cast(kv_seq_len)}); + tester.AddInput("seqlens_k", {batch_size}, {static_cast(sequence_length - 1)}); + tester.AddInput("total_sequence_length", {1}, {static_cast(sequence_length)}); tester.AddOptionalInputEdge(); // cos_cache tester.AddOptionalInputEdge(); // sin_cache @@ -507,10 +506,10 @@ TEST(GroupQueryAttentionTest, OptionalPresent_KVSharedDecode) { tester.AddOptionalOutputEdge(); tester.AddOptionalOutputEdge(); } else { - const int present_size = batch_size * kv_num_heads * kv_seq_len * head_size; - tester.AddOutput("present_key", {batch_size, kv_num_heads, kv_seq_len, head_size}, + const int present_size = batch_size * kv_num_heads * sequence_length * head_size; + tester.AddOutput("present_key", {batch_size, kv_num_heads, sequence_length, head_size}, std::vector(present_size, 0.0f)); - tester.AddOutput("present_value", {batch_size, kv_num_heads, kv_seq_len, head_size}, + tester.AddOutput("present_value", {batch_size, kv_num_heads, sequence_length, head_size}, std::vector(present_size, 0.0f)); } tester.SetOutputTolerance(1e6f); @@ -530,7 +529,7 @@ TEST(GroupQueryAttentionTest, OptionalPresent_KVSharedDecode) { ASSERT_EQ(output_with.size(), output_without.size()); for (size_t i = 0; i < output_with.size(); i++) { EXPECT_NEAR(output_with[i], output_without[i], 1e-5f) - << "KV-shared decode output mismatch at index " << i; + << "KV-shared first-prompt output mismatch at index " << i; } bool all_zero = true; for (float v : output_with) { From 0afa1c9110fee2191858e0446593848f14658777 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 4 May 2026 11:01:52 -0700 Subject: [PATCH 17/32] fix cuda tests --- onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc index 1cf98895f6e1b..667d61d0db8de 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc @@ -289,8 +289,8 @@ Status GroupQueryAttention::ComputeInternal(OpKernelContext* context) cons IAllocatorUniquePtr present_value_scratch; if (present_key_output == nullptr || present_value_output == nullptr) { size_t present_kv_bytes = present_shape.Size() * sizeof(U); - present_key_scratch = GetScratchBuffer(present_kv_bytes, context->GetComputeStream()); - present_value_scratch = GetScratchBuffer(present_kv_bytes, context->GetComputeStream()); + present_key_scratch = GetScratchBuffer(present_kv_bytes, GetComputeStream(context)); + present_value_scratch = GetScratchBuffer(present_kv_bytes, GetComputeStream(context)); } IAllocatorUniquePtr k_buffer; From ab0ddfba92bf452045227919602e0d2742250427 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 4 May 2026 14:45:17 -0700 Subject: [PATCH 18/32] address comments --- .../contrib_ops/cpu/bert/gqa_attention_base.h | 37 ++++---- .../cpu/bert/group_query_attention.cc | 7 +- .../cpu/bert/group_query_attention_helper.h | 16 ++-- .../cuda/bert/group_query_attention.cc | 4 +- .../cuda/bert/group_query_attention_impl.cu | 58 ++++++++++--- .../group_query_attention_op_test.cc | 84 +++++++++++++++++++ 6 files changed, 163 insertions(+), 43 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h b/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h index 60676f11acd3f..223131c804e07 100644 --- a/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h +++ b/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h @@ -74,6 +74,7 @@ class GQAAttentionBase { const bool is_prompt = parameters.is_first_prompt; const int batch_size = parameters.batch_size; const int sequence_length = parameters.sequence_length; + const int kv_sequence_length = parameters.kv_sequence_length; const int total_sequence_length = parameters.total_sequence_length; const int head_size = parameters.head_size; const int hidden_size = parameters.hidden_size; @@ -112,7 +113,7 @@ class GQAAttentionBase { if (gqa_mlas_supported) { ComputeAttentionProbs(static_cast(attention_probs), Q, k, head_sink, seqlens_k->Data(), attention_bias_data, - batch_size, sequence_length, total_sequence_length, attention_bias_shape, seqlen_past_kv_cache, + batch_size, sequence_length, kv_sequence_length, total_sequence_length, attention_bias_shape, seqlen_past_kv_cache, seqlen_present_kv_cache, head_size, past_key_data, present_key_data, output_qk_buffer, past_present_share_buffer, packed_qkv, is_prompt, tp, allocator); @@ -120,12 +121,12 @@ class GQAAttentionBase { const T* v = packed_qkv ? Q + (num_heads_ + kv_num_heads_) * sequence_length * head_size : V; ComputeVxAttentionScore(output->MutableData(), static_cast(attention_probs), v, seqlens_k->Data(), - batch_size, sequence_length, seqlen_past_kv_cache, seqlen_present_kv_cache, head_size, + batch_size, sequence_length, kv_sequence_length, seqlen_past_kv_cache, seqlen_present_kv_cache, head_size, hidden_size, past_value_data, present_value_data, past_present_share_buffer, packed_qkv, is_prompt, tp, allocator); } else { ComputeAttentionProbs(static_cast(attention_probs), Q, k, head_sink, seqlens_k->Data(), attention_bias_data, - batch_size, sequence_length, total_sequence_length, attention_bias_shape, seqlen_past_kv_cache, + batch_size, sequence_length, kv_sequence_length, total_sequence_length, attention_bias_shape, seqlen_past_kv_cache, seqlen_present_kv_cache, head_size, past_key_data, present_key_data, output_qk_buffer, past_present_share_buffer, packed_qkv, is_prompt, tp, allocator); @@ -133,7 +134,7 @@ class GQAAttentionBase { const T* v = packed_qkv ? Q + (num_heads_ + kv_num_heads_) * sequence_length * head_size : V; ComputeVxAttentionScore(output->MutableData(), static_cast(attention_probs), v, seqlens_k->Data(), - batch_size, sequence_length, seqlen_past_kv_cache, seqlen_present_kv_cache, head_size, + batch_size, sequence_length, kv_sequence_length, seqlen_past_kv_cache, seqlen_present_kv_cache, head_size, hidden_size, past_value_data, present_value_data, past_present_share_buffer, packed_qkv, is_prompt, tp, allocator); } @@ -147,15 +148,16 @@ class GQAAttentionBase { // attention_probs(B, N, S, T) = Softmax(attention_probs) // If T is float32, U is float32. If T is float16, U could be float16 or float32. template - void ComputeAttentionProbs(U* attention_probs, // output buffer with size BxNxSxT - const T* Q, // Q data. Its size is BxNxSxH - const T* K, // k data. Its size is BxNxLxH - const T* head_sink, // for smooth softmax. Its size is N. - const int32_t* seqlens_k, // total - 1 sequence lengths tensor - const T* attention_bias, // optional attention bias - const size_t batch_size, // batch size of self-attention - const size_t sequence_length, // sequence length of self-attention (S) - const size_t total_sequence_length, // total sequence length (T) + void ComputeAttentionProbs(U* attention_probs, + const T* Q, + const T* K, + const T* head_sink, + const int32_t* seqlens_k, + const T* attention_bias, + const size_t batch_size, + const size_t sequence_length, + const size_t kv_sequence_length, + const size_t total_sequence_length, const gsl::span attention_bias_shape, // shape of the attention bias const size_t past_buffer_sequence_length, // sequence length of past state const size_t present_buffer_sequence_length, // sequence length of present state @@ -173,7 +175,7 @@ class GQAAttentionBase { : SafeInt(0); const size_t kv_num_heads_factor = num_heads_ / kv_num_heads_; const size_t q_input_chunk_length = sequence_length * head_size; // S x H - const size_t kv_input_chunk_length = sequence_length * head_size; // L x H + const size_t kv_input_chunk_length = kv_sequence_length * head_size; // L x H const size_t past_buff_chunk_length = past_buffer_sequence_length * head_size; // L x H const size_t present_buff_chunk_length = present_buffer_sequence_length * head_size; // T x H @@ -384,9 +386,10 @@ class GQAAttentionBase { const T* V, // V value with size BxN_kvxSxH const int32_t* seqlens_k, // total - 1 sequence lengths tensor const size_t batch_size, // batch size - const size_t sequence_length, // sequence length + const size_t sequence_length, // sequence length of Q + const size_t kv_sequence_length, // sequence length of K/V input const size_t past_buffer_sequence_length, // sequence length in past state - const size_t present_buffer_sequence_length, // sequence length in past state + const size_t present_buffer_sequence_length, // sequence length in present state const size_t head_size, // head size of Q, K, V const size_t hidden_size, // hidden size of Output const T* past_value, // past value only @@ -400,7 +403,7 @@ class GQAAttentionBase { packed_qkv ? SafeInt(num_heads_ + 2 * kv_num_heads_) * sequence_length * head_size : SafeInt(0); const size_t kv_num_heads_factor = num_heads_ / kv_num_heads_; - const size_t kv_input_chunk_length = sequence_length * head_size; // L x H + const size_t kv_input_chunk_length = kv_sequence_length * head_size; // L x H const size_t past_buff_chunk_length = past_buffer_sequence_length * head_size; // L x H const size_t present_buff_chunk_length = present_buffer_sequence_length * head_size; // T x H diff --git a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc index b07ccd371a8ae..b6b266eeac59a 100644 --- a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc @@ -122,8 +122,6 @@ Status GroupQueryAttention::Compute(OpKernelContext* context) const { // Omitting present outputs is only safe when past_key is not provided. // When past_key exists, ConcatStateChunkGQA must build a concatenated // past+current KV buffer in present_key/present_value for the attention GEMMs. - // KV-shared layers (e.g., Gemma 4) legitimately omit present during decode: - // they receive borrowed KV via key/value inputs with no past of their own. if ((present_k == nullptr || present_v == nullptr) && past_key != nullptr) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "present_key and present_value outputs are required when past_key is provided. " @@ -142,6 +140,7 @@ Status GroupQueryAttention::Compute(OpKernelContext* context) const { OrtValue Q; OrtValue K; OrtValue V; + const int kv_sequence_length = parameters.kv_sequence_length; if (packed_qkv) { ORT_RETURN_IF_ERROR(MaybeTransposeToBNSH( allocator, batch_size, num_heads_ + 2 * kv_num_heads_, sequence_length, head_size, query, Q)); @@ -149,9 +148,9 @@ Status GroupQueryAttention::Compute(OpKernelContext* context) const { ORT_RETURN_IF_ERROR(MaybeTransposeToBNSH( allocator, batch_size, num_heads_, sequence_length, head_size, query, Q)); ORT_RETURN_IF_ERROR(MaybeTransposeToBNSH( - allocator, batch_size, kv_num_heads_, sequence_length, head_size, key, K)); + allocator, batch_size, kv_num_heads_, kv_sequence_length, head_size, key, K)); ORT_RETURN_IF_ERROR(MaybeTransposeToBNSH( - allocator, batch_size, kv_num_heads_, sequence_length, head_size, value, V)); + allocator, batch_size, kv_num_heads_, kv_sequence_length, head_size, value, V)); } OrtValue RotaryQKV; diff --git a/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h b/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h index f65568700c0c9..22b974fe0cbe2 100644 --- a/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h +++ b/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h @@ -14,7 +14,8 @@ namespace group_query_attention_helper { template Status Check_Q_K_V(const T* query, const T* key, const T* value, const int num_heads, const int kv_num_heads, - int& batch_size, int& sequence_length, int& q_hidden_size, int& kv_hidden_size, int& head_size) { + int& batch_size, int& sequence_length, int& kv_sequence_length, + int& q_hidden_size, int& kv_hidden_size, int& head_size) { const auto& query_dims = query->Shape().GetDims(); if (query_dims.size() != 3) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'query' is expected to have 3 dimensions, got ", @@ -40,10 +41,8 @@ Status Check_Q_K_V(const T* query, const T* key, const T* value, const int num_h } else if (query_dims[0] != key_dims[0]) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'query' and 'key' shall have same dim 0 (batch size)"); - } else if (query_dims[1] != key_dims[1]) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "Input 'query' and 'key' shall have same dim 1 (sequence length)"); } + kv_sequence_length = static_cast(key_dims[1]); kv_hidden_size = static_cast(key_dims[2]); if (kv_hidden_size % kv_num_heads != 0) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, @@ -61,9 +60,9 @@ Status Check_Q_K_V(const T* query, const T* key, const T* value, const int num_h } else if (query_dims[0] != value_dims[0]) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'query' and 'value' shall have same dim 0 (batch size)"); - } else if (query_dims[1] != value_dims[1]) { + } else if (key_dims[1] != value_dims[1]) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "Input 'query' and 'value' shall have same dim 1 (sequence length)"); + "Input 'key' and 'value' shall have same dim 1 (sequence length)"); } else if (value_dims[2] != kv_hidden_size) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'value' is expected to have same hidden size as key."); } @@ -239,17 +238,19 @@ Status CheckInputs(const T* query, int batch_size = 0; int sequence_length = 0; + int kv_sequence_length = 0; int q_hidden_size = 0; int kv_hidden_size = 0; int head_size = 0; const bool is_packed_qkv = (key == nullptr); if (!is_packed_qkv) { ORT_RETURN_IF_ERROR(Check_Q_K_V(query, key, value, num_heads, kv_num_heads, batch_size, sequence_length, - q_hidden_size, kv_hidden_size, head_size)); + kv_sequence_length, q_hidden_size, kv_hidden_size, head_size)); } else { qkv_format = QKV_BS3NH; ORT_RETURN_IF_ERROR(Check_QKV(query, value, num_heads, kv_num_heads, batch_size, sequence_length, q_hidden_size, kv_hidden_size, head_size)); + kv_sequence_length = sequence_length; } // Check past-present KV @@ -312,6 +313,7 @@ Status CheckInputs(const T* query, GroupQueryAttentionParameters* output_parameters = reinterpret_cast(parameters); output_parameters->batch_size = batch_size; output_parameters->sequence_length = sequence_length; // sequence length of Q + output_parameters->kv_sequence_length = kv_sequence_length; // sequence length of K/V inputs output_parameters->seqlen_past_kv_cache = past_sequence_length; // max sequence length of past kv tensors output_parameters->seqlen_present_kv_cache = present_sequence_length; // max sequence length of present kv tensors output_parameters->total_sequence_length = total_sequence_length; // total sequence length diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc index 667d61d0db8de..9cd8bff7747af 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc @@ -274,8 +274,7 @@ Status GroupQueryAttention::ComputeInternal(OpKernelContext* context) cons } // Omitting present outputs is only safe when past_key is not provided. - // KV-shared layers (e.g., Gemma 4) omit present during decode: they receive - // borrowed KV via key/value inputs with no past of their own. + // When past_key exists, the kernel must concatenate past+current KV into present. if ((present_key_output == nullptr || present_value_output == nullptr) && past_key != nullptr) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "present_key and present_value outputs are required when past_key is provided. " @@ -284,7 +283,6 @@ Status GroupQueryAttention::ComputeInternal(OpKernelContext* context) cons // When present outputs are omitted, allocate internal scratch buffers so the // CUDA kernels (flash attention, MEA, unfused) have a valid KV workspace. - // This keeps behavior consistent with the CPU EP. IAllocatorUniquePtr present_key_scratch; IAllocatorUniquePtr present_value_scratch; if (present_key_output == nullptr || present_value_output == nullptr) { diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu index ebb6a0b0da215..a70750a1071f3 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu @@ -88,6 +88,7 @@ Status PrepareQKV( const int batch_size = parameters.batch_size; const int sequence_length = parameters.sequence_length; + const int kv_sequence_length = parameters.kv_sequence_length; const int num_heads = parameters.num_heads; const int kv_num_heads = parameters.kv_num_heads; const int head_size = parameters.head_size; @@ -123,18 +124,51 @@ Status PrepareQKV( cudaMemcpyDeviceToDevice, stream)); } - ORT_RETURN_IF_ERROR((LaunchUnpackRoPEAppend( - parameters.is_packed_qkv ? reinterpret_cast(data.query) : nullptr, - parameters.is_packed_qkv ? nullptr : reinterpret_cast(data.query), - parameters.is_packed_qkv ? nullptr : reinterpret_cast(data.key), - parameters.is_packed_qkv ? nullptr : reinterpret_cast(data.value), - q_out, k, v, data.k_scale, data.v_scale, - num_heads, kv_num_heads, head_size, sequence_length, batch_size, - max_cache_length, data.past_seq_lens, - reinterpret_cast(data.cos_cache), reinterpret_cast(data.sin_cache), - parameters.rotary_dim, data.position_ids, parameters.rotary_interleaved, - is_cache_bnsh, parameters.k_quant_type, - stream, max_threads_per_block))); + // When kv_sequence_length differs from sequence_length (KV-shared decode), + // we must call LaunchUnpackRoPEAppend separately for Q and K/V since the kernel + // uses a single sequence_length for its thread grid. + if (kv_sequence_length != sequence_length) { + // Process Q only (sequence_length tokens, no K/V) + ORT_RETURN_IF_ERROR((LaunchUnpackRoPEAppend( + nullptr, // no packed_qkv + reinterpret_cast(data.query), // Q input + nullptr, // no K + nullptr, // no V + q_out, k, v, data.k_scale, data.v_scale, + num_heads, kv_num_heads, head_size, sequence_length, batch_size, + max_cache_length, data.past_seq_lens, + reinterpret_cast(data.cos_cache), reinterpret_cast(data.sin_cache), + parameters.rotary_dim, data.position_ids, parameters.rotary_interleaved, + is_cache_bnsh, parameters.k_quant_type, + stream, max_threads_per_block))); + + // Process K/V only (kv_sequence_length tokens, no Q) + ORT_RETURN_IF_ERROR((LaunchUnpackRoPEAppend( + nullptr, // no packed_qkv + nullptr, // no Q + reinterpret_cast(data.key), // K input + reinterpret_cast(data.value), // V input + nullptr, k, v, data.k_scale, data.v_scale, // no Q output + num_heads, kv_num_heads, head_size, kv_sequence_length, batch_size, + max_cache_length, data.past_seq_lens, + reinterpret_cast(data.cos_cache), reinterpret_cast(data.sin_cache), + parameters.rotary_dim, data.position_ids, parameters.rotary_interleaved, + is_cache_bnsh, parameters.k_quant_type, + stream, max_threads_per_block))); + } else { + ORT_RETURN_IF_ERROR((LaunchUnpackRoPEAppend( + parameters.is_packed_qkv ? reinterpret_cast(data.query) : nullptr, + parameters.is_packed_qkv ? nullptr : reinterpret_cast(data.query), + parameters.is_packed_qkv ? nullptr : reinterpret_cast(data.key), + parameters.is_packed_qkv ? nullptr : reinterpret_cast(data.value), + q_out, k, v, data.k_scale, data.v_scale, + num_heads, kv_num_heads, head_size, sequence_length, batch_size, + max_cache_length, data.past_seq_lens, + reinterpret_cast(data.cos_cache), reinterpret_cast(data.sin_cache), + parameters.rotary_dim, data.position_ids, parameters.rotary_interleaved, + is_cache_bnsh, parameters.k_quant_type, + stream, max_threads_per_block))); + } if (q_out != nullptr) { q = reinterpret_cast(q_out); diff --git a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc index 03dcd2e6ed013..7d1334a22ae1b 100644 --- a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc @@ -639,5 +639,89 @@ TEST(GroupQueryAttentionTest, OptionalPresent_CudaOmitMatchesConnected) { EXPECT_FALSE(all_zero) << "CUDA output should not be all zeros"; } +// KV-shared decode: Q has length 1, K/V have full context length (kv_seq_len=8), +// no past, present omitted. Verifies that the output matches when present outputs +// are connected vs omitted for the decode shape. +TEST(GroupQueryAttentionTest, OptionalPresent_KVSharedDecode) { + constexpr int batch_size = 1; + constexpr int q_seq_len = 1; // decode: single token query + constexpr int kv_seq_len = 8; // borrowed KV from source layer + constexpr int num_heads = 2; + constexpr int kv_num_heads = 1; + constexpr int head_size = 8; + constexpr int hidden_size = num_heads * head_size; + constexpr int kv_hidden_size = kv_num_heads * head_size; + + std::vector query_data(batch_size * q_seq_len * hidden_size); + std::vector key_data(batch_size * kv_seq_len * kv_hidden_size); + std::vector value_data(batch_size * kv_seq_len * kv_hidden_size); + for (size_t i = 0; i < query_data.size(); i++) query_data[i] = 0.1f * static_cast(i % 7 + 1); + for (size_t i = 0; i < key_data.size(); i++) key_data[i] = 0.2f * static_cast(i % 5 + 1); + for (size_t i = 0; i < value_data.size(); i++) value_data[i] = 0.3f * static_cast(i % 3 + 1); + + auto run_test = [&](bool omit_present) { + OpTester tester("GroupQueryAttention", 1, onnxruntime::kMSDomain); + tester.AddAttribute("num_heads", static_cast(num_heads)); + tester.AddAttribute("kv_num_heads", static_cast(kv_num_heads)); + + tester.AddInput("query", {batch_size, q_seq_len, hidden_size}, query_data); + tester.AddInput("key", {batch_size, kv_seq_len, kv_hidden_size}, key_data); + tester.AddInput("value", {batch_size, kv_seq_len, kv_hidden_size}, value_data); + + tester.AddOptionalInputEdge(); // past_key + tester.AddOptionalInputEdge(); // past_value + + tester.AddInput("seqlens_k", {batch_size}, {static_cast(kv_seq_len - 1)}); + tester.AddInput("total_sequence_length", {1}, {static_cast(kv_seq_len)}); + + tester.AddOptionalInputEdge(); // cos_cache + tester.AddOptionalInputEdge(); // sin_cache + tester.AddOptionalInputEdge(); // position_ids + tester.AddOptionalInputEdge(); // attention_bias + tester.AddOptionalInputEdge(); // head_sink + + const int output_size = batch_size * q_seq_len * hidden_size; + tester.AddOutput("output", {batch_size, q_seq_len, hidden_size}, + std::vector(output_size, 0.0f)); + + if (omit_present) { + tester.AddOptionalOutputEdge(); + tester.AddOptionalOutputEdge(); + } else { + const int present_size = batch_size * kv_num_heads * kv_seq_len * head_size; + tester.AddOutput("present_key", {batch_size, kv_num_heads, kv_seq_len, head_size}, + std::vector(present_size, 0.0f)); + tester.AddOutput("present_value", {batch_size, kv_num_heads, kv_seq_len, head_size}, + std::vector(present_size, 0.0f)); + } + tester.SetOutputTolerance(1e6f); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); + + auto fetches = tester.GetFetches(); + const float* out = fetches[0].Get().Data(); + return std::vector(out, out + output_size); + }; + + auto output_with = run_test(/*omit_present=*/false); + auto output_without = run_test(/*omit_present=*/true); + + ASSERT_EQ(output_with.size(), output_without.size()); + for (size_t i = 0; i < output_with.size(); i++) { + EXPECT_NEAR(output_with[i], output_without[i], 1e-5f) + << "KV-shared decode output mismatch at index " << i; + } + bool all_zero = true; + for (float v : output_with) { + if (v != 0.0f) { + all_zero = false; + break; + } + } + EXPECT_FALSE(all_zero) << "Output should not be all zeros"; +} + } // namespace test } // namespace onnxruntime From ec507310858a6290234fa736920930233bff17a1 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 4 May 2026 19:27:55 -0700 Subject: [PATCH 19/32] Support KV-shared decode with separate Q/KV sequence lengths --- .../contrib_ops/cpu/bert/gqa_attention_base.h | 5 +- .../cpu/bert/group_query_attention.cc | 6 + .../cuda/bert/group_query_attention.cc | 5 +- .../cuda/bert/group_query_attention_impl.cu | 68 +++-- .../group_query_attention_op_test.cc | 268 ++++++------------ 5 files changed, 145 insertions(+), 207 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h b/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h index 223131c804e07..46c366896b463 100644 --- a/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h +++ b/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h @@ -211,7 +211,8 @@ class GQAAttentionBase { const size_t batch_index = i / num_heads_; const size_t head_index = i % num_heads_; const size_t total_seqlen = SafeInt(seqlens_k[batch_index]) + 1; - const size_t past_seqlen = is_prompt ? 0 : total_seqlen - sequence_length; // Assume no padding sequence length + // past_seqlen is 0 when there is no past (first prompt or KV-shared with no past_key). + const size_t past_seqlen = (is_prompt || past_key == nullptr) ? 0 : total_seqlen - sequence_length; const size_t past_chunk_length = SafeInt(past_seqlen) * head_size; const ptrdiff_t output_offset = SafeInt(i) * sequence_length * present_buffer_sequence_length; @@ -446,7 +447,7 @@ class GQAAttentionBase { const size_t batch_index = i / num_heads_; const size_t head_index = i % num_heads_; const size_t total_seqlen = SafeInt(seqlens_k[batch_index]) + 1; - const size_t past_seqlen = is_prompt ? 0 : total_seqlen - sequence_length; // Assume no padding sequence length + const size_t past_seqlen = (is_prompt || past_value == nullptr) ? 0 : total_seqlen - sequence_length; const size_t past_chunk_length = SafeInt(past_seqlen) * head_size; const T* v; diff --git a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc index b6b266eeac59a..76be4b3b50ff9 100644 --- a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc @@ -159,6 +159,12 @@ Status GroupQueryAttention::Compute(OpKernelContext* context) const { T* q_rotary = Q.GetMutable()->MutableData(); T* k_rotary = packed_qkv ? nullptr : K.GetMutable()->MutableData(); if (do_rotary_) { + // KV-shared decode: K/V already have RoPE from the source layer. + if (kv_sequence_length != sequence_length) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "do_rotary is not supported when query and key have different sequence lengths. " + "Apply RoPE externally before the GQA op for KV-shared layers."); + } ORT_ENFORCE(cos_cache != nullptr && sin_cache != nullptr, "cos_cache and sin_cache must be provided when do_rotary is true"); // Initialize rotary parameters rotary_embedding_helper::RotaryParameters rotary_params = {}; diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc index 9cd8bff7747af..368e71f4e8262 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc @@ -477,13 +477,16 @@ Status GroupQueryAttention::ComputeInternal(OpKernelContext* context) cons data.past_seq_lens = seq_lens_buffer.get(); data.total_seq_lens = seq_lens_buffer.get() + parameters.batch_size; data.padded_seq_lens = data.total_seq_lens + parameters.batch_size; + // For KV-shared decode (no past_key but not first_prompt), treat as first_prompt + // for sequence length computation so past_seq_lens = 0 (no past to offset from). + bool effective_is_first_prompt = parameters.is_first_prompt || (past_key == nullptr); ORT_RETURN_IF_ERROR(LaunchGetSequenceLengths(total_seq_lens_minus_one->Data(), data.past_seq_lens, data.total_seq_lens, data.padded_seq_lens, parameters.batch_size, parameters.sequence_length, - parameters.is_first_prompt, + effective_is_first_prompt, cuda_stream, device_prop.maxThreadsPerBlock)); DUMP_TENSOR_INIT(); diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu index a70750a1071f3..7bbc23d1d143b 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu @@ -103,7 +103,10 @@ Status PrepareQKV( U* v = reinterpret_cast(data.present_value); int max_cache_length = parameters.seqlen_present_kv_cache; - if (!parameters.past_present_share_buffer) { + if (!parameters.past_present_share_buffer && kv_sequence_length != sequence_length) { + // KV-shared decode: Transpose_BSNH_to_BNSH will write every element of the + // present buffer, so skip the memset to save a kernel launch. + } else if (!parameters.past_present_share_buffer) { size_t kv_buffer_size = (size_t)batch_size * kv_num_heads * max_cache_length * head_size * sizeof(U); CUDA_CALL_THROW(cudaMemsetAsync(data.present_key, 0, kv_buffer_size, stream)); CUDA_CALL_THROW(cudaMemsetAsync(data.present_value, 0, kv_buffer_size, stream)); @@ -125,36 +128,43 @@ Status PrepareQKV( } // When kv_sequence_length differs from sequence_length (KV-shared decode), - // we must call LaunchUnpackRoPEAppend separately for Q and K/V since the kernel - // uses a single sequence_length for its thread grid. + // K/V are borrowed from a source layer with the full context length and already + // have RoPE applied. We transpose Q and K/V separately via Transpose_BSNH_to_BNSH + // since they have different sequence lengths. if (kv_sequence_length != sequence_length) { - // Process Q only (sequence_length tokens, no K/V) - ORT_RETURN_IF_ERROR((LaunchUnpackRoPEAppend( - nullptr, // no packed_qkv - reinterpret_cast(data.query), // Q input - nullptr, // no K - nullptr, // no V - q_out, k, v, data.k_scale, data.v_scale, - num_heads, kv_num_heads, head_size, sequence_length, batch_size, - max_cache_length, data.past_seq_lens, - reinterpret_cast(data.cos_cache), reinterpret_cast(data.sin_cache), - parameters.rotary_dim, data.position_ids, parameters.rotary_interleaved, - is_cache_bnsh, parameters.k_quant_type, - stream, max_threads_per_block))); + // KV-shared decode does not support do_rotary or packed QKV — RoPE is applied + // externally before the GQA op, and Q/K/V are separate inputs. + if (parameters.do_rotary) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "do_rotary is not supported when query and key have different sequence lengths. " + "Apply RoPE externally before the GQA op for KV-shared layers."); + } + if (parameters.is_packed_qkv) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Packed QKV is not supported when query and key have different sequence lengths."); + } - // Process K/V only (kv_sequence_length tokens, no Q) - ORT_RETURN_IF_ERROR((LaunchUnpackRoPEAppend( - nullptr, // no packed_qkv - nullptr, // no Q - reinterpret_cast(data.key), // K input - reinterpret_cast(data.value), // V input - nullptr, k, v, data.k_scale, data.v_scale, // no Q output - num_heads, kv_num_heads, head_size, kv_sequence_length, batch_size, - max_cache_length, data.past_seq_lens, - reinterpret_cast(data.cos_cache), reinterpret_cast(data.sin_cache), - parameters.rotary_dim, data.position_ids, parameters.rotary_interleaved, - is_cache_bnsh, parameters.k_quant_type, - stream, max_threads_per_block))); + // Q: use directly from input (already BSNH, no rotary needed) + // q_out is nullptr (no rotary, no packed), so q will point to data.query (set below) + + // K/V: transpose BSNH → BNSH directly into present buffer at offset 0. + // No RoPE needed (already applied by source layer), no append offset (no past). + if constexpr (std::is_same::value) { + ORT_RETURN_IF_ERROR((Transpose_BSNH_to_BNSH( + batch_size, kv_sequence_length, kv_num_heads, head_size, + reinterpret_cast(data.key), + reinterpret_cast(data.present_key), + stream, max_threads_per_block))); + ORT_RETURN_IF_ERROR((Transpose_BSNH_to_BNSH( + batch_size, kv_sequence_length, kv_num_heads, head_size, + reinterpret_cast(data.value), + reinterpret_cast(data.present_value), + stream, max_threads_per_block))); + } else { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "KV-shared decode (query_seq_len != kv_seq_len) with quantized KV cache " + "is not supported. Use non-quantized cache for KV-shared layers."); + } } else { ORT_RETURN_IF_ERROR((LaunchUnpackRoPEAppend( parameters.is_packed_qkv ? reinterpret_cast(data.query) : nullptr, diff --git a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc index 7d1334a22ae1b..3758808cb40b9 100644 --- a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc @@ -312,11 +312,12 @@ TEST(GroupQueryAttentionTest, SeqlensKWrongLength) { // ============================================================================ // Run GQA with the given inputs and return the output tensor as a vector. -// This lets us compare outputs between present-connected and present-omitted runs. +// Supports separate Q and K/V sequence lengths for KV-shared decode scenarios. // When use_cuda=true, runs on CUDA EP instead of CPU EP. static std::vector RunGQAAndGetOutput( int batch_size, - int sequence_length, + int q_seq_len, + int kv_seq_len, const std::vector& query_data, const std::vector& key_data, const std::vector& value_data, @@ -327,20 +328,19 @@ static std::vector RunGQAAndGetOutput( bool use_cuda = false) { const int hidden_size = num_heads * head_size; const int kv_hidden_size = kv_num_heads * head_size; - const int total_seq_len = sequence_length; // first-prompt: no past + const int total_seq_len = kv_seq_len; // no past: total = kv length OpTester tester("GroupQueryAttention", 1, onnxruntime::kMSDomain); tester.AddAttribute("num_heads", static_cast(num_heads)); tester.AddAttribute("kv_num_heads", static_cast(kv_num_heads)); - tester.AddInput("query", {batch_size, sequence_length, hidden_size}, query_data); - tester.AddInput("key", {batch_size, sequence_length, kv_hidden_size}, key_data); - tester.AddInput("value", {batch_size, sequence_length, kv_hidden_size}, value_data); + tester.AddInput("query", {batch_size, q_seq_len, hidden_size}, query_data); + tester.AddInput("key", {batch_size, kv_seq_len, kv_hidden_size}, key_data); + tester.AddInput("value", {batch_size, kv_seq_len, kv_hidden_size}, value_data); tester.AddOptionalInputEdge(); // past_key tester.AddOptionalInputEdge(); // past_value - // First-prompt: seqlens_k = total_seq_len - 1 per GQA convention std::vector seqlens_k_data(batch_size, static_cast(total_seq_len - 1)); tester.AddInput("seqlens_k", {batch_size}, seqlens_k_data); tester.AddInput("total_sequence_length", {1}, {static_cast(total_seq_len)}); @@ -351,9 +351,8 @@ static std::vector RunGQAAndGetOutput( tester.AddOptionalInputEdge(); // attention_bias tester.AddOptionalInputEdge(); // head_sink - // Use a placeholder output with large tolerance — we extract the actual values below. - const int output_size = batch_size * sequence_length * hidden_size; - tester.AddOutput("output", {batch_size, sequence_length, hidden_size}, + const int output_size = batch_size * q_seq_len * hidden_size; + tester.AddOutput("output", {batch_size, q_seq_len, hidden_size}, std::vector(output_size, 0.0f)); if (omit_present) { @@ -376,12 +375,28 @@ static std::vector RunGQAAndGetOutput( } tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); - // Extract the output tensor values auto fetches = tester.GetFetches(); const float* out_data = fetches[0].Get().Data(); return std::vector(out_data, out_data + output_size); } +// Helper: compare two output vectors element-wise and check non-zero. +static void ExpectOutputsMatch(const std::vector& a, const std::vector& b, + float tol, const std::string& label) { + ASSERT_EQ(a.size(), b.size()); + for (size_t i = 0; i < a.size(); i++) { + EXPECT_NEAR(a[i], b[i], tol) << label << " mismatch at index " << i; + } + bool all_zero = true; + for (float v : a) { + if (v != 0.0f) { + all_zero = false; + break; + } + } + EXPECT_FALSE(all_zero) << label << " output should not be all zeros"; +} + // Regression: omitting optional present outputs must not change the attention output // compared to when present outputs are connected (first-prompt, no past KV). TEST(GroupQueryAttentionTest, OptionalPresent_OmittingDoesNotChangeOutput) { @@ -402,28 +417,14 @@ TEST(GroupQueryAttentionTest, OptionalPresent_OmittingDoesNotChangeOutput) { for (size_t i = 0; i < value_data.size(); i++) value_data[i] = 0.3f * static_cast(i % 3 + 1); auto output_with_present = RunGQAAndGetOutput( - batch_size, sequence_length, query_data, key_data, value_data, + batch_size, sequence_length, sequence_length, query_data, key_data, value_data, num_heads, kv_num_heads, head_size, /*omit_present=*/false); auto output_without_present = RunGQAAndGetOutput( - batch_size, sequence_length, query_data, key_data, value_data, + batch_size, sequence_length, sequence_length, query_data, key_data, value_data, num_heads, kv_num_heads, head_size, /*omit_present=*/true); - ASSERT_EQ(output_with_present.size(), output_without_present.size()); - for (size_t i = 0; i < output_with_present.size(); i++) { - EXPECT_NEAR(output_with_present[i], output_without_present[i], 1e-5f) - << "Output mismatch at index " << i; - } - - // Sanity: output should not be all zeros (proves the kernel actually computed something) - bool all_zero = true; - for (float v : output_with_present) { - if (v != 0.0f) { - all_zero = false; - break; - } - } - EXPECT_FALSE(all_zero) << "Output should not be all zeros"; + ExpectOutputsMatch(output_with_present, output_without_present, 1e-5f, "OptionalPresent"); } // Regression (batched): same equivalence check with batch_size > 1 @@ -444,26 +445,20 @@ TEST(GroupQueryAttentionTest, OptionalPresent_BatchedOmitMatchesConnected) { for (size_t i = 0; i < value_data.size(); i++) value_data[i] = 0.35f * static_cast(i % 5 + 1); auto output_with = RunGQAAndGetOutput( - batch_size, sequence_length, query_data, key_data, value_data, + batch_size, sequence_length, sequence_length, query_data, key_data, value_data, num_heads, kv_num_heads, head_size, /*omit_present=*/false); auto output_without = RunGQAAndGetOutput( - batch_size, sequence_length, query_data, key_data, value_data, + batch_size, sequence_length, sequence_length, query_data, key_data, value_data, num_heads, kv_num_heads, head_size, /*omit_present=*/true); - ASSERT_EQ(output_with.size(), output_without.size()); - for (size_t i = 0; i < output_with.size(); i++) { - EXPECT_NEAR(output_with[i], output_without[i], 1e-5f) - << "Batched output mismatch at index " << i; - } + ExpectOutputsMatch(output_with, output_without, 1e-5f, "BatchedOptionalPresent"); } // KV-shared first-prompt: longer sequence with no past, present omitted. -// This simulates a KV-shared layer during prefill where Q, K, V all have -// the full prompt length and no KV cache is maintained. TEST(GroupQueryAttentionTest, OptionalPresent_KVSharedFirstPrompt) { constexpr int batch_size = 1; - constexpr int sequence_length = 8; // full prompt length + constexpr int sequence_length = 8; constexpr int num_heads = 2; constexpr int kv_num_heads = 1; constexpr int head_size = 8; @@ -477,68 +472,14 @@ TEST(GroupQueryAttentionTest, OptionalPresent_KVSharedFirstPrompt) { for (size_t i = 0; i < key_data.size(); i++) key_data[i] = 0.2f * static_cast(i % 5 + 1); for (size_t i = 0; i < value_data.size(); i++) value_data[i] = 0.3f * static_cast(i % 3 + 1); - auto run_test = [&](bool omit_present) { - OpTester tester("GroupQueryAttention", 1, onnxruntime::kMSDomain); - tester.AddAttribute("num_heads", static_cast(num_heads)); - tester.AddAttribute("kv_num_heads", static_cast(kv_num_heads)); - - tester.AddInput("query", {batch_size, sequence_length, hidden_size}, query_data); - tester.AddInput("key", {batch_size, sequence_length, kv_hidden_size}, key_data); - tester.AddInput("value", {batch_size, sequence_length, kv_hidden_size}, value_data); - - tester.AddOptionalInputEdge(); // past_key - tester.AddOptionalInputEdge(); // past_value - - tester.AddInput("seqlens_k", {batch_size}, {static_cast(sequence_length - 1)}); - tester.AddInput("total_sequence_length", {1}, {static_cast(sequence_length)}); - - tester.AddOptionalInputEdge(); // cos_cache - tester.AddOptionalInputEdge(); // sin_cache - tester.AddOptionalInputEdge(); // position_ids - tester.AddOptionalInputEdge(); // attention_bias - tester.AddOptionalInputEdge(); // head_sink - - const int output_size = batch_size * sequence_length * hidden_size; - tester.AddOutput("output", {batch_size, sequence_length, hidden_size}, - std::vector(output_size, 0.0f)); - - if (omit_present) { - tester.AddOptionalOutputEdge(); - tester.AddOptionalOutputEdge(); - } else { - const int present_size = batch_size * kv_num_heads * sequence_length * head_size; - tester.AddOutput("present_key", {batch_size, kv_num_heads, sequence_length, head_size}, - std::vector(present_size, 0.0f)); - tester.AddOutput("present_value", {batch_size, kv_num_heads, sequence_length, head_size}, - std::vector(present_size, 0.0f)); - } - tester.SetOutputTolerance(1e6f); - - std::vector> execution_providers; - execution_providers.push_back(DefaultCpuExecutionProvider()); - tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); - - auto fetches = tester.GetFetches(); - const float* out = fetches[0].Get().Data(); - return std::vector(out, out + output_size); - }; - - auto output_with = run_test(/*omit_present=*/false); - auto output_without = run_test(/*omit_present=*/true); + auto output_with = RunGQAAndGetOutput( + batch_size, sequence_length, sequence_length, query_data, key_data, value_data, + num_heads, kv_num_heads, head_size, /*omit_present=*/false); + auto output_without = RunGQAAndGetOutput( + batch_size, sequence_length, sequence_length, query_data, key_data, value_data, + num_heads, kv_num_heads, head_size, /*omit_present=*/true); - ASSERT_EQ(output_with.size(), output_without.size()); - for (size_t i = 0; i < output_with.size(); i++) { - EXPECT_NEAR(output_with[i], output_without[i], 1e-5f) - << "KV-shared first-prompt output mismatch at index " << i; - } - bool all_zero = true; - for (float v : output_with) { - if (v != 0.0f) { - all_zero = false; - break; - } - } - EXPECT_FALSE(all_zero) << "Output should not be all zeros"; + ExpectOutputsMatch(output_with, output_without, 1e-5f, "KVSharedFirstPrompt"); } // Reject: omitting present outputs when past_key is provided (KV cache concatenation needed) @@ -616,36 +557,20 @@ TEST(GroupQueryAttentionTest, OptionalPresent_CudaOmitMatchesConnected) { for (size_t i = 0; i < value_data.size(); i++) value_data[i] = 0.3f * static_cast(i % 3 + 1); auto output_with = RunGQAAndGetOutput( - batch_size, sequence_length, query_data, key_data, value_data, + batch_size, sequence_length, sequence_length, query_data, key_data, value_data, num_heads, kv_num_heads, head_size, /*omit_present=*/false, /*use_cuda=*/true); - auto output_without = RunGQAAndGetOutput( - batch_size, sequence_length, query_data, key_data, value_data, + batch_size, sequence_length, sequence_length, query_data, key_data, value_data, num_heads, kv_num_heads, head_size, /*omit_present=*/true, /*use_cuda=*/true); - ASSERT_EQ(output_with.size(), output_without.size()); - for (size_t i = 0; i < output_with.size(); i++) { - EXPECT_NEAR(output_with[i], output_without[i], 1e-5f) - << "CUDA output mismatch at index " << i; - } - - bool all_zero = true; - for (float v : output_with) { - if (v != 0.0f) { - all_zero = false; - break; - } - } - EXPECT_FALSE(all_zero) << "CUDA output should not be all zeros"; + ExpectOutputsMatch(output_with, output_without, 1e-5f, "CudaOptionalPresent"); } -// KV-shared decode: Q has length 1, K/V have full context length (kv_seq_len=8), -// no past, present omitted. Verifies that the output matches when present outputs -// are connected vs omitted for the decode shape. +// KV-shared decode: Q_seq=1, KV_seq=8, no past, present omitted. TEST(GroupQueryAttentionTest, OptionalPresent_KVSharedDecode) { constexpr int batch_size = 1; - constexpr int q_seq_len = 1; // decode: single token query - constexpr int kv_seq_len = 8; // borrowed KV from source layer + constexpr int q_seq_len = 1; + constexpr int kv_seq_len = 8; constexpr int num_heads = 2; constexpr int kv_num_heads = 1; constexpr int head_size = 8; @@ -659,68 +584,61 @@ TEST(GroupQueryAttentionTest, OptionalPresent_KVSharedDecode) { for (size_t i = 0; i < key_data.size(); i++) key_data[i] = 0.2f * static_cast(i % 5 + 1); for (size_t i = 0; i < value_data.size(); i++) value_data[i] = 0.3f * static_cast(i % 3 + 1); - auto run_test = [&](bool omit_present) { - OpTester tester("GroupQueryAttention", 1, onnxruntime::kMSDomain); - tester.AddAttribute("num_heads", static_cast(num_heads)); - tester.AddAttribute("kv_num_heads", static_cast(kv_num_heads)); + auto output_with = RunGQAAndGetOutput( + batch_size, q_seq_len, kv_seq_len, query_data, key_data, value_data, + num_heads, kv_num_heads, head_size, /*omit_present=*/false); + auto output_without = RunGQAAndGetOutput( + batch_size, q_seq_len, kv_seq_len, query_data, key_data, value_data, + num_heads, kv_num_heads, head_size, /*omit_present=*/true); - tester.AddInput("query", {batch_size, q_seq_len, hidden_size}, query_data); - tester.AddInput("key", {batch_size, kv_seq_len, kv_hidden_size}, key_data); - tester.AddInput("value", {batch_size, kv_seq_len, kv_hidden_size}, value_data); + ExpectOutputsMatch(output_with, output_without, 1e-5f, "KVSharedDecode"); +} - tester.AddOptionalInputEdge(); // past_key - tester.AddOptionalInputEdge(); // past_value +// CUDA KV-shared decode: Q_seq=1, KV_seq=8, no past, present omitted. +// Exercises the CUDA split path (Transpose_BSNH_to_BNSH for K/V, Q used directly). +// CUDA KV-shared decode: Q_seq=1, KV_seq=8, exercises the CUDA Transpose_BSNH_to_BNSH path. +TEST(GroupQueryAttentionTest, OptionalPresent_CudaKVSharedDecode) { + auto cuda_ep = DefaultCudaExecutionProvider(); + if (!cuda_ep) { + GTEST_SKIP() << "CUDA EP not available"; + } - tester.AddInput("seqlens_k", {batch_size}, {static_cast(kv_seq_len - 1)}); - tester.AddInput("total_sequence_length", {1}, {static_cast(kv_seq_len)}); - - tester.AddOptionalInputEdge(); // cos_cache - tester.AddOptionalInputEdge(); // sin_cache - tester.AddOptionalInputEdge(); // position_ids - tester.AddOptionalInputEdge(); // attention_bias - tester.AddOptionalInputEdge(); // head_sink - - const int output_size = batch_size * q_seq_len * hidden_size; - tester.AddOutput("output", {batch_size, q_seq_len, hidden_size}, - std::vector(output_size, 0.0f)); - - if (omit_present) { - tester.AddOptionalOutputEdge(); - tester.AddOptionalOutputEdge(); - } else { - const int present_size = batch_size * kv_num_heads * kv_seq_len * head_size; - tester.AddOutput("present_key", {batch_size, kv_num_heads, kv_seq_len, head_size}, - std::vector(present_size, 0.0f)); - tester.AddOutput("present_value", {batch_size, kv_num_heads, kv_seq_len, head_size}, - std::vector(present_size, 0.0f)); - } - tester.SetOutputTolerance(1e6f); + constexpr int batch_size = 1; + constexpr int q_seq_len = 1; + constexpr int kv_seq_len = 8; + constexpr int num_heads = 2; + constexpr int kv_num_heads = 1; + constexpr int head_size = 8; + constexpr int hidden_size = num_heads * head_size; + constexpr int kv_hidden_size = kv_num_heads * head_size; - std::vector> execution_providers; - execution_providers.push_back(DefaultCpuExecutionProvider()); - tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); + std::vector query_data(batch_size * q_seq_len * hidden_size); + std::vector key_data(batch_size * kv_seq_len * kv_hidden_size); + std::vector value_data(batch_size * kv_seq_len * kv_hidden_size); + for (size_t i = 0; i < query_data.size(); i++) query_data[i] = 0.1f * static_cast(i % 7 + 1); + for (size_t i = 0; i < key_data.size(); i++) key_data[i] = 0.2f * static_cast(i % 5 + 1); + for (size_t i = 0; i < value_data.size(); i++) value_data[i] = 0.3f * static_cast(i % 3 + 1); - auto fetches = tester.GetFetches(); - const float* out = fetches[0].Get().Data(); - return std::vector(out, out + output_size); - }; + // CUDA with vs without present + auto cuda_with = RunGQAAndGetOutput( + batch_size, q_seq_len, kv_seq_len, query_data, key_data, value_data, + num_heads, kv_num_heads, head_size, /*omit_present=*/false, /*use_cuda=*/true); + auto cuda_without = RunGQAAndGetOutput( + batch_size, q_seq_len, kv_seq_len, query_data, key_data, value_data, + num_heads, kv_num_heads, head_size, /*omit_present=*/true, /*use_cuda=*/true); - auto output_with = run_test(/*omit_present=*/false); - auto output_without = run_test(/*omit_present=*/true); + ExpectOutputsMatch(cuda_with, cuda_without, 1e-4f, "CudaKVSharedDecode"); - ASSERT_EQ(output_with.size(), output_without.size()); - for (size_t i = 0; i < output_with.size(); i++) { - EXPECT_NEAR(output_with[i], output_without[i], 1e-5f) - << "KV-shared decode output mismatch at index " << i; - } - bool all_zero = true; - for (float v : output_with) { - if (v != 0.0f) { - all_zero = false; - break; - } + // Cross-check: CUDA should match CPU + auto cpu_with = RunGQAAndGetOutput( + batch_size, q_seq_len, kv_seq_len, query_data, key_data, value_data, + num_heads, kv_num_heads, head_size, /*omit_present=*/false, /*use_cuda=*/false); + + ASSERT_EQ(cuda_with.size(), cpu_with.size()); + for (size_t i = 0; i < cuda_with.size(); i++) { + EXPECT_NEAR(cuda_with[i], cpu_with[i], 1e-4f) + << "CUDA vs CPU KV-shared decode mismatch at index " << i; } - EXPECT_FALSE(all_zero) << "Output should not be all zeros"; } } // namespace test From 1069d5507e7eb74bb3525968dcbb82ab950e5c10 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 4 May 2026 19:48:36 -0700 Subject: [PATCH 20/32] Address copilot comment: --- .../cpu/bert/group_query_attention_helper.h | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h b/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h index 22b974fe0cbe2..44ec5c5080ad6 100644 --- a/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h +++ b/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h @@ -257,6 +257,13 @@ Status CheckInputs(const T* query, int32_t past_sequence_length = 0; if (past_key != nullptr && past_value != nullptr) { ORT_RETURN_IF_ERROR(CheckPast(past_key, past_value, batch_size, kv_num_heads, head_size, kv_cache_bit_width, past_sequence_length)); + // When past KV exists, Q and K/V must have the same sequence length. + // The KV concat/append paths assume sequence_length == kv_sequence_length. + if (kv_sequence_length != sequence_length) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "query and key must have the same sequence length when past_key is provided. " + "Different Q/K sequence lengths are only supported for KV-shared layers with no past."); + } } else if (past_key != nullptr || past_value != nullptr) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'past_key' and 'past_value' shall be both present or both absent."); @@ -280,6 +287,17 @@ Status CheckInputs(const T* query, return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "total_sequence_length must be positive, got ", total_sequence_length, "."); } + + // When there is no past KV (KV-shared / first-prompt), total_sequence_length + // must not exceed kv_sequence_length — the attention kernel reads up to + // total_sequence_length from the K/V buffer which has kv_sequence_length entries. + if (is_total_seqlen_on_cpu && past_key == nullptr && total_sequence_length > kv_sequence_length) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "total_sequence_length (", total_sequence_length, + ") must not exceed kv_sequence_length (", kv_sequence_length, + ") when past_key is not provided."); + } + int present_sequence_length = std::max(total_sequence_length, past_sequence_length); int rotary_dim = 0; From b1c6271d3edf79217deea9fb487b0e63d3bb29bf Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 4 May 2026 20:37:51 -0700 Subject: [PATCH 21/32] fix cuda pipeline --- .../cuda/bert/group_query_attention_impl.cu | 36 +++++++++++++------ 1 file changed, 26 insertions(+), 10 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu index 7bbc23d1d143b..7e9d0684a71fb 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu @@ -149,17 +149,33 @@ Status PrepareQKV( // K/V: transpose BSNH → BNSH directly into present buffer at offset 0. // No RoPE needed (already applied by source layer), no append offset (no past). + // Transpose_BSNH_to_BNSH accepts half/BFloat16/float, not CUDA native types. if constexpr (std::is_same::value) { - ORT_RETURN_IF_ERROR((Transpose_BSNH_to_BNSH( - batch_size, kv_sequence_length, kv_num_heads, head_size, - reinterpret_cast(data.key), - reinterpret_cast(data.present_key), - stream, max_threads_per_block))); - ORT_RETURN_IF_ERROR((Transpose_BSNH_to_BNSH( - batch_size, kv_sequence_length, kv_num_heads, head_size, - reinterpret_cast(data.value), - reinterpret_cast(data.present_value), - stream, max_threads_per_block))); + static_assert(std::is_same::value || std::is_same::value, + "KV-shared decode transpose only supports __half and __nv_bfloat16."); + if constexpr (std::is_same::value) { + ORT_RETURN_IF_ERROR((Transpose_BSNH_to_BNSH( + batch_size, kv_sequence_length, kv_num_heads, head_size, + reinterpret_cast(data.key), + reinterpret_cast(data.present_key), + stream, max_threads_per_block))); + ORT_RETURN_IF_ERROR((Transpose_BSNH_to_BNSH( + batch_size, kv_sequence_length, kv_num_heads, head_size, + reinterpret_cast(data.value), + reinterpret_cast(data.present_value), + stream, max_threads_per_block))); + } else if constexpr (std::is_same::value) { + ORT_RETURN_IF_ERROR((Transpose_BSNH_to_BNSH( + batch_size, kv_sequence_length, kv_num_heads, head_size, + reinterpret_cast(data.key), + reinterpret_cast(data.present_key), + stream, max_threads_per_block))); + ORT_RETURN_IF_ERROR((Transpose_BSNH_to_BNSH( + batch_size, kv_sequence_length, kv_num_heads, head_size, + reinterpret_cast(data.value), + reinterpret_cast(data.present_value), + stream, max_threads_per_block))); + } } else { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "KV-shared decode (query_seq_len != kv_seq_len) with quantized KV cache " From e4281555d9a76cbf8e3baf95cf787d74cf84d2d2 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Wed, 6 May 2026 21:21:35 +0000 Subject: [PATCH 22/32] [GQA] Support KV-shared layers with empty K/V inputs (kv_sequence_length=0) --- .../contrib_ops/cpu/bert/gqa_attention_base.h | 27 ++++++++++++-- .../cpu/bert/group_query_attention.cc | 23 ++++++------ .../cpu/bert/group_query_attention_helper.h | 13 ++++--- .../cuda/bert/group_query_attention_impl.cu | 36 ++++++++++++++++--- 4 files changed, 76 insertions(+), 23 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h b/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h index 46c366896b463..536ca8fd1db22 100644 --- a/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h +++ b/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h @@ -211,8 +211,20 @@ class GQAAttentionBase { const size_t batch_index = i / num_heads_; const size_t head_index = i % num_heads_; const size_t total_seqlen = SafeInt(seqlens_k[batch_index]) + 1; - // past_seqlen is 0 when there is no past (first prompt or KV-shared with no past_key). - const size_t past_seqlen = (is_prompt || past_key == nullptr) ? 0 : total_seqlen - sequence_length; + // Determine how much data comes from the past buffer. + // - Normal prompt (no past): past_seqlen = 0 + // - Normal decode (past exists, new K appended): past_seqlen = total - seq_len + // - Shared KV (kv_sequence_length=0, past has all data): past_seqlen = total + size_t past_seqlen; + if (past_key == nullptr) { + past_seqlen = 0; + } else if (kv_sequence_length == 0) { + past_seqlen = total_seqlen; // All KV data is in past (shared KV) + } else if (is_prompt) { + past_seqlen = 0; + } else { + past_seqlen = total_seqlen - sequence_length; + } const size_t past_chunk_length = SafeInt(past_seqlen) * head_size; const ptrdiff_t output_offset = SafeInt(i) * sequence_length * present_buffer_sequence_length; @@ -447,7 +459,16 @@ class GQAAttentionBase { const size_t batch_index = i / num_heads_; const size_t head_index = i % num_heads_; const size_t total_seqlen = SafeInt(seqlens_k[batch_index]) + 1; - const size_t past_seqlen = (is_prompt || past_value == nullptr) ? 0 : total_seqlen - sequence_length; + size_t past_seqlen; + if (past_value == nullptr) { + past_seqlen = 0; + } else if (kv_sequence_length == 0) { + past_seqlen = total_seqlen; // All KV data is in past (shared KV) + } else if (is_prompt) { + past_seqlen = 0; + } else { + past_seqlen = total_seqlen - sequence_length; + } const size_t past_chunk_length = SafeInt(past_seqlen) * head_size; const T* v; diff --git a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc index 76be4b3b50ff9..1a084af6fda26 100644 --- a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc @@ -159,8 +159,8 @@ Status GroupQueryAttention::Compute(OpKernelContext* context) const { T* q_rotary = Q.GetMutable()->MutableData(); T* k_rotary = packed_qkv ? nullptr : K.GetMutable()->MutableData(); if (do_rotary_) { - // KV-shared decode: K/V already have RoPE from the source layer. - if (kv_sequence_length != sequence_length) { + // KV-shared decode with empty K/V: only apply RoPE to Q, skip K. + if (kv_sequence_length != sequence_length && kv_sequence_length != 0) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "do_rotary is not supported when query and key have different sequence lengths. " "Apply RoPE externally before the GQA op for KV-shared layers."); @@ -222,19 +222,22 @@ Status GroupQueryAttention::Compute(OpKernelContext* context) const { q_rotary = RotaryQ.GetMutable()->MutableData(); k_rotary = RotaryK.GetMutable()->MutableData(); } - // Run rotary embedding for Q and K + // Run rotary embedding for Q ORT_RETURN_IF_ERROR(RunRotaryEmbedding(tp, rotary_params, q_input, pos_ids_data, cos_cache->Data(), sin_cache->Data(), q_rotary, rotary_interleaved_)); - rotary_params.num_heads = kv_num_heads_; - rotary_params.hidden_size = parameters.kv_hidden_size; - if (!packed_qkv) { - rotary_params.batch_stride = kv_num_heads_ * rotary_params.head_stride; + // Run rotary embedding for K (skip when kv_sequence_length == 0, i.e. shared KV with no new tokens) + if (kv_sequence_length > 0) { + rotary_params.num_heads = kv_num_heads_; + rotary_params.hidden_size = parameters.kv_hidden_size; + if (!packed_qkv) { + rotary_params.batch_stride = kv_num_heads_ * rotary_params.head_stride; + } + ORT_RETURN_IF_ERROR(RunRotaryEmbedding(tp, rotary_params, k_input, + pos_ids_data, cos_cache->Data(), + sin_cache->Data(), k_rotary, rotary_interleaved_)); } - ORT_RETURN_IF_ERROR(RunRotaryEmbedding(tp, rotary_params, k_input, - pos_ids_data, cos_cache->Data(), - sin_cache->Data(), k_rotary, rotary_interleaved_)); // Pack V into rotary QKV buffer if (packed_qkv) { const T* v_input = k_input + kv_num_heads_ * sequence_length * head_size; diff --git a/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h b/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h index 44ec5c5080ad6..86c993d13d643 100644 --- a/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h +++ b/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h @@ -257,12 +257,15 @@ Status CheckInputs(const T* query, int32_t past_sequence_length = 0; if (past_key != nullptr && past_value != nullptr) { ORT_RETURN_IF_ERROR(CheckPast(past_key, past_value, batch_size, kv_num_heads, head_size, kv_cache_bit_width, past_sequence_length)); - // When past KV exists, Q and K/V must have the same sequence length. - // The KV concat/append paths assume sequence_length == kv_sequence_length. - if (kv_sequence_length != sequence_length) { + // When past KV exists, Q and K/V must have the same sequence length, + // UNLESS kv_sequence_length is 0 (shared KV: new K/V are empty, past buffer + // already contains the full shared KV cache — no append needed). + if (kv_sequence_length != sequence_length && kv_sequence_length != 0) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "query and key must have the same sequence length when past_key is provided. " - "Different Q/K sequence lengths are only supported for KV-shared layers with no past."); + "query and key must have the same sequence length when past_key is provided, " + "or key sequence length must be 0 for shared KV (no new KV to append). " + "Got sequence_length=", + sequence_length, ", kv_sequence_length=", kv_sequence_length); } } else if (past_key != nullptr || past_value != nullptr) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu index 7e9d0684a71fb..09daaecd77d46 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu @@ -127,11 +127,37 @@ Status PrepareQKV( cudaMemcpyDeviceToDevice, stream)); } - // When kv_sequence_length differs from sequence_length (KV-shared decode), - // K/V are borrowed from a source layer with the full context length and already - // have RoPE applied. We transpose Q and K/V separately via Transpose_BSNH_to_BNSH - // since they have different sequence lengths. - if (kv_sequence_length != sequence_length) { + // Shared KV path: K/V inputs are empty (kv_sequence_length == 0) and the + // past buffer already contains the full shared KV cache. When + // past_present_share_buffer is true the present buffer aliases past, so we + // only need to process Q (apply RoPE if configured) and skip all K/V work. + if (kv_sequence_length == 0) { + if (parameters.do_rotary && data.cos_cache != nullptr && data.sin_cache != nullptr) { + // Apply RoPE to Q only. Launch the kernel with kv_num_heads=0 so that + // only QUERY head threads are spawned — no KEY/VALUE threads at all. + ORT_RETURN_IF_ERROR((LaunchUnpackRoPEAppend( + nullptr, // packed_qkv + reinterpret_cast(data.query), + nullptr, // key (empty) + nullptr, // value (empty) + q_out, + nullptr, // k_cache (unused) + nullptr, // v_cache (unused) + data.k_scale, data.v_scale, + num_heads, 0 /* kv_num_heads=0: no K/V threads */, head_size, sequence_length, batch_size, + max_cache_length, data.past_seq_lens, + reinterpret_cast(data.cos_cache), reinterpret_cast(data.sin_cache), + parameters.rotary_dim, data.position_ids, parameters.rotary_interleaved, + is_cache_bnsh, parameters.k_quant_type, + stream, max_threads_per_block))); + } + // If do_rotary is false, Q is used directly from data.query (q_out == nullptr). + // K/V present buffers already point to the shared past — no work needed. + } else if (kv_sequence_length != sequence_length) { + // When kv_sequence_length differs from sequence_length (KV-shared decode), + // K/V are borrowed from a source layer with the full context length and already + // have RoPE applied. We transpose Q and K/V separately via Transpose_BSNH_to_BNSH + // since they have different sequence lengths. // KV-shared decode does not support do_rotary or packed QKV — RoPE is applied // externally before the GQA op, and Q/K/V are separate inputs. if (parameters.do_rotary) { From 9d4e8c68a6793accb07a0c690c59ce35c70c1884 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Wed, 6 May 2026 23:08:55 +0000 Subject: [PATCH 23/32] Add unit test and fix documentation --- docs/OperatorKernels.md | 31 +- .../contrib_ops/cpu/bert/gqa_attention_base.h | 18 +- .../cuda/bert/group_query_attention_impl.cu | 9 +- .../group_query_attention_op_test.cc | 327 +++++++++++++++++- 4 files changed, 365 insertions(+), 20 deletions(-) diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index 545c4b8a354c7..7596ab7592b25 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -2,6 +2,14 @@ *This file is automatically generated from the registered kernels by [this script](https://github.com/microsoft/onnxruntime/blob/main/tools/python/gen_opkernel_doc.py). Do not modify directly.* +### Version Notation + +The **OpSet Version** column uses the following notation: + +- `N` — registered only for opset N (e.g., `13`). +- `[N, M]` — registered for opsets N through M inclusive (e.g., `[6, 12]`). +- `N+` — registered for opset N and all later opsets until superseded by a newer kernel registration (e.g., `16+`). + ## Execution Providers - [CPUExecutionProvider](#cpuexecutionprovider) @@ -672,7 +680,8 @@ Do not modify directly.* |||14|**T** = tensor(double), tensor(float), tensor(float16)
**U** = tensor(double), tensor(float), tensor(float16)| |||[9, 13]|**T** = tensor(double), tensor(float), tensor(float16)| |||[7, 8]|**T** = tensor(double), tensor(float), tensor(float16)| -|Cast|*in* input:**T1**
*out* output:**T2**|23+|**T1** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(float4e2m1), tensor(float8e4m3fn), tensor(float8e5m2), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)
**T2** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(float4e2m1), tensor(float8e4m3fn), tensor(float8e5m2), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| +|Cast|*in* input:**T1**
*out* output:**T2**|25+|**T1** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(float4e2m1), tensor(float8e4m3fn), tensor(float8e5m2), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)
**T2** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(float4e2m1), tensor(float8e4m3fn), tensor(float8e5m2), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| +|||[23, 24]|**T1** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(float4e2m1), tensor(float8e4m3fn), tensor(float8e5m2), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)
**T2** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(float4e2m1), tensor(float8e4m3fn), tensor(float8e5m2), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |||[21, 22]|**T1** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(float8e4m3fn), tensor(float8e5m2), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)
**T2** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(float4e2m1), tensor(float8e4m3fn), tensor(float8e5m2), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |||[19, 20]|**T1** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(float8e4m3fn), tensor(float8e5m2), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)
**T2** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(float4e2m1), tensor(float8e4m3fn), tensor(float8e5m2), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |||[13, 18]|**T1** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)
**T2** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(float4e2m1), tensor(float8e4m3fn), tensor(float8e5m2), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| @@ -724,7 +733,8 @@ Do not modify directly.* |DynamicSlice|*in* data:**T**
*in* starts:**Tind**
*in* ends:**Tind**
*in* axes:**Tind**
*out* output:**T**|1+|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)
**Tind** = tensor(int32), tensor(int64)| |Einsum|*in* Inputs:**T**
*out* Output:**T**|12+|**T** = tensor(double), tensor(float), tensor(float16)| |Elu|*in* X:**T**
*out* Y:**T**|6+|**T** = tensor(double), tensor(float), tensor(float16)| -|Equal|*in* A:**T**
*in* B:**T**
*out* C:**T1**|13+|**T** = tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int32), tensor(int64), tensor(uint32), tensor(uint64)
**T1** = tensor(bool)| +|Equal|*in* A:**T**
*in* B:**T**
*out* C:**T1**|19+|**T** = tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int32), tensor(int64), tensor(uint32), tensor(uint64)
**T1** = tensor(bool)| +|||[13, 18]|**T** = tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int32), tensor(int64), tensor(uint32), tensor(uint64)
**T1** = tensor(bool)| |||[11, 12]|**T** = tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int32), tensor(int64), tensor(uint32), tensor(uint64)| |||[7, 10]|**T** = tensor(bool), tensor(int32), tensor(int64)| |Erf|*in* input:**T**
*out* output:**T**|13+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)| @@ -874,7 +884,8 @@ Do not modify directly.* |||[13, 18]|**T1** = tensor(float)
**T2** = tensor(int8), tensor(uint8)| |||[10, 12]|**T1** = tensor(float)
**T2** = tensor(int8), tensor(uint8)| |RMSNormalization|*in* X:**T**
*in* scale:**V**
*out* Y:**V**|23+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)
**V** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)| -|RNN|*in* X:**T**
*in* W:**T**
*in* R:**T**
*in* B:**T**
*in* sequence_lens:**T1**
*in* initial_h:**T**
*out* Y:**T**
*out* Y_h:**T**|14+|**T** = tensor(double), tensor(float), tensor(float16)
**T1** = tensor(int32)| +|RNN|*in* X:**T**
*in* W:**T**
*in* R:**T**
*in* B:**T**
*in* sequence_lens:**T1**
*in* initial_h:**T**
*out* Y:**T**
*out* Y_h:**T**|22+|**T** = tensor(double), tensor(float), tensor(float16)
**T1** = tensor(int32)| +|||[14, 21]|**T** = tensor(double), tensor(float), tensor(float16)
**T1** = tensor(int32)| |||[7, 13]|**T** = tensor(double), tensor(float), tensor(float16)
**T1** = tensor(int32)| |RandomNormal|*out* output:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)| |RandomNormalLike|*in* input:**T1**
*out* output:**T2**|1+|**T1** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(string), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)
**T2** = tensor(double), tensor(float), tensor(float16)| @@ -891,11 +902,13 @@ Do not modify directly.* |||[1, 17]|**T** = tensor(double), tensor(float), tensor(float16)| |ReduceLogSumExp|*in* data:**T**
*in* axes:**tensor(int64)**
*out* reduced:**T**

or

*in* data:**T**
*out* reduced:**T**|18+|**T** = tensor(double), tensor(float), tensor(float16)| |||[1, 17]|**T** = tensor(double), tensor(float), tensor(float16)| -|ReduceMax|*in* data:**T**
*in* axes:**tensor(int64)**
*out* reduced:**T**

or

*in* data:**T**
*out* reduced:**T**|18+|**T** = tensor(double), tensor(float), tensor(float16), tensor(int32), tensor(int64)| -|||[1, 17]|**T** = tensor(double), tensor(float), tensor(float16), tensor(int32), tensor(int64)| +|ReduceMax|*in* data:**T**
*in* axes:**tensor(int64)**
*out* reduced:**T**

or

*in* data:**T**
*out* reduced:**T**|20+|**T** = tensor(double), tensor(float), tensor(float16), tensor(int32), tensor(int64), tensor(int8), tensor(uint8)| +|||[18, 19]|**T** = tensor(double), tensor(float), tensor(float16), tensor(int32), tensor(int64), tensor(int8), tensor(uint8)| +|||[1, 17]|**T** = tensor(double), tensor(float), tensor(float16), tensor(int32), tensor(int64), tensor(int8), tensor(uint8)| |ReduceMean|*in* data:**T**
*in* axes:**tensor(int64)**
*out* reduced:**T**

or

*in* data:**T**
*out* reduced:**T**|18+|**T** = tensor(double), tensor(float), tensor(float16), tensor(int32)| |||[1, 17]|**T** = tensor(double), tensor(float), tensor(float16), tensor(int32)| -|ReduceMin|*in* data:**T**
*in* axes:**tensor(int64)**
*out* reduced:**T**

or

*in* data:**T**
*out* reduced:**T**|18+|**T** = tensor(double), tensor(float), tensor(float16), tensor(int32), tensor(int64), tensor(int8), tensor(uint8)| +|ReduceMin|*in* data:**T**
*in* axes:**tensor(int64)**
*out* reduced:**T**

or

*in* data:**T**
*out* reduced:**T**|20+|**T** = tensor(double), tensor(float), tensor(float16), tensor(int32), tensor(int64), tensor(int8), tensor(uint8)| +|||[18, 19]|**T** = tensor(double), tensor(float), tensor(float16), tensor(int32), tensor(int64), tensor(int8), tensor(uint8)| |||[1, 17]|**T** = tensor(double), tensor(float), tensor(float16), tensor(int32), tensor(int64), tensor(int8), tensor(uint8)| |ReduceProd|*in* data:**T**
*in* axes:**tensor(int64)**
*out* reduced:**T**

or

*in* data:**T**
*out* reduced:**T**|18+|**T** = tensor(double), tensor(float), tensor(float16), tensor(int32)| |||[1, 17]|**T** = tensor(double), tensor(float), tensor(float16), tensor(int32)| @@ -906,7 +919,8 @@ Do not modify directly.* |Relu|*in* X:**T**
*out* Y:**T**|14+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)| |||13|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)| |||[6, 12]|**T** = tensor(double), tensor(float), tensor(float16)| -|Reshape|*in* data:**T**
*in* shape:**tensor(int64)**
*out* reshaped:**T**

or

*in* data:**T**
*out* reshaped:**T**|23+|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(float8e4m3fn), tensor(float8e4m3fnuz), tensor(float8e5m2), tensor(float8e5m2fnuz), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)
**shape** = tensor(int64)| +|Reshape|*in* data:**T**
*in* shape:**tensor(int64)**
*out* reshaped:**T**

or

*in* data:**T**
*out* reshaped:**T**|25+|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(float8e4m3fn), tensor(float8e4m3fnuz), tensor(float8e5m2), tensor(float8e5m2fnuz), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)
**shape** = tensor(int64)| +|||[23, 24]|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(float8e4m3fn), tensor(float8e4m3fnuz), tensor(float8e5m2), tensor(float8e5m2fnuz), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)
**shape** = tensor(int64)| |||[21, 22]|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(float8e4m3fn), tensor(float8e4m3fnuz), tensor(float8e5m2), tensor(float8e5m2fnuz), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)
**shape** = tensor(int64)| |||[19, 20]|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(float8e4m3fn), tensor(float8e4m3fnuz), tensor(float8e5m2), tensor(float8e5m2fnuz), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)
**shape** = tensor(int64)| |||[14, 18]|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)
**shape** = tensor(int64)| @@ -923,7 +937,8 @@ Do not modify directly.* |||[16, 21]|**T1** = tensor(double), tensor(float), tensor(float16)
**T2** = tensor(int64)| |||[10, 15]|**T1** = tensor(double), tensor(float)
**T2** = tensor(int64)| |RotaryEmbedding|*in* X:**T**
*in* cos_cache:**T**
*in* sin_cache:**T**
*in* position_ids:**M**
*out* Y:**T**|23+|**M** = tensor(int64)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| -|Round|*in* X:**T**
*out* Y:**T**|11+|**T** = tensor(double), tensor(float), tensor(float16)| +|Round|*in* X:**T**
*out* Y:**T**|22+|**T** = tensor(double), tensor(float), tensor(float16)| +|||[11, 21]|**T** = tensor(double), tensor(float), tensor(float16)| |ScaledTanh|*in* input:**T**
*out* output:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)| |Scan|*in* initial_state_and_scan_inputs:**V**
*out* final_state_and_scan_outputs:**V**

or

*in* sequence_lens:**I**
*in* initial_state_and_scan_inputs:**V**
*out* final_state_and_scan_outputs:**V**|25+|**V** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(float8e4m3fn), tensor(float8e4m3fnuz), tensor(float8e5m2), tensor(float8e5m2fnuz), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |||[23, 24]|**V** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(float8e4m3fn), tensor(float8e4m3fnuz), tensor(float8e5m2), tensor(float8e5m2fnuz), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| diff --git a/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h b/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h index 536ca8fd1db22..33764ed59a6f7 100644 --- a/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h +++ b/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h @@ -211,19 +211,23 @@ class GQAAttentionBase { const size_t batch_index = i / num_heads_; const size_t head_index = i % num_heads_; const size_t total_seqlen = SafeInt(seqlens_k[batch_index]) + 1; - // Determine how much data comes from the past buffer. - // - Normal prompt (no past): past_seqlen = 0 - // - Normal decode (past exists, new K appended): past_seqlen = total - seq_len - // - Shared KV (kv_sequence_length=0, past has all data): past_seqlen = total + // past_seqlen: how much data to copy from past buffer in ConcatStateChunkGQA. + // causal_past_seqlen: offset for causal masking (seq_causal_length = causal_past_seqlen + seq + 1). + // These differ for shared KV prompt: copy all past data, but causal starts at 0. size_t past_seqlen; + size_t causal_past_seqlen; if (past_key == nullptr) { past_seqlen = 0; + causal_past_seqlen = 0; } else if (kv_sequence_length == 0) { - past_seqlen = total_seqlen; // All KV data is in past (shared KV) + past_seqlen = total_seqlen; // Copy all KV data from past (shared KV) + causal_past_seqlen = is_prompt ? 0 : total_seqlen - sequence_length; } else if (is_prompt) { past_seqlen = 0; + causal_past_seqlen = 0; } else { past_seqlen = total_seqlen - sequence_length; + causal_past_seqlen = past_seqlen; } const size_t past_chunk_length = SafeInt(past_seqlen) * head_size; @@ -317,7 +321,7 @@ class GQAAttentionBase { // compute Softmax U* output_softmax = output; for (size_t seq = 0; seq < sequence_length; seq++) { - size_t seq_causal_length = past_seqlen + seq + 1; + size_t seq_causal_length = causal_past_seqlen + seq + 1; const bool should_apply_local_window = local_window_size_ >= 0 && seq_causal_length > static_cast(local_window_size_); @@ -463,7 +467,7 @@ class GQAAttentionBase { if (past_value == nullptr) { past_seqlen = 0; } else if (kv_sequence_length == 0) { - past_seqlen = total_seqlen; // All KV data is in past (shared KV) + past_seqlen = total_seqlen; } else if (is_prompt) { past_seqlen = 0; } else { diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu index 09daaecd77d46..f52a0b0b536a7 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu @@ -103,9 +103,12 @@ Status PrepareQKV( U* v = reinterpret_cast(data.present_value); int max_cache_length = parameters.seqlen_present_kv_cache; - if (!parameters.past_present_share_buffer && kv_sequence_length != sequence_length) { - // KV-shared decode: Transpose_BSNH_to_BNSH will write every element of the - // present buffer, so skip the memset to save a kernel launch. + if (!parameters.past_present_share_buffer && kv_sequence_length != sequence_length && kv_sequence_length > 0) { + // KV-shared decode (kv_seq != q_seq, kv_seq > 0): Transpose_BSNH_to_BNSH + // will write every element of the present buffer, so skip the memset. + // Note: kv_sequence_length == 0 (shared KV with past) does NOT run the + // transpose path — it copies past→present instead, so memset is still needed + // for the region beyond the copied data. } else if (!parameters.past_present_share_buffer) { size_t kv_buffer_size = (size_t)batch_size * kv_num_heads * max_cache_length * head_size * sizeof(U); CUDA_CALL_THROW(cudaMemsetAsync(data.present_key, 0, kv_buffer_size, stream)); diff --git a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc index 3758808cb40b9..db2b22abbd285 100644 --- a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc @@ -595,8 +595,7 @@ TEST(GroupQueryAttentionTest, OptionalPresent_KVSharedDecode) { } // CUDA KV-shared decode: Q_seq=1, KV_seq=8, no past, present omitted. -// Exercises the CUDA split path (Transpose_BSNH_to_BNSH for K/V, Q used directly). -// CUDA KV-shared decode: Q_seq=1, KV_seq=8, exercises the CUDA Transpose_BSNH_to_BNSH path. +// Exercises the CUDA Transpose_BSNH_to_BNSH path and cross-checks against CPU. TEST(GroupQueryAttentionTest, OptionalPresent_CudaKVSharedDecode) { auto cuda_ep = DefaultCudaExecutionProvider(); if (!cuda_ep) { @@ -641,5 +640,329 @@ TEST(GroupQueryAttentionTest, OptionalPresent_CudaKVSharedDecode) { } } +// --------------------------------------------------------------------------- +// Tests for kv_sequence_length=0 with borrowed past_key/past_value +// (Gemma4 shared KV pattern: empty K/V inputs, all KV data in past buffer) +// --------------------------------------------------------------------------- + +// Helper: run GQA with empty K/V and past_key/past_value (shared KV pattern). +// Returns the attention output. +static std::vector RunGQASharedKV( + int batch_size, + int q_seq_len, + int past_seq_len, + const std::vector& query_data, + const std::vector& past_key_data, + const std::vector& past_value_data, + int num_heads, + int kv_num_heads, + int head_size, + bool use_cuda = false) { + const int hidden_size = num_heads * head_size; + const int total_seq_len = past_seq_len; // all KV data is in past + + OpTester tester("GroupQueryAttention", 1, onnxruntime::kMSDomain); + tester.AddAttribute("num_heads", static_cast(num_heads)); + tester.AddAttribute("kv_num_heads", static_cast(kv_num_heads)); + + // Q: [batch, q_seq_len, hidden_size] + tester.AddInput("query", {batch_size, q_seq_len, hidden_size}, query_data); + // K/V: empty [batch, 0, kv_hidden_size] — kv_sequence_length = 0 + const int kv_hidden_size = kv_num_heads * head_size; + tester.AddInput("key", {batch_size, 0, kv_hidden_size}, {}); + tester.AddInput("value", {batch_size, 0, kv_hidden_size}, {}); + + // past_key/past_value: [batch, kv_num_heads, past_seq_len, head_size] BNSH + tester.AddInput("past_key", {batch_size, kv_num_heads, past_seq_len, head_size}, past_key_data); + tester.AddInput("past_value", {batch_size, kv_num_heads, past_seq_len, head_size}, past_value_data); + + std::vector seqlens_k_data(batch_size, static_cast(total_seq_len - 1)); + tester.AddInput("seqlens_k", {batch_size}, seqlens_k_data); + tester.AddInput("total_sequence_length", {1}, {static_cast(total_seq_len)}); + + tester.AddOptionalInputEdge(); // cos_cache + tester.AddOptionalInputEdge(); // sin_cache + tester.AddOptionalInputEdge(); // position_ids + tester.AddOptionalInputEdge(); // attention_bias + tester.AddOptionalInputEdge(); // head_sink + + const int output_size = batch_size * q_seq_len * hidden_size; + tester.AddOutput("output", {batch_size, q_seq_len, hidden_size}, + std::vector(output_size, 0.0f)); + + // present_key/value: required when past is provided + const int present_size = batch_size * kv_num_heads * past_seq_len * head_size; + tester.AddOutput("present_key", {batch_size, kv_num_heads, past_seq_len, head_size}, + std::vector(present_size, 0.0f)); + tester.AddOutput("present_value", {batch_size, kv_num_heads, past_seq_len, head_size}, + std::vector(present_size, 0.0f)); + + tester.SetOutputTolerance(1e6f); // We compare fetched outputs ourselves + + std::vector> execution_providers; + if (use_cuda) { + execution_providers.push_back(DefaultCudaExecutionProvider()); + } else { + execution_providers.push_back(DefaultCpuExecutionProvider()); + } + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); + + auto fetches = tester.GetFetches(); + const float* out_data = fetches[0].Get().Data(); + return std::vector(out_data, out_data + output_size); +} + +// CPU: kv_sequence_length=0 with past_key/past_value (shared KV decode). +// Validates output matches the equivalent non-empty K/V path. +// CPU: kv_sequence_length=0 with past_key/past_value (shared KV decode). +// Validates output is non-zero (attention over past KV produces valid output). +// Note: cannot compare against RunGQAAndGetOutput because the two paths have +// different causal masking semantics (past_seqlen differs). +TEST(GroupQueryAttentionTest, SharedKV_EmptyKV_WithPast_CPU) { + constexpr int batch_size = 1; + constexpr int q_seq_len = 1; + constexpr int past_seq_len = 8; + constexpr int num_heads = 2; + constexpr int kv_num_heads = 1; + constexpr int head_size = 8; + constexpr int hidden_size = num_heads * head_size; + + std::vector query_data(batch_size * q_seq_len * hidden_size); + std::vector past_key_data(batch_size * kv_num_heads * past_seq_len * head_size); + std::vector past_value_data(batch_size * kv_num_heads * past_seq_len * head_size); + for (size_t i = 0; i < query_data.size(); i++) query_data[i] = 0.1f * static_cast(i % 7 + 1); + for (size_t i = 0; i < past_key_data.size(); i++) past_key_data[i] = 0.2f * static_cast(i % 5 + 1); + for (size_t i = 0; i < past_value_data.size(); i++) past_value_data[i] = 0.3f * static_cast(i % 3 + 1); + + auto output = RunGQASharedKV( + batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, + num_heads, kv_num_heads, head_size, /*use_cuda=*/false); + + // Verify non-zero and no NaN + bool all_zero = true; + for (size_t i = 0; i < output.size(); i++) { + EXPECT_FALSE(std::isnan(output[i])) << "NaN at index " << i; + if (output[i] != 0.0f) all_zero = false; + } + EXPECT_FALSE(all_zero) << "Output should not be all zeros"; +} + +// CPU: kv_sequence_length=0 with past, prompt phase (q_seq_len == total_seq_len). +TEST(GroupQueryAttentionTest, SharedKV_EmptyKV_WithPast_Prompt_CPU) { + constexpr int batch_size = 1; + constexpr int q_seq_len = 8; + constexpr int past_seq_len = 8; + constexpr int num_heads = 2; + constexpr int kv_num_heads = 1; + constexpr int head_size = 8; + constexpr int hidden_size = num_heads * head_size; + + std::vector query_data(batch_size * q_seq_len * hidden_size); + std::vector past_key_data(batch_size * kv_num_heads * past_seq_len * head_size); + std::vector past_value_data(batch_size * kv_num_heads * past_seq_len * head_size); + for (size_t i = 0; i < query_data.size(); i++) query_data[i] = 0.1f * static_cast(i % 7 + 1); + for (size_t i = 0; i < past_key_data.size(); i++) past_key_data[i] = 0.2f * static_cast(i % 5 + 1); + for (size_t i = 0; i < past_value_data.size(); i++) past_value_data[i] = 0.3f * static_cast(i % 3 + 1); + + auto output = RunGQASharedKV( + batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, + num_heads, kv_num_heads, head_size, /*use_cuda=*/false); + + bool all_zero = true; + for (size_t i = 0; i < output.size(); i++) { + EXPECT_FALSE(std::isnan(output[i])) << "NaN at index " << i; + if (output[i] != 0.0f) all_zero = false; + } + EXPECT_FALSE(all_zero) << "Output should not be all zeros"; +} + +// CUDA: kv_sequence_length=0 with past, decode (q_seq=1). +// Cross-checks CUDA against CPU for correctness. +TEST(GroupQueryAttentionTest, SharedKV_EmptyKV_WithPast_CUDA) { + auto cuda_ep = DefaultCudaExecutionProvider(); + if (!cuda_ep) { + GTEST_SKIP() << "CUDA EP not available"; + } + + constexpr int batch_size = 1; + constexpr int q_seq_len = 1; + constexpr int past_seq_len = 8; + constexpr int num_heads = 2; + constexpr int kv_num_heads = 1; + constexpr int head_size = 8; + constexpr int hidden_size = num_heads * head_size; + + std::vector query_data(batch_size * q_seq_len * hidden_size); + std::vector past_key_data(batch_size * kv_num_heads * past_seq_len * head_size); + std::vector past_value_data(batch_size * kv_num_heads * past_seq_len * head_size); + for (size_t i = 0; i < query_data.size(); i++) query_data[i] = 0.1f * static_cast(i % 7 + 1); + for (size_t i = 0; i < past_key_data.size(); i++) past_key_data[i] = 0.2f * static_cast(i % 5 + 1); + for (size_t i = 0; i < past_value_data.size(); i++) past_value_data[i] = 0.3f * static_cast(i % 3 + 1); + + auto cuda_output = RunGQASharedKV( + batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, + num_heads, kv_num_heads, head_size, /*use_cuda=*/true); + + auto cpu_output = RunGQASharedKV( + batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, + num_heads, kv_num_heads, head_size, /*use_cuda=*/false); + + ExpectOutputsMatch(cuda_output, cpu_output, 0.05f, "SharedKV_CUDA_vs_CPU"); +} + +// CPU: kv_sequence_length=0 with past, head_size=64. +TEST(GroupQueryAttentionTest, SharedKV_EmptyKV_WithPast_LargeHeadSize_CPU) { + constexpr int batch_size = 1; + constexpr int q_seq_len = 1; + constexpr int past_seq_len = 4; + constexpr int num_heads = 2; + constexpr int kv_num_heads = 1; + constexpr int head_size = 64; + constexpr int hidden_size = num_heads * head_size; + + std::vector query_data(batch_size * q_seq_len * hidden_size); + std::vector past_key_data(batch_size * kv_num_heads * past_seq_len * head_size); + std::vector past_value_data(batch_size * kv_num_heads * past_seq_len * head_size); + for (size_t i = 0; i < query_data.size(); i++) query_data[i] = 0.1f * static_cast(i % 11 + 1); + for (size_t i = 0; i < past_key_data.size(); i++) past_key_data[i] = 0.2f * static_cast(i % 7 + 1); + for (size_t i = 0; i < past_value_data.size(); i++) past_value_data[i] = 0.3f * static_cast(i % 5 + 1); + + auto output = RunGQASharedKV( + batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, + num_heads, kv_num_heads, head_size, /*use_cuda=*/false); + + bool all_zero = true; + for (size_t i = 0; i < output.size(); i++) { + EXPECT_FALSE(std::isnan(output[i])) << "NaN at index " << i; + if (output[i] != 0.0f) all_zero = false; + } + EXPECT_FALSE(all_zero) << "Output should not be all zeros"; +} + +// CPU: GQA ratio num_heads=8, kv_num_heads=1 (matches Gemma4 config). +TEST(GroupQueryAttentionTest, SharedKV_EmptyKV_WithPast_GQARatio8_CPU) { + constexpr int batch_size = 1; + constexpr int q_seq_len = 1; + constexpr int past_seq_len = 4; + constexpr int num_heads = 8; + constexpr int kv_num_heads = 1; + constexpr int head_size = 16; + constexpr int hidden_size = num_heads * head_size; + + std::vector query_data(batch_size * q_seq_len * hidden_size); + std::vector past_key_data(batch_size * kv_num_heads * past_seq_len * head_size); + std::vector past_value_data(batch_size * kv_num_heads * past_seq_len * head_size); + for (size_t i = 0; i < query_data.size(); i++) query_data[i] = 0.1f * static_cast(i % 13 + 1); + for (size_t i = 0; i < past_key_data.size(); i++) past_key_data[i] = 0.2f * static_cast(i % 7 + 1); + for (size_t i = 0; i < past_value_data.size(); i++) past_value_data[i] = 0.3f * static_cast(i % 5 + 1); + + auto output = RunGQASharedKV( + batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, + num_heads, kv_num_heads, head_size, /*use_cuda=*/false); + + bool all_zero = true; + for (size_t i = 0; i < output.size(); i++) { + EXPECT_FALSE(std::isnan(output[i])) << "NaN at index " << i; + if (output[i] != 0.0f) all_zero = false; + } + EXPECT_FALSE(all_zero) << "Output should not be all zeros"; +} + +// CUDA: kv_sequence_length=0 with past, prompt phase. Cross-checks against CPU. +TEST(GroupQueryAttentionTest, SharedKV_EmptyKV_WithPast_Prompt_CUDA) { + auto cuda_ep = DefaultCudaExecutionProvider(); + if (!cuda_ep) { + GTEST_SKIP() << "CUDA EP not available"; + } + + constexpr int batch_size = 1; + constexpr int q_seq_len = 8; + constexpr int past_seq_len = 8; + constexpr int num_heads = 2; + constexpr int kv_num_heads = 1; + constexpr int head_size = 8; + constexpr int hidden_size = num_heads * head_size; + + std::vector query_data(batch_size * q_seq_len * hidden_size); + std::vector past_key_data(batch_size * kv_num_heads * past_seq_len * head_size); + std::vector past_value_data(batch_size * kv_num_heads * past_seq_len * head_size); + for (size_t i = 0; i < query_data.size(); i++) query_data[i] = 0.1f * static_cast(i % 7 + 1); + for (size_t i = 0; i < past_key_data.size(); i++) past_key_data[i] = 0.2f * static_cast(i % 5 + 1); + for (size_t i = 0; i < past_value_data.size(); i++) past_value_data[i] = 0.3f * static_cast(i % 3 + 1); + + auto cuda_output = RunGQASharedKV( + batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, + num_heads, kv_num_heads, head_size, /*use_cuda=*/true); + auto cpu_output = RunGQASharedKV( + batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, + num_heads, kv_num_heads, head_size, /*use_cuda=*/false); + + ExpectOutputsMatch(cuda_output, cpu_output, 0.05f, "SharedKV_Prompt_CUDA_vs_CPU"); +} + +// CUDA: kv_sequence_length=0 with past, head_size=16 (different from default 8). +TEST(GroupQueryAttentionTest, SharedKV_EmptyKV_WithPast_LargeHeadSize_CUDA) { + auto cuda_ep = DefaultCudaExecutionProvider(); + if (!cuda_ep) { + GTEST_SKIP() << "CUDA EP not available"; + } + + constexpr int batch_size = 1; + constexpr int q_seq_len = 1; + constexpr int past_seq_len = 4; + constexpr int num_heads = 2; + constexpr int kv_num_heads = 1; + constexpr int head_size = 16; + constexpr int hidden_size = num_heads * head_size; + + std::vector query_data(batch_size * q_seq_len * hidden_size); + std::vector past_key_data(batch_size * kv_num_heads * past_seq_len * head_size); + std::vector past_value_data(batch_size * kv_num_heads * past_seq_len * head_size); + for (size_t i = 0; i < query_data.size(); i++) query_data[i] = 0.1f * static_cast(i % 11 + 1); + for (size_t i = 0; i < past_key_data.size(); i++) past_key_data[i] = 0.2f * static_cast(i % 7 + 1); + for (size_t i = 0; i < past_value_data.size(); i++) past_value_data[i] = 0.3f * static_cast(i % 5 + 1); + + auto cuda_output = RunGQASharedKV( + batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, + num_heads, kv_num_heads, head_size, /*use_cuda=*/true); + auto cpu_output = RunGQASharedKV( + batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, + num_heads, kv_num_heads, head_size, /*use_cuda=*/false); + + ExpectOutputsMatch(cuda_output, cpu_output, 0.05f, "SharedKV_LargeHead_CUDA_vs_CPU"); +} + +// CUDA: kv_sequence_length=0 with past, GQA ratio 8:1. Cross-checks against CPU. +TEST(GroupQueryAttentionTest, SharedKV_EmptyKV_WithPast_GQARatio8_CUDA) { + auto cuda_ep = DefaultCudaExecutionProvider(); + if (!cuda_ep) { + GTEST_SKIP() << "CUDA EP not available"; + } + + constexpr int batch_size = 1; + constexpr int q_seq_len = 1; + constexpr int past_seq_len = 4; + constexpr int num_heads = 8; + constexpr int kv_num_heads = 1; + constexpr int head_size = 16; + constexpr int hidden_size = num_heads * head_size; + + std::vector query_data(batch_size * q_seq_len * hidden_size); + std::vector past_key_data(batch_size * kv_num_heads * past_seq_len * head_size); + std::vector past_value_data(batch_size * kv_num_heads * past_seq_len * head_size); + for (size_t i = 0; i < query_data.size(); i++) query_data[i] = 0.1f * static_cast(i % 13 + 1); + for (size_t i = 0; i < past_key_data.size(); i++) past_key_data[i] = 0.2f * static_cast(i % 7 + 1); + for (size_t i = 0; i < past_value_data.size(); i++) past_value_data[i] = 0.3f * static_cast(i % 5 + 1); + + auto cuda_output = RunGQASharedKV( + batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, + num_heads, kv_num_heads, head_size, /*use_cuda=*/true); + auto cpu_output = RunGQASharedKV( + batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, + num_heads, kv_num_heads, head_size, /*use_cuda=*/false); + + ExpectOutputsMatch(cuda_output, cpu_output, 0.15f, "SharedKV_GQA8_CUDA_vs_CPU"); +} + } // namespace test } // namespace onnxruntime From 17a155a4cc3ce35738b5eb8f8da858329c8bdba6 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Fri, 8 May 2026 21:17:01 +0000 Subject: [PATCH 24/32] Fix comments --- .../cpu/bert/group_query_attention.cc | 22 +- .../cpu/bert/group_query_attention_helper.h | 6 +- .../cuda/bert/group_query_attention.cc | 37 +- .../cuda/bert/group_query_attention_impl.cu | 62 +--- .../group_query_attention_op_test.cc | 333 +----------------- 5 files changed, 15 insertions(+), 445 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc index 1a084af6fda26..2a8c26cdb44f3 100644 --- a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc @@ -113,21 +113,6 @@ Status GroupQueryAttention::Compute(OpKernelContext* context) const { Tensor* present_k = context->Output(1, present_k_shape); Tensor* present_v = context->Output(2, present_v_shape); - // present_key and present_value must be both present or both absent. - if ((present_k == nullptr) != (present_v == nullptr)) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "present_key and present_value must be both provided or both omitted."); - } - - // Omitting present outputs is only safe when past_key is not provided. - // When past_key exists, ConcatStateChunkGQA must build a concatenated - // past+current KV buffer in present_key/present_value for the attention GEMMs. - if ((present_k == nullptr || present_v == nullptr) && past_key != nullptr) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "present_key and present_value outputs are required when past_key is provided. " - "Omitting present outputs is only supported when there is no past KV cache."); - } - std::vector output_qk_shape{static_cast(batch_size), static_cast(num_heads_), static_cast(parameters.sequence_length), static_cast(parameters.total_sequence_length)}; Tensor* output_qk = context->Output(3, output_qk_shape); @@ -159,12 +144,7 @@ Status GroupQueryAttention::Compute(OpKernelContext* context) const { T* q_rotary = Q.GetMutable()->MutableData(); T* k_rotary = packed_qkv ? nullptr : K.GetMutable()->MutableData(); if (do_rotary_) { - // KV-shared decode with empty K/V: only apply RoPE to Q, skip K. - if (kv_sequence_length != sequence_length && kv_sequence_length != 0) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "do_rotary is not supported when query and key have different sequence lengths. " - "Apply RoPE externally before the GQA op for KV-shared layers."); - } + // When kv_sequence_length == 0 (shared KV), only Q needs RoPE — K is skipped below. ORT_ENFORCE(cos_cache != nullptr && sin_cache != nullptr, "cos_cache and sin_cache must be provided when do_rotary is true"); // Initialize rotary parameters rotary_embedding_helper::RotaryParameters rotary_params = {}; diff --git a/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h b/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h index efff8fc1fb0e4..c204ba4c43253 100644 --- a/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h +++ b/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h @@ -305,9 +305,11 @@ Status CheckInputs(const T* query, "total_sequence_length must be positive, got ", total_sequence_length, "."); } - // When there is no past KV (KV-shared / first-prompt), total_sequence_length - // must not exceed kv_sequence_length — the attention kernel reads up to + // When there is no past KV (first prompt), total_sequence_length must not + // exceed kv_sequence_length — the attention kernel reads up to // total_sequence_length from the K/V buffer which has kv_sequence_length entries. + // Note: KV-shared layers pass shared KV via past_key (past_key != nullptr), + // so this guard only applies to the first-prompt case. if (is_total_seqlen_on_cpu && past_key == nullptr && total_sequence_length > kv_sequence_length) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "total_sequence_length (", total_sequence_length, diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc index 4375988c3ad20..cc949a4efe3cf 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc @@ -267,30 +267,6 @@ Status GroupQueryAttention::ComputeInternal(OpKernelContext* context) cons Tensor* present_key_output = context->Output(1, present_shape); // present_key Tensor* present_value_output = context->Output(2, present_shape); // present_value - // present_key and present_value must be both present or both absent. - if ((present_key_output == nullptr) != (present_value_output == nullptr)) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "present_key and present_value must be both provided or both omitted."); - } - - // Omitting present outputs is only safe when past_key is not provided. - // When past_key exists, the kernel must concatenate past+current KV into present. - if ((present_key_output == nullptr || present_value_output == nullptr) && past_key != nullptr) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "present_key and present_value outputs are required when past_key is provided. " - "Omitting present outputs is only supported when there is no past KV cache."); - } - - // When present outputs are omitted, allocate internal scratch buffers so the - // CUDA kernels (flash attention, MEA, unfused) have a valid KV workspace. - IAllocatorUniquePtr present_key_scratch; - IAllocatorUniquePtr present_value_scratch; - if (present_key_output == nullptr || present_value_output == nullptr) { - size_t present_kv_bytes = present_shape.Size() * sizeof(U); - present_key_scratch = GetScratchBuffer(present_kv_bytes, GetComputeStream(context)); - present_value_scratch = GetScratchBuffer(present_kv_bytes, GetComputeStream(context)); - } - IAllocatorUniquePtr k_buffer; IAllocatorUniquePtr v_buffer; IAllocatorUniquePtr rotary_buffer; @@ -316,12 +292,8 @@ Status GroupQueryAttention::ComputeInternal(OpKernelContext* context) cons data.past_key = (past_key == nullptr) ? nullptr : reinterpret_cast(past_key->Data()); data.past_value = (past_value == nullptr) ? nullptr : reinterpret_cast(past_value->Data()); - data.present_key = (present_key_output != nullptr) - ? reinterpret_cast(present_key_output->MutableData()) - : reinterpret_cast(present_key_scratch.get()); - data.present_value = (present_value_output != nullptr) - ? reinterpret_cast(present_value_output->MutableData()) - : reinterpret_cast(present_value_scratch.get()); + data.present_key = reinterpret_cast(present_key_output->MutableData()); + data.present_value = reinterpret_cast(present_value_output->MutableData()); // Compute past_present_share_buffer early since it's needed for flash attention path selection. parameters.past_present_share_buffer = (data.past_key != nullptr && data.past_key == data.present_key); @@ -477,16 +449,13 @@ Status GroupQueryAttention::ComputeInternal(OpKernelContext* context) cons data.past_seq_lens = seq_lens_buffer.get(); data.total_seq_lens = seq_lens_buffer.get() + parameters.batch_size; data.padded_seq_lens = data.total_seq_lens + parameters.batch_size; - // For KV-shared decode (no past_key but not first_prompt), treat as first_prompt - // for sequence length computation so past_seq_lens = 0 (no past to offset from). - bool effective_is_first_prompt = parameters.is_first_prompt || (past_key == nullptr); ORT_RETURN_IF_ERROR(LaunchGetSequenceLengths(total_seq_lens_minus_one->Data(), data.past_seq_lens, data.total_seq_lens, data.padded_seq_lens, parameters.batch_size, parameters.sequence_length, - effective_is_first_prompt, + parameters.is_first_prompt, cuda_stream, device_prop.maxThreadsPerBlock)); DUMP_TENSOR_INIT(); diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu index 59dfaa5c5c849..5aaac5c9a1742 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu @@ -103,13 +103,7 @@ Status PrepareQKV( U* v = reinterpret_cast(data.present_value); int max_cache_length = parameters.seqlen_present_kv_cache; - if (!parameters.past_present_share_buffer && kv_sequence_length != sequence_length && kv_sequence_length > 0) { - // KV-shared decode (kv_seq != q_seq, kv_seq > 0): Transpose_BSNH_to_BNSH - // will write every element of the present buffer, so skip the memset. - // Note: kv_sequence_length == 0 (shared KV with past) does NOT run the - // transpose path — it copies past→present instead, so memset is still needed - // for the region beyond the copied data. - } else if (!parameters.past_present_share_buffer) { + if (!parameters.past_present_share_buffer) { size_t kv_buffer_size = (size_t)batch_size * kv_num_heads * max_cache_length * head_size * sizeof(U); CUDA_CALL_THROW(cudaMemsetAsync(data.present_key, 0, kv_buffer_size, stream)); CUDA_CALL_THROW(cudaMemsetAsync(data.present_value, 0, kv_buffer_size, stream)); @@ -156,60 +150,6 @@ Status PrepareQKV( } // If do_rotary is false, Q is used directly from data.query (q_out == nullptr). // K/V present buffers already point to the shared past — no work needed. - } else if (kv_sequence_length != sequence_length) { - // When kv_sequence_length differs from sequence_length (KV-shared decode), - // K/V are borrowed from a source layer with the full context length and already - // have RoPE applied. We transpose Q and K/V separately via Transpose_BSNH_to_BNSH - // since they have different sequence lengths. - // KV-shared decode does not support do_rotary or packed QKV — RoPE is applied - // externally before the GQA op, and Q/K/V are separate inputs. - if (parameters.do_rotary) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "do_rotary is not supported when query and key have different sequence lengths. " - "Apply RoPE externally before the GQA op for KV-shared layers."); - } - if (parameters.is_packed_qkv) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "Packed QKV is not supported when query and key have different sequence lengths."); - } - - // Q: use directly from input (already BSNH, no rotary needed) - // q_out is nullptr (no rotary, no packed), so q will point to data.query (set below) - - // K/V: transpose BSNH → BNSH directly into present buffer at offset 0. - // No RoPE needed (already applied by source layer), no append offset (no past). - // Transpose_BSNH_to_BNSH accepts half/BFloat16/float, not CUDA native types. - if constexpr (std::is_same::value) { - static_assert(std::is_same::value || std::is_same::value, - "KV-shared decode transpose only supports __half and __nv_bfloat16."); - if constexpr (std::is_same::value) { - ORT_RETURN_IF_ERROR((Transpose_BSNH_to_BNSH( - batch_size, kv_sequence_length, kv_num_heads, head_size, - reinterpret_cast(data.key), - reinterpret_cast(data.present_key), - stream, max_threads_per_block))); - ORT_RETURN_IF_ERROR((Transpose_BSNH_to_BNSH( - batch_size, kv_sequence_length, kv_num_heads, head_size, - reinterpret_cast(data.value), - reinterpret_cast(data.present_value), - stream, max_threads_per_block))); - } else if constexpr (std::is_same::value) { - ORT_RETURN_IF_ERROR((Transpose_BSNH_to_BNSH( - batch_size, kv_sequence_length, kv_num_heads, head_size, - reinterpret_cast(data.key), - reinterpret_cast(data.present_key), - stream, max_threads_per_block))); - ORT_RETURN_IF_ERROR((Transpose_BSNH_to_BNSH( - batch_size, kv_sequence_length, kv_num_heads, head_size, - reinterpret_cast(data.value), - reinterpret_cast(data.present_value), - stream, max_threads_per_block))); - } - } else { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "KV-shared decode (query_seq_len != kv_seq_len) with quantized KV cache " - "is not supported. Use non-quantized cache for KV-shared layers."); - } } else { ORT_RETURN_IF_ERROR((LaunchUnpackRoPEAppend( parameters.is_packed_qkv ? reinterpret_cast(data.query) : nullptr, diff --git a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc index 29d9223043af1..074ea2b0f3f32 100644 --- a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc @@ -335,339 +335,18 @@ TEST(GroupQueryAttentionTest, SeqlensKScalarRejected) { /*seqlens_k_shape=*/std::vector{}); } -// ============================================================================ -// Optional present_key/present_value output tests -// ============================================================================ - -// Run GQA with the given inputs and return the output tensor as a vector. -// Supports separate Q and K/V sequence lengths for KV-shared decode scenarios. -// When use_cuda=true, runs on CUDA EP instead of CPU EP. -static std::vector RunGQAAndGetOutput( - int batch_size, - int q_seq_len, - int kv_seq_len, - const std::vector& query_data, - const std::vector& key_data, - const std::vector& value_data, - int num_heads, - int kv_num_heads, - int head_size, - bool omit_present, - bool use_cuda = false) { - const int hidden_size = num_heads * head_size; - const int kv_hidden_size = kv_num_heads * head_size; - const int total_seq_len = kv_seq_len; // no past: total = kv length - - OpTester tester("GroupQueryAttention", 1, onnxruntime::kMSDomain); - tester.AddAttribute("num_heads", static_cast(num_heads)); - tester.AddAttribute("kv_num_heads", static_cast(kv_num_heads)); - - tester.AddInput("query", {batch_size, q_seq_len, hidden_size}, query_data); - tester.AddInput("key", {batch_size, kv_seq_len, kv_hidden_size}, key_data); - tester.AddInput("value", {batch_size, kv_seq_len, kv_hidden_size}, value_data); - - tester.AddOptionalInputEdge(); // past_key - tester.AddOptionalInputEdge(); // past_value - - std::vector seqlens_k_data(batch_size, static_cast(total_seq_len - 1)); - tester.AddInput("seqlens_k", {batch_size}, seqlens_k_data); - tester.AddInput("total_sequence_length", {1}, {static_cast(total_seq_len)}); - - tester.AddOptionalInputEdge(); // cos_cache - tester.AddOptionalInputEdge(); // sin_cache - tester.AddOptionalInputEdge(); // position_ids - tester.AddOptionalInputEdge(); // attention_bias - tester.AddOptionalInputEdge(); // head_sink - - const int output_size = batch_size * q_seq_len * hidden_size; - tester.AddOutput("output", {batch_size, q_seq_len, hidden_size}, - std::vector(output_size, 0.0f)); - - if (omit_present) { - tester.AddOptionalOutputEdge(); // present_key - tester.AddOptionalOutputEdge(); // present_value - } else { - const int present_size = batch_size * kv_num_heads * total_seq_len * head_size; - tester.AddOutput("present_key", {batch_size, kv_num_heads, total_seq_len, head_size}, - std::vector(present_size, 0.0f)); - tester.AddOutput("present_value", {batch_size, kv_num_heads, total_seq_len, head_size}, - std::vector(present_size, 0.0f)); - } - tester.SetOutputTolerance(1e6f); // We compare fetched outputs ourselves - - std::vector> execution_providers; - if (use_cuda) { - execution_providers.push_back(DefaultCudaExecutionProvider()); - } else { - execution_providers.push_back(DefaultCpuExecutionProvider()); - } - tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); - - auto fetches = tester.GetFetches(); - const float* out_data = fetches[0].Get().Data(); - return std::vector(out_data, out_data + output_size); -} - -// Helper: compare two output vectors element-wise and check non-zero. +// Helper to compare two output vectors (non-zero check + element-wise tolerance). static void ExpectOutputsMatch(const std::vector& a, const std::vector& b, - float tol, const std::string& label) { - ASSERT_EQ(a.size(), b.size()); - for (size_t i = 0; i < a.size(); i++) { - EXPECT_NEAR(a[i], b[i], tol) << label << " mismatch at index " << i; - } + float tolerance, const char* label) { + ASSERT_EQ(a.size(), b.size()) << label << ": output size mismatch"; bool all_zero = true; - for (float v : a) { - if (v != 0.0f) { - all_zero = false; - break; - } + for (size_t i = 0; i < a.size(); i++) { + EXPECT_NEAR(a[i], b[i], tolerance) << label << " mismatch at index " << i; + if (a[i] != 0.0f) all_zero = false; } EXPECT_FALSE(all_zero) << label << " output should not be all zeros"; } -// Regression: omitting optional present outputs must not change the attention output -// compared to when present outputs are connected (first-prompt, no past KV). -TEST(GroupQueryAttentionTest, OptionalPresent_OmittingDoesNotChangeOutput) { - constexpr int batch_size = 1; - constexpr int sequence_length = 4; - constexpr int num_heads = 2; - constexpr int kv_num_heads = 1; - constexpr int head_size = 8; - constexpr int hidden_size = num_heads * head_size; - constexpr int kv_hidden_size = kv_num_heads * head_size; - - // Deterministic non-trivial inputs - std::vector query_data(batch_size * sequence_length * hidden_size); - std::vector key_data(batch_size * sequence_length * kv_hidden_size); - std::vector value_data(batch_size * sequence_length * kv_hidden_size); - for (size_t i = 0; i < query_data.size(); i++) query_data[i] = 0.1f * static_cast(i % 7 + 1); - for (size_t i = 0; i < key_data.size(); i++) key_data[i] = 0.2f * static_cast(i % 5 + 1); - for (size_t i = 0; i < value_data.size(); i++) value_data[i] = 0.3f * static_cast(i % 3 + 1); - - auto output_with_present = RunGQAAndGetOutput( - batch_size, sequence_length, sequence_length, query_data, key_data, value_data, - num_heads, kv_num_heads, head_size, /*omit_present=*/false); - - auto output_without_present = RunGQAAndGetOutput( - batch_size, sequence_length, sequence_length, query_data, key_data, value_data, - num_heads, kv_num_heads, head_size, /*omit_present=*/true); - - ExpectOutputsMatch(output_with_present, output_without_present, 1e-5f, "OptionalPresent"); -} - -// Regression (batched): same equivalence check with batch_size > 1 -TEST(GroupQueryAttentionTest, OptionalPresent_BatchedOmitMatchesConnected) { - constexpr int batch_size = 2; - constexpr int sequence_length = 3; - constexpr int num_heads = 2; - constexpr int kv_num_heads = 1; - constexpr int head_size = 8; - constexpr int hidden_size = num_heads * head_size; - constexpr int kv_hidden_size = kv_num_heads * head_size; - - std::vector query_data(batch_size * sequence_length * hidden_size); - std::vector key_data(batch_size * sequence_length * kv_hidden_size); - std::vector value_data(batch_size * sequence_length * kv_hidden_size); - for (size_t i = 0; i < query_data.size(); i++) query_data[i] = 0.15f * static_cast(i % 11 + 1); - for (size_t i = 0; i < key_data.size(); i++) key_data[i] = 0.25f * static_cast(i % 7 + 1); - for (size_t i = 0; i < value_data.size(); i++) value_data[i] = 0.35f * static_cast(i % 5 + 1); - - auto output_with = RunGQAAndGetOutput( - batch_size, sequence_length, sequence_length, query_data, key_data, value_data, - num_heads, kv_num_heads, head_size, /*omit_present=*/false); - - auto output_without = RunGQAAndGetOutput( - batch_size, sequence_length, sequence_length, query_data, key_data, value_data, - num_heads, kv_num_heads, head_size, /*omit_present=*/true); - - ExpectOutputsMatch(output_with, output_without, 1e-5f, "BatchedOptionalPresent"); -} - -// KV-shared first-prompt: longer sequence with no past, present omitted. -TEST(GroupQueryAttentionTest, OptionalPresent_KVSharedFirstPrompt) { - constexpr int batch_size = 1; - constexpr int sequence_length = 8; - constexpr int num_heads = 2; - constexpr int kv_num_heads = 1; - constexpr int head_size = 8; - constexpr int hidden_size = num_heads * head_size; - constexpr int kv_hidden_size = kv_num_heads * head_size; - - std::vector query_data(batch_size * sequence_length * hidden_size); - std::vector key_data(batch_size * sequence_length * kv_hidden_size); - std::vector value_data(batch_size * sequence_length * kv_hidden_size); - for (size_t i = 0; i < query_data.size(); i++) query_data[i] = 0.1f * static_cast(i % 7 + 1); - for (size_t i = 0; i < key_data.size(); i++) key_data[i] = 0.2f * static_cast(i % 5 + 1); - for (size_t i = 0; i < value_data.size(); i++) value_data[i] = 0.3f * static_cast(i % 3 + 1); - - auto output_with = RunGQAAndGetOutput( - batch_size, sequence_length, sequence_length, query_data, key_data, value_data, - num_heads, kv_num_heads, head_size, /*omit_present=*/false); - auto output_without = RunGQAAndGetOutput( - batch_size, sequence_length, sequence_length, query_data, key_data, value_data, - num_heads, kv_num_heads, head_size, /*omit_present=*/true); - - ExpectOutputsMatch(output_with, output_without, 1e-5f, "KVSharedFirstPrompt"); -} - -// Reject: omitting present outputs when past_key is provided (KV cache concatenation needed) -TEST(GroupQueryAttentionTest, OptionalPresent_RejectWithPastKey) { - constexpr int batch_size = 1; - constexpr int sequence_length = 1; - constexpr int past_seq_len = 4; - constexpr int total_seq_len = past_seq_len + sequence_length; - constexpr int num_heads = 2; - constexpr int kv_num_heads = 1; - constexpr int head_size = 8; - constexpr int hidden_size = num_heads * head_size; - constexpr int kv_hidden_size = kv_num_heads * head_size; - - OpTester tester("GroupQueryAttention", 1, onnxruntime::kMSDomain); - tester.AddAttribute("num_heads", static_cast(num_heads)); - tester.AddAttribute("kv_num_heads", static_cast(kv_num_heads)); - - tester.AddInput("query", {batch_size, sequence_length, hidden_size}, - std::vector(batch_size * sequence_length * hidden_size, 1.0f)); - tester.AddInput("key", {batch_size, sequence_length, kv_hidden_size}, - std::vector(batch_size * sequence_length * kv_hidden_size, 0.5f)); - tester.AddInput("value", {batch_size, sequence_length, kv_hidden_size}, - std::vector(batch_size * sequence_length * kv_hidden_size, 0.5f)); - - // Provide past_key/past_value — this triggers the rejection when present is omitted - tester.AddInput("past_key", {batch_size, kv_num_heads, past_seq_len, head_size}, - std::vector(batch_size * kv_num_heads * past_seq_len * head_size, 0.3f)); - tester.AddInput("past_value", {batch_size, kv_num_heads, past_seq_len, head_size}, - std::vector(batch_size * kv_num_heads * past_seq_len * head_size, 0.3f)); - - tester.AddInput("seqlens_k", {batch_size}, {static_cast(total_seq_len - 1)}); - tester.AddInput("total_sequence_length", {1}, {static_cast(total_seq_len)}); - - tester.AddOptionalInputEdge(); // cos_cache - tester.AddOptionalInputEdge(); // sin_cache - tester.AddOptionalInputEdge(); // position_ids - tester.AddOptionalInputEdge(); // attention_bias - tester.AddOptionalInputEdge(); // head_sink - - tester.AddOutput("output", {batch_size, sequence_length, hidden_size}, - std::vector(batch_size * sequence_length * hidden_size, 0.0f)); - tester.AddOptionalOutputEdge(); // present_key — omitted - tester.AddOptionalOutputEdge(); // present_value — omitted - - std::vector> execution_providers; - execution_providers.push_back(DefaultCpuExecutionProvider()); - tester.Run(OpTester::ExpectResult::kExpectFailure, - "present_key and present_value outputs are required when past_key is provided", - {}, nullptr, &execution_providers); -} - -// Regression (CUDA): omitting present outputs on CUDA EP must produce the same -// attention output as when present outputs are connected. The CUDA path allocates -// internal scratch buffers to serve as KV workspace for flash/MEA/unfused kernels. -TEST(GroupQueryAttentionTest, OptionalPresent_CudaOmitMatchesConnected) { - auto cuda_ep = DefaultCudaExecutionProvider(); - if (!cuda_ep) { - GTEST_SKIP() << "CUDA EP not available"; - } - - constexpr int batch_size = 1; - constexpr int sequence_length = 4; - constexpr int num_heads = 2; - constexpr int kv_num_heads = 1; - constexpr int head_size = 8; - constexpr int hidden_size = num_heads * head_size; - constexpr int kv_hidden_size = kv_num_heads * head_size; - - std::vector query_data(batch_size * sequence_length * hidden_size); - std::vector key_data(batch_size * sequence_length * kv_hidden_size); - std::vector value_data(batch_size * sequence_length * kv_hidden_size); - for (size_t i = 0; i < query_data.size(); i++) query_data[i] = 0.1f * static_cast(i % 7 + 1); - for (size_t i = 0; i < key_data.size(); i++) key_data[i] = 0.2f * static_cast(i % 5 + 1); - for (size_t i = 0; i < value_data.size(); i++) value_data[i] = 0.3f * static_cast(i % 3 + 1); - - auto output_with = RunGQAAndGetOutput( - batch_size, sequence_length, sequence_length, query_data, key_data, value_data, - num_heads, kv_num_heads, head_size, /*omit_present=*/false, /*use_cuda=*/true); - auto output_without = RunGQAAndGetOutput( - batch_size, sequence_length, sequence_length, query_data, key_data, value_data, - num_heads, kv_num_heads, head_size, /*omit_present=*/true, /*use_cuda=*/true); - - ExpectOutputsMatch(output_with, output_without, 1e-5f, "CudaOptionalPresent"); -} - -// KV-shared decode: Q_seq=1, KV_seq=8, no past, present omitted. -TEST(GroupQueryAttentionTest, OptionalPresent_KVSharedDecode) { - constexpr int batch_size = 1; - constexpr int q_seq_len = 1; - constexpr int kv_seq_len = 8; - constexpr int num_heads = 2; - constexpr int kv_num_heads = 1; - constexpr int head_size = 8; - constexpr int hidden_size = num_heads * head_size; - constexpr int kv_hidden_size = kv_num_heads * head_size; - - std::vector query_data(batch_size * q_seq_len * hidden_size); - std::vector key_data(batch_size * kv_seq_len * kv_hidden_size); - std::vector value_data(batch_size * kv_seq_len * kv_hidden_size); - for (size_t i = 0; i < query_data.size(); i++) query_data[i] = 0.1f * static_cast(i % 7 + 1); - for (size_t i = 0; i < key_data.size(); i++) key_data[i] = 0.2f * static_cast(i % 5 + 1); - for (size_t i = 0; i < value_data.size(); i++) value_data[i] = 0.3f * static_cast(i % 3 + 1); - - auto output_with = RunGQAAndGetOutput( - batch_size, q_seq_len, kv_seq_len, query_data, key_data, value_data, - num_heads, kv_num_heads, head_size, /*omit_present=*/false); - auto output_without = RunGQAAndGetOutput( - batch_size, q_seq_len, kv_seq_len, query_data, key_data, value_data, - num_heads, kv_num_heads, head_size, /*omit_present=*/true); - - ExpectOutputsMatch(output_with, output_without, 1e-5f, "KVSharedDecode"); -} - -// CUDA KV-shared decode: Q_seq=1, KV_seq=8, no past, present omitted. -// Exercises the CUDA Transpose_BSNH_to_BNSH path and cross-checks against CPU. -TEST(GroupQueryAttentionTest, OptionalPresent_CudaKVSharedDecode) { - auto cuda_ep = DefaultCudaExecutionProvider(); - if (!cuda_ep) { - GTEST_SKIP() << "CUDA EP not available"; - } - - constexpr int batch_size = 1; - constexpr int q_seq_len = 1; - constexpr int kv_seq_len = 8; - constexpr int num_heads = 2; - constexpr int kv_num_heads = 1; - constexpr int head_size = 8; - constexpr int hidden_size = num_heads * head_size; - constexpr int kv_hidden_size = kv_num_heads * head_size; - - std::vector query_data(batch_size * q_seq_len * hidden_size); - std::vector key_data(batch_size * kv_seq_len * kv_hidden_size); - std::vector value_data(batch_size * kv_seq_len * kv_hidden_size); - for (size_t i = 0; i < query_data.size(); i++) query_data[i] = 0.1f * static_cast(i % 7 + 1); - for (size_t i = 0; i < key_data.size(); i++) key_data[i] = 0.2f * static_cast(i % 5 + 1); - for (size_t i = 0; i < value_data.size(); i++) value_data[i] = 0.3f * static_cast(i % 3 + 1); - - // CUDA with vs without present - auto cuda_with = RunGQAAndGetOutput( - batch_size, q_seq_len, kv_seq_len, query_data, key_data, value_data, - num_heads, kv_num_heads, head_size, /*omit_present=*/false, /*use_cuda=*/true); - auto cuda_without = RunGQAAndGetOutput( - batch_size, q_seq_len, kv_seq_len, query_data, key_data, value_data, - num_heads, kv_num_heads, head_size, /*omit_present=*/true, /*use_cuda=*/true); - - ExpectOutputsMatch(cuda_with, cuda_without, 1e-4f, "CudaKVSharedDecode"); - - // Cross-check: CUDA should match CPU - auto cpu_with = RunGQAAndGetOutput( - batch_size, q_seq_len, kv_seq_len, query_data, key_data, value_data, - num_heads, kv_num_heads, head_size, /*omit_present=*/false, /*use_cuda=*/false); - - ASSERT_EQ(cuda_with.size(), cpu_with.size()); - for (size_t i = 0; i < cuda_with.size(); i++) { - EXPECT_NEAR(cuda_with[i], cpu_with[i], 1e-4f) - << "CUDA vs CPU KV-shared decode mismatch at index " << i; - } -} - // --------------------------------------------------------------------------- // Tests for kv_sequence_length=0 with borrowed past_key/past_value // (Gemma4 shared KV pattern: empty K/V inputs, all KV data in past buffer) From 2dd66c41fe0c04c24e6caa079628b66ff01fb30e Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Fri, 8 May 2026 21:28:34 +0000 Subject: [PATCH 25/32] Fix --- .../cpu/bert/group_query_attention_helper.h | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h b/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h index c204ba4c43253..bd5edc37fd444 100644 --- a/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h +++ b/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h @@ -305,18 +305,6 @@ Status CheckInputs(const T* query, "total_sequence_length must be positive, got ", total_sequence_length, "."); } - // When there is no past KV (first prompt), total_sequence_length must not - // exceed kv_sequence_length — the attention kernel reads up to - // total_sequence_length from the K/V buffer which has kv_sequence_length entries. - // Note: KV-shared layers pass shared KV via past_key (past_key != nullptr), - // so this guard only applies to the first-prompt case. - if (is_total_seqlen_on_cpu && past_key == nullptr && total_sequence_length > kv_sequence_length) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "total_sequence_length (", total_sequence_length, - ") must not exceed kv_sequence_length (", kv_sequence_length, - ") when past_key is not provided."); - } - int present_sequence_length = std::max(total_sequence_length, past_sequence_length); int rotary_dim = 0; From 01a1ef6f3c3f36bad308813ace1ed292c7c7d36b Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Fri, 8 May 2026 22:05:32 +0000 Subject: [PATCH 26/32] Apply copilot comments --- .../contrib_ops/cpu/bert/group_query_attention_helper.h | 7 +++++++ onnxruntime/core/graph/contrib_ops/bert_defs.cc | 6 ++---- .../test/contrib_ops/group_query_attention_op_test.cc | 3 +-- 3 files changed, 10 insertions(+), 6 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h b/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h index bd5edc37fd444..731578f1a27e3 100644 --- a/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h +++ b/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h @@ -270,6 +270,13 @@ Status CheckInputs(const T* query, } else if (past_key != nullptr || past_value != nullptr) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'past_key' and 'past_value' shall be both present or both absent."); + } else if (kv_sequence_length != sequence_length) { + // Without past KV, Q and K/V must have the same sequence length. + // Cross-attention (different Q/KV lengths) is not supported by GQA. + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "query and key must have the same sequence length when past_key is not provided. " + "Got sequence_length=", + sequence_length, ", kv_sequence_length=", kv_sequence_length); } // Spec requires 1D shape (batch_size), but older model builders may add unit diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index e8ec04586a9d6..1209446c6a367 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -1323,15 +1323,13 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "present state key with support for format BNSH. When past_key uses same tensor as present_key" "(k-v buffer), it is of length max_sequence_length... otherwise of length past_sequence_length +" "kv_sequence_length.", - "T_CACHE", - OpSchema::Optional) + "T_CACHE") .Output(2, "present_value", "present state value with support for format BNSH. When past_value uses same tensor as present_value" "(k-v buffer), it is of length max_sequence_length... otherwise of length past_sequence_length +" "kv_sequence_length.", - "T_CACHE", - OpSchema::Optional) + "T_CACHE") .Output(3, "output_qk", "Values of QK matrix multiplication, either before or after softmax normalization", diff --git a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc index 074ea2b0f3f32..41d2af652066e 100644 --- a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc @@ -1,6 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include #include #include @@ -419,8 +420,6 @@ static std::vector RunGQASharedKV( return std::vector(out_data, out_data + output_size); } -// CPU: kv_sequence_length=0 with past_key/past_value (shared KV decode). -// Validates output matches the equivalent non-empty K/V path. // CPU: kv_sequence_length=0 with past_key/past_value (shared KV decode). // Validates output is non-zero (attention over past KV produces valid output). // Note: cannot compare against RunGQAAndGetOutput because the two paths have From 24afd1e71e33a3c8a63d1005be7c8d8250fbf445 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Fri, 8 May 2026 22:58:10 +0000 Subject: [PATCH 27/32] revert docs --- docs/ContribOperators.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 45e85fcd9c9d5..9aa44a1600ae6 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -2671,14 +2671,14 @@ This version of the operator has been available since version 1 of the 'com.micr
Scale tensor for past_value.
-#### Outputs (1 - 4) +#### Outputs (3 - 4)
output : T
3D output tensor with shape (batch_size, sequence_length, hidden_size)
-
present_key (optional) : T_CACHE
+
present_key : T_CACHE
present state key with support for format BNSH. When past_key uses same tensor as present_key(k-v buffer), it is of length max_sequence_length... otherwise of length past_sequence_length +kv_sequence_length.
-
present_value (optional) : T_CACHE
+
present_value : T_CACHE
present state value with support for format BNSH. When past_value uses same tensor as present_value(k-v buffer), it is of length max_sequence_length... otherwise of length past_sequence_length +kv_sequence_length.
output_qk (optional) : T
Values of QK matrix multiplication, either before or after softmax normalization
From 29bd4cb986ab7685610f4a52d6e17812a9fdf6fd Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 11 May 2026 04:38:38 +0000 Subject: [PATCH 28/32] Address comments --- .../contrib_ops/cpu/bert/gqa_attention_base.h | 7 + .../cuda/bert/group_query_attention.cc | 9 - .../cuda/bert/group_query_attention_impl.cu | 46 +++-- .../group_query_attention_op_test.cc | 195 ++++++++++++++++++ 4 files changed, 228 insertions(+), 29 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h b/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h index 33764ed59a6f7..ce17ca0dcf358 100644 --- a/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h +++ b/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h @@ -90,6 +90,13 @@ class GQAAttentionBase { ? static_cast(present_key->Shape().GetDims()[2]) : parameters.total_sequence_length; + // Shared KV: total_sequence_length must fit within the past buffer. + if (kv_sequence_length == 0) { + ORT_ENFORCE(total_sequence_length <= seqlen_past_kv_cache, + "total_seqlen (", total_sequence_length, ") exceeds past buffer size (", + seqlen_past_kv_cache, ") in shared KV mode"); + } + // Compute the attention score. bool gqa_mlas_supported = MlasGQASupported(CblasNoTrans, CblasTrans) && MlasGQASupported(CblasNoTrans, CblasNoTrans); diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc index cc949a4efe3cf..a7466b40fe12c 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc @@ -559,15 +559,6 @@ Status GroupQueryAttention::ComputeInternal(OpKernelContext* context) cons std::is_same::value); } - // Validate past_value pointer consistency (past_present_share_buffer was computed early after pointer setup) - if (data.present_value != nullptr) { - if (parameters.past_present_share_buffer) { - ORT_ENFORCE(data.past_value == data.present_value, "past_value and present_value must be the same tensor when past_present_share_buffer is true"); - } else { - ORT_ENFORCE(data.past_value != data.present_value, "past_value and present_value must be different tensors when past_present_share_buffer is false"); - } - } - data.output = reinterpret_cast(output->MutableData()); if (parameters.do_rotary) { diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu index 5aaac5c9a1742..be092056bdc93 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu @@ -125,28 +125,34 @@ Status PrepareQKV( } // Shared KV path: K/V inputs are empty (kv_sequence_length == 0) and the - // past buffer already contains the full shared KV cache. When - // past_present_share_buffer is true the present buffer aliases past, so we - // only need to process Q (apply RoPE if configured) and skip all K/V work. + // past buffer already contains the full shared KV cache. This requires + // past_key/past_value to be provided (with RoPE already applied to K). + // past_present_share_buffer must be true so the present buffer aliases past + // — no copy needed, attention reads directly from the shared KV data. if (kv_sequence_length == 0) { if (parameters.do_rotary && data.cos_cache != nullptr && data.sin_cache != nullptr) { - // Apply RoPE to Q only. Launch the kernel with kv_num_heads=0 so that - // only QUERY head threads are spawned — no KEY/VALUE threads at all. - ORT_RETURN_IF_ERROR((LaunchUnpackRoPEAppend( - nullptr, // packed_qkv - reinterpret_cast(data.query), - nullptr, // key (empty) - nullptr, // value (empty) - q_out, - nullptr, // k_cache (unused) - nullptr, // v_cache (unused) - data.k_scale, data.v_scale, - num_heads, 0 /* kv_num_heads=0: no K/V threads */, head_size, sequence_length, batch_size, - max_cache_length, data.past_seq_lens, - reinterpret_cast(data.cos_cache), reinterpret_cast(data.sin_cache), - parameters.rotary_dim, data.position_ids, parameters.rotary_interleaved, - is_cache_bnsh, parameters.k_quant_type, - stream, max_threads_per_block))); + // Apply RoPE to Q only using the standalone rotary embedding kernel. + // Q is in BSNH format; the kernel writes rotated Q to q_out. + // position_ids_format: 1 = explicit per-token position_ids, 2 = past_seq_lens + s + // When position_ids is null, use format 2 (derives position from past_seq_lens). + const int pos_format = data.position_ids != nullptr ? 1 : 2; + if constexpr (std::is_same::value) { + ORT_RETURN_IF_ERROR((LaunchRotaryEmbeddingKernel( + stream, reinterpret_cast(q_out), reinterpret_cast(data.query), + data.position_ids, data.past_seq_lens, + reinterpret_cast(data.cos_cache), reinterpret_cast(data.sin_cache), + batch_size, sequence_length, num_heads, head_size, parameters.rotary_dim, max_cache_length, + pos_format, parameters.rotary_interleaved, + max_threads_per_block, false /* is_input_bnsh_format: Q is BSNH */))); + } else if constexpr (std::is_same::value) { + ORT_RETURN_IF_ERROR((LaunchRotaryEmbeddingKernel( + stream, reinterpret_cast(q_out), reinterpret_cast(data.query), + data.position_ids, data.past_seq_lens, + reinterpret_cast(data.cos_cache), reinterpret_cast(data.sin_cache), + batch_size, sequence_length, num_heads, head_size, parameters.rotary_dim, max_cache_length, + pos_format, parameters.rotary_interleaved, + max_threads_per_block, false /* is_input_bnsh_format: Q is BSNH */))); + } } // If do_rotary is false, Q is used directly from data.query (q_out == nullptr). // K/V present buffers already point to the shared past — no work needed. diff --git a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc index 41d2af652066e..7487e717cffd4 100644 --- a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc @@ -670,5 +670,200 @@ TEST(GroupQueryAttentionTest, SharedKV_EmptyKV_WithPast_GQARatio8_CUDA) { ExpectOutputsMatch(cuda_output, cpu_output, 0.15f, "SharedKV_GQA8_CUDA_vs_CPU"); } +// --------------------------------------------------------------------------- +// Shared KV tests with do_rotary=1 (Gemma4 primary use case) +// --------------------------------------------------------------------------- + +// Helper: run GQA with empty K/V, past_key/past_value, and do_rotary=1. +// Generates cos/sin caches and position_ids internally. +static std::vector RunGQASharedKVWithRotary( + int batch_size, + int q_seq_len, + int past_seq_len, + const std::vector& query_data, + const std::vector& past_key_data, + const std::vector& past_value_data, + int num_heads, + int kv_num_heads, + int head_size, + bool use_cuda = false) { + const int hidden_size = num_heads * head_size; + const int total_seq_len = past_seq_len; + const int rotary_dim = head_size; // full rotary + const int max_seq_len = past_seq_len + 16; // cos/sin cache length + + OpTester tester("GroupQueryAttention", 1, onnxruntime::kMSDomain); + tester.AddAttribute("num_heads", static_cast(num_heads)); + tester.AddAttribute("kv_num_heads", static_cast(kv_num_heads)); + tester.AddAttribute("do_rotary", static_cast(1)); + + // Q: [batch, q_seq_len, hidden_size] + tester.AddInput("query", {batch_size, q_seq_len, hidden_size}, query_data); + // K/V: empty [batch, 0, kv_hidden_size] + const int kv_hidden_size = kv_num_heads * head_size; + tester.AddInput("key", {batch_size, 0, kv_hidden_size}, {}); + tester.AddInput("value", {batch_size, 0, kv_hidden_size}, {}); + + // past_key/past_value: [batch, kv_num_heads, past_seq_len, head_size] BNSH + tester.AddInput("past_key", {batch_size, kv_num_heads, past_seq_len, head_size}, past_key_data); + tester.AddInput("past_value", {batch_size, kv_num_heads, past_seq_len, head_size}, past_value_data); + + std::vector seqlens_k_data(batch_size, static_cast(total_seq_len - 1)); + tester.AddInput("seqlens_k", {batch_size}, seqlens_k_data); + tester.AddInput("total_sequence_length", {1}, {static_cast(total_seq_len)}); + + // cos_cache/sin_cache: [max_seq_len, rotary_dim / 2] + const int half_rotary = rotary_dim / 2; + std::vector cos_cache(max_seq_len * half_rotary); + std::vector sin_cache(max_seq_len * half_rotary); + for (int pos = 0; pos < max_seq_len; pos++) { + for (int d = 0; d < half_rotary; d++) { + float freq = 1.0f / std::pow(10000.0f, 2.0f * static_cast(d) / static_cast(rotary_dim)); + cos_cache[pos * half_rotary + d] = std::cos(static_cast(pos) * freq); + sin_cache[pos * half_rotary + d] = std::sin(static_cast(pos) * freq); + } + } + tester.AddInput("cos_cache", {max_seq_len, half_rotary}, cos_cache); + tester.AddInput("sin_cache", {max_seq_len, half_rotary}, sin_cache); + + // position_ids: [batch, q_seq_len] — positions for the Q tokens + std::vector position_ids(batch_size * q_seq_len); + for (int b = 0; b < batch_size; b++) { + int past_len = total_seq_len - q_seq_len; + for (int s = 0; s < q_seq_len; s++) { + position_ids[b * q_seq_len + s] = static_cast(past_len + s); + } + } + tester.AddInput("position_ids", {batch_size, q_seq_len}, position_ids); + + tester.AddOptionalInputEdge(); // attention_bias + tester.AddOptionalInputEdge(); // head_sink + + const int output_size = batch_size * q_seq_len * hidden_size; + tester.AddOutput("output", {batch_size, q_seq_len, hidden_size}, + std::vector(output_size, 0.0f)); + + const int present_size = batch_size * kv_num_heads * past_seq_len * head_size; + tester.AddOutput("present_key", {batch_size, kv_num_heads, past_seq_len, head_size}, + std::vector(present_size, 0.0f)); + tester.AddOutput("present_value", {batch_size, kv_num_heads, past_seq_len, head_size}, + std::vector(present_size, 0.0f)); + + tester.SetOutputTolerance(1e6f); + + std::vector> execution_providers; + if (use_cuda) { + execution_providers.push_back(DefaultCudaExecutionProvider()); + } else { + execution_providers.push_back(DefaultCpuExecutionProvider()); + } + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); + + auto fetches = tester.GetFetches(); + const float* out_data = fetches[0].Get().Data(); + return std::vector(out_data, out_data + output_size); +} + +// CPU: shared KV with do_rotary=1 (Q-only RoPE path). +TEST(GroupQueryAttentionTest, SharedKV_EmptyKV_WithPast_Rotary_CPU) { + constexpr int batch_size = 1; + constexpr int q_seq_len = 1; + constexpr int past_seq_len = 8; + constexpr int num_heads = 2; + constexpr int kv_num_heads = 1; + constexpr int head_size = 16; // must be multiple of 16 for rotary + constexpr int hidden_size = num_heads * head_size; + + std::vector query_data(batch_size * q_seq_len * hidden_size); + std::vector past_key_data(batch_size * kv_num_heads * past_seq_len * head_size); + std::vector past_value_data(batch_size * kv_num_heads * past_seq_len * head_size); + for (size_t i = 0; i < query_data.size(); i++) query_data[i] = 0.1f * static_cast(i % 7 + 1); + for (size_t i = 0; i < past_key_data.size(); i++) past_key_data[i] = 0.2f * static_cast(i % 5 + 1); + for (size_t i = 0; i < past_value_data.size(); i++) past_value_data[i] = 0.3f * static_cast(i % 3 + 1); + + auto output = RunGQASharedKVWithRotary( + batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, + num_heads, kv_num_heads, head_size, /*use_cuda=*/false); + + // Output with rotary should differ from without rotary (RoPE changes Q projections) + auto output_no_rotary = RunGQASharedKV( + batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, + num_heads, kv_num_heads, head_size, /*use_cuda=*/false); + + bool all_zero = true; + bool differs_from_no_rotary = false; + for (size_t i = 0; i < output.size(); i++) { + EXPECT_FALSE(std::isnan(output[i])) << "NaN at index " << i; + if (output[i] != 0.0f) all_zero = false; + if (std::abs(output[i] - output_no_rotary[i]) > 1e-6f) differs_from_no_rotary = true; + } + EXPECT_FALSE(all_zero) << "Output should not be all zeros"; + EXPECT_TRUE(differs_from_no_rotary) << "Rotary output should differ from non-rotary output"; +} + +// CUDA: shared KV with do_rotary=1, cross-checked against CPU. +TEST(GroupQueryAttentionTest, SharedKV_EmptyKV_WithPast_Rotary_CUDA) { + auto cuda_ep = DefaultCudaExecutionProvider(); + if (!cuda_ep) { + GTEST_SKIP() << "CUDA EP not available"; + } + + constexpr int batch_size = 1; + constexpr int q_seq_len = 1; + constexpr int past_seq_len = 8; + constexpr int num_heads = 2; + constexpr int kv_num_heads = 1; + constexpr int head_size = 16; + constexpr int hidden_size = num_heads * head_size; + + std::vector query_data(batch_size * q_seq_len * hidden_size); + std::vector past_key_data(batch_size * kv_num_heads * past_seq_len * head_size); + std::vector past_value_data(batch_size * kv_num_heads * past_seq_len * head_size); + for (size_t i = 0; i < query_data.size(); i++) query_data[i] = 0.1f * static_cast(i % 7 + 1); + for (size_t i = 0; i < past_key_data.size(); i++) past_key_data[i] = 0.2f * static_cast(i % 5 + 1); + for (size_t i = 0; i < past_value_data.size(); i++) past_value_data[i] = 0.3f * static_cast(i % 3 + 1); + + auto cuda_output = RunGQASharedKVWithRotary( + batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, + num_heads, kv_num_heads, head_size, /*use_cuda=*/true); + auto cpu_output = RunGQASharedKVWithRotary( + batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, + num_heads, kv_num_heads, head_size, /*use_cuda=*/false); + + ExpectOutputsMatch(cuda_output, cpu_output, 0.05f, "SharedKV_Rotary_CUDA_vs_CPU"); +} + +// CUDA: shared KV + rotary, prompt phase (q_seq_len > 1). +TEST(GroupQueryAttentionTest, SharedKV_EmptyKV_WithPast_Rotary_Prompt_CUDA) { + auto cuda_ep = DefaultCudaExecutionProvider(); + if (!cuda_ep) { + GTEST_SKIP() << "CUDA EP not available"; + } + + constexpr int batch_size = 1; + constexpr int q_seq_len = 4; + constexpr int past_seq_len = 4; + constexpr int num_heads = 2; + constexpr int kv_num_heads = 1; + constexpr int head_size = 16; + constexpr int hidden_size = num_heads * head_size; + + std::vector query_data(batch_size * q_seq_len * hidden_size); + std::vector past_key_data(batch_size * kv_num_heads * past_seq_len * head_size); + std::vector past_value_data(batch_size * kv_num_heads * past_seq_len * head_size); + for (size_t i = 0; i < query_data.size(); i++) query_data[i] = 0.1f * static_cast(i % 7 + 1); + for (size_t i = 0; i < past_key_data.size(); i++) past_key_data[i] = 0.2f * static_cast(i % 5 + 1); + for (size_t i = 0; i < past_value_data.size(); i++) past_value_data[i] = 0.3f * static_cast(i % 3 + 1); + + auto cuda_output = RunGQASharedKVWithRotary( + batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, + num_heads, kv_num_heads, head_size, /*use_cuda=*/true); + auto cpu_output = RunGQASharedKVWithRotary( + batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, + num_heads, kv_num_heads, head_size, /*use_cuda=*/false); + + ExpectOutputsMatch(cuda_output, cpu_output, 0.05f, "SharedKV_Rotary_Prompt_CUDA_vs_CPU"); +} + } // namespace test } // namespace onnxruntime From 85209d6ca619d331aa65a7a7314d387708b0530a Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 11 May 2026 16:52:20 +0000 Subject: [PATCH 29/32] Fix comments --- .../contrib_ops/cuda/bert/group_query_attention.cc | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc index a7466b40fe12c..e9a5976af0957 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc @@ -295,7 +295,11 @@ Status GroupQueryAttention::ComputeInternal(OpKernelContext* context) cons data.present_key = reinterpret_cast(present_key_output->MutableData()); data.present_value = reinterpret_cast(present_value_output->MutableData()); // Compute past_present_share_buffer early since it's needed for flash attention path selection. - parameters.past_present_share_buffer = (data.past_key != nullptr && data.past_key == data.present_key); + bool past_key_shared = (data.past_key != nullptr && data.past_key == data.present_key); + bool past_value_shared = (data.past_value != nullptr && data.past_value == data.present_value); + ORT_ENFORCE(past_key_shared == past_value_shared, + "past_key/present_key and past_value/present_value must be both shared or both separate."); + parameters.past_present_share_buffer = past_key_shared; bool is_inputs_quantized = (k_quant_type_ != KVQuantizationType::NONE) || (v_quant_type_ != KVQuantizationType::NONE); constexpr bool is_int8 = std::is_same::value; @@ -317,6 +321,7 @@ Status GroupQueryAttention::ComputeInternal(OpKernelContext* context) cons (device_prop.major >= 8) && !parameters.is_first_prompt && parameters.sequence_length == 1 && + parameters.kv_sequence_length > 0 && // Shared KV (kv_seq=0) has no new K/V to append parameters.past_present_share_buffer && parameters.softcap == 0.0f && !parameters.use_smooth_softmax && @@ -394,7 +399,7 @@ Status GroupQueryAttention::ComputeInternal(OpKernelContext* context) cons parameters.kv_num_heads); data.use_flash_attention = use_flash_attention; - data.use_flash_attention_fast_decode = use_flash_attention && !disable_flash_decode_ && !parameters.is_first_prompt && parameters.past_present_share_buffer && !is_inputs_quantized; + data.use_flash_attention_fast_decode = use_flash_attention && !disable_flash_decode_ && !parameters.is_first_prompt && parameters.kv_sequence_length > 0 && parameters.past_present_share_buffer && !is_inputs_quantized; if (use_flash_attention) { // Allocate Flash specific buffers (Softmax LSE, Accum) From d819b56ca79cab0ab993e0c8e2a8da44deb6d65f Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 11 May 2026 17:14:21 +0000 Subject: [PATCH 30/32] address copilot comments --- .../contrib_ops/cuda/bert/group_query_attention_impl.cu | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu index be092056bdc93..4b365cb304a43 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu @@ -127,8 +127,10 @@ Status PrepareQKV( // Shared KV path: K/V inputs are empty (kv_sequence_length == 0) and the // past buffer already contains the full shared KV cache. This requires // past_key/past_value to be provided (with RoPE already applied to K). - // past_present_share_buffer must be true so the present buffer aliases past - // — no copy needed, attention reads directly from the shared KV data. + // When past_present_share_buffer is true, present aliases past and no copy + // is needed. When false (e.g., first prompt), the past→present memcpy + // above has already populated the present buffer with the shared KV data. + // In both cases, only Q processing (RoPE if configured) is needed here. if (kv_sequence_length == 0) { if (parameters.do_rotary && data.cos_cache != nullptr && data.sin_cache != nullptr) { // Apply RoPE to Q only using the standalone rotary embedding kernel. From 3b996e42a87315d09cc58d62d93311b709f76a2f Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 11 May 2026 20:29:11 +0000 Subject: [PATCH 31/32] Improve --- .../cuda/bert/group_query_attention_impl.h | 5 + .../group_query_attention_op_test.cc | 245 +++++++++++++++++- 2 files changed, 238 insertions(+), 12 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.h b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.h index 125ab8f76132c..89945b20fcfb3 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.h @@ -118,6 +118,11 @@ struct GQABufferRequirements { } } + // Unfused fallback: needs Q buffer for rotary embedding output. + if (req.qkv_buffer_bytes == 0 && (params.do_rotary || params.is_packed_qkv)) { + req.qkv_buffer_bytes = elem_size * q_elements; + } + return req; } }; diff --git a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc index 7487e717cffd4..2672d162698f5 100644 --- a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc @@ -420,6 +420,70 @@ static std::vector RunGQASharedKV( return std::vector(out_data, out_data + output_size); } +// Helper: run GQA with MLFloat16 tensors for actual CUDA kernel coverage. +// The CUDA GQA kernel only registers for MLFloat16/BFloat16, so float inputs +// fall back to CPU. This helper converts float inputs to fp16. +static std::vector RunGQASharedKVFp16( + int batch_size, + int q_seq_len, + int past_seq_len, + const std::vector& query_data, + const std::vector& past_key_data, + const std::vector& past_value_data, + int num_heads, + int kv_num_heads, + int head_size) { + const int hidden_size = num_heads * head_size; + const int total_seq_len = past_seq_len; + + OpTester tester("GroupQueryAttention", 1, onnxruntime::kMSDomain); + tester.AddAttribute("num_heads", static_cast(num_heads)); + tester.AddAttribute("kv_num_heads", static_cast(kv_num_heads)); + + tester.AddInput("query", {batch_size, q_seq_len, hidden_size}, ToFloat16(query_data)); + const int kv_hidden_size = kv_num_heads * head_size; + tester.AddInput("key", {batch_size, 0, kv_hidden_size}, {}); + tester.AddInput("value", {batch_size, 0, kv_hidden_size}, {}); + + tester.AddInput("past_key", {batch_size, kv_num_heads, past_seq_len, head_size}, ToFloat16(past_key_data)); + tester.AddInput("past_value", {batch_size, kv_num_heads, past_seq_len, head_size}, ToFloat16(past_value_data)); + + std::vector seqlens_k_data(batch_size, static_cast(total_seq_len - 1)); + tester.AddInput("seqlens_k", {batch_size}, seqlens_k_data); + tester.AddInput("total_sequence_length", {1}, {static_cast(total_seq_len)}); + + tester.AddOptionalInputEdge(); // cos_cache + tester.AddOptionalInputEdge(); // sin_cache + tester.AddOptionalInputEdge(); // position_ids + tester.AddOptionalInputEdge(); // attention_bias + tester.AddOptionalInputEdge(); // head_sink + + const int output_size = batch_size * q_seq_len * hidden_size; + tester.AddOutput("output", {batch_size, q_seq_len, hidden_size}, + std::vector(output_size, MLFloat16(0.0f))); + + const int present_size = batch_size * kv_num_heads * past_seq_len * head_size; + tester.AddOutput("present_key", {batch_size, kv_num_heads, past_seq_len, head_size}, + std::vector(present_size, MLFloat16(0.0f))); + tester.AddOutput("present_value", {batch_size, kv_num_heads, past_seq_len, head_size}, + std::vector(present_size, MLFloat16(0.0f))); + + tester.SetOutputTolerance(1e6f); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); + + auto fetches = tester.GetFetches(); + // Convert fp16 output back to float for comparison + const MLFloat16* out_fp16 = fetches[0].Get().Data(); + std::vector result(output_size); + for (int i = 0; i < output_size; i++) { + result[i] = out_fp16[i].ToFloat(); + } + return result; +} + // CPU: kv_sequence_length=0 with past_key/past_value (shared KV decode). // Validates output is non-zero (attention over past KV produces valid output). // Note: cannot compare against RunGQAAndGetOutput because the two paths have @@ -505,9 +569,9 @@ TEST(GroupQueryAttentionTest, SharedKV_EmptyKV_WithPast_CUDA) { for (size_t i = 0; i < past_key_data.size(); i++) past_key_data[i] = 0.2f * static_cast(i % 5 + 1); for (size_t i = 0; i < past_value_data.size(); i++) past_value_data[i] = 0.3f * static_cast(i % 3 + 1); - auto cuda_output = RunGQASharedKV( + auto cuda_output = RunGQASharedKVFp16( batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, - num_heads, kv_num_heads, head_size, /*use_cuda=*/true); + num_heads, kv_num_heads, head_size); auto cpu_output = RunGQASharedKV( batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, @@ -574,6 +638,80 @@ TEST(GroupQueryAttentionTest, SharedKV_EmptyKV_WithPast_GQARatio8_CPU) { EXPECT_FALSE(all_zero) << "Output should not be all zeros"; } +// CPU: shared KV with batch_size > 1. +TEST(GroupQueryAttentionTest, SharedKV_EmptyKV_WithPast_Batched_CPU) { + constexpr int batch_size = 2; + constexpr int q_seq_len = 1; + constexpr int past_seq_len = 4; + constexpr int num_heads = 2; + constexpr int kv_num_heads = 1; + constexpr int head_size = 8; + constexpr int hidden_size = num_heads * head_size; + + std::vector query_data(batch_size * q_seq_len * hidden_size); + std::vector past_key_data(batch_size * kv_num_heads * past_seq_len * head_size); + std::vector past_value_data(batch_size * kv_num_heads * past_seq_len * head_size); + for (size_t i = 0; i < query_data.size(); i++) query_data[i] = 0.1f * static_cast(i % 7 + 1); + for (size_t i = 0; i < past_key_data.size(); i++) past_key_data[i] = 0.2f * static_cast(i % 5 + 1); + for (size_t i = 0; i < past_value_data.size(); i++) past_value_data[i] = 0.3f * static_cast(i % 3 + 1); + + auto output = RunGQASharedKV( + batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, + num_heads, kv_num_heads, head_size, /*use_cuda=*/false); + + bool all_zero = true; + for (size_t i = 0; i < output.size(); i++) { + EXPECT_FALSE(std::isnan(output[i])) << "NaN at index " << i; + if (output[i] != 0.0f) all_zero = false; + } + EXPECT_FALSE(all_zero) << "Output should not be all zeros"; +} + +// Reject: kv_sequence_length=0 without past_key (shared KV requires past). +TEST(GroupQueryAttentionTest, SharedKV_EmptyKV_NoPast_Rejected) { + constexpr int batch_size = 1; + constexpr int sequence_length = 4; + constexpr int num_heads = 2; + constexpr int kv_num_heads = 1; + constexpr int head_size = 8; + constexpr int hidden_size = num_heads * head_size; + constexpr int kv_hidden_size = kv_num_heads * head_size; + + OpTester tester("GroupQueryAttention", 1, onnxruntime::kMSDomain); + tester.AddAttribute("num_heads", static_cast(num_heads)); + tester.AddAttribute("kv_num_heads", static_cast(kv_num_heads)); + + tester.AddInput("query", {batch_size, sequence_length, hidden_size}, + std::vector(batch_size * sequence_length * hidden_size, 1.0f)); + // K/V: empty [B, 0, kv_hidden] — kv_sequence_length = 0 + tester.AddInput("key", {batch_size, 0, kv_hidden_size}, {}); + tester.AddInput("value", {batch_size, 0, kv_hidden_size}, {}); + // No past_key/past_value + tester.AddOptionalInputEdge(); + tester.AddOptionalInputEdge(); + + tester.AddInput("seqlens_k", {batch_size}, {static_cast(sequence_length - 1)}); + tester.AddInput("total_sequence_length", {1}, {static_cast(sequence_length)}); + tester.AddOptionalInputEdge(); // cos_cache + tester.AddOptionalInputEdge(); // sin_cache + tester.AddOptionalInputEdge(); // position_ids + tester.AddOptionalInputEdge(); // attention_bias + tester.AddOptionalInputEdge(); // head_sink + + tester.AddOutput("output", {batch_size, sequence_length, hidden_size}, + std::vector(batch_size * sequence_length * hidden_size, 0.0f)); + tester.AddOutput("present_key", {batch_size, kv_num_heads, sequence_length, head_size}, + std::vector(batch_size * kv_num_heads * sequence_length * head_size, 0.0f)); + tester.AddOutput("present_value", {batch_size, kv_num_heads, sequence_length, head_size}, + std::vector(batch_size * kv_num_heads * sequence_length * head_size, 0.0f)); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectFailure, + "query and key must have the same sequence length", + {}, nullptr, &execution_providers); +} + // CUDA: kv_sequence_length=0 with past, prompt phase. Cross-checks against CPU. TEST(GroupQueryAttentionTest, SharedKV_EmptyKV_WithPast_Prompt_CUDA) { auto cuda_ep = DefaultCudaExecutionProvider(); @@ -596,9 +734,9 @@ TEST(GroupQueryAttentionTest, SharedKV_EmptyKV_WithPast_Prompt_CUDA) { for (size_t i = 0; i < past_key_data.size(); i++) past_key_data[i] = 0.2f * static_cast(i % 5 + 1); for (size_t i = 0; i < past_value_data.size(); i++) past_value_data[i] = 0.3f * static_cast(i % 3 + 1); - auto cuda_output = RunGQASharedKV( + auto cuda_output = RunGQASharedKVFp16( batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, - num_heads, kv_num_heads, head_size, /*use_cuda=*/true); + num_heads, kv_num_heads, head_size); auto cpu_output = RunGQASharedKV( batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, num_heads, kv_num_heads, head_size, /*use_cuda=*/false); @@ -628,9 +766,9 @@ TEST(GroupQueryAttentionTest, SharedKV_EmptyKV_WithPast_LargeHeadSize_CUDA) { for (size_t i = 0; i < past_key_data.size(); i++) past_key_data[i] = 0.2f * static_cast(i % 7 + 1); for (size_t i = 0; i < past_value_data.size(); i++) past_value_data[i] = 0.3f * static_cast(i % 5 + 1); - auto cuda_output = RunGQASharedKV( + auto cuda_output = RunGQASharedKVFp16( batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, - num_heads, kv_num_heads, head_size, /*use_cuda=*/true); + num_heads, kv_num_heads, head_size); auto cpu_output = RunGQASharedKV( batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, num_heads, kv_num_heads, head_size, /*use_cuda=*/false); @@ -660,9 +798,9 @@ TEST(GroupQueryAttentionTest, SharedKV_EmptyKV_WithPast_GQARatio8_CUDA) { for (size_t i = 0; i < past_key_data.size(); i++) past_key_data[i] = 0.2f * static_cast(i % 7 + 1); for (size_t i = 0; i < past_value_data.size(); i++) past_value_data[i] = 0.3f * static_cast(i % 5 + 1); - auto cuda_output = RunGQASharedKV( + auto cuda_output = RunGQASharedKVFp16( batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, - num_heads, kv_num_heads, head_size, /*use_cuda=*/true); + num_heads, kv_num_heads, head_size); auto cpu_output = RunGQASharedKV( batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, num_heads, kv_num_heads, head_size, /*use_cuda=*/false); @@ -764,6 +902,89 @@ static std::vector RunGQASharedKVWithRotary( return std::vector(out_data, out_data + output_size); } +// Helper: run GQA with MLFloat16 tensors + do_rotary=1 for actual CUDA kernel coverage. +static std::vector RunGQASharedKVWithRotaryFp16( + int batch_size, + int q_seq_len, + int past_seq_len, + const std::vector& query_data, + const std::vector& past_key_data, + const std::vector& past_value_data, + int num_heads, + int kv_num_heads, + int head_size) { + const int hidden_size = num_heads * head_size; + const int total_seq_len = past_seq_len; + const int rotary_dim = head_size; + const int max_seq_len = past_seq_len + 16; + + OpTester tester("GroupQueryAttention", 1, onnxruntime::kMSDomain); + tester.AddAttribute("num_heads", static_cast(num_heads)); + tester.AddAttribute("kv_num_heads", static_cast(kv_num_heads)); + tester.AddAttribute("do_rotary", static_cast(1)); + + tester.AddInput("query", {batch_size, q_seq_len, hidden_size}, ToFloat16(query_data)); + const int kv_hidden_size = kv_num_heads * head_size; + tester.AddInput("key", {batch_size, 0, kv_hidden_size}, {}); + tester.AddInput("value", {batch_size, 0, kv_hidden_size}, {}); + + tester.AddInput("past_key", {batch_size, kv_num_heads, past_seq_len, head_size}, ToFloat16(past_key_data)); + tester.AddInput("past_value", {batch_size, kv_num_heads, past_seq_len, head_size}, ToFloat16(past_value_data)); + + std::vector seqlens_k_data(batch_size, static_cast(total_seq_len - 1)); + tester.AddInput("seqlens_k", {batch_size}, seqlens_k_data); + tester.AddInput("total_sequence_length", {1}, {static_cast(total_seq_len)}); + + const int half_rotary = rotary_dim / 2; + std::vector cos_cache(max_seq_len * half_rotary); + std::vector sin_cache(max_seq_len * half_rotary); + for (int pos = 0; pos < max_seq_len; pos++) { + for (int d = 0; d < half_rotary; d++) { + float freq = 1.0f / std::pow(10000.0f, 2.0f * static_cast(d) / static_cast(rotary_dim)); + cos_cache[pos * half_rotary + d] = std::cos(static_cast(pos) * freq); + sin_cache[pos * half_rotary + d] = std::sin(static_cast(pos) * freq); + } + } + tester.AddInput("cos_cache", {max_seq_len, half_rotary}, ToFloat16(cos_cache)); + tester.AddInput("sin_cache", {max_seq_len, half_rotary}, ToFloat16(sin_cache)); + + std::vector position_ids(batch_size * q_seq_len); + for (int b = 0; b < batch_size; b++) { + int past_len = total_seq_len - q_seq_len; + for (int s = 0; s < q_seq_len; s++) { + position_ids[b * q_seq_len + s] = static_cast(past_len + s); + } + } + tester.AddInput("position_ids", {batch_size, q_seq_len}, position_ids); + + tester.AddOptionalInputEdge(); // attention_bias + tester.AddOptionalInputEdge(); // head_sink + + const int output_size = batch_size * q_seq_len * hidden_size; + tester.AddOutput("output", {batch_size, q_seq_len, hidden_size}, + std::vector(output_size, MLFloat16(0.0f))); + + const int present_size = batch_size * kv_num_heads * past_seq_len * head_size; + tester.AddOutput("present_key", {batch_size, kv_num_heads, past_seq_len, head_size}, + std::vector(present_size, MLFloat16(0.0f))); + tester.AddOutput("present_value", {batch_size, kv_num_heads, past_seq_len, head_size}, + std::vector(present_size, MLFloat16(0.0f))); + + tester.SetOutputTolerance(1e6f); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); + + auto fetches = tester.GetFetches(); + const MLFloat16* out_fp16 = fetches[0].Get().Data(); + std::vector result(output_size); + for (int i = 0; i < output_size; i++) { + result[i] = out_fp16[i].ToFloat(); + } + return result; +} + // CPU: shared KV with do_rotary=1 (Q-only RoPE path). TEST(GroupQueryAttentionTest, SharedKV_EmptyKV_WithPast_Rotary_CPU) { constexpr int batch_size = 1; @@ -823,9 +1044,9 @@ TEST(GroupQueryAttentionTest, SharedKV_EmptyKV_WithPast_Rotary_CUDA) { for (size_t i = 0; i < past_key_data.size(); i++) past_key_data[i] = 0.2f * static_cast(i % 5 + 1); for (size_t i = 0; i < past_value_data.size(); i++) past_value_data[i] = 0.3f * static_cast(i % 3 + 1); - auto cuda_output = RunGQASharedKVWithRotary( + auto cuda_output = RunGQASharedKVWithRotaryFp16( batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, - num_heads, kv_num_heads, head_size, /*use_cuda=*/true); + num_heads, kv_num_heads, head_size); auto cpu_output = RunGQASharedKVWithRotary( batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, num_heads, kv_num_heads, head_size, /*use_cuda=*/false); @@ -855,9 +1076,9 @@ TEST(GroupQueryAttentionTest, SharedKV_EmptyKV_WithPast_Rotary_Prompt_CUDA) { for (size_t i = 0; i < past_key_data.size(); i++) past_key_data[i] = 0.2f * static_cast(i % 5 + 1); for (size_t i = 0; i < past_value_data.size(); i++) past_value_data[i] = 0.3f * static_cast(i % 3 + 1); - auto cuda_output = RunGQASharedKVWithRotary( + auto cuda_output = RunGQASharedKVWithRotaryFp16( batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, - num_heads, kv_num_heads, head_size, /*use_cuda=*/true); + num_heads, kv_num_heads, head_size); auto cpu_output = RunGQASharedKVWithRotary( batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, num_heads, kv_num_heads, head_size, /*use_cuda=*/false); From 42af4d76295e38e5c1cc95fb931e0701bb18eb6f Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 11 May 2026 23:27:54 +0000 Subject: [PATCH 32/32] address comments --- .../contrib_ops/cpu/bert/gqa_attention_base.h | 34 +++++++++---------- 1 file changed, 17 insertions(+), 17 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h b/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h index ce17ca0dcf358..a67683f0e77fd 100644 --- a/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h +++ b/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h @@ -155,16 +155,16 @@ class GQAAttentionBase { // attention_probs(B, N, S, T) = Softmax(attention_probs) // If T is float32, U is float32. If T is float16, U could be float16 or float32. template - void ComputeAttentionProbs(U* attention_probs, - const T* Q, - const T* K, - const T* head_sink, - const int32_t* seqlens_k, - const T* attention_bias, - const size_t batch_size, - const size_t sequence_length, - const size_t kv_sequence_length, - const size_t total_sequence_length, + void ComputeAttentionProbs(U* attention_probs, // output probs [B, N, S, T] + const T* Q, // query [B, N, S, H] (BNSH) + const T* K, // key input [B, N_kv, L, H] (BNSH); L=0 for shared KV + const T* head_sink, // smooth softmax sink per head, or nullptr + const int32_t* seqlens_k, // total_sequence_length - 1 per batch + const T* attention_bias, // additive bias [B|1, N|1, S, T], or nullptr + const size_t batch_size, // batch size + const size_t sequence_length, // Q sequence length (new tokens) + const size_t kv_sequence_length, // K/V input sequence length; 0 for shared KV + const size_t total_sequence_length, // total tokens (past + new) const gsl::span attention_bias_shape, // shape of the attention bias const size_t past_buffer_sequence_length, // sequence length of past state const size_t present_buffer_sequence_length, // sequence length of present state @@ -181,10 +181,10 @@ class GQAAttentionBase { packed_qkv ? SafeInt(num_heads_ + 2 * kv_num_heads_) * sequence_length * head_size : SafeInt(0); const size_t kv_num_heads_factor = num_heads_ / kv_num_heads_; - const size_t q_input_chunk_length = sequence_length * head_size; // S x H - const size_t kv_input_chunk_length = kv_sequence_length * head_size; // L x H - const size_t past_buff_chunk_length = past_buffer_sequence_length * head_size; // L x H - const size_t present_buff_chunk_length = present_buffer_sequence_length * head_size; // T x H + const size_t q_input_chunk_length = sequence_length * head_size; + const size_t kv_input_chunk_length = kv_sequence_length * head_size; + const size_t past_buff_chunk_length = past_buffer_sequence_length * head_size; + const size_t present_buff_chunk_length = present_buffer_sequence_length * head_size; if (present_key && !past_present_share_buffer) { memset((void*)present_key, @@ -427,9 +427,9 @@ class GQAAttentionBase { packed_qkv ? SafeInt(num_heads_ + 2 * kv_num_heads_) * sequence_length * head_size : SafeInt(0); const size_t kv_num_heads_factor = num_heads_ / kv_num_heads_; - const size_t kv_input_chunk_length = kv_sequence_length * head_size; // L x H - const size_t past_buff_chunk_length = past_buffer_sequence_length * head_size; // L x H - const size_t present_buff_chunk_length = present_buffer_sequence_length * head_size; // T x H + const size_t kv_input_chunk_length = kv_sequence_length * head_size; + const size_t past_buff_chunk_length = past_buffer_sequence_length * head_size; + const size_t present_buff_chunk_length = present_buffer_sequence_length * head_size; if (present_value && !past_present_share_buffer) { memset((void*)present_value,