[https://nvbugs/6394425][fix] Added _is_effective_dynamic_tree() helper in utils.py that returns True only… by trtllm-agent · Pull Request #15842 · NVIDIA/TensorRT-LLM

trtllm-agent · 2026-07-01T15:10:56Z

Summary

Root cause: When dynamic_tree_max_topK=1 the validator promotes use_dynamic_tree=True and the dispatcher routes to Eagle3OneModelDynamicTreeWorker whose step-0 tree mask, KV rewind, and tree bookkeeping diverge numerically from the equivalent linear chain, lowering the acceptance length vs. the linear Eagle3OneModelWorker.
Fix: Added _is_effective_dynamic_tree() helper in utils.py that returns True only when use_dynamic_tree AND dynamic_tree_max_topK >= 2; used it at all three utils.py gate sites (get_spec_metadata, get_spec_resource_manager, get_spec_worker) so K=1 configs pick the linear metadata / resource manager / worker.
Automated fix generated by repair-bot

Test plan

Verify fix on the same GPU type as the original failure
Check for regressions in related tests

Links

Bug: https://nvbugs/6394425

Summary by CodeRabbit

Bug Fixes
- Improved how speculative decoding chooses the dynamic-tree path, so the feature now activates only when the relevant settings are properly enabled.
- Updated routing decisions for certain Eagle3 workflows to follow the same dynamic-tree behavior, helping keep execution consistent across modes.

When dynamic_tree_max_topK=1 the dynamic tree degenerates into a linear chain and should be equivalent to the linear Eagle3 one-model path, but the dynamic-tree worker's step-0 mask, KV rewind and tree bookkeeping produce a materially lower acceptance length than the linear worker. Add _is_effective_dynamic_tree() in _torch/speculative/utils.py that gates on both use_dynamic_tree and dynamic_tree_max_topK >= 2, and use it consistently in get_spec_metadata, get_spec_resource_manager and get_spec_worker so K=1 configs pick the linear metadata, resource manager and worker. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>

coderabbitai · 2026-07-01T15:14:06Z

📝 Walkthrough

Walkthrough

Adds a helper function _is_effective_dynamic_tree to compute dynamic-tree eligibility from both use_dynamic_tree and dynamic_tree_max_topK. Replaces direct uses of spec_config.use_dynamic_tree in eagle3/eagle3_one_model metadata construction and in resource manager/worker selection routing with this computed value.

Changes

Effective dynamic-tree helper and usage

Layer / File(s)	Summary
Effective dynamic-tree helper function `tensorrt_llm/_torch/speculative/utils.py`	Adds `_is_effective_dynamic_tree(spec_config)`, returning true only when `use_dynamic_tree` is set and `dynamic_tree_max_topK > 1`.
Metadata construction using effective flag `tensorrt_llm/_torch/speculative/utils.py`	Computes `effective_dynamic_tree` for the eagle3 branch and uses `_is_effective_dynamic_tree(spec_config)` to set `is_spec_dec_dynamic_tree` and `use_dynamic_tree` in `Eagle3SpecMetadata` and `Eagle3OneModelSpecMetadata`.
Resource manager and worker selection routing `tensorrt_llm/_torch/speculative/utils.py`	Updates `get_spec_resource_manager` and `get_spec_worker` to route eagle3_one_model dynamic-tree selection through `_is_effective_dynamic_tree(spec_config)` instead of `getattr(spec_config, 'use_dynamic_tree', False)`.

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

Related PRs: None identified.

Suggested labels: speculative-decoding, bug-fix

Suggested reviewers: None identified.

_{🐰 a tiny rabbit hops through logic gates, checking topK before it consents to grow its dynamic tree of thought.}

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is specific and matches the PR's main fix for effective dynamic-tree gating.
Description check	✅ Passed	The description covers the bug, fix, test plan, and link, matching the repository template closely enough.
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.}

coderabbitai

🧹 Nitpick comments (1)

tensorrt_llm/_torch/speculative/utils.py (1)
35-41: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Consider a small unit test for the topK=1 regression.

No test changes are included with this fix. Given this addresses a specific acceptance-length regression bug, a focused unit test asserting _is_effective_dynamic_tree returns False for dynamic_tree_max_topK=1 (and that get_spec_worker/get_spec_resource_manager route to the linear path) would guard against future regressions.
🤖 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/speculative/utils.py` around lines 35 - 41, The fix in
_is_effective_dynamic_tree should be covered with a focused regression test for
the topK=1 case. Add a unit test that verifies _is_effective_dynamic_tree
returns False when use_dynamic_tree is enabled but dynamic_tree_max_topK is 1,
and add a routing test around get_spec_worker and get_spec_resource_manager to
confirm they take the linear Eagle3 one-model path instead of the dynamic tree
path.

🤖 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.

Nitpick comments:
In `@tensorrt_llm/_torch/speculative/utils.py`:
- Around line 35-41: The fix in _is_effective_dynamic_tree should be covered
with a focused regression test for the topK=1 case. Add a unit test that
verifies _is_effective_dynamic_tree returns False when use_dynamic_tree is
enabled but dynamic_tree_max_topK is 1, and add a routing test around
get_spec_worker and get_spec_resource_manager to confirm they take the linear
Eagle3 one-model path instead of the dynamic tree path.

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: d94603a5-b6ac-46f3-afd9-499b6e1033e7

📥 Commits

Reviewing files that changed from the base of the PR and between d567f92 and 8614161.

📒 Files selected for processing (1)

tensorrt_llm/_torch/speculative/utils.py

sunnyqgg

LGTM

zhaoyangwang-nvidia · 2026-07-02T02:50:07Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-02T02:55:59Z

PR_Github #57077 [ run ] triggered by Bot. Commit: 8614161 Link to invocation

tensorrt-cicd · 2026-07-02T08:59:02Z

PR_Github #57077 [ run ] completed with state SUCCESS. Commit: 8614161
/LLM/main/L0_MergeRequest_PR pipeline #45866 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

ziyixiong-nv · 2026-07-03T01:13:01Z

/bot run

tensorrt-cicd · 2026-07-03T01:19:39Z

PR_Github #57285 [ run ] triggered by Bot. Commit: 8614161 Link to invocation

tensorrt-cicd · 2026-07-03T01:28:21Z

PR_Github #57285 [ run ] completed with state FAILURE. Commit: 8614161
/LLM/main/L0_MergeRequest_PR pipeline #46046 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

zhaoyangwang-nvidia · 2026-07-03T02:14:03Z

/bot run

tensorrt-cicd · 2026-07-03T02:19:38Z

PR_Github #57305 [ run ] triggered by Bot. Commit: 8614161 Link to invocation

tensorrt-cicd · 2026-07-03T02:31:25Z

PR_Github #57305 [ run ] completed with state FAILURE. Commit: 8614161
/LLM/main/L0_MergeRequest_PR pipeline #46064 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

xxi-nv · 2026-07-03T09:53:29Z

/bot run

tensorrt-cicd · 2026-07-03T09:58:52Z

PR_Github #57446 [ run ] triggered by Bot. Commit: 8614161 Link to invocation

tensorrt-cicd · 2026-07-03T11:50:39Z

PR_Github #57446 [ run ] completed with state SUCCESS. Commit: 8614161
/LLM/main/L0_MergeRequest_PR pipeline #46183 completed with status: 'SUCCESS'

CI Report

Link to invocation

… (fork PR #1) Supersedes the previous PR NVIDIA#15653 squash with the updated v2kv_mamba branch from #1: - V2 Mamba state snapshot reuse (updated after upstream review) - Fix KV cache v2 test request mocks - Fix V2 page table roles for Mamba pools - Support V2 Mamba hybrid cache in disagg Also carries the upstream main commits pulled in by the v2kv_mamba merge: DSv4 follow-up runtime KV and cache foundations (NVIDIA#15633), DeepGEMM and MegaMoE (NVIDIA#15632), and _is_effective_dynamic_tree fix (NVIDIA#15842). Signed-off-by: nv-guomingz <137257613+nv-guomingz@users.noreply.github.com>

… (fork PR #1) Supersedes the previous PR NVIDIA#15653 squash with the updated v2kv_mamba branch from nv-guomingz#1: - V2 Mamba state snapshot reuse (updated after upstream review) - Fix KV cache v2 test request mocks - Fix V2 page table roles for Mamba pools - Support V2 Mamba hybrid cache in disagg Also carries the upstream main commits pulled in by the v2kv_mamba merge: DSv4 follow-up runtime KV and cache foundations (NVIDIA#15633), DeepGEMM and MegaMoE (NVIDIA#15632), and _is_effective_dynamic_tree fix (NVIDIA#15842). Signed-off-by: nv-guomingz <137257613+nv-guomingz@users.noreply.github.com>

…er in utils.py that returns True only… (NVIDIA#15842) Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>

trtllm-agent requested a review from a team as a code owner July 1, 2026 15:10

trtllm-agent requested a review from sunnyqgg July 1, 2026 15:10

trtllm-agent assigned zhaoyangwang-nvidia Jul 1, 2026

github-actions Bot assigned trtllm-agent Jul 1, 2026

coderabbitai Bot reviewed Jul 1, 2026

View reviewed changes

sunnyqgg approved these changes Jul 2, 2026

View reviewed changes

zhaoyangwang-nvidia merged commit edf63e8 into NVIDIA:main Jul 3, 2026
11 checks passed

Uh oh!

Conversation

trtllm-agent commented Jul 1, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Test plan

Links

Summary by CodeRabbit

Uh oh!

coderabbitai Bot commented Jul 1, 2026

Walkthrough

Changes

❌ Failed checks (1 warning)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

sunnyqgg left a comment

Choose a reason for hiding this comment

Uh oh!

zhaoyangwang-nvidia commented Jul 2, 2026

Uh oh!

tensorrt-cicd commented Jul 2, 2026

Uh oh!

tensorrt-cicd commented Jul 2, 2026

Uh oh!

ziyixiong-nv commented Jul 3, 2026

Uh oh!

tensorrt-cicd commented Jul 3, 2026

Uh oh!

tensorrt-cicd commented Jul 3, 2026

Uh oh!

zhaoyangwang-nvidia commented Jul 3, 2026

Uh oh!

tensorrt-cicd commented Jul 3, 2026

Uh oh!

tensorrt-cicd commented Jul 3, 2026

Uh oh!

xxi-nv commented Jul 3, 2026

Uh oh!

tensorrt-cicd commented Jul 3, 2026

Uh oh!

tensorrt-cicd commented Jul 3, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

trtllm-agent commented Jul 1, 2026 •

edited by coderabbitai Bot

Loading