feat: add contiguity check, contiguous() op, and core fixes - #326
Conversation
- 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>
|
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
Mechanism (why it's decode-specific). The q/k that attention feeds RoPE are non-contiguous transposed views ( 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.) |
Summary
Array::is_contiguous()method to check if an array has row-major contiguous memory layoutcontiguous()op (wrappingmlx_contiguous) so users can make strided views safe foras_slice()AsSliceError::NotContiguouserror variant for explicit contiguity checkingeval()intry_item()— was calling eval twice per.item()call (~15% overhead)StreamOrDevice::default()to respect task-local streams instead of always creating new ones[B,H,L,D]→[B*H,L,D]reshape that broke multi-head attention during decode phase (aligns Rust behavior with Pythonnn.RoPE)Breaking changes
None —
try_as_slice()behavior is unchanged. The contiguity check is opt-in viais_contiguous().Test plan
test_contiguous()andtest_is_contiguous_and_contiguous()inops/shapes🤖 Generated with Claude Code