diff --git a/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h b/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h index adc7b623ec8c4..a67683f0e77fd 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; @@ -85,7 +86,16 @@ 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.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) && @@ -110,7 +120,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); @@ -118,12 +128,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); @@ -131,7 +141,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); } @@ -145,15 +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, // 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, // 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 @@ -170,12 +181,12 @@ 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 = 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 (!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)); @@ -207,7 +218,24 @@ 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: 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; // 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; const ptrdiff_t output_offset = SafeInt(i) * sequence_length * present_buffer_sequence_length; @@ -300,7 +328,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_); @@ -382,9 +410,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 @@ -398,11 +427,11 @@ 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 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 (!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)); @@ -441,7 +470,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 ? 0 : total_seqlen - sequence_length; // Assume no padding sequence length + size_t past_seqlen; + if (past_value == nullptr) { + past_seqlen = 0; + } else if (kv_sequence_length == 0) { + past_seqlen = total_seqlen; + } 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 5698bcb659f20..2a8c26cdb44f3 100644 --- a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc @@ -125,6 +125,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)); @@ -132,9 +133,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; @@ -143,6 +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_) { + // 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 = {}; @@ -200,19 +202,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; @@ -233,7 +238,9 @@ Status GroupQueryAttention::Compute(OpKernelContext* context) const { const T* head_sink_data = (head_sink != nullptr) ? head_sink->Data() : nullptr; // 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(), + 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, past_key, past_value, output, present_k, present_v, output_qk, seqlens_k, parameters, allocator, context); } 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 0269523e0f34e..731578f1a27e3 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,26 +238,45 @@ 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; + 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 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, + // 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, " + "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, "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 @@ -293,6 +311,7 @@ Status CheckInputs(const T* query, return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "total_sequence_length must be positive, got ", total_sequence_length, "."); } + int present_sequence_length = std::max(total_sequence_length, past_sequence_length); int rotary_dim = 0; @@ -326,6 +345,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 dfecc2b810a04..e9a5976af0957 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc @@ -227,6 +227,7 @@ Status GroupQueryAttention::ComputeInternal(OpKernelContext* context) cons attention_bias, head_sink, parameters)); + parameters.local_window_size = local_window_size_; parameters.is_unidirectional = is_unidirectional_; parameters.use_smooth_softmax = use_smooth_softmax_ || head_sink != nullptr; @@ -291,13 +292,14 @@ 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 = 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); + 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; @@ -319,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 && @@ -396,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) @@ -561,13 +564,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 (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 70c58e6b8f764..4b365cb304a43 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,54 @@ 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))); + // 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). + // 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. + // 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. + } 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/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/contrib_ops/webgpu/bert/group_query_attention.cc b/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc index cdf88c2f225e8..e3b91bdbb82f4 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, - context.DeviceLimits().maxComputeInvocationsPerWorkgroup)); + onnxruntime::narrow(context.DeviceLimits().maxComputeInvocationsPerWorkgroup))); params.use_smooth_softmax = use_smooth_softmax_; params.rotary_interleaved = rotary_interleaved_; @@ -239,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."); 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 508d8d0f200ac..2672d162698f5 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 @@ -335,5 +336,755 @@ TEST(GroupQueryAttentionTest, SeqlensKScalarRejected) { /*seqlens_k_shape=*/std::vector{}); } +// Helper to compare two output vectors (non-zero check + element-wise tolerance). +static void ExpectOutputsMatch(const std::vector& a, const std::vector& b, + float tolerance, const char* label) { + ASSERT_EQ(a.size(), b.size()) << label << ": output size mismatch"; + bool all_zero = true; + 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"; +} + +// --------------------------------------------------------------------------- +// 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); +} + +// 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 +// 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 = RunGQASharedKVFp16( + batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, + 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); + + 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"; +} + +// 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(); + 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 = RunGQASharedKVFp16( + batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, + 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); + + 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 = RunGQASharedKVFp16( + batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, + 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); + + 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 = RunGQASharedKVFp16( + batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, + 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); + + 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); +} + +// 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; + 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 = RunGQASharedKVWithRotaryFp16( + batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, + 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); + + 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 = RunGQASharedKVWithRotaryFp16( + batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, + 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); + + ExpectOutputsMatch(cuda_output, cpu_output, 0.05f, "SharedKV_Rotary_Prompt_CUDA_vs_CPU"); +} + } // namespace test } // namespace onnxruntime