From d084f85b5798b3a13c3128c18be11fd29055a66f Mon Sep 17 00:00:00 2001 From: Jeethu Rao Date: Sun, 11 Oct 2026 11:49:38 +0100 Subject: [PATCH 01/34] metal: 1-bit uint32 wide-load qdot over upstream #3161; non-aligned-K qmv_fast tail test (gemma4-qat#282) Patch 1 of the numen-tech/mlx rebase onto upstream main (06eb7483), which carries upstream's final 1-bit affine quantization (ml-explore/mlx#3161, e0408d47). It replaces the fork's five pre-review #3161 snapshot commits (2fb3d1ad..20bf2eb7) with only the deltas we still carry over upstream. - qdot: read the packed codes as uint32 words, 32 values per load, instead of upstream's byte-at-a-time loop. Same little-endian bit order and accumulation order, so bit-exact with the per-byte path. Provenance: mlx-swift bc96a51 "metal: uint32 wide load in 1-bit qmv qdot (large mobile decode win)", sourced in the fork by b7e6d4d2 (numen-tech/gemma4-qat#179). Every 1-bit qdot caller has values_per_thread == 32 and 4-byte aligned weights (qmv_fast, qmv, qmv_quad at D == 128); qdot_safe, qouter and dequantize keep upstream's byte path. - qmv_fast_impl packs_per_thread: upstream already took the fork's `bits <= 2 ? 1 : 2`; restore the fork's comment (numen-tech/mlx#3). - qmv_fast_impl partial tail: keep upstream's lane predicate `aligned_end + simd_lid * vpt < K` and the partial_rows clamp from ml-explore/mlx#4516. The fork's `aligned_end + (simd_lid + 1) * vpt <= K` is equivalent whenever K % vpt == 0, which the host K gate guarantees. Upstream's comment there was stale (it said the 1-bit block is 2048; it is 1024 with 1 pack/thread); replace it with the fork's OOB reasoning, reworded for this predicate. - dispatch_qmv: document why the qmv_quad K == 64 route excludes 1-bit (a quad lane would hold 16 values, less than one uint32 pack). The fork's affine_sym fall-through comment lands with affine_sym (patch 2). Test: test_qmv_fast_tail_non_aligned_k (GPU vs CPU, 1/4-bit, K = 17/33 groups, N = 67/256). It passes before and after the qdot port. Note that on this base the qmv_fast tail is unreachable from the host (the 1-bit K gate is 2048 = 2 * block; other widths gate at block), so these shapes route to qmv. With an experimental (uncommitted) 1-bit K gate of 512 the tail runs, the ported qdot plus upstream's predicate matched the CPU on 1-bit K = 1536/2560 tail shapes, and a predicate letting one lane read past K failed them. --- mlx/backend/metal/kernels/quantized.h | 44 +++++++++++++++++---------- mlx/backend/metal/quantized.cpp | 4 ++- python/tests/test_quantized.py | 17 +++++++++++ 3 files changed, 48 insertions(+), 17 deletions(-) diff --git a/mlx/backend/metal/kernels/quantized.h b/mlx/backend/metal/kernels/quantized.h index 48fe55cd98..838b0d026c 100644 --- a/mlx/backend/metal/kernels/quantized.h +++ b/mlx/backend/metal/kernels/quantized.h @@ -233,16 +233,19 @@ inline U qdot( U accum = 0; if (bits == 1) { - for (int i = 0; i < (values_per_thread / 8); i++) { - uint8_t wb = w[i]; - accum += select(U(0), x_thread[8 * i], bool(wb & 0x01)); - accum += select(U(0), x_thread[8 * i + 1], bool(wb & 0x02)); - accum += select(U(0), x_thread[8 * i + 2], bool(wb & 0x04)); - accum += select(U(0), x_thread[8 * i + 3], bool(wb & 0x08)); - accum += select(U(0), x_thread[8 * i + 4], bool(wb & 0x10)); - accum += select(U(0), x_thread[8 * i + 5], bool(wb & 0x20)); - accum += select(U(0), x_thread[8 * i + 6], bool(wb & 0x40)); - accum += select(U(0), x_thread[8 * i + 7], bool(wb & 0x80)); + // 1-bit: read 32 packed weights as one 32-bit word (one load instead of + // four byte loads). Little-endian bit order and accumulation order are + // those of the per-byte path, so the result is bit-exact with it. Needs + // values_per_thread % 32 == 0 and a 4-byte aligned w, which every 1-bit + // caller (qmv_fast, qmv, qmv_quad with D == 128) provides. + // (numen-tech/gemma4-qat#179: large mobile decode win) + const device uint32_t* w32 = (const device uint32_t*)w; + for (int i = 0; i < (values_per_thread / 32); i++) { + uint32_t wb = w32[i]; + const thread U* xt = x_thread + 32 * i; + for (int b = 0; b < 32; b++) { + accum += select(U(0), xt[b], bool((wb >> b) & 1u)); + } } } @@ -859,8 +862,9 @@ METAL_FUNC void qmv_fast_impl( uint3 tid [[threadgroup_position_in_grid]], uint simd_gid [[simdgroup_index_in_threadgroup]], uint simd_lid [[thread_index_in_simdgroup]]) { - constexpr int packs_per_thread = - bits <= 2 ? 1 : 2; // 1-bit: 1 pack (vpt=32) for occupancy + // numen-tech/mlx#3: 1-bit uses 1 pack/thread (vpt=32) for occupancy + // (~+11% decode). + constexpr int packs_per_thread = bits <= 2 ? 1 : 2; constexpr int num_simdgroups = 2; constexpr int results_per_simdgroup = 4; constexpr int pack_factor = get_pack_factor(); @@ -925,10 +929,18 @@ METAL_FUNC void qmv_fast_impl( x += block_size; } - // Tail block: handles K that is a multiple of 512 (dispatch gate) but not - // of block_size. Only 1-bit has block_size > 512 (it uses 2048), so for - // >=2-bit this is a no-op. Out-of-bounds lanes skip the qdot call entirely; - // their result[row] stays at 0 and contributes 0 to the simd_sum below. + // Partial last block (in_vec_size % block_size != 0), reached only when the + // host's K gate (qmv_fast_k_alignment) is finer than block_size. A lane + // whose first value lies past the end of the row reads nothing (x, weights, + // scales, biases) and skips the row loop: its result[row] stays 0 and it + // still joins the simd_sum below. Masking such lanes instead (zeroed + // x_thread, scale and bias) would still run qdot, whose weight reads land + // past the end of the row and, for the last output rows, past the end of + // the weight buffer. A lane that starts inside the row also ends inside it + // because the K gate keeps in_vec_size a multiple of values_per_thread; a + // gate that did not would need the stricter predicate + // aligned_end + (simd_lid + 1) * values_per_thread <= in_vec_size. + // (numen-tech/gemma4-qat#179, #282) if (aligned_end < in_vec_size) { bool in_bounds = (aligned_end + static_cast(simd_lid) * values_per_thread) < diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index 052dd761c4..43098763b5 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -1852,7 +1852,9 @@ void dispatch_qmv( metal::Device& d, const Stream& s, const std::string& mode) { - // It is a qmv with a small inner dimension so route to qmv_quad kernel + // It is a qmv with a small inner dimension so route to qmv_quad kernel. + // 1-bit needs K == 128: at K == 64 a quad lane holds 16 values, less than + // the one 32-value uint32 pack that qmv_quad_impl and qdot read. if ((K == 128 || (K == 64 && bits >= 2)) && is_power_of_2(bits) && !global_scale) { qmv_quad(x, w, scales, biases, out, group_size, bits, M, N, K, d, s, mode); diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index 3c1bf0bbf8..fd00fbb96c 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -863,6 +863,23 @@ def test_qmv_fast_unaligned_n(self): self.assertEqual(y.shape, y_hat.shape) self.assertLess((y - y_hat).abs().max().item(), tol) + def test_qmv_fast_tail_non_aligned_k(self): + # K is a multiple of the group size but not of qmv_fast's block, and N is not a multiple + # of 8, so the partial tail and the last-row clamp both run. The GPU must match the CPU + # (no tail lanes there), which it cannot do if a lane reads past K or past the last row. + key = mx.random.key(0) + for bits, gs in ((1, 32), (1, 64), (4, 32), (4, 64)): + for K in (gs * 17, gs * 33): + for N in (67, 256): + with self.subTest(bits=bits, gs=gs, K=K, N=N): + k1, k2, key = mx.random.split(key, 3) + w = mx.random.normal((N, K), key=k1) + x = mx.random.normal((1, K), key=k2).astype(mx.float16) + wq, s, b = mx.quantize(w, group_size=gs, bits=bits) + y = mx.quantized_matmul(x, wq, s, b, True, gs, bits, stream=mx.gpu) + y_cpu = mx.quantized_matmul(x, wq, s, b, True, gs, bits, stream=mx.cpu) + self.assertTrue(mx.allclose(y, y_cpu, atol=1e-2, rtol=1e-2)) + @unittest.skipIf(not mx.metal.is_available(), "requires Metal") def test_fp_qmv_fast_unaligned_n(self): for n, k, mode, dtype, batched in product( From aaa6558a2dcf674da7360bd19696825885071927 Mon Sep 17 00:00:00 2001 From: Jeethu Rao Date: Sun, 11 Oct 2026 11:54:53 +0100 Subject: [PATCH 02/34] python/tests: make test_qmv_fast_tail_non_aligned_k honest and future-proof (gemma4-qat#282) The comment claimed the shapes run qmv_fast's partial tail and the partial_rows clamp. They do not on this base: the host K gate is 2048 for 1-bit and 512 for 4-bit, so every shape here routes to plain qmv (or qmv_wide for M = 3 on gen 15+). Say so, and say when it becomes the tail guard: once the routing patch widens the gate to 1-bit K % 512 and affine 4-bit K % 256 (gemma4-qat#192). Add the shapes that the widened gate makes tail-reachable, keeping the existing ones: 1-bit K in {1536, 2560} x gs {32, 64, 128} and 4-bit K in {768, 1280} x gs {32, 64}, with N in {67, 256} and M in {1, 3}. Same GPU-vs-CPU comparison and tolerances (atol = rtol = 1e-2). Skip without Metal like the neighbouring tests; ruff-format. --- python/tests/test_quantized.py | 53 +++++++++++++++++++++++++--------- 1 file changed, 39 insertions(+), 14 deletions(-) diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index fd00fbb96c..bbac58c948 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -863,22 +863,47 @@ def test_qmv_fast_unaligned_n(self): self.assertEqual(y.shape, y_hat.shape) self.assertLess((y - y_hat).abs().max().item(), tol) + @unittest.skipIf(not mx.metal.is_available(), "requires Metal") def test_qmv_fast_tail_non_aligned_k(self): - # K is a multiple of the group size but not of qmv_fast's block, and N is not a multiple - # of 8, so the partial tail and the last-row clamp both run. The GPU must match the CPU - # (no tail lanes there), which it cannot do if a lane reads past K or past the last row. + # GPU vs CPU for K that is a multiple of the group size but not of + # qmv_fast's block, with N = 67 not a multiple of 8. + # + # On the current host K gate (2048 for 1-bit, 512 for 4-bit) no shape + # here reaches the qmv_fast partial tail or the partial_rows clamp: + # they all route to plain qmv (or qmv_wide for M = 3 on gen 15+), so + # this covers plain qmv at non-aligned K. Once the routing patch + # widens the gate (1-bit K % 512, affine 4-bit K % 256; + # numen-tech/gemma4-qat#192), the 1-bit K = 1536/2560 and 4-bit + # K = 768/1280 shapes run the partial tail and, at N = 67, the + # last-row clamp, and this becomes the tail guard: the GPU cannot + # match the CPU if a tail lane reads past K or past the last row. + cases = [ + (bits, gs, K, N, 1) + for bits, gs in ((1, 32), (1, 64), (4, 32), (4, 64)) + for K in (gs * 17, gs * 33) + for N in (67, 256) + ] + cases += [ + (bits, gs, K, N, M) + for bits, gss, Ks in ( + (1, (32, 64, 128), (1536, 2560)), + (4, (32, 64), (768, 1280)), + ) + for gs in gss + for K in Ks + for N in (67, 256) + for M in (1, 3) + ] key = mx.random.key(0) - for bits, gs in ((1, 32), (1, 64), (4, 32), (4, 64)): - for K in (gs * 17, gs * 33): - for N in (67, 256): - with self.subTest(bits=bits, gs=gs, K=K, N=N): - k1, k2, key = mx.random.split(key, 3) - w = mx.random.normal((N, K), key=k1) - x = mx.random.normal((1, K), key=k2).astype(mx.float16) - wq, s, b = mx.quantize(w, group_size=gs, bits=bits) - y = mx.quantized_matmul(x, wq, s, b, True, gs, bits, stream=mx.gpu) - y_cpu = mx.quantized_matmul(x, wq, s, b, True, gs, bits, stream=mx.cpu) - self.assertTrue(mx.allclose(y, y_cpu, atol=1e-2, rtol=1e-2)) + for bits, gs, K, N, M in cases: + with self.subTest(bits=bits, gs=gs, K=K, N=N, M=M): + k1, k2, key = mx.random.split(key, 3) + w = mx.random.normal((N, K), key=k1) + x = mx.random.normal((M, K), key=k2).astype(mx.float16) + wq, s, b = mx.quantize(w, group_size=gs, bits=bits) + y = mx.quantized_matmul(x, wq, s, b, True, gs, bits, stream=mx.gpu) + y_cpu = mx.quantized_matmul(x, wq, s, b, True, gs, bits, stream=mx.cpu) + self.assertTrue(mx.allclose(y, y_cpu, atol=1e-2, rtol=1e-2)) @unittest.skipIf(not mx.metal.is_available(), "requires Metal") def test_fp_qmv_fast_unaligned_n(self): From f44c32a547b1fa2ee8ffc0df5846f9ec1c3108a6 Mon Sep 17 00:00:00 2001 From: Jeethu Rao Date: Sun, 11 Oct 2026 12:10:06 +0100 Subject: [PATCH 03/34] quantized: bias-free symmetric 1/2-bit affine (affine_sym) on upstream main (gemma4-qat#179, #282) --- mlx/backend/cpu/quantized.cpp | 5 + mlx/backend/cuda/quantized/quantized.cpp | 4 + mlx/backend/metal/jit_kernels.cpp | 4 +- mlx/backend/metal/kernels/quantized.h | 149 ++++++++++++++++++++-- mlx/backend/metal/kernels/quantized.metal | 20 ++- mlx/backend/metal/quantized.cpp | 29 ++++- mlx/ops.cpp | 27 +++- mlx/primitives.cpp | 27 +++- python/tests/test_quantized.py | 127 ++++++++++++++++++ tests/ops_tests.cpp | 93 ++++++++++++++ 10 files changed, 460 insertions(+), 25 deletions(-) diff --git a/mlx/backend/cpu/quantized.cpp b/mlx/backend/cpu/quantized.cpp index f029561c78..a9a7d85b71 100644 --- a/mlx/backend/cpu/quantized.cpp +++ b/mlx/backend/cpu/quantized.cpp @@ -943,6 +943,11 @@ void QuantizedMatmul::eval_cpu(const std::vector& inputs, array& out) { encoder.set_input_array(scales); encoder.set_output_array(out); if (mode_ == QuantizationMode::Affine) { + if (inputs.size() < 4) { + throw std::runtime_error( + "[QuantizedMatmul::eval_cpu] Bias-free affine is Metal-only for " + "now; pass biases (e.g. derive_biases) on CPU."); + } auto biases = ensure_row_contiguous(inputs[3], encoder, stream()); encoder.set_input_array(biases); encoder.dispatch([out = array::unsafe_weak_copy(out), diff --git a/mlx/backend/cuda/quantized/quantized.cpp b/mlx/backend/cuda/quantized/quantized.cpp index a69c49d8f8..8cba23dd97 100644 --- a/mlx/backend/cuda/quantized/quantized.cpp +++ b/mlx/backend/cuda/quantized/quantized.cpp @@ -24,6 +24,10 @@ void QuantizedMatmul::eval_gpu(const std::vector& inputs, array& out) { if (inputs.size() > 3) { biases = inputs[3]; } + if (mode_ == QuantizationMode::Affine && !biases) { + throw std::runtime_error( + "[QuantizedMatmul::eval_gpu] Bias-free affine is Metal-only for now."); + } auto supports = [&](auto&& f) { return f( diff --git a/mlx/backend/metal/jit_kernels.cpp b/mlx/backend/metal/jit_kernels.cpp index aad94eba1e..01fcd63452 100644 --- a/mlx/backend/metal/jit_kernels.cpp +++ b/mlx/backend/metal/jit_kernels.cpp @@ -1062,7 +1062,9 @@ MTL::ComputePipelineState* get_quantized_kernel( metal::utils(), metal::gemm(), metal::quantized_utils(), - (mode == "affine") ? metal::quantized() : metal::fp_quantized(), + // "affine" and "affine_sym" (bias-free) both live in the affine source + (mode.rfind("affine", 0) == 0) ? metal::quantized() + : metal::fp_quantized(), template_def); return kernel_source; }); diff --git a/mlx/backend/metal/kernels/quantized.h b/mlx/backend/metal/kernels/quantized.h index 838b0d026c..86fc5282d1 100644 --- a/mlx/backend/metal/kernels/quantized.h +++ b/mlx/backend/metal/kernels/quantized.h @@ -850,7 +850,22 @@ METAL_FUNC void qmv_quad_impl( } } -template +// Bias-free formats derive the bias from the scale: -scale/2 for 1-bit, +// -scale for 2-bit. +template +METAL_FUNC U sym_derived_bias(U scale) { + static_assert( + !bias_free || bits == 1 || bits == 2, + "bias-free (symmetric) path is only defined for 1- and 2-bit formats"); + return bits == 1 ? U(-0.5f) * scale : -scale; +} + +template < + typename T, + int group_size, + int bits, + bool partial_rows = false, + bool bias_free = false> METAL_FUNC void qmv_fast_impl( const device uint32_t* w, const device T* scales, @@ -919,7 +934,7 @@ METAL_FUNC void qmv_fast_impl( const device T* bl = biases + src * in_vec_size_g; U s = sl[0]; - U b = bl[0]; + U b = bias_free ? sym_derived_bias(s) : (U)bl[0]; result[row] += qdot(wl, x_thread, s, b, sum); } @@ -958,7 +973,7 @@ METAL_FUNC void qmv_fast_impl( const device T* bl = biases + src * in_vec_size_g; U s = sl[0]; - U b = bl[0]; + U b = bias_free ? sym_derived_bias(s) : (U)bl[0]; result[row] += qdot(wl, x_thread, s, b, sum); } @@ -973,7 +988,7 @@ METAL_FUNC void qmv_fast_impl( } } -template +template METAL_FUNC void qmv_impl( const device uint32_t* w, const device T* scales, @@ -1035,7 +1050,7 @@ METAL_FUNC void qmv_impl( const device T* bl = biases + row * in_vec_size_g; U s = sl[0]; - U b = bl[0]; + U b = bias_free ? sym_derived_bias(s) : (U)bl[0]; result[row] += qdot(wl, x_thread, s, b, sum); } @@ -1061,7 +1076,7 @@ METAL_FUNC void qmv_impl( const device T* bl = biases + row * in_vec_size_g; U s = sl[0]; - U b = bl[0]; + U b = bias_free ? sym_derived_bias(s) : (U)bl[0]; result[row] += qdot_safe( wl, x_thread, s, b, sum, remaining); } @@ -1096,7 +1111,7 @@ METAL_FUNC void qmv_impl( const device T* bl = biases + row * in_vec_size_g; U s = sl[0]; - U b = bl[0]; + U b = bias_free ? sym_derived_bias(s) : (U)bl[0]; result[row] += qdot(wl, x_thread, s, b, sum); } @@ -1120,7 +1135,7 @@ METAL_FUNC void qmv_impl( const device T* bl = biases + row * in_vec_size_g; U s = sl[0]; - U b = bl[0]; + U b = bias_free ? sym_derived_bias(s) : (U)bl[0]; result[row] += qdot_safe( wl, x_thread, s, b, sum, remaining); } @@ -1869,6 +1884,124 @@ template < simd_lid); } +// Bias-free decode kernels: same buffer layout as affine_qmv* with the biases +// slot (buffer 2) unbound; the bias comes from sym_derived_bias. +template < + typename T, + int group_size, + int bits, + bool batched, + bool partial_rows = false, + bool has_global_scale = false, + int results_per_simdgroup = 4> +[[kernel]] void affine_sym_qmv_fast( + const device uint32_t* w [[buffer(0)]], + const device T* scales [[buffer(1)]], + const device T* x [[buffer(3)]], + device T* y [[buffer(4)]], + const constant int& in_vec_size [[buffer(5)]], + const constant int& out_vec_size [[buffer(6)]], + const constant int& x_batch_ndims [[buffer(7)]], + const constant int* x_shape [[buffer(8)]], + const constant int64_t* x_strides [[buffer(9)]], + const constant int& w_batch_ndims [[buffer(10)]], + const constant int* w_shape [[buffer(11)]], + const constant int64_t* w_strides [[buffer(12)]], + const constant int64_t* s_strides [[buffer(13)]], + uint3 tid [[threadgroup_position_in_grid]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + const device T* biases = scales; // dummy; never dereferenced when bias_free + if (batched) { + int M = x_shape[x_batch_ndims]; + adjust_matrix_offsets( + x, + w, + scales, + biases, + y, + out_vec_size * M, + x_batch_ndims, + x_shape, + x_strides, + w_batch_ndims, + w_shape, + w_strides, + s_strides, + s_strides, + tid); + } + 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, + int bits, + bool batched, + bool partial_rows = false, + bool has_global_scale = false, + int results_per_simdgroup = 4> +[[kernel]] void affine_sym_qmv( + const device uint32_t* w [[buffer(0)]], + const device T* scales [[buffer(1)]], + const device T* x [[buffer(3)]], + device T* y [[buffer(4)]], + const constant int& in_vec_size [[buffer(5)]], + const constant int& out_vec_size [[buffer(6)]], + const constant int& x_batch_ndims [[buffer(7)]], + const constant int* x_shape [[buffer(8)]], + const constant int64_t* x_strides [[buffer(9)]], + const constant int& w_batch_ndims [[buffer(10)]], + const constant int* w_shape [[buffer(11)]], + const constant int64_t* w_strides [[buffer(12)]], + const constant int64_t* s_strides [[buffer(13)]], + uint3 tid [[threadgroup_position_in_grid]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + const device T* biases = scales; // dummy; never dereferenced when bias_free + if (batched) { + int M = x_shape[x_batch_ndims]; + adjust_matrix_offsets( + x, + w, + scales, + biases, + y, + out_vec_size * M, + x_batch_ndims, + x_shape, + x_strides, + w_batch_ndims, + w_shape, + w_strides, + s_strides, + s_strides, + tid); + } + qmv_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 8e46ef97bc..af250cf205 100644 --- a/mlx/backend/metal/kernels/quantized.metal +++ b/mlx/backend/metal/kernels/quantized.metal @@ -186,6 +186,22 @@ instantiate_quantized_types(64, bits) \ instantiate_quantized_types(32, bits) +// Bias-free (symmetric) decode kernels, 1- and 2-bit only. +#define instantiate_quantized_sym_batched(type, group_size, bits) \ + instantiate_quantized_qmv_wrap(affine_sym_qmv_fast, type, group_size, bits, 1) \ + instantiate_quantized_qmv_wrap(affine_sym_qmv_fast, type, group_size, bits, 0) \ + instantiate_quantized_qmv_wrap(affine_sym_qmv, type, group_size, bits, 0) + +#define instantiate_quantized_sym_types(group_size, bits) \ + instantiate_quantized_sym_batched(float, group_size, bits) \ + instantiate_quantized_sym_batched(float16_t, group_size, bits) \ + instantiate_quantized_sym_batched(bfloat16_t, group_size, bits) + +#define instantiate_quantized_sym_groups(bits) \ + instantiate_quantized_sym_types(128, bits) \ + instantiate_quantized_sym_types(64, bits) \ + instantiate_quantized_sym_types(32, bits) + #define instantiate_quantized_all() \ instantiate_quantized_groups(1) \ instantiate_quantized_groups(2) \ @@ -193,6 +209,8 @@ instantiate_quantized_groups(4) \ instantiate_quantized_groups(5) \ instantiate_quantized_groups(6) \ - instantiate_quantized_groups(8) + instantiate_quantized_groups(8) \ + instantiate_quantized_sym_groups(1) \ + instantiate_quantized_sym_groups(2) instantiate_quantized_all() // clang-format on diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index 43098763b5..2db814cb0e 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -29,7 +29,10 @@ auto get_quantized_kernel_wrapped( int bits, Args... args) { std::string template_def; - std::string fname = ((mode == "affine") ? "affine_" : "fp_") + func; + // "affine" -> affine_; "affine_sym" (bias-free) -> affine_sym_; + // other modes use the fp kernels. + std::string fname = + (mode.rfind("affine", 0) == 0) ? (mode + "_" + func) : ("fp_" + func); template_def = get_template_definition( name, fname, type, group_size, bits, std::forward(args)...); return get_quantized_kernel(d, name, template_def, mode); @@ -147,9 +150,10 @@ inline int get_qmv_batch_limit(int D, int O, metal::Device& d) { } // Must match the K step in qmv_fast_impl (kernels/quantized.h): -// pack_factor() * (bits == 2 ? 1 : 2) * SIMD_SIZE +// pack_factor * packs_per_thread * SIMD_SIZE, with one pack per lane for +// bits <= 2. inline int qmv_fast_k_alignment(int bits) { - return get_pack_factor(bits, 32) * (bits == 2 ? 1 : 2) * 32; + return get_pack_factor(bits, 32) * (bits <= 2 ? 1 : 2) * 32; } inline int add_strides_and_shapes( @@ -543,7 +547,11 @@ void qmv( } // affine qmv_wide only beats qmv on gen-15+; fp benefits on every gen. +// affine_sym has no qmv_wide kernel. inline bool use_qmv_wide(const std::string& mode, metal::Device& d) { + if (mode == "affine_sym") { + return false; + } return mode != "affine" || d.get_architecture_gen() >= 15; } @@ -1852,11 +1860,12 @@ void dispatch_qmv( metal::Device& d, const Stream& s, const std::string& mode) { - // It is a qmv with a small inner dimension so route to qmv_quad kernel. + // It is a qmv with a small inner dimension so route to qmv_quad kernel + // (no bias-free variant; affine_sym falls through to the plain qmv). // 1-bit needs K == 128: at K == 64 a quad lane holds 16 values, less than // the one 32-value uint32 pack that qmv_quad_impl and qdot read. if ((K == 128 || (K == 64 && bits >= 2)) && is_power_of_2(bits) && - !global_scale) { + !global_scale && mode != "affine_sym") { qmv_quad(x, w, scales, biases, out, group_size, bits, M, N, K, d, s, mode); return; } @@ -1907,6 +1916,16 @@ 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 && !biases) { + // Bias-free (symmetric) affine: scales-only checkpoints; the bias is + // derived in-kernel. Decode path (1/2-bit qmv, transpose) only for now. + if (bits_ > 2 || !transpose_ || M >= vector_limit) { + throw std::runtime_error( + "[QuantizedMatmul] Bias-free affine currently supports the 1/2-bit " + "decode path (transpose=true, small M) only."); + } + mode = "affine_sym"; + } // It is a matrix matrix product. if (M >= vector_limit) { // Use split-K qmm for small M with transposed weights (non-batched only) diff --git a/mlx/ops.cpp b/mlx/ops.cpp index 01ad2778f8..40dc979fbd 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -4944,11 +4944,28 @@ array quantized_matmul( std::optional bits_ /* = std::nullopt */, const std::string& mode /* = "affine" */, StreamOrDevice s /* = {} */) { - auto [dtype, qmode] = validate_mode_with_type( - "quantized_matmul", scales, biases, std::nullopt, mode); + Dtype dtype = float32; + QuantizationMode qmode; + if (mode == "affine" && !biases) { + // Bias-free (symmetric) affine: 1/2-bit formats where the bias is a fixed + // function of the scale and is derived in-kernel. + qmode = QuantizationMode::Affine; + dtype = scales.dtype(); + if (!issubdtype(dtype, floating)) { + throw std::invalid_argument( + "[quantized_matmul] Bias-free affine requires floating scales."); + } + } else { + std::tie(dtype, qmode) = validate_mode_with_type( + "quantized_matmul", scales, biases, std::nullopt, mode); + } auto [group_size, bits] = quantization_params_from_mode(qmode, group_size_, bits_); + if (qmode == QuantizationMode::Affine && !biases && bits != 1 && bits != 2) { + throw std::invalid_argument( + "[quantized_matmul] Bias-free affine supports bits in {1, 2} only."); + } // Check and extract the quantized matrix shape against x auto [w_inner_dims, w_outer_dims] = extract_quantized_matmul_dims( "quantized_matmul", x, w, scales, biases, transpose, group_size, bits); @@ -4967,8 +4984,10 @@ array quantized_matmul( } std::vector inputs; if (qmode == QuantizationMode::Affine) { - inputs = { - astype(x, dtype), w, astype(scales, dtype), astype(*biases, dtype)}; + inputs = {astype(x, dtype), w, astype(scales, dtype)}; + if (biases) { + inputs.push_back(astype(*biases, dtype)); + } } else { inputs = {x, w, scales}; } diff --git a/mlx/primitives.cpp b/mlx/primitives.cpp index d4035247e7..b9dfd918f8 100644 --- a/mlx/primitives.cpp +++ b/mlx/primitives.cpp @@ -3600,6 +3600,17 @@ std::vector QuantizedMatmul::vjp( const std::vector&) { std::vector vjps; + // Bias-free affine derives the bias from the scale: -scale/2 for 1-bit, + // -scale for 2-bit. The reverse products run on the biased kernels. + bool sym = mode_ == QuantizationMode::Affine && primals.size() == 3; + float sym_c = (bits_ == 1) ? -0.5f : -1.0f; + std::optional biases = std::nullopt; + if (mode_ == QuantizationMode::Affine) { + biases = sym + ? multiply(primals[2], array(sym_c, primals[2].dtype()), stream()) + : primals[3]; + } + // We rely on the fact that w is always 2D so transpose is simple std::optional dsb = std::nullopt; for (auto arg : argnums) { @@ -3609,8 +3620,7 @@ std::vector QuantizedMatmul::vjp( cotangents[0], primals[1], primals[2], - mode_ == QuantizationMode::Affine ? std::optional(primals[3]) - : std::nullopt, + biases, !transpose_, group_size_, bits_, @@ -3642,11 +3652,15 @@ std::vector QuantizedMatmul::vjp( // biases vjps.push_back(sum(*dsb, -1, false, stream())); } else { - // scales + // scales: d(w_hat)/d(scale) is q, or q + c for the derived bias auto wq = dequantize( primals[1], ones_like(primals[2], stream()), - zeros_like(primals[3], stream()), + full( + primals[2].shape(), + sym ? sym_c : 0.0f, + primals[2].dtype(), + stream()), group_size_, bits_, quantization_mode_to_string(mode_), @@ -3673,8 +3687,9 @@ std::vector QuantizedMatmul::jvp( tangents[0], primals[1], primals[2], - mode_ == QuantizationMode::Affine ? std::optional(primals[3]) - : std::nullopt, + (mode_ == QuantizationMode::Affine && primals.size() > 3) + ? std::optional(primals[3]) + : std::nullopt, transpose_, group_size_, bits_, diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index bbac58c948..ab62833cfc 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -1151,6 +1151,133 @@ def test_qmv_wide(self): self.assertEqual(y_q.shape, y_hat.shape) self.assertLess((y_q - y_hat).abs().max(), tol) + @unittest.skipUnless(mx.metal.is_available(), "Bias-free affine is Metal-only") + def test_qmv_affine_sym(self): + def assert_sym_matches_biased(y_sym, y_biased): + # The biased reference may run a different kernel, so allow a few + # ulps of the output dtype. + tol = 4 * mx.finfo(y_biased.dtype).eps * y_biased.abs().max() + diff = (y_sym.astype(mx.float32) - y_biased.astype(mx.float32)).abs() + self.assertLessEqual(diff.max().item(), tol.item()) + + # Bias-free affine (biases=None) must match the biased kernels fed the + # derived biases and a dequantized reference, on both qmv variants. + key = mx.random.key(0) + k1, k2 = mx.random.split(key) + Ms = [1, 2, 3, 8] + Ns = [256, 67] # 67 is a non-multiple of the 8-row output tile + for bits, group_size, K in product( + [1, 2], [32, 64, 128], [64, 128, 512, 1024, 2048] + ): + if K < group_size: + continue + for M, N in product(Ms, Ns): + with self.subTest(M=M, N=N, K=K, group_size=group_size, bits=bits): + x = mx.random.normal(shape=(M, K), key=k1) + w = mx.random.normal(shape=(N, K), key=k2) + w_q, scales, _ = mx.quantize(w, group_size, bits) + sym_biases = -0.5 * scales if bits == 1 else -scales + w_hat = mx.dequantize(w_q, scales, sym_biases, group_size, bits) + y_hat = x @ mx.swapaxes(w_hat, -1, -2) + y_biased = mx.quantized_matmul( + x, w_q, scales, sym_biases, True, group_size, bits + ) + y_sym = mx.quantized_matmul( + x, w_q, scales, None, True, group_size, bits + ) + mx.eval(y_sym) + self.assertEqual(y_sym.shape, y_hat.shape) + self.assertEqual(y_sym.dtype, x.dtype) + assert_sym_matches_biased(y_sym, y_biased) + self.assertLess((y_sym - y_hat).abs().max(), 1e-3) + + # Half-precision activations: the kernels are templated on the + # activation type and derive the bias in fp32 from the fp16/bf16 scale. + for dtype, bits, K in product([mx.float16, mx.bfloat16], [1, 2], [512, 2048]): + for M in [1, 3]: + with self.subTest(dtype=dtype, M=M, K=K, bits=bits): + x = mx.random.normal(shape=(M, K), key=k1).astype(dtype) + w = mx.random.normal(shape=(256, K), key=k2).astype(dtype) + w_q, scales, _ = mx.quantize(w, 64, bits) + sym_biases = -0.5 * scales if bits == 1 else -scales + y_biased = mx.quantized_matmul( + x, w_q, scales, sym_biases, True, 64, bits + ) + y_sym = mx.quantized_matmul(x, w_q, scales, None, True, 64, bits) + self.assertEqual(y_sym.dtype, dtype) + assert_sym_matches_biased(y_sym, y_biased) + + @unittest.skipUnless(mx.metal.is_available(), "Bias-free affine is Metal-only") + def test_qmv_affine_sym_throws(self): + x = mx.random.normal(shape=(1, 512)) + w = mx.random.normal(shape=(64, 512)) + + # Bias-free affine is defined for 1- and 2-bit only. + for bits in [3, 4, 8]: + w_q, scales, _ = mx.quantize(w, 64, bits) + with self.assertRaises(ValueError): + mx.quantized_matmul(x, w_q, scales, None, True, 64, bits) + w_q, scales, _ = mx.quantize(w, 64, 2) + for bits in [0, -1]: + with self.assertRaises(ValueError): + mx.quantized_matmul(x, w_q, scales, None, True, 64, bits) + + w_q, scales, _ = mx.quantize(w, 64, 2) + x_qvm = mx.random.normal(shape=(1, 64)) + # M = 33 is the smallest M above the qmv batch limit on every GPU. + x_qmm_edge = mx.random.normal(shape=(33, 512)) + x_qmm = mx.random.normal(shape=(512, 512)) + # Evaluate the inputs up front: the checks below throw inside eval, and + # an input left inside a failed eval would block the CPU stream. + mx.eval(x, w_q, scales, x_qvm, x_qmm_edge, x_qmm) + + # Non-transposed (qvm) path has no bias-free kernel. + with self.assertRaises(RuntimeError): + mx.eval(mx.quantized_matmul(x_qvm, w_q, scales, None, False, 64, 2)) + # Matrix regime (M >= the device's qmv batch limit) is qmm-only: the + # bias-free path says so instead of silently taking the qmm route. + for x_big in [x_qmm_edge, x_qmm]: + with self.subTest(M=x_big.shape[0]): + with self.assertRaisesRegex(RuntimeError, "Bias-free affine"): + mx.eval(mx.quantized_matmul(x_big, w_q, scales, None, True, 64, 2)) + # CPU has no bias-free kernel either. + with self.assertRaises(RuntimeError): + mx.eval( + mx.quantized_matmul(x, w_q, scales, None, True, 64, 2, stream=mx.cpu) + ) + # Sanity: the very same call on the GPU mat-vec path works. + mx.eval(mx.quantized_matmul(x, w_q, scales, None, True, 64, 2)) + + @unittest.skipUnless(mx.metal.is_available(), "Bias-free affine is Metal-only") + def test_qmv_affine_sym_grad(self): + # The reverse products use the biased kernels with the derived bias, so + # the gradients match the explicitly biased call. + key = mx.random.key(1) + k1, k2, k3 = mx.random.split(key, 3) + for bits, K in product([1, 2], [512, 1024]): + with self.subTest(bits=bits, K=K): + x = mx.random.normal(shape=(2, K), key=k1) + w = mx.random.normal(shape=(64, K), key=k2) + w_q, scales, _ = mx.quantize(w, 64, bits) + cotan = mx.random.normal(shape=(2, 64), key=k3) + + def f_sym(x, scales): + return mx.quantized_matmul(x, w_q, scales, None, True, 64, bits) + + def f_biased(x, scales): + biases = -0.5 * scales if bits == 1 else -scales + return mx.quantized_matmul(x, w_q, scales, biases, True, 64, bits) + + _, vjp_sym = mx.vjp(f_sym, [x, scales], [cotan]) + _, vjp_biased = mx.vjp(f_biased, [x, scales], [cotan]) + for a, b in zip(vjp_sym, vjp_biased): + self.assertEqual(a.shape, b.shape) + self.assertLess((a - b).abs().max().item(), 1e-3) + + _, jvp_sym = mx.jvp(lambda x: f_sym(x, scales), [x], [x]) + _, jvp_biased = mx.jvp(lambda x: f_biased(x, scales), [x], [x]) + self.assertLess((jvp_sym[0] - jvp_biased[0]).abs().max().item(), 1e-3) + def test_qvm(self): key = mx.random.key(0) k1, k2 = mx.random.split(key) diff --git a/tests/ops_tests.cpp b/tests/ops_tests.cpp index cda9523f37..b7dab2163e 100644 --- a/tests/ops_tests.cpp +++ b/tests/ops_tests.cpp @@ -3590,6 +3590,99 @@ TEST_CASE("test quantize dequantize") { } } +// Bias-free affine decode on Metal: the affine_sym_qmv[_fast] kernels load in +// both the metallib and JIT builds and match the derived-bias reference. +TEST_CASE("test bias-free affine quantized_matmul decode") { + if (!metal::is_available()) { + return; + } + const auto gpu = Device::gpu; + const auto cpu = Device::cpu; + + struct Case { + int bits; + int K; + int N; + int group_size; + }; + const Case cases[] = { + {1, 512, 8, 128}, // 1-bit, K % 1024 != 0: generic affine_sym_qmv + {1, 1536, 8, 128}, // ditto (one full 1024 block + 512) + {1, 1024, 8, 128}, // 1-bit fast path (K % 1024 == 0), no tail + {2, 512, 8, 128}, // 2-bit fast path (K % 512 == 0) + {1, 1536, 64, 128}, // generic, several row groups + {1, 512, 8, 64}, // generic, group 64 + {2, 1536, 64, 64}, // 2-bit fast path, group 64, several row groups + {1, 544, 8, 32}, // K % 512 != 0 on either line: generic affine_sym_qmv + {2, 544, 8, 32}, // ditto, 2-bit + }; + + int seed = 0; + for (auto dtype : {float16, float32}) { + for (const auto& c : cases) { + auto w = + random::bits({c.N, c.K * c.bits / 32}, 4, random::key(seed++), cpu); + auto scales = astype( + random::uniform( + 0.5f, + 1.5f, + {c.N, c.K / c.group_size}, + float32, + random::key(seed++), + cpu), + dtype, + cpu); + auto x = astype( + random::normal({1, c.K}, float32, random::key(seed++), cpu), + dtype, + cpu); + + // Reference in fp32 on the CPU, from the same fp16/fp32 scale values and + // the derived bias the kernel uses. + auto scales_f = astype(scales, float32, cpu); + auto biases_f = + multiply(scales_f, array(c.bits == 1 ? -0.5f : -1.0f), cpu); + auto w_hat = dequantize( + w, + scales_f, + biases_f, + c.group_size, + c.bits, + "affine", + std::nullopt, + std::nullopt, + cpu); + auto expected = + matmul(astype(x, float32, cpu), transpose(w_hat, cpu), cpu); + + auto out = quantized_matmul( + x, + w, + scales, + std::nullopt, + /* transpose = */ true, + c.group_size, + c.bits, + "affine", + gpu); + CHECK_EQ(out.dtype(), dtype); + CHECK_EQ(out.shape(), Shape{1, c.N}); + + // The tolerance covers accumulation order and the fp16 output rounding. + float rtol = dtype == float32 ? 1e-4f : 4e-3f; + float scale = max(abs(expected, cpu), cpu).item(); + float max_diff = + max(abs(subtract(astype(out, float32, cpu), expected, cpu), cpu), cpu) + .item(); + INFO( + "bits=" << c.bits << " K=" << c.K << " N=" << c.N + << " group_size=" << c.group_size << " dtype=" << dtype + << " max_diff=" << max_diff << " max|expected|=" << scale); + CHECK(max_diff <= 1e-3f + rtol * scale); + } + } +} + TEST_CASE("test repeat") { auto data = array({13, 3, 16, 6, 14, 4, 15, 5, 11, 1, 12, 2}, {3, 2, 2}); auto repeat_axis_0 = repeat(data, 2, 0); From f4eba1815294d15ea7b4dd47a88bfdcba202bedc Mon Sep 17 00:00:00 2001 From: Jeethu Rao Date: Sun, 11 Oct 2026 12:29:25 +0100 Subject: [PATCH 04/34] quantized: implied-bias affine quantized_matmul/dequantize on upstream main (gemma4-qat#190, #282) --- mlx/backend/common/quantized.h | 16 + mlx/backend/cpu/quantized.cpp | 58 ++ mlx/backend/metal/kernels/quantized.h | 223 ++++++-- mlx/backend/metal/kernels/quantized.metal | 96 +++- mlx/backend/metal/kernels/quantized_nax.h | 95 +++- mlx/backend/metal/kernels/quantized_nax.metal | 40 +- mlx/backend/metal/quantized.cpp | 409 ++++++++++---- mlx/ops.cpp | 92 ++- mlx/ops.h | 23 +- mlx/primitives.cpp | 36 +- python/src/ops.cpp | 17 +- python/tests/test_quantized.py | 529 ++++++++++++++++++ tests/ops_tests.cpp | 342 +++++++++++ 13 files changed, 1757 insertions(+), 219 deletions(-) diff --git a/mlx/backend/common/quantized.h b/mlx/backend/common/quantized.h index 7656da8b8d..240b4930f9 100644 --- a/mlx/backend/common/quantized.h +++ b/mlx/backend/common/quantized.h @@ -1,5 +1,11 @@ // Copyright © 2026 Apple Inc. +#pragma once + +#include + +#include "mlx/array.h" + namespace mlx::core { inline constexpr short get_pack_factor(int bits, int wsize = 8) { @@ -11,4 +17,14 @@ inline constexpr short get_bytes_per_pack(int bits, int wsize = 8) { return power_of_2_bits ? (wsize / 8) : (bits == 5 ? 5 : 3); } +// Implied bias (affine mode): a 0-d `biases` array holds a single factor f +// standing for the per-group bias `scales * T(f)` in the scales' dtype T. +inline bool is_implied_bias(const array& biases) { + return biases.ndim() == 0; +} + +inline bool is_implied_bias(const std::optional& biases) { + return biases.has_value() && is_implied_bias(*biases); +} + } // namespace mlx::core diff --git a/mlx/backend/cpu/quantized.cpp b/mlx/backend/cpu/quantized.cpp index a9a7d85b71..fc3a1813ef 100644 --- a/mlx/backend/cpu/quantized.cpp +++ b/mlx/backend/cpu/quantized.cpp @@ -950,6 +950,64 @@ void QuantizedMatmul::eval_cpu(const std::vector& inputs, array& out) { } auto biases = ensure_row_contiguous(inputs[3], encoder, stream()); encoder.set_input_array(biases); + if (is_implied_bias(biases)) { + // Implied bias: expand the factor f into a temporary `scales * T(f)` + // for the per-group-bias kernels. + array full_biases(scales.shape(), scales.dtype(), nullptr, {}); + full_biases.set_data(allocator::malloc(full_biases.nbytes())); + encoder.add_temporary(full_biases); + encoder.set_output_array(full_biases); + encoder.dispatch([out = array::unsafe_weak_copy(out), + x = array::unsafe_weak_copy(x), + w = array::unsafe_weak_copy(w), + scales = array::unsafe_weak_copy(scales), + factor = array::unsafe_weak_copy(biases), + full_biases = array::unsafe_weak_copy(full_biases), + group_size_ = group_size_, + bits_ = bits_, + transpose_ = transpose_]() mutable { + float f; + switch (factor.dtype()) { + case float32: + f = factor.data()[0]; + break; + case float16: + f = static_cast(factor.data()[0]); + break; + case bfloat16: + f = static_cast(factor.data()[0]); + break; + default: + throw std::runtime_error( + "[QuantizedMatmul::eval_cpu] Implied-bias factor must be a " + "real floating type."); + } + auto fill = [&](auto* dst, const auto* src) { + using T = std::remove_pointer_t; + for (size_t i = 0; i < scales.size(); ++i) { + dst[i] = static_cast(f * static_cast(src[i])); + } + }; + switch (scales.dtype()) { + case float32: + fill(full_biases.data(), scales.data()); + break; + case float16: + fill(full_biases.data(), scales.data()); + break; + case bfloat16: + fill(full_biases.data(), scales.data()); + break; + default: + throw std::runtime_error( + "[QuantizedMatmul::eval_cpu] Only real floating scales are " + "supported with an implied bias."); + } + _qmm_dispatch( + out, x, w, scales, full_biases, group_size_, bits_, transpose_); + }); + return; + } encoder.dispatch([out = array::unsafe_weak_copy(out), x = array::unsafe_weak_copy(x), w = array::unsafe_weak_copy(w), diff --git a/mlx/backend/metal/kernels/quantized.h b/mlx/backend/metal/kernels/quantized.h index 86fc5282d1..d8e8157c4b 100644 --- a/mlx/backend/metal/kernels/quantized.h +++ b/mlx/backend/metal/kernels/quantized.h @@ -665,6 +665,28 @@ inline void dequantize(const device uint8_t* w, U scale, U bias, W w_local) { } } +// Bias-free formats derive the bias from the scale: -scale/2 for 1-bit, +// -scale for 2-bit. +template +METAL_FUNC U sym_derived_bias(U scale) { + static_assert( + !bias_free || bits == 1 || bits == 2, + "bias-free (symmetric) path is only defined for 1- and 2-bit formats"); + return bits == 1 ? U(-0.5f) * scale : -scale; +} + +// Per-group bias for the affine kernels. An implied bias is f * scale, formed +// in T to match a bias materialized in the scales dtype. +template +METAL_FUNC U affine_bias(U scale, T f, T scale_t, const device T* bias) { + static_assert( + !(bias_free && implied_bias), + "bias_free and implied_bias are mutually exclusive"); + return bias_free ? sym_derived_bias(scale) + : implied_bias ? static_cast(f * scale_t) + : (U)bias[0]; +} + template < typename T, short BROWS, @@ -673,7 +695,8 @@ template < short reduction_dim, short tgp_size, short group_size, - short bits> + short bits, + bool implied_bias = false> struct QuantizedBlockLoader { static_assert( BCOLS <= group_size, @@ -712,6 +735,7 @@ struct QuantizedBlockLoader { const device uint8_t* src; const device T* scales; const device T* biases; + const T bias_factor; // implied bias: the single factor f (see affine_bias) QuantizedBlockLoader( const device uint8_t* src_, @@ -734,7 +758,13 @@ struct QuantizedBlockLoader { src(src_ + bi * src_ld * bytes_per_pack / pack_factor + bj * bytes_per_pack), scales(scales_ + bi * src_ld / group_size), - biases(biases_ + bi * src_ld / group_size) {} + biases(implied_bias ? biases_ : biases_ + bi * src_ld / group_size), + bias_factor(implied_bias ? biases_[0] : T(0)) {} + + T load_bias(T scale) const thread { + return affine_bias( + scale, bias_factor, scale, biases); + } void load_unsafe() const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { @@ -742,7 +772,7 @@ struct QuantizedBlockLoader { } T scale = *scales; - T bias = *biases; + T bias = load_bias(scale); for (int i = 0; i < n_reads; i++) { dequantize( src + i * bytes_per_pack, scale, bias, dst + i * pack_factor); @@ -762,7 +792,7 @@ struct QuantizedBlockLoader { } T scale = *scales; - T bias = *biases; + T bias = load_bias(scale); for (int i = 0; i < n_reads; i++) { dequantize( (device uint8_t*)(src + i * bytes_per_pack), @@ -780,20 +810,31 @@ struct QuantizedBlockLoader { if (group_step_cnt == group_steps) { group_step_cnt = 0; scales++; - biases++; + if (!implied_bias) { + biases++; + } } } else { scales++; - biases++; + if (!implied_bias) { + biases++; + } } } else { scales += group_stride; - biases += group_stride; + if (!implied_bias) { + biases += group_stride; + } } } }; -template +template < + typename T, + int group_size, + int bits, + int D, + bool implied_bias = false> METAL_FUNC void qmv_quad_impl( const device uint32_t* w, const device T* scales, @@ -817,6 +858,8 @@ METAL_FUNC void qmv_quad_impl( thread U x_thread[values_per_thread]; thread U result[results_per_quadgroup] = {0}; + const T f = implied_bias ? biases[0] : T(0); + // Adjust positions const int in_vec_size_w = in_vec_size / pack_factor; const int in_vec_size_g = in_vec_size / group_size; @@ -836,7 +879,7 @@ METAL_FUNC void qmv_quad_impl( const device T* bl = biases + row * in_vec_size_g * quads_per_simd; U s = sl[0]; - U b = bl[0]; + U b = affine_bias(s, f, sl[0], bl); if (row * quads_per_simd + out_row < out_vec_size) { result[row] += qdot(wl, x_thread, s, b, sum); } @@ -850,22 +893,13 @@ METAL_FUNC void qmv_quad_impl( } } -// Bias-free formats derive the bias from the scale: -scale/2 for 1-bit, -// -scale for 2-bit. -template -METAL_FUNC U sym_derived_bias(U scale) { - static_assert( - !bias_free || bits == 1 || bits == 2, - "bias-free (symmetric) path is only defined for 1- and 2-bit formats"); - return bits == 1 ? U(-0.5f) * scale : -scale; -} - template < typename T, int group_size, int bits, bool partial_rows = false, - bool bias_free = false> + bool bias_free = false, + bool implied_bias = false> METAL_FUNC void qmv_fast_impl( const device uint32_t* w, const device T* scales, @@ -895,6 +929,8 @@ METAL_FUNC void qmv_fast_impl( thread U x_thread[values_per_thread]; thread U result[results_per_simdgroup] = {0}; + const T f = implied_bias ? biases[0] : T(0); + // Adjust positions const int in_vec_size_w = in_vec_size * bytes_per_pack / pack_factor; const int in_vec_size_g = in_vec_size / group_size; @@ -934,7 +970,7 @@ METAL_FUNC void qmv_fast_impl( const device T* bl = biases + src * in_vec_size_g; U s = sl[0]; - U b = bias_free ? sym_derived_bias(s) : (U)bl[0]; + U b = affine_bias(s, f, sl[0], bl); result[row] += qdot(wl, x_thread, s, b, sum); } @@ -973,7 +1009,7 @@ METAL_FUNC void qmv_fast_impl( const device T* bl = biases + src * in_vec_size_g; U s = sl[0]; - U b = bias_free ? sym_derived_bias(s) : (U)bl[0]; + U b = affine_bias(s, f, sl[0], bl); result[row] += qdot(wl, x_thread, s, b, sum); } @@ -988,7 +1024,12 @@ METAL_FUNC void qmv_fast_impl( } } -template +template < + typename T, + int group_size, + int bits, + bool bias_free = false, + bool implied_bias = false> METAL_FUNC void qmv_impl( const device uint32_t* w, const device T* scales, @@ -1017,6 +1058,8 @@ METAL_FUNC void qmv_impl( thread U x_thread[values_per_thread]; thread U result[results_per_simdgroup] = {0}; + const T f = implied_bias ? biases[0] : T(0); + // Adjust positions const int in_vec_size_w = in_vec_size * bytes_per_pack / pack_factor; const int in_vec_size_g = in_vec_size / group_size; @@ -1050,7 +1093,7 @@ METAL_FUNC void qmv_impl( const device T* bl = biases + row * in_vec_size_g; U s = sl[0]; - U b = bias_free ? sym_derived_bias(s) : (U)bl[0]; + U b = affine_bias(s, f, sl[0], bl); result[row] += qdot(wl, x_thread, s, b, sum); } @@ -1076,7 +1119,7 @@ METAL_FUNC void qmv_impl( const device T* bl = biases + row * in_vec_size_g; U s = sl[0]; - U b = bias_free ? sym_derived_bias(s) : (U)bl[0]; + U b = affine_bias(s, f, sl[0], bl); result[row] += qdot_safe( wl, x_thread, s, b, sum, remaining); } @@ -1111,7 +1154,7 @@ METAL_FUNC void qmv_impl( const device T* bl = biases + row * in_vec_size_g; U s = sl[0]; - U b = bias_free ? sym_derived_bias(s) : (U)bl[0]; + U b = affine_bias(s, f, sl[0], bl); result[row] += qdot(wl, x_thread, s, b, sum); } @@ -1135,7 +1178,7 @@ METAL_FUNC void qmv_impl( const device T* bl = biases + row * in_vec_size_g; U s = sl[0]; - U b = bias_free ? sym_derived_bias(s) : (U)bl[0]; + U b = affine_bias(s, f, sl[0], bl); result[row] += qdot_safe( wl, x_thread, s, b, sum, remaining); } @@ -1152,7 +1195,13 @@ METAL_FUNC void qmv_impl( // Affine analog of fp_qmv_wide. Weights carry a scale and bias per group, so // each group is decoded in 8-value sub-chunks (scale * q + bias, registers // bounded for any group_size) and reused across the vecs_per_tg vectors. -template +template < + typename T, + int group_size, + int bits, + int vecs_per_tg, + int k_lanes, + bool implied_bias = false> METAL_FUNC void qmv_wide_impl( const device uint32_t* w, const device T* scales, @@ -1185,6 +1234,7 @@ METAL_FUNC void qmv_wide_impl( const device uint8_t* wrow = (const device uint8_t*)w + row * in_vec_size_w; const device T* srow = scales + row * in_vec_size_g; const device T* brow = biases + row * in_vec_size_g; + const T f = implied_bias ? biases[0] : T(0); const device T* xv[vecs_per_tg]; for (int v = 0; v < vecs_per_tg; v++) { @@ -1197,7 +1247,8 @@ METAL_FUNC void qmv_wide_impl( // 8-value sub-chunks and reuse each chunk across the streamed vectors. for (int g = k_lane; g < in_vec_size_g; g += k_lanes) { U scale = srow[g]; - U bias = brow[g]; + U bias = affine_bias( + scale, f, srow[g], brow + g); #pragma unroll for (int sc = 0; sc < group_size / sub; sc++) { const int k0 = g * group_size + sc * sub; @@ -1246,7 +1297,11 @@ METAL_FUNC void qmv_wide_impl( } } -template +template < + typename T, + const int group_size, + const int bits, + bool implied_bias = false> METAL_FUNC void qvm_impl( const device uint32_t* w, const device T* scales, @@ -1282,6 +1337,8 @@ METAL_FUNC void qvm_impl( thread U bias = 0; thread U x_local = 0; + const T f = implied_bias ? biases[0] : T(0); + // Adjust positions const int out_vec_size_w = out_vec_size * bytes_per_pack / pack_factor; const int out_vec_size_g = out_vec_size / group_size; @@ -1302,7 +1359,8 @@ METAL_FUNC void qvm_impl( for (int i = 0; i < in_vec_size; i += block_size) { x_local = *x; scale = *scales; - bias = *biases; + bias = affine_bias( + scale, f, *scales, biases); w_local = *((device vec_w*)ws); qouter( (thread uint8_t*)&w_local, x_local, scale, bias, result); @@ -1316,7 +1374,8 @@ METAL_FUNC void qvm_impl( for (int i = block_size; i < in_vec_size; i += block_size) { x_local = *x; scale = *scales; - bias = *biases; + bias = affine_bias( + scale, f, *scales, biases); w_local = *((device vec_w*)ws); qouter( @@ -1330,7 +1389,8 @@ METAL_FUNC void qvm_impl( if (static_cast(simd_lid) < remaining) { x_local = *x; scale = *scales; - bias = *biases; + bias = affine_bias( + scale, f, *scales, biases); w_local = *((device vec_w*)ws); } else { x_local = 0; @@ -1364,7 +1424,8 @@ template < typename U = T, const int BM = 32, const int BK = 32, - const int BN = 32> + const int BN = 32, + bool implied_bias = false> METAL_FUNC void qmm_t_impl( const device uint32_t* w, const device T* scales, @@ -1406,7 +1467,8 @@ METAL_FUNC void qmm_t_impl( 1, WM * WN * SIMD_SIZE, group_size, - bits>; + bits, + implied_bias>; // Set the block const int K_w = K * bytes_per_pack / pack_factor; @@ -1419,7 +1481,9 @@ METAL_FUNC void qmm_t_impl( x += y_row * static_cast(K); wl += y_col * K_w; scales += y_col * K_g; - biases += y_col * K_g; + if (!implied_bias) { + biases += y_col * K_g; + } y += y_row * static_cast(N) + y_col; // Make the x loader and mma operation @@ -1491,7 +1555,8 @@ template < const int bits, const int BM = 32, const int BK = 32, - const int BN = 32> + const int BN = 32, + bool implied_bias = false> METAL_FUNC void qmm_n_impl( const device uint32_t* w, const device T* scales, @@ -1533,7 +1598,8 @@ METAL_FUNC void qmm_n_impl( 0, WM * WN * SIMD_SIZE, group_size, - bits>; + bits, + implied_bias>; auto wl = (const device uint8_t*)w; @@ -1543,7 +1609,9 @@ METAL_FUNC void qmm_n_impl( x += y_row * static_cast(K); wl += y_col * bytes_per_pack / pack_factor; scales += y_col / group_size; - biases += y_col / group_size; + if (!implied_bias) { + biases += y_col / group_size; + } y += y_row * static_cast(N) + y_col; // Make the x loader and mma operation @@ -1714,7 +1782,15 @@ METAL_FUNC void adjust_matrix_offsets( y += tid.z * output_stride; } -template +// `implied_bias`: `biases` holds one factor f and the bias is f * scale. Zero +// b_strides keep the pointer on f. +template < + typename T, + int group_size, + int bits, + int D, + bool batched, + bool implied_bias = false> [[kernel]] void affine_qmv_quad( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -1753,7 +1829,7 @@ template b_strides, tid); } - qmv_quad_impl( + qmv_quad_impl( w, scales, biases, @@ -1773,7 +1849,8 @@ template < bool batched, bool partial_rows = false, bool has_global_scale = false, - int results_per_simdgroup = 4> + int results_per_simdgroup = 4, + bool implied_bias = false> [[kernel]] void affine_qmv_fast( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -1812,7 +1889,13 @@ template < b_strides, tid); } - qmv_fast_impl( + qmv_fast_impl< + T, + group_size, + bits, + partial_rows, + /*bias_free=*/false, + implied_bias>( w, scales, biases, @@ -1832,7 +1915,8 @@ template < bool batched, bool partial_rows = false, bool has_global_scale = false, - int results_per_simdgroup = 4> + int results_per_simdgroup = 4, + bool implied_bias = false> [[kernel]] void affine_qmv( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -1871,7 +1955,7 @@ template < b_strides, tid); } - qmv_impl( + qmv_impl( w, scales, biases, @@ -2008,7 +2092,8 @@ template < int bits, int vecs_per_tg, int k_lanes, - bool batched> + bool batched, + bool implied_bias = false> [[kernel]] void affine_qmv_wide( const device uint32_t* w, const device T* scales, @@ -2047,7 +2132,7 @@ template < b_strides, tid); } - qmv_wide_impl( + qmv_wide_impl( w, scales, biases, @@ -2066,7 +2151,8 @@ template < const int group_size, const int bits, bool batched, - bool has_global_scale = false> + bool has_global_scale = false, + bool implied_bias = false> [[kernel]] void affine_qvm( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -2105,7 +2191,7 @@ template < b_strides, tid); } - qvm_impl( + qvm_impl( w, scales, biases, @@ -2119,7 +2205,12 @@ template < simd_lid); } -template +template < + typename T, + const int group_size, + const int bits, + int split_k = 32, + bool implied_bias = false> [[kernel]] void affine_qvm_split_k( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -2165,7 +2256,7 @@ template // The in_vec_stride is the full K dimension, not the partition size int in_vec_stride = (split_k - 1) * in_vec_size + final_block_size; - qvm_impl( + qvm_impl( w, scales, biases, @@ -2188,7 +2279,8 @@ template < const bool has_global_scale = false, const int BM = 32, const int BK = 32, - const int BN = 32> + const int BN = 32, + bool implied_bias = false> [[kernel]] void affine_qmm_t( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -2235,7 +2327,7 @@ template < b_strides, tid); } - qmm_t_impl( + qmm_t_impl( w, scales, biases, @@ -2260,7 +2352,8 @@ template < const bool aligned_N, const int BM = 32, const int BK = 32, - const int BN = 32> + const int BN = 32, + bool implied_bias = false> [[kernel]] void affine_qmm_t_splitk( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -2291,10 +2384,12 @@ template < auto wl = (const device uint8_t*)w; wl += k_start * bytes_per_pack / pack_factor; scales += k_start / group_size; - biases += k_start / group_size; + if (!implied_bias) { + biases += k_start / group_size; + } y += tid.z * static_cast(split_k_partition_stride); - qmm_t_impl( + qmm_t_impl( (const device uint32_t*)wl, scales, biases, @@ -2320,7 +2415,8 @@ template < bool has_global_scale = false, int BM = 32, int BK = 32, - int BN = 32> + int BN = 32, + bool implied_bias = false> [[kernel]] void affine_qmm_n( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -2369,7 +2465,7 @@ template < tid); } - qmm_n_impl( + qmm_n_impl( w, scales, biases, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); } @@ -2976,7 +3072,12 @@ template } } -template +template < + typename T, + int group_size, + int bits, + bool has_global_scale = false, + bool implied_bias = false> [[kernel]] void affine_dequantize( const device uint8_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -2991,7 +3092,9 @@ template size_t oindex = offset * pack_factor; size_t gindex = oindex / group_size; T scale = scales[gindex]; - T bias = biases[gindex]; + const T f = implied_bias ? biases[0] : T(0); + T bias = affine_bias( + scale, f, scale, biases + gindex); out += oindex; diff --git a/mlx/backend/metal/kernels/quantized.metal b/mlx/backend/metal/kernels/quantized.metal index af250cf205..fc017e9872 100644 --- a/mlx/backend/metal/kernels/quantized.metal +++ b/mlx/backend/metal/kernels/quantized.metal @@ -202,6 +202,96 @@ instantiate_quantized_sym_types(64, bits) \ instantiate_quantized_sym_types(32, bits) +// Implied-bias ("_ib") variants: 1/2/4/8-bit only. The 3/5/6-bit variants +// exist in JIT builds only. Defaulted template args are spelled out. +#define instantiate_quantized_ib_qmv(name, type, group_size, bits, batched, partial_rows) \ + instantiate_kernel( \ + #name "_" #type "_gs_" #group_size "_b_" #bits "_batch_" #batched "_pr_" #partial_rows "_hgs_0_r_4_ib", \ + name, type, group_size, bits, batched, partial_rows, false, 4, true) + +#define instantiate_quantized_ib_qvm(type, group_size, bits, batched) \ + instantiate_kernel( \ + "affine_qvm_" #type "_gs_" #group_size "_b_" #bits "_batch_" #batched "_ib", \ + affine_qvm, type, group_size, bits, batched, false, true) + +#define instantiate_quantized_ib_qmm_n(type, group_size, bits, batched) \ + instantiate_kernel( \ + "affine_qmm_n_" #type "_gs_" #group_size "_b_" #bits "_batch_" #batched "_ib", \ + affine_qmm_n, type, group_size, bits, batched, false, 32, 32, 32, true) + +#define instantiate_quantized_ib_qmm_t(type, group_size, bits, aligned, batched) \ + instantiate_kernel( \ + "affine_qmm_t_" #type "_gs_" #group_size "_b_" #bits "_alN_" #aligned "_batch_" #batched "_ib", \ + affine_qmm_t, type, group_size, bits, aligned, batched, false, 32, 32, 32, true) + +#define instantiate_quantized_ib_quad(type, group_size, bits, D, batched) \ + instantiate_kernel( \ + "affine_qmv_quad_" #type "_gs_" #group_size "_b_" #bits "_d_" #D "_batch_" #batched "_ib", \ + affine_qmv_quad, type, group_size, bits, D, batched, true) + +#define instantiate_quantized_ib_wide(type, group_size, bits, vecs_per_tg, k_lanes, batched) \ + instantiate_kernel( \ + "affine_qmv_wide_" #type "_gs_" #group_size "_b_" #bits "_nv_" #vecs_per_tg "_kl_" #k_lanes "_batch_" #batched "_ib", \ + affine_qmv_wide, type, group_size, bits, vecs_per_tg, k_lanes, batched, true) + +#define instantiate_quantized_ib_split_k(type, group_size, bits, split_k) \ + instantiate_kernel( \ + "affine_qvm_split_k_" #type "_gs_" #group_size "_b_" #bits "_spk_" #split_k "_ib", \ + affine_qvm_split_k, type, group_size, bits, split_k, true) + +#define instantiate_quantized_ib_splitk_qmm(type, group_size, bits, aligned) \ + instantiate_kernel( \ + "affine_qmm_t_splitk_" #type "_gs_" #group_size "_b_" #bits "_alN_" #aligned "_ib", \ + affine_qmm_t_splitk, type, group_size, bits, aligned, 32, 32, 32, true) + +#define instantiate_quantized_ib_dequantize(type, group_size, bits) \ + instantiate_kernel( \ + "affine_dequantize_" #type "_gs_" #group_size "_b_" #bits "_ib", \ + affine_dequantize, type, group_size, bits, false, true) + +#define instantiate_quantized_ib_funcs(type, group_size, bits) \ + instantiate_quantized_ib_dequantize(type, group_size, bits) \ + instantiate_quantized_ib_qmv(affine_qmv_fast, type, group_size, bits, 1, 1) \ + instantiate_quantized_ib_qmv(affine_qmv_fast, type, group_size, bits, 1, 0) \ + instantiate_quantized_ib_qmv(affine_qmv_fast, type, group_size, bits, 0, 1) \ + instantiate_quantized_ib_qmv(affine_qmv_fast, type, group_size, bits, 0, 0) \ + instantiate_quantized_ib_qmv(affine_qmv, type, group_size, bits, 1, 0) \ + instantiate_quantized_ib_qmv(affine_qmv, type, group_size, bits, 0, 0) \ + instantiate_quantized_ib_qvm(type, group_size, bits, 1) \ + instantiate_quantized_ib_qvm(type, group_size, bits, 0) \ + instantiate_quantized_ib_qmm_n(type, group_size, bits, 1) \ + instantiate_quantized_ib_qmm_n(type, group_size, bits, 0) \ + instantiate_quantized_ib_qmm_t(type, group_size, bits, true, 1) \ + instantiate_quantized_ib_qmm_t(type, group_size, bits, true, 0) \ + instantiate_quantized_ib_qmm_t(type, group_size, bits, false, 1) \ + instantiate_quantized_ib_qmm_t(type, group_size, bits, false, 0) \ + instantiate_quantized_ib_quad(type, group_size, bits, 64, 1) \ + instantiate_quantized_ib_quad(type, group_size, bits, 64, 0) \ + instantiate_quantized_ib_quad(type, group_size, bits, 128, 1) \ + instantiate_quantized_ib_quad(type, group_size, bits, 128, 0) \ + instantiate_quantized_ib_wide(type, group_size, bits, 2, 8, 0) \ + instantiate_quantized_ib_wide(type, group_size, bits, 2, 8, 1) \ + instantiate_quantized_ib_wide(type, group_size, bits, 3, 8, 0) \ + instantiate_quantized_ib_wide(type, group_size, bits, 3, 8, 1) \ + instantiate_quantized_ib_wide(type, group_size, bits, 4, 8, 0) \ + instantiate_quantized_ib_wide(type, group_size, bits, 4, 8, 1) \ + instantiate_quantized_ib_wide(type, group_size, bits, 5, 8, 0) \ + instantiate_quantized_ib_wide(type, group_size, bits, 5, 8, 1) \ + instantiate_quantized_ib_split_k(type, group_size, bits, 8) \ + instantiate_quantized_ib_split_k(type, group_size, bits, 32) \ + instantiate_quantized_ib_splitk_qmm(type, group_size, bits, true) \ + instantiate_quantized_ib_splitk_qmm(type, group_size, bits, false) + +#define instantiate_quantized_ib_types(group_size, bits) \ + instantiate_quantized_ib_funcs(float, group_size, bits) \ + instantiate_quantized_ib_funcs(float16_t, group_size, bits) \ + instantiate_quantized_ib_funcs(bfloat16_t, group_size, bits) + +#define instantiate_quantized_ib_groups(bits) \ + instantiate_quantized_ib_types(128, bits) \ + instantiate_quantized_ib_types(64, bits) \ + instantiate_quantized_ib_types(32, bits) + #define instantiate_quantized_all() \ instantiate_quantized_groups(1) \ instantiate_quantized_groups(2) \ @@ -211,6 +301,10 @@ instantiate_quantized_groups(6) \ instantiate_quantized_groups(8) \ instantiate_quantized_sym_groups(1) \ - instantiate_quantized_sym_groups(2) + instantiate_quantized_sym_groups(2) \ + instantiate_quantized_ib_groups(1) \ + instantiate_quantized_ib_groups(2) \ + instantiate_quantized_ib_groups(4) \ + instantiate_quantized_ib_groups(8) instantiate_quantized_all() // clang-format on diff --git a/mlx/backend/metal/kernels/quantized_nax.h b/mlx/backend/metal/kernels/quantized_nax.h index 95f26868fc..c8e91d5ca3 100644 --- a/mlx/backend/metal/kernels/quantized_nax.h +++ b/mlx/backend/metal/kernels/quantized_nax.h @@ -680,7 +680,8 @@ template < short reduction_dim, short tgp_size, short group_size, - short bits> + short bits, + bool implied_bias = false> struct QuantizedBlockLoader { static_assert( BCOLS <= group_size, @@ -713,6 +714,7 @@ struct QuantizedBlockLoader { const device uint8_t* src; const device T* scales; const device T* biases; + const T bias_factor; // implied bias: the single factor f, bias = f * scale QuantizedBlockLoader( const device uint8_t* src_, @@ -735,7 +737,12 @@ struct QuantizedBlockLoader { src(src_ + bi * src_ld * bytes_per_pack / pack_factor + bj * bytes_per_pack), scales(scales_ + bi * src_ld / group_size), - biases(biases_ + bi * src_ld / group_size) {} + biases(implied_bias ? biases_ : biases_ + bi * src_ld / group_size), + bias_factor(implied_bias ? biases_[0] : T(0)) {} + + T load_bias(T scale) const thread { + return implied_bias ? bias_factor * scale : *biases; + } void load_unsafe() const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { @@ -743,7 +750,7 @@ struct QuantizedBlockLoader { } T scale = *scales; - T bias = *biases; + T bias = load_bias(scale); for (int i = 0; i < n_reads; i++) { dequantize( src + i * bytes_per_pack, scale, bias, dst + i * pack_factor); @@ -763,7 +770,7 @@ struct QuantizedBlockLoader { } T scale = *scales; - T bias = *biases; + T bias = load_bias(scale); for (int i = 0; i < n_reads; i++) { dequantize( (device uint8_t*)(src + i * bytes_per_pack), @@ -781,15 +788,21 @@ struct QuantizedBlockLoader { if (group_step_cnt == group_steps) { group_step_cnt = 0; scales++; - biases++; + if (!implied_bias) { + biases++; + } } } else { scales++; - biases++; + if (!implied_bias) { + biases++; + } } } else { scales += group_stride; - biases += group_stride; + if (!implied_bias) { + biases += group_stride; + } } } }; @@ -801,7 +814,8 @@ template < short dst_ld, short reduction_dim, short tgp_size, - short bits> + short bits, + bool implied_bias> struct QuantizedBlockLoader< T, BROWS, @@ -810,7 +824,8 @@ struct QuantizedBlockLoader< reduction_dim, tgp_size, 32, - bits> { + bits, + implied_bias> { MLX_MTL_CONST short group_size = 32; static_assert( @@ -846,6 +861,7 @@ struct QuantizedBlockLoader< const device uint8_t* src; const device T* scales; const device T* biases; + const T bias_factor; // implied bias: the single factor f, bias = f * scale QuantizedBlockLoader( const device uint8_t* src_, @@ -868,7 +884,14 @@ struct QuantizedBlockLoader< src(src_ + bi * src_ld * bytes_per_pack / pack_factor + bj * bytes_per_pack), scales(scales_ + bi * src_ld / group_size + group_id), - biases(biases_ + bi * src_ld / group_size + group_id) {} + biases( + implied_bias ? biases_ + : biases_ + bi * src_ld / group_size + group_id), + bias_factor(implied_bias ? biases_[0] : T(0)) {} + + T load_bias(T scale) const thread { + return implied_bias ? bias_factor * scale : *biases; + } void load_unsafe() const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { @@ -876,7 +899,7 @@ struct QuantizedBlockLoader< } T scale = *scales; - T bias = *biases; + T bias = load_bias(scale); for (int i = 0; i < n_reads; i++) { dequantize( src + i * bytes_per_pack, scale, bias, dst + i * pack_factor); @@ -896,7 +919,7 @@ struct QuantizedBlockLoader< } T scale = *scales; - T bias = *biases; + T bias = load_bias(scale); for (int i = 0; i < n_reads; i++) { dequantize( (device uint8_t*)(src + i * bytes_per_pack), @@ -918,11 +941,15 @@ struct QuantizedBlockLoader< // } // } else { scales += n_groups; - biases += n_groups; + if (!implied_bias) { + biases += n_groups; + } // } } else { scales += group_stride; - biases += group_stride; + if (!implied_bias) { + biases += group_stride; + } } } }; @@ -1029,7 +1056,8 @@ template < const int BK = 64, const int BN = 64, const int WM = 2, - const int WN = 2> + const int WN = 2, + bool implied_bias = false> METAL_FUNC void qmm_t_nax_tgp_impl( const device uint32_t* w, const device T* scales, @@ -1062,7 +1090,8 @@ METAL_FUNC void qmm_t_nax_tgp_impl( 1, WM * WN * SIMD_SIZE, group_size, - bits>; + bits, + implied_bias>; // Set the block const int K_w = K * bytes_per_pack / pack_factor; @@ -1075,7 +1104,9 @@ METAL_FUNC void qmm_t_nax_tgp_impl( x += y_row * static_cast(K); wl += y_col * K_w; scales += y_col * K_g; - biases += y_col * K_g; + if (!implied_bias) { + biases += y_col * K_g; + } y += y_row * static_cast(N) + y_col; // Make the weight loader @@ -1173,7 +1204,8 @@ template < const int BK = 64, const int BN = 64, const int WM = 2, - const int WN = 2> + const int WN = 2, + bool implied_bias = false> METAL_FUNC void qmm_n_nax_tgp_impl( const device uint32_t* w, const device T* scales, @@ -1206,7 +1238,8 @@ METAL_FUNC void qmm_n_nax_tgp_impl( 0, WM * WN * SIMD_SIZE, group_size, - bits>; + bits, + implied_bias>; // Set the block const int y_row = tid.y * BM; @@ -1218,7 +1251,9 @@ METAL_FUNC void qmm_n_nax_tgp_impl( x += y_row * static_cast(K); wl += y_col * bytes_per_pack / pack_factor; scales += y_col / group_size; - biases += y_col / group_size; + if (!implied_bias) { + biases += y_col / group_size; + } y += y_row * static_cast(N) + y_col; // Make the weight loader @@ -1304,7 +1339,8 @@ template < const int BN = 64, const int WM = 2, const int WN = 2, - const bool has_global_scale = false> + const bool has_global_scale = false, + bool implied_bias = false> [[kernel]] void affine_qmm_t_nax( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -1350,7 +1386,17 @@ template < b_strides, tid); } - qmm_t_nax_tgp_impl( + qmm_t_nax_tgp_impl< + T, + group_size, + bits, + aligned_N, + BM, + BK, + BN, + WM, + WN, + implied_bias>( w, scales, biases, x, y, Ws, K, N, M, tid, lid, simd_gid, simd_lid); } @@ -1363,7 +1409,8 @@ template < const int BK = 64, const int BN = 64, const int WM = 2, - const int WN = 2> + const int WN = 2, + bool implied_bias = false> [[kernel]] void affine_qmm_n_nax( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -1410,7 +1457,7 @@ template < tid); } - qmm_n_nax_tgp_impl( + qmm_n_nax_tgp_impl( w, scales, biases, x, y, Ws, K, N, M, tid, lid, simd_gid, simd_lid); } diff --git a/mlx/backend/metal/kernels/quantized_nax.metal b/mlx/backend/metal/kernels/quantized_nax.metal index 4d61232ea7..7d968bbda5 100644 --- a/mlx/backend/metal/kernels/quantized_nax.metal +++ b/mlx/backend/metal/kernels/quantized_nax.metal @@ -101,6 +101,40 @@ instantiate_quantized_types(64, bits) \ instantiate_quantized_types(32, bits) +// Implied-bias ("_ib") variants: 1/2/4/8-bit only. The 3/5/6-bit variants +// exist in JIT builds only. +#define instantiate_quantized_ib_batched(name, type, group_size, bits, bm, bn, bk, wm, wn, batched) \ + instantiate_kernel( \ + #name "_" #type "_gs_" #group_size "_b_" #bits "_bm" #bm "_bn" #bn "_bk" #bk "_wm" #wm "_wn" #wn "_batch_" #batched "_ib", \ + name, type, group_size, bits, batched, bm, bk, bn, wm, wn, true) + +#define instantiate_quantized_ib_aligned_batched(name, type, group_size, bits, bm, bn, bk, wm, wn, aligned, batched) \ + instantiate_kernel( \ + #name "_" #type "_gs_" #group_size "_b_" #bits "_bm" #bm "_bn" #bn "_bk" #bk "_wm" #wm "_wn" #wn "_alN_" #aligned "_batch_" #batched "_ib", \ + name, type, group_size, bits, aligned, batched, bm, bk, bn, wm, wn, false, true) + +#define instantiate_quantized_ib_funcs(type, group_size, bits) \ + instantiate_quantized_ib_batched(affine_qmm_n_nax, type, group_size, bits, 64, 64, 64, 2, 2, 1) \ + instantiate_quantized_ib_batched(affine_qmm_n_nax, type, group_size, bits, 64, 64, 64, 2, 2, 0) \ + instantiate_quantized_ib_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 64, 64, 64, 2, 2, true, 1) \ + instantiate_quantized_ib_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 64, 64, 64, 2, 2, true, 0) \ + instantiate_quantized_ib_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 64, 64, 64, 2, 2, false, 1) \ + instantiate_quantized_ib_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 64, 64, 64, 2, 2, false, 0) \ + instantiate_quantized_ib_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 32, 64, 64, 2, 2, true, 1) \ + instantiate_quantized_ib_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 32, 64, 64, 2, 2, true, 0) \ + instantiate_quantized_ib_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 32, 64, 64, 2, 2, false, 1) \ + instantiate_quantized_ib_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 32, 64, 64, 2, 2, false, 0) + +#define instantiate_quantized_ib_types(group_size, bits) \ + instantiate_quantized_ib_funcs(float, group_size, bits) \ + instantiate_quantized_ib_funcs(float16_t, group_size, bits) \ + instantiate_quantized_ib_funcs(bfloat16_t, group_size, bits) + +#define instantiate_quantized_ib_groups(bits) \ + instantiate_quantized_ib_types(128, bits) \ + instantiate_quantized_ib_types(64, bits) \ + instantiate_quantized_ib_types(32, bits) + #define instantiate_quantized_all() \ instantiate_quantized_groups(1) \ instantiate_quantized_groups(2) \ @@ -108,6 +142,10 @@ instantiate_quantized_groups(4) \ instantiate_quantized_groups(5) \ instantiate_quantized_groups(6) \ - instantiate_quantized_groups(8) + instantiate_quantized_groups(8) \ + instantiate_quantized_ib_groups(1) \ + instantiate_quantized_ib_groups(2) \ + instantiate_quantized_ib_groups(4) \ + instantiate_quantized_ib_groups(8) instantiate_quantized_all() // clang-format on diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index 2db814cb0e..ba441ddb7e 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -55,6 +55,9 @@ auto get_qmm_nax_kernel_wrapped( return get_qmm_nax_kernel(d, name, template_def, mode); } +// Implied-bias kernels have an "_ib" name suffix and a trailing `implied_bias` +// template arg. Only pass that arg on the implied branch. + inline array ensure_row_contiguous(const array& x, metal::Device& d, const Stream& s) { if (!x.flags().row_contiguous) { @@ -73,6 +76,10 @@ inline array ensure_row_contiguous_matrix( if (x.flags().row_contiguous) { return x; } + if (x.ndim() == 0) { + // A 0-d array (an implied-bias factor) is trivially contiguous. + return x; + } if (x.ndim() < 2) { if (x.strides()[0] == 1) { return x; @@ -180,7 +187,14 @@ inline int add_strides_and_shapes( compute_encoder.set_vector_bytes(w.strides(), offset++); compute_encoder.set_vector_bytes(scales.strides(), offset++); if (biases) { - compute_encoder.set_vector_bytes(biases->strides(), offset++); + if (biases->ndim() == 0) { + // Implied-bias factor: zero strides so the batch offset computed by + // adjust_matrix_offsets leaves the pointer on the single factor. + Strides b_strides(std::max(w_batch_ndims, 1), 0); + compute_encoder.set_vector_bytes(b_strides, offset++); + } else { + compute_encoder.set_vector_bytes(biases->strides(), offset++); + } } return offset; @@ -284,6 +298,7 @@ void quantize_impl( auto type_string = dequantize ? get_type_string(out.dtype()) : get_type_string(w_pre.dtype()); auto mode_string = quantization_mode_to_string(mode); + bool implied_bias = dequantize && has_biases && is_implied_bias(inputs[2]); std::string kname; concatenate( kname, @@ -297,15 +312,28 @@ void quantize_impl( if (!has_biases) { concatenate(kname, "_hgs_", has_global_scale ? "true" : "false"); } - auto kernel = get_quantized_kernel_wrapped( - d, - kname, - dequantize ? "dequantize" : "quantize", - mode_string, - type_string, - group_size, - bits, - has_global_scale); + if (implied_bias) { + kname += "_ib"; + } + auto kernel = implied_bias ? get_quantized_kernel_wrapped( + d, + kname, + "dequantize", + mode_string, + type_string, + group_size, + bits, + has_global_scale, + true) + : get_quantized_kernel_wrapped( + d, + kname, + dequantize ? "dequantize" : "quantize", + mode_string, + type_string, + group_size, + bits, + has_global_scale); auto [grid_dims, group_dims] = get_quantize_kernel_dims(kernel, w, out, group_size, bits, dequantize); @@ -445,9 +473,29 @@ void qmv_quad( bits, "_d_", K, - B > 1 ? "_batch_1" : "_batch_0"); - auto kernel = get_quantized_kernel_wrapped( - d, kname, "qmv_quad", mode, type_string, group_size, bits, K, B > 1); + B > 1 ? "_batch_1" : "_batch_0", + is_implied_bias(biases) ? "_ib" : ""); + auto kernel = is_implied_bias(biases) ? get_quantized_kernel_wrapped( + d, + kname, + "qmv_quad", + mode, + type_string, + group_size, + bits, + K, + B > 1, + true) + : get_quantized_kernel_wrapped( + d, + kname, + "qmv_quad", + mode, + type_string, + group_size, + bits, + K, + B > 1); auto& compute_encoder = metal::get_command_encoder(s); compute_encoder.set_compute_pipeline_state(kernel); @@ -512,19 +560,33 @@ void qmv( partial_rows ? "_pr_1" : "_pr_0", global_scale ? "_hgs_1" : "_hgs_0", "_r_", - results_per_simdgroup); - auto kernel = get_quantized_kernel_wrapped( - d, - kname, - (fast ? "qmv_fast" : "qmv"), - mode, - type_string, - group_size, - bits, - B > 1, - partial_rows, - global_scale.has_value(), - results_per_simdgroup); + results_per_simdgroup, + is_implied_bias(biases) ? "_ib" : ""); + auto kernel = is_implied_bias(biases) ? get_quantized_kernel_wrapped( + d, + kname, + (fast ? "qmv_fast" : "qmv"), + mode, + type_string, + group_size, + bits, + B > 1, + partial_rows, + global_scale.has_value(), + results_per_simdgroup, + true) + : get_quantized_kernel_wrapped( + d, + kname, + (fast ? "qmv_fast" : "qmv"), + mode, + type_string, + group_size, + bits, + B > 1, + partial_rows, + global_scale.has_value(), + results_per_simdgroup); auto& compute_encoder = metal::get_command_encoder(s); compute_encoder.set_compute_pipeline_state(kernel); @@ -608,18 +670,31 @@ void qmv_wide( vecs_per_tg, "_kl_", k_lanes, - batched ? "_batch_1" : "_batch_0"); - auto kernel = get_quantized_kernel_wrapped( - d, - kname, - "qmv_wide", - mode, - type_string, - group_size, - bits, - vecs_per_tg, - k_lanes, - batched); + batched ? "_batch_1" : "_batch_0", + is_implied_bias(biases) ? "_ib" : ""); + auto kernel = is_implied_bias(biases) ? get_quantized_kernel_wrapped( + d, + kname, + "qmv_wide", + mode, + type_string, + group_size, + bits, + vecs_per_tg, + k_lanes, + batched, + true) + : get_quantized_kernel_wrapped( + d, + kname, + "qmv_wide", + mode, + type_string, + group_size, + bits, + vecs_per_tg, + k_lanes, + batched); auto& compute_encoder = metal::get_command_encoder(s); compute_encoder.set_compute_pipeline_state(kernel); @@ -727,11 +802,29 @@ void qvm_split_k( "_b_", bits, "_spk_", - split_k); + split_k, + is_implied_bias(biases) ? "_ib" : ""); // Encode and dispatch kernel - auto kernel = get_quantized_kernel_wrapped( - d, kname, "qvm_split_k", mode, type_string, group_size, bits, split_k); + auto kernel = is_implied_bias(biases) ? get_quantized_kernel_wrapped( + d, + kname, + "qvm_split_k", + mode, + type_string, + group_size, + bits, + split_k, + true) + : get_quantized_kernel_wrapped( + d, + kname, + "qvm_split_k", + mode, + type_string, + group_size, + bits, + split_k); compute_encoder.set_compute_pipeline_state(kernel); @@ -754,8 +847,15 @@ void qvm_split_k( compute_encoder.set_vector_bytes(w_strides, c++); compute_encoder.set_vector_bytes(s_strides, c++); if (biases) { - auto b_strides = biases->strides(); - b_strides.insert(b_strides.end() - 2, split_D * biases->shape(-1)); + Strides b_strides; + if (biases->ndim() == 0) { + // Implied-bias factor: zero strides (same rank as the split scales) so + // every split-K slice reads the single factor. + b_strides = Strides(s_strides.size(), 0); + } else { + b_strides = biases->strides(); + b_strides.insert(b_strides.end() - 2, split_D * biases->shape(-1)); + } compute_encoder.set_vector_bytes(b_strides, c++); } compute_encoder.set_bytes(final_block_size, c++); @@ -806,17 +906,29 @@ void qvm( "_b_", bits, B > 1 ? "_batch_1" : "_batch_0", - global_scale ? "_hgs" : ""); - auto kernel = get_quantized_kernel_wrapped( - d, - kname, - "qvm", - mode, - type_string, - group_size, - bits, - B > 1, - global_scale.has_value()); + global_scale ? "_hgs" : "", + is_implied_bias(biases) ? "_ib" : ""); + auto kernel = is_implied_bias(biases) ? get_quantized_kernel_wrapped( + d, + kname, + "qvm", + mode, + type_string, + group_size, + bits, + B > 1, + global_scale.has_value(), + true) + : get_quantized_kernel_wrapped( + d, + kname, + "qvm", + mode, + type_string, + group_size, + bits, + B > 1, + global_scale.has_value()); auto& compute_encoder = metal::get_command_encoder(s); compute_encoder.set_compute_pipeline_state(kernel); @@ -890,41 +1002,76 @@ void qmm_nax( wn, transpose ? (aligned ? "_alN_true" : "_alN_false") : "", batched ? "_batch_1" : "_batch_0", - global_scale ? "_hgs" : ""); + global_scale ? "_hgs" : "", + is_implied_bias(biases) ? "_ib" : ""); std::string template_def; MTL::ComputePipelineState* kernel; + const bool implied_bias = is_implied_bias(biases); if (transpose) { - kernel = get_qmm_nax_kernel_wrapped( - d, - kname, - "qmm_t_nax", - mode, - type_string, - group_size, - bits, - aligned, - batched, - bm, - bk, - bn, - wm, - wn, - global_scale.has_value()); + // An implied bias is a bias, so it never comes with a global scale. + kernel = implied_bias ? get_qmm_nax_kernel_wrapped( + d, + kname, + "qmm_t_nax", + mode, + type_string, + group_size, + bits, + aligned, + batched, + bm, + bk, + bn, + wm, + wn, + /* has_global_scale = */ false, + true) + : get_qmm_nax_kernel_wrapped( + d, + kname, + "qmm_t_nax", + mode, + type_string, + group_size, + bits, + aligned, + batched, + bm, + bk, + bn, + wm, + wn, + global_scale.has_value()); } else { - kernel = get_qmm_nax_kernel_wrapped( - d, - kname, - "qmm_n_nax", - mode, - type_string, - group_size, - bits, - batched, - bm, - bk, - bn, - wm, - wn); + kernel = implied_bias ? get_qmm_nax_kernel_wrapped( + d, + kname, + "qmm_n_nax", + mode, + type_string, + group_size, + bits, + batched, + bm, + bk, + bn, + wm, + wn, + true) + : get_qmm_nax_kernel_wrapped( + d, + kname, + "qmm_n_nax", + mode, + type_string, + group_size, + bits, + batched, + bm, + bk, + bn, + wm, + wn); } auto& compute_encoder = metal::get_command_encoder(s); compute_encoder.set_compute_pipeline_state(kernel); @@ -1128,24 +1275,60 @@ void qmm( bits, transpose ? (aligned ? "_alN_true" : "_alN_false") : "", batched ? "_batch_1" : "_batch_0", - global_scale ? "_hgs" : ""); + global_scale ? "_hgs" : "", + is_implied_bias(biases) ? "_ib" : ""); std::string template_def; MTL::ComputePipelineState* kernel; + const bool implied_bias = is_implied_bias(biases); + // The implied-bias variants spell out has_global_scale (an implied bias is a + // bias, so it never comes with a global scale) and the default tile sizes + // (BM/BK/BN = 32) so the trailing `implied_bias` template argument lands in + // its slot. if (transpose) { - kernel = get_quantized_kernel_wrapped( - d, - kname, - "qmm_t", - mode, - type_string, - group_size, - bits, - aligned, - batched, - global_scale.has_value()); + kernel = implied_bias ? get_quantized_kernel_wrapped( + d, + kname, + "qmm_t", + mode, + type_string, + group_size, + bits, + aligned, + batched, + /* has_global_scale = */ false, + bm, + 32, + bn, + true) + : get_quantized_kernel_wrapped( + d, + kname, + "qmm_t", + mode, + type_string, + group_size, + bits, + aligned, + batched, + global_scale.has_value()); } else { - kernel = get_quantized_kernel_wrapped( - d, kname, "qmm_n", mode, type_string, group_size, bits, batched); + kernel = implied_bias + ? get_quantized_kernel_wrapped( + d, + kname, + "qmm_n", + mode, + type_string, + group_size, + bits, + batched, + /* has_global_scale = */ false, + bm, + 32, + bn, + true) + : get_quantized_kernel_wrapped( + d, kname, "qmm_n", mode, type_string, group_size, bits, batched); } auto& compute_encoder = metal::get_command_encoder(s); compute_encoder.set_compute_pipeline_state(kernel); @@ -1249,9 +1432,31 @@ void qmm_splitk( group_size, "_b_", bits, - aligned ? "_alN_true" : "_alN_false"); - auto kernel = get_quantized_kernel_wrapped( - d, kname, "qmm_t_splitk", mode, type_string, group_size, bits, aligned); + aligned ? "_alN_true" : "_alN_false", + is_implied_bias(biases) ? "_ib" : ""); + // BM/BK/BN = 32 are the kernel defaults, spelled out for the trailing arg. + auto kernel = is_implied_bias(biases) ? get_quantized_kernel_wrapped( + d, + kname, + "qmm_t_splitk", + mode, + type_string, + group_size, + bits, + aligned, + bm, + 32, + bn, + true) + : get_quantized_kernel_wrapped( + d, + kname, + "qmm_t_splitk", + mode, + type_string, + group_size, + bits, + aligned); compute_encoder.set_compute_pipeline_state(kernel); diff --git a/mlx/ops.cpp b/mlx/ops.cpp index 40dc979fbd..dc9f7cc56c 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -9,6 +9,7 @@ #include #include +#include "mlx/backend/common/quantized.h" #include "mlx/backend/cuda/cuda.h" #include "mlx/backend/metal/metal.h" #include "mlx/fast_primitives.h" @@ -121,7 +122,9 @@ void validate_quantized_input( throw std::invalid_argument(msg.str()); } - if (biases && scales.shape() != biases->shape()) { + // An implied-bias factor carries no shape to check. + if (biases && !is_implied_bias(*biases) && + scales.shape() != biases->shape()) { std::ostringstream msg; msg << "[" << tag << "] Scales and biases should have the same shape. " << "Received scales with shape " << scales.shape() @@ -4874,8 +4877,11 @@ std::pair validate_mode_with_type( msg << "[" << tag << "] Biases must be provided for affine quantization."; throw std::invalid_argument(msg.str()); } - auto dtype = result_type(scales, *biases); - if (!issubdtype(dtype, floating)) { + // An implied-bias factor takes the scales dtype and does not promote it. + const bool implied = is_implied_bias(*biases); + auto dtype = implied ? scales.dtype() : result_type(scales, *biases); + if (!issubdtype(dtype, floating) || + !issubdtype(biases->dtype(), floating)) { std::ostringstream msg; msg << "[" << tag << "] Only real floating types are supported but " << "scales.dtype() == " << scales.dtype() @@ -4934,6 +4940,28 @@ void validate_global_scale( } } +// The implied-bias kernels give `scales * T(f)` exactly only with unpromoted +// scales, a prebuilt (1/2/4/8-bit) kernel and a CPU or Metal stream. +bool implied_bias_kernels_exact( + const array& scales, + Dtype kernel_dtype, + int bits, + const Stream& stream) { + return (stream.device == Device::cpu || + (stream.device == Device::gpu && metal::is_available())) && + kernel_dtype == scales.dtype() && + (bits == 1 || bits == 2 || bits == 4 || bits == 8); +} + +array materialize_implied_bias( + const array& scales, + const array& factor, + StreamOrDevice s) { + // The factor is a constant on every path: no gradient flows to it. + auto f = stop_gradient(astype(factor, scales.dtype(), s), s); + return multiply(scales, f, s); +} + array quantized_matmul( const array& x, const array& w, @@ -4986,14 +5014,28 @@ array quantized_matmul( if (qmode == QuantizationMode::Affine) { inputs = {astype(x, dtype), w, astype(scales, dtype)}; if (biases) { - inputs.push_back(astype(*biases, dtype)); + // Materialize a factor that the implied-bias kernels cannot honor. + auto b = (is_implied_bias(*biases) && + !implied_bias_kernels_exact(scales, dtype, bits, to_stream(s))) + ? materialize_implied_bias(scales, *biases, s) + : *biases; + inputs.push_back(astype(b, dtype, s)); } } else { inputs = {x, w, scales}; } if (x.ndim() > 2 && w.ndim() > 2) { + // Do not broadcast an implied-bias factor. + std::optional implied_bias; + if (inputs.size() == 4 && is_implied_bias(inputs[3])) { + implied_bias = std::move(inputs[3]); + inputs.pop_back(); + } inputs = broadcast_arrays(inputs, {-2, -1}, s); + if (implied_bias) { + inputs.push_back(std::move(*implied_bias)); + } } auto out_shape = inputs[0].shape(); out_shape.back() = w_outer_dims; @@ -5440,19 +5482,26 @@ array affine_dequantize( StreamOrDevice s_) { auto wshape = w.shape(); auto sshape = scales.shape(); - auto bshape = biases.shape(); - if (wshape.size() != sshape.size() || wshape.size() != bshape.size()) { + const bool implied_bias = is_implied_bias(biases); + if (wshape.size() != sshape.size() || + (!implied_bias && wshape.size() != biases.ndim())) { throw std::invalid_argument( "[dequantize] Shape of scales and biases does not match the matrix"); } wshape.back() = -1; sshape.back() = -1; - bshape.back() = -1; - - if (wshape != sshape || wshape != bshape) { + if (wshape != sshape) { throw std::invalid_argument( "[dequantize] Shape of scales and biases does not match the matrix"); } + if (!implied_bias) { + auto bshape = biases.shape(); + bshape.back() = -1; + if (wshape != bshape) { + throw std::invalid_argument( + "[dequantize] Shape of scales and biases does not match the matrix"); + } + } // Packing into uint32 int out_size = w.shape(-1) * 32 / bits; @@ -5506,7 +5555,10 @@ array affine_dequantize( wshape.push_back(group_size); w = reshape(w, wshape, s); w = multiply(w, expand_dims(scales, -1, s), s); - w = add(w, expand_dims(biases, -1, s), s); + auto b = is_implied_bias(biases) + ? materialize_implied_bias(scales, biases, s) + : biases; + w = add(w, expand_dims(b, -1, s), s); w = reshape(w, sshape, s); return {w}; @@ -5515,12 +5567,17 @@ array affine_dequantize( if (s.device == Device::gpu) { auto out_shape = w.shape(); out_shape.back() = out_size; + // The kernel reads the implied-bias factor in the scales' dtype. + array b = !implied_bias ? biases + : implied_bias_kernels_exact(scales, scales.dtype(), bits, s) + ? astype(biases, scales.dtype(), s) + : materialize_implied_bias(scales, biases, s); return array( std::move(out_shape), scales.dtype(), std::make_shared( s, fallback, group_size, bits, QuantizationMode::Affine, true), - {w, scales, biases}); + {w, scales, std::move(b)}); } return fallback({w, scales, biases})[0]; } @@ -5740,7 +5797,7 @@ array gather_qmm( const array& x, const array& w, const array& scales, - const std::optional& biases /* = std::nullopt */, + const std::optional& biases_ /* = std::nullopt */, std::optional lhs_indices_ /* = std::nullopt */, std::optional rhs_indices_ /* = std::nullopt */, bool transpose /* = true */, @@ -5756,11 +5813,16 @@ array gather_qmm( "[gather_qmm] Global scale is not supported without indices."); } return quantized_matmul( - x, w, scales, biases, transpose, group_size_, bits_, mode, s); + x, w, scales, biases_, transpose, group_size_, bits_, mode, s); } - auto [out_type, qmode] = - validate_mode_with_type("gather_qmm", scales, biases, std::nullopt, mode); + auto [out_type, qmode] = validate_mode_with_type( + "gather_qmm", scales, biases_, std::nullopt, mode); + // The gather kernels have no implied-bias variant: materialize the factor. + std::optional biases = biases_; + if (is_implied_bias(biases_)) { + biases = materialize_implied_bias(scales, *biases_, s); + } auto [group_size, bits] = quantization_params_from_mode(qmode, group_size_, bits_); auto [w_inner_dims, w_outer_dims] = extract_quantized_matmul_dims( diff --git a/mlx/ops.h b/mlx/ops.h index 017fcbd706..185bca5419 100644 --- a/mlx/ops.h +++ b/mlx/ops.h @@ -1604,7 +1604,14 @@ MLX_API array conv_transpose3d( int groups = 1, StreamOrDevice s = {}); -/** Quantized matmul multiplies x with a quantized matrix w*/ +/** Quantized matmul multiplies x with a quantized matrix w. + * + * In mode "affine", `biases` may be a 0-d array holding an implied-bias + * factor f: it stands for the per-group bias `scales * T(f)`, where T is the + * scales' dtype and f is cast to T first. A factor not exactly representable + * in T is outside this contract (the result is still that of `scales * T(f)`). + * The factor is a constant: its gradient is zero. + */ MLX_API array quantized_matmul( const array& x, const array& w, @@ -1625,7 +1632,12 @@ MLX_API std::vector quantize( const std::optional& global_scale = std::nullopt, StreamOrDevice s = {}); -/** Dequantize a matrix produced by quantize() */ +/** Dequantize a matrix produced by quantize(). + * + * In mode "affine", `biases` may be a 0-d implied-bias factor with the same + * meaning as in quantized_matmul(): bias = `scales * T(f)` in the scales' + * dtype T. + */ MLX_API array dequantize( const array& w, const array& scales, @@ -1655,7 +1667,12 @@ MLX_API array from_fp8(array x, Dtype dtype, StreamOrDevice s = {}); /** Convert a floating point matrix to E4M3 float8. */ MLX_API array to_fp8(array x, StreamOrDevice s = {}); -/** Compute matrix products with matrix-level gather. */ +/** Compute matrix products with matrix-level gather. + * + * In mode "affine", `biases` may be a 0-d implied-bias factor with the same + * meaning as in quantized_matmul(): bias = `scales * T(f)` in the scales' + * dtype T. + */ MLX_API array gather_qmm( const array& x, const array& w, diff --git a/mlx/primitives.cpp b/mlx/primitives.cpp index b9dfd918f8..a6dc0a643f 100644 --- a/mlx/primitives.cpp +++ b/mlx/primitives.cpp @@ -10,6 +10,7 @@ #include #include +#include "mlx/backend/common/quantized.h" #include "mlx/backend/common/utils.h" #include "mlx/fft.h" #include "mlx/linalg.h" @@ -3604,6 +3605,9 @@ std::vector QuantizedMatmul::vjp( // -scale for 2-bit. The reverse products run on the biased kernels. bool sym = mode_ == QuantizationMode::Affine && primals.size() == 3; float sym_c = (bits_ == 1) ? -0.5f : -1.0f; + // Implied bias: primals[3] is a 0-d factor f and bias = f * scale. + bool implied = mode_ == QuantizationMode::Affine && primals.size() == 4 && + is_implied_bias(primals[3]); std::optional biases = std::nullopt; if (mode_ == QuantizationMode::Affine) { biases = sym @@ -3632,6 +3636,9 @@ std::vector QuantizedMatmul::vjp( else if (arg == 1) { throw std::runtime_error( "[QuantizedMatmul::vjp] no gradient wrt the quantized weights."); + } else if (arg == 3 && implied) { + // The implied-bias factor is a constant. + vjps.push_back(zeros_like(primals[3], stream())); } else { if (mode_ != QuantizationMode::Affine) { std::ostringstream msg; @@ -3652,15 +3659,18 @@ std::vector QuantizedMatmul::vjp( // biases vjps.push_back(sum(*dsb, -1, false, stream())); } else { - // scales: d(w_hat)/d(scale) is q, or q + c for the derived bias + // scales: d(w_hat)/d(scale) is q, or q + c for a derived/implied bias + // The factor is a constant, so keep it out of higher-order grads. + array bias_c = sym ? array(sym_c, primals[2].dtype()) + : implied ? astype( + stop_gradient(primals[3], stream()), + primals[2].dtype(), + stream()) + : array(0.0f, primals[2].dtype()); auto wq = dequantize( primals[1], ones_like(primals[2], stream()), - full( - primals[2].shape(), - sym ? sym_c : 0.0f, - primals[2].dtype(), - stream()), + full(primals[2].shape(), bias_c, primals[2].dtype(), stream()), group_size_, bits_, quantization_mode_to_string(mode_), @@ -3679,9 +3689,17 @@ std::vector QuantizedMatmul::jvp( const std::vector& primals, const std::vector& tangents, const std::vector& argnums) { - if (argnums.size() > 1 || argnums[0] != 0) { - throw std::runtime_error( - "[QuantizedMatmul::jvp] No JVP wrt the quantized matrix yet."); + bool implied = mode_ == QuantizationMode::Affine && primals.size() == 4 && + is_implied_bias(primals[3]); + for (auto arg : argnums) { + if (arg != 0 && !(implied && arg == 3)) { + throw std::runtime_error( + "[QuantizedMatmul::jvp] No JVP wrt the quantized matrix yet."); + } + } + if (argnums[0] != 0) { + // The implied-bias factor is a constant. + return {zeros(output_shapes(primals)[0], primals[0].dtype(), stream())}; } return {quantized_matmul( tangents[0], diff --git a/python/src/ops.cpp b/python/src/ops.cpp index 4d83c6293c..494c539430 100644 --- a/python/src/ops.cpp +++ b/python/src/ops.cpp @@ -4802,7 +4802,11 @@ void init_ops(nb::module_& m) { w (array): Quantized matrix packed in unsigned integers scales (array): The scales to use per ``group_size`` elements of ``w`` biases (array, optional): The biases to use per ``group_size`` - elements of ``w``. Default: ``None``. + elements of ``w``. In mode ``"affine"`` a 0-d array is an + implied-bias factor ``f`` standing for ``scales * T(f)`` in the + scales' dtype ``T`` (``f`` is cast to ``T`` first; a factor that is + not exactly representable in ``T`` is outside this contract). The + factor is a constant: its gradient is zero. Default: ``None``. transpose (bool, optional): Defines whether to multiply with the transposed ``w`` or not, namely whether we are performing ``x @ w.T`` or ``x @ w``. Default: ``True``. @@ -4937,7 +4941,9 @@ void init_ops(nb::module_& m) { w (array): Matrix to be dequantized scales (array): The scales to use per ``group_size`` elements of ``w``. biases (array, optional): The biases to use per ``group_size`` - elements of ``w``. Default: ``None``. + elements of ``w``. In mode ``"affine"`` a 0-d array is an + implied-bias factor ``f`` standing for ``scales * T(f)`` in the + scales' dtype ``T`` (see :func:`quantized_matmul`). Default: ``None``. group_size (int, optional): The size of the group in ``w`` that shares a scale and bias. See supported values and defaults in the :ref:`table of quantization modes `. Default: ``None``. @@ -4995,14 +5001,17 @@ void init_ops(nb::module_& m) { all but the last two dimensions) of ``x`` and ``w`` respectively. Note that ``scales`` and ``biases`` must have the same batch dimensions - as ``w`` since they represent the same quantized matrix. + as ``w`` since they represent the same quantized matrix. An implied-bias + factor (a 0-d ``biases``) has no batch dimensions. Args: x (array): Input array w (array): Quantized matrix packed in unsigned integers scales (array): The scales to use per ``group_size`` elements of ``w`` biases (array, optional): The biases to use per ``group_size`` - elements of ``w``. Default: ``None``. + elements of ``w``. In mode ``"affine"`` a 0-d array is an + implied-bias factor ``f`` standing for ``scales * T(f)`` in the + scales' dtype ``T`` (see :func:`quantized_matmul`). Default: ``None``. lhs_indices (array, optional): Integer indices for ``x``. Default: ``None``. rhs_indices (array, optional): Integer indices for ``w``. Default: ``None``. transpose (bool, optional): Defines whether to multiply with the diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index ab62833cfc..e08eb2e40d 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -8,6 +8,7 @@ from itertools import product import mlx.core as mx +import numpy as np import mlx_tests @@ -1278,6 +1279,534 @@ def f_biased(x, scales): _, jvp_biased = mx.jvp(lambda x: f_biased(x, scales), [x], [x]) self.assertLess((jvp_sym[0] - jvp_biased[0]).abs().max().item(), 1e-3) + # Implied bias: a 0-d `biases` array is a factor f and the kernels form + # bias = f * scale per group. The factors of the shipped checkpoints are + # powers of two, so f * scale is exact in every dtype and the implied path + # must reproduce the rebuilt-bias path bit for bit. + IMPLIED_BIAS_CASES = [ + (2, 64, -2.0), # QAT 2-bit + (4, 64, -8.0), # QAT 4-bit + (8, 64, -128.0), # QAT 8-bit + (1, 128, -0.5), # Bonsai 1-bit + (2, 128, -1.0), # Ternary Bonsai 2 2-bit + ] + IMPLIED_BIAS_DTYPES = [mx.float32, mx.float16, mx.bfloat16] + + @staticmethod + def implied_bias_devices(): + return [mx.cpu, mx.gpu] if mx.metal.is_available() else [mx.cpu] + + @staticmethod + def reference_bias(scales, factor): + # The implied-bias contract: a 0-d factor f stands for the materialized + # bias `scales * T(f)` in the scales' storage dtype T (f cast to T first). + return scales * mx.array(factor).astype(scales.dtype) + + def assert_bitwise_equal(self, a, b): + # Unlike assertEqualArray(atol=0), equal inf/nan compare equal. + self.assertEqual(a.dtype, b.dtype) + self.assertEqual(a.shape, b.shape) + mx.eval(a, b) + self.assertTrue( + mx.array_equal(a, b, equal_nan=True).item(), + f"max |diff| = {(a.astype(mx.float32) - b.astype(mx.float32)).abs().max().item()}", + ) + + def assert_implied_bias_matches(self, op, pre, full, factor, *post, **kwargs): + # op(*pre, biases, *post): the 0-d factor must reproduce `full` bitwise. + self.assert_bitwise_equal( + op(*pre, factor, *post, **kwargs), op(*pre, full, *post, **kwargs) + ) + + def test_implied_bias_matmul(self): + key = mx.random.key(3) + k1, k2 = mx.random.split(key) + for dev in self.implied_bias_devices(): + # M = 1..8 (qmv / qmv_quad / qmv_wide), 33 (above every GPU's qmv + # batch limit: qmm / qmm_splitk) and a prefill-sized 512. + Ms = [1, 2, 3, 4, 5, 6, 7, 8, 33, 512] + for dtype, (bits, gs, factor) in product( + self.IMPLIED_BIAS_DTYPES, self.IMPLIED_BIAS_CASES + ): + # K = 64 / 128 route to qmv_quad; 512 / 1536 to qmv_fast (and + # qmv_wide for 2 <= M <= 8 at bits > 1); N = 72 is not a + # multiple of the 32-wide qmm tile. + for K, N in product([64, 128, 512, 1536], [256, 72]): + if K < gs: + continue + w = mx.random.normal(shape=(N, K), key=k2).astype(dtype) + w_q, scales, _ = mx.quantize(w, gs, bits, stream=dev) + full = self.reference_bias(scales, factor) + f = mx.array(factor, dtype) + for M in Ms: + with self.subTest( + dev=dev, dtype=dtype, bits=bits, gs=gs, K=K, N=N, M=M + ): + x = mx.random.normal(shape=(M, K), key=k1).astype(dtype) + self.assert_implied_bias_matches( + mx.quantized_matmul, + (x, w_q, scales), + full, + f, + True, + gs, + bits, + stream=dev, + ) + # Non-transposed (qvm): inner dim N < 1024. + for M in [1, 3]: + with self.subTest( + dev=dev, + dtype=dtype, + bits=bits, + gs=gs, + K=K, + N=N, + M=M, + qvm=True, + ): + xt = mx.random.normal(shape=(M, N), key=k1).astype(dtype) + self.assert_implied_bias_matches( + mx.quantized_matmul, + (xt, w_q, scales), + full, + f, + False, + gs, + bits, + stream=dev, + ) + + @unittest.skipUnless(mx.metal.is_available(), "GPU dispatch paths") + def test_implied_bias_matmul_gpu_paths(self): + # The remaining GPU dispatch routes: qvm_split_k (non-transposed with a + # large inner dim), batched qmv / qmm (3-d x; the 0-d factor is never + # broadcast and its zero batch strides keep it in place), the + # non-aligned qmv (N % 8 != 0) and a factor given in another dtype. + key = mx.random.key(4) + k1, k2 = mx.random.split(key) + for dtype, (bits, gs, factor) in product( + self.IMPLIED_BIAS_DTYPES, self.IMPLIED_BIAS_CASES + ): + f = mx.array(factor, dtype) + + def check(x, w_q, scales, full, transpose=True, factor=f, **sub): + with self.subTest(dtype=dtype, bits=bits, gs=gs, **sub): + self.assert_implied_bias_matches( + mx.quantized_matmul, + (x, w_q, scales), + full, + factor, + transpose, + gs, + bits, + ) + + # qvm_split_k: w is (1536, 256), x @ w_hat with inner dim 1536. + w = mx.random.normal(shape=(1536, 256), key=k2).astype(dtype) + w_q, scales, _ = mx.quantize(w, gs, bits) + full = self.reference_bias(scales, factor) + for M in [1, 3]: + xt = mx.random.normal(shape=(M, 1536), key=k1).astype(dtype) + check(xt, w_q, scales, full, False, M=M, path="qvm_split_k") + + # Batched: x is (B, M, K) over a 2-d w (broadcast to B); qmv for + # M = 1, qmv_quad for K = 128, qmm for M = 64. + for K in [128, 512]: + if K < gs: + continue + w = mx.random.normal(shape=(256, K), key=k2).astype(dtype) + w_q, scales, _ = mx.quantize(w, gs, bits) + full = self.reference_bias(scales, factor) + for M in [1, 3, 64]: + x = mx.random.normal(shape=(3, M, K), key=k1).astype(dtype) + check(x, w_q, scales, full, K=K, M=M, path="batched") + # Batched w as well (B, N, K) with per-batch scales. + wb = mx.random.normal(shape=(2, 256, K), key=k2).astype(dtype) + wb_q, scales_b, _ = mx.quantize(wb, gs, bits) + full_b = self.reference_bias(scales_b, factor) + for M in [1, 64]: + x = mx.random.normal(shape=(2, M, K), key=k1).astype(dtype) + check(x, wb_q, scales_b, full_b, K=K, M=M, path="batched_w") + + # N % 8 != 0: qmv_fast with partial rows where K is fast-aligned + # (2/4/8-bit here), else the generic qmv (1-bit); qmv_wide for + # M = 4 where the device routes affine there. + w = mx.random.normal(shape=(67, 512), key=k2).astype(dtype) + w_q, scales, _ = mx.quantize(w, gs, bits) + full = self.reference_bias(scales, factor) + for M in [1, 4]: + x = mx.random.normal(shape=(M, 512), key=k1).astype(dtype) + check(x, w_q, scales, full, M=M, path="qmv_generic") + + # A float32 factor with half-precision scales adopts the scales' + # dtype (no promotion of the whole matmul to float32). + if dtype != mx.float32: + x = mx.random.normal(shape=(2, 512), key=k1).astype(dtype) + check( + x, w_q, scales, full, factor=mx.array(factor), path="factor_dtype" + ) + + def test_implied_bias_subnormal_scales(self): + # Subnormal fp16 scales (Bonsai-27B-1bit / Ternary Bonsai 2 carry a few): + # f * scale with f = -0.5 and an odd subnormal mantissa is not + # representable in fp16, so the kernels must round the product through + # the scales dtype to stay bit-identical to a materialized bias. + for bits, gs, factor in [(1, 128, -0.5), (2, 128, -1.0), (2, 64, -2.0)]: + K, N = 512, 72 + w = mx.random.randint(0, 2**31 - 1, (N, K * bits // 32)).astype(mx.uint32) + rng = np.random.default_rng(0) + s = rng.uniform(0.01, 1.0, (N, K // gs)).astype(np.float16) + idx = rng.choice(s.size, 200, replace=False) + s.flat[idx] = (rng.integers(1, 1024, 200) * 2.0**-24).astype(np.float16) + scales = mx.array(s) + f = mx.array(factor, dtype=mx.float16) + rebuilt = scales * f + for dev in self.implied_bias_devices(): + with self.subTest(bits=bits, factor=factor, dev=dev): + self.assert_implied_bias_matches( + mx.dequantize, (w, scales), rebuilt, f, gs, bits, stream=dev + ) + for M in [1, 2, 3, 8, 64]: + x = mx.random.normal((M, K)).astype(mx.float16) + self.assert_implied_bias_matches( + mx.quantized_matmul, + (x, w, scales), + rebuilt, + f, + True, + gs, + bits, + stream=dev, + ) + + def test_implied_bias_dequantize(self): + key = mx.random.key(5) + for dev, dtype, (bits, gs, factor) in product( + self.implied_bias_devices(), + self.IMPLIED_BIAS_DTYPES, + self.IMPLIED_BIAS_CASES, + ): + with self.subTest(dev=dev, dtype=dtype, bits=bits, gs=gs): + w = mx.random.normal(shape=(64, 512), key=key).astype(dtype) + w_q, scales, _ = mx.quantize(w, gs, bits, stream=dev) + full = self.reference_bias(scales, factor) + f = mx.array(factor, dtype) + w_full = mx.dequantize(w_q, scales, full, gs, bits, stream=dev) + w_ib = mx.dequantize(w_q, scales, f, gs, bits, stream=dev) + self.assert_bitwise_equal(w_ib, w_full) + # A float32 factor is cast to the scales' dtype. + w_ib32 = mx.dequantize( + w_q, scales, mx.array(factor), gs, bits, stream=dev + ) + self.assert_bitwise_equal(w_ib32, w_full) + # Embedding-style gather: dequantize gathered rows with the + # scalar factor (no gather of a bias table). + idx = mx.array([3, 0, 63, 3]) + e_full = mx.dequantize( + w_q[idx], scales[idx], full[idx], gs, bits, stream=dev + ) + e_ib = mx.dequantize(w_q[idx], scales[idx], f, gs, bits, stream=dev) + self.assert_bitwise_equal(e_ib, e_full) + + def test_implied_bias_gather_qmm(self): + # gather_qmm has no implied-bias kernels: the op materializes the bias, + # so the result equals the per-group-bias call bit for bit. + key = mx.random.key(7) + k1, k2 = mx.random.split(key) + lhs = mx.array([0, 1, 1]) + rhs = mx.array([1, 0, 1]) + for dev, dtype, (bits, gs, factor) in product( + self.implied_bias_devices(), + self.IMPLIED_BIAS_DTYPES, + self.IMPLIED_BIAS_CASES, + ): + with self.subTest(dev=dev, dtype=dtype, bits=bits, gs=gs): + x = mx.random.normal(shape=(2, 3, 512), key=k1).astype(dtype) + w = mx.random.normal(shape=(2, 64, 512), key=k2).astype(dtype) + w_q, scales, _ = mx.quantize(w, gs, bits) + full = self.reference_bias(scales, factor) + f = mx.array(factor, dtype) + self.assert_implied_bias_matches( + mx.gather_qmm, + (x, w_q, scales), + full, + f, + lhs, + rhs, + True, + gs, + bits, + stream=dev, + ) + + def test_implied_bias_grad(self): + # Reverse products wrt x and scales use the implied-bias kernels and + # the same closed form as the explicit-bias call. + key = mx.random.key(6) + k1, k2, k3 = mx.random.split(key, 3) + for bits, gs, factor in self.IMPLIED_BIAS_CASES: + with self.subTest(bits=bits, gs=gs): + x = mx.random.normal(shape=(2, 512), key=k1) + w = mx.random.normal(shape=(64, 512), key=k2) + w_q, scales, _ = mx.quantize(w, gs, bits) + f = mx.array(factor) + cotan = mx.random.normal(shape=(2, 64), key=k3) + + def f_ib(x, scales): + return mx.quantized_matmul(x, w_q, scales, f, True, gs, bits) + + def f_full(x, scales): + return mx.quantized_matmul( + x, w_q, scales, factor * scales, True, gs, bits + ) + + _, vjp_ib = mx.vjp(f_ib, [x, scales], [cotan]) + _, vjp_full = mx.vjp(f_full, [x, scales], [cotan]) + self.assert_bitwise_equal(vjp_ib[0], vjp_full[0]) + # Same closed form, different fp32 accumulation order (the + # explicit path sums the bias term through the multiply's vjp). + self.assertTrue( + mx.allclose(vjp_ib[1], vjp_full[1], rtol=1e-4, atol=1e-3) + ) + + _, jvp_ib = mx.jvp(lambda x: f_ib(x, scales), [x], [x]) + _, jvp_full = mx.jvp(lambda x: f_full(x, scales), [x], [x]) + self.assert_bitwise_equal(jvp_ib[0], jvp_full[0]) + + def test_implied_bias_factor_grad(self): + # The factor gradient is zero where the implied-bias kernels run (4-bit, + # unpromoted scales) and where the bias is materialized (e.g. 3-bit). + key = mx.random.key(8) + k1, k2 = mx.random.split(key) + lhs = mx.array([0, 1]) + rhs = mx.array([1, 0]) + configs = [ + (mx.float32, mx.float32, 4), + (mx.float16, mx.float16, 4), + (mx.float32, mx.float16, 4), + (mx.float32, mx.float32, 3), + ] + for dev, (x_dtype, s_dtype, bits) in product( + self.implied_bias_devices(), configs + ): + factor = -float(2 ** (bits - 1)) + w = mx.random.normal(shape=(2, 64, 512), key=k2).astype(s_dtype) + w_q, scales, _ = mx.quantize(w, 64, bits, stream=dev) + x = mx.random.normal(shape=(2, 1, 512), key=k1).astype(x_dtype) + ops = { + "quantized_matmul": lambda f: mx.quantized_matmul( + x[0], w_q[0], scales[0], f, True, 64, bits, stream=dev + ), + "gather_qmm": lambda f: mx.gather_qmm( + x, w_q, scales, f, lhs, rhs, True, 64, bits, stream=dev + ), + "dequantize": lambda f: mx.dequantize( + w_q[0], scales[0], f, 64, bits, stream=dev + ), + } + for name, op in ops.items(): + with self.subTest( + dev=dev, x_dtype=x_dtype, s_dtype=s_dtype, bits=bits, op=name + ): + f = mx.array(factor) + out = op(f) + _, (df,) = mx.vjp(op, [f], [mx.ones_like(out)]) + self.assert_bitwise_equal(df, mx.zeros_like(f)) + self.assertEqual( + mx.grad(lambda f: op(f).astype(mx.float32).sum())(f).item(), + 0.0, + ) + with self.subTest(dev=dev, x_dtype=x_dtype, s_dtype=s_dtype, bits=bits): + op = ops["quantized_matmul"] + f = mx.array(factor) + (out,), (tf,) = mx.jvp(op, [f], [mx.array(1.0)]) + self.assert_bitwise_equal(tf, mx.zeros_like(out)) + + def test_implied_bias_factor_mixed_grad(self): + # The scales and x gradients must not depend on the factor either, on + # the kernel paths (4-bit) and on the materialized path (3-bit). + key = mx.random.key(9) + k1, k2 = mx.random.split(key) + configs = [ + (mx.float32, mx.float32, 4), + (mx.float16, mx.float16, 4), + (mx.float32, mx.float32, 3), + ] + for dev, (x_dtype, s_dtype, bits) in product( + self.implied_bias_devices(), configs + ): + factor = -float(2 ** (bits - 1)) + w = mx.random.normal(shape=(64, 512), key=k2).astype(s_dtype) + w_q, scales, _ = mx.quantize(w, 64, bits, stream=dev) + x = mx.random.normal(shape=(1, 512), key=k1).astype(x_dtype) + + def loss(x, scales, f): + out = mx.quantized_matmul(x, w_q, scales, f, True, 64, bits, stream=dev) + return out.astype(mx.float32).sum() + + def dscales_sum(f): + return mx.grad(loss, argnums=1)(x, scales, f).astype(mx.float32).sum() + + def dx_sum(f): + return mx.grad(loss, argnums=0)(x, scales, f).astype(mx.float32).sum() + + f = mx.array(factor) + for name, fn in [("scales", dscales_sum), ("x", dx_sum)]: + with self.subTest( + dev=dev, x_dtype=x_dtype, s_dtype=s_dtype, bits=bits, grad=name + ): + self.assertEqual(mx.grad(fn)(f).item(), 0.0) + + def test_implied_bias_dtype_promotion(self): + # fp32 activations over fp16 scales promote the kernel's scales to fp32. + # The bias must still be the fp16 reference `scales * T(f)`: with the + # smallest fp16 subnormal 2^-24 and f = -0.5, that is 0, whereas an + # fp32-formed bias is -2^-25. + tiny = 2.0**-24 + scales = mx.full((8, 1), tiny, dtype=mx.float16) + w = mx.zeros((8, 16), dtype=mx.uint32) + x = mx.ones((1, 128), dtype=mx.float32) + for dev in self.implied_bias_devices(): + for f in [mx.array(-0.5, mx.float16), mx.array(-0.5)]: + with self.subTest(dev=dev, fdtype=f.dtype): + y_full = mx.quantized_matmul( + x, + w, + scales, + self.reference_bias(scales, -0.5), + True, + 128, + 4, + stream=dev, + ) + y_ib = mx.quantized_matmul( + x, w, scales, f, True, 128, 4, stream=dev + ) + self.assertEqual(y_ib.dtype, mx.float32) + self.assert_bitwise_equal(y_ib, y_full) + self.assert_bitwise_equal(y_ib, mx.zeros((1, 8))) + # bf16 activations over fp16 scales promote to fp32 as well. + xb = mx.ones((1, 128), dtype=mx.bfloat16) + for dev in self.implied_bias_devices(): + y_ib = mx.quantized_matmul( + xb, w, scales, mx.array(-0.5, mx.float16), True, 128, 4, stream=dev + ) + self.assertEqual(y_ib.dtype, mx.float32) + self.assert_bitwise_equal(y_ib, mx.zeros((1, 8))) + + def test_implied_bias_cpu_dequantize_rounding(self): + # The CPU fallback forms the bias in the scales' dtype even when the + # factor arrives as fp32: every code is 1, so the reference is + # 2^-24 + T(-0.5 * 2^-24) = 2^-24; an fp32 product would leave 2^-25, + # which the final fp16 cast rounds to 0. + tiny = 2.0**-24 + scales = mx.full((8, 1), tiny, dtype=mx.float16) + w = mx.full((8, 16), 0x11111111, dtype=mx.uint32) + expected = mx.full((8, 128), tiny, dtype=mx.float16) + for dev in self.implied_bias_devices(): + with self.subTest(dev=dev): + d_full = mx.dequantize( + w, scales, self.reference_bias(scales, -0.5), 128, 4, stream=dev + ) + d_ib = mx.dequantize(w, scales, mx.array(-0.5), 128, 4, stream=dev) + self.assert_bitwise_equal(d_ib, d_full) + self.assert_bitwise_equal(d_ib, expected) + + def test_implied_bias_non_representable_factor(self): + # A factor that is not exactly representable in T is outside the + # contract; the documented result is still `scales * T(f)`: with the + # fp16 scale 0.10009765625 and f = 1.0004, T(f) = 1 so the bias is the + # scale itself (T(f * scale) would be the next fp16 value up). + scale = 0.10009765625 + scales = mx.full((8, 1), scale, dtype=mx.float16) + f = mx.array(1.0004) + rebuilt = self.reference_bias(scales, 1.0004) + self.assert_bitwise_equal(rebuilt, scales) + other = (scales.astype(mx.float32) * f).astype(mx.float16) + self.assertFalse(mx.array_equal(rebuilt, other).item()) + w = mx.zeros((8, 16), dtype=mx.uint32) + x = mx.ones((1, 128), dtype=mx.float16) + for dev in self.implied_bias_devices(): + with self.subTest(dev=dev): + d_ib = mx.dequantize(w, scales, f, 128, 4, stream=dev) + self.assert_bitwise_equal( + d_ib, mx.full((8, 128), scale, dtype=mx.float16) + ) + self.assert_implied_bias_matches( + mx.quantized_matmul, + (x, w, scales), + rebuilt, + f, + True, + 128, + 4, + stream=dev, + ) + + def test_implied_bias_non_power_of_two_bits(self): + # 3/5/6-bit have no `_ib` kernels in the non-JIT Metal library, so the + # op materializes the bias and the ordinary kernels run -- on the GPU + # (where a missing kernel name would throw at dispatch) and the CPU. + key = mx.random.key(190) + k1, k2 = mx.random.split(key) + for dev, dtype, bits in product( + self.implied_bias_devices(), self.IMPLIED_BIAS_DTYPES, [3, 5, 6] + ): + factor = -float(2 ** (bits - 1)) + w = mx.random.normal(shape=(72, 512), key=k2).astype(dtype) + w_q, scales, _ = mx.quantize(w, 64, bits) + full = self.reference_bias(scales, factor) + f = mx.array(factor, dtype) + with self.subTest(dev=dev, dtype=dtype, bits=bits, op="dequantize"): + self.assert_implied_bias_matches( + mx.dequantize, (w_q, scales), full, f, 64, bits, stream=dev + ) + for M, transpose in [ + (1, True), + (3, True), + (8, True), + (33, True), + (2, False), + ]: + with self.subTest( + dev=dev, dtype=dtype, bits=bits, M=M, qvm=not transpose + ): + shape = (M, 512) if transpose else (M, 72) + x = mx.random.normal(shape=shape, key=k1).astype(dtype) + self.assert_implied_bias_matches( + mx.quantized_matmul, + (x, w_q, scales), + full, + f, + transpose, + 64, + bits, + stream=dev, + ) + + def test_implied_bias_throws(self): + x = mx.random.normal(shape=(2, 512)) + w = mx.random.normal(shape=(64, 512)) + w_q, scales, _ = mx.quantize(w, 64, 4) + f = mx.array(-8.0) + # Only a 0-d biases array is a factor; a 1-d one is a shape error. + with self.assertRaises(ValueError): + mx.quantized_matmul(x, w_q, scales, mx.array([-8.0]), True, 64, 4) + # The factor must be a real floating type. + with self.assertRaises(ValueError): + mx.quantized_matmul(x, w_q, scales, mx.array(-8), True, 64, 4) + # So must a per-group biases array. + int_biases = mx.zeros(scales.shape, dtype=mx.int32) + with self.assertRaises(ValueError): + mx.quantized_matmul(x, w_q, scales, int_biases, True, 64, 4) + with self.assertRaises(ValueError): + mx.dequantize(w_q, scales, int_biases, 64, 4) + # The fp modes have no biases. + w_q4, scales4 = mx.quantize(w, mode="mxfp4") + with self.assertRaises(ValueError): + mx.quantized_matmul(x, w_q4, scales4, f, True, mode="mxfp4") + def test_qvm(self): key = mx.random.key(0) k1, k2 = mx.random.split(key) diff --git a/tests/ops_tests.cpp b/tests/ops_tests.cpp index b7dab2163e..06792c1d26 100644 --- a/tests/ops_tests.cpp +++ b/tests/ops_tests.cpp @@ -3683,6 +3683,348 @@ TEST_CASE("test bias-free affine quantized_matmul decode") { } } +// Implied bias: a 0-d `biases` array holding a factor f stands for the +// materialized bias `scales * T(f)`, formed in the scales' storage dtype T +// (ops.h). The implied path must reproduce that reference bit for bit. +namespace { +std::vector implied_bias_devices() { + std::vector devices = {Device::cpu}; + if (is_available(Device::gpu)) { + devices.push_back(Device::gpu); + } + return devices; +} +array implied_bias_reference(const array& scales, const array& f) { + auto b = + multiply(scales, astype(f, scales.dtype(), Device::cpu), Device::cpu); + eval(b); + return b; +} +bool same_array(const array& a, const array& b) { + return a.dtype() == b.dtype() && a.shape() == b.shape() && + array_equal(a, b, Device::cpu).item(); +} +} // namespace + +TEST_CASE("test implied-bias affine quantized_matmul and dequantize") { + const auto cpu = Device::cpu; + const auto devices = implied_bias_devices(); + + struct Case { + int bits; + int group_size; + float factor; + }; + const Case cases[] = { + {2, 64, -2.0f}, + {4, 64, -8.0f}, + {8, 64, -128.0f}, + {1, 128, -0.5f}, + {2, 128, -1.0f}, + }; + + int seed = 0; + for (const auto& dev : devices) { + for (auto dtype : {float16, float32}) { + for (const auto& c : cases) { + for (int K : {128, 512, 1536}) { + const int N = 72; + auto w = + random::bits({N, K * c.bits / 32}, 4, random::key(seed++), cpu); + auto scales = astype( + random::uniform( + 0.01f, + 1.0f, + {N, K / c.group_size}, + float32, + random::key(seed++), + cpu), + dtype, + cpu); + auto factor = array(c.factor, dtype); + auto rebuilt = implied_bias_reference(scales, factor); + + // dequantize + auto d_full = dequantize( + w, + scales, + rebuilt, + c.group_size, + c.bits, + "affine", + std::nullopt, + std::nullopt, + dev); + auto d_ib = dequantize( + w, + scales, + factor, + c.group_size, + c.bits, + "affine", + std::nullopt, + std::nullopt, + dev); + CHECK(same_array(d_ib, d_full)); + + for (int M : {1, 2, 3, 5, 8, 33, 128}) { + auto x = astype( + random::normal({M, K}, float32, random::key(seed++), cpu), + dtype, + cpu); + auto o_full = quantized_matmul( + x, + w, + scales, + rebuilt, + true, + c.group_size, + c.bits, + "affine", + dev); + auto o_ib = quantized_matmul( + x, + w, + scales, + factor, + true, + c.group_size, + c.bits, + "affine", + dev); + CHECK(same_array(o_ib, o_full)); + } + // Non-transposed (x @ w): w is (K_in, N_out) packed along N. + { + auto xn = astype( + random::normal({3, N}, float32, random::key(seed++), cpu), + dtype, + cpu); + auto o_full = quantized_matmul( + xn, + w, + scales, + rebuilt, + false, + c.group_size, + c.bits, + "affine", + dev); + auto o_ib = quantized_matmul( + xn, + w, + scales, + factor, + false, + c.group_size, + c.bits, + "affine", + dev); + CHECK(same_array(o_ib, o_full)); + } + } + } + } + } + + // gather_qmm has no implied-bias kernels: the factor is materialized and + // the result equals the per-group-bias call. + { + auto w = random::bits({2, 8, 512 * 4 / 32}, 4, random::key(seed++), cpu); + auto scales = random::uniform( + 0.5f, 1.5f, {2, 8, 8}, float32, random::key(seed++), cpu); + auto x = random::normal({2, 3, 512}, float32, random::key(seed++), cpu); + auto factor = array(-8.0f); + auto rebuilt = implied_bias_reference(scales, factor); + auto lhs = array({0, 1, 1}, int32); + auto rhs = array({1, 0, 1}, int32); + for (const auto& dev : devices) { + auto o_full = gather_qmm( + x, + w, + scales, + rebuilt, + lhs, + rhs, + true, + 64, + 4, + "affine", + std::nullopt, + false, + dev); + auto o_ib = gather_qmm( + x, + w, + scales, + factor, + lhs, + rhs, + true, + 64, + 4, + "affine", + std::nullopt, + false, + dev); + CHECK(same_array(o_ib, o_full)); + } + } + + // A 1-d bias of the wrong shape is rejected. + { + auto w = random::bits({8, 512 * 4 / 32}, 4, random::key(seed++), cpu); + auto scales = + random::uniform(0.5f, 1.5f, {8, 8}, float32, random::key(seed++), cpu); + auto x = random::normal({1, 512}, float32, random::key(seed++), cpu); + CHECK_THROWS_AS( + quantized_matmul( + x, w, scales, array({-8.0f}), true, 64, 4, "affine", cpu), + std::invalid_argument); + } +} + +TEST_CASE("test implied-bias factor rounding contract") { + // Each block below is a case where the bias could be formed elsewhere -- in + // the activations' promoted dtype, in the factor's wider dtype, or through + // an `_ib` kernel the non-JIT library does not ship -- and must still match + // the `scales * T(f)` reference. + const auto cpu = Device::cpu; + const auto devices = implied_bias_devices(); + const auto& reference_bias = implied_bias_reference; + const auto& same = same_array; + + // fp16 scales of 2^-24 (the smallest fp16 subnormal): -0.5 * 2^-24 is not + // representable in fp16 and rounds to 0, while a bias formed in fp32 is + // -2^-25. 4-bit, group size 128, one group per row, 8 rows. + const int bits = 4; + const int gs = 128; + auto tiny = array(static_cast(std::ldexp(1.0, -24)), float16); + auto scales = full({8, 1}, tiny, cpu); + eval(scales); + + // 1. Promotion: fp32 activations over fp16 scales. The kernel sees fp32 + // scales, so the factor is materialized in fp16 first. + { + auto w = zeros({8, 16}, uint32, cpu); + auto x = ones({1, 128}, float32, cpu); + for (auto f : {array(-0.5f, float16), array(-0.5f, float32)}) { + auto rebuilt = reference_bias(scales, f); + CHECK(same(rebuilt, zeros({8, 1}, float16, cpu))); + for (const auto& dev : devices) { + auto o_full = quantized_matmul( + x, w, scales, rebuilt, true, gs, bits, "affine", dev); + auto o_ib = + quantized_matmul(x, w, scales, f, true, gs, bits, "affine", dev); + CHECK_EQ(o_ib.dtype(), float32); + CHECK(same(o_ib, o_full)); + CHECK(same(o_ib, zeros({1, 8}, float32, cpu))); + } + } + } + + // 2. Dequantize with an fp32 factor: every code is 1 (0x11111111), so the + // reference is 1 * 2^-24 + T(-0.5 * 2^-24) = 2^-24 in fp16; a bias + // formed in fp32 gives 2^-24 - 2^-25 = 2^-25, which the final fp16 cast + // rounds to 0 (the CPU fallback used to do exactly that). + { + auto w = full({8, 16}, array(0x11111111u, uint32), cpu); + auto f = array(-0.5f, float32); + auto rebuilt = reference_bias(scales, f); + auto expected = full({8, 128}, tiny, cpu); + for (const auto& dev : devices) { + auto d_full = dequantize( + w, + scales, + rebuilt, + gs, + bits, + "affine", + std::nullopt, + std::nullopt, + dev); + auto d_ib = dequantize( + w, scales, f, gs, bits, "affine", std::nullopt, std::nullopt, dev); + CHECK(same(d_ib, d_full)); + CHECK(same(d_ib, expected)); + } + } + + // 3. A factor that is not exactly representable in T: the reference is + // `scales * T(f)`, not `T(f * scales)`. fp16 scale 0.10009765625 with + // f = 1.0004: T(f) = 1 in fp16, so the bias is the scale itself, whereas + // T(f * scale) is the next fp16 value up (0.10015869140625). + { + auto scale = array(0.10009765625f, float16); + auto s16 = full({8, 1}, scale, cpu); + auto f = array(1.0004f, float32); + auto rebuilt = reference_bias(s16, f); + CHECK(same(rebuilt, s16)); + auto other = + astype(multiply(astype(s16, float32, cpu), f, cpu), float16, cpu); + CHECK_FALSE(array_equal(rebuilt, other, cpu).item()); + auto w = zeros({8, 16}, uint32, cpu); + auto x = ones({1, 128}, float16, cpu); + for (const auto& dev : devices) { + auto d_ib = dequantize( + w, s16, f, gs, bits, "affine", std::nullopt, std::nullopt, dev); + CHECK(same(d_ib, full({8, 128}, scale, cpu))); + auto o_full = + quantized_matmul(x, w, s16, rebuilt, true, gs, bits, "affine", dev); + auto o_ib = quantized_matmul(x, w, s16, f, true, gs, bits, "affine", dev); + CHECK(same(o_ib, o_full)); + } + } + + // 4. Bit widths without a non-JIT `_ib` kernel (3/5/6): the bias is + // materialized and the ordinary kernels run, on both devices, for every + // scales dtype. + { + int seed = 190; + for (int b : {3, 5, 6}) { + const float factor = -static_cast(1 << (b - 1)); + for (auto dtype : {float16, bfloat16, float32}) { + const int K = 512; + const int N = 72; + auto w = random::bits({N, K * b / 32}, 4, random::key(seed++), cpu); + auto sc = astype( + random::uniform( + 0.01f, 1.0f, {N, K / 64}, float32, random::key(seed++), cpu), + dtype, + cpu); + auto f = array(factor, dtype); + auto rebuilt = reference_bias(sc, f); + for (const auto& dev : devices) { + auto d_full = dequantize( + w, sc, rebuilt, 64, b, "affine", std::nullopt, std::nullopt, dev); + auto d_ib = dequantize( + w, sc, f, 64, b, "affine", std::nullopt, std::nullopt, dev); + CHECK(same(d_ib, d_full)); + for (int M : {1, 3, 8, 33}) { + auto x = astype( + random::normal({M, K}, float32, random::key(seed++), cpu), + dtype, + cpu); + auto o_full = + quantized_matmul(x, w, sc, rebuilt, true, 64, b, "affine", dev); + auto o_ib = + quantized_matmul(x, w, sc, f, true, 64, b, "affine", dev); + CHECK(same(o_ib, o_full)); + } + auto xn = astype( + random::normal({2, N}, float32, random::key(seed++), cpu), + dtype, + cpu); + CHECK(same( + quantized_matmul(xn, w, sc, f, false, 64, b, "affine", dev), + quantized_matmul( + xn, w, sc, rebuilt, false, 64, b, "affine", dev))); + } + } + } + } +} + TEST_CASE("test repeat") { auto data = array({13, 3, 16, 6, 14, 4, 15, 5, 11, 1, 12, 2}, {3, 2, 2}); auto repeat_axis_0 = repeat(data, 2, 0); From 939c71ef9aa47e7582a317e00a5fa125b5b6cec7 Mon Sep 17 00:00:00 2001 From: Jeethu Rao Date: Sun, 11 Oct 2026 12:55:50 +0100 Subject: [PATCH 05/34] quantized: qmv_fast for affine 4-bit K % 256 and 1-bit K % 512 on upstream main (gemma4-qat#192, #282) Port of numen-tech/mlx PR #8 (923aef64c, 7d565ec4a) onto upstream main. qmv_fast_k_alignment takes the mode and returns half a kernel block for affine / affine_sym 1-bit and 4-bit (1-bit K % 512, 4-bit K % 256), the whole block otherwise. nvfp4 / mxfp4 share bits == 4 but fp_qmv_fast_impl has no tail, so they keep the whole block; every other width keeps upstream's routing. qmv keeps upstream's gate shape from #4516 (fast = K-aligned, partial_rows = fast && N % 8 != 0); gather_qmv keeps N % 8 == 0. Tail predicate: kept upstream's lane predicate aligned_end + simd_lid * values_per_thread < K and the #4516 partial_rows row clamp, not the fork's aligned_end + (simd_lid + 1) * values_per_thread <= K. The two agree whenever K % values_per_thread == 0, which the host gate guarantees (block / 2 = 16 * values_per_thread, and K % group_size already implies it). Evidence on M3 Max (applegpu_g15s): - GREEN: test_qmv_fast_tail_non_aligned_k and test_qmv_fast_half_block pass; a kernel-name probe (removed) showed the 1-bit K 1536/2560 and 4-bit K 768/1280 M = 1 subtests on affine_qmv_fast (pr_1 at N = 67). - RED: with the predicate loosened to `< K + values_per_thread` (one extra lane reads past K) both tests fail: 20/20 tail-reachable subtests of test_qmv_fast_tail_non_aligned_k and every tail subtest of test_qmv_fast_half_block (165 + 3 gather), errors up to ~78. No read past K shows with upstream's predicate, so the fork's stricter form is not needed; the OOB-reasoning comment stays. Tests: test_qmv_fast_half_block (fork's, with the generic baseline moved to K + gs zero-padded because N % 8 != 0 no longer leaves qmv_fast after #4516; the fast call's x is a view with non-zero memory past K so the test is a tail guard; fp32 allows 4 ulps for accumulation order, measured <= 1.1), K = 768 in test_implied_bias_matmul, the "affine quantized_matmul qmv_fast tail block" doctest (renamed from the fork's "test affine qmv_fast tail block" so the *quantiz* filter runs it), updated bias-free doctest case comments. Also: the stale path="qmv_generic" subTest label is now "n_not_mult_8", and affine_sym_qmv[_fast] note their unused has_global_scale / results_per_simdgroup slots. --- mlx/backend/metal/kernels/quantized.h | 21 ++--- mlx/backend/metal/quantized.cpp | 20 +++-- python/tests/test_quantized.py | 110 ++++++++++++++++++++++---- tests/ops_tests.cpp | 98 +++++++++++++++++++++-- 4 files changed, 212 insertions(+), 37 deletions(-) diff --git a/mlx/backend/metal/kernels/quantized.h b/mlx/backend/metal/kernels/quantized.h index d8e8157c4b..82d191f763 100644 --- a/mlx/backend/metal/kernels/quantized.h +++ b/mlx/backend/metal/kernels/quantized.h @@ -981,15 +981,16 @@ METAL_FUNC void qmv_fast_impl( } // Partial last block (in_vec_size % block_size != 0), reached only when the - // host's K gate (qmv_fast_k_alignment) is finer than block_size. A lane - // whose first value lies past the end of the row reads nothing (x, weights, - // scales, biases) and skips the row loop: its result[row] stays 0 and it - // still joins the simd_sum below. Masking such lanes instead (zeroed - // x_thread, scale and bias) would still run qdot, whose weight reads land - // past the end of the row and, for the last output rows, past the end of - // the weight buffer. A lane that starts inside the row also ends inside it - // because the K gate keeps in_vec_size a multiple of values_per_thread; a - // gate that did not would need the stricter predicate + // host's K gate (qmv_fast_k_alignment) is finer than block_size: affine + // 1/4-bit take K % (block_size / 2) == 0. A lane whose first value lies + // past the end of the row reads nothing (x, weights, scales, biases) and + // skips the row loop: its result[row] stays 0 and it still joins the + // simd_sum below. Masking such lanes instead (zeroed x_thread, scale and + // bias) would still run qdot, whose weight reads land past the end of the + // row and, for the last output rows, past the end of the weight buffer. A + // lane that starts inside the row also ends inside it because the K gate + // keeps in_vec_size a multiple of values_per_thread; a gate that did not + // would need the stricter predicate // aligned_end + (simd_lid + 1) * values_per_thread <= in_vec_size. // (numen-tech/gemma4-qat#179, #282) if (aligned_end < in_vec_size) { @@ -1976,6 +1977,7 @@ template < int bits, bool batched, bool partial_rows = false, + // Unused (shared qmv instantiation macro and kname scheme): bool has_global_scale = false, int results_per_simdgroup = 4> [[kernel]] void affine_sym_qmv_fast( @@ -2034,6 +2036,7 @@ template < int bits, bool batched, bool partial_rows = false, + // Unused (shared qmv instantiation macro and kname scheme): bool has_global_scale = false, int results_per_simdgroup = 4> [[kernel]] void affine_sym_qmv( diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index ba441ddb7e..ef36fb5e9f 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -156,11 +156,17 @@ inline int get_qmv_batch_limit(int D, int O, metal::Device& d) { } } -// Must match the K step in qmv_fast_impl (kernels/quantized.h): -// pack_factor * packs_per_thread * SIMD_SIZE, with one pack per lane for -// bits <= 2. -inline int qmv_fast_k_alignment(int bits) { - return get_pack_factor(bits, 32) * (bits <= 2 ? 1 : 2) * 32; +// Block of qmv_fast_impl (kernels/quantized.h): pack_factor * +// packs_per_thread * SIMD_SIZE, with one pack per lane for bits <= 2 (must +// match the kernel). The affine kernel's partial last block only needs +// K % values_per_thread == 0, which K % group_size already implies; block / 2 +// is the cutoff measured for 1/4-bit, and the other widths stay whole-block +// until benchmarked. fp_qmv_fast_impl has no tail, so the fp modes (which +// share bits == 4) keep the whole block. +inline int qmv_fast_k_alignment(int bits, const std::string& mode) { + int block = get_pack_factor(bits, 32) * (bits <= 2 ? 1 : 2) * 32; + bool affine = mode == "affine" || mode == "affine_sym"; + return (affine && (bits == 1 || bits == 4)) ? block / 2 : block; } inline int add_strides_and_shapes( @@ -537,7 +543,7 @@ void qmv( std::string kname; kname.reserve(64); std::string type_string = get_type_string(x.dtype()); - bool fast = (K % qmv_fast_k_alignment(bits) == 0); + bool fast = (K % qmv_fast_k_alignment(bits, mode) == 0); bool partial_rows = fast && (N % bn != 0); // A narrower output tile reduces register pressure for large // floating-point quantized matrix-vector products on M5 Max GPUs. @@ -1628,7 +1634,7 @@ void gather_qmv( std::string kname; kname.reserve(64); std::string type_string = get_type_string(x.dtype()); - bool fast = N % bn == 0 && K % qmv_fast_k_alignment(bits) == 0; + bool fast = N % bn == 0 && K % qmv_fast_k_alignment(bits, mode) == 0; concatenate( kname, mode + (fast ? "_gather_qmv_fast_" : "_gather_qmv_"), diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index e08eb2e40d..df7bf90a75 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -869,15 +869,13 @@ def test_qmv_fast_tail_non_aligned_k(self): # GPU vs CPU for K that is a multiple of the group size but not of # qmv_fast's block, with N = 67 not a multiple of 8. # - # On the current host K gate (2048 for 1-bit, 512 for 4-bit) no shape - # here reaches the qmv_fast partial tail or the partial_rows clamp: - # they all route to plain qmv (or qmv_wide for M = 3 on gen 15+), so - # this covers plain qmv at non-aligned K. Once the routing patch - # widens the gate (1-bit K % 512, affine 4-bit K % 256; - # numen-tech/gemma4-qat#192), the 1-bit K = 1536/2560 and 4-bit - # K = 768/1280 shapes run the partial tail and, at N = 67, the - # last-row clamp, and this becomes the tail guard: the GPU cannot - # match the CPU if a tail lane reads past K or past the last row. + # The gs * 17 / gs * 33 shapes are on no qmv_fast gate and cover plain + # qmv at non-aligned K. The host gate takes 1-bit K % 512 and affine + # 4-bit K % 256 (half a kernel block; numen-tech/gemma4-qat#192), so + # the 1-bit K = 1536/2560 and 4-bit K = 768/1280 shapes run the + # qmv_fast partial tail and, at N = 67, the last-row clamp (M = 3 only + # where it stays on qmv, i.e. before gen 15): the GPU cannot match the + # CPU if a tail lane reads past K or past the last row. cases = [ (bits, gs, K, N, 1) for bits, gs in ((1, 32), (1, 64), (4, 32), (4, 64)) @@ -1279,6 +1277,86 @@ def f_biased(x, scales): _, jvp_biased = mx.jvp(lambda x: f_biased(x, scales), [x], [x]) self.assertLess((jvp_sym[0] - jvp_biased[0]).abs().max().item(), 1e-3) + @unittest.skipUnless(mx.metal.is_available(), "qmv_fast is a Metal kernel") + def test_qmv_fast_half_block(self): + # Affine 4-bit K % 256 == 0 and 1-bit K % 512 == 0 take qmv_fast with a + # partial last block. The generic qmv baseline runs the same product + # with K padded by one group that x zeroes out: K + gs is on no + # qmv_fast gate. (An extra output row no longer leaves qmv_fast: since + # ml-explore/mlx#4516 N % 8 != 0 takes its partial_rows variant.) The + # fast call's x is a view whose memory past K is that group drawn + # non-zero, so a tail lane that read past K would show. + key = mx.random.key(5) + k1, k2 = mx.random.split(key) + f32 = mx.float32 + cases = [(4, 256), (4, 768), (4, 3840), (1, 512), (1, 1536), (1, 2560)] + dtypes = [f32, mx.float16, mx.bfloat16] + for (bits, K), gs, dtype in product(cases, [32, 64, 128], dtypes): + for sym, B in product([False, True] if bits == 1 else [False], [0, 2]): + with self.subTest(bits=bits, K=K, gs=gs, dtype=dtype, sym=sym, B=B): + N = 64 + Kp = K + gs + lead = (B,) if B else () + w = mx.random.normal(lead + (N, Kp), key=k2).astype(dtype) + w_q, scales, biases = mx.quantize(w, gs, bits) + ref_b = -0.5 * scales if sym else biases + if sym: + biases = None + x = mx.random.normal(lead + (1, Kp), key=k1).astype(dtype) + pad = mx.zeros(lead + (1, gs), dtype) + x0 = mx.concatenate([x[..., :K], pad], axis=-1) + w_hat = mx.dequantize( + w_q, scales.astype(f32), ref_b.astype(f32), gs, bits + ) + ref = x0.astype(f32) @ mx.swapaxes(w_hat, -1, -2) + + def qmm(x, k): + b = None if biases is None else biases[..., : k // gs] + y = mx.quantized_matmul( + x[..., :k], + w_q[..., : k * bits // 32], + scales[..., : k // gs], + b, + True, + gs, + bits, + ) + return y.astype(f32) + + fast, generic = qmm(x, K), qmm(x0, Kp) + err_fast = (fast - ref).abs().max().item() + err_generic = (generic - ref).abs().max().item() + # One ulp of the output dtype covers the rounding of the + # different accumulation order in fp16 / bf16; in fp32 the + # order itself shows (up to ~1.1 ulp here), so allow 4. + n_ulp = 4 if dtype == f32 else 1 + ulp = n_ulp * mx.finfo(dtype).eps * ref.abs().max().item() + self.assertLessEqual(err_fast, err_generic + ulp) + + # gather_qmv shares the kernel. + for dtype in dtypes: + with self.subTest(gather=True, dtype=dtype): + w = mx.random.normal((3, 64, 3840), key=k2).astype(dtype) + w_q, scales, biases = mx.quantize(w, 32, 4) + x = mx.random.normal((2, 1, 3840), key=k1).astype(dtype) + idx = mx.array([2, 0]) + y = mx.gather_qmm( + x, + w_q, + scales, + biases, + rhs_indices=idx, + transpose=True, + group_size=32, + bits=4, + ) + w_hat = mx.dequantize( + w_q, scales.astype(f32), biases.astype(f32), 32, 4 + ) + ref = x.astype(f32) @ mx.swapaxes(w_hat[idx], -1, -2) + tol = 2 * mx.finfo(dtype).eps * ref.abs().max().item() + 1e-4 + self.assertLessEqual((y.astype(f32) - ref).abs().max().item(), tol) + # Implied bias: a 0-d `biases` array is a factor f and the kernels form # bias = f * scale per group. The factors of the shipped checkpoints are # powers of two, so f * scale is exact in every dtype and the implied path @@ -1329,9 +1407,9 @@ def test_implied_bias_matmul(self): self.IMPLIED_BIAS_DTYPES, self.IMPLIED_BIAS_CASES ): # K = 64 / 128 route to qmv_quad; 512 / 1536 to qmv_fast (and - # qmv_wide for 2 <= M <= 8 at bits > 1); N = 72 is not a - # multiple of the 32-wide qmm tile. - for K, N in product([64, 128, 512, 1536], [256, 72]): + # qmv_wide for 2 <= M <= 8 at bits > 1); K = 768 runs the 4-bit + # qmv_fast tail; N = 72 is not a multiple of the 32-wide qmm tile. + for K, N in product([64, 128, 512, 768, 1536], [256, 72]): if K < gs: continue w = mx.random.normal(shape=(N, K), key=k2).astype(dtype) @@ -1429,15 +1507,15 @@ def check(x, w_q, scales, full, transpose=True, factor=f, **sub): x = mx.random.normal(shape=(2, M, K), key=k1).astype(dtype) check(x, wb_q, scales_b, full_b, K=K, M=M, path="batched_w") - # N % 8 != 0: qmv_fast with partial rows where K is fast-aligned - # (2/4/8-bit here), else the generic qmv (1-bit); qmv_wide for - # M = 4 where the device routes affine there. + # N % 8 != 0: qmv_fast with partial rows (K = 512 is fast-aligned + # for every width here); qmv_wide for M = 4 where the device routes + # affine there. w = mx.random.normal(shape=(67, 512), key=k2).astype(dtype) w_q, scales, _ = mx.quantize(w, gs, bits) full = self.reference_bias(scales, factor) for M in [1, 4]: x = mx.random.normal(shape=(M, 512), key=k1).astype(dtype) - check(x, w_q, scales, full, M=M, path="qmv_generic") + check(x, w_q, scales, full, M=M, path="n_not_mult_8") # A float32 factor with half-precision scales adopts the scales' # dtype (no promotion of the whole matmul to float32). diff --git a/tests/ops_tests.cpp b/tests/ops_tests.cpp index 06792c1d26..ef05c1c1c5 100644 --- a/tests/ops_tests.cpp +++ b/tests/ops_tests.cpp @@ -3590,6 +3590,94 @@ TEST_CASE("test quantize dequantize") { } } +// Affine 4-bit K % 256 == 0 and 1-bit K % 512 == 0 take qmv_fast with a +// partial last block. N = 8 puts the last rows at the end of the weight buffer. +TEST_CASE("test affine quantized_matmul qmv_fast tail block") { + if (!metal::is_available()) { + return; + } + const auto gpu = Device::gpu; + const auto cpu = Device::cpu; + + struct Case { + int bits; + int K; + int N; + int group_size; + }; + const Case cases[] = { + {4, 256, 8, 32}, // tail block only + {4, 768, 8, 64}, // one 512 block + tail + {4, 3840, 8, 32}, // Gemma 4 12B hidden size + {4, 3840, 64, 128}, // several row groups + {4, 11008, 8, 128}, + {1, 1536, 8, 64}, // 1-bit, one 1024 block + tail + {4, 4096, 8, 32}, // no tail + }; + + int seed = 0; + for (auto dtype : {float16, bfloat16, float32}) { + for (const auto& c : cases) { + auto w = + random::bits({c.N, c.K * c.bits / 32}, 4, random::key(seed++), cpu); + auto scales = astype( + random::uniform( + 0.01f, + 0.1f, + {c.N, c.K / c.group_size}, + float32, + random::key(seed++), + cpu), + dtype, + cpu); + auto biases = astype( + random::uniform( + -0.5f, + 0.5f, + {c.N, c.K / c.group_size}, + float32, + random::key(seed++), + cpu), + dtype, + cpu); + auto x = astype( + random::normal({1, c.K}, float32, random::key(seed++), cpu), + dtype, + cpu); + + auto w_hat = dequantize( + w, + astype(scales, float32, cpu), + astype(biases, float32, cpu), + c.group_size, + c.bits, + "affine", + std::nullopt, + std::nullopt, + cpu); + auto expected = + matmul(astype(x, float32, cpu), transpose(w_hat, cpu), cpu); + + auto out = quantized_matmul( + x, w, scales, biases, true, c.group_size, c.bits, "affine", gpu); + CHECK_EQ(out.dtype(), dtype); + CHECK_EQ(out.shape(), Shape{1, c.N}); + + float rtol = + dtype == float32 ? 1e-4f : (dtype == float16 ? 4e-3f : 2e-2f); + float scale = max(abs(expected, cpu), cpu).item(); + float max_diff = + max(abs(subtract(astype(out, float32, cpu), expected, cpu), cpu), cpu) + .item(); + INFO( + "bits=" << c.bits << " K=" << c.K << " N=" << c.N + << " group_size=" << c.group_size << " dtype=" << dtype + << " max_diff=" << max_diff << " max|expected|=" << scale); + CHECK(max_diff <= 1e-3f + rtol * scale); + } + } +} + // Bias-free affine decode on Metal: the affine_sym_qmv[_fast] kernels load in // both the metallib and JIT builds and match the derived-bias reference. TEST_CASE("test bias-free affine quantized_matmul decode") { @@ -3606,14 +3694,14 @@ TEST_CASE("test bias-free affine quantized_matmul decode") { int group_size; }; const Case cases[] = { - {1, 512, 8, 128}, // 1-bit, K % 1024 != 0: generic affine_sym_qmv - {1, 1536, 8, 128}, // ditto (one full 1024 block + 512) + {1, 512, 8, 128}, // 1-bit fast path, tail block only + {1, 1536, 8, 128}, // 1-bit fast path, one 1024 block + tail {1, 1024, 8, 128}, // 1-bit fast path (K % 1024 == 0), no tail {2, 512, 8, 128}, // 2-bit fast path (K % 512 == 0) - {1, 1536, 64, 128}, // generic, several row groups - {1, 512, 8, 64}, // generic, group 64 + {1, 1536, 64, 128}, // fast path + tail, several row groups + {1, 512, 8, 64}, // fast path + tail, group 64 {2, 1536, 64, 64}, // 2-bit fast path, group 64, several row groups - {1, 544, 8, 32}, // K % 512 != 0 on either line: generic affine_sym_qmv + {1, 544, 8, 32}, // K % 512 != 0: generic affine_sym_qmv {2, 544, 8, 32}, // ditto, 2-bit }; From 30e08d8f9ba345329b664eeee3a524c0bb59d828 Mon Sep 17 00:00:00 2001 From: Jeethu Rao Date: Sun, 11 Oct 2026 13:02:57 +0100 Subject: [PATCH 06/34] quantized: affine qmv_wide by bit width and batch, single-tile 6-7 row verify (gemma4-qat#191, #282) Port of cddb7bed (route affine qmv_wide by bit width and batch) and numen-tech/mlx PR #9 (071a92a8..43e6613b: run an affine qmv_wide verify of 6 or 7 rows as one tile) onto upstream main. - use_qmv_wide takes bits and M: affine 1-bit, and affine 2-bit at M < 3, stay on qmv; other affine widths take qmv_wide on gen 15+ as before; fp modes unchanged; affine_sym still has no qmv_wide kernel. - qmv_wide runs affine M <= 7 as one tile (vecs_per_tg = M); M >= 8 keeps upstream's ceil(M / 5) tiling. quantized.metal instantiates the affine and implied-bias qmv_wide kernels for vecs_per_tg 6 and 7. No conflict with #4516, #4629 or #4641: none touches use_qmv_wide or the qmv_wide tiling, and #4629's M1 batch limit only applies below gen 15, where affine never takes qmv_wide. Tests: test_qmv_wide gains 1-bit affine (full sweep and tiny shapes); test_qmv_wide_tile_matches_split (M = 6/7/8/11/12 bit-identical to the split calls, verbatim from the fork); the "affine quantized_matmul qmv_wide single tile of 6 and 7 vectors" doctest (renamed from the fork's "test affine qmv_wide single tile of 6 and 7 vectors" so the *quantiz* filter runs it). Routing comments in the tests updated for the 1/2-bit gate. --- mlx/backend/metal/kernels/quantized.metal | 12 ++- mlx/backend/metal/quantized.cpp | 23 ++++- python/tests/test_quantized.py | 57 ++++++++++-- tests/ops_tests.cpp | 107 ++++++++++++++++++++++ 4 files changed, 183 insertions(+), 16 deletions(-) diff --git a/mlx/backend/metal/kernels/quantized.metal b/mlx/backend/metal/kernels/quantized.metal index fc017e9872..682a95744c 100644 --- a/mlx/backend/metal/kernels/quantized.metal +++ b/mlx/backend/metal/kernels/quantized.metal @@ -133,8 +133,8 @@ instantiate_quantized_quad(affine_qmv_quad, type, group_size, bits, 128, 1) \ instantiate_quantized_quad(affine_qmv_quad, type, group_size, bits, 128, 0) -// vecs_per_tg (input-vector tile) 2..5; affine uses k_lanes=8 (more rows per -// simdgroup) where the fp path uses 16. +// vecs_per_tg (input-vector tile) 2..5, affine 2..7; affine uses k_lanes=8 +// (more rows per simdgroup) where the fp path uses 16. #define instantiate_quantized_wide_wrap(name, type, group_size, bits, vecs_per_tg, k_lanes) \ instantiate_quantized_wide(name, type, group_size, bits, vecs_per_tg, k_lanes, 0) \ instantiate_quantized_wide(name, type, group_size, bits, vecs_per_tg, k_lanes, 1) @@ -143,7 +143,9 @@ instantiate_quantized_wide_wrap(affine_qmv_wide, type, group_size, bits, 2, 8) \ instantiate_quantized_wide_wrap(affine_qmv_wide, type, group_size, bits, 3, 8) \ instantiate_quantized_wide_wrap(affine_qmv_wide, type, group_size, bits, 4, 8) \ - instantiate_quantized_wide_wrap(affine_qmv_wide, type, group_size, bits, 5, 8) + instantiate_quantized_wide_wrap(affine_qmv_wide, type, group_size, bits, 5, 8) \ + instantiate_quantized_wide_wrap(affine_qmv_wide, type, group_size, bits, 6, 8) \ + instantiate_quantized_wide_wrap(affine_qmv_wide, type, group_size, bits, 7, 8) #define instantiate_quantized_all_splitk(type, group_size, bits) \ instantiate_quantized_split_k(affine_qvm_split_k, type, group_size, bits, 8) \ @@ -277,6 +279,10 @@ instantiate_quantized_ib_wide(type, group_size, bits, 4, 8, 1) \ instantiate_quantized_ib_wide(type, group_size, bits, 5, 8, 0) \ instantiate_quantized_ib_wide(type, group_size, bits, 5, 8, 1) \ + instantiate_quantized_ib_wide(type, group_size, bits, 6, 8, 0) \ + instantiate_quantized_ib_wide(type, group_size, bits, 6, 8, 1) \ + instantiate_quantized_ib_wide(type, group_size, bits, 7, 8, 0) \ + instantiate_quantized_ib_wide(type, group_size, bits, 7, 8, 1) \ instantiate_quantized_ib_split_k(type, group_size, bits, 8) \ instantiate_quantized_ib_split_k(type, group_size, bits, 32) \ instantiate_quantized_ib_splitk_qmm(type, group_size, bits, true) \ diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index ef36fb5e9f..1cfea56896 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -615,12 +615,22 @@ void qmv( } // affine qmv_wide only beats qmv on gen-15+; fp benefits on every gen. -// affine_sym has no qmv_wide kernel. -inline bool use_qmv_wide(const std::string& mode, metal::Device& d) { +// affine_sym has no qmv_wide kernel. The 1-bit path has so little weight +// traffic that unpacking into registers dominates even when several input +// rows reuse the result. The 2-bit path breaks even after two rows and wins +// once three or more rows share a block. +inline bool +use_qmv_wide(const std::string& mode, int bits, int M, metal::Device& d) { if (mode == "affine_sym") { return false; } - return mode != "affine" || d.get_architecture_gen() >= 15; + if (mode != "affine") { + return true; + } + if (bits == 1 || (bits == 2 && M < 3)) { + return false; + } + return d.get_architecture_gen() >= 15; } // Dispatches qmv_wide (fp modes -> fp_qmv_wide, affine -> affine_qmv_wide): @@ -644,6 +654,11 @@ void qmv_wide( // them. int n_tiles = (M + 4) / 5; // ceil(M / 5); tile size caps at 5 int vecs_per_tg = (M + n_tiles - 1) / n_tiles; + // affine runs M <= 7 as one tile. The result does not depend on the tile. + constexpr int max_affine_single_tile = 7; + if (mode == "affine" && M <= max_affine_single_tile) { + vecs_per_tg = M; + } // k_lanes: lanes reducing K per output row (32/k_lanes rows per simdgroup). // The affine subchunk decode has enough ALU per weight load to favor more @@ -2083,7 +2098,7 @@ void dispatch_qmv( // Small batch so route to qmv_wide, which reuses each weight group across the // M vectors. - if (M >= 2 && use_qmv_wide(mode, d) && !global_scale) { + if (M >= 2 && use_qmv_wide(mode, bits, M, d) && !global_scale) { qmv_wide(x, w, scales, biases, out, group_size, bits, M, N, K, d, s, mode); return; } diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index df7bf90a75..dffe8d8811 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -873,9 +873,9 @@ def test_qmv_fast_tail_non_aligned_k(self): # qmv at non-aligned K. The host gate takes 1-bit K % 512 and affine # 4-bit K % 256 (half a kernel block; numen-tech/gemma4-qat#192), so # the 1-bit K = 1536/2560 and 4-bit K = 768/1280 shapes run the - # qmv_fast partial tail and, at N = 67, the last-row clamp (M = 3 only - # where it stays on qmv, i.e. before gen 15): the GPU cannot match the - # CPU if a tail lane reads past K or past the last row. + # qmv_fast partial tail and, at N = 67, the last-row clamp (M = 3 where + # it stays on qmv: 1-bit, and 4-bit before gen 15): the GPU cannot + # match the CPU if a tail lane reads past K or past the last row. cases = [ (bits, gs, K, N, 1) for bits, gs in ((1, 32), (1, 64), (4, 32), (4, 64)) @@ -1075,8 +1075,9 @@ def test_fp_qmv_large_output(self): def test_qmv_wide(self): tol = 5e-2 if mx.cuda.is_available() else 1e-3 # M in [2, vector_limit) routes to qmv_wide -- except K in {64, 128} - # with power-of-2 bits, which stays on qmv_quad. Check both paths - # against a dequantize-then-matmul reference, with ragged M (token + # with power-of-2 bits, which stays on qmv_quad, and affine 1-bit (and + # 2-bit at M = 2) or affine before gen 15, which stay on qmv. Check all + # paths against a dequantize-then-matmul reference, with ragged M (token # tail) and ragged N (output-tile remainder). B > 1 stacks a distinct # weight matrix per slab and exercises the batched variant. key = mx.random.key(0) @@ -1089,7 +1090,7 @@ def test_qmv_wide(self): # Affine: every bit-width and group size. for group_size, bits, K in product( - [32, 64, 128], [2, 3, 4, 5, 6, 8], [128, 512] + [32, 64, 128], [1, 2, 3, 4, 5, 6, 8], [128, 512] ): for M, N, B in product(Ms, Ns, Bs): with self.subTest(M=M, N=N, K=K, B=B, group_size=group_size, bits=bits): @@ -1123,7 +1124,12 @@ def test_qmv_wide(self): # Tiny shapes (M, K, N): small K and non-multiple output rows. tiny = [(2, 32, 10), (4, 32, 7), (3, 64, 5), (5, 64, 3)] - settings = [(4, 32, "affine"), (6, 32, "affine"), (4, 16, "nvfp4")] + settings = [ + (1, 32, "affine"), + (4, 32, "affine"), + (6, 32, "affine"), + (4, 16, "nvfp4"), + ] for M, K, N in tiny: for bits, group_size, mode in settings: with self.subTest( @@ -1150,6 +1156,38 @@ def test_qmv_wide(self): self.assertEqual(y_q.shape, y_hat.shape) self.assertLess((y_q - y_hat).abs().max(), tol) + @unittest.skipUnless(mx.metal.is_available(), "qmv_wide is a Metal kernel") + def test_qmv_wide_tile_matches_split(self): + # One call over M rows must match the split calls bit for bit. + key = mx.random.key(6) + k1, k2 = mx.random.split(key) + cases = [(bits, gs, None) for bits in [2, 3, 4, 5, 6, 8] for gs in [32, 64]] + cases += [(b, gs, f) for b, gs, f in self.IMPLIED_BIAS_CASES if b > 1] + splits = {6: [3, 3], 7: [4, 3], 8: [4, 4], 11: [4, 4, 3], 12: [4, 4, 4]} + for (bits, gs, factor), dtype, K, B in product( + cases, self.IMPLIED_BIAS_DTYPES, [512, 1536], [1, 2] + ): + lead = (B,) if B > 1 else () + N = 72 + w = mx.random.normal(shape=lead + (N, K), key=k2).astype(dtype) + w_q, scales, biases = mx.quantize(w, gs, bits) + if factor is not None: + biases = mx.array(factor, dtype) + + def qmm(x): + return mx.quantized_matmul(x, w_q, scales, biases, True, gs, bits) + + for M, chunks in splits.items(): + with self.subTest( + bits=bits, gs=gs, factor=factor, dtype=dtype, K=K, B=B, M=M + ): + x = mx.random.normal(shape=lead + (M, K), key=k1).astype(dtype) + parts, r = [], 0 + for c in chunks: + parts.append(qmm(x[..., r : r + c, :])) + r += c + self.assert_bitwise_equal(qmm(x), mx.concatenate(parts, axis=-2)) + @unittest.skipUnless(mx.metal.is_available(), "Bias-free affine is Metal-only") def test_qmv_affine_sym(self): def assert_sym_matches_biased(y_sym, y_biased): @@ -1407,8 +1445,9 @@ def test_implied_bias_matmul(self): self.IMPLIED_BIAS_DTYPES, self.IMPLIED_BIAS_CASES ): # K = 64 / 128 route to qmv_quad; 512 / 1536 to qmv_fast (and - # qmv_wide for 2 <= M <= 8 at bits > 1); K = 768 runs the 4-bit - # qmv_fast tail; N = 72 is not a multiple of the 32-wide qmm tile. + # qmv_wide for 2 <= M <= 8 at bits > 1, from M = 3 at 2-bit); + # K = 768 runs the 4-bit qmv_fast tail; N = 72 is not a + # multiple of the 32-wide qmm tile. for K, N in product([64, 128, 512, 768, 1536], [256, 72]): if K < gs: continue diff --git a/tests/ops_tests.cpp b/tests/ops_tests.cpp index ef05c1c1c5..7cc693e166 100644 --- a/tests/ops_tests.cpp +++ b/tests/ops_tests.cpp @@ -3972,6 +3972,113 @@ TEST_CASE("test implied-bias affine quantized_matmul and dequantize") { } } +// One call over M rows must match the split calls bit for bit and stay close +// to an fp32 reference. +TEST_CASE( + "test affine quantized_matmul qmv_wide single tile of 6 and 7 vectors") { + if (!metal::is_available()) { + return; + } + const auto gpu = Device::gpu; + const auto cpu = Device::cpu; + + struct Case { + int bits; + int group_size; + std::optional factor; // implied bias when set + }; + const Case cases[] = { + {2, 64, std::nullopt}, + {3, 64, std::nullopt}, + {4, 64, std::nullopt}, + {6, 32, std::nullopt}, + {8, 64, std::nullopt}, + {2, 64, -2.0f}, + {4, 64, -8.0f}, + {8, 64, -128.0f}, + }; + struct Split { + int M; + std::vector chunks; + }; + const Split splits[] = {{6, {3, 3}}, {7, {4, 3}}, {8, {4, 4}}}; + + int seed = 0; + for (auto dtype : {float16, bfloat16, float32}) { + for (const auto& c : cases) { + for (int K : {512, 1536}) { + const int N = 72; + auto w = + random::bits({N, K * c.bits / 32}, 4, random::key(seed++), cpu); + auto per_group = [&](float lo, float hi) { + return astype( + random::uniform( + lo, + hi, + {N, K / c.group_size}, + float32, + random::key(seed++), + cpu), + dtype, + cpu); + }; + auto scales = per_group(0.01f, 0.1f); + auto biases = + c.factor ? array(*c.factor, dtype) : per_group(-0.5f, 0.5f); + auto ref_biases = + c.factor ? implied_bias_reference(scales, biases) : biases; + auto w_hat = dequantize( + w, + astype(scales, float32, cpu), + astype(ref_biases, float32, cpu), + c.group_size, + c.bits, + "affine", + std::nullopt, + std::nullopt, + cpu); + auto qmm = [&](const array& x) { + return quantized_matmul( + x, w, scales, biases, true, c.group_size, c.bits, "affine", gpu); + }; + + for (const auto& sp : splits) { + auto x = astype( + random::normal({sp.M, K}, float32, random::key(seed++), cpu), + dtype, + cpu); + std::vector parts; + int r = 0; + for (int n : sp.chunks) { + parts.push_back(qmm(slice(x, {r, 0}, {r + n, K}, cpu))); + r += n; + } + auto tiled = qmm(x); + auto split = concatenate(parts, 0, cpu); + INFO( + "bits=" << c.bits << " group_size=" << c.group_size + << " implied=" << c.factor.has_value() << " K=" << K + << " M=" << sp.M << " dtype=" << dtype); + CHECK_EQ(tiled.shape(), Shape{sp.M, N}); + CHECK(same_array(tiled, split)); + + auto expected = + matmul(astype(x, float32, cpu), transpose(w_hat, cpu), cpu); + float rtol = + dtype == float32 ? 1e-4f : (dtype == float16 ? 4e-3f : 2e-2f); + float scale = max(abs(expected, cpu), cpu).item(); + float max_diff = + max(abs(subtract(astype(tiled, float32, cpu), expected, cpu), + cpu), + cpu) + .item(); + CHECK(max_diff <= 1e-3f + rtol * scale); + } + } + } + } +} + TEST_CASE("test implied-bias factor rounding contract") { // Each block below is a case where the bias could be formed elsewhere -- in // the activations' promoted dtype, in the factor's wider dtype, or through From 8d396ce419e7009ef2bfad1375fc474986a3bf02 Mon Sep 17 00:00:00 2001 From: Jeethu Rao Date: Sun, 11 Oct 2026 13:19:18 +0100 Subject: [PATCH 07/34] metal: dispatch/commit/sync/wait counters (gemma4-qat#181, #282) Process-wide Metal work counters (metal::counters() / metal::reset()): dispatches, commits, explicit stream synchronizations, and host waits on GPU work. Source: numen-tech/mlx PR #3, `git diff 08121fd1 0f4257d0` (544767508, aaefa5005, 2e5d6ccdf, 1fee32b6e, 9cb5069a8). Conflicts with upstream main since v0.32.2: - fence.cpp: upstream dropped ~FenceImpl (fence is now an std::optional) and Fence::wait takes the target value as a parameter. The fork's `Device device` member and the host-wait count in the CPU spin-wait are kept, spinning on upstream's `value`. - tests/CMakeLists.txt: upstream renamed METAL_TEST_SOURCES to GPU_TEST_SOURCES; metal_counters_tests.cpp is appended to the new list. No new host-wait or dispatch site upstream: CommandEncoder:: dispatch_threadgroups/dispatch_threads and CommandEncoder::synchronize are still the only dispatch and waitUntilCompleted sites in mlx/backend/metal, and eval error recovery still flushes via gpu::synchronize(s, false). Tests: ./build/tests/tests -tc='*counters*' -> 4 cases, 26 assertions pass. test_quantized.py + test_fast.py: 102 passed, 2 skipped. --- mlx/backend/cuda/eval.cpp | 2 +- mlx/backend/gpu/eval.h | 3 +- mlx/backend/metal/device.cpp | 41 +++++++++++- mlx/backend/metal/device.h | 5 +- mlx/backend/metal/eval.cpp | 4 +- mlx/backend/metal/event.cpp | 15 ++++- mlx/backend/metal/fence.cpp | 7 +- mlx/backend/metal/metal.h | 12 ++++ mlx/backend/metal/no_metal.cpp | 5 ++ mlx/backend/no_gpu/eval.cpp | 2 +- mlx/scheduler.h | 22 +++++++ mlx/transforms.cpp | 6 +- tests/CMakeLists.txt | 2 +- tests/metal_counters_tests.cpp | 114 +++++++++++++++++++++++++++++++++ 14 files changed, 227 insertions(+), 13 deletions(-) create mode 100644 tests/metal_counters_tests.cpp diff --git a/mlx/backend/cuda/eval.cpp b/mlx/backend/cuda/eval.cpp index 2862f38223..334787dff7 100644 --- a/mlx/backend/cuda/eval.cpp +++ b/mlx/backend/cuda/eval.cpp @@ -83,7 +83,7 @@ void finalize(Stream s) { cu::get_command_encoder(s).commit(); } -void synchronize(Stream s) { +void synchronize(Stream s, bool) { nvtx3::scoped_range r("gpu::synchronize"); cu::get_command_encoder(s).synchronize(); } diff --git a/mlx/backend/gpu/eval.h b/mlx/backend/gpu/eval.h index d75b30abb5..7071231885 100644 --- a/mlx/backend/gpu/eval.h +++ b/mlx/backend/gpu/eval.h @@ -15,7 +15,8 @@ void new_stream(Stream s); void new_thread_unsafe_stream(Stream s); void eval(array& arr); void finalize(Stream s); -void synchronize(Stream s); +// explicit_sync is false for internal flushes (e.g. eval error recovery). +void synchronize(Stream s, bool explicit_sync = true); void clear_streams(); } // namespace mlx::core::gpu diff --git a/mlx/backend/metal/device.cpp b/mlx/backend/metal/device.cpp index 493e8e3153..ea96acea84 100644 --- a/mlx/backend/metal/device.cpp +++ b/mlx/backend/metal/device.cpp @@ -1,5 +1,6 @@ // Copyright © 2023-2024 Apple Inc. +#include #include #include @@ -14,6 +15,7 @@ #include "mlx/backend/metal/event.h" #include "mlx/backend/metal/metal.h" #include "mlx/backend/metal/utils.h" +#include "mlx/scheduler.h" #include "mlx/utils.h" namespace std { @@ -32,6 +34,12 @@ namespace mlx::core::metal { namespace { +// Statistics for counters() (metal.h), not used for synchronization. +std::atomic dispatch_count{0}; +std::atomic commit_count{0}; +std::atomic sync_count{0}; +std::atomic wait_count{0}; + constexpr const char* default_mtllib_path = METAL_PATH; void set_compile_options( @@ -426,6 +434,7 @@ void CommandEncoder::dispatch_threadgroups( MTL::Size group_dims) { maybeInsertBarrier(); buffer_ops_++; + dispatch_count.fetch_add(1, std::memory_order_relaxed); get_command_encoder()->dispatchThreadgroups(grid_dims, group_dims); } @@ -434,6 +443,7 @@ void CommandEncoder::dispatch_threads( MTL::Size group_dims) { maybeInsertBarrier(); buffer_ops_++; + dispatch_count.fetch_add(1, std::memory_order_relaxed); get_command_encoder()->dispatchThreads(grid_dims, group_dims); } @@ -582,16 +592,24 @@ void CommandEncoder::commit(std::function completion) { } }); buffer_->commit(); + commit_count.fetch_add(1, std::memory_order_relaxed); buffer_ = NS::RetainPtr(queue_->commandBufferWithUnretainedReferences()); buffer_ops_ = 0; buffer_sizes_ = 0; } -void CommandEncoder::synchronize() { +void CommandEncoder::synchronize(bool explicit_sync) { auto pool = new_scoped_memory_pool(); auto cbuf = buffer_; // retained end_encoding(); commit(); + if (explicit_sync) { + sync_count.fetch_add(1, std::memory_order_relaxed); + } + // Encoder teardown (clear_streams(), thread exit) is not a host wait. + if (!exiting_) { + count_host_wait(); + } cbuf->waitUntilCompleted(); if (!exiting_) { @@ -996,4 +1014,25 @@ bool is_nax_available() { #endif } +void count_host_wait() { + wait_count.fetch_add(1, std::memory_order_relaxed); +} + +Counters counters() { + return { + dispatch_count.load(std::memory_order_relaxed), + commit_count.load(std::memory_order_relaxed), + sync_count.load(std::memory_order_relaxed), + wait_count.load(std::memory_order_relaxed) + + scheduler::scheduler().gpu_waits()}; +} + +void reset() { + dispatch_count.store(0, std::memory_order_relaxed); + commit_count.store(0, std::memory_order_relaxed); + sync_count.store(0, std::memory_order_relaxed); + wait_count.store(0, std::memory_order_relaxed); + scheduler::scheduler().reset_gpu_waits(); +} + } // namespace mlx::core::metal diff --git a/mlx/backend/metal/device.h b/mlx/backend/metal/device.h index 305a35a772..5bb3359310 100644 --- a/mlx/backend/metal/device.h +++ b/mlx/backend/metal/device.h @@ -99,7 +99,7 @@ class MLX_API CommandEncoder { void signal_event(Event event, uint64_t value); bool needs_commit() const; void commit(std::function completion = nullptr); - void synchronize(); + void synchronize(bool explicit_sync = false); MTL::CommandBuffer* get_command_buffer() const { return buffer_.get(); @@ -262,4 +262,7 @@ NS::SharedPtr new_scoped_memory_pool(); bool is_nax_available(); +// Call before each host wait on GPU work. +void count_host_wait(); + } // namespace mlx::core::metal diff --git a/mlx/backend/metal/eval.cpp b/mlx/backend/metal/eval.cpp index 43a6fb9932..9b492300b6 100644 --- a/mlx/backend/metal/eval.cpp +++ b/mlx/backend/metal/eval.cpp @@ -76,8 +76,8 @@ void finalize(Stream s) { encoder.commit(); } -void synchronize(Stream s) { - metal::get_command_encoder(s).synchronize(); +void synchronize(Stream s, bool explicit_sync) { + metal::get_command_encoder(s).synchronize(explicit_sync); } void clear_streams() { diff --git a/mlx/backend/metal/event.cpp b/mlx/backend/metal/event.cpp index 38a387c9c0..11a6b2ea5f 100644 --- a/mlx/backend/metal/event.cpp +++ b/mlx/backend/metal/event.cpp @@ -45,15 +45,24 @@ Event::Event(Stream stream) : stream_(stream) { void Event::wait() { check_error(); + if (stream_.device == Device::gpu) { + metal::count_host_wait(); + } cast().wait(value()); check_error(); } void Event::wait(Stream stream) { if (stream.device == Device::cpu) { - scheduler::wait_event(stream, *this, [value = value()](Event& self) { - self.cast().wait(value); - }); + scheduler::wait_event( + stream, + *this, + [value = value(), gpu = stream_.device == Device::gpu](Event& self) { + if (gpu) { + metal::count_host_wait(); + } + self.cast().wait(value); + }); } else { auto& encoder = metal::get_command_encoder(stream); encoder.wait_event(*this, value()); diff --git a/mlx/backend/metal/fence.cpp b/mlx/backend/metal/fence.cpp index 819e2bb1b2..209b15a706 100644 --- a/mlx/backend/metal/fence.cpp +++ b/mlx/backend/metal/fence.cpp @@ -10,7 +10,7 @@ namespace mlx::core { struct FenceImpl { - FenceImpl(Stream stream) { + FenceImpl(Stream stream) : device(stream.device) { auto d = metal::device(stream.device).mtl_device(); if (!d->supportsFamily(MTL::GPUFamilyMetal3)) { use_fast = false; @@ -26,6 +26,8 @@ struct FenceImpl { } } + // Device of the stream that updates (signals) the fence. + Device device; bool use_fast{false}; uint32_t count{0}; std::optional fence; @@ -57,6 +59,9 @@ void Fence::wait(Stream stream, const array& x, uint32_t value) { if (stream.device == Device::cpu) { scheduler::enqueue(stream, [fence_ = fence_, value]() mutable { auto& f = *static_cast(fence_.get()); + if (f.device == Device::gpu) { + metal::count_host_wait(); + } while (f.cpu_value()[0] < value) { } }); diff --git a/mlx/backend/metal/metal.h b/mlx/backend/metal/metal.h index d686c2971a..45dc2e0b26 100644 --- a/mlx/backend/metal/metal.h +++ b/mlx/backend/metal/metal.h @@ -2,6 +2,7 @@ #pragma once +#include #include #include #include @@ -27,4 +28,15 @@ MLX_API const MLX_API void set_metallib_path(const std::string& path); MLX_API const std::string& get_metallib_path(); +// Process-wide Metal work counters. syncs counts mx::synchronize() calls, +// waits counts host waits on GPU work. Read or reset only while no eval runs. +struct Counters { + uint64_t dispatches; + uint64_t commits; + uint64_t syncs; + uint64_t waits; +}; +MLX_API Counters counters(); +MLX_API void reset(); + } // namespace mlx::core::metal diff --git a/mlx/backend/metal/no_metal.cpp b/mlx/backend/metal/no_metal.cpp index 3b72ad005a..1609ae93c7 100644 --- a/mlx/backend/metal/no_metal.cpp +++ b/mlx/backend/metal/no_metal.cpp @@ -28,6 +28,11 @@ const std::string& get_metallib_path() { "[metal::get_metallib_path] Cannot get metallib path without metal backend"); } +Counters counters() { + return {0, 0, 0, 0}; +} +void reset() {} + } // namespace metal } // namespace mlx::core diff --git a/mlx/backend/no_gpu/eval.cpp b/mlx/backend/no_gpu/eval.cpp index b08f6ec6fb..816f04b3ff 100644 --- a/mlx/backend/no_gpu/eval.cpp +++ b/mlx/backend/no_gpu/eval.cpp @@ -27,7 +27,7 @@ void finalize(Stream) { throw std::runtime_error("[gpu::finalize] GPU backend is not available"); } -void synchronize(Stream) { +void synchronize(Stream, bool) { throw std::runtime_error("[gpu::synchronize] GPU backend is not available"); } diff --git a/mlx/scheduler.h b/mlx/scheduler.h index 7bce5a9fe6..edea3ce3ad 100644 --- a/mlx/scheduler.h +++ b/mlx/scheduler.h @@ -37,6 +37,9 @@ class MLX_API Scheduler { { std::lock_guard lk(mtx); n_active_tasks_++; + if (stream.device == Device::gpu) { + n_active_gpu_tasks_++; + } } completion_cv.notify_all(); } @@ -45,6 +48,9 @@ class MLX_API Scheduler { { std::lock_guard lk(mtx); n_active_tasks_--; + if (stream.device == Device::gpu) { + n_active_gpu_tasks_--; + } } completion_cv.notify_all(); } @@ -57,18 +63,34 @@ class MLX_API Scheduler { std::unique_lock lk(mtx); int n_tasks_old = n_active_tasks(); if (n_tasks_old > 1) { + int n_gpu_tasks_old = n_active_gpu_tasks_; completion_cv.wait(lk, [this, n_tasks_old] { return this->n_active_tasks() < n_tasks_old; }); + // Count only when a GPU task completed during the wait. + if (n_active_gpu_tasks_ < n_gpu_tasks_old) { + gpu_waits_.fetch_add(1, std::memory_order_relaxed); + } } } + // Blocking wait_for_one() calls that ended on a GPU task completion. + uint64_t gpu_waits() const { + return gpu_waits_.load(std::memory_order_relaxed); + } + + void reset_gpu_waits() { + gpu_waits_.store(0, std::memory_order_relaxed); + } + private: friend Stream mlx::core::new_stream(Device d); StreamThread& get_thread(Stream s); int n_active_tasks_{0}; + int n_active_gpu_tasks_{0}; + std::atomic gpu_waits_{0}; std::unordered_map> threads_; std::shared_mutex threads_mtx_; std::condition_variable completion_cv; diff --git a/mlx/transforms.cpp b/mlx/transforms.cpp index cad704bd43..435678f81e 100644 --- a/mlx/transforms.cpp +++ b/mlx/transforms.cpp @@ -343,7 +343,11 @@ array eval_impl(std::vector outputs, bool async) { } for (auto& s : open_streams) { try { - synchronize(s); + if (s.device == Device::gpu) { + gpu::synchronize(s, /* explicit_sync = */ false); + } else { + synchronize(s); + } } catch (...) { // Preserve the original exception. } diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index aa0d90105a..f3aa6a5935 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -16,7 +16,7 @@ if(MLX_BUILD_CUDA) endif() if(MLX_BUILD_METAL) - list(APPEND GPU_TEST_SOURCES residency_tests.cpp) + list(APPEND GPU_TEST_SOURCES residency_tests.cpp metal_counters_tests.cpp) endif() include(${doctest_SOURCE_DIR}/scripts/cmake/doctest.cmake) diff --git a/tests/metal_counters_tests.cpp b/tests/metal_counters_tests.cpp new file mode 100644 index 0000000000..ea4b01584c --- /dev/null +++ b/tests/metal_counters_tests.cpp @@ -0,0 +1,114 @@ +// Copyright © 2026 Apple Inc. + +#include + +#include "doctest/doctest.h" + +#include "mlx/backend/metal/metal.h" +#include "mlx/mlx.h" + +using namespace mlx::core; + +TEST_CASE("test metal counters count dispatches, commits, syncs and waits") { + // Use the GPU explicitly: DEVICE=cpu makes the CPU the default device. + Device gpu = Device::gpu; + auto a = random::normal({64, 64}, float32, 0.0, 1.0, std::nullopt, gpu); + auto b = random::normal({64, 64}, float32, 0.0, 1.0, std::nullopt, gpu); + eval(a, b); + + metal::reset(); + auto z = metal::counters(); + CHECK_EQ(z.dispatches, 0); + CHECK_EQ(z.commits, 0); + CHECK_EQ(z.syncs, 0); + CHECK_EQ(z.waits, 0); + + auto c = matmul(a, b, gpu); + eval(c); + auto after_eval = metal::counters(); + CHECK_GE(after_eval.dispatches, 1); + CHECK_GE(after_eval.commits, 1); + CHECK_EQ(after_eval.syncs, 0); + // eval() blocks on c's completion event unless the GPU already signaled it + // before the host checked, so at most one wait. + CHECK_LE(after_eval.waits, 1); + + // Re-evaluating (or reading) an evaluated array does not wait again. + eval(c); + CHECK_EQ(metal::counters().waits, after_eval.waits); + + synchronize(default_stream(gpu)); + auto after_sync = metal::counters(); + CHECK_EQ(after_sync.syncs, 1); + CHECK_EQ(after_sync.waits, after_eval.waits + 1); + CHECK_GE(after_sync.commits, after_eval.commits + 1); + CHECK_EQ(after_sync.dispatches, after_eval.dispatches); + + metal::reset(); + auto r = metal::counters(); + CHECK_EQ(r.dispatches, 0); + CHECK_EQ(r.commits, 0); + CHECK_EQ(r.syncs, 0); + CHECK_EQ(r.waits, 0); +} + +TEST_CASE("test metal counters skip CPU-only work") { + Device cpu = Device::cpu; + auto a = random::normal({64, 64}, float32, 0.0, 1.0, std::nullopt, cpu); + auto b = random::normal({64, 64}, float32, 0.0, 1.0, std::nullopt, cpu); + eval(a, b); + + metal::reset(); + auto c = matmul(a, b, cpu); + eval(c); + synchronize(default_stream(cpu)); + auto r = metal::counters(); + CHECK_EQ(r.dispatches, 0); + CHECK_EQ(r.commits, 0); + CHECK_EQ(r.syncs, 0); + CHECK_EQ(r.waits, 0); +} + +TEST_CASE("test metal counters skip encoder teardown") { + // A new thread gets its own GPU stream, so clear_streams() there does not + // touch the streams of the main thread. + std::thread t([]() { + Device gpu = Device::gpu; + auto a = random::normal({64, 64}, float32, 0.0, 1.0, std::nullopt, gpu); + eval(a); + metal::reset(); + clear_streams(); + }); + t.join(); + + auto r = metal::counters(); + CHECK_EQ(r.syncs, 0); + CHECK_EQ(r.waits, 0); +} + +TEST_CASE("test metal counters skip eval error recovery") { + Device gpu = Device::gpu; + auto a = full({1024}, 3.0f, gpu); + auto b = multiply(a, array(2.0f), gpu); + // A Metal compile error makes the kernel throw at eval time. + auto kernel = fast::metal_kernel( + "test_metal_counters_bad_kernel", + {"inp"}, + {"out"}, + "this is not metal code {"); + auto outs = kernel( + {b}, + {b.shape()}, + {b.dtype()}, + {1, 1, 1}, + {1, 1, 1}, + {}, + std::nullopt, + false, + gpu); + + metal::reset(); + CHECK_THROWS(eval(outs[0])); + CHECK_EQ(metal::counters().syncs, 0); + CHECK(all(equal(b, array(6.0f), gpu), gpu).item()); +} From 158d0ce088865e2503fcfabf968bc77c04703e2d Mon Sep 17 00:00:00 2001 From: Jeethu Rao Date: Sun, 11 Oct 2026 13:24:10 +0100 Subject: [PATCH 08/34] metal_kernel: build and hash each kernel variant once, bounded memo (gemma4-qat#189, #282) fast::metal_kernel's closure keeps a per-kernel memo (shared by all copies of the closure, mutex-guarded) from each call variant (input dtypes and pass mode, output dtypes, template argument kind + name + value) to its kernel name, source and source hash. CustomKernel holds the source in a shared_ptr with a precomputed hash. The memo holds at most metal_kernel_max_cached_variants (64) entries and is cleared when full; entries are shared_ptrs, so a pending array whose entry was evicted still evaluates. Source: numen-tech/mlx PR #5, f5428f85c + b2c93bd1d. Merged with upstream #4584 (908148ee9, `-` sanitized in make_template_hash): the memo calls make_template_hash unchanged when it builds a variant, so negative template values still give a valid kernel name, and the memo key encodes the value's bytes, so -1 and 1 stay separate variants. Both behaviors are kept; the only textual conflict was tests/CMakeLists.txt (upstream's METAL_TEST_SOURCES -> GPU_TEST_SOURCES rename). Tests: test_fast.py -k "custom_kernel or metal_kernel" -> 11 passed (incl. #4584's test_custom_metal_kernel_negative_template_int and the fork's test_custom_kernel_many_template_values). ./build/tests/tests -tc='*metal kernel*,*counters*' -> 8 cases, 778 assertions pass. test_quantized.py + test_fast.py: 103 passed, 2 skipped. --- mlx/backend/common/metal_kernel.cpp | 179 ++++++++++++++++++++-------- mlx/backend/cuda/custom_kernel.cpp | 5 +- mlx/backend/metal/custom_kernel.cpp | 6 +- mlx/fast_primitives.h | 43 ++++++- python/tests/test_fast.py | 53 ++++++++ tests/CMakeLists.txt | 3 +- tests/metal_kernel_tests.cpp | 163 +++++++++++++++++++++++++ 7 files changed, 396 insertions(+), 56 deletions(-) create mode 100644 tests/metal_kernel_tests.cpp diff --git a/mlx/backend/common/metal_kernel.cpp b/mlx/backend/common/metal_kernel.cpp index d38a43395f..7e0a28a7ff 100644 --- a/mlx/backend/common/metal_kernel.cpp +++ b/mlx/backend/common/metal_kernel.cpp @@ -1,7 +1,9 @@ // Copyright © 2024 Apple Inc. #include +#include #include +#include #include "mlx/backend/common/compiled.h" #include "mlx/backend/metal/metal.h" @@ -217,6 +219,59 @@ std::string make_template_hash(const std::string& template_def) { return template_hash; } +// The name, source and source hash of one kernel variant. +struct KernelVariant { + std::string name; + std::shared_ptr source; + size_t source_hash; +}; + +// Shared by all copies of one kernel function. Holds at most +// `metal_kernel_max_cached_variants` entries and is cleared when full. Callers +// keep an owning pointer to their entry, so clearing is safe at any time. +struct VariantCache { + std::mutex mutex; + std::unordered_map> + variants; +}; + +// Encodes each call argument that the name and source depend on. Template +// arguments include their kind and name: `int` 1 and `bool` true give the same +// name but a different signature. +std::string variant_key( + const std::vector& inputs, + const std::vector& output_dtypes, + const std::vector>& template_args) { + std::string key; + auto append_bytes = [&key](auto v) { + key.append(reinterpret_cast(&v), sizeof(v)); + }; + for (const auto& arr : inputs) { + key += static_cast(arr.dtype().val()); + key += arr.ndim() == 0 ? 's' + : arr.size() < max_constant_array_size ? 'c' + : 'd'; + } + for (const auto& dtype : output_dtypes) { + key += static_cast(dtype.val()); + } + for (const auto& [name, arg] : template_args) { + key += static_cast(arg.index()); + append_bytes(name.size()); + key += name; + std::visit( + [&](auto v) { + if constexpr (std::is_same_v) { + key += static_cast(v.val()); + } else { + append_bytes(v); + } + }, + arg); + } + return key; +} + } // namespace CustomKernelFunction metal_kernel( @@ -270,7 +325,10 @@ CustomKernelFunction metal_kernel( } } + auto cache = std::make_shared(); + return [=, + cache = std::move(cache), shape_infos = std::move(shape_infos), attributes = std::move(attributes)]( const std::vector& inputs, @@ -298,60 +356,86 @@ CustomKernelFunction metal_kernel( auto s = resolve_metal_kernel_stream(s_); - std::string kernel_name = "custom_kernel_" + name; - std::string template_def = ""; - if (!template_args.empty()) { - template_def = write_template(template_args); - auto template_hash = make_template_hash(template_def); - kernel_name += "_"; - kernel_name += template_hash; + auto key = variant_key(inputs, output_dtypes, template_args); + std::shared_ptr variant; + { + std::lock_guard lock(cache->mutex); + if (auto it = cache->variants.find(key); it != cache->variants.end()) { + variant = it->second; + } } + if (!variant) { + std::string kernel_name = "custom_kernel_" + name; + std::string template_def = ""; + if (!template_args.empty()) { + template_def = write_template(template_args); + auto template_hash = make_template_hash(template_def); + kernel_name += "_"; + kernel_name += template_hash; + } - // The generated source depends on the dtypes of the inputs and outputs - // and on how each input is passed (see `write_signature`). Include them - // in the kernel name so that a given name always maps to the same source. - for (const auto& arr : inputs) { - kernel_name += "_"; - kernel_name += get_type_string(arr.dtype()); - if (arr.ndim() == 0) { - kernel_name += "s"; - } else if (arr.size() < max_constant_array_size) { - kernel_name += "c"; + // The generated source depends on the dtypes of the inputs and outputs + // and on how each input is passed (see `write_signature`). Include them + // in the kernel name so that a given name always maps to the same + // source. + for (const auto& arr : inputs) { + kernel_name += "_"; + kernel_name += get_type_string(arr.dtype()); + if (arr.ndim() == 0) { + kernel_name += "s"; + } else if (arr.size() < max_constant_array_size) { + kernel_name += "c"; + } + } + for (const auto& dtype : output_dtypes) { + kernel_name += "_"; + kernel_name += get_type_string(dtype); } - } - for (const auto& dtype : output_dtypes) { - kernel_name += "_"; - kernel_name += get_type_string(dtype); - } - std::string kernel_source = write_signature( - kernel_name, - header, - source, - input_names, - inputs, - output_names, - output_dtypes, - template_args, - attributes, - shape_infos, - atomic_outputs); + std::string kernel_source = write_signature( + kernel_name, + header, + source, + input_names, + inputs, + output_names, + output_dtypes, + template_args, + attributes, + shape_infos, + atomic_outputs); - if (!template_args.empty()) { - template_def = kernel_name + template_def; - kernel_source += "\ntemplate [[host_name(\""; - kernel_source += kernel_name; - kernel_source += "\")]] [[kernel]] decltype("; - kernel_source += template_def; - kernel_source += ") "; - kernel_source += template_def; - kernel_source += ";\n"; + if (!template_args.empty()) { + template_def = kernel_name + template_def; + kernel_source += "\ntemplate [[host_name(\""; + kernel_source += kernel_name; + kernel_source += "\")]] [[kernel]] decltype("; + kernel_source += template_def; + kernel_source += ") "; + kernel_source += template_def; + kernel_source += ";\n"; + } + + auto source_hash = std::hash{}(kernel_source); + variant = std::make_shared(KernelVariant{ + std::move(kernel_name), + std::make_shared(std::move(kernel_source)), + source_hash}); + std::lock_guard lock(cache->mutex); + if (auto it = cache->variants.find(key); it != cache->variants.end()) { + variant = it->second; + } else { + if (cache->variants.size() >= metal_kernel_max_cached_variants) { + cache->variants.clear(); + } + cache->variants.emplace(std::move(key), variant); + } } if (verbose) { std::cout << "Generated source code for `" << name << "`:" << std::endl << "```" << std::endl - << kernel_source << std::endl + << *variant->source << std::endl << "```" << std::endl; } @@ -360,8 +444,8 @@ CustomKernelFunction metal_kernel( std::move(output_dtypes), std::make_shared( s, - std::move(kernel_name), - std::move(kernel_source), + variant->name, + variant->source, grid, threadgroup, shape_infos, @@ -370,7 +454,8 @@ CustomKernelFunction metal_kernel( std::vector{}, false, 0, - compile_options.serialize()), + compile_options.serialize(), + variant->source_hash), std::move(inputs)); }; } diff --git a/mlx/backend/cuda/custom_kernel.cpp b/mlx/backend/cuda/custom_kernel.cpp index c230656d80..7b33e90de2 100644 --- a/mlx/backend/cuda/custom_kernel.cpp +++ b/mlx/backend/cuda/custom_kernel.cpp @@ -310,14 +310,13 @@ void CustomKernel::eval_gpu( // Compile the custom kernel std::string kernel_name = (is_precompiled_) ? name_ : "mlx::core::cu::" + name_; - std::string module_name = - fmt::format("{}_{:x}", name_, std::hash{}(source_)); + std::string module_name = fmt::format("{}_{:x}", name_, source_hash_); cu::JitModule& mod = cu::get_jit_module( encoder.device(), module_name, [&]() { return std::make_tuple( - is_precompiled_, source_, std::vector{kernel_name}); + is_precompiled_, *source_, std::vector{kernel_name}); }, false); diff --git a/mlx/backend/metal/custom_kernel.cpp b/mlx/backend/metal/custom_kernel.cpp index b73dd72def..2856d30919 100644 --- a/mlx/backend/metal/custom_kernel.cpp +++ b/mlx/backend/metal/custom_kernel.cpp @@ -47,10 +47,10 @@ void CustomKernel::eval_gpu( auto& d = metal::device(s.device); - std::string lib_name = fmt::format( - "{}_{:x}_{}", name_, std::hash{}(source_), compile_options_); + std::string lib_name = + fmt::format("{}_{:x}_{}", name_, source_hash_, compile_options_); auto lib = d.get_library( - lib_name, compile_options_, [this] { return metal::utils() + source_; }); + lib_name, compile_options_, [this] { return metal::utils() + *source_; }); auto kernel = d.get_kernel(name_, lib); auto& compute_encoder = metal::get_command_encoder(s); compute_encoder.set_compute_pipeline_state(kernel); diff --git a/mlx/fast_primitives.h b/mlx/fast_primitives.h index 8117a7a7b6..3e3240970c 100644 --- a/mlx/fast_primitives.h +++ b/mlx/fast_primitives.h @@ -505,6 +505,9 @@ class Quantize : public Custom { using ScalarArg = std::variant; +// The most variants that one metal_kernel function keeps built. +inline constexpr size_t metal_kernel_max_cached_variants = 64; + class CustomKernel : public Primitive { public: CustomKernel( @@ -520,9 +523,40 @@ class CustomKernel : public Primitive { bool is_precompiled, int shared_memory, CompileOptions::Data compile_options = {}) + : CustomKernel( + stream, + std::move(name), + std::make_shared(std::move(source)), + grid, + threadgroup, + std::move(shape_infos), + ensure_row_contiguous, + init_value, + std::move(scalar_arguments), + is_precompiled, + shared_memory, + compile_options) {} + + // Kernels built from one source can share it and its hash. + CustomKernel( + Stream stream, + std::string name, + std::shared_ptr source, + std::tuple grid, + std::tuple threadgroup, + std::vector> shape_infos, + bool ensure_row_contiguous, + std::optional init_value, + std::vector scalar_arguments, + bool is_precompiled, + int shared_memory, + CompileOptions::Data compile_options, + std::optional source_hash = std::nullopt) : Primitive(stream), name_(std::move(name)), source_(std::move(source)), + source_hash_( + source_hash ? *source_hash : std::hash{}(*source_)), grid_(grid), threadgroup_(threadgroup), shape_infos_(std::move(shape_infos)), @@ -545,7 +579,7 @@ class CustomKernel : public Primitive { auto state() const { return std::make_tuple( name_, - source_, + *source_, grid_, threadgroup_, shape_infos_, @@ -557,9 +591,14 @@ class CustomKernel : public Primitive { compile_options_); } + const std::shared_ptr& shared_source() const { + return source_; + } + private: std::string name_; - std::string source_; + std::shared_ptr source_; + size_t source_hash_; std::tuple grid_; std::tuple threadgroup_; std::vector> shape_infos_; diff --git a/python/tests/test_fast.py b/python/tests/test_fast.py index dfdcc18242..6d242bd52f 100644 --- a/python/tests/test_fast.py +++ b/python/tests/test_fast.py @@ -1129,6 +1129,59 @@ def call_kernel(a: mx.array, source): out = call_kernel(a, source) self.assertTrue(mx.array_equal(out, mx.ones_like(out))) + # `int` 1 and `bool` True print the same template value but + # declare a different template parameter. + kernel = mx.fast.metal_kernel( + name="sizeof_kernel", + input_names=["inp"], + output_names=["out"], + source="out[0] = sizeof(V);", + ) + + def call_sizeof(value): + return kernel( + inputs=[a], + grid=(1, 1, 1), + threadgroup=(1, 1, 1), + output_shapes=[(1,)], + output_dtypes=[mx.int32], + template=[("V", value)], + stream=mx.gpu, + )[0].item() + + for _ in range(2): + self.assertEqual(call_sizeof(1), 4) + self.assertEqual(call_sizeof(True), 1) + + @unittest.skipIf(not mx.metal.is_available(), "Metal is not available") + def test_custom_kernel_many_template_values(self): + # More variants than the kernel keeps built. An array made before + # its variant was evicted must still evaluate correctly. + kernel = mx.fast.metal_kernel( + name="value_kernel", + input_names=["inp"], + output_names=["out"], + source="out[0] = V;", + ) + a = mx.zeros((1,), dtype=mx.int32) + + def call_value(value): + return kernel( + inputs=[a], + grid=(1, 1, 1), + threadgroup=(1, 1, 1), + output_shapes=[(1,)], + output_dtypes=[mx.int32], + template=[("V", value)], + stream=mx.gpu, + )[0] + + pending = call_value(0) + for value in range(1, 200): + self.assertEqual(call_value(value).item(), value) + self.assertEqual(pending.item(), 0) + self.assertEqual(call_value(0).item(), 0) + @unittest.skipIf(not mx.metal.is_available(), "Metal is not available") def test_custom_kernel_same_name_different_source_one_eval(self): # Regression test for #3832: two kernels sharing a name but with diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index f3aa6a5935..53f1f0db22 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -16,7 +16,8 @@ if(MLX_BUILD_CUDA) endif() if(MLX_BUILD_METAL) - list(APPEND GPU_TEST_SOURCES residency_tests.cpp metal_counters_tests.cpp) + list(APPEND GPU_TEST_SOURCES residency_tests.cpp metal_counters_tests.cpp + metal_kernel_tests.cpp) endif() include(${doctest_SOURCE_DIR}/scripts/cmake/doctest.cmake) diff --git a/tests/metal_kernel_tests.cpp b/tests/metal_kernel_tests.cpp new file mode 100644 index 0000000000..d97d4760f5 --- /dev/null +++ b/tests/metal_kernel_tests.cpp @@ -0,0 +1,163 @@ +// Copyright © 2026 Apple Inc. + +#include + +#include "doctest/doctest.h" + +#include "mlx/fast_primitives.h" +#include "mlx/mlx.h" + +using namespace mlx::core; + +namespace { + +using TemplateArgs = std::vector>; + +fast::CustomKernelFunction sizeof_kernel() { + return fast::metal_kernel( + "sizeof_kernel", + {"inp"}, + {"out"}, + "out[0] = static_cast(sizeof(V));"); +} + +array call_sizeof( + const fast::CustomKernelFunction& kernel, + const array& inp, + const TemplateArgs& template_args, + StreamOrDevice s = Device::gpu) { + return kernel( + {inp}, + {{1}}, + {int32}, + {1, 1, 1}, + {1, 1, 1}, + template_args, + {}, + false, + s)[0]; +} + +// Read before evaluating `out`: evaluation detaches the primitive. +std::pair name_and_source(const array& out) { + auto state = static_cast(out.primitive()).state(); + return {std::get<0>(state), std::get<1>(state)}; +} + +} // namespace + +TEST_CASE("test metal kernel template argument kinds do not collide") { + auto kernel = sizeof_kernel(); + auto inp = zeros({1}, int32); + + // `int` 1 and `bool` true print the same template value. + auto as_int = call_sizeof(kernel, inp, {{"V", 1}}); + auto as_bool = call_sizeof(kernel, inp, {{"V", true}}); + CHECK_NE(name_and_source(as_int).second, name_and_source(as_bool).second); + eval(as_int, as_bool); + CHECK_EQ(as_int.item(), sizeof(int)); + CHECK_EQ(as_bool.item(), sizeof(bool)); + + // Memoized variants give the same results again. + auto as_bool_again = call_sizeof(kernel, inp, {{"V", true}}); + auto as_int_again = call_sizeof(kernel, inp, {{"V", 1}}); + CHECK_EQ(as_bool_again.item(), sizeof(bool)); + CHECK_EQ(as_int_again.item(), sizeof(int)); +} + +TEST_CASE("test metal kernel memoized source matches a fresh build") { + auto kernel = sizeof_kernel(); + std::vector> variants = { + {array(1, int32), {{"V", 1}}}, + {array(1, float32), {{"V", 1}}}, + {zeros({2}, int32), {{"V", 1}}}, + {zeros({16}, int32), {{"V", 1}}}, + {zeros({16}, int32), {{"V", 2}}}, + {zeros({16}, int32), {{"V", false}}}, + {zeros({16}, float16), {{"V", float16}}}, + {zeros({16}, float16), {{"V", bfloat16}}}, + }; + for (int round = 0; round < 2; ++round) { + for (const auto& [inp, template_args] : variants) { + auto memoized = call_sizeof(kernel, inp, template_args); + auto fresh = call_sizeof(sizeof_kernel(), inp, template_args); + CHECK_EQ(name_and_source(memoized), name_and_source(fresh)); + CHECK_EQ(memoized.item(), fresh.item()); + } + } +} + +TEST_CASE("test metal kernel copies share the memo across threads") { + auto kernel = sizeof_kernel(); + auto s = default_stream(Device::gpu); + auto inp = zeros({1}, int32); + constexpr int n_threads = 8; + constexpr int n_calls = 64; + std::vector> outs(n_threads); + std::vector threads; + for (int t = 0; t < n_threads; ++t) { + threads.emplace_back([&, t, kernel]() { + for (int i = 0; i < n_calls; ++i) { + TemplateArgs template_args; + if ((t + i) % 2 == 0) { + template_args = {{"V", 1}}; + } else { + template_args = {{"V", true}}; + } + outs[t].push_back(call_sizeof(kernel, inp, template_args, s)); + } + }); + } + for (auto& thread : threads) { + thread.join(); + } + for (int t = 0; t < n_threads; ++t) { + eval(outs[t]); + for (int i = 0; i < n_calls; ++i) { + auto expected = (t + i) % 2 == 0 ? sizeof(int) : sizeof(bool); + CHECK_EQ(outs[t][i].item(), expected); + } + } +} + +TEST_CASE("test metal kernel variant memo is bounded") { + auto kernel = fast::metal_kernel( + "value_kernel", {"inp"}, {"out"}, "out[0] = static_cast(V);"); + auto inp = zeros({1}, int32); + auto call = [&](int v) { return call_sizeof(kernel, inp, {{"V", v}}); }; + auto source_of = [](const array& out) { + return static_cast(out.primitive()) + .shared_source(); + }; + + // A repeat call shares the memoized source. + auto pending = call(0); + auto pending_source = source_of(pending); + CHECK_EQ(source_of(call(0)), pending_source); + + constexpr int n_values = 3 * fast::metal_kernel_max_cached_variants + 5; + std::vector> sources; + for (int v = 1; v < n_values; ++v) { + auto out = call(v); + sources.push_back(source_of(out)); + CHECK_EQ(out.item(), v); + } + + // Only the memo holds the sources of evaluated arrays. + size_t live = 0; + for (auto& source : sources) { + live += !source.expired(); + } + CHECK_GT(live, 0); + CHECK_LE(live, fast::metal_kernel_max_cached_variants); + + // The entry for 0 was evicted: a new call builds an equal, new source. + auto rebuilt = call(0); + CHECK_NE(source_of(rebuilt), pending_source); + CHECK_EQ(*source_of(rebuilt), *pending_source); + + // The pending array still owns its source and evaluates correctly. + pending_source.reset(); + CHECK_EQ(pending.item(), 0); + CHECK_EQ(rebuilt.item(), 0); +} From 87e099a58c027a887fefd6e30b2dfeb0cb2a2acf Mon Sep 17 00:00:00 2001 From: Jeethu Rao Date: Sun, 11 Oct 2026 13:31:11 +0100 Subject: [PATCH 09/34] metal: gate NAX off for gen-17 (M5-class) at every NAX dispatch site (#282) Source: numen-tech/mlx 73db231d1 ("gate nax off for gen-17 devices", fork PR #4). M5-class gen-17 parts (applegpu_g17*) compute wrong results in the NAX steel-gemm and qmm_t kernels; the fallback bit-matches stock mlx. is_nax_available() now requires gen >= 18 on every device class (upstream: gen >= (arch == 'p' ? 18 : 17)). The unused `arch` local is dropped instead of the fork's `(void)arch`. The gate is one predicate. Audit of every NAX dispatch decision on upstream main 06eb7483f + patches 1-7 (`git grep -n -iE nax mlx/backend/metal`, including kernels/*.h, jit/ and the JIT/no-JIT kernel getters), each site -> the predicate it asks: matmul.cpp:964 steel_matmul_axpby use_nax (thin #4654 :995, NAX split-K :1050, regular fused :1087) -> is_nax_available() matmul.cpp:2964 gather_mm_rhs_nax -> is_nax_available() matmul.cpp:3043 segmented_mm use_nax (:3060, :3103) -> is_nax_available() quantized.cpp:1243 qmm -> qmm_nax (qmm_t_nax/qmm_n_nax, incl. the _ib and affine_sym variants from patches 2-3) -> is_nax_available() quantized.cpp:1522 gather_qmm -> gather_qmm_nax -> is_nax_available() quantized.cpp:1917 gather_qmm_rhs -> gather_qmm_rhs_nax -> is_nax_available() quantized.cpp:2468 global-scale gather_qqmm K alignment (#4481; the matrix kernels then dispatch via :1917) -> is_nax_available() scaled_dot_product_attention.cpp:29 pre-M5 D512 non-NAX path (#4518) -> !is_nax_available() scaled_dot_product_attention.cpp:412 sdpa_full_self_attention_nax -> is_nax_available() scaled_dot_product_attention.cpp:436 D72/D80 pad to NAX (#4455) -> is_nax_available() scaled_dot_product_attention.cpp:964 D512 support check (#4487) -> is_nax_available() scaled_dot_product_attention.cpp:1436 D512 fused eligibility (#4487) -> is_nax_available() scaled_dot_product_attention.cpp:1445 D256 masked/causal fused NAX (#3842, array mask #4416) -> is_nax_available() scaled_dot_product_attention.cpp:1452 D256 non-NAX fused (#4505) -> !is_nax_available() gated_delta_update.cpp:18, :24 chunk C = 16 (the only value that selects gated_delta_*_nax kernels, :133/:335/:369/:416) only when NAX (#4020, #4565) -> is_nax_available() No other path selects a NAX kernel: the get_*_nax_kernel getters (jit_kernels.cpp, nojit_kernels.cpp) are only called from the functions above, get_sdpa_vjp_nax_kernel has no caller, the MPP tensor_ops code in kernels/ lives only in *_nax.h/.metal, and nothing outside mlx/backend/metal mentions NAX. The remaining raw get_architecture_gen()/get_architecture() uses pick tiles or non-NAX kernels (qmv batch limits quantized.cpp:99-134, nvfp4 narrow qmv quantized.cpp:550, affine qmv_wide gen 15+ quantized.cpp:633, gemv stream matmul.cpp:1416, GEMM_TPARAM/devc tile tunings incl. #4447). This host is an M3 Max (gen 15, no NAX), so the change is a no-op here; the gen-17 runtime check is the iPad M5 leg (gemma4-qat #282 T9/T11). test_quantized/test_fast/test_blas/test_fast_sdpa: 162 passed, 5 skipped. --- mlx/backend/metal/device.cpp | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/mlx/backend/metal/device.cpp b/mlx/backend/metal/device.cpp index ea96acea84..5c54dcb2cc 100644 --- a/mlx/backend/metal/device.cpp +++ b/mlx/backend/metal/device.cpp @@ -1004,9 +1004,16 @@ bool is_nax_available() { can_use_nax = true; } auto& d = metal::device(mlx::core::Device::gpu); - auto arch = d.get_architecture().back(); auto gen = d.get_architecture_gen(); - can_use_nax &= gen >= (arch == 'p' ? 18 : 17); + // gen-17 desktop parts (M5-class, applegpu_g17*) take the nax path under + // the old `gen >= 17` gate but compute WRONG results in the nax steel-gemm + // and qmm_t kernels (measured 2026-07-05: fp16 GEMM max|err| ~4 vs fp32 for + // every shape routed to steel_gemm_fused_nax; quantized qmm_t ~400 abs err; + // non-nax fallback bit-matches stock mlx). Require gen >= 18 on every + // device class until the g17 path is fixed and correctness-gated. Every + // NAX dispatch site (matmul, quantized, sdpa, gated delta) asks this + // predicate, so this one gate keeps gen 17 off all of them. + can_use_nax &= gen >= 18; return can_use_nax; }; static bool is_nax_available_ = _check_nax(); From cffb7f5e30f51c4c85d4d83a2a89691a8b7299a7 Mon Sep 17 00:00:00 2001 From: Jeethu Rao Date: Sun, 11 Oct 2026 13:35:43 +0100 Subject: [PATCH 10/34] metal: keep transcendentals precise unless MathMode::Fast (#282) Source: numen-tech/mlx c73162688 (applied unchanged). On macOS 15 / iOS 18+, set_compile_options set only mathMode, and MTLCompileOptions.mathFloatingPointFunctions (independent, default Fast) stayed Fast, so a Safe/Relaxed runtime compile (JIT kernels, compiled fusions, custom kernels) still emitted the approximate exp/log/sin/cos/pow. It is now Fast only for MathMode::Fast and Precise otherwise; the pre-macOS-15 branch keeps setFastMathEnabled. Checked against upstream #4461 (365bd0fba, "Use precise::exp in Sigmoid so compiled and eager sigmoid agree"): that fix only respells Sigmoid's exp as metal::precise::exp. It does not change compile options, and other runtime-compiled ops still spell unqualified metal:: transcendentals (binary_ops.h LogAddExp exp/log1p, Power pow, complex log/exp/sin/cos/ atan2, ...), so upstream covers one op of what this rule covers. Nothing is dropped. The two compose: an explicit metal::precise:: call stays precise under every mode, including Fast. Tests: test_ops.py -k "exp or log or sin or cos or tanh or erf or softmax or sigmoid or power or sqrt" -> 29 passed. test_compile.py (incl. #4461's test_compile_sigmoid_matches_eager) + test_quantized.py + test_fast.py (incl. test_custom_metal_kernel_math_mode) -> 174 passed, 2 skipped. --- mlx/backend/metal/device.cpp | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/mlx/backend/metal/device.cpp b/mlx/backend/metal/device.cpp index 5c54dcb2cc..93f3fe228c 100644 --- a/mlx/backend/metal/device.cpp +++ b/mlx/backend/metal/device.cpp @@ -59,6 +59,14 @@ void set_compile_options( default: throw std::invalid_argument("[metal::Device] Invalid math mode."); } + // MTLCompileOptions.mathFloatingPointFunctions is independent of mathMode + // and defaults to Fast, so a Safe/Relaxed request would still compile the + // fast (approximate) transcendental variants. Keep them Precise unless + // fast math was explicitly asked for. + mtl_options->setMathFloatingPointFunctions( + compile_options.math_mode == MathMode::Fast + ? MTL::MathFloatingPointFunctionsFast + : MTL::MathFloatingPointFunctionsPrecise); } else { if (compile_options.math_mode == MathMode::Relaxed) { throw std::runtime_error( From 5d2510a7efd448540d092c86c848d38018c2e24e Mon Sep 17 00:00:00 2001 From: Jeethu Rao Date: Sun, 11 Oct 2026 13:40:46 +0100 Subject: [PATCH 11/34] fast: spec_decode_verify op (CPU composition + fused Metal) (#282) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit mx.fast.spec_decode_verify(draft_tokens [B, K], target_logits [B, K+1, V]) -> (n_accepted [B], committed [B, K+1]): greedy speculative verify. The target argmax runs as a regular op; the SpecDecodeVerify primitive does the per-row prefix match in a one-thread-per-row Metal kernel, with an op-composition fallback on CPU / no-GPU / CUDA. Source: 3e1d206bb (v1 composition), 28efbe632 (fused Metal v2), 36030f6f8 (docs), 49840c9e9 (fuzz), e19719180 (test moved into test_fast.py), taken as `git diff 3e1d206b^ e1971918` minus quantized.cpp / test_quantized.py (cddb7bed7 in that range is patch 5). Applied unchanged, except: - add/add conflicts with upstream: kernels/CMakeLists.txt (after the sdpa_vjp / gated_delta build_kernel lines), no_gpu/primitives.cpp (after GatedDeltaUpdate[VJP]::use_fallback), test_fast.py (after #4584's test_custom_metal_kernel_negative_template_int); both sides kept; - clang-format (fast.cpp, fast.h) and ruff 0.16.10 format (test_fast.py) of the new code, whitespace only. No caller in NumenKit or Python; carried as is (spec §2.10). Tests: test_fast.py -k spec -> 1 passed, 76 subtests. test_quantized.py + test_fast.py -> 104 passed, 2 skipped. --- docs/src/python/fast.rst | 1 + mlx/backend/cuda/primitives.cpp | 1 + mlx/backend/metal/CMakeLists.txt | 1 + mlx/backend/metal/kernels/CMakeLists.txt | 1 + mlx/backend/metal/kernels/spec_decode.metal | 40 ++++++++++ mlx/backend/metal/spec_decode.cpp | 50 ++++++++++++ mlx/backend/no_gpu/primitives.cpp | 5 ++ mlx/fast.cpp | 86 +++++++++++++++++++++ mlx/fast.h | 14 ++++ mlx/fast_primitives.h | 30 +++++++ python/src/fast.cpp | 27 +++++++ python/tests/test_fast.py | 82 ++++++++++++++++++++ 12 files changed, 338 insertions(+) create mode 100644 mlx/backend/metal/kernels/spec_decode.metal create mode 100644 mlx/backend/metal/spec_decode.cpp diff --git a/docs/src/python/fast.rst b/docs/src/python/fast.rst index c930c7bb4a..7c23055559 100644 --- a/docs/src/python/fast.rst +++ b/docs/src/python/fast.rst @@ -13,6 +13,7 @@ Fast cross_entropy rope scaled_dot_product_attention + spec_decode_verify metal_kernel cuda_kernel precompiled_cuda_kernel diff --git a/mlx/backend/cuda/primitives.cpp b/mlx/backend/cuda/primitives.cpp index 104c7fd3fc..e91771a1c2 100644 --- a/mlx/backend/cuda/primitives.cpp +++ b/mlx/backend/cuda/primitives.cpp @@ -49,6 +49,7 @@ NO_GPU_MULTI(SVD) NO_GPU(Inverse) NO_GPU_MULTI(Eig) NO_GPU_MULTI(Eigh) +NO_GPU_USE_FALLBACK(fast::SpecDecodeVerify) namespace fast { NO_GPU_MULTI(GatedDeltaUpdate) diff --git a/mlx/backend/metal/CMakeLists.txt b/mlx/backend/metal/CMakeLists.txt index cccaad8a5c..d84d06bf43 100644 --- a/mlx/backend/metal/CMakeLists.txt +++ b/mlx/backend/metal/CMakeLists.txt @@ -150,6 +150,7 @@ target_sources( ${CMAKE_CURRENT_SOURCE_DIR}/scan.cpp ${CMAKE_CURRENT_SOURCE_DIR}/slicing.cpp ${CMAKE_CURRENT_SOURCE_DIR}/softmax.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/spec_decode.cpp ${CMAKE_CURRENT_SOURCE_DIR}/sort.cpp ${CMAKE_CURRENT_SOURCE_DIR}/reduce.cpp ${CMAKE_CURRENT_SOURCE_DIR}/ternary.cpp diff --git a/mlx/backend/metal/kernels/CMakeLists.txt b/mlx/backend/metal/kernels/CMakeLists.txt index 355c763a3a..9e6bd77d71 100644 --- a/mlx/backend/metal/kernels/CMakeLists.txt +++ b/mlx/backend/metal/kernels/CMakeLists.txt @@ -59,6 +59,7 @@ build_kernel(scaled_dot_product_attention sdpa_blocked.h sdpa_vector.h) build_kernel(scaled_dot_product_attention_vjp) build_kernel(gated_delta_update gated_delta_update.h) build_kernel(gated_delta_update_vjp gated_delta_update_vjp.h) +build_kernel(spec_decode) if(MLX_METAL_VERSION GREATER_EQUAL 320) build_kernel(fence) endif() diff --git a/mlx/backend/metal/kernels/spec_decode.metal b/mlx/backend/metal/kernels/spec_decode.metal new file mode 100644 index 0000000000..61b6c66c9b --- /dev/null +++ b/mlx/backend/metal/kernels/spec_decode.metal @@ -0,0 +1,40 @@ +// Copyright © 2024 Apple Inc. + +#include "mlx/backend/metal/kernels/utils.h" + +using namespace metal; + +// Greedy speculative-decoding verify, one thread per batch row. +// draft [B, K] int32 — drafted tokens +// target [B, K+1] int32 — target argmax tokens over [last] + draft +// Writes: +// n_accepted [B] — longest matching prefix length (0..K) +// committed [B, K+1] — draft prefix (< n) then the corrected token at n. +[[kernel]] void spec_decode_verify( + const device int* draft [[buffer(0)]], + const device int* target [[buffer(1)]], + device int* n_accepted [[buffer(2)]], + device int* committed [[buffer(3)]], + constant int& K [[buffer(4)]], + constant int& B [[buffer(5)]], + uint b [[thread_position_in_grid]]) { + if (b >= uint(B)) { + return; + } + int row = int(b); + const device int* d = draft + row * K; + const device int* t = target + row * (K + 1); + device int* c = committed + row * (K + 1); + + int n = K; + for (int j = 0; j < K; ++j) { + if (d[j] != t[j]) { + n = j; + break; + } + } + n_accepted[row] = n; + for (int j = 0; j < K + 1; ++j) { + c[j] = (j < n) ? d[j] : ((j == n) ? t[n] : 0); + } +} diff --git a/mlx/backend/metal/spec_decode.cpp b/mlx/backend/metal/spec_decode.cpp new file mode 100644 index 0000000000..f6134d931d --- /dev/null +++ b/mlx/backend/metal/spec_decode.cpp @@ -0,0 +1,50 @@ +// Copyright © 2024 Apple Inc. + +#include + +#include "mlx/allocator.h" +#include "mlx/backend/metal/device.h" +#include "mlx/backend/metal/utils.h" +#include "mlx/fast_primitives.h" + +namespace mlx::core::fast { + +bool SpecDecodeVerify::use_fallback(Stream s) { + // Fused kernel only on the GPU; CPU uses the op-level composition fallback. + return s.device == Device::cpu; +} + +void SpecDecodeVerify::eval_gpu( + const std::vector& inputs, + std::vector& outputs) { + auto& draft = inputs[0]; // [B, K] int32 + auto& target = inputs[1]; // [B, K+1] int32 + auto& n_accepted = outputs[0]; // [B] int32 + auto& committed = outputs[1]; // [B, K+1] int32 + + n_accepted.set_data(allocator::malloc(n_accepted.nbytes())); + committed.set_data(allocator::malloc(committed.nbytes())); + + int B = draft.shape(0); + int K = draft.shape(1); + + auto& s = stream(); + auto& d = metal::device(s.device); + auto kernel = d.get_kernel("spec_decode_verify"); + auto& compute_encoder = metal::get_command_encoder(s); + compute_encoder.set_compute_pipeline_state(kernel); + compute_encoder.set_input_array(draft, 0); + compute_encoder.set_input_array(target, 1); + compute_encoder.set_output_array(n_accepted, 2); + compute_encoder.set_output_array(committed, 3); + compute_encoder.set_bytes(K, 4); + compute_encoder.set_bytes(B, 5); + + // One thread per batch row (K is tiny; no reduction). + int tgroup = std::min(B, 256); + MTL::Size grid_dims = MTL::Size(B, 1, 1); + MTL::Size group_dims = MTL::Size(tgroup, 1, 1); + compute_encoder.dispatch_threads(grid_dims, group_dims); +} + +} // namespace mlx::core::fast diff --git a/mlx/backend/no_gpu/primitives.cpp b/mlx/backend/no_gpu/primitives.cpp index e3ee598ea3..af278813f3 100644 --- a/mlx/backend/no_gpu/primitives.cpp +++ b/mlx/backend/no_gpu/primitives.cpp @@ -72,6 +72,11 @@ bool fast::GatedDeltaUpdateVJP::use_fallback( return true; } +bool fast::SpecDecodeVerify::use_fallback(Stream s) { + return true; +} +NO_GPU_MULTI(fast::SpecDecodeVerify) + NO_GPU(Abs) NO_GPU(Add) NO_GPU(AddMM) diff --git a/mlx/fast.cpp b/mlx/fast.cpp index 2afb834166..4d1bd5a59f 100644 --- a/mlx/fast.cpp +++ b/mlx/fast.cpp @@ -1245,4 +1245,90 @@ bool ConvertFP8::is_equivalent(const Primitive& other) const { return to_fp8_ == a_other.to_fp8_; } +std::vector spec_decode_verify( + const array& draft_tokens, + const array& target_logits, + StreamOrDevice s_ /* = {} */) { + if (draft_tokens.ndim() != 2) { + throw std::invalid_argument( + "[spec_decode_verify] draft_tokens must be 2D [B, K]."); + } + if (target_logits.ndim() != 3) { + throw std::invalid_argument( + "[spec_decode_verify] target_logits must be 3D [B, K+1, V]."); + } + int B = draft_tokens.shape(0); + int K = draft_tokens.shape(1); + if (target_logits.shape(0) != B) { + throw std::invalid_argument( + "[spec_decode_verify] draft_tokens and target_logits batch size must match."); + } + if (target_logits.shape(1) != K + 1) { + throw std::invalid_argument( + "[spec_decode_verify] target_logits.shape[1] must equal draft_tokens.shape[1] + 1."); + } + auto s = to_stream(s_); + + // Reuse mlx's fast argmax on-device (result stays on the GPU); the fused + // primitive does only the tiny per-row prefix-match. Primitive inputs are + // int32 token arrays, so the kernel needs no large-V reduction. + auto dft = astype(draft_tokens, int32, s); + auto tgt = astype(argmax(target_logits, -1, false, s), int32, s); // [B, K+1] + + // Fallback (CPU path + autodiff): the same verify as a plain op composition, + // operating on (draft_tokens, target_tokens). This is the correctness oracle + // the Metal kernel is checked against. + auto fallback = + [B, K, s](const std::vector& inputs) -> std::vector { + auto d = inputs[0]; // [B, K] int32 + auto t = inputs[1]; // [B, K+1] int32 + auto t_pref = slice(t, Shape{0, 0}, Shape{B, K}, s); + auto mism = not_equal(d, t_pref, s); + auto j = broadcast_to( + reshape(arange(K, int32, s), Shape{1, K}, s), Shape{B, K}, s); + auto n_acc = + min(where(mism, j, full(Shape{B, K}, K, int32, s), s), + /*axis=*/1, + /*keepdims=*/false, + s); // [B] + auto n_acc2 = reshape(n_acc, Shape{B, 1}, s); + auto corrected = take_along_axis(t, n_acc2, /*axis=*/1, s); // [B, 1] + auto j1 = broadcast_to( + reshape(arange(K + 1, int32, s), Shape{1, K + 1}, s), + Shape{B, K + 1}, + s); + auto nacc_b = broadcast_to(n_acc2, Shape{B, K + 1}, s); + auto d_ext = concatenate({d, zeros(Shape{B, 1}, int32, s)}, /*axis=*/1, s); + auto corr_b = broadcast_to(corrected, Shape{B, K + 1}, s); + auto committed = where( + less(j1, nacc_b, s), + d_ext, + where( + equal(j1, nacc_b, s), corr_b, zeros(Shape{B, K + 1}, int32, s), s), + s); + return {n_acc, committed}; + }; + + std::vector prim_inputs = {dft, tgt}; + if (SpecDecodeVerify::use_fallback(s)) { + return fallback(prim_inputs); + } + return array::make_arrays( + {Shape{B}, Shape{B, K + 1}}, + {int32, int32}, + std::make_shared(s, fallback), + prim_inputs); +} + +bool SpecDecodeVerify::is_equivalent(const Primitive&) const { + return true; // no parameters +} + +std::vector SpecDecodeVerify::output_shapes( + const std::vector& inputs) { + int B = inputs[0].shape(0); + int K = inputs[0].shape(1); + return {Shape{B}, Shape{B, K + 1}}; +} + } // namespace mlx::core::fast diff --git a/mlx/fast.h b/mlx/fast.h index 14a46de657..a65246440e 100644 --- a/mlx/fast.h +++ b/mlx/fast.h @@ -117,4 +117,18 @@ MLX_API std::vector precompiled_cuda_kernel( bool ensure_row_contiguous = false, StreamOrDevice s = {}); +/** Fused greedy speculative-decoding verify. + * + * draft_tokens: int32 [B, K] — the K drafted tokens + * target_logits: floating [B, K+1, V] — target logits over [last] + draft + * + * Returns {n_accepted [B] int32, committed [B, K+1] int32}: the longest + * matching greedy prefix length, and (accepted draft prefix ++ corrected + * token). v1 is greedy only (temperature == 0); the caller prepends the seed + * `last` token. */ +MLX_API std::vector spec_decode_verify( + const array& draft_tokens, + const array& target_logits, + StreamOrDevice s = {}); + } // namespace mlx::core::fast diff --git a/mlx/fast_primitives.h b/mlx/fast_primitives.h index 3e3240970c..7636ac1a87 100644 --- a/mlx/fast_primitives.h +++ b/mlx/fast_primitives.h @@ -610,4 +610,34 @@ class CustomKernel : public Primitive { CompileOptions::Data compile_options_; }; +// Fused greedy speculative-decoding verify. Inputs: draft_tokens int32 [B, K] +// and target_tokens int32 [B, K+1] (the target argmax). Outputs: n_accepted +// [B] int32 and committed [B, K+1] int32. The op-level fallback (a plain +// composition) runs on CPU and for autodiff; the Metal eval_gpu runs a trivial +// per-row prefix-match kernel. +class SpecDecodeVerify : public Custom { + public: + SpecDecodeVerify( + Stream stream, + std::function(std::vector)> fallback) + : Custom(stream, std::move(fallback)) {} + + static bool use_fallback(Stream stream); + + void eval_cpu(const std::vector& inputs, std::vector& outputs) + override { + throw std::runtime_error( + "[SpecDecodeVerify] no eval_cpu; CPU uses the op fallback."); + } + void eval_gpu(const std::vector& inputs, std::vector& outputs) + override; + + DEFINE_NAME(SpecDecodeVerify); + bool is_equivalent(const Primitive& other) const override; + std::vector output_shapes(const std::vector& inputs) override; + auto state() const { + return std::make_tuple(nullptr); + } +}; + } // namespace mlx::core::fast diff --git a/python/src/fast.cpp b/python/src/fast.cpp index f9c472ab99..487bc2012f 100644 --- a/python/src/fast.cpp +++ b/python/src/fast.cpp @@ -118,6 +118,33 @@ void init_fast(nb::module_& parent_module) { auto m = parent_module.def_submodule("fast", "mlx.core.fast: fast operations"); + m.def( + "spec_decode_verify", + &mx::fast::spec_decode_verify, + "draft_tokens"_a, + "target_logits"_a, + nb::kw_only(), + "stream"_a = nb::none(), + nb::sig( + "def spec_decode_verify(draft_tokens: array, target_logits: array, *, stream: Union[None, Stream, Device] = None) -> Tuple[array, array]"), + R"pbdoc( + Fused greedy speculative-decoding verify (v1). + + Given the drafted tokens and the target model's logits over + ``[last] + draft``, returns the greedy accept result. + + Args: + draft_tokens (array): int32 ``[B, K]`` drafted tokens. + target_logits (array): ``[B, K+1, V]`` target logits over the + ``[last] + draft`` positions. + + Returns: + tuple(array, array): ``(n_accepted [B] int32, committed [B, K+1] int32)`` + where ``committed[b, :n_accepted[b]]`` is the accepted draft prefix and + ``committed[b, n_accepted[b]]`` is the corrected (bonus) token. The + caller prepends the seed ``last`` token. Greedy only (temperature == 0). + )pbdoc"); + m.def( "rms_norm", &mx::fast::rms_norm, diff --git a/python/tests/test_fast.py b/python/tests/test_fast.py index 6d242bd52f..a10e699b2a 100644 --- a/python/tests/test_fast.py +++ b/python/tests/test_fast.py @@ -2,6 +2,7 @@ import math import os +import random import unittest import mlx.core as mx @@ -1356,6 +1357,87 @@ def test_custom_metal_kernel_negative_template_int(self): mx.eval(out) self.assertEqual(out.tolist(), [-1, -1, -1, -1]) + def test_spec_decode_verify(self): + # Validate mx.fast.spec_decode_verify against a pure-Python oracle that + # mirrors the greedy verify loop. On a Metal build the default device + # exercises the fused kernel and stream=mx.cpu the op-composition + # fallback; both must agree with the oracle. + def oracle(draft, tgt_tokens): + B, K = len(draft), len(draft[0]) + n_acc, committed = [], [] + for b in range(B): + na = K + for j in range(K): + if draft[b][j] != tgt_tokens[b][j]: + na = j + break + n_acc.append(na) + committed.append(draft[b][:na] + [tgt_tokens[b][na]]) + return n_acc, committed + + # One-hot logits make the greedy argmax token unambiguous while still + # driving the same large-V reduction the real caller feeds in. + def onehot_logits(tgt_tokens, V, dtype): + B, L = len(tgt_tokens), len(tgt_tokens[0]) + lg = [[[-10.0] * V for _ in range(L)] for _ in range(B)] + for b in range(B): + for j in range(L): + lg[b][j][tgt_tokens[b][j]] = 10.0 + return mx.array(lg).astype(dtype) + + def check(draft, tgt_tokens, V, dtype, stream): + exp_n, exp_c = oracle(draft, tgt_tokens) + lg = onehot_logits(tgt_tokens, V, dtype) + kw = {"stream": stream} if stream is not None else {} + n, c = mx.fast.spec_decode_verify(mx.array(draft, dtype=mx.int32), lg, **kw) + mx.eval(n, c) + n, c = n.tolist(), c.tolist() + got_c = [c[b][: exp_n[b] + 1] for b in range(len(draft))] + self.assertEqual(n, exp_n) + self.assertEqual(got_c, exp_c) + + cases = [ + ("full-accept", [[5, 6, 7]], [[5, 6, 7, 9]]), + ("mismatch@2", [[5, 6, 7]], [[5, 6, 1, 9]]), + ("reject-all", [[5, 6, 7]], [[2, 6, 7, 9]]), + ("K=1", [[3]], [[3, 4]]), + ( + "batch2-K4", + [[1, 2, 3, 4], [1, 2, 3, 4]], + [[1, 2, 9, 4, 7], [1, 2, 3, 4, 8]], + ), + ("batch3-K1", [[7], [7], [7]], [[7, 1], [2, 3], [7, 9]]), + ] + # stream=None is the default device (fused kernel on Metal); mx.cpu is + # the op-composition fallback. + for path, stream in (("default", None), ("cpu-fallback", mx.cpu)): + for dtype in (mx.float32, mx.float16, mx.bfloat16): + for name, draft, tgt in cases: + with self.subTest(path=path, dtype=str(dtype), case=name): + check(draft, tgt, 32, dtype, stream) + + # Randomized sweep over prefix lengths, batch shapes, and corrected-token + # positions that are awkward to enumerate by hand. + rng = random.Random(7) + for i in range(40): + b, k, v = rng.randint(1, 4), rng.randint(1, 8), rng.randint(8, 96) + target = [[rng.randrange(v) for _ in range(k + 1)] for _ in range(b)] + draft = [ + [ + row[j] if rng.random() < 0.5 else (row[j] + rng.randrange(1, v)) % v + for j in range(k) + ] + for row in target + ] + with self.subTest(sweep=i, b=b, k=k, v=v): + check(draft, target, v, mx.float16, None) + + # target_logits.shape[1] must equal draft.shape[1] + 1. + with self.assertRaises(ValueError): + mx.fast.spec_decode_verify( + mx.zeros((1, 3), dtype=mx.int32), mx.zeros((1, 3, 8)) + ) + if __name__ == "__main__": mlx_tests.MLXTestRunner() From 4783e7bcf3e88d82e280cd6f737ff14f6e0a8df2 Mon Sep 17 00:00:00 2001 From: Jeethu Rao Date: Sun, 11 Oct 2026 13:41:24 +0100 Subject: [PATCH 12/34] python/tests: clarify qmv tail/qmv_wide test comments (#282) Comment-only (gemma4-qat #282 ruling R12, Task 3 review minors): - test_qmv_fast_half_block: the "x is a view whose memory past K is non-zero" claim holds for B = 0 only. At B = 2 ensure_row_contiguous_matrix copies the (2, 1, K) slice, so past K batch 0 reads batch 1's x (still non-zero) and batch 1 reads past the buffer. - test_implied_bias_matmul: affine qmv_wide routing needs gen 15+. --- python/tests/test_quantized.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index dffe8d8811..9a67957eb2 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -1321,9 +1321,11 @@ def test_qmv_fast_half_block(self): # partial last block. The generic qmv baseline runs the same product # with K padded by one group that x zeroes out: K + gs is on no # qmv_fast gate. (An extra output row no longer leaves qmv_fast: since - # ml-explore/mlx#4516 N % 8 != 0 takes its partial_rows variant.) The - # fast call's x is a view whose memory past K is that group drawn - # non-zero, so a tail lane that read past K would show. + # ml-explore/mlx#4516 N % 8 != 0 takes its partial_rows variant.) At + # B = 0 the fast call's x is a view whose memory past K is that group + # drawn non-zero, so a tail lane that read past K would show. At B = 2 + # ensure_row_contiguous_matrix copies the slice: past K, batch 0 reads + # batch 1's x (still non-zero) and batch 1 reads past the buffer. key = mx.random.key(5) k1, k2 = mx.random.split(key) f32 = mx.float32 @@ -1444,8 +1446,9 @@ def test_implied_bias_matmul(self): for dtype, (bits, gs, factor) in product( self.IMPLIED_BIAS_DTYPES, self.IMPLIED_BIAS_CASES ): - # K = 64 / 128 route to qmv_quad; 512 / 1536 to qmv_fast (and - # qmv_wide for 2 <= M <= 8 at bits > 1, from M = 3 at 2-bit); + # K = 64 / 128 route to qmv_quad; 512 / 1536 to qmv_fast (and, + # on gen 15+, affine qmv_wide for 2 <= M <= 8 at bits > 1, from + # M = 3 at 2-bit); # K = 768 runs the 4-bit qmv_fast tail; N = 72 is not a # multiple of the 32-wide qmm tile. for K, N in product([64, 128, 512, 768, 1536], [256, 72]): From 26aaef369b3f60dde1bfef7ef1764b2f98e384be Mon Sep 17 00:00:00 2001 From: Jeethu Rao Date: Sun, 11 Oct 2026 14:14:40 +0100 Subject: [PATCH 13/34] fast: spec_decode_verify materializes non-row-contiguous drafts (#282) Codex round 1, P1 (pre-existing in the fork's 28efbe632). The fused kernel indexes dense rows (draft + row * K, target + row * (K + 1)), but eval_gpu bound its inputs as given, and the wrapper's astype(int32) keeps an int32 view a view. A transposed draft (logically [[1,2],[3,4]]) gave n_accepted [1, 0] on Metal vs [2, 2] on CPU; a broadcast [1,2] -> [2,2] draft made row 1 read past the stored elements; a column slice read the wrong elements. eval_gpu now copies a non-row-contiguous input (draft or target) into a dense temporary (contiguous_copy_gpu + add_temporary, as in gated_delta_update.cpp) before setting the kernel's pipeline. A row-contiguous input is returned untouched before the encoder is used: same path as before, no extra copy or dispatch. Targets come from argmax, which already returns dense rows from strided logits; the check covers them anyway. Test: test_spec_decode_verify_strided_inputs (transposed, broadcast, sliced and mismatching transposed drafts, plus a transposed logits view; Metal and CPU vs explicit expected results). RED (before): 4 failed -- every Metal draft-view subtest, e.g. case='transposed' path='default': Lists differ: [1, 0] != [2, 2] case='broadcast' path='default': Lists differ: [2, 0] != [2, 2] case='sliced' path='default': Lists differ: [2, 0] != [2, 2] case='mismatch-transposed' path='default': [1, 0] != [1, 2] (CPU and the transposed-logits subtests passed.) GREEN: test_fast.py -k spec -> 2 passed, 86 subtests passed. --- mlx/backend/metal/spec_decode.cpp | 24 +++++++++++-- python/tests/test_fast.py | 59 +++++++++++++++++++++++++++++++ 2 files changed, 80 insertions(+), 3 deletions(-) diff --git a/mlx/backend/metal/spec_decode.cpp b/mlx/backend/metal/spec_decode.cpp index f6134d931d..f1bf001510 100644 --- a/mlx/backend/metal/spec_decode.cpp +++ b/mlx/backend/metal/spec_decode.cpp @@ -3,12 +3,29 @@ #include #include "mlx/allocator.h" +#include "mlx/backend/gpu/copy.h" #include "mlx/backend/metal/device.h" #include "mlx/backend/metal/utils.h" #include "mlx/fast_primitives.h" namespace mlx::core::fast { +namespace { + +// The kernel indexes dense rows (row * K, row * (K + 1)). A row-contiguous +// input is bound as is; a view (transposed, sliced, broadcast) is copied +// first and the copy kept alive until the command buffer completes. +array ensure_row_contiguous(const array& x, const Stream& s) { + if (x.flags().row_contiguous) { + return x; + } + array x_copy = contiguous_copy_gpu(x, s); + metal::get_command_encoder(s).add_temporary(x_copy); + return x_copy; +} + +} // namespace + bool SpecDecodeVerify::use_fallback(Stream s) { // Fused kernel only on the GPU; CPU uses the op-level composition fallback. return s.device == Device::cpu; @@ -17,8 +34,10 @@ bool SpecDecodeVerify::use_fallback(Stream s) { void SpecDecodeVerify::eval_gpu( const std::vector& inputs, std::vector& outputs) { - auto& draft = inputs[0]; // [B, K] int32 - auto& target = inputs[1]; // [B, K+1] int32 + auto& s = stream(); + // Copies (if any) are encoded before this kernel's pipeline is set. + array draft = ensure_row_contiguous(inputs[0], s); // [B, K] int32 + array target = ensure_row_contiguous(inputs[1], s); // [B, K+1] int32 auto& n_accepted = outputs[0]; // [B] int32 auto& committed = outputs[1]; // [B, K+1] int32 @@ -28,7 +47,6 @@ void SpecDecodeVerify::eval_gpu( int B = draft.shape(0); int K = draft.shape(1); - auto& s = stream(); auto& d = metal::device(s.device); auto kernel = d.get_kernel("spec_decode_verify"); auto& compute_encoder = metal::get_command_encoder(s); diff --git a/python/tests/test_fast.py b/python/tests/test_fast.py index a10e699b2a..48dc678345 100644 --- a/python/tests/test_fast.py +++ b/python/tests/test_fast.py @@ -1438,6 +1438,65 @@ def check(draft, tgt_tokens, V, dtype, stream): mx.zeros((1, 3), dtype=mx.int32), mx.zeros((1, 3, 8)) ) + def test_spec_decode_verify_strided_inputs(self): + # Drafts and logits that are views, not dense rows: the fused kernel + # must see the logical values, like the CPU composition does. + def onehot(tokens, V=8): + lg = [[[-10.0] * V for _ in row] for row in tokens] + for b, row in enumerate(tokens): + for j, t in enumerate(row): + lg[b][j][t] = 10.0 + return mx.array(lg) + + full = mx.array([[9, 1, 2, 9], [9, 3, 4, 9]], dtype=mx.int32) + logits = onehot([[1, 2, 5], [3, 4, 6]]) + cases = [ + # (name, draft view, target logits, expected n, expected committed) + ( + "transposed", + mx.array([[1, 3], [2, 4]], dtype=mx.int32).T, + logits, + [2, 2], + [[1, 2, 5], [3, 4, 6]], + ), + ( + "broadcast", + mx.broadcast_to(mx.array([[1, 2]], dtype=mx.int32), (2, 2)), + onehot([[1, 2, 5], [1, 2, 6]]), + [2, 2], + [[1, 2, 5], [1, 2, 6]], + ), + ("sliced", full[:, 1:3], logits, [2, 2], [[1, 2, 5], [3, 4, 6]]), + ( + # Dense draft: isolates the logits view (argmax reads strides). + "transposed-logits", + mx.array([[1, 2], [3, 4]], dtype=mx.int32), + mx.swapaxes(mx.swapaxes(logits, 0, 1).astype(mx.float16), 0, 1), + [2, 2], + [[1, 2, 5], [3, 4, 6]], + ), + ( + "mismatch-transposed", + mx.array([[1, 3], [7, 4]], dtype=mx.int32).T, + logits, + [1, 2], + [[1, 2, 0], [3, 4, 6]], + ), + ] + for name, draft, lg, exp_n, exp_c in cases: + for path, stream in (("default", None), ("cpu-fallback", mx.cpu)): + with self.subTest(case=name, path=path): + kw = {"stream": stream} if stream is not None else {} + n, c = mx.fast.spec_decode_verify(draft, lg, **kw) + mx.eval(n, c) + self.assertEqual(n.tolist(), exp_n) + got = c.tolist() + # Positions past the bonus token are unspecified. + self.assertEqual( + [got[b][: exp_n[b] + 1] for b in range(len(exp_n))], + [exp_c[b][: exp_n[b] + 1] for b in range(len(exp_n))], + ) + if __name__ == "__main__": mlx_tests.MLXTestRunner() From e6b7bbe29ba49466cd4a135b7af814d41d8308b0 Mon Sep 17 00:00:00 2001 From: Jeethu Rao Date: Sun, 11 Oct 2026 14:18:10 +0100 Subject: [PATCH 14/34] fast: spec_decode_verify CPU path handles K = 0 like Metal (#282) Codex round 1, P2 (pre-existing in the fork's 3e1d206bb). With K = 0 the CPU / no-GPU composition took min over an empty axis, which ops.cpp rejects, so draft [B, 0] with logits [B, 1, V] threw on stream=mx.cpu while the fused kernel returns a result. Per ruling R16, K = 0 is not rejected: the composition now returns what Metal returns, measured first: n_accepted = 0 for every row and committed [B, 1] = the target argmax token (Metal: B=1 -> [0], [[5]]; B=2 -> [0, 0], [[5], [3]]). Test: test_spec_decode_verify_empty_draft (B = 1 and 2; Metal and CPU vs the explicit result, shapes included). RED (before): 2 failed -- both CPU subtests: ValueError: [min] Cannot min reduce over axis 1 with size 0. (The Metal subtests passed.) GREEN: test_fast.py -k spec -> 3 passed, 90 subtests passed; CPU and Metal both give int32 [0] / [[5]]. --- mlx/fast.cpp | 5 +++++ python/tests/test_fast.py | 19 +++++++++++++++++++ 2 files changed, 24 insertions(+) diff --git a/mlx/fast.cpp b/mlx/fast.cpp index 4d1bd5a59f..1813a68a9e 100644 --- a/mlx/fast.cpp +++ b/mlx/fast.cpp @@ -1282,6 +1282,11 @@ std::vector spec_decode_verify( [B, K, s](const std::vector& inputs) -> std::vector { auto d = inputs[0]; // [B, K] int32 auto t = inputs[1]; // [B, K+1] int32 + if (K == 0) { + // Nothing to accept; committed is the bonus token t[:, 0], as the + // fused kernel writes it. (min over the empty axis below would throw.) + return {zeros(Shape{B}, int32, s), t}; + } auto t_pref = slice(t, Shape{0, 0}, Shape{B, K}, s); auto mism = not_equal(d, t_pref, s); auto j = broadcast_to( diff --git a/python/tests/test_fast.py b/python/tests/test_fast.py index 48dc678345..e93796b6e0 100644 --- a/python/tests/test_fast.py +++ b/python/tests/test_fast.py @@ -1438,6 +1438,25 @@ def check(draft, tgt_tokens, V, dtype, stream): mx.zeros((1, 3), dtype=mx.int32), mx.zeros((1, 3, 8)) ) + def test_spec_decode_verify_empty_draft(self): + # K = 0: nothing to accept, the bonus token is the argmax at the only + # position. The CPU composition must match the fused kernel. + for B, bonus in ((1, [5]), (2, [5, 3])): + lg = [[[-10.0] * 8] for _ in range(B)] + for b, t in enumerate(bonus): + lg[b][0][t] = 10.0 + lg = mx.array(lg) + draft = mx.zeros((B, 0), dtype=mx.int32) + for path, stream in (("default", None), ("cpu-fallback", mx.cpu)): + with self.subTest(B=B, path=path): + kw = {"stream": stream} if stream is not None else {} + n, c = mx.fast.spec_decode_verify(draft, lg, **kw) + mx.eval(n, c) + self.assertEqual(n.shape, (B,)) + self.assertEqual(c.shape, (B, 1)) + self.assertEqual(n.tolist(), [0] * B) + self.assertEqual(c.tolist(), [[t] for t in bonus]) + def test_spec_decode_verify_strided_inputs(self): # Drafts and logits that are views, not dense rows: the fused kernel # must see the logical values, like the CPU composition does. From f54067ddfa0d9aafb0151cc978637ea6eb9a1c62 Mon Sep 17 00:00:00 2001 From: Jeethu Rao Date: Sun, 11 Oct 2026 14:31:14 +0100 Subject: [PATCH 15/34] fast: spec_decode_verify returns empty results for B = 0 without dispatching (#282) Ruling R17 (fix round 1). With B = 0 (drafts [0, K], logits [0, K+1, V], K = 3 and K = 0) the op threw on every backend before reaching the fused kernel: `[argmax] Cannot argmax reduce zero size array.` (also under MTL_DEBUG_LAYER=1 MTL_SHADER_VALIDATION=1, where nothing else is reported: the throw precedes any encode). The op now returns int32 zeros [0] and [0, K+1] right after the shape checks, with no argmax and no SpecDecodeVerify primitive, so nothing is dispatched. B > 0 inputs take exactly the previous path. Test: test_spec_decode_verify_empty_batch (K = 3 and 0; default device and CPU; shape and dtype of both outputs; default vs CPU equal). RED (before): 6 failed -- the 4 path subtests with the argmax ValueError, the 2 compare subtests with KeyError (no result to compare). GREEN: test_fast.py -k spec -> 4 passed, 96 subtests passed; both backends give n (0,) int32 and committed (0, K+1) int32, also clean under the Metal debug/shader-validation layers. --- mlx/fast.cpp | 6 ++++++ python/tests/test_fast.py | 17 +++++++++++++++++ 2 files changed, 23 insertions(+) diff --git a/mlx/fast.cpp b/mlx/fast.cpp index 1813a68a9e..7836624b52 100644 --- a/mlx/fast.cpp +++ b/mlx/fast.cpp @@ -1269,6 +1269,12 @@ std::vector spec_decode_verify( } auto s = to_stream(s_); + // B = 0: nothing to verify. Return empty results without the argmax (which + // rejects a zero-size input) or the fused kernel. + if (B == 0) { + return {zeros(Shape{0}, int32, s), zeros(Shape{0, K + 1}, int32, s)}; + } + // Reuse mlx's fast argmax on-device (result stays on the GPU); the fused // primitive does only the tiny per-row prefix-match. Primitive inputs are // int32 token arrays, so the kernel needs no large-V reduction. diff --git a/python/tests/test_fast.py b/python/tests/test_fast.py index e93796b6e0..bf5f1cbe33 100644 --- a/python/tests/test_fast.py +++ b/python/tests/test_fast.py @@ -1457,6 +1457,23 @@ def test_spec_decode_verify_empty_draft(self): self.assertEqual(n.tolist(), [0] * B) self.assertEqual(c.tolist(), [[t] for t in bonus]) + def test_spec_decode_verify_empty_batch(self): + # B = 0: empty, correctly shaped int32 results on every backend. + for K in (3, 0): + draft = mx.zeros((0, K), dtype=mx.int32) + lg = mx.zeros((0, K + 1, 8)) + results = {} + for path, stream in (("default", None), ("cpu-fallback", mx.cpu)): + with self.subTest(K=K, path=path): + kw = {"stream": stream} if stream is not None else {} + n, c = mx.fast.spec_decode_verify(draft, lg, **kw) + mx.eval(n, c) + self.assertEqual((n.shape, n.dtype), ((0,), mx.int32)) + self.assertEqual((c.shape, c.dtype), ((0, K + 1), mx.int32)) + results[path] = (n.shape, n.dtype, c.shape, c.dtype) + with self.subTest(K=K, compare="default vs cpu"): + self.assertEqual(results["default"], results["cpu-fallback"]) + def test_spec_decode_verify_strided_inputs(self): # Drafts and logits that are views, not dense rows: the fused kernel # must see the logical values, like the CPU composition does. From 438e420d1ef8109b68f1cf4416412dbf253ebaf4 Mon Sep 17 00:00:00 2001 From: Jeethu Rao Date: Sun, 11 Oct 2026 14:34:46 +0100 Subject: [PATCH 16/34] fast: spec_decode_verify takes the composition for K = 0 on every backend (#282) Found while running the R17 B = 0 check under the Metal validation layer (fix round 1; pre-existing in the fork's 28efbe632). With K = 0 and B > 0 the fused kernel was dispatched with an empty draft, which has no buffer, so index 0 was left unbound. Without validation the kernel never reads it (K = 0) and the result was right, but MTL_DEBUG_LAYER=1 MTL_SHADER_VALIDATION=1 aborts: validateComputeFunctionArguments:1043: failed assertion `Compute Function(spec_decode_verify): missing Buffer binding at index 0 for draft[0].' K = 0 now takes the op composition on every backend (n_accepted = 0, committed = the argmax bonus token, the K = 0 result from 52d3ef54b), so the kernel is never bound to an empty draft. K > 0 is unchanged. Test: test_spec_decode_verify_empty_draft (from 52d3ef54b), run under the validation layer. RED (before): pytest -k empty_draft under MTL_DEBUG_LAYER=1 MTL_SHADER_VALIDATION=1 -> rc 134, the assertion above. GREEN: the whole -k spec suite under the validation layer -> 4 passed, 96 subtests (rc 0); without it -> 4 passed, 96 subtests. --- mlx/fast.cpp | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/mlx/fast.cpp b/mlx/fast.cpp index 7836624b52..aff0277798 100644 --- a/mlx/fast.cpp +++ b/mlx/fast.cpp @@ -1321,7 +1321,9 @@ std::vector spec_decode_verify( }; std::vector prim_inputs = {dft, tgt}; - if (SpecDecodeVerify::use_fallback(s)) { + // K = 0 also takes the composition: an empty draft has no buffer to bind + // at the kernel's index 0 (the Metal API validation layer rejects that). + if (K == 0 || SpecDecodeVerify::use_fallback(s)) { return fallback(prim_inputs); } return array::make_arrays( From 329b79a07b4d605fb15e31969dfda2f1d230a4aa Mon Sep 17 00:00:00 2001 From: Jeethu Rao Date: Sun, 11 Oct 2026 15:20:25 +0100 Subject: [PATCH 17/34] metal: condense the gen-17 NAX gate comment (#282) Review feedback on numen-tech/mlx#10 (AGENTS.md code standards: 1-2 line comments). The measurements behind the gate are in 0bf8503e0 and the source commit 73db231d. --- mlx/backend/metal/device.cpp | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/mlx/backend/metal/device.cpp b/mlx/backend/metal/device.cpp index 93f3fe228c..5428c8be7b 100644 --- a/mlx/backend/metal/device.cpp +++ b/mlx/backend/metal/device.cpp @@ -1013,14 +1013,8 @@ bool is_nax_available() { } auto& d = metal::device(mlx::core::Device::gpu); auto gen = d.get_architecture_gen(); - // gen-17 desktop parts (M5-class, applegpu_g17*) take the nax path under - // the old `gen >= 17` gate but compute WRONG results in the nax steel-gemm - // and qmm_t kernels (measured 2026-07-05: fp16 GEMM max|err| ~4 vs fp32 for - // every shape routed to steel_gemm_fused_nax; quantized qmm_t ~400 abs err; - // non-nax fallback bit-matches stock mlx). Require gen >= 18 on every - // device class until the g17 path is fixed and correctness-gated. Every - // NAX dispatch site (matmul, quantized, sdpa, gated delta) asks this - // predicate, so this one gate keeps gen 17 off all of them. + // NAX GEMM and qmm_t kernels give wrong results on gen 17 (M5-class). + // All NAX dispatch sites use this function. can_use_nax &= gen >= 18; return can_use_nax; }; From fc4c06aae03374f5051cb60613a6be5c26877a4a Mon Sep 17 00:00:00 2001 From: Jeethu Rao Date: Sun, 11 Oct 2026 15:20:25 +0100 Subject: [PATCH 18/34] fast: spec_decode_verify rejects non-integer draft tokens (#282) Review feedback on numen-tech/mlx#10: astype(int32) silently truncated floating drafts (1.9 -> 1) and accepted bool. Integer widths are still cast to int32. test_spec_decode_verify_rejects_non_integer_drafts: RED 8 failed (float32/float16/bfloat16/bool x default/cpu), GREEN after the fix. --- mlx/fast.cpp | 6 ++++++ python/src/fast.cpp | 2 +- python/tests/test_fast.py | 15 +++++++++++++++ 3 files changed, 22 insertions(+), 1 deletion(-) diff --git a/mlx/fast.cpp b/mlx/fast.cpp index aff0277798..6e46b86e8b 100644 --- a/mlx/fast.cpp +++ b/mlx/fast.cpp @@ -1257,6 +1257,12 @@ std::vector spec_decode_verify( throw std::invalid_argument( "[spec_decode_verify] target_logits must be 3D [B, K+1, V]."); } + if (!issubdtype(draft_tokens.dtype(), integer)) { + std::ostringstream msg; + msg << "[spec_decode_verify] draft_tokens must be integer token ids but got " + << draft_tokens.dtype() << "."; + throw std::invalid_argument(msg.str()); + } int B = draft_tokens.shape(0); int K = draft_tokens.shape(1); if (target_logits.shape(0) != B) { diff --git a/python/src/fast.cpp b/python/src/fast.cpp index 487bc2012f..a933d1c5a2 100644 --- a/python/src/fast.cpp +++ b/python/src/fast.cpp @@ -134,7 +134,7 @@ void init_fast(nb::module_& parent_module) { ``[last] + draft``, returns the greedy accept result. Args: - draft_tokens (array): int32 ``[B, K]`` drafted tokens. + draft_tokens (array): ``[B, K]`` drafted token ids of an integer type. target_logits (array): ``[B, K+1, V]`` target logits over the ``[last] + draft`` positions. diff --git a/python/tests/test_fast.py b/python/tests/test_fast.py index bf5f1cbe33..df89117f99 100644 --- a/python/tests/test_fast.py +++ b/python/tests/test_fast.py @@ -1474,6 +1474,21 @@ def test_spec_decode_verify_empty_batch(self): with self.subTest(K=K, compare="default vs cpu"): self.assertEqual(results["default"], results["cpu-fallback"]) + def test_spec_decode_verify_rejects_non_integer_drafts(self): + lg = mx.zeros((1, 3, 8)) + for dtype in (mx.float32, mx.float16, mx.bfloat16, mx.bool_): + draft = mx.array([[1.9, 0.0]]).astype(dtype) + for path, stream in (("default", None), ("cpu-fallback", mx.cpu)): + with self.subTest(dtype=dtype, path=path): + kw = {"stream": stream} if stream is not None else {} + with self.assertRaises(ValueError): + mx.fast.spec_decode_verify(draft, lg, **kw) + for dtype in (mx.int64, mx.uint32, mx.int16): + with self.subTest(dtype=dtype): + draft = mx.array([[0, 0]], dtype=dtype) + n, c = mx.fast.spec_decode_verify(draft, lg) + self.assertEqual(n.tolist(), [2]) + def test_spec_decode_verify_strided_inputs(self): # Drafts and logits that are views, not dense rows: the fused kernel # must see the logical values, like the CPU composition does. From 44a6aa981b8553d934ca6c194ad1e13b1cc0bf16 Mon Sep 17 00:00:00 2001 From: Jeethu Rao Date: Mon, 12 Oct 2026 08:50:26 +0100 Subject: [PATCH 19/34] export: serialize SpecDecodeVerify (#282) Review feedback on numen-tech/mlx#10: on Metal with B > 0 and K > 0, mx.export_function of a function calling mx.fast.spec_decode_verify threw "Unable to serialize primitive SpecDecodeVerify" because the PrimitiveFactory did not list it. Registered like the other fast primitives: state() already puts nullptr in the fallback slot, so the imported primitive has no fallback, as RMSNorm, RoPE and ScaledDotProductAttention do after import. The fallback is not needed there: eval_gpu never calls it, and the CPU, K = 0 and B = 0 paths are chosen when the op is built, so they export as plain op compositions. test_export_import_spec_decode_verify (full accept, reject at 0, batch of 2, K = 0, B = 0; imported results equal the eager call): RED 3 failed (the three B > 0, K > 0 rows), GREEN after the fix, also under MTL_DEBUG_LAYER=1 MTL_SHADER_VALIDATION=1. --- mlx/export.cpp | 1 + python/tests/test_export_import.py | 41 ++++++++++++++++++++++++++++++ 2 files changed, 42 insertions(+) diff --git a/mlx/export.cpp b/mlx/export.cpp index c10934d1bd..af88b25dd0 100644 --- a/mlx/export.cpp +++ b/mlx/export.cpp @@ -465,6 +465,7 @@ struct PrimitiveFactory { SERIALIZE_PRIMITIVE(LayerNormVJP), SERIALIZE_PRIMITIVE(RoPE), SERIALIZE_PRIMITIVE(ScaledDotProductAttention), + SERIALIZE_PRIMITIVE(SpecDecodeVerify), SERIALIZE_PRIMITIVE(CustomKernel)}; std::unordered_map name_remap; std::unordered_map stream_map; diff --git a/python/tests/test_export_import.py b/python/tests/test_export_import.py index b8de6ce487..7ddbd702a9 100644 --- a/python/tests/test_export_import.py +++ b/python/tests/test_export_import.py @@ -749,6 +749,47 @@ def fun(x, y, z): imported = mx.import_function(path) self.assertTrue(mx.array_equal(imported(x, y, z)[0], fun(x, y, z))) + def test_export_import_spec_decode_verify(self): + def fun(draft, logits): + return mx.fast.spec_decode_verify(draft, logits) + + def onehot(tokens, V=8): + lg = [[[-10.0] * V for _ in row] for row in tokens] + for b, row in enumerate(tokens): + for j, t in enumerate(row): + lg[b][j][t] = 10.0 + return mx.array(lg) + + # On Metal, B > 0 and K > 0 run the fused primitive; K = 0 and B = 0 + # take the op composition. + cases = [ + ("full-accept", [[5, 6, 7]], [[5, 6, 7, 1]], [3]), + ("reject@0", [[5, 6, 7]], [[2, 6, 7, 1]], [0]), + ("batch2", [[1, 2, 3], [1, 2, 3]], [[1, 4, 3, 5], [1, 2, 3, 6]], [1, 3]), + ("K=0", [[], []], [[5], [3]], [0, 0]), + ] + path = os.path.join(self.test_dir, "spec_decode_verify.mlxfn") + for name, draft, target, expected in cases: + with self.subTest(case=name): + K = len(target[0]) - 1 + draft = mx.array(draft, dtype=mx.int32).reshape(len(target), K) + logits = onehot(target) + mx.export_function(path, fun, draft, logits) + imported = mx.import_function(path) + n, committed = imported(draft, logits) + n_ref, committed_ref = fun(draft, logits) + self.assertEqual(n.tolist(), expected) + self.assertTrue(mx.array_equal(n, n_ref)) + self.assertTrue(mx.array_equal(committed, committed_ref)) + + with self.subTest(case="B=0"): + draft = mx.zeros((0, 3), dtype=mx.int32) + logits = mx.zeros((0, 4, 8)) + mx.export_function(path, fun, draft, logits) + n, committed = mx.import_function(path)(draft, logits) + self.assertEqual((n.shape, n.dtype), ((0,), mx.int32)) + self.assertEqual((committed.shape, committed.dtype), ((0, 4), mx.int32)) + def test_export_matmul_shapeless_mid_dim(self): path = os.path.join(self.test_dir, "matmul_shapeless.mlxfn") From c28faebfba3df4d28d1541b9eeb954ebb8edf96c Mon Sep 17 00:00:00 2001 From: Jeethu Rao Date: Mon, 12 Oct 2026 08:50:26 +0100 Subject: [PATCH 20/34] metal: condense the qmv partial-last-block comment (#282) Review feedback on numen-tech/mlx#10 (AGENTS.md code standards: 1-2 line comments). Comment-only. The block is reached only when the host's K gate (qmv_fast_k_alignment) is finer than block_size: affine 1/4-bit take K % (block_size / 2) == 0. Masking the out-of-row lanes instead (zeroed x_thread, scale and bias) was rejected: qdot would still run, and its weight reads land past the row end and, for the last output rows, past the end of the weight buffer. A K gate that did not keep in_vec_size a multiple of values_per_thread would need the stricter predicate aligned_end + (simd_lid + 1) * values_per_thread <= in_vec_size. The comment came in with 7c1362e04 (numen-tech/gemma4-qat#179, #282). --- mlx/backend/metal/kernels/quantized.h | 15 ++------------- 1 file changed, 2 insertions(+), 13 deletions(-) diff --git a/mlx/backend/metal/kernels/quantized.h b/mlx/backend/metal/kernels/quantized.h index 82d191f763..d5d70a45c5 100644 --- a/mlx/backend/metal/kernels/quantized.h +++ b/mlx/backend/metal/kernels/quantized.h @@ -980,19 +980,8 @@ METAL_FUNC void qmv_fast_impl( x += block_size; } - // Partial last block (in_vec_size % block_size != 0), reached only when the - // host's K gate (qmv_fast_k_alignment) is finer than block_size: affine - // 1/4-bit take K % (block_size / 2) == 0. A lane whose first value lies - // past the end of the row reads nothing (x, weights, scales, biases) and - // skips the row loop: its result[row] stays 0 and it still joins the - // simd_sum below. Masking such lanes instead (zeroed x_thread, scale and - // bias) would still run qdot, whose weight reads land past the end of the - // row and, for the last output rows, past the end of the weight buffer. A - // lane that starts inside the row also ends inside it because the K gate - // keeps in_vec_size a multiple of values_per_thread; a gate that did not - // would need the stricter predicate - // aligned_end + (simd_lid + 1) * values_per_thread <= in_vec_size. - // (numen-tech/gemma4-qat#179, #282) + // Partial last block: lanes starting past the row end read nothing and add 0. + // The K gate keeps in_vec_size a multiple of values_per_thread. if (aligned_end < in_vec_size) { bool in_bounds = (aligned_end + static_cast(simd_lid) * values_per_thread) < From bee9e9d2bac91bf505ee6ad882b11b5d47fe7e63 Mon Sep 17 00:00:00 2001 From: Jeethu Rao Date: Mon, 12 Oct 2026 19:29:32 +0100 Subject: [PATCH 21/34] metal: default D512 vector attention to GPU generation 17+ (#291) On an M3 Max (applegpu_g15s) the D512 sdpa_vector_2pass path from #4459 decodes Gemma 4 E2B (8 query heads, 1 KV head) 3-4% slower than the unfused fallback beyond 1024 keys; forcing the fallback with MLX_SDPA_D512_MIN_KL recovers it. Upstream measured gains on M5 only. Keep the default on generation 17+; elsewhere use the vector kernel only when MLX_SDPA_D512_MIN_KL is set. --- .../metal/scaled_dot_product_attention.cpp | 18 +++++++++++++-- python/tests/test_fast_sdpa.py | 23 +++++++++++++++++++ 2 files changed, 39 insertions(+), 2 deletions(-) diff --git a/mlx/backend/metal/scaled_dot_product_attention.cpp b/mlx/backend/metal/scaled_dot_product_attention.cpp index 5b157c6c5c..b61d0be71d 100644 --- a/mlx/backend/metal/scaled_dot_product_attention.cpp +++ b/mlx/backend/metal/scaled_dot_product_attention.cpp @@ -1018,9 +1018,16 @@ std::tuple has_fused_kernel( } // Head dim 512 uses the generic vector kernel, reading K/V is its main // cost. By default, use it only for one query, GQA factor 8, and no array - // mask. + // mask, and only on GPU generation 17 or newer: on an M3 Max (generation + // 15) it decodes Gemma 4 E2B (8 query heads, 1 KV head) 3-4% slower than + // the fallback beyond 1024 keys. Setting MLX_SDPA_D512_MIN_KL enables it + // on every generation. if (query_head_dim == 512) { - int min_key_sequence_length = env::get_var("MLX_SDPA_D512_MIN_KL", 1024); + int env_min_key_sequence_length = + env::get_var("MLX_SDPA_D512_MIN_KL", -1); + bool use_default = (env_min_key_sequence_length < 0); + int min_key_sequence_length = + use_default ? 1024 : env_min_key_sequence_length; bool always = (min_key_sequence_length == 0); if (!always && query_sequence_length != 1) { msg << "the vector attention kernel for head dim 512 defaults to " @@ -1045,6 +1052,13 @@ std::tuple has_fused_kernel( << key_sequence_length << "."; return {false, msg.str()}; } + int arch_gen = metal::device(s.device).get_architecture_gen(); + if (use_default && arch_gen < 17) { + msg << "the vector attention kernel for head dim 512 defaults to GPU " + << "generation 17 or newer unless MLX_SDPA_D512_MIN_KL is set; got " + << "generation " << arch_gen << "."; + return {false, msg.str()}; + } } } return {true, ""}; diff --git a/python/tests/test_fast_sdpa.py b/python/tests/test_fast_sdpa.py index b2d7e9121e..eeb33fdfbc 100644 --- a/python/tests/test_fast_sdpa.py +++ b/python/tests/test_fast_sdpa.py @@ -1330,6 +1330,29 @@ def make_qkv(qL, kL, D, qH=8, kH=8, dtype=mx.float16): q, k, v, scale=512**-0.5, force_fused=True ) + # Head dim 512 vector attention defaults to GPU generation 17 or newer; + # setting MLX_SDPA_D512_MIN_KL enables it on every generation. + arch = mx.device_info()["architecture"] + arch_gen = int(arch[-3:-1]) if arch[-3:-1].isdigit() else 0 + q, k, v = make_qkv(1, 2048, 512, qH=8, kH=1) + ref = mlx_ref_attn(q, k, v, scale=512**-0.5) + with mlx_tests.scoped_env(MLX_SDPA_D512_MIN_KL=None): + if arch_gen < 17: + with self.assertRaisesRegex(ValueError, "GPU generation 17"): + mx.fast.scaled_dot_product_attention( + q, k, v, scale=512**-0.5, force_fused=True + ) + else: + out = mx.fast.scaled_dot_product_attention( + q, k, v, scale=512**-0.5, force_fused=True + ) + self.assertTrue(mx.allclose(ref, out, atol=1e-3, rtol=1e-3)) + with mlx_tests.scoped_env(MLX_SDPA_D512_MIN_KL="1024"): + out = mx.fast.scaled_dot_product_attention( + q, k, v, scale=512**-0.5, force_fused=True + ) + self.assertTrue(mx.allclose(ref, out, atol=1e-3, rtol=1e-3)) + # No CPU fused kernel. with mx.stream(mx.cpu): q, k, v = make_qkv(8, 128, 8) From ee0b1bcfc1fd12b82846c6b9a684637449bec63b Mon Sep 17 00:00:00 2001 From: Jeethu Rao Date: Mon, 12 Oct 2026 19:35:08 +0100 Subject: [PATCH 22/34] metal: condense the D512 vector-attention gate comment (#291) --- mlx/backend/metal/scaled_dot_product_attention.cpp | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/mlx/backend/metal/scaled_dot_product_attention.cpp b/mlx/backend/metal/scaled_dot_product_attention.cpp index b61d0be71d..12f9f37f4a 100644 --- a/mlx/backend/metal/scaled_dot_product_attention.cpp +++ b/mlx/backend/metal/scaled_dot_product_attention.cpp @@ -1017,11 +1017,8 @@ std::tuple has_fused_kernel( return {false, msg.str()}; } // Head dim 512 uses the generic vector kernel, reading K/V is its main - // cost. By default, use it only for one query, GQA factor 8, and no array - // mask, and only on GPU generation 17 or newer: on an M3 Max (generation - // 15) it decodes Gemma 4 E2B (8 query heads, 1 KV head) 3-4% slower than - // the fallback beyond 1024 keys. Setting MLX_SDPA_D512_MIN_KL enables it - // on every generation. + // cost. By default, use it only for one query, GQA factor 8, no array mask + // and GPU generation 17+; setting MLX_SDPA_D512_MIN_KL lifts the gen gate. if (query_head_dim == 512) { int env_min_key_sequence_length = env::get_var("MLX_SDPA_D512_MIN_KL", -1); From 33ef381f3baab1b6d3be7c79a734ab59e387e24a Mon Sep 17 00:00:00 2001 From: Jeethu Rao Date: Thu, 15 Oct 2026 01:42:55 +0100 Subject: [PATCH 23/34] metal: qmv_quad reads no scale/bias for rows past N on a ragged last tile (gemma4-qat#302) qmv_quad runs ceil(N / 64) tiles with no N % 64 restriction. On a ragged last tile, the quad rows past out_vec_size loaded their scale (and, for an explicit bias, their bias) before the row predicate, reading past the end of the scales/biases buffers. The qdot and the store were already guarded, so the output was never wrong. Move both loads inside the existing predicate; the pointer arithmetic stays outside (forming a pointer reads nothing). v0.32.4's qmv_quad_impl has the same unguarded structure, so this is not an upstream change taken over. test_qmv_quad_ragged_n covers N = 65 (one valid row in tile 2) and N = 72 (eight), at bits 1/2/4/8, gs 32/64, K 64/128 (1-bit K = 128 only), with an explicit and an implied (0-d factor) bias, against a CPU dequantize reference. affine_sym never reaches qmv_quad. Validation: with MTL_DEBUG_LAYER=1 MTL_SHADER_VALIDATION=1 MTL_SHADER_VALIDATION_GLOBAL_MEMORY=1, the unfixed kernel reported nothing, at those shapes nor at fp32 K = 128 gs = 32 N = 1025. A deliberately out-of-bounds custom kernel (1M elements past a 16-element input) was not reported either, so GPU validation does not see out-of-bounds device loads in this setup; the fix rests on the bounds argument above. The test passes before and after the fix (an allclose cannot see a discarded read). --- mlx/backend/metal/kernels/quantized.h | 7 +++-- python/tests/test_quantized.py | 40 +++++++++++++++++++++++++++ 2 files changed, 45 insertions(+), 2 deletions(-) diff --git a/mlx/backend/metal/kernels/quantized.h b/mlx/backend/metal/kernels/quantized.h index d5d70a45c5..d7857753d2 100644 --- a/mlx/backend/metal/kernels/quantized.h +++ b/mlx/backend/metal/kernels/quantized.h @@ -878,9 +878,12 @@ METAL_FUNC void qmv_quad_impl( const device T* sl = scales + row * in_vec_size_g * quads_per_simd; const device T* bl = biases + row * in_vec_size_g * quads_per_simd; - U s = sl[0]; - U b = affine_bias(s, f, sl[0], bl); + // Rows past out_vec_size exist only on a ragged last tile (N % 64 != 0): + // read no scale or bias for them (gemma4-qat#302; the store below was + // already guarded). if (row * quads_per_simd + out_row < out_vec_size) { + U s = sl[0]; + U b = affine_bias(s, f, sl[0], bl); result[row] += qdot(wl, x_thread, s, b, sum); } } diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index 9a67957eb2..df5f1d237a 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -904,6 +904,46 @@ def test_qmv_fast_tail_non_aligned_k(self): y_cpu = mx.quantized_matmul(x, wq, s, b, True, gs, bits, stream=mx.cpu) self.assertTrue(mx.allclose(y, y_cpu, atol=1e-2, rtol=1e-2)) + @unittest.skipIf(not mx.metal.is_available(), "requires Metal") + def test_qmv_quad_ragged_n(self): + # qmv_quad (K = 128, or K = 64 at bits >= 2; power-of-two bits) runs + # ceil(N / 64) tiles with no N % 64 restriction. On a ragged last tile, + # quad rows past N must read no scale or bias (gemma4-qat#302). N = 65 + # leaves one valid row in tile 2, N = 72 leaves eight. Explicit and + # implied (0-d factor) biases are both covered; affine_sym never reaches + # qmv_quad. + key = mx.random.key(7) + cases = [ + (bits, gs, K) for bits in (1, 2, 4, 8) for gs in (32, 64) for K in (64, 128) + ] + for bits, gs, K in cases: + if K == 64 and bits == 1: + continue # 1-bit qmv_quad needs K == 128 + for N in (65, 72): + for implied in (False, True): + with self.subTest(bits=bits, gs=gs, K=K, N=N, implied=implied): + k1, k2, key = mx.random.split(key, 3) + w = mx.random.normal((N, K), key=k1).astype(mx.float16) + x = mx.random.normal((1, K), key=k2).astype(mx.float16) + wq, s, b = mx.quantize(w, group_size=gs, bits=bits) + if implied: + f = -(2.0 ** (bits - 1)) if bits > 1 else -0.5 + b = mx.array(f, mx.float16) + full = s * b + else: + full = b + y = mx.quantized_matmul( + x, wq, s, b, True, gs, bits, stream=mx.gpu + ) + w_ref = mx.dequantize(wq, s, full, gs, bits, stream=mx.cpu) + y_ref = x.astype(mx.float32) @ w_ref.astype(mx.float32).T + self.assertEqual(y.shape, (1, N)) + self.assertTrue( + mx.allclose( + y.astype(mx.float32), y_ref, atol=2e-2, rtol=2e-2 + ) + ) + @unittest.skipIf(not mx.metal.is_available(), "requires Metal") def test_fp_qmv_fast_unaligned_n(self): for n, k, mode, dtype, batched in product( From 4c4d61c5c2153eee0c579562e288fa18b7cc94a2 Mon Sep 17 00:00:00 2001 From: Jeethu Rao Date: Thu, 15 Oct 2026 09:25:50 +0100 Subject: [PATCH 24/34] fast: SpecDecodeVerify rebuilds its fallback after import (#282) Review feedback on numen-tech/mlx#12: an imported SpecDecodeVerify had an empty fallback (state() put nullptr in its slot), so mx.vmap over an imported function threw std::bad_function_call from Custom::vmap. The composition now takes B and K from its input shapes, and the primitive builds its fallback from the stream alone, so import needs no state. The serialized form does not change: nullptr wrote no bytes. Also condenses the class comment to one line. test_export_import_spec_decode_verify vmap case (vmap of the imported function equals vmap of the eager call): RED std::bad_function_call, GREEN after the fix. test_fast + test_quantized + test_vmap + test_export_import 167 passed, 2 skipped; C++ JIT and metallib 311/311. --- mlx/fast.cpp | 86 ++++++++++++++++-------------- mlx/fast_primitives.h | 14 +---- python/tests/test_export_import.py | 12 +++++ 3 files changed, 59 insertions(+), 53 deletions(-) diff --git a/mlx/fast.cpp b/mlx/fast.cpp index 6e46b86e8b..85f0326a0b 100644 --- a/mlx/fast.cpp +++ b/mlx/fast.cpp @@ -1245,6 +1245,44 @@ bool ConvertFP8::is_equivalent(const Primitive& other) const { return to_fp8_ == a_other.to_fp8_; } +// The verify as a plain op composition on (draft_tokens [B, K], target_tokens +// [B, K+1]). It is the CPU path, the fallback for transforms, and the oracle. +static std::vector spec_decode_verify_composition( + const std::vector& inputs, + Stream s) { + int B = inputs[0].shape(0); + int K = inputs[0].shape(1); + auto d = inputs[0]; // [B, K] int32 + auto t = inputs[1]; // [B, K+1] int32 + if (K == 0) { + // Nothing to accept; committed is the bonus token t[:, 0], as the + // fused kernel writes it. (min over the empty axis below would throw.) + return {zeros(Shape{B}, int32, s), t}; + } + auto t_pref = slice(t, Shape{0, 0}, Shape{B, K}, s); + auto mism = not_equal(d, t_pref, s); + auto j = broadcast_to( + reshape(arange(K, int32, s), Shape{1, K}, s), Shape{B, K}, s); + auto n_acc = + min(where(mism, j, full(Shape{B, K}, K, int32, s), s), + /*axis=*/1, + /*keepdims=*/false, + s); // [B] + auto n_acc2 = reshape(n_acc, Shape{B, 1}, s); + auto corrected = take_along_axis(t, n_acc2, /*axis=*/1, s); // [B, 1] + auto j1 = broadcast_to( + reshape(arange(K + 1, int32, s), Shape{1, K + 1}, s), Shape{B, K + 1}, s); + auto nacc_b = broadcast_to(n_acc2, Shape{B, K + 1}, s); + auto d_ext = concatenate({d, zeros(Shape{B, 1}, int32, s)}, /*axis=*/1, s); + auto corr_b = broadcast_to(corrected, Shape{B, K + 1}, s); + auto committed = where( + less(j1, nacc_b, s), + d_ext, + where(equal(j1, nacc_b, s), corr_b, zeros(Shape{B, K + 1}, int32, s), s), + s); + return {n_acc, committed}; +} + std::vector spec_decode_verify( const array& draft_tokens, const array& target_logits, @@ -1287,58 +1325,24 @@ std::vector spec_decode_verify( auto dft = astype(draft_tokens, int32, s); auto tgt = astype(argmax(target_logits, -1, false, s), int32, s); // [B, K+1] - // Fallback (CPU path + autodiff): the same verify as a plain op composition, - // operating on (draft_tokens, target_tokens). This is the correctness oracle - // the Metal kernel is checked against. - auto fallback = - [B, K, s](const std::vector& inputs) -> std::vector { - auto d = inputs[0]; // [B, K] int32 - auto t = inputs[1]; // [B, K+1] int32 - if (K == 0) { - // Nothing to accept; committed is the bonus token t[:, 0], as the - // fused kernel writes it. (min over the empty axis below would throw.) - return {zeros(Shape{B}, int32, s), t}; - } - auto t_pref = slice(t, Shape{0, 0}, Shape{B, K}, s); - auto mism = not_equal(d, t_pref, s); - auto j = broadcast_to( - reshape(arange(K, int32, s), Shape{1, K}, s), Shape{B, K}, s); - auto n_acc = - min(where(mism, j, full(Shape{B, K}, K, int32, s), s), - /*axis=*/1, - /*keepdims=*/false, - s); // [B] - auto n_acc2 = reshape(n_acc, Shape{B, 1}, s); - auto corrected = take_along_axis(t, n_acc2, /*axis=*/1, s); // [B, 1] - auto j1 = broadcast_to( - reshape(arange(K + 1, int32, s), Shape{1, K + 1}, s), - Shape{B, K + 1}, - s); - auto nacc_b = broadcast_to(n_acc2, Shape{B, K + 1}, s); - auto d_ext = concatenate({d, zeros(Shape{B, 1}, int32, s)}, /*axis=*/1, s); - auto corr_b = broadcast_to(corrected, Shape{B, K + 1}, s); - auto committed = where( - less(j1, nacc_b, s), - d_ext, - where( - equal(j1, nacc_b, s), corr_b, zeros(Shape{B, K + 1}, int32, s), s), - s); - return {n_acc, committed}; - }; - std::vector prim_inputs = {dft, tgt}; // K = 0 also takes the composition: an empty draft has no buffer to bind // at the kernel's index 0 (the Metal API validation layer rejects that). if (K == 0 || SpecDecodeVerify::use_fallback(s)) { - return fallback(prim_inputs); + return spec_decode_verify_composition(prim_inputs, s); } return array::make_arrays( {Shape{B}, Shape{B, K + 1}}, {int32, int32}, - std::make_shared(s, fallback), + std::make_shared(s), prim_inputs); } +SpecDecodeVerify::SpecDecodeVerify(Stream stream) + : Custom(stream, [stream](std::vector inputs) { + return spec_decode_verify_composition(inputs, stream); + }) {} + bool SpecDecodeVerify::is_equivalent(const Primitive&) const { return true; // no parameters } diff --git a/mlx/fast_primitives.h b/mlx/fast_primitives.h index 7636ac1a87..586d5099d4 100644 --- a/mlx/fast_primitives.h +++ b/mlx/fast_primitives.h @@ -610,17 +610,10 @@ class CustomKernel : public Primitive { CompileOptions::Data compile_options_; }; -// Fused greedy speculative-decoding verify. Inputs: draft_tokens int32 [B, K] -// and target_tokens int32 [B, K+1] (the target argmax). Outputs: n_accepted -// [B] int32 and committed [B, K+1] int32. The op-level fallback (a plain -// composition) runs on CPU and for autodiff; the Metal eval_gpu runs a trivial -// per-row prefix-match kernel. +// Greedy speculative-decoding verify on int32 [B, K] and [B, K+1] tokens. class SpecDecodeVerify : public Custom { public: - SpecDecodeVerify( - Stream stream, - std::function(std::vector)> fallback) - : Custom(stream, std::move(fallback)) {} + explicit SpecDecodeVerify(Stream stream); static bool use_fallback(Stream stream); @@ -635,9 +628,6 @@ class SpecDecodeVerify : public Custom { DEFINE_NAME(SpecDecodeVerify); bool is_equivalent(const Primitive& other) const override; std::vector output_shapes(const std::vector& inputs) override; - auto state() const { - return std::make_tuple(nullptr); - } }; } // namespace mlx::core::fast diff --git a/python/tests/test_export_import.py b/python/tests/test_export_import.py index 7ddbd702a9..a0d5f654dc 100644 --- a/python/tests/test_export_import.py +++ b/python/tests/test_export_import.py @@ -782,6 +782,18 @@ def onehot(tokens, V=8): self.assertTrue(mx.array_equal(n, n_ref)) self.assertTrue(mx.array_equal(committed, committed_ref)) + with self.subTest(case="vmap"): + # Two independent [1, 3] verifies (accept 1, accept 3). + draft = mx.array([[[1, 2, 3]], [[1, 2, 3]]], dtype=mx.int32) + logits = mx.stack([onehot([[1, 4, 3, 5]]), onehot([[1, 2, 3, 6]])]) + mx.export_function(path, fun, draft[0], logits[0]) + imported = mx.import_function(path) + n, committed = mx.vmap(imported)(draft, logits) + n_ref, committed_ref = mx.vmap(fun)(draft, logits) + self.assertEqual(n.tolist(), [[1], [3]]) + self.assertTrue(mx.array_equal(n, n_ref)) + self.assertTrue(mx.array_equal(committed, committed_ref)) + with self.subTest(case="B=0"): draft = mx.zeros((0, 3), dtype=mx.int32) logits = mx.zeros((0, 4, 8)) From 1e1f993c2d4f2d08b12ddf8b0c43ba21d693a81a Mon Sep 17 00:00:00 2001 From: Jeethu Rao Date: Thu, 15 Oct 2026 10:10:48 +0100 Subject: [PATCH 25/34] fast: spec_decode_verify takes the primitive for every K and on CPU (#282) Review feedback on numen-tech/mlx#12: the CPU path and the K = 0 path built the op composition, whose shapes (Shape{B, K}, arange(K)) are fixed at trace time. Under mx.compile(shapeless=True) a K = 0 first trace then returned n_accepted = [0] for every later shape on CPU and Metal, and a K > 0 CPU trace threw "Slice cannot infer output shapes". SpecDecodeVerify now has a native eval_cpu and the Metal kernel takes K = 0 (target is bound at the empty draft's slot; the kernel reads no draft for K = 0), so every B > 0 call is the one primitive, whose output shapes come from its inputs. The composition stays as the transform fallback and for CUDA GPU streams. B = 0 keeps its early return (argmax rejects a zero-size input). test_spec_decode_verify_shapeless_compile (cpu and gpu, K > 0 and K = 0 traced first, 5 shapes each): RED 13 failed, GREEN after the fix, also under MTL_DEBUG_LAYER=1 MTL_SHADER_VALIDATION=1. test_fast + test_quantized + test_vmap + test_export_import + test_compile 239 passed, 2 skipped; C++ JIT and metallib 311/311. --- mlx/backend/cpu/CMakeLists.txt | 1 + mlx/backend/cpu/spec_decode.cpp | 55 +++++++++++++++++++++++++++++ mlx/backend/cuda/primitives.cpp | 5 ++- mlx/backend/metal/spec_decode.cpp | 6 ++-- mlx/backend/no_cpu/primitives.cpp | 1 + mlx/backend/no_gpu/primitives.cpp | 2 +- mlx/fast.cpp | 7 ++-- mlx/fast_primitives.h | 5 +-- python/tests/test_export_import.py | 4 +-- python/tests/test_fast.py | 56 ++++++++++++++++++++++++------ 10 files changed, 116 insertions(+), 26 deletions(-) create mode 100644 mlx/backend/cpu/spec_decode.cpp diff --git a/mlx/backend/cpu/CMakeLists.txt b/mlx/backend/cpu/CMakeLists.txt index 6192a189a4..3352755e89 100644 --- a/mlx/backend/cpu/CMakeLists.txt +++ b/mlx/backend/cpu/CMakeLists.txt @@ -64,6 +64,7 @@ target_sources( ${CMAKE_CURRENT_SOURCE_DIR}/softmax.cpp ${CMAKE_CURRENT_SOURCE_DIR}/logsumexp.cpp ${CMAKE_CURRENT_SOURCE_DIR}/sort.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/spec_decode.cpp ${CMAKE_CURRENT_SOURCE_DIR}/threefry.cpp ${CMAKE_CURRENT_SOURCE_DIR}/indexing.cpp ${CMAKE_CURRENT_SOURCE_DIR}/luf.cpp diff --git a/mlx/backend/cpu/spec_decode.cpp b/mlx/backend/cpu/spec_decode.cpp new file mode 100644 index 0000000000..89c24f9f73 --- /dev/null +++ b/mlx/backend/cpu/spec_decode.cpp @@ -0,0 +1,55 @@ +// Copyright © 2026 Apple Inc. + +#include "mlx/allocator.h" +#include "mlx/backend/cpu/encoder.h" +#include "mlx/fast_primitives.h" + +namespace mlx::core::fast { + +void SpecDecodeVerify::eval_cpu( + const std::vector& inputs, + std::vector& outputs) { + auto& draft = inputs[0]; // [B, K] int32 + auto& target = inputs[1]; // [B, K+1] int32 + auto& n_accepted = outputs[0]; // [B] int32 + auto& committed = outputs[1]; // [B, K+1] int32 + n_accepted.set_data(allocator::malloc(n_accepted.nbytes())); + committed.set_data(allocator::malloc(committed.nbytes())); + + auto& encoder = cpu::get_command_encoder(stream()); + encoder.set_input_array(draft); + encoder.set_input_array(target); + encoder.set_output_array(n_accepted); + encoder.set_output_array(committed); + encoder.dispatch([draft = array::unsafe_weak_copy(draft), + target = array::unsafe_weak_copy(target), + n_accepted = array::unsafe_weak_copy(n_accepted), + committed = array::unsafe_weak_copy(committed)]() mutable { + int B = draft.shape(0); + int K = draft.shape(1); + auto ds = draft.strides(); + auto ts = target.strides(); + auto* d = draft.data(); + auto* t = target.data(); + auto* n_out = n_accepted.data(); + auto* c_out = committed.data(); + for (int b = 0; b < B; ++b) { + auto dt = [&](int j) { return d[b * ds[0] + j * ds[1]]; }; + auto tt = [&](int j) { return t[b * ts[0] + j * ts[1]]; }; + int n = K; + for (int j = 0; j < K; ++j) { + if (dt(j) != tt(j)) { + n = j; + break; + } + } + n_out[b] = n; + int32_t* c = c_out + b * (K + 1); + for (int j = 0; j < K + 1; ++j) { + c[j] = (j < n) ? dt(j) : ((j == n) ? tt(n) : 0); + } + } + }); +} + +} // namespace mlx::core::fast diff --git a/mlx/backend/cuda/primitives.cpp b/mlx/backend/cuda/primitives.cpp index e91771a1c2..4121e4e5cb 100644 --- a/mlx/backend/cuda/primitives.cpp +++ b/mlx/backend/cuda/primitives.cpp @@ -49,7 +49,10 @@ NO_GPU_MULTI(SVD) NO_GPU(Inverse) NO_GPU_MULTI(Eig) NO_GPU_MULTI(Eigh) -NO_GPU_USE_FALLBACK(fast::SpecDecodeVerify) +bool fast::SpecDecodeVerify::use_fallback(Stream s) { + return s.device == Device::gpu; +} +NO_GPU_MULTI(fast::SpecDecodeVerify) namespace fast { NO_GPU_MULTI(GatedDeltaUpdate) diff --git a/mlx/backend/metal/spec_decode.cpp b/mlx/backend/metal/spec_decode.cpp index f1bf001510..8d38eaa144 100644 --- a/mlx/backend/metal/spec_decode.cpp +++ b/mlx/backend/metal/spec_decode.cpp @@ -27,8 +27,7 @@ array ensure_row_contiguous(const array& x, const Stream& s) { } // namespace bool SpecDecodeVerify::use_fallback(Stream s) { - // Fused kernel only on the GPU; CPU uses the op-level composition fallback. - return s.device == Device::cpu; + return false; } void SpecDecodeVerify::eval_gpu( @@ -51,7 +50,8 @@ void SpecDecodeVerify::eval_gpu( auto kernel = d.get_kernel("spec_decode_verify"); auto& compute_encoder = metal::get_command_encoder(s); compute_encoder.set_compute_pipeline_state(kernel); - compute_encoder.set_input_array(draft, 0); + // K = 0 has no draft buffer, and the kernel reads none: bind target there. + compute_encoder.set_input_array(K > 0 ? draft : target, 0); compute_encoder.set_input_array(target, 1); compute_encoder.set_output_array(n_accepted, 2); compute_encoder.set_output_array(committed, 3); diff --git a/mlx/backend/no_cpu/primitives.cpp b/mlx/backend/no_cpu/primitives.cpp index 2ad6461185..c91ee5effd 100644 --- a/mlx/backend/no_cpu/primitives.cpp +++ b/mlx/backend/no_cpu/primitives.cpp @@ -136,6 +136,7 @@ NO_CPU(View) namespace fast { NO_CPU_MULTI(Quantize) NO_CPU_MULTI(ConvertFP8) +NO_CPU_MULTI(SpecDecodeVerify) } // namespace fast namespace distributed { diff --git a/mlx/backend/no_gpu/primitives.cpp b/mlx/backend/no_gpu/primitives.cpp index af278813f3..1f86d116d9 100644 --- a/mlx/backend/no_gpu/primitives.cpp +++ b/mlx/backend/no_gpu/primitives.cpp @@ -73,7 +73,7 @@ bool fast::GatedDeltaUpdateVJP::use_fallback( } bool fast::SpecDecodeVerify::use_fallback(Stream s) { - return true; + return false; } NO_GPU_MULTI(fast::SpecDecodeVerify) diff --git a/mlx/fast.cpp b/mlx/fast.cpp index 85f0326a0b..163ceaab4f 100644 --- a/mlx/fast.cpp +++ b/mlx/fast.cpp @@ -1246,7 +1246,7 @@ bool ConvertFP8::is_equivalent(const Primitive& other) const { } // The verify as a plain op composition on (draft_tokens [B, K], target_tokens -// [B, K+1]). It is the CPU path, the fallback for transforms, and the oracle. +// [B, K+1]). It is the fallback for transforms and for backends with no kernel. static std::vector spec_decode_verify_composition( const std::vector& inputs, Stream s) { @@ -1326,9 +1326,8 @@ std::vector spec_decode_verify( auto tgt = astype(argmax(target_logits, -1, false, s), int32, s); // [B, K+1] std::vector prim_inputs = {dft, tgt}; - // K = 0 also takes the composition: an empty draft has no buffer to bind - // at the kernel's index 0 (the Metal API validation layer rejects that). - if (K == 0 || SpecDecodeVerify::use_fallback(s)) { + // Every K takes the primitive, so a shapeless compile does not fix K. + if (SpecDecodeVerify::use_fallback(s)) { return spec_decode_verify_composition(prim_inputs, s); } return array::make_arrays( diff --git a/mlx/fast_primitives.h b/mlx/fast_primitives.h index 586d5099d4..8a75db43b3 100644 --- a/mlx/fast_primitives.h +++ b/mlx/fast_primitives.h @@ -618,10 +618,7 @@ class SpecDecodeVerify : public Custom { static bool use_fallback(Stream stream); void eval_cpu(const std::vector& inputs, std::vector& outputs) - override { - throw std::runtime_error( - "[SpecDecodeVerify] no eval_cpu; CPU uses the op fallback."); - } + override; void eval_gpu(const std::vector& inputs, std::vector& outputs) override; diff --git a/python/tests/test_export_import.py b/python/tests/test_export_import.py index a0d5f654dc..ce0f83051b 100644 --- a/python/tests/test_export_import.py +++ b/python/tests/test_export_import.py @@ -760,8 +760,8 @@ def onehot(tokens, V=8): lg[b][j][t] = 10.0 return mx.array(lg) - # On Metal, B > 0 and K > 0 run the fused primitive; K = 0 and B = 0 - # take the op composition. + # B > 0 runs the primitive (K = 0 included); B = 0 returns empty + # arrays without it. cases = [ ("full-accept", [[5, 6, 7]], [[5, 6, 7, 1]], [3]), ("reject@0", [[5, 6, 7]], [[2, 6, 7, 1]], [0]), diff --git a/python/tests/test_fast.py b/python/tests/test_fast.py index df89117f99..5251f33203 100644 --- a/python/tests/test_fast.py +++ b/python/tests/test_fast.py @@ -1360,8 +1360,8 @@ def test_custom_metal_kernel_negative_template_int(self): def test_spec_decode_verify(self): # Validate mx.fast.spec_decode_verify against a pure-Python oracle that # mirrors the greedy verify loop. On a Metal build the default device - # exercises the fused kernel and stream=mx.cpu the op-composition - # fallback; both must agree with the oracle. + # exercises the fused kernel and stream=mx.cpu the CPU kernel; both + # must agree with the oracle. def oracle(draft, tgt_tokens): B, K = len(draft), len(draft[0]) n_acc, committed = [], [] @@ -1409,8 +1409,8 @@ def check(draft, tgt_tokens, V, dtype, stream): ("batch3-K1", [[7], [7], [7]], [[7, 1], [2, 3], [7, 9]]), ] # stream=None is the default device (fused kernel on Metal); mx.cpu is - # the op-composition fallback. - for path, stream in (("default", None), ("cpu-fallback", mx.cpu)): + # the CPU kernel. + for path, stream in (("default", None), ("cpu", mx.cpu)): for dtype in (mx.float32, mx.float16, mx.bfloat16): for name, draft, tgt in cases: with self.subTest(path=path, dtype=str(dtype), case=name): @@ -1440,14 +1440,14 @@ def check(draft, tgt_tokens, V, dtype, stream): def test_spec_decode_verify_empty_draft(self): # K = 0: nothing to accept, the bonus token is the argmax at the only - # position. The CPU composition must match the fused kernel. + # position. The CPU kernel must match the fused kernel. for B, bonus in ((1, [5]), (2, [5, 3])): lg = [[[-10.0] * 8] for _ in range(B)] for b, t in enumerate(bonus): lg[b][0][t] = 10.0 lg = mx.array(lg) draft = mx.zeros((B, 0), dtype=mx.int32) - for path, stream in (("default", None), ("cpu-fallback", mx.cpu)): + for path, stream in (("default", None), ("cpu", mx.cpu)): with self.subTest(B=B, path=path): kw = {"stream": stream} if stream is not None else {} n, c = mx.fast.spec_decode_verify(draft, lg, **kw) @@ -1463,7 +1463,7 @@ def test_spec_decode_verify_empty_batch(self): draft = mx.zeros((0, K), dtype=mx.int32) lg = mx.zeros((0, K + 1, 8)) results = {} - for path, stream in (("default", None), ("cpu-fallback", mx.cpu)): + for path, stream in (("default", None), ("cpu", mx.cpu)): with self.subTest(K=K, path=path): kw = {"stream": stream} if stream is not None else {} n, c = mx.fast.spec_decode_verify(draft, lg, **kw) @@ -1472,13 +1472,47 @@ def test_spec_decode_verify_empty_batch(self): self.assertEqual((c.shape, c.dtype), ((0, K + 1), mx.int32)) results[path] = (n.shape, n.dtype, c.shape, c.dtype) with self.subTest(K=K, compare="default vs cpu"): - self.assertEqual(results["default"], results["cpu-fallback"]) + self.assertEqual(results["default"], results["cpu"]) + + def test_spec_decode_verify_shapeless_compile(self): + # One shapeless compile must serve every K, whichever K it traced first. + def onehot(rows, V=16): + lg = [[[-10.0] * V for _ in row] for row in rows] + for b, row in enumerate(rows): + for j, t in enumerate(row): + lg[b][j][t] = 10.0 + return mx.array(lg) + + cases = [ + ([[1, 2, 3]], [[1, 4, 3, 5]]), + ([[1, 2, 3, 4, 5]], [[1, 2, 3, 4, 5, 6]]), + ([[1, 2], [3, 4]], [[1, 2, 9], [3, 0, 9]]), + ([[]], [[7]]), + ([[], []], [[7], [2]]), + ] + for device in (mx.cpu, mx.gpu): + for first in ("K>0", "K=0"): + order = cases if first == "K>0" else cases[::-1] + fun = mx.compile(mx.fast.spec_decode_verify, shapeless=True) + with mx.stream(device): + for draft, target in order: + K = len(target[0]) - 1 + draft = mx.array(draft, dtype=mx.int32) + draft = draft.reshape(len(target), K) + logits = onehot(target) + with self.subTest(device=device, first=first, K=K): + n, c = fun(draft, logits) + n_ref, c_ref = mx.fast.spec_decode_verify(draft, logits) + self.assertEqual(n.shape, n_ref.shape) + self.assertEqual(c.shape, c_ref.shape) + self.assertTrue(mx.array_equal(n, n_ref)) + self.assertTrue(mx.array_equal(c, c_ref)) def test_spec_decode_verify_rejects_non_integer_drafts(self): lg = mx.zeros((1, 3, 8)) for dtype in (mx.float32, mx.float16, mx.bfloat16, mx.bool_): draft = mx.array([[1.9, 0.0]]).astype(dtype) - for path, stream in (("default", None), ("cpu-fallback", mx.cpu)): + for path, stream in (("default", None), ("cpu", mx.cpu)): with self.subTest(dtype=dtype, path=path): kw = {"stream": stream} if stream is not None else {} with self.assertRaises(ValueError): @@ -1491,7 +1525,7 @@ def test_spec_decode_verify_rejects_non_integer_drafts(self): def test_spec_decode_verify_strided_inputs(self): # Drafts and logits that are views, not dense rows: the fused kernel - # must see the logical values, like the CPU composition does. + # must see the logical values, like the CPU kernel does. def onehot(tokens, V=8): lg = [[[-10.0] * V for _ in row] for row in tokens] for b, row in enumerate(tokens): @@ -1535,7 +1569,7 @@ def onehot(tokens, V=8): ), ] for name, draft, lg, exp_n, exp_c in cases: - for path, stream in (("default", None), ("cpu-fallback", mx.cpu)): + for path, stream in (("default", None), ("cpu", mx.cpu)): with self.subTest(case=name, path=path): kw = {"stream": stream} if stream is not None else {} n, c = mx.fast.spec_decode_verify(draft, lg, **kw) From cbbbe2f5622b12b395432712931e8337edd6331f Mon Sep 17 00:00:00 2001 From: Jeethu Rao Date: Thu, 15 Oct 2026 10:14:32 +0100 Subject: [PATCH 26/34] metal: condense the 1-bit qdot comment (#282) Review feedback on numen-tech/mlx#12 (AGENTS.md: concise comments). Comment only; test_quantized 73 passed, 1 skipped. --- mlx/backend/metal/kernels/quantized.h | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/mlx/backend/metal/kernels/quantized.h b/mlx/backend/metal/kernels/quantized.h index d7857753d2..13ceab6a70 100644 --- a/mlx/backend/metal/kernels/quantized.h +++ b/mlx/backend/metal/kernels/quantized.h @@ -233,12 +233,8 @@ inline U qdot( U accum = 0; if (bits == 1) { - // 1-bit: read 32 packed weights as one 32-bit word (one load instead of - // four byte loads). Little-endian bit order and accumulation order are - // those of the per-byte path, so the result is bit-exact with it. Needs - // values_per_thread % 32 == 0 and a 4-byte aligned w, which every 1-bit - // caller (qmv_fast, qmv, qmv_quad with D == 128) provides. - // (numen-tech/gemma4-qat#179: large mobile decode win) + // One 32-bit load per 32 weights, bit-exact with the byte path. Needs + // values_per_thread % 32 == 0 and a 4-byte aligned w. const device uint32_t* w32 = (const device uint32_t*)w; for (int i = 0; i < (values_per_thread / 32); i++) { uint32_t wb = w32[i]; From c134ee26d066da49c97d15787fb18bb024870411 Mon Sep 17 00:00:00 2001 From: Jeethu Rao Date: Thu, 15 Oct 2026 10:17:23 +0100 Subject: [PATCH 27/34] python/tests: run the shapeless spec_decode_verify test on GPU only with Metal (#282) Codex on 1e1f993c2: mx.stream(mx.gpu) throws on a CPU-only build, and a CUDA GPU stream still takes the shape-dependent composition. CPU coverage stays unconditional. test_fast -k spec_decode 6 passed. --- python/tests/test_fast.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/tests/test_fast.py b/python/tests/test_fast.py index 5251f33203..7110643891 100644 --- a/python/tests/test_fast.py +++ b/python/tests/test_fast.py @@ -1490,7 +1490,9 @@ def onehot(rows, V=16): ([[]], [[7]]), ([[], []], [[7], [2]]), ] - for device in (mx.cpu, mx.gpu): + # CUDA GPU streams still take the shape-dependent composition. + devices = [mx.cpu] + ([mx.gpu] if mx.metal.is_available() else []) + for device in devices: for first in ("K>0", "K=0"): order = cases if first == "K>0" else cases[::-1] fun = mx.compile(mx.fast.spec_decode_verify, shapeless=True) From 560cd028a58b2215a1f656a5c6c22819bb1a3c8c Mon Sep 17 00:00:00 2001 From: Jeethu Rao Date: Thu, 15 Oct 2026 23:14:55 +0100 Subject: [PATCH 28/34] fast: spec_decode_verify builds one graph for every B, including B = 0 (#282) The B = 0 early return recorded two constant empty results, so a shapeless compile first traced with B = 0 returned empty arrays for every later B. B = 0 now takes the same primitive path as any B: the target ids come from an ArgMax primitive built directly (the argmax op rejects any zero-size input), V = 0 is rejected up front, and SpecDecodeVerify's CPU/Metal evals and Metal ArgReduce return early on an empty output instead of dispatching a zero-size grid (a B > 0-first trace replayed at B = 0 tripped the Metal validation layer). CUDA GPU streams keep the composition and its B = 0 early return; that path is not shapeless-safe for a change of B or K. --- mlx/backend/cpu/spec_decode.cpp | 3 ++ mlx/backend/metal/primitives.cpp | 3 ++ mlx/backend/metal/spec_decode.cpp | 3 ++ mlx/fast.cpp | 28 +++++++++++++----- python/tests/test_fast.py | 49 +++++++++++++++++++++++++++++++ 5 files changed, 79 insertions(+), 7 deletions(-) diff --git a/mlx/backend/cpu/spec_decode.cpp b/mlx/backend/cpu/spec_decode.cpp index 89c24f9f73..ce8ff02b3f 100644 --- a/mlx/backend/cpu/spec_decode.cpp +++ b/mlx/backend/cpu/spec_decode.cpp @@ -15,6 +15,9 @@ void SpecDecodeVerify::eval_cpu( auto& committed = outputs[1]; // [B, K+1] int32 n_accepted.set_data(allocator::malloc(n_accepted.nbytes())); committed.set_data(allocator::malloc(committed.nbytes())); + if (n_accepted.size() == 0) { + return; + } auto& encoder = cpu::get_command_encoder(stream()); encoder.set_input_array(draft); diff --git a/mlx/backend/metal/primitives.cpp b/mlx/backend/metal/primitives.cpp index 45929e27dd..45e73b0d16 100644 --- a/mlx/backend/metal/primitives.cpp +++ b/mlx/backend/metal/primitives.cpp @@ -63,6 +63,9 @@ void ArgReduce::eval_gpu(const std::vector& inputs, array& out) { assert(inputs.size() == 1); auto& in = inputs[0]; out.set_data(allocator::malloc(out.nbytes())); + if (out.size() == 0) { + return; + } auto& s = stream(); auto& d = metal::device(s.device); std::string op_name; diff --git a/mlx/backend/metal/spec_decode.cpp b/mlx/backend/metal/spec_decode.cpp index 8d38eaa144..fc1d688d76 100644 --- a/mlx/backend/metal/spec_decode.cpp +++ b/mlx/backend/metal/spec_decode.cpp @@ -42,6 +42,9 @@ void SpecDecodeVerify::eval_gpu( n_accepted.set_data(allocator::malloc(n_accepted.nbytes())); committed.set_data(allocator::malloc(committed.nbytes())); + if (n_accepted.size() == 0) { + return; + } int B = draft.shape(0); int K = draft.shape(1); diff --git a/mlx/fast.cpp b/mlx/fast.cpp index 163ceaab4f..a9cb60a485 100644 --- a/mlx/fast.cpp +++ b/mlx/fast.cpp @@ -1311,23 +1311,37 @@ std::vector spec_decode_verify( throw std::invalid_argument( "[spec_decode_verify] target_logits.shape[1] must equal draft_tokens.shape[1] + 1."); } + if (target_logits.shape(2) == 0) { + throw std::invalid_argument( + "[spec_decode_verify] target_logits must have a non-empty vocab axis."); + } auto s = to_stream(s_); - // B = 0: nothing to verify. Return empty results without the argmax (which - // rejects a zero-size input) or the fused kernel. - if (B == 0) { + // CUDA GPU streams take the composition, which captures B and K as + // constants (arange, full, zeros), so it is not shapeless-safe for any + // change of B or K, B = 0 included. B = 0 returns empty results there + // without the argmax. + bool fallback = SpecDecodeVerify::use_fallback(s); + if (fallback && B == 0) { return {zeros(Shape{0}, int32, s), zeros(Shape{0, K + 1}, int32, s)}; } // Reuse mlx's fast argmax on-device (result stays on the GPU); the fused // primitive does only the tiny per-row prefix-match. Primitive inputs are - // int32 token arrays, so the kernel needs no large-V reduction. + // int32 token arrays, so the kernel needs no large-V reduction. The ArgMax + // primitive is built directly because the argmax op rejects any zero-size + // input, B = 0 included; V > 0 is checked above. auto dft = astype(draft_tokens, int32, s); - auto tgt = astype(argmax(target_logits, -1, false, s), int32, s); // [B, K+1] + auto tgt = array( + Shape{B, K + 1, 1}, + uint32, + std::make_shared(s, ArgReduce::ArgMax, 2), + {target_logits}); + tgt = astype(squeeze(tgt, 2, s), int32, s); // [B, K+1] std::vector prim_inputs = {dft, tgt}; - // Every K takes the primitive, so a shapeless compile does not fix K. - if (SpecDecodeVerify::use_fallback(s)) { + // Every B and K takes the primitive, so a shapeless compile fixes neither. + if (fallback) { return spec_decode_verify_composition(prim_inputs, s); } return array::make_arrays( diff --git a/python/tests/test_fast.py b/python/tests/test_fast.py index 7110643891..91d586a8df 100644 --- a/python/tests/test_fast.py +++ b/python/tests/test_fast.py @@ -1438,6 +1438,14 @@ def check(draft, tgt_tokens, V, dtype, stream): mx.zeros((1, 3), dtype=mx.int32), mx.zeros((1, 3, 8)) ) + # An empty vocab has no argmax, for any B. + for B in (1, 0): + with self.subTest(empty_vocab=B): + with self.assertRaises(ValueError): + mx.fast.spec_decode_verify( + mx.zeros((B, 2), dtype=mx.int32), mx.zeros((B, 3, 0)) + ) + def test_spec_decode_verify_empty_draft(self): # K = 0: nothing to accept, the bonus token is the argmax at the only # position. The CPU kernel must match the fused kernel. @@ -1510,6 +1518,47 @@ def onehot(rows, V=16): self.assertTrue(mx.array_equal(n, n_ref)) self.assertTrue(mx.array_equal(c, c_ref)) + def test_spec_decode_verify_shapeless_compile_batch(self): + # One shapeless compile must serve every B, B = 0 included, whichever + # B it traced first. + V = 16 + + def onehot(rows, K): + lg = [[[-10.0] * V for _ in row] for row in rows] + for b, row in enumerate(rows): + for j, t in enumerate(row): + lg[b][j][t] = 10.0 + return mx.array(lg).reshape(len(rows), K + 1, V) + + # (K, draft, target) per call; B = len(target). + cases = [ + (3, [], []), + (3, [[1, 2, 3]], [[1, 4, 3, 5]]), + (3, [[1, 2, 3], [3, 4, 9]], [[1, 2, 3, 5], [3, 0, 9, 9]]), + (0, [], []), + (0, [[], []], [[7], [2]]), + (1, [[7], [7], [7]], [[7, 1], [2, 3], [7, 9]]), + ] + # CUDA GPU streams still take the shape-dependent composition. + devices = [mx.cpu] + ([mx.gpu] if mx.metal.is_available() else []) + for device in devices: + for first in ("B=0", "B>0"): + order = cases if first == "B=0" else cases[1:] + cases[:1] + fun = mx.compile(mx.fast.spec_decode_verify, shapeless=True) + with mx.stream(device): + for K, draft, target in order: + B = len(target) + draft = mx.array(draft, dtype=mx.int32).reshape(B, K) + logits = onehot(target, K) + with self.subTest(device=device, first=first, B=B, K=K): + n, c = fun(draft, logits) + n_ref, c_ref = mx.fast.spec_decode_verify(draft, logits) + self.assertEqual(n.shape, n_ref.shape) + self.assertEqual(c.shape, c_ref.shape) + self.assertEqual((n.dtype, c.dtype), (mx.int32, mx.int32)) + self.assertTrue(mx.array_equal(n, n_ref)) + self.assertTrue(mx.array_equal(c, c_ref)) + def test_spec_decode_verify_rejects_non_integer_drafts(self): lg = mx.zeros((1, 3, 8)) for dtype in (mx.float32, mx.float16, mx.bfloat16, mx.bool_): From 948213535d96232ea58206de894ece5337746aaf Mon Sep 17 00:00:00 2001 From: Jeethu Rao Date: Thu, 15 Oct 2026 23:36:42 +0100 Subject: [PATCH 29/34] cpu: point the bias-free affine error at the implied-bias factor (#282) derive_biases does not exist. The CPU kernels take a 0-d implied-bias factor, and -0.5 (1-bit) / -1 (2-bit) gives the same bias the Metal bias-free path derives (sym_derived_bias). --- mlx/backend/cpu/quantized.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/mlx/backend/cpu/quantized.cpp b/mlx/backend/cpu/quantized.cpp index fc3a1813ef..66cbe747a7 100644 --- a/mlx/backend/cpu/quantized.cpp +++ b/mlx/backend/cpu/quantized.cpp @@ -946,7 +946,8 @@ void QuantizedMatmul::eval_cpu(const std::vector& inputs, array& out) { if (inputs.size() < 4) { throw std::runtime_error( "[QuantizedMatmul::eval_cpu] Bias-free affine is Metal-only for " - "now; pass biases (e.g. derive_biases) on CPU."); + "now; on CPU pass biases as the 0-d implied-bias factor (-0.5 for " + "1-bit, -1 for 2-bit)."); } auto biases = ensure_row_contiguous(inputs[3], encoder, stream()); encoder.set_input_array(biases); From 7b7e8d2b1f72b3e20529782df140b970824b9fc7 Mon Sep 17 00:00:00 2001 From: Jeethu Rao Date: Thu, 15 Oct 2026 23:36:43 +0100 Subject: [PATCH 30/34] metal: condense the ragged-tile comment in qmv_quad (#282) --- mlx/backend/metal/kernels/quantized.h | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/mlx/backend/metal/kernels/quantized.h b/mlx/backend/metal/kernels/quantized.h index 13ceab6a70..13197bcd8d 100644 --- a/mlx/backend/metal/kernels/quantized.h +++ b/mlx/backend/metal/kernels/quantized.h @@ -874,9 +874,8 @@ METAL_FUNC void qmv_quad_impl( const device T* sl = scales + row * in_vec_size_g * quads_per_simd; const device T* bl = biases + row * in_vec_size_g * quads_per_simd; - // Rows past out_vec_size exist only on a ragged last tile (N % 64 != 0): - // read no scale or bias for them (gemma4-qat#302; the store below was - // already guarded). + // Rows past out_vec_size occur only on a ragged last tile: skip their + // scale and bias reads. if (row * quads_per_simd + out_row < out_vec_size) { U s = sl[0]; U b = affine_bias(s, f, sl[0], bl); From c52689caad322192941677005000ebddb3c64eb9 Mon Sep 17 00:00:00 2001 From: Jeethu Rao Date: Fri, 16 Oct 2026 12:17:47 +0100 Subject: [PATCH 31/34] metal: condense the spec_decode and qmv alignment comments (#282) --- mlx/backend/metal/kernels/spec_decode.metal | 8 ++------ mlx/backend/metal/quantized.cpp | 9 ++------- 2 files changed, 4 insertions(+), 13 deletions(-) diff --git a/mlx/backend/metal/kernels/spec_decode.metal b/mlx/backend/metal/kernels/spec_decode.metal index 61b6c66c9b..0093028939 100644 --- a/mlx/backend/metal/kernels/spec_decode.metal +++ b/mlx/backend/metal/kernels/spec_decode.metal @@ -4,12 +4,8 @@ using namespace metal; -// Greedy speculative-decoding verify, one thread per batch row. -// draft [B, K] int32 — drafted tokens -// target [B, K+1] int32 — target argmax tokens over [last] + draft -// Writes: -// n_accepted [B] — longest matching prefix length (0..K) -// committed [B, K+1] — draft prefix (< n) then the corrected token at n. +// Greedy speculative-decode verify, one thread per row: n_accepted is the +// matching draft prefix length, committed is that prefix then the target token. [[kernel]] void spec_decode_verify( const device int* draft [[buffer(0)]], const device int* target [[buffer(1)]], diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index 1cfea56896..bf37ed3635 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -156,13 +156,8 @@ inline int get_qmv_batch_limit(int D, int O, metal::Device& d) { } } -// Block of qmv_fast_impl (kernels/quantized.h): pack_factor * -// packs_per_thread * SIMD_SIZE, with one pack per lane for bits <= 2 (must -// match the kernel). The affine kernel's partial last block only needs -// K % values_per_thread == 0, which K % group_size already implies; block / 2 -// is the cutoff measured for 1/4-bit, and the other widths stay whole-block -// until benchmarked. fp_qmv_fast_impl has no tail, so the fp modes (which -// share bits == 4) keep the whole block. +// K alignment for qmv_fast_impl; the block must match the kernel. Affine +// 1/4-bit takes half a block (measured), the rest a whole block. inline int qmv_fast_k_alignment(int bits, const std::string& mode) { int block = get_pack_factor(bits, 32) * (bits <= 2 ? 1 : 2) * 32; bool affine = mode == "affine" || mode == "affine_sym"; From 7474d2cfd632f7aeac3f0df74832c252f2a98c08 Mon Sep 17 00:00:00 2001 From: Jeethu Rao Date: Fri, 16 Oct 2026 12:17:47 +0100 Subject: [PATCH 32/34] fast: spec_decode_verify refuses a shapeless trace on the CUDA composition (#282) The CUDA GPU composition fixes B and K at trace time (slice, arange, full, zeros), so a shapeless compile replayed at another B or K gave shape errors or stale results (a B = 0-first trace kept returning empty arrays). It now throws under detail::in_dynamic_tracing(); Metal and CPU take the primitive for every B and K. --- mlx/fast.cpp | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/mlx/fast.cpp b/mlx/fast.cpp index a9cb60a485..84457a00ab 100644 --- a/mlx/fast.cpp +++ b/mlx/fast.cpp @@ -1317,11 +1317,14 @@ std::vector spec_decode_verify( } auto s = to_stream(s_); - // CUDA GPU streams take the composition, which captures B and K as - // constants (arange, full, zeros), so it is not shapeless-safe for any - // change of B or K, B = 0 included. B = 0 returns empty results there - // without the argmax. + // CUDA GPU streams take the composition, which fixes B and K at trace + // time (arange, full, zeros): refuse a shapeless trace, not a stale graph. bool fallback = SpecDecodeVerify::use_fallback(s); + if (fallback && detail::in_dynamic_tracing()) { + throw std::invalid_argument( + "[spec_decode_verify] Shapeless compile is not supported on this " + "stream; compile with shapeless=False."); + } if (fallback && B == 0) { return {zeros(Shape{0}, int32, s), zeros(Shape{0, K + 1}, int32, s)}; } @@ -1340,7 +1343,8 @@ std::vector spec_decode_verify( tgt = astype(squeeze(tgt, 2, s), int32, s); // [B, K+1] std::vector prim_inputs = {dft, tgt}; - // Every B and K takes the primitive, so a shapeless compile fixes neither. + // Off CUDA every B and K takes the primitive, so a shapeless compile fixes + // neither. if (fallback) { return spec_decode_verify_composition(prim_inputs, s); } From 713bf1982892586c562ab42ff496eb298d5a3a89 Mon Sep 17 00:00:00 2001 From: Jeethu Rao Date: Fri, 16 Oct 2026 12:17:47 +0100 Subject: [PATCH 33/34] metal_kernel: build a missing variant under the cache lock (#282) Concurrent first callers of one variant each built its source and hash before the second lookup kept one. The build now runs under cache->mutex, so it happens once. It covers source generation only; the Metal library is compiled later in CustomKernel::eval_gpu. --- mlx/backend/common/metal_kernel.cpp | 22 +++++++++------------- 1 file changed, 9 insertions(+), 13 deletions(-) diff --git a/mlx/backend/common/metal_kernel.cpp b/mlx/backend/common/metal_kernel.cpp index 7e0a28a7ff..90959fcf0e 100644 --- a/mlx/backend/common/metal_kernel.cpp +++ b/mlx/backend/common/metal_kernel.cpp @@ -358,11 +358,11 @@ CustomKernelFunction metal_kernel( auto key = variant_key(inputs, output_dtypes, template_args); std::shared_ptr variant; - { - std::lock_guard lock(cache->mutex); - if (auto it = cache->variants.find(key); it != cache->variants.end()) { - variant = it->second; - } + // Build a missing variant under the lock, so concurrent first callers + // build it once. + std::unique_lock lock(cache->mutex); + if (auto it = cache->variants.find(key); it != cache->variants.end()) { + variant = it->second; } if (!variant) { std::string kernel_name = "custom_kernel_" + name; @@ -421,16 +421,12 @@ CustomKernelFunction metal_kernel( std::move(kernel_name), std::make_shared(std::move(kernel_source)), source_hash}); - std::lock_guard lock(cache->mutex); - if (auto it = cache->variants.find(key); it != cache->variants.end()) { - variant = it->second; - } else { - if (cache->variants.size() >= metal_kernel_max_cached_variants) { - cache->variants.clear(); - } - cache->variants.emplace(std::move(key), variant); + if (cache->variants.size() >= metal_kernel_max_cached_variants) { + cache->variants.clear(); } + cache->variants.emplace(std::move(key), variant); } + lock.unlock(); if (verbose) { std::cout << "Generated source code for `" << name << "`:" << std::endl From bb8fa5facb1a98e2d6ed723927e057590b184e35 Mon Sep 17 00:00:00 2001 From: Jeethu Rao Date: Fri, 16 Oct 2026 13:19:32 +0100 Subject: [PATCH 34/34] Condense the remaining multi-line source comments to two lines (#282) --- mlx/backend/common/metal_kernel.cpp | 16 ++++++---------- mlx/backend/metal/device.cpp | 6 ++---- mlx/backend/metal/quantized.cpp | 17 +++++------------ mlx/backend/metal/spec_decode.cpp | 5 ++--- mlx/fast.cpp | 7 ++----- 5 files changed, 17 insertions(+), 34 deletions(-) diff --git a/mlx/backend/common/metal_kernel.cpp b/mlx/backend/common/metal_kernel.cpp index 90959fcf0e..98c913d369 100644 --- a/mlx/backend/common/metal_kernel.cpp +++ b/mlx/backend/common/metal_kernel.cpp @@ -226,18 +226,16 @@ struct KernelVariant { size_t source_hash; }; -// Shared by all copies of one kernel function. Holds at most -// `metal_kernel_max_cached_variants` entries and is cleared when full. Callers -// keep an owning pointer to their entry, so clearing is safe at any time. +// One per kernel function, cleared when full. Callers own their entry, so a +// clear is safe at any time. struct VariantCache { std::mutex mutex; std::unordered_map> variants; }; -// Encodes each call argument that the name and source depend on. Template -// arguments include their kind and name: `int` 1 and `bool` true give the same -// name but a different signature. +// Encodes each argument that the name and source depend on. Template arguments +// add their kind, because `int` 1 and `bool` true give the same name. std::string variant_key( const std::vector& inputs, const std::vector& output_dtypes, @@ -374,10 +372,8 @@ CustomKernelFunction metal_kernel( kernel_name += template_hash; } - // The generated source depends on the dtypes of the inputs and outputs - // and on how each input is passed (see `write_signature`). Include them - // in the kernel name so that a given name always maps to the same - // source. + // The source depends on the input and output dtypes and on how each + // input is passed, so the name includes them. for (const auto& arr : inputs) { kernel_name += "_"; kernel_name += get_type_string(arr.dtype()); diff --git a/mlx/backend/metal/device.cpp b/mlx/backend/metal/device.cpp index 5428c8be7b..90cd2dcae2 100644 --- a/mlx/backend/metal/device.cpp +++ b/mlx/backend/metal/device.cpp @@ -59,10 +59,8 @@ void set_compile_options( default: throw std::invalid_argument("[metal::Device] Invalid math mode."); } - // MTLCompileOptions.mathFloatingPointFunctions is independent of mathMode - // and defaults to Fast, so a Safe/Relaxed request would still compile the - // fast (approximate) transcendental variants. Keep them Precise unless - // fast math was explicitly asked for. + // mathFloatingPointFunctions defaults to Fast whatever the mathMode, so set + // Precise unless fast math was requested. mtl_options->setMathFloatingPointFunctions( compile_options.math_mode == MathMode::Fast ? MTL::MathFloatingPointFunctionsFast diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index bf37ed3635..b2b9fe9243 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -609,11 +609,8 @@ void qmv( compute_encoder.dispatch_threadgroups(grid_dims, group_dims); } -// affine qmv_wide only beats qmv on gen-15+; fp benefits on every gen. -// affine_sym has no qmv_wide kernel. The 1-bit path has so little weight -// traffic that unpacking into registers dominates even when several input -// rows reuse the result. The 2-bit path breaks even after two rows and wins -// once three or more rows share a block. +// affine qmv_wide beats qmv only on gen-15+, and not at 1-bit or at 2-bit with +// M < 3 (measured); fp gains on every gen. affine_sym has no qmv_wide kernel. inline bool use_qmv_wide(const std::string& mode, int bits, int M, metal::Device& d) { if (mode == "affine_sym") { @@ -1296,10 +1293,8 @@ void qmm( std::string template_def; MTL::ComputePipelineState* kernel; const bool implied_bias = is_implied_bias(biases); - // The implied-bias variants spell out has_global_scale (an implied bias is a - // bias, so it never comes with a global scale) and the default tile sizes - // (BM/BK/BN = 32) so the trailing `implied_bias` template argument lands in - // its slot. + // Implied-bias variants spell out has_global_scale (false) and the default + // tiles (32), so the trailing `implied_bias` argument lands in its slot. if (transpose) { kernel = implied_bias ? get_quantized_kernel_wrapped( d, @@ -2082,9 +2077,7 @@ void dispatch_qmv( const Stream& s, const std::string& mode) { // It is a qmv with a small inner dimension so route to qmv_quad kernel - // (no bias-free variant; affine_sym falls through to the plain qmv). - // 1-bit needs K == 128: at K == 64 a quad lane holds 16 values, less than - // the one 32-value uint32 pack that qmv_quad_impl and qdot read. + // (1-bit needs K == 128 for a full pack per lane; no affine_sym variant). if ((K == 128 || (K == 64 && bits >= 2)) && is_power_of_2(bits) && !global_scale && mode != "affine_sym") { qmv_quad(x, w, scales, biases, out, group_size, bits, M, N, K, d, s, mode); diff --git a/mlx/backend/metal/spec_decode.cpp b/mlx/backend/metal/spec_decode.cpp index fc1d688d76..c8e32120a4 100644 --- a/mlx/backend/metal/spec_decode.cpp +++ b/mlx/backend/metal/spec_decode.cpp @@ -12,9 +12,8 @@ namespace mlx::core::fast { namespace { -// The kernel indexes dense rows (row * K, row * (K + 1)). A row-contiguous -// input is bound as is; a view (transposed, sliced, broadcast) is copied -// first and the copy kept alive until the command buffer completes. +// The kernel indexes dense rows, so copy a view and keep the copy alive until +// the command buffer completes. array ensure_row_contiguous(const array& x, const Stream& s) { if (x.flags().row_contiguous) { return x; diff --git a/mlx/fast.cpp b/mlx/fast.cpp index 84457a00ab..08cf006f21 100644 --- a/mlx/fast.cpp +++ b/mlx/fast.cpp @@ -1329,11 +1329,8 @@ std::vector spec_decode_verify( return {zeros(Shape{0}, int32, s), zeros(Shape{0, K + 1}, int32, s)}; } - // Reuse mlx's fast argmax on-device (result stays on the GPU); the fused - // primitive does only the tiny per-row prefix-match. Primitive inputs are - // int32 token arrays, so the kernel needs no large-V reduction. The ArgMax - // primitive is built directly because the argmax op rejects any zero-size - // input, B = 0 included; V > 0 is checked above. + // Build ArgMax directly, because the argmax op rejects a zero-size input + // (B = 0). The fused kernel then does only the per-row prefix match. auto dft = astype(draft_tokens, int32, s); auto tgt = array( Shape{B, K + 1, 1},