From 1cc01f1d74fddd8a21b82c91bdea8406860af3de Mon Sep 17 00:00:00 2001 From: Guiju Zhang <7135567+cascade812@users.noreply.github.com> Date: Tue, 21 Jul 2026 14:33:25 -0700 Subject: [PATCH 01/15] [None][perf] accumulate encoder work during decode Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 56 +++++++++++++++++-- 1 file changed, 50 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 4f4fc081a1d1..34b0067601bf 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -769,10 +769,10 @@ def __init__( # models, so encoder PP send/recv support is not implemented in the # PyTorch path for now. Reject pp_size > 1. # TODO: Add support for pp + encoder models - is_encoder_decoder = bool( + self.is_encoder_decoder = bool( getattr(getattr(self.model_engine.model, "model_config", None), "is_encoder_decoder", False)) - if is_encoder_decoder: + if self.is_encoder_decoder: if self.dist.pp_size > 1: raise NotImplementedError( "pp_size > 1 is not supported for encoder-decoder models " @@ -829,6 +829,7 @@ def __init__( self.adp_ctx_waiting_iters_count = 0 self.adp_ctx_batching_wait_iters_count = 0 self.batch_wait_iters_count = 0 + self.encoder_batch_wait_iters_count = 0 def on_detected(): # The graceful shutdown path can itself deadlock on collectives @@ -5228,6 +5229,37 @@ def _waiting_requests(self, context_requests: list[LlmRequest], self.batch_wait_iters_count = 0 return context_requests + def _waiting_encoder_requests( + self, encoder_requests: list[LlmRequest], + generation_requests: list[LlmRequest]) -> list[LlmRequest]: + """Accumulate encoder work while an admitted decode batch progresses. + + Encoder-decoder serving can otherwise launch one eager encoder forward + for each replacement request. Use the existing iteration deadline and + token threshold to form a larger encoder microbatch without blocking + the executor thread. Decoder generation continues while the encoder + requests wait. The encoder has its own counter because the resulting + decoder-context requests are already coalesced and must not wait for a + second window. + """ + if not encoder_requests or not generation_requests: + self.encoder_batch_wait_iters_count = 0 + return encoder_requests + + num_scheduled_tokens = sum(request.encoder_output_len + for request in encoder_requests) + num_scheduled_tokens += sum(1 + request.num_draft_tokens + for request in generation_requests) + should_wait = (self.encoder_batch_wait_iters_count + < self.batch_wait_timeout_iters and num_scheduled_tokens + < self.batch_wait_max_tokens_ratio * self.max_num_tokens) + if should_wait: + self.encoder_batch_wait_iters_count += 1 + return [] + + self.encoder_batch_wait_iters_count = 0 + return encoder_requests + @nvtx_range("_schedule") def _schedule(self): if hasattr(self.kv_cache_manager, "prepare_expect_snapshot_points"): @@ -5237,6 +5269,15 @@ def _schedule(self): scheduler_output = self.scheduler.schedule_request( self.active_requests, self.inflight_req_ids) + scheduled_encoder_requests = scheduler_output.encoder_requests + should_batch_encoder_requests = (self.is_encoder_decoder + and not self.enable_attention_dp + and self.enable_batch_waiting) + if should_batch_encoder_requests: + scheduled_encoder_requests = self._waiting_encoder_requests( + scheduler_output.encoder_requests, + scheduler_output.generation_requests) + scheduled_context_requests = scheduler_output.context_requests if self.enable_attention_dp and self.attention_dp_enable_balance: scheduled_context_requests = self._balance_adp_requests( @@ -5244,9 +5285,12 @@ def _schedule(self): scheduler_output.generation_requests) # If no generation requests, no need to wait, to avoid dead waiting - should_check_waiting = not self.enable_attention_dp and self.enable_batch_waiting and len( - scheduler_output.context_requests) > 0 and len( - scheduler_output.generation_requests) > 0 + should_check_waiting = (not self.is_encoder_decoder + and not self.enable_attention_dp + and self.enable_batch_waiting + and len(scheduler_output.context_requests) > 0 + and len( + scheduler_output.generation_requests) > 0) if should_check_waiting: # With KV cache manager V2, scheduling has already grown context request KV cache capacity. Requests dropped # for batch waiting still occupy KV cache and may reduce the batch size available for generation requests. @@ -5265,7 +5309,7 @@ def _schedule(self): num_fitting = len(scheduled_context_requests) scheduled_requests = ScheduledRequests() - scheduled_requests.encoder_requests = scheduler_output.encoder_requests + scheduled_requests.encoder_requests = scheduled_encoder_requests scheduled_requests.reset_context_requests(scheduled_context_requests) scheduled_requests.generation_requests = scheduler_output.generation_requests scheduled_requests.paused_requests = scheduler_output.paused_requests From 6fd8dc85d3eb773a0b15b831fe5c9684859ea5d4 Mon Sep 17 00:00:00 2001 From: Guiju Zhang <7135567+cascade812@users.noreply.github.com> Date: Tue, 21 Jul 2026 16:19:03 -0700 Subject: [PATCH 02/15] [None][perf] prepare mixed decoder batches natively Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com> --- .../nanobind/batch_manager/bindings.cpp | 149 ++++++++++ .../_torch/attention_backend/trtllm.py | 42 +++ .../_torch/pyexecutor/model_engine.py | 267 +++++++++++++++++- 3 files changed, 456 insertions(+), 2 deletions(-) diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp index 202228b394d8..e0263cd3178f 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp @@ -32,6 +32,7 @@ #include "tensorrt_llm/runtime/torchView.h" #include +#include #include #include #include @@ -40,6 +41,7 @@ #include #include #include +#include #include #include @@ -543,6 +545,153 @@ void initBindings(nb::module_& m) "Add new tokens to multiple LLM requests. The tokens vector should contain tokens for beam beam_idx of all " "requests in order."); + m.def( + "prepare_encoder_decoder_inputs", + [](std::vector> const& contextRequests, + std::vector> const& generationRequests, at::Tensor const& inputIds, + at::Tensor const& positionIds, at::Tensor const& sequenceLengths, at::Tensor const& promptLengths, + at::Tensor const& cachedTokenLengths, at::Tensor const& kvLengths, at::Tensor const& encoderKvLengths, + at::Tensor const& previousBatchIndices, SizeType32 positionIdOffset) + { + auto checkIntBuffer = [](at::Tensor const& tensor, char const* name) + { + TLLM_CHECK_WITH_INFO(tensor.device().is_cpu(), "%s must be a CPU tensor", name); + TLLM_CHECK_WITH_INFO(tensor.scalar_type() == at::kInt, "%s must have torch.int32 dtype", name); + TLLM_CHECK_WITH_INFO(tensor.is_contiguous(), "%s must be contiguous", name); + }; + checkIntBuffer(inputIds, "input_ids"); + checkIntBuffer(positionIds, "position_ids"); + checkIntBuffer(sequenceLengths, "sequence_lengths"); + checkIntBuffer(promptLengths, "prompt_lengths"); + checkIntBuffer(cachedTokenLengths, "cached_token_lengths"); + checkIntBuffer(kvLengths, "kv_lengths"); + checkIntBuffer(encoderKvLengths, "encoder_kv_lengths"); + checkIntBuffer(previousBatchIndices, "previous_batch_indices"); + + auto const numSequences = contextRequests.size() + generationRequests.size(); + TLLM_CHECK_WITH_INFO(sequenceLengths.numel() >= static_cast(numSequences), + "sequence_lengths capacity is smaller than the batch"); + TLLM_CHECK_WITH_INFO(promptLengths.numel() >= static_cast(numSequences), + "prompt_lengths capacity is smaller than the batch"); + TLLM_CHECK_WITH_INFO(cachedTokenLengths.numel() >= static_cast(numSequences), + "cached_token_lengths capacity is smaller than the batch"); + TLLM_CHECK_WITH_INFO(kvLengths.numel() >= static_cast(numSequences), + "kv_lengths capacity is smaller than the batch"); + TLLM_CHECK_WITH_INFO(encoderKvLengths.numel() >= static_cast(numSequences), + "encoder_kv_lengths capacity is smaller than the batch"); + TLLM_CHECK_WITH_INFO(previousBatchIndices.numel() >= static_cast(generationRequests.size()), + "previous_batch_indices capacity is smaller than the generation batch"); + + auto* inputIdsPtr = inputIds.data_ptr(); + auto* positionIdsPtr = positionIds.data_ptr(); + auto* sequenceLengthsPtr = sequenceLengths.data_ptr(); + auto* promptLengthsPtr = promptLengths.data_ptr(); + auto* cachedTokenLengthsPtr = cachedTokenLengths.data_ptr(); + auto* kvLengthsPtr = kvLengths.data_ptr(); + auto* encoderKvLengthsPtr = encoderKvLengths.data_ptr(); + auto* previousBatchIndicesPtr = previousBatchIndices.data_ptr(); + + std::vector requestIds; + std::vector encoderSequenceLengths; + std::vector encoderCachedTokenLengths; + requestIds.reserve(numSequences); + encoderSequenceLengths.reserve(numSequences); + encoderCachedTokenLengths.reserve(numSequences); + + SizeType32 numTokens{0}; + SizeType32 numContextTokens{0}; + SizeType32 numPreviousBatchRequests{0}; + SizeType32 cachedKvTokens{0}; + SizeType32 contextKvTokens{0}; + SizeType32 generationKvTokens{0}; + SizeType32 maxKvLength{0}; + SizeType32 contextEncoderKvTokens{0}; + SizeType32 generationEncoderKvTokens{0}; + SizeType32 maxEncoderKvLength{0}; + for (auto const& request : contextRequests) + { + auto const sequenceIdx = requestIds.size(); + auto const begin = request->getContextCurrentPosition(); + auto const chunkSize = request->getContextChunkSize(); + auto const& tokens = request->getTokens(0); + TLLM_CHECK_WITH_INFO(begin + chunkSize <= static_cast(tokens.size()), + "Context chunk exceeds the request token count"); + TLLM_CHECK_WITH_INFO(inputIds.numel() >= static_cast(numTokens + chunkSize), + "input_ids capacity is smaller than the packed context"); + TLLM_CHECK_WITH_INFO(positionIds.numel() >= static_cast(numTokens + chunkSize), + "position_ids capacity is smaller than the packed context"); + + std::copy_n(tokens.data() + begin, chunkSize, inputIdsPtr + numTokens); + std::iota(positionIdsPtr + numTokens, positionIdsPtr + numTokens + chunkSize, begin + positionIdOffset); + sequenceLengthsPtr[sequenceIdx] = chunkSize; + promptLengthsPtr[sequenceIdx] = chunkSize; + cachedTokenLengthsPtr[sequenceIdx] = begin; + auto const kvLength = begin + chunkSize; + kvLengthsPtr[sequenceIdx] = kvLength; + auto const encoderKvLength = request->getEncoderOutputLen(); + encoderKvLengthsPtr[sequenceIdx] = encoderKvLength; + cachedKvTokens += begin; + contextKvTokens += kvLength; + maxKvLength = std::max(maxKvLength, kvLength); + contextEncoderKvTokens += encoderKvLength; + maxEncoderKvLength = std::max(maxEncoderKvLength, encoderKvLength); + numTokens += chunkSize; + numContextTokens += chunkSize; + + requestIds.push_back(request->mRequestId); + encoderSequenceLengths.push_back(encoderKvLength); + encoderCachedTokenLengths.push_back(0); + } + + bool sawDummyRequest{false}; + for (auto const& request : generationRequests) + { + auto const sequenceIdx = requestIds.size(); + auto const isDummy = request->isDummyRequest(); + sawDummyRequest = sawDummyRequest || isDummy; + TLLM_CHECK_WITH_INFO( + isDummy || !sawDummyRequest, "CUDA graph dummy requests must follow real generation requests"); + TLLM_CHECK_WITH_INFO( + positionIds.numel() > numTokens, "position_ids capacity is smaller than the packed batch"); + + auto const pastSeenTokens = request->getMaxBeamNumTokens() - (isDummy ? 1 : 0); + positionIdsPtr[numTokens] = pastSeenTokens + positionIdOffset; + sequenceLengthsPtr[sequenceIdx] = 1; + promptLengthsPtr[sequenceIdx] = request->mPromptLen; + cachedTokenLengthsPtr[sequenceIdx] = pastSeenTokens; + auto const kvLength = pastSeenTokens + 1; + kvLengthsPtr[sequenceIdx] = kvLength; + auto const encoderKvLength = request->getEncoderOutputLen(); + encoderKvLengthsPtr[sequenceIdx] = encoderKvLength; + cachedKvTokens += pastSeenTokens; + generationKvTokens += kvLength; + maxKvLength = std::max(maxKvLength, kvLength); + generationEncoderKvTokens += encoderKvLength; + maxEncoderKvLength = std::max(maxEncoderKvLength, encoderKvLength); + ++numTokens; + + if (!isDummy) + { + TLLM_CHECK_WITH_INFO( + request->mSeqSlot.has_value(), "A real generation request must have a sequence slot"); + previousBatchIndicesPtr[numPreviousBatchRequests++] = request->mSeqSlot.value(); + } + + requestIds.push_back(request->mRequestId); + encoderSequenceLengths.push_back(0); + encoderCachedTokenLengths.push_back(encoderKvLength); + } + + return std::make_tuple(requestIds, encoderSequenceLengths, encoderCachedTokenLengths, numTokens, + numContextTokens, numPreviousBatchRequests, cachedKvTokens, contextKvTokens, generationKvTokens, + maxKvLength, contextEncoderKvTokens, generationEncoderKvTokens, maxEncoderKvLength); + }, + nb::arg("context_requests"), nb::arg("generation_requests"), nb::arg("input_ids"), nb::arg("position_ids"), + nb::arg("sequence_lengths"), nb::arg("prompt_lengths"), nb::arg("cached_token_lengths"), nb::arg("kv_lengths"), + nb::arg("encoder_kv_lengths"), nb::arg("previous_batch_indices"), nb::arg("position_id_offset") = 0, + nb::call_guard(), + "Prepare the persistent CPU input buffers for a simple encoder-decoder batch."); + m.def( "make_decoding_batch_input", [](tb::DecoderInputBuffers& decoderInputBuffers, runtime::decoder::DecoderState& decoderState, diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index 40d15970398b..26ec5db43d15 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -641,6 +641,48 @@ def prepare(self) -> None: host_request_types=self.host_request_types[:self.num_seqs], ) + def prepare_encoder_decoder(self, prompt_lens: torch.Tensor, + kv_lens: torch.Tensor, context_kv_tokens: int, + generation_kv_tokens: int, + max_kv_len: int) -> None: + """Prepare simple encoder-decoder attention from native host buffers.""" + super().prepare() + extra_attrs = get_model_extra_attrs() + if extra_attrs is None: + get_global_attrs().attention_metadata = weakref.ref(self) + + assert self.kv_cache_manager is not None + assert self.draft_kv_cache_manager is None + assert not self.is_spec_decoding_enabled + assert self.kv_cache_params.num_extra_kv_tokens == 0 + assert not self.enable_flash_mla + assert not self.enable_helix + assert not self.enable_context_mla_with_cached_kv + assert self.request_ids is not None + assert max_kv_len <= self.kv_cache_manager.max_seq_len, ( + f"The max KV cache length of input sequences ({max_kv_len}) " + "exceeds the KV cache manager's maximum supported length " + f"({self.kv_cache_manager.max_seq_len}).") + + num_seqs = self.num_seqs + self.prompt_lens_cuda[:num_seqs].copy_(prompt_lens, non_blocking=True) + self.kv_lens_cuda[:num_seqs].copy_(kv_lens, non_blocking=True) + self.host_total_kv_lens[0] = context_kv_tokens + self.host_total_kv_lens[1] = generation_kv_tokens + self.host_request_types[:self.num_contexts].fill_(0) + self.host_request_types[self.num_contexts:num_seqs].fill_(1) + + self.kv_cache_manager.copy_batch_block_offsets( + self.kv_cache_block_offsets, self.request_ids, self.beam_width, + self.num_contexts, num_seqs) + self._bind_runtime_views( + kv_lens_cuda=self.kv_lens_cuda[:num_seqs], + kv_lens=kv_lens, + prompt_lens_cuda=self.prompt_lens_cuda[:num_seqs], + prompt_lens_cpu=prompt_lens, + host_request_types=self.host_request_types[:num_seqs], + ) + def prepare_encoder_only(self) -> None: """Fast path for encoder-only forward (eager + CUDA graph capture).""" extra_attrs = get_model_extra_attrs() diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 3583929a407c..1bdafaab9440 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -21,6 +21,8 @@ from tensorrt_llm._utils import (is_trace_enabled, maybe_pin_memory, nvtx_range, prefer_pinned, release_gc, torch_dtype_to_str, trace_func) +from tensorrt_llm.bindings.internal import \ + batch_manager as batch_manager_bindings from tensorrt_llm.bindings.internal.runtime import TaskLayerModuleConfig from tensorrt_llm.inputs.multimodal import (MultimodalInput, MultimodalParams, MultimodalRuntimeData, @@ -3023,6 +3025,10 @@ def _prepare_enc_dec_cross_attn_inputs( encoder_num_cached_tokens_per_seq: List[int], attn_metadata: AttentionMetadata, resource_manager: Optional[ResourceManager], + encoder_kv_lens: Optional[torch.Tensor] = None, + context_encoder_kv_tokens: int = 0, + generation_encoder_kv_tokens: int = 0, + max_encoder_kv_len: int = 0, ) -> Dict[str, Any]: if not encoder_seq_lens: return {} @@ -3063,6 +3069,19 @@ def _prepare_enc_dec_cross_attn_inputs( packed_encoder_hidden_states = None skip_cross_kv_projection = True + def prepare_cross_metadata( + cross_attn_metadata: AttentionMetadata) -> None: + if encoder_kv_lens is None: + cross_attn_metadata.prepare() + return + assert isinstance(cross_attn_metadata, TrtllmAttentionMetadata) + cross_attn_metadata.prepare_encoder_decoder( + prompt_lens=attn_metadata.prompt_lens, + kv_lens=encoder_kv_lens, + context_kv_tokens=context_encoder_kv_tokens, + generation_kv_tokens=generation_encoder_kv_tokens, + max_kv_len=max_encoder_kv_len) + if attn_metadata.is_cuda_graph and attn_metadata.has_cross_sub_metadata: # Fast path for stable CUDA-graph generation steps: the encoder # KV lengths (kv_lens_cuda) and the frozen prompt lengths @@ -3092,7 +3111,7 @@ def _prepare_enc_dec_cross_attn_inputs( encoder_num_cached_tokens_per_seq= encoder_num_cached_tokens_per_seq, ) - cross_attn_metadata.prepare() + prepare_cross_metadata(cross_attn_metadata) if new_encoder_tokens == 0: # Record this stable state for future fast-path use. self._cross_attn_stable_cached_tokens = list( @@ -3123,7 +3142,7 @@ def _prepare_enc_dec_cross_attn_inputs( else: self._cross_attn_stable_cached_tokens = None self._cross_attn_stable_request_ids = None - cross_attn_metadata.prepare() + prepare_cross_metadata(cross_attn_metadata) return { "encoder_hidden_states": packed_encoder_hidden_states, @@ -3159,6 +3178,242 @@ def _ship_multimodal_indices( inputs['text_token_indices'] = text_token_indices_cpu.to( "cuda", non_blocking=True) + def _can_use_encoder_decoder_input_fast_path( + self, scheduled_requests: ScheduledRequests, + new_tokens_device: Optional[torch.Tensor], + next_draft_tokens_device: Optional[torch.Tensor]) -> bool: + """Return whether the TRT-like persistent input path is sufficient.""" + static_eligible = getattr( + self, '_encoder_decoder_input_fast_path_static_eligible', None) + if static_eligible is None: + static_eligible = ( + self._is_encoder_decoder_model() and not self.enable_spec_decode + and not self.is_draft_model and self.max_beam_width == 1 + and self.sparse_attention_config is None and not self.use_mrope + and not self.enable_attention_dp + and not self.mapping.has_cp_helix() and not self.is_multimodal + and self.lora_model_config is None + and not self.attn_runtime_features.chunked_prefill + and not self.attn_runtime_features.cache_reuse + and not self.attn_runtime_features.has_speculative_draft_tokens) + self._encoder_decoder_input_fast_path_static_eligible = \ + static_eligible + if (not static_eligible or new_tokens_device is None + or next_draft_tokens_device is not None + or self.guided_decoder is not None): + return False + + if scheduled_requests.batch_size == 0: + return False + for request in scheduled_requests.generation_requests: + if request.py_batch_idx is None and not request.is_dummy: + return False + return True + + def _acquire_encoder_decoder_host_buffers(self) -> Dict[str, Any]: + """Acquire pinned staging whose preceding asynchronous copies finished.""" + pool = getattr(self, '_encoder_decoder_host_buffer_pool', None) + if pool is None: + pool = [] + self._encoder_decoder_host_buffer_pool = pool + for buffers in pool: + event = buffers['event'] + if event is None or event.query(): + return buffers + + buffers = { + 'input_ids': + torch.empty(self.max_num_tokens, + dtype=torch.int, + pin_memory=prefer_pinned()), + 'position_ids': + torch.empty(self.max_num_tokens, + dtype=torch.int, + pin_memory=prefer_pinned()), + 'sequence_lengths': + torch.empty(self.batch_size, + dtype=torch.int, + pin_memory=prefer_pinned()), + 'prompt_lengths': + torch.empty(self.batch_size, + dtype=torch.int, + pin_memory=prefer_pinned()), + 'cached_token_lengths': + torch.empty(self.batch_size, + dtype=torch.int, + pin_memory=prefer_pinned()), + 'kv_lengths': + torch.empty(self.batch_size, + dtype=torch.int, + pin_memory=prefer_pinned()), + 'encoder_kv_lengths': + torch.empty(self.batch_size, + dtype=torch.int, + pin_memory=prefer_pinned()), + 'previous_batch_indices': + torch.empty(self.batch_size, + dtype=torch.int, + pin_memory=prefer_pinned()), + 'event': + None, + } + pool.append(buffers) + return buffers + + @nvtx_range("_prepare_encoder_decoder_inputs_fast") + def _prepare_encoder_decoder_inputs_fast( + self, scheduled_requests: ScheduledRequests, + kv_cache_manager: Union[KVCacheManager, KVCacheManagerV2], + attn_metadata: AttentionMetadata, new_tokens_device: torch.Tensor, + resource_manager: Optional[ResourceManager]): + """Prepare a simple BART batch with native collation and reused buffers.""" + buffers = self._acquire_encoder_decoder_host_buffers() + position_id_offset = getattr(self, + '_encoder_decoder_position_id_offset', + None) + if position_id_offset is None: + position_id_offset = self._get_position_id_offset() + self._encoder_decoder_position_id_offset = position_id_offset + (request_ids, encoder_seq_lens, encoder_cached_token_lengths, + total_num_tokens, num_context_tokens, num_previous_batch_requests, + cached_kv_tokens, context_kv_tokens, generation_kv_tokens, max_kv_len, + context_encoder_kv_tokens, generation_encoder_kv_tokens, + max_encoder_kv_len + ) = batch_manager_bindings.prepare_encoder_decoder_inputs( + scheduled_requests.context_requests, + scheduled_requests.generation_requests, + buffers['input_ids'], + buffers['position_ids'], + buffers['sequence_lengths'], + buffers['prompt_lengths'], + buffers['cached_token_lengths'], + buffers['kv_lengths'], + buffers['encoder_kv_lengths'], + buffers['previous_batch_indices'], + position_id_offset, + ) + + num_sequences = scheduled_requests.batch_size + num_generation_requests = scheduled_requests.num_generation_requests + if num_context_tokens: + self.input_ids_cuda[:num_context_tokens].copy_( + buffers['input_ids'][:num_context_tokens], non_blocking=True) + if num_previous_batch_requests: + previous_slots = self.previous_batch_indices_cuda[: + num_previous_batch_requests] + previous_slots.copy_( + buffers['previous_batch_indices'][:num_previous_batch_requests], + non_blocking=True) + new_tokens = new_tokens_device[0, previous_slots, 0] + generation_begin = num_context_tokens + generation_end = generation_begin + num_previous_batch_requests + self.input_ids_cuda[generation_begin:generation_end].copy_( + new_tokens, non_blocking=True) + dummy_begin = num_context_tokens + num_previous_batch_requests + if dummy_begin < total_num_tokens: + self.input_ids_cuda[dummy_begin:total_num_tokens].fill_(0) + + self.position_ids_cuda[:total_num_tokens].copy_( + buffers['position_ids'][:total_num_tokens], non_blocking=True) + final_position_ids = self.position_ids_cuda[: + total_num_tokens].unsqueeze( + 0) + + sequence_lengths = buffers['sequence_lengths'][:num_sequences] + attn_metadata._seq_lens = sequence_lengths + if attn_metadata.is_cuda_graph and attn_metadata._seq_lens_cuda is not None: + attn_metadata._seq_lens_cuda.copy_(sequence_lengths, + non_blocking=True) + else: + attn_metadata._seq_lens_cuda = sequence_lengths.cuda( + non_blocking=True) + attn_metadata._num_contexts = scheduled_requests.num_context_requests + attn_metadata._num_ctx_tokens = num_context_tokens + attn_metadata._num_generations = num_generation_requests + attn_metadata._num_tokens = total_num_tokens + attn_metadata.beam_width = 1 + attn_metadata.request_ids = request_ids + attn_metadata.prompt_lens = buffers['prompt_lengths'][:num_sequences] + attn_metadata.num_chunked_ctx_requests = 0 + attn_metadata.kv_cache_params = KVCacheParams( + use_cache=True, + num_cached_tokens_per_seq=buffers['cached_token_lengths'] + [:num_sequences], + num_extra_kv_tokens=0) + attn_metadata.kv_cache_manager = kv_cache_manager + assert isinstance(attn_metadata, TrtllmAttentionMetadata) + attn_metadata.prepare_encoder_decoder( + prompt_lens=buffers['prompt_lengths'][:num_sequences], + kv_lens=buffers['kv_lengths'][:num_sequences], + context_kv_tokens=context_kv_tokens, + generation_kv_tokens=generation_kv_tokens, + max_kv_len=max_kv_len) + + encoder_hidden_states = [] + for request in scheduled_requests.context_requests: + encoder_output = request.py_encoder_output + if encoder_output is None: + raise RuntimeError( + f"Decoder context request {request.py_request_id} has no " + "encoder output.") + encoder_hidden_states.append(encoder_output) + request.py_batch_idx = request.py_seq_slot + + cross_attention_inputs = self._prepare_enc_dec_cross_attn_inputs( + encoder_hidden_states, + encoder_seq_lens, + encoder_cached_token_lengths, + attn_metadata, + resource_manager, + encoder_kv_lens=buffers['encoder_kv_lengths'][:num_sequences], + context_encoder_kv_tokens=context_encoder_kv_tokens, + generation_encoder_kv_tokens=generation_encoder_kv_tokens, + max_encoder_kv_len=max_encoder_kv_len, + ) + + attn_all_rank_num_tokens = self._get_all_rank_num_tokens(attn_metadata) + padded_num_tokens, can_run_piecewise_cuda_graph, attn_all_rank_num_tokens = self._get_padding_params( + total_num_tokens, scheduled_requests.num_context_requests, + attn_all_rank_num_tokens) + set_per_request_piecewise_cuda_graph_flag(can_run_piecewise_cuda_graph) + attn_metadata.padded_num_tokens = (padded_num_tokens + if padded_num_tokens + != total_num_tokens else None) + + virtual_num_tokens = total_num_tokens + if attn_metadata.padded_num_tokens is not None: + self.input_ids_cuda[total_num_tokens:padded_num_tokens].fill_(0) + self.position_ids_cuda[total_num_tokens:padded_num_tokens].fill_(0) + virtual_num_tokens = padded_num_tokens + final_position_ids = self.position_ids_cuda[: + virtual_num_tokens].unsqueeze( + 0) + + inputs = { + 'attn_metadata': attn_metadata, + 'input_ids': self.input_ids_cuda[:virtual_num_tokens], + 'position_ids': final_position_ids, + 'inputs_embeds': None, + 'multimodal_params': [], + 'resource_manager': resource_manager, + } + inputs.update(cross_attention_inputs) + + self.iter_states[ + 'num_ctx_requests'] = scheduled_requests.num_context_requests + self.iter_states['num_ctx_tokens'] = num_context_tokens + self.iter_states['num_generation_tokens'] = num_generation_requests + self.iter_states['cached_kv_tokens'] = cached_kv_tokens + if not self.is_warmup: + self.previous_request_ids = request_ids[scheduled_requests. + num_context_requests:] + self.has_previous_device_draft = False + + event = torch.cuda.Event() + event.record(torch.cuda.current_stream()) + buffers['event'] = event + return inputs, None + def _can_use_incremental_update( self, scheduled_requests: ScheduledRequests, new_tokens_device: Optional[torch.Tensor], @@ -3837,6 +4092,14 @@ def _prepare_tp_inputs( num_accepted_tokens_device, req_id_to_old_request, resource_manager) + if (type(attn_metadata) is TrtllmAttentionMetadata + and self._can_use_encoder_decoder_input_fast_path( + scheduled_requests, new_tokens_device, + next_draft_tokens_device)): + return self._prepare_encoder_decoder_inputs_fast( + scheduled_requests, kv_cache_manager, attn_metadata, + new_tokens_device, resource_manager) + if self._can_use_steady_gen_fast_prepare(scheduled_requests, new_tokens_device, next_draft_tokens_device, From 433dcb98ed25c2ebf1a9839dfd2f9143e423235f Mon Sep 17 00:00:00 2001 From: Guiju Zhang <7135567+cascade812@users.noreply.github.com> Date: Tue, 21 Jul 2026 18:55:13 -0700 Subject: [PATCH 03/15] [None][perf] reuse stable greedy sampling metadata Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com> --- .../_torch/pyexecutor/sampler/sampler.py | 161 ++++++++++++++++-- .../_torch/sampler/test_torch_sampler.py | 80 +++++++++ 2 files changed, 226 insertions(+), 15 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index bb50aab8c351..b005c2c347b6 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -1114,6 +1114,7 @@ def finish_reasons_list(self) -> FinishReasonsList: @dataclass(kw_only=True) class SampleStateTorch(SampleState[SampleStateTensorsHostTorch, SampleStateTensors]): beam_history_builders: list[BeamHistoryBuilder | None] | None = None + single_step_greedy: bool = False @dataclass(kw_only=True, frozen=True) @@ -2412,6 +2413,9 @@ def __init__(self, args: Args): self._prev_first_finish_reasons_host: list[torch.Tensor | None] = [ None ] * self.max_num_sequences + self._stable_greedy_request_ids: list[int] = [] + self._stable_greedy_seq_slots_host: Optional[torch.Tensor] = None + self._stable_greedy_seq_slots_cuda: Optional[torch.Tensor] = None @staticmethod def _is_draft_batch(requests: list[LlmRequest]) -> bool: @@ -3654,6 +3658,12 @@ def update_requests( self._pending_steps[slot] -= 1 assert state.host is not None + # Reuse sample_async's qualification instead of rechecking every + # request after the asynchronous sample completes. + if state.single_step_greedy: + self._update_requests_single_beam_single_step(state) + return + new_tokens = state.host.new_tokens finish_reasons = state.host.finish_reasons_list() first_finish_reasons = ( @@ -3749,6 +3759,48 @@ def _maybe_build_beam_history(req_idx: int) -> BeamHistory | None: req.py_seq_slot ] = req.py_num_accepted_draft_tokens + @nvtx_range("_update_requests_single_beam_single_step") + def _update_requests_single_beam_single_step(self, state: SampleStateTorch) -> None: + """Update the common greedy, single-token case without draft machinery.""" + assert state.host is not None + requests = [ + request + for request in state.requests + if request.state != LlmRequestState.GENERATION_COMPLETE + ] + if not requests: + return + + all_new_tokens = state.host.new_tokens.tolist() + if len(requests) == len(state.requests): + new_tokens = all_new_tokens + else: + new_tokens = [ + new_token + for request, new_token in zip(state.requests, all_new_tokens) + if request.state != LlmRequestState.GENERATION_COMPLETE + ] + add_new_tokens_to_requests(requests, new_tokens, DEFAULT_BEAM_IDX) + + # sample_async deliberately omits the device finish-reason tensor for + # this qualified path; completion is derived from compact host tokens. + assert state.host.finish_reasons is None + for request, new_token in zip(requests, new_tokens): + # The stable greedy path excludes stop words. Keep EOS ahead of the + # length check so a terminal EOS at the token limit is reported as + # END_ID, matching _handle_stop_criteria. + if new_token == request.py_end_id: + request.finish_by(FinishReason.END_ID, DEFAULT_BEAM_IDX) + elif ( + request.max_beam_num_tokens - request.py_orig_prompt_len + >= request.py_max_new_tokens + or request.max_beam_num_tokens >= self.max_seq_len + ): + request.finish_by(FinishReason.LENGTH, DEFAULT_BEAM_IDX) + request.py_num_accepted_draft_tokens = 0 + request.py_rewind_len = 0 + request.py_decoding_iter += 1 + def _return_log_probs(self, requests: list[LlmRequest]) -> bool: return any(req.py_return_log_probs for req in requests) @@ -3806,6 +3858,7 @@ def sample_async( seq_slots_cuda, seq_lens_cuda, new_tokens_host, + single_step_greedy, ) = self._process_requests( scheduled_requests, model_outputs, @@ -3832,22 +3885,26 @@ def sample_async( # their buffers in the store. # Assume that either all requests are drafts or none are drafts is_draft_batch = requests[0].py_is_draft - finish_reasons_device = self._finish_reasons_handler.write_finish_reasons( - seq_slots_host=seq_slots_host, - is_draft_batch=is_draft_batch, - seq_slots_cuda=seq_slots_cuda, - seq_lens_cuda=seq_lens_cuda, - new_tokens_cuda=new_tokens, - first_finish_reasons_cuda=( - beam_search_store.first_finish_reasons - if beam_search_store is not None - else None - ), - ) - finish_reasons_host = self._copy_to_host(finish_reasons_device) + if not single_step_greedy: + assert seq_lens_host is not None + assert seq_lens_cuda is not None + finish_reasons_device = self._finish_reasons_handler.write_finish_reasons( + seq_slots_host=seq_slots_host, + is_draft_batch=is_draft_batch, + seq_slots_cuda=seq_slots_cuda, + seq_lens_cuda=seq_lens_cuda, + new_tokens_cuda=new_tokens, + first_finish_reasons_cuda=( + beam_search_store.first_finish_reasons + if beam_search_store is not None + else None + ), + ) + finish_reasons_host = self._copy_to_host(finish_reasons_device) if self._use_beam_search: assert beam_search_store is not None + assert seq_lens_cuda is not None first_finish_reasons = beam_search_store.first_finish_reasons first_finish_reasons_host = self._copy_to_host(first_finish_reasons) self._update_original_tokens( @@ -3890,6 +3947,7 @@ def sample_async( ), sampler_event=sampler_event, beam_history_builders=beam_history_builders, + single_step_greedy=single_step_greedy, ) @staticmethod @@ -3910,7 +3968,7 @@ def _fast_greedy_sample_kernel( batch_dest_indices: torch.Tensor, max_beam_width: int, d2t: torch.Tensor | None, - ) -> None: + ) -> torch.Tensor: """Applies fast greedy sampling to the logits. Performs argmax, applies d2t translation if present, and scatters @@ -3929,6 +3987,7 @@ def _fast_greedy_sample_kernel( new_tokens_cuda.view(-1, *new_tokens_cuda.shape[2:]).scatter_( 0, batch_dest_indices_expanded, next_tokens_expanded ) + return next_tokens @staticmethod def _apply_embedding_bias( @@ -4799,10 +4858,80 @@ def _process_requests( new_tokens_cuda: torch.Tensor, num_context_logits_prefix_sum: list[int], ) -> tuple[ - list[LlmRequest], torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor + list[LlmRequest], + torch.Tensor, + Optional[torch.Tensor], + torch.Tensor, + Optional[torch.Tensor], + torch.Tensor, + bool, ]: raw_logits_cuda = model_outputs["logits"] + generation_requests = scheduled_requests.generation_requests + request_ids = [request.py_request_id for request in generation_requests] + has_stable_request_ids = self._stable_greedy_request_ids == request_ids + can_use_stable_greedy_path = ( + bool(generation_requests) + and self.max_beam_width == 1 + and scheduled_requests.num_context_requests == 0 + and len(generation_requests) <= raw_logits_cuda.shape[0] + and model_outputs.get("d2t") is None + and ( + has_stable_request_ids + or all( + not request.is_dummy + and get_draft_token_length(request) == 0 + and request._py_embedding_bias_1d is None + and not getattr(request, "py_bad_words", None) + and not request.py_min_length + and not request.py_return_log_probs + and not request.py_stop_words_list + and _request_strategy(request, vocab_size=2**31) == GREEDY + for request in generation_requests + ) + ) + ) + if can_use_stable_greedy_path: + if has_stable_request_ids: + assert self._stable_greedy_seq_slots_host is not None + assert self._stable_greedy_seq_slots_cuda is not None + seq_slots_host = self._stable_greedy_seq_slots_host + seq_slots_cuda = self._stable_greedy_seq_slots_cuda + else: + maybe_seq_slots = [request.py_seq_slot for request in generation_requests] + assert all(seq_slot is not None for seq_slot in maybe_seq_slots) + seq_slots = [cast(int, seq_slot) for seq_slot in maybe_seq_slots] + seq_slots_host = torch.tensor( + seq_slots, dtype=torch.int32, pin_memory=prefer_pinned() + ) + seq_slots_cuda = seq_slots_host.to( + device="cuda", dtype=torch.int64, non_blocking=True + ) + self._stable_greedy_request_ids = request_ids + self._stable_greedy_seq_slots_host = seq_slots_host + self._stable_greedy_seq_slots_cuda = seq_slots_cuda + + next_tokens = self._fast_greedy_sample_kernel( + raw_logits_cuda[: len(generation_requests)], + new_tokens_cuda, + seq_slots_cuda, + self.max_beam_width, + None, + ) + new_tokens_host = self._copy_to_host(next_tokens) + return ( + generation_requests, + seq_slots_host, + None, + seq_slots_cuda, + None, + new_tokens_host, + True, + ) + + self._stable_greedy_request_ids = [] + sampling_requests, sampling_requests_metadata, logits_cuda = self._select_generated_logits( scheduled_requests, raw_logits_cuda, @@ -4910,6 +5039,7 @@ def _process_requests( seq_slots_cuda, seq_lens_cuda, new_tokens_host, + False, ) # Indexer for accessing tokens in 'logits_cuda', corresponding to the @@ -4963,6 +5093,7 @@ def _process_requests( seq_slots_cuda, seq_lens_cuda, new_tokens_host, + False, ) @override diff --git a/tests/unittest/_torch/sampler/test_torch_sampler.py b/tests/unittest/_torch/sampler/test_torch_sampler.py index e2e0d1bc9402..560e4431e617 100644 --- a/tests/unittest/_torch/sampler/test_torch_sampler.py +++ b/tests/unittest/_torch/sampler/test_torch_sampler.py @@ -42,6 +42,8 @@ get_draft_token_length, ) from tensorrt_llm._torch.pyexecutor.sampler import ( + SampleStateTensorsHostTorch, + SampleStateTorch, TorchSampler, _BatchedSamplingResult, _request_get_sampling_params, @@ -659,6 +661,84 @@ class TestFinishReasons: END_ID = FinishReason.END_ID LENGTH = FinishReason.LENGTH + def test_single_step_greedy_checks_finish_reasons_on_host(self): + sampler = object.__new__(TorchSampler) + sampler.max_seq_len = 20 + sampler._track_pending_steps = False + requests = [ + LlmRequest( + request_id=0, + seq_slot=0, + input_tokens=[2, 0], + max_new_tokens=1, + end_id=2, + sampling_config=SamplingConfig(), + is_streaming=False, + ), + LlmRequest( + request_id=1, + seq_slot=1, + input_tokens=[2, 0], + max_new_tokens=1, + end_id=2, + sampling_config=SamplingConfig(), + is_streaming=False, + ), + ] + new_tokens = torch.tensor([2, 7], dtype=torch.int32) + state = SampleStateTorch( + requests=requests, + device=None, + host=SampleStateTensorsHostTorch( + new_tokens=new_tokens, + finish_reasons=None, + first_finish_reasons=None, + ), + single_step_greedy=True, + ) + + sampler.update_requests(state) + + assert all(request.is_finished for request in requests) + # The first request reaches EOS and length together; EOS takes precedence. + assert not requests[0].is_finished_due_to_length + assert requests[1].is_finished_due_to_length + assert requests[0].get_tokens(0)[-1] == 2 + assert requests[1].get_tokens(0)[-1] == 7 + + def test_single_step_greedy_filters_requests_completed_after_sampling(self): + sampler = object.__new__(TorchSampler) + sampler.max_seq_len = 20 + sampler._track_pending_steps = False + requests = [ + LlmRequest( + request_id=request_id, + seq_slot=request_id, + input_tokens=[2, 0], + max_new_tokens=10, + end_id=2, + sampling_config=SamplingConfig(), + is_streaming=False, + ) + for request_id in range(2) + ] + requests[0].finish_by(FinishReason.LENGTH, 0) + state = SampleStateTorch( + requests=requests, + device=None, + host=SampleStateTensorsHostTorch( + new_tokens=torch.tensor([99, 7], dtype=torch.int32), + finish_reasons=None, + first_finish_reasons=None, + ), + single_step_greedy=True, + ) + + sampler.update_requests(state) + + assert requests[0].get_tokens(0) == [2, 0] + assert requests[1].get_tokens(0)[-1] == 7 + class RequestCase: MAX_NEW_TOKENS = 10 MAX_NUM_SEQUENCES = 128 From f3a8e105f2d5ea6d380ac29a488ed26d0f1a43c7 Mon Sep 17 00:00:00 2001 From: Guiju Zhang <7135567+cascade812@users.noreply.github.com> Date: Sun, 26 Jul 2026 10:17:14 -0700 Subject: [PATCH 04/15] [None][perf] optimize BART encoder-decoder execution Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com> --- ...-launch-path-optimization-opportunities.md | 679 ++++++++++++++++++ optimize_plan_2.md | 494 +++++++++++++ .../_torch/pyexecutor/cuda_graph_runner.py | 307 ++++++-- tensorrt_llm/_torch/pyexecutor/llm_request.py | 3 + .../_torch/pyexecutor/model_engine.py | 586 ++++++++++++--- tensorrt_llm/_torch/pyexecutor/py_executor.py | 256 +++++-- .../_torch/pyexecutor/scheduler/scheduler.py | 3 + .../test_encoder_cuda_graph_runner.py | 225 ++++++ .../test_mixed_decoder_cuda_graph_runner.py | 79 ++ .../_torch/executor/test_py_executor.py | 217 +++++- 10 files changed, 2626 insertions(+), 223 deletions(-) create mode 100644 docs/source/models/bart-decoder-launch-path-optimization-opportunities.md create mode 100644 optimize_plan_2.md create mode 100644 tests/unittest/_torch/executor/test_encoder_cuda_graph_runner.py create mode 100644 tests/unittest/_torch/executor/test_mixed_decoder_cuda_graph_runner.py diff --git a/docs/source/models/bart-decoder-launch-path-optimization-opportunities.md b/docs/source/models/bart-decoder-launch-path-optimization-opportunities.md new file mode 100644 index 000000000000..34dcf94dd7f0 --- /dev/null +++ b/docs/source/models/bart-decoder-launch-path-optimization-opportunities.md @@ -0,0 +1,679 @@ + + +# BART decoder launch-path optimization opportunities + +This note records the remaining host-side bottlenecks observed while profiling +the BART PyTorch encoder-decoder path, along with possible ways to reduce GPU +launch starvation. The proposals are diagnostic follow-up work, not verified +performance improvements, except where a completed experiment is marked +verified. + +## Profile context + +The analysis uses the Nsight Systems report: + +```text +/tmp/bart-encoder-piecewise-profile/piecewise-direct.nsys-rep +``` + +The run used continuous admission with a maximum scheduled decoder batch of 32 +requests. Generation step 208 contained 31 generation requests and no context +requests. It therefore already used the padded batch-32 decoder CUDA graph +bucket. Raising concurrency can amortize host overhead, but does not address +the remaining launch cost at a fixed concurrency. + +A blank region between NVTX ranges means that the CPU work is not covered by an +NVTX annotation. It does not by itself mean that the CPU thread is idle. In +this trace, the gaps before `_update_requests` and `_fetch_new_requests` mostly +overlap forward kernels from the current batch. The interval after +`prepare_resources`, however, substantially overlaps a genuinely idle GPU. + +## Representative launch-path breakdown + +For generation step 208, the host path from the end of `prepare_resources` to +the end of `[Executor] _forward_step` decomposed as follows: + +| Host region | Profiled duration | +| --- | ---: | +| Before `[Executor] _forward_step` | 126.9 us | +| `[Executor]` entry to `_prepare_inputs` | 76.2 us | +| `_prepare_inputs` | 707.4 us | +| `_prepare_inputs` end to `cudaGraphLaunch` | 102.9 us | +| `cudaGraphLaunch` API call | 306.9 us (The experiment confirms the ~300 µs cudaGraphLaunch duration is Nsight tracing overhead.) | +| Remaining host unwind | 36.0 us | + +These are durations under CUDA and NVTX tracing, not unprofiled latency +measurements. In particular, the approximately 280--310 us steady-state +`cudaGraphLaunch` duration is unusually high for a reused graph and may include +substantial profiler overhead. It must be confirmed with an unprofiled host +timer around graph replay, or a less intrusive graph-level trace, before +treating it as an implementation bottleneck. The experiment below provides +that confirmation. + +## Promising approaches + +### Reduce encoder-decoder input preparation + +Cache stable generation metadata for an unchanged batch, including request IDs, +sequence slots, prompt lengths, cross-KV lengths and pointers, and the decoder +graph key. Update only sampled tokens, positions, and KV-length deltas on each +step. + +Fine-grained debug-only NVTX ranges were added to the native encoder-decoder +fast path. Set `TLLM_NVTX_DEBUG=1` to enable them. The profiles are: + +```text +/tmp/bart-encoder-fast-input-profile/native-fast-full-c32-r256.nsys-rep +/tmp/bart-encoder-fast-input-profile/native-fast-fine-c32-r256.nsys-rep +``` + +The first report has coarse ranges with less annotation overhead. Across its +893 generation-only fast-path calls, the representative P50 breakdown was: + +| Fast-path region | P50 | +| --- | ---: | +| Complete encoder-decoder fast path | 513.8 us | +| Input-ID staging | 161.3 us | +| Attention-metadata preparation | 145.9 us | +| Cross-attention preparation | 37.9 us | +| Host-buffer retirement | 35.8 us | +| Position-ID staging | 30.5 us | +| Sequence-length staging | 15.8 us | +| Native request collation | 15.0 us | + +The second report subdivides input-ID and attention-metadata preparation. In +generation step 208, which contained nine generation requests, input-ID +staging took 198.1 us under tracing: + +| Input-ID subregion | Profiled duration | CUDA API work | +| --- | ---: | --- | +| Copy previous batch indices | 33.2 us | One asynchronous copy | +| Gather sampled tokens | 76.9 us | Two kernel launches | +| Copy gathered tokens | 28.2 us | One asynchronous copy | +| Fill graph-padding tokens | 23.3 us | One kernel launch | + +The complete input-ID region issued three kernel launches and two asynchronous +copies. + +#### Sampled-token staging consolidation (Verified: clean end to end +1.0%) + +The sampled-token experiment replaced advanced indexing followed by a copy +with `torch.index_select(..., out=input_ids_cuda_slice)`. It also retains the +device sequence-slot indices while the ordered generation request IDs remain +unchanged. Any non-encoder-decoder preparation path invalidates that cache. + +The optimized profiles are: + +```text +/tmp/bart-encoder-fast-input-profile/sampled-token-direct-c32-r256.nsys-rep +/tmp/bart-encoder-fast-input-profile/sampled-token-direct-cache-c32-r256.nsys-rep +``` + +| Fine-grained P50 | Before | Direct gather | Direct gather plus index reuse | +| --- | ---: | ---: | ---: | +| Complete encoder-decoder fast path | 612.0 us | 560.8 us | 519.1 us | +| Input-ID staging | 194.9 us | 142.8 us | 109.3 us | +| Sampled-token gather/copy | 103.5 us | 58.3 us | 58.9 us | +| Previous-index copy count | 907 | 907 | 202 | + +Index reuse eliminated 705 of 907 previous-index copies. For generation step +208, the sampled-token portion fell from 105.0 to 48.7 us and changed from two +kernel launches plus one asynchronous copy to one kernel launch. A standalone +unprofiled CUDA microbenchmark reduced this operation from 29.1 to 10.6 us at +batch 9 and from 29.6 to 14.3 us at batch 32. + +Because CUDA API tracing inflates the profiled ranges, a diagnostic experiment +temporarily placed `time.perf_counter_ns()` timers around the generation-only +`model_engine.forward()` call and input-ID staging. The timers did not +synchronize the GPU, did not emit per-step output, skipped the first 256 +generation steps, and collected 5,785 samples per run. Both NVTX environment +switches were disabled. The fine-grained `nvtx_range_debug` context managers +were still present as null context managers, however, so this established the +direction of the host-path change but was not a completely +instrumentation-free measurement. + +| Unprofiled host region | Before | Optimized | Reduction | +| --- | ---: | ---: | ---: | +| Input-ID staging, mean | 142.652 us | 94.670 us | 33.6% | +| Input-ID staging, P50 | 131.785 us | 81.866 us | 37.9% | +| Complete generation forward launch, mean | 657.765 us | 615.086 us | 6.5% | +| Complete generation forward launch, P50 | 600.186 us | 550.230 us | 8.3% | + +A fully clean end-to-end experiment then physically removed all added +fine-grained ranges from `model_engine.py` and `trtllm.py`, removed the host +timers, and ran without Nsight. It alternated optimized and pre-change code for +four concurrency-32, 2,048-request runs per version: + +| Order | Version | Mean latency | Makespan | Requests/s | Output tokens/s | +| ---: | --- | ---: | ---: | ---: | ---: | +| 1 | Optimized | 213.794 ms | 13.822200 s | 148.167 | 9953.842 | +| 2 | Before | 217.299 ms | 14.046938 s | 145.797 | 9794.590 | +| 3 | Optimized | 215.025 ms | 13.905088 s | 147.284 | 9894.507 | +| 4 | Before | 217.570 ms | 14.039977 s | 145.869 | 9800.087 | +| 5 | Optimized | 216.035 ms | 13.956041 s | 146.746 | 9855.947 | +| 6 | Before | 215.996 ms | 13.968241 s | 146.618 | 9849.773 | +| 7 | Optimized | 214.392 ms | 13.863271 s | 147.728 | 9924.353 | +| 8 | Before | 217.528 ms | 14.061549 s | 145.645 | 9784.413 | + +| Four-run average | Before | Optimized | Change | +| --- | ---: | ---: | ---: | +| Mean latency | 217.098 ms | 214.812 ms | -1.05% | +| P50 latency | 202.168 ms | 200.133 ms | -1.01% | +| P90 latency | 345.073 ms | 340.892 ms | -1.21% | +| Makespan | 14.029176 s | 13.886650 s | -1.02% | +| Requests/s | 145.982 | 147.481 | +1.03% | +| Output tokens/s | 9807.216 | 9907.162 | +1.02% | + +Two of the eight runs differed slightly in natural-EOS placement, once on +each version; output-token throughput gives the same approximately 1.0% result +after accounting for that small workload variation. The clean comparison +therefore confirms a modest end-to-end gain, while the diagnostic host timers +and Nsight profiles explain where it originates. + +Attention-metadata preparation took 188.7 us in the same step: + +| Attention-metadata subregion | Profiled duration | CUDA API work | +| --- | ---: | --- | +| Stage prompt and KV lengths | 40.7 us | Two asynchronous copies | +| Update host metadata | 23.6 us | Host-only tensor updates | +| Copy KV block offsets | 75.1 us | One asynchronous copy plus event query/record | +| Bind runtime views | 10.0 us | Host-only view binding | + +KV block-offset staging is the largest individual metadata subregion. Reusing +its pinned staging storage and avoiding the copy when the request-to-block +mapping is unchanged should be measured after token staging is consolidated. +This optimization must preserve the existing completion-event lifetime rules +for overlapped scheduling. + +Nested NVTX annotations and CUDA API tracing materially inflate all absolute +durations in the fine-grained report. The ranges establish relative +attribution; they are not unprofiled latency measurements. In particular, +native request collation is already small, while mixed admission makes +cross-attention preparation dominant only on the relatively infrequent +batch-change iterations. + +The representative `_prepare_inputs` range issued 28 CUDA API calls: + +- Six asynchronous copies. +- Three small preparation kernels. +- Three event queries. +- Four event records. + +The remaining calls were stream-state queries and kernel-name lookups recorded +by the profiler. Packing the small metadata transfers into one pinned buffer, +fusing the preparation kernels, or capturing both into the decoder graph would +reduce launch-critical work. + +#### Metadata, retirement, and position staging follow-up + +Three safe changes were prototyped together and separately: + +- Bound KV block-offset staging to the columns required by the batch's maximum + KV length. +- Re-record the completion event associated with an available host-buffer set + instead of constructing another event. +- For decoder CUDA-graph replay, copy position IDs directly from their pinned + host buffer into the graph's static position tensor and skip the otherwise + redundant device-to-device copy in `CUDAGraphRunner.replay()`. + +Reusing the completion event and bounding the block-offset copy were neutral +in two 2,048-request runs: + +| Two-run average | Before | Event + block bound | Change | +| --- | ---: | ---: | ---: | +| Mean latency | 214.292 ms | 214.212 ms | -0.04% | +| Makespan | 13.855197 s | 13.849266 s | -0.04% | +| Requests/s | 147.815 | 147.879 | +0.04% | + +The direct position-ID path by itself was also within run-to-run noise: + +| Two-run average | Before | Direct position staging | Change | +| --- | ---: | ---: | ---: | +| Mean latency | 208.198 ms | 207.623 ms | -0.28% | +| Makespan | 13.458528 s | 13.421465 s | -0.28% | +| Requests/s | 152.172 | 152.597 | +0.28% | + +One pair favored the position change by 0.80%, while the next favored the +baseline by 0.24%. A longer 8,192-request B--O--B run then compared all three +safe changes without Nsight or debug NVTX instrumentation: + +| Order | Version | Mean latency | Makespan | Requests/s | Output tokens/s | +| ---: | --- | ---: | ---: | ---: | ---: | +| 1 | Before | 225.061 ms | 57.744981 s | 141.865 | 10517.278 | +| 2 | Optimized | 227.479 ms | 58.368348 s | 140.350 | 10404.954 | +| 3 | Before | 231.405 ms | 59.377079 s | 137.966 | 10228.189 | +| Baseline average | Before | 228.233 ms | 58.561030 s | 139.916 | 10372.734 | + +The optimized run was 0.31% faster than the average of its surrounding +baselines, but those baselines differed by 2.75%. All three runs generated the +same 607,320 output tokens with the same output hash. The measured change is +therefore not distinguishable from environmental drift, and none of these +three code changes was retained. + +An additional attempt cached prompt lengths and a pinned KV block-offset +snapshot, skipping H2D staging when their host contents appeared unchanged. +That is not safe under the current overlapped metadata ownership contract: the +2,048-request output changed from 137,584 to 150,955 tokens and produced a +different hash. The prototype was removed. Future block-table reuse must use +an explicit cache-manager generation/ownership contract rather than infer +device-buffer validity from equal host contents. + +### Graph replay timing without tracing (Verified: 10 us launch, 47 us runner) + +An environment-injected timer measured graph replay during the +`timed_generate` range of the same concurrency-32, 256-request workload. The +benchmark ran as plain Python without Nsight Systems or CUDA API tracing. The +timer was injected into the MPI executor worker and used +`time.perf_counter_ns()` immediately around both: + +- `torch.cuda.CUDAGraph.replay()`, which is the direct counterpart of the + profiled `cudaGraphLaunch` API call. +- `CUDAGraphRunner.replay()`, which additionally includes graph lookup and + copies of input IDs and position IDs into the graph's static tensors. + +The native extension available to this worktree did not contain +`prepare_encoder_decoder_inputs`, so the runs forced the Python collation +fallback. This matches the earlier eager-versus-piecewise profile setup and +does not change the graph replay implementation being timed. + +| Run and scope | Calls | Mean | P50 | P99 | +| --- | ---: | ---: | ---: | ---: | +| Run 1, CUDA graph replay | 767 | 10.493 us | 10.293 us | 16.677 us | +| Run 2, CUDA graph replay | 767 | 10.331 us | 9.881 us | 20.569 us | +| Run 3, CUDA graph replay | 767 | 10.542 us | 10.209 us | 18.724 us | +| Run 3, complete decoder graph runner | 767 | 47.394 us | 46.718 us | 61.594 us | + +The median cost of an empty timer pair was 66--69 ns. All three runs produced +the same output hash and natural-EOS count. Their end-to-end mean latencies +were 217.267, 218.838, and 224.212 ms, respectively; these absolute values +describe the fallback diagnostic rather than the optimized native collation +path. + +The untraced replay is about 27--30 times shorter than the 280--310 us +steady-state `cudaGraphLaunch` calls in the Nsight trace. Graph launch itself +therefore is not the approximately 300 us bottleneck implied by the traced +timeline. Even the full graph runner is below 50 us at the median, and only +about 10 us of that is the native graph replay. Further work should prioritize +input preparation and the host path leading into the runner rather than graph +executable reuse or upload. + +### Reuse the input-copy completion event + +The model engine constructs and records a new completion event after every +encoder-decoder input preparation. Re-recording the event owned by an +available host-buffer set preserved correctness, but the experiment above +showed no measurable end-to-end benefit. + +### Remove redundant pre-forward Python work + +The overlap loop sorts generation requests every iteration even though the +ordering correction is only needed for disaggregated generation. The +aggregated BART path can skip that list allocation and key-function traversal. + +A qualified generation-only fast path can also skip: + +- Context-token summation. +- Context-logit checks. +- Encoder-output attachment. +- Cache-indirection lookup when it is not used. + +Using an NVTX context directly instead of constructing a decorated nested +function on every iteration removes another small fixed cost. + +### Reduce fixed stream-handoff overhead + +The stream waits around forward implicitly create, record, wait on, and destroy +CUDA events. Reusing preallocated handoff events, or keeping forward and its +dependent sampling work on one stream when KV transfer is inactive, can remove +this fixed overhead. Any change must preserve the ordering required by KV +onboard/offload and asynchronous sampling. + +### Prepare batch-static metadata earlier + +Split decoder input preparation into: + +- A batch-static portion: request ordering, block tables, cross-attention + metadata, and graph selection. +- A token-dependent portion: sampled tokens, positions, and KV-length + increments. + +The batch-static portion can run while the preceding GPU step is executing. +Once its sampled token is available, the launch-critical path should contain +only a compact device update and graph replay. + +### Continue decode while encoder futures are pending (Verified: neutral end to end) + +Encoder-init requests and existing generation requests are independent. The +executor submits encoder forward to one persistent host worker immediately +before decoder forward, keeps the encoder request IDs in the scheduler's +in-flight set, and retains the returned futures in FIFO order. At the start of +each later iteration, the executor polls the oldest future with `done()` and +queries its CUDA completion event. Neither operation waits. Only after both +report completion does the main executor thread publish the encoder output, +transition the request from `ENCODER_INIT` to `CONTEXT_INIT`, and remove its ID +from the in-flight set. This gives request-state mutation and error handling +back to the main executor thread while preventing both duplicate encoder +submission and premature decoder-context admission. + +This removes the prior same-iteration `future.result()` barrier. In +`/tmp/bart-encoder-pending-futures-c32-r256.nsys-rep`, the encoder range from +40.744076 to 40.768364 seconds overlaps host execution of generation-only +decoder steps 211 through 218. Other steady-state encoder ranges similarly +overlap four to nine generation-only decoder forwards. The executor therefore +continues fetch, schedule, input preparation, forward, and sampling work while +the encoder future is pending. + +GPU concurrency remains limited. In the 2.371-second timed window, encoder +stream 21 had 66.106 ms of kernel-busy time, decoder stream 17 had 41.064 ms, +and sampler stream 7 had 14.210 ms. Encoder kernels overlapped decoder kernels +for 0.595 ms and sampler kernels for 0.075 ms, about 1.0% of encoder busy time. +The host pipeline is now independent, but the large BART kernels still consume +most available device resources. + +The compatible temporary native mixed-batch binding used by the earlier +measurements was cleaned before this experiment. The following A/B runs +therefore forced the generic Python input path for both versions; they isolate +the pending-future scheduling change but are not directly comparable to the +native-fast-path throughput above. + +| Pair | Version | Mean latency | Makespan | Requests/s | +| ---: | --- | ---: | ---: | ---: | +| 1 | Blocking encoder | 240.779 ms | 15.563809 s | 131.587 | +| 1 | Pending futures | 243.423 ms | 15.714822 s | 130.323 | +| 2 | Blocking encoder | 242.596 ms | 15.690760 s | 130.523 | +| 2 | Pending futures | 240.687 ms | 15.532974 s | 131.849 | +| Average | Blocking encoder | 241.688 ms | 15.627285 s | 131.055 | +| Average | Pending futures | 242.055 ms | 15.623898 s | 131.086 | + +Average throughput changed by +0.02% and mean latency by +0.15%, both within +run-to-run noise. The blocking baseline produced 137,584 tokens and hash +`93761bbaed28ad0f` in both runs. Pending-future execution changed batch +composition and produced 137,594 and 137,545 tokens; the natural-EOS and +length-stop request counts remained identical. A 64-request token-level check +found only request 27 diverged, beginning at generated token 94, with the same +128-token output length. This is consistent with a close greedy decision +changing under different BF16 batch numerics rather than request/output +misassociation. + +The change removes a real software barrier but does not improve this workload's +end-to-end performance. A material gain still requires coarser encoder replay +(for example, a whole-encoder CUDA graph) or kernels that leave complementary +GPU resources available. + +#### Prioritize decoder kernels over encoder kernels (Rejected) + +An experiment created the encoder-decoder execution stream with CUDA priority +`-1` while retaining priority `0` for the encoder stream. Decoder-only models +kept priority `0`. CUDA stream priority favors pending work on the +higher-priority stream when the GPU scheduler can choose new work, but it +cannot preempt an encoder kernel that is already running. + +The clean benchmark alternated default and high decoder priority at concurrency +32 for 2,048 requests. Both versions used the generic encoder-decoder input +path because the loaded native extension predates +`prepare_encoder_decoder_inputs`. No Nsight or CUDA API tracing was active. + +| Order | Decoder priority | Mean latency | P50 latency | P90 latency | P99 latency | Makespan | Requests/s | Output tokens/s | +| ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| 1 | 0 | 242.986 ms | 227.216 ms | 374.525 ms | 461.013 ms | 15.685588 s | 130.566 | 8764.224 | +| 2 | -1 | 244.408 ms | 228.988 ms | 380.483 ms | 462.689 ms | 15.791230 s | 129.692 | 8711.038 | +| 3 | 0 | 244.749 ms | 229.705 ms | 381.593 ms | 456.817 ms | 15.825666 s | 129.410 | 8689.302 | +| 4 | -1 | 245.734 ms | 231.712 ms | 379.946 ms | 460.800 ms | 15.923400 s | 128.616 | 8644.699 | + +| Two-run average | Priority 0 | Priority -1 | Change | +| --- | ---: | ---: | ---: | +| Mean latency | 243.868 ms | 245.071 ms | +0.49% | +| P50 latency | 228.461 ms | 230.350 ms | +0.83% | +| P90 latency | 378.059 ms | 380.215 ms | +0.57% | +| P99 latency | 458.915 ms | 461.745 ms | +0.62% | +| Makespan | 15.755627 s | 15.857315 s | +0.65% | +| Requests/s | 129.988 | 129.154 | -0.64% | +| Output tokens/s | 8726.763 | 8677.869 | -0.56% | + +Output lengths varied by at most 0.13% between runs; token-normalized +throughput therefore gives the same conclusion as request throughput. The +benchmark records final-response latency rather than per-token inter-token +latency, so it does not exclude a small latency redistribution from existing +generation requests toward replacement encoder requests. It does show that +the proposed priority does not improve end-to-end request latency or +throughput for this workload. The stream-priority change was not retained. + +#### Re-evaluate encoder batch waiting (Current settings retained) + +The pending-future scheduler was adjusted at concurrency 32 because encoder +launch and completion no longer block the decoder host loop. A 256-request +screen covered iteration deadlines 24, 32, 40, 48, and 64 and token-threshold +ratios 0.08, 0.12, 0.1708984375, and 0.25. The existing setting is 48 +iterations and ratio 0.1708984375, which corresponds to 11,200 tokens or +approximately one full batch of 32 average-length encoder inputs. + +The short screen selected 64 iterations and ratio 0.08, but an alternating +2,048-request validation rejected it: it averaged 129.14 requests/s versus +130.71 requests/s for the existing configuration. The less aggressive +64-iteration, 0.12-ratio candidate was positive in two 2,048-request pairs, +improving mean latency and requests/s by 1.50% and output-token throughput by +1.61%. A longer confirmation reduced that result to noise and exposed worse +median and tail latency: + +| Order | Iterations / ratio | Mean latency | P50 latency | P90 latency | P99 latency | Makespan | Requests/s | Output tokens/s | +| ---: | --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| 1 | 48 / 0.1708984375 | 256.200 ms | 245.680 ms | 379.094 ms | 452.209 ms | 32.940426 s | 124.346 | 8883.522 | +| 2 | 64 / 0.12 | 254.405 ms | 247.532 ms | 373.355 ms | 463.342 ms | 32.761208 s | 125.026 | 8936.117 | +| 3 | 48 / 0.1708984375 | 254.777 ms | 243.711 ms | 373.563 ms | 452.560 ms | 32.788736 s | 124.921 | 8928.371 | +| 4 | 64 / 0.12 | 255.028 ms | 247.642 ms | 374.130 ms | 465.251 ms | 32.799612 s | 124.880 | 8925.166 | + +| Two-run average | 48 / 0.1708984375 | 64 / 0.12 | Change | +| --- | ---: | ---: | ---: | +| Mean latency | 255.489 ms | 254.717 ms | -0.30% | +| P50 latency | 244.696 ms | 247.587 ms | +1.18% | +| P90 latency | 376.329 ms | 373.743 ms | -0.69% | +| P99 latency | 452.385 ms | 464.297 ms | +2.63% | +| Makespan | 32.864581 s | 32.780410 s | -0.26% | +| Requests/s | 124.634 | 124.953 | +0.26% | +| Output tokens/s | 8905.947 | 8930.642 | +0.28% | + +Both versions used the generic encoder-decoder input path because the loaded +native extension predates `prepare_encoder_decoder_inputs`; neither used +Nsight or CUDA API tracing. The 0.26--0.28% throughput change is within the +observed run-to-run variation, while the P50 and P99 regressions are larger. +The existing 48-iteration, 0.1708984375-ratio configuration was therefore +retained. + +#### Drain the decoder before admitting another encoder wave (Rejected) + +A stricter scheduling experiment stopped admitting encoder requests whenever +the scheduler had any decoder-context or generation requests. New requests +therefore accumulated until the entire active decoder wave completed, at which +point the scheduler released the waiting encoder requests together. This +formed larger encoder batches, but also made every replacement request wait +for the longest output in the preceding decoder wave. + +The clean concurrency-32 comparison used 1,024 requests, the same dataset and +request order, no Nsight or CUDA API tracing, and the generic encoder-decoder +input path because the loaded native extension predates +`prepare_encoder_decoder_inputs`. + +| Policy | Mean latency | P50 latency | P90 latency | P99 latency | Makespan | Requests/s | Output tokens/s | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| Current bounded encoder accumulation | 225.910 ms | 199.376 ms | 387.849 ms | 472.617 ms | 7.383563 s | 138.686 | 8276.492 | +| Drain decoder completely | 306.773 ms | 285.259 ms | 425.774 ms | 526.375 ms | 10.012829 s | 102.269 | 6109.162 | +| Change | +35.8% | +43.1% | +9.8% | +11.4% | +35.6% | -26.3% | -26.2% | + +The larger encoder waves do not compensate for the head-of-line admission +barrier. In particular, short decoder requests cannot be replaced until the +longest request in the current wave finishes, leaving capacity unused during +the decoder tail. The strict policy was removed; bounded encoder accumulation +continues to admit replacement work while generation proceeds. + +#### Preallocate stable encoder output and prepare decoder context early + +The stronger follow-up publishes stable encoder-output views before submitting +the encoder worker. `ModelEngine` allocates one exactly sized output tensor for +the encoder batch, each request receives a view into that tensor, and the +worker copies the final hidden states into those stable addresses. A CUDA event +is also attached before submission. A separate recorded flag prevents the main +thread from querying or waiting on the event until the worker has enqueued it. +This removes the lifetime hazard that otherwise prevents decoder-context +preparation from starting before encoder completion. + +The scheduler can now capacity-plan those requests in `CONTEXT_INIT` while the +encoder worker is active. The executor removes requests whose encoder event is +not ready from the current decoder batch, retains its generation requests, and +prepares the deferred requests' persistent KV, cross-KV, and sequence-slot +resources in that generation-only iteration. The requests remain in the +executor's in-flight set so they cannot be scheduled twice. Once the recorded +event queries ready, the executor releases them to a later mixed decoder batch. +That mixed batch skips resource allocation already completed by the lookahead +pass and waits once on the shared encoder event before consuming the stable +output views. + +This is host/preparation overlap and launch-bubble removal. It does not attempt +to execute decoder-context kernels before the encoder event is satisfied, and +it retains generation-only decoder iterations while the encoder kernels run. +The optimization is enabled only when every active non-null resource manager is +one of the KV-cache, cross-KV-cache, or sequence-slot managers and a cross-KV +manager is present. Other configurations use the non-blocking encoder future +path without early resource preparation. + +In +`/tmp/bart-encoder-stable-lookahead-early-c32-r256.nsys-rep`, there are 916 +decoder forwards and 928 `prepare_resources` ranges. The 12 additional calls +occur immediately before generation-only forwards while encoder work is still +active. For example, decoder-context preparation runs from 66.385924 to +66.386937 seconds while the encoder runs from 66.381330 to 66.412223 seconds. +The following generation-only forward starts at 66.387195 seconds. When the +encoder event becomes ready, mixed forward 342 performs only a 98 us resource +preparation before starting at 66.414147 seconds. + +Across steady-state encoder admissions, excluding the two startup groups, the +median timings changed as follows: + +| Admission interval | Pending-future baseline | Stable-output lookahead | Change | +| --- | ---: | ---: | ---: | +| Encoder worker return to mixed-forward start | 2.665 ms | 2.405 ms | -9.8% | +| Last encoder GPU operation to first decoder GPU operation | 3.081 ms | 2.900 ms | -5.9% | +| Last encoder GPU operation to first decoder GEMM | 5.400 ms | 5.295 ms | -1.9% | + +One representative pair shows the intended effect more clearly, although the +mixed batch shapes are not identical. The pending-future trace's 13-context, +15-generation batch starts 2.622 ms after its encoder worker returns and its +first decoder input operation begins 3.027 ms after the last encoder GPU +operation. The lookahead trace's 15-context, 12-generation batch reduces those +intervals to 1.924 and 2.253 ms, respectively. + +Clean unprofiled measurements used concurrency 32, 2,048 requests, and forced +the generic encoder-decoder input path in both temporary overlays because the +available compatible native binding predates the current collation interface. +The repeated results remain noisy: + +| Version | Runs | Median mean latency | Median makespan | Median requests/s | +| --- | ---: | ---: | ---: | ---: | +| Pending-future baseline | 5 | 245.028 ms | 15.871215 s | 129.039 | +| Stable-output lookahead | 4 | 242.236 ms | 15.658510 s | 130.794 | + +The medians correspond to -1.14% mean latency, -1.34% makespan, and +1.36% +throughput. Individual optimized runs ranged from 127.942 to 136.162 requests/s, +so this is a small directional gain rather than a statistically decisive +end-to-end improvement. The trace provides the stronger causal result: the +resource work moves earlier and the steady-state admission bubble shrinks, but +most of the path to the first decoder GEMM remains elsewhere in input +preparation and model launch. + +#### Run encoder and decoder context together on the side stream + +A stronger experiment moves the dependent decoder context-only forward into +the encoder worker. The worker executes the encoder and then its decoder +context batch on stream 21, while the main executor continues launching +generation-only CUDA graphs on stream 17. + +The implementation uses a second `PyTorchModelEngine` that shares the BART +model but owns separate decoder input buffers, attention metadata, and graph +runner state. CUDA graphs are disabled for the context engine. KV, cross-KV, +and sequence-slot resources are reserved on the main executor thread because +their managers mutate shared allocation state; decoder input preparation and +model launch remain in the worker. Only one worker context batch can be active, +so its engine and stable encoder-output storage cannot be reused early. + +Two ordering details are required for correctness: + +- The main sampler has reusable device storage and already has one outstanding + generation sample under overlap scheduling. Completed context logits are + therefore sampled only after the prior main-lane sample has been retired. +- A worker context batch is not the main executor's previous batch. + `py_batch_idx` is cleared after its sampled token reaches the host request, + causing the first main-lane generation step to use the explicit host-token + admission path. + +The two engines have distinct destination buffers, but their asynchronous +metadata copies are staged from shared resource-manager state. Reusing that +state before an H2D copy completed corrupted subsequent generation tokens. +Each lane now waits only for its input-copy event after enqueueing the complete +forward. This does not wait for model kernels: host preparation can proceed on +both threads, and the already-enqueued model work remains concurrent on the two +CUDA streams. + +The implementation is restricted to single-rank BART and mBART with overlap +scheduling, KV-cache manager V1, no attention DP, drafting, guided decoding, +KV-cache transfer, connector, or early first-token response. Other +configurations retain the stable-output lookahead path. + +The profile is: + +```text +/tmp/bart-encoder-context-worker-c32-r256.nsys-rep +``` + +During its 2.651-second `timed_generate` interval, stream 21 executes 96.831 ms +of encoder plus context kernels. Stream 17 executes 973.694 ms of decoder CUDA +graphs. Their kernels overlap for 20.641 ms, or 21.3% of stream-21 kernel time. +Steady worker `_run_encoder_context_step` host ranges also overlap seven to ten +main generation `_forward_step` ranges each. The requested host and GPU +concurrency is therefore present, although most stream-21 work still cannot +co-reside with the large decoder kernels. + +Clean unprofiled measurements alternated the stable-output lookahead baseline +and this worker-context version at concurrency 32 with 2,048 requests: + +| Order | Version | Mean latency | P50 latency | Makespan | Requests/s | Output tokens/s | +| ---: | --- | ---: | ---: | ---: | ---: | ---: | +| 1 | Stable-output lookahead | 245.145 ms | 232.250 ms | 15.862079 s | 129.113 | 8671.499 | +| 2 | Worker encoder + context | 264.948 ms | 251.086 ms | 17.171612 s | 119.267 | 8091.727 | +| 3 | Stable-output lookahead | 240.438 ms | 225.812 ms | 15.510191 s | 132.042 | 8869.717 | +| 4 | Worker encoder + context | 263.153 ms | 251.669 ms | 17.003768 s | 120.444 | 8199.771 | + +| Two-run average | Stable-output lookahead | Worker encoder + context | Change | +| --- | ---: | ---: | ---: | +| Mean latency | 242.792 ms | 264.051 ms | +8.76% | +| P50 latency | 229.031 ms | 251.378 ms | +9.76% | +| Makespan | 15.686135 s | 17.087690 s | +8.94% | +| Requests/s | 130.578 | 119.856 | -8.21% | +| Output tokens/s | 8770.608 | 8145.749 | -7.12% | + +Separating context from the main mixed batch changes BF16 batch numerics and +therefore a small number of greedy decisions. The worker runs generated 1.18% +more output tokens on average; output-token throughput still regressed by +7.12%, so output length does not explain the result. The approximately +20.6 ms of hidden stream-21 work is smaller than the cost of separate eager +context launches, sampling/finalization, input-staging fences, and GPU resource +contention. This design proves that the work can overlap, but it is not an +end-to-end performance improvement for the measured workload. + +## Suggested experiment order + +1. Add fine-grained NVTX ranges inside the encoder-decoder input fast path. + (Completed.) +2. Gather sampled tokens directly into the persistent input-ID buffer and + avoid restaging stable previous-batch indices. (Completed; approximately + 1.0% throughput improvement across four fully clean interleaved runs.) +3. Reuse KV block-offset staging storage and skip unchanged block-table copies. + (Attempted; unsafe without an explicit device-buffer ownership contract.) +4. Reuse or ping-pong `_prepare_inputs_event` and measure the host-time change. + (Completed; no measurable end-to-end gain.) +5. Measure broader stable encoder-decoder metadata reuse. (Position staging and + bounded block-table staging completed; no verified end-to-end gain.) + +These experiments determine whether the next large gain is stable metadata +reuse or consolidation of device-side input updates. diff --git a/optimize_plan_2.md b/optimize_plan_2.md new file mode 100644 index 000000000000..027eb7b53e30 --- /dev/null +++ b/optimize_plan_2.md @@ -0,0 +1,494 @@ + + +# BART PyTorch plan to approach legacy TensorRT performance + +## Objective + +Close the remaining BART continuous-admission performance gap between the +optimized PyTorch path and the legacy TensorRT path without changing output +semantics or regressing latency tails. + +In the matched concurrency-32, 1,024-request benchmark, optimized PyTorch +reached 172.6 requests/s and 10,305 output tokens/s, while legacy TensorRT +reached 200.9 requests/s and 11,289 output tokens/s. PyTorch therefore reached +85.9% of TensorRT request throughput and 91.3% of its output-token throughput. + +Recent measurements around the pending-encoder-future experiments used the +generic encoder-decoder input fallback because the loaded native extension +predated `prepare_encoder_decoder_inputs`. Those absolute results must not +replace the native-fast-path baseline above. Rebuild or load a compatible +extension before evaluating progress against legacy TensorRT. + +## Profile-derived hypothesis + +The remaining gap is not primarily caused by a lack of concurrent encoder and +decoder GPU execution. + +The concurrency-32, 256-request profiles contain: + +| Backend | Encoder passes | Decoder passes | +| --- | ---: | ---: | +| Legacy TensorRT | 180 | 681 | +| PyTorch pending-future design | 17 | 932 | + +The backends can generate somewhat different output lengths because BF16 batch +composition affects close greedy decisions, so the pass counts are not a pure +scheduling comparison. Nevertheless, they show a major structural difference: +legacy TensorRT can run cheap replacement encoder microbatches and replenish +the decoder promptly, while PyTorch must coalesce expensive eager encoder +launches into large waves. The resulting PyTorch decoder batches spend more +iterations partially occupied. + +The legacy profile also showed no encoder/decoder, encoder/sampler, or +decoder/sampler kernel overlap. Its performance does not depend on concurrent +execution across those streams. Earlier closed-batch measurements found about +217.9 ms of PyTorch GPU work versus 226.4 ms of legacy TensorRT GPU work. +Consequently, the primary target is cheap and timely decoder replenishment plus +lower per-iteration orchestration, not maximum GPU kernel overlap. + +## Design 1: occupancy-aware encoder replenishment + +### Motivation + +The current encoder admission policy waits until either: + +- Waiting encoder tokens reach + `batch_wait_max_tokens_ratio * max_num_tokens`; or +- `batch_wait_timeout_iters` expires. + +At concurrency 32, the configured token threshold is approximately 32 +average-length encoder inputs. This gives efficient encoder execution, but it +can leave the decoder substantially underfilled while replacements wait. + +The rejected full-drain policy moved in the wrong direction: it waited until +all decoder work completed before releasing another encoder wave. That added a +head-of-line barrier, reduced request throughput by 26.3%, and increased mean +latency by 35.8%. + +### Proposed policy + +Release waiting encoder requests when any of these conditions is true: + +1. The waiting encoder request count reaches a modest microbatch target. +2. The active generation count falls below a decoder low watermark. +3. The oldest encoder request reaches a maximum iteration deadline. +4. There are no active decoder requests. + +Continue accumulating replacements while the decoder remains sufficiently +full. This preserves encoder efficiency during the steady state but refills +decoder capacity before reaching a long, underfilled tail. + +Initial screening matrix: + +| Parameter | Values | +| --- | --- | +| Encoder microbatch target | 4, 8, 12 | +| Decoder low watermark at concurrency 32 | 20, 24, 28 | +| Maximum wait | 16, 24, 48 iterations | + +Record, in addition to end-to-end performance: + +- Encoder-pass count and batch-size distribution. +- Decoder-pass count. +- Generation batch-size distribution and average occupancy. +- Total encoder, decoder, and sampler GPU busy time. +- Delay from request submission to encoder launch. +- Delay from encoder completion to mixed decoder-context launch. + +### Measured result: rejected before small-encoder optimization + +The policy was implemented with the existing pending-future encoder path and +screened at concurrency 32. The proposed default released eight waiting +encoder requests or released earlier when total decoder occupancy reached 24. +Two less aggressive settings were also tested. + +All runs used the same generic encoder-decoder input path because the installed +native extension predates `prepare_encoder_decoder_inputs`. No Nsight, debug +NVTX, or host timers were enabled. + +| Requests | Policy | Mean latency | P50 latency | P90 latency | P99 latency | Requests/s | +| ---: | --- | ---: | ---: | ---: | ---: | ---: | +| 256 | Existing bounded accumulation | 224.531 ms | 189.549 ms | 385.358 ms | 464.158 ms | 134.155 | +| 256 | Target 8 / low watermark 24 | 320.268 ms | 273.850 ms | 608.803 ms | 689.613 ms | 94.240 | +| 256 | Target 12 / low watermark 20 | 285.727 ms | 240.157 ms | 513.444 ms | 582.509 ms | 105.972 | +| 256 | Target 16 / low watermark 16 | 268.005 ms | 237.374 ms | 459.735 ms | 549.034 ms | 113.317 | + +The documented 1,024-request comparison confirmed the short screen: + +| Policy | Mean latency | P50 latency | P90 latency | P99 latency | Makespan | Requests/s | Output tokens/s | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| Existing bounded accumulation | 218.841 ms | 195.114 ms | 375.102 ms | 443.033 ms | 7.121912 s | 143.782 | 8568.205 | +| Target 8 / low watermark 24 | 333.685 ms | 279.505 ms | 639.218 ms | 723.822 ms | 10.822124 s | 94.621 | 5651.201 | +| Change | +52.5% | +43.3% | +70.4% | +63.4% | +52.0% | -34.2% | -34.0% | + +The optimized run produced 61,158 output tokens versus 61,022 for the +baseline, a 0.22% difference that does not explain the regression. Earlier +replenishment increases eager encoder launch frequency and GPU disruption more +than fuller decoder batches save. The implementation was removed. + +Occupancy-aware replenishment should be reconsidered only after design 2 makes +small encoder batches materially cheaper. Its experiment then becomes a useful +way to trade the new microbatch cost against decoder occupancy. + +### Cost-aware follow-up + +If a fixed low-watermark policy is directionally positive, replace the static +threshold with a measured cost decision. Maintain an exponentially weighted +estimate of encoder cost by microbatch size and decoder-step cost by graph +bucket. Dispatch a waiting encoder microbatch when its estimated cost is lower +than the decoder work expected to be saved by replenishing the open lanes. + +The initial policy should remain simple until the profile counters prove that +decoder occupancy predicts end-to-end performance. + +## Design 2: make small encoder batches cheap + +Occupancy-aware admission can succeed only if the additional encoder launches +are inexpensive enough. The desired behavior is not necessarily to reproduce +all 180 legacy encoder passes, but to permit more frequent batches of roughly +1--8 requests without returning to the previously measured batch-one eager +encoder collapse. + +### Whole-encoder CUDA graphs for microbatches + +Prioritize whole-encoder graph capture for batch sizes 1, 2, 4, and 8. Use +exact-shape graph keys initially: + +```text +(batch size, packed token count, maximum sequence length, exact cu_seqlens layout) +``` + +Exact keys avoid the dummy sequence padding and numerical changes observed in +the piecewise graph experiment. Capture graphs lazily and use an LRU limit to +bound static-buffer and graph-executable memory. + +The previous bounded whole-encoder graph prototype improved mean latency by +3.8%, 3.8%, and 6.4% at concurrencies 32, 64, and 128. Its relative benefit +may be larger for the small encoder batches most affected by eager Python and +CUDA launch overhead. + +### Supporting work + +- Pack encoder token and position inputs in native code. +- Reuse stable pinned and device input/output storage for each graph key. +- Preserve the fused residual/layer-normalization and GELU implementations. +- Keep variable-length attention numerically identical; do not pad real + sequences merely to reduce the number of graph keys. +- Benchmark graph lookup, input staging, replay, and output publication + separately from capture. + +Once small encoder execution is faster, rerun the occupancy-aware policy +matrix and progressively lower its microbatch target. + +### Measured result: graph replay helps, but small-batch execution still loses + +An experimental implementation added independent whole-encoder graphs without +replacing the decoder CUDA-graph configuration. It uses exact keys containing +the full sequence-length layout, lazy thread-local capture on the encoder +worker, a 64-entry LRU, and graph-resident pinned/device staging. Real +sequences and token counts are not padded. Runtime capture must be +thread-local: CUDA's default global capture mode otherwise rejects unrelated +sampler synchronization on the main executor thread. + +The first attribution used the same admission schedule on both sides. At +batch one, exact graphs reduced mean latency by 8.0% relative to eager +execution, but the schedule was still far slower than the existing +large-batch policy because batch-one encoder kernels lose too much GPU +efficiency. At batch eight, where this cyclic workload repeatedly reuses two +exact eight-request layouts, graph replay improved the 256-request screen only +slightly: + +| Requests | Encoder schedule | Encoder execution | Mean latency | Requests/s | +| ---: | --- | --- | ---: | ---: | +| 256 | Target 1 / no low watermark | Eager | 509.427 ms | 61.176 | +| 256 | Target 1 / no low watermark | Exact graph | 468.626 ms | 66.505 | +| 256 | Target 8 / no low watermark | Eager | 243.296 ms | 119.891 | +| 256 | Target 8 / no low watermark | Exact graph | 241.235 ms | 122.983 | + +The matched concurrency-32, 1,024-request comparison was negative. Enabling +graphs alone while retaining the existing bounded-accumulation policy did not +help because steady-state encoder batches were normally larger than eight; +the remaining small tail batches paid capture cost without enough replay. +Forcing the best screened target-eight schedule increased encoder frequency +enough to outweigh its slightly cheaper launches. + +| Policy | Mean latency | P50 latency | P90 latency | P99 latency | Makespan | Requests/s | Output tokens/s | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| Existing bounded accumulation | 215.505 ms | 189.828 ms | 368.375 ms | 454.794 ms | 7.046579 s | 145.319 | 8679.106 | +| Exact graphs, existing admission | 219.949 ms | 192.038 ms | 382.436 ms | 459.412 ms | 7.210709 s | 142.011 | 8494.866 | +| Exact graphs, target 8 / no low watermark | 240.486 ms | 199.376 ms | 456.963 ms | 544.848 ms | 7.884365 s | 129.877 | 7758.392 | +| Target-8 change versus baseline | +11.6% | +5.0% | +24.0% | +19.8% | +11.9% | -10.6% | -10.6% | + +The three runs produced 61,158, 61,254, and 61,170 output tokens, +respectively, so output-count variation does not explain the latency result. +As in the design-1 screen, these measurements used the generic +encoder-decoder input path because the loaded native extension predates +`prepare_encoder_decoder_inputs`. + +Exact whole-encoder graphs therefore reduce launch overhead for a fixed small +batch, but not enough to make frequent PyTorch encoder replenishment +competitive. Do not enable this policy by default. A future attempt needs a +material reduction in small-batch kernel time or a graph strategy that +preserves larger encoder batches while removing their host launch overhead. + +### Measured result: `(batch, total tokens, max bucket)` keys + +The cyclic benchmark produces 59 unique `(B,T,S)` triples for +`B in {1,2,4,8}` when maximum sequence length is bucketed in 64-token +increments. The implementation accepts those triples as an allowlist, captures +each graph on its first real batch, and retains up to 64 graphs in an LRU. +Encoder graph metadata uses a separate buffer arena from the concurrently +executing decoder graphs, and shared host staging is retired before the next +replay updates it. + +Individual sequence lengths are runtime graph inputs, not part of the key. +TRTLLM attention rebuilds `cu_seqlens` and padding offsets on the GPU from the +current device sequence-length buffer. `B`, `T`, and the maximum-length bucket +keep tensor extents, workspace requirements, and FMHA launch dimensions stable. + +The earlier cross-layout hang came from inconsistent capture warmup metadata, +not a requirement for exact layouts. Graph metadata initialized its device +sequence lengths to ones, preparation updated only the stable host buffer, and +the H2D copy occurred for the first time inside graph capture. Warmup therefore +ran with real packed-token counts on the host but all-one sequence lengths on +the GPU. Capture now stages input IDs, position IDs, and sequence lengths to +their device buffers before warmup. The graph still captures the H2D copies for +subsequent replays. Exact-layout rejection is removed, so all layouts sharing a +`(B,T,S)` key reuse the same graph. + +Matched concurrency-32, 1,024-request results: + +| Backend / encoder mode | Mean latency | P50 latency | P90 latency | P99 latency | Makespan | Requests/s | Output tokens/s | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| PyTorch eager, target 8 | 245.536 ms | 204.663 ms | 467.781 ms | 521.870 ms | 8.058071 s | 127.078 | 7573.525 | +| PyTorch `(B,T,S)` graph with exact-layout guard | 234.166 ms | 194.920 ms | 453.589 ms | 511.825 ms | 7.685696 s | 133.235 | 7947.751 | +| PyTorch `(B,T,S)` graph with cross-layout replay | 233.765 ms | 192.882 ms | 447.072 ms | 508.841 ms | 7.682285 s | 133.294 | 7945.943 | +| PyTorch exact batch-8 admission at decoder occupancy 24 | 225.955 ms | 187.235 ms | 433.918 ms | 492.370 ms | 7.373542 s | 138.875 | 8274.721 | +| PyTorch exact batch 8, serialized encoder/decoder forward | 225.123 ms | 183.728 ms | 435.884 ms | 485.110 ms | 7.342011 s | 139.471 | 8317.203 | +| PyTorch exact batch 8, overlapped encoder plus mixed decoder graphs | 159.100 ms | 130.335 ms | 306.821 ms | 341.215 ms | 5.229246 s | 195.822 | 11673.959 | +| Legacy TensorRT | 163.359 ms | 127.874 ms | 344.281 ms | 425.151 ms | 5.363881 s | 190.907 | 10783.982 | + +Cross-layout replay did not materially change end-to-end performance versus +the exact-layout guard: mean latency fell by 0.2%, request throughput rose by +0.04%, and makespan fell by 0.04%, all within run-to-run variation. It removes +an invalid hidden key and increases graph coverage, but graph eligibility and +scheduling still dominate the aggregate result. The four PyTorch runs +generated 61,028, 61,084, 61,043, and 61,014 output tokens, respectively; +legacy generated 57,844, so request throughput and latency are the cleaner +cross-backend comparisons. + +The corrected Nsight profile contains 121 encoder batches in the timed +generation window. Eighty-eight batches executed a CUDA graph: 11 captured a +new key and immediately launched it, while 77 were cache-hit replays. +Thirty-three batches used eager fallback. + +The eager fallbacks were caused by scheduler overfill rather than missing +graphs for a supported batch size. Once at least eight encoder requests were +available, admission returned the entire scheduler result, producing batches +of 9 through 12 in steady state plus startup and tail batches of 31 and 5. +Exact batch-8 admission now returns only the first eight requests once decoder +occupancy reaches 24 or lower and leaves excess requests eligible for the next +iteration. A deadline releases a smaller supported power-of-two tail. + +The clean exact-batch-8 run reduced mean latency by 3.3% and increased request +throughput by 4.2% versus cross-layout replay with uncapped admission. Its +profile contained exactly 128 batch-8 encoder steps for 1,024 requests. All +128 were cache-hit graph replays; there were no timed captures and no eager +fallbacks. Encoder dispatch occurred alongside 24 generation requests in 58 +steps. Other steady-state dispatches occurred after the decoder completed +multiple requests between scheduling passes and crossed directly below 24. + +Serializing encoder and decoder forward launches produced no measurable +end-to-end change. The executor launched the encoder on its dedicated stream, +blocked on its completion event, and only then entered decoder forward. +Compared with asynchronous exact-batch-8 admission, mean latency and makespan +were 0.4% lower, request throughput was 0.4% higher, P90 was 0.5% higher, and +the other percentiles moved by less than 2%. This mixed movement is within +single-run variation and indicates that encoder/decoder GPU concurrency was +not contributing material throughput in this workload. + +The final design restores asynchronous encoder/decoder overlap and captures +whole-model CUDA graphs for mixed decoder batches. Its graph key contains the +padded decoder batch size, exact decoder context-query extents, and packed +encoder-hidden-state row count. Cross-attention sequence layouts remain +runtime metadata rather than graph keys. Encoder hidden states are copied into +one graph-stable buffer before replay, so a replacement encoder output never +leaves a captured pointer referring to request-owned storage. + +Mixed graphs are never captured on a live request. Startup first warms all +reachable batch-8 and paired-batch-16 replenishment shapes to finish sizing +the shared attention workspace, then captures the same 61 mixed shapes on a +second pass. An unseen shape falls back to eager execution. This avoids both +repeating live KV-cache writes during capture and invalidating older graph +pointers through a late workspace resize. + +Against the preceding best overlapped exact-batch-8 result, mixed decoder +graphs reduced mean latency by 29.6%, P50 by 30.4%, P90 by 29.3%, P99 by +30.7%, and makespan by 29.1%. Request throughput and output-token throughput +increased by 41.0% and 41.1%, respectively. It also slightly exceeded legacy +TensorRT in this run: mean latency was 2.6% lower, makespan was 2.5% lower, +and request throughput was 2.6% higher. Output-token throughput is not a clean +cross-backend comparison because legacy produced fewer output tokens. + +The final Nsight timed window contained 126 mixed decoder steps, and all 126 +launched a decoder CUDA graph. All 2,205 generation-only decoder steps and all +128 encoder steps also replayed graphs. The single context-only decoder step +remained eager, and there were no CUDA graph captures in the timed window. +Typical mixed `_forward_step` ranges fell from roughly 18--20 ms in the eager +profile to roughly 1.9--2.1 ms with replay. + +The final PyTorch run produced 61,046 output tokens versus 61,014 in the +preceding best run, a 0.05% difference. BF16 continuous admission can change +close greedy decisions when faster completion changes batch composition, so +the output hash is not expected to remain fixed across scheduling changes. + +Raw logs: + +- `/tmp/bart-bts-graphs-eager-target8-pytorch-c32-r1024.log` +- `/tmp/bart-bts-graphs-safe-allow59-pytorch-c32-r1024.log` +- `/tmp/bart-bts-graphs-legacy-c32-r1024.log` +- `/tmp/bart-bts-graphs-cross-layout-fixed-allow59-c32-r1024.nsys-rep` +- `/tmp/bart-bts-graphs-cross-layout-fixed-allow59-c32-r1024.sqlite` +- `/tmp/bart-bts-graphs-exact8-low24-pytorch-c32-r1024.log` +- `/tmp/bart-bts-graphs-exact8-low24-c32-r1024.nsys-rep` +- `/tmp/bart-bts-graphs-exact8-low24-c32-r1024.sqlite` +- `/tmp/bart-bts-graphs-exact8-low24-serialized-pytorch-c32-r1024.log` +- `/tmp/bart-bts-graphs-exact8-low24-serialized-c32-r1024.nsys-rep` +- `/tmp/bart-mixed-decoder-graphs-final-overlap-c32-r1024.log` +- `/tmp/bart-mixed-decoder-graphs-final-overlap-c32-r1024.nsys-rep` +- `/tmp/bart-mixed-decoder-graphs-final-overlap-c32-r1024.sqlite` + +## Design 3: a decoder supergraph + +The current decoder CUDA graph captures model execution, but input staging, +attention-metadata updates, sampling, request updates, stream handoffs, and +completion processing remain separate. Legacy TensorRT hides most of this work +behind one engine enqueue. + +For the qualified single-rank, single-beam, greedy BART path, maintain a +persistent device-side lane table containing: + +- Request and sequence-slot identifiers. +- Current token and position. +- Self-KV block descriptors and lengths. +- Cross-KV descriptors and encoder-output pointers. +- Active, EOS, and length-complete state. + +Capture one decoder supergraph containing: + +```text +sampled-token gather +→ position and KV-length update +→ decoder model +→ greedy argmax +→ next-token scatter +→ EOS and length-completion update +``` + +The next decoder iteration should consume the sampled token directly from +device storage. The CPU should receive a compact completion record and update +only lane admissions, removals, and externally visible request state. Batch +membership changes should be expressed as deltas to persistent lane metadata +rather than a complete rebuild of every active request. + +Use double-buffered launch packets so the batch-static portion of the next +iteration can be prepared while the current graph runs. Once sampled tokens +are available, the launch-critical path should consist of a compact device +update and one graph replay. + +Qualification must initially exclude streaming, beam search, speculative +decoding, guided decoding, LoRA, cache reuse, attention data parallelism, +disaggregated transfer, and other features that require the general path. + +## Design 4: two-token decoder graph unrolling + +After the decoder supergraph is correct and beneficial, capture two dependent +greedy decoder steps in one graph replay: + +```text +decoder step N +→ greedy token N +→ device state update +→ decoder step N+1 +→ greedy token N+1 +``` + +A device finish mask must prevent a token after EOS or the length limit from +being appended. A lane that completes after the first step may still perform +dummy model computation during the second step. + +Start with two steps only. Larger unrolling would save more host launches but +would also delay new-request admission and waste more computation on completed +lanes. This experiment also requires sufficient KV capacity to be reserved +before replay. + +## GPU work ordering + +Do not treat encoder/decoder GPU overlap as a goal by itself. The legacy +profile serialized its encoder, decoder, and sampler kernels, and the PyTorch +pending-future design achieved only about 1% encoder/decoder kernel overlap. + +For occupancy-aware microbatches, compare: + +1. Encoder execution ordered after the current decoder step and before a later + decoder step. +2. The existing independent encoder stream. + +Prefer deterministic serialized ordering if concurrent streams introduce +resource contention. The encoder host worker can still prepare and submit work +asynchronously even if CUDA events serialize its GPU execution. + +Keep decoder context requests in the main mixed decoder batch. Moving encoder +plus decoder-context execution to a side worker produced 20.6 ms of real GPU +overlap but regressed request throughput by 8.2% because separate eager context +launches, sampling, staging fences, and GPU contention cost more than the +hidden work saved. + +## Experiment order + +1. Rebuild or load the compatible native extension and reproduce the optimized + PyTorch and legacy TensorRT baselines with the same workload. +2. Add lightweight counters for encoder-pass sizes, decoder occupancy, and + admission delays. +3. Screen the occupancy-aware encoder policy without changing encoder kernels. +4. Implement exact-shape whole-encoder CUDA graphs for small microbatches. +5. Rerun the occupancy policy matrix and select the best cost/occupancy point. +6. Prototype the qualified decoder supergraph. +7. If a meaningful host bubble remains, test two-token graph unrolling. +8. Confirm every promising result with alternating, fully unprofiled long + runs, followed by Nsight profiling for causal attribution. + +## Validation criteria + +Use the checked-in continuous-admission workload and report both request and +executed-token throughput because natural EOS can differ between batch +compositions. + +A change should be retained only if: + +- It improves both request throughput and output-token throughput outside + observed run-to-run noise. +- Mean, P50, P90, and P99 latency do not show an unacceptable redistribution. +- Natural-EOS and length-stop behavior remains valid. +- Repeated runs show no request/output association errors. +- General configurations fall back without semantic changes. +- The native fast path, rather than the generic compatibility fallback, is + active in the comparison. + +## Designs not to revisit without new evidence + +- Draining every decoder request before admitting another encoder wave. +- Moving decoder-context execution onto the encoder worker. +- Decoder-versus-encoder CUDA stream-priority changes. +- Additional tuning of only the existing fixed token and iteration thresholds. +- Maximizing encoder/decoder kernel concurrency as an independent objective. +- Retaining finished decoder rows as permanent graph-padding lanes. + +The most useful immediate experiment is occupancy-aware replenishment. It +directly tests whether the legacy path's frequent replacement encoding and +lower decoder-pass count explain the remaining gap. If it improves decoder +occupancy but loses the gain to encoder cost, that result becomes a precise +performance requirement for the small-batch whole-encoder graph work. diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index b8b67ed5ded8..d73ba56d4ade 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -1,6 +1,10 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + import bisect import contextlib -from dataclasses import dataclass +from collections import OrderedDict +from dataclasses import dataclass, field from typing import (Any, Callable, Dict, Iterator, List, Optional, Tuple, TypeAlias) @@ -16,7 +20,7 @@ from ..attention_backend.trtllm import TrtllmAttentionMetadata from ..distributed import Distributed from ..expert_statistic import ExpertStatistic -from ..memory_buffer_utils import get_memory_buffers +from ..memory_buffer_utils import Buffers, get_memory_buffers from ..modules.multi_stream_utils import with_multi_stream from ..speculative.eagle3 import Eagle3ResourceManager from ..speculative.interface import SpecMetadata @@ -36,7 +40,8 @@ # as a one-token context chunk to write its cross-KV cache, so enc-dec # dummies need one prompt token plus one generated token. ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM = 2 -KeyType: TypeAlias = Tuple[int, int, bool, bool, bool] +KeyType: TypeAlias = Tuple[int, int, bool, bool, bool, Tuple[int, ...], + Tuple[int, ...]] def _save_spec_decode_capture_state( @@ -111,6 +116,10 @@ class CUDAGraphRunnerConfig: kv_cache_manager_key: Any dynamic_draft_len_mapping: Optional[Dict[int, int]] = None sparse_attention_config: Optional[BaseSparseAttentionConfig] = None + enable_encoder_decoder_mixed_cuda_graph: bool = False + encoder_hidden_size: int = 0 + dtype: Optional[torch.dtype] = None + encoder_decoder_mixed_cuda_graph_encoder_token_counts: Tuple[int, ...] = () class CUDAGraphRunner: @@ -135,6 +144,8 @@ def __init__(self, config: CUDAGraphRunnerConfig): self.spec_config = config.spec_config self.sparse_config = config.sparse_attention_config self.is_encoder_decoder = config.is_encoder_decoder + self.enable_encoder_decoder_mixed_cuda_graph = ( + config.enable_encoder_decoder_mixed_cuda_graph) self.graphs: Dict[KeyType, torch.cuda.CUDAGraph] = {} self.graph_outputs: Dict[KeyType, @@ -163,6 +174,10 @@ def _create_shared_static_tensors(self): token_per_request = runtime_draft_token_buffer_width + 1 max_total_tokens = (self.max_supported_batch_size * self.max_beam_width * token_per_request) + if self.enable_encoder_decoder_mixed_cuda_graph: + # A mixed encoder-decoder batch can contain multiple decoder + # context tokens per request, unlike a pure generation batch. + max_total_tokens = self.config.max_num_tokens max_total_tokens = min(max_total_tokens, self.config.max_num_tokens) self.shared_static_tensors = { @@ -178,6 +193,24 @@ def _create_shared_static_tensors(self): self.shared_static_tensors[ "mrope_delta_read_seq_slots"] = torch.zeros( (max_total_tokens, ), device="cuda", dtype=torch.long) + if self.enable_encoder_decoder_mixed_cuda_graph: + if self.config.encoder_hidden_size <= 0 or self.config.dtype is None: + raise ValueError("Mixed encoder-decoder CUDA graphs require " + "encoder_hidden_size and dtype.") + self.shared_static_tensors["encoder_hidden_states"] = torch.empty( + (self.config.max_num_tokens, self.config.encoder_hidden_size), + device="cuda", + dtype=self.config.dtype, + ) + + def _is_mixed_encoder_decoder_batch(self, batch: ScheduledRequests) -> bool: + return (self.enable_encoder_decoder_mixed_cuda_graph + and batch.num_context_requests > 0 + and batch.num_generation_requests > 0) + + def _can_run_cuda_graph_batch(self, batch: ScheduledRequests) -> bool: + return batch.can_run_cuda_graph or self._is_mixed_encoder_decoder_batch( + batch) def _get_seq_len_mode( self, @@ -251,7 +284,7 @@ def get_graph_key( # Because we will pad the input to 'max_draft_len' length for the first draft layer. draft_len = self.config.original_max_draft_len if spec_resource_manager.is_first_draft else 0 key = (batch_size, draft_len, spec_resource_manager.is_first_draft, - short_seq_len_mode, is_all_greedy_sample) + short_seq_len_mode, is_all_greedy_sample, (), ()) else: # With dynamic spec decode, the draft length may be zero even when enable_spec_decode is True, # so we need to get the draft length from the batch instead of using enable_spec_decode. @@ -261,8 +294,15 @@ def get_graph_key( draft_len = max(draft_len_list) assert len( set(draft_len_list)) == 1, "All draft lengths must be the same" + context_query_lens = tuple( + int(request.context_chunk_size) + for request in batch.context_requests) + encoder_input_lens = (sum( + int(request.encoder_output_len) + for request in batch.context_requests + if not request.py_skip_cross_kv_projection), ) key = (batch_size, draft_len, False, short_seq_len_mode, - is_all_greedy_sample) + is_all_greedy_sample, context_query_lens, encoder_input_lens) return key @staticmethod @@ -303,6 +343,7 @@ def maybe_get_cuda_graph( draft_tokens_cuda: Optional[torch.Tensor] = None, new_tensors_device: Optional[SampleStateTensors] = None, spec_resource_manager: Optional[BaseResourceManager] = None, + allow_mixed_encoder_decoder: bool = False, ) -> Tuple[Optional[Any], Optional[Any], Optional[Tuple[int, int, bool]]]: """ Determines if the current batch can be run with a CUDA graph. @@ -316,7 +357,10 @@ def maybe_get_cuda_graph( if ExpertStatistic.should_record(): return None, None, None - can_run_cuda_graph = batch.can_run_cuda_graph + is_mixed_encoder_decoder = self._is_mixed_encoder_decoder_batch(batch) + can_run_cuda_graph = (batch.can_run_cuda_graph + or (is_mixed_encoder_decoder + and allow_mixed_encoder_decoder)) batch_size = batch.batch_size if self.enabled and self.config.enable_attention_dp and self.config.mapping.tp_size > 1: all_can_graph_batch = self.config.dist.tp_allgather( @@ -347,8 +391,10 @@ def maybe_get_cuda_graph( return self.graph_metadata[key][ "attn_metadata"], self.graph_metadata[key]["spec_metadata"], key - # Graph doesn't exist yet. If on-the-fly capture is not allowed, - # fall back to eager so the caller doesn't need a separate check. + # Capturing a mixed graph on a live batch would execute its KV-cache + # writes during graph warmup/capture and could resize shared attention + # workspace after older graph pointers have been fixed. Only shapes + # captured by the two-pass startup warmup may replay. if not self._capture_allowed: return None, None, None @@ -358,6 +404,15 @@ def maybe_get_cuda_graph( num_sequences_in_batch = batch_size * self.max_beam_width graph_attn_metadata = attn_metadata.create_cuda_graph_metadata( num_sequences_in_batch, False, key[1], self.cuda_graph_meta_buffers) + if is_mixed_encoder_decoder: + context_query_lens = key[5] + generation_query_len = key[1] + 1 + graph_attn_metadata.seq_lens = torch.tensor( + context_query_lens + (generation_query_len, ) * + (num_sequences_in_batch - len(context_query_lens)), + dtype=torch.int, + ) + graph_attn_metadata.num_contexts = len(context_query_lens) assert graph_attn_metadata.is_cuda_graph if enable_spec_decode: @@ -395,6 +450,15 @@ def get_graph_pool(self): """ return self.memory_pool + def _get_num_tokens_for_key(self, key: KeyType) -> int: + batch_size = key[0] + token_per_generation = key[1] + 1 + context_query_lens = key[5] + num_contexts = len(context_query_lens) + return (sum(context_query_lens) + + (batch_size * self.max_beam_width - num_contexts) * + token_per_generation) + def capture(self, key: KeyType, forward_fn: Callable, @@ -406,10 +470,7 @@ def capture(self, # [CUDA graph spec decode padding] # We pad input IDs/position IDs to the maximum draft length (token per request). # We're forced to do this because we cannot reallocate inputs over many graph runs. - max_draft_len = key[1] - token_per_request = max_draft_len + 1 - num_tokens_for_capture = (batch_size * self.max_beam_width * - token_per_request) + num_tokens_for_capture = self._get_num_tokens_for_key(key) sliced_static_tensors = { "input_ids": @@ -429,6 +490,18 @@ def capture(self, capture_inputs = initial_inputs.copy() capture_inputs.update(sliced_static_tensors) + encoder_input_lens = key[6] + num_encoder_tokens = sum(encoder_input_lens) + if num_encoder_tokens: + encoder_hidden_states = initial_inputs.get("encoder_hidden_states") + if encoder_hidden_states is None: + raise RuntimeError("Mixed encoder-decoder CUDA graph capture " + "requires encoder hidden states.") + static_encoder_hidden_states = self.shared_static_tensors[ + "encoder_hidden_states"][:num_encoder_tokens] + static_encoder_hidden_states.copy_(encoder_hidden_states) + capture_inputs[ + "encoder_hidden_states"] = static_encoder_hidden_states attn_metadata = capture_inputs["attn_metadata"] saved_kv_lens_cuda = _save_spec_decode_capture_state( attn_metadata, enable_spec_decode) @@ -509,6 +582,15 @@ def replay(self, key: KeyType, else: static_tensors["position_ids"][:, :seqlen].copy_(position_ids) + num_encoder_tokens = sum(key[6]) + if num_encoder_tokens: + encoder_hidden_states = current_inputs.get("encoder_hidden_states") + if encoder_hidden_states is None: + raise RuntimeError("Mixed encoder-decoder CUDA graph replay " + "requires encoder hidden states.") + static_tensors["encoder_hidden_states"][:num_encoder_tokens].copy_( + encoder_hidden_states) + self.graphs[key].replay() output_ref = self.graph_outputs[key] @@ -519,7 +601,7 @@ def _get_padded_batch(self, batch: ScheduledRequests, runtime_draft_len: int) -> int: kv_cache_manager = resource_manager.get_resource_manager( self.config.kv_cache_manager_key) - can_run_cuda_graph = batch.can_run_cuda_graph + can_run_cuda_graph = self._can_run_cuda_graph_batch(batch) batch_size = batch.batch_size new_batch_size = batch_size @@ -728,6 +810,10 @@ class EncoderCUDAGraphRunnerConfig: max_num_tokens: int max_seq_len: int cuda_graph_mem_pool: Any + dynamic_sequence_layout: bool = False + allow_runtime_capture: bool = False + max_cuda_graphs: int = 0 + capture_keys: List[EncoderKeyType] = field(default_factory=list) class EncoderCUDAGraphRunner: @@ -735,9 +821,11 @@ class EncoderCUDAGraphRunner: Designed for encoder inputs with `input_ids` (flat [total_tokens]) and `seq_lens` ([batch_size]). Encoder CUDA graphs are keyed on the 3-tuple - (padded_batch_size, padded_num_tokens, padded_max_seq_len). + (batch_size, total_tokens, max_seq_len_bucket) for dynamic encoder-decoder + batches. - Restricted to `TrtllmAttentionMetadata` — FlashInfer's per-batch planner state is not compatible with CUDA graph capture/replay. + Restricted to `TrtllmAttentionMetadata`: FlashInfer's per-batch planner + state is not compatible with CUDA graph capture/replay. """ WARMUP_STEPS = 1 @@ -752,8 +840,18 @@ def __init__(self, config: EncoderCUDAGraphRunnerConfig): self.supported_num_tokens = sorted(config.cuda_graph_num_tokens) self.max_supported_num_tokens = config.max_cuda_graph_num_tokens self.supported_seq_lens = sorted(config.cuda_graph_seq_lens) - - self.graphs: Dict[EncoderKeyType, torch.cuda.CUDAGraph] = {} + self.dynamic_sequence_layout = config.dynamic_sequence_layout + self.allow_runtime_capture = config.allow_runtime_capture + self.max_cuda_graphs = config.max_cuda_graphs + self.capture_keys = frozenset(config.capture_keys) + if (self.max_cuda_graphs > 0 + and len(self.capture_keys) > self.max_cuda_graphs): + raise ValueError("Encoder CUDA graph capture key count exceeds " + f"max_cuda_graphs: {len(self.capture_keys)} > " + f"{self.max_cuda_graphs}.") + + self.graphs: OrderedDict[EncoderKeyType, + torch.cuda.CUDAGraph] = OrderedDict() self.graph_outputs: Dict[EncoderKeyType, Callable[[], Optional[Any]]] = {} self.graph_metadata: Dict[EncoderKeyType, Dict[str, Any]] = {} @@ -763,10 +861,12 @@ def __init__(self, config: EncoderCUDAGraphRunnerConfig): self.shared_static_tensors_cpu: Dict[str, torch.Tensor] = {} if self.enabled: self._create_shared_static_tensors() - self.cuda_graph_meta_buffers = get_memory_buffers() + self.cuda_graph_meta_buffers = ( + Buffers() if self.dynamic_sequence_layout else get_memory_buffers()) self._capture_allowed = False self.is_warmup_only = False + self._staging_retirement_event: Optional[torch.cuda.Event] = None # CUDA graph H2D memcpy nodes require pinned host sources. In CC mode # prefer_pinned() is false: pageable host buffers are preferred, so the @@ -775,8 +875,9 @@ def __init__(self, config: EncoderCUDAGraphRunnerConfig): def _create_shared_static_tensors(self): """Allocates static tensors sized for the largest supported num_tokens.""" - max_total_tokens = min(self.max_supported_num_tokens, - self.config.max_num_tokens) + max_total_tokens = ( + self.config.max_num_tokens if self.dynamic_sequence_layout else min( + self.max_supported_num_tokens, self.config.max_num_tokens)) max_batch_size = self.max_supported_batch_size self.shared_static_tensors = { @@ -856,8 +957,17 @@ def get_graph_key( batch_size = len(seq_lens) max_seq_len = max(seq_lens) if batch_size > 0 else 0 + if self.dynamic_sequence_layout: + max_seq_len_bucket = self._round_up(max_seq_len, + self.supported_seq_lens) + key: EncoderKeyType = (batch_size, num_tokens, max_seq_len_bucket) + is_valid = (num_tokens <= self.max_supported_num_tokens + and max_seq_len_bucket > 0) + return key, False, is_valid + key = self._get_valid_graph_key(batch_size, num_tokens, max_seq_len) - _, padded_num_tokens, padded_max_seq_len = key + padded_num_tokens = key[1] + padded_max_seq_len = key[2] is_padding_performed = (padded_num_tokens != num_tokens or padded_max_seq_len != max_seq_len) @@ -870,9 +980,9 @@ def get_graph_key( def allow_capture(self): """Context manager that enables CUDA graph capture. - Capture is disabled by default. On-the-fly captures outside this - context are prevented — unseen keys fall back to eager instead of - incurring a multi-millisecond capture latency spike at runtime. + Static encode-only graphs capture during warmup through this context. + Dynamic encoder-decoder graphs may additionally opt into first-use + runtime capture through ``allow_runtime_capture``. """ self._capture_allowed = True try: @@ -948,17 +1058,23 @@ def maybe_get_cuda_graph( key, is_padding_performed, is_padding_successful = self.get_graph_key( inputs) - _, _, padded_max_seq_len = key + if (self.dynamic_sequence_layout and self.capture_keys + and key not in self.capture_keys): + return None, None + padded_max_seq_len = key[2] if (not self.padding_enabled and is_padding_performed) \ or not is_padding_successful: return None, None if key in self.graph_metadata: + # Every graph key aliases the same host staging buffers. Retire a + # prior graph's captured reads before the caller updates them. + self.retire_staging() return self.graph_metadata[key]["attn_metadata"], key - # New key not yet captured. Only create metadata if capture is - # allowed (warmup time); otherwise fall back to eager. - if not self._capture_allowed: + # New key not yet captured. Create graph metadata only during an + # explicit warmup capture or when first-use runtime capture is enabled. + if not (self._capture_allowed or self.allow_runtime_capture): return None, None if "multi_item_part_lens" in inputs: @@ -998,6 +1114,7 @@ def maybe_get_cuda_graph( graph_attn_metadata.max_seq_len = self.config.max_seq_len graph_attn_metadata.request_ids = list(range(padded_batch_size)) + self.retire_staging() return graph_attn_metadata, key def _contains_nested_tensor(self, x: Any) -> bool: @@ -1010,7 +1127,62 @@ def _contains_nested_tensor(self, x: Any) -> bool: return False def needs_capture(self, key: EncoderKeyType) -> bool: - return self._capture_allowed and key not in self.graphs + return (self._capture_allowed + or self.allow_runtime_capture) and key not in self.graphs + + def _evict_graph_if_needed(self) -> None: + if self.max_cuda_graphs <= 0 or len(self.graphs) < self.max_cuda_graphs: + return + + key, graph = self.graphs.popitem(last=False) + graph.reset() + self.graph_outputs.pop(key, None) + self.graph_metadata.pop(key, None) + + def _stage_inputs(self, key: EncoderKeyType, inputs: Dict[str, + Any]) -> None: + """Stage input and position IDs for capture or replay.""" + padded_num_tokens = key[1] + + # Captured H2D nodes read pinned host buffers. In CC mode, where H2D + # is not captured, stage directly into the graph-resident CUDA buffers. + static_tensors = self.shared_static_tensors_cpu if self._capture_h2d_copy else self.shared_static_tensors + + input_ids = inputs["input_ids"] + if isinstance(input_ids, list): + actual_tokens = len(input_ids) + static_tensors["input_ids"][:actual_tokens].copy_( + torch.tensor(input_ids, dtype=torch.int32)) + elif isinstance(input_ids, torch.Tensor): + actual_tokens = int(input_ids.shape[0]) + static_tensors["input_ids"][:actual_tokens].copy_(input_ids) + else: + raise TypeError(f"Unsupported input_ids type: {type(input_ids)}") + static_tensors["input_ids"][actual_tokens:padded_num_tokens].fill_(0) + + # Auto-generate packed position IDs without allocating one concatenated + # tensor, or copy caller-provided values into the stable staging buffer. + staged_position_ids = static_tensors["position_ids"][0] + position_ids = inputs.get("position_ids") + if position_ids is None: + offset = 0 + for seq_len in inputs["seq_lens"]: + staged_position_ids[offset:offset + seq_len].copy_( + self._arange_max[:seq_len]) + offset += seq_len + else: + if isinstance(position_ids, list): + staged_position_ids[:actual_tokens].copy_( + torch.tensor(position_ids, dtype=torch.int32)) + elif isinstance(position_ids, torch.Tensor): + staged_position_ids[:actual_tokens].copy_( + position_ids.flatten()) + else: + raise TypeError( + f"Unsupported position_ids type: {type(position_ids)}") + offset = actual_tokens + + staged_position_ids[offset:padded_num_tokens].fill_(0) def capture( self, @@ -1019,7 +1191,8 @@ def capture( inputs: Dict[str, Any], ) -> Any: """Warm up and/or capture the forward pass for a graph key.""" - _, padded_num_tokens, _ = key + padded_num_tokens = key[1] + self._evict_graph_if_needed() sliced_static_tensors = { "input_ids": @@ -1042,6 +1215,20 @@ def capture( self.graph_metadata[key] = {"attn_metadata": attn_md} + # Warmup must see the same runtime data as capture. In particular, + # graph metadata initializes _seq_lens_cuda to ones, while + # prepare_encoder_cuda_graph_replay updates its stable host buffer. + # Populate every device input before warmup so packed-token counts and + # sequence boundaries are consistent. + self._stage_inputs(key, inputs) + if self._capture_h2d_copy: + capture_inputs["input_ids"].copy_( + sliced_static_tensors_cpu["input_ids"], non_blocking=True) + capture_inputs["position_ids"].copy_( + sliced_static_tensors_cpu["position_ids"], non_blocking=True) + attn_md._seq_lens_cuda.copy_(attn_md._seq_lens, non_blocking=True) + torch.cuda.current_stream().synchronize() + output = None with with_multi_stream(True), piecewise_cuda_graph(False): # Warmup runs required by CUDA graph semantics. See @@ -1055,7 +1242,9 @@ def capture( return output graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph, pool=self.memory_pool): + with torch.cuda.graph(graph, + pool=self.memory_pool, + capture_error_mode="thread_local"): if self._capture_h2d_copy: # H2D copies for captured inside the graph: at replay # time it re-issues from the pinned static buffer without @@ -1080,65 +1269,33 @@ def capture( self.memory_pool = graph.pool() return graph_output + def retire_staging(self) -> None: + """Wait until a prior replay no longer reads shared staging buffers.""" + if self._staging_retirement_event is not None: + self._staging_retirement_event.synchronize() + self._staging_retirement_event = None + def replay( self, key: EncoderKeyType, inputs: Dict[str, Any], ) -> Any: """Replay a captured graph with current inputs.""" + self.retire_staging() + stored_meta = self.graph_metadata[key] assert inputs["attn_metadata"] is stored_meta["attn_metadata"] - _, padded_num_tokens, _ = key - - # According to prefer_pinned(), CC forces most transfers to be synchronous. - # So we don't put non_blocking=True here. - static_tensors = self.shared_static_tensors_cpu if self._capture_h2d_copy else self.shared_static_tensors - - # input_ids: convert (if list) and write into pinned active region in - # one allocation + one memcpy. Padding region is zero-filled below. - input_ids = inputs["input_ids"] - if isinstance(input_ids, list): - actual_tokens = len(input_ids) - static_tensors["input_ids"][:actual_tokens].copy_( - torch.tensor(input_ids, dtype=torch.int32)) - elif isinstance(input_ids, torch.Tensor): - actual_tokens = int(input_ids.shape[0]) - static_tensors["input_ids"][:actual_tokens].copy_(input_ids) - else: - raise TypeError(f"Unsupported input_ids type: {type(input_ids)}") - static_tensors["input_ids"][actual_tokens:padded_num_tokens].fill_(0) - - # position_ids: pinned buffer is shape [1, max_total_tokens]; use the - # 1-D row view. Auto-generate via the cached arange (zero allocations, - # N small memcpys) or copy user-provided values. - pinned_pos = static_tensors["position_ids"][0] - position_ids = inputs.get("position_ids") - if position_ids is None: - # Pad entries (seq_len=1) get arange[:1] = [0], the correct - # position for a 1-token dummy request. - offset = 0 - for s in inputs["seq_lens"]: - pinned_pos[offset:offset + s].copy_(self._arange_max[:s]) - offset += s - else: - if isinstance(position_ids, list): - pinned_pos[:actual_tokens].copy_( - torch.tensor(position_ids, dtype=torch.int32)) - elif isinstance(position_ids, torch.Tensor): - pinned_pos[:actual_tokens].copy_(position_ids.flatten()) - else: - raise TypeError( - f"Unsupported position_ids type: {type(position_ids)}") - offset = actual_tokens - - pinned_pos[offset:padded_num_tokens].fill_(0) + self._stage_inputs(key, inputs) if not self._capture_h2d_copy: stored_meta["attn_metadata"]._seq_lens_cuda.copy_( stored_meta["attn_metadata"]._seq_lens, non_blocking=True) self.graphs[key].replay() + self.graphs.move_to_end(key) + self._staging_retirement_event = torch.cuda.Event() + self._staging_retirement_event.record(torch.cuda.current_stream()) return self.graph_outputs[key] diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index 88513bc326d6..ab128692086c 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + from copy import copy, deepcopy from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union, cast diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 1bdafaab9440..6e61647d0c1a 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -11,7 +11,7 @@ import weakref from abc import ABC, abstractmethod from contextlib import contextmanager -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union import torch import torch._dynamo.config @@ -257,6 +257,36 @@ def _filter_cuda_graph_seq_lens(cuda_graph_seq_lens: list[int], return result +def _parse_encoder_decoder_cuda_graph_keys( + raw_keys: str) -> List[Tuple[int, int, int]]: + """Parse semicolon-separated ``batch,total_tokens,max_seq_bucket`` keys.""" + if not raw_keys.strip(): + return [] + + keys = [] + for raw_key in raw_keys.split(";"): + fields = raw_key.split(",") + if len(fields) != 3: + raise ValueError( + "TLLM_ENCODER_DECODER_MICROBATCH_CUDA_GRAPH_KEYS entries " + "must use batch_size,total_tokens,max_seq_bucket.") + try: + batch_size, total_tokens, max_seq_bucket = (int(field) + for field in fields) + except ValueError as error: + raise ValueError( + "TLLM_ENCODER_DECODER_MICROBATCH_CUDA_GRAPH_KEYS entries " + "must contain integers.") from error + key = (batch_size, total_tokens, max_seq_bucket) + if any(value <= 0 for value in key): + raise ValueError( + "TLLM_ENCODER_DECODER_MICROBATCH_CUDA_GRAPH_KEYS entries " + "must contain positive integers.") + keys.append(key) + + return sorted(set(keys)) + + _DEEP_GEMM_PDL_CONFIGURED = False @@ -672,12 +702,50 @@ def __init__( self._max_cuda_graph_seq_len = (self._cuda_graph_seq_lens[-1] if self._cuda_graph_seq_lens else 0) + encoder_decoder_graph_max_batch_size = int( + os.environ.get( + "TLLM_ENCODER_DECODER_MICROBATCH_CUDA_GRAPH_MAX_BATCH_SIZE", + "0")) + encoder_decoder_microbatch_cuda_graph_enabled = (os.environ.get( + "TLLM_ENCODER_DECODER_MICROBATCH_CUDA_GRAPH_ENABLED", "1") == "1") + self._enable_encoder_decoder_microbatch_cuda_graph = ( + self._is_encoder_decoder_model() + and encoder_decoder_graph_max_batch_size > 0 + and encoder_decoder_microbatch_cuda_graph_enabled) + encoder_decoder_graph_batch_sizes = [ + batch_size for batch_size in (1, 2, 4, 8) if batch_size <= min( + encoder_decoder_graph_max_batch_size, self.batch_size) + ] + encoder_decoder_graph_keys = _parse_encoder_decoder_cuda_graph_keys( + os.environ.get("TLLM_ENCODER_DECODER_MICROBATCH_CUDA_GRAPH_KEYS", + "")) + encoder_decoder_graph_seq_lens = list( + range(64, self.max_seq_len + 1, 64)) + if (encoder_decoder_graph_seq_lens + and encoder_decoder_graph_seq_lens[-1] < self.max_seq_len): + encoder_decoder_graph_seq_lens.append(self.max_seq_len) + encoder_decoder_graph_keys = [ + key for key in encoder_decoder_graph_keys + if key[0] in encoder_decoder_graph_batch_sizes and key[1] <= + self.max_num_tokens and key[2] in encoder_decoder_graph_seq_lens + ] + mixed_graph_encoder_batch_size = (encoder_decoder_graph_batch_sizes[-1] + if encoder_decoder_graph_batch_sizes + else 0) + mixed_graph_encoder_token_counts = tuple( + sorted({ + total_tokens + for batch_size, total_tokens, _ in encoder_decoder_graph_keys + if batch_size == mixed_graph_encoder_batch_size + })) + self._dynamic_draft_len_mapping = self._compute_dynamic_draft_len_mapping( ) self.previous_batch_indices_cuda = torch.empty((self.max_num_tokens, ), dtype=torch.int, device='cuda') + self._encoder_decoder_staged_request_ids: Optional[List[int]] = None self.input_ids_cuda = torch.empty((self.max_num_tokens, ), dtype=torch.int, device='cuda') @@ -746,6 +814,14 @@ def __init__( self.lora_model_config: Optional[LoraModelConfig] = None self._trtllm_gen_jit_warmup = False + use_encoder_decoder_graph = ( + self._enable_encoder_decoder_microbatch_cuda_graph + and bool(encoder_decoder_graph_batch_sizes)) + enable_encoder_decoder_mixed_cuda_graph = ( + use_encoder_decoder_graph + and self.cuda_graph_config is not None and os.environ.get( + "TLLM_ENCODER_DECODER_MIXED_CUDA_GRAPH_ENABLED", "1") == "1") + # Create config and runner cuda_graph_runner_config = CUDAGraphRunnerConfig( use_cuda_graph=(not self._is_encode_only @@ -770,24 +846,53 @@ def __init__( dist=self.dist, kv_cache_manager_key=self.kv_cache_manager_key, sparse_attention_config=self.sparse_attention_config, + enable_encoder_decoder_mixed_cuda_graph=( + enable_encoder_decoder_mixed_cuda_graph), + encoder_hidden_size=(self._get_enc_dec_hidden_size() + if enable_encoder_decoder_mixed_cuda_graph else + 0), + dtype=(self.dtype + if enable_encoder_decoder_mixed_cuda_graph else None), + encoder_decoder_mixed_cuda_graph_encoder_token_counts=( + mixed_graph_encoder_token_counts), ) self.cuda_graph_runner = CUDAGraphRunner(cuda_graph_runner_config) # Create Encoder CUDA graph config and runner. + encoder_graph_batch_sizes = (encoder_decoder_graph_batch_sizes + if use_encoder_decoder_graph else + self._cuda_graph_batch_sizes) + encoder_graph_max_batch_size = (encoder_graph_batch_sizes[-1] + if encoder_graph_batch_sizes else 0) + encoder_graph_max_num_tokens = (self.max_num_tokens + if use_encoder_decoder_graph else + self._max_cuda_graph_num_tokens) encoder_cuda_graph_runner_config = EncoderCUDAGraphRunnerConfig( - use_cuda_graph=(self._is_encode_only - and self.cuda_graph_config is not None - and bool(self._cuda_graph_num_tokens) - and bool(self._cuda_graph_seq_lens)), - cuda_graph_padding_enabled=self._cuda_graph_padding_enabled, - cuda_graph_batch_sizes=self._cuda_graph_batch_sizes, + use_cuda_graph=(use_encoder_decoder_graph + or (self._is_encode_only + and self.cuda_graph_config is not None + and bool(self._cuda_graph_num_tokens) + and bool(self._cuda_graph_seq_lens))), + cuda_graph_padding_enabled=(False if use_encoder_decoder_graph else + self._cuda_graph_padding_enabled), + cuda_graph_batch_sizes=encoder_graph_batch_sizes, cuda_graph_num_tokens=self._cuda_graph_num_tokens, - cuda_graph_seq_lens=self._cuda_graph_seq_lens, - max_cuda_graph_batch_size=self._max_cuda_graph_batch_size, - max_cuda_graph_num_tokens=self._max_cuda_graph_num_tokens, + cuda_graph_seq_lens=(encoder_decoder_graph_seq_lens + if use_encoder_decoder_graph else + self._cuda_graph_seq_lens), + max_cuda_graph_batch_size=encoder_graph_max_batch_size, + max_cuda_graph_num_tokens=encoder_graph_max_num_tokens, max_num_tokens=self.max_num_tokens, max_seq_len=self.max_seq_len, cuda_graph_mem_pool=self._cuda_graph_mem_pool, + dynamic_sequence_layout=use_encoder_decoder_graph, + allow_runtime_capture=use_encoder_decoder_graph, + max_cuda_graphs=(int( + os.environ.get( + "TLLM_ENCODER_DECODER_MICROBATCH_CUDA_GRAPH_MAX_GRAPHS", + "64")) if use_encoder_decoder_graph else 0), + capture_keys=(encoder_decoder_graph_keys + if use_encoder_decoder_graph else []), ) self.encoder_cuda_graph_runner = EncoderCUDAGraphRunner( encoder_cuda_graph_runner_config) @@ -1763,6 +1868,7 @@ def _run_cuda_graph_warmup(self, resource_manager: ResourceManager): return self._capture_generation_cuda_graphs(resource_manager) + self._capture_mixed_encoder_decoder_cuda_graphs(resource_manager) # Piecewise graphs have separate capture machinery and do not use the # whole-model attention workspace. Capture them only on the second pass. if not self.cuda_graph_runner.is_warmup_only: @@ -1955,6 +2061,91 @@ def _run_capture_pass(force_non_greedy: bool, label: str) -> None: if self.spec_metadata is not None: self.spec_metadata.is_all_greedy_sample = True + def _capture_mixed_encoder_decoder_cuda_graphs( + self, resource_manager: ResourceManager) -> None: + """Warm and capture reachable mixed encoder-decoder graph shapes. + + The first global CUDA-graph pass warms every shape so shared attention + workspace reaches its final size. The second pass captures the same + shapes. Runtime capture is deliberately disabled because graph capture + executes KV-cache writes and must never run against live requests. + """ + runner = self.cuda_graph_runner + if not runner.enable_encoder_decoder_mixed_cuda_graph: + return + + max_encoder_output_len = self._get_max_encoder_output_len( + resource_manager) + encoder_token_counts = ( + runner.config.encoder_decoder_mixed_cuda_graph_encoder_token_counts + or (8 * max_encoder_output_len, )) + context_shapes = [(8, token_count) + for token_count in encoder_token_counts] + if runner.max_supported_batch_size > 16: + paired_token_counts = sorted({ + first + second + for first in encoder_token_counts + for second in encoder_token_counts + }) + context_shapes.extend( + (16, token_count) for token_count in paired_token_counts) + + operation = ("warmup" if runner.is_warmup_only else "capture") + hidden_size = self._get_enc_dec_hidden_size() + for num_contexts, total_encoder_tokens in context_shapes: + if total_encoder_tokens > num_contexts * max_encoder_output_len: + continue + base_encoder_len, remainder = divmod(total_encoder_tokens, + num_contexts) + encoder_output_lens = ([base_encoder_len + 1] * remainder + + [base_encoder_len] * + (num_contexts - remainder)) + if not encoder_output_lens or encoder_output_lens[-1] <= 0: + continue + + for batch_size in runner.supported_batch_sizes: + if batch_size <= num_contexts: + continue + warmup_request = self._create_cuda_graph_warmup_request( + resource_manager, + batch_size, + draft_len=0, + mixed_context_encoder_output_lens=encoder_output_lens) + with self._release_batch_context(warmup_request, + resource_manager) as batch: + if batch is None: + logger.warning( + "Skipping mixed encoder-decoder CUDA graph " + f"{operation}: not enough KV cache space for " + f"batch size={batch_size}.") + continue + + context_requests = batch.context_requests + for request, encoder_output_len in zip( + context_requests, encoder_output_lens): + request.state = LlmRequestState.CONTEXT_INIT + request.context_current_position = 0 + request.context_chunk_size = 2 + request.cached_tokens = 0 + request.py_batch_idx = None + request.py_encoder_output = torch.ones( + (encoder_output_len, hidden_size), + device="cuda", + dtype=self.dtype, + ) + request.py_skip_cross_kv_projection = False + + logger.info("Run mixed encoder-decoder CUDA graph " + f"{operation} for batch size={batch_size}, " + f"context requests={num_contexts}, " + f"packed encoder tokens={total_encoder_tokens}") + self.enable_spec_decode = False + self.runtime_draft_len = 0 + self.forward(batch, + new_tensors_device=None, + resource_manager=resource_manager) + torch.cuda.synchronize() + def _capture_piecewise_cuda_graphs(self, resource_manager: ResourceManager): """Captures piecewise CUDA graphs for context/prefill steps via torch.compile.""" if not (self._torch_compile_piecewise_cuda_graph @@ -2176,11 +2367,13 @@ def _create_warmup_request( return result def _create_cuda_graph_warmup_request( - self, - resource_manager: ResourceManager, - batch_size: int, - draft_len: int, - max_seq_len: int = None) -> Optional[ScheduledRequests]: + self, + resource_manager: ResourceManager, + batch_size: int, + draft_len: int, + max_seq_len: int = None, + mixed_context_encoder_output_lens: Optional[Sequence[int]] = None + ) -> Optional[ScheduledRequests]: """Creates a dummy ScheduledRequests tailored for CUDA graph capture.""" kv_cache_manager = resource_manager.get_resource_manager( self.kv_cache_manager_key) @@ -2203,26 +2396,74 @@ def _create_cuda_graph_warmup_request( max_encoder_output_len = ( self._get_max_encoder_output_len(resource_manager) if is_enc_dec else None) + num_mixed_contexts = len(mixed_context_encoder_output_lens or + ()) if is_enc_dec else 0 + if num_mixed_contexts >= batch_size: + return None - # Add (batch_size - 1) dummy requests with the minimal seq_len. - token_nums = ([ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM] * - (batch_size - 1)) if is_enc_dec else None - encoder_output_lens = ([max_encoder_output_len] * - (batch_size - 1)) if is_enc_dec else None - requests = kv_cache_manager.add_dummy_requests( - list(range(batch_size - 1)), - token_nums=token_nums, - is_gen=True, - max_num_draft_tokens=runtime_draft_token_buffer_width, - kv_reserve_draft_tokens=self.max_draft_loop_tokens, - use_mrope=self.use_mrope, - max_beam_width=self.max_beam_width, - encoder_output_lens=encoder_output_lens, - num_extra_decoding_steps=num_extra_decoding_steps, - draft_kv_cache_manager=draft_kv_cache_manager) + # Add (batch_size - 1) dummy requests with the minimal sequence + # length. Mixed capture must create its context rows as real context + # requests; converting generation dummies afterward leaves their + # native prompt/context bookkeeping at one token. + if mixed_context_encoder_output_lens: + context_request_ids = list(range(num_mixed_contexts)) + context_requests = kv_cache_manager.add_dummy_requests( + context_request_ids, + token_nums=[ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM] * + num_mixed_contexts, + is_gen=False, + max_num_draft_tokens=runtime_draft_token_buffer_width, + kv_reserve_draft_tokens=self.max_draft_loop_tokens, + use_mrope=self.use_mrope, + max_beam_width=self.max_beam_width, + encoder_output_lens=list(mixed_context_encoder_output_lens), + num_extra_decoding_steps=num_extra_decoding_steps, + draft_kv_cache_manager=draft_kv_cache_manager) + if context_requests is None: + return None - if requests is None: - return None + generation_request_ids = list( + range(num_mixed_contexts, batch_size - 1)) + generation_requests = [] + if generation_request_ids: + generation_requests = kv_cache_manager.add_dummy_requests( + generation_request_ids, + token_nums=[ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM] * + len(generation_request_ids), + is_gen=True, + max_num_draft_tokens=runtime_draft_token_buffer_width, + kv_reserve_draft_tokens=self.max_draft_loop_tokens, + use_mrope=self.use_mrope, + max_beam_width=self.max_beam_width, + encoder_output_lens=[max_encoder_output_len] * + len(generation_request_ids), + num_extra_decoding_steps=num_extra_decoding_steps, + draft_kv_cache_manager=draft_kv_cache_manager) + if generation_requests is None: + for request in context_requests: + kv_cache_manager.free_resources(request) + if draft_kv_cache_manager is not None: + draft_kv_cache_manager.free_resources(request) + return None + requests = context_requests + generation_requests + else: + token_nums = ([ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM] * + (batch_size - 1)) if is_enc_dec else None + encoder_output_lens = ([max_encoder_output_len] * + (batch_size - 1)) if is_enc_dec else None + requests = kv_cache_manager.add_dummy_requests( + list(range(batch_size - 1)), + token_nums=token_nums, + is_gen=True, + max_num_draft_tokens=runtime_draft_token_buffer_width, + kv_reserve_draft_tokens=self.max_draft_loop_tokens, + use_mrope=self.use_mrope, + max_beam_width=self.max_beam_width, + encoder_output_lens=encoder_output_lens, + num_extra_decoding_steps=num_extra_decoding_steps, + draft_kv_cache_manager=draft_kv_cache_manager) + if requests is None: + return None def free_warmup_requests() -> None: for r in requests: @@ -2288,14 +2529,26 @@ def free_warmup_requests() -> None: else: max_seq_len_request = max_seq_len_request[0] - # Insert the longest request first to simulate padding for the CUDA graph. - requests.insert(0, max_seq_len_request) - result.generation_requests = requests + if mixed_context_encoder_output_lens: + requests.append(max_seq_len_request) + for request in requests[:num_mixed_contexts]: + request.state = LlmRequestState.CONTEXT_INIT + request.context_current_position = 0 + request.context_chunk_size = 2 + request.cached_tokens = 0 + request.py_batch_idx = None + result.context_requests_last_chunk = requests[:num_mixed_contexts] + result.generation_requests = requests[num_mixed_contexts:] + else: + # Insert the longest request first to simulate padding for the CUDA + # graph. + requests.insert(0, max_seq_len_request) + result.generation_requests = requests if spec_resource_manager is not None: spec_resource_manager.add_dummy_requests( request_ids=list(range(batch_size))) if self._is_encoder_decoder_model(): - if not self._add_cross_dummy_requests(result.generation_requests, + if not self._add_cross_dummy_requests(result.all_requests(), resource_manager): return None return result @@ -3187,8 +3440,11 @@ def _can_use_encoder_decoder_input_fast_path( self, '_encoder_decoder_input_fast_path_static_eligible', None) if static_eligible is None: static_eligible = ( - self._is_encoder_decoder_model() and not self.enable_spec_decode - and not self.is_draft_model and self.max_beam_width == 1 + hasattr(batch_manager_bindings, + "prepare_encoder_decoder_inputs") + and self._is_encoder_decoder_model() + and not self.enable_spec_decode and not self.is_draft_model + and self.max_beam_width == 1 and self.sparse_attention_config is None and not self.use_mrope and not self.enable_attention_dp and not self.mapping.has_cp_helix() and not self.is_multimodal @@ -3294,21 +3550,33 @@ def _prepare_encoder_decoder_inputs_fast( ) num_sequences = scheduled_requests.batch_size + num_context_requests = scheduled_requests.num_context_requests num_generation_requests = scheduled_requests.num_generation_requests + generation_request_ids = request_ids[num_context_requests:] if num_context_tokens: self.input_ids_cuda[:num_context_tokens].copy_( buffers['input_ids'][:num_context_tokens], non_blocking=True) if num_previous_batch_requests: previous_slots = self.previous_batch_indices_cuda[: num_previous_batch_requests] - previous_slots.copy_( - buffers['previous_batch_indices'][:num_previous_batch_requests], - non_blocking=True) - new_tokens = new_tokens_device[0, previous_slots, 0] + staged_request_ids = generation_request_ids[: + num_previous_batch_requests] + # Sequence slots are stable for a request's lifetime, so the + # device indices remain valid while this ordered batch does. + if self._encoder_decoder_staged_request_ids != staged_request_ids: + previous_slots.copy_(buffers['previous_batch_indices'] + [:num_previous_batch_requests], + non_blocking=True) + self._encoder_decoder_staged_request_ids = staged_request_ids generation_begin = num_context_tokens generation_end = generation_begin + num_previous_batch_requests - self.input_ids_cuda[generation_begin:generation_end].copy_( - new_tokens, non_blocking=True) + torch.index_select( + new_tokens_device[0, :, 0], + dim=0, + index=previous_slots, + out=self.input_ids_cuda[generation_begin:generation_end]) + else: + self._encoder_decoder_staged_request_ids = None dummy_begin = num_context_tokens + num_previous_batch_requests if dummy_begin < total_num_tokens: self.input_ids_cuda[dummy_begin:total_num_tokens].fill_(0) @@ -3321,12 +3589,14 @@ def _prepare_encoder_decoder_inputs_fast( sequence_lengths = buffers['sequence_lengths'][:num_sequences] attn_metadata._seq_lens = sequence_lengths - if attn_metadata.is_cuda_graph and attn_metadata._seq_lens_cuda is not None: + if (attn_metadata.is_cuda_graph + and attn_metadata._seq_lens_cuda is not None): attn_metadata._seq_lens_cuda.copy_(sequence_lengths, non_blocking=True) else: attn_metadata._seq_lens_cuda = sequence_lengths.cuda( non_blocking=True) + attn_metadata._num_contexts = scheduled_requests.num_context_requests attn_metadata._num_ctx_tokens = num_context_tokens attn_metadata._num_generations = num_generation_requests @@ -3372,9 +3642,10 @@ def _prepare_encoder_decoder_inputs_fast( ) attn_all_rank_num_tokens = self._get_all_rank_num_tokens(attn_metadata) - padded_num_tokens, can_run_piecewise_cuda_graph, attn_all_rank_num_tokens = self._get_padding_params( - total_num_tokens, scheduled_requests.num_context_requests, - attn_all_rank_num_tokens) + (padded_num_tokens, can_run_piecewise_cuda_graph, + attn_all_rank_num_tokens) = self._get_padding_params( + total_num_tokens, scheduled_requests.num_context_requests, + attn_all_rank_num_tokens) set_per_request_piecewise_cuda_graph_flag(can_run_piecewise_cuda_graph) attn_metadata.padded_num_tokens = (padded_num_tokens if padded_num_tokens @@ -3405,8 +3676,7 @@ def _prepare_encoder_decoder_inputs_fast( self.iter_states['num_generation_tokens'] = num_generation_requests self.iter_states['cached_kv_tokens'] = cached_kv_tokens if not self.is_warmup: - self.previous_request_ids = request_ids[scheduled_requests. - num_context_requests:] + self.previous_request_ids = generation_request_ids self.has_previous_device_draft = False event = torch.cuda.Event() @@ -4086,6 +4356,7 @@ def _prepare_tp_inputs( # defensively so the two fast paths can never interleave if the # gates ever evolve. self._steady_gen_cache = None + self._encoder_decoder_staged_request_ids = None return self._apply_incremental_update( scheduled_requests, kv_cache_manager, attn_metadata, spec_metadata, new_tensors_device, cache_indirection_buffer, @@ -4100,6 +4371,7 @@ def _prepare_tp_inputs( scheduled_requests, kv_cache_manager, attn_metadata, new_tokens_device, resource_manager) + self._encoder_decoder_staged_request_ids = None if self._can_use_steady_gen_fast_prepare(scheduled_requests, new_tokens_device, next_draft_tokens_device, @@ -6375,6 +6647,8 @@ def forward(self, padded_requests.all_requests()) self._sync_group_all_greedy_sample(spec_metadata) + allow_mixed_encoder_decoder_graph = ( + self.cuda_graph_runner.enable_encoder_decoder_mixed_cuda_graph) maybe_attn_metadata, maybe_spec_metadata, key = self.cuda_graph_runner.maybe_get_cuda_graph( padded_requests, enable_spec_decode=self.enable_spec_decode, @@ -6384,6 +6658,7 @@ def forward(self, if self.is_spec_decode else None, new_tensors_device=new_tensors_device, spec_resource_manager=spec_resource_manager, + allow_mixed_encoder_decoder=(allow_mixed_encoder_decoder_graph), ) can_run_graph = key is not None @@ -6630,50 +6905,21 @@ def _forward_step_mm_encoder_only( return result - @nvtx_range("_prepare_tp_inputs_encoder") - def _prepare_tp_inputs_encoder( + def _prepare_encoder_decoder_encoder_inputs( self, - encoder_requests: List[LlmRequest], + encoder_input_ids: List[int], + encoder_position_ids: List[int], + sequence_lengths: List[int], + request_ids: List[int], resource_manager: Optional[ResourceManager] = None, - ): - """Pack encoder-side inputs for an encoder-decoder forward pass. - - Mirrors the no-cache path used by ``mm_encoder_only`` and the - legacy ``EncoderBuffers`` shape contract: ``encoder_input_ids`` - and ``encoder_position_ids`` are concatenated across requests - into a single ``[sum(encoder_output_len)]`` tensor, with one - non-causal :class:`AttentionMetadata` describing the packed - encoder batch. - - The encoder pass does not touch any KV-cache pool. The cross pool is - only written by the decoder's cross-attention on the first context - step. Self-pool blocks for the decoder are reserved on the next - scheduler iteration when the request transitions to ``CONTEXT_INIT``. - """ - if not encoder_requests: - raise ValueError( - "_prepare_tp_inputs_encoder called with no encoder requests") - - encoder_input_ids: List[int] = [] - encoder_position_ids: List[int] = [] - sequence_lengths: List[int] = [] - request_ids: List[int] = [] - - for request in encoder_requests: - tokens = request.encoder_tokens - if tokens is None: - raise ValueError( - f"Encoder request {request.py_request_id} has no " - "encoder_tokens; encoder_input_token_ids must be wired " - "through executor_request_to_llm_request.") - seq_len = len(tokens) - encoder_input_ids.extend(tokens) - encoder_position_ids.extend( - self._apply_position_id_offset(list(range(seq_len)))) - sequence_lengths.append(seq_len) - request_ids.append(request.py_request_id) - + ) -> Dict[str, Any]: + if len(sequence_lengths) != len(request_ids): + raise ValueError("Encoder sequence lengths and request IDs must " + "have the same length.") num_tokens = len(encoder_input_ids) + if num_tokens != len(encoder_position_ids): + raise ValueError("Encoder input IDs and position IDs must have " + "the same length.") assert num_tokens <= self.max_num_tokens, ( f"encoder packed length ({num_tokens}) exceeds max_num_tokens " f"({self.max_num_tokens})") @@ -6708,7 +6954,7 @@ def _prepare_tp_inputs_encoder( dtype=torch.int, pin_memory=prefer_pinned(), ) - encoder_attn_metadata.num_contexts = len(encoder_requests) + encoder_attn_metadata.num_contexts = len(sequence_lengths) encoder_attn_metadata.max_seq_len = self.max_seq_len encoder_attn_metadata.request_ids = request_ids encoder_attn_metadata.prepare_encoder_only() @@ -6719,20 +6965,80 @@ def _prepare_tp_inputs_encoder( encoder_position_ids_t = torch.tensor(encoder_position_ids, dtype=torch.int, pin_memory=prefer_pinned()) + use_graph_staging = ( + self.encoder_cuda_graph_runner.enabled and len(sequence_lengths) + in self.encoder_cuda_graph_runner.supported_batch_sizes) - inputs = { + return { 'encoder_input_ids': - encoder_input_ids_t.to('cuda', non_blocking=True), + (encoder_input_ids_t if use_graph_staging else + encoder_input_ids_t.to('cuda', non_blocking=True)), 'encoder_position_ids': - encoder_position_ids_t.to('cuda', non_blocking=True).unsqueeze(0), + ((encoder_position_ids_t + if use_graph_staging else encoder_position_ids_t.to( + 'cuda', non_blocking=True)).unsqueeze(0)), 'encoder_attn_metadata': encoder_attn_metadata, 'encoder_seq_lens': sequence_lengths, + 'encoder_input_ids_host': + encoder_input_ids_t, + 'encoder_position_ids_host': + encoder_position_ids_t, 'resource_manager': resource_manager, } - return inputs + + @nvtx_range("_prepare_tp_inputs_encoder") + def _prepare_tp_inputs_encoder( + self, + encoder_requests: List[LlmRequest], + resource_manager: Optional[ResourceManager] = None, + ): + """Pack encoder-side inputs for an encoder-decoder forward pass. + + Mirrors the no-cache path used by ``mm_encoder_only`` and the + legacy ``EncoderBuffers`` shape contract: ``encoder_input_ids`` + and ``encoder_position_ids`` are concatenated across requests + into a single ``[sum(encoder_output_len)]`` tensor, with one + non-causal :class:`AttentionMetadata` describing the packed + encoder batch. + + The encoder pass does not touch any KV-cache pool. The cross pool is + only written by the decoder's cross-attention on the first context + step. Self-pool blocks for the decoder are reserved on the next + scheduler iteration when the request transitions to ``CONTEXT_INIT``. + """ + if not encoder_requests: + raise ValueError( + "_prepare_tp_inputs_encoder called with no encoder requests") + + encoder_input_ids: List[int] = [] + encoder_position_ids: List[int] = [] + sequence_lengths: List[int] = [] + request_ids: List[int] = [] + + for request in encoder_requests: + tokens = request.encoder_tokens + if tokens is None: + raise ValueError( + f"Encoder request {request.py_request_id} has no " + "encoder_tokens; encoder_input_token_ids must be wired " + "through executor_request_to_llm_request.") + seq_len = len(tokens) + encoder_input_ids.extend(tokens) + encoder_position_ids.extend( + self._apply_position_id_offset(list(range(seq_len)))) + sequence_lengths.append(seq_len) + request_ids.append(request.py_request_id) + + return self._prepare_encoder_decoder_encoder_inputs( + encoder_input_ids=encoder_input_ids, + encoder_position_ids=encoder_position_ids, + sequence_lengths=sequence_lengths, + request_ids=request_ids, + resource_manager=resource_manager, + ) @nvtx_range("_forward_step_encoder") def _forward_step_encoder( @@ -6787,6 +7093,89 @@ def _forward_step_encoder( ) return encoder_hidden_states + def _forward_step_encoder_cuda_graph( + self, + inputs: Dict[str, Any], + ) -> torch.Tensor: + return self._forward_step_encoder({ + 'encoder_input_ids': + inputs['input_ids'], + 'encoder_position_ids': + inputs.get('position_ids'), + 'encoder_attn_metadata': + inputs['attn_metadata'], + 'resource_manager': + inputs.get('resource_manager'), + }) + + def _forward_encoder_with_cuda_graph( + self, + inputs: Dict[str, Any], + ) -> torch.Tensor: + """Replay a bucketed dynamic-layout graph when the encoder is eligible.""" + input_ids = inputs.get('encoder_input_ids_host') + position_ids = inputs.get('encoder_position_ids_host') + seq_lens = inputs['encoder_seq_lens'] + actual_num_tokens = sum(seq_lens) + runner = self.encoder_cuda_graph_runner + + if input_ids is None or position_ids is None: + return self._forward_step_encoder(inputs) + + runner_inputs = { + 'input_ids': input_ids, + 'position_ids': position_ids, + 'seq_lens': seq_lens, + 'resource_manager': inputs.get('resource_manager'), + } + graph_attn_metadata, key = runner.maybe_get_cuda_graph( + runner_inputs, inputs['encoder_attn_metadata']) + if key is None: + if inputs['encoder_input_ids'].device.type == 'cpu': + inputs = dict(inputs) + inputs['encoder_input_ids'] = inputs['encoder_input_ids'].to( + 'cuda', non_blocking=True) + inputs['encoder_position_ids'] = inputs[ + 'encoder_position_ids'].to('cuda', non_blocking=True) + return self._forward_step_encoder(inputs) + + # Every graph key aliases the same pinned staging allocation. Retire + # the previous captured H2D before updating seq_lens or any other + # shared host input for this replay. + runner.retire_staging() + graph_attn_metadata.prepare_encoder_cuda_graph_replay( + seq_lens, actual_num_tokens) + model_inputs = { + **runner_inputs, + 'attn_metadata': graph_attn_metadata, + } + + moe_load_balancer: MoeLoadBalancer = getattr(self, 'moe_load_balancer', + None) + with with_shared_pool(runner.get_graph_pool()): + capture_outputs = None + if runner.needs_capture(key): + + def capture_forward_fn( + capture_inputs: Dict[str, Any]) -> torch.Tensor: + with MoeLoadBalancerIterContext(moe_load_balancer): + return self._forward_step_encoder_cuda_graph( + capture_inputs) + + capture_outputs = runner.capture(key, capture_forward_fn, + model_inputs) + + if runner.is_warmup_only: + graph_outputs = capture_outputs + else: + with MoeLoadBalancerIterContext(moe_load_balancer): + graph_outputs = runner.replay(key, model_inputs) + + if not isinstance(graph_outputs, torch.Tensor): + raise TypeError("Encoder-decoder CUDA graph replay must return " + "a tensor of encoder hidden states.") + return graph_outputs[:actual_num_tokens].clone() + @nvtx_range("forward_encoder") def forward_encoder( self, @@ -6813,7 +7202,8 @@ def forward_encoder( with torch.inference_mode(): inputs = self._prepare_tp_inputs_encoder( encoder_requests, resource_manager=resource_manager) - encoder_hidden_states = self._forward_step_encoder(inputs) + encoder_hidden_states = self._forward_encoder_with_cuda_graph( + inputs) return encoder_hidden_states, inputs['encoder_seq_lens'] diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 34b0067601bf..ea44ec7b885f 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -8,6 +8,7 @@ import threading import time import traceback +from concurrent.futures import Future, ThreadPoolExecutor from contextlib import contextmanager from enum import IntEnum from queue import Queue @@ -379,6 +380,20 @@ class BatchStatePP(BatchState): microbatch_id: int = -1 +@dataclasses.dataclass +class EncoderStepResult: + hidden_states: torch.Tensor + sequence_lengths: List[int] + ready_event: torch.cuda.Event + + +@dataclasses.dataclass +class PendingEncoderStep: + requests: List[LlmRequest] + future: Future[EncoderStepResult] + result: Optional[EncoderStepResult] = None + + class AsyncTransferManager: """ Handle asynchronous transfer of KV cache after a request has completed. @@ -812,6 +827,11 @@ def __init__( # Ensure the default stream waits for execution_stream to complete # before subsequent operations. torch.cuda.current_stream().wait_stream(self.execution_stream) + self.encoder_launch_executor = (ThreadPoolExecutor( + max_workers=1, thread_name_prefix="encoder-launch") + if self.is_encoder_decoder else None) + self.pending_encoder_steps: List[PendingEncoderStep] = [] + self.is_warmup = False # Snapshot some cumulative KV cache counters so that stats reported to @@ -1452,6 +1472,11 @@ def shutdown(self): # executor loops have already processed the shutdown broadcast and are # no longer driving NCCL, so the send cannot deadlock. self._shutdown_sleep_wakeup_listeners() + encoder_launch_executor = getattr(self, "encoder_launch_executor", None) + if encoder_launch_executor is not None: + encoder_launch_executor.shutdown(wait=True) + self.encoder_launch_executor = None + self.worker_started = False # Release CUDA graphs before resource managers free their GPU memory. # Resource managers (e.g. SuffixAutomatonManager) allocate GPU workspace @@ -3572,6 +3597,7 @@ def _sync_gen_only_benchmark_has_insufficient_kv( def _prepare_and_schedule_batch(self): self._sync_disagg_transfer_made_progress = False + self._poll_encoder_steps() new_requests = self._fetch_and_activate_new_requests() if self.should_stop_processing: return None, None @@ -3991,15 +4017,6 @@ def _executor_loop(self): gpu_forward_end = None gpu_forward_events_from_perf_pool = False - # Run the encoder iteration first. After scatter the - # encoder requests transition to ``CONTEXT_INIT`` and are - # picked up by the next scheduler iteration as decoder - # context. The encoder pass is independent of the decoder - # ``can_queue`` gate, so an iteration with only encoder-init - # requests still makes forward progress. - if scheduled_batch.encoder_requests: - self._run_encoder_step(scheduled_batch.encoder_requests) - can_queue, _ = self._can_queue(scheduled_batch) if can_queue: @@ -4029,6 +4046,9 @@ def _executor_loop(self): self._revert_gen_alloc(scheduled_batch) self._finalize_adp_dummy_allocation(can_queue) + if not can_queue and scheduled_batch.encoder_requests: + self._run_encoder_step(scheduled_batch.encoder_requests) + if can_queue: # init_disagg_gen_requests must be before drafter loop, otherwise draft requests do not have initialized matchers. # init_disagg_gen_requests must be before engine forward, where the prev_seq_slot is updated. @@ -4073,6 +4093,10 @@ def _executor_loop(self): ) gpu_forward_events_from_perf_pool = True + if scheduled_batch.encoder_requests: + self._submit_encoder_step( + scheduled_batch.encoder_requests) + with self.perf_manager.record_perf_events( gpu_forward_start, gpu_forward_end) as fwd_timing: if self.dwdp_manager is not None: @@ -4459,9 +4483,6 @@ def _executor_loop_overlap(self): if not self._is_kv_manager_v2: self._terminate_requests(scheduled_batch.paused_requests) - if scheduled_batch.encoder_requests: - self._run_encoder_step(scheduled_batch.encoder_requests) - gpu_forward_events_from_perf_pool = False can_queue, can_queue_this_rank = self._can_queue( scheduled_batch) @@ -4507,6 +4528,9 @@ def _executor_loop_overlap(self): self._revert_gen_alloc(scheduled_batch) self._finalize_adp_dummy_allocation(can_queue) + if not can_queue and scheduled_batch.encoder_requests: + self._run_encoder_step(scheduled_batch.encoder_requests) + # If the batch is not empty on this rank, but empty on other ranks, # we need to delay the update of the previous batch's sample state, # and let the later iteration to update it. @@ -4572,6 +4596,9 @@ def _executor_loop_overlap(self): gpu_forward_start, gpu_forward_end = self.perf_manager.borrow_forward_timing_events( ) gpu_forward_events_from_perf_pool = True + if scheduled_batch.encoder_requests: + self._submit_encoder_step( + scheduled_batch.encoder_requests) with self.perf_manager.record_perf_events( gpu_forward_start, gpu_forward_end) as fwd_timing: @@ -5231,6 +5258,7 @@ def _waiting_requests(self, context_requests: list[LlmRequest], def _waiting_encoder_requests( self, encoder_requests: list[LlmRequest], + context_requests: list[LlmRequest], generation_requests: list[LlmRequest]) -> list[LlmRequest]: """Accumulate encoder work while an admitted decode batch progresses. @@ -5240,12 +5268,51 @@ def _waiting_encoder_requests( the executor thread. Decoder generation continues while the encoder requests wait. The encoder has its own counter because the resulting decoder-context requests are already coalesced and must not wait for a - second window. + second window. CUDA-graph microbatch admission returns at most one + supported graph batch so scheduler overfill cannot turn a target-eight + batch into an eager batch of nine or more requests. """ - if not encoder_requests or not generation_requests: + if not encoder_requests: self.encoder_batch_wait_iters_count = 0 return encoder_requests + microbatch_graph_max_batch_size = int( + os.environ.get( + "TLLM_ENCODER_DECODER_MICROBATCH_CUDA_GRAPH_MAX_BATCH_SIZE", + "0")) + microbatch_admission_enabled = (os.environ.get( + "TLLM_ENCODER_DECODER_MICROBATCH_ADMISSION_ENABLED", "1") == "1") + if (microbatch_graph_max_batch_size > 0 + and microbatch_admission_enabled): + microbatch_target = min(microbatch_graph_max_batch_size, + self.max_batch_size) + decoder_occupancy = (len(context_requests) + + len(generation_requests)) + decoder_low_watermark = int( + os.environ.get("TLLM_ENCODER_DECODER_MICROBATCH_LOW_WATERMARK", + str(self.max_batch_size - microbatch_target))) + deadline_reached = (self.encoder_batch_wait_iters_count + >= self.batch_wait_timeout_iters) + + if (len(encoder_requests) >= microbatch_target + and decoder_occupancy <= decoder_low_watermark): + self.encoder_batch_wait_iters_count = 0 + return encoder_requests[:microbatch_target] + + if deadline_reached: + supported_batch_sizes = [ + batch_size for batch_size in (1, 2, 4, 8) + if batch_size <= microbatch_target + ] + fallback_batch_size = max( + batch_size for batch_size in supported_batch_sizes + if batch_size <= len(encoder_requests)) + self.encoder_batch_wait_iters_count = 0 + return encoder_requests[:fallback_batch_size] + + self.encoder_batch_wait_iters_count += 1 + return [] + num_scheduled_tokens = sum(request.encoder_output_len for request in encoder_requests) num_scheduled_tokens += sum(1 + request.num_draft_tokens @@ -5276,6 +5343,7 @@ def _schedule(self): if should_batch_encoder_requests: scheduled_encoder_requests = self._waiting_encoder_requests( scheduler_output.encoder_requests, + scheduler_output.context_requests, scheduler_output.generation_requests) scheduled_context_requests = scheduler_output.context_requests @@ -5336,44 +5404,131 @@ def _schedule(self): # micro-batch; this preserves the cross-KV lifecycle and the # dual-pool budget. # --------------------------------------------------------------- - @nvtx_range("_run_encoder_step") - def _run_encoder_step(self, encoder_requests: List[LlmRequest]) -> None: - """Drive one encoder iteration for ``encoder_requests``. + def _submit_encoder_step(self, encoder_requests: List[LlmRequest]) -> None: + """Queue encoder work without blocking the decoder executor thread.""" + executor = self.encoder_launch_executor + if executor is None: + raise RuntimeError("Encoder launch executor is unavailable.") - Runs the encoder stack on the dedicated encoder stream, then - scatters the packed hidden states back onto the per-request - ``py_encoder_output`` field and transitions request state to - ``CONTEXT_INIT`` so the next scheduler pass picks them up as - decoder-context requests. A separate CUDA event is recorded for - each request on the encoder stream; the scheduler queries that - event before admitting the request to a decoder context step. + requests = list(encoder_requests) + for request in requests: + self.inflight_req_ids.insert(request.request_id) + + try: + future = executor.submit(self._run_encoder_step_unchecked, requests) + except Exception: + for request in requests: + self.inflight_req_ids.erase(request.request_id) + raise + + self.pending_encoder_steps.append( + PendingEncoderStep(requests=requests, future=future)) + + @nvtx_range("_poll_encoder_steps") + def _poll_encoder_steps(self) -> None: + """Publish ready encoder batches without waiting for their futures. + + Encoder request IDs stay in ``inflight_req_ids`` from submission + until both the launch worker and its CUDA completion event finish. + Consequently the scheduler cannot submit an encoder request twice or + admit its decoder-context step before the encoder output is ready. """ - if not encoder_requests: + pending_steps = getattr(self, "pending_encoder_steps", None) + if not pending_steps: return + while pending_steps: + pending = pending_steps[0] + if pending.result is None: + if not pending.future.done(): + break + try: + pending.result = pending.future.result() + except Exception as e: + pending_steps.pop(0) + self._finish_failed_encoder_step(pending.requests, e) + continue + + if not pending.result.ready_event.query(): + break + + pending_steps.pop(0) + try: + self._publish_encoder_step(pending.requests, pending.result) + except Exception as e: + self._finish_failed_encoder_step(pending.requests, e) + continue + + for request in pending.requests: + self.inflight_req_ids.erase(request.request_id) + + def _finish_failed_encoder_step(self, encoder_requests: List[LlmRequest], + error: Exception) -> None: + for request in encoder_requests: + self.inflight_req_ids.erase(request.request_id) + traceback.print_exception(error) + error_msg = str(error) + logger.error(f"Encountered an error in encoder forward: {error_msg}") + failed_requests = [ + request for request in encoder_requests + if request.state != LlmRequestState.GENERATION_COMPLETE + ] + if failed_requests: + self._handle_errors(error_msg, requests=failed_requests) + + def _run_encoder_step(self, encoder_requests: List[LlmRequest]) -> None: try: - self.encoder_stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(self.encoder_stream): - encoder_hidden_states, encoder_seq_lens = ( - self.model_engine.forward_encoder( - encoder_requests, - resource_manager=self.resource_manager, - )) + result = self._run_encoder_step_unchecked(encoder_requests) + self._publish_encoder_step(encoder_requests, result) except Exception as e: - traceback.print_exc() - error_msg = str(e) - logger.error( - f"Encountered an error in encoder forward: {error_msg}") - self._handle_errors(error_msg, requests=encoder_requests) - return + self._finish_failed_encoder_step(encoder_requests, e) + + @nvtx_range("_run_encoder_step") + def _run_encoder_step_unchecked( + self, encoder_requests: List[LlmRequest]) -> EncoderStepResult: + """Drive one encoder iteration for ``encoder_requests``. - self._scatter_encoder_output(encoder_requests, encoder_hidden_states, - encoder_seq_lens) - for req in encoder_requests: - req.py_encoder_output_ready_event = torch.cuda.Event() - req.py_encoder_output_ready_event.record(self.encoder_stream) - # TODO(TRTLLM-12339): Honor return_encoder_output once the public - # LLM API shape for returned encoder hidden states is finalized. + Runs the encoder stack on its independent stream and returns packed + output plus a completion event. Request state remains owned by the + main executor thread and is updated by ``_poll_encoder_steps`` only + after this event reports ready. + + The caller submits encoder work immediately before decoder forward so + their independent CPU launch and GPU execution can overlap. Encoder + inputs are staged inside this stream, and request-level completion + events guard the only downstream dependency. Successive encoder + batches remain ordered by the stream itself. + """ + if not encoder_requests: + raise ValueError("Encoder step requires at least one request.") + + torch.cuda.set_device(self.device_id) + with torch.cuda.stream(self.encoder_stream): + encoder_hidden_states, encoder_seq_lens = ( + self.model_engine.forward_encoder( + encoder_requests, + resource_manager=self.resource_manager, + )) + + ready_event = torch.cuda.Event() + ready_event.record(self.encoder_stream) + return EncoderStepResult( + hidden_states=encoder_hidden_states, + sequence_lengths=encoder_seq_lens, + ready_event=ready_event, + ) + + def _publish_encoder_step(self, encoder_requests: List[LlmRequest], + result: EncoderStepResult) -> None: + """Make a completed encoder batch visible to decoder scheduling.""" + self._scatter_encoder_output( + encoder_requests, + result.hidden_states, + result.sequence_lengths, + result.ready_event, + ) + # TODO(TRTLLM-12339): Honor return_encoder_output once the public + # LLM API shape for returned encoder hidden states is finalized. @nvtx_range("_scatter_encoder_output") def _scatter_encoder_output( @@ -5381,6 +5536,7 @@ def _scatter_encoder_output( encoder_requests: List[LlmRequest], encoder_hidden_states: torch.Tensor, encoder_seq_lens: List[int], + ready_event: torch.cuda.Event, ) -> None: """Slice packed encoder hidden states into per-request tensors. @@ -5409,10 +5565,12 @@ def _scatter_encoder_output( offset = 0 for req, seq_len in zip(encoder_requests, encoder_seq_lens): - req.py_encoder_output = encoder_hidden_states[offset:offset + - seq_len] - req.py_skip_cross_kv_projection = False - req.state = LlmRequestState.CONTEXT_INIT + if req.state == LlmRequestState.ENCODER_INIT: + req.py_encoder_output = encoder_hidden_states[offset:offset + + seq_len] + req.py_skip_cross_kv_projection = False + req.py_encoder_output_ready_event = ready_event + req.state = LlmRequestState.CONTEXT_INIT offset += seq_len @nvtx_range("_attach_encoder_output_to_execution_stream") diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py index caa8e3cb3de1..6b36cf898135 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + import dataclasses import inspect from abc import ABC, abstractmethod diff --git a/tests/unittest/_torch/executor/test_encoder_cuda_graph_runner.py b/tests/unittest/_torch/executor/test_encoder_cuda_graph_runner.py new file mode 100644 index 000000000000..f01070f2b42d --- /dev/null +++ b/tests/unittest/_torch/executor/test_encoder_cuda_graph_runner.py @@ -0,0 +1,225 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttentionMetadata +from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import ( + EncoderCUDAGraphRunner, + EncoderCUDAGraphRunnerConfig, +) + + +def _dynamic_layout_runner( + max_cuda_graphs: int = 0, capture_keys: list[tuple[int, int, int]] | None = None +) -> EncoderCUDAGraphRunner: + return EncoderCUDAGraphRunner( + EncoderCUDAGraphRunnerConfig( + use_cuda_graph=False, + cuda_graph_padding_enabled=False, + cuda_graph_batch_sizes=[1, 2, 4, 8], + cuda_graph_num_tokens=[], + cuda_graph_seq_lens=list(range(64, 513, 64)), + max_cuda_graph_batch_size=8, + max_cuda_graph_num_tokens=4096, + max_num_tokens=4096, + max_seq_len=512, + cuda_graph_mem_pool=None, + dynamic_sequence_layout=True, + allow_runtime_capture=True, + max_cuda_graphs=max_cuda_graphs, + capture_keys=capture_keys or [], + ) + ) + + +def test_encoder_graph_key_reuses_total_tokens_and_max_bucket(): + runner = _dynamic_layout_runner() + + key, is_padding_performed, is_valid = runner.get_graph_key( + { + "input_ids": list(range(580)), + "seq_lens": [260, 320], + } + ) + other_layout_key, _, _ = runner.get_graph_key( + { + "input_ids": list(range(580)), + "seq_lens": [284, 296], + } + ) + + assert key == (2, 580, 320) + assert other_layout_key == key + assert not is_padding_performed + assert is_valid + + +def test_encoder_graph_key_distinguishes_max_buckets(): + runner = _dynamic_layout_runner() + + key, _, _ = runner.get_graph_key( + { + "input_ids": [0] * 1400, + "seq_lens": [332, 356, 356, 356], + } + ) + larger_bucket_key, _, _ = runner.get_graph_key( + { + "input_ids": [0] * 1400, + "seq_lens": [260, 260, 440, 440], + } + ) + + assert key == (4, 1400, 384) + assert larger_bucket_key == (4, 1400, 448) + + +def test_bart_microbatch_key_set_fits_graph_cache(): + runner = _dynamic_layout_runner(max_cuda_graphs=64) + sequence_length_cycle = list(range(260, 441, 12)) + keys = set() + + for batch_size in (1, 2, 4, 8): + for start in range(len(sequence_length_cycle)): + sequence_lengths = [ + sequence_length_cycle[(start + offset) % len(sequence_length_cycle)] + for offset in range(batch_size) + ] + key, _, is_valid = runner.get_graph_key( + { + "input_ids": [0] * sum(sequence_lengths), + "seq_lens": sequence_lengths, + } + ) + assert is_valid + keys.add(key) + + assert len(keys) == 59 + assert len(keys) <= runner.max_cuda_graphs + + +def test_encoder_graph_key_rejects_oversized_inputs(): + runner = _dynamic_layout_runner() + + _, _, is_valid = runner.get_graph_key( + { + "input_ids": [0] * 4097, + "seq_lens": [4097], + } + ) + + assert not is_valid + + +def test_encoder_graph_capture_allowlist_must_fit_cache(): + capture_keys = [(1, num_tokens, 64) for num_tokens in range(1, 66)] + + with pytest.raises(ValueError, match="capture key count"): + _dynamic_layout_runner(max_cuda_graphs=64, capture_keys=capture_keys) + + +def test_encoder_graph_reuses_same_key_for_different_sequence_layouts(): + runner = _dynamic_layout_runner() + runner.enabled = True + graph_metadata = object.__new__(TrtllmAttentionMetadata) + key = (2, 580, 320) + runner.graph_metadata[key] = { + "attn_metadata": graph_metadata, + } + + matched_metadata, matched_key = runner.maybe_get_cuda_graph( + { + "input_ids": [0] * 580, + "seq_lens": [260, 320], + }, + graph_metadata, + ) + reused_metadata, reused_key = runner.maybe_get_cuda_graph( + { + "input_ids": [0] * 580, + "seq_lens": [284, 296], + }, + graph_metadata, + ) + + assert matched_metadata is graph_metadata + assert matched_key == key + assert reused_metadata is graph_metadata + assert reused_key == key + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_encoder_graph_capture_stages_warmup_and_replays_new_layout(): + runner = _dynamic_layout_runner() + runner.enabled = True + runner._create_shared_static_tensors() + + key = (2, 5, 64) + seq_lens_host = runner.shared_static_tensors_cpu["seq_lens"][:2] + seq_lens_host.copy_(torch.tensor([2, 3], dtype=torch.int32)) + attn_metadata = SimpleNamespace( + _seq_lens=seq_lens_host, + _seq_lens_cuda=torch.ones(2, device="cuda", dtype=torch.int32), + ) + inputs = { + "input_ids": [10, 11, 12, 13, 14], + "position_ids": [0, 1, 0, 1, 2], + "seq_lens": [2, 3], + "attn_metadata": attn_metadata, + } + warmup_seq_lens = [] + + def forward_fn(capture_inputs): + if not torch.cuda.is_current_stream_capturing(): + warmup_seq_lens.append(capture_inputs["attn_metadata"]._seq_lens_cuda.cpu().tolist()) + return capture_inputs["input_ids"] + capture_inputs["attn_metadata"]._seq_lens_cuda[0] + + runner.capture(key, forward_fn, inputs) + + assert warmup_seq_lens == [[2, 3]] + + first_output = runner.replay(key, inputs) + torch.cuda.synchronize() + torch.testing.assert_close( + first_output, + torch.tensor([12, 13, 14, 15, 16], device="cuda", dtype=torch.int32), + ) + + seq_lens_host.copy_(torch.tensor([1, 4], dtype=torch.int32)) + reused_inputs = { + **inputs, + "seq_lens": [1, 4], + } + reused_output = runner.replay(key, reused_inputs) + torch.cuda.synchronize() + torch.testing.assert_close( + reused_output, + torch.tensor([11, 12, 13, 14, 15], device="cuda", dtype=torch.int32), + ) + + +def test_encoder_graph_lru_evicts_oldest_graph(): + class _Graph: + def __init__(self): + self.was_reset = False + + def reset(self): + self.was_reset = True + + runner = _dynamic_layout_runner(max_cuda_graphs=1) + key = (1, 8, 64) + graph = _Graph() + runner.graphs[key] = graph + runner.graph_outputs[key] = object() + runner.graph_metadata[key] = object() + + runner._evict_graph_if_needed() + + assert graph.was_reset + assert not runner.graphs + assert not runner.graph_outputs + assert not runner.graph_metadata diff --git a/tests/unittest/_torch/executor/test_mixed_decoder_cuda_graph_runner.py b/tests/unittest/_torch/executor/test_mixed_decoder_cuda_graph_runner.py new file mode 100644 index 000000000000..f4412dd23aa7 --- /dev/null +++ b/tests/unittest/_torch/executor/test_mixed_decoder_cuda_graph_runner.py @@ -0,0 +1,79 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace + +from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDAGraphRunner +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests + + +def _mixed_batch() -> ScheduledRequests: + batch = ScheduledRequests() + batch.context_requests_last_chunk = [ + SimpleNamespace( + context_chunk_size=2, + encoder_output_len=260, + py_skip_cross_kv_projection=False, + ), + SimpleNamespace( + context_chunk_size=2, + encoder_output_len=272, + py_skip_cross_kv_projection=False, + ), + ] + batch.generation_requests = [ + SimpleNamespace(py_draft_tokens=[]), + SimpleNamespace(py_draft_tokens=[]), + ] + return batch + + +def _runner() -> CUDAGraphRunner: + runner = object.__new__(CUDAGraphRunner) + runner.config = SimpleNamespace(is_draft_model=False) + runner.sparse_config = None + runner.max_beam_width = 1 + runner.enable_encoder_decoder_mixed_cuda_graph = True + runner.graphs = {} + runner.graph_outputs = {} + runner.graph_metadata = {} + runner.padding_dummy_requests = {} + runner.memory_pool = None + return runner + + +def test_mixed_encoder_decoder_graph_key_captures_dynamic_extents(): + runner = _runner() + + key = runner.get_graph_key(_mixed_batch()) + + assert key == (4, 0, False, False, True, (2, 2), (532,)) + assert runner._get_num_tokens_for_key(key) == 6 + + +def test_mixed_encoder_decoder_graph_key_distinguishes_cached_cross_kv(): + runner = _runner() + batch = _mixed_batch() + batch.context_requests_last_chunk[1].py_skip_cross_kv_projection = True + + key = runner.get_graph_key(batch) + + assert key[6] == (260,) + + +def test_mixed_encoder_decoder_graph_eligibility_requires_both_phases(): + runner = _runner() + batch = _mixed_batch() + + assert runner._is_mixed_encoder_decoder_batch(batch) + + batch.generation_requests = [] + assert not runner._is_mixed_encoder_decoder_batch(batch) + + +def test_mixed_encoder_decoder_graph_never_captures_at_runtime(): + runner = _runner() + runner._capture_allowed = False + key = runner.get_graph_key(_mixed_batch()) + + assert not runner.needs_capture(key) diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 101edffbe155..6340ce22ab9e 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -19,6 +19,7 @@ from unittest.mock import MagicMock, Mock import pytest +import torch from tensorrt_llm._torch.distributed.communicator import ReduceOp from tensorrt_llm._torch.pyexecutor.executor_request_queue import ( @@ -26,7 +27,11 @@ RequestQueueItem, ) from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestState, SamplingConfig -from tensorrt_llm._torch.pyexecutor.py_executor import DisaggTransferAdmissionController, PyExecutor +from tensorrt_llm._torch.pyexecutor.py_executor import ( + DisaggTransferAdmissionController, + EncoderStepResult, + PyExecutor, +) from tensorrt_llm._torch.pyexecutor.resource_manager import NoFreeSlotsError, ResourceManagerType from tensorrt_llm._torch.pyexecutor.scheduler import ( FCFSWaitingQueue, @@ -35,6 +40,17 @@ ) +class _InflightRequestIds: + def __init__(self): + self.ids = set() + + def insert(self, request_id): + self.ids.add(request_id) + + def erase(self, request_id): + self.ids.discard(request_id) + + class MockPyExecutor: """A mock PyExecutor class for testing request handling logic. @@ -117,6 +133,205 @@ def mock_dist(): return mock_dist +def _make_async_encoder_executor(future): + executor = object.__new__(PyExecutor) + executor.encoder_launch_executor = Mock() + executor.encoder_launch_executor.submit.return_value = future + executor.pending_encoder_steps = [] + executor.inflight_req_ids = _InflightRequestIds() + executor._run_encoder_step_unchecked = Mock() + executor._publish_encoder_step = Mock() + executor._handle_errors = Mock() + return executor + + +def _make_encoder_batch_wait_executor(): + executor = object.__new__(PyExecutor) + executor.max_batch_size = 32 + executor.batch_wait_timeout_iters = 48 + executor.encoder_batch_wait_iters_count = 0 + return executor + + +def test_encoder_microbatch_graph_waits_for_target(monkeypatch): + monkeypatch.setenv("TLLM_ENCODER_DECODER_MICROBATCH_CUDA_GRAPH_MAX_BATCH_SIZE", "8") + monkeypatch.setenv("TLLM_ENCODER_DECODER_MICROBATCH_LOW_WATERMARK", "24") + executor = _make_encoder_batch_wait_executor() + encoder_requests = [object()] * 7 + generation_requests = [object()] * 24 + + scheduled = executor._waiting_encoder_requests( + encoder_requests, + [], + generation_requests, + ) + + assert scheduled == [] + assert executor.encoder_batch_wait_iters_count == 1 + + +def test_encoder_microbatch_graph_releases_target(monkeypatch): + monkeypatch.setenv("TLLM_ENCODER_DECODER_MICROBATCH_CUDA_GRAPH_MAX_BATCH_SIZE", "8") + monkeypatch.setenv("TLLM_ENCODER_DECODER_MICROBATCH_LOW_WATERMARK", "24") + executor = _make_encoder_batch_wait_executor() + encoder_requests = [object()] * 8 + + scheduled = executor._waiting_encoder_requests( + encoder_requests, + [], + [object()] * 24, + ) + + assert scheduled == encoder_requests + assert executor.encoder_batch_wait_iters_count == 0 + + +def test_encoder_microbatch_graph_does_not_release_partial_at_low_watermark(monkeypatch): + monkeypatch.setenv("TLLM_ENCODER_DECODER_MICROBATCH_CUDA_GRAPH_MAX_BATCH_SIZE", "8") + monkeypatch.setenv("TLLM_ENCODER_DECODER_MICROBATCH_LOW_WATERMARK", "24") + executor = _make_encoder_batch_wait_executor() + encoder_requests = [object()] + + scheduled = executor._waiting_encoder_requests( + encoder_requests, + [], + [object()] * 24, + ) + + assert scheduled == [] + assert executor.encoder_batch_wait_iters_count == 1 + + +def test_encoder_microbatch_graph_caps_scheduler_overfill(monkeypatch): + monkeypatch.setenv("TLLM_ENCODER_DECODER_MICROBATCH_CUDA_GRAPH_MAX_BATCH_SIZE", "8") + monkeypatch.setenv("TLLM_ENCODER_DECODER_MICROBATCH_LOW_WATERMARK", "24") + executor = _make_encoder_batch_wait_executor() + encoder_requests = [object() for _ in range(12)] + + scheduled = executor._waiting_encoder_requests( + encoder_requests, + [], + [object()] * 20, + ) + + assert scheduled == encoder_requests[:8] + assert executor.encoder_batch_wait_iters_count == 0 + + +def test_encoder_microbatch_graph_releases_supported_tail_at_deadline(monkeypatch): + monkeypatch.setenv("TLLM_ENCODER_DECODER_MICROBATCH_CUDA_GRAPH_MAX_BATCH_SIZE", "8") + monkeypatch.setenv("TLLM_ENCODER_DECODER_MICROBATCH_LOW_WATERMARK", "24") + executor = _make_encoder_batch_wait_executor() + executor.encoder_batch_wait_iters_count = executor.batch_wait_timeout_iters + encoder_requests = [object() for _ in range(7)] + + scheduled = executor._waiting_encoder_requests( + encoder_requests, + [], + [], + ) + + assert scheduled == encoder_requests[:4] + assert executor.encoder_batch_wait_iters_count == 0 + + +def test_pending_encoder_future_is_polled_without_blocking(): + future = Mock() + future.done.return_value = False + executor = _make_async_encoder_executor(future) + request = types.SimpleNamespace(request_id=11, state=LlmRequestState.ENCODER_INIT) + + executor._submit_encoder_step([request]) + executor._poll_encoder_steps() + + future.result.assert_not_called() + executor._publish_encoder_step.assert_not_called() + assert executor.inflight_req_ids.ids == {11} + assert len(executor.pending_encoder_steps) == 1 + executor.encoder_launch_executor.submit.assert_called_once_with( + executor._run_encoder_step_unchecked, + [request], + ) + + +def test_completed_encoder_future_waits_for_cuda_event_without_blocking(): + ready_event = Mock() + ready_event.query.side_effect = [False, True] + result = EncoderStepResult( + hidden_states=torch.empty((1, 2)), + sequence_lengths=[1], + ready_event=ready_event, + ) + future = Mock() + future.done.return_value = True + future.result.return_value = result + executor = _make_async_encoder_executor(future) + request = types.SimpleNamespace(request_id=12, state=LlmRequestState.ENCODER_INIT) + + executor._submit_encoder_step([request]) + executor._poll_encoder_steps() + + future.result.assert_called_once_with() + ready_event.query.assert_called_once_with() + executor._publish_encoder_step.assert_not_called() + assert executor.inflight_req_ids.ids == {12} + assert len(executor.pending_encoder_steps) == 1 + + executor._poll_encoder_steps() + + future.result.assert_called_once_with() + assert ready_event.query.call_count == 2 + executor._publish_encoder_step.assert_called_once_with([request], result) + assert executor.inflight_req_ids.ids == set() + assert executor.pending_encoder_steps == [] + + +def test_publish_encoder_output_does_not_resurrect_completed_request(): + executor = object.__new__(PyExecutor) + active_request = types.SimpleNamespace(state=LlmRequestState.ENCODER_INIT) + completed_request = types.SimpleNamespace(state=LlmRequestState.GENERATION_COMPLETE) + hidden_states = torch.arange(12).reshape(6, 2) + ready_event = Mock() + + executor._scatter_encoder_output( + [active_request, completed_request], + hidden_states, + [2, 4], + ready_event, + ) + + assert active_request.state == LlmRequestState.CONTEXT_INIT + assert active_request.py_encoder_output_ready_event is ready_event + assert torch.equal(active_request.py_encoder_output, hidden_states[:2]) + assert completed_request.state == LlmRequestState.GENERATION_COMPLETE + assert not hasattr(completed_request, "py_encoder_output") + + +def test_attach_encoder_output_records_stream_after_encoder_is_ready(): + executor = object.__new__(PyExecutor) + executor.execution_stream = Mock() + ready_event = Mock() + first_output = Mock() + second_output = Mock() + first_request = types.SimpleNamespace( + py_encoder_output=first_output, + py_encoder_output_ready_event=ready_event, + ) + second_request = types.SimpleNamespace( + py_encoder_output=second_output, + py_encoder_output_ready_event=ready_event, + ) + scheduled_requests = types.SimpleNamespace(context_requests=[first_request, second_request]) + + executor._attach_encoder_output_to_execution_stream(scheduled_requests) + + executor.execution_stream.wait_event.assert_not_called() + first_output.record_stream.assert_called_once_with(executor.execution_stream) + second_output.record_stream.assert_called_once_with(executor.execution_stream) + assert first_request.py_encoder_output_ready_event is None + assert second_request.py_encoder_output_ready_event is None + + @pytest.fixture def mock_executor(mock_dist): """Create a MockPyExecutor instance for testing.""" From ebccb5d1d8637ddbb77a33ed29b63b5209d1c671 Mon Sep 17 00:00:00 2001 From: Guiju Zhang <7135567+cascade812@users.noreply.github.com> Date: Wed, 29 Jul 2026 14:02:11 -0700 Subject: [PATCH 05/15] [None][perf] enable encoder CUDA graphs for encoder-decoder models Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com> --- docs/source/models/encoder-decoder.md | 41 ++- .../_torch/pyexecutor/cuda_graph_runner.py | 208 ++++++++--- .../_torch/pyexecutor/model_engine.py | 326 ++++++++++-------- tensorrt_llm/_torch/pyexecutor/py_executor.py | 104 ++++-- tensorrt_llm/llmapi/llm_args.py | 51 ++- .../usage/llm_args_golden_manifest.json | 58 ++++ .../test_encoder_cuda_graph_runner.py | 242 ++++++++++--- .../test_mixed_decoder_cuda_graph_runner.py | 82 ++++- .../_torch/executor/test_py_executor.py | 94 ++++- .../test_pytorch_model_engine_warmup.py | 50 +++ .../api_stability/references/llm.yaml | 4 + tests/unittest/llmapi/test_llm_args.py | 50 +++ 12 files changed, 987 insertions(+), 323 deletions(-) diff --git a/docs/source/models/encoder-decoder.md b/docs/source/models/encoder-decoder.md index 30fa8d5695b9..788ee9bd27c8 100644 --- a/docs/source/models/encoder-decoder.md +++ b/docs/source/models/encoder-decoder.md @@ -42,7 +42,7 @@ The following table describes the supported and recommended configurations. | Beam search | Yes with V1 | Configure `max_beam_width` when constructing `LLM`, then set `use_beam_search=True` in `SamplingParams`. | | Attention backend | `TRTLLM` | Use this backend for encoder-decoder models. It is required when `tensor_parallel_size > 1`. | | Decoder CUDA graphs | Yes | `CudaGraphConfig` captures decoder work. V1 supports greedy and beam search; V2 supports its single-beam path. | -| Encoder CUDA graphs | No | `EncodeCudaGraphConfig` is disabled for encoder-decoder models. The encoder runs eagerly. | +| Encoder CUDA graphs | Yes | Set `encoder_cuda_graph_config=EncodeCudaGraphConfig(...)` and `encoder_max_batch_size`. The `TRTLLM` attention backend is required. | | Overlap scheduler | Yes | Enabled by default. V1 supports greedy decoding and beam search; V2 remains limited to `max_beam_width=1`. | | Tensor parallelism | Yes | Use `tensor_parallel_size > 1` with `attn_backend="TRTLLM"`. Attention head counts must be divisible by the TP size. | | Pipeline parallelism | No | Keep `pipeline_parallel_size=1`. | @@ -319,12 +319,12 @@ to return only the best hypothesis. Beam search expands decoder-side cache and compute requirements. Include this expansion when sizing the self-attention KV pool and CUDA graph batch sizes. -## Enable decoder CUDA graphs +## Enable encoder and decoder CUDA graphs -Pass `CudaGraphConfig` to capture and replay decoder iterations: +Configure the decoder and encoder graph grids separately: ```python -from tensorrt_llm.llmapi import CudaGraphConfig +from tensorrt_llm.llmapi import CudaGraphConfig, EncodeCudaGraphConfig llm = LLM( @@ -332,10 +332,17 @@ llm = LLM( backend="pytorch", attn_backend="TRTLLM", max_batch_size=8, + encoder_max_batch_size=8, cuda_graph_config=CudaGraphConfig( max_batch_size=8, enable_padding=True, ), + encoder_cuda_graph_config=EncodeCudaGraphConfig( + batch_sizes=[1, 2, 4, 8], + num_tokens=[128, 256, 512, 1024, 2048, 4096], + seq_lens=[128, 256, 512, 1024], + enable_padding=True, + ), kv_cache_config=KvCacheConfig( free_gpu_memory_fraction=0.8, cross_kv_cache_fraction=0.5, @@ -343,14 +350,17 @@ llm = LLM( ) ``` -This configuration captures decoder work only; the encoder continues to run -eagerly. With beam search, graph batch sizes must cover the active decoder -sequences after beam expansion. Padding lets nearby runtime batch sizes reuse a -captured graph. +`cuda_graph_config` controls decoder and mixed decoder graphs. +`encoder_cuda_graph_config` controls encoder-forward graph buckets for batch +size, total packed tokens, and maximum sequence length. The +`encoder_max_batch_size` value is the hard encoder capacity and admission +limit. With beam search, decoder graph batch sizes must cover the active +decoder sequences after beam expansion. -Do not use `EncodeCudaGraphConfig` for an encoder-decoder model. The runtime -warns and disables it. Piecewise CUDA graphs through `TorchCompileConfig` are -also unsupported for this model type. +Passing `EncodeCudaGraphConfig` through `cuda_graph_config` remains unsupported +for encoder-decoder models; pass it through `encoder_cuda_graph_config` +instead. Piecewise CUDA graphs through `TorchCompileConfig` are also +unsupported for this model type. ## Control the overlap scheduler @@ -566,11 +576,12 @@ Check all of the following: - `pipeline_parallel_size=1` and `context_parallel_size=1`. - `enable_attention_dp=False`. -### CUDA graphs do not capture the encoder +### Encoder CUDA graphs fall back to eager execution -This is expected. `CudaGraphConfig` accelerates decoder iterations only. The -encoder path runs eagerly, and `EncodeCudaGraphConfig` is disabled for -encoder-decoder models. +Check that `encoder_cuda_graph_config` and `encoder_max_batch_size` are set, +that the encoder graph buckets cover the request shape, and that +`attn_backend="TRTLLM"`. Unsupported shapes and attention backends fall back to +eager encoder execution. ### Output quality differs from the Hugging Face example diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index d73ba56d4ade..01590f805d2a 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -3,10 +3,9 @@ import bisect import contextlib -from collections import OrderedDict -from dataclasses import dataclass, field -from typing import (Any, Callable, Dict, Iterator, List, Optional, Tuple, - TypeAlias) +from dataclasses import dataclass +from typing import (Any, Callable, Dict, Iterator, List, Optional, Sequence, + Tuple, TypeAlias) import torch @@ -305,6 +304,23 @@ def get_graph_key( is_all_greedy_sample, context_query_lens, encoder_input_lens) return key + def _get_compatible_mixed_encoder_decoder_key(self, + key: KeyType) -> KeyType: + """Round the packed encoder extent up to a captured graph key.""" + if (not self.padding_enabled or self._capture_allowed + or key in self.graph_metadata or len(key[6]) != 1): + return key + + num_encoder_tokens = key[6][0] + compatible_keys = [ + captured_key for captured_key in self.graph_outputs + if captured_key[:6] == key[:6] and len(captured_key[6]) == 1 + and captured_key[6][0] >= num_encoder_tokens + ] + if not compatible_keys: + return key + return min(compatible_keys, key=lambda captured_key: captured_key[6][0]) + @staticmethod def _get_mrope_position_delta(request: Any) -> Optional[Any]: mrope_position_delta = getattr(request, "py_mrope_position_delta", None) @@ -386,6 +402,8 @@ def maybe_get_cuda_graph( return None, None, None key = self.get_graph_key(batch, new_tensors_device, spec_resource_manager, spec_metadata) + if is_mixed_encoder_decoder: + key = self._get_compatible_mixed_encoder_decoder_key(key) if key in self.graph_metadata: return self.graph_metadata[key][ @@ -499,7 +517,15 @@ def capture(self, "requires encoder hidden states.") static_encoder_hidden_states = self.shared_static_tensors[ "encoder_hidden_states"][:num_encoder_tokens] - static_encoder_hidden_states.copy_(encoder_hidden_states) + actual_num_encoder_tokens = encoder_hidden_states.shape[0] + if actual_num_encoder_tokens > num_encoder_tokens: + raise RuntimeError( + "Mixed encoder-decoder CUDA graph capture received " + f"{actual_num_encoder_tokens} encoder tokens for a " + f"{num_encoder_tokens}-token graph.") + static_encoder_hidden_states[:actual_num_encoder_tokens].copy_( + encoder_hidden_states) + static_encoder_hidden_states[actual_num_encoder_tokens:].zero_() capture_inputs[ "encoder_hidden_states"] = static_encoder_hidden_states attn_metadata = capture_inputs["attn_metadata"] @@ -588,8 +614,17 @@ def replay(self, key: KeyType, if encoder_hidden_states is None: raise RuntimeError("Mixed encoder-decoder CUDA graph replay " "requires encoder hidden states.") - static_tensors["encoder_hidden_states"][:num_encoder_tokens].copy_( + actual_num_encoder_tokens = encoder_hidden_states.shape[0] + if actual_num_encoder_tokens > num_encoder_tokens: + raise RuntimeError( + "Mixed encoder-decoder CUDA graph replay received " + f"{actual_num_encoder_tokens} encoder tokens for a " + f"{num_encoder_tokens}-token graph.") + static_encoder_hidden_states = static_tensors[ + "encoder_hidden_states"][:num_encoder_tokens] + static_encoder_hidden_states[:actual_num_encoder_tokens].copy_( encoder_hidden_states) + static_encoder_hidden_states[actual_num_encoder_tokens:].zero_() self.graphs[key].replay() output_ref = self.graph_outputs[key] @@ -810,10 +845,7 @@ class EncoderCUDAGraphRunnerConfig: max_num_tokens: int max_seq_len: int cuda_graph_mem_pool: Any - dynamic_sequence_layout: bool = False - allow_runtime_capture: bool = False - max_cuda_graphs: int = 0 - capture_keys: List[EncoderKeyType] = field(default_factory=list) + encoder_decoder_capture_keys: Optional[Sequence[EncoderKeyType]] = None class EncoderCUDAGraphRunner: @@ -821,8 +853,8 @@ class EncoderCUDAGraphRunner: Designed for encoder inputs with `input_ids` (flat [total_tokens]) and `seq_lens` ([batch_size]). Encoder CUDA graphs are keyed on the 3-tuple - (batch_size, total_tokens, max_seq_len_bucket) for dynamic encoder-decoder - batches. + (padded_batch_size, padded_total_tokens, max_seq_len_bucket) for dynamic + encoder-decoder batches when padding is enabled. Restricted to `TrtllmAttentionMetadata`: FlashInfer's per-batch planner state is not compatible with CUDA graph capture/replay. @@ -840,18 +872,13 @@ def __init__(self, config: EncoderCUDAGraphRunnerConfig): self.supported_num_tokens = sorted(config.cuda_graph_num_tokens) self.max_supported_num_tokens = config.max_cuda_graph_num_tokens self.supported_seq_lens = sorted(config.cuda_graph_seq_lens) - self.dynamic_sequence_layout = config.dynamic_sequence_layout - self.allow_runtime_capture = config.allow_runtime_capture - self.max_cuda_graphs = config.max_cuda_graphs - self.capture_keys = frozenset(config.capture_keys) - if (self.max_cuda_graphs > 0 - and len(self.capture_keys) > self.max_cuda_graphs): - raise ValueError("Encoder CUDA graph capture key count exceeds " - f"max_cuda_graphs: {len(self.capture_keys)} > " - f"{self.max_cuda_graphs}.") - - self.graphs: OrderedDict[EncoderKeyType, - torch.cuda.CUDAGraph] = OrderedDict() + self.is_encoder_decoder = config.encoder_decoder_capture_keys is not None + self.capture_keys = frozenset(config.encoder_decoder_capture_keys or ()) + self._capture_keys_by_batch_size: Dict[int, List[EncoderKeyType]] = {} + for key in sorted(self.capture_keys): + self._capture_keys_by_batch_size.setdefault(key[0], []).append(key) + + self.graphs: Dict[EncoderKeyType, torch.cuda.CUDAGraph] = {} self.graph_outputs: Dict[EncoderKeyType, Callable[[], Optional[Any]]] = {} self.graph_metadata: Dict[EncoderKeyType, Dict[str, Any]] = {} @@ -861,8 +888,8 @@ def __init__(self, config: EncoderCUDAGraphRunnerConfig): self.shared_static_tensors_cpu: Dict[str, torch.Tensor] = {} if self.enabled: self._create_shared_static_tensors() - self.cuda_graph_meta_buffers = ( - Buffers() if self.dynamic_sequence_layout else get_memory_buffers()) + self.cuda_graph_meta_buffers = (Buffers() if self.is_encoder_decoder + else get_memory_buffers()) self._capture_allowed = False self.is_warmup_only = False @@ -876,7 +903,7 @@ def __init__(self, config: EncoderCUDAGraphRunnerConfig): def _create_shared_static_tensors(self): """Allocates static tensors sized for the largest supported num_tokens.""" max_total_tokens = ( - self.config.max_num_tokens if self.dynamic_sequence_layout else min( + self.config.max_num_tokens if self.is_encoder_decoder else min( self.max_supported_num_tokens, self.config.max_num_tokens)) max_batch_size = self.max_supported_batch_size @@ -921,6 +948,76 @@ def _round_up(value: int, supported: List[int]) -> int: return 0 return supported[idx] + def _get_dynamic_capture_key( + self, + batch_size: int, + num_tokens: int, + max_seq_len: int, + allow_batch_padding: bool, + ) -> Optional[EncoderKeyType]: + """Return the smallest compatible dynamic-layout capture key.""" + candidate_batch_sizes = (self.supported_batch_sizes + if allow_batch_padding else [batch_size]) + for padded_batch_size in candidate_batch_sizes: + if padded_batch_size < batch_size: + continue + + required_num_tokens = num_tokens + padded_batch_size - batch_size + for key in self._capture_keys_by_batch_size.get( + padded_batch_size, []): + _, padded_num_tokens, padded_max_seq_len = key + if (padded_num_tokens < required_num_tokens + or padded_num_tokens > self.max_supported_num_tokens + or padded_max_seq_len < max_seq_len + or padded_max_seq_len not in self.supported_seq_lens + or padded_num_tokens + > padded_batch_size * padded_max_seq_len): + continue + return key + + return None + + def get_capture_warmup_sequence_lengths( + self, key: EncoderKeyType) -> Optional[List[int]]: + """Build a real sequence layout whose padded graph key is ``key``. + + A larger max-sequence bucket can be dominated by a smaller bucket at + the same batch size and token count. Such keys can never be selected + for a real batch and return ``None``. + """ + if key not in self.capture_keys: + return None + + batch_size, num_tokens, max_seq_len = key + previous_max_seq_len = max( + (candidate[2] for candidate in self._capture_keys_by_batch_size.get( + batch_size, []) + if candidate[1] == num_tokens and candidate[2] < max_seq_len), + default=0, + ) + actual_max_seq_len = max( + (num_tokens + batch_size - 1) // batch_size, + previous_max_seq_len + 1, + ) + if (actual_max_seq_len > max_seq_len + or actual_max_seq_len + batch_size - 1 > num_tokens): + return None + + if batch_size == 1: + sequence_lengths = [num_tokens] + else: + remaining_tokens = num_tokens - actual_max_seq_len + base, extra = divmod(remaining_tokens, batch_size - 1) + sequence_lengths = [actual_max_seq_len] + sequence_lengths.extend([base + 1] * extra) + sequence_lengths.extend([base] * (batch_size - 1 - extra)) + + selected_key, _, is_valid = self.get_graph_key( + {"seq_lens": sequence_lengths}) + if not is_valid or selected_key != key: + return None + return sequence_lengths + def _get_valid_graph_key(self, batch_size: int, num_tokens: int, max_seq_len: int) -> EncoderKeyType: num_tokens_idx = bisect.bisect_left(self.supported_num_tokens, @@ -957,7 +1054,20 @@ def get_graph_key( batch_size = len(seq_lens) max_seq_len = max(seq_lens) if batch_size > 0 else 0 - if self.dynamic_sequence_layout: + if self.is_encoder_decoder: + if self.padding_enabled and self.capture_keys: + padded_key = self._get_dynamic_capture_key( + batch_size, + num_tokens, + max_seq_len, + allow_batch_padding=False, + ) + if padded_key is None: + return (batch_size, 0, 0), False, False + is_padding_performed = (padded_key[1] != num_tokens + or padded_key[2] != max_seq_len) + return padded_key, is_padding_performed, True + max_seq_len_bucket = self._round_up(max_seq_len, self.supported_seq_lens) key: EncoderKeyType = (batch_size, num_tokens, max_seq_len_bucket) @@ -980,9 +1090,8 @@ def get_graph_key( def allow_capture(self): """Context manager that enables CUDA graph capture. - Static encode-only graphs capture during warmup through this context. - Dynamic encoder-decoder graphs may additionally opt into first-use - runtime capture through ``allow_runtime_capture``. + All encoder graphs are captured during explicit startup warmup through + this context. Unseen runtime keys fall back to eager execution. """ self._capture_allowed = True try: @@ -997,8 +1106,18 @@ def pad_batch(self, inputs: Dict[str, Any], yield inputs return - padded_batch_size = self._round_up(batch_size, - self.supported_batch_sizes) + if self.is_encoder_decoder and self.capture_keys: + seq_lens = inputs['seq_lens'] + padded_key = self._get_dynamic_capture_key( + batch_size, + sum(seq_lens), + max(seq_lens) if seq_lens else 0, + allow_batch_padding=True, + ) + padded_batch_size = padded_key[0] if padded_key is not None else 0 + else: + padded_batch_size = self._round_up(batch_size, + self.supported_batch_sizes) if padded_batch_size == 0 or padded_batch_size == batch_size: yield inputs return @@ -1058,8 +1177,7 @@ def maybe_get_cuda_graph( key, is_padding_performed, is_padding_successful = self.get_graph_key( inputs) - if (self.dynamic_sequence_layout and self.capture_keys - and key not in self.capture_keys): + if self.is_encoder_decoder and key not in self.capture_keys: return None, None padded_max_seq_len = key[2] if (not self.padding_enabled and is_padding_performed) \ @@ -1072,9 +1190,9 @@ def maybe_get_cuda_graph( self.retire_staging() return self.graph_metadata[key]["attn_metadata"], key - # New key not yet captured. Create graph metadata only during an - # explicit warmup capture or when first-use runtime capture is enabled. - if not (self._capture_allowed or self.allow_runtime_capture): + # New key not yet captured. Only create graph metadata during explicit + # startup warmup; unseen runtime keys fall back to eager execution. + if not self._capture_allowed: return None, None if "multi_item_part_lens" in inputs: @@ -1127,17 +1245,7 @@ def _contains_nested_tensor(self, x: Any) -> bool: return False def needs_capture(self, key: EncoderKeyType) -> bool: - return (self._capture_allowed - or self.allow_runtime_capture) and key not in self.graphs - - def _evict_graph_if_needed(self) -> None: - if self.max_cuda_graphs <= 0 or len(self.graphs) < self.max_cuda_graphs: - return - - key, graph = self.graphs.popitem(last=False) - graph.reset() - self.graph_outputs.pop(key, None) - self.graph_metadata.pop(key, None) + return self._capture_allowed and key not in self.graphs def _stage_inputs(self, key: EncoderKeyType, inputs: Dict[str, Any]) -> None: @@ -1192,7 +1300,6 @@ def capture( ) -> Any: """Warm up and/or capture the forward pass for a graph key.""" padded_num_tokens = key[1] - self._evict_graph_if_needed() sliced_static_tensors = { "input_ids": @@ -1293,7 +1400,6 @@ def replay( stored_meta["attn_metadata"]._seq_lens, non_blocking=True) self.graphs[key].replay() - self.graphs.move_to_end(key) self._staging_retirement_event = torch.cuda.Event() self._staging_retirement_event.record(torch.cuda.current_stream()) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 6e61647d0c1a..e5296cb1b53c 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -257,34 +257,15 @@ def _filter_cuda_graph_seq_lens(cuda_graph_seq_lens: list[int], return result -def _parse_encoder_decoder_cuda_graph_keys( - raw_keys: str) -> List[Tuple[int, int, int]]: - """Parse semicolon-separated ``batch,total_tokens,max_seq_bucket`` keys.""" - if not raw_keys.strip(): - return [] - - keys = [] - for raw_key in raw_keys.split(";"): - fields = raw_key.split(",") - if len(fields) != 3: - raise ValueError( - "TLLM_ENCODER_DECODER_MICROBATCH_CUDA_GRAPH_KEYS entries " - "must use batch_size,total_tokens,max_seq_bucket.") - try: - batch_size, total_tokens, max_seq_bucket = (int(field) - for field in fields) - except ValueError as error: - raise ValueError( - "TLLM_ENCODER_DECODER_MICROBATCH_CUDA_GRAPH_KEYS entries " - "must contain integers.") from error - key = (batch_size, total_tokens, max_seq_bucket) - if any(value <= 0 for value in key): - raise ValueError( - "TLLM_ENCODER_DECODER_MICROBATCH_CUDA_GRAPH_KEYS entries " - "must contain positive integers.") - keys.append(key) - - return sorted(set(keys)) +def _build_encoder_decoder_cuda_graph_keys( + batch_sizes: Sequence[int], num_tokens: Sequence[int], + seq_lens: Sequence[int]) -> List[Tuple[int, int, int]]: + """Build all geometrically valid encoder graph bucket combinations.""" + return sorted({(batch_size, total_tokens, max_seq_len) + for batch_size in batch_sizes + for total_tokens in num_tokens + for max_seq_len in seq_lens + if batch_size <= total_tokens <= batch_size * max_seq_len}) _DEEP_GEMM_PDL_CONFIGURED = False @@ -489,13 +470,21 @@ def __init__( self.cuda_graph_config = self.llm_args.cuda_graph_config self._is_encode_only = (self.llm_args.encode_only and not self.llm_args.mm_encoder_only) + if (self._is_encode_only + and isinstance(self.cuda_graph_config, EncodeCudaGraphConfig)): + self.encoder_cuda_graph_config = self.cuda_graph_config + else: + self.encoder_cuda_graph_config = ( + self.llm_args.encoder_cuda_graph_config) if (isinstance(self.cuda_graph_config, EncodeCudaGraphConfig) and self._is_encoder_decoder_model()): logger.warning( "EncodeCudaGraphConfig is not supported for encoder-decoder " - "models. Use DecodeCudaGraphConfig or CudaGraphConfig for " - "decoder CUDA graphs. CUDA graphs will be disabled.") + "models through cuda_graph_config. Use DecodeCudaGraphConfig " + "for cuda_graph_config and configure encoder graphs through " + "encoder_cuda_graph_config. Decoder CUDA graphs will be " + "disabled.") self.cuda_graph_config = None cuda_graph_batch_sizes = self.cuda_graph_config.batch_sizes if self.cuda_graph_config else CudaGraphConfig.model_fields[ @@ -503,25 +492,31 @@ def __init__( cuda_graph_padding_enabled = self.cuda_graph_config.enable_padding if self.cuda_graph_config else CudaGraphConfig.model_fields[ 'enable_padding'].default - # Encode-only CUDA graph detection. Decode configs do not define these - # encoder-specific bucket fields. - cuda_graph_num_tokens = [] - cuda_graph_seq_lens = [] - if isinstance(self.cuda_graph_config, EncodeCudaGraphConfig): - cuda_graph_num_tokens = self.cuda_graph_config.num_tokens or [] - cuda_graph_seq_lens = self.cuda_graph_config.seq_lens or [] - - if (self._is_encode_only and self.cuda_graph_config is not None - and (not cuda_graph_num_tokens or not cuda_graph_seq_lens)): + encoder_cuda_graph_batch_sizes = ( + self.encoder_cuda_graph_config.batch_sizes + if self.encoder_cuda_graph_config is not None else []) + encoder_cuda_graph_num_tokens = ( + self.encoder_cuda_graph_config.num_tokens + if self.encoder_cuda_graph_config is not None else []) + encoder_cuda_graph_seq_lens = (self.encoder_cuda_graph_config.seq_lens + if self.encoder_cuda_graph_config + is not None else []) + encoder_cuda_graph_padding_enabled = ( + self.encoder_cuda_graph_config.enable_padding + if self.encoder_cuda_graph_config is not None else False) + + if (self.encoder_cuda_graph_config is not None + and (not encoder_cuda_graph_num_tokens + or not encoder_cuda_graph_seq_lens)): missing = [] - if not cuda_graph_num_tokens: + if not encoder_cuda_graph_num_tokens: missing.append("num_tokens/max_num_token") - if not cuda_graph_seq_lens: + if not encoder_cuda_graph_seq_lens: missing.append("seq_lens/max_seq_len") logger.warning( - f"encode_only=True with a CudaGraphConfig, but " - f"{' and '.join(missing)} not set. Encoder CUDA graphs " - f"require both. Encoder CUDA graphs will be disabled. " + f"Encoder CUDA graph configuration has " + f"{' and '.join(missing)} unset. Encoder CUDA graphs require " + f"both dimensions and will be disabled. " f"To enable them, specify e.g. " f"EncodeCudaGraphConfig(max_batch_size=64, num_tokens=[128, 256, " f"512], max_seq_len=128, enable_padding=True).") @@ -688,46 +683,53 @@ def __init__( self._max_cuda_graph_batch_size = (self._cuda_graph_batch_sizes[-1] if self._cuda_graph_batch_sizes else 0) + self._encoder_cuda_graph_padding_enabled = ( + encoder_cuda_graph_padding_enabled) + self._encoder_cuda_graph_batch_sizes = (_filter_cuda_graph_batch_sizes( + encoder_cuda_graph_batch_sizes, self.encoder_batch_size, + self.encoder_max_num_tokens, 0, + self._encoder_cuda_graph_padding_enabled) if + encoder_cuda_graph_batch_sizes + else []) + # Encoder CUDA graph bucket lists - self._cuda_graph_num_tokens = _filter_cuda_graph_num_tokens( - cuda_graph_num_tokens, self.max_num_tokens, - self._cuda_graph_padding_enabled) if cuda_graph_num_tokens else [] + self._cuda_graph_num_tokens = (_filter_cuda_graph_num_tokens( + encoder_cuda_graph_num_tokens, self.encoder_max_num_tokens, + self._encoder_cuda_graph_padding_enabled) + if encoder_cuda_graph_num_tokens else []) self._max_cuda_graph_num_tokens = (self._cuda_graph_num_tokens[-1] if self._cuda_graph_num_tokens else 0) - self._cuda_graph_seq_lens = _filter_cuda_graph_seq_lens( - cuda_graph_seq_lens, self.max_seq_len, - self._cuda_graph_padding_enabled) if cuda_graph_seq_lens else [] + self._cuda_graph_seq_lens = (_filter_cuda_graph_seq_lens( + encoder_cuda_graph_seq_lens, self.max_seq_len, + self._encoder_cuda_graph_padding_enabled) + if encoder_cuda_graph_seq_lens else []) self._max_cuda_graph_seq_len = (self._cuda_graph_seq_lens[-1] if self._cuda_graph_seq_lens else 0) - encoder_decoder_graph_max_batch_size = int( - os.environ.get( - "TLLM_ENCODER_DECODER_MICROBATCH_CUDA_GRAPH_MAX_BATCH_SIZE", - "0")) + encoder_max_batch_size = self.llm_args.encoder_max_batch_size encoder_decoder_microbatch_cuda_graph_enabled = (os.environ.get( "TLLM_ENCODER_DECODER_MICROBATCH_CUDA_GRAPH_ENABLED", "1") == "1") self._enable_encoder_decoder_microbatch_cuda_graph = ( self._is_encoder_decoder_model() - and encoder_decoder_graph_max_batch_size > 0 + and encoder_max_batch_size is not None + and self.encoder_cuda_graph_config is not None + and bool(self._cuda_graph_num_tokens) + and bool(self._cuda_graph_seq_lens) and encoder_decoder_microbatch_cuda_graph_enabled) - encoder_decoder_graph_batch_sizes = [ - batch_size for batch_size in (1, 2, 4, 8) if batch_size <= min( - encoder_decoder_graph_max_batch_size, self.batch_size) - ] - encoder_decoder_graph_keys = _parse_encoder_decoder_cuda_graph_keys( - os.environ.get("TLLM_ENCODER_DECODER_MICROBATCH_CUDA_GRAPH_KEYS", - "")) - encoder_decoder_graph_seq_lens = list( - range(64, self.max_seq_len + 1, 64)) - if (encoder_decoder_graph_seq_lens - and encoder_decoder_graph_seq_lens[-1] < self.max_seq_len): - encoder_decoder_graph_seq_lens.append(self.max_seq_len) + encoder_decoder_graph_batch_sizes = ( + self._encoder_cuda_graph_batch_sizes + if self._enable_encoder_decoder_microbatch_cuda_graph else []) + encoder_decoder_graph_keys = _build_encoder_decoder_cuda_graph_keys( + encoder_decoder_graph_batch_sizes, + self._cuda_graph_num_tokens, + self._cuda_graph_seq_lens, + ) encoder_decoder_graph_keys = [ key for key in encoder_decoder_graph_keys if key[0] in encoder_decoder_graph_batch_sizes and key[1] <= - self.max_num_tokens and key[2] in encoder_decoder_graph_seq_lens + self.encoder_max_num_tokens and key[2] in self._cuda_graph_seq_lens ] mixed_graph_encoder_batch_size = (encoder_decoder_graph_batch_sizes[-1] if encoder_decoder_graph_batch_sizes @@ -859,40 +861,28 @@ def __init__( self.cuda_graph_runner = CUDAGraphRunner(cuda_graph_runner_config) # Create Encoder CUDA graph config and runner. - encoder_graph_batch_sizes = (encoder_decoder_graph_batch_sizes - if use_encoder_decoder_graph else - self._cuda_graph_batch_sizes) + encoder_graph_batch_sizes = self._encoder_cuda_graph_batch_sizes encoder_graph_max_batch_size = (encoder_graph_batch_sizes[-1] if encoder_graph_batch_sizes else 0) - encoder_graph_max_num_tokens = (self.max_num_tokens - if use_encoder_decoder_graph else - self._max_cuda_graph_num_tokens) + encoder_graph_max_num_tokens = self._max_cuda_graph_num_tokens encoder_cuda_graph_runner_config = EncoderCUDAGraphRunnerConfig( use_cuda_graph=(use_encoder_decoder_graph or (self._is_encode_only - and self.cuda_graph_config is not None + and self.encoder_cuda_graph_config is not None and bool(self._cuda_graph_num_tokens) and bool(self._cuda_graph_seq_lens))), - cuda_graph_padding_enabled=(False if use_encoder_decoder_graph else - self._cuda_graph_padding_enabled), + cuda_graph_padding_enabled=( + self._encoder_cuda_graph_padding_enabled), cuda_graph_batch_sizes=encoder_graph_batch_sizes, cuda_graph_num_tokens=self._cuda_graph_num_tokens, - cuda_graph_seq_lens=(encoder_decoder_graph_seq_lens - if use_encoder_decoder_graph else - self._cuda_graph_seq_lens), + cuda_graph_seq_lens=self._cuda_graph_seq_lens, max_cuda_graph_batch_size=encoder_graph_max_batch_size, max_cuda_graph_num_tokens=encoder_graph_max_num_tokens, - max_num_tokens=self.max_num_tokens, + max_num_tokens=self.encoder_max_num_tokens, max_seq_len=self.max_seq_len, cuda_graph_mem_pool=self._cuda_graph_mem_pool, - dynamic_sequence_layout=use_encoder_decoder_graph, - allow_runtime_capture=use_encoder_decoder_graph, - max_cuda_graphs=(int( - os.environ.get( - "TLLM_ENCODER_DECODER_MICROBATCH_CUDA_GRAPH_MAX_GRAPHS", - "64")) if use_encoder_decoder_graph else 0), - capture_keys=(encoder_decoder_graph_keys - if use_encoder_decoder_graph else []), + encoder_decoder_capture_keys=(encoder_decoder_graph_keys if + use_encoder_decoder_graph else None), ) self.encoder_cuda_graph_runner = EncoderCUDAGraphRunner( encoder_cuda_graph_runner_config) @@ -1874,6 +1864,62 @@ def _run_cuda_graph_warmup(self, resource_manager: ResourceManager): if not self.cuda_graph_runner.is_warmup_only: self._capture_piecewise_cuda_graphs(resource_manager) + @with_warmup_flag + def _warmup_encoder_decoder_encoder_cuda_graphs( + self, resource_manager: ResourceManager) -> None: + """Capture encoder-decoder encoder graphs on their runtime host thread.""" + runner = self.encoder_cuda_graph_runner + if not runner.enabled or not runner.is_encoder_decoder: + return + + with runner.allow_capture(): + runner.is_warmup_only = True + try: + self._capture_encoder_decoder_encoder_cuda_graphs( + resource_manager) + finally: + runner.is_warmup_only = False + self._capture_encoder_decoder_encoder_cuda_graphs(resource_manager) + + def _capture_encoder_decoder_encoder_cuda_graphs( + self, resource_manager: ResourceManager) -> None: + """Warm up or capture encoder graphs used by encoder-decoder models.""" + runner = self.encoder_cuda_graph_runner + if not runner.enabled or not runner.is_encoder_decoder: + return + + operation = "warmup" if runner.is_warmup_only else "capture" + num_processed = 0 + logger.info( + f"Running encoder-decoder encoder CUDA graph {operation} ...") + for key in sorted(runner.capture_keys, reverse=True): + sequence_lengths = runner.get_capture_warmup_sequence_lengths(key) + if sequence_lengths is None: + continue + + encoder_input_ids = [0] * sum(sequence_lengths) + encoder_position_ids = [] + for sequence_length in sequence_lengths: + encoder_position_ids.extend( + self._apply_position_id_offset(list( + range(sequence_length)))) + inputs = self._prepare_encoder_decoder_encoder_inputs( + encoder_input_ids=encoder_input_ids, + encoder_position_ids=encoder_position_ids, + sequence_lengths=sequence_lengths, + request_ids=list(range(len(sequence_lengths))), + resource_manager=resource_manager, + ) + + logger.info("Encoder-decoder encoder CUDA graph " + f"{operation}: key={key}") + self._forward_encoder_with_cuda_graph(inputs) + torch.cuda.synchronize() + num_processed += 1 + + logger.info("Completed encoder-decoder encoder CUDA graph " + f"{operation} for {num_processed} graph shape(s).") + def _capture_generation_cuda_graphs(self, resource_manager: ResourceManager): """Warm up or capture pure-generation CUDA graph shapes.""" @@ -6410,7 +6456,7 @@ def _capture_encoder_cuda_graphs(self) -> None: if not runner.enabled: return - batch_sizes = sorted(self._cuda_graph_batch_sizes, reverse=True) + batch_sizes = sorted(self._encoder_cuda_graph_batch_sizes, reverse=True) num_tokens_list = sorted(self._cuda_graph_num_tokens) seq_lens_list = sorted(self._cuda_graph_seq_lens) @@ -6920,9 +6966,9 @@ def _prepare_encoder_decoder_encoder_inputs( if num_tokens != len(encoder_position_ids): raise ValueError("Encoder input IDs and position IDs must have " "the same length.") - assert num_tokens <= self.max_num_tokens, ( - f"encoder packed length ({num_tokens}) exceeds max_num_tokens " - f"({self.max_num_tokens})") + assert num_tokens <= self.encoder_max_num_tokens, ( + f"encoder packed length ({num_tokens}) exceeds " + f"encoder_max_num_tokens ({self.encoder_max_num_tokens})") # Build a fresh, no-cache attention metadata for the encoder # pass. We do not reuse ``self.attn_metadata`` because that @@ -6932,9 +6978,9 @@ def _prepare_encoder_decoder_encoder_inputs( pretrained_config=self.model.model_config.pretrained_config) if self.sparse_attention_config is not None else None) encoder_attn_metadata = self.attn_backend.Metadata( - max_num_requests=self.batch_size, - max_num_tokens=self.max_num_tokens, - max_num_sequences=self.batch_size * self.max_beam_width, + max_num_requests=self.encoder_batch_size, + max_num_tokens=self.encoder_max_num_tokens, + max_num_sequences=self.encoder_batch_size * self.max_beam_width, kv_cache_manager=None, mapping=self.mapping, runtime_features=self.attn_runtime_features, @@ -6965,9 +7011,13 @@ def _prepare_encoder_decoder_encoder_inputs( encoder_position_ids_t = torch.tensor(encoder_position_ids, dtype=torch.int, pin_memory=prefer_pinned()) + encoder_graph_runner = self.encoder_cuda_graph_runner + encoder_batch_size = len(sequence_lengths) use_graph_staging = ( - self.encoder_cuda_graph_runner.enabled and len(sequence_lengths) - in self.encoder_cuda_graph_runner.supported_batch_sizes) + encoder_graph_runner.enabled and + (encoder_batch_size in encoder_graph_runner.supported_batch_sizes or + (encoder_graph_runner.padding_enabled and encoder_batch_size + <= encoder_graph_runner.max_supported_batch_size))) return { 'encoder_input_ids': @@ -7128,48 +7178,50 @@ def _forward_encoder_with_cuda_graph( 'seq_lens': seq_lens, 'resource_manager': inputs.get('resource_manager'), } - graph_attn_metadata, key = runner.maybe_get_cuda_graph( - runner_inputs, inputs['encoder_attn_metadata']) - if key is None: - if inputs['encoder_input_ids'].device.type == 'cpu': - inputs = dict(inputs) - inputs['encoder_input_ids'] = inputs['encoder_input_ids'].to( - 'cuda', non_blocking=True) - inputs['encoder_position_ids'] = inputs[ - 'encoder_position_ids'].to('cuda', non_blocking=True) - return self._forward_step_encoder(inputs) - - # Every graph key aliases the same pinned staging allocation. Retire - # the previous captured H2D before updating seq_lens or any other - # shared host input for this replay. - runner.retire_staging() - graph_attn_metadata.prepare_encoder_cuda_graph_replay( - seq_lens, actual_num_tokens) - model_inputs = { - **runner_inputs, - 'attn_metadata': graph_attn_metadata, - } + with runner.pad_batch(runner_inputs, + len(seq_lens)) as padded_runner_inputs: + graph_attn_metadata, key = runner.maybe_get_cuda_graph( + padded_runner_inputs, inputs['encoder_attn_metadata']) + if key is None: + if inputs['encoder_input_ids'].device.type == 'cpu': + inputs = dict(inputs) + inputs['encoder_input_ids'] = inputs[ + 'encoder_input_ids'].to('cuda', non_blocking=True) + inputs['encoder_position_ids'] = inputs[ + 'encoder_position_ids'].to('cuda', non_blocking=True) + return self._forward_step_encoder(inputs) + + # Every graph key aliases the same pinned staging allocation. Retire + # the previous captured H2D before updating seq_lens or any other + # shared host input for this replay. + runner.retire_staging() + graph_attn_metadata.prepare_encoder_cuda_graph_replay( + padded_runner_inputs['seq_lens'], key[1]) + model_inputs = { + **padded_runner_inputs, + 'attn_metadata': graph_attn_metadata, + } - moe_load_balancer: MoeLoadBalancer = getattr(self, 'moe_load_balancer', - None) - with with_shared_pool(runner.get_graph_pool()): - capture_outputs = None - if runner.needs_capture(key): + moe_load_balancer: MoeLoadBalancer = getattr( + self, 'moe_load_balancer', None) + with with_shared_pool(runner.get_graph_pool()): + capture_outputs = None + if runner.needs_capture(key): - def capture_forward_fn( - capture_inputs: Dict[str, Any]) -> torch.Tensor: - with MoeLoadBalancerIterContext(moe_load_balancer): - return self._forward_step_encoder_cuda_graph( - capture_inputs) + def capture_forward_fn( + capture_inputs: Dict[str, Any]) -> torch.Tensor: + with MoeLoadBalancerIterContext(moe_load_balancer): + return self._forward_step_encoder_cuda_graph( + capture_inputs) - capture_outputs = runner.capture(key, capture_forward_fn, - model_inputs) + capture_outputs = runner.capture(key, capture_forward_fn, + model_inputs) - if runner.is_warmup_only: - graph_outputs = capture_outputs - else: - with MoeLoadBalancerIterContext(moe_load_balancer): - graph_outputs = runner.replay(key, model_inputs) + if runner.is_warmup_only: + graph_outputs = capture_outputs + else: + with MoeLoadBalancerIterContext(moe_load_balancer): + graph_outputs = runner.replay(key, model_inputs) if not isinstance(graph_outputs, torch.Tensor): raise TypeError("Encoder-decoder CUDA graph replay must return " diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index ea44ec7b885f..704f59f2e8e4 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -831,6 +831,12 @@ def __init__( max_workers=1, thread_name_prefix="encoder-launch") if self.is_encoder_decoder else None) self.pending_encoder_steps: List[PendingEncoderStep] = [] + if self.encoder_launch_executor is not None: + # CUDA graph capture inherits per-thread CUDA library state. Capture + # on the same single worker that owns every runtime encoder replay. + self.encoder_stream.wait_stream(self.execution_stream) + self.encoder_launch_executor.submit( + self._warmup_encoder_decoder_encoder_cuda_graphs).result() self.is_warmup = False @@ -5276,42 +5282,54 @@ def _waiting_encoder_requests( self.encoder_batch_wait_iters_count = 0 return encoder_requests - microbatch_graph_max_batch_size = int( - os.environ.get( - "TLLM_ENCODER_DECODER_MICROBATCH_CUDA_GRAPH_MAX_BATCH_SIZE", - "0")) - microbatch_admission_enabled = (os.environ.get( - "TLLM_ENCODER_DECODER_MICROBATCH_ADMISSION_ENABLED", "1") == "1") - if (microbatch_graph_max_batch_size > 0 - and microbatch_admission_enabled): - microbatch_target = min(microbatch_graph_max_batch_size, - self.max_batch_size) - decoder_occupancy = (len(context_requests) + - len(generation_requests)) - decoder_low_watermark = int( - os.environ.get("TLLM_ENCODER_DECODER_MICROBATCH_LOW_WATERMARK", - str(self.max_batch_size - microbatch_target))) - deadline_reached = (self.encoder_batch_wait_iters_count - >= self.batch_wait_timeout_iters) - - if (len(encoder_requests) >= microbatch_target - and decoder_occupancy <= decoder_low_watermark): - self.encoder_batch_wait_iters_count = 0 - return encoder_requests[:microbatch_target] - - if deadline_reached: - supported_batch_sizes = [ - batch_size for batch_size in (1, 2, 4, 8) - if batch_size <= microbatch_target - ] - fallback_batch_size = max( - batch_size for batch_size in supported_batch_sizes - if batch_size <= len(encoder_requests)) - self.encoder_batch_wait_iters_count = 0 - return encoder_requests[:fallback_batch_size] - - self.encoder_batch_wait_iters_count += 1 - return [] + encoder_max_batch_size = self.llm_args.encoder_max_batch_size + encoder_cuda_graph_config = self.llm_args.encoder_cuda_graph_config + if (encoder_max_batch_size is not None + and encoder_cuda_graph_config is not None + and bool(encoder_cuda_graph_config.num_tokens) + and bool(encoder_cuda_graph_config.seq_lens)): + encoder_batch_size_limit = min(encoder_max_batch_size, + self.max_batch_size) + configured_batch_sizes = (encoder_cuda_graph_config.batch_sizes + or []) + supported_batch_sizes = [ + batch_size for batch_size in configured_batch_sizes + if batch_size <= encoder_batch_size_limit + ] + if (encoder_cuda_graph_config.enable_padding + and any(batch_size > encoder_batch_size_limit + for batch_size in configured_batch_sizes) and + (not supported_batch_sizes + or supported_batch_sizes[-1] != encoder_batch_size_limit)): + supported_batch_sizes.append(encoder_batch_size_limit) + + if supported_batch_sizes: + microbatch_target = supported_batch_sizes[-1] + decoder_occupancy = (len(context_requests) + + len(generation_requests)) + decoder_low_watermark = (self.max_batch_size - + microbatch_target) + deadline_reached = (self.encoder_batch_wait_iters_count + >= self.batch_wait_timeout_iters) + + if decoder_occupancy <= decoder_low_watermark: + if len(encoder_requests) >= microbatch_target: + self.encoder_batch_wait_iters_count = 0 + return encoder_requests[:microbatch_target] + + if deadline_reached: + releasable_batch_sizes = [ + batch_size for batch_size in supported_batch_sizes + if batch_size <= len(encoder_requests) + ] + fallback_batch_size = ( + releasable_batch_sizes[-1] if releasable_batch_sizes + else min(len(encoder_requests), microbatch_target)) + self.encoder_batch_wait_iters_count = 0 + return encoder_requests[:fallback_batch_size] + + self.encoder_batch_wait_iters_count += 1 + return [] num_scheduled_tokens = sum(request.encoder_output_len for request in encoder_requests) @@ -5404,6 +5422,20 @@ def _schedule(self): # micro-batch; this preserves the cross-KV lifecycle and the # dual-pool budget. # --------------------------------------------------------------- + def _warmup_encoder_decoder_encoder_cuda_graphs(self) -> None: + """Capture encoder graphs on the worker used for runtime replay.""" + warmup = getattr( + self.model_engine, + "_warmup_encoder_decoder_encoder_cuda_graphs", + None, + ) + if not callable(warmup): + return + + torch.cuda.set_device(self.device_id) + with torch.cuda.stream(self.encoder_stream): + warmup(self.resource_manager) + def _submit_encoder_step(self, encoder_requests: List[LlmRequest]) -> None: """Queue encoder work without blocking the decoder executor thread.""" executor = self.encoder_launch_executor diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index a11f79a6ec07..9e1fddc7c2c9 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -4697,6 +4697,15 @@ class TorchLlmArgs(BaseLlmArgs): Note that each CUDA graph can use up to 200 MB of extra memory.", status="beta") + encoder_cuda_graph_config: Optional[EncodeCudaGraphConfig] = Field( + default=None, + description=( + "CUDA graph configuration for the encoder forward pass of an " + "encoder-decoder model. Use `cuda_graph_config` for the decoder " + "and this field for the encoder. Encoder CUDA graphs require " + "`encoder_max_batch_size` to be set."), + status="prototype") + @field_validator('cuda_graph_config', mode='before') @classmethod def infer_cuda_graph_config_mode(cls, v): @@ -4750,21 +4759,21 @@ def init_multimodal_config(cls, v): encoder_max_batch_size: Optional[int] = Field( default=None, - description=( - "Maximum batch size for the multimodal encoder's AttentionMetadata. " - "Falls back to `max_batch_size` when unset. This budget is shared " - "proportionately across all modalities the model encodes, not set " - "per modality; per-modality knobs may be added later."), + description= + ("Maximum encoder batch size. For encoder-decoder models, this also " + "controls encoder microbatch admission and limits encoder CUDA graph " + "batch sizes. For multimodal models, this is the shared " + "AttentionMetadata budget across all encoded modalities. Falls back " + "to `max_batch_size` when unset."), status="prototype") encoder_max_num_tokens: Optional[int] = Field( default=None, description=( - "Maximum number of tokens for the multimodal encoder's " - "AttentionMetadata. Falls back to `max_num_tokens` when unset. This " - "budget is shared proportionately across all modalities the model " - "encodes, not set per modality; per-modality knobs may be added " - "later."), + "Maximum number of encoder tokens. For encoder-decoder models, this " + "limits encoder CUDA graph total-token buckets. For multimodal " + "models, this is the shared AttentionMetadata budget across all " + "encoded modalities. Falls back to `max_num_tokens` when unset."), status="prototype") @field_validator("encoder_max_batch_size", "encoder_max_num_tokens") @@ -4774,6 +4783,28 @@ def validate_encoder_runtime_sizes(cls, v: Optional[int]) -> Optional[int]: raise ValueError("must be a positive integer when set") return v + @model_validator(mode="after") + def validate_encoder_cuda_graph_config(self) -> 'TorchLlmArgs': + if self.encoder_cuda_graph_config is None: + return self + if self.encode_only: + raise ValueError( + "Use cuda_graph_config=EncodeCudaGraphConfig(...) when " + "encode_only=True; encoder_cuda_graph_config is for " + "encoder-decoder models.") + if self.encoder_max_batch_size is None: + raise ValueError( + "encoder_cuda_graph_config requires encoder_max_batch_size.") + missing = [] + if not self.encoder_cuda_graph_config.num_tokens: + missing.append("num_tokens/max_num_token") + if not self.encoder_cuda_graph_config.seq_lens: + missing.append("seq_lens/max_seq_len") + if missing: + raise ValueError("encoder_cuda_graph_config requires " + f"{' and '.join(missing)}.") + return self + attn_backend: str = Field( default='TRTLLM', description="Attention backend to use.", diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index 739db09bfb23..0cfa637cabe4 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -474,6 +474,64 @@ "kind": "value", "path": "encode_only" }, + { + "allowed_values": [], + "annotation": "Optional[List[int]]", + "converter": "", + "kind": "value", + "path": "encoder_cuda_graph_config.batch_sizes" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "encoder_cuda_graph_config.enable_padding" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "encoder_cuda_graph_config.max_batch_size" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "encoder_cuda_graph_config.max_num_token" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "encoder_cuda_graph_config.max_seq_len" + }, + { + "allowed_values": [ + "encode" + ], + "annotation": "Literal['encode']", + "converter": "", + "kind": "categorical", + "path": "encoder_cuda_graph_config.mode" + }, + { + "allowed_values": [], + "annotation": "Optional[List[Annotated[int, Gt(gt=0)]]]", + "converter": "", + "kind": "value", + "path": "encoder_cuda_graph_config.num_tokens" + }, + { + "allowed_values": [], + "annotation": "Optional[List[Annotated[int, Gt(gt=0)]]]", + "converter": "", + "kind": "value", + "path": "encoder_cuda_graph_config.seq_lens" + }, { "allowed_values": [], "annotation": "Optional[int]", diff --git a/tests/unittest/_torch/executor/test_encoder_cuda_graph_runner.py b/tests/unittest/_torch/executor/test_encoder_cuda_graph_runner.py index f01070f2b42d..9a4f399f3b41 100644 --- a/tests/unittest/_torch/executor/test_encoder_cuda_graph_runner.py +++ b/tests/unittest/_torch/executor/test_encoder_cuda_graph_runner.py @@ -11,15 +11,17 @@ EncoderCUDAGraphRunner, EncoderCUDAGraphRunnerConfig, ) +from tensorrt_llm._torch.pyexecutor.model_engine import _build_encoder_decoder_cuda_graph_keys def _dynamic_layout_runner( - max_cuda_graphs: int = 0, capture_keys: list[tuple[int, int, int]] | None = None + capture_keys: list[tuple[int, int, int]] | None = None, + enable_padding: bool = False, ) -> EncoderCUDAGraphRunner: return EncoderCUDAGraphRunner( EncoderCUDAGraphRunnerConfig( use_cuda_graph=False, - cuda_graph_padding_enabled=False, + cuda_graph_padding_enabled=enable_padding, cuda_graph_batch_sizes=[1, 2, 4, 8], cuda_graph_num_tokens=[], cuda_graph_seq_lens=list(range(64, 513, 64)), @@ -28,13 +30,86 @@ def _dynamic_layout_runner( max_num_tokens=4096, max_seq_len=512, cuda_graph_mem_pool=None, - dynamic_sequence_layout=True, - allow_runtime_capture=True, - max_cuda_graphs=max_cuda_graphs, - capture_keys=capture_keys or [], + encoder_decoder_capture_keys=capture_keys or [], + ) + ) + + +def test_encoder_decoder_capture_keys_select_layout_mode(): + encoder_decoder_runner = _dynamic_layout_runner() + encoder_only_runner = EncoderCUDAGraphRunner( + EncoderCUDAGraphRunnerConfig( + use_cuda_graph=False, + cuda_graph_padding_enabled=False, + cuda_graph_batch_sizes=[1], + cuda_graph_num_tokens=[8], + cuda_graph_seq_lens=[8], + max_cuda_graph_batch_size=1, + max_cuda_graph_num_tokens=8, + max_num_tokens=8, + max_seq_len=8, + cuda_graph_mem_pool=None, ) ) + assert encoder_decoder_runner.is_encoder_decoder + assert not encoder_only_runner.is_encoder_decoder + + +def test_build_encoder_decoder_cuda_graph_keys(): + keys = _build_encoder_decoder_cuda_graph_keys( + batch_sizes=[1, 2], + num_tokens=[96, 576, 1056], + seq_lens=[512], + ) + + assert keys == [ + (1, 96, 512), + (2, 96, 512), + (2, 576, 512), + ] + + +def test_bart_encoder_graph_config_builds_feasible_key_grid(): + num_tokens = list(range(96, 4801, 96)) + keys = _build_encoder_decoder_cuda_graph_keys( + batch_sizes=[1, 2, 4, 8], + num_tokens=num_tokens, + seq_lens=[512, 1024], + ) + + assert len(keys) == 201 + assert {total_tokens for batch_size, total_tokens, _ in keys if batch_size == 8} == set( + num_tokens + ) + + +def test_encoder_graph_builds_reachable_startup_warmup_layouts(): + capture_keys = _build_encoder_decoder_cuda_graph_keys( + batch_sizes=[1, 2], + num_tokens=[96, 320], + seq_lens=[256, 512], + ) + runner = _dynamic_layout_runner( + capture_keys=capture_keys, + enable_padding=True, + ) + + for key in capture_keys: + sequence_lengths = runner.get_capture_warmup_sequence_lengths(key) + if sequence_lengths is None: + continue + + selected_key, _, is_valid = runner.get_graph_key({"seq_lens": sequence_lengths}) + assert is_valid + assert selected_key == key + assert len(sequence_lengths) == key[0] + assert sum(sequence_lengths) == key[1] + + assert runner.get_capture_warmup_sequence_lengths((1, 96, 256)) == [96] + assert runner.get_capture_warmup_sequence_lengths((1, 96, 512)) is None + assert runner.get_capture_warmup_sequence_lengths((2, 320, 512)) == [257, 63] + def test_encoder_graph_key_reuses_total_tokens_and_max_bucket(): runner = _dynamic_layout_runner() @@ -78,28 +153,74 @@ def test_encoder_graph_key_distinguishes_max_buckets(): assert larger_bucket_key == (4, 1400, 448) -def test_bart_microbatch_key_set_fits_graph_cache(): - runner = _dynamic_layout_runner(max_cuda_graphs=64) - sequence_length_cycle = list(range(260, 441, 12)) - keys = set() +def test_encoder_graph_key_pads_tokens_and_max_sequence_length(): + runner = _dynamic_layout_runner( + capture_keys=[ + (2, 640, 320), + (2, 640, 384), + (2, 704, 384), + ], + enable_padding=True, + ) + + key, is_padding_performed, is_valid = runner.get_graph_key( + { + "input_ids": [0] * 556, + "seq_lens": [260, 296], + } + ) + + assert key == (2, 640, 320) + assert is_padding_performed + assert is_valid + - for batch_size in (1, 2, 4, 8): - for start in range(len(sequence_length_cycle)): - sequence_lengths = [ - sequence_length_cycle[(start + offset) % len(sequence_length_cycle)] - for offset in range(batch_size) - ] - key, _, is_valid = runner.get_graph_key( - { - "input_ids": [0] * sum(sequence_lengths), - "seq_lens": sequence_lengths, - } - ) - assert is_valid - keys.add(key) +def test_encoder_graph_pad_batch_selects_compatible_capture_key(): + runner = _dynamic_layout_runner( + capture_keys=[ + (4, 350, 192), + (4, 384, 192), + (8, 768, 192), + ], + enable_padding=True, + ) + runner.enabled = True + inputs = { + "input_ids": [0] * 350, + "seq_lens": [100, 120, 130], + } + + with runner.pad_batch(inputs, batch_size=3) as padded_inputs: + assert padded_inputs["seq_lens"] == [100, 120, 130, 1] + assert padded_inputs["input_ids"] is inputs["input_ids"] + key, is_padding_performed, is_valid = runner.get_graph_key(padded_inputs) + + assert key == (4, 384, 192) + assert is_padding_performed + assert is_valid - assert len(keys) == 59 - assert len(keys) <= runner.max_cuda_graphs + +def test_encoder_graph_padding_rejects_incompatible_capture_keys(): + runner = _dynamic_layout_runner( + capture_keys=[ + (4, 320, 128), + (8, 512, 128), + ], + enable_padding=True, + ) + runner.enabled = True + inputs = { + "input_ids": [0] * 350, + "seq_lens": [100, 120, 130], + } + + with runner.pad_batch(inputs, batch_size=3) as padded_inputs: + assert padded_inputs is inputs + key, is_padding_performed, is_valid = runner.get_graph_key(padded_inputs) + + assert key == (3, 0, 0) + assert not is_padding_performed + assert not is_valid def test_encoder_graph_key_rejects_oversized_inputs(): @@ -115,11 +236,14 @@ def test_encoder_graph_key_rejects_oversized_inputs(): assert not is_valid -def test_encoder_graph_capture_allowlist_must_fit_cache(): - capture_keys = [(1, num_tokens, 64) for num_tokens in range(1, 66)] +def test_encoder_graph_only_captures_during_warmup(): + key = (1, 8, 64) + runner = _dynamic_layout_runner(capture_keys=[key]) - with pytest.raises(ValueError, match="capture key count"): - _dynamic_layout_runner(max_cuda_graphs=64, capture_keys=capture_keys) + assert not runner.needs_capture(key) + with runner.allow_capture(): + assert runner.needs_capture(key) + assert not runner.needs_capture(key) def test_encoder_graph_reuses_same_key_for_different_sequence_layouts(): @@ -152,13 +276,42 @@ def test_encoder_graph_reuses_same_key_for_different_sequence_layouts(): assert reused_key == key +def test_encoder_graph_replay_uses_plain_graph_mapping(monkeypatch): + runner = _dynamic_layout_runner() + key = (1, 8, 64) + attn_metadata = object() + expected_output = object() + replay_calls = [] + recorded_streams = [] + current_stream = object() + + runner.graphs[key] = SimpleNamespace(replay=lambda: replay_calls.append(key)) + runner.graph_metadata[key] = {"attn_metadata": attn_metadata} + runner.graph_outputs[key] = expected_output + runner._capture_h2d_copy = True + monkeypatch.setattr(runner, "retire_staging", lambda: None) + monkeypatch.setattr(runner, "_stage_inputs", lambda _key, _inputs: None) + monkeypatch.setattr( + torch.cuda, + "Event", + lambda: SimpleNamespace(record=lambda stream: recorded_streams.append(stream)), + ) + monkeypatch.setattr(torch.cuda, "current_stream", lambda: current_stream) + + output = runner.replay(key, {"attn_metadata": attn_metadata}) + + assert output is expected_output + assert replay_calls == [key] + assert recorded_streams == [current_stream] + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") def test_encoder_graph_capture_stages_warmup_and_replays_new_layout(): runner = _dynamic_layout_runner() runner.enabled = True runner._create_shared_static_tensors() - key = (2, 5, 64) + key = (2, 8, 64) seq_lens_host = runner.shared_static_tensors_cpu["seq_lens"][:2] seq_lens_host.copy_(torch.tensor([2, 3], dtype=torch.int32)) attn_metadata = SimpleNamespace( @@ -186,7 +339,7 @@ def forward_fn(capture_inputs): torch.cuda.synchronize() torch.testing.assert_close( first_output, - torch.tensor([12, 13, 14, 15, 16], device="cuda", dtype=torch.int32), + torch.tensor([12, 13, 14, 15, 16, 2, 2, 2], device="cuda", dtype=torch.int32), ) seq_lens_host.copy_(torch.tensor([1, 4], dtype=torch.int32)) @@ -198,28 +351,5 @@ def forward_fn(capture_inputs): torch.cuda.synchronize() torch.testing.assert_close( reused_output, - torch.tensor([11, 12, 13, 14, 15], device="cuda", dtype=torch.int32), + torch.tensor([11, 12, 13, 14, 15, 1, 1, 1], device="cuda", dtype=torch.int32), ) - - -def test_encoder_graph_lru_evicts_oldest_graph(): - class _Graph: - def __init__(self): - self.was_reset = False - - def reset(self): - self.was_reset = True - - runner = _dynamic_layout_runner(max_cuda_graphs=1) - key = (1, 8, 64) - graph = _Graph() - runner.graphs[key] = graph - runner.graph_outputs[key] = object() - runner.graph_metadata[key] = object() - - runner._evict_graph_if_needed() - - assert graph.was_reset - assert not runner.graphs - assert not runner.graph_outputs - assert not runner.graph_metadata diff --git a/tests/unittest/_torch/executor/test_mixed_decoder_cuda_graph_runner.py b/tests/unittest/_torch/executor/test_mixed_decoder_cuda_graph_runner.py index f4412dd23aa7..8af9118313be 100644 --- a/tests/unittest/_torch/executor/test_mixed_decoder_cuda_graph_runner.py +++ b/tests/unittest/_torch/executor/test_mixed_decoder_cuda_graph_runner.py @@ -3,6 +3,8 @@ from types import SimpleNamespace +import torch + from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDAGraphRunner from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests @@ -30,7 +32,13 @@ def _mixed_batch() -> ScheduledRequests: def _runner() -> CUDAGraphRunner: runner = object.__new__(CUDAGraphRunner) - runner.config = SimpleNamespace(is_draft_model=False) + runner.config = SimpleNamespace( + enable_attention_dp=False, + is_draft_model=False, + use_mrope=False, + ) + runner.enabled = True + runner.padding_enabled = True runner.sparse_config = None runner.max_beam_width = 1 runner.enable_encoder_decoder_mixed_cuda_graph = True @@ -77,3 +85,75 @@ def test_mixed_encoder_decoder_graph_never_captures_at_runtime(): key = runner.get_graph_key(_mixed_batch()) assert not runner.needs_capture(key) + + +def test_mixed_encoder_decoder_graph_key_pads_encoder_extent(): + runner = _runner() + runner._capture_allowed = False + batch = _mixed_batch() + padded_key = (4, 0, False, False, True, (2, 2), (576,)) + graph_attn_metadata = object() + runner.graph_metadata[padded_key] = { + "attn_metadata": graph_attn_metadata, + "spec_metadata": None, + } + runner.graph_outputs[padded_key] = object() + + attn_metadata, spec_metadata, key = runner.maybe_get_cuda_graph( + batch, + enable_spec_decode=False, + attn_metadata=object(), + allow_mixed_encoder_decoder=True, + ) + + assert key == padded_key + assert attn_metadata is graph_attn_metadata + assert spec_metadata is None + + +def test_mixed_encoder_decoder_graph_key_uses_smallest_compatible_extent(): + runner = _runner() + runner._capture_allowed = False + key = runner.get_graph_key(_mixed_batch()) + larger_key = (*key[:6], (672,)) + smallest_key = (*key[:6], (576,)) + incompatible_key = (*key[:5], (2,), (544,)) + runner.graph_outputs = { + larger_key: object(), + smallest_key: object(), + incompatible_key: object(), + } + + assert runner._get_compatible_mixed_encoder_decoder_key(key) == smallest_key + + +def test_mixed_encoder_decoder_replay_zero_pads_encoder_hidden_states(): + runner = _runner() + key = (4, 0, False, False, True, (2, 2), (576,)) + attn_metadata = object() + runner.graph_metadata[key] = { + "attn_metadata": attn_metadata, + "spec_metadata": None, + } + runner.graph_outputs[key] = object() + runner.graphs[key] = SimpleNamespace(replay=lambda: None, reset=lambda: None) + runner.shared_static_tensors = { + "input_ids": torch.zeros(6, dtype=torch.int32), + "position_ids": torch.zeros((1, 6), dtype=torch.int32), + "encoder_hidden_states": torch.ones((576, 2)), + } + encoder_hidden_states = torch.arange(532 * 2, dtype=torch.float32).reshape(532, 2) + + runner.replay( + key, + { + "attn_metadata": attn_metadata, + "input_ids": torch.ones(6, dtype=torch.int32), + "position_ids": torch.ones((1, 6), dtype=torch.int32), + "encoder_hidden_states": encoder_hidden_states, + }, + ) + + staged_encoder_hidden_states = runner.shared_static_tensors["encoder_hidden_states"] + assert torch.equal(staged_encoder_hidden_states[:532], encoder_hidden_states) + assert torch.count_nonzero(staged_encoder_hidden_states[532:]) == 0 diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 6340ce22ab9e..9289de0c8abd 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -16,7 +16,7 @@ import threading import time import types -from unittest.mock import MagicMock, Mock +from unittest.mock import MagicMock, Mock, patch import pytest import torch @@ -145,17 +145,46 @@ def _make_async_encoder_executor(future): return executor -def _make_encoder_batch_wait_executor(): +def _make_encoder_batch_wait_executor(batch_sizes=None, encoder_max_batch_size=8): executor = object.__new__(PyExecutor) executor.max_batch_size = 32 + batch_sizes = batch_sizes or [1, 2, 4, 8] + executor.llm_args = types.SimpleNamespace( + encoder_cuda_graph_config=types.SimpleNamespace( + batch_sizes=batch_sizes, + enable_padding=True, + num_tokens=[96], + seq_lens=[512], + ), + encoder_max_batch_size=encoder_max_batch_size, + ) executor.batch_wait_timeout_iters = 48 executor.encoder_batch_wait_iters_count = 0 return executor -def test_encoder_microbatch_graph_waits_for_target(monkeypatch): - monkeypatch.setenv("TLLM_ENCODER_DECODER_MICROBATCH_CUDA_GRAPH_MAX_BATCH_SIZE", "8") - monkeypatch.setenv("TLLM_ENCODER_DECODER_MICROBATCH_LOW_WATERMARK", "24") +def test_encoder_graph_warmup_uses_runtime_encoder_stream(): + executor = object.__new__(PyExecutor) + executor.device_id = 3 + executor.encoder_stream = Mock() + executor.resource_manager = object() + executor.model_engine = Mock() + stream_context = MagicMock() + + with ( + patch("torch.cuda.set_device") as set_device, + patch("torch.cuda.stream", return_value=stream_context) as cuda_stream, + ): + executor._warmup_encoder_decoder_encoder_cuda_graphs() + + set_device.assert_called_once_with(3) + cuda_stream.assert_called_once_with(executor.encoder_stream) + executor.model_engine._warmup_encoder_decoder_encoder_cuda_graphs.assert_called_once_with( + executor.resource_manager + ) + + +def test_encoder_microbatch_graph_waits_for_target(): executor = _make_encoder_batch_wait_executor() encoder_requests = [object()] * 7 generation_requests = [object()] * 24 @@ -170,9 +199,7 @@ def test_encoder_microbatch_graph_waits_for_target(monkeypatch): assert executor.encoder_batch_wait_iters_count == 1 -def test_encoder_microbatch_graph_releases_target(monkeypatch): - monkeypatch.setenv("TLLM_ENCODER_DECODER_MICROBATCH_CUDA_GRAPH_MAX_BATCH_SIZE", "8") - monkeypatch.setenv("TLLM_ENCODER_DECODER_MICROBATCH_LOW_WATERMARK", "24") +def test_encoder_microbatch_graph_releases_target(): executor = _make_encoder_batch_wait_executor() encoder_requests = [object()] * 8 @@ -186,9 +213,7 @@ def test_encoder_microbatch_graph_releases_target(monkeypatch): assert executor.encoder_batch_wait_iters_count == 0 -def test_encoder_microbatch_graph_does_not_release_partial_at_low_watermark(monkeypatch): - monkeypatch.setenv("TLLM_ENCODER_DECODER_MICROBATCH_CUDA_GRAPH_MAX_BATCH_SIZE", "8") - monkeypatch.setenv("TLLM_ENCODER_DECODER_MICROBATCH_LOW_WATERMARK", "24") +def test_encoder_microbatch_graph_does_not_release_partial_at_low_watermark(): executor = _make_encoder_batch_wait_executor() encoder_requests = [object()] @@ -202,9 +227,7 @@ def test_encoder_microbatch_graph_does_not_release_partial_at_low_watermark(monk assert executor.encoder_batch_wait_iters_count == 1 -def test_encoder_microbatch_graph_caps_scheduler_overfill(monkeypatch): - monkeypatch.setenv("TLLM_ENCODER_DECODER_MICROBATCH_CUDA_GRAPH_MAX_BATCH_SIZE", "8") - monkeypatch.setenv("TLLM_ENCODER_DECODER_MICROBATCH_LOW_WATERMARK", "24") +def test_encoder_microbatch_graph_caps_scheduler_overfill(): executor = _make_encoder_batch_wait_executor() encoder_requests = [object() for _ in range(12)] @@ -218,9 +241,7 @@ def test_encoder_microbatch_graph_caps_scheduler_overfill(monkeypatch): assert executor.encoder_batch_wait_iters_count == 0 -def test_encoder_microbatch_graph_releases_supported_tail_at_deadline(monkeypatch): - monkeypatch.setenv("TLLM_ENCODER_DECODER_MICROBATCH_CUDA_GRAPH_MAX_BATCH_SIZE", "8") - monkeypatch.setenv("TLLM_ENCODER_DECODER_MICROBATCH_LOW_WATERMARK", "24") +def test_encoder_microbatch_graph_releases_supported_tail_at_deadline(): executor = _make_encoder_batch_wait_executor() executor.encoder_batch_wait_iters_count = executor.batch_wait_timeout_iters encoder_requests = [object() for _ in range(7)] @@ -235,6 +256,45 @@ def test_encoder_microbatch_graph_releases_supported_tail_at_deadline(monkeypatc assert executor.encoder_batch_wait_iters_count == 0 +def test_encoder_microbatch_graph_uses_configured_batch_sizes_at_deadline(): + executor = _make_encoder_batch_wait_executor(batch_sizes=[1, 3, 6], encoder_max_batch_size=8) + executor.encoder_batch_wait_iters_count = executor.batch_wait_timeout_iters + encoder_requests = [object() for _ in range(5)] + + scheduled = executor._waiting_encoder_requests( + encoder_requests, + [], + [], + ) + + assert scheduled == encoder_requests[:3] + assert executor.encoder_batch_wait_iters_count == 0 + + +def test_encoder_microbatch_graph_waits_above_low_watermark_after_deadline(): + executor = _make_encoder_batch_wait_executor() + executor.encoder_batch_wait_iters_count = executor.batch_wait_timeout_iters + encoder_requests = [object() for _ in range(8)] + + scheduled = executor._waiting_encoder_requests( + encoder_requests, + [], + [object() for _ in range(25)], + ) + + assert scheduled == [] + assert executor.encoder_batch_wait_iters_count == executor.batch_wait_timeout_iters + 1 + + scheduled = executor._waiting_encoder_requests( + encoder_requests, + [], + [object() for _ in range(24)], + ) + + assert scheduled == encoder_requests + assert executor.encoder_batch_wait_iters_count == 0 + + def test_pending_encoder_future_is_polled_without_blocking(): future = Mock() future.done.return_value = False diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py index 9724ad3d0a5e..930eb75bb596 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py @@ -10,6 +10,7 @@ import contextlib import unittest from dataclasses import dataclass +from types import SimpleNamespace from unittest.mock import patch import torch @@ -174,6 +175,55 @@ def _record(*msg): class TestWarmupCleanup(unittest.TestCase): """Lock in warmup-cleanup behavior introduced by PR #14609 (Plan B).""" + def test_main_cuda_graph_warmup_defers_encoder_decoder_encoder(self): + model_engine = object.__new__(PyTorchModelEngine) + model_engine.cuda_graph_runner = SimpleNamespace( + enabled=True, + is_warmup_only=True, + ) + model_engine.encoder_cuda_graph_runner = SimpleNamespace(enabled=True) + model_engine._torch_compile_piecewise_cuda_graph = False + resource_manager = object() + + with ( + patch.object(model_engine, "_capture_generation_cuda_graphs") as generation, + patch.object(model_engine, "_capture_mixed_encoder_decoder_cuda_graphs") as mixed, + patch.object(model_engine, "_capture_encoder_decoder_encoder_cuda_graphs") as encoder, + ): + model_engine._run_cuda_graph_warmup(resource_manager) + + generation.assert_called_once_with(resource_manager) + mixed.assert_called_once_with(resource_manager) + encoder.assert_not_called() + + def test_encoder_decoder_encoder_warmup_uses_two_passes(self): + model_engine = object.__new__(PyTorchModelEngine) + model_engine.is_warmup = False + + @contextlib.contextmanager + def allow_capture(): + yield + + runner = SimpleNamespace( + enabled=True, + is_encoder_decoder=True, + is_warmup_only=False, + allow_capture=allow_capture, + ) + model_engine.encoder_cuda_graph_runner = runner + resource_manager = object() + warmup_states = [] + + with patch.object( + model_engine, + "_capture_encoder_decoder_encoder_cuda_graphs", + side_effect=lambda _: warmup_states.append(runner.is_warmup_only), + ): + model_engine._warmup_encoder_decoder_encoder_cuda_graphs(resource_manager) + + assert warmup_states == [True, False] + assert not runner.is_warmup_only + def test_empty_cache_fires_immediately_after_autotuner(self): """Change 1 placement: empty_cache must be the call right after _run_autotuner_warmup.""" diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index 7ff208c41ce3..008364ab19f0 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -87,6 +87,10 @@ methods: annotation: Union[tensorrt_llm.llmapi.llm_args.DecodeCudaGraphConfig, tensorrt_llm.llmapi.llm_args.EncodeCudaGraphConfig, NoneType] default: null status: beta + encoder_cuda_graph_config: + annotation: Optional[tensorrt_llm.llmapi.llm_args.EncodeCudaGraphConfig] + default: null + status: prototype multimodal_config: annotation: tensorrt_llm.llmapi.llm_args.MultimodalConfig default: null diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index b5410a8edfdb..597754b33293 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -1511,6 +1511,56 @@ def test_cuda_graph_config_accepts_encoder_config(self): assert args.cuda_graph_config.seq_lens == [8, 32] assert args.cuda_graph_config.max_seq_len == 32 + def test_encoder_decoder_cuda_graph_configs(self): + args = TorchLlmArgs( + model=llama_model_path, + encoder_max_batch_size=4, + cuda_graph_config=DecodeCudaGraphConfig( + batch_sizes=[1, 4], + enable_padding=True, + ), + encoder_cuda_graph_config=EncodeCudaGraphConfig( + batch_sizes=[1, 4], + num_tokens=[16, 64], + seq_lens=[8, 32], + enable_padding=True, + ), + ) + + assert isinstance(args.cuda_graph_config, DecodeCudaGraphConfig) + assert isinstance(args.encoder_cuda_graph_config, EncodeCudaGraphConfig) + assert args.encoder_cuda_graph_config.batch_sizes == [1, 4] + assert args.encoder_cuda_graph_config.num_tokens == [16, 64] + assert args.encoder_cuda_graph_config.seq_lens == [8, 32] + + def test_encoder_cuda_graph_config_requires_encoder_max_batch_size(self): + with pytest.raises(ValidationError, + match=("encoder_cuda_graph_config requires " + "encoder_max_batch_size")): + TorchLlmArgs( + model=llama_model_path, + encoder_cuda_graph_config=EncodeCudaGraphConfig( + batch_sizes=[1, 4], + num_tokens=[16, 64], + seq_lens=[8, 32], + enable_padding=True, + ), + ) + + def test_encoder_cuda_graph_config_requires_shape_dimensions(self): + with pytest.raises( + ValidationError, + match=("encoder_cuda_graph_config requires " + "num_tokens/max_num_token and seq_lens/max_seq_len")): + TorchLlmArgs( + model=llama_model_path, + encoder_max_batch_size=4, + encoder_cuda_graph_config=EncodeCudaGraphConfig( + batch_sizes=[1, 4], + enable_padding=True, + ), + ) + def test_cuda_graph_config_infers_encode_mode_from_raw_dict(self): args = TorchLlmArgs( model=llama_model_path, From 04ca506a91e40932bd166de964ce659d27226144 Mon Sep 17 00:00:00 2001 From: Guiju Zhang <7135567+cascade812@users.noreply.github.com> Date: Wed, 29 Jul 2026 18:03:54 -0700 Subject: [PATCH 06/15] [None][perf] improve encoder-decoder CUDA graph handling Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com> --- ...-launch-path-optimization-opportunities.md | 679 ------------------ docs/source/models/encoder-decoder.md | 8 + optimize_plan_2.md | 494 ------------- .../_torch/pyexecutor/cuda_graph_runner.py | 152 ++-- tensorrt_llm/_torch/pyexecutor/llm_request.py | 3 - .../_torch/pyexecutor/model_engine.py | 199 ++--- .../_torch/pyexecutor/scheduler/scheduler.py | 3 - tensorrt_llm/llmapi/llm_args.py | 12 +- .../usage/llm_args_golden_manifest.json | 7 + .../test_encoder_cuda_graph_runner.py | 355 --------- .../test_mixed_decoder_cuda_graph_runner.py | 159 ---- .../api_stability/references/llm.yaml | 6 +- tests/unittest/llmapi/test_llm_args.py | 16 + 13 files changed, 207 insertions(+), 1886 deletions(-) delete mode 100644 docs/source/models/bart-decoder-launch-path-optimization-opportunities.md delete mode 100644 optimize_plan_2.md delete mode 100644 tests/unittest/_torch/executor/test_encoder_cuda_graph_runner.py delete mode 100644 tests/unittest/_torch/executor/test_mixed_decoder_cuda_graph_runner.py diff --git a/docs/source/models/bart-decoder-launch-path-optimization-opportunities.md b/docs/source/models/bart-decoder-launch-path-optimization-opportunities.md deleted file mode 100644 index 34dcf94dd7f0..000000000000 --- a/docs/source/models/bart-decoder-launch-path-optimization-opportunities.md +++ /dev/null @@ -1,679 +0,0 @@ - - -# BART decoder launch-path optimization opportunities - -This note records the remaining host-side bottlenecks observed while profiling -the BART PyTorch encoder-decoder path, along with possible ways to reduce GPU -launch starvation. The proposals are diagnostic follow-up work, not verified -performance improvements, except where a completed experiment is marked -verified. - -## Profile context - -The analysis uses the Nsight Systems report: - -```text -/tmp/bart-encoder-piecewise-profile/piecewise-direct.nsys-rep -``` - -The run used continuous admission with a maximum scheduled decoder batch of 32 -requests. Generation step 208 contained 31 generation requests and no context -requests. It therefore already used the padded batch-32 decoder CUDA graph -bucket. Raising concurrency can amortize host overhead, but does not address -the remaining launch cost at a fixed concurrency. - -A blank region between NVTX ranges means that the CPU work is not covered by an -NVTX annotation. It does not by itself mean that the CPU thread is idle. In -this trace, the gaps before `_update_requests` and `_fetch_new_requests` mostly -overlap forward kernels from the current batch. The interval after -`prepare_resources`, however, substantially overlaps a genuinely idle GPU. - -## Representative launch-path breakdown - -For generation step 208, the host path from the end of `prepare_resources` to -the end of `[Executor] _forward_step` decomposed as follows: - -| Host region | Profiled duration | -| --- | ---: | -| Before `[Executor] _forward_step` | 126.9 us | -| `[Executor]` entry to `_prepare_inputs` | 76.2 us | -| `_prepare_inputs` | 707.4 us | -| `_prepare_inputs` end to `cudaGraphLaunch` | 102.9 us | -| `cudaGraphLaunch` API call | 306.9 us (The experiment confirms the ~300 µs cudaGraphLaunch duration is Nsight tracing overhead.) | -| Remaining host unwind | 36.0 us | - -These are durations under CUDA and NVTX tracing, not unprofiled latency -measurements. In particular, the approximately 280--310 us steady-state -`cudaGraphLaunch` duration is unusually high for a reused graph and may include -substantial profiler overhead. It must be confirmed with an unprofiled host -timer around graph replay, or a less intrusive graph-level trace, before -treating it as an implementation bottleneck. The experiment below provides -that confirmation. - -## Promising approaches - -### Reduce encoder-decoder input preparation - -Cache stable generation metadata for an unchanged batch, including request IDs, -sequence slots, prompt lengths, cross-KV lengths and pointers, and the decoder -graph key. Update only sampled tokens, positions, and KV-length deltas on each -step. - -Fine-grained debug-only NVTX ranges were added to the native encoder-decoder -fast path. Set `TLLM_NVTX_DEBUG=1` to enable them. The profiles are: - -```text -/tmp/bart-encoder-fast-input-profile/native-fast-full-c32-r256.nsys-rep -/tmp/bart-encoder-fast-input-profile/native-fast-fine-c32-r256.nsys-rep -``` - -The first report has coarse ranges with less annotation overhead. Across its -893 generation-only fast-path calls, the representative P50 breakdown was: - -| Fast-path region | P50 | -| --- | ---: | -| Complete encoder-decoder fast path | 513.8 us | -| Input-ID staging | 161.3 us | -| Attention-metadata preparation | 145.9 us | -| Cross-attention preparation | 37.9 us | -| Host-buffer retirement | 35.8 us | -| Position-ID staging | 30.5 us | -| Sequence-length staging | 15.8 us | -| Native request collation | 15.0 us | - -The second report subdivides input-ID and attention-metadata preparation. In -generation step 208, which contained nine generation requests, input-ID -staging took 198.1 us under tracing: - -| Input-ID subregion | Profiled duration | CUDA API work | -| --- | ---: | --- | -| Copy previous batch indices | 33.2 us | One asynchronous copy | -| Gather sampled tokens | 76.9 us | Two kernel launches | -| Copy gathered tokens | 28.2 us | One asynchronous copy | -| Fill graph-padding tokens | 23.3 us | One kernel launch | - -The complete input-ID region issued three kernel launches and two asynchronous -copies. - -#### Sampled-token staging consolidation (Verified: clean end to end +1.0%) - -The sampled-token experiment replaced advanced indexing followed by a copy -with `torch.index_select(..., out=input_ids_cuda_slice)`. It also retains the -device sequence-slot indices while the ordered generation request IDs remain -unchanged. Any non-encoder-decoder preparation path invalidates that cache. - -The optimized profiles are: - -```text -/tmp/bart-encoder-fast-input-profile/sampled-token-direct-c32-r256.nsys-rep -/tmp/bart-encoder-fast-input-profile/sampled-token-direct-cache-c32-r256.nsys-rep -``` - -| Fine-grained P50 | Before | Direct gather | Direct gather plus index reuse | -| --- | ---: | ---: | ---: | -| Complete encoder-decoder fast path | 612.0 us | 560.8 us | 519.1 us | -| Input-ID staging | 194.9 us | 142.8 us | 109.3 us | -| Sampled-token gather/copy | 103.5 us | 58.3 us | 58.9 us | -| Previous-index copy count | 907 | 907 | 202 | - -Index reuse eliminated 705 of 907 previous-index copies. For generation step -208, the sampled-token portion fell from 105.0 to 48.7 us and changed from two -kernel launches plus one asynchronous copy to one kernel launch. A standalone -unprofiled CUDA microbenchmark reduced this operation from 29.1 to 10.6 us at -batch 9 and from 29.6 to 14.3 us at batch 32. - -Because CUDA API tracing inflates the profiled ranges, a diagnostic experiment -temporarily placed `time.perf_counter_ns()` timers around the generation-only -`model_engine.forward()` call and input-ID staging. The timers did not -synchronize the GPU, did not emit per-step output, skipped the first 256 -generation steps, and collected 5,785 samples per run. Both NVTX environment -switches were disabled. The fine-grained `nvtx_range_debug` context managers -were still present as null context managers, however, so this established the -direction of the host-path change but was not a completely -instrumentation-free measurement. - -| Unprofiled host region | Before | Optimized | Reduction | -| --- | ---: | ---: | ---: | -| Input-ID staging, mean | 142.652 us | 94.670 us | 33.6% | -| Input-ID staging, P50 | 131.785 us | 81.866 us | 37.9% | -| Complete generation forward launch, mean | 657.765 us | 615.086 us | 6.5% | -| Complete generation forward launch, P50 | 600.186 us | 550.230 us | 8.3% | - -A fully clean end-to-end experiment then physically removed all added -fine-grained ranges from `model_engine.py` and `trtllm.py`, removed the host -timers, and ran without Nsight. It alternated optimized and pre-change code for -four concurrency-32, 2,048-request runs per version: - -| Order | Version | Mean latency | Makespan | Requests/s | Output tokens/s | -| ---: | --- | ---: | ---: | ---: | ---: | -| 1 | Optimized | 213.794 ms | 13.822200 s | 148.167 | 9953.842 | -| 2 | Before | 217.299 ms | 14.046938 s | 145.797 | 9794.590 | -| 3 | Optimized | 215.025 ms | 13.905088 s | 147.284 | 9894.507 | -| 4 | Before | 217.570 ms | 14.039977 s | 145.869 | 9800.087 | -| 5 | Optimized | 216.035 ms | 13.956041 s | 146.746 | 9855.947 | -| 6 | Before | 215.996 ms | 13.968241 s | 146.618 | 9849.773 | -| 7 | Optimized | 214.392 ms | 13.863271 s | 147.728 | 9924.353 | -| 8 | Before | 217.528 ms | 14.061549 s | 145.645 | 9784.413 | - -| Four-run average | Before | Optimized | Change | -| --- | ---: | ---: | ---: | -| Mean latency | 217.098 ms | 214.812 ms | -1.05% | -| P50 latency | 202.168 ms | 200.133 ms | -1.01% | -| P90 latency | 345.073 ms | 340.892 ms | -1.21% | -| Makespan | 14.029176 s | 13.886650 s | -1.02% | -| Requests/s | 145.982 | 147.481 | +1.03% | -| Output tokens/s | 9807.216 | 9907.162 | +1.02% | - -Two of the eight runs differed slightly in natural-EOS placement, once on -each version; output-token throughput gives the same approximately 1.0% result -after accounting for that small workload variation. The clean comparison -therefore confirms a modest end-to-end gain, while the diagnostic host timers -and Nsight profiles explain where it originates. - -Attention-metadata preparation took 188.7 us in the same step: - -| Attention-metadata subregion | Profiled duration | CUDA API work | -| --- | ---: | --- | -| Stage prompt and KV lengths | 40.7 us | Two asynchronous copies | -| Update host metadata | 23.6 us | Host-only tensor updates | -| Copy KV block offsets | 75.1 us | One asynchronous copy plus event query/record | -| Bind runtime views | 10.0 us | Host-only view binding | - -KV block-offset staging is the largest individual metadata subregion. Reusing -its pinned staging storage and avoiding the copy when the request-to-block -mapping is unchanged should be measured after token staging is consolidated. -This optimization must preserve the existing completion-event lifetime rules -for overlapped scheduling. - -Nested NVTX annotations and CUDA API tracing materially inflate all absolute -durations in the fine-grained report. The ranges establish relative -attribution; they are not unprofiled latency measurements. In particular, -native request collation is already small, while mixed admission makes -cross-attention preparation dominant only on the relatively infrequent -batch-change iterations. - -The representative `_prepare_inputs` range issued 28 CUDA API calls: - -- Six asynchronous copies. -- Three small preparation kernels. -- Three event queries. -- Four event records. - -The remaining calls were stream-state queries and kernel-name lookups recorded -by the profiler. Packing the small metadata transfers into one pinned buffer, -fusing the preparation kernels, or capturing both into the decoder graph would -reduce launch-critical work. - -#### Metadata, retirement, and position staging follow-up - -Three safe changes were prototyped together and separately: - -- Bound KV block-offset staging to the columns required by the batch's maximum - KV length. -- Re-record the completion event associated with an available host-buffer set - instead of constructing another event. -- For decoder CUDA-graph replay, copy position IDs directly from their pinned - host buffer into the graph's static position tensor and skip the otherwise - redundant device-to-device copy in `CUDAGraphRunner.replay()`. - -Reusing the completion event and bounding the block-offset copy were neutral -in two 2,048-request runs: - -| Two-run average | Before | Event + block bound | Change | -| --- | ---: | ---: | ---: | -| Mean latency | 214.292 ms | 214.212 ms | -0.04% | -| Makespan | 13.855197 s | 13.849266 s | -0.04% | -| Requests/s | 147.815 | 147.879 | +0.04% | - -The direct position-ID path by itself was also within run-to-run noise: - -| Two-run average | Before | Direct position staging | Change | -| --- | ---: | ---: | ---: | -| Mean latency | 208.198 ms | 207.623 ms | -0.28% | -| Makespan | 13.458528 s | 13.421465 s | -0.28% | -| Requests/s | 152.172 | 152.597 | +0.28% | - -One pair favored the position change by 0.80%, while the next favored the -baseline by 0.24%. A longer 8,192-request B--O--B run then compared all three -safe changes without Nsight or debug NVTX instrumentation: - -| Order | Version | Mean latency | Makespan | Requests/s | Output tokens/s | -| ---: | --- | ---: | ---: | ---: | ---: | -| 1 | Before | 225.061 ms | 57.744981 s | 141.865 | 10517.278 | -| 2 | Optimized | 227.479 ms | 58.368348 s | 140.350 | 10404.954 | -| 3 | Before | 231.405 ms | 59.377079 s | 137.966 | 10228.189 | -| Baseline average | Before | 228.233 ms | 58.561030 s | 139.916 | 10372.734 | - -The optimized run was 0.31% faster than the average of its surrounding -baselines, but those baselines differed by 2.75%. All three runs generated the -same 607,320 output tokens with the same output hash. The measured change is -therefore not distinguishable from environmental drift, and none of these -three code changes was retained. - -An additional attempt cached prompt lengths and a pinned KV block-offset -snapshot, skipping H2D staging when their host contents appeared unchanged. -That is not safe under the current overlapped metadata ownership contract: the -2,048-request output changed from 137,584 to 150,955 tokens and produced a -different hash. The prototype was removed. Future block-table reuse must use -an explicit cache-manager generation/ownership contract rather than infer -device-buffer validity from equal host contents. - -### Graph replay timing without tracing (Verified: 10 us launch, 47 us runner) - -An environment-injected timer measured graph replay during the -`timed_generate` range of the same concurrency-32, 256-request workload. The -benchmark ran as plain Python without Nsight Systems or CUDA API tracing. The -timer was injected into the MPI executor worker and used -`time.perf_counter_ns()` immediately around both: - -- `torch.cuda.CUDAGraph.replay()`, which is the direct counterpart of the - profiled `cudaGraphLaunch` API call. -- `CUDAGraphRunner.replay()`, which additionally includes graph lookup and - copies of input IDs and position IDs into the graph's static tensors. - -The native extension available to this worktree did not contain -`prepare_encoder_decoder_inputs`, so the runs forced the Python collation -fallback. This matches the earlier eager-versus-piecewise profile setup and -does not change the graph replay implementation being timed. - -| Run and scope | Calls | Mean | P50 | P99 | -| --- | ---: | ---: | ---: | ---: | -| Run 1, CUDA graph replay | 767 | 10.493 us | 10.293 us | 16.677 us | -| Run 2, CUDA graph replay | 767 | 10.331 us | 9.881 us | 20.569 us | -| Run 3, CUDA graph replay | 767 | 10.542 us | 10.209 us | 18.724 us | -| Run 3, complete decoder graph runner | 767 | 47.394 us | 46.718 us | 61.594 us | - -The median cost of an empty timer pair was 66--69 ns. All three runs produced -the same output hash and natural-EOS count. Their end-to-end mean latencies -were 217.267, 218.838, and 224.212 ms, respectively; these absolute values -describe the fallback diagnostic rather than the optimized native collation -path. - -The untraced replay is about 27--30 times shorter than the 280--310 us -steady-state `cudaGraphLaunch` calls in the Nsight trace. Graph launch itself -therefore is not the approximately 300 us bottleneck implied by the traced -timeline. Even the full graph runner is below 50 us at the median, and only -about 10 us of that is the native graph replay. Further work should prioritize -input preparation and the host path leading into the runner rather than graph -executable reuse or upload. - -### Reuse the input-copy completion event - -The model engine constructs and records a new completion event after every -encoder-decoder input preparation. Re-recording the event owned by an -available host-buffer set preserved correctness, but the experiment above -showed no measurable end-to-end benefit. - -### Remove redundant pre-forward Python work - -The overlap loop sorts generation requests every iteration even though the -ordering correction is only needed for disaggregated generation. The -aggregated BART path can skip that list allocation and key-function traversal. - -A qualified generation-only fast path can also skip: - -- Context-token summation. -- Context-logit checks. -- Encoder-output attachment. -- Cache-indirection lookup when it is not used. - -Using an NVTX context directly instead of constructing a decorated nested -function on every iteration removes another small fixed cost. - -### Reduce fixed stream-handoff overhead - -The stream waits around forward implicitly create, record, wait on, and destroy -CUDA events. Reusing preallocated handoff events, or keeping forward and its -dependent sampling work on one stream when KV transfer is inactive, can remove -this fixed overhead. Any change must preserve the ordering required by KV -onboard/offload and asynchronous sampling. - -### Prepare batch-static metadata earlier - -Split decoder input preparation into: - -- A batch-static portion: request ordering, block tables, cross-attention - metadata, and graph selection. -- A token-dependent portion: sampled tokens, positions, and KV-length - increments. - -The batch-static portion can run while the preceding GPU step is executing. -Once its sampled token is available, the launch-critical path should contain -only a compact device update and graph replay. - -### Continue decode while encoder futures are pending (Verified: neutral end to end) - -Encoder-init requests and existing generation requests are independent. The -executor submits encoder forward to one persistent host worker immediately -before decoder forward, keeps the encoder request IDs in the scheduler's -in-flight set, and retains the returned futures in FIFO order. At the start of -each later iteration, the executor polls the oldest future with `done()` and -queries its CUDA completion event. Neither operation waits. Only after both -report completion does the main executor thread publish the encoder output, -transition the request from `ENCODER_INIT` to `CONTEXT_INIT`, and remove its ID -from the in-flight set. This gives request-state mutation and error handling -back to the main executor thread while preventing both duplicate encoder -submission and premature decoder-context admission. - -This removes the prior same-iteration `future.result()` barrier. In -`/tmp/bart-encoder-pending-futures-c32-r256.nsys-rep`, the encoder range from -40.744076 to 40.768364 seconds overlaps host execution of generation-only -decoder steps 211 through 218. Other steady-state encoder ranges similarly -overlap four to nine generation-only decoder forwards. The executor therefore -continues fetch, schedule, input preparation, forward, and sampling work while -the encoder future is pending. - -GPU concurrency remains limited. In the 2.371-second timed window, encoder -stream 21 had 66.106 ms of kernel-busy time, decoder stream 17 had 41.064 ms, -and sampler stream 7 had 14.210 ms. Encoder kernels overlapped decoder kernels -for 0.595 ms and sampler kernels for 0.075 ms, about 1.0% of encoder busy time. -The host pipeline is now independent, but the large BART kernels still consume -most available device resources. - -The compatible temporary native mixed-batch binding used by the earlier -measurements was cleaned before this experiment. The following A/B runs -therefore forced the generic Python input path for both versions; they isolate -the pending-future scheduling change but are not directly comparable to the -native-fast-path throughput above. - -| Pair | Version | Mean latency | Makespan | Requests/s | -| ---: | --- | ---: | ---: | ---: | -| 1 | Blocking encoder | 240.779 ms | 15.563809 s | 131.587 | -| 1 | Pending futures | 243.423 ms | 15.714822 s | 130.323 | -| 2 | Blocking encoder | 242.596 ms | 15.690760 s | 130.523 | -| 2 | Pending futures | 240.687 ms | 15.532974 s | 131.849 | -| Average | Blocking encoder | 241.688 ms | 15.627285 s | 131.055 | -| Average | Pending futures | 242.055 ms | 15.623898 s | 131.086 | - -Average throughput changed by +0.02% and mean latency by +0.15%, both within -run-to-run noise. The blocking baseline produced 137,584 tokens and hash -`93761bbaed28ad0f` in both runs. Pending-future execution changed batch -composition and produced 137,594 and 137,545 tokens; the natural-EOS and -length-stop request counts remained identical. A 64-request token-level check -found only request 27 diverged, beginning at generated token 94, with the same -128-token output length. This is consistent with a close greedy decision -changing under different BF16 batch numerics rather than request/output -misassociation. - -The change removes a real software barrier but does not improve this workload's -end-to-end performance. A material gain still requires coarser encoder replay -(for example, a whole-encoder CUDA graph) or kernels that leave complementary -GPU resources available. - -#### Prioritize decoder kernels over encoder kernels (Rejected) - -An experiment created the encoder-decoder execution stream with CUDA priority -`-1` while retaining priority `0` for the encoder stream. Decoder-only models -kept priority `0`. CUDA stream priority favors pending work on the -higher-priority stream when the GPU scheduler can choose new work, but it -cannot preempt an encoder kernel that is already running. - -The clean benchmark alternated default and high decoder priority at concurrency -32 for 2,048 requests. Both versions used the generic encoder-decoder input -path because the loaded native extension predates -`prepare_encoder_decoder_inputs`. No Nsight or CUDA API tracing was active. - -| Order | Decoder priority | Mean latency | P50 latency | P90 latency | P99 latency | Makespan | Requests/s | Output tokens/s | -| ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | -| 1 | 0 | 242.986 ms | 227.216 ms | 374.525 ms | 461.013 ms | 15.685588 s | 130.566 | 8764.224 | -| 2 | -1 | 244.408 ms | 228.988 ms | 380.483 ms | 462.689 ms | 15.791230 s | 129.692 | 8711.038 | -| 3 | 0 | 244.749 ms | 229.705 ms | 381.593 ms | 456.817 ms | 15.825666 s | 129.410 | 8689.302 | -| 4 | -1 | 245.734 ms | 231.712 ms | 379.946 ms | 460.800 ms | 15.923400 s | 128.616 | 8644.699 | - -| Two-run average | Priority 0 | Priority -1 | Change | -| --- | ---: | ---: | ---: | -| Mean latency | 243.868 ms | 245.071 ms | +0.49% | -| P50 latency | 228.461 ms | 230.350 ms | +0.83% | -| P90 latency | 378.059 ms | 380.215 ms | +0.57% | -| P99 latency | 458.915 ms | 461.745 ms | +0.62% | -| Makespan | 15.755627 s | 15.857315 s | +0.65% | -| Requests/s | 129.988 | 129.154 | -0.64% | -| Output tokens/s | 8726.763 | 8677.869 | -0.56% | - -Output lengths varied by at most 0.13% between runs; token-normalized -throughput therefore gives the same conclusion as request throughput. The -benchmark records final-response latency rather than per-token inter-token -latency, so it does not exclude a small latency redistribution from existing -generation requests toward replacement encoder requests. It does show that -the proposed priority does not improve end-to-end request latency or -throughput for this workload. The stream-priority change was not retained. - -#### Re-evaluate encoder batch waiting (Current settings retained) - -The pending-future scheduler was adjusted at concurrency 32 because encoder -launch and completion no longer block the decoder host loop. A 256-request -screen covered iteration deadlines 24, 32, 40, 48, and 64 and token-threshold -ratios 0.08, 0.12, 0.1708984375, and 0.25. The existing setting is 48 -iterations and ratio 0.1708984375, which corresponds to 11,200 tokens or -approximately one full batch of 32 average-length encoder inputs. - -The short screen selected 64 iterations and ratio 0.08, but an alternating -2,048-request validation rejected it: it averaged 129.14 requests/s versus -130.71 requests/s for the existing configuration. The less aggressive -64-iteration, 0.12-ratio candidate was positive in two 2,048-request pairs, -improving mean latency and requests/s by 1.50% and output-token throughput by -1.61%. A longer confirmation reduced that result to noise and exposed worse -median and tail latency: - -| Order | Iterations / ratio | Mean latency | P50 latency | P90 latency | P99 latency | Makespan | Requests/s | Output tokens/s | -| ---: | --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | -| 1 | 48 / 0.1708984375 | 256.200 ms | 245.680 ms | 379.094 ms | 452.209 ms | 32.940426 s | 124.346 | 8883.522 | -| 2 | 64 / 0.12 | 254.405 ms | 247.532 ms | 373.355 ms | 463.342 ms | 32.761208 s | 125.026 | 8936.117 | -| 3 | 48 / 0.1708984375 | 254.777 ms | 243.711 ms | 373.563 ms | 452.560 ms | 32.788736 s | 124.921 | 8928.371 | -| 4 | 64 / 0.12 | 255.028 ms | 247.642 ms | 374.130 ms | 465.251 ms | 32.799612 s | 124.880 | 8925.166 | - -| Two-run average | 48 / 0.1708984375 | 64 / 0.12 | Change | -| --- | ---: | ---: | ---: | -| Mean latency | 255.489 ms | 254.717 ms | -0.30% | -| P50 latency | 244.696 ms | 247.587 ms | +1.18% | -| P90 latency | 376.329 ms | 373.743 ms | -0.69% | -| P99 latency | 452.385 ms | 464.297 ms | +2.63% | -| Makespan | 32.864581 s | 32.780410 s | -0.26% | -| Requests/s | 124.634 | 124.953 | +0.26% | -| Output tokens/s | 8905.947 | 8930.642 | +0.28% | - -Both versions used the generic encoder-decoder input path because the loaded -native extension predates `prepare_encoder_decoder_inputs`; neither used -Nsight or CUDA API tracing. The 0.26--0.28% throughput change is within the -observed run-to-run variation, while the P50 and P99 regressions are larger. -The existing 48-iteration, 0.1708984375-ratio configuration was therefore -retained. - -#### Drain the decoder before admitting another encoder wave (Rejected) - -A stricter scheduling experiment stopped admitting encoder requests whenever -the scheduler had any decoder-context or generation requests. New requests -therefore accumulated until the entire active decoder wave completed, at which -point the scheduler released the waiting encoder requests together. This -formed larger encoder batches, but also made every replacement request wait -for the longest output in the preceding decoder wave. - -The clean concurrency-32 comparison used 1,024 requests, the same dataset and -request order, no Nsight or CUDA API tracing, and the generic encoder-decoder -input path because the loaded native extension predates -`prepare_encoder_decoder_inputs`. - -| Policy | Mean latency | P50 latency | P90 latency | P99 latency | Makespan | Requests/s | Output tokens/s | -| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | -| Current bounded encoder accumulation | 225.910 ms | 199.376 ms | 387.849 ms | 472.617 ms | 7.383563 s | 138.686 | 8276.492 | -| Drain decoder completely | 306.773 ms | 285.259 ms | 425.774 ms | 526.375 ms | 10.012829 s | 102.269 | 6109.162 | -| Change | +35.8% | +43.1% | +9.8% | +11.4% | +35.6% | -26.3% | -26.2% | - -The larger encoder waves do not compensate for the head-of-line admission -barrier. In particular, short decoder requests cannot be replaced until the -longest request in the current wave finishes, leaving capacity unused during -the decoder tail. The strict policy was removed; bounded encoder accumulation -continues to admit replacement work while generation proceeds. - -#### Preallocate stable encoder output and prepare decoder context early - -The stronger follow-up publishes stable encoder-output views before submitting -the encoder worker. `ModelEngine` allocates one exactly sized output tensor for -the encoder batch, each request receives a view into that tensor, and the -worker copies the final hidden states into those stable addresses. A CUDA event -is also attached before submission. A separate recorded flag prevents the main -thread from querying or waiting on the event until the worker has enqueued it. -This removes the lifetime hazard that otherwise prevents decoder-context -preparation from starting before encoder completion. - -The scheduler can now capacity-plan those requests in `CONTEXT_INIT` while the -encoder worker is active. The executor removes requests whose encoder event is -not ready from the current decoder batch, retains its generation requests, and -prepares the deferred requests' persistent KV, cross-KV, and sequence-slot -resources in that generation-only iteration. The requests remain in the -executor's in-flight set so they cannot be scheduled twice. Once the recorded -event queries ready, the executor releases them to a later mixed decoder batch. -That mixed batch skips resource allocation already completed by the lookahead -pass and waits once on the shared encoder event before consuming the stable -output views. - -This is host/preparation overlap and launch-bubble removal. It does not attempt -to execute decoder-context kernels before the encoder event is satisfied, and -it retains generation-only decoder iterations while the encoder kernels run. -The optimization is enabled only when every active non-null resource manager is -one of the KV-cache, cross-KV-cache, or sequence-slot managers and a cross-KV -manager is present. Other configurations use the non-blocking encoder future -path without early resource preparation. - -In -`/tmp/bart-encoder-stable-lookahead-early-c32-r256.nsys-rep`, there are 916 -decoder forwards and 928 `prepare_resources` ranges. The 12 additional calls -occur immediately before generation-only forwards while encoder work is still -active. For example, decoder-context preparation runs from 66.385924 to -66.386937 seconds while the encoder runs from 66.381330 to 66.412223 seconds. -The following generation-only forward starts at 66.387195 seconds. When the -encoder event becomes ready, mixed forward 342 performs only a 98 us resource -preparation before starting at 66.414147 seconds. - -Across steady-state encoder admissions, excluding the two startup groups, the -median timings changed as follows: - -| Admission interval | Pending-future baseline | Stable-output lookahead | Change | -| --- | ---: | ---: | ---: | -| Encoder worker return to mixed-forward start | 2.665 ms | 2.405 ms | -9.8% | -| Last encoder GPU operation to first decoder GPU operation | 3.081 ms | 2.900 ms | -5.9% | -| Last encoder GPU operation to first decoder GEMM | 5.400 ms | 5.295 ms | -1.9% | - -One representative pair shows the intended effect more clearly, although the -mixed batch shapes are not identical. The pending-future trace's 13-context, -15-generation batch starts 2.622 ms after its encoder worker returns and its -first decoder input operation begins 3.027 ms after the last encoder GPU -operation. The lookahead trace's 15-context, 12-generation batch reduces those -intervals to 1.924 and 2.253 ms, respectively. - -Clean unprofiled measurements used concurrency 32, 2,048 requests, and forced -the generic encoder-decoder input path in both temporary overlays because the -available compatible native binding predates the current collation interface. -The repeated results remain noisy: - -| Version | Runs | Median mean latency | Median makespan | Median requests/s | -| --- | ---: | ---: | ---: | ---: | -| Pending-future baseline | 5 | 245.028 ms | 15.871215 s | 129.039 | -| Stable-output lookahead | 4 | 242.236 ms | 15.658510 s | 130.794 | - -The medians correspond to -1.14% mean latency, -1.34% makespan, and +1.36% -throughput. Individual optimized runs ranged from 127.942 to 136.162 requests/s, -so this is a small directional gain rather than a statistically decisive -end-to-end improvement. The trace provides the stronger causal result: the -resource work moves earlier and the steady-state admission bubble shrinks, but -most of the path to the first decoder GEMM remains elsewhere in input -preparation and model launch. - -#### Run encoder and decoder context together on the side stream - -A stronger experiment moves the dependent decoder context-only forward into -the encoder worker. The worker executes the encoder and then its decoder -context batch on stream 21, while the main executor continues launching -generation-only CUDA graphs on stream 17. - -The implementation uses a second `PyTorchModelEngine` that shares the BART -model but owns separate decoder input buffers, attention metadata, and graph -runner state. CUDA graphs are disabled for the context engine. KV, cross-KV, -and sequence-slot resources are reserved on the main executor thread because -their managers mutate shared allocation state; decoder input preparation and -model launch remain in the worker. Only one worker context batch can be active, -so its engine and stable encoder-output storage cannot be reused early. - -Two ordering details are required for correctness: - -- The main sampler has reusable device storage and already has one outstanding - generation sample under overlap scheduling. Completed context logits are - therefore sampled only after the prior main-lane sample has been retired. -- A worker context batch is not the main executor's previous batch. - `py_batch_idx` is cleared after its sampled token reaches the host request, - causing the first main-lane generation step to use the explicit host-token - admission path. - -The two engines have distinct destination buffers, but their asynchronous -metadata copies are staged from shared resource-manager state. Reusing that -state before an H2D copy completed corrupted subsequent generation tokens. -Each lane now waits only for its input-copy event after enqueueing the complete -forward. This does not wait for model kernels: host preparation can proceed on -both threads, and the already-enqueued model work remains concurrent on the two -CUDA streams. - -The implementation is restricted to single-rank BART and mBART with overlap -scheduling, KV-cache manager V1, no attention DP, drafting, guided decoding, -KV-cache transfer, connector, or early first-token response. Other -configurations retain the stable-output lookahead path. - -The profile is: - -```text -/tmp/bart-encoder-context-worker-c32-r256.nsys-rep -``` - -During its 2.651-second `timed_generate` interval, stream 21 executes 96.831 ms -of encoder plus context kernels. Stream 17 executes 973.694 ms of decoder CUDA -graphs. Their kernels overlap for 20.641 ms, or 21.3% of stream-21 kernel time. -Steady worker `_run_encoder_context_step` host ranges also overlap seven to ten -main generation `_forward_step` ranges each. The requested host and GPU -concurrency is therefore present, although most stream-21 work still cannot -co-reside with the large decoder kernels. - -Clean unprofiled measurements alternated the stable-output lookahead baseline -and this worker-context version at concurrency 32 with 2,048 requests: - -| Order | Version | Mean latency | P50 latency | Makespan | Requests/s | Output tokens/s | -| ---: | --- | ---: | ---: | ---: | ---: | ---: | -| 1 | Stable-output lookahead | 245.145 ms | 232.250 ms | 15.862079 s | 129.113 | 8671.499 | -| 2 | Worker encoder + context | 264.948 ms | 251.086 ms | 17.171612 s | 119.267 | 8091.727 | -| 3 | Stable-output lookahead | 240.438 ms | 225.812 ms | 15.510191 s | 132.042 | 8869.717 | -| 4 | Worker encoder + context | 263.153 ms | 251.669 ms | 17.003768 s | 120.444 | 8199.771 | - -| Two-run average | Stable-output lookahead | Worker encoder + context | Change | -| --- | ---: | ---: | ---: | -| Mean latency | 242.792 ms | 264.051 ms | +8.76% | -| P50 latency | 229.031 ms | 251.378 ms | +9.76% | -| Makespan | 15.686135 s | 17.087690 s | +8.94% | -| Requests/s | 130.578 | 119.856 | -8.21% | -| Output tokens/s | 8770.608 | 8145.749 | -7.12% | - -Separating context from the main mixed batch changes BF16 batch numerics and -therefore a small number of greedy decisions. The worker runs generated 1.18% -more output tokens on average; output-token throughput still regressed by -7.12%, so output length does not explain the result. The approximately -20.6 ms of hidden stream-21 work is smaller than the cost of separate eager -context launches, sampling/finalization, input-staging fences, and GPU resource -contention. This design proves that the work can overlap, but it is not an -end-to-end performance improvement for the measured workload. - -## Suggested experiment order - -1. Add fine-grained NVTX ranges inside the encoder-decoder input fast path. - (Completed.) -2. Gather sampled tokens directly into the persistent input-ID buffer and - avoid restaging stable previous-batch indices. (Completed; approximately - 1.0% throughput improvement across four fully clean interleaved runs.) -3. Reuse KV block-offset staging storage and skip unchanged block-table copies. - (Attempted; unsafe without an explicit device-buffer ownership contract.) -4. Reuse or ping-pong `_prepare_inputs_event` and measure the host-time change. - (Completed; no measurable end-to-end gain.) -5. Measure broader stable encoder-decoder metadata reuse. (Position staging and - bounded block-table staging completed; no verified end-to-end gain.) - -These experiments determine whether the next large gain is stable metadata -reuse or consolidation of device-side input updates. diff --git a/docs/source/models/encoder-decoder.md b/docs/source/models/encoder-decoder.md index a60d937d6855..ab2158ac7de7 100644 --- a/docs/source/models/encoder-decoder.md +++ b/docs/source/models/encoder-decoder.md @@ -421,6 +421,7 @@ llm = LLM( seq_lens=[128, 256, 512, 1024], enable_padding=True, ), + enable_encoder_decoder_mixed_cuda_graph=True, kv_cache_config=KvCacheConfig( free_gpu_memory_fraction=0.8, cross_kv_cache_fraction=0.5, @@ -435,6 +436,13 @@ size, total packed tokens, and maximum sequence length. The limit. With beam search, decoder graph batch sizes must cover the active decoder sequences after beam expansion. +`enable_encoder_decoder_mixed_cuda_graph` is primarily a performance option. It +reduces CPU launch overhead for decoder iterations that mix newly admitted +context requests with ongoing generation requests. The option defaults to +`True`, but becomes effective only when the encoder and decoder graph +configurations produce usable capture shapes. Set it to `False` to disable +mixed graphs while retaining the separate encoder and decoder CUDA graphs. + Passing `EncodeCudaGraphConfig` through `cuda_graph_config` remains unsupported for encoder-decoder models; pass it through `encoder_cuda_graph_config` instead. Piecewise CUDA graphs through `TorchCompileConfig` are also diff --git a/optimize_plan_2.md b/optimize_plan_2.md deleted file mode 100644 index 027eb7b53e30..000000000000 --- a/optimize_plan_2.md +++ /dev/null @@ -1,494 +0,0 @@ - - -# BART PyTorch plan to approach legacy TensorRT performance - -## Objective - -Close the remaining BART continuous-admission performance gap between the -optimized PyTorch path and the legacy TensorRT path without changing output -semantics or regressing latency tails. - -In the matched concurrency-32, 1,024-request benchmark, optimized PyTorch -reached 172.6 requests/s and 10,305 output tokens/s, while legacy TensorRT -reached 200.9 requests/s and 11,289 output tokens/s. PyTorch therefore reached -85.9% of TensorRT request throughput and 91.3% of its output-token throughput. - -Recent measurements around the pending-encoder-future experiments used the -generic encoder-decoder input fallback because the loaded native extension -predated `prepare_encoder_decoder_inputs`. Those absolute results must not -replace the native-fast-path baseline above. Rebuild or load a compatible -extension before evaluating progress against legacy TensorRT. - -## Profile-derived hypothesis - -The remaining gap is not primarily caused by a lack of concurrent encoder and -decoder GPU execution. - -The concurrency-32, 256-request profiles contain: - -| Backend | Encoder passes | Decoder passes | -| --- | ---: | ---: | -| Legacy TensorRT | 180 | 681 | -| PyTorch pending-future design | 17 | 932 | - -The backends can generate somewhat different output lengths because BF16 batch -composition affects close greedy decisions, so the pass counts are not a pure -scheduling comparison. Nevertheless, they show a major structural difference: -legacy TensorRT can run cheap replacement encoder microbatches and replenish -the decoder promptly, while PyTorch must coalesce expensive eager encoder -launches into large waves. The resulting PyTorch decoder batches spend more -iterations partially occupied. - -The legacy profile also showed no encoder/decoder, encoder/sampler, or -decoder/sampler kernel overlap. Its performance does not depend on concurrent -execution across those streams. Earlier closed-batch measurements found about -217.9 ms of PyTorch GPU work versus 226.4 ms of legacy TensorRT GPU work. -Consequently, the primary target is cheap and timely decoder replenishment plus -lower per-iteration orchestration, not maximum GPU kernel overlap. - -## Design 1: occupancy-aware encoder replenishment - -### Motivation - -The current encoder admission policy waits until either: - -- Waiting encoder tokens reach - `batch_wait_max_tokens_ratio * max_num_tokens`; or -- `batch_wait_timeout_iters` expires. - -At concurrency 32, the configured token threshold is approximately 32 -average-length encoder inputs. This gives efficient encoder execution, but it -can leave the decoder substantially underfilled while replacements wait. - -The rejected full-drain policy moved in the wrong direction: it waited until -all decoder work completed before releasing another encoder wave. That added a -head-of-line barrier, reduced request throughput by 26.3%, and increased mean -latency by 35.8%. - -### Proposed policy - -Release waiting encoder requests when any of these conditions is true: - -1. The waiting encoder request count reaches a modest microbatch target. -2. The active generation count falls below a decoder low watermark. -3. The oldest encoder request reaches a maximum iteration deadline. -4. There are no active decoder requests. - -Continue accumulating replacements while the decoder remains sufficiently -full. This preserves encoder efficiency during the steady state but refills -decoder capacity before reaching a long, underfilled tail. - -Initial screening matrix: - -| Parameter | Values | -| --- | --- | -| Encoder microbatch target | 4, 8, 12 | -| Decoder low watermark at concurrency 32 | 20, 24, 28 | -| Maximum wait | 16, 24, 48 iterations | - -Record, in addition to end-to-end performance: - -- Encoder-pass count and batch-size distribution. -- Decoder-pass count. -- Generation batch-size distribution and average occupancy. -- Total encoder, decoder, and sampler GPU busy time. -- Delay from request submission to encoder launch. -- Delay from encoder completion to mixed decoder-context launch. - -### Measured result: rejected before small-encoder optimization - -The policy was implemented with the existing pending-future encoder path and -screened at concurrency 32. The proposed default released eight waiting -encoder requests or released earlier when total decoder occupancy reached 24. -Two less aggressive settings were also tested. - -All runs used the same generic encoder-decoder input path because the installed -native extension predates `prepare_encoder_decoder_inputs`. No Nsight, debug -NVTX, or host timers were enabled. - -| Requests | Policy | Mean latency | P50 latency | P90 latency | P99 latency | Requests/s | -| ---: | --- | ---: | ---: | ---: | ---: | ---: | -| 256 | Existing bounded accumulation | 224.531 ms | 189.549 ms | 385.358 ms | 464.158 ms | 134.155 | -| 256 | Target 8 / low watermark 24 | 320.268 ms | 273.850 ms | 608.803 ms | 689.613 ms | 94.240 | -| 256 | Target 12 / low watermark 20 | 285.727 ms | 240.157 ms | 513.444 ms | 582.509 ms | 105.972 | -| 256 | Target 16 / low watermark 16 | 268.005 ms | 237.374 ms | 459.735 ms | 549.034 ms | 113.317 | - -The documented 1,024-request comparison confirmed the short screen: - -| Policy | Mean latency | P50 latency | P90 latency | P99 latency | Makespan | Requests/s | Output tokens/s | -| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | -| Existing bounded accumulation | 218.841 ms | 195.114 ms | 375.102 ms | 443.033 ms | 7.121912 s | 143.782 | 8568.205 | -| Target 8 / low watermark 24 | 333.685 ms | 279.505 ms | 639.218 ms | 723.822 ms | 10.822124 s | 94.621 | 5651.201 | -| Change | +52.5% | +43.3% | +70.4% | +63.4% | +52.0% | -34.2% | -34.0% | - -The optimized run produced 61,158 output tokens versus 61,022 for the -baseline, a 0.22% difference that does not explain the regression. Earlier -replenishment increases eager encoder launch frequency and GPU disruption more -than fuller decoder batches save. The implementation was removed. - -Occupancy-aware replenishment should be reconsidered only after design 2 makes -small encoder batches materially cheaper. Its experiment then becomes a useful -way to trade the new microbatch cost against decoder occupancy. - -### Cost-aware follow-up - -If a fixed low-watermark policy is directionally positive, replace the static -threshold with a measured cost decision. Maintain an exponentially weighted -estimate of encoder cost by microbatch size and decoder-step cost by graph -bucket. Dispatch a waiting encoder microbatch when its estimated cost is lower -than the decoder work expected to be saved by replenishing the open lanes. - -The initial policy should remain simple until the profile counters prove that -decoder occupancy predicts end-to-end performance. - -## Design 2: make small encoder batches cheap - -Occupancy-aware admission can succeed only if the additional encoder launches -are inexpensive enough. The desired behavior is not necessarily to reproduce -all 180 legacy encoder passes, but to permit more frequent batches of roughly -1--8 requests without returning to the previously measured batch-one eager -encoder collapse. - -### Whole-encoder CUDA graphs for microbatches - -Prioritize whole-encoder graph capture for batch sizes 1, 2, 4, and 8. Use -exact-shape graph keys initially: - -```text -(batch size, packed token count, maximum sequence length, exact cu_seqlens layout) -``` - -Exact keys avoid the dummy sequence padding and numerical changes observed in -the piecewise graph experiment. Capture graphs lazily and use an LRU limit to -bound static-buffer and graph-executable memory. - -The previous bounded whole-encoder graph prototype improved mean latency by -3.8%, 3.8%, and 6.4% at concurrencies 32, 64, and 128. Its relative benefit -may be larger for the small encoder batches most affected by eager Python and -CUDA launch overhead. - -### Supporting work - -- Pack encoder token and position inputs in native code. -- Reuse stable pinned and device input/output storage for each graph key. -- Preserve the fused residual/layer-normalization and GELU implementations. -- Keep variable-length attention numerically identical; do not pad real - sequences merely to reduce the number of graph keys. -- Benchmark graph lookup, input staging, replay, and output publication - separately from capture. - -Once small encoder execution is faster, rerun the occupancy-aware policy -matrix and progressively lower its microbatch target. - -### Measured result: graph replay helps, but small-batch execution still loses - -An experimental implementation added independent whole-encoder graphs without -replacing the decoder CUDA-graph configuration. It uses exact keys containing -the full sequence-length layout, lazy thread-local capture on the encoder -worker, a 64-entry LRU, and graph-resident pinned/device staging. Real -sequences and token counts are not padded. Runtime capture must be -thread-local: CUDA's default global capture mode otherwise rejects unrelated -sampler synchronization on the main executor thread. - -The first attribution used the same admission schedule on both sides. At -batch one, exact graphs reduced mean latency by 8.0% relative to eager -execution, but the schedule was still far slower than the existing -large-batch policy because batch-one encoder kernels lose too much GPU -efficiency. At batch eight, where this cyclic workload repeatedly reuses two -exact eight-request layouts, graph replay improved the 256-request screen only -slightly: - -| Requests | Encoder schedule | Encoder execution | Mean latency | Requests/s | -| ---: | --- | --- | ---: | ---: | -| 256 | Target 1 / no low watermark | Eager | 509.427 ms | 61.176 | -| 256 | Target 1 / no low watermark | Exact graph | 468.626 ms | 66.505 | -| 256 | Target 8 / no low watermark | Eager | 243.296 ms | 119.891 | -| 256 | Target 8 / no low watermark | Exact graph | 241.235 ms | 122.983 | - -The matched concurrency-32, 1,024-request comparison was negative. Enabling -graphs alone while retaining the existing bounded-accumulation policy did not -help because steady-state encoder batches were normally larger than eight; -the remaining small tail batches paid capture cost without enough replay. -Forcing the best screened target-eight schedule increased encoder frequency -enough to outweigh its slightly cheaper launches. - -| Policy | Mean latency | P50 latency | P90 latency | P99 latency | Makespan | Requests/s | Output tokens/s | -| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | -| Existing bounded accumulation | 215.505 ms | 189.828 ms | 368.375 ms | 454.794 ms | 7.046579 s | 145.319 | 8679.106 | -| Exact graphs, existing admission | 219.949 ms | 192.038 ms | 382.436 ms | 459.412 ms | 7.210709 s | 142.011 | 8494.866 | -| Exact graphs, target 8 / no low watermark | 240.486 ms | 199.376 ms | 456.963 ms | 544.848 ms | 7.884365 s | 129.877 | 7758.392 | -| Target-8 change versus baseline | +11.6% | +5.0% | +24.0% | +19.8% | +11.9% | -10.6% | -10.6% | - -The three runs produced 61,158, 61,254, and 61,170 output tokens, -respectively, so output-count variation does not explain the latency result. -As in the design-1 screen, these measurements used the generic -encoder-decoder input path because the loaded native extension predates -`prepare_encoder_decoder_inputs`. - -Exact whole-encoder graphs therefore reduce launch overhead for a fixed small -batch, but not enough to make frequent PyTorch encoder replenishment -competitive. Do not enable this policy by default. A future attempt needs a -material reduction in small-batch kernel time or a graph strategy that -preserves larger encoder batches while removing their host launch overhead. - -### Measured result: `(batch, total tokens, max bucket)` keys - -The cyclic benchmark produces 59 unique `(B,T,S)` triples for -`B in {1,2,4,8}` when maximum sequence length is bucketed in 64-token -increments. The implementation accepts those triples as an allowlist, captures -each graph on its first real batch, and retains up to 64 graphs in an LRU. -Encoder graph metadata uses a separate buffer arena from the concurrently -executing decoder graphs, and shared host staging is retired before the next -replay updates it. - -Individual sequence lengths are runtime graph inputs, not part of the key. -TRTLLM attention rebuilds `cu_seqlens` and padding offsets on the GPU from the -current device sequence-length buffer. `B`, `T`, and the maximum-length bucket -keep tensor extents, workspace requirements, and FMHA launch dimensions stable. - -The earlier cross-layout hang came from inconsistent capture warmup metadata, -not a requirement for exact layouts. Graph metadata initialized its device -sequence lengths to ones, preparation updated only the stable host buffer, and -the H2D copy occurred for the first time inside graph capture. Warmup therefore -ran with real packed-token counts on the host but all-one sequence lengths on -the GPU. Capture now stages input IDs, position IDs, and sequence lengths to -their device buffers before warmup. The graph still captures the H2D copies for -subsequent replays. Exact-layout rejection is removed, so all layouts sharing a -`(B,T,S)` key reuse the same graph. - -Matched concurrency-32, 1,024-request results: - -| Backend / encoder mode | Mean latency | P50 latency | P90 latency | P99 latency | Makespan | Requests/s | Output tokens/s | -| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | -| PyTorch eager, target 8 | 245.536 ms | 204.663 ms | 467.781 ms | 521.870 ms | 8.058071 s | 127.078 | 7573.525 | -| PyTorch `(B,T,S)` graph with exact-layout guard | 234.166 ms | 194.920 ms | 453.589 ms | 511.825 ms | 7.685696 s | 133.235 | 7947.751 | -| PyTorch `(B,T,S)` graph with cross-layout replay | 233.765 ms | 192.882 ms | 447.072 ms | 508.841 ms | 7.682285 s | 133.294 | 7945.943 | -| PyTorch exact batch-8 admission at decoder occupancy 24 | 225.955 ms | 187.235 ms | 433.918 ms | 492.370 ms | 7.373542 s | 138.875 | 8274.721 | -| PyTorch exact batch 8, serialized encoder/decoder forward | 225.123 ms | 183.728 ms | 435.884 ms | 485.110 ms | 7.342011 s | 139.471 | 8317.203 | -| PyTorch exact batch 8, overlapped encoder plus mixed decoder graphs | 159.100 ms | 130.335 ms | 306.821 ms | 341.215 ms | 5.229246 s | 195.822 | 11673.959 | -| Legacy TensorRT | 163.359 ms | 127.874 ms | 344.281 ms | 425.151 ms | 5.363881 s | 190.907 | 10783.982 | - -Cross-layout replay did not materially change end-to-end performance versus -the exact-layout guard: mean latency fell by 0.2%, request throughput rose by -0.04%, and makespan fell by 0.04%, all within run-to-run variation. It removes -an invalid hidden key and increases graph coverage, but graph eligibility and -scheduling still dominate the aggregate result. The four PyTorch runs -generated 61,028, 61,084, 61,043, and 61,014 output tokens, respectively; -legacy generated 57,844, so request throughput and latency are the cleaner -cross-backend comparisons. - -The corrected Nsight profile contains 121 encoder batches in the timed -generation window. Eighty-eight batches executed a CUDA graph: 11 captured a -new key and immediately launched it, while 77 were cache-hit replays. -Thirty-three batches used eager fallback. - -The eager fallbacks were caused by scheduler overfill rather than missing -graphs for a supported batch size. Once at least eight encoder requests were -available, admission returned the entire scheduler result, producing batches -of 9 through 12 in steady state plus startup and tail batches of 31 and 5. -Exact batch-8 admission now returns only the first eight requests once decoder -occupancy reaches 24 or lower and leaves excess requests eligible for the next -iteration. A deadline releases a smaller supported power-of-two tail. - -The clean exact-batch-8 run reduced mean latency by 3.3% and increased request -throughput by 4.2% versus cross-layout replay with uncapped admission. Its -profile contained exactly 128 batch-8 encoder steps for 1,024 requests. All -128 were cache-hit graph replays; there were no timed captures and no eager -fallbacks. Encoder dispatch occurred alongside 24 generation requests in 58 -steps. Other steady-state dispatches occurred after the decoder completed -multiple requests between scheduling passes and crossed directly below 24. - -Serializing encoder and decoder forward launches produced no measurable -end-to-end change. The executor launched the encoder on its dedicated stream, -blocked on its completion event, and only then entered decoder forward. -Compared with asynchronous exact-batch-8 admission, mean latency and makespan -were 0.4% lower, request throughput was 0.4% higher, P90 was 0.5% higher, and -the other percentiles moved by less than 2%. This mixed movement is within -single-run variation and indicates that encoder/decoder GPU concurrency was -not contributing material throughput in this workload. - -The final design restores asynchronous encoder/decoder overlap and captures -whole-model CUDA graphs for mixed decoder batches. Its graph key contains the -padded decoder batch size, exact decoder context-query extents, and packed -encoder-hidden-state row count. Cross-attention sequence layouts remain -runtime metadata rather than graph keys. Encoder hidden states are copied into -one graph-stable buffer before replay, so a replacement encoder output never -leaves a captured pointer referring to request-owned storage. - -Mixed graphs are never captured on a live request. Startup first warms all -reachable batch-8 and paired-batch-16 replenishment shapes to finish sizing -the shared attention workspace, then captures the same 61 mixed shapes on a -second pass. An unseen shape falls back to eager execution. This avoids both -repeating live KV-cache writes during capture and invalidating older graph -pointers through a late workspace resize. - -Against the preceding best overlapped exact-batch-8 result, mixed decoder -graphs reduced mean latency by 29.6%, P50 by 30.4%, P90 by 29.3%, P99 by -30.7%, and makespan by 29.1%. Request throughput and output-token throughput -increased by 41.0% and 41.1%, respectively. It also slightly exceeded legacy -TensorRT in this run: mean latency was 2.6% lower, makespan was 2.5% lower, -and request throughput was 2.6% higher. Output-token throughput is not a clean -cross-backend comparison because legacy produced fewer output tokens. - -The final Nsight timed window contained 126 mixed decoder steps, and all 126 -launched a decoder CUDA graph. All 2,205 generation-only decoder steps and all -128 encoder steps also replayed graphs. The single context-only decoder step -remained eager, and there were no CUDA graph captures in the timed window. -Typical mixed `_forward_step` ranges fell from roughly 18--20 ms in the eager -profile to roughly 1.9--2.1 ms with replay. - -The final PyTorch run produced 61,046 output tokens versus 61,014 in the -preceding best run, a 0.05% difference. BF16 continuous admission can change -close greedy decisions when faster completion changes batch composition, so -the output hash is not expected to remain fixed across scheduling changes. - -Raw logs: - -- `/tmp/bart-bts-graphs-eager-target8-pytorch-c32-r1024.log` -- `/tmp/bart-bts-graphs-safe-allow59-pytorch-c32-r1024.log` -- `/tmp/bart-bts-graphs-legacy-c32-r1024.log` -- `/tmp/bart-bts-graphs-cross-layout-fixed-allow59-c32-r1024.nsys-rep` -- `/tmp/bart-bts-graphs-cross-layout-fixed-allow59-c32-r1024.sqlite` -- `/tmp/bart-bts-graphs-exact8-low24-pytorch-c32-r1024.log` -- `/tmp/bart-bts-graphs-exact8-low24-c32-r1024.nsys-rep` -- `/tmp/bart-bts-graphs-exact8-low24-c32-r1024.sqlite` -- `/tmp/bart-bts-graphs-exact8-low24-serialized-pytorch-c32-r1024.log` -- `/tmp/bart-bts-graphs-exact8-low24-serialized-c32-r1024.nsys-rep` -- `/tmp/bart-mixed-decoder-graphs-final-overlap-c32-r1024.log` -- `/tmp/bart-mixed-decoder-graphs-final-overlap-c32-r1024.nsys-rep` -- `/tmp/bart-mixed-decoder-graphs-final-overlap-c32-r1024.sqlite` - -## Design 3: a decoder supergraph - -The current decoder CUDA graph captures model execution, but input staging, -attention-metadata updates, sampling, request updates, stream handoffs, and -completion processing remain separate. Legacy TensorRT hides most of this work -behind one engine enqueue. - -For the qualified single-rank, single-beam, greedy BART path, maintain a -persistent device-side lane table containing: - -- Request and sequence-slot identifiers. -- Current token and position. -- Self-KV block descriptors and lengths. -- Cross-KV descriptors and encoder-output pointers. -- Active, EOS, and length-complete state. - -Capture one decoder supergraph containing: - -```text -sampled-token gather -→ position and KV-length update -→ decoder model -→ greedy argmax -→ next-token scatter -→ EOS and length-completion update -``` - -The next decoder iteration should consume the sampled token directly from -device storage. The CPU should receive a compact completion record and update -only lane admissions, removals, and externally visible request state. Batch -membership changes should be expressed as deltas to persistent lane metadata -rather than a complete rebuild of every active request. - -Use double-buffered launch packets so the batch-static portion of the next -iteration can be prepared while the current graph runs. Once sampled tokens -are available, the launch-critical path should consist of a compact device -update and one graph replay. - -Qualification must initially exclude streaming, beam search, speculative -decoding, guided decoding, LoRA, cache reuse, attention data parallelism, -disaggregated transfer, and other features that require the general path. - -## Design 4: two-token decoder graph unrolling - -After the decoder supergraph is correct and beneficial, capture two dependent -greedy decoder steps in one graph replay: - -```text -decoder step N -→ greedy token N -→ device state update -→ decoder step N+1 -→ greedy token N+1 -``` - -A device finish mask must prevent a token after EOS or the length limit from -being appended. A lane that completes after the first step may still perform -dummy model computation during the second step. - -Start with two steps only. Larger unrolling would save more host launches but -would also delay new-request admission and waste more computation on completed -lanes. This experiment also requires sufficient KV capacity to be reserved -before replay. - -## GPU work ordering - -Do not treat encoder/decoder GPU overlap as a goal by itself. The legacy -profile serialized its encoder, decoder, and sampler kernels, and the PyTorch -pending-future design achieved only about 1% encoder/decoder kernel overlap. - -For occupancy-aware microbatches, compare: - -1. Encoder execution ordered after the current decoder step and before a later - decoder step. -2. The existing independent encoder stream. - -Prefer deterministic serialized ordering if concurrent streams introduce -resource contention. The encoder host worker can still prepare and submit work -asynchronously even if CUDA events serialize its GPU execution. - -Keep decoder context requests in the main mixed decoder batch. Moving encoder -plus decoder-context execution to a side worker produced 20.6 ms of real GPU -overlap but regressed request throughput by 8.2% because separate eager context -launches, sampling, staging fences, and GPU contention cost more than the -hidden work saved. - -## Experiment order - -1. Rebuild or load the compatible native extension and reproduce the optimized - PyTorch and legacy TensorRT baselines with the same workload. -2. Add lightweight counters for encoder-pass sizes, decoder occupancy, and - admission delays. -3. Screen the occupancy-aware encoder policy without changing encoder kernels. -4. Implement exact-shape whole-encoder CUDA graphs for small microbatches. -5. Rerun the occupancy policy matrix and select the best cost/occupancy point. -6. Prototype the qualified decoder supergraph. -7. If a meaningful host bubble remains, test two-token graph unrolling. -8. Confirm every promising result with alternating, fully unprofiled long - runs, followed by Nsight profiling for causal attribution. - -## Validation criteria - -Use the checked-in continuous-admission workload and report both request and -executed-token throughput because natural EOS can differ between batch -compositions. - -A change should be retained only if: - -- It improves both request throughput and output-token throughput outside - observed run-to-run noise. -- Mean, P50, P90, and P99 latency do not show an unacceptable redistribution. -- Natural-EOS and length-stop behavior remains valid. -- Repeated runs show no request/output association errors. -- General configurations fall back without semantic changes. -- The native fast path, rather than the generic compatibility fallback, is - active in the comparison. - -## Designs not to revisit without new evidence - -- Draining every decoder request before admitting another encoder wave. -- Moving decoder-context execution onto the encoder worker. -- Decoder-versus-encoder CUDA stream-priority changes. -- Additional tuning of only the existing fixed token and iteration thresholds. -- Maximizing encoder/decoder kernel concurrency as an independent objective. -- Retaining finished decoder rows as permanent graph-padding lanes. - -The most useful immediate experiment is occupancy-aware replenishment. It -directly tests whether the legacy path's frequent replacement encoding and -lower decoder-pass count explain the remaining gap. If it improves decoder -occupancy but loses the gain to encoder cost, that result becomes a precise -performance requirement for the small-batch whole-encoder graph work. diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index 9d1960c239f4..97819bfe4e0b 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -1,11 +1,8 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - import bisect import contextlib from dataclasses import dataclass -from typing import (Any, Callable, Dict, Iterator, List, Optional, Sequence, - Tuple, TypeAlias) +from typing import (Any, Callable, Dict, Iterator, List, Optional, Tuple, + TypeAlias) import torch @@ -116,9 +113,6 @@ class CUDAGraphRunnerConfig: dynamic_draft_len_mapping: Optional[Dict[int, int]] = None sparse_attention_config: Optional[BaseSparseAttentionConfig] = None enable_encoder_decoder_mixed_cuda_graph: bool = False - encoder_hidden_size: int = 0 - dtype: Optional[torch.dtype] = None - encoder_decoder_mixed_cuda_graph_encoder_token_counts: Tuple[int, ...] = () class CUDAGraphRunner: @@ -192,15 +186,33 @@ def _create_shared_static_tensors(self): self.shared_static_tensors[ "mrope_delta_read_seq_slots"] = torch.zeros( (max_total_tokens, ), device="cuda", dtype=torch.long) - if self.enable_encoder_decoder_mixed_cuda_graph: - if self.config.encoder_hidden_size <= 0 or self.config.dtype is None: - raise ValueError("Mixed encoder-decoder CUDA graphs require " - "encoder_hidden_size and dtype.") - self.shared_static_tensors["encoder_hidden_states"] = torch.empty( - (self.config.max_num_tokens, self.config.encoder_hidden_size), - device="cuda", - dtype=self.config.dtype, - ) + + def _get_static_encoder_hidden_states( + self, + encoder_hidden_states: torch.Tensor, + num_encoder_tokens: int, + *, + allow_allocate: bool, + ) -> torch.Tensor: + """Return the stable mixed-graph encoder input, allocating at warmup.""" + if encoder_hidden_states.ndim != 2: + raise RuntimeError( + "Mixed encoder-decoder CUDA graphs require rank-2 packed " + "encoder hidden states.") + + static_encoder_hidden_states = self.shared_static_tensors.get( + "encoder_hidden_states") + if static_encoder_hidden_states is None: + if not allow_allocate: + raise RuntimeError( + "Mixed encoder-decoder CUDA graph replay requires the " + "encoder hidden-state buffer initialized during warmup.") + static_encoder_hidden_states = encoder_hidden_states.new_empty( + (num_encoder_tokens, encoder_hidden_states.shape[1])) + self.shared_static_tensors[ + "encoder_hidden_states"] = static_encoder_hidden_states + + return static_encoder_hidden_states[:num_encoder_tokens] def _is_mixed_encoder_decoder_batch(self, batch: ScheduledRequests) -> bool: return (self.enable_encoder_decoder_mixed_cuda_graph @@ -517,8 +529,12 @@ def capture(self, if encoder_hidden_states is None: raise RuntimeError("Mixed encoder-decoder CUDA graph capture " "requires encoder hidden states.") - static_encoder_hidden_states = self.shared_static_tensors[ - "encoder_hidden_states"][:num_encoder_tokens] + static_encoder_hidden_states = ( + self._get_static_encoder_hidden_states( + encoder_hidden_states, + num_encoder_tokens, + allow_allocate=True, + )) actual_num_encoder_tokens = encoder_hidden_states.shape[0] if actual_num_encoder_tokens > num_encoder_tokens: raise RuntimeError( @@ -622,8 +638,12 @@ def replay(self, key: KeyType, "Mixed encoder-decoder CUDA graph replay received " f"{actual_num_encoder_tokens} encoder tokens for a " f"{num_encoder_tokens}-token graph.") - static_encoder_hidden_states = static_tensors[ - "encoder_hidden_states"][:num_encoder_tokens] + static_encoder_hidden_states = ( + self._get_static_encoder_hidden_states( + encoder_hidden_states, + num_encoder_tokens, + allow_allocate=False, + )) static_encoder_hidden_states[:actual_num_encoder_tokens].copy_( encoder_hidden_states) static_encoder_hidden_states[actual_num_encoder_tokens:].zero_() @@ -847,7 +867,7 @@ class EncoderCUDAGraphRunnerConfig: max_num_tokens: int max_seq_len: int cuda_graph_mem_pool: Any - encoder_decoder_capture_keys: Optional[Sequence[EncoderKeyType]] = None + is_encoder_decoder: bool = False class EncoderCUDAGraphRunner: @@ -874,8 +894,13 @@ def __init__(self, config: EncoderCUDAGraphRunnerConfig): self.supported_num_tokens = sorted(config.cuda_graph_num_tokens) self.max_supported_num_tokens = config.max_cuda_graph_num_tokens self.supported_seq_lens = sorted(config.cuda_graph_seq_lens) - self.is_encoder_decoder = config.encoder_decoder_capture_keys is not None - self.capture_keys = frozenset(config.encoder_decoder_capture_keys or ()) + self.is_encoder_decoder = config.is_encoder_decoder + self.capture_keys: frozenset[EncoderKeyType] = frozenset() + self._capture_sequence_lengths: Dict[EncoderKeyType, List[int]] = {} + if self.is_encoder_decoder: + self._capture_sequence_lengths = ( + self._build_encoder_decoder_capture_layouts()) + self.capture_keys = frozenset(self._capture_sequence_lengths) self._capture_keys_by_batch_size: Dict[int, List[EncoderKeyType]] = {} for key in sorted(self.capture_keys): self._capture_keys_by_batch_size.setdefault(key[0], []).append(key) @@ -950,6 +975,44 @@ def _round_up(value: int, supported: List[int]) -> int: return 0 return supported[idx] + @staticmethod + def build_capture_sequence_lengths(batch_size: int, num_tokens: int, + max_seq_len: int) -> Optional[List[int]]: + """Build a real sequence layout for a configured encoder bucket.""" + if (batch_size <= 0 or num_tokens < batch_size + or num_tokens > batch_size * max_seq_len): + return None + + if batch_size == 1: + return [num_tokens] + + if num_tokens >= max_seq_len + batch_size - 1: + remaining_tokens = num_tokens - max_seq_len + base, extra = divmod(remaining_tokens, batch_size - 1) + return ([max_seq_len] + [base + 1] * extra + [base] * + (batch_size - 1 - extra)) + + return [num_tokens - batch_size + 1] + [1] * (batch_size - 1) + + def _build_encoder_decoder_capture_layouts( + self) -> Dict[EncoderKeyType, List[int]]: + """Derive reachable encoder-decoder capture keys via get_graph_key.""" + capture_layouts: Dict[EncoderKeyType, List[int]] = {} + for batch_size in self.supported_batch_sizes: + for num_tokens in self.supported_num_tokens: + for max_seq_len in self.supported_seq_lens: + sequence_lengths = self.build_capture_sequence_lengths( + batch_size, num_tokens, max_seq_len) + if sequence_lengths is None: + continue + + key, _, is_valid = self.get_graph_key( + {"seq_lens": sequence_lengths}) + if is_valid: + capture_layouts.setdefault(key, sequence_lengths) + + return capture_layouts + def _get_dynamic_capture_key( self, batch_size: int, @@ -981,44 +1044,9 @@ def _get_dynamic_capture_key( def get_capture_warmup_sequence_lengths( self, key: EncoderKeyType) -> Optional[List[int]]: - """Build a real sequence layout whose padded graph key is ``key``. - - A larger max-sequence bucket can be dominated by a smaller bucket at - the same batch size and token count. Such keys can never be selected - for a real batch and return ``None``. - """ - if key not in self.capture_keys: - return None - - batch_size, num_tokens, max_seq_len = key - previous_max_seq_len = max( - (candidate[2] for candidate in self._capture_keys_by_batch_size.get( - batch_size, []) - if candidate[1] == num_tokens and candidate[2] < max_seq_len), - default=0, - ) - actual_max_seq_len = max( - (num_tokens + batch_size - 1) // batch_size, - previous_max_seq_len + 1, - ) - if (actual_max_seq_len > max_seq_len - or actual_max_seq_len + batch_size - 1 > num_tokens): - return None - - if batch_size == 1: - sequence_lengths = [num_tokens] - else: - remaining_tokens = num_tokens - actual_max_seq_len - base, extra = divmod(remaining_tokens, batch_size - 1) - sequence_lengths = [actual_max_seq_len] - sequence_lengths.extend([base + 1] * extra) - sequence_lengths.extend([base] * (batch_size - 1 - extra)) - - selected_key, _, is_valid = self.get_graph_key( - {"seq_lens": sequence_lengths}) - if not is_valid or selected_key != key: - return None - return sequence_lengths + """Return the representative sequence layout for a capture key.""" + sequence_lengths = self._capture_sequence_lengths.get(key) + return list(sequence_lengths) if sequence_lengths is not None else None def _get_valid_graph_key(self, batch_size: int, num_tokens: int, max_seq_len: int) -> EncoderKeyType: diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index 12e786c0b393..8389304f96d6 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - from copy import copy, deepcopy from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union, cast diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 5ed4c0a66564..da7758c55f48 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -257,17 +257,6 @@ def _filter_cuda_graph_seq_lens(cuda_graph_seq_lens: list[int], return result -def _build_encoder_decoder_cuda_graph_keys( - batch_sizes: Sequence[int], num_tokens: Sequence[int], - seq_lens: Sequence[int]) -> List[Tuple[int, int, int]]: - """Build all geometrically valid encoder graph bucket combinations.""" - return sorted({(batch_size, total_tokens, max_seq_len) - for batch_size in batch_sizes - for total_tokens in num_tokens - for max_seq_len in seq_lens - if batch_size <= total_tokens <= batch_size * max_seq_len}) - - _DEEP_GEMM_PDL_CONFIGURED = False @@ -729,38 +718,13 @@ def __init__( self._max_cuda_graph_seq_len = (self._cuda_graph_seq_lens[-1] if self._cuda_graph_seq_lens else 0) - encoder_max_batch_size = self.llm_args.encoder_max_batch_size - encoder_decoder_microbatch_cuda_graph_enabled = (os.environ.get( - "TLLM_ENCODER_DECODER_MICROBATCH_CUDA_GRAPH_ENABLED", "1") == "1") - self._enable_encoder_decoder_microbatch_cuda_graph = ( - self._is_encoder_decoder_model() - and encoder_max_batch_size is not None - and self.encoder_cuda_graph_config is not None - and bool(self._cuda_graph_num_tokens) - and bool(self._cuda_graph_seq_lens) - and encoder_decoder_microbatch_cuda_graph_enabled) - encoder_decoder_graph_batch_sizes = ( - self._encoder_cuda_graph_batch_sizes - if self._enable_encoder_decoder_microbatch_cuda_graph else []) - encoder_decoder_graph_keys = _build_encoder_decoder_cuda_graph_keys( - encoder_decoder_graph_batch_sizes, - self._cuda_graph_num_tokens, - self._cuda_graph_seq_lens, - ) - encoder_decoder_graph_keys = [ - key for key in encoder_decoder_graph_keys - if key[0] in encoder_decoder_graph_batch_sizes and key[1] <= - self.encoder_max_num_tokens and key[2] in self._cuda_graph_seq_lens - ] - mixed_graph_encoder_batch_size = (encoder_decoder_graph_batch_sizes[-1] - if encoder_decoder_graph_batch_sizes - else 0) - mixed_graph_encoder_token_counts = tuple( - sorted({ - total_tokens - for batch_size, total_tokens, _ in encoder_decoder_graph_keys - if batch_size == mixed_graph_encoder_batch_size - })) + # Encoder CUDA graph detection for encoder-decoder models and encoder-only models. + # Decode configs do not define these encoder-specific bucket fields. + use_encoder_cuda_graph = ((self._is_encoder_decoder_model() + or self._is_encode_only()) + and self.encoder_cuda_graph_config is not None + and bool(self._cuda_graph_num_tokens) + and bool(self._cuda_graph_seq_lens)) self._dynamic_draft_len_mapping = self._compute_dynamic_draft_len_mapping( ) @@ -837,15 +801,41 @@ def __init__( self.lora_model_config: Optional[LoraModelConfig] = None self._trtllm_gen_jit_warmup = False - use_encoder_decoder_graph = ( - self._enable_encoder_decoder_microbatch_cuda_graph - and bool(encoder_decoder_graph_batch_sizes)) + # Create the encoder runner first. For encoder-decoder models it derives + # every reachable startup capture key through get_graph_key(). + encoder_graph_batch_sizes = self._encoder_cuda_graph_batch_sizes + encoder_graph_max_batch_size = (encoder_graph_batch_sizes[-1] + if encoder_graph_batch_sizes else 0) + encoder_graph_max_num_tokens = self._max_cuda_graph_num_tokens + encoder_cuda_graph_runner_config = EncoderCUDAGraphRunnerConfig( + use_cuda_graph=use_encoder_cuda_graph, + cuda_graph_padding_enabled=( + self._encoder_cuda_graph_padding_enabled), + cuda_graph_batch_sizes=encoder_graph_batch_sizes, + cuda_graph_num_tokens=self._cuda_graph_num_tokens, + cuda_graph_seq_lens=self._cuda_graph_seq_lens, + max_cuda_graph_batch_size=encoder_graph_max_batch_size, + max_cuda_graph_num_tokens=encoder_graph_max_num_tokens, + max_num_tokens=self.encoder_max_num_tokens, + max_seq_len=self.max_seq_len, + cuda_graph_mem_pool=self._cuda_graph_mem_pool, + is_encoder_decoder=self._is_encoder_decoder_model(), + ) + self.encoder_cuda_graph_runner = EncoderCUDAGraphRunner( + encoder_cuda_graph_runner_config) + + # Once encoder CUDA graphs are usable, enable mixed decoder graphs by + # default unless the user explicitly opts out. + encoder_decoder_cuda_graph_enabled = ( + self.encoder_cuda_graph_runner.enabled + and self.encoder_cuda_graph_runner.is_encoder_decoder + and bool(self.encoder_cuda_graph_runner.capture_keys)) enable_encoder_decoder_mixed_cuda_graph = ( - use_encoder_decoder_graph - and self.cuda_graph_config is not None and os.environ.get( - "TLLM_ENCODER_DECODER_MIXED_CUDA_GRAPH_ENABLED", "1") == "1") + encoder_decoder_cuda_graph_enabled + and self.cuda_graph_config is not None + and self.llm_args.enable_encoder_decoder_mixed_cuda_graph) - # Create config and runner + # Create decoder CUDA graph config and runner. cuda_graph_runner_config = CUDAGraphRunnerConfig( use_cuda_graph=(not self._is_encode_only and self.cuda_graph_config is not None), @@ -871,43 +861,9 @@ def __init__( sparse_attention_config=self.sparse_attention_config, enable_encoder_decoder_mixed_cuda_graph=( enable_encoder_decoder_mixed_cuda_graph), - encoder_hidden_size=(self._get_enc_dec_hidden_size() - if enable_encoder_decoder_mixed_cuda_graph else - 0), - dtype=(self.dtype - if enable_encoder_decoder_mixed_cuda_graph else None), - encoder_decoder_mixed_cuda_graph_encoder_token_counts=( - mixed_graph_encoder_token_counts), ) self.cuda_graph_runner = CUDAGraphRunner(cuda_graph_runner_config) - # Create Encoder CUDA graph config and runner. - encoder_graph_batch_sizes = self._encoder_cuda_graph_batch_sizes - encoder_graph_max_batch_size = (encoder_graph_batch_sizes[-1] - if encoder_graph_batch_sizes else 0) - encoder_graph_max_num_tokens = self._max_cuda_graph_num_tokens - encoder_cuda_graph_runner_config = EncoderCUDAGraphRunnerConfig( - use_cuda_graph=(use_encoder_decoder_graph - or (self._is_encode_only - and self.encoder_cuda_graph_config is not None - and bool(self._cuda_graph_num_tokens) - and bool(self._cuda_graph_seq_lens))), - cuda_graph_padding_enabled=( - self._encoder_cuda_graph_padding_enabled), - cuda_graph_batch_sizes=encoder_graph_batch_sizes, - cuda_graph_num_tokens=self._cuda_graph_num_tokens, - cuda_graph_seq_lens=self._cuda_graph_seq_lens, - max_cuda_graph_batch_size=encoder_graph_max_batch_size, - max_cuda_graph_num_tokens=encoder_graph_max_num_tokens, - max_num_tokens=self.encoder_max_num_tokens, - max_seq_len=self.max_seq_len, - cuda_graph_mem_pool=self._cuda_graph_mem_pool, - encoder_decoder_capture_keys=(encoder_decoder_graph_keys if - use_encoder_decoder_graph else None), - ) - self.encoder_cuda_graph_runner = EncoderCUDAGraphRunner( - encoder_cuda_graph_runner_config) - # Initialize CUDA Graph LoRA manager if LoRA is enabled self.cuda_graph_lora_manager: Optional[CudaGraphLoraManager] = None @@ -2181,23 +2137,35 @@ def _capture_mixed_encoder_decoder_cuda_graphs( max_encoder_output_len = self._get_max_encoder_output_len( resource_manager) - encoder_token_counts = ( - runner.config.encoder_decoder_mixed_cuda_graph_encoder_token_counts - or (8 * max_encoder_output_len, )) - context_shapes = [(8, token_count) - for token_count in encoder_token_counts] - if runner.max_supported_batch_size > 16: - paired_token_counts = sorted({ + context_shapes = {(batch_size, total_tokens) + for batch_size, total_tokens, _ in + self.encoder_cuda_graph_runner.capture_keys} + if not context_shapes: + logger.warning("Skipping mixed encoder-decoder CUDA graph capture: " + "no encoder CUDA graph shapes were captured.") + return + + max_encoder_batch_size = max(batch_size + for batch_size, _ in context_shapes) + max_batch_token_counts = { + total_tokens + for batch_size, total_tokens in context_shapes + if batch_size == max_encoder_batch_size + } + paired_context_count = 2 * max_encoder_batch_size + if runner.max_supported_batch_size > paired_context_count: + paired_token_counts = { first + second - for first in encoder_token_counts - for second in encoder_token_counts - }) - context_shapes.extend( - (16, token_count) for token_count in paired_token_counts) + for first in max_batch_token_counts + for second in max_batch_token_counts + } + context_shapes.update((paired_context_count, token_count) + for token_count in paired_token_counts) operation = ("warmup" if runner.is_warmup_only else "capture") hidden_size = self._get_enc_dec_hidden_size() - for num_contexts, total_encoder_tokens in context_shapes: + for num_contexts, total_encoder_tokens in sorted( + context_shapes, key=lambda shape: shape[1], reverse=True): if total_encoder_tokens > num_contexts * max_encoder_output_len: continue base_encoder_len, remainder = divmod(total_encoder_tokens, @@ -6417,43 +6385,16 @@ def _create_encoder_warmup_inputs( """Synthesize an inputs dict that will bucket exactly at (batch_size, num_tokens, max_seq_len). - Uses two distribution strategies: - - Case A: `total >= max_seq_len + (batch_size - 1)` — one request at - `max_seq_len` tokens, remaining tokens distributed evenly across - the other `batch_size - 1` requests. - - Case B: `total < max_seq_len + (batch_size - 1)` — one request of - `total - (batch_size - 1)` tokens, the rest at 1 token each. - Returns None for infeasible combinations (e.g., batch_size <= 0). """ - if batch_size <= 0 or num_tokens <= 0 or max_seq_len <= 0: + lengths = ( + self.encoder_cuda_graph_runner.build_capture_sequence_lengths( + batch_size, num_tokens, max_seq_len)) + if lengths is None: return None - total = min(num_tokens, batch_size * max_seq_len) - - if batch_size == 1: - lengths = [total] - elif total >= max_seq_len + batch_size - 1: - # Case A - remaining = total - max_seq_len - base = remaining // (batch_size - 1) - extra = remaining % (batch_size - 1) - lengths = [max_seq_len] - lengths += [base + 1] * extra + [base] * (batch_size - 1 - extra) - else: - # Case B - first_len = total - (batch_size - 1) - lengths = [first_len] + [1] * (batch_size - 1) - - # Sanity: every length must be >= 1. - if any(length <= 0 for length in lengths): - return None - - actual_num_tokens = sum(lengths) - input_ids = [0] * actual_num_tokens - inputs: Dict[str, Any] = { - 'input_ids': input_ids, + 'input_ids': [0] * sum(lengths), 'seq_lens': lengths, } return inputs diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py index 6b36cf898135..caa8e3cb3de1 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - import dataclasses import inspect from abc import ABC, abstractmethod diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 1fcbfd9cf668..4023376573bb 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -4930,7 +4930,17 @@ class TorchLlmArgs(BaseLlmArgs): "encoder-decoder model. Use `cuda_graph_config` for the decoder " "and this field for the encoder. Encoder CUDA graphs require " "`encoder_max_batch_size` to be set."), - status="prototype") + status="beta") + + enable_encoder_decoder_mixed_cuda_graph: bool = Field( + default=True, + description=( + "Enable the mixed-batch CUDA graph performance optimization for " + "encoder-decoder models. The graph handles decoder iterations " + "containing both context and generation requests. It is enabled " + "by default when both `cuda_graph_config` and " + "`encoder_cuda_graph_config` produce usable graph shapes."), + status="beta") @field_validator('cuda_graph_config', mode='before') @classmethod diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index 1e92bedcc416..450a4a9ae06b 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -407,6 +407,13 @@ "kind": "value", "path": "enable_early_first_token_response" }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "enable_encoder_decoder_mixed_cuda_graph" + }, { "allowed_values": [], "annotation": "", diff --git a/tests/unittest/_torch/executor/test_encoder_cuda_graph_runner.py b/tests/unittest/_torch/executor/test_encoder_cuda_graph_runner.py deleted file mode 100644 index 7c4222081d25..000000000000 --- a/tests/unittest/_torch/executor/test_encoder_cuda_graph_runner.py +++ /dev/null @@ -1,355 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -from types import SimpleNamespace - -import pytest -import torch - -from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttentionMetadata -from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import ( - EncoderCUDAGraphRunner, - EncoderCUDAGraphRunnerConfig, -) -from tensorrt_llm._torch.pyexecutor.model_engine import _build_encoder_decoder_cuda_graph_keys - - -def _dynamic_layout_runner( - capture_keys: list[tuple[int, int, int]] | None = None, - enable_padding: bool = False, -) -> EncoderCUDAGraphRunner: - return EncoderCUDAGraphRunner( - EncoderCUDAGraphRunnerConfig( - use_cuda_graph=False, - cuda_graph_padding_enabled=enable_padding, - cuda_graph_batch_sizes=[1, 2, 4, 8], - cuda_graph_num_tokens=[], - cuda_graph_seq_lens=list(range(64, 513, 64)), - max_cuda_graph_batch_size=8, - max_cuda_graph_num_tokens=4096, - max_num_tokens=4096, - max_seq_len=512, - cuda_graph_mem_pool=None, - encoder_decoder_capture_keys=capture_keys or [], - ) - ) - - -def test_encoder_decoder_capture_keys_select_layout_mode(): - encoder_decoder_runner = _dynamic_layout_runner() - encoder_only_runner = EncoderCUDAGraphRunner( - EncoderCUDAGraphRunnerConfig( - use_cuda_graph=False, - cuda_graph_padding_enabled=False, - cuda_graph_batch_sizes=[1], - cuda_graph_num_tokens=[8], - cuda_graph_seq_lens=[8], - max_cuda_graph_batch_size=1, - max_cuda_graph_num_tokens=8, - max_num_tokens=8, - max_seq_len=8, - cuda_graph_mem_pool=None, - ) - ) - - assert encoder_decoder_runner.is_encoder_decoder - assert not encoder_only_runner.is_encoder_decoder - - -def test_build_encoder_decoder_cuda_graph_keys(): - keys = _build_encoder_decoder_cuda_graph_keys( - batch_sizes=[1, 2], - num_tokens=[96, 576, 1056], - seq_lens=[512], - ) - - assert keys == [ - (1, 96, 512), - (2, 96, 512), - (2, 576, 512), - ] - - -def test_bart_encoder_graph_config_builds_feasible_key_grid(): - num_tokens = list(range(96, 4801, 96)) - keys = _build_encoder_decoder_cuda_graph_keys( - batch_sizes=[1, 2, 4, 8], - num_tokens=num_tokens, - seq_lens=[512, 1024], - ) - - assert len(keys) == 201 - assert {total_tokens for batch_size, total_tokens, _ in keys if batch_size == 8} == set( - num_tokens - ) - - -def test_encoder_graph_builds_reachable_startup_warmup_layouts(): - capture_keys = _build_encoder_decoder_cuda_graph_keys( - batch_sizes=[1, 2], - num_tokens=[96, 320], - seq_lens=[256, 512], - ) - runner = _dynamic_layout_runner( - capture_keys=capture_keys, - enable_padding=True, - ) - - for key in capture_keys: - sequence_lengths = runner.get_capture_warmup_sequence_lengths(key) - if sequence_lengths is None: - continue - - selected_key, _, is_valid = runner.get_graph_key({"seq_lens": sequence_lengths}) - assert is_valid - assert selected_key == key - assert len(sequence_lengths) == key[0] - assert sum(sequence_lengths) == key[1] - - assert runner.get_capture_warmup_sequence_lengths((1, 96, 256)) == [96] - assert runner.get_capture_warmup_sequence_lengths((1, 96, 512)) is None - assert runner.get_capture_warmup_sequence_lengths((2, 320, 512)) == [257, 63] - - -def test_encoder_graph_key_reuses_total_tokens_and_max_bucket(): - runner = _dynamic_layout_runner() - - key, is_padding_performed, is_valid = runner.get_graph_key( - { - "input_ids": list(range(580)), - "seq_lens": [260, 320], - } - ) - other_layout_key, _, _ = runner.get_graph_key( - { - "input_ids": list(range(580)), - "seq_lens": [284, 296], - } - ) - - assert key == (2, 580, 320) - assert other_layout_key == key - assert not is_padding_performed - assert is_valid - - -def test_encoder_graph_key_distinguishes_max_buckets(): - runner = _dynamic_layout_runner() - - key, _, _ = runner.get_graph_key( - { - "input_ids": [0] * 1400, - "seq_lens": [332, 356, 356, 356], - } - ) - larger_bucket_key, _, _ = runner.get_graph_key( - { - "input_ids": [0] * 1400, - "seq_lens": [260, 260, 440, 440], - } - ) - - assert key == (4, 1400, 384) - assert larger_bucket_key == (4, 1400, 448) - - -def test_encoder_graph_key_pads_tokens_and_max_sequence_length(): - runner = _dynamic_layout_runner( - capture_keys=[ - (2, 640, 320), - (2, 640, 384), - (2, 704, 384), - ], - enable_padding=True, - ) - - key, is_padding_performed, is_valid = runner.get_graph_key( - { - "input_ids": [0] * 556, - "seq_lens": [260, 296], - } - ) - - assert key == (2, 640, 320) - assert is_padding_performed - assert is_valid - - -def test_encoder_graph_pad_batch_selects_compatible_capture_key(): - runner = _dynamic_layout_runner( - capture_keys=[ - (4, 350, 192), - (4, 384, 192), - (8, 768, 192), - ], - enable_padding=True, - ) - runner.enabled = True - inputs = { - "input_ids": [0] * 350, - "seq_lens": [100, 120, 130], - } - - with runner.pad_batch(inputs, batch_size=3) as padded_inputs: - assert padded_inputs["seq_lens"] == [100, 120, 130, 1] - assert padded_inputs["input_ids"] is inputs["input_ids"] - key, is_padding_performed, is_valid = runner.get_graph_key(padded_inputs) - - assert key == (4, 384, 192) - assert is_padding_performed - assert is_valid - - -def test_encoder_graph_padding_rejects_incompatible_capture_keys(): - runner = _dynamic_layout_runner( - capture_keys=[ - (4, 320, 128), - (8, 512, 128), - ], - enable_padding=True, - ) - runner.enabled = True - inputs = { - "input_ids": [0] * 350, - "seq_lens": [100, 120, 130], - } - - with runner.pad_batch(inputs, batch_size=3) as padded_inputs: - assert padded_inputs is inputs - key, is_padding_performed, is_valid = runner.get_graph_key(padded_inputs) - - assert key == (3, 0, 0) - assert not is_padding_performed - assert not is_valid - - -def test_encoder_graph_key_rejects_oversized_inputs(): - runner = _dynamic_layout_runner() - - _, _, is_valid = runner.get_graph_key( - { - "input_ids": [0] * 4097, - "seq_lens": [4097], - } - ) - - assert not is_valid - - -def test_encoder_graph_only_captures_during_warmup(): - key = (1, 8, 64) - runner = _dynamic_layout_runner(capture_keys=[key]) - - assert not runner.needs_capture(key) - with runner.allow_capture(): - assert runner.needs_capture(key) - assert not runner.needs_capture(key) - - -def test_encoder_graph_reuses_same_key_for_different_sequence_layouts(): - key = (2, 580, 320) - runner = _dynamic_layout_runner(capture_keys=[key]) - runner.enabled = True - graph_metadata = object.__new__(TrtllmAttentionMetadata) - runner.graph_metadata[key] = { - "attn_metadata": graph_metadata, - } - - matched_metadata, matched_key = runner.maybe_get_cuda_graph( - { - "input_ids": [0] * 580, - "seq_lens": [260, 320], - }, - graph_metadata, - ) - reused_metadata, reused_key = runner.maybe_get_cuda_graph( - { - "input_ids": [0] * 580, - "seq_lens": [284, 296], - }, - graph_metadata, - ) - - assert matched_metadata is graph_metadata - assert matched_key == key - assert reused_metadata is graph_metadata - assert reused_key == key - - -def test_encoder_graph_replay_uses_plain_graph_mapping(monkeypatch): - runner = _dynamic_layout_runner() - key = (1, 8, 64) - attn_metadata = object() - expected_output = object() - replay_calls = [] - recorded_streams = [] - current_stream = object() - - runner.graphs[key] = SimpleNamespace(replay=lambda: replay_calls.append(key)) - runner.graph_metadata[key] = {"attn_metadata": attn_metadata} - runner.graph_outputs[key] = expected_output - runner._capture_h2d_copy = True - monkeypatch.setattr(runner, "retire_staging", lambda: None) - monkeypatch.setattr(runner, "_stage_inputs", lambda _key, _inputs: None) - monkeypatch.setattr( - torch.cuda, - "Event", - lambda: SimpleNamespace(record=lambda stream: recorded_streams.append(stream)), - ) - monkeypatch.setattr(torch.cuda, "current_stream", lambda: current_stream) - - output = runner.replay(key, {"attn_metadata": attn_metadata}) - - assert output is expected_output - assert replay_calls == [key] - assert recorded_streams == [current_stream] - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -def test_encoder_graph_capture_stages_warmup_and_replays_new_layout(): - runner = _dynamic_layout_runner() - runner.enabled = True - runner._create_shared_static_tensors() - - key = (2, 8, 64) - seq_lens_host = runner.shared_static_tensors_cpu["seq_lens"][:2] - seq_lens_host.copy_(torch.tensor([2, 3], dtype=torch.int32)) - attn_metadata = SimpleNamespace( - _seq_lens=seq_lens_host, - _seq_lens_cuda=torch.ones(2, device="cuda", dtype=torch.int32), - ) - inputs = { - "input_ids": [10, 11, 12, 13, 14], - "position_ids": [0, 1, 0, 1, 2], - "seq_lens": [2, 3], - "attn_metadata": attn_metadata, - } - warmup_seq_lens = [] - - def forward_fn(capture_inputs): - if not torch.cuda.is_current_stream_capturing(): - warmup_seq_lens.append(capture_inputs["attn_metadata"]._seq_lens_cuda.cpu().tolist()) - return capture_inputs["input_ids"] + capture_inputs["attn_metadata"]._seq_lens_cuda[0] - - runner.capture(key, forward_fn, inputs) - - assert warmup_seq_lens == [[2, 3]] - - first_output = runner.replay(key, inputs) - torch.cuda.synchronize() - torch.testing.assert_close( - first_output, - torch.tensor([12, 13, 14, 15, 16, 2, 2, 2], device="cuda", dtype=torch.int32), - ) - - seq_lens_host.copy_(torch.tensor([1, 4], dtype=torch.int32)) - reused_inputs = { - **inputs, - "seq_lens": [1, 4], - } - reused_output = runner.replay(key, reused_inputs) - torch.cuda.synchronize() - torch.testing.assert_close( - reused_output, - torch.tensor([11, 12, 13, 14, 15, 1, 1, 1], device="cuda", dtype=torch.int32), - ) diff --git a/tests/unittest/_torch/executor/test_mixed_decoder_cuda_graph_runner.py b/tests/unittest/_torch/executor/test_mixed_decoder_cuda_graph_runner.py deleted file mode 100644 index 8af9118313be..000000000000 --- a/tests/unittest/_torch/executor/test_mixed_decoder_cuda_graph_runner.py +++ /dev/null @@ -1,159 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -from types import SimpleNamespace - -import torch - -from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDAGraphRunner -from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests - - -def _mixed_batch() -> ScheduledRequests: - batch = ScheduledRequests() - batch.context_requests_last_chunk = [ - SimpleNamespace( - context_chunk_size=2, - encoder_output_len=260, - py_skip_cross_kv_projection=False, - ), - SimpleNamespace( - context_chunk_size=2, - encoder_output_len=272, - py_skip_cross_kv_projection=False, - ), - ] - batch.generation_requests = [ - SimpleNamespace(py_draft_tokens=[]), - SimpleNamespace(py_draft_tokens=[]), - ] - return batch - - -def _runner() -> CUDAGraphRunner: - runner = object.__new__(CUDAGraphRunner) - runner.config = SimpleNamespace( - enable_attention_dp=False, - is_draft_model=False, - use_mrope=False, - ) - runner.enabled = True - runner.padding_enabled = True - runner.sparse_config = None - runner.max_beam_width = 1 - runner.enable_encoder_decoder_mixed_cuda_graph = True - runner.graphs = {} - runner.graph_outputs = {} - runner.graph_metadata = {} - runner.padding_dummy_requests = {} - runner.memory_pool = None - return runner - - -def test_mixed_encoder_decoder_graph_key_captures_dynamic_extents(): - runner = _runner() - - key = runner.get_graph_key(_mixed_batch()) - - assert key == (4, 0, False, False, True, (2, 2), (532,)) - assert runner._get_num_tokens_for_key(key) == 6 - - -def test_mixed_encoder_decoder_graph_key_distinguishes_cached_cross_kv(): - runner = _runner() - batch = _mixed_batch() - batch.context_requests_last_chunk[1].py_skip_cross_kv_projection = True - - key = runner.get_graph_key(batch) - - assert key[6] == (260,) - - -def test_mixed_encoder_decoder_graph_eligibility_requires_both_phases(): - runner = _runner() - batch = _mixed_batch() - - assert runner._is_mixed_encoder_decoder_batch(batch) - - batch.generation_requests = [] - assert not runner._is_mixed_encoder_decoder_batch(batch) - - -def test_mixed_encoder_decoder_graph_never_captures_at_runtime(): - runner = _runner() - runner._capture_allowed = False - key = runner.get_graph_key(_mixed_batch()) - - assert not runner.needs_capture(key) - - -def test_mixed_encoder_decoder_graph_key_pads_encoder_extent(): - runner = _runner() - runner._capture_allowed = False - batch = _mixed_batch() - padded_key = (4, 0, False, False, True, (2, 2), (576,)) - graph_attn_metadata = object() - runner.graph_metadata[padded_key] = { - "attn_metadata": graph_attn_metadata, - "spec_metadata": None, - } - runner.graph_outputs[padded_key] = object() - - attn_metadata, spec_metadata, key = runner.maybe_get_cuda_graph( - batch, - enable_spec_decode=False, - attn_metadata=object(), - allow_mixed_encoder_decoder=True, - ) - - assert key == padded_key - assert attn_metadata is graph_attn_metadata - assert spec_metadata is None - - -def test_mixed_encoder_decoder_graph_key_uses_smallest_compatible_extent(): - runner = _runner() - runner._capture_allowed = False - key = runner.get_graph_key(_mixed_batch()) - larger_key = (*key[:6], (672,)) - smallest_key = (*key[:6], (576,)) - incompatible_key = (*key[:5], (2,), (544,)) - runner.graph_outputs = { - larger_key: object(), - smallest_key: object(), - incompatible_key: object(), - } - - assert runner._get_compatible_mixed_encoder_decoder_key(key) == smallest_key - - -def test_mixed_encoder_decoder_replay_zero_pads_encoder_hidden_states(): - runner = _runner() - key = (4, 0, False, False, True, (2, 2), (576,)) - attn_metadata = object() - runner.graph_metadata[key] = { - "attn_metadata": attn_metadata, - "spec_metadata": None, - } - runner.graph_outputs[key] = object() - runner.graphs[key] = SimpleNamespace(replay=lambda: None, reset=lambda: None) - runner.shared_static_tensors = { - "input_ids": torch.zeros(6, dtype=torch.int32), - "position_ids": torch.zeros((1, 6), dtype=torch.int32), - "encoder_hidden_states": torch.ones((576, 2)), - } - encoder_hidden_states = torch.arange(532 * 2, dtype=torch.float32).reshape(532, 2) - - runner.replay( - key, - { - "attn_metadata": attn_metadata, - "input_ids": torch.ones(6, dtype=torch.int32), - "position_ids": torch.ones((1, 6), dtype=torch.int32), - "encoder_hidden_states": encoder_hidden_states, - }, - ) - - staged_encoder_hidden_states = runner.shared_static_tensors["encoder_hidden_states"] - assert torch.equal(staged_encoder_hidden_states[:532], encoder_hidden_states) - assert torch.count_nonzero(staged_encoder_hidden_states[532:]) == 0 diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index 9bff40b15303..6e9fed489e26 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -94,7 +94,11 @@ methods: encoder_cuda_graph_config: annotation: Optional[tensorrt_llm.llmapi.llm_args.EncodeCudaGraphConfig] default: null - status: prototype + status: beta + enable_encoder_decoder_mixed_cuda_graph: + annotation: bool + default: True + status: beta multimodal_config: annotation: tensorrt_llm.llmapi.llm_args.MultimodalConfig default: null diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 9995d95ba34f..514388826b8f 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -1766,6 +1766,22 @@ def test_encoder_decoder_cuda_graph_configs(self): assert args.encoder_cuda_graph_config.batch_sizes == [1, 4] assert args.encoder_cuda_graph_config.num_tokens == [16, 64] assert args.encoder_cuda_graph_config.seq_lens == [8, 32] + assert args.enable_encoder_decoder_mixed_cuda_graph + + def test_encoder_decoder_mixed_cuda_graph_can_be_disabled(self): + args = TorchLlmArgs( + model=llama_model_path, + encoder_max_batch_size=4, + encoder_cuda_graph_config=EncodeCudaGraphConfig( + batch_sizes=[1, 4], + num_tokens=[16, 64], + seq_lens=[8, 32], + enable_padding=True, + ), + enable_encoder_decoder_mixed_cuda_graph=False, + ) + + assert not args.enable_encoder_decoder_mixed_cuda_graph def test_encoder_cuda_graph_config_requires_encoder_max_batch_size(self): with pytest.raises(ValidationError, From 4f981aca1356db4233eaf7284183fade32c5e675 Mon Sep 17 00:00:00 2001 From: Guiju Zhang <7135567+cascade812@users.noreply.github.com> Date: Thu, 30 Jul 2026 11:10:17 -0700 Subject: [PATCH 07/15] [None][fix] stabilize encoder-decoder CUDA graph capture Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com> --- .../_torch/pyexecutor/model_engine.py | 140 +++++++++--------- .../defs/llmapi/test_llm_api_pytorch_bart.py | 130 ++++++++++++++++ .../defs/llmapi/test_llm_api_pytorch_t5.py | 65 +++++++- .../test_lists/test-db/l0_h100.yml | 2 + .../test_lists/test-db/l0_l40s.yml | 1 - 5 files changed, 266 insertions(+), 72 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index da7758c55f48..9ecaa2ca3420 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -496,6 +496,8 @@ def __init__( cuda_graph_padding_enabled = self.cuda_graph_config.enable_padding if self.cuda_graph_config else CudaGraphConfig.model_fields[ 'enable_padding'].default + # CUDA graph detection for encoder-decoder models and encoder-only models. + # Decode configs do not define these encoder-specific bucket fields. encoder_cuda_graph_batch_sizes = ( self.encoder_cuda_graph_config.batch_sizes if self.encoder_cuda_graph_config is not None else []) @@ -525,6 +527,47 @@ def __init__( f"EncodeCudaGraphConfig(max_batch_size=64, num_tokens=[128, 256, " f"512], max_seq_len=128, enable_padding=True).") + self._cuda_graph_padding_enabled = cuda_graph_padding_enabled + + self._cuda_graph_batch_sizes = _filter_cuda_graph_batch_sizes( + cuda_graph_batch_sizes, self.batch_size, self.max_num_tokens, + self.original_max_total_draft_tokens, + self._cuda_graph_padding_enabled) if cuda_graph_batch_sizes else [] + + self._max_cuda_graph_batch_size = (self._cuda_graph_batch_sizes[-1] if + self._cuda_graph_batch_sizes else 0) + + self._encoder_cuda_graph_padding_enabled = ( + encoder_cuda_graph_padding_enabled) + self._encoder_cuda_graph_batch_sizes = (_filter_cuda_graph_batch_sizes( + encoder_cuda_graph_batch_sizes, self.encoder_batch_size, + self.encoder_max_num_tokens, 0, + self._encoder_cuda_graph_padding_enabled) if + encoder_cuda_graph_batch_sizes + else []) + + # Encoder CUDA graph bucket lists + self._cuda_graph_num_tokens = (_filter_cuda_graph_num_tokens( + encoder_cuda_graph_num_tokens, self.encoder_max_num_tokens, + self._encoder_cuda_graph_padding_enabled) + if encoder_cuda_graph_num_tokens else []) + + self._max_cuda_graph_num_tokens = (self._cuda_graph_num_tokens[-1] if + self._cuda_graph_num_tokens else 0) + self._cuda_graph_seq_lens = (_filter_cuda_graph_seq_lens( + encoder_cuda_graph_seq_lens, self.max_seq_len, + self._encoder_cuda_graph_padding_enabled) + if encoder_cuda_graph_seq_lens else []) + + self._max_cuda_graph_seq_len = (self._cuda_graph_seq_lens[-1] + if self._cuda_graph_seq_lens else 0) + + use_encoder_cuda_graph = ((self._is_encoder_decoder_model() + or self._is_encode_only()) + and self.encoder_cuda_graph_config is not None + and bool(self._cuda_graph_num_tokens) + and bool(self._cuda_graph_seq_lens)) + self.torch_compile_config = self.llm_args.torch_compile_config torch_compile_enabled = bool(self.torch_compile_config is not None) torch_compile_fullgraph = self.torch_compile_config.enable_fullgraph if self.torch_compile_config is not None else TorchCompileConfig.model_fields[ @@ -683,49 +726,6 @@ def __init__( self.iter_states = {} self._cuda_graph_mem_pool = self._torch_compile_backend._graph_pool_handle if self._torch_compile_enabled else None - self._cuda_graph_padding_enabled = cuda_graph_padding_enabled - - self._cuda_graph_batch_sizes = _filter_cuda_graph_batch_sizes( - cuda_graph_batch_sizes, self.batch_size, self.max_num_tokens, - self.original_max_total_draft_tokens, - self._cuda_graph_padding_enabled) if cuda_graph_batch_sizes else [] - - self._max_cuda_graph_batch_size = (self._cuda_graph_batch_sizes[-1] if - self._cuda_graph_batch_sizes else 0) - - self._encoder_cuda_graph_padding_enabled = ( - encoder_cuda_graph_padding_enabled) - self._encoder_cuda_graph_batch_sizes = (_filter_cuda_graph_batch_sizes( - encoder_cuda_graph_batch_sizes, self.encoder_batch_size, - self.encoder_max_num_tokens, 0, - self._encoder_cuda_graph_padding_enabled) if - encoder_cuda_graph_batch_sizes - else []) - - # Encoder CUDA graph bucket lists - self._cuda_graph_num_tokens = (_filter_cuda_graph_num_tokens( - encoder_cuda_graph_num_tokens, self.encoder_max_num_tokens, - self._encoder_cuda_graph_padding_enabled) - if encoder_cuda_graph_num_tokens else []) - - self._max_cuda_graph_num_tokens = (self._cuda_graph_num_tokens[-1] if - self._cuda_graph_num_tokens else 0) - self._cuda_graph_seq_lens = (_filter_cuda_graph_seq_lens( - encoder_cuda_graph_seq_lens, self.max_seq_len, - self._encoder_cuda_graph_padding_enabled) - if encoder_cuda_graph_seq_lens else []) - - self._max_cuda_graph_seq_len = (self._cuda_graph_seq_lens[-1] - if self._cuda_graph_seq_lens else 0) - - # Encoder CUDA graph detection for encoder-decoder models and encoder-only models. - # Decode configs do not define these encoder-specific bucket fields. - use_encoder_cuda_graph = ((self._is_encoder_decoder_model() - or self._is_encode_only()) - and self.encoder_cuda_graph_config is not None - and bool(self._cuda_graph_num_tokens) - and bool(self._cuda_graph_seq_lens)) - self._dynamic_draft_len_mapping = self._compute_dynamic_draft_len_mapping( ) @@ -1884,17 +1884,29 @@ def _warmup_encoder_decoder_encoder_cuda_graphs( self, resource_manager: ResourceManager) -> None: """Capture encoder-decoder encoder graphs on their runtime host thread.""" runner = self.encoder_cuda_graph_runner - if not runner.enabled or not runner.is_encoder_decoder: + if not runner.is_encoder_decoder: + return + + capture = functools.partial( + self._capture_encoder_decoder_encoder_cuda_graphs, + resource_manager, + ) + self._warmup_and_capture_encoder_cuda_graphs(capture) + + def _warmup_and_capture_encoder_cuda_graphs( + self, capture: Callable[[], None]) -> None: + """Warm up every encoder graph shape, then capture those shapes.""" + runner = self.encoder_cuda_graph_runner + if not runner.enabled: return with runner.allow_capture(): runner.is_warmup_only = True try: - self._capture_encoder_decoder_encoder_cuda_graphs( - resource_manager) + capture() finally: runner.is_warmup_only = False - self._capture_encoder_decoder_encoder_cuda_graphs(resource_manager) + capture() def _capture_encoder_decoder_encoder_cuda_graphs( self, resource_manager: ResourceManager) -> None: @@ -2164,6 +2176,11 @@ def _capture_mixed_encoder_decoder_cuda_graphs( operation = ("warmup" if runner.is_warmup_only else "capture") hidden_size = self._get_enc_dec_hidden_size() + model_config = self.model.model_config.pretrained_config + # BART/mBART prepend a forced BOS token after decoder_start; T5 uses + # decoder_start alone. Match the LLM API's decoder-prefix construction. + mixed_context_query_len = (2 if getattr( + model_config, "model_type", None) in ("bart", "mbart") else 1) for num_contexts, total_encoder_tokens in sorted( context_shapes, key=lambda shape: shape[1], reverse=True): if total_encoder_tokens > num_contexts * max_encoder_output_len: @@ -2183,7 +2200,8 @@ def _capture_mixed_encoder_decoder_cuda_graphs( resource_manager, batch_size, draft_len=0, - mixed_context_encoder_output_lens=encoder_output_lens) + mixed_context_encoder_output_lens=encoder_output_lens, + mixed_context_query_len=mixed_context_query_len) with self._release_batch_context(warmup_request, resource_manager) as batch: if batch is None: @@ -2198,7 +2216,7 @@ def _capture_mixed_encoder_decoder_cuda_graphs( context_requests, encoder_output_lens): request.state = LlmRequestState.CONTEXT_INIT request.context_current_position = 0 - request.context_chunk_size = 2 + request.context_chunk_size = mixed_context_query_len request.cached_tokens = 0 request.py_batch_idx = None request.py_encoder_output = torch.ones( @@ -2445,7 +2463,8 @@ def _create_cuda_graph_warmup_request( batch_size: int, draft_len: int, max_seq_len: int = None, - mixed_context_encoder_output_lens: Optional[Sequence[int]] = None + mixed_context_encoder_output_lens: Optional[Sequence[int]] = None, + mixed_context_query_len: int = ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM, ) -> Optional[ScheduledRequests]: """Creates a dummy ScheduledRequests tailored for CUDA graph capture.""" kv_cache_manager = resource_manager.get_resource_manager( @@ -2482,8 +2501,7 @@ def _create_cuda_graph_warmup_request( context_request_ids = list(range(num_mixed_contexts)) context_requests = kv_cache_manager.add_dummy_requests( context_request_ids, - token_nums=[ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM] * - num_mixed_contexts, + token_nums=[mixed_context_query_len] * num_mixed_contexts, is_gen=False, max_num_draft_tokens=runtime_draft_token_buffer_width, kv_reserve_draft_tokens=self.max_draft_loop_tokens, @@ -2617,7 +2635,7 @@ def free_warmup_requests() -> None: for request in requests[:num_mixed_contexts]: request.state = LlmRequestState.CONTEXT_INIT request.context_current_position = 0 - request.context_chunk_size = 2 + request.context_chunk_size = mixed_context_query_len request.cached_tokens = 0 request.py_batch_idx = None result.context_requests_last_chunk = requests[:num_mixed_contexts] @@ -6443,13 +6461,8 @@ def warmup_encoder(self) -> None: # a larger workspace, so the first pass grows the workspace to its # maximum size. The second pass runs the final per-shape warmup and # captures without resizing the workspace. - with self.encoder_cuda_graph_runner.allow_capture(): - self.encoder_cuda_graph_runner.is_warmup_only = True - try: - self._run_cuda_graph_warmup_encoder() - finally: - self.encoder_cuda_graph_runner.is_warmup_only = False - self._run_cuda_graph_warmup_encoder() + self._warmup_and_capture_encoder_cuda_graphs( + self._capture_encoder_cuda_graphs) # Pre-populate the memory pool with max-shape allocations to reduce # fragmentation at runtime. @@ -6497,13 +6510,6 @@ def _run_autotuner_warmup_encoder(self) -> None: f"{len(AutoTuner.get().profiling_cache)}") AutoTuner.get().print_profiling_cache() - def _run_cuda_graph_warmup_encoder(self) -> None: - """Warm up or capture whole-model encode-only CUDA graphs.""" - if not self.encoder_cuda_graph_runner.enabled: - return - - self._capture_encoder_cuda_graphs() - def _capture_encoder_cuda_graphs(self) -> None: """Warm up or capture encoder CUDA graphs for all feasible keys. diff --git a/tests/integration/defs/llmapi/test_llm_api_pytorch_bart.py b/tests/integration/defs/llmapi/test_llm_api_pytorch_bart.py index 6f94d06b6efa..af48a095b2ff 100644 --- a/tests/integration/defs/llmapi/test_llm_api_pytorch_bart.py +++ b/tests/integration/defs/llmapi/test_llm_api_pytorch_bart.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import time from pathlib import Path import pytest @@ -21,6 +22,7 @@ from tensorrt_llm.llmapi import ( LLM, CudaGraphConfig, + EncodeCudaGraphConfig, KvCacheConfig, RequestOutput, SamplingParams, @@ -46,6 +48,7 @@ _MBART_TARGET_LANG = "en_XX" _MBART_SOURCE_TEXT = "Şeful ONU spune că nu există o soluţie militară în Siria." _MAX_NEW_TOKENS = 10 +_CONTINUOUS_ADMISSION_MAX_NEW_TOKENS = 16 _MAX_SEQUENCE_LENGTH = 128 _MAX_KV_TOKENS = 384 _MIN_GPU_MEMORY_MB = 16_000 @@ -303,6 +306,14 @@ def _decoder_cuda_graph_config( ) +class _SleepLogitsProcessor: + def __init__(self, delay_seconds: float) -> None: + self.delay_seconds = delay_seconds + + def __call__(self, req_id, logits, token_ids, stream_ptr, client_id) -> None: + time.sleep(self.delay_seconds) + + def _assert_bart_response( response: RequestOutput, num_return_sequences: int, @@ -593,3 +604,122 @@ def test_bart_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch( request_idx ], ) + + +def test_bart_pytorch_continuous_admission_replays_encoder_and_mixed_cuda_graphs( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Admit encoder work while an older request remains in decoder generation.""" + monkeypatch.setenv("TRTLLM_SKIP_KV_CACHE_ESTIMATION", "1") + monkeypatch.setenv("TLLM_WORKER_USE_SINGLE_PROCESS", "1") + + model_path = _get_model_path(_MODEL_NAME) + tokenizer = AutoTokenizer.from_pretrained(model_path) + first_sampling_params = SamplingParams( + max_tokens=_CONTINUOUS_ADMISSION_MAX_NEW_TOKENS, + temperature=0.0, + ignore_eos=True, + logits_processor=_SleepLogitsProcessor(delay_seconds=0.02), + ) + second_sampling_params = _sampling_params( + num_beams=1, + num_return_sequences=1, + ) + encoder_graph_config = EncodeCudaGraphConfig( + batch_sizes=[1], + num_tokens=[64], + seq_lens=[64], + enable_padding=True, + ) + + with LLM( + model_path, + backend="pytorch", + attn_backend="TRTLLM", + cuda_graph_config=_decoder_cuda_graph_config([2]), + encoder_cuda_graph_config=encoder_graph_config, + enable_encoder_decoder_mixed_cuda_graph=True, + disable_overlap_scheduler=False, + dtype="bfloat16", + enable_chunked_prefill=False, + encoder_max_batch_size=1, + encoder_max_num_tokens=64, + kv_cache_config=KvCacheConfig( + enable_block_reuse=False, + max_tokens=_MAX_KV_TOKENS, + free_gpu_memory_fraction=_FREE_GPU_MEMORY_FRACTION, + cross_kv_cache_fraction=_CROSS_KV_CACHE_FRACTION, + use_kv_cache_manager_v2=False, + ), + max_batch_size=2, + max_beam_width=1, + max_input_len=_MAX_SEQUENCE_LENGTH, + max_num_tokens=_MAX_SEQUENCE_LENGTH, + max_seq_len=_MAX_SEQUENCE_LENGTH, + model_kwargs={"torch_dtype": "bfloat16"}, + scheduler_config=SchedulerConfig(use_python_scheduler=True), + ) as llm: + model_engine = llm._executor.engine.model_engine + encoder_runner = model_engine.encoder_cuda_graph_runner + decoder_runner = model_engine.cuda_graph_runner + + assert encoder_runner.enabled + assert encoder_runner.graphs + captured_mixed_keys = {key for key in decoder_runner.graphs if key[5] and key[6]} + assert captured_mixed_keys + + encoder_replay_keys = [] + decoder_replay_keys = [] + original_encoder_replay = encoder_runner.replay + original_decoder_replay = decoder_runner.replay + + def record_encoder_replay(key, inputs): + encoder_replay_keys.append(key) + return original_encoder_replay(key, inputs) + + def record_decoder_replay(key, inputs): + decoder_replay_keys.append(key) + return original_decoder_replay(key, inputs) + + monkeypatch.setattr(encoder_runner, "replay", record_encoder_replay) + monkeypatch.setattr(decoder_runner, "replay", record_decoder_replay) + + first_response = llm.generate_async( + _SOURCE_TEXT, + sampling_params=first_sampling_params, + streaming=True, + ) + first_stream_step = next(first_response) + assert not first_stream_step.finished + + second_response = llm.generate_async( + _MIXED_ENCODER_SOURCE_TEXTS[1], + sampling_params=second_sampling_params, + streaming=False, + ) + + first_response.result() + second_response.result() + + first_token_ids = _assert_bart_response( + first_response, + num_return_sequences=1, + max_tokens=_CONTINUOUS_ADMISSION_MAX_NEW_TOKENS, + ) + second_token_ids = _assert_bart_response( + second_response, + num_return_sequences=1, + ) + assert first_token_ids[0][:_MAX_NEW_TOKENS] == _EXPECTED_GREEDY_OUTPUT_TOKEN_IDS + _assert_expected_generation( + tokenizer, + second_token_ids, + exact_match=True, + expected_token_ids_by_output=_MIXED_ENCODER_EXPECTED_TOKEN_IDS_BY_REQUEST[1], + ) + + assert len(encoder_replay_keys) >= 2 + assert set(encoder_replay_keys) <= set(encoder_runner.graphs) + replayed_mixed_keys = {key for key in decoder_replay_keys if key[5] and key[6]} + assert replayed_mixed_keys + assert replayed_mixed_keys <= captured_mixed_keys diff --git a/tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py b/tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py index 492b88ff3d26..4c93c0b27c93 100644 --- a/tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py +++ b/tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py @@ -22,6 +22,7 @@ from tensorrt_llm.llmapi import ( LLM, CudaGraphConfig, + EncodeCudaGraphConfig, KvCacheConfig, RequestOutput, SamplingParams, @@ -819,13 +820,16 @@ def test_t5_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch( ) -def test_t5_pytorch_generate_encoder_decoder_mixed_context_generation_batch( +def test_t5_pytorch_continuous_admission_replays_encoder_and_mixed_cuda_graphs( monkeypatch: pytest.MonkeyPatch, ) -> None: + """Continuously admit work and replay encoder and mixed decoder graphs.""" monkeypatch.setenv("TRTLLM_SKIP_KV_CACHE_ESTIMATION", "1") + monkeypatch.setenv("TLLM_WORKER_USE_SINGLE_PROCESS", "1") model_name = "t5-small" model_path = _get_t5_model_path(model_name) + tokenizer = AutoTokenizer.from_pretrained(model_path) first_sampling_params = SamplingParams( max_tokens=_MIXED_CONTEXT_GENERATION_MAX_NEW_TOKENS, temperature=0.0, @@ -836,15 +840,25 @@ def test_t5_pytorch_generate_encoder_decoder_mixed_context_generation_batch( max_tokens=_MAX_NEW_TOKENS, temperature=0.0, ) + encoder_graph_config = EncodeCudaGraphConfig( + batch_sizes=[1], + num_tokens=[64], + seq_lens=[64], + enable_padding=True, + ) with LLM( model_path, backend="pytorch", attn_backend="TRTLLM", cuda_graph_config=_decoder_cuda_graph_config([2]), - disable_overlap_scheduler=True, + encoder_cuda_graph_config=encoder_graph_config, + enable_encoder_decoder_mixed_cuda_graph=True, + disable_overlap_scheduler=False, dtype="bfloat16", enable_chunked_prefill=False, + encoder_max_batch_size=1, + encoder_max_num_tokens=64, kv_cache_config=KvCacheConfig( enable_block_reuse=False, max_tokens=_MAX_KV_TOKENS, @@ -860,6 +874,31 @@ def test_t5_pytorch_generate_encoder_decoder_mixed_context_generation_batch( model_kwargs={"torch_dtype": "bfloat16"}, scheduler_config=SchedulerConfig(use_python_scheduler=True), ) as llm: + model_engine = llm._executor.engine.model_engine + encoder_runner = model_engine.encoder_cuda_graph_runner + decoder_runner = model_engine.cuda_graph_runner + + assert encoder_runner.enabled + assert encoder_runner.graphs + captured_mixed_keys = {key for key in decoder_runner.graphs if key[5] and key[6]} + assert captured_mixed_keys + + encoder_replay_keys = [] + decoder_replay_keys = [] + original_encoder_replay = encoder_runner.replay + original_decoder_replay = decoder_runner.replay + + def record_encoder_replay(key, inputs): + encoder_replay_keys.append(key) + return original_encoder_replay(key, inputs) + + def record_decoder_replay(key, inputs): + decoder_replay_keys.append(key) + return original_decoder_replay(key, inputs) + + monkeypatch.setattr(encoder_runner, "replay", record_encoder_replay) + monkeypatch.setattr(decoder_runner, "replay", record_decoder_replay) + first_response = llm.generate_async( _SOURCE_TEXT, sampling_params=first_sampling_params, @@ -877,9 +916,27 @@ def test_t5_pytorch_generate_encoder_decoder_mixed_context_generation_batch( first_response.result() second_response.result() - _assert_t5_response( + first_token_ids = _assert_t5_response( first_response, num_return_sequences=1, max_tokens=_MIXED_CONTEXT_GENERATION_MAX_NEW_TOKENS, ) - _assert_t5_response(second_response, num_return_sequences=1) + second_token_ids = _assert_t5_response(second_response, num_return_sequences=1) + + expected_token_ids_by_request = _MIXED_ENCODER_OUTPUT_TOKEN_IDS_BY_MODEL_AND_BEAMS[ + (model_name, 1) + ] + assert first_token_ids[0][:_MAX_NEW_TOKENS] == expected_token_ids_by_request[0][0] + _assert_expected_generation( + tokenizer, + second_token_ids, + exact_match=True, + expected_token_ids_by_output=expected_token_ids_by_request[1], + expected_text_fragment=_MIXED_ENCODER_EXPECTED_TEXT_FRAGMENTS_BY_MODEL[model_name][1], + ) + + assert len(encoder_replay_keys) >= 2 + assert set(encoder_replay_keys) <= set(encoder_runner.graphs) + replayed_mixed_keys = {key for key in decoder_replay_keys if key[5] and key[6]} + assert replayed_mixed_keys + assert replayed_mixed_keys <= captured_mixed_keys diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index f4a3fcefef03..a8facffd8a44 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -322,6 +322,7 @@ l0_h100: - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-on-beam2-overlap-bart-large-cnn] - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v2-cuda-graph-on-greedy-overlap-bart-large-cnn] - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch[bf16-kv-v2-decoder-cuda-graph-on-greedy-batch2-bart-large-cnn] + - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_continuous_admission_replays_encoder_and_mixed_cuda_graphs - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-on-beam2-t5-base] - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-on-beam2-t5-large] - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-on-beam2-flan-t5-base] @@ -343,6 +344,7 @@ l0_h100: - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch[bf16-kv-v1-decoder-cuda-graph-on-beam2-batch2-t5-small] - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch[bf16-kv-v1-decoder-cuda-graph-on-beam2-batch2-flan-t5-small] - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch[bf16-kv-v2-decoder-cuda-graph-on-greedy-batch2-t5-small] + - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_continuous_admission_replays_encoder_and_mixed_cuda_graphs - llmapi/test_llm_api_pytorch_whisper.py::test_whisper_pytorch_beam_search[fp32-kv-v1-graphs-off-beam2] - llmapi/test_llm_api_pytorch_whisper.py::test_whisper_pytorch_feature_combinations[fp32-kv-v2-graphs-off-greedy] - llmapi/test_llm_api_pytorch_whisper.py::test_whisper_pytorch_feature_combinations[fp32-kv-v1-graphs-requested-greedy] diff --git a/tests/integration/test_lists/test-db/l0_l40s.yml b/tests/integration/test_lists/test-db/l0_l40s.yml index 18fa6b86be18..1bba99fa161f 100644 --- a/tests/integration/test_lists/test-db/l0_l40s.yml +++ b/tests/integration/test_lists/test-db/l0_l40s.yml @@ -53,7 +53,6 @@ l0_l40s: - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v2-cuda-graph-on-greedy-t5-small] - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-on-greedy-overlap-t5-small] - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch[bf16-kv-v1-decoder-cuda-graph-on-greedy-batch2-t5-small] - - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_mixed_context_generation_batch # Whisper (encoder-decoder) — customer-side deployment targets L40S/H200 - llmapi/test_llm_api_pytorch_whisper.py::test_whisper_pytorch_transcribe_end_to_end - llmapi/test_llm_api_pytorch_whisper.py::test_whisper_pytorch_feature_combinations[bf16-kv-v2-decoder-graphs-on-greedy] From eca64fda1c0b75b051e0f2d4487101b9604eb356 Mon Sep 17 00:00:00 2001 From: Guiju Zhang <7135567+cascade812@users.noreply.github.com> Date: Thu, 30 Jul 2026 21:04:51 -0700 Subject: [PATCH 08/15] [None][fix] stabilize encoder-decoder CUDA graph replay Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com> --- docs/source/models/encoder-decoder.md | 109 ++++---- docs/source/models/supported-models.md | 22 ++ .../_torch/attention_backend/trtllm.py | 12 +- .../_torch/pyexecutor/cuda_graph_runner.py | 245 +++++++++++++++++- .../_torch/pyexecutor/model_engine.py | 41 ++- tensorrt_llm/_torch/pyexecutor/py_executor.py | 30 ++- .../_torch/pyexecutor/sampler/sampler.py | 1 + .../defs/llmapi/test_llm_api_pytorch_bart.py | 118 +++++---- .../defs/llmapi/test_llm_api_pytorch_t5.py | 63 +++-- .../test_lists/test-db/l0_dgx_h100.yml | 1 + .../test_lists/test-db/l0_h100.yml | 2 +- .../_torch/executor/test_py_executor.py | 235 ++++++++--------- .../executor/test_pytorch_model_engine.py | 67 ++++- .../test_pytorch_model_engine_warmup.py | 38 ++- .../_torch/sampler/test_torch_sampler.py | 85 +++--- tests/unittest/llmapi/test_llm_args.py | 89 +++---- 16 files changed, 776 insertions(+), 382 deletions(-) diff --git a/docs/source/models/encoder-decoder.md b/docs/source/models/encoder-decoder.md index ab2158ac7de7..8d09f6816282 100644 --- a/docs/source/models/encoder-decoder.md +++ b/docs/source/models/encoder-decoder.md @@ -45,7 +45,7 @@ The following table describes the supported and recommended configurations. | Beam search | Yes with V1 | Configure `max_beam_width` when constructing `LLM`, then set `use_beam_search=True` in `SamplingParams`. | | Attention backend | `TRTLLM` | Use this backend for encoder-decoder models. It is required when `tensor_parallel_size > 1`. | | Decoder CUDA graphs | Yes, except in FP32 | `CudaGraphConfig` captures decoder work. V1 supports greedy and beam search; V2 supports its single-beam path. FP32 encoder-decoder models decline capture at engine init and log a warning instead of failing. | -| Encoder CUDA graphs | Yes | Set `encoder_cuda_graph_config=EncodeCudaGraphConfig(...)` and `encoder_max_batch_size`. The `TRTLLM` attention backend is required. | +| Encoder CUDA graphs | Yes | Set `encoder_cuda_graph_config=EncodeCudaGraphConfig(...)` and `encoder_max_batch_size`. Usually set `encoder_max_batch_size` lower than `max_batch_size`. The `TRTLLM` attention backend is required. | | Overlap scheduler | Yes | Enabled by default. V1 supports greedy decoding and beam search; V2 remains limited to `max_beam_width=1`. | | Tensor parallelism | Yes | Use `tensor_parallel_size > 1` with `attn_backend="TRTLLM"`. Attention head counts must be divisible by the TP size. | | Pipeline parallelism | No | Keep `pipeline_parallel_size=1`. | @@ -410,14 +410,15 @@ llm = LLM( backend="pytorch", attn_backend="TRTLLM", max_batch_size=8, - encoder_max_batch_size=8, + encoder_max_batch_size=2, + encoder_max_num_tokens=2048, cuda_graph_config=CudaGraphConfig( max_batch_size=8, enable_padding=True, ), encoder_cuda_graph_config=EncodeCudaGraphConfig( - batch_sizes=[1, 2, 4, 8], - num_tokens=[128, 256, 512, 1024, 2048, 4096], + batch_sizes=[1, 2], + num_tokens=[128, 256, 512, 1024, 2048], seq_lens=[128, 256, 512, 1024], enable_padding=True, ), @@ -436,6 +437,13 @@ size, total packed tokens, and maximum sequence length. The limit. With beam search, decoder graph batch sizes must cover the active decoder sequences after beam expansion. +`max_batch_size` controls the total decoder concurrency, while +`encoder_max_batch_size` controls encoder microbatch admission. For better +performance, tune `encoder_max_batch_size`, `encoder_max_num_tokens`, and the +encoder CUDA graph buckets together for the production workload. Start with +`encoder_max_batch_size` smaller than `max_batch_size`, such as 2 versus 8, +then adjust the limits and capture buckets based on benchmark results. + `enable_encoder_decoder_mixed_cuda_graph` is primarily a performance option. It reduces CPU launch overhead for decoder iterations that mix newly admitted context requests with ongoing generation requests. The option defaults to @@ -509,10 +517,22 @@ attn_backend: TRTLLM dtype: bfloat16 disable_overlap_scheduler: false enable_chunked_prefill: false +max_batch_size: 8 +encoder_max_batch_size: 2 +encoder_max_num_tokens: 1024 max_beam_width: 1 max_input_len: 512 max_num_tokens: 2048 max_seq_len: 512 +cuda_graph_config: + max_batch_size: 8 + enable_padding: true +encoder_cuda_graph_config: + batch_sizes: [1, 2] + num_tokens: [128, 256, 512, 1024] + seq_lens: [128, 256, 512] + enable_padding: true +enable_encoder_decoder_mixed_cuda_graph: true kv_cache_config: enable_block_reuse: false free_gpu_memory_fraction: 0.8 @@ -527,7 +547,6 @@ Start the server: ```bash trtllm-serve google/flan-t5-small \ --backend pytorch \ - --max_batch_size 4 \ --config enc-dec-config.yaml ``` @@ -567,68 +586,36 @@ Use these guidelines as a starting point: - Set `max_seq_len` to at least the larger of the maximum encoder input length and maximum decoded sequence length. The current encoder-decoder runtime uses this value while sizing both phases. -- Set `max_num_tokens` high enough for all encoder tokens admitted together and - for the active decoder tokens. This is especially important for mixed-length - batches. -- Increase `max_batch_size` for more concurrent requests. Beam width multiplies - the number of active decoder sequences but not the number of source requests. +- Set `max_num_tokens` high enough for the active decoder tokens. +- Set `encoder_max_num_tokens` high enough for all encoder tokens in one + encoder microbatch. This is especially important for mixed-length batches. +- Increase `max_batch_size` for more concurrent requests. Start with a smaller + `encoder_max_batch_size`, such as 2 when `max_batch_size=8`, to bound encoder + memory and admission cost without reducing decoder concurrency. Beam width + multiplies the number of active decoder sequences but not the number of + source requests. - Tune `free_gpu_memory_fraction` first, then tune `cross_kv_cache_fraction` based on whether the cross-attention or self-attention pool is exhausted. ## Performance -The following benchmarks compare the PyTorch backend with the legacy TensorRT -encoder-decoder path for large-batch inference. The measurements use BF16 on -one H100 80 GB GPU with greedy decoding, an output limit of 128 tokens, and -mixed encoder input lengths from 260 to 440 tokens. The Flan-T5-XL results are -the average of ten timed runs after three warmup runs. The BART results are the -average of 20 timed runs after five warmup runs. Executed-token throughput -includes the terminal EOS token when a sequence emits it. - -The PyTorch configuration uses the `TRTLLM` attention backend, the overlap -scheduler, the Python scheduler, decoder CUDA graphs with padding, KV cache -manager V1, `max_input_len=512`, `max_seq_len=1024`, and -`max_num_tokens=65536`. Block reuse and chunked prefill are disabled. The KV -cache uses `free_gpu_memory_fraction=0.3` and -`cross_kv_cache_fraction=0.5`. - -The legacy TensorRT configuration uses separate BF16 encoder and decoder -engines built for batch size 128 and beam width 1. The encoder supports 512 -input tokens and 65,536 tokens per iteration; the decoder supports a sequence -length of 129. The benchmark runs these engines through `ModelRunnerCpp` with -greedy `top_k=1` decoding and the same KV cache fractions. For BART, the legacy -TensorRT benchmark starts the decoder with token IDs `[2, 0]` and generates at -most 127 more tokens. The PyTorch LLM API applies the same decoder prefix -internally and counts token ID 0 as the first output token; customers do not -need to provide the decoder prefix. Both paths use token ID 2 as EOS and stop -when the model generates it naturally. If a sequence reaches the output limit, -it retains the model-selected final token and reports a length stop instead of -forcing EOS. This setup also lets beam search begin after the shared decoder -prefix without a per-step Python logits processor. - -### Flan-T5-XL - -For Flan-T5-XL, the PyTorch backend performs on par with the legacy TensorRT -path, with slightly lower latency and higher executed-token throughput across -the tested batch sizes. - -| Batch size | Legacy TensorRT latency | PyTorch latency | PyTorch latency improvement over legacy TensorRT | Legacy TensorRT executed tokens/s | PyTorch executed tokens/s | -| ---: | ---: | ---: | ---: | ---: | ---: | -| 32 | 727.6 ms | 706.1 ms | 3.0% | 3,153 | 3,312 | -| 64 | 1,225.0 ms | 1,136.7 ms | 7.2% | 3,863 | 4,184 | -| 128 | 2,056.8 ms | 1,999.3 ms | 2.8% | 4,601 | 4,768 | - -### BART-large-CNN - -For BART-large-CNN, the PyTorch backend has 21.9% to 36.0% higher latency than -the legacy TensorRT path across the tested batch sizes. - -| Batch size | Legacy TensorRT latency | PyTorch latency | PyTorch latency difference | Legacy TensorRT executed tokens/s | PyTorch executed tokens/s | -| ---: | ---: | ---: | ---: | ---: | ---: | -| 32 | 229.9 ms | 280.2 ms | 21.9% slower | 12,611 | 10,662 | -| 64 | 252.2 ms | 343.0 ms | 36.0% slower | 22,007 | 16,209 | -| 128 | 352.3 ms | 472.7 ms | 34.2% slower | 31,544 | 23,518 | +Configure encoder, decoder, and mixed decoder CUDA graphs for the expected +serving workload. With representative capture buckets, the PyTorch backend can +outperform the legacy TensorRT encoder-decoder path while avoiding the engine +build and checkpoint conversion steps. + +For example, a BF16 FLAN-T5 Large serving benchmark on one H100 80 GB GPU used +encoder CUDA graphs, padded decoder CUDA graphs, and mixed encoder-decoder CUDA +graphs. Compared with the legacy TensorRT path, the PyTorch backend delivered +65.8%, 12.6%, and 11.8% higher request throughput at concurrencies 8, 32, and +64, respectively. P99 latency was 51.6%, 31.0%, and 33.5% lower. + +Follow [Enable encoder and decoder CUDA graphs](#enable-encoder-and-decoder-cuda-graphs) +and choose capture buckets that cover the batch sizes, packed encoder token +counts, and sequence lengths expected in production. Capture grids that omit +common runtime shapes fall back to eager execution and can lose these +performance benefits. Performance depends on the model, request distribution, decoding settings, and GPU configuration. Benchmark with a representative workload before deployment. diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md index 195a80d4b4aa..96bbc0993059 100644 --- a/docs/source/models/supported-models.md +++ b/docs/source/models/supported-models.md @@ -6,6 +6,7 @@ The following is a table of supported models for the PyTorch backend: | Architecture | Model | HuggingFace Example | | ------------------------------------ | ---------------------------------- | -------------------------------------------- | | `AfmoeForCausalLM` | Arcee Foundation MoE (Trinity) | `arcee-ai/Trinity-Mini` | +| `BartForConditionalGeneration` | BART | `facebook/bart-large-cnn` | | `BertForSequenceClassification` | BERT-based | `textattack/bert-base-uncased-yelp-polarity` | | `Cohere2ForCausalLM` | Command A | `CohereLabs/c4ai-command-a-03-2025` | | `DeciLMForCausalLM` | Nemotron | `nvidia/Llama-3_1-Nemotron-51B-Instruct` | @@ -33,6 +34,7 @@ The following is a table of supported models for the PyTorch backend: | `LagunaForCausalLM` | Laguna-XS | `poolside/laguna-XS.2` | | `LlamaForCausalLM` | Llama 3.1, Llama 3, Llama 2, LLaMA | `meta-llama/Meta-Llama-3.1-70B` | | `Llama4ForConditionalGeneration` | Llama 4 | `meta-llama/Llama-4-Scout-17B-16E-Instruct` | +| `MBartForConditionalGeneration` | mBART | `facebook/mbart-large-50-many-to-one-mmt` | | `MiniCPMV4_6ForConditionalGeneration` [^14]| MiniCPM-V 4.6 | `openbmb/MiniCPM-V-4.6` | | `MiniMaxM2ForCausalLM` [^5] | MiniMax M2/M2.1/M2.7 | `MiniMaxAI/MiniMax-M2.7` | | `MiniMaxM3SparseForConditionalGeneration` [^12]| MiniMax-M3 | `MiniMaxAI/MiniMax-M3` | @@ -57,6 +59,8 @@ The following is a table of supported models for the PyTorch backend: | `SkyworkR1V2ForConditionalGeneration` [^5] | Skywork R1V2, Skywork SWE | `Skywork/Skywork-R1V2-38B` | | `SmolLM3ForCausalLM` [^5] | SmolLM3 | `HuggingFaceTB/SmolLM3-3B` | | `Step3p7ForConditionalGeneration` [^8]| Step-3.7-Flash | `stepfun-ai/Step-3.7-Flash` | +| `T5ForConditionalGeneration` | T5, Flan-T5, ByT5 | `google/flan-t5-small` | +| `WhisperForConditionalGeneration` | Whisper | `openai/whisper-large-v3` | ## Model-Feature Support Matrix (Key Models) @@ -95,6 +99,24 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl [^13]: The Cosmos 3 family also supports visual generation through the VisualGen API. See [Visual Generation Models](#visual-generation-models). [^14]: Requires `transformers>=5.7.0`: MiniCPM-V 4.6 was upstreamed into transformers as a native model type (`minicpmv4_6`) and the checkpoint ships no remote code (`auto_map`) to fall back on. The Qwen3.5-hybrid text tower runs in BF16. Image, video, and text inputs are supported in this release (video reuses the same NaViT-packed vision path as image via `MiniCPMV4_6InputProcessor`). +# Encoder-Decoder Feature Support Matrix (PyTorch Backend) + +The following capabilities apply to the supported encoder-decoder architectures. For configuration guidance and +limitations, see [Use encoder-decoder models with the PyTorch backend](./encoder-decoder.md). + +| Model Architecture/Feature | Overlap Scheduler | Decoder CUDA Graph | Encoder CUDA Graph | KV Cache Manager V1 | KV Cache Manager V2 | Beam Search | Tensor Parallelism | Pipeline Parallelism | Chunked Prefill | +| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | +| `BartForConditionalGeneration` | Yes | Yes (except FP32) | Yes | Yes | Yes (single beam) | Yes (V1 only) | Yes | No | No (encoder phase) | +| `MBartForConditionalGeneration` | Yes | Yes (except FP32) | Yes | Yes | Yes (single beam) | Yes (V1 only) | Yes | No | No (encoder phase) | +| `T5ForConditionalGeneration` | Yes | Yes (except FP32) | Yes | Yes | Yes (single beam) | Yes (V1 only) | Yes | No | No (encoder phase) | +| `WhisperForConditionalGeneration` | Yes | Yes (except FP32) | No (feature inputs) | Yes | Yes (single beam) | Yes (V1 only) | Yes | No | No (encoder phase) | + +Decoder CUDA graphs support greedy and beam-search decoding with KV cache manager V1 and single-beam decoding with +V2. Encoder CUDA graphs support the token-input BART, mBART, and T5 families; Whisper's feature-driven audio encoder +runs eagerly. Use the `TRTLLM` attention backend for encoder-decoder models; tensor parallelism also requires attention +head counts divisible by the tensor parallel size. Chunked prefill is not supported for the encoder phase, so the +complete encoder input must fit in the iteration token budget. + # Multimodal Feature Support Matrix (PyTorch Backend) | Model Architecture/Feature | Overlap Scheduler | CUDA Graph | Chunked Prefill | Torch Sampler | TLLM C++ Sampler | KV Cache Reuse | Logits Post Processor | EPD Disaggregated Serving | Modality | diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index 68127ce2efd1..56af4c1e56a5 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -673,9 +673,17 @@ def prepare_encoder_decoder(self, prompt_lens: torch.Tensor, self.host_request_types[:self.num_contexts].fill_(0) self.host_request_types[self.num_contexts:num_seqs].fill_(1) + max_blocks = None + if self.kv_cache_manager.tokens_per_block: + max_blocks = ceil_div(max_kv_len, + self.kv_cache_manager.tokens_per_block) self.kv_cache_manager.copy_batch_block_offsets( - self.kv_cache_block_offsets, self.request_ids, self.beam_width, - self.num_contexts, num_seqs) + self.kv_cache_block_offsets, + self.request_ids, + self.beam_width, + self.num_contexts, + num_seqs, + max_blocks=max_blocks) self._bind_runtime_views( kv_lens_cuda=self.kv_lens_cuda[:num_seqs], kv_lens=kv_lens, diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index 97819bfe4e0b..d4567e5c6e8d 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -207,11 +207,20 @@ def _get_static_encoder_hidden_states( raise RuntimeError( "Mixed encoder-decoder CUDA graph replay requires the " "encoder hidden-state buffer initialized during warmup.") + if self.graphs: + raise RuntimeError( + "Mixed encoder-decoder CUDA graph encoder hidden-state " + "buffer cannot be allocated after graph capture.") static_encoder_hidden_states = encoder_hidden_states.new_empty( (num_encoder_tokens, encoder_hidden_states.shape[1])) self.shared_static_tensors[ "encoder_hidden_states"] = static_encoder_hidden_states + if static_encoder_hidden_states.shape[0] < num_encoder_tokens: + raise RuntimeError( + "Mixed encoder-decoder CUDA graph encoder hidden-state buffer " + f"has capacity {static_encoder_hidden_states.shape[0]}, but " + f"{num_encoder_tokens} tokens were requested.") return static_encoder_hidden_states[:num_encoder_tokens] def _is_mixed_encoder_decoder_batch(self, batch: ScheduledRequests) -> bool: @@ -852,6 +861,8 @@ def clear(self): EncoderKeyType: TypeAlias = Tuple[int, int, int] +_ENCODER_SOURCE_SEQ_LENS = "_encoder_source_seq_lens" +_ENCODER_SOURCE_TO_SLOT = "_encoder_source_to_slot" @dataclass @@ -868,6 +879,7 @@ class EncoderCUDAGraphRunnerConfig: max_seq_len: int cuda_graph_mem_pool: Any is_encoder_decoder: bool = False + use_fixed_sequence_slots: bool = False class EncoderCUDAGraphRunner: @@ -895,6 +907,7 @@ def __init__(self, config: EncoderCUDAGraphRunnerConfig): self.max_supported_num_tokens = config.max_cuda_graph_num_tokens self.supported_seq_lens = sorted(config.cuda_graph_seq_lens) self.is_encoder_decoder = config.is_encoder_decoder + self.use_fixed_sequence_slots = config.use_fixed_sequence_slots self.capture_keys: frozenset[EncoderKeyType] = frozenset() self._capture_sequence_lengths: Dict[EncoderKeyType, List[int]] = {} if self.is_encoder_decoder: @@ -996,7 +1009,13 @@ def build_capture_sequence_lengths(batch_size: int, num_tokens: int, def _build_encoder_decoder_capture_layouts( self) -> Dict[EncoderKeyType, List[int]]: - """Derive reachable encoder-decoder capture keys via get_graph_key.""" + """Map each reachable graph key to one physical sequence-slot layout. + + Sequence layout is deliberately not part of ``EncoderKeyType``. Multiple + configured shapes may normalize to the same three-dimensional key, so + the first deterministic layout becomes that graph's fixed slot + capacities. Runtime sequences are assigned to those slots during replay. + """ capture_layouts: Dict[EncoderKeyType, List[int]] = {} for batch_size in self.supported_batch_sizes: for num_tokens in self.supported_num_tokens: @@ -1009,25 +1028,37 @@ def _build_encoder_decoder_capture_layouts( key, _, is_valid = self.get_graph_key( {"seq_lens": sequence_lengths}) if is_valid: + # Capture at most one graph/layout for a normalized key; + # alternative runtime layouts do not create more keys. capture_layouts.setdefault(key, sequence_lengths) return capture_layouts def _get_dynamic_capture_key( self, - batch_size: int, - num_tokens: int, - max_seq_len: int, + sequence_lengths: List[int], allow_batch_padding: bool, ) -> Optional[EncoderKeyType]: - """Return the smallest compatible dynamic-layout capture key.""" + """Return the smallest key whose token bucket and slots fit the batch. + + Keys are ordered from smaller to larger buckets. For fixed-slot replay, + aggregate token and maximum-length checks are insufficient: every + runtime sequence (including batch-padding dummies) must also fit in a + distinct capture-time slot. An incompatible key is skipped in favor of + a larger existing key; no new layout-specific key is created. + """ + batch_size = len(sequence_lengths) + sum(sequence_lengths) + max_seq_len = max(sequence_lengths) if sequence_lengths else 0 candidate_batch_sizes = (self.supported_batch_sizes if allow_batch_padding else [batch_size]) for padded_batch_size in candidate_batch_sizes: if padded_batch_size < batch_size: continue - required_num_tokens = num_tokens + padded_batch_size - batch_size + padded_sequence_lengths = (sequence_lengths + [1] * + (padded_batch_size - batch_size)) + required_num_tokens = sum(padded_sequence_lengths) for key in self._capture_keys_by_batch_size.get( padded_batch_size, []): _, padded_num_tokens, padded_max_seq_len = key @@ -1038,16 +1069,72 @@ def _get_dynamic_capture_key( or padded_num_tokens > padded_batch_size * padded_max_seq_len): continue + if (self.use_fixed_sequence_slots + and self._get_sequence_slot_mapping( + key, padded_sequence_lengths) is None): + # The batch fits the aggregate bucket but not this key's + # individual slot capacities. Try the next captured key. + continue return key return None + def _get_sequence_slot_mapping( + self, + key: EncoderKeyType, + sequence_lengths: List[int], + ) -> Optional[List[int]]: + """Assign each runtime sequence to one compatible physical graph slot. + + The returned list maps source request index to capture slot index. It + changes only physical placement: source request order is retained + separately and restored after replay. + """ + capture_lengths = self._capture_sequence_lengths.get(key) + if (capture_lengths is None + or len(capture_lengths) != len(sequence_lengths)): + return None + + # Preserve physical order when every request already fits its + # corresponding slot, avoiding unnecessary scatter/gather permutation. + if all(sequence_length <= capture_length + for sequence_length, capture_length in zip( + sequence_lengths, capture_lengths)): + return list(range(len(sequence_lengths))) + + # Largest-to-largest matching is sufficient for one-to-one scalar + # capacities: if any sorted request exceeds its paired slot, no + # permutation can make the layout fit. + sequence_order = sorted(range(len(sequence_lengths)), + key=lambda index: + (-sequence_lengths[index], index)) + capture_order = sorted(range(len(capture_lengths)), + key=lambda index: + (-capture_lengths[index], index)) + source_to_slot = [0] * len(sequence_lengths) + for source_index, slot_index in zip(sequence_order, capture_order): + if sequence_lengths[source_index] > capture_lengths[slot_index]: + return None + source_to_slot[source_index] = slot_index + return source_to_slot + def get_capture_warmup_sequence_lengths( self, key: EncoderKeyType) -> Optional[List[int]]: """Return the representative sequence layout for a capture key.""" sequence_lengths = self._capture_sequence_lengths.get(key) return list(sequence_lengths) if sequence_lengths is not None else None + def _get_capture_sequence_offsets(self, key: EncoderKeyType) -> List[int]: + """Return cumulative fixed-slot offsets for a capture layout.""" + offsets = [0] + for sequence_length in self._capture_sequence_lengths[key]: + offsets.append(offsets[-1] + sequence_length) + if offsets[-1] != key[1]: + raise ValueError( + f"Encoder CUDA graph layout for key {key} contains " + f"{offsets[-1]} tokens.") + return offsets + def _get_valid_graph_key(self, batch_size: int, num_tokens: int, max_seq_len: int) -> EncoderKeyType: num_tokens_idx = bisect.bisect_left(self.supported_num_tokens, @@ -1087,9 +1174,7 @@ def get_graph_key( if self.is_encoder_decoder: if self.padding_enabled and self.capture_keys: padded_key = self._get_dynamic_capture_key( - batch_size, - num_tokens, - max_seq_len, + seq_lens, allow_batch_padding=False, ) if padded_key is None: @@ -1139,9 +1224,7 @@ def pad_batch(self, inputs: Dict[str, Any], if self.is_encoder_decoder and self.capture_keys: seq_lens = inputs['seq_lens'] padded_key = self._get_dynamic_capture_key( - batch_size, - sum(seq_lens), - max(seq_lens) if seq_lens else 0, + seq_lens, allow_batch_padding=True, ) padded_batch_size = padded_key[0] if padded_key is not None else 0 @@ -1169,6 +1252,48 @@ def pad_batch(self, inputs: Dict[str, Any], yield padded_inputs + def prepare_encoder_decoder_inputs( + self, + inputs: Dict[str, Any], + key: EncoderKeyType, + source_sequence_lengths: List[int], + ) -> Dict[str, Any]: + """Arrange runtime sequence metadata in capture-time slot order.""" + if not self.is_encoder_decoder: + return inputs + + if not self.use_fixed_sequence_slots: + prepared_inputs = dict(inputs) + prepared_inputs[_ENCODER_SOURCE_SEQ_LENS] = list( + source_sequence_lengths) + return prepared_inputs + + sequence_lengths = inputs["seq_lens"] + if (sequence_lengths[:len(source_sequence_lengths)] + != source_sequence_lengths): + raise ValueError("Encoder source sequence lengths must be the " + "unpadded prefix of graph sequence lengths.") + + source_to_slot = self._get_sequence_slot_mapping(key, sequence_lengths) + if source_to_slot is None: + raise ValueError( + f"Encoder sequence lengths {sequence_lengths} are not " + f"compatible with CUDA graph key {key}.") + + # Attention metadata follows physical slot order, while the packed + # source tensors and the final returned output remain in request order. + slot_sequence_lengths = [0] * len(sequence_lengths) + for source_index, slot_index in enumerate(source_to_slot): + slot_sequence_lengths[slot_index] = sequence_lengths[source_index] + + prepared_inputs = dict(inputs) + prepared_inputs["seq_lens"] = slot_sequence_lengths + prepared_inputs[_ENCODER_SOURCE_SEQ_LENS] = list( + source_sequence_lengths) + prepared_inputs[_ENCODER_SOURCE_TO_SLOT] = source_to_slot[:len( + source_sequence_lengths)] + return prepared_inputs + def maybe_get_cuda_graph( self, inputs: Dict[str, Any], @@ -1259,6 +1384,16 @@ def maybe_get_cuda_graph( # be pinned or pageable; only captured H2D copies require pinned memory. graph_attn_metadata.bind_encoder_cuda_graph_seq_lens( self.shared_static_tensors_cpu["seq_lens"], padded_batch_size) + if self.use_fixed_sequence_slots: + # CUDA graph replay keeps each request in its capture-time token + # slot. Explicit boundaries let attention combine those fixed + # offsets with the per-replay logical sequence lengths above. + capture_offsets = self._get_capture_sequence_offsets(key) + capture_offsets_cuda = torch.tensor(capture_offsets, + dtype=torch.int32, + device="cuda") + graph_attn_metadata.cu_q_seqlens = capture_offsets_cuda + graph_attn_metadata.cu_kv_seqlens = capture_offsets_cuda graph_attn_metadata.max_seq_len = self.config.max_seq_len graph_attn_metadata.request_ids = list(range(padded_batch_size)) @@ -1286,6 +1421,10 @@ def _stage_inputs(self, key: EncoderKeyType, inputs: Dict[str, # is not captured, stage directly into the graph-resident CUDA buffers. static_tensors = self.shared_static_tensors_cpu if self._capture_h2d_copy else self.shared_static_tensors + if self.is_encoder_decoder and _ENCODER_SOURCE_TO_SLOT in inputs: + self._stage_encoder_decoder_inputs(key, inputs, static_tensors) + return + input_ids = inputs["input_ids"] if isinstance(input_ids, list): actual_tokens = len(input_ids) @@ -1322,6 +1461,88 @@ def _stage_inputs(self, key: EncoderKeyType, inputs: Dict[str, staged_position_ids[offset:padded_num_tokens].fill_(0) + def _stage_encoder_decoder_inputs( + self, + key: EncoderKeyType, + inputs: Dict[str, Any], + static_tensors: Dict[str, torch.Tensor], + ) -> None: + """Scatter packed request inputs into fixed capture-time slots.""" + source_sequence_lengths = inputs[_ENCODER_SOURCE_SEQ_LENS] + source_to_slot = inputs[_ENCODER_SOURCE_TO_SLOT] + + input_ids = inputs["input_ids"] + if isinstance(input_ids, list): + source_input_ids = torch.tensor(input_ids, dtype=torch.int32) + elif isinstance(input_ids, torch.Tensor): + source_input_ids = input_ids + else: + raise TypeError(f"Unsupported input_ids type: {type(input_ids)}") + + actual_num_tokens = sum(source_sequence_lengths) + if int(source_input_ids.shape[0]) != actual_num_tokens: + raise ValueError( + "Packed encoder input IDs must match source sequence lengths.") + + position_ids = inputs.get("position_ids") + if isinstance(position_ids, list): + source_position_ids = torch.tensor(position_ids, dtype=torch.int32) + elif isinstance(position_ids, torch.Tensor): + source_position_ids = position_ids.flatten() + elif position_ids is None: + source_position_ids = None + else: + raise TypeError( + f"Unsupported position_ids type: {type(position_ids)}") + if (source_position_ids is not None + and int(source_position_ids.shape[0]) != actual_num_tokens): + raise ValueError("Packed encoder position IDs must match source " + "sequence lengths.") + + static_tensors["input_ids"][:key[1]].zero_() + static_tensors["position_ids"][:, :key[1]].zero_() + + capture_offsets = self._get_capture_sequence_offsets(key) + + source_offset = 0 + for source_index, sequence_length in enumerate(source_sequence_lengths): + slot_index = source_to_slot[source_index] + destination_offset = capture_offsets[slot_index] + source_slice = slice(source_offset, source_offset + sequence_length) + destination_slice = slice(destination_offset, + destination_offset + sequence_length) + static_tensors["input_ids"][destination_slice].copy_( + source_input_ids[source_slice]) + if source_position_ids is None: + static_tensors["position_ids"][0, destination_slice].copy_( + self._arange_max[:sequence_length]) + else: + static_tensors["position_ids"][0, destination_slice].copy_( + source_position_ids[source_slice]) + source_offset += sequence_length + + def restore_encoder_decoder_output( + self, + key: EncoderKeyType, + output: torch.Tensor, + inputs: Dict[str, Any], + ) -> torch.Tensor: + """Compact fixed-slot graph output back into request order.""" + source_sequence_lengths = inputs[_ENCODER_SOURCE_SEQ_LENS] + if _ENCODER_SOURCE_TO_SLOT not in inputs: + return output[:sum(source_sequence_lengths)].clone() + + source_to_slot = inputs[_ENCODER_SOURCE_TO_SLOT] + + capture_offsets = self._get_capture_sequence_offsets(key) + + output_slices = [] + for source_index, sequence_length in enumerate(source_sequence_lengths): + source_offset = capture_offsets[source_to_slot[source_index]] + output_slices.append(output[source_offset:source_offset + + sequence_length]) + return torch.cat(output_slices, dim=0) + def capture( self, key: EncoderKeyType, diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 9ecaa2ca3420..188a8f2685f9 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -563,7 +563,7 @@ def __init__( if self._cuda_graph_seq_lens else 0) use_encoder_cuda_graph = ((self._is_encoder_decoder_model() - or self._is_encode_only()) + or self._is_encode_only) and self.encoder_cuda_graph_config is not None and bool(self._cuda_graph_num_tokens) and bool(self._cuda_graph_seq_lens)) @@ -820,6 +820,10 @@ def __init__( max_seq_len=self.max_seq_len, cuda_graph_mem_pool=self._cuda_graph_mem_pool, is_encoder_decoder=self._is_encoder_decoder_model(), + use_fixed_sequence_slots=(self._is_encoder_decoder_model() + and hasattr( + pretrained_config, + "relative_attention_num_buckets")), ) self.encoder_cuda_graph_runner = EncoderCUDAGraphRunner( encoder_cuda_graph_runner_config) @@ -2176,6 +2180,15 @@ def _capture_mixed_encoder_decoder_cuda_graphs( operation = ("warmup" if runner.is_warmup_only else "capture") hidden_size = self._get_enc_dec_hidden_size() + max_num_encoder_tokens = max( + (total_encoder_tokens + for num_contexts, total_encoder_tokens in context_shapes + if total_encoder_tokens <= num_contexts * max_encoder_output_len + and any(batch_size > num_contexts + for batch_size in runner.supported_batch_sizes)), + default=0) + if max_num_encoder_tokens == 0: + return model_config = self.model.model_config.pretrained_config # BART/mBART prepend a forced BOS token after decoder_start; T5 uses # decoder_start alone. Match the LLM API's decoder-prefix construction. @@ -2226,6 +2239,11 @@ def _capture_mixed_encoder_decoder_cuda_graphs( ) request.py_skip_cross_kv_projection = False + runner._get_static_encoder_hidden_states( + context_requests[0].py_encoder_output, + max_num_encoder_tokens, + allow_allocate=True, + ) logger.info("Run mixed encoder-decoder CUDA graph " f"{operation} for batch size={batch_size}, " f"context requests={num_contexts}, " @@ -3543,19 +3561,19 @@ def _can_use_encoder_decoder_input_fast_path( static_eligible = ( hasattr(batch_manager_bindings, "prepare_encoder_decoder_inputs") - and self._is_encoder_decoder_model() - and not self.enable_spec_decode and not self.is_draft_model + and self._is_encoder_decoder_model() and not self.is_draft_model and self.max_beam_width == 1 and self.sparse_attention_config is None and not self.use_mrope and not self.enable_attention_dp and not self.mapping.has_cp_helix() and not self.is_multimodal - and self.lora_model_config is None and not self.attn_runtime_features.chunked_prefill and not self.attn_runtime_features.cache_reuse and not self.attn_runtime_features.has_speculative_draft_tokens) self._encoder_decoder_input_fast_path_static_eligible = \ static_eligible - if (not static_eligible or new_tokens_device is None + if (not static_eligible or self.enable_spec_decode + or self.lora_model_config is not None + or new_tokens_device is None or next_draft_tokens_device is not None or self.guided_decoder is not None): return False @@ -7309,7 +7327,6 @@ def _forward_encoder_with_cuda_graph( input_ids = inputs.get('encoder_input_ids_host') position_ids = inputs.get('encoder_position_ids_host') seq_lens = inputs['encoder_seq_lens'] - actual_num_tokens = sum(seq_lens) runner = self.encoder_cuda_graph_runner if input_ids is None or position_ids is None: @@ -7338,12 +7355,11 @@ def _forward_encoder_with_cuda_graph( # the previous captured H2D before updating seq_lens or any other # shared host input for this replay. runner.retire_staging() + model_inputs = runner.prepare_encoder_decoder_inputs( + padded_runner_inputs, key, seq_lens) graph_attn_metadata.prepare_encoder_cuda_graph_replay( - padded_runner_inputs['seq_lens'], key[1]) - model_inputs = { - **padded_runner_inputs, - 'attn_metadata': graph_attn_metadata, - } + model_inputs['seq_lens'], key[1]) + model_inputs['attn_metadata'] = graph_attn_metadata moe_load_balancer: MoeLoadBalancer = getattr( self, 'moe_load_balancer', None) @@ -7369,7 +7385,8 @@ def capture_forward_fn( if not isinstance(graph_outputs, torch.Tensor): raise TypeError("Encoder-decoder CUDA graph replay must return " "a tensor of encoder hidden states.") - return graph_outputs[:actual_num_tokens].clone() + return runner.restore_encoder_decoder_output(key, graph_outputs, + model_inputs) @nvtx_range("forward_encoder") def forward_encoder( diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 0561575e5ba7..ac142318de07 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -5375,7 +5375,13 @@ def _waiting_encoder_requests( for request in encoder_requests) num_scheduled_tokens += sum(1 + request.num_draft_tokens for request in generation_requests) - should_wait = (self.encoder_batch_wait_iters_count + has_decoder_work = bool(context_requests or generation_requests) + if not has_decoder_work: + has_decoder_work = any( + request.request_id in self.inflight_req_ids + and request.state != LlmRequestState.ENCODER_INIT + for request in self.active_requests) + should_wait = (has_decoder_work and self.encoder_batch_wait_iters_count < self.batch_wait_timeout_iters and num_scheduled_tokens < self.batch_wait_max_tokens_ratio * self.max_num_tokens) if should_wait: @@ -5477,7 +5483,7 @@ def _warmup_encoder_decoder_encoder_cuda_graphs(self) -> None: warmup(self.resource_manager) def _submit_encoder_step(self, encoder_requests: List[LlmRequest]) -> None: - """Queue encoder work without blocking the decoder executor thread.""" + """Queue encoder work, serializing it with decoder work under TP.""" executor = self.encoder_launch_executor if executor is None: raise RuntimeError("Encoder launch executor is unavailable.") @@ -5486,6 +5492,10 @@ def _submit_encoder_step(self, encoder_requests: List[LlmRequest]) -> None: for request in requests: self.inflight_req_ids.insert(request.request_id) + serialize_tp = self.dist.tp_size > 1 + if serialize_tp: + self.encoder_stream.wait_stream(self.execution_stream) + try: future = executor.submit(self._run_encoder_step_unchecked, requests) except Exception: @@ -5493,6 +5503,22 @@ def _submit_encoder_step(self, encoder_requests: List[LlmRequest]) -> None: self.inflight_req_ids.erase(request.request_id) raise + if serialize_tp: + # Encoder and decoder forwards share TP communicators and + # workspaces. Complete encoder GPU work before the caller launches + # decoder work, while retaining the worker thread affinity needed + # by encoder CUDA graph replay. + try: + result = future.result() + result.ready_event.synchronize() + self._publish_encoder_step(requests, result) + except Exception as e: + self._finish_failed_encoder_step(requests, e) + return + for request in requests: + self.inflight_req_ids.erase(request.request_id) + return + self.pending_encoder_steps.append( PendingEncoderStep(requests=requests, future=future)) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index beb41ce2846c..8d959a643b23 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -3676,6 +3676,7 @@ def _process_requests( and get_draft_token_length(request) == 0 and request._py_embedding_bias_1d is None and not getattr(request, "py_bad_words", None) + and not getattr(request, "py_no_repeat_ngram_size", None) and not request.py_min_length and not request.py_return_log_probs and not request.py_stop_words_list diff --git a/tests/integration/defs/llmapi/test_llm_api_pytorch_bart.py b/tests/integration/defs/llmapi/test_llm_api_pytorch_bart.py index af48a095b2ff..8185497178ed 100644 --- a/tests/integration/defs/llmapi/test_llm_api_pytorch_bart.py +++ b/tests/integration/defs/llmapi/test_llm_api_pytorch_bart.py @@ -606,12 +606,23 @@ def test_bart_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch( ) +@pytest.mark.parametrize( + "tensor_parallel_size", + [ + pytest.param(1, id="tp1"), + pytest.param(2, id="tp2", marks=pytest.mark.skip_less_device(2)), + ], +) def test_bart_pytorch_continuous_admission_replays_encoder_and_mixed_cuda_graphs( monkeypatch: pytest.MonkeyPatch, + tensor_parallel_size: int, ) -> None: - """Admit encoder work while an older request remains in decoder generation.""" + """Preserve mixed-length encoder outputs during continuous admission.""" monkeypatch.setenv("TRTLLM_SKIP_KV_CACHE_ESTIMATION", "1") - monkeypatch.setenv("TLLM_WORKER_USE_SINGLE_PROCESS", "1") + if tensor_parallel_size == 1: + monkeypatch.setenv("TLLM_WORKER_USE_SINGLE_PROCESS", "1") + else: + monkeypatch.delenv("TLLM_WORKER_USE_SINGLE_PROCESS", raising=False) model_path = _get_model_path(_MODEL_NAME) tokenizer = AutoTokenizer.from_pretrained(model_path) @@ -626,8 +637,8 @@ def test_bart_pytorch_continuous_admission_replays_encoder_and_mixed_cuda_graphs num_return_sequences=1, ) encoder_graph_config = EncodeCudaGraphConfig( - batch_sizes=[1], - num_tokens=[64], + batch_sizes=[1, 2], + num_tokens=[64, 128], seq_lens=[64], enable_padding=True, ) @@ -636,14 +647,14 @@ def test_bart_pytorch_continuous_admission_replays_encoder_and_mixed_cuda_graphs model_path, backend="pytorch", attn_backend="TRTLLM", - cuda_graph_config=_decoder_cuda_graph_config([2]), + cuda_graph_config=_decoder_cuda_graph_config([3]), encoder_cuda_graph_config=encoder_graph_config, enable_encoder_decoder_mixed_cuda_graph=True, disable_overlap_scheduler=False, dtype="bfloat16", enable_chunked_prefill=False, - encoder_max_batch_size=1, - encoder_max_num_tokens=64, + encoder_max_batch_size=2, + encoder_max_num_tokens=128, kv_cache_config=KvCacheConfig( enable_block_reuse=False, max_tokens=_MAX_KV_TOKENS, @@ -651,38 +662,41 @@ def test_bart_pytorch_continuous_admission_replays_encoder_and_mixed_cuda_graphs cross_kv_cache_fraction=_CROSS_KV_CACHE_FRACTION, use_kv_cache_manager_v2=False, ), - max_batch_size=2, + max_batch_size=3, max_beam_width=1, max_input_len=_MAX_SEQUENCE_LENGTH, max_num_tokens=_MAX_SEQUENCE_LENGTH, max_seq_len=_MAX_SEQUENCE_LENGTH, model_kwargs={"torch_dtype": "bfloat16"}, scheduler_config=SchedulerConfig(use_python_scheduler=True), + batch_wait_timeout_iters=2, + tensor_parallel_size=tensor_parallel_size, ) as llm: - model_engine = llm._executor.engine.model_engine - encoder_runner = model_engine.encoder_cuda_graph_runner - decoder_runner = model_engine.cuda_graph_runner - - assert encoder_runner.enabled - assert encoder_runner.graphs - captured_mixed_keys = {key for key in decoder_runner.graphs if key[5] and key[6]} - assert captured_mixed_keys - encoder_replay_keys = [] decoder_replay_keys = [] - original_encoder_replay = encoder_runner.replay - original_decoder_replay = decoder_runner.replay + if tensor_parallel_size == 1: + model_engine = llm._executor.engine.model_engine + encoder_runner = model_engine.encoder_cuda_graph_runner + decoder_runner = model_engine.cuda_graph_runner + + assert encoder_runner.enabled + assert encoder_runner.graphs + captured_mixed_keys = {key for key in decoder_runner.graphs if key[5] and key[6]} + assert captured_mixed_keys - def record_encoder_replay(key, inputs): - encoder_replay_keys.append(key) - return original_encoder_replay(key, inputs) + original_encoder_replay = encoder_runner.replay + original_decoder_replay = decoder_runner.replay - def record_decoder_replay(key, inputs): - decoder_replay_keys.append(key) - return original_decoder_replay(key, inputs) + def record_encoder_replay(key, inputs): + encoder_replay_keys.append(key) + return original_encoder_replay(key, inputs) - monkeypatch.setattr(encoder_runner, "replay", record_encoder_replay) - monkeypatch.setattr(decoder_runner, "replay", record_decoder_replay) + def record_decoder_replay(key, inputs): + decoder_replay_keys.append(key) + return original_decoder_replay(key, inputs) + + monkeypatch.setattr(encoder_runner, "replay", record_encoder_replay) + monkeypatch.setattr(decoder_runner, "replay", record_decoder_replay) first_response = llm.generate_async( _SOURCE_TEXT, @@ -692,34 +706,42 @@ def record_decoder_replay(key, inputs): first_stream_step = next(first_response) assert not first_stream_step.finished - second_response = llm.generate_async( - _MIXED_ENCODER_SOURCE_TEXTS[1], - sampling_params=second_sampling_params, - streaming=False, - ) + encoder_replay_count_before_admission = len(encoder_replay_keys) + admitted_responses = [ + llm.generate_async( + source_text, + sampling_params=second_sampling_params, + streaming=False, + ) + for source_text in _MIXED_ENCODER_SOURCE_TEXTS + ] first_response.result() - second_response.result() + for response in admitted_responses: + response.result() first_token_ids = _assert_bart_response( first_response, num_return_sequences=1, max_tokens=_CONTINUOUS_ADMISSION_MAX_NEW_TOKENS, ) - second_token_ids = _assert_bart_response( - second_response, - num_return_sequences=1, - ) assert first_token_ids[0][:_MAX_NEW_TOKENS] == _EXPECTED_GREEDY_OUTPUT_TOKEN_IDS - _assert_expected_generation( - tokenizer, - second_token_ids, - exact_match=True, - expected_token_ids_by_output=_MIXED_ENCODER_EXPECTED_TOKEN_IDS_BY_REQUEST[1], - ) - assert len(encoder_replay_keys) >= 2 - assert set(encoder_replay_keys) <= set(encoder_runner.graphs) - replayed_mixed_keys = {key for key in decoder_replay_keys if key[5] and key[6]} - assert replayed_mixed_keys - assert replayed_mixed_keys <= captured_mixed_keys + for request_idx, response in enumerate(admitted_responses): + token_ids = _assert_bart_response(response, num_return_sequences=1) + _assert_expected_generation( + tokenizer, + token_ids, + exact_match=True, + expected_token_ids_by_output=( + _MIXED_ENCODER_EXPECTED_TOKEN_IDS_BY_REQUEST[request_idx] + ), + ) + + if tensor_parallel_size == 1: + admitted_encoder_keys = encoder_replay_keys[encoder_replay_count_before_admission:] + assert any(key[0] == 2 for key in admitted_encoder_keys) + assert set(encoder_replay_keys) <= set(encoder_runner.graphs) + replayed_mixed_keys = {key for key in decoder_replay_keys if key[5] and key[6]} + assert replayed_mixed_keys + assert replayed_mixed_keys <= captured_mixed_keys diff --git a/tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py b/tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py index 4c93c0b27c93..46f525565ec8 100644 --- a/tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py +++ b/tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py @@ -823,7 +823,7 @@ def test_t5_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch( def test_t5_pytorch_continuous_admission_replays_encoder_and_mixed_cuda_graphs( monkeypatch: pytest.MonkeyPatch, ) -> None: - """Continuously admit work and replay encoder and mixed decoder graphs.""" + """Preserve mixed-length encoder outputs during continuous admission.""" monkeypatch.setenv("TRTLLM_SKIP_KV_CACHE_ESTIMATION", "1") monkeypatch.setenv("TLLM_WORKER_USE_SINGLE_PROCESS", "1") @@ -841,8 +841,8 @@ def test_t5_pytorch_continuous_admission_replays_encoder_and_mixed_cuda_graphs( temperature=0.0, ) encoder_graph_config = EncodeCudaGraphConfig( - batch_sizes=[1], - num_tokens=[64], + batch_sizes=[1, 2], + num_tokens=[64, 128], seq_lens=[64], enable_padding=True, ) @@ -851,14 +851,14 @@ def test_t5_pytorch_continuous_admission_replays_encoder_and_mixed_cuda_graphs( model_path, backend="pytorch", attn_backend="TRTLLM", - cuda_graph_config=_decoder_cuda_graph_config([2]), + cuda_graph_config=_decoder_cuda_graph_config([3]), encoder_cuda_graph_config=encoder_graph_config, enable_encoder_decoder_mixed_cuda_graph=True, disable_overlap_scheduler=False, dtype="bfloat16", enable_chunked_prefill=False, - encoder_max_batch_size=1, - encoder_max_num_tokens=64, + encoder_max_batch_size=2, + encoder_max_num_tokens=128, kv_cache_config=KvCacheConfig( enable_block_reuse=False, max_tokens=_MAX_KV_TOKENS, @@ -866,13 +866,14 @@ def test_t5_pytorch_continuous_admission_replays_encoder_and_mixed_cuda_graphs( cross_kv_cache_fraction=_CROSS_KV_CACHE_FRACTION, use_kv_cache_manager_v2=False, ), - max_batch_size=2, + max_batch_size=3, max_beam_width=1, max_input_len=_MAX_SEQUENCE_LENGTH, max_num_tokens=_MAX_SEQUENCE_LENGTH, max_seq_len=_MAX_SEQUENCE_LENGTH, model_kwargs={"torch_dtype": "bfloat16"}, scheduler_config=SchedulerConfig(use_python_scheduler=True), + batch_wait_timeout_iters=2, ) as llm: model_engine = llm._executor.engine.model_engine encoder_runner = model_engine.encoder_cuda_graph_runner @@ -880,6 +881,7 @@ def test_t5_pytorch_continuous_admission_replays_encoder_and_mixed_cuda_graphs( assert encoder_runner.enabled assert encoder_runner.graphs + assert encoder_runner.use_fixed_sequence_slots captured_mixed_keys = {key for key in decoder_runner.graphs if key[5] and key[6]} assert captured_mixed_keys @@ -899,6 +901,9 @@ def record_decoder_replay(key, inputs): monkeypatch.setattr(encoder_runner, "replay", record_encoder_replay) monkeypatch.setattr(decoder_runner, "replay", record_decoder_replay) + expected_token_ids_by_request = _MIXED_ENCODER_OUTPUT_TOKEN_IDS_BY_MODEL_AND_BEAMS[ + (model_name, 1) + ] first_response = llm.generate_async( _SOURCE_TEXT, sampling_params=first_sampling_params, @@ -907,35 +912,41 @@ def record_decoder_replay(key, inputs): first_stream_step = next(first_response) assert not first_stream_step.finished - second_response = llm.generate_async( - _MIXED_ENCODER_SOURCE_TEXTS[1], - sampling_params=second_sampling_params, - streaming=False, - ) + encoder_replay_count_before_admission = len(encoder_replay_keys) + admitted_responses = [ + llm.generate_async( + source_text, + sampling_params=second_sampling_params, + streaming=False, + ) + for source_text in _MIXED_ENCODER_SOURCE_TEXTS + ] first_response.result() - second_response.result() + for response in admitted_responses: + response.result() first_token_ids = _assert_t5_response( first_response, num_return_sequences=1, max_tokens=_MIXED_CONTEXT_GENERATION_MAX_NEW_TOKENS, ) - second_token_ids = _assert_t5_response(second_response, num_return_sequences=1) - - expected_token_ids_by_request = _MIXED_ENCODER_OUTPUT_TOKEN_IDS_BY_MODEL_AND_BEAMS[ - (model_name, 1) - ] assert first_token_ids[0][:_MAX_NEW_TOKENS] == expected_token_ids_by_request[0][0] - _assert_expected_generation( - tokenizer, - second_token_ids, - exact_match=True, - expected_token_ids_by_output=expected_token_ids_by_request[1], - expected_text_fragment=_MIXED_ENCODER_EXPECTED_TEXT_FRAGMENTS_BY_MODEL[model_name][1], - ) - assert len(encoder_replay_keys) >= 2 + for request_idx, response in enumerate(admitted_responses): + token_ids = _assert_t5_response(response, num_return_sequences=1) + _assert_expected_generation( + tokenizer, + token_ids, + exact_match=True, + expected_token_ids_by_output=expected_token_ids_by_request[request_idx], + expected_text_fragment=( + _MIXED_ENCODER_EXPECTED_TEXT_FRAGMENTS_BY_MODEL[model_name][request_idx] + ), + ) + + admitted_encoder_keys = encoder_replay_keys[encoder_replay_count_before_admission:] + assert any(key[0] == 2 for key in admitted_encoder_keys) assert set(encoder_replay_keys) <= set(encoder_runner.graphs) replayed_mixed_keys = {key for key in decoder_replay_keys if key[5] and key[6]} assert replayed_mixed_keys diff --git a/tests/integration/test_lists/test-db/l0_dgx_h100.yml b/tests/integration/test_lists/test-db/l0_dgx_h100.yml index e36416f7e5fa..aebf9258fcc3 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_h100.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_h100.yml @@ -23,6 +23,7 @@ l0_dgx_h100: # ------------- Encoder-decoder TP tests --------------- - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-on-greedy-tp2-t5-small] - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-on-greedy-tp2-bart-large-cnn] + - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_continuous_admission_replays_encoder_and_mixed_cuda_graphs[tp2] - llmapi/test_llm_api_pytorch_whisper.py::test_whisper_pytorch_feature_combinations[fp32-kv-v1-graphs-off-greedy-tp2] # ------------- Disaggregated serving tests --------------- - accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_eagle3[eagle3_one_model=True-overlap_scheduler=True] diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index 9f342288a47f..f3f9cf03785c 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -323,7 +323,7 @@ l0_h100: - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-on-beam2-overlap-bart-large-cnn] - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v2-cuda-graph-on-greedy-overlap-bart-large-cnn] - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch[bf16-kv-v2-decoder-cuda-graph-on-greedy-batch2-bart-large-cnn] - - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_continuous_admission_replays_encoder_and_mixed_cuda_graphs + - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_continuous_admission_replays_encoder_and_mixed_cuda_graphs[tp1] - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-on-beam2-t5-base] - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-on-beam2-t5-large] - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-on-beam2-flan-t5-base] diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 9289de0c8abd..5e6830c52f3e 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -50,6 +50,9 @@ def insert(self, request_id): def erase(self, request_id): self.ids.discard(request_id) + def __contains__(self, request_id): + return request_id in self.ids + class MockPyExecutor: """A mock PyExecutor class for testing request handling logic. @@ -135,6 +138,7 @@ def mock_dist(): def _make_async_encoder_executor(future): executor = object.__new__(PyExecutor) + executor.dist = types.SimpleNamespace(tp_size=1) executor.encoder_launch_executor = Mock() executor.encoder_launch_executor.submit.return_value = future executor.pending_encoder_steps = [] @@ -163,6 +167,29 @@ def _make_encoder_batch_wait_executor(batch_sizes=None, encoder_max_batch_size=8 return executor +def _make_encoder_fallback_batch_wait_executor(): + executor = object.__new__(PyExecutor) + executor.llm_args = types.SimpleNamespace( + encoder_cuda_graph_config=None, + encoder_max_batch_size=None, + ) + executor.batch_wait_timeout_iters = 48 + executor.encoder_batch_wait_iters_count = 0 + executor.batch_wait_max_tokens_ratio = 0.5 + executor.max_num_tokens = 32 + executor.active_requests = [] + executor.inflight_req_ids = _InflightRequestIds() + return executor + + +def _make_encoder_request(request_id): + return types.SimpleNamespace( + request_id=request_id, + state=LlmRequestState.ENCODER_INIT, + encoder_output_len=4, + ) + + def test_encoder_graph_warmup_uses_runtime_encoder_stream(): executor = object.__new__(PyExecutor) executor.device_id = 3 @@ -184,104 +211,49 @@ def test_encoder_graph_warmup_uses_runtime_encoder_stream(): ) -def test_encoder_microbatch_graph_waits_for_target(): +def test_encoder_microbatch_graph_admission_boundaries(): executor = _make_encoder_batch_wait_executor() encoder_requests = [object()] * 7 - generation_requests = [object()] * 24 - - scheduled = executor._waiting_encoder_requests( - encoder_requests, - [], - generation_requests, - ) - - assert scheduled == [] - assert executor.encoder_batch_wait_iters_count == 1 - - -def test_encoder_microbatch_graph_releases_target(): - executor = _make_encoder_batch_wait_executor() - encoder_requests = [object()] * 8 - scheduled = executor._waiting_encoder_requests( encoder_requests, [], [object()] * 24, ) - - assert scheduled == encoder_requests - assert executor.encoder_batch_wait_iters_count == 0 - - -def test_encoder_microbatch_graph_does_not_release_partial_at_low_watermark(): - executor = _make_encoder_batch_wait_executor() - encoder_requests = [object()] - - scheduled = executor._waiting_encoder_requests( - encoder_requests, - [], - [object()] * 24, - ) - assert scheduled == [] assert executor.encoder_batch_wait_iters_count == 1 - -def test_encoder_microbatch_graph_caps_scheduler_overfill(): executor = _make_encoder_batch_wait_executor() encoder_requests = [object() for _ in range(12)] - scheduled = executor._waiting_encoder_requests( encoder_requests, [], [object()] * 20, ) - assert scheduled == encoder_requests[:8] assert executor.encoder_batch_wait_iters_count == 0 - -def test_encoder_microbatch_graph_releases_supported_tail_at_deadline(): - executor = _make_encoder_batch_wait_executor() - executor.encoder_batch_wait_iters_count = executor.batch_wait_timeout_iters - encoder_requests = [object() for _ in range(7)] - - scheduled = executor._waiting_encoder_requests( - encoder_requests, - [], - [], + executor = _make_encoder_batch_wait_executor( + batch_sizes=[1, 3, 6], + encoder_max_batch_size=8, ) - - assert scheduled == encoder_requests[:4] - assert executor.encoder_batch_wait_iters_count == 0 - - -def test_encoder_microbatch_graph_uses_configured_batch_sizes_at_deadline(): - executor = _make_encoder_batch_wait_executor(batch_sizes=[1, 3, 6], encoder_max_batch_size=8) executor.encoder_batch_wait_iters_count = executor.batch_wait_timeout_iters encoder_requests = [object() for _ in range(5)] - scheduled = executor._waiting_encoder_requests( encoder_requests, [], [], ) - assert scheduled == encoder_requests[:3] assert executor.encoder_batch_wait_iters_count == 0 - -def test_encoder_microbatch_graph_waits_above_low_watermark_after_deadline(): executor = _make_encoder_batch_wait_executor() executor.encoder_batch_wait_iters_count = executor.batch_wait_timeout_iters encoder_requests = [object() for _ in range(8)] - scheduled = executor._waiting_encoder_requests( encoder_requests, [], [object() for _ in range(25)], ) - assert scheduled == [] assert executor.encoder_batch_wait_iters_count == executor.batch_wait_timeout_iters + 1 @@ -295,101 +267,130 @@ def test_encoder_microbatch_graph_waits_above_low_watermark_after_deadline(): assert executor.encoder_batch_wait_iters_count == 0 -def test_pending_encoder_future_is_polled_without_blocking(): - future = Mock() - future.done.return_value = False - executor = _make_async_encoder_executor(future) - request = types.SimpleNamespace(request_id=11, state=LlmRequestState.ENCODER_INIT) +def test_encoder_fallback_distinguishes_inflight_encoder_and_decoder_work(): + executor = _make_encoder_fallback_batch_wait_executor() + encoder_requests = [_make_encoder_request(1)] + inflight_encoder_request = _make_encoder_request(2) + executor.active_requests.append(inflight_encoder_request) + executor.inflight_req_ids.insert(inflight_encoder_request.request_id) - executor._submit_encoder_step([request]) - executor._poll_encoder_steps() + scheduled = executor._waiting_encoder_requests(encoder_requests, [], []) + assert scheduled == encoder_requests + assert executor.encoder_batch_wait_iters_count == 0 - future.result.assert_not_called() - executor._publish_encoder_step.assert_not_called() - assert executor.inflight_req_ids.ids == {11} - assert len(executor.pending_encoder_steps) == 1 - executor.encoder_launch_executor.submit.assert_called_once_with( - executor._run_encoder_step_unchecked, - [request], + decoder_request = types.SimpleNamespace( + request_id=3, + state=LlmRequestState.GENERATION_IN_PROGRESS, ) + executor.active_requests = [decoder_request] + executor.inflight_req_ids.erase(inflight_encoder_request.request_id) + executor.inflight_req_ids.insert(decoder_request.request_id) + + scheduled = executor._waiting_encoder_requests(encoder_requests, [], []) + assert scheduled == [] + assert executor.encoder_batch_wait_iters_count == 1 -def test_completed_encoder_future_waits_for_cuda_event_without_blocking(): +def test_async_encoder_step_lifecycle(): ready_event = Mock() ready_event.query.side_effect = [False, True] result = EncoderStepResult( - hidden_states=torch.empty((1, 2)), - sequence_lengths=[1], + hidden_states=torch.arange(12).reshape(6, 2), + sequence_lengths=[2, 4], ready_event=ready_event, ) future = Mock() - future.done.return_value = True + future.done.side_effect = [False, True] future.result.return_value = result executor = _make_async_encoder_executor(future) - request = types.SimpleNamespace(request_id=12, state=LlmRequestState.ENCODER_INIT) + active_request = types.SimpleNamespace( + request_id=11, + state=LlmRequestState.ENCODER_INIT, + ) + completed_request = types.SimpleNamespace( + request_id=12, + state=LlmRequestState.GENERATION_COMPLETE, + ) + requests = [active_request, completed_request] + executor._publish_encoder_step.side_effect = ( + lambda encoder_requests, encoder_result: PyExecutor._publish_encoder_step( + executor, + encoder_requests, + encoder_result, + ) + ) - executor._submit_encoder_step([request]) + executor._submit_encoder_step(requests) executor._poll_encoder_steps() + future.result.assert_not_called() + executor._publish_encoder_step.assert_not_called() + assert executor.inflight_req_ids.ids == {11, 12} + assert len(executor.pending_encoder_steps) == 1 + executor.encoder_launch_executor.submit.assert_called_once_with( + executor._run_encoder_step_unchecked, + requests, + ) + + executor._poll_encoder_steps() future.result.assert_called_once_with() ready_event.query.assert_called_once_with() executor._publish_encoder_step.assert_not_called() - assert executor.inflight_req_ids.ids == {12} + assert executor.inflight_req_ids.ids == {11, 12} assert len(executor.pending_encoder_steps) == 1 executor._poll_encoder_steps() - future.result.assert_called_once_with() assert ready_event.query.call_count == 2 - executor._publish_encoder_step.assert_called_once_with([request], result) + executor._publish_encoder_step.assert_called_once_with(requests, result) assert executor.inflight_req_ids.ids == set() assert executor.pending_encoder_steps == [] - - -def test_publish_encoder_output_does_not_resurrect_completed_request(): - executor = object.__new__(PyExecutor) - active_request = types.SimpleNamespace(state=LlmRequestState.ENCODER_INIT) - completed_request = types.SimpleNamespace(state=LlmRequestState.GENERATION_COMPLETE) - hidden_states = torch.arange(12).reshape(6, 2) - ready_event = Mock() - - executor._scatter_encoder_output( - [active_request, completed_request], - hidden_states, - [2, 4], - ready_event, - ) - assert active_request.state == LlmRequestState.CONTEXT_INIT assert active_request.py_encoder_output_ready_event is ready_event - assert torch.equal(active_request.py_encoder_output, hidden_states[:2]) + assert torch.equal(active_request.py_encoder_output, result.hidden_states[:2]) assert completed_request.state == LlmRequestState.GENERATION_COMPLETE assert not hasattr(completed_request, "py_encoder_output") - -def test_attach_encoder_output_records_stream_after_encoder_is_ready(): - executor = object.__new__(PyExecutor) executor.execution_stream = Mock() + encoder_output = Mock() + active_request.py_encoder_output = encoder_output + scheduled_requests = types.SimpleNamespace(context_requests=[active_request]) + executor._attach_encoder_output_to_execution_stream(scheduled_requests) + + executor.execution_stream.wait_event.assert_not_called() + encoder_output.record_stream.assert_called_once_with(executor.execution_stream) + assert active_request.py_encoder_output_ready_event is None + + +def test_tp_encoder_step_synchronizes_and_publishes_inline(): + call_order = [] + execution_stream = Mock() + encoder_stream = Mock() + encoder_stream.wait_stream.side_effect = lambda stream: call_order.append("wait_stream") ready_event = Mock() - first_output = Mock() - second_output = Mock() - first_request = types.SimpleNamespace( - py_encoder_output=first_output, - py_encoder_output_ready_event=ready_event, + ready_event.synchronize.side_effect = lambda: call_order.append("synchronize") + result = EncoderStepResult( + hidden_states=torch.empty((1, 2)), + sequence_lengths=[1], + ready_event=ready_event, ) - second_request = types.SimpleNamespace( - py_encoder_output=second_output, - py_encoder_output_ready_event=ready_event, + future = Mock() + future.result.side_effect = lambda: (call_order.append("result"), result)[1] + executor = _make_async_encoder_executor(future) + executor.dist.tp_size = 2 + executor.execution_stream = execution_stream + executor.encoder_stream = encoder_stream + executor._publish_encoder_step.side_effect = lambda requests, encoder_result: call_order.append( + "publish" ) - scheduled_requests = types.SimpleNamespace(context_requests=[first_request, second_request]) + request = types.SimpleNamespace(request_id=13, state=LlmRequestState.ENCODER_INIT) - executor._attach_encoder_output_to_execution_stream(scheduled_requests) + executor._submit_encoder_step([request]) - executor.execution_stream.wait_event.assert_not_called() - first_output.record_stream.assert_called_once_with(executor.execution_stream) - second_output.record_stream.assert_called_once_with(executor.execution_stream) - assert first_request.py_encoder_output_ready_event is None - assert second_request.py_encoder_output_ready_event is None + assert call_order == ["wait_stream", "result", "synchronize", "publish"] + encoder_stream.wait_stream.assert_called_once_with(execution_stream) + assert executor.inflight_req_ids.ids == set() + assert executor.pending_encoder_steps == [] @pytest.fixture diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index 6347275fd640..9bf0f87465d9 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -12,7 +12,8 @@ from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_connector import \ KvCacheConnectorWorker from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import ( - _restore_spec_decode_capture_state, _save_spec_decode_capture_state) + EncoderCUDAGraphRunner, _restore_spec_decode_capture_state, + _save_spec_decode_capture_state) from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest from tensorrt_llm._torch.pyexecutor.model_engine import ( PyTorchModelEngine, _build_request_multimodal_input) @@ -150,6 +151,70 @@ def create_model_engine_and_kvcache(llm_args: TorchLlmArgs = None, class PyTorchModelEngineTestCase(unittest.TestCase): + def test_encoder_cuda_graph_stages_and_restores_fixed_sequence_slots( + self) -> None: + runner = EncoderCUDAGraphRunner.__new__(EncoderCUDAGraphRunner) + runner.is_encoder_decoder = True + runner.use_fixed_sequence_slots = True + runner.supported_batch_sizes = [2] + runner.supported_seq_lens = [512] + runner.max_supported_num_tokens = 1024 + small_key = (2, 512, 512) + compatible_key = (2, 1024, 512) + runner._capture_sequence_lengths = { + small_key: [511, 1], + compatible_key: [512, 512], + } + runner._capture_keys_by_batch_size = { + 2: [small_key, compatible_key], + } + runner._arange_max = torch.arange(1024, dtype=torch.int32) + + self.assertEqual( + runner._get_dynamic_capture_key([200, 300], + allow_batch_padding=False), + compatible_key) + + source_sequence_lengths = [1, 400] + key = runner._get_dynamic_capture_key(source_sequence_lengths, + allow_batch_padding=False) + self.assertEqual(key, small_key) + self.assertEqual(runner._get_capture_sequence_offsets(key), + [0, 511, 512]) + + input_ids = torch.arange(401, dtype=torch.int32) + inputs = runner.prepare_encoder_decoder_inputs( + { + "input_ids": input_ids, + "position_ids": input_ids, + "seq_lens": source_sequence_lengths, + }, + key, + source_sequence_lengths, + ) + self.assertEqual(inputs["seq_lens"], [400, 1]) + self.assertEqual(inputs["_encoder_source_to_slot"], [1, 0]) + + static_tensors = { + "input_ids": torch.empty(512, dtype=torch.int32), + "position_ids": torch.empty((1, 512), dtype=torch.int32), + } + runner._stage_encoder_decoder_inputs(key, inputs, static_tensors) + expected_staged_ids = torch.zeros(512, dtype=torch.int32) + expected_staged_ids[:400] = input_ids[1:] + expected_staged_ids[511] = input_ids[0] + torch.testing.assert_close(static_tensors["input_ids"], + expected_staged_ids) + torch.testing.assert_close(static_tensors["position_ids"][0], + expected_staged_ids) + + fixed_slot_output = torch.arange(512).unsqueeze(1) + restored_output = runner.restore_encoder_decoder_output( + key, fixed_slot_output, inputs) + expected_output = torch.cat( + (fixed_slot_output[511:512], fixed_slot_output[:400])) + torch.testing.assert_close(restored_output, expected_output) + def test_build_request_multimodal_input_skips_when_cache_disabled( self) -> None: request = LlmRequest( diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py index 930eb75bb596..2e3bee0b1fe9 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + """Unit tests for warmup-cleanup behavior in PyTorchModelEngine.warmup(). Locks in that gc.collect() + torch.cuda.empty_cache() fire immediately after @@ -175,29 +178,13 @@ def _record(*msg): class TestWarmupCleanup(unittest.TestCase): """Lock in warmup-cleanup behavior introduced by PR #14609 (Plan B).""" - def test_main_cuda_graph_warmup_defers_encoder_decoder_encoder(self): + def test_encoder_decoder_encoder_warmup_is_deferred_and_uses_two_passes(self): model_engine = object.__new__(PyTorchModelEngine) model_engine.cuda_graph_runner = SimpleNamespace( enabled=True, is_warmup_only=True, ) - model_engine.encoder_cuda_graph_runner = SimpleNamespace(enabled=True) model_engine._torch_compile_piecewise_cuda_graph = False - resource_manager = object() - - with ( - patch.object(model_engine, "_capture_generation_cuda_graphs") as generation, - patch.object(model_engine, "_capture_mixed_encoder_decoder_cuda_graphs") as mixed, - patch.object(model_engine, "_capture_encoder_decoder_encoder_cuda_graphs") as encoder, - ): - model_engine._run_cuda_graph_warmup(resource_manager) - - generation.assert_called_once_with(resource_manager) - mixed.assert_called_once_with(resource_manager) - encoder.assert_not_called() - - def test_encoder_decoder_encoder_warmup_uses_two_passes(self): - model_engine = object.__new__(PyTorchModelEngine) model_engine.is_warmup = False @contextlib.contextmanager @@ -214,13 +201,22 @@ def allow_capture(): resource_manager = object() warmup_states = [] - with patch.object( - model_engine, - "_capture_encoder_decoder_encoder_cuda_graphs", - side_effect=lambda _: warmup_states.append(runner.is_warmup_only), + with ( + patch.object(model_engine, "_capture_generation_cuda_graphs") as generation, + patch.object(model_engine, "_capture_mixed_encoder_decoder_cuda_graphs") as mixed, + patch.object( + model_engine, + "_capture_encoder_decoder_encoder_cuda_graphs", + side_effect=lambda _: warmup_states.append(runner.is_warmup_only), + ) as encoder, ): + model_engine._run_cuda_graph_warmup(resource_manager) + generation.assert_called_once_with(resource_manager) + mixed.assert_called_once_with(resource_manager) + encoder.assert_not_called() model_engine._warmup_encoder_decoder_encoder_cuda_graphs(resource_manager) + assert encoder.call_count == 2 assert warmup_states == [True, False] assert not runner.is_warmup_only diff --git a/tests/unittest/_torch/sampler/test_torch_sampler.py b/tests/unittest/_torch/sampler/test_torch_sampler.py index fe60324c1d31..259a165c15b4 100644 --- a/tests/unittest/_torch/sampler/test_torch_sampler.py +++ b/tests/unittest/_torch/sampler/test_torch_sampler.py @@ -657,13 +657,50 @@ def _uut(res=res): run_test_with_warmup(_test_runner, max_sync_s=0.3) +@force_ampere +def test_greedy_no_repeat_ngram_uses_token_ban_path(): + sampler = TorchSampler( + TorchSampler.Args( + max_seq_len=16, + max_draft_len=0, + max_num_sequences=1, + max_beam_width=1, + max_total_draft_tokens=0, + disable_overlap_scheduler=True, + ) + ) + request = LlmRequest( + request_id=0, + max_new_tokens=4, + input_tokens=[1, 2, 1], + sampling_config=SamplingConfig(), + seq_slot=0, + is_streaming=False, + ) + request.py_no_repeat_ngram_size = 2 + scheduled_requests = ScheduledRequests() + scheduled_requests.generation_requests = [request] + logits = torch.tensor([[0.0, 0.0, 10.0, 9.0]], device="cuda") + + *_, new_tokens_host, single_step_greedy = sampler._process_requests( + scheduled_requests, + {"logits": logits}, + sampler.store.new_tokens, + [0], + ) + torch.cuda.synchronize() + + assert not single_step_greedy + assert new_tokens_host.reshape(-1)[0].item() == 3 + + class TestFinishReasons: NOT_FINISHED = FinishReason.NOT_FINISHED STOP_WORDS = FinishReason.STOP_WORDS END_ID = FinishReason.END_ID LENGTH = FinishReason.LENGTH - def test_single_step_greedy_checks_finish_reasons_on_host(self): + def test_single_step_greedy_updates_finish_reasons_and_filters_completed_requests(self): sampler = object.__new__(TorchSampler) sampler.max_seq_len = 20 sampler._track_pending_steps = False @@ -686,50 +723,23 @@ def test_single_step_greedy_checks_finish_reasons_on_host(self): sampling_config=SamplingConfig(), is_streaming=False, ), - ] - new_tokens = torch.tensor([2, 7], dtype=torch.int32) - state = SampleStateTorch( - requests=requests, - device=None, - host=SampleStateTensorsHostTorch( - new_tokens=new_tokens, - finish_reasons=None, - first_finish_reasons=None, - ), - single_step_greedy=True, - ) - - sampler.update_requests(state) - - assert all(request.is_finished for request in requests) - # The first request reaches EOS and length together; EOS takes precedence. - assert not requests[0].is_finished_due_to_length - assert requests[1].is_finished_due_to_length - assert requests[0].get_tokens(0)[-1] == 2 - assert requests[1].get_tokens(0)[-1] == 7 - - def test_single_step_greedy_filters_requests_completed_after_sampling(self): - sampler = object.__new__(TorchSampler) - sampler.max_seq_len = 20 - sampler._track_pending_steps = False - requests = [ LlmRequest( - request_id=request_id, - seq_slot=request_id, + request_id=2, + seq_slot=2, input_tokens=[2, 0], max_new_tokens=10, end_id=2, sampling_config=SamplingConfig(), is_streaming=False, - ) - for request_id in range(2) + ), ] - requests[0].finish_by(FinishReason.LENGTH, 0) + requests[2].finish_by(FinishReason.LENGTH, 0) + new_tokens = torch.tensor([2, 7, 99], dtype=torch.int32) state = SampleStateTorch( requests=requests, device=None, host=SampleStateTensorsHostTorch( - new_tokens=torch.tensor([99, 7], dtype=torch.int32), + new_tokens=new_tokens, finish_reasons=None, first_finish_reasons=None, ), @@ -738,8 +748,13 @@ def test_single_step_greedy_filters_requests_completed_after_sampling(self): sampler.update_requests(state) - assert requests[0].get_tokens(0) == [2, 0] + assert all(request.is_finished for request in requests) + # The first request reaches EOS and length together; EOS takes precedence. + assert not requests[0].is_finished_due_to_length + assert requests[1].is_finished_due_to_length + assert requests[0].get_tokens(0)[-1] == 2 assert requests[1].get_tokens(0)[-1] == 7 + assert requests[2].get_tokens(0) == [2, 0] class RequestCase: MAX_NEW_TOKENS = 10 diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 514388826b8f..eb6f06597fee 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -1745,7 +1745,13 @@ def test_cuda_graph_config_accepts_encoder_config(self): assert args.cuda_graph_config.seq_lens == [8, 32] assert args.cuda_graph_config.max_seq_len == 32 - def test_encoder_decoder_cuda_graph_configs(self): + def test_encoder_decoder_cuda_graph_user_interface(self): + encoder_config = EncodeCudaGraphConfig( + batch_sizes=[1, 4], + num_tokens=[16, 64], + seq_lens=[8, 32], + enable_padding=True, + ) args = TorchLlmArgs( model=llama_model_path, encoder_max_batch_size=4, @@ -1753,63 +1759,58 @@ def test_encoder_decoder_cuda_graph_configs(self): batch_sizes=[1, 4], enable_padding=True, ), - encoder_cuda_graph_config=EncodeCudaGraphConfig( - batch_sizes=[1, 4], - num_tokens=[16, 64], - seq_lens=[8, 32], - enable_padding=True, - ), + encoder_cuda_graph_config=encoder_config, ) assert isinstance(args.cuda_graph_config, DecodeCudaGraphConfig) assert isinstance(args.encoder_cuda_graph_config, EncodeCudaGraphConfig) assert args.encoder_cuda_graph_config.batch_sizes == [1, 4] - assert args.encoder_cuda_graph_config.num_tokens == [16, 64] - assert args.encoder_cuda_graph_config.seq_lens == [8, 32] assert args.enable_encoder_decoder_mixed_cuda_graph - def test_encoder_decoder_mixed_cuda_graph_can_be_disabled(self): - args = TorchLlmArgs( + disabled_args = TorchLlmArgs( model=llama_model_path, encoder_max_batch_size=4, - encoder_cuda_graph_config=EncodeCudaGraphConfig( - batch_sizes=[1, 4], - num_tokens=[16, 64], - seq_lens=[8, 32], - enable_padding=True, - ), + encoder_cuda_graph_config=encoder_config, enable_encoder_decoder_mixed_cuda_graph=False, ) - assert not args.enable_encoder_decoder_mixed_cuda_graph + assert not disabled_args.enable_encoder_decoder_mixed_cuda_graph - def test_encoder_cuda_graph_config_requires_encoder_max_batch_size(self): - with pytest.raises(ValidationError, - match=("encoder_cuda_graph_config requires " - "encoder_max_batch_size")): - TorchLlmArgs( - model=llama_model_path, - encoder_cuda_graph_config=EncodeCudaGraphConfig( - batch_sizes=[1, 4], - num_tokens=[16, 64], - seq_lens=[8, 32], - enable_padding=True, - ), - ) + def test_encoder_cuda_graph_config_validation(self): + invalid_cases = [ + ( + { + "encoder_cuda_graph_config": + EncodeCudaGraphConfig( + batch_sizes=[1, 4], + num_tokens=[16, 64], + seq_lens=[8, 32], + enable_padding=True, + ), + }, + "encoder_cuda_graph_config requires encoder_max_batch_size", + ), + ( + { + "encoder_max_batch_size": + 4, + "encoder_cuda_graph_config": + EncodeCudaGraphConfig( + batch_sizes=[1, 4], + enable_padding=True, + ), + }, + ("encoder_cuda_graph_config requires " + "num_tokens/max_num_token and seq_lens/max_seq_len"), + ), + ] - def test_encoder_cuda_graph_config_requires_shape_dimensions(self): - with pytest.raises( - ValidationError, - match=("encoder_cuda_graph_config requires " - "num_tokens/max_num_token and seq_lens/max_seq_len")): - TorchLlmArgs( - model=llama_model_path, - encoder_max_batch_size=4, - encoder_cuda_graph_config=EncodeCudaGraphConfig( - batch_sizes=[1, 4], - enable_padding=True, - ), - ) + for kwargs, error_match in invalid_cases: + with pytest.raises(ValidationError, match=error_match): + TorchLlmArgs( + model=llama_model_path, + **kwargs, + ) def test_cuda_graph_config_infers_encode_mode_from_raw_dict(self): args = TorchLlmArgs( From bb70aab8ab4d134b85ff5acabbc282815383d922 Mon Sep 17 00:00:00 2001 From: Guiju Zhang <7135567+cascade812@users.noreply.github.com> Date: Fri, 31 Jul 2026 17:50:00 -0700 Subject: [PATCH 09/15] address comments Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/model_engine.py | 18 +++++++++--------- tensorrt_llm/_torch/pyexecutor/py_executor.py | 13 +++++++++---- .../_torch/pyexecutor/sampler/sampler.py | 8 +++++--- .../_torch/executor/test_py_executor.py | 4 ++-- .../test_pytorch_model_engine_warmup.py | 4 ++-- 5 files changed, 27 insertions(+), 20 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index dc8fe4167fbe..4bc0cd5d73ec 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1865,8 +1865,9 @@ def _run_cuda_graph_warmup(self, resource_manager: ResourceManager): if not self.cuda_graph_runner.is_warmup_only: self._capture_piecewise_cuda_graphs(resource_manager) + @torch.inference_mode() @with_warmup_flag - def _warmup_encoder_decoder_encoder_cuda_graphs( + def _warmup_encoder_cuda_graphs_enc_dec( self, resource_manager: ResourceManager) -> None: """Capture encoder-decoder encoder graphs on their runtime host thread.""" runner = self.encoder_cuda_graph_runner @@ -1874,7 +1875,7 @@ def _warmup_encoder_decoder_encoder_cuda_graphs( return capture = functools.partial( - self._capture_encoder_decoder_encoder_cuda_graphs, + self._capture_encoder_cuda_graphs_enc_dec, resource_manager, ) self._warmup_and_capture_encoder_cuda_graphs(capture) @@ -1894,7 +1895,7 @@ def _warmup_and_capture_encoder_cuda_graphs( runner.is_warmup_only = False capture() - def _capture_encoder_decoder_encoder_cuda_graphs( + def _capture_encoder_cuda_graphs_enc_dec( self, resource_manager: ResourceManager) -> None: """Warm up or capture encoder graphs used by encoder-decoder models.""" runner = self.encoder_cuda_graph_runner @@ -1926,7 +1927,7 @@ def _capture_encoder_decoder_encoder_cuda_graphs( logger.info("Encoder-decoder encoder CUDA graph " f"{operation}: key={key}") - self._forward_encoder_with_cuda_graph(inputs) + self._encoder_forward_enc_dec(inputs) torch.cuda.synchronize() num_processed += 1 @@ -6561,7 +6562,7 @@ def _capture_encoder_cuda_graphs(self) -> None: num_processed = 0 logger.info(f"Running encoder CUDA graph {operation} ...") for bs in batch_sizes: - if bs > self.batch_size: + if bs > self.encoder_batch_size: continue for sl_idx, sl in reversed(list(enumerate(seq_lens_list))): prev_sl = seq_lens_list[sl_idx - 1] if sl_idx > 0 else 0 @@ -7331,11 +7332,11 @@ def _forward_step_encoder_cuda_graph( inputs.get('resource_manager'), }) - def _forward_encoder_with_cuda_graph( + def _encoder_forward_enc_dec( self, inputs: Dict[str, Any], ) -> torch.Tensor: - """Replay a bucketed dynamic-layout graph when the encoder is eligible.""" + """Run the encoder-decoder encoder, using a CUDA graph when eligible.""" input_ids = inputs.get('encoder_input_ids_host') position_ids = inputs.get('encoder_position_ids_host') seq_lens = inputs['encoder_seq_lens'] @@ -7426,8 +7427,7 @@ def forward_encoder( with torch.inference_mode(): inputs = self._prepare_tp_inputs_encoder( encoder_requests, resource_manager=resource_manager) - encoder_hidden_states = self._forward_encoder_with_cuda_graph( - inputs) + encoder_hidden_states = self._encoder_forward_enc_dec(inputs) return encoder_hidden_states, inputs['encoder_seq_lens'] diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 548e39adc8db..05428ba7ac71 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -836,7 +836,7 @@ def __init__( # on the same single worker that owns every runtime encoder replay. self.encoder_stream.wait_stream(self.execution_stream) self.encoder_launch_executor.submit( - self._warmup_encoder_decoder_encoder_cuda_graphs).result() + self._warmup_encoder_cuda_graphs_enc_dec).result() self.is_warmup = False @@ -5471,11 +5471,11 @@ def _schedule(self): # micro-batch; this preserves the cross-KV lifecycle and the # dual-pool budget. # --------------------------------------------------------------- - def _warmup_encoder_decoder_encoder_cuda_graphs(self) -> None: + def _warmup_encoder_cuda_graphs_enc_dec(self) -> None: """Capture encoder graphs on the worker used for runtime replay.""" warmup = getattr( self.model_engine, - "_warmup_encoder_decoder_encoder_cuda_graphs", + "_warmup_encoder_cuda_graphs_enc_dec", None, ) if not callable(warmup): @@ -5579,7 +5579,12 @@ def _finish_failed_encoder_step(self, encoder_requests: List[LlmRequest], def _run_encoder_step(self, encoder_requests: List[LlmRequest]) -> None: try: - result = self._run_encoder_step_unchecked(encoder_requests) + executor = self.encoder_launch_executor + if executor is None: + raise RuntimeError("Encoder launch executor is unavailable.") + future = executor.submit(self._run_encoder_step_unchecked, + encoder_requests) + result = future.result() self._publish_encoder_step(encoder_requests, result) except Exception as e: self._finish_failed_encoder_step(encoder_requests, e) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index ec7e239b21e3..ae55460fd511 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -3859,12 +3859,14 @@ def _process_requests( and scheduled_requests.num_context_requests == 0 and len(generation_requests) <= raw_logits_cuda.shape[0] and model_outputs.get("d2t") is None + and all( + not request.is_dummy and get_draft_token_length(request) == 0 + for request in generation_requests + ) and ( has_stable_request_ids or all( - not request.is_dummy - and get_draft_token_length(request) == 0 - and request._py_embedding_bias_1d is None + request._py_embedding_bias_1d is None and not getattr(request, "py_bad_words", None) and not getattr(request, "py_no_repeat_ngram_size", None) and not request.py_min_length diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 5e6830c52f3e..6bb905e4cfad 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -202,11 +202,11 @@ def test_encoder_graph_warmup_uses_runtime_encoder_stream(): patch("torch.cuda.set_device") as set_device, patch("torch.cuda.stream", return_value=stream_context) as cuda_stream, ): - executor._warmup_encoder_decoder_encoder_cuda_graphs() + executor._warmup_encoder_cuda_graphs_enc_dec() set_device.assert_called_once_with(3) cuda_stream.assert_called_once_with(executor.encoder_stream) - executor.model_engine._warmup_encoder_decoder_encoder_cuda_graphs.assert_called_once_with( + executor.model_engine._warmup_encoder_cuda_graphs_enc_dec.assert_called_once_with( executor.resource_manager ) diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py index 2e3bee0b1fe9..8d95fb9f1a53 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py @@ -206,7 +206,7 @@ def allow_capture(): patch.object(model_engine, "_capture_mixed_encoder_decoder_cuda_graphs") as mixed, patch.object( model_engine, - "_capture_encoder_decoder_encoder_cuda_graphs", + "_capture_encoder_cuda_graphs_enc_dec", side_effect=lambda _: warmup_states.append(runner.is_warmup_only), ) as encoder, ): @@ -214,7 +214,7 @@ def allow_capture(): generation.assert_called_once_with(resource_manager) mixed.assert_called_once_with(resource_manager) encoder.assert_not_called() - model_engine._warmup_encoder_decoder_encoder_cuda_graphs(resource_manager) + model_engine._warmup_encoder_cuda_graphs_enc_dec(resource_manager) assert encoder.call_count == 2 assert warmup_states == [True, False] From 15b2d6796ff4141f5e18fda7cdf5d0b4eeadb978 Mon Sep 17 00:00:00 2001 From: Guiju Zhang <7135567+cascade812@users.noreply.github.com> Date: Mon, 3 Aug 2026 14:39:50 -0700 Subject: [PATCH 10/15] address comment Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com> --- tensorrt_llm/llmapi/llm_args.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 79442e113291..a7b6fa6abcfa 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -4988,7 +4988,7 @@ class TorchLlmArgs(BaseLlmArgs): "encoder-decoder model. Use `cuda_graph_config` for the decoder " "and this field for the encoder. Encoder CUDA graphs require " "`encoder_max_batch_size` to be set."), - status="beta") + status="prototype") enable_encoder_decoder_mixed_cuda_graph: bool = Field( default=True, @@ -4998,7 +4998,7 @@ class TorchLlmArgs(BaseLlmArgs): "containing both context and generation requests. It is enabled " "by default when both `cuda_graph_config` and " "`encoder_cuda_graph_config` produce usable graph shapes."), - status="beta") + status="prototype") @field_validator('cuda_graph_config', mode='before') @classmethod From b7b15eba36736ac1bcf2b116ffb33c50820bf4e6 Mon Sep 17 00:00:00 2001 From: Guiju Zhang <7135567+cascade812@users.noreply.github.com> Date: Tue, 4 Aug 2026 11:36:23 -0700 Subject: [PATCH 11/15] fix failed tests Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py | 9 ++++++--- tensorrt_llm/_torch/pyexecutor/py_executor.py | 2 ++ tests/unittest/_torch/sampler/test_torch_sampler.py | 2 +- 3 files changed, 9 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index 62f0c334a8fb..8239f5ede6b6 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -554,7 +554,9 @@ def get_graph_pool(self): def _get_num_tokens_for_key(self, key: KeyType) -> int: batch_size = key[0] token_per_generation = key[1] + 1 - context_query_lens = key[5] + # Direct capture callers predating mixed encoder-decoder graphs pass + # only the generation-key prefix. A missing suffix means gen-only. + context_query_lens = key[5] if len(key) > 5 else () num_contexts = len(context_query_lens) return (sum(context_query_lens) + (batch_size * self.max_beam_width - num_contexts) * @@ -591,7 +593,7 @@ def capture(self, capture_inputs = initial_inputs.copy() capture_inputs.update(sliced_static_tensors) - encoder_input_lens = key[6] + encoder_input_lens = key[6] if len(key) > 6 else () num_encoder_tokens = sum(encoder_input_lens) if num_encoder_tokens: encoder_hidden_states = initial_inputs.get("encoder_hidden_states") @@ -695,7 +697,8 @@ def replay(self, key: KeyType, else: static_tensors["position_ids"][:, :seqlen].copy_(position_ids) - num_encoder_tokens = sum(key[6]) + encoder_input_lens = key[6] if len(key) > 6 else () + num_encoder_tokens = sum(encoder_input_lens) if num_encoder_tokens: encoder_hidden_states = current_inputs.get("encoder_hidden_states") if encoder_hidden_states is None: diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index e47d45c9ade5..58bef6db5a9a 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -5671,6 +5671,8 @@ def _run_encoder_step(self, encoder_requests: List[LlmRequest]) -> None: future = executor.submit(self._run_encoder_step_unchecked, encoder_requests) result = future.result() + if self.dist.tp_size > 1: + result.ready_event.synchronize() self._publish_encoder_step(encoder_requests, result) except Exception as e: self._finish_failed_encoder_step(encoder_requests, e) diff --git a/tests/unittest/_torch/sampler/test_torch_sampler.py b/tests/unittest/_torch/sampler/test_torch_sampler.py index e9e720335500..65355d2ce723 100644 --- a/tests/unittest/_torch/sampler/test_torch_sampler.py +++ b/tests/unittest/_torch/sampler/test_torch_sampler.py @@ -769,7 +769,7 @@ def test_greedy_no_repeat_ngram_uses_token_ban_path(): seq_slot=0, is_streaming=False, ) - request.py_no_repeat_ngram_size = 2 + setattr(request, "py_no_repeat_ngram_size", 2) scheduled_requests = ScheduledRequests() scheduled_requests.generation_requests = [request] logits = torch.tensor([[0.0, 0.0, 10.0, 9.0]], device="cuda") From 477ac446c34cd2b6af2a61b2ea57f62abfa129f1 Mon Sep 17 00:00:00 2001 From: Guiju Zhang <7135567+cascade812@users.noreply.github.com> Date: Tue, 4 Aug 2026 16:02:25 -0700 Subject: [PATCH 12/15] address failed tests Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com> --- tests/unittest/_torch/sampler/test_torch_sampler.py | 9 ++++++++- tests/unittest/api_stability/references/llm.yaml | 4 ++-- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/tests/unittest/_torch/sampler/test_torch_sampler.py b/tests/unittest/_torch/sampler/test_torch_sampler.py index 65355d2ce723..1edca6b81b96 100644 --- a/tests/unittest/_torch/sampler/test_torch_sampler.py +++ b/tests/unittest/_torch/sampler/test_torch_sampler.py @@ -1739,7 +1739,14 @@ def _mock_filter(self, requests: ScheduledRequests) -> list[LlmRequest]: assert sample_state.sampler_event is not None sample_state.sampler_event.synchronize() assert sample_state.host is not None - new_tokens_tensors.append(sample_state.host.new_tokens.unsqueeze(-1)) + host_new_tokens = sample_state.host.new_tokens + if sample_state.single_step_greedy: + # The stable greedy path copies one token per active request instead of + # the full [step, slot, beam] buffer. This fixture uses dense sequence + # slots, so restore that layout before comparing sampling results. + assert host_new_tokens.shape == (len(sample_state.requests),) + host_new_tokens = host_new_tokens.reshape(1, -1, 1) + new_tokens_tensors.append(host_new_tokens.unsqueeze(-1)) new_tokens = torch.cat(new_tokens_tensors, dim=-1) if num_repeats is None: new_tokens = new_tokens.squeeze(-1) diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index 9351334eae99..962a6813a035 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -98,11 +98,11 @@ methods: encoder_cuda_graph_config: annotation: Optional[tensorrt_llm.llmapi.llm_args.EncodeCudaGraphConfig] default: null - status: beta + status: prototype enable_encoder_decoder_mixed_cuda_graph: annotation: bool default: True - status: beta + status: prototype multimodal_config: annotation: tensorrt_llm.llmapi.llm_args.MultimodalConfig default: null From 47b0116e42cc89c135c03524dae613f0f2aa2daa Mon Sep 17 00:00:00 2001 From: Guiju Zhang <7135567+cascade812@users.noreply.github.com> Date: Wed, 5 Aug 2026 14:08:02 -0700 Subject: [PATCH 13/15] address comments Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com> --- .../_torch/attention_backend/trtllm.py | 10 +- .../_torch/pyexecutor/cuda_graph_runner.py | 152 ++++++++++-------- .../_torch/pyexecutor/model_engine.py | 25 +-- .../_torch/pyexecutor/sampler/penalties.py | 8 +- .../_torch/pyexecutor/sampler/sampler.py | 3 +- .../defs/llmapi/test_llm_api_pytorch_bart.py | 8 +- .../defs/llmapi/test_llm_api_pytorch_t5.py | 8 +- .../executor/test_pytorch_model_engine.py | 99 ++++++++++-- .../_torch/sampler/test_torch_sampler.py | 51 ++++++ 9 files changed, 263 insertions(+), 101 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index 56af4c1e56a5..76bdf984092b 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -642,11 +642,11 @@ def prepare(self) -> None: host_request_types=self.host_request_types[:self.num_seqs], ) - def prepare_encoder_decoder(self, prompt_lens: torch.Tensor, - kv_lens: torch.Tensor, context_kv_tokens: int, - generation_kv_tokens: int, - max_kv_len: int) -> None: - """Prepare simple encoder-decoder attention from native host buffers.""" + def prepare_encoder_decoder_from_precomputed_lengths( + self, prompt_lens: torch.Tensor, kv_lens: torch.Tensor, + context_kv_tokens: int, generation_kv_tokens: int, + max_kv_len: int) -> None: + """Prepare encoder-decoder attention from precomputed lengths.""" super().prepare() extra_attrs = get_model_extra_attrs() if extra_attrs is None: diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index 8239f5ede6b6..e10e55b69132 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -1,8 +1,8 @@ import bisect import contextlib from dataclasses import dataclass -from typing import (Any, Callable, Dict, Iterator, List, Optional, Tuple, - TypeAlias) +from typing import (Any, Callable, Dict, Iterator, List, NamedTuple, Optional, + Tuple, TypeAlias) import torch @@ -36,8 +36,18 @@ # as a one-token context chunk to write its cross-KV cache, so enc-dec # dummies need one prompt token plus one generated token. ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM = 2 -KeyType: TypeAlias = Tuple[int, int, bool, bool, bool, Tuple[int, ...], - Tuple[int, ...]] + + +class KeyType(NamedTuple): + batch_size: int + draft_len: int + is_first_draft: bool + short_seq_len_mode: bool = False + is_all_greedy_sample: bool = True + # Primarily used for mixed batches of encoder-decoder models. + num_contexts: int = 0 + context_query_len: int = 0 + num_encoder_tokens: int = 0 def _save_spec_decode_capture_state( @@ -224,13 +234,13 @@ def _get_static_encoder_hidden_states( return static_encoder_hidden_states[:num_encoder_tokens] def _is_mixed_encoder_decoder_batch(self, batch: ScheduledRequests) -> bool: - return (self.enable_encoder_decoder_mixed_cuda_graph - and batch.num_context_requests > 0 + return (self.is_encoder_decoder and batch.num_context_requests > 0 and batch.num_generation_requests > 0) def _can_run_cuda_graph_batch(self, batch: ScheduledRequests) -> bool: - return batch.can_run_cuda_graph or self._is_mixed_encoder_decoder_batch( - batch) + return (batch.can_run_cuda_graph + or (self.enable_encoder_decoder_mixed_cuda_graph + and self._is_mixed_encoder_decoder_batch(batch))) def _get_seq_len_mode( self, @@ -304,12 +314,11 @@ def get_graph_key( spec_resource_manager: Optional[BaseResourceManager] = None, spec_metadata: Optional[SpecMetadata] = None, promoted_context_request_ids: frozenset[int] = frozenset() - ) -> KeyType: + ) -> Optional[KeyType]: batch_size = batch.batch_size - # Get the sequence length mode. - # Keep the graph-key tuple unchanged; promoted IDs only correct the - # sequence length observed by sparse short/long graph selection. + # Promoted IDs correct the sequence length observed by sparse + # short/long graph selection. short_seq_len_mode = self._get_seq_len_mode( batch, new_tensors_device, promoted_context_request_ids) @@ -326,8 +335,11 @@ def get_graph_key( # If 'is_first_draft' is True, even with tree decoding, the length of draft_len will only be 'max_draft_len', not 'max_total_draft_token'. # Because we will pad the input to 'max_draft_len' length for the first draft layer. draft_len = self.config.original_max_draft_len if spec_resource_manager.is_first_draft else 0 - key = (batch_size, draft_len, spec_resource_manager.is_first_draft, - short_seq_len_mode, is_all_greedy_sample, (), ()) + key = KeyType(batch_size=batch_size, + draft_len=draft_len, + is_first_draft=spec_resource_manager.is_first_draft, + short_seq_len_mode=short_seq_len_mode, + is_all_greedy_sample=is_all_greedy_sample) else: # With dynamic spec decode, the draft length may be zero even when enable_spec_decode is True, # so we need to get the draft length from the batch instead of using enable_spec_decode. @@ -337,33 +349,46 @@ def get_graph_key( draft_len = max(draft_len_list) assert len( set(draft_len_list)) == 1, "All draft lengths must be the same" - context_query_lens = tuple( - int(request.context_chunk_size) - for request in batch.context_requests) - encoder_input_lens = (sum( - int(request.encoder_output_len) - for request in batch.context_requests - if not request.py_skip_cross_kv_projection), ) - key = (batch_size, draft_len, False, short_seq_len_mode, - is_all_greedy_sample, context_query_lens, encoder_input_lens) + context_requests = batch.context_requests + num_contexts = len(context_requests) + context_query_len = 0 + if num_contexts: + context_query_len = int(context_requests[0].context_chunk_size) + if any( + int(request.context_chunk_size) != context_query_len + for request in context_requests[1:]): + return None + num_encoder_tokens = sum( + int(request.encoder_output_len) for request in context_requests + if not request.py_skip_cross_kv_projection) + key = KeyType(batch_size=batch_size, + draft_len=draft_len, + is_first_draft=False, + short_seq_len_mode=short_seq_len_mode, + is_all_greedy_sample=is_all_greedy_sample, + num_contexts=num_contexts, + context_query_len=context_query_len, + num_encoder_tokens=num_encoder_tokens) return key def _get_compatible_mixed_encoder_decoder_key(self, key: KeyType) -> KeyType: """Round the packed encoder extent up to a captured graph key.""" if (not self.padding_enabled or self._capture_allowed - or key in self.graph_metadata or len(key[6]) != 1): + or key in self.graph_metadata or key.num_encoder_tokens == 0): return key - num_encoder_tokens = key[6][0] + key_without_encoder_extent = key._replace(num_encoder_tokens=0) compatible_keys = [ captured_key for captured_key in self.graph_outputs - if captured_key[:6] == key[:6] and len(captured_key[6]) == 1 - and captured_key[6][0] >= num_encoder_tokens + if isinstance(captured_key, KeyType) and captured_key._replace( + num_encoder_tokens=0) == key_without_encoder_extent + and captured_key.num_encoder_tokens >= key.num_encoder_tokens ] if not compatible_keys: return key - return min(compatible_keys, key=lambda captured_key: captured_key[6][0]) + return min(compatible_keys, + key=lambda captured_key: captured_key.num_encoder_tokens) @staticmethod def _get_mrope_position_delta(request: Any) -> Optional[Any]: @@ -403,7 +428,6 @@ def maybe_get_cuda_graph( draft_tokens_cuda: Optional[torch.Tensor] = None, new_tensors_device: Optional[SampleStateTensors] = None, spec_resource_manager: Optional[BaseResourceManager] = None, - allow_mixed_encoder_decoder: bool = False, promoted_context_request_ids: frozenset[int] = frozenset(), ) -> Tuple[Optional[Any], Optional[Any], Optional[KeyType]]: """ @@ -423,20 +447,17 @@ def maybe_get_cuda_graph( return None, None, None is_mixed_encoder_decoder = self._is_mixed_encoder_decoder_batch(batch) - can_run_cuda_graph = (batch.can_run_cuda_graph - or (is_mixed_encoder_decoder - and allow_mixed_encoder_decoder)) + can_run_cuda_graph = self._can_run_cuda_graph_batch(batch) batch_size = batch.batch_size if self.enabled and self.config.enable_attention_dp and self.config.mapping.tp_size > 1: - all_can_graph_batch = self.config.dist.tp_allgather( + graph_batch_info = self.config.dist.tp_allgather( [can_run_cuda_graph, batch_size]) - is_all_gen_only = all(all_can_graph[0] - for all_can_graph in all_can_graph_batch) - all_batch_size_equal = all( - all_gen_only[1] == all_can_graph_batch[0][1] - for all_gen_only in all_can_graph_batch) + all_can_run_cuda_graph = all(rank_info[0] + for rank_info in graph_batch_info) + all_batch_sizes_equal = all(rank_info[1] == graph_batch_info[0][1] + for rank_info in graph_batch_info) - if not is_all_gen_only or not all_batch_size_equal: + if not all_can_run_cuda_graph or not all_batch_sizes_equal: return None, None, None if not self.enabled or not can_run_cuda_graph: @@ -454,6 +475,8 @@ def maybe_get_cuda_graph( key = self.get_graph_key(batch, new_tensors_device, spec_resource_manager, spec_metadata, promoted_context_request_ids) + if key is None: + return None, None, None if is_mixed_encoder_decoder: key = self._get_compatible_mixed_encoder_decoder_key(key) @@ -473,16 +496,17 @@ def maybe_get_cuda_graph( num_sequences_in_batch = batch_size * self.max_beam_width graph_attn_metadata = attn_metadata.create_cuda_graph_metadata( - num_sequences_in_batch, False, key[1], self.cuda_graph_meta_buffers) + num_sequences_in_batch, False, key.draft_len, + self.cuda_graph_meta_buffers) if is_mixed_encoder_decoder: - context_query_lens = key[5] - generation_query_len = key[1] + 1 + generation_query_len = key.draft_len + 1 graph_attn_metadata.seq_lens = torch.tensor( - context_query_lens + (generation_query_len, ) * - (num_sequences_in_batch - len(context_query_lens)), + (key.context_query_len, ) * key.num_contexts + + (generation_query_len, ) * + (num_sequences_in_batch - key.num_contexts), dtype=torch.int, ) - graph_attn_metadata.num_contexts = len(context_query_lens) + graph_attn_metadata.num_contexts = key.num_contexts assert graph_attn_metadata.is_cuda_graph if enable_spec_decode: @@ -552,14 +576,9 @@ def get_graph_pool(self): return self.memory_pool def _get_num_tokens_for_key(self, key: KeyType) -> int: - batch_size = key[0] - token_per_generation = key[1] + 1 - # Direct capture callers predating mixed encoder-decoder graphs pass - # only the generation-key prefix. A missing suffix means gen-only. - context_query_lens = key[5] if len(key) > 5 else () - num_contexts = len(context_query_lens) - return (sum(context_query_lens) + - (batch_size * self.max_beam_width - num_contexts) * + token_per_generation = key.draft_len + 1 + return (key.num_contexts * key.context_query_len + + (key.batch_size * self.max_beam_width - key.num_contexts) * token_per_generation) def capture(self, @@ -569,7 +588,10 @@ def capture(self, enable_spec_decode: bool = False, postprocess_fn: Optional[Callable] = None) -> Any: """Warm up and/or capture the forward pass for a graph key.""" - batch_size = key[0] + # Preserve compatibility with direct callers that still pass the + # original three-field generation-only tuple. + key = KeyType(*key) + batch_size = key.batch_size # [CUDA graph spec decode padding] # We pad input IDs/position IDs to the maximum draft length (token per request). # We're forced to do this because we cannot reallocate inputs over many graph runs. @@ -593,8 +615,7 @@ def capture(self, capture_inputs = initial_inputs.copy() capture_inputs.update(sliced_static_tensors) - encoder_input_lens = key[6] if len(key) > 6 else () - num_encoder_tokens = sum(encoder_input_lens) + num_encoder_tokens = key.num_encoder_tokens if num_encoder_tokens: encoder_hidden_states = initial_inputs.get("encoder_hidden_states") if encoder_hidden_states is None: @@ -628,7 +649,7 @@ def capture(self, def _setup_spec_decoding_and_forward(key: KeyType, forward_fn: Callable, capture_inputs: Dict[str, Any]): - is_first_draft = key[2] + is_first_draft = key.is_first_draft needs_kv_cache_recompute = True if enable_spec_decode and self.config.spec_config.spec_dec_mode.needs_kv_cache_recompute( ) else False if is_first_draft and self.config.is_draft_model and needs_kv_cache_recompute: @@ -671,6 +692,7 @@ def _setup_spec_decoding_and_forward(key: KeyType, forward_fn: Callable, def replay(self, key: KeyType, current_inputs: Dict[str, Any]) -> Optional[torch.Tensor]: """Replays a previously captured graph.""" + key = KeyType(*key) stored_meta = self.graph_metadata[key] assert current_inputs["attn_metadata"] is stored_meta["attn_metadata"] if stored_meta["spec_metadata"] is not None: @@ -697,8 +719,7 @@ def replay(self, key: KeyType, else: static_tensors["position_ids"][:, :seqlen].copy_(position_ids) - encoder_input_lens = key[6] if len(key) > 6 else () - num_encoder_tokens = sum(encoder_input_lens) + num_encoder_tokens = key.num_encoder_tokens if num_encoder_tokens: encoder_hidden_states = current_inputs.get("encoder_hidden_states") if encoder_hidden_states is None: @@ -735,13 +756,13 @@ def _get_padded_batch(self, batch: ScheduledRequests, new_batch_size = batch_size if self.enabled and self.config.enable_attention_dp and self.config.mapping.tp_size > 1: - graph_batch_size = self.config.dist.tp_allgather( + graph_batch_info = self.config.dist.tp_allgather( [can_run_cuda_graph, batch_size]) - all_can_graph = all(graph_batch[0] - for graph_batch in graph_batch_size) - if all_can_graph: - new_batch_size = max(gen_only_batch[1] - for gen_only_batch in graph_batch_size) + all_can_run_cuda_graph = all(rank_info[0] + for rank_info in graph_batch_info) + if all_can_run_cuda_graph: + new_batch_size = max(rank_info[1] + for rank_info in graph_batch_info) if (not self.enabled or not self.padding_enabled or not can_run_cuda_graph @@ -1111,7 +1132,6 @@ def _get_dynamic_capture_key( a larger existing key; no new layout-specific key is created. """ batch_size = len(sequence_lengths) - sum(sequence_lengths) max_seq_len = max(sequence_lengths) if sequence_lengths else 0 candidate_batch_sizes = (self.supported_batch_sizes if allow_batch_padding else [batch_size]) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index c0591b35516c..578dd567b068 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -2305,12 +2305,18 @@ def _capture_mixed_encoder_decoder_cuda_graphs( f"{operation} for batch size={batch_size}, " f"context requests={num_contexts}, " f"packed encoder tokens={total_encoder_tokens}") - self.enable_spec_decode = False - self.runtime_draft_len = 0 - self.forward(batch, - new_tensors_device=None, - resource_manager=resource_manager) - torch.cuda.synchronize() + saved_enable_spec_decode = self.enable_spec_decode + saved_runtime_draft_len = self.runtime_draft_len + try: + self.enable_spec_decode = False + self.runtime_draft_len = 0 + self.forward(batch, + new_tensors_device=None, + resource_manager=resource_manager) + torch.cuda.synchronize() + finally: + self.enable_spec_decode = saved_enable_spec_decode + self.runtime_draft_len = saved_runtime_draft_len def _capture_piecewise_cuda_graphs(self, resource_manager: ResourceManager): """Captures piecewise CUDA graphs for context/prefill steps via torch.compile.""" @@ -3530,7 +3536,7 @@ def prepare_cross_metadata( cross_attn_metadata.prepare() return assert isinstance(cross_attn_metadata, TrtllmAttentionMetadata) - cross_attn_metadata.prepare_encoder_decoder( + cross_attn_metadata.prepare_encoder_decoder_from_precomputed_lengths( prompt_lens=attn_metadata.prompt_lens, kv_lens=encoder_kv_lens, context_kv_tokens=context_encoder_kv_tokens, @@ -3814,7 +3820,7 @@ def _prepare_encoder_decoder_inputs_fast( num_extra_kv_tokens=0) attn_metadata.kv_cache_manager = kv_cache_manager assert isinstance(attn_metadata, TrtllmAttentionMetadata) - attn_metadata.prepare_encoder_decoder( + attn_metadata.prepare_encoder_decoder_from_precomputed_lengths( prompt_lens=buffers['prompt_lengths'][:num_sequences], kv_lens=buffers['kv_lengths'][:num_sequences], context_kv_tokens=context_kv_tokens, @@ -6937,8 +6943,6 @@ def forward(self, padded_graph_requests.all_requests()) self._sync_group_all_greedy_sample(spec_metadata) - allow_mixed_encoder_decoder_graph = ( - self.cuda_graph_runner.enable_encoder_decoder_mixed_cuda_graph) maybe_attn_metadata, maybe_spec_metadata, key = self.cuda_graph_runner.maybe_get_cuda_graph( padded_graph_requests, enable_spec_decode=self.enable_spec_decode, @@ -6948,7 +6952,6 @@ def forward(self, if self.is_spec_decode else None, new_tensors_device=new_tensors_device, spec_resource_manager=spec_resource_manager, - allow_mixed_encoder_decoder=(allow_mixed_encoder_decoder_graph), promoted_context_request_ids=promoted_context_request_ids, ) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/penalties.py b/tensorrt_llm/_torch/pyexecutor/sampler/penalties.py index 4f3ea18a2613..b688a5c696a1 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/penalties.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/penalties.py @@ -31,10 +31,10 @@ from .ops.vanilla import Fusions from .sampler_common import _get_max_beam_width, _unwrap_singleton -__all__ = ["PenaltyHandler", "PenaltyStore"] +__all__ = ["PenaltyHandler", "PenaltyStore", "has_occurrence_penalty"] -def _has_occurrence_penalty(request: LlmRequest) -> bool: +def has_occurrence_penalty(request: LlmRequest) -> bool: sampling_config = request.sampling_config repetition = _unwrap_singleton(sampling_config.repetition_penalty) presence = _unwrap_singleton(sampling_config.presence_penalty) @@ -219,7 +219,7 @@ def validate_request(request: LlmRequest) -> None: Called from ``TorchSampler.validate_request`` (request admission), so a violating request is failed individually instead of aborting the whole batch. """ - if _get_max_beam_width(request) > 1 and _has_occurrence_penalty(request): + if _get_max_beam_width(request) > 1 and has_occurrence_penalty(request): raise ValueError( "TorchSampler does not support repetition, presence, or frequency " "penalties with beam search." @@ -240,7 +240,7 @@ def prepare_for_new_request(self, request: LlmRequest, slot: int) -> None: gathered, so their stale parameters/counts are left untouched. """ was_active = self._slots[slot] is not None - if not (_get_max_beam_width(request) == 1 and _has_occurrence_penalty(request)): + if not (_get_max_beam_width(request) == 1 and has_occurrence_penalty(request)): self._slots[slot] = None if was_active: self._num_active_slots -= 1 diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index ae55460fd511..fa58a09397a1 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -98,7 +98,7 @@ get_logprobs_from_request, store_logprobs_list_to_request, ) -from .penalties import PenaltyHandler +from .penalties import PenaltyHandler, has_occurrence_penalty from .sampler_common import ( DEFAULT_BEAM_IDX, DEFAULT_STEP_IDX, @@ -3869,6 +3869,7 @@ def _process_requests( request._py_embedding_bias_1d is None and not getattr(request, "py_bad_words", None) and not getattr(request, "py_no_repeat_ngram_size", None) + and not has_occurrence_penalty(request) and not request.py_min_length and not request.py_return_log_probs and not request.py_stop_words_list diff --git a/tests/integration/defs/llmapi/test_llm_api_pytorch_bart.py b/tests/integration/defs/llmapi/test_llm_api_pytorch_bart.py index 8185497178ed..4b3194048787 100644 --- a/tests/integration/defs/llmapi/test_llm_api_pytorch_bart.py +++ b/tests/integration/defs/llmapi/test_llm_api_pytorch_bart.py @@ -681,7 +681,9 @@ def test_bart_pytorch_continuous_admission_replays_encoder_and_mixed_cuda_graphs assert encoder_runner.enabled assert encoder_runner.graphs - captured_mixed_keys = {key for key in decoder_runner.graphs if key[5] and key[6]} + captured_mixed_keys = { + key for key in decoder_runner.graphs if key.num_contexts and key.num_encoder_tokens + } assert captured_mixed_keys original_encoder_replay = encoder_runner.replay @@ -742,6 +744,8 @@ def record_decoder_replay(key, inputs): admitted_encoder_keys = encoder_replay_keys[encoder_replay_count_before_admission:] assert any(key[0] == 2 for key in admitted_encoder_keys) assert set(encoder_replay_keys) <= set(encoder_runner.graphs) - replayed_mixed_keys = {key for key in decoder_replay_keys if key[5] and key[6]} + replayed_mixed_keys = { + key for key in decoder_replay_keys if key.num_contexts and key.num_encoder_tokens + } assert replayed_mixed_keys assert replayed_mixed_keys <= captured_mixed_keys diff --git a/tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py b/tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py index 46f525565ec8..5f277567f111 100644 --- a/tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py +++ b/tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py @@ -882,7 +882,9 @@ def test_t5_pytorch_continuous_admission_replays_encoder_and_mixed_cuda_graphs( assert encoder_runner.enabled assert encoder_runner.graphs assert encoder_runner.use_fixed_sequence_slots - captured_mixed_keys = {key for key in decoder_runner.graphs if key[5] and key[6]} + captured_mixed_keys = { + key for key in decoder_runner.graphs if key.num_contexts and key.num_encoder_tokens + } assert captured_mixed_keys encoder_replay_keys = [] @@ -948,6 +950,8 @@ def record_decoder_replay(key, inputs): admitted_encoder_keys = encoder_replay_keys[encoder_replay_count_before_admission:] assert any(key[0] == 2 for key in admitted_encoder_keys) assert set(encoder_replay_keys) <= set(encoder_runner.graphs) - replayed_mixed_keys = {key for key in decoder_replay_keys if key[5] and key[6]} + replayed_mixed_keys = { + key for key in decoder_replay_keys if key.num_contexts and key.num_encoder_tokens + } assert replayed_mixed_keys assert replayed_mixed_keys <= captured_mixed_keys diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index efab11083058..0dcd09aa8e95 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -16,8 +16,8 @@ from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_connector import \ KvCacheConnectorWorker from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import ( - CUDAGraphRunner, EncoderCUDAGraphRunner, _restore_spec_decode_capture_state, - _save_spec_decode_capture_state) + CUDAGraphRunner, EncoderCUDAGraphRunner, KeyType, + _restore_spec_decode_capture_state, _save_spec_decode_capture_state) from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest from tensorrt_llm._torch.pyexecutor.model_engine import ( PyTorchModelEngine, _build_request_multimodal_input, @@ -188,8 +188,7 @@ def _make_request_stub(req_id: int, prompt_len: int = 4) -> SimpleNamespace: def _make_forward_only_engine( - graph_key: tuple[int, int, bool, bool, bool, tuple[int, ...], - tuple[int, ...]] | None, + graph_key: KeyType | None, runner_enabled: bool = True, ) -> tuple[PyTorchModelEngine, Mock, Mock, Mock, dict[str, object]]: engine = object.__new__(PyTorchModelEngine) @@ -611,7 +610,83 @@ def test_graph_key_forwards_promoted_context_ids(self) -> None: runner._get_seq_len_mode.assert_called_once_with( batch, None, promoted_ids) - self.assertEqual(key, (1, 0, False, True, True, (), (0, ))) + self.assertEqual( + key, + KeyType(batch_size=1, + draft_len=0, + is_first_draft=False, + short_seq_len_mode=True, + num_encoder_tokens=0)) + + def test_graph_key_aggregates_encoder_tokens(self) -> None: + runner = Mock() + runner.config = SimpleNamespace(is_draft_model=False) + runner.max_beam_width = 1 + runner._get_seq_len_mode.return_value = False + context = _make_request_stub(1) + context.encoder_output_len = 7 + context.py_skip_cross_kv_projection = False + skipped_context = _make_request_stub(2) + skipped_context.encoder_output_len = 11 + skipped_context.py_skip_cross_kv_projection = True + batch = ScheduledRequests() + batch.context_requests_last_chunk = [context, skipped_context] + batch.generation_requests = [_make_request_stub(3)] + + key = CUDAGraphRunner.get_graph_key(runner, batch) + + assert key is not None + self.assertEqual(key.num_contexts, 2) + self.assertEqual(key.context_query_len, 1) + self.assertEqual(key.num_encoder_tokens, 7) + self.assertEqual(CUDAGraphRunner._get_num_tokens_for_key(runner, key), + 3) + + def test_graph_key_rejects_nonuniform_context_query_lengths(self) -> None: + runner = Mock() + runner.config = SimpleNamespace(is_draft_model=False) + runner._get_seq_len_mode.return_value = False + first_context = _make_request_stub(1) + first_context.encoder_output_len = 7 + first_context.py_skip_cross_kv_projection = False + second_context = _make_request_stub(2) + second_context.context_chunk_size = 2 + second_context.encoder_output_len = 11 + second_context.py_skip_cross_kv_projection = False + batch = ScheduledRequests() + batch.context_requests_last_chunk = [first_context, second_context] + batch.generation_requests = [_make_request_stub(3)] + + key = CUDAGraphRunner.get_graph_key(runner, batch) + + self.assertIsNone(key) + + def test_graph_key_rounds_encoder_tokens_up_to_captured_extent( + self) -> None: + key = KeyType(batch_size=2, + draft_len=0, + is_first_draft=False, + num_contexts=1, + context_query_len=1, + num_encoder_tokens=7) + smaller_key = key._replace(num_encoder_tokens=6) + compatible_key = key._replace(num_encoder_tokens=8) + larger_key = key._replace(num_encoder_tokens=16) + runner = SimpleNamespace( + padding_enabled=True, + _capture_allowed=False, + graph_metadata={}, + graph_outputs={ + smaller_key: object(), + compatible_key: object(), + larger_key: object(), + }, + ) + + actual_key = CUDAGraphRunner._get_compatible_mixed_encoder_decoder_key( + runner, key) + + self.assertEqual(actual_key, compatible_key) def test_graph_lookup_forwards_promoted_context_ids(self) -> None: runner = Mock() @@ -620,7 +695,11 @@ def test_graph_lookup_forwards_promoted_context_ids(self) -> None: enable_attention_dp=False, use_mrope=False, ) - key = (1, 0, False, True, True, (), (0, )) + key = KeyType(batch_size=1, + draft_len=0, + is_first_draft=False, + short_seq_len_mode=True, + num_encoder_tokens=0) graph_attn_metadata = object() graph_spec_metadata = object() runner.get_graph_key.return_value = key @@ -654,7 +733,7 @@ def test_graph_lookup_forwards_promoted_context_ids(self) -> None: (graph_attn_metadata, graph_spec_metadata, key)) def test_forward_commits_candidate_only_on_graph_hit(self) -> None: - key = (2, 0, False, False, True, (), (0, )) + key = KeyType(batch_size=2, draft_len=0, is_first_draft=False) engine, runner, resource_manager, _, outputs = \ _make_forward_only_engine(key) context = _make_request_stub(1) @@ -713,7 +792,7 @@ def test_forward_graph_miss_uses_semantic_eager_batch(self) -> None: def test_zero_runtime_draft_speculation_commits_graph_candidate( self) -> None: - key = (2, 0, False, False, True, (), (0, )) + key = KeyType(batch_size=2, draft_len=0, is_first_draft=False) engine, runner, resource_manager, semantic_attn_metadata, outputs = \ _make_forward_only_engine(key) engine.enable_spec_decode = True @@ -809,7 +888,7 @@ def test_zero_runtime_non_linear_tree_speculation_uses_semantic_eager_batch( runner.replay.assert_not_called() def test_forward_allows_guided_context_logits_on_graph_hit(self) -> None: - key = (1, 0, False, False, True, (), (0, )) + key = KeyType(batch_size=1, draft_len=0, is_first_draft=False) engine, runner, resource_manager, _, outputs = \ _make_forward_only_engine(key) engine.guided_decoder = Mock() @@ -862,7 +941,7 @@ def test_multimodal_graph_miss_preserves_semantic_payload(self) -> None: self.assertIn("multimodal_embedding", multimodal_data) def test_generation_only_forward_does_not_call_new_selector(self) -> None: - key = (1, 0, False, False, True, (), (0, )) + key = KeyType(batch_size=1, draft_len=0, is_first_draft=False) engine, runner, resource_manager, _, _ = _make_forward_only_engine(key) generation = _make_request_stub(2) batch = ScheduledRequests() diff --git a/tests/unittest/_torch/sampler/test_torch_sampler.py b/tests/unittest/_torch/sampler/test_torch_sampler.py index 1edca6b81b96..73861da9d988 100644 --- a/tests/unittest/_torch/sampler/test_torch_sampler.py +++ b/tests/unittest/_torch/sampler/test_torch_sampler.py @@ -786,6 +786,57 @@ def test_greedy_no_repeat_ngram_uses_token_ban_path(): assert new_tokens_host.reshape(-1)[0].item() == 3 +@force_ampere +@pytest.mark.parametrize( + ("penalty_name", "penalty_value"), + [ + pytest.param("repetition_penalty", 100.0, id="repetition"), + pytest.param("presence_penalty", 2.0, id="presence"), + pytest.param("frequency_penalty", 2.0, id="frequency"), + ], +) +def test_greedy_occurrence_penalties_bypass_stable_path(penalty_name: str, penalty_value: float): + sampler = TorchSampler( + TorchSampler.Args( + max_seq_len=16, + max_draft_len=0, + max_num_sequences=1, + max_beam_width=1, + max_total_draft_tokens=0, + disable_overlap_scheduler=True, + ) + ) + request = LlmRequest( + request_id=0, + max_new_tokens=4, + input_tokens=[1], + sampling_config=SamplingConfig( + SamplingParams(**{penalty_name: penalty_value})._get_sampling_config() + ), + seq_slot=0, + is_streaming=False, + ) + + admission = ScheduledRequests() + admission.context_requests_last_chunk = [request] + sampler.setup_sampler_step(admission) + + scheduled_requests = ScheduledRequests() + scheduled_requests.generation_requests = [request] + logits = torch.tensor([[0.0, 10.0, 9.0]], device="cuda") + + *_, new_tokens_host, single_step_greedy = sampler._process_requests( + scheduled_requests, + {"logits": logits}, + sampler.store.new_tokens, + [0], + ) + torch.cuda.synchronize() + + assert not single_step_greedy + assert new_tokens_host.reshape(-1)[0].item() == 2 + + class TestFinishReasons: NOT_FINISHED = FinishReason.NOT_FINISHED STOP_WORDS = FinishReason.STOP_WORDS From 42c8bff12682eb4061004b3274d287f6c0d54e4d Mon Sep 17 00:00:00 2001 From: Guiju Zhang <7135567+cascade812@users.noreply.github.com> Date: Wed, 5 Aug 2026 20:53:24 -0700 Subject: [PATCH 14/15] fix failed test Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com> --- .../defs/kv_cache/test_final_single_token_context_cuda_graph.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/integration/defs/kv_cache/test_final_single_token_context_cuda_graph.py b/tests/integration/defs/kv_cache/test_final_single_token_context_cuda_graph.py index db1b6314aa9f..9ef8cf87232b 100644 --- a/tests/integration/defs/kv_cache/test_final_single_token_context_cuda_graph.py +++ b/tests/integration/defs/kv_cache/test_final_single_token_context_cuda_graph.py @@ -91,7 +91,6 @@ def maybe_get_cuda_graph( draft_tokens_cuda: torch.Tensor | None = None, new_tensors_device: SampleStateTensors | None = None, spec_resource_manager: BaseResourceManager | None = None, - allow_mixed_encoder_decoder: bool = False, promoted_context_request_ids: frozenset[int] = frozenset(), ) -> tuple[Any | None, Any | None, KeyType | None]: # A new decision means the preceding one reached eager execution if it @@ -105,7 +104,6 @@ def maybe_get_cuda_graph( draft_tokens_cuda, new_tensors_device, spec_resource_manager, - allow_mixed_encoder_decoder=allow_mixed_encoder_decoder, promoted_context_request_ids=promoted_context_request_ids, ) if promoted_context_request_ids: From 0789f212c4edcbcfbb9bf6673cffff4ddfd5ad4f Mon Sep 17 00:00:00 2001 From: Guiju Zhang <7135567+cascade812@users.noreply.github.com> Date: Thu, 6 Aug 2026 17:05:55 -0700 Subject: [PATCH 15/15] address comment Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com> --- .../_torch/pyexecutor/sampler/sampler.py | 14 +++-- .../_torch/sampler/test_torch_sampler.py | 61 +++++++++++++++++++ 2 files changed, 71 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index e49d72a82dfd..0ea5b5cf340a 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -1660,6 +1660,7 @@ def __init__(self, args: Args): None ] * self.max_num_sequences self._stable_greedy_request_ids: list[int] = [] + self._stable_greedy_seq_slots: list[int] = [] self._stable_greedy_seq_slots_host: Optional[torch.Tensor] = None self._stable_greedy_seq_slots_cuda: Optional[torch.Tensor] = None @@ -4056,7 +4057,11 @@ def _process_requests( generation_requests = scheduled_requests.generation_requests request_ids = [request.py_request_id for request in generation_requests] - has_stable_request_ids = self._stable_greedy_request_ids == request_ids + maybe_seq_slots = [request.py_seq_slot for request in generation_requests] + has_stable_greedy_batch = ( + self._stable_greedy_request_ids == request_ids + and self._stable_greedy_seq_slots == maybe_seq_slots + ) can_use_stable_greedy_path = ( bool(generation_requests) and self.max_beam_width == 1 @@ -4068,7 +4073,7 @@ def _process_requests( for request in generation_requests ) and ( - has_stable_request_ids + has_stable_greedy_batch or all( request._py_embedding_bias_1d is None and not getattr(request, "py_bad_words", None) @@ -4083,13 +4088,12 @@ def _process_requests( ) ) if can_use_stable_greedy_path: - if has_stable_request_ids: + if has_stable_greedy_batch: assert self._stable_greedy_seq_slots_host is not None assert self._stable_greedy_seq_slots_cuda is not None seq_slots_host = self._stable_greedy_seq_slots_host seq_slots_cuda = self._stable_greedy_seq_slots_cuda else: - maybe_seq_slots = [request.py_seq_slot for request in generation_requests] assert all(seq_slot is not None for seq_slot in maybe_seq_slots) seq_slots = [cast(int, seq_slot) for seq_slot in maybe_seq_slots] seq_slots_host = torch.tensor( @@ -4099,6 +4103,7 @@ def _process_requests( device="cuda", dtype=torch.int64, non_blocking=True ) self._stable_greedy_request_ids = request_ids + self._stable_greedy_seq_slots = seq_slots self._stable_greedy_seq_slots_host = seq_slots_host self._stable_greedy_seq_slots_cuda = seq_slots_cuda @@ -4121,6 +4126,7 @@ def _process_requests( ) self._stable_greedy_request_ids = [] + self._stable_greedy_seq_slots = [] sampling_requests, sampling_requests_metadata, logits_cuda = self._select_generated_logits( scheduled_requests, diff --git a/tests/unittest/_torch/sampler/test_torch_sampler.py b/tests/unittest/_torch/sampler/test_torch_sampler.py index 7be0c48ebac3..4a29f602bba6 100644 --- a/tests/unittest/_torch/sampler/test_torch_sampler.py +++ b/tests/unittest/_torch/sampler/test_torch_sampler.py @@ -751,6 +751,67 @@ def _uut(res=res): run_test_with_warmup(_test_runner, max_sync_s=0.3) +def test_stable_greedy_cache_key_includes_sequence_slots(monkeypatch: pytest.MonkeyPatch): + sampler = object.__new__(TorchSampler) + sampler.max_beam_width = 1 + sampler._stable_greedy_request_ids = [] + sampler._stable_greedy_seq_slots = [] + sampler._stable_greedy_seq_slots_host = None + sampler._stable_greedy_seq_slots_cuda = None + monkeypatch.setattr(sampler, "_copy_to_host", lambda tensor: tensor.clone()) + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.sampler.sampler.prefer_pinned", lambda: False + ) + + original_tensor_to = torch.Tensor.to + + def copy_without_cuda(tensor: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor: + if kwargs.get("device") == "cuda": + kwargs["device"] = "cpu" + return original_tensor_to(tensor, *args, **kwargs) + + monkeypatch.setattr(torch.Tensor, "to", copy_without_cuda) + + logits = torch.tensor([[0.0, 1.0, 2.0]]) + new_tokens = torch.zeros((1, 2, 1), dtype=torch.int32) + requests = [ + LlmRequest( + request_id=0, + max_new_tokens=4, + input_tokens=[1], + sampling_config=SamplingConfig(), + seq_slot=seq_slot, + is_streaming=False, + is_draft=is_draft, + ) + for seq_slot, is_draft in ((0, False), (1, True)) + ] + + for seq_slot, request in enumerate(requests): + scheduled_requests = ScheduledRequests() + scheduled_requests.generation_requests = [request] + + ( + _, + seq_slots_host, + _, + seq_slots_cuda, + _, + _, + single_step_greedy, + ) = sampler._process_requests( + scheduled_requests, + {"logits": logits}, + new_tokens, + [0], + ) + + assert single_step_greedy + assert seq_slots_host.tolist() == [seq_slot] + assert seq_slots_cuda.tolist() == [seq_slot] + assert new_tokens[0, seq_slot, 0].item() == 2 + + @force_ampere def test_greedy_no_repeat_ngram_uses_token_ban_path(): sampler = TorchSampler(