[rcfa]

What fails

On the CPU, qqmm with nvfp4 global scales ignores them and gives no error. Line numbers below are at main = 4a22e4fed.

• qqmm validates global_scale_x and global_scale_w (mlx/ops.cpp:4890-4901: nvfp4 only, one float32 element each, both or neither). It passes them as two more inputs after the scales of w (:4975-4976), and CUDA and Metal read them there.
• QQMatmul::eval_cpu (mlx/backend/cpu/quantized.cpp:1331-1369) reads only inputs[0..2] (:1337-1339):
  • x's quantize-dequantize (:1362, fp_quantize_dequantize :1103) takes each block scale as e4m3(amax / 6). That is the encoding for a global scale of 2688, whatever g_x is.
  • The matmul (:1363) uses the raw e4m3 block scales of w, without g_w / 2688.
• So when g_w is not 2688, the product is 2688 / g_w times too large. And x blocks whose scale is under e4m3's range without g_x become 0.

The constant is F8E4M3_MAX * F4E2M1_MAX = 448 * 6 = 2688. The op-level quantize encodes with 2688 / g (ops.cpp:5149-5155), and the op-level dequantize decodes with g / 2688 (:5455-5458). CUDA and Metal do the same: CUDA fp_quantize.cuh:80-83 and qmm_naive.cuh:205-207, Metal kernels/fp_quantized.h:215-216.

Why no test caught it

python/tests/test_quantized.py:281 passes the global scales to qqmm only on the GPU, with "TODO: Fix qmv with global scale in CPU backend." No C++ test runs qqmm.

The test

A new case in tests/quantized_stream_tests.cpp, on the CPU stream: nvfp4 qqmm with x of shape [1, 32] and w of shape [4, 32]. Every value is on the e2m1 grid, and every group of 16 holds 6 and -6, so every step is exact. It runs four cases:

• A: g_x = 2688, g_w = 336. The block scales of w are stored as 8, and only g_w / 2688 = 1/8 gives w back.
• B: x times 2^-12, g_x = 0.65625 (an encode factor of 4096), g_w = 2688. Without g_x the block scales of x, 2^-12, round to 0 in e4m3.
• C: no global scales.
• D: a float16 x (w's first row), g_x = 2688, g_w = 10.5. The block scales of w are stored as 256 and decoded with exactly 2^-8. The product is [282, 188, 115, 45], exact in float16, but 256 times it is out of float16's range for the first element (72192).

The reference is the op-level dequantize of x's and w's quantization, with the global scales, and a float32 matmul. It does not use the CPU qqmm kernels. The test also checks that this dequantize gives back x and w exactly. All comparisons use array_equal.

Without the fix, A, B and D fail on x86-64 Linux, arm64 Linux and macOS. A gives 8 times the product, B gives 0, and D gives 256 times the product (inf for the first element):

tests/quantized_stream_tests.cpp:269: ERROR: CHECK( array_equal(y, expected, false, s).item<bool>() ) is NOT correct!
  values: CHECK( false )
  logged: c.name := A

(the same for c.name := B and c.name := D; logs bf-gs-redbase-{x86,arm,mac}.tests.log).

The fix

• QQMatmul::eval_cpu reads inputs[3] and inputs[4] when the mode is nvfp4 and both are given. The task reads them, as scale_enc = 2688 / g_x and scale_dec = g_w / 2688. Both are 1 without global scales.
• x's quantize-dequantize encodes each block as the op-level quantize does: s = e4m3(amax / 6 * scale_enc), q = e2m1(x * scale_enc / s), x_hat = s * q / scale_enc, and 0 when s is 0.
• Each block scale of w is multiplied by scale_dec in fp_qmm, fp_qmm_t and fp_qmm_t_simd: T(float(e4m3(s)) * scale_dec).
  • It goes on the block scale, not on the sum. That is the op-level dequantize's order (from_fp8(scales) * (g / 2688), ops.cpp:5455-5458), and the code of Add support for x86 SIMD with Highway ml-explore/mlx#3019 (mlx/backend/cpu/quantized.cpp:1371 at 7c614ee17).
  • On Linux, and for bfloat16 on macOS (there is no max_size<bfloat16_t>), the CPU takes fp_qmm_t, which sums in T. A factor on the sum would carry sums that are 2688 / g_w times too large. In float16 they overflow when g_w is small. bfloat16 has float32's exponent range, so in practice only float16 overflows. The test's case D shows it (see the mutation checks).
  • With scale_dec = 1, the block scale is the same value converted T to float to T. So quantized_matmul and qqmm without global scales compute as before, bit for bit. With scale_enc = 1, x's round trip is also unchanged: x * 1 and / 1 are exact.
• scale_dec defaults to 1 in fp_qmm_dispatch and fp_qmm_dispatch_transpose, so the quantized_matmul and gather_qmm calls are unchanged.

GatherQQMM::eval_cpu still throws NYI, so it does not drop its global scales silently.

Verification

At the branch tip:

Platform	Test cases	Assertions
x86-64 Linux (Clang)	265 passed, 0 failed, 16 skipped (the CPU-JIT cases)	4004
arm64 Linux (Clang)	265 passed, 0 failed, 16 skipped	4004
macOS (Accelerate, Metal off)	281 passed, 0 failed	4097

main passes 264 and 280 cases (3992 and 4085 assertions) in the same builds, so nothing newly fails. Logs: bf-gs-suite2-{x86,arm,mac}.tests.log. The build warnings of quantized.cpp.o are as on main (14 on Linux, 0 on macOS).

The new case alone passes on all three (bf-gs-green2-{x86,arm,mac}.tests.log). For float32 and float16, the Linux runs use fp_qmm_t and the macOS runs use fp_qmm_t_simd, so both kernels are covered.

Mutation checks, the new case alone:

• g_w / 2688 moved from each block scale onto the sum, in fp_qmm_t and fp_qmm_t_simd: only D fails on x86-64 and arm64 Linux, where the float16 sum overflows to inf (bf-gs-mutsum-{x86,arm}.tests.log). It passes on macOS, where fp_qmm_t_simd sums in float (bf-gs-mutsum-mac.tests.log).
• scale_dec left at 1 (g_w dropped, g_x kept), on the first version of the test (cases A, B and C, as named now, before the float16 case D was added): only A failed, on all three platforms (bf-gs-mutgw-{x86,arm,mac}.tests.log).

Upstream

ml-explore/mlx main at 7d2d94a55 has a byte-identical QQMatmul::eval_cpu (quantized.cpp:1346-1384) and the same CPU TODO (python/tests/test_quantized.py:362), so it has the same bug. ml-explore#3019 (open) would fix it there: it applies 2688 / g_x in x's round trip and g_w / 2688 to each block scale.

Merge note

This PR and #30 (fix/cpu-qmm-empty) both add test cases at the end of tests/quantized_stream_tests.cpp, so that file conflicts. git merge-tree against #30's head (a5b93e484) merges mlx/backend/cpu/quantized.cpp cleanly. Whichever merges second rebases.

🤖 Generated with Claude Code