Initial implementation

Signed-off-by: jthomson04 <[email protected]>

Author: jthomson04

components/src/dynamo/common/utils/input_params.py

@@ -0,0 +1,25 @@
+class InputParamManager:
+    def __init__(self, tokenizer):
+        self.tokenizer = tokenizer
+
+    def get_input_param(self, request: dict, use_tokenizer: bool):
+        """
+        Get the input parameter for the request.
+        """
+
+        if use_tokenizer:
+            if self.tokenizer is None:
+                raise ValueError("Tokenizer is not available")
+
+            if "messages" in request:
+                return self.tokenizer.apply_chat_template(
+                    request["messages"], tokenize=False, add_generation_prompt=True
+                )
+            elif "prompt" in request:
+                return request["prompt"]
+            elif "text" in request:
+                return request["text"]
+            else:
+                raise ValueError("No input parameter found in request")
+
+        return request.get("token_ids")

components/src/dynamo/sglang/request_handlers/handler_base.py

@@ -13,6 +13,7 @@
 from sglang.srt.utils import get_local_ip_auto
 
 from dynamo._core import Client, Component, Context
+from dynamo.common.utils.input_params import InputParamManager
 from dynamo.sglang.args import Config
 from dynamo.sglang.publisher import DynamoSglangPublisher
 
@@ -50,6 +51,12 @@ def __init__(
         self.serving_mode = config.serving_mode
         self.skip_tokenizer_init = config.server_args.skip_tokenizer_init
 
+        self.input_param_manager = InputParamManager(
+            self.engine.tokenizer_manager.tokenizer
+            if not self.skip_tokenizer_init
+            else None
+        )
+
     @abstractmethod
     async def generate(self, request: Dict[str, Any], context: Context):
         """Generate response from request.
@@ -68,23 +75,9 @@ def cleanup(self) -> None:
         pass
 
     def _get_input_param(self, request: Dict[str, Any]) -> Dict[str, Any]:
-        """Get the appropriate input parameter for SGLang engine.
-
-        Args:
-            request: Request dict with token_ids or messages.
-
-        Returns:
-            Dict with either input_ids or prompt for engine.
-        """
-        if self.skip_tokenizer_init:
-            return {"input_ids": request["token_ids"]}
-        else:
-            # use sglang's chat templating itself but leave tokenization to the
-            # interal engine's TokenizerManager
-            prompt = self.engine.tokenizer_manager.tokenizer.apply_chat_template(
-                request["messages"], tokenize=False, add_generation_prompt=True
-            )
-            return {"prompt": prompt}
+        return self.input_param_manager.get_input_param(
+            request, use_tokenizer=not self.skip_tokenizer_init
+        )
 
     @staticmethod
     def _generate_bootstrap_room() -> int:

components/src/dynamo/trtllm/main.py

@@ -175,13 +175,25 @@ async def init(runtime: DistributedRuntime, config: Config):
         dynamic_batch_config=dynamic_batch_config,
     )
     modality = getattr(config, "modality", None) or "text"
+    if config.use_trtllm_tokenizer:
+        logging.info(
+            "Using TensorRT-LLM's built in tokenizer. Setting skip_tokenizer_init to False"
+        )
+        skip_tokenizer_init = False
+    else:
+        logging.info(
+            "Using dynamo's built in tokenizer. Setting skip_tokenizer_init to True"
+        )
+        skip_tokenizer_init = True
+
     arg_map = {
         "model": model_path,
         "scheduler_config": scheduler_config,
         "tensor_parallel_size": config.tensor_parallel_size,
         "pipeline_parallel_size": config.pipeline_parallel_size,
         "moe_expert_parallel_size": config.expert_parallel_size,
         "backend": Backend.PYTORCH,
+        "skip_tokenizer_init": skip_tokenizer_init,
         "build_config": build_config,
         "kv_cache_config": kv_cache_config,
         "gpus_per_node": gpus_per_node,
@@ -245,6 +257,8 @@ async def init(runtime: DistributedRuntime, config: Config):
     if hasattr(default_sampling_params, "return_perf_metrics"):
         default_sampling_params.return_perf_metrics = True
     model_input = ModelInput.Tokens
+    if config.use_trtllm_tokenizer:
+        model_input = ModelInput.Text
 
     # Set model type based on disaggregation mode for unified frontend support
     if config.disaggregation_mode == DisaggregationMode.PREFILL:
@@ -275,8 +289,11 @@ async def init(runtime: DistributedRuntime, config: Config):
         )
 
     else:
-        # We already detokenize inside HandlerBase. No need to also do it in TRTLLM.
-        default_sampling_params.detokenize = False
+        if config.use_trtllm_tokenizer:
+            default_sampling_params.detokenize = True
+        else:
+            # We already detokenize inside HandlerBase. No need to also do it in TRTLLM.
+            default_sampling_params.detokenize = False
 
     connector = None
     logging.info("Initializing NIXL Connect.")