diff --git a/tensorrt_llm/serve/tool_parser/deepseekv32_parser.py b/tensorrt_llm/serve/tool_parser/deepseekv32_parser.py index c06c639b92aa..bc2f13a87823 100644 --- a/tensorrt_llm/serve/tool_parser/deepseekv32_parser.py +++ b/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}">', diff --git a/tensorrt_llm/tokenizer/deepseek_v4/tokenizer.py b/tensorrt_llm/tokenizer/deepseek_v4/tokenizer.py index fba84a22ef9c..61c08cb173f4 100644 --- a/tensorrt_llm/tokenizer/deepseek_v4/tokenizer.py +++ b/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, diff --git a/tests/unittest/llmapi/apps/test_tool_parsers.py b/tests/unittest/llmapi/apps/test_tool_parsers.py index 046e85294a89..cd14ac198115 100644 --- a/tests/unittest/llmapi/apps/test_tool_parsers.py +++ b/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", [ diff --git a/tests/unittest/llmapi/test_deepseek_v4_tokenizer.py b/tests/unittest/llmapi/test_deepseek_v4_tokenizer.py index dec508d49cbc..6ca934523832 100644 --- a/tests/unittest/llmapi/test_deepseek_v4_tokenizer.py +++ b/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｜>[0]<｜Assistant｜>" 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 + `` 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｜>" + ) + 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｜>") + 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 rather than opening one. + assert ( + "<｜User｜>hi<｜Assistant｜><｜latest_reminder｜>reminder" + "answer<｜end▁of▁sentence｜>" + ) in prompt + assert prompt.endswith("<｜User｜>again<｜Assistant｜>") + + def test_deepseek_v4_chat_template_renders_action_task_token(): tokenizer = DeepseekV4Tokenizer(_DummyTokenizer())