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/interface.py b/tensorrt_llm/_torch/attention_backend/interface.py index 04fc5a771924..4fe42484bc51 100644 --- a/tensorrt_llm/_torch/attention_backend/interface.py +++ b/tensorrt_llm/_torch/attention_backend/interface.py @@ -455,10 +455,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. @@ -474,11 +473,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). @@ -499,32 +497,88 @@ 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 = None cross_md._seq_lens_kv_cuda = None cross_md.cross = None - cross_md.seq_lens_kv = encoder_seq_lens + 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( + 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]], + *, + base_kv_cache_params: Optional[KVCacheParams], + block_ids_per_seq: Optional[List[list]], + ) -> "AttentionMetadata": + 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: - 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), + 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, - 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, + block_ids_per_seq=block_ids_per_seq, + host_max_attention_window_sizes=( + 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.__post_init__() + + 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/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/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index cdbcd6ec8c29..476583cf6c9d 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] diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index b6969692db5f..e0bd91ac52a8 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 @@ -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] @@ -78,6 +83,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 +113,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, LlmRequest] = {} self.dynamic_draft_len_mapping = config.dynamic_draft_len_mapping self.shared_static_tensors: Dict[str, torch.Tensor] = {} @@ -523,6 +530,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( + 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 @@ -534,10 +549,14 @@ 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=[ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM] + if self.is_encoder_decoder else None, is_gen=True, max_num_draft_tokens=runtime_draft_token_buffer_width, 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: @@ -545,6 +564,11 @@ 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._add_cross_dummy_request( + dummy_request, resource_manager, + dummy_encoder_output_len, draft_kv_cache_manager): + return 0 spec_res_mgr = resource_manager.get_resource_manager( ResourceManagerType.SPEC_RESOURCE_MANAGER) @@ -556,6 +580,44 @@ def _get_padded_batch(self, batch: ScheduledRequests, batch.generation_requests.extend([padding_dummy_request] * padding_size) return padding_size + 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: + 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) + 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( + cross_kv_cache_manager: Any) -> int: + encoder_output_len = 1 + 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: @@ -633,11 +695,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/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 295372c4034a..d016ac88b32b 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -2602,6 +2602,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, ): @@ -2635,8 +2636,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( @@ -2646,6 +2649,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/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 591212348c6e..1dbe57d6f562 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -27,7 +27,7 @@ from tensorrt_llm.inputs.registry import (BaseMultimodalInputProcessor, 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) @@ -69,13 +69,14 @@ 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 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 @@ -249,7 +250,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, @@ -397,6 +398,17 @@ 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) + + 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 + 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 = self.cuda_graph_config.enable_padding if self.cuda_graph_config else CudaGraphConfig.model_fields[ @@ -410,9 +422,6 @@ def __init__( 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 and self.cuda_graph_config is not None and (not cuda_graph_num_tokens or not cuda_graph_seq_lens)): missing = [] @@ -682,6 +691,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, @@ -1019,12 +1029,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(): - logger.info( - "Skipping warmup for encoder-decoder models; warmup dummy " - "requests do not carry encoder output state.") - 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: @@ -1039,12 +1044,13 @@ 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)) log_mem_snapshot("warmup/before_warmup") - 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. @@ -1064,7 +1070,7 @@ def warmup(self, resource_manager: ResourceManager) -> None: # 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) log_mem_snapshot("warmup/after_autotuner") # Release the autotuner's exploration-mode intermediates. The @@ -1445,6 +1451,71 @@ 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: + """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 + + 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._populate_cross_kv_cache(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: @@ -1478,6 +1549,8 @@ 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(): + prepare_cross_batch(batch, resource_manager) self.forward(batch, new_tensors_device=None, resource_manager=resource_manager) @@ -1576,6 +1649,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: @@ -1586,6 +1661,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) @@ -1750,15 +1827,25 @@ def _create_cuda_graph_warmup_request( runtime_tokens_per_gen_step = self.get_runtime_tokens_per_gen_step( draft_len) runtime_draft_token_buffer_width = runtime_tokens_per_gen_step - 1 - - # Add (batch_size - 1) dummy requests with seq_len=1. + 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 the minimal seq_len. + token_nums = ([ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM] * + (batch_size - 1)) if is_enc_dec else None + encoder_output_lens = ([max_encoder_output_len] * + (batch_size - 1)) if is_enc_dec else None requests = kv_cache_manager.add_dummy_requests( list(range(batch_size - 1)), + token_nums=token_nums, is_gen=True, max_num_draft_tokens=runtime_draft_token_buffer_width, kv_reserve_draft_tokens=self.max_draft_loop_tokens, use_mrope=self.use_mrope, max_beam_width=self.max_beam_width, + encoder_output_lens=encoder_output_lens, num_extra_decoding_steps=num_extra_decoding_steps, draft_kv_cache_manager=draft_kv_cache_manager) @@ -1793,7 +1880,7 @@ def free_warmup_requests() -> None: available_tokens = min(available_tokens, draft_available_tokens) token_num = max( - 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)) @@ -1818,6 +1905,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=[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) @@ -1833,8 +1922,105 @@ def free_warmup_requests() -> None: if spec_resource_manager is not None: spec_resource_manager.add_dummy_requests( request_ids=list(range(batch_size))) + if self._is_encoder_decoder_model(): + if not self._add_cross_dummy_requests(result.generation_requests, + resource_manager): + return None return result + 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) + 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 _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( + 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.") + + 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 = [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) + if cross_dummy_requests is not None: + return True + + 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) + 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 _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: + 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_enc_dec_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_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: + 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): @@ -2413,7 +2599,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], @@ -2460,51 +2646,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, ) @@ -3099,7 +3251,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 @@ -3124,7 +3276,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) @@ -3192,7 +3344,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. @@ -3975,14 +4129,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/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 1329a26bdbae..b0d4376c5348 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -661,6 +661,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 @@ -810,31 +842,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 cb048c2d9aaf..f30967275960 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -860,6 +860,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 @@ -893,9 +894,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, @@ -903,7 +905,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 e35486872abc..a3f36b6f7194 100644 --- a/tests/integration/defs/llmapi/test_llm_api_pytorch_bart.py +++ b/tests/integration/defs/llmapi/test_llm_api_pytorch_bart.py @@ -30,28 +30,38 @@ 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." ), ] _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] +_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]], ] @@ -63,10 +73,14 @@ 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 + 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, @@ -76,6 +90,7 @@ def _test_case( num_beams, num_return_sequences, exact_match, + cuda_graph_batch_sizes, kv_cache_dtype, id=f"{feature_id}-{_MODEL_NAME}", ) @@ -91,13 +106,23 @@ 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, 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( @@ -106,9 +131,18 @@ 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( + torch_dtype="bfloat16", + use_kv_cache_manager_v2=False, + enable_cuda_graph=True, + num_beams=2, + num_return_sequences=2, + exact_match=True, + feature_id="bf16-kv-v1-cuda-graph-on-beam2", + ), _test_case( torch_dtype="bfloat16", use_kv_cache_manager_v2=True, @@ -118,6 +152,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", + ), ] @@ -143,17 +186,18 @@ def _mixed_batch_test_case( 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", ), ] + pytestmark = [ pytest.mark.skip_less_device(1), pytest.mark.skip_less_device_memory(_MIN_GPU_MEMORY_MB), @@ -190,16 +234,50 @@ def _sampling_params(num_beams: int, num_return_sequences: int) -> SamplingParam ) -def _cuda_graph_config( - enabled: bool, +def _decoder_cuda_graph_config( batch_sizes: list[int] | None = None, -) -> CudaGraphConfig | None: - return CudaGraphConfig(batch_sizes=batch_sizes or [1]) if enabled else 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], + enable_padding=True, + ) + + +def _assert_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_decoder_cuda_graphs_captured(llm) + 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 _assert_bart_response( response: RequestOutput, num_return_sequences: int, + max_tokens: int = _MAX_NEW_TOKENS, ) -> list[list[int]]: assert response.finished @@ -207,7 +285,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 @@ -224,18 +302,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 @@ -243,19 +316,19 @@ 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, 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") @@ -274,7 +347,9 @@ 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_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, @@ -286,7 +361,7 @@ def _run_bart_pytorch_generate_encoder_decoder( use_kv_cache_manager_v2=use_kv_cache_manager_v2, dtype=kv_cache_dtype, ), - 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,22 +385,29 @@ def _run_bart_pytorch_generate_encoder_decoder( exact_match, expected_output_token_ids_by_output, ) + _assert_decoder_cuda_graph_state( + llm, + enable_cuda_graph, + 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,kv_cache_dtype", + "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, + expected_output_token_ids_by_output: list[list[int]], 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( @@ -337,6 +419,7 @@ def test_bart_pytorch_generate_encoder_decoder_end_to_end( num_beams, num_return_sequences, exact_match, + cuda_graph_batch_sizes, kv_cache_dtype, ) @@ -360,14 +443,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, @@ -408,7 +491,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 b9a68bcac8be..5ea4cffdcf25 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 @@ -95,6 +96,7 @@ "t5-small": [_EXPECTED_TRANSLATION_FRAGMENT, "Buch"], "flan-t5-small": [_EXPECTED_TRANSLATION_FRAGMENT, "Buch"], } +_MIXED_CONTEXT_GENERATION_MAX_NEW_TOKENS = 8 def _test_case( @@ -106,6 +108,7 @@ def _test_case( num_return_sequences: int, exact_match: bool, feature_id: str, + cuda_graph_batch_sizes: list[int] | None = None, marks=None, ): if num_beams == 1: @@ -140,6 +143,7 @@ def _test_case( num_beams, num_return_sequences, exact_match, + cuda_graph_batch_sizes, **param_kwargs, ) @@ -249,6 +253,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( @@ -302,6 +328,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", @@ -402,7 +439,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", @@ -411,7 +448,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", @@ -420,7 +457,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", @@ -429,10 +466,11 @@ 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", ), ] + pytestmark = [ pytest.mark.skip_less_device(1), pytest.mark.skip_less_device_memory(_MIN_GPU_MEMORY_MB), @@ -469,16 +507,63 @@ def _sampling_params(num_beams: int, num_return_sequences: int) -> SamplingParam ) -def _cuda_graph_config( - enabled: bool, +def _decoder_cuda_graph_config( batch_sizes: list[int] | None = None, -) -> CudaGraphConfig | None: - return CudaGraphConfig(batch_sizes=batch_sizes or [1]) if enabled else 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], + enable_padding=True, + ) + + +def _assert_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_decoder_cuda_graphs_captured(llm) + 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 _assert_mixed_context_generation_cuda_graph_state(llm: LLM) -> None: + _assert_decoder_cuda_graphs_captured(llm) + assert llm._executor.engine.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 @@ -486,7 +571,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 @@ -524,6 +609,11 @@ def _assert_expected_generation( assert len({tuple(token_ids) for token_ids in token_ids_by_output}) == len( token_ids_by_output ) + 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_t5_pytorch_generate_encoder_decoder( @@ -536,6 +626,7 @@ def _run_t5_pytorch_generate_encoder_decoder( 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") @@ -552,7 +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), + 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, @@ -563,7 +656,7 @@ def _run_t5_pytorch_generate_encoder_decoder( 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, @@ -587,11 +680,16 @@ def _run_t5_pytorch_generate_encoder_decoder( exact_match, expected_output_token_ids_by_output, ) + _assert_decoder_cuda_graph_state( + llm, + enable_cuda_graph, + cuda_graph_batch_sizes, + ) @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", + "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( @@ -604,6 +702,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: _run_t5_pytorch_generate_encoder_decoder( monkeypatch, @@ -615,6 +714,7 @@ def test_t5_pytorch_generate_encoder_decoder_end_to_end( num_beams, num_return_sequences, exact_match, + cuda_graph_batch_sizes, ) @@ -641,14 +741,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, @@ -702,3 +802,73 @@ 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, ) + + _assert_decoder_cuda_graphs_captured(llm) + + +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=_decoder_cuda_graph_config([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) diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 138d4f2f40b4..16768839b3ac 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -62,9 +62,10 @@ 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-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_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] @@ -86,10 +87,11 @@ 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_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 0b4d0d4ac687..cc9f83b31d43 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -189,9 +189,10 @@ 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-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_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] @@ -213,10 +214,11 @@ 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_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-enable_request_rate] # negative test diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 70b77ea9df90..2a387b76e6d6 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -485,23 +485,7 @@ full:sm100/unittest/trt/quantization/test_weight_only_groupwise_quant_matmul.py full:sm100/unittest/trt/quantization/test_weight_only_quant_matmul.py SKIP (Disable for Blackwell) kv_cache/test_kv_cache_v2_scheduler.py::TestKVCacheV2Llama::test_chunked_prefill_multi_request SKIP (https://nvbugs/6405665) 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_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) 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)