From f0d238a3e843950be7a2325aa94f4f0689c22ec0 Mon Sep 17 00:00:00 2001
From: Satya Kumar Jandhyala <satya.k.jandhyala@gmail.com>
Date: Tue, 21 Jan 2025 14:29:39 -0800
Subject: [PATCH] Removed unnecessary uniforms in GeneratePositionIdsProgram

---
 .../webgpu/bert/group_query_attention.cc       | 18 ++++++------------
 .../webgpu/bert/group_query_attention.h        | 11 ++---------
 2 files changed, 8 insertions(+), 21 deletions(-)

diff --git a/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc b/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc
index 1632dac648..53de4270ff 100644
--- a/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc
+++ b/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc
@@ -86,18 +86,12 @@ Status GeneratePositionIDsProgram::GenerateShaderCode(ShaderHelper& sh) const {
 }
 
 Status GeneratePositionIDs(onnxruntime::webgpu::ComputeContext& context, const WebgpuAttentionParameters& params, const Tensor* seqlens, Tensor* output_tensor) {
-  GeneratePositionIDsProgram program(params);
-  auto output_size = params.is_first_prompt_ ? 1 : params.batch_size_ * params.sequence_length_;
-  program.AddInput(seqlens)
-      .AddOutput(output_tensor)
-      .AddUniformVariables({{static_cast<uint32_t>(params.batch_size_)},
-                            {static_cast<uint32_t>(params.sequence_length_)},
-                            {static_cast<uint32_t>(params.num_heads_)},
-                            {static_cast<uint32_t>(params.head_size_)},
-                            {static_cast<uint32_t>(params.rotary_dim_)},
-                            {static_cast<uint32_t>(params.rotary_interleaved_)},
-                            {static_cast<uint32_t>(params.is_first_prompt_ ? 0 : 1)},
-                            {static_cast<uint32_t>(params.total_sequence_length_)}})
+  GeneratePositionIDsProgram program;
+  auto output_size = params.batch_size_ * params.sequence_length_;
+  program.CacheHint(params.batch_size_, params.sequence_length_)
+      .AddInput({seqlens, ProgramTensorMetadataDependency::Rank})
+      .AddOutput({output_tensor, ProgramTensorMetadataDependency::Rank})
+      .AddUniformVariables({{static_cast<uint32_t>(params.sequence_length_)}})
       .SetDispatchGroupSize((output_size + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE);
   return context.RunProgram(program);
 }
diff --git a/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.h b/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.h
index 4c0b750e45..55127e5017 100644
--- a/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.h
+++ b/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.h
@@ -18,18 +18,11 @@ Status GeneratePositionIDs(onnxruntime::webgpu::ComputeContext& context, bool is
 
 class GeneratePositionIDsProgram final : public Program<GeneratePositionIDsProgram> {
  public:
-  GeneratePositionIDsProgram(const WebgpuAttentionParameters& params) : Program{"GeneratePositionIDs"}, params_(params) {}
+  GeneratePositionIDsProgram() : Program{"GeneratePositionIDs"} {}
 
   Status GenerateShaderCode(ShaderHelper& sh) const override;
 
-  WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"batch_size", ProgramUniformVariableDataType::Uint32},
-                                          {"sequence_length", ProgramUniformVariableDataType::Uint32},
-                                          {"num_heads", ProgramUniformVariableDataType::Uint32},
-                                          {"head_size", ProgramUniformVariableDataType::Uint32},
-                                          {"rotary_embedding_dim", ProgramUniformVariableDataType::Uint32},
-                                          {"interleaved", ProgramUniformVariableDataType::Uint32},
-                                          {"is_first_prompt", ProgramUniformVariableDataType::Uint32},
-                                          {"total_seqlen", ProgramUniformVariableDataType::Uint32});
+  WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"sequence_length", ProgramUniformVariableDataType::Uint32});
 };
 
 class SplitPackedQKVProgram final : public Program<SplitPackedQKVProgram> {