From bbd40f8c136c8fc78abc548713e58f4a362107ef Mon Sep 17 00:00:00 2001 From: Erwin Zhang Date: Wed, 5 Aug 2026 09:38:24 -0400 Subject: [PATCH] Fix int16 overflow in the gather_qmm sorted row bound The NAX sorted gather kernel narrowed the per-simdgroup row and column bounds to short before taking the min: const short sgp_sm = align_M ? SM : min(SM, short(max(0, (M - (y_row + tm))))); For a row count above the int16 range the inner cast wraps negative, so the min returns a negative bound and the tile is loaded with garbage limits. The guard only runs when align_M is false, so both a large M and an M that is not a multiple of BM are required to reach it. Taking the min in int and narrowing afterwards is safe because the result is bounded by SM. This matches fp_quantized_nax.h and every other site. In a quantized MoE the row count is tokens * experts_per_token, so this is reachable at long context. For Qwen3-Coder-30B-A3B (8 experts per token) it corrupts any single prefill whose length is not a multiple of 8, past about 4k tokens. --- mlx/backend/metal/kernels/quantized_nax.h | 8 +++--- python/tests/test_quantized.py | 30 +++++++++++++++++++++++ 2 files changed, 34 insertions(+), 4 deletions(-) diff --git a/mlx/backend/metal/kernels/quantized_nax.h b/mlx/backend/metal/kernels/quantized_nax.h index e67be9a06d..039ccdfe9b 100644 --- a/mlx/backend/metal/kernels/quantized_nax.h +++ b/mlx/backend/metal/kernels/quantized_nax.h @@ -1529,10 +1529,10 @@ template < const short tm = SM * (simd_group_id / WN); const short tn = SN * (simd_group_id % WN); - 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))))); + // Take the min in int and narrow afterwards. Narrowing first overflows for + // M or N above the int16 range, which wraps the bound negative. + const short sgp_sm = align_M ? SM : min(int(SM), max(0, (M - (y_row + tm)))); + const short sgp_sn = align_N ? SN : min(int(SN), max(0, (N - (y_col + tn)))); const bool is_unaligned_sm = align_M ? false : (sgp_sm != SM); const bool is_unaligned_bn = align_N ? false : (tgp_bn != BN); diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index 63254ee9c7..b4f4e9a06e 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -1403,6 +1403,36 @@ def scatter_unsort(x, inv_order, shape=None): self.assertTrue(mx.allclose(y1, y3, atol=tol)) self.assertTrue(mx.allclose(y1, y4, atol=tol)) + def test_gather_qmm_sorted_large_m(self): + # The sorted path bounded the per-simdgroup row count by narrowing to + # short before taking the min, which wraps negative once the row count + # passes the int16 range. The guard only runs when M % BM != 0, so both + # a large M and an unaligned M are needed to reach it. In a MoE this is + # tokens * experts_per_token, which passes 32767 at long context. + E, N, K = 2, 128, 256 + key = mx.random.key(0) + k1, k2, k3 = mx.random.split(key, 3) + on_gpu = mx.default_device() == mx.gpu + dtype = mx.float16 if on_gpu else mx.float32 + + for M in (32704, 32800, 40001): + with self.subTest(M=M): + indices = mx.sort( + (mx.random.uniform(shape=(M,), key=k1) * E).astype(mx.uint32) + ) + x = (mx.random.normal((M, 1, K), key=k2) / K**0.5).astype(dtype) + w = (mx.random.normal((E, N, K), key=k3) / K**0.5).astype(dtype) + wq = mx.quantize(w, group_size=64, bits=8) + + kwargs = dict( + group_size=64, bits=8, transpose=True, rhs_indices=indices + ) + y_sorted = mx.gather_qmm(x, *wq, sorted_indices=True, **kwargs) + y_unsorted = mx.gather_qmm(x, *wq, sorted_indices=False, **kwargs) + + tol = 1e-3 if on_gpu else 1.5e-5 + self.assertLess((y_sorted - y_unsorted).abs().max(), tol) + def test_gather_qmm_grad(self): def gather_qmm_ref(x, w, s, b, lhs, rhs, trans, sort): if lhs is not None: