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.