From ce931010ebfd33ccf5f284892d9c7fd4e3c0fed7 Mon Sep 17 00:00:00 2001 From: jax-0n-git Date: Wed, 24 Jun 2026 13:17:53 -0600 Subject: [PATCH] Fix fp_qmv_impl small-output-dim branch using raw fp8 scale byte The out_vec_size < 8 branch's full-block loop loaded the fp8 scale as a raw byte and passed it to qdot without dequantize_scale, so fp-quantized (mxfp4/mxfp8/nvfp4) matvec with output dim < 8 multiplied by the raw e8m0/e4m3 byte instead of the decoded scale (gross error, ~1e2-1e4 relative). Every other fp scale-load site decodes it, including the remainder loop of this same branch. Add test_fp_qmv_small_non_multiples covering the fp modes at output dim < 8 with K large enough to run the full-block loop (the existing fp tests use output dim >= 8 or K below block_size). Co-Authored-By: Claude Opus 4.8 --- mlx/backend/metal/kernels/fp_quantized.h | 2 +- python/tests/test_quantized.py | 41 ++++++++++++++++++++++++ 2 files changed, 42 insertions(+), 1 deletion(-) diff --git a/mlx/backend/metal/kernels/fp_quantized.h b/mlx/backend/metal/kernels/fp_quantized.h index f4bf438df2..6b43a1d114 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 f30170d44d..9101f97a85 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -482,6 +482,47 @@ def test_fp_qmv(self): self.assertEqual(y_q.shape, y_hat.shape) self.assertLess((y_q - y_hat).abs().max(), 1e-3) + def test_fp_qmv_small_non_multiples(self): + # Regression test for the fp_qmv_impl out_vec_size < 8 branch, whose + # full-block loop loaded the fp8 scale as a raw byte instead of decoding + # it through dequantize_scale (see fp_quantized.h). The bug only bites + # fp-quantized matvec when the output dim is < 8 AND K is large enough to + # run at least one full block (K > block_size), so it slipped past both + # test_fp_qmv (output dim >= 8) and test_qmv_small_non_multiples (K = 32, + # below block_size, exercising only the correct remainder loop). + # K = 512 forces full blocks; N in {1..7} hits the < 8 branch. + key = mx.random.key(0) + k1, k2 = mx.random.split(key) + K = 512 + for mode, group_size, bits in [ + ("mxfp4", 32, 4), + ("mxfp8", 32, 8), + ("nvfp4", 16, 4), + ]: + for M in [1, 2]: + for N in [1, 2, 3, 5, 7]: + with self.subTest(M=M, K=K, N=N, mode=mode): + x = mx.random.normal(shape=(M, K), key=k1) / K**0.5 + w = mx.random.normal(shape=(N, K), key=k2) / K**0.5 + w_q, scales = mx.quantize( + w, group_size=group_size, bits=bits, mode=mode + ) + w_hat = mx.dequantize( + w_q, scales, group_size=group_size, bits=bits, mode=mode + ) + y_q = mx.quantized_matmul( + x, + w_q, + scales, + 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_qvm(self): key = mx.random.key(0) k1, k2 = mx.random.split(key)