Comprehensive fix for Windows MSVC build errors (C2872 std, LNK2001) and thread-safety

README.md

@@ -1,4 +1,15 @@
-# SageAttention
+# SageAttention (Windows Blackwell Fork)
+
+> **Windows Build Fixes**: This fork includes critical fixes for compiling SageAttention on Windows with MSVC. Specifically, it resolves the `C2872: 'std': ambiguous symbol` error by appending `/Zc:preprocessor`, `/DCCCL_IGNORE_MSVC_TRADITIONAL_PREPROCESSOR_WARNING`, `-DUSE_CUDA`, and `-Usmall` to the compilation flags in `setup.py` and `sageattention3_blackwell/setup.py`. It also addresses thread-safety issues during concurrent builds by creating shallow copies of compiler flag lists (`CXX_FLAGS[:]`), preventing `TORCH_EXTENSION_NAME` from being inadvertently overwritten.
+> 
+> ### Pre-built Wheel Requirements
+> The `.whl` files provided in this repository are compiled specifically for the following environment:
+> - **Python Version:** 3.12
+> - **PyTorch Version:** 2.10.0
+> - **CUDA Version:** 13.1 (compatible with cu130/cu131)
+> - **OS:** Windows 10 / Windows 11
+> - **Supported Architectures:** RTX 3000 / 4000 / 5000 Series GPUs (Compute Capability: `sm_80`, `sm_89`, `sm_120a`)
+
 <!-- We are continuously updating more features. You could **Star** and **Watch** our repository to stay updated.
 
 --- -->

sageattention3_blackwell/sageattn3/blackwell/api.cu

@@ -18,6 +18,11 @@
 #include <torch/python.h>
 #include <torch/nn/functional.h>
 #include <ATen/cuda/CUDAContext.h>
+
+#ifdef _WIN32
+#undef small
+#endif
+
 #include <c10/cuda/CUDAGuard.h>
 
 #include <cutlass/numeric_types.h>

sageattention3_blackwell/sageattn3/quantization/fp4_quantization_4d.cu

@@ -17,6 +17,11 @@
 #include <torch/python.h>
 #include <torch/nn/functional.h>
 #include <ATen/cuda/CUDAContext.h>
+
+#ifdef _WIN32
+#undef small
+#endif
+
 #include <c10/cuda/CUDAGuard.h>
 #include <cuda_runtime_api.h>
 #include <cuda_runtime.h>

sageattention3_blackwell/setup.py

@@ -108,6 +108,27 @@ def append_nvcc_threads(nvcc_extra_args):
         "-DCTA256",
         "-DDQINRMEM",
     ]
+
+    cxx_flags = ["-O3", "-std=c++17"]
+
+    import sys
+    if sys.platform == "win32":
+        win_flags = [
+            "/Zc:preprocessor", 
+            "/DCCCL_IGNORE_MSVC_TRADITIONAL_PREPROCESSOR_WARNING",
+            "-DUSE_CUDA", 
+            "-Usmall"
+        ]
+        cxx_flags.extend(win_flags)
+
+        nvcc_win_flags = [
+            "-Xcompiler", "/Zc:preprocessor",
+            "-Xcompiler", "/DCCCL_IGNORE_MSVC_TRADITIONAL_PREPROCESSOR_WARNING",
+            "-DUSE_CUDA", 
+            "-Usmall"
+        ]
+        nvcc_flags.extend(nvcc_win_flags)
+
     include_dirs = [
         repo_dir / "sageattn3",
         cutlass_dir / "include",
@@ -119,7 +140,7 @@ def append_nvcc_threads(nvcc_extra_args):
             name="fp4attn_cuda",
             sources=["sageattn3/blackwell/api.cu"],
             extra_compile_args={
-                "cxx": ["-O3", "-std=c++17"],
+                "cxx": cxx_flags[:],
                 "nvcc": append_nvcc_threads(
                     nvcc_flags + ["-DEXECMODE=0"] + cc_flag
                 ),
@@ -134,7 +155,7 @@ def append_nvcc_threads(nvcc_extra_args):
             name="fp4quant_cuda",
             sources=["sageattn3/quantization/fp4_quantization_4d.cu"],
             extra_compile_args={
-                "cxx": ["-O3", "-std=c++17"],
+                "cxx": cxx_flags[:],
                 "nvcc": append_nvcc_threads(
                     nvcc_flags + ["-DEXECMODE=0"] + cc_flag
                 ),

setup.py

@@ -60,6 +60,21 @@
         "-diag-suppress=174",
     ]
 
+    # Windows specific flags
+    if sys.platform == "win32":
+        CXX_FLAGS += [
+            "/Zc:preprocessor", 
+            "/DCCCL_IGNORE_MSVC_TRADITIONAL_PREPROCESSOR_WARNING",
+            "-DUSE_CUDA", 
+            "-Usmall"
+        ]
+        NVCC_FLAGS += [
+            "-Xcompiler", "/Zc:preprocessor",
+            "-Xcompiler", "/DCCCL_IGNORE_MSVC_TRADITIONAL_PREPROCESSOR_WARNING",
+            "-DUSE_CUDA", 
+            "-Usmall"
+        ]
+
     # Append flags from env if provided
     cxx_append = os.getenv("CXX_APPEND_FLAGS", "").strip()
     if cxx_append:
@@ -179,7 +194,7 @@ def get_nvcc_cuda_version(cuda_dir: str) -> Version:
                     "csrc/qattn/pybind_sm80.cpp",
                     "csrc/qattn/qk_int_sv_f16_cuda_sm80.cu",
                 ],
-                extra_compile_args={"cxx": CXX_FLAGS, "nvcc": NVCC_FLAGS},
+                extra_compile_args={"cxx": CXX_FLAGS[:], "nvcc": NVCC_FLAGS[:]},
             )
         )
 
@@ -197,7 +212,7 @@ def get_nvcc_cuda_version(cuda_dir: str) -> Version:
                     "csrc/qattn/sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf.cu",
                     "csrc/qattn/sm89_qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf.cu",
                 ],
-                extra_compile_args={"cxx": CXX_FLAGS, "nvcc": NVCC_FLAGS},
+                extra_compile_args={"cxx": CXX_FLAGS[:], "nvcc": NVCC_FLAGS[:]},
             )
         )
 
@@ -209,7 +224,7 @@ def get_nvcc_cuda_version(cuda_dir: str) -> Version:
                     "csrc/qattn/pybind_sm90.cpp",
                     "csrc/qattn/qk_int_sv_f8_cuda_sm90.cu",
                 ],
-                extra_compile_args={"cxx": CXX_FLAGS, "nvcc": NVCC_FLAGS},
+                extra_compile_args={"cxx": CXX_FLAGS[:], "nvcc": NVCC_FLAGS[:]},
                 extra_link_args=['-lcuda'],
             )
         )
@@ -218,7 +233,7 @@ def get_nvcc_cuda_version(cuda_dir: str) -> Version:
         CUDAExtension(
             name="sageattention._fused",
             sources=["csrc/fused/pybind.cpp", "csrc/fused/fused.cu"],
-            extra_compile_args={"cxx": CXX_FLAGS, "nvcc": NVCC_FLAGS},
+            extra_compile_args={"cxx": CXX_FLAGS[:], "nvcc": NVCC_FLAGS[:]},
         )
     )