Fix setup.py: use include_dirs and extra_compile_args (correct PyTorch extension API)

src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/setup.py

@@ -1,20 +1,8 @@
 """
 Setup script for CUTLASS NVFP4 block-scaled GEMM PyTorch extension.
-
-Build on the B200 server:
-    python setup.py install
-
-Or build in-place:
-    python setup.py build_ext --inplace
-
-Requires:
-    - CUDA 13.0+ (Blackwell SM100)
-    - CUTLASS headers at CUTLASS_INCLUDE_DIR
-    - PyTorch with CUDA 13.0 support
 """
 
 import os
-import glob
 from setuptools import setup
 from torch.utils.cpp_extension import BuildExtension, CUDAExtension
 
@@ -23,29 +11,21 @@ CUTLASS_INCLUDE_DIR = os.environ.get(
     "CUTLASS_INCLUDE_DIR",
     "/usr/local/lib/python3.12/dist-packages/tilelang/3rdparty/cutlass/include"
 )
+if not os.path.exists(os.path.join(CUTLASS_INCLUDE_DIR, "cutlass", "cutlass.h")):
+    for alt in [
+        "/usr/local/lib/python3.12/dist-packages/tilelang/3rdparty/cutlass/include",
+        "/usr/local/include/cutlass",
+        "/opt/cutlass/include",
+    ]:
+        if os.path.exists(os.path.join(alt, "cutlass", "cutlass.h")):
+            CUTLASS_INCLUDE_DIR = alt
+            break
 
-# Cutlass tools/util include (for some helpers)
 CUTLASS_UTIL_INCLUDE = os.path.join(os.path.dirname(CUTLASS_INCLUDE_DIR), "tools", "util", "include")
 
-extra_include_paths = [CUTLASS_INCLUDE_DIR]
+include_dirs = [CUTLASS_INCLUDE_DIR]
 if os.path.exists(CUTLASS_UTIL_INCLUDE):
-    extra_include_paths.append(CUTLASS_UTIL_INCLUDE)
-
-extra_cuda_cflags = [
-    "-gencode=arch=compute_100a,code=sm_100a",
-    "--expt-relaxed-constexpr",
-    "-DCUTLASS_ENABLE_GEMP_OPERATION=1",
-    "-DCUTLASS_ARCH_SM100_ENABLED=1",
-    "--ptxas-options=-v",
-    "--ptxas-options=-allow-expensive-optimizations=true",
-]
-
-extra_cflags = [
-    "-O3",
-    "-std=c++17",
-    "-DCUTLASS_ENABLE_GEMP_OPERATION=1",
-    "-DCUTLASS_ARCH_SM100_ENABLED=1",
-]
+    include_dirs.append(CUTLASS_UTIL_INCLUDE)
 
 setup(
     name="cutlass_nvfp4_gemm",
@@ -56,9 +36,23 @@ setup(
                 "pytorch_binding.cpp",
                 "cutlass_nvfp4_gemm.cu",
             ],
-            extra_include_paths=extra_include_paths,
-            extra_cuda_cflags=extra_cuda_cflags,
-            extra_cflags=extra_cflags,
+            include_dirs=include_dirs,
+            extra_compile_args={
+                "cxx": [
+                    "-O3",
+                    "-std=c++17",
+                    "-DCUTLASS_ENABLE_GEMP_OPERATION=1",
+                    "-DCUTLASS_ARCH_SM100_ENABLED=1",
+                ],
+                "nvcc": [
+                    "-gencode=arch=compute_100a,code=sm_100a",
+                    "--expt-relaxed-constexpr",
+                    "-DCUTLASS_ENABLE_GEMP_OPERATION=1",
+                    "-DCUTLASS_ARCH_SM100_ENABLED=1",
+                    "--ptxas-options=-v",
+                    "--ptxas-options=-allow-expensive-optimizations=true",
+                ],
+            },
         ),
     ],
     cmdclass={