Repository: microsoft/onnxruntime
File: onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc (C++)

Code before:
      parameters.seqlen_present_kv_cache_,
      parameters.head_size_};
  std::vector<int64_t> 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();

  ORT_ENFORCE(parameters.total_sequence_length_ <= parameters.seqlen_present_kv_cache_, "Total sequence length cannot be greater than the existing KV cache length.");

  Tensor qSplit;
  Tensor kSplit;
  Tensor vSplit;

  Tensor qRotary;
  Tensor kRotary;

  // Use a sliding window if the total sequence exceeds the window's length.
  bool use_sliding_window = (local_window_size_ != -1 && local_window_size_ < parameters.total_sequence_length_);
  bool will_use_flash_attention = false;
  if (head_sink == nullptr && !use_smooth_softmax_ && !use_sliding_window) {
    // Create a temporary parameters copy with is_packed_qkv_ set to false to check if flash attention can be applied after unpacking
    WebgpuAttentionParameters temp_params = parameters;
    temp_params.is_packed_qkv_ = false;
    will_use_flash_attention = CanApplyFlashAttention(attention_bias, present_key, present_value, temp_params, context);
  }

  bool kv_cache_ready = false;
  Tensor* indirect_buffer_ptr = nullptr;
  Tensor indirect_buffer;

  if (parameters.is_packed_qkv_ && do_rotary_) {
    qSplit = context.CreateGPUTensor(query->DataType(), TensorShape({parameters.batch_size_, parameters.sequence_length_, parameters.hidden_size_}));

    // Use the ultimate fused operation when FlashAttention and static KV cache is enabled.
    if (will_use_flash_attention && parameters.past_present_share_buffer_) {
      // Determine if we should use indirect dispatch for graph capture
      const bool use_indirect_dispatch = (parameters.sequence_length_ == 1 &&
                                          seqlen_k != nullptr &&
                                          context.IsGraphCaptureEnabled());

      // Create indirect dispatch buffer if needed
      if (use_indirect_dispatch) {
        const TensorShape indirect_buffer_shape{3};  // 3 uint32 values for dispatch dimensions
        indirect_buffer = context.CreateGPUTensor(DataTypeImpl::GetType<uint32_t>(), indirect_buffer_shape);
        indirect_buffer_ptr = &indirect_buffer;
      }

      // Fused: splitQKV + rotary QK + copy KV cache
      ORT_RETURN_IF_ERROR(RunSplitPackedQKVWithRotaryEmbeddingAndCopyKV(context, parameters,
                                                                        query, seqlen_k,
                                                                        cos_cache, sin_cache,

Review comment (performance):
I mean let RunSplitPackedQKVWithRotaryEmbeddingAndCopyKV invisible for group_query_attetnion.cc. In this way, we don't need to create `indirect_buffer` in advance.
```
if (will_use_flash_attention && parameters.past_present_share_buffer_) {
    return ApplyFlashAttention(query, key, value, attention_bias, output, past_key, present_key, past_value,
                                                 present_value, parameters, context, seqlen_k, cos_cache, sin_cache);  
}

Code after:
      parameters.seqlen_present_kv_cache_,
      parameters.head_size_};
  std::vector<int64_t> 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();

  ORT_ENFORCE(parameters.total_sequence_length_ <= parameters.seqlen_present_kv_cache_, "Total sequence length cannot be greater than the existing KV cache length.");

  Tensor qSplit;
  Tensor kSplit;
  Tensor vSplit;

  Tensor qRotary;
  Tensor kRotary;

  // Use a sliding window if the total sequence exceeds the window's length.
  bool use_sliding_window = (local_window_size_ != -1 && local_window_size_ < parameters.total_sequence_length_);
  bool will_use_flash_attention = false;
  if (head_sink == nullptr && !use_smooth_softmax_ && !use_sliding_window) {
    // Create a temporary parameters copy with is_packed_qkv_ set to false to check if flash attention can be applied after unpacking
    WebgpuAttentionParameters temp_params = parameters;
    temp_params.is_packed_qkv_ = false;
    will_use_flash_attention = CanApplyFlashAttention(attention_bias, present_key, present_value, temp_params, context);
  }

  if (parameters.is_packed_qkv_ && do_rotary_) {
    // Use the ultimate fused operation when FlashAttention and static KV cache is enabled.
    if (will_use_flash_attention && parameters.past_present_share_buffer_) {
      // Directly call ApplyFlashAttention with fused split/rotary/copyKV enabled
      // query points to packed QKV, K and V are nullptr since they're not needed
      return ApplyFlashAttention(query, nullptr, nullptr, attention_bias, output, past_key, present_key, past_value,
                                 present_value, parameters, context, seqlen_k, cos_cache, sin_cache);
    }
    // Fused: splitQKV + rotary QK
    qSplit = context.CreateGPUTensor(query->DataType(), TensorShape({parameters.batch_size_, parameters.sequence_length_, parameters.hidden_size_}));
    kSplit = context.CreateGPUTensor(query->DataType(), TensorShape({parameters.batch_size_, parameters.sequence_length_, parameters.kv_hidden_size_}));
    vSplit = context.CreateGPUTensor(query->DataType(), TensorShape({parameters.batch_size_, parameters.sequence_length_, parameters.kv_hidden_size_}));
    ORT_RETURN_IF_ERROR(RunSplitPackedQKVWithRotaryEmbedding(context, parameters,
                                                             query, seqlen_k,
                                                             cos_cache, sin_cache,
                                                             &qSplit, &kSplit, &vSplit));
    parameters.is_packed_qkv_ = false;
    parameters.qkv_format_ = Q_K_V_BSNH;
    query = &qSplit;
    key = &kSplit;
    value = &vSplit;
  } else {
    if (parameters.is_packed_qkv_) {
      // splitQKV
      qSplit = context.CreateGPUTensor(query->DataType(), TensorShape({parameters.batch_size_, parameters.sequence_length_, parameters.hidden_size_}));