diff --git a/mlx/backend/metal/kernels/quantized.h b/mlx/backend/metal/kernels/quantized.h index 9d9ce368c9..6d70ee1263 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; @@ -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/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/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 74e6980869..e79948ebe0 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -670,6 +670,71 @@ 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()) + + @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)