xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-05-25 14:17:57 +08:00
Code Issues Packages Projects Releases Wiki Activity

[Bugfix] Add strong reference to CUDA pluggable allocator callbacks (#23477)

Browse Source 
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>
This commit is contained in:
22quinn 2025-08-23 21:56:17 -07:00 committed by GitHub
parent 053278a5dc
commit 9dc30b7068
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 9 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
vllm/device_allocator/cumem.py
View File

@ -152,8 +152,13 @@ class CuMemAllocator:
self.pointer_to_data: dict[int, AllocationData] = {} self.pointer_to_data: dict[int, AllocationData] = {}
self.current_tag: str = CuMemAllocator.default_tag self.current_tag: str = CuMemAllocator.default_tag
self.allocator_and_pools: dict[str, Any] = {} self.allocator_and_pools: dict[str, Any] = {}
# Creating strong references to the two callbacks here to prevent
# these ephemeral bound-method objects being garbage collected.
# See discussions in https://github.com/vllm-project/vllm/pull/22724
self.python_malloc_callback = self._python_malloc_callback
self.python_free_callback = self._python_free_callback
def python_malloc_callback(self, allocation_handle: HandleType) -> None: def _python_malloc_callback(self, allocation_handle: HandleType) -> None:
""" """
Internal method to store the allocation data Internal method to store the allocation data
when memory is allocated in the memory pool.""" when memory is allocated in the memory pool."""
@ -162,7 +167,7 @@ class CuMemAllocator:
allocation_handle, self.current_tag) allocation_handle, self.current_tag)
return return
def python_free_callback(self, ptr: int) -> HandleType: def _python_free_callback(self, ptr: int) -> HandleType:
""" """
Internal method to look up the allocation data Internal method to look up the allocation data
when memory is freed in the memory pool.""" when memory is freed in the memory pool."""