Repository: vllm-project/vllm

Issue:
llm.generate hangs for long prompt
I'm using the latest vllm (0.2.1.post1) and `TheBloke/vicuna-13B-v1.5-16K-AWQ` model and it seems that llm.generate will hang when the length of prompt token ids >= 5985.

```
from vllm import LLM, SamplingParams

llm = LLM(model="TheBloke/vicuna-13B-v1.5-16K-AWQ", quantization="awq")
tokenizer = llm.get_tokenizer()

prompt1 = "xxx..."
print(len(tokenizer.encode(prompt1)))  # 5984

sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
outputs = llm.generate([prompt1], sampling_params)
Processed prompts: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:04<00:00,  4.92s/it]

prompt2 = prompt1 + "\n"
print(len(tokenizer.encode(prompt2)))  # 5985

outputs = llm.generate([prompt2], sampling_params)
Processed prompts:   0%|                                                                                                                                                                                                                                | 0/1 [00:00<?, ?it/s]
```


Discussion:
How's your VRAM?   I was having a similar issue and I had to reduce max_model_len a bit to get it to process.

It would chug away for a while then just drop some VRam out without any type of error, then hang and the power usage on the card would drop back towards idle consumption.   Sound similar?

I am having success with prompts up to about 16k tokens now, which is where my limit seems to be sitting currently using 1xA6000 / Yi-34B-200K AWQ quant.
I'm using NVIDIA RTX 3090 (cuda:0) and when the hang happens the occupied VRAM is about 17GB.

```
$ nvidia-smi 
Thu Nov 16 10:01:33 2023       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 520.56.06    Driver Version: 520.56.06    CUDA Version: 11.8     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  NVIDIA GeForce ...  Off  | 00000000:AF:00.0 Off |                  N/A |
| 37%   45C    P2   102W / 350W |  17574MiB / 24576MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   1  NVIDIA GeForce ...  Off  | 00000000:D8:00.0 Off |                  N/A |
| 30%   30C    P8    20W / 350W |   3122MiB / 24576MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
```

Resolving patch:
diff --git a/vllm/core/scheduler.py b/vllm/core/scheduler.py
--- a/vllm/core/scheduler.py
+++ b/vllm/core/scheduler.py
@@ -196,7 +196,8 @@ def _schedule(self) -> SchedulerOutputs:
                 scheduler_outputs = SchedulerOutputs(
                     scheduled_seq_groups=scheduled,
                     prompt_run=True,
-                    num_batched_tokens=len(seq_lens) * max(seq_lens),
+                    num_batched_tokens=len(seq_lens) *
+                    max(seq_lens) if seq_lens else 0,
                     blocks_to_swap_in=blocks_to_swap_in,
                     blocks_to_swap_out=blocks_to_swap_out,
                     blocks_to_copy=blocks_to_copy,


Test patch: