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 4, 17:30 UTC
[Bugfix] Fix workspace resize leaking reserved GPU memory (#39226)
Signed-off-by: root <conway.zhu@cohere.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
czhu-cohere committed
Apr 23, 2026 at 13:50 UTC
7ff65b19003be4955d2d5d1428e7d94d082559d0
1 file changed
+18
-21
vllm/v1/worker/workspace.py
+18
-21
@@ -161,36 +161,33 @@ class WorkspaceManager:
161
"Workspace growth is not allowed after locking."
162
)
163
164
- for ubatch_id in range(self._num_ubatches):
165
- current_workspace = self._current_workspaces[ubatch_id]
166
- if (
167
- current_workspace is None
168
- or self._workspace_size_bytes(current_workspace) < required_bytes
169
- ):
170
- # Delete old tensor before allocating new one to avoid
171
- # memory spike from resize_(). resize_() allocates new
172
- # memory before freeing old, which can cause OOM.
173
- # Must clear the list reference first since local var
174
- # is just a copy of the reference.
175
- self._current_workspaces[ubatch_id] = None
176
- del current_workspace
177
- self._current_workspaces[ubatch_id] = torch.empty(
178
- (required_bytes,), dtype=torch.uint8, device=self._device
179
- )
164
+ # Only resize the requesting ubatch's workspace. Other
165
+ # ubatches resize lazily on their next get_simultaneous call.
166
+ # Resizing all ubatches here would orphan the other ubatch's
167
+ # old tensor when it still holds views into it (DBO leak).
168
+ self._current_workspaces[ubatch_id] = None
169
+ del current_workspace
170
+ # Release the freed segment back to CUDA so the caching
171
+ # allocator can reuse the GPU memory for the larger
172
+ # allocation below. Without this, each resize may leave a
173
+ # dead segment in reserved memory which can cause higher peak
174
+ # memory usage.
175
+ torch.accelerator.empty_cache()
176
+ self._current_workspaces[ubatch_id] = torch.empty(
177
+ (required_bytes,), dtype=torch.uint8, device=self._device
178
+ )
179
+ current_workspace = self._current_workspaces[ubatch_id]
180
181
if envs.VLLM_DEBUG_WORKSPACE:
182
logger.info(
183
"[WORKSPACE DEBUG] Resized workspace from '%s': %.2f MB -> "
184
- "%.2f MB (%d ubatches, total memory %.2f MB)",
184
+ "%.2f MB (ubatch %d)",
185
get_caller_info(),
186
current_size / _MB,
187
required_bytes / _MB,
188
- self._num_ubatches,
189
- required_bytes * self._num_ubatches / _MB,
188
+ ubatch_id,
189
)
190
192
- current_workspace = self._current_workspaces[dbo_current_ubatch_id()]
193
-
191
return current_workspace
192
193