From 05f302b98910348c4c1be2ebfea75347299342c1 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 14 Jul 2026 14:50:09 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- .../graph/ops/glsl/q4gsw_backward.glsl | 104 ++++++++++++++ .../graph/ops/glsl/q4gsw_backward.yaml | 17 +++ .../ops/impl/QuantizedLinearBackward.cpp | 132 ++++++++++++++++++ 3 files changed, 253 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..f959c5a9604 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearBackward.cpp @@ -0,0 +1,132 @@ +/* + * 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