Skip to content

Commit b0b5347

Browse files
[Bugfix] Add NVIDIA HPC SDK support in CUDA detection (#974) (#976)
* [Bugfix] Add NVIDIA HPC SDK support in CUDA detection (#974) Enhanced CUDA detection to recognize NVIDIA HPC SDK installations: - Added path check for nvhpc in nvcc binary path - Added fallback scan for default nvhpc paths: /opt/nvidia/hpc_sdk/Linux_x86_64 - Maintained backward compatibility with standard CUDA installations Verification: - Tested on Ubuntu 24.04 with NVIDIA HPC SDK 25.7 - Confirmed detection works without manual CUDA_HOME or CUDA_PATH setting Fixes #974 * [Bugfix] Fix CUDA home detection logic * [Bugfix] Safely handle None cuda_home during CUDA detection Adds a check for None before validating the CUDA home path to prevent errors when the path is not set. * [Bugfix] Fix CUDA detection edge cases in nvhpc support (#974) - Improved nvhpc path detection logic - Added None check for cuda_home to avoid crashes - Maintained existing CUDA installation compatibility Fixes #974 * chore: rerun CI --------- Co-authored-by: NaNExist <138002947+NaNExist@users.noreply.github.com>
1 parent 0550703 commit b0b5347

File tree

1 file changed

+21
-4
lines changed

1 file changed

+21
-4
lines changed

tilelang/env.py

Lines changed: 21 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -30,17 +30,34 @@ def _find_cuda_home() -> str:
3030
if cuda_home is None:
3131
# Guess #2
3232
nvcc_path = shutil.which("nvcc")
33-
if nvcc_path is not None and "cuda" in nvcc_path.lower():
34-
cuda_home = os.path.dirname(os.path.dirname(nvcc_path))
33+
if nvcc_path is not None:
34+
# Standard CUDA pattern
35+
if "cuda" in nvcc_path.lower():
36+
cuda_home = os.path.dirname(os.path.dirname(nvcc_path))
37+
# NVIDIA HPC SDK pattern
38+
elif "hpc_sdk" in nvcc_path.lower():
39+
# Navigate to the root directory of nvhpc
40+
cuda_home = os.path.dirname(os.path.dirname(os.path.dirname(nvcc_path)))
41+
# Generic fallback for non-standard or symlinked installs
42+
else:
43+
cuda_home = os.path.dirname(os.path.dirname(nvcc_path))
44+
3545
else:
3646
# Guess #3
3747
if sys.platform == 'win32':
3848
cuda_homes = glob.glob('C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v*.*')
3949
cuda_home = '' if len(cuda_homes) == 0 else cuda_homes[0]
4050
else:
41-
cuda_home = '/usr/local/cuda'
42-
if not os.path.exists(cuda_home):
51+
# Linux/macOS
52+
if os.path.exists('/usr/local/cuda'):
53+
cuda_home = '/usr/local/cuda'
54+
elif os.path.exists('/opt/nvidia/hpc_sdk/Linux_x86_64'):
55+
cuda_home = '/opt/nvidia/hpc_sdk/Linux_x86_64'
56+
57+
# Validate found path
58+
if cuda_home is None or not os.path.exists(cuda_home):
4359
cuda_home = None
60+
4461
return cuda_home if cuda_home is not None else ""
4562

4663

0 commit comments

Comments
 (0)