perf: rotate Model Runner V2 output copy streams

[zichen-kuuga]

Summary

Model Runner V2 copies each step's sampled tokens to the host on a single
output_copy_stream (wait_stream(main) → D2H → copy_event.record). On MUSA,
once the next step's wait_stream is queued on that stream, the previous step's
copy_event does not complete until the next step finishes on the GPU (also
observed with query() polling, so it is event completion, not wake-up latency).

Effects:

• every token is published one decode step late → TTFT +1 step;
• get_output alternates between returning instantly and waiting two steps, so
  on every other step the host only launches step N+2 after N+1 drains →
  ~2 ms GPU idle per such step (GPU busy time per step is unchanged).

This patch rotates two output copy streams, one per in-flight step, so a step's
event never has the next step's wait queued behind it.

Evidence (Qwen3-VL-8B-Instruct, TP1)

workload	before (TTFT / TPOT)	after (TTFT / TPOT)
random 2048/1024, rr=1.2, 100 prompts	262 / 40.9 ms	187.56 / 31.05 ms

Decode trace (32 streams): output-event wait 55.4 → 26.3 ms median,
GPU idle 5.24 → 3.39 ms/step, P99 ITL 64.7 → 33.5 ms.
Greedy outputs identical before/after.

[yeahdongcn]

Merge verdict: 1 blocking finding — not ready to merge.

Round status

• Earlier rounds: gh pr view 263 --json reviews and --json comments both return empty arrays, so this is the first review — there is no earlier P0/P1 to re-check.
• Checks: gh pr checks 263 → no checks reported on the 'feat/copy_stream' branch. Per round policy that does not block, but note there is also no CI signal for this patch.
• Head reviewed: 181297374ed253d587b0369f1e321e34c39f2184 — identical to the dispatch snapshot, so nothing was reviewed out of date.
• The change is a two-file patch-series PR: vllm_musa/patches/series/0172-MUSA-rotate-V2-runner-output-copy-streams.patch (new, 48 lines) and the README.md patch-count/summary update. Line references below are into the patch file at the head commit; where I cite the patched file I give upstream line numbers from the pin VLLM_COMMIT=2cf0a6915ce544dc493a0990f2ea38d81601128a (v0.28.0, from third_party/PINS) and say so.

Verified mechanically: the series is now 172 .patch files + README.md (173 dir entries), matching the README's Currently **172 patches**; no quilt series manifest exists, so build_apply.py::series_files falls back to the sorted *.patch glob and 0172 is applied last; git apply --recount -p 1 (the build applier) ignores hunk line counts, and all three hunk contexts match the pinned source with a consistent offset (+192 / +1818 / +1894 in patch headers vs the pinned output_copy_stream block at 194, the AsyncOutput construction at 1782–1790, and return async_output + def take_draft_token_ids at 1861–1863).

P1 — the rotation is hardcoded to two streams, but a V2 runner can have pp_size + 1 steps in flight

vllm_musa/patches/series/0172-MUSA-rotate-V2-runner-output-copy-streams.patch:18-24

+        # MUSA: a copy-stream wait queued for step N+1 delays step N's copy event
+        # until step N+1 finishes, so each in-flight step gets its own stream.
+        self._output_copy_streams = [
+            self.output_copy_stream,
+            torch.cuda.Stream(self.device),
+        ]
+        self._output_copy_stream_idx = 0

Trigger. The number of concurrently live AsyncOutput objects is bounded by vllm_config.max_concurrent_batches:

• vllm/v1/engine/core.py:209 — self.batch_queue_size = vllm_config.max_concurrent_batches, and step_with_batch_queue (pinned core.py:624+) keeps that many unconsumed futures, each created eagerly in the worker (the AsyncOutput.__init__ that queues the copy and records the event runs at sample_tokens call time; get_output() runs only when the queue is popped).
• vllm/config/vllm.py:549-560 — for the V2 model runner with async scheduling this is pp_size + 1.

So with the V2 runner + pipeline_parallel_size >= 2 (async scheduling is on by default for non-pooling models), batch_queue_size >= 3. The round-robin then hands step 1 → streams[1], step 2 → streams[0], step 3 → streams[1], while step 1's AsyncOutput is still unconsumed. Step 3's wait_stream(main) is queued on a stream that still holds step 1's copy_event, which is exactly the condition the patch's own comment and the PR body ("one per in-flight step") say delays step 1's event — so on the PP path the workaround silently loses its effect and tokens are again published behind a later step.

Suggested change. Size the list from the value the same __init__ already reads ~11 lines later for its UVA pools (pinned vllm/v1/worker/gpu/model_runner.py:203, set_default_max_concurrency(vllm_config.max_concurrent_batches)):

self._output_copy_streams = [self.output_copy_stream]
self._output_copy_streams.extend(
    torch.cuda.Stream(self.device)
    for _ in range(max(1, vllm_config.max_concurrent_batches - 1))
)

This is not a correctness regression — at pp_size == 1, max_concurrent_batches is 2 and the current code is exactly right — but the patch's stated invariant does not hold for PP, and the correct bound is already in scope.

P2 — the pooling output path keeps the single shared stream

.../0172-MUSA-rotate-V2-runner-output-copy-streams.patch:32-33 changes only the generation call site:

-            copy_stream=self.output_copy_stream,
+            copy_stream=self._next_output_copy_stream(),

The other consumer of that attribute is untouched: pinned vllm/v1/worker/gpu/model_runner.py:1899-1905 still passes copy_stream=self.output_copy_stream, (line 1904) to AsyncPoolingOutput, whose __init__ repeats the identical pattern (vllm/v1/worker/gpu/async_utils.py:224,241: copy_stream.wait_stream(main_stream) … self.copy_event.record(copy_stream)). Trigger: a pooling engine with pp_size >= 2 (async scheduling is disabled by default for pooling, so max_concurrent_batches = pp_size >= 2), where the same MUSA event behavior applies and the single stream is still reused across in-flight steps. Routing that site through the same helper would close the gap; no test or model-smoke in the PR covers pooling either way.

P2 — the new patch file carries a git format-patch header but no index blob lines

.../0172-MUSA-rotate-V2-runner-output-copy-streams.patch:1-13 has the full header (From 0000000000000000000000000000000000000000 Mon Sep 17 …, From:, Date:, Subject:) and a normal diff --git/--- a//+++ b/ body, but grep -c '^index ' on it returns 0. README.md states the series is generated with git format-patch --no-signature --no-numbered --zero-commit, "keeping index blob lines so git am -3 3-way works across version bumps", and Makefile.sync:14 replays the series with git am -3. Without the index line that patch gets no 3-way fallback at the next VLLM_COMMIT bump and will need manual resolution. The build itself is unaffected (build_apply.py uses git apply --recount -p 1). In a 20-file sample of the series, 0172 is the only file that has a format-patch header and no index line (0145 also lacks one but is a bare --- a/… fragment), so this may just need make -f Makefile.sync format-patches to regenerate the series from its commit rather than a hand-edited file.

P0: no P0 finding.

What I could not verify (assumptions stated)

• The MUSA root cause. I could not confirm on hardware that "once step N+1's wait_stream is queued, step N's copy_event does not complete until step N+1 finishes". AsyncOutput.__init__ only enqueues copy_stream.wait_stream(main_stream) after recording nothing — the record happens last (async_utils.py:138,165) — so on standard stream semantics the already-recorded event should complete independently of later work on that stream; the reported behavior therefore looks like a MUSA/torch_musa stream-wait or event-completion semantic rather than a vLLM bug. I reviewed the patch as reported, not as reproduced, and the patch does do what its description says. A reference to the runtime issue in the PR description (not in shipped comments, per the workspace comment rule) would help a future reader decide whether the workaround can be dropped.
• The performance claim. The 262→187.56 ms TTFT / 40.9→31.05 ms TPOT table and the trace deltas are reported evidence only; this session has no MUSA hardware and I ran no benchmark. The exact launch configuration behind the table is not stated — in particular whether --async-scheduling was on, which is what makes max_concurrent_batches = 2 here. Stating the flags would let a reader reproduce it.
• CI. Quoted in full above: no checks reported on the 'feat/copy_stream' branch — so no build/apply check of the series and no model smoke was read for this PR.

auto-review round 1/3 · verdict: blockers

[yeahdongcn]

Merge verdict: no blocking findings — ready to merge.

No P0 or P1 finding remains at this head, so this ends auto-review for this pull request (round 2 of 3; nothing further will be reviewed here, including later pushes).

Round status

• Head reviewed: fcf25c3bdb10dfc0c40ee70aee63d076d70937b8 ("perf: size V2 output copy streams by max_concurrent_batches"), parent 181297374ed2 — i.e. the branch moved after round 1, and this is the fix push, not the commit round 1 saw.
• Earlier rounds: one review exists (yeahdongcn, round 1, at head 181297374ed253d587b0369f1e321e34c39f2184): one P1 and two P2. All three are addressed at this head, verified in the diff, not taken on trust:
  • P1, rotation hardcoded to two streams while a V2 runner can have pp_size + 1 in flight → the pool is now sized by max_concurrent_batches (vllm_musa/patches/series/0172-MUSA-rotate-V2-runner-output-copy-streams.patch:27-32).
  • P2, pooling output still on the single shared stream → the AsyncPoolingOutput site now calls self._next_output_copy_stream() (same patch, line 62).
  • P2, new patch carried a format-patch header but no index blob line → the regenerated file now has index 79953a7..02ae53b 100644 (line 17).
• Checks: gh pr checks 263 → no checks reported on the 'feat/copy_stream' branch (exit 1, stderr as quoted). Per round policy that does not block, but there is still no CI signal for this patch.
• This round: no P0, no P1; two P2 nits below.

P2 — the README still says "two" copy streams, but the head sizes the ring by max_concurrent_batches

vllm_musa/patches/series/README.md:49-51

per-layer launch count off the eager prefill path. Model Runner V2 rotates two
output copy streams so a step's sampled tokens are not held back by the next
step's copy-stream wait. The fused TileLang `hc_head` is

"two" was accurate for the round-1 revision. At this head the count is max(1, vllm_config.max_concurrent_batches - 1) + 1, which is 2 only when pipeline_parallel_size == 1 (pinned vllm/config/vllm.py:550-560 returns pp_size + 1 for the V2 runner with async scheduling), so on a PP deployment the sentence understates the change; it is also inconsistent with the patch's own commit message, which correctly says "sized by max_concurrent_batches". Suggest "rotates one output copy stream per in-flight step".

P2 — the rotation starts at index 1, so self.output_copy_stream is skipped for the first step

.../0172-MUSA-rotate-V2-runner-output-copy-streams.patch:49-52

+    def _next_output_copy_stream(self) -> torch.cuda.Stream:
+        streams = self._output_copy_streams
+        self._output_copy_stream_idx = (self._output_copy_stream_idx + 1) % len(streams)
+        return streams[self._output_copy_stream_idx]

The counter is incremented before it is used, so the first output copy of a process runs on _output_copy_streams[1] and the pre-existing output_copy_stream at index 0 is not touched until the second step. This is harmless — the cyclic alignment is constant, so consecutive steps still get consecutive streams and the window argument below is unaffected — but it does not match the natural reading of "streams[0] first". Suggested form: stream = streams[self._output_copy_stream_idx % len(streams)], then increment, or start the counter at -1.

What I verified (so the above is not an assumption)

• The stream count is the right bound, not just a bigger number. At the moment step j builds its AsyncOutput, an older object is unconsumed only if its step is still in flight, and the engine allows at most max_concurrent_batches of those: vllm/v1/engine/core.py:209 sets batch_queue_size = vllm_config.max_concurrent_batches, core.py:647 asserts len(batch_queue) < self.batch_queue_size before submitting a new step, and consumption is AsyncModelRunnerOutput.get_output() (the worker's enqueue_output, vllm/v1/executor/multiproc_executor.py:983-986). So at most N-1 older outputs are live, the only older step sharing the new step's stream is j−N, and that one is already consumed. That closes the round-1 P1 concern on the merits for pp >= 2 as well as pp == 1.
• Both call sites are rotated, and the pooling path keeps a correct ring size. AsyncPoolingOutput uses the identical copy_stream.wait_stream(main_stream) → copies → copy_event.record(copy_stream) pattern (vllm/v1/worker/gpu/async_utils.py:223-241). Async scheduling is off by default for pooling, so max_concurrent_batches == pp_size there and the ring degenerates to the previous single stream at pp == 1 — no behaviour change on the common case.
• Nothing else in the tree consumes runner.output_copy_stream. A full-tree search at this head returns only this patch plus pre-existing V1-runner references (0091-...patch:178 comment context, and draft_token_ids_copy_stream in 0150-...patch:43), so no other code is left holding the un-rotated stream.
• The patch anchors to the pinned base. Each hunk's pre-image appears verbatim in upstream vllm/v1/worker/gpu/model_runner.py at VLLM_COMMIT=2cf0a6915ce544dc493a0990f2ea38d81601128a (third_party/PINS) at one consistent offset: +0 for the __init__ hunk (upstream 192–197, output_copy_stream at 194) and +34 for the three later hunks (upstream 1784→1818, 1860→1894, 1901→1935). The +34 matches the earlier series patches that touch this file (0091 at 416, 0095 at 1345, 0097 at 1236–1422, 0119 at 1241), so the series should still apply.
• Series bookkeeping needs no extra edit. 172 *.patch files, unique sorted 4-digit prefixes, the new file starts with the zero-commit From header and From: musa <musa@local>, and README.md:20 says 172 — the three repo tests that police this (tests/test_musa_sync.py::test_series_readme_count_matches_directory, tests/test_manifest_completeness.py::test_series_readme_count_matches_disk, test_series_uses_documented_prefixes_and_canonical_metadata) all pass by construction. In particular no manifest row is needed: manifest.py::_series_entries() derives one entry per series/*.patch from its first diff --git target, so 0172 is auto-classified as category 1 / pre-install and test_series_bijection_with_manifest still holds. I checked this instead of assuming it.
• Comment style. The new block is a 3-line # MUSA: rationale with no ticket, patch-number, or history narration, and both the patch subject and the comment describe the current behaviour only.

What I could not verify (stated as assumptions)

• The MUSA root cause. I have no MUSA hardware and ran nothing. AsyncOutput.__init__ records copy_event last, after the copies (async_utils.py:137-165), so under ordinary stream semantics an already-recorded event should complete independently of later work queued on the same stream; the reported behaviour therefore looks like a MUSA/torch_musa event-completion or stream-wait semantic difference, not a vLLM bug. I reviewed the patch as a workaround that does what its description says, and I cannot confirm the underlying runtime behaviour. A reference to the runtime issue in the PR description (not in the shipped comment, per the comment rule) would let a future reader judge when the workaround can be dropped.
• The performance claim. 262 → 187.56 ms TTFT and 40.9 → 31.05 ms TPOT, and the trace deltas, are author-reported; I ran no benchmark and no trace was attached. The launch configuration behind the table is not stated — in particular whether --async-scheduling was on (that is what makes max_concurrent_batches == 2 and therefore two streams on the measured pp=1 config), plus hardware, driver, and batch shape. Stating them would make the gain reproducible.
• CI. Quoted above in full: no checks reported on the 'feat/copy_stream' branch — no build/apply check of the series and no model smoke was read for this PR.
• One unrelated observation, not caused by this PR: tests/test_musa_sync.py::test_report hardcodes its census ("174 divergences", '1': 112), which matches neither the pre-PR tree (209 total, category 1 = 147) nor this head (210 total, category 1 = 148). That staleness predates the PR; the only census delta this PR introduces is the one new category-1 patch, which is what the derived manifest expects.

auto-review round 2/3 · verdict: clean