#4479 mx.dequantize returns garbage on CPU past 2^31 unpacked bits (int size in cpu/binary.h)

What happens

On a CPU stream, a contiguous elementwise binary op silently stops writing its output once one side reaches 2^31 elements. No exception, no warning. The op returns an array that holds whatever the output buffer held before: zeros when the buffer was freshly allocated, the unmodified input when mlx donated the input's buffer.

I found it through mx.dequantize. Affine dequantize has no CPU kernel: affine_dequantize builds the fast::Quantize primitive only for a GPU stream and otherwise evaluates an op-by-op fallback (mlx/ops.cpp:5341 and :5345, v0.32.2). At a bit width that is not a power of two, that fallback widens every packed uint32 word to one uint32 per bit and masks it, so the full size ops in that branch run over numel * bits elements; the bitwise_and at mlx/ops.cpp:5314 is the VectorScalar case that reads a.data_size(), so it is the likely one, though I did not instrument which of them drops its writes. Past 2^31 such an op writes nothing. When the input buffer is donated to the output, the raw shifted words survive into the sum and get scaled by the group scale, which fits the magnitudes below. A 6-bit round trip of values in [-1, 1] then comes back off by 7e7, in an array that looks perfectly ordinary.

Repro

# /// script
# requires-python = ">=3.12"
# dependencies = ["mlx==0.32.2"]
# ///
import mlx.core as mx

mx.set_default_device(mx.cpu)

# 1. Any contiguous elementwise binary op on the CPU stream. About 4 GB peak, under a second.
for cols in (2**30 - 1, 2**30):
    a = mx.full((2, cols), 3, dtype=mx.uint8)
    out = mx.add(a, mx.array(1, mx.uint8))
    mx.eval(out)
    a0, a1 = int(a[0, 0].item()), int(a[-1, -1].item())
    o0, o1 = int(out[0, 0].item()), int(out[-1, -1].item())
    print(f"add         n={2 * cols:>13,}  in={a0} {a1}  out={o0} {o1}  want out 4 4")
    del a, out
    mx.clear_cache()

# 2. What that does to mx.dequantize at a bit width that is not a power of two. ~15 GB peak.
for lead in (218, 219):
    n = lead * 640 * 2560
    w = mx.sin(mx.arange(n, dtype=mx.float32).reshape(lead, 640, 2560) * 0.017)
    w = w.astype(mx.bfloat16)
    wq, scales, biases = mx.quantize(w, group_size=64, bits=6)
    deq = mx.dequantize(wq, scales, biases, group_size=64, bits=6)
    err = mx.max(mx.abs(deq.astype(mx.float32) - w.astype(mx.float32))).item()
    print(f"dequantize  numel*bits={n * 6:>13,}  max|deq-w|={err:.3e}  want <= 0.02")
    del w, wq, scales, biases, deq
    mx.clear_cache()

Output here, on both 0.32.2 and 0.32.0, byte for byte the same:

add         n=2,147,483,646  in=3 3  out=4 4  want out 4 4
add         n=2,147,483,648  in=3 3  out=0 0  want out 4 4
dequantize  numel*bits=2,143,027,200  max|deq-w|=1.758e-02  want <= 0.02
dequantize  numel*bits=2,152,857,600  max|deq-w|=7.025e+07  want <= 0.02

11 seconds, 15.0 GiB peak RSS. The input of the add case reads back as 3 at both ends at every size, so the array and the indexing are fine and only the output is untouched.

Observed against expected

The affine round trip of an input in [-1, 1] at 6 bits cannot be off by more than one quantization step, 0.032. Below the boundary the repro above measures 0.0176. One row of 640x2560 further up it measures 7.0e7, which is 2.2e9 steps; the magnitude of the garbage depends on the input (I have seen 1.2e7 to 7e7 for the same shape), the failure does not.

The add probe puts the boundary exactly at 2^31 = 2,147,483,648. The dequantize sweep below brackets it and shows that it is numel * bits that decides, not numel and not the tensor rank. The sweep uses a different input generator from the repro above, so its in-tolerance figures differ slightly (0.0156 against 0.0176 on the same shape); the pass/fail split does not:

shape	bits	group	numel * bits	max abs err
(218, 640, 2560)	6	64	2,143,027,200	1.56e-02
(219, 640, 2560)	6	64	2,152,857,600	6.97e+07
(436, 640, 2560)	3	64	2,143,027,200	1.29e-01
(437, 640, 2560)	3	64	2,147,942,400	6.54e+08
(262, 640, 2560)	5	64	2,146,304,000	3.12e-02
(263, 640, 2560)	5	64	2,154,496,000	1.42e+08
(140160, 2560)	6	64	2,152,857,600	6.97e+07
(219, 640, 2560)	6	32	2,152,857,600	3.54e+07
(219, 640, 2560)	6	128	2,152,857,600	1.21e+08
(656, 640, 2560)	2	64	2,149,580,800	2.58e-01
(328, 640, 2560)	4	64	2,149,580,800	6.25e-02
(164, 640, 2560)	8	64	2,149,580,800	7.81e-03

Widths 3, 5 and 6 break above the line and are correct below it. Widths 2, 4 and 8 are correct at a comparable bit count above the line (2,149,580,800) because they take the other branch of the fallback, whose widest intermediate holds numel elements (537,395,200 for the 4-bit row, far below 2^31) and whose full size ops are a broadcast multiply and add at a collapsed rank of 3 or less, where the count comes from the shape. binary_op_dispatch_dims (binary.h:115) narrows to int as well, from a.size() at :245, :257, :269 and :281; at collapsed rank 4 or more that count drives the loop at :138, so I would expect the same failure there, untested. Running the same case in float32 instead of bfloat16, weights and scales alike, changes nothing (6.96e+07). mx.quantize looks correct: it has a real CPU kernel, and dequantizing the same packed arrays on a GPU stream is within tolerance whichever device did the quantize (1.56e-02).

One more measurement that pins the mechanism. At 2^32 + 4 elements the same mx.add writes exactly four values and leaves the rest untouched:

cpu  n=2^32+4  first six outputs=[4, 4, 4, 4, 0, 0]  out[-1,-1]=0

A count that wraps to 4 explains that. A clamp or an allocation limit does not.

Where it comes from

mlx/backend/cpu/binary.h, identical at tag v0.32.2 and on the default branch today. The three contiguous functors take the element count as int:

19:  void operator()(const T* a, const T* b, U* dst, int size) {   // VectorScalar
39:  void operator()(const T* a, const T* b, U* dst, int size) {   // ScalarVector
59:  void operator()(const T* a, const T* b, U* dst, int size) {   // VectorVector

and the callers hand them a size_t:

167:    ScalarVector<Op>{}(a_ptr, b_ptr, out_ptr, b.data_size());
173:    VectorScalar<Op>{}(a_ptr, b_ptr, out_ptr, a.data_size());
179:    VectorVector<Op>{}(a_ptr, b_ptr, out_ptr, a.size());

At 2^31 the narrowing makes size negative, so while (size >= N) and while (size-- > 0) both fail on their first test and the functor returns having written nothing. Between 2^32 and 2^32 + 2^31 the value wraps back to a small positive number and a short prefix gets written, which is the four elements above.

The neighbouring headers do not narrow: unary.h:27 keeps auto size = a.data_size() and ternary.h:61 declares size_t size, so binary.h looks like the one that was missed.

I have not tested a patch, and I do not know whether widening those three signatures to size_t is enough on its own or whether the SIMD tail wants the same treatment elsewhere.

Versions and hardware

mlx 0.32.2 (newest on PyPI, uploaded 2026-08-25) and mlx 0.32.0, same numbers on both. mlx-metal 0.32.2, numpy 2.5.3, CPython 3.12.12, installed with uv into an empty project. Apple M3 Ultra, 256 GB, macOS 26.6.2 (25G83). mlx/backend/cpu/binary.h on the default branch is byte identical to the v0.32.2 file as of today, so I expect main to reproduce as well.

Why it matters in practice

Quantized MoE checkpoints put whole expert stacks in one 3-D tensor of several hundred million elements, which at 6 bits is above this line, so a CPU side round trip check on one reports a failure that is not in the data, and code that dequantizes it on the CPU gets a plausible looking array that is wrong by seven orders of magnitude. The workaround is easy once you know: run it on the GPU, or chunk below 2^31 / bits elements. Finding out is the expensive part, because nothing about the result says it failed.

Same family of 32-bit count and index overflows, all on the Metal side and all closed: #3248 (conv_general above 2^31 output elements), #2894 (SDPA boolean mask above 2^31), #3979 (conv2d returning all zeros for large inputs), #3836 (strided copy picking the int32 kernel by element count), #3327 (arrays above 2^31 elements failing materialization). I did not find an open issue covering the CPU binary path.