Skip to content

feat: add contiguity check, contiguous() op, and core fixes - #326

Open
ymote wants to merge 1 commit into
oxiglade:mainfrom
ymote:fix/core-improvements
Open

feat: add contiguity check, contiguous() op, and core fixes#326
ymote wants to merge 1 commit into
oxiglade:mainfrom
ymote:fix/core-improvements

Conversation

@ymote

@ymote ymote commented Mar 15, 2026

Copy link
Copy Markdown

Summary

  • Add Array::is_contiguous() method to check if an array has row-major contiguous memory layout
  • Add contiguous() op (wrapping mlx_contiguous) so users can make strided views safe for as_slice()
  • Add AsSliceError::NotContiguous error variant for explicit contiguity checking
  • Fix duplicate eval() in try_item() — was calling eval twice per .item() call (~15% overhead)
  • Fix StreamOrDevice::default() to respect task-local streams instead of always creating new ones
  • Fix RoPE reshape bug: remove unnecessary [B,H,L,D][B*H,L,D] reshape that broke multi-head attention during decode phase (aligns Rust behavior with Python nn.RoPE)

Breaking changes

None — try_as_slice() behavior is unchanged. The contiguity check is opt-in via is_contiguous().

Test plan

  • All 582 existing tests pass
  • New tests: test_contiguous() and test_is_contiguous_and_contiguous() in ops/shapes

🤖 Generated with Claude Code

- Add `Array::is_contiguous()` to check row-major memory layout
- Add `contiguous()` op wrapping `mlx_contiguous` for making strided
  views safe for `as_slice()`
- Add `AsSliceError::NotContiguous` variant for explicit checking
- Fix duplicate `eval()` call in `try_item()` (was calling twice)
- Fix `StreamOrDevice::default()` to use task-local stream instead
  of ignoring task-local streams via `Stream::new()`
- Fix RoPE: remove unnecessary reshape from [B,H,L,D] to [B*H,L,D]
  that broke multi-head attention (aligns with Python nn.RoPE)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@Ogilthorp3

Copy link
Copy Markdown

Independent downstream confirmation that this RoPE fix resolves a real, severe decode bug — with byte-exact evidence.

Context: a Qwen2.5-Coder-14B-4bit server built on mlx-rs (0.25.3). Dense models (Qwen2, Mistral) produced prefix-correct-then-degenerate output on any multi-token completion during incremental decode; single-token answers were always fine.

This PR's RoPE change is the fix. Removing the [B,H,L,D][B*H,L,D] reshape round-trip in Rope::forward — and nothing else — makes our decode byte-exact to Python mlx_lm at greedy temp=0, per-position argmax:

Case Before After this PR
fibonacci, 40 tok first argmax flip at decode pos 16 → repetition collapse 40/40 match
palindrome, ~19 tok garble 19/19 match
with fused QKV + MLP projections garble 40/40 match

Mechanism (why it's decode-specific). The q/k that attention feeds RoPE are non-contiguous transposed views (reshape(B,L,H,D).transpose(0,2,1,3) → logical [B,H,L,D], physical [B,L,H,D]). The reshape([-1, L, D]) round-trip forces a contiguous copy of that strided view, which is mishandled at L=1 and corrupts the values — so prefill (T_q=N) survives but decode (T_q=1) garbles, exactly as the PR title says. fast::rope handles strided inputs correctly on its own, so deleting the reshape is the right fix.

Diagnostic note for others hitting this: dense / no-QK-norm models (Qwen2, Mistral) garble visibly; a QK-norm model (Qwen3.5/3.6 MoE) looks clean through the same path because the per-head Q/K RMSNorm renormalizes the corrupted tail and masks it. So "my MoE model works fine" does not mean you're unaffected.

This is a correctness fix for every dense model doing incremental decode — would love to see it merged. Happy to contribute the argmax-parity harness (Rust dump + Python reference + diff) as a regression test if useful.

(Note: I previously mis-attributed this same downstream symptom to the KV-cache repros in #338; I've corrected that on-thread. The cause is here, in RoPE.)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants