[Performance]: [MRV2][PP]: Post sampled-result receives closer to T+PP consumption

[chengchengpei]

Proposal to improve performance

Summary

Model Runner V2 pipeline parallelism sends sampled results from the last PP
rank back to the non-last ranks. The CUDA path uses a sibling process group and
a side stream, as introduced by
#42187. However, each
non-last rank posts its receive in local feedback step T, while that result
is not consumed until T + pp_size.

If the last rank has not produced the result, the receiver-side NCCL broadcast
kernel may remain active across multiple intervening model executions. A side
stream removes default-stream ordering, but it does not isolate GPU execution
or communication resources. The pending kernel may therefore interfere with
model work. If its event is still incomplete at T + pp_size, the main stream
must also wait before applying the sampled result.

I propose evaluating an opt-in, just-in-time receiver policy:

• keep the last rank's source broadcast at its current point;
• save the receive metadata at step T;
• post non-last-rank receives closer to consumption, initially at
  T + pp_size - 1; and
• preserve the current T + pp_size event wait as the correctness fallback.

For pp_size=4:

Current:
  T       post recv [possibly resident----------------------] consume at T+4

Candidate:
  T       save receive metadata
  T+3                                      post recv [-------] consume at T+4

This changes receive placement only. It does not weaken the autoregressive
dependency or allow decode to proceed without the sampled token.

This issue currently presents a source-based proposal and measurement plan.
A public-model reproducer and sanitized trace summary are still pending.

Current source behavior

On public commit
d591d1d5:

• PPHandler.receive()
  immediately posts sampled-token and count broadcasts on
  broadcast_stream;
• the non-last-rank FIFO is pre-seeded with pp_size entries; and
• get_prev_sampled_outputs()
  later consumes the slot and calls main_stream.wait_event(slot.event).

The delay is a local PP feedback-step cadence, not a fixed wall-clock interval.
Variable stage times can therefore leave a result incomplete at consumption.

Why a side stream may not be sufficient

PR #42187 correctly removes stream and communicator serialization from normal
PP traffic. An unmatched receiver-side collective may nevertheless remain
resident while waiting for the last rank to publish its data.

Closed PR
#47686 describes related
early receiver-side LL-kernel spinning and proposes packing feedback fields
into fewer collectives. Packing reduces the number of kernels; later receive
placement would reduce their lifetime. The ideas are complementary.

This proposal distinguishes two costs:

1. Pending-receive interference: a receiver collective remains active while
   independent model work runs.
2. Correctness wait: the main stream waits at T + pp_size when the result
   is genuinely late.

Just-in-time posting primarily targets the first cost. It must not hide or
weaken the second.

Why queue depth is a separate mechanism

MRV2 has two related depths:

EngineCore async PP queue:              pp_size + 1 batches
Sampled-result feedback/consume delay:  pp_size steps

The extra EngineCore entry overlaps host scheduling with worker execution. It
does not add a PP compute stage, move the feedback consume point, or shorten an
already-posted receive.

Suggested public reproducer

Use a public PP-compatible model and fixed synthetic prompts. Compare otherwise
identical workloads:

1. max_tokens=1: the request finishes after final-prefill sampling, so no
   sampled result is needed for another decode step;
2. max_tokens=2: the first sampled token must return before the next decode
   step.

Controls:

• Model Runner V2;
• pipeline_parallel_size > 1;
• asynchronous scheduling;
• fixed input length and --ignore-eos;
• prefix caching disabled;
• identical concurrency and request count; and
• enough requests to exclude startup and drain effects.

Capture Nsight Systems on at least the first and last PP ranks. Record:

• receive-post time on each non-last rank;
• source-broadcast time on the last rank;
• receive readiness at the next step and at T + pp_size;
• time blocked by main_stream.wait_event();
• lifetime of the selected NCCL broadcast kernel; and
• model-kernel duration with and without a pending receive.

Compare immediate and deferred receive placement while holding all other
settings fixed. Only public commands, public-model inputs, sanitized
collect_env.py output, and redacted summary measurements should be added to
this issue.

Proposed implementation scope

The first implementation should be small and guarded:

1. Keep the last rank's source broadcast unchanged.
2. At feedback step T, save receive shape, request-slot mapping, liveness
   mask, and generation snapshot without posting the receive.
3. Post every non-last rank's receive in FIFO order at a common logical delay,
   initially pp_size - 1.
4. Consume the same slot at T + pp_size and retain the existing event wait.
5. Preserve collective ordering through ordinary, empty, skipped, and finite
   drain steps.
6. Enable the experiment only on a validated CUDA/NCCL path; retain immediate
   placement elsewhere.

Correctness invariants and risks

The implementation must preserve:

• identical collective order, shape, row count, and dtype on all ranks;
• association between step-T feedback and its saved request-slot mapping;
• generation/liveness checks across abort, preemption, and slot reuse;
• the existing T + pp_size decode-eligibility cadence;
• non-final chunked-prefill behavior;
• no unnecessary receive for requests that finish after the first sample;
• explicit drain before shutdown or communicator destruction; and
• safe CUDA-graph warmup and capture boundaries.

Receiver deferral can move contention to the source rank if the source
broadcast is posted much earlier than its receivers. Evaluation must therefore
measure both source- and receiver-side kernel residency. Other risks include
collective-order mismatch on empty steps, insufficient lead time on slower
topologies, abort cleanup, and backend-specific stream/event behavior.

Acceptance criteria

• Identical deterministic outputs to current main.
• No deadlock or collective mismatch during startup, steady state, empty
  steps, abort, preemption, finite drain, or shutdown.
• Lower receiver-side NCCL residency without harmful source-side residency.
• No increase in hard waits at T + pp_size.
• No material regression at low concurrency or for one-output requests.
• Reproduced throughput and latency results across multiple allocations.
• Validation on at least PP2 and PP4 before considering a default.

Related upstream work

• #42187, merged: introduced
  the current MRV2 feedback cadence and sibling stream/process group.
• #47686, closed: proposed
  packing feedback fields and documented early receiver-side LL-kernel
  spinning.
• #50853, open RFC:
  broader MRV2 PP composition and deferred-result transport.
• #50410, open prototype:
  balances admission across PP batches; it changes work availability rather
  than receive placement.
• #51650, merged: overlaps
  sampled-token broadcast with compute on XPU.

I did not find another open issue or PR specifically proposing delayed posting
of the CUDA MRV2 sampled-result receive.

Report of performance regression

This is not presented as a regression. PR #42187 is necessary and improves PP
overlap. This proposal targets possible residual contention when asynchronous
receives are posted substantially before source data exists.

Misc discussion on performance

The intended experiment changes only receiver timing. The sampled-result
dependency, consume cadence, and correctness wait remain unchanged.

Current environment

A sanitized collect_env.py result will be added with the public reproducer.
No private hostnames, filesystem paths, scheduler identifiers, credentials,
model artifacts, or internal traces are included here.

Validation status

• Searched related public issues and PRs.
• Confirmed the referenced source behavior at the pinned public commit.
• Separated receive placement from queue depth and admission balancing.
• Add a public-model synthetic reproducer.
• Add sanitized collect_env.py output and a redacted trace summary.
• Rebase the prototype onto current main.
• Add tests for FIFO ordering, empty steps, abort/reuse, and finite drain.
• Validate PP2 and PP4 before proposing an upstream default.