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 11, 08:45 UTC
[BugFix] Fix memory spike in workspace allocation (#30744)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Lucas Wilkinson committed
Dec 16, 2025 at 09:46 UTC
00a8d7628c202f580d5230eaa7fe94338a0549f5
2 files changed
+13
-3
.buildkite/test-pipeline.yaml
+2
@@ -1223,6 +1223,8 @@ steps:
1223
# FIXIT: find out which code initialize cuda before running the test
1224
# before the fix, we need to use spawn to test it
1225
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
1226
+ # Alot of these tests are on the edge of OOMing
1227
+ - export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
1228
# There is some Tensor Parallelism related processing logic in LoRA that
1229
# requires multi-GPU testing for validation.
1230
- pytest -v -s -x lora/test_chatglm3_tp.py
vllm/v1/worker/workspace.py
+11
-3
@@ -145,12 +145,20 @@ class WorkspaceManager:
145
146
for ubatch_id in range(self._num_ubatches):
147
current_workspace = self._current_workspaces[ubatch_id]
148
- if current_workspace is None:
148
+ if (
149
+ current_workspace is None
150
+ or self._workspace_size_bytes(current_workspace) < required_bytes
151
+ ):
152
+ # Delete old tensor before allocating new one to avoid
153
+ # memory spike from resize_(). resize_() allocates new
154
+ # memory before freeing old, which can cause OOM.
155
+ # Must clear the list reference first since local var
156
+ # is just a copy of the reference.
157
+ self._current_workspaces[ubatch_id] = None
158
+ del current_workspace
159
self._current_workspaces[ubatch_id] = torch.empty(
160
(required_bytes,), dtype=torch.uint8, device=self._device
161
)
152
- elif self._workspace_size_bytes(current_workspace) < required_bytes:
153
- current_workspace.resize_(required_bytes)
162
163
if envs.VLLM_DEBUG_WORKSPACE:
164
logger.info(