From e347083a49a08e70878bef7daec97c4ea18625c6 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Tue, 27 Oct 2026 01:37:30 +0000 Subject: [PATCH 1/5] [None][fix] serve: fail a request whose generation degenerates into a token run A model can fall into repeating a token it never writes in valid output (its pad token, an unused placeholder) until its token budget runs out. The tokens are special, so the client receives an empty or truncated reply as a normal answer and keeps it in the conversation. A serving extension now declares such tokens (ServingExtension.degenerate_tokens): DeepSeek-V4.1 its pad token and its last vocabulary id, Kimi-K3 its [PAD]. The server adds a stop word of 20 copies of each id to every generation request, matched by token id after the request's own stop words. A generation that stops on one fails with a RequestError instead of reporting a stop, so none of its output is returned: - non-streamed requests get HTTP 500; - chat streams, also on the Harmony route, end with an error event and [DONE]; - Anthropic Messages streams end with an error event; - Responses streams end with error and response.failed, without output. Models that declare nothing are unchanged. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- tensorrt_llm/executor/result.py | 31 +- tensorrt_llm/executor/utils.py | 27 ++ tensorrt_llm/sampling_params.py | 40 +- tensorrt_llm/serve/anthropic_adapter.py | 27 +- tensorrt_llm/serve/extensions/__init__.py | 4 +- tensorrt_llm/serve/extensions/deepseek_v41.py | 49 +++ tensorrt_llm/serve/extensions/kimi_k3.py | 20 +- tensorrt_llm/serve/openai_server.py | 105 ++++- tensorrt_llm/serve/serving_extensions.py | 57 ++- .../integration/test_lists/test-db/l0_cpu.yml | 2 + .../executor/test_degenerate_token_runs.py | 241 ++++++++++ .../apps/test_degenerate_generation_errors.py | 411 ++++++++++++++++++ .../llmapi/apps/test_serving_extensions.py | 93 ++++ 13 files changed, 1075 insertions(+), 32 deletions(-) create mode 100644 tensorrt_llm/serve/extensions/deepseek_v41.py create mode 100644 tests/unittest/executor/test_degenerate_token_runs.py create mode 100644 tests/unittest/llmapi/apps/test_degenerate_generation_errors.py diff --git a/tensorrt_llm/executor/result.py b/tensorrt_llm/executor/result.py index 145a0275dd0a..a285d50b34f9 100644 --- a/tensorrt_llm/executor/result.py +++ b/tensorrt_llm/executor/result.py @@ -10,7 +10,7 @@ from dataclasses import dataclass, field from queue import Empty, Queue from typing import (TYPE_CHECKING, Any, Callable, Dict, List, Literal, - NamedTuple, Optional, TypeAlias, Union) + NamedTuple, NoReturn, Optional, TypeAlias, Union) from weakref import WeakMethod import torch @@ -27,9 +27,10 @@ from ..metrics import MetricNames, MetricsCollector, RequestEventTiming from ..metrics.perf_utils import \ process_req_perf_metrics as _process_req_perf_metrics -from ..sampling_params import LogprobParams, SamplingParams +from ..sampling_params import DegenerateTokenRun, LogprobParams, SamplingParams from .postprocessor_hook import PostProcessorHook, apply_post_processor_hook -from .utils import (EngineDeadError, ErrorResponse, has_event_loop, +from .utils import (EngineDeadError, ErrorResponse, RequestError, + degenerate_generation_message, has_event_loop, is_llm_response) if TYPE_CHECKING: @@ -494,6 +495,9 @@ def _handle_sequence(self, for stop_reason, stop_ids in self.sampling_params._get_stop_reasons_and_words( ): if output.token_ids[-len(stop_ids):] == stop_ids: + if isinstance(stop_reason, DegenerateTokenRun): + self._fail_degenerate_generation( + stop_reason, output) output.stop_reason = stop_reason if not self.sampling_params.include_stop_str_in_output: self._trim_stop_word_outputs(output, len(stop_ids)) @@ -522,6 +526,27 @@ def _handle_sequence(self, if self._done: self.do_tracing(output, req_perf_metrics_dict) + def _fail_degenerate_generation(self, run: DegenerateTokenRun, + output: CompletionOutput) -> NoReturn: + """Fail a request whose generation stopped on a degenerate token run. + + The run means the generation has broken down, so nothing of it is a + reply: the request fails without output, and every later wait on this + result raises the same error. A client that kept the output instead + would carry a truncated reply into every later turn. + """ + message = degenerate_generation_message(run.token_id, run.length) + logger.warning(f"Request {self.id}: {message} " + f"({len(output.token_ids)} tokens generated)") + if not self._done: + # Other sequences of the request are still generating, for output + # that is discarded with this one. + self.abort() + error = RequestError(message) + self._error_msg = message + self._terminal_error = error + raise error + @staticmethod def _trim_stop_word_outputs(output: CompletionOutput, num_stop_ids: int) -> None: diff --git a/tensorrt_llm/executor/utils.py b/tensorrt_llm/executor/utils.py index f488a03682dc..f0c682cb5965 100644 --- a/tensorrt_llm/executor/utils.py +++ b/tensorrt_llm/executor/utils.py @@ -169,6 +169,33 @@ def is_context_length_exceeded_message(message: str) -> bool: return bool(_CONTEXT_LENGTH_EXCEEDED_RE.search(message)) +def degenerate_generation_message(token_id: int, run_length: int) -> str: + """Client-facing message for a generation failed on a degenerate token run. + + See ``SamplingParams._set_degenerate_token_runs``. Keep in sync with + ``is_degenerate_generation_message``. + """ + return (f"Generation failed: the output degenerated into a run of token id " + f"{token_id} ({run_length} in a row), which the model does not " + "write in valid output. The output was discarded; retry the " + "request.") + + +_DEGENERATE_GENERATION_RE = re.compile( + r"Generation failed: the output degenerated into a run of token id \d+ " + r"\(\d+ in a row\)") + + +def is_degenerate_generation_message(message: str) -> bool: + """True if ``message`` contains a ``degenerate_generation_message``. + + Substring search for the same reason as + ``is_context_length_exceeded_message``: the message reaches the server as + plain text, prefixed when it comes from a postprocessing worker. + """ + return bool(_DEGENERATE_GENERATION_RE.search(message)) + + class EngineDeadError(RuntimeError): """Raised by pending and new requests once the engine is known dead. diff --git a/tensorrt_llm/sampling_params.py b/tensorrt_llm/sampling_params.py index c133f71946f4..79b7ceef49d8 100644 --- a/tensorrt_llm/sampling_params.py +++ b/tensorrt_llm/sampling_params.py @@ -101,6 +101,18 @@ class LogprobParams(NamedTuple): prompt_logprobs_simple_format: bool = False +class DegenerateTokenRun(NamedTuple): + """``length`` consecutive copies of ``token_id`` that end a generation as failed. + + The stop reason ``SamplingParams._get_stop_reasons_and_words`` reports for a + run set with ``SamplingParams._set_degenerate_token_runs``. A generation that + stops on one has broken down, so it is reported as an error, never as a stop. + """ + + token_id: int + length: int + + class LogprobMode(StrEnum): RAW = "raw" """ @@ -288,6 +300,8 @@ class SamplingParams: stop_token_ids: Optional[List[int]] = None include_stop_str_in_output: bool = False _stop_word_ids: Optional[List[List[int]]] = field(default=None, init=False, repr=False) + # (token id, run length) pairs set by _set_degenerate_token_runs. + _degenerate_token_runs: Tuple[Tuple[int, int], ...] = field(default=(), init=False, repr=False) embedding_bias: Optional[torch.Tensor] = None logits_processor: Optional[Union[LogitsProcessor, List[LogitsProcessor]]] = None @@ -628,22 +642,37 @@ def _get_bad_words(self) -> List[List[int]]: ) return words + self._bad_word_ids + def _set_degenerate_token_runs(self, token_ids: Iterable[int], run_length: int) -> None: + """Fail the generation once it writes one of ``token_ids`` ``run_length`` times in a row. + + For servers of models that declare such tokens. Each run is an extra + stop word matched by token id, after the request's own stop words, and + reported as a ``DegenerateTokenRun`` stop reason. + """ + if run_length < 1: + raise ValueError(f"run_length must be at least 1, got {run_length}") + self._degenerate_token_runs = tuple( + (token_id, run_length) for token_id in dict.fromkeys(token_ids) + ) + def _get_stop_words(self) -> List[List[int]]: words = [] if self.stop_token_ids: words = [[i] for i in self.stop_token_ids] - if self.stop is None: - return words - else: + if self.stop is not None: if self._stop_word_ids is None: raise RuntimeError( f"{self.__class__.__name__}.stop ({self.stop}) is not processed by tokenizer, " "please call the setup method." ) - return words + self._stop_word_ids + words = words + self._stop_word_ids + # Last, so that a stop word of the request wins where both match. + return words + [[token_id] * length for token_id, length in self._degenerate_token_runs] - def _get_stop_reasons_and_words(self) -> List[Tuple[Union[str, int], List[List[int]]]]: + def _get_stop_reasons_and_words( + self, + ) -> List[Tuple[Union[str, int, DegenerateTokenRun], List[List[int]]]]: stop_reasons = [] if self.stop_token_ids is not None: stop_reasons.extend(self.stop_token_ids) @@ -652,6 +681,7 @@ def _get_stop_reasons_and_words(self) -> List[Tuple[Union[str, int], List[List[i stop_reasons.append(self.stop) else: stop_reasons.extend(self.stop) + stop_reasons.extend(DegenerateTokenRun(*run) for run in self._degenerate_token_runs) stop_words = self._get_stop_words() return list(zip(stop_reasons, stop_words)) diff --git a/tensorrt_llm/serve/anthropic_adapter.py b/tensorrt_llm/serve/anthropic_adapter.py index 489f4d1e0a32..6729a2c03639 100644 --- a/tensorrt_llm/serve/anthropic_adapter.py +++ b/tensorrt_llm/serve/anthropic_adapter.py @@ -1073,6 +1073,15 @@ async def _iter_openai_sse_lines( yield buffer.rstrip("\r") +def _openai_stream_error_message(payload: Any) -> Optional[str]: + """The message of an OpenAI in-stream error frame, or None for any other payload.""" + if not isinstance(payload, dict) or "error" not in payload: + return None + error = payload["error"] + message = error.get("message") if isinstance(error, dict) else None + return message if isinstance(message, str) and message else "Internal server error" + + async def reframe_openai_stream( openai_sse: AsyncIterator[Union[str, bytes]], model: str ) -> AsyncIterator[str]: @@ -1103,8 +1112,17 @@ async def reframe_openai_stream( done = True continue try: - chunk = ChatCompletionStreamResponse(**json.loads(data)) - except (json.JSONDecodeError, ValueError) as e: + payload = json.loads(data) + except json.JSONDecodeError as e: + raise AnthropicResponseError("Malformed upstream stream chunk") from e + upstream_error = _openai_stream_error_message(payload) + if upstream_error is not None: + # The request failed after its stream began: end the message + # with an error rather than as a finished turn. + raise AnthropicResponseError(upstream_error) + try: + chunk = ChatCompletionStreamResponse(**payload) + except ValueError as e: raise AnthropicResponseError("Malformed upstream stream chunk") from e for frame in reframer.process_chunk(chunk): yield frame @@ -1122,8 +1140,9 @@ async def reframe_openai_stream( ) except AnthropicResponseError as e: # Raised only on upstream framing this module can describe precisely, - # so the text originates here and carries no internal detail: pass it - # through rather than flattening it to a generic message. + # or on an upstream error frame, whose message the chat route already + # shows its own clients. Either way the text carries no internal + # detail: pass it through rather than flattening it to a generic message. logger.error(f"Anthropic stream failed: {e}") for frame in reframer.error(str(e)): yield frame diff --git a/tensorrt_llm/serve/extensions/__init__.py b/tensorrt_llm/serve/extensions/__init__.py index fcfb9988d2b4..0a6a97160ef8 100644 --- a/tensorrt_llm/serve/extensions/__init__.py +++ b/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, nemotron +from . import deepseek_v41, gpt_oss, kimi_k3, nemotron -__all__ = ["gpt_oss", "kimi_k3", "nemotron"] +__all__ = ["deepseek_v41", "gpt_oss", "kimi_k3", "nemotron"] diff --git a/tensorrt_llm/serve/extensions/deepseek_v41.py b/tensorrt_llm/serve/extensions/deepseek_v41.py new file mode 100644 index 000000000000..b97b68e3f173 --- /dev/null +++ b/tensorrt_llm/serve/extensions/deepseek_v41.py @@ -0,0 +1,49 @@ +# 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 DeepSeek-V4.1 checkpoints (``model_type == "deepseek_v41"``). + +Prompts are rendered by the ``deepseek_v41`` custom tokenizer, so the only +serving behavior declared here is which token runs mark a generation that has +broken down. +""" + +from typing import Any, Optional + +from tensorrt_llm.serve.serving_extensions import ( + DegenerateTokens, + ServingExtension, + last_vocabulary_id, + register_serving_extension, +) + + +@register_serving_extension(model_types=("deepseek_v41", "deepseek_v41_text")) +class DeepseekV41ServingExtension(ServingExtension): + """DeepSeek-V4.1 serving behavior.""" + + def degenerate_tokens(self, hf_config: Any) -> Optional[DegenerateTokens]: + """``<｜▁pad▁｜>`` and the last vocabulary id, ``<|place_holder_mm_span_0442|>``. + + Neither belongs in generated text: the pad token fills batches, and the + last vocabulary id is a multimodal placeholder that no prompt uses. The + pad id comes from the model config, not the tokenizer, whose pad token + is ``<｜end▁of▁sentence｜>``. + """ + pad_token_id = getattr(hf_config, "pad_token_id", None) + if isinstance(pad_token_id, bool) or not isinstance(pad_token_id, int) or pad_token_id < 0: + pad_token_id = None + candidates = (pad_token_id, last_vocabulary_id(hf_config)) + token_ids = tuple(dict.fromkeys(i for i in candidates if i is not None)) + return DegenerateTokens(token_ids=token_ids) if token_ids else None diff --git a/tensorrt_llm/serve/extensions/kimi_k3.py b/tensorrt_llm/serve/extensions/kimi_k3.py index db24c04ea604..daff27051a49 100644 --- a/tensorrt_llm/serve/extensions/kimi_k3.py +++ b/tensorrt_llm/serve/extensions/kimi_k3.py @@ -17,7 +17,8 @@ Implements the Kimi/Moonshot API semantics layered on the OpenAI chat surface: message-level (dynamic) tool validation, the optional pinned sampling-parameter policy, default streaming usage, request-field to chat-template-kwarg -derivation, and the K3 XTML placement of guided-decoding constraints. +derivation, and the K3 XTML placement of guided-decoding constraints. It also +declares ``[PAD]`` as the token whose run marks a broken-down generation. """ import os @@ -29,7 +30,12 @@ ChatCompletionRequest, StreamOptions, ) -from tensorrt_llm.serve.serving_extensions import ServingExtension, register_serving_extension +from tensorrt_llm.serve.serving_extensions import ( + DegenerateTokens, + ServingExtension, + last_vocabulary_id, + register_serving_extension, +) # Valid function-tool name for message-level tools: no leading digit, word # chars/dash only, at most 256 chars. @@ -267,3 +273,13 @@ def structured_output_format( ], "stop_after_first": True, } + + def degenerate_tokens(self, hf_config: Any) -> Optional[DegenerateTokens]: + """``[PAD]``, the last vocabulary id, which K3 never writes in valid output. + + Taken from the vocabulary size rather than ``pad_token_id``, which the + composite config defaults to 0 when the checkpoint leaves it out. The + text "[PAD]" in code the model writes is ordinary tokens, not this id. + """ + pad_token_id = last_vocabulary_id(hf_config) + return None if pad_token_id is None else DegenerateTokens(token_ids=(pad_token_id,)) diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 8c4ba0f9315a..19e19b64a2e1 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -45,7 +45,9 @@ from tensorrt_llm.executor.request import DEFAULT_REQUEST_PRIORITY from tensorrt_llm.executor.rpc import RPCError from tensorrt_llm.executor.utils import (CONTEXT_LENGTH_EXCEEDED_CODE, - is_context_length_exceeded_message) + RequestError, + is_context_length_exceeded_message, + is_degenerate_generation_message) from tensorrt_llm.inputs import prompt_inputs from tensorrt_llm.inputs.data import TokensPrompt from tensorrt_llm.inputs.media_io import BaseMediaIO @@ -82,7 +84,8 @@ AnthropicMessagesRequest) from tensorrt_llm.serve.chat_tokenization import tokenize_harmony_chat_request from tensorrt_llm.serve.chat_utils import (load_chat_template, - parse_chat_messages_coroutines) + parse_chat_messages_coroutines, + resolve_top_level_model_type) from tensorrt_llm.serve.cluster_storage import create_cluster_storage_client from tensorrt_llm.serve.conversation_id import resolve_request_conversation_id from tensorrt_llm.serve.disagg_auth import ( @@ -134,7 +137,9 @@ from tensorrt_llm.serve.rl_control_auth import validate_rl_control_request from tensorrt_llm.serve.runtime_control_auth import ( RuntimeControlReplayCache, validate_runtime_control_request) -from tensorrt_llm.serve.serving_extensions import load_builtin_extensions +from tensorrt_llm.serve.serving_extensions import (DegenerateTokens, + get_serving_extension, + load_builtin_extensions) from tensorrt_llm.serve.tool_parser.tool_parser_factory import ToolParserFactory from tensorrt_llm.serve.visual_gen_metrics import ( build_visual_gen_server_timings, build_visual_gen_timing_headers) @@ -313,6 +318,28 @@ def _record_generator_termination(generator) -> None: )) +def _sse_error_frame(message: str) -> str: + """An in-stream error event; OpenAI clients raise it as an API error.""" + error_data = json.dumps({ + "error": { + "message": message, + "type": "server_error", + "code": None, + "param": None, + } + }) + return f"data: {error_data}\n\n" + + +def _degenerate_stream_end(error: RequestError) -> List[str]: + """The frames that end a chat stream whose generation broke down. + + The response has started, so the failure is reported in-stream, and the + client drops what it has received instead of keeping a truncated reply. + """ + return [_sse_error_frame(str(error)), "data: [DONE]\n\n"] + + def _build_tool_strict_guided_decoding_params(tools, tool_parser_name): """Build GuidedDecodingParams with structural tags for tools with strict=True. @@ -1227,14 +1254,51 @@ def _vocab_size(self) -> Optional[int]: return self.tokenizer.tokenizer.vocab_size return None + @functools.cached_property + def _degenerate_tokens(self) -> Optional[DegenerateTokens]: + """The degenerate tokens the served model's extension declares, if any. + + Computed on first use, so servers built without ``__init__`` (unit + tests) read the ``model_config`` they were given. + """ + hf_config = getattr(self, "model_config", None) + if hf_config is None: + return None + degenerate = get_serving_extension( + resolve_top_level_model_type(hf_config)).degenerate_tokens( + hf_config) + if degenerate is not None: + logger.info( + f"A generation that writes one of token ids " + f"{list(degenerate.token_ids)} {degenerate.run_length} times in " + "a row fails without output.") + return degenerate + + def _stop_on_degenerate_tokens(self, + sampling_params: SamplingParams) -> None: + """Fail the request on a run of the model's degenerate tokens. + + A no-op for models that declare none (see + ``ServingExtension.degenerate_tokens``). + """ + degenerate = self._degenerate_tokens + if degenerate is not None: + sampling_params._set_degenerate_token_runs(degenerate.token_ids, + degenerate.run_length) + @staticmethod def create_error_response( message: str, err_type: str = "BadRequestError", status_code: HTTPStatus = HTTPStatus.BAD_REQUEST) -> Response: - # Over-length prompts are detected from the message text because they - # cross the executor IPC boundary as plain strings (RequestError), so - # no exception type survives to here. + # Over-length prompts and degenerate generations are detected from the + # message text because they cross the executor IPC boundary as plain + # strings (RequestError), so no exception type survives to here. + if is_degenerate_generation_message(message): + # The server's failure, not the request's: a 5xx, which clients + # retry, whatever the caller's default. + err_type = "InternalServerError" + status_code = HTTPStatus.INTERNAL_SERVER_ERROR code: Union[int, str] = status_code.value if is_context_length_exceeded_message(message): code = CONTEXT_LENGTH_EXCEEDED_CODE @@ -2087,6 +2151,12 @@ async def chat_stream_generator( yield "data: [DONE]\n\n" await self._extract_metrics(res, raw_request) nvtx_mark("generation ends") + except RequestError as e: + if not is_degenerate_generation_message(str(e)): + logger.error(traceback.format_exc()) + raise + for frame in _degenerate_stream_end(e): + yield frame except Exception: # Only failures are logged: a consumer that stops reading closes # the stream with GeneratorExit, and a client that disconnects @@ -2116,6 +2186,7 @@ async def chat_stream_generator( gather_generation_logits, reasoning_parser=self.generator.args.reasoning_parser, backend=self.generator.args.backend) + self._stop_on_degenerate_tokens(sampling_params) if required_tool_grammar: _validate_required_tool_grammar_request( request, @@ -3006,15 +3077,7 @@ async def completion_generator(promise: RequestOutput, # chunk, so we cannot change the status code. Yield # an SSE error event so the stream terminates cleanly # instead of breaking the HTTP connection. - error_data = json.dumps({ - "error": { - "message": str(e), - "type": "server_error", - "code": None, - "param": None, - } - }) - yield f"data: {error_data}\n\n" + yield _sse_error_frame(str(e)) yield "data: [DONE]\n\n" async def merge_generators(generators: List[AsyncIterator[Any]]): @@ -3069,6 +3132,7 @@ async def generator_wrapper(generator: AsyncIterator[Any]): gather_generation_logits=self.generator.args. gather_generation_logits, backend=self.generator.args.backend) + self._stop_on_degenerate_tokens(sampling_params) add_thinking_budget_logits_processor( sampling_params, reasoning_parser=self.generator.args.reasoning_parser, @@ -3192,6 +3256,12 @@ async def create_streaming_generator(promise: RequestOutput, await self._extract_metrics(res, raw_request) except asyncio.CancelledError: raise + except RequestError as e: + if not is_degenerate_generation_message(str(e)): + logger.error(traceback.format_exc()) + raise + for frame in _degenerate_stream_end(e): + yield frame except Exception: logger.error(traceback.format_exc()) raise @@ -3243,6 +3313,7 @@ async def create_streaming_generator(promise: RequestOutput, sampling_params = request.to_sampling_params( vocab_size=self._logit_bias_vocab_size(), reasoning_parser="gpt_oss") + self._stop_on_degenerate_tokens(sampling_params) if sampling_params.thinking_token_budget is not None: raise ValueError( "thinking_token_budget is not supported by the Harmony " @@ -3462,6 +3533,10 @@ async def create_streaming_generator(promise: RequestOutput, chat_template=self.chat_template if not self.use_harmony else None, ) + # A run fails the request before any item it would have closed: + # the stream ends in response.failed (guard_responses_stream), a + # non-streamed request in an error response. + self._stop_on_degenerate_tokens(sampling_params) streaming_processor = None if request.stream: diff --git a/tensorrt_llm/serve/serving_extensions.py b/tensorrt_llm/serve/serving_extensions.py index ef6a129e73db..e5907a37914e 100644 --- a/tensorrt_llm/serve/serving_extensions.py +++ b/tensorrt_llm/serve/serving_extensions.py @@ -34,7 +34,8 @@ import enum import importlib -from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Literal, Optional +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Literal, Optional, Tuple if TYPE_CHECKING: from tensorrt_llm.serve.openai_protocol import ChatCompletionRequest @@ -42,6 +43,11 @@ RequiredToolChoiceMode = Literal["template", "grammar"] +# Copies of one degenerate token in a row that end a generation: far more than +# any valid output repeats a token it never writes, and few enough decoding steps +# to waste before the request fails. +DEGENERATE_TOKEN_RUN_LENGTH = 20 + class OutputMode(enum.Enum): """How a model's generated output is turned into an API response.""" @@ -52,6 +58,47 @@ class OutputMode(enum.Enum): HARMONY_TOKENS = "harmony_tokens" +@dataclass(frozen=True) +class DegenerateTokens: + """Token ids whose run marks a generation that has broken down. + + A model can fall into repeating a token it never writes in valid output (its + pad token, an unused placeholder) until its token budget runs out. Such tokens + are special, so they vanish from the detokenized text and the client receives + an empty or truncated reply as if it were complete. ``run_length`` consecutive + copies of one of ``token_ids`` end the generation, and the request fails + without output instead. Matching is by token id: text that merely spells such + a token is ordinary tokens and never matches. + """ + + token_ids: Tuple[int, ...] + run_length: int = DEGENERATE_TOKEN_RUN_LENGTH + + def __post_init__(self) -> None: + if not self.token_ids: + raise ValueError("DegenerateTokens needs at least one token id") + if any( + isinstance(token_id, bool) or not isinstance(token_id, int) or token_id < 0 + for token_id in self.token_ids + ): + raise ValueError(f"Token ids must be non-negative integers, got {self.token_ids}") + if self.run_length < 1: + raise ValueError(f"run_length must be at least 1, got {self.run_length}") + + +def last_vocabulary_id(hf_config: Any) -> Optional[int]: + """The highest token id of the model's vocabulary, or None if not configured. + + Read from the text configuration of a composite (multimodal) config. + """ + get_text_config = getattr(hf_config, "get_text_config", None) + text_config = get_text_config() if callable(get_text_config) else hf_config + vocab_size = getattr(text_config, "vocab_size", None) + if isinstance(vocab_size, bool) or not isinstance(vocab_size, int) or vocab_size < 1: + return None + return vocab_size - 1 + + class ServingExtension: """Per-model serving behavior consulted by the OpenAI-compatible server. @@ -122,6 +169,14 @@ def structured_output_format( """ return None + def degenerate_tokens(self, hf_config: Any) -> Optional[DegenerateTokens]: + """Tokens whose run marks a broken-down generation; see :class:`DegenerateTokens`. + + ``hf_config`` is the checkpoint's loaded configuration. ``None``, the + default, leaves the stop conditions of every request unchanged. + """ + return None + _BY_MODEL_TYPE: Dict[str, ServingExtension] = {} _BY_REASONING_PARSER: Dict[str, ServingExtension] = {} diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index ba52d1078918..aa8a20ac4a00 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -185,6 +185,7 @@ l0_cpu: - unittest/executor/test_stats_serializer.py - unittest/executor/test_spec_dec_perf_metrics.py - unittest/executor/test_stop_word_side_output_trim.py + - unittest/executor/test_degenerate_token_runs.py - unittest/executor/test_ray_stub.py - unittest/executor/test_generation_request_array_prompt.py - unittest/grpc/openengine/test_coordinator.py @@ -209,6 +210,7 @@ l0_cpu: - unittest/llmapi/apps/test_chat_template_kwargs_guard.py - unittest/llmapi/apps/test_chat_utils.py - unittest/llmapi/apps/test_context_length_error.py + - unittest/llmapi/apps/test_degenerate_generation_errors.py - unittest/llmapi/apps/test_harmony_channel_validation.py - unittest/llmapi/apps/test_harmony_parsing.py::TestStripIncompleteMessagesReporting - unittest/llmapi/apps/test_kimi_serve_extensions.py diff --git a/tests/unittest/executor/test_degenerate_token_runs.py b/tests/unittest/executor/test_degenerate_token_runs.py new file mode 100644 index 000000000000..e752db1c171c --- /dev/null +++ b/tests/unittest/executor/test_degenerate_token_runs.py @@ -0,0 +1,241 @@ +# 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. +"""A run of a degenerate token stops the generation and fails the request. + +A server sets the runs with ``SamplingParams._set_degenerate_token_runs`` for a +model that declares degenerate tokens. Each run is a stop word of token ids, so +the engine stops on it as on any stop word, and the result layer then fails the +request instead of reporting a stop. CPU-only: the engine's stop-word rule is +exercised through its host-side reference, ``meet_stop_token_criteria``. +""" + +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch.pyexecutor.sampler.sampler_features import meet_stop_token_criteria +from tensorrt_llm.bindings import executor as tllm +from tensorrt_llm.executor.result import GenerationResultBase +from tensorrt_llm.executor.utils import ( + RequestError, + degenerate_generation_message, + is_context_length_exceeded_message, + is_degenerate_generation_message, +) +from tensorrt_llm.sampling_params import DegenerateTokenRun, SamplingParams + +pytestmark = pytest.mark.cpu_only + +END_ID = 1 +PAD = 2 +PLACEHOLDER = 129279 +# Servers use 20; a short run keeps the sequences readable. +RUN = 4 +PROMPT = (5, 6, 7) + + +def _params(**kwargs) -> SamplingParams: + params = SamplingParams(end_id=END_ID, **kwargs) + params._set_degenerate_token_runs([PAD, PLACEHOLDER], RUN) + return params + + +class _Request: + """The request state ``meet_stop_token_criteria`` reads.""" + + def __init__(self, params: SamplingParams): + self.py_stop_words_list = params._get_stop_words() + self.tokens = list(PROMPT) + + def get_tokens(self, beam_idx: int) -> list: + return self.tokens + + +def _engine_stop(params: SamplingParams, generated: list) -> int | None: + """Generate ``generated`` token by token; how many were written when the engine stopped.""" + request = _Request(params) + for count, token in enumerate(generated, start=1): + request.tokens.append(token) + if meet_stop_token_criteria(request, token, 0): + return count + return None + + +def _finish(result: GenerationResultBase, generated: list, finish_reason) -> None: + """Hand the result layer the engine's final response for ``generated``.""" + response = SimpleNamespace( + output_token_ids=[list(generated)], + cum_log_probs=None, + log_probs=None, + generation_logits=None, + per_pos_drafted=None, + per_pos_accepted=None, + spec_dec_totals=None, + time_breakdown_metrics=None, + request_perf_metrics=None, + additional_context_outputs=None, + additional_generation_outputs=None, + ) + result._handle_sequence([finish_reason], response, 0) + + +def _finished_result(params: SamplingParams) -> GenerationResultBase: + result = GenerationResultBase(id=11, sampling_params=params) + result._done = True + return result + + +# -------------------------------------------------------------------------- +# The runs are stop words, matched by id +# -------------------------------------------------------------------------- + + +def test_runs_are_stop_words_after_the_requests_own(): + params = _params(stop_token_ids=[7]) + assert params._get_stop_words() == [[7], [PAD] * RUN, [PLACEHOLDER] * RUN] + assert params._get_stop_reasons_and_words() == [ + (7, [7]), + (DegenerateTokenRun(PAD, RUN), [PAD] * RUN), + (DegenerateTokenRun(PLACEHOLDER, RUN), [PLACEHOLDER] * RUN), + ] + + +def test_without_runs_the_stop_words_are_unchanged(): + params = SamplingParams(end_id=END_ID, stop_token_ids=[7]) + assert params._get_stop_words() == [[7]] + assert params._get_stop_reasons_and_words() == [(7, [7])] + assert SamplingParams(end_id=END_ID)._get_stop_words() == [] + + +def test_runs_stay_ids_when_stop_strings_are_tokenized(): + """A stop string spelling a degenerate token is encoded as text; the run is not.""" + tokenizer = SimpleNamespace(encode=lambda text, add_special_tokens=False: [91, 47, 93]) + params = _params(stop="") + params._setup(tokenizer, None, None) + assert params._get_stop_words() == [[91, 47, 93], [PAD] * RUN, [PLACEHOLDER] * RUN] + assert [reason for reason, _ in params._get_stop_reasons_and_words()] == [ + "", + DegenerateTokenRun(PAD, RUN), + DegenerateTokenRun(PLACEHOLDER, RUN), + ] + + +def test_duplicate_ids_give_one_run_each(): + params = SamplingParams(end_id=END_ID) + params._set_degenerate_token_runs([PAD, PAD, PLACEHOLDER], RUN) + assert params._get_stop_words() == [[PAD] * RUN, [PLACEHOLDER] * RUN] + + +def test_a_run_must_be_at_least_one_token(): + with pytest.raises(ValueError, match="run_length"): + SamplingParams(end_id=END_ID)._set_degenerate_token_runs([PAD], 0) + + +@pytest.mark.parametrize("token", [PAD, PLACEHOLDER]) +def test_the_engine_stops_at_the_threshold(token): + assert _engine_stop(_params(), [10, 11] + [token] * (RUN + 5)) == 2 + RUN + + +def test_a_shorter_run_does_not_stop_the_engine(): + generated = [10] + [PAD] * (RUN - 1) + [11] + [PLACEHOLDER] * (RUN - 1) + [12] + assert _engine_stop(_params(), generated) is None + + +def test_text_spelling_a_degenerate_token_does_not_stop_the_engine(): + """Ordinary tokens that read like the pad token are not the pad token.""" + assert _engine_stop(_params(), [91, 47, 93] * (RUN + 1)) is None + + +def test_without_runs_the_engine_does_not_stop(): + assert _engine_stop(SamplingParams(end_id=END_ID), [PAD] * (RUN * 5)) is None + + +# -------------------------------------------------------------------------- +# The result layer fails a generation that stopped on a run +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize("token", [PAD, PLACEHOLDER]) +def test_a_generation_stopped_on_a_run_fails(token): + result = _finished_result(_params()) + + with pytest.raises(RequestError) as excinfo: + _finish(result, [10, 11] + [token] * RUN, tllm.FinishReason.STOP_WORDS) + + message = degenerate_generation_message(token, RUN) + assert str(excinfo.value) == message + # The failure sticks: the result reports it and every later wait raises it. + assert result.error == message + assert result._terminal_error is excinfo.value + # It is the request's failure, so nothing stopped normally. + assert result._outputs[0].stop_reason is None + + +def test_the_requests_own_stop_token_wins_over_a_run(): + params = SamplingParams(end_id=END_ID, stop_token_ids=[PAD]) + params._set_degenerate_token_runs([PAD], 1) + result = _finished_result(params) + + _finish(result, [10, PAD], tllm.FinishReason.STOP_WORDS) + + output = result._outputs[0] + assert (output.finish_reason, output.stop_reason, output.token_ids) == ("stop", PAD, [10]) + assert result.error is None + + +def test_a_shorter_run_completes_normally(): + result = _finished_result(_params()) + generated = [10] + [PAD] * (RUN - 1) + [11] + + _finish(result, generated, tllm.FinishReason.END_ID) + + output = result._outputs[0] + assert (output.finish_reason, output.token_ids) == ("stop", generated) + assert result.error is None + + +def test_a_failed_sequence_aborts_the_rest_of_the_request(): + """With several sequences, the others are cancelled with the failed one.""" + result = GenerationResultBase(id=11, sampling_params=_params()) + result._done = False + + with pytest.raises(RequestError): + _finish(result, [10] + [PAD] * RUN, tllm.FinishReason.STOP_WORDS) + + assert result.aborted() + + +# -------------------------------------------------------------------------- +# The message the server recognizes +# -------------------------------------------------------------------------- + + +def test_message_names_the_token_and_the_run(): + message = degenerate_generation_message(PLACEHOLDER, 20) + assert f"token id {PLACEHOLDER} (20 in a row)" in message + assert "retry" in message + + +def test_detector_matches_exact_and_wrapped_messages(): + message = degenerate_generation_message(PAD, 20) + assert is_degenerate_generation_message(message) + # A postprocessing worker forwards it with a prefix. + assert is_degenerate_generation_message(f"Postprocessing error: {message}") + assert not is_context_length_exceeded_message(message) + + +def test_detector_rejects_other_errors(): + assert not is_degenerate_generation_message("Generation failed: request cancelled") + assert not is_degenerate_generation_message("") diff --git a/tests/unittest/llmapi/apps/test_degenerate_generation_errors.py b/tests/unittest/llmapi/apps/test_degenerate_generation_errors.py new file mode 100644 index 000000000000..8484b435a712 --- /dev/null +++ b/tests/unittest/llmapi/apps/test_degenerate_generation_errors.py @@ -0,0 +1,411 @@ +# 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. +"""A degenerate generation reaches the client of every API as an error. + +A model whose serving extension declares degenerate tokens gets their runs as +stop words on every generation request, and a generation that stops on one +fails with a RequestError carrying ``degenerate_generation_message`` (see +test_degenerate_token_runs.py). These tests drive the real route handlers on a +server built with ``object.__new__``: the engine is a stub whose result fails +the way the result layer fails it, in the server process or, prefixed, in a +postprocessing worker. CPU-only. +""" + +import json +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import openai +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from tensorrt_llm.executor.utils import RequestError, degenerate_generation_message +from tensorrt_llm.serve import openai_server as openai_server_module +from tensorrt_llm.serve import serving_extensions +from tensorrt_llm.serve.openai_protocol import ResponsesRequest +from tensorrt_llm.serve.openai_server import OpenAIServer +from tensorrt_llm.serve.render import chat as render_chat_module +from tensorrt_llm.serve.serving_extensions import ( + DEGENERATE_TOKEN_RUN_LENGTH, + DegenerateTokens, + ServingExtension, + register_serving_extension, +) + +pytestmark = pytest.mark.cpu_only + +MODEL = "test-model" +DECLARED_MODEL_TYPE = "degenerate-errors-test-model" +PAD = 2 +LAST_ID = 129279 +RUN = DEGENERATE_TOKEN_RUN_LENGTH +MESSAGE = degenerate_generation_message(LAST_ID, RUN) +# How the failure arrives when a postprocessing worker built the response. +WORKER_MESSAGE = f"Postprocessing error: {MESSAGE}" + + +@pytest.fixture(autouse=True) +def declaring_extension(): + """A model type whose extension declares two degenerate tokens.""" + + @register_serving_extension(model_types=(DECLARED_MODEL_TYPE,)) + class DeclaringExtension(ServingExtension): + def degenerate_tokens(self, hf_config): + return DegenerateTokens(token_ids=(hf_config.pad_token_id, hf_config.vocab_size - 1)) + + try: + yield + finally: + serving_extensions._BY_MODEL_TYPE.pop(DECLARED_MODEL_TYPE, None) + + +class _DeclaredConfig: + model_type = DECLARED_MODEL_TYPE + pad_token_id = PAD + vocab_size = LAST_ID + 1 + + +class _UndeclaredConfig: + model_type = "degenerate-errors-test-undeclared" + pad_token_id = PAD + vocab_size = LAST_ID + 1 + + +def _engine_stop_words(sampling_params): + """The stop words the engine gets, once the LLM has set the params up.""" + tokenizer = SimpleNamespace( + eos_token_id=1, pad_token_id=None, encode=lambda text, add_special_tokens=False: [99] + ) + sampling_params._setup(tokenizer, None, None) + return sampling_params._get_stop_words() + + +def _step(*frames): + """A streamed step whose frames a postprocessing worker already formatted.""" + return SimpleNamespace(outputs=[SimpleNamespace(_postprocess_result=list(frames))]) + + +class _FailingPromise: + """A generation result that yields ``steps``, then fails on a degenerate run.""" + + def __init__(self, steps, message): + self._steps = list(steps) + self._message = message + self.prompt_token_ids = [1, 2, 3] + self.finished = False + + def __aiter__(self): + return self + + async def __anext__(self): + if not self._steps: + self._fail() + return self._steps.pop(0) + + async def aresult(self): + self._fail() + + def __await__(self): + return self.aresult().__await__() + + def _fail(self): + self.finished = True + raise RequestError(self._message) + + +@pytest.fixture +def make_server(monkeypatch): + """A server for ``config`` whose engine records each request and fails it.""" + + async def fake_apply_chat_template(**kwargs): + return "rendered prompt" + + monkeypatch.setattr(render_chat_module, "async_apply_chat_template", fake_apply_chat_template) + + async def fake_responses_preprocess(*, request, **kwargs): + # Rendering is not under test here; the stop runs are added after it. + return [1, 2, 3], request.to_sampling_params() + + monkeypatch.setattr( + openai_server_module, "responses_api_request_preprocess", fake_responses_preprocess + ) + monkeypatch.setattr( + openai_server_module, + "tokenize_harmony_chat_request", + lambda request, harmony_adapter: [1, 2, 3], + ) + + def make(config=_DeclaredConfig, *, postproc_workers=0, steps=(), message=MESSAGE): + submitted = [] + + def generate_async(*, inputs, sampling_params, **kwargs): + submitted.append(sampling_params) + return _FailingPromise(steps, message) + + server = object.__new__(OpenAIServer) + server.model = MODEL + server.allow_request_chat_template = False + server.model_config = config() + server.processor = None + server.tokenizer = SimpleNamespace(tokenizer=SimpleNamespace(vocab_size=config.vocab_size)) + server.chat_template = None + server.tool_parser = None + server.tool_call_id_type = "random" + server.multimodal_server_config = None + server.use_harmony = False + # Serves /harmony/v1/chat/completions, the chat route of gpt-oss models. + server.harmony_adapter = SimpleNamespace(get_stop_tokens=lambda: [200002]) + server.enable_store = False + server.conversation_store = None + server._input_proc_executor = None + server._adjusted_steady_clock = SimpleNamespace(now=lambda: 0.0) + server.generator = SimpleNamespace( + args=SimpleNamespace( + gather_generation_logits=False, + reasoning_parser=None, + backend="pytorch", + guided_decoding_backend=None, + num_postprocess_workers=postproc_workers, + ), + generate_async=generate_async, + input_processor=lambda prompt, sampling_params: ([1, 2, 3], None), + ) + server.await_disconnected = AsyncMock() + + app = FastAPI() + app.add_api_route("/v1/chat/completions", server.openai_chat, methods=["POST"]) + app.add_api_route("/harmony/v1/chat/completions", server.chat_harmony, methods=["POST"]) + app.add_api_route("/v1/completions", server.openai_completion, methods=["POST"]) + app.add_api_route("/v1/messages", server.anthropic_messages, methods=["POST"]) + app.add_api_route("/v1/responses", server.openai_responses, methods=["POST"]) + # A Responses stream re-raises its failure after the terminal events. + client = TestClient(app, raise_server_exceptions=False) + return SimpleNamespace(server=server, client=client, submitted=submitted) + + return make + + +def _chat(**overrides): + body = {"model": MODEL, "messages": [{"role": "user", "content": "hi"}], "max_tokens": 64} + body.update(overrides) + return body + + +def _sse_data(text): + return [line[len("data: ") :] for line in text.splitlines() if line.startswith("data: ")] + + +def _sse_events(text): + """(event, data) pairs of an SSE body whose frames name their event.""" + events = [] + for frame in text.split("\n\n"): + fields = dict(line.split(": ", 1) for line in frame.splitlines() if ": " in line) + if "event" in fields: + events.append((fields["event"], json.loads(fields["data"]))) + return events + + +# -------------------------------------------------------------------------- +# The declared runs reach the engine; undeclared models are unaffected +# -------------------------------------------------------------------------- + + +# Every route that generates, with the body of a minimal request. +ROUTES = [ + ("/v1/chat/completions", _chat()), + ("/harmony/v1/chat/completions", _chat()), + ("/v1/completions", {"model": MODEL, "prompt": "hi", "max_tokens": 64}), + ("/v1/responses", {"model": MODEL, "input": "hi"}), +] +CHAT_ROUTES = ["/v1/chat/completions", "/harmony/v1/chat/completions"] + + +@pytest.mark.parametrize(("path", "body"), ROUTES) +def test_declared_runs_are_stop_words_of_every_request(make_server, path, body): + server = make_server() + server.client.post(path, json=body) + [sampling_params] = server.submitted + assert _engine_stop_words(sampling_params)[-2:] == [[PAD] * RUN, [LAST_ID] * RUN] + + +@pytest.mark.parametrize( + ("path", "body", "stop_words"), + [ + ("/v1/chat/completions", _chat(stop_token_ids=[7]), [[7]]), + ("/harmony/v1/chat/completions", _chat(), [[200002]]), + ("/v1/completions", {"model": MODEL, "prompt": "hi", "max_tokens": 64}, []), + ("/v1/responses", {"model": MODEL, "input": "hi"}, []), + ], +) +def test_models_that_declare_nothing_keep_their_stop_words(make_server, path, body, stop_words): + server = make_server(_UndeclaredConfig) + server.client.post(path, json=body) + [sampling_params] = server.submitted + assert sampling_params._degenerate_token_runs == () + assert _engine_stop_words(sampling_params) == stop_words + + +# -------------------------------------------------------------------------- +# Each API's client gets an error +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize("message", [MESSAGE, WORKER_MESSAGE]) +@pytest.mark.parametrize(("path", "body"), ROUTES) +def test_non_streamed_requests_fail_with_a_server_error(make_server, path, body, message): + response = make_server(message=message).client.post(path, json=body) + assert response.status_code == 500 + error = response.json() + assert (error["type"], error["code"]) == ("InternalServerError", 500) + assert MESSAGE in error["message"] + + +def test_openai_client_raises_for_a_non_streamed_chat(make_server): + client = openai.OpenAI( + api_key="unused", + base_url="http://testserver/v1", + http_client=make_server().client, + max_retries=0, + ) + with pytest.raises(openai.InternalServerError, match="degenerated"): + client.chat.completions.create(model=MODEL, messages=[{"role": "user", "content": "hi"}]) + + +@pytest.mark.parametrize("path", CHAT_ROUTES) +@pytest.mark.parametrize( + ("postproc_workers", "steps", "message"), + [ + pytest.param(0, (), MESSAGE, id="in_server"), + pytest.param(1, (_step('data: {"id": "chunk-1"}\n\n'),), WORKER_MESSAGE, id="worker"), + ], +) +def test_a_chat_stream_ends_with_an_error_event( + make_server, path, postproc_workers, steps, message +): + server = make_server(postproc_workers=postproc_workers, steps=steps, message=message) + response = server.client.post(path, json=_chat(stream=True)) + + assert response.status_code == 200 + data = _sse_data(response.text) + # Whatever was streamed before the run stays as it was ... + assert data[: len(steps)] == ['{"id": "chunk-1"}'][: len(steps)] + # ... and the stream ends with an error rather than a finished reply. + error = json.loads(data[-2])["error"] + assert error["type"] == "server_error" + assert MESSAGE in error["message"] + assert data[-1] == "[DONE]" + assert not any('"finish_reason"' in chunk for chunk in data) + + +def test_openai_client_raises_for_a_streamed_chat(make_server): + client = openai.OpenAI( + api_key="unused", + base_url="http://testserver/v1", + http_client=make_server().client, + max_retries=0, + ) + stream = client.chat.completions.create( + model=MODEL, messages=[{"role": "user", "content": "hi"}], stream=True + ) + with pytest.raises(openai.APIError, match="degenerated"): + list(stream) + + +@pytest.mark.parametrize( + ("postproc_workers", "message"), + [(0, MESSAGE), (1, WORKER_MESSAGE)], + ids=["in_server", "worker"], +) +def test_an_anthropic_stream_ends_with_an_error_event(make_server, postproc_workers, message): + server = make_server(postproc_workers=postproc_workers, message=message) + body = { + "model": MODEL, + "max_tokens": 64, + "messages": [{"role": "user", "content": "hi"}], + "stream": True, + } + response = server.client.post("/v1/messages", json=body) + + assert response.status_code == 200 + events = _sse_events(response.text) + assert events[-1][0] == "error" + assert events[-1][1]["error"]["type"] == "api_error" + assert MESSAGE in events[-1][1]["error"]["message"] + assert "message_stop" not in [name for name, _ in events] + + +def test_a_non_streamed_anthropic_request_fails_with_an_api_error(make_server): + body = {"model": MODEL, "max_tokens": 64, "messages": [{"role": "user", "content": "hi"}]} + response = make_server().client.post("/v1/messages", json=body) + + assert response.status_code == 500 + error = response.json() + assert error["type"] == "error" + assert error["error"]["type"] == "api_error" + assert MESSAGE in error["error"]["message"] + + +@pytest.mark.parametrize( + ("postproc_workers", "steps", "message"), + [ + pytest.param(0, (), MESSAGE, id="in_server"), + pytest.param( + 1, + ( + _step( + "event: response.output_item.added\n" + 'data: {"type": "response.output_item.added", "sequence_number": 2}\n\n' + ), + ), + WORKER_MESSAGE, + id="worker", + ), + ], +) +@pytest.mark.asyncio +async def test_a_responses_stream_ends_failed_without_output( + make_server, postproc_workers, steps, message +): + server = make_server(postproc_workers=postproc_workers, steps=steps, message=message) + request = ResponsesRequest(model=MODEL, input="hi", stream=True) + raw_request = SimpleNamespace(state=SimpleNamespace(no_client_connection=True), headers={}) + + response = await server.server.openai_responses(request, raw_request) + + assert response.status_code == 200 + # The route re-raises the failure after the terminal events; the frames + # sent before that are what the client receives. (An HTTP test client may + # drop a streamed body whose app raised, so the body is read here.) + frames = [] + with pytest.raises(RequestError): + async for frame in response.body_iterator: + frames.append(frame if isinstance(frame, str) else frame.decode()) + text = "".join(frames) + names = [name for name, _ in _sse_events(text)] + # The item opened before the run is never completed, so a client that keeps + # items on response.output_item.done keeps nothing of this turn. + assert names == [ + "response.created", + "response.in_progress", + *(["response.output_item.added"] if steps else []), + "error", + "response.failed", + ] + failed = _sse_events(text)[-1][1]["response"] + assert failed["status"] == "failed" + assert failed["output"] == [] + assert failed["error"]["code"] == "server_error" + assert MESSAGE in failed["error"]["message"] diff --git a/tests/unittest/llmapi/apps/test_serving_extensions.py b/tests/unittest/llmapi/apps/test_serving_extensions.py index 51b7bea0172c..cde78b805719 100644 --- a/tests/unittest/llmapi/apps/test_serving_extensions.py +++ b/tests/unittest/llmapi/apps/test_serving_extensions.py @@ -21,6 +21,8 @@ from tensorrt_llm.serve import serving_extensions from tensorrt_llm.serve.serving_extensions import ( + DEGENERATE_TOKEN_RUN_LENGTH, + DegenerateTokens, ServingExtension, apply_model_chat_extensions, register_serving_extension, @@ -776,6 +778,7 @@ def test_default_hooks_match_the_generic_path(self) -> None: assert extension.serialize_tool(request.tools[0]) == request.tools[0].model_dump() assert extension.dynamic_tools(request.messages) == [] assert extension.prompt_tokens_excluded_from_usage(request) == 0 + assert extension.degenerate_tokens(SimpleNamespace(pad_token_id=0, vocab_size=8)) is None def test_load_builtin_extensions_keeps_its_private_alias(self) -> None: assert ( @@ -822,3 +825,93 @@ def test_kimi_k3_excludes_the_generation_channel_opener( def test_gpt_oss_declares_harmony_output(self) -> None: extension = serving_extensions.get_serving_extension("gpt_oss") assert extension.output_mode() is serving_extensions.OutputMode.HARMONY_TOKENS + + +class _CompositeConfig: + """A multimodal config: the vocabulary is on its text config.""" + + def __init__(self, vocab_size, **fields): + self.text_config = SimpleNamespace(vocab_size=vocab_size) + self.__dict__.update(fields) + + def get_text_config(self): + return self.text_config + + +class TestDegenerateTokens: + """Which models declare degenerate tokens, and which ids they resolve to.""" + + def test_deepseek_v41_resolves_from_the_checkpoint_config(self, tmp_path) -> None: + """Loaded the way the server loads it, from the fields the release checkpoint ships.""" + import json + + from tensorrt_llm._torch.pyexecutor.config_utils import load_pretrained_config + from tensorrt_llm.serve.chat_utils import resolve_top_level_model_type + + (tmp_path / "config.json").write_text( + json.dumps( + { + "architectures": ["DeepseekV41ForCausalLM"], + "model_type": "deepseek_v41", + "bos_token_id": 0, + "eos_token_id": 1, + "pad_token_id": 2, + "text_config": {"model_type": "deepseek_v41_text", "vocab_size": 129280}, + "vision_config": {"model_type": "deepseek_v41_vision"}, + } + ) + ) + config = load_pretrained_config(str(tmp_path)) + extension = serving_extensions.get_serving_extension(resolve_top_level_model_type(config)) + + # <｜▁pad▁｜> and <|place_holder_mm_span_0442|>. + assert extension.degenerate_tokens(config) == DegenerateTokens( + token_ids=(2, 129279), run_length=DEGENERATE_TOKEN_RUN_LENGTH + ) + + @pytest.mark.parametrize( + ("config", "token_ids"), + [ + pytest.param(_CompositeConfig(129280, pad_token_id=2), (2, 129279), id="composite"), + pytest.param( + SimpleNamespace(pad_token_id=2, vocab_size=129280), (2, 129279), id="flat" + ), + pytest.param(_CompositeConfig(129280), (129279,), id="no_pad_token"), + pytest.param(_CompositeConfig(3, pad_token_id=2), (2,), id="pad_is_the_last_id"), + pytest.param(SimpleNamespace(pad_token_id=2), (2,), id="no_vocabulary_size"), + ], + ) + def test_deepseek_v41_ids(self, config, token_ids) -> None: + extension = serving_extensions.get_serving_extension("deepseek_v41") + assert extension.degenerate_tokens(config).token_ids == token_ids + + def test_deepseek_v41_without_ids_declares_nothing(self) -> None: + extension = serving_extensions.get_serving_extension("deepseek_v41") + assert extension.degenerate_tokens(SimpleNamespace()) is None + + def test_kimi_k3_declares_its_last_vocabulary_id(self) -> None: + """[PAD]; the composite config's pad_token_id defaults to 0 and is not read.""" + extension = serving_extensions.get_serving_extension("kimi_k3") + assert extension.degenerate_tokens(_CompositeConfig(163840, pad_token_id=0)) == ( + DegenerateTokens(token_ids=(163839,)) + ) + assert extension.degenerate_tokens(SimpleNamespace()) is None + + @pytest.mark.parametrize("model_type", ["qwen3", "gpt_oss", "nemotron_h", None]) + def test_other_models_declare_nothing(self, model_type) -> None: + extension = serving_extensions.get_serving_extension(model_type) + config = _CompositeConfig(129280, pad_token_id=2) + assert extension.degenerate_tokens(config) is None + + @pytest.mark.parametrize( + ("fields", "match"), + [ + pytest.param({"token_ids": ()}, "at least one", id="no_ids"), + pytest.param({"token_ids": (-1,)}, "non-negative", id="negative_id"), + pytest.param({"token_ids": (True,)}, "non-negative", id="bool_id"), + pytest.param({"token_ids": (2,), "run_length": 0}, "run_length", id="empty_run"), + ], + ) + def test_declaration_is_validated(self, fields, match) -> None: + with pytest.raises(ValueError, match=match): + DegenerateTokens(**fields) From 9120d1f139dd9c45c5b855d3e07a671f28322f59 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Sun, 1 Nov 2026 21:10:58 +0000 Subject: [PATCH 2/5] [None][fix] exclude prompt tokens from degenerate stop runs Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/llm_request.py | 6 ++ .../_torch/pyexecutor/request_utils.py | 1 + .../pyexecutor/sampler/finish_reasons.py | 25 +++++++- .../_torch/pyexecutor/sampler/ops/triton.py | 20 ++++++- .../pyexecutor/sampler/sampler_features.py | 8 ++- tensorrt_llm/executor/base_worker.py | 5 ++ .../_torch/pyexecutor/test_request_utils.py | 29 +++++++++ .../sampler/test_finish_reasons_fused.py | 60 ++++++++++++++++++- .../executor/test_degenerate_token_runs.py | 42 +++++++++++++ 9 files changed, 186 insertions(+), 10 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index e10c636250ff..ffb7c6d138e8 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -707,6 +707,7 @@ def __init__( additional_outputs: Optional[List[str]] = None, return_perf_metrics: bool = False, stop_words_list: list[list[int]] | None = None, + num_generated_only_stop_words: int = 0, llm_request: Optional[ tensorrt_llm.bindings.internal.batch_manager.LlmRequest] = None, parent_request: Optional["LlmRequest"] = None, @@ -853,6 +854,8 @@ def __init__( # Keep py_stop_words_list as a python list, rather than a tensor. self.py_stop_words_list = stop_words_list + # Only the trailing server-added rules exclude prompt tokens. + self.py_num_generated_only_stop_words = num_generated_only_stop_words self.py_logprobs_mode = LogprobMode( logprobs_mode) # handle passed a raw string @@ -1242,6 +1245,7 @@ def create_child_request(self, child_id: int) -> None: return_perf_metrics=self.return_perf_metrics, return_encoder_output=self.py_return_encoder_output, stop_words_list=deepcopy(self.py_stop_words_list), + num_generated_only_stop_words=self.py_num_generated_only_stop_words, num_logprobs=self.py_num_logprobs, use_chunked_generation_logits=self.py_use_chunked_generation_logits, logits_chunk_size=self.py_logits_chunk_size, @@ -1679,6 +1683,8 @@ def executor_request_to_llm_request( end_id=executor_request.end_id, embedding_bias=executor_request.embedding_bias, stop_words_list=stop_words_list, + num_generated_only_stop_words=getattr( + executor_request, "py_num_generated_only_stop_words", 0), position_ids=position_ids, multimodal_hashes=multimodal_hashes, multimodal_positions=multimodal_positions, diff --git a/tensorrt_llm/_torch/pyexecutor/request_utils.py b/tensorrt_llm/_torch/pyexecutor/request_utils.py index aa8f5636bd6c..ed9f10b4eb7f 100644 --- a/tensorrt_llm/_torch/pyexecutor/request_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/request_utils.py @@ -521,6 +521,7 @@ def _collect_py_objects(self, new_requests: List) -> Tuple: "py_multimodal_data", "py_scheduling_params", "py_num_logprobs", + "py_num_generated_only_stop_words", "py_disaggregated_params", "py_lora_path", "py_lora_ckpt_source", diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/finish_reasons.py b/tensorrt_llm/_torch/pyexecutor/sampler/finish_reasons.py index aeba2cd4043b..9499761f74e2 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/finish_reasons.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/finish_reasons.py @@ -249,7 +249,12 @@ def prepare_for_new_request(self, request: LlmRequest) -> None: for word in words: flat_tokens.extend(word) offsets.append(len(flat_tokens)) - history_offset = len(offsets) + len(flat_tokens) + generated_only_start = len(words) - getattr(request, "py_num_generated_only_stop_words", 0) + min_lengths = [ + request.py_orig_prompt_len + len(word) if index >= generated_only_start else 0 + for index, word in enumerate(words) + ] + history_offset = len(offsets) + len(flat_tokens) + len(min_lengths) window_length = max_length - 1 + self._max_tokens if max_length > 1 else 0 history_size = window_length * self._max_beam_width host = torch.zeros( @@ -258,7 +263,7 @@ def prepare_for_new_request(self, request: LlmRequest) -> None: device="cpu", pin_memory=prefer_pinned(), ) - host[:history_offset] = torch.tensor(offsets + flat_tokens, dtype=torch.int32) + host[:history_offset] = torch.tensor(offsets + flat_tokens + min_lengths, dtype=torch.int32) if window_length: history = host[history_offset : history_offset + history_size].view( window_length, self._max_beam_width @@ -537,6 +542,7 @@ def _write_finish_reasons( num_accepted_tokens[stop_word_indices] if isinstance(num_accepted_tokens, torch.Tensor) else num_accepted_tokens, + seq_lens=seq_lens[stop_word_indices], ), FinishReason.STOP_WORDS.value, ) @@ -616,6 +622,8 @@ def _are_stop_words( seq_slots: torch.Tensor, tokens: torch.Tensor, num_accepted_tokens: torch.Tensor | int, + *, + seq_lens: torch.Tensor | None = None, ) -> torch.Tensor: """Checks if the tokens are stop words @@ -626,6 +634,8 @@ def _are_stop_words( Shape: [max_tokens, len(requests), max_beam_width] num_accepted_tokens: The number of accepted draft tokens +1 for each request. Shape: [len(requests)] if torch.Tensor + seq_lens: Sequence lengths before the newly generated tokens, shaped [requests]. + Required for generated-only rules. Returns: A tensor where each element is True if the sequence at the corresponding token ends with a stop word, False otherwise @@ -633,7 +643,11 @@ def _are_stop_words( """ self.record_stop_word_stream() return compact_stop_words( - self.store.stop_words_metadata_cuda, seq_slots, tokens, num_accepted_tokens + self.store.stop_words_metadata_cuda, + seq_slots, + tokens, + num_accepted_tokens, + seq_lens=seq_lens, ) @nvtx_range("_are_stop_words_single_token") @@ -642,6 +656,8 @@ def _are_stop_words_single_token( seq_slots: torch.Tensor, tokens: torch.Tensor, _num_accepted_tokens: torch.Tensor | int, + *, + seq_lens: torch.Tensor | None = None, ) -> torch.Tensor: """Checks if the tokens are stop words (single token per stop word only) @@ -651,6 +667,8 @@ def _are_stop_words_single_token( tokens: A buffer containing the newly generated tokens. Shape: [max_tokens, len(requests), max_beam_width] _num_accepted_tokens: Unused + seq_lens: Sequence lengths before the newly generated tokens, shaped [requests]. + Required for generated-only rules. Returns: A tensor where each element is True if the sequence at the corresponding token ends with a stop word, False otherwise @@ -663,4 +681,5 @@ def _are_stop_words_single_token( tokens, _num_accepted_tokens, single_token_only=True, + seq_lens=seq_lens, ) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/ops/triton.py b/tensorrt_llm/_torch/pyexecutor/sampler/ops/triton.py index 7b482faf0737..71fec4284ac5 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/ops/triton.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/ops/triton.py @@ -192,6 +192,8 @@ def _compact_stop_match_kernel( slots, tokens, output, + seq_lens, + HAS_SEQ_LENS: tl.constexpr, token_stride_step, token_stride_request, token_stride_beam, @@ -224,7 +226,13 @@ def _compact_stop_match_kernel( begin = tl.load(base + rule, valid, other=0) end = tl.load(base + rule + 1, valid, other=0) length = end - begin - matches = valid & (length > 0) + min_lengths = base + rule_count + 1 + tl.load(base + rule_count) + min_length = tl.load(min_lengths + rule, valid, other=0) + if HAS_SEQ_LENS: + enough_tokens = tl.load(seq_lens + request) + step + 1 >= min_length + else: + enough_tokens = min_length == 0 + matches = valid & (length > 0) & enough_tokens for offset in range(tl.max(length, 0)): active = valid & (offset < length) expected = tl.load(base + rule_count + 1 + begin + offset, active, other=0) @@ -245,12 +253,14 @@ def compact_stop_words( accepted: torch.Tensor | int, *, single_token_only: bool = False, + seq_lens: torch.Tensor | None = None, ) -> torch.Tensor: """Match packed per-request rules without a rule-by-token comparison tensor. Metadata points to int32 allocations containing rule offsets, rule tokens, - and two rolling-history windows. The caller retains these allocations and - records their use on the current stream because they are indirect inputs. + per-rule minimum sequence lengths, and two rolling-history windows. The caller + retains these allocations and records their use on the current stream because + they are indirect inputs. Args: metadata: Contiguous CUDA int64 tensor of shape [max_num_sequences, 5]. @@ -265,6 +275,8 @@ def compact_stop_words( one, used to advance the rolling history. Unused for single-token rules. single_token_only: Whether all rules contain exactly one token, allowing history updates to be skipped. + seq_lens: CUDA sequence lengths before these tokens, shaped [requests]. + Required to match generated-only rules; ordinary rules need no lengths. Returns: CUDA bool tensor of shape [steps, requests, beams], true where a sequence @@ -293,6 +305,8 @@ def compact_stop_words( slots, tokens, result, + seq_lens, + seq_lens is not None, *tokens.stride(), requests, steps, diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler_features.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler_features.py index b66f2d9fc18b..ed7b53292d30 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler_features.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler_features.py @@ -358,8 +358,12 @@ def meet_stop_token_criteria( # Slow path: at least one multi-token stop word exists tokens = request.get_tokens(beam_idx) - for stop_word in stop_words_list: - if len(stop_word) > len(tokens): + generated_only_start = len(stop_words_list) - request.py_num_generated_only_stop_words + for index, stop_word in enumerate(stop_words_list): + available = len(tokens) + if index >= generated_only_start: + available -= request.py_orig_prompt_len + if len(stop_word) > available: continue if tokens[-len(stop_word) :] == stop_word: return True diff --git a/tensorrt_llm/executor/base_worker.py b/tensorrt_llm/executor/base_worker.py index 217d3736cc40..5142ae8f65b5 100644 --- a/tensorrt_llm/executor/base_worker.py +++ b/tensorrt_llm/executor/base_worker.py @@ -483,6 +483,11 @@ def _deduce_max_tokens(request: GenerationRequest) -> int: disagg_request_id=disagg_request_id, cache_salt=request.cache_salt, priority=request.priority) + # Server-added degenerate runs are the trailing stop rules and + # must never consume a suffix of the input prompt. + executor_request.py_num_generated_only_stop_words = ( + 0 if request.sampling_params.ignore_eos else len( + request.sampling_params._degenerate_token_runs)) executor_request.py_original_end_id = request.sampling_params.end_id executor_request.py_num_logprobs = request.sampling_params.logprobs executor_request.py_lora_path = py_lora_path diff --git a/tests/unittest/_torch/pyexecutor/test_request_utils.py b/tests/unittest/_torch/pyexecutor/test_request_utils.py index 92aa4c531667..fe18cbd3e8e9 100644 --- a/tests/unittest/_torch/pyexecutor/test_request_utils.py +++ b/tests/unittest/_torch/pyexecutor/test_request_utils.py @@ -1746,3 +1746,32 @@ def test_sync_clock_aux_rides_along(self): (values,), _ = other.broadcast_int64.call_args assert list(values) == [0, 0, -1] assert b.rank0_aux == pytest.approx(0.17) and b.rank0_now == pytest.approx(1.0) + + +def test_generated_only_stop_rules_survive_broadcast_and_child_cloning() -> None: + """Rule provenance must reach every rank and every return sequence.""" + source = trtllm.Request( + input_token_ids=[2] * 19, + max_tokens=32, + stop_words=[[7, 2], [2] * 20], + sampling_config=trtllm.SamplingConfig(num_return_sequences=2), + ) + source.py_num_generated_only_stop_words = 1 + source.py_conversation_params = None + objects = RequestBroadcaster._collect_py_objects(None, [RequestQueueItem(42, source)]) + restored = trtllm.Request( + input_token_ids=[2] * 19, + max_tokens=32, + stop_words=[[7, 2], [2] * 20], + sampling_config=trtllm.SamplingConfig(num_return_sequences=2), + ) + attach_py_objects_to_requests( + [RequestQueueItem(42, restored)], pickle.loads(pickle.dumps(objects)) + ) + request = executor_request_to_llm_request( + 42, restored, child_req_ids=[43], exclude_last_generation_logits=False + ) + for candidate in (request, request.child_requests[0]): + assert candidate.py_num_generated_only_stop_words == 1 + assert candidate.py_stop_words_list == [[7, 2], [2] * 20] + assert candidate.py_orig_prompt_len == 19 diff --git a/tests/unittest/_torch/sampler/test_finish_reasons_fused.py b/tests/unittest/_torch/sampler/test_finish_reasons_fused.py index 65cf23eddac0..5637d4a6c474 100644 --- a/tests/unittest/_torch/sampler/test_finish_reasons_fused.py +++ b/tests/unittest/_torch/sampler/test_finish_reasons_fused.py @@ -145,7 +145,7 @@ def test_compact_stop_words_history_and_slot_reuse() -> None: assert handler.store.stop_words_metadata_cuda[8].tolist() == [ replacement.tensor.data_ptr(), 2, - 8, + 10, steps + 2, beams, ] @@ -254,7 +254,7 @@ def admit(slot: int, words: list[list[int]], prompt: list[int]) -> SimpleNamespa ) # Metadata plus actual rules and per-request windows; no 1024 x 1024 x 256 tensor. assert total_bytes < 64 * 1024 - assert handler._stop_word_allocations[3].tensor.numel() == 2 + 256 + 2 * 256 + assert handler._stop_word_allocations[3].tensor.numel() == 2 + 256 + 1 + 2 * 256 slots = torch.tensor([2], device="cuda", dtype=torch.int64) tokens = torch.zeros((1, 1025, 1), device="cuda", dtype=torch.int32) @@ -314,3 +314,59 @@ def admit(slot: int, words: list[list[int]], prompt: list[int]) -> SimpleNamespa handler.release_request(long) assert not handler._stop_word_allocations assert not handler.has_stop_word_history + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("steps", [1, 3]) +def test_generated_only_stop_words_preserve_prompt_spanning_rules(steps: int) -> None: + """Production-length runs ignore prompt tokens, including rejected draft suffixes.""" + from types import SimpleNamespace + + handler = FinishReasonsHandler( + max_num_sequences=3, max_beam_width=1, max_tokens=steps, max_seq_len=256 + ) + slots_host = torch.tensor([0, 2], dtype=torch.int32) + slots = slots_host.cuda().long() + handler.setup_new_request_handling() + for slot in (0, 2): + request = SimpleNamespace( + py_seq_slot=slot, + py_stop_words_list=[[2] * 20] * (1 if slot == 0 else 2), + py_num_generated_only_stop_words=1, + py_orig_prompt_len=19, + py_max_new_tokens=100, + py_end_id=99, + is_context_only_request=False, + py_beam_width=1, + get_tokens=lambda: [[2] * 19], + ) + handler.prepare_for_new_request(request) + handler.update_for_new_request( + seq_slots_cuda_long=slots, + max_lengths_cuda=torch.full((2,), 119, device="cuda", dtype=torch.int32), + end_ids_cuda=torch.full((2,), 99, device="cuda", dtype=torch.int32), + seq_slots_host=slots_host, + ) + tokens = torch.full((steps, 4, 1), 2, device="cuda", dtype=torch.int32) + generated = 0 + previous_accepted = 1 + iteration = 0 + while generated < 20: + handler.store.num_accepted_draft_tokens_host[slots_host] = previous_accepted - 1 + reasons = handler.write_finish_reasons( + seq_slots_host=slots_host, + seq_slots_cuda=slots, + seq_lens_cuda=torch.full((2,), 19 + generated, device="cuda", dtype=torch.int32), + new_tokens_cuda=tokens, + ).cpu() + expected = [ + FinishReason.STOP_WORDS.value + if generated + step + 1 >= 20 + else FinishReason.NOT_FINISHED.value + for step in range(steps) + ] + assert reasons[:, 0, 0].tolist() == expected + assert reasons[:, 2, 0].tolist() == [FinishReason.STOP_WORDS.value] * steps + previous_accepted = 1 if iteration % 2 == 0 else steps + generated += previous_accepted + iteration += 1 diff --git a/tests/unittest/executor/test_degenerate_token_runs.py b/tests/unittest/executor/test_degenerate_token_runs.py index e752db1c171c..9a7f23cb5812 100644 --- a/tests/unittest/executor/test_degenerate_token_runs.py +++ b/tests/unittest/executor/test_degenerate_token_runs.py @@ -57,6 +57,8 @@ class _Request: def __init__(self, params: SamplingParams): self.py_stop_words_list = params._get_stop_words() + self.py_num_generated_only_stop_words = len(params._degenerate_token_runs) + self.py_orig_prompt_len = len(PROMPT) self.tokens = list(PROMPT) def get_tokens(self, beam_idx: int) -> list: @@ -239,3 +241,43 @@ def test_detector_matches_exact_and_wrapped_messages(): def test_detector_rejects_other_errors(): assert not is_degenerate_generation_message("Generation failed: request cancelled") assert not is_degenerate_generation_message("") + + +@pytest.mark.parametrize("ordinary_rule", [False, True]) +def test_production_run_does_not_count_prompt_suffix(ordinary_rule: bool) -> None: + """Only server-added runs exclude a pre-tokenized prompt suffix.""" + params = SamplingParams(end_id=END_ID) + params._set_degenerate_token_runs([PAD], 20) + tokens = [PAD] * 19 + request = SimpleNamespace( + py_stop_words_list=([[PAD] * 20] if ordinary_rule else []) + params._get_stop_words(), + py_num_generated_only_stop_words=1, + py_orig_prompt_len=len(tokens), + get_tokens=lambda beam_idx: tokens, + ) + for generated in range(1, 21): + tokens.append(PAD) + assert meet_stop_token_criteria(request, PAD) == (ordinary_rule or generated >= 20) + result = _finished_result(params) + with pytest.raises(RequestError, match="degenerate"): + _finish(result, [PAD] * 20, tllm.FinishReason.STOP_WORDS) + + +@pytest.mark.parametrize("ignore_eos", [False, True]) +def test_worker_marks_only_server_added_stop_rules(ignore_eos: bool) -> None: + from tensorrt_llm.executor.base_worker import BaseWorker + from tensorrt_llm.executor.request import GenerationRequest + + params = SamplingParams(end_id=END_ID, max_tokens=32, stop_token_ids=[7], ignore_eos=ignore_eos) + params._set_degenerate_token_runs([PAD], 20) + request = GenerationRequest(prompt_token_ids=[PAD] * 19, sampling_params=params) + request.set_id(42) + submitted: list[tllm.Request] = [] + worker = SimpleNamespace( + _lora_manager=None, + max_seq_len=256, + engine=SimpleNamespace(enqueue_request=lambda value: submitted.append(value) or 42), + ) + assert BaseWorker._enqueue_request(worker, request) == 42 + assert submitted[0].py_num_generated_only_stop_words == (0 if ignore_eos else 1) + assert submitted[0].stop_words == ([] if ignore_eos else [[7], [PAD] * 20]) From 2922cfffcfd202a5b035f0ed982a5566e1eee2f9 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Sun, 1 Nov 2026 22:35:38 +0000 Subject: [PATCH 3/5] [None][fix] retain degenerate detection across pauses and length limits Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- .../pyexecutor/sampler/finish_reasons.py | 3 +- .../pyexecutor/sampler/sampler_features.py | 2 +- tensorrt_llm/executor/result.py | 11 +++ .../sampler/test_finish_reasons_fused.py | 50 +++++++++++++ .../executor/test_degenerate_token_runs.py | 71 ++++++++++++++++++- 5 files changed, 134 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/finish_reasons.py b/tensorrt_llm/_torch/pyexecutor/sampler/finish_reasons.py index 9499761f74e2..692aee045a42 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/finish_reasons.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/finish_reasons.py @@ -250,8 +250,9 @@ def prepare_for_new_request(self, request: LlmRequest) -> None: flat_tokens.extend(word) offsets.append(len(flat_tokens)) generated_only_start = len(words) - getattr(request, "py_num_generated_only_stop_words", 0) + # The native original boundary survives pause/recompute. min_lengths = [ - request.py_orig_prompt_len + len(word) if index >= generated_only_start else 0 + request.orig_prompt_len + len(word) if index >= generated_only_start else 0 for index, word in enumerate(words) ] history_offset = len(offsets) + len(flat_tokens) + len(min_lengths) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler_features.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler_features.py index ed7b53292d30..3d3b02583ad5 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler_features.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler_features.py @@ -362,7 +362,7 @@ def meet_stop_token_criteria( for index, stop_word in enumerate(stop_words_list): available = len(tokens) if index >= generated_only_start: - available -= request.py_orig_prompt_len + available -= request.orig_prompt_len if len(stop_word) > available: continue if tokens[-len(stop_word) :] == stop_word: diff --git a/tensorrt_llm/executor/result.py b/tensorrt_llm/executor/result.py index 5b65345d8b41..6f905af1aa0a 100644 --- a/tensorrt_llm/executor/result.py +++ b/tensorrt_llm/executor/result.py @@ -520,6 +520,17 @@ def _handle_sequence(self, break elif finish_reasons[src_idx] == tllm.FinishReason.LENGTH: output.finish_reason = 'length' + # Length takes precedence in both samplers, but a completed + # degenerate run still invalidates the generated response. + if (self.sampling_params._degenerate_token_runs + and not self.sampling_params.ignore_eos): + for stop_reason, stop_ids in self.sampling_params._get_stop_reasons_and_words( + ): + if output.token_ids[-len(stop_ids):] == stop_ids: + if isinstance(stop_reason, DegenerateTokenRun): + self._fail_degenerate_generation( + stop_reason, output) + break elif finish_reasons[src_idx] == tllm.FinishReason.TIMED_OUT: output.finish_reason = 'timeout' # For disaggregated serving, finish reason might be NOT_FINISHED which is ok diff --git a/tests/unittest/_torch/sampler/test_finish_reasons_fused.py b/tests/unittest/_torch/sampler/test_finish_reasons_fused.py index 5637d4a6c474..ccd73d586080 100644 --- a/tests/unittest/_torch/sampler/test_finish_reasons_fused.py +++ b/tests/unittest/_torch/sampler/test_finish_reasons_fused.py @@ -334,6 +334,7 @@ def test_generated_only_stop_words_preserve_prompt_spanning_rules(steps: int) -> py_stop_words_list=[[2] * 20] * (1 if slot == 0 else 2), py_num_generated_only_stop_words=1, py_orig_prompt_len=19, + orig_prompt_len=19, py_max_new_tokens=100, py_end_id=99, is_context_only_request=False, @@ -370,3 +371,52 @@ def test_generated_only_stop_words_preserve_prompt_spanning_rules(steps: int) -> previous_accepted = 1 if iteration % 2 == 0 else steps generated += previous_accepted iteration += 1 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("steps", [1, 3]) +def test_generated_only_stop_words_survive_native_recompute(steps: int) -> None: + from tensorrt_llm._torch.pyexecutor.request_utils import executor_request_to_llm_request + from tensorrt_llm.bindings import executor as tllm + + source = tllm.Request(input_token_ids=[2] * 19, max_tokens=40, stop_words=[[2] * 20], end_id=99) + source.py_num_generated_only_stop_words = 1 + request = executor_request_to_llm_request( + 42, source, child_req_ids=[], exclude_last_generation_logits=False + ) + for _ in range(10): + request.add_new_token(2, 0) + request.reset_for_recompute(256) + assert request.py_orig_prompt_len == 29 + assert request.orig_prompt_len == 19 + request.py_seq_slot = 0 + handler = FinishReasonsHandler( + max_num_sequences=1, max_beam_width=1, max_tokens=steps, max_seq_len=256 + ) + slots_host = torch.tensor([0], dtype=torch.int32) + slots = slots_host.cuda().long() + handler.setup_new_request_handling() + handler.prepare_for_new_request(request) + handler.update_for_new_request( + seq_slots_cuda_long=slots, + max_lengths_cuda=torch.tensor([59], device="cuda", dtype=torch.int32), + end_ids_cuda=torch.tensor([99], device="cuda", dtype=torch.int32), + seq_slots_host=slots_host, + ) + tokens = torch.full((steps, 2, 1), 2, device="cuda", dtype=torch.int32) + for generated in range(10, 20): + # Accept one token per step; speculative suffixes are discarded. + handler.store.num_accepted_draft_tokens_host[slots_host] = 0 + reasons = handler.write_finish_reasons( + seq_slots_host=slots_host, + seq_slots_cuda=slots, + seq_lens_cuda=torch.tensor([19 + generated], device="cuda", dtype=torch.int32), + new_tokens_cuda=tokens, + ).cpu() + expected = [ + FinishReason.STOP_WORDS.value + if generated + step + 1 >= 20 + else FinishReason.NOT_FINISHED.value + for step in range(steps) + ] + assert reasons[:, 0, 0].tolist() == expected diff --git a/tests/unittest/executor/test_degenerate_token_runs.py b/tests/unittest/executor/test_degenerate_token_runs.py index 9a7f23cb5812..262de27fb6be 100644 --- a/tests/unittest/executor/test_degenerate_token_runs.py +++ b/tests/unittest/executor/test_degenerate_token_runs.py @@ -25,7 +25,11 @@ import pytest -from tensorrt_llm._torch.pyexecutor.sampler.sampler_features import meet_stop_token_criteria +from tensorrt_llm._torch.pyexecutor.request_utils import executor_request_to_llm_request +from tensorrt_llm._torch.pyexecutor.sampler.sampler_features import ( + handle_stop_criteria, + meet_stop_token_criteria, +) from tensorrt_llm.bindings import executor as tllm from tensorrt_llm.executor.result import GenerationResultBase from tensorrt_llm.executor.utils import ( @@ -59,6 +63,7 @@ def __init__(self, params: SamplingParams): self.py_stop_words_list = params._get_stop_words() self.py_num_generated_only_stop_words = len(params._degenerate_token_runs) self.py_orig_prompt_len = len(PROMPT) + self.orig_prompt_len = len(PROMPT) self.tokens = list(PROMPT) def get_tokens(self, beam_idx: int) -> list: @@ -253,6 +258,7 @@ def test_production_run_does_not_count_prompt_suffix(ordinary_rule: bool) -> Non py_stop_words_list=([[PAD] * 20] if ordinary_rule else []) + params._get_stop_words(), py_num_generated_only_stop_words=1, py_orig_prompt_len=len(tokens), + orig_prompt_len=len(tokens), get_tokens=lambda beam_idx: tokens, ) for generated in range(1, 21): @@ -281,3 +287,66 @@ def test_worker_marks_only_server_added_stop_rules(ignore_eos: bool) -> None: assert BaseWorker._enqueue_request(worker, request) == 42 assert submitted[0].py_num_generated_only_stop_words == (0 if ignore_eos else 1) assert submitted[0].stop_words == ([] if ignore_eos else [[7], [PAD] * 20]) + + +def test_generated_run_survives_native_pause_and_recompute() -> None: + source = tllm.Request( + input_token_ids=[PAD] * 19, max_tokens=40, stop_words=[[PAD] * 20], end_id=END_ID + ) + source.py_num_generated_only_stop_words = 1 + request = executor_request_to_llm_request( + 42, source, child_req_ids=[], exclude_last_generation_logits=False + ) + for _ in range(10): + request.add_new_token(PAD, 0) + assert not meet_stop_token_criteria(request, PAD) + request.reset_for_recompute(256) + assert request.orig_prompt_len == 19 + assert request.py_orig_prompt_len == 29 + for count in range(1, 11): + request.add_new_token(PAD, 0) + assert meet_stop_token_criteria(request, PAD) == (count == 10) + + +@pytest.mark.parametrize("limit", ["max_tokens", "max_seq_len"]) +@pytest.mark.parametrize("ignored", [False, True]) +def test_generated_run_at_length_boundary_fails(limit: str, ignored: bool) -> None: + params = SamplingParams( + end_id=END_ID, max_tokens=20 if limit == "max_tokens" else 40, ignore_eos=ignored + ) + params._set_degenerate_token_runs([PAD], 20) + source = tllm.Request( + input_token_ids=list(PROMPT), + max_tokens=params.max_tokens, + stop_words=[] if ignored else params._get_stop_words(), + end_id=END_ID, + ) + source.py_num_generated_only_stop_words = 0 if ignored else 1 + request = executor_request_to_llm_request( + 42, source, child_req_ids=[], exclude_last_generation_logits=False + ) + max_seq_len = len(PROMPT) + 20 if limit == "max_seq_len" else 256 + for count in range(1, 21): + request.add_new_token(PAD, 0) + assert handle_stop_criteria(request, PAD, max_seq_len=max_seq_len, beam_idx=0) == ( + count == 20 + ) + result = _finished_result(params) + assert request.is_finished_due_to_length + reason = tllm.FinishReason.LENGTH + if ignored: + _finish(result, [PAD] * 20, reason) + assert result._outputs[0].finish_reason == "length" + else: + with pytest.raises(RequestError, match="degenerate"): + _finish(result, [PAD] * 20, reason) + + +def test_user_stop_at_length_boundary_keeps_length_result() -> None: + params = SamplingParams(end_id=END_ID, max_tokens=20, stop_token_ids=[PAD]) + params._set_degenerate_token_runs([PAD], 20) + result = _finished_result(params) + _finish(result, [PAD] * 20, tllm.FinishReason.LENGTH) + assert result._outputs[0].finish_reason == "length" + assert result._outputs[0].token_ids == [PAD] * 20 + assert result.error is None From d675add3094e9b2fbb3161f197e4c6b61af04b72 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Sun, 1 Nov 2026 23:50:35 +0000 Subject: [PATCH 4/5] [None][fix] count only generated tokens at length limits Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- tensorrt_llm/executor/result.py | 17 +++++- .../executor/test_degenerate_token_runs.py | 53 ++++++++++++++++++- 2 files changed, 68 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/executor/result.py b/tensorrt_llm/executor/result.py index 6f905af1aa0a..1e58147075a0 100644 --- a/tensorrt_llm/executor/result.py +++ b/tensorrt_llm/executor/result.py @@ -524,10 +524,25 @@ def _handle_sequence(self, # degenerate run still invalidates the generated response. if (self.sampling_params._degenerate_token_runs and not self.sampling_params.ignore_eos): + num_generated_tokens = output.length - len( + decoder_output_prefix) + if (not self.sampling_params.exclude_input_from_output + and not getattr(self, "_streaming", False)): + # Non-streamed output can contain the original prompt. + # Postprocessing records carry its count, not its tokens. + num_prompt_tokens = len( + getattr(self, "prompt_token_ids", ())) + if (self.postproc_params is not None and + self.postproc_params.postproc_args is not None): + num_prompt_tokens = (self.postproc_params. + postproc_args.num_prompt_tokens + or num_prompt_tokens) + num_generated_tokens -= num_prompt_tokens for stop_reason, stop_ids in self.sampling_params._get_stop_reasons_and_words( ): if output.token_ids[-len(stop_ids):] == stop_ids: - if isinstance(stop_reason, DegenerateTokenRun): + if (isinstance(stop_reason, DegenerateTokenRun) + and num_generated_tokens >= len(stop_ids)): self._fail_degenerate_generation( stop_reason, output) break diff --git a/tests/unittest/executor/test_degenerate_token_runs.py b/tests/unittest/executor/test_degenerate_token_runs.py index 262de27fb6be..190e7ecf71d3 100644 --- a/tests/unittest/executor/test_degenerate_token_runs.py +++ b/tests/unittest/executor/test_degenerate_token_runs.py @@ -31,7 +31,9 @@ meet_stop_token_criteria, ) from tensorrt_llm.bindings import executor as tllm -from tensorrt_llm.executor.result import GenerationResultBase +from tensorrt_llm.executor.postproc_worker import PostprocArgs, PostprocParams, PostprocWorker +from tensorrt_llm.executor.request import GenerationRequest +from tensorrt_llm.executor.result import GenerationResult, GenerationResultBase from tensorrt_llm.executor.utils import ( RequestError, degenerate_generation_message, @@ -350,3 +352,52 @@ def test_user_stop_at_length_boundary_keeps_length_result() -> None: assert result._outputs[0].finish_reason == "length" assert result._outputs[0].token_ids == [PAD] * 20 assert result.error is None + + +@pytest.mark.parametrize("postproc_worker", [False, True]) +@pytest.mark.parametrize("output_mode", ["prompt", "generated", "streaming"]) +@pytest.mark.parametrize("generated_count", [1, 20]) +def test_length_limited_run_counts_only_generated_tokens( + postproc_worker: bool, output_mode: str, generated_count: int +) -> None: + """Returned prompt tokens do not contribute to a server-added run.""" + prompt = [PAD] * 19 + streaming = output_mode == "streaming" + params = SamplingParams( + end_id=END_ID, + max_tokens=generated_count, + exclude_input_from_output=output_mode == "generated", + ) + params._set_degenerate_token_runs([PAD], 20) + if postproc_worker: + result = PostprocWorker.default_record_creator( + PostprocWorker.Input( + rsp=SimpleNamespace(client_id=11), + sampling_params=params, + postproc_params=PostprocParams( + postproc_args=PostprocArgs(num_prompt_tokens=len(prompt)) + ), + streaming=streaming, + ), + tokenizer=None, + ) + else: + request = GenerationRequest( + prompt_token_ids=prompt, sampling_params=params, streaming=streaming + ) + request.set_id(11) + result = GenerationResult(request) + generated = [PAD] * generated_count + if streaming and generated_count > 1: + _finish(result, generated[:-1], tllm.FinishReason.NOT_FINISHED) + generated = generated[-1:] + returned = prompt + generated if output_mode == "prompt" else generated + result._done = True + if generated_count == 20: + with pytest.raises(RequestError, match="degenerate"): + _finish(result, returned, tllm.FinishReason.LENGTH) + else: + _finish(result, returned, tllm.FinishReason.LENGTH) + assert result._outputs[0].finish_reason == "length" + assert result._outputs[0].token_ids == returned + assert result.error is None From 55140dade39dbd0b0e025c15c643edcae34462d8 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Mon, 2 Nov 2026 01:04:22 +0000 Subject: [PATCH 5/5] [None][fix] initialize prompt counts for direct submissions Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- tensorrt_llm/executor/result.py | 6 ++ .../executor/test_degenerate_token_runs.py | 93 ++++++++++++++++++- 2 files changed, 98 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/executor/result.py b/tensorrt_llm/executor/result.py index 1e58147075a0..ae8e931a13ea 100644 --- a/tensorrt_llm/executor/result.py +++ b/tensorrt_llm/executor/result.py @@ -1185,6 +1185,12 @@ def __init__( postproc_params=generation_request.postproc_params, ) self._generation_request = generation_request + if self.postproc_params is not None: + args = self.postproc_params.postproc_args + if args is not None and args.num_prompt_tokens is None: + # Postprocessing records carry the prompt count, not its tokens. + args.num_prompt_tokens = len( + generation_request.prompt_token_ids) self._streaming = generation_request.streaming self._disaggregated_params = disaggregated_params # minimal sampling params needed for logprob calculation diff --git a/tests/unittest/executor/test_degenerate_token_runs.py b/tests/unittest/executor/test_degenerate_token_runs.py index 190e7ecf71d3..5300deaed377 100644 --- a/tests/unittest/executor/test_degenerate_token_runs.py +++ b/tests/unittest/executor/test_degenerate_token_runs.py @@ -21,7 +21,9 @@ exercised through its host-side reference, ``meet_stop_token_criteria``. """ -from types import SimpleNamespace +from copy import deepcopy +from types import MethodType, SimpleNamespace +from unittest.mock import MagicMock import pytest @@ -401,3 +403,92 @@ def test_length_limited_run_counts_only_generated_tokens( assert result._outputs[0].finish_reason == "length" assert result._outputs[0].token_ids == returned assert result.error is None + + +@pytest.mark.parametrize("executor_kind", ["worker", "proxy"]) +@pytest.mark.parametrize("generated_count", [1, 20]) +def test_direct_submit_populates_prompt_count_for_postprocessing( + executor_kind: str, generated_count: int +) -> None: + from tensorrt_llm.executor.base_worker import BaseWorker + from tensorrt_llm.executor.proxy import GenerationExecutorProxy + + prompt = [PAD] * 19 + params = SamplingParams( + end_id=END_ID, max_tokens=generated_count, exclude_input_from_output=False + ) + params._set_degenerate_token_runs([PAD], 20) + request = GenerationRequest( + prompt_token_ids=prompt, + sampling_params=params, + postproc_params=PostprocParams(postproc_args=PostprocArgs()), + ) + executor = MagicMock( + rank=0, + _results={}, + _client_id_to_request_id={}, + _engine_dead=False, + _fatal_error=None, + ) + executor._handle_background_error = MethodType(lambda self: None, executor) + executor._get_next_client_id.return_value = 11 + executor._get_logprob_params.return_value = None + records = [] + + def dispatch(submitted: GenerationRequest) -> int: + records.append( + PostprocWorker.default_record_creator( + PostprocWorker.Input( + rsp=SimpleNamespace(client_id=submitted.id), + sampling_params=submitted.sampling_params, + postproc_params=deepcopy(submitted.postproc_params), + streaming=False, + ), + tokenizer=None, + ) + ) + return 42 + + executor._enqueue_request.side_effect = dispatch + executor._send_request.side_effect = dispatch + executor_class = BaseWorker if executor_kind == "worker" else GenerationExecutorProxy + executor_class.submit(executor, request) + (record,) = records + record._done = True + returned = prompt + [PAD] * generated_count + if generated_count == 20: + with pytest.raises(RequestError, match="degenerate"): + _finish(record, returned, tllm.FinishReason.LENGTH) + else: + _finish(record, returned, tllm.FinishReason.LENGTH) + assert record._outputs[0].finish_reason == "length" + assert record._outputs[0].token_ids == returned + assert record.error is None + assert record.postproc_params.postproc_args.num_prompt_tokens == len(prompt) + + +@pytest.mark.parametrize("num_prompt_tokens", [0, 7]) +def test_direct_result_preserves_explicit_prompt_count(num_prompt_tokens: int) -> None: + request = GenerationRequest( + prompt_token_ids=list(PROMPT), + sampling_params=_params(), + postproc_params=PostprocParams( + postproc_args=PostprocArgs(num_prompt_tokens=num_prompt_tokens) + ), + ).set_id(11) + result = GenerationResult(request) + assert result.postproc_params.postproc_args.num_prompt_tokens == num_prompt_tokens + + +def test_generate_async_still_sets_prompt_count() -> None: + from tensorrt_llm.executor.executor import GenerationExecutor + + executor = MagicMock() + executor.submit.side_effect = lambda request: GenerationResult(request.set_id(11)) + result = GenerationExecutor.generate_async( + executor, + prompt_token_ids=list(PROMPT), + sampling_params=_params(), + postproc_params=PostprocParams(postproc_args=PostprocArgs(num_prompt_tokens=1)), + ) + assert result.postproc_params.postproc_args.num_prompt_tokens == len(PROMPT)