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

[builtbycnob]

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.

Related

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.

[DengZhiyuan-math]

Thanks for the detailed report. I’d like to investigate a fix, provided no one is already working on it.

I plan to start with the smaller CPU uint8 reproduction, then check the element-count handling in the three contiguous functors and binary_op_dispatch_dims, including the relevant broadcast and strided paths.

I’ll also investigate regression coverage that exercises the affected count-handling paths without requiring multi-GB allocations in routine CI.

Please let me know if this overlaps with work already in progress.

[builtbycnob]

Thanks for the offer. I will take this one: a PR with the element-count types fixed in the three contiguous functors and in binary_op_dispatch_dims, plus whatever sibling in the CPU backend carries the same int count, is in preparation and should be up within a day or two. A regression test at the boundary needs at least 2 GiB per operand, so I will propose it as an opt-in test and leave its CI placement to the maintainers. A review, and a run of the uint8 case on other hardware, would be welcome.

[Prudctual]

Opened a PR to use size_t for contiguous binary op element counts, fixing the CPU overflow beyond 2^31 elements.

[builtbycnob]

Prudctual's #4517 fixes the three contiguous functors, so I am not opening the PR I announced above. One sibling is left in the same header, binary_op_dispatch_dims (binary.h:115, int size fed from a.size() at :245, :257, :269, :281); noted on the PR with line numbers against main at c948334a.