Add padded-buffer KV corruption reproductions

[i386]

Summary

Adds minimal reproductions for a lower-level padded-buffer corruption issue observed while investigating KV cache growth. This PR is intentionally diagnostic only; it does not attempt a fix.

What is included

• synthetic KVCache vs ConcatKeyValueCache reproduction tests in mlx-lm/src/cache.rs
• a direct large-prefix try_index_mut reproduction
• a direct concatenate(prefix, token, zero_pad) reproduction
• a control test showing concatenate(prefix, token) works

Key finding

The failure reproduces below the Qwen2/model layer. Appending or concatenating into a padded-capacity buffer can yield a live prefix view whose tail is zeroed where the appended token should be.

Current status

• diagnostic / draft only
• no proposed fix in this PR
• intended to isolate the bug before deciding where the real fix belongs

[Ogilthorp3]

Independent confirmation of this bug from a downstream project (a Qwen2.5-Coder-14B-4bit server built on mlx-rs 0.25.3), plus two findings that may help localize it.

Confirmed against a real model + reference, not just synthetic arrays. Per-position top-k logit dump at greedy temp=0, mlx-rs vs Python mlx_lm 0.31.2 on the same checkpoint:

• Incremental decode (T_q=1, persistent ConcatKeyValueCache): first argmax flip vs Python at decode position 16, cascading into degenerate output.
• Control — re-prefilling the whole sequence each step (T_q=N, fresh cache, one graph): byte-exact to Python at every position.
• The divergence is born in layer 0: token embedding is identical, but the layer-0 output L2 norm is ~15–20% smaller in the incremental path (consistent with this issue's "tail zeroed where the appended token should be" — attention reads a zero where the new token's K/V should be). MLX 0.30.6 and 0.31.2 behave identically.

Finding 1 — it is not limited to padded buffers. It also reproduces with the default lazy concatenate path when the appended token is a non-contiguous (transposed) view — i.e. the real attention layout [B, H, L, D] produced by .transpose([0,2,1,3]), not a contiguous from_slice token. This is why the concatenate(prefix, token.deep_clone()) control passes (its token is contiguous) but the pattern does not transfer to a live cache: deep_clone() on the strided view reads its base buffer contiguously and corrupts it. Neither a lazy + 0 de-alias nor an eager deep_clone/host round-trip of the strided token fixes it.

Finding 2 — architecture-dependent visibility (useful as a diagnostic). Dense, no-QK-norm models (Qwen2, Mistral) garble visibly. A QK-norm model (Qwen3.5/3.6 MoE) appears clean through the same cache — the per-head RMSNorm on Q/K renormalizes the zeroed tail and masks the corruption. So the corruption is present model-agnostically (as you note, below the model layer); QK-norm just hides it.

Also consistent with this issue: the output is evaluation-order-dependent — forcing the cache to materialize each step (eager eval, a pre-allocated slice_update buffer, or a host round-trip) makes it worse (divergence moves earlier), while the lazy concat is less-wrong; neither matches Python. That order-dependence is itself the UB signature.

Happy to share the argmax-parity harness (Rust dump + Python reference dump + diff) and the per-layer norm probe if useful. Thanks for the diagnostic repros here — they matched our independent investigation exactly.

[Ogilthorp3]

Correction to my earlier comment above. I attributed our downstream dense-model decode garble to this KV-cache / padded-buffer corruption. That was wrong, and I want to set the record straight.

The actual root cause in our project was the RoPE reshape round-trip in mlx-rs/src/nn/positional_encoding.rs — the [B,H,L,D]→[B*H,L,D]→[B,H,L,D] wrapper around fast::rope — fixed by #326. Removing that reshape (and nothing else) makes our decode byte-exact to Python mlx_lm.

Why it looked like cache corruption: RoPE mangles the non-contiguous transposed-view q/k during decode (T_q=1); the mangled K is then written into the cache, so the damage presents one layer downstream as "the cache returned a zeroed/corrupted tail." My layer-0 norm-shrink evidence was real — just localized one layer too low. The transposed-view sensitivity I described is real too, but its live manifestation was the RoPE reshape, not concatenate.

Your padded-buffer repros here may still capture a separate latent contiguity issue worth keeping; I just can't claim they were the cause of the symptom I reported. Apologies for the noise — and thanks for the diagnostics that kept me looking at the contiguity angle, which is ultimately where it was.