From 12588047fd21917bf5ecf46bd677061a79e5c61e Mon Sep 17 00:00:00 2001 From: biondizzle Date: Wed, 13 May 2026 23:17:30 +0000 Subject: [PATCH] Fix setup.py: use include_dirs and extra_compile_args (correct PyTorch extension API) --- .../cutlass_nvfp4_gemm/setup.py | 62 +++++++++---------- 1 file changed, 28 insertions(+), 34 deletions(-) diff --git a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/setup.py b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/setup.py index f82528e0..d68fd16b 100644 --- a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/setup.py +++ b/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={