[TRTLLM-14388][refactor] Remove 2 model spec dec drafting loops

.pre-commit-config.yaml

@@ -319,7 +319,6 @@ 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 |
@@ -1086,7 +1085,6 @@ 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 |

legacy-files.txt

@@ -311,7 +311,6 @@ 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

pyproject.toml

@@ -368,7 +368,6 @@ 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",

ruff-legacy.toml

@@ -328,7 +328,6 @@ 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/attention_backend/interface.py

@@ -445,7 +445,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/metadata.py

@@ -851,7 +851,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,
@@ -864,7 +863,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/sparse/rocket/cache_manager.py

@@ -123,7 +123,6 @@ def add_dummy_requests(
         use_mrope: bool = False,
         max_beam_width: int = 1,
         encoder_output_lens: Optional[List[int]] = None,
-        num_extra_decoding_steps: int = 0,
         draft_kv_cache_manager=None,
     ):
         requests = super().add_dummy_requests(
@@ -136,7 +135,6 @@ def add_dummy_requests(
             use_mrope=use_mrope,
             max_beam_width=max_beam_width,
             encoder_output_lens=encoder_output_lens,
-            num_extra_decoding_steps=num_extra_decoding_steps,
             draft_kv_cache_manager=draft_kv_cache_manager,
         )
         if prepare_resource:

tensorrt_llm/_torch/attention_backend/trtllm.py

@@ -1187,7 +1187,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,
@@ -1201,7 +1200,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
@@ -1329,61 +1327,26 @@ 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 only)
             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()
                         or spec_metadata.spec_dec_mode.is_eagle3_one_model()
                         ), "Tree decoding is only supported for Eagle3 now"
+                assert not getattr(spec_metadata, 'is_draft_model', False), (
+                    "Static tree spec-dec params are only prepared for the target model"
+                )
 
-                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
+                # 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: 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/pyexecutor/kv_cache_manager_v2.py

@@ -3357,7 +3357,6 @@ def add_dummy_requests(
         use_mrope: bool = False,
         max_beam_width: int = 1,
         encoder_output_lens: Optional[List[int]] = None,
-        num_extra_decoding_steps: int = 0,
         draft_kv_cache_manager: Optional["BaseResourceManager"] = None,
     ):
         _kv_draft = (
@@ -3430,7 +3429,7 @@ def release_resources(
                     release_resources(req)
                     return None
                 kv_cache.stop_committing()
-                dummy_capacity = token_num + self.num_extra_kv_tokens + num_extra_decoding_steps
+                dummy_capacity = token_num + self.num_extra_kv_tokens
                 if is_gen and not materialize_history:
                     kv_cache.enable_swa_scratch_reuse = False
                 # Need to hint the committed history to activate stale-block

tensorrt_llm/_torch/pyexecutor/llm_request.py

@@ -1029,7 +1029,6 @@ def __init__(
         # Whether the request is for the first forward of the draft model.
         self.py_is_first_draft = is_first_draft
         self.d2t = None
-        self.py_draft_use_greedy_sampling = False
         self.py_disable_speculative_decoding = False
 
         # Chunked logits parameters

tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py

@@ -2488,11 +2488,6 @@ def add_dummy_requests(
         use_mrope: bool = False,
         max_beam_width: int = 1,
         encoder_output_lens: Optional[List[int]] = None,
-        # For capturable drafting loops. During normal inference, the draft model always
-        # has enough KV cache space to fit all of our draft tokens. During warmup, however,
-        # we need to make the KV cache manager aware that multiple autoregressive steps will
-        # occur.
-        num_extra_decoding_steps: int = 0,
         draft_kv_cache_manager: Optional[KVCacheManager] = None,
     ) -> List[LlmRequest]:
         requests = super().add_dummy_requests(
@@ -2505,7 +2500,6 @@ def add_dummy_requests(
             use_mrope=use_mrope,
             max_beam_width=max_beam_width,
             encoder_output_lens=encoder_output_lens,
-            num_extra_decoding_steps=num_extra_decoding_steps,
             draft_kv_cache_manager=draft_kv_cache_manager,
         )
         if requests:
@@ -3551,7 +3545,6 @@ def add_dummy_requests(
         use_mrope: bool = False,
         max_beam_width: int = 1,
         encoder_output_lens: Optional[List[int]] = None,
-        num_extra_decoding_steps: int = 0,
         draft_kv_cache_manager: Optional[BaseResourceManager] = None,
     ) -> List[LlmRequest]:
         requests = super().add_dummy_requests(
@@ -3564,7 +3557,6 @@ def add_dummy_requests(
             use_mrope=use_mrope,
             max_beam_width=max_beam_width,
             encoder_output_lens=encoder_output_lens,
-            num_extra_decoding_steps=num_extra_decoding_steps,
             draft_kv_cache_manager=draft_kv_cache_manager,
         )
         if requests and prepare_resource: