feat: add DeepSeek-V4-Flash support - #676
Conversation
f38b87c to
00e489b
Compare
|
Pushed an update after auditing how the pieces behave once they are wired into the server rather than exercised on their own. Two integration bugs came out of it, plus a benchmark.
Tool call arguments arrive decoded. With native format preserved, Both are covered by regression tests. I also added a test that chains the reasoning parser into the tool parser the way the server chains them, since testing them separately misses exactly the handoff where the markup crosses over. Benchmark ( Batching costs nothing per token — each request holds independent parser state and totals scale linearly to 16 streams. The DSML parser's own cost is flat. The combined single-stream figure does grow with output length, and it is not this PR's code: Rebased onto current |
|
@Thump604 — review requested. Two notes for reviewing this one. The DSML tool parser is a scanner rather than a regex because the The streaming logic is the part most worth a careful look. Related: #677 fixes two bugs in |
00e489b to
4ebb270
Compare
|
Update after validating this against a live model end to end. Which mlx-lm PR to pair this with: #1192. I tried three that add
#1189 and #1195 emit Verified against #1192 on DeepSeek-V4-Flash-0731 MXFP4 (283.8 B params, M3 Ultra), through this branch's encoder and parsers: Over HTTP with Nothing in the diff changed — it is the same encoder and parsers as before, now with the model side pinned down. |
36aea0d to
8181e62
Compare
|
Update: getting this model actually serving turned up three engine bugs, now fixed in this branch. Flagging them here since they are not really DeepSeek-specific and reviewers may want them split out. 1. Model load and generation ran on different threadsMLX streams exist only in the thread that created them, and an array with pending primitives carries the stream those primitives were built on. Symptoms differed but the cause was one:
The scattered Side effect worth knowing: generation no longer blocks the event loop, so genuinely concurrent requests now reach admission control where the default 2. The prefix cache could never hit on this architectureThe scheduler stored one entry keyed by Fixed by also snapshotting the cache while it still covers exactly the prompt — before the first generated token is appended — and storing that under the prompt tokens. Such an entry is reusable by strict-prefix match with no trimming at all. Two traps: mlx-lm attaches Exact matches are deliberately skipped. The scheduler re-feeds 3.
|
| cache correctness, 5 scenarios × 2 turns | 10/10 character-identical cold vs warm |
| non-stream suite | 30/30, 0 stream errors |
| streaming suite | 18/18 |
| tool calls | 6/6 non-stream, 3/3 streaming |
| concurrency | 4/4, 54.3 tok/s aggregate vs ~31 single-stream |
| long shared prefix | 4.19s → 1.09s |
| chunked prefill, 8413-token prefill + short request 1s later | 1024/chunk: short 27.3s · 256/chunk: short 12.1s |
| repo suite | 2412 passed |
8181e62 to
acaad47
Compare
|
Follow-up: the post-prefill snapshot I added above had a memory bug that only showed under a real agentic client, now fixed. It was stored with The failure lands inside the model during a generation step, the request is aborted, and the client just sees the connection stall. Two changes:
Verified with an 8-turn tool-calling conversation over a ~33k-token context:
Worth noting for reviewers: the byte-based limit in |
21bf1d5 to
12dcc4e
Compare
Thump604
left a comment
Thump604
left a comment
There was a problem hiding this comment.
The DeepSeek encoder/parser work is valuable, but the current branch is no longer reviewable or safe to merge as one feature: 26 files / +3,089 lines now combine DeepSeek prompt+DSML support with SimpleEngine and BatchedEngine thread ownership, prefix-cache snapshot/accounting policy, and chunked-prefill behavior. The embedded SimpleEngine lifecycle is also the earlier implementation superseded by #679's stop/drain corrections, and #648 already carries the narrow chunked-prefill API update.
Please split this into independently reviewable changes:
- DeepSeek V4 encoder, DSML tool/reasoning parsers, registration/CLI, goldens, and their focused tests.
- SimpleEngine ownership via #679; remove the duplicate implementation here.
- A separate BatchedEngine owner-thread/error-handling PR with its own reproduction and lifecycle tests.
- Separate prefix-cache changes, with explicit coverage for non-trimmable/rotating caches and memory ownership.
- Reuse/rebase onto #648 for chunked prefill rather than carrying another copy.
Then rebase the DeepSeek slice after #677 and rerun the combined streaming/tool replay. Finally, update the dependency story: the PR currently identifies mlx-lm #1192 as the usable implementation, but #1192 is closed and unmerged, so the feature needs an installable upstream version/active dependency plus an end-to-end reproduction before merge. CI being green on this combined branch does not close those scope and dependency blockers.
12dcc4e to
c37e588
Compare
|
Rebased onto current One conflict, in That is only the mechanical blocker. On the substance of your review: The split is fair and I am not going to argue it. The branch grew engine changes because that was the only way to get the model serving at all, but you are right that "DeepSeek encoder + DSML parsers" and "engine thread ownership + prefix-cache policy + chunked prefill" are four different review problems wearing one hat, and the SimpleEngine copy here is genuinely superseded by #679's stop/drain work. I will restructure along your lines: reuse #679 for SimpleEngine, rebase chunked prefill onto your #648 rather than carrying a second copy, and take BatchedEngine's owner thread and the prefix-cache snapshot out as their own PRs with their own reproductions. The DeepSeek slice then rebases after #677 and gets the combined streaming/tool replay re-run. The dependency blocker is real and I do not have an answer for it yet. I re-checked upstream today: no mlx-lm PR adding Given that, my preference is to land the parts that stand on their own — they are useful independently of whether DeepSeek ever arrives — and hold the DeepSeek slice until upstream settles. If you would rather I keep it open as a tracking branch or close it until then, say which and I will do that; I will not close it on my own. |
c37e588 to
dde8feb
Compare
dde8feb to
583c5cb
Compare
|
Split done, along the lines you set out. This PR is now the DeepSeek slice alone: 14 files, +2333 instead of 26 files / +3089.
Notes on three of them. #684 is stacked on #679, because it needs On #648: I dropped my chunked-prefill change rather than rebasing onto yours, because yours is strictly better. Mine mapped the budget onto #683 and #684 both got the coverage you asked for, and writing #684's tests caught me repeating the exact mistake waybarrios flagged on #679: my first version replicated the engine's load call in the test body, so reverting Merge order for this PR: it depends on nothing, but #677 changes the reasoning base class The dependency blocker stands unchanged, and I am not going to pretend otherwise: no mlx-lm PR adding |
Thump604
left a comment
Thump604
left a comment
There was a problem hiding this comment.
The split is good. This is reviewable now, and the engine, cache, and chunked-prefill scope is gone. Two blockers remain. The only model implementation validated end to end is closed and unmerged mlx-lm #1192, while released mlx-lm and current main still have no DeepSeek V4 model. Please update the PR body and rerun end to end once there is an installable dependency. Also, buffered streaming drops text before a tool call when the text and complete DSML block arrive in one delta; the parser returns only tool_calls. Please cover that case and preserve both. Keeping changes requested until those are fixed.
|
Keep it open as the tracking branch. The split was the right move, and the standalone engine/cache work can move independently. This one should stay held until there is an installable DeepSeek V4 dependency and the buffered text-before-tool-call case is fixed. |
MLX streams exist only in the thread that created them, and `BatchGenerator` captures `generation_stream` into `self._stream` when it is built. So the thread that loads the model has to be the thread that drives `scheduler.step`. The two batched paths step on different threads, so they load on different threads: - AsyncEngineCore steps on a worker. `BatchedEngine` loaded inline on the event loop (issue waybarrios#407), so the two never matched. The symptom does not look like a stream problem: `batch_generator.next()` raises "There is no Stream(gpu, N) in current thread", `EngineCore` falls back to stepping on the model thread, hits the same error there, and — because that fallback fires only once — then spins on the error. What an operator sees is `running=1`, the step counter climbing into the millions, and not one token emitted. It reads as a scheduler hang, which is how I first misdiagnosed it. - MLLM never reaches AsyncEngineCore. `_start_mllm` drives MLLMScheduler, whose `_process_loop` calls `step()` on the event loop with no executor hop, so an MLLM model has to be built there — as it was before this change. `_model_load_executor()` states that policy in one place, and `start()` and ResidencyManager both read it. `EngineCore` steps on a supplied worker and does not shut it down, since that thread owns the loaded model and outlives the engine loop; without one it still creates and retires its own, unchanged. Stacked on waybarrios#679 for `run_blocking_startup_work(executor=...)`. It cleans up by itself once waybarrios#679 lands. Split out of waybarrios#676 on review. Verified on DeepSeek-V4-Flash-0731 (283.8B MXFP4, M3 Ultra), which could not emit a single token through the batched path before this: 4/4 concurrent requests, 54.3 tok/s aggregate against ~31 single-stream. `tests/test_batched_engine_owner_thread.py` drives the engine's real `start()` rather than replicating its load call, and covers both paths: MLLM load and MLLMScheduler stepping land on the event loop, non-MLLM load and EngineCore stepping land on the worker, and ResidencyManager honours the same split. Confirmed by mutation — pinning MLLM to the worker, having lifecycle ignore the policy, dropping `executor=` from the load, and shutting down a caller-supplied worker are each caught by a distinct test. Repo suite: 2282 passed. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The scheduler stored one prefix-cache entry per request, keyed by
``prompt + output``. Every later query is *shorter* than such a key, so reuse
requires trimming the generated tail off first. Models with sliding-window or
pooled KV cannot do that: ``RotatingKVCache.is_trimmable()`` is
``offset < max_size``, so once the ring buffer wraps the older KV is physically
overwritten, and a pooling cache merges tokens into windows it cannot split.
Zero hits were guaranteed by construction on those architectures, and the
entries were not free: each holds a full-length KV copy, and in an agentic loop
where every turn extends the previous prompt they accumulate. Measured 45
entries of a ~46k-token cache. Metal runs out of *buffers* long before the
byte budget notices:
RuntimeError: [metal::malloc] Resource limit (499000) exceeded.
Two changes:
- Snapshot the cache once it covers the prompt and store that. Such an entry is
reusable by strict prefix match with no trimming at all.
- Skip the completion-time ``prompt + output`` entry when the cache is not
trimmable, since it can never be reused.
Four things the implementation has to get right:
- mlx-lm attaches ``prompt_cache`` only to the response carrying a
``finish_reason``, so mid-generation it is None. The per-sequence cache is
pulled from the live batch via ``extract_cache(idx)``.
- The snapshot must be a real copy. Both cache types write into their buffers
in place, so a snapshot that aliases them is rewritten by the generation it
is supposed to predate.
- It is stored with ``evict_prefixes=True``, or each turn adds another
full-length copy rather than replacing the entry it supersedes.
- The key must name exactly the tokens the snapshot holds. The snapshot is
taken while processing the response carrying the first generated token, and
the batch has already fed that token through the cache — measured
``prompt_len=5, cache_offset=6`` on a real scheduler run. Storing that under
``prompt_token_ids`` left every warm reuse one token ahead of its key.
Trimming the overshoot off is not available here, so the key is extended
instead; the extra token is the first token of the reply, which the next
turn's prompt also contains, so the entry still matches by strict prefix.
A key that cannot be named exactly means no entry at all.
The snapshot destination mirrors the live cache objects rather than calling
``make_prompt_cache``. A plain ``KVCache`` destination cannot take a
``RotatingKVCache``'s state or meta_state; the assignment raised, the broad
handler logged a warning, and nothing was stored on precisely the
configurations this feature targets. Deriving the destination from
``config.max_kv_size`` instead is also wrong, which measurement showed:
``_create_batch_generator`` does not pass ``max_kv_size`` to
``BatchGenerator``, so with it configured the live layers were still plain
``KVCache``.
Exact matches are deliberately skipped rather than used. The scheduler re-feeds
``prompt[-1:]`` on those, which duplicates that token in the KV cache when the
entry covers the whole key — measured as the same prompt returning a different
answer.
``memory_cache.py`` could not measure these caches either: ``CacheList.state``
is nested, and the two-way unpack raised a ValueError that was swallowed, so
every such entry was accounted as 0 bytes and the byte-based LRU never evicted.
Verified on DeepSeek-V4-Flash-0731 (283.8B MXFP4, M3 Ultra), an 8-turn
tool-calling conversation over a ~33k-token context: cache entries 45 -> 1,
Metal buffer errors 7 -> 0, 7/8 hits with 84 tokens prefilled instead of 33k,
long shared prefix 4.19s -> 1.09s. Re-checked end to end on Qwen3-0.6B through
the real scheduler after the alignment fix: identical greedy output cold and
warm, key length equal to cache coverage, and a prefix hit on the extended
second-turn prompt.
Split out of waybarrios#676 on review. Repo suite: 2310 passed.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The scheduler stored one prefix-cache entry per request, keyed by
``prompt + output``. Every later query is *shorter* than such a key, so reuse
requires trimming the generated tail off first. Models with sliding-window or
pooled KV cannot do that: ``RotatingKVCache.is_trimmable()`` is
``offset < max_size``, so once the ring buffer wraps the older KV is physically
overwritten, and a pooling cache merges tokens into windows it cannot split.
Zero hits were guaranteed by construction on those architectures, and the
entries were not free: each holds a full-length KV copy, and in an agentic loop
where every turn extends the previous prompt they accumulate. Measured 45
entries of a ~46k-token cache. Metal runs out of *buffers* long before the
byte budget notices:
RuntimeError: [metal::malloc] Resource limit (499000) exceeded.
Two changes:
- Snapshot the cache once it covers the prompt and store that. Such an entry is
reusable by strict prefix match with no trimming at all.
- Skip the completion-time ``prompt + output`` entry when the cache is not
trimmable, since it can never be reused.
Four things the implementation has to get right:
- mlx-lm attaches ``prompt_cache`` only to the response carrying a
``finish_reason``, so mid-generation it is None. The per-sequence cache is
pulled from the live batch via ``extract_cache(idx)``.
- The snapshot must be a real copy. Both cache types write into their buffers
in place, so a snapshot that aliases them is rewritten by the generation it
is supposed to predate.
- It is stored with ``evict_prefixes=True``, or each turn adds another
full-length copy rather than replacing the entry it supersedes.
- The key must name exactly the tokens the snapshot holds. The snapshot is
taken while processing the response carrying the first generated token, and
the batch has already fed that token through the cache — measured
``prompt_len=5, cache_offset=6`` on a real scheduler run. Storing that under
``prompt_token_ids`` left every warm reuse one token ahead of its key.
Trimming the overshoot off is not available here, so the key is extended
instead; the extra token is the first token of the reply, which the next
turn's prompt also contains, so the entry still matches by strict prefix.
A key that cannot be named exactly means no entry at all.
The snapshot destination mirrors the live cache objects rather than calling
``make_prompt_cache``. A plain ``KVCache`` destination cannot take a
``RotatingKVCache``'s state or meta_state; the assignment raised, the broad
handler logged a warning, and nothing was stored on precisely the
configurations this feature targets. Deriving the destination from
``config.max_kv_size`` instead is also wrong, which measurement showed:
``_create_batch_generator`` does not pass ``max_kv_size`` to
``BatchGenerator``, so with it configured the live layers were still plain
``KVCache``.
Exact matches are deliberately skipped rather than used. The scheduler re-feeds
``prompt[-1:]`` on those, which duplicates that token in the KV cache when the
entry covers the whole key — measured as the same prompt returning a different
answer.
``memory_cache.py`` could not measure these caches either: ``CacheList.state``
is nested, and the two-way unpack raised a ValueError that was swallowed, so
every such entry was accounted as 0 bytes and the byte-based LRU never evicted.
Verified on DeepSeek-V4-Flash-0731 (283.8B MXFP4, M3 Ultra), an 8-turn
tool-calling conversation over a ~33k-token context: cache entries 45 -> 1,
Metal buffer errors 7 -> 0, 7/8 hits with 84 tokens prefilled instead of 33k,
long shared prefix 4.19s -> 1.09s. Re-checked end to end on Qwen3-0.6B through
the real scheduler after the alignment fix: identical greedy output cold and
warm, key length equal to cache coverage, and a prefix hit on the extended
second-turn prompt.
Split out of waybarrios#676 on review. Repo suite: 2310 passed.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
DeepSeek-V4 has no Jinja chat template, so without a prompt encoder the prompt
is built by plain concatenation and the model sees a format it was never
trained on. It also emits tool calls in its own DSML markup rather than JSON.
This adds the encoder and both parsers, plus registration and CLI wiring.
The DSML tool parser is a scanner rather than a regex, because the
`string="true|false"` attribute means a string parameter may legitimately
contain quotes, angle brackets or a JSON-looking payload;
`TestParameterTyping::test_string_value_may_contain_markup_like_text` pins that.
The streaming path is the subtle part. `<｜DSML｜tool_calls>` has no token id of
its own — only the bare `｜DSML｜` does — so the marker always straddles delta
boundaries, and two obvious implementations are both wrong: detecting
completion against the delta rather than the accumulated text drops the calls
entirely, and emitting marker fragments as they arrive leaks markup to the
client and then repeats the whole marker. Both are covered at chunk sizes 1
through 128.
Two integration bugs that only appear once this is wired into the server:
- `SUPPORTS_NATIVE_TOOL_FORMAT` must be True. With False,
`extract_multimodal_content` flattens `role="tool"` into
`"[Tool Result (id)]: ..."` and assistant `tool_calls` into
`"[Calling tool: name(...)]"` *before* the encoder runs, so a multi-turn tool
conversation reaches the model as prose and the encoder's own
`<tool_result>`/DSML handling never fires.
- With native format preserved, `api/utils.py` json-loads `arguments` in place.
The encoder loaded it again, and `json.loads` on a mapping raises, so every
parameter collapsed into one bogus `arguments` entry — the model saw a
malformed call in its own history. It now accepts either form.
Benchmark (`benchmarks/bench_deepseek_v4.py`), both serving paths:
Prompt encoding 4 messages -> 0.022 ms 130 messages -> 0.335 ms
Single stream 100 tok -> 0.0061 ms/tok 5000 tok -> 0.0380 ms/tok
DSML parser alone 100 tok -> 0.0018 ms/tok 5000 tok -> 0.0028 ms/tok
Batched decode 1 concurrent -> 0.0107 16 concurrent -> 0.0107 ms/tok
Batching costs nothing per token; each request holds independent parser state.
Verified end to end on DeepSeek-V4-Flash-0731 MXFP4 (283.8B, M3 Ultra): 11/11
over HTTP with `--tool-call-parser deepseek_v4 --reasoning-parser deepseek_v4`
— `finish_reason=tool_calls`, arguments as JSON objects, reasoning split into
`reasoning_content`, no DSML in user-visible content, parallel calls intact,
and a tool-result round trip where the model uses the returned values.
Scope note: this was previously one branch carrying engine changes as well.
Those are now waybarrios#679 (SimpleEngine ownership), waybarrios#684 (BatchedEngine owner thread)
and waybarrios#683 (prefix cache on non-trimmable KV), and chunked prefill is dropped in
favour of waybarrios#648. This PR is the DeepSeek slice alone.
Merge order: this depends on nothing, but waybarrios#677 changes the reasoning base class
this parser inherits from. Checked merged with waybarrios#677 rather than only alongside
it: 307 targeted tests and 2479 repo-wide, clean.
Repo suite on this branch alone: 2439 passed.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
|
Understood on keeping it open as the tracking branch — that matches how I read the dependency situation too. Fixed the second blocker. Reproduced it first: So the response depended on nothing but how the model's output happened to be chunked, which is the worst shape for a bug like this to have — token-level streaming hides it, and it only surfaces when a whole response lands at once. The cause was a deliberate choice of mine, and the comment in the code even argued for it: "only a whole response arriving as a single delta would force the choice, and there the calls win." That was wrong. The text is user-visible assistant content and dropping it silently is data loss, not a tie-break. It also could not be fixed in the parser alone. The premise I had built on was real — Here, the parser now flushes the buffered head with the calls when the block opens and closes in the same delta, and ahead of them otherwise. New Repo suite on this branch: 2447 passed. That leaves only the dependency blocker, which is unchanged: still no merged mlx-lm PR adding |
583c5cb to
961b402
Compare
The scheduler stored one prefix-cache entry per request, keyed by
``prompt + output``. Every later query is *shorter* than such a key, so reuse
requires trimming the generated tail off first. Models with sliding-window or
pooled KV cannot do that: ``RotatingKVCache.is_trimmable()`` is
``offset < max_size``, so once the ring buffer wraps the older KV is physically
overwritten, and a pooling cache merges tokens into windows it cannot split.
Zero hits were guaranteed by construction on those architectures, and the
entries were not free: each holds a full-length KV copy, and in an agentic loop
where every turn extends the previous prompt they accumulate. Measured 45
entries of a ~46k-token cache. Metal runs out of *buffers* long before the
byte budget notices:
RuntimeError: [metal::malloc] Resource limit (499000) exceeded.
Two changes:
- Snapshot the cache once it covers the prompt and store that. Such an entry is
reusable by strict prefix match with no trimming at all.
- Skip the completion-time ``prompt + output`` entry when the cache is not
trimmable, since it can never be reused.
Four things the implementation has to get right:
- mlx-lm attaches ``prompt_cache`` only to the response carrying a
``finish_reason``, so mid-generation it is None. The per-sequence cache is
pulled from the live batch via ``extract_cache(idx)``.
- The snapshot must be a real copy. Both cache types write into their buffers
in place, so a snapshot that aliases them is rewritten by the generation it
is supposed to predate.
- It is stored with ``evict_prefixes=True``, or each turn adds another
full-length copy rather than replacing the entry it supersedes.
- The key must name exactly the tokens the snapshot holds. The snapshot is
taken while processing the response carrying the first generated token, and
the batch has already fed that token through the cache — measured
``prompt_len=5, cache_offset=6`` on a real scheduler run. Storing that under
``prompt_token_ids`` left every warm reuse one token ahead of its key.
Trimming the overshoot off is not available here, so the key is extended
instead; the extra token is the first token of the reply, which the next
turn's prompt also contains, so the entry still matches by strict prefix.
A key that cannot be named exactly means no entry at all.
The snapshot destination mirrors the live cache objects rather than calling
``make_prompt_cache``. A plain ``KVCache`` destination cannot take a
``RotatingKVCache``'s state or meta_state; the assignment raised, the broad
handler logged a warning, and nothing was stored on precisely the
configurations this feature targets. Deriving the destination from
``config.max_kv_size`` instead is also wrong, which measurement showed:
``_create_batch_generator`` does not pass ``max_kv_size`` to
``BatchGenerator``, so with it configured the live layers were still plain
``KVCache``.
Exact matches are deliberately skipped rather than used. The scheduler re-feeds
``prompt[-1:]`` on those, which duplicates that token in the KV cache when the
entry covers the whole key — measured as the same prompt returning a different
answer.
``memory_cache.py`` could not measure these caches either: ``CacheList.state``
is nested, and the two-way unpack raised a ValueError that was swallowed, so
every such entry was accounted as 0 bytes and the byte-based LRU never evicted.
Verified on DeepSeek-V4-Flash-0731 (283.8B MXFP4, M3 Ultra), an 8-turn
tool-calling conversation over a ~33k-token context: cache entries 45 -> 1,
Metal buffer errors 7 -> 0, 7/8 hits with 84 tokens prefilled instead of 33k,
long shared prefix 4.19s -> 1.09s. Re-checked end to end on Qwen3-0.6B through
the real scheduler after the alignment fix: identical greedy output cold and
warm, key length equal to cache coverage, and a prefix hit on the extended
second-turn prompt.
Split out of waybarrios#676 on review. Repo suite: 2310 passed.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…es (#683) The scheduler stored one prefix-cache entry per request, keyed by ``prompt + output``. Every later query is *shorter* than such a key, so reuse requires trimming the generated tail off first. Models with sliding-window or pooled KV cannot do that: ``RotatingKVCache.is_trimmable()`` is ``offset < max_size``, so once the ring buffer wraps the older KV is physically overwritten, and a pooling cache merges tokens into windows it cannot split. Zero hits were guaranteed by construction on those architectures, and the entries were not free: each holds a full-length KV copy, and in an agentic loop where every turn extends the previous prompt they accumulate. Measured 45 entries of a ~46k-token cache. Metal runs out of *buffers* long before the byte budget notices: RuntimeError: [metal::malloc] Resource limit (499000) exceeded. Two changes: - Snapshot the cache once it covers the prompt and store that. Such an entry is reusable by strict prefix match with no trimming at all. - Skip the completion-time ``prompt + output`` entry when the cache is not trimmable, since it can never be reused. Four things the implementation has to get right: - mlx-lm attaches ``prompt_cache`` only to the response carrying a ``finish_reason``, so mid-generation it is None. The per-sequence cache is pulled from the live batch via ``extract_cache(idx)``. - The snapshot must be a real copy. Both cache types write into their buffers in place, so a snapshot that aliases them is rewritten by the generation it is supposed to predate. - It is stored with ``evict_prefixes=True``, or each turn adds another full-length copy rather than replacing the entry it supersedes. - The key must name exactly the tokens the snapshot holds. The snapshot is taken while processing the response carrying the first generated token, and the batch has already fed that token through the cache — measured ``prompt_len=5, cache_offset=6`` on a real scheduler run. Storing that under ``prompt_token_ids`` left every warm reuse one token ahead of its key. Trimming the overshoot off is not available here, so the key is extended instead; the extra token is the first token of the reply, which the next turn's prompt also contains, so the entry still matches by strict prefix. A key that cannot be named exactly means no entry at all. The snapshot destination mirrors the live cache objects rather than calling ``make_prompt_cache``. A plain ``KVCache`` destination cannot take a ``RotatingKVCache``'s state or meta_state; the assignment raised, the broad handler logged a warning, and nothing was stored on precisely the configurations this feature targets. Deriving the destination from ``config.max_kv_size`` instead is also wrong, which measurement showed: ``_create_batch_generator`` does not pass ``max_kv_size`` to ``BatchGenerator``, so with it configured the live layers were still plain ``KVCache``. Exact matches are deliberately skipped rather than used. The scheduler re-feeds ``prompt[-1:]`` on those, which duplicates that token in the KV cache when the entry covers the whole key — measured as the same prompt returning a different answer. ``memory_cache.py`` could not measure these caches either: ``CacheList.state`` is nested, and the two-way unpack raised a ValueError that was swallowed, so every such entry was accounted as 0 bytes and the byte-based LRU never evicted. Verified on DeepSeek-V4-Flash-0731 (283.8B MXFP4, M3 Ultra), an 8-turn tool-calling conversation over a ~33k-token context: cache entries 45 -> 1, Metal buffer errors 7 -> 0, 7/8 hits with 84 tokens prefilled instead of 33k, long shared prefix 4.19s -> 1.09s. Re-checked end to end on Qwen3-0.6B through the real scheduler after the alignment fix: identical greedy output cold and warm, key length equal to cache coverage, and a prefix hit on the extended second-turn prompt. Split out of #676 on review. Repo suite: 2310 passed. Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
The scheduler stored one prefix-cache entry per request, keyed by
``prompt + output``. Every later query is *shorter* than such a key, so reuse
requires trimming the generated tail off first. Models with sliding-window or
pooled KV cannot do that: ``RotatingKVCache.is_trimmable()`` is
``offset < max_size``, so once the ring buffer wraps the older KV is physically
overwritten, and a pooling cache merges tokens into windows it cannot split.
Zero hits were guaranteed by construction on those architectures, and the
entries were not free: each holds a full-length KV copy, and in an agentic loop
where every turn extends the previous prompt they accumulate. Measured 45
entries of a ~46k-token cache. Metal runs out of *buffers* long before the
byte budget notices:
RuntimeError: [metal::malloc] Resource limit (499000) exceeded.
Two changes:
- Snapshot the cache once it covers the prompt and store that. Such an entry is
reusable by strict prefix match with no trimming at all.
- Skip the completion-time ``prompt + output`` entry when the cache is not
trimmable, since it can never be reused.
Four things the implementation has to get right:
- mlx-lm attaches ``prompt_cache`` only to the response carrying a
``finish_reason``, so mid-generation it is None. The per-sequence cache is
pulled from the live batch via ``extract_cache(idx)``.
- The snapshot must be a real copy. Both cache types write into their buffers
in place, so a snapshot that aliases them is rewritten by the generation it
is supposed to predate.
- It is stored with ``evict_prefixes=True``, or each turn adds another
full-length copy rather than replacing the entry it supersedes.
- The key must name exactly the tokens the snapshot holds. The snapshot is
taken while processing the response carrying the first generated token, and
the batch has already fed that token through the cache — measured
``prompt_len=5, cache_offset=6`` on a real scheduler run. Storing that under
``prompt_token_ids`` left every warm reuse one token ahead of its key.
Trimming the overshoot off is not available here, so the key is extended
instead; the extra token is the first token of the reply, which the next
turn's prompt also contains, so the entry still matches by strict prefix.
A key that cannot be named exactly means no entry at all.
The snapshot destination mirrors the live cache objects rather than calling
``make_prompt_cache``. A plain ``KVCache`` destination cannot take a
``RotatingKVCache``'s state or meta_state; the assignment raised, the broad
handler logged a warning, and nothing was stored on precisely the
configurations this feature targets. Deriving the destination from
``config.max_kv_size`` instead is also wrong, which measurement showed:
``_create_batch_generator`` does not pass ``max_kv_size`` to
``BatchGenerator``, so with it configured the live layers were still plain
``KVCache``.
Exact matches are deliberately skipped rather than used. The scheduler re-feeds
``prompt[-1:]`` on those, which duplicates that token in the KV cache when the
entry covers the whole key — measured as the same prompt returning a different
answer.
``memory_cache.py`` could not measure these caches either: ``CacheList.state``
is nested, and the two-way unpack raised a ValueError that was swallowed, so
every such entry was accounted as 0 bytes and the byte-based LRU never evicted.
Verified on DeepSeek-V4-Flash-0731 (283.8B MXFP4, M3 Ultra), an 8-turn
tool-calling conversation over a ~33k-token context: cache entries 45 -> 1,
Metal buffer errors 7 -> 0, 7/8 hits with 84 tokens prefilled instead of 33k,
long shared prefix 4.19s -> 1.09s. Re-checked end to end on Qwen3-0.6B through
the real scheduler after the alignment fix: identical greedy output cold and
warm, key length equal to cache coverage, and a prefix hit on the extended
second-turn prompt.
Split out of waybarrios#676 on review. Repo suite: 2310 passed.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
MLX streams exist only in the thread that created them, and `BatchGenerator` captures `generation_stream` into `self._stream` when it is built. So the thread that loads the model has to be the thread that drives `scheduler.step`. The two batched paths step on different threads, so they load on different threads: - AsyncEngineCore steps on a worker. `BatchedEngine` loaded inline on the event loop (issue waybarrios#407), so the two never matched. The symptom does not look like a stream problem: `batch_generator.next()` raises "There is no Stream(gpu, N) in current thread", `EngineCore` falls back to stepping on the model thread, hits the same error there, and — because that fallback fires only once — then spins on the error. What an operator sees is `running=1`, the step counter climbing into the millions, and not one token emitted. It reads as a scheduler hang, which is how I first misdiagnosed it. - MLLM never reaches AsyncEngineCore. `_start_mllm` drives MLLMScheduler, whose `_process_loop` calls `step()` on the event loop with no executor hop, so an MLLM model has to be built there — as it was before this change. `_model_load_executor()` states that policy in one place, and `start()` and ResidencyManager both read it. `EngineCore` steps on a supplied worker and does not shut it down, since that thread owns the loaded model and outlives the engine loop; without one it still creates and retires its own, unchanged. Stacked on waybarrios#679 for `run_blocking_startup_work(executor=...)`. It cleans up by itself once waybarrios#679 lands. Split out of waybarrios#676 on review. Verified on DeepSeek-V4-Flash-0731 (283.8B MXFP4, M3 Ultra), which could not emit a single token through the batched path before this: 4/4 concurrent requests, 54.3 tok/s aggregate against ~31 single-stream. `tests/test_batched_engine_owner_thread.py` drives the engine's real `start()` rather than replicating its load call, and covers both paths: MLLM load and MLLMScheduler stepping land on the event loop, non-MLLM load and EngineCore stepping land on the worker, and ResidencyManager honours the same split. Confirmed by mutation — pinning MLLM to the worker, having lifecycle ignore the policy, dropping `executor=` from the load, and shutting down a caller-supplied worker are each caught by a distinct test. Repo suite: 2282 passed. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
|
Pushed follow-up patch
Local verification: 502 compatibility tests and 172 dedicated V4 tests passed. CI is running. |
|
Pushed |
3 participants
DeepSeek-V4 ships no Jinja
chat_template— itstokenizer_config.jsoncarries only BOS/EOS/pad — so the prompt has to be built programmatically. Today the model would either raise in_apply_chat_templateor fall back to naive"role: content"concatenation, and its DSML tool calls would go unparsed. This adds the three pieces needed to serve it.Prompt encoder
vllm_mlx/utils/deepseek_v4_encoding.pyports the referenceencoding_dsv4.pypublished with the weights:System content is bare text with no wrapper; roles are delimited solely by
<｜User｜>/<｜Assistant｜>. The turn closes on<think>in thinking mode or</think>in chat mode, which suppresses reasoning. Tool schemas render into the system message, and tool results fold into the preceding user turn as<tool_result>blocks — V4 has notoolrole.reasoning_effortis a text prefix on the whole conversation rather than a token or a sampling parameter, anddrop_thinkingis forced off when tools are present, because the model needs to see why it made the earlier calls.vllm_mlx/utils/tokenizer.pyinstalls the encoder by overridingapply_chat_templateon the tokenizer whenmodel_typeisdeepseek_v4. That fixes all three call sites at once —engine/batched.py,engine/simple.pyandmodels/llm.pyall reach the template through that one method — without touching any of them. It mirrors what upstream vLLM does invllm/tokenizers/deepseek_v4.py.DSML tool parser
V4 emits its own markup rather than JSON:
The
stringattribute carries the type —"true"is a raw string,"false"is JSON. That is why this is a scanner and not a regex overname=valuepairs: a string parameter may legitimately contain quotes, angle brackets or a JSON-looking payload. Registered asdeepseek_v4/dsmland wired intoAutoToolParser; the existingDeepSeekToolParserhandles the V3/R1<｜tool▁calls▁begin｜>format and shares nothing with this one.Reasoning parser
Extends the R1 parser, which already tolerates the missing opening
<think>that the encoder's prompt implies. What V4 adds is that a tool call must follow completed reasoning, so an opening<｜DSML｜tool_calls>terminates the reasoning block even when</think>never arrives. Without that, the entire DSML payload is swallowed as reasoning and the caller sees no tool call at all.Streaming
Both parsers track how much of the accumulated text they have emitted and withhold a tail that could still grow into a marker.
<｜DSML｜tool_calls>is assembled from several tokens — only the bare｜DSML｜has an id of its own — so it always straddles delta boundaries. Two failure modes follow from that, and both are covered by tests: detecting completion against the delta rather than the accumulated text loses the calls entirely, and emitting marker fragments as they arrive leaks markup into the user-visible stream and then repeats the whole marker once it is recognised.Scope
Model loading is deliberately not part of this. vllm-mlx defines no model architectures —
models/only wrapsmlx_lm.load()— anddeepseek_v4lands in mlx-lm via ml-explore/mlx-lm#1189, which is open and active. Everything here is text processing and starts working the moment that merges. There is precedent for staging it this way:text_model_from_vlm.pyalready reports"Cannot import mlx_lm TextModel (need PR #990)".Continuous batching will likely want its own patch alongside
patches/gemma4_mllm.pyandpatches/glm4v_moe_mllm.py, since V4 has MLA with a compressor, a sparse indexer and mHC. That is best written once the model actually runs.Testing
tests/test_deepseek_v4_encoding.pyasserts against prompts frozen from the reference implementation, which cannot be vendored here. During development the port was also run differentially against that reference across 832 combinations of conversation shape, thinking mode, effort level and tool usage — all identical.Verified against the real model, DeepSeek-V4-Flash-0731 MXFP4 with 167 GB resident on an M3 Ultra, across 23 scenarios: chat, thinking, every
reasoning_effortlevel, single and parallel tool calls, tool-result round trips and multi-turn history. Prompts were built by this encoder and the completions run back through these parsers — 127/128 checks pass. The one failure is a scenario truncated at the token cap before it emitted</think>, where treating the output as content is the correct behaviour. The model's DSML matched the specification including thestring="true|false"type flag, and parallel calls came back intact:The three new test files are added to the CI whitelist in
.github/workflows/ci.yml, without which they would never run.