From d6870c1329f6fc5f49df12071861808501777a1b Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Wed, 12 Aug 2026 12:16:57 -0700 Subject: [PATCH 1/5] perf(metal): faster iq4_nl mat-vec at decode shapes --- CHANGELOG.md | 4 ++ .../metal/kernels/kq_quantized_legacy.h | 65 ++++++++++++++++--- 2 files changed, 59 insertions(+), 10 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d476007..3ba864e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,10 @@ 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. + ## [0.3.11] ### Changed diff --git a/metal/mlx/backend/metal/kernels/kq_quantized_legacy.h b/metal/mlx/backend/metal/kernels/kq_quantized_legacy.h index f667ebe..3980977 100644 --- a/metal/mlx/backend/metal/kernels/kq_quantized_legacy.h +++ b/metal/mlx/backend/metal/kernels/kq_quantized_legacy.h @@ -2367,9 +2367,24 @@ template 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 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(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( + 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++) { From d73e4cb087642046285b0ccb1b0c2761f21be02a Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Wed, 12 Aug 2026 12:16:57 -0700 Subject: [PATCH 2/5] test(matmul): cover the plain qmv path for 32-block codecs --- tests/test_matmul_synth.py | 36 ++++++++++++++++++++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/tests/test_matmul_synth.py b/tests/test_matmul_synth.py index 8fbb4ef..76c10a1 100644 --- a/tests/test_matmul_synth.py +++ b/tests/test_matmul_synth.py @@ -84,3 +84,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}" + ) From e28592c1c1364bc5a38b4478e862ec6286bb7cd2 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Wed, 12 Aug 2026 12:26:07 -0700 Subject: [PATCH 3/5] perf(metal): faster iq4_xs mat-vec at decode shapes --- CHANGELOG.md | 2 ++ .../backend/metal/kernels/kq_quantized_iq.h | 34 +++++++++++++------ 2 files changed, 25 insertions(+), 11 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3ba864e..e8e81c4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,8 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ### 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. ## [0.3.11] diff --git a/metal/mlx/backend/metal/kernels/kq_quantized_iq.h b/metal/mlx/backend/metal/kernels/kq_quantized_iq.h index 383b7f9..7c301d4 100644 --- a/metal/mlx/backend/metal/kernels/kq_quantized_iq.h +++ b/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(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( + 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; } From 1d6af737631b0eb5b1c1bd36bd91595d47f961a4 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Wed, 12 Aug 2026 12:38:46 -0700 Subject: [PATCH 4/5] test(matmul): exercise the qmm path in the synthetic codec sweep --- tests/test_matmul_synth.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/test_matmul_synth.py b/tests/test_matmul_synth.py index 76c10a1..8139eaa 100644 --- a/tests/test_matmul_synth.py +++ b/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)) From 37221cd6f2b7b1259ed9140074974e9d83b0746b Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Wed, 12 Aug 2026 13:03:14 -0700 Subject: [PATCH 5/5] perf(metal): faster iq4_xs dequant in the nax prefill loader --- CHANGELOG.md | 3 +++ .../backend/metal/kernels/kq_quantized_nax.h | 27 ++++++++++++++----- 2 files changed, 23 insertions(+), 7 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e8e81c4..b6e187e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,6 +16,9 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). 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] diff --git a/metal/mlx/backend/metal/kernels/kq_quantized_nax.h b/metal/mlx/backend/metal/kernels/kq_quantized_nax.h index 2ce294a..6b2555b 100644 --- a/metal/mlx/backend/metal/kernels/kq_quantized_nax.h +++ b/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( + 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])); + } } }