From fc323e5045c3adb583cb59b9b84079c9df78f400 Mon Sep 17 00:00:00 2001 From: Magnus Lundstedt Date: Sat, 10 Oct 2026 20:15:31 +0200 Subject: [PATCH] Fix split-K test compatibility on M5, M6 --- python/tests/test_quantized.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index 298dd5be44..e42a7823ca 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -665,8 +665,10 @@ def test_qmm_splitk_precision(self): if mx.default_device() == mx.cpu: self.skipTest("Covers GPU kernels only") K, N = 2560, 1024 + # Keep M at or above get_qmv_batch_limit on all devices. Below it the + # qmv kernel runs, which dequantizes in float32 and not in the input type. for mode, dtype, M in product( - ["affine", "mxfp4"], [mx.bfloat16, mx.float16], [16, 33, 65] + ["affine", "mxfp4"], [mx.bfloat16, mx.float16], [25, 33, 65] ): with self.subTest(mode=mode, dtype=dtype, M=M): k1, k2 = mx.random.split(mx.random.key(0))