vLLM/vllm/v1/worker/workspace.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

import inspect
import os
from collections.abc import Callable, Hashable, Iterator
from contextlib import AbstractContextManager, contextmanager, nullcontext
from contextvars import ContextVar
from itertools import accumulate
from math import prod
from typing import TypeVar, cast

import torch

import vllm.envs as envs
from vllm.logger import init_logger
from vllm.utils.math_utils import round_up
from vllm.v1.worker.ubatching import dbo_current_ubatch_id

logger = init_logger(__name__)
T = TypeVar("T")


def _compute_bytes(shape: tuple[int, ...], dtype: torch.dtype) -> int:
    return prod(shape) * dtype.itemsize


# Constants
_MB = 1024**2
_GiB = 1024**3

# Global workspace manager instance
_manager: "WorkspaceManager | None" = None
_workspace_lane: ContextVar[int] = ContextVar("vllm_workspace_lane", default=0)


@contextmanager
def use_workspace_lane(lane: int) -> Iterator[None]:
    """Select an independent workspace owner for this execution context."""
    if lane < 0:
        raise ValueError(f"Workspace lane must be non-negative, got {lane}.")
    token = _workspace_lane.set(lane)
    try:
        yield
    finally:
        _workspace_lane.reset(token)


class WorkspaceManager:
    """Manager for workspace allocation.

    Manages shared scratch and keyed persistent resources per ``(ubatch, lane)``.
    Can be locked to prevent growth or new resources during execution.
    """

    def __init__(
        self,
        device: torch.device,
        num_ubatches: int | None = None,
        num_lanes: int = 1,
        alloc_context: Callable[[], AbstractContextManager] = nullcontext,
    ):
        self._device = device
        self._alloc_context = alloc_context
        # Cache num ubatches at init based on configuration (default to 1)
        self._num_ubatches = num_ubatches if num_ubatches is not None else 1
        if num_lanes < 1:
            raise ValueError(f"num_lanes must be at least one, got {num_lanes}.")
        self._num_lanes = num_lanes
        self._current_workspaces: list[torch.Tensor | None] = [None] * (
            self._num_ubatches * self._num_lanes
        )
        self._persistent_resources: list[dict[Hashable, object]] = [
            {} for _ in self._current_workspaces
        ]
        self._locked: bool = False

    @staticmethod
    def _workspace_size_bytes(workspace: torch.Tensor | None) -> int:
        """Get size of workspace in bytes."""
        if workspace is None:
            return 0
        return workspace.numel() * workspace.element_size()

    def lock(self) -> None:
        """Lock the workspace to prevent growth or new persistent resources.

        Larger scratch requests and new persistent resources raise an assertion
        error. Existing resources retain their storage throughout execution.
        """
        self._locked = True
        if envs.VLLM_DEBUG_WORKSPACE:
            logger.info(
                "[WORKSPACE DEBUG] Workspace locked. Current sizes: %s",
                [
                    self._workspace_size_bytes(ws) / _MB
                    for ws in self._current_workspaces
                    if ws is not None
                ],
            )

    def unlock(self) -> None:
        """Unlock the workspace to allow growth.

        This is used during elastic EP scaling when the workspace size
        needs to grow due to changes in the number of experts.
        """
        self._locked = False
        if envs.VLLM_DEBUG_WORKSPACE:
            logger.info(
                "[WORKSPACE DEBUG] Workspace unlocked. Current sizes: %s",
                [
                    self._workspace_size_bytes(ws) / _MB
                    for ws in self._current_workspaces
                    if ws is not None
                ],
            )

    def is_locked(self) -> bool:
        """Check if workspace is locked."""
        return self._locked

    def _get_workspace_id(self) -> int:
        lane = _workspace_lane.get()
        if lane >= self._num_lanes:
            raise RuntimeError(
                f"Workspace lane {lane} is not configured; manager has "
                f"{self._num_lanes} lane(s)."
            )
        return dbo_current_ubatch_id() * self._num_lanes + lane

    def get_persistent_resource(self, key: Hashable, factory: Callable[[], T]) -> T:
        """Return the resource cached for ``key`` in the current ubatch and lane.

        The first request calls ``factory`` and stores its result for the lifetime
        of this manager. Later requests return the same object without calling
        the factory again. Each ubatch/lane pair has a separate cache. Once the
        manager is locked, requesting a key missing from that cache raises an
        AssertionError.

        Keys must identify the resource type and configuration. Callers must not
        resize or replace tensor storage held by a cached resource. Uses of the
        same resource must not overlap; include the CUDA stream in the key when
        streams can execute concurrently.

        With torch.compile, call this inside a runtime custom op so resource
        selection happens during execution.
        """
        resources = self._persistent_resources[self._get_workspace_id()]
        if key not in resources:
            if self._locked:
                raise AssertionError(
                    f"Workspace is locked but {key!r} was not allocated during "
                    "warmup for the current ubatch and lane."
                )
            resources[key] = factory()
        return cast(T, resources[key])

    def get_persistent(
        self,
        key: Hashable,
        shape: tuple[int, ...],
        dtype: torch.dtype,
        *,
        zero_init: bool = False,
    ) -> torch.Tensor:
        """Get a fixed-shape tensor, optionally zeroed on first allocation only."""
        factory = torch.zeros if zero_init else torch.empty
        workspace = self.get_persistent_resource(
            (torch.Tensor, key),
            lambda: factory(shape, dtype=dtype, device=self._device),
        )
        if workspace.shape != shape or workspace.dtype != dtype:
            raise ValueError(
                f"Persistent workspace {key!r} has shape={workspace.shape}, "
                f"dtype={workspace.dtype}; requested shape={shape}, dtype={dtype}."
            )
        return workspace

    def get_simultaneous(
        self, *shapes_and_dtypes: tuple[tuple[int, ...], torch.dtype]
    ) -> list[torch.Tensor]:
        """Get multiple workspace tensors simultaneously from a single allocation.

        Args:
            *shapes_and_dtypes: One or more (shape, dtype) tuples.

        Returns:
            List of tensor views into the workspace buffer, one per shape/dtype pair.

        """
        actual_bytes = [_compute_bytes(s, d) for s, d in shapes_and_dtypes]
        aligned_bytes = [round_up(actual, 256) for actual in actual_bytes]
        total_bytes = sum(aligned_bytes)

        # Calculate cumulative offsets using itertools.accumulate
        offsets = list(accumulate([0] + aligned_bytes[:-1]))

        current_workspace = self._ensure_workspace_size(total_bytes)

        return [
            current_workspace[offsets[i] : offsets[i] + actual_bytes[i]]
            .view(shapes_and_dtypes[i][1])
            .reshape(shapes_and_dtypes[i][0])
            for i in range(len(shapes_and_dtypes))
        ]

    def _ensure_workspace_size(self, required_bytes: int) -> torch.Tensor:
        """Ensure workspace is allocated and large enough, return current workspace.

        Args:
            required_bytes: The number of bytes required.

        Returns:
            The current workspace tensor.

        """
        workspace_id = self._get_workspace_id()
        current_workspace = self._current_workspaces[workspace_id]
        current_size = self._workspace_size_bytes(current_workspace)

        if current_size < required_bytes:

            def get_caller_info() -> str:
                """Find first frame outside WorkspaceManager."""
                curr_frame = inspect.currentframe()
                if curr_frame is None:
                    return "unknown"
                # Walk up the stack skipping WorkspaceManager frames
                curr_frame = curr_frame.f_back
                while curr_frame is not None:
                    if curr_frame.f_code.co_qualname.startswith("WorkspaceManager."):
                        curr_frame = curr_frame.f_back
                        continue
                    filename = os.path.basename(curr_frame.f_code.co_filename)
                    return (
                        f"{filename}:{curr_frame.f_lineno}:{curr_frame.f_code.co_name}"
                    )
                return "unknown"

            if self._locked:
                raise AssertionError(
                    f"Workspace is locked but allocation from '{get_caller_info()}' "
                    f"requires {required_bytes / _MB:.2f} MB, current size is "
                    f"{current_size / _MB:.2f} MB. "
                    "Workspace growth is not allowed after locking."
                )

            # Only resize the requesting ubatch/lane workspace. Other slots
            # resize lazily on their next get_simultaneous call.
            # Resizing all ubatches here would orphan the other ubatch's
            # old tensor when it still holds views into it (DBO leak).
            self._current_workspaces[workspace_id] = None
            del current_workspace
            # Release the freed segment back to CUDA so the caching
            # allocator can reuse the GPU memory for the larger
            # allocation below. Without this, each resize may leave a
            # dead segment in reserved memory which can cause higher peak
            # memory usage.
            torch.accelerator.empty_cache()
            # In sleep mode, entering the pool also trims outgrown buffers.
            with self._alloc_context():
                self._current_workspaces[workspace_id] = torch.empty(
                    (required_bytes,), dtype=torch.uint8, device=self._device
                )
            current_workspace = self._current_workspaces[workspace_id]

            if envs.VLLM_DEBUG_WORKSPACE:
                logger.info(
                    "[WORKSPACE DEBUG] Resized workspace from '%s': %.2f MB -> "
                    "%.2f MB (ubatch %d, lane %d)",
                    get_caller_info(),
                    current_size / _MB,
                    required_bytes / _MB,
                    workspace_id // self._num_lanes,
                    workspace_id % self._num_lanes,
                )

        return current_workspace


def is_workspace_manager_initialized() -> bool:
    """Check if workspace manager has been initialized.

    Returns:
        True if workspace manager is initialized, False otherwise.

    """
    return _manager is not None


def current_workspace_manager() -> "WorkspaceManager":
    """Get the current workspace manager instance.

    Raises:
        AssertionError: If workspace manager has not been initialized.

    """
    assert _manager is not None, (
        "WorkspaceManager not initialized. Call init_workspace_manager() "
        "with a device before using workspace functions."
    )
    return _manager


def init_workspace_manager(
    device: torch.device,
    num_ubatches: int | None = None,
    num_lanes: int = 1,
    alloc_context: Callable[[], AbstractContextManager] = nullcontext,
) -> None:
    """Initialize the workspace manager with a device.

    Must be called before using any workspace functions. Typically called
    from GPUModelRunner.__init__.

    Args:
        device: The device to allocate workspace on.
        num_ubatches: Number of workspace ubatch slots. Defaults to 1.
        num_lanes: Number of independent execution lanes per ubatch. Defaults to 1.
        alloc_context: Factory for the context scratch is allocated in.

    """
    global _manager
    if _manager is not None:
        logger.warning(
            "WorkspaceManager already initialized on device %s, "
            "reinitializing on device %s",
            _manager._device,
            device,
        )
    _manager = WorkspaceManager(device, num_ubatches, num_lanes, alloc_context)


def lock_workspace() -> None:
    """Lock the workspace to prevent growth or new persistent resources.

    Larger scratch requests and new persistent resources raise an AssertionError.
    All required resources must be initialized in each execution slot during
    warmup, so their storage remains fixed during execution.

    Example:
        # During initialization
        init_workspace_manager(device)
        reserve_workspace(shape1, dtype1)
        reserve_workspace(shape2, dtype2)

        # Lock after warmup/profiling
        lock_workspace()

        # Now all get_workspace calls must fit in pre-allocated size

    """
    current_workspace_manager().lock()


def unlock_workspace() -> None:
    """Unlock the workspace to allow growth.

    This is used during elastic EP scaling when the workspace size
    needs to grow due to changes in the number of experts.
    After scaling operations complete, lock_workspace() should be
    called again to prevent unexpected allocations.
    """
    current_workspace_manager().unlock()


def reset_workspace_manager() -> None:
    """Reset the workspace manager to uninitialized state.

    This is primarily intended for testing purposes to allow tests
    to reinitialize the workspace manager cleanly.
    """
    global _manager
    _manager = None