diff --git a/mlx/backend/metal/kernels/fp_quantized.h b/mlx/backend/metal/kernels/fp_quantized.h index 8d6740db5b..5dd81b4f1e 100644 --- a/mlx/backend/metal/kernels/fp_quantized.h +++ b/mlx/backend/metal/kernels/fp_quantized.h @@ -440,7 +440,7 @@ METAL_FUNC void fp_qmv_impl( auto wl = (const device uint8_t*)(ws + row * in_vec_size_w); const device auto* sl = scales + row * in_vec_size_g; - uint8_t s = sl[0]; + U s = dequantize_scale(sl[0]); result[row] += qdot(wl, x_thread, s); } diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index 2850c7c357..f031d2aa92 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -435,7 +435,7 @@ def test_fp_qmv(self): key = mx.random.key(0) k1, k2 = mx.random.split(key) tests = product( - [256, 512, 67], # M + [256, 512, 67, 5, 7], # M -- 5, 7 exercise out_vec_size < 8 (#3762) [64, 256], # N [0, 1, 3, 8], # B )