[TRTLLM-13176][feat] Enables CUDA graph execution for PyTorch encoder-decoder models

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.

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,17 +512,28 @@ 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:
                 return 0
             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:
@@ -601,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.
     """

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 = []