[BUG] Metal softmax / logsumexp / rms_norm / layer_norm drop rows when grid width exceeds 2^32 #4624

☑️ I understand it is strictly prohibited to use AI to write issues.

Describe the bug

On GPU, softmax, logsumexp, mx.fast.rms_norm, and mx.fast.layer_norm fail to compute some or all rows when the row count is large. The result is silently wrong (no exception and no error), just returning zeros.

With a last axis of 128 or smaller, this starts at 2^27 rows (layer_norm: 256 or smaller). Not due to a large allocation since an fp16 tensor of shape (2^27, 1) requires only 0.5 GiB.

Of R rows, only ((R × tg) mod 2^32) / tg size rows are actually computed*. That matches a grid width wrapping at 2^32. *tg stands for thread group.

To Reproduce

Run one condition per process because, if R = 2^27 - 1 runs first in the same process, the previous buffer may stay in the output and only part of the failure shows up.

import mlx.core as mx
import numpy as np

R = 2**27
x = mx.ones((R, 1), dtype=mx.float16)
y = mx.softmax(x, axis=-1)
mx.eval(y)
print(np.array(y)[[0, 1, -1], 0])

Results vary according to conditions:

• R = 2^27 − 1: [1. 1. 1.] (correct: every row is 1)
• R = 2^27: [0. 0. 0.] (wrong: every row is 0)
• R = 2^27 + 1: [1. 0. 0.] (partially wrong: only row 0 is 1 and the rest are 0)

The same happens for logsumexp, rms_norm, and layer_norm. dtype (float16 / bfloat16 / float32) does not move the boundary.

Expected behavior

Every row should be computed correctly, or an exception should be raised.

Desktop (please complete the following information):

• OS Version: macOS 27.0 (Mac Studio, M2 Ultra, 192 GB)
• Version: 0.32.3 (PyPI).
• Also reproduces on a local build of main (65be047)

Additional context

We get to this issue through the sweep for differential testing around 2^32, extending from PyTorch to MLX (my issue for PyTorch: pytorch/pytorch#197636).

With MTL_DEBUG_LAYER=1, R = 2^27 stops on this assertion (R = 2^27 − 1 does not):

threadsPerGrid.width(4294967296) must be <= UINT32_MAX.

Same pattern elsewhere: mx.fast.cross_entropy (logits of shape (2^27, 2), all rows 0) and the layer_norm gradient (x of shape (2^27, 2), incorrect) hit the same assertion.

Related:

#3784 and PR #4437: the same class of bug in reductions. PR #4437 only fixes reduce.cpp and does not touch softmax or the norms.
#1299: fixed index types for softmax and norms on large arrays; grid width was out of scope.