fix(deepseek_v4): with-cache decode parity + regression tests (@anerjy)

[jackneil]

Summary

Fixes the @anerjy decode-cache bug on ml-explore#1189 and ships the deterministic regression tests that prove it. Original bug: with-cache autoregressive decode produced different logits at S=1 than re-prefilling the same token sequence, so greedy generation visibly corrupted after the first token ('{"verdict":"yes"}' → '{"ver":"yes"}').

Investigation went one layer deeper than the original bisection (which had ruled out caches, rope, compressor, sinks, and mHC residual). Per-layer instrumentation exposed a chain of four mutually-reinforcing defects in the compressed-attention + sliding-window paths. Each fix unmasked the next.

Root causes (all in mlx_lm/models/deepseek_v4.py)

1. Bool-dtype prepend mask was zero-False in V4Attention.__call__ — for a bool mask, mx.zeros(... dtype=bool) is False = "block", silently masking out the entire compressed pool from S>1 (prefill / re-prefill) attention while S=1 decode passed mask=None and attended it. Asymmetric behaviour ⇒ Path A ≠ Path B.
2. Ratio-4 cross-chunk overlap not reproducible at decode emission in CompressedKVCache.accumulate — the prefill compressor with overlap=True produces row K (K≥1) using chunk K-1's first-half features + chunk K's second-half features (via _overlap_transform). Decode-time emission only saw the current chunk, missing the first-half. Fixed by carrying a _prev_chunk slot and, when overlap is enabled, calling compressor(prev || cur)[:, 1:2] to recreate prefill semantics exactly.
3. Non-causal pool prepend mask in V4Attention.__call__ — once Bug 1 was inverted to "include compressed pool", non-last queries in re-prefill saw pool rows summarising future tokens, while Path B's incrementally-grown pool only contained already-emittable rows. Fixed with comp_visible = (q_pos + 1) >= (k_idx + 1) * r as the prepend mask.
4. Sliding-window cutoff missing on no-cache forward in DeepseekV4Model.__call__ — create_attention_mask was called without window_size=, so re-prefill mask was causal-only while RotatingKVCache physically dropped keys outside the window. Fixed by passing window_size=self.args.sliding_window.

_prev_chunk is also propagated through merge / filter / extend / extract / nbytes for batched-cache parity.

Verification

pytest tests/test_deepseek_v4_decode_cache.py -v --timeout=300
  7 passed  (3 logit-parity at atol=1e-5, 3 argmax-parity, meta-state round-trip)

pytest tests/test_models.py --timeout=300
  83 passed, 2 skipped  (no regressions in V4 / cache / other models)

The original bisection log (.bisect-results/v4-bisect-results.md) and the regression tests remain in-repo as permanent guards. The xfail markers from the handoff revision are removed (the tests are now hard guards).

Diagnostic from the smallest failing case (kept here for the record)

prompt_len=4 BEFORE this fix:

no-cache argmax:   [461, 326, 726, 459, 702, 102, 676, 979]
with-cache argmax: [461, 326,  10, 102, 461,  65, 449, 102]
                          ^-- divergence at index 2 (cache.offset 5→6)

After the fix, both paths produce byte-identical token sequences across all parameterized prompt lengths (4 / 8 / 12) and 8 decode steps each.

Test plan

• test_deepseek_v4_decode_cache_logit_parity[4|8|12] — PASS
• test_deepseek_v4_decode_cache_argmax_parity[4|8|12] — PASS
• test_rotating_kv_cache_meta_state_round_trip_post_wrap — PASS
• tests/test_models.py -k deepseek_v4 — 8 PASS, 1 skipped (no regressions)
• tests/test_models.py -k cache — 5 PASS

Known follow-up (not in this PR)

CompressedKVCache.state / meta_state only proxy local's arrays — _pool / _buf / _buf_count / _prev_chunk are not yet round-tripped through save_prompt_cache / load_prompt_cache. Pre-existing gap, orthogonal to this fix; deserves its own commit + serialization round-trip test.

Reported-by: @anerjy
Co-Authored-By: Claude Opus 4.7 (1M context) noreply@anthropic.com

[anerjy]

Thank you, @jackneil 🙏

Just verified your fix on mlx-community/DeepSeek-V4-Flash-bf16 against my original reproducer:

NO-CACHE  : '{"verrid":"yes"}<eos>...'
WITH-CACHE: '{"verrid":"yes"}<eos>...'
MATCH: True ✅

Byte-for-byte parity between cache-on and cache-off paths — exactly what we needed.

Going one diagnostic layer deeper than my (admittedly incomplete) bisection paid off. I had ruled out the surface culprits and concluded the divergence had to be ULP-level numeric drift amplified by mHC + Sinkhorn. Forcing fp32 SDPA + fp32 Sinkhorn didn't move the needle — that's when I tapped out and posted the "this might not be a single fixable bug" follow-up.

You found four mutually-reinforcing defects I didn't see, and the chain of unmasking is genuinely beautiful:

1. Bool-dtype prepend mask zero=False — the prefill path silently masking out the entire compressed pool while decode (mask=None) attended it. So Path A and Path B were structurally inequivalent before any numerics even entered the picture. This is the kind of bug that hides in plain sight — mx.zeros(..., dtype=bool) reading as "block" is exactly the inverted convention I'd never have looked for.
2. Ratio-4 cross-chunk overlap unreproducible at decode emission — the prefill compressor's _overlap_transform carries chunk K-1's first-half features into row K, but decode emission only saw the current chunk. The _prev_chunk carry is the right shape of fix and clearly the result of careful reading of _overlap_transform.
3. Non-causal pool prepend mask — once Fix DeepSeek V4 checkpoint loading and mixed quantization #1 was inverted, the next layer of bug surfaced: non-last queries in re-prefill could see pool rows summarising future tokens. comp_visible = (q_pos + 1) >= (k_idx + 1) * r is exactly the right rule. Layered fixes that each unmask the next is the hardest debugging mode.
4. Sliding-window cutoff missing on no-cache forward — RotatingKVCache physically dropping keys outside the window while create_attention_mask was causal-only. This one's especially sneaky because the cache's physical behavior and the mask's causality were each individually correct.

The 7 regression tests + the bisection log staying in-repo as permanent guards is the right call — these failure modes are exactly the kind that an unrelated future refactor would re-break silently.

Heads-up on the batch path for whichever maintainer picks this up: mlx_lm.server's batch generator passes an mx.array offset (from BatchRotatingKVCache) into V4Attention.__call__, and the q_pos = mx.arange(offset, offset + S).reshape(S, 1) line at the new causal-pool mask path crashes with TypeError: arange() incompatible argument types. Single-request flow (the one I tested in the repro) is unaffected. I'm working around it locally by hiding cache.merge via a property descriptor so the server falls back to the non-batch path, but that's a band-aid. Probably wants a if isinstance(offset, mx.array): offset = offset.max().item() (or the proper batched broadcast) in the same hunk that introduced q_pos.