[None][feat] Support Nemotron dynamic-tree MTP decoding#15582
Conversation
|
/bot run --disable-fail-fast |
|
PR_Github #55465 [ run ] triggered by Bot. Commit: |
📝 WalkthroughWalkthroughThis PR implements greedy dynamic-tree one-model MTP-Eagle speculative decoding. It adds new ChangesMTP-Eagle Dynamic-Tree Speculative Decoding
Sequence DiagramsequenceDiagram
rect rgba(100, 149, 237, 0.5)
note over MTPEagleDynamicTreeWorker: Target-Verify Step
MTPEagleDynamicTreeWorker->>MTPEagleDynamicTreeWorker: sample_and_accept_draft_tokens
MTPEagleDynamicTreeWorker->>verify_dynamic_tree_greedy_out_packed: tree logits, packed mask
verify_dynamic_tree_greedy_out_packed-->>MTPEagleDynamicTreeWorker: accepted tokens, lengths, indices
MTPEagleDynamicTreeWorker->>MambaHybridCacheManager: update_mamba_states(accepted_leaf_positions)
MambaHybridCacheManager->>MambaHybridCacheManager: SSM + conv Triton promotion w/ accepted_positions
end
rect rgba(60, 179, 113, 0.5)
note over MTPEagleDynamicTreeWorker: Draft-Build Step
MTPEagleDynamicTreeWorker->>MTPEagleDynamicTreeWorker: _prepare_attn_metadata_for_spec_dec (snapshot)
loop MTP draft layers
MTPEagleDynamicTreeWorker->>NemotronHMTPDecoderLayer: forward(spec_metadata, mamba_metadata)
NemotronHMTPDecoderLayer->>Mamba2Mixer: forward w/ retrieve_next_token/sibling/parent
Mamba2Mixer->>causal_conv1d_update_triton: dynamic-tree link tensors
Mamba2Mixer->>selective_state_update_native: ssu_extra(retrieve_parent_token)
MTPEagleDynamicTreeWorker->>MTPEagleDynamicTreeWorker: sample top-K, update_draft_tokens_and_scores
end
MTPEagleDynamicTreeWorker->>SpecTreeManager: build_dynamic_tree → scatter slot_storage
MTPEagleDynamicTreeWorker->>MTPEagleDynamicTreeWorker: _restore_attn_metadata_from_spec_dec
end
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (3)
tests/unittest/_torch/speculative/hw_agnostic/test_dynamic_tree_slot_storage.py (1)
23-77: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd coverage for mask/position fallback and invalidation.
Coverage in
tests/unittest/_torch/speculative/hw_agnostic/test_dynamic_tree_slot_storage.pyis insufficient for this PR: the tests cover retrieve links, but not the changed no-treeposition_offsets/packed_mask, dummy-slot K-ary mask, ormark_invalid()restoration. Please add those cases in this file; no follow-up outside this PR should be needed.As per path instructions,
tests/**: “Act as a QA engineer reviewing test changes and coverage for TensorRT-LLM.”🤖 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 `@tests/unittest/_torch/speculative/hw_agnostic/test_dynamic_tree_slot_storage.py` around lines 23 - 77, The current tests in DynamicTreeSlotStorage only validate next-link chaining and miss the new no-tree fallback and invalidation behavior. Extend test_dynamic_tree_slot_storage to cover position_offsets and packed_mask for no-tree rows, verify the dummy-slot K-ary mask path, and add a case that calls mark_invalid() and checks the original state is restored. Use the existing DynamicTreeSlotStorage helper methods and slot_ids-based setup so the new assertions exercise the same no-tree and real-tree branches.Source: Path instructions
tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py (1)
72-77: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy liftAdd missing type annotations on the new methods.
Several new methods omit return annotations, and many helper parameters remain untyped. This makes the new dynamic-tree surface harder to validate with static tooling.
As per coding guidelines, “Always annotate functions. Make the return type
Noneif the function does not return anything.”Also applies to: 209-210, 240-241, 267-270, 272-286, 288-300, 302-331, 367-375, 426-432, 434-436, 490-497, 546-546, 637-658, 660-729, 735-745, 845-853, 958-972, 1163-1173, 1238-1245, 1263-1303
🤖 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/mtp_dynamic_tree.py` around lines 72 - 77, Add the missing type annotations throughout the new dynamic-tree surface in mtp_dynamic_tree.py: every new function/method should have explicit return types (use None for non-returning methods), and any currently untyped helper parameters should be annotated. Update the affected methods in the MTP dynamic tree implementation, including the __init__ constructor and the other new helpers listed in the review, so the full API is statically type-checkable and consistent with the coding guidelines.Source: Coding guidelines
tensorrt_llm/_torch/pyexecutor/model_engine.py (1)
1-15: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRemove the newly added header from this
pyexecutorPython file.This path follows a no-header convention for Python modules, so adding this block creates inconsistency.
Suggested diff
-# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License.Based on learnings, files under
tensorrt_llm/_torch/pyexecutor/should not add NVIDIA copyright headers.🤖 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/pyexecutor/model_engine.py` around lines 1 - 15, Remove the newly added license header from the pyexecutor Python module because this package follows the no-header convention for Python files. Update the top of model_engine.py to match other files under tensorrt_llm/_torch/pyexecutor/ and ensure no NVIDIA copyright block is present.Source: Learnings
🤖 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 `@cpp/tensorrt_llm/thop/attentionOp.cpp`:
- Line 1183: The `force_prepare_spec_dec_tree_mask` assignment in `AttentionOp`
is happening before the `op_cache` lookup, but `AttentionOp::data()` does not
include that field, so cached instances can reuse a stale value. Update the
cache key by adding `mForcePrepareSpecDecTreeMask` to `AttentionOp::data()` (and
any matching equality/hash logic used by the cache) so the pre-cache
`AttentionOp` is selected/initialized with the correct flag instead of
discarding it on cache hits.
In `@tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py`:
- Around line 797-820: The attention metadata mutation around
_forward_draft_loop in mtp_dynamic_tree should always be undone even when draft
execution fails. Wrap the call sequence that sets
force_prepare_spec_dec_tree_mask and updates all_rank_num_tokens in a
try/finally so _restore_attn_metadata_from_spec_dec and the original
attn_metadata fields are restored unconditionally. Keep the fix localized to the
speculative decoding path in the same method that prepares draft tree metadata.
---
Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/model_engine.py`:
- Around line 1-15: Remove the newly added license header from the pyexecutor
Python module because this package follows the no-header convention for Python
files. Update the top of model_engine.py to match other files under
tensorrt_llm/_torch/pyexecutor/ and ensure no NVIDIA copyright block is present.
In `@tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py`:
- Around line 72-77: Add the missing type annotations throughout the new
dynamic-tree surface in mtp_dynamic_tree.py: every new function/method should
have explicit return types (use None for non-returning methods), and any
currently untyped helper parameters should be annotated. Update the affected
methods in the MTP dynamic tree implementation, including the __init__
constructor and the other new helpers listed in the review, so the full API is
statically type-checkable and consistent with the coding guidelines.
In
`@tests/unittest/_torch/speculative/hw_agnostic/test_dynamic_tree_slot_storage.py`:
- Around line 23-77: The current tests in DynamicTreeSlotStorage only validate
next-link chaining and miss the new no-tree fallback and invalidation behavior.
Extend test_dynamic_tree_slot_storage to cover position_offsets and packed_mask
for no-tree rows, verify the dummy-slot K-ary mask path, and add a case that
calls mark_invalid() and checks the original state is restored. Use the existing
DynamicTreeSlotStorage helper methods and slot_ids-based setup so the new
assertions exercise the same no-tree and real-tree branches.
🪄 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: e7d46d6c-db64-470c-aa0a-ae60be69d872
📒 Files selected for processing (22)
cpp/tensorrt_llm/common/attentionOp.cppcpp/tensorrt_llm/common/attentionOp.hcpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/xqaParams.hcpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.hcpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaRunnerParams.hcpp/tensorrt_llm/kernels/xqaDispatcher.cppcpp/tensorrt_llm/nanobind/thop/bindings.cppcpp/tensorrt_llm/thop/attentionOp.cppcpp/tensorrt_llm/thop/attentionOp.hexamples/llm-api/quickstart_advanced.pytensorrt_llm/_torch/attention_backend/fmha/fallback.pytensorrt_llm/_torch/attention_backend/trtllm.pytensorrt_llm/_torch/models/modeling_nemotron_h.pytensorrt_llm/_torch/modules/mamba/mamba2_mixer.pytensorrt_llm/_torch/pyexecutor/mamba_cache_manager.pytensorrt_llm/_torch/pyexecutor/model_engine.pytensorrt_llm/_torch/speculative/mtp.pytensorrt_llm/_torch/speculative/mtp_dynamic_tree.pytensorrt_llm/_torch/speculative/spec_tree_manager.pytensorrt_llm/_torch/speculative/utils.pytensorrt_llm/llmapi/llm_args.pytests/unittest/_torch/speculative/hw_agnostic/test_dynamic_tree_slot_storage.py
sunnyqgg
changed the title
|
PR_Github #55465 [ run ] completed with state
|
b9ba07f to
40043da
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #56326 [ run ] triggered by Bot. Commit: |
|
PR_Github #56326 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #57098 [ run ] triggered by Bot. Commit: |
|
PR_Github #57098 [ run ] completed with state
|
Signed-off-by: qgai <qgai@nvidia.com>
Signed-off-by: qgai <qgai@nvidia.com>
Signed-off-by: qgai <qgai@nvidia.com>
Signed-off-by: qgai <qgai@nvidia.com>
Signed-off-by: qgai <qgai@nvidia.com>
98a7f63 to
024d237
Compare
|
/bot run --disable-fail-fast |
1 similar comment
|
/bot run --disable-fail-fast |
|
PR_Github #60163 [ run ] triggered by Bot. Commit: |
|
PR_Github #60163 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
Signed-off-by: qgai <qgai@nvidia.com>
|
PR_Github #60210 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #60262 [ run ] triggered by Bot. Commit: |
|
PR_Github #60210 [ run ] completed with state |
|
PR_Github #60262 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #60332 [ run ] triggered by Bot. Commit: |
|
PR_Github #60332 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #60849 [ run ] triggered by Bot. Commit: |
|
PR_Github #60849 [ run ] completed with state |
|
Heads-up: this PR appears to have merged with a semantic conflict against a guard that landed on main after this branch's CI ran, and current main tip is import-broken as a result.
Repro: Effect: every PR CI run (which tests the PR merged with latest main) currently fails at test collection for anything importing the PyTorch backend. The likely fix is renaming the override to |
…d conflict (drop after upstream fix) MTPEagleDynamicTreeWorker overrides SpecWorkerBase.forward, which the __init_subclass__ guard from NVIDIA#16382 rejects at class-definition time, breaking all backend imports on current main (see NVIDIA#15582 (comment)). Rename the override to _forward_impl per the guard's contract; SpecWorkerBase.forward forwards *args/**kwargs unchanged. Temporary carry so this PR's CI can run; becomes a no-op once the equivalent fix lands on main. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
Signed-off-by: qgai <qgai@nvidia.com>
9 participants
Summary
Changes
tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py: add dynamic-tree MTP state and mask management.tensorrt_llm/_torch/speculative/mtp.py,spec_tree_manager.py, andutils.py: integrate dynamic-tree MTP flow.tensorrt_llm/_torch/models/modeling_nemotron_h.py, Mamba cache/model engine, and attention backend files: propagate dynamic-tree metadata.tests/unittest/_torch/speculative/hw_agnostic/test_dynamic_tree_slot_storage.py: add slot storage regression coverage.Test plan
pytest tests/unittest/_torch/speculative/hw_agnostic/test_dynamic_tree_slot_storage.pySummary by CodeRabbit
New Features
Bug Fixes