vllm.utils.gpu_sync_debug ¶

def _active_check_mode() -> str | None:
    """The mode the calling thread is being checked in, if any."""
    if not _sync_check_enabled or _allow_depth.get() or torch.compiler.is_compiling():
        return None
    return _checking.get()

@contextmanager
def _checked_region(mode: str):
    """Police syncs on this thread for the duration of the block.

    The debug mode is armed per call so that syncs outside a checked region
    emit nothing, and refcounted because `execute_model` and `sample_tokens`
    nest. "warn" rather than "error" because torch's error mode raises on
    whichever thread synced, with no way to exempt one.
    """
    global _arm_count, _saved_sync_debug_mode
    _install_warning_hook()
    with _arm_lock:
        if _arm_count == 0:
            _saved_sync_debug_mode = torch.cuda.get_sync_debug_mode()
            torch.cuda.set_sync_debug_mode("warn")
        _arm_count += 1
    token = _checking.set(mode)
    try:
        yield
    finally:
        _checking.reset(token)
        with _arm_lock:
            _arm_count -= 1
            if _arm_count == 0:
                torch.cuda.set_sync_debug_mode(_saved_sync_debug_mode)

def _cpu_copy_stall_reason(
    cpu_tensor: torch.Tensor, device: torch.device | None
) -> str | None:
    """Why a `non_blocking` copy between `cpu_tensor` and a CUDA tensor would
    silently block the host, or None if it is genuinely asynchronous.

    `device` describes the other side of the copy; a non-CUDA `device` means
    the copy does not involve the GPU at all.
    """
    if device is None or device.type != "cuda" or cpu_tensor.device.type != "cpu":
        return None
    if cpu_tensor.numel() == 0:
        return None  # Nothing to copy; torch issues no CUDA call at all.
    if not _is_dense(cpu_tensor):
        return "the CPU tensor is not densely laid out"
    if not cpu_tensor.is_pinned():
        return "the CPU tensor is not pinned"
    return None

def _install_compile_time_sync_suppressors() -> None:
    """Allow the syncs torch's compile passes perform.

    Warmup-time compiles run before the gate flips, but lazy ones fire inside
    `execute_model`.
    """
    global _compile_time_suppressors_installed
    if _compile_time_suppressors_installed:
        return
    _compile_time_suppressors_installed = True

    try:
        from torch._inductor.fx_passes import joint_graph as _jg

        orig = _jg.joint_graph_passes
        wrapped = _suppressing(orig)
        # `compile_fx` imports this by value, so patching the defining module
        # alone misses that rebind; patch every compile-time module still
        # holding the original.
        _jg.joint_graph_passes = wrapped
        for name, mod in list(sys.modules.items()):
            if (
                mod is not None
                and name.startswith(
                    ("torch._inductor", "torch._functorch", "torch._dynamo")
                )
                and getattr(mod, "joint_graph_passes", None) is orig
            ):
                setattr(mod, "joint_graph_passes", wrapped)  # noqa: B010
    except Exception:  # pragma: no cover
        pass

    try:
        # Inductor builds its cudagraph tree lazily, so `deferred_cudagraphify`
        # and the `capture_begin` sync inside it can fire during
        # `execute_model`. It resolves `cudagraphify` as a module global at
        # call time, so patching the attribute is enough.
        from torch._inductor import cudagraph_trees as _ct

        _ct.cudagraphify = _suppressing(_ct.cudagraphify)
    except Exception:  # pragma: no cover
        pass

def _install_copy_checkers() -> None:
    """Wrap the Python-level CPU<->CUDA transfer entry points with the checks
    above.

    The patch is process-global, but it only reports on threads inside a
    checked region, matching the scoping of the torch-level check. Copies
    issued from C++ (e.g. inside compiled graphs) bypass it.
    """
    global _copy_checkers_installed
    if _copy_checkers_installed:
        return
    _copy_checkers_installed = True
    torch.Tensor.to = _checked_to  # type: ignore[method-assign]
    torch.Tensor.cuda = _checked_cuda  # type: ignore[method-assign]
    torch.Tensor.copy_ = _checked_copy_  # type: ignore[method-assign]

def _install_warning_hook() -> None:
    """(Re)install the hook and a filter that lets torch's warning reach it.

    Done per checked call because pytest runs each test inside
    `warnings.catch_warnings()`, which restores both `showwarning` and
    `filters`. The hook is left in place afterwards: outside a checked call
    the debug mode is disarmed, so torch emits nothing for it to see.
    """
    global _prev_showwarning, _sync_filter_head
    if warnings.showwarning is not _sync_warning_hook:
        _prev_showwarning = warnings.showwarning
        warnings.showwarning = _sync_warning_hook
    # "always" so the warning survives filtering and isn't deduplicated by
    # `__warningregistry__`. `filterwarnings` prepends, so only re-assert it
    # once ours is no longer in front.
    if warnings.filters[:1] != [_sync_filter_head]:
        warnings.filterwarnings(
            "always", message=_TORCH_SYNC_WARNING, category=UserWarning
        )
        _sync_filter_head = warnings.filters[0]

def _is_dense(t: torch.Tensor) -> bool:
    """Whether `t`'s elements fill a contiguous storage range, possibly
    permuted (e.g. a transposed matrix). Such layouts can be copied with
    pitched cudaMemcpy2D/3DAsync directly from/to pinned memory, whereas
    gapped layouts (e.g. strided slices) stage through a pageable temp."""
    if t.is_contiguous():
        # The common case, answered by a memoized TensorImpl flag.
        return True
    expected_stride = 1
    for stride, size in sorted(zip(t.stride(), t.shape)):
        if size <= 1:
            continue
        if stride != expected_stride:
            return False
        expected_stride *= size
    return True

def _suppressing(fn):
    """Allow the syncs `fn` performs on its calling thread."""

    @functools.wraps(fn)
    def wrapper(*args, **kwargs):
        # Not `gpu_sync_allowed()`, which no-ops while
        # `torch.compiler.is_compiling()` -- exactly when these run.
        with _allow_syncs():
            return fn(*args, **kwargs)

    return wrapper

def _sync_warning_hook(message, category, filename, lineno, file=None, line=None):
    """Turn torch's sync warning into an error, but only on a thread that is
    being checked and is outside any allow region."""
    if _TORCH_SYNC_WARNING in str(message):
        mode = _checking.get()
        if mode is None or _allow_depth.get():
            return None
        if mode == "error":
            raise RuntimeError(SYNC_ERROR_MESSAGE)
    return _prev_showwarning(message, category, filename, lineno, file, line)

def enable_gpu_sync_check() -> None:
    """Flip the sync-check gate on, once per worker, after warmup."""
    if _SYNC_CHECK_MODE is None:
        return
    global _sync_check_enabled
    _sync_check_enabled = True
    _install_compile_time_sync_suppressors()
    _install_copy_checkers()

def gpu_sync_allowed(first_only: bool = False):
    """Allow GPU<->CPU syncs inside the `with` block, on this thread only.

    With `first_only`, only the first entry from a given call site
    (filename, lineno) is allowed, so later syncs there are still reported.
    """
    if _SYNC_CHECK_MODE is None or torch.compiler.is_compiling():
        return _NOOP_CM
    if first_only:
        frame = sys._getframe(1)
        key = (frame.f_code.co_filename, frame.f_lineno)
        if key in _GPU_SYNC_ALLOWED_FIRST_SEEN:
            return _NOOP_CM
        _GPU_SYNC_ALLOWED_FIRST_SEEN.add(key)
    return _allow_syncs()

def with_gpu_sync_check(fn):
    """Report GPU<->CPU syncs performed by `fn` on its calling thread.

    Active only once `enable_gpu_sync_check()` has flipped the gate. Other
    threads are never policed, so deliberate syncs there (e.g. the EPLB
    transfer worker) are unaffected.
    """
    if (mode := _SYNC_CHECK_MODE) is None:
        return fn

    @functools.wraps(fn)
    def wrapper(*args, **kwargs):
        if not _sync_check_enabled:
            return fn(*args, **kwargs)
        with _checked_region(mode):
            return fn(*args, **kwargs)

    return wrapper