Make static KV cache work. #23061
Make static KV cache work. #23061
Conversation
| @@ -436,7 +436,7 @@ Status ApplyAttention(const Tensor* Q, const Tensor* K, const Tensor* V, const T | |||
| WebgpuAttentionParameters& parameters, onnxruntime::webgpu::ComputeContext& context, const Tensor* seqlen_k) { | |||
| const int output_count = std::min({context.OutputCount(), 1 + (past_key != nullptr ? 1 : 0) + (past_value != nullptr ? 1 : 0)}); | |||
| const int past_sequence_length = output_count > 1 ? parameters.past_sequence_length_ : 0; | |||
| const int total_sequence_length = past_sequence_length + parameters.kv_sequence_length_; | |||
| const int total_sequence_length = parameters.is_gqa_ && parameters.past_present_share_buffer_ ? parameters.seqlen_present_kv_cache_ : (past_sequence_length + parameters.kv_sequence_length_); | |||
There was a problem hiding this comment.
For GQA, total_sequence_length is read from node input like
seqlen_present_kv_cache is the max buffer length, when past and present share buffer.
There was a problem hiding this comment.
That is correct. This name is a bit confusing because we are using this variable total_sequence_length in different ways to work for both MHA and GQA. I am trying to avoid code duplication. My intention is to use the same implementation of Attention for other variations we want to support so that we get the benefit of any optimizations for all Attention related operators. This way we can limit the binary size.
There was a problem hiding this comment.
Here we are assigning present_sequence_length, which is seqlen_present_kv_cache in GQA parameters. The WebGPU implementation uses CheckInputs implementation in onnxruntime\contrib_ops\cpu\bert\group_query_attention_helper.h
There was a problem hiding this comment.
parameters.total_sequence_length_ should also work - it gets set to *((*total_seqlen).template Data<int32_t>()) in CheckInputs
There was a problem hiding this comment.
The use of the name total_sequence_length is generalized in that sense that it corresponds to the present key and present value buffer sequence_lengths.
There was a problem hiding this comment.
that does not make sense: total_sequence_length is not total_sequence_length ?
|
Please add a test case. |
Description
Fix for GQA static KV cache
Motivation and Context