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

[cascade812]

Description

This PR enables CUDA graph execution for the decoder generation path of PyTorch encoder-decoder models, with initial coverage for T5 and BART.

The encoder prepass for encoder-decoder generation remains eager. Encoder CUDA graph capture/replay is preserved only for the encode-only / encoder-only path.

Key changes:

• Enables decoder CUDA graph warmup and replay for PyTorch encoder-decoder generation.
• Uses the existing CudaGraphConfig decoder-style configuration for encoder-decoder models.
• Adds encoder-decoder decoder graph warmup support, including dummy request setup, dummy encoder output state, and cross-KV cache projection so decoder capture has valid cross-attention state.
• Updates cross-attention metadata handling so decoder CUDA graph replay can refresh encoder sequence lengths, cached token counts, request IDs, and cross-KV state safely.
• Keeps existing encode-only encoder CUDA graph behavior intact.
• Updates BART and T5 integration coverage for greedy, beam search, KV cache manager v1/v2, mixed encoder lengths, and mixed context/generation scheduling.

Why decoder-only CUDA graphs for encoder-decoder models

We evaluated enabling CUDA graph capture for both the encoder forward pass and the decoder generation pass. The measured end-to-end latency gain from encoder+decoder CUDA graphs was almost the same as decoder-only CUDA graphs.

That result is expected for the target encoder-decoder serving path:

• The encoder prepass runs once per request, while the decoder runs repeatedly for generated tokens.
• The decoder path dominates CUDA launch overhead during generation, so decoder graph replay captures most of the practical gain.
• Encoder CUDA graph support for encoder-decoder generation requires extra shape buckets over (batch_size, encoder_num_tokens, encoder_seq_len), warmup inputs, padding behavior, and cross-KV/cache-state interactions.
• The incremental performance benefit from graphing the encoder forward pass is very limited relative to that added API and runtime complexity.

For that reason this PR intentionally does not enable CUDA graph support for the encoder forward pass of encoder-decoder models. Users should enable CUDA graphs for the decoder path only.

Customer Usage: Enabling Decoder CUDA Graph for Encoder-Decoder Models

For T5/BART-style PyTorch encoder-decoder models, use the existing CudaGraphConfig to configure decoder CUDA graph batch buckets.

from tensorrt_llm.llmapi import (
    LLM,
    CudaGraphConfig,
    KvCacheConfig,
    SchedulerConfig,
)

llm = LLM(
    model="path/to/t5-or-bart",
    backend="pytorch",
    attn_backend="TRTLLM",
    cuda_graph_config=CudaGraphConfig(
        batch_sizes=[1, 2, 4],
        enable_padding=True,
    ),
    disable_overlap_scheduler=True,
    enable_chunked_prefill=False,
    max_batch_size=4,
    max_seq_len=128,
    kv_cache_config=KvCacheConfig(
        cross_kv_cache_fraction=0.5,
    ),
    scheduler_config=SchedulerConfig(use_python_scheduler=True),
)

EncodeCudaGraphConfig remains for encode-only models. It is not supported for encoder-decoder generation.

Test Coverage

• python3 -m py_compile tensorrt_llm/_torch/pyexecutor/model_engine.py tensorrt_llm/llmapi/llm_args.py tensorrt_llm/llmapi/__init__.py tensorrt_llm/llmapi/llm_utils.py tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py tests/integration/defs/llmapi/test_llm_api_pytorch_bart.py tests/unittest/llmapi/test_llm_args.py
• git diff --check
• PYTHONPATH=/home/scratch.guijuz_coreai/TensorRT-LLM pytest tests/unittest/llmapi/test_llm_args.py::TestTorchLlmArgsCudaGraphSettings -q
• PYTHONPATH=/home/scratch.guijuz_coreai/TensorRT-LLM pytest tests/unittest/llmapi/test_llm_encode.py -q
• PYTHONPATH=/home/scratch.guijuz_coreai/TensorRT-LLM pytest --collect-only tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py tests/integration/defs/llmapi/test_llm_api_pytorch_bart.py -q

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

[coderabbitai]

📝 Walkthrough

Walkthrough

Encoder-decoder CUDA graph handling now threads encoder lengths through metadata and dummy requests, adds in-place cross-metadata updates, extends warmup and replay paths for encoder-decoder models, and adds BART and T5 integration coverage for batch-size-aware and mixed-context execution.

Changes

Encoder-decoder CUDA graph support

Layer / File(s)	Summary
Cross-attention metadata and request shapes tensorrt_llm/_torch/attention_backend/interface.py, tensorrt_llm/_torch/pyexecutor/resource_manager.py, tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py	AttentionMetadata cross-metadata now accepts list lengths, refreshes cross sub-metadata in place, and dummy request builders in both KV cache managers carry encoder output lengths.
Encoder-decoder warmup padding tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py, tensorrt_llm/_torch/pyexecutor/model_engine.py	CUDA-graph padding and warmup paths compute dummy encoder output lengths, allocate cross-KV resources, thread encoder-output lengths through dummy requests, and free cross-KV resources on failure.
Encoder graph enablement and capture setup tensorrt_llm/_torch/pyexecutor/model_engine.py, tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py	Encoder CUDA-graph eligibility now distinguishes encoder-decoder mode, the runner config carries that mode, and encoder-specific warmup capture configs are computed from it.
Encoder input packing and replay tensorrt_llm/_torch/pyexecutor/model_engine.py	Cross-attention input prep switches between create and update metadata, packed encoder inputs are rebuilt for encoder graphs, and encoder forward routes through CUDA-graph replay helpers.
Executor guards tensorrt_llm/_torch/pyexecutor/py_executor.py	PyExecutor now rejects piecewise CUDA-graph compilation and speculative decoding with encoder-decoder CUDA graphs while removing the previous blanket CUDA-graph rejection.
BART encoder-decoder coverage tests/integration/defs/llmapi/test_llm_api_pytorch_bart.py, tests/unittest/_torch/helpers.py	BART integration coverage adds batch-size-aware CUDA graph configs, runner-state assertions, encoder-decoder end-to-end cases, and mock runner support.
T5 encoder-decoder coverage tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py	T5 integration coverage adds batch-size-aware encode graph configs, mixed-context async generation, and CUDA graph state assertions.

Estimated code review effort

🎯 5 (Critical) | ⏱️ ~90+ minutes

Suggested reviewers

• yuxianq
• danielafrimi
• jiaganc
• lfr-0531
• yizhang-nv

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 33.80% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is specific, concise, and accurately summarizes the main change.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description check	✅ Passed	The PR description follows the template well, explains the change, motivation, usage, and test coverage, and includes the checklist.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 2

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (2)

tests/integration/defs/llmapi/test_llm_api_pytorch_bart.py (2)

125-150: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win

Add batch-size-aware coverage for the other enabled BART CUDA graph paths.

Coverage is currently insufficient for beam-search and KVCacheManagerV2 batch-size-aware CUDA graph behavior: these enabled cases default to batch size [1], so only the v1 greedy case exercises the [2] padding/dummy-request path. Consider adding cuda_graph_batch_sizes=[2] to these enabled cases or adding separate list entries for beam2 and kv-v2.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/integration/defs/llmapi/test_llm_api_pytorch_bart.py` around lines 125
- 150, The BART CUDA graph test matrix in test_llm_api_pytorch_bart.py is
missing batch-size-aware coverage for the enabled beam-search and
KVCacheManagerV2 paths, so only the existing v1 greedy case exercises the batch
size [2] padding/dummy-request behavior. Update the relevant _test_case entries
to include cuda_graph_batch_sizes=[2] or add separate feature_id variants for
the beam2 and kv-v2 enabled CUDA graph cases, so the enabled paths are validated
with non-default batch sizing.

Source: Path instructions

---

327-356: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win

Wire kv_cache_dtype into KvCacheConfig or remove it from the parameter set.

kv_cache_dtype is accepted and forwarded into this helper, but it is not applied when constructing KvCacheConfig, so any non-default parametrized case would silently still run with "auto".

Proposed fix

         kv_cache_config=KvCacheConfig(
             enable_block_reuse=False,
             max_tokens=_MAX_KV_TOKENS,
             free_gpu_memory_fraction=_FREE_GPU_MEMORY_FRACTION,
             cross_kv_cache_fraction=_CROSS_KV_CACHE_FRACTION,
             use_kv_cache_manager_v2=use_kv_cache_manager_v2,
+            dtype=kv_cache_dtype,
         ),

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/integration/defs/llmapi/test_llm_api_pytorch_bart.py` around lines 327
- 356, The helper currently accepts kv_cache_dtype but never uses it when
building KvCacheConfig, so non-default test cases still run with the default
value. Update the LLM setup in test_llm_api_pytorch_bart.py so the
kv_cache_dtype argument is wired into KvCacheConfig, or remove kv_cache_dtype
from the helper signature and any callers if it is intentionally unused; use the
existing _sampling_params and KvCacheConfig construction in the test helper to
locate the change.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Around line 572-577: The encoder-decoder CUDA-graph guard in PyExecutor
currently only blocks speculative decoding when model_engine.spec_config is
present, but _prepare_and_schedule_batch() can also enable speculation via
self.drafter, so the unsupported path can still be reached. Update the check in
py_executor.py to reject CUDA graph whenever cuda_graph_config is set and either
spec_config or drafter-based speculation is enabled, using the existing
PyExecutor model_engine/self.drafter logic, and keep the NotImplementedError
path covering both speculative sources.

In `@tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py`:
- Around line 861-866: The mixed-context test only checks response shape, so it
can miss swapped or incorrect encoder context. Update the test around the second
request in this T5 mixed-context case to validate the actual decoded content,
using the existing _assert_t5_response helper and the
first_response/second_response values, by asserting prompt-specific expected
fragments or at minimum that the two generated outputs differ and align with
their intended source texts.

---

Outside diff comments:
In `@tests/integration/defs/llmapi/test_llm_api_pytorch_bart.py`:
- Around line 125-150: The BART CUDA graph test matrix in
test_llm_api_pytorch_bart.py is missing batch-size-aware coverage for the
enabled beam-search and KVCacheManagerV2 paths, so only the existing v1 greedy
case exercises the batch size [2] padding/dummy-request behavior. Update the
relevant _test_case entries to include cuda_graph_batch_sizes=[2] or add
separate feature_id variants for the beam2 and kv-v2 enabled CUDA graph cases,
so the enabled paths are validated with non-default batch sizing.
- Around line 327-356: The helper currently accepts kv_cache_dtype but never
uses it when building KvCacheConfig, so non-default test cases still run with
the default value. Update the LLM setup in test_llm_api_pytorch_bart.py so the
kv_cache_dtype argument is wired into KvCacheConfig, or remove kv_cache_dtype
from the helper signature and any callers if it is intentionally unused; use the
existing _sampling_params and KvCacheConfig construction in the test helper to
locate the change.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: e9c34301-d9a0-4d51-9cb9-a029750a6062

📥 Commits

Reviewing files that changed from the base of the PR and between 33bf59e and f228561.

📒 Files selected for processing (9)

• tensorrt_llm/_torch/attention_backend/interface.py
• tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
• tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
• tensorrt_llm/_torch/pyexecutor/model_engine.py
• tensorrt_llm/_torch/pyexecutor/py_executor.py
• tensorrt_llm/_torch/pyexecutor/resource_manager.py
• tests/integration/defs/llmapi/test_llm_api_pytorch_bart.py
• tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py
• tests/unittest/_torch/helpers.py

[nvpohanh]

@lowsfer @yizhang-nv could you review this? thanks!

[cascade812]

/bot run

[tensorrt-cicd]

PR_Github #56138 [ run ] triggered by Bot. Commit: a1f69df Link to invocation

[tensorrt-cicd]

PR_Github #56138 [ run ] completed with state FAILURE. Commit: a1f69df
/LLM/main/L0_MergeRequest_PR pipeline #45000 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[cascade812]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #56308 [ run ] triggered by Bot. Commit: da4dba2 Link to invocation

[tensorrt-cicd]

PR_Github #56308 [ run ] completed with state SUCCESS. Commit: da4dba2
/LLM/main/L0_MergeRequest_PR pipeline #45162 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[cascade812]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #57097 [ run ] triggered by Bot. Commit: 6f6b0bf Link to invocation

[cascade812]

/bot run

[tensorrt-cicd]

PR_Github #57301 [ run ] triggered by Bot. Commit: 6f6b0bf Link to invocation

[tensorrt-cicd]

PR_Github #57097 [ run ] completed with state ABORTED. Commit: 6f6b0bf

Link to invocation

[tensorrt-cicd]

PR_Github #57301 [ run ] completed with state FAILURE. Commit: 6f6b0bf
/LLM/main/L0_MergeRequest_PR pipeline #46065 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[cascade812]

/bot run

[tensorrt-cicd]

PR_Github #57361 [ run ] triggered by Bot. Commit: 6f6b0bf Link to invocation

[tensorrt-cicd]

PR_Github #57361 [ run ] completed with state FAILURE. Commit: 6f6b0bf
/LLM/main/L0_MergeRequest_PR pipeline #46113 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[cascade812]

/bot run

[cascade812]

/bot run

[tensorrt-cicd]

PR_Github #57505 [ run ] triggered by Bot. Commit: 21e0d6e Link to invocation

[tensorrt-cicd]

PR_Github #57505 [ run ] completed with state FAILURE. Commit: 21e0d6e
/LLM/main/L0_MergeRequest_PR pipeline #46237 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[cascade812]

/bot run

[tensorrt-cicd]

PR_Github #57511 [ run ] triggered by Bot. Commit: 21e0d6e Link to invocation

[tensorrt-cicd]

PR_Github #57511 [ run ] completed with state SUCCESS. Commit: 21e0d6e
/LLM/main/L0_MergeRequest_PR pipeline #46243 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[cascade812]

/bot run

[tensorrt-cicd]

PR_Github #57528 [ run ] triggered by Bot. Commit: 21e0d6e Link to invocation

[tensorrt-cicd]

PR_Github #57528 [ run ] completed with state SUCCESS. Commit: 21e0d6e
/LLM/main/L0_MergeRequest_PR pipeline #46259 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

CI Report

Link to invocation

[cascade812]

/bot skip --comment "pipeline 46259 passed"

[tensorrt-cicd]

PR_Github #57584 [ skip ] triggered by Bot. Commit: 66d89db Link to invocation

[tensorrt-cicd]

PR_Github #57584 [ skip ] completed with state SUCCESS. Commit: 66d89db
Skipping testing for commit 66d89db

Link to invocation