[TRTLLM-6291] feat: Add user-provided speculative decoding support (#5204)

Signed-off-by: Robin Kobus <[email protected]>

Author: Funatiq

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -400,6 +400,7 @@ def teardown_managers(self, resources: Dict) -> None:
 
 
 def create_py_executor_instance(
+        *,
         dist,
         resources,
         mapping,

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -343,6 +343,7 @@ class PyTorchModelEngine(ModelEngine):
 
     def __init__(
         self,
+        *,
         model_path: str,
         pytorch_backend_config: PyTorchConfig,
         batch_size: int = 8,

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -168,13 +168,13 @@ def __init__(self,
                  sampler: Sampler,
                  dist: Distributed,
                  max_num_sequences: int,
-                 drafter: Drafter = None,
+                 drafter: Optional[Drafter] = None,
                  disable_overlap_scheduler: bool = False,
                  max_input_len: int = 2048,
                  max_batch_size: int = 8,
                  max_beam_width: int = 1,
                  max_draft_tokens: int = 0,
-                 kv_cache_transceiver: KvCacheTransceiver = None,
+                 kv_cache_transceiver: Optional[KvCacheTransceiver] = None,
                  draft_model_engine: Optional[ModelEngine] = None,
                  garbage_collection_gen0_threshold: Optional[int] = None,
                  start_worker: bool = True):
@@ -922,8 +922,7 @@ def _executor_loop(self):
                         self._prepare_draft_tokens(scheduled_batch)
 
                     if self.drafter is not None:
-                        self.drafter.prepare_draft_tokens(
-                            scheduled_batch, sample_state)
+                        self.drafter.prepare_draft_tokens(scheduled_batch)
 
                     if self.kv_cache_transceiver:
                         # For generation requests which have completed KV cache transfer

tensorrt_llm/_torch/pyexecutor/py_executor_creator.py

@@ -227,8 +227,8 @@ def create_py_executor(
     with mem_monitor.observe_creation_stage(
             _ExecutorCreationStage.MODEL_ENGINE_MAIN):
         model_engine = PyTorchModelEngine(
-            checkpoint_dir,
-            pytorch_backend_config,
+            model_path=checkpoint_dir,
+            pytorch_backend_config=pytorch_backend_config,
             batch_size=executor_config.max_batch_size,
             max_beam_width=executor_config.max_beam_width,
             max_num_tokens=executor_config.max_num_tokens,
@@ -250,8 +250,8 @@ def create_py_executor(
             draft_spec_config.max_draft_tokens = 0
 
             draft_model_engine = PyTorchModelEngine(
-                spec_config.draft_model_path,
-                pytorch_backend_config,
+                model_path=spec_config.draft_model_path,
+                pytorch_backend_config=pytorch_backend_config,
                 batch_size=executor_config.max_batch_size,
                 max_beam_width=executor_config.max_beam_width,
                 max_num_tokens=executor_config.max_num_tokens,
@@ -358,24 +358,36 @@ def create_py_executor(
                 if estimating_kv_cache else _ExecutorCreationStage.KV_CACHE):
             kv_cache_creator.build_managers(resources)
 
+    # Drafter for speculative decoding
+    with mem_monitor.observe_creation_stage(_ExecutorCreationStage.DRAFTER):
+        drafter = get_spec_drafter(model_engine)
+
     # Resource managers for speculative decoding
     spec_resource_manager = get_spec_resource_manager(model_engine,
-                                                      draft_model_engine)
+                                                      draft_model_engine,
+                                                      drafter)
     if spec_resource_manager is not None:
         resources[
             ResourceManagerType.SPEC_RESOURCE_MANAGER] = spec_resource_manager
 
-    # Drafter for speculative decoding
-    with mem_monitor.observe_creation_stage(_ExecutorCreationStage.DRAFTER):
-        drafter = get_spec_drafter(model_engine, spec_resource_manager)
-
     with mem_monitor.observe_creation_stage(
             _ExecutorCreationStage.INIT_EXTRA_RESOURCES
             if estimating_kv_cache else _ExecutorCreationStage.EXTRA_RESOURCES):
         py_executor = create_py_executor_instance(
-            dist, resources, mapping, pytorch_backend_config, executor_config,
-            ctx_chunk_config, model_engine, draft_model_engine, False, sampler,
-            drafter, lora_config, garbage_collection_gen0_threshold)
+            dist=dist,
+            resources=resources,
+            mapping=mapping,
+            pytorch_backend_config=pytorch_backend_config,
+            executor_config=executor_config,
+            ctx_chunk_config=ctx_chunk_config,
+            model_engine=model_engine,
+            draft_model_engine=draft_model_engine,
+            start_worker=False,
+            sampler=sampler,
+            drafter=drafter,
+            lora_config=lora_config,
+            garbage_collection_gen0_threshold=garbage_collection_gen0_threshold,
+        )
 
     if estimating_kv_cache:
         assert kv_cache_creator is not None
@@ -404,10 +416,21 @@ def create_py_executor(
         with mem_monitor.observe_creation_stage(
                 _ExecutorCreationStage.EXTRA_RESOURCES):
             py_executor = create_py_executor_instance(
-                dist, resources, mapping, pytorch_backend_config,
-                executor_config, ctx_chunk_config, model_engine,
-                draft_model_engine, False, sampler, drafter, lora_config,
-                garbage_collection_gen0_threshold)
+                dist=dist,
+                resources=resources,
+                mapping=mapping,
+                pytorch_backend_config=pytorch_backend_config,
+                executor_config=executor_config,
+                ctx_chunk_config=ctx_chunk_config,
+                model_engine=model_engine,
+                draft_model_engine=draft_model_engine,
+                start_worker=False,
+                sampler=sampler,
+                drafter=drafter,
+                lora_config=lora_config,
+                garbage_collection_gen0_threshold=
+                garbage_collection_gen0_threshold,
+            )
 
     py_executor.start_worker()
     return py_executor

tensorrt_llm/_torch/speculative/__init__.py

@@ -3,6 +3,7 @@
 from .interface import SpecConfig, SpecMetadata
 from .mtp import MTPConfig, MTPEagleWorker, MTPSpecMetadata, MTPWorker
 from .ngram import NGramConfig, NGramDrafter, NGramPoolManager
+from .user_provided import UserProvidedConfig
 from .utils import (get_num_spec_layers, get_spec_decoder, get_spec_drafter,
                     get_spec_metadata, get_spec_resource_manager,
                     get_spec_worker)
@@ -20,6 +21,7 @@
     "NGramPoolManager",
     "SpecConfig",
     "SpecMetadata",
+    "UserProvidedConfig",
     "get_num_spec_layers",
     "get_spec_decoder",
     "get_spec_drafter",

tensorrt_llm/_torch/speculative/drafter.py

@@ -2,7 +2,6 @@
 from typing import Optional
 
 from ..pyexecutor.resource_manager import BaseResourceManager
-from ..pyexecutor.sampler import SampleState
 from ..pyexecutor.scheduler import ScheduledRequests
 
 
@@ -18,7 +17,6 @@ def __init__(
     def prepare_draft_tokens(
         self,
         scheduled_requests: ScheduledRequests,
-        state: SampleState,
     ) -> None:
         """
         Prepare the drafter tokens for the forward computation this step.

tensorrt_llm/_torch/speculative/interface.py

@@ -17,6 +17,7 @@ class SpeculativeDecodingMode(IntEnum):
     EAGLE3_ONE_MODEL = auto()
     NGRAM = auto()
     DRAFT_TARGET = auto()
+    USER_PROVIDED = auto()
     NONE = auto()
 
     def is_mtp(self):
@@ -37,6 +38,9 @@ def is_eagle3_one_model(self):
     def is_ngram(self):
         return self == SpeculativeDecodingMode.NGRAM
 
+    def is_user_provided(self):
+        return self == SpeculativeDecodingMode.USER_PROVIDED
+
     def is_none(self):
         return self == SpeculativeDecodingMode.NONE
 
@@ -74,7 +78,7 @@ def has_spec_decoder(self):
         return self.is_mtp() or self.is_eagle3() or self.is_eagle3_one_model()
 
     def has_spec_drafter(self):
-        return self.is_ngram()
+        return self.is_ngram() or self.is_user_provided()
 
     def extend_ctx(self, attention_backend: Type[AttentionBackend]):
         """
@@ -86,7 +90,8 @@ def extend_ctx(self, attention_backend: Type[AttentionBackend]):
         # Fixme: only trtllm attention backend supports eagle3 generation-phase kernels on blackwell.
         return ((self.is_eagle3() or self.is_draft_target())
                 and not (isinstance(attention_backend, TrtllmAttention)
-                         and get_sm_version() == 100)) or self.is_ngram()
+                         and get_sm_version() == 100)
+                ) or self.is_ngram() or self.is_user_provided()
 
     def attention_need_spec_dec_mode(self):
         """
@@ -185,6 +190,9 @@ class SpecMetadata:
     # if spec-dec tree wouldn't be changed at all, the mask won't be computed every step.
     is_spec_dec_dynamic_tree: bool = False
 
+    def __post_init__(self):
+        pass
+
     def prepare(self):
         """
         Hook to be called before the forward step of the model.

tensorrt_llm/_torch/speculative/ngram.py

@@ -7,10 +7,9 @@
 
 from ..pyexecutor.llm_request import *
 from ..pyexecutor.resource_manager import BaseResourceManager
-from ..pyexecutor.sampler import SampleState
 from ..pyexecutor.scheduler import ScheduledRequests
 from .drafter import Drafter
-from .interface import SpecConfig, SpecMetadata, SpeculativeDecodingMode
+from .interface import SpecConfig, SpeculativeDecodingMode
 
 
 @dataclass
@@ -40,16 +39,6 @@ def update_from_model_config(self, model_config):
         pass
 
 
-@dataclass
-class NGramSpecMetadata(SpecMetadata):
-    """
-    Metadata for NGram.
-    """
-
-    def __post_init__(self) -> None:
-        return
-
-
 class NGramPoolManager(BaseResourceManager):
     """
     Drafter for NGram. This class maintains the pattern-matches pairs for NGram drafter.
@@ -212,12 +201,8 @@ def __init__(
     def prepare_draft_tokens(
         self,
         scheduled_requests: ScheduledRequests,
-        state: SampleState,
     ) -> None:
 
-        if state is None:  # Skip the first step
-            return
-
         for request in sorted(scheduled_requests.generation_requests,
                               key=lambda r: r.py_batch_idx):
             # Add new token to a copy of the generated tokens to find new daft tokens

tensorrt_llm/_torch/speculative/user_provided.py

@@ -0,0 +1,26 @@
+from dataclasses import dataclass
+from typing import Optional
+
+from tensorrt_llm._torch.speculative.drafter import Drafter
+
+from .interface import SpecConfig, SpeculativeDecodingMode
+
+
+@dataclass
+class UserProvidedConfig(SpecConfig):
+    """
+    Configuration for user provided speculative decoding.
+    """
+    # The name of speculative decoding.
+    spec_dec_name = "USER_PROVIDED"
+
+    num_extra_kv_tokens: int = 0
+    max_draft_tokens: int = 0
+    drafter: Optional[Drafter] = None
+
+    def __post_init__(self) -> None:
+        self.spec_dec_mode = SpeculativeDecodingMode.from_string(
+            self.spec_dec_name)
+
+    def update_from_model_config(self, model_config):
+        pass

tensorrt_llm/_torch/speculative/utils.py

@@ -1,13 +1,13 @@
 from tensorrt_llm._torch.pyexecutor.sampler import TorchSampler
-from tensorrt_llm._torch.speculative.interface import SpecConfig
+from tensorrt_llm._torch.speculative.interface import SpecConfig, SpecMetadata
 
 from .draft_target import DraftTargetSpecMetadata
 from .eagle3 import (Eagle3OneModelSampler, Eagle3OneModelSpecMetadata,
                      Eagle3OneModelWorker, Eagle3ResourceManager,
                      Eagle3SpecMetadata)
 from .mtp import (MTPEagleWorker, MTPHiddenStatesManager, MTPSampler,
                   MTPSpecMetadata, MTPWorker)
-from .ngram import NGramDrafter, NGramPoolManager, NGramSpecMetadata
+from .ngram import NGramDrafter, NGramPoolManager
 
 
 def get_spec_metadata(spec_config,
@@ -50,16 +50,19 @@ def get_spec_metadata(spec_config,
             spec_dec_mode=spec_config.spec_dec_mode,
             max_num_requests=max_num_requests,
         )
-    if spec_config.spec_dec_mode.is_ngram():
-        return NGramSpecMetadata(
+    if spec_config.spec_dec_mode.is_ngram(
+    ) or spec_config.spec_dec_mode.is_user_provided():
+        return SpecMetadata(
             max_draft_tokens=spec_config.max_draft_tokens,
             spec_dec_mode=spec_config.spec_dec_mode,
             max_num_requests=max_num_requests,
         )
     return None
 
 
-def get_spec_resource_manager(model_engine, draft_model_engine=None):
+def get_spec_resource_manager(model_engine,
+                              draft_model_engine=None,
+                              drafter=None):
     spec_config = model_engine.spec_config
     if spec_config is None:
         return None
@@ -95,8 +98,9 @@ def get_spec_resource_manager(model_engine, draft_model_engine=None):
             max_seq_len,
             max_num_tokens,
         )
-    if spec_dec_mode.is_ngram():
-        return NGramPoolManager(spec_config, max_num_requests)
+    if spec_dec_mode.is_ngram() or spec_dec_mode.is_user_provided():
+        assert drafter is not None, "Drafter is required for ngram or user provided speculative decoding."
+        return drafter.spec_resource_manager
     return None
 
 
@@ -113,12 +117,16 @@ def get_spec_decoder(sampler_args: TorchSampler.Args, spec_config: SpecConfig):
         f"Unsupported speculative decoding mode: {spec_config.spec_dec_mode}")
 
 
-def get_spec_drafter(model_engine, spec_resource_manager=None):
+def get_spec_drafter(model_engine):
     spec_config = model_engine.spec_config
+    max_num_requests = model_engine.batch_size
     if spec_config is None:
         return None
     if spec_config.spec_dec_mode.is_ngram():
-        return NGramDrafter(spec_config, spec_resource_manager)
+        return NGramDrafter(spec_config,
+                            NGramPoolManager(spec_config, max_num_requests))
+    if spec_config.spec_dec_mode.is_user_provided():
+        return spec_config.drafter
     return None