| 1 | // A CUDAPluggableAllocator based on cumem* APIs. |
| 2 | // Important: allocation size, CUdeviceptr and CUmemGenericAllocationHandle* |
| 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 <cstdint> |
| 13 | #include <cerrno> |
| 14 | #include <climits> |
| 15 | |
| 16 | // Default chunk size 256MB for ROCm. Can be overridden at runtime by the |
| 17 | // environment variable VLLM_ROCM_SLEEP_MEM_CHUNK_SIZE, specified in megabytes |
| 18 | // (MB). The env value is parsed with strtoull as an integer number of MB |
| 19 | // (decimal or 0x hex). The parsed MB value is converted to bytes. If |
| 20 | // parsing fails, the value is 0, or the multiplication would overflow, |
| 21 | // the default (256MB) is used. |
| 22 | static const unsigned long long DEFAULT_MEMCREATE_CHUNK_SIZE = |
| 23 | (256ULL * 1024ULL * 1024ULL); |
| 24 | |
| 25 | static unsigned long long get_memcreate_chunk_size() { |
| 26 | const char* env = getenv("VLLM_ROCM_SLEEP_MEM_CHUNK_SIZE"); |
| 27 | if (!env) return DEFAULT_MEMCREATE_CHUNK_SIZE; |
| 28 | char* endptr = nullptr; |
| 29 | errno = 0; |
| 30 | unsigned long long val_mb = strtoull(env, &endptr, 0); |
| 31 | if (endptr == env || errno != 0) { |
| 32 | // parsing failed, fallback to default |
| 33 | return DEFAULT_MEMCREATE_CHUNK_SIZE; |
| 34 | } |
| 35 | if (val_mb == 0) return DEFAULT_MEMCREATE_CHUNK_SIZE; |
| 36 | |
| 37 | const unsigned long long MB = 1024ULL * 1024ULL; |
| 38 | // guard against overflow when converting MB -> bytes |
| 39 | if (val_mb > (ULLONG_MAX / MB)) { |
| 40 | return DEFAULT_MEMCREATE_CHUNK_SIZE; |
| 41 | } |
| 42 | return val_mb * MB; |
| 43 | } |
| 44 | |
| 45 | static inline unsigned long long my_min(unsigned long long a, |
| 46 | unsigned long long b) { |
| 47 | return a < b ? a : b; |
| 48 | } |
| 49 | |
| 50 | static CUresult reserve_rocm_address(CUdeviceptr* d_mem, size_t size, |
| 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 | } |
| 56 | |
| 57 | // Some ROCm stacks can report OOM while reserving VA with an explicit |
| 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, addr, 0); |
| 62 | if (status != CUresult(0)) { |
| 63 | return status; |
| 64 | } |
| 65 | if (((std::uintptr_t)(*d_mem) % alignment) == 0) { |
| 66 | return status; |
| 67 | } |
| 68 | |
| 69 | (void)cuMemAddressFree(*d_mem, size); |
| 70 | return hipErrorNotSupported; |
| 71 | } |
| 72 | |
| 73 | static const char* PYARGS_PARSE = "KKKO"; |
| 74 | #endif |
| 75 | |
| 76 | extern "C" { |
| 77 | |
| 78 | #define PY_SSIZE_T_CLEAN |
| 79 | #include <Python.h> |
| 80 | |
| 81 | #include <sys/types.h> |
| 82 | |
| 83 | char error_msg[10240]; // 10KB buffer to store error messages |
| 84 | CUresult no_error = CUresult(0); |
| 85 | CUresult error_code = no_error; // store error code |
| 86 | |
| 87 | #define CUDA_CHECK(condition) \ |
| 88 | do { \ |
| 89 | CUresult error = condition; \ |
| 90 | if (error != 0) { \ |
| 91 | error_code = error; \ |
| 92 | char* error_string; \ |
| 93 | cuGetErrorString(error, (const char**)&error_string); \ |
| 94 | snprintf(error_msg, sizeof(error_msg), "CUDA Error: %s at %s:%d", \ |
| 95 | error_string, __FILE__, __LINE__); \ |
| 96 | std::cerr << error_msg << std::endl; \ |
| 97 | } \ |
| 98 | } while (0) |
| 99 | |
| 100 | // Global references to Python callables |
| 101 | // NOTE: this is borrowed reference, so we don't need to DECREF them. |
| 102 | // This brings the limitation that the allocator needs to be singleton. |
| 103 | static PyObject* g_python_malloc_callback = nullptr; |
| 104 | static PyObject* g_python_free_callback = nullptr; |
| 105 | |
| 106 | // --------------------------------------------------------------------------- |
| 107 | // Helper functions: |
| 108 | |
| 109 | void ensure_context(unsigned long long device) { |
| 110 | CUcontext pctx; |
| 111 | CUDA_CHECK(cuCtxGetCurrent(&pctx)); |
| 112 | if (!pctx) { |
| 113 | // Ensure device context. |
| 114 | CUDA_CHECK(cuDevicePrimaryCtxRetain(&pctx, device)); |
| 115 | CUDA_CHECK(cuCtxSetCurrent(pctx)); |
| 116 | } |
| 117 | } |
| 118 | |
| 119 | void create_and_map(unsigned long long device, ssize_t size, CUdeviceptr d_mem, |
| 120 | #ifndef USE_ROCM |
| 121 | CUmemGenericAllocationHandle* p_memHandle) { |
| 122 | #else |
| 123 | CUmemGenericAllocationHandle** p_memHandle, |
| 124 | unsigned long long* chunk_sizes, size_t num_chunks) { |
| 125 | #endif |
| 126 | ensure_context(device); |
| 127 | // Define memory allocation properties |
| 128 | CUmemAllocationProp prop = {}; |
| 129 | prop.type = CU_MEM_ALLOCATION_TYPE_PINNED; |
| 130 | prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE; |
| 131 | prop.location.id = device; |
| 132 | prop.allocFlags.compressionType = CU_MEM_ALLOCATION_COMP_NONE; |
| 133 | |
| 134 | #ifndef USE_ROCM |
| 135 | int flag = 0; |
| 136 | CUresult rdma_result = cuDeviceGetAttribute( |
| 137 | &flag, CU_DEVICE_ATTRIBUTE_GPU_DIRECT_RDMA_WITH_CUDA_VMM_SUPPORTED, |
| 138 | device); |
| 139 | if (rdma_result == CUDA_SUCCESS && |
| 140 | flag) { // support GPUDirect RDMA if possible |
| 141 | prop.allocFlags.gpuDirectRDMACapable = 1; |
| 142 | } |
| 143 | int fab_flag = 0; |
| 144 | CUresult fab_result = cuDeviceGetAttribute( |
| 145 | &fab_flag, CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED, device); |
| 146 | if (fab_result == CUDA_SUCCESS && |
| 147 | fab_flag) { // support fabric handle if possible |
| 148 | prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_FABRIC; |
| 149 | } |
| 150 | #endif |
| 151 | |
| 152 | #ifndef USE_ROCM |
| 153 | // Allocate memory using cuMemCreate |
| 154 | CUresult ret = (CUresult)cuMemCreate(p_memHandle, size, &prop, 0); |
| 155 | if (ret) { |
| 156 | if (fab_flag && |
| 157 | (ret == CUDA_ERROR_NOT_PERMITTED || ret == CUDA_ERROR_NOT_SUPPORTED)) { |
| 158 | // Fabric allocation may fail without multi-node nvlink, |
| 159 | // fallback to POSIX file descriptor |
| 160 | prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR; |
| 161 | CUDA_CHECK(cuMemCreate(p_memHandle, size, &prop, 0)); |
| 162 | } else { |
| 163 | CUDA_CHECK(ret); |
| 164 | } |
| 165 | } |
| 166 | if (error_code != 0) { |
| 167 | return; |
| 168 | } |
| 169 | CUDA_CHECK(cuMemMap(d_mem, size, 0, *p_memHandle, 0)); |
| 170 | if (error_code != 0) { |
| 171 | return; |
| 172 | } |
| 173 | #else |
| 174 | for (auto i = 0; i < num_chunks; ++i) { |
| 175 | CUDA_CHECK(cuMemCreate(p_memHandle[i], chunk_sizes[i], &prop, 0)); |
| 176 | if (error_code != 0) { |
| 177 | // Clean up previously created handles |
| 178 | for (auto j = 0; j < i; ++j) { |
| 179 | cuMemRelease(*(p_memHandle[j])); |
| 180 | } |
| 181 | return; |
| 182 | } |
| 183 | } |
| 184 | unsigned long long allocated_size = 0; |
| 185 | for (auto i = 0; i < num_chunks; ++i) { |
| 186 | void* map_addr = (void*)((uintptr_t)d_mem + allocated_size); |
| 187 | CUDA_CHECK(cuMemMap(map_addr, chunk_sizes[i], 0, *(p_memHandle[i]), 0)); |
| 188 | if (error_code != 0) { |
| 189 | // unmap previously mapped chunks |
| 190 | unsigned long long unmapped_size = 0; |
| 191 | for (auto j = 0; j < i; ++j) { |
| 192 | void* unmap_addr = (void*)((uintptr_t)d_mem + unmapped_size); |
| 193 | cuMemUnmap(unmap_addr, chunk_sizes[j]); |
| 194 | unmapped_size += chunk_sizes[j]; |
| 195 | } |
| 196 | // release all created handles |
| 197 | for (auto j = 0; j < num_chunks; ++j) { |
| 198 | cuMemRelease(*(p_memHandle[j])); |
| 199 | } |
| 200 | return; |
| 201 | } |
| 202 | allocated_size += chunk_sizes[i]; |
| 203 | } |
| 204 | #endif |
| 205 | |
| 206 | CUmemAccessDesc accessDesc = {}; |
| 207 | accessDesc.location.type = CU_MEM_LOCATION_TYPE_DEVICE; |
| 208 | accessDesc.location.id = device; |
| 209 | accessDesc.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; |
| 210 | |
| 211 | CUDA_CHECK(cuMemSetAccess(d_mem, size, &accessDesc, 1)); |
| 212 | if (error_code != 0) { |
| 213 | return; |
| 214 | } |
| 215 | |
| 216 | #ifdef USE_ROCM |
| 217 | // ROCm-only: zero the freshly (re)mapped region. |
| 218 | // |
| 219 | // On some amdgpu driver versions the physical VRAM backing a buffer is not |
| 220 | // scrubbed on release, so the pages handed back here (e.g. on wake_up after a |
| 221 | // sleep that discarded a tag) can still contain stale data from a previously |
| 222 | // freed allocation -- possibly from another process. Callers of the device |
| 223 | // allocator expect newly mapped memory to be well defined; in particular |
| 224 | // hybrid Mamba / gated-delta-net (GDN) models keep a persisted conv/recurrent |
| 225 | // state cache, and a new sequence that reads an unzeroed slot will propagate |
| 226 | // garbage/NaN, collapsing generation to token id 0 (decoded as "!"). |
| 227 | // |
| 228 | // This is scoped to the ROCm path on purpose: the CUDA path is unchanged. |
| 229 | // Use CUDA_CHECK so a failure sets the global error_code; the Python caller |
| 230 | // (python_create_and_map) only inspects error_code, so returning without |
| 231 | // setting it would make a failed memset look successful and still hand back |
| 232 | // unzeroed memory. |
| 233 | CUDA_CHECK((CUresult)hipMemset(reinterpret_cast<void*>(d_mem), 0, size)); |
| 234 | if (error_code != 0) { |
| 235 | return; |
| 236 | } |
| 237 | #endif |
| 238 | // std::cout << "create_and_map: device=" << device << ", size=" << size << ", |
| 239 | // d_mem=" << d_mem << ", p_memHandle=" << p_memHandle << std::endl; |
| 240 | } |
| 241 | |
| 242 | void unmap_and_release(unsigned long long device, ssize_t size, |
| 243 | CUdeviceptr d_mem, |
| 244 | #ifndef USE_ROCM |
| 245 | CUmemGenericAllocationHandle* p_memHandle) { |
| 246 | #else |
| 247 | CUmemGenericAllocationHandle** p_memHandle, |
| 248 | unsigned long long* chunk_sizes, size_t num_chunks) { |
| 249 | #endif |
| 250 | // std::cout << "unmap_and_release: device=" << device << ", size=" << size << |
| 251 | // ", d_mem=" << d_mem << ", p_memHandle=" << p_memHandle << std::endl; |
| 252 | ensure_context(device); |
| 253 | #ifndef USE_ROCM |
| 254 | CUDA_CHECK(cuMemUnmap(d_mem, size)); |
| 255 | if (error_code != 0) { |
| 256 | return; |
| 257 | } |
| 258 | CUDA_CHECK(cuMemRelease(*p_memHandle)); |
| 259 | if (error_code != 0) { |
| 260 | return; |
| 261 | } |
| 262 | #else |
| 263 | unsigned long long allocated_size = 0; |
| 264 | CUresult first_error = no_error; |
| 265 | |
| 266 | for (auto i = 0; i < num_chunks; ++i) { |
| 267 | void* map_addr = (void*)((uintptr_t)d_mem + allocated_size); |
| 268 | CUresult status = cuMemUnmap(map_addr, chunk_sizes[i]); |
| 269 | if (status != no_error && first_error == no_error) { |
| 270 | first_error = status; |
| 271 | } |
| 272 | allocated_size += chunk_sizes[i]; |
| 273 | } |
| 274 | |
| 275 | for (auto i = 0; i < num_chunks; ++i) { |
| 276 | CUresult status = cuMemRelease(*(p_memHandle[i])); |
| 277 | if (status != no_error && first_error == no_error) { |
| 278 | first_error = status; |
| 279 | } |
| 280 | } |
| 281 | |
| 282 | if (first_error != no_error) { |
| 283 | CUDA_CHECK(first_error); |
| 284 | } |
| 285 | #endif |
| 286 | } |
| 287 | |
| 288 | PyObject* create_tuple_from_c_integers(unsigned long long a, |
| 289 | unsigned long long b, |
| 290 | unsigned long long c, |
| 291 | unsigned long long d) { |
| 292 | // Create a new tuple of size 4 |
| 293 | PyObject* tuple = PyTuple_New(4); |
| 294 | if (!tuple) { |
| 295 | return NULL; // Return NULL on failure |
| 296 | } |
| 297 | |
| 298 | // Convert integers to Python objects and set them in the tuple |
| 299 | PyTuple_SetItem( |
| 300 | tuple, 0, |
| 301 | PyLong_FromUnsignedLongLong(a)); // Steals reference to the PyLong |
| 302 | PyTuple_SetItem(tuple, 1, PyLong_FromUnsignedLongLong(b)); |
| 303 | PyTuple_SetItem(tuple, 2, PyLong_FromUnsignedLongLong(c)); |
| 304 | PyTuple_SetItem(tuple, 3, PyLong_FromUnsignedLongLong(d)); |
| 305 | |
| 306 | // Note: PyTuple_SetItem "steals" a reference to each object, |
| 307 | // so we do not need to Py_DECREF the PyLong objects explicitly. |
| 308 | |
| 309 | return tuple; // Return the created tuple |
| 310 | } |
| 311 | |
| 312 | PyObject* create_tuple_from_c_mixed(unsigned long long a, unsigned long long b, |
| 313 | unsigned long long c, |
| 314 | CUmemGenericAllocationHandle** vec, |
| 315 | unsigned long long* chunk_sizes, |
| 316 | size_t num_chunks) { |
| 317 | PyObject* tuple = PyTuple_New(4); |
| 318 | if (!tuple) { |
| 319 | return NULL; |
| 320 | } |
| 321 | |
| 322 | // PyObject* list = PyList_New(vec.size()); |
| 323 | PyObject* list = PyList_New(num_chunks); |
| 324 | for (auto i = 0; i < num_chunks; ++i) { |
| 325 | PyObject* addr_size_pair = PyTuple_New(2); |
| 326 | PyObject* addr = PyLong_FromUnsignedLongLong((unsigned long long)(vec[i])); |
| 327 | PyObject* size = |
| 328 | PyLong_FromUnsignedLongLong((unsigned long long)(chunk_sizes[i])); |
| 329 | PyTuple_SetItem(addr_size_pair, 0, addr); |
| 330 | PyTuple_SetItem(addr_size_pair, 1, size); |
| 331 | PyList_SetItem(list, i, addr_size_pair); |
| 332 | } |
| 333 | |
| 334 | PyTuple_SetItem(tuple, 0, PyLong_FromUnsignedLongLong(a)); |
| 335 | PyTuple_SetItem(tuple, 1, PyLong_FromUnsignedLongLong(b)); |
| 336 | PyTuple_SetItem(tuple, 2, PyLong_FromUnsignedLongLong(c)); |
| 337 | PyTuple_SetItem(tuple, 3, list); |
| 338 | |
| 339 | return tuple; |
| 340 | } |
| 341 | |
| 342 | // --------------------------------------------------------------------------- |
| 343 | // Our exported C functions that call Python: |
| 344 | |
| 345 | // use CUstream instead of cudaStream_t, to avoid including cuda_runtime_api.h |
| 346 | void* my_malloc(ssize_t size, int device, CUstream stream) { |
| 347 | ensure_context(device); |
| 348 | |
| 349 | // first allocation, align the size, and reserve an address, and also allocate |
| 350 | // a CUmemGenericAllocationHandle |
| 351 | |
| 352 | // Define memory allocation properties |
| 353 | CUmemAllocationProp prop = {}; |
| 354 | prop.type = CU_MEM_ALLOCATION_TYPE_PINNED; |
| 355 | prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE; |
| 356 | prop.location.id = device; |
| 357 | prop.allocFlags.compressionType = CU_MEM_ALLOCATION_COMP_NONE; |
| 358 | |
| 359 | // Check if the allocation is supported |
| 360 | size_t granularity; |
| 361 | CUDA_CHECK(cuMemGetAllocationGranularity(&granularity, &prop, |
| 362 | CU_MEM_ALLOC_GRANULARITY_MINIMUM)); |
| 363 | if (error_code != 0) { |
| 364 | return nullptr; |
| 365 | } |
| 366 | size_t alignedSize = ((size + granularity - 1) / granularity) * granularity; |
| 367 | |
| 368 | CUdeviceptr d_mem; |
| 369 | #ifndef USE_ROCM |
| 370 | CUDA_CHECK(cuMemAddressReserve(&d_mem, alignedSize, 0, 0, 0)); |
| 371 | if (error_code != 0) { |
| 372 | return nullptr; |
| 373 | } |
| 374 | #else |
| 375 | CUDA_CHECK(reserve_rocm_address(&d_mem, alignedSize, granularity)); |
| 376 | if (error_code != 0) { |
| 377 | return nullptr; |
| 378 | } |
| 379 | #endif |
| 380 | |
| 381 | #ifndef USE_ROCM |
| 382 | // allocate the CUmemGenericAllocationHandle |
| 383 | CUmemGenericAllocationHandle* p_memHandle = |
| 384 | (CUmemGenericAllocationHandle*)malloc( |
| 385 | sizeof(CUmemGenericAllocationHandle)); |
| 386 | #else |
| 387 | // Make sure chunk size is aligned with hardware granularity. The base |
| 388 | // chunk size can be configured via environment variable |
| 389 | // ``VLLM_ROCM_SLEEP_MEM_CHUNK_SIZE``; otherwise |
| 390 | // DEFAULT_MEMCREATE_CHUNK_SIZE is used. |
| 391 | size_t base_chunk = (size_t)get_memcreate_chunk_size(); |
| 392 | size_t aligned_chunk_size = |
| 393 | ((base_chunk + granularity - 1) / granularity) * granularity; |
| 394 | size_t num_chunks = |
| 395 | (alignedSize + aligned_chunk_size - 1) / aligned_chunk_size; |
| 396 | CUmemGenericAllocationHandle** p_memHandle = |
| 397 | (CUmemGenericAllocationHandle**)malloc( |
| 398 | num_chunks * sizeof(CUmemGenericAllocationHandle*)); |
| 399 | unsigned long long* chunk_sizes = |
| 400 | (unsigned long long*)malloc(num_chunks * sizeof(unsigned long long)); |
| 401 | for (auto i = 0; i < num_chunks; ++i) { |
| 402 | p_memHandle[i] = (CUmemGenericAllocationHandle*)malloc( |
| 403 | sizeof(CUmemGenericAllocationHandle)); |
| 404 | if (p_memHandle[i] == nullptr) { |
| 405 | std::cerr << "ERROR: malloc failed for p_memHandle[" << i << "].\n"; |
| 406 | for (auto j = 0; j < i; ++j) { |
| 407 | free(p_memHandle[j]); |
| 408 | } |
| 409 | free(p_memHandle); |
| 410 | free(chunk_sizes); |
| 411 | return nullptr; |
| 412 | } |
| 413 | chunk_sizes[i] = (unsigned long long)my_min( |
| 414 | (unsigned long long)(alignedSize - i * aligned_chunk_size), |
| 415 | (unsigned long long)aligned_chunk_size); |
| 416 | } |
| 417 | #endif |
| 418 | |
| 419 | if (!g_python_malloc_callback) { |
| 420 | std::cerr << "ERROR: g_python_malloc_callback not set.\n"; |
| 421 | return nullptr; |
| 422 | } |
| 423 | |
| 424 | // Acquire GIL (not in stable ABI officially, but often works) |
| 425 | PyGILState_STATE gstate = PyGILState_Ensure(); |
| 426 | |
| 427 | #ifndef USE_ROCM |
| 428 | PyObject* arg_tuple = create_tuple_from_c_integers( |
| 429 | (unsigned long long)device, (unsigned long long)alignedSize, |
| 430 | (unsigned long long)d_mem, (unsigned long long)p_memHandle); |
| 431 | #else |
| 432 | PyObject* arg_tuple = create_tuple_from_c_mixed( |
| 433 | (unsigned long long)device, (unsigned long long)alignedSize, |
| 434 | (unsigned long long)d_mem, p_memHandle, chunk_sizes, num_chunks); |
| 435 | #endif |
| 436 | |
| 437 | // Call g_python_malloc_callback |
| 438 | PyObject* py_result = |
| 439 | PyObject_CallFunctionObjArgs(g_python_malloc_callback, arg_tuple, NULL); |
| 440 | Py_DECREF(arg_tuple); |
| 441 | |
| 442 | if (!py_result) { |
| 443 | PyErr_Print(); |
| 444 | PyGILState_Release(gstate); |
| 445 | return nullptr; |
| 446 | } |
| 447 | |
| 448 | PyGILState_Release(gstate); |
| 449 | |
| 450 | // do the final mapping |
| 451 | #ifndef USE_ROCM |
| 452 | create_and_map(device, alignedSize, d_mem, p_memHandle); |
| 453 | #else |
| 454 | create_and_map(device, alignedSize, d_mem, p_memHandle, chunk_sizes, |
| 455 | num_chunks); |
| 456 | free(chunk_sizes); |
| 457 | #endif |
| 458 | |
| 459 | if (error_code != 0) { |
| 460 | // free address and the handle |
| 461 | CUDA_CHECK(cuMemAddressFree(d_mem, alignedSize)); |
| 462 | #ifndef USE_ROCM |
| 463 | free(p_memHandle); |
| 464 | #else |
| 465 | for (size_t i = 0; i < num_chunks; ++i) { |
| 466 | free(p_memHandle[i]); |
| 467 | } |
| 468 | free(p_memHandle); |
| 469 | #endif |
| 470 | return nullptr; |
| 471 | } |
| 472 | |
| 473 | return (void*)d_mem; |
| 474 | } |
| 475 | |
| 476 | // use CUstream instead of cudaStream_t, to avoid including cuda_runtime_api.h |
| 477 | void my_free(void* ptr, ssize_t size, int device, CUstream stream) { |
| 478 | // get memory handle from the pointer |
| 479 | if (!g_python_free_callback) { |
| 480 | std::cerr << "ERROR: g_python_free_callback not set.\n"; |
| 481 | return; |
| 482 | } |
| 483 | |
| 484 | // Acquire GIL (not in stable ABI officially, but often works) |
| 485 | PyGILState_STATE gstate = PyGILState_Ensure(); |
| 486 | |
| 487 | PyObject* py_ptr = |
| 488 | PyLong_FromUnsignedLongLong(reinterpret_cast<unsigned long long>(ptr)); |
| 489 | |
| 490 | PyObject* py_result = |
| 491 | PyObject_CallFunctionObjArgs(g_python_free_callback, py_ptr, NULL); |
| 492 | |
| 493 | if (!py_result || !PyTuple_Check(py_result) || PyTuple_Size(py_result) != 4) { |
| 494 | PyErr_SetString(PyExc_TypeError, "Expected a tuple of size 4"); |
| 495 | Py_XDECREF(py_result); |
| 496 | Py_XDECREF(py_ptr); |
| 497 | return; |
| 498 | } |
| 499 | |
| 500 | unsigned long long recv_device, recv_size; |
| 501 | unsigned long long recv_d_mem; |
| 502 | #ifndef USE_ROCM |
| 503 | unsigned long long recv_p_memHandle; |
| 504 | #else |
| 505 | PyObject* recv_p_memHandle; |
| 506 | #endif |
| 507 | // Unpack the tuple into four C integers |
| 508 | if (!PyArg_ParseTuple(py_result, PYARGS_PARSE, &recv_device, &recv_size, |
| 509 | &recv_d_mem, &recv_p_memHandle)) { |
| 510 | // PyArg_ParseTuple sets an error if it fails |
| 511 | Py_XDECREF(py_result); |
| 512 | Py_XDECREF(py_ptr); |
| 513 | return; |
| 514 | } |
| 515 | |
| 516 | // For ROCm, copy the Python list of (addr,size) pairs into C arrays while |
| 517 | // holding the GIL. Then release the GIL and call the unmap/release helper |
| 518 | // using the copied arrays. This avoids calling PyList_* APIs without the |
| 519 | // GIL (which is undefined behavior and can crash when called from other |
| 520 | // threads). |
| 521 | CUdeviceptr d_mem = (CUdeviceptr)recv_d_mem; |
| 522 | #ifdef USE_ROCM |
| 523 | Py_ssize_t num_chunks = PyList_Size(recv_p_memHandle); |
| 524 | CUmemGenericAllocationHandle** p_memHandle = |
| 525 | (CUmemGenericAllocationHandle**)malloc( |
| 526 | num_chunks * sizeof(CUmemGenericAllocationHandle*)); |
| 527 | if (p_memHandle == nullptr) { |
| 528 | Py_DECREF(py_ptr); |
| 529 | Py_DECREF(py_result); |
| 530 | PyGILState_Release(gstate); |
| 531 | std::cerr << "ERROR: malloc failed for p_memHandle in my_free." |
| 532 | << std::endl; |
| 533 | return; |
| 534 | } |
| 535 | unsigned long long* chunk_sizes = |
| 536 | (unsigned long long*)malloc(num_chunks * sizeof(unsigned long long)); |
| 537 | if (chunk_sizes == nullptr) { |
| 538 | free(p_memHandle); |
| 539 | Py_DECREF(py_ptr); |
| 540 | Py_DECREF(py_result); |
| 541 | PyGILState_Release(gstate); |
| 542 | std::cerr << "ERROR: malloc failed for chunk_sizes in my_free." |
| 543 | << std::endl; |
| 544 | return; |
| 545 | } |
| 546 | for (Py_ssize_t i = 0; i < num_chunks; ++i) { |
| 547 | PyObject* item = PyList_GetItem(recv_p_memHandle, i); |
| 548 | PyObject* addr_py = PyTuple_GetItem(item, 0); |
| 549 | PyObject* size_py = PyTuple_GetItem(item, 1); |
| 550 | p_memHandle[i] = |
| 551 | (CUmemGenericAllocationHandle*)PyLong_AsUnsignedLongLong(addr_py); |
| 552 | chunk_sizes[i] = (unsigned long long)PyLong_AsUnsignedLongLong(size_py); |
| 553 | } |
| 554 | |
| 555 | // Drop temporary Python refs, then release the GIL before calling into |
| 556 | // non-Python APIs. |
| 557 | Py_DECREF(py_ptr); |
| 558 | Py_DECREF(py_result); |
| 559 | PyGILState_Release(gstate); |
| 560 | |
| 561 | // An empty chunk list means this allocation is asleep: its physical chunks |
| 562 | // were already unmapped and released by sleep(), but the virtual address is |
| 563 | // still held as a placeholder reservation. Skip unmap/release (freeing the |
| 564 | // placeholder address happens below). |
| 565 | if (num_chunks > 0) { |
| 566 | unmap_and_release(device, size, d_mem, p_memHandle, chunk_sizes, |
| 567 | num_chunks); |
| 568 | } |
| 569 | #else |
| 570 | // Non-ROCm path: simple integer handle already extracted; drop temporary |
| 571 | // Python refs while still holding the GIL, then release it. |
| 572 | Py_DECREF(py_ptr); |
| 573 | Py_DECREF(py_result); |
| 574 | PyGILState_Release(gstate); |
| 575 | |
| 576 | CUmemGenericAllocationHandle* p_memHandle = |
| 577 | (CUmemGenericAllocationHandle*)recv_p_memHandle; |
| 578 | unmap_and_release(device, size, d_mem, p_memHandle); |
| 579 | #endif |
| 580 | |
| 581 | // Free the virtual address. On ROCm this also covers an asleep allocation, |
| 582 | // whose placeholder reservation made by sleep() is still held here. |
| 583 | CUDA_CHECK(cuMemAddressFree(d_mem, size)); |
| 584 | #ifndef USE_ROCM |
| 585 | free(p_memHandle); |
| 586 | #else |
| 587 | // Only awake allocations have per-chunk handles to free. |
| 588 | for (auto i = 0; i < num_chunks; ++i) { |
| 589 | free(p_memHandle[i]); |
| 590 | } |
| 591 | free(p_memHandle); |
| 592 | free(chunk_sizes); |
| 593 | #endif |
| 594 | } |
| 595 | |
| 596 | // --------------------------------------------------------------------------- |
| 597 | // Python extension boilerplate: |
| 598 | |
| 599 | // Python-exposed function: init_module(python_malloc, python_free) |
| 600 | static PyObject* py_init_module(PyObject* self, PyObject* args) { |
| 601 | PyObject* malloc_callback = nullptr; |
| 602 | PyObject* free_callback = nullptr; |
| 603 | |
| 604 | if (!PyArg_ParseTuple(args, "OO", &malloc_callback, &free_callback)) { |
| 605 | return nullptr; |
| 606 | } |
| 607 | |
| 608 | if (!PyCallable_Check(malloc_callback) || !PyCallable_Check(free_callback)) { |
| 609 | PyErr_SetString(PyExc_TypeError, "Both arguments must be callables"); |
| 610 | return nullptr; |
| 611 | } |
| 612 | |
| 613 | // Save the Python callables |
| 614 | // This module does not handle GC of these objects, so they must be kept alive |
| 615 | // outside of this module. |
| 616 | g_python_malloc_callback = malloc_callback; |
| 617 | g_python_free_callback = free_callback; |
| 618 | |
| 619 | Py_RETURN_NONE; |
| 620 | } |
| 621 | |
| 622 | static PyObject* python_unmap_and_release(PyObject* self, PyObject* args) { |
| 623 | if (!args || !PyTuple_Check(args) || PyTuple_Size(args) != 4) { |
| 624 | PyErr_SetString(PyExc_TypeError, "Expected a tuple of size 4"); |
| 625 | return nullptr; |
| 626 | } |
| 627 | |
| 628 | unsigned long long recv_device, recv_size; |
| 629 | unsigned long long recv_d_mem; |
| 630 | #ifndef USE_ROCM |
| 631 | unsigned long long recv_p_memHandle; |
| 632 | #else |
| 633 | PyObject* recv_p_memHandle; |
| 634 | #endif |
| 635 | // Unpack the tuple into four C integers |
| 636 | if (!PyArg_ParseTuple(args, PYARGS_PARSE, &recv_device, &recv_size, |
| 637 | &recv_d_mem, &recv_p_memHandle)) { |
| 638 | // PyArg_ParseTuple sets an error if it fails |
| 639 | return nullptr; |
| 640 | } |
| 641 | |
| 642 | CUdeviceptr d_mem_ptr = (CUdeviceptr)recv_d_mem; |
| 643 | #ifndef USE_ROCM |
| 644 | CUmemGenericAllocationHandle* p_memHandle = |
| 645 | (CUmemGenericAllocationHandle*)recv_p_memHandle; |
| 646 | |
| 647 | unmap_and_release(recv_device, recv_size, d_mem_ptr, p_memHandle); |
| 648 | #else |
| 649 | if (!PyList_Check(recv_p_memHandle)) { |
| 650 | PyErr_SetString(PyExc_TypeError, |
| 651 | "Expected a list for the 4th argument on ROCm"); |
| 652 | return nullptr; |
| 653 | } |
| 654 | Py_ssize_t num_chunks = PyList_Size(recv_p_memHandle); |
| 655 | if (num_chunks < 0) { |
| 656 | return nullptr; // PyList_Size sets an exception on error. |
| 657 | } |
| 658 | CUmemGenericAllocationHandle** p_memHandle = |
| 659 | (CUmemGenericAllocationHandle**)malloc( |
| 660 | num_chunks * sizeof(CUmemGenericAllocationHandle*)); |
| 661 | if (p_memHandle == nullptr) { |
| 662 | PyErr_SetString(PyExc_MemoryError, "malloc failed for p_memHandle"); |
| 663 | return nullptr; |
| 664 | } |
| 665 | unsigned long long* chunk_sizes = |
| 666 | (unsigned long long*)malloc(num_chunks * sizeof(unsigned long long)); |
| 667 | if (chunk_sizes == nullptr) { |
| 668 | free(p_memHandle); |
| 669 | PyErr_SetString(PyExc_MemoryError, "malloc failed for chunk_sizes"); |
| 670 | return nullptr; |
| 671 | } |
| 672 | for (Py_ssize_t i = 0; i < num_chunks; ++i) { |
| 673 | PyObject* item = PyList_GetItem(recv_p_memHandle, i); |
| 674 | if (item == nullptr || !PyTuple_Check(item) || PyTuple_Size(item) != 2) { |
| 675 | free(p_memHandle); |
| 676 | free(chunk_sizes); |
| 677 | PyErr_SetString( |
| 678 | PyExc_TypeError, |
| 679 | "List items must be tuples of size 2 (handle_addr, size)"); |
| 680 | return nullptr; |
| 681 | } |
| 682 | PyObject* addr_py = PyTuple_GetItem(item, 0); |
| 683 | PyObject* size_py = PyTuple_GetItem(item, 1); |
| 684 | if (addr_py == nullptr || size_py == nullptr) { |
| 685 | free(p_memHandle); |
| 686 | free(chunk_sizes); |
| 687 | return nullptr; // PyTuple_GetItem sets an exception |
| 688 | } |
| 689 | p_memHandle[i] = |
| 690 | (CUmemGenericAllocationHandle*)PyLong_AsUnsignedLongLong(addr_py); |
| 691 | if (PyErr_Occurred()) { |
| 692 | free(p_memHandle); |
| 693 | free(chunk_sizes); |
| 694 | return nullptr; |
| 695 | } |
| 696 | chunk_sizes[i] = (unsigned long long)PyLong_AsUnsignedLongLong(size_py); |
| 697 | if (PyErr_Occurred()) { |
| 698 | free(p_memHandle); |
| 699 | free(chunk_sizes); |
| 700 | return nullptr; |
| 701 | } |
| 702 | } |
| 703 | |
| 704 | unmap_and_release(recv_device, recv_size, d_mem_ptr, p_memHandle, chunk_sizes, |
| 705 | num_chunks); |
| 706 | |
| 707 | // On ROCm/Linux, physical VRAM is only reclaimed once the virtual address |
| 708 | // range is freed; hipMemUnmap + hipMemRelease alone leave the memory |
| 709 | // resident (see ROCm#6021). Free the address to release physical memory, |
| 710 | // then immediately re-reserve the SAME address as an empty placeholder so |
| 711 | // the regular allocator cannot hand it out while we sleep. wake_up remaps |
| 712 | // physical chunks into this placeholder. |
| 713 | if (error_code == no_error) { |
| 714 | CUDA_CHECK(cuMemAddressFree(d_mem_ptr, recv_size)); |
| 715 | if (error_code == no_error) { |
| 716 | CUdeviceptr reserved = 0; |
| 717 | CUDA_CHECK(reserve_rocm_address(&reserved, recv_size, /*alignment=*/0, |
| 718 | d_mem_ptr)); |
| 719 | if (error_code == no_error && reserved != d_mem_ptr) { |
| 720 | (void)cuMemAddressFree(reserved, recv_size); |
| 721 | snprintf(error_msg, sizeof(error_msg), |
| 722 | "failed to re-reserve placeholder address on sleep " |
| 723 | "(requested %#llx, got %#llx)", |
| 724 | (unsigned long long)d_mem_ptr, (unsigned long long)reserved); |
| 725 | error_code = CUresult(1); |
| 726 | } |
| 727 | } |
| 728 | } |
| 729 | |
| 730 | free(p_memHandle); |
| 731 | free(chunk_sizes); |
| 732 | #endif |
| 733 | |
| 734 | if (error_code != 0) { |
| 735 | error_code = no_error; |
| 736 | PyErr_SetString(PyExc_RuntimeError, error_msg); |
| 737 | return nullptr; |
| 738 | } |
| 739 | |
| 740 | Py_RETURN_NONE; |
| 741 | } |
| 742 | |
| 743 | static PyObject* python_create_and_map(PyObject* self, PyObject* args) { |
| 744 | if (!args || !PyTuple_Check(args) || PyTuple_Size(args) != 4) { |
| 745 | PyErr_SetString(PyExc_TypeError, "Expected a tuple of size 4"); |
| 746 | return nullptr; |
| 747 | } |
| 748 | |
| 749 | unsigned long long recv_device, recv_size; |
| 750 | unsigned long long recv_d_mem; |
| 751 | #ifndef USE_ROCM |
| 752 | unsigned long long recv_p_memHandle; |
| 753 | #else |
| 754 | PyObject* recv_p_memHandle; |
| 755 | #endif |
| 756 | // Unpack the tuple into four C integers |
| 757 | if (!PyArg_ParseTuple(args, PYARGS_PARSE, &recv_device, &recv_size, |
| 758 | &recv_d_mem, &recv_p_memHandle)) { |
| 759 | // PyArg_ParseTuple sets an error if it fails |
| 760 | return nullptr; |
| 761 | } |
| 762 | |
| 763 | CUdeviceptr d_mem_ptr = (CUdeviceptr)recv_d_mem; |
| 764 | #ifndef USE_ROCM |
| 765 | CUmemGenericAllocationHandle* p_memHandle = |
| 766 | (CUmemGenericAllocationHandle*)recv_p_memHandle; |
| 767 | |
| 768 | create_and_map(recv_device, recv_size, d_mem_ptr, p_memHandle); |
| 769 | #else |
| 770 | Py_ssize_t num_chunks = PyList_Size(recv_p_memHandle); |
| 771 | CUmemGenericAllocationHandle** p_memHandle = |
| 772 | (CUmemGenericAllocationHandle**)malloc( |
| 773 | num_chunks * sizeof(CUmemGenericAllocationHandle*)); |
| 774 | if (p_memHandle == nullptr) { |
| 775 | PyErr_SetString(PyExc_MemoryError, "malloc failed for p_memHandle"); |
| 776 | return nullptr; |
| 777 | } |
| 778 | unsigned long long* chunk_sizes = |
| 779 | (unsigned long long*)malloc(num_chunks * sizeof(unsigned long long)); |
| 780 | if (chunk_sizes == nullptr) { |
| 781 | free(p_memHandle); |
| 782 | PyErr_SetString(PyExc_MemoryError, "malloc failed for chunk_sizes"); |
| 783 | return nullptr; |
| 784 | } |
| 785 | for (auto i = 0; i < num_chunks; ++i) { |
| 786 | PyObject* item = PyList_GetItem(recv_p_memHandle, i); |
| 787 | PyObject* addr_py = PyTuple_GetItem(item, 0); |
| 788 | PyObject* size_py = PyTuple_GetItem(item, 1); |
| 789 | p_memHandle[i] = |
| 790 | (CUmemGenericAllocationHandle*)PyLong_AsUnsignedLongLong(addr_py); |
| 791 | chunk_sizes[i] = PyLong_AsUnsignedLongLong(size_py); |
| 792 | } |
| 793 | |
| 794 | // Address already reserved as a placeholder by sleep(); just remap chunks. |
| 795 | create_and_map(recv_device, recv_size, d_mem_ptr, p_memHandle, chunk_sizes, |
| 796 | num_chunks); |
| 797 | |
| 798 | free(p_memHandle); |
| 799 | free(chunk_sizes); |
| 800 | #endif |
| 801 | |
| 802 | if (error_code != 0) { |
| 803 | error_code = no_error; |
| 804 | PyErr_SetString(PyExc_RuntimeError, error_msg); |
| 805 | return nullptr; |
| 806 | } |
| 807 | |
| 808 | Py_RETURN_NONE; |
| 809 | } |
| 810 | |
| 811 | static PyMethodDef module_methods[] = { |
| 812 | {"init_module", (PyCFunction)py_init_module, METH_VARARGS, |
| 813 | "Initialize module with python_malloc and python_free callables."}, |
| 814 | {"python_create_and_map", (PyCFunction)python_create_and_map, METH_VARARGS, |
| 815 | "Create and map memory on the device."}, |
| 816 | {"python_unmap_and_release", (PyCFunction)python_unmap_and_release, |
| 817 | METH_VARARGS, "Unmap and release memory on the device."}, |
| 818 | {NULL, NULL, 0, NULL} // sentinel |
| 819 | }; |
| 820 | |
| 821 | static struct PyModuleDef cumem_allocator_module = { |
| 822 | PyModuleDef_HEAD_INIT, "cumem_allocator", |
| 823 | "cumem-based allocator for CUDAPluggableAllocator", -1, module_methods}; |
| 824 | |
| 825 | PyMODINIT_FUNC PyInit_cumem_allocator(void) { |
| 826 | // Initialize the module |
| 827 | PyObject* module = PyModule_Create(&cumem_allocator_module); |
| 828 | if (!module) { |
| 829 | return NULL; |
| 830 | } |
| 831 | return module; |
| 832 | } |
| 833 | } // extern "C" |