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

Code before:
       gsl::narrow_cast<uint32_t>(head_size),
       1u});

  const auto k_input_output_strides = std::vector<uint32_t>(
      {gsl::narrow_cast<uint32_t>(key_in->Shape().SizeFromDimension(1)),
       gsl::narrow_cast<uint32_t>(hidden_size_k),
       gsl::narrow_cast<uint32_t>(head_size),
       1u});

  const uint32_t total = q_domain_size + k_domain_size;

  FusedQKRotaryEmbeddingProgram program(params.rotary_interleaved_);
  program
      .CacheHint(params.rotary_interleaved_)
      .AddInputs({
          {query_in, ProgramTensorMetadataDependency::Rank},
          {key_in, ProgramTensorMetadataDependency::Rank},
          {&pos_ids, ProgramTensorMetadataDependency::Rank},
          {cos_cache, ProgramTensorMetadataDependency::Rank},
          {sin_cache, ProgramTensorMetadataDependency::Rank},
      })
      .AddOutputs({
          {query_out, ProgramTensorMetadataDependency::Rank},
          {key_out, ProgramTensorMetadataDependency::Rank},
      })
      .SetDispatchGroupSize((total + 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_global_strides)},
          {gsl::make_span(k_input_output_strides)},
          {q_domain_size},
          {k_domain_size},
      })
      .AddIndices(TensorShape{1, 1});

  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);

Review comment (performance):
Maybe we can optimize it by `total = q_domain_size` not `total = q_domain_size + k_domain_size`. And get the updated Q and K in one pass not one after another in shader since Q has shape [batch_size, sequence_length, num_head, head_size] and K has shape [batch_size, sequence_length, num_head_k, head_size] and num_head_k is always less than or equal to num_head.
Please see comments in shader part.

Code after:
  const auto k_input_output_strides = std::vector<uint32_t>(
      {gsl::narrow_cast<uint32_t>(key_in->Shape().SizeFromDimension(1)),
       gsl::narrow_cast<uint32_t>(hidden_size_k),
       gsl::narrow_cast<uint32_t>(head_size),
       1u});

  // Dispatch computations only over the Q domain, and fuse K write operations using a head-index-based condition.
  FusedQKRotaryEmbeddingProgram program(params.rotary_interleaved_);
  program
      .CacheHint(params.rotary_interleaved_)
      .AddInputs({
          {query_in, ProgramTensorMetadataDependency::Rank},
          {key_in, ProgramTensorMetadataDependency::Rank},
          {&pos_ids, ProgramTensorMetadataDependency::Rank},
          {cos_cache, ProgramTensorMetadataDependency::Rank},
          {sin_cache, ProgramTensorMetadataDependency::Rank},
      })
      .AddOutputs({
          {query_out, ProgramTensorMetadataDependency::Rank},
          {key_out, ProgramTensorMetadataDependency::Rank},
      })
      .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},
      })
      .AddIndices(TensorShape{1, 1});

  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 = {};