[#9241][feat] AutoDeploy: Support Eagle3 Speculative Decoding

[govind-ramnarayan]

@coderabbitai summary

Description

Fixes: #9241.

Implements Eagle3 speculative decoding for AutoDeploy.

Layers are captured by detecting residual nodes after hidden states and saving those buffers. This thanks to @lucaslie .

Those hidden states are fed to the Eagle3 Resource Manager to coordinate with the draft model running in the PyTorch backend.

We only support the two-model speculative decoding flow still. Additionally, we have the overlap scheduler and chain drafter both disabled, as these both cause issues. These will need to be resolved when we want a performant 2-model solution.

Additionally, only tp=1 is currently supported. A different embedding function implementation will need to be used for tp > 1.

Test Coverage

• Added to integration test for testing correctness with Eagle3.
• Verified via print statements that distribution of num accepted tokens approximately matches PyTorch backend.

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)

• 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

/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...

Provide a user friendly way for developers to interact with a Jenkins server.

Run /bot [-h|--help] to print this help message.

See details below for each supported subcommand.

Details

run [--reuse-test (optional)pipeline-id --disable-fail-fast --skip-test --stage-list "A10-PyTorch-1, xxx" --gpu-type "A30, H100_PCIe" --test-backend "pytorch, cpp" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" --detailed-log --debug(experimental)]

Launch build/test pipelines. All previously running jobs will be killed.

--reuse-test (optional)pipeline-id (OPTIONAL) : Allow the new pipeline to reuse build artifacts and skip successful test stages from a specified pipeline or the last pipeline if no pipeline-id is indicated. If the Git commit ID has changed, this option will be always ignored. The DEFAULT behavior of the bot is to reuse build artifacts and successful test results from the last pipeline.

--disable-reuse-test (OPTIONAL) : Explicitly prevent the pipeline from reusing build artifacts and skipping successful test stages from a previous pipeline. Ensure that all builds and tests are run regardless of previous successes.

--disable-fail-fast (OPTIONAL) : Disable fail fast on build/tests/infra failures.

--skip-test (OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.

--stage-list "A10-PyTorch-1, xxx" (OPTIONAL) : Only run the specified test stages. Examples: "A10-PyTorch-1, xxx". Note: Does NOT update GitHub check status.

--gpu-type "A30, H100_PCIe" (OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.

--test-backend "pytorch, cpp" (OPTIONAL) : Skip test stages which don't match the specified backends. Only support [pytorch, cpp, tensorrt, triton]. Examples: "pytorch, cpp" (does not run test stages with tensorrt or triton backend). Note: Does NOT update GitHub pipeline status.

--only-multi-gpu-test (OPTIONAL) : Only run the multi-GPU tests. Note: Does NOT update GitHub check status.

--disable-multi-gpu-test (OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.

--add-multi-gpu-test (OPTIONAL) : Force run the multi-GPU tests in addition to running L0 pre-merge pipeline.

--post-merge (OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.

--extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" (OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx".

--detailed-log (OPTIONAL) : Enable flushing out all logs to the Jenkins console. This will significantly increase the log volume and may slow down the job.

--debug (OPTIONAL) : Experimental feature. Enable access to the CI container for debugging purpose. Note: Specify exactly one stage in the stage-list parameter to access the appropriate container environment. Note: Does NOT update GitHub check status.

For guidance on mapping tests to stage names, see docs/source/reference/ci-overview.md
and the scripts/test_to_stage_mapping.py helper.

kill

kill

Kill all running builds associated with pull request.

skip

skip --comment COMMENT

Skip testing for latest commit on pull request. --comment "Reason for skipping build/test" is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

reuse-pipeline

reuse-pipeline

Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

[govind-ramnarayan]

/bot run --disable-fail-fast

[coderabbitai]

📝 Walkthrough

Walkthrough

This pull request adds Eagle3 speculative decoding support to AutoDeploy with hidden state capture capability. It introduces configuration transforms, a hidden state manager, custom operations for residual add caching, and updates the executor to integrate hidden state capture during model execution.

Changes

Cohort / File(s)	Summary
Configuration and transform registration tensorrt_llm/_torch/auto_deploy/config/default.yaml	Added two instances of detect_hidden_states_for_capture transform (stage: pattern_matcher) and insert_cached_residual_add transform (stage: cache_init, backend: cached_residual_add) to enable hidden state capture in the transform pipeline.
Configuration and validation logic tensorrt_llm/_torch/auto_deploy/llm_args.py	Added Set type import, EagleDecodingConfig import, and default_eagle3_layers_to_capture() helper function. Introduced draft_checkpoint_loader field to AutoDeployConfig and added two @model_validator methods to populate eagle3_layers_to_capture from speculative config in both AutoDeployConfig and LlmArgs.
Hidden state capture operations and transforms tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py	New module implementing custom residual add operations (residual_add_for_capture, cached_residual_add) with fake variants, metadata preparation function, and three transform/descriptor classes: DetectHiddenStatesForCaptureConfig, DetectHiddenStatesForCapture, CachedResidualAdd, and InsertCachedResidualAdd.
Executor integration and resource management tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py	Added ADHiddenStateManager class for managing hidden state capture, share_embedding_weights() helper for weight sharing, updated function signatures in forward() and _compute_logits() to accept resource_manager, modified create_draft_model_engine_maybe() to accept target_engine instead of engine, and added imports (MethodType, F, Eagle3ResourceManager, EagleDecodingConfig, RequestList).
Test infrastructure updates tests/integration/defs/examples/test_ad_speculative_decoding.py	Added EagleDecodingConfig import, new model constants (EAGLE_MODEL_SUBPATH, LLAMA_BASE_SUBPATH, draft length constants), extended get_model_paths() to return three paths, added make_spec_config() helper to create appropriate decoding configs, updated run_with_autodeploy() signature to accept spec_dec_mode, and parametrized test to run with both draft_target and eagle modes.
Test configuration tests/integration/test_lists/test-db/l0_h100.yml	Updated AutoDeploy test entries to replace test_autodeploy_spec_dec[1] and test_autodeploy_spec_dec[4] with parametrized variants test_autodeploy_spec_dec[1-draft_target] and test_autodeploy_spec_dec[4-eagle].

Sequence Diagram(s)

sequenceDiagram
    participant Config as Config & Validators
    participant Executor as ADEngine
    participant ResourceMgr as ADHiddenStateManager
    participant ModelExec as Model Execution
    participant Capture as Hidden State Capture

    Config->>Config: Apply model validators
    Config->>Config: Populate eagle3_layers_to_capture
    Config->>ResourceMgr: Initialize with layers_to_capture
    
    Executor->>ResourceMgr: Call prepare_hidden_states_capture()
    ResourceMgr->>ResourceMgr: Locate hidden_state buffers by layer
    ResourceMgr->>ResourceMgr: Prepare write indices for capture
    
    Executor->>ModelExec: Execute forward pass
    ModelExec->>ModelExec: Transform pipeline applies cached_residual_add ops
    ModelExec-->>Executor: Return logits
    
    Executor->>ResourceMgr: Call capture_hidden_states()
    ResourceMgr->>Capture: Extract hidden states from buffers
    Capture-->>ResourceMgr: Store captured states at indices

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~50 minutes

• ADHiddenStateManager implementation: New class with multiple methods (_get_hidden_state_buffers, prepare_hidden_states_capture, capture_hidden_states) requiring careful validation of buffer access patterns and index management.
• hidden_states.py module: New file containing custom operations, metadata preparation logic, and transform classes with graph manipulation and fake operation implementations; requires understanding of the transform pipeline and cache initialization patterns.
• Model validator logic: Two new @model_validator methods with conditional population of eagle3_layers_to_capture based on speculative config type; logic propagation across AutoDeployConfig and LlmArgs should be verified for consistency.
• Function signature propagation: Changes to _compute_logits() and forward() signatures to thread resource_manager through the executor; verify all call sites are updated and semantics are preserved.
• Embedding weight sharing: New share_embedding_weights() helper using MethodType for dynamic binding; ensure weight aliasing semantics are correct and no unintended side effects occur.

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 31.71% which is insufficient. The required threshold is 80.00%.	You can run @coderabbitai generate docstrings to improve docstring coverage.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	The PR implements Eagle3 speculative decoding support for AutoDeploy with hidden state capture and coordination with the draft model, directly addressing the core requirement in #9421 to support speculative decoding in AutoDeploy without requiring TRTLLM attention backend.
Out of Scope Changes check	✅ Passed	All changes are focused on Eagle3 speculative decoding implementation: config additions, hidden state capture transforms, resource management, and related test updates. No unrelated changes detected.
Title check	✅ Passed	The pull request title clearly summarizes the main change: adding Eagle3 speculative decoding support to AutoDeploy, which is the primary feature being introduced.
Description check	✅ Passed	The PR description provides a clear explanation of the issue (#9241), the solution (implementing Eagle3 speculative decoding), and technical details about hidden state capture and limitations (tp=1, scheduler issues). However, the Test Coverage section is incomplete with only mentions of integration tests and print statement verification.

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 9

🧹 Nitpick comments (4)

tensorrt_llm/_torch/auto_deploy/llm_args.py (1)

166-169: Consider a more specific type for draft_checkpoint_loader.

Optional[object] is very generic. If there's a base class or protocol for checkpoint loaders, using that would improve type safety and IDE support.

tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py (1)

156-156: Add strict=True to zip() for safety.

Using strict=True ensures ordered_requests and seq_lens have matching lengths, providing an early error if there's a mismatch.

-        for request, seq_len in zip(ordered_requests, seq_lens):
+        for request, seq_len in zip(ordered_requests, seq_lens, strict=True):

tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py (2)

161-161: Prefix unused variable with underscore.

unprocessed_linear_nodes is not used. Per Python convention and static analysis, prefix it with an underscore.

-        layer_subgraphs, unprocessed_linear_nodes = get_all_layer_subgraphs(gm)
+        layer_subgraphs, _unprocessed_linear_nodes = get_all_layer_subgraphs(gm)

---

177-181: Prefer next(iter(...)) for single-element access.

Using next(iter(res_node.users)) is more efficient than creating a list for single element access.

             while len(res_node.users) == 1:
-                user_node = list(res_node.users)[0]
+                user_node = next(iter(res_node.users))
                 if not is_op(user_node, torch.ops.aten.add):
                     break
                 res_node = user_node

📜 Review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 2c0293c and c10c1ce.

📒 Files selected for processing (6)

• tensorrt_llm/_torch/auto_deploy/config/default.yaml (2 hunks)
• tensorrt_llm/_torch/auto_deploy/llm_args.py (4 hunks)
• tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py (12 hunks)
• tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py (1 hunks)
• tests/integration/defs/examples/test_ad_speculative_decoding.py (6 hunks)
• tests/integration/test_lists/test-db/l0_h100.yml (1 hunks)

🧰 Additional context used

📓 Path-based instructions (2)

**/*.py

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

**/*.py: The code developed for TensorRT-LLM should conform to Python 3.8+
Indent Python code with 4 spaces; do not use tabs
Always maintain the namespace when importing in Python, even if only one class or function from a module is used (e.g., use from package.subpackage import foo and then foo.SomeClass() instead of from package.subpackage.foo import SomeClass)
Python filenames should use snake_case (e.g., some_file.py)
Python class names should use PascalCase (e.g., class SomeClass)
Python function and method names should use snake_case (e.g., def my_awesome_function():)
Python local variable names should use snake_case, with prefix k for variable names that start with a number (e.g., k_99th_percentile = ...)
Python global variables should use upper snake_case with prefix G (e.g., G_MY_GLOBAL = ...)
Python constants should use upper snake_case (e.g., MY_CONSTANT = ...)
Avoid shadowing variables declared in an outer scope in Python
Initialize all externally visible members of a Python class in the constructor
For Python interfaces that may be used outside a file, prefer docstrings over comments
Python comments should be reserved for code within a function, or interfaces that are local to a file
Use Google style docstrings for Python classes and functions, which can be parsed by Sphinx
Python attributes and variables can be documented inline with type and description (e.g., self.x = 5 followed by """<type>: Description of 'x'""" )
Avoid using reflection in Python when functionality can be easily achieved without reflection
When using try-except blocks in Python, limit the except clause to the smallest set of specific errors possible instead of catching all exceptions
When using try-except blocks in Python to handle multiple possible variable types (duck-typing), keep the body of the try as small as possible and use the else block to implement the logic

Files:

• tensorrt_llm/_torch/auto_deploy/llm_args.py
• tests/integration/defs/examples/test_ad_speculative_decoding.py
• tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py
• tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py

**/*.{cpp,h,cu,py}

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

All TensorRT-LLM Open Source Software code files should contain an NVIDIA copyright header that includes the current year at the top

Files:

• tensorrt_llm/_torch/auto_deploy/llm_args.py
• tests/integration/defs/examples/test_ad_speculative_decoding.py
• tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py
• tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py

🧠 Learnings (9)

📚 Learning: 2025-08-26T09:37:10.463Z

Learnt from: jiaganc
Repo: NVIDIA/TensorRT-LLM PR: 7031
File: tensorrt_llm/bench/dataclasses/configuration.py:90-104
Timestamp: 2025-08-26T09:37:10.463Z
Learning: In TensorRT-LLM, the `get_pytorch_perf_config()` method returns `self.pytorch_config` which can contain default `cuda_graph_config` values, so `llm_args` may already have this config before the extra options processing.

Applied to files:

• tensorrt_llm/_torch/auto_deploy/llm_args.py
• tests/integration/defs/examples/test_ad_speculative_decoding.py
• tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py

📚 Learning: 2025-08-26T09:37:10.463Z

Learnt from: jiaganc
Repo: NVIDIA/TensorRT-LLM PR: 7031
File: tensorrt_llm/bench/dataclasses/configuration.py:90-104
Timestamp: 2025-08-26T09:37:10.463Z
Learning: In TensorRT-LLM's bench configuration, the `get_pytorch_perf_config()` method returns `self.pytorch_config` which is a Dict[str, Any] that can contain default values including `cuda_graph_config`, making the fallback `llm_args["cuda_graph_config"]` safe to use.

Applied to files:

• tensorrt_llm/_torch/auto_deploy/llm_args.py
• tests/integration/defs/examples/test_ad_speculative_decoding.py
• tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py

📚 Learning: 2025-07-28T17:06:08.621Z

Learnt from: moraxu
Repo: NVIDIA/TensorRT-LLM PR: 6303
File: tests/integration/test_lists/qa/examples_test_list.txt:494-494
Timestamp: 2025-07-28T17:06:08.621Z
Learning: In TensorRT-LLM testing, it's common to have both CLI flow tests (test_cli_flow.py) and PyTorch API tests (test_llm_api_pytorch.py) for the same model. These serve different purposes: CLI flow tests validate the traditional command-line workflow, while PyTorch API tests validate the newer LLM API backend. Both are legitimate and should coexist.

Applied to files:

• tensorrt_llm/_torch/auto_deploy/llm_args.py
• tests/integration/test_lists/test-db/l0_h100.yml

📚 Learning: 2025-08-06T13:58:07.506Z

Learnt from: galagam
Repo: NVIDIA/TensorRT-LLM PR: 6487
File: tests/unittest/_torch/auto_deploy/unit/singlegpu/test_ad_trtllm_bench.py:1-12
Timestamp: 2025-08-06T13:58:07.506Z
Learning: In TensorRT-LLM, test files (files under tests/ directories) do not require NVIDIA copyright headers, unlike production source code files. Test files typically start directly with imports, docstrings, or code.

Applied to files:

• tensorrt_llm/_torch/auto_deploy/llm_args.py

📚 Learning: 2025-11-24T17:09:17.870Z

Learnt from: CR
Repo: NVIDIA/TensorRT-LLM PR: 0
File: CODING_GUIDELINES.md:0-0
Timestamp: 2025-11-24T17:09:17.870Z
Learning: Applies to **/*.py : The code developed for TensorRT-LLM should conform to Python 3.8+

Applied to files:

• tensorrt_llm/_torch/auto_deploy/llm_args.py

📚 Learning: 2025-09-09T09:40:45.658Z

Learnt from: fredricz-20070104
Repo: NVIDIA/TensorRT-LLM PR: 7645
File: tests/integration/test_lists/qa/llm_function_core.txt:648-648
Timestamp: 2025-09-09T09:40:45.658Z
Learning: In TensorRT-LLM test lists, it's common and intentional for the same test to appear in multiple test list files when they serve different purposes (e.g., llm_function_core.txt for comprehensive core functionality testing and llm_function_core_sanity.txt for quick sanity checks). This duplication allows tests to be run in different testing contexts.

Applied to files:

• tensorrt_llm/_torch/auto_deploy/llm_args.py
• tests/integration/test_lists/test-db/l0_h100.yml

📚 Learning: 2025-08-26T09:49:04.956Z

Learnt from: pengbowang-nv
Repo: NVIDIA/TensorRT-LLM PR: 7192
File: tests/integration/test_lists/test-db/l0_dgx_b200.yml:56-72
Timestamp: 2025-08-26T09:49:04.956Z
Learning: In TensorRT-LLM test configuration files, the test scheduling system handles wildcard matching with special rules that prevent duplicate test execution even when the same tests appear in multiple yaml files with overlapping GPU wildcards (e.g., "*b200*" and "*gb200*").

Applied to files:

• tests/integration/test_lists/test-db/l0_h100.yml

📚 Learning: 2025-11-14T11:22:03.729Z

Learnt from: nzmora-nvidia
Repo: NVIDIA/TensorRT-LLM PR: 9163
File: tensorrt_llm/_torch/auto_deploy/custom_ops/quant.py:107-113
Timestamp: 2025-11-14T11:22:03.729Z
Learning: In TensorRT-LLM AutoDeploy custom ops, when adding hardware capability checks to select between kernel implementations (e.g., cuBLAS vs. CUDA kernel), use descriptive variable names that identify the specific GPU architectures or families being targeted (e.g., `is_blackwell_geforce_or_ada`) rather than generic names like `enable_cuda_core`. This makes it clear that the code is selecting an implementation path based on hardware capabilities, not enabling/disabling hardware features.

Applied to files:

• tests/integration/test_lists/test-db/l0_h100.yml

📚 Learning: 2025-10-20T17:07:18.745Z

Learnt from: nvchenghaoz
Repo: NVIDIA/TensorRT-LLM PR: 8469
File: tensorrt_llm/_torch/auto_deploy/models/patches/nemotron_h.py:98-116
Timestamp: 2025-10-20T17:07:18.745Z
Learning: In NemotronH models (tensorrt_llm/_torch/auto_deploy/models/patches/nemotron_h.py), the gate (self.gate) returns topk_indices and topk_weights that are already in the correct shape to be passed directly to torch_ops.auto_deploy.torch_moe without needing to reshape them when hidden_states is flattened.

Applied to files:

• tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py

🧬 Code graph analysis (3)

tensorrt_llm/_torch/auto_deploy/llm_args.py (1)

tensorrt_llm/llmapi/llm_args.py (4)

• Field (63-90)
• BaseLlmArgs (1719-2143)
• EagleDecodingConfig (774-891)
• _ParallelConfig (460-525)

tests/integration/defs/examples/test_ad_speculative_decoding.py (1)

tensorrt_llm/llmapi/llm_args.py (7)

• DraftTargetDecodingConfig (995-1008)
• EagleDecodingConfig (774-891)
• spec_dec_mode (727-734)
• spec_dec_mode (869-874)
• spec_dec_mode (923-926)
• spec_dec_mode (1055-1062)
• speculative_model_dir (1983-1984)

tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py (14)

tests/unittest/_torch/thop/parallel/test_custom_ops.py (1)

• custom_ops (37-42)

tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py (4)

• AttentionDescriptor (761-889)
• AttentionRegistry (892-915)
• MHACallable (728-732)
• PrepareMetadataCallable (735-745)

tensorrt_llm/_torch/auto_deploy/llm.py (1)

• factory (110-113)

tensorrt_llm/_torch/auto_deploy/shim/interface.py (1)

• CachedSequenceInterface (11-92)

cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.h (1)

• utils (102-105)

tensorrt_llm/_torch/auto_deploy/utils/node_utils.py (2)

• get_all_layer_subgraphs (426-477)
• is_op (198-221)

tensorrt_llm/_torch/attention_backend/flashinfer.py (1)

• page_size (200-204)

cpp/tests/unit_tests/kernels/fusedMoeCommKernelTest.cpp (1)

• skipped (60-66)

tensorrt_llm/_torch/auto_deploy/transform/library/fused_moe.py (1)

• target (602-603)

tensorrt_llm/_torch/auto_deploy/utils/pattern_matcher.py (1)

• call_function (249-276)

tensorrt_llm/builder.py (1)

• default (45-50)

tensorrt_llm/functional.py (1)

• replace_all_uses_with (556-573)

tensorrt_llm/llmapi/utils.py (1)

• empty (388-389)

tests/unittest/_torch/modeling/test_modeling_out_of_tree.py (1)

• max_num_tokens (63-66)

🪛 Ruff (0.14.8)

tensorrt_llm/_torch/auto_deploy/llm_args.py

50-50: Avoid specifying long messages outside the exception class

(TRY003)

tests/integration/defs/examples/test_ad_speculative_decoding.py

65-65: Avoid specifying long messages outside the exception class

(TRY003)

tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py

142-144: Avoid specifying long messages outside the exception class

(TRY003)

---

156-156: zip() without an explicit strict= parameter

Add explicit value for parameter strict=

(B905)

---

167-170: Avoid specifying long messages outside the exception class

(TRY003)

---

812-814: Avoid specifying long messages outside the exception class

(TRY003)

---

817-819: Avoid specifying long messages outside the exception class

(TRY003)

tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py

63-63: Unused function argument: hidden_states_cache

(ARG001)

---

138-138: Unused method argument: cm

(ARG002)

---

139-139: Unused method argument: factory

(ARG002)

---

140-140: Unused method argument: shared_config

(ARG002)

---

161-161: Unpacked variable unprocessed_linear_nodes is never used

Prefix it with an underscore or any other dummy variable pattern

(RUF059)

---

178-178: Prefer next(iter(res_node.users)) over single element slice

Replace with next(iter(res_node.users))

(RUF015)

---

193-193: Loop control variable layer_number not used within loop body

(B007)

---

238-238: Unused class method argument: cache_config

(ARG003)

---

249-249: Unused class method argument: source_attn_node

(ARG003)

---

253-253: Unused class method argument: source_attn_node

(ARG003)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)

• GitHub Check: Pre-commit Check

🔇 Additional comments (16)

tests/integration/test_lists/test-db/l0_h100.yml (1)

120-121: LGTM!

The test entries correctly reflect the updated parametrization with spec_dec_mode values (draft_target and eagle), matching the test function signature changes in test_ad_speculative_decoding.py.

tensorrt_llm/_torch/auto_deploy/config/default.yaml (2)

78-79: LGTM!

The detect_hidden_states_for_capture transform is appropriately placed in the pattern_matcher stage alongside other pattern detection transforms.

---

165-167: LGTM!

The insert_cached_residual_add transform is correctly placed in the cache_init stage, consistent with the other insert_cached_* transforms that handle caching mechanisms.

tensorrt_llm/_torch/auto_deploy/llm_args.py (1)

11-18: LGTM!

Import organization follows the coding guidelines by maintaining the namespace (importing from ...llmapi.llm_args).

tests/integration/defs/examples/test_ad_speculative_decoding.py (3)

120-122: Workaround note: bypassing Pydantic validation.

The comment indicates this bypasses validation for eagle3_layers_to_capture. This is acceptable for tests, but ensure the production code path (in LlmArgs.default_eagle3_layers_to_capture validator) handles this correctly.

---

53-65: LGTM!

The make_spec_config helper cleanly encapsulates speculative config creation for both modes. The eagle3_one_model=False setting correctly configures two-model Eagle3 mode for this test.

---

140-141: LGTM!

The parametrization correctly tests draft_target mode with batch_size=1 and eagle mode with batch_size=4, providing coverage for both speculative decoding paths.

tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py (5)

114-131: LGTM on ADHiddenStateManager initialization.

The class properly extends Eagle3ResourceManager and pre-allocates the hidden_state_write_indices tensor. The initialization order (calling _get_hidden_state_buffers before super().__init__) is necessary to determine dtype and hidden_size for the parent constructor.

---

640-659: TP > 1 limitation is documented.

The assertion at line 650 and the TODO comment at line 653 correctly document that this embedding weight sharing approach is limited to tp <= 1. This is acceptable for the initial implementation.

Ensure this limitation is documented in user-facing documentation or release notes for Eagle3 support in AutoDeploy.

---

809-814: LGTM!

The speculative decoding mode check correctly extends to include Eagle3 mode alongside draft target mode.

---

577-582: LGTM!

The hidden state preparation is correctly gated by checking both the existence of spec_resource_manager and its type being ADHiddenStateManager.

---

595-601: LGTM!

Hidden state capture is correctly triggered after model execution, with appropriate type checking.

tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py (4)

34-42: LGTM!

The residual_add_for_capture custom op correctly wraps torch.ops.aten.add and serves as a placeholder marker for detection during graph transformation.

---

114-123: Configuration looks reasonable.

The TODO (lines 117-121) acknowledges a design consideration about extracting layer indices from Eagle checkpoints. This is a valid future improvement but doesn't block the current implementation.

---

210-255: LGTM!

The CachedResidualAdd descriptor correctly implements the AttentionDescriptor interface. The cache initializer properly allocates based on max_num_tokens and hidden_size from node metadata. Unused arguments in interface methods are expected per the abstract contract.

---

203-207: Verify is_clean and has_valid_shapes semantics.

When cnt > 0 (transform successfully matched and modified nodes), both is_clean and has_valid_shapes are set to False. Please verify this aligns with the TransformInfo contract—typically is_clean=False indicates the graph was modified, but has_valid_shapes=False might incorrectly signal shape validation failure rather than serving as a flag for successful transformation completion. Confirm the intended semantics of these fields by reviewing the TransformInfo class definition and checking how other transforms in the library set these fields when reporting successful modifications.

[tensorrt-cicd]

PR_Github #27772 [ run ] triggered by Bot. Commit: c10c1ce

[govind-ramnarayan]

/bot kill

[govind-ramnarayan]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #27787 [ kill ] triggered by Bot. Commit: 95bfa11

[tensorrt-cicd]

PR_Github #27787 [ kill ] completed with state SUCCESS. Commit: 95bfa11
Successfully killed previous jobs for commit 95bfa11

[tensorrt-cicd]

PR_Github #27788 [ run ] triggered by Bot. Commit: 95bfa11

[tensorrt-cicd]

PR_Github #27788 [ run ] completed with state SUCCESS. Commit: 95bfa11
/LLM/main/L0_MergeRequest_PR pipeline #21207 completed with status: 'FAILURE'

[govind-ramnarayan]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #27919 [ run ] triggered by Bot. Commit: f9d2253

[tensorrt-cicd]

PR_Github #27919 [ run ] completed with state SUCCESS. Commit: f9d2253
/LLM/main/L0_MergeRequest_PR pipeline #21319 completed with status: 'FAILURE'

[govind-ramnarayan]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #28430 [ run ] triggered by Bot. Commit: 80a7c3f

[tensorrt-cicd]

PR_Github #28430 [ run ] completed with state SUCCESS. Commit: 80a7c3f
/LLM/main/L0_MergeRequest_PR pipeline #21759 completed with status: 'FAILURE'

[govind-ramnarayan]

/bot run --disable-fail-fast

[govind-ramnarayan]

/bot run

[tensorrt-cicd]

PR_Github #29844 [ run ] triggered by Bot. Commit: d8a1cba

[tensorrt-cicd]

PR_Github #29844 [ run ] completed with state SUCCESS. Commit: d8a1cba
/LLM/main/L0_MergeRequest_PR pipeline #22946 completed with status: 'FAILURE'

⚠️ 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

[govind-ramnarayan]

/bot run

[tensorrt-cicd]

PR_Github #29848 [ run ] triggered by Bot. Commit: 131a588

[tensorrt-cicd]

PR_Github #29848 [ run ] completed with state SUCCESS. Commit: 131a588
/LLM/main/L0_MergeRequest_PR pipeline #22950 completed with status: 'FAILURE'

⚠️ 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

[govind-ramnarayan]

/bot run --reuse-test

[tensorrt-cicd]

PR_Github #29854 [ run ] triggered by Bot. Commit: 131a588

[tensorrt-cicd]

PR_Github #29854 [ run ] completed with state SUCCESS. Commit: 131a588
/LLM/main/L0_MergeRequest_PR pipeline #22956 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.