From 5e6a9af0e635db42fa9f6ef9b7497f8dd0c8798a Mon Sep 17 00:00:00 2001 From: Magnus Lundstedt Date: Thu, 15 Oct 2026 15:14:19 +0200 Subject: [PATCH 1/2] Add optional int8 path for quantized matmul with affine weights --- docs/src/usage/environment_variables.rst | 7 + docs/src/usage/precision.rst | 23 ++ mlx/backend/metal/kernels/quantized_nax.h | 346 ++++++++++++++++++ mlx/backend/metal/kernels/quantized_nax.metal | 47 +++ mlx/backend/metal/quantized.cpp | 163 +++++++++ mlx/ops.cpp | 8 +- mlx/primitives.cpp | 3 +- mlx/primitives.h | 12 +- mlx/utils.h | 5 + python/tests/test_quantized.py | 128 +++++++ 10 files changed, 737 insertions(+), 5 deletions(-) diff --git a/docs/src/usage/environment_variables.rst b/docs/src/usage/environment_variables.rst index b43c053d52..bea8fb94a5 100644 --- a/docs/src/usage/environment_variables.rst +++ b/docs/src/usage/environment_variables.rst @@ -26,6 +26,13 @@ General on supported hardware. The default is ``1``. Set it to ``0`` to keep these operations in full ``float32`` precision. See :doc:`precision`. +.. envvar:: MLX_QMM_INT8 + + Run affine quantized matmuls on 8-bit integer weights and activations on + supported hardware. The default is ``0``. Set it to ``1``, ``4`` or ``2`` + to enable it, from most to least precise. Other values are ignored. See + :doc:`precision`. + Distributed ----------- diff --git a/docs/src/usage/precision.rst b/docs/src/usage/precision.rst index 31c3f2efea..09a6e27f42 100644 --- a/docs/src/usage/precision.rst +++ b/docs/src/usage/precision.rst @@ -31,3 +31,26 @@ To disable Winograd, set :envvar:`MLX_CONV_WINOGRAD` to ``0``: .. code-block:: shell MLX_CONV_WINOGRAD=0 python my_script.py + +8-bit Quantized Matmul +---------------------- + +Quantized matmuls dequantize the weights to the type of the input and +multiply in floating point. On hardware with dedicated +matrix-multiplication units MLX can instead round the weights and the input +to 8 bits and multiply those, which is faster and less precise. This is off +by default. + +For affine quantized weights set :envvar:`MLX_QMM_INT8`. The weights and the +input are rounded to 8-bit integers with one scale per row. With ``1`` the +scale of the input is refined for every block of 64 elements, with ``4`` for +every four blocks, and ``2`` uses the row scale alone. Other values keep the +default path: + +.. code-block:: shell + + MLX_QMM_INT8=1 python my_script.py + +Expect a relative difference of one to a few percent from the floating point +result. Only products of 64 rows or more with transposed weights take this +path. The variable is read each time a quantized matmul is built. diff --git a/mlx/backend/metal/kernels/quantized_nax.h b/mlx/backend/metal/kernels/quantized_nax.h index 4ee92c3854..3347d1695f 100644 --- a/mlx/backend/metal/kernels/quantized_nax.h +++ b/mlx/backend/metal/kernels/quantized_nax.h @@ -2133,3 +2133,349 @@ template < }); }); } + +// Staged int8 qmm: weights and activations are rounded to int8 with one scale +// per row and the product is accumulated in int32. + +// Fixed point bits of the weight fold w8 = (q * a + b) >> QMM_I8_FRAC +MLX_MTL_CONST int QMM_I8_FRAC = 16; + +// Unpack 4 consecutive codes, p points at the start of a pack +template +inline int4 qmm_i8_unpack4(const device uint8_t* p, int i) { + static_assert( + bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || + bits == 8, + "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + + constexpr uint mask = (1u << bits) - 1u; + + uint w; + if (bits == 4) { + w = *(const device uint16_t*)(p + 2 * i); + } else if (bits == 8) { + w = *(const device uint32_t*)(p + 4 * i); + } else { + const int offset = 4 * bits * i; + p += offset >> 3; + w = p[0]; + if (bits > 2) { + w |= uint(p[1]) << 8; + } + if (bits > 4) { + w |= uint(p[2]) << 16; + } + w >>= (offset & 7); + } + + return int4( + w & mask, + (w >> bits) & mask, + (w >> (2 * bits)) & mask, + (w >> (3 * bits)) & mask); +} + +// The scale of each row of the weights and its fixed point reciprocal +template +[[kernel]] void affine_qmm_i8_row_scales( + const device T* scales [[buffer(0)]], + const device T* biases [[buffer(1)]], + device float* w_scale [[buffer(2)]], + device float* w_inv [[buffer(3)]], + const constant int& K_g [[buffer(4)]], + uint tid [[thread_position_in_grid]]) { + constexpr float q_max = float((1 << bits) - 1); + + scales += size_t(tid) * K_g; + biases += size_t(tid) * K_g; + + float extent = 0.0f; + for (int g = 0; g < K_g; g++) { + const float s = float(scales[g]); + const float b = float(biases[g]); + extent = max(extent, max(fabs(b), fabs(fma(s, q_max, b)))); + } + + const float sw = max(extent * (1.0f / 127.0f), 1e-30f); + w_scale[tid] = sw; + w_inv[tid] = float(1 << QMM_I8_FRAC) / sw; +} + +// Quantize each row of x to int8, one threadgroup per row. A block of 64 has +// the row scale / 2^e, one e per x_span blocks. Output [M/64][K/64][64][64]. +template +[[kernel]] void affine_qmm_i8_quantize_x( + const device T* x [[buffer(0)]], + device int8_t* x8 [[buffer(1)]], + device float* x_scale [[buffer(2)]], + device uint8_t* x_exp [[buffer(3)]], + const constant int& K [[buffer(4)]], + const constant int& M [[buffer(5)]], + const constant int& max_exp [[buffer(6)]], + const constant int& x_span [[buffer(7)]], + 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]], + uint n_simd [[simdgroups_per_threadgroup]]) { + constexpr int BM = 64; + constexpr int BK = 64; + + threadgroup float shared[33]; + + const int row = tid.y; + const int n_blocks = K / BK; + const size_t tile = size_t(row / BM) * n_blocks + lid; + device char4* q = (device char4*)(x8 + (tile * BM + row % BM) * BK); + + // The exponent of the first block is stored twice so that the matmul can + // read the change from the block before + x_exp += size_t(row) * (n_blocks + 1); + + if (row >= M) { + for (int i = 0; i < BK / 4; i++) { + q[i] = char4(0); + } + x_exp[lid + 1] = 0; + if (lid == 0) { + x_exp[0] = 0; + } + return; + } + + x += size_t(row) * K + lid * BK; + + float vals[BK]; + float block_max = 0.0f; + for (int i = 0; i < BK; i++) { + vals[i] = float(x[i]); + block_max = max(block_max, fabs(vals[i])); + } + const float simd_row_max = simd_max(block_max); + if (simd_lid == 0) { + shared[simd_gid] = simd_row_max; + } + + // The blocks that share an exponent are threads next to each other + float span_max = block_max; + if (x_span >= 2) { + span_max = max(span_max, simd_shuffle_xor(span_max, ushort(1))); + } + if (x_span >= 4) { + span_max = max(span_max, simd_shuffle_xor(span_max, ushort(2))); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (lid == 0) { + float row_max = shared[0]; + for (uint i = 1; i < n_simd; i++) { + row_max = max(row_max, shared[i]); + } + shared[32] = max(row_max * (1.0f / 127.0f), 1e-12f); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + const float sx = shared[32]; + const float inv = 1.0f / sx; + + // Halve the scale of these blocks while their values still fit in int8 + float t = span_max * inv; + int e = 0; + while (e < max_exp && t * 2.0f < 127.5f) { + t *= 2.0f; + e++; + } + const float gain = float(1 << e); + + for (int i = 0; i < BK / 4; i++) { + const float4 v = + float4(vals[4 * i], vals[4 * i + 1], vals[4 * i + 2], vals[4 * i + 3]); + q[i] = char4(clamp(int4(round((v * inv) * gain)), -127, 127)); + } + + x_exp[lid + 1] = e; + if (lid == 0) { + x_exp[0] = e; + } + // The results are in the units of the last block + if (lid == n_blocks - 1) { + x_scale[row] = sx / gain; + } +} + +// Same structure as qmm_t_nax_tgp_impl, with int8 tiles and an int32 sum. +// x_span is 0 for one scale per row of x, or the blocks sharing an exponent. +template +[[kernel]] void affine_qmm_t_nax_i8( + const device uint32_t* w [[buffer(0)]], + const device T* scales [[buffer(1)]], + const device T* biases [[buffer(2)]], + const device float* w_inv [[buffer(3)]], + const device float* w_scale [[buffer(4)]], + const device int8_t* x8 [[buffer(5)]], + const device float* x_scale [[buffer(6)]], + device T* y [[buffer(7)]], + const device uint8_t* x_exp [[buffer(8)]], + const constant int& K [[buffer(9)]], + const constant int& N [[buffer(10)]], + const constant int& M [[buffer(11)]], + uint3 tid [[threadgroup_position_in_grid]], + uint lid [[thread_index_in_threadgroup]], + uint simd_gid [[simdgroup_index_in_threadgroup]]) { + constexpr int BM = 64; + constexpr int BN = 64; + constexpr int BK = 64; + constexpr int WM = 2; + constexpr int WN = 2; + constexpr int BK_padded = BK + 16; + + // Each thread folds VPT weights of one row, all in one quantization group + constexpr int VPT = (BN * BK) / (WM * WN * SIMD_SIZE); + static_assert( + VPT % 8 == 0 && BK % VPT == 0 && group_size % VPT == 0, + "A thread stages whole packs of the weights within one group"); + + constexpr int pack_factor = get_pack_factor(); + constexpr int bytes_per_pack = get_bytes_per_pack(); + + threadgroup uint Ws_storage[(BN * BK_padded) / 4]; + threadgroup int8_t* Ws = (threadgroup int8_t*)Ws_storage; + + constexpr short SM = BM / WM; + constexpr short SN = BN / WN; + constexpr short SK = 32; + + constexpr short TM = SM / 16; + constexpr short TN = SN / 16; + constexpr short TK = SK / 16; + + const int y_row = tid.y * BM; + const int y_col = tid.x * BN; + const short tm = SM * (simd_gid / WN); + const short tn = SN * (simd_gid % WN); + + // The part of the weight tile this thread stages + const int nl = lid / (BK / VPT); + const int kl = (lid % (BK / VPT)) * VPT; + const int n = y_col + nl; + const bool valid = n < N; + + threadgroup int8_t* dst = Ws + nl * BK_padded + kl; + const device uint8_t* wl = (const device uint8_t*)w; + float inv = 0.0f; + if (valid) { + wl += size_t(n) * (K / pack_factor) * bytes_per_pack; + scales += size_t(n) * (K / group_size); + biases += size_t(n) * (K / group_size); + inv = w_inv[n]; + } else { + for (int i = 0; i < VPT / 4; i++) { + *(threadgroup uint*)(dst + 4 * i) = 0; + } + } + + // x8 has one BM x BK block per row tile and BK columns of K + x8 += (size_t(tid.y) * (K / BK) * BM + tm) * BK; + + // A thread has 8 values of the accumulator for each of 4 rows + constexpr int exp_ld = BaseNAXFrag::kElemRowsJump; + const int x_exp_stride = K / BK + 1; + x_exp += size_t(y_row + tm + BaseNAXFrag::get_coord().y) * x_exp_stride; + + NAXTile Dtile; + Dtile.clear(); + + for (int k = 0; k < K; k += BK) { + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (valid) { + const int kk = k + kl; + const int a = int(rint(float(scales[kk / group_size]) * inv)); + const int b = int(rint(float(biases[kk / group_size]) * inv)) + + (1 << (QMM_I8_FRAC - 1)); + const device uint8_t* p = wl + (kk / pack_factor) * bytes_per_pack; + + STEEL_PRAGMA_UNROLL + for (int i = 0; i < VPT / 4; i++) { + const int4 q = qmm_i8_unpack4(p, i); + const int4 v = (q * a + b) >> QMM_I8_FRAC; + *(threadgroup uint*)(dst + 4 * i) = as_type(char4(v)); + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Move the accumulator of each row to the scale of this block of x. A + // fragment has 4 values for each of 2 rows. + constexpr int shift_every = x_span > 0 ? x_span : 1; + const bool shift = x_span > 0 && (k / BK) % shift_every == 0; + STEEL_PRAGMA_UNROLL + for (short i = 0; shift && i < TM; i++) { + typedef metal::vec frag_t; + const device uint8_t* e0 = + x_exp + size_t(i * BaseNAXFrag::kFragRows) * x_exp_stride + k / BK; + const device uint8_t* e1 = e0 + size_t(exp_ld) * x_exp_stride; + const int s0 = int(e0[1]) - int(e0[0]); + const int s1 = int(e1[1]) - int(e1[0]); + const int l0 = max(s0, 0); + const int l1 = max(s1, 0); + const int r0 = max(-s0, 0); + const int r1 = max(-s1, 0); + const frag_t left = frag_t(l0, l0, l0, l0, l1, l1, l1, l1); + const frag_t right = frag_t(r0, r0, r0, r0, r1, r1, r1, r1); + const frag_t rounding = (frag_t(1) << right) >> 1; + STEEL_PRAGMA_UNROLL + for (short j = 0; j < TN; j++) { + Dtile.frag_at(i, j) = + ((Dtile.frag_at(i, j) << left) + rounding) >> right; + } + } + + STEEL_PRAGMA_NO_UNROLL + for (int kk1 = 0; kk1 < BK; kk1 += SK) { + NAXTile Atile; + NAXTile Btile; + + volatile int compiler_barrier; + + Atile.load(x8 + kk1, BK); + Btile.template load(Ws + tn * BK_padded + kk1); + + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); + + (void)compiler_barrier; + } + + x8 += BM * BK; + } + + // Guarded stores are slow, only use them at the edges of the output + const bool interior = (y_row + tm + SM <= M) && (y_col + tn + SN <= N); + dispatch_bool(interior, [&](auto kInterior) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < TM; i++) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < TN; j++) { + const auto frag = Dtile.frag_at(i, j); + STEEL_PRAGMA_UNROLL + for (short e = 0; e < BaseNAXFrag::kElemsPerFrag; e++) { + const short2 pos = BaseNAXFrag::get_coord(e); + const int r = y_row + tm + i * BaseNAXFrag::kFragRows + pos.y; + const int c = y_col + tn + j * BaseNAXFrag::kFragCols + pos.x; + if constexpr (kInterior.value) { + y[size_t(r) * N + c] = + static_cast(float(frag[e]) * x_scale[r] * w_scale[c]); + } else if (r < M && c < N) { + y[size_t(r) * N + c] = + static_cast(float(frag[e]) * x_scale[r] * w_scale[c]); + } + } + } + } + }); +} diff --git a/mlx/backend/metal/kernels/quantized_nax.metal b/mlx/backend/metal/kernels/quantized_nax.metal index 868f0a1935..bf359ab27e 100644 --- a/mlx/backend/metal/kernels/quantized_nax.metal +++ b/mlx/backend/metal/kernels/quantized_nax.metal @@ -112,4 +112,51 @@ instantiate_quantized_all_bits(32) instantiate_quantized_all_bits(64) instantiate_quantized_all_bits(128) + +// Staged int8 qmm +#define instantiate_qmm_i8_span(type, group_size, bits, span) \ + instantiate_kernel( \ + "affine_qmm_t_nax_i8_" #type "_gs_" #group_size "_b_" #bits \ + "_xs_" #span, \ + affine_qmm_t_nax_i8, \ + type, \ + group_size, \ + bits, \ + span) + +#define instantiate_qmm_i8(type, group_size, bits) \ + instantiate_qmm_i8_span(type, group_size, bits, 0) \ + instantiate_qmm_i8_span(type, group_size, bits, 1) \ + instantiate_qmm_i8_span(type, group_size, bits, 4) + +#define instantiate_qmm_i8_groups(type, bits) \ + instantiate_kernel( \ + "affine_qmm_i8_row_scales_" #type "_b_" #bits, \ + affine_qmm_i8_row_scales, \ + type, \ + bits) \ + instantiate_qmm_i8(type, 128, bits) \ + instantiate_qmm_i8(type, 64, bits) \ + instantiate_qmm_i8(type, 32, bits) + +#define instantiate_qmm_i8_types(bits) \ + instantiate_qmm_i8_groups(float, bits) \ + instantiate_qmm_i8_groups(float16_t, bits) \ + instantiate_qmm_i8_groups(bfloat16_t, bits) + +#define instantiate_qmm_i8_quantize(type) \ + instantiate_kernel( \ + "affine_qmm_i8_quantize_x_" #type, \ + affine_qmm_i8_quantize_x, \ + type) + +instantiate_qmm_i8_types(2) +instantiate_qmm_i8_types(3) +instantiate_qmm_i8_types(4) +instantiate_qmm_i8_types(5) +instantiate_qmm_i8_types(6) +instantiate_qmm_i8_types(8) +instantiate_qmm_i8_quantize(float) +instantiate_qmm_i8_quantize(float16_t) +instantiate_qmm_i8_quantize(bfloat16_t) // clang-format on diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index 5e8c4add3f..49af595727 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -1025,6 +1025,147 @@ void qmm_nax( MTL::Size(ceildiv(N, bn), ceildiv(M, bm), B), group_dims); } +// x @ w.T with int8 weights and int8 activations +void qmm_nax_i8( + const array& x, + const array& w, + const array& scales, + const array& biases, + array& out, + int group_size, + int bits, + int M, + int N, + int K, + int x_span, + metal::Device& d, + const Stream& s) { + // Same tiling as qmm_nax + int wm = 2; + int wn = 2; + int bm = 64; + int bn = 64; + int bk = 64; + int M_pad = ((M + bm - 1) / bm) * bm; + int n_blocks = K / bk; + + // x_span is 0 for one scale per row of x, or the number of blocks that + // share an exponent. The blocks of a row must divide into such groups. + if (x_span > 1 && n_blocks % x_span != 0) { + x_span = 1; + } + + // The scale of x can be halved max_exp times for a block. With all the + // other blocks of the row at the row scale the int32 sum does not overflow. + int max_exp = 0; + if (x_span > 0) { + // INT32_MAX / (64 * 127 * 127) + constexpr int budget = 2080; + while (max_exp < 7 && (n_blocks - 1) * (2 << max_exp) + 1 <= budget) { + max_exp++; + } + } + + auto& compute_encoder = metal::get_command_encoder(s); + auto make_temporary = [&compute_encoder](Shape shape, Dtype dtype) { + array t(std::move(shape), dtype, nullptr, {}); + t.set_data(allocator::malloc(t.nbytes())); + compute_encoder.add_temporary(t); + return t; + }; + array x8 = make_temporary({M_pad, K}, int8); + array x_scale = make_temporary({M}, float32); + array x_exp = make_temporary({M_pad, n_blocks + 1}, uint8); + array w_scale = make_temporary({N}, float32); + array w_inv = make_temporary({N}, float32); + + std::string type_string = get_type_string(x.dtype()); + auto get_kernel = [&d](const std::string& kname, std::string template_def) { + return get_qmm_nax_kernel(d, kname, template_def, "affine"); + }; + + { + std::string kname; + concatenate(kname, "affine_qmm_i8_row_scales_", type_string, "_b_", bits); + auto kernel = get_kernel( + kname, + get_template_definition( + kname, "affine_qmm_i8_row_scales", type_string, bits)); + compute_encoder.set_compute_pipeline_state(kernel); + compute_encoder.set_input_array(scales, 0); + compute_encoder.set_input_array(biases, 1); + compute_encoder.set_output_array(w_scale, 2); + compute_encoder.set_output_array(w_inv, 3); + int K_g = K / group_size; + compute_encoder.set_bytes(K_g, 4); + MTL::Size grid_dims(N, 1, 1); + MTL::Size group_dims( + std::min(N, kernel->maxTotalThreadsPerThreadgroup()), 1, 1); + compute_encoder.dispatch_threads(grid_dims, group_dims); + } + + { + std::string kname; + concatenate(kname, "affine_qmm_i8_quantize_x_", type_string); + auto kernel = get_kernel( + kname, + get_template_definition( + kname, "affine_qmm_i8_quantize_x", type_string)); + compute_encoder.set_compute_pipeline_state(kernel); + compute_encoder.set_input_array(x, 0); + compute_encoder.set_output_array(x8, 1); + compute_encoder.set_output_array(x_scale, 2); + compute_encoder.set_output_array(x_exp, 3); + compute_encoder.set_bytes(K, 4); + compute_encoder.set_bytes(M, 5); + compute_encoder.set_bytes(max_exp, 6); + compute_encoder.set_bytes(x_span, 7); + MTL::Size group_dims(n_blocks, 1, 1); + MTL::Size grid_dims(1, M_pad, 1); + compute_encoder.dispatch_threadgroups(grid_dims, group_dims); + } + + { + std::string kname; + kname.reserve(64); + concatenate( + kname, + "affine_qmm_t_nax_i8_", + type_string, + "_gs_", + group_size, + "_b_", + bits, + "_xs_", + x_span); + auto kernel = get_kernel( + kname, + get_template_definition( + kname, + "affine_qmm_t_nax_i8", + type_string, + group_size, + bits, + x_span)); + compute_encoder.set_compute_pipeline_state(kernel); + compute_encoder.set_input_array(w, 0); + compute_encoder.set_input_array(scales, 1); + compute_encoder.set_input_array(biases, 2); + compute_encoder.set_input_array(w_inv, 3); + compute_encoder.set_input_array(w_scale, 4); + compute_encoder.set_input_array(x8, 5); + compute_encoder.set_input_array(x_scale, 6); + compute_encoder.set_output_array(out, 7); + compute_encoder.set_input_array(x_exp, 8); + compute_encoder.set_bytes(K, 9); + compute_encoder.set_bytes(N, 10); + compute_encoder.set_bytes(M, 11); + MTL::Size group_dims(32, wn, wm); + MTL::Size grid_dims((N + bn - 1) / bn, M_pad / bm, 1); + compute_encoder.dispatch_threadgroups(grid_dims, group_dims); + } +} + void gather_qmm_nax( const array& x, const array& w, @@ -2164,6 +2305,28 @@ void QuantizedMatmul::eval_gpu(const std::vector& inputs, array& out) { int M = non_batched ? x.size() / K : x.shape(-2); int N = out.shape(-1); + // The int8 qmm needs 2 bits or more, K / 64 threads per row and tiles of 64 + // rows + if (compute_8bit_ && transpose_ && mode_ == QuantizationMode::Affine && + biases && bits_ > 1 && non_batched && metal::is_nax_available() && + K % 64 == 0 && K / 64 <= 1024 && M >= 64) { + qmm_nax_i8( + x, + w, + scales, + *biases, + out, + group_size_, + bits_, + M, + N, + K, + compute_8bit_ == 2 ? 0 : (compute_8bit_ == 4 ? 4 : 1), + d, + s); + return; + } + int vector_limit = transpose_ ? get_qmv_batch_limit(K, N, d) : 4; auto mode = quantization_mode_to_string(mode_); diff --git a/mlx/ops.cpp b/mlx/ops.cpp index 2d456997c8..69bf91ff58 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -4981,11 +4981,17 @@ array quantized_matmul( } auto out_shape = inputs[0].shape(); out_shape.back() = w_outer_dims; + int compute_8bit = 0; + if (transpose && qmode == QuantizationMode::Affine) { + // Only 1, 2 and 4 are modes. Other values keep the default path. + int mode = env::qmm_int8(); + compute_8bit = (mode == 1 || mode == 2 || mode == 4) ? mode : 0; + } return array( std::move(out_shape), dtype, std::make_shared( - to_stream(s), group_size, bits, qmode, transpose), + to_stream(s), group_size, bits, qmode, transpose, compute_8bit), std::move(inputs)); } diff --git a/mlx/primitives.cpp b/mlx/primitives.cpp index 2f5bc77831..0eb4a0f52a 100644 --- a/mlx/primitives.cpp +++ b/mlx/primitives.cpp @@ -3702,7 +3702,8 @@ std::vector QuantizedMatmul::jvp( bool QuantizedMatmul::is_equivalent(const Primitive& other) const { const QuantizedMatmul& qm_other = static_cast(other); return group_size_ == qm_other.group_size_ && bits_ == qm_other.bits_ && - mode_ == qm_other.mode_ && transpose_ == qm_other.transpose_; + mode_ == qm_other.mode_ && transpose_ == qm_other.transpose_ && + compute_8bit_ == qm_other.compute_8bit_; } std::vector QuantizedMatmul::output_shapes( diff --git a/mlx/primitives.h b/mlx/primitives.h index 68995b7503..26605e6dc8 100644 --- a/mlx/primitives.h +++ b/mlx/primitives.h @@ -1642,12 +1642,14 @@ class QuantizedMatmul : public UnaryPrimitive { int group_size, int bits, QuantizationMode mode, - bool transpose) + bool transpose, + int compute_8bit = 0) : UnaryPrimitive(stream), group_size_(group_size), bits_(bits), mode_(mode), - transpose_(transpose) {} + transpose_(transpose), + compute_8bit_(compute_8bit) {} void eval_cpu(const std::vector& inputs, array& out) override; void eval_gpu(const std::vector& inputs, array& out) override; @@ -1658,7 +1660,8 @@ class QuantizedMatmul : public UnaryPrimitive { bool is_equivalent(const Primitive& other) const override; std::vector output_shapes(const std::vector& inputs) override; auto state() const { - return std::make_tuple(group_size_, bits_, mode_, transpose_); + return std::make_tuple( + group_size_, bits_, mode_, transpose_, compute_8bit_); } private: @@ -1666,6 +1669,9 @@ class QuantizedMatmul : public UnaryPrimitive { int bits_; QuantizationMode mode_; bool transpose_; + // Staged 8-bit compute: int8 with a scale of x per block (1), per row (2) + // or per 4 blocks (4) + int compute_8bit_; }; class QQMatmul : public UnaryPrimitive { diff --git a/mlx/utils.h b/mlx/utils.h index 4b7e6d12d9..dc3e4a3b80 100644 --- a/mlx/utils.h +++ b/mlx/utils.h @@ -227,6 +227,11 @@ inline bool enable_tf32() { return enable_tf32_; } +// Not cached, each quantized matmul reads it when it is built +inline int qmm_int8() { + return get_var("MLX_QMM_INT8", 0); +} + inline int nccl_timeout(int default_value) { static int nccl_timeout = get_var("MLX_NCCL_TIMEOUT", default_value); return nccl_timeout; diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index 429d583f62..366bb69a8c 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -536,6 +536,134 @@ def test_qmm_large_dims(self): self.assertLess((y_q - y_hat).abs().max(), tol) @unittest.skipIf("CI" in os.environ, "too slow in CI") + def test_qmm_int8(self): + # With MLX_QMM_INT8 set, eligible affine qmms run on int8 where the + # matrix kernels support it. The result stays close to the float one. + key = mx.random.key(0) + k1, k2 = mx.random.split(key) + dtype = mx.float16 if (mx.default_device() == mx.gpu) else mx.float32 + tests = product( + [128, 64, 32], # group_size + [2, 3, 4, 5, 6, 8], # bits + [(64, 128, 256), (100, 96, 512), (257, 200, 1024)], # M, N, K + ) + for group_size, bits, (M, N, K) in tests: + with self.subTest(shape=(M, N, K), group_size=group_size, bits=bits): + x = (mx.random.normal(shape=(M, K), key=k1) / K**0.5).astype(dtype) + w = (mx.random.normal(shape=(N, K), key=k2) / K**0.5).astype(dtype) + w_q, scales, biases = mx.quantize(w, group_size, bits) + + ys = [] + # 1 and 4: a scale per block or per 4 blocks of x, 2: per row + for flag in ("0", "1", "2", "4"): + # The switch is read when the op is built + with mlx_tests.scoped_env(MLX_QMM_INT8=flag): + y = mx.quantized_matmul( + x, w_q, scales, biases, True, group_size, bits + ) + mx.eval(y) + ys.append(y.astype(mx.float32)) + for y in ys[1:]: + self.assertEqual(ys[0].shape, y.shape) + err = mx.linalg.norm(y - ys[0]) / mx.linalg.norm(ys[0]) + self.assertLess(err.item(), 5e-2) + + def test_qmm_int8_reference(self): + # Compare with x times the dequantized weights. The bound is for + # Gaussian inputs, heavy tailed inputs measure up to 7%. + key = mx.random.key(0) + k1, k2 = mx.random.split(key) + dtype = mx.float16 if (mx.default_device() == mx.gpu) else mx.float32 + tests = product( + ["1", "2", "4"], # MLX_QMM_INT8 + [128, 64, 32], # group_size + [4, 8], # bits + [(64, 128, 256), (100, 96, 512), (257, 200, 1024)], # M, N, K + ) + for flag, group_size, bits, (M, N, K) in tests: + with self.subTest( + flag=flag, group_size=group_size, bits=bits, shape=(M, N, K) + ): + x = (mx.random.normal(shape=(M, K), key=k1) / K**0.5).astype(dtype) + w = (mx.random.normal(shape=(N, K), key=k2) / K**0.5).astype(dtype) + w_q, scales, biases = mx.quantize(w, group_size, bits) + w_hat = mx.dequantize(w_q, scales, biases, group_size, bits) + y_ref = x.astype(mx.float32) @ w_hat.astype(mx.float32).T + with mlx_tests.scoped_env(MLX_QMM_INT8=flag): + y = mx.quantized_matmul( + x, w_q, scales, biases, True, group_size, bits + ) + err = mx.linalg.norm(y.astype(mx.float32) - y_ref) + self.assertLess((err / mx.linalg.norm(y_ref)).item(), 5e-2) + + def test_qmm_int8_semantics(self): + key = mx.random.key(0) + k1, k2 = mx.random.split(key) + dtype = mx.float16 if (mx.default_device() == mx.gpu) else mx.float32 + + def operands(M, N, K): + x = (mx.random.normal(shape=(M, K), key=k1) / K**0.5).astype(dtype) + w = (mx.random.normal(shape=(N, K), key=k2) / K**0.5).astype(dtype) + return x, w + + def qmm(x, q, flag, group_size=64, bits=4, transpose=True): + # The switch is read when the op is built + with mlx_tests.scoped_env(MLX_QMM_INT8=flag): + y = mx.quantized_matmul(x, *q, transpose, group_size, bits) + mx.eval(y) + return y + + # Unset, 0 and values that are not modes keep the default path + x, w = operands(128, 128, 512) + q = mx.quantize(w, 64, 4) + y_off = qmm(x, q, "0") + for flag in (None, "3", "8", "-1"): + self.assertTrue(mx.array_equal(qmm(x, q, flag), y_off).item()) + + # Fewer than 64 rows, K that is not a multiple of 64, 1-bit weights and + # weights that are not transposed keep the default path + cases = [] + for M, N, K in [(63, 128, 512), (1, 128, 512), (128, 128, 96)]: + x, w = operands(M, N, K) + cases.append((x, mx.quantize(w, 32, 4), 32, 4, True)) + x, w = operands(128, 128, 512) + cases.append((x, mx.quantize(w, 64, 1), 64, 1, True)) + w = (mx.random.normal(shape=(512, 128), key=k2) / 512**0.5).astype(dtype) + cases.append((x, mx.quantize(w, 64, 4), 64, 4, False)) + for i, (x, q, group_size, bits, transpose) in enumerate(cases): + y_off = qmm(x, q, "0", group_size, bits, transpose) + for flag in ("1", "2", "4"): + with self.subTest(case=i, flag=flag): + y = qmm(x, q, flag, group_size, bits, transpose) + self.assertTrue(mx.array_equal(y, y_off).item()) + + # A row does not depend on the rows after it and a column not on the + # columns after it. The default path rounds differently by shape, so + # the rows are only close where the path does not run. + x, w = operands(128, 128, 512) + q = mx.quantize(w, 64, 4) + staged = not mx.array_equal(qmm(x, q, "1"), qmm(x, q, "0")).item() + for flag in ("1", "2", "4"): + for M, N, K in [(257, 128, 512), (300, 200, 512)]: + with self.subTest(flag=flag, shape=(M, N, K)): + x, w = operands(M, N, K) + q = mx.quantize(w, 64, 4) + y = qmm(x, q, flag) + y_part = qmm(x[:256], q, flag) + err = mx.linalg.norm(y[:256].astype(mx.float32) - y_part) + self.assertLess((err / mx.linalg.norm(y_part)).item(), 1e-3) + if staged: + self.assertTrue(mx.array_equal(y[:256], y_part).item()) + with self.subTest(flag=flag, columns=True): + x, w = operands(128, 200, 512) + q = mx.quantize(w, 64, 4) + y = qmm(x, q, flag) + y_part = qmm(x, [a[:192] for a in q], flag) + err = mx.linalg.norm(y[:, :192].astype(mx.float32) - y_part) + self.assertLess((err / mx.linalg.norm(y_part)).item(), 1e-3) + if staged: + self.assertTrue(mx.array_equal(y[:, :192], y_part).item()) + def test_qmm_non_transposed(self): # The non-transposed matmul (w is [K, N]) is reachable mainly from the # vjp of a quantized linear layer, so it gets much less coverage than From c8a13121de316129c6f91904d4ad43486a6c546e Mon Sep 17 00:00:00 2001 From: Magnus Lundstedt Date: Mon, 26 Oct 2026 17:04:34 +0100 Subject: [PATCH 2/2] Remove undefined behavior for partial SIMD groups in affine_qmm_i8_quantize_x --- mlx/backend/metal/kernels/quantized_nax.h | 29 ++++++++++++++++------- mlx/backend/metal/quantized.cpp | 3 ++- 2 files changed, 22 insertions(+), 10 deletions(-) diff --git a/mlx/backend/metal/kernels/quantized_nax.h b/mlx/backend/metal/kernels/quantized_nax.h index 3347d1695f..11d4f6c0ab 100644 --- a/mlx/backend/metal/kernels/quantized_nax.h +++ b/mlx/backend/metal/kernels/quantized_nax.h @@ -2203,6 +2203,8 @@ template // Quantize each row of x to int8, one threadgroup per row. A block of 64 has // the row scale / 2^e, one e per x_span blocks. Output [M/64][K/64][64][64]. +// One thread per block; the threadgroup is rounded up to whole SIMD groups and +// the threads past the last block only take part in the reductions. template [[kernel]] void affine_qmm_i8_quantize_x( const device T* x [[buffer(0)]], @@ -2225,6 +2227,7 @@ template const int row = tid.y; const int n_blocks = K / BK; + const bool active = lid < uint(n_blocks); const size_t tile = size_t(row / BM) * n_blocks + lid; device char4* q = (device char4*)(x8 + (tile * BM + row % BM) * BK); @@ -2233,12 +2236,14 @@ template x_exp += size_t(row) * (n_blocks + 1); if (row >= M) { - for (int i = 0; i < BK / 4; i++) { - q[i] = char4(0); - } - x_exp[lid + 1] = 0; - if (lid == 0) { - x_exp[0] = 0; + if (active) { + for (int i = 0; i < BK / 4; i++) { + q[i] = char4(0); + } + x_exp[lid + 1] = 0; + if (lid == 0) { + x_exp[0] = 0; + } } return; } @@ -2247,9 +2252,11 @@ template float vals[BK]; float block_max = 0.0f; - for (int i = 0; i < BK; i++) { - vals[i] = float(x[i]); - block_max = max(block_max, fabs(vals[i])); + if (active) { + for (int i = 0; i < BK; i++) { + vals[i] = float(x[i]); + block_max = max(block_max, fabs(vals[i])); + } } const float simd_row_max = simd_max(block_max); if (simd_lid == 0) { @@ -2275,6 +2282,10 @@ template } threadgroup_barrier(mem_flags::mem_threadgroup); + if (!active) { + return; + } + const float sx = shared[32]; const float inv = 1.0f / sx; diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index 49af595727..35d38c6f4e 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -1120,7 +1120,8 @@ void qmm_nax_i8( compute_encoder.set_bytes(M, 5); compute_encoder.set_bytes(max_exp, 6); compute_encoder.set_bytes(x_span, 7); - MTL::Size group_dims(n_blocks, 1, 1); + // Whole SIMD groups: the kernel reduces across them + MTL::Size group_dims(((n_blocks + 31) / 32) * 32, 1, 1); MTL::Size grid_dims(1, M_pad, 1); compute_encoder.dispatch_threadgroups(grid_dims, group_dims); }