diff --git a/tensorrt_llm/_torch/pyexecutor/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler.py index 9ea7153f542d..0828917bfec4 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler.py @@ -1422,7 +1422,7 @@ def prepare_for_new_request(self, request: LlmRequest) -> None: def update_for_new_request( self, *, - seq_slots_cuda: torch.Tensor, + seq_slots_cuda_long: torch.Tensor, max_lengths_cuda: torch.Tensor, end_ids_cuda: torch.Tensor, seq_slots_host: torch.Tensor, @@ -1436,8 +1436,8 @@ def update_for_new_request( need to be re-processed. Args: - seq_slots_cuda: The sequence slots of the processed requests. Used for accessing device buffers. - Shape: [len(requests)] + 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. @@ -1450,8 +1450,8 @@ def update_for_new_request( temp_data = self._temp_data store = self.store - store.max_lengths_cuda[seq_slots_cuda] = max_lengths_cuda - store.end_ids_cuda[seq_slots_cuda] = end_ids_cuda + store.max_lengths_cuda.index_copy_(0, seq_slots_cuda_long, max_lengths_cuda) + store.end_ids_cuda.index_copy_(0, seq_slots_cuda_long, end_ids_cuda) store.max_stop_word_lengths_host[seq_slots_host] = torch.tensor( temp_data.max_stop_word_lengths, device="cpu", dtype=torch.int32 ) @@ -2036,6 +2036,14 @@ class BeamSearchStore: """Shape: batch_size, beam_width, sequence_length Usage: Stores the original tokens for each beam. This is used to recover the original tokens for each beam when streaming is enabled""" + seq_offsets: torch.Tensor + """Shape: (max_num_sequences,), dtype int64 + Usage: Cached `arange(max_num_sequences) * max_beam_width` used by + ``beam_search_sampling_batch`` to flatten (batch_idx, beam_idx) pairs.""" + beam_idx_arange: torch.Tensor + """Shape: (max_beam_width,), dtype int32 + Usage: Cached `arange(max_beam_width)` used as the scatter source in the + per-step ``cache_indirection.scatter_``.""" @dataclass(kw_only=True) class LogProbsStore: @@ -2103,6 +2111,11 @@ def _create_store(self) -> Store: predecessor_beams = int_tensor(self.CACHE_INDIRECTION_SHAPE[:-1]) original_tokens = int_tensor(self.CACHE_INDIRECTION_SHAPE) first_finish_reasons = int_tensor(self.CACHE_INDIRECTION_SHAPE[:-1]) + seq_offsets = ( + torch.arange(self.max_num_sequences, device="cuda", dtype=torch.int64) + * self.max_beam_width + ) + beam_idx_arange = torch.arange(self.max_beam_width, device="cuda", dtype=torch.int32) beam_search_store = self.BeamSearchStore( cache_indirection=cache_indirection, cache_indirection_buffer=cache_indirection_buffer, @@ -2110,6 +2123,8 @@ def _create_store(self) -> Store: predecessor_beams=predecessor_beams, original_tokens=original_tokens, first_finish_reasons=first_finish_reasons, + seq_offsets=seq_offsets, + beam_idx_arange=beam_idx_arange, ) return self.Store( new_tokens=new_tokens, @@ -2746,8 +2761,12 @@ def setup_sampler_step(self, scheduled_requests: ScheduledRequests) -> None: seq_slots_tensor_cuda = full_list_tensor_cuda[0] max_lens_tensor_cuda = full_list_tensor_cuda[1] end_ids_tensor_cuda = full_list_tensor_cuda[2] + + # Cast to int64 once for downstream ``index_copy_`` / ``index_fill_`` calls. + seq_slots_tensor_cuda_long = seq_slots_tensor_cuda.long() + self._finish_reasons_handler.update_for_new_request( - seq_slots_cuda=seq_slots_tensor_cuda, + seq_slots_cuda_long=seq_slots_tensor_cuda_long, max_lengths_cuda=max_lens_tensor_cuda, end_ids_cuda=end_ids_tensor_cuda, seq_slots_host=seq_slots_tensor_host, @@ -2760,7 +2779,7 @@ def setup_sampler_step(self, scheduled_requests: ScheduledRequests) -> None: self._prepare_beam_search( beam_search_store, self.store.log_probs_store, - seq_slots=seq_slots_tensor_cuda, + seq_slots_long=seq_slots_tensor_cuda_long, max_prompt_len=max_prompt_len, ) @@ -2768,60 +2787,27 @@ def setup_sampler_step(self, scheduled_requests: ScheduledRequests) -> None: def _prepare_beam_search( beam_search_store: BeamSearchStore, log_probs_store: LogProbsStore, - seq_slots: torch.Tensor, + seq_slots_long: torch.Tensor, max_prompt_len: int, ) -> None: """Prepare the beam search buffers for the requests If the last context chunk is being processed, - initialize/reset the buffers for the request + initialize/reset the buffers for the request. + + ``seq_slots_long`` must be int64 (required by ``index_fill_``). """ - cache_indirection = beam_search_store.cache_indirection - cache_indirection[seq_slots, :, :max_prompt_len] = torch.zeros( - (1), - dtype=cache_indirection.dtype, - device=cache_indirection.device, - ) - cum_log_probs = beam_search_store.cum_log_probs - cum_log_probs[seq_slots] = torch.zeros( - (1,), - dtype=cum_log_probs.dtype, - device=cum_log_probs.device, - ) - sampled_log_probs = log_probs_store.sampled_log_probs - sampled_log_probs[seq_slots] = torch.zeros( - (1,), - dtype=sampled_log_probs.dtype, - device=sampled_log_probs.device, - ) - sampled_log_prob_ranks = log_probs_store.sampled_log_prob_ranks - sampled_log_prob_ranks[seq_slots] = torch.zeros( - (1,), - dtype=sampled_log_prob_ranks.dtype, - device=sampled_log_prob_ranks.device, - ) - predecessor_beams = beam_search_store.predecessor_beams - predecessor_beams[seq_slots] = torch.zeros( - (1,), - dtype=predecessor_beams.dtype, - device=predecessor_beams.device, - ) - first_finish_reasons = beam_search_store.first_finish_reasons - first_finish_reasons[seq_slots] = ( - torch.tensor( - FinishReason.NOT_FINISHED.value, - pin_memory=prefer_pinned(), - dtype=first_finish_reasons.dtype, - ) - .to(first_finish_reasons.device, non_blocking=True) - .unsqueeze(0) + beam_search_store.cache_indirection.narrow(2, 0, max_prompt_len).index_fill_( + 0, seq_slots_long, 0 ) - original_tokens = beam_search_store.original_tokens - original_tokens[seq_slots] = torch.zeros( - (1,), - dtype=original_tokens.dtype, - device=original_tokens.device, + beam_search_store.cum_log_probs.index_fill_(0, seq_slots_long, 0) + log_probs_store.sampled_log_probs.index_fill_(0, seq_slots_long, 0) + log_probs_store.sampled_log_prob_ranks.index_fill_(0, seq_slots_long, 0) + beam_search_store.predecessor_beams.index_fill_(0, seq_slots_long, 0) + beam_search_store.first_finish_reasons.index_fill_( + 0, seq_slots_long, FinishReason.NOT_FINISHED.value ) + beam_search_store.original_tokens.index_fill_(0, seq_slots_long, 0) @torch.inference_mode() def _process_draft_tokens_rejection_sampling( @@ -3231,6 +3217,8 @@ def _add_metadata_to_grouped_requests( ), # Should be on device for beam search finished_beams=beam_search_store.first_finish_reasons, predecessor_beams=beam_search_store.predecessor_beams, + seq_offsets=beam_search_store.seq_offsets, + beam_idx_arange=beam_search_store.beam_idx_arange, ) elif metadata_type is None: metadata = None diff --git a/tensorrt_llm/_torch/pyexecutor/sampling_utils.py b/tensorrt_llm/_torch/pyexecutor/sampling_utils.py index d61815d8cfcf..6df7048c2c68 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampling_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/sampling_utils.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -64,6 +64,11 @@ class BeamSearchMetadata(StrategyMetadata): seq_lens: torch.Tensor finished_beams: torch.Tensor predecessor_beams: torch.Tensor + # Pre-computed indexer constants sliced (not allocated) per call. + # seq_offsets[i] = i * max_beam_width, shape (max_num_sequences,), int64. + # beam_idx_arange[j] = j, shape (max_beam_width,), int32. + seq_offsets: torch.Tensor + beam_idx_arange: torch.Tensor @dataclass(frozen=True, kw_only=True) @@ -361,10 +366,9 @@ def beam_search_sampling_batch( max_beam_width = beam_search_args.finished_beams.size(1) finished_beams = beam_search_args.finished_beams[beam_search_args.seq_slots].view(-1) - offset_predecessor_beam = predecessor_beam + ( - torch.arange(predecessor_beam.size(0), device=predecessor_beam.device).unsqueeze(1) - * max_beam_width - ) + offset_predecessor_beam = predecessor_beam + beam_search_args.seq_offsets[ + : predecessor_beam.size(0) + ].unsqueeze(1) finished_beams = finished_beams[offset_predecessor_beam] beam_search_args.finished_beams[beam_search_args.seq_slots] = finished_beams.view( batch_size, max_beam_width @@ -385,19 +389,15 @@ def beam_search_sampling_batch( out=cache_indirection, ) - # Prepare target values - target_values = ( - torch.arange( - beam_width_out * batch_size, device=cache_indirection.device, dtype=torch.int32 - ) - % beam_width_out - ) - # seq lens is of shape (batch_size), we assume all beams have the same seq len # therefore we can use expand index = beam_search_args.seq_lens.view(-1, 1, 1).expand(-1, beam_width_out, 1) # index is of shape (batch_size, beam_width, 1) - src = target_values.view(batch_size, beam_width_out, 1) + src = ( + beam_search_args.beam_idx_arange[:beam_width_out] + .view(1, beam_width_out, 1) + .expand(batch_size, beam_width_out, 1) + ) # src is of shape (batch_size, beam_width, 1) # cache_indirection is of shape (batch_size, beam_width, max_seq_len) cache_indirection.scatter_(2, index, src) diff --git a/tests/unittest/_torch/sampler/test_beam_search.py b/tests/unittest/_torch/sampler/test_beam_search.py index 607bfe341202..c582810aeda8 100644 --- a/tests/unittest/_torch/sampler/test_beam_search.py +++ b/tests/unittest/_torch/sampler/test_beam_search.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -488,6 +488,10 @@ def test_beam_search_sampling_batch_basic(): new_log_probs_result = new_log_probs.clone() predecessor_beams_result = predecessor_beams.clone() + # Pre-computed per-step constants (matches production cache contract). + seq_offsets = torch.arange(max_batch_size, dtype=torch.int64) * beam_width + beam_idx_arange = torch.arange(beam_width, dtype=torch.int32) + # Create BeamSearchMetadata beam_search_args = BeamSearchMetadata( cache_indirection=cache_indirection_result, @@ -498,6 +502,8 @@ def test_beam_search_sampling_batch_basic(): finished_beams=finished_beams_result, new_log_probs=new_log_probs_result, predecessor_beams=predecessor_beams_result, + seq_offsets=seq_offsets, + beam_idx_arange=beam_idx_arange, ) # Run beam search sampling