Fix DeepSeek V4/0731 target-path correctness and fail closed on unsupported state

[PhilipJohnBasile]

Summary

This is a focused correctness and safety patch on top of feat/deepseek-v4. It does not claim that ml-explore#1189's advertised DeepSeek-V4 feature set is complete.

Supported contract

The supported target path is deliberately narrow: one stream, initial prefill, then single-token cached decode.

• Fix Python <=3.13 loading by importing Any.
• Match the official per-layer RoPE regime, position compressed rows at their chunk starts, and make compressed-pool visibility causal for boolean and additive masks.
• Preserve official ratio-4 overlap and ratio-128 compressor state across initial prefill and cached decode.
• Keep all DeepSeek-V4 singleton cache layers in their serial representation; generic batching capability now fails closed before any scheduler or live-cache mutation.
• Preserve singleton lifecycle state across merge/filter/extract and sequential scheduler reuse.
• Detect the official 0731 three-stage DSpark layout, load the target model only, and drop incompatible mtp.* weights with a warning instead of instantiating conventional MTP.
• Keep custom Sinkhorn Metal paths lazy and opt-in until numerical parity is established.
• Reject speculative rollback at the prefill-time capability gate because compressed-cache rollback is not atomic yet.

Explicit fail-closed boundaries

• Generic fresh-cache and continuous multi-stream insertion reject before mutating active scheduler/cache state.
• Multi-stream compressed-cache merge and cache extension reject.
• Prompt-cache serialization and KV-cache quantization reject.
• Chunked-prefill continuation after compressor-state initialization rejects.
• External speculative decoding rejects until atomic compressed-cache rollback exists.

Verification

On the final frozen diff:

• Focused DeepSeek/batching selector: 23 passed, 1 skipped, 3 subtests passed.
• Prompt-cache suite: 22 passed.
• Deterministic exactness covers compressed-only, realistic mixed [0, 128, 4, 0], and all-uncompressed [0, 0, 0, 0] topologies: serial and BatchGenerator tokens, logprobs, cache types, and cache state match.
• Continuous insertion is rejected atomically; the active stream remains unchanged and a later sequential request succeeds.
• None, [], and [None] cache arguments, callable/nested cache capabilities, and ordinary batchable models have regression coverage.
• compileall and git diff --check pass.
• A fresh independent Codex Sol-max review of the complete final diff returned APPROVE after three earlier reviews rejected concrete cache-lifecycle defects that were then fixed and retested.

The untouched mixed_quant_predicate_builder path parser still fails at mlx_lm/convert.py:57; it was reproduced as a pre-existing base defect and is not presented as passing evidence here.

No real-checkpoint generation, long-context quality, or performance claim is made by this PR.

Known limitations

• DSpark inference/speculation is not implemented; 0731 is target-only here.
• Cached decode is supported only for the qualified single-stream path.
• The current top-k sparse-attention/indexer path is not source-faithful to the official persistent, per-query indexer.
• Generic fresh-cache and multi-stream batching are unsupported and fail closed.
• Prompt-cache serialization and KV-cache quantization are unsupported and fail closed.
• Chunked-prefill continuation after compressor-state initialization is unsupported and fails closed.
• External speculative decoding is unsupported and fails closed until atomic compressed-cache rollback exists.
• Real-checkpoint generation and long-context output quality have not been qualified.
• No performance claim is made.

Prior work and credit

This patch consolidates and extends findings already reported on ml-explore#1189:

• onchainengineer's PR #5: per-layer and compressed-pool RoPE.
• onchainengineer's PR #6: causal compressed-pool visibility.
• onchainengineer's PR #7: default-off Sinkhorn Metal path.
• onchainengineer's PR #8: fail loudly on incompatible compressed-cache batching.
• OmarB97's 0731 reference-parity findings.
• ncdrone's Python import report and real-checkpoint confirmation.

The additional work here covers official compressor rolling-state parity, vector cache offsets, serial-exact singleton lifecycle behavior across mixed cache topologies, generic non-mutating batching capability preflight, DSpark target-only detection, lazy opt-in behavior, and fail-closed speculative rollback.

[PhilipJohnBasile]

Current review handoff for frozen head b759c7fc9cb1c5dfb70f2e56da002bd718b6410d:

The PR is now ready for review and the deliberately narrow contract is unchanged: one stream, initial prefill, then single-token cached decode. Generic fresh-cache/continuous multi-stream insertion, prompt-cache serialization, KV-cache quantization, chunked-prefill continuation after compressor initialization, external speculation, DSpark, real-checkpoint qualification, long-context quality, and performance remain unsupported or unqualified.

The fail-closed batching guard remains non-mutating: capability is checked before a second stream is admitted or an active stream advances. Exact-head evidence remains attached:

• focused DeepSeek/batching: 23 passed, 1 skipped, 3 subtests passed
• prompt-cache: 22 passed
• compileall and git diff --check: passed
• independent complete-diff review: no blocker within the narrow contract

@Thump604 — a formal review request has now been submitted. Your prior COMMENTED review said you would be comfortable approving this exact frozen diff once the PR was ready and the evidence was attached; both conditions are now satisfied.

[Thump604]

Independent external review only: I am not a maintainer or collaborator on machiabeli/mlx-lm-1 or ml-explore/mlx-lm, so this should not be read as project approval or merge authority.

I reviewed frozen head b759c7f against base 63a2662. The revised scope now matches the implementation: singleton initial prefill plus single-token cached decode, with unsupported overlapping insertion, cache extension, serialization, quantization, chunked continuation, and speculative rollback failing closed. I specifically checked that the second-stream batching preflight occurs before scheduler queues or active prompt/generation caches are mutated, and that sequential reuse remains possible after the singleton completes.

Local independent verification on this head:

• all 18 newly added focused regressions passed, including 3 subtests
• prompt-cache suite: 22 passed
• git diff --check: clean

I found no blocking defect within that deliberately narrow contract. The broader DeepSeek V4 feature claims in upstream ml-explore#1189 remain unqualified and should not inherit this result. Because this PR is still a draft and has no reported GitHub checks, I am submitting a formal COMMENTED review rather than bypassing the draft gate with APPROVE. Once it is marked ready and its checks/evidence are attached, I would be comfortable approving this exact frozen diff.

[PhilipJohnBasile]

Exact-head verification receipt - 2026-08-09

Frozen candidate:

• head: b759c7fc9cb1c5dfb70f2e56da002bd718b6410d
• base: 63a26625c7ba2ffb8159ff430e630321446c7df4
• worktree: clean and equal to the PR head
• environment: macOS 27.0 (26A5388g), arm64, Python 3.13.14, MLX 0.32.0, pytest 9.1.1

Focused DeepSeek V4 and batching boundary suite:

PYTHONDONTWRITEBYTECODE=1 PYTHONPATH=tests /private/tmp/mlx-lm-pr10-venv/bin/python -B -m pytest -q -p no:cacheprovider tests/test_models.py -k "(deepseek_v4 or hc_sinkhorn_metal_is_opt_in) and not mixed_quant_preserves"
..............s.........                                              [100%]
23 passed, 1 skipped, 79 deselected, 3 subtests passed in 0.86s

Prompt-cache regression suite:

PYTHONDONTWRITEBYTECODE=1 PYTHONPATH=tests /private/tmp/mlx-lm-pr10-venv/bin/python -B -m pytest -q -p no:cacheprovider tests/test_prompt_cache.py
......................                                                   [100%]
22 passed in 0.91s

Static verification:

python -m compileall -q mlx_lm/generate.py mlx_lm/models/deepseek_v4.py mlx_lm/models/sinkhorn.py tests/test_models.py
git diff --check 63a26625c7ba2ffb8159ff430e630321446c7df4...b759c7fc9cb1c5dfb70f2e56da002bd718b6410d
PASS

No GitHub Actions jobs can run in this fork: both jobs in .github/workflows/pull_request.yml are explicitly gated by github.repository == "ml-explore/mlx-lm". This receipt does not present local validation as upstream CI.

Independent review:

• GitHub review by @Thump604 on this exact head found no blocker within the narrow, fail-closed contract and offered to submit APPROVE once the PR is ready and evidence is attached.
• A separate read-only Sol review of the complete frozen diff likewise found no code defect or unresolved thread blocking this PR.

The branch remains frozen. Any code commit or material base movement invalidates this receipt and requires rerunning the gate.

[PhilipJohnBasile]

Ready for review at frozen head b759c7f.

@Thump604, the exact-head verification receipt is attached above and the branch remains unchanged. When convenient, please submit the formal APPROVE you offered for this frozen diff.

@machiabeli, GitHub reports the PR cleanly mergeable into feat/deepseek-v4. There are no unresolved review threads or code blockers within the deliberately narrow single-stream contract. The fork workflow cannot emit Actions checks because its jobs are gated to ml-explore/mlx-lm; the complete local receipt and independent reviews are linked in this thread.