diff --git a/backends/vulkan/test/op_tests/CMakeLists.txt b/backends/vulkan/test/op_tests/CMakeLists.txt index 4f2a0fb94a4..6623d2cde44 100644 --- a/backends/vulkan/test/op_tests/CMakeLists.txt +++ b/backends/vulkan/test/op_tests/CMakeLists.txt @@ -130,6 +130,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 707d47ba50f..783ecfa3925 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 = [