[https://nvbugspro.nvidia.com/bug/5243740][fix] deduce default max_tokens for trtllm-serve (#4265)

* Deduce default max_tokens for trtllm-serve

Signed-off-by: Pengyun Lin <[email protected]>

* Improve executor_config.max_seq_len assignment in TRT workflow

Signed-off-by: Pengyun Lin <[email protected]>

* Enhance error message

Signed-off-by: Pengyun Lin <[email protected]>

* Add deduced max_tokens test

Signed-off-by: Pengyun Lin <[email protected]>

---------

Signed-off-by: Pengyun Lin <[email protected]>

Author: LinPoly

tensorrt_llm/executor/worker.py

@@ -394,11 +394,34 @@ def _enqueue_request(self, request: GenerationRequest) -> int:
                 )
 
         assert request.id is not None
+
+        def _deduce_max_tokens(request: GenerationRequest,
+                               executor_config: tllm.ExecutorConfig) -> int:
+            if request.sampling_params.max_tokens:
+                return request.sampling_params.max_tokens
+            # deduce max_tokens when it's not set by user
+            query_token_len = len(
+                request.query_token_ids) if request.query_token_ids else 0
+            cp_size = 1 if (not hasattr(executor_config, "mapping")
+                            or executor_config.mapping.cp_size
+                            is None) else executor_config.mapping.cp_size
+            if not hasattr(executor_config, "max_seq_len"):
+                raise RuntimeError(
+                    "max_tokens for sampling is not set and cannot be deduced")
+            splited_prompt_len = int(len(prompt_token_ids) / cp_size)
+            default_max_tokens = executor_config.max_seq_len - splited_prompt_len - query_token_len
+            if default_max_tokens < 0:
+                raise ValueError(
+                    f"Deduced max_tokens {default_max_tokens} is less than 0, because"
+                    f"prompt length {splited_prompt_len} plus query length {query_token_len} "
+                    f"is larger than max_seq_len {executor_config.max_seq_len}")
+            return default_max_tokens
+
         try:
             executor_request = tllm.Request(
                 client_id=request.id,
                 input_token_ids=prompt_token_ids,
-                max_tokens=request.sampling_params.max_tokens,
+                max_tokens=_deduce_max_tokens(request, self._executor_config),
                 streaming=request.streaming,
                 sampling_config=request.sampling_params._get_sampling_config(),
                 end_id=-1 if request.sampling_params.ignore_eos else

tensorrt_llm/llmapi/llm.py

@@ -486,7 +486,7 @@ def _check_arguments(self, prompt_len: int, query_len: int,
 
         if (not self.args.enable_chunked_prefill) and (
                 prompt_len / self.args.parallel_config.cp_size + query_len +
-                sampling_params.max_tokens > max_seq_len):
+            (sampling_params.max_tokens or 0) > max_seq_len):
             raise ValueError(
                 f"The sum of prompt length ({prompt_len/self.args.parallel_config.cp_size}) and query length ({query_len}) max_tokens ({sampling_params.max_tokens}) should not exceed "
                 f"max_seq_len ({build_config.max_seq_len})")
@@ -542,6 +542,14 @@ def _build_model(self):
             max_batch_size=max_batch_size,
             max_num_tokens=max_num_tokens,
             gather_generation_logits=self.args.gather_generation_logits)
+        if self.args.backend is None:
+            # also set executor_config.max_seq_len in TRT workflow, to deduce default max_tokens
+            if max_seq_len is not None:
+                executor_config.max_seq_len = max_seq_len
+            else:
+                engine_config = EngineConfig.from_json_file(self._engine_dir /
+                                                            "config.json")
+                executor_config.max_seq_len = engine_config.build_config.max_seq_len
         if self.args.kv_cache_config is not None:
             executor_config.kv_cache_config = PybindMirror.maybe_to_pybind(
                 self.args.kv_cache_config)

tensorrt_llm/serve/openai_protocol.py

@@ -156,7 +156,7 @@ class CompletionRequest(OpenAIBaseModel):
     frequency_penalty: Optional[float] = 0.0
     logit_bias: Optional[Dict[str, float]] = None
     logprobs: Optional[int] = None
-    max_tokens: Optional[int] = 16
+    max_tokens: Optional[int] = None
     n: int = 1
     presence_penalty: Optional[float] = 0.0
     seed: Optional[int] = Field(default=None)
@@ -426,7 +426,7 @@ class ChatCompletionRequest(OpenAIBaseModel):
     logit_bias: Optional[Dict[str, float]] = None
     logprobs: Optional[int] = None
     top_logprobs: Optional[int] = 0
-    max_completion_tokens: int = Field(default=16,
+    max_completion_tokens: int = Field(default=None,
                                        validation_alias='max_tokens')
     n: Optional[int] = 1
     presence_penalty: Optional[float] = 0.0

tests/unittest/llmapi/apps/_test_openai_chat.py

@@ -130,6 +130,18 @@ def test_single_chat_session(client: openai.OpenAI, model_name: str):
     )
     assert legacy.choices[0].message.content \
         == chat_completion.choices[0].message.content
+    # test deduced max_tokens
+    chat_completion = client.chat.completions.create(
+        model=model_name,
+        messages=messages,
+        temperature=0.0,
+        logprobs=False,
+    )
+    assert chat_completion.id is not None
+    assert len(chat_completion.choices) == 1
+    message = chat_completion.choices[0].message
+    assert message.content is not None
+    assert message.role == "assistant"
 
 
 def test_single_chat_session_with_logprobs(client: openai.OpenAI,
@@ -458,6 +470,7 @@ def test_custom_role(client: openai.OpenAI, model_name: str):
             "content": "what is 1+1?",
         }],  # type: ignore
         temperature=0.0,
+        max_completion_tokens=16,
         seed=0)
 
     resp2 = client.chat.completions.create(
@@ -470,6 +483,7 @@ def test_custom_role(client: openai.OpenAI, model_name: str):
             }]
         }],  # type: ignore
         temperature=0.0,
+        max_completion_tokens=16,
         seed=0)
 
     content1 = resp1.choices[0].message.content