fix

.github/workflows/test_cli_cuda_tensorrt_llm.yaml

@@ -44,7 +44,7 @@ jobs:
 
       - name: Install dependencies
         run: |
-          pip install -e .[testing,tesnsorrt-llm]
+          pip install -e .[testing]
 
       - name: Run tests
         run: |
@@ -84,7 +84,7 @@ jobs:
 
       - name: Install dependencies
         run: |
-          pip install -e .[testing,tesnsorrt-llm]
+          pip install -e .[testing]
 
       - name: Run tests (sequential)
         run: |

optimum_benchmark/backends/tensorrt_llm/backend.py

@@ -55,7 +55,6 @@ def download_pretrained_model(self) -> None:
     def prepare_generation_config(self) -> None:
         self.generation_config.eos_token_id = None
         self.generation_config.pad_token_id = None
-
         model_cache_folder = f"models/{self.config.model}".replace("/", "--")
         model_cache_path = f"{HUGGINGFACE_HUB_CACHE}/{model_cache_folder}"
         snapshot_file = f"{model_cache_path}/refs/{self.config.model_kwargs.get('revision', 'main')}"
@@ -91,7 +90,6 @@ def create_no_weights_model(self) -> None:
             self.logger.info("\t+ Modifying generation config for fixed length generation")
             self.generation_config.eos_token_id = None
             self.generation_config.pad_token_id = None
-            self.logger.info("\t+ Saving new pretrained generation config")
             self.generation_config.save_pretrained(save_directory=self.no_weights_model)
 
     def load_trtllm_with_no_weights(self) -> None:
@@ -128,6 +126,12 @@ def trtllm_kwargs(self):
         if self.config.gpus_per_node is not None:
             kwargs["gpus_per_node"] = self.config.gpus_per_node
 
+        if self.config.max_input_len is not None:
+            kwargs["max_input_len"] = self.config.max_input_len
+
+        if self.config.max_output_len is not None:
+            kwargs["max_output_len"] = self.config.max_output_len
+
         if self.config.max_batch_size is not None:
             kwargs["max_batch_size"] = self.config.max_batch_size
 

optimum_benchmark/backends/tensorrt_llm/config.py

@@ -22,6 +22,8 @@ class TRTLLMConfig(BackendConfig):
     use_fp8: Optional[bool] = None
     world_size: Optional[int] = None
     gpus_per_node: Optional[int] = None
+    max_input_len: Optional[int] = None
+    max_output_len: Optional[int] = None
     max_batch_size: Optional[int] = None
     max_new_tokens: Optional[int] = None
     max_prompt_length: Optional[int] = None