A high-throughput and memory-efficient inference and serving engine for LLMs
[ROCm][CI] Stabilize sleep-mode memory release (#43022)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Jun 9, 2026 at 03:51 UTC
2385e140d6965c279b48f9c50f51e2fddfcf66bd
1 file changed
+25
-1
csrc/cumem_allocator.cpp
+25
-1
@@ -9,6 +9,7 @@
9
static const char* PYARGS_PARSE = "KKKK";
10
#else
11
#include <cstdlib>
12
+ #include <cstdint>
13
#include <cerrno>
14
#include <climits>
15
@@ -46,6 +47,29 @@ static inline unsigned long long my_min(unsigned long long a,
47
return a < b ? a : b;
48
}
49
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);
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, 0, 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
@@ -325,7 +349,7 @@ void* my_malloc(ssize_t size, int device, CUstream stream) {
349
return nullptr;
350
}
351
#else
328
- CUDA_CHECK(cuMemAddressReserve(&d_mem, alignedSize, granularity, 0, 0));
352
+ CUDA_CHECK(reserve_rocm_address(&d_mem, alignedSize, granularity));
353
if (error_code != 0) {
354
return nullptr;
355
}