[ExecuTorch][WebGPU] Dynamic resize hook for linear_q4gsw

JCNTH · JCNTH · commit efb58ca1906c · 2026-07-04T10:35:16.000-07:00
Pull Request resolved: #20576 **Make the 4-bit quantized linear serve any live M (rows) from one graph, so a dynamic prefill+decode graph computes correct-size outputs.** **Problem:** `linear_q4gsw` baked its dispatch count, `params.M`, and output shape at `build()` for the max M. On a dynamic-shape graph at a smaller live M (e.g. decode M=1 vs prefill M=S) it would over-dispatch and leave the output sized at the max. **Solution:** - Before: one fixed dispatch sized for the build-time M. - After: a tensor resize hook on the input recomputes the live M from `cur_dims`, rewrites `params.M`, updates the dispatch `workgroup_count_x` for the SAME kernel chosen at build (bicol GEMV / shmem GEMM / register-tiled), and sets the output `cur_dims` (= input dims with the last dim replaced by N). Inert until the input is resized. **Implementation:** - The build-time kernel select (bicol GEMV for M==1, else shmem GEMM for large K/N, else register-tiled) is fixed at build; the hook re-runs `compute_q4gsw_workgroup_count` for whichever of the three the build chose and rewrites the param UBO + output dims for the live M — it does not switch kernels (runtime M-switching is a separate optimization). - `own_uniform_buffer` keeps the param UBO alive so the hook can rewrite it. - Mirrors Vulkan `resize_q4gsw_linear_node` (recompute M-derived dispatch each execute). **Constraints:** Behavior-neutral on static graphs (hook fires only when the input's live M differs from the max). No kernel/WGSL/numerics change. Runtime M-based kernel switching is deliberately out of scope (a later opt diff). Co-authored-with: Claude Code. ghstack-source-id: 399812825 @exported-using-ghexport Differential Revision: [D109906094](https://our.internmc.facebook.com/intern/diff/D109906094/)
diff --git a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp
@@ -18,6 +18,8 @@
 #include <cstdint>
 #include <cstring>
 #include <stdexcept>
+#include <string>
+#include <vector>
 
 namespace executorch::backends::webgpu {
 
@@ -50,6 +52,62 @@ constexpr int64_t kQ4gswShmemTileN = 32;
 constexpr uint32_t kQ4gswShmemMinDim = 4096u;
 constexpr uint32_t kQ4gswShmemNMinDim = 2048u;
 
+// Workgroup count for a linear_q4gsw dispatch (bicol GEMV / shmem GEMM / tiled
+// GEMM), with the range/limit guards shared by the build-time path and the
+// resize hook. use_gemv/use_shmem_gemm are the build-time routing decision (the
+// shader/pipeline is fixed at build); the resize hook re-runs this with live m.
+uint32_t compute_q4gsw_workgroup_count(
+    WGPUDevice device,
+    bool use_gemv,
+    bool use_shmem_gemm,
+    uint32_t m,
+    uint32_t n,
+    uint32_t wg_size,
+    const char* op_name) {
+  if (use_gemv) {
+    // bicol: fixed 64 lanes, 2 output columns/workgroup, grid-strided over
+    // ceil(N/2) column-pairs (M == 1 on this decode path).
+    const uint64_t pairs = (static_cast<uint64_t>(n) + 1u) / 2u;
+    if (pairs == 0u || pairs > UINT32_MAX) {
+      throw std::runtime_error(
+          std::string("WebGPU ") + op_name + ": N/2 out of range");
+    }
+    const uint32_t wgc =
+        utils::clamp_workgroup_count(device, static_cast<uint32_t>(pairs));
+    if (wgc == 0u) {
+      throw std::runtime_error(
+          std::string("WebGPU ") + op_name + ": zero GEMV dispatch");
+    }
+    return wgc;
+  }
+  if (use_shmem_gemm) {
+    // shmem GEMM: one workgroup per tile, no grid-stride -> throw over limit.
+    const int64_t total_wgs = utils::div_up<int64_t>(m, kQ4gswShmemTileM) *
+        utils::div_up<int64_t>(n, kQ4gswShmemTileN);
+    WGPULimits limits = {};
+    const uint32_t max_wgs =
+        wgpuDeviceGetLimits(device, &limits) == WGPUStatus_Success &&
+            limits.maxComputeWorkgroupsPerDimension > 0
+        ? limits.maxComputeWorkgroupsPerDimension
+        : 65535u;
+    if (total_wgs > static_cast<int64_t>(max_wgs)) {
+      throw std::runtime_error(
+          std::string("WebGPU ") + op_name +
+          ": shmem GEMM tile count exceeds the 1D dispatch limit");
+    }
+    return static_cast<uint32_t>(total_wgs);
+  }
+  const int64_t total_tiles = utils::div_up<int64_t>(m, kQ4gswTileM) *
+      utils::div_up<int64_t>(n, kQ4gswTileN);
+  if (total_tiles > static_cast<int64_t>(UINT32_MAX)) {
+    throw std::runtime_error(
+        std::string("WebGPU ") + op_name +
+        ": tile count exceeds the 1D dispatch limit");
+  }
+  return utils::compute_1d_workgroup_count(
+      device, static_cast<uint32_t>(total_tiles), wg_size, op_name);
+}
+
 // et_vk.linear_q4gsw args: [in, weight, scales, group_size, bias, out].
 void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   const int in_id = args.at(0);
@@ -137,44 +195,8 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   const char* shader_src = use_gemv ? kQ4gswLinearCoop4BicolWGSL
       : use_shmem_gemm              ? kQ4gswLinearGemmShmemWGSL
                                     : kQ4gswLinearWGSL;
-  uint32_t workgroup_count;
-  if (use_gemv) {
-    // bicol: fixed 64 lanes, 2 output columns/workgroup, grid-strided over
-    // ceil(N/2) column-pairs (M == 1 on this decode path).
-    const uint64_t pairs = (static_cast<uint64_t>(N) + 1u) / 2u;
-    if (pairs == 0u || pairs > UINT32_MAX) {
-      throw std::runtime_error("WebGPU linear_q4gsw: N/2 out of range");
-    }
-    workgroup_count =
-        utils::clamp_workgroup_count(device, static_cast<uint32_t>(pairs));
-    if (workgroup_count == 0u) {
-      throw std::runtime_error("WebGPU linear_q4gsw: zero GEMV dispatch");
-    }
-  } else if (use_shmem_gemm) {
-    // shmem GEMM: one workgroup per tile, no grid-stride -> throw over limit.
-    const int64_t total_wgs = utils::div_up<int64_t>(M, kQ4gswShmemTileM) *
-        utils::div_up<int64_t>(N, kQ4gswShmemTileN);
-    WGPULimits limits = {};
-    const uint32_t max_wgs =
-        wgpuDeviceGetLimits(device, &limits) == WGPUStatus_Success &&
-            limits.maxComputeWorkgroupsPerDimension > 0
-        ? limits.maxComputeWorkgroupsPerDimension
-        : 65535u;
-    if (total_wgs > static_cast<int64_t>(max_wgs)) {
-      throw std::runtime_error(
-          "WebGPU linear_q4gsw: shmem GEMM tile count exceeds the 1D dispatch limit");
-    }
-    workgroup_count = static_cast<uint32_t>(total_wgs);
-  } else {
-    const int64_t total_tiles = utils::div_up<int64_t>(M, kQ4gswTileM) *
-        utils::div_up<int64_t>(N, kQ4gswTileN);
-    if (total_tiles > static_cast<int64_t>(UINT32_MAX)) {
-      throw std::runtime_error(
-          "WebGPU linear_q4gsw: tile count exceeds the 1D dispatch limit");
-    }
-    workgroup_count = utils::compute_1d_workgroup_count(
-        device, static_cast<uint32_t>(total_tiles), wg_size, "linear_q4gsw");
-  }
+  const uint32_t workgroup_count = compute_q4gsw_workgroup_count(
+      device, use_gemv, use_shmem_gemm, M, N, wg_size, "linear_q4gsw");
 
   // Optional bias: real buffer if present, else a dummy for the fixed layout.
   uint32_t has_bias = 0;
@@ -287,12 +309,78 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   bg_desc.entries = bg_entries;
   WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);
 
-  graph.add_dispatch({pipeline, bind_group, workgroup_count, "linear_q4gsw"});
+  const size_t dispatch_idx = graph.add_dispatch(
+      {pipeline, bind_group, workgroup_count, "linear_q4gsw"});
+
+  // Dynamic shapes: recompute dispatch + params.M for the live M. use_gemv and
+  // use_shmem_gemm are captured (routing is fixed at build); the helper re-runs
+  // the same path's workgroup-count formula with the live m.
+  graph.add_tensor_resize_hook(
+      in_id,
+      [in_id,
+       out_id,
+       M,
+       K,
+       N,
+       K_packed,
+       gs,
+       padded_N,
+       has_bias,
+       wg_size,
+       use_gemv,
+       use_shmem_gemm,
+       dispatch_idx,
+       uniform_buffer](WebGPUGraph& g) {
+        const auto& d = g.cur_dims(in_id);
+        if (d.empty()) {
+          throw std::runtime_error(
+              "WebGPU linear_q4gsw(resize): empty input dims");
+        }
+        const uint64_t numel = utils::numel_of(d);
+        if (numel % static_cast<uint64_t>(K) != 0u) {
+          throw std::runtime_error(
+              "WebGPU linear_q4gsw(resize): live input numel not a multiple "
+              "of K");
+        }
+        const uint32_t m =
+            static_cast<uint32_t>(numel / static_cast<uint64_t>(K));
+        if (m == 0u) {
+          throw std::runtime_error("WebGPU linear_q4gsw(resize): live M == 0");
+        }
+        // Buffers/bind-groups were sized for the build-time max M; a larger
+        // live M would write out of bounds.
+        if (m > M) {
+          throw std::runtime_error(
+              "WebGPU linear_q4gsw(resize): live M exceeds the build-time max");
+        }
+        const uint32_t wgc = compute_q4gsw_workgroup_count(
+            g.device(),
+            use_gemv,
+            use_shmem_gemm,
+            m,
+            N,
+            wg_size,
+            "linear_q4gsw(resize)");
+        Q4gswParams p = {};
+        p.M = m;
+        p.N = N;
+        p.K = K;
+        p.K_packed = K_packed;
+        p.group_size = gs;
+        p.padded_N = padded_N;
+        p.has_bias = has_bias;
+        wgpuQueueWriteBuffer(g.queue(), uniform_buffer, 0, &p, sizeof(p));
+        g.dispatch_at(dispatch_idx).workgroup_count_x = wgc;
+        std::vector<int64_t> od(d.begin(), d.end());
+        od.back() = static_cast<int64_t>(N);
+        g.set_cur_dims(out_id, od);
+      });
 
   wgpuShaderModuleRelease(shader);
   wgpuBindGroupLayoutRelease(bgl);
   wgpuPipelineLayoutRelease(pipeline_layout);
-  wgpuBufferRelease(uniform_buffer);
+  // Graph owns it so the resize hook can rewrite it; freed in the dtor.
+  graph.own_uniform_buffer(uniform_buffer);
 }
 
 } // namespace
