MLXNN.RoPE produces row-asymmetric output for byte-identical row inputs at [B>1, H, 1, D]

[ivan-digital]

Summary

MLXNN.RoPE (which dispatches into MLXFast.RoPE) on a [B, H, 1, D] input where row 0 == row 1 produces row-asymmetric output. The same op on [B, H, T, D] with T > 1 is row-symmetric. The same op with B = 1 is deterministic and correct.

This blocks any straightforward [B,1,H] batched autoregressive decode in mlx-swift — every decoder layer's attention applies RoPE per step at T=1, and the resulting per-row drift compounds through the layers fast enough that argmax tokens diverge within a few decode steps.

Minimal reproducer (no model, no quantization)

import MLX
import MLXNN
import MLXRandom

MLXRandom.seed(42)

let rope = MLXNN.RoPE(dimensions: 128, traditional: false, base: 1_000_000)

// Row-symmetric input: row 0 byte-equal to row 1.
let baseRow = MLXRandom.normal([1, 8, 1, 128])  // [1, H, T=1, D]
let input = concatenated([baseRow, baseRow], axis: 0)  // [2, H, T=1, D]

// Sanity: input rows are identical.
let i0 = input[0..<1, 0..., 0..., 0...].asArray(Float.self)
let i1 = input[1..<2, 0..., 0..., 0...].asArray(Float.self)
let inputDiff = zip(i0, i1).map { abs($0 - $1) }.max() ?? 0
assert(inputDiff == 0)

let output = rope(input, offset: 0)

let r0 = output[0..<1, 0..., 0..., 0...].asArray(Float.self)
let r1 = output[1..<2, 0..., 0..., 0...].asArray(Float.self)
let rowDiff = zip(r0, r1).map { abs($0 - $1) }.max() ?? 0

print("RoPE B=2 T=1 offset=0 row-diff: \(rowDiff)")
// Expected: 0.0
// Actual:   5.379... (varies by seed but always non-zero)

Diagnostic matrix

Shape	Op	Offset	row 0 vs row 1 maxDiff
[2, 1, 1024]	quantizedMatmul (4-bit, group 64)	—	0.0 ✓
[2, 148, 1024]	quantizedMatmul (4-bit, group 64)	—	0.0 ✓
[2, 8, 1, 128]	qProj → reshape → RMSNorm → transpose	—	0.0 ✓
[2, 8, 1, 128]	RoPE	0	5.38 ✗
[2, 8, 1, 128]	RoPE	148	7.52 ✗
[2, 8, 2, 128]	RoPE	148	0.0 ✓ (T>1 control)
[1, 8, 1, 128]	RoPE (called twice)	148	0.0 ✓ (B=1 deterministic)

The trigger is B > 1 AND T = 1 — independent of offset, independent of preceding ops (the bug shows up on a freshly constructed tensor that never went through quantizedMatmul or RMSNorm).

Why it matters

This silently corrupts any batched autoregressive decode loop in mlx-swift, because:

1. Decoder forward pass at decode step has shape [B, H, T=1, D] (one new query per row).
2. RoPE is applied to every layer's Q and K at offset = current_cache_length.
3. The per-row drift introduced by RoPE compounds through ~28 layers per step.
4. Argmax over the LM head picks different tokens for different rows within a few steps even when rows started identical.

Surface symptom: feeding two identical chunks through a batched decode loop produces correct output for row 0 and a truncated / different output for row 1.

We hit this in soniqo/speech-swift PR ml-explore/mlx-swift-examples#453's experimental [B,1,H] batched decoder for Qwen3-ASR. With B=2 and two identical 10s audio chunks, row 0 transcribes the full sentence and row 1 stops after the first few tokens. Cache layout (per-step concatenated vs. mlx-lm-style pre-allocated BatchKVCache) does not change the symptom — both fail in the same way because the upstream cause is RoPE itself.

What we ruled out

• ✅ Determinism: same op, same input, run twice → byte-identical (so it's not a kernel-level race).
• ✅ Cache layout: mlx-lm-style BatchKVCache (pre-allocated, indexed-slice writes) does not fix it.
• ✅ Quantization: triggers without any quantized op — pure RoPE on a freshly constructed tensor reproduces.
• ✅ Position offset: offset = 0 reproduces (~5.4 row diff), so it's not specific to non-zero positions.
• ✅ T > 1 is fine: same op at T = 2 is row-symmetric.
• ✅ B = 1 is fine: same op at B = 1 is deterministic and matches per-row B=2 row 0.

Probable cause

A T = 1 fast path in MLXFast.RoPE's Metal kernel that doesn't broadcast / iterate the batch dim correctly. Apple's mlx core changelog mentioned "No copy batch rope" recently — possibly related.

Workaround we ship

Per-row B=1 decoder forwards (no kernel-level B amortisation). Acceptable for our 11% throughput win on token-sync batching, but blocks the larger speedup that true [B,1,H] would unlock.

Environment

• macOS 26.4 (Tahoe), arm64, Apple M5 Pro
• Xcode 17 stable, Swift 6.2
• mlx-swift 0.31.3 (revision 61b9e01)
• Reproducible against MLXNN.RoPE and MLXFast.RoPE (RoPE wraps Fast)

Cross-references

• mlx-lm BatchKVCache — Python pattern for batched generation. Likely also affected upstream once mlx-swift-lm wires up batched decode with quantized models.
• "The Hidden Problem With MLX" — documented batch invariance failures in floating-point Metal kernels; this finding extends the picture: also RoPE at T=1, also affects models with quantized weights when surrounded by bf16 RoPE.

[davidkoski]

OK, I need to see if mlx python reproduces it -- if so we should move the issue to mlx as it is in the core.

[davidkoski]

import mlx.core as mx
import mlx.nn as nn

mx.random.seed(42)

rope = nn.RoPE(dims=128, traditional=False, base=1_000_000)

# Row-symmetric input: row 0 byte-equal to row 1.
base_row = mx.random.normal([1, 8, 1, 128])  # [1, H, T=1, D]
x = mx.concatenate([base_row, base_row], axis=0)  # [2, H, T=1, D]

# Sanity: input rows are identical.
input_diff = mx.abs(x[0:1] - x[1:2]).max().item()
assert input_diff == 0, f"Input rows differ: {input_diff}"

output = rope(x, offset=0)

row_diff = mx.abs(output[0:1] - output[1:2]).max().item()

print(f"RoPE B=2 T=1 offset=0 row-diff: {row_diff}")
# Expected: 0.0
# Actual:   non-zero (same bug as Swift)

repros with 0.31.2

[ivan-digital]

Thanks for confirming. Should we file this at ml-explore/mlx ourselves with the Python repro you posted, or are you handling the move?

[davidkoski]

I asked somebody to move it.

[ivan-digital]

Hi @davidkoski — bumping this since the move went through (thanks). No urgency on a fix, but flagging that downstream work is still sitting behind the gate.

Context for prioritization: we have a batched-decode optimization in our ASR pipeline (soniqo/speech-swift) gated behind an env var (QWEN3_ASR_EXPERIMENTAL_BATCH_DECODE) waiting on this. The fall-back path — per-row serial decoder forwards with batched host token sync — ships fine and gives a meaningful share of the win, so this isn't blocking us shipping. But removing the gate would unlock the rest of the speedup once RoPE goes row-symmetric at [B>1, H, T=1, D].

If a fix is in flight I'd be happy to validate against the regression probe we kept in soniqo/speech-swift#236 — it's XCTExpectFailure-pinned, so the day the upstream goes symmetric the test flips red in our CI and we get the signal to drop the gate. If it would help, I can also take a stab at a patch — the Python repro you posted is precise enough that the surface looks tractable.

Thanks again for confirming the bug.

[zcbenz]

This seems to have been fixed in the main branch, I have no idea which RP fixed it though.

[zcbenz]

Actually this had been fixed by #3498.