[FIX] Resolve JIT GPU backend by explicit signals, not stray /opt/rocm#652
Conversation
The JIT extension backend selector decided HIP-vs-CUDA solely by probing for a ROCm home, whose default guess succeeds whenever /opt/rocm merely exists. On a CUDA/NVIDIA host that carries a stray /opt/rocm directory this misclassified the backend as HIP even when nvcc was present and hipcc was not, breaking CUDA source builds. Replace that heuristic with an ordered, evidence-based resolver: 1. honor the explicit TVM_FFI_GPU_BACKEND override; 2. use PyTorch build signals (torch.version.hip / torch.version.cuda) when PyTorch is importable and conclusive; 3. fall back to an available nvcc (CUDA); 4. then an available hipcc (HIP); 5. otherwise raise a clear error naming the override and required compilers. Backend selection no longer consults the ROCm home directory, so a stray /opt/rocm can never force HIP. Toolkit-home, compiler, include-path, flag, and architecture discovery are unchanged and run only after the backend is resolved, so CUDA and HIP state never mix.
Inline the PyTorch build-signal check directly in _detect_gpu_backend and drop the separate helper. When no override, Torch signal, or compiler evidence resolves the backend, default to CUDA instead of raising, and select HIP only from an available hipcc without nvcc. A stray /opt/rocm still never forces HIP.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the GPU backend auto-detection logic in _detect_gpu_backend to resolve the backend using PyTorch build signals and compiler presence (nvcc and hipcc) rather than relying on the existence of /opt/rocm. The review feedback suggests improving the robustness of the detection by checking for nvcc and hipcc in directories specified by CUDA_HOME/CUDA_PATH and ROCM_HOME/ROCM_PATH when they are not in the system PATH, and safely accessing torch.version to avoid potential AttributeErrors.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
Resolve nvcc/hipcc from PATH and from the CUDA_HOME/CUDA_PATH and ROCM_HOME/ROCM_PATH toolkit homes, so a stray hipcc on PATH no longer wins when an env-specified CUDA install is present but nvcc is off PATH. Guard the torch.version lookup with getattr so a minimal or mocked torch cannot raise.
2 participants
Fixes #651.
The JIT extension backend selector decided CUDA-vs-HIP by probing for a ROCm
home, whose fallback returns
/opt/rocmwhenever that directory merely exists.A CUDA/NVIDIA host carrying a stray
/opt/rocmwas therefore misclassified asHIP even with
nvccpresent and nohipcc._detect_gpu_backend()now resolves the backend from explicit signals, inorder: the
TVM_FFI_GPU_BACKENDoverride, then PyTorch build signals(
torch.version.hip/torch.version.cuda), then an availablehipccwhen nonvccis present. Detection defaults to CUDA when otherwise inconclusive, andthe existence of
/opt/rocmalone never forces HIP. Toolkit-home, compiler,include-path, flag, and architecture discovery are unchanged and run only after
the backend is resolved, so CUDA and HIP state never mix.