diff --git a/backends/vulkan/custom_ops_lib.py b/backends/vulkan/custom_ops_lib.py index 4364f67123d..313daf65439 100644 --- a/backends/vulkan/custom_ops_lib.py +++ b/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) diff --git a/backends/vulkan/op_registry.py b/backends/vulkan/op_registry.py index e8f22c8661d..5e280f82152 100644 --- a/backends/vulkan/op_registry.py +++ b/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_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 ## ####################### diff --git a/backends/webgpu/CMakeLists.txt b/backends/webgpu/CMakeLists.txt index 2fcfab98f49..bd63e254242 100644 --- a/backends/webgpu/CMakeLists.txt +++ b/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/QuantizedLinearDw.cpp + runtime/ops/quantized_linear/QuantizedLinearRequant.cpp ) add_library(webgpu_backend ${WEBGPU_SRCS}) diff --git a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinearBackward.cpp b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinearBackward.cpp new file mode 100644 index 00000000000..9f0ac36dfbd --- /dev/null +++ b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinearBackward.cpp @@ -0,0 +1,208 @@ +/* + * 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 +#include +#include +#include + +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +struct Q4gswBackwardParams { + uint32_t M; + uint32_t N; + uint32_t K; + uint32_t K_packed; + uint32_t group_size; + uint32_t padded_N; + uint32_t has_bias; + uint32_t _pad; +}; +static_assert(sizeof(Q4gswBackwardParams) == 32, "params must be 32 bytes"); + +// linear_q4gsw_backward: d_x[M,K] = d_out[M,N] @ dequant(W)[N,K]. +void q4gsw_backward_impl(WebGPUGraph& graph, const std::vector& args) { + const int dout_id = args.at(0); + const int weight_id = args.at(1); + const int scales_id = args.at(2); + const int group_size_id = args.at(3); + const int dx_id = args.at(4); + + WGPUDevice device = graph.device(); + const auto& dout = graph.get_tensor(dout_id); + const auto& weight = graph.get_tensor(weight_id); + const auto& scales = graph.get_tensor(scales_id); + const auto& dx = graph.get_tensor(dx_id); + + if (weight.dims.size() != 2 || scales.dims.size() != 2 || dx.dims.empty() || + dout.dims.empty()) { + throw std::runtime_error("q4gsw_backward: bad tensor ranks"); + } + const uint32_t N = static_cast(weight.dims[0]); + const uint32_t K_packed = static_cast(weight.dims[1]); + const uint32_t K = static_cast(dx.dims.back()); + if (N == 0 || K == 0) { + throw std::runtime_error("q4gsw_backward: N or K == 0"); + } + uint64_t dx_numel = 1; + for (int64_t d : dx.dims) { + dx_numel *= static_cast(d); + } + const uint32_t M = static_cast(dx_numel / K); + const uint32_t num_groups = static_cast(scales.dims[0]); + const uint32_t padded_N = static_cast(scales.dims[1]); + + if (graph.get_value_type(group_size_id) != WebGPUGraph::ValueType::Int) { + throw std::runtime_error("q4gsw_backward: group_size must be Int"); + } + const int64_t group_size = graph.get_int(group_size_id); + if (group_size <= 0) { + throw std::runtime_error("q4gsw_backward: group_size must be positive"); + } + const uint32_t gs = static_cast(group_size); + + // fp32 + shape guards (mirror the forward's byte checks). + if (dx.nbytes != dx_numel * sizeof(float)) { + throw std::runtime_error("q4gsw_backward: d_x fp32-only"); + } + if (dout.nbytes != static_cast(M) * N * sizeof(float)) { + throw std::runtime_error("q4gsw_backward: d_out fp32/shape mismatch"); + } + if (scales.nbytes != + static_cast(num_groups) * padded_N * sizeof(float)) { + throw std::runtime_error("q4gsw_backward: scales fp32/shape mismatch"); + } + if (weight.nbytes != static_cast(N) * K_packed) { + throw std::runtime_error("q4gsw_backward: weight byte-size mismatch"); + } + if (K_packed != (K + 1u) / 2u) { + throw std::runtime_error("q4gsw_backward: K_packed != ceil(K/2)"); + } + if (num_groups < (K + gs - 1u) / gs || padded_N < N) { + throw std::runtime_error("q4gsw_backward: scales too small"); + } + + Q4gswBackwardParams params = {}; + params.M = M; + params.N = N; + params.K = K; + params.K_packed = K_packed; + params.group_size = gs; + params.padded_N = padded_N; + + const uint32_t wg_size = + utils::clamp_workgroup_size(device, kQ4gswBackwardWorkgroupSizeX); + const uint64_t tiles = static_cast((M + 3u) / 4u) * ((K + 3u) / 4u); + if (tiles > UINT32_MAX) { + throw std::runtime_error("q4gsw_backward: tile count exceeds u32"); + } + const uint32_t workgroup_count = utils::compute_1d_workgroup_count( + device, static_cast(tiles), wg_size, "q4gsw_backward"); + + WGPUBufferDescriptor uniform_desc = {}; + uniform_desc.size = sizeof(Q4gswBackwardParams); + uniform_desc.usage = WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst; + uniform_desc.mappedAtCreation = true; + WGPUBuffer uniform_buffer = wgpuDeviceCreateBuffer(device, &uniform_desc); + std::memcpy( + wgpuBufferGetMappedRange(uniform_buffer, 0, sizeof(Q4gswBackwardParams)), + ¶ms, + sizeof(Q4gswBackwardParams)); + wgpuBufferUnmap(uniform_buffer); + graph.add_uniform_buffer_bytes(sizeof(Q4gswBackwardParams)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kQ4gswBackwardWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + WGPUBindGroupLayoutEntry entries[5] = {}; + entries[0].binding = 0; + entries[0].visibility = WGPUShaderStage_Compute; + entries[0].buffer.type = WGPUBufferBindingType_Storage; + for (uint32_t i = 1; i <= 3; i++) { + entries[i].binding = i; + entries[i].visibility = WGPUShaderStage_Compute; + entries[i].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + } + entries[4].binding = 4; + entries[4].visibility = WGPUShaderStage_Compute; + entries[4].buffer.type = WGPUBufferBindingType_Uniform; + + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 5; + 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(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[5] = {}; + bg_entries[0].binding = 0; + bg_entries[0].buffer = dx.buffer; + bg_entries[0].size = dx.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 = weight.buffer; + bg_entries[2].size = weight.nbytes; + bg_entries[3].binding = 3; + bg_entries[3].buffer = scales.buffer; + bg_entries[3].size = scales.nbytes; + bg_entries[4].binding = 4; + bg_entries[4].buffer = uniform_buffer; + bg_entries[4].size = sizeof(Q4gswBackwardParams); + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 5; + bg_desc.entries = bg_entries; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + graph.add_dispatch({pipeline, bind_group, workgroup_count, "q4gsw_backward"}); + + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + wgpuBufferRelease(uniform_buffer); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(et_vk.linear_q4gsw_backward.default, q4gsw_backward_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinearDw.cpp b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinearDw.cpp new file mode 100644 index 00000000000..628a0f140c8 --- /dev/null +++ b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinearDw.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 +#include +#include +#include + +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +struct Q4gswDwParams { + uint32_t M; + uint32_t N; + uint32_t K; + uint32_t _pad; +}; +static_assert(sizeof(Q4gswDwParams) == 16, "params must be 16 bytes"); + +// STE weight gradient d_W[N,K] = d_out^T @ x. +void q4gsw_dw_impl(WebGPUGraph& graph, const std::vector& 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("q4gsw_dw: bad tensor ranks"); + } + const uint32_t N = static_cast(dw.dims[0]); + const uint32_t K = static_cast(dw.dims[1]); + if (N == 0 || K == 0) { + throw std::runtime_error("q4gsw_dw: N or K == 0"); + } + + uint64_t dout_numel = 1; + for (int64_t d : dout.dims) { + dout_numel *= static_cast(d); + } + uint64_t x_numel = 1; + for (int64_t d : x.dims) { + x_numel *= static_cast(d); + } + if (static_cast(dout.dims.back()) != N || + static_cast(x.dims.back()) != K) { + throw std::runtime_error("q4gsw_dw: d_out/x last dim mismatch"); + } + const uint32_t M = static_cast(dout_numel / N); + if (dout_numel % N != 0 || x_numel % K != 0 || x_numel / K != M) { + throw std::runtime_error("q4gsw_dw: M mismatch across d_out/x"); + } + + // fp32-only byte-size guards (mirror the forward's byte checks). + if (dw.nbytes != static_cast(N) * K * sizeof(float) || + dout.nbytes != dout_numel * sizeof(float) || + x.nbytes != x_numel * sizeof(float)) { + throw std::runtime_error("q4gsw_dw: fp32-only (byte-size mismatch)"); + } + + Q4gswDwParams params = {}; + params.M = M; + params.N = N; + params.K = K; + + const uint32_t wg_size = + utils::clamp_workgroup_size(device, kQ4gswDwWorkgroupSizeX); + const uint64_t tiles = + utils::div_up(N, 4u) * utils::div_up(K, 4u); + if (tiles > UINT32_MAX) { + throw std::runtime_error("q4gsw_dw: tile count exceeds u32"); + } + const uint32_t workgroup_count = utils::compute_1d_workgroup_count( + device, static_cast(tiles), wg_size, "q4gsw_dw"); + + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(params)); + graph.add_uniform_buffer_bytes(sizeof(params)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kQ4gswDwWGSL, 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(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, "q4gsw_dw"}); + + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + graph.own_uniform_buffer(uniform_buffer); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(et_vk.linear_q4gsw_dw.default, q4gsw_dw_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinearRequant.cpp b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinearRequant.cpp new file mode 100644 index 00000000000..3c75d63165a --- /dev/null +++ b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinearRequant.cpp @@ -0,0 +1,192 @@ +/* + * 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 +#include +#include +#include + +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +struct Q4gswRequantParams { + uint32_t N; + uint32_t K; + uint32_t K_packed; + uint32_t group_size; + uint32_t padded_N; + uint32_t num_words; + uint32_t _pad0; + uint32_t _pad1; +}; +static_assert(sizeof(Q4gswRequantParams) == 32, "params must be 32 bytes"); + +// STE re-quant + int4 pack at a frozen per-group scale (only the codes move). +void q4gsw_requant_impl(WebGPUGraph& graph, const std::vector& args) { + const int latent_id = args.at(0); + const int scales_id = args.at(1); + const int group_size_id = args.at(2); + const int packed_id = args.at(3); + + WGPUDevice device = graph.device(); + const auto& latent = graph.get_tensor(latent_id); + const auto& scales = graph.get_tensor(scales_id); + const auto& packed = graph.get_tensor(packed_id); + + if (latent.dims.size() != 2 || scales.dims.size() != 2 || + packed.dims.size() != 2) { + throw std::runtime_error("q4gsw_requant: bad tensor ranks"); + } + const uint32_t N = static_cast(latent.dims[0]); + const uint32_t K = static_cast(latent.dims[1]); + const uint32_t K_packed = static_cast(packed.dims[1]); + const uint32_t num_groups = static_cast(scales.dims[0]); + const uint32_t padded_N = static_cast(scales.dims[1]); + if (N == 0 || K == 0) { + throw std::runtime_error("q4gsw_requant: N or K == 0"); + } + if (static_cast(packed.dims[0]) != N) { + throw std::runtime_error("q4gsw_requant: packed rows != N"); + } + if (K_packed != (K + 1u) / 2u) { + throw std::runtime_error("q4gsw_requant: K_packed != ceil(K/2)"); + } + if ((static_cast(N) * K_packed) % 4u != 0u) { + throw std::runtime_error( + "q4gsw_requant: N*K_packed must be a multiple of 4 (u32-packed)"); + } + + if (graph.get_value_type(group_size_id) != WebGPUGraph::ValueType::Int) { + throw std::runtime_error("q4gsw_requant: group_size must be Int"); + } + const int64_t group_size = graph.get_int(group_size_id); + if (group_size <= 0) { + throw std::runtime_error("q4gsw_requant: group_size must be positive"); + } + const uint32_t gs = static_cast(group_size); + if (num_groups < (K + gs - 1u) / gs || padded_N < N) { + throw std::runtime_error("q4gsw_requant: scales dims too small for K/N"); + } + + if (latent.nbytes != static_cast(N) * K * sizeof(float)) { + throw std::runtime_error("q4gsw_requant: latent fp32-only"); + } + if (scales.nbytes != + static_cast(num_groups) * padded_N * sizeof(float)) { + throw std::runtime_error("q4gsw_requant: scales fp32/shape mismatch"); + } + if (packed.nbytes != static_cast(N) * K_packed) { + throw std::runtime_error("q4gsw_requant: packed byte-size mismatch"); + } + + const uint32_t num_words = + static_cast((static_cast(N) * K_packed) / 4u); + + Q4gswRequantParams params = {}; + params.N = N; + params.K = K; + params.K_packed = K_packed; + params.group_size = gs; + params.padded_N = padded_N; + params.num_words = num_words; + + const uint32_t wg_size = + utils::clamp_workgroup_size(device, kQ4gswRequantWorkgroupSizeX); + const uint32_t workgroup_count = utils::compute_1d_workgroup_count( + device, num_words, wg_size, "q4gsw_requant"); + + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(params)); + graph.add_uniform_buffer_bytes(sizeof(params)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kQ4gswRequantWGSL, 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(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 = packed.buffer; + bg_entries[0].size = packed.nbytes; + bg_entries[1].binding = 1; + bg_entries[1].buffer = latent.buffer; + bg_entries[1].size = latent.nbytes; + bg_entries[2].binding = 2; + bg_entries[2].buffer = scales.buffer; + bg_entries[2].size = scales.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, "q4gsw_requant"}); + + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + graph.own_uniform_buffer(uniform_buffer); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(et_vk.q4gsw_requant.default, q4gsw_requant_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_backward.wgsl b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_backward.wgsl new file mode 100644 index 00000000000..569d9450b5b --- /dev/null +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_backward.wgsl @@ -0,0 +1,84 @@ +// Mirrors q4gsw_linear.wgsl dequant; contracts over N (forward over K). + +@group(0) @binding(0) var t_dx: array; // [M, K] +@group(0) @binding(1) var t_dout: array; // [M, N] +@group(0) @binding(2) var t_weight: array; +@group(0) @binding(3) var t_scales: array; + +struct Params { + M: u32, + N: u32, + K: u32, + K_packed: u32, + group_size: u32, + padded_N: u32, + has_bias: u32, + _pad: u32, +} +@group(0) @binding(4) var params: Params; + +override wg_size: u32 = 64u; + +const TM: u32 = 4u; +const TK: u32 = 4u; +const TILE_ELEMS: u32 = TM * TK; + +@compute @workgroup_size(wg_size, 1, 1) +fn main(@builtin(global_invocation_id) gid: vec3) { + let nrt = (params.M + TM - 1u) / TM; + let nkt = (params.K + TK - 1u) / TK; + let tiles = nrt * nkt; + if (gid.x >= tiles) { + return; + } + let row_tile = gid.x / nkt; + let col_tile = gid.x % nkt; + let m0 = row_tile * TM; + let k0 = col_tile * TK; + + var acc: array; + for (var i: u32 = 0u; i < TILE_ELEMS; i = i + 1u) { + acc[i] = 0.0; + } + + var n: u32 = 0u; + loop { + if (n >= params.N) { + break; + } + // Load TM d_out for column n once; reused across TK k columns. + var dout_reg: array; + for (var ml: u32 = 0u; ml < TM; ml = ml + 1u) { + let m_eff = min(m0 + ml, params.M - 1u); + dout_reg[ml] = t_dout[m_eff * params.N + n]; + } + for (var kl: u32 = 0u; kl < TK; kl = kl + 1u) { + let k_eff = min(k0 + kl, params.K - 1u); + let byte_idx = n * params.K_packed + (k_eff >> 1u); + let word = t_weight[byte_idx >> 2u]; + let b = (word >> ((byte_idx & 3u) * 8u)) & 0xFFu; + var nib: u32; + if ((k_eff & 1u) == 0u) { + nib = b & 0x0Fu; + } else { + nib = (b >> 4u) & 0x0Fu; + } + let q = f32(i32(nib) - 8); + let dq = q * t_scales[(k_eff / params.group_size) * params.padded_N + n]; + for (var ml: u32 = 0u; ml < TM; ml = ml + 1u) { + acc[ml * TK + kl] = acc[ml * TK + kl] + dout_reg[ml] * dq; + } + } + n = n + 1u; + } + + for (var ml: u32 = 0u; ml < TM; ml = ml + 1u) { + let m = m0 + ml; + for (var kl: u32 = 0u; kl < TK; kl = kl + 1u) { + let k = k0 + kl; + if (m < params.M && k < params.K) { + t_dx[m * params.K + k] = acc[ml * TK + kl]; + } + } + } +} diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_backward_wgsl.h b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_backward_wgsl.h new file mode 100644 index 00000000000..8882fff0b09 --- /dev/null +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_backward_wgsl.h @@ -0,0 +1,108 @@ +/* + * 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 + +namespace executorch::backends::webgpu { + +// @generated from q4gsw_backward.wgsl - DO NOT EDIT. +// wgsl-sha256: eb7da436c89fc3cd8c062af5450d11f28fd4fc5905cc38d043fc3584252438b2 +inline constexpr const char* kQ4gswBackwardWGSL = R"( +// Mirrors q4gsw_linear.wgsl dequant; contracts over N (forward over K). + +@group(0) @binding(0) var t_dx: array; // [M, K] +@group(0) @binding(1) var t_dout: array; // [M, N] +@group(0) @binding(2) var t_weight: array; +@group(0) @binding(3) var t_scales: array; + +struct Params { + M: u32, + N: u32, + K: u32, + K_packed: u32, + group_size: u32, + padded_N: u32, + has_bias: u32, + _pad: u32, +} +@group(0) @binding(4) var params: Params; + +override wg_size: u32 = 64u; + +const TM: u32 = 4u; +const TK: u32 = 4u; +const TILE_ELEMS: u32 = TM * TK; + +@compute @workgroup_size(wg_size, 1, 1) +fn main(@builtin(global_invocation_id) gid: vec3) { + let nrt = (params.M + TM - 1u) / TM; + let nkt = (params.K + TK - 1u) / TK; + let tiles = nrt * nkt; + if (gid.x >= tiles) { + return; + } + let row_tile = gid.x / nkt; + let col_tile = gid.x % nkt; + let m0 = row_tile * TM; + let k0 = col_tile * TK; + + var acc: array; + for (var i: u32 = 0u; i < TILE_ELEMS; i = i + 1u) { + acc[i] = 0.0; + } + + var n: u32 = 0u; + loop { + if (n >= params.N) { + break; + } + // Load TM d_out for column n once; reused across TK k columns. + var dout_reg: array; + for (var ml: u32 = 0u; ml < TM; ml = ml + 1u) { + let m_eff = min(m0 + ml, params.M - 1u); + dout_reg[ml] = t_dout[m_eff * params.N + n]; + } + for (var kl: u32 = 0u; kl < TK; kl = kl + 1u) { + let k_eff = min(k0 + kl, params.K - 1u); + let byte_idx = n * params.K_packed + (k_eff >> 1u); + let word = t_weight[byte_idx >> 2u]; + let b = (word >> ((byte_idx & 3u) * 8u)) & 0xFFu; + var nib: u32; + if ((k_eff & 1u) == 0u) { + nib = b & 0x0Fu; + } else { + nib = (b >> 4u) & 0x0Fu; + } + let q = f32(i32(nib) - 8); + let dq = q * t_scales[(k_eff / params.group_size) * params.padded_N + n]; + for (var ml: u32 = 0u; ml < TM; ml = ml + 1u) { + acc[ml * TK + kl] = acc[ml * TK + kl] + dout_reg[ml] * dq; + } + } + n = n + 1u; + } + + for (var ml: u32 = 0u; ml < TM; ml = ml + 1u) { + let m = m0 + ml; + for (var kl: u32 = 0u; kl < TK; kl = kl + 1u) { + let k = k0 + kl; + if (m < params.M && k < params.K) { + t_dx[m * params.K + k] = acc[ml * TK + kl]; + } + } + } +} +)"; + +inline constexpr uint32_t kQ4gswBackwardWorkgroupSizeX = 64; +inline constexpr uint32_t kQ4gswBackwardWorkgroupSizeY = 1; +inline constexpr uint32_t kQ4gswBackwardWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_dw.wgsl b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_dw.wgsl new file mode 100644 index 00000000000..d7a68cd5d40 --- /dev/null +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_dw.wgsl @@ -0,0 +1,69 @@ +// STE weight-gradient d_W[N,K] = sum_m d_out[m,N]*x[m,K] (operands f32). + +@group(0) @binding(0) var t_dw: array; // [N, K] +@group(0) @binding(1) var t_dout: array; // [M, N] +@group(0) @binding(2) var t_x: array; // [M, K] + +struct Params { + M: u32, + N: u32, + K: u32, + _pad: u32, +} +@group(0) @binding(3) var params: Params; + +override wg_size: u32 = 64u; + +const TN: u32 = 4u; +const TK: u32 = 4u; +const TILE_ELEMS: u32 = TN * TK; + +@compute @workgroup_size(wg_size, 1, 1) +fn main(@builtin(global_invocation_id) gid: vec3) { + let nnt = (params.N + TN - 1u) / TN; + let nkt = (params.K + TK - 1u) / TK; + let tiles = nnt * nkt; + if (gid.x >= tiles) { + return; + } + let row_tile = gid.x / nkt; + let col_tile = gid.x % nkt; + let n0 = row_tile * TN; + let k0 = col_tile * TK; + + var acc: array; + for (var i: u32 = 0u; i < TILE_ELEMS; i = i + 1u) { + acc[i] = 0.0; + } + + var m: u32 = 0u; + loop { + if (m >= params.M) { + break; + } + // Load the TN d_out values for row m once; reused across all TK k columns. + var dout_reg: array; + for (var nl: u32 = 0u; nl < TN; nl = nl + 1u) { + let n_eff = min(n0 + nl, params.N - 1u); + dout_reg[nl] = t_dout[m * params.N + n_eff]; + } + for (var kl: u32 = 0u; kl < TK; kl = kl + 1u) { + let k_eff = min(k0 + kl, params.K - 1u); + let xv = t_x[m * params.K + k_eff]; + for (var nl: u32 = 0u; nl < TN; nl = nl + 1u) { + acc[nl * TK + kl] = acc[nl * TK + kl] + dout_reg[nl] * xv; + } + } + m = m + 1u; + } + + for (var nl: u32 = 0u; nl < TN; nl = nl + 1u) { + let n = n0 + nl; + for (var kl: u32 = 0u; kl < TK; kl = kl + 1u) { + let k = k0 + kl; + if (n < params.N && k < params.K) { + t_dw[n * params.K + k] = acc[nl * TK + kl]; + } + } + } +} diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_dw_wgsl.h b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_dw_wgsl.h new file mode 100644 index 00000000000..77c3eda1b8e --- /dev/null +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_dw_wgsl.h @@ -0,0 +1,93 @@ +/* + * 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 + +namespace executorch::backends::webgpu { + +// @generated from q4gsw_dw.wgsl - DO NOT EDIT. +// wgsl-sha256: be50075764ccd87e0d0c86124f13fd753aa618f994c9a5e1b8b631516aa51f84 +inline constexpr const char* kQ4gswDwWGSL = R"( +// STE weight-gradient d_W[N,K] = sum_m d_out[m,N]*x[m,K] (operands f32). + +@group(0) @binding(0) var t_dw: array; // [N, K] +@group(0) @binding(1) var t_dout: array; // [M, N] +@group(0) @binding(2) var t_x: array; // [M, K] + +struct Params { + M: u32, + N: u32, + K: u32, + _pad: u32, +} +@group(0) @binding(3) var params: Params; + +override wg_size: u32 = 64u; + +const TN: u32 = 4u; +const TK: u32 = 4u; +const TILE_ELEMS: u32 = TN * TK; + +@compute @workgroup_size(wg_size, 1, 1) +fn main(@builtin(global_invocation_id) gid: vec3) { + let nnt = (params.N + TN - 1u) / TN; + let nkt = (params.K + TK - 1u) / TK; + let tiles = nnt * nkt; + if (gid.x >= tiles) { + return; + } + let row_tile = gid.x / nkt; + let col_tile = gid.x % nkt; + let n0 = row_tile * TN; + let k0 = col_tile * TK; + + var acc: array; + for (var i: u32 = 0u; i < TILE_ELEMS; i = i + 1u) { + acc[i] = 0.0; + } + + var m: u32 = 0u; + loop { + if (m >= params.M) { + break; + } + // Load the TN d_out values for row m once; reused across all TK k columns. + var dout_reg: array; + for (var nl: u32 = 0u; nl < TN; nl = nl + 1u) { + let n_eff = min(n0 + nl, params.N - 1u); + dout_reg[nl] = t_dout[m * params.N + n_eff]; + } + for (var kl: u32 = 0u; kl < TK; kl = kl + 1u) { + let k_eff = min(k0 + kl, params.K - 1u); + let xv = t_x[m * params.K + k_eff]; + for (var nl: u32 = 0u; nl < TN; nl = nl + 1u) { + acc[nl * TK + kl] = acc[nl * TK + kl] + dout_reg[nl] * xv; + } + } + m = m + 1u; + } + + for (var nl: u32 = 0u; nl < TN; nl = nl + 1u) { + let n = n0 + nl; + for (var kl: u32 = 0u; kl < TK; kl = kl + 1u) { + let k = k0 + kl; + if (n < params.N && k < params.K) { + t_dw[n * params.K + k] = acc[nl * TK + kl]; + } + } + } +} +)"; + +inline constexpr uint32_t kQ4gswDwWorkgroupSizeX = 64; +inline constexpr uint32_t kQ4gswDwWorkgroupSizeY = 1; +inline constexpr uint32_t kQ4gswDwWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_requant.wgsl b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_requant.wgsl new file mode 100644 index 00000000000..1f1af914630 --- /dev/null +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_requant.wgsl @@ -0,0 +1,54 @@ +// STE re-quant + int4 pack: round(latent/scale) in [-8,7], 2 nibbles/byte. + +@group(0) @binding(0) var t_packed: array; +@group(0) @binding(1) var t_latent: array; // [N, K] +@group(0) @binding(2) var t_scales: array; + +struct Params { + N: u32, + K: u32, + K_packed: u32, + group_size: u32, + padded_N: u32, + num_words: u32, + _pad0: u32, + _pad1: u32, +} +@group(0) @binding(3) var params: Params; + +override wg_size: u32 = 64u; + +fn quant_nibble(n: u32, k: u32) -> u32 { + let s = t_scales[(k / params.group_size) * params.padded_N + n]; + var q: i32 = 0; + if (s != 0.0) { + q = i32(round(t_latent[n * params.K + k] / s)); + } + q = clamp(q, -8, 7); + return u32(q + 8) & 0xFu; +} + +@compute @workgroup_size(wg_size, 1, 1) +fn main(@builtin(global_invocation_id) gid: vec3) { + let word_idx = gid.x; + if (word_idx >= params.num_words) { + return; + } + var word: u32 = 0u; + for (var bi: u32 = 0u; bi < 4u; bi = bi + 1u) { + let byte_idx = word_idx * 4u + bi; + let n = byte_idx / params.K_packed; + let byte_in_row = byte_idx % params.K_packed; + let k_lo = byte_in_row * 2u; + var b: u32 = 0u; + if (n < params.N && k_lo < params.K) { + b = quant_nibble(n, k_lo); + let k_hi = k_lo + 1u; + if (k_hi < params.K) { + b = b | (quant_nibble(n, k_hi) << 4u); + } + } + word = word | (b << (bi * 8u)); + } + t_packed[word_idx] = word; +} diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_requant_wgsl.h b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_requant_wgsl.h new file mode 100644 index 00000000000..84f30d7e1b8 --- /dev/null +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_requant_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 + +namespace executorch::backends::webgpu { + +// @generated from q4gsw_requant.wgsl - DO NOT EDIT. +// wgsl-sha256: 66cbf8061af4580be1773acd2e9d11d539d1a9e4a3fc8d08bd1d4def42e4d643 +inline constexpr const char* kQ4gswRequantWGSL = R"( +// STE re-quant + int4 pack: round(latent/scale) in [-8,7], 2 nibbles/byte. + +@group(0) @binding(0) var t_packed: array; +@group(0) @binding(1) var t_latent: array; // [N, K] +@group(0) @binding(2) var t_scales: array; + +struct Params { + N: u32, + K: u32, + K_packed: u32, + group_size: u32, + padded_N: u32, + num_words: u32, + _pad0: u32, + _pad1: u32, +} +@group(0) @binding(3) var params: Params; + +override wg_size: u32 = 64u; + +fn quant_nibble(n: u32, k: u32) -> u32 { + let s = t_scales[(k / params.group_size) * params.padded_N + n]; + var q: i32 = 0; + if (s != 0.0) { + q = i32(round(t_latent[n * params.K + k] / s)); + } + q = clamp(q, -8, 7); + return u32(q + 8) & 0xFu; +} + +@compute @workgroup_size(wg_size, 1, 1) +fn main(@builtin(global_invocation_id) gid: vec3) { + let word_idx = gid.x; + if (word_idx >= params.num_words) { + return; + } + var word: u32 = 0u; + for (var bi: u32 = 0u; bi < 4u; bi = bi + 1u) { + let byte_idx = word_idx * 4u + bi; + let n = byte_idx / params.K_packed; + let byte_in_row = byte_idx % params.K_packed; + let k_lo = byte_in_row * 2u; + var b: u32 = 0u; + if (n < params.N && k_lo < params.K) { + b = quant_nibble(n, k_lo); + let k_hi = k_lo + 1u; + if (k_hi < params.K) { + b = b | (quant_nibble(n, k_hi) << 4u); + } + } + word = word | (b << (bi * 8u)); + } + t_packed[word_idx] = word; +} +)"; + +inline constexpr uint32_t kQ4gswRequantWorkgroupSizeX = 64; +inline constexpr uint32_t kQ4gswRequantWorkgroupSizeY = 1; +inline constexpr uint32_t kQ4gswRequantWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu