Fix nvfp4 quantize and dequantize on the CPU for element counts that are not a multiple of 32

[kapellirohith]

Problem

mx.quantize and mx.dequantize in nvfp4 mode reject shapes on the CPU that the GPU accepts. Same op, same input, two backends disagree.

Base d9077d831
Chip M3 Pro
macOS 26.6.1

import mlx.core as mx

w = mx.random.normal(shape=(33, 16))
print([tuple(a.shape) for a in mx.quantize(w, mode="nvfp4", stream=mx.gpu)])
mx.eval(mx.quantize(w, mode="nvfp4", stream=mx.cpu))

Before:

[(33, 2), (33, 1)]
[reshape] Cannot reshape array of size 528 into shape (16,4,8).

After: both succeed, and the scales and dequantized values agree.

dequantize fails the same way on the same shapes (Cannot reshape array of size 66 into shape (16,4)). The minimal case is (1, 16), which asks for a reshape into (0, 4, 8), a literal zero dimension.

Cause

The 4-bit packing in the composed fp fallbacks groups the flattened values into fixed blocks of 32 before packing them into uint32 words:

• ops.cpp:5189, in fp_quantize: reshape(wq, {-1, 4, 8}, s)
• ops.cpp:5435, in fp_dequantize: view(reshape(out, {-1, 4}, s), int8, s)

Both require the flattened element count to be a multiple of 32. fp_quantize reshapes to {-1, group_size} first (ops.cpp:5136), so that count is R * K for a matrix of R rows and K columns.

nvfp4 is the only mode with group_size 16, so it is the only one whose element count can be a multiple of 16 but not of 32. mxfp4 and mxfp8 use group_size 32, and every affine group size is a multiple of 32, so none of them can reach this. Within nvfp4 the failing set is exactly those inputs whose flattened element count is divisible by 16 but not by 32.

This inconsistency was noted in passing in #3912 while fixing the GPU side of the same group_size = 16 tail. This is the op-layer half of it.

Fix

Group by the 8 nibbles that actually share a word, and view the packed words as bytes without grouping them first.

On the quantize side the 8 is load-bearing: it matches shifts, which is power(2, arange(0, 32, 4)) and therefore 8 wide, and the sum is over the last axis. The 4 was arbitrary. Row-major flattening makes {-1, 8} and {-1, 4, 8} produce the same flat word sequence before the final reshape(new_shape).

On the dequantize side gather builds its output shape as indices.shape followed by slice_sizes, so gather(lut, idx, 0, {1}, s) appends a trailing axis of size 1 and concatenate({lo, hi}, -1, s) interleaves the low and high nibble per byte regardless of the leading shape. The grouping reshape was only there to give view a shape; view rescales the last axis on its own, so it is not needed. View::eval_cpu (mlx/backend/cpu/primitives.cpp:367) handles a non-contiguous input itself: it shares the buffer with rescaled leading strides when the last axis is contiguous or the array is row contiguous, and otherwise makes a contiguous copy. Dropping the reshape therefore does not change what the byte view walks.

Both forms walk the same flat sequence, so the result is unchanged wherever the old form was legal.

Testing

All numbers measured against d9077d831.

Both hunks are required. Applying either one alone leaves the new test failing 6/6 on both devices, so neither changed line is dead.

check	before	after
new test, both devices, JIT and no-JIT	18 of 24 subtests error	pass
grid, nvfp4 quantize, CPU-only failures	240	0
grid, nvfp4 dequantize, CPU-only failures	240	0
grid, mxfp4 / mxfp8 / affine	0	0
hashed outputs changed	n/a	0 of 3954
hashed outputs newly succeeding	n/a	198
hashed outputs regressed	n/a	0
python suite, GPU, no-JIT	n/a	868 OK, 5 skipped
python suite, GPU, JIT	n/a	868 OK, 5 skipped
test_quantized, JIT	n/a	39/39
C++ suite, both devices	n/a	278/278
pre-commit	n/a	clean

The grid is 4 modes x 2 ops x rows 1 to 40 x K in {16, 32, 48, 64, 80, 96, 112, 128} x {float32, float16, bfloat16}. That is 3840 records, 1680 skipped because K % group_size != 0, and 2160 evaluated. The 240 failures per op are 80 per dtype, which is the 20 odd row counts times the 4 K values where K / 16 is odd, exactly the set predicted above. No shape fails on the Metal backend before or after.

The byte-identity check hashes every quantize and dequantize output on both builds across 15 row counts, 9 K values, 5 higher-rank shapes, 4 modes, 3 dtypes and both backends. Of the 3954 outputs that exist on both, none changed. The 198 that newly appear are 99 quantize and 99 dequantize, all nvfp4, all CPU, which is 33 tail shapes times 3 dtypes times 2 ops. Nothing regressed.

Under DEVICE=cpu the suite has 8 failures. They are identical to a pristine rebuild of the same base: 6 test_conv2d_winograd_batch_tiling subtests, test_fft_too_large, and test_export_import.test_leaks. All pre-existing and unrelated.

The equivalence argument was also checked directly rather than only reasoned about, by rebuilding both graph forms with mx ops and comparing them: 13 quantize and 26 dequantize comparisons on contiguous inputs, plus 7 non-contiguous packed views (row slice, strided rows, reversed, column slice, strided last axis, swapaxes, broadcast). Zero mismatches.

CUDA validation is source-level only. No NVIDIA hardware was used.

Tests

test_nvfp4_element_count_alignment covers (1, 16), (3, 16), (33, 16), (5, 80), (7, 112) and (1, 1, 16), which include the zero-dimension minimal case and a higher-rank shape, plus the controls (2, 16) and (1, 32) whose count is already a multiple of 32, each in float32, float16 and bfloat16. Inputs are built from the fp4 lookup table with the group maximum forced to 6, so the scale is exact and the round trip is exact, matching the construction the neighbouring test_nvfp4_quantize_dequantize already uses. The test asserts the CPU round trip is exact and, where a GPU is available, that the CPU and GPU scales and dequantized values are equal.

It asserts on scales and dequantized values rather than on packed code bytes. The CPU argmin picks lut index 0 (+0.0) where the GPU encoder keeps the sign bit and emits index 8 (-0.0), on shapes that already worked as well as new ones, so packed bytes are not a valid equality target on either backend.

Checklist

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

[zcbenz]

Please hold on creating new PRs before we review your previous ones, we can't review as fast as you create because: they are really minor problems that no ones cares in practice, it takes a lot of time to verify whether the fix is correct, and we can't trust any word in the PR description.