Repository: vllm-project/vllm
Pull request: [Platform] Add custom default max tokens
FIX: https://github.com/vllm-project/vllm-spyre/issues/148

Currently the default max tokens for the completions API is set to `max_model_len - prompt_len`. The changes in this PR make so that when a platform needs to use a different value for `default_max_tokens` it can be altered simply by overriding the `maybe_update_max_tokens` method in the class `Plataform`. When it is not needed it returns the current default. Edit: typo in commit message: `class Plataform` is meant to be `class Platform`.

Problem statement:
Incorrect default max_completion_tokens being set
### Changes

Trying to run vLLM with `VLLM_SPYRE_WARMUP_NEW_TOKENS` set to X number always results in error:

```
No applicable warmup shape exists for combination of prompt length (Z tokens) and maximum number of output tokens to be generated (Y tokens)
```

Where `Z`  is input token count and `Y` is `Z + VLLM_SPYRE_WARMUP_NEW_TOKENS`

Hints from issue discussion:
To add some more context: This is observed when using the Chat API while not setting `max_completions_token` or `max_tokens` in the request body.

Turns out that the OpenAI frontend code is where the default value for `max_completions_token`/`max_tokens` is set.
For the completions API, it defaults to 16 ([REF](https://github.com/vllm-project/vllm/blob/a944f8ede7361a5233e112a575ff77c4aaa268a5/vllm/entrypoints/openai/protocol.py#L747)), so is typically overwritten in the incoming request.
For the chat API, it defaults to `max_model_len - prompt_len` ([REF](https://github.com/vllm-project/vllm/blob/a944f8ede7361a5233e112a575ff77c4aaa268a5/vllm/entrypoints/openai/serving_chat.py#L213-L214)). The setting of the default max tokens happens before the execution gets to any vllm-spyre code 🤔

The code in `platform.py` overrides `max_model_len` to be the sum of the max prompt size and generation length [REF](https://github.com/vllm-project/vllm-spyre/blob/main/vllm_spyre/platform.py#L83-L88), so the only way that leaving max tokens unset works is if the incoming request has tokens equal to the largest max prompt size.
Since the default value for generation tokens is set in the OpenAI frontend code, the only fix seems to involve altering the vllm source code:

If a new method is implemented in the class `Plataform` in `vllm/platforms/interface.py`:
```python
    def maybe_update_default_max_tokens(self, prompt_len: int, default_max_tokens: int) -> int:
        return default_max_tokens
```


It can be override in `vllm_spyre/platform.py`:
```python
    @classmethod
    def maybe_update_default_max_tokens(self, prompt_len: int, default_max_tokens: int) -> int:
        max_new_tokens = 1
        for shape in self._warmup_shapes:
            if prompt_len <= shape['prompt_length']:
                max_new_tokens = max(max_new_tokens, shape['new_tokens'])
        
        return max_new_tokens
```

Now, in `vllm/entrypoints/openai/serving_chat.py`, in the class `OpenAIServingChat`, instead of using:
```python
     default_max_tokens = self.max_model_len - len(engine_prompt["prompt_token_ids"])
```


use:
```python
     default_max_tokens = current_platform.maybe_update_default_max_tokens(
                    prompt_len=prompt_len,
                    default_max_tokens=self.max_model_len - len(engine_prompt["prompt_token_ids"])
                    )
```

This way, any plataform can override this method and use a different default if needed, without altering the vllm default value.

Commit message: trigger ci

Signed-off-by: Gabriel Marinho <gmarinho@ibm.com>

Patch to review:
diff --git a/vllm/entrypoints/openai/protocol.py b/vllm/entrypoints/openai/protocol.py
index 3b5281962b2d..d9d112730150 100644
--- a/vllm/entrypoints/openai/protocol.py
+++ b/vllm/entrypoints/openai/protocol.py
@@ -229,7 +229,6 @@ class ChatCompletionRequest(OpenAIBaseModel):
     logit_bias: Optional[dict[str, float]] = None
     logprobs: Optional[bool] = False
     top_logprobs: Optional[int] = 0
-    # TODO(#9845): remove max_tokens when field is removed from OpenAI API
     max_tokens: Optional[int] = Field(
         default=None,
         deprecated=
@@ -433,23 +432,10 @@ class ChatCompletionRequest(OpenAIBaseModel):
     }
 
     def to_beam_search_params(
-            self,
-            default_max_tokens: int,
-            default_sampling_params: Optional[dict] = None
-    ) -> BeamSearchParams:
-        # TODO(#9845): remove max_tokens when field is removed from OpenAI API
-        max_tokens = self.max_completion_tokens or self.max_tokens
+            self, max_tokens: int,
+            default_sampling_params: dict) -> BeamSearchParams:
 
-        if default_sampling_params is None:
-            default_sampling_params = {}
         n = self.n if self.n is not None else 1
-
-        # Use minimum of context window, user request & server limit.
-        max_tokens = min(
-            val for val in (default_max_tokens, max_tokens,
-                            default_sampling_params.get("max_tokens", None))
-            if val is not None)
-
         if (temperature := self.temperature) is None:
             temperature = default_sampling_params.get(
                 "temperature", self._DEFAULT_SAMPLING_PARAMS["temperature"])
@@ -465,21 +451,10 @@ def to_beam_search_params(
 
     def to_sampling_params(
         self,
-        default_max_tokens: int,
+        max_tokens: int,
         logits_processor_pattern: Optional[str],
-        default_sampling_params: Optional[dict] = None,
+        default_sampling_params: dict,
     ) -> SamplingParams:
-        # TODO(#9845): remove max_tokens when field is removed from OpenAI API
-        max_tokens = self.max_completion_tokens or self.max_tokens
-
-        if default_sampling_params is None:
-            default_sampling_params = {}
-
-        # Use minimum of context window, user request & server limit.
-        max_tokens = min(
-            val for val in (default_max_tokens, max_tokens,
-                            default_sampling_params.get("max_tokens", None))
-            if val is not None)
 
         # Default parameters
         if (repetition_penalty := self.repetition_penalty) is None:
@@ -900,22 +875,15 @@ class CompletionRequest(OpenAIBaseModel):
     }
 
     def to_beam_search_params(
-            self,
-            default_max_tokens: int,
-            default_sampling_params: Optional[dict] = None
+        self,
+        max_tokens: int,
+        default_sampling_params: Optional[dict] = None,
     ) -> BeamSearchParams:
-        max_tokens = self.max_tokens
 
         if default_sampling_params is None:
             default_sampling_params = {}
         n = self.n if self.n is not None else 1
 
-        # Use minimum of context window, user request & server limit.
-        max_tokens = min(
-            val for val in (default_max_tokens, max_tokens,
-                            default_sampling_params.get("max_tokens", None))
-            if val is not None)
-
         if (temperature := self.temperature) is None:
             temperature = default_sampling_params.get("temperature", 1.0)
 
@@ -930,21 +898,14 @@ def to_beam_search_params(
 
     def to_sampling_params(
         self,
-        default_max_tokens: int,
+        max_tokens: int,
         logits_processor_pattern: Optional[str],
         default_sampling_params: Optional[dict] = None,
     ) -> SamplingParams:
-        max_tokens = self.max_tokens
 
         if default_sampling_params is None:
             default_sampling_params = {}
 
-        # Use minimum of context window, user request & server limit.
-        max_tokens = min(
-            val for val in (default_max_tokens, max_tokens,
-                            default_sampling_params.get("max_tokens", None))
-            if val is not None)
-
         # Default parameters
         if (repetition_penalty := self.repetition_penalty) is None:
             repetition_penalty = default_sampling_params.get(
@@ -1816,7 +1777,7 @@ def to_sampling_params(
             self,
             default_max_tokens: int,
             default_sampling_params: Optional[dict] = None) -> SamplingParams:
-        # TODO(#9845): remove max_tokens when field is removed from OpenAI API
+
         max_tokens = default_max_tokens
 
         if default_sampling_params is None:
@@ -2032,7 +1993,7 @@ def to_sampling_params(
             self,
             default_max_tokens: int,
             default_sampling_params: Optional[dict] = None) -> SamplingParams:
-        # TODO(#9845): remove max_tokens when field is removed from OpenAI API
+
         max_tokens = default_max_tokens
 
         if default_sampling_params is None:
diff --git a/vllm/entrypoints/openai/serving_chat.py b/vllm/entrypoints/openai/serving_chat.py
index 10aced83b60b..64bf7d9e34f8 100644
--- a/vllm/entrypoints/openai/serving_chat.py
+++ b/vllm/entrypoints/openai/serving_chat.py
@@ -34,6 +34,7 @@
 from vllm.entrypoints.openai.tool_parsers import ToolParser, ToolParserManager
 from vllm.entrypoints.openai.tool_parsers.mistral_tool_parser import (
     MistralToolCall)
+from vllm.entrypoints.utils import get_max_tokens
 from vllm.logger import init_logger
 from vllm.outputs import CompletionOutput, RequestOutput
 from vllm.reasoning import ReasoningParser, ReasoningParserManager
@@ -215,15 +216,22 @@ async def create_chat_completion(
         try:
             for i, engine_prompt in enumerate(engine_prompts):
                 sampling_params: Union[SamplingParams, BeamSearchParams]
-                default_max_tokens = self.max_model_len - len(
-                    engine_prompt["prompt_token_ids"])
+
+                if self.default_sampling_params is None:
+                    self.default_sampling_params = {}
+
+                max_tokens = get_max_tokens(
+                    max_model_len=self.max_model_len,
+                    request=request,
+                    input_length=len(engine_prompt["prompt_token_ids"]),
+                    default_sampling_params=self.default_sampling_params)
+
                 if request.use_beam_search:
                     sampling_params = request.to_beam_search_params(
-                        default_max_tokens, self.default_sampling_params)
+                        max_tokens, self.default_sampling_params)
                 else:
                     sampling_params = request.to_sampling_params(
-                        default_max_tokens,
-                        self.model_config.logits_processor_pattern,
+                        max_tokens, self.model_config.logits_processor_pattern,
                         self.default_sampling_params)
 
                 self._log_inputs(request_id,
diff --git a/vllm/entrypoints/openai/serving_completion.py b/vllm/entrypoints/openai/serving_completion.py
index a19fde8d70a8..52c7b88cbf32 100644
--- a/vllm/entrypoints/openai/serving_completion.py
+++ b/vllm/entrypoints/openai/serving_completion.py
@@ -30,6 +30,7 @@
                                                     clamp_prompt_logprobs,
                                                     is_text_tokens_prompt)
 from vllm.entrypoints.openai.serving_models import OpenAIServingModels
+from vllm.entrypoints.utils import get_max_tokens
 from vllm.inputs.data import (EmbedsPrompt, TokensPrompt, is_embeds_prompt,
                               is_tokens_prompt)
 from vllm.logger import init_logger
@@ -157,15 +158,22 @@ async def create_completion(
                     input_length = len(engine_prompt["prompt_token_ids"])
                 else:
                     assert_never(engine_prompt)
-                default_max_tokens = self.max_model_len - input_length
+
+                if self.default_sampling_params is None:
+                    self.default_sampling_params = {}
+
+                max_tokens = get_max_tokens(
+                    max_model_len=self.max_model_len,
+                    request=request,
+                    input_length=input_length,
+                    default_sampling_params=self.default_sampling_params)
 
                 if request.use_beam_search:
                     sampling_params = request.to_beam_search_params(
-                        default_max_tokens, self.default_sampling_params)
+                        max_tokens, self.default_sampling_params)
                 else:
                     sampling_params = request.to_sampling_params(
-                        default_max_tokens,
-                        self.model_config.logits_processor_pattern,
+                        max_tokens, self.model_config.logits_processor_pattern,
                         self.default_sampling_params)
 
                 request_id_item = f"{request_id}-{i}"
diff --git a/vllm/entrypoints/utils.py b/vllm/entrypoints/utils.py
index 16ba2b4531ac..582396391edc 100644
--- a/vllm/entrypoints/utils.py
+++ b/vllm/entrypoints/utils.py
@@ -4,13 +4,17 @@
 import asyncio
 import functools
 import os
-from typing import Any, Optional
+import sys
+from typing import Any, Optional, Union
 
 from fastapi import Request
 from fastapi.responses import JSONResponse, StreamingResponse
 from starlette.background import BackgroundTask, BackgroundTasks
 
+from vllm.entrypoints.openai.protocol import (ChatCompletionRequest,
+                                              CompletionRequest)
 from vllm.logger import init_logger
+from vllm.platforms import current_platform
 
 logger = init_logger(__name__)
 
@@ -175,7 +179,6 @@ def _validate_truncation_size(
 
 
 def show_filtered_argument_or_group_from_help(parser, subcommand_name):
-    import sys
 
     # Only handle --help=<keyword> for the current subcommand.
     # Since subparser_init() runs for all subcommands during CLI setup,
@@ -231,3 +234,17 @@ def show_filtered_argument_or_group_from_help(parser, subcommand_name):
             print(f"\nNo group or parameter matching '{search_keyword}'")
             print("Tip: use `--help=listgroup` to view all groups.")
             sys.exit(1)
+
+
+def get_max_tokens(max_model_len: int, request: Union[ChatCompletionRequest,
+                                                      CompletionRequest],
+                   input_length: int, default_sampling_params: dict) -> int:
+
+    max_tokens = getattr(request, "max_completion_tokens", request.max_tokens)
+    default_max_tokens = max_model_len - input_length
+    max_output_tokens = current_platform.get_max_output_tokens(input_length)
+
+    return min(val
+               for val in (default_max_tokens, max_tokens, max_output_tokens,
+                           default_sampling_params.get("max_tokens"))
+               if val is not None)
diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py
index 0f08bf986333..567d5cbf503f 100644
--- a/vllm/platforms/interface.py
+++ b/vllm/platforms/interface.py
@@ -4,6 +4,7 @@
 import os
 import platform
 import random
+import sys
 from datetime import timedelta
 from platform import uname
 from typing import TYPE_CHECKING, NamedTuple, Optional, Union
@@ -164,6 +165,9 @@ def is_neuron(self) -> bool:
     def is_out_of_tree(self) -> bool:
         return self._enum == PlatformEnum.OOT
 
+    def get_max_output_tokens(self, prompt_len: int) -> int:
+        return sys.maxsize
+
     def is_cuda_alike(self) -> bool:
         """Stateless version of [torch.cuda.is_available][]."""
         return self._enum in (PlatformEnum.CUDA, PlatformEnum.ROCM)


Code review comments:
diff_hunk: @@ -231,3 +234,17 @@ def show_filtered_argument_or_group_from_help(parser, subcommand_name):
             print(f"\nNo group or parameter matching '{search_keyword}'")
             print("Tip: use `--help=listgroup` to view all groups.")
             sys.exit(1)
+
+
+def get_max_tokens(max_model_len: int, request: Union[ChatCompletionRequest,
+                                                      CompletionRequest],
+                   input_length: int, default_sampling_params: dict) -> int:
+
+    max_tokens = getattr(request, "max_completion_tokens", request.max_tokens)
path: vllm/entrypoints/utils.py
text: @user1:
This will fix the test error:
```suggestion
max_tokens = getattr(request, "max_completion_tokens", None) or request.max_tokens
```
If the attribute exists but is None, getattr will return None and not the default.