diff --git a/backends/vulkan/test/op_tests/CMakeLists.txt b/backends/vulkan/test/op_tests/CMakeLists.txt index 5f733deedaa..4f2a0fb94a4 100644 --- a/backends/vulkan/test/op_tests/CMakeLists.txt +++ b/backends/vulkan/test/op_tests/CMakeLists.txt @@ -126,6 +126,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 83f9b39a9a4..707d47ba50f 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 = [