vllm.v1.worker.workspace

Classes:

• WorkspaceManager –

  Manager for workspace allocation.

Functions:

• current_workspace_manager –

  Get the current workspace manager instance.

• init_workspace_manager –

  Initialize the workspace manager with a device.

• is_workspace_manager_initialized –

  Check if workspace manager has been initialized.

• lock_workspace –

  Lock the workspace to prevent growth or new persistent resources.

• reset_workspace_manager –

  Reset the workspace manager to uninitialized state.

• unlock_workspace –

  Unlock the workspace to allow growth.

• use_workspace_lane –

  Select an independent workspace owner for this execution context.

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.

Methods:

• get_persistent –

  Get a fixed-shape tensor, optionally zeroed on first allocation only.

• get_persistent_resource –

  Return the resource cached for key in the current ubatch and lane.

• get_simultaneous –

  Get multiple workspace tensors simultaneously from a single allocation.

• is_locked –

  Check if workspace is locked.

• lock –

  Lock the workspace to prevent growth or new persistent resources.

• unlock –

  Unlock the workspace to allow growth.

Source code in vllm/v1/worker/workspace.py

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:
                    # TODO: This only catches instance methods (self), missing
                    # classmethods and staticmethods. Once Python 3.11+ is the
                    # minimum supported version, use co_qualname instead:
                    #   qualname = curr_frame.f_code.co_qualname
                    #   if qualname.startswith("WorkspaceManager."):
                    if isinstance(curr_frame.f_locals.get("self"), 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

_ensure_workspace_size(required_bytes)

Ensure workspace is allocated and large enough, return current workspace.

Parameters:

• required_bytes

  (int) –

  The number of bytes required.

Returns:

• Tensor –

  The current workspace tensor.

Source code in vllm/v1/worker/workspace.py

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:
                # TODO: This only catches instance methods (self), missing
                # classmethods and staticmethods. Once Python 3.11+ is the
                # minimum supported version, use co_qualname instead:
                #   qualname = curr_frame.f_code.co_qualname
                #   if qualname.startswith("WorkspaceManager."):
                if isinstance(curr_frame.f_locals.get("self"), 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

_workspace_size_bytes(workspace) staticmethod

Get size of workspace in bytes.

Source code in vllm/v1/worker/workspace.py

@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()

get_persistent(key, shape, dtype, *, zero_init=False)

Get a fixed-shape tensor, optionally zeroed on first allocation only.

Source code in vllm/v1/worker/workspace.py

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

get_persistent_resource(key, factory)

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.

Source code in vllm/v1/worker/workspace.py

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])

get_simultaneous(*shapes_and_dtypes)

Get multiple workspace tensors simultaneously from a single allocation.

Parameters:

• *shapes_and_dtypes

  (tuple[tuple[int, ...], dtype], default: () ) –

  One or more (shape, dtype) tuples.

Returns:

• list[Tensor] –

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

Source code in vllm/v1/worker/workspace.py

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))
    ]

is_locked()

Check if workspace is locked.

Source code in vllm/v1/worker/workspace.py

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

lock()

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.

Source code in vllm/v1/worker/workspace.py

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
            ],
        )

unlock()

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.

Source code in vllm/v1/worker/workspace.py

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
            ],
        )

current_workspace_manager()

Get the current workspace manager instance.

Raises:

• AssertionError –

  If workspace manager has not been initialized.

Source code in vllm/v1/worker/workspace.py

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

init_workspace_manager(device, num_ubatches=None, num_lanes=1, alloc_context=nullcontext)

Initialize the workspace manager with a device.

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

Parameters:

• device

  (device) –

  The device to allocate workspace on.

• num_ubatches

  (int | None, default: None ) –

  Number of workspace ubatch slots. Defaults to 1.

• num_lanes

  (int, default: 1 ) –

  Number of independent execution lanes per ubatch. Defaults to 1.

• alloc_context

  (Callable[[], AbstractContextManager], default: nullcontext ) –

  Factory for the context scratch is allocated in.

Source code in vllm/v1/worker/workspace.py

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)

is_workspace_manager_initialized()

Check if workspace manager has been initialized.

Returns:

• bool –

  True if workspace manager is initialized, False otherwise.

Source code in vllm/v1/worker/workspace.py

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

lock_workspace()

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

Source code in vllm/v1/worker/workspace.py

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()

reset_workspace_manager()

Reset the workspace manager to uninitialized state.

This is primarily intended for testing purposes to allow tests to reinitialize the workspace manager cleanly.

Source code in vllm/v1/worker/workspace.py

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

unlock_workspace()

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.

Source code in vllm/v1/worker/workspace.py

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()

use_workspace_lane(lane)

Select an independent workspace owner for this execution context.

Source code in vllm/v1/worker/workspace.py

@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)