@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 68afd78897

[Bugfix][ROCm] Fix cumem sleep and teardown (#46203)

Signed-off-by: pei.zhang <pei.zhang@amd.com> Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Matthew Wong <Matthew.Wong2@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>

peizhang56 committed Jun 23, 2026 at 11:45 UTC 68afd7889723c1538af7993b6e568b01945ce43a

7 files changed +130 -18

.buildkite/test-amd.yaml

+1 -1

@@ -632,7 +632,7 @@ steps:

632 #----------------------------------------------------- mi300 · basic_correctness -----------------------------------------------------#

633

634 - label: Basic Correctness # TBD

635 - timeout_in_minutes: 180

635 + timeout_in_minutes: 40

636 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]

637 agent_pool: mi300_1

638 fast_check: true

.buildkite/test_areas/basic_correctness.yaml

@@ -16,3 +16,9 @@ steps:

16 - pytest -v -s basic_correctness/test_mem.py

17 - pytest -v -s basic_correctness/test_basic_correctness.py

18 - pytest -v -s basic_correctness/test_cpu_offload.py

19 + mirror:

20 + amd:

21 + device: mi325_1

22 + timeout_in_minutes: 40

23 + depends_on:

24 + - image-build-amd

csrc/cumem_allocator.cpp

+38 -5

@@ -48,8 +48,8 @@ static inline unsigned long long my_min(unsigned long long a,

48 }

50 static CUresult reserve_rocm_address(CUdeviceptr* d_mem, size_t size,

51 - size_t alignment) {

52 - CUresult status = cuMemAddressReserve(d_mem, size, alignment, 0, 0);

51 + size_t alignment, CUdeviceptr addr = 0) {

52 + CUresult status = cuMemAddressReserve(d_mem, size, alignment, addr, 0);

53 if (status == CUresult(0) || alignment == 0) {

54 return status;

55 }

@@ -58,7 +58,7 @@ static CUresult reserve_rocm_address(CUdeviceptr* d_mem, size_t size,

58 // alignment even when physical VRAM is free. Let HIP choose the default

59 // alignment, then verify that the returned address still satisfies the

60 // requested alignment before accepting it.

61 - status = cuMemAddressReserve(d_mem, size, 0, 0, 0);

61 + status = cuMemAddressReserve(d_mem, size, 0, addr, 0);

62 if (status != CUresult(0)) {

63 return status;

64 }

@@ -535,7 +535,14 @@ void my_free(void* ptr, ssize_t size, int device, CUstream stream) {

535 Py_DECREF(py_result);

536 PyGILState_Release(gstate);

537

538 - unmap_and_release(device, size, d_mem, p_memHandle, chunk_sizes, num_chunks);

538 + // An empty chunk list means this allocation is asleep: its physical chunks

539 + // were already unmapped and released by sleep(), but the virtual address is

540 + // still held as a placeholder reservation. Skip unmap/release (freeing the

541 + // placeholder address happens below).

542 + if (num_chunks > 0) {

543 + unmap_and_release(device, size, d_mem, p_memHandle, chunk_sizes,

544 + num_chunks);

545 + }

546 #else

547 // Non-ROCm path: simple integer handle already extracted; drop temporary

548 // Python refs while still holding the GIL, then release it.

@@ -548,11 +555,13 @@ void my_free(void* ptr, ssize_t size, int device, CUstream stream) {

555 unmap_and_release(device, size, d_mem, p_memHandle);

556 #endif

557

551 - // free address and the handle

558 + // Free the virtual address. On ROCm this also covers an asleep allocation,

559 + // whose placeholder reservation made by sleep() is still held here.

560 CUDA_CHECK(cuMemAddressFree(d_mem, size));

561 #ifndef USE_ROCM

562 free(p_memHandle);

563 #else

564 + // Only awake allocations have per-chunk handles to free.

565 for (auto i = 0; i < num_chunks; ++i) {

566 free(p_memHandle[i]);

567 }

@@ -672,6 +681,29 @@ static PyObject* python_unmap_and_release(PyObject* self, PyObject* args) {

681 unmap_and_release(recv_device, recv_size, d_mem_ptr, p_memHandle, chunk_sizes,

682 num_chunks);

683

684 + // On ROCm/Linux, physical VRAM is only reclaimed once the virtual address

685 + // range is freed; hipMemUnmap + hipMemRelease alone leave the memory

686 + // resident (see ROCm#6021). Free the address to release physical memory,

687 + // then immediately re-reserve the SAME address as an empty placeholder so

688 + // the regular allocator cannot hand it out while we sleep. wake_up remaps

689 + // physical chunks into this placeholder.

690 + if (error_code == no_error) {

691 + CUDA_CHECK(cuMemAddressFree(d_mem_ptr, recv_size));

692 + if (error_code == no_error) {

693 + CUdeviceptr reserved = 0;

694 + CUDA_CHECK(reserve_rocm_address(&reserved, recv_size, /*alignment=*/0,

695 + d_mem_ptr));

696 + if (error_code == no_error && reserved != d_mem_ptr) {

697 + (void)cuMemAddressFree(reserved, recv_size);

698 + snprintf(error_msg, sizeof(error_msg),

699 + "failed to re-reserve placeholder address on sleep "

700 + "(requested %#llx, got %#llx)",

701 + (unsigned long long)d_mem_ptr, (unsigned long long)reserved);

702 + error_code = CUresult(1);

703 + }

704 + }

705 + }

706 +

707 free(p_memHandle);

708 free(chunk_sizes);

709 #endif

@@ -736,6 +768,7 @@ static PyObject* python_create_and_map(PyObject* self, PyObject* args) {

768 chunk_sizes[i] = PyLong_AsUnsignedLongLong(size_py);

769 }

770

771 + // Address already reserved as a placeholder by sleep(); just remap chunks.

772 create_and_map(recv_device, recv_size, d_mem_ptr, p_memHandle, chunk_sizes,

773 num_chunks);

774

vllm/device_allocator/__init__.py

+4 -2

@@ -3,14 +3,15 @@

4 import dataclasses

5 from contextlib import AbstractContextManager

6 -from typing import Protocol

6 +from typing import Protocol, TypeAlias

8 import torch

10 from vllm.platforms import current_platform

12 # py_device, py_size_or_aligned_size, py_ptr, py_handle

13 -HandleType = tuple[int, int, int, int]

13 +# py_handle has type list[int] on ROCm and int otherwise

14 +HandleType: TypeAlias = tuple[int, int, int, list[int] | int]

17 @dataclasses.dataclass

@@ -18,6 +19,7 @@ class AllocationData:

19 handle: HandleType

20 tag: str

21 cpu_backup_tensor: torch.Tensor | None = None

22 + is_asleep: bool = False

25 class MemAllocator(Protocol):

vllm/device_allocator/cumem.py

+67 -1

@@ -8,6 +8,7 @@

8 # both of them failed because of cuda context mismatch.

9 # not sure why, they are created from a different context.

10 # the only successful approach is to call cuda driver API in C.

11 +import atexit

12 import gc

13 import os

14 from collections.abc import Callable, Iterator

@@ -18,6 +19,7 @@ import torch

20 from vllm.device_allocator import AllocationData, HandleType

21 from vllm.logger import init_logger

22 +from vllm.platforms import current_platform

23 from vllm.utils.system_utils import find_loaded_library

24 from vllm.utils.torch_utils import PIN_MEMORY

@@ -115,8 +117,21 @@ class CuMemAllocator:

117 assert cumem_available, "cumem allocator is not available"

118 if CuMemAllocator.instance is None:

119 CuMemAllocator.instance = CuMemAllocator()

120 + # Ensure MemPool/allocator wrappers are released before interpreter

121 + # finalization tears down PyTorch allocator internals.

122 + atexit.register(CuMemAllocator._shutdown_singleton)

123 return CuMemAllocator.instance

124

125 + @staticmethod

126 + def _shutdown_singleton() -> None:

127 + instance = CuMemAllocator.instance

128 + if instance is None:

129 + return

130 + try:

131 + instance.release_pools()

132 + except Exception:

133 + logger.exception("CuMemAllocator singleton shutdown failed")

134 +

135 def __init__(self):

136 self.pointer_to_data: dict[int, AllocationData] = {}

137 self.current_tag: str = CuMemAllocator.default_tag

@@ -127,6 +142,45 @@ class CuMemAllocator:

142 self.python_malloc_callback = self._python_malloc_callback

143 self.python_free_callback = self._python_free_callback

144

145 + def release_pools(self) -> None:

146 + """Drop Python references to MemPool/pluggable allocators eagerly.

147 +

148 + A cumem ``MemPool`` outlives the ``use_memory_pool`` context (a strong

149 + reference is kept in ``allocator_and_pools`` to work around

150 + pytorch/pytorch#146431), and a captured CUDA graph can keep it alive

151 + longer still. ``MemPool`` only holds a non-owning pointer to the

152 + allocator, whose owning reference lives in the Python

153 + ``CUDAPluggableAllocator``. If both are instead dropped during

154 + interpreter shutdown, GC may finalize the allocator first; the eventual

155 + ``~MemPool`` -> ``emptyCache`` -> ``release_block`` then makes a virtual

156 + call into the freed allocator -- aborting the process with "pure virtual

157 + method called" (pytorch/pytorch#145168).

158 +

159 + Release the kept-alive pools before interpreter finalization, and keep

160 + the pluggable allocator wrappers alive while MemPool destructors run.

161 + This is safe to call more than once.

162 + """

163 + if not self.allocator_and_pools:

164 + return

165 +

166 + pool_entries = list(self.allocator_and_pools.values())

167 + self.allocator_and_pools.clear()

168 +

169 + mem_pools = [entry[0] for entry in pool_entries]

170 + allocators = [entry[1] for entry in pool_entries]

171 + pool_entries.clear()

172 +

173 + # Phase 1: drop MemPool refs while allocators are still strongly held.

174 + mem_pools.clear()

175 + gc.collect()

176 +

177 + # Phase 2: now it is safe to release allocator wrappers.

178 + allocators.clear()

179 +

180 + def close(self) -> None:

181 + """Compatibility alias for deterministic pool release."""

182 + self.release_pools()

183 +

184 def _python_malloc_callback(self, allocation_handle: HandleType) -> None:

185 """

186 Internal method to store the allocation data

@@ -150,6 +204,14 @@ class CuMemAllocator:

204 data = self.pointer_to_data.pop(ptr)

205 if data.cpu_backup_tensor is not None:

206 data.cpu_backup_tensor = None

207 + if data.is_asleep and current_platform.is_rocm():

208 + # On ROCm, sleep() already unmapped and released this allocation's

209 + # physical chunks and holds its virtual address as a placeholder

210 + # reservation. Return a handle with an empty chunk list so the C

211 + # extension skips unmap/release (avoiding a double-free) while

212 + # still freeing the placeholder address.

213 + device, size, d_mem, _ = data.handle

214 + return (device, size, d_mem, [])

215 # Drain pending kernels before the C extension's cuMemUnmap.

216 # The pluggable allocator path doesn't defer reclaim like the

217 # regular caching allocator, so without this, in-flight work

@@ -201,7 +263,10 @@ class CuMemAllocator:

263 cpu_ptr = cpu_backup_tensor.data_ptr()

264 libcudart.cudaMemcpy(cpu_ptr, ptr, size_in_bytes)

265 data.cpu_backup_tensor = cpu_backup_tensor

204 - unmap_and_release(handle)

266 + try:

267 + unmap_and_release(handle)

268 + finally:

269 + data.is_asleep = True

270

271 logger.info(

272 "CuMemAllocator: sleep freed %.2f GiB memory in total, of which "

@@ -230,6 +295,7 @@ class CuMemAllocator:

295 if tags is None or data.tag in tags:

296 handle = data.handle

297 create_and_map(handle)

298 + data.is_asleep = False

299 if data.cpu_backup_tensor is not None:

300 cpu_backup_tensor = data.cpu_backup_tensor

301 if cpu_backup_tensor is not None:

vllm/distributed/device_communicators/cuda_wrapper.py

+6 -9

@@ -104,15 +104,12 @@ class CudaRTLibrary:

104

105 def __init__(self, so_file: str | None = None):

106 if so_file is None:

107 - so_file = find_loaded_library("libcudart")

108 - if so_file is None:

109 - # libcudart is not loaded in the current process, try hip

110 - so_file = find_loaded_library("libamdhip64")

111 - # should be safe to assume now that we are using ROCm

112 - # as the following assertion should error out if the

113 - # libhiprtc library is also not loaded

114 - if so_file is None:

115 - so_file = envs.VLLM_CUDART_SO_PATH # fallback to env var

107 + so_file = (

108 + find_loaded_library(

109 + "libamdhip64" if current_platform.is_rocm() else "libcudart"

110 + )

111 + or envs.VLLM_CUDART_SO_PATH # fallback to env var

112 + )

113 assert so_file is not None, (

114 "libcudart is not loaded in the current process, "

115 "try setting VLLM_CUDART_SO_PATH"

vllm/v1/worker/gpu_worker.py

@@ -1185,6 +1185,14 @@ class Worker(WorkerBase):

1185 if model_runner := getattr(self, "model_runner", None):

1186 model_runner.shutdown()

1187

1188 + # Release kept-alive cumem pools while the pluggable allocator wrappers

1189 + # and callbacks are still alive, so MemPool teardown is not deferred to

1190 + # interpreter finalization (pytorch/pytorch#145168).

1191 + from vllm.device_allocator.cumem import CuMemAllocator

1192 +

1193 + if CuMemAllocator.instance is not None:

1194 + CuMemAllocator.instance.release_pools()

1195 +

1196 def elastic_ep_execute(self, execute_method: str, *args, **kwargs):

1197 return self.elastic_ep_executor.execute(execute_method, *args, **kwargs)

1198