From 1eef1d155c96cdb87fa8e7ff995facf8ca26d369 Mon Sep 17 00:00:00 2001 From: Awni Hannun Date: Mon, 22 Dec 2025 20:45:19 -0800 Subject: [PATCH] Metal/CPU nvfp4 and mxfp8 (#2946) --- mlx/backend/cpu/quantized.cpp | 224 ++++++++++----- mlx/backend/metal/kernels/fp4.h | 31 +-- mlx/backend/metal/kernels/fp8.h | 34 ++- mlx/backend/metal/kernels/fp_quantized.h | 258 +++++++++--------- mlx/backend/metal/kernels/fp_quantized.metal | 140 +++++----- mlx/backend/metal/kernels/fp_quantized_nax.h | 143 +++++----- .../metal/kernels/fp_quantized_nax.metal | 61 +++-- mlx/backend/metal/quantized.cpp | 21 +- mlx/primitives.cpp | 16 +- python/mlx/nn/layers/embedding.py | 8 +- python/mlx/nn/layers/linear.py | 9 +- python/mlx/nn/layers/quantized.py | 61 +++-- python/tests/cuda_skip.py | 4 +- python/tests/test_quantized.py | 186 +++++++------ 14 files changed, 660 insertions(+), 536 deletions(-) diff --git a/mlx/backend/cpu/quantized.cpp b/mlx/backend/cpu/quantized.cpp index f7665851..b21310bc 100644 --- a/mlx/backend/cpu/quantized.cpp +++ b/mlx/backend/cpu/quantized.cpp @@ -14,7 +14,7 @@ namespace mlx::core { namespace { -const static float MXFP4_LUT[16] = { +const static float FP4_LUT[16] = { +0.0f, +0.5f, +1.0f, @@ -32,15 +32,19 @@ const static float MXFP4_LUT[16] = { -4.0f, -6.0f}; -template +template static inline T dequantize_scale(uint8_t s) { - using FOrI = union { - bfloat16_t f; - uint16_t i; - }; - FOrI out; - out.i = (s == 0 ? 0x40 : (static_cast(s) << 7)); - return static_cast(out.f); + if constexpr (group_size == 16) { + return static_cast(detail::FromFP8{}(s)); + } else { + using FOrI = union { + bfloat16_t f; + uint16_t i; + }; + FOrI out; + out.i = (s == 0 ? 0x40 : (static_cast(s) << 7)); + return static_cast(out.f); + } } inline constexpr short get_pack_factor(int bits, int wsize = 8) { @@ -437,8 +441,8 @@ void _qmm_dispatch( } } -template -void mxfp4_qmm( +template +void fp_qmm( T* result, const T* x, const uint32_t* w, @@ -446,8 +450,7 @@ void mxfp4_qmm( int M, int N, int K) { - constexpr int group_size = 32; - constexpr int pack_factor = get_pack_factor(4, 8); + constexpr int pack_factor = get_pack_factor(bits, 8); constexpr int packs_in_group = group_size / pack_factor; for (int m = 0; m < M; m++) { @@ -461,25 +464,27 @@ void mxfp4_qmm( T xi = *x++; for (int n = 0; n < N; n += group_size) { - T scale = dequantize_scale(*scales_local++); + T scale = dequantize_scale(*scales_local++); for (int ng = 0; ng < packs_in_group; ng++) { - uint8_t wi = *w_local++; -#pragma clang loop unroll(full) - for (int p = 0; p < pack_factor; p++) { + if constexpr (bits == 4) { (*result_local++) += - xi * scale * static_cast(MXFP4_LUT[wi & 0xf]); - wi >>= 4; + xi * scale * static_cast(FP4_LUT[w_local[0] & 0xf]); + (*result_local++) += + xi * scale * static_cast(FP4_LUT[(w_local[0] >> 4) & 0xf]); + } else { + (*result_local++) += + xi * scale * static_cast(detail::FromFP8{}(w_local[0])); } + w_local++; } } } - result += N; } } -template -void mxfp4_qmm_t( +template +void fp_qmm_t( T* result, const T* x, const uint32_t* w, @@ -487,8 +492,7 @@ void mxfp4_qmm_t( int M, int N, int K) { - constexpr int group_size = 32; - constexpr int pack_factor = get_pack_factor(4, 8); + constexpr int pack_factor = get_pack_factor(bits, 8); constexpr int packs_in_group = group_size / pack_factor; for (int m = 0; m < M; m++) { @@ -499,16 +503,19 @@ void mxfp4_qmm_t( const T* x_local = x; T sum = 0; for (int k = 0; k < K; k += group_size) { - T scale = dequantize_scale(*scales_local++); + T scale = dequantize_scale(*scales_local++); T gsum = 0; for (int kw = 0; kw < packs_in_group; kw++) { - uint8_t wi = *w_local++; -#pragma clang loop unroll(full) - for (int p = 0; p < pack_factor; p++) { - gsum += (*x_local++) * static_cast(MXFP4_LUT[wi & 0xf]); - wi >>= 4; + if constexpr (bits == 4) { + gsum += (*x_local++) * static_cast(FP4_LUT[w_local[0] & 0xf]); + gsum += + (*x_local++) * static_cast(FP4_LUT[(w_local[0] >> 4) & 0xf]); + } else { + gsum += + (*x_local++) * static_cast(detail::FromFP8{}(w_local[0])); } + w_local++; } sum += scale * gsum; } @@ -520,9 +527,9 @@ void mxfp4_qmm_t( } } -template -simd::Simd mxfp4_extract_bits_simd(const uint32_t* w) { - if constexpr (S == 8) { +template +simd::Simd fp_extract_bits_simd(const uint32_t* w) { + if constexpr (S == 8 && bits == 4) { constexpr std::array shifts_ = {{0, 4, 8, 12, 16, 20, 24, 28}}; auto shifts(*(simd::Simd*)&shifts_); auto wi = simd::Simd(*w); @@ -530,17 +537,20 @@ simd::Simd mxfp4_extract_bits_simd(const uint32_t* w) { wi = wi & 0xf; simd::Simd w_out; for (int i = 0; i < S; ++i) { - w_out[i] = MXFP4_LUT[wi[i]]; + w_out[i] = FP4_LUT[wi[i]]; } return w_out; + } else if constexpr (S == 8 && bits == 8) { + auto w_out = simd::load(reinterpret_cast(w)); + return detail::FromFP8{}(w_out); } else { // Appease compiler.. but should never get here throw std::runtime_error("Unsupported combination for simd qmm."); } } -template -void mxfp4_qmm_t_simd( +template +void fp_qmm_t_simd( T* result, const T* x, const uint32_t* w, @@ -548,8 +558,7 @@ void mxfp4_qmm_t_simd( int M, int N, int K) { - constexpr int group_size = 32; - constexpr int pack_factor = 32 / 4; + constexpr int pack_factor = get_pack_factor(bits, 32); constexpr int packs_in_group = group_size / pack_factor; constexpr int S = simd::max_size; static_assert( @@ -564,12 +573,12 @@ void mxfp4_qmm_t_simd( simd::Simd acc(0); auto x_local = x; for (int k = 0; k < K; k += group_size) { - T scale = dequantize_scale(*scales_local++); + T scale = dequantize_scale(*scales_local++); simd::Simd g_acc(0); for (int kw = 0; kw < packs_in_group; kw += packs_per_simd) { // Extract bits - auto wf = mxfp4_extract_bits_simd(w_local); + auto wf = fp_extract_bits_simd(w_local); w_local += packs_per_simd; simd::Simd x_simd = simd::load(x_local); g_acc = g_acc + x_simd * wf; @@ -585,8 +594,8 @@ void mxfp4_qmm_t_simd( } } -template -void mxfp4_qmm_dispatch_transpose( +template +void fp_qmm_dispatch_transpose( T* result, const T* x, const uint32_t* w, @@ -598,17 +607,17 @@ void mxfp4_qmm_dispatch_transpose( if (transposed_w) { // the simd size must be a multiple of the number of elements per word if constexpr (simd::max_size % 8 == 0) { - mxfp4_qmm_t_simd(result, x, w, scales, M, N, K); + fp_qmm_t_simd(result, x, w, scales, M, N, K); } else { - mxfp4_qmm_t(result, x, w, scales, M, N, K); + fp_qmm_t(result, x, w, scales, M, N, K); } } else { - mxfp4_qmm(result, x, w, scales, M, N, K); + fp_qmm(result, x, w, scales, M, N, K); } } -template -void mxfp4_qmm_dispatch_typed( +template +void fp_qmm_dispatch_mode( array& out, const array& x, const array& w, @@ -626,7 +635,7 @@ void mxfp4_qmm_dispatch_typed( auto w_ptr = w.data(); auto scales_ptr = scales.data(); for (int i = 0; i < batch_size; i++) { - mxfp4_qmm_dispatch_transpose( + fp_qmm_dispatch_transpose( out_ptr + i * M * N, x_ptr + elem_to_loc(i * M * K, x.shape(), x.strides()), w_ptr + elem_to_loc(i * w_els, w.shape(), w.strides()), @@ -638,21 +647,44 @@ void mxfp4_qmm_dispatch_typed( } } -void mxfp4_qmm_dispatch( +template +void fp_qmm_dispatch_typed( array& out, const array& x, const array& w, const array& scales, + int group_size, + int bits, + bool transposed_w) { + if (bits == 8) { + fp_qmm_dispatch_mode(out, x, w, scales, transposed_w); + } else if (group_size == 32) { + fp_qmm_dispatch_mode(out, x, w, scales, transposed_w); + } else { + fp_qmm_dispatch_mode(out, x, w, scales, transposed_w); + } +} + +void fp_qmm_dispatch( + array& out, + const array& x, + const array& w, + const array& scales, + int group_size, + int bits, bool transposed_w) { switch (x.dtype()) { case bfloat16: - mxfp4_qmm_dispatch_typed(out, x, w, scales, transposed_w); + fp_qmm_dispatch_typed( + out, x, w, scales, group_size, bits, transposed_w); break; case float16: - mxfp4_qmm_dispatch_typed(out, x, w, scales, transposed_w); + fp_qmm_dispatch_typed( + out, x, w, scales, group_size, bits, transposed_w); break; case float32: - mxfp4_qmm_dispatch_typed(out, x, w, scales, transposed_w); + fp_qmm_dispatch_typed( + out, x, w, scales, group_size, bits, transposed_w); break; default: throw std::invalid_argument( @@ -765,9 +797,8 @@ void _bs_qmm_dispatch( "[quantized_matmul] only floating types are supported"); } } - -template -void mxfp4_bs_qmm_dispatch_typed( +template +void fp_bs_qmm_dispatch_mode( array& out, const array& x, const array& w, @@ -794,7 +825,7 @@ void mxfp4_bs_qmm_dispatch_typed( i, lhs_indices.shape(), lhs_indices.strides())]; int w_idx = rhs_indices_ptr[elem_to_loc( i, rhs_indices.shape(), rhs_indices.strides())]; - mxfp4_qmm_dispatch_transpose( + fp_qmm_dispatch_transpose( out_ptr + i * M * N, x_ptr + elem_to_loc(x_idx * M * K, x.shape(), x.strides()), w_ptr + elem_to_loc(w_idx * w_els, w.shape(), w.strides()), @@ -807,26 +838,75 @@ void mxfp4_bs_qmm_dispatch_typed( } } -void mxfp4_bs_qmm_dispatch( +template +void fp_bs_qmm_dispatch_typed( array& out, const array& x, const array& w, const array& scales, const array& lhs_indices, const array& rhs_indices, + int group_size, + int bits, + bool transposed_w) { + if (bits == 8) { + fp_bs_qmm_dispatch_mode( + out, x, w, scales, lhs_indices, rhs_indices, transposed_w); + } else if (group_size == 32) { + fp_bs_qmm_dispatch_mode( + out, x, w, scales, lhs_indices, rhs_indices, transposed_w); + } else { + fp_bs_qmm_dispatch_mode( + out, x, w, scales, lhs_indices, rhs_indices, transposed_w); + } +} + +void fp_bs_qmm_dispatch( + array& out, + const array& x, + const array& w, + const array& scales, + const array& lhs_indices, + const array& rhs_indices, + int group_size, + int bits, bool transposed_w) { switch (x.dtype()) { case float32: - mxfp4_bs_qmm_dispatch_typed( - out, x, w, scales, lhs_indices, rhs_indices, transposed_w); + fp_bs_qmm_dispatch_typed( + out, + x, + w, + scales, + lhs_indices, + rhs_indices, + group_size, + bits, + transposed_w); break; case float16: - mxfp4_bs_qmm_dispatch_typed( - out, x, w, scales, lhs_indices, rhs_indices, transposed_w); + fp_bs_qmm_dispatch_typed( + out, + x, + w, + scales, + lhs_indices, + rhs_indices, + group_size, + bits, + transposed_w); break; case bfloat16: - mxfp4_bs_qmm_dispatch_typed( - out, x, w, scales, lhs_indices, rhs_indices, transposed_w); + fp_bs_qmm_dispatch_typed( + out, + x, + w, + scales, + lhs_indices, + rhs_indices, + group_size, + bits, + transposed_w); break; default: throw std::invalid_argument( @@ -881,8 +961,10 @@ void QuantizedMatmul::eval_cpu(const std::vector& inputs, array& out) { x = array::unsafe_weak_copy(x), w = array::unsafe_weak_copy(w), scales = array::unsafe_weak_copy(scales), + group_size_ = group_size_, + bits_ = bits_, transpose_ = transpose_]() mutable { - mxfp4_qmm_dispatch(out, x, w, scales, transpose_); + fp_qmm_dispatch(out, x, w, scales, group_size_, bits_, transpose_); }); } } @@ -953,9 +1035,19 @@ void GatherQMM::eval_cpu(const std::vector& inputs, array& out) { scales = array::unsafe_weak_copy(scales), lhs_indices = array::unsafe_weak_copy(lhs_indices), rhs_indices = array::unsafe_weak_copy(rhs_indices), + group_size_ = group_size_, + bits_ = bits_, transpose_ = transpose_]() mutable { - mxfp4_bs_qmm_dispatch( - out, x, w, scales, lhs_indices, rhs_indices, transpose_); + fp_bs_qmm_dispatch( + out, + x, + w, + scales, + lhs_indices, + rhs_indices, + group_size_, + bits_, + transpose_); }); } } diff --git a/mlx/backend/metal/kernels/fp4.h b/mlx/backend/metal/kernels/fp4.h index e701adc5..25642f20 100644 --- a/mlx/backend/metal/kernels/fp4.h +++ b/mlx/backend/metal/kernels/fp4.h @@ -1,23 +1,5 @@ #pragma once -constexpr constant static float FP4_LUT[16] = { - +0.0f, - +0.5f, - +1.0f, - +1.5f, - +2.0f, - +3.0f, - +4.0f, - +6.0f, - -0.0f, - -0.5f, - -1.0f, - -1.5f, - -2.0f, - -3.0f, - -4.0f, - -6.0f}; - struct fp4_e2m1 { fp4_e2m1(float x) { if (metal::isnan(x)) { @@ -48,11 +30,18 @@ struct fp4_e2m1 { bits |= sign_bit; } - operator float() { + operator float16_t() { half converted = as_type(ushort((bits & 7) << 9)); converted *= 16384.0; - converted = bits & 8 ? -converted : converted; - return converted; + return bits & 8 ? -converted : converted; + } + + operator float() { + return static_cast(this->operator float16_t()); + } + + operator bfloat16_t() { + return static_cast(this->operator float16_t()); } uint8_t bits; diff --git a/mlx/backend/metal/kernels/fp8.h b/mlx/backend/metal/kernels/fp8.h index 34816b42..60d34be6 100644 --- a/mlx/backend/metal/kernels/fp8.h +++ b/mlx/backend/metal/kernels/fp8.h @@ -29,24 +29,20 @@ struct fp8_e4m3 { bits |= static_cast(sign >> 24); } + operator float16_t() { + uint16_t v = (bits & 127) << 7; + half converted = as_type(v); + converted *= 256.0; + auto sign = bits & 128; + return (sign ? -converted : converted); + } + + operator bfloat16_t() { + return static_cast(this->operator float16_t()); + } + operator float() { - // From PyTorch: - // https://github.com/pytorch/pytorch/blob/e3643e1e0e923f0fc063dfab6f45c956d568919d/c10/util/Float8_e4m3fn.h#L46 - uint32_t w = static_cast(bits) << 24; - uint32_t sign = w & 0x80000000; - uint32_t nonsign = w & 0x7FFFFFFF; - - uint32_t renorm_shift = metal::clz(nonsign); - renorm_shift = renorm_shift > 4 ? renorm_shift - 4 : 0; - - int32_t inf_nan_mask = - (static_cast(nonsign + 0x01000000) >> 8) & 0x7F800000; - int32_t zero_mask = static_cast(nonsign - 1) >> 31; - uint32_t result = sign | - ((((nonsign << renorm_shift >> 4) + ((0x78 - renorm_shift) << 23)) | - inf_nan_mask) & - ~zero_mask); - return as_type(result); + return static_cast(this->operator float16_t()); } uint8_t bits; @@ -74,8 +70,10 @@ struct fp8_e8m0 { uint16_t out = (bits == 0 ? 0x40 : (static_cast(bits) << 7)); return as_type(out); } + operator float() { - return static_cast(this->operator bfloat16_t()); + uint32_t out = (bits == 0 ? 0x400000 : (static_cast(bits) << 23)); + return as_type(out); } uint8_t bits; diff --git a/mlx/backend/metal/kernels/fp_quantized.h b/mlx/backend/metal/kernels/fp_quantized.h index cae1bbd9..2bd52fa7 100644 --- a/mlx/backend/metal/kernels/fp_quantized.h +++ b/mlx/backend/metal/kernels/fp_quantized.h @@ -17,9 +17,9 @@ using namespace metal; MLX_MTL_CONST int SIMD_SIZE = 32; MLX_MTL_CONST int QUAD_SIZE = 4; -template +template inline constexpr short get_pack_factor() { - return wsize / 4; + return wsize / bits; } template @@ -27,9 +27,14 @@ inline constexpr short get_bytes_per_pack() { return wsize / 8; } -template +template static inline T dequantize_scale(uint8_t s) { - return T(*(thread fp8_e8m0*)(&s)); + if constexpr (group_size == 16) { + // Use nv scale + return T(*(thread fp8_e4m3*)(&s)); + } else { + return T(*(thread fp8_e8m0*)(&s)); + } } template @@ -43,34 +48,29 @@ struct Quantize { } }; -template +template struct Dequantize { - float operator()(uint8_t x) { - if (bits == 8) { - return float(*(thread fp8_e4m3*)(&x)); + U operator()(uint8_t x) { + if constexpr (bits == 8) { + return U(*(thread fp8_e4m3*)(&x)); } else { - return float(*(thread fp4_e2m1*)(&x)); + return U(*(thread fp4_e2m1*)(&x)); } } }; template inline void load_vector(const device T* x, thread U* x_thread) { - for (int i = 0; i < values_per_thread; i += 4) { +#pragma unroll + for (int i = 0; i < values_per_thread; i++) { x_thread[i] = x[i]; - x_thread[i + 1] = x[i + 1]; - x_thread[i + 2] = x[i + 2]; - x_thread[i + 3] = x[i + 3]; } } template inline void load_vector_safe(const device T* x, thread U* x_thread, int N) { - for (int i = 0; i < N; i += 4) { + for (int i = 0; i < N; i++) { x_thread[i] = x[i]; - x_thread[i + 1] = x[i + 1]; - x_thread[i + 2] = x[i + 2]; - x_thread[i + 3] = x[i + 3]; } for (int i = N; i < values_per_thread; i++) { @@ -78,53 +78,70 @@ inline void load_vector_safe(const device T* x, thread U* x_thread, int N) { } } -template +template inline U qdot(const device uint8_t* w, const thread U* x_thread, U scale) { U accum = 0; - const device uint16_t* ws = (const device uint16_t*)w; - for (int i = 0; i < (values_per_thread / 4); i++) { - accum += - (x_thread[4 * i] * Dequantize<4>{}(ws[i]) + - x_thread[4 * i + 1] * Dequantize<4>{}(ws[i] >> 4) + - x_thread[4 * i + 2] * Dequantize<4>{}(ws[i] >> 8) + - x_thread[4 * i + 3] * Dequantize<4>{}(ws[i] >> 12)); + if constexpr (bits == 4) { + const device uint16_t* ws = (const device uint16_t*)w; + for (int i = 0; i < (values_per_thread / 4); i++) { + accum += + (x_thread[4 * i] * Dequantize<4>{}(ws[i]) + + x_thread[4 * i + 1] * Dequantize<4>{}(ws[i] >> 4) + + x_thread[4 * i + 2] * Dequantize<4>{}(ws[i] >> 8) + + x_thread[4 * i + 3] * Dequantize<4>{}(ws[i] >> 12)); + } + } else { + for (int i = 0; i < values_per_thread; i++) { + accum += x_thread[i] * Dequantize<8>{}(w[i]); + } } + return scale * accum; } -template +template inline U qdot_safe(const device uint8_t* w, const thread U* x_thread, U scale, int N) { U accum = 0; - const device uint16_t* ws = (const device uint16_t*)w; - for (int i = 0; i < (N / 4); i++) { - accum += - (x_thread[4 * i] * Dequantize<4>{}(ws[i]) + - x_thread[4 * i + 1] * Dequantize<4>{}(ws[i] >> 4) + - x_thread[4 * i + 2] * Dequantize<4>{}(ws[i] >> 8) + - x_thread[4 * i + 3] * Dequantize<4>{}(ws[i] >> 12)); + if constexpr (bits == 4) { + const device uint16_t* ws = (const device uint16_t*)w; + for (int i = 0; i < (N / 4); i++) { + accum += + (x_thread[4 * i] * Dequantize<4>{}(ws[i]) + + x_thread[4 * i + 1] * Dequantize<4>{}(ws[i] >> 4) + + x_thread[4 * i + 2] * Dequantize<4>{}(ws[i] >> 8) + + x_thread[4 * i + 3] * Dequantize<4>{}(ws[i] >> 12)); + } + } else { + for (int i = 0; i < N; i++) { + accum += x_thread[i] * Dequantize<8>{}(w[i]); + } } return scale * accum; } -template +template inline void qouter(const thread uint8_t* w, U x, U scale, thread U* result) { - for (int i = 0; i < (values_per_thread / 2); i++) { - result[2 * i] += x * scale * Dequantize<4>{}(w[i]); - result[2 * i + 1] += x * scale * Dequantize<4>{}(w[i] >> 4); + if constexpr (bits == 4) { + for (int i = 0; i < (values_per_thread / 2); i++) { + result[2 * i] += x * scale * Dequantize<4>{}(w[i]); + result[2 * i + 1] += x * scale * Dequantize<4>{}(w[i] >> 4); + } + } else { + for (int i = 0; i < values_per_thread; i++) { + result[i] += x * scale * Dequantize<8>{}(w[i]); + } } } -template -inline void dequantize( - const device uint8_t* w, - U scale, - threadgroup U* w_local, - const threadgroup U* lut) { - for (int i = 0; i < (N / 2); i++) { - w_local[2 * i] = scale * lut[w[i] & 0xf]; - w_local[2 * i + 1] = scale * lut[(w[i] >> 4) & 0xf]; +template +inline void dequantize(uint8_t w, U scale, threadgroup U* w_local) { + if constexpr (bits == 4) { + w_local[0] = scale * Dequantize<4, U>{}(w); + w_local[1] = scale * Dequantize<4, U>{}(w >> 4); + } else { + w_local[0] = scale * Dequantize<8, U>{}(w); } } @@ -135,21 +152,20 @@ template < short dst_ld, short reduction_dim, short tgp_size, - short group_size> + short group_size, + short bits> struct QuantizedBlockLoader { - static_assert( - BCOLS <= group_size, - "The group size should be larger than the columns"); - static_assert( - group_size % BCOLS == 0, - "The group size should be divisible by the columns"); - - MLX_MTL_CONST short pack_factor = get_pack_factor<8>(); + MLX_MTL_CONST short pack_factor = get_pack_factor<8, bits>(); MLX_MTL_CONST short bytes_per_pack = get_bytes_per_pack(); MLX_MTL_CONST short BCOLS_PACKED = BCOLS / pack_factor; MLX_MTL_CONST short n_reads = (BCOLS_PACKED * BROWS < tgp_size) ? 1 : (BCOLS_PACKED * BROWS) / tgp_size; - MLX_MTL_CONST short group_steps = group_size / BCOLS; + MLX_MTL_CONST short group_steps = group_size < BCOLS ? 1 : group_size / BCOLS; + MLX_MTL_CONST short scale_step = group_size < BCOLS ? BCOLS / group_size : 1; + + static_assert( + (n_reads * pack_factor) <= group_size, + "The number of reads per thread must be less than the group size."); const int src_ld; const int tile_stride; @@ -163,14 +179,12 @@ struct QuantizedBlockLoader { threadgroup T* dst; const device uint8_t* src; const device uint8_t* scales; - threadgroup T* lut; QuantizedBlockLoader( const device uint8_t* src_, const device uint8_t* scales_, const int src_ld_, threadgroup T* dst_, - threadgroup T* lut_, ushort simd_group_id [[simdgroup_index_in_threadgroup]], ushort simd_lane_id [[thread_index_in_simdgroup]]) : src_ld(src_ld_), @@ -185,23 +199,19 @@ struct QuantizedBlockLoader { dst(dst_ + bi * dst_ld + bj * pack_factor), src(src_ + bi * src_ld * bytes_per_pack / pack_factor + bj * bytes_per_pack), - scales(scales_ + bi * src_ld / group_size), - lut(lut_) { - if (simd_group_id == 0 && simd_lane_id < 16) { - lut[simd_lane_id] = static_cast(FP4_LUT[simd_lane_id]); - } - threadgroup_barrier(mem_flags::mem_threadgroup); - } + scales( + scales_ + bi * src_ld / group_size + + (bj * pack_factor) / group_size) {} void load_unsafe() const { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { return; } - T scale = dequantize_scale(*scales); + T scale = dequantize_scale(*scales); for (int i = 0; i < n_reads; i++) { - dequantize( - src + i * bytes_per_pack, scale, dst + i * pack_factor, lut); + dequantize( + src[i * bytes_per_pack], scale, dst + i * pack_factor); } } @@ -224,13 +234,10 @@ struct QuantizedBlockLoader { return; } - T scale = dequantize_scale(*scales); + T scale = dequantize_scale(*scales); for (int i = 0; i < n_reads; i++) { - dequantize( - (device uint8_t*)(src + i * bytes_per_pack), - scale, - dst + i * pack_factor, - lut); + dequantize( + src[i * bytes_per_pack], scale, dst + i * pack_factor); } } @@ -244,7 +251,7 @@ struct QuantizedBlockLoader { scales++; } } else { - scales++; + scales += scale_step; } } else { scales += group_stride; @@ -264,10 +271,13 @@ METAL_FUNC void fp_qmv_quad_impl( uint quad_gid [[quadgroup_index_in_threadgroup]], uint quad_lid [[thread_index_in_quadgroup]]) { constexpr int quads_per_simd = SIMD_SIZE / QUAD_SIZE; - constexpr int pack_factor = 8; + constexpr int pack_factor = get_pack_factor<32, bits>(); constexpr int values_per_thread = D / QUAD_SIZE; + constexpr int steps_per_thread = + values_per_thread < group_size ? 1 : values_per_thread / group_size; + constexpr int values_per_step = values_per_thread / steps_per_thread; constexpr int packs_per_thread = values_per_thread / pack_factor; - constexpr int scale_step_per_thread = group_size / values_per_thread; + constexpr int packs_per_step = values_per_step / pack_factor; constexpr int results_per_quadgroup = 8; typedef float U; @@ -281,7 +291,8 @@ METAL_FUNC void fp_qmv_quad_impl( const int out_row = tid.y * quads_per_simd * results_per_quadgroup + quad_gid; w += out_row * in_vec_size_w + quad_lid * packs_per_thread; - scales += out_row * in_vec_size_g + quad_lid / scale_step_per_thread; + scales += + out_row * in_vec_size_g + (quad_lid * values_per_thread) / group_size; x += tid.x * in_vec_size + quad_lid * values_per_thread; y += tid.x * out_vec_size + out_row; @@ -290,10 +301,15 @@ METAL_FUNC void fp_qmv_quad_impl( for (int row = 0; row < results_per_quadgroup; row++) { auto wl = (const device uint8_t*)(w + row * in_vec_size_w * quads_per_simd); const device uint8_t* sl = scales + row * in_vec_size_g * quads_per_simd; - - U s = dequantize_scale(sl[0]); - if (row * quads_per_simd + out_row < out_vec_size) { - result[row] += qdot(wl, x_thread, s); +#pragma unroll + for (int k = 0; k < steps_per_thread; ++k) { + U s = dequantize_scale(sl[0]); + if (row * quads_per_simd + out_row < out_vec_size) { + result[row] += qdot( + wl, x_thread + k * values_per_step, s); + } + sl++; + wl += (sizeof(uint32_t) / sizeof(uint8_t)) * packs_per_step; } } @@ -319,7 +335,7 @@ METAL_FUNC void fp_qmv_fast_impl( constexpr int packs_per_thread = 2; constexpr int num_simdgroups = 2; constexpr int results_per_simdgroup = 4; - constexpr int pack_factor = get_pack_factor<32>(); + constexpr int pack_factor = get_pack_factor<32, bits>(); constexpr int bytes_per_pack = get_bytes_per_pack<32>(); constexpr int values_per_thread = pack_factor * packs_per_thread; constexpr int block_size = values_per_thread * SIMD_SIZE; @@ -349,8 +365,8 @@ METAL_FUNC void fp_qmv_fast_impl( auto wl = (const device uint8_t*)(ws + row * in_vec_size_w); const device auto* sl = scales + row * in_vec_size_g; - U s = dequantize_scale(sl[0]); - result[row] += qdot(wl, x_thread, s); + U s = dequantize_scale(sl[0]); + result[row] += qdot(wl, x_thread, s); } ws += block_size * bytes_per_pack / pack_factor; @@ -380,7 +396,7 @@ METAL_FUNC void fp_qmv_impl( constexpr int num_simdgroups = 2; constexpr int results_per_simdgroup = 4; constexpr int packs_per_thread = 1; - constexpr int pack_factor = get_pack_factor<32>(); + constexpr int pack_factor = get_pack_factor<32, bits>(); constexpr int bytes_per_pack = get_bytes_per_pack<32>(); constexpr int values_per_thread = pack_factor * packs_per_thread; @@ -423,7 +439,7 @@ METAL_FUNC void fp_qmv_impl( const device auto* sl = scales + row * in_vec_size_g; uint8_t s = sl[0]; - result[row] += qdot(wl, x_thread, s); + result[row] += qdot(wl, x_thread, s); } ws += block_size * bytes_per_pack / pack_factor; @@ -441,8 +457,8 @@ METAL_FUNC void fp_qmv_impl( auto wl = (const device uint8_t*)(ws + row * in_vec_size_w); const device auto* sl = scales + row * in_vec_size_g; - U s = dequantize_scale(sl[0]); - result[row] += qdot(wl, x_thread, s); + U s = dequantize_scale(sl[0]); + result[row] += qdot(wl, x_thread, s); } } @@ -470,8 +486,8 @@ METAL_FUNC void fp_qmv_impl( auto wl = (const device uint8_t*)(ws + row * in_vec_size_w); const device auto* sl = scales + row * in_vec_size_g; - U s = dequantize_scale(sl[0]); - result[row] += qdot(wl, x_thread, s); + U s = dequantize_scale(sl[0]); + result[row] += qdot(wl, x_thread, s); } ws += block_size * bytes_per_pack / pack_factor; @@ -489,9 +505,9 @@ METAL_FUNC void fp_qmv_impl( auto wl = (const device uint8_t*)(ws + row * in_vec_size_w); const device auto* sl = scales + row * in_vec_size_g; - U s = dequantize_scale(sl[0]); + U s = dequantize_scale(sl[0]); result[row] += - qdot_safe(wl, x_thread, s, remaining); + qdot_safe(wl, x_thread, s, remaining); } } for (int row = 0; row < results_per_simdgroup; row++) { @@ -515,10 +531,10 @@ METAL_FUNC void fp_qvm_impl( uint simd_gid [[simdgroup_index_in_threadgroup]], uint simd_lid [[thread_index_in_simdgroup]]) { constexpr int num_simdgroups = 2; - constexpr int pack_factor = get_pack_factor<32>(); + constexpr int pack_factor = get_pack_factor<32, bits>(); constexpr int bytes_per_pack = get_bytes_per_pack(); - constexpr int tn = 32 / pack_factor; + constexpr int tn = group_size / pack_factor; constexpr int block_size = SIMD_SIZE; using W_T = uint32_t; @@ -537,6 +553,7 @@ METAL_FUNC void fp_qvm_impl( // 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; + // 32 * (tid.y * 2 + simd_gid) int out_col = pack_factor * tn * (tid.y * num_simdgroups + simd_gid); ws += out_col * bytes_per_pack / pack_factor + simd_lid * out_vec_size_w; scales += out_col / group_size + simd_lid * out_vec_size_g; @@ -552,9 +569,9 @@ METAL_FUNC void fp_qvm_impl( if (remaining == 0) { for (int i = 0; i < in_vec_size; i += block_size) { x_local = *x; - scale = dequantize_scale(*scales); + scale = dequantize_scale(*scales); w_local = *((device vec_w*)ws); - qouter( + qouter( (thread uint8_t*)&w_local, x_local, scale, result); x += block_size; @@ -564,10 +581,10 @@ METAL_FUNC void fp_qvm_impl( } else { for (int i = block_size; i < in_vec_size; i += block_size) { x_local = *x; - scale = dequantize_scale(*scales); + scale = dequantize_scale(*scales); w_local = *((device vec_w*)ws); - qouter( + qouter( (thread uint8_t*)&w_local, x_local, scale, result); x += block_size; @@ -576,13 +593,13 @@ METAL_FUNC void fp_qvm_impl( } if (static_cast(simd_lid) < remaining) { x_local = *x; - scale = dequantize_scale(*scales); + scale = dequantize_scale(*scales); w_local = *((device vec_w*)ws); } else { x_local = 0; scale = 0; } - qouter( + qouter( (thread uint8_t*)&w_local, x_local, scale, result); } @@ -622,8 +639,7 @@ METAL_FUNC void fp_qmm_t_impl( uint3 tid [[threadgroup_position_in_grid]], uint lid [[thread_index_in_threadgroup]], uint simd_gid [[simdgroup_index_in_threadgroup]], - uint simd_lid [[thread_index_in_simdgroup]], - threadgroup T* lut) { + uint simd_lid [[thread_index_in_simdgroup]]) { static_assert(BK >= SIMD_SIZE, "BK should be larger than SIMD_SIZE"); static_assert(BK % SIMD_SIZE == 0, "BK should be divisible by SIMD_SIZE"); @@ -631,7 +647,7 @@ METAL_FUNC void fp_qmm_t_impl( constexpr int WM = 2; constexpr int WN = 2; - constexpr int pack_factor = get_pack_factor<8>(); + constexpr int pack_factor = get_pack_factor<8, bits>(); constexpr int bytes_per_pack = get_bytes_per_pack(); constexpr int BK_padded = (BK + 16 / sizeof(T)); @@ -648,7 +664,8 @@ METAL_FUNC void fp_qmm_t_impl( BK_padded, 1, WM * WN * SIMD_SIZE, - group_size>; + group_size, + bits>; // Set the block const int K_w = K * bytes_per_pack / pack_factor; @@ -667,7 +684,7 @@ METAL_FUNC void fp_qmm_t_impl( const short num_els = min(BM, M - y_row); const short num_outs = min(BN, N - y_col); loader_x_t loader_x(x, K, Xs, simd_gid, simd_lid); - loader_w_t loader_w(wl, scales, K, Ws, lut, simd_gid, simd_lid); + loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid); mma_t mma_op(simd_gid, simd_lid); if (num_els < BM) { @@ -746,8 +763,7 @@ METAL_FUNC void fp_qmm_n_impl( uint3 tid [[threadgroup_position_in_grid]], uint lid [[thread_index_in_threadgroup]], uint simd_gid [[simdgroup_index_in_threadgroup]], - uint simd_lid [[thread_index_in_simdgroup]], - threadgroup T* lut) { + uint simd_lid [[thread_index_in_simdgroup]]) { static_assert(BK >= SIMD_SIZE, "BK should be larger than SIMD_SIZE"); static_assert(BK % SIMD_SIZE == 0, "BK should be divisible by SIMD_SIZE"); @@ -755,7 +771,7 @@ METAL_FUNC void fp_qmm_n_impl( constexpr int WM = 2; constexpr int WN = 2; - constexpr int pack_factor = get_pack_factor<8>(); + constexpr int pack_factor = get_pack_factor<8, bits>(); constexpr int bytes_per_pack = get_bytes_per_pack(); constexpr int BK_padded = (BK + 16 / sizeof(T)); @@ -773,7 +789,8 @@ METAL_FUNC void fp_qmm_n_impl( BN_padded, 0, WM * WN * SIMD_SIZE, - group_size>; + group_size, + bits>; auto wl = (const device uint8_t*)w; @@ -788,7 +805,7 @@ METAL_FUNC void fp_qmm_n_impl( // Make the x loader and mma operation const short num_els = min(BM, M - y_row); loader_x_t loader_x(x, K, Xs, simd_gid, simd_lid); - loader_w_t loader_w(wl, scales, N, Ws, lut, simd_gid, simd_lid); + loader_w_t loader_w(wl, scales, N, Ws, simd_gid, simd_lid); mma_t mma_op(simd_gid, simd_lid); if (num_els < BM) { @@ -1178,7 +1195,6 @@ template < threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BN * BK_padded]; - threadgroup T lut[16]; if (batched) { adjust_matrix_offsets( @@ -1197,7 +1213,7 @@ template < tid); } fp_qmm_t_impl( - w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid, lut); + w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); } template < @@ -1234,7 +1250,6 @@ template < threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BK * BN_padded]; - threadgroup T lut[16]; if (batched) { adjust_matrix_offsets( @@ -1254,7 +1269,7 @@ template < } fp_qmm_n_impl( - w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid, lut); + w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); } template @@ -1443,7 +1458,6 @@ template < threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BN * BK_padded]; - threadgroup T lut[16]; adjust_matrix_offsets( x, @@ -1466,7 +1480,7 @@ template < s_strides, tid); fp_qmm_t_impl( - w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid, lut); + w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); } template < @@ -1508,7 +1522,6 @@ template < threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BK * BN_padded]; - threadgroup T lut[16]; adjust_matrix_offsets( x, @@ -1531,7 +1544,7 @@ template < s_strides, tid); fp_qmm_n_impl( - w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid, lut); + w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); } template < @@ -1556,11 +1569,10 @@ template < uint3 tid [[threadgroup_position_in_grid]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint simd_lane_id [[thread_index_in_simdgroup]]) { - constexpr int pack_factor = get_pack_factor<8>(); + constexpr int pack_factor = get_pack_factor<8, bits>(); constexpr int bytes_per_pack = get_bytes_per_pack(); constexpr int BK_padded = (BK + 16 / sizeof(T)); constexpr int BN_padded = (BN + 16 / sizeof(T)); - threadgroup T lut[16]; using mma_t = mlx::steel::BlockMMA< T, @@ -1583,7 +1595,8 @@ template < transpose ? BK_padded : BN_padded, transpose, WM * WN * SIMD_SIZE, - group_size>; + group_size, + bits>; threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[transpose ? BN * BK_padded : BK * BN_padded]; @@ -1648,7 +1661,6 @@ template < scales + index * stride_s, transpose ? K : N, Ws, - lut, simd_group_id, simd_lane_id); diff --git a/mlx/backend/metal/kernels/fp_quantized.metal b/mlx/backend/metal/kernels/fp_quantized.metal index 091174d9..51ed8ca5 100644 --- a/mlx/backend/metal/kernels/fp_quantized.metal +++ b/mlx/backend/metal/kernels/fp_quantized.metal @@ -6,68 +6,68 @@ #include "mlx/backend/metal/kernels/quantized_utils.h" #include "mlx/backend/metal/kernels/fp_quantized.h" -#define instantiate_quantized(mode, name, type) \ +#define instantiate_quantized(mode, name, type, group_size, bits) \ instantiate_kernel( \ - #mode "_" #name "_" #type "_gs_32_b_4", \ + #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits, \ fp_ ## name, \ type, \ - 32, \ - 4) + group_size, \ + bits) -#define instantiate_quantized_batched(mode, name, type, batched) \ +#define instantiate_quantized_batched(mode, name, type, batched, group_size, bits) \ instantiate_kernel( \ - #mode "_" #name "_" #type "_gs_32_b_4_batch_" #batched, \ + #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_batch_" #batched, \ fp_ ## name, \ type, \ - 32, \ - 4, \ + group_size, \ + bits, \ batched) -#define instantiate_quantized_aligned(mode, name, type, aligned) \ +#define instantiate_quantized_aligned(mode, name, type, aligned, group_size, bits) \ instantiate_kernel( \ - #mode "_" #name "_" #type "_gs_32_b_4_alN_" #aligned, \ + #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_alN_" #aligned, \ fp_ ## name, \ type, \ - 32, \ - 4, \ + group_size, \ + bits, \ aligned) -#define instantiate_quantized_aligned_batched(mode, name, type, aligned, batched) \ +#define instantiate_quantized_aligned_batched(mode, name, type, aligned, batched, group_size, bits) \ instantiate_kernel( \ - #mode "_" #name "_" #type "_gs_32_b_4_alN_" #aligned "_batch_" #batched, \ + #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_alN_" #aligned "_batch_" #batched, \ fp_ ## name, \ type, \ - 32, \ - 4, \ + group_size, \ + bits, \ aligned, \ batched) -#define instantiate_quantized_quad(mode, name, type, D, batched) \ +#define instantiate_quantized_quad(mode, name, type, D, batched, group_size, bits) \ instantiate_kernel( \ - #mode "_" #name "_" #type "_gs_32_b_4_d_" #D "_batch_" #batched, \ + #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_d_" #D "_batch_" #batched, \ fp_ ## name, \ type, \ - 32, \ - 4, \ + group_size, \ + bits, \ D, \ batched) -#define instantiate_quantized_split_k(mode, name, type, split_k) \ +#define instantiate_quantized_split_k(mode, name, type, split_k, group_size, bits) \ instantiate_kernel( \ - #mode "_" #name "_" #type "_gs_32_b_4_spk_" #split_k, \ + #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_spk_" #split_k, \ fp_ ## name, \ type, \ - 32, \ - 4, \ + group_size, \ + bits, \ split_k) -#define instantiate_gather_qmm_rhs(func, name, type, bm, bn, bk, wm, wn, transpose) \ +#define instantiate_gather_qmm_rhs(func, name, type, bm, bn, bk, wm, wn, transpose, mode, group_size, bits) \ instantiate_kernel( \ - #name "_" #type "_gs_32_b_4_bm_" #bm "_bn_" #bn "_bk_" #bk "_wm_" #wm "_wn_" #wn, \ + #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_bm_" #bm "_bn_" #bn "_bk_" #bk "_wm_" #wm "_wn_" #wn, \ func, \ type, \ - 32, \ - 4, \ + group_size, \ + bits, \ bm, \ bn, \ bk, \ @@ -75,43 +75,43 @@ wn, \ transpose) -#define instantiate_quantized_batched_wrap(mode, name, type) \ - instantiate_quantized_batched(mode, name, type, 1) \ - instantiate_quantized_batched(mode, name, type, 0) +#define instantiate_quantized_batched_wrap(name, type, mode, group_size, bits) \ + instantiate_quantized_batched(mode, name, type, 1, group_size, bits) \ + instantiate_quantized_batched(mode, name, type, 0, group_size, bits) -#define instantiate_quantized_all_batched(type) \ - instantiate_quantized_batched_wrap(mxfp4, qmv_fast, type) \ - instantiate_quantized_batched_wrap(mxfp4, qmv, type) \ - instantiate_quantized_batched_wrap(mxfp4, qvm, type) \ - instantiate_quantized_batched_wrap(mxfp4, qmm_n, type) +#define instantiate_quantized_all_batched(type, mode, group_size, bits) \ + instantiate_quantized_batched_wrap(qmv_fast, type, mode, group_size, bits) \ + instantiate_quantized_batched_wrap(qmv, type, mode, group_size, bits) \ + instantiate_quantized_batched_wrap(qvm, type, mode, group_size, bits) \ + instantiate_quantized_batched_wrap(qmm_n, type, mode, group_size, bits) -#define instantiate_quantized_all_single(type) \ - instantiate_quantized(mxfp4, gather_qmv_fast, type) \ - instantiate_quantized(mxfp4, gather_qmv, type) \ - instantiate_quantized(mxfp4, gather_qvm, type) \ - instantiate_quantized(mxfp4, gather_qmm_n, type) +#define instantiate_quantized_all_single(type, mode, group_size, bits) \ + instantiate_quantized(mode, gather_qmv_fast, type, group_size, bits) \ + instantiate_quantized(mode, gather_qmv, type, group_size, bits) \ + instantiate_quantized(mode, gather_qvm, type, group_size, bits) \ + instantiate_quantized(mode, gather_qmm_n, type, group_size, bits) -#define instantiate_quantized_all_aligned(type) \ - instantiate_quantized_aligned(mxfp4, gather_qmm_t, type, true) \ - instantiate_quantized_aligned(mxfp4, gather_qmm_t, type, false) \ - instantiate_quantized_aligned_batched(mxfp4, qmm_t, type, true, 1) \ - instantiate_quantized_aligned_batched(mxfp4, qmm_t, type, true, 0) \ - instantiate_quantized_aligned_batched(mxfp4, qmm_t, type, false, 1) \ - instantiate_quantized_aligned_batched(mxfp4, qmm_t, type, false, 0) +#define instantiate_quantized_all_aligned(type, mode, group_size, bits) \ + instantiate_quantized_aligned(mode, gather_qmm_t, type, true, group_size, bits) \ + instantiate_quantized_aligned(mode, gather_qmm_t, type, false, group_size, bits) \ + instantiate_quantized_aligned_batched(mode, qmm_t, type, true, 1, group_size, bits) \ + instantiate_quantized_aligned_batched(mode, qmm_t, type, true, 0, group_size, bits) \ + instantiate_quantized_aligned_batched(mode, qmm_t, type, false, 1, group_size, bits) \ + instantiate_quantized_aligned_batched(mode, qmm_t, type, false, 0, group_size, bits) -#define instantiate_quantized_all_quad(type) \ - instantiate_quantized_quad(mxfp4, qmv_quad, type, 64, 1) \ - instantiate_quantized_quad(mxfp4, qmv_quad, type, 64, 0) \ - instantiate_quantized_quad(mxfp4, qmv_quad, type, 128, 1) \ - instantiate_quantized_quad(mxfp4, qmv_quad, type, 128, 0) +#define instantiate_quantized_all_quad(type, mode, group_size, bits) \ + instantiate_quantized_quad(mode, qmv_quad, type, 64, 1, group_size, bits) \ + instantiate_quantized_quad(mode, qmv_quad, type, 64, 0, group_size, bits) \ + instantiate_quantized_quad(mode, qmv_quad, type, 128, 1, group_size, bits) \ + instantiate_quantized_quad(mode, qmv_quad, type, 128, 0, group_size, bits) -#define instantiate_quantized_all_splitk(type) \ - instantiate_quantized_split_k(mxfp4, qvm_split_k, type, 8) \ - instantiate_quantized_split_k(mxfp4, qvm_split_k, type, 32) +#define instantiate_quantized_all_splitk(type, mode, group_size, bits) \ + instantiate_quantized_split_k(mode, qvm_split_k, type, 8, group_size, bits) \ + instantiate_quantized_split_k(mode, qvm_split_k, type, 32, group_size, bits) -#define instantiate_quantized_all_rhs(type) \ - instantiate_gather_qmm_rhs(fp_gather_qmm_rhs, mxfp4_gather_qmm_rhs_nt, type, 16, 32, 32, 1, 2, true) \ - instantiate_gather_qmm_rhs(fp_gather_qmm_rhs, mxfp4_gather_qmm_rhs_nn, type, 16, 32, 32, 1, 2, false) +#define instantiate_quantized_all_rhs(type, mode, group_size, bits) \ + instantiate_gather_qmm_rhs(fp_gather_qmm_rhs, gather_qmm_rhs_nt, type, 16, 32, 32, 1, 2, true, mode, group_size, bits) \ + instantiate_gather_qmm_rhs(fp_gather_qmm_rhs, gather_qmm_rhs_nn, type, 16, 32, 32, 1, 2, false, mode, group_size, bits) #define instantiate_quantize_dequantize(type, mode, group_size, bits) \ instantiate_kernel( \ @@ -127,19 +127,19 @@ group_size, \ bits) -#define instantiate_quantize_dequantize_modes(type) \ - instantiate_quantize_dequantize(type, mxfp4, 32, 4) \ - instantiate_quantize_dequantize(type, nvfp4, 16, 4) \ - instantiate_quantize_dequantize(type, mxfp8, 32, 8) +#define instantiate_quantized_modes(type, mode, group_size, bits) \ + instantiate_quantized_all_batched(type, mode, group_size, bits) \ + instantiate_quantized_all_single(type, mode, group_size, bits) \ + instantiate_quantized_all_quad(type, mode, group_size, bits) \ + instantiate_quantized_all_splitk(type, mode, group_size, bits) \ + instantiate_quantized_all_aligned(type, mode, group_size, bits) \ + instantiate_quantized_all_rhs(type, mode, group_size, bits) \ + instantiate_quantize_dequantize(type, mode, group_size, bits) #define instantiate_quantized_types(type) \ - instantiate_quantized_all_batched(type) \ - instantiate_quantized_all_quad(type) \ - instantiate_quantized_all_splitk(type) \ - instantiate_quantized_all_single(type) \ - instantiate_quantized_all_aligned(type) \ - instantiate_quantized_all_rhs(type) \ - instantiate_quantize_dequantize_modes(type) + instantiate_quantized_modes(type, nvfp4, 16, 4) \ + instantiate_quantized_modes(type, mxfp8, 32, 8) \ + instantiate_quantized_modes(type, mxfp4, 32, 4) instantiate_quantized_types(float) instantiate_quantized_types(bfloat16_t) diff --git a/mlx/backend/metal/kernels/fp_quantized_nax.h b/mlx/backend/metal/kernels/fp_quantized_nax.h index 28c4ec82..80e1c4c2 100644 --- a/mlx/backend/metal/kernels/fp_quantized_nax.h +++ b/mlx/backend/metal/kernels/fp_quantized_nax.h @@ -17,9 +17,9 @@ using namespace metal; MLX_MTL_CONST int SIMD_SIZE = 32; MLX_MTL_CONST int QUAD_SIZE = 4; -template +template inline constexpr short get_pack_factor() { - return wsize / 4; + return wsize / bits; } template @@ -27,15 +27,20 @@ inline constexpr short get_bytes_per_pack() { return wsize / 8; } -template +template static inline T dequantize_scale(uint8_t s) { - return T(*(thread fp8_e8m0*)(&s)); + if constexpr (group_size == 16) { + // Use nv scale + return T(*(thread fp8_e4m3*)(&s)); + } else { + return T(*(thread fp8_e8m0*)(&s)); + } } template struct Quantize { uint8_t operator()(float x) { - if constexpr (bits == 8) { + if (bits == 8) { return fp8_e4m3(x).bits; } else { return fp4_e2m1(x).bits; @@ -43,26 +48,24 @@ struct Quantize { } }; -template +template struct Dequantize { - float operator()(uint8_t x) { + U operator()(uint8_t x) { if constexpr (bits == 8) { - return float(*(thread fp8_e4m3*)(&x)); + return U(*(thread fp8_e4m3*)(&x)); } else { - return float(*(thread fp4_e2m1*)(&x)); + return U(*(thread fp4_e2m1*)(&x)); } } }; -template -inline void dequantize( - const device uint8_t* w, - U scale, - threadgroup U* w_local, - const threadgroup U* lut) { - for (int i = 0; i < (N / 2); i++) { - w_local[2 * i] = scale * lut[w[i] & 0xf]; - w_local[2 * i + 1] = scale * lut[(w[i] >> 4) & 0xf]; +template +inline void dequantize(uint8_t w, U scale, threadgroup U* w_local) { + if constexpr (bits == 4) { + w_local[0] = scale * Dequantize<4, U>{}(w); + w_local[1] = scale * Dequantize<4, U>{}(w >> 4); + } else { + w_local[0] = scale * Dequantize<8, U>{}(w); } } @@ -73,22 +76,21 @@ template < short dst_ld, short reduction_dim, short tgp_size, - short group_size> + short group_size, + short bits> struct QuantizedBlockLoader { - static_assert( - BCOLS % group_size == 0, - "The group size should be divisible by the columns"); - - MLX_MTL_CONST short pack_factor = get_pack_factor<8>(); + MLX_MTL_CONST short pack_factor = get_pack_factor<8, bits>(); MLX_MTL_CONST short bytes_per_pack = get_bytes_per_pack(); MLX_MTL_CONST short BCOLS_PACKED = BCOLS / pack_factor; MLX_MTL_CONST short n_reads = (BCOLS_PACKED * BROWS < tgp_size) ? 1 : (BCOLS_PACKED * BROWS) / tgp_size; - MLX_MTL_CONST short n_groups = BCOLS / group_size; - static_assert( - (BCOLS_PACKED / n_reads) == n_groups, - "Other configurations are not yet supported"); + MLX_MTL_CONST short n_reads_per_scale = (n_reads * pack_factor) <= group_size + ? n_reads + : (group_size / pack_factor); + MLX_MTL_CONST short n_steps_per_read = n_reads / n_reads_per_scale; + + MLX_MTL_CONST short n_groups = BCOLS / group_size; const int src_ld; const int tile_stride; @@ -103,14 +105,12 @@ struct QuantizedBlockLoader { threadgroup T* dst; const device uint8_t* src; const device uint8_t* scales; - threadgroup T* lut; QuantizedBlockLoader( const device uint8_t* src_, const device uint8_t* scales_, const int src_ld_, threadgroup T* dst_, - threadgroup T* lut_, ushort simd_group_id [[simdgroup_index_in_threadgroup]], ushort simd_lane_id [[thread_index_in_simdgroup]]) : src_ld(src_ld_), @@ -125,23 +125,21 @@ struct QuantizedBlockLoader { dst(dst_ + bi * dst_ld + bj * pack_factor), src(src_ + bi * src_ld * bytes_per_pack / pack_factor + bj * bytes_per_pack), - scales(scales_ + bi * src_ld / group_size + group_id), - lut(lut_) { - if (simd_group_id == 0 && simd_lane_id < 16) { - lut[simd_lane_id] = static_cast(FP4_LUT[simd_lane_id]); - } - threadgroup_barrier(mem_flags::mem_threadgroup); - } + scales(scales_ + bi * src_ld / group_size + group_id) {} void load_unsafe() const { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { return; } - T scale = dequantize_scale(*scales); - for (int i = 0; i < n_reads; i++) { - dequantize( - src + i * bytes_per_pack, scale, dst + i * pack_factor, lut); + int k = 0; + for (int i = 0; i < n_steps_per_read; i++) { + T scale = dequantize_scale(scales[i]); + for (int j = 0; j < n_reads_per_scale; j++) { + dequantize( + src[k * bytes_per_pack], scale, dst + k * pack_factor); + k++; + } } } @@ -164,28 +162,21 @@ struct QuantizedBlockLoader { return; } - T scale = dequantize_scale(*scales); - for (int i = 0; i < n_reads; i++) { - dequantize( - (device uint8_t*)(src + i * bytes_per_pack), - scale, - dst + i * pack_factor, - lut); + int k = 0; + for (int i = 0; i < n_steps_per_read; i++) { + T scale = dequantize_scale(scales[i]); + for (int j = 0; j < n_reads_per_scale; j++) { + dequantize( + src[k * bytes_per_pack], scale, dst + k * pack_factor); + k++; + } } } void next() { src += tile_stride; if (reduction_dim == 1) { - // if (group_steps > 1) { - // group_step_cnt++; - // if (group_step_cnt == group_steps) { - // group_step_cnt = 0; - // scales++; - // } - // } else { scales += n_groups; - // } } else { scales += n_groups * group_stride; } @@ -217,14 +208,13 @@ METAL_FUNC void fp_qmm_t_impl( uint3 tid [[threadgroup_position_in_grid]], uint lid [[thread_index_in_threadgroup]], uint simd_gid [[simdgroup_index_in_threadgroup]], - uint simd_lid [[thread_index_in_simdgroup]], - threadgroup Wtype* lut) { + uint simd_lid [[thread_index_in_simdgroup]]) { static_assert(BK >= SIMD_SIZE, "BK should be larger than SIMD_SIZE"); static_assert(BK % SIMD_SIZE == 0, "BK should be divisible by SIMD_SIZE"); (void)lid; - constexpr int pack_factor = get_pack_factor<8>(); + constexpr int pack_factor = get_pack_factor<8, bits>(); constexpr int bytes_per_pack = get_bytes_per_pack(); constexpr int BK_padded = (BK + 16 / sizeof(Wtype)); @@ -237,7 +227,8 @@ METAL_FUNC void fp_qmm_t_impl( BK_padded, 1, WM * WN * SIMD_SIZE, - group_size>; + group_size, + bits>; // Set the block const int K_w = K * bytes_per_pack / pack_factor; @@ -253,7 +244,7 @@ METAL_FUNC void fp_qmm_t_impl( y += y_row * static_cast(N) + y_col; // Make the weight loader - loader_w_t loader_w(wl, scales, K, Ws, lut, simd_gid, simd_lid); + loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid); constexpr short UM = 16; constexpr short UN = 32; @@ -369,15 +360,14 @@ METAL_FUNC void fp_qmm_n_impl( uint3 tid [[threadgroup_position_in_grid]], uint lid [[thread_index_in_threadgroup]], uint simd_gid [[simdgroup_index_in_threadgroup]], - uint simd_lid [[thread_index_in_simdgroup]], - threadgroup Wtype* lut) { + uint simd_lid [[thread_index_in_simdgroup]]) { static_assert(BK >= SIMD_SIZE, "BK should be larger than SIMD_SIZE"); static_assert(BK % SIMD_SIZE == 0, "BK should be divisible by SIMD_SIZE"); (void)lid; (void)M; - constexpr int pack_factor = get_pack_factor<8>(); + constexpr int pack_factor = get_pack_factor<8, bits>(); constexpr int bytes_per_pack = get_bytes_per_pack(); constexpr int BN_padded = (BN + 16 / sizeof(T)); @@ -389,7 +379,8 @@ METAL_FUNC void fp_qmm_n_impl( BN_padded, 0, WM * WN * SIMD_SIZE, - group_size>; + group_size, + bits>; // Set the block const int K_w = K * bytes_per_pack / pack_factor; @@ -407,7 +398,7 @@ METAL_FUNC void fp_qmm_n_impl( // Make the x loader and mma operation // const short num_els = min(BM, M - y_row); // const short num_outs = min(BN, N - y_col); - loader_w_t loader_w(wl, scales, K, Ws, lut, simd_gid, simd_lid); + loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid); constexpr short UM = 16; constexpr short UN = 32; @@ -596,7 +587,6 @@ template < constexpr int BK_padded = (BK + 16 / sizeof(Wtype)); threadgroup Wtype Ws[BN * BK_padded]; - threadgroup Wtype lut[16]; if (batched) { adjust_matrix_offsets( @@ -615,7 +605,7 @@ template < tid); } fp_qmm_t_impl( - w, scales, x, y, Ws, K, N, M, tid, lid, simd_gid, simd_lid, lut); + w, scales, x, y, Ws, K, N, M, tid, lid, simd_gid, simd_lid); } template < @@ -655,7 +645,6 @@ template < threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BK * BN_padded]; - threadgroup T lut[16]; if (batched) { adjust_matrix_offsets( @@ -675,7 +664,7 @@ template < } fp_qmm_n_impl( - w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid, lut); + w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); } template < @@ -719,7 +708,6 @@ template < constexpr int BK_padded = (BK + 16 / sizeof(Wtype)); threadgroup Wtype Ws[BN * BK_padded]; - threadgroup Wtype lut[16]; adjust_matrix_offsets( x, @@ -742,7 +730,7 @@ template < s_strides, tid); fp_qmm_t_impl( - w, scales, x, y, Ws, K, N, M, tid, lid, simd_gid, simd_lid, lut); + w, scales, x, y, Ws, K, N, M, tid, lid, simd_gid, simd_lid); } template < @@ -787,7 +775,6 @@ template < threadgroup T Xs[BM * BK_padded]; threadgroup T Ws[BK * BN_padded]; - threadgroup T lut[16]; adjust_matrix_offsets( x, @@ -810,7 +797,7 @@ template < s_strides, tid); fp_qmm_n_impl( - w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid, lut); + w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); } template < @@ -836,13 +823,11 @@ template < uint3 tid [[threadgroup_position_in_grid]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint simd_lane_id [[thread_index_in_simdgroup]]) { - constexpr int pack_factor = get_pack_factor<8>(); + constexpr int pack_factor = get_pack_factor<8, bits>(); constexpr int bytes_per_pack = get_bytes_per_pack(); constexpr int BK_padded = (BK + 16 / sizeof(Wtype)); constexpr int BN_padded = (BN + 16 / sizeof(Wtype)); - threadgroup Wtype lut[16]; - using loader_w_t = QuantizedBlockLoader< Wtype, transpose ? BN : BK, @@ -850,7 +835,8 @@ template < transpose ? BK_padded : BN_padded, transpose, WM * WN * SIMD_SIZE, - group_size>; + group_size, + bits>; threadgroup Wtype Ws[transpose ? BN * BK_padded : BK * BN_padded]; @@ -947,7 +933,6 @@ template < scales + index * stride_s, transpose ? K : N, Ws, - lut, simd_group_id, simd_lane_id); diff --git a/mlx/backend/metal/kernels/fp_quantized_nax.metal b/mlx/backend/metal/kernels/fp_quantized_nax.metal index bd2df2b7..4d65a384 100644 --- a/mlx/backend/metal/kernels/fp_quantized_nax.metal +++ b/mlx/backend/metal/kernels/fp_quantized_nax.metal @@ -8,41 +8,41 @@ #include "mlx/backend/metal/kernels/fp_quantized_nax.h" -#define instantiate_quantized_batched(mode, name, type, bm, bn, bk, wm, wn, batched) \ +#define instantiate_quantized_batched(mode, name, type, bm, bn, bk, wm, wn, batched, group_size, bits) \ instantiate_kernel( \ - #mode "_" #name "_" #type "_gs_32_b_4_bm" #bm "_bn" #bn "_bk" #bk "_wm" #wm "_wn" #wn "_batch_" #batched, \ + #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_bm" #bm "_bn" #bn "_bk" #bk "_wm" #wm "_wn" #wn "_batch_" #batched, \ fp_ ## name, \ type, \ - 32, \ - 4, \ + group_size, \ + bits, \ batched) -#define instantiate_quantized_aligned(mode, name, type, bm, bn, bk, wm, wn, aligned) \ +#define instantiate_quantized_aligned(mode, name, type, bm, bn, bk, wm, wn, aligned, group_size, bits) \ instantiate_kernel( \ - #mode "_" #name "_" #type "_gs_32_b_4_bm" #bm "_bn" #bn "_bk" #bk "_wm" #wm "_wn" #wn "_alN_" #aligned, \ + #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_bm" #bm "_bn" #bn "_bk" #bk "_wm" #wm "_wn" #wn "_alN_" #aligned, \ fp_ ## name, \ type, \ - 32, \ - 4, \ + group_size, \ + bits, \ aligned) -#define instantiate_quantized_aligned_batched(mode, name, type, bm, bn, bk, wm, wn, aligned, batched) \ +#define instantiate_quantized_aligned_batched(mode, name, type, bm, bn, bk, wm, wn, aligned, batched, group_size, bits) \ instantiate_kernel( \ - #mode "_" #name "_" #type "_gs_32_b_4_bm" #bm "_bn" #bn "_bk" #bk "_wm" #wm "_wn" #wn "_alN_" #aligned "_batch_" #batched, \ + #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_bm" #bm "_bn" #bn "_bk" #bk "_wm" #wm "_wn" #wn "_alN_" #aligned "_batch_" #batched, \ fp_ ## name, \ type, \ - 32, \ - 4, \ + group_size, \ + bits, \ aligned, \ batched) -#define instantiate_gather_qmm_rhs(func, name, type, bm, bn, bk, wm, wn, transpose) \ +#define instantiate_gather_qmm_rhs(func, name, type, bm, bn, bk, wm, wn, transpose, mode, group_size, bits) \ instantiate_kernel( \ - #name "_" #type "_gs_32_b_4_bm_" #bm "_bn_" #bn "_bk_" #bk "_wm_" #wm "_wn_" #wn, \ + #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_bm_" #bm "_bn_" #bn "_bk_" #bk "_wm_" #wm "_wn_" #wn, \ func, \ type, \ - 32, \ - 4, \ + group_size, \ + bits, \ bm, \ bn, \ bk, \ @@ -51,22 +51,27 @@ transpose) -#define instantiate_quantized_all_aligned(type) \ - instantiate_quantized_aligned(mxfp4, gather_qmm_t_nax, type, 64, 64, 64, 2, 2, true) \ - instantiate_quantized_aligned(mxfp4, gather_qmm_t_nax, type, 64, 64, 64, 2, 2, false) \ - instantiate_quantized_aligned_batched(mxfp4, qmm_t_nax, type, 64, 64, 64, 2, 2, true, 1) \ - instantiate_quantized_aligned_batched(mxfp4, qmm_t_nax, type, 64, 64, 64, 2, 2, true, 0) \ - instantiate_quantized_aligned_batched(mxfp4, qmm_t_nax, type, 64, 64, 64, 2, 2, false, 1) \ - instantiate_quantized_aligned_batched(mxfp4, qmm_t_nax, type, 64, 64, 64, 2, 2, false, 0) +#define instantiate_quantized_all_aligned(type, mode, group_size, bits) \ + instantiate_quantized_aligned(mode, gather_qmm_t_nax, type, 64, 64, 64, 2, 2, true, group_size, bits) \ + instantiate_quantized_aligned(mode, gather_qmm_t_nax, type, 64, 64, 64, 2, 2, false, group_size, bits) \ + instantiate_quantized_aligned_batched(mode, qmm_t_nax, type, 64, 64, 64, 2, 2, true, 1, group_size, bits) \ + instantiate_quantized_aligned_batched(mode, qmm_t_nax, type, 64, 64, 64, 2, 2, true, 0, group_size, bits) \ + instantiate_quantized_aligned_batched(mode, qmm_t_nax, type, 64, 64, 64, 2, 2, false, 1, group_size, bits) \ + instantiate_quantized_aligned_batched(mode, qmm_t_nax, type, 64, 64, 64, 2, 2, false, 0, group_size, bits) -#define instantiate_quantized_all_rhs(type) \ - instantiate_gather_qmm_rhs(fp_gather_qmm_rhs_nax, mxfp4_gather_qmm_rhs_nax_nt, type, 64, 64, 64, 2, 2, true) \ - instantiate_gather_qmm_rhs(fp_gather_qmm_rhs_nax, mxfp4_gather_qmm_rhs_nax_nn, type, 64, 64, 64, 2, 2, false) +#define instantiate_quantized_all_rhs(type, mode, group_size, bits) \ + instantiate_gather_qmm_rhs(fp_gather_qmm_rhs_nax, gather_qmm_rhs_nax_nt, type, 64, 64, 64, 2, 2, true, mode, group_size, bits) \ + instantiate_gather_qmm_rhs(fp_gather_qmm_rhs_nax, gather_qmm_rhs_nax_nn, type, 64, 64, 64, 2, 2, false, mode, group_size, bits) + +#define instantiate_quantized_modes(type, mode, group_size, bits) \ + instantiate_quantized_all_aligned(type, mode, group_size, bits) \ + instantiate_quantized_all_rhs(type, mode, group_size, bits) #define instantiate_quantized_types(type) \ - instantiate_quantized_all_aligned(type) \ - instantiate_quantized_all_rhs(type) + instantiate_quantized_modes(type, nvfp4, 16, 4) \ + instantiate_quantized_modes(type, mxfp8, 32, 8) \ + instantiate_quantized_modes(type, mxfp4, 32, 4) instantiate_quantized_types(float) instantiate_quantized_types(bfloat16_t) diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index 723bccc6..badc8883 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -315,9 +315,10 @@ void qvm_split_k( int B = out.size() / M / N; B *= split_k; - int bn = 64; - int bk = 32; - MTL::Size group_dims = MTL::Size(bk, 2, 1); + constexpr int num_simdgroups = 2; + constexpr int bk = 32; + int bn = std::min(group_size, 32) * num_simdgroups; + MTL::Size group_dims = MTL::Size(bk, num_simdgroups, 1); MTL::Size grid_dims = MTL::Size(M, N / bn, B); auto x_shape = x.shape(); @@ -431,9 +432,10 @@ void qvm( const std::string& mode) { int B = out.size() / M / N; - int bn = 64; - int bk = 32; - MTL::Size group_dims(bk, 2, 1); + constexpr int num_simdgroups = 2; + constexpr int bk = 32; + int bn = std::min(group_size, 32) * num_simdgroups; + MTL::Size group_dims(bk, num_simdgroups, 1); MTL::Size grid_dims(M, (N + bn - 1) / bn, B); std::string kname; @@ -944,9 +946,10 @@ void gather_qvm( const std::string& mode) { int B = out.size() / M / N; - int bn = 64; - int bk = 32; - MTL::Size group_dims(bk, 2, 1); + constexpr int num_simdgroups = 2; + constexpr int bk = 32; + int bn = std::min(group_size, 32) * num_simdgroups; + MTL::Size group_dims(bk, num_simdgroups, 1); MTL::Size grid_dims(M, (N + bn - 1) / bn, B); std::string kname; diff --git a/mlx/primitives.cpp b/mlx/primitives.cpp index a48f8bad..fd3bda37 100644 --- a/mlx/primitives.cpp +++ b/mlx/primitives.cpp @@ -3397,9 +3397,11 @@ std::vector QuantizedMatmul::vjp( throw std::runtime_error( "[QuantizedMatmul::vjp] no gradient wrt the quantized weights."); } else { - if (mode_ == QuantizationMode::Mxfp4) { - throw std::invalid_argument( - "[QuantizedMatmul::vjp] no gradient wrt scales with mxfp4 quantization."); + if (mode_ != QuantizationMode::Affine) { + std::ostringstream msg; + msg << "[QuantizedMatmul::vjp] no gradient wrt scales in " + << quantization_mode_to_string(mode_) << " quantization."; + throw std::invalid_argument(msg.str()); } if (!dsb) { int ndim = primals[1].ndim(); @@ -3606,9 +3608,11 @@ std::vector GatherQMM::vjp( throw std::runtime_error( "[GatherQMM::vjp] no gradient wrt the quantized weights."); } else { - if (mode_ == QuantizationMode::Mxfp4) { - throw std::invalid_argument( - "[GatherQMM::vjp] no gradient wrt scales with mxfp4 quantization."); + if (mode_ != QuantizationMode::Affine) { + std::ostringstream msg; + msg << "[GatherQMM::vjp] no gradient wrt scales in " + << quantization_mode_to_string(mode_) << " quantization."; + throw std::invalid_argument(msg.str()); } if (!dsb) { diff --git a/python/mlx/nn/layers/embedding.py b/python/mlx/nn/layers/embedding.py index 1edf7e3a..891627be 100644 --- a/python/mlx/nn/layers/embedding.py +++ b/python/mlx/nn/layers/embedding.py @@ -1,6 +1,7 @@ # Copyright © 2023-2024 Apple Inc. import math +from typing import Optional import mlx.core as mx from mlx.nn.layers.base import Module @@ -39,6 +40,11 @@ class Embedding(Module): """ return x @ self.weight.T - def to_quantized(self, group_size: int = 64, bits: int = 4, mode: str = "affine"): + def to_quantized( + self, + group_size: Optional[int] = None, + bits: Optional[int] = None, + mode: str = "affine", + ): """Return a :obj:`QuantizedEmbedding` layer that approximates this embedding layer.""" return QuantizedEmbedding.from_embedding(self, group_size, bits, mode) diff --git a/python/mlx/nn/layers/linear.py b/python/mlx/nn/layers/linear.py index 84a4d832..8f00ff18 100644 --- a/python/mlx/nn/layers/linear.py +++ b/python/mlx/nn/layers/linear.py @@ -1,7 +1,7 @@ # Copyright © 2023 Apple Inc. import math -from typing import Any +from typing import Any, Optional import mlx.core as mx from mlx.nn.layers.base import Module @@ -70,7 +70,12 @@ class Linear(Module): x = x @ self["weight"].T return x - def to_quantized(self, group_size: int = 64, bits: int = 4, mode: str = "affine"): + def to_quantized( + self, + group_size: Optional[int] = None, + bits: Optional[int] = None, + mode: str = "affine", + ): """Return a :obj:`QuantizedLinear` layer that approximates this layer.""" return QuantizedLinear.from_linear(self, group_size, bits, mode) diff --git a/python/mlx/nn/layers/quantized.py b/python/mlx/nn/layers/quantized.py index c308e884..f762847e 100644 --- a/python/mlx/nn/layers/quantized.py +++ b/python/mlx/nn/layers/quantized.py @@ -8,10 +8,21 @@ from mlx.nn.layers.base import Module from mlx.utils import tree_map_with_path +def _defaults_for_mode(mode, group_size, bits): + mode_defaults = { + "affine": (64, 4), + "mxfp4": (32, 4), + "nvfp4": (16, 4), + "mxfp8": (32, 8), + } + default_group_size, default_bits = mode_defaults[mode] + return group_size or default_group_size, bits or default_bits + + def quantize( model: Module, - group_size: int = 64, - bits: int = 4, + group_size: int = None, + bits: int = None, *, mode: str = "affine", class_predicate: Optional[Callable[[str, Module], Union[bool, dict]]] = None, @@ -24,10 +35,10 @@ def quantize( Args: model (mlx.nn.Module): The model whose leaf modules may be quantized. - group_size (int): The quantization group size (see - :func:`mlx.core.quantize`). Default: ``64``. - bits (int): The number of bits per parameter (see - :func:`mlx.core.quantize`). Default: ``4``. + group_size (Optional[int]): The quantization group size (see + :func:`mlx.core.quantize`). Default: ``None``. + bits (Optional[int]): The number of bits per parameter (see + :func:`mlx.core.quantize`). Default: ``None``. mode (str): The quantization method to use (see :func:`mlx.core.quantize`). Default: ``"affine"``. class_predicate (Optional[Callable]): A callable which receives the @@ -72,10 +83,10 @@ class QuantizedEmbedding(Module): num_embeddings (int): How many possible discrete tokens can we embed. Usually called the vocabulary size. dims (int): The dimensionality of the embeddings. - group_size (int, optional): The group size to use for the quantized - weight. See :func:`~mlx.core.quantize`. Default: ``64``. - bits (int, optional): The bit width to use for the quantized weight. - See :func:`~mlx.core.quantize`. Default: ``4``. + group_size (Optional[int]): The group size to use for the quantized + weight. See :func:`~mlx.core.quantize`. Default: ``None``. + bits (Optional[int]): The bit width to use for the quantized weight. + See :func:`~mlx.core.quantize`. Default: ``None``. mode (str): The quantization method to use (see :func:`mlx.core.quantize`). Default: ``"affine"``. """ @@ -84,15 +95,14 @@ class QuantizedEmbedding(Module): self, num_embeddings: int, dims: int, - group_size: int = 64, - bits: int = 4, + group_size: int = None, + bits: int = None, mode: str = "affine", ): super().__init__() # Quantization config - self.group_size = group_size - self.bits = bits + self.group_size, self.bits = _defaults_for_mode(mode, group_size, bits) self.mode = mode # Initialize the quantized weight @@ -147,8 +157,8 @@ class QuantizedEmbedding(Module): def from_embedding( cls, embedding_layer: Module, - group_size: int = 64, - bits: int = 4, + group_size: int = None, + bits: int = None, mode: str = "affine", ): """Create a :obj:`QuantizedEmbedding` layer from an :obj:`Embedding` layer.""" @@ -179,10 +189,10 @@ class QuantizedLinear(Module): output_dims (int): The dimensionality of the output features. bias (bool, optional): If set to ``False`` then the layer will not use a bias. Default: ``True``. - group_size (int, optional): The group size to use for the quantized - weight. See :func:`~mlx.core.quantize`. Default: ``64``. - bits (int, optional): The bit width to use for the quantized weight. - See :func:`~mlx.core.quantize`. Default: ``4``. + group_size (Optional[int]): The group size to use for the quantized + weight. See :func:`~mlx.core.quantize`. Default: ``None``. + bits (Optional[int]): The bit width to use for the quantized weight. + See :func:`~mlx.core.quantize`. Default: ``None``. mode (str): The quantization method to use (see :func:`mlx.core.quantize`). Default: ``"affine"``. """ @@ -192,15 +202,14 @@ class QuantizedLinear(Module): input_dims: int, output_dims: int, bias: bool = True, - group_size: int = 64, - bits: int = 4, + group_size: int = None, + bits: int = None, mode: str = "affine", ): super().__init__() # Quantization config - self.group_size = group_size - self.bits = bits + self.group_size, self.bits = _defaults_for_mode(mode, group_size, bits) self.mode = mode # Initialize the quantized weight @@ -249,8 +258,8 @@ class QuantizedLinear(Module): def from_linear( cls, linear_layer: Module, - group_size: int = 64, - bits: int = 4, + group_size: int = None, + bits: int = None, mode: str = "affine", ): """Create a :obj:`QuantizedLinear` layer from a :obj:`Linear` layer.""" diff --git a/python/tests/cuda_skip.py b/python/tests/cuda_skip.py index 2d5ef804..f9e6fe8e 100644 --- a/python/tests/cuda_skip.py +++ b/python/tests/cuda_skip.py @@ -48,8 +48,8 @@ cuda_skip = { "TestQuantized.test_qmm_shapes", "TestQuantized.test_qmm_vjp", "TestQuantized.test_qmv", - "TestQuantized.test_mxfp4_qmv", - "TestQuantized.test_mxfp4_qvm", + "TestQuantized.test_fp_qmv", + "TestQuantized.test_fp_qvm", "TestQuantized.test_qvm", "TestQuantized.test_qvm_splitk", "TestQuantized.test_small_matrix", diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index 87b1f2ff..9e581970 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -289,7 +289,7 @@ class TestQuantized(mlx_tests.MLXTestCase): [128, 64, 32], # group_size [2, 3, 4, 5, 6, 8], # bits [256, 512, 67], # M - [64, 128], # N + [64, 256], # N [0, 1, 3, 8], # B ) for group_size, bits, M, N, B in tests: @@ -309,33 +309,34 @@ class TestQuantized(mlx_tests.MLXTestCase): self.assertEqual(y_q.shape, y_hat.shape) self.assertLess((y_q - y_hat).abs().max(), 1e-3) - def test_mxfp4_qmv(self): + def test_fp_qmv(self): key = mx.random.key(0) k1, k2 = mx.random.split(key) tests = product( [256, 512, 67], # M - [64, 128], # N + [64, 256], # N [0, 1, 3, 8], # B ) + modes = ["mxfp4", "nvfp4", "mxfp8"] for M, N, B in tests: - with self.subTest(shape=(B, M, N), group_size=32): - x_shape = (3, 1, N) if B == 0 else (B, 1, N) - w_shape = (M, N) if B == 0 else (B, M, N) - x = mx.random.normal(shape=x_shape, key=k1) - w = mx.random.normal(shape=w_shape, key=k2) - w_q, scales = mx.quantize(w, group_size=32, mode="mxfp4") - w_hat = mx.dequantize(w_q, scales, group_size=32, mode="mxfp4") - y_q = mx.quantized_matmul( - x, - w_q, - scales, - transpose=True, - group_size=32, - mode="mxfp4", - ) - y_hat = x @ mx.swapaxes(w_hat, -1, -2) - self.assertEqual(y_q.shape, y_hat.shape) - self.assertLess((y_q - y_hat).abs().max(), 1e-3) + for mode in modes: + with self.subTest(shape=(B, M, N), mode=mode): + x_shape = (3, 1, N) if B == 0 else (B, 1, N) + w_shape = (M, N) if B == 0 else (B, M, N) + x = mx.random.normal(shape=x_shape, key=k1) + w = mx.random.normal(shape=w_shape, key=k2) + w_q, scales = mx.quantize(w, mode=mode) + w_hat = mx.dequantize(w_q, scales, mode=mode) + y_q = mx.quantized_matmul( + x, + w_q, + scales, + transpose=True, + mode=mode, + ) + y_hat = x @ mx.swapaxes(w_hat, -1, -2) + self.assertEqual(y_q.shape, y_hat.shape) + self.assertLess((y_q - y_hat).abs().max(), 1e-3) def test_qvm(self): key = mx.random.key(0) @@ -402,7 +403,7 @@ class TestQuantized(mlx_tests.MLXTestCase): self.assertEqual(y_q.shape, y_hat.shape) self.assertLess((y_q - y_hat).abs().max(), 2e-3) - def test_mxfp4_qvm(self): + def test_fp_qvm(self): key = mx.random.key(0) k1, k2 = mx.random.split(key) tests = product( @@ -413,26 +414,27 @@ class TestQuantized(mlx_tests.MLXTestCase): # Add a splitk tests = list(tests) tests.append((128, 16384, 0)) + modes = ["mxfp4", "nvfp4", "mxfp8"] for M, N, B in tests: - with self.subTest(shape=(B, M, N)): - x_shape = (1, N) if B == 0 else (B, 1, N) - w_shape = (N, M) if B == 0 else (B, N, M) - x = mx.random.normal(shape=x_shape, key=k1) - w = mx.random.normal(shape=w_shape, key=k2) - w_q, scales = mx.quantize(w, group_size=32, mode="mxfp4") - w_hat = mx.dequantize(w_q, scales, group_size=32, mode="mxfp4") - y_q = mx.quantized_matmul( - x, - w_q, - scales, - transpose=False, - group_size=32, - mode="mxfp4", - ) - y_hat = x @ w_hat - self.assertEqual(y_q.shape, y_hat.shape) - self.assertLess((y_q - y_hat).abs().max(), 2e-3) + for mode in modes: + with self.subTest(shape=(B, M, N), mode=mode): + x_shape = (1, N) if B == 0 else (B, 1, N) + w_shape = (N, M) if B == 0 else (B, N, M) + x = mx.random.normal(shape=x_shape, key=k1) + w = mx.random.normal(shape=w_shape, key=k2) + w_q, scales = mx.quantize(w, mode=mode) + w_hat = mx.dequantize(w_q, scales, mode=mode) + y_q = mx.quantized_matmul( + x, + w_q, + scales, + transpose=False, + mode=mode, + ) + y_hat = x @ w_hat + self.assertEqual(y_q.shape, y_hat.shape) + self.assertLess((y_q - y_hat).abs().max(), 2e-3) def test_mode_error_cases(self): w = mx.random.normal(shape=(256, 256)) @@ -626,7 +628,7 @@ class TestQuantized(mlx_tests.MLXTestCase): self.assertLess((y_q - y_hat).abs().max(), 1e-3) def test_gather_qmm(self): - def quantize(w, transpose=True, group_size=64, bits=4, mode="affine"): + def quantize(w, transpose=True, group_size=None, bits=None, mode="affine"): if mode == "affine": qw, s, b = mx.quantize(w, group_size=group_size, bits=bits, mode=mode) else: @@ -647,8 +649,8 @@ class TestQuantized(mlx_tests.MLXTestCase): lhs_indices=None, rhs_indices=None, transpose=True, - group_size=64, - bits=4, + group_size=None, + bits=None, mode="affine", ): with self.subTest( @@ -737,9 +739,22 @@ class TestQuantized(mlx_tests.MLXTestCase): "lhs_indices": (0,), "batch_B": (3,), "rhs_indices": (2, 1), - "group_size": 32, + "mode": "nvfp4", + }, + { + "batch_A": (1,), + "lhs_indices": (0,), + "batch_B": (3,), + "rhs_indices": (2, 1), "mode": "mxfp4", }, + { + "batch_A": (1,), + "lhs_indices": (0,), + "batch_B": (3,), + "rhs_indices": (2, 1), + "mode": "mxfp8", + }, ) for kwargs in inputs: @@ -753,24 +768,24 @@ class TestQuantized(mlx_tests.MLXTestCase): test_shape(32, 512, 32, transpose=False, **kwargs) test_shape(1, 512, 32, transpose=False, **kwargs) - def test_qmm_mxfp4_type(self): + def test_qmm_fp_type(self): indices = mx.array([[2], [0], [1]], dtype=mx.uint32) - for t in [mx.bfloat16, mx.float16, mx.float32]: - x = mx.random.normal((32, 256)).astype(t) + modes = ["mxfp8", "mxfp4"] + for mode in modes: + for t in [mx.bfloat16, mx.float16, mx.float32]: + x = mx.random.normal((32, 256)).astype(t) - w = mx.random.normal((32, 256)) - wq, s = mx.quantize(w, mode="mxfp4", bits=4, group_size=32) - out = mx.quantized_matmul(x, wq, s, mode="mxfp4", group_size=32, bits=4) - self.assertEqual(out.dtype, t) + w = mx.random.normal((32, 256)) + wq, s = mx.quantize(w, mode=mode) + out = mx.quantized_matmul(x, wq, s, mode=mode) + self.assertEqual(out.dtype, t) - w = mx.random.normal((4, 32, 256)) - wq, s = mx.quantize(w, mode="mxfp4", bits=4, group_size=32) + w = mx.random.normal((4, 32, 256)) + wq, s = mx.quantize(w, mode=mode) - out = mx.gather_qmm( - x, wq, s, rhs_indices=indices, mode="mxfp4", group_size=32, bits=4 - ) - self.assertEqual(out.dtype, t) + out = mx.gather_qmm(x, wq, s, rhs_indices=indices, mode=mode) + self.assertEqual(out.dtype, t) def test_gather_matmul_grad(self): def quantize(w, transpose=True, group_size=64, bits=4): @@ -802,14 +817,14 @@ class TestQuantized(mlx_tests.MLXTestCase): self.assertTrue(mx.allclose(g1, g2, atol=1e-4)) def test_gather_qmm_sorted(self): - def quantize(w, transpose=True, bits=4, group_size=64, mode="affine"): + def quantize(w, transpose=True, group_size=None, mode="affine"): if mode == "affine": - qw, s, b = mx.quantize(w, group_size=group_size, bits=bits, mode=mode) + qw, s, b = mx.quantize(w, group_size=group_size, mode=mode) else: - qw, s = mx.quantize(w, group_size=group_size, bits=bits, mode=mode) + qw, s = mx.quantize(w, mode=mode) b = None - w_hat = mx.dequantize(qw, s, b, group_size=group_size, bits=bits, mode=mode) + w_hat = mx.dequantize(qw, s, b, group_size=group_size, mode=mode) if transpose: w_hat = w_hat.swapaxes(-1, -2) return w_hat, qw, s, b @@ -831,11 +846,15 @@ class TestQuantized(mlx_tests.MLXTestCase): # L, K, D, E, I, transpose (32, 512, 512, 4, 2, True, "affine"), (32, 512, 544, 4, 2, True, "mxfp4"), + (32, 512, 544, 4, 2, True, "nvfp4"), + (32, 512, 544, 4, 2, True, "mxfp8"), (133, 512, 512, 4, 2, True, "affine"), (133, 512, 555, 4, 2, True, "affine"), (133, 512, 512, 4, 2, True, "affine"), (64, 512, 512, 4, 2, False, "affine"), (64, 512, 544, 4, 2, False, "mxfp4"), + (64, 512, 544, 4, 2, False, "nvfp4"), + (64, 512, 544, 4, 2, False, "mxfp8"), (133, 512, 512, 4, 2, False, "affine"), (133, 512, 544, 4, 2, False, "affine"), (133, 512, 555, 4, 2, False, "affine"), @@ -848,8 +867,8 @@ class TestQuantized(mlx_tests.MLXTestCase): for L, K, D, E, I, transpose, mode in parameters: with self.subTest(L=L, K=K, D=D, E=E, I=I, transpose=transpose, mode=mode): - if mode == "mxfp4": - group_size = 32 + if mode != "affine": + group_size = None dtype = ( mx.bfloat16 if (mx.default_device() == mx.gpu) else mx.float32 ) @@ -984,36 +1003,33 @@ class TestQuantized(mlx_tests.MLXTestCase): num_ds = (out_up - out_down) / (2 * eps) self.assertAlmostEqual(dparams[p][idx], num_ds, delta=2e-2) - def test_mxfp4_vjp_scales_throws(self): + def test_fp_vjp_scales_throws(self): mx.random.seed(0) x = mx.random.normal(shape=(2, 512)) w = mx.random.normal(shape=(512, 512)) - wq, s = mx.quantize(w, bits=4, group_size=32, mode="mxfp4") + for mode in ["mxfp4", "mxfp8", "nvfp4"]: + wq, s = mx.quantize(w, mode=mode) - def mm(s, x, wq): - return mx.quantized_matmul( - x, wq, s, bits=4, group_size=32, mode="mxfp4" - ).sum() + def mm(s, x, wq): + return mx.quantized_matmul(x, wq, s, mode=mode).sum() - # Should raise - with self.assertRaises(ValueError): - ds = mx.grad(mm)(s, x, wq) + # Should raise + with self.assertRaises(ValueError): + ds = mx.grad(mm)(s, x, wq) - rhs_indices = mx.array(0) - with self.assertRaises(ValueError): + rhs_indices = mx.array(0) + with self.assertRaises(ValueError): - def gmm(s, x, wq): - return mx.gather_qmm( - x, - wq, - s, - rhs_indices=rhs_indices, - bits=4, - group_size=32, - mode="mxfp4", - ).sum() + def gmm(s, x, wq): + return mx.gather_qmm( + x, + wq, + s, + rhs_indices=rhs_indices, + mode=mode, + ).sum() - ds = mx.grad(gmm)(s, x, wq) + ds = mx.grad(gmm)(s, x, wq) def test_quantize_strided(self): N = 64