From 9127f4622f463f69b2d740bfeafbf23432bd9cb8 Mon Sep 17 00:00:00 2001 From: Ryan Guo Date: Thu, 21 Aug 2025 14:05:03 -0700 Subject: [PATCH] Fix windows build for pytorch 2.9 Fixes #242. I'm no expert on this and used Gemini/Chatgpt to figure this out -- basically, the latest PyTorch introduced some more aggressive checks on build flags, which ended up discovering the pre-existing issue in sage attention's `setup.py`. This patch fixes that. --- setup.py | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/setup.py b/setup.py index 5e4779dd..cb838351 100644 --- a/setup.py +++ b/setup.py @@ -46,6 +46,36 @@ "-diag-suppress=174", # suppress the specific warning ] +# Platform-specific compiler flags. +if os.name == 'nt': # Windows + CXX_FLAGS = ["/Zi", "/O2", "/openmp", "/std:c++17", "/DENABLE_BF16", "/MD"] + NVCC_FLAGS = [ + "-O3", + "-std=c++17", + "--use-local-env", # specific to windows NVCC builds + "-U__CUDA_NO_HALF_OPERATORS__", + "-U__CUDA_NO_HALF_CONVERSIONS__", + "--use_fast_math", + "--threads=8", + "-Xptxas=-v", + "-diag-suppress=174", + # NVCC forwards flags to MSVC differently (-Xcompiler is needed). + "-Xcompiler", "/openmp", + "-Xcompiler", "/std:c++17", + ] +else: # Linux, macOS, etc. + CXX_FLAGS = ["-g", "-O3", "-fopenmp", "-lgomp", "-std=c++17", "-DENABLE_BF16"] + NVCC_FLAGS = [ + "-O3", + "-std=c++17", + "-U__CUDA_NO_HALF_OPERATORS__", + "-U__CUDA_NO_HALF_CONVERSIONS__", + "--use_fast_math", + "--threads=8", + "-Xptxas=-v", + "-diag-suppress=174", + ] + ABI = 1 if torch._C._GLIBCXX_USE_CXX11_ABI else 0 CXX_FLAGS += [f"-D_GLIBCXX_USE_CXX11_ABI={ABI}"] NVCC_FLAGS += [f"-D_GLIBCXX_USE_CXX11_ABI={ABI}"]