DeepSeek-V4-Flash: native deepseek_v4 AR backend

[davidtai]

Adds a native MLX backend for DeepSeek-V4-Flash (model_type: deepseek_v4), loading the published mlx-community/DeepSeek-V4-Flash-4bit / -2bit-DQ checkpoints directly through mtplx serve. No deepseek_v4 loader exists in mlx-lm; this is a from-scratch port of the reference deepseek-ai/DeepSeek-V4-Flash inference code.

What V4 adds over V3.2, and what this implements

• Hyper-Connections: the residual stream becomes 4 parallel copies mixed through a Sinkhorn-derived doubly-stochastic matrix around every block.
• Compressed Sparse Attention: layers carry a second, compressed KV cache built by learned gated pooling (overlapping ratio-4 windows with a top-512 indexer, or non-overlapping ratio-128 windows), alongside a 128-token sliding window. Both prefill and single-token streaming decode maintain the compressed lanes incrementally; the indexer's top-k filter engages in both paths once compressed positions exceed index_topk.
• Grouped output-LoRA on the attention output projection (8 groups, rank 1024).
• Hash-routed MoE for the first 3 layers (token-id expert lookup instead of scored top-k), sqrtsoftplus scoring elsewhere.

Attention itself is MQA-shaped MLA (64 query heads over one 512-dim latent with per-head sink logits). The decode-side compressor/indexer state machines are adapted from Salvatore Sanfilippo's ds4 engine (MIT), noted inline where copied; everything is expressed in stock MLX ops.

The runtime change is minimal: MTP-declaring checkpoints that ship no draft weights (the mlx-community conversions drop the MTP block) now degrade to AR with a clear message instead of failing at bind, gated on a shard-index probe that fails closed for every existing model.

Verification

• Reference parity: every new-math component is unit-gated against a transcription of the reference math, and the whole forward is gated layer-by-layer against a golden covering every layer type (final logits max_rel 1.8e-6, argmax 160/160 at the gated config).
• Decode == prefill: step decode reproduces the one-shot forward across partial windows, chunked prefill, window eviction, compress-window boundary crossings, and batching (max_rel ≤ 4e-6, argmax exact per step). The sparse-regime tests force n_comp > index_topk and cover the dense→sparse crossing mid-decode.
• Gates are mutation-tested: 18 deliberately injected state-machine and indexer bugs (off-by-one window keep, stale frontiers, dropped Hadamard, wrong tie-break, and so on) each fail the suites.
• Real weights: on the 2bit-DQ checkpoint (M5 Max), all 2,610 tensors bind with zero missing/extra; greedy generation is coherent at short and 3.3K-token context, and the streamed decode path matches the one-shot argmax 128/128 (dense regime) with the sparse regime's single divergence root-caused to a discrete top-k near-tie rather than a state bug. Repro scripts for these runs are included under scripts/.
• Suite: 39 new tests; the documented pre-PR selection plus every test file touching the modified modules passes (795 tests, 0 failures, 2 pre-existing environment skips).

Known limits (documented in-code)

• Throughput is unoptimized; the output-LoRA path currently re-dequantizes per token and dominates decode cost.
• Indexer/attention score tensors grow with context in one-shot prefill; the smoke harness documents a context ceiling pending chunked scoring.
• Compressed rows are retained rather than evicted, matching the reference (the filter bounds what is attended, not what is stored).
• MTP speculative decode for this model is in progress separately; this PR is the AR backend.

🤖 Generated with Claude Code

[davidtai]

Speculative decode on DeepSeek-V4-Flash: the MTP draft head, wired and measured

DeepSeek-V4-Flash ships one multi-token-prediction block upstream as mtp.0.*, but the published mlx-community conversions (2bit-DQ and 4bit) ship without it while leaving num_nextn_predict_layers: 1 in the config. This adds the block and the speculative lane that drives it. The block, DeepseekV4MTP, is a DeepseekV4DecoderLayer subclass, exactly as the reference's MTPBlock subclasses Block: it has its own attention, 256-expert MoE and two Hyper-Connection blocks, and shares only the embedding and lm_head with the trunk. There is no sidecar and no env var. The block lives at mtp.0.* in the module tree, so a checkpoint that ships those tensors binds through the ordinary sanitize -> quantize -> load_weights(strict=True) path. sanitize() lets the weights decide, so a conversion without the head still drops cleanly into the existing degrade-to-autoregressive branch.

The one real risk in speculating on this backend is the cache: a rejected draft that does not fully un-decode leaves every later token conditioning on state the committed prefix never produced. DeepseekV4Cache.trim() rewinds all three lanes, by three mechanisms picked for exactness rather than cost. Emitted compressed rows truncate, since a row is a pure function of one completed window. The compressor frontiers rebuild from a bounded journal of their own projected rows, because a rewind across an emission boundary needs rows the frontier already dropped. The sliding window retains rollback_capacity extra rows, because eviction is irreversible; retention and attendance are separate, so the deeper buffer cannot change the forward. Asked to exceed that bound, trim raises rather than clamping. Exact trim is what lets the engine's existing rejection repair, which corrects by trimming state rather than snapshotting it, serve this backend with no bespoke snapshot/restore path.

Benchmark

bench/deepseek-v4/mtpk-2bitdq-20260731: four arms off one load in one guarded window, with an unrecorded 8-token AR warmup ahead of the measured arms. Arms are only ever compared inside a window; this box carries roughly 15-20% run-to-run drift across windows, which swamps the depth-to-depth gaps here (K=2 to K=3 is 5%). B=1, greedy, 328-token code prompt, 256 decode tokens, capture_commit / stock verify core / committed history policy. Model: DeepSeek-V4-Flash-2bit-DQ trunk plus the mxfp4/bf16 MTP bank, merged; 93.4 GiB active after load. M5 Max, macOS 26.5.2, MLX 0.31.2. Arms run through mtplx.generation over a real MTPLXRuntime, so the receipt covers runtime.load()'s injector chain and the generation machine, not only the model.

arm	tok/s	ms/tok	x AR	committed tok / verify call	peak GiB	spec == AR
AR	4.534	220.6	1.00	n/a	95.17	—
K=1	8.042	124.3	1.77	2.08	95.18	PASS
K=2	9.865	101.4	2.18	2.69	95.18	PASS
K=3	10.356	96.6	2.28	2.94	95.18	PASS

Per-depth acceptance:

arm	d1	d2	d3
K=1	0.933 (112/120)	—	—
K=2	0.913 (84/92)	0.620 (57/92)	—
K=3	0.905 (76/84)	0.667 (56/84)	0.202 (17/84)

spec == AR is a hard gate in the harness: it sets exit status 1, because a K arm whose committed tokens differ from the AR arm's has a lossy rollback and its tok/s means nothing. All three arms emitted the identical 256-token id sequence as the AR control (spec_equals_ar.pass, first_divergence_index: null), and the decoded text matches character for character.

Depth beyond 3

This window ran K=1..3 only. A later sweep to K=5 (bench/deepseek-v4/goal-ab-20260731) shows the curve turning over after K=3 — in-window there, K=3 peaks at 25.9 tok/s, K=2 is level at 25.6, then K=4 falls to 22.6 and K=5 to 21.7 — and the counters say why: the draft chain runs to full depth unconditionally (drafted_by_depth is flat, [87,87,87,86,86] at K=5) while conditional acceptance collapses, with d5 accepting 0/86 and d4 3/86.

Two things about that sweep, both of which cut against leaning on it. It ran on a later branch that merges an unrelated decode-path change (cached o-LoRA dequant plus bf16 activation storage) into this lane; that code is not in this PR, and its absolute tok/s are not comparable to the arms above. More importantly, every K arm in that sweep failed spec == AR. bf16 activation storage removes the precision headroom that had been keeping batch-width-dependent rounding below the argmax. Every divergence was inspected and they are semantically equivalent code variants at near-ties rather than corruption, but the gate fails as written, and settling it needs a task eval rather than a byte comparison. So treat the depth curve there as a speed observation on different code. K>3 has never been run in this PR's configuration, and I am not claiming the turnover transfers — only that d3 acceptance is already down at 20% here, which leaves a fourth step little to recover.

Building the model directory

Speculation needs a model directory that carries the draft head, which no published conversion does. scripts/deepseek_v4_build_mtp_model.py builds one: it hardlinks every file of the MLX trunk snapshot (zero new bytes, the HF cache is never written to), translates mtp.0.* from the upstream deepseek-ai/DeepSeek-V4-Flash shard onto the MLX module tree as one extra shard, and rewrites only model.safetensors.index.json and config.json. It records source shard bytes and sha256 in an mtp_provenance block.

The default --bank exact repacks rather than re-quantizes. Routed experts stay FP4: MLX's mxfp4 layout is byte-identical to the source's, so the payload is viewed through rather than decoded and re-encoded, with mx.dequantize(mode="mxfp4") asserted equal to the reference LUT decode per tensor. The FP8 e4m3 x e8m0 dense projections are written as plain bf16, checked per tensor to hold every value exactly (it does on this checkpoint), with a float32 fallback and a printed warning if that ever fails. --bank affine-q8 reproduces the earlier lossy bank and is kept as an acceptance A/B arm. scripts/deepseek_v4_mtp_bind_check.py binds the draft head plus the embedding and lm_head it shares, so the load path can be gated without a full-model window.

Load the merged directory normally. is_deepseek_v4_mtp_config is the first branch of runtime.load()'s injector chain; the ordering is defensive rather than load-bearing, since the V3 predicate keys on model_type in {deepseek_v3, deepseek_v32, glm_moe_dsa} and cannot match a deepseek_v4 config. The injection only publishes the surface, since the block is already bound from the checkpoint's paths. An AR-only conversion detects as the same deepseek-v4 arch and keeps running target-only, unchanged.

Caveats

• Depth > 1 is an MTPLX extension. The reference defines exactly one MTP block and predicts one token ahead. K=2 and K=3 chain that single block against its own output, which the reference does not specify and upstream DeepSeek does not claim.
• One prompt, 256 tokens, B=1, greedy, one window. This measures speed and committed-sequence identity. No HumanEval or MBPP was run against this lane, so "the draft head does not change what the model writes" is established here as sequence identity on one completion.
• The spec tests run on a shrunk seeded model on CPU with an untrained draft head. They gate the engine wiring and the rollback arithmetic, which is what they are for; real-weight coverage of this lane is the single bench window above.
• spec == AR means the committed sequence is identical, which is the lane's documented invariant. Logits are not required to be bitwise equal: draft and verify are batch-shaped forwards, so the committed row's KV is projected inside a K+1-wide GEMM rather than alone, and that difference only stays below the argmax while there is precision headroom to absorb it.
• trim raises on one real serving path. The session bank's near-prefix restore (generation._trim_cache_to_offset) trims to an arbitrary matched prefix, which can exceed rollback_capacity (default 64). On this backend the rows are physically gone, so it raises instead of returning the False that would let the caller fall back to a cold prefill. Serving V4 behind a session bank needs either a rollback_capacity sized for it or a max_rollback pre-check in that caller. Recorded in the trim docstring, not fixed here.
• deepseek-v4 still carries support_level="experimental-native-ar-only". That field routes a checkpoint with no draft head to the AR-only verdict, which is the majority case as long as the published conversions ship without the block; an MTP-bearing artifact is resolved dynamically by the family gate instead. The entry's notes say so, but it is the one field in this PR that reads against the headline, so it is worth your call. deepseek-v4-mtp stays recognized-backend-pending because that entry describes vLLM's split checkpoint layout, and MTPLX has no loader that assembles a target from two repos.
• Reproducing the benchmark is expensive. It needs the upstream shard plus a local mlx-community trunk snapshot, and produces a ~96 GiB directory that peaks at 95.18 GiB against this box's 100 GiB wired limit. The bind-check script covers the draft head alone and runs anywhere.

Tests

tests/test_deepseek_v4_mtp.py (21) gates the block against a NumPy transcription of the reference MTPBlock at a shrunk config, with a 1e-5 relative bound and 2e-7 observed, plus four transcription mutations that must all move the logits past the bound. tests/test_deepseek_v4_spec.py (31) gates rollback exactness bit-for-bit across a ratio-4 emission boundary, the ratio-128 boundary and the indexer's dense-to-sparse threshold; spec == AR at K=1/2/3 through the real generation machine; all four verify strategies with the verify snapshot on and off; and four rollback defects, three of which run through both the cache-state comparison and 24 further decode steps, with a fourth (a window holding no rollback margin) that must refuse rather than half-rewind.

The open question I would put to you is whether K>1 chaining belongs in the default path or behind a flag, given the d2/d3 acceptance curve and that the reference only defines depth 1.

[davidtai]

Decode path: reference-faithful bf16 activation storage + o-LoRA weight modes

Two changes on top of the speculative lane, plus the gates and the receipts.

What it is

1. Activation dtype (60f41d5) — this is the speed. The backend stored three
activations in fp32 that the reference stores at the model dtype. Because
mx.concatenate and mx.matmul promote, one of them was enough: the promotion
starts at _apply_interleaved_rope on layer 0 and never unwinds, dragging the KV
cache, both attention matmuls, the o-LoRA einsum (which then upcasts wo_a — 128
MiB of fp32 weight per layer per token) and the entire residual stream up with it,
on every layer. bf16 is the reference's behaviour, not a quantisation of it:

• apply_rotary_emb rotates x.float() and copies the result back into the
  caller's own bf16 tensor — inference/model.py L234/L243
• the compressor pools in fp32 ("compression need fp32", L321-322) but casts the
  row back before the norm, kv = self.norm(kv.to(dtype)) L362 — and
  rotate_activation then asserts the row is bf16, L249
• sparse_attn is declared q/kv/o: BF16 with fp32 accumulator fragments, and
  casts the probability block to BF16 before the PV gemm — inference/kernel.py
  L295-297, L305, L340

Softmax stays fp32 (the reference keeps acc_s/scores_max/sum_exp in fp32
fragments); only storage moved. MTPLX_DSV4_FP32_ACTIVATIONS=1 restores the old
path as an A/B arm. Every cast is a no-op at fp32, which is where both parity
goldens and the decode oracle were captured — asserted with mx.array_equal, so no
golden and no tolerance moved.

2. o-LoRA weight modes (2dd7a12). wo_a is static ([8192, 4096]) and
_o_lora re-ran mx.dequantize on it every call — per token, per layer, 43 layers.
MTPLX_DSV4_O_LORA now selects: cached (default, dequantise once; bit-identical,
gated with mx.array_equal through both the trunk and the MTP draft block),
dequant (the old behaviour, kept as control and oracle), gather_qmm (the
quantised block-diagonal matmul the reference flags and declines at L538-539;
env-gated, not bit-identical). Cache is keyed on the identity of the quantised
tensors, so load_weights/update/set_dtype invalidate it, and it hangs off a
plain object so it never reaches the weight tree.

Honest attribution: the cached dequant is not where the speed came from.
Measured at fp32 storage it is 4.534 → 4.627 tok/s AR (+2.1%), inside this box's
15-20% cross-window drift — not distinguishable from zero, because at fp32 the
einsum promotes wo_a anyway. It costs +2.69 GiB resident. It ships because it is
bit-identical and removes real redundant work. bf16 storage carries the 3.5×.

Speed (bench/deepseek-v4/goal-ab-20260731)

Real 2bit-DQ checkpoint + mxfp4/bf16 MTP bank, B=1, greedy, 328-token code prompt,
256 decode tokens, capture_commit/stock/committed, 8-token unrecorded warmup.
All rows below are in-window (one guarded window, arms interleaved B→A→C so
drift pushes both controls the same way).

config	o_lora	acts	arm	tok/s	ms/tok	peak GiB
B (baseline)	dequant	fp32	AR	4.534	220.6	95.17
B	dequant	fp32	K=3	9.530	104.9	95.18
A (shipped default)	cached	bf16	AR	15.954	62.7	96.97
A	cached	bf16	K=1	22.634	44.2	96.97
A	cached	bf16	K=2	25.646	39.0	96.98
A	cached	bf16	K=3 (best)	25.856	38.7	96.98
A	cached	bf16	K=4	22.633	44.2	96.98
A	cached	bf16	K=5	21.704	46.1	96.98
C (env-gated)	gather_qmm	bf16	AR	16.146	61.9	94.31
C (env-gated)	gather_qmm	bf16	K=3	26.762	37.4	94.30

AR 4.534 → 15.954 = 3.52×; K=3 9.530 → 25.856. Stacked with speculative
decode, baseline to best shipped-default arm is 5.70×. Config B reproduces the
pre-change window exactly (same tok/s, same committed-text sha, same accept counts),
so the baseline is anchored rather than asserted. gather_qmm is a strictly better
speed and memory point (+1.2% AR / +3.5% K3, −2.66 GiB — it gives back the entire
cache cost) but its quality is unproven, so it stays off by default.

Quality (bench/deepseek-v4/quality-eval-20260731)

Paired, same prompts, same decode config, candidate (bf16) vs fp32-activation
control, EvalPlus, greedy:

arm	dataset	n	base	plus
control (fp32)	HumanEval	164	85.4%	82.3%
candidate (bf16)	HumanEval	164	86.0%	81.7%
candidate (bf16)	MBPP	378	84.1%	66.9%
control (fp32)	MBPP	378	pending	pending

Exact McNemar on the discordants:

• HumanEval base: 85.4% → 86.0% (+0.6 pts), b=7 c=8, p = 1.0000
• HumanEval plus: 82.3% → 81.7% (−0.6 pts), b=9 c=8, p = 1.0000

Discordants are symmetric in both directions (7↔8, 9↔8) — this is noise, not a
shift.

The near-tie disclosure

spec ≡ AR byte identity does not hold at bf16 defaults, and that is disclosed
rather than papered over. The mechanism is structural, not a rollback bug: draft
and verify are batch-shaped forwards, so the committed row's KV is projected inside
a K+1-wide GEMM rather than alone. The backend's documented invariant has always
been committed-sequence exactness, not bitwise-identical logits. At fp32 the
precision headroom absorbed that difference; at bf16 it reaches the argmax on
near-tied tokens.

Every divergence in the goal window was inspected. They are semantically equivalent
completions at near-ties — for name, deps vs for package_name, deps; two
equivalent ways to compute remaining; an inlined cycle walk vs a helper call —
and the same mechanism separates bf16-AR from fp32-AR, so it is not a spec-lane
regression. Several distinct arms land on the same alternative text. The fp32
arms (config D) still pass byte-identity at K=2 and K=3, exactly as before.

16f604b reframes the harness accordingly: the comparison always runs and is always
recorded, only the exit status moves. Byte identity stays a hard gate under
MTPLX_DSV4_FP32_ACTIVATIONS=1 (the diagnostic lane where it does hold) or
--require-exact; at bf16 defaults divergence is reported as data — count, compared
length, first index, both tokens — in the log, the summary column and the JSON. The
receipt records which lane ran and whether identity was enforced, so a status-0
receipt can never be read as "spec==AR held" when it was never asked to.

The in-repo spec gates are unaffected and that is now asserted rather than assumed:
their shrunk model never calls set_dtype, so it is all-fp32 and every cast is a
no-op there. 31/31 stay green.

Acceptance criterion

Per the maintainer's ruling for this lane: byte-exactness is not the bar; a
measured-minimal quality delta is. On that bar this clears — McNemar p=1.0 on both
HumanEval metrics with symmetric discordants, for 3.52× AR and 2.7× at K=3 — with the
fp32 escape hatch retained for anyone who needs byte identity as a diagnostic.

Gates

139 deepseek_v4 tests green (was 101): decode 7, dtypes 13, indexer 14, loader 3,
mtp 21, new_math 7, o_lora 19, parity 2, spec 37, swiglu_clamp 16. Upstream-touched
coverage 189/189. scripts/hygiene_scan.sh exit 0.

The MTP draft block reuses DeepseekV4Attention, so it inherits both changes by
construction — and is gated for it directly (656bc96) rather than by inheritance
argument: five attention instances each hold their own dequantised wo_a (not a
neighbour's, and no two equal, so the check can't pass vacuously), cached == dequant
is bit-for-bit through Model.mtp_forward including the cache-hit call, and the draft
path holds bf16 storage plus the fp32 escape hatch. A module-level singleton cache, a
rounded cache copy, and a fp32-pinned _store_dtype each fail these gates.

[davidtai]

Quality follow-up: MBPP paired result (completes the bf16 verdict)

The paired MBPP control arm from the previous comment. Same setup — bf16 candidate vs the fp32-activation control, EvalPlus, greedy, same prompts and decode config.

arm	HumanEval+ base/plus (n=164)	MBPP+ base/plus
control (fp32)	85.4 / 82.3	88.5 / 73.2
candidate (bf16)	86.0 / 81.7	88.0 / 71.6

Exact McNemar on the discordants (b = control-pass/candidate-fail, c = the reverse):

• HumanEval base: +0.6 pts, b=7 c=8, p = 1.0000
• HumanEval plus: −0.6 pts, b=9 c=8, p = 1.0000
• MBPP base: −0.5 pts, b=5 c=4, p = 1.0000
• MBPP plus: −1.6 pts, b=9 c=6, p = 0.6072

Every cell's discordants are near-symmetric — bf16 reshuffles a handful of near-ties in both directions with no net damage, which is the signature the byte-gate near-ties predicted and the opposite of a regression. Worst case is −1.6 points against a >3-point failure bar; nothing significant, nothing cratered. On the "minimal quality delta" bar, bf16 passes on both benchmarks.

One honest bound: the MBPP control arm is scored on the paired first 183 of 378 problems — it was stopped early to return a shared GPU to production serving, and the pairing is exact (candidate scored on the same 183 for this comparison; its full-378 absolutes were 84.1 / 66.9). HumanEval is final at 164. The MBPP discordant balance (b/c within one of each other on both metrics) is not close to the threshold where the remaining problems could flip the verdict, but I'll post the full-378 McNemar when the control arm resumes on an idle window, for completeness.

Receipts: bench/deepseek-v4/quality-eval-20260731-VERDICT.txt (+ -ANALYSIS.json, per-arm sample/eval trees). gather_qmm (the env-gated MTPLX_DSV4_O_LORA=gather_qmm mode, +3.5% K3 / −2.66 GiB) was not quality-tested here and stays off by default; its task-eval arm rides a later window.

[davidtai]

Decode dispatch reduction: fused CSA attention + Hyper-Connection tape collapse

Two structural levers on the DeepSeek-V4 decode path, both env-gated with the fast
path on by default, both measured on the real 2-bit-DQ checkpoint.

(A) Fused CSA attention — attn_sink becomes one all-zero-V extra KV column so
the per-head sink logit rides an ordinary softmax denominator, and the whole block
is a single mx.softmax(precise=True) (fp32 max + fp32 accumulate, bf16 in/out).
This drops the materialized fp32 score block the dense path wrote twice over.
MTPLX_DSV4_ATTN ∈ {dense,fused,sdpa}, default fused. The sdpa arm is kept and
gated exact but is not default: MLX only instantiates fused SDPA kernels for
head dims 64/96/128/(192)/256, and DeepSeek-V4's MLA latent is 512 wide, so on
every MLX on this box sdpa takes the unfused fallback — fused is that fallback
minus two copies.

(B) Hyper-Connection / Sinkhorn tape collapse — derive the fp32 fn/base/scale
once (not per call), fold three affines into one, and run the whole pre as a
module-level pure function so mx.compile holds one tape shared across all 87
HC modules. MTPLX_DSV4_HC_COMPILE, default 1, with _HC_COMPILE_MAX_ROWS=32
(bounds MLX's tape-cache scan and the >32-row precision drift).

Why dispatch count ≈ wall clock here

mlx-profiler census, per bf16 decode step:

	dispatches	command buffers
before	19,809	384
after (default)	14,639	288
	−26.1%	−25.0%

The load-bearing finding: host encode is exposed — 56–59 ms of encode against
32–35 ms of GPU execution, ~2.9 µs/dispatch, size-independent. So removing
~5,200 dispatches removes host encode roughly one-for-one.

Measured window (bench/deepseek-v4/kernel-a-ab-20260801)

B=1 greedy, 328-tok prompt, 256 decode, drift-bracketed, MTPLX_DSV4_O_LORA=cached
held fixed:

arm	AR tok/s	K=3 spec tok/s
before (dense, hc0)	17.37	26.31
after (fused, hc1)	22.80 (+31.3%)	30.88 (+17.4%)

+5.43 AR tok/s = 13.7 ms/token removed, 91% of the ~15 ms the census projected.
The mlx-0.32 arm is a clean null — once host encode is gone, decode is GPU-forward
bound, not qmm-ALU bound. No regression; all arms' decode text stayed coherent.

Scope (honest)

This trims host encode only. The remaining gap to the ≥40 goal (K=3 30.9 → 40 =
1.29×) is GPU-forward / verify-width, not dispatch — a different lever.

Two deferrals worth flagging, neither done here:

• Sinkhorn is activation-dependent, so comb cannot be precomputed at load
  (mixes = F.linear(x, hc_fn)·rsqrt from the layer's own hidden state; that's
  also why hc_fn ships [24, hc·dim]).
• The Sinkhorn reduce_sums are the floor: 3,678 of the remaining 14,639
  dispatches (25%). mx.compile doesn't fuse reductions; only a hand
  mx.fast.metal_kernel doing the whole 4×4 chain would take those 3,678 → 87.

Gates

tests/test_deepseek_v4_*.py: 163 passed (139 pre-existing + 24 new in
test_deepseek_v4_kernel_paths.py), green in all six arms
{fused,dense,sdpa} × hc_compile{0,1}. Coverage includes fp32 parity vs the
dense oracle (one-shot + streaming, dense + sparse regimes, and the
MTPLX_DSV4_FP32_ACTIVATIONS=1 escape hatch), compiled==eager bit-identity at
decode shape, and a weight-rebind mutation guard on the static HC cache. Census is
reproducible via scripts/deepseek_v4_dispatch_census.py.

[davidtai]

Stage 4 — shape-fixed Sinkhorn kernel: +29.3% AR, 32.5 tok/s K3

This stage targets DeepSeek-V4's actual Hyper-Connection shape rather than a generic reduction: hc=4, a [4,4] fp32 comb, exactly 20 alternating normalization iterations, and 86 pre calls per decode step. The stock graph launches 39 reductions, 39 divides, and one softmax per call. The Metal lane keeps the same operation order in registers and emits one kernel per call.

The port is opt-in (MTPLX_DSV4_SINKHORN_KERNEL=1) and remains off by default. Route selection is now installed once when HyperConnection is constructed: CPU/no-Metal explicitly installs the stock callable; the GPU lane accepts only hc=4, iters=20, eps=1e-6 and fails before generation otherwise. The enabled token path has no environment/device/eligibility checks and no silent fallback. Fused attention remains the default.

Real-weight end-to-end window

One guarded load, B=1 greedy, 328-token prompt + 256 generated tokens, 8-token warmup, capture_commit / stock verify / committed history. Model: DeepSeek-V4-Flash 2-bit-DQ trunk plus the lossless mxfp4/bf16 MTP bank. Held fixed: cached o-LoRA, HC compile on, fused attention. M5 Max, macOS 26.5.2, MLX 0.31.2.

arm	AR tok/s	K3 tok/s	K3 committed/verify	peak GiB
control: fused + Sinkhorn off	22.318	30.440	2.844	96.89
Sinkhorn kernel on	28.858	32.497	2.723	97.12
control repeat / drift bracket	22.881	30.903	2.844	96.89

• AR: 22.318 → 28.858 = +29.3% against the first paired control, or +27.7% against the two-control mean. Pure AR now clears the DS4 27 tok/s comparison point.
• K3: 30.440 → 32.497 = +6.8% in this window; 32.497 tok/s is the best K3 result in the matrix.
• All outputs remained coherent (unique-line ratio 1.00, longest identical-line run 1). The kernel arm's spec != AR is the already documented bf16 near-tie diagnostic, not corruption; byte identity is not being claimed.

MLX-profiler attribution

This is structural evidence, not the end-to-end speed measurement above. In the profiler census, the Sinkhorn floor collapses 6,794 stock dispatches to 86 kernel dispatches; whole-step dispatches move 14,639 → 7,845, command buffers 288 → 155. The profiler build suggested 42 → 22 ms host encode, but only the guarded real-weight table above is used for the tok/s claim.

The benchmark ran on combined revision 3f0aa06; its Sinkhorn tree is identical to source revision 57f54dd, with the device guard from 8067f9b. Stage-4 ports that implementation and adds construction-time routing. The scrubbed immutable receipt records the command shape, arm data, provenance, drift, coherence, and SHA-256 manifest: DeepSeek-V4 Sinkhorn stage-4 receipt.

Gates

• 163 DeepSeek-V4 tests passed with the flag unset on explicit CPU.
• 163 passed with MTPLX_DSV4_SINKHORN_KERNEL=1 on explicit CPU, proving the stock CPU route remains safe.
• 18 Metal tests collect, including GPU route engagement, extreme positive/negative fp32 combs, and one-ULP near-equal maxima at the existing <=1e-6 plus argmax-exact parity bar.
• Diff checks, compileall, and Ruff passed. No dependency or lockfile change.

Stage-4 head: 6977736.

[davidtai]

Shape-specific MLA disposition: rejected; no MLA code added to PR.

This was a clean, immutable local capture on one model load: the canonical raw
328-token Stage-4 prompt plus 256 decoded tokens, B=1 greedy, AR and K3,
capture_commit / stock / committed, with 8 warmup tokens. The profiler was
off during TPS measurement. Two fused Sinkhorn-off controls bracketed the MLA
candidates; all 14 measured arms completed 256 tokens and passed the
coherence, memory, and route-engagement gates. This is n=1 per arm/mode, so
there is no per-arm median; the control-bracket drift was 0.213% AR and 0.160%
K3.

The exact target workload is MLA at B1/H64/KV1/D512, bf16 storage and fp32
accumulation. A complete MLX-profiler geometry sweep chose HG8/TK8. That shape
reduces 64 KV reads to 8; it does not make a false 64-to-1 claim.

arm	AR tok/s	K3 tok/s
fused_sink0_control_a	22.9152	30.9666
generic_mla_sink0	22.4299	30.9738
shaped_mla_sink0	21.7033	30.9516
fused_sink1	28.9257	32.6521
generic_mla_sink1	28.5335	31.8245
shaped_mla_sink1	27.4052	31.7870
fused_sink0_control_b	22.8665	31.0161

Against the fused Sinkhorn-off control-bracket mean, shaped MLA + sink0 is
-5.188% AR / -0.128% K3, with AR parity only 11/256. Against fused + Sinkhorn,
shaped MLA + sink1 is -5.257% AR / -2.649% K3, with parity 221/256 AR and
81/256 K3. Both shaped and generic MLA are rejected and will not be ported.

The positive result is that fused + Sinkhorn remains the winner at 28.9257 AR /
32.6521 K3, independently corroborating the shipped Stage-4 receipt/comment
(whose official prior-best K3 was 32.497). This is not a new stage or success
claim.

For audit, local immutable receipt identifiers are: combined receipt SHA256
cf83d44ab958788c7e7bb591430fc7236eb5e4dc8c66e217607910da1ad4fcac; geometry
receipt SHA256
259315bd2bfa85627dc0e92eb7390deb2ca424f160f1ba106f7632e904d7cd35; capture-code
HEAD 206bb9799bf7a92fd55dbd97946211aa55ff7540; canonical prompt SHA
ee94397faa812c91d5f1a0ee17c5bb6ca6032883653591dd33d4cfddb737ac33. They are
local immutable receipt identifiers if not reachable upstream.

One failed first attempt is disclosed: it was a pre-model prompt-contract
rejection, produced zero benchmark rows, was fixed by restoring the exact
Stage-4 prompt, and is excluded from the performance data.

PR disposition unchanged: only Sinkhorn Stage 4 is proposed; MLA remains out.