RotaryPositionalEncoding::forward corrupts values for length-1 sequences with non-zero offset (incremental decode)

[max-lt]

Title

RotaryPositionalEncoding::forward corrupts values for length-1 sequences with non-zero offset (incremental decode)

Body

Summary

nn::RotaryPositionalEncoding::forward produces incorrect (not just imprecise — qualitatively wrong) output when called on a 4D [B, heads, L, head_dim] tensor with L == 1 and a non-zero offset, i.e. exactly the shape used by every single-token step of incremental/streaming decoding after the first. Full-sequence calls (L > 1, offset == 0, e.g. prefill) are unaffected.

Root cause

The forward impl does:

fn forward(&mut self, input: Input) -> Result<Self::Output, Self::Error> {
    let RopeInput { x, offset } = input.into();
    let shape = x.shape();
    let x = x.reshape(&[-1, x.dim(-2), x.dim(-1)])?;   // collapse [B, heads, L, head_dim] -> 3D
    let x = crate::fast::rope(x, self.dimensions, self.traditional, self.base, self.scale, offset, None)?;
    x.reshape(shape)                                     // restore original shape
}

When x is the result of .reshape(...).transpose_axes(&[0, 2, 1, 3]) (the standard "split heads" pattern — non-contiguous in memory), this collapse-then-restore reshape silently scrambles values specifically when L == 1 and offset != 0. Calling fast::rope directly on the original 4D tensor (no reshape) produces correct output.

I don't have a minimal Metal-kernel-level explanation for why the reshape breaks specifically at L=1, offset≠0 (rather than erroring, or breaking uniformly) — happy to help narrow it down further if useful, but wanted to file this now since it's a silent-corruption bug (no panic, no shape mismatch, plausible-looking output) that's easy to hit and hard to detect for anyone doing GQA/MQA incremental decoding with a growing KV cache (i.e. any autoregressive transformer decoder).

How I found it

Porting a Voxtral speech model, I hit a bug where the model got stuck predicting one fixed token forever regardless of audio content — every intermediate tensor (encoder output, embeddings, prefill logits) matched a Python/MLX reference almost exactly, but generation diverged starting at the second decode step. Bisecting by dumping tensor stats (mean/min/max) at each stage in both implementations isolated it to immediately before/after the RoPE call in the decoder's self-attention: pre-rope Q/K matched the reference almost exactly, post-rope did not.

Repro sketch

use mlx_rs::nn::{RopeBuilder, Module};
use mlx_rs::fast;

let head_dim = 128;
let mut rope = RopeBuilder::new(head_dim).traditional(true).base(1_000_000.0).build()?;

// q: [1, n_heads, 1, head_dim] — the shape used at any decode step after the first
let offset = 39; // some non-zero position

let via_wrapper = rope.forward(mlx_rs::nn::RopeInputBuilder::new(&q).offset(offset).build()?)?;
let via_raw = fast::rope(&q, head_dim, true, 1_000_000.0, 1.0, offset, None)?;

// via_wrapper and via_raw differ significantly (not float-rounding-level) for this shape;
// they agree when L > 1 or offset == 0.

Workaround

Bypass the wrapper, call the raw op directly on the un-reshaped 4D tensor:

let q = mlx_rs::fast::rope(&q, head_dim, /* traditional */ true, rope_theta, /* scale */ 1.0, offset, None)?;

Environment

• mlx-rs git rev f4aa309c (main branch)
• macOS 26.6, Apple M-series (Metal backend)
• Reproduced consistently across multiple test audio files, not a one-off flake

Happy to provide a smaller standalone repro / write a regression test if useful.

[max-lt]

Filed a minimal, standalone repro: https://github.com/max-lt/mlx-rs-rope-bug-repro

Running it narrows the diagnosis down further than my original report above — worth correcting:

• Contiguity doesn't matter. A directly-built (contiguous) [B, heads, 1, head_dim] array diverges just as much as the non-contiguous "split heads" (reshape + transpose_axes) one.
• offset doesn't matter. L == 1, offset == 0 diverges too, not just non-zero offsets.
• traditional mode doesn't matter. Both traditional = true and the default false are affected.
• The trigger is purely L == 1. Every L > 1 case checked matches fast::rope to float precision.
• It's not actually a bug in the Rope/wrapper's reshape logic. Calling fast::rope directly — no wrapper involved at all — on a manually reshaped-to-3D [N, 1, head_dim] view of some data diverges from calling it on the native 4D [B, heads, 1, head_dim] array holding the identical underlying data. Same op, same data, only the input rank differs, and the results differ.

So this looks like fast::rope itself (the raw op/kernel) returns different results depending on the input array's rank specifically when the sequence-length axis is 1 — the RotaryPositionalEncoding::forward wrapper's collapse-to-3D just happens to be a common way to trigger it, since it's exactly what every incremental-decode step hits. The workaround is unchanged (call fast::rope directly at the tensor's original rank), but the root cause is lower-level than a wrapper reshape mistake.

Repro output and details in the linked repo's README.

[max-lt]

Root cause found — it's an upstream MLX Metal dispatch bug, already fixed in MLX ≥ v0.32.0.

Digging into the MLX source bundled through mlx-c (MLX v0.30.6), RoPE::eval_gpu in mlx/backend/metal/rope.cpp has a fast "single" path taken when T == 1 && input is row-contiguous && offset is scalar — exactly what the wrapper's collapse-to-3D produces (reshaping a non-contiguous input makes a contiguous copy, which is why contiguity of the original tensor doesn't matter in the repro).

The bug: that path's thread grid is (dims/2, N, 1), where N is the product of dims 1..ndim-2 — excluding dim 0:

• 4D [B=1, H, 1, D] → N = H → all H matrices rotated → correct. (This is why calling fast::rope at native 4D rank works.)
• 3D [B·H, 1, D] → the product is empty, N = 1 → only the first of the B·H matrices is rotated; the rest of the output buffer is left unrotated or uninitialized. Silent corruption, O(1) diffs — matching the repro output.
• Corollary the repro doesn't cover: 4D with B > 1 at L == 1 is broken too (the grid ignores B), no reshape needed at all.

L > 1 or non-contiguous inputs take the general kernel, whose grid correctly includes the batch (dim2 = B·⌈N/4⌉) — which is why every L > 1 case matches exactly, and why offset / traditional have no effect on the outcome.

Fixed upstream in ml-explore/mlx#3498 ("Fix rope single token multiple sequences", commit ml-explore/mlx@76a977c) — the fix is exactly dim1 = B * N in the single-path grid. First release containing it is MLX v0.32.0. Note the version chain: mlx-c pinned here bundles MLX v0.30.6 (buggy), and even current mlx-c main pins v0.31.2 (still buggy), so a plain mlx-c bump doesn't pick up the fix yet.

Suggested mlx-rs-side fix, independent of the pin: drop the collapse-to-3D reshape in RotaryPositionalEncoding::forward (mlx-rs/src/nn/positional_encoding.rs) and call fast::rope at the input's native rank — this is what MLX's Python nn.RoPE does (no reshape). That fixes the common B = 1 incremental-decode case even with the bundled buggy kernel, and saves a copy. B > 1 at L == 1 remains broken until the bundled MLX reaches ≥ v0.32.0. Happy to send a PR.