Skip to content

viable/strict/1789084224

@jeffdaily jeffdaily tagged this 10 Sep 20:09
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>
Assets 2
Loading