[TRTLLM] always setting request output length (#1258)

deepjavalibrary · Nov 3, 2023 · 1d51b1f · 1d51b1f
1 parent bac045d
commit 1d51b1f
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/engines/python/setup/djl_python/rolling_batch/trtllm_rolling_batch.py b/engines/python/setup/djl_python/rolling_batch/trtllm_rolling_batch.py
@@ -38,8 +38,8 @@ def reset(self):
         super().reset()
 
     def translate_triton_params(self, parameters):
-        if "max_new_tokens" in parameters.keys():
-            parameters["request_output_len"] = parameters.pop("max_new_tokens")
+        parameters["max_new_tokens"] = parameters.get("max_new_tokens", 128)
+        parameters["request_output_len"] = parameters.pop("max_new_tokens")
         if "top_k" in parameters.keys():
             parameters["runtime_top_k"] = parameters.pop("top_k")
         if "top_p" in parameters.keys():