[https://nvbugs/6478723][fix] Rename forward -> _forward_impl AND stash the extra dynamic-tree state on…

[trtllm-agent]

Summary

• Root cause: MTPEagleDynamicTreeWorker overrode SpecWorkerBase.forward, which init_subclass forbids (nvbugs/6442074) — subclasses must implement _forward_impl so the base-class forward can guarantee spec-dec attn-metadata cleanup on failure. Class-creation TypeError blocked pytest collection.
• Fix: Rename forward -> _forward_impl AND stash the extra dynamic-tree state on self, then override _ensure_spec_dec_state_restored to restore it — matching the PARD/DFlash pattern. This addresses both the immediate collection failure and closes the state-leak gap the base class explicitly documents for subclasses with extra transient state.
• 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/6478723

Dev Engineer Review

• Renamed forward to _forward_impl to comply with SpecWorkerBase.__init_subclass__, allowing centralized failure cleanup.
• Preserved the existing profiler label and implementation signature.
• Added restoration of all_rank_num_tokens and force_prepare_spec_dec_tree_mask after tolerated failures, preventing dynamic-tree state leakage.
• Clears the saved snapshot after restoration.
• No configuration or test-list changes.

QA Engineer Review

No test changes.

[coderabbitai]

Walkthrough

MTPEagleDynamicTreeWorker now snapshots additional speculative decoding attention state before dynamic-tree execution and restores it during cleanup. Its execution method is renamed from forward to _forward_impl.

Changes

Dynamic tree state handling

Layer / File(s)	Summary
State capture and execution entry point tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py	Adds saved extra state for attention metadata, renames the worker entry point to _forward_impl, and snapshots token counts and the tree mask before draft-tree processing.
Cleanup state restoration tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py	Restores the saved metadata through _ensure_spec_dec_state_restored and the draft-loop cleanup path, then clears the snapshot.

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

Possibly related PRs

• NVIDIA/TensorRT-LLM#16724: Renames the same worker method from forward to _forward_impl.
• NVIDIA/TensorRT-LLM#16727: Shares the method rename and related saved/restored speculative decoding state changes.
• NVIDIA/TensorRT-LLM#16731: Extends the same _forward_impl rename with attention and speculative decoding state restoration.

Suggested reviewers: zhaoyangwang-nvidia

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The description is useful, but it does not follow the required template sections for Description, Test Coverage, and PR Checklist.	Add the missing template sections: @coderabbitai summary, Description, Test Coverage, and PR Checklist.
Docstring Coverage	⚠️ Warning	Docstring coverage is 25.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is clearly about the same fix: renaming forward and restoring extra dynamic-tree state.
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/mtp_dynamic_tree.py (1)

238-250: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Add type annotations to these overrides. tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py#L238-L250 should annotate both metadata parameters and return None; tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py#L661-L671 should annotate its parameters and return type consistently with the worker contract.

🤖 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 238 - 250,
Add type annotations to the _ensure_spec_dec_state_restored override parameters
using the appropriate attention and speculative metadata types, and annotate its
return as None. Also update the override at
tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py lines 661-671 with parameter
and return annotations consistent with the worker contract.

Source: Coding guidelines

🤖 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/mtp_dynamic_tree.py`:
- Around line 238-250: Add type annotations to the
_ensure_spec_dec_state_restored override parameters using the appropriate
attention and speculative metadata types, and annotate its return as None. Also
update the override at tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py lines
661-671 with parameter and return annotations consistent with the worker
contract.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 1ce8df88-e0ff-4a29-9d1e-ea3a3c7469a9

📥 Commits

Reviewing files that changed from the base of the PR and between 9cd6a49 and 2872dcb.

📒 Files selected for processing (1)

• tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py