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

Code before:
  }
  // Adjust dimensions for vectorization
  const auto half_rotary_embedding_dim_vec = half_rotary_embedding_dim / components;
  const auto head_size_vec = head_size / components;

  // Dispatch: batch_size * sequence_length * num_heads * (half_rotary_dim + need_copy_dim)
  // work_per_head = half_rotary_dim + (head_size - 2 * half_rotary_dim)
  //               = head_size - half_rotary_dim
  const auto work_per_head = head_size_vec - half_rotary_embedding_dim_vec;
  auto dispatch_size = static_cast<uint32_t>(params.batch_size_ * params.sequence_length_ * params.num_heads_ * work_per_head);

  // Extract present_sequence_length from present_key tensor shape
  const uint32_t present_sequence_length = gsl::narrow_cast<uint32_t>(present_key->Shape()[2]);

  // Determine if we should use graph capture features
  const bool use_seqlen_k = (seqlen_k != nullptr && context.IsGraphCaptureEnabled());
  const bool prepare_indirect_dispatch = (indirect_buffer != nullptr);

  constexpr uint32_t tile_size = 64;

  SplitPackedQKVWithRotaryEmbeddingAndCopyKVProgram program(params.rotary_interleaved_, use_seqlen_k, prepare_indirect_dispatch);
  program
      .CacheHint(params.rotary_interleaved_, use_seqlen_k, prepare_indirect_dispatch)
      .AddInput({packedQKV, ProgramTensorMetadataDependency::Rank, components})
      .AddInputs({
          {seqlen_k, ProgramTensorMetadataDependency::Rank},
          {cos_cache, ProgramTensorMetadataDependency::Rank, components},
          {sin_cache, ProgramTensorMetadataDependency::Rank, components},
      });

  if (use_seqlen_k) {
    program.AddInput({seqlen_k, ProgramTensorMetadataDependency::None});
  }

  program.AddOutputs({{query, ProgramTensorMetadataDependency::Rank, components},
                      {present_key, ProgramTensorMetadataDependency::Rank, components},
                      {present_value, ProgramTensorMetadataDependency::Rank, components}});

  if (prepare_indirect_dispatch) {
    program.AddOutput({indirect_buffer, ProgramTensorMetadataDependency::None});
  }

  program.AddUniformVariables({
      {static_cast<uint32_t>(params.sequence_length_)},
      {static_cast<uint32_t>(params.hidden_size_ / components)},
      {static_cast<uint32_t>(params.kv_hidden_size_ / components)},
      {static_cast<uint32_t>(params.num_heads_)},
      {static_cast<uint32_t>(params.kv_num_heads_)},
      {head_size_vec},
      {half_rotary_embedding_dim_vec},
      {present_sequence_length},

Review comment (refactor):
Since you always need seqlen_k in your shader, `use_seqlen_k` can be removed. It's not necessary to add duplicated `seqlen_k` in line 222-224.

Code after:
      .CacheHint(params.rotary_interleaved_)
      .AddInputs({
          {query_in, ProgramTensorMetadataDependency::TypeAndRank},
          {key_in, ProgramTensorMetadataDependency::Rank},
          {seqlen_k, ProgramTensorMetadataDependency::TypeAndRank},
          {cos_cache, ProgramTensorMetadataDependency::Rank},
          {sin_cache, ProgramTensorMetadataDependency::Rank},
      })
      .AddOutputs({
          {query_out, ProgramTensorMetadataDependency::None},
          {key_out, ProgramTensorMetadataDependency::None},
      })
      .SetDispatchGroupSize((q_domain_size + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE)
      .AddUniformVariables({
          {params.scale_},
          {gsl::make_span(q_global_dims)},
          {gsl::make_span(q_global_strides)},
          {gsl::make_span(q_input_output_strides)},
          {gsl::make_span(k_global_dims)},
          {gsl::make_span(k_input_output_strides)},
          {q_domain_size},
      });

  return context.RunProgram(program);
}

Status GroupQueryAttention::ComputeInternal(onnxruntime::webgpu::ComputeContext& context) const {
  const Tensor* query = context.Input<Tensor>(0);
  const Tensor* key = context.Input<Tensor>(1);
  const Tensor* value = context.Input<Tensor>(2);
  const Tensor* past_key = context.Input<Tensor>(3);
  const Tensor* past_value = context.Input<Tensor>(4);
  const Tensor* seqlen_k = context.Input<Tensor>(5);
  const Tensor* total_seqlen_tensor = context.Input<Tensor>(6);
  const Tensor* cos_cache = context.Input<Tensor>(7);
  const Tensor* sin_cache = context.Input<Tensor>(8);
  const Tensor* position_ids = context.Input<Tensor>(9);  // TODO: support sliding window
  const Tensor* attention_bias = context.Input<Tensor>(10);
  const Tensor* head_sink = context.Input<Tensor>(11);

  GroupQueryAttentionParameters params = {};
  ORT_RETURN_IF_ERROR(group_query_attention_helper::CheckInputs(query,
                                                                key,
                                                                value,
                                                                past_key,
                                                                past_value,
                                                                cos_cache,
                                                                sin_cache,
                                                                &params,
                                                                num_heads_,
                                                                kv_num_heads_,