[Bugfix][V2] Wait for the prior encoder-only step before reusing its input buffers

[gty111]

Purpose

An encoder-only instance can overwrite the input metadata of a step that is
still running on the device, which surfaces as
CUDA error: an illegal memory access was encountered at an unrelated
stream_synchronize.

UvaBufferPool hands out pinned host buffers that the device reads in place,
and holds exactly max_concurrent_batches of them
(model_runner.py, set_default_max_concurrency). A slot therefore comes up for
reuse as soon as the device falls one step behind, and copy_to_uva() writes
into it with no synchronisation — by design, because the batch queue bounds how
many steps are in flight.

That bound only reaches the device if each step has a point where the host waits
for it. Sampling instances do: AsyncOutput records copy_event on a copy
stream that has done copy_stream.wait_stream(main_stream), so the event is
ordered after every kernel of that step, and get_output() blocks on it before
the step's future resolves. The batch queue pops step N during step N+1 while
step N's slot is not reused until step N+2, leaving one step of margin.

An encoder-only step has no such point. It returns
make_empty_encoder_model_runner_output(), built on the host from
scheduler_output, with no D2H copy and no wait, so its future resolves as soon
as the host has finished launching. The host can then run arbitrarily far ahead
while slots keep being recycled every max_concurrent_batches steps.
_apply_write_kernel then loads a stale row_idx from a recycled slot, computes
row_ptr = base + row_idx * row_stride + start_idx, and stores through it.

Grepping the V2 step path confirms that copy_event is the only host/device
synchronisation per step: the three torch.accelerator.synchronize() calls in
gpu/model_runner.py are in profile_run() and shutdown(). One of them is
already if self.is_encoder_only: torch.accelerator.synchronize() — the profile
path needed the same barrier the step path is missing.

V1 is unaffected: synchronize_input_prep() waits and records around input
preparation for every step, so it never depends on the output path for ordering.

The fix arms one event for encoder-only instances, waits on it before the step's
first host write and records it after the last one. Sampling instances keep a
synchronisation-free input path (encoder_input_reuse_event stays None).

Test Plan

Repro: two GB200 nodes, Qwen3.5-35B-A3B, disaggregated encoder
(--mm-encoder-only producer on one node, consumer on the other), V2 runner on
both, MuirBench through the EPD proxy, 8 client processes x 96 concurrent
requests, GPU core dumps enabled (CUDA_ENABLE_COREDUMP_ON_EXCEPTION=1).

.venv/bin/python -m pytest tests/v1/worker/test_gpu_block_table.py \
    tests/v1/worker/test_gpu_input_batch.py -q
pre-commit run ruff-check ruff-format check-torch-cuda-call --files \
    vllm/v1/worker/gpu/model_runner.py
pre-commit run mypy-3.12 --hook-stage manual --files \
    vllm/v1/worker/gpu/model_runner.py

No unit test: the change is a device event in the step loop, and covering it
would mean standing up a full GPUModelRunner. V1's equivalent
(synchronize_input_prep) is likewise covered only end to end. The evidence is
the repro below.

Test Result

• pytest: 19 passed. ruff-check / ruff-format / check-torch-cuda-call /
  mypy-3.12: all Passed.
• Before: the encoder dies mid-run, reproducibly. Four runs, all with
  _apply_write_kernel in the log: 9,532 encoder requests with the mooncake EC
  connector, then 15,275 / 15,279 / 23,542 with a CPU/NIXL one — 7m29s to 8m52s
  of load each. A GPU core dump puts the fault at that kernel's tl.store, all
  32 lanes of the warp on one invalid base address (a garbage row_idx, not an
  index merely out of range), in the last block of the grid in both dumps — the
  block whose pid-indexed metadata is loaded latest.
• After: 20 passes, 103,409 encoder requests, 24,000/24,000 requests OK,
  no core dump, engine alive at the end. That clears the observed crash window by
  4.4x in requests and ~2.4x in wall clock.
• Accuracy: MuirBench 0.5900–0.5975 across the 20 passes, matching pre-fix
  runs — no output change.
• Throughput: plateau 20.21 req/s, stable over the last five passes
  (20.70 / 20.92 / 20.59 / 20.12 / 20.11), against 18.4–18.7 for the same
  configuration before the fix. The wait only blocks when the device is already
  a step behind, which the pool sizing already implies.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[WoosukKwon]

Can you be more exact about when the race condition happens?

[WoosukKwon]

Basically, the idea behind the current design is that we can avoid race condition without such a barrier because async scheduling bounds the maximum concurrency. If I understand correctly, race condition shouldn't happen unless we accidentally reuse the buffer for more than one data in the same step.