[None][perf] optimize encoder-decoder PyTorch performance - #16706
[None][perf] optimize encoder-decoder PyTorch performance#16706cascade812 wants to merge 13 commits into
Conversation
Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com>
Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com>
Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com>
Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com>
Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com>
Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com>
Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com>
Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com>
Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com>
cascade812
added
the
api-compatible
ChatGPT Codex Connector
WalkthroughChangesThe PR adds encoder and mixed encoder-decoder CUDA graph support for the PyTorch backend. It adds persistent input packing, cross-attention metadata preparation, graph capture and replay, asynchronous encoder scheduling, and single-step greedy sampling. It also adds configuration validation, documentation, and unit and integration coverage. Encoder-decoder graph execution
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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: 5
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tensorrt_llm/_torch/pyexecutor/model_engine.py (1)
6526-6535: 🚀 Performance & Scalability | 🟠 Major | ⚡ Quick winEncoder capture loop still gates on the decoder
batch_size.
self._encoder_cuda_graph_batch_sizesis already filtered againstself.encoder_batch_size(Lines 542-547), but Line 6534 skips any bucket larger thanself.batch_size. Whenencoder_max_batch_size > max_batch_size, the largest encoder graphs are silently never captured while the runtime admission/replay path still allows those encoder batch sizes, so those iterations fall back to eager.🐛 Proposed fix
for bs in batch_sizes: - if bs > self.batch_size: + if bs > self.encoder_batch_size: continue🤖 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 `@tensorrt_llm/_torch/pyexecutor/model_engine.py` around lines 6526 - 6535, Update the batch-size guard in the encoder CUDA graph loop within the warmup/capture flow to compare each bs against self.encoder_batch_size instead of the decoder self.batch_size. Preserve the existing behavior for buckets within the encoder limit so all runtime-admitted encoder batch sizes can be captured.
🧹 Nitpick comments (13)
tensorrt_llm/_torch/pyexecutor/sampler/sampler.py (3)
1463-1465: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueMinor: use
X | Noneinstead ofOptional[X]for new fields.The rest of this diff (e.g.,
d2t: torch.Tensor | Noneat line 2909) uses PEP 604 syntax; these three new fields useOptional[...]instead, which is inconsistent within the same change.As per coding guidelines,
**/*.pyshould "prefer built-in generic types and|."♻️ Proposed fix
- self._stable_greedy_request_ids: list[int] = [] - self._stable_greedy_seq_slots_host: Optional[torch.Tensor] = None - self._stable_greedy_seq_slots_cuda: Optional[torch.Tensor] = None + self._stable_greedy_request_ids: list[int] = [] + self._stable_greedy_seq_slots_host: torch.Tensor | None = None + self._stable_greedy_seq_slots_cuda: torch.Tensor | None = None🤖 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 `@tensorrt_llm/_torch/pyexecutor/sampler/sampler.py` around lines 1463 - 1465, Update the type annotations for _stable_greedy_seq_slots_host and _stable_greedy_seq_slots_cuda to use torch.Tensor | None instead of Optional[torch.Tensor], matching the PEP 604 style used elsewhere in the change; leave _stable_greedy_request_ids unchanged.
2908-2934: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocstring is stale after the return-value change.
_fast_greedy_sample_kernelnow returnsnext_tokens, but the docstring still describes it as purely in-place ("All operations are in-place").📝 Proposed doc tweak
"""Applies fast greedy sampling to the logits. Performs argmax, applies d2t translation if present, and scatters - tokens into the output buffer. All operations are in-place. + tokens into the output buffer (in-place) and returns the pre-scatter + per-request `next_tokens` tensor for callers that need it directly. """🤖 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 `@tensorrt_llm/_torch/pyexecutor/sampler/sampler.py` around lines 2908 - 2934, Update the docstring of _fast_greedy_sample_kernel to state that it computes and returns the translated greedy tokens while scattering them into the output buffer, removing the inaccurate claim that all operations are in-place.
3663-3706: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winCached stable-greedy path skips re-validating per-request invariants once
has_stable_request_idsis true.When the exact same
request_idslist recurs,can_use_stable_greedy_pathshort-circuits onhas_stable_request_idsand never re-checksget_draft_token_length(request) == 0, stop words, bad words,min_length,return_log_probs, embedding bias, orGREEDYstrategy for those requests — it only re-validates them the very first time a given request set enters this path. It also reusesself._stable_greedy_seq_slots_host/_cudawithout re-verifying they still match each request's currentpy_seq_slot.If any of these ever become mutable mid-life for a request that keeps the same
py_request_id(e.g., draft tokens re-appearing, or a seq-slot reassignment while the ID list happens to repeat),raw_logits_cuda[: len(generation_requests)]would silently misalign rows to requests, corrupting sampled tokens without any error. Based on theSpeculationGate/py_disable_speculative_decodingcontext inpy_executor.py, this transition looks one-directional in practice today, but nothing in this function enforces or documents that invariant.Consider adding a cheap re-check (e.g., asserting draft length stays 0, or comparing
py_seq_slotagainst the cached tensor) even on the cached branch, so a future change that violates the assumption fails loudly instead of silently corrupting output.#!/bin/bash # Description: Check whether draft-token/stop-word/min-length/log-prob state can change # mid-life for a request whose py_request_id is retained across steps. rg -nP -C4 'py_disable_speculative_decoding|SpeculationGate' tensorrt_llm/_torch/pyexecutor/py_executor.py tensorrt_llm/_torch/pyexecutor/llm_request.py🤖 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 `@tensorrt_llm/_torch/pyexecutor/sampler/sampler.py` around lines 3663 - 3706, Update the cached branch of can_use_stable_greedy_path so has_stable_request_ids does not bypass validation: re-check the same per-request greedy invariants and verify each current py_seq_slot matches the cached stable slot tensors before reusing them. If any invariant or slot mapping differs, reject the stable path and rebuild or use the normal path, preventing silent row-to-request misalignment.tests/unittest/_torch/sampler/test_torch_sampler.py (1)
666-743: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winTest coverage summary.
- Added:
test_single_step_greedy_checks_finish_reasons_on_host,test_single_step_greedy_filters_requests_completed_after_sampling— both targetTorchSampler.update_requests/_update_requests_single_beam_single_step, covering EOS-over-LENGTH precedence and correct skipping of already-GENERATION_COMPLETErequests.- Not covered by any test in this diff:
TorchSampler._process_requests's newcan_use_stable_greedy_path/has_stable_request_idscaching logic (sampler.py lines 3653-3726) — e.g., that the fast path is taken/declined correctly on first entry vs. cached entry, that it's invalidated when the request set changes, or that cachedseq_slotsstay consistent across iterations for the same request IDs.- Test list applicability: this is a
tests/unittest/file, not undertests/integration/test_lists/**, so CI/QA list membership doesn't apply here.- Verdict: needs follow-up — the
update_requestsconsumer side is well covered, but the higher-complexity caching/eligibility logic in_process_requeststhat producessingle_step_greedyhas no direct test in this diff.As per path instructions, "Always produce a test coverage summary, even if no issues are found" and "A coverage verdict: sufficient, insufficient, or needs follow-up."
🤖 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/unittest/_torch/sampler/test_torch_sampler.py` around lines 666 - 743, The new caching and eligibility logic in TorchSampler._process_requests, particularly can_use_stable_greedy_path and has_stable_request_ids, lacks direct tests. Add tests covering initial fast-path eligibility, reuse of cached eligibility, invalidation when request IDs change, and preservation of cached seq_slots for unchanged request IDs; retain the existing update_requests coverage.Source: Path instructions
tests/integration/defs/llmapi/test_llm_api_pytorch_bart.py (1)
309-316: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate
_SleepLogitsProcessoracross BART and T5 integration tests. Both files define the identical helper class for delaying generation during continuous-admission tests; the shared root cause is the lack of a common test-utility module for encoder-decoder CUDA-graph integration test helpers.
tests/integration/defs/llmapi/test_llm_api_pytorch_bart.py#L309-L316: move_SleepLogitsProcessorinto a shared helper module (e.g. alongside_get_model_path/_assert_bart_response-style shared utilities) and import it here.tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py#L569-L574: remove this duplicate definition and import the shared_SleepLogitsProcessorinstead.🤖 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 309 - 316, Move the duplicate _SleepLogitsProcessor from tests/integration/defs/llmapi/test_llm_api_pytorch_bart.py:309-316 into the shared encoder-decoder CUDA-graph test utility module and import it in BART. Remove the duplicate definition from tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py:569-574 and import the shared _SleepLogitsProcessor there.cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp (1)
582-728: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider validating
input_idscapacity for the generation segment too.The generation loop only guards
position_idscapacity (Line 688-689). Hostinput_idsslots for generation rows are intentionally left unwritten (they are filled device-side viaindex_selectin_prepare_encoder_decoder_inputs_fast), so this is not a bug today — but the asymmetry is easy to misread. A short comment stating that generationinput_idsare produced on device would make the contract explicit.🤖 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 `@cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp` around lines 582 - 728, Add a brief comment beside the generation-loop capacity checks in prepare_encoder_decoder_inputs explaining that generation input_ids slots are intentionally not written because _prepare_encoder_decoder_inputs_fast produces them device-side via index_select; do not add an input_ids capacity check.tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py (1)
897-906: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueNote the initialization-order dependency on
get_graph_key.
_build_encoder_decoder_capture_layouts()callsget_graph_key()whileself.capture_keysis still the empty frozenset, so it deliberately falls through to the_round_upbucketing branch rather than the dynamic_get_dynamic_capture_keybranch it later enables. That is load-bearing but invisible at the call site; a one-line comment would prevent a future reorder (e.g. assigningcapture_keysbefore the layout build) from silently changing which keys are derivable.Also applies to: 978-1014
🤖 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 `@tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py` around lines 897 - 906, Add a concise comment at the initialization sequence around `_build_encoder_decoder_capture_layouts()` explaining that it calls `get_graph_key()` while `self.capture_keys` is intentionally empty, forcing `_round_up` bucketing; preserve this ordering and assign `capture_keys` only after layout construction.tensorrt_llm/_torch/pyexecutor/py_executor.py (1)
5640-5645: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider logging the skip.
Dropping a completed encoder result silently makes the "request advanced past ENCODER_INIT" case invisible; a
logger.debugwith the request id and observed state would make future triage much cheaper.🤖 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 `@tensorrt_llm/_torch/pyexecutor/py_executor.py` around lines 5640 - 5645, Add a logger.debug call in the encoder-result handling branch surrounding req.state == LlmRequestState.ENCODER_INIT to record when a completed encoder result is skipped because the request has already advanced, including the request id and observed state; preserve the existing processing for ENCODER_INIT requests.tensorrt_llm/_torch/pyexecutor/model_engine.py (3)
549-563: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueNaming: encoder-only buckets stored under generic
_cuda_graph_*names.
self._cuda_graph_num_tokens/_cuda_graph_seq_lensnow hold strictly encoder values while the sibling batch-size list is_encoder_cuda_graph_batch_sizes. Renaming to_encoder_cuda_graph_num_tokens/_encoder_cuda_graph_seq_lens(and the_max_*variants) would keep the encoder bucket set self-describing and avoid future confusion with decoder graph state.🤖 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 `@tensorrt_llm/_torch/pyexecutor/model_engine.py` around lines 549 - 563, Rename the encoder-only bucket attributes in the model engine initialization from _cuda_graph_num_tokens and _cuda_graph_seq_lens to _encoder_cuda_graph_num_tokens and _encoder_cuda_graph_seq_lens, including their _max_* counterparts. Update all references to these attributes consistently while leaving decoder graph state unchanged.
3570-3618: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winPinned-buffer pool is unbounded.
Every acquire that finds no completed event allocates and appends another full set of pinned buffers (two
max_num_tokens-sized plus sixbatch_size-sized). Under any sustained backlog of un-completed copy events this grows without bound and pinned memory is never released. Consider capping the pool (e.g., a small N) and falling back toevent.synchronize()on the oldest entry once the cap is reached.🤖 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 `@tensorrt_llm/_torch/pyexecutor/model_engine.py` around lines 3570 - 3618, The _acquire_encoder_decoder_host_buffers pool currently grows without bound when all copy events remain pending. Cap the pool at a small fixed size, and when no completed buffer is available at the cap, synchronize the oldest entry’s event before reusing and returning it; preserve immediate reuse for entries whose event is None or already complete.
2214-2227: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDuplicated mixed-context request re-initialization.
The same state/positional/chunk resets are applied to the mixed context requests both here and in
_create_cuda_graph_warmup_request(Lines 2635-2641). Only the encoder-output/py_skip_cross_kv_projectionfields are unique to this site. Consider having the request builder own all of the bookkeeping (passing the per-request encoder output length) so the two sites cannot drift.🤖 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 `@tensorrt_llm/_torch/pyexecutor/model_engine.py` around lines 2214 - 2227, Move the shared mixed-context request initialization from the loop in the current batch setup into _create_cuda_graph_warmup_request, passing each request’s encoder output length so that builder initializes state, positions, chunk size, cached tokens, batch index, encoder output, and cross-KV projection flags. Remove the duplicated resets from the surrounding loop while preserving the per-request encoder output length behavior.tests/unittest/_torch/executor/test_py_executor.py (2)
298-393: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a failing-future test.
_poll_encoder_stepshas an error path (future.result()raising →_finish_failed_encoder_step) that clearsinflight_req_ids, pops the pending step, and calls_handle_errors._make_async_encoder_executoralready mocks_handle_errors, so a test withfuture.result.side_effect = RuntimeError(...)asserting the ids are released and the step is popped would be a few lines and covers the only path that can otherwise leak in-flight IDs.🤖 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/unittest/_torch/executor/test_py_executor.py` around lines 298 - 393, Add a test covering the failure path in _poll_encoder_steps by configuring the mocked future’s result() to raise RuntimeError. Submit and poll an encoder request, then assert _handle_errors is called, inflight_req_ids is cleared, pending_encoder_steps is empty, and no encoder output is published.Source: Path instructions
187-296: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUntested branch: the non-CUDA-graph fallback of
_waiting_encoder_requests.All six microbatch tests configure
encoder_cuda_graph_configwithnum_tokens/seq_lens, so they only exercise the graph-microbatch branch. The token-ratio fallback (py_executor.pyLines 5374-5386, reached whenencoder_max_batch_sizeor the graph shapes are unset) has no coverage — and that is the branch where a lone encoder request can dead-wait. Adding one test withencoder_cuda_graph_config=Noneplusbatch_wait_max_tokens_ratio/max_num_tokensset would pin that behavior.🤖 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/unittest/_torch/executor/test_py_executor.py` around lines 187 - 296, Add a unit test for the non-CUDA-graph fallback of _waiting_encoder_requests by configuring encoder_cuda_graph_config=None with batch_wait_max_tokens_ratio and max_num_tokens set. Use a lone encoder request and assert the fallback releases or waits according to the intended token-ratio behavior, including that it does not dead-wait; keep the existing graph-microbatch tests unchanged.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/attention_backend/trtllm.py`:
- Around line 668-678: Update the metadata preparation method containing the
host buffer assignments and kv_cache_manager.copy_batch_block_offsets call to
pass the configured max_blocks argument. Also refresh the persistent
prompt_lens_cpu and kv_lens host buffers from the current prompt_lens and
kv_lens inputs, not only their runtime views, so LoRA consumers read current
values after prepare_encoder_decoder().
In `@tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py`:
- Around line 190-215: Update _get_static_encoder_hidden_states and the
mixed-graph warmup flow so the shared encoder_hidden_states buffer is allocated
once using the maximum num_encoder_tokens across all mixed capture keys, rather
than the first key encountered. Ensure later capture and replay requests
validate that the existing buffer covers the requested extent, without
reallocating after CUDA graphs have been captured.
In `@tensorrt_llm/_torch/pyexecutor/model_engine.py`:
- Around line 565-569: In the use_encoder_cuda_graph condition, update the
_is_encode_only reference to use the boolean attribute directly rather than
calling it as a method. Preserve the existing short-circuit logic and all other
CUDA graph eligibility checks.
- Around line 3540-3557: Update the eligibility logic around
_encoder_decoder_input_fast_path_static_eligible so only
engine-lifetime-invariant predicates remain cached, while enable_spec_decode and
lora_model_config are evaluated on every call. Preserve the existing fast-path
requirements and combine the cached immutable result with fresh mutable-state
checks before using the path, so later speculation or LoRA configuration changes
take effect.
In `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Around line 5374-5386: Update the fallback scheduling logic around
`_waiting_requests` and the shown `should_wait` calculation so encoder requests
are released immediately when there are no generation requests or decoder work
in flight. Preserve the existing token-based and timeout-based waiting behavior
when decoder work exists, while avoiding batch-wait delay for a lone idle
encoder request.
---
Outside diff comments:
In `@tensorrt_llm/_torch/pyexecutor/model_engine.py`:
- Around line 6526-6535: Update the batch-size guard in the encoder CUDA graph
loop within the warmup/capture flow to compare each bs against
self.encoder_batch_size instead of the decoder self.batch_size. Preserve the
existing behavior for buckets within the encoder limit so all runtime-admitted
encoder batch sizes can be captured.
---
Nitpick comments:
In `@cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp`:
- Around line 582-728: Add a brief comment beside the generation-loop capacity
checks in prepare_encoder_decoder_inputs explaining that generation input_ids
slots are intentionally not written because _prepare_encoder_decoder_inputs_fast
produces them device-side via index_select; do not add an input_ids capacity
check.
In `@tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py`:
- Around line 897-906: Add a concise comment at the initialization sequence
around `_build_encoder_decoder_capture_layouts()` explaining that it calls
`get_graph_key()` while `self.capture_keys` is intentionally empty, forcing
`_round_up` bucketing; preserve this ordering and assign `capture_keys` only
after layout construction.
In `@tensorrt_llm/_torch/pyexecutor/model_engine.py`:
- Around line 549-563: Rename the encoder-only bucket attributes in the model
engine initialization from _cuda_graph_num_tokens and _cuda_graph_seq_lens to
_encoder_cuda_graph_num_tokens and _encoder_cuda_graph_seq_lens, including their
_max_* counterparts. Update all references to these attributes consistently
while leaving decoder graph state unchanged.
- Around line 3570-3618: The _acquire_encoder_decoder_host_buffers pool
currently grows without bound when all copy events remain pending. Cap the pool
at a small fixed size, and when no completed buffer is available at the cap,
synchronize the oldest entry’s event before reusing and returning it; preserve
immediate reuse for entries whose event is None or already complete.
- Around line 2214-2227: Move the shared mixed-context request initialization
from the loop in the current batch setup into _create_cuda_graph_warmup_request,
passing each request’s encoder output length so that builder initializes state,
positions, chunk size, cached tokens, batch index, encoder output, and cross-KV
projection flags. Remove the duplicated resets from the surrounding loop while
preserving the per-request encoder output length behavior.
In `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Around line 5640-5645: Add a logger.debug call in the encoder-result handling
branch surrounding req.state == LlmRequestState.ENCODER_INIT to record when a
completed encoder result is skipped because the request has already advanced,
including the request id and observed state; preserve the existing processing
for ENCODER_INIT requests.
In `@tensorrt_llm/_torch/pyexecutor/sampler/sampler.py`:
- Around line 1463-1465: Update the type annotations for
_stable_greedy_seq_slots_host and _stable_greedy_seq_slots_cuda to use
torch.Tensor | None instead of Optional[torch.Tensor], matching the PEP 604
style used elsewhere in the change; leave _stable_greedy_request_ids unchanged.
- Around line 2908-2934: Update the docstring of _fast_greedy_sample_kernel to
state that it computes and returns the translated greedy tokens while scattering
them into the output buffer, removing the inaccurate claim that all operations
are in-place.
- Around line 3663-3706: Update the cached branch of can_use_stable_greedy_path
so has_stable_request_ids does not bypass validation: re-check the same
per-request greedy invariants and verify each current py_seq_slot matches the
cached stable slot tensors before reusing them. If any invariant or slot mapping
differs, reject the stable path and rebuild or use the normal path, preventing
silent row-to-request misalignment.
In `@tests/integration/defs/llmapi/test_llm_api_pytorch_bart.py`:
- Around line 309-316: Move the duplicate _SleepLogitsProcessor from
tests/integration/defs/llmapi/test_llm_api_pytorch_bart.py:309-316 into the
shared encoder-decoder CUDA-graph test utility module and import it in BART.
Remove the duplicate definition from
tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py:569-574 and import the
shared _SleepLogitsProcessor there.
In `@tests/unittest/_torch/executor/test_py_executor.py`:
- Around line 298-393: Add a test covering the failure path in
_poll_encoder_steps by configuring the mocked future’s result() to raise
RuntimeError. Submit and poll an encoder request, then assert _handle_errors is
called, inflight_req_ids is cleared, pending_encoder_steps is empty, and no
encoder output is published.
- Around line 187-296: Add a unit test for the non-CUDA-graph fallback of
_waiting_encoder_requests by configuring encoder_cuda_graph_config=None with
batch_wait_max_tokens_ratio and max_num_tokens set. Use a lone encoder request
and assert the fallback releases or waits according to the intended token-ratio
behavior, including that it does not dead-wait; keep the existing
graph-microbatch tests unchanged.
In `@tests/unittest/_torch/sampler/test_torch_sampler.py`:
- Around line 666-743: The new caching and eligibility logic in
TorchSampler._process_requests, particularly can_use_stable_greedy_path and
has_stable_request_ids, lacks direct tests. Add tests covering initial fast-path
eligibility, reuse of cached eligibility, invalidation when request IDs change,
and preservation of cached seq_slots for unchanged request IDs; retain the
existing update_requests coverage.
🪄 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: 06e13f9e-fb86-4b22-aa1b-f4b1179159a0
📒 Files selected for processing (18)
cpp/tensorrt_llm/nanobind/batch_manager/bindings.cppdocs/source/models/encoder-decoder.mdtensorrt_llm/_torch/attention_backend/trtllm.pytensorrt_llm/_torch/pyexecutor/cuda_graph_runner.pytensorrt_llm/_torch/pyexecutor/model_engine.pytensorrt_llm/_torch/pyexecutor/py_executor.pytensorrt_llm/_torch/pyexecutor/sampler/sampler.pytensorrt_llm/llmapi/llm_args.pytensorrt_llm/usage/llm_args_golden_manifest.jsontests/integration/defs/llmapi/test_llm_api_pytorch_bart.pytests/integration/defs/llmapi/test_llm_api_pytorch_t5.pytests/integration/test_lists/test-db/l0_h100.ymltests/integration/test_lists/test-db/l0_l40s.ymltests/unittest/_torch/executor/test_py_executor.pytests/unittest/_torch/executor/test_pytorch_model_engine_warmup.pytests/unittest/_torch/sampler/test_torch_sampler.pytests/unittest/api_stability/references/llm.yamltests/unittest/llmapi/test_llm_args.py
💤 Files with no reviewable changes (1)
- tests/integration/test_lists/test-db/l0_l40s.yml
Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com>
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py (1)
448-456: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winUse the expanded generation count for mixed CUDA-graph metadata
With one context request, one generation request, and
max_beam_width=2, the batch has three sequences, but lines 448–456 build four sequence lengths. Use(batch_size - len(context_query_lens)) * self.max_beam_widthfor the generation tail. Apply the same count in_get_num_tokens_for_key.🤖 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 `@tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py` around lines 448 - 456, Mixed encoder-decoder CUDA-graph metadata uses an incorrect generation tail count when beam width expands the batch. Update the sequence-length construction in the mixed branch to use the expanded generation count based on batch size, context count, and max_beam_width, and apply the same calculation in _get_num_tokens_for_key.
🧹 Nitpick comments (4)
tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py (1)
1050-1052: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRemove the discarded
sum(sequence_lengths)statement.Line 1051 computes a total and discards it. The authoritative total is
required_num_tokensat line 1061, which is recomputed per candidate batch size. The dead statement suggests that an unpadded aggregate check exists at this point, and a later edit could be anchored on it.♻️ Proposed cleanup
batch_size = len(sequence_lengths) - sum(sequence_lengths) max_seq_len = max(sequence_lengths) if sequence_lengths else 0🤖 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 `@tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py` around lines 1050 - 1052, Remove the standalone discarded sum(sequence_lengths) expression from the batch setup near batch_size and max_seq_len, leaving required_num_tokens as the authoritative total used by the candidate batch-size logic.tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py (1)
878-902: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExtract the shared replay-recording setup into a helper.
tests/integration/defs/llmapi/test_llm_api_pytorch_bart.pylines 675-699 contain the same runner lookup,key[5] and key[6]mixed-key filter, and replay-recording wrappers. The magic tuple indices are duplicated in both files. Move this setup into a shared test helper so a change toKeyTypeneeds one update.🤖 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 878 - 902, Extract the duplicated replay-recording setup from the test around encoder_runner and decoder_runner into a shared test helper, reusing it from both the T5 test and the corresponding BART test. The helper should perform runner lookup, filter decoder graphs using the mixed-key condition, wrap both replay methods, and return the captured replay keys; keep the key-index logic centralized there.tensorrt_llm/_torch/pyexecutor/model_engine.py (1)
7164-7166: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winConsider replacing
assertwith an explicit exception for the capacity check.
assert num_tokens <= self.encoder_max_num_tokensis stripped when Python runs with-O. If that happens, an oversized feature-driven encoder batch silently proceeds instead of failing, and downstream attention metadata is sized toself.encoder_max_num_tokens, which can misbehave for larger inputs. The same pattern exists elsewhere in this file (for example lines 5121, 5714, 7095-7097), so a full fix is a broader, separate cleanup; flagging here since this instance is new in this diff.🤖 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 `@tensorrt_llm/_torch/pyexecutor/model_engine.py` around lines 7164 - 7166, Replace the capacity-check assert in the encoder batch path with an explicit runtime exception that is always enforced, preserving the existing oversized-input message and failure behavior. Update only this check near the encoder metadata sizing logic; do not broaden the change to other assertions in the file.tensorrt_llm/_torch/pyexecutor/py_executor.py (1)
5486-5486: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winNarrow the blind
except Exceptionin the TP-serialization path.Static analysis flags
except Exception as e:at line 5515 (Ruff BLE001). As per coding guidelines,**/*.{py,cpp,cc,cxx,h,hpp}files must "Catch specific exceptions instead of using broad or bare exception handling such asexcept:." This is the same top-level "boundary" pattern used elsewhere in this file (for example_forward_step,_sample_async) to keep the executor loop alive and route failures through_handle_errors/_finish_failed_encoder_step, so a full fix likely needs a broader, coordinated pass rather than a point fix here.Also applies to: 5495-5521
🤖 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 `@tensorrt_llm/_torch/pyexecutor/py_executor.py` at line 5486, Replace the broad `except Exception as e` in the TP serialization path around the encoder work queue with specific exception types covering the expected failures, while preserving routing through `_handle_errors` and `_finish_failed_encoder_step`. Review the analogous boundary handling in `_forward_step` and `_sample_async` and coordinate the change so executor-loop liveness and existing failure behavior remain unchanged without retaining a blind catch-all.Sources: Coding guidelines, Linters/SAST tools
🤖 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 `@docs/source/models/supported-models.md`:
- Line 102: Change the “Encoder-Decoder Feature Support Matrix (PyTorch
Backend)” heading from level 1 to level 2, matching the existing matrix section
and preserving “Supported Models” as the document’s only top-level heading.
In `@tests/integration/defs/llmapi/test_llm_api_pytorch_bart.py`:
- Around line 741-743: The encoder microbatch assertions incorrectly require
coalescing into a single key with key[0] == 2, which can fail when fallback
admission releases requests separately. Update
tests/integration/defs/llmapi/test_llm_api_pytorch_bart.py lines 741-743 and
tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py lines 948-949 to use an
admission-independent assertion that the admitted encoder keys jointly cover
both requests, or increase batch_wait_timeout_iters so coalescing is guaranteed;
apply the same behavior in both tests.
In `@tests/unittest/_torch/executor/test_pytorch_model_engine.py`:
- Around line 185-209: Update the test input in the staging assertions to use
non-zero-starting token IDs instead of torch.arange(401), ensuring input_ids[0]
is non-zero. Keep the expected staging logic and assertions unchanged so
expected_staged_ids[511] verifies that slot 1 receives the first sequence.
---
Outside diff comments:
In `@tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py`:
- Around line 448-456: Mixed encoder-decoder CUDA-graph metadata uses an
incorrect generation tail count when beam width expands the batch. Update the
sequence-length construction in the mixed branch to use the expanded generation
count based on batch size, context count, and max_beam_width, and apply the same
calculation in _get_num_tokens_for_key.
---
Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py`:
- Around line 1050-1052: Remove the standalone discarded sum(sequence_lengths)
expression from the batch setup near batch_size and max_seq_len, leaving
required_num_tokens as the authoritative total used by the candidate batch-size
logic.
In `@tensorrt_llm/_torch/pyexecutor/model_engine.py`:
- Around line 7164-7166: Replace the capacity-check assert in the encoder batch
path with an explicit runtime exception that is always enforced, preserving the
existing oversized-input message and failure behavior. Update only this check
near the encoder metadata sizing logic; do not broaden the change to other
assertions in the file.
In `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Line 5486: Replace the broad `except Exception as e` in the TP serialization
path around the encoder work queue with specific exception types covering the
expected failures, while preserving routing through `_handle_errors` and
`_finish_failed_encoder_step`. Review the analogous boundary handling in
`_forward_step` and `_sample_async` and coordinate the change so executor-loop
liveness and existing failure behavior remain unchanged without retaining a
blind catch-all.
In `@tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py`:
- Around line 878-902: Extract the duplicated replay-recording setup from the
test around encoder_runner and decoder_runner into a shared test helper, reusing
it from both the T5 test and the corresponding BART test. The helper should
perform runner lookup, filter decoder graphs using the mixed-key condition, wrap
both replay methods, and return the captured replay keys; keep the key-index
logic centralized there.
🪄 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: 6cd063c3-9ea8-46ca-8afd-b346f153faa2
📒 Files selected for processing (16)
docs/source/models/encoder-decoder.mddocs/source/models/supported-models.mdtensorrt_llm/_torch/attention_backend/trtllm.pytensorrt_llm/_torch/pyexecutor/cuda_graph_runner.pytensorrt_llm/_torch/pyexecutor/model_engine.pytensorrt_llm/_torch/pyexecutor/py_executor.pytensorrt_llm/_torch/pyexecutor/sampler/sampler.pytests/integration/defs/llmapi/test_llm_api_pytorch_bart.pytests/integration/defs/llmapi/test_llm_api_pytorch_t5.pytests/integration/test_lists/test-db/l0_dgx_h100.ymltests/integration/test_lists/test-db/l0_h100.ymltests/unittest/_torch/executor/test_py_executor.pytests/unittest/_torch/executor/test_pytorch_model_engine.pytests/unittest/_torch/executor/test_pytorch_model_engine_warmup.pytests/unittest/_torch/sampler/test_torch_sampler.pytests/unittest/llmapi/test_llm_args.py
🚧 Files skipped from review as they are similar to previous changes (5)
- tests/integration/test_lists/test-db/l0_h100.yml
- tests/unittest/llmapi/test_llm_args.py
- tests/unittest/_torch/sampler/test_torch_sampler.py
- tensorrt_llm/_torch/attention_backend/trtllm.py
- tensorrt_llm/_torch/pyexecutor/sampler/sampler.py
|
/bot run |
|
PR_Github #62953 [ run ] triggered by Bot. Commit: |
|
PR_Github #62953 [ run ] completed with state
|
Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com>
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (6)
tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py (3)
1082-1083: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRemove the dead
sum(sequence_lengths)statement.Line 1082 computes a value and discards it. The needed total is recomputed per padded batch size at line 1092.
♻️ Proposed cleanup
batch_size = len(sequence_lengths) - sum(sequence_lengths) max_seq_len = max(sequence_lengths) if sequence_lengths else 0🤖 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 `@tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py` around lines 1082 - 1083, Remove the standalone sum(sequence_lengths) expression in the sequence-length handling code, leaving the max_seq_len calculation and the existing per-padded-batch total computation unchanged.
330-345: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueKey lookup and capture bookkeeping use different dictionaries.
_get_compatible_mixed_encoder_decoder_keyscansself.graph_outputs, butmaybe_get_cuda_graphresolves a hit fromself.graph_metadata. In warmup-only modecapturepopulatesgraph_metadataand returns before it writesgraph_outputs, so compatible-key selection sees no candidates while metadata exists. Scanself.graph_metadatafor consistency, or state whygraph_outputsis the intended source.🤖 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 `@tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py` around lines 330 - 345, Update _get_compatible_mixed_encoder_decoder_key to scan self.graph_metadata, matching the dictionary used by maybe_get_cuda_graph for key resolution. Preserve the existing compatibility filters and minimum encoder-token selection, and avoid relying on graph_outputs during warmup-only capture.
170-174: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueSimplify the redundant clamp.
Line 173 assigns
self.config.max_num_tokens, and line 174 clamps the same value againstself.config.max_num_tokens. The clamp is a no-op in the mixed path. Keep the branch, or fold both lines into onemin/maxexpression for clarity.🤖 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 `@tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py` around lines 170 - 174, Remove the redundant max_total_tokens clamp in the mixed CUDA graph setup around enable_encoder_decoder_mixed_cuda_graph, since that branch already assigns self.config.max_num_tokens. Preserve the existing branch behavior while simplifying the assignment so max_total_tokens is computed only once.tensorrt_llm/_torch/pyexecutor/model_engine.py (2)
2239-2243: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueExpose a public runner method instead of calling
_get_static_encoder_hidden_states.
_capture_mixed_encoder_decoder_cuda_graphsreaches into a private member ofCUDAGraphRunnerto pre-allocate the mixed-graph encoder buffer. The runner already raises when the buffer is missing at replay time, so the allocation is a real contract between the two classes. Add a public method on the runner (for examplereserve_encoder_hidden_states(num_tokens, dtype, hidden_size)) and call that here, so the buffer lifecycle stays owned by the runner.🤖 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 `@tensorrt_llm/_torch/pyexecutor/model_engine.py` around lines 2239 - 2243, Replace the direct `_get_static_encoder_hidden_states` call in `_capture_mixed_encoder_decoder_cuda_graphs` with a public `CUDAGraphRunner` method for reserving the encoder hidden-state buffer, passing the required token count, dtype, and hidden size. Implement the method on `CUDAGraphRunner` so buffer allocation remains owned by the runner while preserving the existing pre-allocation behavior.
531-545: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRename the encoder bucket attributes for consistency.
self._cuda_graph_num_tokens,self._max_cuda_graph_num_tokens,self._cuda_graph_seq_lens, andself._max_cuda_graph_seq_lenare all derived fromencoder_cuda_graph_*inputs, but their names omit theencoderprefix thatself._encoder_cuda_graph_batch_sizesandself._encoder_cuda_graph_padding_enableduse. A reader of_capture_encoder_cuda_graphs(lines 6556-6558) seesself._encoder_cuda_graph_batch_sizesnext toself._cuda_graph_num_tokensand cannot tell that both are encoder buckets. Rename them to_encoder_cuda_graph_num_tokens,_max_encoder_cuda_graph_num_tokens,_encoder_cuda_graph_seq_lens, and_max_encoder_cuda_graph_seq_len.🤖 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 `@tensorrt_llm/_torch/pyexecutor/model_engine.py` around lines 531 - 545, Rename the encoder-derived bucket attributes throughout the class: replace _cuda_graph_num_tokens with _encoder_cuda_graph_num_tokens, _max_cuda_graph_num_tokens with _max_encoder_cuda_graph_num_tokens, _cuda_graph_seq_lens with _encoder_cuda_graph_seq_lens, and _max_cuda_graph_seq_len with _max_encoder_cuda_graph_seq_len. Update every reference, including _capture_encoder_cuda_graphs, while preserving their existing initialization and behavior.tensorrt_llm/_torch/pyexecutor/sampler/sampler.py (1)
3842-3850: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider a named return type for
_process_requests.The return value is now a 7-element positional tuple with two
Optionalentries and a trailingbool. ANamedTupleor dataclass would make each field self-describing at both the definition and the call site, and would prevent ordering mistakes when the tuple grows again.🤖 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 `@tensorrt_llm/_torch/pyexecutor/sampler/sampler.py` around lines 3842 - 3850, Introduce a named return type for _process_requests, such as a NamedTuple or dataclass, with descriptive fields for all seven returned values, including the Optional entries and trailing bool. Update _process_requests and its callers to construct and access this type by field name while preserving the existing values and behavior.
🤖 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 1868-1880: Add `@torch.inference_mode`() to
_warmup_encoder_decoder_encoder_cuda_graphs so encoder-decoder encoder CUDA
graph capture runs with autograd disabled, matching the existing capture and
replay paths. Preserve the current warmup and capture flow unchanged.
- Around line 6556-6558: Update the encoder graph-capture eligibility guard near
the batch-size list initialization to compare against self.encoder_batch_size
rather than the general max_batch_size, while preserving the existing encoder
bucket sorting and capture behavior.
In `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Around line 5349-5375: Update the encoder scheduling logic around
deadline_reached so encoder requests are released when the wait deadline is
reached even if decoder_occupancy exceeds decoder_low_watermark. Preserve the
existing preferred microbatch selection and supported-batch-size fallback, but
ensure the busy-decoder path cannot continue incrementing
encoder_batch_wait_iters_count and returning [] indefinitely.
- Around line 4098-4100: Update the not-can-queue encoder fallback around
_run_encoder_step to dispatch graph-capable encoder requests through
encoder_launch_executor, allowing forward_encoder’s captured-graph replay to run
on the owning worker. Preserve the existing completion signaling and error
propagation when submitting and awaiting this executor work.
In `@tensorrt_llm/_torch/pyexecutor/sampler/sampler.py`:
- Around line 3855-3877: Update the can_use_stable_greedy_path logic around
has_stable_request_ids so volatile per-iteration checks remain unconditional:
require every generation request to be non-dummy and have
get_draft_token_length(request) == 0 even when request IDs are unchanged. Keep
the remaining static feature checks cached behind the stable-request-ID
shortcut, preserving the existing greedy-path conditions for those properties.
---
Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py`:
- Around line 1082-1083: Remove the standalone sum(sequence_lengths) expression
in the sequence-length handling code, leaving the max_seq_len calculation and
the existing per-padded-batch total computation unchanged.
- Around line 330-345: Update _get_compatible_mixed_encoder_decoder_key to scan
self.graph_metadata, matching the dictionary used by maybe_get_cuda_graph for
key resolution. Preserve the existing compatibility filters and minimum
encoder-token selection, and avoid relying on graph_outputs during warmup-only
capture.
- Around line 170-174: Remove the redundant max_total_tokens clamp in the mixed
CUDA graph setup around enable_encoder_decoder_mixed_cuda_graph, since that
branch already assigns self.config.max_num_tokens. Preserve the existing branch
behavior while simplifying the assignment so max_total_tokens is computed only
once.
In `@tensorrt_llm/_torch/pyexecutor/model_engine.py`:
- Around line 2239-2243: Replace the direct `_get_static_encoder_hidden_states`
call in `_capture_mixed_encoder_decoder_cuda_graphs` with a public
`CUDAGraphRunner` method for reserving the encoder hidden-state buffer, passing
the required token count, dtype, and hidden size. Implement the method on
`CUDAGraphRunner` so buffer allocation remains owned by the runner while
preserving the existing pre-allocation behavior.
- Around line 531-545: Rename the encoder-derived bucket attributes throughout
the class: replace _cuda_graph_num_tokens with _encoder_cuda_graph_num_tokens,
_max_cuda_graph_num_tokens with _max_encoder_cuda_graph_num_tokens,
_cuda_graph_seq_lens with _encoder_cuda_graph_seq_lens, and
_max_cuda_graph_seq_len with _max_encoder_cuda_graph_seq_len. Update every
reference, including _capture_encoder_cuda_graphs, while preserving their
existing initialization and behavior.
In `@tensorrt_llm/_torch/pyexecutor/sampler/sampler.py`:
- Around line 3842-3850: Introduce a named return type for _process_requests,
such as a NamedTuple or dataclass, with descriptive fields for all seven
returned values, including the Optional entries and trailing bool. Update
_process_requests and its callers to construct and access this type by field
name while preserving the existing values and behavior.
🪄 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: a4f3eb06-e625-43cf-8dd0-73b03c490ca7
📒 Files selected for processing (11)
docs/source/models/supported-models.mdtensorrt_llm/_torch/pyexecutor/cuda_graph_runner.pytensorrt_llm/_torch/pyexecutor/model_engine.pytensorrt_llm/_torch/pyexecutor/py_executor.pytensorrt_llm/_torch/pyexecutor/sampler/sampler.pytensorrt_llm/llmapi/llm_args.pytests/integration/test_lists/test-db/l0_h100.ymltests/unittest/_torch/executor/test_pytorch_model_engine.pytests/unittest/_torch/sampler/test_torch_sampler.pytests/unittest/api_stability/references/llm.yamltests/unittest/llmapi/test_llm_args.py
🚧 Files skipped from review as they are similar to previous changes (7)
- tests/unittest/api_stability/references/llm.yaml
- tests/integration/test_lists/test-db/l0_h100.yml
- tests/unittest/llmapi/test_llm_args.py
- docs/source/models/supported-models.md
- tests/unittest/_torch/executor/test_pytorch_model_engine.py
- tests/unittest/_torch/sampler/test_torch_sampler.py
- tensorrt_llm/llmapi/llm_args.py
Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com>
…-perf Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com> # Conflicts: # tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py # tensorrt_llm/_torch/pyexecutor/model_engine.py # tests/unittest/_torch/executor/test_pytorch_model_engine.py
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tensorrt_llm/_torch/pyexecutor/py_executor.py (1)
4100-4102: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winSerialize the fallback encoder step under TP.
ScheduledRequests.batch_sizeexcludes encoder requests, so encoder-only batches reach this fallback in both executor loops. Whenself.dist.tp_size > 1,_run_encoder_stepmust callself.encoder_stream.wait_stream(self.execution_stream)before submission, as_submit_encoder_stepdoes. Otherwise, a previous decoder forward can overlap encoder collectives on shared TP communicators and workspaces.🤖 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 `@tensorrt_llm/_torch/pyexecutor/py_executor.py` around lines 4100 - 4102, Update the fallback encoder-step path around _run_encoder_step so that, when self.dist.tp_size > 1, it waits on self.execution_stream via self.encoder_stream.wait_stream before submitting encoder work. Match the synchronization behavior in _submit_encoder_step and preserve the existing behavior for single-rank TP.
♻️ Duplicate comments (1)
tensorrt_llm/_torch/pyexecutor/py_executor.py (1)
5373-5399: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winThe CUDA-graph branch still has no deadline when the decoder stays busy.
deadline_reachedis computed on Line 5379 but is only read inside thedecoder_occupancy <= decoder_low_watermarkbranch. When occupancy stays above the watermark, Lines 5398-5399 increment the counter and return[]on every iteration with no bound.
decoder_low_watermarkisself.max_batch_size - microbatch_target. When the largest supported graph batch size equalsmax_batch_size, the watermark is0. One ongoing generation request then keepsdecoder_occupancy > 0for its whole lifetime, so admitted encoder requests are never released and their decoder-context steps never start.batch_wait_timeout_itersdoes not bound this case.The token-threshold branch below now guards against a related stall through
has_decoder_work(Lines 5405-5410). Apply an equivalent bound here.🐛 Proposed fix to bound the wait by the existing iteration deadline
- if decoder_occupancy <= decoder_low_watermark: - if len(encoder_requests) >= microbatch_target: - self.encoder_batch_wait_iters_count = 0 - return encoder_requests[:microbatch_target] - - if deadline_reached: - releasable_batch_sizes = [ - batch_size for batch_size in supported_batch_sizes - if batch_size <= len(encoder_requests) - ] - fallback_batch_size = ( - releasable_batch_sizes[-1] if releasable_batch_sizes - else min(len(encoder_requests), microbatch_target)) - self.encoder_batch_wait_iters_count = 0 - return encoder_requests[:fallback_batch_size] + if (decoder_occupancy <= decoder_low_watermark + and len(encoder_requests) >= microbatch_target): + self.encoder_batch_wait_iters_count = 0 + return encoder_requests[:microbatch_target] + + if deadline_reached: + releasable_batch_sizes = [ + batch_size for batch_size in supported_batch_sizes + if batch_size <= len(encoder_requests) + ] + fallback_batch_size = ( + releasable_batch_sizes[-1] if releasable_batch_sizes + else min(len(encoder_requests), microbatch_target)) + self.encoder_batch_wait_iters_count = 0 + return encoder_requests[:fallback_batch_size] self.encoder_batch_wait_iters_count += 1 return []🤖 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 `@tensorrt_llm/_torch/pyexecutor/py_executor.py` around lines 5373 - 5399, Update the CUDA-graph batching flow around deadline_reached and the decoder_occupancy <= decoder_low_watermark branch so the existing batch-wait deadline also releases encoder_requests when decoder occupancy remains above the watermark. Reuse the existing releasable_batch_sizes, fallback_batch_size, counter reset, and return behavior from the deadline path, while preserving normal waiting when the deadline has not been reached.
🧹 Nitpick comments (4)
tensorrt_llm/_torch/pyexecutor/model_engine.py (3)
4906-4909: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winThree new sites read one prompt token through the full-vector accessor.
request.get_tokens(0)marshals the entireO(prompt_len)VecTokensinto a Python list of boxed ints, and each site then indexes a single element.request.get_tokens_range(0, begin, end)copies only the requested window; this file already documents that cost difference at Line 4676.
tensorrt_llm/_torch/pyexecutor/model_engine.py#L4906-L4909: in the extend-request loop, replacerequest.get_tokens(0)[request.context_current_position]withrequest.get_tokens_range(0, position, position + 1)[0]using a localposition = request.context_current_position.tensorrt_llm/_torch/pyexecutor/model_engine.py#L5064-L5067: apply the same replacement in the generation loop.tensorrt_llm/_torch/pyexecutor/model_engine.py#L3433-L3436: apply the same replacement in_is_final_multimodal_context_decode_compatiblebefore calling_prepare_multimodal_indices.🤖 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 `@tensorrt_llm/_torch/pyexecutor/model_engine.py` around lines 4906 - 4909, Replace full-vector token access with a one-token range fetch in all three sites: the extend-request loop at tensorrt_llm/_torch/pyexecutor/model_engine.py:4906-4909, the generation loop at tensorrt_llm/_torch/pyexecutor/model_engine.py:5064-5067, and _is_final_multimodal_context_decode_compatible at tensorrt_llm/_torch/pyexecutor/model_engine.py:3433-3436. In each location, assign the relevant position to a local variable and use get_tokens_range(0, position, position + 1)[0] before preserving the existing downstream behavior.
7256-7258: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winRaise an exception instead of asserting the encoder packed length.
assertstatements are removed when Python runs with-O. This check guards a runtime-reachable input size againstencoder_max_num_tokens, which also sizes the encoder attention metadata buffers (Line 7217). Without the check, an oversized packed batch writes past those buffers instead of failing.The sibling check in
_prepare_tp_inputs_encoder_features(Line 7325) has the same problem.🛡️ Proposed change
- assert num_tokens <= self.encoder_max_num_tokens, ( - f"encoder packed length ({num_tokens}) exceeds " - f"encoder_max_num_tokens ({self.encoder_max_num_tokens})") + if num_tokens > self.encoder_max_num_tokens: + raise ValueError( + f"encoder packed length ({num_tokens}) exceeds " + f"encoder_max_num_tokens ({self.encoder_max_num_tokens})")🤖 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 `@tensorrt_llm/_torch/pyexecutor/model_engine.py` around lines 7256 - 7258, Replace the runtime assert in the encoder input preparation flow with an unconditional exception when num_tokens exceeds self.encoder_max_num_tokens, preserving the existing diagnostic details. Apply the same change to the sibling check in _prepare_tp_inputs_encoder_features so both paths reject oversized packed batches even under optimized Python execution.
579-593: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRename the encoder-only CUDA graph bucket fields with an
encoderprefix.Rename the four fields and update their consumers. This makes their encoder scope explicit and matches
_encoder_cuda_graph_batch_sizes.🤖 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 `@tensorrt_llm/_torch/pyexecutor/model_engine.py` around lines 579 - 593, Rename the encoder-only fields _cuda_graph_num_tokens, _max_cuda_graph_num_tokens, _cuda_graph_seq_lens, and _max_cuda_graph_seq_len to use an encoder prefix, and update every consumer to reference the renamed fields consistently, matching the naming style of _encoder_cuda_graph_batch_sizes.tests/unittest/_torch/executor/test_pytorch_model_engine.py (1)
1-1961: 📐 Maintainability & Code Quality | 🔵 TrivialTest coverage summary (QA).
This is a unit test file under
tests/unittest/, so test-list registration intests/integration/test_lists/does not apply.Added/modified test functions:
SingleTokenContextGraphBatchTestCase:test_generation_only_is_identity,test_eligible_batch_has_independent_lists_and_stable_order,test_structural_fallbacks_return_semantic_batch,test_context_shape_and_mode_fallback_matrix,test_context_logits_use_final_token_graph_candidate,test_generation_only_request_in_context_list_falls_back,test_generation_shape_fallback_matrix,test_mixed_one_and_two_token_contexts_fall_back_together,test_mrope_delta_is_supported_by_decode_provider,test_multimodal_context_requires_compatible_decode_token,test_multimodal_pending_event_is_rechecked,test_multimodal_decode_compatibility_uses_final_prompt_token,test_sparse_sequence_mode_uses_promoted_context_cursor,test_graph_key_forwards_promoted_context_ids,test_graph_lookup_forwards_promoted_context_ids,test_forward_commits_candidate_only_on_graph_hit,test_forward_graph_miss_uses_semantic_eager_batch,test_zero_runtime_draft_speculation_commits_graph_candidate,test_zero_runtime_draft_speculation_graph_miss_is_semantic_eager,test_zero_runtime_non_linear_tree_speculation_uses_semantic_eager_batch,test_forward_allows_guided_context_logits_on_graph_hit,test_multimodal_graph_miss_preserves_semantic_payload,test_generation_only_forward_does_not_call_new_selector,test_global_incompatibilities_bypass_candidate_selection.PyTorchModelEngineTestCaseadditions:test_promoted_context_uses_prompt_token_during_overlap,test_promoted_context_precedes_speculative_overlap_generation,test_promoted_mrope_context_uses_decode_state_contract,test_cuda_graph_replay_observes_execution_stream_dependency.Coverage verdict: sufficient. The new tests cover single-token context promotion eligibility (structural, shape, multimodal, MRoPE, speculative fallback matrices), CUDA-graph key/lookup forwarding of promoted context IDs,
PyTorchModelEngine.forwardgraph-hit/miss/speculative dispatch, promoted-context input preparation (including MRoPE deltas and KV-cache accounting), and CUDA-graph replay's cross-stream dependency on restored KV data. One pre-existing gap remains: the vacuous assertion flagged above (lines 968-992) is still unaddressed.🤖 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/unittest/_torch/executor/test_pytorch_model_engine.py` around lines 1 - 1961, The vacuous assertion in SingleTokenContextGraphBatchTestCase remains ineffective. Replace it with an assertion that verifies the intended result of the test scenario, or remove it if no meaningful condition can be checked; preserve the surrounding test’s existing behavior and coverage.
🤖 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.
Outside diff comments:
In `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Around line 4100-4102: Update the fallback encoder-step path around
_run_encoder_step so that, when self.dist.tp_size > 1, it waits on
self.execution_stream via self.encoder_stream.wait_stream before submitting
encoder work. Match the synchronization behavior in _submit_encoder_step and
preserve the existing behavior for single-rank TP.
---
Duplicate comments:
In `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Around line 5373-5399: Update the CUDA-graph batching flow around
deadline_reached and the decoder_occupancy <= decoder_low_watermark branch so
the existing batch-wait deadline also releases encoder_requests when decoder
occupancy remains above the watermark. Reuse the existing
releasable_batch_sizes, fallback_batch_size, counter reset, and return behavior
from the deadline path, while preserving normal waiting when the deadline has
not been reached.
---
Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/model_engine.py`:
- Around line 4906-4909: Replace full-vector token access with a one-token range
fetch in all three sites: the extend-request loop at
tensorrt_llm/_torch/pyexecutor/model_engine.py:4906-4909, the generation loop at
tensorrt_llm/_torch/pyexecutor/model_engine.py:5064-5067, and
_is_final_multimodal_context_decode_compatible at
tensorrt_llm/_torch/pyexecutor/model_engine.py:3433-3436. In each location,
assign the relevant position to a local variable and use get_tokens_range(0,
position, position + 1)[0] before preserving the existing downstream behavior.
- Around line 7256-7258: Replace the runtime assert in the encoder input
preparation flow with an unconditional exception when num_tokens exceeds
self.encoder_max_num_tokens, preserving the existing diagnostic details. Apply
the same change to the sibling check in _prepare_tp_inputs_encoder_features so
both paths reject oversized packed batches even under optimized Python
execution.
- Around line 579-593: Rename the encoder-only fields _cuda_graph_num_tokens,
_max_cuda_graph_num_tokens, _cuda_graph_seq_lens, and _max_cuda_graph_seq_len to
use an encoder prefix, and update every consumer to reference the renamed fields
consistently, matching the naming style of _encoder_cuda_graph_batch_sizes.
In `@tests/unittest/_torch/executor/test_pytorch_model_engine.py`:
- Around line 1-1961: The vacuous assertion in
SingleTokenContextGraphBatchTestCase remains ineffective. Replace it with an
assertion that verifies the intended result of the test scenario, or remove it
if no meaningful condition can be checked; preserve the surrounding test’s
existing behavior and coverage.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 056d41a7-ef59-44e3-ad7c-6397d96be505
📒 Files selected for processing (12)
tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.pytensorrt_llm/_torch/pyexecutor/model_engine.pytensorrt_llm/_torch/pyexecutor/py_executor.pytensorrt_llm/_torch/pyexecutor/sampler/sampler.pytensorrt_llm/llmapi/llm_args.pytensorrt_llm/usage/llm_args_golden_manifest.jsontests/integration/defs/kv_cache/test_final_single_token_context_cuda_graph.pytests/integration/test_lists/test-db/l0_dgx_h100.ymltests/integration/test_lists/test-db/l0_h100.ymltests/unittest/_torch/executor/test_py_executor.pytests/unittest/_torch/executor/test_pytorch_model_engine.pytests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py
🚧 Files skipped from review as they are similar to previous changes (8)
- tests/integration/test_lists/test-db/l0_dgx_h100.yml
- tests/integration/test_lists/test-db/l0_h100.yml
- tensorrt_llm/usage/llm_args_golden_manifest.json
- tensorrt_llm/llmapi/llm_args.py
- tensorrt_llm/_torch/pyexecutor/sampler/sampler.py
- tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
- tests/unittest/_torch/executor/test_py_executor.py
- tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py
|
/bot run |
|
PR_Github #63174 [ run ] triggered by Bot. Commit: |
|
PR_Github #63174 [ run ] completed with state
|
3 participants
Summary
Why
BART continuous-admission serving repeatedly mixes replacement encoder requests with active decoder requests. The existing path launched small encoder batches, rebuilt mixed decoder metadata in Python, and reconstructed sampling/finish metadata on every greedy decode step. It also ran eligible encoder and mixed decoder work eagerly instead of replaying fixed-shape CUDA graphs.
These changes reduce per-iteration CPU launch and device-to-host overhead while keeping decoder generation active during encoder admission. The CUDA-graph capture layout is derived from the encoder runner's effective capture keys so padding and runtime replay use the same buckets.
User interface
For encoder-decoder models,
encoder_cuda_graph_config=EncodeCudaGraphConfig(...)enables encoder-forward CUDA graphs and defines batch-size, total packed-token, and maximum sequence-length buckets.encoder_max_batch_sizeremains the hard encoder capacity and admission limit.enable_encoder_decoder_mixed_cuda_graphcontrols the mixed-batch decoder optimization. It defaults toTruebut becomes effective only when both the encoder and decoder graph configurations produce usable capture shapes. Set it toFalseto retain separate encoder and decoder CUDA graphs while disabling mixed-batch graphs.Performance
BART-large-CNN
This comparison uses the same deterministic natural-length CNN/DailyMail workload on both backends: 1,024 unique validation articles sampled without replacement using seed 0, encoder lengths of 93–1,022 tokens, and no truncation.
PyTorch improves throughput by 12.7–20.9% and reduces mean latency by 11.0–17.2% compared with legacy TensorRT.
More BART workload, latency, distribution, and configuration details
Detailed latency percentiles
Encoder input-length distribution
Lengths include the tokenizer's special tokens.
Configuration
PyTorch configuration:
encoder_max_batch_size=2at concurrency 8 and8at concurrency 32/64[1, 2]at concurrency 8 and[1, 2, 4, 8]at concurrency 32/64[512, 1024]Legacy TensorRT configuration:
Generated outputs were not bit-identical: PyTorch averaged approximately 73.37 output tokens per request, while legacy TensorRT averaged 72.32–72.41. Latency and throughput are end-to-end request measurements and are not normalized to identical output-token counts.
FLAN-T5 Large
This comparison uses
google/flan-t5-largeand a deterministic 1,024-request Super-NaturalInstructions workload derived from the officialallenai/natural-instructionsdefault/testsplit. Inputs longer than 512 tokens and reference outputs longer than 128 tokens are rejected rather than truncated.With encoder and mixed encoder-decoder CUDA graphs enabled, PyTorch improves request throughput by 11.8–57.8% and reduces mean latency by 9.4–36.6% compared with legacy TensorRT.
More T5 workload, latency, output-check, and configuration details
Each result is one clean run of 1,024 requests after an untimed warmup of one concurrency-sized request window. Timing covers closed-loop request submission through receipt of the final output.
Detailed throughput and latency measurements
Output checks
Both APIs used greedy decoding with a maximum of 128 generated tokens. The benchmark normalizes the legacy decoder-start and EOS conventions before counting or hashing outputs.
An eight-request encoder-graph smoke test produced exactly the same greedy token sequences as eager execution. Full-run cross-backend token counts differ by at most 34 tokens (0.36%) because BF16 execution and batching change a small number of near-tie decoding decisions. This is a performance benchmark, not a task-accuracy evaluation.
Workload
Prompts use the following form:
Selection is deterministic with seed 0:
google/flan-t5-largetokenizerThe final workload covers 116 test tasks and 12 task categories. During selection, 7,086 candidate instances were rejected for exceeding 512 input tokens and five were rejected for exceeding 128 reference-output tokens.
Encoder input-length distribution
Reference output-length distribution
Configuration
google/flan-t5-large, TP=1, PP=1PyTorch decoder CUDA graph batch sizes:
PyTorch encoder CUDA graph buckets:
Mixed encoder-decoder CUDA graphs are enabled, with encoder-token buckets derived from the encoder runner's captured keys. At concurrency 64, the encoder batch-eight capture bucket is omitted because the full relative-attention layout set exceeds the 80 GB GPU during capture; encoder admission therefore uses batch-four microbatches.
The legacy TensorRT path reuses one BF16 encoder/decoder engine pair built with maximum batch size 64:
Both legacy engines use BF16 BERT-attention, GPT-attention, and GEMM plugins with input-padding removal enabled. Context FMHA is disabled because the legacy T5 implementation does not support T5 relative attention bias through that path.
Encoder CUDA graph correctness validation
Before the relative-position correction, graph replay reused relative-position bias from capture instead of rebuilding it for replayed sequence lengths. For the same eight greedy requests, that produced 158 tokens, seven EOS stops, and one 128-token length stop, versus the eager reference's 32 tokens, eight EOS stops, and no length stops.
After the fix, the graph-enabled smoke run exactly matches eager execution: 32 generated tokens, eight EOS stops, and no length stops. All three full graph-enabled runs also have the expected 1,021 natural EOS stops and three length stops.
Validation
Dev Engineer Review
encoder_cuda_graph_configandenable_encoder_decoder_mixed_cuda_graph.QA Engineer Review
TorchLlmArgstests for encoder graph configuration validation and mixed-graph settings.l0_h100.yml: adds BART and T5 continuous-admission tests.l0_dgx_h100.yml: adds the two-GPU BART continuous-admission test.l0_l40s.yml: removes the superseded T5 mixed-context test.