From d5ca4495c0c068bc617102841b9322d378fac8ea Mon Sep 17 00:00:00 2001 From: guyueh1 <140554423+guyueh1@users.noreply.github.com> Date: Sun, 18 May 2025 18:43:48 -0700 Subject: [PATCH] Make `TORCH_CUDA_ARCH_LIST` as an environment variable (#167) * Add 10.0 to TORCH_CUDA_ARCH_LIST Signed-off-by: Guyue Huang * Revert csrc/CMakeLists change; in setup.py make TORCH_CUDA_ARCH_LIST configurable Signed-off-by: Guyue Huang --------- Signed-off-by: Guyue Huang --- setup.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/setup.py b/setup.py index f9df917..af8e4e6 100644 --- a/setup.py +++ b/setup.py @@ -10,7 +10,8 @@ if __name__ == '__main__': print(f'NVSHMEM directory: {nvshmem_dir}') # TODO: currently, we only support Hopper architecture, we may add Ampere support later - os.environ['TORCH_CUDA_ARCH_LIST'] = '9.0' + if os.getenv('TORCH_CUDA_ARCH_LIST', None) is None: + os.environ['TORCH_CUDA_ARCH_LIST'] = '9.0' cxx_flags = ['-O3', '-Wno-deprecated-declarations', '-Wno-unused-variable', '-Wno-sign-compare', '-Wno-reorder', '-Wno-attributes'] nvcc_flags = ['-O3', '-Xcompiler', '-O3', '-rdc=true', '--ptxas-options=--register-usage-level=10']