[ExecuTorch][WebGPU] Add Gemma 4 MTP operator and route support

backends/vulkan/custom_ops_lib.py

@@ -350,6 +350,42 @@ def linear_q4gsw_backward(ctx, grad_out):
     setup_context=linear_q4gsw_setup_context,
 )
 
+########################
+## scatter_src_unique ##
+########################
+
+
+def scatter_src_unique_impl(
+    self: torch.Tensor,
+    dim: int,
+    index: torch.Tensor,
+    src: torch.Tensor,
+) -> torch.Tensor:
+    normalized_dim = dim if dim >= 0 else dim + self.dim()
+    if normalized_dim != self.dim() - 1:
+        raise ValueError("scatter_src_unique requires the final dimension")
+    if not isinstance(index, FakeTensor):
+        flattened = index.detach().reshape(-1)
+        if torch.unique(flattened).numel() != flattened.numel():
+            raise ValueError("scatter_src_unique requires unique destinations")
+    return torch.scatter(self, dim, index, src)
+
+
+def scatter_src_unique_meta(
+    self: torch.Tensor,
+    dim: int,
+    index: torch.Tensor,
+    src: torch.Tensor,
+) -> torch.Tensor:
+    return torch.empty_like(self)
+
+
+name = "scatter_src_unique"
+lib.define(f"{name}(Tensor self, int dim, Tensor index, Tensor src) -> Tensor")
+lib.impl(name, scatter_src_unique_impl, "CompositeExplicitAutograd")
+lib.impl(name, scatter_src_unique_meta, "Meta")
+scatter_src_unique_op = getattr(getattr(torch.ops, namespace), name)
+
 name = "linear_dq8ca_q4gsw"
 lib.define(
     f"""

backends/webgpu/runtime/WebGPUShaderRegistry.cpp

@@ -23,6 +23,7 @@
 #include <executorch/backends/webgpu/runtime/ops/binary_op/binary_minimum_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/binary_op/binary_mul_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/binary_op/binary_pow_wgsl.h>
+#include <executorch/backends/webgpu/runtime/ops/binary_op/binary_sub_int32_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/binary_op/binary_sub_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/bitwise_not/bitwise_not_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/bmm/bmm_tiled_wgsl.h>
@@ -101,6 +102,7 @@
 #include <executorch/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_wgsl.h>
+#include <executorch/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_m3_shared_bicol_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/quantized_linear/q4gsw_qkv_bk64_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/quantized_linear/q4gsw_requant_wgsl.h>
@@ -116,6 +118,8 @@
 #include <executorch/backends/webgpu/runtime/ops/rope/apply_rotary_emb_interleaved_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/rope/rotary_embedding_hf_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/rope/rotary_embedding_wgsl.h>
+#include <executorch/backends/webgpu/runtime/ops/scatter/scatter_unique_indices_wgsl.h>
+#include <executorch/backends/webgpu/runtime/ops/scatter/scatter_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/sdpa/sdpa_compute_attn_weights_half_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/sdpa/sdpa_compute_attn_weights_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/sdpa/sdpa_compute_out_half_wgsl.h>
@@ -136,6 +140,7 @@
 #include <executorch/backends/webgpu/runtime/ops/to_copy/to_copy_bool_to_float_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/to_copy/to_copy_float_to_int_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/to_copy/to_copy_int_to_float_wgsl.h>
+#include <executorch/backends/webgpu/runtime/ops/topk/topk_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/abs_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/clamp_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/unary/cos_wgsl.h>
@@ -161,7 +166,7 @@
 namespace executorch::backends::webgpu {
 namespace {
 
-constexpr std::array<WebGPUShaderInfo, 143> kShaderRegistry = {{
+constexpr std::array<WebGPUShaderInfo, 148> kShaderRegistry = {{
     {
         "abs",
         kAbsWGSL,
@@ -288,6 +293,13 @@ constexpr std::array<WebGPUShaderInfo, 143> kShaderRegistry = {{
         kBinarySubWorkgroupSizeY,
         kBinarySubWorkgroupSizeZ,
     },
+    {
+        "binary_sub_int32",
+        kBinarySubInt32WGSL,
+        kBinarySubInt32WorkgroupSizeX,
+        kBinarySubInt32WorkgroupSizeY,
+        kBinarySubInt32WorkgroupSizeZ,
+    },
     {
         "bitwise_not",
         kBitwiseNotWGSL,
@@ -792,6 +804,13 @@ constexpr std::array<WebGPUShaderInfo, 143> kShaderRegistry = {{
         kQ4gswLinearGemmSteelHalfPwdqF16accWorkgroupSizeY,
         kQ4gswLinearGemmSteelHalfPwdqF16accWorkgroupSizeZ,
     },
+    {
+        "q4gsw_linear_m3_shared_bicol",
+        kQ4gswLinearM3SharedBicolWGSL,
+        kQ4gswLinearM3SharedBicolWorkgroupSizeX,
+        kQ4gswLinearM3SharedBicolWorkgroupSizeY,
+        kQ4gswLinearM3SharedBicolWorkgroupSizeZ,
+    },
     {
         "q4gsw_qkv_bk64",
         kQ4gswQkvBk64WGSL,
@@ -967,6 +986,20 @@ constexpr std::array<WebGPUShaderInfo, 143> kShaderRegistry = {{
         kRsqrtWorkgroupSizeY,
         kRsqrtWorkgroupSizeZ,
     },
+    {
+        "scatter",
+        kScatterWGSL,
+        kScatterWorkgroupSizeX,
+        kScatterWorkgroupSizeY,
+        kScatterWorkgroupSizeZ,
+    },
+    {
+        "scatter_unique_indices",
+        kScatterUniqueIndicesWGSL,
+        kScatterUniqueIndicesWorkgroupSizeX,
+        kScatterUniqueIndicesWorkgroupSizeY,
+        kScatterUniqueIndicesWorkgroupSizeZ,
+    },
     {
         "sdpa_compute_attn_weights",
         kSdpaComputeAttnWeightsWGSL,
@@ -1128,6 +1161,13 @@ constexpr std::array<WebGPUShaderInfo, 143> kShaderRegistry = {{
         kToCopyIntToFloatWorkgroupSizeY,
         kToCopyIntToFloatWorkgroupSizeZ,
     },
+    {
+        "topk",
+        kTopkWGSL,
+        kTopkWorkgroupSizeX,
+        kTopkWorkgroupSizeY,
+        kTopkWorkgroupSizeZ,
+    },
     {
         "update_cache",
         kUpdateCacheWGSL,

backends/webgpu/runtime/ops/binary_op/binary_op.wgsl

@@ -1,6 +1,6 @@
-@group(0) @binding(0) var<storage, read> input1: array<f32>;
-@group(0) @binding(1) var<storage, read> input2: array<f32>;
-@group(0) @binding(2) var<storage, read_write> output: array<f32>;
+@group(0) @binding(0) var<storage, read> input1: array<${SCALAR_TYPE}>;
+@group(0) @binding(1) var<storage, read> input2: array<${SCALAR_TYPE}>;
+@group(0) @binding(2) var<storage, read_write> output: array<${SCALAR_TYPE}>;
 
 struct TensorMeta {
   ndim: u32,
@@ -14,12 +14,12 @@ struct TensorMeta {
 
 override wg_size: u32 = 64u;
 $if USE_ALPHA:
-  override alpha: f32 = 1.0;
+  override alpha: ${ALPHA_TYPE} = ${ALPHA_DEFAULT};
 
 $if INLINE:
   @compute @workgroup_size(wg_size, 1, 1)
 $else:
-  fn op(a: f32, b: f32) -> f32 {
+  fn op(a: ${SCALAR_TYPE}, b: ${SCALAR_TYPE}) -> ${SCALAR_TYPE} {
     return ${OP_EXPR};
   }
 

backends/webgpu/runtime/ops/binary_op/binary_op.yaml

@@ -5,13 +5,22 @@ binary_op:
     INLINE: 0
     SAME_EXPR: input1[idx] + input2[idx]
     BROADCAST_EXPR: input1[l1] + input2[l2]
+    SCALAR_TYPE: f32
+    ALPHA_TYPE: f32
+    ALPHA_DEFAULT: 1.0
   shader_variants:
     - NAME: binary_div
       OP_EXPR: a / b
       USE_ALPHA: 0
     - NAME: binary_sub
       OP_EXPR: a - alpha * b
       USE_ALPHA: 1
+    - NAME: binary_sub_int32
+      OP_EXPR: bitcast<i32>(bitcast<u32>(a) - bitcast<u32>(alpha) * bitcast<u32>(b))
+      USE_ALPHA: 1
+      SCALAR_TYPE: i32
+      ALPHA_TYPE: i32
+      ALPHA_DEFAULT: 1i
     - NAME: binary_minimum
       USE_ALPHA: 0
       INLINE: 1

backends/webgpu/runtime/ops/binary_op/binary_sub_int32_wgsl.h

@@ -0,0 +1,78 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <cstdint>
+
+namespace executorch::backends::webgpu {
+
+// @generated from binary_op.wgsl - DO NOT EDIT.
+// wgsl-sha256: 134151da070a891e539f6ede5974310c623a9a7d379c90e69f91bec56ddc9b29
+inline constexpr const char* kBinarySubInt32WGSL = R"(
+@group(0) @binding(0) var<storage, read> input1: array<i32>;
+@group(0) @binding(1) var<storage, read> input2: array<i32>;
+@group(0) @binding(2) var<storage, read_write> output: array<i32>;
+
+struct TensorMeta {
+  ndim: u32,
+  numel: u32,
+  sizes: array<vec4<u32>, 2>,
+  strides: array<vec4<u32>, 2>,
+}
+@group(0) @binding(3) var<uniform> out_meta: TensorMeta;
+@group(0) @binding(4) var<uniform> in1_meta: TensorMeta;
+@group(0) @binding(5) var<uniform> in2_meta: TensorMeta;
+
+override wg_size: u32 = 64u;
+override alpha: i32 = 1i;
+
+fn op(a: i32, b: i32) -> i32 {
+  return bitcast<i32>(bitcast<u32>(a) - bitcast<u32>(alpha) * bitcast<u32>(b));
+}
+
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+    // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel).
+    let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
+    if (idx >= out_meta.numel) {
+        return;
+    }
+
+    var same = true;
+    for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) {
+        if (in1_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u] ||
+            in2_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u]) {
+            same = false;
+        }
+    }
+    if (same) {
+        output[idx] = op(input1[idx], input2[idx]);
+        return;
+    }
+
+    var rem = idx;
+    var l1: u32 = 0u;
+    var l2: u32 = 0u;
+    for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) {
+        let coord = rem / out_meta.strides[d >> 2u][d & 3u];
+        rem = rem % out_meta.strides[d >> 2u][d & 3u];
+        l1 = l1 + min(coord, in1_meta.sizes[d >> 2u][d & 3u] - 1u) * in1_meta.strides[d >> 2u][d & 3u];
+        l2 = l2 + min(coord, in2_meta.sizes[d >> 2u][d & 3u] - 1u) * in2_meta.strides[d >> 2u][d & 3u];
+    }
+    output[idx] = op(input1[l1], input2[l2]);
+}
+)";
+
+inline constexpr uint32_t kBinarySubInt32WorkgroupSizeX = 64;
+inline constexpr uint32_t kBinarySubInt32WorkgroupSizeY = 1;
+inline constexpr uint32_t kBinarySubInt32WorkgroupSizeZ = 1;
+
+} // namespace executorch::backends::webgpu

backends/webgpu/runtime/ops/boolean_op/BooleanOp.cpp

@@ -72,8 +72,8 @@ void dispatch_bool_op(
   const uint32_t n_words = (numel + 3u) / 4u;
 
   uint32_t wg_size = utils::clamp_workgroup_size(device, wg_size_x);
-  uint32_t workgroup_count =
-      utils::compute_1d_workgroup_count(device, n_words, wg_size, op_name);
+  const utils::WgCount workgroup_count =
+      utils::compute_2d_workgroup_count(device, n_words, wg_size, op_name);
 
   WGPUConstantEntry wg_size_constant = {};
   wg_size_constant.key = {"wg_size", WGPU_STRLEN};
@@ -98,8 +98,12 @@ void dispatch_bool_op(
       &wg_size_constant,
       1);
 
-  const size_t dispatch_idx =
-      graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count});
+  const size_t dispatch_idx = graph.add_dispatch(
+      {bundle.pipeline,
+       bundle.bind_group,
+       workgroup_count.x,
+       op_name,
+       workgroup_count.y});
 
   WGPUBuffer p_buf = params_buf;
   auto resize =
@@ -114,8 +118,10 @@ void dispatch_bool_op(
         BoolOpParams p = {n, scalar, 0u, 0u};
         wgpuQueueWriteBuffer(g.queue(), p_buf, 0, &p, sizeof(p));
         const uint32_t nw = (n + 3u) / 4u;
-        g.dispatch_at(dispatch_idx).workgroup_count_x =
-            utils::compute_1d_workgroup_count(g.device(), nw, wg_size, op_name);
+        const utils::WgCount workgroups =
+            utils::compute_2d_workgroup_count(g.device(), nw, wg_size, op_name);
+        g.dispatch_at(dispatch_idx).workgroup_count_x = workgroups.x;
+        g.dispatch_at(dispatch_idx).workgroup_count_y = workgroups.y;
       };
   graph.add_tensor_resize_hook(self_id, resize);
 

backends/webgpu/runtime/ops/boolean_op/boolean_op.wgsl

@@ -18,8 +18,10 @@ fn elem_bool(i: u32) -> bool {
 
 // One thread per output u32 word packs 4 bool bytes -> no inter-thread race.
 @compute @workgroup_size(wg_size, 1, 1)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let word_idx = gid.x;
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+    let word_idx = gid.x + gid.y * (num_workgroups.x * wg_size);
     let n_words = (params.num_elements + 3u) / 4u;
     if (word_idx >= n_words) {
         return;

backends/webgpu/runtime/ops/boolean_op/compare_eq_wgsl.h

@@ -13,7 +13,7 @@
 namespace executorch::backends::webgpu {
 
 // @generated from boolean_op.wgsl - DO NOT EDIT.
-// wgsl-sha256: 2a3c203d0255086a6e67a9e7cb08538858e8e6cb6a5542f6acebf04b8ccca6b7
+// wgsl-sha256: 558bc966cc511d239ed4901644f515cae9b639bbc4150689c5b2954599417ac9
 inline constexpr const char* kCompareEqWGSL = R"(
 @group(0) @binding(0) var<storage, read> input: array<f32>;
 @group(0) @binding(1) var<storage, read_write> output: array<u32>;
@@ -35,8 +35,10 @@ fn elem_bool(i: u32) -> bool {
 
 // One thread per output u32 word packs 4 bool bytes -> no inter-thread race.
 @compute @workgroup_size(wg_size, 1, 1)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let word_idx = gid.x;
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+    let word_idx = gid.x + gid.y * (num_workgroups.x * wg_size);
     let n_words = (params.num_elements + 3u) / 4u;
     if (word_idx >= n_words) {
         return;

backends/webgpu/runtime/ops/boolean_op/compare_ge_wgsl.h

@@ -13,7 +13,7 @@
 namespace executorch::backends::webgpu {
 
 // @generated from boolean_op.wgsl - DO NOT EDIT.
-// wgsl-sha256: 5fb319a54ed666644f118119639dd8cd33256ea5e5163c4c1392ce6e84b369cc
+// wgsl-sha256: d080502cc35fe57711b986bf0abca83bf365de59d9180601717836c35cab37da
 inline constexpr const char* kCompareGeWGSL = R"(
 @group(0) @binding(0) var<storage, read> input: array<f32>;
 @group(0) @binding(1) var<storage, read_write> output: array<u32>;
@@ -35,8 +35,10 @@ fn elem_bool(i: u32) -> bool {
 
 // One thread per output u32 word packs 4 bool bytes -> no inter-thread race.
 @compute @workgroup_size(wg_size, 1, 1)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let word_idx = gid.x;
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+    let word_idx = gid.x + gid.y * (num_workgroups.x * wg_size);
     let n_words = (params.num_elements + 3u) / 4u;
     if (word_idx >= n_words) {
         return;

backends/webgpu/runtime/ops/boolean_op/compare_gt_wgsl.h

@@ -13,7 +13,7 @@
 namespace executorch::backends::webgpu {
 
 // @generated from boolean_op.wgsl - DO NOT EDIT.
-// wgsl-sha256: b4a35be8a34774a47be457130125336b5accf8d24615514f3fe1559159644491
+// wgsl-sha256: 8ab67c8415a30f2258ef64c109e94ef608a9b4c309f2812c8362ed626d0289c4
 inline constexpr const char* kCompareGtWGSL = R"(
 @group(0) @binding(0) var<storage, read> input: array<f32>;
 @group(0) @binding(1) var<storage, read_write> output: array<u32>;
@@ -35,8 +35,10 @@ fn elem_bool(i: u32) -> bool {
 
 // One thread per output u32 word packs 4 bool bytes -> no inter-thread race.
 @compute @workgroup_size(wg_size, 1, 1)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let word_idx = gid.x;
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+    let word_idx = gid.x + gid.y * (num_workgroups.x * wg_size);
     let n_words = (params.num_elements + 3u) / 4u;
     if (word_idx >= n_words) {
         return;