[None][chore] Remove dead code in model engine

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -1768,8 +1768,7 @@ def _create_kv_cache_manager(
             layer_mask=spec_dec_layer_mask,
             is_disagg=self._is_disagg,
             disable_overlap_scheduler=self._disable_overlap_scheduler,
-            kv_events_config=None
-            if estimating_kv_cache or model_engine.is_draft_model else
+            kv_events_config=None if estimating_kv_cache else
             self._llm_args.kv_cache_config.kv_events_config,
             cold_page_codec_provider=cold_page_codec_provider,
             joint_kv_cache_reuse=self._joint_kv_cache_reuse,
@@ -2754,7 +2753,7 @@ def _create_kv_cache_manager(
         # Optional overrides for one-model draft case (when model_engine is None)
         model_config: Optional[ModelConfig] = None,
         dtype: Optional[torch.dtype] = None,
-        is_draft: Optional[bool] = None,
+        is_draft: bool = False,
         layer_mask: Optional[List[bool]] = None,
         num_layers: Optional[int] = None,
         num_kv_heads: Optional[Union[int, List[int]]] = None,
@@ -2802,9 +2801,6 @@ def _create_kv_cache_manager(
     if dtype is None:
         dtype = model_engine.dtype
 
-    if is_draft is None:
-        is_draft = model_engine.is_draft_model
-
     if kv_cache_type is None:
         kv_cache_type = tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF
 

tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py

@@ -18,7 +18,6 @@
 from ..memory_buffer_utils import Buffers, get_memory_buffers
 from ..modules.multi_stream_utils import with_multi_stream
 from ..moe.expert_statistic import ExpertStatistic
-from ..speculative.eagle3 import Eagle3ResourceManager
 from ..speculative.interface import SpecMetadata
 from ..speculative.spec_sampler_base import SampleStateTensorsSpec
 from ..speculative.utils import get_draft_kv_cache_manager
@@ -127,7 +126,6 @@ class CUDAGraphRunnerConfig:
     use_mrope: bool
     original_max_draft_len: int
     original_max_total_draft_tokens: int
-    is_draft_model: bool
     enable_attention_dp: bool
     is_encoder_decoder: bool
     batch_size: int
@@ -379,52 +377,37 @@ def get_graph_key(
                        else self._resolve_sample_type(
                            batch, promoted_context_request_ids))
 
-        if self.config.is_draft_model and spec_resource_manager is not None and isinstance(
-                spec_resource_manager, Eagle3ResourceManager):
-            # If 'is_first_draft' is True, even with tree decoding, the length of draft_len will only be 'max_draft_len', not 'max_total_draft_token'.
-            # Because we will pad the input to 'max_draft_len' length for the first draft layer.
-            draft_len = self.config.original_max_draft_len if spec_resource_manager.is_first_draft else 0
-            key = KeyType(batch_size=batch_size,
-                          draft_len=draft_len,
-                          is_first_draft=spec_resource_manager.is_first_draft,
-                          short_seq_len_mode=short_seq_len_mode,
-                          is_all_greedy_sample=is_all_greedy_sample,
-                          sample_type=sample_type,
-                          peft_cache_data_type=peft_cache_data_type,
-                          use_lora_graph=use_lora_graph)
-        else:
-            # With dynamic spec decode, the draft length may be zero even when enable_spec_decode is True,
-            # so we need to get the draft length from the batch instead of using enable_spec_decode.
-            draft_len_list = []
-            for request in batch.generation_requests:
-                draft_len_list.append(len(request.py_draft_tokens))
-            draft_len = max(draft_len_list)
-            assert len(
-                set(draft_len_list)) == 1, "All draft lengths must be the same"
-            context_requests = batch.context_requests
-            num_contexts = len(context_requests)
-            context_query_len = 0
-            if num_contexts:
-                context_query_len = int(context_requests[0].context_chunk_size)
-                if any(
-                        int(request.context_chunk_size) != context_query_len
-                        for request in context_requests[1:]):
-                    return None
-            num_encoder_tokens = sum(
-                int(request.encoder_output_len) for request in context_requests
-                if not request.py_skip_cross_kv_projection)
-            key = KeyType(batch_size=batch_size,
-                          draft_len=draft_len,
-                          is_first_draft=False,
-                          short_seq_len_mode=short_seq_len_mode,
-                          is_all_greedy_sample=is_all_greedy_sample,
-                          sample_type=sample_type,
-                          num_contexts=num_contexts,
-                          context_query_len=context_query_len,
-                          num_encoder_tokens=num_encoder_tokens,
-                          peft_cache_data_type=peft_cache_data_type,
-                          use_lora_graph=use_lora_graph)
-        return key
+        # With dynamic spec decode, the draft length may be zero even when enable_spec_decode is True,
+        # so we need to get the draft length from the batch instead of using enable_spec_decode.
+        draft_len_list = []
+        for request in batch.generation_requests:
+            draft_len_list.append(len(request.py_draft_tokens))
+        draft_len = max(draft_len_list)
+        assert len(
+            set(draft_len_list)) == 1, "All draft lengths must be the same"
+        context_requests = batch.context_requests
+        num_contexts = len(context_requests)
+        context_query_len = 0
+        if num_contexts:
+            context_query_len = int(context_requests[0].context_chunk_size)
+            if any(
+                    int(request.context_chunk_size) != context_query_len
+                    for request in context_requests[1:]):
+                return None
+        num_encoder_tokens = sum(
+            int(request.encoder_output_len) for request in context_requests
+            if not request.py_skip_cross_kv_projection)
+        return KeyType(batch_size=batch_size,
+                       draft_len=draft_len,
+                       is_first_draft=False,
+                       short_seq_len_mode=short_seq_len_mode,
+                       is_all_greedy_sample=is_all_greedy_sample,
+                       sample_type=sample_type,
+                       num_contexts=num_contexts,
+                       context_query_len=context_query_len,
+                       num_encoder_tokens=num_encoder_tokens,
+                       peft_cache_data_type=peft_cache_data_type,
+                       use_lora_graph=use_lora_graph)
 
     def set_capture_sample_type(self,
                                 sample_type: Optional[SampleType]) -> None:

tensorrt_llm/_torch/pyexecutor/engine/metadata.py

@@ -99,7 +99,6 @@ def update_spec_metadata(
     *,
     runtime_draft_len: int,
     runtime_tokens_per_gen_step: int,
-    is_draft_model: bool,
     attention_backend: type[AttentionBackend],
     original_max_draft_len: int,
     original_max_total_draft_tokens: int,
@@ -109,10 +108,7 @@ def update_spec_metadata(
     spec_metadata.runtime_draft_len = runtime_draft_len
     spec_metadata.runtime_tokens_per_gen_step = runtime_tokens_per_gen_step
 
-    is_spec_dec_mode = spec_metadata.spec_dec_mode.attention_need_spec_dec_mode(
-        is_draft_model,
-        attention_backend,
-    )
+    is_spec_dec_mode = spec_metadata.spec_dec_mode.attention_need_spec_dec_mode(attention_backend)
     # Parallel-draft modes advertise their full generation width rather than a
     # conventional draft length, so attention needs the total-token capacity.
     if spec_metadata.spec_dec_mode.is_parallel_draft():
@@ -122,7 +118,7 @@ def update_spec_metadata(
         max_draft_len = original_max_draft_len
         max_total_draft_tokens = spec_dec_max_total_draft_tokens
 
-    if spec_tree_manager is not None and not is_draft_model:
+    if spec_tree_manager is not None:
         spec_tree_manager.slot_storage.fill_all_slot_ids(
             scheduled_requests.context_requests,
             scheduled_requests.generation_requests,

tensorrt_llm/_torch/pyexecutor/engine/runners/encoder.py

@@ -80,7 +80,6 @@ def create(
         attention_runtime_features: AttentionRuntimeFeatures,
         enable_autotuner: bool,
         is_encoder_decoder: bool,
-        draft_model: bool,
     ) -> Self:
         """Resolve encoder settings and construct the concrete runner config."""
         batch_sizes = list(graph_config.batch_sizes or []) if graph_config is not None else []
@@ -107,12 +106,7 @@ def create(
         feature_shape = None
         feature_dtype = None
         fixed_seq_len = None
-        if (
-            graph_config is not None
-            and not draft_model
-            and is_encoder_decoder
-            and model_graph_spec is not None
-        ):
+        if graph_config is not None and is_encoder_decoder and model_graph_spec is not None:
             if mapping.tp_size > 1:
                 logger.warning(
                     "Feature-mode encoder CUDA graphs require TP=1; the encoder phase stays eager."
@@ -226,7 +220,6 @@ def create(
         attention_backend: type[AttentionBackend],
         attention_runtime_features: AttentionRuntimeFeatures,
         enable_autotuner: bool,
-        draft_model: bool,
     ) -> Self:
         return super().create(
             model=model,
@@ -241,7 +234,6 @@ def create(
             attention_runtime_features=attention_runtime_features,
             enable_autotuner=enable_autotuner,
             is_encoder_decoder=False,
-            draft_model=draft_model,
         )
 
 

tensorrt_llm/_torch/pyexecutor/engine/runners/no_kv_cache.py

@@ -45,7 +45,6 @@ class NoKVCacheRunnerConfig(RunnerConfig):
     prefill_cuda_graph_num_tokens: list[int]
     mm_encoder_cache_enabled: bool
     spec_config: DecodingBaseConfig | None
-    is_draft_model: bool
     num_seq_slots: int | None
     original_max_draft_len: int
     original_max_total_draft_tokens: int
@@ -107,7 +106,6 @@ def setup_spec_metadata(
             runner_config.max_batch_size,
             max_num_tokens=runner_config.max_num_tokens,
             spec_resource_manager=spec_resource_manager,
-            is_draft_model=runner_config.is_draft_model,
             max_seq_len=runner_config.max_seq_len,
             num_seq_slots=runner_config.num_seq_slots,
         )
@@ -121,7 +119,6 @@ def setup_spec_metadata(
             runtime_tokens_per_gen_step=spec_config.get_runtime_tokens_per_gen_step(
                 runtime_draft_len
             ),
-            is_draft_model=runner_config.is_draft_model,
             attention_backend=runner_config.attention_backend,
             original_max_draft_len=runner_config.original_max_draft_len,
             original_max_total_draft_tokens=(runner_config.original_max_total_draft_tokens),