[TRTLLM-14778][perf] Add feature-mode encoder CUDA graphs for fixed-shape encoders (Whisper) - #17030
[TRTLLM-14778][perf] Add feature-mode encoder CUDA graphs for fixed-shape encoders (Whisper)#17030pranav-nvidia wants to merge 9 commits into
Conversation
27454f4 to
e0e0af0
Compare
An encoder that consumes fixed-shape per-request features emits the same number of positions for every request, so its graph key is the batch size alone and the token-shaped num_tokens / seq_lens buckets do not apply. encoder_cuda_graph_config becomes a discriminated union on mode so each encoder kind accepts only the buckets it has. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
…ture encoders Whisper's encoder takes a 30 s-padded waveform per request, so the runner swaps its packed-token static tensors for an input_features buffer keyed on batch size. Capture goes through the shared two-pass warmup helper and runs on a dedicated stream, because encoder replay is device-concurrent with decoder replay. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
Covers capture and replay across the configured encoder batch sizes, the eager fallback for an uncaptured size, the config/model mismatch branches, and encoder microbatch admission with the feature config enabled and declined. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
e0e0af0 to
7618a71
Compare
…, not a new config Encoder-graph capture keyed on batch size alone applies to an encoder whose input is a fixed-shape per-request feature tensor, which is a property of the model rather than a choice the caller makes. Detect it from encoder_graph_spec() and drop the separate config type, so encoder_cuda_graph_config keeps its existing shape and the token buckets a feature encoder derives become optional. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
…ranches Delete five unreachable branches, fold the feature and token capture setups into one parameterized capture epilogue, and route feature-mode warmup through the existing enc-dec driver, which captures on the worker owning runtime replay. Feature capture now uses capture_error_mode="thread_local" like the token path, and a feature model with no fitting batch size disables the runner outright. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
…redown Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> # Conflicts: # tests/unittest/_torch/executor/test_py_executor.py
pranav-nvidia
added
the
api-compatible
pranav-nvidia
changed the title
WalkthroughWhisper now provides fixed encoder graph metadata. The PyTorch engine and CUDA graph runner support fixed-shape feature inputs, asynchronous staging, capture, replay, padding, and eager fallback. Configuration, scheduling, documentation, unit tests, and integration tests cover the new path. ChangesWhisper encoder CUDA graphs
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant WhisperModel
participant ModelEngine
participant FeaturePacker
participant EncoderCUDAGraphRunner
participant Encoder
WhisperModel->>ModelEngine: provide encoder_graph_spec()
ModelEngine->>EncoderCUDAGraphRunner: configure fixed feature graph
ModelEngine->>FeaturePacker: pack CPU feature tensors
FeaturePacker->>EncoderCUDAGraphRunner: stage padded feature batch
EncoderCUDAGraphRunner->>Encoder: replay captured encoder graph
Encoder-->>ModelEngine: return encoder outputs
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: 6
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)
5528-5533: 🚀 Performance & Scalability | 🟠 Major | ⚡ Quick winKeep feature admission limited to runner-supported batch sizes.
When
is_feature_encoderis true, Lines 5528-5533 can appendencoder_batch_size_limiteven when it is not inrunner.supported_batch_sizes. For example, runner sizes[1, 2, 4, 16]with a limit of8produce target8. The executor then waits for and submits a size with no captured feature graph. This defeats graph admission and can force eager execution.
tensorrt_llm/_torch/pyexecutor/py_executor.py#L5528-L5533: Do not append a synthetic padded size for feature mode. Select only resolved runner sizes at or below the limit.tests/unittest/_torch/executor/test_py_executor.py#L300-L327: Add[1, 2, 4, 16]with an encoder limit of8, and assert admission selects4.Proposed fix
- if (encoder_cuda_graph_config.enable_padding + if (not is_feature_encoder + and encoder_cuda_graph_config.enable_padding and any(batch_size > encoder_batch_size_limit for batch_size in configured_batch_sizes) and (not supported_batch_sizes or supported_batch_sizes[-1] != encoder_batch_size_limit)):🤖 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 5528 - 5533, In tensorrt_llm/_torch/pyexecutor/py_executor.py lines 5528-5533, update the feature-encoder batch-size admission logic to never append synthetic encoder_batch_size_limit values; select only resolved runner-supported sizes at or below the limit, preserving the largest eligible size. In tests/unittest/_torch/executor/test_py_executor.py lines 300-327, add the runner sizes [1, 2, 4, 16] with encoder limit 8 and assert admission selects 4.
🧹 Nitpick comments (5)
tensorrt_llm/_torch/models/modeling_whisper.py (1)
657-666: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueReuse the cached decoder prompt instead of rebuilding it.
__init__already storesself._decoder_prompt = self._build_decoder_prompt()at Line 625, and_resolve_decoder_promptreturns a copy of that list. Reading the cached list keeps the reported prefix length tied to the exact prompt requests receive, and avoids a secondget_decoder_prompt_idscall.♻️ Proposed refactor
- return len(self._build_decoder_prompt()) + return len(self._decoder_prompt)🤖 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/models/modeling_whisper.py` around lines 657 - 666, Update get_decoder_prefix_len to return the length of the cached self._decoder_prompt initialized in __init__, instead of calling _build_decoder_prompt again. Preserve the existing prefix-length behavior while ensuring it matches the prompt returned by _resolve_decoder_prompt.tensorrt_llm/_torch/pyexecutor/model_engine.py (2)
8001-8004: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winConsider caching the eager encoder attention metadata across forwards.
_make_encoder_attn_metadataconstructs a freshTrtllmAttentionMetadataand runsprepare_encoder_only()on every encoder forward, including graph hits wheremaybe_get_cuda_graphdiscards the object and returns the stored graph-resident metadata. In feature mode the padded layout is fully determined bypadded_batch_size(always[fixed] * padded_batch_size), so one cached object per batch size would suffice.
_set_up_attn_metadataalready applies this pattern for the encoder-only case at Lines 3109-3128.🤖 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 8001 - 8004, Cache encoder eager attention metadata in the encoder-forward path using the existing pattern from _set_up_attn_metadata, keyed by padded_batch_size. Reuse the cached TrtllmAttentionMetadata when the padded layout is unchanged, while preserving _make_encoder_attn_metadata for new batch sizes and passing the reused object to maybe_get_cuda_graph.
3710-3716: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePrefer
functools.cached_propertyover thehasattrsentinel.This class already uses
@functools.cached_propertyfor_mm_encoder_cache_enabledat Line 1081. A cached property removes thehasattrcheck and makes the "computed once,Noneis a valid result" contract explicit.♻️ Proposed refactor
- def _model_encoder_graph_spec(self): - """The model's fixed-shape encoder contract, or None. Queried once.""" - if not hasattr(self, "_cached_model_encoder_graph_spec"): - # torch.compile wraps the model; the spec is on the original. - model = getattr(self.model, "_orig_mod", self.model) - spec_fn = getattr(model, "encoder_graph_spec", None) - self._cached_model_encoder_graph_spec = (spec_fn() if spec_fn - is not None else None) - return self._cached_model_encoder_graph_spec + `@functools.cached_property` + def _model_encoder_graph_spec(self): + """The model's fixed-shape encoder contract, or None. Queried once.""" + # torch.compile wraps the model; the spec is on the original. + model = getattr(self.model, "_orig_mod", self.model) + spec_fn = getattr(model, "encoder_graph_spec", None) + return spec_fn() if spec_fn is not None else NoneThis changes the call sites from
self._model_encoder_graph_spec()toself._model_encoder_graph_specat Lines 574, 902, and 3733, and intests/unittest/_torch/executor/test_pytorch_model_engine.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/model_engine.py` around lines 3710 - 3716, Replace the lazy `_cached_model_encoder_graph_spec` sentinel logic with a `functools.cached_property` named `_model_encoder_graph_spec`, preserving the original model lookup and allowing `None` as a cached result. Update all call sites, including the referenced tests, to access the property without invoking it.tests/unittest/_torch/executor/test_pytorch_model_engine.py (1)
1121-1165: 📐 Maintainability & Code Quality | 🔵 TrivialTest coverage summary (required by path instructions for
tests/**).
Changed test functions in this file:
- Added:
test_feature_encoder_capture_keys_are_all_reachable,test_encoder_graph_spec_selection.- Added helpers:
_feature_encoder_runner,_encoder_spec_engine.- Removed or modified: none.
Test-list placement: these are unit tests under
tests/unittest/_torch/executor/. They run throughpytest tests/unittest/per the coding guidelines and do not belong intests/integration/test_lists/test-db/ortests/integration/test_lists/qa/.Coverage verdict: insufficient for the changed production surface. The two added tests cover capture-key derivation and spec selection, both pure-CPU. The following new behaviour in this PR has no unit coverage:
PyTorchModelEngine.__init__batch-size capping for feature mode (bs_cap,model_engine.pyLines 881-896), including the empty-result eager fallback. This is CPU-only logic reachable through a constructed engine object and is directly testable with the existing_encoder_spec_enginestyle fixture.EncoderCUDAGraphRunner.pad_batchfeature-mode branch (cuda_graph_runner.pyLines 1506-1521), specifically theMAX_FEATURE_PADDING_RATIOeager fallback and thefixed_seq_lenpad entries._feature_encoder_runnerbuilds a capture-disabled runner, butpad_batchreturns early whenself.enabledis false, so a second helper withuse_cuda_graph=Trueand stubbed static tensors would be needed.Add these two cases to this file. Both are deterministic and require no GPU.
_encoder_spec_engineandtest_encoder_graph_spec_selectionthemselves look correct:_Modelsupplies exactly the attributes_encoder_graph_specreads, and the conditional method definition inside the class body is valid.The path instructions state: "If the change includes test-code files (outside tests/integration/test_lists/), the summary must include: 1. Which test functions were added, modified, or removed. 2. Whether each changed test is listed in the appropriate test list files [...] 3. A coverage verdict: sufficient, insufficient, or needs follow-up."
Do you want me to generate the two missing unit tests?
🤖 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 1121 - 1165, Add deterministic CPU unit tests in this file for the uncovered feature-mode batch-size capping in PyTorchModelEngine.__init__, including the empty-result eager fallback, and for EncoderCUDAGraphRunner.pad_batch with CUDA graphs enabled, covering MAX_FEATURE_PADDING_RATIO fallback and fixed_seq_len padding. Reuse the existing helper patterns, stubbing static tensors as needed so no GPU is required, and update the required test coverage summary to list the additions, test-list placement, and coverage verdict.Source: Path instructions
tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py (1)
1178-1217: 🚀 Performance & Scalability | 🔵 TrivialPinned host footprint scales with
max_supported_batch_sizeand is allocated eagerly.Feature mode allocates
FEATURE_MIRROR_SLOTSpinned host buffers of shape(max_supported_batch_size, *feature_shape)plus one device buffer of the same shape, at construction, regardless of the batch sizes actually served. For Whisper withfeature_shape=(480000,), fp32, andmax_supported_batch_size=8this is roughly 30 MB pinned host and 15 MB device. Pinned memory is a process-wide resource shared with the decoder path's staging buffers and the KV-cache transfer paths.Consider recording the peak pinned allocation in the existing
log_mem_snapshotwarmup checkpoints so the footprint is visible whenmax_supported_batch_sizegrows.🤖 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 1178 - 1217, Update the feature-mode initialization around self._feature_mirrors and self.shared_static_tensors to record the peak pinned-memory allocation during the existing warmup checkpoints via log_mem_snapshot. Ensure snapshots capture the footprint after buffers are allocated and when max_supported_batch_size grows, without changing the current buffer allocation or mirror 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/models/modeling_whisper.py`:
- Line 894: Annotate the encoder graph spec contract across all affected
methods: update modeling_whisper.py lines 894-894 so encoder_graph_spec returns
tuple[tuple[int, ...], torch.dtype, int], and update model_engine.py lines
3708-3743 so _model_encoder_graph_spec returns tuple[tuple[int, ...],
torch.dtype, int] | None and _encoder_graph_spec returns tuple[tuple[int, ...] |
None, torch.dtype | None, int | None].
In `@tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py`:
- Around line 1506-1521: Update the comment in the feature-mode padding branch
near MAX_FEATURE_PADDING_RATIO to document that the 9/8 bound rejects every
non-exact batch size with consecutive power-of-two buckets, including the
default [1, 2] configuration, making enable_padding effectively inert there;
note that padding is only reachable with denser bucket lists, while preserving
the existing logic.
In `@tensorrt_llm/_torch/pyexecutor/model_engine.py`:
- Around line 7660-7679: Before replacing the staging buffer in the reallocation
branch of the encoder feature staging logic, synchronize the existing
_encoder_feature_staging_event when present so any prior H2D completes before
the old pinned tensor is released. Preserve a single
_encoder_feature_copy_stream across reallocations by creating it only if it does
not already exist, rather than replacing the stream each time.
- Around line 881-896: Update the bs_cap calculation in the encoder CUDA-graph
setup to use the raw token-budget division without a max(1, ...) floor, and cap
it with self.encoder_batch_size instead of self.batch_size. Preserve the
existing filtering and warning fallback so no fitting batch sizes disables
encoder CUDA graphs and keeps the encoder step eager.
In `@tests/integration/defs/llmapi/test_llm_api_pytorch_whisper.py`:
- Around line 291-296: Update the comment above the assertions in the encoder
CUDA graph test to state only that graph capture occurred, not that serving
requests replayed a captured graph. Keep the existing assertions unchanged, and
do not add replay instrumentation or assertions; runtime replay should be
tracked separately through EncoderCUDAGraphRunner._replay_features.
In `@tests/unittest/llmapi/test_llm_args.py`:
- Around line 1886-1916: Extend test_encoder_cuda_graph_config_validation with a
case that sets encode_only=True alongside encoder_cuda_graph_config and asserts
TorchLlmArgs raises ValidationError. Use a valid encoder CUDA graph
configuration so the assertion specifically covers the encode_only rejection
branch in validate_encoder_cuda_graph_config.
---
Outside diff comments:
In `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Around line 5528-5533: In tensorrt_llm/_torch/pyexecutor/py_executor.py lines
5528-5533, update the feature-encoder batch-size admission logic to never append
synthetic encoder_batch_size_limit values; select only resolved runner-supported
sizes at or below the limit, preserving the largest eligible size. In
tests/unittest/_torch/executor/test_py_executor.py lines 300-327, add the runner
sizes [1, 2, 4, 16] with encoder limit 8 and assert admission selects 4.
---
Nitpick comments:
In `@tensorrt_llm/_torch/models/modeling_whisper.py`:
- Around line 657-666: Update get_decoder_prefix_len to return the length of the
cached self._decoder_prompt initialized in __init__, instead of calling
_build_decoder_prompt again. Preserve the existing prefix-length behavior while
ensuring it matches the prompt returned by _resolve_decoder_prompt.
In `@tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py`:
- Around line 1178-1217: Update the feature-mode initialization around
self._feature_mirrors and self.shared_static_tensors to record the peak
pinned-memory allocation during the existing warmup checkpoints via
log_mem_snapshot. Ensure snapshots capture the footprint after buffers are
allocated and when max_supported_batch_size grows, without changing the current
buffer allocation or mirror behavior.
In `@tensorrt_llm/_torch/pyexecutor/model_engine.py`:
- Around line 8001-8004: Cache encoder eager attention metadata in the
encoder-forward path using the existing pattern from _set_up_attn_metadata,
keyed by padded_batch_size. Reuse the cached TrtllmAttentionMetadata when the
padded layout is unchanged, while preserving _make_encoder_attn_metadata for new
batch sizes and passing the reused object to maybe_get_cuda_graph.
- Around line 3710-3716: Replace the lazy `_cached_model_encoder_graph_spec`
sentinel logic with a `functools.cached_property` named
`_model_encoder_graph_spec`, preserving the original model lookup and allowing
`None` as a cached result. Update all call sites, including the referenced
tests, to access the property without invoking it.
In `@tests/unittest/_torch/executor/test_pytorch_model_engine.py`:
- Around line 1121-1165: Add deterministic CPU unit tests in this file for the
uncovered feature-mode batch-size capping in PyTorchModelEngine.__init__,
including the empty-result eager fallback, and for
EncoderCUDAGraphRunner.pad_batch with CUDA graphs enabled, covering
MAX_FEATURE_PADDING_RATIO fallback and fixed_seq_len padding. Reuse the existing
helper patterns, stubbing static tensors as needed so no GPU is required, and
update the required test coverage summary to list the additions, test-list
placement, and coverage verdict.
🪄 Autofix
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: aa61dbd5-e474-4187-aadf-cc69c2da4ce8
📒 Files selected for processing (11)
tensorrt_llm/_torch/models/modeling_whisper.pytensorrt_llm/_torch/pyexecutor/cuda_graph_runner.pytensorrt_llm/_torch/pyexecutor/model_engine.pytensorrt_llm/_torch/pyexecutor/py_executor.pytensorrt_llm/llmapi/llm_args.pytests/integration/defs/llmapi/test_llm_api_pytorch_whisper.pytests/integration/test_lists/test-db/l0_l40s.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/llmapi/test_llm_args.py
| # Capture must actually have happened: a silent fallback to the eager | ||
| # encoder path would otherwise pass every output assertion above. | ||
| assert encoder_runner.enabled | ||
| assert encoder_runner.graphs | ||
| assert encoder_runner.feature_mode | ||
| assert encoder_runner.is_encoder_decoder |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
The comment overstates what these assertions prove.
assert encoder_runner.graphs proves that warmup captured at least one graph. It does not prove that any serving request replayed one. EncoderCUDAGraphRunner.pad_batch falls back to eager whenever the padded batch exceeds MAX_FEATURE_PADDING_RATIO, and _maybe_forward_encoder_graph returns None on any feature-shape, dtype, or encoder_output_len mismatch. Both fallbacks leave encoder_runner.graphs populated, so the "silent fallback to the eager encoder path" the comment rules out would still pass every assertion here.
The PR objectives already list "asserting runtime replay" as a follow-up. Until that lands, narrow the comment so it claims capture only.
A replay counter incremented in EncoderCUDAGraphRunner._replay_features would make this assertable, matching the guidance that confirming an executed path needs production-side instrumentation rather than existing public fields. Do you want me to open an issue to track adding it?
✏️ Proposed comment change
- # Capture must actually have happened: a silent fallback to the eager
- # encoder path would otherwise pass every output assertion above.
+ # Capture must actually have happened during warmup. Note this does NOT
+ # prove any request replayed a graph: `pad_batch` and
+ # `_maybe_forward_encoder_graph` can route every request to the eager
+ # encoder path while `graphs` stays populated. Asserting runtime replay
+ # needs a replay counter on the runner (tracked as a follow-up).📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| # Capture must actually have happened: a silent fallback to the eager | |
| # encoder path would otherwise pass every output assertion above. | |
| assert encoder_runner.enabled | |
| assert encoder_runner.graphs | |
| assert encoder_runner.feature_mode | |
| assert encoder_runner.is_encoder_decoder | |
| # Capture must actually have happened during warmup. Note this does NOT | |
| # prove any request replayed a graph: `pad_batch` and | |
| # `_maybe_forward_encoder_graph` can route every request to the eager | |
| # encoder path while `graphs` stays populated. Asserting runtime replay | |
| # needs a replay counter on the runner (tracked as a follow-up). | |
| assert encoder_runner.enabled | |
| assert encoder_runner.graphs | |
| assert encoder_runner.feature_mode | |
| assert encoder_runner.is_encoder_decoder |
🤖 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_whisper.py` around lines
291 - 296, Update the comment above the assertions in the encoder CUDA graph
test to state only that graph capture occurred, not that serving requests
replayed a captured graph. Keep the existing assertions unchanged, and do not
add replay instrumentation or assertions; runtime replay should be tracked
separately through EncoderCUDAGraphRunner._replay_features.
`_filter_cuda_graph_batch_sizes` now takes the per-request token cost, so a feature encoder's `fixed_seq_len` budget is applied by the same filter the token path uses instead of a second cap layered on afterwards. A cap below one request returns no batch sizes, making the "nothing fits, stay eager" fallback reachable. Also annotate the spec helpers and retire the in-flight H2D before the feature staging buffer is reallocated. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
A populated `graphs` only proves warmup captured; `pad_batch` and the shape checks in `_maybe_forward_encoder_graph` can route every request to the eager encoder without emptying it. Count feature replays on the runner and assert it. Also cover the `encode_only` rejection and document that the 12.5% padding bound is unreachable for power-of-two bucket lists. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
Two things worth addressing before merge, both small; the rest is body-level context.
The replay assertion doesn't prove what it claims. During engine warmup, _warmup_and_capture_encoder_cuda_graphs runs the capture pass with is_warmup_only=False, and _feature_encoder_graph_forward calls runner.replay() right after runner.capture() — so num_feature_replays == len(graphs) before the first real request. The integration test's > 0 check therefore passes even if every runtime batch silently falls back to eager, which is exactly the failure mode the counter was introduced to catch. Details inline.
The dedicated capture stream and unshared memory pool apply to the token-path (T5/BART) encoder runner too, not just feature graphs. cuda_graph_mem_pool=None at model_engine.py:948 and the stream= argument in the shared capture() change behavior for existing enc-dec deployments: encoder graphs no longer share the decoder's pool (extra device memory) and capture on a per-runner stream. The in-code rationale (concurrent encoder/decoder replay violating torch's pool-sharing contract; cuBLAS split-K semaphore coupling) reads as a correctness fix for the existing path, but the PR description frames both as feature-mode decisions. Please state explicitly that this changes T5/BART encoder graph capture/allocation, and whether it was validated against the existing l0 enc-dec cases.
Behavior change for previously-invalid configs: a token-encoder model (T5) configured with batch_sizes only used to fail Pydantic validation; it now gets a log warning at engine init and runs eager (model_engine.py:594). That converts a loud config error into silent perf degradation for non-feature models. Consider raising in the engine when a token-encoder model has a config with missing buckets, keeping the warning path only for feature-capable models.
Docs: docs/source/models/encoder-decoder.md (CUDA-graphs section, line ~419) still shows num_tokens/seq_lens as required, and the EncodeCudaGraphConfig field descriptions don't mention that a fixed-shape feature encoder derives them from the model. The Whisper section should show the batch-sizes-only enablement snippet from the PR description.
Verified the l0_l40s swap: kv-v2 decoder-graph coverage does remain on l0_h100 (bf16-kv-v2-decoder-graphs-on-greedy), so the substitution is sound.
| # `_maybe_forward_encoder_graph` can route every request to the eager | ||
| # encoder while `graphs` stays populated, and that silent fallback would | ||
| # pass every output assertion above. Only the replay counter rules it out. | ||
| assert encoder_runner.num_feature_replays > 0 |
There was a problem hiding this comment.
This assertion is satisfied by warmup alone. The capture pass in _warmup_and_capture_encoder_cuda_graphs runs with is_warmup_only=False, and _feature_encoder_graph_forward calls runner.replay() immediately after runner.capture() — so num_feature_replays == len(runner.graphs) before the first real request is served. If runtime batches all fall back to eager (the silent-fallback case the comment above describes), this still passes.
Assert against the warmup baseline instead:
assert encoder_runner.num_feature_replays > len(encoder_runner.graphs)Each captured key gets exactly one capture-pass replay, so anything beyond len(graphs) must be runtime traffic.
| # [1, 2] the Whisper integration test configures, `enable_padding` | ||
| # is therefore inert and only exact batch sizes replay - which is | ||
| # what `_waiting_encoder_requests` forms microbatches to hit. | ||
| # Padding only becomes reachable with a dense list such as |
There was a problem hiding this comment.
The [1, 2, 3, 4] example is wrong: with that list every admissible batch size (1–4) is an exact bucket, so padded_batch_size == batch_size and pad_batch yields unchanged — padding never fires there either. Under the 9/8 bound, padding is only reachable when a non-bucket batch size has the next bucket within 12.5%, which first happens at batch 8 → bucket 9 (9 ≤ 8 × 1.125). Since this comment is how readers (and the admission-path unit tests) reason about when padding can occur, fix the example, e.g. "reachable only when a non-bucket batch has the next bucket within 12.5%, e.g. a batch of 8 padding to a configured bucket of 9".
| for request in encoder_requests: | ||
| f = request.py_encoder_input_features | ||
| if (f is None or int(request.encoder_output_len) != fixed | ||
| or f.shape[1:] != runner.config.feature_shape |
There was a problem hiding this comment.
f.shape[1:] != runner.config.feature_shape accepts any leading dimension, but _replay_features writes f.shape[0] rows per request into a mirror sized (max_batch_size, *feature_shape) while seq_lens/padded_batch_size count one row per request. A feature tensor of shape (2, 480000) would pass this gate, double-count rows, and overflow the mirror slice (or scatter the wrong rows) at bucket-sized batches. Whisper always produces (1, n_samples) today, so make the gate enforce the actual contract:
or f.shape != (1, *runner.config.feature_shape)| if missing: | ||
| raise ValueError("encoder_cuda_graph_config requires " | ||
| f"{' and '.join(missing)}.") | ||
| # `num_tokens` / `seq_lens` are checked by the model engine rather than |
There was a problem hiding this comment.
Moving this check out of the validator changes the failure mode for token-encoder models: a T5/BART config with only batch_sizes used to raise ValidationError at args construction; it now produces a log warning at engine init (model_engine.py:594) and silently runs eager. For a feature encoder that's the new valid path, but for token models it converts a config error into an easy-to-miss perf regression. Since the engine knows the encoder kind at that point, consider raising there when _model_encoder_graph_spec() is None and the buckets are missing, preserving the old hard failure for the models that still require them.
Assert encoder replay against the warmup baseline: the capture pass replays each key once, so `> 0` was satisfied before any request was served. Require a token encoder's bucket lists at engine init instead of warning and running eager. Reject multi-row feature tensors, which would overrun the replay mirror. Stop admitting feature batches at a size the runner never captured. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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/encoder-decoder.md`:
- Around line 440-448: Update the model examples in the encoder bucket
documentation to describe Whisper as using a fixed-shape `input_features` tensor
after audio preprocessing, replacing the inaccurate reference to its audio
waveform; leave the surrounding capture behavior unchanged.
🪄 Autofix
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: 11acd091-2858-4296-baf4-7ce086ba0684
📒 Files selected for processing (8)
docs/source/models/encoder-decoder.mdtensorrt_llm/_torch/pyexecutor/cuda_graph_runner.pytensorrt_llm/_torch/pyexecutor/model_engine.pytensorrt_llm/_torch/pyexecutor/py_executor.pytensorrt_llm/llmapi/llm_args.pytests/integration/defs/llmapi/test_llm_api_pytorch_whisper.pytests/unittest/_torch/executor/test_py_executor.pytests/unittest/_torch/executor/test_pytorch_model_engine.py
🚧 Files skipped from review as they are similar to previous changes (7)
- tensorrt_llm/_torch/pyexecutor/py_executor.py
- tests/unittest/_torch/executor/test_py_executor.py
- tensorrt_llm/llmapi/llm_args.py
- tests/integration/defs/llmapi/test_llm_api_pytorch_whisper.py
- tests/unittest/_torch/executor/test_pytorch_model_engine.py
- tensorrt_llm/_torch/pyexecutor/model_engine.py
- tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
| Which encoder buckets you must supply depends on the model. A text encoder, | ||
| such as BART or T5, packs a variable number of tokens per request, so | ||
| `num_tokens` and `seq_lens` are part of its key space and are required; leaving | ||
| either unset is an error at engine construction. An encoder whose input is a | ||
| fixed-shape per-request feature tensor, such as Whisper's audio waveform, | ||
| produces the same number of encoder positions for every request, so both lists | ||
| follow from the model and are derived rather than configured. For those models | ||
| `batch_sizes` alone enables capture, and any `num_tokens` or `seq_lens` you set | ||
| is ignored: |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Describe Whisper's feature tensor instead of its waveform.
Whisper's encoder receives input_features, not the raw audio waveform. Replace “Whisper's audio waveform” with “Whisper's fixed-shape input_features tensor after audio preprocessing.”
Proposed wording
- such as Whisper's audio waveform,
+ such as Whisper's fixed-shape `input_features` tensor after audio preprocessing,📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| Which encoder buckets you must supply depends on the model. A text encoder, | |
| such as BART or T5, packs a variable number of tokens per request, so | |
| `num_tokens` and `seq_lens` are part of its key space and are required; leaving | |
| either unset is an error at engine construction. An encoder whose input is a | |
| fixed-shape per-request feature tensor, such as Whisper's audio waveform, | |
| produces the same number of encoder positions for every request, so both lists | |
| follow from the model and are derived rather than configured. For those models | |
| `batch_sizes` alone enables capture, and any `num_tokens` or `seq_lens` you set | |
| is ignored: | |
| Which encoder buckets you must supply depends on the model. A text encoder, | |
| such as BART or T5, packs a variable number of tokens per request, so | |
| `num_tokens` and `seq_lens` are part of its key space and are required; leaving | |
| either unset is an error at engine construction. An encoder whose input is a | |
| fixed-shape per-request feature tensor, such as Whisper's fixed-shape `input_features` tensor after audio preprocessing, | |
| produces the same number of encoder positions for every request, so both lists | |
| follow from the model and are derived rather than configured. For those models | |
| `batch_sizes` alone enables capture, and any `num_tokens` or `seq_lens` you set | |
| is ignored: |
🤖 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 `@docs/source/models/encoder-decoder.md` around lines 440 - 448, Update the
model examples in the encoder bucket documentation to describe Whisper as using
a fixed-shape `input_features` tensor after audio preprocessing, replacing the
inaccurate reference to its audio waveform; leave the surrounding capture
behavior unchanged.
| f"{' and '.join(missing)}.") | ||
| # `num_tokens` / `seq_lens` are checked by the model engine rather than | ||
| # here: an encoder whose input is a fixed-shape per-request feature | ||
| # tensor derives both from the model, and only the engine knows which |
There was a problem hiding this comment.
This changes which TorchLlmArgs values are accepted, so please run python3 scripts/generate_llm_args_golden_manifest.py and commit the updated tensorrt_llm/usage/llm_args_golden_manifest.json in this PR; the repository requires schema and telemetry-manifest changes to land together.
fredricz-20070104
left a comment
fredricz-20070104
left a comment
There was a problem hiding this comment.
Review summary - CONCERNS
Verdict: No hard correctness blocker is visible from the diff, and the feature path is reasonably tested, but two things should be resolved before merge: this quietly changes encoder-graph allocation/capture for existing T5/BART models, and the feature path was only validated locally on SM120.
Concerns
-
[MAJOR]
tensorrt_llm/_torch/pyexecutor/model_engine.py:927- pool/stream change hits existing T5/BART, not just feature mode- What is wrong:
cuda_graph_mem_pool=None(wasself._cuda_graph_mem_pool) plus the newstream=self._get_capture_stream()in the sharedcapture()apply to every encoder-decoder model. Encoder graphs no longer share the decoder's memory pool and now capture on a per-runner stream. - How it fails: a T5/BART deployment tuned around the shared pool now allocates a second independent graph pool at engine init; on a memory-tight box that can OOM at warmup/capture where it previously fit. The l0 enc-dec cases prove functional correctness but do not bound the extra device memory.
- Suggested fix: state explicitly in the PR that this changes T5/BART encoder graph allocation/capture, confirm the l0 enc-dec cases were re-run with the memory delta measured, and — if the unshared pool is only needed for concurrent feature-mode replay — gate the pool split on
feature_modeso the token path keeps sharing the decoder pool.
- What is wrong:
-
[MAJOR]
tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py:2018- feature replay validated only on SM120- What is wrong: the feature replay machinery (single device
input_featuresbuffer, double-buffered pinned mirrors, per-mirror events, eager stream-ordered H2D issued right beforegraphs[key].replay()) is new and concurrency-sensitive. - How it fails: correctness depends on the H2D and the replay staying on the same stream; behavior can differ under a different driver/allocator. The PR states local validation was SM120 only and lists L40S/H100 CI as a follow-up.
- Suggested fix: require the L40S (and ideally H100) CI run to be green before merge rather than treating the SM120 local run as a stand-in.
- What is wrong: the feature replay machinery (single device
Minor notes (non-blocking)
docs/source/models/encoder-decoder.md:442- Whisper's fixed-shape input is the preprocessedinput_featurestensor, not its "audio waveform".tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py:1917-_prepare_feature_capturedocstring says buckets "share one pinned mirror" butFEATURE_MIRROR_SLOTS = 2(double-buffered).
QA view
- Test coverage: partial - the integration case
bf16-kv-v1-encoder-graphs-on-greedychecks exact greedy token ids ANDnum_feature_replays > len(graphs), which correctly rules out silent eager fallback and validates output. Uncovered: the T5/BART pool/stream change has no dedicated test (only pre-existing functional l0 cases), and the eager_pack_encoder_featuresstaging/event path has no unit test. - SM coverage: architecture-independent (no arch guards in the diff); tests run on L40S (SM89) via the l0 list, author validated locally on SM120 only — L40S/H100 CI not yet confirmed.
- Test code: no significant issues; the
> len(graphs)baseline correctly excludes the warmup replay. Unit stubs setfeature_modeby hand onSimpleNamespace, mildly brittle but acceptable. - Test time: small - the l0 list swaps one decoder-graphs case for one encoder-graphs case (net zero new entries); the swapped case is slightly heavier because it also captures encoder graphs.
- Needs
/qa-verify: yes - shared encoder-graph infrastructure changes for existing T5/BART, and the feature path was only validated on SM120; QA should confirm the l0 T5/BART and Whisper encoder-graph cases on L40S (ideally H100) and watch for encoder-graph memory regressions.
Possible new issues
- Extra unshared graph pool for all enc-dec models → possible OOM on previously-fitting T5/BART configs.
- New
_pack_encoder_featuresstaging path (dedicated copy stream + event over a reused pinned buffer) replacestorch.cat(...).to('cuda')for all eager Whisper runs; a mis-ordered event could let a batch overwrite staging still in flight. - Single device feature buffer refilled by an eager H2D immediately before replay — correct only while both stay on the same stream.
What I could not verify
- Actual device-memory delta of the unshared encoder pool for T5/BART.
- Runtime replay correctness and mirror/event ordering on L40S/H100 (only SM120 local run reported).
- Whether
torch.cuda.graphstream handoff from the warmup stream is fully ordered in this codebase's torch version (assumed viawait_streamon capture entry).
Automated review by NVCortex Lite, run by @fredricz-20070104.
fredricz-20070104
left a comment
fredricz-20070104
left a comment
There was a problem hiding this comment.
Review summary - Approve (non-blocking)
Approving so this is not blocked on me. The points raised in my review comment above are non-blocking — please read them and address what you agree with before merging.
Worth doing before this is relied on: This changes shared encoder-graph infrastructure (pool/stream) for existing T5/BART paths, and the feature-mode replay path was only validated locally on SM120 with L40S/H100 CI deferred by the author. A human QA should confirm the l0 T5/BART and Whisper encoder-graph cases pass on L40S (and ideally H100), and watch for encoder-graph memory-footprint regressions.
Automated review by NVCortex Lite, run by @fredricz-20070104.
| missing.append("seq_lens/max_seq_len") | ||
| if not missing: | ||
| return | ||
| raise ValueError( |
There was a problem hiding this comment.
Upgrading this warning to an error breaks encode-only setups. Previously the encoder-only model would run eagerly with a warning.
| if self.encoder_max_batch_size is None: | ||
| raise ValueError( | ||
| "encoder_cuda_graph_config requires encoder_max_batch_size.") | ||
| missing = [] |
There was a problem hiding this comment.
Validation now fires after model load, in workers — and never on AutoDeploy.
The deleted pydantic fail-fast used to reject a bad config instantly, client-side. Now a T5/BART user with forgotten buckets spawns workers, loads weights for minutes, then gets a worker-side ValueError; on backend='_autodeploy' (whose LlmArgs inherits TorchLlmArgs but never constructs this engine) the invalid config is accepted silently.
| """ | ||
| fixed_seq_len = int(self.config.max_source_positions) | ||
| hop_length = int(self.model.encoder.log_mel.hop_length) | ||
| n_samples = fixed_seq_len * 2 * hop_length |
There was a problem hiding this comment.
encoder_graph_spec window derivation can permanently disagree with the processor.
The spec computes max_source_positions * 2 * hop_length; the processor pads to the extractor's n_samples and validates with truncating division. A checkpoint with n_samples=480100 passes validation, captures graphs for (1, 480000), and then every runtime batch fails the shape guard forever — capture time and graph memory spent, zero replays, no warning.
can the spec be derived from the extractor's n_samples (same source the processor uses), or assert spec-vs-processor agreement at startup and decline feature mode loudly on mismatch.
| -(i + 1) | ||
| for i in range(len(padded_seq_lens) - len(request_ids)) | ||
| ] | ||
| eager_attn_metadata = self._make_encoder_attn_metadata( |
There was a problem hiding this comment.
The hot path rebuilds throwaway attention metadata on every graph hit.
Suggested fix: check key in runner.graph_metadata first and only build eager metadata on a miss.
| attn_backend="TRTLLM", | ||
| max_batch_size=8, | ||
| encoder_max_batch_size=8, | ||
| encoder_cuda_graph_config=EncodeCudaGraphConfig(batch_sizes=[1, 2, 4, 8]), |
There was a problem hiding this comment.
The example configures batch_sizes=[1,2,4,8] on top of settings whose token budget (max_num_tokens=3000 // 1500) silently drops buckets 4 and 8 — with default enable_padding=False there's not even a warning, and admission then targets batch 2, halving the concurrency the example's encoder_max_batch_size=8 implies. The drop rule is stated two lines below the example. One-line fix (encoder_max_num_tokens: 12000 or batch_sizes=[1,2]), but it's the feature's onboarding path, so it ships confusion to every doc-follower.
6 participants
Dev Engineer Review
encoder_graph_spec().encoder_cuda_graph_configvalues.QA Engineer Review
Test changes
Added or updated coverage for:
The Whisper integration coverage is mapped in
tests/integration/test_lists/test-db/l0_l40s.yml. The feature-encoder unit tests are not explicitly mapped to test-list entries.Verdict: needs follow-up.
Description
Encoder-decoder encoder CUDA graphs already exist for packed-token encoders such as T5 and BART (#16706). This PR extends that machinery to encoders whose input is a fixed-shape per-request feature tensor, starting with Whisper.
The difference is the graph key. A token encoder's key depends on the packed token count and sequence lengths, so those buckets have to be configured. A feature encoder emits a fixed number of encoder positions per request whatever the input, so its key degenerates to the batch size and both bucket lists are derived from the model.
Models opt in by declaring
encoder_graph_spec()returning(feature_shape, dtype, fixed_seq_len)— the model selects the mode, not the config.encoder_cuda_graph_configtherefore acceptsbatch_sizeson its own;num_tokens/seq_lensare no longer required at config validation, a relaxation that leaves every previously-valid config valid. TP > 1, models that do not declare a spec, and draft models stay on the eager path.Four decisions worth a reviewer's attention:
encoder_stream, device-concurrent with decoder replay, and torch's pool-sharing contract assumes replays from a shared pool are not concurrent.Mixed encoder/decoder capture now takes the decoder prefix length from the input processor (Whisper forces 4 tokens) rather than the BART/T5 heuristic; a mismatch makes every mixed batch miss its graph silently. Feature capture uses
capture_error_mode="thread_local", matching the token path.Enable with:
Test Coverage
test_whisper_pytorch_feature_combinations[bf16-kv-v1-encoder-graphs-on-greedy]transcribes at batch 1 and 2 with exact pinned greedy token ids, and additionally asserts the encoder graphs actually captured — without that, a silent fallback to the eager encoder would pass every output check. It replaces the existing L40S pre-merge decoder-only Whisper case rather than adding an invocation, since encoder graphs exercise decoder graphs too; KV-v2 decoder coverage stays on H100.Local validation on SM120, pre-merge tree: Whisper 9 passed / 1 skipped (TP2 needs 2 devices), T5 15 passed, BART 11 passed,
tests/unittest/_torch/executor/1366 passed,test_llm_args.py283 passed. Pre-commit passes, including after the merge with main.Follow-ups:
PR Checklist
GitHub Bot Help
To see a list of available CI bot commands, comment
/bot help.