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