From 23d7924080f818739b042c759d1947a22fa64122 Mon Sep 17 00:00:00 2001 From: Mergen Nachin Date: Mon, 21 Sep 2026 14:14:58 -0400 Subject: [PATCH] [Vulkan] Reduce 4w decode overhead with specialization Eliminate the scale-cache check when each worker advances to a different quantization group, and specialize logical activation K and weight N for FP16 decode. Keep forced-kernel specialization consistent with production. Cover decode/prefill transitions and padded weights, including poisoned spare input capacity. --- .../glsl/q4gsw_linear_gemv_coop__w_4x8.glsl | 18 +- .../runtime/graph/ops/impl/Q4gswLinear.cpp | 14 +- .../custom_ops/impl/TestFpaQ4gswLinear.cpp | 14 +- .../vulkan/test/vulkan_compute_api_test.cpp | 198 ++++++++++++++++++ 4 files changed, 229 insertions(+), 15 deletions(-) diff --git a/backends/vulkan/runtime/graph/ops/glsl/q4gsw_linear_gemv_coop__w_4x8.glsl b/backends/vulkan/runtime/graph/ops/glsl/q4gsw_linear_gemv_coop__w_4x8.glsl index 014c1c2bf70..6a4fa1b15a4 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/q4gsw_linear_gemv_coop__w_4x8.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/q4gsw_linear_gemv_coop__w_4x8.glsl @@ -109,11 +109,11 @@ ${layout_declare_ubo(B, "ivec4", "input_sizes")} layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; ${layout_declare_spec_const(C, "int", "apply_bias", "0")} -// Aligned with the rest of the q4gsw_linear shader family. K is unused here -// (the local one derived from input_sizes shadows it); kept to share -// descriptor + spec-constant layout. +// FP16 decode specializes logical K and weight N. ${layout_declare_spec_const(C, "int", "K", "1024")} ${layout_declare_spec_const(C, "int", "group_size", "32")} +$if DTYPE == "half": + ${layout_declare_spec_const(C, "int", "weight_N", "0")} // Shared memory for the cooperative reduction. Each lane writes 8 partial // floats (one per N row in the n8 tile = 2 vec4) at the end of its K loop; @@ -151,8 +151,11 @@ void main() { // Per-group base offset into the shared-mem partial-sum slabs. const int group_slab_base = group_id * WGS; - const int N = output_sizes.x; - const int K = input_sizes.x; + $if DTYPE == "half": + const int N = weight_N > 0 ? weight_N : output_sizes.x; + $else: + const int N = output_sizes.x; + const int K = input_sizes.x; const int N4 = (N + 3) / 4; const int N2 = N / 2; const int K4 = K / 4; // texels along K @@ -200,9 +203,10 @@ void main() { // the per-group tree reduction stays well-defined for valid peer groups. const int K4_eff = group_valid ? K4 : 0; for (int k4 = lid; k4 < K4_eff; k4 += WGS) { - // Update scales when crossing into a new group. + // A stride spanning a whole quantization group cannot reuse scales; + // specialize out the cache check in that case. const int group_idx = k4 / blocks_per_group; - if (group_idx != cur_group) { + if (blocks_per_group <= WGS || group_idx != cur_group) { sc_a_lo = load_scale_pair(n2_a_lo, group_idx, N2); sc_a_hi = load_scale_pair(n2_a_hi, group_idx, N2); sc_b_lo = load_scale_pair(n2_b_lo, group_idx, N2); diff --git a/backends/vulkan/runtime/graph/ops/impl/Q4gswLinear.cpp b/backends/vulkan/runtime/graph/ops/impl/Q4gswLinear.cpp index cf8afc8c8c1..50842545fab 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q4gswLinear.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Q4gswLinear.cpp @@ -431,14 +431,22 @@ void add_q4gsw_linear_nc_coop_gemv_node( const ValueRef packed_scales, const ValueRef packed_bias, const uint32_t apply_bias, - const uint32_t K_val, const uint32_t group_size_val, const ValueRef output) { const vkapi::ScalarType in_dtype = graph.dtype_of(fp_input); + // Weight packing may pad K beyond the logical input width. + const uint32_t K_val = + utils::safe_downcast(graph.size_at(-1, fp_input)); TmpTensor dummy_transposed_input( &graph, {}, in_dtype, utils::kBuffer, utils::kWidthPacked); + vkapi::SpecVarList spec_vars = {apply_bias, K_val, group_size_val}; + if (in_dtype == vkapi::kHalf) { + spec_vars.append(utils::safe_downcast( + graph.size_at(-2, weight_data))); + } + graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, pick_q4gsw_nc_coop_shader, @@ -453,7 +461,7 @@ void add_q4gsw_linear_nc_coop_gemv_node( vkapi::kRead}}, {graph.sizes_ubo(output), graph.sizes_ubo(fp_input)}, {}, - {apply_bias, K_val, group_size_val}, + spec_vars, {weight_data, fp_input}, resize_q4gsw_linear_node)); } @@ -533,7 +541,6 @@ void add_q4gsw_linear_w_4x8_node( packed_scales, packed_bias, apply_bias, - K_val, static_cast(group_size_val), output); } @@ -628,7 +635,6 @@ void add_q4gsw_linear_tin_w_4x8_node( packed_scales, packed_bias, apply_bias, - K_val, static_cast(group_size_val), output); } diff --git a/backends/vulkan/test/custom_ops/impl/TestFpaQ4gswLinear.cpp b/backends/vulkan/test/custom_ops/impl/TestFpaQ4gswLinear.cpp index a9e5ad96e95..ab150874ad9 100644 --- a/backends/vulkan/test/custom_ops/impl/TestFpaQ4gswLinear.cpp +++ b/backends/vulkan/test/custom_ops/impl/TestFpaQ4gswLinear.cpp @@ -520,9 +520,8 @@ void add_q4gsw_linear_coop_kc_forced_node( const int64_t group_size_val = graph.extract_scalar(group_size_ref); - std::vector weight_sizes = graph.sizes_of(weight_data); - const int64_t K = weight_sizes.at(1) * 2; - const uint32_t K_val = static_cast(K); + const uint32_t K_val = + utils::safe_downcast(graph.size_at(-1, fp_input)); const ValueRef packed_weight_kc = prepack_q4_w_4x8_nc_buffer(graph, weight_data); @@ -587,6 +586,13 @@ void add_q4gsw_linear_coop_kc_forced_node( VK_THROW("add_q4gsw_linear_coop_kc_forced_node: non-coop kind"); } + vkapi::SpecVarList spec_vars = { + apply_bias, K_val, static_cast(group_size_val)}; + if (in_dtype == vkapi::kHalf) { + spec_vars.append(utils::safe_downcast( + graph.size_at(-2, weight_data))); + } + graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, pick_shader, @@ -601,7 +607,7 @@ void add_q4gsw_linear_coop_kc_forced_node( vkapi::kRead}}, {graph.sizes_ubo(output), graph.sizes_ubo(fp_input)}, {}, - {apply_bias, K_val, static_cast(group_size_val)}, + spec_vars, {weight_data, fp_input}, resize_q4gsw_linear_node)); } diff --git a/backends/vulkan/test/vulkan_compute_api_test.cpp b/backends/vulkan/test/vulkan_compute_api_test.cpp index a0ca49cff6b..3a8bff7718a 100644 --- a/backends/vulkan/test/vulkan_compute_api_test.cpp +++ b/backends/vulkan/test/vulkan_compute_api_test.cpp @@ -8,8 +8,10 @@ #include +#include #include #include +#include #include #include @@ -2590,6 +2592,202 @@ TEST( test_quantize_and_pack_handles_dynamic_row_counts(128, {2u, 1u, 32u}); } +template +void test_q4gsw_decode_prefill(vkapi::ScalarType dtype) { + constexpr int64_t K = 384; + constexpr int64_t max_m = 8; + for (const auto storage : {utils::kBuffer, utils::kTexture3D}) { + for (const int64_t N : {12, 1024, 1028, 2048, 4100, 8192}) { + for (const int64_t group_size : {32, 128}) { + for (const bool has_bias : {false, true}) { + SCOPED_TRACE( + ::testing::Message() + << "N=" << N << " group_size=" << group_size + << " bias=" << has_bias << " storage=" << int(storage)); + GraphConfig config; + config.expect_dynamic_shapes = true; + config.enable_querypool = true; + ComputeGraph graph(config); + + std::mt19937 rng(419); + std::uniform_real_distribution distribution(-1.0f, 1.0f); + std::vector weights(N * K / 2); + for (auto& value : weights) { + value = static_cast(rng()); + } + std::vector scales(K / group_size * N); + for (auto& value : scales) { + value = T(0.01f + 0.1f * std::abs(distribution(rng))); + } + std::vector bias(N); + for (auto& value : bias) { + value = T(distribution(rng)); + } + const auto input = graph.add_input_tensor( + {max_m, K}, dtype, storage, utils::kWidthPacked); + const auto output = + graph.add_tensor({max_m, N}, dtype, storage, utils::kWidthPacked); + VK_GET_OP_FN("et_vk.q4gsw_linear.default") + (graph, + {input.value, + graph.add_tensorref({N, K / 2}, vkapi::kByte, weights.data()), + graph.add_tensorref({K / group_size, N}, dtype, scales.data()), + graph.add_scalar(group_size), + has_bias ? graph.add_tensorref({N}, dtype, bias.data()) + : graph.add_none(), + output}); + const auto staging = graph.set_output_tensor(output); + graph.prepare(); + graph.prepack(); + + for (const int64_t M : {1, 8, 1, 3, 1}) { + SCOPED_TRACE(M); + graph.resize_input(0, {M, K}); + graph.propagate_resize(); + ASSERT_EQ(graph.sizes_of(output), std::vector({M, N})); + std::vector x(M * K); + for (auto& value : x) { + value = T(distribution(rng)); + } + graph.maybe_cast_and_copy_into_staging( + input.staging, x.data(), x.size(), dtype); + graph.execute(); + std::vector actual(M * N); + graph.maybe_cast_and_copy_from_staging( + staging, actual.data(), actual.size(), dtype); + double squared_error = 0.0; + double squared_reference = 0.0; + for (int64_t m = 0; m < M; ++m) { + for (int64_t n = 0; n < N; ++n) { + double expected = has_bias ? float(bias[n]) : 0.0; + double magnitude = std::abs(expected); + for (int64_t k = 0; k < K; ++k) { + const int q = + ((weights[n * K / 2 + k / 2] >> (4 * (k % 2))) & 15) - 8; + const double product = double(float(x[m * K + k])) * q * + float(scales[(k / group_size) * N + n]); + expected += product; + magnitude += std::abs(product); + } + const double half_tolerance = + storage == utils::kBuffer ? 5e-4 : 2e-3; + const double tolerance = dtype == vkapi::kHalf + ? half_tolerance * (1.0 + std::abs(expected)) + + (M == 1 ? 0.0 : 5e-3 * magnitude) + : 2e-5 * (1.0 + std::abs(expected)); + ASSERT_NEAR(float(actual[m * N + n]), expected, tolerance) + << "m=" << m << " n=" << n; + const double error = float(actual[m * N + n]) - expected; + squared_error += error * error; + squared_reference += expected * expected; + } + } + EXPECT_LT( + std::sqrt(squared_error / squared_reference), + dtype == vkapi::kHalf ? (M == 1 ? 2e-3 : 1e-2) : 2e-5); + graph.context()->querypool().extract_results(); + bool has_gemv = false; + bool has_gemm = false; + for (const auto& result : + graph.context()->querypool().get_shader_timestamp_data()) { + has_gemv |= result.kernel_name.find("q4gsw_linear_gemv") != + std::string::npos; + has_gemm |= result.kernel_name.find("q4gsw_linear_gemm") != + std::string::npos; + } + EXPECT_EQ(has_gemv, M == 1); + EXPECT_EQ(has_gemm, M != 1); + } + } + } + } + } +} + +TEST(VulkanComputeGraphTest, q4gsw_decode_prefill_float) { + test_q4gsw_decode_prefill(vkapi::kFloat); +} + +TEST(VulkanComputeGraphTest, q4gsw_decode_prefill_half) { + if (!api::context()->adapter_ptr()->supports_16bit_storage_buffers() || + !api::context()->adapter_ptr()->supports_float16_shader_types()) { + GTEST_SKIP(); + } + test_q4gsw_decode_prefill(vkapi::kHalf); +} + +template +void test_q4gsw_decode_with_padded_weights(vkapi::ScalarType dtype) { + constexpr int64_t K = 12; + constexpr int64_t padded_K = 16; + constexpr int64_t N = 128; + for (const auto storage : {utils::kBuffer, utils::kTexture3D}) { + SCOPED_TRACE(int(storage)); + GraphConfig config; + config.expect_dynamic_shapes = true; + ComputeGraph graph(config); + const auto input = graph.add_input_tensor( + {1, padded_K}, dtype, storage, utils::kWidthPacked); + graph.resize_input(0, {1, K}); + const auto output = + graph.add_tensor({1, N}, dtype, storage, utils::kWidthPacked); + + // Match export's zero padding before packing pairs of signed 4-bit weights. + std::vector weights(N * padded_K / 2, 0x88); + for (int64_t n = 0; n < N; ++n) { + std::fill_n(weights.begin() + n * padded_K / 2, K / 2, 0x99); + } + std::vector scales(N, T(1.0f)); + VK_GET_OP_FN("et_vk.q4gsw_linear.default") + (graph, + {input.value, + graph.add_tensorref({N, padded_K / 2}, vkapi::kByte, weights.data()), + graph.add_tensorref({1, N}, dtype, scales.data()), + graph.add_scalar(K), + graph.add_none(), + output}); + const auto staging = graph.set_output_tensor(output); + graph.prepare(); + graph.prepack(); + + // Poison spare input capacity using only the staging node. Zero weights + // must not hide activation reads beyond logical K. + graph.resize_input(0, {1, padded_K}); + graph.propagate_resize(); + std::vector x(padded_K, T(std::numeric_limits::quiet_NaN())); + graph.maybe_cast_and_copy_into_staging( + input.staging, x.data(), x.size(), dtype); + graph.context()->set_cmd(); + graph.execute_nodes().front()->encode(&graph); + graph.context()->submit_cmd_to_gpu(); + graph.context()->wait_for_queue(); + + graph.resize_input(0, {1, K}); + graph.propagate_resize(); + std::fill_n(x.begin(), K, T(1.0f)); + graph.maybe_cast_and_copy_into_staging(input.staging, x.data(), K, dtype); + graph.execute(); + std::vector actual(N); + graph.maybe_cast_and_copy_from_staging( + staging, actual.data(), actual.size(), dtype); + for (int64_t n = 0; n < N; ++n) { + ASSERT_EQ(float(actual[n]), float(K)) << "n=" << n; + } + } +} + +TEST(VulkanComputeGraphTest, q4gsw_decode_with_padded_weights_float) { + test_q4gsw_decode_with_padded_weights(vkapi::kFloat); +} + +TEST(VulkanComputeGraphTest, q4gsw_decode_with_padded_weights_half) { + if (!api::context()->adapter_ptr()->supports_16bit_storage_buffers() || + !api::context()->adapter_ptr()->supports_float16_shader_types()) { + GTEST_SKIP(); + } + test_q4gsw_decode_with_padded_weights(vkapi::kHalf); +} + #define CREATE_WEIGHT_TENSOR(name, sizes, dtype, val) \ std::vector data_##name(utils::multiply_integers(sizes)); \ std::fill(data_##name.begin(), data_##name.end(), val); \