Fix int16 overflow in the gather_qmm sorted row bound

[erwinzhang7]

Fixes #3856.

A quantized MoE silently returns wrong results for a single long prefill on M5. The same tokens fed through in chunks are correct, so the corruption is invisible unless you compare the two.

This is independent of #4009. I verified that by running the issue's repro on 0.32.0 and on a build carrying #4009, and the outputs were identical, so the two defects do not interact.

Root cause

In affine_gather_qmm_rhs_nax:

const short sgp_sm = align_M ? SM : min(SM, short(max(0, (M - (y_row + tm)))));
const short sgp_sn = align_N ? SN : min(SN, short(max(0, (N - (y_col + tn)))));

The narrowing to short happens inside the min. Once M passes the int16 range the inner expression wraps negative, the min returns that negative value, and the tile is then loaded with a garbage row bound. Taking the min in int and narrowing after is safe, since the result is bounded by SM.

Two details make this hard to hit by accident:

1. The unaligned branch only runs when align_M is false, so M % BM != 0 is required as well as M being large.
2. The line directly above it already does it in the safe order, which is why nothing else in the kernel misbehaves:

const short tgp_bm = align_M ? BM : short(min(BM, M - y_row));

fp_quantized_nax.h:864 and every other sgp_sm site in the tree already use the min(int(SM), ...) form. These two lines were the only outliers.

Trigger condition, and a correction to the issue title

The issue reports the trigger as sequence length % 32 != 0. The actual condition is the kernel row count M > 32767 and M % 64 != 0.

In SwitchGLU the rows are sorted and flattened, so M = tokens * experts_per_token. Qwen3-Coder-30B-A3B routes to 8 experts per token, giving M = 8L, so M % 64 != 0 reduces to L % 8 != 0. The lengths sampled in the issue (16000, 16032, 16064 clean; 16065, 16066 corrupt) happen to satisfy both % 32 and % 8, so they cannot distinguish the two rules. L % 8 is the correct one.

Evidence

Synthetic, no model required. gather_qmm with sorted_indices=True against sorted_indices=False on the same already sorted input, E=4, N=256, K=512, group_size 64, 8 bit. K is aligned here so #4009 is not involved. Max difference relative to mean(abs(unsorted)):

M	M % 64	M > 32767	before	after
4000	32	no	0.0035	0.0035
32704	0	no	0.0035	0.0035
32760	56	no	0.0035	0.0035
32768	0	yes	0.0035	0.0035
32800	32	yes	4.894	0.0035
32832	0	yes	0.0035	0.0035
40000	0	yes	0.0035	0.0035
40001	1	yes	6.517	0.0035
65600	0	yes	0.0035	0.0035

Every corrupt case satisfies both conditions and no clean case satisfies both.

End to end on Qwen3-Coder-30B-A3B-Instruct 8 bit, M5 Max, comparing a single prefill against 2048 token chunks. Max absolute difference in the final position logits:

L	L % 8	before	after
8192	0	0.156	0.156
12288	0	0.352	0.352
16000	0	0.367	0.367
16032	0	0.250	0.250
16064	0	0.281	0.281
16065	1	9.875	0.156
16066	2	8.875	0.227

Worst KV cache difference at the two corrupt lengths drops from 43.75 and 43.0 to 20.0, matching the aligned baseline.

One note on reading that table: a residual argmax flip remains at L=16000, unchanged by this patch. It is not corruption. At that position the top two logits are exactly tied in fp16 (both 9.6250, gap 0.0000), so which one wins is decided by ordinary accumulation noise between the one shot and chunked paths. Every length now sits in the same 0.15 to 0.37 band.

Test

Added test_gather_qmm_sorted_large_m, covering M of 32704 (under the int16 range, as a control), 32800 and 40001. It fails on 0.32.0 on the latter two and passes with this change.

test_quantized.py, test_blas.py, test_nn.py and test_ops.py pass, 279 tests.

Performance

No measurable change. The edit only widens a scalar min computed once per thread, and the difference is well inside the run to run variance of a gather_qmm benchmark.

[zcbenz]

Closing since there are already PRs with similar changes.