fix(nn): apply RoPE on the input shape, not a reshaped 3D tensor - #357
Open
sergey-scherbina wants to merge 1 commit into
Open
fix(nn): apply RoPE on the input shape, not a reshaped 3D tensor#357sergey-scherbina wants to merge 1 commit into
sergey-scherbina wants to merge 1 commit into
Conversation
`RotaryPositionalEncoding::forward` reshaped [B, n_heads, L, head_dim] to [-1, L, head_dim] before `mx.fast.rope`. For single-position input (L == 1, i.e. every decode step) the resulting [B*n_heads, 1, head_dim] shape triggers a fast-rope bug that rotates only the first batch row, leaving every head past the first un-rotated -> corrupt decode queries -> garbage generation. Prefill (L > 1) is unaffected, so this only shows up during incremental decoding. Python mlx_lm applies `mx.fast.rope` to the 4D tensor directly and is correct. Fix: drop the reshape and apply RoPE on the input shape. Adds a regression test that feeds identical heads at a non-zero offset and asserts every head is rotated alike (a single-head reference, broadcast).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
RotaryPositionalEncoding::forwardreshapes[B, n_heads, L, head_dim]to[-1, L, head_dim]before callingmx.fast.rope, then reshapes back. Forsingle-position input (
L == 1, i.e. every autoregressive decode step) thecollapsed
[B*n_heads, 1, head_dim]shape hits afast::ropebug that rotatesonly the first row along the leading axis — every head past the first is left
un-rotated. The result is corrupt decode queries and garbage generation.
Prefill (
L > 1) is unaffected, so this only surfaces during incrementaldecoding, which makes it easy to miss.
Python
mlx_lmappliesmx.fast.ropeto the 4D tensor directly and is correct.Fix
Drop the reshape and apply RoPE on the input shape directly.
Test
Adds
test_rope_rotates_all_heads_when_single_position: feeds four identicalheads at a non-zero offset and asserts every head comes out identical to a
single-head reference (RoPE depends only on position, so identical heads must
rotate identically) and that the rotation is non-trivial. Fails before this
change (heads 1–3 stay un-rotated), passes after.