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

Code before:
  // Indices helper (reuse existing pattern)
  const auto& dummy_indices = shader.AddIndices("dummy_indices", ShaderUsage::None);

  const auto interleaved_str = interleaved_ ? "true" : "false";

  // We flatten dispatch space: first q_domain_size items map to Q rotation domain, next k_domain_size to K.
  shader.MainFunctionBody()
      << "  let total = uniforms.q_domain_size + uniforms.k_domain_size;\n"
      << "  if (global_idx >= total) { return; }\n"
      << "  let half_rotary_dim = uniforms.cos_cache_shape[1];\n"
      // Query branch
      << "  if (global_idx < uniforms.q_domain_size) {\n"
      << "    let q_idx = global_idx;\n"
      << "    let bsnh = q_idx / uniforms.q_global_stride % uniforms.q_global_shape;\n"
      << "    if (bsnh[3] < half_rotary_dim) {\n"
      << "      let pos_ids_idx = " << position_ids.BroadcastedIndicesToOffset("bsnh.xy", dummy_indices) << ";\n"
      << "      let position_id = u32(" << position_ids.GetByOffset("pos_ids_idx") << ") + select(0u, bsnh[1], pos_ids_idx == 0u);\n"
      << "      let i = dot(bsnh, uniforms.q_input_output_stride) + select(0u, bsnh[3], " << interleaved_str << ");\n"
      << "      let j = i + select(half_rotary_dim, 1u, " << interleaved_str << ");\n"
      << "      let cos_v = " << cos_cache.GetByIndices("vec2<u32>(position_id, bsnh[3])") << ";\n"
      << "      let sin_v = " << sin_cache.GetByIndices("vec2<u32>(position_id, bsnh[3])") << ";\n"
      << "      let re = " << q_input.GetByOffset("i") << " * cos_v - " << q_input.GetByOffset("j") << " * sin_v;\n"
      << "      " << q_output.SetByOffset("i", "re") << "\n"
      << "      let im = " << q_input.GetByOffset("i") << " * sin_v + " << q_input.GetByOffset("j") << " * cos_v;\n"
      << "      " << q_output.SetByOffset("j", "im") << "\n"
      << "    } else {\n"
      << "      let k = dot(bsnh, uniforms.q_input_output_stride) + half_rotary_dim;\n"
      << "      " << q_output.SetByOffset("k", q_input.GetByOffset("k")) << "\n"
      << "    }\n"
      << "    return;\n"
      << "  }\n"
      // Key branch
      << "  let k_linear = global_idx - uniforms.q_domain_size;\n"
      << "  if (k_linear < uniforms.k_domain_size) {\n"
      << "    let bsnh = k_linear / uniforms.k_global_stride % uniforms.k_global_shape;\n"
      << "    if (bsnh[3] < half_rotary_dim) {\n"
      << "      let pos_ids_idx = " << position_ids.BroadcastedIndicesToOffset("bsnh.xy", dummy_indices) << ";\n"
      << "      let position_id = u32(" << position_ids.GetByOffset("pos_ids_idx") << ") + select(0u, bsnh[1], pos_ids_idx == 0u);\n"
      << "      let i = dot(bsnh, uniforms.k_input_output_stride) + select(0u, bsnh[3], " << interleaved_str << ");\n"
      << "      let j = i + select(half_rotary_dim, 1u, " << interleaved_str << ");\n"
      << "      let cos_v = " << cos_cache.GetByIndices("vec2<u32>(position_id, bsnh[3])") << ";\n"
      << "      let sin_v = " << sin_cache.GetByIndices("vec2<u32>(position_id, bsnh[3])") << ";\n"
      << "      let re = " << k_input.GetByOffset("i") << " * cos_v - " << k_input.GetByOffset("j") << " * sin_v;\n"
      << "      " << k_output.SetByOffset("i", "re") << "\n"
      << "      let im = " << k_input.GetByOffset("i") << " * sin_v + " << k_input.GetByOffset("j") << " * cos_v;\n"
      << "      " << k_output.SetByOffset("j", "im") << "\n"
      << "    } else {\n"
      << "      let k2 = dot(bsnh, uniforms.k_input_output_stride) + half_rotary_dim;\n"
      << "      " << k_output.SetByOffset("k2", k_input.GetByOffset("k2")) << "\n"
      << "    }\n"
      << "  }\n";

Review comment (suggestion):
I think the `if - else` is not necessary since `bsnh[3]` must be smaller than `half_rotary_dim` here because that q_global_shape[3] = half_rotary_dim = head_size / 2 in GQA.

Code after:
  // Indices helper
  const auto& dummy_indices = shader.AddIndices("dummy_indices", ShaderUsage::None);

  const auto interleaved_str = interleaved_ ? "true" : "false";

  shader.MainFunctionBody()
      << "  if (global_idx >= uniforms.q_domain_size) { return; }\n"
      << "  let half_rotary_dim = uniforms.cos_cache_shape[1];\n"
      << "  let bsnh = global_idx / uniforms.q_global_stride % uniforms.q_global_shape;\n"
      << "  if (bsnh[3] < half_rotary_dim) {\n"
      << "    let pos_ids_idx = " << position_ids.BroadcastedIndicesToOffset("bsnh.xy", dummy_indices) << ";\n"
      << "    let position_id = u32(" << position_ids.GetByOffset("pos_ids_idx") << ") + select(0u, bsnh[1], pos_ids_idx == 0u);\n"
      << "    let cos_v = " << cos_cache.GetByIndices("vec2<u32>(position_id, bsnh[3])") << ";\n"
      << "    let sin_v = " << sin_cache.GetByIndices("vec2<u32>(position_id, bsnh[3])") << ";\n"
      << "    let qi = dot(bsnh, uniforms.q_input_output_stride) + select(0u, bsnh[3], " << interleaved_str << ");\n"
      << "    let qj = qi + select(half_rotary_dim, 1u, " << interleaved_str << ");\n"
      << "    let q_re = " << q_input.GetByOffset("qi") << " * cos_v - " << q_input.GetByOffset("qj") << " * sin_v;\n"
      << "    " << q_output.SetByOffset("qi", "q_re") << "\n"
      << "    let q_im = " << q_input.GetByOffset("qi") << " * sin_v + " << q_input.GetByOffset("qj") << " * cos_v;\n"
      << "    " << q_output.SetByOffset("qj", "q_im") << "\n"
      // Conditionally process Key (only for heads that exist in K domain)
      << "    if (bsnh[2] < uniforms.k_global_shape[2]) {\n"
      << "      let ki = dot(bsnh, uniforms.k_input_output_stride) + select(0u, bsnh[3], " << interleaved_str << ");\n"
      << "      let kj = ki + select(half_rotary_dim, 1u, " << interleaved_str << ");\n"
      << "      let k_re = " << k_input.GetByOffset("ki") << " * cos_v - " << k_input.GetByOffset("kj") << " * sin_v;\n"
      << "      " << k_output.SetByOffset("ki", "k_re") << "\n"
      << "      let k_im = " << k_input.GetByOffset("ki") << " * sin_v + " << k_input.GetByOffset("kj") << " * cos_v;\n"
      << "      " << k_output.SetByOffset("kj", "k_im") << "\n"
      << "    }\n"
      << "  } else {\n"
      << "    let qk = dot(bsnh, uniforms.q_input_output_stride) + half_rotary_dim;\n"
      << "    " << q_output.SetByOffset("qk", q_input.GetByOffset("qk")) << "\n"
      // Conditionally process Key (only for heads that exist in K domain)
      << "    if (bsnh[2] < uniforms.k_global_shape[2]) {\n"
      << "      let kk = dot(bsnh, uniforms.k_input_output_stride) + half_rotary_dim;\n"
      << "      " << k_output.SetByOffset("kk", k_input.GetByOffset("kk")) << "\n"
      << "    }\n"
      << "  }\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));
  num_heads_ = static_cast<int>(info.GetAttrOrDefault<int64_t>("num_heads", 0));
  interleaved_ = (info.GetAttrOrDefault<int64_t>("interleaved", 0) == 1);
  is_packed_batching_ = (info.GetAttrOrDefault<int64_t>("is_packed_batching", 0) == 1);
}

Status RotaryEmbedding::ComputeInternal(onnxruntime::webgpu::ComputeContext& context) const {
  const auto* input = context.Input<Tensor>(0);