Repository: vllm-project/vllm

Issue:
[Bug]: prompt_logprobs=0 raises AssertionError
### Your current environment

PyTorch version: 2.3.0+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.4 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: version 3.29.2
Libc version: glibc-2.35

Python version: 3.10.12 (main, Nov 20 2023, 15:14:05) [GCC 11.4.0] (64-bit runtime)
...

### 🐛 Describe the bug

With vllm==0.4.3, `prompt_logprobs=0` raises AssertionError on the shape of (pruned) logits.

```python
import vllm

llm = vllm.LLM("facebook/opt-125m")
llm.generate(["Hello, my name is"], vllm.SamplingParams(prompt_logprobs=0))
```
raises
```
Traceback (most recent call last):
  File "/home/kataoka/venv/lib/python3.10/site-packages/IPython/core/interactiveshell.py", line 3553, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "/tmp/ipykernel_2972011/3028385135.py", line 1, in <module>
    llm.generate(["Hello, my name is"], vllm.SamplingParams(prompt_logprobs=0))
  File "/home/kataoka/venv/lib/python3.10/site-packages/vllm/utils.py", line 672, in inner
    return fn(*args, **kwargs)
  File "/home/kataoka/venv/lib/python3.10/site-packages/vllm/entrypoints/llm.py", line 300, in generate
    outputs = self._run_engine(use_tqdm=use_tqdm)
  File "/home/kataoka/venv/lib/python3.10/site-packages/vllm/entrypoints/llm.py", line 552, in _run_engine
    step_outputs = self.llm_engine.step()
  File "/home/kataoka/venv/lib/python3.10/site-packages/vllm/engine/llm_engine.py", line 772, in step
    output = self.model_executor.execute_model(
  File "/home/kataoka/venv/lib/python3.10/site-packages/vllm/executor/gpu_executor.py", line 91, in execute_model
    output = self.driver_worker.execute_model(execute_model_req)
  File "/home/kataoka/venv/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/home/kataoka/venv/lib/python3.10/site-packages/vllm/worker/worker.py", line 272, in execute_model
    output = self.model_runner.execute_model(seq_group_metadata_list,
  File "/home/kataoka/venv/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/home/kataoka/venv/lib/python3.10/site-packages/vllm/worker/model_runner.py", line 738, in execute_model
    output = self.model.sample(
  File "/home/kataoka/venv/lib/python3.10/site-packages/vllm/model_executor/models/opt.py", line 323, in sample
    next_tokens = self.sampler(logits, sampling_metadata)
  File "/home/kataoka/venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/kataoka/venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/kataoka/venv/lib/python3.10/site-packages/vllm/model_executor/layers/sampler.py", line 63, in forward
    logits = _apply_min_tokens_penalty(logits, sampling_metadata)
  File "/home/kataoka/venv/lib/python3.10/site-packages/vllm/model_executor/layers/sampler.py", line 196, in _apply_min_tokens_penalty
    assert logits_applied == logits.shape[0]
AssertionError
```

As a consequence, `logprobs=0` cannot be used with `echo=True` in OpenAI-compatible API:
```python
import requests

requests.post("http://localhost:8000/v1/completions", json=dict(
    model="facebook/opt-125m",
    prompt="Hello, my name is",
    max_tokens=10,
    logprobs=0,
    echo=True,
))
```
to `python -m vllm.entrypoints.openai.api_server --model facebook/opt-125m` tries to process
```
INFO 06-03 18:19:58 async_llm_engine.py:553] Received request cmpl-9a5c07c312af4bb2a16290754987c2a8-0: prompt: 'Hello, my name is', params: SamplingParams(n=1, best_of=1, presence_penalty=0.0, frequency_penalty=0.0, repetition_penalty=1.0, temperature=1.0, top_p=1.0, top_k=-1, min_p=0.0, seed=None, use_beam_search=False, length_penalty=1.0, early_stopping=False, stop=[], stop_token_ids=[], include_stop_str_in_output=False, ignore_eos=False, max_tokens=10, min_tokens=0, logprobs=0, prompt_logprobs=0, skip_special_tokens=True, spaces_between_special_tokens=True, truncate_prompt_tokens=None), prompt_token_ids: [2, 31414, 6, 127, 766, 16], lora_request: None.
```
and causes 500 Internal Server Error from the same AssertionError.



Discussion:
It seems #4309 caused the degradation.

`llm.generate(["Hello, my name is"], vllm.SamplingParams(prompt_logprobs=0))`
- pass: v0.4.0.post1, v0.4.1, 603ad848~
- fail: 603ad848 v0.4.2, v0.4.3

Resolving patch:
diff --git a/vllm/entrypoints/openai/serving_completion.py b/vllm/entrypoints/openai/serving_completion.py
--- a/vllm/entrypoints/openai/serving_completion.py
+++ b/vllm/entrypoints/openai/serving_completion.py
@@ -312,7 +312,7 @@ def request_output_to_completion_response(
                 elif request.echo and request.max_tokens > 0:
                     token_ids = prompt_token_ids + output.token_ids
                     top_logprobs = (prompt_logprobs + output.logprobs
-                                    if request.logprobs else None)
+                                    if request.logprobs is not None else None)
                     output_text = prompt_text + output.text
                 else:
                     token_ids = output.token_ids
diff --git a/vllm/model_executor/sampling_metadata.py b/vllm/model_executor/sampling_metadata.py
--- a/vllm/model_executor/sampling_metadata.py
+++ b/vllm/model_executor/sampling_metadata.py
@@ -233,7 +233,7 @@ def _prepare_seq_groups(
         logits = hidden_states[selected_token_indices]
         """
 
-        if sampling_params.prompt_logprobs:
+        if sampling_params.prompt_logprobs is not None:
             selected_token_indices.extend(
                 range(model_output_idx, model_output_idx + prompt_logprob_len))
         model_output_idx += prompt_logprob_len
diff --git a/vllm/worker/model_runner.py b/vllm/worker/model_runner.py
--- a/vllm/worker/model_runner.py
+++ b/vllm/worker/model_runner.py
@@ -427,7 +427,7 @@ def _prepare_model_input(
                     [lora_id] *
                     (query_len if seq_group_metadata.sampling_params
                      and seq_group_metadata.sampling_params.prompt_logprobs
-                     else 1))
+                     is not None else 1))
 
                 mm_data = seq_group_metadata.multi_modal_data
                 if mm_data is not None:


Test patch:
diff --git a/tests/entrypoints/test_openai_server.py b/tests/entrypoints/test_openai_server.py
--- a/tests/entrypoints/test_openai_server.py
+++ b/tests/entrypoints/test_openai_server.py
@@ -224,7 +224,7 @@ async def test_zero_logprobs(server, client: openai.AsyncOpenAI,
     assert choice.logprobs is not None
     assert choice.logprobs.token_logprobs is not None
     assert choice.logprobs.top_logprobs is not None
-    assert len(choice.logprobs.top_logprobs[0]) <= 1
+    assert len(choice.logprobs.top_logprobs[0]) == 1
 
 
 @pytest.mark.asyncio
@@ -246,7 +246,7 @@ async def test_some_logprobs(server, client: openai.AsyncOpenAI,
     assert choice.logprobs is not None
     assert choice.logprobs.token_logprobs is not None
     assert choice.logprobs.top_logprobs is not None
-    assert len(choice.logprobs.top_logprobs[0]) <= 6
+    assert 5 <= len(choice.logprobs.top_logprobs[0]) <= 6
 
 
 @pytest.mark.asyncio
@@ -1032,8 +1032,9 @@ async def test_guided_grammar(server, client: openai.AsyncOpenAI):
     "model_name",
     [MODEL_NAME, "zephyr-lora", "zephyr-lora2"],
 )
+@pytest.mark.parametrize("logprobs_arg", [1, 0])
 async def test_echo_logprob_completion(server, client: openai.AsyncOpenAI,
-                                       model_name: str):
+                                       model_name: str, logprobs_arg: int):
     tokenizer = get_tokenizer(tokenizer_name=MODEL_NAME)
     # test using text and token IDs
     for prompt in ("Hello, my name is", [0, 0, 0, 0, 0]):
@@ -1042,7 +1043,7 @@ async def test_echo_logprob_completion(server, client: openai.AsyncOpenAI,
                                                      max_tokens=5,
                                                      temperature=0.0,
                                                      echo=True,
-                                                     logprobs=1)
+                                                     logprobs=logprobs_arg)
 
         prompt_text = tokenizer.decode(prompt) if isinstance(prompt,
                                                              list) else prompt
@@ -1055,6 +1056,9 @@ async def test_echo_logprob_completion(server, client: openai.AsyncOpenAI,
                 and logprobs.token_logprobs[0] is None)
         assert (len(logprobs.top_logprobs) > 5
                 and logprobs.top_logprobs[0] is None)
+        for top_logprobs in logprobs.top_logprobs[1:]:
+            assert max(logprobs_arg,
+                       1) <= len(top_logprobs) <= logprobs_arg + 1
         assert len(logprobs.tokens) > 5