Repository navigation
quantized_matmul(transpose=False) wrong results on Metal (gs 32/64, bits 2/4/8) #4201
Description
Bug: quantized_matmul(..., transpose=False) returns wrong results on Metal
Summary
On current main (31b5cbb09), the non-transposed path of mx.quantized_matmul produces garbage on Metal. transpose=True is fine; transpose=False fails for group_size 32 and 64 and bits 2/4/8. This also shows up as 48 failing subtests in TestQuantized::test_qmm.
Repro (M5 Max, MLX built from main, Metal backend)
import mlx.core as mx
import numpy as np
mx.random.seed(1)
M, N, K = 8, 128, 128
x = mx.random.normal((M, K)).astype(mx.float16)
w = mx.random.normal((N, K)).astype(mx.float16)
for bits in (4, 8):
for gs in (32, 64):
wq, s, b = mx.quantize(w, group_size=gs, bits=bits)
refq = mx.dequantize(wq, s, b, group_size=gs, bits=bits)
y_ref = (x.astype(mx.float32) @ refq.astype(mx.float32).T)
for tp in (False, True):
y_q = mx.quantized_matmul(x, wq, s, b, transpose=tp, group_size=gs, bits=bits)
mx.eval(y_q, y_ref)
err = np.abs(np.array(y_ref) - np.array(y_q)).max()
print(f"bits={bits} gs={gs} transpose={tp}: max_err={err:.2e}")Output:
bits=4 gs=32 transpose=False: max_err=5.90e+02
bits=4 gs=32 transpose=True : max_err=2.92e-02
bits=4 gs=64 transpose=False: max_err=6.79e+01
bits=4 gs=64 transpose=True : max_err=2.87e-02
bits=8 gs=32 transpose=False: max_err=8.95e+03
bits=8 gs=32 transpose=True : max_err=1.87e-02
bits=8 gs=64 transpose=False: max_err=5.07e+01
bits=8 gs=64 transpose=True : max_err=1.61e-02
Notes / hypothesis
- Native an already
transpose=Truepath is correct, so this is scoped to the non-transposed quantized matmul Metal implementation (the path reworked by Fix and enable non-transposed NAX qmm #4051). It is not caused by the open Fix sorted gather_qmm NAX row overflow above 32K #3922 fix (verified by running the same suite with and without it). TestQuantized::test_qmmFAILED/SUBFAILED shapes: {M∈(8,33,64)} x {N∈(128,256)} x {K∈(128,256)}, bits 2/4/8, group_size 32, transpose=False (same class at gs=64).
Happy to dig into the kernel (local Metal build available) and bring a fix PR.
Activity
PhilipJohnBasile commentedon Aug 12, 2026
Root cause found and fixed in PR #4202.
Root cause: the non-transposed NAX qmm path (#4051) tiles 64 columns wide but dequantizes a 64-column block with one per-group scale, so group_size < 64 (e.g. 32) scales columns beyond the first group incorrectly -> garbage/NaN on Metal. transpose=True unaffected.
Fix: gate the non-transposed NAX path on group_size >= 64; smaller group sizes fall back to the legacy kernel.
Verified on a local Metal build (M5 Max, macOS 27): gs=32/64/128 all ~1e-2 vs reference;
test_quantizedwent from 48 failed subtests to 36 passed / 3079 subtests passed.Fix non-transposed NAX quantized matmul at group_size < 64
PhilipJohnBasile commentedon Aug 12, 2026
Resolved with the root-cause fix in PR #4202.
Root cause (in #4051's non-transposed NAX path): the group_size=32
QuantizedBlockLoaderspecialization advanced the scale/bias pointer byn_groups * group_strideper reduction tile. With the scale plane laid out as [K, N/32] and a 64-row tile = 2N scale entries, the double step read out of bounds from the second K-tile on — so gs=32, transpose=False returned garbage/NaN on Metal for any K > 64 (K=64, a single tile, was correct; gs=64/128 and transpose=True use other paths and were fine).Fix: advance by
group_stridein the reduction_dim==0 branch (also fixed the identical pattern in the fp8/MX loader). Addedtest_qmm_non_transposed_group_size_lt_64(single-vs-multi K-tile). Verified on a local Metal build: test_quantized went 48 failed subtests -> 37 passed / 3106 subtests passed.