From d1c0e9bc80320d52ab557987ea8e4c722549d008 Mon Sep 17 00:00:00 2001 From: katlun-lgtm Date: Sun, 5 Jul 2026 10:02:45 -0400 Subject: [PATCH] Fix fp quantized matvec for output dim < 8 (#3762) fp_qmv_impl's out_vec_size < 8 branch (taken when the weight matrix has fewer than one full output tile) read the fp8 group scale as a raw byte in its full-block loop and passed it straight to qdot, instead of decoding it with dequantize_scale like every sibling site -- including the remainder loop of the same branch. This corrupted mxfp4/mxfp8/nvfp4 matvec whenever the output dimension is < 8 (rel_err ~2e2). Affine quantization uses different kernels and is unaffected. Also extend test_fp_qmv with small output dims (M = 5, 7) which exercise the previously-untested out_vec_size < 8 path. --- mlx/backend/metal/kernels/fp_quantized.h | 2 +- python/tests/test_quantized.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) 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 )