From 584525391d783b5e40f44e172e3405d59aa07631 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 21 Jul 2026 00:29:37 -0700 Subject: [PATCH 1/6] [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 1fce414e5b381daa9014654a59cbfb7d82c1b531 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 21 Jul 2026 00:29:37 -0700 Subject: [PATCH 2/6] [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 47375909f70d87bb22bd42d50dbb3b4a2bff228d Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 21 Jul 2026 00:29:38 -0700 Subject: [PATCH 3/6] [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 e334a4d188fbda8ffe25f1cc649bb0a7633adeb1 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 21 Jul 2026 00:29:38 -0700 Subject: [PATCH 4/6] [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 2e76c0c62f40b7dcf850d32d42d35225088554d9 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 21 Jul 2026 00:29:38 -0700 Subject: [PATCH 5/6] [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 70c6fd28b3bee7a1374667573b1a2cb3bba76370 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 21 Jul 2026 00:29:39 -0700 Subject: [PATCH 6/6] [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).