fix compile group norm [v1/v2] (#1927)

AlongWY · web-flow · commit 5c8860f1dc79 · 2025-10-03T15:54:42.000-07:00
* fix group norm

* merge upstream code
diff --git a/setup.py b/setup.py
@@ -512,31 +512,36 @@ def check_cudnn_version_and_warn(global_option: str, required_cudnn_version: int
     )
 
     # CUDA group norm V2 is tested on SM100
-    if bare_metal_version >= Version("12.8"):
-        arch_flags = ["-gencode=arch=compute_100,code=sm_100"]
-    else:
-        arch_flags = ["-gencode=arch=compute_90,code=compute_90"]
+    if bare_metal_version >= Version("12.4"):
+        if bare_metal_version >= Version("12.8"):
+            arch_flags = [
+                "-gencode=arch=compute_90,code=sm_90",
+                "-gencode=arch=compute_100,code=sm_100",
+                "-gencode=arch=compute_120,code=compute_120",
+            ]
+        else:
+            arch_flags = ["-gencode=arch=compute_90,code=compute_90"]
 
-    ext_modules.append(
-        CUDAExtension(
-            name="group_norm_v2_cuda",
-            sources=[
-                "apex/contrib/csrc/group_norm_v2/gn.cpp",
-                "apex/contrib/csrc/group_norm_v2/gn_cuda.cu",
-                "apex/contrib/csrc/group_norm_v2/gn_utils.cpp",
-            ] + glob.glob("apex/contrib/csrc/group_norm_v2/gn_cuda_inst_*.cu"),
-            extra_compile_args={
-                "cxx": ["-O2"],
-                "nvcc": [
-                    "-O2", "--use_fast_math", "--ftz=false",
-                    "-U__CUDA_NO_HALF_CONVERSIONS__",
-                    "-U__CUDA_NO_HALF_OPERATORS__",
-                    "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
-                    "-U__CUDA_NO_BFLOAT16_OPERATORS__",
-                ] + arch_flags,
-            },
+        ext_modules.append(
+            CUDAExtension(
+                name="group_norm_v2_cuda",
+                sources=[
+                    "apex/contrib/csrc/group_norm_v2/gn.cpp",
+                    "apex/contrib/csrc/group_norm_v2/gn_cuda.cu",
+                    "apex/contrib/csrc/group_norm_v2/gn_utils.cpp",
+                ] + glob.glob("apex/contrib/csrc/group_norm_v2/gn_cuda_inst_*.cu"),
+                extra_compile_args={
+                    "cxx": ["-O2"],
+                    "nvcc": [
+                        "-O2", "--use_fast_math", "--ftz=false",
+                        "-U__CUDA_NO_HALF_CONVERSIONS__",
+                        "-U__CUDA_NO_HALF_OPERATORS__",
+                        "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
+                        "-U__CUDA_NO_BFLOAT16_OPERATORS__",
+                    ] + arch_flags,
+                },
+            )
         )
-    )
 
 if has_flag("--index_mul_2d", "APEX_INDEX_MUL_2D"):
     if "--index_mul_2d" in sys.argv:
