Fix fp quantized matmul corruption when the quantized dim is not a multiple of 32

[kapellirohith]

Problem

nvfp4 is the only quantization mode whose group size (16) lets the quantized
dimension of a matmul legally be a multiple of 16 but not of 32.
mx.quantize(w, mode="nvfp4") accepts K = 1040, and the fp quantized Metal
kernels tile that dimension by 32 without bounding the 16-wide tail. Whenever it
happens, every matrix-sized quantized_matmul / gather_qmm on the GPU
silently returns corrupted results. The CPU backend and the vector (decode)
kernels handle the same shapes correctly, so a model can decode perfectly and
corrupt during prefill, which is about as quiet as corruption gets.

Minimal reproducer (M3 Pro, macOS 25.5, main @ 39d9a8a):

import mlx.core as mx

M, K, N = 50, 1040, 128  # K % 32 == 16; M large enough for the tiled kernel
x = mx.random.normal((M, K)).astype(mx.float16)
w = mx.random.normal((N, K)).astype(mx.float16)
wq, s = mx.quantize(w, mode="nvfp4")

out = mx.quantized_matmul(x, wq, s, transpose=True, mode="nvfp4")
ref = x @ mx.dequantize(wq, s, mode="nvfp4").T
print((out - ref).abs().max())   # ~40; 72% of outputs wrong

stream=mx.cpu: 1.2e-3. M=1 (vector kernel): 1.1e-3. Aligned K=1024: 3e-3.
Only the GPU matrix path with K % 32 == 16 is wrong.

The three defects

All in mlx/backend/metal/kernels/fp_quantized.h, introduced with the Metal
nvfp4 support in #2946.

1. fp_qmm_t_impl ran its K loop past K_eff (for (int k = 0; k < K_eff; k += BK) with full-BK loads), so the final iteration read 16 columns of the
   next output row's packed weights and scales (past the buffer for the last
   row) into the accumulator. Metal shader validation on the unpatched kernel:
   Invalid device load at offset 2129920 … "nvfp4_qmm_t_float16_t_gs_16_b_4_alN_true_batch_0"
   (the weight buffer is exactly 2129920 bytes). Also backs the tiled
   gather_qmm, so the unsorted MoE path fails identically.

2. fp_qmm_n_impl stored full 32-column tiles for a 16-column tail. The
   store clipped rows only (store_result_safe(y, N, short2(BN, num_els))), so
   with N % 32 == 16 the last column tile stored 32 columns into rows only N
   long: the 16 spilled columns land on the next row's first 16 outputs (racing
   the neighbouring threadgroup) and past the end of the output buffer for the
   last row. Unpatched: 200 identical runs give 61 distinct results,
   corruption sits exactly at columns 0 to 15, and shader validation flags
   Invalid device store at offset 1064976 … "nvfp4_qmm_n_float16_t_gs_16_b_4_batch_0".

3. QuantizedBlockLoader::load_safe masked the wrong axis (if (reduction_dim == 1 && bi >= src_tile_dim.x)): bi is the thread's row,
   src_tile_dim.x the column-direction bound. In the K-remainder step of the
   sorted-MoE kernel fp_gather_qmm_rhs (tile_w = short2(k_remain=16, tgp_bn)), this zeroed the threads owning output columns 16 to 31 of each tile
   and silently dropped their last 16 K-elements. Positive ID: on the unpatched
   kernel the wrong columns are exactly those ≡ 16..31 (mod 32), and their
   values equal ref - x[:, K-16:] @ w[e][:, K-16:].T to within 0.05 (fp16
   rounding): the dot product minus exactly its tail. In the fp_qmm_t call
   sites the same compare is bi >= 32, dead code, so their intended column
   masking never ran.

Fix

• fp_qmm_t_impl: iterate K_eff / BK full tiles, then one bounded
  remainder tile, mirroring fp_qmm_n_impl's existing K handling. The split-K
  kernel fp_qmm_t_splitk passes K_eff = k_partition_size, which Fix incorrect nvfp4 quantized_matmul through the split-K path #3854 keeps
  a multiple of 32, so num_k == 0 and split-K is unchanged.
• fp_qmm_n_impl: take num_outs = min(BN, N - y_col), clip the w loads to
  it, and store with store_result_safe(y, N, short2(num_outs, num_els)).
  Removes the race and the out-of-bounds store.
• QuantizedBlockLoader::load_safe (fp_quantized.h and the affine twin in
  quantized.h): mask rows against .y and the thread's columns against .x.
  A thread owns n_reads * pack_factor contiguous columns and every legal
  partial extent is a whole number of quantization groups, so a thread is never
  straddled; the new static_assert(group_size % (n_reads * pack_factor) == 0)
  in both loaders pins that invariant.

A quantized dim is always a multiple of group_size, so it can only end in a
partial tile when group_size does not divide the block. Only nvfp4 does;
mxfp4 and mxfp8 have group_size == 32 == BK == BN, and every affine group
size is a multiple of the block. The remainder tile and the column bound are
therefore selected with if constexpr on group_size % BK != 0,
group_size % BN != 0 and group_size % BCOLS != 0, so they are discarded at
compile time for those modes.

This also removes an out-of-bounds read that affine and mxfp hit today

The row half of the mask is not nvfp4-only. In qmm_t_impl the w loader is
instantiated with BROWS = BN and reduction_dim = 1, and the call is
load_safe(short2(BK, num_outs)). The old compare was bi >= src_tile_dim.x,
that is bi >= BK == 32, while bi ranges over [0, BROWS) == [0, 32), so it
never fired: whenever N % BN != 0 the last tile read weight and scale rows
past N. The store already clipped those columns, so results were correct;
this is a read past the end of the buffer, not corruption. It reproduces on
affine, which is the most used mode:

import mlx.core as mx

M, K, N = 64, 32768, 100  # N % 32 != 0, M >= 33 for the tiled kernel
x = mx.random.normal((M, K)).astype(mx.float16)
w = mx.random.normal((N, K)).astype(mx.float16)
wq, s, b = mx.quantize(w, group_size=32, bits=4)
mx.eval(x, wq, s, b)
mx.eval(mx.quantized_matmul(x, wq, s, b, transpose=True, group_size=32, bits=4))

K = 32768 makes the packed weight buffer an exact multiple of the page size,
so the over-read crosses the allocation. Under Metal shader validation on
39d9a8a:

Invalid device load at offset 213056, executing kernel function:
"affine_qmm_t_splitk_float16_t_gs_32_b_4_alN_false"

With this change the same run is clean, as are the two nvfp4 shapes below. The
same applies to mxfp4 and mxfp8 with N % BN != 0.

On the loader mask: bi is a row index and bj a packed column index, and
every call site passes src_tile_dim as (valid columns, valid rows), swapping
which physical dim that is to match the layout. fp_qmm_t_impl passes
short2(BK, num_outs) with BROWS = BN, fp_qmm_n_impl passes
short2(num_outs, BK) with BROWS = BK, and fp_gather_qmm_rhs picks
short2(k_remain, tgp_bn) or short2(tgp_bn, k_remain) on transpose. That is
the convention steel's BlockLoader::load_safe already uses
(src_tile_dim - short2(bj, bi)), so the row index tests against .y in every
case and the mask does not depend on reduction_dim.

Validation

M3 Pro (applegpu_g15s), macOS 25.5, against main @ 39d9a8a. "ref" is fp64 on
the MLX-dequantized weights (bit-identical CPU vs GPU dequantize, asserted).

check	unpatched	patched
t=1, M=50, K=1040/4112 (3 seeds)	rel err 0.33 to 0.69, 46 to 72% of outputs >5% off	1.2e-3
t=0, M=50, K=512, N=1040	rel err 5.2 at cols 0 to 15	1.7e-3
t=0 determinism, 200 identical runs	61 distinct	1 (also 1/200 across 4 families × 3 seeds)
sorted gather, T=64, K=1040	cols ≡16..31(32) = ref minus K-tail (±0.05)	1.2e-3
unsorted gather, M=50, K=1040	rel err 0.82	2.0e-3
split-K (fp_qmm_t_splitk), K∈{2048,4096,8192}	(unaffected; K_eff a mult of 32)	2.6 to 3.9e-3
vjp of quantized_matmul + gather_qmm, both transposes	n/a	1.5e-6 to 8.1e-6
shader validation, nvfp4 t=1 read / t=0 store	invalid load @2129920 / invalid store @1064976	0 faults
shader validation, affine N % BN != 0	invalid load @213056	0 faults
added regression tests	17/19 subtests fail	pass
python/tests/test_quantized.py	n/a	36/36
full python test suite	n/a	804 OK, 3 skipped
deterministic misalignment sweep (all 4 kernels, fp32, 1e-3 gate)	n/a	289/289
randomized differential campaign (affine 2 to 8 bits gs 32/64/128 / mxfp4 / mxfp8 / nvfp4 × fp16/bf16/fp32 × both transposes × contiguous+strided × plain/gather/sorted-MoE, GPU vs fp64)	n/a	4515 cases, 0 GPU failures

The campaign's CPU cross-check uses a looser gate for fp16 and bf16, since the
CPU backend accumulates those differently over a large K; that is not code this
PR touches.

The added tests use M ≥ 33; get_qmv_batch_limit caps the vector/matrix
threshold at 32, so the tiled kernels are exercised on every Apple GPU family
(M1 to M4, including Max/Ultra). 17 of the 19 subtests fail on 39d9a8a and all
pass with the fix. The two that pass on both are controls: the vjp with
transpose=False has an aligned quantized dim, and the sorted gather with
transpose=False runs the loader with reduction_dim == 0, where the old row
compare was already correct.

To confirm the if constexpr guards cost the other modes nothing, I compiled
fp_quantized.metal against the base headers and against this branch and
compared every emitted kernel body after normalizing metadata ids. Of the 306
gs_32 kernels, 253 are byte-identical and none gained a single
instruction; 53 shrank, for a net of −141 instructions, while the +12950
instructions added by the fix land entirely on the nvfp4 kernels. The 53 that
shrank are exactly the reduction_dim == 1 families (qmm_t, qmm_t_splitk,
gather_qmm_t, gather_qmm_rhs_nt), where the row mask now short-circuits the
rows it used to read out of bounds.

Perf (aligned fast path, min-of-medians, ms, before -> after)

                              aligned         unaligned gather
nvfp4  qmm_t  512x4096x4096   4.247 -> 4.191
nvfp4  qmm_n  512x4096x4096   4.175 -> 4.151
mxfp4  qmm_t  512x4096x4096   4.274 -> 4.215
mxfp4  qmm_n  512x4096x4096   4.208 -> 4.157
mxfp8  qmm_t  512x4096x4096   4.340 -> 4.305
mxfp8  qmm_n  512x4096x4096   4.262 -> 4.209
affine qmm_t  512x4096x4096   4.217 -> 4.230
affine qmm_n  512x4096x4096   4.203 -> 4.146
nvfp4  gather_rhs 512/8/2048  1.514 -> 1.454
mxfp4  gather_rhs 512/8/2048  1.514 -> 1.453
mxfp4  gather_rhs nt T=500 K=2080 N=2080   1.619 -> 1.556
mxfp8  gather_rhs nt T=500 K=2080 N=2080   1.700 -> 1.637
mxfp4  gather_rhs nn T=500 K=2080 N=2080   1.636 -> 1.578
mxfp8  gather_rhs nn T=500 K=2080 N=2080   1.695 -> 1.627

Nothing regresses. The mxfp4, mxfp8 and affine rows are the ones that matter
for the if constexpr guards above, and the unaligned gather rows are the
kernels where the row mask now short-circuits work the base build was doing on
rows it read out of bounds.

Why fix the kernel rather than reject the dimension

group_size = 16 advertises that quantized dimensions which are multiples of 16
are supported; silently corrupting a legal, documented input is an API-contract
violation the kernel should honor rather than an input to reject at the op layer.
The inconsistency already shows: mx.quantize(w, mode="nvfp4", stream=mx.cpu)
throws a reshape error when w.size % 32 != 0 while the GPU path accepts the
same tensor and later corrupts. One backend rejects what the other silently
mishandles. Making the kernels correct for every legal group_size=16 dimension
resolves both.

Provenance

Introduced with the Metal nvfp4 kernels in #2946. #3854 narrowed only the
split-K route: for K ≡ 16 (mod 32) no split_k makes K % (split_k * 32) == 0,
so those shapes always take its split_k = 1 fallback straight into plain
qmm() -> fp_qmm_t_impl, the unbounded loop fixed here. The reproducer above
(M=50, K=1040, B=1, transpose=True) is exactly that path, so the shapes #3854
meant to protect were still corrupted.

Related but distinct: #3856 reports affine gather_qmm corruption at large unaligned
row counts (n > 32768, n % 64 != 0) on M5, a row-count (M) bug. This is a
quantized-dimension (K/N) bug specific to nvfp4; #3856's own repro does not reproduce
against this branch, and this fix leaves its K % 64 == 0 shapes on the same NAX
kernel, so the two do not overlap.

Tests

test_fp_qmm_non_multiple_of_32 covers K=1040/528 (transpose=True) and N=1040
(transpose=False, run twice and compared bit for bit to catch the store race),
each in fp32, fp16 and bf16, plus mxfp4 and mxfp8 at the same block-unaligned
shapes as a control that the if constexpr guards did not change them.
test_fp_gather_qmm_non_multiple_of_32 covers sorted and unsorted gather_qmm
in the same three dtypes, with the tail on K (transpose=True) and on N
(transpose=False). test_fp_qmm_non_multiple_of_32_vjp covers the backward of
both transposes. Inputs are scaled by 1 / sqrt(K) so the dot products are
O(1); at magnitude 32 a single fp16 ulp is 0.03 and would swamp the tolerance.

[PhilipJohnBasile]

Coordination note from M5/NAX testing in #3922:

This PR remains the authoritative fix for NVFP4's legal 16-wide tails and the required plain-kernel bounded load/store work. I deliberately did not duplicate those changes.

The independent M5 matrix clarified the merge interaction:

• current plain sorted kernels are clean for MXFP4/MXFP8 tails and for NVFP4 when K is a multiple of 32;
• NVFP4 K % 32 == 16 requires this PR's deeper plain-kernel repairs;
• Fix sorted gather_qmm NAX row overflow above 32K #3922 corrects affine NAX M/K tails and routes group-size-32 FP partial N/K tiles to the bounded path;
• after both land, the shared dispatch condition should retain this PR's NVFP4 fallback coverage, keep corrected affine tails on NAX, and preserve the group-size-32 ragged-N fallback from Fix sorted gather_qmm NAX row overflow above 32K #3922.

Independent native M5 tests for #3922 kept affine/MXFP4/MXFP8 tail errors below 4.7e-4; untouched NVFP4 non-64 K remains failing, which confirms the scope split rather than contradicting this PR.

Security scope for the related reports: the observed behavior is same-process incorrect tensor values and, in #3856, same-process allocator-pool reuse. We found no cross-process disclosure, arbitrary code execution, sandbox escape, or other trust-boundary crossing; current evidence supports numerical correctness bugs, not cybersecurity vulnerabilities.

[kapellirohith]

Thanks for the heads up, and for running this on M5. That's the piece I couldn't test.

Scope split looks right to me. This PR fixes the plain kernels for NVFP4 dims that are
16 mod 32 (group_size=16 is the only way to get one). The only dispatch change here is
the K % 64 == 0 guard on gather_qmm_rhs_nax, which was the one NAX entry missing it.
qmm_nax and gather_qmm_nax both already had it. So your NVFP4 non-64 K result is the
case that guard covers: those shapes get routed to the plain fp_gather_qmm_rhs this PR
repairs.

We both touch quantized.cpp and test_quantized.py so there's a conflict either way.
Happy to rebase if #3922 lands first.

[PhilipJohnBasile]

Thanks — this confirms both the scope split and the observed M5 NVFP4 path. The cleanest sequence is #3922 first, then rebase #3912 as offered: #3922 supplies the sorted gather_qmm boundary/correctness foundation, while #3912 owns the group-size-16 plain-kernel dispatch. If maintainers land #3912 first, I’ll rebase #3922 instead. No additional code change is needed from this finding.

[zcbenz]

You can putting this inside a if constexpr (group_size == 16) to avoid overhead on other quantization modes.

[zcbenz]

Similarly here if constexpr (group_size == 16) can be used to choose load_unsafe for faster load.

[zcbenz]

Should the condition depend on reduction_dim, i.e. bi >= src_tile_dim.y when reduction_dim == 0, and bi >= src_tile_dim.x when reduction_dim == 1?

[zcbenz]

I'm not sure on this, should drop it if you can not test it.

[kapellirohith]

All four addressed, rebased on main.

if constexpr on the K remainder tile. A quantized dim is always a multiple of group_size, so it can only end in a partial tile when group_size does not divide the block. Only nvfp4 can: mxfp4 and mxfp8 have group_size == 32 == BK == BN. The remainder tile is now behind if constexpr (group_size % BK != 0). I wrote the condition against BK rather than a literal 16 so it stays correct if another group size is added, and it degrades safely either way: a hypothetical group size that did not divide the thread's read span would trip the static_assert at compile time rather than miscompute.

Same on the qmm_n side, with if constexpr (group_size % BN != 0), and I did the same for the column half of the mask inside the loader with group_size % BCOLS != 0, which was the part that had been costing the other modes a compare. To check the guards rather than assume them, I compiled fp_quantized.metal against the base headers and against this branch and diffed every emitted kernel body after normalizing metadata ids. Of the 306 gs_32 kernels, 253 are byte-identical and none gained a single instruction; 53 shrank, net −141 instructions, while the +12950 instructions the fix adds land entirely on nvfp4. Timings match: mxfp4 qmm_t 4.274 to 4.215 ms, mxfp8 qmm_t 4.340 to 4.305, affine qmm_n 4.203 to 4.146, mxfp4 unaligned sorted gather 1.619 to 1.556.

Whether the mask should depend on reduction_dim. It should not, and the row half is not nvfp4-only. bi is a row index and bj a packed column index, and every call site passes src_tile_dim as (valid columns, valid rows), swapping which physical dim that is to match the layout: fp_qmm_t_impl passes short2(BK, num_outs) with BROWS = BN, fp_qmm_n_impl passes short2(num_outs, BK) with BROWS = BK, and gemm_loop_unaligned picks short2(tgp_bk, tgp_bn) or short2(tgp_bn, tgp_bk) on transpose. That is the convention steel's BlockLoader::load_safe already uses, src_tile_dim - short2(bj, bi). I built the reduction_dim dependent version to be sure: for reduction_dim == 1 the K remainder call short2(k_remain=16, num_outs) then zeroes every thread with bi >= 16, dropping the K tail for output rows 16 and up. qmm_t M=50 K=1040 N=128 goes to 6.9e-1 relative error and sorted gather_qmm to 6.1e-1, the wrong columns are again exactly those congruent to 16..31 mod 32, and 4 subtests fail. reduction_dim == 0 is unaffected either way, since .y is already the row bound there.

The same reasoning is why the quantized.h hunk is in this PR, and it is worth more than I gave it credit for in the description. With reduction_dim == 1 the old compare was bi >= src_tile_dim.x, that is bi >= BK == 32, while bi ranges over [0, BROWS) == [0, 32), so it never fired. Whenever N % BN != 0 the last tile read weight and scale rows past N. Affine hits this today. On 39d9a8a with affine gs=32, M=64, K=32768, N=100, Metal shader validation reports Invalid device load at offset 213056, executing kernel function: "affine_qmm_t_splitk_float16_t_gs_32_b_4_alN_false", and it is clean with this change. Results were never wrong there since the store already clipped those columns, so it is a read past the end of the buffer rather than corruption, but it is real and it also applies to mxfp4 and mxfp8.

NAX dispatch guard. Dropped, as asked. is_nax_available() needs gen >= 17 and this is an M3 Pro at gen 15, so I cannot exercise that path and quantized.cpp is untouched now. To be clear about what that leaves: nvfp4 with a quantized dim that is an odd multiple of 16 is still broken on NAX, because fp_quantized_nax.h has its own copy of the loader with the same mask and this PR does not touch it. That is a separate fix for someone with M5 hardware.

Revalidated on the result: three reproducers clean, 289/289 misalignment sweep, 4515 randomized differential cases across all modes, bit widths and dtypes with no failures, test_quantized.py 36/36, full suite 804 passing, shader validation clean on the three shapes that faulted before. The added tests now cover fp32, fp16 and bf16, the tail on both K and N, the vjp, and mxfp4/mxfp8 controls; 17 of the 19 subtests fail on 39d9a8a and all pass here.

[erwinzhang7]

Hey, please see #4009. Thanks!