Fix incorrect nvfp4 quantized_matmul through the split-K path

[metascroy]

Summary

nvfp4 quantized matmuls that take the split-K path (qmm_splitk / fp_qmm_t_splitk) produce incorrect results — non-uniform ~2x error and NaN/inf. affine is unaffected.

Root cause

qmm_splitk caps split_k by the quantization-group count but not by the kernel's K-tile width BK (=32):

split_k = std::min(split_k, K / group_size);
while (split_k > 1 && (K % (split_k * group_size) != 0)) split_k--;

For nvfp4 (group_size == 16) this can pick a per-partition K smaller than BK. Example: K=64, group_size=16 → split_k=4, k_partition_size=16 < BK=32.

The qmm_t kernels tile K in BK-wide chunks and do not bound the K dimension — load_safe/load_unsafe only bound M/N, and QuantizedBlockLoader advances scales by scale_step = BCOLS/group_size per BK step. So when a partition is smaller than BK, the single tile load reads a full BK-wide slice that spills past the partition into the next group's packed weights and fp8 scales (and past the buffer on the last partition → NaN/inf). affine never hits this because group_size >= 32 keeps partitions >= BK (or split_k collapses to 1 and falls back to qmm).

Fix

Require each K partition to be a whole number of BK-wide tiles as well as whole quantization groups, by aligning split_k to max(group_size, BK):

int k_align = group_size > 32 ? group_size : 32; // BK
split_k = std::min(split_k, K / k_align);
while (split_k > 1 && (K % (split_k * k_align) != 0)) split_k--;

This only changes behavior for group_size < 32 (i.e. nvfp4); for group_size >= 32 it is identical to today. It preserves split-K where valid (e.g. the K=64 case still runs 2-way with partition=32=BK) and otherwise falls back to qmm via the existing split_k <= 1 path.

Testing

• nvfp4 quantized_matmul with M >= vector_limit, transpose=True, single batch (e.g. x=[32,64], w=[128,64], group_size=16): incorrect (NaN + ~2x error) before, matches the reference after.
• affine and larger group_size unchanged.

Open questions for reviewers

• BK is hardcoded as 32 here to match the qmm_t_splitk template default (consistent with the existing local bm/bn = 32). Happy to use a named constant / template-derived BK if preferred.
• max(group_size, BK) is correct because all supported group sizes (16/32/64/128) are in a divisor relationship with 32; std::lcm(group_size, BK) would be the fully general form.
• Should mxfp4 (same fp_qmm_t_splitk kernel) get a regression test too? Its default group_size=32 avoids the sub-BK partition.

[pierre427]

Independent confirmation of both the bug and the fix, plus scoping data suggesting the trigger
is broader than the sub-BK partition case — with a regression-test implication.

We hit this in a kernel-dispatch seam audit on an M5 (~25k quantized_matmul /
gather_qmm (shape × quant-mode) cells arbitrated against an fp32-dequantized dense
reference) and root-caused it to the same missing BK alignment in qmm_splitk before
finding this PR.

The corruption triggers for any k_partition_size ≡ 16 (mod 32), not only partitions
smaller than BK. Example: K = 544 → split_k = 2, k_partition_size = 272. The
K-loop consumes whole BK = 32 tiles, so each partition's final tile overreads 16 columns
into the next partition (double-accumulated after the split-K reduce — the ~2× cells), and
the last partition reads 16 columns past the buffer (finite garbage, or inf/NaN on
batched shapes — e.g. B=4, m=33, K=2080 produced inf cells in our sweep).

Corrupt-cell predicate, verified exhaustively across the sweep on both 0.32.0 (PyPI) and
0.32.1.dev20260716 (fp16 and fp32 activations corrupt identically; 3D row-contiguous
batches collapse to M = B·m): with G = K/16 and
split_k = largest s ≤ min(max(1, 512/(⌈M/32⌉·⌈N/32⌉)), G) dividing G, output is
corrupt iff the matmul path is taken (M above the qmv batch threshold), split_k > 1,
and G/split_k is odd — i.e. exactly k_partition_size % 32 == 16.

Minimal repro (M5, both versions above):

import mlx.core as mx

mx.random.seed(7)
w = mx.random.normal((1024, 544)).astype(mx.float32) / 544**0.5
w_q, scales = mx.quantize(w, mode="nvfp4")
w_ref = mx.dequantize(w_q, scales, mode="nvfp4").astype(mx.float32)

mx.random.seed(8)
x = mx.random.normal((32, 544)).astype(mx.float16)
y = mx.quantized_matmul(x, w_q, scales, transpose=True, mode="nvfp4")
ref = x.astype(mx.float32) @ w_ref.T
print(float(mx.abs(y.astype(mx.float32) - ref).max()))
# 1.10 (unit-scale outputs)     K=512 aligned control: 0.0036
# m=8 (vector path) is clean:   0.0010

Since every k_partition_size is a multiple of 16, partition % 32 ∈ {0, 16} — so the
alignment guard here removes exactly the corrupt set; nothing survives it. For the
regression test it may be worth covering a partition > BK case like K=544 in addition
to the sub-BK K=64 case, so a future looser guard can't pass with only the small shape.

Happy to re-run the full seam sweep on an M5 with this patch applied if useful.

[angeloskath]

Thanks for the fix!