[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)

[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.

[stepnoy]

Thanks for the writeup. Yes, please open the PR — both pieces are real on our side too.

I re-checked them today against our patched branch (M3 Ultra 512 GB, mlx-lm 0.31.3 / mlx 0.32.0). No weights needed; a 64-token prefill into a CompressedKVCache with both pools populated is enough to show it.

• merge([a, b]) carries the attention pool over as (2, 16, 16), but the indexer slot is gone entirely and _abs_pos isn't set as an attribute at all. The next accumulate() on the merged cache dies with AttributeError: 'CompressedKVCache' object has no attribute '_abs_pos'. extract(0) drops the same two.
• state returns exactly the two arrays self.local returns, so no pools. meta_state comes back as ('0', '128', '64', '64') — the RotatingKVCache meta, nothing about _abs_pos. save_prompt_cache doesn't complain about any of this: it writes the incomplete snapshot, and you only find out at load time, as KeyError: 'CompressedKVCache'. That's your globals() point in cache.py.

On mixed-quant we hit it from the other direction. named_modules() hands back mtp.0.block.ffn.shared_experts.down_proj first, so layer_location comes out as 1, and then every model.layers.N path raises ValueError: invalid literal for int() with base 10: 'layers' — 89 of them on a small config. Same root cause: the offset is derived once and applied to every path. We convert with our own script for unrelated reasons (strict=True chokes on the MTP heads), so it doesn't block us, but it explains another failure we'd written off.

If it helps, we can run your branch on our two 512 GB M3 Ultras and send back traces or numbers. I can also share the repro script — it builds a tiny Compressor and exercises merge/extract/state without loading a checkpoint, runs in about a second, so it would drop straight into a regression test.

[stepnoy]

Here are the two repro scripts, in case they save you the setup. Neither loads a checkpoint — the first builds a tiny Compressor and drives the pools by hand, the second builds a small deepseek_v4 config. Both finish in about a second on an M3 Ultra, against our patched branch with mlx-lm 0.31.3 / mlx 0.32.0.

repro_cache_state.py — merge / extract / state / prompt-cache round trip (currently 8 failures)

"""Repro for the CompressedKVCache gaps discussed in ml-explore/mlx-lm#1189.

Loads no weights: builds a tiny Compressor and drives the pools by hand, so it
runs in about a second and is cheap enough for a regression test.

    python repro_cache_state.py
"""
import os
import tempfile

import mlx.core as mx
from mlx_lm.models.cache import load_prompt_cache, save_prompt_cache
from mlx_lm.models.deepseek_v4 import CompressedKVCache, Compressor, ModelArgs

FAILED = []


def check(name, ok, detail=""):
    print(f"  {'OK  ' if ok else 'FAIL'}  {name}" + (f"   {detail}" if detail else ""))
    if not ok:
        FAILED.append(name)


args = ModelArgs(
    hidden_size=32, num_hidden_layers=6, head_dim=16, index_head_dim=8,
    qk_rope_head_dim=4, vocab_size=64,
)
RATIO = 4
comp_main = Compressor(args, compress_ratio=RATIO, head_dim=16)  # attention pool
comp_idx = Compressor(args, compress_ratio=RATIO, head_dim=8)    # indexer pool


def fresh(n_tokens):
    """A cache after an n_tokens prefill, with both pools and the local window filled."""
    c = CompressedKVCache(max_size=128)
    x = mx.random.normal((1, n_tokens, args.hidden_size)).astype(mx.bfloat16)
    c.accumulate(x, comp_main, slot="main")
    c.accumulate(x, comp_idx, slot="idx")
    k = mx.random.normal((1, 1, n_tokens, args.head_dim)).astype(mx.bfloat16)
    v = mx.random.normal((1, 1, n_tokens, args.head_dim)).astype(mx.bfloat16)
    c.update_and_fetch(k, v)
    mx.eval(c.pool, c._alt_slots["idx"]._pool, c.local.keys)
    return c


def rows(slot):
    return 0 if slot._pool is None else slot._pool.shape[1]


print("=== 0. after prefill the two pools agree")
c0 = fresh(64)
check("main rows == idx rows", rows(c0) == rows(c0._alt_slots["idx"]),
      f"main={rows(c0)} idx={rows(c0._alt_slots['idx'])} _abs_pos={c0._abs_pos}")

print("\n=== 1. merge: what survives into the batched cache")
a, b = fresh(64), fresh(64)
idx_rows_before = rows(a._alt_slots["idx"])
m = CompressedKVCache.merge([a, b])
check("attention pool carried over", m._pool is not None and m._pool.shape[0] == 2,
      f"shape={None if m._pool is None else m._pool.shape}")
has_alt = getattr(m, "_alt_slots", None) is not None
check("indexer pool carried over", has_alt,
      f"{idx_rows_before} rows before -> _alt_slots {'present' if has_alt else 'MISSING'}")
check("_abs_pos carried over", getattr(m, "_abs_pos", None) == a._abs_pos,
      f"{a._abs_pos} before -> {getattr(m, '_abs_pos', '<no attribute>')}")

print("\n=== 2. next accumulate on the merged cache")
x = mx.random.normal((2, RATIO, args.hidden_size)).astype(mx.bfloat16)
try:
    m.accumulate(x, comp_main, slot="main")
    m.accumulate(x, comp_idx, slot="idx")
    mx.eval(m.pool, m._alt_slots["idx"]._pool)
    check("indexer_rows == pool_rows", rows(m) == rows(m._alt_slots["idx"]),
          f"pool_rows={rows(m)} indexer_rows={rows(m._alt_slots['idx'])}")
except Exception as e:
    check("accumulate does not crash", False, f"{type(e).__name__}: {e}")

print("\n=== 3. extract: back out of the batch")
c = fresh(64)
one = c.extract(0)
check("indexer pool carried over", getattr(one, "_alt_slots", None) is not None)
check("_abs_pos carried over", getattr(one, "_abs_pos", None) == c._abs_pos,
      f"{c._abs_pos} before -> {getattr(one, '_abs_pos', '<no attribute>')}")

print("\n=== 4. state / meta_state: what a prompt-cache snapshot would hold")
c = fresh(64)
st, local_only = c.state, c.local.state
n_state = len(st) if isinstance(st, (list, tuple)) else 1
n_local = len(local_only) if isinstance(local_only, (list, tuple)) else 1
check("state holds more than the local window", n_state != n_local,
      f"state yields {n_state} array(s), exactly as many as self.local")
check("meta_state knows about _abs_pos", "abs" in str(c.meta_state).lower(),
      f"meta_state={c.meta_state}")

print("\n=== 5. prompt-cache round trip")
path = os.path.join(tempfile.mkdtemp(), "cache.safetensors")
try:
    save_prompt_cache(path, [c])
    print("  save_prompt_cache: went through")
    try:
        load_prompt_cache(path)
        check("load_prompt_cache", True)
    except Exception as e:
        check("load_prompt_cache", False, f"{type(e).__name__}: {e}")
except Exception as e:
    check("save_prompt_cache", False, f"{type(e).__name__}: {e}")

print("\n" + ("ALL CLEAN" if not FAILED else f"{len(FAILED)} FAILED: " + ", ".join(FAILED)))

Output here:

=== 0. after prefill the two pools agree
  OK    main rows == idx rows   main=16 idx=16 _abs_pos=64

=== 1. merge: what survives into the batched cache
  OK    attention pool carried over   shape=(2, 16, 16)
  FAIL  indexer pool carried over   16 rows before -> _alt_slots MISSING
  FAIL  _abs_pos carried over   64 before -> <no attribute>

=== 2. next accumulate on the merged cache
  FAIL  accumulate does not crash   AttributeError: 'CompressedKVCache' object has no attribute '_abs_pos'

=== 3. extract: back out of the batch
  FAIL  indexer pool carried over
  FAIL  _abs_pos carried over   64 before -> <no attribute>

=== 4. state / meta_state: what a prompt-cache snapshot would hold
  FAIL  state holds more than the local window   state yields 2 array(s), exactly as many as self.local
  FAIL  meta_state knows about _abs_pos   meta_state=('0', '128', '64', '64')

=== 5. prompt-cache round trip
  save_prompt_cache: went through
  FAIL  load_prompt_cache   KeyError: 'CompressedKVCache'

repro_mixed_quant_mtp.py — the layer-index offset vs. MTP blocks

"""Repro for mixed_quant_predicate_builder vs. MTP blocks (ml-explore/mlx-lm#1189).

The layer-index offset is derived once, from the first down_proj in the model,
and then applied to every path.

    python repro_mixed_quant_mtp.py
"""
import mlx.nn as nn
from mlx_lm.convert import mixed_quant_predicate_builder
from mlx_lm.models.deepseek_v4 import Model, ModelArgs

args = ModelArgs(
    hidden_size=64, num_hidden_layers=8, head_dim=16, index_head_dim=8,
    qk_rope_head_dim=4, vocab_size=128, n_routed_experts=4, num_experts_per_tok=2,
    moe_intermediate_size=32, index_n_heads=2, num_attention_heads=4,
    q_lora_rank=16, o_lora_rank=16, o_groups=1, num_nextn_predict_layers=1,
)
model = Model(args)

down_keys = [k for k, _ in model.named_modules() if "down_proj" in k]
print("first path with down_proj:", down_keys[0])
for layer_location, part in enumerate(down_keys[0].split(".")):
    if part.isdigit():
        break
print("layer_location derived from it:", layer_location)
print("paths with down_proj:", len(down_keys),
      "of which MTP:", len([k for k in down_keys if k.startswith("mtp")]))

pred = mixed_quant_predicate_builder("mixed_3_6", model)
bad = []
for path, module in model.named_modules():
    if not isinstance(module, (nn.Linear, nn.Embedding)):
        continue
    try:
        pred(path, module)
    except Exception as e:
        bad.append((path, f"{type(e).__name__}: {e}"))

print(f"\npaths the predicate raises on: {len(bad)}")
for path, err in bad[:5]:
    print(f"  {path}\n     {err}")

Output here:

first path with down_proj: mtp.0.block.ffn.shared_experts.down_proj
layer_location derived from it: 1
paths with down_proj: 18 of which MTP: 2

paths the predicate raises on: 89
  model.layers.7.ffn.shared_experts.down_proj
     ValueError: invalid literal for int() with base 10: 'layers'

The pool field names are ours (_alt_slots, _abs_pos), so the first script will need renaming against your tree, but the shape of the checks should carry over.

[avlp12]

Found three deviations from the reference implementation while testing this
branch on long inputs. Sharing with fixes for the first two, since together
they explain the degraded generation some users may be seeing on longer
prompts.

1. rope/YaRN assignment. The reference Attention.__init__ builds one
rope per layer:

if self.compress_ratio:
    original_seq_len, rope_theta = args.original_seq_len, args.compress_rope_theta
else:
    # disable YaRN and use base rope_theta in pure sliding-window attention
    original_seq_len, rope_theta = 0, args.rope_theta

i.e. compressed layers (41 of 43) rotate main Q/K with compress_rope_theta
(160000) + YaRN, and the ratio-0 sliding-window layers (0, 1, 42) use base
theta with YaRN disabled. This branch applies rope_theta (10000) + YaRN to
main Q/K on every layer — the self.compress_rope instance is created but
never called. Empirically, on my M3 Ultra with the 4-bit community quant,
real-document prompts are fine at 1.5K, marginal at 2.2K, and fully broken at
4.9K ("? 2:2:2:..." style output), while the same prompts through llama.cpp's
deepseek4 arch stay coherent — which is what localized it to this branch.
(Other deviations below plausibly contribute past ~2K as well, e.g.
index_topk=512 × ratio 4 puts the indexer's activation threshold right at
2048, so I'm attributing by differential measurement rather than claiming a
single cause.)

Fix (one conditional):

self.rope = DeepseekV4RoPE(
    self.rope_head_dim,
    args.compress_rope_theta if self.compress_ratio else args.rope_theta,
    args.rope_scaling if self.compress_ratio else None,
)

After this, greedy generation over 4.9K- and 19K-token documents is coherent
and on-topic (occasional CJK token slips remain at 19K — see item 3).

2. Compressed-pool visibility during prefill. The prefill path prepends
the compressed pool with an all-zeros additive mask, so every query in a chunk
can attend to pool rows that summarize positions after it. All three
implementations I checked (the reference model.py, HF transformers'
modeling_deepseek_v4, and FreeToken's port) clamp per query: row i is
visible to position p iff i < (p + 1) // ratio (the pool row's last source
token must be ≤ p). Besides the mid-chunk state corruption this causes, it
also makes prefill and decode inconsistent (the decode path only ever sees
completed pool rows, so teacher-forced evals of the prefill path read
optimistic). I verified with a fused kernel that implements the reference
visibility — outputs become clean once both (1) and (2) are fixed.

3. Pool rows are never rotated. The reference applies compress-theta rope
to each pool row at its block-start position before it enters the K/V stream
(freqs_cis[:cutoff:ratio] in the reference Compressor; same in transformers
and FreeToken), so query·pool-row dot products encode relative distance. This
branch's Compressor returns the normed rows without any rope — which looks
like what self.compress_rope was created for. I haven't patched this one
yet; it plausibly accounts for the residual long-range slips I still see at
19K after fixing (1) and (2).

There are a few further indexer-path deviations (per-query masked top-k vs a
single shared top-k over the mean, relu+linear weighting vs sigmoid, and the
S=1 decode path bypassing top-k entirely so the pool is fully attended past
index_topk × ratio tokens) — happy to detail those, share benchmarks, or
test other configurations.

[avlp12]

Following up on item 3 (pool rows never rotated) — I've now patched and
validated it, and it was indeed the cause of the residual long-range slips.

The fix, mirroring the reference Compressor: after the RMSNorm, rotate the
last rope_head_dim dims of each pool row with the compress-theta + YaRN
rope at the row's block-start position (i * ratio). Since the positions are
strided rather than contiguous, I computed the angles directly from the
rope's inv_freq (same interleaved pairing as the main Q/K rope):

pos = (base_row + mx.arange(P)) * ratio          # block-start positions
theta = pos[:, None] * rope.inv_freq[None, :]
# rotate ckv[..., -rope_head_dim:] with cos/sin of theta (interleaved pairs)

(Equivalently, a rope instance with its frequencies pre-scaled by ratio can
be applied with a contiguous row offset — same math.) Using the pool-row
count as base_row also keeps positions correct across chunked prefill. I
additionally replaced the all-zeros pool mask on the non-kernel prefill path
with the per-query clamp from item 2.

Result on the same 19K-token document task where I previously saw occasional
CJK token slips after fixing items 1–2: zero CJK slips over 250 generated
tokens, coherent and on-topic (4.9K unchanged and clean; single-box prefill
throughput unaffected, 422–473 tok/s on my M3 Ultra). So the three items
together fully account for the degradation I could measure on this branch.

[nastya236]

Thanks everyone for the work here.
However, this change is too large for a single pull request. I'd propose splitting it into multiple reviewable changes.
We've also introduced a policy on AI use in PR descriptions and discussions (#1770), which this pull request doesn't follow. You're very welcome to open a new PR that does.