diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index f80b3aaed8c1..a0423b3a926a 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/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) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py b/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py index 1f7f02755e56..ad82cdb02a41 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py +++ b/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,7 +1537,7 @@ 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. @@ -1567,7 +1545,7 @@ def build_metadata( 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 diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/finish_reasons.py b/tensorrt_llm/_torch/pyexecutor/sampler/finish_reasons.py index f538f8a9ac10..aeba2cd4043b 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/finish_reasons.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/finish_reasons.py @@ -16,30 +16,36 @@ Owns the stop criteria evaluated on the device each step -- end-id match, max-length, and stop words -- along with the per-slot buffers they need -(stop-word rules, the past-token window) and their resize/refresh logic. +(packed stop-word rules and per-request past-token windows). ``TorchSampler`` holds one :class:`FinishReasonsHandler` and drives it -through request admission and the per-step write. +through request admission, release, and the per-step write. """ from dataclasses import dataclass +from typing import NamedTuple import torch from tensorrt_llm._utils import nvtx_range, prefer_pinned from tensorrt_llm.bindings.executor import FinishReason -from ...utils import torch_multi_arange from ..llm_request import LlmRequest -from .ops.triton import MAX_FUSED_ELEMENTS_PER_REQUEST, fused_write_finish_reasons +from .ops.triton import ( + MAX_FUSED_ELEMENTS_PER_REQUEST, + compact_stop_words, + fused_write_finish_reasons, +) from .sampler_common import int_tensor __all__ = ["FinishReasonsHandler"] -class FinishReasonsHandler: - _EMPTY_STOP_WORD_TOKEN_ID: int = -2 - _PAD_STOP_WORD_TOKEN_ID: int = -1 +class _StopWordAllocation(NamedTuple): + request: LlmRequest + tensor: torch.Tensor + +class FinishReasonsHandler: @dataclass(kw_only=True) class _FinishReasonsStore: """Auxiliary data structures used for finish reasons handling.""" @@ -58,12 +64,10 @@ class _FinishReasonsStore: end_ids_cuda: torch.Tensor """Shape: batch_size Usage: Stores the end ids for each request""" - stop_words_cuda: torch.Tensor - """Shape: [max_num_stop_words, max_stop_word_length, batch_size] - Usage: Stores the stop words for each request as a padded tensor.""" - past_tokens_cuda: torch.Tensor - """Shape: [max_stop_word_length,batch_size, beam_width] - Usage: Stores the last max_stop_word_length tokens for each beam.""" + stop_words_metadata_cuda: torch.Tensor + """Shape: [batch_size, 5], int64. + Per-slot allocation pointer, rule count, history offset, window length, + and beam stride. The handler owns the pointed-to allocations.""" max_stop_word_lengths_host: torch.Tensor """Shape: [batch_size] Usage: Stores the size of the longest stop word for each request.""" @@ -74,25 +78,30 @@ class _FinishReasonsStore: def __init__( self, *, - max_stop_word_length: int, - max_num_stop_words: int, max_num_sequences: int, max_beam_width: int, max_tokens: int, max_seq_len: int, ): self._update_sizes( - max_stop_word_length=max_stop_word_length, - max_num_stop_words=max_num_stop_words, max_num_sequences=max_num_sequences, max_beam_width=max_beam_width, max_tokens=max_tokens, max_seq_len=max_seq_len, ) + self._stop_word_allocations: dict[int, _StopWordAllocation] = {} + self._stop_word_history_slots: set[int] = set() + self._stop_word_allocation_generation = 0 + self._stop_word_stream_generations: dict[torch.cuda.Stream, int] = {} self._setup_store() self._setup_helper_tensors() self._temp_data: FinishReasonsHandler._TemporaryData = self._TemporaryData() + @property + def has_stop_word_history(self) -> bool: + """Whether any retained request needs a multi-token history window.""" + return bool(self._stop_word_history_slots) + @property def _use_speculative_decoding(self) -> bool: return self._max_tokens > 1 @@ -108,8 +117,6 @@ def new_end_ids(self) -> list[int]: def _update_sizes( self, *, - max_stop_word_length: int, - max_num_stop_words: int, max_num_sequences: int, max_beam_width: int, max_tokens: int, @@ -120,25 +127,13 @@ def _update_sizes( Sets member variables to store the current sizes. These sizes are used to initialize the buffer tensors. """ - self._max_stop_word_length: int = max_stop_word_length - self._max_num_stop_words: int = max_num_stop_words self._max_num_sequences: int = max_num_sequences self._max_beam_width: int = max_beam_width self._max_tokens: int = max_tokens self._max_seq_len: int = max_seq_len - self._stop_words_shape: tuple[int, int, int] = ( - self._max_num_stop_words, - self._max_stop_word_length, - self._max_num_sequences, - ) self._fused_tile_fits: bool = ( self._max_tokens * self._max_beam_width <= MAX_FUSED_ELEMENTS_PER_REQUEST ) - self._past_tokens_shape: tuple[int, int, int] = ( - self._max_stop_word_length - 1 + self._max_tokens, - self._max_num_sequences, - self._max_beam_width, - ) def _setup_store(self) -> None: """Sets up the store for the finish reasons handler by initializing all buffer tensors.""" @@ -147,9 +142,10 @@ def _setup_store(self) -> None: ) max_lengths_cuda = int_tensor((self._max_num_sequences,)) end_ids_cuda = int_tensor((self._max_num_sequences,)) - stop_words_cuda = int_tensor(self._stop_words_shape) - past_tokens_cuda = int_tensor(self._past_tokens_shape) - max_stop_word_lengths_host = torch.empty( + stop_words_metadata_cuda = torch.zeros( + (self._max_num_sequences, 5), dtype=torch.int64, device="cuda" + ) + max_stop_word_lengths_host = torch.zeros( self._max_num_sequences, device="cpu", dtype=torch.int32 ) num_accepted_draft_tokens_host = torch.empty( @@ -159,93 +155,45 @@ def _setup_store(self) -> None: finish_reasons_cuda=finish_reasons_cuda, max_lengths_cuda=max_lengths_cuda, end_ids_cuda=end_ids_cuda, - stop_words_cuda=stop_words_cuda, - past_tokens_cuda=past_tokens_cuda, + stop_words_metadata_cuda=stop_words_metadata_cuda, max_stop_word_lengths_host=max_stop_word_lengths_host, num_accepted_draft_tokens_host=num_accepted_draft_tokens_host, ) def _setup_helper_tensors(self) -> None: - # Helper tensors for finish_reasons: - """Preallocate buffer needed for torch.nonzero_static(..., out=finish_reasons_nonzero_static_buffer). - See `def _write_reason`.""" - # setup local buffer for max tokens checking + """Preallocate the token offsets used by the maximum-length check.""" self._max_tokens_offset_cuda: torch.Tensor = torch.arange( 1, self._max_tokens + 1, device="cuda", dtype=torch.int32 ).view(-1, 1, 1) - self._stop_words_index_offset_cuda: torch.Tensor = torch.arange( - max(0, self._max_stop_word_length - 1), device="cuda" - ).unsqueeze(1) - - self._past_token_buffer_cuda: torch.Tensor = torch.empty( - self._past_tokens_shape, device="cuda", dtype=torch.int32 - ) - starts = torch.arange(self._max_tokens, device="cuda") - ends = starts + self._max_stop_word_length - self._multi_arange_indexing: torch.Tensor = torch_multi_arange( - ends=ends, - starts=starts, - output_length=self._max_tokens * self._max_stop_word_length, - ) - - def _resize_stop_word_buffers(self) -> None: - self._stop_words_index_offset_cuda = torch.arange( - max(0, self._max_stop_word_length - 1), device="cuda" - ).unsqueeze(1) - - self._past_tokens_shape = ( - self._max_stop_word_length - 1 + self._max_tokens, - self._max_num_sequences, - self._max_beam_width, - ) - self._stop_words_shape = ( - self._max_num_stop_words, - self._max_stop_word_length, - self._max_num_sequences, - ) - self._past_token_buffer_cuda = torch.empty( - self._past_tokens_shape, device="cuda", dtype=torch.int32 - ) - starts = torch.arange(self._max_tokens, device="cuda") - ends = starts + self._max_stop_word_length - self._multi_arange_indexing = torch_multi_arange( - ends=ends, - starts=starts, - output_length=self._max_tokens * self._max_stop_word_length, - ) - # resize the stop words buffer if necessary - # if the sizes are constant, this does nothing - store = self.store - _ = store.stop_words_cuda.resize_(self._stop_words_shape) - _ = store.past_tokens_cuda.resize_(self._past_tokens_shape) - @dataclass(kw_only=True) class _TemporaryData: """Data structure to store the temporary data during setup_sampler_step for new requests""" def __init__(self) -> None: - # list of device tensors - self.stop_words_cuda_list: list[torch.Tensor] = [] - self.past_tokens_cuda_list: list[torch.Tensor] = [] - # list of integers - self.stop_word_seq_slots: list[int] = [] + self.metadata: list[list[int]] = [] self.max_lens: list[int] = [] self.end_ids: list[int] = [] self.max_stop_word_lengths: list[int] = [] - # integers - self.total_max_length: int = 0 - self.total_max_num_stop_words: int = 0 def clear(self) -> None: - self.stop_words_cuda_list = [] - self.past_tokens_cuda_list = [] - self.stop_word_seq_slots = [] - self.max_lens = [] - self.end_ids = [] - self.max_stop_word_lengths = [] - self.total_max_length = 0 - self.total_max_num_stop_words = 0 + self.metadata.clear() + self.max_lens.clear() + self.end_ids.clear() + self.max_stop_word_lengths.clear() + + def release_request(self, request: LlmRequest) -> None: + """Release a request's allocation before its sequence slot is returned.""" + slot = request.py_seq_slot + if slot is None: + return + allocation = self._stop_word_allocations.get(slot) + if allocation is None or allocation.request is not request: + return + del self._stop_word_allocations[slot] + self._stop_word_history_slots.discard(slot) + # Outstanding kernels may still read the metadata on another stream. + # Admission overwrites host and device metadata before the slot is reused. def setup_new_request_handling(self) -> None: """Setup the new request handling for the finish reasons handler @@ -284,23 +232,61 @@ def prepare_for_new_request(self, request: LlmRequest) -> None: end_id = -1 self._temp_data.end_ids.append(end_id) - if (stop_words_list := request.py_stop_words_list) is not None: - assert (seq_slot := request.py_seq_slot) is not None - self._temp_data.stop_word_seq_slots.append(seq_slot) - extracted_stop_words_cuda, max_length, num_stop_words = self._extract_stop_words( - stop_words_list - ) - self._temp_data.stop_words_cuda_list.append(extracted_stop_words_cuda) - self._temp_data.past_tokens_cuda_list.append(self._get_past_tokens(request)) - self._temp_data.total_max_length = max(self._temp_data.total_max_length, max_length) - self._temp_data.total_max_num_stop_words = max( - self._temp_data.total_max_num_stop_words, num_stop_words - ) - self._temp_data.max_stop_word_lengths.append(max_length) - else: - # max stop word length is used to determine if a request has stop words - # explicitly set it to 0 here to avoid stale data from previous requests + assert (seq_slot := request.py_seq_slot) is not None + self._stop_word_allocations.pop(seq_slot, None) + self._stop_word_history_slots.discard(seq_slot) + words = request.py_stop_words_list or [] + if not words: + self._temp_data.metadata.append([0, 0, 0, 0, self._max_beam_width]) self._temp_data.max_stop_word_lengths.append(0) + return + if any(not word for word in words): + raise ValueError("Stop words must contain at least one token") + + max_length = max(map(len, words)) + offsets = [0] + flat_tokens: list[int] = [] + for word in words: + flat_tokens.extend(word) + offsets.append(len(flat_tokens)) + history_offset = len(offsets) + len(flat_tokens) + 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( + history_offset + 2 * history_size, + dtype=torch.int32, + device="cpu", + pin_memory=prefer_pinned(), + ) + host[:history_offset] = torch.tensor(offsets + flat_tokens, dtype=torch.int32) + if window_length: + history = host[history_offset : history_offset + history_size].view( + window_length, self._max_beam_width + ) + # The first sampling step shifts by one before appending new tokens. + for beam, tokens in enumerate(request.get_tokens()[: self._max_beam_width]): + count = min(max_length - 1, len(tokens)) + if count: + history[max_length - count : max_length, beam] = torch.tensor( + tokens[-count:], dtype=torch.int32 + ) + allocation = host.to("cuda", non_blocking=True) + self._stop_word_allocations[seq_slot] = _StopWordAllocation( + request=request, tensor=allocation + ) + self._stop_word_allocation_generation += 1 + if window_length: + self._stop_word_history_slots.add(seq_slot) + self._temp_data.metadata.append( + [ + allocation.data_ptr(), + len(words), + history_offset, + window_length, + self._max_beam_width, + ] + ) + self._temp_data.max_stop_word_lengths.append(max_length) def update_for_new_request( self, @@ -309,28 +295,8 @@ def update_for_new_request( max_lengths_cuda: torch.Tensor, end_ids_cuda: torch.Tensor, seq_slots_host: torch.Tensor, - all_sampling_requests: list[LlmRequest], ) -> None: - """Update tensors of this store with the new request data. - - If stop words are present, also update the stop words buffers. - If the new stop words exceed either the current max_num_stop_words or max_stop_word_length values, - a resize of the stop words buffers is triggered. If a resize is necessary, all requests in the batch - need to be re-processed. - - Args: - seq_slots_cuda_long: The sequence slots of the processed requests, as int64 - CUDA indices (required by ``index_copy_``). Shape: [len(requests)] - max_lengths_cuda: The maximum lengths for each request. - Shape: [len(requests)] - end_ids_cuda: The end ids for each request. - Shape: [len(requests)] - seq_slots_host: The sequence slots of the processed requests. Used for accessing host buffers. - Shape: [len(requests)] - all_sampling_requests: If a resize of the stop words related buffers is necessary, all sampling requests - need to be re-processed. - """ - + """Publish new-request metadata without resizing existing requests.""" temp_data = self._temp_data store = self.store store.max_lengths_cuda.index_copy_(0, seq_slots_cuda_long, max_lengths_cuda) @@ -339,196 +305,27 @@ def update_for_new_request( temp_data.max_stop_word_lengths, device="cpu", dtype=torch.int32 ) - # Handle stop words only if any new ones are added - if temp_data.stop_word_seq_slots: - self._update_stop_words_buffer( - all_sampling_requests, - temp_data.total_max_length, - temp_data.total_max_num_stop_words, - temp_data.stop_words_cuda_list, - temp_data.past_tokens_cuda_list, - temp_data.stop_word_seq_slots, - ) - - def _maybe_resize_stop_words_buffer( - self, total_max_length: int, total_max_num_stop_words: int - ) -> bool: - """Checks if the stop words buffer needs to be resized and resizes it if necessary - - If the total maximum length or number of stop words exceeds the current maximum values, - the stop words buffer is resized to the new maximum values. - - Args: - total_max_length: The maximum length of the stop words in this batch. - total_max_num_stop_words: The maximum number of stop words of a request in this batch. - Returns: - True if the stop words buffer needs to be resized, False otherwise. - """ - if ( - total_max_length > self._max_stop_word_length - or total_max_num_stop_words > self._max_num_stop_words - ): - self._max_stop_word_length = max(total_max_length, self._max_stop_word_length) - self._max_num_stop_words = max(total_max_num_stop_words, self._max_num_stop_words) - self._resize_stop_word_buffers() - return True - return False - - def _reprocess_stop_words_buffer( - self, requests: list[LlmRequest] - ) -> tuple[list[torch.Tensor], list[torch.Tensor], list[int]]: - """Reprocesses the stop words buffer with the new maximum values - - If a resize of the stop words buffer is necessary, all requests in the batch need to be re-processed. - - Args: - requests: The requests to reprocess the stop words buffer for. - Returns: - stop_words_cuda_list: A list of device tensors containing the stop words per request with stop words. - past_tokens_cuda_list: A list of device tensors containing the past tokens per request with stop words. - stop_word_seq_slots: A list of sequence slot indices (int) per request with stop words. - """ - stop_words_cuda_list: list[torch.Tensor] = [] - past_tokens_cuda_list: list[torch.Tensor] = [] - stop_word_seq_slots: list[int] = [] - # Rerun with the new size. Set the stop words and past tokens for all the requests. - for request in requests: - if (stop_words_list := request.py_stop_words_list) is not None: - extracted_stop_words_cuda, _, _ = self._extract_stop_words(stop_words_list) - assert (seq_slot := request.py_seq_slot) is not None - stop_word_seq_slots.append(seq_slot) - stop_words_cuda_list.append(extracted_stop_words_cuda) - past_tokens_cuda_list.append(self._get_past_tokens(request)) - return stop_words_cuda_list, past_tokens_cuda_list, stop_word_seq_slots - - def _update_stop_words_buffer( - self, - all_sampling_requests: list[LlmRequest], - total_max_length: int, - total_max_num_stop_words: int, - stop_words_cuda_list: list[torch.Tensor], - past_tokens_cuda_list: list[torch.Tensor], - stop_word_seq_slots: list[int], - ) -> None: - """Updates the stop words buffer with the new maximum values - - Args: - all_sampling_requests: If a resize of the stop words related buffers is necessary, all sampling requests - need to be re-processed. - total_max_length: The maximum length of the stop words in this batch. - total_max_num_stop_words: The maximum number of stop words of a request in this batch. - stop_words_cuda_list: A list of device tensors containing the stop words per request with stop words. - past_tokens_cuda_list: A list of device tensors containing the past tokens per request with stop words. - stop_word_seq_slots: Sequence slot index (int) per request with stop words; - same order as the lists above. - """ - # Potentially resize the buffers and update - # stop_words, past_tokens and stop_word_seq_slots - # In case of a resize all requests in the batch need to be re-processed - if self._maybe_resize_stop_words_buffer(total_max_length, total_max_num_stop_words): - stop_words_cuda_list, past_tokens_cuda_list, stop_word_seq_slots = ( - self._reprocess_stop_words_buffer(all_sampling_requests) + store.num_accepted_draft_tokens_host[seq_slots_host] = 0 + metadata = ( + torch.tensor( + temp_data.metadata, dtype=torch.int64, device="cpu", pin_memory=prefer_pinned() ) - - # Host Tensor for host access of self.store.num_accepted_draft_tokens - stop_word_seq_slots_tensor_host = torch.tensor( - stop_word_seq_slots, device="cpu", dtype=torch.int32, pin_memory=prefer_pinned() - ) - # Device Tensor for device access of self.store.stop_words and self.store.past_tokens - stop_word_seq_slots_tensor_cuda = stop_word_seq_slots_tensor_host.to( - device="cuda", non_blocking=True - ) - # stop_word_seq_slots x max_num_stop_words x max_stop_word_length - stop_words_cuda_tensor = torch.stack(stop_words_cuda_list) - # stop_word_seq_slots x max_stop_word_length x beam_width - past_tokens_cuda_tensor = torch.stack(past_tokens_cuda_list) - - store = self.store - # Reset the accepted tokens buffer for the stop word sequence slots - store.num_accepted_draft_tokens_host[stop_word_seq_slots_tensor_host] = 0 - - store.stop_words_cuda[..., stop_word_seq_slots_tensor_cuda] = ( - stop_words_cuda_tensor.permute(1, 2, 0) - ) - # Past tokens will be shifted by 1 to the left on their first sampling iteration - # We need to consider this here. - store.past_tokens_cuda[1 : self._max_stop_word_length, stop_word_seq_slots_tensor_cuda] = ( - past_tokens_cuda_tensor.permute(1, 0, 2) + .reshape(-1, 5) + .to("cuda", non_blocking=True) ) + store.stop_words_metadata_cuda.index_copy_(0, seq_slots_cuda_long, metadata) - def _extract_stop_words( - self, stop_words_list: list[list[int]] - ) -> tuple[torch.Tensor, int, int]: - """Extract the stop words and size information from the stop words list - - Processes the stop words list and stores the stop words in a padded device tensor. - Stop words shorter than FinishReasonsHandler.max_stop_word_length - are padded with _PAD_STOP_WORD_TOKEN_ID to max_stop_word_length. - Unused stop word slots are padded with _EMPTY_STOP_WORD_TOKEN_ID. - This function additionally returns the maximum stop word length and the number of stop words - in the processed stop words list. - - - Args: - stop_words_list: A list of stop sequences, each a list of token ids. - - Returns: - stop_words: A padded device tensor containing the stop words - Shape: [max_num_stop_words, max_stop_word_length] - max_stop_word_length: The maximum stop word length in the stop words list - num_stop_words: The number of stop words in the stop words list - """ - stop_words_host = torch.empty( - self._max_num_stop_words, - self._max_stop_word_length, - device="cpu", - dtype=torch.int32, - ) - _ = stop_words_host.fill_(self._EMPTY_STOP_WORD_TOKEN_ID) - max_stop_word_length = 0 - num_stop_words = 0 - for idx, word in enumerate(stop_words_list): - length = len(word) - max_stop_word_length = max(max_stop_word_length, length) - num_stop_words += 1 - # skip processing if either the length or the index is greater than the current max values. - # These will be updated outside this function. - if length > self._max_stop_word_length or idx >= self._max_num_stop_words: - continue - stop_words_host[idx, -length:] = torch.tensor(word, dtype=torch.int32) - stop_words_host[idx, :-length] = self._PAD_STOP_WORD_TOKEN_ID - return ( - stop_words_host.to("cuda", non_blocking=True), - max_stop_word_length, - num_stop_words, - ) - - def _get_past_tokens(self, request: LlmRequest) -> torch.Tensor: - """Get the past tokens from the request and return the past tokens device tensor - - Args: - request: The request to get the past tokens for - - Returns: - past_tokens: The past tokens device tensor - Shape: [max_stop_word_length - 1, max_beam_width] - """ - past_tokens_host = torch.zeros( - max(0, self._max_stop_word_length - 1), - self._max_beam_width, - device="cpu", - dtype=torch.int32, - pin_memory=prefer_pinned(), - ) - tokens = request.get_tokens() - for beam_idx in range(self._max_beam_width): - max_len = min(past_tokens_host.shape[0], len(tokens[beam_idx])) - past_tokens_host[past_tokens_host.shape[0] - max_len :, beam_idx] = torch.tensor( - tokens[beam_idx][len(tokens[beam_idx]) - max_len :], - device="cpu", - dtype=torch.int32, - ) - return past_tokens_host.to("cuda", non_blocking=True) + def record_stop_word_stream(self) -> None: + """Register streams that consume the indirectly referenced allocations.""" + stream = torch.cuda.current_stream() + generation = self._stop_word_allocation_generation + if self._stop_word_stream_generations.get(stream) == generation: + return + for allocation in self._stop_word_allocations.values(): + allocation.tensor.record_stream(stream) + # The allocator records completion events when storage is released; + # repeated uses on an already registered stream need no registration. + self._stop_word_stream_generations[stream] = generation def write_finish_reasons( self, @@ -834,67 +631,10 @@ def _are_stop_words( ends with a stop word, False otherwise Shape: [max_tokens, len(requests), max_beam_width] """ - store = self.store - # num_words, len_words, batch_size - # unsqueeze the beam_width dimension to match the past tokens tensor - stop_words = ( - store.stop_words_cuda[..., seq_slots] - .unsqueeze(3) - .expand(-1, -1, -1, self._max_beam_width) - ) - # Get the past tokens - # num_steps, batch_size, beam_width - past_tokens_batch = store.past_tokens_cuda[:, seq_slots] - # Shift the past tokens to the left by the number of accepted draft tokens - - full_tokens = self._past_token_buffer_cuda[:, : seq_slots.shape[0]] - - index_tensor = ( - (self._stop_words_index_offset_cuda + num_accepted_tokens) - .unsqueeze(2) - .expand(-1, seq_slots.shape[0], self._max_beam_width) - ) - _ = torch.gather( - past_tokens_batch, - dim=0, - index=index_tensor, - out=full_tokens[: index_tensor.shape[0]], - ) - # Fill in the new tokens at the end of the past tokens buffer - full_tokens[-self._max_tokens :] = tokens - # short words are padded with _PAD_STOP_WORD_TOKEN_ID, so we need to mask them - mask = stop_words == self._PAD_STOP_WORD_TOKEN_ID - matches = torch.empty( - ( - self._max_tokens, - stop_words.shape[0], - stop_words.shape[1], - stop_words.shape[2], - stop_words.shape[3], - ), - device="cuda", - dtype=torch.bool, - ) - - # Get the comparison sequence for each step - full_tokens_for_match = full_tokens[self._multi_arange_indexing].view( - self._max_tokens, - 1, # Unsqueeze on dimension 1 to match the num_stop_words dimension of stop words - self._max_stop_word_length, - seq_slots.shape[0], - self._max_beam_width, + self.record_stop_word_stream() + return compact_stop_words( + self.store.stop_words_metadata_cuda, seq_slots, tokens, num_accepted_tokens ) - # Unsqueeze on dimension 0 to match the max_tokens dimension of full tokens - stop_words_for_match = stop_words.unsqueeze(0) - _ = torch.eq(full_tokens_for_match, stop_words_for_match, out=matches) - # Mask the padding tokens - _ = matches.masked_fill_(mask.unsqueeze(0).expand(self._max_tokens, -1, -1, -1, -1), True) - # Update the past tokens storage for the next iteration - store.past_tokens_cuda[:, seq_slots] = full_tokens - # Return the result - word_len_dim = 2 - num_words_dim = 1 - return torch.any(matches.all(dim=word_len_dim), dim=num_words_dim) @nvtx_range("_are_stop_words_single_token") def _are_stop_words_single_token( @@ -916,14 +656,11 @@ def _are_stop_words_single_token( ends with a stop word, False otherwise Shape: [max_tokens, len(requests), max_beam_width] """ - per_step = torch.zeros( - (self._max_tokens, seq_slots.shape[0], self._max_beam_width), - dtype=torch.bool, - device="cuda", + self.record_stop_word_stream() + return compact_stop_words( + self.store.stop_words_metadata_cuda, + seq_slots, + tokens, + _num_accepted_tokens, + single_token_only=True, ) - # num_words, 1, batch_size - stop_words = self.store.stop_words_cuda[:, -1:, seq_slots].unsqueeze(3) - full_tokens = tokens.unsqueeze(0) - matches = full_tokens == stop_words - _ = torch.any(matches, dim=0, out=per_step) - return per_step diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/ops/triton.py b/tensorrt_llm/_torch/pyexecutor/sampler/ops/triton.py index 15dd01d2a002..06d0c9e50299 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/ops/triton.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/ops/triton.py @@ -2,8 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 """Triton sampling kernels. -Every op writes into caller-allocated tensors and launches one program per -request, so the launch shape depends only on ``max_tokens`` and the beam width. +Finish-reason and compact stop-word operations for the sampler. """ import functools @@ -12,7 +11,12 @@ import triton import triton.language as tl -__all__ = ["MAX_FUSED_ELEMENTS_PER_REQUEST", "fused_write_finish_reasons"] +__all__ = [ + "MAX_FUSED_ELEMENTS_PER_REQUEST", + "fused_write_finish_reasons", + "compact_stop_words", + "reorder_compact_stop_word_history", +] # One Triton program owns all token/beam elements for one request. MAX_FUSED_ELEMENTS_PER_REQUEST = 1024 @@ -119,3 +123,232 @@ def fused_write_finish_reasons( BLOCK=block, num_warps=num_warps, ) + + +@triton.jit +def _compact_stop_history_kernel( + metadata, + slots, + tokens, + accepted, + predecessors, + token_stride_step, + token_stride_request, + token_stride_beam, + predecessor_stride, + MAX_TOKENS: tl.constexpr, + ACCEPTED_TENSOR: tl.constexpr, + MODE: tl.constexpr, + BLOCK: tl.constexpr, +): + request = tl.program_id(0) + beam = tl.program_id(1) + slot = tl.load(slots + request) + row = metadata + slot * 5 + window = tl.load(row + 3) + if window > 0: + base = tl.load(row).to(tl.pointer_type(tl.int32)) + history = base + tl.load(row + 2) + width = tl.load(row + 4) + scratch = history + window * width + if MODE == 0: + if ACCEPTED_TENSOR: + shift = tl.load(accepted + request) + else: + shift = accepted + if MODE == 2: + parent = tl.load(predecessors + request * predecessor_stride + beam) + for start in range(tl.cdiv(window, BLOCK)): + position = start * BLOCK + tl.arange(0, BLOCK) + valid = position < window + if MODE == 0: + retained = window - MAX_TOKENS + old = tl.load( + history + (position + shift) * width + beam, + valid & (position < retained), + other=0, + ) + new = tl.load( + tokens + + (position - retained) * token_stride_step + + request * token_stride_request + + beam * token_stride_beam, + valid & (position >= retained), + other=0, + ) + value = tl.where(position < retained, old, new) + elif MODE == 1: + value = tl.load(scratch + position * width + beam, valid, other=0) + else: + value = tl.load(history + position * width + parent, valid, other=0) + if MODE == 1: + tl.store(history + position * width + beam, value, valid) + else: + tl.store(scratch + position * width + beam, value, valid) + + +@triton.jit +def _compact_stop_match_kernel( + metadata, + slots, + tokens, + output, + token_stride_step, + token_stride_request, + token_stride_beam, + REQUESTS, + MAX_TOKENS: tl.constexpr, + BEAMS: tl.constexpr, + BLOCK_RULES: tl.constexpr, +): + request = tl.program_id(0) + step_beam = tl.program_id(1) + step = step_beam // BEAMS + beam = step_beam % BEAMS + slot = tl.load(slots + request) + row = metadata + slot * 5 + base = tl.load(row).to(tl.pointer_type(tl.int32)) + rule_count = tl.load(row + 1) + window = tl.load(row + 3) + width = tl.load(row + 4) + scratch = base + tl.load(row + 2) + window * width + new_token = tl.load( + tokens + + step * token_stride_step + + request * token_stride_request + + beam * token_stride_beam + ) + found = tl.full((), False, tl.int1) + for block in range(tl.cdiv(rule_count, BLOCK_RULES)): + rule = block * BLOCK_RULES + tl.arange(0, BLOCK_RULES) + valid = rule < rule_count + begin = tl.load(base + rule, valid, other=0) + end = tl.load(base + rule + 1, valid, other=0) + length = end - begin + matches = valid & (length > 0) + for offset in range(tl.max(length, 0)): + active = valid & (offset < length) + expected = tl.load(base + rule_count + 1 + begin + offset, active, other=0) + if window > 0: + position = window - MAX_TOKENS + step - length + 1 + offset + actual = tl.load(scratch + position * width + beam, active, other=0) + else: + actual = tl.full((BLOCK_RULES,), 0, tl.int32) + new_token + matches = matches & (~active | (expected == actual)) + found = found | (tl.sum(matches.to(tl.int32), 0) > 0) + tl.store(output + (step * REQUESTS + request) * BEAMS + beam, found) + + +def compact_stop_words( + metadata: torch.Tensor, + slots: torch.Tensor, + tokens: torch.Tensor, + accepted: torch.Tensor | int, + *, + single_token_only: bool = False, +) -> 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. + + Args: + metadata: Contiguous CUDA int64 tensor of shape [max_num_sequences, 5]. + Columns are allocation address, rule count, history offset in int32 + elements, history window length, and allocated beam width. + slots: Contiguous CUDA int32 or int64 tensor of shape [requests], mapping + batch requests to metadata rows. + tokens: CUDA int32 or int64 tensor of shape [steps, requests, beams]. + All three strides are passed to the kernels; contiguity is not required. + accepted: Contiguous CUDA int32 or int64 tensor of shape [requests], or + a Python int shared by all requests. Counts accepted draft tokens plus + 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. + + Returns: + CUDA bool tensor of shape [steps, requests, beams], true where a sequence + ends with a stop word. + """ + steps, requests, beams = tokens.shape + result = torch.empty(tokens.shape, device=tokens.device, dtype=torch.bool) + if requests == 0: + return result + if not single_token_only: + _compact_stop_history_kernel[(requests, beams)]( + metadata, + slots, + tokens, + accepted, + slots, + *tokens.stride(), + 0, + steps, + isinstance(accepted, torch.Tensor), + 0, + 128, + ) + _compact_stop_match_kernel[(requests, steps * beams)]( + metadata, + slots, + tokens, + result, + *tokens.stride(), + requests, + steps, + beams, + 32, + ) + if not single_token_only: + _compact_stop_history_kernel[(requests, beams)]( + metadata, + slots, + tokens, + accepted, + slots, + *tokens.stride(), + 0, + steps, + isinstance(accepted, torch.Tensor), + 1, + 128, + ) + return result + + +def reorder_compact_stop_word_history( + metadata: torch.Tensor, slots: torch.Tensor, predecessors: torch.Tensor +) -> None: + """Reorder packed histories through scratch storage, including duplicate parents. + + Args: + metadata: Contiguous CUDA int64 tensor of shape [max_num_sequences, 5], + with the column layout described in :func:`compact_stop_words`. + The caller retains the pointed-to allocations and records their use + on the current stream. + slots: Contiguous CUDA int32 or int64 tensor of shape [requests], mapping + batch requests to metadata rows. + predecessors: CUDA int32 or int64 tensor of shape [requests, beams], + containing each beam's parent index. The last dimension must have + unit stride because only stride(0) is passed to the kernel. + """ + requests, beams = predecessors.shape + if requests == 0: + return + for mode in (2, 1): + _compact_stop_history_kernel[(requests, beams)]( + metadata, + slots, + predecessors, + 1, + predecessors, + 0, + 0, + 0, + predecessors.stride(0), + 1, + False, + mode, + 128, + ) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index 1725fa49228e..abd35e1afbf2 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -27,7 +27,7 @@ from abc import ABC, abstractmethod from collections.abc import Hashable, Iterable from dataclasses import dataclass, field -from typing import Any, Dict, Final, Generic, List, Optional, TypeAlias, TypeVar, cast +from typing import Any, Dict, Generic, List, Optional, TypeAlias, TypeVar, cast import torch @@ -180,6 +180,9 @@ class Sampler(ABC, Generic[GenericSampleState]): def setup_sampler_step(self, scheduled_requests: ScheduledRequests) -> None: pass + def release_request(self, request: LlmRequest) -> None: + """Release sampler-owned state before the request's slot is returned.""" + def get_cache_indirection(self) -> torch.Tensor | None: return None @@ -497,9 +500,6 @@ class SampleStateTorch(SampleState[SampleStateTensorsHostTorch, SampleStateTenso class TorchSampler(Sampler[SampleStateTorch], AsyncWorkerMixin): - DEFAULT_MAX_STOP_WORD_LENGTH: Final[int] = 20 - DEFAULT_MAX_STOP_WORDS: Final[int] = 10 - SampleState = SampleStateTorch @override @@ -660,8 +660,6 @@ def __init__(self, args: Args): # One row per sequence slot plus the scratch row (dummy_slot_row). self._slot_params = SlotSamplingParams(self.max_num_sequences + 1, device="cuda") self._finish_reasons_handler = FinishReasonsHandler( - max_stop_word_length=self.DEFAULT_MAX_STOP_WORD_LENGTH, - max_num_stop_words=self.DEFAULT_MAX_STOP_WORDS, max_num_sequences=self.max_num_sequences, max_beam_width=self.max_beam_width, max_tokens=self.max_tokens, @@ -1261,6 +1259,10 @@ def validate_request(self, request: LlmRequest) -> None: # path (see beam_search_sampling_batch_cba); the mode only selects # the done verdict computed there. + @override + def release_request(self, request: LlmRequest) -> None: + self._finish_reasons_handler.release_request(request) + @override @nvtx_range("setup_sampler_step") def setup_sampler_step(self, scheduled_requests: ScheduledRequests) -> None: @@ -1325,7 +1327,6 @@ def setup_sampler_step(self, scheduled_requests: ScheduledRequests) -> None: max_lengths_cuda=max_lens_tensor_cuda, end_ids_cuda=end_ids_tensor_cuda, seq_slots_host=seq_slots_tensor_host, - all_sampling_requests=new_requests + scheduled_requests.generation_requests, ) self._slot_params.admit(new_requests, seq_slots_tensor_cuda_long) @@ -1901,6 +1902,12 @@ def _sample_batched_by_strategy( ) ) else: + stop_words_metadata_cuda: torch.Tensor | None = None + if self._finish_reasons_handler.has_stop_word_history: + self._finish_reasons_handler.record_stop_word_stream() + stop_words_metadata_cuda = ( + self._finish_reasons_handler.store.stop_words_metadata_cuda + ) beam_step = BeamSearchStep.from_params( group.beam_params, cuda_device, return_probs=group.needs_probs ) @@ -1918,7 +1925,7 @@ def _sample_batched_by_strategy( ..., DEFAULT_STEP_IDX ], end_ids_cuda=self._finish_reasons_handler.store.end_ids_cuda, - past_tokens_cuda=self._finish_reasons_handler.store.past_tokens_cuda, + stop_words_metadata_cuda=stop_words_metadata_cuda, ), ) group_temperature_cuda = beam_step.temperature diff --git a/tests/unittest/_torch/disaggregation/test_benchmark_disagg.py b/tests/unittest/_torch/disaggregation/test_benchmark_disagg.py index b1101197c92f..1e37f0eee436 100644 --- a/tests/unittest/_torch/disaggregation/test_benchmark_disagg.py +++ b/tests/unittest/_torch/disaggregation/test_benchmark_disagg.py @@ -1278,6 +1278,7 @@ def _make_executor( from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor ex = object.__new__(PyExecutor) + ex.sampler = Mock(spec=["release_request"]) ex.benchmark_req_queues_size = 8 ex.kv_cache_transceiver = Mock() # The fill-gate requests model "waiting for KV" / "generating", not @@ -1581,6 +1582,8 @@ def test_a_stalled_fill_fails_every_active_request_through_the_real_error_path( assert request.state == LlmRequestState.GENERATION_COMPLETE freed = [call.args[0] for call in ex.resource_manager.free_resources.call_args_list] assert freed == requests + released = [call.args[0] for call in ex.sampler.release_request.call_args_list] + assert released == requests assert ex._fatal_error is None and not ex.is_shutdown def test_a_fitting_init_request_keeps_the_real_error_path_quiet(self, monkeypatch): @@ -1594,6 +1597,7 @@ def test_a_fitting_init_request_keeps_the_real_error_path_quiet(self, monkeypatc assert ex.responses == {} assert ex.active_requests == requests ex.resource_manager.free_resources.assert_not_called() + ex.sampler.release_request.assert_not_called() def test_the_fail_fast_wakes_a_consumer_asleep_in_await_responses(self, monkeypatch): """The error responses are published under the response condition and diff --git a/tests/unittest/_torch/disaggregation/test_disagg_coordinator_errors.py b/tests/unittest/_torch/disaggregation/test_disagg_coordinator_errors.py index 0ab83cc39d2d..944271cf6cba 100644 --- a/tests/unittest/_torch/disaggregation/test_disagg_coordinator_errors.py +++ b/tests/unittest/_torch/disaggregation/test_disagg_coordinator_errors.py @@ -370,6 +370,7 @@ def _executor_with_real_error_path(monkeypatch) -> tuple: monkeypatch.delenv("TRTLLM_DISAGG_BENCHMARK_GEN_ONLY", raising=False) monkeypatch.delenv("TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP", raising=False) executor = object.__new__(PyExecutor) + executor.sampler = Mock(spec=["release_request"]) executor.active_requests = [] executor.canceled_req_ids = [] executor._pending_transfer_responses = [] diff --git a/tests/unittest/_torch/executor/test_multimodal_scheduler.py b/tests/unittest/_torch/executor/test_multimodal_scheduler.py index 6b1830c83395..0cfa5a6ee9d8 100644 --- a/tests/unittest/_torch/executor/test_multimodal_scheduler.py +++ b/tests/unittest/_torch/executor/test_multimodal_scheduler.py @@ -1047,6 +1047,7 @@ def test_terminate_request_releases_multimodal_cache_references_idempotently(): executor.global_rank = 0 executor.model_engine = SimpleNamespace(mm_encoder_cache=cache) executor.draft_model_engine = None + executor.sampler = Mock() executor.resource_manager = SimpleNamespace(free_resources=freed.append) executor._prefetched_request_ids = {request.py_request_id} executor._disagg_coordinator = Mock() diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index c104133ae8a5..6c698a26266d 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -830,15 +830,24 @@ def finish_gen_receive(req): def test_free_request_resources_releases_target_and_draft_model_state() -> None: target, draft = Mock(), Mock() + events = [] + sampler = Mock() + sampler.release_request.side_effect = lambda request: events.append(request.py_seq_slot) executor = types.SimpleNamespace( + sampler=sampler, model_engine=types.SimpleNamespace(model=target), draft_model_engine=types.SimpleNamespace(model=draft), resource_manager=Mock(), _prefetched_request_ids={41}, disagg=Mock(), ) - request = types.SimpleNamespace(py_request_id=41) + request = types.SimpleNamespace(py_request_id=41, py_seq_slot=3) + executor.resource_manager.free_resources.side_effect = lambda req: setattr( + req, "py_seq_slot", None + ) PyExecutor._free_request_resources(executor, request) + sampler.release_request.assert_called_once_with(request) + assert events == [3] target.release_request_state.assert_called_once_with(41) draft.release_request_state.assert_called_once_with(41) assert not executor._prefetched_request_ids diff --git a/tests/unittest/_torch/executor/test_recompute_pause.py b/tests/unittest/_torch/executor/test_recompute_pause.py index e83684cf9cae..7c5e3c417f18 100644 --- a/tests/unittest/_torch/executor/test_recompute_pause.py +++ b/tests/unittest/_torch/executor/test_recompute_pause.py @@ -23,6 +23,7 @@ def __init__(self, request_id: int = 7) -> None: def _make_executor(handler: Mock | None) -> PyExecutor: executor = object.__new__(PyExecutor) + executor.sampler = Mock() executor.model_engine = None executor.draft_model_engine = None executor._disagg_pp_termination_handler = handler @@ -59,6 +60,7 @@ def test_recompute_pause_skips_request_finished_during_overlap() -> None: executor._pause_recompute_paused_requests(scheduled_batch) executor.resource_manager.free_resources.assert_called_once_with(request) + executor.sampler.release_request.assert_called_once_with(request) request.reset_for_recompute.assert_not_called() @@ -79,6 +81,7 @@ def test_recompute_pause_defers_reset_until_pp_consensus() -> None: handler.terminate.assert_called_once_with(request) assert events == [] + executor.sampler.release_request.assert_not_called() assert request.py_request_id in executor._pending_recompute_pause_ids assert request.py_request_id in executor.inflight_req_ids assert request.py_request_id in executor.result_wait_queues @@ -88,6 +91,7 @@ def test_recompute_pause_defers_reset_until_pp_consensus() -> None: executor._on_disagg_pp_termination(request) + executor.sampler.release_request.assert_called_once_with(request) assert events == ["free", "reset"] assert request.py_request_id not in executor._pending_recompute_pause_ids assert request.py_request_id not in executor.inflight_req_ids @@ -106,6 +110,7 @@ def test_terminal_request_does_not_recompute_after_pp_consensus() -> None: executor._on_disagg_pp_termination(request) executor.resource_manager.free_resources.assert_called_once_with(request) + executor.sampler.release_request.assert_called_once_with(request) request.reset_for_recompute.assert_not_called() assert request.py_request_id not in executor._pending_recompute_pause_ids assert request.py_request_id not in executor.inflight_req_ids diff --git a/tests/unittest/_torch/sampler/test_beam_search.py b/tests/unittest/_torch/sampler/test_beam_search.py index 2aaa75c9a312..b2277790cdfb 100644 --- a/tests/unittest/_torch/sampler/test_beam_search.py +++ b/tests/unittest/_torch/sampler/test_beam_search.py @@ -1777,10 +1777,12 @@ def test_beam_search_cba_reorders_stop_window(): m.cache_indirection[0, b, :] = b m.cum_log_probs[0] = torch.tensor([-5.0, -1.0]) # window rows distinguishable per beam - stop_window = torch.zeros((3, 1, K), dtype=torch.int32) + allocation = torch.zeros(2 * 3 * K, dtype=torch.int32) + stop_window = allocation[:3 * K].view(3, 1, K) stop_window[:, 0, 0] = 100 stop_window[:, 0, 1] = 200 - m.stop_past_tokens = stop_window + m.stop_words_metadata = torch.tensor([[allocation.data_ptr(), 0, 0, 3, K]], + dtype=torch.int64) # beam1 dominates: both slots descend from beam 1 logits = torch.full((K, vocab), -50.0) diff --git a/tests/unittest/_torch/sampler/test_finish_reasons_fused.py b/tests/unittest/_torch/sampler/test_finish_reasons_fused.py index 1b5b6f3c6663..65cf23eddac0 100644 --- a/tests/unittest/_torch/sampler/test_finish_reasons_fused.py +++ b/tests/unittest/_torch/sampler/test_finish_reasons_fused.py @@ -15,8 +15,6 @@ def test_fused_matches_tensor_ops(max_tokens: int, max_beam_width: int) -> None: torch.manual_seed(0) max_num_sequences, end_id = 5, 99 handler = FinishReasonsHandler( - max_stop_word_length=1, - max_num_stop_words=1, max_num_sequences=max_num_sequences, max_beam_width=max_beam_width, max_tokens=max_tokens, @@ -58,3 +56,261 @@ def run() -> torch.Tensor: torch.testing.assert_close(fused, run(), rtol=0, atol=0) assert fused[0, 4, 0].item() == FinishReason.END_ID.value assert fused[0, 0, 0].item() == FinishReason.LENGTH.value + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_compact_stop_words_history_and_slot_reuse() -> None: + """Packed matching agrees with a token-list oracle across accepted prefixes.""" + from types import SimpleNamespace + + from tensorrt_llm._torch.pyexecutor.sampler.ops.triton import reorder_compact_stop_word_history + + for steps, beams in ((1, 1), (3, 1), (1, 3)): + handler = FinishReasonsHandler( + max_num_sequences=64, + max_beam_width=beams, + max_tokens=steps, + max_seq_len=512, + ) + # Non-contiguous slots and different rule counts / lengths in the same batch. + words = [[[7], [2, 3, 4], [0, 1, 2, 3, 4]], [[9], [8]], [[4, 5]]] + prompts = [[0, 1, 2, 3], [8, 2], [2, 3, 4]] + slots_host = torch.tensor([5, 1, 8], dtype=torch.int32) + slots = slots_host.cuda().long() + requests = [ + SimpleNamespace( + py_seq_slot=slot, + py_stop_words_list=rules, + py_orig_prompt_len=len(prompt), + py_max_new_tokens=64, + py_end_id=99, + is_context_only_request=False, + py_beam_width=beams, + is_finished=False, + get_tokens=lambda prompt=prompt: [prompt[:] for _ in range(beams)], + ) + for slot, rules, prompt in zip(slots_host.tolist(), words, prompts) + ] + handler.setup_new_request_handling() + for request in requests: + handler.prepare_for_new_request(request) + handler.update_for_new_request( + seq_slots_cuda_long=slots, + max_lengths_cuda=torch.full_like(slots, 100, dtype=torch.int32), + end_ids_cuda=torch.full_like(slots, 99, dtype=torch.int32), + seq_slots_host=slots_host, + ) + histories = [[prompt[:] for _ in range(beams)] for prompt in prompts] + previous = None + for iteration in range(4): + accepted_host = [1, steps, max(1, steps - 1)] + if steps > 1 and iteration == 2: + accepted_host[0] = steps + if previous is not None: + for request_idx in range(3): + for beam in range(beams): + histories[request_idx][beam].extend( + previous[: accepted_host[request_idx], request_idx, beam].tolist() + ) + if steps > 1 and iteration == 2: + retained = [handler._stop_word_allocations[slot] for slot in (5, 1)] + retained_data = [entry.tensor.clone() for entry in retained] + retained_metadata = handler.store.stop_words_metadata_cuda[slots[:2]].clone() + handler.release_request(requests[2]) + words[2] = [[6, 7], [5, 6, 7]] + requests[2] = SimpleNamespace( + **{ + **vars(requests[2]), + "py_stop_words_list": words[2], + "py_orig_prompt_len": 2, + "get_tokens": lambda: [[5, 6] for _ in range(beams)], + } + ) + handler.setup_new_request_handling() + handler.prepare_for_new_request(requests[2]) + handler.update_for_new_request( + seq_slots_cuda_long=slots[2:], + max_lengths_cuda=torch.full_like(slots[2:], 100, dtype=torch.int32), + end_ids_cuda=torch.full_like(slots[2:], 99, dtype=torch.int32), + seq_slots_host=slots_host[2:], + ) + for slot, entry, before in zip((5, 1), retained, retained_data): + assert handler._stop_word_allocations[slot] is entry + torch.testing.assert_close(entry.tensor, before) + torch.testing.assert_close( + handler.store.stop_words_metadata_cuda[slots[:2]], retained_metadata + ) + replacement = handler._stop_word_allocations[8] + assert replacement.request is requests[2] + assert handler.store.stop_words_metadata_cuda[8].tolist() == [ + replacement.tensor.data_ptr(), + 2, + 8, + steps + 2, + beams, + ] + assert handler.store.max_stop_word_lengths_host[8].item() == 3 + histories[2] = [[5, 6] for _ in range(beams)] + accepted_host[2] = 1 + if beams > 1: + parents_host = torch.tensor([[2, 0, 0], [1, 1, 0], [0, 2, 1]]) + reorder_compact_stop_word_history( + handler.store.stop_words_metadata_cuda, slots, parents_host.cuda() + ) + histories = [ + [history[parent][:] for parent in parents] + for history, parents in zip(histories, parents_host.tolist()) + ] + tokens_host = ( + torch.arange(steps * 3 * beams).reshape(steps, 3, beams) + 4 + iteration + ) % 10 + if beams > 1 and iteration < 2: + tokens_host[0, 2] = torch.tensor([5, 7, 4] if iteration == 0 else [5, 5, 5]) + if steps > 1 and iteration == 2: + tokens_host[0, 2] = 7 + tokens = tokens_host.cuda().int() + accepted = torch.tensor(accepted_host, dtype=torch.int32, device="cuda") + if iteration == 0: + accepted.fill_(1) + actual = handler._are_stop_words(slots, tokens, accepted).cpu() + expected = torch.zeros_like(tokens_host, dtype=torch.bool) + for step in range(steps): + for request_idx in range(3): + for beam in range(beams): + prefix = ( + histories[request_idx][beam] + + tokens_host[: step + 1, request_idx, beam].tolist() + ) + expected[step, request_idx, beam] = any( + prefix[-len(word) :] == word for word in words[request_idx] + ) + torch.testing.assert_close(actual, expected) + previous = tokens_host + + # Reusing a slot without stop words must clear its metadata and ownership. + requests[0].py_stop_words_list = None + handler.setup_new_request_handling() + handler.prepare_for_new_request(requests[0]) + handler.update_for_new_request( + seq_slots_cuda_long=slots[:1], + max_lengths_cuda=torch.tensor([100], device="cuda", dtype=torch.int32), + end_ids_cuda=torch.tensor([99], device="cuda", dtype=torch.int32), + seq_slots_host=slots_host[:1], + ) + assert 5 not in handler._stop_word_allocations + assert handler.store.max_stop_word_lengths_host[5].item() == 0 + assert handler.store.stop_words_metadata_cuda[5, :4].count_nonzero().item() == 0 + requests[1].is_finished = True + handler.release_request(requests[1]) + assert 1 not in handler._stop_word_allocations + assert 8 in handler._stop_word_allocations + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_compact_stop_words_storage_and_finish_precedence() -> None: + """An outlier does not expand other slots or reset their device histories.""" + from types import SimpleNamespace + + handler = FinishReasonsHandler( + max_num_sequences=1024, + max_beam_width=1, + max_tokens=1, + max_seq_len=4096, + ) + + def admit(slot: int, words: list[list[int]], prompt: list[int]) -> SimpleNamespace: + request = SimpleNamespace( + py_seq_slot=slot, + py_stop_words_list=words, + py_orig_prompt_len=len(prompt), + py_max_new_tokens=1000, + py_end_id=99, + is_context_only_request=False, + py_beam_width=1, + is_finished=False, + get_tokens=lambda: [prompt], + ) + handler.setup_new_request_handling() + handler.prepare_for_new_request(request) + handler.update_for_new_request( + seq_slots_cuda_long=torch.tensor([slot], device="cuda", dtype=torch.int64), + max_lengths_cuda=torch.tensor([10], device="cuda", dtype=torch.int32), + end_ids_cuda=torch.tensor([99], device="cuda", dtype=torch.int32), + seq_slots_host=torch.tensor([slot]), + ) + return request + + assert not handler.has_stop_word_history + many = admit(2, [[i] for i in range(1024)], [1, 2, 3]) + assert not handler.has_stop_word_history + allocation = handler._stop_word_allocations[2].tensor + snapshot = allocation.clone() + long = admit(3, [list(range(256))], list(range(255))) + assert handler.has_stop_word_history + assert handler._stop_word_allocations[2].tensor is allocation + torch.testing.assert_close(allocation, snapshot) + total_bytes = handler.store.stop_words_metadata_cuda.nbytes + sum( + entry.tensor.nbytes for entry in handler._stop_word_allocations.values() + ) + # 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 + + slots = torch.tensor([2], device="cuda", dtype=torch.int64) + tokens = torch.zeros((1, 1025, 1), device="cuda", dtype=torch.int32) + for token, seq_len, expected in ( + (7, 1, FinishReason.STOP_WORDS), + (7, 9, FinishReason.LENGTH), + (99, 9, FinishReason.END_ID), + ): + tokens[0, 2, 0] = token + handler.write_finish_reasons( + seq_slots_host=slots.cpu(), + seq_slots_cuda=slots, + seq_lens_cuda=torch.tensor([seq_len], device="cuda"), + new_tokens_cuda=tokens, + ) + assert handler.store.finish_reasons_cuda[0, 2, 0].item() == expected.value + + long_match = handler._are_stop_words( + torch.tensor([3], dtype=torch.int64, device="cuda"), + torch.tensor([[[255]]], dtype=torch.int32, device="cuda"), + 1, + ) + assert long_match.item() + + primary_stream = torch.cuda.current_stream() + alternate_stream = torch.cuda.Stream() + for stream in (primary_stream, alternate_stream, primary_stream): + transient = admit(4, [[7, 8]], [7]) + stream.wait_stream(primary_stream) + with torch.cuda.stream(stream): + match_slots = torch.tensor([4], dtype=torch.int64, device="cuda") + match_tokens = torch.tensor([[[8]]], dtype=torch.int32, device="cuda") + first_match = handler._are_stop_words(match_slots, match_tokens, 1) + second_match = handler._are_stop_words(match_slots, match_tokens, 1) + primary_stream.wait_stream(stream) + assert first_match.item() + assert not second_match.item() + transient.is_finished = True + handler.release_request(transient) + assert 4 not in handler._stop_word_allocations + + old = admit(4, [[7, 8]], [7]) + replacement = admit(4, [[8, 9]], [8]) + replacement_allocation = handler._stop_word_allocations[4].tensor + handler.release_request(old) + assert handler._stop_word_allocations[4].tensor is replacement_allocation + handler.release_request(replacement) + + # Releasing one request preserves other live requests and is idempotent. + many.is_finished = True + handler.release_request(many) + handler.release_request(many) + assert 2 not in handler._stop_word_allocations + assert 3 in handler._stop_word_allocations + handler.release_request(long) + long.py_seq_slot = None + handler.release_request(long) + assert not handler._stop_word_allocations + assert not handler.has_stop_word_history diff --git a/tests/unittest/_torch/sampler/test_torch_sampler.py b/tests/unittest/_torch/sampler/test_torch_sampler.py index 954cdb07ca02..d9402df918d9 100644 --- a/tests/unittest/_torch/sampler/test_torch_sampler.py +++ b/tests/unittest/_torch/sampler/test_torch_sampler.py @@ -1395,152 +1395,6 @@ def raising_single_token_stop_words_ctx(expect_raise: bool) -> Generator[None, N ) run_test_with_warmup(uut_provider_with_stop_words, max_sync_s=0.5) - @classmethod - def check_resize_and_update_stop_words_buffer( - cls, - old_stop_words: torch.Tensor, - old_past_tokens: torch.Tensor, - new_stop_words: torch.Tensor, - new_past_tokens: torch.Tensor, - seq_slots_to_compare: torch.Tensor, - num_draft_tokens: int, - ): - old_num_stop_words = old_stop_words.shape[0] - old_stop_word_length = old_stop_words.shape[1] - - old_past_token_length = old_past_tokens.shape[0] - - # These sizes should not change after the resize - assert old_stop_words.shape[2] == new_stop_words.shape[2] - assert old_past_tokens.shape[1] == new_past_tokens.shape[1] - assert old_past_tokens.shape[2] == new_past_tokens.shape[2] - - assert ( - new_stop_words[-old_num_stop_words:, -old_stop_word_length:, seq_slots_to_compare] - == old_stop_words[..., seq_slots_to_compare] - ).all() - # initial fill has an offset of 1, as there will be a shift happening during sample_async - assert ( - new_past_tokens[-old_past_token_length:-num_draft_tokens, seq_slots_to_compare, :] - == old_past_tokens[:-num_draft_tokens, seq_slots_to_compare, :] - ).all() - - @classmethod - def test_stop_words_buffer_resize(cls, monkeypatch: pytest.MonkeyPatch): - @contextmanager - def check_resize_ctx() -> Generator[None, None, None]: - with monkeypatch.context() as patch_ctx: - setup_sampler_step_orig = TorchSampler.setup_sampler_step - - def setup_sampler_step_with_size_check(self, scheduled_requests: ScheduledRequests): - # RequestCase.build calls setup_sampler_step to fill the buffers for all context requests - # Move the context requests to the generation requests - scheduled_requests.generation_requests = scheduled_requests.context_requests - # Add a request that enforces a resize - scheduled_requests.context_requests_last_chunk = [ - cls.RequestCase( - prompt=[1], - stop_words_list=[ - [x for x in range(2 * TorchSampler.DEFAULT_MAX_STOP_WORD_LENGTH)] - ], - new_tokens=[4], - finish_reasons=[cls.NOT_FINISHED], - ).request - ] - # Store the old stop words and past tokens for comparison - old_stop_words = self.store.stop_words.clone() - old_past_tokens = self.store.past_tokens.clone() - # Call setup sampler step to trigger the resize - setup_sampler_step_orig(self, scheduled_requests) - - # Check if sizes are correct - assert self.store.stop_words.shape[0] == TorchSampler.DEFAULT_MAX_STOP_WORDS - assert ( - self.store.stop_words.shape[1] - == 2 * TorchSampler.DEFAULT_MAX_STOP_WORD_LENGTH - ) - assert self.max_tokens == 1 - assert ( - self.store.past_tokens.shape[0] - == 2 * TorchSampler.DEFAULT_MAX_STOP_WORD_LENGTH - 1 + self.max_tokens - ) - # Check if values are added correctly - seq_slots_to_compare_cuda = torch.Tensor( - [ - scheduled_requests.generation_requests[x].py_seq_slot - for x in range(len(scheduled_requests.generation_requests)) - ] - ).to(device="cuda", dtype=torch.int32, non_blocking=True) - cls.check_resize_and_update_stop_words_buffer( - old_stop_words, - old_past_tokens, - self.store.stop_words, - self.store.past_tokens, - seq_slots_to_compare_cuda, - self.max_draft_len, - ) - - patch_ctx.setattr( - TorchSampler, "setup_sampler_step", setup_sampler_step_with_size_check - ) - yield - - # The test adds one more request - num_requests = 8 - uut_provider_with_resize_on_demand = cls.RequestCase.build( - [ - cls.RequestCase( - prompt=[1 + x], - stop_words_list=[[x + 100]], - new_tokens=[x + 4, x + 3, x + 2], - finish_reasons=[cls.NOT_FINISHED, cls.NOT_FINISHED, cls.NOT_FINISHED], - ) - for x in range(num_requests) - ], - extra_context=lambda: check_resize_ctx(), - ) - run_test_with_warmup(uut_provider_with_resize_on_demand, max_sync_s=0.5) - - @classmethod - def test_stop_words_buffer_resizes_on_demand(cls, monkeypatch: pytest.MonkeyPatch): - @contextmanager - def check_resize_ctx() -> Generator[None, None, None]: - with monkeypatch.context() as patch_ctx: - setup_sampler_step_orig = TorchSampler.setup_sampler_step - - def setup_sampler_step_with_size_check(self, scheduled_requests: ScheduledRequests): - setup_sampler_step_orig(self, scheduled_requests) - assert self.store.stop_words.shape[0] == TorchSampler.DEFAULT_MAX_STOP_WORDS - assert ( - self.store.stop_words.shape[1] - == 2 * TorchSampler.DEFAULT_MAX_STOP_WORD_LENGTH - ) - assert self.max_tokens == 1 - assert ( - self.store.past_tokens.shape[0] - == 2 * TorchSampler.DEFAULT_MAX_STOP_WORD_LENGTH - 1 + self.max_tokens - ) - - patch_ctx.setattr( - TorchSampler, "setup_sampler_step", setup_sampler_step_with_size_check - ) - yield - - uut_provider_with_resize_on_demand = cls.RequestCase.build( - [ - cls.RequestCase( - prompt=[1], - stop_words_list=[ - [x for x in range(2 * TorchSampler.DEFAULT_MAX_STOP_WORD_LENGTH)] - ], - new_tokens=[4], - finish_reasons=[cls.NOT_FINISHED], - ) - ], - extra_context=lambda: check_resize_ctx(), - ) - run_test_with_warmup(uut_provider_with_resize_on_demand, max_sync_s=None) - @staticmethod def _all_beams_finished_reference(row: torch.Tensor, beam_width: int) -> bool: """The per-request reduction the batched prefix count replaces.""" diff --git a/tests/unittest/disaggregated/test_transfer_lifecycle_gpu.py b/tests/unittest/disaggregated/test_transfer_lifecycle_gpu.py index 871715350073..46a9334c10c0 100644 --- a/tests/unittest/disaggregated/test_transfer_lifecycle_gpu.py +++ b/tests/unittest/disaggregated/test_transfer_lifecycle_gpu.py @@ -543,6 +543,7 @@ def observe_fatal(event: QuiescenceFatalEvent) -> None: assert full_pool["free"] == full_pool["available"] == full_pool["evictable"] == 0 # No model, scheduler loop, response transport, or async-send manager is initialized. executor = object.__new__(PyExecutor) + executor.sampler = Mock(spec=["release_request"]) executor.model_engine = None executor.draft_model_engine = None executor.kv_cache_transceiver = transceiver