From 77ab49475e1a43829882a1dada35e69c809c57c3 Mon Sep 17 00:00:00 2001 From: Yanzhao Wang Date: Tue, 15 Sep 2026 00:23:20 -0700 Subject: [PATCH 1/2] Accumulate affine QMV input sums in float load_vector and load_vector_safe add the inputs of each sub-byte chunk in the input type before widening the sum. With bfloat16 or float16 inputs that intermediate rounds, so the bias contribution bias * sum(x) is off even when the exact sum is representable in float. Widen each input before adding it, for every sub-byte width (2, 3, 4, 5 and 6 bits); the 8-bit branch already accumulates one value at a time in float. quantized_nax.h carries unused copies of both functions and gets the same change so the two headers stay in sync. The test builds inputs whose running sum is not representable in the half precision types and checks the exact result for several output sizes, input sizes, group sizes, bit widths and dtypes. --- mlx/backend/metal/kernels/quantized.h | 28 ++++++++++----------- mlx/backend/metal/kernels/quantized_nax.h | 28 ++++++++++----------- python/tests/test_quantized.py | 30 +++++++++++++++++++++++ 3 files changed, 58 insertions(+), 28 deletions(-) diff --git a/mlx/backend/metal/kernels/quantized.h b/mlx/backend/metal/kernels/quantized.h index 9d9ce368c9..5deaad811c 100644 --- a/mlx/backend/metal/kernels/quantized.h +++ b/mlx/backend/metal/kernels/quantized.h @@ -36,7 +36,7 @@ inline U load_vector(const device T* x, thread U* x_thread) { if (bits == 2) { for (int i = 0; i < values_per_thread; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 4.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -46,8 +46,8 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 3) { for (int i = 0; i < values_per_thread; i += 8) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] + - x[i + 6] + x[i + 7]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 8.0f; x_thread[i + 2] = x[i + 2] / 64.0f; @@ -61,7 +61,7 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 4) { for (int i = 0; i < values_per_thread; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 16.0f; x_thread[i + 2] = x[i + 2] / 256.0f; @@ -71,8 +71,8 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 5) { for (int i = 0; i < values_per_thread; i += 8) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] + - x[i + 6] + x[i + 7]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 32.0f; x_thread[i + 2] = x[i + 2] / 4.0f; @@ -86,7 +86,7 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 6) { for (int i = 0; i < values_per_thread; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 64.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -115,7 +115,7 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { if (bits == 2) { for (int i = 0; i < N; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 4.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -125,8 +125,8 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 3) { for (int i = 0; i < N; i += 8) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] + - x[i + 6] + x[i + 7]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 8.0f; @@ -141,7 +141,7 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 4) { for (int i = 0; i < N; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 16.0f; x_thread[i + 2] = x[i + 2] / 256.0f; @@ -151,8 +151,8 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 5) { for (int i = 0; i < N; i += 8) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] + - x[i + 6] + x[i + 7]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 32.0f; x_thread[i + 2] = x[i + 2] / 4.0f; @@ -166,7 +166,7 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 6) { for (int i = 0; i < N; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 64.0f; x_thread[i + 2] = x[i + 2] / 16.0f; diff --git a/mlx/backend/metal/kernels/quantized_nax.h b/mlx/backend/metal/kernels/quantized_nax.h index 54a37482b6..296390d5b1 100644 --- a/mlx/backend/metal/kernels/quantized_nax.h +++ b/mlx/backend/metal/kernels/quantized_nax.h @@ -39,7 +39,7 @@ inline U load_vector(const device T* x, thread U* x_thread) { if (bits == 2) { for (int i = 0; i < values_per_thread; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 4.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -49,8 +49,8 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 3) { for (int i = 0; i < values_per_thread; i += 8) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] + - x[i + 6] + x[i + 7]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 8.0f; x_thread[i + 2] = x[i + 2] / 64.0f; @@ -64,7 +64,7 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 4) { for (int i = 0; i < values_per_thread; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 16.0f; x_thread[i + 2] = x[i + 2] / 256.0f; @@ -74,8 +74,8 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 5) { for (int i = 0; i < values_per_thread; i += 8) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] + - x[i + 6] + x[i + 7]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 32.0f; x_thread[i + 2] = x[i + 2] / 4.0f; @@ -89,7 +89,7 @@ inline U load_vector(const device T* x, thread U* x_thread) { else if (bits == 6) { for (int i = 0; i < values_per_thread; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 64.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -118,7 +118,7 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { if (bits == 2) { for (int i = 0; i < N; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 4.0f; x_thread[i + 2] = x[i + 2] / 16.0f; @@ -128,8 +128,8 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 3) { for (int i = 0; i < N; i += 8) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] + - x[i + 6] + x[i + 7]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 8.0f; @@ -144,7 +144,7 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 4) { for (int i = 0; i < N; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 16.0f; x_thread[i + 2] = x[i + 2] / 256.0f; @@ -154,8 +154,8 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 5) { for (int i = 0; i < N; i += 8) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3] + x[i + 4] + x[i + 5] + - x[i + 6] + x[i + 7]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]) + U(x[i + 4]) + + U(x[i + 5]) + U(x[i + 6]) + U(x[i + 7]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 32.0f; x_thread[i + 2] = x[i + 2] / 4.0f; @@ -169,7 +169,7 @@ inline U load_vector_safe(const device T* x, thread U* x_thread, int N) { else if (bits == 6) { for (int i = 0; i < N; i += 4) { - sum += x[i] + x[i + 1] + x[i + 2] + x[i + 3]; + sum += U(x[i]) + U(x[i + 1]) + U(x[i + 2]) + U(x[i + 3]); x_thread[i] = x[i]; x_thread[i + 1] = x[i + 1] / 64.0f; x_thread[i + 2] = x[i + 2] / 16.0f; diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index 74e6980869..58d87d4ef5 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -670,6 +670,36 @@ def test_qmv(self): self.assertEqual(y_q.shape, y_hat.shape) self.assertLess((y_q - y_hat).abs().max(), 1e-3) + @unittest.skipIf(not mx.metal.is_available(), "requires Metal") + def test_qmv_affine_bias_sum_precision(self): + # The bias contributes bias * sum(x), and a sum accumulated in the input + # type loses the small values: the sum of [1, 1 / denom, -1, 0] cancels + # to zero. + for n, k, gs, dtype, sign, bits in product( + [1, 4, 8, 12], + [64, 96, 512], + [32, 64], + [mx.bfloat16, mx.float16], + [-1, 1], + [2, 3, 4, 5, 6, 8], + ): + if k % gs: + continue + with self.subTest(n=n, k=k, gs=gs, dtype=dtype, sign=sign, bits=bits): + # 1 / denom is exactly half an ulp at 1.0, so 1 + 1 / denom + # rounds back to 1 (ties-to-even) in the input type. + denom = 2048 if dtype == mx.float16 else 256 + x = mx.tile(mx.array([1, 1 / denom, -1, 0], dtype), k // 4) + x = (sign * x).reshape(1, k) + # Zero scales and unit biases dequantize to all ones, so the + # product is exactly sum(x) = sign * k / (4 * denom). + q = mx.zeros((n, k * bits // 32), mx.uint32) + scales = mx.zeros((n, k // gs), dtype) + biases = mx.ones((n, k // gs), dtype) + y = mx.quantized_matmul(x, q, scales, biases, group_size=gs, bits=bits) + expected = mx.full((1, n), sign * k / (4 * denom), dtype) + self.assertTrue(mx.array_equal(y, expected).item()) + def test_fp_qmv(self): key = mx.random.key(0) k1, k2 = mx.random.split(key) From 7ee2bd0b99e6c80e047ff6e445806ab930d3ccc7 Mon Sep 17 00:00:00 2001 From: Yanzhao Wang Date: Tue, 15 Sep 2026 00:23:20 -0700 Subject: [PATCH 2/2] Use the fast qmv kernel for affine outputs not divisible by 8 When the input size is aligned for qmv_fast but the output size is not a multiple of 8, qmv falls back to the generic kernel. For outputs below 8 rows its branch bounds-checks every row inside the reduction loop, which makes each call several times slower on the GPU than the fast kernel. For outputs of 8 rows or more that are not a multiple of 8 the generic kernel takes its other branch, which moves the last tile back over rows it has already computed and so needs no per-row check either. That branch is still the slower one: it reads a narrower block per thread than qmv_fast for every width but 2 bits, and it keeps the guarded tail that an aligned input does not need. Give qmv_fast_impl a partial_rows flag: rows of the last SIMD group that fall past the output reuse the weights of the last valid output row and are not stored, so the reduction loop stays free of per-row checks. qmv uses the new affine_qmv_fast_rows kernel for aligned affine inputs whose output size is not a multiple of 8; other routes are unchanged. --- mlx/backend/metal/kernels/quantized.h | 84 +++++++++++++++++++++-- mlx/backend/metal/kernels/quantized.metal | 1 + mlx/backend/metal/quantized.cpp | 11 ++- python/tests/test_quantized.py | 35 ++++++++++ 4 files changed, 123 insertions(+), 8 deletions(-) diff --git a/mlx/backend/metal/kernels/quantized.h b/mlx/backend/metal/kernels/quantized.h index 5deaad811c..6d70ee1263 100644 --- a/mlx/backend/metal/kernels/quantized.h +++ b/mlx/backend/metal/kernels/quantized.h @@ -753,7 +753,7 @@ METAL_FUNC void qmv_quad_impl( } } -template +template METAL_FUNC void qmv_fast_impl( const device uint32_t* w, const device T* scales, @@ -787,6 +787,18 @@ METAL_FUNC void qmv_fast_impl( const int out_row = tid.y * (num_simdgroups * results_per_simdgroup) + simd_gid * results_per_simdgroup; + // With partial rows the output size need not be a multiple of 8. Rows of the + // last SIMD-group that fall past the output reuse the weights of the last + // valid output row, so the reduction loop needs no per-row bounds check, and + // are not stored. + int last_row = results_per_simdgroup - 1; + if constexpr (partial_rows) { + if (out_row >= out_vec_size) { + return; + } + last_row = min(last_row, out_vec_size - 1 - out_row); + } + ws += out_row * in_vec_size_w + simd_lid * packs_per_thread * bytes_per_pack; scales += out_row * in_vec_size_g + simd_lid / scale_step_per_thread; biases += out_row * in_vec_size_g + simd_lid / scale_step_per_thread; @@ -797,9 +809,13 @@ METAL_FUNC void qmv_fast_impl( U sum = load_vector(x, x_thread); for (int row = 0; row < results_per_simdgroup; row++) { - auto wl = (const device uint8_t*)(ws + row * in_vec_size_w); - const device T* sl = scales + row * in_vec_size_g; - const device T* bl = biases + row * in_vec_size_g; + int src = row; + if constexpr (partial_rows) { + src = min(row, last_row); + } + auto wl = (const device uint8_t*)(ws + src * in_vec_size_w); + const device T* sl = scales + src * in_vec_size_g; + const device T* bl = biases + src * in_vec_size_g; U s = sl[0]; U b = bl[0]; @@ -814,7 +830,7 @@ METAL_FUNC void qmv_fast_impl( for (int row = 0; row < results_per_simdgroup; row++) { result[row] = simd_sum(result[row]); - if (simd_lid == 0) { + if (simd_lid == 0 && row <= last_row) { y[row] = static_cast(result[row]); } } @@ -1655,6 +1671,64 @@ template < simd_lid); } +template < + typename T, + int group_size, + int bits, + bool batched, + bool has_global_scale = false, + int results_per_simdgroup = 4> +[[kernel]] void affine_qmv_fast_rows( + const device uint32_t* w [[buffer(0)]], + const device T* scales [[buffer(1)]], + const device T* biases [[buffer(2)]], + const device T* x [[buffer(3)]], + device T* y [[buffer(4)]], + const constant int& in_vec_size [[buffer(5)]], + const constant int& out_vec_size [[buffer(6)]], + const constant int& x_batch_ndims [[buffer(7)]], + const constant int* x_shape [[buffer(8)]], + const constant int64_t* x_strides [[buffer(9)]], + const constant int& w_batch_ndims [[buffer(10)]], + const constant int* w_shape [[buffer(11)]], + const constant int64_t* w_strides [[buffer(12)]], + const constant int64_t* s_strides [[buffer(13)]], + const constant int64_t* b_strides [[buffer(14)]], + uint3 tid [[threadgroup_position_in_grid]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + if (batched) { + int M = x_shape[x_batch_ndims]; + adjust_matrix_offsets( + x, + w, + scales, + biases, + y, + out_vec_size * M, + x_batch_ndims, + x_shape, + x_strides, + w_batch_ndims, + w_shape, + w_strides, + s_strides, + b_strides, + tid); + } + qmv_fast_impl( + w, + scales, + biases, + x, + y, + in_vec_size, + out_vec_size, + tid, + simd_gid, + simd_lid); +} + template < typename T, int group_size, diff --git a/mlx/backend/metal/kernels/quantized.metal b/mlx/backend/metal/kernels/quantized.metal index 069482cbaf..6d0642c136 100644 --- a/mlx/backend/metal/kernels/quantized.metal +++ b/mlx/backend/metal/kernels/quantized.metal @@ -92,6 +92,7 @@ #define instantiate_quantized_all_batched(type, group_size, bits) \ instantiate_quantized_batched_wrap(affine_qmv_fast, type, group_size, bits) \ + instantiate_quantized_batched_wrap(affine_qmv_fast_rows, type, group_size, bits) \ instantiate_quantized_batched_wrap(affine_qmv, type, group_size, bits) \ instantiate_quantized_batched_wrap(affine_qvm, type, group_size, bits) \ instantiate_quantized_batched_wrap(affine_qmm_n, type, group_size, bits) diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index 55608183dc..45747b82be 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -481,7 +481,12 @@ void qmv( std::string kname; kname.reserve(64); std::string type_string = get_type_string(x.dtype()); - bool fast = N % bn == 0 && K % qmv_fast_k_alignment(bits) == 0; + bool aligned = K % qmv_fast_k_alignment(bits) == 0; + bool fast = N % bn == 0 && aligned; + // Affine outputs that are not a multiple of 8 can still use the fast kernel + // when the input is aligned: its last SIMD-group covers the remaining rows. + bool fast_rows = !fast && aligned && mode == "affine" && !global_scale; + const char* func = fast ? "qmv_fast" : (fast_rows ? "qmv_fast_rows" : "qmv"); // A narrower output tile reduces register pressure for large // floating-point quantized matrix-vector products on M5 Max GPUs. bool use_narrow_qmv = fast && N >= 4096 && d.get_architecture_gen() == 17 && @@ -493,7 +498,7 @@ void qmv( concatenate( kname, - mode + (fast ? "_qmv_fast_" : "_qmv_"), + mode + "_" + func + "_", type_string, "_gs_", group_size, @@ -505,7 +510,7 @@ void qmv( auto kernel = get_quantized_kernel_wrapped( d, kname, - (fast ? "qmv_fast" : "qmv"), + func, mode, type_string, group_size, diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index 58d87d4ef5..e79948ebe0 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -700,6 +700,41 @@ def test_qmv_affine_bias_sum_precision(self): expected = mx.full((1, n), sign * k / (4 * denom), dtype) self.assertTrue(mx.array_equal(y, expected).item()) + @unittest.skipIf(not mx.metal.is_available(), "requires Metal") + def test_qmv_fast_rows(self): + # Output sizes that are not a multiple of 8 with an aligned input size. + # "vector" is one input row, "batched" two weight batches, and "rows" + # is M = 3, which reaches this kernel only on GPUs before gen 15, since + # later ones route M >= 2 to qmv_wide. + for n, k, bits, gs, dtype, layout in product( + [1, 3, 5, 12, 17], + [512, 1024], + [2, 3, 4, 5, 6, 8], + [64], + [mx.float32, mx.bfloat16], + ["vector", "batched", "rows"], + ): + with self.subTest(n=n, k=k, bits=bits, gs=gs, dtype=dtype, layout=layout): + key = mx.random.key(n * 7 + k + bits) + if layout == "batched": + w_shape, x_shape = (2, n, k), (2, 1, k) + elif layout == "rows": + w_shape, x_shape = (n, k), (3, k) + else: + w_shape, x_shape = (n, k), (1, k) + w = mx.random.normal(w_shape, key=key) / k**0.5 + x = mx.random.normal(x_shape, key=mx.random.split(key)[0]) + x = (x / k**0.5).astype(dtype) + q, s, b = mx.quantize(w.astype(dtype), group_size=gs, bits=bits) + w_hat = mx.dequantize(q, s, b, group_size=gs, bits=bits) + y_hat = x.astype(mx.float32) @ mx.swapaxes( + w_hat.astype(mx.float32), -1, -2 + ) + y = mx.quantized_matmul(x, q, s, b, group_size=gs, bits=bits) + tol = 1e-3 if dtype == mx.float32 else 1.5e-3 + self.assertEqual(y.shape, y_hat.shape) + self.assertLess((y - y_hat).abs().max().item(), tol) + def test_fp_qmv(self): key = mx.random.key(0) k1, k2 = mx.random.split(key)