[None][fix] two silent DeepSeek-V4 defects: prompt handoff and withheld stream text

tensorrt_llm/serve/tool_parser/deepseekv32_parser.py

@@ -308,6 +308,37 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami
             logger.error(f"Error in parse_streaming_increment: {e}")
             return StreamingParseResult(normal_text=normal_text + current_text)
 
+    def finish(self, tools: List[Tool]) -> StreamingParseResult:
+        """Release text the stream ended on while it was still withheld.
+
+        ``parse_streaming_increment`` holds the buffer back whenever its tail
+        could still grow into a DSML delimiter, and releases it as soon as the
+        next chunk resolves the ambiguity. When generation stops instead --
+        max_tokens, a stop string, an abort -- nothing resolves it, and the
+        base class's no-op ``finish`` drops whatever was held. A response
+        ending on ``<`` loses every character buffered since the last release,
+        silently: the request succeeds and the tail of the answer is missing.
+
+        Only a buffer with no tool-call section is released. Text that precedes
+        a section is already streamed by the increment path, so what remains
+        there is partial DSML markup rather than content, and surfacing that
+        raw would trade one defect for another.
+        """
+        buffer = self._buffer
+        if not buffer:
+            return StreamingParseResult()
+
+        start_tokens = [self.bot_token, self._INVOKE_HEADER_PREFIX]
+        if any(idx != -1 for idx in map(buffer.find, start_tokens)):
+            return StreamingParseResult()
+
+        self._buffer = ""
+        # The same delimiters the increment path strips before emitting, so a
+        # stream that ended just after one does not show it to the caller.
+        for token in (self.eot_token, self.invoke_end_token, self._eos_token):
+            buffer = buffer.replace(token, "")
+        return StreamingParseResult(normal_text=buffer)
+
     def structure_info(self) -> _GetInfoFunc:
         return lambda name: StructureInfo(
             begin=f'<｜DSML｜invoke name="{name}">',

tensorrt_llm/tokenizer/deepseek_v4/tokenizer.py

@@ -275,6 +275,20 @@ def _render_user_content(message: dict[str, Any]) -> str:
     return "\n\n".join(parts)
 
 
+def _only_reminders_follow(index: int, messages: list[dict[str, Any]]) -> bool:
+    """True when every message after `index` is a ``latest_reminder``.
+
+    A reminder carries its own token but never emits a turn boundary, so when
+    reminders are all that remains there is no assistant message for the
+    preceding user turn to hand the floor to. Deferring the handoff past them
+    keeps the prompt ending on it; emitting it early leaves the prompt ending
+    inside the reminder text with the assistant turn already open, and the
+    model continues the document instead of answering.
+    """
+    rest = messages[index + 1 :]
+    return bool(rest) and all(m.get("role") == "latest_reminder" for m in rest)
+
+
 def _render_message(
     index: int,
     messages: list[dict[str, Any]],
@@ -367,16 +381,55 @@ def _render_message(
             prompt += ASSISTANT_TOKEN
             prompt += THINKING_START_TOKEN if thinking_mode == "thinking" else THINKING_END_TOKEN
         prompt += VALID_TASKS[task]
-    elif role in ("user", "developer") and (next_role == "assistant" or add_generation_prompt):
+    elif role in ("user", "developer") and (
+        next_role == "assistant"
+        or (add_generation_prompt and not _only_reminders_follow(index, messages))
+    ):
         prompt += ASSISTANT_TOKEN
         if thinking_mode == "thinking" and (not drop_thinking or index >= last_user_index):
             prompt += THINKING_START_TOKEN
         else:
             prompt += THINKING_END_TOKEN
+    elif role == "latest_reminder" and next_role is None and add_generation_prompt:
+        # Deferred from the user turn above so the prompt ends on the boundary
+        # rather than inside the reminder text.
+        prompt += ASSISTANT_TOKEN
+        prompt += THINKING_START_TOKEN if thinking_mode == "thinking" else THINKING_END_TOKEN
 
     return prompt
 
 
+def _map_trailing_system_to_reminder(
+    messages: list[dict[str, Any]],
+) -> list[dict[str, Any]]:
+    """Re-role non-leading ``system`` messages as ``latest_reminder``.
+
+    DeepSeek-V4 encodes the system prompt *positionally*: the opening bare-text
+    span before the first ``<｜User｜>`` is the system slot, so ``_render_message``
+    emits system content with no role token at all. That is correct at the front
+    and degenerate anywhere else -- a mid-conversation system message renders as
+    unmarked text floating between two turns.
+
+    The format already has the right slot for out-of-band mid-conversation
+    content: ``latest_reminder``, which carries its own token. Clients that append
+    transient system messages (task reminders, background-task notifications) map
+    onto exactly that, so route them there instead of emitting bare text.
+    """
+    mapped: list[dict[str, Any]] = []
+    in_leading_system_run = True
+    for message in messages:
+        role = message.get("role")
+        if role == "system" and not in_leading_system_run:
+            reminder = dict(message)
+            reminder["role"] = "latest_reminder"
+            mapped.append(reminder)
+            continue
+        if role != "system":
+            in_leading_system_run = False
+        mapped.append(message)
+    return mapped
+
+
 def _encode_messages(
     messages: list[dict[str, Any]],
     thinking_mode: str,
@@ -386,6 +439,7 @@ def _encode_messages(
 ) -> str:
     messages = _merge_tool_messages(messages)
     messages = _sort_tool_results_by_call_order(messages)
+    messages = _map_trailing_system_to_reminder(messages)
 
     effective_drop_thinking = drop_thinking
     if any(message.get("tools") for message in messages):
@@ -419,12 +473,6 @@ def from_pretrained(
         revision: str | None = None,
         **kwargs,
     ) -> "DeepseekV4Tokenizer":
-        # AutoTokenizer resolves the checkpoint config first, and the
-        # deepseek_v4 model_type is invisible to stock transformers; the
-        # registration is an import side effect of tensorrt_llm._torch.configs
-        # (previously guaranteed by the eager `import tensorrt_llm`).
-        import tensorrt_llm._torch.configs  # noqa: F401
-
         tokenizer = AutoTokenizer.from_pretrained(
             path_or_repo_id,
             *args,

tests/unittest/llmapi/apps/test_tool_parsers.py

@@ -1897,6 +1897,66 @@ def test_deepseek_streaming_preserves_withheld_text(
                                          sample_tools).normal_text == expected
 
 
+@pytest.mark.parametrize("parser_cls", [DeepSeekV32Parser, DeepSeekV4Parser])
+@pytest.mark.parametrize(
+    "deltas",
+    [
+        # Generation stops on a bare `<`, which is a prefix of every DSML
+        # delimiter, so the increment path is still withholding when the
+        # stream ends.
+        ["The condition is a <"],
+        # Several characters into a delimiter rather than one.
+        ["cost is 3", "<｜D"],
+        # The withheld run spans more than the last delta.
+        ["Here is the analysis. ", "The threshold is <"],
+    ],
+)
+def test_deepseek_streaming_emits_withheld_text_when_the_stream_ends(
+        sample_tools: list[ChatCompletionToolsParam],
+        parser_cls: type[BaseToolParser], deltas: list[str]) -> None:
+    """A stream that stops mid-ambiguity must still deliver what was held.
+
+    Withholding is released by the next chunk, so a stream that ends instead --
+    max_tokens, a stop string, an abort -- has nothing to release it. Without a
+    `finish` that flushes, the tail is dropped silently: the request succeeds
+    and the end of the answer is simply missing.
+    """
+    parser = parser_cls()
+
+    streamed = "".join(
+        parser.parse_streaming_increment(delta, sample_tools).normal_text
+        for delta in deltas)
+    streamed += parser.finish(sample_tools).normal_text
+
+    assert streamed == "".join(deltas)
+
+
+@pytest.mark.parametrize("parser_cls", [DeepSeekV32Parser, DeepSeekV4Parser])
+def test_deepseek_finish_leaves_a_truncated_tool_call_alone(
+        sample_tools: list[ChatCompletionToolsParam],
+        parser_cls: type[BaseToolParser]) -> None:
+    """Flushing must not turn half a tool call into visible text.
+
+    Content before a tool-call section is already streamed by the increment
+    path, so a buffer holding a section holds markup rather than content.
+    Emitting it would trade a dropped tail for DSML leaking into the answer.
+    """
+    parser = parser_cls()
+
+    # Built from the parser's own tokens: V3.2 opens a section with
+    # `function_calls` and V4 with `tool_calls`, so a hard-coded one is not a
+    # section start for the other parser and the buffer would be flushed as
+    # ordinary text.
+    truncated = f'Let me read it. {parser.bot_token}<｜DSML｜invoke name="Re'
+    streamed = parser.parse_streaming_increment(truncated,
+                                                sample_tools).normal_text
+    finished = parser.finish(sample_tools)
+
+    assert streamed == "Let me read it. "
+    assert finished.normal_text == ""
+    assert finished.calls == []
+
+
 @pytest.mark.parametrize(
     "parser_cls, tool_call_text",
     [

tests/unittest/llmapi/test_deepseek_v4_tokenizer.py

@@ -284,6 +284,82 @@ def test_deepseek_v4_chat_template_renders_developer_tools_and_latest_reminder()
     assert "<｜User｜><tool_result>[0]</tool_result><｜Assistant｜><think>" in prompt
 
 
+def test_deepseek_v4_chat_template_ends_on_handoff_after_trailing_reminder():
+    """A trailing reminder must not swallow the generation prompt.
+
+    Clients append transient system messages (task nags, background-task
+    notifications) after the last user turn; those are re-roled to
+    `latest_reminder`. The reminder itself never emits a turn boundary, so
+    without deferring the handoff the prompt ends inside the reminder text with
+    `<think>` already open, and the model continues the document rather than
+    answering.
+    """
+    tokenizer = DeepseekV4Tokenizer(_DummyTokenizer())
+
+    prompt = tokenizer.apply_chat_template(
+        [
+            {"role": "system", "content": "sys"},
+            {"role": "user", "content": "hi"},
+            {"role": "system", "content": "reminder"},
+        ],
+        tokenize=False,
+        enable_thinking=True,
+    )
+
+    assert prompt == (
+        "<｜begin▁of▁sentence｜>sys<｜User｜>hi<｜latest_reminder｜>reminder<｜Assistant｜><think>"
+    )
+    assert prompt.count("<｜Assistant｜>") == 1
+
+
+def test_deepseek_v4_chat_template_defers_handoff_past_every_trailing_reminder():
+    tokenizer = DeepseekV4Tokenizer(_DummyTokenizer())
+
+    prompt = tokenizer.apply_chat_template(
+        [
+            {"role": "system", "content": "sys"},
+            {"role": "user", "content": "hi"},
+            {"role": "system", "content": "one"},
+            {"role": "system", "content": "two"},
+        ],
+        tokenize=False,
+        enable_thinking=True,
+    )
+
+    assert prompt.endswith("<｜Assistant｜><think>")
+    assert prompt.count("<｜Assistant｜>") == 1
+    assert "<｜latest_reminder｜>one<｜latest_reminder｜>two" in prompt
+
+
+def test_deepseek_v4_chat_template_keeps_handoff_before_mid_conversation_reminder():
+    """A reminder with a turn after it keeps the existing layout.
+
+    Here the handoff belongs to the assistant message that follows, and the
+    reminder is injected inside that assistant's thinking block.
+    """
+    tokenizer = DeepseekV4Tokenizer(_DummyTokenizer())
+
+    prompt = tokenizer.apply_chat_template(
+        [
+            {"role": "system", "content": "sys"},
+            {"role": "user", "content": "hi"},
+            {"role": "system", "content": "reminder"},
+            {"role": "assistant", "content": "answer", "reasoning": "why"},
+            {"role": "user", "content": "again"},
+        ],
+        tokenize=False,
+        enable_thinking=True,
+    )
+
+    # apply_chat_template drops historical reasoning by default, which also
+    # closes the completed turn with </think> rather than opening one.
+    assert (
+        "<｜User｜>hi<｜Assistant｜></think><｜latest_reminder｜>reminder"
+        "answer<｜end▁of▁sentence｜>"
+    ) in prompt
+    assert prompt.endswith("<｜User｜>again<｜Assistant｜><think>")
+
+
 def test_deepseek_v4_chat_template_renders_action_task_token():
     tokenizer = DeepseekV4Tokenizer(_DummyTokenizer())