[Bugfix] Honor generation controls in the DeepSeek V4.1 tokenizer

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｜><think>"
 
 
 @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｜>" + ("<think>" if thinking else "</think>")
+
+    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｜>" + ("<think>" if thinking else "</think>")
+    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｜><think>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)

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 = {

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