Enhance CUDA architecture support in setup.py by allowing user-defined architectures via environment variable. Refactor GPU capability checks and streamline NVCC flags for SM89 and SM90 extensions. Improve build process by creating separate output directories for extensions.

Author: GACLove

setup.py

@@ -66,15 +66,22 @@ def get_nvcc_cuda_version(cuda_dir: str) -> Version:
     nvcc_cuda_version = parse(output[release_idx].split(",")[0])
     return nvcc_cuda_version
 
-# Iterate over all GPUs on the current machine. Also you can modify this part to specify the architecture if you want to build for specific GPU architectures.
 compute_capabilities = set()
-device_count = torch.cuda.device_count()
-for i in range(device_count):
-    major, minor = torch.cuda.get_device_capability(i)
-    if major < 8:
-        warnings.warn(f"skipping GPU {i} with compute capability {major}.{minor}")
-        continue
-    compute_capabilities.add(f"{major}.{minor}")
+cuda_architectures = os.environ.get("CUDA_ARCHITECTURES")
+if cuda_architectures is not None:
+    for arch in cuda_architectures.split(","):
+        arch = arch.strip()
+        if arch:
+            compute_capabilities.add(arch)
+else:
+    #Iterate over all GPUs on the current machine. Also you can modify this part to specify the architecture if you want to build for specific GPU architectures.
+    device_count = torch.cuda.device_count()
+    for i in range(device_count):
+        major, minor = torch.cuda.get_device_capability(i)
+        if major < 8:
+            warnings.warn(f"skipping GPU {i} with compute capability {major}.{minor}")
+            continue
+        compute_capabilities.add(f"{major}.{minor}")
 
 nvcc_cuda_version = get_nvcc_cuda_version(CUDA_HOME)
 if not compute_capabilities:
@@ -119,54 +126,96 @@ def get_nvcc_cuda_version(cuda_dir: str) -> Version:
 ext_modules = []
 
 if HAS_SM80 or HAS_SM86 or HAS_SM89 or HAS_SM90 or HAS_SM120:
-    qattn_extension = CUDAExtension(
+    sm80_sources = [
+        "csrc/qattn/pybind_sm80.cpp",
+        "csrc/qattn/qk_int_sv_f16_cuda_sm80.cu",
+    ]
+
+    qattn_extension_sm80 = CUDAExtension(
         name="sageattention._qattn_sm80",
-        sources=[
-            "csrc/qattn/pybind_sm80.cpp",
-            "csrc/qattn/qk_int_sv_f16_cuda_sm80.cu",
-        ],
+        sources=sm80_sources,
         extra_compile_args={
             "cxx": CXX_FLAGS,
             "nvcc": NVCC_FLAGS,
         },
     )
-    ext_modules.append(qattn_extension)
+    ext_modules.append(qattn_extension_sm80)
 
 if HAS_SM89 or HAS_SM120:
-    qattn_extension = CUDAExtension(
+    sm89_sources = [
+        "csrc/qattn/pybind_sm89.cpp",
+        "csrc/qattn/sm89_qk_int8_sv_f8_accum_f32_attn_inst_buf.cu",
+        "csrc/qattn/sm89_qk_int8_sv_f8_accum_f16_attn_inst_buf.cu",
+        "csrc/qattn/sm89_qk_int8_sv_f8_accum_f32_attn.cu",
+        "csrc/qattn/sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn.cu",
+        "csrc/qattn/sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn.cu",
+        "csrc/qattn/sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf.cu",
+        "csrc/qattn/sm89_qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf.cu"
+        #"csrc/qattn/qk_int_sv_f8_cuda_sm89.cu",
+    ]
+    
+    sm89_nvcc_flags = [flag for flag in NVCC_FLAGS]
+    
+    filtered_flags = []
+    skip_next = False
+    for i, flag in enumerate(sm89_nvcc_flags):
+        if skip_next:
+            skip_next = False
+            continue
+        if flag == "-gencode":
+            if i + 1 < len(sm89_nvcc_flags):
+                arch_flag = sm89_nvcc_flags[i + 1]
+                if "compute_89" in arch_flag or "compute_90" in arch_flag or "compute_120" in arch_flag:
+                    filtered_flags.append(flag)
+                    filtered_flags.append(arch_flag)
+                skip_next = True
+        elif flag not in ["-gencode"]:
+            filtered_flags.append(flag)
+    
+    qattn_extension_sm89 = CUDAExtension(
         name="sageattention._qattn_sm89",
-        sources=[
-            "csrc/qattn/pybind_sm89.cpp",
-            "csrc/qattn/sm89_qk_int8_sv_f8_accum_f32_attn_inst_buf.cu",
-            "csrc/qattn/sm89_qk_int8_sv_f8_accum_f16_attn_inst_buf.cu",
-            "csrc/qattn/sm89_qk_int8_sv_f8_accum_f32_attn.cu",
-            "csrc/qattn/sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn.cu",
-            "csrc/qattn/sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn.cu",
-            "csrc/qattn/sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf.cu",
-            "csrc/qattn/sm89_qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf.cu"
-            #"csrc/qattn/qk_int_sv_f8_cuda_sm89.cu",
-        ],
+        sources=sm89_sources,
         extra_compile_args={
             "cxx": CXX_FLAGS,
-            "nvcc": NVCC_FLAGS,
+            "nvcc": filtered_flags if filtered_flags else NVCC_FLAGS,
         },
     )
-    ext_modules.append(qattn_extension)
+    ext_modules.append(qattn_extension_sm89)
 
 if HAS_SM90:
-    qattn_extension = CUDAExtension(
+    sm90_sources = [
+        "csrc/qattn/pybind_sm90.cpp",
+        "csrc/qattn/qk_int_sv_f8_cuda_sm90.cu",
+    ]
+    
+    sm90_nvcc_flags = [flag for flag in NVCC_FLAGS]
+    
+    filtered_flags = []
+    skip_next = False
+    for i, flag in enumerate(sm90_nvcc_flags):
+        if skip_next:
+            skip_next = False
+            continue
+        if flag == "-gencode":
+            if i + 1 < len(sm90_nvcc_flags):
+                arch_flag = sm90_nvcc_flags[i + 1]
+                if "compute_90" in arch_flag or "compute_120" in arch_flag:
+                    filtered_flags.append(flag)
+                    filtered_flags.append(arch_flag)
+                skip_next = True
+        elif flag not in ["-gencode"]:
+            filtered_flags.append(flag)
+    
+    qattn_extension_sm90 = CUDAExtension(
         name="sageattention._qattn_sm90",
-        sources=[
-            "csrc/qattn/pybind_sm90.cpp",
-            "csrc/qattn/qk_int_sv_f8_cuda_sm90.cu",
-        ],
+        sources=sm90_sources,
         extra_compile_args={
             "cxx": CXX_FLAGS,
-            "nvcc": NVCC_FLAGS,
+            "nvcc": filtered_flags if filtered_flags else NVCC_FLAGS,
         },
         extra_link_args=['-lcuda'],
     )
-    ext_modules.append(qattn_extension)
+    ext_modules.append(qattn_extension_sm90)
 
 # Fused kernels.
 fused_extension = CUDAExtension(
@@ -208,15 +257,23 @@ def compile_new(*args, **kwargs):
                         **kwargs,
                         "output_dir": os.path.join(
                             kwargs["output_dir"],
-                            self.thread_ext_name_map[threading.current_thread().ident]),
+                            self.thread_ext_name_map.get(threading.current_thread().ident, "default")),
                     })
                 self.compiler.compile = compile_new
                 self.compiler._compile_separate_output_dir = True
         self.thread_ext_name_map[threading.current_thread().ident] = ext.name
-        objects = super().build_extension(ext)
+        
+        original_build_temp = self.build_temp
+        self.build_temp = os.path.join(original_build_temp, ext.name.replace(".", "_"))
+        os.makedirs(self.build_temp, exist_ok=True)
+        
+        try:
+            objects = super().build_extension(ext)
+        finally:
+            self.build_temp = original_build_temp
+        
         return objects
 
-
 setup(
     name='sageattention', 
     version='2.2.0',