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_retained –

  Get a scratch view at the end of the workspace that can outlive requests.

• 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
        ]
        # Per slot, (key, start) of the tail span get_retained last handed out;
        # cleared when a scratch request reaches into it or on reallocation.
        self._retained: list[tuple[Hashable, int] | None] = [
            None 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 once the
        manager is locked. 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, _VIEW_ALIGNMENT) 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)
        workspace_id = self._get_workspace_id()
        retained = self._retained[workspace_id]
        if retained is not None and total_bytes > retained[1]:
            self._retained[workspace_id] = None

        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 get_retained(
        self, key: Hashable, shape: tuple[int, ...], dtype: torch.dtype
    ) -> tuple[torch.Tensor, bool]:
        """Get a scratch view at the end of the workspace that can outlive requests.

        ``get_simultaneous`` hands out views from the start of the workspace, so
        requests smaller than the workspace leave its end alone. This view spans
        the last bytes of the workspace, and the returned flag says whether they
        still hold what the previous ``get_retained`` left there: that call used
        the same key and size, no ``get_simultaneous`` request since reached into
        the span, and the workspace was not reallocated in between. What the
        bytes mean is the caller's to track; only the latest key is remembered.

        Args:
            key: Identifies the caller; another key's call takes the span over.
            shape: Shape of the view.
            dtype: Dtype of the view.

        Returns:
            The view and whether its contents are intact.

        """
        num_bytes = _compute_bytes(shape, dtype)
        workspace_id = self._get_workspace_id()
        current_workspace = self._ensure_workspace_size(num_bytes)
        start = round_down(current_workspace.numel() - num_bytes, _VIEW_ALIGNMENT)
        intact = self._retained[workspace_id] == (key, start)
        self._retained[workspace_id] = (key, start)
        view = current_workspace[start : start + num_bytes].view(dtype).reshape(shape)
        return view, intact

    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
            self._retained[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:
                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
        self._retained[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 once the manager is locked. 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 once the
    manager is locked. 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_retained(key, shape, dtype)

Get a scratch view at the end of the workspace that can outlive requests.

get_simultaneous hands out views from the start of the workspace, so requests smaller than the workspace leave its end alone. This view spans the last bytes of the workspace, and the returned flag says whether they still hold what the previous get_retained left there: that call used the same key and size, no get_simultaneous request since reached into the span, and the workspace was not reallocated in between. What the bytes mean is the caller's to track; only the latest key is remembered.

Parameters:

• key

  (Hashable) –

  Identifies the caller; another key's call takes the span over.

• shape

  (tuple[int, ...]) –

  Shape of the view.

• dtype

  (dtype) –

  Dtype of the view.

Returns:

• tuple[Tensor, bool] –

  The view and whether its contents are intact.

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

def get_retained(
    self, key: Hashable, shape: tuple[int, ...], dtype: torch.dtype
) -> tuple[torch.Tensor, bool]:
    """Get a scratch view at the end of the workspace that can outlive requests.

    ``get_simultaneous`` hands out views from the start of the workspace, so
    requests smaller than the workspace leave its end alone. This view spans
    the last bytes of the workspace, and the returned flag says whether they
    still hold what the previous ``get_retained`` left there: that call used
    the same key and size, no ``get_simultaneous`` request since reached into
    the span, and the workspace was not reallocated in between. What the
    bytes mean is the caller's to track; only the latest key is remembered.

    Args:
        key: Identifies the caller; another key's call takes the span over.
        shape: Shape of the view.
        dtype: Dtype of the view.

    Returns:
        The view and whether its contents are intact.

    """
    num_bytes = _compute_bytes(shape, dtype)
    workspace_id = self._get_workspace_id()
    current_workspace = self._ensure_workspace_size(num_bytes)
    start = round_down(current_workspace.numel() - num_bytes, _VIEW_ALIGNMENT)
    intact = self._retained[workspace_id] == (key, start)
    self._retained[workspace_id] = (key, start)
    view = current_workspace[start : start + num_bytes].view(dtype).reshape(shape)
    return view, intact

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, _VIEW_ALIGNMENT) 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)
    workspace_id = self._get_workspace_id()
    retained = self._retained[workspace_id]
    if retained is not None and total_bytes > retained[1]:
        self._retained[workspace_id] = None

    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)