From e77cbf76e36f82572f09c5458137bcf2aa31134d Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 21 Jul 2026 00:29:33 -0700 Subject: [PATCH 01/15] [ExecuTorch][WebGPU] Add et_vk.fused_ce (fused cross-entropy) to the WebGPU backend MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/20933 Add `et_vk.fused_ce` — a fused cross-entropy (loss + dlogits in one op) for the on-device training tail. Key changes: - `runtime/ops/fused_ce/` — fused-CE WGSL kernel (+ the `reduce` shader it depends on) + handler - `custom_ops_lib.py` — `fused_ce` custom-op def + `register_autograd` - `op_registry.py` — `OpFeatures` - `CMakeLists.txt` `WEBGPU_SRCS` — wire the source Training-only, WebGPU-only custom op registered under the shared Vulkan partitioner; no Vulkan kernel yet. Co-authored-with: Claude Code. ghstack-source-id: 405026078 @exported-using-ghexport Differential Revision: [D111755132](https://our.internmc.facebook.com/intern/diff/D111755132/) --- backends/vulkan/custom_ops_lib.py | 50 +++ backends/vulkan/op_registry.py | 14 + backends/webgpu/CMakeLists.txt | 1 + .../webgpu/runtime/ops/fused_ce/FusedCe.cpp | 285 ++++++++++++++++++ .../webgpu/runtime/ops/fused_ce/fused_ce.wgsl | 74 +++++ .../runtime/ops/fused_ce/fused_ce_wgsl.h | 98 ++++++ 6 files changed, 522 insertions(+) create mode 100644 backends/webgpu/runtime/ops/fused_ce/FusedCe.cpp create mode 100644 backends/webgpu/runtime/ops/fused_ce/fused_ce.wgsl create mode 100644 backends/webgpu/runtime/ops/fused_ce/fused_ce_wgsl.h diff --git a/backends/vulkan/custom_ops_lib.py b/backends/vulkan/custom_ops_lib.py index 4f5343700a6..cfa32282236 100644 --- a/backends/vulkan/custom_ops_lib.py +++ b/backends/vulkan/custom_ops_lib.py @@ -1153,6 +1153,56 @@ def rms_norm_impl( rms_norm_op = getattr(getattr(torch.ops, namespace), name) +######################## +## fused_ce (training) ## +######################## + + +def fused_ce_impl( + logits: torch.Tensor, + labels: torch.Tensor, + n_valid: float, +) -> tuple[torch.Tensor, torch.Tensor]: + mask = labels >= 0 + safe = labels.clamp(min=0).long() + lse = torch.logsumexp(logits, dim=-1) + picked = logits.gather(-1, safe[:, None]).squeeze(-1) + loss = torch.where(mask, (lse - picked) / n_valid, torch.zeros_like(lse)).sum() + softmax = torch.softmax(logits, dim=-1) + onehot = torch.nn.functional.one_hot(safe, logits.shape[-1]).to(logits.dtype) + dlogits = torch.where( + mask[:, None], (softmax - onehot) / n_valid, torch.zeros_like(softmax) + ) + return loss, dlogits + + +def fused_ce_meta( + logits: torch.Tensor, + labels: torch.Tensor, + n_valid: float, +) -> tuple[torch.Tensor, torch.Tensor]: + return logits.new_empty([]), torch.empty_like(logits) + + +def fused_ce_setup_context(ctx, inputs, output) -> None: + ctx.save_for_backward(output[1]) + + +def fused_ce_backward(ctx, grad_loss, grad_dlogits): + (dlogits,) = ctx.saved_tensors + return grad_loss * dlogits, None, None + + +name = "fused_ce" +lib.define(f"{name}(Tensor logits, Tensor labels, float n_valid) -> (Tensor, Tensor)") +lib.impl(name, fused_ce_impl, "CompositeExplicitAutograd") +lib.impl(name, fused_ce_meta, "Meta") +torch.library.register_autograd( + f"{namespace}::{name}", fused_ce_backward, setup_context=fused_ce_setup_context +) +fused_ce_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, diff --git a/backends/vulkan/op_registry.py b/backends/vulkan/op_registry.py index db78283e920..06c70ba1a04 100644 --- a/backends/vulkan/op_registry.py +++ b/backends/vulkan/op_registry.py @@ -1755,6 +1755,20 @@ def register_rms_norm(): ) +# ============================================================================= +# FusedCe.cpp (training) +# ============================================================================= + + +@update_features(exir_ops.edge.et_vk.fused_ce.default) +def register_fused_ce(): + return OpFeatures( + inputs_storage=utils.CONTIGUOUS_ANY, + inputs_dtypes=[utils.FP_T, utils.INT_T, utils.NONE_T], + outputs_dtypes=[utils.FP_T, utils.FP_T], + ) + + @update_features( [ exir_ops.edge.aten.ne.Scalar, diff --git a/backends/webgpu/CMakeLists.txt b/backends/webgpu/CMakeLists.txt index 39cead85ae4..62dc0e69c6d 100644 --- a/backends/webgpu/CMakeLists.txt +++ b/backends/webgpu/CMakeLists.txt @@ -54,6 +54,7 @@ set(WEBGPU_SRCS runtime/ops/index/Index.cpp runtime/ops/sdpa_fd_decode/SdpaFdDecode.cpp runtime/ops/mm/Mm.cpp + runtime/ops/fused_ce/FusedCe.cpp runtime/ops/log_softmax/LogSoftmax.cpp runtime/ops/softmax/Softmax.cpp runtime/ops/bmm/Bmm.cpp diff --git a/backends/webgpu/runtime/ops/fused_ce/FusedCe.cpp b/backends/webgpu/runtime/ops/fused_ce/FusedCe.cpp new file mode 100644 index 00000000000..e51052f978d --- /dev/null +++ b/backends/webgpu/runtime/ops/fused_ce/FusedCe.cpp @@ -0,0 +1,285 @@ +/* + * 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 +#include + +namespace executorch::backends::webgpu { + +namespace { + +// Uniform layout matching the fused_ce.wgsl Params struct (16-byte aligned). +struct FusedCeParams { + uint32_t vocab; + uint32_t n_rows; + float n_valid; + float _pad0; +}; +static_assert(sizeof(FusedCeParams) == 16, "FusedCeParams must be 16 bytes"); + +// Mirror reduce.wgsl Params (file-local in Reduce.cpp; re-declared here). +struct ReduceParams { + uint32_t outer; + uint32_t r; + uint32_t inner; + uint32_t is_mean; +}; +static_assert(sizeof(ReduceParams) == 16, "ReduceParams must be 16 bytes"); + +WGPUShaderModule make_shader(WGPUDevice device, const char* wgsl) { + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {wgsl, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + return wgpuDeviceCreateShaderModule(device, &shader_desc); +} + +WGPUBuffer create_uniform( + WebGPUGraph& graph, + WGPUDevice device, + const void* data, + size_t size) { + WGPUBufferDescriptor desc = {}; + desc.size = size; + desc.usage = WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst; + desc.mappedAtCreation = true; + WGPUBuffer buffer = wgpuDeviceCreateBuffer(device, &desc); + std::memcpy(wgpuBufferGetMappedRange(buffer, 0, size), data, size); + wgpuBufferUnmap(buffer); + graph.add_uniform_buffer_bytes(size); + return buffer; +} + +// out valuelist packs the 2-tuple (loss, dlogits) as one id. +void fused_ce_impl(WebGPUGraph& graph, const std::vector& args) { + const int logits_id = args.at(0); + const int labels_id = args.at(1); + const int n_valid_id = args.at(2); + const std::vector& outs = graph.get_value_list(args.at(3)); + if (outs.size() != 2) { + throw std::runtime_error( + "WebGPU fused_ce: expected 2 outputs (loss, dlogits)"); + } + const int loss_id = outs.at(0); + const int dlogits_id = outs.at(1); + + WGPUDevice device = graph.device(); + const auto& logits = graph.get_tensor(logits_id); + const auto& labels = graph.get_tensor(labels_id); + const auto& dlogits = graph.get_tensor(dlogits_id); + const auto& loss = graph.get_tensor(loss_id); + + if (logits.dims.size() != 2) { + throw std::runtime_error("WebGPU fused_ce: logits must be 2D [N, V]"); + } + const uint64_t n_rows = static_cast(logits.dims[0]); + const uint64_t vocab = static_cast(logits.dims[1]); + const uint64_t numel = n_rows * vocab; + + if (dlogits.dims != logits.dims) { + throw std::runtime_error( + "WebGPU fused_ce: dlogits shape must match logits"); + } + if (logits.nbytes != numel * sizeof(float) || + dlogits.nbytes != numel * sizeof(float)) { + throw std::runtime_error("WebGPU fused_ce: logits/dlogits fp32-only"); + } + if (labels.nbytes != n_rows * sizeof(int32_t)) { + throw std::runtime_error("WebGPU fused_ce: labels must be int32 [N]"); + } + if (loss.nbytes != sizeof(float)) { + throw std::runtime_error("WebGPU fused_ce: loss must be a scalar [1]"); + } + if (graph.get_value_type(n_valid_id) != WebGPUGraph::ValueType::Double) { + throw std::runtime_error("WebGPU fused_ce: n_valid must be a float scalar"); + } + const double n_valid = graph.get_double(n_valid_id); + if (n_valid <= 0.0) { + throw std::runtime_error("WebGPU fused_ce: n_valid must be positive"); + } + if (n_rows > utils::queried_max_workgroups(device)) { + throw std::runtime_error("WebGPU fused_ce: n_rows exceeds dispatch limit"); + } + + WGPUBuffer loss_partial = graph.create_scratch_buffer(n_rows * sizeof(float)); + + // one workgroup per row + FusedCeParams ce_params = {}; + ce_params.vocab = static_cast(vocab); + ce_params.n_rows = static_cast(n_rows); + ce_params.n_valid = static_cast(n_valid); + const uint32_t ce_wg = + utils::clamp_workgroup_size(device, kFusedCeWorkgroupSizeX); + WGPUBuffer ce_uniform = + create_uniform(graph, device, &ce_params, sizeof(ce_params)); + WGPUShaderModule ce_shader = make_shader(device, kFusedCeWGSL); + + WGPUBindGroupLayoutEntry ce_entries[5] = {}; + ce_entries[0].binding = 0; + ce_entries[0].visibility = WGPUShaderStage_Compute; + ce_entries[0].buffer.type = WGPUBufferBindingType_Storage; + ce_entries[1].binding = 1; + ce_entries[1].visibility = WGPUShaderStage_Compute; + ce_entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + ce_entries[2].binding = 2; + ce_entries[2].visibility = WGPUShaderStage_Compute; + ce_entries[2].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + ce_entries[3].binding = 3; + ce_entries[3].visibility = WGPUShaderStage_Compute; + ce_entries[3].buffer.type = WGPUBufferBindingType_Storage; + ce_entries[4].binding = 4; + ce_entries[4].visibility = WGPUShaderStage_Compute; + ce_entries[4].buffer.type = WGPUBufferBindingType_Uniform; + + WGPUBindGroupLayoutDescriptor ce_bgl_desc = {}; + ce_bgl_desc.entryCount = 5; + ce_bgl_desc.entries = ce_entries; + WGPUBindGroupLayout ce_bgl = + wgpuDeviceCreateBindGroupLayout(device, &ce_bgl_desc); + + WGPUPipelineLayoutDescriptor ce_pl_desc = {}; + ce_pl_desc.bindGroupLayoutCount = 1; + ce_pl_desc.bindGroupLayouts = &ce_bgl; + WGPUPipelineLayout ce_pl = + wgpuDeviceCreatePipelineLayout(device, &ce_pl_desc); + + WGPUConstantEntry ce_wg_const = {}; + ce_wg_const.key = {"wg_size", WGPU_STRLEN}; + ce_wg_const.value = static_cast(ce_wg); + + WGPUComputePipelineDescriptor ce_pipe_desc = {}; + ce_pipe_desc.layout = ce_pl; + ce_pipe_desc.compute.module = ce_shader; + ce_pipe_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + ce_pipe_desc.compute.constantCount = 1; + ce_pipe_desc.compute.constants = &ce_wg_const; + WGPUComputePipeline ce_pipe = + wgpuDeviceCreateComputePipeline(device, &ce_pipe_desc); + + WGPUBindGroupEntry ce_bg[5] = {}; + ce_bg[0].binding = 0; + ce_bg[0].buffer = dlogits.buffer; + ce_bg[0].size = dlogits.nbytes; + ce_bg[1].binding = 1; + ce_bg[1].buffer = logits.buffer; + ce_bg[1].size = logits.nbytes; + ce_bg[2].binding = 2; + ce_bg[2].buffer = labels.buffer; + ce_bg[2].size = labels.nbytes; + ce_bg[3].binding = 3; + ce_bg[3].buffer = loss_partial; + ce_bg[3].size = n_rows * sizeof(float); + ce_bg[4].binding = 4; + ce_bg[4].buffer = ce_uniform; + ce_bg[4].size = sizeof(ce_params); + + WGPUBindGroupDescriptor ce_bg_desc = {}; + ce_bg_desc.layout = ce_bgl; + ce_bg_desc.entryCount = 5; + ce_bg_desc.entries = ce_bg; + WGPUBindGroup ce_bind_group = wgpuDeviceCreateBindGroup(device, &ce_bg_desc); + + graph.add_dispatch( + {ce_pipe, ce_bind_group, static_cast(n_rows), "fused_ce"}); + + wgpuShaderModuleRelease(ce_shader); + wgpuBindGroupLayoutRelease(ce_bgl); + wgpuPipelineLayoutRelease(ce_pl); + wgpuBufferRelease(ce_uniform); + + // reduce loss_partial[N] -> loss[1] (reuses reduce.wgsl) + ReduceParams r_params = {}; + r_params.outer = 1u; + r_params.r = static_cast(n_rows); + r_params.inner = 1u; + r_params.is_mean = 0u; // mean already folded into loss_partial + const uint32_t r_wg = + utils::clamp_workgroup_size(device, kReduceWorkgroupSizeX); + const uint32_t r_wgc = + utils::compute_1d_workgroup_count(device, 1u, r_wg, "fused_ce_reduce"); + WGPUBuffer r_uniform = + create_uniform(graph, device, &r_params, sizeof(r_params)); + WGPUShaderModule r_shader = make_shader(device, kReduceWGSL); + + WGPUBindGroupLayoutEntry r_entries[3] = {}; + r_entries[0].binding = 0; + r_entries[0].visibility = WGPUShaderStage_Compute; + r_entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + r_entries[1].binding = 1; + r_entries[1].visibility = WGPUShaderStage_Compute; + r_entries[1].buffer.type = WGPUBufferBindingType_Storage; + r_entries[2].binding = 2; + r_entries[2].visibility = WGPUShaderStage_Compute; + r_entries[2].buffer.type = WGPUBufferBindingType_Uniform; + + WGPUBindGroupLayoutDescriptor r_bgl_desc = {}; + r_bgl_desc.entryCount = 3; + r_bgl_desc.entries = r_entries; + WGPUBindGroupLayout r_bgl = + wgpuDeviceCreateBindGroupLayout(device, &r_bgl_desc); + + WGPUPipelineLayoutDescriptor r_pl_desc = {}; + r_pl_desc.bindGroupLayoutCount = 1; + r_pl_desc.bindGroupLayouts = &r_bgl; + WGPUPipelineLayout r_pl = wgpuDeviceCreatePipelineLayout(device, &r_pl_desc); + + WGPUConstantEntry r_wg_const = {}; + r_wg_const.key = {"wg_size", WGPU_STRLEN}; + r_wg_const.value = static_cast(r_wg); + + WGPUComputePipelineDescriptor r_pipe_desc = {}; + r_pipe_desc.layout = r_pl; + r_pipe_desc.compute.module = r_shader; + r_pipe_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + r_pipe_desc.compute.constantCount = 1; + r_pipe_desc.compute.constants = &r_wg_const; + WGPUComputePipeline r_pipe = + wgpuDeviceCreateComputePipeline(device, &r_pipe_desc); + + WGPUBindGroupEntry r_bg[3] = {}; + r_bg[0].binding = 0; + r_bg[0].buffer = loss_partial; + r_bg[0].size = n_rows * sizeof(float); + r_bg[1].binding = 1; + r_bg[1].buffer = loss.buffer; + r_bg[1].size = loss.nbytes; + r_bg[2].binding = 2; + r_bg[2].buffer = r_uniform; + r_bg[2].size = sizeof(r_params); + + WGPUBindGroupDescriptor r_bg_desc = {}; + r_bg_desc.layout = r_bgl; + r_bg_desc.entryCount = 3; + r_bg_desc.entries = r_bg; + WGPUBindGroup r_bind_group = wgpuDeviceCreateBindGroup(device, &r_bg_desc); + + graph.add_dispatch({r_pipe, r_bind_group, r_wgc, "fused_ce_reduce"}); + + wgpuShaderModuleRelease(r_shader); + wgpuBindGroupLayoutRelease(r_bgl); + wgpuPipelineLayoutRelease(r_pl); + wgpuBufferRelease(r_uniform); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(et_vk.fused_ce.default, fused_ce_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/fused_ce/fused_ce.wgsl b/backends/webgpu/runtime/ops/fused_ce/fused_ce.wgsl new file mode 100644 index 00000000000..5910e5e9576 --- /dev/null +++ b/backends/webgpu/runtime/ops/fused_ce/fused_ce.wgsl @@ -0,0 +1,74 @@ + +@group(0) @binding(0) var dlogits: array; +@group(0) @binding(1) var logits: array; +@group(0) @binding(2) var labels: array; +@group(0) @binding(3) var loss_partial: array; + +struct Params { vocab: u32, n_rows: u32, n_valid: f32, pad: f32 } +@group(0) @binding(4) var params: Params; + +override wg_size: u32 = 256u; +var red_m: array; +var red_l: array; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(workgroup_id) wg: vec3, + @builtin(local_invocation_id) lid: vec3) { + let row = wg.x; + if (row >= params.n_rows) { return; } + let tid = lid.x; + let V = params.vocab; + let base = row * V; + let lbl = labels[row]; + + if (lbl < 0) { + for (var j = tid; j < V; j = j + wg_size) { dlogits[base + j] = 0.0; } + if (tid == 0u) { loss_partial[row] = 0.0; } + return; + } + + // -3.4e38 not 3.4028235e38 (Tint f32 overflow) + var m = -3.4e38; + var l = 0.0; + for (var j = tid; j < V; j = j + wg_size) { + let x = logits[base + j]; + if (x > m) { + l = l * exp(m - x) + 1.0; + m = x; + } else { + l = l + exp(x - m); + } + } + red_m[tid] = m; + red_l[tid] = l; + workgroupBarrier(); + for (var s = wg_size / 2u; s > 0u; s = s >> 1u) { + if (tid < s) { + let ma = red_m[tid]; + let la = red_l[tid]; + let mb = red_m[tid + s]; + let lb = red_l[tid + s]; + let mm = max(ma, mb); + red_m[tid] = mm; + red_l[tid] = la * exp(ma - mm) + lb * exp(mb - mm); + } + workgroupBarrier(); + } + let row_max = red_m[0]; + let denom = red_l[0]; + workgroupBarrier(); + + let inv = 1.0 / denom; + let scale = 1.0 / params.n_valid; + if (tid == 0u) { + let lse = row_max + log(denom); + loss_partial[row] = (lse - logits[base + u32(lbl)]) * scale; + } + workgroupBarrier(); + for (var j = tid; j < V; j = j + wg_size) { + var g = exp(logits[base + j] - row_max) * inv * scale; + if (j == u32(lbl)) { g = g - scale; } + dlogits[base + j] = g; + } +} diff --git a/backends/webgpu/runtime/ops/fused_ce/fused_ce_wgsl.h b/backends/webgpu/runtime/ops/fused_ce/fused_ce_wgsl.h new file mode 100644 index 00000000000..6ed1f709b25 --- /dev/null +++ b/backends/webgpu/runtime/ops/fused_ce/fused_ce_wgsl.h @@ -0,0 +1,98 @@ +/* + * 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 fused_ce.wgsl - DO NOT EDIT. +// wgsl-sha256: 416484a64dc4f31940ca63965e28b5747b34cd6b3a64e141ffaa5835ec28891d +inline constexpr const char* kFusedCeWGSL = R"( + +@group(0) @binding(0) var dlogits: array; +@group(0) @binding(1) var logits: array; +@group(0) @binding(2) var labels: array; +@group(0) @binding(3) var loss_partial: array; + +struct Params { vocab: u32, n_rows: u32, n_valid: f32, pad: f32 } +@group(0) @binding(4) var params: Params; + +override wg_size: u32 = 256u; +var red_m: array; +var red_l: array; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(workgroup_id) wg: vec3, + @builtin(local_invocation_id) lid: vec3) { + let row = wg.x; + if (row >= params.n_rows) { return; } + let tid = lid.x; + let V = params.vocab; + let base = row * V; + let lbl = labels[row]; + + if (lbl < 0) { + for (var j = tid; j < V; j = j + wg_size) { dlogits[base + j] = 0.0; } + if (tid == 0u) { loss_partial[row] = 0.0; } + return; + } + + // -3.4e38 not 3.4028235e38 (Tint f32 overflow) + var m = -3.4e38; + var l = 0.0; + for (var j = tid; j < V; j = j + wg_size) { + let x = logits[base + j]; + if (x > m) { + l = l * exp(m - x) + 1.0; + m = x; + } else { + l = l + exp(x - m); + } + } + red_m[tid] = m; + red_l[tid] = l; + workgroupBarrier(); + for (var s = wg_size / 2u; s > 0u; s = s >> 1u) { + if (tid < s) { + let ma = red_m[tid]; + let la = red_l[tid]; + let mb = red_m[tid + s]; + let lb = red_l[tid + s]; + let mm = max(ma, mb); + red_m[tid] = mm; + red_l[tid] = la * exp(ma - mm) + lb * exp(mb - mm); + } + workgroupBarrier(); + } + let row_max = red_m[0]; + let denom = red_l[0]; + workgroupBarrier(); + + let inv = 1.0 / denom; + let scale = 1.0 / params.n_valid; + if (tid == 0u) { + let lse = row_max + log(denom); + loss_partial[row] = (lse - logits[base + u32(lbl)]) * scale; + } + workgroupBarrier(); + for (var j = tid; j < V; j = j + wg_size) { + var g = exp(logits[base + j] - row_max) * inv * scale; + if (j == u32(lbl)) { g = g - scale; } + dlogits[base + j] = g; + } +} +)"; + +inline constexpr uint32_t kFusedCeWorkgroupSizeX = 256; +inline constexpr uint32_t kFusedCeWorkgroupSizeY = 1; +inline constexpr uint32_t kFusedCeWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From 3e45164f8227d3f09e67696c8f71b8f37adb79ed Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 21 Jul 2026 00:29:33 -0700 Subject: [PATCH 02/15] [ExecuTorch][WebGPU] Tests for the fused_ce WebGPU ops MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/20934 Export-delegation + fp64 golden tests for the fused_ce ops. Key changes: - `test/ops/test_fused_ce.py` — fp64 loss + dlogits goldens (incl. masked-label) + delegation checks — fp64 library-reference goldens + delegation checks (≥2 cases/op). Co-authored-with: Claude Code. ghstack-source-id: 405026080 @exported-using-ghexport Differential Revision: [D111755119](https://our.internmc.facebook.com/intern/diff/D111755119/) --- backends/webgpu/test/ops/test_fused_ce.py | 130 ++++++++++++++++++++++ 1 file changed, 130 insertions(+) create mode 100644 backends/webgpu/test/ops/test_fused_ce.py diff --git a/backends/webgpu/test/ops/test_fused_ce.py b/backends/webgpu/test/ops/test_fused_ce.py new file mode 100644 index 00000000000..13d2131c64a --- /dev/null +++ b/backends/webgpu/test/ops/test_fused_ce.py @@ -0,0 +1,130 @@ +# 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. + +"""Fused cross-entropy training op (`et_vk.fused_ce`) export + fp64 golden. + +`fused_ce(logits[M,V], labels[M], n_valid) -> (loss, dlogits[M,V])` computes the +mean-over-valid CE loss and its gradient in one op (labels < 0 are ignored/pad). +Golden is the fp64 reference (`logsumexp - picked`, `softmax - onehot`), the same math +torch's cross_entropy uses; the native test reconstructs the deterministic inputs. +""" + +import os +import unittest +from dataclasses import dataclass + +import numpy as np +import torch + +from executorch.backends.vulkan import VulkanPartitioner +from executorch.exir import to_edge_transform_and_lower + + +@dataclass(frozen=True) +class CeConfig: + name: str + m: int # rows (valid + pad positions) + v: int # vocab + n_pad: int = 0 # trailing rows with label = -1 (ignored) + + +# Mirrored by the C++ kFusedCeConfigs table. Llama-3.2-1B vocab = 128256. +CONFIGS = [ + CeConfig("tiny", 4, 32), + CeConfig("masked", 8, 128, n_pad=3), # some ignored labels + CeConfig("llama_vocab", 16, 128256), # real vocab width +] + + +def _inputs(cfg: CeConfig): + """Deterministic logits [M,V] + labels [M] (last n_pad = -1); reconstructable in C++.""" + flat = np.arange(cfg.m * cfg.v, dtype=np.int64) + logits = torch.from_numpy( + (((flat % 23) - 11).astype(np.float32) / np.float32(8.0)).reshape(cfg.m, cfg.v) + ) + labels = torch.from_numpy((np.arange(cfg.m, dtype=np.int64) * 7 + 3) % cfg.v) + if cfg.n_pad: + labels[cfg.m - cfg.n_pad :] = -1 + n_valid = float(max(1, cfg.m - cfg.n_pad)) + return logits, labels, n_valid + + +def _fp64_golden(logits: torch.Tensor, labels: torch.Tensor, n_valid: float): + mask = labels >= 0 + safe = labels.clamp(min=0).long() + lg = logits.double() + lse = torch.logsumexp(lg, dim=-1) + picked = lg.gather(-1, safe[:, None]).squeeze(-1) + loss = torch.where(mask, (lse - picked) / n_valid, torch.zeros_like(lse)).sum() + softmax = torch.softmax(lg, dim=-1) + onehot = torch.nn.functional.one_hot(safe, logits.shape[-1]).double() + dlogits = torch.where( + mask[:, None], (softmax - onehot) / n_valid, torch.zeros_like(softmax) + ) + return loss.to(torch.float32), dlogits.to(torch.float32) + + +class _CeModule(torch.nn.Module): + def forward(self, logits, labels, n_valid): + return torch.ops.et_vk.fused_ce(logits, labels, n_valid) + + +def _export(logits, labels, n_valid): + ep = torch.export.export(_CeModule(), (logits, labels, n_valid)) + return to_edge_transform_and_lower( + ep, partitioner=[VulkanPartitioner()] + ).to_executorch() + + +class TestFusedCe(unittest.TestCase): + def test_export_delegates(self) -> None: + for cfg in CONFIGS: + if cfg.v > 1024: # width-independent; skip the 128k fixture + continue + with self.subTest(config=cfg.name): + logits, labels, n_valid = _inputs(cfg) + et = _export(logits, labels, n_valid) + found = any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ) + self.assertTrue(found, f"no VulkanBackend delegate in {cfg.name}") + + def test_op_matches_fp64_golden(self) -> None: + for cfg in CONFIGS: + if cfg.v > 1024: + continue + with self.subTest(config=cfg.name): + logits, labels, n_valid = _inputs(cfg) + loss, dlogits = torch.ops.et_vk.fused_ce(logits, labels, n_valid) + g_loss, g_dlogits = _fp64_golden(logits, labels, n_valid) + torch.testing.assert_close(loss, g_loss, atol=5e-4, rtol=1e-3) + torch.testing.assert_close(dlogits, g_dlogits, atol=5e-4, rtol=1e-3) + + +def export_fused_ce_model(cfg: CeConfig, pte_path: str, golden_path: str) -> None: + logits, labels, n_valid = _inputs(cfg) + et = _export(logits, labels, n_valid) + with open(pte_path, "wb") as f: + f.write(et.buffer) + g_loss, g_dlogits = _fp64_golden(logits, labels, n_valid) + # loss scalar then dlogits, both raw LE fp32 + np.concatenate([g_loss.reshape(1).numpy(), g_dlogits.reshape(-1).numpy()]).astype( + " None: + for cfg in CONFIGS: + pte = os.path.join(out_dir, f"fused_ce_{cfg.name}.pte") + golden = os.path.join(out_dir, f"fused_ce_{cfg.name}.golden.bin") + export_fused_ce_model(cfg, pte, golden) + + +if __name__ == "__main__": + unittest.main() From 70ff4e047c04d2a925fa778a633da167774eb5cf Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 21 Jul 2026 00:29:33 -0700 Subject: [PATCH 03/15] [ExecuTorch][WebGPU] Add et_vk.adamw_step (on-GPU AdamW optimizer step) to the WebGPU backend MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/20935 Add `et_vk.adamw_step` — an on-GPU AdamW optimizer step (delegated in-place parameter update) for on-device training. Key changes: - `runtime/ops/adamw/` — elementwise AdamW-update WGSL kernel + handler - `custom_ops_lib.py` — `adamw_step` custom-op def (mutating, `Tensor(a!)`) - `op_registry.py` — `OpFeatures` (string key `et_vk::adamw_step`) - `CMakeLists.txt` `WEBGPU_SRCS` — wire the source Delegating the optimizer step to the GPU is new to ExecuTorch — its optimizers are host-side C++ (`extension/training/optimizer/adamw.cpp`). Training-only WebGPU custom op under the shared Vulkan partitioner; no Vulkan kernel yet. Co-authored-with: Claude Code. ghstack-source-id: 405026081 @exported-using-ghexport Differential Revision: [D111755116](https://our.internmc.facebook.com/intern/diff/D111755116/) --- backends/vulkan/custom_ops_lib.py | 52 +++++ backends/vulkan/op_registry.py | 8 + backends/webgpu/CMakeLists.txt | 1 + .../webgpu/runtime/ops/adamw/AdamwStep.cpp | 181 ++++++++++++++++++ .../webgpu/runtime/ops/adamw/adamw_step.wgsl | 41 ++++ .../runtime/ops/adamw/adamw_step_wgsl.h | 65 +++++++ 6 files changed, 348 insertions(+) create mode 100644 backends/webgpu/runtime/ops/adamw/AdamwStep.cpp create mode 100644 backends/webgpu/runtime/ops/adamw/adamw_step.wgsl create mode 100644 backends/webgpu/runtime/ops/adamw/adamw_step_wgsl.h diff --git a/backends/vulkan/custom_ops_lib.py b/backends/vulkan/custom_ops_lib.py index cfa32282236..a074597466a 100644 --- a/backends/vulkan/custom_ops_lib.py +++ b/backends/vulkan/custom_ops_lib.py @@ -1203,6 +1203,58 @@ def fused_ce_backward(ctx, grad_loss, grad_dlogits): fused_ce_op = getattr(getattr(torch.ops, namespace), name) +########################### +## adamw_step (training) ## +########################### + + +def adamw_step_impl( + param: torch.Tensor, + m: torch.Tensor, + v: torch.Tensor, + grad: torch.Tensor, + lr: float, + beta1: float, + beta2: float, + eps: float, + weight_decay: float, + bias_correction1: float, + bias_correction2: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + param.mul_(1.0 - lr * weight_decay) + m.mul_(beta1).add_(grad, alpha=1.0 - beta1) + v.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2) + mhat = m / bias_correction1 + denom = (v / bias_correction2).sqrt() + eps + param.addcdiv_(mhat, denom, value=-lr) + return param, m, v + + +def adamw_step_meta( + param: torch.Tensor, + m: torch.Tensor, + v: torch.Tensor, + grad: torch.Tensor, + lr: float, + beta1: float, + beta2: float, + eps: float, + weight_decay: float, + bias_correction1: float, + bias_correction2: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return param, m, v + + +name = "adamw_step" +lib.define( + f"{name}(Tensor(a!) param, Tensor(b!) m, Tensor(c!) v, Tensor grad, float lr, float beta1, float beta2, float eps, float weight_decay, float bias_correction1, float bias_correction2) -> (Tensor(a!), Tensor(b!), Tensor(c!))" +) +lib.impl(name, adamw_step_impl, "CompositeExplicitAutograd") +lib.impl(name, adamw_step_meta, "Meta") +adamw_step_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, diff --git a/backends/vulkan/op_registry.py b/backends/vulkan/op_registry.py index 06c70ba1a04..78270be4373 100644 --- a/backends/vulkan/op_registry.py +++ b/backends/vulkan/op_registry.py @@ -1798,6 +1798,14 @@ def register_logical_not(): ) +@update_features("et_vk::adamw_step") +def register_adamw_step(): + return OpFeatures( + inputs_storage=utils.CONTIGUOUS_ANY, + inputs_dtypes=utils.FP_T, + ) + + @update_features(exir_ops.edge.et_vk.linear_dW.default) def register_linear_dW(): return OpFeatures( diff --git a/backends/webgpu/CMakeLists.txt b/backends/webgpu/CMakeLists.txt index 62dc0e69c6d..ff5a0fc172a 100644 --- a/backends/webgpu/CMakeLists.txt +++ b/backends/webgpu/CMakeLists.txt @@ -69,6 +69,7 @@ set(WEBGPU_SRCS runtime/ops/dim_order/DimOrder.cpp runtime/ops/linear/Linear.cpp runtime/ops/embedding/Embedding.cpp + runtime/ops/adamw/AdamwStep.cpp runtime/ops/quantized_linear/LinearDw.cpp runtime/ops/quantized_linear/QuantizedLinearRequant.cpp ) diff --git a/backends/webgpu/runtime/ops/adamw/AdamwStep.cpp b/backends/webgpu/runtime/ops/adamw/AdamwStep.cpp new file mode 100644 index 00000000000..4a78694d40b --- /dev/null +++ b/backends/webgpu/runtime/ops/adamw/AdamwStep.cpp @@ -0,0 +1,181 @@ +/* + * 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 AdamwStepParams { + uint32_t numel; + uint32_t _pad0; + uint32_t _pad1; + uint32_t _pad2; + float lr; + float beta1; + float beta2; + float eps; + float weight_decay; + float bias_correction1; + float bias_correction2; + float _pad3; +}; +static_assert(sizeof(AdamwStepParams) == 48, "params must be 48 bytes"); + +// AdamW step over an fp32 latent (elementwise, in place); mirrors torch AdamW. +void adamw_step_impl(WebGPUGraph& graph, const std::vector& args) { + const int param_id = args.at(0); + const int m_id = args.at(1); + const int v_id = args.at(2); + const int grad_id = args.at(3); + + WGPUDevice device = graph.device(); + const auto& param = graph.get_tensor(param_id); + const auto& m = graph.get_tensor(m_id); + const auto& v = graph.get_tensor(v_id); + const auto& grad = graph.get_tensor(grad_id); + + uint64_t numel = 1; + for (int64_t d : param.dims) { + numel *= static_cast(d); + } + if (param.dims.empty() || numel == 0) { + throw std::runtime_error("adamw_step: empty param"); + } + const uint64_t bytes = numel * sizeof(float); + if (param.nbytes != bytes || m.nbytes != bytes || v.nbytes != bytes || + grad.nbytes != bytes) { + throw std::runtime_error( + "adamw_step: param/m/v/grad must be fp32 and same numel"); + } + if (numel > UINT32_MAX) { + throw std::runtime_error("adamw_step: numel exceeds u32"); + } + + auto scalar = [&](int id, const char* name) -> float { + if (graph.get_value_type(id) != WebGPUGraph::ValueType::Double) { + throw std::runtime_error( + std::string("adamw_step: ") + name + " must be a float scalar"); + } + return static_cast(graph.get_double(id)); + }; + + AdamwStepParams params = {}; + params.numel = static_cast(numel); + params.lr = scalar(args.at(4), "lr"); + params.beta1 = scalar(args.at(5), "beta1"); + params.beta2 = scalar(args.at(6), "beta2"); + params.eps = scalar(args.at(7), "eps"); + params.weight_decay = scalar(args.at(8), "weight_decay"); + params.bias_correction1 = scalar(args.at(9), "bias_correction1"); + params.bias_correction2 = scalar(args.at(10), "bias_correction2"); + if (params.bias_correction1 == 0.0f || params.bias_correction2 == 0.0f) { + throw std::runtime_error("adamw_step: bias corrections must be non-zero"); + } + + const uint32_t wg_size = + utils::clamp_workgroup_size(device, kAdamwStepWorkgroupSizeX); + const uint32_t workgroup_count = utils::compute_1d_workgroup_count( + device, params.numel, wg_size, "adamw_step"); + + 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 = {kAdamwStepWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + WGPUBindGroupLayoutEntry entries[5] = {}; + for (uint32_t i = 0; i <= 2; i++) { + entries[i].binding = i; + entries[i].visibility = WGPUShaderStage_Compute; + entries[i].buffer.type = WGPUBufferBindingType_Storage; + } + entries[3].binding = 3; + entries[3].visibility = WGPUShaderStage_Compute; + entries[3].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 = param.buffer; + bg_entries[0].size = param.nbytes; + bg_entries[1].binding = 1; + bg_entries[1].buffer = m.buffer; + bg_entries[1].size = m.nbytes; + bg_entries[2].binding = 2; + bg_entries[2].buffer = v.buffer; + bg_entries[2].size = v.nbytes; + bg_entries[3].binding = 3; + bg_entries[3].buffer = grad.buffer; + bg_entries[3].size = grad.nbytes; + bg_entries[4].binding = 4; + bg_entries[4].buffer = uniform_buffer; + bg_entries[4].size = sizeof(params); + + 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, "adamw_step"}); + + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + graph.own_uniform_buffer(uniform_buffer); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(et_vk.adamw_step.default, adamw_step_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/adamw/adamw_step.wgsl b/backends/webgpu/runtime/ops/adamw/adamw_step.wgsl new file mode 100644 index 00000000000..f6a5ad58943 --- /dev/null +++ b/backends/webgpu/runtime/ops/adamw/adamw_step.wgsl @@ -0,0 +1,41 @@ + +@group(0) @binding(0) var t_param: array; +@group(0) @binding(1) var t_m: array; +@group(0) @binding(2) var t_v: array; +@group(0) @binding(3) var t_grad: array; + +struct Params { + numel: u32, + _pad0: u32, + _pad1: u32, + _pad2: u32, + lr: f32, + beta1: f32, + beta2: f32, + eps: f32, + weight_decay: f32, + bias_correction1: f32, + bias_correction2: f32, + _pad3: f32, +} +@group(0) @binding(4) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main(@builtin(global_invocation_id) gid: vec3) { + let i = gid.x; + if (i >= params.numel) { + return; + } + let g = t_grad[i]; + var p = t_param[i]; + p = p - params.lr * params.weight_decay * p; + let m = params.beta1 * t_m[i] + (1.0 - params.beta1) * g; + let v = params.beta2 * t_v[i] + (1.0 - params.beta2) * g * g; + t_m[i] = m; + t_v[i] = v; + let mhat = m / params.bias_correction1; + let vhat = v / params.bias_correction2; + t_param[i] = p - params.lr * mhat / (sqrt(vhat) + params.eps); +} diff --git a/backends/webgpu/runtime/ops/adamw/adamw_step_wgsl.h b/backends/webgpu/runtime/ops/adamw/adamw_step_wgsl.h new file mode 100644 index 00000000000..38cbda98c6b --- /dev/null +++ b/backends/webgpu/runtime/ops/adamw/adamw_step_wgsl.h @@ -0,0 +1,65 @@ +/* + * 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 adamw_step.wgsl - DO NOT EDIT. +// wgsl-sha256: 0957c04168872db5e2b39cf5f26beefba27f3b5514ec69cece4de5145e97156f +inline constexpr const char* kAdamwStepWGSL = R"( + +@group(0) @binding(0) var t_param: array; +@group(0) @binding(1) var t_m: array; +@group(0) @binding(2) var t_v: array; +@group(0) @binding(3) var t_grad: array; + +struct Params { + numel: u32, + _pad0: u32, + _pad1: u32, + _pad2: u32, + lr: f32, + beta1: f32, + beta2: f32, + eps: f32, + weight_decay: f32, + bias_correction1: f32, + bias_correction2: f32, + _pad3: f32, +} +@group(0) @binding(4) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main(@builtin(global_invocation_id) gid: vec3) { + let i = gid.x; + if (i >= params.numel) { + return; + } + let g = t_grad[i]; + var p = t_param[i]; + p = p - params.lr * params.weight_decay * p; + let m = params.beta1 * t_m[i] + (1.0 - params.beta1) * g; + let v = params.beta2 * t_v[i] + (1.0 - params.beta2) * g * g; + t_m[i] = m; + t_v[i] = v; + let mhat = m / params.bias_correction1; + let vhat = v / params.bias_correction2; + t_param[i] = p - params.lr * mhat / (sqrt(vhat) + params.eps); +} +)"; + +inline constexpr uint32_t kAdamwStepWorkgroupSizeX = 64; +inline constexpr uint32_t kAdamwStepWorkgroupSizeY = 1; +inline constexpr uint32_t kAdamwStepWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From f0941578b06c41efdee08e663a3f3175dbbdfcc1 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 21 Jul 2026 00:29:34 -0700 Subject: [PATCH 04/15] [ExecuTorch][WebGPU] Tests for the adamw WebGPU ops MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/20936 Export-delegation + fp64 golden tests for the adamw (`adamw_step`) ops. Key changes: - `test/ops/test_adamw.py` — fp64 in-place param/m/v goldens (incl. non-zero weight-decay) + delegation checks — fp64 library-reference goldens + delegation checks (≥2 cases/op). Co-authored-with: Claude Code. ghstack-source-id: 405026082 @exported-using-ghexport Differential Revision: [D111755127](https://our.internmc.facebook.com/intern/diff/D111755127/) --- backends/webgpu/test/ops/test_adamw.py | 142 +++++++++++++++++++++++++ 1 file changed, 142 insertions(+) create mode 100644 backends/webgpu/test/ops/test_adamw.py diff --git a/backends/webgpu/test/ops/test_adamw.py b/backends/webgpu/test/ops/test_adamw.py new file mode 100644 index 00000000000..b15b4d25bce --- /dev/null +++ b/backends/webgpu/test/ops/test_adamw.py @@ -0,0 +1,142 @@ +# 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. + +"""AdamW optimizer step (`et_vk.adamw_step`) export + fp64 golden. + +`adamw_step(param, m, v, grad, lr, beta1, beta2, eps, weight_decay, bc1, bc2)` +updates the fp32 latent in place: decoupled weight decay, then the bias-corrected +Adam moment update. `bc1`/`bc2` (= 1 - beta^t) are host-precomputed so the kernel +carries no step counter. The op mutates and returns `param`/`m`/`v` (aliased), so +export wraps it in `auto_functionalized`, which VulkanPartitioner tags by name (the +same mutating-op path as `update_cache`). Golden is the fp64 reference, computed +independently so a lossy fp32 op impl cannot fake-pass. +""" + +from __future__ import annotations + +import unittest +from dataclasses import dataclass + +import torch + +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.exir import to_edge_transform_and_lower + +# torch.optim.AdamW defaults; bc1/bc2 are 1 - beta^step (host-precomputed). +LR = 1e-3 +BETA1 = 0.9 +BETA2 = 0.999 +EPS = 1e-8 + + +@dataclass(frozen=True) +class AdamwConfig: + name: str + numel: int + weight_decay: float = 0.01 + step: int = 1 + + +CONFIGS = [ + AdamwConfig("small", 64), + AdamwConfig("no_wd", 256, weight_decay=0.0), # wd=0 -> the plain Adam path + AdamwConfig("later_step", 1000, step=10), # bias correction well past t=1 +] + + +def _inputs(cfg: AdamwConfig): + """Deterministic fp32 param/m/v/grad + the host bias corrections.""" + g = torch.Generator().manual_seed(0) + param = torch.randn(cfg.numel, generator=g, dtype=torch.float32) + m = torch.randn(cfg.numel, generator=g, dtype=torch.float32) * 0.1 + v = torch.rand(cfg.numel, generator=g, dtype=torch.float32) * 0.01 + grad = torch.randn(cfg.numel, generator=g, dtype=torch.float32) + bc1 = 1.0 - BETA1**cfg.step + bc2 = 1.0 - BETA2**cfg.step + return param, m, v, grad, bc1, bc2 + + +def _fp64_golden(param, m, v, grad, wd, bc1, bc2): + """fp64 truth for one AdamW step; mirrors adamw_step.wgsl exactly.""" + p = param.double() + g = grad.double() + p = p - LR * wd * p + m64 = BETA1 * m.double() + (1.0 - BETA1) * g + v64 = BETA2 * v.double() + (1.0 - BETA2) * g * g + mhat = m64 / bc1 + vhat = v64 / bc2 + p = p - LR * mhat / (torch.sqrt(vhat) + EPS) + return p.float(), m64.float(), v64.float() + + +class _AdamwModule(torch.nn.Module): + def __init__(self, wd: float, bc1: float, bc2: float) -> None: + super().__init__() + self.wd = wd + self.bc1 = bc1 + self.bc2 = bc2 + + def forward(self, param, m, v, grad): + return torch.ops.et_vk.adamw_step( + param, m, v, grad, LR, BETA1, BETA2, EPS, self.wd, self.bc1, self.bc2 + ) + + +def _export(cfg: AdamwConfig): + param, m, v, grad, bc1, bc2 = _inputs(cfg) + ep = torch.export.export( + _AdamwModule(cfg.weight_decay, bc1, bc2), (param, m, v, grad) + ) + return to_edge_transform_and_lower( + ep, partitioner=[VulkanPartitioner()] + ).to_executorch() + + +def _delegates(et) -> bool: + return any( + d.id == "VulkanBackend" + for plan in et.executorch_program.execution_plan + for d in plan.delegates + ) + + +class TestAdamwStep(unittest.TestCase): + def test_export_delegates(self) -> None: + for cfg in CONFIGS: + with self.subTest(config=cfg.name): + et = _export(cfg) + self.assertTrue( + _delegates(et), f"no VulkanBackend delegate in {cfg.name}" + ) + + def test_op_matches_fp64_golden(self) -> None: + for cfg in CONFIGS: + with self.subTest(config=cfg.name): + param, m, v, grad, bc1, bc2 = _inputs(cfg) + g_param, g_m, g_v = _fp64_golden( + param, m, v, grad, cfg.weight_decay, bc1, bc2 + ) + # Op mutates in place; clone so the golden saw the originals. + out_param, out_m, out_v = torch.ops.et_vk.adamw_step( + param.clone(), + m.clone(), + v.clone(), + grad, + LR, + BETA1, + BETA2, + EPS, + cfg.weight_decay, + bc1, + bc2, + ) + torch.testing.assert_close(out_param, g_param, atol=5e-4, rtol=1e-3) + torch.testing.assert_close(out_m, g_m, atol=5e-4, rtol=1e-3) + torch.testing.assert_close(out_v, g_v, atol=5e-4, rtol=1e-3) + + +if __name__ == "__main__": + unittest.main() From a968838a06c1cb1c2f7dacecf9a230c3e2fe6f11 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 21 Jul 2026 00:29:34 -0700 Subject: [PATCH 05/15] [ExecuTorch][Vulkan] Add et_vk.adamw_step kernel (on-GPU AdamW optimizer step) Pull Request resolved: https://github.com/pytorch/executorch/pull/20937 Vulkan GLSL kernel + handler for the training custom op `et_vk.adamw_step` (elementwise in-place AdamW parameter update). Mirrors `impl/BinaryOp.cpp` elementwise structure + `unary_op.glsl` push-constant scalars; param/m/v are read-modify-write (`kReadWrite`). AOT registration already exists under the shared Vulkan partitioner. ghstack-source-id: 405026087 @exported-using-ghexport Differential Revision: [D111761777](https://our.internmc.facebook.com/intern/diff/D111761777/) --- .../runtime/graph/ops/glsl/adamw_step.glsl | 54 ++++++++ .../runtime/graph/ops/glsl/adamw_step.yaml | 17 +++ .../runtime/graph/ops/impl/AdamwStep.cpp | 126 ++++++++++++++++++ 3 files changed, 197 insertions(+) create mode 100644 backends/vulkan/runtime/graph/ops/glsl/adamw_step.glsl create mode 100644 backends/vulkan/runtime/graph/ops/glsl/adamw_step.yaml create mode 100644 backends/vulkan/runtime/graph/ops/impl/AdamwStep.cpp diff --git a/backends/vulkan/runtime/graph/ops/glsl/adamw_step.glsl b/backends/vulkan/runtime/graph/ops/glsl/adamw_step.glsl new file mode 100644 index 00000000000..ccb7b34bcb7 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/adamw_step.glsl @@ -0,0 +1,54 @@ +/* + * 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. + */ + +#version 450 core + +#define PRECISION ${PRECISION} + +#define T ${buffer_scalar_type(DTYPE)} + +${define_active_storage_type(STORAGE)} + +${define_required_extensions(STORAGE, DTYPE)} + +layout(std430) buffer; + +${layout_declare_tensor(B, "rw", "t_param", DTYPE, STORAGE)} +${layout_declare_tensor(B, "rw", "t_m", DTYPE, STORAGE)} +${layout_declare_tensor(B, "rw", "t_v", DTYPE, STORAGE)} +${layout_declare_tensor(B, "r", "t_grad", DTYPE, STORAGE)} + +layout(push_constant) uniform restrict Block { + int numel; + float lr; + float beta1; + float beta2; + float eps; + float weight_decay; + float bias_correction1; + float bias_correction2; +}; + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +void main() { + const int i = int(gl_GlobalInvocationID.x); + if (i >= numel) { + return; + } + T g = t_grad[i]; + T p = t_param[i]; + p = p - lr * weight_decay * p; + T m = beta1 * t_m[i] + (1.0 - beta1) * g; + T v = beta2 * t_v[i] + (1.0 - beta2) * g * g; + t_m[i] = m; + t_v[i] = v; + T mhat = m / bias_correction1; + T vhat = v / bias_correction2; + t_param[i] = p - lr * mhat / (sqrt(vhat) + eps); +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/adamw_step.yaml b/backends/vulkan/runtime/graph/ops/glsl/adamw_step.yaml new file mode 100644 index 00000000000..114a257cfad --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/adamw_step.yaml @@ -0,0 +1,17 @@ +# 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. + +adamw_step: + parameter_names_with_default_values: + DTYPE: float + STORAGE: buffer + generate_variant_forall: + STORAGE: + - VALUE: buffer + DTYPE: + - VALUE: float + shader_variants: + - NAME: adamw_step diff --git a/backends/vulkan/runtime/graph/ops/impl/AdamwStep.cpp b/backends/vulkan/runtime/graph/ops/impl/AdamwStep.cpp new file mode 100644 index 00000000000..489e95a773d --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/AdamwStep.cpp @@ -0,0 +1,126 @@ +/* + * 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 + +namespace vkcompute { + +void check_adamw_step_args( + ComputeGraph& graph, + const ValueRef param, + const ValueRef m, + const ValueRef v, + const ValueRef grad) { + VK_CHECK_COND(graph.dtype_of(param) == vkapi::kFloat); + VK_CHECK_COND(graph.dtype_of(m) == vkapi::kFloat); + VK_CHECK_COND(graph.dtype_of(v) == vkapi::kFloat); + VK_CHECK_COND(graph.dtype_of(grad) == vkapi::kFloat); + + const int32_t numel = graph.numel_of(param); + VK_CHECK_COND(graph.numel_of(m) == numel); + VK_CHECK_COND(graph.numel_of(v) == numel); + VK_CHECK_COND(graph.numel_of(grad) == numel); +} + +void add_adamw_step_node( + ComputeGraph& graph, + const ValueRef param, + const ValueRef m, + const ValueRef v, + const ValueRef grad, + const ValueRef lr, + const ValueRef beta1, + const ValueRef beta2, + const ValueRef eps, + const ValueRef weight_decay, + const ValueRef bias_correction1, + const ValueRef bias_correction2) { + check_adamw_step_args(graph, param, m, v, grad); + + const float bc1 = graph.extract_scalar(bias_correction1); + const float bc2 = graph.extract_scalar(bias_correction2); + VK_CHECK_COND(bc1 != 0.0f); + VK_CHECK_COND(bc2 != 0.0f); + + // Split into <=16-byte push-constant entries (PushConstantData.h limit). + const std::array pc0 = { + graph.extract_scalar(lr), + graph.extract_scalar(beta1), + graph.extract_scalar(beta2), + graph.extract_scalar(eps)}; + const std::array pc1 = { + graph.extract_scalar(weight_decay), bc1, bc2}; + + std::string kernel_name("adamw_step"); + add_storage_type_suffix(kernel_name, graph.storage_type_of(param)); + add_dtype_suffix(kernel_name, graph.dtype_of(param)); + + graph.execute_nodes().emplace_back(new DynamicDispatchNode( + graph, + VK_KERNEL_FROM_STR(kernel_name), + default_pick_global_wg_size, + default_pick_local_wg_size, + // Inputs and Outputs + {{param, vkapi::kReadWrite}, + {m, vkapi::kReadWrite}, + {v, vkapi::kReadWrite}, + {grad, vkapi::kRead}}, + // Shader params buffers + {}, + // Push Constants + {graph.numel_pc_of(param), + PushConstantDataInfo(pc0.data(), sizeof(pc0)), + PushConstantDataInfo(pc1.data(), sizeof(pc1))}, + // Specialization Constants + {}, + // Resize Args + {}, + // Resizing Logic + nullptr)); +} + +void adamw_step(ComputeGraph& graph, const std::vector& args) { + int arg_idx = 0; + const ValueRef param = args[arg_idx++]; + const ValueRef m = args[arg_idx++]; + const ValueRef v = args[arg_idx++]; + const ValueRef grad = args[arg_idx++]; + const ValueRef lr = args[arg_idx++]; + const ValueRef beta1 = args[arg_idx++]; + const ValueRef beta2 = args[arg_idx++]; + const ValueRef eps = args[arg_idx++]; + const ValueRef weight_decay = args[arg_idx++]; + const ValueRef bias_correction1 = args[arg_idx++]; + const ValueRef bias_correction2 = args[arg_idx++]; + add_adamw_step_node( + graph, + param, + m, + v, + grad, + lr, + beta1, + beta2, + eps, + weight_decay, + bias_correction1, + bias_correction2); +} + +REGISTER_OPERATORS { + VK_REGISTER_OP(et_vk.adamw_step.default, adamw_step); +} + +} // namespace vkcompute From ca19549683e3a29c0cba234276950a1e6b5dbca6 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 21 Jul 2026 00:29:35 -0700 Subject: [PATCH 06/15] [ExecuTorch][Vulkan] Tests for et_vk.adamw_step Pull Request resolved: https://github.com/pytorch/executorch/pull/20938 Vulkan op-test golden for `et_vk.adamw_step` (ATen library-computed reference of the CPU-eager update; in-place param/m/v vs the kernel). 3 cases incl. non-zero weight-decay + a large 1-D tail. Wires `adamw_step_test` in `targets.bzl` + `CMakeLists.txt`. ghstack-source-id: 405030805 @exported-using-ghexport Differential Revision: [D111761775](https://our.internmc.facebook.com/intern/diff/D111761775/) --- backends/vulkan/test/op_tests/CMakeLists.txt | 3 + .../vulkan/test/op_tests/adamw_step_test.cpp | 191 ++++++++++++++++++ backends/vulkan/test/op_tests/targets.bzl | 6 + 3 files changed, 200 insertions(+) create mode 100644 backends/vulkan/test/op_tests/adamw_step_test.cpp diff --git a/backends/vulkan/test/op_tests/CMakeLists.txt b/backends/vulkan/test/op_tests/CMakeLists.txt index 5e8991f8e50..47170373185 100644 --- a/backends/vulkan/test/op_tests/CMakeLists.txt +++ b/backends/vulkan/test/op_tests/CMakeLists.txt @@ -116,6 +116,9 @@ if(TARGET vulkan_backend AND LIB_TORCH) quantized_linear_test ${CMAKE_CURRENT_SOURCE_DIR}/quantized_linear_test.cpp test_utils ) + vulkan_op_test( + adamw_step_test ${CMAKE_CURRENT_SOURCE_DIR}/adamw_step_test.cpp test_utils + ) # Only build generated op tests if a path to tags.yaml and # native_functions.yaml is provided. These files are required for codegen. diff --git a/backends/vulkan/test/op_tests/adamw_step_test.cpp b/backends/vulkan/test/op_tests/adamw_step_test.cpp new file mode 100644 index 00000000000..e2db8592712 --- /dev/null +++ b/backends/vulkan/test/op_tests/adamw_step_test.cpp @@ -0,0 +1,191 @@ +/* + * 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 "test_utils.h" + +#include +#include + +// +// Reference implementation: an ATen transcription of the CPU-eager +// adamw_step_impl golden in backends/vulkan/custom_ops_lib.py. +// + +std::tuple adamw_reference_impl( + const at::Tensor& param_in, + const at::Tensor& m_in, + const at::Tensor& v_in, + const at::Tensor& grad, + const double lr, + const double beta1, + const double beta2, + const double eps, + const double weight_decay, + const double bias_correction1, + const double bias_correction2) { + at::Tensor param = param_in.clone(); + at::Tensor m = m_in.clone(); + at::Tensor v = v_in.clone(); + + param.mul_(1.0 - lr * weight_decay); + m.mul_(beta1).add_(grad, 1.0 - beta1); + v.mul_(beta2).addcmul_(grad, grad, 1.0 - beta2); + at::Tensor mhat = m / bias_correction1; + at::Tensor denom = (v / bias_correction2).sqrt() + eps; + param.addcdiv_(mhat, denom, -lr); + + return std::make_tuple(param, m, v); +} + +// +// Test harness +// + +void test_vulkan_adamw_step( + const std::vector& sizes, + const double lr, + const double beta1, + const double beta2, + const double eps, + const double weight_decay, + const int step) { + at::manual_seed(0); + + const double bias_correction1 = 1.0 - std::pow(beta1, step); + const double bias_correction2 = 1.0 - std::pow(beta2, step); + + const auto options = at::device(at::kCPU).dtype(at::kFloat); + at::Tensor param = at::randn(sizes, options); + at::Tensor m = at::randn(sizes, options); + at::Tensor v = at::rand(sizes, options); // second moment is non-negative + at::Tensor grad = at::randn(sizes, options); + + at::Tensor ref_param; + at::Tensor ref_m; + at::Tensor ref_v; + std::tie(ref_param, ref_m, ref_v) = adamw_reference_impl( + param, + m, + v, + grad, + lr, + beta1, + beta2, + eps, + weight_decay, + bias_correction1, + bias_correction2); + + using namespace vkcompute; + + GraphConfig config; + ComputeGraph graph(config); + + ValueRef r_param = graph.add_tensor(sizes, vkapi::kFloat, utils::kBuffer); + ValueRef r_m = graph.add_tensor(sizes, vkapi::kFloat, utils::kBuffer); + ValueRef r_v = graph.add_tensor(sizes, vkapi::kFloat, utils::kBuffer); + ValueRef r_grad = graph.add_tensor(sizes, vkapi::kFloat, utils::kBuffer); + + ValueRef staging_param = graph.set_input_tensor(r_param); + ValueRef staging_m = graph.set_input_tensor(r_m); + ValueRef staging_v = graph.set_input_tensor(r_v); + ValueRef staging_grad = graph.set_input_tensor(r_grad); + + VK_GET_OP_FN("et_vk.adamw_step.default") + (graph, + { + r_param, + r_m, + r_v, + r_grad, + graph.add_scalar(lr), + graph.add_scalar(beta1), + graph.add_scalar(beta2), + graph.add_scalar(eps), + graph.add_scalar(weight_decay), + graph.add_scalar(bias_correction1), + graph.add_scalar(bias_correction2), + }); + + ValueRef staging_param_out = graph.set_output_tensor(r_param); + ValueRef staging_m_out = graph.set_output_tensor(r_m); + ValueRef staging_v_out = graph.set_output_tensor(r_v); + + graph.prepare(); + graph.prepack(); + + graph.maybe_cast_and_copy_into_staging( + staging_param, param.const_data_ptr(), param.numel(), vkapi::kFloat); + graph.maybe_cast_and_copy_into_staging( + staging_m, m.const_data_ptr(), m.numel(), vkapi::kFloat); + graph.maybe_cast_and_copy_into_staging( + staging_v, v.const_data_ptr(), v.numel(), vkapi::kFloat); + graph.maybe_cast_and_copy_into_staging( + staging_grad, grad.const_data_ptr(), grad.numel(), vkapi::kFloat); + + graph.execute(); + + at::Tensor vk_param = at::empty_like(param); + at::Tensor vk_m = at::empty_like(m); + at::Tensor vk_v = at::empty_like(v); + + graph.maybe_cast_and_copy_from_staging( + staging_param_out, + vk_param.mutable_data_ptr(), + vk_param.numel(), + vkapi::kFloat); + graph.maybe_cast_and_copy_from_staging( + staging_m_out, vk_m.mutable_data_ptr(), vk_m.numel(), vkapi::kFloat); + graph.maybe_cast_and_copy_from_staging( + staging_v_out, vk_v.mutable_data_ptr(), vk_v.numel(), vkapi::kFloat); + + ASSERT_TRUE(at::allclose(ref_param, vk_param, /*rtol=*/1e-4, /*atol=*/1e-4)); + ASSERT_TRUE(at::allclose(ref_m, vk_m, /*rtol=*/1e-4, /*atol=*/1e-4)); + ASSERT_TRUE(at::allclose(ref_v, vk_v, /*rtol=*/1e-4, /*atol=*/1e-4)); +} + +TEST(VulkanAdamwStepTest, test_adamw_step_no_weight_decay) { + test_vulkan_adamw_step( + /*sizes=*/{4, 8}, + /*lr=*/1e-3, + /*beta1=*/0.9, + /*beta2=*/0.999, + /*eps=*/1e-8, + /*weight_decay=*/0.0, + /*step=*/1); +} + +TEST(VulkanAdamwStepTest, test_adamw_step_with_weight_decay) { + test_vulkan_adamw_step( + /*sizes=*/{3, 5, 7}, + /*lr=*/1e-2, + /*beta1=*/0.9, + /*beta2=*/0.999, + /*eps=*/1e-8, + /*weight_decay=*/0.01, + /*step=*/10); +} + +TEST(VulkanAdamwStepTest, test_adamw_step_large_1d) { + test_vulkan_adamw_step( + /*sizes=*/{1000}, + /*lr=*/1e-3, + /*beta1=*/0.9, + /*beta2=*/0.999, + /*eps=*/1e-8, + /*weight_decay=*/0.05, + /*step=*/5); +} diff --git a/backends/vulkan/test/op_tests/targets.bzl b/backends/vulkan/test/op_tests/targets.bzl index 0a9b0743415..8aab725417a 100644 --- a/backends/vulkan/test/op_tests/targets.bzl +++ b/backends/vulkan/test/op_tests/targets.bzl @@ -180,6 +180,12 @@ def define_common_targets(is_fbcode = False): ":test_utils", ] ) + define_test_targets( + "adamw_step_test", + extra_deps = [ + ":test_utils", + ] + ) define_test_targets( "rms_norm_test", extra_deps = [ From bd2695f4aee793230a29941dc67918f01118966c Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 21 Jul 2026 00:29:35 -0700 Subject: [PATCH 07/15] [ExecuTorch][Vulkan] Add et_vk.linear_dW kernel (dense fp32 weight-grad GEMM) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/20939 Vulkan GLSL kernel + handler for the training custom op `et_vk.linear_q4gsw_dw` — a dense fp32 GEMM `dW[N,K] = d_out^T @ x` (contracts over M; NOT quantized despite the name). Mirrors `impl/Matmul.cpp` `add_matmul_tiled_node` structure + `sizes_ubo` for dynamic shapes; the M-contraction tile math is transcribed from the golden-verified reference. AOT registration already exists under the shared Vulkan partitioner. ghstack-source-id: 405026089 @exported-using-ghexport Differential Revision: [D111761779](https://our.internmc.facebook.com/intern/diff/D111761779/) --- .../runtime/graph/ops/glsl/linear_dW.glsl | 78 ++++++++++++ .../runtime/graph/ops/glsl/linear_dW.yaml | 17 +++ .../runtime/graph/ops/impl/LinearDw.cpp | 112 ++++++++++++++++++ 3 files changed, 207 insertions(+) create mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dW.glsl create mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_dW.yaml create mode 100644 backends/vulkan/runtime/graph/ops/impl/LinearDw.cpp diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dW.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_dW.glsl new file mode 100644 index 00000000000..1a8ed94031b --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dW.glsl @@ -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. + */ + +#version 450 core + +#define PRECISION ${PRECISION} + +#define T ${texel_load_component_type(DTYPE, STORAGE)} + +#define TILE_N 4 +#define TILE_K 4 + +${define_required_extensions(STORAGE, DTYPE)} + +layout(std430) buffer; + +${layout_declare_tensor(B, "w", "t_dw", DTYPE, STORAGE, is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_dout", DTYPE, STORAGE, is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_x", DTYPE, STORAGE, is_scalar_array=True)} + +${layout_declare_ubo(B, "ivec4", "dout_sizes")} +${layout_declare_ubo(B, "ivec4", "x_sizes")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +void main() { + // dW[N, K] = sum_m d_out[m, N] * x[m, K]; contraction over the flattened M. + const int N = dout_sizes.x; + const int M = dout_sizes.y * dout_sizes.z * dout_sizes.w; + const int K = x_sizes.x; + + const int nnt = (N + TILE_N - 1) / TILE_N; + const int nkt = (K + TILE_K - 1) / TILE_K; + const int tiles = nnt * nkt; + + const int tile_idx = int(gl_GlobalInvocationID.x); + if (tile_idx >= tiles) { + return; + } + + const int n0 = (tile_idx / nkt) * TILE_N; + const int k0 = (tile_idx % nkt) * TILE_K; + + float acc[TILE_N * TILE_K]; + for (int i = 0; i < TILE_N * TILE_K; ++i) { + acc[i] = 0.0; + } + + for (int m = 0; m < M; ++m) { + float dout_reg[TILE_N]; + for (int nl = 0; nl < TILE_N; ++nl) { + const int n_eff = min(n0 + nl, N - 1); + dout_reg[nl] = float(t_dout[m * N + n_eff]); + } + for (int kl = 0; kl < TILE_K; ++kl) { + const int k_eff = min(k0 + kl, K - 1); + const float xv = float(t_x[m * K + k_eff]); + for (int nl = 0; nl < TILE_N; ++nl) { + acc[nl * TILE_K + kl] += dout_reg[nl] * xv; + } + } + } + + for (int nl = 0; nl < TILE_N; ++nl) { + const int n = n0 + nl; + for (int kl = 0; kl < TILE_K; ++kl) { + const int k = k0 + kl; + if (n < N && k < K) { + t_dw[n * K + k] = T(acc[nl * TILE_K + kl]); + } + } + } +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dW.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_dW.yaml new file mode 100644 index 00000000000..a2d38101f3d --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dW.yaml @@ -0,0 +1,17 @@ +# 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. + +linear_dW: + parameter_names_with_default_values: + DTYPE: float + STORAGE: buffer + generate_variant_forall: + STORAGE: + - VALUE: buffer + DTYPE: + - VALUE: float + shader_variants: + - NAME: linear_dW diff --git a/backends/vulkan/runtime/graph/ops/impl/LinearDw.cpp b/backends/vulkan/runtime/graph/ops/impl/LinearDw.cpp new file mode 100644 index 00000000000..1ab5e6976ca --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/LinearDw.cpp @@ -0,0 +1,112 @@ +/* + * 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 + +namespace vkcompute { + +void resize_linear_dW_node( + ComputeGraph* graph, + const std::vector& args, + const std::vector& resize_args) { + (void)resize_args; + const ValueRef dW = args.at(0).refs.at(0); + const ValueRef d_out = args.at(1).refs.at(0); + const ValueRef x = args.at(1).refs.at(1); + const int64_t N = graph->size_at(-1, d_out); + const int64_t K = graph->size_at(-1, x); + graph->virtual_resize(dW, {N, K}); +} + +utils::uvec3 linear_dW_global_wg_size( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const std::vector& args, + const std::vector& resize_args) { + (void)shader; + (void)resize_args; + const ValueRef dW = args.at(0).refs.at(0); + const uint32_t N = graph->size_at(-2, dW); + const uint32_t K = graph->size_at(-1, dW); + const uint32_t tiles = utils::div_up_4(N) * utils::div_up_4(K); + return {tiles, 1u, 1u}; +} + +utils::uvec3 linear_dW_local_wg_size( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const utils::uvec3& global_workgroup_size, + const std::vector& args, + const std::vector& resize_args) { + (void)graph; + (void)shader; + (void)global_workgroup_size; + (void)args; + (void)resize_args; + return {64u, 1u, 1u}; +} + +void linear_dW(ComputeGraph& graph, const std::vector& args) { + int32_t i = 0; + const ValueRef d_out = args.at(i++); + const ValueRef x = args.at(i++); + const ValueRef dW = args.at(i); + + VK_CHECK_COND(graph.dtype_of(d_out) == vkapi::kFloat); + VK_CHECK_COND(graph.dtype_of(x) == vkapi::kFloat); + VK_CHECK_COND(graph.dtype_of(dW) == vkapi::kFloat); + VK_CHECK_COND(graph.is_buffer_storage(d_out)); + VK_CHECK_COND(graph.is_buffer_storage(x)); + VK_CHECK_COND(graph.is_buffer_storage(dW)); + + const int64_t N = graph.size_at(-1, d_out); + const int64_t K = graph.size_at(-1, x); + VK_CHECK_COND(N > 0 && K > 0); + VK_CHECK_COND(graph.numel_of(d_out) % N == 0); + VK_CHECK_COND(graph.numel_of(x) % K == 0); + VK_CHECK_COND(graph.numel_of(d_out) / N == graph.numel_of(x) / K); + + const uint32_t tiles = utils::div_up_4(static_cast(N)) * + utils::div_up_4(static_cast(K)); + VK_CHECK_COND( + (tiles + 63u) / 64u <= 65535u, + "linear_dW: tile count exceeds max workgroup count"); + + std::string kernel_name = "linear_dW"; + kernel_name.reserve(kShaderNameReserve); + add_storage_type_suffix(kernel_name, graph.storage_type_of(dW)); + add_dtype_suffix(kernel_name, graph.dtype_of(dW)); + + graph.execute_nodes().emplace_back(new DynamicDispatchNode( + graph, + VK_KERNEL_FROM_STR(kernel_name), + linear_dW_global_wg_size, + linear_dW_local_wg_size, + // Inputs and Outputs + {{dW, vkapi::kWrite}, {{d_out, x}, vkapi::kRead}}, + // Shader params buffers + {graph.sizes_ubo(d_out), graph.sizes_ubo(x)}, + // Push Constants + {}, + // Specialization Constants + {}, + // Resize Args + {}, + // Resizing Logic + resize_linear_dW_node)); +} + +REGISTER_OPERATORS { + VK_REGISTER_OP(et_vk.linear_dW.default, linear_dW); +} + +} // namespace vkcompute From 27f60688114915bc4ff3b30a2dd2e960ba876042 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 21 Jul 2026 00:29:35 -0700 Subject: [PATCH 08/15] [ExecuTorch][Vulkan] Tests for et_vk.linear_dW Pull Request resolved: https://github.com/pytorch/executorch/pull/20940 Vulkan op-test golden for `et_vk.linear_q4gsw_dw` (ATen `d_out.reshape(-1,N).t() @ x.reshape(-1,K)` reference). 3 cases: tile-aligned, non-tile-multiple (partial-tile clamp), and >2D leading-dim flatten. Wires `linear_q4gsw_dw_test` in `targets.bzl` + `CMakeLists.txt`. ghstack-source-id: 405026090 @exported-using-ghexport Differential Revision: [D111761776](https://our.internmc.facebook.com/intern/diff/D111761776/) --- backends/vulkan/test/op_tests/CMakeLists.txt | 3 + .../vulkan/test/op_tests/linear_dW_test.cpp | 110 ++++++++++++++++++ backends/vulkan/test/op_tests/targets.bzl | 6 + 3 files changed, 119 insertions(+) create mode 100644 backends/vulkan/test/op_tests/linear_dW_test.cpp diff --git a/backends/vulkan/test/op_tests/CMakeLists.txt b/backends/vulkan/test/op_tests/CMakeLists.txt index 47170373185..4a39c860d33 100644 --- a/backends/vulkan/test/op_tests/CMakeLists.txt +++ b/backends/vulkan/test/op_tests/CMakeLists.txt @@ -119,6 +119,9 @@ if(TARGET vulkan_backend AND LIB_TORCH) vulkan_op_test( adamw_step_test ${CMAKE_CURRENT_SOURCE_DIR}/adamw_step_test.cpp test_utils ) + vulkan_op_test( + linear_dW_test ${CMAKE_CURRENT_SOURCE_DIR}/linear_dW_test.cpp test_utils + ) # Only build generated op tests if a path to tags.yaml and # native_functions.yaml is provided. These files are required for codegen. diff --git a/backends/vulkan/test/op_tests/linear_dW_test.cpp b/backends/vulkan/test/op_tests/linear_dW_test.cpp new file mode 100644 index 00000000000..1d9f31328ed --- /dev/null +++ b/backends/vulkan/test/op_tests/linear_dW_test.cpp @@ -0,0 +1,110 @@ +/* + * 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 "test_utils.h" + +// +// Reference Implementation +// + +// Golden: dW[N, K] = d_out^T @ x, contracting over the flattened leading dims. +// Mirrors the CPU-eager linear_dW_impl in custom_ops_lib.py. +at::Tensor linear_dW_reference_impl( + const at::Tensor& d_out, + const at::Tensor& x) { + const int64_t N = d_out.size(-1); + const int64_t K = x.size(-1); + return d_out.reshape({-1, N}).t().matmul(x.reshape({-1, K})).contiguous(); +} + +// +// Test function +// + +void test_vulkan_linear_dW_impl( + const std::vector& d_out_sizes, + const std::vector& x_sizes, + const vkcompute::utils::StorageType storage = vkcompute::utils::kBuffer) { + at::Tensor d_out = + at::rand(d_out_sizes, at::device(at::kCPU).dtype(at::kFloat)); + at::Tensor x = at::rand(x_sizes, at::device(at::kCPU).dtype(at::kFloat)); + + at::Tensor dW_ref = linear_dW_reference_impl(d_out, x); + + // Build Vulkan graph + using namespace vkcompute; + + GraphConfig config; + ComputeGraph graph(config); + + IOValueRef r_d_out = graph.add_input_tensor( + d_out.sizes().vec(), from_at_scalartype(d_out.scalar_type()), storage); + IOValueRef r_x = graph.add_input_tensor( + x.sizes().vec(), from_at_scalartype(x.scalar_type()), storage); + + const ValueRef r_dW = graph.add_tensor( + dW_ref.sizes().vec(), from_at_scalartype(dW_ref.scalar_type()), storage); + + VK_GET_OP_FN("et_vk.linear_dW.default") + (graph, {r_d_out.value, r_x.value, r_dW}); + + ValueRef staging_out = graph.set_output_tensor(r_dW); + + graph.prepare(); + graph.prepack(); + graph.propagate_resize(); + + graph.maybe_cast_and_copy_into_staging( + r_d_out.staging, + d_out.const_data_ptr(), + d_out.numel(), + from_at_scalartype(d_out.scalar_type())); + graph.maybe_cast_and_copy_into_staging( + r_x.staging, + x.const_data_ptr(), + x.numel(), + from_at_scalartype(x.scalar_type())); + + graph.execute(); + + at::Tensor vk_dW = at::empty_like(dW_ref); + graph.maybe_cast_and_copy_from_staging( + staging_out, + vk_dW.mutable_data_ptr(), + vk_dW.numel(), + from_at_scalartype(vk_dW.scalar_type())); + + ASSERT_TRUE(at::allclose(vk_dW, dW_ref, 1e-3, 1e-3)); +} + +// Tile-aligned 2D shapes (M, N, K all multiples of 4). +TEST(VulkanLinearDwTest, test_tile_aligned) { + test_vulkan_linear_dW_impl( + /*d_out_sizes=*/{8, 16}, /*x_sizes=*/{8, 32}); +} + +// Non-tile-multiple shapes (M, N, K each not a multiple of 4) to exercise the +// partial-tile min()-clamp paths in the shader. +TEST(VulkanLinearDwTest, test_non_tile_multiple) { + test_vulkan_linear_dW_impl( + /*d_out_sizes=*/{5, 6}, /*x_sizes=*/{5, 10}); +} + +// Leading dims > 2D: M is the flattened product of all leading dims. +TEST(VulkanLinearDwTest, test_leading_dims_flatten) { + test_vulkan_linear_dW_impl( + /*d_out_sizes=*/{2, 3, 16}, /*x_sizes=*/{2, 3, 32}); +} diff --git a/backends/vulkan/test/op_tests/targets.bzl b/backends/vulkan/test/op_tests/targets.bzl index 8aab725417a..f3e6d42bc54 100644 --- a/backends/vulkan/test/op_tests/targets.bzl +++ b/backends/vulkan/test/op_tests/targets.bzl @@ -186,6 +186,12 @@ def define_common_targets(is_fbcode = False): ":test_utils", ] ) + define_test_targets( + "linear_dW_test", + extra_deps = [ + ":test_utils", + ] + ) define_test_targets( "rms_norm_test", extra_deps = [ From 63fd90be39002f46242425f05a8152040ac41eb4 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 21 Jul 2026 00:29:36 -0700 Subject: [PATCH 09/15] [ExecuTorch][Vulkan] Add et_vk.fused_ce kernels (fused cross-entropy loss + dlogits) Pull Request resolved: https://github.com/pytorch/executorch/pull/20941 Vulkan GLSL kernels + handler for the training custom op `et_vk.fused_ce` (fused CE: per-row online-softmax loss + dlogits, then an N->1 loss sum). Per-row shared-memory reduction structured like `glsl/reduce_per_row_buffer.glsl` (`NWORKERS=64`, one workgroup/row, barrier tree-combine); two dispatch nodes share `loss_partial`, ordered by the runtime per-`vTensor` barrier (same pattern as `impl/SDPA.cpp` QK->softmax->AV). AOT registration already exists under the shared Vulkan partitioner. ghstack-source-id: 405059103 @exported-using-ghexport Differential Revision: [D111761780](https://our.internmc.facebook.com/intern/diff/D111761780/) --- .../runtime/graph/ops/glsl/fused_ce.glsl | 110 ++++++++++ .../runtime/graph/ops/glsl/fused_ce.yaml | 15 ++ .../runtime/graph/ops/glsl/fused_ce_sum.glsl | 55 +++++ .../runtime/graph/ops/glsl/fused_ce_sum.yaml | 15 ++ .../vulkan/runtime/graph/ops/impl/FusedCe.cpp | 197 ++++++++++++++++++ 5 files changed, 392 insertions(+) create mode 100644 backends/vulkan/runtime/graph/ops/glsl/fused_ce.glsl create mode 100644 backends/vulkan/runtime/graph/ops/glsl/fused_ce.yaml create mode 100644 backends/vulkan/runtime/graph/ops/glsl/fused_ce_sum.glsl create mode 100644 backends/vulkan/runtime/graph/ops/glsl/fused_ce_sum.yaml create mode 100644 backends/vulkan/runtime/graph/ops/impl/FusedCe.cpp diff --git a/backends/vulkan/runtime/graph/ops/glsl/fused_ce.glsl b/backends/vulkan/runtime/graph/ops/glsl/fused_ce.glsl new file mode 100644 index 00000000000..044301d5330 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/fused_ce.glsl @@ -0,0 +1,110 @@ +/* + * 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. + */ + +#version 450 core + +${define_required_extensions(STORAGE, DTYPE)} + +#define PRECISION ${PRECISION} +#define T ${buffer_scalar_type(DTYPE)} + +${define_active_storage_type(STORAGE)} + +layout(std430) buffer; + +${layout_declare_tensor(B, "w", "t_dlogits", DTYPE, "buffer")} +${layout_declare_tensor(B, "w", "t_loss_partial", DTYPE, "buffer")} +${layout_declare_tensor(B, "r", "t_logits", DTYPE, "buffer")} +${layout_declare_tensor(B, "r", "t_labels", "int", "buffer")} + +${layout_declare_ubo(B, "int", "vocab")} +${layout_declare_ubo(B, "int", "n_rows")} +${layout_declare_ubo(B, "float", "n_valid")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +#define NWORKERS 64 + +shared float red_m[NWORKERS]; +shared float red_l[NWORKERS]; + +// Fused cross-entropy: one workgroup per row cooperatively reduces the vocab +// dimension with a single-pass online softmax (running max + rescaled running +// sum), then writes per-row loss and dlogits. Rows with label < 0 are masked. +void main() { + const uint row = gl_GlobalInvocationID.y; + if (int(row) >= n_rows) { + return; + } + + const uint tid = gl_LocalInvocationID.x; + const uint V = uint(vocab); + const uint base = row * V; + const int lbl = t_labels[row]; + + if (lbl < 0) { + for (uint j = tid; j < V; j += NWORKERS) { + t_dlogits[base + j] = T(0); + } + if (tid == 0u) { + t_loss_partial[row] = T(0); + } + return; + } + + // Single read pass: maintain a running max m and the sum l of exp(x - m), + // rescaling l whenever a larger value updates m. Finite -3.4e38 init. + float m = -3.4e38; + float l = 0.0; + for (uint j = tid; j < V; j += NWORKERS) { + float x = float(t_logits[base + j]); + if (x > m) { + l = l * exp(m - x) + 1.0; + m = x; + } else { + l = l + exp(x - m); + } + } + red_m[tid] = m; + red_l[tid] = l; + memoryBarrierShared(); + barrier(); + + // Tree-combine the (m, l) pairs: m becomes the max, l is rescaled to it. + for (uint s = NWORKERS / 2u; s > 0u; s >>= 1u) { + if (tid < s) { + float ma = red_m[tid]; + float la = red_l[tid]; + float mb = red_m[tid + s]; + float lb = red_l[tid + s]; + float mm = max(ma, mb); + red_m[tid] = mm; + red_l[tid] = la * exp(ma - mm) + lb * exp(mb - mm); + } + memoryBarrierShared(); + barrier(); + } + + const float row_max = red_m[0]; + const float denom = red_l[0]; + const float inv = 1.0 / denom; + const float scale = 1.0 / n_valid; + + if (tid == 0u) { + const float lse = row_max + log(denom); + t_loss_partial[row] = T((lse - float(t_logits[base + uint(lbl)])) * scale); + } + + for (uint j = tid; j < V; j += NWORKERS) { + float g = exp(float(t_logits[base + j]) - row_max) * inv * scale; + if (j == uint(lbl)) { + g = g - scale; + } + t_dlogits[base + j] = T(g); + } +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/fused_ce.yaml b/backends/vulkan/runtime/graph/ops/glsl/fused_ce.yaml new file mode 100644 index 00000000000..8cf9cdac04e --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/fused_ce.yaml @@ -0,0 +1,15 @@ +# 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. + +fused_ce: + parameter_names_with_default_values: + DTYPE: float + STORAGE: buffer + generate_variant_forall: + DTYPE: + - VALUE: float + shader_variants: + - NAME: fused_ce_buffer diff --git a/backends/vulkan/runtime/graph/ops/glsl/fused_ce_sum.glsl b/backends/vulkan/runtime/graph/ops/glsl/fused_ce_sum.glsl new file mode 100644 index 00000000000..c61820e074c --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/fused_ce_sum.glsl @@ -0,0 +1,55 @@ +/* + * 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. + */ + +#version 450 core + +${define_required_extensions(STORAGE, DTYPE)} + +#define PRECISION ${PRECISION} +#define T ${buffer_scalar_type(DTYPE)} + +${define_active_storage_type(STORAGE)} + +layout(std430) buffer; + +${layout_declare_tensor(B, "w", "t_loss", DTYPE, "buffer")} +${layout_declare_tensor(B, "r", "t_loss_partial", DTYPE, "buffer")} + +${layout_declare_ubo(B, "int", "n_rows")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +#define NWORKERS 64 + +shared float red[NWORKERS]; + +// Self-contained [N] -> [1] tree-sum of the per-row losses in one workgroup, so +// fused_ce carries no cross-op reduce dependency. +void main() { + const uint tid = gl_LocalInvocationID.x; + + float s = 0.0; + for (uint j = tid; j < uint(n_rows); j += NWORKERS) { + s += float(t_loss_partial[j]); + } + red[tid] = s; + memoryBarrierShared(); + barrier(); + + for (uint k = NWORKERS / 2u; k > 0u; k >>= 1u) { + if (tid < k) { + red[tid] += red[tid + k]; + } + memoryBarrierShared(); + barrier(); + } + + if (tid == 0u) { + t_loss[0] = T(red[0]); + } +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/fused_ce_sum.yaml b/backends/vulkan/runtime/graph/ops/glsl/fused_ce_sum.yaml new file mode 100644 index 00000000000..cc1bd16cebb --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/fused_ce_sum.yaml @@ -0,0 +1,15 @@ +# 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. + +fused_ce_sum: + parameter_names_with_default_values: + DTYPE: float + STORAGE: buffer + generate_variant_forall: + DTYPE: + - VALUE: float + shader_variants: + - NAME: fused_ce_sum_buffer diff --git a/backends/vulkan/runtime/graph/ops/impl/FusedCe.cpp b/backends/vulkan/runtime/graph/ops/impl/FusedCe.cpp new file mode 100644 index 00000000000..f66e0e4fc46 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/FusedCe.cpp @@ -0,0 +1,197 @@ +/* + * 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 + +namespace vkcompute { + +using namespace utils; + +utils::uvec3 fused_ce_global_wg_size( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const std::vector& args, + const std::vector& resize_args) { + (void)shader; + (void)resize_args; + const ValueRef loss_partial = args.at(0).refs.at(1); + return { + 1u, utils::safe_downcast(graph->numel_of(loss_partial)), 1u}; +} + +utils::uvec3 fused_ce_local_wg_size( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const utils::uvec3& global_workgroup_size, + const std::vector& args, + const std::vector& resize_args) { + (void)graph; + (void)shader; + (void)global_workgroup_size; + (void)args; + (void)resize_args; + return {64u, 1u, 1u}; +} + +utils::uvec3 fused_ce_sum_global_wg_size( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const std::vector& args, + const std::vector& resize_args) { + (void)graph; + (void)shader; + (void)args; + (void)resize_args; + return {1u, 1u, 1u}; +} + +utils::uvec3 fused_ce_sum_local_wg_size( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const utils::uvec3& global_workgroup_size, + const std::vector& args, + const std::vector& resize_args) { + (void)graph; + (void)shader; + (void)global_workgroup_size; + (void)args; + (void)resize_args; + return {64u, 1u, 1u}; +} + +void resize_fused_ce_node( + ComputeGraph* graph, + const std::vector& args, + const std::vector& resize_args) { + (void)resize_args; + const ValueRef dlogits = args.at(0).refs.at(0); + const ValueRef loss_partial = args.at(0).refs.at(1); + const ValueRef logits = args.at(1).refs.at(0); + + const std::vector logits_sizes = graph->sizes_of(logits); + graph->virtual_resize(dlogits, logits_sizes); + graph->virtual_resize(loss_partial, {logits_sizes.at(0)}); +} + +void resize_fused_ce_sum_node( + ComputeGraph* graph, + const std::vector& args, + const std::vector& resize_args) { + (void)resize_args; + const ValueRef loss = args.at(0).refs.at(0); + // Rank-0 scalar, matching fused_ce_meta's logits.new_empty([]). + graph->virtual_resize(loss, {}); +} + +void fused_ce(ComputeGraph& graph, const std::vector& args) { + int arg_idx = 0; + const ValueRef logits = args[arg_idx++]; + const ValueRef labels = args[arg_idx++]; + const ValueRef n_valid_ref = args[arg_idx++]; + const ValueRef out_tuple_ref = args[arg_idx]; + + // Read loss/dlogits via short-lived get_value_list() temporaries: no get_*() + // pointer may be alive when loss_partial is added below, else adding Values + // can reallocate values_ and check_no_active_value_ptrs() throws. + const ValueRef loss = graph.get_value_list(out_tuple_ref)->at(0); + const ValueRef dlogits = graph.get_value_list(out_tuple_ref)->at(1); + + VK_CHECK_COND( + graph.is_buffer_storage(logits) && graph.is_buffer_storage(dlogits), + "fused_ce: logits and dlogits must use buffer storage"); + VK_CHECK_COND( + graph.dim_of(logits) == 2, "fused_ce: logits must be 2D [n_rows, vocab]"); + VK_CHECK_COND( + graph.sizes_of(dlogits) == graph.sizes_of(logits), + "fused_ce: dlogits must match logits shape"); + VK_CHECK_COND( + graph.dtype_of(labels) == vkapi::kInt, "fused_ce: labels must be int32"); + VK_CHECK_COND(graph.dim_of(labels) == 1, "fused_ce: labels must be 1D [N]"); + VK_CHECK_COND( + graph.size_at(0, labels) == graph.size_at(0, logits), + "fused_ce: labels length must equal number of rows"); + VK_CHECK_COND(graph.numel_of(loss) == 1, "fused_ce: loss must be a scalar"); + + const int32_t n_rows = graph.size_at(0, logits); + const int32_t vocab = graph.size_at(1, logits); + const float n_valid = graph.extract_scalar(n_valid_ref); + + // One workgroup per row; the Vulkan spec guarantees maxComputeWorkGroupCount + // >= 65535 per dimension. + VK_CHECK_COND( + n_rows <= 65535, "fused_ce: n_rows exceeds max workgroup count"); + + // Per-row loss contributions, reduced to the scalar loss by the sum node. + TmpTensor loss_partial( + &graph, {n_rows}, graph.dtype_of(logits), utils::kBuffer); + + std::string kernel_name = "fused_ce"; + kernel_name.reserve(kShaderNameReserve); + add_storage_type_suffix(kernel_name, graph.storage_type_of(dlogits)); + add_dtype_suffix(kernel_name, graph.dtype_of(dlogits)); + + graph.execute_nodes().emplace_back(new DynamicDispatchNode( + graph, + VK_KERNEL_FROM_STR(kernel_name), + fused_ce_global_wg_size, + fused_ce_local_wg_size, + // Inputs and Outputs + {{{dlogits, loss_partial}, vkapi::kWrite}, + {{logits, labels}, vkapi::kRead}}, + // Shader params buffers + {graph.create_params_buffer(vocab), + graph.create_params_buffer(n_rows), + graph.create_params_buffer(n_valid)}, + // Push Constants + {}, + // Specialization Constants + {}, + // Resize Args + {}, + // Resizing Logic + resize_fused_ce_node)); + + std::string sum_kernel_name = "fused_ce_sum"; + sum_kernel_name.reserve(kShaderNameReserve); + add_storage_type_suffix(sum_kernel_name, graph.storage_type_of(loss)); + add_dtype_suffix(sum_kernel_name, graph.dtype_of(loss)); + + // Separate dispatch node so the loss_partial write (above) is ordered before + // this read: the graph inserts a per-tensor pipeline barrier between them. + graph.execute_nodes().emplace_back(new DynamicDispatchNode( + graph, + VK_KERNEL_FROM_STR(sum_kernel_name), + fused_ce_sum_global_wg_size, + fused_ce_sum_local_wg_size, + // Inputs and Outputs + {{loss, vkapi::kWrite}, {loss_partial, vkapi::kRead}}, + // Shader params buffers + {graph.create_params_buffer(n_rows)}, + // Push Constants + {}, + // Specialization Constants + {}, + // Resize Args + {}, + // Resizing Logic + resize_fused_ce_sum_node)); +} + +REGISTER_OPERATORS { + VK_REGISTER_OP(et_vk.fused_ce.default, fused_ce); +} + +} // namespace vkcompute From 5dcdaa1ea21e398843937ae110be9a07fb30db39 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 21 Jul 2026 00:29:37 -0700 Subject: [PATCH 10/15] [ExecuTorch][Vulkan] Tests for et_vk.fused_ce Pull Request resolved: https://github.com/pytorch/executorch/pull/20942 Vulkan op-test golden for `et_vk.fused_ce` (ATen transcription of the CPU-eager loss/dlogits reference). 3 cases: all-valid, a masked label (`label<0`), and vocab >> workers (strided per-row reduce). Wires `fused_ce_test` in `targets.bzl` + `CMakeLists.txt`. ghstack-source-id: 405059109 @exported-using-ghexport Differential Revision: [D111761774](https://our.internmc.facebook.com/intern/diff/D111761774/) --- backends/vulkan/test/op_tests/CMakeLists.txt | 3 + .../vulkan/test/op_tests/fused_ce_test.cpp | 145 ++++++++++++++++++ backends/vulkan/test/op_tests/targets.bzl | 6 + 3 files changed, 154 insertions(+) create mode 100644 backends/vulkan/test/op_tests/fused_ce_test.cpp diff --git a/backends/vulkan/test/op_tests/CMakeLists.txt b/backends/vulkan/test/op_tests/CMakeLists.txt index 4a39c860d33..048371a35bb 100644 --- a/backends/vulkan/test/op_tests/CMakeLists.txt +++ b/backends/vulkan/test/op_tests/CMakeLists.txt @@ -122,6 +122,9 @@ if(TARGET vulkan_backend AND LIB_TORCH) vulkan_op_test( linear_dW_test ${CMAKE_CURRENT_SOURCE_DIR}/linear_dW_test.cpp test_utils ) + vulkan_op_test( + fused_ce_test ${CMAKE_CURRENT_SOURCE_DIR}/fused_ce_test.cpp test_utils + ) # Only build generated op tests if a path to tags.yaml and # native_functions.yaml is provided. These files are required for codegen. diff --git a/backends/vulkan/test/op_tests/fused_ce_test.cpp b/backends/vulkan/test/op_tests/fused_ce_test.cpp new file mode 100644 index 00000000000..fa2a8ae22bc --- /dev/null +++ b/backends/vulkan/test/op_tests/fused_ce_test.cpp @@ -0,0 +1,145 @@ +/* + * 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 "test_utils.h" + +#include +#include + +// +// Reference implementation: transcription of the CPU-eager fused_ce_impl +// (backends/vulkan/custom_ops_lib.py), computed with ATen (library golden). +// + +std::pair fused_ce_reference( + const at::Tensor& logits, + const at::Tensor& labels, + double n_valid) { + at::Tensor mask = labels.ge(0); + at::Tensor safe = labels.clamp_min(0).to(at::kLong); + at::Tensor lse = at::logsumexp(logits, -1); + at::Tensor picked = logits.gather(-1, safe.unsqueeze(-1)).squeeze(-1); + at::Tensor loss = + at::where(mask, (lse - picked) / n_valid, at::zeros_like(lse)).sum(); + at::Tensor softmax = at::softmax(logits, -1); + at::Tensor onehot = at::one_hot(safe, logits.size(-1)).to(logits.dtype()); + at::Tensor dlogits = at::where( + mask.unsqueeze(-1), + (softmax - onehot) / n_valid, + at::zeros_like(softmax)); + return {loss, dlogits}; +} + +void test_vulkan_fused_ce( + const int n_rows, + const int vocab, + const std::vector& labels_data, + const double n_valid) { + at::manual_seed(0); + + at::Tensor logits = + at::rand({n_rows, vocab}, at::device(at::kCPU).dtype(at::kFloat)) * 4.0 - + 2.0; + at::Tensor labels = + at::from_blob( + const_cast(labels_data.data()), {n_rows}, at::kInt) + .clone(); + + auto ref = fused_ce_reference(logits, labels, n_valid); + at::Tensor ref_loss = ref.first; + at::Tensor ref_dlogits = ref.second; + + using namespace vkcompute; + + GraphConfig config; + ComputeGraph graph(config); + + IOValueRef r_logits = graph.add_input_tensor( + logits.sizes().vec(), vkapi::kFloat, utils::kBuffer); + IOValueRef r_labels = + graph.add_input_tensor(labels.sizes().vec(), vkapi::kInt, utils::kBuffer); + + const ValueRef r_n_valid = graph.add_scalar(n_valid); + + const ValueRef r_loss = graph.add_tensor({}, vkapi::kFloat, utils::kBuffer); + const ValueRef r_dlogits = + graph.add_tensor({n_rows, vocab}, vkapi::kFloat, utils::kBuffer); + const ValueRef r_out = graph.add_value_list({r_loss, r_dlogits}); + + VK_GET_OP_FN("et_vk.fused_ce.default") + (graph, {r_logits.value, r_labels.value, r_n_valid, r_out}); + + ValueRef staging_loss = graph.set_output_tensor(r_loss); + ValueRef staging_dlogits = graph.set_output_tensor(r_dlogits); + + graph.prepare(); + graph.prepack(); + graph.propagate_resize(); + + graph.maybe_cast_and_copy_into_staging( + r_logits.staging, logits.const_data_ptr(), logits.numel(), vkapi::kFloat); + graph.maybe_cast_and_copy_into_staging( + r_labels.staging, labels.const_data_ptr(), labels.numel(), vkapi::kInt); + + graph.execute(); + + at::Tensor vk_loss = at::zeros({}, at::device(at::kCPU).dtype(at::kFloat)); + graph.maybe_cast_and_copy_from_staging( + staging_loss, vk_loss.mutable_data_ptr(), 1, vkapi::kFloat); + + at::Tensor vk_dlogits = + at::zeros({n_rows, vocab}, at::device(at::kCPU).dtype(at::kFloat)); + graph.maybe_cast_and_copy_from_staging( + staging_dlogits, + vk_dlogits.mutable_data_ptr(), + vk_dlogits.numel(), + vkapi::kFloat); + + const double atol = 1e-4; + const double rtol = 1e-4; + + const bool loss_ok = at::allclose(ref_loss, vk_loss, rtol, atol); + const bool dlogits_ok = at::allclose(ref_dlogits, vk_dlogits, rtol, atol); + + if (!loss_ok || !dlogits_ok) { + std::cout << "fused_ce mismatch: n_rows=" << n_rows << " vocab=" << vocab + << " n_valid=" << n_valid << std::endl; + std::cout << "loss ref=" << ref_loss.item() + << " vk=" << vk_loss.item() << std::endl; + std::cout << "max dlogits diff=" + << at::max(at::abs(ref_dlogits - vk_dlogits)).item() + << std::endl; + } + ASSERT_TRUE(loss_ok); + ASSERT_TRUE(dlogits_ok); +} + +TEST(VulkanFusedCeTest, all_valid_small) { + test_vulkan_fused_ce( + /*n_rows=*/4, /*vocab=*/8, /*labels=*/{0, 3, 7, 1}, /*n_valid=*/4.0); +} + +TEST(VulkanFusedCeTest, masked_label) { + // One masked row (label < 0): contributes 0 to loss and 0 gradient. + test_vulkan_fused_ce( + /*n_rows=*/4, /*vocab=*/8, /*labels=*/{2, -1, 5, 0}, /*n_valid=*/3.0); +} + +TEST(VulkanFusedCeTest, large_vocab_strided_reduce) { + // vocab >> NWORKERS exercises the strided per-row reduction. + test_vulkan_fused_ce( + /*n_rows=*/3, /*vocab=*/200, /*labels=*/{17, -1, 199}, /*n_valid=*/2.0); +} diff --git a/backends/vulkan/test/op_tests/targets.bzl b/backends/vulkan/test/op_tests/targets.bzl index f3e6d42bc54..1b053343460 100644 --- a/backends/vulkan/test/op_tests/targets.bzl +++ b/backends/vulkan/test/op_tests/targets.bzl @@ -192,6 +192,12 @@ def define_common_targets(is_fbcode = False): ":test_utils", ] ) + define_test_targets( + "fused_ce_test", + extra_deps = [ + ":test_utils", + ] + ) define_test_targets( "rms_norm_test", extra_deps = [ From d27eb128e340377ed96bfe03dcb77b69bd665048 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 21 Jul 2026 00:29:37 -0700 Subject: [PATCH 11/15] [ExecuTorch][Vulkan] Add et_vk.linear_q4gsw_backward kernel (4-bit input-grad) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/20943 **Adds the Vulkan `et_vk.linear_q4gsw_backward` kernel** — the input-gradient of the frozen 4-bit `linear_q4gsw` base, for on-device adapter training. Computes `d_x[M, K] = d_out[M, N] @ dequant(W)[N, K]`, contracting over N. **Problem:** `et_vk.linear_q4gsw_backward` is registered in the shared Vulkan partitioner (`custom_ops_lib.py` + `op_registry.py`) but Vulkan had no runtime kernel, so the op could not run. **Solution:** a 4M x 4K register-tiled GLSL kernel that reads the SAME W_4X8 block-packed weight the forward reads, re-deriving the nibble/scale addressing so the training loop stays consistent with the forward with no re-pack. `dequant(W[n, k]) = (code - 8) * scale`. Key changes: - `glsl/q4gsw_backward.{glsl,yaml}` — buffer x float; W_4X8 nibble unpack mirroring `glsl/q4gsw_linear_gemm__w_4x8.glsl` (even-N low nibble, odd-N high, `N4_padded` ivec4 stride, `[num_groups, N]` scales). - `impl/QuantizedLinearBackward.cpp` — reuses `prepack_q4_w_4x8_nc_buffer` + `prepack_q4_scales` from the forward; 1D tile dispatch `ceil(M/4) * ceil(K/4)` with a workgroup-count guard; `group_size` specialization constant. **Constraints:** buffer storage, fp32; `N % 4 == 0`, `K % 4 == 0`, `group_size % 4 == 0` (matches the forward prepack). Weight/scale layout identical to the forward. ghstack-source-id: 405059111 @exported-using-ghexport Differential Revision: [D111797529](https://our.internmc.facebook.com/intern/diff/D111797529/) --- .../graph/ops/glsl/q4gsw_backward.glsl | 104 ++++++++++++++ .../graph/ops/glsl/q4gsw_backward.yaml | 17 +++ .../ops/impl/QuantizedLinearBackward.cpp | 130 ++++++++++++++++++ 3 files changed, 251 insertions(+) create mode 100644 backends/vulkan/runtime/graph/ops/glsl/q4gsw_backward.glsl create mode 100644 backends/vulkan/runtime/graph/ops/glsl/q4gsw_backward.yaml create mode 100644 backends/vulkan/runtime/graph/ops/impl/QuantizedLinearBackward.cpp diff --git a/backends/vulkan/runtime/graph/ops/glsl/q4gsw_backward.glsl b/backends/vulkan/runtime/graph/ops/glsl/q4gsw_backward.glsl new file mode 100644 index 00000000000..d89bb83feed --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/q4gsw_backward.glsl @@ -0,0 +1,104 @@ +/* + * 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. + */ + +#version 450 core + +#define PRECISION ${PRECISION} + +#define T ${texel_load_component_type(DTYPE, STORAGE)} + +#define TILE_M 4 +#define TILE_K 4 + +${define_required_extensions(STORAGE, DTYPE)} +${define_required_extensions("buffer", DTYPE)} + +layout(std430) buffer; + +${layout_declare_tensor(B, "w", "t_dx", DTYPE, STORAGE, is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_dout", DTYPE, STORAGE, is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_q4_weights", "int", "buffer", is_scalar_array=False, vec_size=4)} +${layout_declare_tensor(B, "r", "t_scales", DTYPE, "buffer", is_scalar_array=True)} + +${layout_declare_ubo(B, "ivec4", "dout_sizes")} +${layout_declare_ubo(B, "ivec4", "dx_sizes")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +${layout_declare_spec_const(C, "int", "group_size", "32")} + +// d_x[M, K] = d_out[M, N] @ dequant(W)[N, K], contracting over N. +// dequant(W[n, k]) = (code - 8) * scale, with code read from the same W_4X8 +// block-packed weight the forward reads (mirrors q4gsw_linear_gemm__w_4x8.glsl). +void main() { + const int N = dout_sizes.x; + const int M = dout_sizes.y * dout_sizes.z * dout_sizes.w; + const int K = dx_sizes.x; + + const int nmt = (M + TILE_M - 1) / TILE_M; + const int nkt = (K + TILE_K - 1) / TILE_K; + const int tiles = nmt * nkt; + + const int tile_idx = int(gl_GlobalInvocationID.x); + if (tile_idx >= tiles) { + return; + } + + const int m0 = (tile_idx / nkt) * TILE_M; + const int k0 = (tile_idx % nkt) * TILE_K; + + // K and N are multiples of 4 (prepack guarantees), so k0 is 4-aligned: the + // tile's 4 K lanes are byte b = kl of one k4 group and share one scale group. + const int k4 = k0 >> 2; + const int N4 = (N + 3) >> 2; + const int N4_padded = (N4 + 1) & ~1; + const int N8 = N4_padded >> 1; + const int group = k0 / group_size; + + float acc[TILE_M * TILE_K]; + for (int i = 0; i < TILE_M * TILE_K; ++i) { + acc[i] = 0.0; + } + + for (int n = 0; n < N; ++n) { + float dout_reg[TILE_M]; + for (int ml = 0; ml < TILE_M; ++ml) { + const int m_eff = min(m0 + ml, M - 1); + dout_reg[ml] = float(t_dout[m_eff * N + n]); + } + + // W_4X8 address for column n: ivec4 at (k4, n8); component by (n4 parity, + // n-in-tile half); low/high nibble by n parity (even-N low, odd-N high). + const int n4 = n >> 2; + const int ni = n & 3; + const int n8 = n4 >> 1; + const int comp = (n4 & 1) * 2 + (ni >> 1); + const int nib_hi = (ni & 1) * 4; + const ivec4 w_block = t_q4_weights[k4 * N8 + n8]; + const int w_int = w_block[comp]; + const float scale = float(t_scales[group * N + n]); + + for (int kl = 0; kl < TILE_K; ++kl) { + const int code = int((uint(w_int) >> (8 * kl + nib_hi)) & 0xFu); + const float dq = float(code - 8) * scale; + for (int ml = 0; ml < TILE_M; ++ml) { + acc[ml * TILE_K + kl] += dout_reg[ml] * dq; + } + } + } + + for (int ml = 0; ml < TILE_M; ++ml) { + const int m = m0 + ml; + for (int kl = 0; kl < TILE_K; ++kl) { + const int k = k0 + kl; + if (m < M && k < K) { + t_dx[m * K + k] = T(acc[ml * TILE_K + kl]); + } + } + } +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/q4gsw_backward.yaml b/backends/vulkan/runtime/graph/ops/glsl/q4gsw_backward.yaml new file mode 100644 index 00000000000..b2bb8ea1173 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/q4gsw_backward.yaml @@ -0,0 +1,17 @@ +# 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. + +q4gsw_backward: + parameter_names_with_default_values: + DTYPE: float + STORAGE: buffer + generate_variant_forall: + STORAGE: + - VALUE: buffer + DTYPE: + - VALUE: float + shader_variants: + - NAME: q4gsw_backward diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearBackward.cpp b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearBackward.cpp new file mode 100644 index 00000000000..493daa47126 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearBackward.cpp @@ -0,0 +1,130 @@ +/* + * 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 + +namespace vkcompute { + +// Resize d_x to d_out.shape[:-1] + (K,), mirroring linear_q4gsw_backward_meta. +// extra_args = { weight_data, d_out }. +void resize_linear_q4gsw_backward_node( + ComputeGraph* graph, + const std::vector& args, + const std::vector& extra_args) { + const ValueRef d_x = args.at(0).refs.at(0); + const ValueRef weight_data = extra_args.at(0); + const ValueRef d_out = extra_args.at(1); + const int64_t K = graph->sizes_of(weight_data).at(1) * 2; + std::vector new_sizes = graph->sizes_of(d_out); + new_sizes.back() = K; + graph->virtual_resize(d_x, new_sizes); +} + +utils::uvec3 linear_q4gsw_backward_global_wg_size( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const std::vector& args, + const std::vector& resize_args) { + (void)shader; + (void)resize_args; + const ValueRef d_x = args.at(0).refs.at(0); + const uint32_t K = graph->size_at(-1, d_x); + const uint32_t M = utils::safe_downcast(graph->numel_of(d_x) / K); + const uint32_t tiles = utils::div_up_4(M) * utils::div_up_4(K); + return {tiles, 1u, 1u}; +} + +utils::uvec3 linear_q4gsw_backward_local_wg_size( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const utils::uvec3& global_workgroup_size, + const std::vector& args, + const std::vector& resize_args) { + (void)graph; + (void)shader; + (void)global_workgroup_size; + (void)args; + (void)resize_args; + return {64u, 1u, 1u}; +} + +void linear_q4gsw_backward( + ComputeGraph& graph, + const std::vector& args) { + int32_t i = 0; + const ValueRef d_out = args.at(i++); + const ValueRef weight_data = args.at(i++); + const ValueRef weight_scales_data = args.at(i++); + const ValueRef group_size_ref = args.at(i++); + const ValueRef d_x = args.at(i++); + + VK_CHECK_COND(graph.dtype_of(d_out) == vkapi::kFloat); + VK_CHECK_COND(graph.dtype_of(d_x) == vkapi::kFloat); + VK_CHECK_COND(graph.is_buffer_storage(d_out)); + VK_CHECK_COND(graph.is_buffer_storage(d_x)); + + const vkapi::ScalarType in_dtype = graph.dtype_of(d_out); + const int64_t group_size_val = graph.extract_scalar(group_size_ref); + VK_CHECK_COND(group_size_val > 0 && group_size_val % 4 == 0); + + const std::vector weight_sizes = graph.sizes_of(weight_data); + const int64_t N = weight_sizes.at(0); + const int64_t K = weight_sizes.at(1) * 2; + VK_CHECK_COND(N > 0 && K > 0); + VK_CHECK_COND(N % 4 == 0 && K % 4 == 0); + VK_CHECK_COND(K % group_size_val == 0); + VK_CHECK_COND(graph.size_at(-1, d_out) == N); + + const ValueRef packed_weight = prepack_q4_w_4x8_nc_buffer(graph, weight_data); + const ValueRef packed_scales = + prepack_q4_scales(graph, weight_scales_data, in_dtype); + + const uint32_t M = utils::safe_downcast(graph.numel_of(d_out) / N); + const uint32_t tiles = + utils::div_up_4(M) * utils::div_up_4(static_cast(K)); + VK_CHECK_COND( + (tiles + 63u) / 64u <= 65535u, + "linear_q4gsw_backward: tile count exceeds max workgroup count"); + + std::string kernel_name = "q4gsw_backward"; + kernel_name.reserve(kShaderNameReserve); + add_storage_type_suffix(kernel_name, graph.storage_type_of(d_x)); + add_dtype_suffix(kernel_name, graph.dtype_of(d_x)); + + graph.execute_nodes().emplace_back(new DynamicDispatchNode( + graph, + VK_KERNEL_FROM_STR(kernel_name), + linear_q4gsw_backward_global_wg_size, + linear_q4gsw_backward_local_wg_size, + // Inputs and Outputs + {{d_x, vkapi::kWrite}, + {{d_out, packed_weight, packed_scales}, vkapi::kRead}}, + // Shader params buffers + {graph.sizes_ubo(d_out), graph.sizes_ubo(d_x)}, + // Push Constants + {}, + // Specialization Constants + {static_cast(group_size_val)}, + // Resize Args + {weight_data, d_out}, + // Resizing Logic + resize_linear_q4gsw_backward_node)); +} + +REGISTER_OPERATORS { + VK_REGISTER_OP(et_vk.linear_q4gsw_backward.default, linear_q4gsw_backward); +} + +} // namespace vkcompute From 190414c580e188ba3df1310c23aa4ae21a93a971 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 21 Jul 2026 00:29:38 -0700 Subject: [PATCH 12/15] [ExecuTorch][Vulkan] Tests for et_vk.linear_q4gsw_backward MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/20944 **Correctness tests for the Vulkan `et_vk.linear_q4gsw_backward` kernel** (stacked above the op diff). **Coverage:** the golden `d_x = d_out @ dequant(W)` is computed with ATen (`matmul` over the library-dequantized weight), mirroring the CPU-eager `linear_q4gsw_backward_impl` in `custom_ops_lib.py` — never a hand-rolled matmul. The packed weight is fed as a constant `tensorref` so the op exercises the real `prepack_q4_w_4x8_nc_buffer` W_4X8 path. Cases: - `test_tile_aligned` — single group, tile-aligned M/N/K. - `test_grouped` — multiple quantization groups along K. - `test_odd_n4_partial_m` — `N % 8 != 0` (odd N4 -> padded W_4X8 stride) plus a partial-M tile (exercises the `min()` clamps). Also wires `quantized_linear_backward_test` into `targets.bzl` + `CMakeLists.txt` (mirrors the sibling `linear_q4gsw_dw_test`). ghstack-source-id: 405059117 @exported-using-ghexport Differential Revision: [D111797530](https://our.internmc.facebook.com/intern/diff/D111797530/) --- backends/vulkan/test/op_tests/CMakeLists.txt | 4 + .../quantized_linear_backward_test.cpp | 160 ++++++++++++++++++ backends/vulkan/test/op_tests/targets.bzl | 6 + 3 files changed, 170 insertions(+) create mode 100644 backends/vulkan/test/op_tests/quantized_linear_backward_test.cpp diff --git a/backends/vulkan/test/op_tests/CMakeLists.txt b/backends/vulkan/test/op_tests/CMakeLists.txt index 048371a35bb..0e78b23c184 100644 --- a/backends/vulkan/test/op_tests/CMakeLists.txt +++ b/backends/vulkan/test/op_tests/CMakeLists.txt @@ -125,6 +125,10 @@ if(TARGET vulkan_backend AND LIB_TORCH) vulkan_op_test( fused_ce_test ${CMAKE_CURRENT_SOURCE_DIR}/fused_ce_test.cpp test_utils ) + vulkan_op_test( + quantized_linear_backward_test + ${CMAKE_CURRENT_SOURCE_DIR}/quantized_linear_backward_test.cpp test_utils + ) # Only build generated op tests if a path to tags.yaml and # native_functions.yaml is provided. These files are required for codegen. diff --git a/backends/vulkan/test/op_tests/quantized_linear_backward_test.cpp b/backends/vulkan/test/op_tests/quantized_linear_backward_test.cpp new file mode 100644 index 00000000000..8135c339f99 --- /dev/null +++ b/backends/vulkan/test/op_tests/quantized_linear_backward_test.cpp @@ -0,0 +1,160 @@ +/* + * 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 "test_utils.h" + +// +// Reference Implementation +// + +// Pack unpacked [N, K] codes (0..15) into the flat [N, K/2] uint8 weight the +// forward's prepack consumes: even-K in the low nibble, odd-K in the high. +at::Tensor pack_codes_flat(const at::Tensor& codes) { + const int64_t N = codes.size(0); + const int64_t K = codes.size(1); + at::Tensor packed = + at::empty({N, K / 2}, at::device(at::kCPU).dtype(at::kByte)); + auto ca = codes.accessor(); + auto pa = packed.accessor(); + for (int64_t n = 0; n < N; ++n) { + for (int64_t kb = 0; kb < K / 2; ++kb) { + const int lo = ca[n][2 * kb] & 0xF; + const int hi = ca[n][2 * kb + 1] & 0xF; + pa[n][kb] = static_cast(lo | (hi << 4)); + } + } + return packed; +} + +// Golden d_x[M, K] = d_out[M, N] @ dequant(W)[N, K], with +// dequant(W[n, k]) = (code(n, k) - 8) * scales[k / group_size, n]. +// Mirrors the CPU-eager linear_q4gsw_backward_impl in custom_ops_lib.py. +at::Tensor linear_q4gsw_backward_reference_impl( + const at::Tensor& d_out, + const at::Tensor& codes, + const at::Tensor& scales, + const int64_t group_size) { + const int64_t N = codes.size(0); + const int64_t K = codes.size(1); + const at::Tensor group_idx = + at::arange(K, at::device(at::kCPU).dtype(at::kLong)) + .div(group_size, "floor"); + const at::Tensor scale_full = + scales.t().contiguous().index_select(1, group_idx); // [N, K] + const at::Tensor dequant_w = + (codes.to(at::kFloat) - 8.0) * scale_full; // [N, K] + const at::Tensor d_x_flat = d_out.reshape({-1, N}).matmul(dequant_w); + std::vector out_shape = d_out.sizes().vec(); + out_shape.back() = K; + return d_x_flat.reshape(out_shape).contiguous(); // d_out[..., :-1] + [K] +} + +// +// Test function +// + +void test_vulkan_linear_q4gsw_backward_impl( + const std::vector& d_out_sizes, + const int64_t K, + const int64_t group_size) { + const int64_t N = d_out_sizes.back(); + const int64_t num_groups = K / group_size; + + at::Tensor codes = + at::randint(0, 16, {N, K}, at::device(at::kCPU).dtype(at::kInt)); + at::Tensor scales = + at::rand({num_groups, N}, at::device(at::kCPU).dtype(at::kFloat)) + 0.5; + at::Tensor packed = pack_codes_flat(codes); + at::Tensor d_out = + at::rand(d_out_sizes, at::device(at::kCPU).dtype(at::kFloat)); + + at::Tensor d_x_ref = + linear_q4gsw_backward_reference_impl(d_out, codes, scales, group_size); + + using namespace vkcompute; + + GraphConfig config; + ComputeGraph graph(config); + + ValueRef r_weights = graph.add_tensorref( + packed.sizes().vec(), + from_at_scalartype(packed.scalar_type()), + packed.const_data_ptr()); + ValueRef r_scales = graph.add_tensorref( + scales.sizes().vec(), + from_at_scalartype(scales.scalar_type()), + scales.const_data_ptr()); + + IOValueRef r_d_out = graph.add_input_tensor( + d_out.sizes().vec(), + from_at_scalartype(d_out.scalar_type()), + utils::kBuffer); + const ValueRef r_group_size = graph.add_scalar(group_size); + const ValueRef r_d_x = graph.add_tensor( + d_x_ref.sizes().vec(), + from_at_scalartype(d_x_ref.scalar_type()), + utils::kBuffer); + + VK_GET_OP_FN("et_vk.linear_q4gsw_backward.default") + (graph, {r_d_out.value, r_weights, r_scales, r_group_size, r_d_x}); + + ValueRef staging_out = graph.set_output_tensor(r_d_x); + + graph.prepare(); + graph.prepack(); + graph.propagate_resize(); + + graph.maybe_cast_and_copy_into_staging( + r_d_out.staging, + d_out.const_data_ptr(), + d_out.numel(), + from_at_scalartype(d_out.scalar_type())); + + graph.execute(); + + at::Tensor vk_d_x = at::empty_like(d_x_ref); + graph.maybe_cast_and_copy_from_staging( + staging_out, + vk_d_x.mutable_data_ptr(), + vk_d_x.numel(), + from_at_scalartype(vk_d_x.scalar_type())); + + ASSERT_TRUE(at::allclose(vk_d_x, d_x_ref, 1e-3, 1e-3)); +} + +// Tile-aligned single-group shapes. +TEST(VulkanLinearQ4gswBackwardTest, test_tile_aligned) { + test_vulkan_linear_q4gsw_backward_impl( + /*d_out_sizes=*/{8, 16}, /*K=*/32, /*group_size=*/32); +} + +// Multiple quantization groups along K. +TEST(VulkanLinearQ4gswBackwardTest, test_grouped) { + test_vulkan_linear_q4gsw_backward_impl( + /*d_out_sizes=*/{8, 32}, /*K=*/64, /*group_size=*/32); +} + +// N not a multiple of 8 (odd N4 -> padded W_4X8 stride) plus partial-M tile. +TEST(VulkanLinearQ4gswBackwardTest, test_odd_n4_partial_m) { + test_vulkan_linear_q4gsw_backward_impl( + /*d_out_sizes=*/{5, 12}, /*K=*/16, /*group_size=*/16); +} + +// Leading dims > 2D: M is the flattened product of all leading dims. +TEST(VulkanLinearQ4gswBackwardTest, test_leading_dims_flatten) { + test_vulkan_linear_q4gsw_backward_impl( + /*d_out_sizes=*/{2, 3, 16}, /*K=*/32, /*group_size=*/32); +} diff --git a/backends/vulkan/test/op_tests/targets.bzl b/backends/vulkan/test/op_tests/targets.bzl index 1b053343460..aa3a4f3648b 100644 --- a/backends/vulkan/test/op_tests/targets.bzl +++ b/backends/vulkan/test/op_tests/targets.bzl @@ -198,6 +198,12 @@ def define_common_targets(is_fbcode = False): ":test_utils", ] ) + define_test_targets( + "quantized_linear_backward_test", + extra_deps = [ + ":test_utils", + ] + ) define_test_targets( "rms_norm_test", extra_deps = [ From a6cbbcae729a7d48aa13b59f2b0e988904986135 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 21 Jul 2026 00:29:38 -0700 Subject: [PATCH 13/15] [ExecuTorch][Vulkan] Add et_vk.q4gsw_requant kernel (STE re-quant to W_4X8) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/20945 **Adds the Vulkan `et_vk.q4gsw_requant` kernel** — straight-through re-quant of fp32 latent weights into the frozen-scale 4-bit codes, for on-device weight training. Writes the codes directly in the forward's W_4X8 layout so no per-step re-pack is needed. **Problem:** `et_vk.q4gsw_requant` is registered in the shared Vulkan partitioner but Vulkan had no runtime kernel. **Solution:** a GLSL kernel that quantizes `round(latent / scale)` clamped to `[-8, 7]` and packs the codes in the W_4X8 block layout the forward reads, reusing the exact byte-pair convention of `glsl/pack_q4_linear_weight__w_4x8.glsl` (even-N low nibble, odd-N high; one ivec4 per 4K x 8N block at `(k4, n8)`; OOB upper tile = the bias-zero `0x88888888`). A zero scale yields code 8 (no divide-by-zero), matching the eager reference. Key changes: - `glsl/q4gsw_requant.{glsl,yaml}` — buffer x float; 2D dispatch over `(k4, n8)`. - `impl/QuantizedLinearRequant.cpp` — prepacks the constant scales; output is the W_4X8 int buffer `[K4 * N4_padded * 2]`; `group_size` spec constant; dispatch-grid guard. **Design note (for review):** per the layout decision, requant writes W_4X8 to match the Vulkan forward. Today the forward and backward each prepack their own flat `[N, K/2]` weight internally, so nothing yet consumes an externally-produced W_4X8 buffer — closing the training loop needs a forward path that reads a mutable pre-packed weight (the weight-lifecycle follow-up). The op's AOT meta in `custom_ops_lib.py` still describes the flat `[N, K/2]` output and should be reconciled with this W_4X8 output. **Constraints:** buffer storage, fp32 latent/scales; `N % 4 == 0`, `K % 4 == 0`, `group_size % 4 == 0`. ghstack-source-id: 405059118 @exported-using-ghexport Differential Revision: [D111797527](https://our.internmc.facebook.com/intern/diff/D111797527/) --- .../runtime/graph/ops/glsl/q4gsw_requant.glsl | 97 +++++++++++++ .../runtime/graph/ops/glsl/q4gsw_requant.yaml | 17 +++ .../graph/ops/impl/QuantizedLinearRequant.cpp | 131 ++++++++++++++++++ 3 files changed, 245 insertions(+) create mode 100644 backends/vulkan/runtime/graph/ops/glsl/q4gsw_requant.glsl create mode 100644 backends/vulkan/runtime/graph/ops/glsl/q4gsw_requant.yaml create mode 100644 backends/vulkan/runtime/graph/ops/impl/QuantizedLinearRequant.cpp diff --git a/backends/vulkan/runtime/graph/ops/glsl/q4gsw_requant.glsl b/backends/vulkan/runtime/graph/ops/glsl/q4gsw_requant.glsl new file mode 100644 index 00000000000..896da905648 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/q4gsw_requant.glsl @@ -0,0 +1,97 @@ +/* + * 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. + */ + +#version 450 core + +#define PRECISION ${PRECISION} + +${define_required_extensions(STORAGE, DTYPE)} + +layout(std430) buffer; + +${layout_declare_tensor(B, "w", "t_packed", "int", "buffer", is_scalar_array=False, vec_size=4)} +${layout_declare_tensor(B, "r", "t_latent", DTYPE, STORAGE, is_scalar_array=True)} +${layout_declare_tensor(B, "r", "t_scales", DTYPE, "buffer", is_scalar_array=True)} + +${layout_declare_ubo(B, "ivec4", "latent_sizes")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +${layout_declare_spec_const(C, "int", "group_size", "32")} + +// STE re-quant of fp32 latent [N, K] into the W_4X8 block-packed 4-bit codes +// the forward reads (mirrors pack_q4_linear_weight__w_4x8.glsl). Each thread +// writes one 4K x 8N block (an ivec4) at (k4, n8). +uint quant_nibble(const int n, const int k, const int N, const int K) { + const float s = float(t_scales[(k / group_size) * N + n]); + // roundEven + clamp-in-float match torch.round (half-to-even) then clamp. + float qf = 0.0; + if (s != 0.0) { + qf = clamp(roundEven(float(t_latent[n * K + k]) / s), -8.0, 7.0); + } + return uint(int(qf) + 8) & 0xFu; +} + +// Pack a 4K x 4N tile into the (packed_x, packed_y) int pair. Byte b holds one +// (even-N low nibble, odd-N high nibble) pair at K = k4*4 + b; rows N0,N1 go to +// packed_x, rows N2,N3 to packed_y. +void pack_tile( + out uint packed_x, + out uint packed_y, + const int k4, + const int n4, + const int N, + const int K) { + packed_x = 0u; + packed_y = 0u; + for (int ni = 0; ni < 4; ++ni) { + const int n = n4 * 4 + ni; + for (int b = 0; b < 4; ++b) { + const uint code = quant_nibble(n, k4 * 4 + b, N, K); + const int shift = 8 * b + (ni & 1) * 4; + if (ni < 2) { + packed_x |= code << shift; + } else { + packed_y |= code << shift; + } + } + } +} + +void main() { + const int k4 = int(gl_GlobalInvocationID.x); + const int n8 = int(gl_GlobalInvocationID.y); + + const int K = latent_sizes.x; + const int N = latent_sizes.y; + const int K4 = K >> 2; + const int N4 = (N + 3) >> 2; + const int N8 = (N4 + 1) >> 1; + + if (k4 >= K4 || n8 >= N8) { + return; + } + + const int n4_a = 2 * n8; + const int n4_b = n4_a + 1; + + uint packed_x_a = 0u; + uint packed_y_a = 0u; + // OOB upper tile (odd N4 boundary) is the bias-zero pattern, matching the + // forward's prepack padding so the whole block is always readable. + uint packed_x_b = 0x88888888u; + uint packed_y_b = 0x88888888u; + + pack_tile(packed_x_a, packed_y_a, k4, n4_a, N, K); + if (n4_b < N4) { + pack_tile(packed_x_b, packed_y_b, k4, n4_b, N, K); + } + + t_packed[k4 * N8 + n8] = ivec4( + int(packed_x_a), int(packed_y_a), int(packed_x_b), int(packed_y_b)); +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/q4gsw_requant.yaml b/backends/vulkan/runtime/graph/ops/glsl/q4gsw_requant.yaml new file mode 100644 index 00000000000..275bec5bae2 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/q4gsw_requant.yaml @@ -0,0 +1,17 @@ +# 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. + +q4gsw_requant: + parameter_names_with_default_values: + DTYPE: float + STORAGE: buffer + generate_variant_forall: + STORAGE: + - VALUE: buffer + DTYPE: + - VALUE: float + shader_variants: + - NAME: q4gsw_requant diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearRequant.cpp b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearRequant.cpp new file mode 100644 index 00000000000..ab7336fcc33 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearRequant.cpp @@ -0,0 +1,131 @@ +/* + * 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 + +namespace vkcompute { + +// Resize the packed output to the W_4X8 int buffer size K4 * N4_padded * 2, +// matching prepack_q4_w_4x8_nc_buffer's layout. extra_args = { latent }. +void resize_q4gsw_requant_node( + ComputeGraph* graph, + const std::vector& args, + const std::vector& extra_args) { + const ValueRef packed = args.at(0).refs.at(0); + const ValueRef latent = extra_args.at(0); + const std::vector latent_sizes = graph->sizes_of(latent); + const int64_t N = latent_sizes.at(0); + const int64_t K = latent_sizes.at(1); + const int64_t K4 = K / 4; + const int64_t N4 = N / 4; + const int64_t N4_padded = (N4 + 1) & ~int64_t{1}; + graph->virtual_resize(packed, {K4 * N4_padded * 2}); +} + +utils::uvec3 q4gsw_requant_global_wg_size( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const std::vector& args, + const std::vector& resize_args) { + (void)shader; + (void)resize_args; + const ValueRef latent = args.at(1).refs.at(0); + const std::vector latent_sizes = graph->sizes_of(latent); + const uint32_t N = utils::safe_downcast(latent_sizes.at(0)); + const uint32_t K = utils::safe_downcast(latent_sizes.at(1)); + const uint32_t K4 = K / 4u; + const uint32_t N4 = (N + 3u) / 4u; + const uint32_t N8 = (N4 + 1u) / 2u; + return {K4, N8, 1u}; +} + +utils::uvec3 q4gsw_requant_local_wg_size( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const utils::uvec3& global_workgroup_size, + const std::vector& args, + const std::vector& resize_args) { + (void)graph; + (void)shader; + (void)global_workgroup_size; + (void)args; + (void)resize_args; + return {8u, 8u, 1u}; +} + +void q4gsw_requant(ComputeGraph& graph, const std::vector& args) { + int32_t i = 0; + const ValueRef latent = args.at(i++); + const ValueRef scales = args.at(i++); + const ValueRef group_size_ref = args.at(i++); + const ValueRef packed = args.at(i++); + + VK_CHECK_COND(graph.dtype_of(latent) == vkapi::kFloat); + VK_CHECK_COND(graph.dtype_of(scales) == vkapi::kFloat); + VK_CHECK_COND(graph.dtype_of(packed) == vkapi::kInt); + VK_CHECK_COND(graph.is_buffer_storage(latent)); + VK_CHECK_COND(graph.is_buffer_storage(packed)); + + const int64_t group_size_val = graph.extract_scalar(group_size_ref); + VK_CHECK_COND(group_size_val > 0 && group_size_val % 4 == 0); + + const std::vector latent_sizes = graph.sizes_of(latent); + VK_CHECK_COND(latent_sizes.size() == 2); + const int64_t N = latent_sizes.at(0); + const int64_t K = latent_sizes.at(1); + VK_CHECK_COND(N > 0 && K > 0); + VK_CHECK_COND(N % 4 == 0 && K % 4 == 0); + VK_CHECK_COND(K % group_size_val == 0); + + const uint32_t K4 = utils::safe_downcast(K / 4); + const uint32_t N4 = utils::safe_downcast(N / 4); + const uint32_t N8 = (N4 + 1u) / 2u; + VK_CHECK_COND( + K4 <= 65535u && N8 <= 65535u, + "q4gsw_requant: dispatch grid exceeds max workgroup count"); + + // Scales are a frozen constant; materialize them to a GPU buffer once. + const ValueRef packed_scales = + prepack_standard(graph, scales, utils::kBuffer, utils::kWidthPacked); + + std::string kernel_name = "q4gsw_requant"; + kernel_name.reserve(kShaderNameReserve); + add_storage_type_suffix(kernel_name, graph.storage_type_of(latent)); + add_dtype_suffix(kernel_name, graph.dtype_of(latent)); + + graph.execute_nodes().emplace_back(new DynamicDispatchNode( + graph, + VK_KERNEL_FROM_STR(kernel_name), + q4gsw_requant_global_wg_size, + q4gsw_requant_local_wg_size, + // Inputs and Outputs + {{packed, vkapi::kWrite}, {{latent, packed_scales}, vkapi::kRead}}, + // Shader params buffers + {graph.sizes_ubo(latent)}, + // Push Constants + {}, + // Specialization Constants + {static_cast(group_size_val)}, + // Resize Args + {latent}, + // Resizing Logic + resize_q4gsw_requant_node)); +} + +REGISTER_OPERATORS { + VK_REGISTER_OP(et_vk.q4gsw_requant.default, q4gsw_requant); +} + +} // namespace vkcompute From 0d2f4eebc0eae3dc10ab12140506749a815be158 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 21 Jul 2026 00:29:38 -0700 Subject: [PATCH 14/15] [ExecuTorch][Vulkan] Tests for et_vk.q4gsw_requant MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/20946 **Correctness tests for the Vulkan `et_vk.q4gsw_requant` kernel** (stacked above the op diff). **Coverage:** the golden codes are computed with ATen (`round`/`clamp`, zero-scale -> code 8), mirroring `quant_nibble`, then packed into the expected W_4X8 int buffer with a small bit-packing reference (data-reshaping only, no hand-rolled math). The kernel output is compared int-for-int against that buffer, which locks the exact byte layout the forward reads. The latent is built as `code * scale` so `round()` is unambiguous (no `.5` tie-break divergence). Cases: - `test_tile_aligned` — single group, tile-aligned N/K. - `test_grouped` — multiple quantization groups along K. - `test_odd_n4` — `N % 8 != 0` (odd N4 -> padded stride + bias-zero OOB tile). - `test_zero_scale` — a zero scale must yield code 8, not a divide-by-zero. Also wires `q4gsw_requant_test` into `targets.bzl` + `CMakeLists.txt`. ghstack-source-id: 405059122 @exported-using-ghexport Differential Revision: [D111797526](https://our.internmc.facebook.com/intern/diff/D111797526/) --- backends/vulkan/test/op_tests/CMakeLists.txt | 4 + .../test/op_tests/q4gsw_requant_test.cpp | 195 ++++++++++++++++++ backends/vulkan/test/op_tests/targets.bzl | 6 + 3 files changed, 205 insertions(+) create mode 100644 backends/vulkan/test/op_tests/q4gsw_requant_test.cpp diff --git a/backends/vulkan/test/op_tests/CMakeLists.txt b/backends/vulkan/test/op_tests/CMakeLists.txt index 0e78b23c184..0f8456accf5 100644 --- a/backends/vulkan/test/op_tests/CMakeLists.txt +++ b/backends/vulkan/test/op_tests/CMakeLists.txt @@ -129,6 +129,10 @@ if(TARGET vulkan_backend AND LIB_TORCH) quantized_linear_backward_test ${CMAKE_CURRENT_SOURCE_DIR}/quantized_linear_backward_test.cpp test_utils ) + vulkan_op_test( + q4gsw_requant_test ${CMAKE_CURRENT_SOURCE_DIR}/q4gsw_requant_test.cpp + test_utils + ) # Only build generated op tests if a path to tags.yaml and # native_functions.yaml is provided. These files are required for codegen. diff --git a/backends/vulkan/test/op_tests/q4gsw_requant_test.cpp b/backends/vulkan/test/op_tests/q4gsw_requant_test.cpp new file mode 100644 index 00000000000..a0233bc2706 --- /dev/null +++ b/backends/vulkan/test/op_tests/q4gsw_requant_test.cpp @@ -0,0 +1,195 @@ +/* + * 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 "test_utils.h" + +// +// Reference Implementation +// + +// Pack [N, K] codes (0..15) into the W_4X8 block-packed int buffer the forward +// reads. Mirrors pack_q4_linear_weight__w_4x8.glsl: one ivec4 per (k4, n8), +// byte b holds an (even-N low nibble, odd-N high nibble) pair at K = k4*4 + b. +std::vector pack_codes_w4x8(const at::Tensor& codes) { + const int64_t N = codes.size(0); + const int64_t K = codes.size(1); + const int64_t K4 = K / 4; + const int64_t N4 = N / 4; + const int64_t N4_padded = (N4 + 1) & ~int64_t{1}; + const int64_t N8 = N4_padded / 2; + std::vector buf(K4 * N4_padded * 2, 0); + auto ca = codes.accessor(); + + auto pack_tile = [&](int64_t k4, int64_t n4, uint32_t& px, uint32_t& py) { + px = 0u; + py = 0u; + for (int ni = 0; ni < 4; ++ni) { + const int64_t n = n4 * 4 + ni; + for (int b = 0; b < 4; ++b) { + const uint32_t code = static_cast(ca[n][k4 * 4 + b] & 0xF); + const int shift = 8 * b + (ni & 1) * 4; + if (ni < 2) { + px |= code << shift; + } else { + py |= code << shift; + } + } + } + }; + + for (int64_t k4 = 0; k4 < K4; ++k4) { + for (int64_t n8 = 0; n8 < N8; ++n8) { + const int64_t n4_a = 2 * n8; + const int64_t n4_b = n4_a + 1; + uint32_t px_a, py_a, px_b = 0x88888888u, py_b = 0x88888888u; + pack_tile(k4, n4_a, px_a, py_a); + if (n4_b < N4) { + pack_tile(k4, n4_b, px_b, py_b); + } + const int64_t base = (k4 * N8 + n8) * 4; + buf[base + 0] = static_cast(px_a); + buf[base + 1] = static_cast(py_a); + buf[base + 2] = static_cast(px_b); + buf[base + 3] = static_cast(py_b); + } + } + return buf; +} + +// +// Test function +// + +void test_vulkan_q4gsw_requant_impl( + const int64_t N, + const int64_t K, + const int64_t group_size, + const bool with_zero_scale) { + const int64_t num_groups = K / group_size; + + at::Tensor scales = + at::rand({num_groups, N}, at::device(at::kCPU).dtype(at::kFloat)) + 0.5; + if (with_zero_scale) { + scales.index_put_({0, 0}, 0.0); + } + + const at::Tensor group_idx = + at::arange(K, at::device(at::kCPU).dtype(at::kLong)) + .div(group_size, "floor"); + const at::Tensor scale_full = + scales.t().contiguous().index_select(1, group_idx); // [N, K] + + // Deterministic quotient targets, each >=0.2 from any .5 tie, so GPU fp32 + // division (~2.5 ULP, not correctly rounded) and the CPU golden round + // identically. Covers round both directions and clamp past [-8, 7]. + const std::vector pattern = { + 0.3f, -0.4f, 2.7f, -3.3f, 6.4f, -6.4f, 13.2f, -21.7f}; + const at::Tensor pat = + at::tensor(pattern, at::device(at::kCPU).dtype(at::kFloat)); + const at::Tensor q_idx = + at::arange(N * K, at::device(at::kCPU).dtype(at::kLong)) + .remainder(static_cast(pattern.size())); + const at::Tensor target_q = pat.index_select(0, q_idx).reshape({N, K}); + at::Tensor latent = target_q * scale_full; + + // Golden codes, mirroring quant_nibble: q=0 where scale==0, else roundEven + // (matches at::round half-to-even); clamp to [-8, 7]; code = (q + 8) & 0xF. + const at::Tensor nonzero = scale_full != 0; + const at::Tensor safe = + at::where(nonzero, scale_full, at::ones_like(scale_full)); + at::Tensor q = at::round(latent / safe); + q = at::where(nonzero, q, at::zeros_like(q)); + q = at::clamp(q, -8, 7); + const at::Tensor golden_codes = + (q.to(at::kInt) + 8).bitwise_and(0xF); // [N, K] in 0..15 + + const std::vector expected = pack_codes_w4x8(golden_codes); + + using namespace vkcompute; + + GraphConfig config; + ComputeGraph graph(config); + + IOValueRef r_latent = graph.add_input_tensor( + latent.sizes().vec(), + from_at_scalartype(latent.scalar_type()), + utils::kBuffer); + ValueRef r_scales = graph.add_tensorref( + scales.sizes().vec(), + from_at_scalartype(scales.scalar_type()), + scales.const_data_ptr()); + const ValueRef r_group_size = graph.add_scalar(group_size); + + const int64_t N4 = N / 4; + const int64_t N4_padded = (N4 + 1) & ~int64_t{1}; + const ValueRef r_packed = + graph.add_tensor({(K / 4) * N4_padded * 2}, vkapi::kInt, utils::kBuffer); + + VK_GET_OP_FN("et_vk.q4gsw_requant.default") + (graph, {r_latent.value, r_scales, r_group_size, r_packed}); + + ValueRef staging_out = graph.set_output_tensor(r_packed); + + graph.prepare(); + graph.prepack(); + graph.propagate_resize(); + + graph.maybe_cast_and_copy_into_staging( + r_latent.staging, + latent.const_data_ptr(), + latent.numel(), + from_at_scalartype(latent.scalar_type())); + + graph.execute(); + + at::Tensor vk_packed = at::empty( + {static_cast(expected.size())}, + at::device(at::kCPU).dtype(at::kInt)); + graph.maybe_cast_and_copy_from_staging( + staging_out, + vk_packed.mutable_data_ptr(), + vk_packed.numel(), + from_at_scalartype(vk_packed.scalar_type())); + + auto va = vk_packed.accessor(); + for (size_t i = 0; i < expected.size(); ++i) { + ASSERT_EQ(va[i], expected[i]) << "mismatch at packed int " << i; + } +} + +// Tile-aligned single-group. +TEST(VulkanQ4gswRequantTest, test_tile_aligned) { + test_vulkan_q4gsw_requant_impl( + /*N=*/16, /*K=*/32, /*group_size=*/32, /*with_zero_scale=*/false); +} + +// Multiple quantization groups along K. +TEST(VulkanQ4gswRequantTest, test_grouped) { + test_vulkan_q4gsw_requant_impl( + /*N=*/32, /*K=*/64, /*group_size=*/32, /*with_zero_scale=*/false); +} + +// N not a multiple of 8 (odd N4 -> padded stride + bias-zero OOB tile). +TEST(VulkanQ4gswRequantTest, test_odd_n4) { + test_vulkan_q4gsw_requant_impl( + /*N=*/12, /*K=*/16, /*group_size=*/16, /*with_zero_scale=*/false); +} + +// A zero scale must produce the bias-zero code (8), not a divide-by-zero. +TEST(VulkanQ4gswRequantTest, test_zero_scale) { + test_vulkan_q4gsw_requant_impl( + /*N=*/16, /*K=*/32, /*group_size=*/32, /*with_zero_scale=*/true); +} diff --git a/backends/vulkan/test/op_tests/targets.bzl b/backends/vulkan/test/op_tests/targets.bzl index aa3a4f3648b..8284c5c9aca 100644 --- a/backends/vulkan/test/op_tests/targets.bzl +++ b/backends/vulkan/test/op_tests/targets.bzl @@ -204,6 +204,12 @@ def define_common_targets(is_fbcode = False): ":test_utils", ] ) + define_test_targets( + "q4gsw_requant_test", + extra_deps = [ + ":test_utils", + ] + ) define_test_targets( "rms_norm_test", extra_deps = [ From d9ef877a3aaf9c61f5328617d94773d45de0f3bc Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 21 Jul 2026 00:29:39 -0700 Subject: [PATCH 15/15] [ExecuTorch][WebGPU] Glob runtime/ops sources in CMakeLists MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/21005 Replace the hand-maintained explicit `WEBGPU_SRCS` op-handler list with `file(GLOB WEBGPU_OP_SRCS CONFIGURE_DEPENDS runtime/ops/*/*.cpp)` so adding a new op no longer requires editing this file (addresses review feedback). The five `runtime/*.cpp` sources and `runtime/ops/OperatorRegistry.cpp` (which sits directly under `ops/`, not a per-op subdir) stay explicit. `CONFIGURE_DEPENDS` re-globs at build time when op sources are added or removed. The glob resolves to exactly the 40 op handlers the explicit list enumerated (verified set-equal) — no op added or dropped, and static-init registration is order-independent under `--whole-archive`. Co-authored-with: Claude Code. ghstack-source-id: 405059133 @exported-using-ghexport Differential Revision: [D112482039](https://our.internmc.facebook.com/intern/diff/D112482039/) --- backends/webgpu/CMakeLists.txt | 56 ++++++---------------------------- 1 file changed, 10 insertions(+), 46 deletions(-) diff --git a/backends/webgpu/CMakeLists.txt b/backends/webgpu/CMakeLists.txt index ff5a0fc172a..78d9c4f30dd 100644 --- a/backends/webgpu/CMakeLists.txt +++ b/backends/webgpu/CMakeLists.txt @@ -26,54 +26,18 @@ if(NOT TARGET vulkan_schema) endif() set(WEBGPU_SRCS - runtime/WebGPUBackend.cpp - runtime/WebGPUGraph.cpp - runtime/WebGPUDelegateHeader.cpp - runtime/WebGPUDevice.cpp - runtime/WebGPUQueryPool.cpp - runtime/ops/OperatorRegistry.cpp - runtime/ops/add/BinaryOp.cpp - runtime/ops/rms_norm/RmsNorm.cpp - runtime/ops/update_cache/UpdateCache.cpp - 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 - runtime/ops/prepack/Prepack.cpp - runtime/ops/view_copy/ViewCopy.cpp - runtime/ops/select/Select.cpp - runtime/ops/sigmoid/UnaryOp.cpp - runtime/ops/squeeze/Squeeze.cpp - runtime/ops/unsqueeze/Unsqueeze.cpp - runtime/ops/slice/Slice.cpp - runtime/ops/permute/Permute.cpp - runtime/ops/cat/Cat.cpp - runtime/ops/index/Index.cpp - runtime/ops/sdpa_fd_decode/SdpaFdDecode.cpp - runtime/ops/mm/Mm.cpp - runtime/ops/fused_ce/FusedCe.cpp - runtime/ops/log_softmax/LogSoftmax.cpp - runtime/ops/softmax/Softmax.cpp - runtime/ops/bmm/Bmm.cpp - runtime/ops/reduce/Reduce.cpp - runtime/ops/div/BinaryOp.cpp - runtime/ops/sub/BinaryOp.cpp - runtime/ops/where/Where.cpp - runtime/ops/boolean_op/BooleanOp.cpp - runtime/ops/gather/Gather.cpp - runtime/ops/expand_copy/ExpandCopy.cpp - runtime/ops/fill/Fill.cpp - runtime/ops/dim_order/DimOrder.cpp - runtime/ops/linear/Linear.cpp - runtime/ops/embedding/Embedding.cpp - runtime/ops/adamw/AdamwStep.cpp - runtime/ops/quantized_linear/LinearDw.cpp - runtime/ops/quantized_linear/QuantizedLinearRequant.cpp + runtime/WebGPUBackend.cpp runtime/WebGPUGraph.cpp + runtime/WebGPUDelegateHeader.cpp runtime/WebGPUDevice.cpp + runtime/WebGPUQueryPool.cpp runtime/ops/OperatorRegistry.cpp ) +# Op handlers: glob so adding an op needs no CMakeLists edit. CONFIGURE_DEPENDS +# re-globs at build time when op sources are added or removed. +file(GLOB WEBGPU_OP_SRCS CONFIGURE_DEPENDS + "${CMAKE_CURRENT_SOURCE_DIR}/runtime/ops/*/*.cpp" +) +list(APPEND WEBGPU_SRCS ${WEBGPU_OP_SRCS}) + add_library(webgpu_backend ${WEBGPU_SRCS}) # Verify committed *_wgsl.h match their *.wgsl (drift fails the build).