[Bug]: TRACKING ISSUE: CUDA OOM with Logprobs

[robertgshaw2-neuralmagic]

Your current environment

The output of `python collect_env.py`

🐛 Describe the bug

vLLM has an issue where we can go OOM if too many logprobs are requested.

The reason that this happens is that there are three sources of memory usage:

• Model weights
• KV caches
• Activations

When determining the KV cache size, we calculate peak memory running a long prefill without logprobs

If a prompt requests many logprobs, however, this is an additional source of memory usage which is not considered during warmup and can cause OOM because we have nothing in scheduler to prevent this

We have received several examples of this:

• https://github.com/vllm-project/vllm/issues/5890
• https://github.com/vllm-project/vllm/issues/5060 << some of the OOM issues in AsyncEngineDeadError
• Anyone running MMLU or ARC in lm-eval-harness

Attempt to fix this:

• https://github.com/vllm-project/vllm/pull/5355

I am working on a design to address this issue

[neubig]

Hi @robertgshaw2-neuralmagic , I was wondering if there were any ongoing attempts to resolve this issue, or if #5355 seems like an acceptable fix? Having this fixed would be very helpful to us!

[binxuan]

In my case with large vocab size (200k+) and long seq length (8k+), the logits sort is the most memory consuming part and can easily trigger OOM on one single GPU. Is it possible to do some kind of sequence parallel to distribute it to all TP workers?