[TRTLLM-14388][refactor] BREAKING: Remove 2-model code and tests

.pre-commit-config.yaml

@@ -323,10 +323,8 @@ common-files: &common_files |
         tensorrt_llm/_torch/speculative/__init__.py |
         tensorrt_llm/_torch/speculative/auto_heuristic.py |
         tensorrt_llm/_torch/speculative/drafter.py |
-        tensorrt_llm/_torch/speculative/drafting_loops.py |
         tensorrt_llm/_torch/speculative/eagle3.py |
         tensorrt_llm/_torch/speculative/interface.py |
-        tensorrt_llm/_torch/speculative/model_drafter.py |
         tensorrt_llm/_torch/speculative/mtp.py |
         tensorrt_llm/_torch/speculative/ngram.py |
         tensorrt_llm/_torch/speculative/save_hidden_state.py |
@@ -1093,10 +1091,8 @@ legacy-files: &legacy_files |
         tensorrt_llm/_torch/speculative/__init__.py |
         tensorrt_llm/_torch/speculative/auto_heuristic.py |
         tensorrt_llm/_torch/speculative/drafter.py |
-        tensorrt_llm/_torch/speculative/drafting_loops.py |
         tensorrt_llm/_torch/speculative/eagle3.py |
         tensorrt_llm/_torch/speculative/interface.py |
-        tensorrt_llm/_torch/speculative/model_drafter.py |
         tensorrt_llm/_torch/speculative/mtp.py |
         tensorrt_llm/_torch/speculative/ngram.py |
         tensorrt_llm/_torch/speculative/save_hidden_state.py |

docs/source/blogs/tech_blog/blog06_Llama4_maverick_eagle_guide.md

@@ -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_autotuner: false\nenable_attention_dp: false\nenable_min_latency: true\ncuda_graph_config:\n  max_batch_size: 8\nspeculative_config:\n  decoding_type: Eagle3\n  max_draft_len: 3\n  speculative_model_dir: /config/models/eagle\n  eagle3_one_model: true\nkv_cache_config:\n  enable_block_reuse: false' > c.yaml && \
+        -c "echo -e 'enable_autotuner: false\nenable_attention_dp: false\nenable_min_latency: true\ncuda_graph_config:\n  max_batch_size: 8\nspeculative_config:\n  decoding_type: Eagle3\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 \

docs/source/features/torch_compile_and_piecewise_cuda_graph.md

@@ -88,18 +88,7 @@ Even with Piecewise CUDA Graph enabled, you may still observe bubbles in the con
 
 Torch compile cannot work with multi-ModelEngine config.
 
-1. Speculative Decoding in Two-Model Style
-
-``` yaml
-speculative_config:
-  decoding_type: "MTP"
-  mtp_eagle_one_model: False # Not supported
-
-speculative_config:
-  decoding_type: "Eagle3"
-  eagle3_one_model: False # Not supported
-```
-
+1. NGram or User-Provided Speculative Decoding
 2. Multimodal Model Family
 
 ## Development Guide

docs/source/legacy/advanced/speculative-decoding.md

@@ -175,7 +175,7 @@ Similarly to ReDrafter, TensorRT-LLM implements the EAGLE model such that logits
 
 ### Disaggregated Serving
 
-[Disaggregated Serving](https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/features/disaggregated-service.md) with EAGLE3 using the two model approach is supported in the Pytorch backend. Please refer to the following [Dynamo example](https://github.com/ai-dynamo/dynamo/blob/main/examples/tensorrt_llm/llama4_plus_eagle.md) on how to run EAGLE3 with Disaggregated Serving for Llama 4 Maverick.
+[Disaggregated Serving](https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/features/disaggregated-service.md) with EAGLE3 is supported in the Pytorch backend. Please refer to the following [Dynamo example](https://github.com/ai-dynamo/dynamo/blob/main/examples/tensorrt_llm/llama4_plus_eagle.md) on how to run EAGLE3 with Disaggregated Serving for Llama 4 Maverick.
 
 ## Lookahead Decoding
 

examples/auto_deploy/model_registry/configs/super_v3_mtp.yaml

@@ -15,7 +15,6 @@ kv_cache_config:
 speculative_config:
   decoding_type: MTP
   num_nextn_predict_layers: 6
-  mtp_eagle_one_model: true
 attention_dp_config:
   enable_balance: true
   batching_wait_iters: 10

examples/llm-api/llm_speculative_decoding.py

@@ -34,8 +34,7 @@ def run_MTP(model: Optional[str] = None):
 def run_Eagle3():
     spec_config = Eagle3DecodingConfig(
         max_draft_len=3,
-        speculative_model="yuhuili/EAGLE3-LLaMA3.1-Instruct-8B",
-        eagle3_one_model=True)
+        speculative_model="yuhuili/EAGLE3-LLaMA3.1-Instruct-8B")
 
     kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.8)
 

examples/llm-api/quickstart_advanced.py

@@ -210,9 +210,6 @@ def add_llm_args(parser):
     parser.add_argument('--spec_decode_max_draft_len', type=int, default=1)
     parser.add_argument('--draft_model_dir', type=str, default=None)
     parser.add_argument('--max_matching_ngram_size', type=int, default=5)
-    parser.add_argument('--use_one_model',
-                        default=True,
-                        action=argparse.BooleanOptionalAction)
     parser.add_argument('--eagle_choices', type=str, default=None)
     parser.add_argument('--use_dynamic_tree',
                         default=False,
@@ -301,15 +298,12 @@ def setup_llm(args, **kwargs):
     ) if args.spec_decode_algo is not None else None
 
     if spec_decode_algo == 'MTP':
-        if not args.use_one_model:
-            print("Running MTP eagle with two model style.")
         spec_config = MTPDecodingConfig(
             max_draft_len=args.spec_decode_max_draft_len,
             use_relaxed_acceptance_for_thinking=args.
             use_relaxed_acceptance_for_thinking,
             relaxed_topk=args.relaxed_topk,
             relaxed_delta=args.relaxed_delta,
-            mtp_eagle_one_model=args.use_one_model,
             use_dynamic_tree=args.use_dynamic_tree,
             dynamic_tree_max_topK=args.dynamic_tree_max_topK,
             max_total_draft_tokens=args.max_total_draft_tokens,
@@ -318,7 +312,6 @@ def setup_llm(args, **kwargs):
         spec_config = Eagle3DecodingConfig(
             max_draft_len=args.spec_decode_max_draft_len,
             speculative_model=args.draft_model_dir,
-            eagle3_one_model=args.use_one_model,
             eagle_choices=args.eagle_choices,
             use_dynamic_tree=args.use_dynamic_tree,
             dynamic_tree_max_topK=args.dynamic_tree_max_topK,

examples/models/core/deepseek_v4/README.md

@@ -182,8 +182,7 @@ python quickstart_advanced.py \
   --max_seq_len 4096 \
   --kv_cache_fraction 0.5 \
   --spec_decode_algo MTP \
-  --spec_decode_max_draft_len N \
-  --use_one_model
+  --spec_decode_max_draft_len N
 ```
 
 `N` is the number of draft tokens to predict. Start with `N=1` for bring-up, then increase it after

examples/models/core/exaone/README.md

@@ -197,8 +197,7 @@ python ../../../llm-api/quickstart_advanced.py \
     --moe_ep_size 8 \
     --enable_attention_dp \
     --spec_decode_algo MTP \
-    --spec_decode_max_draft_len N \
-    --use_one_model
+    --spec_decode_max_draft_len N
 ```
 
 `N` is the number of MTP modules. When `N` is equal to `0`, which means that MTP is not used (default). When `N` is greater than `0`, which means that `N` MTP modules are enabled. In the current implementation, the weight of each MTP module is shared.

examples/models/core/mistral_large_3/README.md

@@ -46,7 +46,6 @@ mpirun -n 1 --allow-run-as-root --oversubscribe python3 examples/llm-api/quickst
     --checkpoint_format mistral \
     --spec_decode_algo EAGLE3 \
     --spec_decode_max_draft_len 1 \
-    --use_one_model \
     --draft_model_dir ${mistral_large_3_eagle_model_path} \
     --eagle3_model_arch mistral_large3 \
     --moe_backend TRTLLM

examples/models/core/nemotron/README_nemotron_super_v3.md

@@ -182,7 +182,6 @@ python3 examples/llm-api/quickstart_advanced.py \
     --moe_backend=TRTLLM \
     --spec_decode_algo=MTP \
     --spec_decode_max_draft_len=3 \
-    --use_one_model \
     --tp_size=8 \
     --moe_ep_size 8 \
     --apply_chat_template

legacy-files.txt

@@ -315,10 +315,8 @@ tensorrt_llm/_torch/shared_tensor/shared_tensor.py
 tensorrt_llm/_torch/speculative/__init__.py
 tensorrt_llm/_torch/speculative/auto_heuristic.py
 tensorrt_llm/_torch/speculative/drafter.py
-tensorrt_llm/_torch/speculative/drafting_loops.py
 tensorrt_llm/_torch/speculative/eagle3.py
 tensorrt_llm/_torch/speculative/interface.py
-tensorrt_llm/_torch/speculative/model_drafter.py
 tensorrt_llm/_torch/speculative/mtp.py
 tensorrt_llm/_torch/speculative/ngram.py
 tensorrt_llm/_torch/speculative/save_hidden_state.py

pyproject.toml

@@ -372,10 +372,8 @@ exclude = [
     "tensorrt_llm/_torch/speculative/__init__.py",
     "tensorrt_llm/_torch/speculative/auto_heuristic.py",
     "tensorrt_llm/_torch/speculative/drafter.py",
-    "tensorrt_llm/_torch/speculative/drafting_loops.py",
     "tensorrt_llm/_torch/speculative/eagle3.py",
     "tensorrt_llm/_torch/speculative/interface.py",
-    "tensorrt_llm/_torch/speculative/model_drafter.py",
     "tensorrt_llm/_torch/speculative/mtp.py",
     "tensorrt_llm/_torch/speculative/ngram.py",
     "tensorrt_llm/_torch/speculative/save_hidden_state.py",

ruff-legacy.toml

@@ -332,10 +332,8 @@ include = [
     "tensorrt_llm/_torch/speculative/__init__.py",
     "tensorrt_llm/_torch/speculative/auto_heuristic.py",
     "tensorrt_llm/_torch/speculative/drafter.py",
-    "tensorrt_llm/_torch/speculative/drafting_loops.py",
     "tensorrt_llm/_torch/speculative/eagle3.py",
     "tensorrt_llm/_torch/speculative/interface.py",
-    "tensorrt_llm/_torch/speculative/model_drafter.py",
     "tensorrt_llm/_torch/speculative/mtp.py",
     "tensorrt_llm/_torch/speculative/ngram.py",
     "tensorrt_llm/_torch/speculative/save_hidden_state.py",

tensorrt_llm/_torch/attention_backend/interface.py

@@ -440,7 +440,6 @@ def update_spec_dec_param(
             is_spec_dec_dynamic_tree,
             max_draft_len,
             max_total_draft_tokens,
-            model_is_wrapped: bool = False,
             spec_metadata: Optional['SpecMetadata'] = None,
             spec_tree_manager: Optional['SpecTreeManager'] = None,
             num_contexts: int = 0):

tensorrt_llm/_torch/attention_backend/sparse/dsa.py

@@ -1234,7 +1234,6 @@ def update_spec_dec_param(
         is_spec_dec_dynamic_tree,
         max_draft_len,
         max_total_draft_tokens,
-        model_is_wrapped: bool = False,
         spec_metadata: Optional['SpecMetadata'] = None,
         spec_tree_manager: Optional['SpecTreeManager'] = None,
         num_contexts: int = 0,
@@ -1246,7 +1245,6 @@ def update_spec_dec_param(
                                       is_spec_dec_dynamic_tree,
                                       max_draft_len,
                                       max_total_draft_tokens,
-                                      model_is_wrapped,
                                       spec_metadata,
                                       spec_tree_manager,
                                       num_contexts=num_contexts)

tensorrt_llm/_torch/attention_backend/trtllm.py

@@ -972,7 +972,6 @@ def update_spec_dec_param(
         is_spec_dec_dynamic_tree,
         max_draft_len,
         max_total_draft_tokens,
-        model_is_wrapped: bool = False,
         spec_metadata: Optional['SpecMetadata'] = None,
         spec_tree_manager: Optional['SpecTreeManager'] = None,
         num_contexts: int = 0,
@@ -986,7 +985,6 @@ def update_spec_dec_param(
             is_spec_dec_dynamic_tree: bool, whether using dynamic tree.
             max_draft_len: int, the number of the draft layers.
             max_total_draft_tokens: int, the number of all nodes in the tree (except the root).
-            model_is_wrapped: Optional[bool] = False, whether the drafter model is wrapped (i.e, CDL).
             spec_metadata: Optional['SpecMetadata'] = None, the metadata of the spec-dec.
             spec_tree_manager: Optional['SpecTreeManager'] = None, the spec_tree_manager for draft token tree.
             num_contexts: int = 0, the number of context (prefill) requests in the
@@ -1109,60 +1107,23 @@ def update_spec_dec_param(
                 self.spec_decoding_generation_lengths[:batch_size].fill_(n_dt)
                 cpp_query_len = n_dt
 
-            # Case 2/3: static tree
+            # Case 2: static tree (target model verification)
             elif self.is_spec_dec_tree and not self.is_spec_dec_dynamic_tree and spec_metadata is not None:
                 assert spec_metadata.spec_dec_mode.is_eagle3(
                 ), "Tree decoding is only supported for Eagle3 now"
 
-                is_target_model = not getattr(spec_metadata, 'is_draft_model',
-                                              False)
-
-                # Case 2: static tree and target model
-                if is_target_model:
-                    # For the target model, we update the spec-dec parameters with the spec_tree_manager, which is prepared in advance.
-                    self.spec_decoding_position_offsets[:batch_size, :].copy_(
-                        spec_tree_manager.spec_dec_position_offsets[0, :],
-                        non_blocking=True)
-                    self.spec_decoding_packed_mask[:batch_size, :, :].copy_(
-                        spec_tree_manager.spec_dec_packed_mask[0, :, :],
-                        non_blocking=True)
-                    self.spec_decoding_generation_lengths[:batch_size].fill_(
-                        spec_tree_manager.max_total_draft_tokens + 1)
-
-                # Case 3: static tree and the first drafter layer
-                else:
-                    assert model_is_wrapped == True, "The drafter model should be wrapped"
-                    # The first drafter layer will take the padded tokens as input (padding to the max_draft_len + 1)
-                    # But the spec-dec parameters are still in the shape of max_total_draft_tokens + 1.
-                    # Considering that these spec-dec params are accessed consecutively (without padding) in the attention Op,
-                    # we need to write them consecutively when setting them.
-                    # For the next drafter layers, we will prepare these spec-dec params in the drafting loops.
-
-                    # position_offsets
-                    position_offset = torch.arange(
-                        max_draft_len + 1,
-                        dtype=torch.int,
-                        device='cpu',
-                        pin_memory=prefer_pinned()).repeat(batch_size)
-                    self.spec_decoding_position_offsets.reshape(
-                        -1)[:(max_draft_len + 1) * batch_size].copy_(
-                            position_offset, non_blocking=True)
-                    # packed_mask
-                    dummy_idx = torch.arange(max_draft_len + 1)
-                    spec_decoding_packed_mask = torch.pow(
-                        2, dummy_idx + 1) - 1  # [max_draft_len + 1]
-                    spec_decoding_packed_mask = spec_decoding_packed_mask.repeat(
-                        batch_size)  # [batch_size * (max_draft_len + 1)]
-                    self.spec_decoding_packed_mask.reshape(
-                        -1)[:(max_draft_len + 1) * batch_size].copy_(
-                            spec_decoding_packed_mask, non_blocking=True)
-                    self.generate_spec_decoding_generation_length(
-                        runtime_draft_len=max_draft_len)
-                    if (self.spec_decoding_position_offsets is not None
-                            and self.spec_decoding_position_offsets.dim() == 1):
-                        cpp_query_len = max_draft_len + 1
-
-            # Case 4: linear tree
+                # Update the spec-dec parameters with the SpecTreeManager prepared
+                # in advance for target-model verification of the static tree.
+                self.spec_decoding_position_offsets[:batch_size, :].copy_(
+                    spec_tree_manager.spec_dec_position_offsets[0, :],
+                    non_blocking=True)
+                self.spec_decoding_packed_mask[:batch_size, :, :].copy_(
+                    spec_tree_manager.spec_dec_packed_mask[0, :, :],
+                    non_blocking=True)
+                self.spec_decoding_generation_lengths[:batch_size].fill_(
+                    spec_tree_manager.max_total_draft_tokens + 1)
+
+            # Case 3: linear tree
             else:
                 # Currently dynamic draft length is only supported for linear tree
                 # Dynamic draft length needs position offsets and packed mask to be shaped for each runtime draft length.

tensorrt_llm/_torch/auto_deploy/llm_args.py

@@ -116,24 +116,18 @@ def validate_supported_speculative_config(self):
             return self
 
         if isinstance(spec_config, MTPDecodingConfig):
-            if not spec_config.mtp_eagle_one_model or spec_config.use_mtp_vanilla:
+            if spec_config.use_mtp_vanilla:
                 raise ValueError(
-                    "AutoDeploy only supports MTP speculative decoding with "
-                    "mtp_eagle_one_model=True and use_mtp_vanilla=False "
-                    f"(got mtp_eagle_one_model={spec_config.mtp_eagle_one_model}, "
-                    f"use_mtp_vanilla={spec_config.use_mtp_vanilla})."
+                    "AutoDeploy only supports EAGLE-style MTP speculative "
+                    "decoding with use_mtp_vanilla=False "
+                    f"(got use_mtp_vanilla={spec_config.use_mtp_vanilla})."
                 )
         elif isinstance(spec_config, EagleDecodingConfig):
-            if not spec_config.eagle3_one_model:
-                raise ValueError(
-                    "AutoDeploy only supports Eagle speculative decoding with "
-                    f"eagle3_one_model=True (got eagle3_one_model={spec_config.eagle3_one_model})."
-                )
+            pass
         else:
             raise ValueError(
                 "AutoDeploy only supports speculative decoding via "
-                "MTPDecodingConfig(mtp_eagle_one_model=True) or "
-                "EagleDecodingConfig(eagle3_one_model=True)."
+                "MTPDecodingConfig or EagleDecodingConfig."
             )
 
         self.model_factory = "eagle_one_model"
@@ -148,8 +142,8 @@ def setup_hidden_state_capture(self):
         if isinstance(spec_config, MTPDecodingConfig):
             if spec_config.max_draft_len is None:
                 raise ValueError(
-                    "MTPDecodingConfig.max_draft_len must not be None when mtp_eagle_one_model is "
-                    "enabled. Ensure num_nextn_predict_layers is set in the model config."
+                    "MTPDecodingConfig.max_draft_len must not be None. "
+                    "Ensure num_nextn_predict_layers is set in the model config."
                 )
             capture_layers = {-1}
         else:

tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py

@@ -1100,10 +1100,9 @@ def instantiate_sampler(
     max_draft_len = 0 if spec_config is None else spec_config.max_draft_len
     max_total_draft_tokens = 0 if spec_config is None else spec_config.tokens_per_gen_step - 1
 
-    # One-model spec dec: model performs sampling internally, returns pre-computed tokens
+    # One-engine spec dec: model performs sampling internally, returns pre-computed tokens
     if spec_config is not None and (
-        spec_config.spec_dec_mode.is_eagle3_one_model()
-        or spec_config.spec_dec_mode.is_mtp_eagle_one_model()
+        spec_config.spec_dec_mode.is_eagle3() or spec_config.spec_dec_mode.is_mtp_eagle()
     ):
         sampler_args = TorchSampler.Args(
             max_seq_len=ad_config.max_seq_len,

tensorrt_llm/_torch/models/modeling_auto.py

@@ -34,12 +34,6 @@ def _resolve_class(config: ModelConfig) -> Optional[Type]:
                                             "")  # Strip the appended EAGLE3
             model_arch = "EAGLE3" + model_arch
 
-        if model_arch in (
-                "DeepseekV3ForCausalLM", "Glm4MoeForCausalLM",
-                "ExaoneMoEForCausalLM"
-        ) and config.spec_config is not None and config.spec_config.max_draft_len == 0:
-            model_arch = "MTPDraftModelForCausalLM"
-
         return MODEL_CLASS_MAPPING.get(model_arch)
 
     @staticmethod