Add small-batch quantized matvec kernel (qmv_wide) (#3764)

Author: jessegross

mlx/backend/metal/kernels/fp_quantized.h

@@ -525,6 +525,125 @@ METAL_FUNC void fp_qmv_impl(
   }
 }
 
+// Quantized matrix-vector for a small batch of input vectors, the M in
+// [2, vector_limit) band between qmv (M==1) and qmm. Each thread owns one
+// output row that k_lanes lanes reduce over K; the vecs_per_tg vectors are
+// streamed so each weight group is dequantized once and reused across them.
+template <typename T, int group_size, int bits, int vecs_per_tg, int k_lanes>
+METAL_FUNC void fp_qmv_wide_impl(
+    const device uint32_t* w,
+    const device uint8_t* scales,
+    const device T* x,
+    device T* y,
+    const constant int& in_vec_size,
+    const constant int& out_vec_size,
+    const constant int& M,
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  constexpr int num_simdgroups = 2;
+  constexpr int results_per_simdgroup = SIMD_SIZE / k_lanes;
+  constexpr int pack_factor = get_pack_factor<32, bits>();
+  constexpr int bytes_per_pack = get_bytes_per_pack<32>();
+  constexpr int nf4 = group_size / 4; // float4 lanes per quant group
+
+  typedef float U;
+
+  const short k_lane =
+      simd_lid % k_lanes; // this lane's slot in the K reduction
+  const short sg_row = simd_lid / k_lanes; // which output row of the simdgroup
+
+  const int out_row = tid.y * (results_per_simdgroup * num_simdgroups) +
+      results_per_simdgroup * simd_gid + sg_row;
+  const int vec0 = tid.x * vecs_per_tg; // first input vector handled here
+
+  const int row = min(out_row, out_vec_size - 1);
+
+  const int in_vec_size_w = in_vec_size * bytes_per_pack / pack_factor;
+  const int in_vec_size_g = in_vec_size / group_size;
+  const device uint8_t* wrow = (const device uint8_t*)w + row * in_vec_size_w;
+  const device uint8_t* srow = scales + row * in_vec_size_g;
+
+  // One device pointer per streamed vector; the clamp keeps an out-of-range
+  // tail slot reading a valid row (it is never written below).
+  const device T* xv[vecs_per_tg];
+  for (int v = 0; v < vecs_per_tg; v++) {
+    xv[v] = x + min(vec0 + v, M - 1) * in_vec_size;
+  }
+
+  U result[vecs_per_tg] = {0};
+
+  // Each lane reduces a strided subset of the row's quant groups: one group is
+  // dequantized to float4 lanes and dot()'d against each streamed vector. One
+  // group per iteration keeps weight-register pressure low for occupancy.
+  for (int g = k_lane; g < in_vec_size_g; g += k_lanes) {
+    const int k0 = g * group_size;
+    U s = dequantize_scale<U, group_size>(srow[g]);
+    const device uint8_t* wg = wrow + k0 * bytes_per_pack / pack_factor;
+
+    float4 w4[nf4];
+    if constexpr (bits == 4) {
+      const device uint16_t* wq = (const device uint16_t*)wg;
+#pragma unroll
+      for (int i = 0; i < nf4; i++) {
+        w4[i] = float4(
+            Dequantize<4>{}(wq[i]),
+            Dequantize<4>{}(wq[i] >> 4),
+            Dequantize<4>{}(wq[i] >> 8),
+            Dequantize<4>{}(wq[i] >> 12));
+      }
+    } else {
+#pragma unroll
+      for (int i = 0; i < nf4; i++) {
+        w4[i] = float4(
+            Dequantize<8>{}(wg[4 * i]),
+            Dequantize<8>{}(wg[4 * i + 1]),
+            Dequantize<8>{}(wg[4 * i + 2]),
+            Dequantize<8>{}(wg[4 * i + 3]));
+      }
+    }
+
+#pragma unroll
+    for (int v = 0; v < vecs_per_tg; v++) {
+      const device vec<T, 4>* xv4 = (const device vec<T, 4>*)(xv[v] + k0);
+      float acc = 0;
+#pragma unroll
+      for (int j = 0; j < nf4; j++) {
+        acc += dot(w4[j], float4(xv4[j]));
+      }
+      result[v] += s * acc;
+    }
+  }
+
+  // Reduce each vector's partial over its k_lanes with a shuffle ladder:
+  // simd_sum would mix the results_per_simdgroup rows a simdgroup spans.
+  for (int v = 0; v < vecs_per_tg; v++) {
+    if constexpr (k_lanes >= 32) {
+      result[v] += simd_shuffle_down(result[v], 16);
+    }
+    if constexpr (k_lanes >= 16) {
+      result[v] += simd_shuffle_down(result[v], 8);
+    }
+    if constexpr (k_lanes >= 8) {
+      result[v] += simd_shuffle_down(result[v], 4);
+    }
+    if constexpr (k_lanes >= 4) {
+      result[v] += simd_shuffle_down(result[v], 2);
+    }
+    if constexpr (k_lanes >= 2) {
+      result[v] += simd_shuffle_down(result[v], 1);
+    }
+  }
+
+  if (k_lane == 0 && out_row < out_vec_size) {
+    for (int v = 0; v < vecs_per_tg; v++) {
+      if (vec0 + v < M) {
+        y[(vec0 + v) * out_vec_size + out_row] = static_cast<T>(result[v]);
+      }
+    }
+  }
+}
+
 template <typename T, const int group_size, int bits>
 METAL_FUNC void fp_qvm_impl(
     const device uint32_t* w,
@@ -1087,6 +1206,51 @@ template <typename T, const int group_size, int bits, bool batched>
       w, scales, x, y, in_vec_size, out_vec_size, tid, simd_gid, simd_lid);
 }
 
+template <
+    typename T,
+    int group_size,
+    int bits,
+    int vecs_per_tg,
+    int k_lanes,
+    bool batched>
+[[kernel]] void fp_qmv_wide(
+    const device uint32_t* w,
+    const device uint8_t* scales,
+    const device T* x,
+    device T* y,
+    const constant int& in_vec_size,
+    const constant int& out_vec_size,
+    const constant int& M,
+    const constant int& x_batch_ndims,
+    const constant int* x_shape,
+    const constant int64_t* x_strides,
+    const constant int& w_batch_ndims,
+    const constant int* w_shape,
+    const constant int64_t* w_strides,
+    const constant int64_t* s_strides,
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  if (batched) {
+    adjust_matrix_offsets(
+        x,
+        w,
+        scales,
+        y,
+        out_vec_size * M,
+        x_batch_ndims,
+        x_shape,
+        x_strides,
+        w_batch_ndims,
+        w_shape,
+        w_strides,
+        s_strides,
+        tid);
+  }
+  fp_qmv_wide_impl<T, group_size, bits, vecs_per_tg, k_lanes>(
+      w, scales, x, y, in_vec_size, out_vec_size, M, tid, simd_gid, simd_lid);
+}
+
 template <typename T, const int group_size, int bits, bool batched>
 [[kernel]] void fp_qvm(
     const device uint32_t* w,

mlx/backend/metal/kernels/fp_quantized.metal

@@ -52,6 +52,17 @@
       D,       \
       batched)
 
+#define instantiate_quantized_wide(mode, name, type, vecs_per_tg, k_lanes, group_size, bits, batched) \
+  instantiate_kernel( \
+      #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_nv_" #vecs_per_tg "_kl_" #k_lanes "_batch_" #batched, \
+      fp_ ## name,    \
+      type,    \
+      group_size,      \
+      bits,       \
+      vecs_per_tg,       \
+      k_lanes,       \
+      batched)
+
 #define instantiate_quantized_split_k(mode, name, type, split_k, group_size, bits) \
   instantiate_kernel( \
       #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_spk_" #split_k, \
@@ -105,6 +116,17 @@
   instantiate_quantized_quad(mode, qmv_quad, type, 128, 1, group_size, bits) \
   instantiate_quantized_quad(mode, qmv_quad, type, 128, 0, group_size, bits)
 
+// vecs_per_tg (input-vector tile) 2..5; the fp path uses k_lanes=16.
+#define instantiate_quantized_wide_wrap(mode, name, type, vecs_per_tg, k_lanes, group_size, bits) \
+  instantiate_quantized_wide(mode, name, type, vecs_per_tg, k_lanes, group_size, bits, 0)         \
+  instantiate_quantized_wide(mode, name, type, vecs_per_tg, k_lanes, group_size, bits, 1)
+
+#define instantiate_quantized_all_wide(type, mode, group_size, bits) \
+  instantiate_quantized_wide_wrap(mode, qmv_wide, type, 2, 16, group_size, bits) \
+  instantiate_quantized_wide_wrap(mode, qmv_wide, type, 3, 16, group_size, bits) \
+  instantiate_quantized_wide_wrap(mode, qmv_wide, type, 4, 16, group_size, bits) \
+  instantiate_quantized_wide_wrap(mode, qmv_wide, type, 5, 16, group_size, bits)
+
 #define instantiate_quantized_all_splitk(type, mode, group_size, bits) \
   instantiate_quantized_split_k(mode, qvm_split_k, type, 8, group_size, bits) \
   instantiate_quantized_split_k(mode, qvm_split_k, type, 32, group_size, bits) \
@@ -139,6 +161,7 @@
   instantiate_quantized_all_batched(type, mode, group_size, bits) \
   instantiate_quantized_all_single(type, mode, group_size, bits)  \
   instantiate_quantized_all_quad(type, mode, group_size, bits)    \
+  instantiate_quantized_all_wide(type, mode, group_size, bits)    \
   instantiate_quantized_all_splitk(type, mode, group_size, bits)  \
   instantiate_quantized_all_aligned(type, mode, group_size, bits) \
   instantiate_quantized_all_rhs(type, mode, group_size, bits)     \