[Bugfix][DP] Synchronize the device on pause completion

[aoshen02]

Addresses the third problem in #51476: pause_generation() resolves before the device is idle — and in DP, the pause path is itself what makes the device busy. Companion PR #52957 (independent, no shared commits) removes the DP dummy-batch burst; this PR makes the pause completion contract correct regardless of how much work is in flight.

Purpose

1. Pause completion is a CPU-side predicate. The pause future resolves off has_work() and the DP consensus is a gloo CPU all-reduce; nothing on the pause path touches a CUDA stream. sleep() is safe only because GPUWorker.sleep synchronizes before unmapping — bare pause is sleep minus the synchronizing half. RL callers pause, run their own weight-sync kernels/NCCL, then sleep; that work races whatever the engine launched last. Reported from production, surviving today only on the incidental ~0.5s between pause and sleep.

2. The pause path launches unwaited GPU work. DPEngineCoreProc._pause_complete kick-starts engines_running = True so every rank reaches the consensus all-reduce in _has_global_unfinished_reqs — and the fixed every-32-step cadence then forces 32 dummy forwards per rank per pause call on an idle engine. The final dummy batch is launched in the consensus iteration itself, is not in batch_queue, and is invisible to has_work(); one iteration later the future resolves while it can still be executing. wait_for_requests_to_drain() cannot compensate: the coordinator only tracks request-driven waves, so during the burst it publishes "not running" and drain returns instantly.

Measured (DP=2 + EP, PowerMoE-3b, dummy batch inflated to ~110ms to model a large-MoE step): pause_generation() on an idle engine takes 3.40s and returns with the device still busy.

Change

A device barrier at pause completion. WorkerBase.synchronize_device() synchronizes through the device-generic torch.accelerator API whenever an accelerator is registered — the same barrier sleep() relies on — so CUDA, XPU and any other torch.accelerator backend get it without an override, and it stays a no-op only where there is genuinely no accelerator (CPU); backends outside torch.accelerator must override with their own wait. EngineCore._finish_pause(): collective_rpc("synchronize_device"), then cache reset — synchronizing first also stops clear_cache=True from resetting state the last launched batch may still be using. Called from all three completion paths (in-proc, fast path, deferred idle callback). collective_rpc waits for all workers, which also closes execute_dummy_batch waiting only for the output rank.

The barrier does not shrink the burst itself: an idle-engine DP pause still runs 32 dummy forwards, it just waits for them honestly. #52957 cuts the burst to a single step with a one-condition cadence change; the two PRs compose but neither depends on the other.

The same contract is applied to the mirror transition: GPUWorker.wake_up now ends with torch.accelerator.synchronize(). sleep synchronizes on both sides of the memory transition, wake_up on neither — its tail work (buffer restores after level-2 sleep, post_kv_cache_wake_up's re-init kernels) was unwaited, so a caller could start same-device work (e.g. RLHF weight sync right after wake) racing it. Covered by the existing sleep/wake e2e cases in the regression suite below.

Reproduction (scripts included)

Requires 2 GPUs and vllm/vllm-openai:nightly. Save the three files below into a directory $W, then:

# baseline: the bug
docker run --rm --gpus '"device=0,1"' --ipc=host --shm-size=16g \
  -v ~/.cache/huggingface:/root/.cache/huggingface -v $W:/patch \
  --entrypoint bash vllm/vllm-openai:nightly /patch/gpurepro.sh instr

# with this PR: check out the branch, then
git diff main...HEAD -- vllm/ > $W/fix.patch
docker run --rm --gpus '"device=0,1"' --ipc=host --shm-size=16g \
  -v ~/.cache/huggingface:/root/.cache/huggingface -v $W:/patch \
  --entrypoint bash vllm/vllm-openai:nightly /patch/gpurepro.sh fixed

(If NCCL fails at init on a CUDA-13.1 driver, add -e NCCL_NVLS_ENABLE=0.)

gpupatch.py — instrumentation only: inflates the dummy batch to ~110ms to model a large-MoE step, and adds a device_idle() probe RPC

#!/usr/bin/env python3
"""Usage: gpupatch.py <site-packages> — instruments a vLLM install."""
import sys

site = sys.argv[1]

def sub(path, old, new):
    with open(path) as f:
        src = f.read()
    if new in src:
        return
    assert old in src and src.count(old) == 1, f"pattern not found/unique in {path}"
    with open(path, "w") as f:
        f.write(src.replace(old, new))

OLD = """\
    def execute_dummy_batch(self) -> None:
        num_tokens = getattr(self.model_runner, "uniform_decode_query_len", 1)
        self.model_runner._dummy_run(num_tokens, uniform_decode=True)
"""
NEW = """\
    def execute_dummy_batch(self) -> None:
        num_tokens = getattr(self.model_runner, "uniform_decode_query_len", 1)
        self.model_runner._dummy_run(num_tokens, uniform_decode=True)
        import os as _os
        _cycles = int(_os.environ.get("REPRO_SLEEP_CYCLES", "0"))
        if _cycles:
            # Simulate a large-MoE forward: keep the stream busy after launch.
            torch.cuda._sleep(_cycles)

    def device_idle(self) -> bool:
        return bool(torch.cuda.current_stream().query())
"""
sub(f"{site}/vllm/v1/worker/gpu_worker.py", OLD, NEW)
print("instrumented")

repro_gpu.py — the scenario: pause / drain / probe

import asyncio
import time

from vllm import SamplingParams
from vllm.engine.arg_utils import AsyncEngineArgs
from vllm.v1.engine.async_llm import AsyncLLM


async def main():
    engine = AsyncLLM.from_engine_args(
        AsyncEngineArgs(
            model="ibm-research/PowerMoE-3b",
            data_parallel_size=2,
            enable_expert_parallel=True,
            gpu_memory_utilization=0.5,
            max_model_len=512,
            enforce_eager=True,
        )
    )

    async def gen(rid):
        async for _ in engine.generate(
            "The capital of France is", SamplingParams(max_tokens=8), rid
        ):
            pass

    async def probe(tag):
        idle = await engine.collective_rpc("device_idle")
        print(f"PROBE {tag}: device_idle={idle}", flush=True)
        return all(idle)

    await gen("warmup")
    await asyncio.sleep(6)
    await probe("settled")  # sanity: True when quiet

    # scenario 1: pause on a quiet engine
    t0 = time.perf_counter()
    await engine.pause_generation(mode="abort")
    t1 = time.perf_counter()
    await probe("after_pause1_returned")
    print(f"S1 pause took {t1 - t0:.2f}s on an idle engine", flush=True)
    await engine.wait_for_requests_to_drain(60)
    await probe("after_drain1_returned")

    # scenario 2: second pause+drain right before sleep (the production recipe)
    t0 = time.perf_counter()
    await engine.pause_generation(mode="abort")
    t1 = time.perf_counter()
    await probe("after_pause2_returned")
    await engine.wait_for_requests_to_drain(60)
    await probe("after_drain2_returned")
    print(f"S2 second pause took {t1 - t0:.2f}s", flush=True)

    # scenario 3: drain called while the pause burst is mid-flight
    await engine.resume_generation()
    await gen("rearm")
    await asyncio.sleep(6)
    task = asyncio.create_task(engine.pause_generation(mode="abort"))
    await asyncio.sleep(0.8)
    await engine.wait_for_requests_to_drain(60)
    idle = await probe("drain_during_pause_burst")
    print(f"S3 drain returned while pause_done={task.done()}, idle={idle}", flush=True)
    await task

    engine.shutdown()


asyncio.run(main())

gpurepro.sh — container entrypoint

#!/bin/bash
# usage: gpurepro.sh instr|fixed
set -x
SITE=/usr/local/lib/python3.12/dist-packages
if [ "$1" = "fixed" ]; then
  (cd "$SITE" && patch -p1 --forward) < /patch/fix.patch || exit 89
fi
python3 /patch/gpupatch.py "$SITE" || exit 90
export REPRO_SLEEP_CYCLES=200000000   # ~110ms on H200
export VLLM_LOGGING_LEVEL=WARNING
python3 /patch/repro_gpu.py

Output on `0.26.1rc1.dev942` (H200):

probe	baseline	this PR
device idle when engine settled	True	True
device idle when pause_generation() returns	False	True
pause_generation() wall time on an idle engine	3.40s	3.40s (burst kept; #52957 cuts it to 0.17s)
device idle when drain returns	False	True
second pause+drain before sleep	3.40s, busy	3.40s, idle
drain mid-pause	"Engines are idle" while pause_done=False, device busy	pause already complete, device idle

Test Plan

Three unit tests in tests/v1/engine/test_engine_core.py (no GPU): the pause-ordering contracts — fast path and deferred path run synchronize → cache reset → future resolution, and overlapping pause calls each complete behind their own barrier.

Two adversarial DP tests for edges the pause path crosses that had no coverage: mode="wait" draining through the DP consensus path, and pausing a sleeping engine — whose ranks skip dummy batches yet must still reach consensus, with the barrier running against unmapped memory.

Regression: tests/v1/distributed/test_async_llm_dp.py -k "dp_pause or dp_sleep or prefill_schedule" — includes the anti-#38009 staggered-race and barrier-deadlock guards and the #51481 late-request pair.

python -m pytest tests/v1/engine/test_engine_core.py -k "pause_" -v
python -m pytest tests/v1/distributed/test_async_llm_dp.py -k "dp_pause or dp_sleep or prefill_schedule" -v

Test Result

vllm/vllm-openai:nightly (0.26.1rc1.dev942), A/B on H200s, before = unpatched image:

suite	before	after
test_engine_core.py unit tests (new)	3 failed	3 passed
test_async_llm_dp.py (14 cases incl. the 2 new adversarial)	14 passed	14 passed
GPU probe repro	table above	table above
offline LLM sleep/wake smoke, MPClient and in-proc	—	both pass

Essential Elements

• Not a duplicate: no open PR gives pause a device barrier (Add /pause/step fast pause + barrier for online weight update #31741 adds a different, step-bounded pause endpoint and does not synchronize the device). Filed under [RFC]: DP pause/sleep correctness: coordinator latch, request admission, device-idle completion, and sleep/pause layering #51476; [Core] Sync DP state on the first step of a wave #52957 is the independent cadence companion.
• Control-plane only; no effect on model output, so no eval run is included.
• AI assistance was used (the design was adversarially reviewed with a second model, which also surfaced the non-torch.accelerator backend hole fixed in the base implementation). The submitter has reviewed each changed line and the runs above.

[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.

[aoshen02]

/ci run

[github-actions]

✅ Triggered Buildkite CI #84723 for commit 6526c9cb4c64.

[aoshen02]

/ci run

[github-actions]

✅ Triggered Buildkite CI #84728 for commit d2ae48cd181c.

[aoshen02]

/ci run

[github-actions]

✅ Triggered Buildkite CI #84838 for commit c4b35e3af75e.

[aoshen02]

/ci run

[github-actions]

✅ Triggered Buildkite CI #84995 for commit 56e8f6e78673.

[aoshen02]

/ci run

[github-actions]

✅ Triggered Buildkite CI #85010 for commit 0c1bf2c2d7a2.

[njhill]

Thanks @aoshen02