Silent numerical corruption in single long forward on quantized MoE when sequence length % 32 != 0 (Qwen3-Coder-30B-A3B-8bit, M5)

Summary

A single model(ids) forward over a long prompt on a quantized MoE model
produces silently corrupted activations/KV whenever the forward's sequence
length is not a multiple of 32. Feeding the identical token ids through
the identical model in 2048-token chunks (what mlx_lm.generate_step does)
gives a different — and behaviorally correct — result. The divergence is
deterministic, content-independent (reproduces with random token ids), large
(fp16 K-cache entries off by up to ~86; final-position logits off by ~17–23;
greedy argmax flips), and big enough to break downstream behavior: a
needle-in-haystack QA that scores 100% recall with chunked prefill scores 0%
with a one-shot prefill at N≈16k.

This looks like a sibling of the 0.31.x quantized-matmul alignment bugs
(gather_qmm >32k-row MoE corruption; unaligned-N quantized_matmul with
m ≥ 10, m % 32 ≠ 0 — both fixed in 0.32.0). The MoE expert path appears to
still have an unaligned-M case: dense models (Qwen3-8B/32B, same quant
recipes, same lengths) are behaviorally clean in the same A/B, and the first
divergent tensor we observe is layer 1's K — i.e. the corruption enters
through layer 0's block output. Since dense models exercise the same
attention path at these shapes cleanly, the MoE FFN (gather_qmm /
gather-QMM-family kernels) is the prime suspect.

Environment

mlx 0.32.0.dev20260708+1ac2fbac and 0.32.1.dev20260715+cf3ce39c
(both reproduce; source builds with NAX kernels enabled,
MACOSX_DEPLOYMENT_TARGET=26.2)
Apple M5, macOS 26 (Darwin 25.5.0), Python 3.12
Model: mlx-community/Qwen3-Coder-30B-A3B-Instruct-8bit (affine 8-bit,
MoE 30B-A3B). Dense mlx-community/Qwen3-8B-4bit / Qwen3-32B-4bit do
not reproduce behaviorally at the same lengths.
mlx-lm used only to load the model and build the KV cache.

Repro

# python repro.py            -> single length
# python repro.py --bisect   -> alignment scan
import argparse
import mlx.core as mx
from mlx_lm import load
from mlx_lm.models.cache import make_prompt_cache

def cache_diff(model, ids, chunk=2048):
    def build(step):
        cache = make_prompt_cache(model)
        for s in range(0, len(ids), step):
            logits = model(mx.array([ids[s:s+step]]), cache=cache)
            mx.eval(logits, [c.state for c in cache])
        return cache, logits[0, -1]
    c1, l1 = build(len(ids))   # one-shot
    c2, l2 = build(chunk)      # chunked
    worst_k = 0.0
    for a, b in zip(c1, c2):
        if getattr(a, "keys", None) is None:
            continue
        n = min(a.offset, b.offset)
        dk = float(mx.max(mx.abs(a.state[0][..., :n, :] - b.state[0][..., :n, :])).item())
        worst_k = max(worst_k, dk)
    dl = float(mx.max(mx.abs(l1 - l2)).item())
    same = int(mx.argmax(l1).item()) == int(mx.argmax(l2).item())
    return worst_k, dl, same

ap = argparse.ArgumentParser()
ap.add_argument("--model", default="mlx-community/Qwen3-Coder-30B-A3B-Instruct-8bit")
ap.add_argument("--bisect", action="store_true")
args = ap.parse_args()

model, tok = load(args.model)
mx.random.seed(0)
ids = mx.random.randint(0, 50000, (17000,)).tolist()
lengths = [8192, 12288, 16000, 16032, 16064, 16065, 16066, 16067, 16068,
           16128, 16384, 16900] if args.bisect else [16068]
print(f"mlx {mx.__version__}")
print(f"{'N':>7} {'worstK':>8} {'logitdiff':>9}  argmax")
for L in lengths:
    wk, dl, ok = cache_diff(model, ids[:L], 2048)
    print(f"{L:>7} {wk:>8.3f} {dl:>9.3f}  {'OK' if ok else 'DIVERGED'}", flush=True)

Observed (random ids, seed 0, deterministic across runs)

mlx 0.32.0.dev20260708+1ac2fbac        (0.32.1.dev20260715: same pattern)
      N   worstK logitdiff  argmax
   8192   15.125     0.562  OK          # 32-aligned
  12288   15.125     1.172  OK          # 32-aligned
  16000   15.125     0.182  OK          # 32-aligned
  16032   15.125     0.344  OK          # 32-aligned
  16064   26.000     0.109  OK          # 32-aligned
  16065   32.000     7.562  DIVERGED    # N % 32 == 1
  16066   32.500     8.312  DIVERGED    # N % 32 == 2
  16067   32.000     8.375  DIVERGED    # N % 32 == 3
  16068   32.000    17.500  DIVERGED    # N % 32 == 4
  16128   26.000     1.148  OK          # 32-aligned
  16384   26.000     0.188  OK          # 32-aligned
  16900   32.000     9.438  DIVERGED    # N % 32 == 4

Every diverging length is unaligned; every 32-aligned length matches argmax
with small logit deltas. Unaligned lengths below 16k already show large
numeric error without flipping the argmax at the final position (e.g. on
0.32.1.dev20260715: N=8193 → worstK 131, logitdiff 17.0, argmax OK; N=12289 →
worstK 54, logitdiff 4.3, argmax OK), so short-prompt spot checks miss it.

Notes:

The constant worstK ≈ 15–26 at aligned lengths is presumably benign
kernel-tiling numerics amplified by outlier channels (argmax stable, logit
deltas ≪ 1 in most rows); the unaligned case is qualitatively different.
Layer-0 K is identical between the two paths (it depends only on the
embeddings); divergence first appears in layer-1 K and compounds through
all 48 layers.
Real-workload impact: with a real 16,068-token needle-in-haystack prompt,
the one-shot path's K diverges by up to 86, final logits by 23.4, and
greedy QA recall flips 1.0 → 0.0 (the model emits repetitive
off-task text). Chunked prefill of the same ids answers correctly. This
cost us a multi-day false negative in a KV-eviction study before we traced
it to the prefill path.

Workaround

Chunk all prefills (any chunk size that keeps each forward 32-aligned or
short works; we use 2048).

We can test candidate fixes, run finer-grained bisects, or collect kernel-level
traces on M5 hardware on request.

Additional scoping data (M5, mlx 0.32.0.dev20260708, same standalone repro as
the issue body — random ids, N=16068, one-shot vs 2048-chunked prefill).
Models with rotating-window caches were compared on the last 64 positions'
logits instead of raw K tensors.

model	expert quant mode	max logit diff	argmax
Qwen3-Coder-30B-A3B-Instruct-8bit	affine 8-bit	7.4	diverged
Qwen3-Coder-Next-MLX-4bit	affine 4-bit	10.8	diverged (37/64 positions)
Qwen3-Next-80B-A3B-Instruct-4bit	affine 4-bit	3.3	diverged
Laguna-XS-2.1-8bit	affine 8-bit	17.8	diverged (60/64 positions)
DeepSeek-V2-Lite-Chat-4bit (MLA)	affine 4-bit	6.7	survived
Nemotron-Super-120B-5bit (hybrid)	affine 5-bit	7.1	survived
gpt-oss-20b	mxfp4	0.34	ok (2/64 near-tie flips)
gpt-oss-88B variant, q6 affine EXCEPT experts	mxfp4 experts	0.38	ok (1/64)
Qwen3-Coder-Next abliterated mxfp4-gs32	mxfp4	1.28	ok (6/64)

Observations:

Every affine-expert model corrupts (4/5/8-bit; standard MoE, MLA,
GDN-hybrid, and mamba-hybrid attention all affected — the common factor is
the expert weights' quant mode, not the attention architecture).
Every mxfp4-expert model is at FP-noise level at the same lengths.
The Qwen3-Coder-Next pair isolates the variable: identical architecture,
affine-4bit experts diverge (10.8 / 37 flips) while mxfp4 experts are near
clean (1.28 / 6 flips).
All of these route expert matmuls through the same SwitchGLU →
mx.gather_qmm path in mlx-lm, so the data points at the affine-mode
gather_qmm kernel for long unaligned row counts specifically.

Happy to run more cells (other bit-widths / group sizes) if useful.

Update with a sharper repro and a scope correction: this is alive in 0.32.0 release (PyPI), not just 0.32.1.dev, and the trigger is cleaner than the per-model token counts suggested.

Going through the sorted-indices gather_qmm path directly (fp32-dequantized dense reference as arbiter, M5, mlx 0.32.0):

Trigger: flattened gathered row count n with n > 32768 and n % 64 != 0. In an MoE forward n = tokens × top_k, which unifies the earlier reports: at top_k=2 it manifests as tokens%32, at top_k=4 as tokens%16 — the invariant is row alignment, not sequence length.
Affine 4-bit and 8-bit: corrupt on both 0.31.2 and 0.32.0.
mxfp4: clean at all tested shapes on 0.32.0 — presumably why the fix that landed there looked complete.
Boundary is exact in our tests: n = 32768 and n = 40000 clean; n = 32802 and n = 40002 corrupt.

Minimal repro, pure mlx (no mlx-lm), mirrors the call shape of an MoE expert matmul:

import mlx.core as mx

print("mlx", mx.__version__)


def trial(n, bits, group_size, mode, E=8, N=64, K=64):
    mx.random.seed(0)
    w = mx.random.normal((E, N, K)).astype(mx.float16)
    q = tuple(mx.quantize(w, group_size=group_size, bits=bits, mode=mode))
    x = (mx.random.normal((n, 1, K)) * 0.5).astype(mx.float16)
    idx = mx.sort(mx.random.randint(0, E, (n,)).astype(mx.uint32))
    out = mx.gather_qmm(
        x, *q, rhs_indices=idx, transpose=True,
        group_size=group_size, bits=bits, mode=mode, sorted_indices=True,
    )
    wd = mx.dequantize(*q, group_size=group_size, bits=bits, mode=mode)
    ref = x.astype(mx.float32) @ wd[idx].swapaxes(-1, -2).astype(mx.float32)
    err = mx.abs(out.astype(mx.float32) - ref).max(axis=(1, 2))
    print(
        f"mode={mode:6s} bits={bits} n={n} (n%64={n % 64:2d}): "
        f"max|err|={err.max().item():.4f}  bad rows={(err > 0.05).sum().item()}/{n}"
    )


for n in (32768, 32802, 40002):
    trial(n, bits=4, group_size=32, mode="affine")
trial(32802, bits=8, group_size=64, mode="affine")
trial(32802, bits=4, group_size=32, mode="mxfp4")

Output on mlx 0.32.0 (M5, macOS 25.5):

mode=affine bits=4 n=32768 (n%64= 0): max|err|=0.0078  bad rows=0/32768
mode=affine bits=4 n=32802 (n%64=34): max|err|=16.9303  bad rows=64/32802
mode=affine bits=4 n=40002 (n%64= 2): max|err|=21.7688  bad rows=7264/40002
mode=affine bits=8 n=32802 (n%64=34): max|err|=16.5858  bad rows=64/32802
mode=mxfp4  bits=4 n=32802 (n%64=34): max|err|=0.0077  bad rows=0/32802

Workaround we're deploying in mlx-lm meanwhile: pad the sorted rows to a multiple of 64 (provably output-neutral — the unsort indexes only original rows). Happy to PR the repro as a test case here if useful.

@pierre427

I’d like to investigate and implement the kernel-side fix for the affine gather_qmm tail corruption. I’ll keep the change scoped to the minimal n > 32768 && n % 64 != 0 reproducer, with aligned controls and affine 4/8-bit coverage.

@pierre427, are you already working on the implementation, or would you be willing to validate a candidate patch on your M5? I’ll avoid a Python-level padding workaround unless maintainers prefer that direction.

Thanks. I’m not currently implementing the MLX kernel-side fix; the work I opened is the downstream mlx-lm pad workaround in ml-explore/mlx-lm#1585.

Happy to validate a candidate kernel patch on M5. I already have fp32-reference arbitration for the n > 32768 && n % 64 != 0 sorted gather_qmm case, with affine 4/8-bit coverage and aligned controls, so I can run a focused before/after check quickly.

@wgu9

Root cause found — this is an int16 overflow in affine_gather_qmm_rhs_nax, the one
sibling that #3631 missed.

mlx/backend/metal/kernels/quantized_nax.h#L1532-L1535:

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)))));

M - (y_row + tm) — the rows remaining below this tile, where M is the flattened
gathered row count (the n in my scoping comments) — is cast to short before the
min. When align_M == false (n % 64 != 0, BM = 64) and a tile sits ≥ 32768 rows
from the end, the cast wraps negative; a negative sgp_sm zeroes the A-tile and the store
path stores nothing. Those output rows are never written and expose stale allocator
memory — which is why the corruption reads as garbage rather than a consistent numeric
error, and why it varies run to run.

#3631 fixed this exact pattern in three siblings (fp_qmm_t_impl, the fp gather-rhs
kernel, and affine qmm_t_nax_tgp_impl — the "before" code in its fp_quantized_nax.h
second hunk is character-identical to the lines above) but missed this fourth site. That
is also why mxfp4 tests clean while affine corrupts. A twin lurks in sgp_sn two lines
down (needs N > 32768 with N % 64 != 0 — harder to hit, same one-line fix).

The mechanism accounts for every measured property:

n % 64 == 0 escapes — align_M short-circuits the cast entirely.
Corruption arrives in 64-row blocks, leading rows first — exactly the BM=64 tiles whose
distance to the end exceeds 32767. Verified: n = 32769 → rows 0–31 bad; n = 32833 →
rows 0–95 bad; n = 32896 → clean.
Apparent bit-width/version dependence in my earlier scoping was an artifact: unwritten
rows hold whatever the recycled MTLBuffer last contained, which is sometimes
coincidentally plausible. A regression test should poison the output buffer (or compare
two runs) rather than trust one lucky read.

One important note for anyone trying to reproduce upstream: this kernel only dispatches
when NAX is available (M5-generation GPU on macOS 26.2+) and the sorted fast path is
taken; on M1–M4 the steel fallback runs and is correct. The bug cannot be reproduced on
M1–M4 hardware.

@wgu9 — so the candidate patch should be exactly #3631's transform applied to these two
lines (min(int(SM), max(0, ...)) / min(int(SN), max(0, ...))). The M5 validation
offer stands: I have fp32-reference arbitration for the sorted gather_qmm case ready
with affine 4/8-bit coverage and aligned controls, plus the three n values above as the
row-block prediction check.

Confirmed on M5 Max (mlx 0.32.0, macOS 27.0) — and the boundary condition in the report holds exactly on this chip generation too.

Self-contained synthetic repro (16 experts, K=N=64, gather_qmm with sorted_indices=True, fp32 dequantized reference, per-row batched shapes):

mode=affine bits=8 group_size=32          mode=mxfp4 bits=4 group_size=32
  len      max_abs                          len      max_abs
32768     0.015549   (n%64==0, clean)     32768     0.015450
32769    27.001427   CORRUPT              32769     0.015610
32800    40.453659   CORRUPT              32800     0.015579
32802    33.777496   CORRUPT              32802     0.015579
40000     0.015610   (n%64==0, clean)     40000     0.015476

Three observations from this run:

The trigger is precisely n > 32768 AND n % 64 != 0 — 32768 and 40000 (both %64==0) sit at the ~0.0155 quantization-noise floor; the three non-aligned lengths jump 3 orders of magnitude. Not gradual degradation — a hard tile/boundary effect.
mxfp4 is unaffected here — all lengths at noise floor. Whatever the defect is, it lives in the affine gather_qmm path's handling of the ragged tail past the 32768 boundary, not in the shared gather machinery.
Reproduces with tiny synthetic weights (<2GB RAM, ~a minute) — no model download needed; happy to share the script.

Our stake, for context: we validate large quantized MoE models through this exact path, so we checked exposure — for anyone doing the same, padding row counts to a multiple of 64 is a working mitigation until the kernel fix.

fix(patches): reroute sorted gather_qmm around defective M5 NAX kerne…

@jundot

Heads-up for anyone working this kernel: a second, independent defect in affine_gather_qmm_rhs_nax is now filed as #3887 — an !align_K tail bound that corrupts the sorted-rhs path for K % 64 != 0 (root-caused by @jundot in omlx#2267; verified on M5 against 0.32.0). Two things that differ from this issue: the trigger axis is the reduction dim rather than row count, and the mxfp4 variant is affected too.

@wgu9 — if your fix for the int16 row-offset overflow here is still in flight, one PR could cover both defects plus the sgp_sn twin noted above; the M5 validation-sweep offer stands for the combined patch.

@pierre427

@pierre427 Good cross-link, and agreed they are distinct defects: this one keys on row count (seqlen % 32 != 0), #3887 on the reduction dim (K % 64 != 0), and the mxfp4 variant being affected there widens the blast radius beyond what this issue covers.

No objection to one PR covering both plus the sgp_sn twin if @wgu9's int16 row-offset fix is still in flight — they are the same kernel and splitting them would just mean two validation sweeps. The repro in this issue (Qwen3-Coder-30B-A3B-8bit, M5, single long forward with seqlen % 32 != 0) is available for validating a combined patch; it fails loudly enough to be a decent regression case since the corruption is silent rather than a crash.

Opened #3922 with the kernel-side fix and M5 regression coverage.

The patch clamps the remaining row/column counts in int before narrowing, matching the earlier #3631 sibling fix. Native M5 Max tests cover affine 4/8-bit at the exact 32768/32769 seam and several ragged tails, with two allocator-poison values and a dense fp32 reference. Independent post-build validation found max error below 4.7e-4 across the declared affine/MXFP matrix; the full 34-test quantized suite and forced-gen16 controls also pass.

Security scope: this can leave tensor rows unwritten and expose contents reused from MLX's same-process Metal allocator pool. Current evidence shows numerical corruption within the same process. We found no cross-process disclosure, arbitrary code execution, sandbox escape, or other trust-boundary crossing, so this is a correctness bug rather than a cybersecurity vulnerability.

fix(mlx): reroute sorted gather_qmm around the defective NAX rhs kernel

fix(mlx): reroute sorted gather_qmm around the defective NAX rhs kern…

fix: segment oversized M5 sorted gather calls (#3432)

perf(moe): single-dispatch native NAX gather for oversized sorted MoE…

Navigation Menu

Repository navigation

Silent numerical corruption in single long forward on quantized MoE when sequence length % 32 != 0 (Qwen3-Coder-30B-A3B-8bit, M5) #3856

Summary

Environment

Repro

Observed (random ids, seed 0, deterministic across runs)

Workaround

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Silent numerical corruption in single long forward on quantized MoE when sequence length % 32 != 0 (Qwen3-Coder-30B-A3B-8bit, M5) #3856

Description

Summary

Environment

Repro

Observed (random ids, seed 0, deterministic across runs)

Workaround

Activity

pierre427 commented on Jul 17, 2026

pierre427 commented on Jul 18, 2026

wgu9 commented on Jul 20, 2026

pierre427 commented on Jul 20, 2026

pierre427 commented on Jul 20, 2026

PhilipJohnBasile commented on Jul 21, 2026

pierre427 commented on Jul 21, 2026

PhilipJohnBasile commented on Jul 22, 2026

PhilipJohnBasile commented on Jul 26, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions