[njhill]

Overview

In Model Runner V2, accumulate the prompt logprobs of a prompt chunked across steps in a pinned host buffer, copied each step on the output copy stream, instead of keeping every step's rows on the device until the final chunk and concatenating them. Follow-up to #64841 and #64914.

Claims

• Device memory for chunked prompt logprobs no longer grows with prompt length: the device holds at most one step's rows. With full-vocab logprobs (prompt_logprobs=-1) over a 152k vocab, each prompt token costs ~1.2 MB (int32 ids + fp32 scores), i.e. ~10 GB for an 8k-token prompt before this change.
• Removes the final LogprobsTensors.cat, which briefly held both the per-step tensors and their concatenation.
• Matches the MRV1 runner, which already accumulates in-progress prompt logprobs on the CPU.

Validation

• An earlier version of this change passed the prompt-logprobs tests on GB200: tests/v1/sample/test_logprobs.py prompt-logprobs selection, including chunked prefill, 16 passed / 7 skipped (the skips are the by-design prefix-caching cases).
• The current version (row offset derived from prompt position instead of a stored count) has only been lint/mypy checked so far; draft until CI runs.

Details

• Buffer: on the first step that produces rows for a chunked request, allocate LogprobsTensors.empty_cpu(rows_left, width), where rows_left = prompt_len - 1 - max(prompt_logprob_start, num_computed_prefill). Each step writes at offset = len(buffer) - rows_left, so the final chunk fills the buffer exactly and it is returned without slicing.
• Copies: after the loop, copy_stream.wait_stream(main_stream), then non-blocking D2H copies on the output copy stream, with record_stream on the device sources. The step's output copy (AsyncOutput) runs on the same stream later and its completion event covers these copies; LogprobsTensors.to_cpu_nonblocking already passes CPU tensors through.
• No new syncs: pinned non-blocking copies only; the host never waits on them inside compute_prompt_logprobs.
• Stale entries: add_request drops any leftover buffer for a reused request id.
• Aggregation stays in the worker: the scheduler asserts that no prompt logprobs are emitted on non-final chunk steps.

---

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