Repository navigation
Conversation
|
BTW, I found another bug in the sampler while testing this feature. When some sequences are preempted (which is common when the size of the KV cache is insufficient), the This bug appears to be unrelated to my changes, as it also occurs in the original code. I would like to fix it, but I am currently busy looking for a research internship, so I may work on it later. |
I think this will be fixed by #5349 |
|
The modifications to the scheduler logic of chunked prefills have issues, causing the tests to fail. Let me fix this now and submit a new commit shortly. |
|
Thanks for taking this on! |
|
I have carefully reviewed the failed checks and confirmed that they are not caused by my changes. @robertgshaw2-neuralmagic could you please review this PR? Thank you! |
|
@Conless I tried using this PR to perform an evaluation on MMLU (where it is easy to run into the existing OOM issue because of the large number of logprobs required) and this PR seems to completely prevent that evaluation from working now. For context, the previous method of failure on main was just if the batch size was too large or memory utilization was too high, we would hit OOM trying to allocate dynamic structures to hold the logprobs. The key issue seems to be the triggering of: If I reduce the I thought this PR would enable the scheduler to make space for logprobs and not "over-allocate" for subsequent requests coming in. Or at least make vLLM aware of the "leftover" memory it has available and do some work of managing that in tandem with kv cache blocks to decide how to schedule the requests coming in. I believe this current method of having a fixed number is not sufficient and doesn't scale when we do have enough memory available to run the request. The following are steps to reproduce. Install lm-evaluation-harness: git clone https://github.com/EleutherAI/lm-evaluation-harness.git
pip install -e lm-evaluation-harnessEval: |
|
@mgoin Sorry for the late reply. I was busy handling other stuff these days.
To handle long requests with prompt_logprobs, you may increase
I'm not sure if I fully understand your point. Actually, I don't think this design conflicts with your idea. The current logic is as follows:
A possible improvement is to balance the KV cache and the number of logprobs in each iteration. However, that would be quite difficult since the KV cache is allocated statically. |
|
Hi thanks for the great work, it will fix a lot of OOM errors when we run prompt logprobs. Do you have any ETA on when it will get merged? |
Sorry I’m busy with some other works these days. As what is mentioned above, to implement a more fine-grained management, which may be expected by maintainers, requires some more work in memory management. Thus it may be hard to finish it soon. If this bug also bothers you, we may work on it together. I’m looking forward to collaborators! |
|
This pull request has been automatically marked as stale because it has not had any activity within 90 days. It will be automatically closed if no further activity occurs within 30 days. Leave a comment if you feel this pull request should remain open. Thank you! |
|
This pull request has merge conflicts that must be resolved before it can be |
github-actions
Bot
added
unstale
`_get_prompt_logprobs_dict` materialized the full [num_scheduled_tokens, vocab_size] logits tensor in one shot and then upcast it to float32 via `log_softmax`. Both tensors are live at the same time, so the transient peak is roughly `num_scheduled_tokens * vocab_size * (2 + 4)` bytes: at 8192 scheduled tokens and a 152k vocabulary that is ~7.5 GiB. That peak is never observed by memory profiling. `profile_run` calls `_dummy_sampler_run`, which builds its `SamplingMetadata` with `max_num_logprobs=None`, so the logprobs path is not exercised and the float32 score tensor is never allocated. KV cache capacity is therefore sized against a peak that excludes prompt logprobs, and the allocation has to come out of whatever is left at serving time. This is the mechanism described in vllm-project#5907: "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." That issue was closed as stale and its fix attempt (vllm-project#5355) was closed unmerged; the behaviour is still present on main. Rather than change profiling, bound the allocation directly, which is the "chunked logits processing" approach proposed and endorsed in the over row chunks of `VLLM_PROMPT_LOGPROBS_CHUNK_SIZE` (default 1024), and each chunk's results are copied into the existing per-request CPU tensors before the next chunk is materialized. The peak becomes independent of `max_num_batched_tokens`. Behaviour is unchanged: same logits, same scores, same destination rows, same `logprobs_mode` handling. Only the number of rows materialized at once differs. The default of 1024 leaves the common case (short prompts, small `max_num_batched_tokens`) at a single chunk. Signed-off-by: Michel Belleau <michel.belleau@malaiwah.com>
As mentioned in #5067 , current profiler and scheduler logic do not consider the VRAM usage of computing log probabilities, which can consume up to 7.8 GB VRAM. This PR introduces a new engine arg called
max_num_batched_logprobs, with its default value tomax_seqsfor simplicity, and take it into account in the profiler and scheduler logic.Specifically, in$k$ prompt are set to $k=\lceil\frac{max\ num\ batched\ logprobs}{tokens}\rceil$ . Then in
profile_run, sampling params of the firstprompt_logprobs=1, wherescheduler.py, a new variable calledlogprob_budgetis added intoSchedulingBudgetto track the number of tokens that require logprobs.A simple unit test is added into
test_logprobs.pyfor testing this new engine arg.FIX #5067
FIX #1532
FIX #5147
FIX #5274
FIX #5550
PR Checklist (Click to Expand)
Thank you for your contribution to vLLM! Before submitting the pull request, please ensure the PR meets the following criteria. This helps vLLM maintain the code quality and improve the efficiency of the review process.
PR Title and Classification
Only specific types of PRs will be reviewed. The PR title is prefixed appropriately to indicate the type of change. Please use one of the following:
[Bugfix]for bug fixes.[CI/Build]for build or continuous integration improvements.[Doc]for documentation fixes and improvements.[Model]for adding a new model or improving an existing model. Model name should appear in the title.[Frontend]For changes on the vLLM frontend (e.g., OpenAI API server,LLMclass, etc.)[Kernel]for changes affecting CUDA kernels or other compute kernels.[Core]for changes in the core vLLM logic (e.g.,LLMEngine,AsyncLLMEngine,Scheduler, etc.)[Hardware][Vendor]for hardware-specific changes. Vendor name should appear in the prefix (e.g.,[Hardware][AMD]).[Misc]for PRs that do not fit the above categories. Please use this sparingly.Note: If the PR spans more than one category, please include all relevant prefixes.
Code Quality
The PR need to meet the following code quality standards:
format.shto format your code.docs/source/if the PR modifies the user-facing behaviors of vLLM. It helps vLLM user understand and utilize the new features or changes.Notes for Large Changes
Please keep the changes as concise as possible. For major architectural changes (>500 LOC excluding kernel/data/config/test), we would expect a GitHub issue (RFC) discussing the technical design and justification. Otherwise, we will tag it with
rfc-requiredand might not go through the PR.What to Expect for the Reviews
The goal of the vLLM team is to be a transparent reviewing machine. We would like to make the review process transparent and efficient and make sure no contributor feel confused or frustrated. However, the vLLM team is small, so we need to prioritize some PRs over others. Here is what you can expect from the review process:
action-requiredlabel on the PR if there are changes required. The contributor should address the comments and ping the reviewer to re-review the PR.Thank You
Finally, thank you for taking the time to read these guidelines and for your interest in contributing to vLLM. Your contributions make vLLM a great tool for everyone!