From 5dd4a628dde34d5b5737067f131a2499f672f684 Mon Sep 17 00:00:00 2001 From: Matt <156021403+mawong-amd@users.noreply.github.com> Date: Wed, 7 Oct 2026 21:26:12 -0700 Subject: [PATCH] [ROCm][Bugfix] Map sleep-mode cuMem memory for the host (#60367) Signed-off-by: Matthew Wong Co-authored-by: Claude Opus 5.5 (1M context) --- csrc/cumem_allocator.cpp | 23 +++++++++++++++---- .../memory/sleep_mode/test_sleep_mode.py | 1 - 2 files changed, 18 insertions(+), 6 deletions(-) diff --git a/csrc/cumem_allocator.cpp b/csrc/cumem_allocator.cpp index 720c990450..6965d2124a 100644 --- a/csrc/cumem_allocator.cpp +++ b/csrc/cumem_allocator.cpp @@ -203,12 +203,25 @@ void create_and_map(unsigned long long device, ssize_t size, CUdeviceptr d_mem, } #endif - CUmemAccessDesc accessDesc = {}; - accessDesc.location.type = CU_MEM_LOCATION_TYPE_DEVICE; - accessDesc.location.id = device; - accessDesc.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; +#if defined(USE_ROCM) && HIP_VERSION >= 70200000 + // Also map for the host, as PyTorch does for its expandable segments: on + // ROCm >= 7.2, Tensor.item() on a large-BAR device dereferences VRAM + // directly from the host, which segfaults on a device-only mapping. + constexpr int num_desc = 2; +#else + constexpr int num_desc = 1; +#endif + CUmemAccessDesc accessDesc[num_desc] = {}; + accessDesc[0].location.type = CU_MEM_LOCATION_TYPE_DEVICE; + accessDesc[0].location.id = device; + accessDesc[0].flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; +#if defined(USE_ROCM) && HIP_VERSION >= 70200000 + accessDesc[1].location.type = hipMemLocationTypeHost; + accessDesc[1].location.id = 0; // ignored + accessDesc[1].flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; +#endif - CUDA_CHECK(cuMemSetAccess(d_mem, size, &accessDesc, 1)); + CUDA_CHECK(cuMemSetAccess(d_mem, size, accessDesc, num_desc)); if (error_code != 0) { return; } diff --git a/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py b/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py index 73af6e3910..9f6180335a 100644 --- a/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py +++ b/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py @@ -323,7 +323,6 @@ def _attn_kv_scales(model) -> dict[str, list[float]]: } -@pytest.mark.skipif(not current_platform.is_cuda(), reason="Segfaults on ROCm.") @requires_fp8 @create_new_process_for_each_test() def test_deep_sleep_compressed_tensors_kv_scales(