[None][fix] serve: take Kimi-K3 chat prompt ids from its renderer; count_tokens matches usage

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.
 - <span style="color: red">**[BREAKING CHANGE] Removed obsolete C++ executor orchestration helpers.**</span> 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.
 - <span style="color: red">**[BREAKING CHANGE] Unified multimodal processor metadata.**</span>

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,

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.

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:

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",

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.