perf(metal): vectorize iq4_nl and iq4_xs dequant loads

CHANGELOG.md

@@ -11,6 +11,15 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
   copies the RoPE tail through unchanged. Fuses the DeepSeek-V4 compressor
   emit-path quantization, which has no f16 cache step, into one dispatch.
 
+### Changed
+- IQ4_NL decode is faster, because each mat-vec lane now reads eight weights
+  instead of one. IQ4_NL no longer trails the other 4-bit codecs.
+- IQ4_XS decode is faster, because each mat-vec lane now reads a quant byte
+  once and uses both nibbles instead of dropping one.
+- IQ4_XS prefill is faster, because the tensor-core loader now reads the
+  quant bytes with vector loads and uses both nibbles of each byte. IQ4_XS
+  prefill no longer trails the other 4-bit codecs.
+
 ## [0.3.11]
 
 ### Changed

metal/mlx/backend/metal/kernels/kq_quantized_iq.h

@@ -1058,33 +1058,45 @@ METAL_FUNC void kq_iq4_xs_qmv_impl(
   const int nb = in_vec_size / KQ_IQ4_XS_SUPERBLOCK;
   x += tid.x * in_vec_size;
   y += tid.x * out_vec_size;
+  // A byte holds two weights of the sub-block. The low nibble gives weight b
+  // and the high nibble gives weight b + 16. Each lane takes four bytes and
+  // uses both nibbles, so it still covers eight weights but reads each byte
+  // once. The old shape read a byte for each weight and dropped one nibble,
+  // which doubled the load count.
+  //
+  // A block is 136 bytes and the quants start at byte 8. Both divide by four,
+  // so one uint load is safe for every block.
+  constexpr int bytes_per_lane = vpt / 2;
   const int s = simd_lid / 4; // sub-block
-  const int o = (simd_lid % 4) * 8; // offset within sub-block
-  const bool is_high = o >= 16;
-  const int byte0 = is_high ? (o - 16) : o;
+  const int q0 = int(simd_lid % 4) * bytes_per_lane; // first byte of the lane
   U result[results_per_simdgroup] = {0};
   for (int ib = 0; ib < nb; ib++) {
-    U xt[vpt];
+    U xlo[bytes_per_lane];
+    U xhi[bytes_per_lane];
+    const device T* xb = x + ib * KQ_IQ4_XS_SUPERBLOCK + s * 32 + q0;
 #pragma unroll
-    for (int i = 0; i < vpt; i++) {
-      xt[i] = U(x[ib * KQ_IQ4_XS_SUPERBLOCK + simd_lid * vpt + i]);
+    for (int i = 0; i < bytes_per_lane; i++) {
+      xlo[i] = U(xb[i]);
+      xhi[i] = U(xb[16 + i]);
     }
     for (int row = 0; row < active_rows; row++) {
       const device uint8_t* sb = w +
           static_cast<int64_t>(out_row + row) * row_bytes +
           ib * KQ_IQ4_XS_BLOCK_BYTES;
       const U d = U(float(*(const device half*)sb));
-      const uint16_t scales_h = uint16_t(sb[2]) | (uint16_t(sb[3]) << 8);
+      const uint scales_h = uint(*(const device ushort*)(sb + 2));
       const device uint8_t* scales_l = sb + KQ_IQ4_XS_SCALESL_OFFSET;
       const int ls = ((scales_l[s / 2] >> (4 * (s & 1))) & 0xf) |
           (((scales_h >> (2 * s)) & 3) << 4);
       const U dl = d * U(ls - 32);
-      const device uint8_t* qs = sb + KQ_IQ4_XS_QS_OFFSET + s * 16 + byte0;
+      const uint qb = *reinterpret_cast<const device uint*>(
+          sb + KQ_IQ4_XS_QS_OFFSET + s * 16 + q0);
       U partial = 0;
 #pragma unroll
-      for (int i = 0; i < vpt; i++) {
-        const int nib = is_high ? (qs[i] >> 4) : (qs[i] & 0xf);
-        partial += xt[i] * U(kvalues_iq4nl[nib]);
+      for (int i = 0; i < bytes_per_lane; i++) {
+        const uint b = (qb >> (8 * i)) & 0xFF;
+        partial += xlo[i] * U(kvalues_iq4nl[b & 0x0F]);
+        partial += xhi[i] * U(kvalues_iq4nl[b >> 4]);
       }
       result[row] += dl * partial;
     }

metal/mlx/backend/metal/kernels/kq_quantized_legacy.h

@@ -2367,9 +2367,24 @@ template <typename T, short r1ptg, short nsg, short nxpsg>
       w, x, y, in_vec_size, out_vec_size, tgpig, tiisg, sgitg);
 }
 
-// IQ4_NL mat-vec: one impl for both qmv and qmv_fast (the LUT decode is cheap,
-// so there's no separate aligned fast path). Lane `simd_lid` owns weight `lane`
-// of every 32-block; simd_sum reduces the 32 lanes.
+// IQ4_NL mat-vec: one impl for both qmv and qmv_fast. The LUT decode is cheap,
+// so there is no separate aligned fast path.
+//
+// Four lanes cover one 32-weight block, and a simdgroup covers eight blocks per
+// step. Each lane owns four qs bytes and reads both nibbles of each byte, which
+// gives the lane eight weights. simd_sum then reduces the 32 lanes.
+//
+// The earlier version gave each lane one weight. That version read one byte and
+// one scale per lane per row-block, which is 64 load instructions for each
+// 32-weight block. This version reads two ushorts and one scale for each group
+// of four lanes, which is 12 load instructions for the same block. The kernel
+// is instruction-bound, not bandwidth-bound, so the load count sets the speed.
+// The iq4_xs kernel in kq_quantized_iq.h uses the same shape.
+//
+// Blocks are 18 bytes and qs starts at byte 2, so every qs address is a
+// multiple of 2. Therefore the lane can use ushort loads. It cannot use uint
+// loads, because an 18-byte block makes every second block address odd for a
+// 4-byte type.
 template <typename T, int group_size, int bits, int results_per_simdgroup = 4>
 METAL_FUNC void kq_iq4_nl_qmv_impl(
     const device uint8_t* w,
@@ -2383,6 +2398,9 @@ METAL_FUNC void kq_iq4_nl_qmv_impl(
   static_assert(group_size == KQ_IQ4_NL_GROUP, "IQ4_NL requires gs=32");
   static_assert(bits == 4, "IQ4_NL requires bits=4");
   constexpr int num_simdgroups = 2;
+  constexpr int lanes_per_block = 4; // four lanes cover one 32-weight block
+  constexpr int blocks_per_step = 32 / lanes_per_block; // eight blocks a step
+  constexpr int bytes_per_lane = KQ_IQ4_NL_GROUP / 2 / lanes_per_block; // four
   typedef float U;
   const int out_row = tid.y * (num_simdgroups * results_per_simdgroup) +
       simd_gid * results_per_simdgroup;
@@ -2394,19 +2412,46 @@ METAL_FUNC void kq_iq4_nl_qmv_impl(
   const int nb = in_vec_size / KQ_IQ4_NL_GROUP;
   x += tid.x * in_vec_size;
   y += tid.x * out_vec_size;
-  const bool is_high = simd_lid >= 16;
-  const int byteidx = is_high ? int(simd_lid) - 16 : int(simd_lid);
+  // The lane picks its block within the step, then its byte group in that
+  // block. Byte `byte0 + k` holds weight `byte0 + k` in the low nibble and
+  // weight `byte0 + k + 16` in the high nibble.
+  const int lane_blk = int(simd_lid) / lanes_per_block; // 0..7
+  const int byte0 = (int(simd_lid) % lanes_per_block) * bytes_per_lane;
   U result[results_per_simdgroup] = {0};
-  for (int ib = 0; ib < nb; ib++) {
-    const U xv = U(x[ib * KQ_IQ4_NL_GROUP + simd_lid]);
+  for (int base = 0; base < nb; base += blocks_per_step) {
+    const int ib = base + lane_blk;
+    // The last step can run past the end. Those lanes add nothing. Every lane
+    // still reaches the simd_sum below, because `base` is the same in all
+    // lanes.
+    if (ib >= nb) {
+      continue;
+    }
+    // Hold the lane's eight x values in registers. All rows below reuse them.
+    U xt[2 * bytes_per_lane];
+    const device T* xb = x + ib * KQ_IQ4_NL_GROUP + byte0;
+#pragma unroll
+    for (int k = 0; k < bytes_per_lane; k++) {
+      xt[k] = U(xb[k]);
+      xt[bytes_per_lane + k] = U(xb[KQ_IQ4_NL_GROUP / 2 + k]);
+    }
     for (int row = 0; row < active_rows; row++) {
       const device uint8_t* blk = w +
           static_cast<int64_t>(out_row + row) * row_bytes +
           ib * KQ_IQ4_NL_BLOCK_BYTES;
       const U d = U(float(*(const device half*)blk));
-      const uint8_t b = blk[KQ_IQ4_NL_QS_OFFSET + byteidx];
-      const int nib = is_high ? (b >> 4) : (b & 0x0F);
-      result[row] += d * U(kvalues_iq4nl[nib]) * xv;
+      const device ushort* qw = reinterpret_cast<const device ushort*>(
+          blk + KQ_IQ4_NL_QS_OFFSET + byte0);
+      const uint qbytes = uint(qw[0]) | (uint(qw[1]) << 16);
+      // One scale multiply for the whole group. The earlier version scaled
+      // every weight, so this version also rounds less.
+      U partial = 0;
+#pragma unroll
+      for (int k = 0; k < bytes_per_lane; k++) {
+        const uint b = (qbytes >> (8 * k)) & 0xFF;
+        partial += U(kvalues_iq4nl[b & 0x0F]) * xt[k];
+        partial += U(kvalues_iq4nl[b >> 4]) * xt[bytes_per_lane + k];
+      }
+      result[row] += d * partial;
     }
   }
   for (int row = 0; row < results_per_simdgroup; row++) {

metal/mlx/backend/metal/kernels/kq_quantized_nax.h

@@ -2527,14 +2527,27 @@ struct KqNaxIq4_xsBlockLoader {
     const int ls = ((scales_l[sb / 2] >> (4 * (sb & 1))) & 0xf) |
         (((scales_h >> (2 * sb)) & 3) << 4);
     const float dl = d * float(ls - 32);
-    const device uint8_t* qs = src + KQ_IQ4_XS_QS_OFFSET + sb * 16;
+    // A byte holds two weights of the sub-block. The low nibble gives
+    // weight b and the high nibble gives weight b + 16. Read the 16 quant
+    // bytes with two uint2 loads and use both nibbles, instead of one byte
+    // load per weight that drops a nibble. The quants of sub-block sb start
+    // at byte 8 + 16 * sb of a 136-byte block. Both terms divide by 8, so
+    // the uint2 loads are aligned.
+    const device uint2* qw = reinterpret_cast<const device uint2*>(
+        src + KQ_IQ4_XS_QS_OFFSET + sb * 16);
+    const uint2 q01 = qw[0];
+    const uint2 q23 = qw[1];
+    const uint qword[4] = {q01.x, q01.y, q23.x, q23.y};
 #pragma unroll
-    for (short i = 0; i < n_reads; i++) {
-      const int p = i;
-      const bool is_high = p >= 16;
-      const int b = is_high ? (p - 16) : p;
-      const int nib = is_high ? (qs[b] >> 4) : (qs[b] & 0xf);
-      dst[i] = T(dl * float(kvalues_iq4nl[nib]));
+    for (short v = 0; v < 4; v++) {
+      const uint q = qword[v];
+#pragma unroll
+      for (short k = 0; k < 4; k++) {
+        const short i = 4 * v + k;
+        const uint b = (q >> (8 * k)) & 0xFFu;
+        dst[i] = T(dl * float(kvalues_iq4nl[b & 0x0Fu]));
+        dst[16 + i] = T(dl * float(kvalues_iq4nl[b >> 4]));
+      }
     }
   }
 

tests/test_matmul_synth.py

@@ -53,7 +53,10 @@
 # ggml marks these imatrix-required; kq.quantize rejects them without one.
 REQ_IMAT = {"iq2_xxs", "iq2_xs", "iq1_s"}
 N, K = 256, 512
-MS = (1, 2, 3, 4, 8)
+# 1 takes qmv, 2 through 8 take the verify mv_ext kernel, and 64 takes
+# qmm. Without a large M the block loaders never run, so the prefill
+# path of every codec goes untested.
+MS = (1, 2, 3, 4, 8, 64)
 
 
 @pytest.mark.parametrize("codec", list(CODECS))
@@ -84,3 +87,39 @@ def test_matmul_synth(codec):
         assert max_rel < 5e-2 or max_abs < 1e-2, (
             f"{codec} M={M}: max_rel={max_rel:.3e} max_abs={max_abs:.3e}"
         )
+
+
+# Codecs with a 32-weight block. They accept any K that 32 divides, so they
+# reach the plain qmv path. The superblock codecs need K that 256 divides,
+# which always takes the fast path.
+BLOCK32 = ("q4_0", "q4_1", "q5_0", "q5_1", "q8_0", "iq4_nl")
+# 288 is 9 blocks of 32. It is not a multiple of qmv_fast_k_align(), so the
+# kernel must guard the tail. K=512 above only covers the fast path.
+K_TAIL = 288
+
+
+@pytest.mark.parametrize("codec", BLOCK32)
+def test_matmul_synth_tail(codec):
+    """Validate the plain qmv path, where the block count leaves a tail."""
+    gtype = CODECS[codec]
+    rng = np.random.default_rng(0)
+    w_np = (rng.standard_normal((N, K_TAIL)) * 0.1).astype(np.float32)
+    wq, _ = kq.quantize(mx.array(w_np), codec)
+    mx.eval(wq)
+    packed = np.ascontiguousarray(np.array(wq).astype(np.uint8))
+    w = mx.array(packed)
+    scales = mx.zeros((1,), dtype=mx.uint8)
+    deq = quants.dequantize(packed, gtype).astype(np.float32)
+    for M in MS:
+        x_np = (rng.standard_normal((M, K_TAIL)) * 0.1).astype(np.float32)
+        x = mx.array(x_np).astype(mx.float16)
+        got = kq.quantized_matmul(x, w, scales, codec, transpose=True)
+        mx.eval(got)
+        g = np.array(got).astype(np.float32)
+        r = np.array(x).astype(np.float32) @ deq.T
+        diff = np.abs(g - r)
+        max_abs = float(diff.max())
+        max_rel = float((diff / (np.abs(r) + 1e-3)).max())
+        assert max_rel < 5e-2 or max_abs < 1e-2, (
+            f"{codec} M={M}: max_rel={max_rel:.3e} max_abs={max_abs:.3e}"
+        )