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

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,

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.

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,
         )

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]

tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py

@@ -23,14 +23,19 @@
 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
 from .scheduler import ScheduledRequests
 
 # 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,17 +549,26 @@ 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:
                 return 0
             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.
     """

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

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,
         )