From 8131369ca42177a53568a5a9ab0fb8deec39aef1 Mon Sep 17 00:00:00 2001 From: Erwann Millon Date: Mon, 24 Mar 2025 19:41:40 +0000 Subject: [PATCH 1/2] expose TORCH_CUDA_ARCH_LIST --- setup.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/setup.py b/setup.py index dd1015d6..4f55f565 100644 --- a/setup.py +++ b/setup.py @@ -34,6 +34,7 @@ SUPPORTED_ARCHS = {"8.0", "8.6", "8.9", "9.0", "12.0"} # Compiler flags. +TORCH_CUDA_ARCH_LIST = os.environ.get("TORCH_CUDA_ARCH_LIST", "") CXX_FLAGS = ["-g", "-O3", "-fopenmp", "-lgomp", "-std=c++17", "-DENABLE_BF16"] NVCC_FLAGS = [ "-O3", @@ -76,6 +77,11 @@ def get_nvcc_cuda_version(cuda_dir: str) -> Version: continue compute_capabilities.add(f"{major}.{minor}") +if TORCH_CUDA_ARCH_LIST: + for arch in TORCH_CUDA_ARCH_LIST.split(","): + compute_capabilities.add(arch) + + nvcc_cuda_version = get_nvcc_cuda_version(CUDA_HOME) if not compute_capabilities: raise RuntimeError("No GPUs found. Please specify the target GPU architectures or build on a machine with GPUs.") From adda1972da605e948e8f4ac35eddae3a0e362588 Mon Sep 17 00:00:00 2001 From: Erwann Millon Date: Tue, 25 Mar 2025 00:01:34 +0000 Subject: [PATCH 2/2] override automatically detected cuda architectures if TORCH_CUDA_ARCH_LIST is passed --- setup.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/setup.py b/setup.py index 4f55f565..d3360850 100644 --- a/setup.py +++ b/setup.py @@ -78,9 +78,7 @@ def get_nvcc_cuda_version(cuda_dir: str) -> Version: compute_capabilities.add(f"{major}.{minor}") if TORCH_CUDA_ARCH_LIST: - for arch in TORCH_CUDA_ARCH_LIST.split(","): - compute_capabilities.add(arch) - + compute_capabilities = set(TORCH_CUDA_ARCH_LIST.split(",")) nvcc_cuda_version = get_nvcc_cuda_version(CUDA_HOME) if not compute_capabilities: