Describe the bug
The non-transposed quantized_matmul GPU path (qvm_split_k) uses floor
division (N / bn) for the grid's N-dimension where every sibling dispatcher
uses ceil division ((N + bn - 1) / bn). When the output dimension N is
not divisible by bn (= 64), entire output columns are never computed:
- N < 64 (e.g. N = 32):
N / bn == 0 → zero threadgroups in Y → the kernel
never launches → the entire output is zeros.
- N % 64 ≠ 0 (e.g. N = 96):
96 / 64 == 1 → the last 32 columns are never
written → partial garbage in the trailing columns.
The CPU backend is correct. There is no error or crash.
POC
import numpy as np
import mlx.core as mx
K, N = 1024, 32 # K>=1024 -> qvm_split_k; N<64 -> bug
rng = np.random.RandomState(0)
w_data = rng.randint(0, 16, (K, N)).astype(np.uint32) # 4-bit weights
w = mx.array(sum(w_data[:, i::8] << (4 * i) for i in range(8)).astype(np.uint32))
s = mx.array(np.full((K, N // 32), 0.1, np.float32)) # scales, group_size=32
b = mx.array(np.zeros((K, N // 32), np.float32)) # biases
x = mx.array(rng.randn(1, K).astype(np.float32))
out = mx.quantized_matmul(x, w, s, b, transpose=False) # M=1 -> vector path
assert out.any(), "BUG: GPU returned all zeros"
print("OK:", np.array(out)[0, :4])quantized_matmul(transpose=False, M=1): max abs diff GPU vs CPU
K N GPU-CPU GPU all-zero verdict
------------------------------------------------------------
512 32 0.0000 False ok
512 64 0.0000 False ok
...
1024 32 79.7963 True WRONG ← ALL ZEROS
1024 64 0.0000 False ok
1024 96 51.2256 False WRONG ← partial garbage
1024 128 0.0000 False ok
2048 32 81.4160 True WRONG ← ALL ZEROS
2048 64 0.0001 False ok
2048 96 83.6282 False WRONG ← partial garbage
2048 128 0.0001 False ok
BUG REPRODUCED:
K=1024, N=32: GPU-CPU diff=79.80 (ALL ZEROS)
K=1024, N=96: GPU-CPU diff=51.23 (partial garbage)
Describe the bug
The non-transposed
quantized_matmulGPU path (qvm_split_k) uses floordivision (
N / bn) for the grid's N-dimension where every sibling dispatcheruses ceil division (
(N + bn - 1) / bn). When the output dimensionNisnot divisible by
bn(= 64), entire output columns are never computed:N / bn == 0→ zero threadgroups in Y → the kernelnever launches → the entire output is zeros.
96 / 64 == 1→ the last 32 columns are neverwritten → partial garbage in the trailing columns.
The CPU backend is correct. There is no error or crash.
POC