Repository navigation
[MRV2][PP] Defer sampled-result receives - #53948
chengchengpei wants to merge 9 commits into
Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
9cdea89 to
62ff6ff
Compare
|
This pull request has merge conflicts that must be resolved before it can be |
929df86 to
bccca9f
Compare
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Team Run ID: 📥 CommitsReviewing files that changed from the base of the PR and between 42c50a73c4a25d4c14572936a010fd24a1608214 and 5e19e157001a5a68c6f91d9af26db15323857357. 📒 Files selected for processing (2)
Included review availability: Your plan provides up to 10 included reviews per hour; 9 remain after this review. WalkthroughAdds configurable deferred pipeline-parallel sampled-token receives. The change validates environment settings, queues and schedules receives, handles idle and shutdown flushing, integrates GPU execution paths, and adds configuration, environment, and PP utility tests. ChangesDeferred PP sampled-token receives
Estimated code review effort: 4 (Complex) | ~45 minutes Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant GPUModelRunner
participant PPHandler
participant PPBroadcast
GPUModelRunner->>PPHandler: queue receive slot
PPHandler->>PPHandler: advance receive queue
PPHandler->>PPBroadcast: launch sampled, combined, and draft broadcasts
PPHandler->>GPUModelRunner: wait for receive event
Merge Risk: ⚪ Minimal · up to This adds an opt-in deferred pipeline-parallel receive path while retaining immediate receives by default. Supported configurations are validated and the covered lifecycle paths show no current merge-blocking risk. 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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 |
5756974 to
d0096e6
Compare
|
This pull request has merge conflicts that must be resolved before it can be |
|
This pull request has merge conflicts that must be resolved before it can be |
7212070 to
32db8ce
Compare
2d1b125 to
1839e51
Compare
|
/ci run |
|
✅ Triggered Buildkite CI #91940 for commit |
yewentao256
left a comment
yewentao256
left a comment
There was a problem hiding this comment.
Thanks, Generally LGTM, also want an approval from @njhill or @WoosukKwon
|
This pull request has merge conflicts that must be resolved before it can be |
Signed-off-by: Chengcheng Pei <5881383+chengchengpei@users.noreply.github.com>
Signed-off-by: Chengcheng Pei <5881383+chengchengpei@users.noreply.github.com>
Signed-off-by: Chengcheng Pei <5881383+chengchengpei@users.noreply.github.com>
Signed-off-by: Chengcheng Pei <5881383+chengchengpei@users.noreply.github.com>
Signed-off-by: Chengcheng Pei <5881383+chengchengpei@users.noreply.github.com>
Signed-off-by: Chengcheng Pei <5881383+chengchengpei@users.noreply.github.com>
Signed-off-by: Chengcheng Pei <5881383+chengchengpei@users.noreply.github.com>
1839e51 to
20d6804
Compare
|
/ci run |
|
✅ Triggered Buildkite CI #92748 for commit |
@yewentao256 @njhill @WoosukKwon Can someone approve? thanks |
Keep the state-neutral post-update warmup in warmup.py and preserve both upstream prefill and deferred-receive regression coverage after rebasing. Co-authored-by: Codex <noreply@openai.com> Signed-off-by: Chengcheng Pei <5881383+chengchengpei@users.noreply.github.com>
20d6804 to
c8d9fb1
Compare
CI selector (shadow): 133 test steps (174 jobs) instead of 75 (91 jobs)Shadow mode: this changes nothing about what CI runs. It shows what the evidence-based selector would pick for this PR, next to today's rules. How it works. Feedback welcome: reply here if it would skip a step this change needs, or runs something unrelated.
Selector would run (133)
Would skip (today's rules run them) (14)
Would add (today's rules do not run them) (72)
AMD mirrors: would skip (8)
AMD mirrors: would add (64)
11 changed files · base |
Keep async scheduling opt-in for callers using the original three arguments. Existing explicit GPU configuration still enables deferred receives only on CUDA after sync-free warmup. Cover legacy calls, synchronous scheduling and non-CUDA fallback. Co-authored-by: Codex <noreply@openai.com> Signed-off-by: Chengcheng Pei <5881383+chengchengpei@users.noreply.github.com>
|
/ci run |
|
✅ Triggered Buildkite CI #92751 for commit |
yewentao256
left a comment
yewentao256
left a comment
There was a problem hiding this comment.
Thanks for the work!
I took a further look, the code diff can be further simplified so that it is easier to land, idealy < 400 LOC.
Please also take a check at runtime collective_rpc, I am not sure if there might be a dead lock.
Thanks for the reviews. Repeated rebases and validation have become too costly for me to continue incrementally. I’m pausing active work until the relevant reviewers can consolidate the remaining blocking feedback and agree on the required validation. I can then decide whether to complete one final revision. |
|
Thanks @chengchengpei that sounds like a good plan. Apologies for the delays, there is a bit of a review bottleneck now especially since it's much quicker to create PRs than review them. I will try to look at this properly asap. My high-level comment is that (as @yewentao256 said) it still needs quite a bit of simplification, but will try to respond with some more concrete suggestions soon. We did recently merge this fix #58542 which avoids this problem entirely in the prefill-only case, which may be most common. |
sure. thanks. take your time. we deploy this PR (with latest main branch, including #58542) internally because this can improve throughput by about 10-40% on our workload. hope it can be merged so we can upgrade our internal deployment. |
|
This pull request has merge conflicts that must be resolved before it can be |
This branch has not been deployed
5 participants
Purpose
Related to #53810.
On non-last pipeline-parallel stages, sampled-result receives are currently
posted as soon as their buffers are allocated, although the results are not
consumed until a later scheduler step. For small broadcasts, the receiver-side
NCCL kernel can remain resident while model kernels execute.
This change delays those receives on CUDA Model Runner V2 workers until the
last safe pipeline step and places them after that step's model work. The last
PP rank remains the immediate broadcast source. A result produced at step
Tis still consumed at
T + PP; only the receiver launch moves.Activation and scope
The behavior is automatic when all of these are true:
Synchronous scheduling, CPU, ROCm, single-stage execution, eager execution
without JIT warmup, and diagnostic modes that may synchronize the device retain
immediate receive posting. There are no environment-variable gates.
async_schedulingis resolved fromNoneto a concrete boolean duringVllmConfigvalidation, beforePPHandleris constructed.Correctness and lifecycle guarantees
after warmup finishes.
FIFO collective order.
queuing them on the potentially backlogged broadcast stream, so later steps
cannot overwrite an in-flight payload.
traffic. Its broadcasts are serialized on one dedicated stream; receiver
deferral can leave the head send waiting, but does not make all queued sends
simultaneously resident.
pause/sleep boundaries before device synchronization. The zero-token worker
flush and model-runner launch are intentionally redundant: the latter also
protects direct runner invocations.
launchedflag because their eventrecording is a no-op.
idx_mapping=-1sentinel and a zero accepted-token count, preventing model-specific state
handlers such as RecoverSSM from committing stale state.
including Mamba sampled-count/alignment kernels and RecoverSSM commit
specializations. The all-
-1mapping and zero accepted counts make thiswarmup state-neutral.
Sender-side validation envelope
The H100 W1 runs used CUDA 13.3 and NCCL 2.30.7 with TP8/PP4, concurrency 710,
and no speculative draft tokens. The two per-step broadcasts were about 5.7
KiB each. NCCL protocol selection was left automatic and was not printed by
the logs, so this PR does not claim a specific LL/LL128/Simple protocol or a
validated payload-size cutoff. Larger speculative payloads retain the same
ordering and flush invariants in tests, but still need hardware performance
validation.
Validation
Focused tests cover receive cadence, immediate/deferred launch idempotency,
FIFO recovery, speculative draft propagation, invalidated-row masking, idle
flushing, CPU no-op events, warmup specialization and state neutrality, and
pause-before-synchronize ordering.
A matched same-base H100 TP8/PP4 A/B was run on a behavior-equivalent
no-speculation revision of this PR with 4,000 requests, and isolated cold caches. Both arms completed
4,000/4,000 requests with zero failures:
Automatic deferral improved output and total throughput by 38.5%, reduced
mean TTFT by 41.2%, and reduced mean TPOT by 25.7%. The deferred run had no
errors, tracebacks, NCCL warnings, or CUDA errors. This run also exercised the
immediate last-rank sender under the validated payload envelope without a hang
or observed throughput regression.
The current head additionally hardens mapping dtypes, logical padded-buffer
views, speculative draft snapshots, runtime-sync gating, RecoverSSM warmup, and
invalidated-row handling. Those safeguards are covered by focused tests and
static checks; they have not yet been rerun in the cluster A/B.
AI assistance was used to analyze failures and prepare the changes; the final
implementation and results should be reviewed by the author.