[TRTLLM-17184][fix] Avoid global stop-word buffer growth in the PyTorch sampler

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -9107,6 +9107,7 @@ def _terminate_request(self, request: LlmRequest) -> None:
 
     def _free_request_resources(self, request: LlmRequest) -> None:
         """Release execution resources without removing response routing."""
+        self.sampler.release_request(request)
         self.resource_manager.free_resources(request)
         self._prefetched_request_ids.discard(request.py_request_id)
         self.disagg.forget_request(request.py_request_id)

tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py

@@ -43,6 +43,7 @@
 from ..llm_request import LlmRequest, LlmRequestState
 from .logprobs import LogProbsStore
 from .ops.flashinfer import radix_topk_op
+from .ops.triton import reorder_compact_stop_word_history
 from .sampler_common import _get_beam_width_out, _request_get_sampling_params, int_tensor
 from .sampler_features import _SideStreamCopier
 
@@ -330,14 +331,8 @@ class BeamSearchMetadata:
     pending_harvest: torch.Tensor
     predecessor_beams: torch.Tensor
     beam_idx_arange: torch.Tensor
-    stop_past_tokens: Optional[torch.Tensor] = None
-    """[max_stop_word_length, max_num_sequences, max_beam_width] int32, the
-    finish handler's rolling stop-word window (FinishReasonsHandler store).
-    The beam axis is reordered by the
-    step's predecessor beams so the handler's stop-word matching stays correct
-    across beam swaps. When None (tests without stop words), the reorder is
-    skipped — multi-token stop-word matching would then be unreliable across
-    beam swaps."""
+    stop_words_metadata: Optional[torch.Tensor] = None
+    """Per-slot compact stop-word allocations, including the rolling history."""
     cba: Optional[CBAState] = None
     """Candidate-Beams-Array state, present for every beam-search request
     regardless of its early_stopping mode; None until the first one."""
@@ -523,7 +518,7 @@ def beam_candidate_topk(
 class CBAStepResult(NamedTuple):
     """Return of ``_cba_step_math``. A NamedTuple (not a dataclass) so it is a
     valid output of the torch.compile'd fullgraph function while still naming
-    the eleven tensors the caller writes back."""
+    the tensors the caller writes back."""
 
     slot_pred: torch.Tensor
     """[bs, num_beams] int32, predecessor beam of each continuing slot."""
@@ -545,9 +540,6 @@ class CBAStepResult(NamedTuple):
     """[bs, pool_width, snap_len] float32, merged CBA pool per-token log-probs."""
     done: torch.Tensor
     """[bs] bool, per-slot beam-search termination verdict."""
-    reordered_window: Optional[torch.Tensor]
-    """Stop-word window reordered to follow the beam swap, or None when no
-    stop-word window is present."""
 
 
 def _cba_step_math(
@@ -573,7 +565,6 @@ def _cba_step_math(
     cba_lengths: torch.Tensor,
     cba_tokens: torch.Tensor,
     cba_log_probs: torch.Tensor,
-    stop_past_tokens: Optional[torch.Tensor],
     # static (graph-specializing) parameters
     beam_width_in: int,
     num_beams: int,
@@ -740,15 +731,6 @@ def _take(src: torch.Tensor) -> torch.Tensor:
     ignore_attainable = early_stopping == BeamSearchEarlyStop.TRUE
     done = pool_full & (ignore_attainable | (min_kept >= best_attainable))
 
-    # Reorder the finish handler's rolling stop-word window to follow the
-    # beam swap (matching stays correct across swaps).
-    reordered_window = None
-    if stop_past_tokens is not None:
-        window = stop_past_tokens[:, slots, :beam_width_in]
-        reordered_window = torch.gather(
-            window, 2, slot_pred.long().unsqueeze(0).expand(window.size(0), -1, -1)
-        )
-
     return CBAStepResult(
         slot_pred=slot_pred,
         slot_tok=slot_tok,
@@ -760,7 +742,6 @@ def _take(src: torch.Tensor) -> torch.Tensor:
         merged_tokens=merged_tokens,
         merged_lps=merged_lps,
         done=done,
-        reordered_window=reordered_window,
     )
 
 
@@ -939,7 +920,6 @@ def beam_search_sampling_batch_cba(
         merged_tokens,
         merged_lps,
         done,
-        reordered_window,
     ) = step_fn(
         cand_cum=cand_cum,
         cand_pred=cand_pred,
@@ -961,7 +941,6 @@ def beam_search_sampling_batch_cba(
         cba_lengths=cba.cba_lengths,
         cba_tokens=cba.cba_tokens,
         cba_log_probs=cba.cba_log_probs,
-        stop_past_tokens=args.stop_past_tokens,
         beam_width_in=beam_width_in,
         num_beams=num_beams,
         early_stopping=early_stopping,
@@ -999,9 +978,8 @@ def beam_search_sampling_batch_cba(
         torch.full_like(prev_reasons, FinishReason.LENGTH.value),
         prev_reasons,
     )
-    stop_window = args.stop_past_tokens
-    if reordered_window is not None and stop_window is not None:
-        stop_window[:, slots, :num_beams] = reordered_window
+    if args.stop_words_metadata is not None:
+        reorder_compact_stop_word_history(args.stop_words_metadata, slots, slot_pred)
 
     # --- Beam-slot state updates (same contract as beam_search_sampling_batch_cba).
     args.predecessor_beams[slots, :num_beams] = slot_pred
@@ -1559,15 +1537,15 @@ def build_metadata(
         num_requests: int,
         new_log_probs: torch.Tensor,
         end_ids_cuda: torch.Tensor,
-        past_tokens_cuda: torch.Tensor,
+        stop_words_metadata_cuda: torch.Tensor | None,
     ) -> BeamSearchMetadata:
         """Build the beam-search metadata for a beam-search group.
 
         Assembles the per-step view the CBA op consumes: the persistent store
         tensors plus the group's slot/length tensors and the host-derived
         ``max_gen_len`` bound. Mirrors ``TopPDecayHandler.build_metadata``.
 
-        ``new_log_probs``, ``end_ids_cuda`` and ``past_tokens_cuda`` are owned
+        ``new_log_probs``, ``end_ids_cuda`` and ``stop_words_metadata_cuda`` are owned
         by the log-probs and finish-reason handlers respectively and are passed
         in rather than reached for, keeping this module independent of them.
         """
@@ -1598,7 +1576,7 @@ def build_metadata(
             pending_harvest=store.pending_harvest,
             predecessor_beams=store.predecessor_beams,
             beam_idx_arange=store.beam_idx_arange,
-            stop_past_tokens=past_tokens_cuda,
+            stop_words_metadata=stop_words_metadata_cuda,
             # None unless a beam-search request has been
             # admitted; the CBA tensors are not allocated before that.
             cba=None