[TRTLLM-14778][perf] Add feature-mode encoder CUDA graphs for fixed-shape encoders (Whisper)

[pranav-nvidia]

Dev Engineer Review

• Adds fixed-shape feature-encoder CUDA graph support for Whisper.
• Derives graph shapes and batch-size keys from encoder_graph_spec().
• Uses dedicated streams, memory pools, and pinned staging buffers.
• Performs feature H2D copies outside captured graphs.
• Clones graph outputs before request scattering.
• Applies a 12.5% padding overhead limit.
• Keeps unsupported cases on the eager path.
• Allows batch-size-only encoder_cuda_graph_config values for fixed-shape feature encoders.
• Uses the input processor for decoder prefix length.
• Updates encoder-decoder documentation and Whisper CI coverage.
• Configuration and test-list changes match the stated scope.
• Follow-up remains for L40S/H100 CI coverage.

QA Engineer Review

Test changes

Added or updated coverage for:

• Whisper encoder graph configuration and replay state in test_llm_api_pytorch_whisper.py.
• Feature-mode batch admission and disabled-runner behavior in test_py_executor.py.
• Encoder graph specification, bucket validation, padding limits, fallback behavior, and metadata lookup in test_pytorch_model_engine.py.
• Encoder-decoder warmup configuration in test_pytorch_model_engine_warmup.py.
• Batch-size-only encoder graph validation in test_llm_args.py.

Whisper integration coverage is mapped in tests/integration/test_lists/test-db/l0_l40s.yml. The unit-test changes are not explicitly mapped in a test-list file.

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_config therefore accepts batch_sizes on its own. num_tokens/seq_lens are no longer required by config validation, which cannot tell the two kinds of encoder apart; they are checked at engine init instead, where the model is loaded and its encoder kind is known, and a token encoder that omits them still raises. TP > 1, models that do not declare a spec, and draft models stay on the eager path.

Four decisions worth a reviewer's attention:

• The input H2D is deliberately not captured into the graph. Mirrors are pooled across buckets and rotated (FEATURE_MIRROR_SLOTS, currently two) rather than owned per bucket, and consecutive encoder batches can be enqueued back to back, so a captured copy would read a mirror after the host had already rotated onto it. Replay issues an eager stream-ordered H2D guarded by per-mirror events instead.
• Feature graphs capture on their own stream, and the encoder runner is never handed the decoder's pool. Encoder replay runs on encoder_stream, device-concurrent with decoder replay, and torch's pool-sharing contract assumes replays from a shared pool are not concurrent. The capture stream half changes the token path (T5/BART) too: without an explicit stream=, torch.cuda.graph captures on a process-wide singleton stream shared with the decoder graphs, which couples the two graph sets through stream-keyed cuBLAS scratch. The pool half costs nothing — _cuda_graph_mem_pool is None for the engine's life, so both runners already allocated their own pool at first capture, and t5-small measures identical peaks either way (628.0 MiB reserved, 596.0 MiB allocated). The literal None makes non-sharing a requirement rather than a coincidence.
• Graph output is cloned before request scatter. The executor holds views of the result across scheduler iterations, and a later replay of the same bucket would clobber them.
• Padding falls back to eager past a 12.5% overhead bound. Each Whisper pad slot is a full 1500-position encoder forward, unlike the 1-token pads of the token path, so unbounded power-of-two padding would regress large ragged batches.

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:

LLM(
    model=...,
    encoder_max_batch_size=8,
    # 8 buckets x 1500 encoder positions; left at the default, the token
    # budget silently drops the larger buckets.
    encoder_max_num_tokens=12000,
    encoder_cuda_graph_config=EncodeCudaGraphConfig(
        batch_sizes=[1, 2, 4, 8], enable_padding=True
    ),
)

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 num_feature_replays > len(encoder_runner.graphs). That baseline excludes the one replay each key gets during the capture pass, so it proves runtime replay rather than mere capture — without it, a silent fallback to the eager encoder would pass every output check. The case 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.

Unit coverage for the feature path is in test_pytorch_model_engine.py (graph spec selection, bucket validation, batch-size capping against the encoder token budget, the 12.5% padding bound, captured-metadata hit and miss) and test_py_executor.py (feature-mode microbatch admission, including that it never targets an uncaptured batch size).

Local validation on SM120 (RTX PRO 6000 Blackwell), re-run on the tree as merged with current main:

Suite	Result
Whisper integration	9 passed, 1 skipped
T5 + BART integration	27 passed, 18 skipped
tests/unittest/_torch/executor/ -k encoder	95 passed, 15 subtests
tests/unittest/_torch/executor/test_py_executor.py	146 passed
test_pytorch_model_engine.py + test_py_executor.py + test_pytorch_model_engine_warmup.py	223 passed, 41 subtests

Skips are TP2 cases needing two devices in a single-GPU container. The T5 + BART run is what covers the capture-stream change on the token path.

An earlier revision of this branch was additionally built and run on a B200 (SM100): Whisper 9 passed / 1 skipped and the encoder unit tests green, with bf16-kv-v1-encoder-graphs-on-greedy passing there too, so feature-mode capture and replay are not SM120-only. The B200 T5/BART sweep on that box fails on an NVRTC JIT include path (could not open source file "cuda.h") which reproduces identically on pristine main with CUDA graphs off, so it is a container configuration issue rather than anything in this PR.

Standing against the legacy TensorRT backend

Measured 2026-08-26 on this branch as merged with main, plus #17531 — these are the two PRs benchmarked together, so the figures are the current standing of the Whisper PyTorch path, not an attribution of this PR alone. fp16, greedy, overlap scheduler on, encoder CUDA graphs on, product worker topology; block shape pinned to the legacy harness. Legacy denominators are prior measurements on the same harness and hosts.

platform	model	bs	legacy RTFx	this branch + #17531	% of legacy
SM120	whisper-tiny	1 / 8 / 32	673 / 2338 / 4344	283 / 1044 / 1515	42 % / 45 % / 35 %
SM120	whisper-large-v3	1 / 8 / 32	79.0 / 232 / 344	70.1 / 261 / 355	89 % / 112 % / 103 %
B200	whisper-tiny	1 / 8 / 32	380 / 1283 / 2888	228 / 758 / 1188	60 % / 59 % / 41 %
B200	whisper-large-v3	1 / 8 / 32	73.6 / 251 / 530	72.9 / 315 / 577	99 % / 125 % / 109 %

large-v3 exceeds legacy at batch 8 and 32 on both platforms. whisper-tiny remains behind and is worst at batch 32; that gap is host-side Python in the executor loop, not encoder work, and is tracked separately. WER was identical to the reference in all 39 runs (tiny 0.1250, large-v3 0.0423), so nothing here is a speed-for-accuracy trade.

Caveats worth carrying: the B200 legacy column is a single run with no recorded range, and the B200 arm was measured on a shared node under foreign load (host loadavg median 15 vs SM120's 3.9), so the B200 percentages are softer than the SM120 ones — whisper-tiny at batch 32 especially, being the most host-bound cell in the matrix. SM120 whisper-tiny at batch 32 is a known bimodal cell; the figure quoted is the median over 10 rounds.

Follow-up:

• CI on L40S/H100 — local hardware is SM120 and B200, neither of which stands in for the L40S the l0 list actually uses.

PR Checklist

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

GitHub Bot Help

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

[coderabbitai]

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

[coderabbitai]

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

[coderabbitai]

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

[pranav-nvidia]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #69544 [ run ] triggered by Bot. Commit: fef5a1a Link to invocation

[tensorrt-cicd]

PR_Github #69544 [ run ] completed with state SUCCESS. Commit: fef5a1a
/LLM/main/L0_MergeRequest_PR pipeline #56868 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

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

CI Agent Failure Analysis

Link to invocation

[pranav-nvidia]

/bot run

[tensorrt-cicd]

PR_Github #69786 [ run ] triggered by Bot. Commit: fef5a1a Link to invocation

[tensorrt-cicd]

PR_Github #69786 [ run ] completed with state SUCCESS. Commit: fef5a1a
/LLM/main/L0_MergeRequest_PR pipeline #57084 completed with status: 'UNSTABLE'

CI Report

⚠️ Multi-GPU Label Required:
Multi-GPU tests require the ci: full pre-merge approved label on this PR. Ask a member of NVIDIA/trt-llm-ci-approvers to add the label, then re-trigger CI with the same bot command (no rebase needed).

⚠️ Action Required:

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

Link to invocation

[pranav-nvidia]

/bot run

[tensorrt-cicd]

PR_Github #70021 [ run ] triggered by Bot. Commit: fef5a1a Link to invocation

[tensorrt-cicd]

PR_Github #70021 [ run ] completed with state SUCCESS. Commit: fef5a1a
/LLM/main/L0_MergeRequest_PR pipeline #57302 completed with status: 'SUCCESS'

CI Report

Link to invocation