@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 4ca5cd5740
Fork 0

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 13, 08:15 UTC

[Core][AMD] Migrate fully transparent sleep mode to ROCm platform (#12695)

Signed-off-by: Hollow Man <hollowman@opensuse.org> Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com> Co-authored-by: tjtanaa <tunjian.tan@embeddedllm.com> Co-authored-by: kliuae <kuanfu.liu@embeddedllm.com>

ℍ𝕠𝕝𝕝𝕠𝕨 𝕄𝕒𝕟 committed Nov 13, 2025 at 01:24 UTC 4ca5cd5740c0cd7788cdfa8b7ec6a27335607a48
11 files changed +582 -31
CMakeLists.txt
+27 -3
@@ -39,6 +39,13 @@ set(PYTHON_SUPPORTED_VERSIONS "3.10" "3.11" "3.12" "3.13")
39 # Supported AMD GPU architectures.
40 set(HIP_SUPPORTED_ARCHS "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1030;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151")
41
42 +# ROCm installation prefix. Default to /opt/rocm but allow override via
43 +# -DROCM_PATH=/your/rocm/path when invoking cmake.
44 +if(NOT DEFINED ROCM_PATH)
45 + set(ROCM_PATH "/opt/rocm" CACHE PATH "ROCm installation prefix")
46 +else()
47 + set(ROCM_PATH ${ROCM_PATH} CACHE PATH "ROCm installation prefix" FORCE)
48 +endif()
49 #
50 # Supported/expected torch versions for CUDA/ROCm.
51 #
@@ -237,10 +244,27 @@ set_gencode_flags_for_srcs(
244 SRCS "${VLLM_CUMEM_EXT_SRC}"
245 CUDA_ARCHS "${CUDA_ARCHS}")
246
240 -if(VLLM_GPU_LANG STREQUAL "CUDA")
247 +if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
248 message(STATUS "Enabling cumem allocator extension.")
242 - # link against cuda driver library
243 - list(APPEND CUMEM_LIBS CUDA::cuda_driver)
249 + if(VLLM_GPU_LANG STREQUAL "CUDA")
250 + # link against cuda driver library
251 + list(APPEND CUMEM_LIBS CUDA::cuda_driver)
252 + else()
253 + # link against rocm driver library. Prefer an absolute path to
254 + # libamdhip64.so inside ${ROCM_PATH}/lib if available, otherwise fall
255 + # back to linking by name "amdhip64".
256 + find_library(AMDHIP64_LIB
257 + NAMES amdhip64 libamdhip64.so
258 + PATHS ${ROCM_PATH}/lib
259 + NO_DEFAULT_PATH)
260 + if(AMDHIP64_LIB)
261 + message(STATUS "Found libamdhip64 at ${AMDHIP64_LIB}")
262 + list(APPEND CUMEM_LIBS ${AMDHIP64_LIB})
263 + else()
264 + message(WARNING "libamdhip64 not found in ${ROCM_PATH}/lib; falling back to linking 'amdhip64' by name")
265 + list(APPEND CUMEM_LIBS amdhip64)
266 + endif()
267 + endif()
268 define_extension_target(
269 cumem_allocator
270 DESTINATION vllm
csrc/cumem_allocator.cpp
+392 -17
@@ -3,14 +3,58 @@
3 // need to be unsigned long long
4 #include <iostream>
5
6 +#include "cumem_allocator_compat.h"
7 +
8 +#ifndef USE_ROCM
9 +static const char* PYARGS_PARSE = "KKKK";
10 +#else
11 + #include <cstdlib>
12 + #include <cerrno>
13 + #include <climits>
14 +
15 +// Default chunk size 256MB for ROCm. Can be overridden at runtime by the
16 +// environment variable VLLM_ROCM_SLEEP_MEM_CHUNK_SIZE, specified in megabytes
17 +// (MB). The env value is parsed with strtoull as an integer number of MB
18 +// (decimal or 0x hex). The parsed MB value is converted to bytes. If
19 +// parsing fails, the value is 0, or the multiplication would overflow,
20 +// the default (256MB) is used.
21 +static const unsigned long long DEFAULT_MEMCREATE_CHUNK_SIZE =
22 + (256ULL * 1024ULL * 1024ULL);
23 +
24 +static unsigned long long get_memcreate_chunk_size() {
25 + const char* env = getenv("VLLM_ROCM_SLEEP_MEM_CHUNK_SIZE");
26 + if (!env) return DEFAULT_MEMCREATE_CHUNK_SIZE;
27 + char* endptr = nullptr;
28 + errno = 0;
29 + unsigned long long val_mb = strtoull(env, &endptr, 0);
30 + if (endptr == env || errno != 0) {
31 + // parsing failed, fallback to default
32 + return DEFAULT_MEMCREATE_CHUNK_SIZE;
33 + }
34 + if (val_mb == 0) return DEFAULT_MEMCREATE_CHUNK_SIZE;
35 +
36 + const unsigned long long MB = 1024ULL * 1024ULL;
37 + // guard against overflow when converting MB -> bytes
38 + if (val_mb > (ULLONG_MAX / MB)) {
39 + return DEFAULT_MEMCREATE_CHUNK_SIZE;
40 + }
41 + return val_mb * MB;
42 +}
43 +
44 +static inline unsigned long long my_min(unsigned long long a,
45 + unsigned long long b) {
46 + return a < b ? a : b;
47 +}
48 +
49 +static const char* PYARGS_PARSE = "KKKO";
50 +#endif
51 +
52 extern "C" {
53
54 #define PY_SSIZE_T_CLEAN
55 #include <Python.h>
56
57 #include <sys/types.h>
12 -#include <cuda_runtime_api.h>
13 -#include <cuda.h>
58
59 char error_msg[10240]; // 10KB buffer to store error messages
60 CUresult no_error = CUresult(0);
@@ -49,7 +93,12 @@ void ensure_context(unsigned long long device) {
93 }
94
95 void create_and_map(unsigned long long device, ssize_t size, CUdeviceptr d_mem,
96 +#ifndef USE_ROCM
97 CUmemGenericAllocationHandle* p_memHandle) {
98 +#else
99 + CUmemGenericAllocationHandle** p_memHandle,
100 + unsigned long long* chunk_sizes, size_t num_chunks) {
101 +#endif
102 ensure_context(device);
103 // Define memory allocation properties
104 CUmemAllocationProp prop = {};
@@ -58,6 +107,7 @@ void create_and_map(unsigned long long device, ssize_t size, CUdeviceptr d_mem,
107 prop.location.id = device;
108 prop.allocFlags.compressionType = CU_MEM_ALLOCATION_COMP_NONE;
109
110 +#ifndef USE_ROCM
111 // Allocate memory using cuMemCreate
112 CUDA_CHECK(cuMemCreate(p_memHandle, size, &prop, 0));
113 if (error_code != 0) {
@@ -67,6 +117,39 @@ void create_and_map(unsigned long long device, ssize_t size, CUdeviceptr d_mem,
117 if (error_code != 0) {
118 return;
119 }
120 +#else
121 + for (auto i = 0; i < num_chunks; ++i) {
122 + CUDA_CHECK(cuMemCreate(p_memHandle[i], chunk_sizes[i], &prop, 0));
123 + if (error_code != 0) {
124 + // Clean up previously created handles
125 + for (auto j = 0; j < i; ++j) {
126 + cuMemRelease(*(p_memHandle[j]));
127 + }
128 + return;
129 + }
130 + }
131 + unsigned long long allocated_size = 0;
132 + for (auto i = 0; i < num_chunks; ++i) {
133 + void* map_addr = (void*)((uintptr_t)d_mem + allocated_size);
134 + CUDA_CHECK(cuMemMap(map_addr, chunk_sizes[i], 0, *(p_memHandle[i]), 0));
135 + if (error_code != 0) {
136 + // unmap previously mapped chunks
137 + unsigned long long unmapped_size = 0;
138 + for (auto j = 0; j < i; ++j) {
139 + void* unmap_addr = (void*)((uintptr_t)d_mem + unmapped_size);
140 + cuMemUnmap(unmap_addr, chunk_sizes[j]);
141 + unmapped_size += chunk_sizes[j];
142 + }
143 + // release all created handles
144 + for (auto j = 0; j < num_chunks; ++j) {
145 + cuMemRelease(*(p_memHandle[j]));
146 + }
147 + return;
148 + }
149 + allocated_size += chunk_sizes[i];
150 + }
151 +#endif
152 +
153 CUmemAccessDesc accessDesc = {};
154 accessDesc.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
155 accessDesc.location.id = device;
@@ -82,10 +165,16 @@ void create_and_map(unsigned long long device, ssize_t size, CUdeviceptr d_mem,
165
166 void unmap_and_release(unsigned long long device, ssize_t size,
167 CUdeviceptr d_mem,
168 +#ifndef USE_ROCM
169 CUmemGenericAllocationHandle* p_memHandle) {
170 +#else
171 + CUmemGenericAllocationHandle** p_memHandle,
172 + unsigned long long* chunk_sizes, size_t num_chunks) {
173 +#endif
174 // std::cout << "unmap_and_release: device=" << device << ", size=" << size <<
175 // ", d_mem=" << d_mem << ", p_memHandle=" << p_memHandle << std::endl;
176 ensure_context(device);
177 +#ifndef USE_ROCM
178 CUDA_CHECK(cuMemUnmap(d_mem, size));
179 if (error_code != 0) {
180 return;
@@ -94,6 +183,30 @@ void unmap_and_release(unsigned long long device, ssize_t size,
183 if (error_code != 0) {
184 return;
185 }
186 +#else
187 + unsigned long long allocated_size = 0;
188 + CUresult first_error = no_error;
189 +
190 + for (auto i = 0; i < num_chunks; ++i) {
191 + void* map_addr = (void*)((uintptr_t)d_mem + allocated_size);
192 + CUresult status = cuMemUnmap(map_addr, chunk_sizes[i]);
193 + if (status != no_error && first_error == no_error) {
194 + first_error = status;
195 + }
196 + allocated_size += chunk_sizes[i];
197 + }
198 +
199 + for (auto i = 0; i < num_chunks; ++i) {
200 + CUresult status = cuMemRelease(*(p_memHandle[i]));
201 + if (status != no_error && first_error == no_error) {
202 + first_error = status;
203 + }
204 + }
205 +
206 + if (first_error != no_error) {
207 + CUDA_CHECK(first_error);
208 + }
209 +#endif
210 }
211
212 PyObject* create_tuple_from_c_integers(unsigned long long a,
@@ -120,6 +233,36 @@ PyObject* create_tuple_from_c_integers(unsigned long long a,
233 return tuple; // Return the created tuple
234 }
235
236 +PyObject* create_tuple_from_c_mixed(unsigned long long a, unsigned long long b,
237 + unsigned long long c,
238 + CUmemGenericAllocationHandle** vec,
239 + unsigned long long* chunk_sizes,
240 + size_t num_chunks) {
241 + PyObject* tuple = PyTuple_New(4);
242 + if (!tuple) {
243 + return NULL;
244 + }
245 +
246 + // PyObject* list = PyList_New(vec.size());
247 + PyObject* list = PyList_New(num_chunks);
248 + for (auto i = 0; i < num_chunks; ++i) {
249 + PyObject* addr_size_pair = PyTuple_New(2);
250 + PyObject* addr = PyLong_FromUnsignedLongLong((unsigned long long)(vec[i]));
251 + PyObject* size =
252 + PyLong_FromUnsignedLongLong((unsigned long long)(chunk_sizes[i]));
253 + PyTuple_SetItem(addr_size_pair, 0, addr);
254 + PyTuple_SetItem(addr_size_pair, 1, size);
255 + PyList_SetItem(list, i, addr_size_pair);
256 + }
257 +
258 + PyTuple_SetItem(tuple, 0, PyLong_FromUnsignedLongLong(a));
259 + PyTuple_SetItem(tuple, 1, PyLong_FromUnsignedLongLong(b));
260 + PyTuple_SetItem(tuple, 2, PyLong_FromUnsignedLongLong(c));
261 + PyTuple_SetItem(tuple, 3, list);
262 +
263 + return tuple;
264 +}
265 +
266 // ---------------------------------------------------------------------------
267 // Our exported C functions that call Python:
268
@@ -147,14 +290,55 @@ void* my_malloc(ssize_t size, int device, CUstream stream) {
290 size_t alignedSize = ((size + granularity - 1) / granularity) * granularity;
291
292 CUdeviceptr d_mem;
293 +#ifndef USE_ROCM
294 CUDA_CHECK(cuMemAddressReserve(&d_mem, alignedSize, 0, 0, 0));
295 if (error_code != 0) {
296 return nullptr;
297 }
298 +#else
299 + CUDA_CHECK(cuMemAddressReserve(&d_mem, alignedSize, granularity, 0, 0));
300 + if (error_code != 0) {
301 + return nullptr;
302 + }
303 +#endif
304 +
305 +#ifndef USE_ROCM
306 // allocate the CUmemGenericAllocationHandle
307 CUmemGenericAllocationHandle* p_memHandle =
308 (CUmemGenericAllocationHandle*)malloc(
309 sizeof(CUmemGenericAllocationHandle));
310 +#else
311 + // Make sure chunk size is aligned with hardware granularity. The base
312 + // chunk size can be configured via environment variable
313 + // ``VLLM_ROCM_SLEEP_MEM_CHUNK_SIZE``; otherwise
314 + // DEFAULT_MEMCREATE_CHUNK_SIZE is used.
315 + size_t base_chunk = (size_t)get_memcreate_chunk_size();
316 + size_t aligned_chunk_size =
317 + ((base_chunk + granularity - 1) / granularity) * granularity;
318 + size_t num_chunks =
319 + (alignedSize + aligned_chunk_size - 1) / aligned_chunk_size;
320 + CUmemGenericAllocationHandle** p_memHandle =
321 + (CUmemGenericAllocationHandle**)malloc(
322 + num_chunks * sizeof(CUmemGenericAllocationHandle*));
323 + unsigned long long* chunk_sizes =
324 + (unsigned long long*)malloc(num_chunks * sizeof(unsigned long long));
325 + for (auto i = 0; i < num_chunks; ++i) {
326 + p_memHandle[i] = (CUmemGenericAllocationHandle*)malloc(
327 + sizeof(CUmemGenericAllocationHandle));
328 + if (p_memHandle[i] == nullptr) {
329 + std::cerr << "ERROR: malloc failed for p_memHandle[" << i << "].\n";
330 + for (auto j = 0; j < i; ++j) {
331 + free(p_memHandle[j]);
332 + }
333 + free(p_memHandle);
334 + free(chunk_sizes);
335 + return nullptr;
336 + }
337 + chunk_sizes[i] = (unsigned long long)my_min(
338 + (unsigned long long)(alignedSize - i * aligned_chunk_size),
339 + (unsigned long long)aligned_chunk_size);
340 + }
341 +#endif
342
343 if (!g_python_malloc_callback) {
344 std::cerr << "ERROR: g_python_malloc_callback not set.\n";
@@ -164,9 +348,15 @@ void* my_malloc(ssize_t size, int device, CUstream stream) {
348 // Acquire GIL (not in stable ABI officially, but often works)
349 PyGILState_STATE gstate = PyGILState_Ensure();
350
351 +#ifndef USE_ROCM
352 PyObject* arg_tuple = create_tuple_from_c_integers(
353 (unsigned long long)device, (unsigned long long)alignedSize,
354 (unsigned long long)d_mem, (unsigned long long)p_memHandle);
355 +#else
356 + PyObject* arg_tuple = create_tuple_from_c_mixed(
357 + (unsigned long long)device, (unsigned long long)alignedSize,
358 + (unsigned long long)d_mem, p_memHandle, chunk_sizes, num_chunks);
359 +#endif
360
361 // Call g_python_malloc_callback
362 PyObject* py_result =
@@ -182,7 +372,27 @@ void* my_malloc(ssize_t size, int device, CUstream stream) {
372 PyGILState_Release(gstate);
373
374 // do the final mapping
375 +#ifndef USE_ROCM
376 create_and_map(device, alignedSize, d_mem, p_memHandle);
377 +#else
378 + create_and_map(device, alignedSize, d_mem, p_memHandle, chunk_sizes,
379 + num_chunks);
380 + free(chunk_sizes);
381 +#endif
382 +
383 + if (error_code != 0) {
384 + // free address and the handle
385 + CUDA_CHECK(cuMemAddressFree(d_mem, alignedSize));
386 +#ifndef USE_ROCM
387 + free(p_memHandle);
388 +#else
389 + for (size_t i = 0; i < num_chunks; ++i) {
390 + free(p_memHandle[i]);
391 + }
392 + free(p_memHandle);
393 +#endif
394 + return nullptr;
395 + }
396
397 return (void*)d_mem;
398 }
@@ -206,36 +416,96 @@ void my_free(void* ptr, ssize_t size, int device, CUstream stream) {
416
417 if (!py_result || !PyTuple_Check(py_result) || PyTuple_Size(py_result) != 4) {
418 PyErr_SetString(PyExc_TypeError, "Expected a tuple of size 4");
419 + Py_XDECREF(py_result);
420 + Py_XDECREF(py_ptr);
421 return;
422 }
423
424 unsigned long long recv_device, recv_size;
213 - unsigned long long recv_d_mem, recv_p_memHandle;
425 + unsigned long long recv_d_mem;
426 +#ifndef USE_ROCM
427 + unsigned long long recv_p_memHandle;
428 +#else
429 + PyObject* recv_p_memHandle;
430 +#endif
431 // Unpack the tuple into four C integers
215 - if (!PyArg_ParseTuple(py_result, "KKKK", &recv_device, &recv_size,
432 + if (!PyArg_ParseTuple(py_result, PYARGS_PARSE, &recv_device, &recv_size,
433 &recv_d_mem, &recv_p_memHandle)) {
434 // PyArg_ParseTuple sets an error if it fails
435 + Py_XDECREF(py_result);
436 + Py_XDECREF(py_ptr);
437 return;
438 }
439
221 - PyGILState_Release(gstate);
440 + // For ROCm, copy the Python list of (addr,size) pairs into C arrays while
441 + // holding the GIL. Then release the GIL and call the unmap/release helper
442 + // using the copied arrays. This avoids calling PyList_* APIs without the
443 + // GIL (which is undefined behavior and can crash when called from other
444 + // threads).
445 + CUdeviceptr d_mem = (CUdeviceptr)recv_d_mem;
446 +#ifdef USE_ROCM
447 + Py_ssize_t num_chunks = PyList_Size(recv_p_memHandle);
448 + CUmemGenericAllocationHandle** p_memHandle =
449 + (CUmemGenericAllocationHandle**)malloc(
450 + num_chunks * sizeof(CUmemGenericAllocationHandle*));
451 + if (p_memHandle == nullptr) {
452 + Py_DECREF(py_ptr);
453 + Py_DECREF(py_result);
454 + PyGILState_Release(gstate);
455 + std::cerr << "ERROR: malloc failed for p_memHandle in my_free."
456 + << std::endl;
457 + return;
458 + }
459 + unsigned long long* chunk_sizes =
460 + (unsigned long long*)malloc(num_chunks * sizeof(unsigned long long));
461 + if (chunk_sizes == nullptr) {
462 + free(p_memHandle);
463 + Py_DECREF(py_ptr);
464 + Py_DECREF(py_result);
465 + PyGILState_Release(gstate);
466 + std::cerr << "ERROR: malloc failed for chunk_sizes in my_free."
467 + << std::endl;
468 + return;
469 + }
470 + for (Py_ssize_t i = 0; i < num_chunks; ++i) {
471 + PyObject* item = PyList_GetItem(recv_p_memHandle, i);
472 + PyObject* addr_py = PyTuple_GetItem(item, 0);
473 + PyObject* size_py = PyTuple_GetItem(item, 1);
474 + p_memHandle[i] =
475 + (CUmemGenericAllocationHandle*)PyLong_AsUnsignedLongLong(addr_py);
476 + chunk_sizes[i] = (unsigned long long)PyLong_AsUnsignedLongLong(size_py);
477 + }
478
223 - // recv_size == size
224 - // recv_device == device
479 + // Drop temporary Python refs, then release the GIL before calling into
480 + // non-Python APIs.
481 + Py_DECREF(py_ptr);
482 + Py_DECREF(py_result);
483 + PyGILState_Release(gstate);
484
226 - // Free memory
485 + unmap_and_release(device, size, d_mem, p_memHandle, chunk_sizes, num_chunks);
486 +#else
487 + // Non-ROCm path: simple integer handle already extracted; drop temporary
488 + // Python refs while still holding the GIL, then release it.
489 + Py_DECREF(py_ptr);
490 + Py_DECREF(py_result);
491 + PyGILState_Release(gstate);
492
228 - CUdeviceptr d_mem = (CUdeviceptr)recv_d_mem;
493 CUmemGenericAllocationHandle* p_memHandle =
494 (CUmemGenericAllocationHandle*)recv_p_memHandle;
495 unmap_and_release(device, size, d_mem, p_memHandle);
496 +#endif
497
498 // free address and the handle
499 CUDA_CHECK(cuMemAddressFree(d_mem, size));
235 - if (error_code != 0) {
236 - return;
500 +#ifndef USE_ROCM
501 + free(p_memHandle);
502 +#else
503 + for (auto i = 0; i < num_chunks; ++i) {
504 + free(p_memHandle[i]);
505 }
506 free(p_memHandle);
507 + free(chunk_sizes);
508 +#endif
509 }
510
511 // ---------------------------------------------------------------------------
@@ -271,19 +541,87 @@ static PyObject* python_unmap_and_release(PyObject* self, PyObject* args) {
541 }
542
543 unsigned long long recv_device, recv_size;
274 - unsigned long long recv_d_mem, recv_p_memHandle;
544 + unsigned long long recv_d_mem;
545 +#ifndef USE_ROCM
546 + unsigned long long recv_p_memHandle;
547 +#else
548 + PyObject* recv_p_memHandle;
549 +#endif
550 // Unpack the tuple into four C integers
276 - if (!PyArg_ParseTuple(args, "KKKK", &recv_device, &recv_size, &recv_d_mem,
277 - &recv_p_memHandle)) {
551 + if (!PyArg_ParseTuple(args, PYARGS_PARSE, &recv_device, &recv_size,
552 + &recv_d_mem, &recv_p_memHandle)) {
553 // PyArg_ParseTuple sets an error if it fails
554 return nullptr;
555 }
556
557 CUdeviceptr d_mem_ptr = (CUdeviceptr)recv_d_mem;
558 +#ifndef USE_ROCM
559 CUmemGenericAllocationHandle* p_memHandle =
560 (CUmemGenericAllocationHandle*)recv_p_memHandle;
561
562 unmap_and_release(recv_device, recv_size, d_mem_ptr, p_memHandle);
563 +#else
564 + if (!PyList_Check(recv_p_memHandle)) {
565 + PyErr_SetString(PyExc_TypeError,
566 + "Expected a list for the 4th argument on ROCm");
567 + return nullptr;
568 + }
569 + Py_ssize_t num_chunks = PyList_Size(recv_p_memHandle);
570 + if (num_chunks < 0) {
571 + return nullptr; // PyList_Size sets an exception on error.
572 + }
573 + CUmemGenericAllocationHandle** p_memHandle =
574 + (CUmemGenericAllocationHandle**)malloc(
575 + num_chunks * sizeof(CUmemGenericAllocationHandle*));
576 + if (p_memHandle == nullptr) {
577 + PyErr_SetString(PyExc_MemoryError, "malloc failed for p_memHandle");
578 + return nullptr;
579 + }
580 + unsigned long long* chunk_sizes =
581 + (unsigned long long*)malloc(num_chunks * sizeof(unsigned long long));
582 + if (chunk_sizes == nullptr) {
583 + free(p_memHandle);
584 + PyErr_SetString(PyExc_MemoryError, "malloc failed for chunk_sizes");
585 + return nullptr;
586 + }
587 + for (Py_ssize_t i = 0; i < num_chunks; ++i) {
588 + PyObject* item = PyList_GetItem(recv_p_memHandle, i);
589 + if (item == nullptr || !PyTuple_Check(item) || PyTuple_Size(item) != 2) {
590 + free(p_memHandle);
591 + free(chunk_sizes);
592 + PyErr_SetString(
593 + PyExc_TypeError,
594 + "List items must be tuples of size 2 (handle_addr, size)");
595 + return nullptr;
596 + }
597 + PyObject* addr_py = PyTuple_GetItem(item, 0);
598 + PyObject* size_py = PyTuple_GetItem(item, 1);
599 + if (addr_py == nullptr || size_py == nullptr) {
600 + free(p_memHandle);
601 + free(chunk_sizes);
602 + return nullptr; // PyTuple_GetItem sets an exception
603 + }
604 + p_memHandle[i] =
605 + (CUmemGenericAllocationHandle*)PyLong_AsUnsignedLongLong(addr_py);
606 + if (PyErr_Occurred()) {
607 + free(p_memHandle);
608 + free(chunk_sizes);
609 + return nullptr;
610 + }
611 + chunk_sizes[i] = (unsigned long long)PyLong_AsUnsignedLongLong(size_py);
612 + if (PyErr_Occurred()) {
613 + free(p_memHandle);
614 + free(chunk_sizes);
615 + return nullptr;
616 + }
617 + }
618 +
619 + unmap_and_release(recv_device, recv_size, d_mem_ptr, p_memHandle, chunk_sizes,
620 + num_chunks);
621 +
622 + free(p_memHandle);
623 + free(chunk_sizes);
624 +#endif
625
626 if (error_code != 0) {
627 error_code = no_error;
@@ -301,19 +639,56 @@ static PyObject* python_create_and_map(PyObject* self, PyObject* args) {
639 }
640
641 unsigned long long recv_device, recv_size;
304 - unsigned long long recv_d_mem, recv_p_memHandle;
642 + unsigned long long recv_d_mem;
643 +#ifndef USE_ROCM
644 + unsigned long long recv_p_memHandle;
645 +#else
646 + PyObject* recv_p_memHandle;
647 +#endif
648 // Unpack the tuple into four C integers
306 - if (!PyArg_ParseTuple(args, "KKKK", &recv_device, &recv_size, &recv_d_mem,
307 - &recv_p_memHandle)) {
649 + if (!PyArg_ParseTuple(args, PYARGS_PARSE, &recv_device, &recv_size,
650 + &recv_d_mem, &recv_p_memHandle)) {
651 // PyArg_ParseTuple sets an error if it fails
652 return nullptr;
653 }
654
655 CUdeviceptr d_mem_ptr = (CUdeviceptr)recv_d_mem;
656 +#ifndef USE_ROCM
657 CUmemGenericAllocationHandle* p_memHandle =
658 (CUmemGenericAllocationHandle*)recv_p_memHandle;
659
660 create_and_map(recv_device, recv_size, d_mem_ptr, p_memHandle);
661 +#else
662 + Py_ssize_t num_chunks = PyList_Size(recv_p_memHandle);
663 + CUmemGenericAllocationHandle** p_memHandle =
664 + (CUmemGenericAllocationHandle**)malloc(
665 + num_chunks * sizeof(CUmemGenericAllocationHandle*));
666 + if (p_memHandle == nullptr) {
667 + PyErr_SetString(PyExc_MemoryError, "malloc failed for p_memHandle");
668 + return nullptr;
669 + }
670 + unsigned long long* chunk_sizes =
671 + (unsigned long long*)malloc(num_chunks * sizeof(unsigned long long));
672 + if (chunk_sizes == nullptr) {
673 + free(p_memHandle);
674 + PyErr_SetString(PyExc_MemoryError, "malloc failed for chunk_sizes");
675 + return nullptr;
676 + }
677 + for (auto i = 0; i < num_chunks; ++i) {
678 + PyObject* item = PyList_GetItem(recv_p_memHandle, i);
679 + PyObject* addr_py = PyTuple_GetItem(item, 0);
680 + PyObject* size_py = PyTuple_GetItem(item, 1);
681 + p_memHandle[i] =
682 + (CUmemGenericAllocationHandle*)PyLong_AsUnsignedLongLong(addr_py);
683 + chunk_sizes[i] = PyLong_AsUnsignedLongLong(size_py);
684 + }
685 +
686 + create_and_map(recv_device, recv_size, d_mem_ptr, p_memHandle, chunk_sizes,
687 + num_chunks);
688 +
689 + free(p_memHandle);
690 + free(chunk_sizes);
691 +#endif
692
693 if (error_code != 0) {
694 error_code = no_error;
csrc/cumem_allocator_compat.h new
+109
@@ -0,0 +1,109 @@
1 +#pragma once
2 +
3 +#ifdef USE_ROCM
4 +////////////////////////////////////////
5 +// For compatibility with CUDA and ROCm
6 +////////////////////////////////////////
7 + #include <hip/hip_runtime_api.h>
8 +
9 +extern "C" {
10 + #ifndef CUDA_SUCCESS
11 + #define CUDA_SUCCESS hipSuccess
12 + #endif // CUDA_SUCCESS
13 +
14 +// https://rocm.docs.amd.com/projects/HIPIFY/en/latest/tables/CUDA_Driver_API_functions_supported_by_HIP.html
15 +typedef unsigned long long CUdevice;
16 +typedef hipDeviceptr_t CUdeviceptr;
17 +typedef hipError_t CUresult;
18 +typedef hipCtx_t CUcontext;
19 +typedef hipStream_t CUstream;
20 +typedef hipMemGenericAllocationHandle_t CUmemGenericAllocationHandle;
21 +typedef hipMemAllocationGranularity_flags CUmemAllocationGranularity_flags;
22 +typedef hipMemAllocationProp CUmemAllocationProp;
23 +typedef hipMemAccessDesc CUmemAccessDesc;
24 +
25 + #define CU_MEM_ALLOCATION_TYPE_PINNED hipMemAllocationTypePinned
26 + #define CU_MEM_LOCATION_TYPE_DEVICE hipMemLocationTypeDevice
27 + #define CU_MEM_ACCESS_FLAGS_PROT_READWRITE hipMemAccessFlagsProtReadWrite
28 + #define CU_MEM_ALLOC_GRANULARITY_MINIMUM hipMemAllocationGranularityMinimum
29 +
30 + // https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__TYPES.html
31 + #define CU_MEM_ALLOCATION_COMP_NONE 0x0
32 +
33 +// Error Handling
34 +// https://docs.nvidia.com/cuda/archive/11.4.4/cuda-driver-api/group__CUDA__ERROR.html
35 +CUresult cuGetErrorString(CUresult hipError, const char** pStr) {
36 + *pStr = hipGetErrorString(hipError);
37 + return CUDA_SUCCESS;
38 +}
39 +
40 +// Context Management
41 +// https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__CTX.html
42 +CUresult cuCtxGetCurrent(CUcontext* ctx) {
43 + // This API is deprecated on the AMD platform, only for equivalent cuCtx
44 + // driver API on the NVIDIA platform.
45 + return hipCtxGetCurrent(ctx);
46 +}
47 +
48 +CUresult cuCtxSetCurrent(CUcontext ctx) {
49 + // This API is deprecated on the AMD platform, only for equivalent cuCtx
50 + // driver API on the NVIDIA platform.
51 + return hipCtxSetCurrent(ctx);
52 +}
53 +
54 +// Primary Context Management
55 +// https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__PRIMARY__CTX.html
56 +CUresult cuDevicePrimaryCtxRetain(CUcontext* ctx, CUdevice dev) {
57 + return hipDevicePrimaryCtxRetain(ctx, dev);
58 +}
59 +
60 +// Virtual Memory Management
61 +// https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__VA.html
62 +CUresult cuMemAddressFree(CUdeviceptr ptr, size_t size) {
63 + return hipMemAddressFree(ptr, size);
64 +}
65 +
66 +CUresult cuMemAddressReserve(CUdeviceptr* ptr, size_t size, size_t alignment,
67 + CUdeviceptr addr, unsigned long long flags) {
68 + return hipMemAddressReserve(ptr, size, alignment, addr, flags);
69 +}
70 +
71 +CUresult cuMemCreate(CUmemGenericAllocationHandle* handle, size_t size,
72 + const CUmemAllocationProp* prop,
73 + unsigned long long flags) {
74 + return hipMemCreate(handle, size, prop, flags);
75 +}
76 +
77 +CUresult cuMemGetAllocationGranularity(
78 + size_t* granularity, const CUmemAllocationProp* prop,
79 + CUmemAllocationGranularity_flags option) {
80 + return hipMemGetAllocationGranularity(granularity, prop, option);
81 +}
82 +
83 +CUresult cuMemMap(CUdeviceptr dptr, size_t size, size_t offset,
84 + CUmemGenericAllocationHandle handle,
85 + unsigned long long flags) {
86 + return hipMemMap(dptr, size, offset, handle, flags);
87 +}
88 +
89 +CUresult cuMemRelease(CUmemGenericAllocationHandle handle) {
90 + return hipMemRelease(handle);
91 +}
92 +
93 +CUresult cuMemSetAccess(CUdeviceptr ptr, size_t size,
94 + const CUmemAccessDesc* desc, size_t count) {
95 + return hipMemSetAccess(ptr, size, desc, count);
96 +}
97 +
98 +CUresult cuMemUnmap(CUdeviceptr ptr, size_t size) {
99 + return hipMemUnmap(ptr, size);
100 +}
101 +} // extern "C"
102 +
103 +#else
104 +////////////////////////////////////////
105 +// Import CUDA headers for NVIDIA GPUs
106 +////////////////////////////////////////
107 + #include <cuda_runtime_api.h>
108 + #include <cuda.h>
109 +#endif
docs/features/sleep_mode.md
+5 -1
@@ -11,7 +11,7 @@ Key benefits:
11 - **Fine-grained control**: Optionally wake up only model weights or KV cache to avoid OOM during weight updates.
12
13 !!! note
14 - This feature is only supported on CUDA platform.
14 + This feature is now supported on CUDA and ROCm platform.
15
16 !!! note
17 For more information, see this [Blog Post](https://blog.vllm.ai/2025/10/26/sleep-mode.html).
@@ -116,3 +116,7 @@ curl -X POST 'http://localhost:8000/wake_up?tags=kv_cache'
116
117 !!! note
118 These endpoints are only available when passing `VLLM_SERVER_DEV_MODE=1`.
119 +
120 +## Limitation
121 +
122 +On ROCm, the virtual memory allocation on ROCm is done through chunked memory allocation. You can control the chunk size through `VLLM_ROCM_SLEEP_MEM_CHUNK_SIZE` (in MB). The default value is set at 256MB. The larger the chunk size the faster the performance. However, setting it too large will cause OOM. So if you encounter OOM when using sleep mode. Try reducing the chunk size. It is recommended to define the chunk size as a power of 2.
setup.py
+3 -1
@@ -208,6 +208,8 @@ class cmake_build_ext(build_ext):
208 # Make sure we use the nvcc from CUDA_HOME
209 if _is_cuda():
210 cmake_args += [f"-DCMAKE_CUDA_COMPILER={CUDA_HOME}/bin/nvcc"]
211 + elif _is_hip():
212 + cmake_args += [f"-DROCM_PATH={ROCM_HOME}"]
213
214 other_cmake_args = os.environ.get("CMAKE_ARGS")
215 if other_cmake_args:
@@ -628,6 +630,7 @@ ext_modules = []
630
631 if _is_cuda() or _is_hip():
632 ext_modules.append(CMakeExtension(name="vllm._moe_C"))
633 + ext_modules.append(CMakeExtension(name="vllm.cumem_allocator"))
634
635 if _is_hip():
636 ext_modules.append(CMakeExtension(name="vllm._rocm_C"))
@@ -643,7 +646,6 @@ if _is_cuda():
646 ext_modules.append(
647 CMakeExtension(name="vllm._flashmla_extension_C", optional=True)
648 )
646 - ext_modules.append(CMakeExtension(name="vllm.cumem_allocator"))
649
650 if _build_custom_ops():
651 ext_modules.append(CMakeExtension(name="vllm._C"))
tests/basic_correctness/test_cumem.py
+5 -4
@@ -8,12 +8,13 @@ import torch
8
9 from vllm import LLM, AsyncEngineArgs, AsyncLLMEngine, SamplingParams
10 from vllm.device_allocator.cumem import CuMemAllocator
11 +from vllm.platforms import current_platform
12 from vllm.utils.mem_constants import GiB_bytes
13
14 from ..utils import create_new_process_for_each_test
15
16
16 -@create_new_process_for_each_test()
17 +@create_new_process_for_each_test("fork" if not current_platform.is_rocm() else "spawn")
18 def test_python_error():
19 """
20 Test if Python error occurs when there's low-level
@@ -39,7 +40,7 @@ def test_python_error():
40 allocator.wake_up()
41
42
42 -@create_new_process_for_each_test()
43 +@create_new_process_for_each_test("fork" if not current_platform.is_rocm() else "spawn")
44 def test_basic_cumem():
45 # some tensors from default memory pool
46 shape = (1024, 1024)
@@ -72,7 +73,7 @@ def test_basic_cumem():
73 assert torch.allclose(output, torch.ones_like(output) * 3)
74
75
75 -@create_new_process_for_each_test()
76 +@create_new_process_for_each_test("fork" if not current_platform.is_rocm() else "spawn")
77 def test_cumem_with_cudagraph():
78 allocator = CuMemAllocator.get_instance()
79 with allocator.use_memory_pool():
@@ -117,7 +118,7 @@ def test_cumem_with_cudagraph():
118 assert torch.allclose(y, x + 1)
119
120
120 -@create_new_process_for_each_test()
121 +@create_new_process_for_each_test("fork" if not current_platform.is_rocm() else "spawn")
122 @pytest.mark.parametrize(
123 "model",
124 [
vllm/config/model.py
+2 -1
@@ -264,7 +264,8 @@ class ModelConfig:
264 merged with the default config from the model. If used with
265 `--generation-config vllm`, only the override parameters are used."""
266 enable_sleep_mode: bool = False
267 - """Enable sleep mode for the engine (only cuda platform is supported)."""
267 + """Enable sleep mode for the engine (only cuda and
268 + hip platforms are supported)."""
269 model_impl: str | ModelImpl = "auto"
270 """Which implementation of the model to use:\n
271 - "auto" will try to use the vLLM implementation, if it exists, and fall
vllm/device_allocator/cumem.py
+1 -1
@@ -63,7 +63,7 @@ try:
63 libcudart = CudaRTLibrary()
64 cumem_available = True
65 except ModuleNotFoundError:
66 - # rocm platform does not support cumem allocator
66 + # only cuda and rocm platforms support cumem allocator
67 init_module = None
68 python_create_and_map = None
69 python_unmap_and_release = None
vllm/distributed/device_communicators/cuda_wrapper.py
+28 -2
@@ -14,6 +14,7 @@ import torch # noqa
14
15 import vllm.envs as envs
16 from vllm.logger import init_logger
17 +from vllm.platforms import current_platform
18
19 logger = init_logger(__name__)
20
@@ -105,6 +106,20 @@ class CudaRTLibrary:
106 ),
107 ]
108
109 + # https://rocm.docs.amd.com/projects/HIPIFY/en/latest/tables/CUDA_Runtime_API_functions_supported_by_HIP.html # noqa
110 + cuda_to_hip_mapping = {
111 + "cudaSetDevice": "hipSetDevice",
112 + "cudaDeviceSynchronize": "hipDeviceSynchronize",
113 + "cudaDeviceReset": "hipDeviceReset",
114 + "cudaGetErrorString": "hipGetErrorString",
115 + "cudaMalloc": "hipMalloc",
116 + "cudaFree": "hipFree",
117 + "cudaMemset": "hipMemset",
118 + "cudaMemcpy": "hipMemcpy",
119 + "cudaIpcGetMemHandle": "hipIpcGetMemHandle",
120 + "cudaIpcOpenMemHandle": "hipIpcOpenMemHandle",
121 + }
122 +
123 # class attribute to store the mapping from the path to the library
124 # to avoid loading the same library multiple times
125 path_to_library_cache: dict[str, Any] = {}
@@ -117,7 +132,13 @@ class CudaRTLibrary:
132 if so_file is None:
133 so_file = find_loaded_library("libcudart")
134 if so_file is None:
120 - so_file = envs.VLLM_CUDART_SO_PATH # fallback to env var
135 + # libcudart is not loaded in the current process, try hip
136 + so_file = find_loaded_library("libamdhip64")
137 + # should be safe to assume now that we are using ROCm
138 + # as the following assertion should error out if the
139 + # libhiprtc library is also not loaded
140 + if so_file is None:
141 + so_file = envs.VLLM_CUDART_SO_PATH # fallback to env var
142 assert so_file is not None, (
143 "libcudart is not loaded in the current process, "
144 "try setting VLLM_CUDART_SO_PATH"
@@ -130,7 +151,12 @@ class CudaRTLibrary:
151 if so_file not in CudaRTLibrary.path_to_dict_mapping:
152 _funcs = {}
153 for func in CudaRTLibrary.exported_functions:
133 - f = getattr(self.lib, func.name)
154 + f = getattr(
155 + self.lib,
156 + CudaRTLibrary.cuda_to_hip_mapping[func.name]
157 + if current_platform.is_rocm()
158 + else func.name,
159 + )
160 f.restype = func.restype
161 f.argtypes = func.argtypes
162 _funcs[func.name] = f
vllm/envs.py
+5
@@ -18,6 +18,7 @@ if TYPE_CHECKING:
18 VLLM_RINGBUFFER_WARNING_INTERVAL: int = 60
19 VLLM_NCCL_SO_PATH: str | None = None
20 LD_LIBRARY_PATH: str | None = None
21 + VLLM_ROCM_SLEEP_MEM_CHUNK_SIZE: int = 256
22 VLLM_V1_USE_PREFILL_DECODE_ATTENTION: bool = False
23 VLLM_FLASH_ATTN_VERSION: int | None = None
24 LOCAL_RANK: int = 0
@@ -520,6 +521,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
521 # when `VLLM_NCCL_SO_PATH` is not set, vllm will try to find the nccl
522 # library file in the locations specified by `LD_LIBRARY_PATH`
523 "LD_LIBRARY_PATH": lambda: os.environ.get("LD_LIBRARY_PATH", None),
524 + # flag to control the chunk size (in MB) for sleeping memory allocations under ROCm
525 + "VLLM_ROCM_SLEEP_MEM_CHUNK_SIZE": lambda: int(
526 + os.environ.get("VLLM_ROCM_SLEEP_MEM_CHUNK_SIZE", "256")
527 + ),
528 # Use separate prefill and decode kernels for V1 attention instead of
529 # the unified triton kernel.
530 "VLLM_V1_USE_PREFILL_DECODE_ATTENTION": lambda: (
vllm/platforms/interface.py
+5 -1
@@ -171,7 +171,11 @@ class Platform:
171 return self._enum in (PlatformEnum.CUDA, PlatformEnum.ROCM)
172
173 def is_sleep_mode_available(self) -> bool:
174 - return self._enum == PlatformEnum.CUDA
174 + # TODO: Actually only mi3xx has the sleep mode support now
175 + # for ROCm, but currently we don't have a way to detect the
176 + # exact GPU model statelessly here. So we return True for
177 + # all ROCm platforms for now.
178 + return self._enum in (PlatformEnum.CUDA, PlatformEnum.ROCM)
179
180 @classmethod
181 def device_id_to_physical_device_id(cls, device_id: int):