[quanliu1991]

Purpose

Fixes #57759

An eager break that queues asynchronous work on a side CUDA stream and returns without joining (e.g. the pinned-host PLE UVA prefetch launched by start_prefetch in vllm/models/qwen4_exp/nvidia/ngram_embedding.py) leaves that work in flight when a later capture segment calls capture_end(). Depending on timing, one of two things happens:

1. Intermittent crash: the capture is invalidated and all workers raise torch.AcceleratorError: CUDA error: operation not permitted at capture_end() (the cudaErrorStreamCaptureUnjoined/Invalidated family). This is load-dependent — it succeeds on an idle host and fails under host load, because host load slows the PCIe UVA reads and widens the race window. Observed on 4x L20 TP4 with --engram-config '{"cpu_offload": true}' and --compilation-config '{"cudagraph_mode": "FULL_AND_PIECEWISE"}'.
2. Silent corruption (found while writing the regression test): the side-stream kernels leak past the capture and interleave with post-capture work.

Fix: join all outstanding device work in the legal window — after the previous segment has ended and before the next capture_begin() — by calling torch.cuda.synchronize() at the top of BreakableCUDAGraphCapture._begin_segment().

Two properties worth highlighting:

• Placement matters: the sync must live in _begin_segment, not _end_segment. _end_segment runs while the segment's capture is still active (the sync would then itself be illegal and fails deterministically — verified on a rig while developing this fix). Only the between-segments window is valid.
• Zero steady-state cost: breakable capture happens once at startup (profile_cudagraph_memory / capture phase); serving-path replay is untouched.

Test Plan

New regression test test_eager_break_side_stream_work_is_joined_before_next_capture in tests/v1/cudagraph/test_breakable_cudagraph.py mimics the PLE prefetch pattern: the eager-break fn forks a side stream, queues 256–2048 async kernels, and returns without joining back.

# regression test sketch (full test in the diff)
side_stream = torch.cuda.Stream()

def prefetch_like():
    side_stream.wait_stream(torch.cuda.current_stream())
    with torch.cuda.stream(side_stream):
        for _ in range(256):
            y.add_(1.0)   # no join back

cap = BreakableCUDAGraphCapture()
with cap:
    x.add_(1.0)
    cap.add_eager(prefetch_like)
    x.add_(1.0)  # unfixed: capture_end invalidated / kernels leak

Test Result

Verified on L20 (driver 535.183.01, torch 2.13.0+cu130):

• Unfixed code: the regression test fails deterministically on the first run — the side-stream adds leak past the capture and corrupt post-capture zeroing (518 instead of 256, i.e. leaked kernels re-adding after the zeroing), no synthetic host load needed.
• Fixed code: the regression test passes; replay produces exact results across repeated runs.
• End-to-end (4x L20 TP4, Qwen3.8-Flash-Next-FP8 with cpu_offload + FULL_AND_PIECEWISE): patched build survives restarts under synthetic host load (CPU spinners + concurrent direct reads of the checkpoint) that reproduces the original production failure (cudaErrorNotPermitted at FULL capture 1/2 on all 4 workers); unpatched nightly failed under the same conditions.
• Performance unchanged: TPOT C1 12.3 ms / C16 26.9 ms, identical to unpatched (the sync is startup-only).

---

Note on authorship: this fix was developed with AI assistance (GLM5.3, via Cursor — see the Co-authored-by trailer). I reviewed every changed line of code and personally ran all the tests described above on our 4x L20 rig, including the production-reproduction restarts under synthetic load.

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

[quanliu1991]

@peakcrosser7 @ZJY0516 @Isotr0py This fixes a capture race we hit while running #54371 in production. Since you reviewed the PLE UVA prefetch path, could you please take a look when you have time? Thanks!

[peakcrosser7]

Thanks for the PR!

The change looks reasonable to me, though I have limited experience with CUDA graph capture and am not sure whether unfinished side-stream work alone explains the capture_end() error.

One concern about the regression test: I think the main stream should wait for the side stream before y.zero_(). Otherwise, the test has a race independent of graph capture, and a capture-time sync doesn’t replace proper stream dependencies during replay.

[ZJY0516]

This fixes a capture race we hit while running #54371 in production

Could you share your config? We didn't see related issue recently

And I don't want it to be specialized for prefetch

[quanliu1991]

@ZJY0516

1. BASE_IMAGE: docker.das-security.cn/hb/vllm-openai:v0.29.0

2. Install UVA PLE-offload([Qwen4Exp] Support UVA PLE-offload and Engram tensor parallelism #54371) nightlywheel

curl -L -o /tmp/vllm-0.2.1.dev58+g9854b580-cp38-abi3-manylinux_2_28_x86_64.whl \
  "https://wheels.vllm.ai/9854b580dfe518171eed3a6ad76cbd402ca2bdfa/vllm-0.2.1.dev58%2Bg9854b580d-cp38-abi3-manylinux_2_28_x86_64.whl"
pip uninstall -y vllm
pip install --no-deps "/tmp/vllm-0.2.1.dev58+g9854b580-cp38-abi3-manylinux_2_28_x86_64.whl"

3. Start serve

#!/bin/bash
# UVA PLE-offload  nightly 9854b580, PR #54371
# use --engram-config cpu_offload  replace VLLM_PLE_CPU_OFFLOAD
set -euo pipefail
docker rm -f qwen38fn-uva 2>/dev/null || true
docker run -d --gpus all --name qwen38fn-uva \
  --shm-size 16G -p 34004:8000 \
  -e CUDA_VISIBLE_DEVICES=2,3,4,5 \
  -e VLLM_ENGINE_READY_TIMEOUT_S=1800 \
  docker.das-security.cn/hb/vllm-openai:qwen38-flash-next-uva-v1 \
  Qwen/Qwen3.8-Flash-Next-FP8 \
  --language-model-only \
  --max-model-len 256k \
  --max-num-seqs 32 \
  --enable-prefix-caching \
  --no-enable-flashinfer-autotune \
  --tensor-parallel-size 4 \
  --moe-backend triton \
  --gpu-memory-utilization 0.95 \
  --enable-auto-tool-choice \
  --tool-call-parser qwen3_xml \
  --reasoning-parser qwen3 \
  --served-model-name qwen3.8-flash-next-fp8 \
  --engram-config '{"cpu_offload": true}'

[ZJY0516]

I don't see any issue using 1b9fa3e

[quanliu1991]

I will further confirm whether there is a stable reproduction method for #57759 .

(Worker_TP1 pid=711) ERROR [multiproc_executor.py:1053] Traceback (most recent call last):
  File ".../vllm/v1/executor/multiproc_executor.py", line 1045, in _execute_worker_rpc
  File ".../vllm/v1/worker/gpu_worker.py", line 586, in determine_available_memory
    cudagraph_memory_estimate = self.model_runner.profile_cudagraph_memory()
  File ".../vllm/v1/worker/gpu/cudagraph_utils.py", line 922, in profile_cudagraph_memory
    measured = int(runner.capture_model(profile_only=True))
  File ".../vllm/v1/worker/gpu/cudagraph_utils.py", line 748, in capture
  File ".../torch/cuda/graphs.py", line 477, in __exit__
    self.cuda_graph.capture_end()
  File ".../torch/cuda/graphs.py", line 169, in capture_end
torch.AcceleratorError: CUDA error: operation not permitted
Search for `cudaErrorNotPermitted` in https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__TYPES.html for more information.

[quanliu1991]

Closing without merge.

Re-reading the code and re-running the experiments, this fix does not target the failure it claims to fix (see my closing comment on #57759):

• The reported crash is raised in the FULL torch.cuda.graph(...) branch of cudagraph_utils.py, which never calls BreakableCUDAGraphCapture._begin_segment(). A sync there cannot affect that path, and torch.cuda.graph.__enter__ already synchronizes before capture_begin.
• In the breakable path, start_prefetch and _finalize_prefetch are both eager breaks, so fork and join both land in eager gaps. I could not make an unjoined-in-gap fork fail under any load in a standalone repro, so the extra synchronize() in _begin_segment is redundant rather than protective.
• @peakcrosser7's point about the regression test is also correct: the test as written has a race independent of graph capture and would not prove the fix.

I agree with @ZJY0516 that the breakable capture code should not be specialized for the prefetch pattern. If I get a reliable reproduction I will open a new PR with a fix that actually addresses the failing path.

Thanks for the reviews.