[ExecuTorch][WebGPU] Add q4gsw training kernels (backward, weight-grad, requant) to the WebGPU backend

Pull Request resolved: #20931

Add the 4-bit fine-tuning weight-path custom ops — `et_vk.linear_q4gsw_backward` (input grad dX), `et_vk.linear_dW` (weight grad dW — a generic fp32 weight-gradient matmul, not q4gsw-specific), and `et_vk.q4gsw_requant` (STE re-quant of fp32 latents → 4-bit codes) — that close the on-device training loop through a frozen 4-bit base.

Key changes:
- `runtime/ops/quantized_linear/{QuantizedLinearBackward,QuantizedLinearRequant}.cpp` + `LinearDw.cpp` + `q4gsw_{backward,requant}.wgsl` + `linear_dW.wgsl` — the three training kernels
- `custom_ops_lib.py` — 3 custom-op defs (`linear_q4gsw_backward` + `register_autograd` on `linear_q4gsw`, `linear_dW`, `q4gsw_requant`)
- `op_registry.py` — 3 `OpFeatures`
- `CMakeLists.txt` `WEBGPU_SRCS` — wire the sources

Training-only, WebGPU-only custom ops registered under the shared Vulkan partitioner; Vulkan has no kernels for them yet (a separate Vulkan-kernel track covers that). Guards mirror the sibling `Dw`/`Requant` ops (scales-rank + `K==0`).

Co-authored-with: Claude Code.
ghstack-source-id: 405026076
@exported-using-ghexport

Differential Revision: [D111755130](https://our.internmc.facebook.com/intern/diff/D111755130/)

Author: JCNTH

backends/vulkan/custom_ops_lib.py

@@ -289,6 +289,67 @@ def linear_dq8ca_q4gsw(
 lib.impl(name, linear_q4gsw, "CompositeExplicitAutograd")
 linear_qc4w_op = getattr(getattr(torch.ops, namespace), name)
 
+
+# Backward of linear_q4gsw wrt input (for on-device LoRA training through a frozen
+# 4-bit base): d_x = d_out @ dequant(W). Reference impl extracts dequant(W) via the
+# forward on an identity so it is layout-agnostic; the runtime dispatches this op to
+# the tiled q4gsw_backward WGSL kernel (contracts over N).
+def linear_q4gsw_backward_impl(
+    d_out: torch.Tensor,
+    weights: torch.Tensor,
+    weight_scales: torch.Tensor,
+    group_size: int,
+) -> torch.Tensor:
+    in_features = int(weights.shape[1]) * 2
+    eye = torch.eye(in_features, dtype=d_out.dtype, device=d_out.device)
+    w_t = linear_q4gsw(eye, weights, weight_scales, group_size)  # [in, out]
+    return d_out @ w_t.t()  # [M, out] @ [out, in] = [M, in]
+
+
+def linear_q4gsw_backward_meta(
+    d_out: torch.Tensor,
+    weights: torch.Tensor,
+    weight_scales: torch.Tensor,
+    group_size: int,
+) -> torch.Tensor:
+    return d_out.new_empty(d_out.shape[:-1] + (int(weights.shape[1]) * 2,))
+
+
+name = "linear_q4gsw_backward"
+lib.define(
+    f"{name}(Tensor d_out, Tensor weights, Tensor weight_scales, int group_size) -> Tensor"
+)
+lib.impl(name, linear_q4gsw_backward_impl, "CompositeExplicitAutograd")
+lib.impl(name, linear_q4gsw_backward_meta, "Meta")
+linear_q4gsw_backward_op = getattr(getattr(torch.ops, namespace), name)
+
+
+def linear_q4gsw_setup_context(ctx, inputs, output) -> None:
+    _x, weights, weight_scales, group_size, _bias = inputs
+    ctx.save_for_backward(weights, weight_scales)
+    ctx.group_size = group_size
+
+
+def linear_q4gsw_backward(ctx, grad_out):
+    weights, weight_scales = ctx.saved_tensors
+    d_x = torch.ops.et_vk.linear_q4gsw_backward(
+        grad_out, weights, weight_scales, ctx.group_size
+    )
+    return (
+        d_x,
+        None,
+        None,
+        None,
+        None,
+    )  # grads for (x, weights, scales, group_size, bias)
+
+
+torch.library.register_autograd(
+    f"{namespace}::linear_q4gsw",
+    linear_q4gsw_backward,
+    setup_context=linear_q4gsw_setup_context,
+)
+
 name = "linear_dq8ca_q4gsw"
 lib.define(
     f"""
@@ -1090,3 +1151,69 @@ def rms_norm_impl(
 lib.define(f"{name}(Tensor x, Tensor weight, float eps) -> Tensor")
 lib.impl(name, rms_norm_impl, "CompositeExplicitAutograd")
 rms_norm_op = getattr(getattr(torch.ops, namespace), name)
+
+
+# STE weight gradient d_out^T @ x through the frozen 4-bit linear_q4gsw base.
+def linear_dW_impl(
+    d_out: torch.Tensor,
+    x: torch.Tensor,
+) -> torch.Tensor:
+    return d_out.reshape(-1, d_out.shape[-1]).t() @ x.reshape(-1, x.shape[-1])
+
+
+def linear_dW_meta(
+    d_out: torch.Tensor,
+    x: torch.Tensor,
+) -> torch.Tensor:
+    return d_out.new_empty((d_out.shape[-1], x.shape[-1]))
+
+
+name = "linear_dW"
+lib.define(f"{name}(Tensor d_out, Tensor x) -> Tensor")
+lib.impl(name, linear_dW_impl, "CompositeExplicitAutograd")
+lib.impl(name, linear_dW_meta, "Meta")
+linear_dW_op = getattr(getattr(torch.ops, namespace), name)
+
+
+##################
+## q4gsw_requant ##
+##################
+
+
+# STE re-quant of fp32 latent weights into the frozen-scale 4-bit codes.
+def q4gsw_requant_impl(
+    latent: torch.Tensor,
+    scales: torch.Tensor,
+    group_size: int,
+) -> torch.Tensor:
+    n, k = latent.shape
+    group_idx = torch.arange(k, device=latent.device) // group_size
+    scale_full = scales.t()[:, group_idx]  # [N, K]: scales[k // group_size, n]
+    nonzero = scale_full != 0
+    safe = torch.where(nonzero, scale_full, torch.ones_like(scale_full))
+    q = torch.round(latent / safe)
+    q = torch.where(nonzero, q, torch.zeros_like(q))
+    codes = (torch.clamp(q, -8, 7).to(torch.int32) + 8) & 0xF  # [N, K] in 0..15
+    k_packed = (k + 1) // 2
+    packed = torch.zeros((n, k_packed), dtype=torch.uint8, device=latent.device)
+    packed[:, :] = codes[:, 0::2].to(torch.uint8)
+    if k > 1:
+        high = codes[:, 1::2].to(torch.uint8)
+        packed[:, : high.shape[1]] |= high << 4
+    return packed
+
+
+def q4gsw_requant_meta(
+    latent: torch.Tensor,
+    scales: torch.Tensor,
+    group_size: int,
+) -> torch.Tensor:
+    n, k = latent.shape
+    return latent.new_empty((n, (k + 1) // 2), dtype=torch.uint8)
+
+
+name = "q4gsw_requant"
+lib.define(f"{name}(Tensor latent, Tensor scales, int group_size) -> Tensor")
+lib.impl(name, q4gsw_requant_impl, "CompositeExplicitAutograd")
+lib.impl(name, q4gsw_requant_meta, "Meta")
+q4gsw_requant_op = getattr(getattr(torch.ops, namespace), name)

backends/vulkan/op_registry.py

@@ -462,6 +462,15 @@ def register_quantizedlinear_cpp_ops():
     )
 
 
+@update_features(exir_ops.edge.et_vk.linear_q4gsw_backward.default)
+def register_linear_q4gsw_backward():
+    return OpFeatures(
+        inputs_storage=utils.CONTIGUOUS_ANY,
+        inputs_dtypes=utils.FP_T,
+        supports_prepacking=True,
+    )
+
+
 @update_features(exir_ops.edge.et_vk.linear_dq8ca_q4gsw.default)
 def register_linear_dq8ca_q4gsw():
     return OpFeatures(
@@ -1775,6 +1784,23 @@ def register_logical_not():
     )
 
 
+@update_features(exir_ops.edge.et_vk.linear_dW.default)
+def register_linear_dW():
+    return OpFeatures(
+        inputs_storage=utils.CONTIGUOUS_ANY,
+        inputs_dtypes=utils.FP_T,
+        supports_prepacking=True,
+    )
+
+
+@update_features(exir_ops.edge.et_vk.q4gsw_requant.default)
+def register_q4gsw_requant():
+    return OpFeatures(
+        inputs_storage=utils.CONTIGUOUS_ANY,
+        inputs_dtypes=utils.FP_T,
+    )
+
+
 #######################
 ## Utility functions ##
 #######################

backends/webgpu/CMakeLists.txt

@@ -38,6 +38,7 @@ set(WEBGPU_SRCS
     runtime/ops/sdpa/Sdpa.cpp
     runtime/ops/select_as_symint/SelectAsSymint.cpp
     runtime/ops/quantized_linear/QuantizedLinear.cpp
+    runtime/ops/quantized_linear/QuantizedLinearBackward.cpp
     runtime/ops/mul/BinaryOp.cpp
     runtime/ops/embedding_q4gsw/EmbeddingQ4gsw.cpp
     runtime/ops/rope/RotaryEmbedding.cpp
@@ -67,6 +68,8 @@ set(WEBGPU_SRCS
     runtime/ops/dim_order/DimOrder.cpp
     runtime/ops/linear/Linear.cpp
     runtime/ops/embedding/Embedding.cpp
+    runtime/ops/quantized_linear/LinearDw.cpp
+    runtime/ops/quantized_linear/QuantizedLinearRequant.cpp
 )
 
 add_library(webgpu_backend ${WEBGPU_SRCS})

backends/webgpu/runtime/ops/quantized_linear/LinearDw.cpp

@@ -0,0 +1,173 @@
+/*
+ * 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.
+ */
+
+#include <executorch/backends/webgpu/runtime/WebGPUGraph.h>
+#include <executorch/backends/webgpu/runtime/WebGPUUtils.h>
+#include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
+#include <executorch/backends/webgpu/runtime/ops/quantized_linear/linear_dW_wgsl.h>
+
+#include <webgpu/webgpu.h>
+
+#include <cstdint>
+#include <cstring>
+#include <stdexcept>
+
+namespace executorch::backends::webgpu {
+
+namespace {
+
+struct LinearDwParams {
+  uint32_t M;
+  uint32_t N;
+  uint32_t K;
+  uint32_t _pad;
+};
+static_assert(sizeof(LinearDwParams) == 16, "params must be 16 bytes");
+
+// STE weight gradient d_W[N,K] = d_out^T @ x.
+void linear_dW_impl(WebGPUGraph& graph, const std::vector<int>& args) {
+  const int dout_id = args.at(0);
+  const int x_id = args.at(1);
+  const int dw_id = args.at(2);
+
+  WGPUDevice device = graph.device();
+  const auto& dout = graph.get_tensor(dout_id);
+  const auto& x = graph.get_tensor(x_id);
+  const auto& dw = graph.get_tensor(dw_id);
+
+  if (dw.dims.size() != 2 || dout.dims.empty() || x.dims.empty()) {
+    throw std::runtime_error("linear_dW: bad tensor ranks");
+  }
+  const uint32_t N = static_cast<uint32_t>(dw.dims[0]);
+  const uint32_t K = static_cast<uint32_t>(dw.dims[1]);
+  if (N == 0 || K == 0) {
+    throw std::runtime_error("linear_dW: N or K == 0");
+  }
+
+  uint64_t dout_numel = 1;
+  for (int64_t d : dout.dims) {
+    dout_numel *= static_cast<uint64_t>(d);
+  }
+  uint64_t x_numel = 1;
+  for (int64_t d : x.dims) {
+    x_numel *= static_cast<uint64_t>(d);
+  }
+  if (static_cast<uint32_t>(dout.dims.back()) != N ||
+      static_cast<uint32_t>(x.dims.back()) != K) {
+    throw std::runtime_error("linear_dW: d_out/x last dim mismatch");
+  }
+  const uint32_t M = static_cast<uint32_t>(dout_numel / N);
+  if (dout_numel % N != 0 || x_numel % K != 0 || x_numel / K != M) {
+    throw std::runtime_error("linear_dW: M mismatch across d_out/x");
+  }
+
+  // fp32-only byte-size guards (mirror the forward's byte checks).
+  if (dw.nbytes != static_cast<uint64_t>(N) * K * sizeof(float) ||
+      dout.nbytes != dout_numel * sizeof(float) ||
+      x.nbytes != x_numel * sizeof(float)) {
+    throw std::runtime_error("linear_dW: fp32-only (byte-size mismatch)");
+  }
+
+  LinearDwParams params = {};
+  params.M = M;
+  params.N = N;
+  params.K = K;
+
+  const uint32_t wg_size =
+      utils::clamp_workgroup_size(device, kLinearDwWorkgroupSizeX);
+  const uint64_t tiles =
+      utils::div_up<uint64_t>(N, 4u) * utils::div_up<uint64_t>(K, 4u);
+  if (tiles > UINT32_MAX) {
+    throw std::runtime_error("linear_dW: tile count exceeds u32");
+  }
+  const uint32_t workgroup_count = utils::compute_1d_workgroup_count(
+      device, static_cast<uint32_t>(tiles), wg_size, "linear_dW");
+
+  WGPUBuffer uniform_buffer =
+      utils::make_uniform(device, &params, sizeof(params));
+  graph.add_uniform_buffer_bytes(sizeof(params));
+
+  WGPUShaderSourceWGSL wgsl_desc = {};
+  wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL;
+  wgsl_desc.code = {kLinearDwWGSL, WGPU_STRLEN};
+  WGPUShaderModuleDescriptor shader_desc = {};
+  shader_desc.nextInChain = &wgsl_desc.chain;
+  WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc);
+
+  WGPUBindGroupLayoutEntry entries[4] = {};
+  entries[0].binding = 0;
+  entries[0].visibility = WGPUShaderStage_Compute;
+  entries[0].buffer.type = WGPUBufferBindingType_Storage;
+  for (uint32_t i = 1; i <= 2; i++) {
+    entries[i].binding = i;
+    entries[i].visibility = WGPUShaderStage_Compute;
+    entries[i].buffer.type = WGPUBufferBindingType_ReadOnlyStorage;
+  }
+  entries[3].binding = 3;
+  entries[3].visibility = WGPUShaderStage_Compute;
+  entries[3].buffer.type = WGPUBufferBindingType_Uniform;
+
+  WGPUBindGroupLayoutDescriptor bgl_desc = {};
+  bgl_desc.entryCount = 4;
+  bgl_desc.entries = entries;
+  WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc);
+
+  WGPUPipelineLayoutDescriptor pl_desc = {};
+  pl_desc.bindGroupLayoutCount = 1;
+  pl_desc.bindGroupLayouts = &bgl;
+  WGPUPipelineLayout pipeline_layout =
+      wgpuDeviceCreatePipelineLayout(device, &pl_desc);
+
+  WGPUConstantEntry wg_size_constant = {};
+  wg_size_constant.key = {"wg_size", WGPU_STRLEN};
+  wg_size_constant.value = static_cast<double>(wg_size);
+
+  WGPUComputePipelineDescriptor pipeline_desc = {};
+  pipeline_desc.layout = pipeline_layout;
+  pipeline_desc.compute.module = shader;
+  pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN};
+  pipeline_desc.compute.constantCount = 1;
+  pipeline_desc.compute.constants = &wg_size_constant;
+  WGPUComputePipeline pipeline =
+      wgpuDeviceCreateComputePipeline(device, &pipeline_desc);
+
+  WGPUBindGroupEntry bg_entries[4] = {};
+  bg_entries[0].binding = 0;
+  bg_entries[0].buffer = dw.buffer;
+  bg_entries[0].size = dw.nbytes;
+  bg_entries[1].binding = 1;
+  bg_entries[1].buffer = dout.buffer;
+  bg_entries[1].size = dout.nbytes;
+  bg_entries[2].binding = 2;
+  bg_entries[2].buffer = x.buffer;
+  bg_entries[2].size = x.nbytes;
+  bg_entries[3].binding = 3;
+  bg_entries[3].buffer = uniform_buffer;
+  bg_entries[3].size = sizeof(params);
+
+  WGPUBindGroupDescriptor bg_desc = {};
+  bg_desc.layout = bgl;
+  bg_desc.entryCount = 4;
+  bg_desc.entries = bg_entries;
+  WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);
+
+  graph.add_dispatch({pipeline, bind_group, workgroup_count, "linear_dW"});
+
+  wgpuShaderModuleRelease(shader);
+  wgpuBindGroupLayoutRelease(bgl);
+  wgpuPipelineLayoutRelease(pipeline_layout);
+  graph.own_uniform_buffer(uniform_buffer);
+}
+
+} // namespace
+
+WEBGPU_REGISTER_OPERATORS {
+  WEBGPU_REGISTER_OP(et_vk.linear_dW.default, linear_dW_impl);
+}
+
+} // namespace executorch::backends::webgpu