[refactor] Clean up drafter/resource manager creation logic

Signed-off-by: Mike Iovine <[email protected]>

Author: mikeiovine

tensorrt_llm/_torch/pyexecutor/py_executor_creator.py

@@ -358,18 +358,19 @@ 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
+    # For user-specified drafters, use extra_resource_managers in PyTorchBackend config
+    # to provide a resource manager if required.
     spec_resource_manager = get_spec_resource_manager(model_engine,
-                                                      draft_model_engine,
-                                                      drafter)
+                                                      draft_model_engine)
     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):

tensorrt_llm/_torch/speculative/drafter.py

@@ -1,18 +1,10 @@
 from abc import ABC, abstractmethod
-from typing import Optional
 
-from ..pyexecutor.resource_manager import BaseResourceManager
 from ..pyexecutor.scheduler import ScheduledRequests
 
 
 class Drafter(ABC):
 
-    def __init__(
-        self,
-        spec_resource_manager: Optional[BaseResourceManager] = None,
-    ):
-        self.spec_resource_manager = spec_resource_manager
-
     @abstractmethod
     def prepare_draft_tokens(
         self,

tensorrt_llm/_torch/speculative/ngram.py

@@ -194,8 +194,9 @@ def __init__(
         spec_config: SpecConfig,
         ngram_pool_manager: NGramPoolManager = None,
     ):
+        super().__init__()
         assert ngram_pool_manager is not None, "NGram needs a resource manager to maintain the pool."
-        super().__init__(spec_resource_manager=ngram_pool_manager)
+        self.spec_resource_manager = ngram_pool_manager
         self.max_num_draft_tokens = spec_config.max_draft_tokens
 
     def prepare_draft_tokens(

tensorrt_llm/_torch/speculative/user_provided.py

@@ -1,6 +1,7 @@
 from dataclasses import dataclass
 from typing import Optional
 
+from tensorrt_llm._torch.pyexecutor.resource_manager import BaseResourceManager
 from tensorrt_llm._torch.speculative.drafter import Drafter
 
 from .interface import SpecConfig, SpeculativeDecodingMode
@@ -17,10 +18,17 @@ class UserProvidedConfig(SpecConfig):
     num_extra_kv_tokens: int = 0
     max_draft_tokens: int = 0
     drafter: Optional[Drafter] = None
+    # For convenience, this will default to drafter.spec_resource_manager if such
+    # an attribute exists and resource_manager has not been explicitly specified.
+    resource_manager: Optional[BaseResourceManager] = None
 
     def __post_init__(self) -> None:
         self.spec_dec_mode = SpeculativeDecodingMode.from_string(
             self.spec_dec_name)
 
+        if self.resource_manager is None and self.drafter is not None and hasattr(
+                self.drafter, "spec_resource_manager"):
+            self.resource_manager = self.drafter.spec_resource_manager
+
     def update_from_model_config(self, model_config):
         pass

tensorrt_llm/_torch/speculative/utils.py

@@ -60,9 +60,7 @@ def get_spec_metadata(spec_config,
     return None
 
 
-def get_spec_resource_manager(model_engine,
-                              draft_model_engine=None,
-                              drafter=None):
+def get_spec_resource_manager(model_engine, draft_model_engine=None):
     spec_config = model_engine.spec_config
     if spec_config is None:
         return None
@@ -98,9 +96,10 @@ def get_spec_resource_manager(model_engine,
             max_seq_len,
             max_num_tokens,
         )
-    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
+    if spec_dec_mode.is_ngram():
+        return NGramPoolManager(spec_config, max_num_requests)
+    if spec_dec_mode.is_user_provided():
+        return spec_config.resource_manager
     return None
 
 
@@ -117,16 +116,13 @@ 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):
+def get_spec_drafter(model_engine, spec_resource_manager):
     spec_config = model_engine.spec_config
-    max_num_requests = model_engine.batch_size
+    model_engine.batch_size
     if spec_config is None:
         return None
     if spec_config.spec_dec_mode.is_ngram():
-        return NGramDrafter(spec_config,
-                            NGramPoolManager(spec_config, max_num_requests))
-    if spec_config.spec_dec_mode.is_user_provided():
-        return spec_config.drafter
+        return NGramDrafter(spec_config, spec_resource_manager)
     return None
 
 

tensorrt_llm/llmapi/llm_args.py

@@ -278,8 +278,9 @@ def from_dict(cls, data: dict):
 
 
 class UserProvidedDecodingConfig(DecodingBaseConfig):
-    # Type should be Drafter, but it leads to circular import
-    drafter: object
+    # Cannot use real type annotations due to circular imports
+    drafter: object  # Type is Drafter
+    resource_manager: object = None  # Type is Optional[ResourceManager]
 
     @classmethod
     def from_dict(cls, data: dict):
@@ -1401,7 +1402,8 @@ def validate_speculative_config(self):
                 from tensorrt_llm._torch.speculative import UserProvidedConfig
                 self.speculative_config = UserProvidedConfig(
                     max_draft_tokens=self.speculative_config.max_draft_len,
-                    drafter=self.speculative_config.drafter)
+                    drafter=self.speculative_config.drafter,
+                    resource_manager=self.speculative_config.resource_manager)
                 self.build_config.speculative_decoding_mode = SpeculativeDecodingMode.USER_PROVIDED
                 self.build_config.max_draft_len = self.speculative_config.max_draft_tokens
             else: