From 1ef25498f88eb3d0b0be74951364445f2334b87c Mon Sep 17 00:00:00 2001 From: bri-prism <288398250+bri-prism@users.noreply.github.com> Date: Sun, 27 Sep 2026 22:06:29 -0700 Subject: [PATCH] metal: read FP16 and BF16 affine constants directly for FP32 inputs quantized_matmul used to cast an FP32 input's FP16 or BF16 scales and biases to FP32 before every call: two extra kernels per projection, or a second FP32 copy of every constant kept resident. On Metal the op now keeps them narrow. The one-row qmv_fast kernel reads them as they are and widens each value in registers (a new affine_qmv_fast_mixed entry). Every other route widens them inside eval_gpu first, as the op did. Outputs are bitwise identical to widening first; other backends and CPU streams keep the old cast. --- mlx/backend/metal/kernels/quantized.h | 44 +++++++-- mlx/backend/metal/kernels/quantized.metal | 28 +++++- mlx/backend/metal/quantized.cpp | 64 ++++++++++++ mlx/ops.cpp | 9 +- mlx/primitives.cpp | 1 + python/tests/test_export_import.py | 23 +++++ python/tests/test_quantized.py | 114 ++++++++++++++++++++++ 7 files changed, 274 insertions(+), 9 deletions(-) diff --git a/mlx/backend/metal/kernels/quantized.h b/mlx/backend/metal/kernels/quantized.h index 63fffb1d69..7903a42cd0 100644 --- a/mlx/backend/metal/kernels/quantized.h +++ b/mlx/backend/metal/kernels/quantized.h @@ -847,11 +847,16 @@ METAL_FUNC void qmv_quad_impl( } } -template +template < + typename T, + int group_size, + int bits, + bool partial_rows = false, + typename S = T> METAL_FUNC void qmv_fast_impl( const device uint32_t* w, - const device T* scales, - const device T* biases, + const device S* scales, + const device S* biases, const device T* x, device T* y, const constant int& in_vec_size, @@ -911,8 +916,8 @@ METAL_FUNC void qmv_fast_impl( src = min(row, last_row); } auto wl = (const device uint8_t*)(ws + src * in_vec_size_w); - const device T* sl = scales + src * in_vec_size_g; - const device T* bl = biases + src * in_vec_size_g; + const device S* sl = scales + src * in_vec_size_g; + const device S* bl = biases + src * in_vec_size_g; U s = sl[0]; U b = bl[0]; @@ -942,8 +947,8 @@ METAL_FUNC void qmv_fast_impl( src = min(row, last_row); } auto wl = (const device uint8_t*)(ws + src * in_vec_size_w); - const device T* sl = scales + src * in_vec_size_g; - const device T* bl = biases + src * in_vec_size_g; + const device S* sl = scales + src * in_vec_size_g; + const device S* bl = biases + src * in_vec_size_g; U s = sl[0]; U b = bl[0]; @@ -1739,6 +1744,31 @@ template quad_lid); } +template +[[kernel]] void affine_qmv_fast_mixed( + const device uint32_t* w [[buffer(0)]], + const device S* scales [[buffer(1)]], + const device S* biases [[buffer(2)]], + const device float* x [[buffer(3)]], + device float* y [[buffer(4)]], + const constant int& in_vec_size [[buffer(5)]], + const constant int& out_vec_size [[buffer(6)]], + uint3 tid [[threadgroup_position_in_grid]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + qmv_fast_impl( + w, + scales, + biases, + x, + y, + in_vec_size, + out_vec_size, + tid, + simd_gid, + simd_lid); +} + template < typename T, int group_size, diff --git a/mlx/backend/metal/kernels/quantized.metal b/mlx/backend/metal/kernels/quantized.metal index 324185bf1b..ea70dc329b 100644 --- a/mlx/backend/metal/kernels/quantized.metal +++ b/mlx/backend/metal/kernels/quantized.metal @@ -226,4 +226,30 @@ instantiate_quantized_all_groups(float) instantiate_quantized_all_groups(float16_t) instantiate_quantized_all_groups(bfloat16_t) -// clang-format on + +#define instantiate_qmv_mixed(type, group_size, bits, partial_rows) \ + instantiate_kernel( \ + "affine_qmv_fast_mixed_" #type "_gs_" #group_size "_b_" #bits "_pr_" #partial_rows, \ + affine_qmv_fast_mixed, type, group_size, bits, partial_rows) + +#define instantiate_qmv_mixed_rows(type, group_size, bits) \ + instantiate_qmv_mixed(type, group_size, bits, 0) \ + instantiate_qmv_mixed(type, group_size, bits, 1) + +#define instantiate_qmv_mixed_bits(type, group_size) \ + instantiate_qmv_mixed_rows(type, group_size, 1) \ + instantiate_qmv_mixed_rows(type, group_size, 2) \ + instantiate_qmv_mixed_rows(type, group_size, 3) \ + instantiate_qmv_mixed_rows(type, group_size, 4) \ + instantiate_qmv_mixed_rows(type, group_size, 5) \ + instantiate_qmv_mixed_rows(type, group_size, 6) \ + instantiate_qmv_mixed_rows(type, group_size, 8) + +#define instantiate_qmv_mixed_groups(type) \ + instantiate_qmv_mixed_bits(type, 32) \ + instantiate_qmv_mixed_bits(type, 64) \ + instantiate_qmv_mixed_bits(type, 128) + +instantiate_qmv_mixed_groups(float16_t) +instantiate_qmv_mixed_groups(bfloat16_t) + // clang-format on diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index 5e8c4add3f..0d924eb771 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -2142,6 +2142,53 @@ void dispatch_qmv( mode); } +void qmv_mixed( + const array& x, + const array& w, + const array& scales, + const array& biases, + array& out, + int group_size, + int bits, + int N, + int K, + metal::Device& d, + const Stream& s) { + std::string kname; + auto scale_type = get_type_string(scales.dtype()); + bool partial_rows = N % 8 != 0; + concatenate( + kname, + "affine_qmv_fast_mixed_", + scale_type, + "_gs_", + group_size, + "_b_", + bits, + "_pr_", + partial_rows); + auto kernel = get_quantized_kernel_wrapped( + d, + kname, + "qmv_fast_mixed", + "affine", + scale_type, + group_size, + bits, + partial_rows); + auto& encoder = metal::get_command_encoder(s); + encoder.set_compute_pipeline_state(kernel); + encoder.set_input_array(w, 0); + encoder.set_input_array(scales, 1); + encoder.set_input_array(biases, 2); + encoder.set_input_array(x, 3); + encoder.set_output_array(out, 4); + encoder.set_bytes(K, 5); + encoder.set_bytes(N, 6); + encoder.dispatch_threadgroups( + MTL::Size(1, (N + 7) / 8, 1), MTL::Size(32, 2, 1)); +} + void QuantizedMatmul::eval_gpu(const std::vector& inputs, array& out) { auto& s = stream(); auto& d = metal::device(s.device); @@ -2167,6 +2214,23 @@ void QuantizedMatmul::eval_gpu(const std::vector& inputs, array& out) { int vector_limit = transpose_ ? get_qmv_batch_limit(K, N, d) : 4; auto mode = quantization_mode_to_string(mode_); + if (mode_ == QuantizationMode::Affine && scales.dtype() != x.dtype()) { + if (transpose_ && non_batched && M == 1 && + K % qmv_fast_k_alignment(bits_) == 0) { + qmv_mixed(x, w, scales, *biases, out, group_size_, bits_, N, K, d, s); + return; + } + // Other kernels require constants with the activation dtype. + auto widen = [&](const array& a) { + array copy(a.shape(), x.dtype(), nullptr, {}); + copy_gpu(a, copy, CopyType::General, s); + metal::get_command_encoder(s).add_temporary(copy); + return copy; + }; + scales = widen(scales); + biases = widen(*biases); + } + // For few rows use qmv_mma when N/K is large enough. if (transpose_ && non_batched && (M >= 6) && (M <= 16) && use_qmv_mma(x, w, scales, mode_, M, N, K, d)) { diff --git a/mlx/ops.cpp b/mlx/ops.cpp index 26e7e8caa3..528069f420 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -4966,7 +4966,14 @@ array quantized_matmul( throw std::invalid_argument(msg.str()); } std::vector inputs; - if (qmode == QuantizationMode::Affine) { + bool narrow_constants = qmode == QuantizationMode::Affine && + x.dtype() == float32 && + (scales.dtype() == float16 || scales.dtype() == bfloat16) && + biases->dtype() == scales.dtype() && to_stream(s).device == Device::gpu && + metal::is_available(); + if (narrow_constants) { + inputs = {x, w, scales, *biases}; + } else if (qmode == QuantizationMode::Affine) { inputs = { astype(x, dtype, s), w, diff --git a/mlx/primitives.cpp b/mlx/primitives.cpp index 2f5bc77831..f9b4f801c2 100644 --- a/mlx/primitives.cpp +++ b/mlx/primitives.cpp @@ -3673,6 +3673,7 @@ std::vector QuantizedMatmul::vjp( wq = unflatten(wq, -1, {-1, group_size_}, stream()); vjps.push_back(sum(multiply(*dsb, wq, stream()), -1, false, stream())); } + vjps.back() = astype(vjps.back(), primals[arg].dtype(), stream()); } } return vjps; diff --git a/python/tests/test_export_import.py b/python/tests/test_export_import.py index b8de6ce487..d3ba4fd972 100644 --- a/python/tests/test_export_import.py +++ b/python/tests/test_export_import.py @@ -12,6 +12,29 @@ class TestExportImport(mlx_tests.MLXTestCase): + @unittest.skipIf(not mx.metal.is_available(), "requires Metal") + def test_quantized_matmul_narrow_constants(self): + for dtype in [mx.float16, mx.bfloat16]: + for m in [1, 3]: + x = mx.random.normal((m, 1024)) + wq, scales, biases = mx.quantize( + mx.random.normal((17, 1024)).astype(dtype), group_size=64, bits=4 + ) + + def fun(x, scales, biases): + return mx.quantized_matmul(x, wq, scales, biases, True, 64, 4) + + path = os.path.join(self.test_dir, "qmm-narrow.mlxfn") + mx.export_function(path, fun, x, scales, biases) + imported = mx.import_function(path) + fresh = mx.random.normal(x.shape) + (actual,) = imported(fresh, scales, biases) + expected = fun( + fresh, scales.astype(mx.float32), biases.astype(mx.float32) + ) + self.assertEqual(actual.dtype, mx.float32) + self.assertTrue(mx.array_equal(actual, expected)) + @classmethod def setUpClass(cls): cls.test_dir_fid = tempfile.TemporaryDirectory() diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index 4de260181b..a256d7014c 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -683,6 +683,120 @@ def test_qmm_small_m_block(self): ) self.assertLess((y_ref - y).abs().max(), 1e-3) + @unittest.skipIf(not mx.metal.is_available(), "requires Metal") + def test_qmm_narrow_affine_constants(self): + for dtype, bits, gs, n in product( + [mx.float16, mx.bfloat16], [1, 2, 3, 4, 5, 6, 8], [32, 64, 128], [1, 16, 17] + ): + k = 1536 if bits in [3, 6] else 2048 + with self.subTest(dtype=dtype, bits=bits, gs=gs, n=n): + key_x, key_w = mx.random.split(mx.random.key(7)) + x = mx.random.normal((1, k), key=key_x) + w = mx.random.normal((n, k), key=key_w).astype(dtype) + wq, scales, biases = mx.quantize(w, group_size=gs, bits=bits) + actual = mx.quantized_matmul( + x, wq, scales, biases, group_size=gs, bits=bits, stream=mx.gpu + ) + expected = mx.quantized_matmul( + x, + wq, + scales.astype(mx.float32), + biases.astype(mx.float32), + group_size=gs, + bits=bits, + stream=mx.gpu, + ) + self.assertEqual(actual.dtype, mx.float32) + self.assertTrue(mx.array_equal(actual, expected)) + + def test_qmm_narrow_affine_fallbacks(self): + for dtype, m, k, transpose in product( + [mx.float16, mx.bfloat16], [1, 2, 33], [128, 768], [False, True] + ): + with self.subTest(dtype=dtype, m=m, k=k, transpose=transpose): + x = mx.random.normal((m, k)) + shape = (128, k) if transpose else (k, 128) + wq, scales, biases = mx.quantize( + mx.random.normal(shape).astype(dtype), group_size=32, bits=4 + ) + for device in [mx.cpu, mx.default_device()]: + actual = mx.quantized_matmul( + x, wq, scales, biases, transpose, 32, 4, stream=device + ) + expected = mx.quantized_matmul( + x, + wq, + scales.astype(mx.float32), + biases.astype(mx.float32), + transpose, + 32, + 4, + stream=device, + ) + self.assertTrue(mx.array_equal(actual, expected)) + + def test_qmm_narrow_affine_layouts(self): + for dtype in [mx.float16, mx.bfloat16]: + wq, scales, biases = mx.quantize( + mx.random.normal((2, 34, 1024)).astype(dtype), group_size=64, bits=2 + ) + wq, scales, biases = wq[:, ::2], scales[:, ::2], biases[:, ::2] + x = mx.random.normal((2, 1, 2048))[..., ::2] + for inputs in [ + (x, wq, scales, biases), + (x[0, 0], wq[0], scales[0], biases[0]), + ]: + a, w, s, b = inputs + actual = mx.quantized_matmul(a, w, s, b, True, 64, 2) + expected = mx.quantized_matmul( + a, w, s.astype(mx.float32), b.astype(mx.float32), True, 64, 2 + ) + self.assertTrue(mx.array_equal(actual, expected)) + + @unittest.skipIf(not mx.metal.is_available(), "requires Metal") + def test_qmm_narrow_affine_transforms(self): + for dtype in [mx.float16, mx.bfloat16]: + x = mx.random.normal((1, 1024)) + wq, scales, biases = mx.quantize( + mx.random.normal((9, 1024)).astype(dtype), group_size=64, bits=4 + ) + + def narrow(x, scales, biases): + return mx.quantized_matmul(x, wq, scales, biases, True, 64, 4) + + def widened(x, scales, biases): + return narrow(x, scales.astype(mx.float32), biases.astype(mx.float32)) + + actual = mx.grad(lambda x, s, b: narrow(x, s, b).sum(), (0, 1, 2))( + x, scales, biases + ) + expected = mx.grad(lambda x, s, b: widened(x, s, b).sum(), (0, 1, 2))( + x, scales, biases + ) + for a, e in zip(actual, expected): + self.assertEqual(a.dtype, e.dtype) + self.assertTrue(mx.array_equal(a, e)) + + compiled = mx.compile(narrow) + x_new = mx.random.normal(x.shape, key=mx.random.key(29)) + s_new = scales * mx.array(0.75, dtype=dtype) + b_new = biases + mx.array(0.03125, dtype=dtype) + self.assertTrue( + mx.array_equal( + compiled(x_new, s_new, b_new), widened(x_new, s_new, b_new) + ) + ) + compiled_grad = mx.compile( + mx.grad(lambda x, s, b: narrow(x, s, b).sum(), (0, 1, 2)) + ) + actual = compiled_grad(x_new, s_new, b_new) + expected = mx.grad(lambda x, s, b: widened(x, s, b).sum(), (0, 1, 2))( + x_new, s_new, b_new + ) + for a, e in zip(actual, expected): + self.assertEqual(a.dtype, e.dtype) + self.assertTrue(mx.array_equal(a, e)) + @unittest.skipIf(not mx.metal.is_available(), "requires Metal") def test_qmm_splitk_precision(self): # Split-K sums the partitions in float32 and rounds the result to the