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(