feat: add DeepSeek-V4 (Pro/Flash) model support

[machiabeli]

Summary

Adds model_type: deepseek_v4 support for DeepSeek-V4-Pro (1.6T/49B active) and DeepSeek-V4-Flash (284B/13B active), released April 22, 2026.

V4-novel architecture features implemented:

• Manifold-constrained Hyper-Connections (mHC) — replaces residual connections with a 4-copy hidden state (hc_mult=4). Each block reduces via learned pre weights, applies its sub-layer, then expands via post + doubly-stochastic comb matrix (20-iteration Sinkhorn-Knopp normalization on the Birkhoff polytope). Pure-MLX implementation; Metal kernel follow-up planned.
• Hash-routed MoE — first num_hash_layers (3) layers use a deterministic tid2eid table instead of learned gating, for stable early-layer routing.
• sqrtsoftplus scoring — sqrt(softplus(x)) expert scoring function (new in V4).
• Sliding-window + compressed KV attention — per-layer compress_ratios array (0 = pure window, 4 = light, 128 = heavy). Compressor module with learned gated pooling + APE. Indexer params loaded (topk sparse dispatch planned for v0.2).
• FP8 e4m3 block dequant — 128×128 blocks with ue8m0 scales, dequantized to bf16 at load time via mx.from_fp8.
• HyperHead — final sigmoid-weighted reduction from hc_mult copies to 1 before lm_head.
• Full weight sanitization — tested against real HF checkpoint key names (69K keys, 64 shards).
• Pipeline + distributed sharding via PipelineMixin.

What works now (v0.1):

• Config parsing from HF config.json ✓
• Full model init (284.2B params for Flash) ✓
• Prefill + cached decode ✓
• Hash routing + score-based routing ✓
• Sinkhorn doubly-stochastic verification ✓
• Weight sanitize mapping (FP8 dequant, expert stacking, name remapping) ✓

Planned for v0.2:

• Topk sparse attention dispatch (currently attends to all compressed KV rows)
• attn_sink integration via SDPA sinks= kwarg
• Metal-accelerated Sinkhorn kernel (4×4 per-token, batched)
• Generation quality validation on real V4-Flash weights

Memory estimates (V4-Flash):

Format	Size	Fits on
FP8 (native)	~160 GB	—
bf16	~568 GB	2-3 node cluster
q4	~142 GB	M3 Ultra 512GB (single node)

Test plan

• ModelArgs.from_dict(config) with real V4-Flash config
• Forward pass (prefill + decode) with random weights
• Hash-routed gate (layers 0-2) + score-based gate (layers 3+)
• Sinkhorn comb matrix is doubly stochastic (row/col sums verified)
• sanitize() correctly remaps all 69K checkpoint keys
• End-to-end generation with real V4-Flash weights (in progress — downloading)
• Benchmark: TTFT + tok/s on M3 Ultra

🤖 Generated with Claude Code

[Thump604]

I tested this against the DeepSeek-V4-Flash config and full checkpoint, then pushed the support branch here: https://github.com/Thump604/mlx-lm/tree/deepseek-v4-support-fixes

I also opened the same changes as a draft PR against the contributor branch so the review diff stays on top of this PR: machiabeli#1

What I changed on top of this PR:

• matched the DeepSeek V4 RoPE path more closely, including inverse output rotation for the post-attention rope dims
• passed attention sinks into SDPA
• decoded e8m0 block scales for FP8 and FP4 paths
• unpacked routed expert FP4 weights instead of treating the int8 tensors as ordinary int8 weights
• fixed ratio-4 overlap compressor shapes from the real checkpoint, for example ape (4, 1024) and wkv/wgate (1024, 4096)
• made mixed_3_6 preserve DeepSeek-specific attention, compressor, indexer, embedding, shared-expert, and output-critical paths at higher precision
• added a DeepSeek V4 loader path for F8_E8M0 scale metadata by reinterpreting those one-byte scale tensors as raw uint8 exponent bytes before sanitizer decode
• mapped the official final norm.weight key to model.norm.weight
• kept RoPE inv_freq derived from config instead of exposing it as a checkpoint parameter
• added a tokenizer fallback for unknown model configs so conversion does not require a Transformers release with native deepseek_v4 config support

Validation so far:

PYTHONPATH=/Users/David/code/mlx-lm-deepseek-v4-review \
  /opt/ai-runtime/venv-live/bin/python -m pytest -q tests/test_models.py tests/test_tokenizers.py
# 86 passed, 1 skipped, 5 warnings, 57 subtests passed

Full checkpoint conversion now completes for DeepSeek-V4-Flash with this local recipe:

mlx_lm.convert \
  --hf-path /Volumes/Lexar/hf-staging/deepseek-ai/DeepSeek-V4-Flash \
  --mlx-path /Volumes/Lexar/mlx_models/DeepSeek-V4-Flash-MLX-Q3-mixed-gs128-affine \
  --quantize \
  --quant-predicate mixed_3_6 \
  --q-group-size 128 \
  --q-mode affine \
  --trust-remote-code
# reported 3.808 bits per weight

The converted artifact lazy-loads cleanly in MLX. I would still keep this gated before merge on sparse compressed attention and learned indexer parity. The current path can convert and load the real Flash checkpoint, but it is not yet a complete correctness claim against the reference implementation.

[Blaizzy]

Awesome work guys!

Additionally, I would like to note that there are missing tests in test_models.py

[Blaizzy]

In general, look at deepeseek v3 implementation in this repo for inspiration around dequantization, predicates, and rope 👌🏽

[Thump604]

I pushed the follow-up branch updates here: https://github.com/Thump604/mlx-lm/tree/deepseek-v4-support-fixes

New commits since my last note:

• 6908736 casts DeepSeek V4 attention sinks to the query dtype before MLX SDPA. Without this, quantized generation fails with Type of sinks must promote to output type bfloat16.
• 9c990f4 fixes quantized grouped wo_a output projection. The dense path can reshape wo_a.weight, but a quantized QuantizedLinear stores packed input columns, so the output projection must slice output rows per group and call mx.quantized_matmul instead of reshaping packed weights as dense tensors.

Validation on the branch:

PYTHONPATH=/Users/David/code/mlx-lm-deepseek-v4-review /opt/ai-runtime/venv-live/bin/python -m pytest -q tests/test_models.py tests/test_tokenizers.py
87 passed, 1 skipped, 5 warnings, 57 subtests passed

Conversion evidence from the official deepseek-ai/DeepSeek-V4-Flash revision 6e763230a9d263eca2023f1d4a5ce1bfe126cf48:

• Q3 mixed artifact uploaded: https://huggingface.co/Thump604/DeepSeek-V4-Flash-MLX-Q3-mixed-gs128-affine
• Q3 recipe: mixed_3_6, affine, group size 128, effective 3.808 bpw, 28 shards, indexed tensor size 135,346,422,876 bytes
• Q3 lazy-loads on a 128 GB Mac Studio, but I do not consider it generation-qualified locally. A one-token smoke crossed my memory safety boundary with heavy swap activity and was killed.
• Q2 mixed artifact uploaded: https://huggingface.co/Thump604/DeepSeek-V4-Flash-MLX-Q2-mixed-gs128-affine
• Q2 recipe: mixed_2_6, affine, group size 128, effective 2.992 bpw, 23 shards, indexed tensor size 106,355,393,628 bytes
• Q2 raw generation smoke completed with --max-tokens 2 --max-kv-size 1024: output the name, prompt 4 tokens at 7.488 tok/s, generation 2 tokens at 19.182 tok/s, 54.59s real, 74.5GB max RSS, 106.94GB peak footprint, zero swaps.

I still treat full sparse compressed-attention/indexer parity as unproven. The branch now has conversion, lazy-load, and Q2 smoke-generation evidence, but not a production-quality or long-context claim.

[Thump604]

I validated the current PR head after it pulled in my support branch:

PYTHONPATH=/Users/David/code/mlx-lm-pr1189 \
  /opt/ai-runtime/venv-live/bin/python -m pytest -q tests/test_models.py tests/test_tokenizers.py

Result: 87 passed, 1 skipped, 5 warnings, 57 subtests passed.

I do not have another local patch to push right now. The branch now has the important unit coverage I wanted to see for V4 shape handling, RoPE inverse, FP4/FP8/E8M0 loading, mixed quantization predicate, and quantized grouped output projection.

My remaining caveat is the same claim boundary as before: this is strong conversion/load/unit-test evidence, but I would still avoid framing it as complete long-context or production-quality sparse compressed-attention/indexer parity until that path is compared against the reference implementation. I also would not treat the uploaded Q2/Q3 artifacts as local 128 GB runtime candidates; they are useful conversion artifacts and support evidence, not a quality-qualified serving lane for that class of machine.

[Thump604]

I validated the new 469433c head locally after the indexer top-k update:

/opt/ai-runtime/venv-live/bin/python -m pytest -q tests/test_models.py tests/test_tokenizers.py -k 'deepseek_v4 or deepseek_unknown'
# 8 passed, 81 deselected, 4 warnings

The new test coverage exercises the indexer shape/bounds, full forward on the compressed attention path, and prefill to decode cache transition. The claim boundary still looks right to me: this is now stronger unit-level coverage for compressed sparse attention prefill/top-k selection, but not yet a full production parity claim because the incremental compressor decode path is still deferred.

No additional local patch from me right now. I still would not recommend treating the uploaded Q2/Q3 artifacts as 128 GB Mac runtime candidates. They remain useful conversion artifacts and support evidence for this PR, not a local serving lane for that class of machine.

[machiabeli]

Thanks @Blaizzy — test_deepseek_v4 landed in 37e0e3c (pushed ~11 min after your comment, must have crossed in flight). Covers the standard model_test_runner matrix (fp32/fp16 forward, prompt cache prefill + single-token decode, batch > 1, deepcopy), with V4-specific dims exercising mHC, hash-routed MoE (num_hash_layers=2 of 4), and o_groups split.

Writing the test caught a real bug the smoke test hid: MoEGate hash branch returned flat (B*S, top_k) inds while the non-hash branch returns (B, S, top_k); SwitchGLU broadcast happened to work at B=1 and failed at B>1. Fixed in 6a71579.

Also looking at V3's quant/predicate/rope patterns per your suggestion — @Thump604 has a draft PR against this branch (machiabeli#1) that adds mixed_3_6 predicate, inverse output RoPE, attention sinks through SDPA, and F8_E8M0 block-scale decoding. Pulling that in shortly.

[machiabeli]

Merged @Thump604's draft PR (machiabeli#1) into this branch at 45665f8.

Landed:

• DeepseekV4RoPE with inverse output rotation for post-attention rope dims
• attn_sink integrated into SDPA (with dtype cast fix for bf16)
• F8_E8M0 block scale decode — native FP8/FP4 loading now works without upcasted bf16 staging
• Routed expert FP4 weight unpacking
• Ratio-4 Compressor shapes validated against real checkpoint
• mixed_3_6 quant predicate preserving attn/compressor/indexer/embedding/shared-expert/output paths at higher precision
• Quantized grouped output projection via per-group slice + mx.quantized_matmul (replaces dequant+reshape)
• Tokenizer fallback (cleaner version of the same approach I had — now keeping Thump604's)

Plus our earlier perf work stays on top:

• Fused Metal kernel for mHC Sinkhorn (1.83x end-to-end decode measured on M3 Ultra: 11 to 20 tok/s at B=1)
• Hash-routing reshape fix for B>1 MoE

Test matrix now at 6 V4 tests, all passing: test_deepseek_v4, test_deepseek_v4_rope_inverse, test_deepseek_v4_quantized_grouped_output_projection, test_deepseek_v4_sanitize_unpacks_fp4_experts, test_deepseek_v4_sanitize_dequantizes_fp8_blocks, test_deepseek_v4_loads_e8m0_scales_as_uint8.

Ran 6 tests in 0.596s — OK

Thump604 has already uploaded validated Q2/Q3 mixed artifacts to HF; I've uploaded a (now outdated) Q4 at mlx-community/DeepSeek-V4-Flash-4bit and will re-quantize from the native FP8 source once V4-Pro finishes downloading.

[machiabeli]

v0.2: Indexer topk for compressed sparse attention

Pushed 469433c — implements the Indexer forward pass that was deferred in v0.1.

What changed:

For ratio-4 layers (20 of 43 in Flash, 29 of 61 in Pro), the indexer now runs a lightweight scoring pass over all compressed KV rows (using its own compressor at index_head_dim=128) and selects the top-k rows before attention:

• Flash: topk=512 from ~S/4 compressed rows
• Pro: topk=1024 from ~S/4 compressed rows

Selected rows are gathered and prepended to the sliding-window KV before SDPA with a zero-mask prefix (compressed rows are always "past"). This reduces per-layer attention cost from O(S/4) to O(topk) — roughly 500x at 1M context.

Also in this commit:

• hc_post einsum replaced with direct matmul (@): eliminates einsum overhead on the hot decode path (called 86x/token for Flash, 122x/token for Pro)
• Sinkhorn Metal kernel: fixed constant-vs-device address space cast failure on small inputs by switching from pointer variables to base+offset indexing

Test coverage:

8 passed in 4.08s

New test_deepseek_v4_indexer_topk validates:

• Indexer output shape and index bounds
• Full model forward with 32-token prefill exercising the compressed attention path
• Prefill to decode transition with cache

Live generation validated:

Generation: 48 tokens, 21.863 tokens-per-sec
Peak memory: 160.190 GB

DeepSeek-V4-Flash-4bit on single M3 Ultra (512GB).

What is next (v0.3):

• Compressor decode path (emit compressed rows incrementally during generation)
• V4-Pro config validation (61 layers, 384 experts, 16 o_groups, q_lora_rank=1536)
• Fused Metal kernel for grouped output projection (16 sequential quantized_matmul calls to 1)

[onchainengineer]

Hi — I've been validating this PR end-to-end against the antirez/ds4 reference C+Metal implementation (which is logit-validated against the official DeepSeek-V4 release) as a numerical oracle, on the DeepSeek-V4-Flash-4bit MLX checkpoint (M3 Ultra 256 GB, 256-bit MLX vs ds4's q4-imatrix GGUF — different quants, so we measure symptom-level agreement: leading-word match, non-empty output, no pathological char/word loops).

Pre-fix: everything ≥ ~210 tokens collapses into pathological repetition. Threshold lines up exactly with sliding_window=128. Inside the window the model works fine; the moment a query has to look further it produces output like "4.4.4.4.4.", "Which is which? Which is which?", "# The\n\n## The\n\n## The..." and so on.

Root cause (sourced from ds4 — see ds4.c:4817-4860 layer_rope_freq_base and ds4.c:6580-6581 compressed-pool RoPE):

1. Per-layer RoPE base is wrong. The 40 compressed layers were trained with compress_rope_theta=160000 + YaRN (factor=16); the 3 dense layers with rope_theta=10000 and NO YaRN. The PR currently uses (10000 + YaRN) for all 43 layers. Symptom is invisible inside the local window (q_pos − k_pos deltas are tiny).
2. Compressed pool keys never RoPE'd. The compress_rope is allocated in V4Attention.__init__ but never invoked; pool keys carry no positional signal so attention against them is geometrically incoherent.
3. Compressed-pool mask block-zeros the entire pool. comp_mask = mx.zeros(comp_shape, dtype=mask.dtype) with dtype == bool is all-False == "block" in MLX SDPA convention. Pool was never attended to at all.

I've opened 5 atomic draft PRs against machiabeli:feat/deepseek-v4 covering these + two adjacent issues raised in your review thread (batched-server KV cache "stale pool rows" + S=1 SDPA-with-sinks mlx#3452 ULP drift):

• fix(deepseek_v4): per-layer RoPE base + compressed-pool RoPE at chunk start machiabeli/mlx-lm-1#5 — per-layer RoPE base + chunk-start compressed-pool RoPE (the smoking gun)
• fix(deepseek_v4): mask compressed-pool columns as causal-allow, not bool-zero machiabeli/mlx-lm-1#6 — causal compressed-pool mask
• fix(deepseek_v4): default HC Sinkhorn Metal kernel to off machiabeli/mlx-lm-1#7 — HC Sinkhorn Metal kernel default-off (PR review thread)
• fix(deepseek_v4): fail loudly on mismatched-state CompressedKVCache merge/extend machiabeli/mlx-lm-1#8 — fail loudly on batched CompressedKVCache mismatch
• fix(deepseek_v4): pad S=1 query to bypass SDPA-with-sinks mlx#3452 machiabeli/mlx-lm-1#9 — S=1 SDPA workaround for mlx#3452

Verification with all 5 applied: 7/7 prompts pass against the ds4 oracle, spanning 5 tokens → 3500 tokens. The most striking case: at ~1960 tokens, post-fix mlx-lm produces "...</think>DeepSeek V4 Flash introduces Multi-head Latent Attention with Decoupled" — the post-</think> text matches ds4's first words verbatim.

Each PR has its own commit, comment trail, and test plan; machiabeli#5 and machiabeli#6 must land together to fix the ≥128-token collapse (they're logically coupled but touch different code paths). Happy to add unit tests, rebase, or restructure however helps the merge — let me know.

cc @pcuenca @awni

[aidiffuser]

We've been running this PR's deepseek_v4 on a 2× Mac Studio (512 GB) MLX
cluster since the 0731 release dropped, at contexts up to 256k. It works
beautifully — thank you! Two small loader fixes from profiling it hard, both
model-output-preserving, plus one likely root-cause for #1332.

1. Decode slows down linearly with generated tokens (likely the cause of #1332)

Decode starts ~29 tok/s and decays linearly (~4.3 µs/tok per generated token
on M-series; 28.8 → 17.7 by 6k generated, worse forever after). sample
profiles show mlx::core::detail::CompilerCache::find + memcmp growing
from 12% → 29% of decode-thread time between token ~2k and ~7k.

Cause: the @mx.compile'd rope helpers (_attn_qkv_partial_rope,
_attn_inv_rope_flatten, _compressor_norm_strided_rope) receive the cache
offset as a python int. mx.compile treats scalars as cache-key
constants, so every generated token mints a brand-new cache entry (plus a
re-trace), and lookup is a linear scan. This also grows host memory without
bound (~3 entries retained per token, forever) — which looks exactly like
the unbounded residency growth reported in #1332.

Fix — pass the offset as a 0-d array (a traced input → one cache entry):

def _offset_input(offset):
    """Wrap python-int cache offsets as 0-d arrays before mx.compile'd fns:
    an array is a traced input (one compile-cache entry); a python int is a
    per-value constant (one entry per token during decode)."""
    if isinstance(offset, mx.array):
        return offset
    return mx.array(offset, dtype=mx.int32)

-        q, kv = _attn_qkv_partial_rope(q, kv, offset, rd, self.rope.freqs)
+        q, kv = _attn_qkv_partial_rope(
+            q, kv, _offset_input(offset), rd, self.rope.freqs
+        )

-        o = _attn_inv_rope_flatten(
-            o, offset, rd, self.rope.freqs, self.n_heads * self.head_dim
-        )
+        o = _attn_inv_rope_flatten(
+            o, _offset_input(offset), rd, self.rope.freqs,
+            self.n_heads * self.head_dim,
+        )

             new_pooled = _compressor_norm_strided_rope(
                 new_pooled,
                 self.norm.weight,
                 self.norm.eps,
-                int(pool_base) // ratio,
+                _offset_input(int(pool_base) // ratio),

Verified: temp-0 outputs are byte-identical before/after; decode is flat
(~31 tok/s at 6k generated where stock had decayed to 17.7), and prefill
gains ~15% too (it was paying one cache entry per chunk offset per layer).
The batch path already passes array offsets into these same functions, so
in-graph support is already exercised.

2. Top-k membership mask materializes [B, S, topk, T]

In V4Attention, the indexer-membership test builds a broadcast compare:

selected = (indexer_topk[..., None] == k_range[None, None, None, :]).any(axis=-2)

At a 4096-token prefill chunk against a 32k-row pool (128k ctx) that's a
[1, 4096, 512, 32768] bool ≈ 34 GB per layer per chunk, immediately
reduced away. A scatter builds the identical mask with O(S·topk) writes:

-                    k_range = mx.arange(compressed_len, dtype=mx.int32)
-                    selected = (
-                        indexer_topk[..., None] == k_range[None, None, None, :]
-                    ).any(axis=-2)[:, None, :, :]
+                    selected = mx.put_along_axis(
+                        mx.zeros((B, S, compressed_len), dtype=mx.bool_),
+                        indexer_topk.astype(mx.int64),
+                        mx.ones(indexer_topk.shape, dtype=mx.bool_),
+                        axis=-1,
+                    )[:, None, :, :]

Equality is exact (mx.array_equal verified incl. duplicate indices). On our
cluster this plus a fused indexer scorer took 64k-context prefill from 191 →
335 tok/s and cut peak transient memory ~4×; the scatter alone carries most
of the speed win and applies to any hardware.

The diffs above are the complete change in both cases — if you'd rather
review them as PRs against this branch, say so and they'll appear as-is.
Full measurement notes: https://gist.github.com/aidiffuser/fef1890680e4eed60d3902511b02a696.

Co-authored with Claude Fable 5 (Anthropic)

[OmarB97]

We ran this PR's deepseek_v4 as a 2× M4 Max (128 GB) pipeline over Thunderbolt (mlx ring backend) against the 0731 release in its original reference layout (the MXFP4 MLX conversion), and differentially validated every layer kind against the reference implementation that ships with the checkpoint — real weights, real prompt, per-module diffs. Three model-level fixes came out of it, plus loader/pipeline work. Branch with everything: https://github.com/OmarB97/mlx-lm/tree/ds4-0731-2mac — happy to split any of it into PRs against this branch on request.

1. RoPE regime is per layer kind (ratio-4 attention cosine 0.926 → 0.999)

The reference initializes one rope per layer: sparse layers (compress_ratio > 0) rotate q/k/o with compress_rope_theta (160000) with YaRN, dense layers use rope_theta (10000) with YaRN disabled. The PR ropes every layer's main q/k at rope_theta+YaRN — the wrong base on ~40 of 43 Flash layers. After matching the regime (and using the same instance for the compressed pool), ratio-4 layer attention went from cos 0.926 to 0.999 vs reference. This is also a plausible root cause for the CJK periodic token-drop reports — the earlier two-instance change moved the wrong theta around rather than fixing it.

2. Compressed pool rows must be roped at birth

The reference applies rope.apply_positions(kv[..., -rd:], positions) to each pooled row at its window-start position (0, r, 2r, …) before caching. The PR concatenates un-roped pool rows into K. After roping at birth (prefill + the decode rolling-state path), our pool rows match the reference bit-close (cos 1.00000).

3. Compressed-row visibility mask: boolean masks invert/erase it

create_attention_mask(..., return_array=True) returns a boolean keep-mask here. The compressed-mask block builds mx.zeros(comp_shape, dtype=mask.dtype) — which coerces to all-False, i.e. compressed rows are entirely invisible during prefill (and a float 0/-inf block coerces inverted). Fix is convention-aware and causal per the reference (get_compress_topk_idxs): query at position p sees row j iff j < (p+1)//ratio; decode sees all rows.

Also on the branch

• PipelineMixin.pipeline() drops a layer on uneven splits (43 layers over 2 ranks leaves layer 21 on no rank) — exact per-position partition; plus DeepseekV4Model must retighten num_layers after pipeline() or 2-rank runs index past the truncated list.
• Loader support for reference-named quantized checkpoints (the 0731 MXFP4 conversion): expert stacking must carry scales/biases (weight-only stacking leaves switch_mlp unquantizable → raw gather_mm shape errors), per-module quantization-config lookups need pre-sanitize aliases (attention otherwise wraps with the wrong mode and strict=False silently loads nothing), and sharded_load's per-rank file resolution can't match post-sanitize names against a reference-named index.
• sharded_load now materializes weights layer-by-layer on the CPU stream — monolithic Metal command buffers for ~80 GiB cold materializations intermittently never complete on macOS 26.5.x (parked forever in CommandEncoder::commit); unified memory makes CPU-stream loading free.

Verification

• Per-layer attention cosine vs the bundled reference implementation: dense 0.9992, ratio-4 0.9991, ratio-128 0.9994 (residual ≈ the reference's act-quant QAT sims, which we don't apply).
• 2-rank pipeline logits are bit-identical (max|Δ| = 0.000000) to single-process on a small random-weight fixture.
• The PR's deepseek_v4 test suite passes (plus new partition/sync tests); two small test-infra fixes included (typing.Any import, mixed-quant path parsing).

Co-authored with Claude Fable 5 (Anthropic).

[ncdrone]

Found a small blocker while running this PR on Apple Silicon: mlx_lm/models/deepseek_v4.py uses Optional[Any] in the signature at line 1104, but Any is missing from the typing import on line 13:

from typing import Dict, List, Optional   # <- no Any

On Python <= 3.13 this raises NameError: name 'Any' is not defined at module import, so the arch fails to load entirely. It goes unnoticed on Python 3.14 because PEP 649 lazy annotations defer evaluation.

One-line fix:

from typing import Any, Dict, List, Optional

Verified at 63a2662: import fails on 3.12.13, succeeds on 3.14.3; with the import fixed it loads fine on 3.12. Happy to open a patch PR against the branch if useful.

(Context: running DeepSeek-V4-Flash-2bit-DQ through a benchmarking harness on an M4 Max — the arch itself has been generating fine for us under this PR once past the import.)

[ncdrone]

Follow-up to my earlier comment (the Any import): after getting past the import, DeepSeek-V4-Flash-2bit-DQ generated incoherent output on Apple Silicon at 63a2662 — token soup degenerating into repeated # within a few hundred tokens, on both greedy and sampled decoding.

Applying the two open PRs on the fork branch fixed it completely: machiabeli#5 (per-layer RoPE base + chunk-start RoPE for the compressed pool) and machiabeli#6 (causal allow-mask for pool columns). With both applied, generation is fully coherent on an M4 Max 128GB (2bit-DQ quant, ~90 GB resident, py3.14) — verified across greedy, temp 1.0, templated and raw prompts.

So for anyone hitting garbage output on this PR: it's the compressed-attention RoPE + pool mask, and the fixes already exist on the fork. Would be great to see #5/#6 folded into this PR.

[PhilipJohnBasile]

I opened a consolidated draft patch against feat/deepseek-v4: machiabeli#10

It folds in and explicitly credits the RoPE, compressed-mask, Sinkhorn, cache-safety, 0731 parity, and Python import findings already reported by @onchainengineer, @OmarB97, and @ncdrone. It also adds official compressor rolling-state parity tests, vector-offset mask handling, 0731 DSpark target-only detection, and a fail-closed speculative-rollback gate.

The draft is intentionally narrow: it does not claim full 0731 parity, DSpark support, a source-faithful indexer, prompt-cache serialization, generic batching/KV-cache quantization, long-context quality, or performance. The focused CPU suite is 11/11, and an independent final review approved the frozen diff.

[Thump604]

@PhilipJohnBasile Thanks for consolidating this and for keeping the limitations explicit. I inspected draft PR machiabeli#10 at 9deff61 against 63a2662. git diff --check is clean, and a broader CPU selection (deepseek_v4 or hc_sinkhorn) reached 17 passed / 1 skipped; its only head failure was the untouched mixed_quant_predicate_builder path parser. On the exact base, that test fails even earlier at the missing Any import, so the draft correctly removes the first blocker and exposes the pre-existing converter failure.

The consolidation is a good stabilization vehicle, but it should not yet be treated as evidence that #1189’s advertised feature set is complete. Please reconcile this PR’s summary/checklist with the draft’s narrower contract before promotion: qualify cached decode as the supported single-stream target path; do not claim source-faithful top-k sparse attention; and explicitly call out unsupported generic fresh-cache batching, prompt-cache serialization, DSpark speculation, real-checkpoint generation, long-context quality, and performance. Where an unsupported batching or cache path can still be entered, I would prefer an explicit fail-closed capability guard rather than relying only on a known-limitations paragraph.

Also, the draft currently has no GitHub checks or recorded reviews, so the “independent final review” is not verifiable from the PR. Please attach that evidence or have the reviewer submit it formally once the diff is frozen. Keeping #10 as a draft until those scope claims and gates are reconciled is the right posture.

[PhilipJohnBasile]

Thanks for the concrete review. I pushed the revised consolidation draft at PhilipJohnBasile/mlx-lm@b759c7f and kept machiabeli#10 in draft.

The draft body now matches the narrower contract: one stream, initial prefill, then single-token cached decode. It does not claim source-faithful top-k sparse attention, generic batching, prompt-cache serialization, DSpark, real-checkpoint generation, long-context quality, or performance.

The unsupported overlapping-insertion path now fails closed before scheduler/cache mutation. The final regression matrix includes compressed-only, realistic mixed [0,128,4,0], and all-uncompressed topologies; atomic second-insert rejection; later sequential reuse; and nested/callable cache capability handling. Focused final tests are 23 passed / 1 skipped plus 22 prompt-cache tests; compilation and diff checks pass.

An independent Codex Sol-max reviewer approved the complete frozen diff after earlier passes rejected several concrete cache-lifecycle defects. I have described that as independent AI review, not formal GitHub review, and asked @Thump604 to submit a formal review on #10 if the final scope and gates are acceptable.

#10 remains a stabilization vehicle, not evidence that #1189's advertised feature set is complete. Please reconcile this PR's summary/checklist with that narrower contract before promotion.

[Thump604]

Independent external review only: I am not a maintainer or collaborator on ml-explore/mlx-lm, and this is offered as review assistance rather than project approval.

I reviewed the frozen stabilization draft machiabeli#10 at b759c7f. Within its explicitly narrow contract—one stream, initial prefill, then single-token cached decode—I found no blocking defect. Its non-batchable capability preflight rejects overlapping insertion before scheduler queues or active cache state are mutated; the singleton cache lifecycle and sequential reuse are covered. My local verification passed all 18 newly added focused regressions (plus 3 subtests), the 22-test prompt-cache suite, and git diff --check.

That result does not validate this PR as currently advertised. #1189 remains conflicting and its summary/checklist still claims a broader feature set than #10 qualifies, including general cached decode and model support while source-faithful indexer behavior, generic batching, DSpark speculation, real-checkpoint generation, long-context quality, and performance remain unsupported or unverified.

Recommendation: hold #1189 until the stabilization work is integrated or rebased, the PR body/checklist is narrowed to the proven contract, conflicts are resolved, and project CI/reviewer evidence is present.

[stepnoy]

Repro: progressive corruption starting at ~99 prompt tokens (boundary at sliding_window), checkpoint ruled out

Environment: 2× Mac Studio M3 Ultra 512 GB (repros identically on a single node and on a 2-node JACCL TP cluster), macOS 26.6.1, Python 3.12, mlx==0.32.0, this branch @ 63a26625c7 (current head).
Checkpoint: mlx-community/DeepSeek-V4-Flash-mxfp4, all 41 files sha256-verified against lfs.oid.

Symptom

Needle-in-haystack ladder (non-repetitive procedurally-generated filler, chat template applied, greedy decoding, needle = a password planted mid-prompt):

prompt tokens	result
16 (17×23)	✅ 391
~67	✅ The maintenance password is PELICAN-7734.
~99	❌ PELIC-7731 — password garbled, sentence structure intact
~131	❌ echoes filler text
~163 … ~6100	❌ fully degenerate repetition loops

The corruption is progressive, not a cliff — at 99 tokens the model still finds the needle but retrieves it corrupted. The onset sits right at sliding_window: 128, far below index_topk: 512, so the indexer/top-k path is not involved.

Ruled out

• Checkpoint / weights — the same-source Vontra/DeepSeek-V4-Flash-0731-MXFP4-MLX build run through its bundled standalone runtime retrieves the needle perfectly at ~800 prompt tokens (PELICAN-7734.). Weights family is healthy.
• Quantization — consistent with @anerjy's 04-28 finding that the long-prompt failure is quant-blind (bf16 fails identically).
• Repetitive-filler artifact — repro'd with fully diverse filler sentences; same degeneration.
• Prefill chunking — prefill_step_size 32/64 changes the output slightly but does not fix it.
• Sharding — single-node and 2-node TP degrade identically, so it's not shard().

Side note: head does not import on Python 3.12

MTPBlock (line ~1104) uses Any in a class-body annotation without importing it → NameError: name 'Any' is not defined at module import on 3.12 (lazy annotations on 3.14 mask this). One-line fix: add Any to the typing import.

Question

@aidiffuser — since it works for you at 256k on the same hardware and the same head: which checkpoint (and Python version) are you running? Given the ladder above we'd expect any run through mlx_lm on 3.12 with an mxfp4 community checkpoint to break at ~100 tokens, so whatever differs between our setups should pinpoint the bug.

Happy to run any diagnostic build/patch against this ladder — it takes minutes on our side.

Co-authored with Claude (Anthropic).

[aidiffuser]

@stepnoy — great ladder, and the answer to your question turns out to matter a lot: we are not running current head. Our fork predates the recent rewrite (no MTPBlock, and a single DeepseekV4Cache rather than the CompressedKVCache/RotatingKVCache pair), which is also why the Any NameError never hit us. So our "works at 256k" vouches for an earlier revision, not for 63a2662 — your bisect is inside this PR's history, not between our setups.

For completeness: Python 3.13.12, mlx 0.32.0.dev20260609, mlx-lm 0.31.3, and a checkpoint we converted ourselves from the official DeepSeek-V4-Flash-0731 release (fp4 experts reinterpreted losslessly to mxfp4, everything else affine-8 g64) — not the community mxfp4 build.

That said, we went looking in head with our working implementation as an oracle, and the ladder is explained by one line.

The compressed/global branch is masked OFF during every prefill

deepseek_v4.py (head @ 63a2662), V4Attention.__call__:

if mask is not None:
    comp_shape = list(mask.shape)
    comp_shape[-1] = n_comp
    comp_mask = mx.zeros(comp_shape, dtype=mask.dtype)   # <-- here
    mask = mx.concatenate([comp_mask, mask], axis=-1)

mask here comes from create_attention_mask(..., return_array=True) → RotatingKVCache.make_mask → create_causal_mask, which returns a boolean array (linds >= rinds). For a boolean mask, MLX's SDPA treats True as attend and False as masked out. So mx.zeros(..., dtype=bool) is all-False: every compressed row is forbidden to every query.

Verified on our machine (mlx 0.32.0.dev, Metal):

create_causal_mask(99, 0, window_size=128).dtype   # mlx.core.bool
mx.zeros((99, 24), dtype=mx.bool_)                 # all False
# SDPA bool semantics, 3 keys, one row each:
#   mask [True, False, False] -> output == v[0]
#   mask [False, False, True] -> output == v[2]

And end-to-end: building head's exact mask for S=99, then scaling the pool keys by 50×, changes the attention output by 4.17e-07 against mean|out| = 0.26 — pure reassociation noise. The pool contributes exactly nothing during prefill.

There's a second half that makes it worse: at decode S == 1, make_mask returns None, so the if mask is not None guard is skipped and the pool is attended. Every prompt token (and every pool row built from those tokens) is computed by a pure 128-token sliding-window model, and then generation reads a pool the prefill was never conditioned on.

This reproduces your ladder exactly. With the global branch dead, a query at position p sees only the last sliding_window = 128 raw tokens, and the needle at prompt position p_n stays visible only while S_prompt + g − p_n ≤ 127 (g = tokens generated so far):

prompt	needle visible for	observed
16, 67	the whole answer	verbatim ✅
~99	first ~35–50 generated tokens	starts right, then the window slides past the needle mid-answer → PELIC-7731, local coherence intact
~131	already out of window at step 0	echoes the tail filler
~163+	window contains only filler	degenerate repetition

It also explains your "ruled out" list: it's independent of weights, of quantization, and of sharding, and prefill_step_size only perturbs offsets without un-blocking anything.

The fix is a real visibility mask instead of a constant. Compressed row j pools raw tokens [j*ratio, (j+1)*ratio − 1], so it is visible to a query at raw position p iff (j+1)*ratio <= p+1. Ours builds exactly that and concatenates it (AND-ed with the indexer's top-k membership when the indexer is active).

Two more in the same branch, likely latent behind the first

While the pool is blocked in prefill these can't show up, so they're worth checking right after the mask fix:

1. self.compress_rope is constructed and never called. In head it appears only at its definition (~line 823); there is no call site, so pool rows carry no positional rotation at all. Our compressor applies a strided RoPE to each pooled entry as it is produced.
2. Compress layers rotate q/k with the wrong base. The config carries rope_theta = 10000 and compress_rope_theta = 160000 (plus YaRN); head's compress layers use the former. Ours selects compress_rope_theta for those layers.

Send a patch and we'll validate it against a needle ladder on our side — 2× M3 Ultra, single-node and 2-node TP, minutes per run. We can also run your exact ladder against head on our own checkpoint, which would confirm the repro is checkpoint-independent from our end.

Co-authored with Claude Fable 5 (Anthropic).

[stepnoy]

@aidiffuser — your diagnosis was exact, thank you. Patched head accordingly; three of your items are confirmed fixed by the ladder, one gap remains, localized by a probe. Diff below; validation runs welcome.

What the patch does (in dependency order)

1. Real visibility mask for pool rows (your §"masked OFF"): row j visible to query at absolute position p iff (j+1)*ratio <= p+1; handles boolean and additive masks. Confirms your table — with just this fix, degeneration turned into coherent-but-pool-blind output.
2. RoPE bases per reference: dense (ratio 0) layers use plain rope_theta, YaRN disabled; compress layers use YaRN + compress_rope_theta for q/k/inverse-output; pool rows rotated at their window starts j*ratio via a new apply_positions on DeepseekV4RoPE. (This was your latent MLX-my-repo - a no-code way to create MLX quants (Q4/ Q8) #1/ValueError: Model type gemma3 not supported. #2 — both real; without them the model "saw" the pool but produced garbage of a different flavor.)
3. Indexer: bound to the layer rope (q at absolute positions — note DeepseekV4RoPE.__call__ rotates the first dims, so the last-rope.dims slice must be passed explicitly), rows at j*ratio; causality applied before ranking; scoring rewritten to reference semantics (raw weights_proj output with folded scale, no sigmoid, ReLU on logits, sum over heads — head had sigmoid+mean); per-query top-k via a membership mask AND-ed into the visibility mask (the shared prompt-wide gather drops mid-prompt rows — your scatter snippet gave this away); decode (S==1) gathers its own top-k so the softmax isn't diluted over the full pool.

Results (needle ladder, mlx-community/DeepSeek-V4-Flash-mxfp4, single node)

prompt tokens	before	after
~267 / ~467 / ~866 / ~1567	❌ degenerate	✅ verbatim
~4101	❌	✅ verbatim
~6133	❌	⚠️ coherent refusal ("cannot determine")

Decode throughput unchanged (~30–32 tok/s). The reference standalone runtime (Vontra) retrieves the same 6133-token needle, so the remaining gap is implementation, not model capacity.

The remaining gap, localized

Instrumented Indexer.__call__ on the 6133 case: the needle's pool row (765 of 1533) is absent from the final query's top-512 in 0 of 42 indexer invocations. With ~512/1533 per-layer inclusion, random noise would miss all 42 with p ≈ 4e-8 — the ranking systematically down-ranks it, while at ≤4101 tokens (pool ≤ 1025) retrieval survives. Things already replicated from the reference: ReLU-before-weighting, no sigmoid, folded index_head_dim^-0.5 * n_heads^-0.5 scale, sum over heads, fp32 scoring, causal mask before argpartition. Not replicated: the Hadamard rotate_activation + fp4 activation-sim pair (orthogonal-invariant for the dot product, so it should only matter through quantization noise) — and whatever else your working indexer does differently.

Could you share your indexer scoring path (or diff it against head's)? That's the one remaining piece between head and a clean ladder, and your fork predates the rewrite so the divergence should be easy to spot on your side.

Full diff vs head @ 63a2662 (208 lines)

--- mlx_lm/models/deepseek_v4.py.orig	2026-08-12 19:02:31
+++ mlx_lm/models/deepseek_v4.py	2026-08-13 11:07:47
@@ -238,6 +238,26 @@
     @property
     def inv_freq(self):
         return self._inv_freq[0]
+
+    def apply_positions(self, x: mx.array, positions: mx.array):
+        """Rotate the last self.dims dims of x at explicit absolute positions.
+
+        x: [..., T, D]; positions: [T] (float or int). Interleaved-pair
+        convention identical to the array-offset branch of __call__.
+        """
+        dtype = x.dtype
+        theta = positions.astype(mx.float32)[:, None] * self.inv_freq[None, :]
+        shape = (1,) * (x.ndim - 2) + (x.shape[-2], self.dims // 2)
+        cos = mx.cos(theta).reshape(shape).astype(dtype)
+        sin = mx.sin(theta).reshape(shape).astype(dtype)
+        rot = x[..., : self.dims].reshape(*x.shape[:-1], self.dims // 2, 2)
+        x0, x1 = rot[..., 0], rot[..., 1]
+        rotated = mx.stack(
+            [x0 * cos - x1 * sin, x0 * sin + x1 * cos], axis=-1
+        ).reshape(*x.shape[:-1], self.dims)
+        if self.dims < x.shape[-1]:
+            return mx.concatenate([rotated, x[..., self.dims :]], axis=-1)
+        return rotated
 
     def __call__(self, x: mx.array, offset: int = 0, inverse: bool = False):
         dtype = x.dtype
@@ -819,10 +839,16 @@
         # attention rotation to the wrong base on compressed layers, manifesting
         # as periodic token drops in CJK (cf. Shinka-Man's report on #1192,
         # fixed there in @Blaizzy/mlx-lm@b78ccb1).
-        self.rope = DeepseekV4RoPE(self.rope_head_dim, args.rope_theta, args.rope_scaling)
-        self.compress_rope = DeepseekV4RoPE(
-            self.rope_head_dim, args.compress_rope_theta, args.rope_scaling,
-        )
+        # Reference semantics: dense (ratio==0) layers use plain RoPE with
+        # rope_theta and YaRN disabled; compress layers use YaRN with
+        # compress_rope_theta for q/k/output *and* the pooled rows.
+        if self.compress_ratio:
+            self.rope = DeepseekV4RoPE(
+                self.rope_head_dim, args.compress_rope_theta, args.rope_scaling,
+            )
+        else:
+            self.rope = DeepseekV4RoPE(self.rope_head_dim, args.rope_theta, None)
+        self.compress_rope = self.rope
 
         # Compressor / Indexer — present only when compress_ratio > 0
         if self.compress_ratio:
@@ -901,6 +927,8 @@
 
         # --- Compressed sparse attention ---
         compressed_k = compressed_v = None
+        pool_topk_idx = None
+        pool_member_mask = None
         if self.compress_ratio:
             comp_cache = cache if isinstance(cache, CompressedKVCache) else None
             if comp_cache is not None:
@@ -912,15 +940,43 @@
                 pool = None
 
             if pool is not None:
+                # Pool rows carry no positional rotation in the cache; rotate
+                # row j at its window-start position j*ratio (rope dims are the
+                # last rope_head_dim of each row, as in the raw K path).
+                n_pool = pool.shape[1]
+                pool_pos = mx.arange(n_pool, dtype=mx.float32) * self.compress_ratio
+                rd_ = self.rope_head_dim
+                pool = mx.concatenate(
+                    [pool[..., :-rd_],
+                     self.rope.apply_positions(pool[..., -rd_:], pool_pos)],
+                    axis=-1,
+                )
                 ckv = pool
                 if hasattr(self, "indexer") and ckv.shape[1] > self.args.index_topk:
-                    topk_idx = self.indexer(x, qr)
-                    if topk_idx is not None:
+                    topk_idx = self.indexer(x, qr, offset=offset, rope=self.rope)
+                    if topk_idx is not None and S == 1:
+                        # Decode: a single query — gather its top-k rows so the
+                        # softmax is not diluted over the full pool (reference
+                        # restricts to index_topk at decode too).
+                        sel = topk_idx[:, 0, :]                    # [B, k]
                         idx = mx.broadcast_to(
-                            topk_idx[:, :, None],
-                            (B, topk_idx.shape[1], self.head_dim),
+                            sel[:, :, None], (B, sel.shape[1], ckv.shape[-1])
                         )
                         ckv = mx.take_along_axis(ckv, idx, axis=1)
+                    elif topk_idx is not None and S > 1:
+                        # Keep the full pool; restrict attention per query via
+                        # a membership mask (visibility AND top-k), built in
+                        # the mask branch below. Gathering one shared row set
+                        # for the whole prompt drops mid-prompt rows.
+                        member = mx.put_along_axis(
+                            mx.zeros(
+                                (B, S, ckv.shape[1]), dtype=mx.bool_
+                            ),
+                            topk_idx.astype(mx.int64),
+                            mx.ones(topk_idx.shape, dtype=mx.bool_),
+                            axis=-1,
+                        )
+                        pool_member_mask = member
                 compressed_k = ckv[:, None, :, :]
                 compressed_v = compressed_k
 
@@ -934,9 +990,36 @@
             v = mx.concatenate([compressed_v, v], axis=2)
             n_comp = compressed_k.shape[2]
             if mask is not None:
+                # Visibility mask for pool rows. Row j pools raw tokens
+                # [j*ratio, (j+1)*ratio - 1], so it is visible to a query at
+                # absolute position p iff (j+1)*ratio <= p + 1. A constant
+                # zeros() here is all-False for boolean masks and silently
+                # disables the entire global branch during prefill.
+                ratio = self.compress_ratio
+                q_pos = offset + mx.arange(S)                     # [S]
+                j_idx = mx.arange(n_comp)                          # [n_comp]
+                visible = (j_idx[None, :] + 1) * ratio <= q_pos[:, None] + 1
+                if pool_member_mask is not None:
+                    visible = mx.logical_and(
+                        visible[None, :, :], pool_member_mask
+                    )
+                # Broadcast to the mask's leading dims, then match dtype
+                # semantics: boolean -> keep; additive -> 0 / -inf.
                 comp_shape = list(mask.shape)
                 comp_shape[-1] = n_comp
-                comp_mask = mx.zeros(comp_shape, dtype=mask.dtype)
+                while visible.ndim > len(comp_shape) and visible.shape[0] == 1:
+                    visible = visible.squeeze(0)
+                while visible.ndim < len(comp_shape):
+                    visible = mx.expand_dims(visible, axis=-3)
+                visible = mx.broadcast_to(visible, comp_shape)
+                if mask.dtype == mx.bool_:
+                    comp_mask = visible
+                else:
+                    comp_mask = mx.where(
+                        visible,
+                        mx.zeros(comp_shape, dtype=mask.dtype),
+                        mx.full(comp_shape, float("-inf"), dtype=mask.dtype),
+                    )
                 mask = mx.concatenate([comp_mask, mask], axis=-1)
 
         out = scaled_dot_product_attention(
@@ -990,6 +1073,8 @@
         self,
         x: mx.array,
         q_intermediate: mx.array,
+        offset: int = 0,
+        rope=None,
     ) -> Optional[mx.array]:
         """Score compressed rows and return topk indices.
 
@@ -1012,16 +1097,46 @@
         q = q.reshape(B, S, self.n_heads, self.head_dim)
         q = q.transpose(0, 2, 1, 3)
 
-        scores = (q @ ck[:, None].transpose(0, 1, 3, 2)) * self.scale
-
-        hw = mx.sigmoid(self.weights_proj(x))
-        hw = hw.transpose(0, 2, 1)[..., None]
-        scores = scores * hw
+        ratio = self.compressor.ratio
+        if rope is not None:
+            # Reference binds the layer rope to the indexer and its compressor:
+            # queries rotate at their absolute positions, pooled rows at their
+            # window starts j*ratio. Rotary dims are the LAST rope.dims —
+            # __call__ rotates the first dims of its input, so pass the slice.
+            q = mx.concatenate(
+                [q[..., : -rope.dims], rope(q[..., -rope.dims :], offset=offset)],
+                axis=-1,
+            )
+            row_pos = mx.arange(n_compressed, dtype=mx.float32) * ratio
+            ck = mx.concatenate(
+                [ck[..., : -rope.dims],
+                 rope.apply_positions(ck[..., -rope.dims :], row_pos)],
+                axis=-1,
+            )
 
-        agg = scores.sum(axis=2).mean(axis=1)
+        # Reference scoring: raw ReLU(q . ck) weighted by weights_proj output
+        # WITHOUT sigmoid, with the softmax scale folded into the weights, and
+        # summed (not averaged) over heads. fp32 to keep the ranking stable.
+        scores = mx.maximum(
+            (q.astype(mx.float32) @ ck[:, None].transpose(0, 1, 3, 2).astype(mx.float32)),
+            0.0,
+        )
+        hw = self.weights_proj(x).astype(mx.float32) * (
+            self.scale * self.n_heads ** -0.5
+        )
+        hw = hw.transpose(0, 2, 1)[..., None]
+        agg = (scores * hw).sum(axis=1)                    # [B, S, n]
 
+        # Causality: row j is only scoreable by queries at absolute position
+        # p with (j+1)*ratio <= p+1; forbid the rest before ranking.
+        q_pos = offset + mx.arange(S)
+        vis = (mx.arange(n_compressed)[None, :] + 1) * ratio <= q_pos[:, None] + 1
+        agg = mx.where(
+            vis[None, :, :], agg, mx.full(agg.shape, float("-inf"), dtype=agg.dtype)
+        )
+
         topk = min(self.index_topk, n_compressed)
-        return mx.argpartition(-agg, kth=topk - 1, axis=-1)[:, :topk]
+        return mx.argpartition(-agg, kth=topk - 1, axis=-1)[..., :topk]
 
 
 # --------------------------------------------------------------------------- #

Co-authored with Claude (Anthropic).

[aidiffuser]

@stepnoy — excellent turnaround, and your patch matches our semantics on every point you listed. Here is our indexer scoring path verbatim, then the one structural difference that I think explains the remaining gap.

Our Indexer.__call__

def __call__(self, x, qr, cache, offset):
    B, S, _ = x.shape
    rd = self.rope_head_dim

    # NOTE: the indexer's own compressed buffer, cache-backed and keyed
    # separately from the attention pool (_K_IDX vs _K_COMP), same ratio.
    idx_kv = self.compressor(x, cache, offset, key=_K_IDX)
    if idx_kv is None or idx_kv.shape[1] == 0:
        return None

    q = self.wq_b(qr).reshape(B, S, self.n_heads, self.head_dim)
    q = mx.concatenate(
        [q[..., :-rd], self.rope(q[..., -rd:], offset=offset)], axis=-1
    )
    per_head_weights = self.weights_proj(x) * (
        self.softmax_scale * (self.n_heads ** -0.5)
    )                                              # no sigmoid; scale folded

    score = mx.einsum("bshd,btd->bsht", q.astype(idx_kv.dtype), idx_kv)
    score = mx.maximum(score, 0)                                    # ReLU
    score = (per_head_weights[:, :, None, :] @ score).squeeze(2)    # sum over heads

    # variable-length batch rows only; no causal mask here — causality is
    # applied later, AND-ed with this membership in the attention mask.
    pool_lengths = cache.pooled_lengths(_K_IDX)
    if pool_lengths is not None:
        lengths_a = mx.array(pool_lengths, dtype=mx.int32)
        valid = mx.arange(idx_kv.shape[1])[None, None, :] < lengths_a[:, None, None]
        score = mx.where(valid, score, mx.array(-1e30, dtype=score.dtype))

    k = min(self.index_topk, idx_kv.shape[1])
    return mx.argpartition(-score, kth=k - 1, axis=-1)[..., :k].astype(mx.int32)

Scoring-wise this is now identical to your patch (raw weights_proj, folded index_head_dim^-0.5 * n_heads^-0.5, ReLU on the logits, sum over heads, per-query top-k). Two notes on your open questions:

• Hadamard rotate_activation + fp4 activation-sim: we don't implement either, and we retrieve at 256k. So it isn't required for correct ranking — as you reasoned, the rotation cancels in the dot product, and skipping the quantization sim can only reduce noise.
• RoPE: our pooled rows are rotated inside the compressor as they are produced, at their absolute window start (mx.fast.rope(..., offset=pool_base // ratio, scale=ratio), which yields position (pool_base//ratio + i) * ratio), and the indexer's queries at their absolute positions with the same layer rope. Equivalent to your apply_positions(..., arange(n)*ratio) when the pool starts at 0.

The structural difference: our indexer's pool is cache-backed, so its index space is global

Head's indexer recomputes its pool from the current call's hidden states:

ck = self.compressor(x)          # stateless, prefill-only, keep = (S//r)*r

Ours reads an accumulated buffer (_K_IDX branch of the same cache that holds _K_COMP), carrying the sub-window tail and the previous window needed by the ratio-4 overlap transform across calls. That matters because the indices this function returns are used to select rows of the attention pool, which is global. The two index spaces only coincide when the indexer sees the entire sequence in one call.

So the discriminating check on your side is one line at the failing call:

print(ck.shape[1], compressed_k.shape[2])   # indexer rows vs attention-pool rows

If your 6133 run prefills in chunks (you mentioned experimenting with prefill_step_size 32/64), then for the final chunk ck holds only that chunk's rows — e.g. 509 for a 2037-token tail — while the attention pool holds 1533. Every index the indexer can emit is then < 509, and row 765 is unreachable by construction, in every layer, at every invocation: exactly your "0 of 42" with p ≈ 4e-8 against the noise hypothesis. It would also explain why ≤4101 survives if those runs happen to fit a single prefill call, and why the standalone reference retrieves the same needle.

If instead the two numbers are equal at the failing call, that theory is dead and the next things I would compare against ours are, in order: whether the ratio-4 overlap transform carries the previous window across call boundaries (head pads it with -inf/0 at every call start — for us that history comes from the cache's prev_kv/prev_gate), and whether tail tokens dropped by keep = (S//r)*r shift window alignment on subsequent calls (ours buffers the tail instead of dropping it).

Our full compressor and cache classes are a couple hundred lines; if the check above points at the index space, I can post the state-carrying parts.

Also, since it is cheap for us: we can run your patched head against a needle ladder on our own checkpoint (2× M3 Ultra, single node and 2-node TP) — that would confirm the repro and any fix are checkpoint-independent. Say the word and we will queue it.

Co-authored with Claude Fable 5 (Anthropic).

[stepnoy]

@aidiffuser — your theory was right, and the discriminating check settled it in one run. Thank you; that saved us from digging further into the scoring path, where the bug was not.

The check

Instrumented both numbers at every indexer invocation, 6042-token prompt, stock generate_step (prefill_step_size=2048):

S=2048 offset=2048  indexer_rows=512  pool_rows=1024
S=1945 offset=4096  indexer_rows=486  pool_rows=1510
S=1    offset=6041  indexer_rows=0    pool_rows=1510

Exactly as you predicted. On the final prefill chunk the indexer can only emit indices < 486, while the needle sits at pool row ~755 — unreachable in every layer, at every invocation. "0 of 42" was not a systematic down-ranking; it was arithmetic. And at decode Compressor.__call__ returns zero rows for S == 1, so head's indexer is inert during generation altogether.

That also explains the ≤4101 survivals: those runs put the needle in the last chunk.

Fix

Three changes, in the order they were needed:

1. Indexer pool moved into the cache, as a second slot alongside the attention pool (CompressedKVCache._slot("idx")). Its compressor has a different row width but the same ratio, so the row counts coincide and the two index spaces are identical by construction.
2. Accumulation made unconditional. My first attempt kept it inside the existing pool > index_topk guard, which is false on the first chunk — so the indexer pool ran a constant 512 rows behind the attention pool. Worth flagging because the needle was still retrieved in that state: a constant offset in index space still produces plausible output, and only the probe showed it. Any fix here should assert indexer_rows == pool_rows rather than trust the ladder.
3. Chunk tail and absolute position. keep = (S//r)*r dropped the tail of each chunk, and _abs_pos = S (rather than +=) reset the APE origin to the start of the last chunk. The tail is now buffered and replayed into the next chunk and into the decode rolling state.

Ladder after the fix — Flash, single M3 Ultra, greedy

prompt tokens	pool rows	prefill tok/s	decode tok/s	needle
1 503	375	447.3	29.1	✅
3 995	998	423.5	27.0	✅
6 042	1 510	413.3	26.6	✅ (was the ceiling)
11 827	2 956	392.4	25.8	✅
23 664	5 916	351.8	24.2	✅
47 427	11 856	288.8	23.1	✅
94 864	23 716	209.2	21.4	✅
189 827	47 456	129.9	18.5	✅

Ceiling not found — 190K is simply where I stopped. Long-form generation is clean too: worst of three runs on different prompts, 1% 4-gram repetition, no degeneration.

The fix is not configuration-specific — confirmed on V4-Pro-0813

You asked (reasonably) whether patches tuned on Flash would transfer. They do. I ran the same code, unmodified, on DeepSeek-V4-Pro-0813 — a rather different shape:

	Flash	Pro-0813
layers	43	61
routed experts	256	384
index_topk	512	1024
compress_ratios head	0, 0, 4, 128, …	128, 128, 4, 128, …
MXFP4 size	141 GiB	792 GiB

Note the third row in particular: Pro has no dense (ratio == 0) prefix layers, so it never takes that branch.

Probe across a full run: 205 indexer invocations, 0 mismatches between indexer_rows and pool_rows. Needle retrieved verbatim at 4 463 / 9 743 / 20 303 tokens (2×M3 Ultra, tensor-parallel). Before the fix this checkpoint produced garbage past a few thousand tokens.

Two things others may hit on Pro-0813

1. It ships no chat_template — not in the HF release, not in a converted build. It carries its own encoder instead, encoding/encoding_dsv4.py:

from encoding_dsv4 import encode_messages
prompt = encode_messages(msgs, thinking_mode="chat")   # "chat" | "thinking"
# <｜begin▁of▁sentence｜>…<｜User｜>…<｜Assistant｜>

Feeding it through apply_chat_template (or a default template) yields fluent nonsense that looks like a model defect. This cost me a day before I read the model's own README — worth a line in the PR docs.

2. Loading 792 GiB trips the Metal watchdog. sharded_load ends with a single mx.eval(model.parameters()); on a model this size that command buffer times out:

[METAL] Command buffer execution failed: Caused GPU Timeout Error
(00000002:kIOGPUCommandBufferCallbackErrorTimeout)

It reads like OOM but isn't — it dies at 398–409 GiB against a 448 GiB iogpu.wired_limit_mb, and Metal is not wedged afterwards. Materializing per layer fixes it:

model, config = load_model(path, lazy=True, strict=False)
model.shard(mx.distributed.init())
for layer in model.model.layers:
    mx.eval(layer.parameters())

With that, Pro loads in ~220 s at 404.5 GiB per rank. Might be worth doing inside sharded_load for large models generally.

Remaining divergence, and your offer

_overlap_transform folds the previous window into each row at prefill, but the decode rolling path never carries that history — so rows built during generation are computed differently from rows built during prefill. That is presumably what your prev_kv/prev_gate in the cache are for. It does not affect needle retrieval (the ladders above are clean), but it is the obvious suspect for long-generation drift, and I have not touched it. If the offer of the state-carrying parts of your compressor and cache still stands, that is the piece I would take.

And yes — please do queue the run against your own checkpoint, single node and 2-node TP. Independent confirmation on a checkpoint I cannot see would be worth more than anything further I can produce here.

Patch is against 63a26625 and includes the earlier six fixes. One unrelated nit: that commit annotates MTPBlock.cache with Any without importing it, so the module does not import as-is.

Co-authored with Claude (Anthropic).

[nh13]

Hit this PR debugging DeepSeek-V4-Flash locally and independently reproduced the same failures on 63a2662 (M3 Ultra, mlx-community/DeepSeek-V4-Flash-4bit): long-context collapse past ~200 tokens, i32 -> i drops from the compressed-KV decode cache, the indexer going inert at S==1, and the chunked-prefill remainder/_abs_pos reset. I landed on the same fixes @stepnoy and @aidiffuser describe (compress-RoPE base + block-position pool rows, indexer pool accumulated in-cache, per-window remainder/overlap carry, and #6's causal pool mask), so this is independent confirmation of the diagnosis. The debugging in this thread is excellent.

Two areas that look orthogonal to the generation work, in case they help:

• CompressedKVCache as a complete cache. Implemented state/meta_state/from_state so save_prompt_cache/load_prompt_cache round-trip the compressed and indexer pools (verified bit-exact vs a live cache, including a rotated window), and carried all fields through merge/filter/extend/extract so batched (continuous-batching) caches do not crash on the next accumulate. This also surfaced a general cache.py gap: load_prompt_cache resolves classes via globals(), so a cache class defined in a model file (like CompressedKVCache) raises KeyError. Fixable by saving module-qualified names with a bare-name fallback.
• convert.py mixed-quant. mixed_quant_predicate_builder derives one layer-index offset from the first down_proj, which for DeepSeek is an MTP block (mtp.0.block...), so model.layers.N paths then parse "layers" as an int and raise. Deriving the index per path fixes it (likely affects any MTP model, e.g. deepseek_v3).

Happy to open a PR (this branch, or main for the general two) with regression tests if useful. Not trying to step on the in-flight work, just offering the complementary pieces.