perf(SAMPLE-LOGPROBS): the logprobs D2H stages through pageable host memory while a pinned allocator sits in the same file
Nobody has claimed this yet.
- Dominant language
- C++
- Stars
- 423
- Forks
- 53
- Avg merge
- 20h 26m
- Merged PRs (30d)
- 310
Description
Row: SAMPLE-LOGPROBS (.agents/engine-matrix.md:142), under roadmap row ROAD-V1-C7 (.agents/roadmap_v1.md:177). Both resolved by line number. The row is scoped to the payload's correctness, not its copy cost, so whoever picks this up may reasonably decide it is out of that row's scope and needs its own; this issue does not assume it is in scope.
Found by wave PORTQ-7 (#2717) re-deriving PORT-NOW entry [277], upstream 6c18a54648 vllm#54299. Nothing was executed, and this issue is about a bandwidth lever, so nothing here is measured.
Read the caveat first, because it decides whether this is worth doing
The specific defect upstream fixes does not exist here. Upstream's bug is an H2D/D2H copy from a non-pinned tensor silently degrading to a synchronous pageable fallback. In this tree both logprobs staging copies are followed immediately by an explicit b.Synchronize(q) (src/vllm/v1/sample/sampler.cpp:414 and :296) — the sync is deliberate and visible, not a hidden stall.
What remains is that the copies go through the driver's bounce buffer rather than DMA. That is a throughput lever, not a correctness or hidden-stall bug.
Read as "add pin_memory= to three torch.empty calls", the answer here is NOT_APPLICABLE. Read as "the logprobs host staging should be pinned", it is portable work. This issue takes the second reading, and says so rather than hiding the choice.
The facts, which are not in doubt
The live staging buffers are pageable. std::vector<float> raw_logprobs is sized at src/vllm/v1/sample/sampler.cpp:405 and is the destination of b.Copy(q, raw_logprobs.data(), logits.data, ...) at :412; processed_out has the same shape at :289 / :293. Both are [n * vocab] f32 — tens of megabytes per step at a 150k vocabulary — copied device-to-host.
A pinned-host allocator already exists and is already used in this very file. backend->AllocPinned(bytes) at sampler.cpp:252 for the greedy-argmax scratch, and again at src/vllm/v1/worker/gpu/async_output.cpp:43. So the work is reuse, not invention.
The file the record line names is a dead end. LogprobsTensors::empty_cpu (src/vllm/v1/outputs.cpp:17-19) does allocate the three CPU buffers upstream pins, and they are plain std::vector on the pageable heap (include/vllm/v1/outputs.h:42-46) — but they are host-only aggregation buffers, never a DMA endpoint, and empty_cpu has zero callers: grep -rn 'empty_cpu' src/ include/ returns only the definition (outputs.cpp:10) and the declaration (outputs.h:50), with GatherLogprobs in the same probe returning many hits as a built-in positive control. Pinning those would buy nothing.
The commit has fourteen files. Ten are per-model index-tensor constructions that build on device instead of staging through the host; those are per-model questions outside this issue. The FlashInfer backend half is not carried here at all.
What is missing
A pinned, grow-only host staging buffer for the logprobs D2H, replacing the two per-step std::vector destinations — exactly the GreedyArgmaxScratch pattern already at sampler.cpp:225-255, reused rather than reinvented.
Size
~40-70 product lines: one scratch struct extended or generalised, two call sites. Plus a test asserting the destination pointer came from AllocPinned, and a same-binary A/B on the copy. Call it ~120 lines with tests.
The measurement that should gate this, and was not run
Measure the D2H at a realistic n x vocab with and without pinning on one card. If the delta is inside noise — plausible, because the copy may be dwarfed by the host-side top-k in GatherLogprobs (sampler.cpp:99-160, a partial_sort per row) — this issue should be closed rather than ported. It needs a GPU lease, which this wave did not hold.
Pre-pin
LogprobsTensors.empty_cpu existed unpinned at 5559679229 (vllm/v1/outputs.py:121-132), and this tree mirrors that pre-commit shape faithfully.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Read the logprobs staging paths in src/vllm/v1/sample/sampler.cpp:289-293 and :405-412, then compare them with the GreedyArgmaxScratch pinned allocation at :225-255. Measure the same-binary D2H copy with and without pinned staging at a realistic n × vocab size before changing code. Done means a measurable benefit, a test that verifies the destination uses AllocPinned, and no change if the result is within noise.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- backend, performance
- Issue type
- Refactor
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100