DeepSeek/FlashMLA
1
0
Fork 0
mirror of https://github.com/deepseek-ai/FlashMLA synced 2025-04-23 15:44:07 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Merge pull request #9 from homorunner/main

Browse Source 
support Windows build
This commit is contained in:
Jiashi Li 2025-02-24 13:21:58 +08:00 committed by GitHub
commit bcb90f2afd
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 8 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
setup.py
View File

@ -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")
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(
"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__",