From 20a24eaea536510c19ca01b78ace3c0207124027 Mon Sep 17 00:00:00 2001 From: Alin Huzmezan <32289581+ahuzmeza@users.noreply.github.com> Date: Wed, 28 Oct 2026 21:59:54 +0200 Subject: [PATCH 1/3] [Metal] Add affine qmv_mma for 3-16 rows --- mlx/backend/metal/kernels/quantized.h | 183 ++++++++++++++++++++++ mlx/backend/metal/kernels/quantized.metal | 32 ++++ mlx/backend/metal/quantized.cpp | 46 ++++-- 3 files changed, 250 insertions(+), 11 deletions(-) diff --git a/mlx/backend/metal/kernels/quantized.h b/mlx/backend/metal/kernels/quantized.h index 63fffb1d69..dadafb7cc3 100644 --- a/mlx/backend/metal/kernels/quantized.h +++ b/mlx/backend/metal/kernels/quantized.h @@ -1916,6 +1916,189 @@ template < simd_lid); } +// Few rows of x times transposed affine weights, as fp_qmv_mma. 4 and 8 bits +// only: a lane's 32 values are nibble or byte aligned and decode through half. +template +[[kernel]] void affine_qmv_mma( + const device uint32_t* w, + const device T* scales, + const device T* biases, + const device T* x, + device T* y, + const constant int& K, + const constant int& N, + const constant int& M, + uint3 tid [[threadgroup_position_in_grid]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + static_assert(BN % 8 == 0, "BN should be divisible by 8"); + static_assert(TM == 1 || TM == 2, "qmv_mma takes at most 16 rows"); + static_assert( + (bits == 4 || bits == 8) && + (group_size == 32 || group_size == 64 || group_size == 128), + "qmv_mma reads 4 and 8 bits at group sizes 32, 64 and 128"); + + using Frag = mlx::steel::BaseMMAFrag; + using frag_t = typename Frag::frag_type; + constexpr int BK = 128; + constexpr short TN = BN / 8; + constexpr int pack_factor = 8 / bits; + constexpr short loads = bits / 4; + + const short2 coord = Frag::get_coord(simd_lid); + const short fn = coord.x; + const short fm = coord.y; + const int row_bytes = K / pack_factor; + const int row_scales = K / group_size; + + const int y_col = tid.x * BN; + auto wl = (const device uint8_t*)w + + (y_col + fm) * static_cast(row_bytes) + 16 * fn / pack_factor; + const int64_t s_off = + (y_col + fm) * static_cast(row_scales) + 16 * fn / group_size; + scales += s_off; + biases += s_off; + x += fn * static_cast(K) + 16 * fm; + + bool x_rows[TM][2]; + frag_t C[TN][TM]; + STEEL_PRAGMA_UNROLL + for (short mt = 0; mt < TM; mt++) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < 2; j++) { + x_rows[mt][j] = 8 * mt + fn + j < M; + } + STEEL_PRAGMA_UNROLL + for (short t = 0; t < TN; t++) { + C[t][mt] = frag_t(0); + } + } + + for (int k = simd_gid * BK; k < K; k += WK * BK) { + vec xs[TM][2][2]; + STEEL_PRAGMA_UNROLL + for (short mt = 0; mt < TM; mt++) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < 2; j++) { + const device T* xr = x + (8 * mt + j) * static_cast(K) + k; + STEEL_PRAGMA_UNROLL + for (short h = 0; h < 2; h++) { + xs[mt][j][h] = x_rows[mt][j] ? *(const device vec*)(xr + 8 * h) + : vec(0); + } + } + } + uint4 ws[TN][loads]; + float sc[TN]; + float bs[TN]; + STEEL_PRAGMA_UNROLL + for (short t = 0; t < TN; t++) { + const device uint8_t* wt = + wl + 8 * t * static_cast(row_bytes) + k / pack_factor; + STEEL_PRAGMA_UNROLL + for (short l = 0; l < loads; l++) { + ws[t][l] = *(const device uint4*)(wt + 16 * l); + } + const int64_t si = + 8 * t * static_cast(row_scales) + k / group_size; + sc[t] = float(scales[si]); + bs[t] = float(biases[si]); + } + + STEEL_PRAGMA_UNROLL + for (short g = 0; g < 4; g++) { + frag_t b[4][TM]; + STEEL_PRAGMA_UNROLL + for (short i = 0; i < 4; i++) { + const short f = 4 * g + i; + STEEL_PRAGMA_UNROLL + for (short mt = 0; mt < TM; mt++) { + b[i][mt] = frag_t( + float(xs[mt][0][f / 8][f % 8]), float(xs[mt][1][f / 8][f % 8])); + } + } + STEEL_PRAGMA_UNROLL + for (short t = 0; t < TN; t++) { + // p[i] holds weights 4g+i and 16+4g+i of the lane. With 0x6400 set + // each half reads as 1024 + q, and taking 1024 off again is exact. + uint p[4]; + if constexpr (bits == 4) { + const uint lo = ws[t][0][g / 2]; + const uint hi = ws[t][0][2 + g / 2]; + const uint q = (g % 2) == 0 ? (lo & 0xffff) | (hi << 16) + : (lo >> 16) | (hi & 0xffff0000); + STEEL_PRAGMA_UNROLL + for (short i = 0; i < 4; i++) { + p[i] = (q >> (4 * i)) & 0x000f000f; + } + } else { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < 4; i++) { + p[i] = ((ws[t][0][g] >> (8 * i)) & 0xff) | + (((ws[t][1][g] >> (8 * i)) & 0xff) << 16); + } + } + STEEL_PRAGMA_UNROLL + for (short i = 0; i < 4; i++) { + const half2 v = as_type(p[i] | 0x64006400) - half2(1024.0h); + frag_t a( + fma(float(v.x), sc[t], bs[t]), fma(float(v.y), sc[t], bs[t])); + STEEL_PRAGMA_UNROLL + for (short mt = 0; mt < TM; mt++) { + Frag::mma(C[t][mt], a, b[i][mt], C[t][mt]); + } + } + } + } + } + + constexpr int tile_elems = TN * TM * 2; + threadgroup float Ps[(WK - 1) * tile_elems * 32]; + if (simd_gid > 0) { + threadgroup float* p = Ps + (simd_gid - 1) * tile_elems * 32 + simd_lid; + STEEL_PRAGMA_UNROLL + for (short t = 0; t < TN; t++) { + STEEL_PRAGMA_UNROLL + for (short mt = 0; mt < TM; mt++) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < 2; j++) { + p[((t * TM + mt) * 2 + j) * 32] = C[t][mt][j]; + } + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (simd_gid == 0) { + for (int s = 0; s < WK - 1; s++) { + const threadgroup float* p = Ps + s * tile_elems * 32 + simd_lid; + STEEL_PRAGMA_UNROLL + for (short t = 0; t < TN; t++) { + STEEL_PRAGMA_UNROLL + for (short mt = 0; mt < TM; mt++) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < 2; j++) { + C[t][mt][j] += p[((t * TM + mt) * 2 + j) * 32]; + } + } + } + } + STEEL_PRAGMA_UNROLL + for (short mt = 0; mt < TM; mt++) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < 2; j++) { + if (x_rows[mt][j]) { + device T* yr = + y + (8 * mt + fn + j) * static_cast(N) + y_col + fm; + STEEL_PRAGMA_UNROLL + for (short t = 0; t < TN; t++) { + yr[8 * t] = static_cast(C[t][mt][j]); + } + } + } + } + } +} + template < typename T, int group_size, diff --git a/mlx/backend/metal/kernels/quantized.metal b/mlx/backend/metal/kernels/quantized.metal index 324185bf1b..d9d8e487ae 100644 --- a/mlx/backend/metal/kernels/quantized.metal +++ b/mlx/backend/metal/kernels/quantized.metal @@ -195,6 +195,35 @@ instantiate_qmv_wide_batched(affine_qmv_wide, type, group_size, bits, 7, 8) \ instantiate_qmv_wide_batched(affine_qmv_wide, type, group_size, bits, 8, 8) +/////////////////////////////////////////////////////////////////////////////// +// affine_qmv_mma +/////////////////////////////////////////////////////////////////////////////// + +#define instantiate_qmv_mma(name, type, group_size, bits, wk, bn, tm) \ + instantiate_kernel( \ + #name "_" #type "_gs_" #group_size "_b_" #bits "_wk_" #wk "_bn_" #bn "_tm_" #tm, \ + name, \ + type, \ + group_size, \ + bits, \ + wk, \ + bn, \ + tm) + +#define instantiate_qmv_mma_all(type, group_size, bits) \ + instantiate_qmv_mma(affine_qmv_mma, type, group_size, bits, 4, 16, 1) \ + instantiate_qmv_mma(affine_qmv_mma, type, group_size, bits, 4, 16, 2) \ + instantiate_qmv_mma(affine_qmv_mma, type, group_size, bits, 4, 32, 1) + +#define instantiate_qmv_mma_all_bits(type, group_size) \ + instantiate_qmv_mma_all(type, group_size, 4) \ + instantiate_qmv_mma_all(type, group_size, 8) + +#define instantiate_qmv_mma_all_groups(type) \ + instantiate_qmv_mma_all_bits(type, 32) \ + instantiate_qmv_mma_all_bits(type, 64) \ + instantiate_qmv_mma_all_bits(type, 128) + /////////////////////////////////////////////////////////////////////////////// // Instantiation /////////////////////////////////////////////////////////////////////////////// @@ -226,4 +255,7 @@ instantiate_quantized_all_groups(float) instantiate_quantized_all_groups(float16_t) instantiate_quantized_all_groups(bfloat16_t) +instantiate_qmv_mma_all_groups(float) +instantiate_qmv_mma_all_groups(float16_t) +instantiate_qmv_mma_all_groups(bfloat16_t) // clang-format on diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index 40ee53a9ea..b2f94cf297 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -641,7 +641,15 @@ inline bool use_qmv_mma( int N, int K, metal::Device& d) { - if (K % 128 != 0) { + bool nax = metal::is_nax_available(); + bool affine = mode == QuantizationMode::Affine; + // affine_qmv_mma beats qmv from 3 rows below gen 15 and qmv_wide from 4 + // rows above; the other paths from 6. + int min_rows = 6; + if (affine && !nax) { + min_rows = d.get_architecture_gen() >= 15 ? 4 : 3; + } + if (M < min_rows || M > 16 || K % 128 != 0) { return false; } // Require aligned vector load. @@ -650,12 +658,23 @@ inline bool use_qmv_mma( return false; } // Route for large K/N. - bool nax = metal::is_nax_available(); - if (mode == QuantizationMode::Affine) { - // Only qmv_nax reads affine weights, with 3 to 8 bits. - return nax && (bits >= 3 && bits != 7) && (group_size >= 32) && - (x.dtype() == float16 || x.dtype() == bfloat16) && (N >= 512) && - (K >= 1024) && (N % 32 == 0); + if (affine) { + if (nax) { + // qmv_nax reads affine weights with 3 to 8 bits. + return (bits >= 3 && bits != 7) && (group_size >= 32) && + (x.dtype() == float16 || x.dtype() == bfloat16) && (N >= 512) && + (K >= 1024) && (N % 32 == 0); + } + // Ultra keeps the vector kernels. + if (d.get_architecture().back() == 'd') { + return false; + } + // Only 4 and 8 bits are instantiated; affine pays off from a smaller N * K + // than fp. + return (bits == 4 || bits == 8) && + (group_size == 32 || group_size == 64 || group_size == 128) && + (N >= 512) && (K >= 512) && (int64_t(N) * K >= (int64_t(1) << 21)) && + (N % 16 == 0); } else { // fp quantization if (nax) { // Below 2048, nvfp4 with few rows is faster on qmv_wide, which runs for @@ -690,8 +709,14 @@ void qmv_mma( const Stream& s, const std::string& mode) { bool nax = metal::is_nax_available(); - int wk = nax || d.get_architecture_gen() >= 15 ? 8 : 4; - int bn = nax ? 32 : 16; + bool affine = mode == "affine"; + // Affine stays at wk 4 on gen-15+ too (wk 8 was no faster on M4). + int wk = nax || (!affine && d.get_architecture_gen() >= 15) ? 8 : 4; + int tm = M > 8 ? 2 : 1; + // Below gen 15, affine with one row tile is faster at 32 columns (M2 Pro). + bool affine_bn32 = + affine && tm == 1 && d.get_architecture_gen() < 15 && N % 32 == 0; + int bn = (nax || affine_bn32) ? 32 : 16; std::string type_string = get_type_string(x.dtype()); std::string kname; kname.reserve(64); @@ -713,7 +738,6 @@ void qmv_mma( kernel = get_qmm_nax_kernel_wrapped( d, kname, "qmv_nax", mode, type_string, group_size, bits, wk, bn); } else { - int tm = M > 8 ? 2 : 1; concatenate(kname, "_tm_", tm); kernel = get_quantized_kernel_wrapped( d, kname, "qmv_mma", mode, type_string, group_size, bits, wk, bn, tm); @@ -2184,7 +2208,7 @@ void QuantizedMatmul::eval_gpu(const std::vector& inputs, array& out) { auto mode = quantization_mode_to_string(mode_); // For few rows use qmv_mma when N/K is large enough. - if (transpose_ && non_batched && (M >= 6) && (M <= 16) && + if (transpose_ && non_batched && use_qmv_mma(x, w, scales, mode_, group_size_, bits_, M, N, K, d)) { qmv_mma(x, w, scales, biases, out, group_size_, bits_, M, N, K, d, s, mode); return; From 8b018e97a1a368244fb693bee3332ddc937c90b2 Mon Sep 17 00:00:00 2001 From: Alin Huzmezan <32289581+ahuzmeza@users.noreply.github.com> Date: Wed, 28 Oct 2026 21:59:54 +0200 Subject: [PATCH 2/3] [Metal] Test affine qmv_mma --- python/tests/test_quantized.py | 51 ++++++++++++++++++++++++++++++++++ 1 file changed, 51 insertions(+) diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index 3186aef956..b913163fa0 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -884,6 +884,57 @@ def rel_err(got, expected): out = mx.quantized_matmul(xv, w_q, sv, mode=mode) self.assertLess(rel_err(out, xv @ w_hat.T), tol) + @unittest.skipIf(not mx.metal.is_available(), "requires Metal") + def test_qmm_few_rows_affine(self): + if mx.default_device() == mx.cpu: + self.skipTest("Covers GPU kernels only") + mx.random.seed(0) + + def rel_err(got, expected): + d = mx.abs(got.astype(mx.float32) - expected.astype(mx.float32)) + return d.max() / mx.abs(expected.astype(mx.float32)).max() + + shapes = [ + (4224, 4128), + (4224, 4112), # N % 32 == 16 + (512, 4096), # smallest K: one K step per simdgroup + ] + dtypes = [mx.bfloat16, mx.float16, mx.float32] + for (K, N), dtype in product(shapes, dtypes): + tol = 1e-5 if dtype == mx.float32 else 3e-2 + # 16 extra rows so that w, scales and biases can be sliced + w = mx.random.normal((N + 16, K), dtype=dtype) + x = (mx.random.normal((19, K)) / K**0.5).astype(dtype) + # Rows around the qmv_mma window and its row tilings; offsets + # above 0 are row slices of x and of w (nonzero buffer offsets) + cases = [(M, 0, 0) for M in [3, 4, 8, 9, 16, 17]] + cases += [(8, 16, 2), (16, 16, 2)] + xo = (mx.random.normal((8 * K + 1,)) / K**0.5).astype(dtype) + for bits, gs in product([4, 8], [32, 64, 128]): + w_q, scales, biases = mx.quantize(w, group_size=gs, bits=bits) + w_hat = mx.dequantize(w_q, scales, biases, group_size=gs, bits=bits) + y_ref = x @ w_hat.T + for M, w0, x0 in cases: + rows = slice(w0, w0 + N) + with self.subTest(N=N, bits=bits, gs=gs, dtype=dtype, M=M, w0=w0): + out = mx.quantized_matmul( + x[x0 : x0 + M], + w_q[rows], + scales[rows], + biases[rows], + group_size=gs, + bits=bits, + ) + self.assertLess(rel_err(out, y_ref[x0 : x0 + M, rows]), tol) + + # A view at an odd offset must not reach the Metal vector loads + xv = xo[1:].reshape(8, K) + with self.subTest(N=N, dtype=dtype, offset="x"): + out = mx.quantized_matmul( + xv, w_q, scales, biases, group_size=gs, bits=bits + ) + self.assertLess(rel_err(out, xv @ w_hat.T), tol) + def test_qmm_vjp(self): key = mx.random.key(0) k1, k2 = mx.random.split(key) From 125e62c39fcb764059e85955e84709027d24fd67 Mon Sep 17 00:00:00 2001 From: Cheng Date: Mon, 2 Nov 2026 13:32:22 -0800 Subject: [PATCH 3/3] nit --- mlx/backend/metal/kernels/quantized.h | 2 - mlx/backend/metal/kernels/quantized.metal | 16 ++---- mlx/backend/metal/quantized.cpp | 52 ++++++++++++-------- python/tests/test_quantized.py | 60 +++++++++-------------- 4 files changed, 57 insertions(+), 73 deletions(-) diff --git a/mlx/backend/metal/kernels/quantized.h b/mlx/backend/metal/kernels/quantized.h index dadafb7cc3..5357bafbd4 100644 --- a/mlx/backend/metal/kernels/quantized.h +++ b/mlx/backend/metal/kernels/quantized.h @@ -1916,8 +1916,6 @@ template < simd_lid); } -// Few rows of x times transposed affine weights, as fp_qmv_mma. 4 and 8 bits -// only: a lane's 32 values are nibble or byte aligned and decode through half. template [[kernel]] void affine_qmv_mma( const device uint32_t* w, diff --git a/mlx/backend/metal/kernels/quantized.metal b/mlx/backend/metal/kernels/quantized.metal index d9d8e487ae..fa945c9b19 100644 --- a/mlx/backend/metal/kernels/quantized.metal +++ b/mlx/backend/metal/kernels/quantized.metal @@ -215,15 +215,6 @@ instantiate_qmv_mma(affine_qmv_mma, type, group_size, bits, 4, 16, 2) \ instantiate_qmv_mma(affine_qmv_mma, type, group_size, bits, 4, 32, 1) -#define instantiate_qmv_mma_all_bits(type, group_size) \ - instantiate_qmv_mma_all(type, group_size, 4) \ - instantiate_qmv_mma_all(type, group_size, 8) - -#define instantiate_qmv_mma_all_groups(type) \ - instantiate_qmv_mma_all_bits(type, 32) \ - instantiate_qmv_mma_all_bits(type, 64) \ - instantiate_qmv_mma_all_bits(type, 128) - /////////////////////////////////////////////////////////////////////////////// // Instantiation /////////////////////////////////////////////////////////////////////////////// @@ -245,7 +236,9 @@ instantiate_quantized_all(type, group_size, 4) \ instantiate_quantized_all(type, group_size, 5) \ instantiate_quantized_all(type, group_size, 6) \ - instantiate_quantized_all(type, group_size, 8) + instantiate_quantized_all(type, group_size, 8) \ + instantiate_qmv_mma_all(type, group_size, 4) \ + instantiate_qmv_mma_all(type, group_size, 8) #define instantiate_quantized_all_groups(type) \ instantiate_quantized_all_bits(type, 32) \ @@ -255,7 +248,4 @@ instantiate_quantized_all_groups(float) instantiate_quantized_all_groups(float16_t) instantiate_quantized_all_groups(bfloat16_t) -instantiate_qmv_mma_all_groups(float) -instantiate_qmv_mma_all_groups(float16_t) -instantiate_qmv_mma_all_groups(bfloat16_t) // clang-format on diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index b2f94cf297..c7213ae33c 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -649,7 +649,7 @@ inline bool use_qmv_mma( if (affine && !nax) { min_rows = d.get_architecture_gen() >= 15 ? 4 : 3; } - if (M < min_rows || M > 16 || K % 128 != 0) { + if ((M < min_rows) || (K % 128 != 0)) { return false; } // Require aligned vector load. @@ -661,20 +661,19 @@ inline bool use_qmv_mma( if (affine) { if (nax) { // qmv_nax reads affine weights with 3 to 8 bits. - return (bits >= 3 && bits != 7) && (group_size >= 32) && + return (bits >= 3 && bits != 7) && (x.dtype() == float16 || x.dtype() == bfloat16) && (N >= 512) && (K >= 1024) && (N % 32 == 0); + } else { // non-nax + // Ultra keeps the vector kernels. + if (d.get_architecture().back() == 'd') { + return false; + } + // Only 4 and 8 bits are instantiated; affine pays off from a smaller N * + // K than fp. + return (bits == 4 || bits == 8) && (N >= 512) && (K >= 512) && + (int64_t(N) * K >= (int64_t(1) << 21)) && (N % 16 == 0); } - // Ultra keeps the vector kernels. - if (d.get_architecture().back() == 'd') { - return false; - } - // Only 4 and 8 bits are instantiated; affine pays off from a smaller N * K - // than fp. - return (bits == 4 || bits == 8) && - (group_size == 32 || group_size == 64 || group_size == 128) && - (N >= 512) && (K >= 512) && (int64_t(N) * K >= (int64_t(1) << 21)) && - (N % 16 == 0); } else { // fp quantization if (nax) { // Below 2048, nvfp4 with few rows is faster on qmv_wide, which runs for @@ -709,14 +708,25 @@ void qmv_mma( const Stream& s, const std::string& mode) { bool nax = metal::is_nax_available(); - bool affine = mode == "affine"; - // Affine stays at wk 4 on gen-15+ too (wk 8 was no faster on M4). - int wk = nax || (!affine && d.get_architecture_gen() >= 15) ? 8 : 4; - int tm = M > 8 ? 2 : 1; - // Below gen 15, affine with one row tile is faster at 32 columns (M2 Pro). - bool affine_bn32 = - affine && tm == 1 && d.get_architecture_gen() < 15 && N % 32 == 0; - int bn = (nax || affine_bn32) ? 32 : 16; + bool affine = (mode == "affine"); + int wk = 4; + int bn = 16; + int tm = (M > 8) ? 2 : 1; + if (nax) { + wk = 8; + bn = 32; + } else { // non-nax + // Affine stays at wk 4 on gen-15+ too (wk 8 was no faster on M4). + if (!affine && (d.get_architecture_gen() >= 15)) { + wk = 8; + } + // Below gen 15, affine with one row tile is faster at 32 columns (M2 Pro). + if (affine && (d.get_architecture_gen() < 15) && (tm == 1) && + (N % 32 == 0)) { + bn = 32; + } + } + std::string type_string = get_type_string(x.dtype()); std::string kname; kname.reserve(64); @@ -2208,7 +2218,7 @@ void QuantizedMatmul::eval_gpu(const std::vector& inputs, array& out) { auto mode = quantization_mode_to_string(mode_); // For few rows use qmv_mma when N/K is large enough. - if (transpose_ && non_batched && + if (transpose_ && non_batched && (M <= 16) && use_qmv_mma(x, w, scales, mode_, group_size_, bits_, M, N, K, d)) { qmv_mma(x, w, scales, biases, out, group_size_, bits_, M, N, K, d, s, mode); return; diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index b913163fa0..39d794a190 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -849,7 +849,7 @@ def rel_err(got, expected): self.assertLess(rel_err(out, x @ w_hat.swapaxes(-1, -2)), 3e-2) @unittest.skipIf(not mx.metal.is_available(), "requires Metal") - def test_qmm_few_rows(self): + def test_fp_qmm_few_rows(self): if mx.default_device() == mx.cpu: self.skipTest("Covers GPU kernels only") mx.random.seed(0) @@ -885,7 +885,7 @@ def rel_err(got, expected): self.assertLess(rel_err(out, xv @ w_hat.T), tol) @unittest.skipIf(not mx.metal.is_available(), "requires Metal") - def test_qmm_few_rows_affine(self): + def test_affine_qmm_few_rows(self): if mx.default_device() == mx.cpu: self.skipTest("Covers GPU kernels only") mx.random.seed(0) @@ -894,46 +894,32 @@ def rel_err(got, expected): d = mx.abs(got.astype(mx.float32) - expected.astype(mx.float32)) return d.max() / mx.abs(expected.astype(mx.float32)).max() - shapes = [ - (4224, 4128), - (4224, 4112), # N % 32 == 16 - (512, 4096), # smallest K: one K step per simdgroup - ] - dtypes = [mx.bfloat16, mx.float16, mx.float32] - for (K, N), dtype in product(shapes, dtypes): - tol = 1e-5 if dtype == mx.float32 else 3e-2 - # 16 extra rows so that w, scales and biases can be sliced - w = mx.random.normal((N + 16, K), dtype=dtype) - x = (mx.random.normal((19, K)) / K**0.5).astype(dtype) - # Rows around the qmv_mma window and its row tilings; offsets - # above 0 are row slices of x and of w (nonzero buffer offsets) - cases = [(M, 0, 0) for M in [3, 4, 8, 9, 16, 17]] - cases += [(8, 16, 2), (16, 16, 2)] - xo = (mx.random.normal((8 * K + 1,)) / K**0.5).astype(dtype) - for bits, gs in product([4, 8], [32, 64, 128]): + for (K, N), M, bits, gs, dtype in product( + [ + (4224, 4128), + (512, 4096), # smallest K: one K step per simdgroup + ], + [3, 8, 9], + [4, 8], + [32, 64, 128], + [mx.bfloat16, mx.float16, mx.float32], + ): + with self.subTest(M=M, N=N, K=K, bits=bits, gs=gs, dtype=dtype): + tol = 1e-5 if dtype == mx.float32 else 3e-2 + x = (mx.random.normal((M, K)) / K**0.5).astype(dtype) + w = mx.random.normal((N, K), dtype=dtype) w_q, scales, biases = mx.quantize(w, group_size=gs, bits=bits) w_hat = mx.dequantize(w_q, scales, biases, group_size=gs, bits=bits) y_ref = x @ w_hat.T - for M, w0, x0 in cases: - rows = slice(w0, w0 + N) - with self.subTest(N=N, bits=bits, gs=gs, dtype=dtype, M=M, w0=w0): - out = mx.quantized_matmul( - x[x0 : x0 + M], - w_q[rows], - scales[rows], - biases[rows], - group_size=gs, - bits=bits, - ) - self.assertLess(rel_err(out, y_ref[x0 : x0 + M, rows]), tol) - - # A view at an odd offset must not reach the Metal vector loads - xv = xo[1:].reshape(8, K) - with self.subTest(N=N, dtype=dtype, offset="x"): out = mx.quantized_matmul( - xv, w_q, scales, biases, group_size=gs, bits=bits + x, + w_q, + scales, + biases, + group_size=gs, + bits=bits, ) - self.assertLess(rel_err(out, xv @ w_hat.T), tol) + self.assertLess(rel_err(out, y_ref), tol) def test_qmm_vjp(self): key = mx.random.key(0)