From 796ea434624334505eb8eabe7c4a1dfdb169aa30 Mon Sep 17 00:00:00 2001 From: Guiju Zhang <7135567+cascade812@users.noreply.github.com> Date: Wed, 24 Jun 2026 15:51:41 -0700 Subject: [PATCH 1/9] add cuda graph support for decoder generation request Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com> --- .../_torch/attention_backend/interface.py | 67 ++++- .../_torch/pyexecutor/cuda_graph_runner.py | 59 +++- .../_torch/pyexecutor/kv_cache_manager_v2.py | 8 +- .../_torch/pyexecutor/model_engine.py | 272 ++++++++++++++---- tensorrt_llm/_torch/pyexecutor/py_executor.py | 57 ++-- .../_torch/pyexecutor/resource_manager.py | 11 +- .../defs/llmapi/test_llm_api_pytorch_bart.py | 100 ++++++- .../defs/llmapi/test_llm_api_pytorch_t5.py | 87 +++++- tests/unittest/_torch/helpers.py | 1 + 9 files changed, 554 insertions(+), 108 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/interface.py b/tensorrt_llm/_torch/attention_backend/interface.py index c1630bbf5041..b92a059f0b50 100644 --- a/tensorrt_llm/_torch/attention_backend/interface.py +++ b/tensorrt_llm/_torch/attention_backend/interface.py @@ -433,10 +433,9 @@ def update_helix_param( def create_cross_metadata( self, - encoder_seq_lens: torch.Tensor, + encoder_seq_lens: List[int], cross_kv_cache_manager: Union[KVCacheManager, KVCacheManagerV2, None] = None, - *, encoder_num_cached_tokens_per_seq: Optional[List[int]] = None, ) -> "AttentionMetadata": """Build a sub-metadata instance for cross-attention. @@ -452,11 +451,10 @@ def create_cross_metadata( ``self.cross``); callers can attach it to ``self.cross`` if desired. Args: - encoder_seq_lens: Per-request encoder sequence length (CPU - int32 tensor). On the first decoder context step this is - the full encoder length; on generation steps it should be - ``0`` (no new K/V tokens to add to the cross pool — the - encoder K/V are already cached). + encoder_seq_lens: Per-request encoder sequence lengths. On the + first decoder context step this is the full encoder length; + on generation steps it should be ``0`` (no new K/V tokens to + add to the cross pool — the encoder K/V are already cached). cross_kv_cache_manager: KV cache manager for the cross pool. When ``None``, the returned metadata uses the stateless (no-KV-cache) path (suitable for unit tests). @@ -469,6 +467,8 @@ def create_cross_metadata( with ``seq_lens_kv`` set to ``encoder_seq_lens`` so that ``is_cross`` becomes ``True``. """ + encoder_seq_lens_tensor = torch.tensor(encoder_seq_lens, + dtype=torch.int) cross_md = copy.copy(self) cross_md._saved_tensors = {} if self.is_cuda_graph: @@ -478,20 +478,17 @@ def create_cross_metadata( # cannot overwrite self-attention sequence lengths. cross_md.cuda_graph_buffers = Buffers() cross_md.kv_cache_manager = cross_kv_cache_manager - cross_md._seq_lens_kv = None cross_md._seq_lens_kv_cuda = None cross_md.cross = None - cross_md.seq_lens_kv = encoder_seq_lens + cross_md.seq_lens_kv = encoder_seq_lens_tensor if encoder_num_cached_tokens_per_seq is not None: - from ..metadata import KVCacheParams base_params = self.kv_cache_params cross_md.kv_cache_params = KVCacheParams( use_cache=base_params.use_cache if base_params is not None else (cross_kv_cache_manager is not None), num_cached_tokens_per_seq=list( encoder_num_cached_tokens_per_seq), - block_ids_per_seq=base_params.block_ids_per_seq - if base_params is not None else None, + block_ids_per_seq=None, host_max_attention_window_sizes=base_params. host_max_attention_window_sizes if base_params is not None else None, @@ -503,6 +500,52 @@ def create_cross_metadata( cross_md.__post_init__() return cross_md + def update_cross_metadata( + self, + encoder_seq_lens: List[int], + cross_kv_cache_manager: Union[KVCacheManager, KVCacheManagerV2, None], + encoder_num_cached_tokens_per_seq: Optional[List[int]] = None, + ) -> "AttentionMetadata": + """Refresh an existing CUDA graph cross-attention sub-metadata.""" + if not self.has_cross_sub_metadata: + raise RuntimeError( + "CUDA graph cross-attention metadata has not been initialized.") + + cross_md = self.cross + encoder_seq_lens_tensor = torch.tensor(encoder_seq_lens, + dtype=torch.int) + cross_md.kv_cache_manager = cross_kv_cache_manager + cross_md._seq_lens = self.seq_lens + cross_md._seq_lens_cuda = self.seq_lens_cuda + cross_md.seq_lens_kv = encoder_seq_lens_tensor + + # Cross-attention keeps decoder-side prompt lengths for the Q-side + # context metadata. Encoder-side lengths are represented by + # seq_lens_kv and kv_cache_params.num_cached_tokens_per_seq. + cross_md.prompt_lens = self.prompt_lens + + if encoder_num_cached_tokens_per_seq is not None: + base_params = cross_md.kv_cache_params + cross_md.kv_cache_params = KVCacheParams( + use_cache=(base_params.use_cache if base_params is not None else + (cross_md.kv_cache_manager is not None)), + num_cached_tokens_per_seq=list( + encoder_num_cached_tokens_per_seq), + block_ids_per_seq=(base_params.block_ids_per_seq + if base_params is not None else None), + host_max_attention_window_sizes=( + base_params.host_max_attention_window_sizes + if base_params is not None else None), + host_sink_token_length=(base_params.host_sink_token_length + if base_params is not None else None), + num_extra_kv_tokens=(base_params.num_extra_kv_tokens + if base_params is not None else 0), + ) + + cross_md.request_ids = self.request_ids + cross_md.num_contexts = self.num_contexts + return cross_md + def update_for_spec_dec(self) -> None: """ Hook to be called during forward when using spec-dec one-model mode. diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index 478639035c62..883a82c677a5 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -78,6 +78,7 @@ class CUDAGraphRunnerConfig: original_max_total_draft_tokens: int is_draft_model: bool enable_attention_dp: bool + is_encoder_decoder: bool batch_size: int mapping: Optional[Mapping] dist: Optional[Distributed] @@ -107,13 +108,14 @@ def __init__(self, config: CUDAGraphRunnerConfig): self.max_beam_width = config.max_beam_width self.spec_config = config.spec_config self.sparse_config = config.sparse_attention_config + self.is_encoder_decoder = config.is_encoder_decoder self.graphs: Dict[KeyType, torch.cuda.CUDAGraph] = {} self.graph_outputs: Dict[KeyType, Callable[[], Optional[torch.Tensor]]] = {} self.graph_metadata: Dict[KeyType, Dict[str, Any]] = {} self.memory_pool = config.cuda_graph_mem_pool - self.padding_dummy_requests: Dict[int, "Request"] = {} + self.padding_dummy_requests: Dict[int, Any] = {} self.dynamic_draft_len_mapping = config.dynamic_draft_len_mapping self.shared_static_tensors: Dict[str, torch.Tensor] = {} @@ -491,6 +493,14 @@ def _get_padded_batch(self, batch: ScheduledRequests, # respect the requirement just in case that changes in the future. # Use per-draft-len dummy requests for dynamic draft length support. if runtime_draft_len not in self.padding_dummy_requests: + dummy_encoder_output_len = None + if self.is_encoder_decoder: + cross_kv_cache_manager = resource_manager.get_resource_manager( + ResourceManagerType.CROSS_KV_CACHE_MANAGER) + if cross_kv_cache_manager is None: + return 0 + dummy_encoder_output_len = self._get_padding_dummy_encoder_output_len( + batch, cross_kv_cache_manager) # Get draft KV cache manager only for one-model speculative decoding. # In two-model mode, each model has its own KV cache manager, so @@ -502,10 +512,13 @@ def _get_padded_batch(self, batch: ScheduledRequests, dummy_request_id = CUDA_GRAPH_DUMMY_REQUEST_ID - runtime_draft_len dummy_request = kv_cache_manager.add_dummy_requests( [dummy_request_id], + token_nums=[2] if self.is_encoder_decoder else None, is_gen=True, max_num_draft_tokens=runtime_draft_len, use_mrope=self.config.use_mrope, max_beam_width=self.config.max_beam_width, + encoder_output_lens=[dummy_encoder_output_len] + if dummy_encoder_output_len is not None else None, draft_kv_cache_manager=draft_kv_cache_manager) if dummy_request is None: @@ -513,6 +526,14 @@ def _get_padded_batch(self, batch: ScheduledRequests, else: dummy_request = dummy_request[0] dummy_request.is_cuda_graph_dummy = True + if self.is_encoder_decoder: + if not self._prepare_encoder_decoder_padding_dummy( + dummy_request, resource_manager, + dummy_encoder_output_len): + kv_cache_manager.free_resources(dummy_request) + if draft_kv_cache_manager is not None: + draft_kv_cache_manager.free_resources(dummy_request) + return 0 spec_res_mgr = resource_manager.get_resource_manager( ResourceManagerType.SPEC_RESOURCE_MANAGER) @@ -524,6 +545,42 @@ def _get_padded_batch(self, batch: ScheduledRequests, batch.generation_requests.extend([padding_dummy_request] * padding_size) return padding_size + def _prepare_encoder_decoder_padding_dummy( + self, dummy_request: Any, resource_manager: ResourceManager, + encoder_output_len: int) -> bool: + cross_kv_cache_manager = resource_manager.get_resource_manager( + ResourceManagerType.CROSS_KV_CACHE_MANAGER) + if cross_kv_cache_manager is None: + return False + + dummy_request.py_encoder_output = None + dummy_request.py_skip_cross_kv_projection = True + + encoder_output_lens = [encoder_output_len] + cross_dummy_requests = cross_kv_cache_manager.add_dummy_requests( + request_ids=[dummy_request.py_request_id], + token_nums=encoder_output_lens, + is_gen=True, + max_beam_width=self.config.max_beam_width, + encoder_output_lens=encoder_output_lens) + return cross_dummy_requests is not None + + @staticmethod + def _get_padding_dummy_encoder_output_len( + batch: ScheduledRequests, cross_kv_cache_manager: Any) -> int: + encoder_output_len = 1 + for request in batch.generation_requests: + request_encoder_output_len = getattr(request, "encoder_output_len", + None) + if request_encoder_output_len is not None: + encoder_output_len = max(1, int(request_encoder_output_len)) + break + + max_seq_len = getattr(cross_kv_cache_manager, "max_seq_len", None) + if max_seq_len is not None: + encoder_output_len = min(encoder_output_len, int(max_seq_len)) + return encoder_output_len + def _round_up_batch_size(self, batch_size: int) -> int: """Finds the smallest supported graph batch size >= the given size.""" if not self.supported_batch_sizes: diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 02e48743465c..ac6330fa5bc6 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -1501,6 +1501,7 @@ def add_dummy_requests( kv_reserve_draft_tokens: Optional[int] = None, use_mrope: bool = False, max_beam_width: int = 1, + encoder_output_lens: Optional[List[int]] = None, num_extra_decoding_steps: int = 0, draft_kv_cache_manager: Optional["BaseResourceManager"] = None, ): @@ -1534,8 +1535,10 @@ def release_resources( token_num = token_nums[i] if token_nums is not None else 1 + max_num_draft_tokens # token_num - 1 is the past history length in generation. history_hint = max(0, token_num - 1) if is_gen else None - # TODO: support cross attention - encoder_input_tokens = None + encoder_output_len = encoder_output_lens[i] if encoder_output_lens is not None else None + encoder_input_tokens = ( + [1] * encoder_output_len if encoder_output_len is not None else None + ) # Using 1 instead of 0 prevents NaN during warmup in e.g. Deepseek input_tokens = [1 for _ in range(token_num)] req = LlmRequest( @@ -1545,6 +1548,7 @@ def release_resources( sampling_config=SamplingConfig(sampling_params._get_sampling_config()), is_streaming=False, encoder_input_tokens=encoder_input_tokens, + encoder_output_len=encoder_output_len, ) req.is_dummy_request = True req.paged_kv_block_ids = [] diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index ff1e5f8a5306..b8b60ec82894 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -25,7 +25,7 @@ strip_mm_data_for_generation) from tensorrt_llm.inputs.registry import (create_input_processor, create_input_processor_with_hash) -from tensorrt_llm.llmapi.llm_args import (CudaGraphConfig, +from tensorrt_llm.llmapi.llm_args import (CudaGraphConfig, DecodingBaseConfig, EncodeCudaGraphConfig, SeqLenAwareSparseAttentionConfig, TorchCompileConfig, TorchLlmArgs) @@ -71,7 +71,7 @@ from .guided_decoder import CapturableGuidedDecoder from .kv_cache_manager_v2 import KVCacheManagerV2 from .layerwise_nvtx_marker import LayerwiseNvtxMarker -from .llm_request import (LlmRequest, get_draft_token_length, +from .llm_request import (LlmRequest, LlmRequestState, get_draft_token_length, get_multimodal_embedding_lengths) from .mamba_cache_manager import MambaHybridCacheManager from .model_loader import ModelLoader, _construct_checkpoint_loader @@ -230,7 +230,7 @@ def __init__( mapping: Optional[Mapping] = None, attn_runtime_features: Optional[AttentionRuntimeFeatures] = None, dist: Optional[Distributed] = None, - spec_config: Optional["DecodingBaseConfig"] = None, + spec_config: Optional[DecodingBaseConfig] = None, is_draft_model: bool = False, drafting_loop_wrapper: Optional[Callable[[torch.nn.Module], torch.nn.Module]] = None, @@ -553,6 +553,16 @@ def __init__( self.encoder_attn_metadata = None self.spec_metadata = None self.iter_states = {} + self.cuda_graph_phase_stats = { + "encoder_eager": 0, + "encoder_graph": 0, + "decoder_context_eager": 0, + "decoder_context_graph": 0, + "decoder_generation_eager": 0, + "decoder_generation_graph": 0, + "decoder_mixed_eager": 0, + "decoder_mixed_graph": 0, + } 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 @@ -659,6 +669,7 @@ def __init__( original_max_total_draft_tokens, is_draft_model=self.is_draft_model, enable_attention_dp=self.enable_attention_dp, + is_encoder_decoder=self._is_encoder_decoder_model(), batch_size=self.batch_size, mapping=self.mapping, dist=self.dist, @@ -799,6 +810,30 @@ def moe_load_balancer_iter_info(self, value: Tuple[bool, bool]): moe_load_balancer.set_iter_info(enable_statistic=value[0], enable_update_weights=value[1]) + def get_cuda_graph_phase_stats(self) -> Dict[str, int]: + return dict(self.cuda_graph_phase_stats) + + def _record_encoder_cuda_graph_phase(self, used_cuda_graph: bool) -> None: + if self.is_warmup: + return + key = "encoder_graph" if used_cuda_graph else "encoder_eager" + self.cuda_graph_phase_stats[key] += 1 + + def _record_decoder_cuda_graph_phase(self, + scheduled_requests: ScheduledRequests, + used_cuda_graph: bool) -> None: + if self.is_warmup: + return + + suffix = "graph" if used_cuda_graph else "eager" + if scheduled_requests.num_context_requests > 0: + self.cuda_graph_phase_stats[f"decoder_context_{suffix}"] += 1 + if scheduled_requests.num_generation_requests > 0: + self.cuda_graph_phase_stats[f"decoder_generation_{suffix}"] += 1 + if (scheduled_requests.num_context_requests > 0 + and scheduled_requests.num_generation_requests > 0): + self.cuda_graph_phase_stats[f"decoder_mixed_{suffix}"] += 1 + @property def use_beam_search(self): return self.max_beam_width > 1 @@ -995,9 +1030,9 @@ def warmup(self, resource_manager: ResourceManager) -> None: self.cuda_graph_runner.padding_dummy_requests = {} if self._is_encoder_decoder_model(): - logger.info( - "Skipping warmup for encoder-decoder models; warmup dummy " - "requests do not carry encoder output state.") + AutoTuner.get() + with self.cuda_graph_runner.allow_capture(): + self._run_cuda_graph_warmup(resource_manager) return if self.mapping.cp_size > 1: @@ -1383,6 +1418,9 @@ def _run_capture_pass(force_non_greedy: bool, label: str) -> None: self._update_draft_inference_state_for_warmup( batch, draft_len > 0, resource_manager) self.runtime_draft_len = draft_len + if self._is_encoder_decoder_model(): + self._prepare_encoder_decoder_cuda_graph_warmup_batch( + batch, resource_manager) self.forward(batch, new_tensors_device=None, resource_manager=resource_manager) @@ -1481,6 +1519,8 @@ def _release_batch_context(self, batch: Optional[ScheduledRequests], self.kv_cache_manager_key) draft_kv_cache_manager = self._get_draft_kv_cache_manager( resource_manager) + cross_kv_cache_manager = resource_manager.get_resource_manager( + ResourceManagerType.CROSS_KV_CACHE_MANAGER) spec_resource_manager = resource_manager.get_resource_manager( ResourceManagerType.SPEC_RESOURCE_MANAGER) try: @@ -1491,6 +1531,8 @@ def _release_batch_context(self, batch: Optional[ScheduledRequests], kv_cache_manager.free_resources(req) if draft_kv_cache_manager is not None: draft_kv_cache_manager.free_resources(req) + if cross_kv_cache_manager is not None: + cross_kv_cache_manager.free_resources(req) if spec_resource_manager is not None: spec_resource_manager.free_resources(req) @@ -1652,15 +1694,26 @@ def _create_cuda_graph_warmup_request( result = ScheduledRequests() num_extra_decoding_steps = self._get_num_extra_decoding_steps() + is_encoder_decoder = self._is_encoder_decoder_model() + dummy_encoder_output_len = ( + self._get_encoder_decoder_dummy_encoder_output_len(resource_manager) + if is_encoder_decoder else None) # Add (batch_size - 1) dummy requests with seq_len=1. + short_token_nums = ([2] * + (batch_size - 1)) if is_encoder_decoder else None + short_encoder_output_lens = ( + [dummy_encoder_output_len] * + (batch_size - 1)) if is_encoder_decoder else None requests = kv_cache_manager.add_dummy_requests( list(range(batch_size - 1)), + token_nums=short_token_nums, is_gen=True, max_num_draft_tokens=draft_len, kv_reserve_draft_tokens=self.max_draft_loop_tokens, use_mrope=self.use_mrope, max_beam_width=self.max_beam_width, + encoder_output_lens=short_encoder_output_lens, num_extra_decoding_steps=num_extra_decoding_steps, draft_kv_cache_manager=draft_kv_cache_manager) @@ -1689,7 +1742,7 @@ def _create_cuda_graph_warmup_request( available_tokens = min(available_tokens, draft_available_tokens) token_num = max( - 1, + 2 if is_encoder_decoder else 1, min( available_tokens, max_seq_len - 1 - get_num_extra_kv_tokens(self.spec_config) - _kv_draft)) @@ -1714,6 +1767,8 @@ def _create_cuda_graph_warmup_request( kv_reserve_draft_tokens=self.max_draft_loop_tokens, use_mrope=self.use_mrope, max_beam_width=self.max_beam_width, + encoder_output_lens=[dummy_encoder_output_len] + if is_encoder_decoder else None, num_extra_decoding_steps=num_extra_decoding_steps, draft_kv_cache_manager=draft_kv_cache_manager) @@ -1732,8 +1787,155 @@ def _create_cuda_graph_warmup_request( 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._allocate_encoder_decoder_dummy_cross_kv( + result.generation_requests, resource_manager): + for request in result.generation_requests: + kv_cache_manager.free_resources(request) + if draft_kv_cache_manager is not None: + draft_kv_cache_manager.free_resources(request) + if spec_resource_manager is not None: + spec_resource_manager.free_resources(request) + return None return result + def _get_encoder_decoder_dummy_encoder_output_len( + self, resource_manager: ResourceManager) -> int: + cross_kv_cache_manager = resource_manager.get_resource_manager( + ResourceManagerType.CROSS_KV_CACHE_MANAGER) + max_encoder_output_len = int(self.max_seq_len) + if cross_kv_cache_manager is not None: + max_encoder_output_len = min( + max_encoder_output_len, + int( + getattr(cross_kv_cache_manager, "max_seq_len", + max_encoder_output_len))) + return max(1, max_encoder_output_len) + + def _allocate_encoder_decoder_dummy_cross_kv( + self, requests: List[LlmRequest], + resource_manager: ResourceManager) -> bool: + if not requests: + return True + cross_kv_cache_manager = resource_manager.get_resource_manager( + ResourceManagerType.CROSS_KV_CACHE_MANAGER) + if cross_kv_cache_manager is None: + raise RuntimeError("Encoder-decoder CUDA graph warmup requires " + "ResourceManagerType.CROSS_KV_CACHE_MANAGER.") + + encoder_output_len = self._get_encoder_decoder_dummy_encoder_output_len( + resource_manager) + for request in requests: + request.py_encoder_output = None + request.py_skip_cross_kv_projection = True + + encoder_output_lens = [encoder_output_len] * len(requests) + cross_dummy_requests = cross_kv_cache_manager.add_dummy_requests( + request_ids=[request.py_request_id for request in requests], + token_nums=encoder_output_lens, + is_gen=True, + max_beam_width=1, + encoder_output_lens=encoder_output_lens) + return cross_dummy_requests is not None + + def _prepare_encoder_decoder_cuda_graph_warmup_batch( + self, batch: ScheduledRequests, + resource_manager: ResourceManager) -> None: + if not batch.generation_requests: + return + + encoder_output_len = self._get_encoder_decoder_dummy_encoder_output_len( + resource_manager) + hidden_size = self._get_encoder_decoder_hidden_size() + saved_request_state = [] + for request in batch.generation_requests: + saved_request_state.append( + (request, request.py_encoder_output, + request.py_skip_cross_kv_projection, request.state, + request.py_batch_idx, request._cached_tokens, + request._cached_tokens_set)) + request.py_encoder_output = torch.ones( + (encoder_output_len, hidden_size), + device="cuda", + dtype=self.dtype) + request.py_skip_cross_kv_projection = False + request.state = LlmRequestState.CONTEXT_INIT + request.context_current_position = 0 + request.context_chunk_size = 1 + + projection_batch = ScheduledRequests() + projection_batch.reset_context_requests(batch.generation_requests) + kv_cache_manager = resource_manager.get_resource_manager( + self.kv_cache_manager_key) + draft_kv_cache_manager = self._get_draft_kv_cache_manager( + resource_manager) + attn_metadata = self._set_up_attn_metadata(kv_cache_manager, + draft_kv_cache_manager) + with self.no_cuda_graph(): + projection_inputs, _ = self._prepare_inputs( + projection_batch, + kv_cache_manager, + attn_metadata, + spec_metadata=None, + new_tensors_device=None, + resource_manager=resource_manager, + maybe_graph=False) + self._project_encoder_decoder_warmup_cross_kv(projection_inputs) + torch.cuda.synchronize() + + for (request, encoder_output, skip_cross_kv_projection, state, + batch_idx, cached_tokens, + cached_tokens_set) in saved_request_state: + request.py_encoder_output = encoder_output + request.py_skip_cross_kv_projection = skip_cross_kv_projection + request.state = state + if state == LlmRequestState.GENERATION_IN_PROGRESS: + request.context_current_position = request.prompt_len + request.py_batch_idx = batch_idx + request._cached_tokens = cached_tokens + request._cached_tokens_set = cached_tokens_set + + def _project_encoder_decoder_warmup_cross_kv( + self, inputs: Dict[str, Any]) -> None: + encoder_hidden_states = inputs.get("encoder_hidden_states") + cross_attn_metadata = inputs.get("cross_attn_metadata") + if encoder_hidden_states is None or cross_attn_metadata is None: + return + + decoder = getattr(self._get_top_level_model(), "decoder", None) + layers = getattr(decoder, "layers", None) + if layers is None: + raise RuntimeError("Encoder-decoder CUDA graph warmup requires a " + "decoder with cross-attention layers.") + + attn_metadata = inputs["attn_metadata"] + hidden_states = torch.ones( + (attn_metadata.num_tokens, self._get_encoder_decoder_hidden_size()), + device=encoder_hidden_states.device, + dtype=encoder_hidden_states.dtype) + for layer in layers: + cross_attn = getattr(layer, "cross_attn", None) + if cross_attn is None: + raise RuntimeError( + "Encoder-decoder CUDA graph warmup requires every decoder " + "layer to expose a cross_attn module.") + cross_attn(hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + attn_metadata=attn_metadata, + cross_attn_metadata=cross_attn_metadata, + skip_cross_kv_projection=False) + + def _get_encoder_decoder_hidden_size(self) -> int: + config = self.model.model_config.pretrained_config + hidden_size = getattr(config, "hidden_size", None) + if hidden_size is None: + hidden_size = getattr(config, "d_model", None) + if hidden_size is None: + raise RuntimeError( + "Encoder-decoder CUDA graph warmup could not infer encoder " + "hidden size from the model config.") + return int(hidden_size) + def _update_draft_inference_state_for_warmup( self, batch: ScheduledRequests, is_first_draft: bool, resource_manager: ResourceManager): @@ -2323,51 +2525,17 @@ def _prepare_encoder_decoder_cross_attention_inputs( packed_encoder_hidden_states = None skip_cross_kv_projection = True - encoder_seq_lens_tensor = torch.tensor(encoder_seq_lens, - dtype=torch.int, - pin_memory=prefer_pinned()) - - def update_cross_metadata( - cross_attn_metadata: AttentionMetadata) -> AttentionMetadata: - base_params = attn_metadata.kv_cache_params - cross_attn_metadata.kv_cache_manager = cross_kv_cache_manager - cross_attn_metadata._seq_lens = attn_metadata.seq_lens - cross_attn_metadata._seq_lens_cuda = attn_metadata.seq_lens_cuda - cross_attn_metadata.cross = cross_attn_metadata - cross_attn_metadata.seq_lens_kv = encoder_seq_lens_tensor - if encoder_num_cached_tokens_per_seq is not None: - use_cache = (base_params.use_cache if base_params is not None - else (cross_kv_cache_manager is not None)) - block_ids_per_seq = (base_params.block_ids_per_seq - if base_params is not None else None) - host_max_attention_window_sizes = ( - base_params.host_max_attention_window_sizes - if base_params is not None else None) - host_sink_token_length = (base_params.host_sink_token_length - if base_params is not None else None) - num_extra_kv_tokens = (base_params.num_extra_kv_tokens - if base_params is not None else 0) - cross_attn_metadata.kv_cache_params = KVCacheParams( - use_cache=use_cache, - num_cached_tokens_per_seq=list( - encoder_num_cached_tokens_per_seq), - block_ids_per_seq=block_ids_per_seq, - host_max_attention_window_sizes= - host_max_attention_window_sizes, - host_sink_token_length=host_sink_token_length, - num_extra_kv_tokens=num_extra_kv_tokens, - ) - cross_attn_metadata.request_ids = attn_metadata.request_ids - cross_attn_metadata.prompt_lens = attn_metadata.prompt_lens - cross_attn_metadata.num_contexts = attn_metadata.num_contexts - return cross_attn_metadata - if attn_metadata.is_cuda_graph and attn_metadata.has_cross_sub_metadata: - cross_attn_metadata = update_cross_metadata(attn_metadata.cross) + cross_attn_metadata = attn_metadata.update_cross_metadata( + encoder_seq_lens=encoder_seq_lens, + cross_kv_cache_manager=cross_kv_cache_manager, + encoder_num_cached_tokens_per_seq= + encoder_num_cached_tokens_per_seq, + ) else: cross_attn_metadata = attn_metadata.create_cross_metadata( - encoder_seq_lens=encoder_seq_lens_tensor, cross_kv_cache_manager=cross_kv_cache_manager, + encoder_seq_lens=encoder_seq_lens, encoder_num_cached_tokens_per_seq= encoder_num_cached_tokens_per_seq, ) @@ -3013,7 +3181,9 @@ def append_cross_attention_state(request: LlmRequest, append_cross_attention_state( request, project_encoder_output=not request.py_skip_cross_kv_projection - and not getattr(request, "is_dummy", False)) + and + (not getattr(request, "is_dummy", False) + or getattr(request, "py_encoder_output", None) is not None)) # Embed mask is required only for partial iterations (chunked # prefill or KV-cache reuse); full-prefill degrades gracefully. @@ -4785,6 +4955,7 @@ def encoder_forward(self, inputs: Dict[str, Any], ) if self.encoder_attn_metadata is None else self.encoder_attn_metadata graph_attn_metadata, key = self.encoder_cuda_graph_runner.maybe_get_cuda_graph( padded_inputs, attn_metadata) + self._record_encoder_cuda_graph_phase(key is not None) # Unpad seq_lens when fallback to eager path. if key is None: padded_inputs['seq_lens'] = padded_inputs[ @@ -4958,6 +5129,8 @@ def forward(self, ) can_run_graph = key is not None + self._record_decoder_cuda_graph_phase(scheduled_requests, + can_run_graph) if can_run_graph: attn_metadata = maybe_attn_metadata spec_metadata = maybe_spec_metadata @@ -5377,6 +5550,7 @@ def forward_encoder( with torch.inference_mode(): inputs = self._prepare_tp_inputs_encoder( encoder_requests, resource_manager=resource_manager) + self._record_encoder_cuda_graph_phase(False) encoder_hidden_states = self._forward_step_encoder(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 904964acb8a3..1e48d3d0e89b 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -544,6 +544,38 @@ def __init__( # The scheduler will avoid scheduling requests that are already in flight. self.inflight_req_ids = ReqIdsSet() + # Encoder-decoder models execute the encoder and decoder in separate + # iterations. The encoder branch lives in ``_executor_loop`` only; + # ``_executor_loop_overlap`` has not been threaded yet. Reject + # pp_size > 1 for parity with the legacy TRT path (Encoder PP support + # is intentionally out of scope for this port). + is_encoder_decoder = bool( + getattr(getattr(self.model_engine.model, "model_config", None), + "is_encoder_decoder", False)) + if is_encoder_decoder: + if self.dist.pp_size > 1: + raise NotImplementedError( + "pp_size > 1 is not supported for encoder-decoder models " + "in the PyTorch flow; encoder send/recv hooks are out of " + "scope. Set pp_size=1 to run T5/BART/mBART.") + if not self.disable_overlap_scheduler: + raise NotImplementedError( + "Overlap scheduler is not yet wired for encoder-decoder " + "models. Set disable_overlap_scheduler=True for " + "encoder-decoder runs.") + if getattr(self.model_engine, "_torch_compile_piecewise_cuda_graph", + False): + raise NotImplementedError( + "Piecewise CUDA graph is not supported for " + "encoder-decoder models. Disable " + "torch_compile_config.enable_piecewise_cuda_graph.") + if (getattr(self.model_engine, "cuda_graph_config", None) + is not None and getattr(self.model_engine, "spec_config", + None) is not None): + raise NotImplementedError( + "Speculative decoding with CUDA graph is not supported " + "for encoder-decoder models.") + # Synchronize all ranks before warmup. This prevents a PP # communication deadlock when ranks on the last PP stage are delayed # by heavy initialisation (e.g. guided-decoder / llguidance tokenizer @@ -678,31 +710,6 @@ def on_detected(): self._disagg_pp_termination_handler = DisaggPPTerminationHandler( self.dist, self._do_terminate_request) - # Encoder-decoder models execute the encoder and decoder in separate - # iterations. The encoder branch lives in ``_executor_loop`` only; - # ``_executor_loop_overlap`` has not been threaded yet. Reject - # pp_size > 1 for parity with the legacy TRT path (Encoder PP support - # is intentionally out of scope for this port). - is_encoder_decoder = bool( - getattr(getattr(self.model_engine.model, "model_config", None), - "is_encoder_decoder", False)) - if is_encoder_decoder: - if self.dist.pp_size > 1: - raise NotImplementedError( - "pp_size > 1 is not supported for encoder-decoder models " - "in the PyTorch flow; encoder send/recv hooks are out of " - "scope. Set pp_size=1 to run T5/BART/mBART.") - if not self.disable_overlap_scheduler: - raise NotImplementedError( - "Overlap scheduler is not yet wired for encoder-decoder " - "models. Set disable_overlap_scheduler=True for " - "encoder-decoder runs.") - if getattr(self.model_engine, "cuda_graph_config", - None) is not None: - raise NotImplementedError( - "CUDA graph is not supported for encoder-decoder models. " - "Disable cuda_graph_config for encoder-decoder runs.") - if self.dist.pp_size > 1: self.event_loop = self._executor_loop_pp # `TLLM_PP_ASYNC_BROADCAST_SAMPLE_STATE` controls whether to broadcast the sample state asynchronously. diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index e4ad48050606..e6d26f02568a 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -821,6 +821,7 @@ def add_dummy_requests( kv_reserve_draft_tokens: Optional[int] = None, use_mrope: bool = False, max_beam_width: int = 1, + encoder_output_lens: Optional[List[int]] = None, # For capturable drafting loops. During normal inference, the draft model always # has enough KV cache space to fit all of our draft tokens. During warmup, however, # we need to make the KV cache manager aware that multiple autoregressive steps will @@ -854,9 +855,10 @@ def add_dummy_requests( # in _prepare_tp_inputs; need token_num >= 2 so that doesn't go negative. if self.mapping.has_cp_helix(): token_num = max(token_num, 2) - encoder_input_tokens = [ - 1 - ] * token_num if self.impl.cross_kv else None + encoder_output_len = (encoder_output_lens[i] + if encoder_output_lens is not None else None) + encoder_input_tokens = ([1] * encoder_output_len + if encoder_output_len is not None else None) # Using 1 instead of 0 prevents NaN during warmup in e.g. Deepseek req = LlmRequest(request_id=req_id, max_new_tokens=1, @@ -864,7 +866,8 @@ def add_dummy_requests( sampling_config=SamplingConfig( sampling_params._get_sampling_config()), is_streaming=False, - encoder_input_tokens=encoder_input_tokens) + encoder_input_tokens=encoder_input_tokens, + encoder_output_len=encoder_output_len) req.is_dummy_request = True req.paged_kv_block_ids = [] if prepare_resource: 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 3ae0e3652a3d..be54838c739b 100644 --- a/tests/integration/defs/llmapi/test_llm_api_pytorch_bart.py +++ b/tests/integration/defs/llmapi/test_llm_api_pytorch_bart.py @@ -63,6 +63,8 @@ def _test_case( num_return_sequences: int, exact_match: bool, feature_id: str, + cuda_graph_batch_sizes: list[int] | None = None, + kv_cache_dtype: str = "auto", ): expected_output_token_ids = [_EXPECTED_GREEDY_OUTPUT_TOKEN_IDS] if num_beams == 1 else None assert not exact_match or expected_output_token_ids is not None @@ -75,6 +77,8 @@ def _test_case( num_beams, num_return_sequences, exact_match, + cuda_graph_batch_sizes, + kv_cache_dtype, id=f"{feature_id}-{_MODEL_NAME}", ) @@ -89,6 +93,16 @@ def _test_case( exact_match=True, feature_id="bf16-kv-v1-cuda-graph-off-greedy", ), + _test_case( + torch_dtype="bfloat16", + use_kv_cache_manager_v2=False, + enable_cuda_graph=True, + num_beams=1, + num_return_sequences=1, + exact_match=True, + cuda_graph_batch_sizes=[2], + feature_id="bf16-kv-v1-cuda-graph-on-greedy", + ), _test_case( torch_dtype="float16", use_kv_cache_manager_v2=False, @@ -107,6 +121,15 @@ def _test_case( exact_match=False, feature_id="bf16-kv-v1-cuda-graph-off-beam2", ), + _test_case( + torch_dtype="bfloat16", + use_kv_cache_manager_v2=False, + enable_cuda_graph=True, + num_beams=2, + num_return_sequences=2, + exact_match=False, + feature_id="bf16-kv-v1-cuda-graph-on-beam2", + ), _test_case( torch_dtype="bfloat16", use_kv_cache_manager_v2=True, @@ -116,6 +139,15 @@ def _test_case( exact_match=True, feature_id="bf16-kv-v2-cuda-graph-off-greedy", ), + _test_case( + torch_dtype="bfloat16", + use_kv_cache_manager_v2=True, + enable_cuda_graph=True, + num_beams=1, + num_return_sequences=1, + exact_match=True, + feature_id="bf16-kv-v2-cuda-graph-on-greedy", + ), ] @@ -192,7 +224,31 @@ def _cuda_graph_config( enabled: bool, batch_sizes: list[int] | None = None, ) -> CudaGraphConfig | None: - return CudaGraphConfig(batch_sizes=batch_sizes or [1]) if enabled else None + if not enabled: + return None + return CudaGraphConfig(batch_sizes=batch_sizes or [1], enable_padding=True) + + +def _assert_cuda_graphs_captured(llm: LLM, enabled: bool) -> None: + if not enabled: + return + model_engine = llm._executor.engine.model_engine + assert model_engine.cuda_graph_runner.graphs + + +def _assert_cuda_graph_padding_used(llm: LLM, batch_sizes: list[int] | None) -> None: + if batch_sizes is None: + return + model_engine = llm._executor.engine.model_engine + assert model_engine.cuda_graph_runner.padding_dummy_requests + + +def _enable_trtllm_gen_attention(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("TRTLLM_ENABLE_TRTLLM_GEN_ATTENTION", "1") + + from tensorrt_llm._torch.attention_backend import trtllm + + monkeypatch.setattr(trtllm, "_TRTLLM_ENABLE_TRTLLM_GEN_ATTENTION", True) def _assert_bart_response( @@ -250,7 +306,7 @@ def _assert_expected_generation( "enable_cuda_graph,num_beams,num_return_sequences,exact_match", _TEST_CASES, ) -def test_bart_pytorch_generate_encoder_decoder_end_to_end( +def _run_bart_pytorch_generate_encoder_decoder( monkeypatch: pytest.MonkeyPatch, expected_output_token_ids_by_output: list[list[int]] | None, torch_dtype: str, @@ -259,6 +315,8 @@ def test_bart_pytorch_generate_encoder_decoder_end_to_end( num_beams: int, num_return_sequences: int, exact_match: bool, + cuda_graph_batch_sizes: list[int] | None, + kv_cache_dtype: str = "auto", ) -> None: monkeypatch.setenv("TLLM_WORKER_USE_SINGLE_PROCESS", "1") monkeypatch.setenv("TRTLLM_SKIP_KV_CACHE_ESTIMATION", "1") @@ -275,7 +333,7 @@ def test_bart_pytorch_generate_encoder_decoder_end_to_end( model_path, backend="pytorch", attn_backend="TRTLLM", - cuda_graph_config=_cuda_graph_config(enable_cuda_graph), + cuda_graph_config=_cuda_graph_config(enable_cuda_graph, cuda_graph_batch_sizes), disable_overlap_scheduler=True, dtype=torch_dtype, enable_chunked_prefill=False, @@ -286,7 +344,7 @@ def test_bart_pytorch_generate_encoder_decoder_end_to_end( cross_kv_cache_fraction=_CROSS_KV_CACHE_FRACTION, use_kv_cache_manager_v2=use_kv_cache_manager_v2, ), - max_batch_size=1, + max_batch_size=max(cuda_graph_batch_sizes or [1]), max_beam_width=num_beams, max_input_len=_MAX_SEQUENCE_LENGTH, max_num_tokens=_MAX_SEQUENCE_LENGTH, @@ -310,6 +368,40 @@ def test_bart_pytorch_generate_encoder_decoder_end_to_end( exact_match, expected_output_token_ids_by_output, ) + _assert_cuda_graphs_captured(llm, enable_cuda_graph) + _assert_cuda_graph_padding_used(llm, cuda_graph_batch_sizes) + + +@pytest.mark.parametrize( + "expected_output_token_ids_by_output,torch_dtype,use_kv_cache_manager_v2," + "enable_cuda_graph,num_beams,num_return_sequences,exact_match,cuda_graph_batch_sizes," + "kv_cache_dtype", + _TEST_CASES, +) +def test_bart_pytorch_generate_encoder_decoder_end_to_end( + monkeypatch: pytest.MonkeyPatch, + expected_output_token_ids_by_output: list[list[int]] | None, + torch_dtype: str, + use_kv_cache_manager_v2: bool, + enable_cuda_graph: bool, + num_beams: int, + num_return_sequences: int, + exact_match: bool, + cuda_graph_batch_sizes: list[int] | None, + kv_cache_dtype: str, +) -> None: + _run_bart_pytorch_generate_encoder_decoder( + monkeypatch, + expected_output_token_ids_by_output, + torch_dtype, + use_kv_cache_manager_v2, + enable_cuda_graph, + num_beams, + num_return_sequences, + exact_match, + cuda_graph_batch_sizes, + kv_cache_dtype, + ) @pytest.mark.parametrize( 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 c42bb0a40057..4601d88a1a9b 100644 --- a/tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py +++ b/tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py @@ -106,7 +106,8 @@ def _test_case( num_return_sequences: int, exact_match: bool, feature_id: str, - marks=(), + cuda_graph_batch_sizes: list[int] | None = None, + marks=None, ): if num_beams == 1: expected_output_token_ids = ( @@ -127,6 +128,10 @@ def _test_case( assert not exact_match or expected_output_token_ids is not None + param_kwargs = {"id": f"{feature_id}-{model_name}"} + if marks is not None: + param_kwargs["marks"] = marks + return pytest.param( model_name, expected_output_token_ids, @@ -136,8 +141,8 @@ def _test_case( num_beams, num_return_sequences, exact_match, - id=f"{feature_id}-{model_name}", - marks=marks, + cuda_graph_batch_sizes, + **param_kwargs, ) @@ -246,6 +251,28 @@ def _test_case( exact_match=True, feature_id="bf16-kv-v1-cuda-graph-off-greedy", ), + _test_case( + model_name="t5-small", + torch_dtype="bfloat16", + use_kv_cache_manager_v2=False, + enable_cuda_graph=True, + num_beams=1, + num_return_sequences=1, + exact_match=True, + cuda_graph_batch_sizes=[2], + feature_id="bf16-kv-v1-cuda-graph-on-greedy", + ), + _test_case( + model_name="t5-small", + torch_dtype="bfloat16", + use_kv_cache_manager_v2=False, + enable_cuda_graph=True, + num_beams=2, + num_return_sequences=2, + exact_match=False, + cuda_graph_batch_sizes=[2], + feature_id="bf16-kv-v1-cuda-graph-on-beam2", + ), # Precision coverage for beam search. KVCacheManagerV2 currently requires # max_beam_width == 1, so beam-search precision coverage uses v1. _test_case( @@ -299,6 +326,17 @@ def _test_case( exact_match=True, feature_id="bf16-kv-v2-cuda-graph-off-greedy", ), + _test_case( + model_name="t5-small", + torch_dtype="bfloat16", + use_kv_cache_manager_v2=True, + enable_cuda_graph=True, + num_beams=1, + num_return_sequences=1, + exact_match=True, + cuda_graph_batch_sizes=[2], + feature_id="bf16-kv-v2-cuda-graph-on-greedy", + ), _test_case( model_name="t5-small", torch_dtype="float16", @@ -524,12 +562,7 @@ def _assert_expected_generation( assert token_ids_by_output == expected_token_ids_by_output -@pytest.mark.parametrize( - "model_name,expected_output_token_ids_by_output,torch_dtype,use_kv_cache_manager_v2," - "enable_cuda_graph,num_beams,num_return_sequences,exact_match", - _TEST_CASES, -) -def test_t5_pytorch_generate_encoder_decoder_end_to_end( +def _run_t5_pytorch_generate_encoder_decoder( monkeypatch: pytest.MonkeyPatch, model_name: str, expected_output_token_ids_by_output: list[list[int]] | None, @@ -539,6 +572,7 @@ def test_t5_pytorch_generate_encoder_decoder_end_to_end( num_beams: int, num_return_sequences: int, exact_match: bool, + cuda_graph_batch_sizes: list[int] | None, ) -> None: monkeypatch.setenv("TLLM_WORKER_USE_SINGLE_PROCESS", "1") monkeypatch.setenv("TRTLLM_SKIP_KV_CACHE_ESTIMATION", "1") @@ -555,7 +589,7 @@ def test_t5_pytorch_generate_encoder_decoder_end_to_end( model_path, backend="pytorch", attn_backend="TRTLLM", - cuda_graph_config=_cuda_graph_config(enable_cuda_graph), + cuda_graph_config=_cuda_graph_config(enable_cuda_graph, cuda_graph_batch_sizes), disable_overlap_scheduler=True, dtype=torch_dtype, enable_chunked_prefill=False, @@ -566,7 +600,7 @@ def test_t5_pytorch_generate_encoder_decoder_end_to_end( cross_kv_cache_fraction=_CROSS_KV_CACHE_FRACTION, use_kv_cache_manager_v2=use_kv_cache_manager_v2, ), - max_batch_size=1, + max_batch_size=max(cuda_graph_batch_sizes or [1]), max_beam_width=num_beams, max_input_len=_MAX_SEQUENCE_LENGTH, max_num_tokens=_MAX_SEQUENCE_LENGTH, @@ -592,6 +626,37 @@ def test_t5_pytorch_generate_encoder_decoder_end_to_end( ) +@pytest.mark.parametrize( + "model_name,expected_output_token_ids_by_output,torch_dtype,use_kv_cache_manager_v2," + "enable_cuda_graph,num_beams,num_return_sequences,exact_match,cuda_graph_batch_sizes", + _TEST_CASES, +) +def test_t5_pytorch_generate_encoder_decoder_end_to_end( + monkeypatch: pytest.MonkeyPatch, + model_name: str, + expected_output_token_ids_by_output: list[list[int]] | None, + torch_dtype: str, + use_kv_cache_manager_v2: bool, + enable_cuda_graph: bool, + num_beams: int, + num_return_sequences: int, + exact_match: bool, + cuda_graph_batch_sizes: list[int] | None, +) -> None: + _run_t5_pytorch_generate_encoder_decoder( + monkeypatch, + model_name, + expected_output_token_ids_by_output, + torch_dtype, + use_kv_cache_manager_v2, + enable_cuda_graph, + num_beams, + num_return_sequences, + exact_match, + cuda_graph_batch_sizes, + ) + + @pytest.mark.parametrize( "model_name,expected_output_token_ids_by_request,torch_dtype,use_kv_cache_manager_v2," "num_beams,num_return_sequences,exact_match", diff --git a/tests/unittest/_torch/helpers.py b/tests/unittest/_torch/helpers.py index 3f7cfad7c0b5..7f167077effa 100644 --- a/tests/unittest/_torch/helpers.py +++ b/tests/unittest/_torch/helpers.py @@ -252,6 +252,7 @@ def create_mock_cuda_graph_runner(batch_size: int, use_mrope: bool = False): original_max_draft_len=0, original_max_total_draft_tokens=0, is_draft_model=False, + is_encoder_decoder=False, mapping=Mapping(), dist=None, kv_cache_manager_key=ResourceManagerType.KV_CACHE_MANAGER) From f228561492eea81c5e9d3edac2e25e6a856c371d Mon Sep 17 00:00:00 2001 From: Guiju Zhang <7135567+cascade812@users.noreply.github.com> Date: Thu, 25 Jun 2026 13:06:46 -0700 Subject: [PATCH 2/9] enable cuda graph for encoder pass Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com> --- .../_torch/pyexecutor/cuda_graph_runner.py | 8 +- .../_torch/pyexecutor/model_engine.py | 488 ++++++++++++------ .../defs/llmapi/test_llm_api_pytorch_bart.py | 48 +- .../defs/llmapi/test_llm_api_pytorch_t5.py | 130 ++++- 4 files changed, 490 insertions(+), 184 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index 883a82c677a5..a9b37615fe40 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -658,11 +658,11 @@ class EncoderCUDAGraphRunnerConfig: class EncoderCUDAGraphRunner: - """CUDA graph runner for models using encode_only path. + """CUDA graph runner for no-cache encoder forward passes. - Designed for the `LLM.encode()` API — consumes raw inputs dicts with - `input_ids` (flat [total_tokens]), `seq_lens` ([batch_size]). Encoder CUDA graphs - are keyed on the 3-tuple (padded_batch_size, padded_num_tokens, padded_max_seq_len) + 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). Restricted to `TrtllmAttentionMetadata` — FlashInfer's per-batch planner state is not compatible with CUDA graph capture/replay. """ diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index b8b60ec82894..8e187738a509 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -383,7 +383,8 @@ def __init__( self._is_encode_only = (self.llm_args.encode_only and not self.llm_args.mm_encoder_only) - if (self._is_encode_only and self.cuda_graph_config is not None + if ((self._is_encode_only or self._is_encoder_decoder_model()) + and isinstance(self.cuda_graph_config, EncodeCudaGraphConfig) and (not cuda_graph_num_tokens or not cuda_graph_seq_lens)): missing = [] if not cuda_graph_num_tokens: @@ -391,7 +392,7 @@ def __init__( if not cuda_graph_seq_lens: missing.append("seq_lens/max_seq_len") logger.warning( - f"encode_only=True with a CudaGraphConfig, but " + f"Encoder CUDA graph config is set, but " f"{' and '.join(missing)} not set. Encoder CUDA graphs " f"require both. Encoder CUDA graphs will be disabled. " f"To enable them, specify e.g. " @@ -553,16 +554,6 @@ def __init__( self.encoder_attn_metadata = None self.spec_metadata = None self.iter_states = {} - self.cuda_graph_phase_stats = { - "encoder_eager": 0, - "encoder_graph": 0, - "decoder_context_eager": 0, - "decoder_context_graph": 0, - "decoder_generation_eager": 0, - "decoder_generation_graph": 0, - "decoder_mixed_eager": 0, - "decoder_mixed_graph": 0, - } 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 @@ -589,6 +580,12 @@ def __init__( self._max_cuda_graph_seq_len = (self._cuda_graph_seq_lens[-1] if self._cuda_graph_seq_lens else 0) + self._enable_encoder_cuda_graph = ( + (self._is_encode_only or self._is_encoder_decoder_model()) + and self.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( ) @@ -680,10 +677,7 @@ def __init__( # Create Encoder CUDA graph config and runner. 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)), + use_cuda_graph=self._enable_encoder_cuda_graph, cuda_graph_padding_enabled=self._cuda_graph_padding_enabled, cuda_graph_batch_sizes=self._cuda_graph_batch_sizes, cuda_graph_num_tokens=self._cuda_graph_num_tokens, @@ -810,30 +804,6 @@ def moe_load_balancer_iter_info(self, value: Tuple[bool, bool]): moe_load_balancer.set_iter_info(enable_statistic=value[0], enable_update_weights=value[1]) - def get_cuda_graph_phase_stats(self) -> Dict[str, int]: - return dict(self.cuda_graph_phase_stats) - - def _record_encoder_cuda_graph_phase(self, used_cuda_graph: bool) -> None: - if self.is_warmup: - return - key = "encoder_graph" if used_cuda_graph else "encoder_eager" - self.cuda_graph_phase_stats[key] += 1 - - def _record_decoder_cuda_graph_phase(self, - scheduled_requests: ScheduledRequests, - used_cuda_graph: bool) -> None: - if self.is_warmup: - return - - suffix = "graph" if used_cuda_graph else "eager" - if scheduled_requests.num_context_requests > 0: - self.cuda_graph_phase_stats[f"decoder_context_{suffix}"] += 1 - if scheduled_requests.num_generation_requests > 0: - self.cuda_graph_phase_stats[f"decoder_generation_{suffix}"] += 1 - if (scheduled_requests.num_context_requests > 0 - and scheduled_requests.num_generation_requests > 0): - self.cuda_graph_phase_stats[f"decoder_mixed_{suffix}"] += 1 - @property def use_beam_search(self): return self.max_beam_width > 1 @@ -1031,6 +1001,8 @@ def warmup(self, resource_manager: ResourceManager) -> None: if self._is_encoder_decoder_model(): AutoTuner.get() + with self.encoder_cuda_graph_runner.allow_capture(): + self._run_cuda_graph_warmup_encoder_decoder() with self.cuda_graph_runner.allow_capture(): self._run_cuda_graph_warmup(resource_manager) return @@ -4786,6 +4758,28 @@ def _create_encoder_warmup_inputs( } return inputs + def _create_encoder_decoder_encoder_warmup_inputs( + self, batch_size: int, num_tokens: int, + max_seq_len: int) -> Optional[Dict[str, Any]]: + raw_inputs = self._create_encoder_warmup_inputs(batch_size, num_tokens, + max_seq_len) + if raw_inputs is None: + return None + + sequence_lengths = raw_inputs['seq_lens'] + encoder_position_ids: List[int] = [] + for seq_len in sequence_lengths: + encoder_position_ids.extend( + self._apply_position_id_offset(list(range(seq_len)))) + + return self._prepare_encoder_decoder_encoder_inputs( + encoder_input_ids=raw_inputs['input_ids'], + encoder_position_ids=encoder_position_ids, + sequence_lengths=sequence_lengths, + request_ids=list(range(batch_size)), + resource_manager=None, + ) + @contextlib.contextmanager def no_encoder_cuda_graph(self): """Temporarily disable the encoder CUDA graph runner.""" @@ -4881,25 +4875,21 @@ def _run_cuda_graph_warmup_encoder(self) -> None: self._capture_encoder_cuda_graphs() - def _capture_encoder_cuda_graphs(self) -> None: - """Capture whole-model encoder CUDA graphs for all feasible keys. - - Feasibility filter (also used in source): - nt >= prev_sl + bs (enough tokens for this sl bucket) - prev_nt < bs * sl (not enough tokens for a smaller nt bucket) - nt <= bs * sl (num tokens should not exceed total possible in batch) - sl <= nt (seq len should not exceed num tokens) - """ - runner = self.encoder_cuda_graph_runner - if not runner.enabled: + def _run_cuda_graph_warmup_encoder_decoder(self) -> None: + """Captures encoder CUDA graphs for encoder-decoder models.""" + if not self.encoder_cuda_graph_runner.enabled: return + self._capture_encoder_decoder_encoder_cuda_graphs() + + def _get_encoder_cuda_graph_warmup_configs( + self) -> List[Tuple[int, int, int]]: + """Return feasible (batch_size, num_tokens, max_seq_len) graph keys.""" batch_sizes = sorted(self._cuda_graph_batch_sizes, reverse=True) num_tokens_list = sorted(self._cuda_graph_num_tokens) seq_lens_list = sorted(self._cuda_graph_seq_lens) - num_captured = 0 - logger.info("Capturing encoder CUDA graphs ...") + warmup_configs: List[Tuple[int, int, int]] = [] for bs in batch_sizes: if bs > self.batch_size: continue @@ -4914,92 +4904,162 @@ def _capture_encoder_cuda_graphs(self) -> None: if nt > bs * sl or sl > nt: continue - inputs = self._create_encoder_warmup_inputs(bs, nt, sl) - if inputs is None: - continue + warmup_configs.append((bs, nt, sl)) - logger.info(f"Encoder CUDA graph capture: " - f"bs={bs}, nt={nt}, sl={sl}") - self.encoder_forward(inputs) - torch.cuda.synchronize() - num_captured += 1 + return warmup_configs + + def _capture_encoder_cuda_graphs(self) -> None: + """Capture whole-model encoder CUDA graphs for all feasible keys. + + Feasibility filter (also used in source): + nt >= prev_sl + bs (enough tokens for this sl bucket) + prev_nt < bs * sl (not enough tokens for a smaller nt bucket) + nt <= bs * sl (num tokens should not exceed total possible in batch) + sl <= nt (seq len should not exceed num tokens) + """ + runner = self.encoder_cuda_graph_runner + if not runner.enabled: + return + + num_captured = 0 + logger.info("Capturing encoder CUDA graphs ...") + for bs, nt, sl in self._get_encoder_cuda_graph_warmup_configs(): + inputs = self._create_encoder_warmup_inputs(bs, nt, sl) + if inputs is None: + continue + + logger.info(f"Encoder CUDA graph capture: " + f"bs={bs}, nt={nt}, sl={sl}") + self.encoder_forward(inputs) + torch.cuda.synchronize() + num_captured += 1 logger.info(f"Captured {num_captured} encoder CUDA graph(s).") - @torch.inference_mode() - @with_model_extra_attrs(lambda self: self.model.extra_attrs) - @nvtx_range("encoder_forward") - def encoder_forward(self, inputs: Dict[str, Any], - **kwargs) -> Dict[str, Any]: - """Direct tensor-level forward for encode-only path. + def _capture_encoder_decoder_encoder_cuda_graphs(self) -> None: + """Capture encoder-stack CUDA graphs for encoder-decoder models.""" + runner = self.encoder_cuda_graph_runner + if not runner.enabled: + return - Bypasses ScheduledRequests/LlmRequest entirely. Takes a raw inputs - dict, attempts encoder CUDA graph capture/replay if enabled, otherwise falls - back to eager execution. + num_captured = 0 + logger.info("Capturing encoder-decoder encoder CUDA graphs ...") + for bs, nt, sl in self._get_encoder_cuda_graph_warmup_configs(): + inputs = self._create_encoder_decoder_encoder_warmup_inputs( + bs, nt, sl) + if inputs is None: + continue - Args: - inputs: Dict with 'input_ids' and 'seq_lens' (required), plus - any model-specific kwargs (token_type_ids, inputs_embeds, etc.). + logger.info(f"Encoder-decoder encoder CUDA graph capture: " + f"bs={bs}, nt={nt}, sl={sl}") + self._forward_encoder_with_cuda_graph(inputs) + torch.cuda.synchronize() + num_captured += 1 - Returns: - Dict with 'logits' tensor and any other model outputs. - """ + logger.info( + f"Captured {num_captured} encoder-decoder encoder CUDA graph(s).") + + def _run_encoder_forward_with_cuda_graph( + self, + runner_inputs: Dict[str, Any], + attn_metadata: Any, + prepare_model_inputs: Callable[ + [Dict[str, Any], Optional[Any], Optional[Any]], Dict[str, Any]], + forward_fn: Callable[[Dict[str, Any]], Any], + postprocess_graph_outputs: Callable[[Any], Any], + ) -> Any: + """Run a no-cache encoder forward with shared CUDA graph handling.""" moe_load_balancer: MoeLoadBalancer = getattr(self, 'moe_load_balancer', None) - batch_size = len(inputs['seq_lens']) + batch_size = len(runner_inputs['seq_lens']) with self.encoder_cuda_graph_runner.pad_batch( - inputs, batch_size) as padded_inputs: - attn_metadata = self._set_up_attn_metadata( - kv_cache_manager=None - ) if self.encoder_attn_metadata is None else self.encoder_attn_metadata + runner_inputs, batch_size) as padded_inputs: graph_attn_metadata, key = self.encoder_cuda_graph_runner.maybe_get_cuda_graph( padded_inputs, attn_metadata) - self._record_encoder_cuda_graph_phase(key is not None) - # Unpad seq_lens when fallback to eager path. - if key is None: - padded_inputs['seq_lens'] = padded_inputs[ - 'seq_lens'][:batch_size] - model_inputs = self._prepare_encoder_inputs( - padded_inputs, - attn_metadata=graph_attn_metadata, - padded_num_tokens=key[1] if key is not None else None) - forward_kwargs = { - "gather_ids": None, - "gather_context_logits": False, - **kwargs, - } with with_shared_pool( self.encoder_cuda_graph_runner.get_graph_pool()): if key is None: + eager_inputs = dict(padded_inputs) + eager_inputs['seq_lens'] = eager_inputs[ + 'seq_lens'][:batch_size] + model_inputs = prepare_model_inputs(eager_inputs, None, + None) with MoeLoadBalancerIterContext(moe_load_balancer): - # Eager path — no graph for this bucket. - return self._forward_step(model_inputs, - **forward_kwargs) + return forward_fn(model_inputs) + model_inputs = prepare_model_inputs(padded_inputs, + graph_attn_metadata, key) if self.encoder_cuda_graph_runner.needs_capture(key): - def forward_fn( - capture_inputs: Dict[str, Any]) -> Dict[str, Any]: - capture_inputs = capture_inputs.copy() - forward_kwargs = capture_inputs.pop("_forward_kwargs") + def capture_forward_fn( + capture_inputs: Dict[str, Any]) -> Any: with MoeLoadBalancerIterContext(moe_load_balancer): - return self._forward_step(capture_inputs, - **forward_kwargs) + return forward_fn(capture_inputs) self.encoder_cuda_graph_runner.capture( - key, forward_fn, { - **model_inputs, "_forward_kwargs": forward_kwargs - }) + key, capture_forward_fn, model_inputs) with MoeLoadBalancerIterContext(moe_load_balancer): graph_outputs = self.encoder_cuda_graph_runner.replay( - key, { - **model_inputs, "_forward_kwargs": forward_kwargs - }) + key, model_inputs) + + return postprocess_graph_outputs(graph_outputs) + + @torch.inference_mode() + @with_model_extra_attrs(lambda self: self.model.extra_attrs) + @nvtx_range("encoder_forward") + def encoder_forward(self, inputs: Dict[str, Any], + **kwargs) -> Dict[str, Any]: + """Run the tensor-level encode-only forward path. + + This is the ``LLM.encode()``/``encode_only`` path. It bypasses + ``ScheduledRequests`` and ``LlmRequest`` entirely, consumes a raw + tensor-style inputs dict, and returns model outputs such as logits. + It is separate from :meth:`forward_encoder`, which consumes scheduled + encoder-decoder requests and returns encoder hidden states for the + decoder. + + Attempts encoder CUDA graph capture/replay when enabled; otherwise + runs the same forward eagerly. + + Args: + inputs: Dict with 'input_ids' and 'seq_lens' (required), plus + any model-specific kwargs (token_type_ids, inputs_embeds, etc.). + + Returns: + Dict with 'logits' tensor and any other model outputs. + """ + batch_size = len(inputs['seq_lens']) + attn_metadata = self._set_up_attn_metadata( + kv_cache_manager=None + ) if self.encoder_attn_metadata is None else self.encoder_attn_metadata + forward_kwargs = { + "gather_ids": None, + "gather_context_logits": False, + **kwargs, + } + + def prepare_model_inputs(padded_inputs: Dict[str, Any], + graph_attn_metadata: Any, + key: Optional[Any]) -> Dict[str, Any]: + model_inputs = self._prepare_encoder_inputs( + padded_inputs, + attn_metadata=graph_attn_metadata, + padded_num_tokens=key[1] if key is not None else None) + return { + **model_inputs, + "_forward_kwargs": forward_kwargs, + } + + def forward_fn(model_inputs: Dict[str, Any]) -> Dict[str, Any]: + model_inputs = model_inputs.copy() + model_forward_kwargs = model_inputs.pop("_forward_kwargs") + return self._forward_step(model_inputs, **model_forward_kwargs) - # Return a clone to avoid sharing data_ptr with the static buffers. + def postprocess_graph_outputs( + graph_outputs: Dict[str, Any]) -> Dict[str, Any]: outputs = {} for name, value in graph_outputs.items(): if isinstance(value, torch.Tensor): @@ -5011,6 +5071,14 @@ def forward_fn( return outputs + return self._run_encoder_forward_with_cuda_graph( + runner_inputs=inputs, + attn_metadata=attn_metadata, + prepare_model_inputs=prepare_model_inputs, + forward_fn=forward_fn, + postprocess_graph_outputs=postprocess_graph_outputs, + ) + @torch.inference_mode() @with_model_extra_attrs(lambda self: self.model.extra_attrs) def forward(self, @@ -5129,8 +5197,6 @@ def forward(self, ) can_run_graph = key is not None - self._record_decoder_cuda_graph_phase(scheduled_requests, - can_run_graph) if can_run_graph: attn_metadata = maybe_attn_metadata spec_metadata = maybe_spec_metadata @@ -5367,50 +5433,22 @@ 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})") @@ -5445,7 +5483,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() @@ -5457,7 +5495,7 @@ def _prepare_tp_inputs_encoder( dtype=torch.int, pin_memory=prefer_pinned()) - inputs = { + return { 'encoder_input_ids': encoder_input_ids_t.to('cuda', non_blocking=True), 'encoder_position_ids': @@ -5466,10 +5504,64 @@ def _prepare_tp_inputs_encoder( encoder_attn_metadata, 'encoder_seq_lens': sequence_lengths, + 'encoder_input_ids_host': + encoder_input_ids, + 'encoder_position_ids_host': + encoder_position_ids, '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( @@ -5524,13 +5616,91 @@ 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: + """Run the encoder stack, replaying an encoder CUDA graph when possible.""" + 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) + + if input_ids is None or position_ids is None: + moe_load_balancer: MoeLoadBalancer = getattr( + self, 'moe_load_balancer', None) + with MoeLoadBalancerIterContext(moe_load_balancer): + 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'), + } + + def prepare_model_inputs(padded_inputs: Dict[str, Any], + graph_attn_metadata: Any, + key: Optional[Any]) -> Dict[str, Any]: + if key is None: + return inputs + + padded_num_tokens = key[1] + graph_attn_metadata.prepare_encoder_cuda_graph_replay( + padded_inputs['seq_lens'], padded_num_tokens) + return { + **padded_inputs, + 'attn_metadata': graph_attn_metadata, + } + + def forward_fn(model_inputs: Dict[str, Any]) -> torch.Tensor: + if 'attn_metadata' in model_inputs: + return self._forward_step_encoder_cuda_graph(model_inputs) + return self._forward_step_encoder(model_inputs) + + def postprocess_graph_outputs(graph_outputs: Any) -> torch.Tensor: + 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 self._run_encoder_forward_with_cuda_graph( + runner_inputs=runner_inputs, + attn_metadata=inputs['encoder_attn_metadata'], + prepare_model_inputs=prepare_model_inputs, + forward_fn=forward_fn, + postprocess_graph_outputs=postprocess_graph_outputs, + ) + @nvtx_range("forward_encoder") def forward_encoder( self, encoder_requests: List[LlmRequest], resource_manager: Optional[ResourceManager] = None, ) -> Tuple[torch.Tensor, List[int]]: - """Run the encoder stack for ``encoder_requests``. + """Run the encoder-decoder encoder pass for scheduled requests. + + This is the encoder-side prepass for T5/BART-style + encoder-decoder generation. It consumes ``LlmRequest`` objects with + ``encoder_tokens``, packs them into one no-cache encoder batch, and + returns hidden states for decoder cross-attention. It is separate from + :meth:`encoder_forward`, which is the raw tensor ``LLM.encode()`` path + and returns model outputs such as logits. Returns a tuple ``(encoder_hidden_states, encoder_seq_lens)`` where the hidden states tensor is shaped @@ -5550,8 +5720,8 @@ def forward_encoder( with torch.inference_mode(): inputs = self._prepare_tp_inputs_encoder( encoder_requests, resource_manager=resource_manager) - self._record_encoder_cuda_graph_phase(False) - 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/tests/integration/defs/llmapi/test_llm_api_pytorch_bart.py b/tests/integration/defs/llmapi/test_llm_api_pytorch_bart.py index be54838c739b..abd9d46f07c8 100644 --- a/tests/integration/defs/llmapi/test_llm_api_pytorch_bart.py +++ b/tests/integration/defs/llmapi/test_llm_api_pytorch_bart.py @@ -21,6 +21,7 @@ from tensorrt_llm.llmapi import ( LLM, CudaGraphConfig, + EncodeCudaGraphConfig, KvCacheConfig, RequestOutput, SamplingParams, @@ -223,24 +224,37 @@ def _sampling_params(num_beams: int, num_return_sequences: int) -> SamplingParam def _cuda_graph_config( enabled: bool, batch_sizes: list[int] | None = None, -) -> CudaGraphConfig | None: +) -> CudaGraphConfig | EncodeCudaGraphConfig | None: if not enabled: return None - return CudaGraphConfig(batch_sizes=batch_sizes or [1], enable_padding=True) + return EncodeCudaGraphConfig( + batch_sizes=batch_sizes or [1], + max_num_token=_MAX_SEQUENCE_LENGTH, + max_seq_len=_MAX_SEQUENCE_LENGTH, + enable_padding=True, + ) -def _assert_cuda_graphs_captured(llm: LLM, enabled: bool) -> None: - if not enabled: - return +def _assert_encoder_decoder_cuda_graph_state( + llm: LLM, + enabled: bool, + batch_sizes: list[int] | None, +) -> None: model_engine = llm._executor.engine.model_engine - assert model_engine.cuda_graph_runner.graphs - -def _assert_cuda_graph_padding_used(llm: LLM, batch_sizes: list[int] | None) -> None: - if batch_sizes is None: + if not enabled: + assert not model_engine.encoder_cuda_graph_runner.enabled + assert not model_engine.cuda_graph_runner.enabled + assert not model_engine.encoder_cuda_graph_runner.graphs + assert not model_engine.cuda_graph_runner.graphs return - model_engine = llm._executor.engine.model_engine - assert model_engine.cuda_graph_runner.padding_dummy_requests + + assert model_engine.encoder_cuda_graph_runner.enabled + assert model_engine.encoder_cuda_graph_runner.graphs + assert model_engine.cuda_graph_runner.enabled + assert model_engine.cuda_graph_runner.graphs + if batch_sizes is not None: + assert model_engine.cuda_graph_runner.padding_dummy_requests def _enable_trtllm_gen_attention(monkeypatch: pytest.MonkeyPatch) -> None: @@ -301,11 +315,6 @@ def _assert_expected_generation( assert token_ids_by_output == expected_token_ids_by_output -@pytest.mark.parametrize( - "expected_output_token_ids_by_output,torch_dtype,use_kv_cache_manager_v2," - "enable_cuda_graph,num_beams,num_return_sequences,exact_match", - _TEST_CASES, -) def _run_bart_pytorch_generate_encoder_decoder( monkeypatch: pytest.MonkeyPatch, expected_output_token_ids_by_output: list[list[int]] | None, @@ -368,8 +377,11 @@ def _run_bart_pytorch_generate_encoder_decoder( exact_match, expected_output_token_ids_by_output, ) - _assert_cuda_graphs_captured(llm, enable_cuda_graph) - _assert_cuda_graph_padding_used(llm, cuda_graph_batch_sizes) + _assert_encoder_decoder_cuda_graph_state( + llm, + enable_cuda_graph, + cuda_graph_batch_sizes, + ) @pytest.mark.parametrize( 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 4601d88a1a9b..263768365db0 100644 --- a/tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py +++ b/tests/integration/defs/llmapi/test_llm_api_pytorch_t5.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, @@ -95,6 +97,7 @@ "t5-small": [_EXPECTED_TRANSLATION_FRAGMENT, "Buch"], "flan-t5-small": [_EXPECTED_TRANSLATION_FRAGMENT, "Buch"], } +_MIXED_CONTEXT_GENERATION_MAX_NEW_TOKENS = 8 def _test_case( @@ -509,13 +512,61 @@ def _sampling_params(num_beams: int, num_return_sequences: int) -> SamplingParam def _cuda_graph_config( enabled: bool, batch_sizes: list[int] | None = None, -) -> CudaGraphConfig | None: - return CudaGraphConfig(batch_sizes=batch_sizes or [1]) if enabled else None +) -> CudaGraphConfig | EncodeCudaGraphConfig | None: + if not enabled: + return None + return EncodeCudaGraphConfig( + batch_sizes=batch_sizes or [1], + max_num_token=_MAX_SEQUENCE_LENGTH, + max_seq_len=_MAX_SEQUENCE_LENGTH, + enable_padding=True, + ) + + +def _assert_encoder_decoder_cuda_graph_state( + llm: LLM, + enabled: bool, + batch_sizes: list[int] | None, +) -> None: + model_engine = llm._executor.engine.model_engine + + if not enabled: + assert not model_engine.encoder_cuda_graph_runner.enabled + assert not model_engine.cuda_graph_runner.enabled + assert not model_engine.encoder_cuda_graph_runner.graphs + assert not model_engine.cuda_graph_runner.graphs + return + + assert model_engine.encoder_cuda_graph_runner.enabled + assert model_engine.encoder_cuda_graph_runner.graphs + assert model_engine.cuda_graph_runner.enabled + assert model_engine.cuda_graph_runner.graphs + if batch_sizes is not None: + assert model_engine.cuda_graph_runner.padding_dummy_requests + + +def _assert_mixed_context_generation_cuda_graph_state(llm: LLM) -> None: + model_engine = llm._executor.engine.model_engine + + assert model_engine.cuda_graph_runner.enabled + assert model_engine.cuda_graph_runner.graphs + assert model_engine.encoder_cuda_graph_runner.enabled + assert model_engine.encoder_cuda_graph_runner.graphs + assert model_engine.cuda_graph_runner.padding_dummy_requests + + +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_t5_response( response: RequestOutput, num_return_sequences: int, + max_tokens: int = _MAX_NEW_TOKENS, ) -> list[list[int]]: assert response.finished @@ -523,7 +574,7 @@ def _assert_t5_response( token_ids_by_output = [] for output in response.outputs: assert output.token_ids is not None - assert 0 < len(output.token_ids) <= _MAX_NEW_TOKENS + assert 0 < len(output.token_ids) <= max_tokens token_ids_by_output.append(output.token_ids) return token_ids_by_output @@ -624,6 +675,11 @@ def _run_t5_pytorch_generate_encoder_decoder( exact_match, expected_output_token_ids_by_output, ) + _assert_encoder_decoder_cuda_graph_state( + llm, + enable_cuda_graph, + cuda_graph_batch_sizes, + ) @pytest.mark.parametrize( @@ -741,3 +797,71 @@ def test_t5_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch( expected_token_ids_by_output=expected_token_ids, expected_text_fragment=expected_text_fragment, ) + + +def test_t5_pytorch_generate_encoder_decoder_mixed_context_generation_batch( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("TLLM_WORKER_USE_SINGLE_PROCESS", "1") + monkeypatch.setenv("TRTLLM_SKIP_KV_CACHE_ESTIMATION", "1") + + model_name = "t5-small" + model_path = _get_t5_model_path(model_name) + first_sampling_params = SamplingParams( + max_tokens=_MIXED_CONTEXT_GENERATION_MAX_NEW_TOKENS, + temperature=0.0, + ignore_eos=True, + logits_processor=_SleepLogitsProcessor(delay_seconds=0.02), + ) + second_sampling_params = SamplingParams( + max_tokens=_MAX_NEW_TOKENS, + temperature=0.0, + ) + + with LLM( + model_path, + backend="pytorch", + attn_backend="TRTLLM", + cuda_graph_config=_cuda_graph_config(True, [2]), + disable_overlap_scheduler=True, + dtype="bfloat16", + enable_chunked_prefill=False, + 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: + 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() + + _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) + _assert_mixed_context_generation_cuda_graph_state(llm) From a1f69dfc111636f86e28131495efceb7acd29fff Mon Sep 17 00:00:00 2001 From: Guiju Zhang <7135567+cascade812@users.noreply.github.com> Date: Fri, 26 Jun 2026 14:55:43 -0700 Subject: [PATCH 3/9] update Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com> --- tensorrt_llm/_torch/attention_backend/fmha/fallback.py | 2 +- tensorrt_llm/_torch/attention_backend/trtllm.py | 10 ++++++++++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/attention_backend/fmha/fallback.py b/tensorrt_llm/_torch/attention_backend/fmha/fallback.py index c13a3298dafa..73c5778379d8 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/fallback.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/fallback.py @@ -94,7 +94,7 @@ def forward( block_ids_per_seq=metadata.block_ids_per_seq, tokens_per_block=metadata.tokens_per_block, max_num_requests=metadata.max_num_requests, - beam_width=metadata.beam_width, + beam_width=metadata.effective_beam_width, use_paged_context_fmha=metadata.use_paged_context_fmha, helix_position_offsets=metadata.helix_position_offsets, helix_is_inactive_rank=metadata.helix_is_inactive_rank, diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index 63f92fa71728..b0af5852536e 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -80,6 +80,16 @@ class TrtllmAttentionMetadata(AttentionMetadata): # when beam search is enabled. beam_width: int = 1 + @property + def effective_beam_width(self) -> int: + # Only use this for the fallback kernel's beam_width argument. + # Cross-attention reads request-scoped encoder K/V that is written once + # and reused unchanged by every decoder beam. Metadata preparation still + # uses beam_width to expand cross block-offset rows to decoder-sequence + # scope, but the fallback kernel should treat the cross K/V cache as + # non-beam-packed. + return 1 if self.is_cross else self.beam_width + # TrtllmAttention needs to know the max sequence length. # Implemented as a property to support no cache mode. max_seq_len: Optional[int] From f5a527dd77c4a659c462da14382b103b4b4b9dd9 Mon Sep 17 00:00:00 2001 From: Guiju Zhang <7135567+cascade812@users.noreply.github.com> Date: Sun, 28 Jun 2026 21:26:33 -0700 Subject: [PATCH 4/9] fix ci and update interface Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com> --- .../_torch/attention_backend/sparse/rocket.py | 2 + .../_torch/pyexecutor/mamba_cache_manager.py | 2 + .../_torch/pyexecutor/model_engine.py | 185 ++++++++--- tensorrt_llm/llmapi/__init__.py | 10 +- tensorrt_llm/llmapi/llm_args.py | 29 +- tensorrt_llm/llmapi/llm_utils.py | 11 +- .../defs/llmapi/test_llm_api_pytorch_bart.py | 306 +++++++++++++++--- .../defs/llmapi/test_llm_api_pytorch_t5.py | 238 ++++++++++++-- .../test_lists/test-db/l0_b200.yml | 16 +- .../test_lists/test-db/l0_h100.yml | 16 +- tests/integration/test_lists/waives.txt | 18 -- 11 files changed, 687 insertions(+), 146 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/rocket.py b/tensorrt_llm/_torch/attention_backend/sparse/rocket.py index 76874307c6d4..400dff2a5259 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/rocket.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/rocket.py @@ -1027,6 +1027,7 @@ def add_dummy_requests( kv_reserve_draft_tokens: Optional[int] = None, use_mrope: bool = False, max_beam_width: int = 1, + encoder_output_lens: Optional[List[int]] = None, num_extra_decoding_steps: int = 0, draft_kv_cache_manager=None, ): @@ -1039,6 +1040,7 @@ def add_dummy_requests( kv_reserve_draft_tokens=kv_reserve_draft_tokens, use_mrope=use_mrope, max_beam_width=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, ) diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 454dfa08fba7..d4d54d7df905 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -1804,6 +1804,7 @@ def add_dummy_requests( kv_reserve_draft_tokens: Optional[int] = None, use_mrope: bool = False, max_beam_width: int = 1, + encoder_output_lens: Optional[List[int]] = None, # For capturable drafting loops. During normal inference, the draft model always # has enough KV cache space to fit all of our draft tokens. During warmup, however, # we need to make the KV cache manager aware that multiple autoregressive steps will @@ -1820,6 +1821,7 @@ def add_dummy_requests( kv_reserve_draft_tokens=kv_reserve_draft_tokens, use_mrope=use_mrope, max_beam_width=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, ) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 75ff2795b964..4507c55e12ff 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -28,6 +28,7 @@ create_input_processor_with_hash) from tensorrt_llm.llmapi.llm_args import (CudaGraphConfig, DecodingBaseConfig, EncodeCudaGraphConfig, + EncoderDecoderCudaGraphConfig, SeqLenAwareSparseAttentionConfig, TorchCompileConfig, TorchLlmArgs) from tensorrt_llm.logger import logger @@ -368,37 +369,76 @@ def __init__( self._init_model_capacity() self.cuda_graph_config = self.llm_args.cuda_graph_config - cuda_graph_batch_sizes = self.cuda_graph_config.batch_sizes if self.cuda_graph_config else CudaGraphConfig.model_fields[ + + # Split config into decoder and encoder parts. + # EncoderDecoderCudaGraphConfig: explicit sub-configs for each pass. + # EncodeCudaGraphConfig: encode-only models only; no decoder graphs. + # DecodeCudaGraphConfig / None: decoder only; no encoder graphs. + if isinstance(self.cuda_graph_config, EncoderDecoderCudaGraphConfig): + _decode_cfg = self.cuda_graph_config.decoder + _encode_cfg = self.cuda_graph_config.encoder + elif isinstance(self.cuda_graph_config, EncodeCudaGraphConfig): + if self._is_encoder_decoder_model(): + logger.warning( + "EncodeCudaGraphConfig is not supported for encoder-decoder " + "models. Use EncoderDecoderCudaGraphConfig instead. " + "Encoder and decoder CUDA graphs will be disabled.") + self.cuda_graph_config = None + _decode_cfg = None + _encode_cfg = None + else: + _decode_cfg = None + _encode_cfg = self.cuda_graph_config + else: + _decode_cfg = self.cuda_graph_config + _encode_cfg = None + + # Decoder CUDA graph batch-size buckets. + cuda_graph_batch_sizes = _decode_cfg.batch_sizes if _decode_cfg else CudaGraphConfig.model_fields[ 'batch_sizes'].default - cuda_graph_padding_enabled = self.cuda_graph_config.enable_padding if self.cuda_graph_config else CudaGraphConfig.model_fields[ + cuda_graph_padding_enabled = _decode_cfg.enable_padding if _decode_cfg else CudaGraphConfig.model_fields[ 'enable_padding'].default + # Encoder CUDA graph batch-size buckets (may differ from decoder's). + encoder_cuda_graph_batch_sizes = _encode_cfg.batch_sizes if _encode_cfg else [] + # 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 _encode_cfg is not None: + cuda_graph_num_tokens = _encode_cfg.num_tokens or [] + cuda_graph_seq_lens = _encode_cfg.seq_lens or [] self._is_encode_only = (self.llm_args.encode_only and not self.llm_args.mm_encoder_only) if ((self._is_encode_only or self._is_encoder_decoder_model()) - and isinstance(self.cuda_graph_config, EncodeCudaGraphConfig) + and _encode_cfg is not None and (not cuda_graph_num_tokens or not cuda_graph_seq_lens)): missing = [] if not cuda_graph_num_tokens: missing.append("num_tokens/max_num_token") if not cuda_graph_seq_lens: missing.append("seq_lens/max_seq_len") + if self._is_encode_only: + example = ( + "EncodeCudaGraphConfig(max_batch_size=64, " + "num_tokens=[128, 256, 512], " + "max_seq_len=128, enable_padding=True) for encode-only models" + ) + else: + example = ( + "EncoderDecoderCudaGraphConfig(" + "encoder=EncodeCudaGraphConfig(" + "max_batch_size=64, num_tokens=[128, 256, 512], " + "max_seq_len=128, enable_padding=True)) for encoder-decoder models" + ) logger.warning( f"Encoder CUDA graph config is set, but " f"{' and '.join(missing)} not set. Encoder CUDA graphs " f"require both. Encoder CUDA graphs 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).") + f"To enable them, specify e.g. {example}.") self.torch_compile_config = self.llm_args.torch_compile_config torch_compile_enabled = bool(self.torch_compile_config is not None) @@ -558,6 +598,10 @@ def __init__( 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 + # Encoder pass may have independent padding setting (EncoderDecoderCudaGraphConfig). + self._encoder_cuda_graph_padding_enabled = ( + _encode_cfg.enable_padding + if _encode_cfg is not None else self._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, @@ -567,25 +611,39 @@ def __init__( self._max_cuda_graph_batch_size = (self._cuda_graph_batch_sizes[-1] if self._cuda_graph_batch_sizes else 0) + # Encoder batch-size buckets are independent from the decoder's. + self._encoder_cuda_graph_batch_sizes = _filter_cuda_graph_batch_sizes( + encoder_cuda_graph_batch_sizes, + self.batch_size, + self.max_num_tokens, + 0, # no draft tokens on the encoder pass + self._encoder_cuda_graph_padding_enabled + ) if encoder_cuda_graph_batch_sizes else [] + + self._max_encoder_cuda_graph_batch_size = ( + self._encoder_cuda_graph_batch_sizes[-1] + if self._encoder_cuda_graph_batch_sizes else 0) + # 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._encoder_cuda_graph_padding_enabled + ) if 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 [] + cuda_graph_seq_lens, self.max_seq_len, self. + _encoder_cuda_graph_padding_enabled) if 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) - self._enable_encoder_cuda_graph = ( - (self._is_encode_only or self._is_encoder_decoder_model()) - and self.cuda_graph_config is not None - and bool(self._cuda_graph_num_tokens) - and bool(self._cuda_graph_seq_lens)) + self._enable_encoder_cuda_graph = ((self._is_encode_only + or self._is_encoder_decoder_model()) + and _encode_cfg 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( ) @@ -679,11 +737,11 @@ def __init__( # Create Encoder CUDA graph config and runner. encoder_cuda_graph_runner_config = EncoderCUDAGraphRunnerConfig( use_cuda_graph=self._enable_encoder_cuda_graph, - cuda_graph_padding_enabled=self._cuda_graph_padding_enabled, - cuda_graph_batch_sizes=self._cuda_graph_batch_sizes, + cuda_graph_padding_enabled=self._encoder_cuda_graph_padding_enabled, + cuda_graph_batch_sizes=self._encoder_cuda_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_batch_size=self._max_encoder_cuda_graph_batch_size, max_cuda_graph_num_tokens=self._max_cuda_graph_num_tokens, max_num_tokens=self.max_num_tokens, max_seq_len=self.max_seq_len, @@ -4939,7 +4997,7 @@ def _run_cuda_graph_warmup_encoder_decoder(self) -> None: def _get_encoder_cuda_graph_warmup_configs( self) -> List[Tuple[int, int, int]]: """Return feasible (batch_size, num_tokens, max_seq_len) graph keys.""" - 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) @@ -4960,7 +5018,16 @@ def _get_encoder_cuda_graph_warmup_configs( warmup_configs.append((bs, nt, sl)) - return warmup_configs + if self._encoder_cuda_graph_padding_enabled: + for bs in batch_sizes: + if bs > self.batch_size: + continue + for sl in reversed(seq_lens_list): + nt = bs * sl + if nt <= self._max_cuda_graph_num_tokens: + warmup_configs.append((bs, nt, sl)) + + return list(dict.fromkeys(warmup_configs)) def _capture_encoder_cuda_graphs(self) -> None: """Capture whole-model encoder CUDA graphs for all feasible keys. @@ -5811,6 +5878,25 @@ def load_weights_from_target_model(self, if callable(loader): loader(target_model) + @staticmethod + def _apply_logits_processors(request, logits_processors, logits_tensor, + beam_width, token_ids, logits_row_offset): + logits_rows = logits_tensor[logits_row_offset:logits_row_offset + + beam_width] + # Reshape to align w/ the shape used in the TRT backend, + # so the same logit processors can be used across both backends. + logits_rows = logits_rows.view(beam_width, 1, -1) + for lp in logits_processors: + lp_params = inspect.signature(lp).parameters + + assert 4 <= len(lp_params) <= 5, ( + "Logit post processor signature must match the `LogitsProcessor` interface " + "defined in `tensorrtllm.sampling_params`.") + lp(request.py_request_id, logits_rows, token_ids, None, None) + + logits_tensor[logits_row_offset:logits_row_offset + + beam_width] = logits_rows.view(beam_width, -1) + def _execute_logit_post_processors(self, scheduled_requests: ScheduledRequests, outputs: dict): @@ -5823,32 +5909,37 @@ def _execute_logit_post_processors(self, # TODO: support models that don't return outputs as dict return - num_ctx_req = scheduled_requests.num_context_requests logits_tensor = outputs["logits"] - for idx, request in enumerate(scheduled_requests.all_requests()): - logits_processors = getattr(request, "py_logits_post_processors", - None) - if not logits_processors: - continue - - token_ids = request.get_tokens(0) - if idx < num_ctx_req and request.py_orig_prompt_len < len( - token_ids): - # Skip as we only need to apply logit processor on the last context request - continue - - logits_row = logits_tensor[idx] - # Reshape to align w/ the shape used in the TRT backend, - # so the same logit processors can be used across both backends. - logits_row = logits_row.view(1, 1, -1) - token_ids = [token_ids] - for lp in logits_processors: - lp_params = inspect.signature(lp).parameters + logits_row_offset = 0 + request_groups = ( + (scheduled_requests.context_requests, True), + (scheduled_requests.generation_requests, False), + ) - assert 4 <= len(lp_params) <= 5, ( - "Logit post processor signature must match the `LogitsProcessor` interface " - "defined in `tensorrtllm.sampling_params`.") - lp(request.py_request_id, logits_row, token_ids, None, None) + for requests, is_context_request in request_groups: + for request in requests: + if is_context_request: + beam_width = 1 + else: + beam_width = request.get_beam_width_by_iter( + for_next_iteration=False) + + logits_processors = getattr(request, + "py_logits_post_processors", None) + if logits_processors: + token_ids = ([request.get_tokens(0)] + if is_context_request else [ + request.get_tokens(beam_idx) + for beam_idx in range(beam_width) + ]) + if (is_context_request + and request.py_orig_prompt_len < len(token_ids[0])): + # Skip as we only need to apply logit processor on the last context request + logits_row_offset += beam_width + continue - logits_tensor[idx] = logits_row.view(-1) + self._apply_logits_processors(request, logits_processors, + logits_tensor, beam_width, + token_ids, logits_row_offset) + logits_row_offset += beam_width diff --git a/tensorrt_llm/llmapi/__init__.py b/tensorrt_llm/llmapi/__init__.py index 78251d1ae978..9583dbc64a94 100644 --- a/tensorrt_llm/llmapi/__init__.py +++ b/tensorrt_llm/llmapi/__init__.py @@ -14,10 +14,11 @@ DeepSeekV4SparseAttentionConfig, DFlashDecodingConfig, DraftTargetDecodingConfig, DynamicBatchConfig, Eagle3DecodingConfig, EagleDecodingConfig, - EncodeCudaGraphConfig, ExtendedRuntimePerfKnobConfig, - KvCacheConfig, LlmArgs, LookaheadDecodingConfig, - MedusaDecodingConfig, MiniMaxM3SparseAttentionConfig, - MoeConfig, MTPDecodingConfig, NGramDecodingConfig, + EncodeCudaGraphConfig, EncoderDecoderCudaGraphConfig, + ExtendedRuntimePerfKnobConfig, KvCacheConfig, LlmArgs, + LookaheadDecodingConfig, MedusaDecodingConfig, + MiniMaxM3SparseAttentionConfig, MoeConfig, + MTPDecodingConfig, NGramDecodingConfig, PARDDecodingConfig, PrometheusMetricsConfig, ReorderRequestPolicyConfig, RocketSparseAttentionConfig, SADecodingConfig, SAEnhancerConfig, @@ -46,6 +47,7 @@ 'CudaGraphConfig', 'DecodeCudaGraphConfig', 'EncodeCudaGraphConfig', + 'EncoderDecoderCudaGraphConfig', 'MoeConfig', 'LookaheadDecodingConfig', 'MedusaDecodingConfig', diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 0be9e847ce43..afdc7ae172cc 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -452,11 +452,30 @@ def _generate_cuda_graph_seq_lens(max_seq_len: int, return sizes +class EncoderDecoderCudaGraphConfig(StrictBaseModel): + """CUDA graph configuration for encoder-decoder models. + + Provides independent configs for the encoder and decoder passes, allowing + separate control of batch sizes, token buckets, and padding for each. + """ + + mode: Literal["encoder_decoder"] = Field( + default="encoder_decoder", description="CUDA graph configuration mode.") + + encoder: EncodeCudaGraphConfig = Field( + description="CUDA graph configuration for the encoder pass.") + + decoder: DecodeCudaGraphConfig = Field( + default_factory=DecodeCudaGraphConfig, + description="CUDA graph configuration for the decoder pass.") + + # For CudaGraphConfig's backward compatibility CudaGraphConfig = DecodeCudaGraphConfig CudaGraphConfigType: TypeAlias = Annotated[ - Union[DecodeCudaGraphConfig, EncodeCudaGraphConfig], + Union[DecodeCudaGraphConfig, EncodeCudaGraphConfig, + EncoderDecoderCudaGraphConfig], Field(discriminator="mode"), ] @@ -4507,8 +4526,12 @@ def infer_cuda_graph_config_mode(cls, v): "num_tokens", "max_num_token", "seq_lens", "max_seq_len" } v = dict(v) - v["mode"] = "encode" if any(k in v and v[k] not in (None, 0) - for k in encoder_keys) else "decode" + if "encoder" in v or "decoder" in v: + v["mode"] = "encoder_decoder" + elif any(k in v and v[k] not in (None, 0) for k in encoder_keys): + v["mode"] = "encode" + else: + v["mode"] = "decode" return v multimodal_config: MultimodalConfig = Field( diff --git a/tensorrt_llm/llmapi/llm_utils.py b/tensorrt_llm/llmapi/llm_utils.py index 6b9fc63511fa..ff22f46807b5 100644 --- a/tensorrt_llm/llmapi/llm_utils.py +++ b/tensorrt_llm/llmapi/llm_utils.py @@ -38,11 +38,11 @@ from .llm_args import (CalibConfig, CudaGraphConfig, DecodeCudaGraphConfig, DraftTargetDecodingConfig, Eagle3DecodingConfig, EagleDecodingConfig, EncodeCudaGraphConfig, - KvCacheConfig, LlmArgs, LookaheadDecodingConfig, - MedusaDecodingConfig, MTPDecodingConfig, - NGramDecodingConfig, SchedulerConfig, TorchLlmArgs, - UserProvidedDecodingConfig, _ModelFormatKind, - _ModelWrapper, _ParallelConfig, + EncoderDecoderCudaGraphConfig, KvCacheConfig, LlmArgs, + LookaheadDecodingConfig, MedusaDecodingConfig, + MTPDecodingConfig, NGramDecodingConfig, SchedulerConfig, + TorchLlmArgs, UserProvidedDecodingConfig, + _ModelFormatKind, _ModelWrapper, _ParallelConfig, update_llm_args_with_extra_dict, update_llm_args_with_extra_options) # yapf: enable @@ -1021,6 +1021,7 @@ class LlmBuildStats: 'CudaGraphConfig', 'DecodeCudaGraphConfig', 'EncodeCudaGraphConfig', + 'EncoderDecoderCudaGraphConfig', 'KvCacheConfig', 'CachedModelLoader', 'EagleDecodingConfig', 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 8fce56ccf90a..605997c7ad51 100644 --- a/tests/integration/defs/llmapi/test_llm_api_pytorch_bart.py +++ b/tests/integration/defs/llmapi/test_llm_api_pytorch_bart.py @@ -21,7 +21,9 @@ from tensorrt_llm.llmapi import ( LLM, CudaGraphConfig, + DecodeCudaGraphConfig, EncodeCudaGraphConfig, + EncoderDecoderCudaGraphConfig, KvCacheConfig, RequestOutput, SamplingParams, @@ -31,28 +33,67 @@ from ..conftest import llm_models_root _SOURCE_TEXT = ( - "Summarize: NVIDIA builds fast inference software for large language models. " - "TensorRT-LLM supports encoder-decoder models such as BART and T5." + "Summarize: The engineering team released a faster inference service on Monday. " + "The update improves batching, lowers latency, and adds detailed monitoring for operators." ) _MIXED_ENCODER_SOURCE_TEXTS = [ _SOURCE_TEXT, ( - "Summarize: The city opened a new public library on Monday. Residents said " - "the library has quiet rooms, computer access, and a large children section." + "Summarize: The company opened a training center on Monday. Managers said " + "the center adds classrooms, simulation labs, and career coaching for workers." ), ] +_FIXED_SEQ_LENS_SOURCE_TEXTS = [ + ( + "Summarize: The engineering team released a faster inference service on Monday. " + "The update improves batching, lowers latency, and adds detailed monitoring for operators." + ), + ( + "Summarize: The city opened a new public library on Monday. Residents said the library " + "has quiet rooms, computer access, and a large children section." + ), + ( + "Summarize: The company opened a training center on Monday. Managers said the center " + "adds classrooms, simulation labs, and career coaching for local workers." + ), + ( + "Summarize: The museum opened a science exhibit on Monday. Curators said the exhibit " + "has space models, interactive lessons, and workshops for local children." + ), +] +_FIXED_SEQ_LENS_EXPECTED_TOKEN_IDS_BY_REQUEST = [ + # "The update improves batching, lowers latency" + [[0, 133, 2935, 15296, 14398, 154, 6, 32222, 35940, 2]], + # "Residents said the library has quiet rooms," + [[0, 35129, 26, 5, 5560, 34, 5128, 5351, 6, 2]], + # "The company opened a training center on Monday" + [[0, 133, 138, 1357, 10, 1058, 1312, 15, 302, 2]], + # "The museum opened a science exhibit on Monday" + [[0, 133, 5707, 1357, 10, 2866, 8483, 15, 302, 2]], +] _MODEL_NAME = "bart-large-cnn" -_MAX_NEW_TOKENS = 8 +_MAX_NEW_TOKENS = 10 _MAX_SEQUENCE_LENGTH = 128 _MAX_KV_TOKENS = 384 _MIN_GPU_MEMORY_MB = 16_000 _FREE_GPU_MEMORY_FRACTION = 0.2 _CROSS_KV_CACHE_FRACTION = 0.5 -_EXPECTED_GREEDY_OUTPUT_TOKEN_IDS = [0, 565, 35354, 13963, 12, 6006, 448, 2] -_EXPECTED_TEXT_FRAGMENT = "TensorRT" -_MIXED_ENCODER_EXPECTED_TEXT_FRAGMENTS = [ - _EXPECTED_TEXT_FRAGMENT, - "library", +# "The update improves batching, lowers latency" +_EXPECTED_GREEDY_OUTPUT_TOKEN_IDS = [0, 133, 2935, 15296, 14398, 154, 6, 32222, 35940, 2] +_FIXED_SEQ_LENS_MAX_NEW_TOKENS = 10 +_EXPECTED_BEAM_OUTPUT_TOKEN_IDS_BY_BEAMS = { + 2: [ + # "The update improves batching, lowers latency" + [0, 133, 2935, 15296, 14398, 154, 6, 32222, 35940, 2], + # "The update improves batching, lowers" + [0, 0, 133, 2935, 15296, 14398, 154, 6, 32222, 2], + ], +} +_MIXED_ENCODER_EXPECTED_TOKEN_IDS_BY_REQUEST = [ + # "The update improves batching, lowers latency" + [[0, 133, 2935, 15296, 14398, 154, 6, 32222, 35940, 2]], + # "The company opened a training center on Monday" + [[0, 133, 138, 1357, 10, 1058, 1312, 15, 302, 2]], ] @@ -67,8 +108,11 @@ def _test_case( cuda_graph_batch_sizes: list[int] | None = None, kv_cache_dtype: str = "auto", ): - expected_output_token_ids = [_EXPECTED_GREEDY_OUTPUT_TOKEN_IDS] if num_beams == 1 else None - assert not exact_match or expected_output_token_ids is not None + expected_output_token_ids = ( + [_EXPECTED_GREEDY_OUTPUT_TOKEN_IDS] + if num_beams == 1 + else _EXPECTED_BEAM_OUTPUT_TOKEN_IDS_BY_BEAMS[num_beams] + ) return pytest.param( expected_output_token_ids, @@ -119,7 +163,7 @@ def _test_case( enable_cuda_graph=False, num_beams=2, num_return_sequences=2, - exact_match=False, + exact_match=True, feature_id="bf16-kv-v1-cuda-graph-off-beam2", ), _test_case( @@ -128,7 +172,7 @@ def _test_case( enable_cuda_graph=True, num_beams=2, num_return_sequences=2, - exact_match=False, + exact_match=True, feature_id="bf16-kv-v1-cuda-graph-on-beam2", ), _test_case( @@ -168,20 +212,54 @@ def _mixed_batch_test_case( ) +def _fixed_seq_lens_shape_test_case( + torch_dtype: str, + use_kv_cache_manager_v2: bool, + batch_size: int, + feature_id: str, + kv_cache_dtype: str = "auto", +): + return pytest.param( + torch_dtype, + use_kv_cache_manager_v2, + 1, + 1, + batch_size, + kv_cache_dtype, + id=f"{feature_id}-{_MODEL_NAME}", + ) + + _MIXED_BATCH_TEST_CASES = [ _mixed_batch_test_case( torch_dtype="bfloat16", use_kv_cache_manager_v2=False, num_beams=1, num_return_sequences=1, - feature_id="bf16-kv-v1-cuda-graph-off-greedy-batch2", + feature_id="bf16-kv-v1-decoder-cuda-graph-on-greedy-batch2", ), _mixed_batch_test_case( torch_dtype="bfloat16", use_kv_cache_manager_v2=True, num_beams=1, num_return_sequences=1, - feature_id="bf16-kv-v2-cuda-graph-off-greedy-batch2", + feature_id="bf16-kv-v2-decoder-cuda-graph-on-greedy-batch2", + ), +] + + +_FIXED_SEQ_LENS_SHAPE_TEST_CASES = [ + _fixed_seq_lens_shape_test_case( + torch_dtype="bfloat16", + use_kv_cache_manager_v2=False, + batch_size=2, + feature_id="bf16-kv-v1-encoder-cuda-graph-fixed-shape-batch2", + ), + _fixed_seq_lens_shape_test_case( + torch_dtype="bfloat16", + use_kv_cache_manager_v2=False, + batch_size=4, + feature_id="bf16-kv-v1-encoder-cuda-graph-fixed-shape-batch4", ), ] @@ -223,14 +301,33 @@ def _sampling_params(num_beams: int, num_return_sequences: int) -> SamplingParam def _cuda_graph_config( enabled: bool, + num_tokens: list[int], + seq_lens: list[int], batch_sizes: list[int] | None = None, -) -> CudaGraphConfig | EncodeCudaGraphConfig | None: +) -> EncoderDecoderCudaGraphConfig | None: if not enabled: return None - return EncodeCudaGraphConfig( + return EncoderDecoderCudaGraphConfig( + encoder=EncodeCudaGraphConfig( + num_tokens=num_tokens, + seq_lens=seq_lens, + enable_padding=True, + ), + decoder=DecodeCudaGraphConfig( + batch_sizes=batch_sizes or [1], + enable_padding=True, + ), + ) + + +def _decoder_cuda_graph_config( + batch_sizes: list[int] | None = None, +) -> CudaGraphConfig: + # CudaGraphConfig is decode-only. It keeps encoder CUDA graphs disabled, + # which is what mixed encoder-length tests want while still covering + # decoder graph capture/replay. + return CudaGraphConfig( batch_sizes=batch_sizes or [1], - max_num_token=_MAX_SEQUENCE_LENGTH, - max_seq_len=_MAX_SEQUENCE_LENGTH, enable_padding=True, ) @@ -249,12 +346,27 @@ def _assert_encoder_decoder_cuda_graph_state( assert not model_engine.cuda_graph_runner.graphs return + _assert_encoder_decoder_cuda_graphs_captured(llm) + if batch_sizes is not None: + assert model_engine.cuda_graph_runner.padding_dummy_requests + + +def _assert_encoder_decoder_cuda_graphs_captured(llm: LLM) -> None: + model_engine = llm._executor.engine.model_engine + assert model_engine.encoder_cuda_graph_runner.enabled assert model_engine.encoder_cuda_graph_runner.graphs assert model_engine.cuda_graph_runner.enabled assert model_engine.cuda_graph_runner.graphs - if batch_sizes is not None: - assert model_engine.cuda_graph_runner.padding_dummy_requests + + +def _assert_decoder_cuda_graphs_captured(llm: LLM) -> None: + model_engine = llm._executor.engine.model_engine + + assert not model_engine.encoder_cuda_graph_runner.enabled + assert not model_engine.encoder_cuda_graph_runner.graphs + assert model_engine.cuda_graph_runner.enabled + assert model_engine.cuda_graph_runner.graphs def _enable_trtllm_gen_attention(monkeypatch: pytest.MonkeyPatch) -> None: @@ -268,6 +380,7 @@ def _enable_trtllm_gen_attention(monkeypatch: pytest.MonkeyPatch) -> None: def _assert_bart_response( response: RequestOutput, num_return_sequences: int, + max_tokens: int = _MAX_NEW_TOKENS, ) -> list[list[int]]: assert response.finished @@ -275,7 +388,7 @@ def _assert_bart_response( token_ids_by_output = [] for output in response.outputs: assert output.token_ids is not None - assert 0 < len(output.token_ids) <= _MAX_NEW_TOKENS + assert 0 < len(output.token_ids) <= max_tokens token_ids_by_output.append(output.token_ids) return token_ids_by_output @@ -292,18 +405,13 @@ def _assert_expected_generation( tokenizer, token_ids_by_output: list[list[int]], exact_match: bool, - expected_token_ids_by_output: list[list[int]] | None, - expected_text_fragment: str | None = _EXPECTED_TEXT_FRAGMENT, + expected_token_ids_by_output: list[list[int]], ) -> None: decoded_text_by_output = [ tokenizer.decode(token_ids, skip_special_tokens=True) for token_ids in token_ids_by_output ] assert all(decoded_text_by_output) - if expected_token_ids_by_output is None: - if expected_text_fragment is not None: - assert all(expected_text_fragment in text for text in decoded_text_by_output) - else: - assert token_ids_by_output[0] == expected_token_ids_by_output[0] + assert token_ids_by_output[0] == expected_token_ids_by_output[0] if len(token_ids_by_output) > 1: assert len({tuple(token_ids) for token_ids in token_ids_by_output}) == len( token_ids_by_output @@ -311,13 +419,12 @@ def _assert_expected_generation( if not exact_match: return - assert expected_token_ids_by_output is not None assert token_ids_by_output == expected_token_ids_by_output def _run_bart_pytorch_generate_encoder_decoder( monkeypatch: pytest.MonkeyPatch, - expected_output_token_ids_by_output: list[list[int]] | None, + expected_output_token_ids_by_output: list[list[int]], torch_dtype: str, use_kv_cache_manager_v2: bool, enable_cuda_graph: bool, @@ -332,6 +439,8 @@ def _run_bart_pytorch_generate_encoder_decoder( model_path = _get_bart_model_path() tokenizer = AutoTokenizer.from_pretrained(model_path) + encoder_seq_len = len(tokenizer(_SOURCE_TEXT).input_ids) + encoder_num_tokens = encoder_seq_len case_id = ( f"model={_MODEL_NAME}, dtype={torch_dtype}, kv_v2={use_kv_cache_manager_v2}, " f"cuda_graph={enable_cuda_graph}, beams={num_beams}, returns={num_return_sequences}, " @@ -343,7 +452,12 @@ def _run_bart_pytorch_generate_encoder_decoder( model_path, backend="pytorch", attn_backend="TRTLLM", - cuda_graph_config=_cuda_graph_config(enable_cuda_graph, cuda_graph_batch_sizes), + cuda_graph_config=_cuda_graph_config( + enable_cuda_graph, + [encoder_num_tokens], + [encoder_seq_len], + cuda_graph_batch_sizes, + ), disable_overlap_scheduler=True, dtype=torch_dtype, enable_chunked_prefill=False, @@ -394,7 +508,7 @@ def _run_bart_pytorch_generate_encoder_decoder( ) def test_bart_pytorch_generate_encoder_decoder_end_to_end( monkeypatch: pytest.MonkeyPatch, - expected_output_token_ids_by_output: list[list[int]] | None, + expected_output_token_ids_by_output: list[list[int]], torch_dtype: str, use_kv_cache_manager_v2: bool, enable_cuda_graph: bool, @@ -418,6 +532,117 @@ def test_bart_pytorch_generate_encoder_decoder_end_to_end( ) +@pytest.mark.parametrize( + "torch_dtype,use_kv_cache_manager_v2,num_beams,num_return_sequences,batch_size,kv_cache_dtype", + _FIXED_SEQ_LENS_SHAPE_TEST_CASES, +) +def test_bart_pytorch_generate_encoder_decoder_fixed_encoder_lengths_batch( + monkeypatch: pytest.MonkeyPatch, + torch_dtype: str, + use_kv_cache_manager_v2: bool, + num_beams: int, + num_return_sequences: int, + batch_size: int, + kv_cache_dtype: str, +) -> None: + monkeypatch.setenv("TLLM_WORKER_USE_SINGLE_PROCESS", "1") + monkeypatch.setenv("TRTLLM_SKIP_KV_CACHE_ESTIMATION", "1") + + model_path = _get_bart_model_path() + tokenizer = AutoTokenizer.from_pretrained(model_path) + source_texts = _FIXED_SEQ_LENS_SOURCE_TEXTS[:batch_size] + assert len(source_texts) == batch_size + expected_token_ids_by_request = _FIXED_SEQ_LENS_EXPECTED_TOKEN_IDS_BY_REQUEST[:batch_size] + encoder_seq_lens_by_request = [ + len(tokenizer(source_text).input_ids) for source_text in source_texts + ] + assert len(set(encoder_seq_lens_by_request)) == 1 + encoder_seq_len = encoder_seq_lens_by_request[0] + encoder_num_tokens = sum(encoder_seq_lens_by_request) + graph_key = (batch_size, encoder_num_tokens, encoder_seq_len) + assert num_beams == 1 + assert num_return_sequences == 1 + sampling_params = SamplingParams( + max_tokens=_FIXED_SEQ_LENS_MAX_NEW_TOKENS, + temperature=0.0, + ) + max_num_tokens = max(_MAX_SEQUENCE_LENGTH, encoder_num_tokens) + kv_cache_max_tokens = max( + _MAX_KV_TOKENS, + batch_size * _MAX_SEQUENCE_LENGTH * 2, + ) + case_id = ( + f"model={_MODEL_NAME}, dtype={torch_dtype}, kv_v2={use_kv_cache_manager_v2}, " + f"encoder_cuda_graph=True, beams={num_beams}, returns={num_return_sequences}, " + f"kv_dtype={kv_cache_dtype}, fixed_encoder_seq_lens=True, batch_size={batch_size}" + ) + llm_kwargs = { + "backend": "pytorch", + "attn_backend": "TRTLLM", + "disable_overlap_scheduler": True, + "dtype": torch_dtype, + "enable_chunked_prefill": False, + "kv_cache_config": KvCacheConfig( + enable_block_reuse=False, + max_tokens=kv_cache_max_tokens, + free_gpu_memory_fraction=_FREE_GPU_MEMORY_FRACTION, + cross_kv_cache_fraction=_CROSS_KV_CACHE_FRACTION, + use_kv_cache_manager_v2=use_kv_cache_manager_v2, + dtype=kv_cache_dtype, + ), + "max_batch_size": batch_size, + "max_beam_width": num_beams, + "max_input_len": _MAX_SEQUENCE_LENGTH, + "max_num_tokens": max_num_tokens, + "max_seq_len": _MAX_SEQUENCE_LENGTH, + "model_kwargs": {"torch_dtype": torch_dtype}, + "scheduler_config": SchedulerConfig(use_python_scheduler=True), + } + + with LLM( + model_path, + cuda_graph_config=_cuda_graph_config( + True, + [encoder_num_tokens], + [encoder_seq_len], + [batch_size], + ), + **llm_kwargs, + ) as llm: + model_engine = llm._executor.engine.model_engine + assert graph_key in model_engine.encoder_cuda_graph_runner.graphs + + responses = llm.generate( + source_texts, + sampling_params=sampling_params, + use_tqdm=False, + ) + assert len(responses) == batch_size + + token_ids_by_request = [] + for request_idx, response in enumerate(responses): + token_ids = _assert_bart_response( + response, + num_return_sequences=num_return_sequences, + max_tokens=_FIXED_SEQ_LENS_MAX_NEW_TOKENS, + ) + token_ids_by_request.append(token_ids) + _print_generated_text( + tokenizer, + f"{case_id}, request={request_idx}", + "output", + token_ids, + ) + _assert_expected_generation( + tokenizer, + token_ids, + exact_match=True, + expected_token_ids_by_output=expected_token_ids_by_request[request_idx], + ) + + assert token_ids_by_request == expected_token_ids_by_request + + @pytest.mark.parametrize( "torch_dtype,use_kv_cache_manager_v2,num_beams,num_return_sequences", _MIXED_BATCH_TEST_CASES, @@ -437,14 +662,14 @@ def test_bart_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch( sampling_params = _sampling_params(num_beams, num_return_sequences) case_id = ( f"model={_MODEL_NAME}, dtype={torch_dtype}, kv_v2={use_kv_cache_manager_v2}, " - f"cuda_graph=False, beams={num_beams}, returns={num_return_sequences}, " + f"decoder_cuda_graph=True, beams={num_beams}, returns={num_return_sequences}, " "mixed_encoder_lengths=True, batch_size=2" ) with LLM( model_path, backend="pytorch", attn_backend="TRTLLM", - cuda_graph_config=None, + cuda_graph_config=_decoder_cuda_graph_config([2]), disable_overlap_scheduler=True, dtype=torch_dtype, enable_chunked_prefill=False, @@ -485,7 +710,10 @@ def test_bart_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch( _assert_expected_generation( tokenizer, token_ids, - exact_match=False, - expected_token_ids_by_output=None, - expected_text_fragment=_MIXED_ENCODER_EXPECTED_TEXT_FRAGMENTS[request_idx], + exact_match=True, + expected_token_ids_by_output=_MIXED_ENCODER_EXPECTED_TOKEN_IDS_BY_REQUEST[ + request_idx + ], ) + + _assert_decoder_cuda_graphs_captured(llm) 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 464cea22ee3b..1e1e6e8a3e42 100644 --- a/tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py +++ b/tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py @@ -22,7 +22,9 @@ from tensorrt_llm.llmapi import ( LLM, CudaGraphConfig, + DecodeCudaGraphConfig, EncodeCudaGraphConfig, + EncoderDecoderCudaGraphConfig, KvCacheConfig, RequestOutput, SamplingParams, @@ -36,6 +38,18 @@ _SOURCE_TEXT, "translate English to German: The book is on the table.", ] +_FIXED_SEQ_LENS_SOURCE_TEXTS = [ + "translate English to German: The book is on table.", + "translate English to German: The school is very clean.", + "translate English to German: The car is very small.", + "translate English to German: The food is very good.", +] +_FIXED_SEQ_LENS_EXPECTED_TOKEN_IDS_BY_REQUEST = [ + [644, 4675, 229, 219], + [316, 12853, 229, 1319], + [644, 2040, 229, 1319], + [644, 11722, 229, 1319], +] _MAX_NEW_TOKENS = 4 _MAX_SEQUENCE_LENGTH = 64 _MAX_KV_TOKENS = 256 @@ -100,6 +114,24 @@ _MIXED_CONTEXT_GENERATION_MAX_NEW_TOKENS = 8 +def _fixed_seq_lens_shape_test_case( + model_name: str, + torch_dtype: str, + use_kv_cache_manager_v2: bool, + batch_size: int, + feature_id: str, +): + return pytest.param( + model_name, + torch_dtype, + use_kv_cache_manager_v2, + 1, + 1, + batch_size, + id=f"{feature_id}-{model_name}", + ) + + def _test_case( model_name: str, torch_dtype: str, @@ -440,7 +472,7 @@ def _mixed_batch_test_case( num_beams=2, num_return_sequences=2, exact_match=False, - feature_id="bf16-kv-v1-cuda-graph-off-beam2-batch2", + feature_id="bf16-kv-v1-decoder-cuda-graph-on-beam2-batch2", ), _mixed_batch_test_case( model_name="flan-t5-small", @@ -449,7 +481,7 @@ def _mixed_batch_test_case( num_beams=2, num_return_sequences=2, exact_match=False, - feature_id="bf16-kv-v1-cuda-graph-off-beam2-batch2", + feature_id="bf16-kv-v1-decoder-cuda-graph-on-beam2-batch2", ), _mixed_batch_test_case( model_name="t5-small", @@ -458,7 +490,7 @@ def _mixed_batch_test_case( num_beams=1, num_return_sequences=1, exact_match=True, - feature_id="bf16-kv-v1-cuda-graph-off-greedy-batch2", + feature_id="bf16-kv-v1-decoder-cuda-graph-on-greedy-batch2", ), _mixed_batch_test_case( model_name="t5-small", @@ -467,7 +499,25 @@ def _mixed_batch_test_case( num_beams=1, num_return_sequences=1, exact_match=True, - feature_id="bf16-kv-v2-cuda-graph-off-greedy-batch2", + feature_id="bf16-kv-v2-decoder-cuda-graph-on-greedy-batch2", + ), +] + + +_FIXED_SEQ_LENS_SHAPE_TEST_CASES = [ + _fixed_seq_lens_shape_test_case( + model_name="t5-small", + torch_dtype="bfloat16", + use_kv_cache_manager_v2=False, + batch_size=2, + feature_id="bf16-kv-v1-encoder-cuda-graph-fixed-shape-batch2", + ), + _fixed_seq_lens_shape_test_case( + model_name="t5-small", + torch_dtype="bfloat16", + use_kv_cache_manager_v2=False, + batch_size=4, + feature_id="bf16-kv-v1-encoder-cuda-graph-fixed-shape-batch4", ), ] @@ -509,14 +559,33 @@ def _sampling_params(num_beams: int, num_return_sequences: int) -> SamplingParam def _cuda_graph_config( enabled: bool, + num_tokens: list[int], + seq_lens: list[int], batch_sizes: list[int] | None = None, -) -> CudaGraphConfig | EncodeCudaGraphConfig | None: +) -> EncoderDecoderCudaGraphConfig | None: if not enabled: return None - return EncodeCudaGraphConfig( + return EncoderDecoderCudaGraphConfig( + encoder=EncodeCudaGraphConfig( + num_tokens=num_tokens, + seq_lens=seq_lens, + enable_padding=True, + ), + decoder=DecodeCudaGraphConfig( + batch_sizes=batch_sizes or [1], + enable_padding=True, + ), + ) + + +def _decoder_cuda_graph_config( + batch_sizes: list[int] | None = None, +) -> CudaGraphConfig: + # CudaGraphConfig is decode-only. It keeps encoder CUDA graphs disabled, + # which is what mixed encoder-length tests want while still covering + # decoder graph capture/replay. + return CudaGraphConfig( batch_sizes=batch_sizes or [1], - max_num_token=_MAX_SEQUENCE_LENGTH, - max_seq_len=_MAX_SEQUENCE_LENGTH, enable_padding=True, ) @@ -535,22 +604,32 @@ def _assert_encoder_decoder_cuda_graph_state( assert not model_engine.cuda_graph_runner.graphs return + _assert_encoder_decoder_cuda_graphs_captured(llm) + if batch_sizes is not None: + assert model_engine.cuda_graph_runner.padding_dummy_requests + + +def _assert_encoder_decoder_cuda_graphs_captured(llm: LLM) -> None: + model_engine = llm._executor.engine.model_engine + assert model_engine.encoder_cuda_graph_runner.enabled assert model_engine.encoder_cuda_graph_runner.graphs assert model_engine.cuda_graph_runner.enabled assert model_engine.cuda_graph_runner.graphs - if batch_sizes is not None: - assert model_engine.cuda_graph_runner.padding_dummy_requests -def _assert_mixed_context_generation_cuda_graph_state(llm: LLM) -> None: +def _assert_decoder_cuda_graphs_captured(llm: LLM) -> None: model_engine = llm._executor.engine.model_engine + assert not model_engine.encoder_cuda_graph_runner.enabled + assert not model_engine.encoder_cuda_graph_runner.graphs assert model_engine.cuda_graph_runner.enabled assert model_engine.cuda_graph_runner.graphs - assert model_engine.encoder_cuda_graph_runner.enabled - assert model_engine.encoder_cuda_graph_runner.graphs - assert model_engine.cuda_graph_runner.padding_dummy_requests + + +def _assert_mixed_context_generation_cuda_graph_state(llm: LLM) -> None: + _assert_encoder_decoder_cuda_graphs_captured(llm) + assert llm._executor.engine.model_engine.cuda_graph_runner.padding_dummy_requests class _SleepLogitsProcessor: @@ -634,6 +713,8 @@ def _run_t5_pytorch_generate_encoder_decoder( model_path = _get_t5_model_path(model_name) tokenizer = AutoTokenizer.from_pretrained(model_path) + encoder_seq_len = len(tokenizer(_SOURCE_TEXT).input_ids) + encoder_num_tokens = encoder_seq_len case_id = ( f"model={model_name}, dtype={torch_dtype}, kv_v2={use_kv_cache_manager_v2}, " f"cuda_graph={enable_cuda_graph}, beams={num_beams}, returns={num_return_sequences}" @@ -644,7 +725,12 @@ def _run_t5_pytorch_generate_encoder_decoder( model_path, backend="pytorch", attn_backend="TRTLLM", - cuda_graph_config=_cuda_graph_config(enable_cuda_graph, cuda_graph_batch_sizes), + cuda_graph_config=_cuda_graph_config( + enable_cuda_graph, + [encoder_num_tokens], + [encoder_seq_len], + cuda_graph_batch_sizes, + ), disable_overlap_scheduler=True, dtype=torch_dtype, enable_chunked_prefill=False, @@ -717,6 +803,108 @@ def test_t5_pytorch_generate_encoder_decoder_end_to_end( ) +@pytest.mark.parametrize( + "model_name,torch_dtype,use_kv_cache_manager_v2,num_beams,num_return_sequences,batch_size", + _FIXED_SEQ_LENS_SHAPE_TEST_CASES, +) +def test_t5_pytorch_generate_encoder_decoder_fixed_encoder_lengths_batch( + monkeypatch: pytest.MonkeyPatch, + model_name: str, + torch_dtype: str, + use_kv_cache_manager_v2: bool, + num_beams: int, + num_return_sequences: int, + batch_size: int, +) -> None: + monkeypatch.setenv("TLLM_WORKER_USE_SINGLE_PROCESS", "1") + monkeypatch.setenv("TRTLLM_SKIP_KV_CACHE_ESTIMATION", "1") + + model_path = _get_t5_model_path(model_name) + tokenizer = AutoTokenizer.from_pretrained(model_path) + source_texts = _FIXED_SEQ_LENS_SOURCE_TEXTS[:batch_size] + assert len(source_texts) == batch_size + expected_token_ids_by_request = _FIXED_SEQ_LENS_EXPECTED_TOKEN_IDS_BY_REQUEST[:batch_size] + encoder_seq_lens_by_request = [ + len(tokenizer(source_text).input_ids) for source_text in source_texts + ] + assert len(set(encoder_seq_lens_by_request)) == 1 + encoder_seq_len = encoder_seq_lens_by_request[0] + encoder_num_tokens = sum(encoder_seq_lens_by_request) + sampling_params = _sampling_params(num_beams, num_return_sequences) + max_num_tokens = max( + _MAX_SEQUENCE_LENGTH, + encoder_num_tokens + batch_size * (_MAX_NEW_TOKENS + 1), + ) + case_id = ( + f"model={model_name}, dtype={torch_dtype}, kv_v2={use_kv_cache_manager_v2}, " + f"encoder_cuda_graph=True, beams={num_beams}, returns={num_return_sequences}, " + f"fixed_encoder_seq_lens=True, batch_size={batch_size}" + ) + llm_kwargs = { + "backend": "pytorch", + "attn_backend": "TRTLLM", + "disable_overlap_scheduler": True, + "dtype": torch_dtype, + "enable_chunked_prefill": False, + "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=use_kv_cache_manager_v2, + ), + "max_batch_size": batch_size, + "max_beam_width": num_beams, + "max_input_len": _MAX_SEQUENCE_LENGTH, + "max_num_tokens": max_num_tokens, + "max_seq_len": _MAX_SEQUENCE_LENGTH, + "model_kwargs": {"torch_dtype": torch_dtype}, + "scheduler_config": SchedulerConfig(use_python_scheduler=True), + } + + with LLM( + model_path, + cuda_graph_config=_cuda_graph_config( + True, + [encoder_num_tokens], + [encoder_seq_len], + [batch_size], + ), + **llm_kwargs, + ) as llm: + responses = llm.generate( + source_texts, + sampling_params=sampling_params, + use_tqdm=False, + ) + assert len(responses) == batch_size + + token_ids_by_request = [] + for request_idx, response in enumerate(responses): + token_ids = _assert_t5_response( + response, + num_return_sequences=num_return_sequences, + ) + token_ids_by_request.append(token_ids) + _print_generated_text( + tokenizer, + f"{case_id}, request={request_idx}", + "output", + token_ids, + ) + _assert_expected_generation( + tokenizer, + token_ids, + exact_match=True, + expected_token_ids_by_output=[expected_token_ids_by_request[request_idx]], + ) + + assert token_ids_by_request == [ + [expected_token_ids] for expected_token_ids in expected_token_ids_by_request + ] + _assert_encoder_decoder_cuda_graphs_captured(llm) + + @pytest.mark.parametrize( "model_name,expected_output_token_ids_by_request,torch_dtype,use_kv_cache_manager_v2," "num_beams,num_return_sequences,exact_match", @@ -740,14 +928,14 @@ def test_t5_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch( sampling_params = _sampling_params(num_beams, num_return_sequences) case_id = ( f"model={model_name}, dtype={torch_dtype}, kv_v2={use_kv_cache_manager_v2}, " - f"cuda_graph=False, beams={num_beams}, returns={num_return_sequences}, " + f"decoder_cuda_graph=True, beams={num_beams}, returns={num_return_sequences}, " "mixed_encoder_lengths=True, batch_size=2" ) with LLM( model_path, backend="pytorch", attn_backend="TRTLLM", - cuda_graph_config=None, + cuda_graph_config=_decoder_cuda_graph_config([2]), disable_overlap_scheduler=True, dtype=torch_dtype, enable_chunked_prefill=False, @@ -802,6 +990,8 @@ def test_t5_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch( expected_text_fragment=expected_text_fragment, ) + _assert_decoder_cuda_graphs_captured(llm) + def test_t5_pytorch_generate_encoder_decoder_mixed_context_generation_batch( monkeypatch: pytest.MonkeyPatch, @@ -811,6 +1001,13 @@ def test_t5_pytorch_generate_encoder_decoder_mixed_context_generation_batch( model_name = "t5-small" model_path = _get_t5_model_path(model_name) + tokenizer = AutoTokenizer.from_pretrained(model_path) + encoder_seq_lens_by_request = [ + len(tokenizer(source_text).input_ids) for source_text in _MIXED_ENCODER_SOURCE_TEXTS + ] + encoder_num_tokens = sum(encoder_seq_lens_by_request) + encoder_num_tokens_list = sorted({*encoder_seq_lens_by_request, encoder_num_tokens}) + encoder_seq_lens = sorted(set(encoder_seq_lens_by_request)) first_sampling_params = SamplingParams( max_tokens=_MIXED_CONTEXT_GENERATION_MAX_NEW_TOKENS, temperature=0.0, @@ -826,7 +1023,12 @@ def test_t5_pytorch_generate_encoder_decoder_mixed_context_generation_batch( model_path, backend="pytorch", attn_backend="TRTLLM", - cuda_graph_config=_cuda_graph_config(True, [2]), + cuda_graph_config=_cuda_graph_config( + True, + encoder_num_tokens_list, + encoder_seq_lens, + [2], + ), disable_overlap_scheduler=True, dtype="bfloat16", enable_chunked_prefill=False, diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 6730a0f1b716..7036a199c600 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -67,8 +67,10 @@ l0_b200: - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_end_to_end[fp16-kv-v1-cuda-graph-off-greedy-bart-large-cnn] - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-off-beam2-bart-large-cnn] - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v2-cuda-graph-off-greedy-bart-large-cnn] - - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch[bf16-kv-v1-cuda-graph-off-greedy-batch2-bart-large-cnn] - - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch[bf16-kv-v2-cuda-graph-off-greedy-batch2-bart-large-cnn] + - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch[bf16-kv-v1-decoder-cuda-graph-on-greedy-batch2-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_generate_encoder_decoder_fixed_encoder_lengths_batch[bf16-kv-v1-encoder-cuda-graph-fixed-shape-batch2-bart-large-cnn] + - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_fixed_encoder_lengths_batch[bf16-kv-v1-encoder-cuda-graph-fixed-shape-batch4-bart-large-cnn] - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-off-beam2-t5-small0] - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-off-beam2-flan-t5-small] - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-off-beam2-t5-base] @@ -90,10 +92,12 @@ l0_b200: - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[fp16-kv-v2-cuda-graph-off-greedy-flan-t5-small] - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[fp32-kv-v2-cuda-graph-off-greedy-flan-t5-small] - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v2-cuda-graph-off-greedy-byt5-small] - - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch[bf16-kv-v1-cuda-graph-off-beam2-batch2-t5-small] - - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch[bf16-kv-v1-cuda-graph-off-beam2-batch2-flan-t5-small] - - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch[bf16-kv-v1-cuda-graph-off-greedy-batch2-t5-small] - - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch[bf16-kv-v2-cuda-graph-off-greedy-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-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-v1-decoder-cuda-graph-on-greedy-batch2-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_generate_encoder_decoder_fixed_encoder_lengths_batch[bf16-kv-v1-encoder-cuda-graph-fixed-shape-batch2-t5-small] + - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_fixed_encoder_lengths_batch[bf16-kv-v1-encoder-cuda-graph-fixed-shape-batch4-t5-small] - test_e2e.py::test_ptp_quickstart_advanced[Llama3.1-8B-NVFP4-nvfp4-quantized/Meta-Llama-3.1-8B] - test_e2e.py::test_ptp_quickstart_advanced[Llama3.1-8B-FP8-llama-3.1-model/Llama-3.1-8B-Instruct-FP8] - test_e2e.py::test_ptp_quickstart_advanced_mtp[DeepSeek-V3-Lite-BF16-DeepSeek-V3-Lite/bf16] diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index c26137c1ce09..d7cc8f22af78 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -179,8 +179,10 @@ l0_h100: - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_end_to_end[fp16-kv-v1-cuda-graph-off-greedy-bart-large-cnn] - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-off-beam2-bart-large-cnn] - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v2-cuda-graph-off-greedy-bart-large-cnn] - - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch[bf16-kv-v1-cuda-graph-off-greedy-batch2-bart-large-cnn] - - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch[bf16-kv-v2-cuda-graph-off-greedy-batch2-bart-large-cnn] + - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch[bf16-kv-v1-decoder-cuda-graph-on-greedy-batch2-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_generate_encoder_decoder_fixed_encoder_lengths_batch[bf16-kv-v1-encoder-cuda-graph-fixed-shape-batch2-bart-large-cnn] + - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_fixed_encoder_lengths_batch[bf16-kv-v1-encoder-cuda-graph-fixed-shape-batch4-bart-large-cnn] - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-off-beam2-t5-small0] - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-off-beam2-flan-t5-small] - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-off-beam2-t5-base] @@ -202,10 +204,12 @@ l0_h100: - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[fp16-kv-v2-cuda-graph-off-greedy-flan-t5-small] - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[fp32-kv-v2-cuda-graph-off-greedy-flan-t5-small] - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v2-cuda-graph-off-greedy-byt5-small] - - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch[bf16-kv-v1-cuda-graph-off-beam2-batch2-t5-small] - - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch[bf16-kv-v1-cuda-graph-off-beam2-batch2-flan-t5-small] - - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch[bf16-kv-v1-cuda-graph-off-greedy-batch2-t5-small] - - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch[bf16-kv-v2-cuda-graph-off-greedy-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-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-v1-decoder-cuda-graph-on-greedy-batch2-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_generate_encoder_decoder_fixed_encoder_lengths_batch[bf16-kv-v1-encoder-cuda-graph-fixed-shape-batch2-t5-small] + - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_fixed_encoder_lengths_batch[bf16-kv-v1-encoder-cuda-graph-fixed-shape-batch4-t5-small] - test_e2e.py::test_trtllm_bench_iteration_log[PyTorch-streaming-meta-llama/Llama-3.1-8B-llama-3.1-model/Meta-Llama-3.1-8B] - test_e2e.py::test_trtllm_bench_iteration_log[PyTorch-non-streaming-meta-llama/Llama-3.1-8B-llama-3.1-model/Meta-Llama-3.1-8B] - test_e2e.py::test_trtllm_bench_request_rate_and_concurrency[enable_concurrency-] diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index afb767029b1d..e92f4a673065 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -292,24 +292,6 @@ full:sm100/unittest/trt/quantization SKIP (Disable for Blackwell) full:sm100/unittest/trt/quantization/test_weight_only_groupwise_quant_matmul.py SKIP (Disable for Blackwell) full:sm100/unittest/trt/quantization/test_weight_only_quant_matmul.py SKIP (Disable for Blackwell) kv_cache/test_prefix_aware_scheduling.py::TestServePrefixAwareScheduling::test_multi_round_qa_shared_prefix_smoke SKIP (https://nvbugs/6266306) -llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-off-beam2-bart-large-cnn] SKIP (https://nvbugs/6340115) -llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-off-greedy-bart-large-cnn] SKIP (https://nvbugs/6340115) -llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_end_to_end[fp16-kv-v1-cuda-graph-off-greedy-bart-large-cnn] SKIP (https://nvbugs/6340115) -llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch[bf16-kv-v2-cuda-graph-off-greedy-batch2-bart-large-cnn] SKIP (https://nvbugs/6340115) -llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-off-beam2-flan-t5-xl] SKIP (https://nvbugs/6340115) -llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-off-beam2-t5-base] SKIP (https://nvbugs/6340115) -llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-off-beam2-t5-small0] SKIP (https://nvbugs/6340115) -llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-off-greedy-t5-small] SKIP (https://nvbugs/6340115) -llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v2-cuda-graph-off-greedy-flan-t5-small] SKIP (https://nvbugs/6340115) -llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v2-cuda-graph-off-greedy-t5-small] SKIP (https://nvbugs/6340115) -llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[fp16-kv-v1-cuda-graph-off-beam2-flan-t5-small] SKIP (https://nvbugs/6340115) -llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[fp16-kv-v2-cuda-graph-off-greedy-flan-t5-small] SKIP (https://nvbugs/6340115) -llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[fp32-kv-v1-cuda-graph-off-beam2-flan-t5-small] SKIP (https://nvbugs/6340115) -llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[fp32-kv-v1-cuda-graph-off-beam2-t5-small] SKIP (https://nvbugs/6340115) -llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[fp32-kv-v2-cuda-graph-off-greedy-t5-small] SKIP (https://nvbugs/6340115) -llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch[bf16-kv-v1-cuda-graph-off-beam2-batch2-flan-t5-small] SKIP (https://nvbugs/6340115) -llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch[bf16-kv-v1-cuda-graph-off-beam2-batch2-t5-small] SKIP (https://nvbugs/6340115) -llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch[bf16-kv-v2-cuda-graph-off-greedy-batch2-t5-small] SKIP (https://nvbugs/6340115) llmapi/test_llm_examples.py::test_llmapi_speculative_decoding_eagle3 SKIP (https://nvbugs/6075431) llmapi/test_llm_examples.py::test_llmapi_tensorrt_engine SKIP (https://nvbugs/5820553) perf/test_perf.py::test_perf[bart_large_cnn-plugin-float16-bs:8-input_output_len:60,20] SKIP # (https://nvidia.slack.com/archives/C059LSY62BT/p1704525727177449) From 5fcfc2915d8d262cec28aca3eb03923b7175b89f Mon Sep 17 00:00:00 2001 From: Guiju Zhang <7135567+cascade812@users.noreply.github.com> Date: Mon, 29 Jun 2026 14:29:20 -0700 Subject: [PATCH 5/9] disable encoder forward cuda graph support Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com> --- .../_torch/pyexecutor/model_engine.py | 503 +++++------------- tensorrt_llm/llmapi/__init__.py | 10 +- tensorrt_llm/llmapi/llm_args.py | 30 +- tensorrt_llm/llmapi/llm_utils.py | 11 +- .../defs/llmapi/test_llm_api_pytorch_bart.py | 223 +------- .../defs/llmapi/test_llm_api_pytorch_t5.py | 215 +------- .../test_lists/test-db/l0_b200.yml | 4 - .../test_lists/test-db/l0_h100.yml | 4 - 8 files changed, 157 insertions(+), 843 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index c340959e63f1..d5a7047d215d 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -28,7 +28,6 @@ create_input_processor_with_hash) from tensorrt_llm.llmapi.llm_args import (CudaGraphConfig, DecodingBaseConfig, EncodeCudaGraphConfig, - EncoderDecoderCudaGraphConfig, SeqLenAwareSparseAttentionConfig, TorchCompileConfig, TorchLlmArgs) from tensorrt_llm.logger import logger @@ -369,76 +368,44 @@ def __init__( self._init_model_capacity() 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) - # Split config into decoder and encoder parts. - # EncoderDecoderCudaGraphConfig: explicit sub-configs for each pass. - # EncodeCudaGraphConfig: encode-only models only; no decoder graphs. - # DecodeCudaGraphConfig / None: decoder only; no encoder graphs. - if isinstance(self.cuda_graph_config, EncoderDecoderCudaGraphConfig): - _decode_cfg = self.cuda_graph_config.decoder - _encode_cfg = self.cuda_graph_config.encoder - elif isinstance(self.cuda_graph_config, EncodeCudaGraphConfig): - if self._is_encoder_decoder_model(): - logger.warning( - "EncodeCudaGraphConfig is not supported for encoder-decoder " - "models. Use EncoderDecoderCudaGraphConfig instead. " - "Encoder and decoder CUDA graphs will be disabled.") - self.cuda_graph_config = None - _decode_cfg = None - _encode_cfg = None - else: - _decode_cfg = None - _encode_cfg = self.cuda_graph_config - else: - _decode_cfg = self.cuda_graph_config - _encode_cfg = None + 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.") + self.cuda_graph_config = None - # Decoder CUDA graph batch-size buckets. - cuda_graph_batch_sizes = _decode_cfg.batch_sizes if _decode_cfg else CudaGraphConfig.model_fields[ + cuda_graph_batch_sizes = self.cuda_graph_config.batch_sizes if self.cuda_graph_config else CudaGraphConfig.model_fields[ 'batch_sizes'].default - cuda_graph_padding_enabled = _decode_cfg.enable_padding if _decode_cfg else CudaGraphConfig.model_fields[ + cuda_graph_padding_enabled = self.cuda_graph_config.enable_padding if self.cuda_graph_config else CudaGraphConfig.model_fields[ 'enable_padding'].default - # Encoder CUDA graph batch-size buckets (may differ from decoder's). - encoder_cuda_graph_batch_sizes = _encode_cfg.batch_sizes if _encode_cfg else [] - # Encode-only CUDA graph detection. Decode configs do not define these # encoder-specific bucket fields. cuda_graph_num_tokens = [] cuda_graph_seq_lens = [] - if _encode_cfg is not None: - cuda_graph_num_tokens = _encode_cfg.num_tokens or [] - cuda_graph_seq_lens = _encode_cfg.seq_lens or [] + 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 [] - self._is_encode_only = (self.llm_args.encode_only - and not self.llm_args.mm_encoder_only) - - if ((self._is_encode_only or self._is_encoder_decoder_model()) - and _encode_cfg is not None + 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)): missing = [] if not cuda_graph_num_tokens: missing.append("num_tokens/max_num_token") if not cuda_graph_seq_lens: missing.append("seq_lens/max_seq_len") - if self._is_encode_only: - example = ( - "EncodeCudaGraphConfig(max_batch_size=64, " - "num_tokens=[128, 256, 512], " - "max_seq_len=128, enable_padding=True) for encode-only models" - ) - else: - example = ( - "EncoderDecoderCudaGraphConfig(" - "encoder=EncodeCudaGraphConfig(" - "max_batch_size=64, num_tokens=[128, 256, 512], " - "max_seq_len=128, enable_padding=True)) for encoder-decoder models" - ) logger.warning( - f"Encoder CUDA graph config is set, but " + 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"To enable them, specify e.g. {example}.") + 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).") self.torch_compile_config = self.llm_args.torch_compile_config torch_compile_enabled = bool(self.torch_compile_config is not None) @@ -600,10 +567,6 @@ def __init__( 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 - # Encoder pass may have independent padding setting (EncoderDecoderCudaGraphConfig). - self._encoder_cuda_graph_padding_enabled = ( - _encode_cfg.enable_padding - if _encode_cfg is not None else self._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, @@ -613,40 +576,20 @@ def __init__( self._max_cuda_graph_batch_size = (self._cuda_graph_batch_sizes[-1] if self._cuda_graph_batch_sizes else 0) - # Encoder batch-size buckets are independent from the decoder's. - self._encoder_cuda_graph_batch_sizes = _filter_cuda_graph_batch_sizes( - encoder_cuda_graph_batch_sizes, - self.batch_size, - self.max_num_tokens, - 0, # no draft tokens on the encoder pass - self._encoder_cuda_graph_padding_enabled - ) if encoder_cuda_graph_batch_sizes else [] - - self._max_encoder_cuda_graph_batch_size = ( - self._encoder_cuda_graph_batch_sizes[-1] - if self._encoder_cuda_graph_batch_sizes else 0) - # Encoder CUDA graph bucket lists self._cuda_graph_num_tokens = _filter_cuda_graph_num_tokens( cuda_graph_num_tokens, self.max_num_tokens, - self._encoder_cuda_graph_padding_enabled - ) if cuda_graph_num_tokens else [] + self._cuda_graph_padding_enabled) if 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. - _encoder_cuda_graph_padding_enabled) if cuda_graph_seq_lens else [] + cuda_graph_seq_lens, self.max_seq_len, + self._cuda_graph_padding_enabled) if 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) - self._enable_encoder_cuda_graph = ((self._is_encode_only - or self._is_encoder_decoder_model()) - and _encode_cfg 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( ) @@ -738,12 +681,15 @@ def __init__( # Create Encoder CUDA graph config and runner. encoder_cuda_graph_runner_config = EncoderCUDAGraphRunnerConfig( - use_cuda_graph=self._enable_encoder_cuda_graph, - cuda_graph_padding_enabled=self._encoder_cuda_graph_padding_enabled, - cuda_graph_batch_sizes=self._encoder_cuda_graph_batch_sizes, + 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, 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_encoder_cuda_graph_batch_size, + max_cuda_graph_batch_size=self._max_cuda_graph_batch_size, max_cuda_graph_num_tokens=self._max_cuda_graph_num_tokens, max_num_tokens=self.max_num_tokens, max_seq_len=self.max_seq_len, @@ -1064,8 +1010,6 @@ def warmup(self, resource_manager: ResourceManager) -> None: if self._is_encoder_decoder_model(): AutoTuner.get() - with self.encoder_cuda_graph_runner.allow_capture(): - self._run_cuda_graph_warmup_encoder_decoder() with self.cuda_graph_runner.allow_capture(): self._run_cuda_graph_warmup(resource_manager) return @@ -1490,7 +1434,7 @@ def _run_capture_pass(force_non_greedy: bool, label: str) -> None: batch, draft_len > 0, resource_manager) self.runtime_draft_len = draft_len if self._is_encoder_decoder_model(): - self._prepare_encoder_decoder_cuda_graph_warmup_batch( + self._prepare_decoder_cuda_graph_warmup_batch_for_encoder_decoder_model( batch, resource_manager) self.forward(batch, new_tensors_device=None, @@ -1912,7 +1856,7 @@ def _allocate_encoder_decoder_dummy_cross_kv( encoder_output_lens=encoder_output_lens) return cross_dummy_requests is not None - def _prepare_encoder_decoder_cuda_graph_warmup_batch( + def _prepare_decoder_cuda_graph_warmup_batch_for_encoder_decoder_model( self, batch: ScheduledRequests, resource_manager: ResourceManager) -> None: if not batch.generation_requests: @@ -4910,28 +4854,6 @@ def _create_encoder_warmup_inputs( } return inputs - def _create_encoder_decoder_encoder_warmup_inputs( - self, batch_size: int, num_tokens: int, - max_seq_len: int) -> Optional[Dict[str, Any]]: - raw_inputs = self._create_encoder_warmup_inputs(batch_size, num_tokens, - max_seq_len) - if raw_inputs is None: - return None - - sequence_lengths = raw_inputs['seq_lens'] - encoder_position_ids: List[int] = [] - for seq_len in sequence_lengths: - encoder_position_ids.extend( - self._apply_position_id_offset(list(range(seq_len)))) - - return self._prepare_encoder_decoder_encoder_inputs( - encoder_input_ids=raw_inputs['input_ids'], - encoder_position_ids=encoder_position_ids, - sequence_lengths=sequence_lengths, - request_ids=list(range(batch_size)), - resource_manager=None, - ) - @contextlib.contextmanager def no_encoder_cuda_graph(self): """Temporarily disable the encoder CUDA graph runner.""" @@ -5027,17 +4949,10 @@ def _run_cuda_graph_warmup_encoder(self) -> None: self._capture_encoder_cuda_graphs() - def _run_cuda_graph_warmup_encoder_decoder(self) -> None: - """Captures encoder CUDA graphs for encoder-decoder models.""" - if not self.encoder_cuda_graph_runner.enabled: - return - - self._capture_encoder_decoder_encoder_cuda_graphs() - def _get_encoder_cuda_graph_warmup_configs( self) -> List[Tuple[int, int, int]]: """Return feasible (batch_size, num_tokens, max_seq_len) graph keys.""" - batch_sizes = sorted(self._encoder_cuda_graph_batch_sizes, reverse=True) + batch_sizes = sorted(self._cuda_graph_batch_sizes, reverse=True) num_tokens_list = sorted(self._cuda_graph_num_tokens) seq_lens_list = sorted(self._cuda_graph_seq_lens) @@ -5058,7 +4973,7 @@ def _get_encoder_cuda_graph_warmup_configs( warmup_configs.append((bs, nt, sl)) - if self._encoder_cuda_graph_padding_enabled: + if self._cuda_graph_padding_enabled: for bs in batch_sizes: if bs > self.batch_size: continue @@ -5097,130 +5012,79 @@ def _capture_encoder_cuda_graphs(self) -> None: logger.info(f"Captured {num_captured} encoder CUDA graph(s).") - def _capture_encoder_decoder_encoder_cuda_graphs(self) -> None: - """Capture encoder-stack CUDA graphs for encoder-decoder models.""" - runner = self.encoder_cuda_graph_runner - if not runner.enabled: - return - - num_captured = 0 - logger.info("Capturing encoder-decoder encoder CUDA graphs ...") - for bs, nt, sl in self._get_encoder_cuda_graph_warmup_configs(): - inputs = self._create_encoder_decoder_encoder_warmup_inputs( - bs, nt, sl) - if inputs is None: - continue + @torch.inference_mode() + @with_model_extra_attrs(lambda self: self.model.extra_attrs) + @nvtx_range("encoder_forward") + def encoder_forward(self, inputs: Dict[str, Any], + **kwargs) -> Dict[str, Any]: + """Direct tensor-level forward for encode-only path. - logger.info(f"Encoder-decoder encoder CUDA graph capture: " - f"bs={bs}, nt={nt}, sl={sl}") - self._forward_encoder_with_cuda_graph(inputs) - torch.cuda.synchronize() - num_captured += 1 + Bypasses ScheduledRequests/LlmRequest entirely. Takes a raw inputs + dict, attempts encoder CUDA graph capture/replay if enabled, otherwise falls + back to eager execution. - logger.info( - f"Captured {num_captured} encoder-decoder encoder CUDA graph(s).") + Args: + inputs: Dict with 'input_ids' and 'seq_lens' (required), plus + any model-specific kwargs (token_type_ids, inputs_embeds, etc.). - def _run_encoder_forward_with_cuda_graph( - self, - runner_inputs: Dict[str, Any], - attn_metadata: Any, - prepare_model_inputs: Callable[ - [Dict[str, Any], Optional[Any], Optional[Any]], Dict[str, Any]], - forward_fn: Callable[[Dict[str, Any]], Any], - postprocess_graph_outputs: Callable[[Any], Any], - ) -> Any: - """Run a no-cache encoder forward with shared CUDA graph handling.""" + Returns: + Dict with 'logits' tensor and any other model outputs. + """ moe_load_balancer: MoeLoadBalancer = getattr(self, 'moe_load_balancer', None) - batch_size = len(runner_inputs['seq_lens']) + batch_size = len(inputs['seq_lens']) with self.encoder_cuda_graph_runner.pad_batch( - runner_inputs, batch_size) as padded_inputs: + inputs, batch_size) as padded_inputs: + attn_metadata = self._set_up_attn_metadata( + kv_cache_manager=None + ) if self.encoder_attn_metadata is None else self.encoder_attn_metadata graph_attn_metadata, key = self.encoder_cuda_graph_runner.maybe_get_cuda_graph( padded_inputs, attn_metadata) + # Unpad seq_lens when fallback to eager path. + if key is None: + padded_inputs['seq_lens'] = padded_inputs[ + 'seq_lens'][:batch_size] + model_inputs = self._prepare_encoder_inputs( + padded_inputs, + attn_metadata=graph_attn_metadata, + padded_num_tokens=key[1] if key is not None else None) + forward_kwargs = { + "gather_ids": None, + "gather_context_logits": False, + **kwargs, + } with with_shared_pool( self.encoder_cuda_graph_runner.get_graph_pool()): if key is None: - eager_inputs = dict(padded_inputs) - eager_inputs['seq_lens'] = eager_inputs[ - 'seq_lens'][:batch_size] - model_inputs = prepare_model_inputs(eager_inputs, None, - None) with MoeLoadBalancerIterContext(moe_load_balancer): - return forward_fn(model_inputs) + # Eager path — no graph for this bucket. + return self._forward_step(model_inputs, + **forward_kwargs) - model_inputs = prepare_model_inputs(padded_inputs, - graph_attn_metadata, key) if self.encoder_cuda_graph_runner.needs_capture(key): - def capture_forward_fn( - capture_inputs: Dict[str, Any]) -> Any: + def forward_fn( + capture_inputs: Dict[str, Any]) -> Dict[str, Any]: + capture_inputs = capture_inputs.copy() + forward_kwargs = capture_inputs.pop("_forward_kwargs") with MoeLoadBalancerIterContext(moe_load_balancer): - return forward_fn(capture_inputs) + return self._forward_step(capture_inputs, + **forward_kwargs) self.encoder_cuda_graph_runner.capture( - key, capture_forward_fn, model_inputs) + key, forward_fn, { + **model_inputs, "_forward_kwargs": forward_kwargs + }) with MoeLoadBalancerIterContext(moe_load_balancer): graph_outputs = self.encoder_cuda_graph_runner.replay( - key, model_inputs) - - return postprocess_graph_outputs(graph_outputs) - - @torch.inference_mode() - @with_model_extra_attrs(lambda self: self.model.extra_attrs) - @nvtx_range("encoder_forward") - def encoder_forward(self, inputs: Dict[str, Any], - **kwargs) -> Dict[str, Any]: - """Run the tensor-level encode-only forward path. - - This is the ``LLM.encode()``/``encode_only`` path. It bypasses - ``ScheduledRequests`` and ``LlmRequest`` entirely, consumes a raw - tensor-style inputs dict, and returns model outputs such as logits. - It is separate from :meth:`forward_encoder`, which consumes scheduled - encoder-decoder requests and returns encoder hidden states for the - decoder. - - Attempts encoder CUDA graph capture/replay when enabled; otherwise - runs the same forward eagerly. - - Args: - inputs: Dict with 'input_ids' and 'seq_lens' (required), plus - any model-specific kwargs (token_type_ids, inputs_embeds, etc.). - - Returns: - Dict with 'logits' tensor and any other model outputs. - """ - batch_size = len(inputs['seq_lens']) - attn_metadata = self._set_up_attn_metadata( - kv_cache_manager=None - ) if self.encoder_attn_metadata is None else self.encoder_attn_metadata - forward_kwargs = { - "gather_ids": None, - "gather_context_logits": False, - **kwargs, - } - - def prepare_model_inputs(padded_inputs: Dict[str, Any], - graph_attn_metadata: Any, - key: Optional[Any]) -> Dict[str, Any]: - model_inputs = self._prepare_encoder_inputs( - padded_inputs, - attn_metadata=graph_attn_metadata, - padded_num_tokens=key[1] if key is not None else None) - return { - **model_inputs, - "_forward_kwargs": forward_kwargs, - } - - def forward_fn(model_inputs: Dict[str, Any]) -> Dict[str, Any]: - model_inputs = model_inputs.copy() - model_forward_kwargs = model_inputs.pop("_forward_kwargs") - return self._forward_step(model_inputs, **model_forward_kwargs) + key, { + **model_inputs, "_forward_kwargs": forward_kwargs + }) - def postprocess_graph_outputs( - graph_outputs: Dict[str, Any]) -> Dict[str, Any]: + # Return a clone to avoid sharing data_ptr with the static buffers. outputs = {} for name, value in graph_outputs.items(): if isinstance(value, torch.Tensor): @@ -5232,14 +5096,6 @@ def postprocess_graph_outputs( return outputs - return self._run_encoder_forward_with_cuda_graph( - runner_inputs=inputs, - attn_metadata=attn_metadata, - prepare_model_inputs=prepare_model_inputs, - forward_fn=forward_fn, - postprocess_graph_outputs=postprocess_graph_outputs, - ) - @torch.inference_mode() @with_model_extra_attrs(lambda self: self.model.extra_attrs) def forward(self, @@ -5596,22 +5452,50 @@ def _forward_step_mm_encoder_only( return result - def _prepare_encoder_decoder_encoder_inputs( + @nvtx_range("_prepare_tp_inputs_encoder") + def _prepare_tp_inputs_encoder( self, - encoder_input_ids: List[int], - encoder_position_ids: List[int], - sequence_lengths: List[int], - request_ids: List[int], + encoder_requests: List[LlmRequest], resource_manager: Optional[ResourceManager] = None, - ) -> Dict[str, Any]: - if len(sequence_lengths) != len(request_ids): - raise ValueError("Encoder sequence lengths and request IDs must " - "have the same length.") + ): + """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) 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})") @@ -5646,7 +5530,7 @@ def _prepare_encoder_decoder_encoder_inputs( dtype=torch.int, pin_memory=prefer_pinned(), ) - encoder_attn_metadata.num_contexts = len(sequence_lengths) + encoder_attn_metadata.num_contexts = len(encoder_requests) encoder_attn_metadata.max_seq_len = self.max_seq_len encoder_attn_metadata.request_ids = request_ids encoder_attn_metadata.prepare() @@ -5658,7 +5542,7 @@ def _prepare_encoder_decoder_encoder_inputs( dtype=torch.int, pin_memory=prefer_pinned()) - return { + inputs = { 'encoder_input_ids': encoder_input_ids_t.to('cuda', non_blocking=True), 'encoder_position_ids': @@ -5667,64 +5551,10 @@ def _prepare_encoder_decoder_encoder_inputs( encoder_attn_metadata, 'encoder_seq_lens': sequence_lengths, - 'encoder_input_ids_host': - encoder_input_ids, - 'encoder_position_ids_host': - encoder_position_ids, 'resource_manager': resource_manager, } - - @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, - ) + return inputs @nvtx_range("_forward_step_encoder") def _forward_step_encoder( @@ -5779,91 +5609,13 @@ 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: - """Run the encoder stack, replaying an encoder CUDA graph when possible.""" - 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) - - if input_ids is None or position_ids is None: - moe_load_balancer: MoeLoadBalancer = getattr( - self, 'moe_load_balancer', None) - with MoeLoadBalancerIterContext(moe_load_balancer): - 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'), - } - - def prepare_model_inputs(padded_inputs: Dict[str, Any], - graph_attn_metadata: Any, - key: Optional[Any]) -> Dict[str, Any]: - if key is None: - return inputs - - padded_num_tokens = key[1] - graph_attn_metadata.prepare_encoder_cuda_graph_replay( - padded_inputs['seq_lens'], padded_num_tokens) - return { - **padded_inputs, - 'attn_metadata': graph_attn_metadata, - } - - def forward_fn(model_inputs: Dict[str, Any]) -> torch.Tensor: - if 'attn_metadata' in model_inputs: - return self._forward_step_encoder_cuda_graph(model_inputs) - return self._forward_step_encoder(model_inputs) - - def postprocess_graph_outputs(graph_outputs: Any) -> torch.Tensor: - 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 self._run_encoder_forward_with_cuda_graph( - runner_inputs=runner_inputs, - attn_metadata=inputs['encoder_attn_metadata'], - prepare_model_inputs=prepare_model_inputs, - forward_fn=forward_fn, - postprocess_graph_outputs=postprocess_graph_outputs, - ) - @nvtx_range("forward_encoder") def forward_encoder( self, encoder_requests: List[LlmRequest], resource_manager: Optional[ResourceManager] = None, ) -> Tuple[torch.Tensor, List[int]]: - """Run the encoder-decoder encoder pass for scheduled requests. - - This is the encoder-side prepass for T5/BART-style - encoder-decoder generation. It consumes ``LlmRequest`` objects with - ``encoder_tokens``, packs them into one no-cache encoder batch, and - returns hidden states for decoder cross-attention. It is separate from - :meth:`encoder_forward`, which is the raw tensor ``LLM.encode()`` path - and returns model outputs such as logits. + """Run the encoder stack for ``encoder_requests``. Returns a tuple ``(encoder_hidden_states, encoder_seq_lens)`` where the hidden states tensor is shaped @@ -5883,8 +5635,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._forward_step_encoder(inputs) return encoder_hidden_states, inputs['encoder_seq_lens'] diff --git a/tensorrt_llm/llmapi/__init__.py b/tensorrt_llm/llmapi/__init__.py index ea8d8abbdf6f..96032c23c227 100644 --- a/tensorrt_llm/llmapi/__init__.py +++ b/tensorrt_llm/llmapi/__init__.py @@ -14,11 +14,10 @@ DeepSeekV4SparseAttentionConfig, DFlashDecodingConfig, DraftTargetDecodingConfig, DynamicBatchConfig, Eagle3DecodingConfig, EagleDecodingConfig, - EncodeCudaGraphConfig, EncoderDecoderCudaGraphConfig, - ExtendedRuntimePerfKnobConfig, KvCacheConfig, LlmArgs, - LookaheadDecodingConfig, MedusaDecodingConfig, - MiniMaxM3SparseAttentionConfig, MoeConfig, - MTPDecodingConfig, NGramDecodingConfig, + EncodeCudaGraphConfig, ExtendedRuntimePerfKnobConfig, + KvCacheConfig, LlmArgs, LookaheadDecodingConfig, + MedusaDecodingConfig, MiniMaxM3SparseAttentionConfig, + MoeConfig, MTPDecodingConfig, NGramDecodingConfig, PARDDecodingConfig, PrometheusMetricsConfig, ReorderRequestPolicyConfig, RocketSparseAttentionConfig, SADecodingConfig, SAEnhancerConfig, @@ -47,7 +46,6 @@ 'CudaGraphConfig', 'DecodeCudaGraphConfig', 'EncodeCudaGraphConfig', - 'EncoderDecoderCudaGraphConfig', 'MoeConfig', 'LookaheadDecodingConfig', 'MedusaDecodingConfig', diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 0feaca79ebe9..6685b95993a6 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -452,30 +452,11 @@ def _generate_cuda_graph_seq_lens(max_seq_len: int, return sizes -class EncoderDecoderCudaGraphConfig(StrictBaseModel): - """CUDA graph configuration for encoder-decoder models. - - Provides independent configs for the encoder and decoder passes, allowing - separate control of batch sizes, token buckets, and padding for each. - """ - - mode: Literal["encoder_decoder"] = Field( - default="encoder_decoder", description="CUDA graph configuration mode.") - - encoder: EncodeCudaGraphConfig = Field( - description="CUDA graph configuration for the encoder pass.") - - decoder: DecodeCudaGraphConfig = Field( - default_factory=DecodeCudaGraphConfig, - description="CUDA graph configuration for the decoder pass.") - - # For CudaGraphConfig's backward compatibility CudaGraphConfig = DecodeCudaGraphConfig CudaGraphConfigType: TypeAlias = Annotated[ - Union[DecodeCudaGraphConfig, EncodeCudaGraphConfig, - EncoderDecoderCudaGraphConfig], + Union[DecodeCudaGraphConfig, EncodeCudaGraphConfig], Field(discriminator="mode"), ] @@ -4606,14 +4587,17 @@ class TorchLlmArgs(BaseLlmArgs): @field_validator('cuda_graph_config', mode='before') @classmethod def infer_cuda_graph_config_mode(cls, v): + if isinstance(v, dict) and ("encoder" in v or "decoder" in v): + raise ValueError( + "Composite encoder-decoder CUDA graph configs have been " + "removed. Use CudaGraphConfig or DecodeCudaGraphConfig for " + "decoder CUDA graphs.") if isinstance(v, dict) and "mode" not in v: encoder_keys = { "num_tokens", "max_num_token", "seq_lens", "max_seq_len" } v = dict(v) - if "encoder" in v or "decoder" in v: - v["mode"] = "encoder_decoder" - elif any(k in v and v[k] not in (None, 0) for k in encoder_keys): + if any(k in v and v[k] not in (None, 0) for k in encoder_keys): v["mode"] = "encode" else: v["mode"] = "decode" diff --git a/tensorrt_llm/llmapi/llm_utils.py b/tensorrt_llm/llmapi/llm_utils.py index ce6c4016eeb0..e01043406944 100644 --- a/tensorrt_llm/llmapi/llm_utils.py +++ b/tensorrt_llm/llmapi/llm_utils.py @@ -38,11 +38,11 @@ from .llm_args import (CalibConfig, CudaGraphConfig, DecodeCudaGraphConfig, DraftTargetDecodingConfig, Eagle3DecodingConfig, EagleDecodingConfig, EncodeCudaGraphConfig, - EncoderDecoderCudaGraphConfig, KvCacheConfig, LlmArgs, - LookaheadDecodingConfig, MedusaDecodingConfig, - MTPDecodingConfig, NGramDecodingConfig, SchedulerConfig, - TorchLlmArgs, UserProvidedDecodingConfig, - _ModelFormatKind, _ModelWrapper, _ParallelConfig, + KvCacheConfig, LlmArgs, LookaheadDecodingConfig, + MedusaDecodingConfig, MTPDecodingConfig, + NGramDecodingConfig, SchedulerConfig, TorchLlmArgs, + UserProvidedDecodingConfig, _ModelFormatKind, + _ModelWrapper, _ParallelConfig, update_llm_args_with_extra_dict, update_llm_args_with_extra_options) # yapf: enable @@ -1058,7 +1058,6 @@ class LlmBuildStats: 'CudaGraphConfig', 'DecodeCudaGraphConfig', 'EncodeCudaGraphConfig', - 'EncoderDecoderCudaGraphConfig', 'KvCacheConfig', 'CachedModelLoader', 'EagleDecodingConfig', 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 605997c7ad51..126342be75d5 100644 --- a/tests/integration/defs/llmapi/test_llm_api_pytorch_bart.py +++ b/tests/integration/defs/llmapi/test_llm_api_pytorch_bart.py @@ -21,9 +21,6 @@ from tensorrt_llm.llmapi import ( LLM, CudaGraphConfig, - DecodeCudaGraphConfig, - EncodeCudaGraphConfig, - EncoderDecoderCudaGraphConfig, KvCacheConfig, RequestOutput, SamplingParams, @@ -43,34 +40,6 @@ "the center adds classrooms, simulation labs, and career coaching for workers." ), ] -_FIXED_SEQ_LENS_SOURCE_TEXTS = [ - ( - "Summarize: The engineering team released a faster inference service on Monday. " - "The update improves batching, lowers latency, and adds detailed monitoring for operators." - ), - ( - "Summarize: The city opened a new public library on Monday. Residents said the library " - "has quiet rooms, computer access, and a large children section." - ), - ( - "Summarize: The company opened a training center on Monday. Managers said the center " - "adds classrooms, simulation labs, and career coaching for local workers." - ), - ( - "Summarize: The museum opened a science exhibit on Monday. Curators said the exhibit " - "has space models, interactive lessons, and workshops for local children." - ), -] -_FIXED_SEQ_LENS_EXPECTED_TOKEN_IDS_BY_REQUEST = [ - # "The update improves batching, lowers latency" - [[0, 133, 2935, 15296, 14398, 154, 6, 32222, 35940, 2]], - # "Residents said the library has quiet rooms," - [[0, 35129, 26, 5, 5560, 34, 5128, 5351, 6, 2]], - # "The company opened a training center on Monday" - [[0, 133, 138, 1357, 10, 1058, 1312, 15, 302, 2]], - # "The museum opened a science exhibit on Monday" - [[0, 133, 5707, 1357, 10, 2866, 8483, 15, 302, 2]], -] _MODEL_NAME = "bart-large-cnn" _MAX_NEW_TOKENS = 10 _MAX_SEQUENCE_LENGTH = 128 @@ -80,7 +49,6 @@ _CROSS_KV_CACHE_FRACTION = 0.5 # "The update improves batching, lowers latency" _EXPECTED_GREEDY_OUTPUT_TOKEN_IDS = [0, 133, 2935, 15296, 14398, 154, 6, 32222, 35940, 2] -_FIXED_SEQ_LENS_MAX_NEW_TOKENS = 10 _EXPECTED_BEAM_OUTPUT_TOKEN_IDS_BY_BEAMS = { 2: [ # "The update improves batching, lowers latency" @@ -212,24 +180,6 @@ def _mixed_batch_test_case( ) -def _fixed_seq_lens_shape_test_case( - torch_dtype: str, - use_kv_cache_manager_v2: bool, - batch_size: int, - feature_id: str, - kv_cache_dtype: str = "auto", -): - return pytest.param( - torch_dtype, - use_kv_cache_manager_v2, - 1, - 1, - batch_size, - kv_cache_dtype, - id=f"{feature_id}-{_MODEL_NAME}", - ) - - _MIXED_BATCH_TEST_CASES = [ _mixed_batch_test_case( torch_dtype="bfloat16", @@ -248,21 +198,6 @@ def _fixed_seq_lens_shape_test_case( ] -_FIXED_SEQ_LENS_SHAPE_TEST_CASES = [ - _fixed_seq_lens_shape_test_case( - torch_dtype="bfloat16", - use_kv_cache_manager_v2=False, - batch_size=2, - feature_id="bf16-kv-v1-encoder-cuda-graph-fixed-shape-batch2", - ), - _fixed_seq_lens_shape_test_case( - torch_dtype="bfloat16", - use_kv_cache_manager_v2=False, - batch_size=4, - feature_id="bf16-kv-v1-encoder-cuda-graph-fixed-shape-batch4", - ), -] - pytestmark = [ pytest.mark.skip_less_device(1), pytest.mark.skip_less_device_memory(_MIN_GPU_MEMORY_MB), @@ -299,27 +234,6 @@ def _sampling_params(num_beams: int, num_return_sequences: int) -> SamplingParam ) -def _cuda_graph_config( - enabled: bool, - num_tokens: list[int], - seq_lens: list[int], - batch_sizes: list[int] | None = None, -) -> EncoderDecoderCudaGraphConfig | None: - if not enabled: - return None - return EncoderDecoderCudaGraphConfig( - encoder=EncodeCudaGraphConfig( - num_tokens=num_tokens, - seq_lens=seq_lens, - enable_padding=True, - ), - decoder=DecodeCudaGraphConfig( - batch_sizes=batch_sizes or [1], - enable_padding=True, - ), - ) - - def _decoder_cuda_graph_config( batch_sizes: list[int] | None = None, ) -> CudaGraphConfig: @@ -332,7 +246,7 @@ def _decoder_cuda_graph_config( ) -def _assert_encoder_decoder_cuda_graph_state( +def _assert_decoder_cuda_graph_state( llm: LLM, enabled: bool, batch_sizes: list[int] | None, @@ -346,20 +260,11 @@ def _assert_encoder_decoder_cuda_graph_state( assert not model_engine.cuda_graph_runner.graphs return - _assert_encoder_decoder_cuda_graphs_captured(llm) + _assert_decoder_cuda_graphs_captured(llm) if batch_sizes is not None: assert model_engine.cuda_graph_runner.padding_dummy_requests -def _assert_encoder_decoder_cuda_graphs_captured(llm: LLM) -> None: - model_engine = llm._executor.engine.model_engine - - assert model_engine.encoder_cuda_graph_runner.enabled - assert model_engine.encoder_cuda_graph_runner.graphs - assert model_engine.cuda_graph_runner.enabled - assert model_engine.cuda_graph_runner.graphs - - def _assert_decoder_cuda_graphs_captured(llm: LLM) -> None: model_engine = llm._executor.engine.model_engine @@ -439,8 +344,6 @@ def _run_bart_pytorch_generate_encoder_decoder( model_path = _get_bart_model_path() tokenizer = AutoTokenizer.from_pretrained(model_path) - encoder_seq_len = len(tokenizer(_SOURCE_TEXT).input_ids) - encoder_num_tokens = encoder_seq_len case_id = ( f"model={_MODEL_NAME}, dtype={torch_dtype}, kv_v2={use_kv_cache_manager_v2}, " f"cuda_graph={enable_cuda_graph}, beams={num_beams}, returns={num_return_sequences}, " @@ -452,12 +355,9 @@ def _run_bart_pytorch_generate_encoder_decoder( model_path, backend="pytorch", attn_backend="TRTLLM", - cuda_graph_config=_cuda_graph_config( - enable_cuda_graph, - [encoder_num_tokens], - [encoder_seq_len], - cuda_graph_batch_sizes, - ), + cuda_graph_config=_decoder_cuda_graph_config(cuda_graph_batch_sizes) + if enable_cuda_graph + else None, disable_overlap_scheduler=True, dtype=torch_dtype, enable_chunked_prefill=False, @@ -493,7 +393,7 @@ def _run_bart_pytorch_generate_encoder_decoder( exact_match, expected_output_token_ids_by_output, ) - _assert_encoder_decoder_cuda_graph_state( + _assert_decoder_cuda_graph_state( llm, enable_cuda_graph, cuda_graph_batch_sizes, @@ -532,117 +432,6 @@ def test_bart_pytorch_generate_encoder_decoder_end_to_end( ) -@pytest.mark.parametrize( - "torch_dtype,use_kv_cache_manager_v2,num_beams,num_return_sequences,batch_size,kv_cache_dtype", - _FIXED_SEQ_LENS_SHAPE_TEST_CASES, -) -def test_bart_pytorch_generate_encoder_decoder_fixed_encoder_lengths_batch( - monkeypatch: pytest.MonkeyPatch, - torch_dtype: str, - use_kv_cache_manager_v2: bool, - num_beams: int, - num_return_sequences: int, - batch_size: int, - kv_cache_dtype: str, -) -> None: - monkeypatch.setenv("TLLM_WORKER_USE_SINGLE_PROCESS", "1") - monkeypatch.setenv("TRTLLM_SKIP_KV_CACHE_ESTIMATION", "1") - - model_path = _get_bart_model_path() - tokenizer = AutoTokenizer.from_pretrained(model_path) - source_texts = _FIXED_SEQ_LENS_SOURCE_TEXTS[:batch_size] - assert len(source_texts) == batch_size - expected_token_ids_by_request = _FIXED_SEQ_LENS_EXPECTED_TOKEN_IDS_BY_REQUEST[:batch_size] - encoder_seq_lens_by_request = [ - len(tokenizer(source_text).input_ids) for source_text in source_texts - ] - assert len(set(encoder_seq_lens_by_request)) == 1 - encoder_seq_len = encoder_seq_lens_by_request[0] - encoder_num_tokens = sum(encoder_seq_lens_by_request) - graph_key = (batch_size, encoder_num_tokens, encoder_seq_len) - assert num_beams == 1 - assert num_return_sequences == 1 - sampling_params = SamplingParams( - max_tokens=_FIXED_SEQ_LENS_MAX_NEW_TOKENS, - temperature=0.0, - ) - max_num_tokens = max(_MAX_SEQUENCE_LENGTH, encoder_num_tokens) - kv_cache_max_tokens = max( - _MAX_KV_TOKENS, - batch_size * _MAX_SEQUENCE_LENGTH * 2, - ) - case_id = ( - f"model={_MODEL_NAME}, dtype={torch_dtype}, kv_v2={use_kv_cache_manager_v2}, " - f"encoder_cuda_graph=True, beams={num_beams}, returns={num_return_sequences}, " - f"kv_dtype={kv_cache_dtype}, fixed_encoder_seq_lens=True, batch_size={batch_size}" - ) - llm_kwargs = { - "backend": "pytorch", - "attn_backend": "TRTLLM", - "disable_overlap_scheduler": True, - "dtype": torch_dtype, - "enable_chunked_prefill": False, - "kv_cache_config": KvCacheConfig( - enable_block_reuse=False, - max_tokens=kv_cache_max_tokens, - free_gpu_memory_fraction=_FREE_GPU_MEMORY_FRACTION, - cross_kv_cache_fraction=_CROSS_KV_CACHE_FRACTION, - use_kv_cache_manager_v2=use_kv_cache_manager_v2, - dtype=kv_cache_dtype, - ), - "max_batch_size": batch_size, - "max_beam_width": num_beams, - "max_input_len": _MAX_SEQUENCE_LENGTH, - "max_num_tokens": max_num_tokens, - "max_seq_len": _MAX_SEQUENCE_LENGTH, - "model_kwargs": {"torch_dtype": torch_dtype}, - "scheduler_config": SchedulerConfig(use_python_scheduler=True), - } - - with LLM( - model_path, - cuda_graph_config=_cuda_graph_config( - True, - [encoder_num_tokens], - [encoder_seq_len], - [batch_size], - ), - **llm_kwargs, - ) as llm: - model_engine = llm._executor.engine.model_engine - assert graph_key in model_engine.encoder_cuda_graph_runner.graphs - - responses = llm.generate( - source_texts, - sampling_params=sampling_params, - use_tqdm=False, - ) - assert len(responses) == batch_size - - token_ids_by_request = [] - for request_idx, response in enumerate(responses): - token_ids = _assert_bart_response( - response, - num_return_sequences=num_return_sequences, - max_tokens=_FIXED_SEQ_LENS_MAX_NEW_TOKENS, - ) - token_ids_by_request.append(token_ids) - _print_generated_text( - tokenizer, - f"{case_id}, request={request_idx}", - "output", - token_ids, - ) - _assert_expected_generation( - tokenizer, - token_ids, - exact_match=True, - expected_token_ids_by_output=expected_token_ids_by_request[request_idx], - ) - - assert token_ids_by_request == expected_token_ids_by_request - - @pytest.mark.parametrize( "torch_dtype,use_kv_cache_manager_v2,num_beams,num_return_sequences", _MIXED_BATCH_TEST_CASES, 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 1e1e6e8a3e42..5ea4cffdcf25 100644 --- a/tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py +++ b/tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py @@ -22,9 +22,6 @@ from tensorrt_llm.llmapi import ( LLM, CudaGraphConfig, - DecodeCudaGraphConfig, - EncodeCudaGraphConfig, - EncoderDecoderCudaGraphConfig, KvCacheConfig, RequestOutput, SamplingParams, @@ -38,18 +35,6 @@ _SOURCE_TEXT, "translate English to German: The book is on the table.", ] -_FIXED_SEQ_LENS_SOURCE_TEXTS = [ - "translate English to German: The book is on table.", - "translate English to German: The school is very clean.", - "translate English to German: The car is very small.", - "translate English to German: The food is very good.", -] -_FIXED_SEQ_LENS_EXPECTED_TOKEN_IDS_BY_REQUEST = [ - [644, 4675, 229, 219], - [316, 12853, 229, 1319], - [644, 2040, 229, 1319], - [644, 11722, 229, 1319], -] _MAX_NEW_TOKENS = 4 _MAX_SEQUENCE_LENGTH = 64 _MAX_KV_TOKENS = 256 @@ -114,24 +99,6 @@ _MIXED_CONTEXT_GENERATION_MAX_NEW_TOKENS = 8 -def _fixed_seq_lens_shape_test_case( - model_name: str, - torch_dtype: str, - use_kv_cache_manager_v2: bool, - batch_size: int, - feature_id: str, -): - return pytest.param( - model_name, - torch_dtype, - use_kv_cache_manager_v2, - 1, - 1, - batch_size, - id=f"{feature_id}-{model_name}", - ) - - def _test_case( model_name: str, torch_dtype: str, @@ -504,23 +471,6 @@ def _mixed_batch_test_case( ] -_FIXED_SEQ_LENS_SHAPE_TEST_CASES = [ - _fixed_seq_lens_shape_test_case( - model_name="t5-small", - torch_dtype="bfloat16", - use_kv_cache_manager_v2=False, - batch_size=2, - feature_id="bf16-kv-v1-encoder-cuda-graph-fixed-shape-batch2", - ), - _fixed_seq_lens_shape_test_case( - model_name="t5-small", - torch_dtype="bfloat16", - use_kv_cache_manager_v2=False, - batch_size=4, - feature_id="bf16-kv-v1-encoder-cuda-graph-fixed-shape-batch4", - ), -] - pytestmark = [ pytest.mark.skip_less_device(1), pytest.mark.skip_less_device_memory(_MIN_GPU_MEMORY_MB), @@ -557,27 +507,6 @@ def _sampling_params(num_beams: int, num_return_sequences: int) -> SamplingParam ) -def _cuda_graph_config( - enabled: bool, - num_tokens: list[int], - seq_lens: list[int], - batch_sizes: list[int] | None = None, -) -> EncoderDecoderCudaGraphConfig | None: - if not enabled: - return None - return EncoderDecoderCudaGraphConfig( - encoder=EncodeCudaGraphConfig( - num_tokens=num_tokens, - seq_lens=seq_lens, - enable_padding=True, - ), - decoder=DecodeCudaGraphConfig( - batch_sizes=batch_sizes or [1], - enable_padding=True, - ), - ) - - def _decoder_cuda_graph_config( batch_sizes: list[int] | None = None, ) -> CudaGraphConfig: @@ -590,7 +519,7 @@ def _decoder_cuda_graph_config( ) -def _assert_encoder_decoder_cuda_graph_state( +def _assert_decoder_cuda_graph_state( llm: LLM, enabled: bool, batch_sizes: list[int] | None, @@ -604,20 +533,11 @@ def _assert_encoder_decoder_cuda_graph_state( assert not model_engine.cuda_graph_runner.graphs return - _assert_encoder_decoder_cuda_graphs_captured(llm) + _assert_decoder_cuda_graphs_captured(llm) if batch_sizes is not None: assert model_engine.cuda_graph_runner.padding_dummy_requests -def _assert_encoder_decoder_cuda_graphs_captured(llm: LLM) -> None: - model_engine = llm._executor.engine.model_engine - - assert model_engine.encoder_cuda_graph_runner.enabled - assert model_engine.encoder_cuda_graph_runner.graphs - assert model_engine.cuda_graph_runner.enabled - assert model_engine.cuda_graph_runner.graphs - - def _assert_decoder_cuda_graphs_captured(llm: LLM) -> None: model_engine = llm._executor.engine.model_engine @@ -628,7 +548,7 @@ def _assert_decoder_cuda_graphs_captured(llm: LLM) -> None: def _assert_mixed_context_generation_cuda_graph_state(llm: LLM) -> None: - _assert_encoder_decoder_cuda_graphs_captured(llm) + _assert_decoder_cuda_graphs_captured(llm) assert llm._executor.engine.model_engine.cuda_graph_runner.padding_dummy_requests @@ -713,8 +633,6 @@ def _run_t5_pytorch_generate_encoder_decoder( model_path = _get_t5_model_path(model_name) tokenizer = AutoTokenizer.from_pretrained(model_path) - encoder_seq_len = len(tokenizer(_SOURCE_TEXT).input_ids) - encoder_num_tokens = encoder_seq_len case_id = ( f"model={model_name}, dtype={torch_dtype}, kv_v2={use_kv_cache_manager_v2}, " f"cuda_graph={enable_cuda_graph}, beams={num_beams}, returns={num_return_sequences}" @@ -725,12 +643,9 @@ def _run_t5_pytorch_generate_encoder_decoder( model_path, backend="pytorch", attn_backend="TRTLLM", - cuda_graph_config=_cuda_graph_config( - enable_cuda_graph, - [encoder_num_tokens], - [encoder_seq_len], - cuda_graph_batch_sizes, - ), + cuda_graph_config=_decoder_cuda_graph_config(cuda_graph_batch_sizes) + if enable_cuda_graph + else None, disable_overlap_scheduler=True, dtype=torch_dtype, enable_chunked_prefill=False, @@ -765,7 +680,7 @@ def _run_t5_pytorch_generate_encoder_decoder( exact_match, expected_output_token_ids_by_output, ) - _assert_encoder_decoder_cuda_graph_state( + _assert_decoder_cuda_graph_state( llm, enable_cuda_graph, cuda_graph_batch_sizes, @@ -803,108 +718,6 @@ def test_t5_pytorch_generate_encoder_decoder_end_to_end( ) -@pytest.mark.parametrize( - "model_name,torch_dtype,use_kv_cache_manager_v2,num_beams,num_return_sequences,batch_size", - _FIXED_SEQ_LENS_SHAPE_TEST_CASES, -) -def test_t5_pytorch_generate_encoder_decoder_fixed_encoder_lengths_batch( - monkeypatch: pytest.MonkeyPatch, - model_name: str, - torch_dtype: str, - use_kv_cache_manager_v2: bool, - num_beams: int, - num_return_sequences: int, - batch_size: int, -) -> None: - monkeypatch.setenv("TLLM_WORKER_USE_SINGLE_PROCESS", "1") - monkeypatch.setenv("TRTLLM_SKIP_KV_CACHE_ESTIMATION", "1") - - model_path = _get_t5_model_path(model_name) - tokenizer = AutoTokenizer.from_pretrained(model_path) - source_texts = _FIXED_SEQ_LENS_SOURCE_TEXTS[:batch_size] - assert len(source_texts) == batch_size - expected_token_ids_by_request = _FIXED_SEQ_LENS_EXPECTED_TOKEN_IDS_BY_REQUEST[:batch_size] - encoder_seq_lens_by_request = [ - len(tokenizer(source_text).input_ids) for source_text in source_texts - ] - assert len(set(encoder_seq_lens_by_request)) == 1 - encoder_seq_len = encoder_seq_lens_by_request[0] - encoder_num_tokens = sum(encoder_seq_lens_by_request) - sampling_params = _sampling_params(num_beams, num_return_sequences) - max_num_tokens = max( - _MAX_SEQUENCE_LENGTH, - encoder_num_tokens + batch_size * (_MAX_NEW_TOKENS + 1), - ) - case_id = ( - f"model={model_name}, dtype={torch_dtype}, kv_v2={use_kv_cache_manager_v2}, " - f"encoder_cuda_graph=True, beams={num_beams}, returns={num_return_sequences}, " - f"fixed_encoder_seq_lens=True, batch_size={batch_size}" - ) - llm_kwargs = { - "backend": "pytorch", - "attn_backend": "TRTLLM", - "disable_overlap_scheduler": True, - "dtype": torch_dtype, - "enable_chunked_prefill": False, - "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=use_kv_cache_manager_v2, - ), - "max_batch_size": batch_size, - "max_beam_width": num_beams, - "max_input_len": _MAX_SEQUENCE_LENGTH, - "max_num_tokens": max_num_tokens, - "max_seq_len": _MAX_SEQUENCE_LENGTH, - "model_kwargs": {"torch_dtype": torch_dtype}, - "scheduler_config": SchedulerConfig(use_python_scheduler=True), - } - - with LLM( - model_path, - cuda_graph_config=_cuda_graph_config( - True, - [encoder_num_tokens], - [encoder_seq_len], - [batch_size], - ), - **llm_kwargs, - ) as llm: - responses = llm.generate( - source_texts, - sampling_params=sampling_params, - use_tqdm=False, - ) - assert len(responses) == batch_size - - token_ids_by_request = [] - for request_idx, response in enumerate(responses): - token_ids = _assert_t5_response( - response, - num_return_sequences=num_return_sequences, - ) - token_ids_by_request.append(token_ids) - _print_generated_text( - tokenizer, - f"{case_id}, request={request_idx}", - "output", - token_ids, - ) - _assert_expected_generation( - tokenizer, - token_ids, - exact_match=True, - expected_token_ids_by_output=[expected_token_ids_by_request[request_idx]], - ) - - assert token_ids_by_request == [ - [expected_token_ids] for expected_token_ids in expected_token_ids_by_request - ] - _assert_encoder_decoder_cuda_graphs_captured(llm) - - @pytest.mark.parametrize( "model_name,expected_output_token_ids_by_request,torch_dtype,use_kv_cache_manager_v2," "num_beams,num_return_sequences,exact_match", @@ -1001,13 +814,6 @@ def test_t5_pytorch_generate_encoder_decoder_mixed_context_generation_batch( model_name = "t5-small" model_path = _get_t5_model_path(model_name) - tokenizer = AutoTokenizer.from_pretrained(model_path) - encoder_seq_lens_by_request = [ - len(tokenizer(source_text).input_ids) for source_text in _MIXED_ENCODER_SOURCE_TEXTS - ] - encoder_num_tokens = sum(encoder_seq_lens_by_request) - encoder_num_tokens_list = sorted({*encoder_seq_lens_by_request, encoder_num_tokens}) - encoder_seq_lens = sorted(set(encoder_seq_lens_by_request)) first_sampling_params = SamplingParams( max_tokens=_MIXED_CONTEXT_GENERATION_MAX_NEW_TOKENS, temperature=0.0, @@ -1023,12 +829,7 @@ def test_t5_pytorch_generate_encoder_decoder_mixed_context_generation_batch( model_path, backend="pytorch", attn_backend="TRTLLM", - cuda_graph_config=_cuda_graph_config( - True, - encoder_num_tokens_list, - encoder_seq_lens, - [2], - ), + cuda_graph_config=_decoder_cuda_graph_config([2]), disable_overlap_scheduler=True, dtype="bfloat16", enable_chunked_prefill=False, diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 9f5b60e783fb..115fdfa3d482 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -66,8 +66,6 @@ l0_b200: - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v2-cuda-graph-off-greedy-bart-large-cnn] - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch[bf16-kv-v1-decoder-cuda-graph-on-greedy-batch2-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_generate_encoder_decoder_fixed_encoder_lengths_batch[bf16-kv-v1-encoder-cuda-graph-fixed-shape-batch2-bart-large-cnn] - - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_fixed_encoder_lengths_batch[bf16-kv-v1-encoder-cuda-graph-fixed-shape-batch4-bart-large-cnn] - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-off-beam2-t5-small0] - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-off-beam2-flan-t5-small] - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-off-beam2-t5-base] @@ -93,8 +91,6 @@ l0_b200: - 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-v1-decoder-cuda-graph-on-greedy-batch2-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_generate_encoder_decoder_fixed_encoder_lengths_batch[bf16-kv-v1-encoder-cuda-graph-fixed-shape-batch2-t5-small] - - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_fixed_encoder_lengths_batch[bf16-kv-v1-encoder-cuda-graph-fixed-shape-batch4-t5-small] - test_e2e.py::test_ptp_quickstart_advanced[Llama3.1-8B-NVFP4-nvfp4-quantized/Meta-Llama-3.1-8B] - test_e2e.py::test_ptp_quickstart_advanced[Llama3.1-8B-FP8-llama-3.1-model/Llama-3.1-8B-Instruct-FP8] - test_e2e.py::test_ptp_quickstart_advanced_mtp[DeepSeek-V3-Lite-BF16-DeepSeek-V3-Lite/bf16] diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index b59fde9c51e5..2e68a17100f5 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -178,8 +178,6 @@ l0_h100: - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v2-cuda-graph-off-greedy-bart-large-cnn] - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch[bf16-kv-v1-decoder-cuda-graph-on-greedy-batch2-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_generate_encoder_decoder_fixed_encoder_lengths_batch[bf16-kv-v1-encoder-cuda-graph-fixed-shape-batch2-bart-large-cnn] - - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_fixed_encoder_lengths_batch[bf16-kv-v1-encoder-cuda-graph-fixed-shape-batch4-bart-large-cnn] - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-off-beam2-t5-small0] - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-off-beam2-flan-t5-small] - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-off-beam2-t5-base] @@ -205,8 +203,6 @@ 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-flan-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_encoder_lengths_batch[bf16-kv-v2-decoder-cuda-graph-on-greedy-batch2-t5-small] - - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_fixed_encoder_lengths_batch[bf16-kv-v1-encoder-cuda-graph-fixed-shape-batch2-t5-small] - - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_fixed_encoder_lengths_batch[bf16-kv-v1-encoder-cuda-graph-fixed-shape-batch4-t5-small] - test_e2e.py::test_trtllm_bench_iteration_log[PyTorch-streaming-meta-llama/Llama-3.1-8B-llama-3.1-model/Meta-Llama-3.1-8B] - test_e2e.py::test_trtllm_bench_iteration_log[PyTorch-non-streaming-meta-llama/Llama-3.1-8B-llama-3.1-model/Meta-Llama-3.1-8B] - test_e2e.py::test_trtllm_bench_request_rate_and_concurrency[enable_concurrency-] From d746e171e312f5adc64d3cc6a57edf088d781702 Mon Sep 17 00:00:00 2001 From: Guiju Zhang <7135567+cascade812@users.noreply.github.com> Date: Mon, 29 Jun 2026 15:09:18 -0700 Subject: [PATCH 6/9] clean up Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com> --- .../_torch/pyexecutor/model_engine.py | 64 +++++++------------ tensorrt_llm/llmapi/llm_args.py | 11 +--- 2 files changed, 25 insertions(+), 50 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index d5a7047d215d..36f6bd9bee31 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -4949,14 +4949,25 @@ def _run_cuda_graph_warmup_encoder(self) -> None: self._capture_encoder_cuda_graphs() - def _get_encoder_cuda_graph_warmup_configs( - self) -> List[Tuple[int, int, int]]: - """Return feasible (batch_size, num_tokens, max_seq_len) graph keys.""" + def _capture_encoder_cuda_graphs(self) -> None: + """Capture whole-model encoder CUDA graphs for all feasible keys. + + Feasibility filter (also used in source): + nt >= prev_sl + bs (enough tokens for this sl bucket) + prev_nt < bs * sl (not enough tokens for a smaller nt bucket) + nt <= bs * sl (num tokens should not exceed total possible in batch) + sl <= nt (seq len should not exceed num tokens) + """ + runner = self.encoder_cuda_graph_runner + if not runner.enabled: + return + batch_sizes = sorted(self._cuda_graph_batch_sizes, reverse=True) num_tokens_list = sorted(self._cuda_graph_num_tokens) seq_lens_list = sorted(self._cuda_graph_seq_lens) - warmup_configs: List[Tuple[int, int, int]] = [] + num_captured = 0 + logger.info("Capturing encoder CUDA graphs ...") for bs in batch_sizes: if bs > self.batch_size: continue @@ -4971,44 +4982,15 @@ def _get_encoder_cuda_graph_warmup_configs( if nt > bs * sl or sl > nt: continue - warmup_configs.append((bs, nt, sl)) - - if self._cuda_graph_padding_enabled: - for bs in batch_sizes: - if bs > self.batch_size: - continue - for sl in reversed(seq_lens_list): - nt = bs * sl - if nt <= self._max_cuda_graph_num_tokens: - warmup_configs.append((bs, nt, sl)) - - return list(dict.fromkeys(warmup_configs)) - - def _capture_encoder_cuda_graphs(self) -> None: - """Capture whole-model encoder CUDA graphs for all feasible keys. - - Feasibility filter (also used in source): - nt >= prev_sl + bs (enough tokens for this sl bucket) - prev_nt < bs * sl (not enough tokens for a smaller nt bucket) - nt <= bs * sl (num tokens should not exceed total possible in batch) - sl <= nt (seq len should not exceed num tokens) - """ - runner = self.encoder_cuda_graph_runner - if not runner.enabled: - return - - num_captured = 0 - logger.info("Capturing encoder CUDA graphs ...") - for bs, nt, sl in self._get_encoder_cuda_graph_warmup_configs(): - inputs = self._create_encoder_warmup_inputs(bs, nt, sl) - if inputs is None: - continue + inputs = self._create_encoder_warmup_inputs(bs, nt, sl) + if inputs is None: + continue - logger.info(f"Encoder CUDA graph capture: " - f"bs={bs}, nt={nt}, sl={sl}") - self.encoder_forward(inputs) - torch.cuda.synchronize() - num_captured += 1 + logger.info(f"Encoder CUDA graph capture: " + f"bs={bs}, nt={nt}, sl={sl}") + self.encoder_forward(inputs) + torch.cuda.synchronize() + num_captured += 1 logger.info(f"Captured {num_captured} encoder CUDA graph(s).") diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 6685b95993a6..4290b82bfb76 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -4587,20 +4587,13 @@ class TorchLlmArgs(BaseLlmArgs): @field_validator('cuda_graph_config', mode='before') @classmethod def infer_cuda_graph_config_mode(cls, v): - if isinstance(v, dict) and ("encoder" in v or "decoder" in v): - raise ValueError( - "Composite encoder-decoder CUDA graph configs have been " - "removed. Use CudaGraphConfig or DecodeCudaGraphConfig for " - "decoder CUDA graphs.") if isinstance(v, dict) and "mode" not in v: encoder_keys = { "num_tokens", "max_num_token", "seq_lens", "max_seq_len" } v = dict(v) - if any(k in v and v[k] not in (None, 0) for k in encoder_keys): - v["mode"] = "encode" - else: - v["mode"] = "decode" + v["mode"] = "encode" if any(k in v and v[k] not in (None, 0) + for k in encoder_keys) else "decode" return v multimodal_config: MultimodalConfig = Field( From ffcc48b24c4fa7a30db06cb65fa8796d7d4e1ab3 Mon Sep 17 00:00:00 2001 From: Guiju Zhang <7135567+cascade812@users.noreply.github.com> Date: Tue, 30 Jun 2026 18:06:44 -0700 Subject: [PATCH 7/9] address comments Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com> --- .../_torch/attention_backend/interface.py | 89 +++---- .../_torch/pyexecutor/cuda_graph_runner.py | 28 ++- .../_torch/pyexecutor/model_engine.py | 218 +++++++++--------- .../defs/llmapi/test_llm_api_pytorch_bart.py | 10 +- .../test_lists/test-db/l0_b200.yml | 1 + .../test_lists/test-db/l0_h100.yml | 1 + 6 files changed, 178 insertions(+), 169 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/interface.py b/tensorrt_llm/_torch/attention_backend/interface.py index 8d68cbb195f1..1e6bb6837e71 100644 --- a/tensorrt_llm/_torch/attention_backend/interface.py +++ b/tensorrt_llm/_torch/attention_backend/interface.py @@ -474,8 +474,6 @@ def create_cross_metadata( with ``seq_lens_kv`` set to ``encoder_seq_lens`` so that ``is_cross`` becomes ``True``. """ - encoder_seq_lens_tensor = torch.tensor(encoder_seq_lens, - dtype=torch.int) cross_md = copy.copy(self) cross_md._saved_tensors = {} if self.is_cuda_graph: @@ -484,41 +482,29 @@ def create_cross_metadata( # CUDA graph metadata buffers separate so preparing cross metadata # cannot overwrite self-attention sequence lengths. cross_md.cuda_graph_buffers = Buffers() - cross_md.kv_cache_manager = cross_kv_cache_manager cross_md._seq_lens_kv_cuda = None cross_md.cross = None - cross_md.seq_lens_kv = encoder_seq_lens_tensor - if encoder_num_cached_tokens_per_seq is not None: - base_params = self.kv_cache_params - cross_md.kv_cache_params = KVCacheParams( - use_cache=base_params.use_cache if base_params is not None else - (cross_kv_cache_manager is not None), - num_cached_tokens_per_seq=list( - encoder_num_cached_tokens_per_seq), - block_ids_per_seq=None, - host_max_attention_window_sizes=base_params. - host_max_attention_window_sizes - if base_params is not None else None, - host_sink_token_length=base_params.host_sink_token_length - if base_params is not None else None, - num_extra_kv_tokens=base_params.num_extra_kv_tokens - if base_params is not None else 0, - ) + self._update_cross_metadata( + cross_md, + encoder_seq_lens, + cross_kv_cache_manager, + encoder_num_cached_tokens_per_seq, + base_kv_cache_params=self.kv_cache_params, + block_ids_per_seq=None, + ) cross_md.__post_init__() return cross_md - def update_cross_metadata( + def _update_cross_metadata( self, + cross_md: "AttentionMetadata", encoder_seq_lens: List[int], cross_kv_cache_manager: Union[KVCacheManager, KVCacheManagerV2, None], - encoder_num_cached_tokens_per_seq: Optional[List[int]] = None, + encoder_num_cached_tokens_per_seq: Optional[List[int]], + *, + base_kv_cache_params: Optional[KVCacheParams], + block_ids_per_seq: Optional[List[list]], ) -> "AttentionMetadata": - """Refresh an existing CUDA graph cross-attention sub-metadata.""" - if not self.has_cross_sub_metadata: - raise RuntimeError( - "CUDA graph cross-attention metadata has not been initialized.") - - cross_md = self.cross encoder_seq_lens_tensor = torch.tensor(encoder_seq_lens, dtype=torch.int) cross_md.kv_cache_manager = cross_kv_cache_manager @@ -532,27 +518,52 @@ def update_cross_metadata( cross_md.prompt_lens = self.prompt_lens if encoder_num_cached_tokens_per_seq is not None: - base_params = cross_md.kv_cache_params cross_md.kv_cache_params = KVCacheParams( - use_cache=(base_params.use_cache if base_params is not None else - (cross_md.kv_cache_manager is not None)), + use_cache=(base_kv_cache_params.use_cache + if base_kv_cache_params is not None else + (cross_kv_cache_manager is not None)), num_cached_tokens_per_seq=list( encoder_num_cached_tokens_per_seq), - block_ids_per_seq=(base_params.block_ids_per_seq - if base_params is not None else None), + block_ids_per_seq=block_ids_per_seq, host_max_attention_window_sizes=( - base_params.host_max_attention_window_sizes - if base_params is not None else None), - host_sink_token_length=(base_params.host_sink_token_length - if base_params is not None else None), - num_extra_kv_tokens=(base_params.num_extra_kv_tokens - if base_params is not None else 0), + base_kv_cache_params.host_max_attention_window_sizes + if base_kv_cache_params is not None else None), + host_sink_token_length=( + base_kv_cache_params.host_sink_token_length + if base_kv_cache_params is not None else None), + num_extra_kv_tokens=(base_kv_cache_params.num_extra_kv_tokens if + base_kv_cache_params is not None else 0), ) cross_md.request_ids = self.request_ids cross_md.num_contexts = self.num_contexts return cross_md + def update_cross_metadata( + self, + encoder_seq_lens: List[int], + cross_kv_cache_manager: Union[KVCacheManager, KVCacheManagerV2, None], + encoder_num_cached_tokens_per_seq: Optional[List[int]] = None, + ) -> "AttentionMetadata": + """Refresh an existing CUDA graph cross-attention sub-metadata.""" + if not self.has_cross_sub_metadata: + raise RuntimeError( + "CUDA graph cross-attention metadata has not been initialized.") + + cross_md = self.cross + assert cross_md is not None + base_kv_cache_params = cross_md.kv_cache_params + block_ids_per_seq = (base_kv_cache_params.block_ids_per_seq + if base_kv_cache_params is not None else None) + return self._update_cross_metadata( + cross_md, + encoder_seq_lens, + cross_kv_cache_manager, + encoder_num_cached_tokens_per_seq, + base_kv_cache_params=base_kv_cache_params, + block_ids_per_seq=block_ids_per_seq, + ) + def update_for_spec_dec(self) -> None: """ Hook to be called during forward when using spec-dec one-model mode. diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index ebf405556b3d..0433ada60352 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -23,7 +23,7 @@ from ..speculative.spec_sampler_base import SampleStateTensorsSpec from ..speculative.utils import get_draft_kv_cache_manager from ..utils import make_weak_ref, piecewise_cuda_graph -from .llm_request import get_draft_token_length +from .llm_request import LlmRequest, get_draft_token_length from .resource_manager import (BaseResourceManager, ResourceManager, ResourceManagerType) from .sampler import SampleStateTensors @@ -115,7 +115,7 @@ def __init__(self, config: CUDAGraphRunnerConfig): Callable[[], Optional[torch.Tensor]]] = {} self.graph_metadata: Dict[KeyType, Dict[str, Any]] = {} self.memory_pool = config.cuda_graph_mem_pool - self.padding_dummy_requests: Dict[int, Any] = {} + self.padding_dummy_requests: Dict[int, LlmRequest] = {} self.dynamic_draft_len_mapping = config.dynamic_draft_len_mapping self.shared_static_tensors: Dict[str, torch.Tensor] = {} @@ -552,12 +552,9 @@ def _get_padded_batch(self, batch: ScheduledRequests, dummy_request = dummy_request[0] dummy_request.is_cuda_graph_dummy = True if self.is_encoder_decoder: - if not self._prepare_encoder_decoder_padding_dummy( + if not self._add_cross_dummy_request( dummy_request, resource_manager, - dummy_encoder_output_len): - kv_cache_manager.free_resources(dummy_request) - if draft_kv_cache_manager is not None: - draft_kv_cache_manager.free_resources(dummy_request) + dummy_encoder_output_len, draft_kv_cache_manager): return 0 spec_res_mgr = resource_manager.get_resource_manager( @@ -570,9 +567,10 @@ def _get_padded_batch(self, batch: ScheduledRequests, batch.generation_requests.extend([padding_dummy_request] * padding_size) return padding_size - def _prepare_encoder_decoder_padding_dummy( - self, dummy_request: Any, resource_manager: ResourceManager, - encoder_output_len: int) -> bool: + def _add_cross_dummy_request( + self, dummy_request: LlmRequest, resource_manager: ResourceManager, + encoder_output_len: int, + draft_kv_cache_manager: Optional[BaseResourceManager]) -> bool: cross_kv_cache_manager = resource_manager.get_resource_manager( ResourceManagerType.CROSS_KV_CACHE_MANAGER) if cross_kv_cache_manager is None: @@ -588,7 +586,15 @@ def _prepare_encoder_decoder_padding_dummy( is_gen=True, max_beam_width=self.config.max_beam_width, encoder_output_lens=encoder_output_lens) - return cross_dummy_requests is not None + if cross_dummy_requests is not None: + return True + + kv_cache_manager = resource_manager.get_resource_manager( + self.config.kv_cache_manager_key) + kv_cache_manager.free_resources(dummy_request) + if draft_kv_cache_manager is not None: + draft_kv_cache_manager.free_resources(dummy_request) + return False @staticmethod def _get_padding_dummy_encoder_output_len( diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 36f6bd9bee31..7dffd516b94b 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1008,12 +1008,7 @@ def warmup(self, resource_manager: ResourceManager) -> None: # Reset the global cuda graph dummy requests in warmup. self.cuda_graph_runner.padding_dummy_requests = {} - if self._is_encoder_decoder_model(): - AutoTuner.get() - with self.cuda_graph_runner.allow_capture(): - self._run_cuda_graph_warmup(resource_manager) - return - + is_enc_dec = self._is_encoder_decoder_model() if self.mapping.cp_size > 1: cp_type = self.mapping.cp_config.get("cp_type", None) if cp_type != CpType.HELIX: @@ -1028,11 +1023,12 @@ def warmup(self, resource_manager: ResourceManager) -> None: AutoTuner.get() can_run_general_warmup = ( - not self.is_draft_model and not self.mapping.has_cp_helix() - and self.guided_decoder is None + not is_enc_dec and not self.is_draft_model + and not self.mapping.has_cp_helix() and self.guided_decoder is None and not isinstance(kv_cache_manager, MambaHybridCacheManager)) - self._run_attention_warmup(resource_manager, can_run_general_warmup) + if not is_enc_dec: + self._run_attention_warmup(resource_manager, can_run_general_warmup) if can_run_general_warmup: # Specialize torch.compile graphs across the key input shapes before CUDA graph capture. @@ -1051,7 +1047,7 @@ def warmup(self, resource_manager: ResourceManager) -> None: torch.cuda.empty_cache() # Autotuner warmup uses context-only requests. Helix CP # is decode-only and runs into issues with autotuner warmup. - if not self.mapping.has_cp_helix(): + if not is_enc_dec and not self.mapping.has_cp_helix(): self._run_autotuner_warmup(resource_manager) # Release the autotuner's exploration-mode intermediates. The # exploration leftovers are pure waste that hide tens of GiB from @@ -1400,6 +1396,62 @@ def _capture_generation_cuda_graphs(self, else: max_seq_len_list = [effective_max_seq_len] + def prepare_cross_batch(batch: ScheduledRequests, + resource_manager: ResourceManager) -> None: + if not batch.generation_requests: + return + + max_encoder_output_len = self._get_max_encoder_output_len( + resource_manager) + hidden_size = self._get_enc_dec_hidden_size() + saved_request_state = [] + for request in batch.generation_requests: + saved_request_state.append( + (request, request.py_encoder_output, + request.py_skip_cross_kv_projection, request.state, + request.py_batch_idx, request._cached_tokens, + request._cached_tokens_set)) + request.py_encoder_output = torch.ones( + (max_encoder_output_len, hidden_size), + device="cuda", + dtype=self.dtype) + request.py_skip_cross_kv_projection = False + request.state = LlmRequestState.CONTEXT_INIT + request.context_current_position = 0 + request.context_chunk_size = 1 + + projection_batch = ScheduledRequests() + projection_batch.reset_context_requests(batch.generation_requests) + kv_cache_manager = resource_manager.get_resource_manager( + self.kv_cache_manager_key) + draft_kv_cache_manager = self._get_draft_kv_cache_manager( + resource_manager) + attn_metadata = self._set_up_attn_metadata(kv_cache_manager, + draft_kv_cache_manager) + with self.no_cuda_graph(): + projection_inputs, _ = self._prepare_inputs( + projection_batch, + kv_cache_manager, + attn_metadata, + spec_metadata=None, + new_tensors_device=None, + resource_manager=resource_manager, + maybe_graph=False) + self._project_enc_dec_warmup_cross_kv(projection_inputs) + torch.cuda.synchronize() + + for (request, encoder_output, skip_cross_kv_projection, state, + batch_idx, cached_tokens, + cached_tokens_set) in saved_request_state: + request.py_encoder_output = encoder_output + request.py_skip_cross_kv_projection = skip_cross_kv_projection + request.state = state + if state == LlmRequestState.GENERATION_IN_PROGRESS: + request.context_current_position = request.prompt_len + request.py_batch_idx = batch_idx + request._cached_tokens = cached_tokens + request._cached_tokens_set = cached_tokens_set + def _run_capture_pass(force_non_greedy: bool, label: str) -> None: spec_metadata = self.spec_metadata if force_non_greedy and spec_metadata is not None: @@ -1434,8 +1486,7 @@ def _run_capture_pass(force_non_greedy: bool, label: str) -> None: batch, draft_len > 0, resource_manager) self.runtime_draft_len = draft_len if self._is_encoder_decoder_model(): - self._prepare_decoder_cuda_graph_warmup_batch_for_encoder_decoder_model( - batch, resource_manager) + prepare_cross_batch(batch, resource_manager) self.forward(batch, new_tensors_device=None, resource_manager=resource_manager) @@ -1709,26 +1760,24 @@ def _create_cuda_graph_warmup_request( result = ScheduledRequests() num_extra_decoding_steps = self._get_num_extra_decoding_steps() - is_encoder_decoder = self._is_encoder_decoder_model() - dummy_encoder_output_len = ( - self._get_encoder_decoder_dummy_encoder_output_len(resource_manager) - if is_encoder_decoder else None) + is_enc_dec = self._is_encoder_decoder_model() + max_encoder_output_len = ( + self._get_max_encoder_output_len(resource_manager) + if is_enc_dec else None) # Add (batch_size - 1) dummy requests with seq_len=1. - short_token_nums = ([2] * - (batch_size - 1)) if is_encoder_decoder else None - short_encoder_output_lens = ( - [dummy_encoder_output_len] * - (batch_size - 1)) if is_encoder_decoder else None + token_nums = ([2] * (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=short_token_nums, + token_nums=token_nums, is_gen=True, max_num_draft_tokens=draft_len, kv_reserve_draft_tokens=self.max_draft_loop_tokens, use_mrope=self.use_mrope, max_beam_width=self.max_beam_width, - encoder_output_lens=short_encoder_output_lens, + encoder_output_lens=encoder_output_lens, num_extra_decoding_steps=num_extra_decoding_steps, draft_kv_cache_manager=draft_kv_cache_manager) @@ -1763,7 +1812,7 @@ def free_warmup_requests() -> None: available_tokens = min(available_tokens, draft_available_tokens) token_num = max( - 2 if is_encoder_decoder else 1, + 2 if is_enc_dec else 1, min( available_tokens, max_seq_len - 1 - get_num_extra_kv_tokens(self.spec_config) - _kv_draft)) @@ -1788,8 +1837,8 @@ def free_warmup_requests() -> None: kv_reserve_draft_tokens=self.max_draft_loop_tokens, use_mrope=self.use_mrope, max_beam_width=self.max_beam_width, - encoder_output_lens=[dummy_encoder_output_len] - if is_encoder_decoder else None, + encoder_output_lens=[max_encoder_output_len] + if is_enc_dec else None, num_extra_decoding_steps=num_extra_decoding_steps, draft_kv_cache_manager=draft_kv_cache_manager) @@ -1806,19 +1855,13 @@ def free_warmup_requests() -> None: spec_resource_manager.add_dummy_requests( request_ids=list(range(batch_size))) if self._is_encoder_decoder_model(): - if not self._allocate_encoder_decoder_dummy_cross_kv( - result.generation_requests, resource_manager): - for request in result.generation_requests: - kv_cache_manager.free_resources(request) - if draft_kv_cache_manager is not None: - draft_kv_cache_manager.free_resources(request) - if spec_resource_manager is not None: - spec_resource_manager.free_resources(request) + if not self._add_cross_dummy_requests(result.generation_requests, + resource_manager): return None return result - def _get_encoder_decoder_dummy_encoder_output_len( - self, resource_manager: ResourceManager) -> int: + def _get_max_encoder_output_len(self, + resource_manager: ResourceManager) -> int: cross_kv_cache_manager = resource_manager.get_resource_manager( ResourceManagerType.CROSS_KV_CACHE_MANAGER) max_encoder_output_len = int(self.max_seq_len) @@ -1830,9 +1873,8 @@ def _get_encoder_decoder_dummy_encoder_output_len( max_encoder_output_len))) return max(1, max_encoder_output_len) - def _allocate_encoder_decoder_dummy_cross_kv( - self, requests: List[LlmRequest], - resource_manager: ResourceManager) -> bool: + def _add_cross_dummy_requests(self, requests: List[LlmRequest], + resource_manager: ResourceManager) -> bool: if not requests: return True cross_kv_cache_manager = resource_manager.get_resource_manager( @@ -1841,80 +1883,37 @@ def _allocate_encoder_decoder_dummy_cross_kv( raise RuntimeError("Encoder-decoder CUDA graph warmup requires " "ResourceManagerType.CROSS_KV_CACHE_MANAGER.") - encoder_output_len = self._get_encoder_decoder_dummy_encoder_output_len( + max_encoder_output_len = self._get_max_encoder_output_len( resource_manager) for request in requests: request.py_encoder_output = None request.py_skip_cross_kv_projection = True - encoder_output_lens = [encoder_output_len] * len(requests) + encoder_output_lens = [max_encoder_output_len] * len(requests) cross_dummy_requests = cross_kv_cache_manager.add_dummy_requests( request_ids=[request.py_request_id for request in requests], token_nums=encoder_output_lens, is_gen=True, max_beam_width=1, encoder_output_lens=encoder_output_lens) - return cross_dummy_requests is not None - - def _prepare_decoder_cuda_graph_warmup_batch_for_encoder_decoder_model( - self, batch: ScheduledRequests, - resource_manager: ResourceManager) -> None: - if not batch.generation_requests: - return + if cross_dummy_requests is not None: + return True - encoder_output_len = self._get_encoder_decoder_dummy_encoder_output_len( - resource_manager) - hidden_size = self._get_encoder_decoder_hidden_size() - saved_request_state = [] - for request in batch.generation_requests: - saved_request_state.append( - (request, request.py_encoder_output, - request.py_skip_cross_kv_projection, request.state, - request.py_batch_idx, request._cached_tokens, - request._cached_tokens_set)) - request.py_encoder_output = torch.ones( - (encoder_output_len, hidden_size), - device="cuda", - dtype=self.dtype) - request.py_skip_cross_kv_projection = False - request.state = LlmRequestState.CONTEXT_INIT - request.context_current_position = 0 - request.context_chunk_size = 1 - - projection_batch = ScheduledRequests() - projection_batch.reset_context_requests(batch.generation_requests) kv_cache_manager = resource_manager.get_resource_manager( self.kv_cache_manager_key) draft_kv_cache_manager = self._get_draft_kv_cache_manager( resource_manager) - attn_metadata = self._set_up_attn_metadata(kv_cache_manager, - draft_kv_cache_manager) - with self.no_cuda_graph(): - projection_inputs, _ = self._prepare_inputs( - projection_batch, - kv_cache_manager, - attn_metadata, - spec_metadata=None, - new_tensors_device=None, - resource_manager=resource_manager, - maybe_graph=False) - self._project_encoder_decoder_warmup_cross_kv(projection_inputs) - torch.cuda.synchronize() - - for (request, encoder_output, skip_cross_kv_projection, state, - batch_idx, cached_tokens, - cached_tokens_set) in saved_request_state: - request.py_encoder_output = encoder_output - request.py_skip_cross_kv_projection = skip_cross_kv_projection - request.state = state - if state == LlmRequestState.GENERATION_IN_PROGRESS: - request.context_current_position = request.prompt_len - request.py_batch_idx = batch_idx - request._cached_tokens = cached_tokens - request._cached_tokens_set = cached_tokens_set - - def _project_encoder_decoder_warmup_cross_kv( - self, inputs: Dict[str, Any]) -> None: + spec_resource_manager = resource_manager.get_resource_manager( + ResourceManagerType.SPEC_RESOURCE_MANAGER) + for request in requests: + kv_cache_manager.free_resources(request) + if draft_kv_cache_manager is not None: + draft_kv_cache_manager.free_resources(request) + if spec_resource_manager is not None: + spec_resource_manager.free_resources(request) + return False + + def _project_enc_dec_warmup_cross_kv(self, inputs: Dict[str, Any]) -> None: encoder_hidden_states = inputs.get("encoder_hidden_states") cross_attn_metadata = inputs.get("cross_attn_metadata") if encoder_hidden_states is None or cross_attn_metadata is None: @@ -1928,7 +1927,7 @@ def _project_encoder_decoder_warmup_cross_kv( attn_metadata = inputs["attn_metadata"] hidden_states = torch.ones( - (attn_metadata.num_tokens, self._get_encoder_decoder_hidden_size()), + (attn_metadata.num_tokens, self._get_enc_dec_hidden_size()), device=encoder_hidden_states.device, dtype=encoder_hidden_states.dtype) for layer in layers: @@ -1943,7 +1942,7 @@ def _project_encoder_decoder_warmup_cross_kv( cross_attn_metadata=cross_attn_metadata, skip_cross_kv_projection=False) - def _get_encoder_decoder_hidden_size(self) -> int: + def _get_enc_dec_hidden_size(self) -> int: config = self.model.model_config.pretrained_config hidden_size = getattr(config, "hidden_size", None) if hidden_size is None: @@ -2496,7 +2495,7 @@ def _apply_position_id_offset(self, position_ids: List[int]) -> List[int]: return position_ids return [position_id + offset for position_id in position_ids] - def _prepare_encoder_decoder_cross_attention_inputs( + def _prepare_enc_dec_cross_attn_inputs( self, encoder_hidden_states: List[torch.Tensor], encoder_seq_lens: List[int], @@ -3106,7 +3105,7 @@ def _prepare_tp_inputs( # requests, whose outputs are discarded. mrope_dummy_seq_slot = self.max_num_tokens * self.mapping.pp_size num_accepted_draft_tokens = [] # per request - is_encoder_decoder = self._is_encoder_decoder_model() + is_enc_dec = self._is_encoder_decoder_model() cross_encoder_hidden_states: List[torch.Tensor] = [] cross_encoder_seq_lens: List[int] = [ ] # new encoder K/V tokens per decoder sequence @@ -3131,7 +3130,7 @@ def _prepare_tp_inputs( def append_cross_attention_state(request: LlmRequest, project_encoder_output: bool, repeat: int = 1) -> None: - if not is_encoder_decoder: + if not is_enc_dec: return encoder_output_len = int(request.encoder_output_len) @@ -3960,14 +3959,13 @@ def previous_seq_slots_device(): if hasattr(self.model.model_config.pretrained_config, 'chunk_size'): attn_metadata.mamba_chunk_size = self.model.model_config.pretrained_config.chunk_size attn_metadata.prepare() - cross_attention_inputs = ( - self._prepare_encoder_decoder_cross_attention_inputs( - cross_encoder_hidden_states, - cross_encoder_seq_lens, - cross_encoder_cached_tokens_per_seq, - attn_metadata, - resource_manager, - ) if is_encoder_decoder else {}) + cross_attention_inputs = (self._prepare_enc_dec_cross_attn_inputs( + cross_encoder_hidden_states, + cross_encoder_seq_lens, + cross_encoder_cached_tokens_per_seq, + attn_metadata, + resource_manager, + ) if is_enc_dec else {}) peft_cache_manager = resource_manager and resource_manager.get_resource_manager( ResourceManagerType.PEFT_CACHE_MANAGER) 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 126342be75d5..a3f36b6f7194 100644 --- a/tests/integration/defs/llmapi/test_llm_api_pytorch_bart.py +++ b/tests/integration/defs/llmapi/test_llm_api_pytorch_bart.py @@ -122,7 +122,7 @@ def _test_case( enable_cuda_graph=False, num_beams=1, num_return_sequences=1, - exact_match=False, + exact_match=True, feature_id="fp16-kv-v1-cuda-graph-off-greedy", ), _test_case( @@ -274,14 +274,6 @@ def _assert_decoder_cuda_graphs_captured(llm: LLM) -> None: assert model_engine.cuda_graph_runner.graphs -def _enable_trtllm_gen_attention(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("TRTLLM_ENABLE_TRTLLM_GEN_ATTENTION", "1") - - from tensorrt_llm._torch.attention_backend import trtllm - - monkeypatch.setattr(trtllm, "_TRTLLM_ENABLE_TRTLLM_GEN_ATTENTION", True) - - def _assert_bart_response( response: RequestOutput, num_return_sequences: int, diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 226cedfcc81b..176da6615bf9 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -90,6 +90,7 @@ l0_b200: - 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-v1-decoder-cuda-graph-on-greedy-batch2-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_generate_encoder_decoder_mixed_context_generation_batch - test_e2e.py::test_ptp_quickstart_advanced[Llama3.1-8B-NVFP4-nvfp4-quantized/Meta-Llama-3.1-8B] - test_e2e.py::test_ptp_quickstart_advanced[Llama3.1-8B-FP8-llama-3.1-model/Llama-3.1-8B-Instruct-FP8] - test_e2e.py::test_ptp_quickstart_advanced_mtp[DeepSeek-V3-Lite-BF16-DeepSeek-V3-Lite/bf16] diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index 2e68a17100f5..f6cc50ed9eb3 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -203,6 +203,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-flan-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_encoder_lengths_batch[bf16-kv-v2-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 - test_e2e.py::test_trtllm_bench_iteration_log[PyTorch-streaming-meta-llama/Llama-3.1-8B-llama-3.1-model/Meta-Llama-3.1-8B] - test_e2e.py::test_trtllm_bench_iteration_log[PyTorch-non-streaming-meta-llama/Llama-3.1-8B-llama-3.1-model/Meta-Llama-3.1-8B] - test_e2e.py::test_trtllm_bench_request_rate_and_concurrency[enable_concurrency-] From 34f955ae9484452e85c7ae071e7c7fe0faa3822a Mon Sep 17 00:00:00 2001 From: Guiju Zhang <7135567+cascade812@users.noreply.github.com> Date: Tue, 30 Jun 2026 18:19:56 -0700 Subject: [PATCH 8/9] minor update Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/model_engine.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 7dffd516b94b..92a5c9ab2890 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1437,7 +1437,7 @@ def prepare_cross_batch(batch: ScheduledRequests, new_tensors_device=None, resource_manager=resource_manager, maybe_graph=False) - self._project_enc_dec_warmup_cross_kv(projection_inputs) + self._populate_cross_kv_cache(projection_inputs) torch.cuda.synchronize() for (request, encoder_output, skip_cross_kv_projection, state, @@ -1913,7 +1913,7 @@ def _add_cross_dummy_requests(self, requests: List[LlmRequest], spec_resource_manager.free_resources(request) return False - def _project_enc_dec_warmup_cross_kv(self, inputs: Dict[str, Any]) -> None: + def _populate_cross_kv_cache(self, inputs: Dict[str, Any]) -> None: encoder_hidden_states = inputs.get("encoder_hidden_states") cross_attn_metadata = inputs.get("cross_attn_metadata") if encoder_hidden_states is None or cross_attn_metadata is None: From 4425e9a94e7b82b8db441218a83a188dddc25d45 Mon Sep 17 00:00:00 2001 From: Guiju Zhang <7135567+cascade812@users.noreply.github.com> Date: Wed, 1 Jul 2026 20:20:26 -0700 Subject: [PATCH 9/9] address comments Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com> --- .../_torch/pyexecutor/cuda_graph_runner.py | 19 ++-- .../_torch/pyexecutor/model_engine.py | 93 ++++++++----------- .../test_lists/test-db/l0_b200.yml | 1 + .../test_lists/test-db/l0_h100.yml | 1 + tests/integration/test_lists/waives.txt | 2 + 5 files changed, 53 insertions(+), 63 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index 0433ada60352..a316b28bac61 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -31,6 +31,11 @@ # A large prime number used for dummy request IDs to avoid collisions CUDA_GRAPH_DUMMY_REQUEST_ID = (1 << 64) - 1 +# Gen dummies get prompt_len = token_num - 1. Before capturing enc-dec decode +# graphs, prepare_cross_batch temporarily runs each dummy generation request +# 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] @@ -525,7 +530,7 @@ def _get_padded_batch(self, batch: ScheduledRequests, if cross_kv_cache_manager is None: return 0 dummy_encoder_output_len = self._get_padding_dummy_encoder_output_len( - batch, cross_kv_cache_manager) + cross_kv_cache_manager) # Get draft KV cache manager only for one-model speculative decoding. # In two-model mode, each model has its own KV cache manager, so @@ -537,7 +542,8 @@ def _get_padded_batch(self, batch: ScheduledRequests, dummy_request_id = CUDA_GRAPH_DUMMY_REQUEST_ID - runtime_draft_len dummy_request = kv_cache_manager.add_dummy_requests( [dummy_request_id], - token_nums=[2] if self.is_encoder_decoder else None, + token_nums=[ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM] + if self.is_encoder_decoder else None, is_gen=True, max_num_draft_tokens=runtime_draft_len, use_mrope=self.config.use_mrope, @@ -598,15 +604,8 @@ def _add_cross_dummy_request( @staticmethod def _get_padding_dummy_encoder_output_len( - batch: ScheduledRequests, cross_kv_cache_manager: Any) -> int: + cross_kv_cache_manager: Any) -> int: encoder_output_len = 1 - for request in batch.generation_requests: - request_encoder_output_len = getattr(request, "encoder_output_len", - None) - if request_encoder_output_len is not None: - encoder_output_len = max(1, int(request_encoder_output_len)) - break - max_seq_len = getattr(cross_kv_cache_manager, "max_seq_len", None) if max_seq_len is not None: encoder_output_len = min(encoder_output_len, int(max_seq_len)) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 92a5c9ab2890..89755ad1ab75 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -66,7 +66,8 @@ set_per_request_piecewise_cuda_graph_flag, set_torch_compiling, with_model_extra_attrs) from .config_utils import is_mla -from .cuda_graph_runner import (CUDAGraphRunner, CUDAGraphRunnerConfig, +from .cuda_graph_runner import (ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM, + CUDAGraphRunner, CUDAGraphRunnerConfig, EncoderCUDAGraphRunner, EncoderCUDAGraphRunnerConfig) from .guided_decoder import CapturableGuidedDecoder @@ -1398,6 +1399,15 @@ def _capture_generation_cuda_graphs(self, def prepare_cross_batch(batch: ScheduledRequests, resource_manager: ResourceManager) -> None: + """Populate dummy gen requests' cross-KV cache before capture. + + Dummy generation requests used for graph capture never ran a + context step, so their cross-KV cache blocks are uninitialized + and captured kernels would read garbage. Temporarily switch each + request to a one-token context chunk with a fake encoder output + to run just the cross-KV projection (via _populate_cross_kv_cache), + then restore generation state for the actual capture. + """ if not batch.generation_requests: return @@ -1765,8 +1775,9 @@ def _create_cuda_graph_warmup_request( self._get_max_encoder_output_len(resource_manager) if is_enc_dec else None) - # Add (batch_size - 1) dummy requests with seq_len=1. - token_nums = ([2] * (batch_size - 1)) if is_enc_dec else 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( @@ -1812,7 +1823,7 @@ def free_warmup_requests() -> None: available_tokens = min(available_tokens, draft_available_tokens) token_num = max( - 2 if is_enc_dec else 1, + ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM if is_enc_dec else 1, min( available_tokens, max_seq_len - 1 - get_num_extra_kv_tokens(self.spec_config) - _kv_draft)) @@ -5651,25 +5662,6 @@ def load_weights_from_target_model(self, if callable(loader): loader(target_model) - @staticmethod - def _apply_logits_processors(request, logits_processors, logits_tensor, - beam_width, token_ids, logits_row_offset): - logits_rows = logits_tensor[logits_row_offset:logits_row_offset + - beam_width] - # Reshape to align w/ the shape used in the TRT backend, - # so the same logit processors can be used across both backends. - logits_rows = logits_rows.view(beam_width, 1, -1) - for lp in logits_processors: - lp_params = inspect.signature(lp).parameters - - assert 4 <= len(lp_params) <= 5, ( - "Logit post processor signature must match the `LogitsProcessor` interface " - "defined in `tensorrtllm.sampling_params`.") - lp(request.py_request_id, logits_rows, token_ids, None, None) - - logits_tensor[logits_row_offset:logits_row_offset + - beam_width] = logits_rows.view(beam_width, -1) - def _execute_logit_post_processors(self, scheduled_requests: ScheduledRequests, outputs: dict): @@ -5682,40 +5674,35 @@ def _execute_logit_post_processors(self, # TODO: support models that don't return outputs as dict return + num_ctx_req = scheduled_requests.num_context_requests logits_tensor = outputs["logits"] - logits_row_offset = 0 - request_groups = ( - (scheduled_requests.context_requests, True), - (scheduled_requests.generation_requests, False), - ) + for idx, request in enumerate(scheduled_requests.all_requests()): + logits_processors = getattr(request, "py_logits_post_processors", + None) + if not logits_processors: + continue - for requests, is_context_request in request_groups: - for request in requests: - if is_context_request: - beam_width = 1 - else: - beam_width = request.get_beam_width_by_iter( - for_next_iteration=False) - - logits_processors = getattr(request, - "py_logits_post_processors", None) - if logits_processors: - token_ids = ([request.get_tokens(0)] - if is_context_request else [ - request.get_tokens(beam_idx) - for beam_idx in range(beam_width) - ]) - if (is_context_request - and request.py_orig_prompt_len < len(token_ids[0])): - # Skip as we only need to apply logit processor on the last context request - logits_row_offset += beam_width - continue + token_ids = request.get_tokens(0) + if idx < num_ctx_req and request.py_orig_prompt_len < len( + token_ids): + # Skip as we only need to apply logit processor on the last context request + continue + + logits_row = logits_tensor[idx] + # Reshape to align w/ the shape used in the TRT backend, + # so the same logit processors can be used across both backends. + logits_row = logits_row.view(1, 1, -1) + token_ids = [token_ids] + for lp in logits_processors: + lp_params = inspect.signature(lp).parameters + + assert 4 <= len(lp_params) <= 5, ( + "Logit post processor signature must match the `LogitsProcessor` interface " + "defined in `tensorrtllm.sampling_params`.") + lp(request.py_request_id, logits_row, token_ids, None, None) - self._apply_logits_processors(request, logits_processors, - logits_tensor, beam_width, - token_ids, logits_row_offset) - logits_row_offset += beam_width + logits_tensor[idx] = logits_row.view(-1) def wait_for_input_copy(self): """ diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 176da6615bf9..2533a1630f58 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -62,6 +62,7 @@ l0_b200: - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-off-greedy-bart-large-cnn] - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_end_to_end[fp16-kv-v1-cuda-graph-off-greedy-bart-large-cnn] - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-off-beam2-bart-large-cnn] + - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-on-beam2-bart-large-cnn] - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v2-cuda-graph-off-greedy-bart-large-cnn] - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch[bf16-kv-v1-decoder-cuda-graph-on-greedy-batch2-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] diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index be4f9bad5ecf..922392eb1b40 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -175,6 +175,7 @@ l0_h100: - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-off-greedy-bart-large-cnn] - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_end_to_end[fp16-kv-v1-cuda-graph-off-greedy-bart-large-cnn] - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-off-beam2-bart-large-cnn] + - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-on-beam2-bart-large-cnn] - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v2-cuda-graph-off-greedy-bart-large-cnn] - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch[bf16-kv-v1-decoder-cuda-graph-on-greedy-batch2-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] diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 4d355d0ef136..5f3d30b2438f 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -313,6 +313,8 @@ full:sm100/unittest/trt/model/test_mamba.py SKIP (Disable for Blackwell) full:sm100/unittest/trt/quantization SKIP (Disable for Blackwell) full:sm100/unittest/trt/quantization/test_weight_only_groupwise_quant_matmul.py SKIP (Disable for Blackwell) full:sm100/unittest/trt/quantization/test_weight_only_quant_matmul.py SKIP (Disable for Blackwell) +llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-off-beam2-bart-large-cnn] SKIP (https://nvbugs/6340115) +llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-on-beam2-bart-large-cnn] SKIP (https://nvbugs/6340115) llmapi/test_llm_examples.py::test_llmapi_speculative_decoding_eagle3 SKIP (https://nvbugs/6075431) llmapi/test_llm_examples.py::test_llmapi_tensorrt_engine SKIP (https://nvbugs/5820553) perf/test_perf.py::test_perf[bart_large_cnn-plugin-float16-bs:8-input_output_len:60,20] SKIP # (https://nvidia.slack.com/archives/C059LSY62BT/p1704525727177449)