Fix sorted gather_qmm NAX row overflow above 32K

[PhilipJohnBasile]

Summary

Fixes #3856.

The sorted-RHS affine NAX kernel narrowed the remaining row count to short
before clamping it to the SIMD-group tile size. For ragged workloads with
M > 32768, that intermediate conversion can wrap negative, causing early
tiles to leave output rows unwritten.

This keeps the remaining-row arithmetic in int through the clamp and narrows
only the final value. The regression test covers the exact signed-short seam,
an aligned control above it, a dense fp32 oracle, and two allocator-poison
values.

Scope

This recut intentionally contains only the failure class unique to #3856.

• Ragged-K handling from [BUG] gather_qmm sorted-rhs path corrupt for K % 64 != 0 on M5/NAX: !align_K tail bounds the load with BK instead of the K remainder (affine + mxfp4) #3887 remains with Fix sorted gather_qmm on ragged K #4009.
• NVFP4/plain-kernel work remains with Fix fp quantized matmul corruption when the quantized dim is not a multiple of 32 #3912.
• quantized.cpp, K/N-tail dispatch, and loader behavior are unchanged.
• The adjacent sgp_sn expression is unchanged: its consumer is reached only
  on the final partial-N tile, where the remaining extent is already below 64.

Validation

Rebased onto upstream main at
47a1bd613ebb67ed03da764c6ffedb8df9939997; candidate head
bd8c825c2cad3d5ae3a6462f599c1a2320eeb9dd. The editable build reports
mlx 0.32.1.dev20260810+bd8c825c2 and applegpu_g17s on Apple M5 Max.

• focused native NAX regression: passed;
• negative control from the original frozen receipt: restoring only the old
  pre-clamp short conversion failed exactly at M=32769 (31.2725 maximum
  error versus the 0.05 limit), while controls remained clean;
• full GPU TestQuantized: 36/36 passed;
• forced applegpu_g16s fallback control: passed;
• clang-format, Black, isort, and git diff --check: passed;
• independent frozen-diff review: approved with no P0-P2 findings.

No performance claim is made; the change only corrects the row bound used by
the existing kernel.

Attribution

Philip John Basile is the primary author. Development and review used OpenAI
Codex assistance, disclosed in the commit trailer.

[PhilipJohnBasile]

Independent final-patch review on the same M5 Max completed after the exact-source rebuild:

• affine 4/8-bit, K={32,64,96,160}: max error <= 4.25e-4;
• MXFP4: <= 4.07e-4; MXFP8: <= 4.67e-4;
• N={64,65,80,95,96,97,127,128} at K=64: affine/MXFP4/MXFP8 <= 2.44e-4;
• aligned NVFP4 controls and those N tails: <= 2.44e-4;
• focused new tests pass and git show --check is clean.

The reviewer also confirmed the declared scope boundary: untouched NVFP4 non-64 K tails still fail on current main/this branch and remain assigned to #3912. This PR does not claim otherwise.

[vvsotnikov]

I can confirm this works on M5 Max, thank you for the fix!

LLM-generated report:

Tested on M5 Max, macOS 26.5, MLX 0.32.0, a MoE model with 2-bit experts.

I sent tokens through the model twice: once in a single forward pass, once in
chunks of 2048. Both must give the same result. I compared the output values at
the last token and recorded the largest difference.

The model sends 10 rows to the gather kernel per token, so 4001 tokens make
40010 rows (not a multiple of 64) and 4000 tokens make 40000 rows (a multiple
of 64).

Tokens	Rows multiple of 64	Before fix	After fix
4001	no	5.16	1.65
4000	yes	1.9 to 2.5	—

A difference of about 2 is normal, because the two runs add numbers in a
different order. 5.16 is the error. After the fix it drops to 1.65, so the
error is gone.

Prefill speed is 1903 tokens per second before and after the fix.

[erwinzhang7]

Hey, please see #4009. Thanks!

[PhilipJohnBasile]

Thank you to the MLX maintainers and everyone who has taken the time to test and review these fixes. I really appreciate the opportunity to contribute. I’ll keep the remaining overlap and correctness questions in view so the patch is as straightforward as possible to evaluate.

[PhilipJohnBasile]

Recut onto current main at 9cad37fed8472eba79990d7205b8c093d3080446.

The PR is now intentionally limited to #3856's unique sorted-RHS NAX row-bound
overflow. The earlier K/N-tail, loader, dispatch, and quantized.cpp changes
have been removed; #3887 remains with #4009, and the NVFP4/plain-kernel work
remains with #3912. Earlier comments describing the combined patch refer to the
superseded d86fc875 head.

On M5 Max (applegpu_g17s), the focused test passed on the fixed tree, failed
only at M=32769 after restoring the old pre-clamp short conversion
(31.2725 max error versus 0.05), and passed again after restoration. The
full GPU TestQuantized suite passed 35/35, and the forced gen16 fallback
control passed.

The current remote diff is exactly three files: the one-line kernel correction,
its focused regression, and the contributor acknowledgment requested by MLX's
README.

[PhilipJohnBasile]

Rebased onto the latest upstream main (8c28c38) after the recent merges. The current head is 41aba66. The diff remains intentionally limited to the three files already described, with no conflicts or whitespace errors; the focused M5 Max regression and full TestQuantized validation were completed on the unchanged patch before this rebase. Ready for CI approval and maintainer review.

[PhilipJohnBasile]

Rebase and exact-head verification are complete at bd8c825c2cad3d5ae3a6462f599c1a2320eeb9dd on upstream main 47a1bd613ebb67ed03da764c6ffedb8df9939997.

The frozen candidate was independently reviewed after the M5 Max rerun. Verdict: APPROVE, with no P0-P2 findings. The reviewer confirmed that the patch remains identical across the rebase, the int clamp makes the final short narrowing safe for both current SM=16 and SM=32 instantiations, and open #4009/#3912 cover separate failure classes.

Exact-head checks:

• focused native NAX regression: pass;
• complete TestQuantized: 36/36 pass;
• forced MLX_METAL_GPU_ARCH=applegpu_g16s fallback: pass;
• clang-format, Black, isort, and git diff --check: pass.

One non-blocking review note: the >32K regression selects BM=64 with E=16; the broader suite exercises the newly merged BM=32 dispatch at smaller sizes. No performance claim is made.

[PhilipJohnBasile]

@erwinzhang7 — thank you for the measurement write-up; the four-failure-class matrix is the best documentation this area has had, and keeping each fix as its own commit is the right call now that squash-merges are in play.

Two status points from the #3922 side:

1. The Normalize biases before encoding in gather_qmm_rhs #4056 overlap you flagged is resolved in the current head. My rebase base (upstream main @ 47a1bd613) postdates Normalize biases before encoding in gather_qmm_rhs #4056's merge, and the patch on top is still exactly the three files already described: the one-line kernel correction, its focused regression, and the acknowledgment. Verified on the rebased head: focused native NAX regression passes, complete TestQuantized 36/36, forced applegpu_g16s fallback passes. So whichever of Fix sorted gather_qmm NAX row overflow above 32K #3922/Fix sorted gather_qmm on ragged K #4009 lands second only needs the partial-tile guard rebased, and I'll produce that the moment the other one merges.
2. On the open question in your notes — why the FP column term is unsafe for mx formats under a narrowed guard, and what it would take for ragged group-size-32 FP to stay on NAX (the ~1.75x end state): I'd like to pick that up as a follow-up PR once the family lands. It reads like a clean self-contained study (error ramp of the column term vs K for mxfp4/mxfp8, versus the affine path) rather than more scope on any of the three pending PRs.

@zcbenz — for review-sequence comfort: #3912 is hardware-agnostic and already 28/28 on upstream CI, so it can land anytime; #3922 and #4009 are NAX-only and touch the same partial-tile guard, so whichever you approve first leaves the other an easy rebase. Happy to re-scope, split, or repackage however is easiest to verify.

[erwinzhang7]

Just checked, and the overlap is indeed closed.

The column term follow-up is all yours. I scoped it out of #4009 and wasn't planning to come back to it. The 1.75x figure came out of my ragged-K sweep, so you may wanna measure it again. Thanks for checking though!

Agreed on sequencing, I can rebase #4009 if yours lands first.

[pierre427]

Independent validation on M5 Max (applegpu_g17s, macOS 26.5.2) — we filed #3856, so here is
a check on hardware that reproduces it. Built the merge-base (47a1bd613) and this PR's head
(bd8c825c2) from source with MACOSX_DEPLOYMENT_TARGET=26.2.

The added regression test does its job. On the merge-base it fails at exactly one subtest —
M=32769, max_error 31.27 (that is the -31.0 poison value showing through unwritten output
rows) — while M=32767, 32768 and 32832 pass. That is the short-overflow predicate exactly:
M > 32768 and M % 64 != 0. On this PR's head the test passes.

Full TestQuantized on the head: 36/36 OK.

Extra coverage worth having, since #4023 landed after this PR was written. #4023 picks BM
from rows-per-expert, so the BM=32 path is only exercised at small rows-per-expert while the

32K regression at E=16 selects BM=64. Sweeping E at M=32769 (N=K=64, fp16, same poisoned
idiom as the new test):

E	rows/expert	merge-base 47a1bd613	this PR bd8c825c2
16	~2048	max_err 47.29	1e-4
64	~512	max_err 47.26	1e-4
256	~128	max_err 31.29	1e-4
512	~64	max_err 31.32	1e-4

Corrupt in all four on the merge-base, clean in all four here — so the fix holds across both BM
dispatch regimes, not just the one the new test pins.

NAX-only, as expected. The dispatch guard
(quantized.cpp, is_nax_available() && transpose && (enable_tf32() || x.dtype() != float32))
means fp32 activations with MLX_ENABLE_TF32=0 bypass the NAX path entirely; that same
M=32769 shape is clean on the merge-base through the fallback. Consistent with this being
unreproducible on non-NAX hardware.

Scope check, no action implied: the K % 64 tail (#3887) is untouched by this PR, as
intended given #4009 — K=96 → max_err 1241, K=160 → 261 on both the merge-base and this
head, with aligned K clean at 1e-4. The split between the two PRs is clean.

One correction to our own earlier analysis: when we filed #3856 we flagged a suspected twin at
the sgp_sn line just below this fix (same short-cast on the N axis, predicted to need
N > 32768 && N % 64 != 0). We could not reproduce it — N = 32769 and N = 32833 are
clean on both builds. So the M axis looks like the only live one, and nothing here argues for
widening this PR's scope.

[PhilipJohnBasile]

Rebased onto current main (31b5cbb09), which now includes #3941 and #4051. The kernel change is unchanged: the one-line int-clamp on sgp_sm (see is_unaligned_sm downstream). I reviewed the changed region against the current quantized_nax.h: #3941/#4051 touch the non-transposed qmm_n_nax_tgp_impl and the empty-output skip paths, not this sorted-gather kernel, so the patch is unaffected.

One note on local verification: my toolchain is Command Line Tools only (no full Xcode), so the Metal shader compile and TestQuantized can't run here — relying on project CI for the Metal pass. Ready for review.

[PhilipJohnBasile]

Local full-Metal verification is now complete (this was the remaining gap — the previous runs depended on project CI).

• Built mlx from source on the PR head: mlx 0.32.1.dev+34924ded1, Metal enabled, M5 Max, macOS 27.0, Xcode 27 / clang 21.
• test_gather_qmm_sorted_nax_large_m: 4/4 subtests pass (M = 32767/32768/32769/32832 with allocator-poisoned outputs).
• TestQuantized: 36 passed / 3079 subtests passed on this head.

Note: the 48 test_qmm failures seen on current main at group_size=32, transpose=False are a separate defect fixed by #4202 (non-transposed NAX quantized matmul); they are independent of this PR — confirmed by running the same suite with and without this change. The patch here remains only the sorted gather_qmm row-overflow clamp + regression + acknowledgment.