[TRTLLM-6352][feat] Migrate EAGLE3 and draft/target speculation to Drafter

Signed-off-by: ziyixiong-nv <[email protected]>

Author: ziyixiong-nv

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -30,7 +30,7 @@
 from tensorrt_llm.logger import logger
 
 from ..distributed import Distributed
-from ..speculative.drafter import Drafter
+from ..speculative.drafter import Drafter, create_drafter
 from .kv_cache_transceiver import KvCacheTransceiver
 from .llm_request import (ExecutorRequest, LlmRequest, LlmRequestState,
                           LlmResponse, executor_request_to_llm_request)
@@ -959,8 +959,20 @@ def _executor_loop(self):
                             scheduled_batch)
 
                     self.resource_manager.prepare_resources(scheduled_batch)
-                    if self.draft_model_engine is not None:
-                        self._prepare_draft_tokens(scheduled_batch)
+                    if self.draft_model_engine is not None and self.drafter is None:
+                        spec_resource_manager = self.resource_manager.get_resource_manager(
+                            ResourceManagerType.SPEC_RESOURCE_MANAGER)
+                        self.drafter = create_drafter(
+                            spec_decoding_mode=self.model_engine.spec_config.
+                            spec_dec_mode,
+                            spec_config=self.model_engine.spec_config,
+                            draft_model_engine=self.draft_model_engine,
+                            max_draft_tokens=self.max_draft_tokens,
+                            draft_seq_slot_manager=self.draft_seq_slot_manager,
+                            sampler=self.sampler,
+                            resource_manager=self.resource_manager,
+                            spec_resource_manager=spec_resource_manager,
+                        )
 
                     if self.drafter is not None:
                         self.drafter.prepare_draft_tokens(scheduled_batch)

tensorrt_llm/_torch/speculative/drafter.py

@@ -1,9 +1,11 @@
 from abc import ABC, abstractmethod
 
 from ..pyexecutor.scheduler import ScheduledRequests
+from .interface import SpeculativeDecodingMode
 
 
 class Drafter(ABC):
+    """Abstract base class for all drafter implementations."""
 
     @abstractmethod
     def prepare_draft_tokens(
@@ -12,5 +14,36 @@ def prepare_draft_tokens(
     ) -> None:
         """
         Prepare the drafter tokens for the forward computation this step.
+
+        Args:
+            scheduled_requests: The scheduled requests for this iteration
         """
         raise NotImplementedError
+
+
+def create_drafter(spec_decoding_mode: SpeculativeDecodingMode,
+                   **kwargs) -> Drafter:
+    """
+    Factory function to create the appropriate drafter based on the mode.
+
+    Args:
+        spec_decoding_mode: The speculative decoding mode
+        **kwargs: Additional arguments for drafter construction
+
+    Returns:
+        Drafter: The appropriate drafter instance
+
+    Raises:
+        ValueError: If the speculative decoding mode is not supported
+    """
+    match spec_decoding_mode:
+        case SpeculativeDecodingMode.NGRAM:
+            from .ngram import NGramDrafter
+            return NGramDrafter(**kwargs)
+        case SpeculativeDecodingMode.EAGLE3 | SpeculativeDecodingMode.DRAFT_TARGET:
+            # Import here to avoid circular import
+            from .model_drafter import ModelDrafter
+            return ModelDrafter(**kwargs)
+        case _:
+            raise ValueError(
+                f"Unsupported speculative decoding mode: {spec_decoding_mode}")

tensorrt_llm/_torch/speculative/eagle3.py

@@ -117,29 +117,40 @@ def prepare(self):
         # hidden state space before the target model forward.
         start_idx = 0
         if not self.is_draft_model:
-            for req_id, seq_len in zip(self.request_ids, self.seq_lens):
-                slot_id = self.eagle3_resource_manager.slot_manager.get_slot(
-                    req_id)
-                self.eagle3_resource_manager.start_indices[slot_id] = start_idx
-                start_idx += seq_len
+            if self.request_ids is not None and self.seq_lens is not None:
+                for req_id, seq_len in zip(self.request_ids, self.seq_lens):
+                    slot_id = self.eagle3_resource_manager.slot_manager.get_slot(
+                        req_id
+                    ) if self.eagle3_resource_manager is not None else None
+                    if self.eagle3_resource_manager is not None and slot_id is not None:
+                        self.eagle3_resource_manager.start_indices[
+                            slot_id] = start_idx
+                        start_idx += seq_len
         # Prepare hidden states gather ids
         hidden_states_read_indices = []
         hidden_states_write_indices = []
-        for req_id, seq_len in zip(self.request_ids, self.seq_lens):
-            slot_id = self.eagle3_resource_manager.slot_manager.get_slot(req_id)
-            start_idx = self.eagle3_resource_manager.start_indices[slot_id]
-            # If this is the first draft or the target model forward, we need to
-            # read/write all of the hidden states, otherwise, only read the last token
-            if is_first_draft or not self.is_draft_model:
-                hidden_states_read_indices.extend(
-                    list(range(start_idx, start_idx + seq_len)))
-                hidden_states_write_indices.extend(
-                    list(range(start_idx, start_idx + seq_len)))
-            else:
-                old_seq_len = self.eagle3_resource_manager.seq_lens[slot_id]
-                hidden_states_read_indices.append(start_idx + old_seq_len - 1)
-                hidden_states_write_indices.append(start_idx + seq_len - 1)
-            self.eagle3_resource_manager.seq_lens[slot_id] = seq_len
+        if self.request_ids is not None and self.seq_lens is not None:
+            for req_id, seq_len in zip(self.request_ids, self.seq_lens):
+                if self.eagle3_resource_manager is not None:
+                    slot_id = self.eagle3_resource_manager.slot_manager.get_slot(
+                        req_id)
+                    start_idx = self.eagle3_resource_manager.start_indices[
+                        slot_id]
+                    # If this is the first draft or the target model forward, we need to
+                    # read/write all of the hidden states, otherwise, only read the last token
+                    if is_first_draft or not self.is_draft_model:
+                        hidden_states_read_indices.extend(
+                            list(range(start_idx, start_idx + seq_len)))
+                        hidden_states_write_indices.extend(
+                            list(range(start_idx, start_idx + seq_len)))
+                    else:
+                        old_seq_len = self.eagle3_resource_manager.seq_lens[
+                            slot_id]
+                        hidden_states_read_indices.append(start_idx +
+                                                          old_seq_len - 1)
+                        hidden_states_write_indices.append(start_idx + seq_len -
+                                                           1)
+                    self.eagle3_resource_manager.seq_lens[slot_id] = seq_len
         # Prepare hidden states gather ids
         self.hidden_states_read_indices_host = torch.tensor(
             hidden_states_read_indices, dtype=torch.long, pin_memory=True)