Update

[ghstack-poisoned]

Author: JCNTH

backends/vulkan/custom_ops_lib.py

@@ -335,7 +335,13 @@ def linear_q4gsw_backward(ctx, grad_out):
     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)
+    return (
+        d_x,
+        None,
+        None,
+        None,
+        None,
+    )  # grads for (x, weights, scales, group_size, bias)
 
 
 torch.library.register_autograd(
@@ -1197,7 +1203,6 @@ def fused_ce_backward(ctx, grad_loss, grad_dlogits):
 fused_ce_op = getattr(getattr(torch.ops, namespace), name)
 
 
-
 ###########################
 ## adamw_step (training) ##
 ###########################
@@ -1272,7 +1277,6 @@ def linear_q4gsw_dw_meta(
 linear_q4gsw_dw_op = getattr(getattr(torch.ops, namespace), name)
 
 
-
 ##################
 ## q4gsw_requant ##
 ##################

backends/vulkan/op_registry.py

@@ -1769,8 +1769,6 @@ def register_fused_ce():
     )
 
 
-
-
 @update_features(
     [
         exir_ops.edge.aten.ne.Scalar,
@@ -1789,8 +1787,6 @@ def register_compare_scalar_ops():
     )
 
 
-
-
 @update_features(exir_ops.edge.aten.logical_not.default)
 def register_logical_not():
     return OpFeatures(
@@ -1801,7 +1797,6 @@ def register_logical_not():
     )
 
 
-
 @update_features("et_vk::adamw_step")
 def register_adamw_step():
     return OpFeatures(
@@ -1819,8 +1814,6 @@ def register_linear_q4gsw_dw():
     )
 
 
-
-
 @update_features(exir_ops.edge.et_vk.q4gsw_requant.default)
 def register_q4gsw_requant():
     return OpFeatures(

backends/webgpu/runtime/ops/compare/Compare.cpp

@@ -69,7 +69,8 @@ void compare_impl(
   const size_t out_bind_size = (out_tensor.nbytes + 3) & ~size_t(3);
   const uint32_t n_words = (numel + 3u) / 4u;
 
-  uint32_t wg_size = utils::clamp_workgroup_size(device, kCompareWorkgroupSizeX);
+  uint32_t wg_size =
+      utils::clamp_workgroup_size(device, kCompareWorkgroupSizeX);
   uint32_t workgroup_count =
       utils::compute_1d_workgroup_count(device, n_words, wg_size, op_name);
 
@@ -139,20 +140,21 @@ void compare_impl(
       graph.add_dispatch({pipeline, bind_group, workgroup_count});
 
   WGPUBuffer p_buf = params_buf;
-  auto cmp_resize = [self_id, out_id, mode, scalar, wg_size, dispatch_idx,
-                     p_buf, op_name](WebGPUGraph& g) {
-    const auto& d = g.cur_dims(self_id);
-    uint32_t n = 1u;
-    for (auto x : d) {
-      n *= static_cast<uint32_t>(x);
-    }
-    g.set_cur_dims(out_id, d);
-    CompareParams p = {n, mode, scalar, 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);
-  };
+  auto cmp_resize =
+      [self_id, out_id, mode, scalar, wg_size, dispatch_idx, p_buf, op_name](
+          WebGPUGraph& g) {
+        const auto& d = g.cur_dims(self_id);
+        uint32_t n = 1u;
+        for (auto x : d) {
+          n *= static_cast<uint32_t>(x);
+        }
+        g.set_cur_dims(out_id, d);
+        CompareParams p = {n, mode, scalar, 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);
+      };
   graph.add_tensor_resize_hook(self_id, cmp_resize);
 
   wgpuShaderModuleRelease(shader);

backends/webgpu/runtime/ops/embedding/Embedding.cpp

@@ -54,7 +54,8 @@ void embedding_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   if (out.nbytes != out_numel * sizeof(float) ||
       weight.nbytes % sizeof(float) != 0 ||
       indices.nbytes != index_numel * sizeof(int32_t)) {
-    throw std::runtime_error("embedding: fp32 weight/out + i32 indices required");
+    throw std::runtime_error(
+        "embedding: fp32 weight/out + i32 indices required");
   }
   if (out_numel != index_numel * dim) {
     throw std::runtime_error("embedding: out numel != num_indices * dim");

backends/webgpu/runtime/ops/fill/Fill.cpp

@@ -144,20 +144,26 @@ void add_fill(
 
 void full_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   add_fill(
-      graph, args.at(args.size() - 1), read_scalar(graph, args.at(1), "full"),
+      graph,
+      args.at(args.size() - 1),
+      read_scalar(graph, args.at(1), "full"),
       "full");
 }
 
 void full_like_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   add_fill(
-      graph, args.at(args.size() - 1),
-      read_scalar(graph, args.at(1), "full_like"), "full_like");
+      graph,
+      args.at(args.size() - 1),
+      read_scalar(graph, args.at(1), "full_like"),
+      "full_like");
 }
 
 void scalar_tensor_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   add_fill(
-      graph, args.at(args.size() - 1),
-      read_scalar(graph, args.at(0), "scalar_tensor"), "scalar_tensor");
+      graph,
+      args.at(args.size() - 1),
+      read_scalar(graph, args.at(0), "scalar_tensor"),
+      "scalar_tensor");
 }
 
 } // namespace

backends/webgpu/runtime/ops/log_softmax/LogSoftmax.cpp

@@ -29,7 +29,9 @@ struct LogSoftmaxParams {
   uint32_t inner;
   uint32_t pad_;
 };
-static_assert(sizeof(LogSoftmaxParams) == 16, "LogSoftmaxParams must be 16 bytes");
+static_assert(
+    sizeof(LogSoftmaxParams) == 16,
+    "LogSoftmaxParams must be 16 bytes");
 
 // Decompose dims into [outer, R, inner] for a reduction along `dim`.
 void decompose(
@@ -85,7 +87,8 @@ void log_softmax_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   for (int64_t d : in.dims) {
     numel *= static_cast<uint64_t>(d);
   }
-  if (in.nbytes != numel * sizeof(float) || out.nbytes != numel * sizeof(float)) {
+  if (in.nbytes != numel * sizeof(float) ||
+      out.nbytes != numel * sizeof(float)) {
     throw std::runtime_error(
         "WebGPU log_softmax: fp32-only (byte-size mismatch)");
   }
@@ -205,7 +208,9 @@ void log_softmax_impl(WebGPUGraph& graph, const std::vector<int>& args) {
         wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p));
         g.dispatch_at(dispatch_idx).workgroup_count_x =
             utils::compute_1d_workgroup_count(
-                g.device(), static_cast<uint32_t>(live_lines), wg_size,
+                g.device(),
+                static_cast<uint32_t>(live_lines),
+                wg_size,
                 "log_softmax(resize)");
         g.set_cur_dims(out_id, std::vector<int64_t>(d.begin(), d.end()));
       });

backends/webgpu/runtime/ops/logical_not/LogicalNot.cpp

@@ -120,20 +120,21 @@ void logical_not_impl(WebGPUGraph& graph, const std::vector<int>& args) {
       graph.add_dispatch({pipeline, bind_group, workgroup_count});
 
   WGPUBuffer p_buf = params_buf;
-  auto resize = [self_id, out_id, wg_size, dispatch_idx, p_buf](
-                    WebGPUGraph& g) {
-    const auto& d = g.cur_dims(self_id);
-    uint32_t n = 1u;
-    for (auto x : d) {
-      n *= static_cast<uint32_t>(x);
-    }
-    g.set_cur_dims(out_id, d);
-    LogicalNotParams p = {n, 0u, 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, "logical_not");
-  };
+  auto resize =
+      [self_id, out_id, wg_size, dispatch_idx, p_buf](WebGPUGraph& g) {
+        const auto& d = g.cur_dims(self_id);
+        uint32_t n = 1u;
+        for (auto x : d) {
+          n *= static_cast<uint32_t>(x);
+        }
+        g.set_cur_dims(out_id, d);
+        LogicalNotParams p = {n, 0u, 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, "logical_not");
+      };
   graph.add_tensor_resize_hook(self_id, resize);
 
   wgpuShaderModuleRelease(shader);

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

@@ -81,7 +81,8 @@ void q4gsw_backward_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   if (dout.nbytes != static_cast<uint64_t>(M) * N * sizeof(float)) {
     throw std::runtime_error("q4gsw_backward: d_out fp32/shape mismatch");
   }
-  if (scales.nbytes != static_cast<uint64_t>(num_groups) * padded_N * sizeof(float)) {
+  if (scales.nbytes !=
+      static_cast<uint64_t>(num_groups) * padded_N * sizeof(float)) {
     throw std::runtime_error("q4gsw_backward: scales fp32/shape mismatch");
   }
   if (weight.nbytes != static_cast<uint64_t>(N) * K_packed) {

backends/webgpu/runtime/ops/reduce/Reduce.cpp

@@ -227,7 +227,9 @@ void reduce_impl(
         wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p));
         g.dispatch_at(dispatch_idx).workgroup_count_x =
             utils::compute_1d_workgroup_count(
-                g.device(), static_cast<uint32_t>(live_outputs), wg_size,
+                g.device(),
+                static_cast<uint32_t>(live_outputs),
+                wg_size,
                 "reduce(resize)");
         // Propagate reduced output dims for downstream resize hooks.
         int64_t nd = static_cast<int64_t>(d.size());

backends/webgpu/runtime/ops/softmax/Softmax.cpp

@@ -173,8 +173,7 @@ void softmax_impl(WebGPUGraph& graph, const std::vector<int>& args) {
 
   // Dynamic shapes: recompute the decomposition for the reduced dim + dispatch.
   WGPUBuffer params_buf = uniform_buffer;
-  const uint64_t build_numel =
-      static_cast<uint64_t>(outer) * r * inner;
+  const uint64_t build_numel = static_cast<uint64_t>(outer) * r * inner;
   graph.add_tensor_resize_hook(
       in_id,
       [in_id, out_id, dim, build_numel, wg_size, dispatch_idx, params_buf](
@@ -198,7 +197,9 @@ void softmax_impl(WebGPUGraph& graph, const std::vector<int>& args) {
         wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p));
         g.dispatch_at(dispatch_idx).workgroup_count_x =
             utils::compute_1d_workgroup_count(
-                g.device(), static_cast<uint32_t>(live_lines), wg_size,
+                g.device(),
+                static_cast<uint32_t>(live_lines),
+                wg_size,
                 "softmax(resize)");
         g.set_cur_dims(out_id, d);
       });