Please bump the vendored mlx-core so the batched single-token RoPE fix (ml-explore/mlx#3498) reaches mlx-swift

[magicnight]

Hi team, first — thanks a lot for mlx-swift. I've been building an in-process MLX inference app on top of it and it's been really solid.

I ran into a bug while working on continuous batching. When I decode a single token for a batch ([B, H, 1, D], scalar offset, B > 1), MLXNN.RoPE / MLXFast.RoPE only computes row 0 correctly — the other rows come out as garbage / NaN, looks like uninitialized memory. At B = 1 everything is fine, so I guess that's why it stayed hidden for so long.

After digging a bit I realized it's not a new problem — it's the same as ml-explore/mlx#3496 (Swift side) and ml-explore/mlx#3494 (Python side), and it was already fixed in mlx-core by ml-explore/mlx#3498 (76a977ca — the single-token Metal path was launching an N grid instead of B*N). That fix went into mlx-core 0.32.0.

The thing is, it hasn't reached mlx-swift yet: 0.31.6 (the latest) still pins Source/Cmlx/mlx at v0.31.1 (ce45c52), which is before the fix. So anyone doing batched decode on mlx-swift still hits it today.

Small repro on 0.31.6 (Apple Silicon, Metal):

let rope = RoPE(dimensions: 8, base: 10000)
let r = MLXRandom.normal([1, 1, 8]).asType(.float16)
let x = concatenated([r.expandedDimensions(axis: 0),
                      r.expandedDimensions(axis: 0)], axis: 0)  // [B=2,1,1,8], two identical rows
let y = rope(x, offset: 1)
// abs(y - y[0]).max() -> NaN / nonzero, expected 0.0
// Works with an array offset MLXArray([1, 1]), or on mlx-core >= 0.32.0

Would it be possible to bump Source/Cmlx/mlx to a commit that includes ml-explore/mlx#3498 (mlx-core >= 0.32.0), or cherry-pick those few lines onto the current pin? If it helps I'm happy to open the PR myself.

Thanks again for all the work on this.

[davidkoski]

This next bump may require a bit of extra work -- there are some changes to the threading model with regard to Streams. I think that the updates on the mlx side will be easy to adopt, but not 100% sure.

Anyway, yes, as soon as there is a new mlx-c I will work on picking it up here. I will talk to the other mlx folks and see when we can do that.

[magicnight]

Quick status note in case it's useful: ml-explore/mlx#3498 shipped in mlx v0.32.0 (Jul 7). mlx-c hasn't picked that up yet though — main is still on mlx 0.31.2 and the latest tag (v0.6.0) is on 0.31.1 — so the bump you're waiting for looks like it's specifically mlx-c -> 0.32.0. No rush on my side; the app runs a patched fork in the meantime.

[xocialize]

Another data point for whenever the mlx-c bump becomes possible — no urgency implied, just noting that the same bump carries a second fix, in case it's useful for weighing priority.

mlx v0.32.0 also contains ml-explore/mlx#3810, a fix for the NAX split-K GEMM kernel (steel_gemm_splitk_nax). We reported that one as ml-explore/mlx#3797: the host dispatch passed the output dtype where the kernel template needed the input dtype, so the JIT-compiled kernel read bf16 input memory as float and loaded out of bounds. @davidkoski reproduced it here and caught it via Metal shader validation.

The reason it may be worth weighting: it only affects the JIT path, and mlx-swift builds with MLX_METAL_JIT=ON (tools/update-mlx.sh), so it reaches mlx-swift users but not the PyPI wheel, which is AOT-compiled. On the same machine and commit, MLX_METAL_JIT=OFF produced cos 0.9999986 against an fp32 reference while MLX_METAL_JIT=ON produced NaN/1e35 garbage across every size we tried.

The part that makes it awkward downstream is that it fails silently: no crash, no NaN guard, just wrong values from ordinary half-precision matmuls once they cross the split-K dispatch window (large K, batch 1). We hit it as degraded model output that looked like a parity bug in our own port before we traced it to the kernel. We're carrying a row-chunking workaround, so nothing is blocked on our end — I just didn't want the second fix to be invisible when the bindings work gets scheduled.

Happy to help verify once a bump lands; we have a machine that reproduced it deterministically.

[magicnight]

Second data point on the same bump, from a different machine: I reproduced the NAX split-K issue @xocialize described above, and confirmed the fix.

On an M5 Max running macOS 26.6, against the mlx vendored by mlx-swift 0.31.6, a bf16 matmul at M = N = 2048, K = 12288 — just inside the split-K dispatch window — scored cosine 0.0 against a float32 reference of the same operands and contained non-finite values. The same matmul at K = 8192, one step outside the window, was correct. So the divergence tracks the dispatch, not bf16 precision.

Cherry-picking ml-explore/mlx#3810's four-line diff onto that vendored mlx makes both cases pass, with no other change. My full test suite (real-Metal, 819 tests including per-architecture parity fixtures) is unaffected either way, which matches the narrow trigger conditions — it needs NAX, a non-float32 input, batch 1, M*N >= 2048^2, K >= 10240 and K >= 3*max(M, N).

Nothing blocked on my end; I'm carrying the cherry-pick alongside the mlx#3498 one until the bindings work lands. Just adding an independent confirmation in case it helps weigh the bump.