[BUG] fast.cross_entropy VJP returns wrong gradients for large non-contiguous inputs on Metal (0.32.3 regression) #4686

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

Describe the bug
On the GPU, when the logits are a large non-contiguous view, such as logits[:, :-1], the loss value is correct, but the gradient of mx.fast.cross_entropy is wrong. The gradient is wrong on some calls and not others: in the run below the first call returned all zeros and later calls returned a norm of 6218.6 instead of 34.352. In other runs the first call was right and later calls were wrong.

The failure needs all of the following:

• GPU only: the CPU is correct.
• Non-contiguous input only: the same values passed through mx.contiguous first are correct, and so is the unsliced array.
• Large arrays only: 37 MB of logits is fine; 149 MB and up fails. It also fails at 690 MB (10 × 119 × 151,936, a Qwen3 vocabulary).
• Independent of TF32: MLX_ENABLE_TF32=0 doesn't change it.
• Only the fused kernel: mx.logsumexp, mx.take_along_axis, mx.max, mx.softmax, and a hand-written logsumexp - take_along_axis are all correct on the same slice.

nn.losses.cross_entropy uses mx.fast.cross_entropy on the GPU, so the standard next-token loss cross_entropy(logits[:, :-1], targets[:, 1:]) gets these gradients. Fine-tuning Qwen3-0.6B with it collapsed within a few steps, while the reported loss looked normal.

It's a regression: 0.32.2 is correct, 0.32.3 is not. It may come from the Metal kernels added in #4520. In CrossEntropyVJP::eval_gpu, a non-row-contiguous input is copied with ensure_row_contiguous into an encoder temporary and the output shares that buffer (donate_in, then out.copy_shared_buffer(in)). Only the failing case goes down that path. I haven't confirmed that this is the mechanism.

To Reproduce

import mlx.core as mx

print(mx.__version__, mx.device_info()["device_name"])
mx.random.seed(0)
B, L, V = 10, 119, 32768  # 149 MB of float32 logits
x = mx.random.normal((B, L, V))
y = mx.random.randint(0, V, (B, L))

def loss_sliced(x):
    return mx.fast.cross_entropy(x[:, :-1], y[:, 1:]).sum()  # non-contiguous view

def loss_contiguous(x):
    return mx.fast.cross_entropy(mx.contiguous(x[:, :-1]), y[:, 1:]).sum()

for f in (loss_sliced, loss_contiguous):
    norms = []
    for _ in range(3):
        g = mx.grad(f)(x)
        norms.append(round(mx.linalg.norm(g.flatten()).item(), 3))
        f(x).item()  # a plain forward call between gradient calls
    print(f.__name__, "loss", round(f(x).item(), 2), "grad norms", norms)

Expected behavior

0.32.2 Apple M5 Pro
loss_sliced loss 12904.3 grad norms [34.352, 34.352, 34.352]
loss_contiguous loss 12904.3 grad norms [34.352, 34.352, 34.352]

Actual behavior
0.32.3 Apple M5 Pro
loss_sliced loss 12904.3 grad norms [0.0, 6218.577, 6218.577]
loss_contiguous loss 12904.3 grad norms [34.352, 34.352, 34.352]

Desktop:

• OS Version: macOS 26.5.1 (25F80)
• Version: mlx 0.32.3 (correct on 0.32.2); Apple M5 Pro

Additional context
Not tested on main: I couldn't build from source on this machine. mlx/backend/metal/cross_entropy.cpp on main has the same VJP code as 0.32.3. A workaround is to shift the targets instead of slicing the logits, so the kernel always receives a contiguous array.