A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 13:30 UTC
[Bug] Fix workspace manager `_current_workspaces` size (#38853)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Apr 3, 2026 at 21:29 UTC
cab4064cd5a64e3a13ac2af96aa346439b8a6b8f
1 file changed
+5
-3
vllm/v1/worker/workspace.py
+5
-3
@@ -31,7 +31,7 @@ _manager: "WorkspaceManager | None" = None
31
class WorkspaceManager:
32
"""Manager for workspace allocation.
33
34
- Manages workspace buffers for DBO (Dual Batch Overlap) execution.
34
+ Manages one workspace buffer per active ubatch slot.
35
Can be locked to prevent further growth during execution.
36
"""
37
@@ -39,7 +39,9 @@ class WorkspaceManager:
39
self._device = device
40
# Cache num ubatches at init based on configuration (default to 1)
41
self._num_ubatches = num_ubatches if num_ubatches is not None else 1
42
- self._current_workspaces: list[torch.Tensor | None] = [None, None]
42
+ self._current_workspaces: list[torch.Tensor | None] = [
43
+ None
44
+ ] * self._num_ubatches
45
self._locked: bool = False
46
47
@staticmethod
@@ -224,7 +226,7 @@ def init_workspace_manager(
226
227
Args:
228
device: The device to allocate workspace on.
227
- num_ubatches: Number of micro-batches. Defaults to 1.
229
+ num_ubatches: Number of workspace ubatch slots. Defaults to 1.
230
"""
231
global _manager
232
if _manager is not None: