Skip to content

[PD] Batch the decode-side KV preallocation and stage its lengths through pinned memory - #38955

Open
hnyls2002 wants to merge 1 commit into
mainfrom
lsyin/pd-prealloc-pinned
Open

hnyls2002 wants to merge 1 commit into
mainfrom
lsyin/pd-prealloc-pinned

Conversation

@hnyls2002

@hnyls2002 hnyls2002 commented Sep 10, 2026

Copy link
Copy Markdown
Collaborator

pop_preallocated / resume_retracted_reqs now decide admission per request on the host, allocate every admitted request's KV in one alloc_extend / alloc_extend_swa_tail call, and only then publish the transfer metadata. The per-request torch.tensor(..., device=cuda) copies that used to synchronize the scheduler stream on every arriving request are gone: lengths go up once through pinned memory (pinned_int64_pair). alloc_extend_swa_tail takes swa_tail_lens per request instead of a single swa_tail_len with a bs == 1 assert.

Follows #38159; the DSV4 NPU allocator keeps its per-request hooks.


CI States

Latest PR Test (Base): ❌ Run #34829774543
Latest PR Test (Extra): ❌ Run #34829774291
Latest PR Test (AMD ROCm 10): ❌ Run #34829774408

Co-authored-by: Jingshu Chen <96519051+elmo2019@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants