[Bugfix][V2] record_stream idx_mapping in the PP draft broadcast - #55745
Conversation
`PPHandler.broadcast_drafts` gathers the draft block on `broadcast_stream` using
`input_batch.idx_mapping`, which the model runner allocates on the main stream
every step (`async_copy_to_gpu` -> `torch.empty_like`). `wait_stream` orders the
gather after the H2D copy, so the indices are correct at launch, but nothing
tells the caching allocator the block is still in use on another stream. The
last rank does not retain `idx_mapping` -- unlike `receive`, which parks it in a
`PendingRecv` -- so it is freed at end of step and the allocator can hand the
block to the next allocation while the gather is still pending. The indices are
then overwritten with unrelated data:
ATen/native/cuda/IndexKernel.cu:111: operator(): block: [0,0,0], thread:
[0,0,0] Assertion `-sizes[i] <= index && index < sizes[i] &&
"index out of bounds"` failed.
usually surfacing via the `pp_broadcast` watchdog as `CUDA error: device-side
assert triggered`, killing EngineCore.
`broadcast`, the sampled-token sibling, already record_streams every main-stream
input it sends. Apply the same guard to the one input `broadcast_drafts` missed.
`draft_tokens` needs none -- it is `RequestStates`' persistent buffer.
Triggering it requires PP>=2, a drafter, and async launches together, which is
why it has gone unnoticed: `CUDA_LAUNCH_BLOCKING=1` serialises the gather ahead
of any reuse and hides it completely.
Tests: Qwen3.6-35B-A3B-AWQ, GSM8K 200 prompts, PP=2, TP=1, MTP k=1, V2 runner,
2x H200, max_tokens=2048, same tree and venv either side of the patch:
before: device-side assert 6/6 runs, 148s to 258s in
after: 200/200 clean 4/4 runs, mean_acceptance_length
1.9380 / 1.9373 / 1.9354 / 1.9357
Acceptance is unchanged within noise against the two configurations that were
already clean -- 1.9352 with CUDA_LAUNCH_BLOCKING=1 and 1.9375 at PP=1 -- as
expected for a change that only affects allocator lifetime.
tests/v1/worker/test_pp_utils.py: 4 passed.
Signed-off-by: Chris Eastwood <chris.eastwood@pwn4g3.dev>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Team Run ID: 📒 Files selected for processing (1)
Included review availability: Your plan provides up to 10 included reviews per hour; 9 remain after this review. 📝 SummarySummary by CodeRabbit
Walkthrough
ChangesBroadcast stream safety
Estimated code review effort: 2 (Simple) | ~5 minutes Merge Risk: ⚪ Minimal · up to This fixes premature CUDA allocator reuse of draft-broadcast indices, preventing pipeline-parallel speculative decoding failures without changing draft-token handling. The current change is ready to merge. Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
@yewentao256 @njhill This issue blocks full verification that I have been running for closing out #46994 |
|
👋 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. 🚀 |
yewentao256
left a comment
yewentao256
left a comment
There was a problem hiding this comment.
LGTM, thanks for the work!
|
/ci run |
|
✅ @eastwood-c, CI is now available for this PR.
|
|
✅ Triggered Buildkite CI #87720 for commit |
Signed-off-by: Nick Hill <nickhill123@gmail.com>
njhill
left a comment
njhill
left a comment
There was a problem hiding this comment.
Thanks @eastwood-c. I made a small adjustment, including removing another record_stream which is unnecessary (since the send tensor is allocated and used on the same stream).
|
/ci run |
|
✅ Triggered Buildkite CI #87949 for commit |
|
@njhill Nice catch, I thought there was something different about the targets so I hadn't removed that. I'll take some time and keep familiarizing myself with the structures. |
Resolve pp_utils.py: upstream vllm-project#55745 re-added a broadcast_drafts near the top of PPHandler; keep this branch's single definition (the one with the disabled/max_sample_len guards and the main-stream gather). Co-authored-by: Kimi Code <noreply@moonshot.cn> Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>
…m-project#55745) Signed-off-by: Chris Eastwood <chris.eastwood@pwn4g3.dev> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com> Signed-off-by: Jyotirmoy Roy <jyotirmoyroy649@gmail.com>
3 participants
Purpose
PPHandler.broadcast_draftsreads a main-stream tensor on the broadcast stream without arecord_streamguard, so the CUDA caching allocator can recycle that tensor's memory while the gather is still in flight. The recycled indices then land out of bounds:usually surfacing through the
pp_broadcastwatchdog asCUDA error: device-side assert triggered, which kills EngineCore mid-run.The site, in
vllm/v1/worker/gpu/pp_utils.py:input_batch.idx_mappingis allocated on the main stream every step (async_copy_to_gpu->torch.empty_like). Thewait_streamorders the gather after the H2D copy, so the indices are correct at launch; what is missing is the other half of the contract. The last rank never retainsidx_mapping-- unlikereceive, which parks it in aPendingRecv-- so it is freed at end of step and the allocator may hand the block to the next allocation while the gather is still pending onbroadcast_stream.broadcast, the sampled-token sibling a few lines below, alreadyrecord_streams every main-stream input it sends. This applies the same guard to the one inputbroadcast_draftsmissed.draft_tokensneeds none -- it isRequestStates' persistent buffer, not a per-step allocation.Reproducing it requires PP>=2, a drafter, and async launches all at once. That is why it has gone unnoticed since #50514: a PP run without a speculator never calls
broadcast_drafts, andCUDA_LAUNCH_BLOCKING=1hides it completely by serialising the gather ahead of any reuse. EAGLE3 + PP is exposed to it, and so is MTP + PP once a drafter is reachable there.Not a duplicate: no open PR or issue touches
pp_utils.py,broadcast_drafts, or this assert. Searchedpp_utils,record_stream broadcast_drafts, andIndexKernel index out of bounds pipeline parallelacross open PRs and issues.Test Plan
Unit:
End-to-end, the configuration that reproduces the assert -- GSM8K, 200 prompts, greedy,
max_tokens=2048,VLLM_USE_V2_MODEL_RUNNER=1, 2x H200:Run before and after the patch on the same tree, same venv, same box, otherwise idle. Acceptance is measured as
1 + accepted/draftsper prompt, mirroringtests/v1/e2e/spec_decode/acceptance_rates/.Test Result
tests/v1/worker/test_pp_utils.py: 4 passed.End-to-end, same cell before and after:
mean_acceptance_lengthafter the fix: 1.9380, 1.9373, 1.9354, 1.9357.Unchanged within noise against the two configurations that were already clean, as expected for a change that only affects allocator lifetime:
CUDA_LAUNCH_BLOCKING=1(assert masked)The same assert was also observed on zai-org/GLM-5.2-FP8 at TP=4/PP=2 with MTP k=1 and k=2, so it is not specific to one model or quantisation.
AI assistance was used for this change: to investigate the failure, isolate the faulting path, and draft the patch and this description. The change and its test results have been reviewed line by line before submission.