NVIDIA · mikeiovine · Jul 10, 2025 · Jul 1, 2025 · Jul 8, 2025 · Jul 10, 2025
@@ -68,7 +68,7 @@ docker run -d --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 \
     -p 8000:8000 --gpus=all -e "TRTLLM_ENABLE_PDL=1" \
     -v /path/to/maverick:/config/models/maverick -v /path/to/eagle:/config/models/eagle \
     docker.io/<username>/tensorrt_llm:main sh \
-        -c "echo -e 'enable_attention_dp: false\nenable_min_latency: true\nenable_autotuner: false\ncuda_graph_config:\n  max_batch_size: 8\nspeculative_config:\n  decoding_type: Eagle\n  max_draft_len: 3\n  pytorch_weights_path: /config/models/eagle\nkv_cache_config:\n  enable_block_reuse: false' > c.yaml && \
+        -c "echo -e 'enable_attention_dp: false\nenable_min_latency: true\nenable_autotuner: false\ncuda_graph_config:\n  max_batch_size: 8\nspeculative_config:\n  decoding_type: Eagle\n  max_draft_len: 3\n  speculative_model_dir: /config/models/eagle\nkv_cache_config:\n  enable_block_reuse: false' > c.yaml && \
         TRT_LLM_DISABLE_LOAD_WEIGHTS_IN_PARALLEL=True \
         trtllm-serve /config/models/maverick \
             --host 0.0.0.0 --port 8000 \

@@ -108,7 +108,7 @@ When using EAGLE-2, please enable `--eagle_use_dynamic_tree`, which indicates wh
 - In EagleNet2, the `N` output nodes of EagleNet1 are expanded, and each node expands `N` new draft tokens. Therefore, this layer also has a total of `N * N` draft tokens. And select the top `N` as the output of this layer.
 - Etc.
 
-Finally, after `num_eagle_layer` EagleNets, `N + N * N * (num_eagle_layer - 1)` draft tokens are generated. We will rebuild the final tree based on all draft tokens and their scores. The final generated tree will have `min(N + N * N * (num_eagle_layer - 1), max_draft_tokens)` nodes.
+Finally, after `num_eagle_layer` EagleNets, `N + N * N * (num_eagle_layer - 1)` draft tokens are generated. We will rebuild the final tree based on all draft tokens and their scores. The final generated tree will have `min(N + N * N * (num_eagle_layer - 1), max_draft_len)` nodes.
 
 
 

@@ -23,12 +23,12 @@ def main():
     model = "lmsys/vicuna-7b-v1.3"
 
     # The end user can customize the eagle decoding configuration by specifying the
-    # speculative_model, max_draft_len, num_eagle_layers, max_non_leaves_per_layer, eagle_choices
+    # speculative_model_dir, max_draft_len, num_eagle_layers, max_non_leaves_per_layer, eagle_choices
     # greedy_sampling,posterior_threshold, use_dynamic_tree and dynamic_tree_max_topK
     # with the EagleDecodingConfig class
 
     speculative_config = EagleDecodingConfig(
-        speculative_model="yuhuili/EAGLE-Vicuna-7B-v1.3",
+        speculative_model_dir="yuhuili/EAGLE-Vicuna-7B-v1.3",
         max_draft_len=63,
         num_eagle_layers=4,
         max_non_leaves_per_layer=10,

@@ -23,12 +23,12 @@ def main():
     model = "lmsys/vicuna-7b-v1.3"
 
     # The end user can customize the eagle decoding configuration by specifying the
-    # speculative_model, max_draft_len, num_eagle_layers, max_non_leaves_per_layer, eagle_choices
+    # speculative_model_dir, max_draft_len, num_eagle_layers, max_non_leaves_per_layer, eagle_choices
     # greedy_sampling,posterior_threshold, use_dynamic_tree and dynamic_tree_max_topK
     # with the EagleDecodingConfig class
 
     speculative_config = EagleDecodingConfig(
-        speculative_model="yuhuili/EAGLE-Vicuna-7B-v1.3",
+        speculative_model_dir="yuhuili/EAGLE-Vicuna-7B-v1.3",
         max_draft_len=63,
         num_eagle_layers=4,
         max_non_leaves_per_layer=10,

@@ -48,10 +48,10 @@ def run_medusa_decoding(use_modelopt_ckpt=False, model_dir=None):
         model = "lmsys/vicuna-7b-v1.3"
 
         # The end user can customize the medusa decoding configuration by specifying the
-        # speculative_model, max_draft_len, medusa heads num and medusa choices
+        # speculative_model_dir, max_draft_len, medusa heads num and medusa choices
         # with the MedusaDecodingConfig class
         speculative_config = MedusaDecodingConfig(
-                                        speculative_model="FasterDecoding/medusa-vicuna-7b-v1.3",
+                                        speculative_model_dir="FasterDecoding/medusa-vicuna-7b-v1.3",
                                         max_draft_len=63,
                                         num_medusa_heads=4,
                                         medusa_choices=[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], \

@@ -35,7 +35,7 @@ def run_MTP(model: Optional[str] = None):
 def run_Eagle3():
     spec_config = EagleDecodingConfig(
         max_draft_len=3,
-        pytorch_weights_path="yuhuili/EAGLE3-LLaMA3.1-Instruct-8B",
+        speculative_model_dir="yuhuili/EAGLE3-LLaMA3.1-Instruct-8B",
         eagle3_one_model=True)
 
     llm = LLM(
@@ -50,7 +50,7 @@ def run_Eagle3():
 
 def run_ngram():
     spec_config = NGramDecodingConfig(
-        prompt_lookup_num_tokens=3,
+        max_draft_len=3,
         max_matching_ngram_size=3,
         is_keep_all=True,
         is_use_oldest=True,

@@ -169,15 +169,15 @@ def setup_llm(args):
     elif spec_decode_algo == "EAGLE3":
         spec_config = EagleDecodingConfig(
             max_draft_len=args.spec_decode_nextn,
-            pytorch_weights_path=args.draft_model_dir,
+            speculative_model_dir=args.draft_model_dir,
             eagle3_one_model=args.use_one_model)
     elif spec_decode_algo == "DRAFT_TARGET":
         spec_config = DraftTargetDecodingConfig(
             max_draft_len=args.spec_decode_nextn,
-            pytorch_weights_path=args.draft_model_dir)
+            speculative_model_dir=args.draft_model_dir)
     elif spec_decode_algo == "NGRAM":
         spec_config = NGramDecodingConfig(
-            prompt_lookup_num_tokens=args.spec_decode_nextn,
+            max_draft_len=args.spec_decode_nextn,
             max_matching_ngram_size=args.max_matching_ngram_size,
             is_keep_all=True,
             is_use_oldest=True,

@@ -261,8 +261,8 @@ def create_autodeploy_executor(executor_config: ExecutorConfig, checkpoint_dir:
 
     max_num_sequences = ad_config.max_batch_size * dist_mapping.pp_size
     # some derivative properties
-    max_draft_tokens = (
-        0 if ad_config.speculative_config is None else ad_config.speculative_config.max_draft_tokens
+    max_draft_len = (
+        0 if ad_config.speculative_config is None else ad_config.speculative_config.max_draft_len
     )
 
     # initialize model engine
@@ -299,7 +299,7 @@ def create_autodeploy_executor(executor_config: ExecutorConfig, checkpoint_dir:
     # correctly for models as needed.
     sampler_args = TorchSampler.Args(
         max_seq_len=ad_config.max_seq_len,
-        max_draft_tokens=max_draft_tokens,
+        max_draft_len=max_draft_len,
         max_num_sequences=max_num_sequences,
         max_beam_width=executor_config.max_beam_width,
         enable_mixed_sampler=ad_config.enable_mixed_sampler,
@@ -317,7 +317,7 @@ def create_autodeploy_executor(executor_config: ExecutorConfig, checkpoint_dir:
         disable_overlap_scheduler=ad_config.disable_overlap_scheduler,
         max_input_len=ad_config.max_input_len,
         max_batch_size=ad_config.max_batch_size,
-        max_draft_tokens=max_draft_tokens,
+        max_draft_len=max_draft_len,
         max_beam_width=ad_config.max_beam_width,
     )
     return py_executor
@@ -70,7 +70,7 @@ class ModelConfig(Generic[TConfig]):
     # to support mixed quantization.
     skip_create_weights_in_init: bool = False
 
-    spec_config: Optional["SpecConfig"] = None
+    spec_config: Optional["DecodingBaseConfig"] = None
     lora_config: Optional["LoraConfig"] = None
 
     is_generation: bool = True

@@ -340,7 +340,7 @@ def __init__(self, model: TModel, model_config: ModelConfig[TConfig]):
                 model_config, 'spec_config', None
         ) and model_config.spec_config.spec_dec_mode.use_one_engine():
             draft_config = ModelConfig.from_pretrained(
-                model_config.spec_config.draft_model_path,
+                model_config.spec_config.speculative_model_dir,
                 trust_remote_code=True,
                 attn_backend=model_config.attn_backend,
                 moe_backend=model_config.moe_backend,

@@ -157,10 +157,10 @@ def _get_token_num_for_estimation(self) -> int:
         if not pytorch_backend_config.disable_overlap_scheduler:
             num_extra_tokens_per_seq = num_extra_tokens_per_seq + 1
             if spec_cfg is not None:
-                num_extra_tokens_per_seq += spec_cfg.max_draft_tokens
+                num_extra_tokens_per_seq += spec_cfg.max_draft_len
 
         if spec_cfg is not None:
-            num_extra_tokens_per_seq += spec_cfg.max_draft_tokens
+            num_extra_tokens_per_seq += spec_cfg.max_draft_len
             num_extra_tokens_per_seq += spec_cfg.num_extra_kv_tokens
         for req in self._dummy_reqs:
             num_req_tokens = len(req.input_token_ids) + num_extra_tokens_per_seq
@@ -538,7 +538,7 @@ def create_py_executor_instance(
         disable_overlap_scheduler,
         max_batch_size=executor_config.max_batch_size,
         max_beam_width=executor_config.max_beam_width,
-        max_draft_tokens=spec_config.max_draft_tokens
+        max_draft_len=spec_config.max_draft_len
         if spec_config is not None else 0,
         kv_cache_transceiver=kv_cache_transceiver,
         draft_model_engine=draft_model_engine,
@@ -549,11 +549,11 @@ def create_py_executor_instance(
 def create_torch_sampler_args(executor_config: ExecutorConfig, mapping: Mapping,
                               *, max_seq_len: int, enable_mixed_sampler: bool):
     max_num_sequences = executor_config.max_batch_size * mapping.pp_size
-    max_draft_tokens = (0 if executor_config.speculative_config is None else
-                        executor_config.speculative_config.max_draft_tokens)
+    max_draft_len = (0 if executor_config.speculative_config is None else
+                     executor_config.speculative_config.max_draft_len)
     return TorchSampler.Args(
         max_seq_len=max_seq_len,
-        max_draft_tokens=max_draft_tokens,
+        max_draft_len=max_draft_len,
         max_num_sequences=max_num_sequences,
         max_beam_width=executor_config.max_beam_width,
         enable_mixed_sampler=enable_mixed_sampler,

@@ -8,7 +8,6 @@
 from ...logger import logger
 from ...mapping import Mapping
 from ..model_config import MoeLoadBalancerConfig
-from ..speculative import SpecConfig
 from .resource_manager import BaseResourceManager
 
 
@@ -110,7 +109,7 @@ def update_executor_config(
         pytorch_backend_config: Optional[PyTorchConfig] = None,
         mapping: Optional[Mapping] = None,
         build_config: Optional[BuildConfig] = None,
-        speculative_config: Optional[SpecConfig] = None,
+        speculative_config: Optional["DecodingBaseConfig"] = None,
         hf_model_dir: Optional[str] = None,
         max_input_len: Optional[int] = None,
         max_seq_len: Optional[int] = None):

@@ -53,7 +53,7 @@ def __init__(
         # [CUDA graph spec decode padding]
         # We pad input IDs/position IDs to the maximum draft length (token per request).
         # We're forced to do this because we cannot reallocate inputs over many graph runs.
-        token_per_request = spec_metadata.max_draft_tokens + 1 if spec_metadata is not None else 1
+        token_per_request = spec_metadata.max_draft_len + 1 if spec_metadata is not None else 1
 
         # Using ones instead of zeros prevents NaNs in e.g. Deepseek
         self.input_ids = torch.ones((batch_size * token_per_request, ),

@@ -51,7 +51,7 @@
                                      run_concurrently, timing)
 from ..modules.fused_moe.moe_load_balancer import (
     MoeLoadBalancer, MoeLoadBalancerIterContext, maybe_create_moe_load_balancer)
-from ..speculative import SpecConfig, SpecMetadata, get_spec_metadata
+from ..speculative import SpecMetadata, get_spec_metadata
 from ..utils import (get_model_extra_attrs, set_torch_compiling,
                      with_model_extra_attrs)
 from .config import LoadFormat, PyTorchConfig
@@ -353,7 +353,7 @@ def __init__(
         mapping: Optional[Mapping] = None,
         attn_runtime_features: Optional[AttentionRuntimeFeatures] = None,
         dist: Optional[MPIDist] = None,
-        spec_config: Optional[SpecConfig] = None,
+        spec_config: Optional["DecodingBaseConfig"] = None,
         guided_decoding_config: Optional[GuidedDecodingConfig] = None,
         lora_config: Optional[LoraConfig] = None,
         is_draft_model: bool = False,
@@ -456,7 +456,7 @@ def __init__(
         if self.is_spec_decode:
             self.spec_metadata = None
             self.spec_config.update_from_model_config(self.model.config)
-            max_num_draft_tokens = self.spec_config.max_draft_tokens * batch_size
+            max_num_draft_tokens = self.spec_config.max_draft_len * batch_size
             self.draft_tokens_cuda = torch.empty((max_num_draft_tokens, ),
                                                  dtype=torch.int,
                                                  device='cuda')
@@ -472,7 +472,7 @@ def __init__(
                                                              device='cuda')
             self.without_logits = self.spec_config.spec_dec_mode.without_logits(
             )
-            self.max_draft_len = spec_config.max_draft_tokens
+            self.max_draft_len = spec_config.max_draft_len
         else:
             self.without_logits = False
             self.max_draft_len = 0
@@ -858,6 +858,7 @@ def _set_up_spec_metadata(
         if no_cache:
             return get_spec_metadata(
                 self.spec_config,
+                self.model.config,
                 self.batch_size,
                 max_num_tokens=self.max_num_tokens,
                 spec_resource_manager=spec_resource_manager,
@@ -867,6 +868,7 @@ def _set_up_spec_metadata(
             return self.spec_metadata
         self.spec_metadata = get_spec_metadata(
             self.spec_config,
+            self.model.config,
             self.batch_size,
             max_num_tokens=self.max_num_tokens,
             spec_resource_manager=spec_resource_manager,
@@ -951,7 +953,7 @@ def _round_up_batch_size(self, batch_size: int) -> int:
     def _maybe_get_cuda_graph(
         self,
         batch: ScheduledRequests,
-        spec_config: Optional[SpecConfig] = None
+        spec_config: Optional["DecodingBaseConfig"] = None
     ) -> Optional[DecodingCUDAGraphRunner]:
         """
         Get a CUDA graph runner or return None (e.g. if CUDA graphs are disabled
@@ -961,7 +963,7 @@ def _maybe_get_cuda_graph(
         if ExpertStatistic.set_iter(self.iter_counter):
             return None
 
-        spec_max_draft_tokens = spec_config.max_draft_tokens if self.is_spec_decode else 0
+        spec_max_draft_tokens = spec_config.max_draft_len if self.is_spec_decode else 0
         can_run_cuda_graph = batch.can_run_cuda_graph
         batch_size = len(batch.generation_requests)
         if self._run_cuda_graphs and self.enable_attention_dp and self.mapping.tp_size > 1:
@@ -1078,7 +1080,8 @@ def init_meta_tensor(t: torch.Tensor):
 
                 if self.spec_config is not None and self.spec_config.spec_dec_mode.need_load_draft_weights(
                 ):
-                    weights = load_weights(self.spec_config.draft_model_path)
+                    weights = load_weights(
+                        self.spec_config.speculative_model_dir)
                     model.load_draft_weights(weights)
 
             elif load_format == LoadFormat.DUMMY:
@@ -1261,9 +1264,8 @@ def _prepare_tp_inputs(
         extend_requests += extend_dummy_requests
 
         if not self._disable_overlap_scheduler and self.is_spec_decode:
-            spec_dec_mode = self.spec_config.spec_dec_mode
-            assert spec_dec_mode.support_overlap_scheduler(
-            ), f"{self.spec_config.spec_dec_name} does not support overlap scheduler"
+            assert self.spec_config.spec_dec_mode.support_overlap_scheduler(
+            ), f"{self.spec_config.decoding_type} does not support overlap scheduler"
 
         # will contain previous batch indices of generation requests
         previous_batch_indices = []
@@ -2078,7 +2080,7 @@ def forward(
                 spec_metadata.spec_dec_mode.attention_need_spec_dec_mode(),
                 spec_metadata.is_spec_dec_tree,
                 spec_metadata.is_spec_dec_dynamic_tree,
-                spec_metadata.max_draft_tokens)
+                spec_metadata.max_draft_len)
         else:
             spec_metadata = None
 

@@ -173,7 +173,7 @@ def __init__(self,
                  max_input_len: int = 2048,
                  max_batch_size: int = 8,
                  max_beam_width: int = 1,
-                 max_draft_tokens: int = 0,
+                 max_draft_len: int = 0,
                  kv_cache_transceiver: Optional[KvCacheTransceiver] = None,
                  draft_model_engine: Optional[ModelEngine] = None,
                  garbage_collection_gen0_threshold: Optional[int] = None,
@@ -207,7 +207,7 @@ def __init__(self,
         self.active = True
         self.next_req_id = max_batch_size  # The first max_batch_size request IDs are reserved for dummy requests
         self.max_beam_width = max_beam_width
-        self.max_draft_tokens = max_draft_tokens
+        self.max_draft_len = max_draft_len
         self.print_log = model_engine.pytorch_backend_config.print_iter_log
         self.enable_iter_perf_stats = model_engine.pytorch_backend_config.enable_iter_perf_stats
         self.enable_iter_req_stats = model_engine.pytorch_backend_config.enable_iter_req_stats
@@ -979,7 +979,7 @@ def _prepare_draft_requests(self):
                                      LlmRequestState.DISAGG_GENERATION_INIT):
                     continue
                 req.py_last_draft_tokens = req.py_draft_tokens
-                max_draft_len = self.model_engine.spec_config.max_draft_tokens
+                max_draft_len = self.model_engine.spec_config.max_draft_len
 
                 if max_draft_len > 0:
                     req.py_draft_tokens = [0] * max_draft_len
@@ -1523,7 +1523,7 @@ def _pad_attention_dp_dummy_request(self):
                 request_ids=[0],
                 is_gen=not self.has_context_request,
                 prepare_resource=not self.has_context_request,
-                max_num_draft_tokens=self.max_draft_tokens,
+                max_num_draft_tokens=self.max_draft_len,
             )[0]
             llm_request.is_attention_dp_dummy = True
             spec_resource_manager = self.resource_manager.get_resource_manager(
@@ -1871,15 +1871,15 @@ def _process_decoded_tokens(draft_batch):
             # this? Just needs proper kernel support.
             def _pad_to_max_draft_tokens():
                 for req in scheduled_requests.generation_requests:
-                    max_draft_tokens = self.max_draft_tokens
+                    max_draft_len = self.max_draft_len
                     num_draft_tokens = len(req.py_draft_tokens)
                     req.py_draft_tokens.extend(
-                        0 for _ in range(max_draft_tokens - num_draft_tokens))
+                        0 for _ in range(max_draft_len - num_draft_tokens))
 
             draft_batch.generation_requests = draft_batch.context_requests + draft_batch.generation_requests
             draft_batch.context_requests = []
 
-            for i in range(self.max_draft_tokens - 1):
+            for i in range(self.max_draft_len - 1):
                 if len(draft_batch.generation_requests) == 0:
                     break
-Original file line number
+Diff line change
@@ Expand Up @@
     - In EagleNet2, the `N` output nodes of EagleNet1 are expanded, and each node expands `N` new draft tokens. Therefore, this layer also has a total of `N * N` draft tokens. And select the top `N` as the output of this layer.
     - Etc.
-    Finally, after `num_eagle_layer` EagleNets, `N + N * N * (num_eagle_layer - 1)` draft tokens are generated. We will rebuild the final tree based on all draft tokens and their scores. The final generated tree will have `min(N + N * N * (num_eagle_layer - 1), max_draft_tokens)` nodes.
+    Finally, after `num_eagle_layer` EagleNets, `N + N * N * (num_eagle_layer - 1)` draft tokens are generated. We will rebuild the final tree based on all draft tokens and their scores. The final generated tree will have `min(N + N * N * (num_eagle_layer - 1), max_draft_len)` nodes.
@@ Expand Down @@