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

Code before:
  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},
      {tile_size},
      {static_cast<uint32_t>(dispatch_size)},
  });

  program.SetDispatchGroupSize((dispatch_size + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE);
  return context.RunProgram(program);
}

// Fused Q/K rotary embedding

Review comment (suggestion):
Please use `TypeAndRank`. Otherwise, you can't distinguish f16 from f32 shader. It seems that I missed this one for your previous PRs. Please also correct them.
```suggestion
  program.AddOutputs({{query, ProgramTensorMetadataDependency::TypeAndRank, components},
```

Code after:
          {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_,
                                                                seqlen_k,
                                                                total_seqlen_tensor,
                                                                scale_,
                                                                softcap_));
  params.use_smooth_softmax = use_smooth_softmax_;

  ORT_RETURN_IF_ERROR(group_query_attention_helper::CheckCustomAttentionInputs(position_ids,
                                                                               attention_bias,
                                                                               head_sink,