[None][fix] serve: fail a request whose generation degenerates into a token run

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,

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",

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,14 +635,20 @@ 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
             Shape: [max_tokens, len(requests), max_beam_width]
         """
         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,
         )

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,

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

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

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

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.