[rcfa]

What fails

On x86-64, the CPU quantized matmuls crash the process with SIGFPE when M, N or K is 0. On arm64 they do not crash, but with K = 0 the output is never written. Line numbers below are at main = 4a22e4fed.

• _qmm_dispatch_typed (mlx/backend/cpu/quantized.cpp:398, affine) and fp_qmm_dispatch_mode (:636, mxfp4, nvfp4, mxfp8) count the matrices as x.size() / (K * M).
  • With x of shape [0, K] or [M, 0] that is 0 / 0. x86's div traps.
  • arm64's gives 0, so no kernel runs. An empty output needs none, but a K = 0 output ([M, N], an empty sum) is left as allocated.
  • QQMatmul reaches :636 through qqmm (:1363).
• Each operand's offset comes from elem_to_loc (mlx/backend/common/utils.h:16-25), which divides by every dimension. So any operand with a 0 dimension traps on x86 too:
  • in the two dispatchers (:408-411, :645-647);
  • in the gather paths of gather_qmm (:728-739, :829-838).
• quantized_matmul with a 2-D x and a 3-D w whose batch is 0: validation allows it, the output is [M, N] (batch dimensions broadcast only when x has them too, mlx/ops.cpp:4831-4835), K > 0, and elem_to_loc(0, w.shape()) divides by 0 (quantized.cpp:409).
• quantized_matmul's check for its mixed bfloat16/float16 path (mlx/ops.cpp:4802) computes x.size() / x.shape(-1). For a bfloat16 x with float16 scales and biases, group 64 and 4 or 8 bits, a K = 0 input crashes while the op is built, on any device.

Why no test caught it

test empty gather output constructs without dividing by zero builds a gather_qmm with N = 0, but it does not evaluate it. No test evaluates a quantized matmul with M, N or K of 0.

The test

Five new cases in tests/quantized_stream_tests.cpp, on the CPU stream. The first four check the output's shape and that it is all zeros:

• quantized_matmul, affine and mxfp4, transposed and not, with x of shape [0, 64] and [2, 0].
• The mixed bfloat16/float16 quantized_matmul with K = 0.
• qqmm (mxfp4) with K = 0, and with N = 0.
• gather_qmm, affine and mxfp4, with M = 0, N = 0 and K = 0.
• quantized_matmul, affine and mxfp4, with x of shape [2, 64] and w of shape [0, 64, 8]: it must throw std::invalid_argument.

Before the K = 0 case of the first test, clear_cache() runs and an array of ones of the output's size is evaluated and dropped, so the allocator hands that buffer to the output: an output that is not written then reads as ones, not as fresh zeros. That case compares the output with a scalar 0, so that no other array takes the buffer. The mixed-path test also checks the output dtype, float32 (the promoted path).

Without the fix each case crashes on x86-64 Linux. The first failure:

tests/quantized_stream_tests.cpp:207: FATAL ERROR: test case CRASHED: SIGFPE - Floating point error signal

Before the first test freed a buffer of ones, the cases passed on arm64 Linux and on macOS: the unwritten K = 0 outputs happened to read as zeros. On arm64 and macOS the new case fails without the fix: it does not throw (bf-qe-vred5-{arm,mac}.tests.log).

The fix

• QuantizedMatmul, GatherQMM and QQMatmul::eval_cpu first call one helper. For an empty output it sets the data and returns, as other CPU primitives do (out.set_data(allocator::malloc(0))). For K = 0 it fills the output with zeros, as Matmul::eval_cpu does (matmul.cpp:111-118). Both return before any elem_to_loc.
• The two dispatchers count matrices as out.size() / (M * N), which does not depend on K.
• The mixed-path check requires x.shape(-1) > 0, as gather_qmm's check requires experts > 0. A K = 0 input takes the promoted path.
• Empty and K = 0 shapes now return before QQMatmul's NYI throw (quantized.cpp:1367) and the dtype throw of its dispatch_quantize_dequantize (:1153).
• quantized_matmul rejects a w whose batch is empty when x has no batch dimensions (std::invalid_argument, on every device). Batch dimensions broadcast only when both x and w have more than two. Otherwise the output has the shape of x, with N as its last dimension, and comes from one matrix of w: the CPU computes one output matrix and reads w's first matrix. With a batch of 0, w has no matrix, and K > 0, so the op's rules define no value for the [M, N] output. Zeros would not be an empty sum here, and an output of [0, M, N] would contradict the same rule for a batch of 3. So the shape is invalid, not a case to handle.

The other divisions in quantized.cpp cannot see a zero: group_size is validated, and the rest are compile-time constants.

Verification

At the branch tip:

Platform	Test cases	Assertions
x86-64 Linux (Clang)	269 passed, 0 failed, 16 skipped (the CPU-JIT cases)	4033
arm64 Linux (Clang)	269 passed, 0 failed, 16 skipped	4033
macOS (Accelerate, Metal off)	285 passed, 0 failed	4126

main passes 264 and 280 cases (3992 and 4085 assertions) in the same builds, so nothing newly fails. Logs: bf-qe-vsuite-{x86,arm,mac}.tests.log. The one new build warning is libstdc++'s result_of deprecation for the new encoder.dispatch lambda, which every dispatch lambda in the file also gets.

The test was mutation-checked, each case run alone. On x86-64, each removed guard crashes with SIGFPE in the cases it protects:

• The empty-output return: the [0, 64] case, qqmm with N = 0, and gather_qmm with M = 0.
• The K = 0 zero fill: the K = 0 case of each of the four tests.
• ops.cpp's x.shape(-1) > 0: the mixed-path case, while the op is built.
• ops.cpp's empty-batch check: the new case (bf-qe-vmutR-x86.tests.log).

Logs: bf-qe-mutA{1..4}-x86.tests.log, bf-qe-mutK{1..4}-x86.tests.log (the first four cases, before the buffer of ones was added).

With the zero fill removed but the K = 0 return kept, the K = 0 case of the first test fails on all three platforms, four times (affine and mxfp4, transposed and not), with no crash (bf-qe-vmutZ-{x86,arm,mac}.tests.log).

Upstream

ml-explore/mlx main at 32a183e12 has the same two divisions (quantized.cpp:403 and :641), the same elem_to_loc offsets, the same gather paths and the same QQMatmul route (:1345, :1377), so the same crashes. Its ops.cpp has no mixed bfloat16/float16 path, so :4802 is this fork's alone. The CPU part will be reported there with a link to this PR.

🤖 Generated with Claude Code