[Vulkan] Reduce 4w decode overhead with specialization

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);

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<uint32_t>(graph.size_at<int64_t>(-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<uint32_t>(
+        graph.size_at<int64_t>(-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<uint32_t>(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<uint32_t>(group_size_val),
       output);
 }

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<int64_t>(group_size_ref);
 
-  std::vector<int64_t> weight_sizes = graph.sizes_of(weight_data);
-  const int64_t K = weight_sizes.at(1) * 2;
-  const uint32_t K_val = static_cast<uint32_t>(K);
+  const uint32_t K_val =
+      utils::safe_downcast<uint32_t>(graph.size_at<int64_t>(-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<uint32_t>(group_size_val)};
+  if (in_dtype == vkapi::kHalf) {
+    spec_vars.append(utils::safe_downcast<uint32_t>(
+        graph.size_at<int64_t>(-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<uint32_t>(group_size_val)},
+      spec_vars,
       {weight_data, fp_input},
       resize_q4gsw_linear_node));
 }

backends/vulkan/test/vulkan_compute_api_test.cpp

@@ -8,8 +8,10 @@
 
 #include <gtest/gtest.h>
 
+#include <algorithm>
 #include <bitset>
 #include <iomanip>
+#include <limits>
 #include <utility>
 #include <vector>
 
@@ -2590,6 +2592,202 @@ TEST(
   test_quantize_and_pack_handles_dynamic_row_counts(128, {2u, 1u, 32u});
 }
 
+template <typename T>
+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<float> distribution(-1.0f, 1.0f);
+          std::vector<uint8_t> weights(N * K / 2);
+          for (auto& value : weights) {
+            value = static_cast<uint8_t>(rng());
+          }
+          std::vector<T> scales(K / group_size * N);
+          for (auto& value : scales) {
+            value = T(0.01f + 0.1f * std::abs(distribution(rng)));
+          }
+          std::vector<T> 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<int64_t>(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<int64_t>({M, N}));
+            std::vector<T> 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<T> 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<float>(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<executorch::aten::Half>(vkapi::kHalf);
+}
+
+template <typename T>
+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<uint8_t> 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<T> 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<int64_t>(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<T> x(padded_K, T(std::numeric_limits<float>::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<T> 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<float>(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<executorch::aten::Half>(vkapi::kHalf);
+}
+
 #define CREATE_WEIGHT_TENSOR(name, sizes, dtype, val)              \
   std::vector<float> data_##name(utils::multiply_integers(sizes)); \
   std::fill(data_##name.begin(), data_##name.end(), val);          \