From 21f4c25c91e8ab28f726487a334a8128d207dad4 Mon Sep 17 00:00:00 2001 From: ahmed xijiaat <52128022+xijiaat@users.noreply.github.com> Date: Mon, 21 Sep 2026 14:34:22 +0800 Subject: [PATCH] [Bugfix] Honor generation controls in the DeepSeek V4.1 tokenizer Forward add_generation_prompt and continue_final_message to the independent V4.1 encoder. Preserve historical turn separators and EOS while honoring final assistant continuation, including partial reasoning. Add focused regression coverage to the existing tokenizer tests. Co-authored-by: OpenAI Codex Signed-off-by: ahmed xijiaat <52128022+xijiaat@users.noreply.github.com> --- tests/tokenizers_/test_deepseek_v41.py | 82 +++++++++++++++++++++++- vllm/tokenizers/deepseek_v41.py | 2 + vllm/tokenizers/deepseek_v41_encoding.py | 34 +++++++++- 3 files changed, 114 insertions(+), 4 deletions(-) diff --git a/tests/tokenizers_/test_deepseek_v41.py b/tests/tokenizers_/test_deepseek_v41.py index 13c0375fbb0a..643296e00535 100644 --- a/tests/tokenizers_/test_deepseek_v41.py +++ b/tests/tokenizers_/test_deepseek_v41.py @@ -29,7 +29,8 @@ def test_reference_encoder_fixtures(case_id): messages = data["messages"] if isinstance(data, dict) else data tools = data.get("tools") if isinstance(data, dict) else None expected = (FIXTURES / "deepseek_v41" / f"test_output_{case_id}.txt").read_text() - assert render(messages, tools=tools) == expected + assert render(messages, tools=tools, add_generation_prompt=False) == expected + assert render(messages, tools=tools) == expected + "<｜Assistant｜>" @pytest.mark.parametrize( @@ -185,3 +186,82 @@ def image_part(color): def test_unsupported_media_is_a_request_error(): with pytest.raises(ValueError, match="text and image content only"): render([{"role": "user", "content": [{"type": "input_audio"}]}]) + + +@pytest.mark.parametrize("thinking", [False, True]) +@pytest.mark.parametrize("final_role", ["user", "developer", "system", "assistant"]) +def test_generation_prompt_only_controls_final_turn(thinking, final_role): + messages = [ + {"role": "user", "content": "first question"}, + {"role": "assistant", "content": "first answer"}, + {"role": "user", "content": "second question"}, + {"role": final_role, "content": "final message"}, + ] + original = copy.deepcopy(messages) + closed = render(messages, thinking=thinking, add_generation_prompt=False) + opened = render(messages, thinking=thinking, add_generation_prompt=True) + cue = "<｜Assistant｜>" + ("" if thinking else "") + + assert opened == closed + cue + assert "first answer<｜end▁of▁sentence｜>" in closed + assert "<｜User｜>first question<｜Assistant｜>" in closed + assert messages == original + + +@pytest.mark.parametrize("thinking", [False, True]) +def test_system_only_generation_prompt(thinking): + messages = [{"role": "system", "content": "Be helpful."}] + closed = render(messages, thinking=thinking, add_generation_prompt=False) + cue = "<｜Assistant｜>" + ("" if thinking else "") + assert render(messages, thinking=thinking) == closed + cue + + +@pytest.mark.parametrize("thinking", [False, True]) +def test_continue_final_message_preserves_historical_eos(thinking): + messages = [ + {"role": "user", "content": "first question"}, + {"role": "assistant", "content": "first answer"}, + {"role": "user", "content": "second question"}, + {"role": "assistant", "content": "The answer is "}, + ] + original = copy.deepcopy(messages) + closed = render(messages, thinking=thinking, add_generation_prompt=False) + continued = render( + messages, + thinking=thinking, + add_generation_prompt=False, + continue_final_message=True, + ) + assert closed == continued + "<｜end▁of▁sentence｜>" + assert continued.endswith("The answer is ") + assert "first answer<｜end▁of▁sentence｜>" in continued + assert messages == original + + +def test_continue_final_message_through_chat_request(): + request = ChatCompletionRequest( + model="deepseek-ai/DeepSeek-V4.1-Flash", + messages=[ + {"role": "user", "content": "Finish this sentence."}, + {"role": "assistant", "content": "The answer is "}, + ], + continue_final_message=True, + ) + kwargs = request.build_chat_params(None, "auto").get_apply_chat_template_kwargs() + assert render(request.messages, **kwargs).endswith("The answer is ") + + +def test_continue_final_reasoning_keeps_thinking_open(): + messages = [ + {"role": "user", "content": "Explain your answer."}, + {"role": "assistant", "content": None, "reasoning": "Let me consider "}, + ] + assert render(messages, continue_final_message=True).endswith( + "<｜Assistant｜>Let me consider " + ) + + +@pytest.mark.parametrize("messages", [[], [{"role": "user", "content": "hello"}]]) +def test_continue_final_message_requires_assistant(messages): + with pytest.raises(ValueError, match="final assistant"): + render(messages, continue_final_message=True) diff --git a/vllm/tokenizers/deepseek_v41.py b/vllm/tokenizers/deepseek_v41.py index a1449e59e203..c6cce506076c 100644 --- a/vllm/tokenizers/deepseek_v41.py +++ b/vllm/tokenizers/deepseek_v41.py @@ -90,6 +90,8 @@ def apply_chat_template( thinking_mode="thinking" if thinking else "chat", drop_thinking=kwargs.get("drop_thinking", True), reasoning_effort=effort, + add_generation_prompt=kwargs.get("add_generation_prompt", True), + continue_final_message=kwargs.get("continue_final_message", False), ) if kwargs.get("tokenize", True): tokenizer_kwargs = { diff --git a/vllm/tokenizers/deepseek_v41_encoding.py b/vllm/tokenizers/deepseek_v41_encoding.py index 841de385cb61..65a25300903c 100644 --- a/vllm/tokenizers/deepseek_v41_encoding.py +++ b/vllm/tokenizers/deepseek_v41_encoding.py @@ -312,6 +312,8 @@ def render_message( thinking_mode: str, drop_thinking: bool = True, reasoning_effort: Union[str, int, None] = None, + add_generation_prompt: bool = True, + continue_final_message: bool = False, ) -> str: """ Render a single message at the given index into its V4.1 encoded string form. @@ -330,7 +332,8 @@ def render_message( 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) + is_last = index == len(messages) - 1 + wo_eos = msg.get("wo_eos", False) or (is_last and continue_final_message) if tools: tools = tools_from_openai_format(tools) @@ -446,6 +449,14 @@ def render_message( thinking_part = ( thinking_template.format(reasoning_content=rc) + thinking_end_token ) + if ( + is_last + and continue_final_message + and rc + and not summary_content + and not tc_content + ): + thinking_part = rc else: thinking_part = "" @@ -464,6 +475,10 @@ def render_message( else: raise NotImplementedError(f"Unknown role: {role}") + # Historical assistant separators are needed even without a generation prompt. + if is_last and not add_generation_prompt: + return prompt + # Append transition tokens based on what follows if index + 1 < len(messages) and messages[index + 1].get("role") not in [ "assistant", @@ -492,8 +507,10 @@ def render_message( ) prompt += task_sp_token - elif messages[index].get("role") in ["user", "developer"] or ( - messages[index].get("role") == "system" and index > 0 + elif ( + messages[index].get("role") in ["user", "developer"] + or (messages[index].get("role") == "system" and index > 0) + or (is_last and role in ["system", "assistant"]) ): # Normal generation: append Assistant + thinking token # (mid-conversation system messages also trigger the assistant header) @@ -537,8 +554,17 @@ def encode_messages( drop_thinking: bool = True, add_default_bos_token: bool = True, reasoning_effort: Union[str, int, None] = None, + add_generation_prompt: bool = True, + continue_final_message: bool = False, ) -> str: """Encode preprocessed (text-only) messages into the V4.1 prompt format.""" + if continue_final_message: + if not messages or messages[-1].get("role") != "assistant": + raise ValueError( + "continue_final_message requires a final assistant message" + ) + add_generation_prompt = False + context = context if context else [] # Preprocess: merge tool messages and sort tool results @@ -572,6 +598,8 @@ def encode_messages( thinking_mode=thinking_mode, drop_thinking=effective_drop_thinking, reasoning_effort=reasoning_effort, + add_generation_prompt=add_generation_prompt, + continue_final_message=continue_final_message, ) return prompt