[https://nvbugs/6932723][feat] Support required tool choice for Nemotron Super 3.5 VL (#20422)

Signed-off-by: Yueh-Ting Chen <yuehtingc@nvidia.com>

Author: eopXD

tensorrt_llm/llmapi/reasoning_parser.py

@@ -768,6 +768,10 @@ class NemotronV3ReasoningParser(DeepSeekR1Parser):
     The `enable_thinking` flag is read from `chat_template_kwargs`.
     """
 
+    resolves_thinking_from_prompt = True
+    reasoning_start = "<think>"
+    reasoning_end = "</think>"
+
     def __init__(self,
                  *,
                  reasoning_at_start: bool = True,
@@ -803,32 +807,35 @@ def _maybe_swap_content(
         return result
 
     def parse_delta(self, delta_text: str) -> ReasoningParserResult:
-        """Wraps the parent parse_delta to also treat `<tool_call>` as an
-        implicit end-of-reasoning marker.  When the model omits `</think>`
-        before generating a tool call, the tag would otherwise be absorbed
-        into reasoning_content and the downstream tool parser would never
-        see it (NVBug 6082303).
-
-        `<tool_call>` is a special token that always arrives as a single
-        atomic delta, so we only need to check `delta_text` (not the
-        parent's internal buffer)."""
-        if (self.in_reasoning and self._tool_call_start in delta_text
-                and self.reasoning_end not in self._buffer):
-            remaining = self._buffer
+        """End reasoning at the first closing marker or tool-call opener."""
+        buffered = self._buffer + delta_text
+        tool_idx = buffered.find(self._tool_call_start)
+        end_idx = buffered.find(self.reasoning_end)
+        if self.in_reasoning and tool_idx != -1 and (end_idx == -1
+                                                     or tool_idx < end_idx):
             self._buffer = ""
             self.in_reasoning = False
-            # Guaranteed non-negative: guarded by `in delta_text` above.
-            tool_idx = delta_text.find(self._tool_call_start)
-            reasoning = remaining + delta_text[:tool_idx]
-            content = delta_text[tool_idx:]
+            reasoning = buffered[:tool_idx]
+            content = buffered[tool_idx:]
             if self._force_nonempty_content:
                 self._found_closing_tag = True
                 self._accumulated_reasoning = ""
             return ReasoningParserResult(content=content,
                                          reasoning_content=reasoning)
 
+        pending = ""
+        if self.in_reasoning and end_idx == -1 and tool_idx == -1:
+            for size in range(
+                    min(len(buffered),
+                        len(self._tool_call_start) - 1), 0, -1):
+                if buffered.endswith(self._tool_call_start[:size]):
+                    pending = buffered[-size:]
+                    self._buffer = ""
+                    delta_text = buffered[:-size]
+                    break
         was_in_reasoning = self.in_reasoning
         result = super().parse_delta(delta_text)
+        self._buffer += pending
         if self._force_nonempty_content:
             if result.reasoning_content:
                 self._accumulated_reasoning += result.reasoning_content
@@ -868,14 +875,19 @@ def finish(self) -> ReasoningParserResult:
         return ReasoningParserResult()
 
     def parse(self, text: str) -> ReasoningParserResult:
-        result = super().parse(text)
-        tc = (result.reasoning_content.find(self._tool_call_start)
-              if result.reasoning_content else -1)
-        if tc != -1:
-            result = ReasoningParserResult(
-                content=result.reasoning_content[tc:] + result.content,
-                reasoning_content=result.reasoning_content[:tc])
-        return self._maybe_swap_content(result)
+        start = 0 if self.reasoning_at_start else text.find(
+            self.reasoning_start)
+        if start != -1:
+            if not self.reasoning_at_start:
+                start += len(self.reasoning_start)
+            tool_idx = text.find(self._tool_call_start, start)
+            end_idx = text.find(self.reasoning_end, start)
+            if tool_idx != -1 and (end_idx == -1 or tool_idx < end_idx):
+                return self._maybe_swap_content(
+                    ReasoningParserResult(
+                        content=text[tool_idx:],
+                        reasoning_content=text[start:tool_idx]))
+        return self._maybe_swap_content(super().parse(text))
 
 
 @register_reasoning_parser("gemma4")

tensorrt_llm/serve/extensions/__init__.py

@@ -25,6 +25,6 @@
 turn imports the registry module. Nothing here may import ``openai_server``.
 """
 
-from . import gpt_oss, kimi_k3
+from . import gpt_oss, kimi_k3, nemotron
 
-__all__ = ["gpt_oss", "kimi_k3"]
+__all__ = ["gpt_oss", "kimi_k3", "nemotron"]

tensorrt_llm/serve/extensions/nemotron.py

@@ -0,0 +1,29 @@
+# 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.
+"""Serving extension for Nemotron checkpoints using required-tool grammars."""
+
+from tensorrt_llm.serve.serving_extensions import (
+    RequiredToolChoiceMode,
+    ServingExtension,
+    register_serving_extension,
+)
+
+
+@register_serving_extension(model_types=("nemotron_h", "nemotron_h_omni"))
+class NemotronServingExtension(ServingExtension):
+    """Use mandatory tool-call grammars for native Nemotron chat requests."""
+
+    def required_tool_choice_mode(self) -> RequiredToolChoiceMode:
+        return "grammar"

tensorrt_llm/serve/openai_server.py

@@ -229,6 +229,34 @@ def _warn_unresolvable_thinking_once(reasoning_parser: str) -> None:
         "build that relays 'resolved_thinking'.")
 
 
+def _validate_required_tool_grammar_request(
+        request: ChatCompletionRequest, sampling_params: SamplingParams, *,
+        guided_decoding_backend: Optional[str], tool_parser: Optional[str],
+        reasoning_parser: Optional[str]) -> None:
+    """Validate inputs for mandatory tool-call grammar construction."""
+    if guided_decoding_backend != "xgrammar":
+        raise ValueError("Required tool grammars need the xgrammar backend.")
+    if not tool_parser:
+        raise ValueError("Required tool grammars need a tool parser.")
+    if not reasoning_parser or not ReasoningParserFactory.resolves_thinking_from_prompt(
+            reasoning_parser):
+        raise ValueError(
+            "Required tool grammars need a reasoning parser that resolves its "
+            "mode from the rendered prompt. Configure the server with "
+            "--reasoning_parser nemotron-v3 (or a compatible parser alias).")
+    if not request.tools:
+        raise ValueError("tool_choice='required' requires request-level tools.")
+    if (request.prompt_token_ids is not None or request.prompt_token_ids_b64
+            or not request.add_generation_prompt
+            or request.disaggregated_params is not None):
+        raise ValueError("Required tool grammars need chat messages with "
+                         "add_generation_prompt=True on a native server.")
+    if sampling_params.guided_decoding is not None:
+        raise ValueError(
+            "tool_choice='required' cannot be combined with response_format "
+            "or other guided-decoding settings.")
+
+
 def _configure_parser_special_token_decoding(
         sampling_params: SamplingParams, reasoning_parser_name: Optional[str],
         tool_parser_name: Optional[str], has_tools: bool) -> None:
@@ -2044,6 +2072,9 @@ async def chat_stream_generator(
             extension = render_resources.extension
             tool_dicts = prepare_chat_request(request,
                                               render_resources).tool_dicts
+            required_mode = extension.required_tool_choice_mode()
+            required_tool_grammar = (request.tool_choice == "required"
+                                     and required_mode == "grammar")
             conversation: List[ConversationMessage] = []
             # Pass the model vocabulary size so ``logit_bias`` can be
             # expanded into an embedding bias tensor in the sampler.
@@ -2056,6 +2087,14 @@ async def chat_stream_generator(
                 gather_generation_logits,
                 reasoning_parser=self.generator.args.reasoning_parser,
                 backend=self.generator.args.backend)
+            if required_tool_grammar:
+                _validate_required_tool_grammar_request(
+                    request,
+                    sampling_params,
+                    guided_decoding_backend=self.generator.args.
+                    guided_decoding_backend,
+                    tool_parser=self.tool_parser,
+                    reasoning_parser=self.generator.args.reasoning_parser)
             # Pre-render, so the mode may be unresolved. Safe because the
             # boundary lookup reads the parser class attributes, not the
             # branch `__init__` picked.
@@ -2173,7 +2212,7 @@ async def chat_stream_generator(
                 # via structural tags (only if response_format doesn't already
                 # set guided decoding). Dynamic tools must be in the grammar
                 # too, or their calls would be masked out.
-                if sampling_params.guided_decoding is None:
+                if sampling_params.guided_decoding is None and not required_tool_grammar:
                     if (forced_tool_name is not None
                             and _parser_extracts_forced_tool_calls(
                                 self.tool_parser)):
@@ -2330,14 +2369,14 @@ async def chat_stream_generator(
             if request.mm_processor_kwargs:
                 prompt["mm_processor_kwargs"] = request.mm_processor_kwargs
 
+            thinking = None
             postproc_args.reasoning_parser = self.generator.args.reasoning_parser
             # Templates that prefill <think>/</think> leave the marker in the
             # prompt, so the request kwargs alone cannot tell the parser which
             # mode was rendered. Take it from the prompt instead.
             if postproc_args.reasoning_parser and (
                     ReasoningParserFactory.resolves_thinking_from_prompt(
                         postproc_args.reasoning_parser)):
-                thinking = None
                 if rendered_prompt and request.add_generation_prompt:
                     thinking = ReasoningParserFactory.resolve_prefilled_thinking(
                         postproc_args.reasoning_parser, rendered_prompt)
@@ -2356,6 +2395,26 @@ async def chat_stream_generator(
                         "thinking": thinking,
                         "enable_thinking": thinking,
                     }
+            if required_tool_grammar:
+                if thinking is None:
+                    raise ValueError(
+                        "Cannot resolve required-tool reasoning state from the rendered prompt."
+                    )
+                parser = ToolParserFactory.create_tool_parser(self.tool_parser)
+                build_required = functools.partial(
+                    parser.build_required_structural_tag,
+                    request.tools,
+                    thinking=thinking)
+                required_format = await asyncio.get_running_loop(
+                ).run_in_executor(self._input_proc_executor, build_required)
+                if required_format is None:
+                    raise ValueError(
+                        "The configured tool parser does not support required tools."
+                    )
+                sampling_params.guided_decoding = GuidedDecodingParams(
+                    structural_tag=required_format.model_dump_json(
+                        by_alias=True))
+                postproc_args.tool_parser_grammar_constrained = True
             postproc_args.tool_parser = self.tool_parser
             postproc_args.tool_call_id_type = self.tool_call_id_type
             if conversation and conversation[-1].get(

tensorrt_llm/serve/postprocess_handlers.py

@@ -103,6 +103,7 @@ class ChatPostprocArgs(PostprocArgs):
     last_message_content: Optional[str] = None
     reasoning_parser: Optional[str] = None
     tool_parser: Optional[str] = None
+    tool_parser_grammar_constrained: bool = False
     reasoning_parser_dict: dict[int, BaseReasoningParser] = field(
         default_factory=dict)
     tool_parser_dict: dict[int, BaseToolParser] = field(default_factory=dict)
@@ -299,7 +300,8 @@ def apply_tool_parser(args: ChatPostprocArgs,
         if output_index not in args.tool_parser_dict:
             args.tool_parser_dict[
                 output_index] = ToolParserFactory.create_tool_parser(
-                    args.tool_parser)
+                    args.tool_parser,
+                    grammar_constrained=args.tool_parser_grammar_constrained)
         tool_parser = args.tool_parser_dict[output_index]
 
     if tool_parser is not None and tools is not None:

tensorrt_llm/serve/render/chat.py

@@ -92,7 +92,7 @@ def prepare_chat_request(request: "ChatCompletionRequest", res: RenderResources)
     ensure_request_chat_template_allowed(request, res.allow_request_chat_template)
     extension = res.extension
     extension.apply_chat_extensions(request)
-    if request.tool_choice == "required" and not extension.allows_required_tool_choice():
+    if request.tool_choice == "required" and extension.required_tool_choice_mode() is None:
         # Schema-accepting "required" everywhere but enforcing it only for models
         # whose extension allows it would silently degrade to "auto" elsewhere;
         # reject loudly for models that cannot honor it. Still a ValueError for the

tensorrt_llm/serve/responses_utils.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import asyncio
@@ -1445,6 +1446,8 @@ def _apply_tool_parser(
     text: str,
     streaming: bool,
     tool_parser_dict: Optional[dict[int, BaseToolParser]] = None,
+    finished: bool = False,
+    partial: bool = False,
 ) -> Tuple[str, list[ToolCallItem]]:
     tool_parser: Optional[BaseToolParser] = None
     if tool_parser_id is not None and tools is not None:
@@ -1460,10 +1463,15 @@ def _apply_tool_parser(
 
     if tool_parser is not None and tools is not None:
         if not streaming:
-            result = tool_parser.detect_and_parse(text, tools)
+            result = (tool_parser.parse_partial(text, tools)
+                      if partial else tool_parser.detect_and_parse(text, tools))
         else:
             result = tool_parser.parse_streaming_increment(text, tools)
         normal_text, calls = result.normal_text, result.calls
+        if streaming and finished:
+            tail = tool_parser.finish(tools)
+            normal_text += tail.normal_text
+            calls += tail.calls
         if not streaming:
             warn_if_tool_call_unparsed(tool_parser_id, tool_parser, text, calls)
     else:
@@ -2159,6 +2167,7 @@ def _should_send_done_events(
             text=full_text,
             streaming=False,
             tool_parser_dict=tool_parser_dict,
+            partial=not finished_generation,
         )
 
     # Detect reasoning -> text transition
@@ -2278,15 +2287,15 @@ def check_parser(parser_id: Optional[str],
         chat_template_kwargs=reasoning_chat_template_kwargs(request),
     )
 
-    if delta_text:
-        # TODO(JunyiXu-nv): handle tool calls in streaming mode
+    if delta_text or finished_generation:
         delta_text, calls = _apply_tool_parser(
             tool_parser_id=tool_parser_id,
             tools=available_tools,
             output_index=output_idx,
             text=delta_text,
             streaming=True,
             tool_parser_dict=tool_parser_dict,
+            finished=finished_generation,
         )
 
     _responses_debug_log(

tensorrt_llm/serve/serving_extensions.py

@@ -25,7 +25,7 @@
   (:func:`structured_output_format_for`).
 
 Built-in extensions live in :mod:`tensorrt_llm.serve.extensions`; importing
-that package registers them. Both lookup functions import it on first use so
+that package registers them. The hook lookup functions import it on first use so
 the registry is populated wherever it is consulted (``openai_protocol`` reaches
 it from ``to_sampling_params`` without going through ``openai_server``). The
 import is deferred to call time because the built-ins import
@@ -34,12 +34,15 @@
 
 import enum
 import importlib
-from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Optional
+from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Literal, Optional
 
 if TYPE_CHECKING:
     from tensorrt_llm.serve.openai_protocol import ChatCompletionRequest
 
 
+RequiredToolChoiceMode = Literal["template", "grammar"]
+
+
 class OutputMode(enum.Enum):
     """How a model's generated output is turned into an API response."""
 
@@ -103,6 +106,10 @@ def apply_chat_extensions(self, request) -> None:
         ``model_type`` this extension is registered for.
         """
 
+    def required_tool_choice_mode(self) -> Optional[RequiredToolChoiceMode]:
+        """Return the required-tool policy; boolean permission selects template enforcement."""
+        return "template" if self.allows_required_tool_choice() else None
+
     def structured_output_format(
         self, content: Dict[str, Any], chat_template_kwargs: Optional[Dict[str, Any]]
     ) -> Optional[Dict[str, Any]]:
@@ -189,3 +196,10 @@ def structured_output_format_for(
     _load_builtin_extensions()
     extension = _BY_REASONING_PARSER.get(reasoning_parser) if reasoning_parser else None
     return None if extension is None else extension.structured_output_format
+
+
+def required_tool_choice_mode(model_type: Optional[str]) -> Optional[RequiredToolChoiceMode]:
+    """Return required-tool handling for native chat, or None if unsupported."""
+    _load_builtin_extensions()
+    extension = _BY_MODEL_TYPE.get(model_type) if model_type else None
+    return None if extension is None else extension.required_tool_choice_mode()