diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index 7b384d2b052d..0da0225976cf 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, @@ -849,6 +850,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..692aee045a42 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/finish_reasons.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/finish_reasons.py @@ -249,7 +249,13 @@ 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) + # The native original boundary survives pause/recompute. + min_lengths = [ + 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) 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 +264,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 +543,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 +623,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 +635,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 +644,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 +657,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 +668,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 +682,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..3d3b02583ad5 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.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/tensorrt_llm/executor/result.py b/tensorrt_llm/executor/result.py index 43bcee12ff40..ae8e931a13ea 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: @@ -510,12 +511,41 @@ 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)) 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): + 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) + and num_generated_tokens >= len(stop_ids)): + 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 @@ -538,6 +568,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: @@ -1134,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/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 49b28ebc679a..a4aca06a351b 100644 --- a/tensorrt_llm/sampling_params.py +++ b/tensorrt_llm/sampling_params.py @@ -103,6 +103,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" """ @@ -290,6 +302,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 @@ -636,22 +650,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) @@ -660,6 +689,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 267b77ce2d1d..ceeb6f8578b1 100644 --- a/tensorrt_llm/serve/anthropic_adapter.py +++ b/tensorrt_llm/serve/anthropic_adapter.py @@ -1199,6 +1199,15 @@ async def read_all() -> None: await asyncio.gather(reader, return_exceptions=True) +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]: @@ -1251,8 +1260,17 @@ def until_ping() -> Optional[float]: 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 @@ -1271,8 +1289,9 @@ def until_ping() -> Optional[float]: ) 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/deepseek_v41.py b/tensorrt_llm/serve/extensions/deepseek_v41.py index bfe71e5ef500..4409ad5524e6 100644 --- a/tensorrt_llm/serve/extensions/deepseek_v41.py +++ b/tensorrt_llm/serve/extensions/deepseek_v41.py @@ -12,22 +12,41 @@ # 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. +"""Serving extension for DeepSeek-V4.1 checkpoints (``model_type == "deepseek_v41"``). -The checkpoint's prompt renderer (``tensorrt_llm/tokenizer/deepseek_v41``) -reads tool namespaces itself, so a tool declared inside a namespace is handed -to it in that form rather than flattened. +Preserve namespaces for the checkpoint prompt renderer and declare token runs +that mark a generation that has broken down. """ -from typing import Any, Dict +from typing import Any, Dict, Optional from tensorrt_llm.serve.openai_protocol import NamespacedFunctionToolParam -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, +) @register_serving_extension(model_types=("deepseek_v41", "deepseek_v41_text")) class DeepseekV41ServingExtension(ServingExtension): - """Tool namespaces in the form the DeepSeek-V4.1 checkpoint takes them.""" + """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 def serialize_tool(self, tool) -> Dict[str, Any]: """Dump a namespaced tool with its namespace, any other tool as usual. 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 2c3fffad6c10..43c9bcef101a 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, @@ -84,7 +86,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 ( @@ -144,7 +147,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) @@ -323,6 +328,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. @@ -1148,14 +1175,51 @@ def postproc_worker_enabled(self) -> bool: return True if self.generator.args.num_postprocess_workers > 0 else False + @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 @@ -2094,6 +2158,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 @@ -2123,6 +2193,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, @@ -3013,15 +3084,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]]): @@ -3076,6 +3139,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, @@ -3199,6 +3263,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 @@ -3250,6 +3320,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 " @@ -3481,6 +3552,10 @@ async def create_streaming_generator(promise: RequestOutput, reasoning_parser_name=self.generator.args.reasoning_parser, tool_parser_name=self.tool_parser, has_tools=bool(request.tools)) + # 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 77e4cbdba451..be7f1fa6a380 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -224,6 +224,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 @@ -258,6 +259,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 - unittest/llmapi/apps/test_kimi_serve_extensions.py 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..ccd73d586080 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,109 @@ 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, + 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 + + +@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 new file mode 100644 index 000000000000..5300deaed377 --- /dev/null +++ b/tests/unittest/executor/test_degenerate_token_runs.py @@ -0,0 +1,494 @@ +# 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 copy import deepcopy +from types import MethodType, SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +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.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, + 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.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: + 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("") + + +@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), + 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]) + + +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 + + +@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 + + +@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) 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 0df03139c884..e0778162568c 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, @@ -784,6 +786,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 ( @@ -895,3 +898,93 @@ def test_deepseek_v41_keeps_a_name_its_encoder_rejects_flattened( tool = self._namespaced_tool(namespace=namespace, name=name) assert deepseek_v41.serialize_tool(tool) == tool.model_dump() + + +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)