[None][feat] Draft Target One Model#10671
Conversation
|
Depends on #10502 |
📝 WalkthroughWalkthroughThis pull request adds support for a new speculative decoding mode called "DraftTarget One Model," enabling speculative decoding with a single model engine. The implementation includes configuration options, new worker and sampler classes, interface updates, utility routing, and model integration to orchestrate the complete workflow. Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant DraftTargetOneModelWorker
participant DraftModel
participant TargetModel
participant Sampler
Client->>DraftTargetOneModelWorker: forward(input_ids, logits, ...)
DraftTargetOneModelWorker->>Sampler: sample_and_accept_draft_tokens()
Sampler->>DraftTargetOneModelWorker: accepted_tokens, num_accepted
loop for each draft step (max_draft_len)
DraftTargetOneModelWorker->>DraftTargetOneModelWorker: prepare_drafter_inputs()
DraftTargetOneModelWorker->>DraftModel: forward(draft_input_ids, ...)
DraftModel->>DraftTargetOneModelWorker: draft_logits
DraftTargetOneModelWorker->>DraftTargetOneModelWorker: draft_decoder(draft_logits)
DraftTargetOneModelWorker->>DraftTargetOneModelWorker: update_attn_metadata()
end
DraftTargetOneModelWorker->>Client: return next_tokens, next_lens, overlapped_inputs
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes 🚥 Pre-merge checks | ✅ 1 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (1 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 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: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tensorrt_llm/llmapi/llm_args.py (1)
1074-1105: Update validation forDraftTargetDecodingConfigone-model mode inTorchLlmArgs.validate_speculative_config.The validation at line 3054 unconditionally requires
speculative_model_dirfor allDraftTargetDecodingConfiginstances. However, for one-model mode (draft_target_one_model=True), where draft and target models share the same engine,speculative_model_dirmay be optional ifnum_draft_layersis explicitly provided. The runtime code inmodeling_speculative.pytreatsspeculative_model_diras optional for one-model mode (conditional check at line 920). Update the validation logic to allowspeculative_model_dirto beNonewhendraft_target_one_model=Trueandnum_draft_layersis explicitly set.
🤖 Fix all issues with AI agents
In `@tensorrt_llm/_torch/models/modeling_speculative.py`:
- Around line 918-931: When
spec_config.spec_dec_mode.is_draft_target_one_model() is true ensure
speculative_model_dir is validated before attempting to load the draft config:
check spec_config.speculative_model_dir and raise an explicit error or assertion
if missing, so self.draft_config is never left None; perform this validation
right before calling ModelConfig.from_pretrained(...) (the block that assigns
self.draft_config) and include a clear message referencing
speculative_model_dir, and ensure get_draft_model (where
AutoModelForCausalLM.from_config(draft_config) is used) can assume draft_config
is present.
- Around line 868-869: When spec_dec_mode.is_draft_target_one_model() is true,
ensure spec_config.speculative_model_dir is provided before calling
AutoModelForCausalLM.from_config(draft_config); validate that draft_config (or
self.draft_config) is non-None and raise a clear ValueError (or similar) if
spec_config.speculative_model_dir is missing, or move/assign draft_config
earlier when spec_config.speculative_model_dir is truthy so
AutoModelForCausalLM.from_config(draft_config) never receives None; reference
spec_dec_mode.is_draft_target_one_model(),
AutoModelForCausalLM.from_config(draft_config), and
spec_config.speculative_model_dir to locate the change.
In `@tensorrt_llm/_torch/speculative/draft_target.py`:
- Around line 1-7: Add the required NVIDIA copyright header (including the year
of latest meaningful modification) to the top of the module to comply with
project guidelines; insert the standard multi-line header block before the
module docstring in tensorrt_llm._torch.speculative.draft_target (i.e., at the
very top of this file) so the header precedes the existing module docstring and
includes the proper copyright line and any required SPDX or license identifier.
In `@tensorrt_llm/llmapi/llm_args.py`:
- Around line 1075-1083: Remove the duplicate allow_advanced_sampling field from
DraftTargetDecodingConfig since it is already declared in the parent
DecodingBaseConfig; locate the allow_advanced_sampling attribute declaration
inside the DraftTargetDecodingConfig block (near draft_target_one_model and
num_draft_layers) and delete that line so the class inherits the field from
DecodingBaseConfig without shadowing it.
🧹 Nitpick comments (2)
tensorrt_llm/_torch/speculative/draft_target.py (2)
109-118: Unusedhidden_statesparameter.The
hidden_statesparameter is declared but never used in the method body. According to the docstring, it represents "Hidden states from the target model," but DraftTarget doesn't require hidden states from the target model (as stated in the class docstring). Consider prefixing with underscore to indicate intentional non-use, or remove if the interface allows.Suggested fix
def forward( self, input_ids, position_ids, - hidden_states, + hidden_states, # noqa: ARG002 - unused, kept for interface compatibility logits, attn_metadata, spec_metadata, draft_model, ):
344-352: Consider documenting or removing unused parameters.The
num_accepted_tokensanddraft_modelparameters are declared but never used in the method body. If these are for future use or interface compatibility with other workers, consider adding a brief comment. Otherwise, they could be removed to reduce confusion.
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (6)
tensorrt_llm/_torch/models/modeling_speculative.pytensorrt_llm/_torch/speculative/__init__.pytensorrt_llm/_torch/speculative/draft_target.pytensorrt_llm/_torch/speculative/interface.pytensorrt_llm/_torch/speculative/utils.pytensorrt_llm/llmapi/llm_args.py
🧰 Additional context used
📓 Path-based instructions (2)
**/*.py
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
**/*.py: The code developed for TensorRT-LLM should conform to Python 3.8+
Indent Python code with 4 spaces. Do not use tabs
Always maintain the namespace when importing Python modules, even if only one class or function from a module is used
Python filenames should use snake_case (e.g.,some_file.py)
Python classes should use PascalCase (e.g.,class SomeClass)
Python functions and methods should use snake_case (e.g.,def my_awesome_function():)
Python local variables should use snake_case, with prefixkfor variable names that start with a number (e.g.,k_99th_percentile)
Python global variables should use upper snake_case with prefixG(e.g.,G_MY_GLOBAL)
Python constants should use upper snake_case (e.g.,MY_CONSTANT)
Avoid shadowing variables declared in an outer scope in Python
Initialize all externally visible members of a Python class in the constructor
For Python interfaces that may be used outside a file, prefer docstrings over comments
Use comments in Python for code within a function, or interfaces that are local to a file
Use Google-style docstrings for Python classes and functions, which can be parsed by Sphinx
Python attributes and variables can be documented inline with the format"""<type>: Description"""
Avoid using reflection in Python when functionality can be easily achieved without reflection
When using try-except blocks in Python, limit the except clause to the smallest set of errors possible
When using try-except blocks in Python to handle multiple possible variable types (duck-typing), keep the body of the try as small as possible and use the else block for the main logic
Files:
tensorrt_llm/_torch/speculative/__init__.pytensorrt_llm/_torch/models/modeling_speculative.pytensorrt_llm/_torch/speculative/interface.pytensorrt_llm/_torch/speculative/utils.pytensorrt_llm/_torch/speculative/draft_target.pytensorrt_llm/llmapi/llm_args.py
**/*.{cpp,cc,cxx,h,hpp,hxx,cu,cuh,py}
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
All TensorRT-LLM source files (.cpp, .h, .cu, .py, and other source files) should contain an NVIDIA copyright header with the year of latest meaningful modification
Files:
tensorrt_llm/_torch/speculative/__init__.pytensorrt_llm/_torch/models/modeling_speculative.pytensorrt_llm/_torch/speculative/interface.pytensorrt_llm/_torch/speculative/utils.pytensorrt_llm/_torch/speculative/draft_target.pytensorrt_llm/llmapi/llm_args.py
🧠 Learnings (4)
📓 Common learnings
Learnt from: venkywonka
Repo: NVIDIA/TensorRT-LLM PR: 6029
File: .github/pull_request_template.md:45-53
Timestamp: 2025-08-27T17:50:13.264Z
Learning: For PR templates in TensorRT-LLM, avoid suggesting changes that would increase developer overhead, such as converting plain bullets to mandatory checkboxes. The team prefers guidance-style bullets that don't require explicit interaction to reduce friction in the PR creation process.
📚 Learning: 2025-08-14T15:38:01.771Z
Learnt from: MatthiasKohl
Repo: NVIDIA/TensorRT-LLM PR: 6904
File: cpp/tensorrt_llm/pybind/thop/bindings.cpp:55-57
Timestamp: 2025-08-14T15:38:01.771Z
Learning: In TensorRT-LLM Python bindings, tensor parameter collections like mla_tensor_params and spec_decoding_tensor_params are kept as required parameters without defaults to maintain API consistency, even when it might affect backward compatibility.
Applied to files:
tensorrt_llm/llmapi/llm_args.py
📚 Learning: 2025-08-26T09:37:10.463Z
Learnt from: jiaganc
Repo: NVIDIA/TensorRT-LLM PR: 7031
File: tensorrt_llm/bench/dataclasses/configuration.py:90-104
Timestamp: 2025-08-26T09:37:10.463Z
Learning: In TensorRT-LLM, the `get_pytorch_perf_config()` method returns `self.pytorch_config` which can contain default `cuda_graph_config` values, so `llm_args` may already have this config before the extra options processing.
Applied to files:
tensorrt_llm/llmapi/llm_args.py
📚 Learning: 2026-01-06T03:07:15.754Z
Learnt from: CR
Repo: NVIDIA/TensorRT-LLM PR: 0
File: CODING_GUIDELINES.md:0-0
Timestamp: 2026-01-06T03:07:15.754Z
Learning: Applies to **/*.{md,rst} : When documenting CLI commands for TensorRT-LLM tools like `trtllm-serve`, `trtllm-bench`, or `trtllm-eval`, prefer using `--config` over `--extra_llm_api_options` for specifying configuration files
Applied to files:
tensorrt_llm/llmapi/llm_args.py
🧬 Code graph analysis (2)
tensorrt_llm/_torch/speculative/__init__.py (1)
tensorrt_llm/_torch/speculative/draft_target.py (2)
DraftTargetOneModelSpecMetadata(27-73)DraftTargetOneModelWorker(87-394)
tensorrt_llm/llmapi/llm_args.py (2)
tensorrt_llm/_torch/speculative/interface.py (1)
SpeculativeDecodingMode(41-189)tensorrt_llm/models/modeling_utils.py (1)
SpeculativeDecodingMode(91-127)
🪛 Ruff (0.14.11)
tensorrt_llm/_torch/speculative/draft_target.py
55-55: Unused method argument: layer_id
(ARG002)
113-113: Unused method argument: hidden_states
(ARG002)
349-349: Unused method argument: num_accepted_tokens
(ARG002)
352-352: Unused method argument: draft_model
(ARG002)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Pre-commit Check
🔇 Additional comments (21)
tensorrt_llm/_torch/speculative/draft_target.py (5)
26-73: LGTM!The
DraftTargetOneModelSpecMetadataclass correctly implements theSpecMetadatainterface. The unusedlayer_idparameter inis_layer_captureis required by the interface contract, and the no-op implementations for hidden state capture are intentional as documented in the docstrings.
76-84: LGTM!Clean delegation to
MTPSamplerwith appropriate configuration.
87-107: LGTM!The worker initialization and property are implemented correctly.
269-325: LGTM!The token sampling and draft verification logic is well-implemented. The use of
torch.cumprodto count consecutive matching tokens is an efficient pattern for speculative decoding verification.
327-342: LGTM!Greedy decoding for draft tokens is appropriate as noted in the docstring.
tensorrt_llm/_torch/speculative/__init__.py (2)
2-3: LGTM!The new exports are properly added following the existing import patterns.
16-17: LGTM!The
__all__list is correctly updated with the new public entities.tensorrt_llm/_torch/speculative/interface.py (4)
49-49: LGTM!New enum member
DRAFT_TARGET_ONE_MODELis correctly added.
70-72: LGTM!
use_one_engine()correctly includes the newis_draft_target_one_model()check.
89-91: LGTM!New predicate
is_draft_target_one_model()follows the established pattern.
95-134: LGTM!All predicate methods are consistently updated to include
is_draft_target_one_model()where appropriate, following the same patterns asis_eagle3_one_model().tensorrt_llm/_torch/models/modeling_speculative.py (3)
23-23: LGTM!Import addition for
AutoModelForCausalLMis appropriate for the new draft model instantiation path.
940-946: LGTM!The epilogue extra_attrs merging logic is correctly extended to handle
is_draft_target_one_model()alongside the existing Eagle3 path.
1028-1031: LGTM!Correctly skips
load_weights_from_target_modelfor DraftTarget one-model since the draft model operates independently with its own embeddings and weights.tensorrt_llm/_torch/speculative/utils.py (5)
10-12: LGTM!Import of new DraftTarget classes follows the existing import pattern.
84-92: LGTM!
get_spec_metadatacorrectly handles the newis_draft_target_one_model()mode with appropriate parameters.
187-188: LGTM!
get_spec_decodercorrectly returnsDraftTargetOneModelSamplerfor the new mode.
233-237: LGTM!
get_num_spec_layerscorrectly handlesis_draft_target_one_model()with thenum_draft_layersconfig field and appropriate fallback.
249-250: LGTM!
get_spec_workercorrectly returnsDraftTargetOneModelWorkerfor the new mode.tensorrt_llm/llmapi/llm_args.py (2)
809-812: LGTM!Good addition of a default implementation returning 0. This follows the established pattern where subclasses like
EagleDecodingConfigandMTPDecodingConfigoverride this property when they need layer capture.
1098-1105: LGTM — The logic correctly returnsDRAFT_TARGET_ONE_MODELwhendraft_target_one_modelis enabled, otherwise falls back toDRAFT_TARGET. This follows the same pattern asEagleDecodingConfig.spec_dec_modeandMTPDecodingConfig.spec_dec_mode. The enum value is properly defined inSpeculativeDecodingMode.
✏️ Tip: You can disable this entire section by setting review_details to false in your review settings.
| elif spec_dec_mode.is_draft_target_one_model(): | ||
| return AutoModelForCausalLM.from_config(draft_config) |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# First, let's examine the file around lines 868-869 and 920-931
wc -l tensorrt_llm/_torch/models/modeling_speculative.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 118
🏁 Script executed:
#!/bin/bash
# Read the relevant sections of the file
echo "=== Lines 860-880 (context around 868-869) ==="
sed -n '860,880p' tensorrt_llm/_torch/models/modeling_speculative.py
echo -e "\n=== Lines 910-940 (context around 920-931) ==="
sed -n '910,940p' tensorrt_llm/_torch/models/modeling_speculative.py
echo -e "\n=== Lines 1-50 (to understand class structure and imports) ==="
head -50 tensorrt_llm/_torch/models/modeling_speculative.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 5025
🏁 Script executed:
#!/bin/bash
# Search for DRAFT_TARGET_ONE_MODEL and validation around speculative_model_dir
rg -n "DRAFT_TARGET_ONE_MODEL|speculative_model_dir" tensorrt_llm/_torch/models/modeling_speculative.py -B 2 -A 2Repository: NVIDIA/TensorRT-LLM
Length of output: 1178
🏁 Script executed:
#!/bin/bash
# Look for the SpecDecOneEngineForCausalLM class definition and __init__ method
rg -n "class SpecDecOneEngineForCausalLM" tensorrt_llm/_torch/models/modeling_speculative.py -A 30Repository: NVIDIA/TensorRT-LLM
Length of output: 2046
🏁 Script executed:
#!/bin/bash
# Search for validation or assertions that check speculative_model_dir
rg -n "assert|raise|if.*speculative_model_dir" tensorrt_llm/_torch/models/modeling_speculative.py -B 1 -A 1Repository: NVIDIA/TensorRT-LLM
Length of output: 2371
🏁 Script executed:
#!/bin/bash
# Search for spec_config class definition and validation
rg -n "class.*SpecConfig|class.*SpeculativeConfig" --type py -A 10 | head -50Repository: NVIDIA/TensorRT-LLM
Length of output: 45
🏁 Script executed:
#!/bin/bash
# Look for any validation of spec_config or speculative_model_dir
rg -n "def.*validate|spec_config.*validate|speculative_model_dir.*validate" --type py -B 2 -A 5 | head -80Repository: NVIDIA/TensorRT-LLM
Length of output: 6767
🏁 Script executed:
#!/bin/bash
# Check if speculative_model_dir is required when draft_target_one_model is set
rg -n "DRAFT_TARGET_ONE_MODEL|draft_target_one_model" --type py -B 5 -A 5 | head -100Repository: NVIDIA/TensorRT-LLM
Length of output: 8307
🏁 Script executed:
#!/bin/bash
# Look for where the spec_config is created/validated in the initialization chain
rg -n "spec_config.*=|SpecConfig\(" tensorrt_llm/_torch/models/modeling_speculative.py | head -30Repository: NVIDIA/TensorRT-LLM
Length of output: 908
Validate speculative_model_dir is provided when using DRAFT_TARGET_ONE_MODEL mode.
AutoModelForCausalLM.from_config(draft_config) will fail with None if speculative_model_dir is not provided. The code at lines 920-921 only sets self.draft_config when spec_config.speculative_model_dir is truthy; otherwise, draft_config remains None and the call at line 868 will fail. Add validation to ensure speculative_model_dir is required when DRAFT_TARGET_ONE_MODEL mode is used.
🤖 Prompt for AI Agents
In `@tensorrt_llm/_torch/models/modeling_speculative.py` around lines 868 - 869,
When spec_dec_mode.is_draft_target_one_model() is true, ensure
spec_config.speculative_model_dir is provided before calling
AutoModelForCausalLM.from_config(draft_config); validate that draft_config (or
self.draft_config) is non-None and raise a clear ValueError (or similar) if
spec_config.speculative_model_dir is missing, or move/assign draft_config
earlier when spec_config.speculative_model_dir is truthy so
AutoModelForCausalLM.from_config(draft_config) never receives None; reference
spec_dec_mode.is_draft_target_one_model(),
AutoModelForCausalLM.from_config(draft_config), and
spec_config.speculative_model_dir to locate the change.
| elif spec_config.spec_dec_mode.is_draft_target_one_model(): | ||
| # Load the draft model config for DraftTarget one-model | ||
| if spec_config.speculative_model_dir: | ||
| self.draft_config = ModelConfig.from_pretrained( | ||
| spec_config.speculative_model_dir, | ||
| trust_remote_code=True, | ||
| attn_backend=model_config.attn_backend, | ||
| moe_backend=model_config.moe_backend, | ||
| mapping=model_config.mapping, | ||
| spec_config=None, # Draft model doesn't need spec_config | ||
| max_num_tokens=model_config.max_num_tokens, | ||
| moe_max_num_tokens=model_config.moe_max_num_tokens) | ||
| self.draft_config.quant_config.kv_cache_quant_algo = \ | ||
| model_config.quant_config.kv_cache_quant_algo |
There was a problem hiding this comment.
Consider adding validation for required speculative_model_dir.
For is_draft_target_one_model(), the draft model is loaded from speculative_model_dir. If this directory is not provided, self.draft_config remains None, which will cause get_draft_model to fail when calling AutoModelForCausalLM.from_config(draft_config). Consider adding an explicit check or assertion.
Suggested validation
elif spec_config.spec_dec_mode.is_draft_target_one_model():
# Load the draft model config for DraftTarget one-model
- if spec_config.speculative_model_dir:
+ if not spec_config.speculative_model_dir:
+ raise ValueError(
+ "speculative_model_dir is required for DraftTarget one-model mode"
+ )
+ else:
self.draft_config = ModelConfig.from_pretrained(🤖 Prompt for AI Agents
In `@tensorrt_llm/_torch/models/modeling_speculative.py` around lines 918 - 931,
When spec_config.spec_dec_mode.is_draft_target_one_model() is true ensure
speculative_model_dir is validated before attempting to load the draft config:
check spec_config.speculative_model_dir and raise an explicit error or assertion
if missing, so self.draft_config is never left None; perform this validation
right before calling ModelConfig.from_pretrained(...) (the block that assigns
self.draft_config) and include a clear message referencing
speculative_model_dir, and ensure get_draft_model (where
AutoModelForCausalLM.from_config(draft_config) is used) can assume draft_config
is present.
| """ | ||
| DraftTarget One-Model Speculative Decoding Implementation. | ||
|
|
||
| This module implements a one-model approach for DraftTarget speculative decoding, | ||
| where the draft and target models share the same model engine. The draft model | ||
| layers are integrated into the target model's KV cache and run in a single forward pass. | ||
| """ |
There was a problem hiding this comment.
Missing NVIDIA copyright header.
As per the coding guidelines, all TensorRT-LLM source files should contain an NVIDIA copyright header with the year of latest meaningful modification.
Suggested header
+# SPDX-FileCopyrightText: Copyright (c) 2022-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.
+
"""
DraftTarget One-Model Speculative Decoding Implementation.📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| """ | |
| DraftTarget One-Model Speculative Decoding Implementation. | |
| This module implements a one-model approach for DraftTarget speculative decoding, | |
| where the draft and target models share the same model engine. The draft model | |
| layers are integrated into the target model's KV cache and run in a single forward pass. | |
| """ | |
| # SPDX-FileCopyrightText: Copyright (c) 2022-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. | |
| """ | |
| DraftTarget One-Model Speculative Decoding Implementation. | |
| This module implements a one-model approach for DraftTarget speculative decoding, | |
| where the draft and target models share the same model engine. The draft model | |
| layers are integrated into the target model's KV cache and run in a single forward pass. | |
| """ |
🤖 Prompt for AI Agents
In `@tensorrt_llm/_torch/speculative/draft_target.py` around lines 1 - 7, Add the
required NVIDIA copyright header (including the year of latest meaningful
modification) to the top of the module to comply with project guidelines; insert
the standard multi-line header block before the module docstring in
tensorrt_llm._torch.speculative.draft_target (i.e., at the very top of this
file) so the header precedes the existing module docstring and includes the
proper copyright line and any required SPDX or license identifier.
IzzyPutterman
changed the title
mikeiovine
left a comment
mikeiovine
left a comment
There was a problem hiding this comment.
Can we add a test? The existing 2 model test can just be converted to use this
| batch_indices = torch.arange(num_seqs, dtype=torch.int, device="cpu", pin_memory=True) | ||
| self.batch_indices_cuda[:num_seqs].copy_(batch_indices, non_blocking=True) | ||
| # Adjust num_tokens for generation phase | ||
| self.num_tokens -= (self.num_generations) * self.max_draft_len |
There was a problem hiding this comment.
| self.num_tokens -= (self.num_generations) * self.max_draft_len | |
| self.num_tokens -= self.num_generations * self.max_draft_len |
nit
tensorrt_llm/llmapi/llm_args.py
Outdated
| # If None, will be inferred from the speculative_model_dir config. | ||
| num_draft_layers: Optional[int] = None | ||
| # Enable advanced sampling (temperature, top-k, top-p) for one-model mode. | ||
| allow_advanced_sampling: bool = False |
There was a problem hiding this comment.
This should probably be in the base class (I think I moved it there? If not on trunk then definitely in the pending PR that adds MTP sampling support)
IzzyPutterman
force-pushed
the
iputterman/draft-target-one-model
branch
from
25dd26b to
8e6bfdb
Compare
|
/bot run |
|
PR_Github #33021 [ run ] triggered by Bot. Commit: |
IzzyPutterman
force-pushed
the
iputterman/draft-target-one-model
branch
from
9962dc0 to
5246af9
Compare
|
/bot run |
|
PR_Github #33057 [ run ] triggered by Bot. Commit: |
|
PR_Github #33057 [ run ] completed with state
|
IzzyPutterman
force-pushed
the
iputterman/draft-target-one-model
branch
from
5246af9 to
4c5bfb0
Compare
|
/bot run |
|
PR_Github #34282 [ run ] triggered by Bot. Commit: |
|
PR_Github #34282 [ run ] completed with state
|
|
broken atm, do not merge |
IzzyPutterman
force-pushed
the
iputterman/draft-target-one-model
branch
2 times, most recently
from
2bdeb4f to
92ded5d
Compare
|
/bot run |
|
PR_Github #34675 [ run ] triggered by Bot. Commit: |
IzzyPutterman
force-pushed
the
iputterman/draft-target-one-model
branch
from
92ded5d to
a567743
Compare
|
This PR will be blocked until #10502 |
IzzyPutterman
force-pushed
the
iputterman/draft-target-one-model
branch
from
6c2c3cf to
0d22210
Compare
|
/bot run |
|
PR_Github #35415 [ run ] triggered by Bot. Commit: |
|
PR_Github #35415 [ run ] completed with state
|
IzzyPutterman
force-pushed
the
iputterman/draft-target-one-model
branch
from
0d22210 to
f76646b
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #35423 [ run ] triggered by Bot. Commit: |
IzzyPutterman
force-pushed
the
iputterman/draft-target-one-model
branch
from
f76646b to
5cfe03a
Compare
|
PR_Github #35423 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #35516 [ run ] triggered by Bot. Commit: |
|
PR_Github #35516 [ run ] completed with state
|
| return self.is_mtp_one_model() or self.is_eagle3_one_model( | ||
| ) or self.is_draft_target_one_model() |
There was a problem hiding this comment.
We can simplify this to is_one_engine(). Don't have to do it now though
tensorrt_llm/llmapi/llm_args.py
Outdated
| # Invalidate cached spec_dec_mode in case it was already accessed | ||
| self.speculative_config.__dict__.pop('spec_dec_mode', None) |
There was a problem hiding this comment.
IMO it is simpler to just remove the cached_property decorator. Accessing the attribute is very cheap
| """ | ||
|
|
||
| # The max number of tokens | ||
| max_num_tokens: int = 0 |
There was a problem hiding this comment.
I think we already have an attribute similar to this in the base class? max_total_tokens or something like that
Signed-off-by: Izzy Putterman <iputterman@nvidia.com>
Signed-off-by: Izzy Putterman <iputterman@nvidia.com>
Signed-off-by: Izzy Putterman <iputterman@nvidia.com>
Signed-off-by: Izzy Putterman <iputterman@nvidia.com>
IzzyPutterman
force-pushed
the
iputterman/draft-target-one-model
branch
2 times, most recently
from
9356b54 to
e8fddca
Compare
|
/bot run |
|
PR_Github #36711 [ run ] triggered by Bot. Commit: |
Signed-off-by: Izzy Putterman <iputterman@nvidia.com>
IzzyPutterman
force-pushed
the
iputterman/draft-target-one-model
branch
from
e8fddca to
109c9a7
Compare
|
PR_Github #36711 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #36819 [ run ] triggered by Bot. Commit: |
|
PR_Github #36819 [ run ] completed with state
|
|
close in favor of: #11758 |
3 participants
Summary by CodeRabbit
✏️ Tip: You can customize this high-level summary in your review settings.
Description
Test Coverage
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)
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
/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...Provide a user friendly way for developers to interact with a Jenkins server.
Run
/bot [-h|--help]to print this help message.See details below for each supported subcommand.
Details
run [--reuse-test (optional)pipeline-id --disable-fail-fast --skip-test --stage-list "A10-PyTorch-1, xxx" --gpu-type "A30, H100_PCIe" --test-backend "pytorch, cpp" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" --detailed-log --debug(experimental)]Launch build/test pipelines. All previously running jobs will be killed.
--reuse-test (optional)pipeline-id(OPTIONAL) : Allow the new pipeline to reuse build artifacts and skip successful test stages from a specified pipeline or the last pipeline if no pipeline-id is indicated. If the Git commit ID has changed, this option will be always ignored. The DEFAULT behavior of the bot is to reuse build artifacts and successful test results from the last pipeline.--disable-reuse-test(OPTIONAL) : Explicitly prevent the pipeline from reusing build artifacts and skipping successful test stages from a previous pipeline. Ensure that all builds and tests are run regardless of previous successes.--disable-fail-fast(OPTIONAL) : Disable fail fast on build/tests/infra failures.--skip-test(OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.--stage-list "A10-PyTorch-1, xxx"(OPTIONAL) : Only run the specified test stages. Examples: "A10-PyTorch-1, xxx". Note: Does NOT update GitHub check status.--gpu-type "A30, H100_PCIe"(OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.--test-backend "pytorch, cpp"(OPTIONAL) : Skip test stages which don't match the specified backends. Only support [pytorch, cpp, tensorrt, triton]. Examples: "pytorch, cpp" (does not run test stages with tensorrt or triton backend). Note: Does NOT update GitHub pipeline status.--only-multi-gpu-test(OPTIONAL) : Only run the multi-GPU tests. Note: Does NOT update GitHub check status.--disable-multi-gpu-test(OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.--add-multi-gpu-test(OPTIONAL) : Force run the multi-GPU tests in addition to running L0 pre-merge pipeline.--post-merge(OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.--extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx"(OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx".--detailed-log(OPTIONAL) : Enable flushing out all logs to the Jenkins console. This will significantly increase the log volume and may slow down the job.--debug(OPTIONAL) : Experimental feature. Enable access to the CI container for debugging purpose. Note: Specify exactly one stage in thestage-listparameter to access the appropriate container environment. Note: Does NOT update GitHub check status.For guidance on mapping tests to stage names, see
docs/source/reference/ci-overview.mdand the
scripts/test_to_stage_mapping.pyhelper.kill
killKill all running builds associated with pull request.
skip
skip --comment COMMENTSkip testing for latest commit on pull request.
--comment "Reason for skipping build/test"is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.reuse-pipeline
reuse-pipelineReuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.