diff --git a/mlx/backend/metal/kernels/fp_quantized.h b/mlx/backend/metal/kernels/fp_quantized.h index 0591532f..5c5e4b2e 100644 --- a/mlx/backend/metal/kernels/fp_quantized.h +++ b/mlx/backend/metal/kernels/fp_quantized.h @@ -434,7 +434,9 @@ METAL_FUNC void fp_qmv_impl( for (; k < in_vec_size - block_size; k += block_size) { load_vector(x, x_thread); - for (int row = 0; out_row + row < out_vec_size; row++) { + for (int row = 0; + row < results_per_simdgroup && out_row + row < out_vec_size; + row++) { auto wl = (const device uint8_t*)(ws + row * in_vec_size_w); const device auto* sl = scales + row * in_vec_size_g; @@ -453,7 +455,9 @@ METAL_FUNC void fp_qmv_impl( if (remaining > 0) { load_vector_safe(x, x_thread, remaining); - for (int row = 0; out_row + row < out_vec_size; row++) { + for (int row = 0; + row < results_per_simdgroup && out_row + row < out_vec_size; + row++) { auto wl = (const device uint8_t*)(ws + row * in_vec_size_w); const device auto* sl = scales + row * in_vec_size_g; @@ -462,7 +466,9 @@ METAL_FUNC void fp_qmv_impl( } } - for (int row = 0; out_row + row < out_vec_size; row++) { + for (int row = 0; + row < results_per_simdgroup && out_row + row < out_vec_size; + row++) { result[row] = simd_sum(result[row]); if (simd_lid == 0) { y[row] = static_cast(result[row]); diff --git a/mlx/backend/metal/kernels/quantized.h b/mlx/backend/metal/kernels/quantized.h index bf639814..5ac4c6e1 100644 --- a/mlx/backend/metal/kernels/quantized.h +++ b/mlx/backend/metal/kernels/quantized.h @@ -867,7 +867,9 @@ METAL_FUNC void qmv_impl( for (; k < in_vec_size - block_size; k += block_size) { U sum = load_vector(x, x_thread); - for (int row = 0; out_row + row < out_vec_size; row++) { + for (int row = 0; + row < results_per_simdgroup && out_row + row < out_vec_size; + 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; @@ -891,19 +893,23 @@ METAL_FUNC void qmv_impl( U sum = load_vector_safe( x, x_thread, remaining); - for (int row = 0; out_row + row < out_vec_size; row++) { + for (int row = 0; + row < results_per_simdgroup && out_row + row < out_vec_size; + 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; U s = sl[0]; U b = bl[0]; - result[row] += - qdot(wl, x_thread, s, b, sum); + result[row] += qdot_safe( + wl, x_thread, s, b, sum, remaining); } } - for (int row = 0; out_row + row < out_vec_size; row++) { + for (int row = 0; + row < results_per_simdgroup && out_row + row < out_vec_size; + row++) { result[row] = simd_sum(result[row]); if (simd_lid == 0) { y[row] = static_cast(result[row]); diff --git a/python/tests/cuda_skip.py b/python/tests/cuda_skip.py index 47ed4a61..20793d5c 100644 --- a/python/tests/cuda_skip.py +++ b/python/tests/cuda_skip.py @@ -50,6 +50,7 @@ cuda_skip = { "TestQuantized.test_fp_qvm", "TestQuantized.test_qvm", "TestQuantized.test_qvm_splitk", + "TestQuantized.test_qmv_small_non_multiples", "TestQuantized.test_small_matrix", "TestQuantized.test_throw", "TestQuantized.test_vjp_scales_biases", diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index 6e4e9eee..ba108e95 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -681,6 +681,65 @@ class TestQuantized(mlx_tests.MLXTestCase): self.assertEqual(y_q.shape, y_hat.shape) self.assertLess((y_q - y_hat).abs().max(), 1e-3) + def test_qmv_small_non_multiples(self): + # Test very small K and N dimensions (e.g., [MxK] x [NxK].T = [MxN]) + # Each tuple is (M, K, N) representing input rows, weight cols, weight rows + test_cases = [ + (1, 32, 3), + (2, 32, 10), + (1, 32, 5), + (4, 32, 7), + ] + + # Test different quantization settings (bits, group_size, mode) + quantization_settings = [ + (4, 32, "affine"), + (6, 32, "affine"), + (4, 16, "nvfp4"), + ] + + for M, K, N in test_cases: + for bits, group_size, mode in quantization_settings: + # Test without batch dimension + with self.subTest( + M=M, + K=K, + N=N, + batch=None, + group_size=group_size, + bits=bits, + mode=mode, + ): + w = mx.random.normal(shape=(N, K)) + w_q, *sb = mx.quantize( + w, + group_size=group_size, + bits=bits, + mode=mode, + ) + w_hat = mx.dequantize( + w_q, + *sb, + group_size=group_size, + bits=bits, + mode=mode, + ) + + # Test qmv/qmm_t (transpose=True): [MxK] @ [NxK].T = [MxN] + x = mx.random.normal(shape=(M, K)) + y_q = mx.quantized_matmul( + x, + w_q, + *sb, + transpose=True, + group_size=group_size, + bits=bits, + mode=mode, + ) + y_hat = x @ mx.swapaxes(w_hat, -1, -2) + self.assertEqual(y_q.shape, y_hat.shape) + self.assertLess((y_q - y_hat).abs().max(), 1e-3) + def test_gather_qmm(self): def quantize(w, transpose=True, group_size=None, bits=None, mode="affine"): if mode == "affine": @@ -957,7 +1016,7 @@ class TestQuantized(mlx_tests.MLXTestCase): group_size=group_size, mode=mode, transpose=transpose, - rhs_indices=indices + rhs_indices=indices, ) xs, idx, inv_order = gather_sort(x, indices) y3 = mx.gather_mm(xs, w, rhs_indices=idx, sorted_indices=True) @@ -969,7 +1028,7 @@ class TestQuantized(mlx_tests.MLXTestCase): mode=mode, rhs_indices=idx, transpose=transpose, - sorted_indices=True + sorted_indices=True, ) y3 = scatter_unsort(y3, inv_order, indices.shape) y4 = scatter_unsort(y4, inv_order, indices.shape)