[None][feat] Add encoder CUDA graph support to llm.encode()#14326
Conversation
|
/bot run |
📝 WalkthroughWalkthroughThis PR enables CUDA graph capture and replay for encoder-only ChangesEncoder-only CUDA Graph Capture and Replay
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
Actionable comments posted: 7
🤖 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/cuda_graph_runner.py`:
- Around line 885-915: The capture currently only refreshes input_ids,
position_ids and attn_metadata.seq_lens before replay, leaving other per-call
dynamic kwargs in capture_inputs (e.g., token_type_ids, inputs_embeds,
attention_mask) frozen; update the capture/replay logic in the CUDA graph
creation around capture_inputs / sliced_static_tensors /
sliced_static_tensors_cpu and forward_fn so that all dynamic per-call inputs are
refreshed before graph replay: either detect extra dynamic kwargs and fall back
to eager execution, or allocate pinned/static GPU buffers for each dynamic field
present and perform non_blocking.copy_ into capture_inputs for those fields
(similar to input_ids/position_ids handling, and also update attn_md fields),
and ensure the same change is applied to the corresponding replay code paths
later (the section around lines 924-984 mentioned in the review). Ensure checks
respect self._capture_h2d_copy and that graph_metadata still records
attn_metadata correctly.
- Around line 916-921: The code assigns self.graphs[key] before checking for
nested-tensor outputs, which can leave a half-registered graph if
self._contains_nested_tensor(output) raises; move the assignment of
self.graphs[key] so it occurs after the nested-tensor check (or ensure you
remove self.graphs[key] on error) and only set self.graph_outputs[key] =
make_weak_ref(output) after the check passes; adjust the logic in the method
that builds/stores graphs (the block that references self.graphs[key],
self._contains_nested_tensor, and self.graph_outputs[key]) so an exception won’t
leave an inconsistent entry.
In `@tensorrt_llm/_torch/pyexecutor/model_engine.py`:
- Around line 4243-4246: The returned graph outputs other than "logits" are
still backed by static replay buffers and must be cloned before returning;
update the return construction (where outputs = dict(graph_outputs) and logits
is sliced/cloned) to: for each key,value in graph_outputs (including non-logit
tensors) slice to batch_size if applicable and call .clone()/.detach() so no
tensor shares the underlying static buffer; specifically modify the code around
outputs = dict(graph_outputs) / outputs["logits"] = ... to iterate graph_outputs
items, clone/detach each tensor (and slice by batch_size when per-batch) and
place into outputs so future encode() calls cannot overwrite returned tensors.
In `@tensorrt_llm/llmapi/llm_args.py`:
- Around line 127-151: Make the encoder bucket list fields reject empty lists
and non-positive values: change num_tokens: Optional[List[int]] to
Optional[List[PositiveInt]] with Field(min_length=1, default=None,
description=...), change seq_lens similarly to Optional[List[PositiveInt]] with
Field(min_length=1, default=None,...), and change max_num_token and max_seq_len
from NonNegativeInt to PositiveInt (and update their default/validation so 0 is
not allowed). Apply the same adjustments to the analogous fields referenced
later (lines ~184-220) so list types require PositiveInt elements and lists have
min_length=1, and ensure validation logic that compares num_tokens/seq_lens to
max_num_token/max_seq_len remains consistent.
In `@tensorrt_llm/llmapi/llm.py`:
- Around line 757-759: The encode() signature and docstring are incorrect for
return_raw_logits: when return_raw_logits=True it returns a torch.Tensor, not
EncoderOutput(s); add explicit typing overloads using typing.overload above
encode()—one overload for return_raw_logits: Literal[True] -> torch.Tensor and
one for return_raw_logits: Literal[False] (or omitted) -> Union[EncoderOutput,
List[EncoderOutput]]—then update the implementation signature to return
Union[EncoderOutput, List[EncoderOutput], torch.Tensor] and revise the encode()
docstring to document the return_raw_logits behavior and the tensor return path;
apply the same overload+docstring change to the other similar encode variants
referenced (the other overload blocks for encode).
- Around line 851-855: The current guard in LLM.encode that raises
NotImplementedError when engine.encoder_cuda_graph_runner.enabled blocks all
non-reserved model_kwargs (filtered_kwargs) including documented supported
inputs like token_type_ids and inputs_embeds; change the check to whitelist
these allowed encoder inputs instead of rejecting everything: in the LLM.encode
code path that references filtered_kwargs and
engine.encoder_cuda_graph_runner.enabled, build an allowed_encoder_kwargs set
(e.g., {"token_type_ids", "inputs_embeds"}), subtract that from filtered_kwargs,
and only raise NotImplementedError if any unsupported keys remain (include the
remaining sorted keys in the error message as before).
In `@tests/integration/defs/accuracy/test_llm_api_pytorch_encode.py`:
- Around line 173-176: When running the CUDA-graph variants created with
LLM(model_path, encode_only=True, ..., cuda_graph_config=cgc), add an explicit
assertion that CUDA-graph capture/replay is actually active instead of relying
on logits differences: after calling outs = llm.encode(PROMPTS) (and before
comparing logits) assert that cgc is None or llm.is_cuda_graph_active() (or
llm.backend == "TRTLLM" if your code uses backend names), and if
is_cuda_graph_active() does not exist add a small accessor on LLM (e.g.,
is_cuda_graph_active or cuda_graph_active property) that returns whether
cuda_graph_config was activated so the test fails if execution silently fell
back to eager mode; apply the same pattern to the other occurrences around the
encode calls (the blocks at the other line ranges noted).
🪄 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: 7dddcd35-5603-4689-a484-b0b2bf1c7c1e
📒 Files selected for processing (20)
cpp/tensorrt_llm/nanobind/thop/bindings.cppcpp/tensorrt_llm/thop/attentionOp.cppcpp/tensorrt_llm/thop/attentionOp.htensorrt_llm/_torch/attention_backend/flashinfer.pytensorrt_llm/_torch/attention_backend/interface.pytensorrt_llm/_torch/attention_backend/trtllm.pytensorrt_llm/_torch/attention_backend/trtllm_gen.pytensorrt_llm/_torch/pyexecutor/cuda_graph_runner.pytensorrt_llm/_torch/pyexecutor/encoder_executor.pytensorrt_llm/_torch/pyexecutor/model_engine.pytensorrt_llm/llmapi/llm.pytensorrt_llm/llmapi/llm_args.pytensorrt_llm/llmapi/mm_encoder.pytests/integration/defs/accuracy/test_llm_api_pytorch_encode.pytests/integration/test_lists/qa/llm_function_core.txttests/integration/test_lists/qa/llm_function_l20.txttests/integration/test_lists/qa/llm_function_rtx6k.txttests/integration/test_lists/test-db/l0_a100.ymltests/unittest/llmapi/test_llm_args.pytests/unittest/llmapi/test_llm_encode.py
|
PR_Github #49268 [ run ] triggered by Bot. Commit: |
|
/bot kill |
|
PR_Github #49271 [ kill ] triggered by Bot. Commit: |
|
PR_Github #49268 [ run ] completed with state |
|
PR_Github #49271 [ kill ] completed with state |
dbb7997 to
adc3500
Compare
|
/bot run |
|
PR_Github #49272 [ run ] triggered by Bot. Commit: |
|
/bot kill |
|
/bot run --disable-fail-fast |
|
PR_Github #49278 [ kill ] triggered by Bot. Commit: |
|
PR_Github #49279 [ run ] triggered by Bot. Commit: |
|
PR_Github #49278 [ kill ] completed with state |
adc3500 to
97d9495
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #49754 [ run ] triggered by Bot. Commit: |
|
PR_Github #49754 [ run ] completed with state
|
Signed-off-by: tingyangk <tingyangk@nvidia.com>
Signed-off-by: tingyangk <tingyangk@nvidia.com>
Signed-off-by: tingyangk <tingyangk@nvidia.com>
Signed-off-by: tingyangk <tingyangk@nvidia.com>
97d9495 to
e1fc709
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #51421 [ run ] triggered by Bot. Commit: |
|
PR_Github #51421 [ run ] completed with state
|
|
/bot run |
|
PR_Github #51521 [ run ] triggered by Bot. Commit: |
|
PR_Github #51521 [ run ] completed with state
|
|
/bot run |
2ez4bz
left a comment
There was a problem hiding this comment.
Approving changes to tensorrt_llm/llmapi/mm_encoder.py.
|
/bot run |
|
PR_Github #51687 [ run ] triggered by Bot. Commit: |
|
PR_Github #51687 [ run ] completed with state |
Summary by CodeRabbit
Release Notes
New Features
return_raw_logitsparameter to encoder API for direct logits output.Improvements
Tests
Description
The PyTorch backend encoder forward path
llm.encode()currently runs in pure eager mode. Prefill-only encoder workloads (BERT classification, decoder-as-encoder) are dominated by launch overhead at small batch sizes or short sequence lengths.This PR introduces encoder CUDA graph infrastructure to the encode-only path. The main advantage is that, unlike piecewise CUDA graphs, it captures the entire model, including the attention layers, into one CUDA graph per bucket.
It adds a dedicated
EncoderCUDAGraphRunneralongside the decoder runner incuda_graph_runner.py, wires graph capture/replay intoPyTorchModelEngine.warmup_encoder()and the encode forward step, and surfaces the configuration through the existingCudaGraphConfigonTorchLlmArgs.Key design points:
Composite key. Encoder graphs use a 3-tuple
(padded_batch_size, padded_num_tokens, padded_max_seq_len)because two batches with the samenum_tokenscan still differ inmax(seq_lens)for context-only batches, and FMHA launch parameters depend on the latter. A newmax_context_q_len_overridefield — plumbed through to thethop.attentionop — locks the FMHA kernel launch parameters across replays even when the actual batch has shorter sequences.Stable attention metadata.
TrtllmAttentionMetadatagainsbind_encoder_cuda_graph_seq_lens()andprepare_encoder_cuda_graph_replay(). When CC (confidential computing) is off, pinned host buffers are used: per-replayseq_lensupdates and model inputs are CPU memcpys into pinned host buffers, and the host-to-device copies are captured in the graph to minimize CPU work in the replay path. When CC is on, pinned host buffers are not preferred, so the replay path updates the persistent GPU buffers directly before graph replay.Attention backend constraints. Only
TRTLLMattention metadata is eligible. Non-TRTLLM backends, including FlashInfer, fall back to eager because their planner/wrapper state is per batch and is not replay-safe.Input constraints. Encoder CUDA graphs currently support text-token inputs only. Additional encoder inputs such as
inputs_embedsor other model-specific tensors remain supported by the eagerllm.encode()path, but are not captured/replayed byEncoderCUDAGraphRunneryet.New API surface
CudaGraphConfig(tensorrt_llm/llmapi/llm_args.py) gains four encoder-only fields. They mirror the existingbatch_sizes / max_batch_sizepair:num_tokensOptional[List[int]]Nonemax_num_tokenNonNegativeInt0num_tokensis provided, must equalmax(num_tokens); otherwise the list is auto-generated from this.seq_lensOptional[List[int]]Nonemax_seq_lenNonNegativeInt0max_num_token, forseq_lens.Leaving all four at their defaults keeps encoder CUDA graphs disabled, even when
batch_sizesis set — graphs only activate when the user explicitly opts in to encoder bucketing.Usage
Minimal example — BERT classification with encoder CUDA graphs enabled:
The grid here captures graphs for the feasible graph buckets of the cross-product
{1, 4} × {16, 64} × {8, 32}; at inference timeenable_padding=Truerounds each incoming batch up to the nearest supported bucket. Decoder-only models can also use the encode-only path withencode_only=Trueand the same config.Performance
BERT 110M benchmark
baseline_generate: runs with oldllm.generate()pathbaseline_encode: runs with newllm.encode()path in eager modecuda-graph: runs with newllm.encode()path + encoder CUDA graph enabledtorch-compile+cuda-graph: runs with newllm.encode()path + TorchCompileConfig & encoder CUDA graph enabledgenerate: uses originalllm.generate()latency as baseline; shows the speedup from moving tollm.encode().encode: uses eagerllm.encode()latency as baseline; shows the speedup from encoder CUDA graphs.baseline_generatebaseline_encodecuda-graphtorch-compile+cuda-graphp50p99generateencodeQwen3-0.6B benchmark: PyTorch
llm.encode()variants vs legacy TRT pathmax_num_tokens=768for shorter (ISL <= 512) sequence lengthsmax_num_tokens=16384for longer (ISL > 512) sequence lengthseager,cuda-graph, ortorch-compile+cuda-graphand compares it with TRT; negative means the PyTorch variant is faster than TRT.cuda-graph(ms)torch-compile+cuda-graph(ms)Test Coverage
Integration tests
Add new tests to
tests/integration/defs/accuracy/test_llm_api_pytorch_encode.py:TestEncoderEncode(encoder-only architectures):test_encoder_encode_cuda_graph_matches_eager_logits[bert-yelp]— Tight numerical bound (rtol=atol=1e-3) between graph replay and eager output on the same model, to catch subtle launch-parameter or metadata drift that HF-tolerance wouldn't flag.TestDecoderEncode(decoder-as-encoder):test_decoder_encode_cuda_graph_matches_eager_logits[tinyllama-1.1b]— Tight graph-vs-eager bound for the single-prefill decoder path.Unit tests
Add new tests to
tests/unittest/llmapi/test_llm_encode.py:test_encode_cuda_graph_and_return_raw_logits— exercises capture + replay end-to-end and verifiesreturn_raw_logitsoutput wrapping on the graph path.test_trtllm_gen_is_supported_rejects_none_kv_cache_manager— asserts thetrtllm_genbackend refuses to be selected without akv_cache_manager.tests/unittest/llmapi/test_llm_args.pyadds coverage for the newCudaGraphConfigfields (auto-derivation ofmax_num_token/max_seq_len, mismatch error paths, sorted-list normalization).CI wiring
The new integration tests are added to:
tests/integration/test_lists/test-db/l0_a100.ymltests/integration/test_lists/qa/llm_function_core.txttests/integration/test_lists/qa/llm_function_l20.txttests/integration/test_lists/qa/llm_function_rtx6k.txtCC: @schetlur-nv @nvrohanv @amukkara @symphonylyh
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.