Skip to content

[https://nvbugs/5829097][fix] Re-init TRTLLM sampler to use sample stream in multi-stream cases. (cherry-pick from https://github.com/NVIDIA/TensorRT-LLM/pull/10918) #10971

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yuxianq merged 1 commit into from
Jan 26, 2026
Merged

[https://nvbugs/5829097][fix] Re-init TRTLLM sampler to use sample stream in multi-stream cases. (cherry-pick from https://github.com/NVIDIA/TensorRT-LLM/pull/10918) #10971

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 12 additions & 3 deletions tensorrt_llm/_torch/pyexecutor/py_executor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,7 @@
from .model_engine import ModelEngine
from .resource_manager import ResourceManager
from .sampler import (AsyncWorkerMixin, Sampler, SamplerEvent, SampleState,
SampleStateTensors)
SampleStateTensors, TRTLLMSampler)
from .scheduler import (RequestScheduler, ScheduledRequests,
SerializableSchedulerOutput)

Expand Down Expand Up @@ -252,9 +252,19 @@ def __init__(self,
self.send_schedule_handler = None
self.pp_scheduler_max_retry_count = int(
os.environ.get("TLLM_PP_SCHEDULER_MAX_RETRY_COUNT", 10))
self.pp_multi_stream_sample = os.environ.get(
"TRTLLM_PP_MULTI_STREAM_SAMPLE", "1") == "1"
self.sample_stream = torch.cuda.Stream()
self.start_sample_event = torch.cuda.Event()
self.finish_sample_event = torch.cuda.Event()
if (self.dist.pp_size > 1 and self.pp_multi_stream_sample
and isinstance(self.sampler, TRTLLMSampler)):
# TRTLLM sampler uses default stream for store and algorithms.
# To enable multi-stream sampling, we need to re-initialize
# the sampler store and algorithms on the sample stream.
with torch.cuda.stream(self.sample_stream):
self.sampler._initialize_store()
self.sampler._instantiate_algorithms()

# Set of request IDs that are currently in flight across all micro batches.
# The scheduler will avoid scheduling requests that are already in flight.
Expand Down Expand Up @@ -1094,8 +1104,7 @@ def _executor_loop_pp(self):
guided_decoder_failed_requests = self.guided_decoder.execute(
batch_outputs['logits'])

if os.environ.get("TRTLLM_PP_MULTI_STREAM_SAMPLE",
"1") == "1":
if self.pp_multi_stream_sample:
# Wait for the previous sample to finish.
self.finish_sample_event.wait()
# Copy the batch outputs as sampler inputs
Expand Down
Loading