[None][chore] Remove two-model speculative decoding leftovers

docs/source/features/sampling.md

@@ -204,8 +204,7 @@ the sampling parameters and the random seed and offset per row. A seeded request
 therefore depends only on its own seed and how far it has decoded, not on where it sits in the
 batch. The sampled tokens can still differ across batch sizes, because the kernel picks its
 algorithm by the number of rows. The kernel computes the complete set of token sampling
-probabilities (after top-k / top-p masking etc.) only when the user requests them or speculative
-decoding (rejection sampling) requires them.
+probabilities (after top-k / top-p masking etc.) only when the user requests them.
 
 Moreover, Torch Sampler internally batches requests with compatible sampling parameters. This
 can greatly reduce the overall latency of the sampling step when request batches are comprised

docs/source/llm-api/index.md

@@ -127,7 +127,6 @@ The `py_executor` property contains the timed scopes in `create_py_executor()`,
 |-----------------|-------|
 | `config_and_checkpoint_loader_initialization_seconds` | Initialize the model config and checkpoint loader. |
 | `model_engine_creation_seconds` | Construct the main model engine, which includes loading model weights. |
-| `draft_model_engine_creation_seconds` | Construct the separate draft model engine in the deprecated two-model MTP setting. |
 | `guided_decoder_creation_seconds` | Construct guided-decoding resources. |
 | `sampler_creation_seconds` | Construct the sampler. |
 | `initial_kv_cache_creation_seconds` | Construct the temporary KV cache used to estimate final capacity. |
@@ -145,8 +144,6 @@ The two model-engine properties contain times for various `PyTorchModelEngine` w
 `py_executor.initial_py_executor_creation_seconds_for_kv_cache_estimation` timing.
 Likewise, `final_model_engine` measures the timings for the final model engine creation as part of
 the `py_executor.final_py_executor_creation_seconds` timing.
-The deprecated two-model speculative decoding setting may produce the additional properties
-`initial_draft_model_engine` and `final_draft_model_engine`.
 
 Each model-engine property contains the warmup scopes below. Which scopes appear depends on the model
 and configuration; for example, encoder-decoder and context-parallel configurations can skip some
@@ -174,7 +171,7 @@ The `model_loader` property contains timings for loading the main LLM weights, i
 `checkpoint_finalization_seconds` for checkpoint-session synchronization and cleanup after weight population.
 If a draft model is used, additional fields `draft_checkpoint_preparation_seconds`,
 `draft_weight_population_seconds`, and `draft_checkpoint_finalization_seconds` will appear.
-A `draft_model_loader` property can also appear in the deprecated two-model MTP setting.
+A `draft_model_loader` property appears when a user-provided drafter exposes a `draft_model_engine`.
 The `py_executor.model_engine_creation_seconds` timing includes the total `model_loader` timing.
 
 | Metric | Description |

pyproject.toml

@@ -877,7 +877,6 @@ files = [
     "tensorrt_llm/_torch/pyexecutor/sampler/logprobs.py",
     "tensorrt_llm/_torch/pyexecutor/sampler/sampler_features.py",
     "tensorrt_llm/_torch/pyexecutor/sampler/seed_manager.py",
-    "tensorrt_llm/_torch/pyexecutor/sampler/two_model_spec_dec.py",
     "tensorrt_llm/_torch/pyexecutor/sampler/finish_reasons.py",
     "tensorrt_llm/_torch/pyexecutor/sampler/top_p_decay.py",
     "tensorrt_llm/_torch/pyexecutor/sampler/penalties.py",

tensorrt_llm/_torch/models/_arch_index.py

@@ -244,7 +244,7 @@ def is_builtin_zoo_module(module_name: str) -> bool:
 # model builder (registered via ``@register_draft_model``). Keyed by the enum
 # member *name* rather than the enum itself so this module keeps importing
 # nothing. Modes absent from this table have no one-engine draft model to build
-# (two-model / drafter-loop modes such as NGRAM, SA and USER_PROVIDED).
+# (model-free or drafter-loop modes such as NGRAM, SA and USER_PROVIDED).
 #
 # Adding a speculative decoding mode
 # ----------------------------------

tensorrt_llm/_torch/models/modeling_gpt_oss.py

@@ -564,10 +564,7 @@ def get_preferred_kv_cache_manager_version(cls,
         statically dividing memory between them.
 
         The preference is adopted only when the user leaves
-        ``kv_cache_config.use_kv_cache_manager_v2`` at ``"auto"``. Two-model
-        speculative decoding demotes it to V1 because V2 sizes both the target
-        and draft KV cache managers from the full budget; an explicit ``True``
-        is rejected by ``llm_utils._resolve_kv_cache_manager_v2_auto``.
+        ``kv_cache_config.use_kv_cache_manager_v2`` at ``"auto"``.
         """
         return "V2"
 

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -901,7 +901,6 @@ def __init__(
         self,
         *,
         model_engine: PyTorchModelEngine,
-        draft_model_engine: Optional[PyTorchModelEngine],
         mapping: Mapping,
         net_max_seq_len: int,
         kv_connector_manager: Optional[KvCacheConnectorManager],
@@ -921,7 +920,6 @@ def __init__(
         skip_est: bool = False,
     ):
         self._model_engine = model_engine
-        self._draft_model_engine = draft_model_engine
         self._mapping = mapping
         self._kv_cache_config = kv_cache_config
         self._max_kv_tokens_in = self._kv_cache_config.max_tokens
@@ -1254,17 +1252,6 @@ def scaled(cost: CacheCost) -> CacheCost:
             # fields would silently drop it.
             return dataclasses.replace(cost, slope=cost.slope * helix_cp_scale)
 
-        if self._draft_model_engine is not None:
-            draft_model_config = self._draft_model_engine.model.model_config
-            draft_kv_cache_manager_cls = self._get_model_kv_cache_manager_cls(
-                self._draft_model_engine, kv_cache_config)
-            return scaled(
-                self._per_manager_cache_cost(
-                    draft_kv_cache_manager_cls,
-                    draft_model_config,
-                    kv_cache_config,
-                    mapping=draft_mapping,
-                    include_request_cost=include_request_cost))
         if use_separate_draft_kv_cache:
             # One-model draft with separate KV cache layout.
             # Pass num_layers explicitly since the HF config may report a
@@ -1829,8 +1816,7 @@ def configure_kv_cache_capacity(self,
                 # get kv cache stats for both model and draft model
                 kv_stats = py_executor.resource_manager.resource_managers.get(
                     ResourceManagerType.KV_CACHE_MANAGER).get_kv_cache_stats()
-                # Get draft KV cache stats if present (either from two-model mode or one-model
-                # mode with separate draft KV cache)
+                # Get draft KV cache stats if present (one-model mode with a separate draft KV cache)
                 draft_kv_cache_manager = py_executor.resource_manager.resource_managers.get(
                     ResourceManagerType.DRAFT_KV_CACHE_MANAGER)
                 kv_stats_draft = draft_kv_cache_manager.get_kv_cache_stats(
@@ -2550,7 +2536,6 @@ def _bound_unsplit_v2_draft_by_tokens(
         successful split (which already sizes both managers) is unaffected.
         """
         if (not self._is_kv_cache_manager_v2
-                or self._draft_model_engine is not None
                 or not self._should_create_separate_draft_kv_cache()
                 or target_kv_cache_config.max_tokens is not None):
             return draft_kv_cache_config
@@ -2994,9 +2979,8 @@ def build_managers(self,
             self_kv_cache_config, cross_kv_cache_config = self._split_kv_cache_budget_for_cross(
             )
 
-        has_separate_one_model_draft = (
-            self._draft_model_engine is None
-            and self._should_create_separate_draft_kv_cache())
+        has_separate_one_model_draft = self._should_create_separate_draft_kv_cache(
+        )
         # One term per axis: a draft pool exists to pair with, block reuse is on,
         # and the span is known. Anything outside keeps its unpaired path.
         self._joint_kv_cache_reuse = (has_separate_one_model_draft and
@@ -3014,11 +2998,8 @@ def build_managers(self,
                 cross_kv_cache_config)
 
         # Split combined KV cache budgets before creating managers.
-        has_draft = (
-            self._draft_model_engine is not None  # two-model
-            or has_separate_one_model_draft)  # one-model
         draft_kv_cache_config = None
-        if has_draft:
+        if has_separate_one_model_draft:
             # The GPU split applies when each manager sizes its pools from
             # max_gpu_total_bytes (V2 and V1 VSWA). V1 non-VSWA and estimation
             # size GPU pools from a shared max_tokens instead.
@@ -3061,32 +3042,13 @@ def build_managers(self,
         if not estimating_kv_cache and kv_cache_manager is not None:
             kv_cache_manager.fp8_ctx_mla_kv_len_cap = self._fp8_ctx_mla_kv_len_cap
 
-        if (not estimating_kv_cache and self._kv_connector_manager is not None
-                and self._draft_model_engine is not None):
-            raise NotImplementedError(
-                "Connector manager is not supported for draft model.")
-
         draft_kv_cache_manager = None
         draft_build_kv_cache_config = (draft_kv_cache_config
                                        if draft_kv_cache_config is not None else
                                        self_kv_cache_config)
 
-        # Two-model speculative decoding: draft model has separate engine
-        if self._draft_model_engine is not None:
-            if (self._is_kv_cache_manager_v2
-                    and draft_kv_cache_config is not None):
-                # Offload budgets are divided per manager, GPU budgets are not.
-                assert (draft_kv_cache_config.max_gpu_total_bytes ==
-                        self_kv_cache_config.max_gpu_total_bytes), (
-                            "KVCacheManagerV2 does not support two-model "
-                            "speculative decoding with separate draft GPU "
-                            "budgets.")
-            draft_kv_cache_manager = self._create_kv_cache_manager(
-                self._draft_model_engine,
-                estimating_kv_cache,
-                kv_cache_config_override=draft_build_kv_cache_config)
         # One-model speculative decoding with different KV layouts
-        elif self._should_create_separate_draft_kv_cache():
+        if has_separate_one_model_draft:
             draft_kv_cache_manager = self._create_one_model_draft_kv_cache_manager(
                 original_max_seq_len,
                 estimating_kv_cache,

tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py

@@ -1092,8 +1092,6 @@ def _get_or_create_padding_dummy(
                 cross_kv_cache_manager)
 
         # Get draft KV cache manager only for one-model speculative decoding.
-        # In two-model mode, each model has its own KV cache manager, so
-        # draft_kv_cache_manager should be None.
         draft_kv_cache_manager = get_draft_kv_cache_manager(
             self.spec_config, resource_manager)
 

tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py

@@ -2270,10 +2270,9 @@ def _run_capture_pass(
         # branch's variants would spend warmup on graphs whose keys no batch can
         # ever produce.
         #
-        # The predicate is use_one_engine() rather than has_spec_decoder(): the
-        # two-model eagle3 / mtp_eagle modes also "have a spec decoder", but
-        # get_spec_decoder hands them a plain TorchSampler, so they belong on
-        # the TorchSampler branch and do need the FAST graphs.
+        # The predicate is use_one_engine(): the drafter-based modes (NGram,
+        # user-provided) get a plain TorchSampler from get_spec_decoder, so they
+        # belong on the TorchSampler branch and do need the FAST graphs.
         uses_own_spec_decoder = (
             self._config.spec_config is not None
             and self._config.spec_config.spec_dec_mode.use_one_engine()

tensorrt_llm/_torch/pyexecutor/guided_decoder.py

@@ -86,8 +86,7 @@ def from_llm_request(cls, request: LlmRequest):
         return cls(
             guided_decoding_params=request.py_guided_decoding_params,
             request_id=request.py_request_id,
-            seq_slot=(request.py_target_seq_slot
-                      if request.py_is_draft else request.py_seq_slot),
+            seq_slot=request.py_seq_slot,
             prev_seq_slot=request.py_batch_idx,
             is_context_init_state=request.is_context_init_state,
             is_last_context_chunk=(request.is_context_init_state
@@ -100,7 +99,6 @@ def from_llm_request(cls, request: LlmRequest):
                 and request.py_decoding_iter == 1
                 and request.py_batch_idx is None),
             new_token=request.get_last_tokens(0),
-            is_draft=request.py_is_draft,
             draft_tokens=request.py_draft_tokens,
             num_accepted_draft_tokens=request.py_num_accepted_draft_tokens)
 

tensorrt_llm/_torch/pyexecutor/llm_request.py

@@ -857,9 +857,7 @@ def __init__(
             stop_words_list: list[list[int]] | None = None,
             llm_request: Optional[
                 tensorrt_llm.bindings.internal.batch_manager.LlmRequest] = None,
-            is_draft: bool = False,
             seq_slot: Optional[int] = None,
-            target_seq_slot: Optional[int] = None,
             num_logprobs: int = 0,
             use_chunked_generation_logits: bool = True,
             logits_chunk_size: int = 8,
@@ -983,11 +981,6 @@ def __init__(
         self.py_additional_outputs = additional_outputs
 
         self.py_beam_width = cast(int, self.sampling_config.beam_width)
-        self.py_is_draft = is_draft
-        # If the request is a draft request, target_seq_slot is the sequence slot ID of its target request.
-        self.py_target_seq_slot = target_seq_slot
-        self.use_draft_model = is_draft
-        self.d2t = None
         self.py_disable_speculative_decoding = False
 
         # Chunked logits parameters