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 14, 05:30 UTC
[Bugfix] Add strong reference to CUDA pluggable allocator callbacks (#23477)
Signed-off-by: 22quinn <33176974+22quinn@users.noreply.github.com> Signed-off-by: youkaichao <youkaichao@gmail.com> Co-authored-by: Eric Marcus <eric.marcus@kaiko.ai> Co-authored-by: youkaichao <youkaichao@gmail.com>
22quinn committed
Aug 23, 2025 at 21:56 UTC
9dc30b7068ae07ceca89663e9f8403d00217256d
1 file changed
+9
-4
vllm/device_allocator/cumem.py
+9
-4
@@ -152,8 +152,13 @@ class CuMemAllocator:
152
self.pointer_to_data: dict[int, AllocationData] = {}
153
self.current_tag: str = CuMemAllocator.default_tag
154
self.allocator_and_pools: dict[str, Any] = {}
155
+ # Creating strong references to the two callbacks here to prevent
156
+ # these ephemeral bound-method objects being garbage collected.
157
+ # See discussions in https://github.com/vllm-project/vllm/pull/22724
158
+ self.python_malloc_callback = self._python_malloc_callback
159
+ self.python_free_callback = self._python_free_callback
160
156
- def python_malloc_callback(self, allocation_handle: HandleType) -> None:
161
+ def _python_malloc_callback(self, allocation_handle: HandleType) -> None:
162
"""
163
Internal method to store the allocation data
164
when memory is allocated in the memory pool."""
@@ -162,7 +167,7 @@ class CuMemAllocator:
167
allocation_handle, self.current_tag)
168
return
169
165
- def python_free_callback(self, ptr: int) -> HandleType:
170
+ def _python_free_callback(self, ptr: int) -> HandleType:
171
"""
172
Internal method to look up the allocation data
173
when memory is freed in the memory pool."""
@@ -212,9 +217,9 @@ class CuMemAllocator:
217
def wake_up(self, tags: Optional[list[str]] = None) -> None:
218
"""
219
Wake up the allocator from sleep mode.
215
- All data that is previously offloaded will be loaded back to GPU
220
+ All data that is previously offloaded will be loaded back to GPU
221
memory, and the rest of the data will have empty memory.
217
-
222
+
223
:param tags: The tags of the memory allocation that will be loaded
224
back to GPU memory. If None, all memory allocation will be loaded
225
back to GPU memory.