Incorrect default max_completion_tokens being set

[gkumbhat]

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

[tjohnson31415]

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), so is typically overwritten in the incoming request.
For the chat API, it defaults to max_model_len - prompt_len (REF). 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, so the only way that leaving max tokens unset works is if the incoming request has tokens equal to the largest max prompt size.

[gmarinho2]

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:

    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:

    @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:

     default_max_tokens = self.max_model_len - len(engine_prompt["prompt_token_ids"])

use:

     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.