Replace RunRotaryEmbeddingQOnly with standalone RotaryEmbeddingWithOffsetProgram

Introduce RotaryEmbeddingWithOffsetProgram as a separate class that computes
position from a uniform offset (position_offset + sequence_index) instead of
requiring a position_ids tensor input. This avoids the need for RangeProgram
dispatch and keeps the original RotaryEmbeddingProgram untouched.

Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>

Author: feich-ms

onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc

@@ -183,22 +183,56 @@ Status RunFusedQKRotaryEmbedding(onnxruntime::webgpu::ComputeContext& context,
   return context.RunProgram(program);
 }
 
-// Apply rotary embedding to Q only. Reuses RunFusedQKRotaryEmbedding with a 1-element
-// dummy K output and query_in as dummy K input. The shader skips K because k_global_dims[2]=0.
-Status RunRotaryEmbeddingQOnly(onnxruntime::webgpu::ComputeContext& context,
-                               const WebgpuAttentionParameters& params,
-                               const Tensor* query_in,
-                               const Tensor* seqlen_k,
-                               const Tensor* cos_cache,
-                               const Tensor* sin_cache,
-                               Tensor* query_out) {
-  Tensor k_dummy_out = context.CreateGPUTensor(query_in->DataType(), TensorShape({1}));
-  // Temporarily patch kv_num_heads to 0 so RunFusedQKRotaryEmbedding builds k_global_shape[2]=0.
-  WebgpuAttentionParameters params_q_only = params;
-  params_q_only.kv_num_heads_ = 0;
-  params_q_only.kv_hidden_size_ = 0;
-  return RunFusedQKRotaryEmbedding(context, params_q_only, query_in, query_in,
-                                   seqlen_k, cos_cache, sin_cache, query_out, &k_dummy_out);
+// Apply rotary embedding to a single tensor using RotaryEmbeddingWithOffsetProgram.
+// Position for each token = past_sequence_length + sequence_index.
+Status RunRotaryEmbedding(onnxruntime::webgpu::ComputeContext& context,
+                          const Tensor* input,
+                          const Tensor* cos_cache,
+                          const Tensor* sin_cache,
+                          Tensor* output,
+                          int batch_size,
+                          int sequence_length,
+                          int hidden_size,
+                          int head_size,
+                          int past_sequence_length,
+                          float scale,
+                          bool rotary_interleaved) {
+  const auto half_rotary_embedding_dim = gsl::narrow_cast<uint32_t>(cos_cache->Shape()[1]);
+  const auto num_heads = hidden_size / head_size;
+
+  const TensorShape global_shape({static_cast<int64_t>(batch_size),
+                                  static_cast<int64_t>(sequence_length),
+                                  static_cast<int64_t>(num_heads),
+                                  static_cast<int64_t>(head_size - half_rotary_embedding_dim)});
+  const auto rank = global_shape.NumDimensions();
+  std::vector<uint32_t> global_dims(rank);
+  std::vector<uint32_t> global_strides(rank);
+  for (size_t j = 0; j < rank; ++j) {
+    global_dims[j] = gsl::narrow_cast<uint32_t>(global_shape[j]);
+    global_strides[j] = gsl::narrow_cast<uint32_t>(global_shape.SizeFromDimension(j + 1));
+  }
+
+  const auto output_size = gsl::narrow_cast<uint32_t>(global_shape.Size());
+  const auto input_output_strides = std::vector<uint32_t>({
+      gsl::narrow_cast<uint32_t>(input->Shape().SizeFromDimension(1)),
+      gsl::narrow_cast<uint32_t>(hidden_size),
+      gsl::narrow_cast<uint32_t>(head_size),
+      1u});
+
+  RotaryEmbeddingWithOffsetProgram program(rotary_interleaved);
+  program
+      .CacheHint(rotary_interleaved)
+      .AddInputs({{input, ProgramTensorMetadataDependency::TypeAndRank},
+                  {cos_cache, ProgramTensorMetadataDependency::Rank},
+                  {sin_cache, ProgramTensorMetadataDependency::Rank}})
+      .AddOutput({output, ProgramTensorMetadataDependency::None})
+      .SetDispatchGroupSize((output_size + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE)
+      .AddUniformVariables({{scale},
+                            {gsl::make_span(global_dims)},
+                            {gsl::make_span(global_strides)},
+                            {gsl::make_span(input_output_strides)},
+                            {static_cast<uint32_t>(past_sequence_length)}});
+  return context.RunProgram(program);
 }
 
 Status GroupQueryAttention::ComputeInternal(onnxruntime::webgpu::ComputeContext& context) const {
@@ -299,10 +333,12 @@ Status GroupQueryAttention::ComputeInternal(onnxruntime::webgpu::ComputeContext&
     if (do_rotary_) {
       // Apply RoPE to Q only — K doesn't need rotation since we reuse another layer's already-rotated KV cache.
       qRotary = context.CreateGPUTensor(query->DataType(), query->Shape());
-      ORT_RETURN_IF_ERROR(RunRotaryEmbeddingQOnly(context, parameters,
-                                                  query, seqlen_k,
-                                                  cos_cache, sin_cache,
-                                                  &qRotary));
+      ORT_RETURN_IF_ERROR(RunRotaryEmbedding(context,
+                                                  query, cos_cache, sin_cache, &qRotary,
+                                                  parameters.batch_size_, parameters.sequence_length_,
+                                                  parameters.hidden_size_, parameters.head_size_,
+                                                  parameters.past_sequence_length_,
+                                                  parameters.scale_, parameters.rotary_interleaved_));
       query = &qRotary;
     }
   } else if (parameters.is_packed_qkv_ && do_rotary_) {

onnxruntime/contrib_ops/webgpu/bert/rotary_embedding.cc

@@ -134,6 +134,38 @@ Status FusedQKRotaryEmbeddingProgram::GenerateShaderCode(ShaderHelper& shader) c
   return Status::OK();
 }
 
+Status RotaryEmbeddingWithOffsetProgram::GenerateShaderCode(ShaderHelper& shader) const {
+  const auto& input = shader.AddInput("input", ShaderUsage::UseUniform);
+  const auto& cos_cache = shader.AddInput("cos_cache", ShaderUsage::UseUniform);
+  const auto& sin_cache = shader.AddInput("sin_cache", ShaderUsage::UseUniform);
+  const auto& output = shader.AddOutput("output", ShaderUsage::UseUniform);
+  const auto interleaved_str = interleaved_ ? "true" : "false";
+  shader.MainFunctionBody() << "  let half_rotary_emb_dim = uniforms.cos_cache_shape[1];\n"
+                               "  let bsnh = global_idx / uniforms.global_stride % uniforms.global_shape;\n"
+                               "  let size = uniforms.global_shape[0] * uniforms.global_stride[0];\n"
+                               "  if (global_idx >= size) { return; }\n"
+                               "  if (bsnh[3] < half_rotary_emb_dim) {\n"
+                               "    let position_id = uniforms.position_offset + bsnh[1];\n"
+                            << "    let i = dot(bsnh, uniforms.input_output_stride) + select(0, bsnh[3], " << interleaved_str << ");\n"
+                            << "    let j = i + select(half_rotary_emb_dim, 1, " << interleaved_str << ");\n"
+                               "    let max_position = uniforms.cos_cache_shape[0];\n"
+                               "    if (position_id >= max_position) {\n"
+                            << "      " << output.SetByOffset("i", input.GetByOffset("i")) << "\n"
+                            << "      " << output.SetByOffset("j", input.GetByOffset("j")) << "\n"
+                               "    } else {\n"
+                            << "      let re = " << input.GetByOffset("i") << " * " << cos_cache.GetByIndices("vec2<u32>(position_id, bsnh[3])") << " - " << input.GetByOffset("j") << " * " << sin_cache.GetByIndices("vec2<u32>(position_id, bsnh[3])") << ";\n"
+                            << "      " << output.SetByOffset("i", "re") << "\n"
+                            << "      let im = " << input.GetByOffset("i") << " * " << sin_cache.GetByIndices("vec2<u32>(position_id, bsnh[3])") << " + " << input.GetByOffset("j") << " * " << cos_cache.GetByIndices("vec2<u32>(position_id, bsnh[3])") << ";\n"
+                            << "      " << output.SetByOffset("j", "im") << "\n"
+                               "    }\n"
+                            << "  } else { \n"
+                               "    let k = dot(bsnh, uniforms.input_output_stride) + half_rotary_emb_dim;\n"
+                            << "    " << output.SetByOffset("k", input.GetByOffset("k")) << "\n"
+                            << "  }";
+
+  return Status::OK();
+}
+
 RotaryEmbedding::RotaryEmbedding(const OpKernelInfo& info) : WebGpuKernel(info) {
   scale_ = info.GetAttrOrDefault<float>("scale", 1.0);
   rotary_embedding_dim_ = static_cast<int>(info.GetAttrOrDefault<int64_t>("rotary_embedding_dim", 0));

onnxruntime/contrib_ops/webgpu/bert/rotary_embedding.h

@@ -29,6 +29,23 @@ class RotaryEmbeddingProgram final : public Program<RotaryEmbeddingProgram> {
   const bool interleaved_;
 };
 
+class RotaryEmbeddingWithOffsetProgram final : public Program<RotaryEmbeddingWithOffsetProgram> {
+ public:
+  RotaryEmbeddingWithOffsetProgram(bool interleaved)
+      : Program{"RotaryEmbeddingWithOffset"}, interleaved_{interleaved} {}
+
+  Status GenerateShaderCode(ShaderHelper& sh) const override;
+
+  WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"scale", ProgramUniformVariableDataType::Float32},
+                                          {"global_shape", ProgramUniformVariableDataType::Uint32},
+                                          {"global_stride", ProgramUniformVariableDataType::Uint32},
+                                          {"input_output_stride", ProgramUniformVariableDataType::Uint32},
+                                          {"position_offset", ProgramUniformVariableDataType::Uint32});
+
+ private:
+  const bool interleaved_;
+};
+
 class FusedQKRotaryEmbeddingProgram final : public Program<FusedQKRotaryEmbeddingProgram> {
  public:
   FusedQKRotaryEmbeddingProgram(bool interleaved) : Program{"FusedQKRotaryEmbedding"}, interleaved_{interleaved} {}