[TRTLLM-14388][refactor] Remove 2 model spec dec drafting loops - #17562
Conversation
mikeiovine
added
the
api-compatible
|
/bot run |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughChangesThe PR removes static drafting-loop and chain-drafter execution paths. It updates speculative-decoding APIs, adds occurrence-penalty handling, changes graph and cache sizing, centralizes request-state resets, and updates related tests and configuration. Speculative decoding simplification
Estimated code review effort: 4 (Complex) | ~45 minutes Merge Risk: 🟠 High · up to The current head removes obsolete drafting-loop code, but adjacent speculative-decoding and replay changes still have unresolved correctness issues: requested penalties can be skipped, and prefill replay can duplicate returned context logits; required configuration metadata is also incomplete. These can change generated results or returned outputs, so the PR is not merge-ready until the correctness and metadata issues are addressed or explicitly accepted. Sequence Diagram(s)sequenceDiagram
participant PyExecutor
participant PyTorchModelEngine
participant ModelDrafter
participant DraftModel
participant SpecMetadata
participant Sampler
PyExecutor->>PyTorchModelEngine: initialize target and draft engines
PyExecutor->>ModelDrafter: prepare dynamic draft requests
ModelDrafter->>DraftModel: run dynamic draft passes
DraftModel->>ModelDrafter: return draft outputs and d2t
ModelDrafter->>Sampler: generate draft probabilities
SpecMetadata->>Sampler: apply occurrence penalties
PyTorchModelEngine->>PyExecutor: capture graphs and allocate warmup requests
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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)
tensorrt_llm/_torch/pyexecutor/py_executor_creator.py (1)
626-631: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRemove the now-unreachable overlap-scheduler warning.
Lines 626-631 set
llm_args.disable_overlap_scheduler = Truefor every value ofhas_draft_model_enginethat is true. The later block at lines 650-653 requireshas_draft_model_engine and not llm_args.disable_overlap_scheduler, so it can never run. Delete it to avoid a misleading log path.♻️ Proposed removal of the dead warning block
- if has_draft_model_engine and not llm_args.disable_overlap_scheduler: - logger.warning( - "Overlap scheduler is enabled for two-model speculative decoding. Rejection sampling will fallback to greedy sampling." - ) - max_seq_len = model_engine_max_seq_lenAlso applies to: 650-653
🤖 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/py_executor_creator.py` around lines 626 - 631, Remove the has_draft_model_engine block that logs the overlap-scheduler warning and sets llm_args.disable_overlap_scheduler in the surrounding executor-creation flow. Also remove the now-unreachable later branch guarded by has_draft_model_engine and not llm_args.disable_overlap_scheduler, while preserving all other scheduler logic.tests/unittest/_torch/speculative/test_eagle3.py (2)
477-496: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winPrune the parameterization rows that always skip.
Line 504 skips unconditionally when
use_one_modelisFalse. Most rows in this list set the fifth value toFalse, so those runs only ever reportSKIPPED. That hides the real covered matrix and costs collection time on every CI run. Since this pull request removes the two-model path, drop theuse_one_modelcolumn and keep only the one-model rows.🤖 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/test_eagle3.py` around lines 477 - 496, Update the parameterization list for the speculative model test by removing the obsolete use_one_model column and deleting rows that represent the removed two-model path. Retain only one-model cases, adjust each remaining row to the new parameter shape, and preserve the mocked HuggingFace auto-download case.
504-505: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePrefer a module-level or collection-time skip over a body skip.
The body skip runs after fixtures and after
torch.cudaaccess in later lines is set up. If the two-model path stays deprecated, express the restriction in the parameterization instead, so the skipped cases never reach the test body.🤖 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/test_eagle3.py` around lines 504 - 505, Move the deprecation skip for the two-model Eagle3 cases out of the test body and into the test parameterization or collection-time configuration. Update the parameterization associated with use_one_model so deprecated two-model cases are marked skipped before fixtures and torch.cuda setup run, then remove the body-level pytest.skip.
🤖 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 `@tensorrt_llm/_torch/attention_backend/trtllm.py`:
- Around line 1194-1210: Update the static-tree branch in update_spec_dec_param
to handle spec_metadata.is_draft_model before the target-only assertion. Add the
explicit draft-model metadata path required for draft forwards, then retain the
existing target-model spec_tree_manager updates and assertion for non-draft
metadata.
In `@tensorrt_llm/_torch/pyexecutor/model_engine.py`:
- Around line 1971-1973: Update the is_draft_model branch to derive draft_len
from the runtime draft width used by ModelDrafter, rather than
self.max_total_draft_tokens. Preserve the existing cuda_graph_batch_sizes
mapping while ensuring draft requests without py_draft_tokens produce the
zero-width graph key required for replay.
---
Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/py_executor_creator.py`:
- Around line 626-631: Remove the has_draft_model_engine block that logs the
overlap-scheduler warning and sets llm_args.disable_overlap_scheduler in the
surrounding executor-creation flow. Also remove the now-unreachable later branch
guarded by has_draft_model_engine and not llm_args.disable_overlap_scheduler,
while preserving all other scheduler logic.
In `@tests/unittest/_torch/speculative/test_eagle3.py`:
- Around line 477-496: Update the parameterization list for the speculative
model test by removing the obsolete use_one_model column and deleting rows that
represent the removed two-model path. Retain only one-model cases, adjust each
remaining row to the new parameter shape, and preserve the mocked HuggingFace
auto-download case.
- Around line 504-505: Move the deprecation skip for the two-model Eagle3 cases
out of the test body and into the test parameterization or collection-time
configuration. Update the parameterization associated with use_one_model so
deprecated two-model cases are marked skipped before fixtures and torch.cuda
setup run, then remove the body-level pytest.skip.
🪄 Autofix
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: ca308652-32bc-4416-bde4-e11211178b17
📒 Files selected for processing (25)
.pre-commit-config.yamllegacy-files.txtpyproject.tomlruff-legacy.tomltensorrt_llm/_torch/attention_backend/interface.pytensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.pytensorrt_llm/_torch/attention_backend/sparse/rocket/cache_manager.pytensorrt_llm/_torch/attention_backend/trtllm.pytensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.pytensorrt_llm/_torch/pyexecutor/llm_request.pytensorrt_llm/_torch/pyexecutor/mamba_cache_manager.pytensorrt_llm/_torch/pyexecutor/model_engine.pytensorrt_llm/_torch/pyexecutor/py_executor_creator.pytensorrt_llm/_torch/pyexecutor/resource_manager.pytensorrt_llm/_torch/pyexecutor/sampler/sampler.pytensorrt_llm/_torch/speculative/drafting_loops.pytensorrt_llm/_torch/speculative/interface.pytensorrt_llm/_torch/speculative/model_drafter.pytensorrt_llm/llmapi/llm_args.pytests/unittest/_torch/executor/test_pytorch_model_engine.pytests/unittest/_torch/modeling/test_modeling_llama.pytests/unittest/_torch/speculative/hw_agnostic/test_draft_token_prepare_for_generation.pytests/unittest/_torch/speculative/hw_agnostic/test_draft_token_tree_sampling.pytests/unittest/_torch/speculative/test_eagle3.pytests/unittest/llmapi/test_llm_args.py
💤 Files with no reviewable changes (16)
- ruff-legacy.toml
- tensorrt_llm/llmapi/llm_args.py
- tests/unittest/_torch/speculative/hw_agnostic/test_draft_token_prepare_for_generation.py
- legacy-files.txt
- tensorrt_llm/_torch/attention_backend/sparse/rocket/cache_manager.py
- tensorrt_llm/_torch/speculative/drafting_loops.py
- tests/unittest/_torch/executor/test_pytorch_model_engine.py
- .pre-commit-config.yaml
- tensorrt_llm/_torch/pyexecutor/resource_manager.py
- tests/unittest/_torch/speculative/hw_agnostic/test_draft_token_tree_sampling.py
- tests/unittest/_torch/modeling/test_modeling_llama.py
- tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py
- tensorrt_llm/_torch/pyexecutor/llm_request.py
- tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py
- pyproject.toml
- tensorrt_llm/_torch/attention_backend/interface.py
|
PR_Github #65609 [ run ] triggered by Bot. Commit: |
|
PR_Github #67966 [ run ] triggered by Bot. Commit: |
|
PR_Github #67966 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #68012 [ run ] triggered by Bot. Commit: |
|
PR_Github #68012 [ run ] completed with state
|
|
/bot run |
|
PR_Github #68074 [ run ] triggered by Bot. Commit: |
|
PR_Github #68074 [ run ] completed with state
|
|
/bot run |
|
PR_Github #68372 [ run ] triggered by Bot. Commit: |
|
PR_Github #68372 [ run ] completed with state
|
|
/bot run |
|
PR_Github #68414 [ run ] triggered by Bot. Commit: |
|
PR_Github #68414 [ run ] completed with state
|
|
/bot run |
|
PR_Github #68428 [ run ] triggered by Bot. Commit: |
|
PR_Github #68428 [ run ] completed with state
|
Signed-off-by: allisonlim-nv <allim@nvidia.com>
Signed-off-by: Allison Lim <allim@nvidia.com>
92ee5e6 to
2fd438d
Compare
|
/bot run |
|
PR_Github #68852 [ run ] triggered by Bot. Commit: |
|
PR_Github #68852 [ run ] completed with state
|
|
/bot run |
|
PR_Github #68923 [ run ] triggered by Bot. Commit: |
|
PR_Github #68923 [ run ] completed with state |
10 participants
Description
2 model based spec dec is now disabled after #17366.
This PR removes some dead code. Drafting loops were a component used by 2 model eagle/MTP only. No usable codepath is affected by this change.
Test Coverage
Existing tests.
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-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin 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.Dev Engineer Review
QA Engineer Review
Modified or removed test coverage:
test_draft_token_static_tree_prepare_for_generation.test_draft_token_static_tree_sampling.test_llama_eagle3.test_eagle3_cdl_sampling.num_extra_decoding_stepscoverage fromtests/unittest/llmapi/test_llm_args.py.tests/unittest/_torch/modeling/test_modeling_llama.py.tests/unittest/_torch/executor/test_pytorch_model_engine.py.test_llama_eagle3remains covered bytests/integration/test_lists/test-db/l0_h100.ymlandl0_b200.yml. No matching test-list entries were identified for the other modified or removed unit tests.Verdict: needs follow-up.