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

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_q4gsw_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_q4gsw_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_q4gsw_dw"
+lib.define(f"{name}(Tensor d_out, Tensor x) -> Tensor")
+lib.impl(name, linear_q4gsw_dw_impl, "CompositeExplicitAutograd")
+lib.impl(name, linear_q4gsw_dw_meta, "Meta")
+linear_q4gsw_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(
@@ -1774,6 +1783,23 @@ def register_logical_not():
     )
 
 
+@update_features(exir_ops.edge.et_vk.linear_q4gsw_dw.default)
+def register_linear_q4gsw_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
@@ -68,6 +69,8 @@ set(WEBGPU_SRCS
     runtime/ops/linear/Linear.cpp
     runtime/ops/embedding/Embedding.cpp
     runtime/ops/logical_not/LogicalNot.cpp
+    runtime/ops/quantized_linear/QuantizedLinearDw.cpp
+    runtime/ops/quantized_linear/QuantizedLinearRequant.cpp
 )
 
 add_library(webgpu_backend ${WEBGPU_SRCS})