torch.cuda._use_uvm() already works on ROCm: it imports cuda.bindings.runtime directly, and hip-python's interop package (PyPI: hip-python-interop) provides a HIP-backed cuda.bindings with the cuda-python 12+ module layout, verified live against HIP 7.15 (cudaMallocManaged / cudaMemAdvise / cudaFree behave with the (err, *outs) convention, and its error enums compare equal across the cuda/hip enum classes so cuda_python_error_check works unchanged). The only blockers were test-infra: TEST_CUDA_PYTHON_BINDINGS required torch.version.cuda, and the four TestMemPool UVM tests carried a blanket "UVM is not supported on ROCm" skip.
TEST_CUDA_PYTHON_BINDINGS now accepts cuda.bindings on ROCm only when it is the HIP-backed flavor, detected via hip-python's HIP_PYTHON marker attribute (import-safe, no driver call; NVIDIA's cuda-bindings on a ROCm box remains excluded). The four skipIfRocm lines are removed; on machines without hip-python-interop, including current ROCm CI images, the tests skip via requires_cuda_python_bindings exactly as CUDA machines without cuda-python do. Adding hip-python/hip-python-interop to the ROCm CI images is deferred until wheels matching the CI ROCm version publish (ROCm 10.0.0 bindings are in progress per ROCm/hip-python#109).
Test plan, on gfx950 (ROCm 10.0, HIP 7.15): with hip-python-interop + hip-python installed, all four test_use_uvm* tests pass, 25 consecutive runs clean; without them, all four skip with "requires cuda-python (cuda.bindings)".
This PR was authored with the assistance of an AI agent (Claude). The analysis and test results were verified by hand.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/196492
Approved by: https://github.com/jeffdaily
Co-authored-by: Jeff Daily <jeff.daily@amd.com>