Fix cc

setup.py

@@ -89,7 +89,9 @@ def get_generator_flag():
     return generator_flag
 
 
-def get_compute_capabilities():
+def get_compute_capabilities(
+    compute_capabilities={75, 80, 86, 89, 90}
+):
     capability_flags = []
 
     if CUDA_VERSION:
@@ -103,7 +105,6 @@ def get_compute_capabilities():
                 )
 
         # Figure out compute capability
-        compute_capabilities = {75, 80, 86, 89, 90}
         for cap in compute_capabilities:
             capability_flags += ["-gencode", f"arch=compute_{cap},code=sm_{cap}"]
 
@@ -175,10 +176,23 @@ def get_extra_link_args():
                 "awq_ext/vllm/moe_alig_block.cu",
                 "awq_ext/vllm/activation.cu",
                 "awq_ext/vllm/topk_softmax_kernels.cu",
+            ],
+            extra_compile_args=extra_compile_args,
+        )
+    )
+
+    # only compatible with ampere
+    arch_flags = get_compute_capabilities({80, 86, 89, 90})
+    extra_compile_args_v2 = get_extra_compile_args(arch_flags, generator_flags)
+
+    extensions.append(
+        CUDAExtension(
+            "awq_v2_ext",
+            [
                 "awq_ext/quantization_new/gemv/gemv_cuda.cu",
                 "awq_ext/quantization_new/gemm/gemm_cuda.cu",
             ],
-            extra_compile_args=extra_compile_args,
+            extra_compile_args=extra_compile_args_v2,
         )
     )