Updating "triton_cli" in order to support TensorRT-LLM 1.1.0 (#124)

Co-authored-by: Yingge He <157551214+yinggeh@users.noreply.github.com>

Author: mc-nv

src/triton_cli/parser.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-# Copyright 2023-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
 # modification, are permitted provided that the following conditions
@@ -66,7 +66,7 @@
     # Public
     "gpt2": "hf:gpt2",
     "opt125m": "hf:facebook/opt-125m",
-    "mistral-7b": "hf:mistralai/Mistral-7B-v0.1",
+    "mistral-7b": "hf:mistralai/Mistral-7B-Instruct-v0.1",
     "falcon-7b": "hf:tiiuae/falcon-7b",
 }
 

src/triton_cli/repository.py

@@ -1,4 +1,4 @@
-# Copyright 2023-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
 # modification, are permitted provided that the following conditions
@@ -323,13 +323,23 @@ def __generate_vllm_model(self, huggingface_id: str):
 
     def __generate_ngc_model(self, name: str, source: str):
         engines_path = ENGINE_DEST_PATH + "/" + source
+        # Find the actual engine directory that contains config.json
+        actual_engine_dir = self.__find_engine_directory(engines_path)
         parse_and_substitute(
-            str(self.repo), name, engines_path, engines_path, "auto", dry_run=False
+            str(self.repo),
+            name,
+            actual_engine_dir,
+            actual_engine_dir,
+            "auto",
+            dry_run=False,
         )
 
     def __generate_trtllm_model(self, name: str, huggingface_id: str):
         engines_path = ENGINE_DEST_PATH + "/" + name
-        engines = [engine for engine in Path(engines_path).glob("*.engine")]
+        # Search for engine files recursively since they might be in subdirectories
+        engines = list(Path(engines_path).glob("*.engine")) + list(
+            Path(engines_path).glob("*/*.engine")
+        )
         if engines:
             logger.warning(
                 f"Found existing engine(s) at {engines_path}, skipping build."
@@ -343,21 +353,64 @@ def __generate_trtllm_model(self, name: str, huggingface_id: str):
             p.start()
             p.join()
 
+        # Find the actual engine directory that contains config.json
+        # When using workspace parameter, TRT-LLM creates a subdirectory
+        actual_engine_dir = self.__find_engine_directory(engines_path)
+
         # NOTE: In every case, the TRT LLM template should be filled in with values.
         # If the model exists, the CLI will raise an exception when creating the model repo.
         # If a user clears the model repo, they won't need to re-build the engines,
         # but they will still need to modify the TRT LLM template.
         parse_and_substitute(
             triton_model_dir=str(self.repo),
             bls_model_name=name,
-            engine_dir=engines_path,
-            token_dir=engines_path,
+            engine_dir=actual_engine_dir,
+            token_dir=actual_engine_dir,
             token_type="auto",
             dry_run=False,
         )
 
+    def __find_engine_directory(self, workspace_path: str) -> str:
+        """
+        Find the actual engine directory that contains config.json.
+        When using the workspace parameter, TRT-LLM creates a subdirectory structure.
+        This method searches for config.json and returns its parent directory.
+        """
+        workspace_path = Path(workspace_path)
+
+        # First check if config.json exists directly in the workspace path
+        if (workspace_path / "config.json").exists():
+            return str(workspace_path)
+
+        # Search for config.json in subdirectories (up to 2 levels deep)
+        for config_file in workspace_path.glob("*/config.json"):
+            logger.info(f"Found engine directory at {config_file.parent}")
+            return str(config_file.parent)
+
+        for config_file in workspace_path.glob("*/*/config.json"):
+            logger.info(f"Found engine directory at {config_file.parent}")
+            return str(config_file.parent)
+
+        # If no config.json found, return the original path and let the error surface
+        logger.warning(
+            f"Could not find config.json in {workspace_path} or its subdirectories. "
+            f"Returning original path."
+        )
+        return str(workspace_path)
+
     def __build_trtllm_engine(self, huggingface_id: str, engines_path: Path):
-        from tensorrt_llm import LLM, BuildConfig
+        # Ensure engines_path is a Path object
+        engines_path = Path(engines_path)
+
+        # Import from _tensorrt_engine to force TensorRT backend (not PyTorch)
+        # The PyTorch backend doesn't support workspace parameter
+        try:
+            from tensorrt_llm._tensorrt_engine import LLM
+        except ImportError:
+            # Fallback to regular import for newer versions
+            from tensorrt_llm import LLM
+
+        from tensorrt_llm import BuildConfig
 
         # NOTE: Given config.json, can read from 'build_config' section and from_dict
         config = BuildConfig()
@@ -367,13 +420,31 @@ def __build_trtllm_engine(self, huggingface_id: str, engines_path: Path):
         # config.max_seq_len = 8192
         # config.max_batch_size = 256
 
-        engine = LLM(huggingface_id, build_config=config)
-        # TODO: Investigate if LLM is internally saving a copy to a temp dir
-        engine.save(str(engines_path))
+        # Create the workspace directory if it doesn't exist
+        # TensorRT-LLM will create a temp subdir inside this workspace
+        engines_path.mkdir(parents=True, exist_ok=True)
+
+        # Build engine to target directory using workspace parameter (TensorRT backend only)
+        engine = LLM(huggingface_id, build_config=config, workspace=str(engines_path))
+
+        # For newer API versions with save() method, call it to ensure engine is properly saved
+        # In older versions, the workspace parameter should have already placed the engine there
+        if hasattr(engine, "save") and callable(getattr(engine, "save")):
+            try:
+                engine.save(str(engines_path))
+                logger.debug(
+                    f"Called save() method to ensure engine is at {engines_path}"
+                )
+            except Exception as e:
+                # If save fails, workspace parameter should have already placed it correctly
+                logger.debug(
+                    f"save() call failed (engine may already be in workspace): {e}"
+                )
 
         # The new trtllm(v0.17.0+) requires explicit calling shutdown to shutdown
         # the mpi blocking thread, or the engine process won't exit
-        engine.shutdown()
+        if hasattr(engine, "shutdown") and callable(getattr(engine, "shutdown")):
+            engine.shutdown()
 
     def __create_model_repository(
         self, name: str, version: int = 1, backend: str = None

src/triton_cli/server/server_utils.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 
-# Copyright 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -210,9 +210,11 @@ def _parse_world_size(self) -> int:
             engine(s).
         """
         assert self._is_trtllm_model, "World size cannot be parsed from a model repository that does not contain a TRT LLM model."
+        engine_path = self._get_engine_path(self._trtllm_model_config_path)
+        engine_config_path = engine_path / "config.json"
+
+        # Try to read config.json if it exists, otherwise use defaults
         try:
-            engine_path = self._get_engine_path(self._trtllm_model_config_path)
-            engine_config_path = engine_path / "config.json"
             with open(engine_config_path) as json_data:
                 data = json.load(json_data)
                 # FIXME: Revert handling using 'build_config' as the key when gpt migrates to using unified builder
@@ -226,8 +228,13 @@ def _parse_world_size(self) -> int:
                 tp = int(config.get("tensor_parallel", 1))
                 pp = int(config.get("pipeline_parallel", 1))
                 return tp * pp
-        except OSError:
-            raise Exception(f"Unable to open {engine_config_path}")
+        except FileNotFoundError:
+            # If config.json doesn't exist, use default world size
+            # Default tensor_parallel=1, pipeline_parallel=1, so world_size=1
+            logger.warning(
+                f"{engine_config_path} not found, using default world_size=1"
+            )
+            return 1
 
 
 class VLLMUtils:

src/triton_cli/templates/trt_llm/tensorrt_llm/config.pbtxt

@@ -1,4 +1,4 @@
-# Copyright 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
 # modification, are permitted provided that the following conditions
@@ -756,3 +756,9 @@ parameters: {
     string_value: "${guided_decoding_backend}"
   }
 }
+parameters: {
+  key: "xgrammar_tokenizer_info_path"
+  value: {
+    string_value: "${xgrammar_tokenizer_info_path}"
+  }
+}

src/triton_cli/templates/trt_llm/tensorrt_llm_bls/config.pbtxt

@@ -1,4 +1,4 @@
-# Copyright 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
 # modification, are permitted provided that the following conditions
@@ -299,6 +299,18 @@ input [
     data_type: TYPE_STRING
     dims: [ 1 ]
     optional: true
+  },
+  {
+    name: "return_num_input_tokens"
+    data_type: TYPE_BOOL
+    dims: [ 1 ]
+    optional: true
+  },
+  {
+    name: "return_num_output_tokens"
+    data_type: TYPE_BOOL
+    dims: [ 1 ]
+    optional: true
   }
 ]
 output [
@@ -351,6 +363,16 @@ output [
     name: "kv_cache_alloc_total_blocks"
     data_type: TYPE_INT32
     dims: [ 1 ]
+  },
+  {
+    name: "num_input_tokens"
+    data_type: TYPE_UINT32
+    dims: [ 1 ]
+  },
+  {
+    name: "num_output_tokens"
+    data_type: TYPE_UINT32
+    dims: [ 1 ]
   }
 ]
 

src/triton_cli/trt_llm/engine_config_parser.py

@@ -1,4 +1,4 @@
-# Copyright 2023-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
 # modification, are permitted provided that the following conditions
@@ -35,22 +35,33 @@ def parse_and_substitute(
     triton_model_dir, bls_model_name, engine_dir, token_dir, token_type, dry_run
 ):
     json_path = engine_dir + "/config.json"
-    with open(json_path) as j:
-        config_file = json.load(j)
+
+    # Try to read config.json if it exists, otherwise use defaults
+    config_file = None
+    try:
+        with open(json_path) as j:
+            config_file = json.load(j)
+    except FileNotFoundError:
+        print(f"Warning: {json_path} not found, using default configuration")
 
     config_dict = {}
     # These fields will cause parsing issues when parsing model config if not
     # replaced, so replace with sensible defaults.
 
-    # FIXME: Revert handling using 'build_config' as the key when gpt migrates to using unified builder
-    build_config_key = (
-        "builder_config"
-        if config_file.get("builder_config") is not None
-        else "build_config"
-    )
-    config_dict["triton_max_batch_size"] = config_file[build_config_key][
-        "max_batch_size"
-    ]
+    # Get max_batch_size from config.json if available, otherwise use default
+    if config_file:
+        # FIXME: Revert handling using 'build_config' as the key when gpt migrates to using unified builder
+        build_config_key = (
+            "builder_config"
+            if config_file.get("builder_config") is not None
+            else "build_config"
+        )
+        config_dict["triton_max_batch_size"] = config_file[build_config_key][
+            "max_batch_size"
+        ]
+    else:
+        # Default max_batch_size when config.json is not available
+        config_dict["triton_max_batch_size"] = 256
 
     config_dict["logits_datatype"] = "TYPE_FP32"
     config_dict["triton_backend"] = "tensorrtllm"  # or python
@@ -66,6 +77,9 @@ def parse_and_substitute(
     config_dict["engine_dir"] = engine_dir
     config_dict["tokenizer_dir"] = token_dir
     config_dict["tokenizer_type"] = token_type
+    # Disable guided decoding by default (xgrammar requires additional configuration)
+    config_dict["guided_decoding_backend"] = ""
+    config_dict["xgrammar_tokenizer_info_path"] = ""
 
     config_dict["max_queue_delay_microseconds"] = 0
     # Default echo = False