Skip to content

GitLab

Unverified Commit bcb90f2a authored by Jiashi Li's avatar Jiashi Li Committed by GitHub
Browse files

Merge pull request #9 from homorunner/main 

support Windows build
parents dd1161e3 e62bdb4d
Hide whitespace changes
Inline Side-by-side
Showing with 8 additions and 1 deletion
setup.py
View file @ bcb90f2a
...@@ -8,6 +8,7 @@ from setuptools import setup, find_packages ...@@ -8,6 +8,7 @@ from setuptools import setup, find_packages
from torch.utils.cpp_extension import ( from torch.utils.cpp_extension import (
BuildExtension, BuildExtension,
CUDAExtension, CUDAExtension,
IS_WINDOWS,
) )
...@@ -24,6 +25,11 @@ cc_flag.append("arch=compute_90a,code=sm_90a") ...@@ -24,6 +25,11 @@ cc_flag.append("arch=compute_90a,code=sm_90a")
this_dir = os.path.dirname(os.path.abspath(__file__)) this_dir = os.path.dirname(os.path.abspath(__file__))
if IS_WINDOWS:
cxx_args = ["/O2", "/std:c++17", "/DNDEBUG", "/W0"]
else:
cxx_args = ["-O3", "-std=c++17", "-DNDEBUG", "-Wno-deprecated-declarations"]
ext_modules = [] ext_modules = []
ext_modules.append( ext_modules.append(
CUDAExtension( CUDAExtension(
...@@ -33,12 +39,13 @@ ext_modules.append( ...@@ -33,12 +39,13 @@ ext_modules.append(
"csrc/flash_fwd_mla_bf16_sm90.cu", "csrc/flash_fwd_mla_bf16_sm90.cu",
], ],
extra_compile_args={ extra_compile_args={
"cxx": ["-O3", "-std=c++17", "-DNDEBUG", "-Wno-deprecated-declarations"], "cxx": cxx_args,
"nvcc": append_nvcc_threads( "nvcc": append_nvcc_threads(
[ [
"-O3", "-O3",
"-std=c++17", "-std=c++17",
"-DNDEBUG", "-DNDEBUG",
"-D_USE_MATH_DEFINES",
"-Wno-deprecated-declarations", "-Wno-deprecated-declarations",
"-U__CUDA_NO_HALF_OPERATORS__", "-U__CUDA_NO_HALF_OPERATORS__",
"-U__CUDA_NO_HALF_CONVERSIONS__", "-U__CUDA_NO_HALF_CONVERSIONS__",
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment