From 6d38b9580d8670e976e3a94265bad7f3fd7bc242 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 21 Jul 2026 00:29:38 -0700 Subject: [PATCH 1/2] [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 4b65b5b81dba9be00d5dc6f60fd9a1e479d886f8 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 21 Jul 2026 00:29:39 -0700 Subject: [PATCH 2/2] [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).