Update on "[ET-VK][Ops] linear_qta8a_qga4w_qta8o impl and shaders"

# Operator Description

The linear_qta8a_qga4w_qta8o operator implements a quantized linear transformation that enables efficient neural network inference through dynamic quantization. This operator performs matrix multiplication between quantized 8-bit activations and 4-bit grouped quantized weights, producing quantized 8-bit outputs.

The quantization scheme follows the standard affine mapping where `real_value = scale * (quantized_value - zero_point)`. Input activations use 8-bit signed integers with per-token scale and zero-point parameters, while weights employ 4-bit quantization with group-wise parameters.

# Implementation Architecture

The operator provides two distinct computational approaches optimized for different matrix multiplication scenarios: the TILED algorithm for general matrix-matrix multiplication (GEMM) and the COOPERATIVE algorithm for matrix-vector multiplication (GEMV).

## TILED Algorithm (GEMM Cases)

The tiled implementation processes the output matrix in rectangular blocks. Each thread is responsible for calculating a tile of output values, typically processing 3 rows and 2 columns worth of results in each iteration. The algorithm operates by having each thread load blocks of quantized weights and activations, perform integer arithmetic accumulation, and then apply the necessary scaling operations.

Weight data is pre-packed in a specialized format where two 4-bit values are stored in each byte. Each thread loads multiple weight elements simultaneously and unpacks them during computation. The quantization parameters for weights are organized by groups, where each group of consecutive weight elements shares the same scale and zero-point values.

## COOPERATIVE Algorithm (GEMV Cases)

The cooperative implementation uses shared memory and thread cooperation where this approach uses workgroups of 64 threads arranged as 8 groups of 8 workers each. The key insight is that GEMV operations have limited parallelism in the output dimension but substantial parallelism in the reduction dimension, making cooperative reduction strategies more effective than independent thread computation.

Each group of 8 worker threads collaboratively computes a portion of the output vector. The workers divide the reduction work along the input feature dimension, with each worker processing every 8th element in a strided pattern.

# Future Performance Improvements
- Making use of dotPacked4x8EXT (this requires upgrading glslc and vulkan)
- Fixed point math for pure integer operations
- Might be more performant to avoid preloading tensors
- Might also be more performant to avoid excessive register overhead by defining the ivec4 within each block operation (allowing more threads to be more register intensive)

Differential Revision: [D77173441](https://our.internmc.facebook.com/intern/diff/D77173441/)

[ghstack-poisoned]

Author: morelos

backends/vulkan/runtime/graph/ops/glsl/linear_qta8a_qga4w_coop.glsl

@@ -157,15 +157,15 @@ void main() {
       // Preload A (quantized input) - keep as quantized integers
       [[unroll]] for (int r = 0; r < TILE_ROWS; ++r) {
         $if IN_STORAGE == "buffer":
-          mat1_quantized[r] = ivec4(t_mat1[((out_row + r) * mat1_sizes.x + k) >> 2] - t_input_zero_point[int(out_row) + r]);
+          mat1_quantized[r] = t_mat1[((out_row + r) * mat1_sizes.x + k) >> 2] - t_input_zero_point[int(out_row) + r];
         $else:
-          mat1_quantized[r] = ivec4(texelFetch(t_mat1, ivec3(k >> 2, out_row + r, 0), 0) - t_input_zero_point[int(out_row) + r]);
-
-        input_sums[r] += mat1_quantized[r].x + mat1_quantized[r].y + mat1_quantized[r].z + mat1_quantized[r].w;
+          mat1_quantized[r] = texelFetch(t_mat1, ivec3(k >> 2, out_row + r, 0), 0) - t_input_zero_point[int(out_row) + r];
       }
 
       // Accumulate in integer arithmetic: (input_quantized - input_zero_point) * (weight_quantized - weight_zero_point)
       [[unroll]] for (int r = 0; r < TILE_ROWS; ++r) {
+        input_sums[r] += mat1_quantized[r].x + mat1_quantized[r].y + mat1_quantized[r].z + mat1_quantized[r].w;
+
         int32_sums[r][0] +=   mat1_quantized[r].x * qmat2_quantized[0][0]
                             + mat1_quantized[r].y * qmat2_quantized[1][0]
                             + mat1_quantized[r].z * qmat2_quantized[2][0]

backends/vulkan/runtime/graph/ops/glsl/linear_qta8a_qga4w_tiled.glsl

@@ -129,7 +129,7 @@ void main() {
       // Preload B (weights) - keep as quantized integers
       [[unroll]] for (int r = 0; r < 4; ++r) {
         $if WEIGHT_STORAGE == "buffer":
-          const uvec4 packed_weight_tex = t_qmat2[(k + r) * qmat2_stride + gl_GlobalInvocationID.x];
+          const u8vec4 packed_weight_tex = t_qmat2[(k + r) * qmat2_stride + gl_GlobalInvocationID.x];
         $else:
           const uvec4 packed_weight_tex = texelFetch(
               t_qmat2,
@@ -144,15 +144,15 @@ void main() {
       // Preload A (quantized input) - keep as quantized integers
       [[unroll]] for (int r = 0; r < TILE_ROWS; ++r) {
         $if IN_STORAGE == "buffer":
-          mat1_quantized[r] = ivec4(t_mat1[((out_row + r) * mat1_sizes.x + k) >> 2] - t_input_zero_point[int(out_row) + r]);
+          mat1_quantized[r] = t_mat1[((out_row + r) * mat1_sizes.x + k) >> 2] - t_input_zero_point[int(out_row) + r];
         $else:
-          mat1_quantized[r] = ivec4(texelFetch(t_mat1, ivec3(k >> 2, out_row + r, 0), 0) - t_input_zero_point[int(out_row) + r]);
-
-        input_sums[r] += mat1_quantized[r].x + mat1_quantized[r].y + mat1_quantized[r].z + mat1_quantized[r].w;
+          mat1_quantized[r] = texelFetch(t_mat1, ivec3(k >> 2, out_row + r, 0), 0) - t_input_zero_point[int(out_row) + r];
       }
 
       // Accumulate in integer arithmetic: (input_quantized - input_zero_point) * (weight_quantized - weight_zero_point)
       [[unroll]] for (int r = 0; r < TILE_ROWS; ++r) {
+        input_sums[r] += mat1_quantized[r].x + mat1_quantized[r].y + mat1_quantized[r].z + mat1_quantized[r].w;
+
         int32_sums[r][0] +=   mat1_quantized[r].x * qmat2_quantized[0][0]
                             + mat1_quantized[r].y * qmat2_quantized[1][0]
                             + mat1_quantized[r].z * qmat2_quantized[2][0]

backends/vulkan/runtime/graph/ops/glsl/quantized_linear_utils.glslh

@@ -0,0 +1,161 @@
+/*
+ * 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.
+ */
+
+/*
+ * Utility macros for quantized linear operations.
+ * This header provides large code block replacements for common patterns
+ * in quantized linear shaders.
+ */
+
+/*
+ * Load quantization scales and zeros for both texel columns
+ * Handles both buffer and texture storage in one macro
+ */
+#define LOAD_SCALES_AND_ZEROS(storage_type, scales_tensor, zeros_tensor, block_idx, out_col_texel_idx, qparams_stride) \
+  do { \
+    $if storage_type == "buffer": \
+      scales[0] = scales_tensor[block_idx * qparams_stride + out_col_texel_idx]; \
+      scales[1] = scales_tensor[block_idx * qparams_stride + out_col_texel_idx + 1]; \
+      zeros[0] = vec4(zeros_tensor[block_idx * qparams_stride + out_col_texel_idx]); \
+      zeros[1] = vec4(zeros_tensor[block_idx * qparams_stride + out_col_texel_idx + 1]); \
+    $else: \
+      scales[0] = texelFetch(scales_tensor, ivec3(out_col_texel_idx, 0, block_idx), 0); \
+      scales[1] = texelFetch(scales_tensor, ivec3(out_col_texel_idx + 1, 0, block_idx), 0); \
+      zeros[0] = vec4(texelFetch(zeros_tensor, ivec3(out_col_texel_idx, 1, block_idx), 0)); \
+      zeros[1] = vec4(texelFetch(zeros_tensor, ivec3(out_col_texel_idx + 1, 1, block_idx), 0)); \
+  } while(false)
+
+/*
+ * Load quantization parameters for qga4w shaders (combined scales/zeros tensor)
+ */
+#define LOAD_QGA4W_PARAMS(storage_type, qparams_tensor, block_idx, out_col_texel_idx, qparams_y_stride, qparams_z_stride) \
+  do { \
+    $if storage_type == "buffer": \
+      scales[0] = qparams_tensor[block_idx * qparams_z_stride + out_col_texel_idx]; \
+      zeros[0] = qparams_tensor[block_idx * qparams_z_stride + out_col_texel_idx + qparams_y_stride]; \
+      scales[1] = qparams_tensor[block_idx * qparams_z_stride + out_col_texel_idx + 1]; \
+      zeros[1] = qparams_tensor[block_idx * qparams_z_stride + out_col_texel_idx + 1 + qparams_y_stride]; \
+    $else: \
+      scales[0] = texelFetch(qparams_tensor, ivec3(out_col_texel_idx, 0, block_idx), 0); \
+      zeros[0] = texelFetch(qparams_tensor, ivec3(out_col_texel_idx, 1, block_idx), 0); \
+      scales[1] = texelFetch(qparams_tensor, ivec3(out_col_texel_idx + 1, 0, block_idx), 0); \
+      zeros[1] = texelFetch(qparams_tensor, ivec3(out_col_texel_idx + 1, 1, block_idx), 0); \
+  } while(false)
+
+/*
+ * Preload and extract 4-bit weights (keep quantized for qta8a shaders)
+ */
+#define PRELOAD_WEIGHTS_QUANTIZED(storage_type, weight_tensor, k, col_idx, stride) \
+  do { \
+    [[unroll]] for (int r = 0; r < 4; ++r) { \
+      $if storage_type == "buffer": \
+        const uvec4 packed_weight_tex = weight_tensor[(k + r) * stride + col_idx]; \
+      $else: \
+        const uvec4 packed_weight_tex = texelFetch(weight_tensor, ivec2(col_idx, k + r), 0); \
+      qmat2_quantized[r][0] = ivec4((packed_weight_tex & 0xF0) >> 4) - 8; \
+      qmat2_quantized[r][1] = ivec4(packed_weight_tex & 0x0F) - 8; \
+    } \
+  } while(false)
+
+/*
+ * Preload 4-bit weights (keep quantized initially, then dequantize)
+ */
+#define PRELOAD_AND_DEQUANTIZE_WEIGHTS(storage_type, weight_tensor, k, col_idx, stride, scales_arr, zeros_arr) \
+  do { \
+    ivec4 qmat2_quantized_temp[4][2]; \
+    [[unroll]] for (int r = 0; r < 4; ++r) { \
+      $if storage_type == "buffer": \
+        const uvec4 packed_weight_tex = weight_tensor[(k + r) * stride + col_idx]; \
+      $else: \
+        const uvec4 packed_weight_tex = texelFetch(weight_tensor, ivec2(col_idx, k + r), 0); \
+      qmat2_quantized_temp[r][0] = ivec4((packed_weight_tex & 0xF0) >> 4) - 8; \
+      qmat2_quantized_temp[r][1] = ivec4(packed_weight_tex & 0x0F) - 8; \
+    } \
+    [[unroll]] for (int r = 0; r < 4; ++r) { \
+      qmat2[r][0] = VEC4_T(qmat2_quantized_temp[r][0]) * scales_arr[0] + zeros_arr[0]; \
+      qmat2[r][1] = VEC4_T(qmat2_quantized_temp[r][1]) * scales_arr[1] + zeros_arr[1]; \
+    } \
+  } while(false)
+
+/*
+ * Preload input matrix for quantized int8 inputs
+ */
+#define PRELOAD_QUANTIZED_INPUT(storage_type, input_tensor, out_row, k, input_sizes, zero_point_tensor, input_sums_arr) \
+  do { \
+    [[unroll]] for (int r = 0; r < TILE_ROWS; ++r) { \
+      $if storage_type == "buffer": \
+        mat1_quantized[r] = ivec4(input_tensor[((out_row + r) * input_sizes.x + k) >> 2] - zero_point_tensor[int(out_row) + r]); \
+      $else: \
+        mat1_quantized[r] = ivec4(texelFetch(input_tensor, ivec3(k >> 2, out_row + r, 0), 0) - zero_point_tensor[int(out_row) + r]); \
+      input_sums_arr[r] += mat1_quantized[r].x + mat1_quantized[r].y + mat1_quantized[r].z + mat1_quantized[r].w; \
+    } \
+  } while(false)
+
+/*
+ * Preload input matrix for float inputs
+ */
+#define PRELOAD_FLOAT_INPUT(storage_type, input_tensor, out_row, k, input_sizes) \
+  do { \
+    [[unroll]] for (int r = 0; r < TILE_ROWS; ++r) { \
+      $if storage_type == "buffer": \
+        mat1[r] = input_tensor[((out_row + r) * input_sizes.x + k) >> 2]; \
+      $else: \
+        mat1[r] = texelFetch(input_tensor, ivec3(k >> 2, out_row + r, 0), 0); \
+    } \
+  } while(false)
+
+/*
+ * Store final output results for both texel columns
+ */
+#define STORE_FINAL_OUTPUT(storage_type, output_tensor, out_row, out_col, out_col_texel_idx, out_sizes, results) \
+  do { \
+    [[unroll]] for (int r = 0; r < TILE_ROWS; ++r) { \
+      $if storage_type == "buffer": \
+        if (out_row + r < out_sizes.y) { \
+          output_tensor[((out_row + r) * out_sizes.x + out_col) >> 2] = results[r][0]; \
+          output_tensor[((out_row + r) * out_sizes.x + out_col + 4) >> 2] = results[r][1]; \
+        } \
+      $else: \
+        imageStore(output_tensor, ivec3(out_col_texel_idx, out_row + r, 0), results[r][0]); \
+        imageStore(output_tensor, ivec3(out_col_texel_idx + 1, out_row + r, 0), results[r][1]); \
+    } \
+  } while(false)
+
+/*
+ * Matrix multiplication accumulation for quantized int8 inputs
+ */
+#define ACCUMULATE_QUANTIZED_MATMUL(mat1_quantized_arr, qmat2_quantized_arr, int32_sums_arr) \
+  do { \
+    [[unroll]] for (int r = 0; r < TILE_ROWS; ++r) { \
+      int32_sums_arr[r][0] += mat1_quantized_arr[r].x * qmat2_quantized_arr[0][0] \
+                            + mat1_quantized_arr[r].y * qmat2_quantized_arr[1][0] \
+                            + mat1_quantized_arr[r].z * qmat2_quantized_arr[2][0] \
+                            + mat1_quantized_arr[r].w * qmat2_quantized_arr[3][0]; \
+      int32_sums_arr[r][1] += mat1_quantized_arr[r].x * qmat2_quantized_arr[0][1] \
+                            + mat1_quantized_arr[r].y * qmat2_quantized_arr[1][1] \
+                            + mat1_quantized_arr[r].z * qmat2_quantized_arr[2][1] \
+                            + mat1_quantized_arr[r].w * qmat2_quantized_arr[3][1]; \
+    } \
+  } while(false)
+
+/*
+ * Matrix multiplication accumulation for float inputs
+ */
+#define ACCUMULATE_FLOAT_MATMUL(mat1_arr, qmat2_arr, sums_arr) \
+  do { \
+    [[unroll]] for (int r = 0; r < TILE_ROWS; ++r) { \
+      sums_arr[r][0] += mat1_arr[r].x * qmat2_arr[0][0] \
+                      + mat1_arr[r].y * qmat2_arr[1][0] \
+                      + mat1_arr[r].z * qmat2_arr[2][0] \
+                      + mat1_arr[r].w * qmat2_arr[3][0]; \
+      sums_arr[r][1] += mat1_arr[r].x * qmat2_arr[0][1] \
+                      + mat1_arr[r].y * qmat2_arr[1][1] \
+                      + mat1_arr[r].z * qmat2_arr[2][1] \
+                      + mat1_arr[r].w * qmat2_arr[3][1]; \
+    } \
+  } while(false)