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..d247042aed 100644 --- a/mlx/backend/cpu/quantized.cpp +++ b/mlx/backend/cpu/quantized.cpp @@ -950,6 +950,66 @@ 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: `biases` is a single factor f, exact in the scales' + // dtype T by contract. The CPU kernels index a per-group bias array, so + // materialize f * scales into a temporary: the float product of two T + // values is exact, so rounding it to T matches `scales * T(f)`. + 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 033723d46b..4c34efaabb 100644 --- a/mlx/backend/metal/kernels/quantized.h +++ b/mlx/backend/metal/kernels/quantized.h @@ -660,6 +660,29 @@ 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. Implied bias: `biases` holds a +// single factor f and the bias is f * scale, formed in T so it matches a bias +// materialized in the scales dtype; otherwise the stored bias is read. +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, @@ -668,7 +691,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, @@ -701,6 +725,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_, @@ -723,7 +748,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) { @@ -731,7 +762,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); @@ -758,7 +789,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), @@ -776,20 +807,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, @@ -813,6 +855,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; @@ -832,7 +876,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); } @@ -846,17 +890,12 @@ 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 +template < + typename T, + int group_size, + int bits, + bool bias_free = false, + bool implied_bias = false> METAL_FUNC void qmv_fast_impl( const device uint32_t* w, const device T* scales, @@ -884,6 +923,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; @@ -907,7 +948,7 @@ METAL_FUNC void qmv_fast_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); } @@ -931,7 +972,7 @@ METAL_FUNC void qmv_fast_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); } @@ -946,7 +987,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, @@ -975,6 +1021,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; @@ -1008,7 +1056,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); } @@ -1034,7 +1082,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); } @@ -1069,7 +1117,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); } @@ -1093,7 +1141,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); } @@ -1110,7 +1158,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, @@ -1143,6 +1197,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++) { @@ -1155,7 +1210,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; @@ -1204,7 +1260,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, @@ -1240,6 +1300,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; @@ -1260,7 +1322,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); @@ -1274,7 +1337,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( @@ -1288,7 +1352,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; @@ -1321,7 +1386,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> METAL_FUNC void qmm_t_impl( const device uint32_t* w, const device T* scales, @@ -1363,7 +1429,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; @@ -1376,7 +1443,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 @@ -1448,7 +1517,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, @@ -1490,7 +1560,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; @@ -1500,7 +1571,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 @@ -1671,7 +1744,17 @@ METAL_FUNC void adjust_matrix_offsets( y += tid.z * output_stride; } -template +// `implied_bias` (last template parameter of every affine_* kernel below): +// the biases buffer holds a single factor f and the per-group bias is f * +// scale, formed in-kernel. The host binds a 0-d array and zero b_strides, so +// adjust_matrix_offsets leaves the pointer on the factor. +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)]], @@ -1710,7 +1793,7 @@ template b_strides, tid); } - qmv_quad_impl( + qmv_quad_impl( w, scales, biases, @@ -1729,7 +1812,8 @@ template < int bits, bool batched, 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)]], @@ -1768,7 +1852,7 @@ template < b_strides, tid); } - qmv_fast_impl( + qmv_fast_impl( w, scales, biases, @@ -1787,7 +1871,8 @@ template < const int bits, bool batched, 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)]], @@ -1826,7 +1911,7 @@ template < b_strides, tid); } - qmv_impl( + qmv_impl( w, scales, biases, @@ -1961,7 +2046,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, @@ -2000,7 +2086,7 @@ template < b_strides, tid); } - qmv_wide_impl( + qmv_wide_impl( w, scales, biases, @@ -2019,7 +2105,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)]], @@ -2058,7 +2145,7 @@ template < b_strides, tid); } - qvm_impl( + qvm_impl( w, scales, biases, @@ -2072,7 +2159,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)]], @@ -2118,7 +2210,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, @@ -2140,7 +2232,8 @@ template < const bool batched, 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)]], @@ -2187,7 +2280,7 @@ template < b_strides, tid); } - qmm_t_impl( + qmm_t_impl( w, scales, biases, @@ -2212,7 +2305,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)]], @@ -2243,10 +2337,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, @@ -2272,7 +2368,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)]], @@ -2321,7 +2418,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); } @@ -2973,7 +3070,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)]], @@ -2988,7 +3090,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 91dcaf7362..78dfc2a607 100644 --- a/mlx/backend/metal/kernels/quantized.metal +++ b/mlx/backend/metal/kernels/quantized.metal @@ -186,6 +186,99 @@ instantiate_quantized_sym_types(64, bits) \ instantiate_quantized_sym_types(32, bits) +// Implied-bias ("_ib") variants: `biases` is a single factor f and the +// kernel forms bias = f * scale per group. The trailing `true` is the +// kernels' `implied_bias` template argument; the defaulted parameters before +// it (has_global_scale = false, results_per_simdgroup = 4, BM/BK/BN = 32) are +// spelled out, matching the host's template definition. Only the bit widths +// of the shipped power-of-two-factor checkpoints (1, 2, 4, 8) are built into +// the metallib; 3/5/6-bit implied-bias kernels exist in the JIT build only. +#define instantiate_quantized_ib_qmv(name, type, group_size, bits, batched) \ + instantiate_kernel( \ + #name "_" #type "_gs_" #group_size "_b_" #bits "_batch_" #batched "_ib", \ + name, type, group_size, bits, batched, 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, 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) \ + instantiate_quantized_ib_qmv(affine_qmv_fast, type, group_size, bits, 0) \ + instantiate_quantized_ib_qmv(affine_qmv, type, group_size, bits, 1) \ + instantiate_quantized_ib_qmv(affine_qmv, type, group_size, bits, 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) \ @@ -195,6 +288,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 a8e6384b07..8251ea245b 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); @@ -770,7 +777,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), @@ -788,15 +795,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; + } } } }; @@ -808,7 +821,8 @@ template < short dst_ld, short reduction_dim, short tgp_size, - short bits> + short bits, + bool implied_bias> struct QuantizedBlockLoader< T, BROWS, @@ -817,7 +831,8 @@ struct QuantizedBlockLoader< reduction_dim, tgp_size, 32, - bits> { + bits, + implied_bias> { MLX_MTL_CONST short group_size = 32; static_assert( @@ -853,6 +868,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_, @@ -875,7 +891,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) { @@ -883,7 +906,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); @@ -910,7 +933,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), @@ -932,11 +955,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; + } } } }; @@ -1043,7 +1070,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, @@ -1076,7 +1104,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; @@ -1089,7 +1118,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 @@ -1187,7 +1218,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, @@ -1220,7 +1252,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; @@ -1232,7 +1265,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 @@ -1317,7 +1352,8 @@ template < const int BK = 32, const int BN = 64, const int WM = 2, - const int WN = 2> + const int WN = 2, + bool implied_bias = false> [[kernel]] void affine_qmm_t_nax( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -1363,7 +1399,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); } @@ -1376,7 +1422,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)]], @@ -1423,7 +1470,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..7088f0cb2b 100644 --- a/mlx/backend/metal/kernels/quantized_nax.metal +++ b/mlx/backend/metal/kernels/quantized_nax.metal @@ -101,6 +101,42 @@ instantiate_quantized_types(64, bits) \ instantiate_quantized_types(32, bits) +// Implied-bias ("_ib") variants: `biases` is a single factor f and the +// kernel forms bias = f * scale per group (trailing `implied_bias = true`). +// Built for the shipped power-of-two-factor bit widths (1, 2, 4, 8) only; +// 3/5/6-bit implied-bias NAX kernels exist in the JIT build 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, 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 +144,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 a5a5dc0b71..051ba103b7 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -30,8 +30,8 @@ auto get_quantized_kernel_wrapped( std::string template_def; // "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); + 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); @@ -54,6 +54,11 @@ auto get_qmm_nax_kernel_wrapped( return get_qmm_nax_kernel(d, name, template_def, mode); } +// Implied-bias kernels (is_implied_bias, mlx/backend/common/quantized.h) are +// selected by an "_ib" kernel-name suffix and a trailing `implied_bias = true` +// template arg. The affine_sym and fp kernels have no such parameter, so the +// trailing arg is passed only on the implied branch. + inline array ensure_row_contiguous(const array& x, metal::Device& d, const Stream& s) { if (!x.flags().row_contiguous) { @@ -69,6 +74,10 @@ inline array ensure_row_contiguous_matrix( const array& x, metal::Device& d, const Stream& s) { + 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; @@ -176,7 +185,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; @@ -280,6 +296,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, @@ -293,15 +310,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); @@ -441,9 +471,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); @@ -505,18 +555,31 @@ void qmv( bits, use_narrow_qmv ? "_r_2" : "", B > 1 ? "_batch_1" : "_batch_0", - global_scale ? "_hgs" : ""); - auto kernel = get_quantized_kernel_wrapped( - d, - kname, - (fast ? "qmv_fast" : "qmv"), - mode, - type_string, - group_size, - bits, - B > 1, - global_scale.has_value(), - results_per_simdgroup); + global_scale ? "_hgs" : "", + 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, + 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, + global_scale.has_value(), + results_per_simdgroup); auto& compute_encoder = metal::get_command_encoder(s); compute_encoder.set_compute_pipeline_state(kernel); @@ -610,18 +673,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); @@ -719,11 +795,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); @@ -746,8 +840,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++); @@ -798,17 +899,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); @@ -879,40 +992,73 @@ void qmm_nax( "_wn", wn, transpose ? (aligned ? "_alN_true" : "_alN_false") : "", - batched ? "_batch_1" : "_batch_0"); + batched ? "_batch_1" : "_batch_0", + 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); + 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, + true) + : get_qmm_nax_kernel_wrapped( + d, + kname, + "qmm_t_nax", + mode, + type_string, + group_size, + bits, + aligned, + batched, + bm, + bk, + bn, + wm, + wn); } 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); @@ -1098,23 +1244,56 @@ void qmm( "_b_", bits, transpose ? (aligned ? "_alN_true" : "_alN_false") : "", - batched ? "_batch_1" : "_batch_0"); + batched ? "_batch_1" : "_batch_0", + 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 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); + kernel = implied_bias ? get_quantized_kernel_wrapped( + d, + kname, + "qmm_t", + mode, + type_string, + group_size, + bits, + aligned, + batched, + bm, + 32, + bn, + true) + : get_quantized_kernel_wrapped( + d, + kname, + "qmm_t", + mode, + type_string, + group_size, + bits, + aligned, + batched); } 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); @@ -1203,9 +1382,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 d7710484a9..139486303c 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() @@ -4721,8 +4724,13 @@ 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 (0-d biases) adopts the scales dtype instead of + // promoting it: `mx.array(-8.0)` is float32 but the bias it implies lives + // in the scales' precision (exactly, as the factor is a power of two). + 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() @@ -4781,6 +4789,41 @@ void validate_global_scale( } } +// Implied bias (a 0-d `biases` array holding a factor f): the contract is that +// it stands for the materialized per-group bias `scales * T(f)`, formed in the +// scales' storage dtype T -- f is cast to T first, so a factor that is not +// exactly representable in T is outside the contract (the reference is still +// `scales * T(f)`, not `T(f * scales)`). The `_ib` kernels form `T(f) * +// T(scale)` in place of loading the bias, which honors that definition exactly +// only when +// - the kernel sees the scales in T (the activations' dtype did not promote +// them; a promoted kernel would round the product in the wider dtype), and +// - the non-JIT Metal library ships the `_ib` instantiation (1/2/4/8-bit; the +// 3/5/6-bit implied kernels exist in JIT builds only, and a prebuilt +// metallib has no fallback for a missing name), and +// - the stream runs on the CPU (QuantizedMatmul::eval_cpu expands the factor +// itself) or on Metal (the `_ib` kernels); the CUDA kernels have no +// implied-bias variant and would index the 0-d factor as a per-group bias. +// Everywhere else the bias is materialized by `materialize_implied_bias` and +// the ordinary per-group-bias path runs. +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) { + return multiply(scales, astype(factor, scales.dtype(), s), s); +} + array quantized_matmul( const array& x, const array& w, @@ -4833,14 +4876,31 @@ array quantized_matmul( if (qmode == QuantizationMode::Affine) { inputs = {astype(x, dtype), w, astype(scales, dtype)}; if (biases) { - inputs.push_back(astype(*biases, dtype)); + // An implied-bias factor the kernels cannot honor exactly (see + // implied_bias_kernels_exact) is materialized in the scales' storage + // dtype here and takes the ordinary per-group-bias path. + 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) { + // An implied-bias factor (0-d) is never broadcast: the kernels read the + // scalar once and form `factor * scale` per group. + 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; @@ -5287,19 +5347,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; @@ -5353,7 +5420,13 @@ 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); + // Implied bias: form `scales * T(f)` per group in the scales' dtype (the + // factor is cast first so a wider factor cannot promote the product past + // the rounding a materialized bias went through). + 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}; @@ -5362,12 +5435,20 @@ array affine_dequantize( if (s.device == Device::gpu) { auto out_shape = w.shape(); out_shape.back() = out_size; + // The dequantize kernel reads the implied-bias factor in the scales' + // dtype, so cast a 0-d factor given in another dtype; when the non-JIT + // library has no `_ib` kernel for this bit width, materialize the bias + // instead (see implied_bias_kernels_exact). + 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]; } @@ -5587,7 +5668,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 */, @@ -5598,11 +5679,16 @@ array gather_qmm( StreamOrDevice s /* = {} */) { if (!lhs_indices_ && !rhs_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 f597753b1e..4eb2933acd 100644 --- a/mlx/ops.h +++ b/mlx/ops.h @@ -1563,7 +1563,13 @@ 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)`). + */ MLX_API array quantized_matmul( const array& x, const array& w, @@ -1584,7 +1590,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, diff --git a/mlx/primitives.cpp b/mlx/primitives.cpp index fa09eff3d2..00d62fdaf1 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" @@ -3540,6 +3541,10 @@ 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, bias = f * scale. The reverse + // product for x passes the factor straight through (the kernels honor it). + bool implied = mode_ == QuantizationMode::Affine && primals.size() == 4 && + is_implied_bias(primals[3]); std::optional biases = std::nullopt; if (mode_ == QuantizationMode::Affine) { biases = sym @@ -3586,17 +3591,20 @@ std::vector QuantizedMatmul::vjp( } if (arg == 3) { // biases + if (implied) { + throw std::runtime_error( + "[QuantizedMatmul::vjp] no gradient wrt an implied-bias factor."); + } 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 + array bias_c = sym ? array(sym_c, primals[2].dtype()) + : implied ? astype(primals[3], 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_), diff --git a/python/src/ops.cpp b/python/src/ops.cpp index 51dd5743da..40f761d478 100644 --- a/python/src/ops.cpp +++ b/python/src/ops.cpp @@ -4645,7 +4645,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). + 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``. @@ -4780,7 +4784,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``. diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index a64c0f0af3..fd00ff3d35 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 @@ -1009,6 +1010,460 @@ 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: the generic (non-fast) qmv. + 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; there is no gradient + # wrt the factor itself. + 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]) + + with self.assertRaises(RuntimeError): + mx.eval( + mx.vjp( + lambda f: mx.quantized_matmul( + x, w_q, scales, f, True, gs, bits + ), + [f], + [cotan], + )[1] + ) + + 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 eec8b5c091..df60bbca9a 100644 --- a/tests/ops_tests.cpp +++ b/tests/ops_tests.cpp @@ -3487,6 +3487,324 @@ 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", false, dev); + auto o_ib = gather_qmm( + x, w, scales, factor, lhs, rhs, true, 64, 4, "affine", 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);