Fix sorted gather_mm activation row stride

[metascroy]

Fix sorted gather_mm with singleton-dimension inputs

Summary

Fix the activation row stride used by the specialized sorted RHS gather_mm implementations.

Both the Steel and NAX paths flatten the leading activation dimensions into M, so consecutive rows are K elements apart. However, they currently derive lda from the original second-to-last dimension, which can have stride 1 when it is a singleton dimension introduced by expand_dims.

Reproduction

import mlx.core as mx

mx.random.seed(42)

x = mx.random.normal((2, 64))
w = mx.random.normal((4, 64, 128))
indices = mx.array([3, 3], dtype=mx.int32)

a = mx.expand_dims(x, -2)  # [2, 1, 64]
expected = mx.matmul(a, w[indices])

sorted_output = mx.gather_mm(
    a,
    w,
    rhs_indices=indices,
    sorted_indices=True,
)
unsorted_output = mx.gather_mm(
    a,
    w,
    rhs_indices=indices,
    sorted_indices=False,
)

mx.eval(expected, sorted_output, unsorted_output)

print("sorted:", mx.max(mx.abs(sorted_output - expected)).item())
print("unsorted:", mx.max(mx.abs(unsorted_output - expected)).item())

Before this change, MLX 0.32 produces:

sorted: 33.21953201293945
unsorted: 0.0

The first row is correct, while subsequent rows are read using the wrong offset.

Root cause

The sorted RHS paths calculate:

int K = a.shape(-1);
int M = a.size() / K;
int lda = a.strides()[a.ndim() - 2];

For the [2, 1, 64] view produced by expand_dims, the singleton dimension can have stride 1. The kernel consequently reads the second row from a.flatten()[1:65] instead of a.flatten()[64:128].

mx.contiguous(a) does not reliably avoid the issue because MLX can reuse the buffer and preserve its strides.

Fix

The activation is row-contiguous before this calculation, and the leading dimensions are flattened into rows of length K. Therefore, both the Steel and NAX paths should use:

int lda = K;

With this change, the sorted result agrees with the reference within floating-point tolerance.

[AxelNoun]

Confirming root cause and isolation. We discovered this through the ExecuTorch MLX delegate:

gather_mm(sorted_indices=True) was returning max_diff ≈ 33 on inputs shaped [T, 1, K] (produced by weight.transpose(-1,-2) → expand_dims in the MoE prefill path). Your fix is the exact root cause.
We isolated it through a methodical CI sweep on macOS runners without hardware access:

•	Ruled out image-level causes: the same libmlx.a in a standalone binary was correct; the exact CMake link closure (same objects, same flags) was correct.
•	Ruled out stream/execution context: synthetic arrays on the delegate’s own stream were correct.
•	Ruled out data provenance: re-materializing the other inputs didn’t help, re-materializing a alone fixed it.
•	Pinpointed the view: a freshly allocated [T, K] re-viewed through expand_dims failed identically.

Your analysis is sound: the stride of a singleton dimension is arbitrary by design, so deriving lda from it is unsafe. Using lda = K directly is the right fix.

Impact: This unblocks MoE prefill on Apple Silicon. We tested locally on Qwen3.5 with MoE expert sorting: throughput improved from 73 to 80 tok/sec on decode.

The ExecutorTorch PR (#20685) includes a workaround (materializing a) while this lands upstream. Once it merges, the workaround can be dropped and the tests will be clean.