Skip to content

[CPU] Fix accumulation precision in affine quantized matmul #4803

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
zcbenz merged 2 commits into from
Nov 5, 2026
+103 −36
Merged

[CPU] Fix accumulation precision in affine quantized matmul #4803

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
1e15461
[CPU] Fix accumulation precision in affine quantized matmul
erwinzhang7 Oct 30, 2026
b434bdb
[CPU] Use MLX allocator for quantized matmul accumulator
erwinzhang7 Nov 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
105 changes: 69 additions & 36 deletions mlx/backend/cpu/quantized.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,19 @@ array ensure_row_contiguous(
}
};

// Row accumulator for the non-transposed affine kernel, taken from the
// allocator so it is served by the buffer cache. Empty when transposed.
array qmm_accumulator(
const array& out,
bool transposed_w,
cpu::CommandEncoder& encoder) {
array acc({transposed_w ? 0 : out.shape(-1)}, float32, nullptr, {});
acc.set_data(allocator::malloc(acc.nbytes()));
encoder.add_temporary(acc);
encoder.set_input_array(acc);
return acc;
}

const static float FP4_LUT[16] = {
+0.0f,
+0.5f,
Expand Down Expand Up @@ -103,7 +116,8 @@ void _qmm(
const T* biases,
int M,
int N,
int K) {
int K,
float* acc) {
constexpr int bitmask = (1 << bits) - 1;
constexpr int pack_factor = get_pack_factor(bits, 8);
constexpr int bytes_per_pack = get_bytes_per_pack(bits);
Expand All @@ -114,15 +128,15 @@ void _qmm(
const T* scales_local = scales;
const T* biases_local = biases;

std::fill(result, result + N, 0);
std::fill(acc, acc + N, 0.0f);

for (int k = 0; k < K; k++) {
T* result_local = result;
float* result_local = acc;
T xi = *x++;

for (int n = 0; n < N; n += group_size) {
T scale = *scales_local++;
T bias = *biases_local++;
float scale = *scales_local++;
float bias = *biases_local++;
for (int ng = 0; ng < packs_in_group; ng++) {
if constexpr (bits == 3 || bits == 5 || bits == 6) {
T wl[pack_factor];
Expand All @@ -148,6 +162,7 @@ void _qmm(
}
}

std::copy(acc, acc + N, result);
result += N;
}
}
Expand Down Expand Up @@ -175,10 +190,10 @@ void _qmm_t(

for (int n = 0; n < N; n++) {
const T* x_local = x;
T sum = 0;
float sum = 0;
for (int k = 0; k < K; k += group_size) {
T scale = *scales_local++;
T bias = *biases_local++;
float scale = *scales_local++;
float bias = *biases_local++;

for (int kw = 0; kw < packs_in_group; kw++) {
if constexpr (bits == 3 || bits == 5 || bits == 6) {
Expand Down Expand Up @@ -294,7 +309,8 @@ void _qmm_dispatch_transpose(
int M,
int N,
int K,
bool transposed_w) {
bool transposed_w,
float* acc) {
if (transposed_w) {
// the simd size must be a multiple of the number of elements per word
if constexpr (32 % bits == 0 && simd::max_size<T> % (32 / bits) == 0) {
Expand All @@ -303,7 +319,7 @@ void _qmm_dispatch_transpose(
_qmm_t<T, bits, group_size>(result, x, w, scales, biases, M, N, K);
}
} else {
_qmm<T, bits, group_size>(result, x, w, scales, biases, M, N, K);
_qmm<T, bits, group_size>(result, x, w, scales, biases, M, N, K, acc);
}
}

Expand All @@ -318,19 +334,20 @@ void _qmm_dispatch_group(
int N,
int K,
int group_size,
bool transposed_w) {
bool transposed_w,
float* acc) {
switch (group_size) {
case 32:
_qmm_dispatch_transpose<T, bits, 32>(
result, x, w, scales, biases, M, N, K, transposed_w);
result, x, w, scales, biases, M, N, K, transposed_w, acc);
break;
case 64:
_qmm_dispatch_transpose<T, bits, 64>(
result, x, w, scales, biases, M, N, K, transposed_w);
result, x, w, scales, biases, M, N, K, transposed_w, acc);
break;
case 128:
_qmm_dispatch_transpose<T, bits, 128>(
result, x, w, scales, biases, M, N, K, transposed_w);
result, x, w, scales, biases, M, N, K, transposed_w, acc);
break;
default:
throw std::invalid_argument(
Expand All @@ -350,35 +367,36 @@ void _qmm_dispatch_typed(
int K,
int group_size,
int bits,
bool transposed_w) {
bool transposed_w,
float* acc) {
switch (bits) {
case 1:
_qmm_dispatch_group<T, 1>(
result, x, w, scales, biases, M, N, K, group_size, transposed_w);
result, x, w, scales, biases, M, N, K, group_size, transposed_w, acc);
break;
case 2:
_qmm_dispatch_group<T, 2>(
result, x, w, scales, biases, M, N, K, group_size, transposed_w);
result, x, w, scales, biases, M, N, K, group_size, transposed_w, acc);
break;
case 3:
_qmm_dispatch_group<T, 3>(
result, x, w, scales, biases, M, N, K, group_size, transposed_w);
result, x, w, scales, biases, M, N, K, group_size, transposed_w, acc);
break;
case 4:
_qmm_dispatch_group<T, 4>(
result, x, w, scales, biases, M, N, K, group_size, transposed_w);
result, x, w, scales, biases, M, N, K, group_size, transposed_w, acc);
break;
case 5:
_qmm_dispatch_group<T, 5>(
result, x, w, scales, biases, M, N, K, group_size, transposed_w);
result, x, w, scales, biases, M, N, K, group_size, transposed_w, acc);
break;
case 6:
_qmm_dispatch_group<T, 6>(
result, x, w, scales, biases, M, N, K, group_size, transposed_w);
result, x, w, scales, biases, M, N, K, group_size, transposed_w, acc);
break;
case 8:
_qmm_dispatch_group<T, 8>(
result, x, w, scales, biases, M, N, K, group_size, transposed_w);
result, x, w, scales, biases, M, N, K, group_size, transposed_w, acc);
break;
default:
throw std::invalid_argument(
Expand All @@ -395,7 +413,8 @@ void _qmm_dispatch_typed(
const array& biases,
int bits,
int group_size,
bool transposed_w) {
bool transposed_w,
array& acc) {
int K = x.shape(-1);
int M = x.ndim() > 1 ? x.shape(-2) : 1;
int N = out.shape(-1);
Expand All @@ -420,7 +439,8 @@ void _qmm_dispatch_typed(
K,
bits,
group_size,
transposed_w);
transposed_w,
acc.data<float>());
}
}

Expand All @@ -432,19 +452,20 @@ void _qmm_dispatch(
const array& biases,
int bits,
int group_size,
bool transposed_w) {
bool transposed_w,
array& acc) {
switch (x.dtype()) {
case float32:
_qmm_dispatch_typed<float>(
out, x, w, scales, biases, bits, group_size, transposed_w);
out, x, w, scales, biases, bits, group_size, transposed_w, acc);
break;
case float16:
_qmm_dispatch_typed<float16_t>(
out, x, w, scales, biases, bits, group_size, transposed_w);
out, x, w, scales, biases, bits, group_size, transposed_w, acc);
break;
case bfloat16:
_qmm_dispatch_typed<bfloat16_t>(
out, x, w, scales, biases, bits, group_size, transposed_w);
out, x, w, scales, biases, bits, group_size, transposed_w, acc);
break;
default:
throw std::invalid_argument(
Expand Down Expand Up @@ -714,7 +735,8 @@ void _bs_qmm_dispatch_typed(
const array& rhs_indices,
int bits,
int group_size,
bool transposed_w) {
bool transposed_w,
array& acc) {
int K = x.shape(-1);
int M = x.shape(-2);
int N = out.shape(-1);
Expand Down Expand Up @@ -748,7 +770,8 @@ void _bs_qmm_dispatch_typed(
K,
bits,
group_size,
transposed_w);
transposed_w,
acc.data<float>());
}
}

Expand All @@ -762,7 +785,8 @@ void _bs_qmm_dispatch(
const array& rhs_indices,
int bits,
int group_size,
bool transposed_w) {
bool transposed_w,
array& acc) {
switch (x.dtype()) {
case float32:
_bs_qmm_dispatch_typed<float>(
Expand All @@ -775,7 +799,8 @@ void _bs_qmm_dispatch(
rhs_indices,
bits,
group_size,
transposed_w);
transposed_w,
acc);
break;
case float16:
_bs_qmm_dispatch_typed<float16_t>(
Expand All @@ -788,7 +813,8 @@ void _bs_qmm_dispatch(
rhs_indices,
bits,
group_size,
transposed_w);
transposed_w,
acc);
break;
case bfloat16:
_bs_qmm_dispatch_typed<bfloat16_t>(
Expand All @@ -801,7 +827,8 @@ void _bs_qmm_dispatch(
rhs_indices,
bits,
group_size,
transposed_w);
transposed_w,
acc);
break;
default:
throw std::invalid_argument(
Expand Down Expand Up @@ -946,15 +973,18 @@ void QuantizedMatmul::eval_cpu(const std::vector<array>& inputs, array& out) {
if (mode_ == QuantizationMode::Affine) {
auto biases = ensure_row_contiguous(inputs[3], encoder, stream());
encoder.set_input_array(biases);
auto acc = qmm_accumulator(out, transpose_, encoder);
encoder.dispatch([out = array::unsafe_weak_copy(out),
x = array::unsafe_weak_copy(x),
w = array::unsafe_weak_copy(w),
scales = array::unsafe_weak_copy(scales),
biases = array::unsafe_weak_copy(biases),
acc = array::unsafe_weak_copy(acc),
group_size_ = group_size_,
bits_ = bits_,
transpose_ = transpose_]() mutable {
_qmm_dispatch(out, x, w, scales, biases, group_size_, bits_, transpose_);
_qmm_dispatch(
out, x, w, scales, biases, group_size_, bits_, transpose_, acc);
});
} else {
encoder.dispatch([out = array::unsafe_weak_copy(out),
Expand Down Expand Up @@ -1006,13 +1036,15 @@ void GatherQMM::eval_cpu(const std::vector<array>& inputs, array& out) {
if (mode_ == QuantizationMode::Affine) {
auto biases = ensure_row_contiguous_last_dims(inputs[3]);
encoder.set_input_array(biases);
auto acc = qmm_accumulator(out, transpose_, encoder);
encoder.dispatch([out = array::unsafe_weak_copy(out),
x = array::unsafe_weak_copy(x),
w = array::unsafe_weak_copy(w),
scales = array::unsafe_weak_copy(scales),
biases = array::unsafe_weak_copy(biases),
lhs_indices = array::unsafe_weak_copy(lhs_indices),
rhs_indices = array::unsafe_weak_copy(rhs_indices),
acc = array::unsafe_weak_copy(acc),
group_size_ = group_size_,
bits_ = bits_,
transpose_ = transpose_]() mutable {
Expand All @@ -1026,7 +1058,8 @@ void GatherQMM::eval_cpu(const std::vector<array>& inputs, array& out) {
rhs_indices,
group_size_,
bits_,
transpose_);
transpose_,
acc);
});
} else {
encoder.dispatch([out = array::unsafe_weak_copy(out),
Expand Down
34 changes: 34 additions & 0 deletions python/tests/test_quantized.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1092,6 +1092,40 @@ def test_qmv_affine_bias_sum_precision(self):
expected = mx.full((1, n), sign * k / (4 * denom), dtype)
self.assertTrue(mx.array_equal(y, expected).item())

def test_qmm_accumulation_precision(self):
# x repeats [1, eps, -1, 0] and each block of four weights along K is
# equal, so each block adds eps * w. A sum in the input type loses
# eps * w next to w. In float32 every partial sum is exact.
N, K = 128, 256
for dtype, transpose, gs, bits in product(
[mx.bfloat16, mx.float16],
[True, False],
[32, 64, 128],
[2, 4, 8],
):
with self.subTest(dtype=dtype, transpose=transpose, gs=gs, bits=bits):
eps = 1 / 256 if dtype == mx.bfloat16 else 1 / 2048
x = mx.tile(mx.array([1, eps, -1, 0], dtype), K // 4)
x = mx.stack([x, -x])
# All bits set give the code 2**bits - 1, and the bias moves w to v.
if transpose:
rows, cols = N, K
v = mx.arange(rows)[:, None] + 3 * mx.arange(cols // gs)
else:
rows, cols = K, N
v = mx.arange(rows)[:, None] // 4 + 3 * mx.arange(cols // gs)
v = v % 8 + 1
q = mx.full((rows, cols * bits // 32), 0xFFFFFFFF, mx.uint32)
scales = mx.ones(v.shape, dtype)
biases = (v - (2**bits - 1)).astype(dtype)
w = mx.dequantize(q, scales, biases, group_size=gs, bits=bits)
w = w.astype(mx.float32)
y_ref = x.astype(mx.float32) @ (w.T if transpose else w)
y = mx.quantized_matmul(
x, q, scales, biases, transpose=transpose, group_size=gs, bits=bits
)
self.assertTrue(mx.array_equal(y, y_ref.astype(dtype)).item())

@unittest.skipIf(not mx.metal.is_available(), "requires Metal")
def test_qmv_fast_unaligned_n(self):
# Output sizes that are not a multiple of 8 with an aligned input size.
Expand Down
Loading