Conversation
nh13
force-pushed
the
dsv4/01-core-model
branch
3 times, most recently
from
August 28, 2026 23:57
5ab3413 to
d186f45
Compare
Save each cache entry's class so a model-specific cache subclass (DeepSeek-V4's CompressedKVCache) round-trips through a saved prompt cache. Built-in caches keep a bare name (readable by older mlx-lm); only out-of-module caches are stored module-qualified and imported on load, restricted to mlx_lm modules so loading a cache file cannot import arbitrary code.
… width Mixed-bit quantization reads an optional `mixed_quant_extra_high` tuple from the model and keeps those module names at the high bit width, mirroring the existing `cast_predicate` hook. It is empty by default, so behavior is unchanged for models that do not set it. The layer index also falls back to the first numeric path component when the fixed position is non-numeric, so module paths of varying depth no longer raise.
Self-contained model file for DeepSeek-V4-Flash: MLA attention with a compressed-sparse KV path and indexer top-k, hash-routed and bias-corrected softmax MoE, manifold-constrained Hyper-Connections (Sinkhorn-projected mHC), and dual-theta partial RoPE. Loads community MLX checkpoints (FP8/FP4 dequant, grouped output projection, pre-stacked wo_a) and drops the multi-token- prediction heads, which base inference does not consume. The compressed-attention and indexer correctness was informed by issues raised by @avlp12 and @stepnoy on ml-explore#1189.
Cover the RoPE inverse and sliding-window (no-YaRN) RoPE, the compressed-KV decode path and prompt-cache round-trip, indexer top-k (roped and unroped), the grouped output projection, mixed-quant path preservation, FP8/FP4 sanitize, and the mHC recombination direction (a reference-value test pinning comb^T @ residual). Also register deepseek_v4 in the all-models smoke test.
transformers registers the deepseek_v4 config and tokenizer as of 5.8.0, which the AutoTokenizer load path needs.
nh13
force-pushed
the
dsv4/01-core-model
branch
from
August 29, 2026 00:31
d186f45 to
70819eb
Compare
…uential BatchRotatingKVCache stores `offset` as an mx.array that update_and_fetch() increments in place. V4Attention captured it once and reused it after the cache update for the inverse RoPE on the attention output, so the aliased array was corrupted from the query position to the post-update value, misrotating the output. Batched generation (mlx_lm.server's continuous batcher) diverged from sequential decode from the first token; sequential decode uses an int offset and was unaffected. Snapshot the offset when it is an mx.array. Add a regression test asserting exact batched/sequential decode parity across sliding and compressed layers.
… exists CompressedKVCache.trim() rolls back only the local sliding-window cache; the compressed pool and indexer pool cannot be un-compressed and the raw hidden states behind them are not retained. A prompt-cache trim therefore left stale compressed rows for rolled-back tokens, corrupting sparse attention on reuse. is_trimmable() delegated to the local cache and returned True for short sequences, so the corruption was reachable via trim_prompt_cache. Report is_trimmable() == False once any compressed or indexer pool/buffer has accumulated, so trim_prompt_cache skips the cache instead of corrupting it (a fresh/empty cache stays trimmable). Add a regression test.
Under concurrent serving (mlx_lm.server --decode-concurrency > 1) requests batch into B>1 through the continuous-batching engine, whose caches store `offset` as a per-batch mx.array [B]. A compressed-context prefill builds the compressed-column mask with that offset, but _compressed_mask assumed a scalar: (offset + arange(S_q)) put the batch dim on the S_q axis and the broadcast raised, wedging the generation engine. Build a per-batch [B, 1, S_q, n_comp] mask when offset is an array; the scalar path (and its memoization) is unchanged. B>1 indexer top-k stays NotImplementedError at the call site. Regression test asserts a B=2 compressed prefill matches per-sequence sequential decode.
Contributor
Author
|
DeepSeek-V4 tool calling (DSML parser, chat template, and grammar-constrained decoding) is prepared on top of this branch: nh13#7. It's a preview in my fork, not for merge here yet — flagging so it's on your radar while this base model PR is reviewed. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This adds
deepseek_v4.pyto enable DeepSeek-V4-Flash to run with mlx-lm. MTP is omitted to make this PR focused for review, and I've opened the follow-on PR from my own fork: nh13#1The architecture includes:
sqrtsoftplusto score/weight the MoE instead of softmax.It also adds:
convert.py: a model-agnostic mixed-quant hookcache.py: cache-class serializationtransformersfloor somlx_lm.load()can build the tokenizerThis was verified against the Hugging Face reference implementation and by generating correct output. This is not a logit-for-logit CUDA comparison.
Acknowledgements
Credit to @avlp12 and @stepnoy (#1189) for the compressed-attention/indexer work.
Requirements
Model information
Supported checkpoints
Verification
Verification command:
mlx_lm.generate --model mlx-community/DeepSeek-V4-Flash-4bit -p "The secret to baking a good cake is" -m 1024output:
Prompt: 8 tokens, 31.128 tokens-per-sec
Generation: 1024 tokens, 28.994 tokens-per-sec
Peak memory: 151.797 GB