diff --git a/docs/source/release-notes.md b/docs/source/release-notes.md index 39dd4fb80f68..40c0edecc508 100644 --- a/docs/source/release-notes.md +++ b/docs/source/release-notes.md @@ -19,6 +19,8 @@ All published functionality in the Release Notes has been fully tested and verif ### API Changes +- Eligible text-only Kimi-K3 chat requests use the renderer’s native token IDs; `/v1/messages/count_tokens` excludes the same prompt tokens as response usage (three fewer for the default Kimi-K3 template). Requests containing media, server or tokenizer templates, forced tool-call prefixes, added special tokens, or prompt truncation retain the text path. + - `/v1/responses`: a response cut short by `max_output_tokens` now reports `incomplete_details.reason` `max_output_tokens`, and when streamed it ends with `response.incomplete` instead of `response.completed`, as in the OpenAI API. - **[BREAKING CHANGE] Removed obsolete C++ executor orchestration helpers.** The internal `RequestWithId` type, its `serializeReqWithIds` / `deserializeReqWithIds` functions, all production `insertRequestInOrder` overloads, and the unused `requestUtils.h` and `orchestratorUtils.h` helpers have been removed. These helpers belonged to the retired C++ executor orchestration path; downstream code linking their exported symbols must stop using them and rebuild. The current `Request` / `Response` interfaces and serialization remain available. Executor factory and worker constructors no longer accept the unused `engine` and `batched_logits_processor` arguments; callers should omit them and continue passing model configuration through `llm_args` and `hf_model_dir`. Unused decoder, request-cancellation, executor, legacy-model, orchestration, and speculative-decoding MPI tag names have also been removed; values 1–23, 127–129, and 1023–1025 are retired. The default, KV-cache event, and context-transfer tags retain their existing values. Internal `LogprobParams` and `LogProbsResult` tuples now contain only prompt-logprob fields; generation logprobs remain supplied by the sampler. The unused `TLLM_EXECUTOR_PERIODICAL_RESP_IN_AWAIT` environment variable has been removed. - **[BREAKING CHANGE] Unified multimodal processor metadata.** diff --git a/tensorrt_llm/inputs/utils.py b/tensorrt_llm/inputs/utils.py index 1360861ed743..a0144e9b2ac8 100644 --- a/tensorrt_llm/inputs/utils.py +++ b/tensorrt_llm/inputs/utils.py @@ -4,6 +4,7 @@ import asyncio import base64 from collections import defaultdict +from functools import partial from io import BytesIO from pathlib import Path from typing import (Any, Collection, Coroutine, Dict, List, Optional, Tuple, @@ -625,6 +626,24 @@ def _has_python_chat_template(tokenizer: PreTrainedTokenizerBase) -> bool: is not PreTrainedTokenizerBase.apply_chat_template) +def uses_tokenizer_native_renderer(tokenizer: Union[TransformersTokenizer, + TokenizerBase], + processor: ProcessorMixin, + chat_template: Optional[str]) -> bool: + """Whether `apply_chat_template` renders with the tokenizer's own renderer. + + It does when no Jinja template is declared anywhere (the call, the + processor or the tokenizer) and the tokenizer overrides + `apply_chat_template` in Python, as the kimi_k3 checkpoint tokenizer does. + """ + if isinstance(tokenizer, TransformersTokenizer): + tokenizer = tokenizer.tokenizer + return (chat_template is None + and getattr(processor, "chat_template", None) is None + and getattr(tokenizer, "chat_template", None) is None + and _has_python_chat_template(tokenizer)) + + def resolve_hf_chat_template( tokenizer: TokenizerBase, processor: ProcessorMixin, @@ -785,10 +804,7 @@ def apply_chat_template( if isinstance(tokenizer, TransformersTokenizer): tokenizer = tokenizer.tokenizer # we need the TokenizerBase for apply_chat_template - if (chat_template is None - and getattr(processor, "chat_template", None) is None - and getattr(tokenizer, "chat_template", None) is None - and _has_python_chat_template(tokenizer)): + if uses_tokenizer_native_renderer(tokenizer, processor, chat_template): native_kwargs = dict(chat_template_kwargs or {}) if documents is not None: native_kwargs["documents"] = documents @@ -822,7 +838,17 @@ def apply_chat_template( mm_placeholder_count) # STRING path: placeholders already inserted in content by caller - result = tokenizer.apply_chat_template( + template_renderer = tokenizer.apply_chat_template + if (model_type in ("kimi_k3", "kimi_linear") + and isinstance(tokenizer, PreTrainedTokenizerBase) + and callable(getattr(tokenizer, "_encode_chat_segments", None))): + # Kimi-K3's segment renderer (also used by its kimi_linear text model) + # ignores Jinja templates. Other overrides, e.g. Kimi-K2.5's, delegate + # to Jinja after preprocessing tools and must remain in the call path. + template_renderer = partial(PreTrainedTokenizerBase.apply_chat_template, + tokenizer) + + result = template_renderer( conversation=conversation, tokenize=enable_tokenize, return_dict=False, diff --git a/tensorrt_llm/serve/extensions/kimi_k3.py b/tensorrt_llm/serve/extensions/kimi_k3.py index efe6dd84a35c..fe50f0f54ef4 100644 --- a/tensorrt_llm/serve/extensions/kimi_k3.py +++ b/tensorrt_llm/serve/extensions/kimi_k3.py @@ -141,7 +141,7 @@ def validate_kimi_dynamic_tools(request: ChatCompletionRequest) -> None: class KimiK3ServingExtension(ServingExtension): """Kimi/Moonshot API semantics for kimi_k3 chat requests.""" - render_version = 2 + render_version = 3 # Trailing tokens of the native K3 generation prompt (the channel opener # ``<|open|>think|response<|sep|>``) that are not reported as prompt usage. diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 7e7afa6256fc..a9daa1d7156f 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -127,8 +127,9 @@ from tensorrt_llm.serve.render import (HarmonyConversionError, RenderResources, UnsupportedRenderError, mount_render_endpoints, - prepare_chat_request, render_chat, - render_chat_for_length, + prepare_chat_request, + prompt_tokens_excluded_from_usage, + render_chat, render_chat_for_length, render_conversation) from tensorrt_llm.serve.responses_utils import (ClosingStreamingResponse, ConversationHistoryStore, @@ -2300,14 +2301,9 @@ async def chat_stream_generator( self._apply_spec_decode_stats_opt_in(postproc_args) # Some models' prompt-token accounting excludes trailing rendered # tokens (kimi_k3: the generation channel opener); the model still - # sees the full prompt. + # sees the full prompt. count_tokens leaves out the same tokens. postproc_args.num_prompt_tokens_offset = ( - render_resources.prompt_tokens_excluded_from_usage( - add_generation_prompt=request.add_generation_prompt, - pre_tokenized=request.prompt_token_ids is not None - or request.prompt_token_ids_b64 is not None, - request_chat_template=request.chat_template, - )) + prompt_tokens_excluded_from_usage(request, render_resources)) if dynamic_tool_params: # The tool parser must see dynamic tools to recognize their # calls in the model output. @@ -2368,7 +2364,12 @@ async def chat_stream_generator( else: # Render (and append the forced named-tool prefix, see # ``_build_forced_tool_call_decoding``) while the media loads. - # Tokenization stays with the engine's input processor. + # Tokenization stays with the engine's input processor, except + # that a tokenizer-native renderer's own ids are taken: it + # keeps control-token strings in message content as text, and + # encoding its rendered string would not. Only for text-only + # prompts, since the input processor expands media placeholders + # in the text. render_task = render_conversation( render_resources, conversation=conversation, @@ -2381,6 +2382,10 @@ async def chat_stream_generator( injected_chat_template_kwargs=request. injected_chat_template_kwargs, forced_prefix=forced_tool_begin_prefix, + prefer_native_ids=not any( + message.get("media") for message in conversation), + add_special_tokens=request.add_special_tokens, + truncate_prompt_tokens=request.truncate_prompt_tokens, ) rendered, (mm_data, mm_embeddings) = await asyncio.gather( render_task, mm_coroutines) @@ -2641,7 +2646,11 @@ def _count() -> int: # Media, an own-processor model or an unconvertible Harmony # request: count the text, tools and template without the media. rendered = render_chat_for_length(chat_request, resources) - return len(rendered.token_ids) + # Reported like the chat route's prompt usage, which leaves out + # the rendered tokens the model extension excludes from it. + excluded = prompt_tokens_excluded_from_usage( + chat_request, resources) + return max(0, len(rendered.token_ids) - excluded) input_tokens = await asyncio.to_thread(_count) except Exception: diff --git a/tensorrt_llm/serve/render/__init__.py b/tensorrt_llm/serve/render/__init__.py index 9729715f60b8..8d7f44ba549e 100644 --- a/tensorrt_llm/serve/render/__init__.py +++ b/tensorrt_llm/serve/render/__init__.py @@ -25,6 +25,7 @@ PreparedChat, arender_chat, prepare_chat_request, + prompt_tokens_excluded_from_usage, render_chat, render_chat_for_length, render_conversation, @@ -70,6 +71,7 @@ "harmony_enabled_for", "mount_render_endpoints", "prepare_chat_request", + "prompt_tokens_excluded_from_usage", "render_chat", "render_chat_for_length", "render_conversation", diff --git a/tensorrt_llm/serve/render/chat.py b/tensorrt_llm/serve/render/chat.py index b53e5b4bbacd..8ad826b9820f 100644 --- a/tensorrt_llm/serve/render/chat.py +++ b/tensorrt_llm/serve/render/chat.py @@ -21,7 +21,8 @@ 1. :func:`prepare_chat_request`: model-extension preprocessing and tool dicts; 2. template rendering with the server template, processor and extension rules; 3. the forced named-tool prefix, when the caller supplies one; -4. optional tokenization through the engine's own text tokenizer. +4. optional tokenization: the tokenizer-native renderer's own ids when they are the + prompt, otherwise the engine's own text tokenizer. :func:`render_conversation` is steps 2-4 for a caller that already parsed the conversation (the chat route overlaps media loading with rendering). @@ -39,9 +40,9 @@ from tensorrt_llm.inputs.tokenization import tokenize_prompt from tensorrt_llm.inputs.utils import ( - _has_python_chat_template, apply_chat_template, async_apply_chat_template, + uses_tokenizer_native_renderer, ) from tensorrt_llm.sampling_params import SamplingParams from tensorrt_llm.serve.chat_utils import parse_chat_messages_coroutines @@ -121,6 +122,32 @@ def _template_kwargs( ) +def _prompt_is_native_ids( + res: RenderResources, + *, + chat_template: Optional[str], + forced_prefix: Optional[str], + add_special_tokens: bool, + truncate_prompt_tokens: Optional[int], +) -> bool: + """Whether the prompt is the tokenizer-native renderer's own token ids. + + Such a renderer (see ``uses_tokenizer_native_renderer``) tokenizes segment by segment: its + structural markers become control tokens and message content stays ordinary text, so a + control-token string inside user or tool text stays text. Encoding its rendered string + instead turns that string into a control token, or into an id past the vocabulary. Its ids + are the prompt only when nothing is appended to them, added around them or cut from them. + """ + return ( + forced_prefix is None + and not add_special_tokens + and truncate_prompt_tokens is None + and uses_tokenizer_native_renderer( + res.tokenizer, res.processor, chat_template or res.default_chat_template + ) + ) + + def tokenize_rendered( res: RenderResources, text: str, @@ -210,19 +237,24 @@ async def render_conversation( injected_chat_template_kwargs: Any = None, forced_prefix: Optional[str] = None, tokenize: bool = False, + prefer_native_ids: bool = False, add_special_tokens: bool = False, truncate_prompt_tokens: Optional[int] = None, encode_rendered: Optional[Callable[[str, Any], List[int]]] = None, ) -> RenderedPrompt: - """Render an already-parsed conversation without blocking the event loop.""" - # A tokenizer-native template tokenizes segment by segment, so control-token strings inside - # user or tool text stay text. Encoding the finished string would turn them into special - # tokens, so keep the tokenizer's own ids when nothing is added after the template. - native_ids = ( - tokenize - and forced_prefix is None - and truncate_prompt_tokens is None - and _has_python_chat_template(getattr(res.tokenizer, "tokenizer", res.tokenizer)) + """Render an already-parsed conversation without blocking the event loop. + + ``tokenize`` returns token ids. ``prefer_native_ids`` is for a caller whose engine tokenizes + the text itself (with ``add_special_tokens`` and ``truncate_prompt_tokens``): it returns the + tokenizer-native renderer's own ids when they are the prompt (see + :func:`_prompt_is_native_ids`), and the text otherwise. + """ + native_ids = (tokenize or prefer_native_ids) and _prompt_is_native_ids( + res, + chat_template=chat_template, + forced_prefix=forced_prefix, + add_special_tokens=add_special_tokens, + truncate_prompt_tokens=truncate_prompt_tokens, ) rendered = await async_apply_chat_template( **_template_kwargs( @@ -397,6 +429,15 @@ def render_chat( # Nothing here loads media; close the coroutine instead of leaving it to be reported # as never awaited. mm_coroutine.close() + # The chat route executes the tokenizer-native renderer's own ids when they are the prompt + # (see ``render_conversation``), so these are the ids to count and route on as well. + native_ids = tokenize and _prompt_is_native_ids( + res, + chat_template=request.chat_template, + forced_prefix=None, + add_special_tokens=request.add_special_tokens, + truncate_prompt_tokens=request.truncate_prompt_tokens, + ) rendered = apply_chat_template( **_template_kwargs( res, @@ -408,7 +449,7 @@ def render_chat( chat_template=request.chat_template, chat_template_kwargs=request.chat_template_kwargs, injected_chat_template_kwargs=request.injected_chat_template_kwargs, - enable_tokenize=False, + enable_tokenize=native_ids, ) ) result = _finish( @@ -483,6 +524,24 @@ def render_chat_for_length( return render_chat(text_only, plain, encode_rendered=encode_rendered) +def prompt_tokens_excluded_from_usage( + request: "ChatCompletionRequest", res: RenderResources +) -> int: + """Rendered prompt tokens the chat route leaves out of the reported prompt usage. + + ``RenderResources.prompt_tokens_excluded_from_usage`` asked with how ``request`` is + rendered (kimi_k3: the trailing generation channel opener); the model still sees those + tokens. The chat route, its preparation context and count_tokens all ask here, so they + leave out the same tokens. + """ + return res.prompt_tokens_excluded_from_usage( + add_generation_prompt=request.add_generation_prompt, + pre_tokenized=request.prompt_token_ids is not None + or request.prompt_token_ids_b64 is not None, + request_chat_template=request.chat_template, + ) + + def preparation_context( request: "ChatCompletionRequest", res: RenderResources, rendered: RenderedPrompt ) -> Dict[str, Any]: @@ -493,12 +552,7 @@ def preparation_context( with the ids (see ``PreparedContext``). ``request`` is the request as rendered, so it does not yet carry ``prompt_token_ids``. """ - excluded = res.prompt_tokens_excluded_from_usage( - add_generation_prompt=request.add_generation_prompt, - pre_tokenized=request.prompt_token_ids is not None - or request.prompt_token_ids_b64 is not None, - request_chat_template=request.chat_template, - ) + excluded = prompt_tokens_excluded_from_usage(request, res) thinking: Dict[str, bool] = {} if rendered.template_text and request.add_generation_prompt: # Deferred: the reasoning parsers module is only needed here. diff --git a/tests/unittest/inputs/test_chat_template_dispatch.py b/tests/unittest/inputs/test_chat_template_dispatch.py index 33cd90f9d17b..fdb156caf598 100644 --- a/tests/unittest/inputs/test_chat_template_dispatch.py +++ b/tests/unittest/inputs/test_chat_template_dispatch.py @@ -365,6 +365,34 @@ def apply_chat_template(self, **_): class TestPythonChatTemplate: + @pytest.mark.parametrize("model_type", ["kimi_k25", "kimi_linear"]) + @pytest.mark.parametrize("template_source", ["tokenizer", "request"]) + def test_jinja_override_preserves_tool_preprocessing( + self, template_source: str, model_type: str + ) -> None: + class ToolPreprocessingTokenizer(PreTrainedTokenizerBase): + def apply_chat_template(self, conversation, tools=None, **kwargs): + # Kimi-K2.5/K2.6 preprocess tools before delegating to the Jinja renderer. + kwargs["tools_ts_str"] = f"declare function {tools[0]['function']['name']}();" + return super().apply_chat_template(conversation, tools=tools, **kwargs) + + template = "SELECTED:{{ tools_ts_str }}|{{ messages[0].content }}" + tokenizer = ToolPreprocessingTokenizer( + chat_template=template if template_source == "tokenizer" else "DEFAULT" + ) + result = apply_chat_template( + model_type=model_type, + tokenizer=TransformersTokenizer(tokenizer), + processor=None, + conversation=[ConversationMessage(role="user", content="hello", media=[])], + add_generation_prompt=True, + mm_placeholder_counts=[{}], + tools=[{"type": "function", "function": {"name": "echo"}}], + chat_template=template if template_source == "request" else None, + ) + + assert result == "SELECTED:declare function echo();|hello" + @pytest.mark.parametrize( ("enable_tokenize", "expected"), [(False, "native-rendered"), (True, [1, 2, 3])], @@ -441,6 +469,44 @@ def apply_chat_template(self, conversation, **kwargs): _, captured_kwargs = inner_tokenizer.captured assert captured_kwargs["chat_template"] == template + def test_the_native_renderer_is_used_only_without_any_jinja_template(self): + from tensorrt_llm.inputs.utils import uses_tokenizer_native_renderer + + class NativeTokenizer: + chat_template = None + + def apply_chat_template(self, conversation, **kwargs): + return "native-rendered" + + class OverrideWithJinjaTokenizer(NativeTokenizer): + # A Python override next to a Jinja template renders the Jinja template. + chat_template = "{{ messages }}" + + class DefaultTokenizer: + chat_template = None + apply_chat_template = PreTrainedTokenizerBase.apply_chat_template + + class JinjaTokenizer(DefaultTokenizer): + chat_template = "{{ messages }}" + + class Processor: + chat_template = None + + class JinjaProcessor: + chat_template = "{{ messages }}" + + native = NativeTokenizer() + assert uses_tokenizer_native_renderer(native, None, None) + assert uses_tokenizer_native_renderer(TransformersTokenizer(native), Processor(), None) + assert not uses_tokenizer_native_renderer(native, None, "{{ messages }}") + assert not uses_tokenizer_native_renderer(native, JinjaProcessor(), None) + assert not uses_tokenizer_native_renderer(OverrideWithJinjaTokenizer(), None, None) + assert not uses_tokenizer_native_renderer(DefaultTokenizer(), None, None) + # The wrapper's own apply_chat_template only forwards to the tokenizer it wraps. + assert not uses_tokenizer_native_renderer( + TransformersTokenizer(JinjaTokenizer()), None, None + ) + def test_missing_template_still_errors_without_native_override(self): class DefaultTokenizer: chat_template = None diff --git a/tests/unittest/llmapi/apps/test_openai_chat_multimodal_prompt.py b/tests/unittest/llmapi/apps/test_openai_chat_multimodal_prompt.py index a2ab67b6677c..4d86199a41fb 100644 --- a/tests/unittest/llmapi/apps/test_openai_chat_multimodal_prompt.py +++ b/tests/unittest/llmapi/apps/test_openai_chat_multimodal_prompt.py @@ -16,7 +16,8 @@ A prompt the server rendered through the chat template already carries one placeholder per media item, so it goes to the input processor marked -``chat_template_applied``; a pre-tokenized request does not. +``chat_template_applied``; a pre-tokenized request does not. A text-only prompt +that a tokenizer-native renderer builds goes as that renderer's own ids. """ import base64 @@ -146,3 +147,182 @@ def generate_async(*, inputs, **kwargs): (item,) = prompt["mm_item_order"] assert (item["modality"], item["index"]) == ("image", 0) assert item["placeholder"] in _StubTokenizer.RENDERED + + +_TEST_VLM_MODEL_TYPE = "test_native_renderer_vlm" + + +class _TestVlmConfig: + model_type = _TEST_VLM_MODEL_TYPE + vocab_size = 1024 + + +@pytest.fixture +def test_vlm_placeholders(): + from tensorrt_llm.inputs.registry import ( + MULTIMODAL_PLACEHOLDER_REGISTRY, + MultimodalPlaceholderMetadata, + MultimodalPlaceholderPlacement, + ) + + MULTIMODAL_PLACEHOLDER_REGISTRY.set_placeholder_metadata( + _TEST_VLM_MODEL_TYPE, + MultimodalPlaceholderMetadata( + placeholder_map={"image": ""}, + placeholder_placement=MultimodalPlaceholderPlacement.BEFORE_TEXT, + placeholders_separator="", + ), + ) + yield + MULTIMODAL_PLACEHOLDER_REGISTRY.remove_placeholder_metadata(_TEST_VLM_MODEL_TYPE) + + +class _RecordingTokenizer: + """Renders in Python, records whether it was asked for ids, and gives its own ids if so.""" + + NATIVE_IDS = [5, 6, 7] + + def __init__(self, chat_template=None): + self.tokenizer = self + self.chat_template = chat_template + self.vocab_size = _TestVlmConfig.vocab_size + self.name_or_path = None + self.tokenize_requests = [] + + def get_chat_template(self, chat_template=None, tools=None): + return chat_template or self.chat_template + + def apply_chat_template(self, conversation, **kwargs): + self.tokenize_requests.append(kwargs["tokenize"]) + if kwargs["tokenize"]: + return list(self.NATIVE_IDS) + return "".join(str(message["content"]) for message in conversation) + + def encode(self, text, **kwargs): + return [1, 2, 3] + + +def _engine_inputs( + tmp_path, tokenizer, content, server_chat_template=None, tool_parser=None, **fields +): + """What ``openai_chat`` hands to the engine for one user message.""" + captured = {} + + def generate_async(*, inputs, **kwargs): + captured["inputs"] = inputs + return SimpleNamespace(request_id=1, prompt_token_ids=[7, 11], finished=True) + + server = make_stub_openai_server( + tmp_path, "test-vlm", _TestVlmConfig(), generate_async, tokenizer=tokenizer + ) + server.chat_template = server_chat_template + server.tool_parser = tool_parser + if tool_parser is not None: + # A named tool_choice is decoded under a JSON schema. + server.generator.args.guided_decoding_backend = "xgrammar" + server.await_disconnected = AsyncMock() + server._create_chat_response = AsyncMock( + return_value=SimpleNamespace(model_dump=lambda: {"ok": True}) + ) + app = FastAPI() + app.add_api_route("/v1/chat/completions", server.openai_chat, methods=["POST"]) + body = { + "model": "test-vlm", + "messages": [{"role": "user", "content": content}], + "max_tokens": 4, + **fields, + } + response = TestClient(app).post("/v1/chat/completions", json=body) + assert response.status_code == 200, response.text + return captured["inputs"] + + +def test_a_text_only_prompt_reaches_the_engine_as_the_native_renderers_ids( + tmp_path, test_vlm_placeholders +): + tokenizer = _RecordingTokenizer() + + prompt = _engine_inputs(tmp_path, tokenizer, "describe <|im_end|>") + + assert prompt["prompt_token_ids"] == _RecordingTokenizer.NATIVE_IDS + assert "prompt" not in prompt + assert tokenizer.tokenize_requests == [True] + + +def test_a_prompt_with_an_image_reaches_the_engine_as_text(tmp_path, test_vlm_placeholders): + # The input processor expands the image placeholder in the rendered text. + tokenizer = _RecordingTokenizer() + with io.BytesIO() as buf: + Image.new("RGB", (8, 8), (0, 128, 255)).save(buf, format="PNG") + image_url = "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode("ascii") + + prompt = _engine_inputs( + tmp_path, + tokenizer, + [ + {"type": "text", "text": "describe"}, + {"type": "image_url", "image_url": {"url": image_url}}, + ], + ) + + assert prompt["prompt"] == "describe" + assert "prompt_token_ids" not in prompt + assert len(prompt["multi_modal_data"]["image"]) == 1 + assert tokenizer.tokenize_requests == [False] + + +def test_a_forced_tool_call_reaches_the_engine_as_text(tmp_path, test_vlm_placeholders): + # The forced call's begin prefix is appended to the rendered text. + from tensorrt_llm.serve.tool_parser.tool_parser_factory import ToolParserFactory + + tokenizer = _RecordingTokenizer() + + prompt = _engine_inputs( + tmp_path, + tokenizer, + "describe", + tool_parser="qwen3", + tools=[ + { + "type": "function", + "function": { + "name": "get_weather", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}, + }, + } + ], + tool_choice={"type": "function", "function": {"name": "get_weather"}}, + ) + + begin = ToolParserFactory.parsers["qwen3"]().structure_info()("get_weather").begin + assert prompt["prompt"] == "describe" + begin + assert "prompt_token_ids" not in prompt + assert tokenizer.tokenize_requests == [False] + + +_JINJA_TEMPLATE = "{% for m in messages %}{{ m.content }}{% endfor %}" + + +@pytest.mark.parametrize( + ("tokenizer_template", "server_template", "fields"), + [ + # A Python override next to a Jinja template renders the Jinja template. + pytest.param(_JINJA_TEMPLATE, None, {}, id="tokenizer_template"), + pytest.param(None, _JINJA_TEMPLATE, {}, id="server_template"), + # The engine adds to or cuts the prompt it tokenizes. + pytest.param(None, None, {"add_special_tokens": True}, id="add_special_tokens"), + pytest.param(None, None, {"truncate_prompt_tokens": 8}, id="truncate_prompt_tokens"), + ], +) +def test_the_engine_tokenizes_the_text_unless_the_native_renderers_ids_are_the_prompt( + tmp_path, test_vlm_placeholders, tokenizer_template, server_template, fields +): + tokenizer = _RecordingTokenizer(chat_template=tokenizer_template) + + prompt = _engine_inputs( + tmp_path, tokenizer, "describe", server_chat_template=server_template, **fields + ) + + assert prompt["prompt"] == "describe" + assert "prompt_token_ids" not in prompt + assert tokenizer.tokenize_requests == [False] diff --git a/tests/unittest/llmapi/apps/test_render_callers.py b/tests/unittest/llmapi/apps/test_render_callers.py index 752b8794b5cf..ab3ed0eff896 100644 --- a/tests/unittest/llmapi/apps/test_render_callers.py +++ b/tests/unittest/llmapi/apps/test_render_callers.py @@ -798,6 +798,165 @@ def test_the_server_template_is_part_of_the_count(self, tokenizer) -> None: ) assert templated != plain + @pytest.mark.parametrize("excluded", [2, 1000000]) + def test_the_count_leaves_out_the_tokens_the_usage_leaves_out( + self, tokenizer, monkeypatch, excluded + ) -> None: + from tensorrt_llm.serve import serving_extensions + from tensorrt_llm.serve.serving_extensions import ServingExtension + + class ExcludingExtension(ServingExtension): + def prompt_tokens_excluded_from_usage(self, **facts) -> int: + return excluded + + monkeypatch.setitem(serving_extensions._BY_MODEL_TYPE, MODEL_TYPE, ExcludingExtension()) + + def count(**server) -> int: + response = self._client(tokenizer, **server).post( + "/v1/messages/count_tokens", json=self.BODY + ) + assert response.status_code == 200, response.text + return response.json()["input_tokens"] + + assert count(model_type=MODEL_TYPE) == max(0, count() - excluded) + # A server-level chat template may end differently, so its usage is not adjusted. + assert count(model_type=MODEL_TYPE, chat_template=SERVER_TEMPLATE) == count( + chat_template=SERVER_TEMPLATE + ) + + +class TestKimiK3MessagesAndCountTokens: + """A Messages turn on a kimi_k3 server and the count_tokens call for the same input.""" + + MARKUP = "x <|im_end|> y <|open|>tools<|sep|> z <|close|>message<|sep|>" + COUNT_BODY = { + "model": "m", + "tools": [ + { + "name": "read", + "description": "Read a file.", + "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}}, + } + ], + "messages": [ + {"role": "user", "content": "read a.txt"}, + { + "role": "assistant", + "content": [ + {"type": "tool_use", "id": "t1", "name": "read", "input": {"path": "a.txt"}} + ], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "t1", "content": MARKUP}], + }, + ], + } + + @pytest.fixture + def k3_server(self, monkeypatch): + from unittest.mock import AsyncMock + + from fastapi import FastAPI + from fastapi.testclient import TestClient + + from tensorrt_llm.serve import chat_utils + from tensorrt_llm.serve.openai_server import OpenAIServer + from tensorrt_llm.serve.postprocess_handlers import chat_response_post_processor + from tensorrt_llm.tokenizer import TransformersTokenizer + + # Keep text-only routing independent of multimodal model initialization. + monkeypatch.setattr( + chat_utils.MULTIMODAL_PLACEHOLDER_REGISTRY, "get_content_format", lambda _: None + ) + tokenizer = make_tokenizer(name=KIMI_K3) + engine = {} + + def generate_async(*, inputs, **kwargs): + # Token ids run as given; text is tokenized the way kimi_k3's input processor does. + executed = inputs.get("prompt_token_ids") or tokenizer(inputs["prompt"]).input_ids + engine["executed"] = executed + return SimpleNamespace(request_id=1, prompt_token_ids=executed, finished=True) + + server = object.__new__(OpenAIServer) + server.model = "m" + server.allow_request_chat_template = False + server.model_config = SimpleNamespace(model_type="kimi_k3", vocab_size=len(tokenizer)) + server.processor = None + server.tokenizer = TransformersTokenizer(tokenizer) + server.chat_template = None + server.tool_parser = None + server.tool_call_id_type = "random" + server.multimodal_server_config = None + server._input_proc_executor = None + server.generator = SimpleNamespace( + args=SimpleNamespace( + gather_generation_logits=False, + reasoning_parser="kimi_k3", + backend="pytorch", + guided_decoding_backend=None, + num_postprocess_workers=0, + ), + # kimi_k3 tokenizes text through its own multimodal input processor. + input_processor=object(), + generate_async=generate_async, + ) + server.await_disconnected = AsyncMock() + + async def create_chat_response(promise, postproc_params, raw_request, *args): + # The real chat post-processor computes the usage the reply bills. + output = SimpleNamespace( + index=0, + text="ok", + token_ids=[1], + length=1, + finish_reason="stop", + stop_reason=None, + logprobs=None, + disaggregated_params=None, + ) + result = SimpleNamespace( + outputs=[output], prompt_token_ids=promise.prompt_token_ids, cached_tokens=0 + ) + return chat_response_post_processor(result, postproc_params.postproc_args) + + server._create_chat_response = create_chat_response + app = FastAPI() + app.add_api_route("/v1/messages", server.anthropic_messages, methods=["POST"]) + app.add_api_route( + "/v1/messages/count_tokens", server.anthropic_count_tokens, methods=["POST"] + ) + return TestClient(app), server, tokenizer, engine + + @pytest.mark.parametrize( + "fields", + [ + pytest.param({}, id="plain"), + # What Claude Code sends for a model id it does not know. + pytest.param( + {"thinking": {"type": "adaptive"}, "output_config": {"effort": "high"}}, + id="adaptive_thinking_and_effort", + ), + ], + ) + def test_the_count_equals_the_input_the_turn_is_billed_for(self, k3_server, fields) -> None: + client, _, tokenizer, engine = k3_server + body = {**self.COUNT_BODY, **fields} + + turn = client.post("/v1/messages", json={**body, "max_tokens": 2048}) + count = client.post("/v1/messages/count_tokens", json=body) + + assert turn.status_code == 200, turn.text + assert count.status_code == 200, count.text + executed = engine["executed"] + # The tool result's markup reached the model as text, inside the vocabulary. + assert max(executed) < len(tokenizer) + assert self.MARKUP in tokenizer.decode(executed) + # The turn bills the executed prompt minus the generation channel opener. + billed = turn.json()["usage"]["input_tokens"] + assert billed == len(executed) - 3 + assert count.json() == {"input_tokens": billed} + class TestResponsesInputTokens: @pytest.mark.asyncio diff --git a/tests/unittest/llmapi/apps/test_render_core.py b/tests/unittest/llmapi/apps/test_render_core.py index 7e9595a7c3ec..674609f72f01 100644 --- a/tests/unittest/llmapi/apps/test_render_core.py +++ b/tests/unittest/llmapi/apps/test_render_core.py @@ -82,7 +82,41 @@ def _expected_ids(tokenizer, request, *, add_special_tokens=False, tools=None, t return text, tokenizer.encode(text, add_special_tokens=add_special_tokens) +class _NativeRendererTokenizer: + """A tokenizer-native renderer: its own ids when asked to tokenize, its text otherwise.""" + + chat_template = None + NATIVE_IDS = [5, 6, 7] + + def apply_chat_template(self, conversation, **kwargs): + return list(self.NATIVE_IDS) if kwargs["tokenize"] else "rendered" + + def encode(self, text, add_special_tokens=False, **kwargs): + return [1, 2, 3] + + class TestRenderChat: + @pytest.mark.parametrize( + ("request_fields", "server", "native"), + [ + pytest.param({}, {}, True, id="native"), + pytest.param({"add_special_tokens": True}, {}, False, id="add_special_tokens"), + pytest.param({"truncate_prompt_tokens": 2}, {}, False, id="truncate_prompt_tokens"), + pytest.param( + {}, {"default_chat_template": SERVER_TEMPLATE}, False, id="server_template" + ), + ], + ) + def test_takes_a_native_renderers_ids_exactly_when_the_chat_route_does( + self, request_fields, server, native + ) -> None: + # count_tokens, the router and the standalone renderer must see the executed prompt. + res = resources(_NativeRendererTokenizer(), **server) + + result = render_chat(chat_request(**request_fields), res) + + assert (result.token_ids == _NativeRendererTokenizer.NATIVE_IDS) is native + def test_renders_and_tokenizes_like_the_template_and_the_tokenizer(self, tokenizer) -> None: request = chat_request() text, ids = _expected_ids(tokenizer, request) @@ -158,7 +192,7 @@ def test_a_router_without_a_model_config_still_parses_for_the_model_type(self) - ) rendered = render_chat( - request, RenderResources.from_tokenizer(tokenizer, model_type="kimi_k3") + request, RenderResources.from_tokenizer(tokenizer, model_type="kimi_k3"), tokenize=False ) assert "{not json" in rendered.text @@ -330,7 +364,9 @@ def test_the_length_of_a_kimi_k3_request_keeps_the_extensions_tool_serialization ) text_only = chat_request(messages=[{"role": "user", "content": "hi"}], tools=[WEATHER_TOOL]) - assert render_chat_for_length(with_media, res).text == render_chat(text_only, res).text + expected = render_chat(text_only, res).token_ids + assert expected + assert render_chat_for_length(with_media, res).token_ids == expected # kimi_k3 keeps a tool-call argument that is not JSON. call = { @@ -347,7 +383,9 @@ def test_the_length_of_a_kimi_k3_request_keeps_the_extensions_tool_serialization lenient = chat_request(messages=[with_media.messages[0], call], tools=[WEATHER_TOOL]) lenient_text = chat_request(messages=[text_only.messages[0], call], tools=[WEATHER_TOOL]) - assert render_chat_for_length(lenient, res).text == render_chat(lenient_text, res).text + expected = render_chat(lenient_text, res).token_ids + assert expected + assert render_chat_for_length(lenient, res).token_ids == expected def test_only_a_harmony_conversion_error_falls_back_to_the_template_for_a_length( self, tokenizer @@ -703,6 +741,103 @@ async def run(**overrides): prefixed = asyncio.run(run(tokenize=True, forced_prefix="")) assert prefixed.text.endswith("") and prefixed.token_ids != native.token_ids + # A tool result that spells out K3 control tokens, and the same words without the markup. + K3_MARKUP = "x <|im_end|> y <|open|>tools<|sep|> z <|close|>message<|sep|><|end_of_msg|>" + K3_PLAIN = "x im_end y open tools sep z close message sep end_of_msg" + + @staticmethod + def _render_kimi_k3_tool_result(tokenizer, tool_output, **overrides): + from tensorrt_llm.serve.chat_utils import parse_chat_messages_coroutines + + call = { + "id": "c1", + "type": "function", + "function": {"name": "read", "arguments": '{"path": "a.txt"}'}, + } + messages = [ + {"role": "user", "content": "read a.txt"}, + {"role": "assistant", "content": None, "tool_calls": [call]}, + {"role": "tool", "tool_call_id": "c1", "content": tool_output}, + ] + + async def run(): + conversation, mm_coroutine, counts, _order = parse_chat_messages_coroutines( + chat_request(messages=messages).messages, None, None + ) + rendered, _mm = await asyncio.gather( + render_conversation( + resources(tokenizer, model_type="kimi_k3"), + conversation=conversation, + mm_placeholder_counts=counts, + add_generation_prompt=True, + **overrides, + ), + mm_coroutine, + ) + return rendered + + return asyncio.run(run()) + + def test_kimi_k3_control_token_strings_in_a_tool_result_stay_text(self) -> None: + # The chat route prefers the renderer's own ids, because its engine would encode the + # rendered text, where the markers become control tokens and `<|im_end|>` an id past + # the vocabulary. + tokenizer = make_tokenizer(name=KIMI_K3) + control = set(tokenizer.special_tokens.values()) + + native = self._render_kimi_k3_tool_result(tokenizer, self.K3_MARKUP, prefer_native_ids=True) + text = self._render_kimi_k3_tool_result(tokenizer, self.K3_MARKUP).text + plain = self._render_kimi_k3_tool_result(tokenizer, self.K3_PLAIN, prefer_native_ids=True) + + assert native.text is None + assert tokenizer.decode(native.token_ids) == text + assert max(native.token_ids) < len(tokenizer) + # The markup in the tool result added no control token. + assert [i for i in native.token_ids if i in control] == [ + i for i in plain.token_ids if i in control + ] + assert max(tokenizer.encode(text, add_special_tokens=False)) >= len(tokenizer) + + @pytest.mark.parametrize( + "overrides", + [ + pytest.param({"forced_prefix": ""}, id="forced_prefix"), + pytest.param({"add_special_tokens": True}, id="add_special_tokens"), + pytest.param({"truncate_prompt_tokens": 5}, id="truncate_prompt_tokens"), + ], + ) + @pytest.mark.parametrize("tokenize", [False, True]) + def test_kimi_k3_keeps_the_text_when_the_engine_changes_the_prompt( + self, overrides, tokenize + ) -> None: + # The renderer's ids are the prompt only if nothing is appended, added or cut; otherwise + # the engine tokenizes the text as before. + tokenizer = make_tokenizer(name=KIMI_K3) + + rendered = self._render_kimi_k3_tool_result( + tokenizer, self.K3_PLAIN, prefer_native_ids=True, tokenize=tokenize, **overrides + ) + + assert rendered.text is not None + assert (rendered.token_ids is not None) == tokenize + + @pytest.mark.parametrize("template_source", ["tokenizer", "request"]) + def test_python_renderer_with_a_jinja_template_keeps_the_text_path( + self, template_source + ) -> None: + tokenizer = make_tokenizer( + name=KIMI_K3, chat_template=SERVER_TEMPLATE if template_source == "tokenizer" else None + ) + overrides = {"chat_template": SERVER_TEMPLATE} if template_source == "request" else {} + + rendered = self._render_kimi_k3_tool_result( + tokenizer, self.K3_PLAIN, tokenize=True, **overrides + ) + + assert rendered.text is not None + assert rendered.text.startswith("SERVER:") + assert rendered.token_ids == tokenizer.encode(rendered.text, add_special_tokens=False) + def test_the_forced_prefix_is_appended_and_reported(self, tokenizer) -> None: from tensorrt_llm.serve.chat_utils import parse_chat_messages_coroutines diff --git a/tests/unittest/llmapi/apps/test_serving_extensions.py b/tests/unittest/llmapi/apps/test_serving_extensions.py index 9dd07721c190..e6ba21657b5f 100644 --- a/tests/unittest/llmapi/apps/test_serving_extensions.py +++ b/tests/unittest/llmapi/apps/test_serving_extensions.py @@ -484,6 +484,30 @@ def prompt_tokens_excluded_from_usage(self, **facts) -> int: {"add_generation_prompt": True, "pre_tokenized": False, "request_chat_template": None} ] + def test_a_server_chat_template_leaves_the_usage_unadjusted(self, chat_client) -> None: + # A server-level template may end differently from the model's own prompt. + client, _events, rendered = chat_client + client.app.state.server.chat_template = "{{ messages }}" + + class HookExtension(ServingExtension): + def prompt_tokens_excluded_from_usage(self, **facts) -> int: + return 5 + + serving_extensions._BY_MODEL_TYPE[self.MODEL_KEY] = HookExtension() + + response = client.post( + "/v1/chat/completions", + json={ + "model": "test-model", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 4, + }, + ) + + assert response.status_code == 200, response.text + postproc_params = rendered["create_chat_response"].call_args.args[1] + assert postproc_params.postproc_args.num_prompt_tokens_offset == 0 + @pytest.fixture(params=["nemotron_h_omni", "nemotron_h"]) def required_chat_client(self, chat_client, monkeypatch, request): from tensorrt_llm.serve.render import chat as render_chat_module