[https://nvbugs/6284101][fix] Fix guided decoding for Harmony parser

tensorrt_llm/evaluate/interface.py

@@ -26,6 +26,8 @@
 import tensorrt_llm.profiler as profiler
 
 from ..llmapi import RequestOutput
+from ..llmapi.reasoning_parser import (HARMONY_REASONING_PARSER,
+                                       ReasoningParserFactory)
 from ..logger import logger
 from ..sampling_params import SamplingParams
 
@@ -66,6 +68,48 @@ def get_model_context(llm: Any) -> tuple[str, str]:
     return str(model_dir), str(model_type)
 
 
+def extract_final_content_from_generation(
+    output: RequestOutput,
+    *,
+    reasoning_parser: Optional[str] = None,
+) -> str:
+    """Return the scoreable final-answer text from a raw generation.
+
+    Handles plain JSON text, Harmony token transcripts, and normal reasoning
+    parser formats like `<think>...</think>{...}` for eval-only scoring.
+    """
+    # Eval scoring has historically used the first completion choice.
+    text = output.outputs[0].text
+
+    if reasoning_parser is None:
+        # Plain models keep the raw output; never guess a parser from text.
+        return text
+
+    if reasoning_parser.lower() == HARMONY_REASONING_PARSER:
+        # Harmony preserves final-channel boundaries in token ids, not text.
+        try:
+            from tensorrt_llm.serve.harmony_adapter import get_harmony_adapter
+
+            parsed = get_harmony_adapter().harmony_output_to_openai(
+                output.outputs[0].token_ids)
+            content = parsed.get("content")
+            if not parsed.get("_harmony_parsing_failed") and isinstance(
+                    content, str):
+                return content
+        except (ImportError, AttributeError, TypeError, RuntimeError,
+                ValueError):
+            return text
+        return text
+
+    try:
+        # Normal reasoning formats expose final content through their parser.
+        parser = ReasoningParserFactory.create_reasoning_parser(
+            reasoning_parser)
+        return parser.parse(text).content or text
+    except (AttributeError, TypeError, ValueError):
+        return text
+
+
 class Evaluator(ABC):
 
     def __init__(self,

tensorrt_llm/evaluate/json_mode_eval.py

@@ -14,7 +14,7 @@
 # limitations under the License.
 import json
 import os
-from typing import Iterable, List, Optional
+from typing import Iterable, List, Optional, Union
 
 import click
 import datasets
@@ -23,9 +23,31 @@
 
 from .. import LLM as PyTorchLLM
 from ..llmapi import RequestOutput
+from ..llmapi.reasoning_parser import resolve_guided_decoding_reasoning_parser
 from ..logger import logger
 from ..sampling_params import GuidedDecodingParams, SamplingParams
-from .interface import Evaluator
+from .interface import Evaluator, extract_final_content_from_generation
+
+JsonValue = Union[dict[str, object], list[object], str, int, float, bool, None]
+
+
+def _load_json_from_generation(
+        output: RequestOutput,
+        reasoning_parser: Optional[str] = None) -> JsonValue:
+    text = output.outputs[0].text
+    try:
+        # Plain/non-reasoning evals already return the constrained JSON as raw
+        # text, so keep the existing fast path and semantics unchanged.
+        return json.loads(text)
+    except json.JSONDecodeError as original_error:
+        # Reasoning models may return a raw transcript such as Harmony channels
+        # or `<think>...</think>{...}`. Score the final answer content when it
+        # can be extracted; otherwise preserve the original JSON failure.
+        final_content = extract_final_content_from_generation(
+            output, reasoning_parser=reasoning_parser)
+        if final_content == text:
+            raise original_error
+        return json.loads(final_content)
 
 
 class JsonModeEval(Evaluator):
@@ -44,6 +66,7 @@ def __init__(self,
                          apply_chat_template=apply_chat_template,
                          system_prompt=system_prompt,
                          output_dir=output_dir)
+        self._reasoning_parser: Optional[str] = None
         if dataset_path is None:
             dataset_path = "NousResearch/json-mode-eval"
         self.data = datasets.load_dataset(dataset_path,
@@ -55,6 +78,21 @@ def __init__(self,
         else:
             self.num_samples = min(num_samples, self.data.num_rows)
 
+    def evaluate(self,
+                 llm: PyTorchLLM,
+                 sampling_params: Optional[SamplingParams] = None,
+                 streaming: bool = False) -> float:
+        # Resolve once from the LLM rather than guessing Harmony from every
+        # failed JSON sample. Plain models therefore keep strict raw-text
+        # scoring, while GPT-OSS uses its known final-channel token framing.
+        model_type = getattr(getattr(llm, "_hf_model_config", None),
+                             "model_type", None)
+        self._reasoning_parser = resolve_guided_decoding_reasoning_parser(
+            getattr(getattr(llm, "args", None), "reasoning_parser", None),
+            model_type,
+        )
+        return super().evaluate(llm, sampling_params, streaming)
+
     def generate_samples(self) -> Iterable[tuple]:
         for i, sample in enumerate(self.data):
             if i >= self.num_samples:
@@ -75,7 +113,8 @@ def compute_score(self, outputs: List[RequestOutput], references: List[str],
         all_corrections, all_grammar_corrections = [], []
         for output, ref, schema in zip(outputs, references, schemas):
             try:
-                output_json = json.loads(output.outputs[0].text)
+                output_json = _load_json_from_generation(
+                    output, self._reasoning_parser)
                 jsonschema.validate(output_json, json.loads(schema))
             except (json.JSONDecodeError, jsonschema.ValidationError):
                 all_corrections.append(False)

tensorrt_llm/llmapi/llm.py

@@ -45,6 +45,9 @@
 from .llm_utils import (CachedModelLoader, KvCacheRetentionConfig,
                         LlmBuildStats, ModelLoader)
 from .mpi_session import MpiPoolSession, external_mpi_comm_available
+from .reasoning_parser import (
+    adapt_guided_decoding_params_for_reasoning_parser,
+    resolve_raw_guided_decoding_reasoning_parser)
 from .thinking_budget import add_thinking_budget_logits_processor
 from .tokenizer import TokenizerBase
 # TODO[chunweiy]: move the following symbols back to utils scope, and remove the following import
@@ -1390,6 +1393,24 @@ def _prepare_sampling_params(
                                        self._generation_config)
             self._configure_bart_decoder_prefix(sampling_params)
             self._add_whisper_suppress_tokens_logits_processor(sampling_params)
+            if sampling_params.guided_decoding is not None:
+                reasoning_format_for_guided_decoding = (
+                    resolve_raw_guided_decoding_reasoning_parser(
+                        self.args.reasoning_parser,
+                        getattr(self._hf_model_config, "model_type", None),
+                        self.args.guided_decoding_backend,
+                    ))
+                if reasoning_format_for_guided_decoding is not None:
+                    # SamplingParams carries a caller-provided content
+                    # constraint, but not how this model frames reasoning and
+                    # final output. Add that model-aware framing here so
+                    # xgrammar applies the guide only to Harmony's final
+                    # channel. Other formats preserve the original guide.
+                    sampling_params.guided_decoding = (
+                        adapt_guided_decoding_params_for_reasoning_parser(
+                            sampling_params.guided_decoding,
+                            reasoning_format_for_guided_decoding,
+                        ))
             add_thinking_budget_logits_processor(
                 sampling_params,
                 reasoning_parser=self.args.reasoning_parser,

tensorrt_llm/llmapi/reasoning_parser.py

@@ -20,7 +20,8 @@
 from pathlib import Path
 from typing import Any, ClassVar, Optional, Type
 
-from tensorrt_llm import logger
+from tensorrt_llm.logger import logger
+from tensorrt_llm.sampling_params import GuidedDecodingParams
 
 
 @dataclass
@@ -32,6 +33,15 @@ class ReasoningParserResult:
 # Enough of the rendered prompt's tail to hold a prefilled marker and any
 # trailing whitespace, without copying a prompt that may be very long.
 _PROMPT_TAIL_CHARS = 64
+HARMONY_REASONING_PARSER = "gpt_oss"
+HARMONY_FINAL_CHANNEL_TRIGGER = "<|start|>assistant<|channel|>final<|message|>"
+
+# Unlike normal reasoning parsers, Harmony is selected automatically by the
+# serving stack from the model type. Raw LLM requests need the same default
+# when they arrive with pre-built SamplingParams.
+_DEFAULT_GUIDED_DECODING_REASONING_PARSER_BY_MODEL_TYPE = {
+    "gpt_oss": HARMONY_REASONING_PARSER,
+}
 
 
 def register_reasoning_parser(*keys: str, **default_kwargs):
@@ -196,6 +206,127 @@ def parse_delta(self, delta_text: str) -> ReasoningParserResult:
         return ReasoningParserResult(content=delta_text)
 
 
+def resolve_guided_decoding_reasoning_parser(
+    reasoning_parser: Optional[str],
+    model_type: Optional[str],
+) -> Optional[str]:
+    """Resolve the reasoning format used to scope a request's guide.
+
+    An explicitly configured parser always wins. The model-type dispatch only
+    supplies formats that serving already selects implicitly, currently the
+    GPT-OSS Harmony protocol. Other raw LLM models therefore retain their
+    pre-existing guided-decoding behavior unless a parser was configured.
+    """
+    if reasoning_parser is not None:
+        return reasoning_parser
+    return _DEFAULT_GUIDED_DECODING_REASONING_PARSER_BY_MODEL_TYPE.get(
+        model_type)
+
+
+def resolve_raw_guided_decoding_reasoning_parser(
+    reasoning_parser: Optional[str],
+    model_type: Optional[str],
+    guided_decoding_backend: Optional[str],
+) -> Optional[str]:
+    """Resolve the reasoning format that raw LLM requests should adapt.
+
+    Raw LLM only needs new final-content scoping for Harmony, and structural
+    tags are supported by xgrammar only. Explicit normal parsers and
+    llguidance retain their pre-existing raw guided-decoding behavior.
+    """
+    resolved_parser = resolve_guided_decoding_reasoning_parser(
+        reasoning_parser, model_type)
+    if (resolved_parser is not None
+            and resolved_parser.lower() == HARMONY_REASONING_PARSER
+            and guided_decoding_backend == "xgrammar"):
+        return HARMONY_REASONING_PARSER
+    return None
+
+
+def _normalize_json_schema_for_structural_tag(json_schema: Any) -> Any:
+    """Convert supported schema representations to structural-tag JSON."""
+    if hasattr(json_schema, "model_json_schema"):
+        json_schema = json_schema.model_json_schema()
+    if isinstance(json_schema, str):
+        json_schema = json.loads(json_schema)
+    return json_schema
+
+
+def _guided_decoding_content(
+        guided_decoding_params: GuidedDecodingParams) -> Optional[dict]:
+    """Translate an ordinary guide into structural-tag content."""
+    if guided_decoding_params.json is not None:
+        json_schema = _normalize_json_schema_for_structural_tag(
+            guided_decoding_params.json)
+        return {"type": "json_schema", "json_schema": json_schema}
+    if guided_decoding_params.json_object:
+        return {"type": "json_schema", "json_schema": {"type": "object"}}
+    if guided_decoding_params.regex is not None:
+        return {"type": "regex", "pattern": guided_decoding_params.regex}
+    if guided_decoding_params.grammar is not None:
+        return {"type": "grammar", "grammar": guided_decoding_params.grammar}
+    return None
+
+
+def adapt_guided_decoding_params_for_reasoning_parser(
+    guided_decoding_params: Optional[GuidedDecodingParams],
+    reasoning_parser: Optional[str],
+) -> Optional[GuidedDecodingParams]:
+    """Scope a guide to final content while leaving reasoning unconstrained.
+
+    Normal reasoning formats use a reasoning-tag sequence. Harmony instead
+    activates the guide when the assistant's final channel begins. Existing
+    structural-tag guides are already fully specified and remain untouched.
+    """
+    if guided_decoding_params is None or reasoning_parser is None:
+        return guided_decoding_params
+    if guided_decoding_params.structural_tag is not None:
+        return guided_decoding_params
+
+    content = _guided_decoding_content(guided_decoding_params)
+    if content is None:
+        return guided_decoding_params
+
+    if reasoning_parser.lower() == HARMONY_REASONING_PARSER:
+        stag_format = {
+            "type":
+            "triggered_tags",
+            "triggers": [HARMONY_FINAL_CHANNEL_TRIGGER],
+            "tags": [{
+                "begin": HARMONY_FINAL_CHANNEL_TRIGGER,
+                "content": content,
+                "end": "",
+            }],
+            "stop_after_first":
+            True,
+        }
+    else:
+        parser = ReasoningParserFactory.create_reasoning_parser(
+            reasoning_parser)
+        stag_format = {
+            "type":
+            "sequence",
+            "elements": [
+                {
+                    "type": "tag",
+                    "begin": parser.reasoning_start,
+                    "content": {
+                        "type": "any_text"
+                    },
+                    "end": parser.reasoning_end,
+                },
+                content,
+            ],
+        }
+
+    structural_tag = {
+        "type": "structural_tag",
+        "format": stag_format,
+    }
+    return GuidedDecodingParams(
+        structural_tag=json.dumps(structural_tag, separators=(",", ":")))
+
+
 @register_reasoning_parser("deepseek-r1", reasoning_at_start=True)
 @register_reasoning_parser("qwen3")
 # Qwen3.5 (and forced-thinking Qwen3 variants) use a chat template that