[TRTLLM-17136][feat] Serve Muse Spark Open: reasoning and tool parsers, effort scale, native tools

tensorrt_llm/inputs/chat_template_guard.py

@@ -4,7 +4,7 @@
 
 import hashlib
 import os
-from collections.abc import Collection
+from collections.abc import Collection, Iterable
 from functools import lru_cache
 
 import jinja2
@@ -118,6 +118,36 @@ def _warn_unused_key(template_hash: str, key: str) -> None:
     )
 
 
+def unreferenced_chat_template_kwargs(
+    template_source: str | None,
+    keys: Iterable[str],
+) -> frozenset[str]:
+    """Subset of `keys` the active template never reads.
+
+    The same reference analysis the validator runs, exposed as a query so
+    that code which *constructs* template kwargs (rather than forwarding a
+    caller's) can avoid injecting a key the active template never reads: a
+    key is flagged when the template does not reference it and the renderer
+    does not always consume it. For a template the analysis cannot see into
+    (parse failure, includes/imports) every key outside the always-consumed
+    surface is flagged: constructing code can safely omit an unverifiable
+    key and leave the template its own default, whereas the validator (which
+    faces caller-supplied kwargs it cannot remove) skips instead. A
+    non-string `template_source` flags nothing, matching the validator's
+    skip.
+    """
+    if not isinstance(template_source, str):
+        return frozenset()
+    referenced = _referenced_template_variables(template_source)
+    if referenced is None:
+        return frozenset(key for key in keys if key not in ALWAYS_ALLOWED_CHAT_TEMPLATE_KWARGS)
+    return frozenset(
+        key
+        for key in keys
+        if key not in referenced and key not in ALWAYS_ALLOWED_CHAT_TEMPLATE_KWARGS
+    )
+
+
 def validate_chat_template_kwargs(
     template_source: str | None,
     chat_template_kwargs: dict[str, object] | None,

tensorrt_llm/inputs/utils.py

@@ -6,8 +6,8 @@
 from collections import defaultdict
 from io import BytesIO
 from pathlib import Path
-from typing import (Any, Collection, Coroutine, Dict, List, Optional, Tuple,
-                    TypedDict, Union)
+from typing import (Any, Collection, Coroutine, Dict, List, NamedTuple,
+                    Optional, Tuple, TypedDict, Union)
 from urllib.parse import urlparse
 
 import numpy as np
@@ -659,6 +659,78 @@ def resolve_hf_chat_template(
     return None
 
 
+class _RenderDispatch(NamedTuple):
+    """The renderer `apply_chat_template` uses for a set of inputs.
+
+    `kind` is `"deepseek_v41"`, `"deepseek"` (the other DeepSeek
+    tokenizers), `"passthrough"` (the model skips chat templates),
+    `"native"` (a tokenizer-native python template) or `"jinja"`.
+    `tokenizer` is the object that renderer is called on, and `template`
+    is the resolved source for `"jinja"` (None when no template resolves).
+    """
+    kind: str
+    tokenizer: Any
+    template: Optional[str] = None
+
+
+def _resolve_render_dispatch(
+    *,
+    model_type: Optional[str],
+    tokenizer: Union[TransformersTokenizer, TokenizerBase],
+    processor: ProcessorMixin,
+    chat_template: Optional[str],
+    tools: Optional[list[dict[str, Any]]],
+) -> _RenderDispatch:
+    """Pick the renderer, in the order `apply_chat_template` checks them."""
+    if isinstance(tokenizer, DeepseekV41Tokenizer):
+        return _RenderDispatch("deepseek_v41", tokenizer)
+    if isinstance(tokenizer, (DeepseekV32Tokenizer, DeepseekV4Tokenizer)):
+        return _RenderDispatch("deepseek", tokenizer)
+    # Before the tokenizer, processor and template are needed: the registry
+    # may already know this model skips chat templates entirely.
+    if model_type and MULTIMODAL_PLACEHOLDER_REGISTRY.get_content_format(
+            model_type) == ContentFormat.PASSTHROUGH:
+        return _RenderDispatch("passthrough", tokenizer)
+    if isinstance(tokenizer, TransformersTokenizer):
+        # The TokenizerBase renders the template.
+        tokenizer = tokenizer.tokenizer
+    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)):
+        return _RenderDispatch("native", tokenizer)
+    return _RenderDispatch(
+        "jinja", tokenizer,
+        resolve_hf_chat_template(tokenizer, processor, chat_template, tools))
+
+
+def rendered_hf_chat_template_source(
+    *,
+    tokenizer: Union[TransformersTokenizer, TokenizerBase],
+    processor: ProcessorMixin,
+    chat_template: Optional[str] = None,
+    tools: Optional[list[dict[str, Any]]] = None,
+    model_type: Optional[str] = None,
+) -> Optional[str]:
+    """The Jinja source `apply_chat_template` will render for these inputs.
+
+    Returns `None` when `apply_chat_template` dispatches to a renderer
+    with no Jinja source to analyze (the DeepSeek custom tokenizers, a model
+    that skips chat templates, a tokenizer-native python template) or when no
+    template resolves. Callers that construct chat_template_kwargs use this
+    to run the same referenced-kwargs analysis the render-site guard applies,
+    so they can avoid injecting a kwarg the active template never reads.
+    """
+    dispatch = _resolve_render_dispatch(model_type=model_type,
+                                        tokenizer=tokenizer,
+                                        processor=processor,
+                                        chat_template=chat_template,
+                                        tools=tools)
+    if dispatch.kind != "jinja" or not isinstance(dispatch.template, str):
+        return None
+    return dispatch.template
+
+
 def _resolve_content_format(model_type: str,
                             chat_template: Optional[str]) -> ContentFormat:
     """Determine the content format for the given model and template.
@@ -747,7 +819,13 @@ def apply_chat_template(
     - STRING: keeps flattened text with pre-inserted placeholders
     """
 
-    if isinstance(tokenizer, DeepseekV41Tokenizer):
+    dispatch = _resolve_render_dispatch(model_type=model_type,
+                                        tokenizer=tokenizer,
+                                        processor=processor,
+                                        chat_template=chat_template,
+                                        tools=tools)
+
+    if dispatch.kind == "deepseek_v41":
         # The native renderer deep-copies messages; omit pending media coroutines
         # and retain flattened content with its interleaved image placeholders.
         messages = [{
@@ -763,9 +841,7 @@ def apply_chat_template(
                                              **native_kwargs)
 
     # Handle DeepSeek tokenizers with custom chat templates.
-    if isinstance(
-            tokenizer,
-        (DeepseekV32Tokenizer, DeepseekV4Tokenizer, DeepseekV41Tokenizer)):
+    if dispatch.kind == "deepseek":
         prompt = tokenizer.apply_chat_template(
             messages=conversation,
             tools=tools,
@@ -775,20 +851,11 @@ def apply_chat_template(
             return tokenizer.encode(prompt)
         return prompt
 
-    # Check for PASSTHROUGH early — before we need tokenizer/processor/template.
-    # The registry may already know this model skips chat templates entirely.
-    registry_format = MULTIMODAL_PLACEHOLDER_REGISTRY.get_content_format(
-        model_type)
-    if registry_format == ContentFormat.PASSTHROUGH:
+    if dispatch.kind == "passthrough":
         return "".join([conv["content"] for conv in conversation])
 
-    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)):
+    tokenizer = dispatch.tokenizer
+    if dispatch.kind == "native":
         native_kwargs = dict(chat_template_kwargs or {})
         if documents is not None:
             native_kwargs["documents"] = documents
@@ -800,8 +867,7 @@ def apply_chat_template(
             **native_kwargs,
         )
 
-    hf_chat_template = resolve_hf_chat_template(tokenizer, processor,
-                                                chat_template, tools)
+    hf_chat_template = dispatch.template
     if hf_chat_template is None:
         raise ValueError(
             "No chat template found for the given tokenizer and tools.")

tensorrt_llm/llmapi/reasoning/__init__.py

@@ -0,0 +1,19 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Model-specific reasoning parser implementations.
+
+Each module defines a parser registered with
+`tensorrt_llm.llmapi.reasoning_parser.register_reasoning_parser`.
+"""