[None][perf] offload chat template rendering in serving

Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>

Author: yechank-nvidia

tensorrt_llm/inputs/utils.py

@@ -705,6 +705,37 @@ def apply_chat_template(
     return result
 
 
+async def async_apply_chat_template(
+    *,
+    model_type: str,
+    tokenizer: Union[TransformersTokenizer, TokenizerBase],
+    processor: ProcessorMixin,
+    conversation: list[ConversationMessage],
+    add_generation_prompt: bool,
+    mm_placeholder_counts: list[dict[str, int]],
+    tools: Optional[list[dict[str, Any]]] = None,
+    documents: Optional[list[dict[str, str]]] = None,
+    chat_template: Optional[str] = None,
+    chat_template_kwargs: Optional[dict[str, Any]] = None,
+    enable_tokenize: bool = False,
+) -> (str | List[str]):
+    """Apply chat template without blocking the event loop."""
+    return await asyncio.to_thread(
+        apply_chat_template,
+        model_type=model_type,
+        tokenizer=tokenizer,
+        processor=processor,
+        conversation=conversation,
+        add_generation_prompt=add_generation_prompt,
+        mm_placeholder_counts=mm_placeholder_counts,
+        tools=tools,
+        documents=documents,
+        chat_template=chat_template,
+        chat_template_kwargs=chat_template_kwargs,
+        enable_tokenize=enable_tokenize,
+    )
+
+
 def default_multimodal_input_loader(
     *,
     tokenizer: Optional[Union[TransformersTokenizer, TokenizerBase]],

tensorrt_llm/serve/openai_server.py

@@ -35,7 +35,8 @@
 from tensorrt_llm.inputs import prompt_inputs
 from tensorrt_llm.inputs.data import TokensPrompt
 from tensorrt_llm.inputs.multimodal import MultimodalServerConfig
-from tensorrt_llm.inputs.utils import ConversationMessage, apply_chat_template
+from tensorrt_llm.inputs.utils import (ConversationMessage,
+                                       async_apply_chat_template)
 from tensorrt_llm.llmapi import DisaggregatedParams as LlmDisaggregatedParams
 from tensorrt_llm.llmapi import MultimodalEncoder, SchedulingParams, tracing
 from tensorrt_llm.llmapi.disagg_utils import (DisaggClusterConfig,
@@ -1261,7 +1262,7 @@ async def chat_stream_generator(
             if request.prompt_token_ids is not None:
                 prompt = request.prompt_token_ids
             else:
-                prompt: str = apply_chat_template(
+                prompt_task = async_apply_chat_template(
                     model_type=resolve_top_level_model_type(self.model_config),
                     tokenizer=self.tokenizer,
                     processor=self.processor,
@@ -1273,9 +1274,12 @@ async def chat_stream_generator(
                     chat_template=request.chat_template or self.chat_template,
                     chat_template_kwargs=request.chat_template_kwargs or {},
                 )
+                prompt, (mm_data, mm_embeddings) = await asyncio.gather(
+                    prompt_task, mm_coroutines)
             prompt = prompt_inputs(prompt)
 
-            mm_data, mm_embeddings = await mm_coroutines
+            if request.prompt_token_ids is not None:
+                mm_data, mm_embeddings = await mm_coroutines
             if mm_data:
                 prompt["multi_modal_data"] = mm_data
             if mm_embeddings:
@@ -1414,7 +1418,7 @@ async def create_mm_embedding_response(promise: RequestOutput):
             if request.prompt_token_ids is not None:
                 prompt = request.prompt_token_ids
             else:
-                prompt: str = apply_chat_template(
+                prompt_task = async_apply_chat_template(
                     model_type=resolve_top_level_model_type(self.model_config),
                     tokenizer=self.tokenizer,
                     processor=self.processor,
@@ -1426,9 +1430,12 @@ async def create_mm_embedding_response(promise: RequestOutput):
                     chat_template=request.chat_template,
                     chat_template_kwargs=request.chat_template_kwargs or {},
                 )
+                prompt, (mm_data, mm_embeddings) = await asyncio.gather(
+                    prompt_task, mm_coroutines)
             prompt = prompt_inputs(prompt)
 
-            mm_data, mm_embeddings = await mm_coroutines
+            if request.prompt_token_ids is not None:
+                mm_data, mm_embeddings = await mm_coroutines
             if mm_embeddings:
                 raise ValueError("Cannot use multimodal embeddings as input")
             if mm_data is not None:

tensorrt_llm/serve/resource_governor.py

@@ -19,6 +19,7 @@
 LLM/Proxy/Worker chain.
 """
 
+import asyncio
 import traceback
 from http import HTTPStatus
 from typing import Callable, List, Optional
@@ -27,7 +28,7 @@
 from starlette.responses import JSONResponse, Response
 
 from tensorrt_llm.executor.request import TruncateKVCacheRequest
-from tensorrt_llm.inputs.utils import ConversationMessage, apply_chat_template
+from tensorrt_llm.inputs.utils import ConversationMessage, async_apply_chat_template
 from tensorrt_llm.logger import logger
 from tensorrt_llm.serve.chat_utils import parse_chat_messages_coroutines
 from tensorrt_llm.serve.openai_protocol import KVCacheTruncateRequest
@@ -86,7 +87,7 @@ def _put_or_unavailable(self, request: TruncateKVCacheRequest) -> Optional[Respo
         queue.put(request)
         return None
 
-    def _convert_messages(
+    async def _convert_messages(
         self,
         messages,
         tool_dicts,
@@ -97,20 +98,24 @@ def _convert_messages(
     ) -> List[int]:
         """Convert chat messages to token IDs via chat template + tokenization."""
         conversation: List[ConversationMessage] = []
-        conversation, _, __ = parse_chat_messages_coroutines(messages, self.model_config, None)
-        return apply_chat_template(
+        conversation, mm_coroutines, mm_placeholder_counts = parse_chat_messages_coroutines(
+            messages, self.model_config, None
+        )
+        token_task = async_apply_chat_template(
             model_type=self.model_config.model_type,
             tokenizer=self.tokenizer,
             processor=self.processor,
             conversation=conversation,
             add_generation_prompt=add_generation_prompt,
-            mm_placeholder_counts=[],
+            mm_placeholder_counts=mm_placeholder_counts,
             tools=tool_dicts,
             documents=documents,
             chat_template=chat_template,
             chat_template_kwargs=chat_template_kwargs or {},
             enable_tokenize=True,
         )
+        token_ids, _ = await asyncio.gather(token_task, mm_coroutines)
+        return token_ids
 
     async def _truncate_kv_cache(self, request: KVCacheTruncateRequest) -> Response:
         try:
@@ -120,7 +125,7 @@ async def _truncate_kv_cache(self, request: KVCacheTruncateRequest) -> Response:
             chat_template_kwargs = request.chat_template_kwargs or {}
 
             messages_to_retain = (
-                self._convert_messages(
+                await self._convert_messages(
                     request.messages_to_retain,
                     tool_dicts,
                     request.add_generation_prompt,
@@ -133,7 +138,7 @@ async def _truncate_kv_cache(self, request: KVCacheTruncateRequest) -> Response:
             )
 
             messages = (
-                self._convert_messages(
+                await self._convert_messages(
                     request.messages,
                     tool_dicts,
                     request.add_generation_prompt,

tensorrt_llm/serve/responses_utils.py

@@ -42,7 +42,7 @@
 
 from tensorrt_llm.bindings import steady_clock_now
 from tensorrt_llm.executor import GenerationResult
-from tensorrt_llm.inputs.utils import apply_chat_template
+from tensorrt_llm.inputs.utils import async_apply_chat_template
 from tensorrt_llm.llmapi import SamplingParams
 from tensorrt_llm.llmapi.llm import RequestOutput
 from tensorrt_llm.llmapi.reasoning_parser import (BaseReasoningParser,
@@ -821,13 +821,11 @@ async def _create_input_tokens(
 
     conversation, mm_coroutines, mm_placeholder_counts = parse_chat_messages_coroutines(
         messages, model_config)
-    mm_data = await mm_coroutines
-
     tools_dict = [
         tool.model_dump()
         for tool in _get_chat_completion_function_tools(request.tools)
     ]
-    token_ids = apply_chat_template(
+    token_task = async_apply_chat_template(
         model_type=resolve_top_level_model_type(model_config),
         tokenizer=tokenizer,
         processor=processor,
@@ -837,6 +835,7 @@ async def _create_input_tokens(
         mm_placeholder_counts=mm_placeholder_counts,
         enable_tokenize=True,
     )
+    token_ids, mm_data = await asyncio.gather(token_task, mm_coroutines)
 
     return token_ids, mm_data
 

tests/unittest/inputs/test_chat_template_dispatch.py

@@ -2,6 +2,8 @@
 # Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 """Tests for content-format-driven chat template dispatch and placeholder handling."""
 
+import threading
+
 import pytest
 
 from tensorrt_llm.inputs.content_format import ContentFormat
@@ -15,6 +17,7 @@
     _build_openai_content,
     _resolve_content_format,
     add_multimodal_placeholders,
+    async_apply_chat_template,
     interleave_mm_placeholders,
 )
 
@@ -324,3 +327,33 @@ def test_excess_existing_placeholders_preserved(self):
         )
         assert result == text
         assert result.count("<image>") == 3
+
+
+class TestAsyncApplyChatTemplate:
+    @pytest.mark.asyncio
+    async def test_runs_in_worker_thread(self):
+        event_loop_thread_id = threading.current_thread().ident
+
+        class TrackingTokenizer:
+            def __init__(self):
+                self.worker_thread_id = None
+
+            def apply_chat_template(self, **_):
+                self.worker_thread_id = threading.current_thread().ident
+                return "rendered"
+
+        tokenizer = TrackingTokenizer()
+
+        result = await async_apply_chat_template(
+            model_type="test_string_model",
+            tokenizer=tokenizer,
+            processor=None,
+            conversation=[ConversationMessage(role="user", content="hello", media=[])],
+            add_generation_prompt=True,
+            mm_placeholder_counts=[{}],
+            chat_template="{{ messages }}",
+        )
+
+        assert result == "rendered"
+        assert tokenizer.worker_thread_id is not None
+        assert tokenizer.worker_thread_id != event_loop_thread_id