Add DeepSeek-v4 (Flash/Pro)

[Blaizzy]

Note: Please install this transformers PR from source to avoid tokenizer bugs.

pip install git+https://github.com/huggingface/transformers.git@refs/pull/45643/head

Weights here:
https://huggingface.co/collections/mlx-community/deepseek-v4

[Blaizzy]

You can now run it on a 256GB Mac by keeping a experts in 4bit!

We could do 5bit since it's much better than 4bit right now. I'm open to opinions @angeloskath

[Blaizzy]

It's faster now!

[machiabeli]

Hey @Blaizzy — just flagging some technical notes since we're both working on V4 support and PR #1189 landed ~10 hours earlier with significant overlap:

Compressed attention mask direction (line 770-773):
The mask padding for compressed KV rows uses mx.ones, but create_attention_mask returns negative values for blocked positions. Padding with ones would block attention to compressed rows rather than allow it. PR #1189 uses mx.zeros here.

Sinkhorn normalization:
The Python loop path (line 222-226) dispatches ~40 kernel launches per call (softmax + iters x sum + div). PR #1189 has a fused Metal kernel that does this in a single register-resident dispatch — benchmarked at 3.5-5.7x faster on micro, 1.83x end-to-end.

sqrtsoftplus numerical stability:
nn.softplus(x) can overflow for large scores. PR #1189 uses mx.logaddexp(scores, zeros) which is log-sum-exp stable.

Happy to coordinate if the maintainers want to consolidate into one PR. Our implementation has live generation validation at 21.86 tok/s on M3 Ultra (DeepSeek-V4-Flash-4bit, 160GB peak).

[Blaizzy]

Hey @machiabeli, thanks!

Yes, same person who left the earlier feedback, good to connect properly.

I've been poking at this in parallel and landed on something close to the source numerically with minimal changes, but there's definitely room to combine approaches. A PR from you on the compressed attention mask, Sinkhorn norm, and sqrt-softplus would be really welcome, happy to review and merge what works best.

Or I can cherry pick and add you as a co-author.

[Blaizzy]

The goal here is to preserve mxfp4 expert quant since MLX supports it. So I made the quantize_config key in the config class default to that, and these changes help prequantized models load properly.

It can be done via predicate but couldn't find an elegant way of doing it.

Note: it doesn't affect any model.

[Blaizzy]

Alternative is to dequant -> requant similar to how we do with FP8.

[ivaniguarans]

Hey @Blaizzy, thanks for your work!

I tested this branch on an M5 Max 128 GB, trying to requantize DeepSeek-V4-Flash from FP8 source to lower-precision affine (2-3 bpp) via mlx_lm.convert.

The current sanitize in mlx_lm/models/deepseek_v4.py (Model.sanitize()) correctly converts FP8 weights to MXFP8 quantized format for native inference:

# deepseek_v4.py, Model.sanitize(), line 1082-1084 — works for inference
elif weight.dtype == mx.uint8:
    new_weights[k + "s"] = mx.repeat(mx.repeat(v, 4, -1), 128, 0)
    new_weights[wk] = weight.view(mx.uint32)

But there's no path in Model.sanitize() to dequantize FP8 weights back to BF16 for requantization. The naive approach — weight.astype(mx.bfloat16) — treats uint8 bytes as integers (0–255) instead of decoding them as FP8 E4M3 floats. This produces weight values ~20,000× too large (mean ~20000 vs expected ~0.016), causing immediate numerical explosion in layer 0.

The fix is to use mx.dequantize with MXFP8 mode, which properly decodes the FP8 values:

elif weight.dtype == mx.uint8:
    # Dequantize FP8 E4M3 weights with E8M0 block scales to BF16
    m, n = weight.shape
    sm, sn = v.shape
    w32 = weight.view(mx.uint32)
    target_sn = w32.shape[-1] // 8  # MXFP8 group_size=32 → 8 uint32 per scale
    s_expanded = mx.repeat(mx.repeat(v, target_sn // sn, -1), m // sm, 0)
    new_weights[wk] = mx.dequantize(w32, s_expanded, group_size=32, bits=8, mode="mxfp8")

Same pattern applies to the FP4 expert branch (Model.sanitize(), line 1074-1081) — use mx.dequantize(..., mode="mxfp4") instead of manual reinterpretation:

# FP4 experts → BF16
w32 = weight.view(mx.uint32)
new_weights[wk] = mx.dequantize(w32, v, group_size=32, bits=4, mode="mxfp4").astype(mx.bfloat16)

These dequantization paths would go behind a flag (e.g., triggered when convert is called with -q on an FP8 source) so the existing MXFP8 inference path stays untouched.

After this fix, mlx_lm.convert -q --q-bits 2 --q-group-size 64 produces a working 85 GB model (2.56 bpp) from the 149 GB FP8 source. The model generates coherent output for short queries. (2-bit is too aggressive for sustained quality on this architecture, but the dequantization itself is correct — verified by checking weight statistics and layer-0 activation norms.)

[spicyneuron]

I did several rounds of brute force Codex / Opus comparisons between this branch vs transformers and vLLM implementations. Results below.

I'm testing out rough fixes on my own server to see if it helps with the looping / long context issues. But in the meantime, anyone else want to validate?

---

Cross-reference: pr-1192 vs. vLLM and HF transformers

Comparison sources:

• vLLM vllm/model_executor/models/deepseek_v4.py and vllm/model_executor/layers/deepseek_compressor.py
• HF modeling_deepseek_v4.py and configuration_deepseek_v4.py

Topic	pr-1192	vLLM	HF transformers	Note
LocalAttention RoPE scaling	Passes None for rope_scaling on local layers (no YaRN).	Single get_rope per layer; YaRN/mscale apply uniformly, only rope_theta swaps between compress_rope_theta and rope_theta (L1009-1031).	One shared DeepseekV4RotaryEmbedding(config); sliding/local layers (compressor=None) read the same main rope-type that carries scaling (L740-763).	pr-1192 deviates — local layers should use config.rope_scaling.
MoE gate dtype	logits = x @ self.weight.T in activation dtype.	router_logits_dtype=torch.float32 (L853).	F.linear(flat.float(), self.weight.float()) in both TopK and Hash routers (L979, L1010).	pr-1192 deviates — gate must be fp32.
Compressor softmax precision	_simple_compress_kv (L286-290) casts gate logits to fp32 before softmax, but omits precise=True; _overlap_compress_kv (L293-310) casts ape down to gate.dtype (bf16/fp16 from wgate), so its softmax input stays low precision even though it passes precise=True.	Compressor state is fp32 throughout (assert self.dtype == torch.float32, L143).	HCA, Indexer, and CSA all do softmax(..., dtype=torch.float32) (L419, L520, L629-L630).	pr-1192's two compress paths trade off precision in opposite directions; neither matches HF/vLLM's uniformly fp32 compressor softmax.
CSA cross-call overlap state	_overlap_compress_kv prepends a zero/-inf row at every call — overlap is reset each chunk, breaking long-context generation.	Compressor stores prior kv_a/gate_a slice in cache state and pulls it forward across calls.	Same — overlap is part of CSA layer state, not zeroed per-call.	pr-1192 is functionally incorrect for chunked prefill or multi-call generation; the new DeepseekV4PoolingCache carries overlap_kv/overlap_gate across calls.
Pooling cache storage growth	Each update_and_fetch does mx.concatenate([self.pooled, px], axis=1), accumulating lazy-graph nodes until Metal's per-command-buffer resource limit aborts long runs.	N/A (PyTorch/CUDA — no equivalent constraint).	N/A (PyTorch — no equivalent constraint).	MLX-specific. pr-1192 should step-allocate a backing buffer (e.g. 256-token chunks) and expose pooled as a logical-size view into it.

[spicyneuron]

Superseding my previous comment on model looping, I've isolated a script that reliably reproduces the issue at ~4000 tokens on my Mac Studio M3.

Could someone else try it?

# just in case
uv cache prune

# serve this branch
uvx --from git+https://github.com/Blaizzy/mlx-lm/mlx-lm@pc/add-deepseekv4flash-model \
  mlx_lm.server \
  --max-tokens 32000 \
  --model mlx-community/DeepSeek-V4-Flash-8bit

# and then
uv run repro_v4_loop_minimal.py --base-url http://localhost:8080 --min-p 0.0

# still reproduces, but later
uv run repro_v4_loop_minimal.py --base-url http://localhost:8080 --min-p 0.05

Runs sometimes fail with the RuntimeError: [metal::malloc] Resource limit (499000) I mentioned above but simply restarting the server and trying again works to bypass the issue (eventually).

---

EDIT: Tentative fix in this branch, but I'd be more confident in it if others also can reproduce the bug, and then see it go away when using:

uvx --from git+https://github.com/spicyneuron/mlx-lm/mlx-lm@fix-ds4-cache-reuse \
  mlx_lm.server \
  --max-tokens 32000 \
  --model mlx-community/DeepSeek-V4-Flash-8bit

[kidroca]

@spicyneuron

Superseding my previous comment on model looping, I've isolated a script that reliably reproduces the issue at ~4000 tokens on my Mac Studio M3.

Could someone else try it?

I gave a try on a machine I have access to, I pretty much confirm same behavior, though I couldn't recreate with min-p 0.05 - I gave it just a few tries though

---

mlx-lm DeepSeek V4 Flash repro

I ran the repro on a Mac Studio M3 Ultra and can confirm that the baseline branch reproduces the loop, while fix-ds4-cache-reuse fixes the case I could reproduce.

One command difference: the original command in the comment did not work for me because the git URL had an extra /mlx-lm path segment:

git+https://github.com/Blaizzy/mlx-lm/mlx-lm@pc/add-deepseekv4flash-model

I used this corrected baseline command instead:

uvx --from git+https://github.com/Blaizzy/mlx-lm@pc/add-deepseekv4flash-model \
  mlx_lm.server \
  --max-tokens 32000 \
  --model mlx-community/DeepSeek-V4-Flash-8bit

And this corrected fix-branch command:

uvx --from git+https://github.com/spicyneuron/mlx-lm@fix-ds4-cache-reuse \
  mlx_lm.server \
  --max-tokens 32000 \
  --model mlx-community/DeepSeek-V4-Flash-8bit

Baseline, --min-p 0.0:

seed=0 min_p=0.0
turn=1 tokens=4000 status=LOOP
loop='consequential—was consequential—was consequential—was consequential—was consequential—was consequential—was consequential—was consequential—was' x671

Fix branch, --min-p 0.0:

seed=0 min_p=0.0
turn=1 tokens=4000 status=ok
turn=2 tokens=4000 status=ok
no loop reproduced

For --min-p 0.05, I hit the known Metal resource-limit/server-disconnect path on one attempt:

RuntimeError: [metal::malloc] Resource limit (499000) exceeded.
http.client.RemoteDisconnected: Remote end closed connection without response

After restarting the server and retrying, the run completed without a loop (2 times):

seed=0 min_p=0.05
turn=1 tokens=4000 status=ok
turn=2 tokens=4000 status=ok
no loop reproduced

Summary from my run:

• baseline reproduces the loop with min_p=0.0
• fix-ds4-cache-reuse removes that reproduced loop
• min_p=0.05 did not reproduce the loop for me after retrying past the Metal resource-limit failure

[bojiang]

Heads up: the mlx-community DeepSeek-V4-Flash conversions (the mlx-community/deepseek-ai-DeepSeek-V4-Flash-{2,3,4,6,8}bit and -fp16 repos, all uploaded 2026-04-24/25) ship weights in a format that the current sanitize doesn't recognize. Two specific gaps observed on a 4-bit load:

1. embed.{scales,biases} / head.{scales,biases} aren't in top_remap — only the .weight keys are, so the quantized scales/biases never get renamed to model.embed_tokens.* / lm_head.*.
2. Experts are pre-stacked under ffn.experts.w{1,2,3}.{weight,scales,biases} instead of per-expert ffn.experts.E.w{1,2,3}.{...}. The existing expert-stacking block expects the per-expert form and so silently leaves the pre-stacked tensors un-renamed.
3. wo_a is split into o_groups separate Linear modules under attn.wo_a.{0..o_groups-1}.{weight,scales,biases} instead of a single 2D tensor. The current wo_a reshape only handles a 2D-tensor input.

A minimal patch on top of `pc/add-deepseekv4flash-model` that covers all three:

```python

extend the existing top_remap dict with:

"embed.scales": "model.embed_tokens.scales",
"embed.biases": "model.embed_tokens.biases",
"head.scales": "lm_head.scales",
"head.biases": "lm_head.biases",

before the existing per-expert stacking loop, rename pre-stacked experts:

for layer_idx in range(n_layers):
prefix = f"model.layers.{layer_idx}.ffn.experts"
for src, dst in (("w1", "gate_proj"), ("w2", "down_proj"), ("w3", "up_proj")):
for suffix in ("weight", "scales", "biases"):
old_k = f"{prefix}.{src}.{suffix}"
new_k = f"model.layers.{layer_idx}.ffn.switch_mlp.{dst}.{suffix}"
if old_k in weights and new_k not in weights:
weights[new_k] = weights.pop(old_k)

before the existing 2D->3D wo_a reshape, stack split wo_a shards:

for layer_idx in range(n_layers):
prefix = f"model.layers.{layer_idx}.attn.wo_a"
for suffix in ("weight", "scales", "biases"):
key0 = f"{prefix}.0.{suffix}"
if key0 in weights:
stacked = [weights.pop(f"{prefix}.{g}.{suffix}") for g in range(self.args.o_groups)]
weights[f"{prefix}.{suffix}"] = mx.stack(stacked, axis=0)
```

Verified with both 3-bit (124 GB on disk) and 4-bit (149 GB on disk) on a M2 Ultra (192 GB unified). Both load (peak 124.6 GB and 160.2 GB respectively) and generate at ~33–35 tok/s. Separately, neither of those repos ships a `chat_template`; copying the one from `inferencerlabs/DeepSeek-V4-Flash-MLX-2.8bit-EXP` fixes runaway USER/ASSISTANT continuation in `mlx_lm.server`.

Happy to open a PR against your branch if useful.

[snagnever]

Several people in this thread are hitting [metal::malloc] Resource limit (499000) exceeded during long or looping generations on Apple Silicon (and the command queue then wedges until a restart). I tracked down the cause and have a fix.

Root cause: an unbounded per‑decode‑step leak of live Metal buffers in the DeepSeek‑V4 attention caches. PoolingCache (concatenate‑grow) and RotatingKVCache (sliced assignment), plus their Batch* variants, build per‑step lazy graphs that are never detached — the cache holds the head of the chain, so every prior step's intermediate array and its backing Metal buffer stays resident. resource_limit (499000) is a count of live resident buffers, not bytes, so the count climbs ~one per layer per step and hits the cap at ~11.3K tokens regardless of prompt length (≈ 499000 / 43 layers). Only ~2–3 GB is actually leaked at that point — it's the count that's exhausted, which is why no set_memory_limit/set_wired_limit knob helps.

• Root cause + reproducer + full diagnostic path: DeepSeek-V4 (Flash/Pro) on Apple Silicon: unbounded Metal residency growth during decode → [metal::malloc] Resource limit (499000) exceeded after ~11K tokens #1332
• Fix (one hunk in DeepseekV4Model.__call__, targets this branch): Fix DeepSeek-V4 Metal residency leak: materialize per-layer cache state each forward Blaizzy/mlx-lm#25

The fix materializes the per‑layer cache state once per forward pass (mx.eval), which detaches the chains and keeps the live‑buffer count bounded. Verified on M4 Max 128 GB / mlx 0.31.2: forced 20K‑token generation clean to 19,989 (was OOM at 11,314), 31.3 tok/s (no regression), and a 300‑request knowledge‑bench soak on a single long‑lived server with 0 OOMs.

@Blaizzy — happy to adjust the placement if you'd prefer to fold it into this PR directly.

[spicyneuron]

So the code review burden on this model is clearly huge. But I can attest that DS 4 Flash runs incredibly well on Apple hardware, and it feels like a waste that this isn't available to more people.

Open question to @angeloskath @Blaizzy @kernelpool @pcuenca and others — how can we move this forward?

On my end, I was able to patch my way to a stable implementation based on this PR, using transformers and vllm as references. No Resource limit (499000), no catastrophic looping, and tool calling works beautifully.

Code likely suffers from LLM-isms, but I'll put it forward as a rough implementation that others might be able to streamline or borrow from. I had Codex squash the history into a cleaner commit story, so each fix should be reviewable in isolation: Blaizzy/mlx-lm@pc/add-deepseekv4flash-model...spicyneuron:mlx-lm:pr-1192-ds4-clean

[hehua2008]

Are there any talents still interested in this PR?
Looking forward to your amazing work! Thanks～
I miss @awni ‘s hard work😭

[siddharthjthapa]

@hehua2008 Have you tried @antirez ‘s https://github.com/antirez/ds4 ?

[hehua2008]

Thank you for reminding. I will try it.
But have you found that the Apple MLX team has almost ignored this project? This project has not released an updated version for 2 months! The great @awni has left Apple MLX to join Anthropic. Now only volunteers are submitting PR. It shows that Apple's investment in AI is getting less and less, which is different from what they advertise. This is what makes me very disappointed……

[spicyneuron]

Local LLMs are still a niche use case, and AI-coded PRs have added a lot of noise to many open source projects. So it's not surprising things are slow.

But without a core maintainer to champion things, complex PRs are effectively dead ends. So at this point, best we can do is show interest, try to make code review easy, and run our own forks if we need something faster.

[Blaizzy]

Closing this PR for now