vllm.v1.conf_compute_utils

Keep non-blocking host<->device copies non-blocking under NVIDIA CC.

Under bounce-buffer Confidential Computing, cudaMemcpyAsync is host-synchronous: it returns only after the copy and all prior work on its stream complete. install() patches torch's non_blocking copy entry points so that:

• H2D runs on an idle per-device staging stream, so the host waits only for the transfer. The current stream then waits on the staging stream. A copy into an existing tensor lands in a staging tensor and is moved by a D2D on the current stream, so it is ordered exactly like the original copy. This applies to every caller.
• D2H issued on a stream created with new_external_d2h_stream() is issued by a worker thread once the source is ready. Event / stream / device synchronize additionally wait for the D2H copies submitted before the event was recorded (or before the call). Using such a stream asserts that its D2H sources aren't overwritten until the copy is synced (later compute-stream work is never ordered after a side stream anyway), and that their completion is only awaited through the patched torch.cuda / torch.Event methods. D2H on any other stream keeps the original, host-blocking copy, since e.g. a torch.accelerator.current_stream().synchronize() can't be patched. A registered stream is only ever ordered after a producer stream by wait_stream() and then carries D2H copies, so the events that mark its position (the copies' readiness, and events recorded on it such as the caller's copy event) are recorded on that producer stream instead. Under CC an event recorded on an otherwise idle stream behind a cross-stream wait was observed to fire only once the compute stream's next graph had finished, which delayed every sampled token by a full step; an event recorded on the producer stream fires on time.

Copies made while torch.compile is tracing or a CUDA graph is being captured are left alone, as are H2D copies into an existing tensor that would need a staging buffer larger than _MAX_STAGED_BYTES (e.g. KV cache swaps).

Functions:

• install –

  Install the patches. Call once per worker process, only under CC.

• new_external_d2h_stream –

  Create a stream whose non_blocking D2H copies go through the worker

• register_inplace_h2d_target –

  Let non_blocking H2D copies into tensor (or a view of it) land in

_capture_originals()

Save the torch.cuda methods the patches wrap, once.

Source code in vllm/v1/conf_compute_utils.py

def _capture_originals() -> None:
    """Save the torch.cuda methods the patches wrap, once."""
    global _orig_cuda_event_record, _orig_cuda_event_synchronize
    global _orig_stream_synchronize, _orig_stream_wait_stream, _orig_stream_query
    global _orig_cuda_synchronize, _orig_accel_synchronize
    if _orig_cuda_event_record is not None:
        return
    _orig_cuda_event_record = torch.cuda.Event.record
    _orig_cuda_event_synchronize = torch.cuda.Event.synchronize
    _orig_stream_synchronize = torch.cuda.Stream.synchronize
    _orig_stream_wait_stream = torch.cuda.Stream.wait_stream
    _orig_stream_query = torch.cuda.Stream.query
    _orig_cuda_synchronize = torch.cuda.synchronize
    _orig_accel_synchronize = torch.accelerator.synchronize

_d2h(src, out=None)

Returns None if the copy should fall back to the original method.

Source code in vllm/v1/conf_compute_utils.py

def _d2h(src: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor | None:
    """Returns None if the copy should fall back to the original method."""
    cur = torch.cuda.current_stream(src.device)
    if cur.cuda_stream not in _d2h_streams:
        return None
    if out is None:
        out = torch.empty(src.shape, dtype=src.dtype, pin_memory=True)
    ready = torch.cuda.Event(blocking=True)
    _orig_cuda_event_record(ready, _producer_of.get(cur.cuda_stream, cur))

    def run() -> None:
        # Sleep until the producer is done, so the host-synchronous copy only
        # pays for its own transfer (and doesn't hold driver locks meanwhile).
        _orig_cuda_event_synchronize(ready)
        side = _d2h_stream(src.device.index)
        # Entering the stream also switches this thread to its device.
        # Inference mode is thread-local, and out may be an inference tensor;
        # an in-place copy into a normal tensor is allowed in inference mode.
        with torch.inference_mode(), torch.cuda.stream(side):
            _orig_copy_(out, src, non_blocking=True)
            # The closure keeps src alive until here; this keeps the caching
            # allocator from reusing it before the copy has run.
            src.record_stream(side)
        _orig_stream_synchronize(side)

    assert _d2h_executor is not None
    fut = _d2h_executor.submit(run)
    with _pending_lock:
        _pending.append(fut)
    return out

_new_external_stream(device)

A stream created outside torch's per-device pool, so torch.cuda.Stream() can never hand out its handle again (the pool is 32 streams, round-robin).

Source code in vllm/v1/conf_compute_utils.py

def _new_external_stream(device: torch.device) -> torch.cuda.Stream:
    """A stream created outside torch's per-device pool, so `torch.cuda.Stream()`
    can never hand out its handle again (the pool is 32 streams, round-robin)."""
    from vllm.distributed.device_communicators.cuda_wrapper import CudaRTLibrary

    with torch.cuda.device(device):
        handle = CudaRTLibrary().cudaStreamCreateWithFlags(_CUDA_STREAM_NON_BLOCKING)
    return torch.cuda.ExternalStream(handle, device=device)

_patch_event_methods(cls)

Wrap cls's record/synchronize/query so a recorded event also covers the D2H copies submitted before it was recorded.

Source code in vllm/v1/conf_compute_utils.py

def _patch_event_methods(cls: Any) -> dict[str, Any]:
    """Wrap `cls`'s record/synchronize/query so a recorded event also covers
    the D2H copies submitted before it was recorded."""
    orig_record, orig_synchronize, orig_query = (
        cls.record,
        cls.synchronize,
        cls.query,
    )

    def record(self, stream=None):
        self._cc_pending = _pending_snapshot()
        stream = _record_target(stream)
        if stream is None:
            return orig_record(self)
        return orig_record(self, stream)

    def synchronize(self):
        for f in getattr(self, "_cc_pending", ()):
            f.result()
        return orig_synchronize(self)

    def query(self):
        pending = getattr(self, "_cc_pending", ())
        return all(f.done() for f in pending) and orig_query(self)

    return {"record": record, "synchronize": synchronize, "query": query}

_record_target(stream)

Where an event recorded on stream (None: the current stream) goes: the producer stream of a registered D2H stream, otherwise stream.

Source code in vllm/v1/conf_compute_utils.py

def _record_target(stream):
    """Where an event recorded on `stream` (None: the current stream) goes:
    the producer stream of a registered D2H stream, otherwise `stream`."""
    s = stream if stream is not None else torch.cuda.current_stream()
    target = _producer_of.get(s.native_handle, stream)
    if target is not stream and torch.cuda.is_current_stream_capturing():
        raise RuntimeError(
            "Confidential Computing: an event is being recorded during CUDA graph "
            f"capture on a D2H copy stream (handle {s.native_handle:#x}). Copy "
            "streams must not be shared with capture streams."
        )
    return target

install()

Install the patches. Call once per worker process, only under CC.

torch.Event is an immutable C type, so it is replaced with a subclass: events created through torch.Event(...) afterwards are covered, but events created in C (e.g. torch.Stream.record_event()) are not, and isinstance(e, torch.Event) becomes False for those and for torch.cuda.Event instances (which subclass the original). Only D2H copies on registered streams depend on these patches.

Source code in vllm/v1/conf_compute_utils.py

def install() -> None:
    """Install the patches. Call once per worker process, only under CC.

    torch.Event is an immutable C type, so it is replaced with a subclass:
    events created through `torch.Event(...)` afterwards are covered, but
    events created in C (e.g. `torch.Stream.record_event()`) are not, and
    `isinstance(e, torch.Event)` becomes False for those and for
    `torch.cuda.Event` instances (which subclass the original). Only D2H
    copies on registered streams depend on these patches.
    """
    global _d2h_executor
    if _d2h_executor is not None:
        return
    _capture_originals()
    _d2h_executor = ThreadPoolExecutor(1, thread_name_prefix="vllm-cc-d2h")
    logger.info(
        "NVIDIA Confidential Computing: issuing non-blocking H2D copies on a "
        "staging stream and D2H copies from a worker thread."
    )
    torch.Tensor.copy_ = _copy_
    torch.Tensor.to = _to
    torch.Tensor.cuda = _cuda
    for name, fn in _patch_event_methods(torch.cuda.Event).items():
        setattr(torch.cuda.Event, name, fn)
    torch.Event = type(
        "Event", (_orig_event_cls,), _patch_event_methods(_orig_event_cls)
    )
    torch.cuda.Stream.synchronize = _stream_synchronize
    torch.cuda.Stream.wait_stream = _stream_wait_stream
    torch.cuda.Stream.query = _stream_query
    torch.cuda.synchronize = _cuda_synchronize
    torch.accelerator.synchronize = _accel_synchronize

new_external_d2h_stream(device=None)

Create a stream whose non_blocking D2H copies go through the worker thread once install() has run (see the module docstring). Callers use this through async_utils.new_d2h_stream().

The stream is created outside torch's stream pool: the patches key on the raw handle, and a pool stream handed out later (e.g. by graph_capture() or aux_stream()) must never share it, or its events would be redirected to the producer stream. The stream lives for the life of the process.

Such a stream must only carry copies: readiness is tracked on the stream it last waited on via wait_stream(), so a kernel launched on the copy stream itself is not ordered before the copies of its output. Compute such values on the producer stream first.

Source code in vllm/v1/conf_compute_utils.py

def new_external_d2h_stream(
    device: torch.device | int | None = None,
) -> torch.cuda.Stream:
    """Create a stream whose non_blocking D2H copies go through the worker
    thread once `install()` has run (see the module docstring). Callers use
    this through `async_utils.new_d2h_stream()`.

    The stream is created outside torch's stream pool: the patches key on the
    raw handle, and a pool stream handed out later (e.g. by `graph_capture()`
    or `aux_stream()`) must never share it, or its events would be redirected
    to the producer stream. The stream lives for the life of the process.

    Such a stream must only carry copies: readiness is tracked on the stream
    it last waited on via `wait_stream()`, so a kernel launched on the copy
    stream itself is not ordered before the copies of its output. Compute such
    values on the producer stream first.
    """
    dev = torch.device("cuda") if device is None else torch.device(device)
    if dev.index is None:
        dev = torch.device("cuda", torch.cuda.current_device())
    stream = _new_external_stream(dev)
    assert stream.cuda_stream not in _d2h_streams, "duplicate D2H stream handle"
    _d2h_streams.add(stream.cuda_stream)
    return stream

register_inplace_h2d_target(tensor)

Let non_blocking H2D copies into tensor (or a view of it) land in place on the staging stream once install() has run: no staging tensor and no D2D. The caller guarantees that no GPU reader of tensor is still in flight whenever it is rewritten (e.g. a round-robin pool slot whose readers retired). The staging stream is ordered once after the current stream, so work that initialized tensor (e.g. the zero fill of a fresh allocation) completes before the first in-place upload. The registration lapses with tensor, so a later tensor reusing its memory is not matched.

Source code in vllm/v1/conf_compute_utils.py

def register_inplace_h2d_target(tensor: torch.Tensor) -> None:
    """Let non_blocking H2D copies into `tensor` (or a view of it) land in
    place on the staging stream once `install()` has run: no staging tensor
    and no D2D. The caller guarantees that no GPU reader of `tensor` is still
    in flight whenever it is rewritten (e.g. a round-robin pool slot whose
    readers retired). The staging stream is ordered once after the current
    stream, so work that initialized `tensor` (e.g. the zero fill of a fresh
    allocation) completes before the first in-place upload. The registration
    lapses with `tensor`, so a later tensor reusing its memory is not matched.
    """
    from vllm.platforms import current_platform

    if not tensor.is_cuda or not current_platform.is_confidential_compute():
        return
    ptr = tensor.untyped_storage().data_ptr()
    cur = torch.cuda.current_stream(tensor.device)
    _staging_stream(cur.device_index).wait_stream(cur)
    _inplace_h2d_storages.add(ptr)
    weakref.finalize(tensor, _inplace_h2d_storages.discard, ptr)