[None][perf] Improve PyTorch encoder-decoder support and performance - #16158
Conversation
Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com>
|
/bot run --disable-fast-fail |
📝 WalkthroughWalkthroughThis PR enables the overlap scheduler for encoder-decoder models. It adds a ChangesOverlap scheduler support for encoder-decoder
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant ExecutorLoopOverlap
participant ModelEngine
participant LogitPostProcessors
participant GuidedDecoder
ExecutorLoopOverlap->>ModelEngine: run encoder step (encoder_requests)
ExecutorLoopOverlap->>ModelEngine: forward(scheduled_batch)
ModelEngine-->>ExecutorLoopOverlap: batch_outputs
alt defer_logits_post_processors enabled
ExecutorLoopOverlap->>LogitPostProcessors: _execute_logit_post_processors(scheduled_batch, batch_outputs)
LogitPostProcessors-->>ExecutorLoopOverlap: processed outputs
end
ExecutorLoopOverlap->>GuidedDecoder: proceed with sampling
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py (1)
407-455: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winSame overlap coverage gaps as the BART file.
New overlap cases are limited to
t5-small,use_kv_cache_manager_v2=False, andtensor_parallel_size=1. Consider adding an overlap case withuse_kv_cache_manager_v2=True(greedy path) and one withtensor_parallel_size=2, since overlap-scheduler gating depends onpp_size, not TP size.As with the BART file, none of these tests assert that the overlap scheduler was actually engaged — they only compare outputs to non-overlap baselines. Consider adding an executor-state assertion (paralleling
_assert_decoder_cuda_graph_state) for TP=1 runs to guard against a silent fallback regression.Coverage verdict: adequate for basic output-parity regression detection, but insufficient to catch a silent overlap-scheduler wiring regression or KV-cache-v2/TP interactions; recommend as follow-up if not addressed here.
🤖 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_t5.py` around lines 407 - 455, The overlap-scheduler coverage in the T5 integration tests is incomplete because the new cases only exercise `t5-small` with `use_kv_cache_manager_v2=False` and `tensor_parallel_size=1`. Add at least one overlap case through the existing `_test_case` helper for the `use_kv_cache_manager_v2=True` greedy path and another with `tensor_parallel_size=2`, since gating depends on `pp_size` rather than TP size. Also add an executor-state assertion for TP=1 runs, similar to `_assert_decoder_cuda_graph_state`, so the tests verify that the overlap scheduler actually engaged instead of only comparing outputs.Source: Path instructions
tests/integration/defs/llmapi/test_llm_api_pytorch_bart.py (1)
173-216: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winOverlap test coverage is a good start but has two gaps worth a follow-up.
The new overlap cases only cover
use_kv_cache_manager_v2=Falseandtensor_parallel_size=1. Given the PR also touches the KV-cache-manager-v2 path indirectly (cross-attention fast path caching) and TP behavior is independent of the overlap-scheduler gating (onlypp_sizedisables overlap), consider adding:
- one overlap case with
use_kv_cache_manager_v2=True(greedy only, matching existing v2 constraints), and- one overlap case with
tensor_parallel_size=2(mirroring the existing TP2 cases at lines 218-239).Also, none of these tests directly assert the overlap scheduler was actually engaged (e.g. checking
llm._executor.event_loopor an equivalent state), they only compare final outputs to the non-overlap baseline. A wiring regression that silently falls back to the non-overlap loop would still pass. Consider adding an assertion similar to_assert_decoder_cuda_graph_state(lines 325-341) that inspects executor state for TP=1 runs.Coverage verdict: current overlap coverage is adequate for basic regression detection but insufficient to catch a silent fallback to the non-overlap path or KV-cache-v2/TP interactions; recommend tracking the above as follow-up if not addressed here.
🤖 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 173 - 216, Add follow-up overlap-scheduler coverage in the BART integration test cases: extend the existing `_test_case` overlap block to include one `use_kv_cache_manager_v2=True` greedy case and one `tensor_parallel_size=2` case, matching the constraints already used elsewhere in `test_llm_api_pytorch_bart.py`. Also add an assertion in the TP=1 overlap tests that explicitly verifies the overlap path was engaged by inspecting the same executor state used by `_assert_decoder_cuda_graph_state` or an equivalent `llm._executor` event-loop/scheduler indicator, so a silent fallback to the non-overlap path cannot pass.Source: Path instructions
🤖 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/model_engine.py`:
- Around line 2663-2697: Gate the CUDA-graph cross-attn fast path on request
identity as well as cached-token counts. In the `is_stable_gen_step` branch
inside `model_engine.py`, the reused `attn_metadata.cross` must not skip
`prepare()` when `request_ids` have changed or been reordered, because
`VanillaAttentionMetadata.prepare()` rebuilds `block_ids_per_seq` from them.
Update the stability check around `attn_metadata.is_cuda_graph`,
`has_cross_sub_metadata`, and `_cross_attn_stable_cached_tokens` so that any
request swap or order change falls back to `update_cross_metadata()` and
`prepare()` instead of reusing stale cross metadata.
---
Nitpick comments:
In `@tests/integration/defs/llmapi/test_llm_api_pytorch_bart.py`:
- Around line 173-216: Add follow-up overlap-scheduler coverage in the BART
integration test cases: extend the existing `_test_case` overlap block to
include one `use_kv_cache_manager_v2=True` greedy case and one
`tensor_parallel_size=2` case, matching the constraints already used elsewhere
in `test_llm_api_pytorch_bart.py`. Also add an assertion in the TP=1 overlap
tests that explicitly verifies the overlap path was engaged by inspecting the
same executor state used by `_assert_decoder_cuda_graph_state` or an equivalent
`llm._executor` event-loop/scheduler indicator, so a silent fallback to the
non-overlap path cannot pass.
In `@tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py`:
- Around line 407-455: The overlap-scheduler coverage in the T5 integration
tests is incomplete because the new cases only exercise `t5-small` with
`use_kv_cache_manager_v2=False` and `tensor_parallel_size=1`. Add at least one
overlap case through the existing `_test_case` helper for the
`use_kv_cache_manager_v2=True` greedy path and another with
`tensor_parallel_size=2`, since gating depends on `pp_size` rather than TP size.
Also add an executor-state assertion for TP=1 runs, similar to
`_assert_decoder_cuda_graph_state`, so the tests verify that the overlap
scheduler actually engaged instead of only comparing outputs.
🪄 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: d3fd0068-131f-437c-a1e7-36e5866bc1bd
📒 Files selected for processing (4)
tensorrt_llm/_torch/pyexecutor/model_engine.pytensorrt_llm/_torch/pyexecutor/py_executor.pytests/integration/defs/llmapi/test_llm_api_pytorch_bart.pytests/integration/defs/llmapi/test_llm_api_pytorch_t5.py
|
PR_Github #58330 Bot args parsing error: usage: /bot [-h] |
|
/bot run --disable-fail-fast |
|
PR_Github #58371 [ run ] triggered by Bot. Commit: |
|
/bot kill |
|
PR_Github #58376 [ kill ] triggered by Bot. Commit: |
|
PR_Github #58371 [ run ] completed with state |
|
PR_Github #58376 [ kill ] completed with state |
Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com>
|
/bot run |
|
PR_Github #59013 [ run ] triggered by Bot. Commit: |
|
PR_Github #59013 [ run ] completed with state
|
cascade812
changed the title
cascade812
changed the title
|
/bot run --disable-fail-fast |
|
PR_Github #59040 [ run ] triggered by Bot. Commit: |
Wanli-Jiang
left a comment
Wanli-Jiang
left a comment
There was a problem hiding this comment.
LGTM for model changes. Better to ask review from request input/output response changes, it is a bit complicated for me.
Tabrizian
left a comment
Tabrizian
left a comment
There was a problem hiding this comment.
Minor comments, otherwise LGTM.
Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #60070 [ run ] triggered by Bot. Commit: |
|
PR_Github #60070 [ run ] completed with state
|
|
/bot run |
|
PR_Github #60090 [ run ] triggered by Bot. Commit: |
|
PR_Github #60090 [ run ] completed with state
|
|
/bot run |
|
PR_Github #60122 [ run ] triggered by Bot. Commit: |
|
PR_Github #60122 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #60130 [ run ] triggered by Bot. Commit: |
|
PR_Github #60130 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #60186 [ run ] triggered by Bot. Commit: |
|
PR_Github #60186 [ run ] completed with state
|
|
/bot run |
|
PR_Github #60195 [ run ] triggered by Bot. Commit: |
|
PR_Github #60195 [ run ] completed with state
|
|
/bot run |
|
PR_Github #60267 [ run ] triggered by Bot. Commit: |
|
PR_Github #60267 [ run ] completed with state |
9 participants
Description
This PR broadens and optimizes encoder-decoder support in the PyTorch backend for T5, BART, and mBART. It enables overlap-scheduler execution, removes recurring decode overhead, improves T5 kernels, aligns BART and mBART generation behavior with the legacy TensorRT encoder-decoder path (which differs slightly from Hugging Face
generate()at the output-token limit; see item 4), and adds customer-facing documentation and test coverage.Runtime changes
Enable the overlap scheduler for encoder-decoder models (
py_executor.py)_executor_loop_overlapbefore decoder forward execution, so encoder-decoder models now support the PyTorch backend's default overlap scheduling.Reduce recurring decode overhead (
model_engine.py)update_cross_metadata()+prepare()calls (repeated tensor allocation and host-to-device copies) and refresh only the decoder-side references. The cached state is invalidated whenever the batch composition changes or execution leaves the CUDA-graph path.prepare_encoder_only()instead of the fullprepare().Improve T5 kernels (
modeling_t5.py)gelu_newMLPs in FP16/BF16.Model and generation behavior changes
Seed BART/mBART forced BOS in the decoder prompt instead of a logits processor (
llm.py,sampling_params.py,base_worker.py,result.py)forced_bos_token_id, the token is now placed directly in the internal decoder prompt ([decoder_start_token_id, forced_bos_token_id]) rather than forced by the per-step_BartForcedTokensLogitsProcessor, which is removed. This eliminates per-step host work and keeps the decode path CUDA-graph and overlap friendly.SamplingParams._decoder_output_token_prefixrecords the moved token; the result layer restores it to the user-visible token IDs and logprobs, andmax_tokensdeduction subtracts the prefix. The forced BOS still counts againstmax_tokens, so BART/mBART requests with forced BOS requiremax_tokens >= 2(validated with a clear error).forced_eos_token_idis no longer injected when a sequence reachesmax_tokens; the model-selected final token is preserved andfinish_reason="length"is reported. EOS remains a natural stopping token viaend_id. This matches the legacy TensorRT encoder-decoder path but differs from Hugging Facegenerate(), which forces EOS as the final token when a sequence reaches the output limit.Complete mBART model support (
modeling_bart.py)ACT2FN, while preserving BART post-norm behavior.Add an encoder-decoder user guide (
docs/source/models/encoder-decoder.md)Performance
The following benchmarks compare the PyTorch backend with the legacy TensorRT encoder-decoder path for large-batch inference. The measurements use BF16 on one H100 80 GB GPU with greedy decoding, natural EOS stopping, an output limit of 128 tokens, and mixed encoder input lengths from 260 to 440 tokens. The Flan-T5-XL results are the average of 10 timed runs after 3 warmup runs; the BART results are the average of 20 timed runs after 5 warmup runs. Executed-token throughput includes the terminal EOS token for both paths.
Flan-T5-XL
For Flan-T5-XL, the PyTorch backend performs on par with the legacy TensorRT path, with slightly lower latency and higher executed-token throughput across the tested batch sizes.
BART-large-CNN
For BART-large-CNN, the PyTorch backend has 21.9% to 36.0% higher latency than the legacy TensorRT path across the tested batch sizes (down from 67% to 155% before the overlap scheduler and cross-attention fast path in this PR). The remaining gap is dominated by host-side input preparation and sampling; work to close it is in progress and will come in a separate PR.
Test Coverage
l0_h100.yml; moved the single-GPU encoder-decoder suite froml0_b200.ymltol0_l40s.yml; removed the TP2 cases froml0_dgx_b200.ymlthat are already covered inl0_dgx_h100.yml.Summary by CodeRabbit
New Features
Bug Fixes
Tests