From 961b4021102f3ea532791c735bbadf698b5ad572 Mon Sep 17 00:00:00 2001 From: Jan Hilgard Date: Sat, 8 Aug 2026 10:59:05 +0200 Subject: [PATCH 1/3] feat: add DeepSeek-V4-Flash support MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 ``/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 #679 (SimpleEngine ownership), #684 (BatchedEngine owner thread) and #683 (prefix cache on non-trimmable KV), and chunked prefill is dropped in favour of #648. This PR is the DeepSeek slice alone. Merge order: this depends on nothing, but #677 changes the reasoning base class this parser inherits from. Checked merged with #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) --- benchmarks/bench_deepseek_v4.py | 214 ++++++ tests/deepseek_v4_golden_prompts.py | 79 ++ tests/test_deepseek_v4_encoding.py | 329 ++++++++ tests/test_deepseek_v4_reasoning.py | 141 ++++ tests/test_deepseek_v4_tool_parser.py | 420 ++++++++++ vllm_mlx/cli.py | 5 +- vllm_mlx/reasoning/__init__.py | 2 + vllm_mlx/reasoning/deepseek_v4_parser.py | 130 ++++ vllm_mlx/server.py | 3 + vllm_mlx/tool_parsers/__init__.py | 3 + vllm_mlx/tool_parsers/auto_tool_parser.py | 12 + .../tool_parsers/deepseek_v4_tool_parser.py | 297 +++++++ vllm_mlx/utils/deepseek_v4_encoding.py | 726 ++++++++++++++++++ vllm_mlx/utils/tokenizer.py | 31 + 14 files changed, 2391 insertions(+), 1 deletion(-) create mode 100644 benchmarks/bench_deepseek_v4.py create mode 100644 tests/deepseek_v4_golden_prompts.py create mode 100644 tests/test_deepseek_v4_encoding.py create mode 100644 tests/test_deepseek_v4_reasoning.py create mode 100644 tests/test_deepseek_v4_tool_parser.py create mode 100644 vllm_mlx/reasoning/deepseek_v4_parser.py create mode 100644 vllm_mlx/tool_parsers/deepseek_v4_tool_parser.py create mode 100644 vllm_mlx/utils/deepseek_v4_encoding.py diff --git a/benchmarks/bench_deepseek_v4.py b/benchmarks/bench_deepseek_v4.py new file mode 100644 index 000000000..c418ea4a6 --- /dev/null +++ b/benchmarks/bench_deepseek_v4.py @@ -0,0 +1,214 @@ +"""Benchmark: DeepSeek-V4 prompt encoding and streaming parser overhead. + +Covers both serving paths. Single-stream is one sequence decoded token by +token, where the parser sits directly in the latency path. Batch is N +concurrent sequences interleaved, the shape continuous batching produces: each +request carries its own parser state, so per-token cost is paid N times per +decode step and any per-token work that scales with output length compounds. + +The thing to watch for is quadratic scaling. A parser that rescans the whole +accumulated text on every delta is O(N²) over a generation, which is invisible +on short replies and dominant on long ones — so the ms/tok column is the one +that matters, not the totals. It is currently flat for the DSML tool parser and +grows for the reasoning path; see the notes the script prints at the end. + +Usage: + python benchmarks/bench_deepseek_v4.py +""" + +import time + +from vllm_mlx.reasoning.deepseek_v4_parser import DeepSeekV4ReasoningParser +from vllm_mlx.tool_parsers.deepseek_v4_tool_parser import DeepSeekV4ToolParser +from vllm_mlx.utils.deepseek_v4_encoding import apply_chat_template + +D = "｜DSML｜" + +TOOLS = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather for a city.", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string"}, + "days": {"type": "integer"}, + }, + "required": ["city"], + }, + }, + } +] + + +def reasoning_tokens(n: int) -> list[str]: + """A plain thinking turn: N reasoning tokens, then a short answer.""" + return ( + [f"step{i} " for i in range(n)] + + [""] + + [f"word{i} " for i in range(10)] + ) + + +def tool_call_tokens(n: int) -> list[str]: + """A tool-calling turn: N reasoning tokens, then DSML markup. + + The markup is split the way the tokenizer splits it — ｜DSML｜ has an id of + its own, the surrounding punctuation does not — so the parser sees the same + fragmentation it sees in production. + """ + tokens = [f"step{i} " for i in range(n)] + ["", "\n\n"] + tokens += ["<", D, "tool_calls", ">", "\n"] + tokens += ["<", D, "invoke", ' name="get_weather"', ">", "\n"] + tokens += [ + "<", + D, + "parameter", + ' name="city"', + ' string="true"', + ">", + "Prague", + "", + "\n", + ] + tokens += ["", "\n"] + tokens += [""] + return tokens + + +def bench_stream(make_tokens, n_tokens: int, streams: int) -> tuple[float, int]: + """Interleave `streams` concurrent sequences, one delta each per step. + + Returns (total ms, total deltas). With streams=1 this is the single-stream + latency path; above that it is the batched decode step. + """ + token_lists = [make_tokens(n_tokens) for _ in range(streams)] + parsers = [] + for _ in range(streams): + reasoner = DeepSeekV4ReasoningParser() + reasoner.reset_state() + tools = DeepSeekV4ToolParser() + tools.reset() + parsers.append((reasoner, tools, {"acc": "", "tool_acc": ""})) + + steps = max(len(t) for t in token_lists) + deltas = 0 + start = time.perf_counter() + for step in range(steps): + for stream in range(streams): + tokens = token_lists[stream] + if step >= len(tokens): + continue + reasoner, tools, state = parsers[stream] + delta = tokens[step] + previous, state["acc"] = state["acc"], state["acc"] + delta + deltas += 1 + + message = reasoner.extract_reasoning_streaming( + previous, state["acc"], delta + ) + if message is None or not message.content: + continue + prev_tool = state["tool_acc"] + state["tool_acc"] = prev_tool + message.content + tools.extract_tool_calls_streaming( + prev_tool, state["tool_acc"], message.content + ) + return (time.perf_counter() - start) * 1000, deltas + + +def bench_tool_parser_only(make_tokens, n_tokens: int) -> tuple[float, int]: + """The DSML parser without the reasoning parser in front of it. + + Isolates how much of the per-token cost is the tool parser's own work. + """ + tokens = make_tokens(n_tokens) + parser = DeepSeekV4ToolParser() + parser.reset() + + accumulated = "" + start = time.perf_counter() + for delta in tokens: + previous, accumulated = accumulated, accumulated + delta + parser.extract_tool_calls_streaming(previous, accumulated, delta) + return (time.perf_counter() - start) * 1000, len(tokens) + + +def bench_encoder(turns: int, repeats: int = 200) -> float: + """Prompt build cost for a conversation of `turns` user/assistant pairs.""" + conversation = [{"role": "system", "content": "You are a helpful assistant."}] + for i in range(turns): + conversation.append({"role": "user", "content": f"Question number {i}?"}) + conversation.append( + { + "role": "assistant", + "content": f"Answer number {i}.", + "reasoning_content": f"Thinking about question {i} at some length.", + } + ) + conversation.append({"role": "user", "content": "And finally?"}) + + start = time.perf_counter() + for _ in range(repeats): + apply_chat_template(conversation, tools=TOOLS) + return (time.perf_counter() - start) * 1000 / repeats + + +def main(): + print("DeepSeek-V4 encoder and parser benchmark") + print("=" * 68) + + print("\nPrompt encoding (per call, tools attached)") + for turns in (1, 4, 16, 64): + ms = bench_encoder(turns) + print(f" {turns * 2 + 2:>4} messages -> {ms:>8.3f} ms") + + for label, make_tokens in ( + ("plain thinking turn", reasoning_tokens), + ("tool-calling turn", tool_call_tokens), + ): + print(f"\nSingle stream, {label}") + for n in (100, 500, 1000, 2000, 5000): + ms, deltas = bench_stream(make_tokens, n, streams=1) + print( + f" {n:>5} reasoning tokens -> {ms:>8.2f} ms total, " + f"{ms / deltas:>7.4f} ms/tok" + ) + + print("\nDSML tool parser alone, tool-calling turn") + for n in (100, 500, 1000, 2000, 5000): + ms, deltas = bench_tool_parser_only(tool_call_tokens, n) + print( + f" {n:>5} reasoning tokens -> {ms:>8.2f} ms total, " + f"{ms / deltas:>7.4f} ms/tok" + ) + + print("\nBatched decode, tool-calling turn, 1000 reasoning tokens each") + for streams in (1, 2, 4, 8, 16): + ms, deltas = bench_stream(tool_call_tokens, 1000, streams=streams) + print( + f" {streams:>3} concurrent -> {ms:>8.2f} ms total, " + f"{ms / deltas:>7.4f} ms/tok" + ) + + print("\nReading the numbers:") + print(" At 50 tok/s the per-token budget is 20 ms, so anything under") + print(" 0.1 ms/tok is noise. Batching adds no per-token cost — each") + print(" request carries independent parser state and the totals scale") + print(" linearly with the number of streams.") + print() + print(" The single-stream ms/tok does grow with output length. That comes") + print(" from BaseThinkingReasoningParser, which searches the accumulated") + print(" text for its start and end tags on every delta while the reasoning") + print(" block is open; the DSML tool parser on its own stays flat. It is") + print(" 0.2% of the decode budget even at 5000 tokens, but it is quadratic,") + print(" so it is worth fixing in the base class rather than per model.") + + +if __name__ == "__main__": + main() diff --git a/tests/deepseek_v4_golden_prompts.py b/tests/deepseek_v4_golden_prompts.py new file mode 100644 index 000000000..5a73dfba1 --- /dev/null +++ b/tests/deepseek_v4_golden_prompts.py @@ -0,0 +1,79 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Golden prompts for the DeepSeek-V4 encoder. + +Generated from the reference ``encoding_dsv4.py`` published with the model +weights. The reference is not vendored into this repository, so its output is +frozen here: any drift in our encoder shows up as a failing equality assertion +rather than as a subtly malformed prompt at inference time. + +Regenerate only against the reference implementation, never by pasting in what +our own encoder currently produces. +""" + +GOLDEN_PROMPTS = { + "single_user_thinking": ( + '[{"role": "user", "content": "What is 2+2?"}]', + {"thinking_mode": "thinking"}, + "<｜begin▁of▁sentence｜><｜User｜>What is 2+2?<｜Assistant｜>", + ), + "single_user_chat": ( + '[{"role": "user", "content": "What is 2+2?"}]', + {"thinking_mode": "chat"}, + "<｜begin▁of▁sentence｜><｜User｜>What is 2+2?<｜Assistant｜>", + ), + "system_and_user": ( + '[{"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "What is 2+2?"}]', + {"thinking_mode": "thinking"}, + "<｜begin▁of▁sentence｜>You are a helpful assistant.<｜User｜>What is 2+2?<｜Assistant｜>", + ), + "effort_high": ( + '[{"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "What is 2+2?"}]', + {"thinking_mode": "thinking", "reasoning_effort": "high"}, + "<｜begin▁of▁sentence｜>Reasoning Effort: Absolute maximum with no shortcuts permitted.\nYou MUST be very thorough in your thinking and comprehensively decompose the problem to resolve the root cause, rigorously stress-testing your logic against all potential paths, edge cases, and adversarial scenarios.\nExplicitly write out your entire deliberation process, documenting every intermediate step, considered alternative, and rejected hypothesis to ensure absolutely no assumption is left unchecked.\n\nYou are a helpful assistant.<｜User｜>What is 2+2?<｜Assistant｜>", + ), + "effort_max": ( + '[{"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "What is 2+2?"}]', + {"thinking_mode": "thinking", "reasoning_effort": "max"}, + "<｜begin▁of▁sentence｜>Reasoning Effort: Beyond maximum — exhaustive, relentless, and uncompromising.\nYou MUST reason with the utmost depth and rigor, leaving absolutely nothing to chance: exhaustively decompose the problem into its most fundamental components, trace every causal chain to its root, and resolve the underlying cause rather than any surface symptom.\nDo not stop reasoning until you have independently verified the solution from multiple angles and are certain that no assumption remains unchecked and no error remains undiscovered.\n\nYou are a helpful assistant.<｜User｜>What is 2+2?<｜Assistant｜>", + ), + "effort_low": ( + '[{"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "What is 2+2?"}]', + {"thinking_mode": "thinking", "reasoning_effort": "low"}, + "<｜begin▁of▁sentence｜>You are a helpful assistant.<｜User｜>What is 2+2?<｜Assistant｜>", + ), + "effort_ignored_in_chat": ( + '[{"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "What is 2+2?"}]', + {"thinking_mode": "chat", "reasoning_effort": "max"}, + "<｜begin▁of▁sentence｜>You are a helpful assistant.<｜User｜>What is 2+2?<｜Assistant｜>", + ), + "multiturn_drop": ( + '[{"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "What is 2+2?"}, {"role": "assistant", "content": "4", "reasoning_content": "2 plus 2 is 4."}, {"role": "user", "content": "And 3+3?"}]', + {"thinking_mode": "thinking"}, + "<｜begin▁of▁sentence｜>You are a helpful assistant.<｜User｜>What is 2+2?<｜Assistant｜>4<｜end▁of▁sentence｜><｜User｜>And 3+3?<｜Assistant｜>", + ), + "multiturn_keep": ( + '[{"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "What is 2+2?"}, {"role": "assistant", "content": "4", "reasoning_content": "2 plus 2 is 4."}, {"role": "user", "content": "And 3+3?"}]', + {"thinking_mode": "thinking", "drop_thinking": False}, + "<｜begin▁of▁sentence｜>You are a helpful assistant.<｜User｜>What is 2+2?<｜Assistant｜>2 plus 2 is 4.4<｜end▁of▁sentence｜><｜User｜>And 3+3?<｜Assistant｜>", + ), + "no_bos": ( + '[{"role": "user", "content": "What is 2+2?"}]', + {"thinking_mode": "thinking", "add_default_bos_token": False}, + "<｜User｜>What is 2+2?<｜Assistant｜>", + ), + "tools_declared": ( + '[{"role": "system", "content": "You are a helpful assistant.", "tools": [{"type": "function", "function": {"name": "get_weather", "description": "Get the weather for a city", "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}}}]}, {"role": "user", "content": "What is 2+2?"}]', + {"thinking_mode": "thinking"}, + '<｜begin▁of▁sentence｜>You are a helpful assistant.\n\n## Tools\n\nYou have access to a set of tools to help answer the user\'s question. You can invoke tools by writing a "<｜DSML｜tool_calls>" block like the following:\n\n<｜DSML｜tool_calls>\n<｜DSML｜invoke name="$TOOL_NAME">\n<｜DSML｜parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE\n...\n\n<｜DSML｜invoke name="$TOOL_NAME2">\n...\n\n\n\nString parameters should be specified as is and set `string="true"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string="false"`.\n\nIf thinking_mode is enabled (triggered by ), you MUST output your complete reasoning inside ... BEFORE any tool calls or final response.\n\nOtherwise, output directly after with tool calls or final response.\n\n### Available Tool Schemas\n\n{"name": "get_weather", "description": "Get the weather for a city", "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}}\n\nYou MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.\n<｜User｜>What is 2+2?<｜Assistant｜>', + ), + "tool_roundtrip": ( + '[{"role": "system", "content": "You are a helpful assistant.", "tools": [{"type": "function", "function": {"name": "get_weather", "description": "Get the weather for a city", "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}}}]}, {"role": "user", "content": "What is 2+2?"}, {"role": "assistant", "content": "", "reasoning_content": "I need the weather.", "tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "get_weather", "arguments": "{\\"city\\": \\"Prague\\", \\"days\\": 3}"}}]}, {"role": "tool", "tool_call_id": "call_1", "content": "{\\"temp\\": 21}"}]', + {"thinking_mode": "thinking"}, + '<｜begin▁of▁sentence｜>You are a helpful assistant.\n\n## Tools\n\nYou have access to a set of tools to help answer the user\'s question. You can invoke tools by writing a "<｜DSML｜tool_calls>" block like the following:\n\n<｜DSML｜tool_calls>\n<｜DSML｜invoke name="$TOOL_NAME">\n<｜DSML｜parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE\n...\n\n<｜DSML｜invoke name="$TOOL_NAME2">\n...\n\n\n\nString parameters should be specified as is and set `string="true"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string="false"`.\n\nIf thinking_mode is enabled (triggered by ), you MUST output your complete reasoning inside ... BEFORE any tool calls or final response.\n\nOtherwise, output directly after with tool calls or final response.\n\n### Available Tool Schemas\n\n{"name": "get_weather", "description": "Get the weather for a city", "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}}\n\nYou MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.\n<｜User｜>What is 2+2?<｜Assistant｜>I need the weather.\n\n<｜DSML｜tool_calls>\n<｜DSML｜invoke name="get_weather">\n<｜DSML｜parameter name="city" string="true">Prague\n<｜DSML｜parameter name="days" string="false">3\n\n<｜end▁of▁sentence｜><｜User｜>{"temp": 21}<｜Assistant｜>', + ), + "two_tool_results_merge": ( + '[{"role": "system", "content": "You are a helpful assistant.", "tools": [{"type": "function", "function": {"name": "get_weather", "description": "Get the weather for a city", "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}}}]}, {"role": "user", "content": "What is 2+2?"}, {"role": "assistant", "content": "", "reasoning_content": "I need the weather.", "tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "get_weather", "arguments": "{\\"city\\": \\"Prague\\", \\"days\\": 3}"}}]}, {"role": "tool", "tool_call_id": "call_1", "content": "{\\"temp\\": 21}"}, {"role": "tool", "tool_call_id": "call_1", "content": "{\\"temp\\": 21}"}]', + {"thinking_mode": "thinking"}, + '<｜begin▁of▁sentence｜>You are a helpful assistant.\n\n## Tools\n\nYou have access to a set of tools to help answer the user\'s question. You can invoke tools by writing a "<｜DSML｜tool_calls>" block like the following:\n\n<｜DSML｜tool_calls>\n<｜DSML｜invoke name="$TOOL_NAME">\n<｜DSML｜parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE\n...\n\n<｜DSML｜invoke name="$TOOL_NAME2">\n...\n\n\n\nString parameters should be specified as is and set `string="true"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string="false"`.\n\nIf thinking_mode is enabled (triggered by ), you MUST output your complete reasoning inside ... BEFORE any tool calls or final response.\n\nOtherwise, output directly after with tool calls or final response.\n\n### Available Tool Schemas\n\n{"name": "get_weather", "description": "Get the weather for a city", "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}}\n\nYou MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.\n<｜User｜>What is 2+2?<｜Assistant｜>I need the weather.\n\n<｜DSML｜tool_calls>\n<｜DSML｜invoke name="get_weather">\n<｜DSML｜parameter name="city" string="true">Prague\n<｜DSML｜parameter name="days" string="false">3\n\n<｜end▁of▁sentence｜><｜User｜>{"temp": 21}\n\n{"temp": 21}<｜Assistant｜>', + ), +} diff --git a/tests/test_deepseek_v4_encoding.py b/tests/test_deepseek_v4_encoding.py new file mode 100644 index 000000000..0b0b3e1e1 --- /dev/null +++ b/tests/test_deepseek_v4_encoding.py @@ -0,0 +1,329 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the DeepSeek-V4 prompt encoder. + +DeepSeek-V4 has no Jinja chat template — the prompt is built programmatically — +so these tests assert against prompts frozen from the reference implementation +shipped with the model weights (see ``deepseek_v4_golden_prompts.py``). A silent +divergence here would not raise anywhere; it would just degrade generation. +""" + +import json + +import pytest + +from vllm_mlx.utils.deepseek_v4_encoding import ( + BOS_TOKEN, + THINKING_END_TOKEN, + THINKING_START_TOKEN, + apply_chat_template, + encode_arguments_to_dsml, + encode_messages, + install, + merge_tool_messages, + resolve_thinking, +) + +from .deepseek_v4_golden_prompts import GOLDEN_PROMPTS + +TOOLS = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather for a city", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + } +] + + +@pytest.mark.parametrize("name", sorted(GOLDEN_PROMPTS)) +def test_matches_reference_implementation(name): + """Our encoder reproduces the reference byte for byte.""" + messages_json, kwargs, expected = GOLDEN_PROMPTS[name] + assert encode_messages(json.loads(messages_json), **kwargs) == expected + + +class TestPromptShape: + def test_thinking_mode_opens_reasoning(self): + prompt = encode_messages( + [{"role": "user", "content": "Hi"}], thinking_mode="thinking" + ) + assert prompt.endswith(THINKING_START_TOKEN) + assert prompt.startswith(BOS_TOKEN) + + def test_chat_mode_closes_reasoning(self): + """Chat mode pre-closes so the model skips reasoning.""" + prompt = encode_messages( + [{"role": "user", "content": "Hi"}], thinking_mode="chat" + ) + assert prompt.endswith(THINKING_END_TOKEN) + + def test_system_message_is_bare_text(self): + """There is no system wrapper token — content follows BOS directly.""" + prompt = encode_messages( + [{"role": "system", "content": "SYS"}, {"role": "user", "content": "Hi"}], + thinking_mode="thinking", + ) + assert prompt.startswith(BOS_TOKEN + "SYS<｜User｜>Hi") + + def test_bos_can_be_suppressed(self): + prompt = encode_messages( + [{"role": "user", "content": "Hi"}], + thinking_mode="thinking", + add_default_bos_token=False, + ) + assert not prompt.startswith(BOS_TOKEN) + + def test_invalid_thinking_mode_rejected(self): + with pytest.raises(ValueError, match="thinking_mode"): + encode_messages([{"role": "user", "content": "Hi"}], thinking_mode="nope") + + +class TestReasoningEffort: + @pytest.mark.parametrize( + "effort,marker", + [("high", "Absolute maximum"), ("max", "Beyond maximum")], + ) + def test_prefix_is_prepended(self, effort, marker): + prompt = encode_messages( + [{"role": "user", "content": "Hi"}], + thinking_mode="thinking", + reasoning_effort=effort, + ) + assert marker in prompt + # The prefix sits after BOS but before the conversation. + assert prompt.index(marker) < prompt.index("<｜User｜>") + + def test_low_adds_nothing(self): + prompt = encode_messages( + [{"role": "user", "content": "Hi"}], + thinking_mode="thinking", + reasoning_effort="low", + ) + assert "Reasoning Effort" not in prompt + + def test_no_prefix_in_chat_mode(self): + """Effort only shapes reasoning, and chat mode has none.""" + prompt = encode_messages( + [{"role": "user", "content": "Hi"}], + thinking_mode="chat", + reasoning_effort="max", + ) + assert "Reasoning Effort" not in prompt + + def test_unknown_effort_rejected_by_encoder(self): + with pytest.raises(ValueError, match="reasoning effort"): + encode_messages( + [{"role": "user", "content": "Hi"}], + thinking_mode="thinking", + reasoning_effort="turbo", + ) + + +class TestResolveThinking: + @pytest.mark.parametrize( + "kwargs,expected", + [ + ({}, ("thinking", None)), + ({"enable_thinking": False}, ("chat", None)), + ({"enable_thinking": True}, ("thinking", None)), + ({"reasoning_effort": "none"}, ("chat", None)), + ({"reasoning_effort": "low"}, ("thinking", "low")), + ({"reasoning_effort": "medium"}, ("thinking", "high")), + ({"reasoning_effort": "high"}, ("thinking", "high")), + ({"reasoning_effort": "max"}, ("thinking", "max")), + ({"reasoning_effort": "xhigh"}, ("thinking", "max")), + ({"thinking_mode": "chat"}, ("chat", None)), + ({"enable_thinking": False, "reasoning_effort": "max"}, ("chat", None)), + ], + ) + def test_mapping(self, kwargs, expected): + assert resolve_thinking(**kwargs) == expected + + def test_unknown_effort_falls_back_to_high(self): + """A typo must not silently disable reasoning.""" + assert resolve_thinking(reasoning_effort="turbo") == ("thinking", "high") + + def test_invalid_thinking_mode_rejected(self): + with pytest.raises(ValueError, match="thinking_mode"): + resolve_thinking(thinking_mode="bogus") + + +class TestTools: + def test_schema_lands_in_system_message(self): + prompt = apply_chat_template( + [{"role": "system", "content": "SYS"}, {"role": "user", "content": "Hi"}], + tools=TOOLS, + ) + assert "## Tools" in prompt + assert "get_weather" in prompt + assert prompt.index("## Tools") < prompt.index("<｜User｜>") + + def test_system_message_synthesised_when_absent(self): + prompt = apply_chat_template([{"role": "user", "content": "Hi"}], tools=TOOLS) + assert "## Tools" in prompt + + def test_existing_declaration_wins(self): + """A conversation that already declares tools is left alone.""" + conversation = [ + {"role": "system", "content": "SYS", "tools": TOOLS}, + {"role": "user", "content": "Hi"}, + ] + assert apply_chat_template(conversation) == apply_chat_template( + conversation, tools=TOOLS + ) + + def test_tools_keep_reasoning_history(self): + """drop_thinking is forced off when tools are in play. + + The model has to see why it made the earlier calls, so the reasoning + that produced them must survive into the next turn. + """ + messages = [ + {"role": "system", "content": "S", "tools": TOOLS}, + {"role": "user", "content": "Q"}, + {"role": "assistant", "content": "A", "reasoning_content": "BECAUSE"}, + {"role": "user", "content": "Q2"}, + ] + prompt = encode_messages(messages, thinking_mode="thinking", drop_thinking=True) + assert "BECAUSE" in prompt + + def test_reasoning_dropped_without_tools(self): + messages = [ + {"role": "system", "content": "S"}, + {"role": "user", "content": "Q"}, + {"role": "assistant", "content": "A", "reasoning_content": "STALE"}, + {"role": "user", "content": "Q2"}, + ] + prompt = encode_messages(messages, thinking_mode="thinking", drop_thinking=True) + assert "STALE" not in prompt + + +class TestToolResults: + def test_tool_role_becomes_user_block(self): + """V4 has no tool role — results ride inside the user turn.""" + merged = merge_tool_messages( + [ + {"role": "user", "content": "Q"}, + {"role": "assistant", "content": "", "tool_calls": []}, + {"role": "tool", "tool_call_id": "c1", "content": "RESULT"}, + ] + ) + assert not any(m["role"] == "tool" for m in merged) + assert merged[-1]["role"] == "user" + assert merged[-1]["content_blocks"][0]["content"] == "RESULT" + + def test_result_rendered_as_tool_result_tag(self): + prompt = encode_messages( + [ + {"role": "user", "content": "Q"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "c1", + "type": "function", + "function": {"name": "f", "arguments": "{}"}, + } + ], + }, + {"role": "tool", "tool_call_id": "c1", "content": "42"}, + ], + thinking_mode="thinking", + ) + assert "42" in prompt + + def test_arguments_accepted_as_decoded_mapping(self): + """``api/utils.py`` decodes ``arguments`` in place for native format. + + Re-decoding an already-decoded mapping used to collapse every parameter + into a single bogus ``arguments`` entry, so the model saw a malformed + call in its own history. + """ + as_string = encode_arguments_to_dsml( + {"name": "f", "arguments": json.dumps({"city": "Prague", "days": 3})} + ) + as_mapping = encode_arguments_to_dsml( + {"name": "f", "arguments": {"city": "Prague", "days": 3}} + ) + assert as_mapping == as_string + assert 'name="city" string="true"' in as_mapping + assert 'name="days" string="false"' in as_mapping + + def test_unparseable_arguments_are_not_dropped(self): + rendered = encode_arguments_to_dsml({"name": "f", "arguments": "not json"}) + assert 'name="arguments"' in rendered + assert "not json" in rendered + + def test_results_sorted_by_call_order(self): + """Clients may answer out of order; the model expects call order.""" + messages = [ + {"role": "user", "content": "Q"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "first", + "type": "function", + "function": {"name": "f", "arguments": "{}"}, + }, + { + "id": "second", + "type": "function", + "function": {"name": "g", "arguments": "{}"}, + }, + ], + }, + {"role": "tool", "tool_call_id": "second", "content": "SECOND"}, + {"role": "tool", "tool_call_id": "first", "content": "FIRST"}, + ] + prompt = encode_messages(messages, thinking_mode="thinking") + assert prompt.index("FIRST") < prompt.index("SECOND") + + +class TestInstall: + class _FakeTokenizer: + def __init__(self): + self.chat_template = None + + def encode(self, text): + return [len(text)] + + def test_overrides_apply_chat_template(self): + tokenizer = self._FakeTokenizer() + install(tokenizer) + prompt = tokenizer.apply_chat_template([{"role": "user", "content": "Hi"}]) + assert prompt.startswith(BOS_TOKEN) + assert prompt.endswith(THINKING_START_TOKEN) + + def test_is_idempotent(self): + tokenizer = self._FakeTokenizer() + install(tokenizer) + first = tokenizer.apply_chat_template + install(tokenizer) + assert tokenizer.apply_chat_template is first + + def test_tokenize_returns_ids(self): + tokenizer = self._FakeTokenizer() + install(tokenizer) + assert tokenizer.apply_chat_template( + [{"role": "user", "content": "Hi"}], tokenize=True + ) == [len(tokenizer.apply_chat_template([{"role": "user", "content": "Hi"}]))] + + def test_tolerates_unknown_kwargs(self): + """Callers pass add_generation_prompt and friends; none may blow up.""" + tokenizer = self._FakeTokenizer() + install(tokenizer) + prompt = tokenizer.apply_chat_template( + [{"role": "user", "content": "Hi"}], + add_generation_prompt=True, + some_future_kwarg=123, + ) + assert isinstance(prompt, str) diff --git a/tests/test_deepseek_v4_reasoning.py b/tests/test_deepseek_v4_reasoning.py new file mode 100644 index 000000000..c70858302 --- /dev/null +++ b/tests/test_deepseek_v4_reasoning.py @@ -0,0 +1,141 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the DeepSeek-V4 reasoning parser. + +V4 shares ````/```` with R1, so the interesting cases are the +ones R1 does not have: the prompt closes on ```` (leaving the opening tag +out of the output), and a tool call implicitly ends the reasoning block. +""" + +import pytest + +from vllm_mlx.reasoning import get_parser, list_parsers +from vllm_mlx.reasoning.deepseek_v4_parser import DeepSeekV4ReasoningParser + +D = "｜DSML｜" +TOOL_CALLS_START = f"<{D}tool_calls>" + +TOOL_BLOCK = ( + f"{TOOL_CALLS_START}\n" + f'<{D}invoke name="f">\n' + f'<{D}parameter name="a" string="true">1\n' + f"\n" + f"" +) + + +@pytest.fixture +def parser(): + p = DeepSeekV4ReasoningParser() + p.reset_state() + return p + + +class TestRegistration: + def test_registered(self): + assert "deepseek_v4" in list_parsers() + + def test_resolves_to_v4_parser(self): + assert get_parser("deepseek_v4") is DeepSeekV4ReasoningParser + + +class TestExtraction: + def test_implicit_opening_tag(self, parser): + """The encoder ends the prompt with , so output starts inside it.""" + reasoning, content = parser.extract_reasoning("weighing itAnswer.") + assert reasoning == "weighing it" + assert content == "Answer." + + def test_explicit_tags(self, parser): + reasoning, content = parser.extract_reasoning( + "weighing itAnswer." + ) + assert reasoning == "weighing it" + assert content == "Answer." + + def test_no_tags_is_pure_content(self, parser): + reasoning, content = parser.extract_reasoning("Just an answer.") + assert reasoning is None + assert content == "Just an answer." + + +class TestToolCallInteraction: + def test_tool_call_closes_unterminated_reasoning(self, parser): + """Without this the whole DSML payload would be swallowed as reasoning.""" + reasoning, content = parser.extract_reasoning(f"I should call it{TOOL_BLOCK}") + assert reasoning == "I should call it" + assert content.startswith(TOOL_CALLS_START) + + def test_markup_is_left_for_the_tool_parser(self, parser): + _, content = parser.extract_reasoning(f"thinking{TOOL_BLOCK}") + assert TOOL_BLOCK in content + + def test_explicit_open_tag_with_tool_call(self, parser): + reasoning, content = parser.extract_reasoning(f"thinking{TOOL_BLOCK}") + assert reasoning == "thinking" + assert content.startswith(TOOL_CALLS_START) + + def test_closed_reasoning_before_tool_call(self, parser): + """When came first it wins — the split is already unambiguous.""" + reasoning, content = parser.extract_reasoning( + f"thinkingCalling.\n\n{TOOL_BLOCK}" + ) + assert reasoning == "thinking" + assert content.startswith("Calling.") + assert TOOL_CALLS_START in content + + +class TestStreaming: + @staticmethod + def _stream(text, chunk): + parser = DeepSeekV4ReasoningParser() + parser.reset_state() + previous, reasoning, content = "", [], [] + for i in range(0, len(text), chunk): + delta = text[i : i + chunk] + current = previous + delta + message = parser.extract_reasoning_streaming(previous, current, delta) + if message is not None: + if message.reasoning: + reasoning.append(message.reasoning) + if message.content: + content.append(message.content) + previous = current + return "".join(reasoning), "".join(content) + + @pytest.mark.parametrize("chunk", [1, 3, 7, 16, 64]) + def test_split_matches_non_streaming(self, chunk): + text = "weighing itAnswer." + reasoning, content = self._stream(text, chunk) + expected_reasoning, expected_content = ( + DeepSeekV4ReasoningParser().extract_reasoning(text) + ) + assert reasoning.strip() == (expected_reasoning or "") + assert content.strip() == (expected_content or "") + + @pytest.mark.parametrize("chunk", [1, 3, 7, 16, 64]) + def test_tool_call_ends_reasoning_mid_stream(self, chunk): + reasoning, content = self._stream(f"deciding{TOOL_BLOCK}", chunk) + assert reasoning.strip() == "deciding" + assert TOOL_CALLS_START in content + assert D not in reasoning + + @pytest.mark.parametrize("chunk", [1, 3, 7, 16, 64]) + def test_nothing_is_lost(self, chunk): + """Every character must land in exactly one of the two channels.""" + text = f"deciding{TOOL_BLOCK}" + reasoning, content = self._stream(text, chunk) + assert len(reasoning) + len(content) == len(text) + + @pytest.mark.parametrize("chunk", [1, 2, 3, 4, 5, 6, 7]) + def test_split_think_tag_does_not_leak(self, chunk): + """A split ```` must not leak fragments into reasoning. + + V4 has its own ids for the think tags, so they normally arrive whole; + this pins the behaviour for any detokenizer that splits them, which is + the same failure mode the multi-token DSML marker hits for real. + """ + reasoning, content = self._stream("weighing itAnswer.", chunk) + assert reasoning.strip() == "weighing it" + assert content.strip() == "Answer." + for fragment in (""): + assert fragment not in reasoning diff --git a/tests/test_deepseek_v4_tool_parser.py b/tests/test_deepseek_v4_tool_parser.py new file mode 100644 index 000000000..a5782faa6 --- /dev/null +++ b/tests/test_deepseek_v4_tool_parser.py @@ -0,0 +1,420 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the DeepSeek-V4 DSML tool call parser.""" + +import json + +import pytest + +from vllm_mlx.tool_parsers import ToolParserManager +from vllm_mlx.tool_parsers.deepseek_v4_tool_parser import ( + TOOL_CALLS_END, + TOOL_CALLS_START, + DeepSeekV4ToolParser, +) + +D = "｜DSML｜" + + +def invoke(name: str, *params: str) -> str: + body = "\n".join(params) + return f'<{D}invoke name="{name}">\n{body}\n' + + +def param(name: str, value: str, *, is_str: bool) -> str: + flag = "true" if is_str else "false" + return f'<{D}parameter name="{name}" string="{flag}">{value}' + + +def block(*invokes: str) -> str: + return f"{TOOL_CALLS_START}\n" + "\n".join(invokes) + f"\n{TOOL_CALLS_END}" + + +@pytest.fixture +def parser(): + p = DeepSeekV4ToolParser() + p.reset() + return p + + +def args_of(result, index=0): + return json.loads(result.tool_calls[index]["arguments"]) + + +class TestRegistration: + @pytest.mark.parametrize("name", ["deepseek_v4", "dsml"]) + def test_registered_under(self, name): + assert ToolParserManager.get_tool_parser(name) is DeepSeekV4ToolParser + + def test_distinct_from_v3_parser(self): + """V3/R1 use <｜tool▁calls▁begin｜> + fenced JSON; V4 shares nothing.""" + v3 = ToolParserManager.get_tool_parser("deepseek") + assert v3 is not DeepSeekV4ToolParser + + def test_declares_native_tool_format(self): + """The encoder consumes role="tool" and tool_calls directly. + + Declaring otherwise makes the server flatten them into + "[Tool Result (id)]: ..." and "[Calling tool: name(...)]" before the + encoder ever runs, so the model would see a shape it was never trained + on instead of blocks and DSML. + """ + assert DeepSeekV4ToolParser.supports_native_format() is True + + +class TestExtraction: + def test_single_call(self, parser): + text = "Let me check.\n\n" + block( + invoke("get_weather", param("city", "Prague", is_str=True)) + ) + result = parser.extract_tool_calls(text) + assert result.tools_called + assert result.content == "Let me check." + assert result.tool_calls[0]["name"] == "get_weather" + assert args_of(result) == {"city": "Prague"} + + def test_parallel_calls(self, parser): + text = block( + invoke("search", param("q", "mlx", is_str=True)), + invoke("get_weather", param("city", "Brno", is_str=True)), + ) + result = parser.extract_tool_calls(text) + assert [tc["name"] for tc in result.tool_calls] == ["search", "get_weather"] + assert args_of(result, 0) == {"q": "mlx"} + assert args_of(result, 1) == {"city": "Brno"} + + def test_call_ids_are_unique(self, parser): + text = block( + invoke("a", param("x", "1", is_str=True)), + invoke("b", param("x", "2", is_str=True)), + ) + result = parser.extract_tool_calls(text) + assert result.tool_calls[0]["id"] != result.tool_calls[1]["id"] + + def test_no_markup_is_plain_content(self, parser): + result = parser.extract_tool_calls("Just an answer.") + assert not result.tools_called + assert result.content == "Just an answer." + + def test_content_after_reasoning_only(self, parser): + """Reasoning belongs to the reasoning parser, not to content.""" + text = "deliberatingChecking.\n\n" + block( + invoke("f", param("x", "1", is_str=True)) + ) + result = parser.extract_tool_calls(text) + assert result.content == "Checking." + + +class TestParameterTyping: + """``string="true"`` is verbatim, ``string="false"`` is JSON.""" + + def test_string_parameter_kept_raw(self, parser): + result = parser.extract_tool_calls( + block(invoke("f", param("s", "42", is_str=True))) + ) + assert args_of(result) == {"s": "42"} + + def test_non_string_parameter_decoded(self, parser): + result = parser.extract_tool_calls( + block(invoke("f", param("n", "42", is_str=False))) + ) + assert args_of(result) == {"n": 42} + + @pytest.mark.parametrize( + "raw,expected", + [ + ("3", 3), + ("3.5", 3.5), + ("true", True), + ("false", False), + ("null", None), + ("[1, 2]", [1, 2]), + ('{"a": {"b": [1]}}', {"a": {"b": [1]}}), + ], + ) + def test_json_value_kinds(self, parser, raw, expected): + result = parser.extract_tool_calls( + block(invoke("f", param("v", raw, is_str=False))) + ) + assert args_of(result) == {"v": expected} + + def test_string_value_may_contain_markup_like_text(self, parser): + """A string parameter is why this cannot be a regex over name=value.""" + tricky = 'has "quotes" & and {"json": 1}' + result = parser.extract_tool_calls( + block(invoke("f", param("s", tricky, is_str=True))) + ) + assert args_of(result) == {"s": tricky} + + def test_unparseable_json_falls_back_to_raw(self, parser): + result = parser.extract_tool_calls( + block(invoke("f", param("v", "{not json", is_str=False))) + ) + assert args_of(result) == {"v": "{not json"} + + def test_multiple_parameters(self, parser): + result = parser.extract_tool_calls( + block( + invoke( + "f", + param("a", "x", is_str=True), + param("b", "2", is_str=False), + param("c", "[true]", is_str=False), + ) + ) + ) + assert args_of(result) == {"a": "x", "b": 2, "c": [True]} + + def test_call_without_parameters(self, parser): + result = parser.extract_tool_calls( + f"{TOOL_CALLS_START}\n" + f'<{D}invoke name="ping">\n' + f"\n{TOOL_CALLS_END}" + ) + assert result.tools_called + assert args_of(result) == {} + + +class TestMalformedInput: + """Half-generated markup must degrade, never raise.""" + + def test_truncated_after_start_marker(self, parser): + result = parser.extract_tool_calls(f"text{TOOL_CALLS_START}\n") + assert not result.tools_called + assert result.content + + def test_truncated_mid_invoke_header(self, parser): + result = parser.extract_tool_calls(f'{TOOL_CALLS_START}\n<{D}invoke name="f') + assert not result.tools_called + + def test_truncated_parameter_is_dropped(self, parser): + text = ( + f"{TOOL_CALLS_START}\n" + f'<{D}invoke name="f">\n' + f'<{D}parameter name="a" string="true">x\n' + f'<{D}parameter name="b" string="true">unterminated' + ) + result = parser.extract_tool_calls(text) + assert result.tools_called + assert args_of(result) == {"a": "x"} + + def test_missing_end_marker_still_parses(self, parser): + text = f"{TOOL_CALLS_START}\n" + invoke("f", param("a", "1", is_str=True)) + result = parser.extract_tool_calls(text) + assert result.tools_called + assert args_of(result) == {"a": "1"} + + +class TestStreaming: + """Deltas arrive token-sized, so every marker can straddle a boundary.""" + + FULL = "Checking.\n\n" + block( + invoke( + "get_weather", + param("city", "Prague", is_str=True), + param("days", "3", is_str=False), + ), + invoke("search", param("q", "mlx", is_str=True)), + ) + + @staticmethod + def _stream(text, chunk): + parser = DeepSeekV4ToolParser() + parser.reset() + previous, content, tool_deltas = "", [], [] + for i in range(0, len(text), chunk): + delta = text[i : i + chunk] + current = previous + delta + result = parser.extract_tool_calls_streaming(previous, current, delta) + if result: + if "content" in result: + content.append(result["content"]) + if "tool_calls" in result: + tool_deltas.append(result["tool_calls"]) + previous = current + return "".join(content), tool_deltas + + @pytest.mark.parametrize("chunk", [1, 2, 3, 5, 7, 13, 31, 64, 128]) + def test_matches_non_streaming(self, chunk): + content, tool_deltas = self._stream(self.FULL, chunk) + expected = DeepSeekV4ToolParser().extract_tool_calls(self.FULL) + + assert len(tool_deltas) == 1, "tool calls must be emitted exactly once" + streamed = [ + (tc["function"]["name"], tc["function"]["arguments"]) + for tc in tool_deltas[0] + ] + assert streamed == [(tc["name"], tc["arguments"]) for tc in expected.tool_calls] + assert content.strip() == "Checking." + + @pytest.mark.parametrize("chunk", [1, 2, 3, 5, 7, 13, 31, 64, 128]) + def test_markup_never_leaks_into_content(self, chunk): + content, _ = self._stream(self.FULL, chunk) + assert D not in content + + def test_indices_are_sequential(self): + _, tool_deltas = self._stream(self.FULL, 5) + assert [tc["index"] for tc in tool_deltas[0]] == [0, 1] + + def test_partial_marker_is_not_swallowed(self): + """Text that merely looks like the marker must still be delivered.""" + parser = DeepSeekV4ToolParser() + parser.reset() + previous, out = "", [] + for delta in ["hello <", "｜not-dsml", " after"]: + current = previous + delta + result = parser.extract_tool_calls_streaming(previous, current, delta) + if result and "content" in result: + out.append(result["content"]) + previous = current + assert "".join(out) == "hello <｜not-dsml after" + + def test_plain_text_streams_through(self): + content, tool_deltas = self._stream("no tools here at all", 4) + assert content == "no tools here at all" + assert tool_deltas == [] + + def test_reset_clears_state_between_requests(self): + parser = DeepSeekV4ToolParser() + parser.reset() + parser.extract_tool_calls_streaming("", "hello <", "hello <") + parser.reset() + result = parser.extract_tool_calls_streaming("", "plain", "plain") + assert result == {"content": "plain"} + + +class TestAutoDetection: + def test_auto_parser_routes_dsml(self): + from vllm_mlx.tool_parsers.auto_tool_parser import AutoToolParser + + text = block(invoke("f", param("a", "1", is_str=True))) + result = AutoToolParser().extract_tool_calls(text) + assert result.tools_called + assert result.tool_calls[0]["name"] == "f" + + +class TestChainedWithReasoningParser: + """The server feeds the reasoning parser's content to the tool parser. + + Testing the two in isolation misses everything that can go wrong between + them, which is where the markup actually crosses over. + """ + + FULL = "deciding what to callChecking.\n\n" + block( + invoke( + "get_weather", + param("city", "Prague", is_str=True), + param("days", "3", is_str=False), + ), + invoke("search", param("q", "mlx", is_str=True)), + ) + + @staticmethod + def _pipeline(text, chunk): + from vllm_mlx.reasoning.deepseek_v4_parser import DeepSeekV4ReasoningParser + + reasoner = DeepSeekV4ReasoningParser() + reasoner.reset_state() + tools = DeepSeekV4ToolParser() + tools.reset() + + accumulated, tool_acc = "", "" + reasoning, content, calls = [], [], [] + for i in range(0, len(text), chunk): + delta = text[i : i + chunk] + previous, accumulated = accumulated, accumulated + delta + msg = reasoner.extract_reasoning_streaming(previous, accumulated, delta) + if msg is None: + continue + if msg.reasoning: + reasoning.append(msg.reasoning) + if not msg.content: + continue + prev_tool, tool_acc = tool_acc, tool_acc + msg.content + result = tools.extract_tool_calls_streaming( + prev_tool, tool_acc, msg.content + ) + if result is None: + continue + if "tool_calls" in result: + calls.append(result["tool_calls"]) + elif result.get("content"): + content.append(result["content"]) + return "".join(reasoning), "".join(content), calls + + @pytest.mark.parametrize("chunk", [1, 3, 8, 25, 64]) + def test_calls_survive_the_chain(self, chunk): + _, _, calls = self._pipeline(self.FULL, chunk) + expected = DeepSeekV4ToolParser().extract_tool_calls(self.FULL) + assert len(calls) == 1 + assert [ + (tc["function"]["name"], tc["function"]["arguments"]) for tc in calls[0] + ] == [(tc["name"], tc["arguments"]) for tc in expected.tool_calls] + + @pytest.mark.parametrize("chunk", [1, 3, 8, 25, 64]) + def test_no_markup_reaches_the_client(self, chunk): + reasoning, content, _ = self._pipeline(self.FULL, chunk) + for channel in (reasoning, content): + assert D not in channel + assert "" not in channel + assert "" not in channel + + @pytest.mark.parametrize("chunk", [1, 3, 8, 25, 64]) + def test_reasoning_and_content_land_in_the_right_channel(self, chunk): + reasoning, content, _ = self._pipeline(self.FULL, chunk) + assert reasoning.strip() == "deciding what to call" + assert content.strip() == "Checking." + + +class TestTextAndCallsInOneDelta: + """A block that opens and closes in one delta must not eat the text before it. + + Chunked arrival always worked, which is exactly why this hid: the streaming + tests drive the parser chunk by chunk, and only a whole response arriving as + a single delta reaches the branch that used to skip the flush. The result + was an assistant message whose text disappeared based on nothing but how the + model's output happened to be chunked. + """ + + SAMPLE = "Checking.\n\n" + block( + invoke("get_weather", param("city", "Prague", is_str=True)) + ) + + @staticmethod + def _drive(deltas): + parser = DeepSeekV4ToolParser() + previous, content, calls = "", [], 0 + for delta in deltas: + current = previous + delta + result = parser.extract_tool_calls_streaming(previous, current, delta) + if result: + if result.get("content"): + content.append(result["content"]) + if result.get("tool_calls"): + calls += len(result["tool_calls"]) + previous = current + return "".join(content), calls + + def test_single_delta_keeps_both(self): + content, calls = self._drive([self.SAMPLE]) + + assert calls == 1 + assert "Checking." in content, ( + "text preceding the tool call was dropped because the block opened " + "and closed in the same delta" + ) + + @pytest.mark.parametrize("chunk", [1, 3, 4, 17, 64, 4096]) + def test_result_does_not_depend_on_chunk_size(self, chunk): + whole = self._drive([self.SAMPLE]) + chunked = self._drive( + [self.SAMPLE[i : i + chunk] for i in range(0, len(self.SAMPLE), chunk)] + ) + + assert ( + whole == chunked + ), f"chunk={chunk} changed the response: {whole!r} vs {chunked!r}" + + def test_text_is_not_emitted_twice(self): + """The head is flushed once, whether it rides along or goes ahead.""" + content, _ = self._drive([self.SAMPLE]) + assert content.count("Checking.") == 1 diff --git a/vllm_mlx/cli.py b/vllm_mlx/cli.py index b01420a19..472e11c43 100644 --- a/vllm_mlx/cli.py +++ b/vllm_mlx/cli.py @@ -1354,6 +1354,8 @@ def create_parser() -> argparse.ArgumentParser: "harmony", "gpt-oss", "deepseek", + "deepseek_v4", + "dsml", "kimi", "granite", "nemotron", @@ -1366,7 +1368,8 @@ def create_parser() -> argparse.ArgumentParser: help=( "Select the tool call parser for the model. Options: " "auto (auto-detect), mistral, qwen, qwen3_coder, llama, hermes, " - "harmony, gpt-oss, deepseek, gemma4, kimi, granite, nemotron, " + "harmony, gpt-oss, deepseek, deepseek_v4, dsml, gemma4, kimi, " + "granite, nemotron, " "xlam, functionary, glm47, minimax. " "Required for --enable-auto-tool-choice." ), diff --git a/vllm_mlx/reasoning/__init__.py b/vllm_mlx/reasoning/__init__.py index 01b86696b..bf6efde78 100644 --- a/vllm_mlx/reasoning/__init__.py +++ b/vllm_mlx/reasoning/__init__.py @@ -76,6 +76,7 @@ def list_parsers() -> list[str]: def _register_builtin_parsers(): """Register built-in parsers.""" from .deepseek_r1_parser import DeepSeekR1ReasoningParser + from .deepseek_v4_parser import DeepSeekV4ReasoningParser from .gemma4_parser import Gemma4ReasoningParser from .glm4_parser import Glm4ReasoningParser from .gpt_oss_parser import GptOssReasoningParser @@ -86,6 +87,7 @@ def _register_builtin_parsers(): register_parser("qwen3", Qwen3ReasoningParser) register_parser("deepseek_r1", DeepSeekR1ReasoningParser) + register_parser("deepseek_v4", DeepSeekV4ReasoningParser) register_parser("gpt_oss", GptOssReasoningParser) register_parser("harmony", HarmonyReasoningParser) register_parser("gemma4", Gemma4ReasoningParser) diff --git a/vllm_mlx/reasoning/deepseek_v4_parser.py b/vllm_mlx/reasoning/deepseek_v4_parser.py new file mode 100644 index 000000000..0da7ccf0d --- /dev/null +++ b/vllm_mlx/reasoning/deepseek_v4_parser.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Reasoning parser for DeepSeek-V4 (Pro/Flash). + +V4 shares ````/```` with R1, and the prompt encoder closes the +generation prompt on ````, so the opening tag is usually absent from the +output — the lenient R1 behaviour this builds on already covers that. + +What V4 adds is the interaction with tool calls: the model must finish reasoning +before it may call a tool, so an opening ``<｜DSML｜tool_calls>`` marker +terminates the reasoning block even when ```` never arrives. Without +this, a tool-calling turn would have its entire DSML payload swallowed as +reasoning and the caller would see no tool call at all. +""" + +from ..utils.deepseek_v4_encoding import ( + DSML_TOKEN, + TOOL_CALLS_START, + partial_marker_len, +) +from .base import DeltaMessage +from .deepseek_r1_parser import DeepSeekR1ReasoningParser + +__all__ = ["DeepSeekV4ReasoningParser", "DSML_TOKEN", "TOOL_CALLS_START"] + + +class DeepSeekV4ReasoningParser(DeepSeekR1ReasoningParser): + """Reasoning parser for DeepSeek-V4. + + Example:: + + Input: "weighing options<｜DSML｜tool_calls>\\n<｜DSML｜invoke ..." + Output: reasoning="weighing options", + content="<｜DSML｜tool_calls>\\n<｜DSML｜invoke ..." + + The tool markup is deliberately left in ``content`` — extracting it is the + tool parser's job, and it needs the markup intact. + """ + + def __init__(self, tokenizer=None): + super().__init__(tokenizer) + # Characters of the accumulated text already emitted on either channel, + # and whether the stream has crossed into tool markup. + self._emitted_len: int = 0 + self._in_tool_markup: bool = False + + def reset_state(self): + """Reset state machine for a new streaming request.""" + super().reset_state() + self._emitted_len = 0 + self._in_tool_markup = False + + def extract_reasoning( + self, + model_output: str, + ) -> tuple[str | None, str | None]: + """Extract reasoning, honouring the implicit close before a tool call.""" + tool_idx = model_output.find(TOOL_CALLS_START) + if tool_idx != -1: + end_idx = model_output.find(self.end_token) + if end_idx == -1 or end_idx > tool_idx: + # A tool call opened while still reasoning: everything before it + # is reasoning, the markup and beyond is content. + reasoning = model_output[:tool_idx] + if reasoning.startswith(self.start_token): + reasoning = reasoning[len(self.start_token) :] + content = model_output[tool_idx:] + return reasoning.strip() or None, content or None + + return super().extract_reasoning(model_output) + + def extract_reasoning_streaming( + self, + previous_text: str, + current_text: str, + delta_text: str, + ) -> DeltaMessage | None: + """Stream reasoning, closing it when tool markup opens. + + Tracks how much of the accumulated text has been emitted rather than + trusting the delta, because the marker straddles delta boundaries: a + parser that streams each fragment as it arrives leaks markup into the + reasoning channel and then repeats the whole marker as content once it + recognises it. + """ + if self._in_tool_markup: + new = current_text[self._emitted_len :] + self._emitted_len = len(current_text) + return DeltaMessage(content=new) if new else None + + marker_idx = current_text.find(TOOL_CALLS_START) + end_idx = current_text.find(self.end_token) + reasoning_open = end_idx == -1 + + if marker_idx != -1 and (reasoning_open or end_idx > marker_idx): + # The marker completed while reasoning was still open. Withholding + # its prefix above guarantees nothing past marker_idx went out yet. + head = current_text[self._emitted_len : marker_idx] + tail = current_text[marker_idx:] + self._emitted_len = len(current_text) + self._in_tool_markup = True + self._phase = "content" + self._content_started = True + return DeltaMessage(reasoning=head or None, content=tail or None) + + # Hold back a tail that could still grow into one of the markers. + # <｜DSML｜tool_calls> is assembled from several tokens, so it always + # straddles deltas; the think tags have their own ids and normally + # arrive whole, but are covered too because a detokenizer that splits + # them would otherwise leak fragments into the reasoning channel and + # emit the remainder as content. + hold = ( + partial_marker_len( + current_text, TOOL_CALLS_START, self.start_token, self.end_token + ) + if reasoning_open + else 0 + ) + limit = len(current_text) - hold + if limit <= self._emitted_len: + return None + + effective_previous = current_text[: self._emitted_len] + effective_current = current_text[:limit] + effective_delta = current_text[self._emitted_len : limit] + self._emitted_len = limit + + return super().extract_reasoning_streaming( + effective_previous, effective_current, effective_delta + ) diff --git a/vllm_mlx/server.py b/vllm_mlx/server.py index f066593e9..e91b6121d 100644 --- a/vllm_mlx/server.py +++ b/vllm_mlx/server.py @@ -972,6 +972,9 @@ def _thinking_disabled(request, chat_kwargs: dict | None = None) -> bool: # streaming gates scan, so the harmony tokens remain present. "<|channel|>commentary", "<|call|>", + # DeepSeek-V4 DSML. Listed as the bare token rather than the full + # <｜DSML｜tool_calls> marker so a partial chunk still trips the gate. + "｜DSML｜", ) _STREAMING_BARE_BRACKET_MARKER = re.compile(r"\[\w+\(\{") _STREAMING_BARE_BRACKET_PARTIAL = re.compile(r"\[\w+\($") diff --git a/vllm_mlx/tool_parsers/__init__.py b/vllm_mlx/tool_parsers/__init__.py index 6b5d3462a..acc0d7371 100644 --- a/vllm_mlx/tool_parsers/__init__.py +++ b/vllm_mlx/tool_parsers/__init__.py @@ -13,6 +13,7 @@ - gemma4/gemma_4: Google Gemma 4 models (<|tool_call>call:name{} format) - hermes/nous: Hermes/NousResearch models - deepseek/deepseek_v3/deepseek_r1: DeepSeek models (unicode tokens) +- deepseek_v4/dsml: DeepSeek-V4 models (<｜DSML｜invoke name="..."> markup) - kimi/kimi_k2/moonshot: Kimi/Moonshot models - granite/granite3: IBM Granite models - nemotron/nemotron3: NVIDIA Nemotron models @@ -48,6 +49,7 @@ # Import parsers to register them from .auto_tool_parser import AutoToolParser from .deepseek_tool_parser import DeepSeekToolParser +from .deepseek_v4_tool_parser import DeepSeekV4ToolParser from .functionary_tool_parser import FunctionaryToolParser from .gemma4_tool_parser import Gemma4ToolParser from .granite_tool_parser import GraniteToolParser @@ -105,6 +107,7 @@ def get_parser_stop_tokens( "LlamaToolParser", "HermesToolParser", "DeepSeekToolParser", + "DeepSeekV4ToolParser", "KimiToolParser", "GraniteToolParser", "NemotronToolParser", diff --git a/vllm_mlx/tool_parsers/auto_tool_parser.py b/vllm_mlx/tool_parsers/auto_tool_parser.py index c759e33d7..f2fa4225a 100644 --- a/vllm_mlx/tool_parsers/auto_tool_parser.py +++ b/vllm_mlx/tool_parsers/auto_tool_parser.py @@ -16,6 +16,10 @@ ToolParser, ToolParserManager, ) +from .deepseek_v4_tool_parser import ( + TOOL_CALLS_START as DSML_TOOL_CALLS_START, +) +from .deepseek_v4_tool_parser import DeepSeekV4ToolParser from .gemma4_tool_parser import Gemma4ToolParser @@ -74,6 +78,14 @@ def extract_tool_calls( if result.tools_called: return result + # 1b. Try DeepSeek-V4 DSML. The marker carries the model's private + # ｜DSML｜ token, so a false positive is not realistic. + if DSML_TOOL_CALLS_START in model_output: + dsml_parser = DeepSeekV4ToolParser() + result = dsml_parser.extract_tool_calls(model_output, request) + if result.tools_called: + return result + # 2. Try Mistral format if self.MISTRAL_TOKEN in model_output: parts = model_output.split(self.MISTRAL_TOKEN) diff --git a/vllm_mlx/tool_parsers/deepseek_v4_tool_parser.py b/vllm_mlx/tool_parsers/deepseek_v4_tool_parser.py new file mode 100644 index 000000000..3b7fe7b88 --- /dev/null +++ b/vllm_mlx/tool_parsers/deepseek_v4_tool_parser.py @@ -0,0 +1,297 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Tool call parser for DeepSeek-V4 (Pro/Flash) DSML markup. + +DeepSeek-V4 does not emit JSON function calls like V3/R1 — it uses its own +markup language, DSML:: + + <｜DSML｜tool_calls> + <｜DSML｜invoke name="get_weather"> + <｜DSML｜parameter name="city" string="true">Prague + <｜DSML｜parameter name="days" string="false">3 + + + +The ``string`` attribute carries the type: ``"true"`` means the value is a raw +string, ``"false"`` means it is JSON (number, bool, array or object). That +distinction is why this cannot be a regex over ``name=value`` pairs — a string +parameter may legitimately contain ``"``, ``<`` or a JSON-looking payload. + +``DeepSeekToolParser`` in ``deepseek_tool_parser.py`` handles the V3/R1 format +(``<｜tool▁calls▁begin｜>`` plus fenced JSON) and shares nothing with this one. +""" + +import json +import re +import uuid +from collections.abc import Sequence +from typing import Any + +from ..utils.deepseek_v4_encoding import ( + DSML_TOKEN, + TOOL_CALLS_END, + TOOL_CALLS_START, + partial_marker_len, +) +from .abstract_tool_parser import ( + ExtractedToolCallInformation, + ToolParser, + ToolParserManager, +) + +INVOKE_START = f"<{DSML_TOKEN}invoke" +INVOKE_END = f"" +PARAM_START = f"<{DSML_TOKEN}parameter" +PARAM_END = f"" + +THINKING_END = "" + +# Attribute headers are bounded and well-formed; only the *values* are +# free-form, so a regex is safe here and a scanner is used for the rest. +_INVOKE_HEADER_RE = re.compile(r'\s*name="(?P[^"]*)"\s*>') +_PARAM_HEADER_RE = re.compile( + r'\s*name="(?P[^"]*)"\s+string="(?Ptrue|false)"\s*>' +) + + +def generate_tool_id() -> str: + """Generate a unique tool call ID.""" + return f"call_{uuid.uuid4().hex[:8]}" + + +@ToolParserManager.register_module(["deepseek_v4", "dsml"]) +class DeepSeekV4ToolParser(ToolParser): + """Parse DeepSeek-V4 DSML tool calls. + + Example:: + + <｜DSML｜tool_calls> + <｜DSML｜invoke name="search"> + <｜DSML｜parameter name="q" string="true">mlx + + + + Malformed markup is never fatal: whatever parses becomes a tool call and the + remainder is returned as content, because a half-generated call must not + take the server down. + """ + + # The encoder consumes role="tool" messages and assistant tool_calls + # directly — folding results into blocks and rendering calls + # back as DSML. Declaring this False would make the server flatten them to + # "[Tool Result (id)]: ..." and "[Calling tool: name(...)]" first + # (api/utils.py), so the model would see a shape it was never trained on + # and the encoder's own handling would never run. + SUPPORTS_NATIVE_TOOL_FORMAT = True + + def __init__(self, tokenizer=None): + super().__init__(tokenizer) + # Trailing text withheld because it might be the start of the opening + # marker; reset per request via reset(). + self._pending: str = "" + # How much of the accumulated text has already been streamed as + # content, so the text preceding a tool block is emitted exactly once. + self._emitted_len: int = 0 + + def extract_tool_calls( + self, model_output: str, request: dict[str, Any] | None = None + ) -> ExtractedToolCallInformation: + """Extract DSML tool calls from a complete response.""" + start = model_output.find(TOOL_CALLS_START) + if start == -1: + return ExtractedToolCallInformation( + tools_called=False, + tool_calls=[], + content=self.strip_think_tags(model_output) or None, + ) + + content = self._content_before(model_output, start) + block_start = start + len(TOOL_CALLS_START) + end = model_output.find(TOOL_CALLS_END, block_start) + block = model_output[block_start : end if end != -1 else len(model_output)] + + tool_calls = self._parse_invokes(block) + if not tool_calls: + # Marker present but nothing parsed — surface the raw text rather + # than silently dropping the model's output. + return ExtractedToolCallInformation( + tools_called=False, + tool_calls=[], + content=self.strip_think_tags(model_output) or None, + ) + + return ExtractedToolCallInformation( + tools_called=True, tool_calls=tool_calls, content=content + ) + + def _content_before(self, text: str, tool_start: int) -> str | None: + """Return the assistant content that precedes the tool call block. + + A tool call inside a reasoning block implicitly ends it, so anything up + to and including ```` belongs to the reasoning parser, not here. + """ + head = text[:tool_start] + think_end = head.rfind(THINKING_END) + if think_end != -1: + head = head[think_end + len(THINKING_END) :] + return head.strip() or None + + def _parse_invokes(self, block: str) -> list[dict[str, Any]]: + """Parse every ``invoke`` element inside a tool_calls block.""" + tool_calls: list[dict[str, Any]] = [] + pos = 0 + + while True: + inv = block.find(INVOKE_START, pos) + if inv == -1: + break + + header = _INVOKE_HEADER_RE.match(block, inv + len(INVOKE_START)) + if header is None: + # Unparseable header: skip this marker and keep looking. + pos = inv + len(INVOKE_START) + continue + + name = header.group("name") + body_start = header.end() + body_end = block.find(INVOKE_END, body_start) + body = block[body_start : body_end if body_end != -1 else len(block)] + + arguments = self._parse_parameters(body) + tool_calls.append( + { + "id": generate_tool_id(), + "type": "function", + "name": name, + "arguments": json.dumps(arguments, ensure_ascii=False), + } + ) + + if body_end == -1: + break + pos = body_end + len(INVOKE_END) + + return tool_calls + + def _parse_parameters(self, body: str) -> dict[str, Any]: + """Parse ``parameter`` elements into a plain argument dict. + + ``string="true"`` keeps the value verbatim; ``string="false"`` is JSON + and gets decoded, falling back to the raw string if the model emitted + something that does not parse. + """ + arguments: dict[str, Any] = {} + pos = 0 + + while True: + par = body.find(PARAM_START, pos) + if par == -1: + break + + header = _PARAM_HEADER_RE.match(body, par + len(PARAM_START)) + if header is None: + pos = par + len(PARAM_START) + continue + + value_start = header.end() + value_end = body.find(PARAM_END, value_start) + if value_end == -1: + # Truncated parameter — drop it rather than guess its value. + break + + raw = body[value_start:value_end] + if header.group("is_str") == "true": + value: Any = raw + else: + try: + value = json.loads(raw) + except (ValueError, json.JSONDecodeError): + value = raw + + arguments[header.group("name")] = value + pos = value_end + len(PARAM_END) + + return arguments + + def _format_streaming(self, result: ExtractedToolCallInformation) -> dict[str, Any]: + """Render extracted tool calls into the streaming delta shape.""" + return { + "tool_calls": [ + { + "index": i, + "id": tc["id"], + "type": "function", + "function": { + "name": tc["name"], + "arguments": tc["arguments"], + }, + } + for i, tc in enumerate(result.tool_calls) + ] + } + + def reset(self) -> None: + """Reset parser state for a new request.""" + super().reset() + self._pending = "" + self._emitted_len = 0 + + def extract_tool_calls_streaming( + self, + previous_text: str, + current_text: str, + delta_text: str, + previous_token_ids: Sequence[int] | None = None, + current_token_ids: Sequence[int] | None = None, + delta_token_ids: Sequence[int] | None = None, + request: dict[str, Any] | None = None, + ) -> dict[str, Any] | None: + """Stream DSML output. + + DSML exposes no usable partial state — a parameter's type is only known + once its closing tag arrives — so once the block opens everything is + buffered and the calls are emitted in one delta when the block closes. + + Before that, content passes through, except for a tail that could still + grow into the opening marker. That tail is held in ``_pending`` and + released as soon as the next delta proves it was ordinary text. + """ + block_start = current_text.find(TOOL_CALLS_START) + if block_start != -1: + closed = ( + TOOL_CALLS_END in current_text and TOOL_CALLS_END not in previous_text + ) + # Flush any text preceding the block that has not been streamed + # yet. When the block spans several deltas this goes out on its + # own, ahead of the calls. When the whole response arrives as one + # delta there is no later delta to flush into, so the text rides + # along with the calls rather than being dropped — losing it made + # the response depend on how the output happened to be chunked. + head = "" + if self._emitted_len < block_start: + head = current_text[self._emitted_len : block_start] + self._emitted_len = block_start + self._pending = "" + if head and not closed: + return {"content": head} + + # The closing marker frequently straddles delta boundaries, so + # completion is detected against the accumulated text, not the + # delta. Comparing with previous_text makes it fire exactly once. + if closed: + result = self.extract_tool_calls(current_text, request) + if result.tools_called: + self._pending = "" + self._emitted_len = len(current_text) + formatted = self._format_streaming(result) + if head: + formatted = {**formatted, "content": head} + return formatted + return None + + text = self._pending + delta_text + hold = partial_marker_len(text) + self._pending = text[len(text) - hold :] if hold else "" + emit = text[: len(text) - hold] if hold else text + self._emitted_len += len(emit) + return {"content": emit} if emit else None diff --git a/vllm_mlx/utils/deepseek_v4_encoding.py b/vllm_mlx/utils/deepseek_v4_encoding.py new file mode 100644 index 000000000..d5d2e8d29 --- /dev/null +++ b/vllm_mlx/utils/deepseek_v4_encoding.py @@ -0,0 +1,726 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Prompt encoding for DeepSeek-V4 (Pro/Flash). + +DeepSeek-V4 ships no Jinja ``chat_template`` — its ``tokenizer_config.json`` +carries only BOS/EOS/pad. The prompt is built programmatically instead, which is +what this module does. It is a port of the reference ``encoding_dsv4.py`` +published alongside the model weights; upstream vLLM solves the same problem the +same way in ``vllm/tokenizers/deepseek_v4_encoding.py``. + +Format in brief:: + + <｜begin▁of▁sentence｜>{system}<｜User｜>{question}<｜Assistant｜> + +The system message is bare text with no wrapper — roles are delimited solely by +``<｜User｜>`` and ``<｜Assistant｜>``. A turn ends with ```` in thinking +mode (the model then reasons) or ```` in chat mode (reasoning +suppressed). There is no ``tool`` role: tool results are merged into the +preceding user turn as ```` blocks. +""" + +import copy +import json +import logging +from typing import Any + +logger = logging.getLogger(__name__) + +BOS_TOKEN = "<｜begin▁of▁sentence｜>" +EOS_TOKEN = "<｜end▁of▁sentence｜>" +THINKING_START_TOKEN = "" +THINKING_END_TOKEN = "" +DSML_TOKEN = "｜DSML｜" + +USER_SP_TOKEN = "<｜User｜>" +ASSISTANT_SP_TOKEN = "<｜Assistant｜>" +LATEST_REMINDER_SP_TOKEN = "<｜latest_reminder｜>" + +# Special tokens for DeepSeek-internal classification tasks. Not reachable +# through the OpenAI API surface, but render_message() honours them so the +# encoder stays a faithful port. +DS_TASK_SP_TOKENS = { + "action": "<｜action｜>", + "query": "<｜query｜>", + "authority": "<｜authority｜>", + "domain": "<｜domain｜>", + "title": "<｜title｜>", + "read_url": "<｜read_url｜>", +} +VALID_TASKS = set(DS_TASK_SP_TOKENS) + +TOOL_CALLS_BLOCK_NAME = "tool_calls" + +# Markers delimiting a tool call block. Defined here, with the rest of the wire +# format, so the encoder and both parsers cannot drift apart. +TOOL_CALLS_START = f"<{DSML_TOKEN}{TOOL_CALLS_BLOCK_NAME}>" +TOOL_CALLS_END = f"" + + +def partial_marker_len(text: str, *markers: str) -> int: + """Length of the trailing run of ``text`` that is still a prefix of one of + ``markers`` (the tool-call start marker by default). + + Streaming deltas are token-sized. ``<｜DSML｜tool_calls>`` is not a single + token — only the bare ``｜DSML｜`` is — so the marker always arrives in + pieces. A parser that emits those pieces as they come leaks markup into the + user-visible stream and then repeats the whole marker once it recognises + it. Withholding this many trailing characters avoids both. + + Returns 0 when the tail cannot grow into any of the markers. + """ + markers = markers or (TOOL_CALLS_START,) + longest = max(len(m) for m in markers) + for size in range(min(len(text), longest - 1), 0, -1): + tail = text[-size:] + if any(m.startswith(tail) for m in markers): + return size + return 0 + + +ASSISTANT_MSG_TEMPLATE = "{reasoning}{content}{tool_calls}" + EOS_TOKEN +ASSISTANT_MSG_WO_EOS_TEMPLATE = "{reasoning}{content}{tool_calls}" +TOOL_CALL_TEMPLATE = ( + '<{dsml_token}invoke name="{name}">\n{arguments}\n' +) +TOOL_CALLS_TEMPLATE = ( + "<{dsml_token}{tc_block_name}>\n{tool_calls}\n" +) +TOOL_OUTPUT_TEMPLATE = "{content}" +RESPONSE_FORMAT_TEMPLATE = ( + "## Response Format:\n\nYou MUST strictly adhere to the following " + "schema to reply:\n{schema}" +) + +# Reasoning effort is a plain text prefix prepended to the whole conversation, +# not a token and not a sampling parameter. "low" is the default and adds +# nothing; the prefix only applies in thinking mode. +REASONING_EFFORT_PROMPTS: dict[str, str] = { + "low": "", + "high": ( + "Reasoning Effort: Absolute maximum with no shortcuts permitted.\n" + "You MUST be very thorough in your thinking and comprehensively " + "decompose the problem to resolve the root cause, rigorously " + "stress-testing your logic against all potential paths, edge cases, " + "and adversarial scenarios.\n" + "Explicitly write out your entire deliberation process, documenting " + "every intermediate step, considered alternative, and rejected " + "hypothesis to ensure absolutely no assumption is left unchecked.\n\n" + ), + "max": ( + "Reasoning Effort: Beyond maximum — exhaustive, relentless, and " + "uncompromising.\n" + "You MUST reason with the utmost depth and rigor, leaving absolutely " + "nothing to chance: exhaustively decompose the problem into its most " + "fundamental components, trace every causal chain to its root, and " + "resolve the underlying cause rather than any surface symptom.\n" + "Do not stop reasoning until you have independently verified the " + "solution from multiple angles and are certain that no assumption " + "remains unchecked and no error remains undiscovered.\n\n" + ), +} +DEFAULT_REASONING_EFFORT = "low" + +TOOLS_TEMPLATE = """## Tools + +You have access to a set of tools to help answer the user's question. You can invoke tools by writing a "<{dsml_token}tool_calls>" block like the following: + +<{dsml_token}tool_calls> +<{dsml_token}invoke name="$TOOL_NAME"> +<{dsml_token}parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE +... + +<{dsml_token}invoke name="$TOOL_NAME2"> +... + + + +String parameters should be specified as is and set `string="true"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string="false"`. + +If thinking_mode is enabled (triggered by {thinking_start_token}), you MUST output your complete reasoning inside {thinking_start_token}...{thinking_end_token} BEFORE any tool calls or final response. + +Otherwise, output directly after {thinking_end_token} with tool calls or final response. + +### Available Tool Schemas + +{tool_schemas} + +You MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls. +""" + + +def to_json(value: Any) -> str: + """Serialize to JSON, falling back to ASCII escaping if needed.""" + try: + return json.dumps(value, ensure_ascii=False) + except (TypeError, ValueError): + return json.dumps(value, ensure_ascii=True) + + +def tools_from_openai_format(tools: list[dict]) -> list[dict]: + """Strip the OpenAI ``{"type": "function", "function": {...}}`` wrapper.""" + return [tool["function"] if "function" in tool else tool for tool in tools] + + +def tool_calls_from_openai_format(tool_calls: list[dict]) -> list[dict]: + return [ + { + "name": tc["function"]["name"], + "arguments": tc["function"]["arguments"], + } + for tc in tool_calls + ] + + +def encode_arguments_to_dsml(tool_call: dict[str, str]) -> str: + """Render one tool call's arguments as DSML ``parameter`` elements. + + ``string="true"`` marks a raw string value; anything else is JSON-encoded + and marked ``string="false"``. Arguments that do not parse as JSON are + wrapped under an ``arguments`` key rather than dropped. + + Accepts ``arguments`` either as the JSON string the OpenAI wire format uses + or as an already-decoded mapping: ``api/utils.py`` decodes it in place when + native tool format is preserved, and json-loading that again would collapse + every parameter into one bogus ``arguments`` entry. + """ + template = ( + '<{dsml_token}parameter name="{key}" string="{is_str}">' + "{value}" + ) + + raw = tool_call["arguments"] + if isinstance(raw, (dict, list)): + arguments = raw + else: + try: + arguments = json.loads(raw) + except (TypeError, ValueError, json.JSONDecodeError): + arguments = {"arguments": raw} + + if not isinstance(arguments, dict): + arguments = {"arguments": arguments} + + parts = [] + for key, value in arguments.items(): + is_str = isinstance(value, str) + parts.append( + template.format( + dsml_token=DSML_TOKEN, + key=key, + is_str="true" if is_str else "false", + value=value if is_str else to_json(value), + ) + ) + return "\n".join(parts) + + +def render_tools(tools: list[dict[str, Any]]) -> str: + """Render tool schemas into the block that goes into the system message.""" + return TOOLS_TEMPLATE.format( + tool_schemas="\n".join(to_json(t) for t in tools), + dsml_token=DSML_TOKEN, + thinking_start_token=THINKING_START_TOKEN, + thinking_end_token=THINKING_END_TOKEN, + ) + + +def find_last_user_index(messages: list[dict[str, Any]]) -> int: + """Index of the last user/developer message, or -1 if there is none.""" + for idx in range(len(messages) - 1, -1, -1): + if messages[idx].get("role") in ("user", "developer"): + return idx + return -1 + + +def render_message( + index: int, + messages: list[dict[str, Any]], + thinking_mode: str, + drop_thinking: bool = True, + reasoning_effort: str | None = None, +) -> str: + """Render a single message into its encoded form. + + Args: + index: Position of the message to render. + messages: The full conversation (needed for look-ahead and for + locating the last user turn). + thinking_mode: ``"thinking"`` or ``"chat"``. + drop_thinking: Drop reasoning content from turns before the last user + message. + reasoning_effort: ``"low"`` (default), ``"high"`` or ``"max"``. Only + applied at index 0 and only in thinking mode. + """ + if not 0 <= index < len(messages): + raise IndexError(f"index {index} out of range for {len(messages)} messages") + if thinking_mode not in ("chat", "thinking"): + raise ValueError(f"Invalid thinking_mode `{thinking_mode}`") + + prompt = "" + msg = messages[index] + last_user_idx = find_last_user_index(messages) + + role = msg.get("role") + content = msg.get("content") + tools = msg.get("tools") + response_format = msg.get("response_format") + tool_calls = msg.get("tool_calls") + reasoning_content = msg.get("reasoning_content") + wo_eos = msg.get("wo_eos", False) + + if tools: + tools = tools_from_openai_format(tools) + if tool_calls: + tool_calls = tool_calls_from_openai_format(tool_calls) + + reasoning_effort = reasoning_effort or DEFAULT_REASONING_EFFORT + if reasoning_effort not in REASONING_EFFORT_PROMPTS: + raise ValueError( + f"Invalid reasoning effort: {reasoning_effort}, expected one of " + f"{list(REASONING_EFFORT_PROMPTS)}" + ) + if index == 0 and thinking_mode == "thinking": + prompt += REASONING_EFFORT_PROMPTS[reasoning_effort] + + if role == "system": + prompt += content or "" + if tools: + prompt += "\n\n" + render_tools(tools) + if response_format: + prompt += "\n\n" + RESPONSE_FORMAT_TEMPLATE.format( + schema=to_json(response_format) + ) + + elif role == "developer": + if not content: + raise ValueError(f"Invalid message for role `{role}`: {msg}") + prompt += USER_SP_TOKEN + content + if tools: + prompt += "\n\n" + render_tools(tools) + if response_format: + prompt += "\n\n" + RESPONSE_FORMAT_TEMPLATE.format( + schema=to_json(response_format) + ) + + elif role == "user": + prompt += USER_SP_TOKEN + content_blocks = msg.get("content_blocks") + if content_blocks: + prompt += "\n\n".join(_render_content_blocks(content_blocks)) + else: + prompt += content or "" + + elif role == "latest_reminder": + prompt += LATEST_REMINDER_SP_TOKEN + (content or "") + + elif role == "tool": + raise NotImplementedError( + "deepseek_v4 has no tool role; preprocess with merge_tool_messages()" + ) + + elif role == "assistant": + thinking_part = "" + tc_content = "" + + if tool_calls: + rendered = [ + TOOL_CALL_TEMPLATE.format( + dsml_token=DSML_TOKEN, + name=tc.get("name"), + arguments=encode_arguments_to_dsml(tc), + ) + for tc in tool_calls + ] + tc_content = "\n\n" + TOOL_CALLS_TEMPLATE.format( + dsml_token=DSML_TOKEN, + tool_calls="\n".join(rendered), + tc_block_name=TOOL_CALLS_BLOCK_NAME, + ) + + # A message following a task carries the task's output, which has no + # thinking block of its own. + prev_has_task = index - 1 >= 0 and messages[index - 1].get("task") is not None + + if thinking_mode == "thinking" and not prev_has_task: + if not drop_thinking or index > last_user_idx: + thinking_part = (reasoning_content or "") + THINKING_END_TOKEN + + template = ASSISTANT_MSG_WO_EOS_TEMPLATE if wo_eos else ASSISTANT_MSG_TEMPLATE + prompt += template.format( + reasoning=thinking_part, + content=content or "", + tool_calls=tc_content, + ) + else: + raise NotImplementedError(f"Unknown role: {role}") + + # Transition tokens are only appended when this message is the last one, or + # when an assistant turn follows it. + next_role = messages[index + 1].get("role") if index + 1 < len(messages) else None + if next_role is not None and next_role not in ("assistant", "latest_reminder"): + return prompt + + task = msg.get("task") + if task is not None: + if task not in VALID_TASKS: + raise ValueError( + f"Invalid task: '{task}'. Valid tasks are: {sorted(VALID_TASKS)}" + ) + if task != "action": + prompt += DS_TASK_SP_TOKENS[task] + else: + prompt += ASSISTANT_SP_TOKEN + prompt += ( + THINKING_START_TOKEN + if thinking_mode == "thinking" + else THINKING_END_TOKEN + ) + prompt += DS_TASK_SP_TOKENS[task] + + elif role in ("user", "developer"): + prompt += ASSISTANT_SP_TOKEN + if thinking_mode == "thinking" and ( + not drop_thinking or index >= last_user_idx + ): + prompt += THINKING_START_TOKEN + else: + prompt += THINKING_END_TOKEN + + return prompt + + +def _render_content_blocks(content_blocks: list[dict[str, Any]]) -> list[str]: + """Render user content blocks (interleaved text and tool results).""" + parts = [] + for block in content_blocks: + block_type = block.get("type") + if block_type == "text": + parts.append(block.get("text", "")) + elif block_type == "tool_result": + tool_content = block.get("content", "") + if isinstance(tool_content, list): + text_parts = [ + ( + b.get("text", "") + if b.get("type") == "text" + else f"[Unsupported {b.get('type')}]" + ) + for b in tool_content + ] + tool_content = "\n\n".join(text_parts) + parts.append(TOOL_OUTPUT_TEMPLATE.format(content=tool_content)) + else: + parts.append(f"[Unsupported {block_type}]") + return parts + + +def merge_tool_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: + """Fold ``role: tool`` messages into the preceding user turn. + + DeepSeek-V4 has no standalone tool role — results are carried as + ```` blocks inside a user message. Consecutive tool results + accumulate into one user turn. + """ + merged: list[dict[str, Any]] = [] + + for msg in messages: + msg = copy.deepcopy(msg) + role = msg.get("role") + + if role == "tool": + block = { + "type": "tool_result", + "tool_use_id": msg.get("tool_call_id", ""), + "content": msg.get("content", ""), + } + if ( + merged + and merged[-1].get("role") == "user" + and "content_blocks" in merged[-1] + ): + merged[-1]["content_blocks"].append(block) + else: + merged.append({"role": "user", "content_blocks": [block]}) + elif role == "user": + block = {"type": "text", "text": msg.get("content", "")} + if ( + merged + and merged[-1].get("role") == "user" + and "content_blocks" in merged[-1] + and merged[-1].get("task") is None + ): + merged[-1]["content_blocks"].append(block) + else: + new_msg: dict[str, Any] = { + "role": "user", + "content": msg.get("content", ""), + "content_blocks": [block], + } + for key in ("task", "wo_eos", "mask"): + if key in msg: + new_msg[key] = msg[key] + merged.append(new_msg) + else: + merged.append(msg) + + return merged + + +def sort_tool_results_by_call_order( + messages: list[dict[str, Any]], +) -> list[dict[str, Any]]: + """Reorder tool results to match the order of the calls that produced them. + + Clients may return results out of order; the model was trained on results + that follow call order. + """ + call_order: dict[str, int] = {} + + for msg in messages: + role = msg.get("role") + if role == "assistant" and msg.get("tool_calls"): + call_order = {} + for idx, tc in enumerate(msg["tool_calls"]): + tc_id = tc.get("id") or tc.get("function", {}).get("id", "") + if tc_id: + call_order[tc_id] = idx + + elif role == "user" and msg.get("content_blocks"): + tool_blocks = [ + b for b in msg["content_blocks"] if b.get("type") == "tool_result" + ] + if len(tool_blocks) > 1 and call_order: + ordered = sorted( + tool_blocks, + key=lambda b: call_order.get(b.get("tool_use_id", ""), 0), + ) + it = iter(ordered) + msg["content_blocks"] = [ + next(it) if b.get("type") == "tool_result" else b + for b in msg["content_blocks"] + ] + + return messages + + +def _drop_thinking_messages( + messages: list[dict[str, Any]], +) -> list[dict[str, Any]]: + """Strip stale reasoning from turns before the last user message. + + User/system/tool/reminder turns survive untouched, as does everything from + the last user message onward. Earlier assistant turns keep their content but + lose ``reasoning_content``; earlier developer turns are dropped entirely. + """ + last_user_idx = find_last_user_index(messages) + keep_roles = { + "user", + "system", + "tool", + "latest_reminder", + "direct_search_results", + } + result = [] + + for idx, msg in enumerate(messages): + role = msg.get("role") + if role in keep_roles or idx >= last_user_idx: + result.append(msg) + elif role == "assistant": + msg = copy.copy(msg) + msg.pop("reasoning_content", None) + result.append(msg) + + return result + + +def encode_messages( + messages: list[dict[str, Any]], + thinking_mode: str, + context: list[dict[str, Any]] | None = None, + drop_thinking: bool = True, + add_default_bos_token: bool = True, + reasoning_effort: str | None = None, +) -> str: + """Encode a conversation into a DeepSeek-V4 prompt string. + + Args: + messages: Conversation in OpenAI format. A ``tool`` role is accepted and + folded into the preceding user turn. + thinking_mode: ``"thinking"`` (prompt ends with ````) or + ``"chat"`` (ends with ````, suppressing reasoning). + context: Already-established prefix turns. When given, BOS is not + emitted again. + drop_thinking: Drop reasoning from turns before the last user message. + Forced off when any message defines tools, because tool-calling + depends on the reasoning that produced earlier calls. + add_default_bos_token: Emit BOS at the start of the conversation. + reasoning_effort: ``"low"``/``None``, ``"high"`` or ``"max"``. Thinking + mode only. + + Returns: + The prompt string, ready to tokenize. + """ + context = context or [] + + messages = merge_tool_messages(messages) + messages = sort_tool_results_by_call_order(context + messages)[len(context) :] + if context: + context = merge_tool_messages(context) + context = sort_tool_results_by_call_order(context) + + full_messages = context + messages + + prompt = BOS_TOKEN if add_default_bos_token and not context else "" + + # Tool-calling conversations need their full reasoning history: the model + # has to see why it made the earlier calls. + effective_drop_thinking = drop_thinking + if any(m.get("tools") for m in full_messages): + effective_drop_thinking = False + + if thinking_mode == "thinking" and effective_drop_thinking: + full_messages = _drop_thinking_messages(full_messages) + num_to_render = len(full_messages) - len(_drop_thinking_messages(context)) + context_len = len(full_messages) - num_to_render + else: + num_to_render = len(messages) + context_len = len(context) + + for idx in range(num_to_render): + prompt += render_message( + idx + context_len, + full_messages, + thinking_mode=thinking_mode, + drop_thinking=effective_drop_thinking, + reasoning_effort=reasoning_effort, + ) + + return prompt + + +# --------------------------------------------------------------------------- +# OpenAI API adaptation +# --------------------------------------------------------------------------- + +# OpenAI exposes reasoning_effort as low/medium/high; DeepSeek-V4 defines +# low/high/max prompts plus "no thinking at all". "medium" has no distinct +# prompt of its own, so it maps onto "high" — as does any unrecognised value, +# which keeps a typo from silently disabling reasoning. +_EFFORT_ALIASES = { + "low": "low", + "minimal": "low", + "medium": "high", + "high": "high", + "max": "max", + "xhigh": "max", +} + + +def resolve_thinking( + enable_thinking: bool | None = None, + reasoning_effort: str | None = None, + thinking_mode: str | None = None, +) -> tuple[str, str | None]: + """Map OpenAI-style knobs onto ``(thinking_mode, reasoning_effort)``. + + ``reasoning_effort="none"`` and ``enable_thinking=False`` both select chat + mode, in which the prompt is closed with ```` and the model skips + reasoning entirely. In that mode the effort prefix is meaningless and is + dropped. + """ + if thinking_mode is not None: + if thinking_mode not in ("chat", "thinking"): + raise ValueError(f"Invalid thinking_mode `{thinking_mode}`") + mode = thinking_mode + elif reasoning_effort == "none" or enable_thinking is False: + mode = "chat" + else: + mode = "thinking" + + if mode == "chat" or reasoning_effort in (None, "none"): + return mode, None + + effort = _EFFORT_ALIASES.get(str(reasoning_effort).lower()) + if effort is None: + logger.warning( + "Unknown reasoning_effort %r for deepseek_v4, treating as 'high'", + reasoning_effort, + ) + effort = "high" + return mode, effort + + +def _attach_tools( + conversation: list[dict[str, Any]], tools: list[dict] | None +) -> list[dict[str, Any]]: + """Put the tool schemas where the encoder expects them: on a system turn. + + OpenAI passes tools as a top-level parameter, but DeepSeek-V4 renders them + from a message field. A conversation that already declares tools on one of + its messages is left alone. + """ + if not tools or any(m.get("tools") for m in conversation): + return conversation + + conversation = [dict(m) for m in conversation] + for msg in conversation: + if msg.get("role") == "system": + msg["tools"] = tools + return conversation + + return [{"role": "system", "content": "", "tools": tools}, *conversation] + + +def apply_chat_template( + conversation: list[dict[str, Any]], + tools: list[dict] | None = None, + enable_thinking: bool | None = None, + reasoning_effort: str | None = None, + thinking_mode: str | None = None, + drop_thinking: bool = True, + add_default_bos_token: bool = True, + **_ignored: Any, +) -> str: + """Build a DeepSeek-V4 prompt from an OpenAI-format conversation. + + Signature-compatible with ``tokenizer.apply_chat_template`` for the kwargs + vllm-mlx actually passes. ``add_generation_prompt`` is accepted and ignored: + the encoder always closes on the assistant prefix, which is the only mode + the model was trained for. + """ + mode, effort = resolve_thinking(enable_thinking, reasoning_effort, thinking_mode) + conversation = _attach_tools(conversation, tools) + return encode_messages( + conversation, + thinking_mode=mode, + drop_thinking=drop_thinking, + add_default_bos_token=add_default_bos_token, + reasoning_effort=effort, + ) + + +def install(tokenizer: Any) -> Any: + """Route ``tokenizer.apply_chat_template`` through the V4 encoder. + + DeepSeek-V4 carries no Jinja template, so the stock path either raises or + falls back to naive ``"role: content"`` concatenation. Overriding the method + on the tokenizer fixes every caller at once — the two engines and + ``models/llm.py`` all reach the template through this one method. + + Idempotent; returns the same tokenizer for convenience. + """ + if getattr(tokenizer, "_deepseek_v4_encoding_installed", False): + return tokenizer + + def _apply(conversation, tools=None, tokenize=False, **kwargs): + prompt = apply_chat_template(conversation, tools=tools, **kwargs) + if tokenize: + return tokenizer.encode(prompt) + return prompt + + tokenizer.apply_chat_template = _apply + tokenizer._deepseek_v4_encoding_installed = True + logger.info("[deepseek_v4] installed programmatic chat template encoder") + return tokenizer diff --git a/vllm_mlx/utils/tokenizer.py b/vllm_mlx/utils/tokenizer.py index 9d200ab9f..498c3f24d 100644 --- a/vllm_mlx/utils/tokenizer.py +++ b/vllm_mlx/utils/tokenizer.py @@ -49,6 +49,32 @@ def _needs_strict_false(model_name: str) -> bool: return False +def _model_type(model_name: str) -> str: + """Read model_type from config.json, empty string if unavailable.""" + from mlx_lm.utils import _download, load_config + + try: + config = load_config(_download(model_name)) + except Exception: + return "" + return config.get("text_config", config).get("model_type", "") or "" + + +def _install_custom_chat_template(model_name: str, tokenizer): + """Give models without a Jinja chat template a programmatic encoder. + + DeepSeek-V4 ships no ``chat_template``, so the stock path would either raise + or fall back to naive ``"role: content"`` concatenation. Patching the + tokenizer covers every caller at once. + """ + if tokenizer is None or _model_type(model_name) != "deepseek_v4": + return tokenizer + + from .deepseek_v4_encoding import install as install_deepseek_v4 + + return install_deepseek_v4(tokenizer) + + def load_model_with_fallback(model_name: str, tokenizer_config: dict = None): """ Load model and tokenizer with fallback for non-standard tokenizers. @@ -60,6 +86,11 @@ def load_model_with_fallback(model_name: str, tokenizer_config: dict = None): Returns: Tuple of (model, tokenizer) """ + model, tokenizer = _load_model_with_fallback(model_name, tokenizer_config) + return model, _install_custom_chat_template(model_name, tokenizer) + + +def _load_model_with_fallback(model_name: str, tokenizer_config: dict = None): from mlx_lm import load tokenizer_config = tokenizer_config or {} From f19475dcd3e0ac3c699aff1aa3dac4f40faa73d0 Mon Sep 17 00:00:00 2001 From: Wayner Barrios Date: Wed, 26 Aug 2026 00:32:42 -0400 Subject: [PATCH 2/3] Fix DeepSeek V4 streaming lifecycle and reasoning profiles (#676) --- .github/workflows/ci.yml | 3 + README.es.md | 4 +- README.fr.md | 4 +- README.md | 4 +- README.zh.md | 4 +- docs/development/architecture.md | 2 +- docs/es/guides/mcp-tools.md | 2 +- docs/fr/guides/mcp-tools.md | 2 +- docs/guides/mcp-tools.md | 2 +- docs/guides/reasoning.md | 18 ++ docs/guides/tool-calling.md | 63 ++++- docs/zh/guides/mcp-tools.md | 2 +- tests/test_deepseek_v4_encoding.py | 59 ++++- tests/test_deepseek_v4_reasoning.py | 10 + tests/test_deepseek_v4_tool_parser.py | 40 ++++ tests/test_responses_api.py | 69 ++++++ tests/test_server.py | 217 ++++++++++++++++++ vllm_mlx/reasoning/deepseek_v4_parser.py | 14 ++ vllm_mlx/server.py | 123 +++++++--- vllm_mlx/tool_parsers/auto_tool_parser.py | 65 +++++- .../tool_parsers/deepseek_v4_tool_parser.py | 25 ++ vllm_mlx/utils/deepseek_v4_encoding.py | 118 ++++++++-- vllm_mlx/utils/tokenizer.py | 2 +- 23 files changed, 787 insertions(+), 65 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 17af2bf6c..05f1261d4 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -94,6 +94,9 @@ jobs: tests/test_gemma4_openai_format.py \ tests/test_gemma4_streaming_edge.py \ tests/test_gemma4_tool_parser.py \ + tests/test_deepseek_v4_encoding.py \ + tests/test_deepseek_v4_reasoning.py \ + tests/test_deepseek_v4_tool_parser.py \ tests/test_minimax_tool_calling.py \ tests/test_qwen3_xml_parser.py \ tests/test_qwen3_xml_registration.py \ diff --git a/README.es.md b/README.es.md index c78e921d9..7305aa2c5 100644 --- a/README.es.md +++ b/README.es.md @@ -46,7 +46,7 @@ claude ### APIs - **Compatible con OpenAI**: `/v1/chat/completions`, `/v1/completions`, `/v1/embeddings`, `/v1/rerank`, `/v1/responses` - **Compatible con Anthropic**: `/v1/messages` (streaming, tool use, system prompts) -- **MCP Tool Calling**: 12 parsers (OpenAI, Anthropic, Gemini, Qwen, DeepSeek, Gemma y más) +- **MCP Tool Calling**: 19 parsers (OpenAI, Anthropic, Gemini, Qwen, DeepSeek, Gemma y más) - **Salida estructurada**: JSON Schema vía `response_format` (lm-format-enforcer) ### Throughput y memoria @@ -64,7 +64,7 @@ claude - **STT**: familia Whisper con RTF hasta 197x en M4 Max ### Razonamiento y avanzado -- **Extracción de razonamiento**: Qwen3, DeepSeek-R1 (`--reasoning-parser`) +- **Extracción de razonamiento**: Qwen3, DeepSeek-R1, DeepSeek-V4 (`--reasoning-parser`) - **Reducción de expertos MoE**: `--moe-top-k` para +7-16% en Qwen3-30B-A3B - **Decodificación especulativa**: `--mtp` para Qwen3-Next - **Prefill disperso**: `--spec-prefill` basado en atención para reducir TTFT diff --git a/README.fr.md b/README.fr.md index 1155cba81..2a938770d 100644 --- a/README.fr.md +++ b/README.fr.md @@ -46,7 +46,7 @@ claude ### APIs - **Compatible OpenAI** : `/v1/chat/completions`, `/v1/completions`, `/v1/embeddings`, `/v1/rerank`, `/v1/responses` - **Compatible Anthropic** : `/v1/messages` (streaming, tool use, system prompts) -- **MCP Tool Calling** : 12 parsers (OpenAI, Anthropic, Gemini, Qwen, DeepSeek, Gemma et plus) +- **MCP Tool Calling** : 19 parsers (OpenAI, Anthropic, Gemini, Qwen, DeepSeek, Gemma et plus) - **Sortie structurée** : JSON Schema via `response_format` (lm-format-enforcer) ### Débit et mémoire @@ -64,7 +64,7 @@ claude - **STT** : famille Whisper avec RTF jusqu'à 197x sur M4 Max ### Raisonnement et avancé -- **Extraction du raisonnement** : Qwen3, DeepSeek-R1 (`--reasoning-parser`) +- **Extraction du raisonnement** : Qwen3, DeepSeek-R1, DeepSeek-V4 (`--reasoning-parser`) - **Réduction d'experts MoE** : `--moe-top-k` pour +7-16% sur Qwen3-30B-A3B - **Décodage spéculatif** : `--mtp` pour Qwen3-Next - **Prefill creux** : `--spec-prefill` basé sur l'attention pour réduire le TTFT diff --git a/README.md b/README.md index 0051f807d..bb2a26629 100644 --- a/README.md +++ b/README.md @@ -46,7 +46,7 @@ claude ### APIs - **OpenAI-compatible**: `/v1/chat/completions`, `/v1/completions`, `/v1/embeddings`, `/v1/rerank`, `/v1/responses` - **Anthropic-compatible**: `/v1/messages` (streaming, tool use, system prompts) -- **MCP Tool Calling**: 12 parsers (OpenAI, Anthropic, Gemini, Qwen, DeepSeek, Gemma, and more) +- **MCP Tool Calling**: 19 parsers (OpenAI, Anthropic, Gemini, Qwen, DeepSeek, Gemma, and more) - **Structured output**: JSON Schema via `response_format` (lm-format-enforcer) ### Throughput & memory @@ -64,7 +64,7 @@ claude - **STT**: Whisper family with RTF up to 197x on M4 Max ### Reasoning & advanced -- **Reasoning extraction**: Qwen3, DeepSeek-R1 (`--reasoning-parser`) +- **Reasoning extraction**: Qwen3, DeepSeek-R1, DeepSeek-V4 (`--reasoning-parser`) - **MoE expert reduction**: `--moe-top-k` for +7-16% on Qwen3-30B-A3B - **Speculative decoding**: `--mtp` for Qwen3-Next - **Sparse prefill**: attention-based `--spec-prefill` for TTFT reduction diff --git a/README.zh.md b/README.zh.md index c3d52dea2..7da2e6958 100644 --- a/README.zh.md +++ b/README.zh.md @@ -46,7 +46,7 @@ claude ### API - **兼容 OpenAI**：`/v1/chat/completions`、`/v1/completions`、`/v1/embeddings`、`/v1/rerank`、`/v1/responses` - **兼容 Anthropic**：`/v1/messages`（流式、工具调用、system prompts） -- **MCP 工具调用**：12 种解析器（OpenAI、Anthropic、Gemini、Qwen、DeepSeek、Gemma 等） +- **MCP 工具调用**：19 种解析器（OpenAI、Anthropic、Gemini、Qwen、DeepSeek、Gemma 等） - **结构化输出**：通过 `response_format` 的 JSON Schema（基于 lm-format-enforcer） ### 吞吐与内存 @@ -64,7 +64,7 @@ claude - **STT**：Whisper 系列，M4 Max 上 RTF 最高可达 197 倍 ### 推理与高级功能 -- **思维链提取**：Qwen3、DeepSeek-R1（`--reasoning-parser`） +- **思维链提取**：Qwen3、DeepSeek-R1、DeepSeek-V4（`--reasoning-parser`） - **MoE 专家裁剪**：`--moe-top-k`，Qwen3-30B-A3B 上 +7-16% - **投机解码**：`--mtp`，用于 Qwen3-Next - **稀疏 prefill**：基于注意力的 `--spec-prefill`，降低 TTFT diff --git a/docs/development/architecture.md b/docs/development/architecture.md index 7965f868b..9997a254b 100644 --- a/docs/development/architecture.md +++ b/docs/development/architecture.md @@ -135,7 +135,7 @@ vllm_mlx/ ├── models/ │ ├── llm.py # MLXLanguageModel │ └── mllm.py # MLXMultimodalLM -├── tool_parsers/ # Tool call parsers (12 formats) +├── tool_parsers/ # Tool call parsers (19 formats) ├── reasoning_parsers/ # Reasoning parsers (qwen3, deepseek_r1) ├── server.py # FastAPI server ├── engine_core.py # AsyncEngineCore diff --git a/docs/es/guides/mcp-tools.md b/docs/es/guides/mcp-tools.md index 6577cf3b9..0fefd1e54 100644 --- a/docs/es/guides/mcp-tools.md +++ b/docs/es/guides/mcp-tools.md @@ -228,7 +228,7 @@ python examples/mcp_chat.py ## Formatos de herramientas soportados -vllm-mlx soporta 12 tool call parsers que cubren todas las familias de modelos principales. Consulta [Tool Calling](tool-calling.md) para ver la lista completa de parsers, alias y ejemplos. +vllm-mlx soporta 19 tool call parsers que cubren todas las familias de modelos principales. Consulta [Tool Calling](tool-calling.md) para ver la lista completa de parsers, alias y ejemplos. ## Seguridad diff --git a/docs/fr/guides/mcp-tools.md b/docs/fr/guides/mcp-tools.md index 77b97a434..f54b0d6fb 100644 --- a/docs/fr/guides/mcp-tools.md +++ b/docs/fr/guides/mcp-tools.md @@ -228,7 +228,7 @@ python examples/mcp_chat.py ## Formats d'outils pris en charge -vllm-mlx prend en charge 12 tool call parsers couvrant toutes les grandes familles de modèles. Voir [Tool Calling](tool-calling.md) pour la liste complète des parsers, alias et exemples. +vllm-mlx prend en charge 19 tool call parsers couvrant toutes les grandes familles de modèles. Voir [Tool Calling](tool-calling.md) pour la liste complète des parsers, alias et exemples. ## Sécurité diff --git a/docs/guides/mcp-tools.md b/docs/guides/mcp-tools.md index 605ae6823..126028ceb 100644 --- a/docs/guides/mcp-tools.md +++ b/docs/guides/mcp-tools.md @@ -228,7 +228,7 @@ python examples/mcp_chat.py ## Supported Tool Formats -vllm-mlx supports 12 tool call parsers covering all major model families. See [Tool Calling](tool-calling.md) for the full list of parsers, aliases, and examples. +vllm-mlx supports 19 tool call parsers covering all major model families. See [Tool Calling](tool-calling.md) for the full list of parsers, aliases, and examples. ## Security diff --git a/docs/guides/reasoning.md b/docs/guides/reasoning.md index 45d0ed206..085affd53 100644 --- a/docs/guides/reasoning.md +++ b/docs/guides/reasoning.md @@ -131,6 +131,24 @@ For DeepSeek-R1 models that may omit the opening `` tag. vllm-mlx serve mlx-community/DeepSeek-R1-Distill-Qwen-7B-4bit --reasoning-parser deepseek_r1 ``` +### DeepSeek-V4 Parser (`deepseek_v4`) + +DeepSeek V4 Flash can end an implicit reasoning block by opening a DSML tool +call, even when `` is absent. Use the V4 reasoning parser together with +the V4 tool parser so the DSML tail is routed to structured tool-call output: + +```bash +vllm-mlx serve deepseek-ai/DeepSeek-V4-Flash-0731 \ + --reasoning-parser deepseek_v4 \ + --enable-auto-tool-choice \ + --tool-call-parser deepseek_v4 +``` + +The encoder detects the published preview and 0731 reasoning-effort profiles. +For OpenAI requests, unspecified effort selects `high`; `minimal`, `low`, and +`medium` select `low`; `high` and `xhigh` select `high`; and `max` selects +`max`. The preview profile normalizes `low` to its `high` behavior. + ## How It Works The reasoning parser uses text-based detection to identify thinking tags in the model output. During streaming, it tracks the current position in the output to correctly route each token to either `reasoning` or `content`. diff --git a/docs/guides/tool-calling.md b/docs/guides/tool-calling.md index c927f7546..8327bb024 100644 --- a/docs/guides/tool-calling.md +++ b/docs/guides/tool-calling.md @@ -54,12 +54,19 @@ Use `--tool-call-parser` to select a parser for your model family: | `auto` | | Any model | Auto-detects format (tries all parsers) | | `mistral` | | Mistral, Devstral | `[TOOL_CALLS]` JSON array | | `qwen` | `qwen3` | Qwen, Qwen3 | `` XML or `[Calling tool:]` | +| `qwen3_xml` | `qwen3.5`, `qwen3_coder` | Qwen 3.5, Qwen3 Coder | Typed streaming XML | | `llama` | `llama3`, `llama4` | Llama 3.x, 4.x | `<\|python_tag\|>` JSON, bare JSON, or `` tags | | `hermes` | `nous` | Hermes, NousResearch | `` JSON in XML | | `deepseek` | `deepseek_v3`, `deepseek_r1` | DeepSeek V3, R1 | Unicode delimiters | +| `deepseek_v4` | `dsml` | DeepSeek V4 Flash / Flash-0731 | DSML typed parameters | +| `gemma4` | | Gemma 4 | `<\|tool_call>` call blocks | | `kimi` | `kimi_k2`, `moonshot` | Kimi K2, Moonshot | `<\|tool_call_begin\|>` tokens | | `granite` | `granite3` | IBM Granite 3.x, 4.x | `<\|tool_call\|>` or `` | | `nemotron` | `nemotron3` | NVIDIA Nemotron | `` | +| `minimax` | `minimax_m2` | MiniMax M2 | `` XML | +| `harmony` | `gpt-oss` | GPT-OSS | Harmony commentary control tokens | +| `poolside_v1` | | Poolside v1 / Laguna | Typed ``/`` XML | +| `step3p5` | `step` | Step 3.5 | `` XML | | `xlam` | | Salesforce xLAM | JSON with `tool_calls` array | | `functionary` | `meetkai` | MeetKai Functionary | Multiple function blocks | | `glm47` | `glm4` | GLM-4.7, GLM-4.7-Flash | `` with ``/`` XML | @@ -102,6 +109,48 @@ vllm-mlx serve mlx-community/DeepSeek-V3-0324-4bit \ --enable-auto-tool-choice --tool-call-parser deepseek ``` +DeepSeek V4 uses a different wire format and must use both V4 parsers. The +programmatic encoder supports the published `DeepSeek-V4-Flash` preview and +the `DeepSeek-V4-Flash-0731` profile: + +```bash +vllm-mlx serve deepseek-ai/DeepSeek-V4-Flash-0731 \ + --reasoning-parser deepseek_v4 \ + --enable-auto-tool-choice \ + --tool-call-parser deepseek_v4 +``` + +The preview profile accepts `high` and `max`; the 0731 profile accepts `low`, +`high`, and `max`. OpenAI `reasoning_effort` values are normalized as follows: +unspecified → `high`, `minimal`/`low`/`medium` → `low`, `high`/`xhigh` → +`high`, and `max` → `max`. On the preview checkpoint, `low` normalizes to +`high`. Streaming buffers partial DSML markers and emits exactly one structured +tool call after the block is complete; truncated markup is returned as text. + +```python +from openai import OpenAI + +client = OpenAI(base_url="http://localhost:8000/v1", api_key="unused") +response = client.chat.completions.create( + model="deepseek-ai/DeepSeek-V4-Flash-0731", + messages=[{"role": "user", "content": "What is the weather in Prague?"}], + reasoning_effort="high", + tools=[{ + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + }], +) +print(response.choices[0].message.tool_calls[0].function.arguments) +``` + ### IBM Granite ```bash @@ -152,12 +201,14 @@ vllm-mlx serve mlx-community/Qwen3-4B-4bit \ ``` The auto parser tries formats in this order: -1. Mistral (`[TOOL_CALLS]`) -2. Qwen bracket (`[Calling tool:]`) -3. Nemotron (``) -4. Qwen/Hermes XML (`{...}`) -5. Llama (`{...}`) -6. Raw JSON +1. Gemma 4 (`<|tool_call>call:name...`) +2. DeepSeek V4 DSML (`<｜DSML｜tool_calls>`) +3. Mistral (`[TOOL_CALLS]`) +4. Qwen bracket (`[Calling tool:]`) +5. Nemotron (``) +6. Qwen/Hermes XML (`{...}`) +7. Llama (`{...}`) +8. Raw JSON ## Streaming Tool Calls diff --git a/docs/zh/guides/mcp-tools.md b/docs/zh/guides/mcp-tools.md index 3b52c0ab2..a9b197149 100644 --- a/docs/zh/guides/mcp-tools.md +++ b/docs/zh/guides/mcp-tools.md @@ -228,7 +228,7 @@ python examples/mcp_chat.py ## 支持的工具格式 -vllm-mlx 支持 12 种 tool call parser，覆盖所有主流模型系列。完整的 parser 列表、别名及示例请参见 [Tool Calling](tool-calling.md)。 +vllm-mlx 支持 19 种 tool call parser，覆盖所有主流模型系列。完整的 parser 列表、别名及示例请参见 [Tool Calling](tool-calling.md)。 ## 安全性 diff --git a/tests/test_deepseek_v4_encoding.py b/tests/test_deepseek_v4_encoding.py index 0b0b3e1e1..7624374bf 100644 --- a/tests/test_deepseek_v4_encoding.py +++ b/tests/test_deepseek_v4_encoding.py @@ -16,6 +16,7 @@ THINKING_END_TOKEN, THINKING_START_TOKEN, apply_chat_template, + detect_reasoning_effort_profile, encode_arguments_to_dsml, encode_messages, install, @@ -129,15 +130,16 @@ class TestResolveThinking: @pytest.mark.parametrize( "kwargs,expected", [ - ({}, ("thinking", None)), + ({}, ("thinking", "high")), ({"enable_thinking": False}, ("chat", None)), - ({"enable_thinking": True}, ("thinking", None)), + ({"enable_thinking": True}, ("thinking", "high")), ({"reasoning_effort": "none"}, ("chat", None)), ({"reasoning_effort": "low"}, ("thinking", "low")), - ({"reasoning_effort": "medium"}, ("thinking", "high")), + ({"reasoning_effort": "minimal"}, ("thinking", "low")), + ({"reasoning_effort": "medium"}, ("thinking", "low")), ({"reasoning_effort": "high"}, ("thinking", "high")), ({"reasoning_effort": "max"}, ("thinking", "max")), - ({"reasoning_effort": "xhigh"}, ("thinking", "max")), + ({"reasoning_effort": "xhigh"}, ("thinking", "high")), ({"thinking_mode": "chat"}, ("chat", None)), ({"enable_thinking": False, "reasoning_effort": "max"}, ("chat", None)), ], @@ -153,6 +155,55 @@ def test_invalid_thinking_mode_rejected(self): with pytest.raises(ValueError, match="thinking_mode"): resolve_thinking(thinking_mode="bogus") + @pytest.mark.parametrize( + "effort,expected", + [ + (None, "high"), + ("minimal", "high"), + ("low", "high"), + ("medium", "high"), + ("high", "high"), + ("xhigh", "high"), + ("max", "max"), + ("unknown", "high"), + ], + ) + def test_preview_profile_mapping(self, effort, expected): + assert resolve_thinking( + reasoning_effort=effort, reasoning_effort_profile="preview" + ) == ("thinking", expected) + + @pytest.mark.parametrize( + "profile,effort,marker", + [ + ("preview", "high", None), + ("preview", "max", "Absolute maximum"), + ("official", "low", None), + ("official", "high", "Absolute maximum"), + ("official", "max", "Beyond maximum"), + ], + ) + def test_profile_golden_prefix(self, profile, effort, marker): + prompt = apply_chat_template( + [{"role": "user", "content": "Hi"}], + reasoning_effort=effort, + reasoning_effort_profile=profile, + ) + if marker is None: + assert "Reasoning Effort:" not in prompt + else: + assert marker in prompt + + def test_detects_published_model_profiles(self): + assert ( + detect_reasoning_effort_profile("deepseek-ai/DeepSeek-V4-Flash") + == "preview" + ) + assert ( + detect_reasoning_effort_profile("deepseek-ai/DeepSeek-V4-Flash-0731") + == "official" + ) + class TestTools: def test_schema_lands_in_system_message(self): diff --git a/tests/test_deepseek_v4_reasoning.py b/tests/test_deepseek_v4_reasoning.py index c70858302..dbac6a4a6 100644 --- a/tests/test_deepseek_v4_reasoning.py +++ b/tests/test_deepseek_v4_reasoning.py @@ -139,3 +139,13 @@ def test_split_think_tag_does_not_leak(self, chunk): assert content.strip() == "Answer." for fragment in (""): assert fragment not in reasoning + + def test_finalize_releases_ordinary_partial_marker(self, parser): + """A final ordinary ``<`` is text, not a marker to drop at EOS.""" + message = parser.extract_reasoning_streaming("", "2 <", "2 <") + final = parser.finalize_stream() + + assert message is not None + assert message.reasoning == "2 " + assert final is not None + assert final.reasoning == "<" diff --git a/tests/test_deepseek_v4_tool_parser.py b/tests/test_deepseek_v4_tool_parser.py index a5782faa6..163b9f5af 100644 --- a/tests/test_deepseek_v4_tool_parser.py +++ b/tests/test_deepseek_v4_tool_parser.py @@ -281,6 +281,24 @@ def test_reset_clears_state_between_requests(self): result = parser.extract_tool_calls_streaming("", "plain", "plain") assert result == {"content": "plain"} + @pytest.mark.parametrize("text", ["2 <", f'{TOOL_CALLS_START}\n<{D}invoke name="f']) + def test_finalize_releases_unparsed_suffix(self, text): + """EOS must not discard a partial marker or malformed DSML block.""" + parser = DeepSeekV4ToolParser() + previous, content = "", [] + for delta in text: + current = previous + delta + result = parser.extract_tool_calls_streaming(previous, current, delta) + if result and result.get("content"): + content.append(result["content"]) + previous = current + + result = parser.finalize_streaming(previous) + if result and result.get("content"): + content.append(result["content"]) + + assert "".join(content) == text + class TestAutoDetection: def test_auto_parser_routes_dsml(self): @@ -291,6 +309,28 @@ def test_auto_parser_routes_dsml(self): assert result.tools_called assert result.tool_calls[0]["name"] == "f" + def test_auto_parser_routes_character_streaming_dsml(self): + from vllm_mlx.tool_parsers.auto_tool_parser import AutoToolParser + + text = "Checking.\n\n" + block(invoke("f", param("a", "1", is_str=True))) + parser = AutoToolParser() + parser.reset() + previous, content, calls = "", [], [] + for delta in text: + current = previous + delta + result = parser.extract_tool_calls_streaming(previous, current, delta) + if result: + if result.get("content"): + content.append(result["content"]) + if result.get("tool_calls"): + calls.append(result["tool_calls"]) + previous = current + + assert "".join(content).strip() == "Checking." + assert D not in "".join(content) + assert len(calls) == 1 + assert calls[0][0]["function"]["name"] == "f" + class TestChainedWithReasoningParser: """The server feeds the reasoning parser's content to the tool parser. diff --git a/tests/test_responses_api.py b/tests/test_responses_api.py index 9e5cd966e..df375f4cb 100644 --- a/tests/test_responses_api.py +++ b/tests/test_responses_api.py @@ -790,6 +790,75 @@ def test_streaming_response_bracket_tool_call_does_not_leak_text( assert len(function_call_deltas) == 1 assert function_call_deltas[0]["delta"] == '{"a": 1, "b": 2}' + def test_deepseek_v4_reasoning_routes_dsml_through_tool_parser( + self, client, monkeypatch + ): + import vllm_mlx.server as srv + + d = "｜DSML｜" + engine = _mock_engine(_output("unused")) + engine._stream_outputs = [ + _stream_output("thinkingChecking.\n\n"), + _stream_output("<"), + _stream_output(d), + _stream_output("tool_calls>\n"), + _stream_output(f'<{d}invoke name="get_weather">\n'), + _stream_output( + f'<{d}parameter name="city" string="true">Prague' f"\n" + ), + _stream_output(f"\n"), + _stream_output(f"", finish_reason="stop"), + ] + srv._engine = engine + monkeypatch.setattr(srv, "_enable_auto_tool_choice", True) + monkeypatch.setattr(srv, "_tool_call_parser", "deepseek_v4") + monkeypatch.setattr(srv, "_tool_parser_instance", None) + monkeypatch.setattr(srv, "_reasoning_parser_name", "deepseek_v4") + monkeypatch.setattr(srv, "_reasoning_parser", None) + + with client.stream( + "POST", + "/v1/responses", + json={ + "model": "test-model", + "input": "Check the weather", + "stream": True, + "tools": [ + { + "type": "function", + "name": "get_weather", + "parameters": {"type": "object"}, + } + ], + }, + ) as resp: + body = "".join(resp.iter_text()) + + events = _parse_sse_events(body) + text_deltas = [ + payload["delta"] + for event_type, payload in events + if event_type == "response.output_text.delta" + ] + function_items = [ + payload["item"] + for event_type, payload in events + if event_type == "response.output_item.added" + and payload["item"]["type"] == "function_call" + ] + completed = next( + payload["response"] + for event_type, payload in events + if event_type == "response.completed" + ) + + assert resp.status_code == 200 + assert "".join(text_deltas).strip() == "Checking." + assert d not in "".join(text_deltas) + assert d not in completed["output_text"] + assert len(function_items) == 1 + assert function_items[0]["name"] == "get_weather" + def test_streaming_response_without_tools_keeps_llama_shaped_json( self, client, monkeypatch ): diff --git a/tests/test_server.py b/tests/test_server.py index 617c12bce..8ab479a75 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -2239,6 +2239,223 @@ async def stream_chat(self, messages, **kwargs): "total_tokens": 12, } + @pytest.mark.anyio + @pytest.mark.parametrize("parser_name", ["deepseek_v4", "auto"]) + async def test_deepseek_v4_split_marker_never_leaks(self, monkeypatch, parser_name): + """The request-local DSML parser must see text before the split marker.""" + from vllm_mlx.engine.base import GenerationOutput + from vllm_mlx.server import ( + ChatCompletionRequest, + Message, + stream_chat_completion, + ) + import vllm_mlx.server as server + + d = "｜DSML｜" + deltas = [ + "Checking.\n\n", + "<", + d, + "tool_calls>\n", + f'<{d}invoke name="get_weather">\n', + f'<{d}parameter name="city" string="true">Prague\n', + f"\n", + f"", + ] + + class FakeEngine: + model_name = "fake-engine" + tokenizer = None + + async def stream_chat(self, messages, **kwargs): + for index, delta in enumerate(deltas): + finished = index == len(deltas) - 1 + yield GenerationOutput( + text="", + new_text=delta, + finished=finished, + finish_reason="stop" if finished else None, + ) + + monkeypatch.setattr(server, "_model_name", "served-model") + monkeypatch.setattr(server, "_reasoning_parser", None) + monkeypatch.setattr(server, "_enable_auto_tool_choice", True) + monkeypatch.setattr(server, "_tool_call_parser", parser_name) + monkeypatch.setattr(server, "_tool_parser_instance", None) + + request = ChatCompletionRequest( + model="served-model", + messages=[Message(role="user", content="hi")], + tools=[ + { + "type": "function", + "function": { + "name": "get_weather", + "parameters": {"type": "object"}, + }, + } + ], + stream=True, + ) + chunks = [ + chunk + async for chunk in stream_chat_completion( + FakeEngine(), request.messages, request + ) + ] + payloads = [ + json.loads(chunk.removeprefix("data: ").strip()) + for chunk in chunks + if chunk != "data: [DONE]\n\n" + ] + content = "".join( + payload["choices"][0]["delta"].get("content") or "" + for payload in payloads + if payload["choices"] + ) + tool_payloads = [ + payload + for payload in payloads + if payload["choices"] and payload["choices"][0]["delta"].get("tool_calls") + ] + + assert content.strip() == "Checking." + assert d not in content + assert len(tool_payloads) == 1 + assert ( + tool_payloads[0]["choices"][0]["delta"]["tool_calls"][0]["function"]["name"] + == "get_weather" + ) + + @pytest.mark.anyio + async def test_deepseek_v4_truncated_dsml_is_visible_at_eos(self, monkeypatch): + from vllm_mlx.engine.base import GenerationOutput + from vllm_mlx.server import ( + ChatCompletionRequest, + Message, + stream_chat_completion, + ) + import vllm_mlx.server as server + + d = "｜DSML｜" + truncated = f'Before <{d}tool_calls>\n<{d}invoke name="f' + model_output = "thinking" + truncated + + class FakeEngine: + model_name = "fake-engine" + tokenizer = None + + async def stream_chat(self, messages, **kwargs): + yield GenerationOutput( + text="", + new_text=model_output, + finished=False, + ) + yield GenerationOutput( + text="", + new_text="", + finished=True, + finish_reason="length", + ) + + monkeypatch.setattr(server, "_model_name", "served-model") + monkeypatch.setattr(server, "_reasoning_parser_name", "deepseek_v4") + monkeypatch.setattr(server, "_reasoning_parser", None) + monkeypatch.setattr(server, "_enable_auto_tool_choice", True) + monkeypatch.setattr(server, "_tool_call_parser", "deepseek_v4") + monkeypatch.setattr(server, "_tool_parser_instance", None) + + request = ChatCompletionRequest( + model="served-model", + messages=[Message(role="user", content="hi")], + tools=[ + { + "type": "function", + "function": { + "name": "f", + "parameters": {"type": "object"}, + }, + } + ], + stream=True, + ) + chunks = [ + chunk + async for chunk in stream_chat_completion( + FakeEngine(), request.messages, request + ) + ] + payloads = [ + json.loads(chunk.removeprefix("data: ").strip()) + for chunk in chunks + if chunk != "data: [DONE]\n\n" + ] + content = "".join( + payload["choices"][0]["delta"].get("content") or "" + for payload in payloads + if payload["choices"] + ) + + assert content == truncated + assert payloads[-1]["choices"][0]["finish_reason"] == "length" + + @pytest.mark.anyio + async def test_deepseek_v4_reasoning_flushes_final_lt_and_terminal( + self, monkeypatch + ): + from vllm_mlx.engine.base import GenerationOutput + from vllm_mlx.server import ( + ChatCompletionRequest, + Message, + stream_chat_completion, + ) + import vllm_mlx.server as server + + class FakeEngine: + model_name = "fake-engine" + tokenizer = None + + async def stream_chat(self, messages, **kwargs): + yield GenerationOutput(text="", new_text="2 ", finished=False) + yield GenerationOutput( + text="", + new_text="<", + finished=True, + finish_reason="stop", + ) + + monkeypatch.setattr(server, "_model_name", "served-model") + monkeypatch.setattr(server, "_reasoning_parser_name", "deepseek_v4") + monkeypatch.setattr(server, "_reasoning_parser", None) + monkeypatch.setattr(server, "_enable_auto_tool_choice", False) + monkeypatch.setattr(server, "_tool_call_parser", None) + monkeypatch.setattr(server, "_tool_parser_instance", None) + + request = ChatCompletionRequest( + model="served-model", + messages=[Message(role="user", content="hi")], + stream=True, + ) + chunks = [ + chunk + async for chunk in stream_chat_completion( + FakeEngine(), request.messages, request + ) + ] + payloads = [ + json.loads(chunk.removeprefix("data: ").strip()) + for chunk in chunks + if chunk != "data: [DONE]\n\n" + ] + reasoning = "".join( + payload["choices"][0]["delta"].get("reasoning_content") or "" + for payload in payloads + if payload["choices"] + ) + + assert reasoning == "2 <" + assert payloads[-1]["choices"][0]["finish_reason"] == "stop" + @pytest.mark.anyio async def test_stream_without_parser_flags_keeps_plain_text(self, monkeypatch): """Generic streaming fallback should not interfere with normal text.""" diff --git a/vllm_mlx/reasoning/deepseek_v4_parser.py b/vllm_mlx/reasoning/deepseek_v4_parser.py index 0da7ccf0d..7328b941c 100644 --- a/vllm_mlx/reasoning/deepseek_v4_parser.py +++ b/vllm_mlx/reasoning/deepseek_v4_parser.py @@ -43,12 +43,14 @@ def __init__(self, tokenizer=None): # and whether the stream has crossed into tool markup. self._emitted_len: int = 0 self._in_tool_markup: bool = False + self._current_text: str = "" def reset_state(self): """Reset state machine for a new streaming request.""" super().reset_state() self._emitted_len = 0 self._in_tool_markup = False + self._current_text = "" def extract_reasoning( self, @@ -83,6 +85,7 @@ def extract_reasoning_streaming( reasoning channel and then repeats the whole marker as content once it recognises it. """ + self._current_text = current_text if self._in_tool_markup: new = current_text[self._emitted_len :] self._emitted_len = len(current_text) @@ -128,3 +131,14 @@ def extract_reasoning_streaming( return super().extract_reasoning_streaming( effective_previous, effective_current, effective_delta ) + + def finalize_stream(self) -> DeltaMessage | None: + """Release an incomplete marker prefix retained at end of stream.""" + if self._emitted_len >= len(self._current_text): + return None + + pending = self._current_text[self._emitted_len :] + self._emitted_len = len(self._current_text) + if self._in_tool_markup or self._phase == "content": + return DeltaMessage(content=pending) + return DeltaMessage(reasoning=pending) diff --git a/vllm_mlx/server.py b/vllm_mlx/server.py index 1f10701b7..7cf5fecee 100644 --- a/vllm_mlx/server.py +++ b/vllm_mlx/server.py @@ -2674,7 +2674,7 @@ def _start_reasoning_item() -> list[str]: tool_parser = _get_streaming_tool_parser(chat_request, engine) tool_accumulated_text = "" - tool_markup_possible = False + tool_markup_possible = _requires_eager_tool_streaming(tool_parser) async for output in engine.stream_chat(messages=messages, **chat_kwargs): last_output = output @@ -2689,7 +2689,10 @@ def _start_reasoning_item() -> list[str]: use_reasoning = reasoning_parser and not _thinking_disabled( request, chat_kwargs ) - if not delta_text and not (use_reasoning and output_finished): + if not delta_text and not ( + (use_reasoning and output_finished) + or (tool_parser and tool_markup_possible and output_finished) + ): continue previous_text = raw_accumulated_text @@ -2704,7 +2707,10 @@ def _start_reasoning_item() -> list[str]: finished=output_finished, ) if delta_msg is None: - continue + if output_finished and tool_parser and tool_markup_possible: + delta_msg = DeltaMessage() + else: + continue if delta_msg.reasoning: for event in _start_reasoning_item(): @@ -2722,10 +2728,29 @@ def _start_reasoning_item() -> list[str]: ) sequence += 1 - if delta_msg.content: + content = delta_msg.content or "" + if tool_parser and (content or (output_finished and tool_markup_possible)): + tool_accumulated_text, tool_result = _extract_streaming_tool_delta( + tool_parser, + tool_accumulated_text, + content, + tool_request_context, + ) + if output_finished and ( + tool_result is None or _requires_eager_tool_streaming(tool_parser) + ): + tool_result = _finalize_streaming_tool_result( + tool_parser, tool_accumulated_text, tool_result + ) + if tool_result is None or "tool_calls" in tool_result: + content = "" + else: + content = tool_result.get("content", "") + + if content: for event in _start_text_item(): yield event - accumulated_text += delta_msg.content + accumulated_text += content yield _responses_sse_event( "response.output_text.delta", ResponseOutputTextDeltaEvent( @@ -2733,7 +2758,7 @@ def _start_reasoning_item() -> list[str]: item_id=text_item_id, output_index=text_output_index, content_index=0, - delta=delta_msg.content, + delta=content, ), ) sequence += 1 @@ -2762,9 +2787,11 @@ def _start_reasoning_item() -> list[str]: delta_text, tool_request_context, ) - if tool_result is None and output.finished: + if output.finished and ( + tool_result is None or _requires_eager_tool_streaming(tool_parser) + ): tool_result = _finalize_streaming_tool_result( - tool_parser, tool_accumulated_text + tool_parser, tool_accumulated_text, tool_result ) if tool_result is None: continue @@ -3224,6 +3251,11 @@ def _streaming_tool_markup_possible(text: str, tool_parser=None) -> bool: ) +def _requires_eager_tool_streaming(tool_parser) -> bool: + """Return whether a parser must receive every delta from response start.""" + return bool(getattr(tool_parser, "REQUIRES_EAGER_STREAMING", False)) + + def _streaming_tool_markup_possible_after_delta( accumulated_text: str, delta_text: str, tool_parser=None ) -> bool: @@ -3242,12 +3274,28 @@ def _streaming_tool_markup_possible_after_delta( return _streaming_tool_markup_possible(check_text, tool_parser) -def _finalize_streaming_tool_result(tool_parser, current_text: str): +def _finalize_streaming_tool_result( + tool_parser, current_text: str, result: dict | None = None +): """Let parsers resolve an ambiguous prefix at end of generation.""" finalize = getattr(tool_parser, "finalize_streaming", None) if finalize is None: - return None - return finalize(current_text) + return result + final = finalize(current_text) + if final is None: + return result + if result is None: + return final + + merged = dict(result) + if final.get("content"): + merged["content"] = (merged.get("content") or "") + final["content"] + if final.get("tool_calls"): + merged["tool_calls"] = [ + *(merged.get("tool_calls") or []), + *final["tool_calls"], + ] + return merged def load_embedding_model( @@ -6028,8 +6076,8 @@ async def _stream_anthropic_messages( # Tool call streaming suppression — prevents raw tool markup from leaking # as text_delta events. Mirrors the OpenAI streaming path logic. tool_accumulated_text = "" - tool_markup_possible = False tool_parser = _get_streaming_tool_parser(openai_request, engine) + tool_markup_possible = _requires_eager_tool_streaming(tool_parser) tool_request_context = openai_request.model_dump() try: @@ -6046,7 +6094,10 @@ async def _stream_anthropic_messages( if hasattr(output, "completion_tokens") and output.completion_tokens: completion_tokens = output.completion_tokens - if not delta_text and not (use_reasoning and output_finished): + if not delta_text and not ( + (use_reasoning and output_finished) + or (tool_parser and tool_markup_possible and output_finished) + ): continue # Filter special tokens @@ -6084,9 +6135,12 @@ async def _stream_anthropic_messages( tool_request_context, ) ) - if tool_result is None and output.finished: + if output.finished and ( + tool_result is None + or _requires_eager_tool_streaming(tool_parser) + ): tool_result = _finalize_streaming_tool_result( - tool_parser, tool_accumulated_text + tool_parser, tool_accumulated_text, tool_result ) if tool_result is None: # Inside tool markup, so suppress this delta. @@ -6115,7 +6169,10 @@ async def _stream_anthropic_messages( ) if delta_msg is None: - continue + if output_finished and tool_parser and tool_markup_possible: + delta_msg = DeltaMessage() + else: + continue if delta_msg.reasoning: if not thinking_block_started: @@ -6148,9 +6205,12 @@ async def _stream_anthropic_messages( tool_request_context, ) ) - if tool_result is None and output.finished: + if output.finished and ( + tool_result is None + or _requires_eager_tool_streaming(tool_parser) + ): tool_result = _finalize_streaming_tool_result( - tool_parser, tool_accumulated_text + tool_parser, tool_accumulated_text, tool_result ) if tool_result is None: # Inside tool markup, so suppress this delta. @@ -6397,8 +6457,8 @@ async def stream_chat_completion( tool_parser = None tool_accumulated_text = "" tool_calls_detected = False - tool_markup_possible = False # Fast path: skip parsing until markers appear tool_parser = _get_streaming_tool_parser(request, engine) + tool_markup_possible = _requires_eager_tool_streaming(tool_parser) # Whether any emitted chunk carried a terminal finish_reason. The engine's # finished=True output can be swallowed by a parser `continue` below (e.g. # a bare end-of-turn token arriving after a completed tool call); without @@ -6439,8 +6499,11 @@ async def stream_chat_completion( ) if delta_msg is None: - # Skip this chunk (e.g., token itself) - continue + if output_finished and tool_parser and tool_markup_possible: + delta_msg = DeltaMessage() + else: + # Skip this chunk (e.g., token itself) + continue content = delta_msg.content reasoning = delta_msg.reasoning @@ -6459,7 +6522,9 @@ async def stream_chat_completion( reasoning = None # Tool call parsing on content portion - if tool_parser and content: + if tool_parser and ( + content or (output_finished and tool_markup_possible) + ): if ( not tool_markup_possible and not _streaming_tool_markup_possible_after_delta( @@ -6480,9 +6545,12 @@ async def stream_chat_completion( ) ) - if tool_result is None and output.finished: + if output.finished and ( + tool_result is None + or _requires_eager_tool_streaming(tool_parser) + ): tool_result = _finalize_streaming_tool_result( - tool_parser, tool_accumulated_text + tool_parser, tool_accumulated_text, tool_result ) if tool_result is None: @@ -6628,9 +6696,12 @@ async def stream_chat_completion( ) ) - if tool_result is None and output.finished: + if output.finished and ( + tool_result is None + or _requires_eager_tool_streaming(tool_parser) + ): tool_result = _finalize_streaming_tool_result( - tool_parser, tool_accumulated_text + tool_parser, tool_accumulated_text, tool_result ) if tool_result is None: diff --git a/vllm_mlx/tool_parsers/auto_tool_parser.py b/vllm_mlx/tool_parsers/auto_tool_parser.py index f2fa4225a..8af8a2a10 100644 --- a/vllm_mlx/tool_parsers/auto_tool_parser.py +++ b/vllm_mlx/tool_parsers/auto_tool_parser.py @@ -61,6 +61,19 @@ class AutoToolParser(ToolParser): ) BARE_BRACKET_PATTERN = re.compile(r"\[(\w+)\((\{.*?\})\)\]", re.DOTALL) BARE_BRACKET_PARTIAL_PATTERN = re.compile(r"\[\w+\($") + # Auto includes DSML, whose opening marker must be buffered from its first + # ``<`` rather than after the complete marker becomes visible. + REQUIRES_EAGER_STREAMING = True + + def __init__(self, tokenizer=None): + super().__init__(tokenizer) + self._dsml_parser = DeepSeekV4ToolParser(tokenizer) + self._legacy_stream_text = "" + + def reset(self) -> None: + super().reset() + self._dsml_parser.reset() + self._legacy_stream_text = "" def extract_tool_calls( self, model_output: str, request: dict[str, Any] | None = None @@ -361,7 +374,7 @@ def _parse_raw_json_tool_calls(self, text: str) -> list[dict[str, Any]]: return tool_calls - def extract_tool_calls_streaming( + def _extract_legacy_tool_calls_streaming( self, previous_text: str, current_text: str, @@ -424,3 +437,53 @@ def extract_tool_calls_streaming( } return None + + def extract_tool_calls_streaming( + self, + previous_text: str, + current_text: str, + delta_text: str, + previous_token_ids: Sequence[int] | None = None, + current_token_ids: Sequence[int] | None = None, + delta_token_ids: Sequence[int] | None = None, + request: dict[str, Any] | None = None, + ) -> dict[str, Any] | None: + """Run DSML buffering before the legacy auto-detection formats.""" + dsml_result = self._dsml_parser.extract_tool_calls_streaming( + previous_text, + current_text, + delta_text, + previous_token_ids, + current_token_ids, + delta_token_ids, + request, + ) + if dsml_result is None or "tool_calls" in dsml_result: + return dsml_result + + safe_delta = dsml_result.get("content", "") + safe_previous = self._legacy_stream_text + self._legacy_stream_text += safe_delta + return self._extract_legacy_tool_calls_streaming( + safe_previous, + self._legacy_stream_text, + safe_delta, + previous_token_ids, + current_token_ids, + delta_token_ids, + request, + ) + + def finalize_streaming(self, current_text: str) -> dict[str, Any] | None: + """Flush DSML prefixes or truncated blocks at end of generation.""" + result = self._dsml_parser.finalize_streaming(current_text) + if result is None or "tool_calls" in result: + return result + + safe_delta = result.get("content", "") + safe_previous = self._legacy_stream_text + self._legacy_stream_text += safe_delta + legacy = self._extract_legacy_tool_calls_streaming( + safe_previous, self._legacy_stream_text, safe_delta + ) + return legacy if legacy is not None else {"content": safe_delta} diff --git a/vllm_mlx/tool_parsers/deepseek_v4_tool_parser.py b/vllm_mlx/tool_parsers/deepseek_v4_tool_parser.py index 3b7fe7b88..eff8497c2 100644 --- a/vllm_mlx/tool_parsers/deepseek_v4_tool_parser.py +++ b/vllm_mlx/tool_parsers/deepseek_v4_tool_parser.py @@ -83,6 +83,10 @@ class DeepSeekV4ToolParser(ToolParser): # (api/utils.py), so the model would see a shape it was never trained on # and the encoder's own handling would never run. SUPPORTS_NATIVE_TOOL_FORMAT = True + # DSML's opening marker is tokenized as ``<``, ``｜DSML｜`` and + # ``tool_calls>``. The routing layer therefore cannot wait for a complete + # marker before invoking this state machine. + REQUIRES_EAGER_STREAMING = True def __init__(self, tokenizer=None): super().__init__(tokenizer) @@ -295,3 +299,24 @@ def extract_tool_calls_streaming( emit = text[: len(text) - hold] if hold else text self._emitted_len += len(emit) return {"content": emit} if emit else None + + def finalize_streaming(self, current_text: str) -> dict[str, Any] | None: + """Resolve withheld text when generation ends. + + A complete invoke without the outer closing marker is parsed just as + it is in non-streaming mode. Anything else is returned byte-for-byte + so max-token truncation cannot silently discard model output. + """ + if self._emitted_len >= len(current_text): + return None + + result = self.extract_tool_calls(current_text) + if result.tools_called: + self._pending = "" + self._emitted_len = len(current_text) + return self._format_streaming(result) + + content = current_text[self._emitted_len :] + self._pending = "" + self._emitted_len = len(current_text) + return {"content": content} if content else None diff --git a/vllm_mlx/utils/deepseek_v4_encoding.py b/vllm_mlx/utils/deepseek_v4_encoding.py index d5d2e8d29..224f8a727 100644 --- a/vllm_mlx/utils/deepseek_v4_encoding.py +++ b/vllm_mlx/utils/deepseek_v4_encoding.py @@ -20,8 +20,10 @@ """ import copy +import ast import json import logging +from pathlib import Path from typing import Any logger = logging.getLogger(__name__) @@ -120,6 +122,15 @@ def partial_marker_len(text: str, *markers: str) -> int: ), } DEFAULT_REASONING_EFFORT = "low" +REASONING_EFFORT_PROFILES: dict[str, dict[str, str]] = { + # The first published Flash checkpoint accepts only high/max. Its max + # prefix became the official profile's high prefix. + "preview": { + "high": "", + "max": REASONING_EFFORT_PROMPTS["high"], + }, + "official": REASONING_EFFORT_PROMPTS, +} TOOLS_TEMPLATE = """## Tools @@ -239,6 +250,7 @@ def render_message( thinking_mode: str, drop_thinking: bool = True, reasoning_effort: str | None = None, + reasoning_effort_profile: str = "official", ) -> str: """Render a single message into its encoded form. @@ -274,14 +286,21 @@ def render_message( if tool_calls: tool_calls = tool_calls_from_openai_format(tool_calls) - reasoning_effort = reasoning_effort or DEFAULT_REASONING_EFFORT - if reasoning_effort not in REASONING_EFFORT_PROMPTS: + if reasoning_effort_profile not in REASONING_EFFORT_PROFILES: + raise ValueError( + f"Invalid reasoning effort profile: {reasoning_effort_profile}" + ) + effort_prompts = REASONING_EFFORT_PROFILES[reasoning_effort_profile] + reasoning_effort = reasoning_effort or ( + DEFAULT_REASONING_EFFORT if reasoning_effort_profile == "official" else "high" + ) + if reasoning_effort not in effort_prompts: raise ValueError( f"Invalid reasoning effort: {reasoning_effort}, expected one of " - f"{list(REASONING_EFFORT_PROMPTS)}" + f"{list(effort_prompts)}" ) if index == 0 and thinking_mode == "thinking": - prompt += REASONING_EFFORT_PROMPTS[reasoning_effort] + prompt += effort_prompts[reasoning_effort] if role == "system": prompt += content or "" @@ -542,6 +561,7 @@ def encode_messages( drop_thinking: bool = True, add_default_bos_token: bool = True, reasoning_effort: str | None = None, + reasoning_effort_profile: str = "official", ) -> str: """Encode a conversation into a DeepSeek-V4 prompt string. @@ -595,6 +615,7 @@ def encode_messages( thinking_mode=thinking_mode, drop_thinking=effective_drop_thinking, reasoning_effort=reasoning_effort, + reasoning_effort_profile=reasoning_effort_profile, ) return prompt @@ -604,17 +625,15 @@ def encode_messages( # OpenAI API adaptation # --------------------------------------------------------------------------- -# OpenAI exposes reasoning_effort as low/medium/high; DeepSeek-V4 defines -# low/high/max prompts plus "no thinking at all". "medium" has no distinct -# prompt of its own, so it maps onto "high" — as does any unrecognised value, -# which keeps a typo from silently disabling reasoning. +# Match the DeepSeek-V4 OpenAI wrapper in vLLM. The official 0731 profile has +# low/high/max; the earlier preview exposes high/max and normalizes low to high. _EFFORT_ALIASES = { "low": "low", "minimal": "low", - "medium": "high", + "medium": "low", "high": "high", "max": "max", - "xhigh": "max", + "xhigh": "high", } @@ -622,6 +641,7 @@ def resolve_thinking( enable_thinking: bool | None = None, reasoning_effort: str | None = None, thinking_mode: str | None = None, + reasoning_effort_profile: str = "official", ) -> tuple[str, str | None]: """Map OpenAI-style knobs onto ``(thinking_mode, reasoning_effort)``. @@ -639,19 +659,77 @@ def resolve_thinking( else: mode = "thinking" - if mode == "chat" or reasoning_effort in (None, "none"): + if mode == "chat": return mode, None - effort = _EFFORT_ALIASES.get(str(reasoning_effort).lower()) + if reasoning_effort_profile not in REASONING_EFFORT_PROFILES: + raise ValueError( + f"Invalid reasoning effort profile: {reasoning_effort_profile}" + ) + + if reasoning_effort is None: + effort = "high" + else: + effort = _EFFORT_ALIASES.get(str(reasoning_effort).lower()) + if effort is None: logger.warning( "Unknown reasoning_effort %r for deepseek_v4, treating as 'high'", reasoning_effort, ) effort = "high" + if reasoning_effort_profile == "preview" and effort == "low": + effort = "high" return mode, effort +def _profile_from_encoder_source(path: Path) -> str | None: + """Detect the checkpoint encoder profile using SGLang's stable symbols.""" + encoder = path / "encoding" / "encoding_dsv4.py" + if not encoder.is_file(): + return None + try: + tree = ast.parse(encoder.read_text(encoding="utf-8")) + except (OSError, SyntaxError, UnicodeError): + return None + + assignments = { + target.id: node.value + for node in tree.body + if isinstance(node, (ast.Assign, ast.AnnAssign)) + for target in (node.targets if isinstance(node, ast.Assign) else [node.target]) + if isinstance(target, ast.Name) + } + + try: + default = ast.literal_eval(assignments["DEFAULT_REASONING_EFFORT"]) + prompts = ast.literal_eval(assignments["REASONING_EFFORT_PROMPTS"]) + except (KeyError, ValueError, TypeError): + default = None + prompts = None + if ( + default == "low" + and isinstance(prompts, dict) + and {"low", "high", "max"}.issubset(prompts) + ): + return "official" + if "REASONING_EFFORT_MAX" in assignments: + return "preview" + return None + + +def detect_reasoning_effort_profile(model_name: str | None) -> str: + """Identify the preview or official Flash prompt profile without network I/O.""" + if model_name: + profile = _profile_from_encoder_source(Path(model_name).expanduser()) + if profile is not None: + return profile + if "deepseek-v4-flash-0731" in model_name.lower(): + return "official" + # SGLang also falls back to preview when checkpoint metadata is absent. + return "preview" + + def _attach_tools( conversation: list[dict[str, Any]], tools: list[dict] | None ) -> list[dict[str, Any]]: @@ -678,6 +756,7 @@ def apply_chat_template( tools: list[dict] | None = None, enable_thinking: bool | None = None, reasoning_effort: str | None = None, + reasoning_effort_profile: str = "official", thinking_mode: str | None = None, drop_thinking: bool = True, add_default_bos_token: bool = True, @@ -690,7 +769,12 @@ def apply_chat_template( the encoder always closes on the assistant prefix, which is the only mode the model was trained for. """ - mode, effort = resolve_thinking(enable_thinking, reasoning_effort, thinking_mode) + mode, effort = resolve_thinking( + enable_thinking, + reasoning_effort, + thinking_mode, + reasoning_effort_profile, + ) conversation = _attach_tools(conversation, tools) return encode_messages( conversation, @@ -698,10 +782,11 @@ def apply_chat_template( drop_thinking=drop_thinking, add_default_bos_token=add_default_bos_token, reasoning_effort=effort, + reasoning_effort_profile=reasoning_effort_profile, ) -def install(tokenizer: Any) -> Any: +def install(tokenizer: Any, model_name: str | None = None) -> Any: """Route ``tokenizer.apply_chat_template`` through the V4 encoder. DeepSeek-V4 carries no Jinja template, so the stock path either raises or @@ -714,7 +799,12 @@ def install(tokenizer: Any) -> Any: if getattr(tokenizer, "_deepseek_v4_encoding_installed", False): return tokenizer + profile = detect_reasoning_effort_profile( + model_name or getattr(tokenizer, "name_or_path", None) + ) + def _apply(conversation, tools=None, tokenize=False, **kwargs): + kwargs.setdefault("reasoning_effort_profile", profile) prompt = apply_chat_template(conversation, tools=tools, **kwargs) if tokenize: return tokenizer.encode(prompt) diff --git a/vllm_mlx/utils/tokenizer.py b/vllm_mlx/utils/tokenizer.py index 498c3f24d..9e12b00d2 100644 --- a/vllm_mlx/utils/tokenizer.py +++ b/vllm_mlx/utils/tokenizer.py @@ -72,7 +72,7 @@ def _install_custom_chat_template(model_name: str, tokenizer): from .deepseek_v4_encoding import install as install_deepseek_v4 - return install_deepseek_v4(tokenizer) + return install_deepseek_v4(tokenizer, model_name=model_name) def load_model_with_fallback(model_name: str, tokenizer_config: dict = None): From 295fda3dde0e5bcc43a5957d071d2d19757f94c6 Mon Sep 17 00:00:00 2001 From: Wayner Barrios Date: Wed, 26 Aug 2026 01:06:10 -0400 Subject: [PATCH 3/3] Finalize empty DeepSeek V4 streaming deltas safely (#676) --- tests/test_chat_template_kwargs.py | 90 ++++++++++++++++++++++++++++++ vllm_mlx/server.py | 12 +++- 2 files changed, 99 insertions(+), 3 deletions(-) diff --git a/tests/test_chat_template_kwargs.py b/tests/test_chat_template_kwargs.py index 18c87a24b..137a542cc 100644 --- a/tests/test_chat_template_kwargs.py +++ b/tests/test_chat_template_kwargs.py @@ -493,3 +493,93 @@ async def fake_stream_chat(messages, **kwargs): assert thinking == "thinking\n<{d}invoke name="f' + + async def fake_stream_chat(messages, **kwargs): + yield GenerationOutput( + text="", + new_text=("thinking" if use_reasoning else "") + truncated, + finished=False, + ) + yield GenerationOutput( + text="", + new_text="", + finished=True, + finish_reason="length", + prompt_tokens=3, + completion_tokens=2, + ) + + engine = MagicMock(stream_chat=fake_stream_chat, tokenizer=None) + messages = [{"role": "user", "content": "Think"}] + tools = [ + { + "type": "function", + "function": { + "name": "f", + "parameters": {"type": "object"}, + }, + } + ] + openai_request = srv.ChatCompletionRequest( + model="test-model", + messages=[srv.Message(**messages[0])], + max_tokens=8, + tools=tools, + ) + anthropic_request = srv.AnthropicRequest( + model="test-model", + max_tokens=8, + messages=messages, + tools=[ + { + "name": "f", + "input_schema": {"type": "object"}, + } + ], + ) + prepared = srv.PreparedChatInvocation( + messages=messages, + chat_kwargs={}, + response_format=None, + json_logits_processor=None, + ) + monkeypatch.setattr( + srv, "_reasoning_parser_name", "deepseek_v4" if use_reasoning else None + ) + monkeypatch.setattr(srv, "_reasoning_parser", None) + monkeypatch.setattr(srv, "_enable_auto_tool_choice", True) + monkeypatch.setattr(srv, "_tool_call_parser", "deepseek_v4") + monkeypatch.setattr(srv, "_tool_parser_instance", None) + monkeypatch.setattr(srv, "_model_name", "test-model") + + body = "".join( + [ + chunk + async for chunk in srv._stream_anthropic_messages( + engine, openai_request, anthropic_request, prepared + ) + ] + ) + events = [ + json.loads(line.removeprefix("data: ")) + for line in body.splitlines() + if line.startswith("data: ") + ] + text = "".join( + event["delta"]["text"] + for event in events + if event["type"] == "content_block_delta" + and event["delta"]["type"] == "text_delta" + ) + + assert text == truncated + assert events[-1]["type"] == "message_stop" diff --git a/vllm_mlx/server.py b/vllm_mlx/server.py index 7cf5fecee..72d8902ba 100644 --- a/vllm_mlx/server.py +++ b/vllm_mlx/server.py @@ -6102,7 +6102,10 @@ async def _stream_anthropic_messages( # Filter special tokens filtered = SPECIAL_TOKENS_PATTERN.sub("", delta_text) - if not filtered and not (use_reasoning and output_finished): + if not filtered and not ( + (use_reasoning and output_finished) + or (tool_parser and tool_markup_possible and output_finished) + ): continue if not use_reasoning: @@ -6180,8 +6183,10 @@ async def _stream_anthropic_messages( thinking_block_started = True yield f"event: content_block_delta\ndata: {json.dumps({'type': 'content_block_delta', 'index': thinking_index, 'delta': {'type': 'thinking_delta', 'thinking': delta_msg.reasoning}})}\n\n" - if delta_msg.content: - content_to_emit = delta_msg.content + content_to_emit = delta_msg.content or "" + if content_to_emit or ( + tool_parser and output_finished and tool_markup_possible + ): # Filter tool call markup during streaming if tool_parser and ( @@ -6510,6 +6515,7 @@ async def stream_chat_completion( content, reasoning = _promote_streaming_response_format_delta( content, reasoning, request ) + content = content or "" # Some models (e.g. MiniMax) wrap tool calls in # blocks, so reasoning parser captures tool call XML as