Skip to content

[None][fix] Remove redundant residual bound check in EAGLE3 hidden-state capture#16165

Open
eopXD wants to merge 1 commit into
NVIDIA:mainfrom
eopXD:eopXD/eagle3-remove-redundant-residual-check
Open

[None][fix] Remove redundant residual bound check in EAGLE3 hidden-state capture#16165
eopXD wants to merge 1 commit into
NVIDIA:mainfrom
eopXD:eopXD/eagle3-remove-redundant-residual-check

Conversation

@eopXD

@eopXD eopXD commented Jul 9, 2026

Copy link
Copy Markdown
Collaborator

Description

Follow-up to review feedback on #15708. During review, @lancelly noted that the residual bound check in Eagle3OneModelSpecMetadata.maybe_capture_hidden_states is redundant with the preceding hidden_states check. This PR removes it.

In every caller, residual and hidden_states are produced by the same decoder layer and therefore share their leading (token) dimension, so once the num_tokens > hidden_states.shape[0] check passes, num_tokens <= residual.shape[0] is already guaranteed — the residual branch was unreachable dead code. It is replaced with a short comment explaining why no separate check is needed, so it is not re-added.

The two remaining checks are intentionally kept because they guard different tensors:

  • the buffer-capacity check (self.hidden_states.shape[0]) protects the preallocated capture buffer that inplace_slice_copy writes into (the overrun the original fix addresses);
  • the available-hidden-states check (hidden_states.shape[0]) protects the input read.

No behavior change.

Test Coverage

No new tests. The existing GPU-gated test_eagle3_one_model_capture_uses_real_token_count in tests/unittest/_torch/speculative/test_eagle3.py exercises the retained residual path (hidden_states[:4] + residual[:4]); the removed branch was never reachable, so no test regresses. This is a pure dead-code removal on the CUDA-only maybe_capture_hidden_statesinplace_slice_copy path, covered by the same CI stages that passed on #15708.

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-compatible or api-breaking. For api-breaking, include BREAKING in 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.

Summary by CodeRabbit

  • Bug Fixes
    • Improved reliability when capturing hidden states during speculative decoding by simplifying internal bounds handling.
    • Reduced the chance of unnecessary runtime checks while preserving safe token-limited copying behavior.

…ate capture

The residual bound check in Eagle3OneModelSpecMetadata.maybe_capture_hidden_states
is unreachable: residual and hidden_states are always produced by the same
decoder layer and share their leading (token) dimension, so the preceding
num_tokens > hidden_states.shape[0] check already guarantees
num_tokens <= residual.shape[0]. Remove the dead branch and leave a note so it
is not re-added.

Follow-up to review feedback on NVIDIA#15708.

Co-Authored-By: Yueh-Ting Chen <yueh.ting.chen@gmail.com>
Signed-off-by: Yueh-Ting Chen <yuehtingc@nvidia.com>
@eopXD eopXD requested a review from a team as a code owner July 9, 2026 03:38
@eopXD eopXD requested review from cascade812 and lancelly July 9, 2026 03:38
@coderabbitai

coderabbitai Bot commented Jul 9, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: a80a9a3a-70ff-4262-a87a-7341a0e2d9ca

📥 Commits

Reviewing files that changed from the base of the PR and between 60583d2 and 2690e4c.

📒 Files selected for processing (1)
  • tensorrt_llm/_torch/speculative/eagle3.py

📝 Walkthrough

Walkthrough

Removed a redundant explicit runtime bounds check on residual.shape[0] within Eagle3OneModelSpecMetadata.maybe_capture_hidden_states, relying instead on the pre-existing hidden_states-based validation, with added comments explaining the assumption before slicing residual[:num_tokens].

Changes

Eagle3 hidden states capture bounds check removal

Layer / File(s) Summary
Simplify residual bounds validation
tensorrt_llm/_torch/speculative/eagle3.py
Removed the explicit RuntimeError check for residual.shape[0] vs num_tokens, replacing it with comments asserting the earlier hidden_states-based bound check already guarantees safe slicing.

Estimated code review effort: 2 (Simple) | ~5 minutes

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Title check ✅ Passed The title is concise and accurately summarizes the main change: removing a redundant residual bound check.
Description check ✅ Passed The description follows the template with Description, Test Coverage, and PR Checklist sections and explains the change clearly.
Docstring Coverage ✅ Passed No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Comment @coderabbitai help to get the list of available commands.

@eopXD

eopXD commented Jul 9, 2026

Copy link
Copy Markdown
Collaborator Author

/bot skip --comment "This is a trvial change that addresses comments from my previous merge request. This should be safe to land."

@eopXD eopXD enabled auto-merge (squash) July 9, 2026 06:05
@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #58392 [ skip ] triggered by Bot. Commit: 2690e4c Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #58392 [ skip ] completed with state SUCCESS. Commit: 2690e4c
Skipping testing for commit 2690e4c

Link to invocation

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants