[njhill]

Overview

Reduce the peak GPU memory of Model Runner V2's chunked prompt-logprobs computation, which dominates step memory when a request asks for full-vocab prompt logprobs.

Claims

• compute_prompt_logprobs_with_chunking no longer holds every chunk's results twice: chunks are written into outputs preallocated after the first chunk instead of being concatenated at the end.
• Prompt-logprob token ids are returned as int32 (vocab ids fit), taking ids + fp32 scores from 12 to 8 bytes per entry.
• Together: full-vocab prompt logprobs for a 2048-token prompt of a 152k-vocab model (e.g. Qwen2-1.5B) peak at ~2.3 GiB of outputs plus one chunk's working set, instead of ~7 GiB.
• Single-chunk prompts (≤1024 tokens) still return the chunk directly, with no extra copy beyond the int32 cast of the ids.

Validation

On GB200 (1 GPU):

pytest tests/v1/sample/test_logprobs.py -k "chunked_prompt_logprobs_match_unchunked or test_all_logprobs or prompt_logprobs_with_chunking_and_preemption or test_prompt_logprobs_mode or test_get_logprobs_and_prompt_logprobs"
14 passed, 7 skipped

The new test_chunked_prompt_logprobs_match_unchunked (full vocab and top-3) runs 2×1024+3 rows through the chunked path and checks ids, logprobs and ranks against a single unchunked compute_topk_scores pass.

Details

• Motivation: with a measured KV cache sizing warmup ([Core] Extensible KV cache #56492), the warmup's largest prefill with max_logprobs equal to the vocab size reserved ~7.5 GiB for this transient on an 18 GB MIG slice, leaving no room for KV cache (tests/entrypoints/openai/chat_completion/test_chat_echo.py, which sets --max-logprobs 151936). The breakdown: ~3.5 GiB of accumulated chunk outputs, the same again for the final torch.cat, plus one chunk's fp32 logits.
• int32 ids are safe downstream: the tensors are copied to CPU and converted with .tolist(); V1's LogprobsTensors.empty_cpu already uses int32 ids.
• The per-chunk torch.topk indices are still int64; they are bounded by the chunk size and cast when copied into the output.
• The CPU worker's copy of this function (vllm/v1/worker/cpu/logprob.py) is left unchanged.
• Not addressed: when one prompt is split across scheduler steps, PromptLogprobsWorker still keeps each step's results on device until the last step and concatenates them.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

🤖 Generated with Claude Code

https://claude.ai/code/session_01Xgy3B2NXY4GnQaw2WwsqpC

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[njhill]

/ci run

[github-actions]

✅ Triggered Buildkite CI #98905 for commit 40de35b9bd3c.

ℹ️ This PR is 15 commits behind upstream main. Merged upstream main into this branch as 40de35b9bd3c before starting CI. Pull the branch before pushing more changes.

[njhill]

/ci run

[github-actions]

✅ Triggered Buildkite CI #98907 for commit 9ae59d18e1d3.

ℹ️ This PR is 1 commit behind upstream main. Merged upstream main into this branch as 9ae59d18e1d3 before starting CI. Pull the branch before pushing more changes.