Prebuilt Windows wheel and the full build harness for SageAttention 2.2.0 on
RTX 50-series (Blackwell, sm_120a), built against CUDA 13.3 and a PyTorch 2.14
nightly, with an aggressively optimized toolchain (clang thinLTO host codegen +
lld-link, -Xptxas -O3 device codegen, sm_120a SASS with no PTX).
Measured 5.77–5.86× faster than torch.nn.functional.scaled_dot_product_attention
on an RTX 5070 at B=2 H=24 L=4096 D=128 fp16 — with cosine similarity ≥ 0.9992 against a
float64 reference.
The interesting part of this repo is probably not the wheel — it's the four
non-obvious toolchain blockers documented below. A naive pip install . on this
configuration fails at every one of them.
1. Will the prebuilt wheel work for you?
Python C extensions that link libtorch are tightly coupled to the exact torch build.
The wheel here matches:
If you get ImportError: DLL load failed — that is expected on a different torch
nightly. The C++ ABI drifts between nightlies. Build from source (§3); it takes a few
minutes and the harness is the point of this repo.
The wheel is compiled SASS-only for sm_120a (no embedded PTX), so it will not
JIT onto Hopper, Ada, or Ampere. Those need a rebuild with SAGE_ARCH changed.
Verify what you got:
python verify_sage.py # run it from OUTSIDE a SageAttention source checkout
2. Four things that break this build
Each was isolated with a minimal smoke compile rather than inferred. If you're building
any CUDA extension on CUDA 13 + a recent torch nightly, you will hit these too.
2.1 torch 2.14 requires C++20, SageAttention asks for C++17
torch/utils/cpp_extension.py hardcodes -std=c++20. SageAttention's setup.py passes
-std=c++17. When nvcc's EDG frontend gets C++17, it rejects torch's own headers:
c10/core/TensorImpl.h(2952): error: data member initializer is not allowed
c10/core/AutogradState.h(84): error: data member initializer is not allowed
Those are C++20 NSDMIs inside bitfields. The fix is to strip stale -std=//std: flags
from both command lines and force C++20.
2.2 CUDA 13.3's CCCL rejects MSVC's traditional preprocessor
Any .cu that includes torch headers dies with:
cccl/cuda/std/__cccl/preprocessor.h(23): fatal error C1189: #error:
MSVC/cl.exe with traditional preprocessor is used.
clang-cl /clang:-flto=thin writes objects beginning 42 43 c0 de (BC\xc0\xde) — LLVM
bitcode — not COFF (64 86). link.exe then fails:
pybind.obj : fatal error LNK1107: invalid or corrupt file: cannot read at 0xA21F4
lld-link.exe reads bitcode and COFF in a single link and honours /OPT:REF/OPT:ICF=10. Mixing clang bitcode host objects with nvcc COFF device objects links
cleanly and imports fine, because both are MSVC-ABI and both /MD.
2.4 nvcc refuses clang as a host compiler on Windows
So clang can only own the pure-C++ pybind glue; nvcc's host pass stays on cl.exe.
-march=znver3 and thinLTO therefore apply to the glue, not to .cu host code. This
costs essentially nothing — the kernels are device code, and ptxas gets -O3.
Bonus: /fp:fast makes infinity undefined behaviour under clang
clang's /fp:fast implies -ffinite-math-only, and it warns on torch's headers:
warning: use of infinity is undefined behavior due to the currently enabled
floating-point options [-Wnan-infinity-disabled]
Attention masking depends on real -inf. The build passes
-fno-finite-math-only -fhonor-infinities -fhonor-nans to keep inf/nan semantics while
retaining reassociation and reciprocal math. SageAttention's own csrc/ kernels contain
no inf/numeric_limits sentinels, so device-side --use_fast_math is safe.
3. Building from source
Requirements
Visual Studio 2022 — CUDA 13.x does not support VS 2026. Build Tools is enough.
CUDA Toolkit 13.x
LLVM/clang for Windows (releases) —
must include clang-cl.exe and lld-link.exe. Tested with 22.1.8.
PyTorch nightly matching your CUDA, plus ninja and packaging.
-dlcm=ca caches global loads in L1; --allow-expensive-optimizations lets ptxas spend
more compile time on scheduling and register allocation.
5. How the build harness works
Worth understanding before you adapt it, because the obvious approach silently fails.
pip install . runs setup.py in a child process. Patching
torch.utils.cpp_extension in the parent has no effect on the process that actually
compiles — flags set that way vanish without any error. So _sage_opt_patch.py is
injected via PYTHONPATH + sitecustomize.py, gated on SAGE_OPT_BUILD=1.
Use sitecustomize, not usercustomize: virtualenvs typically report
site.ENABLE_USER_SITE == False, so usercustomize is never imported.
Two independent seams need patching, because ninja does not link on Windows. The
generated build.ninja has an empty ldflags and no link rule at all — setuptools spawns
self.linker itself.
_write_ninja_file — fix the -std flags, append device/host flags, and swap the
compiler. Note torch bakes the literal string cl into the compile rule
(command = cl /showIncludes ...) and never expands $cxx, so rewriting the
cxx = variable alone does nothing.
Compiler.link / Compiler.initialize — repoint at lld-link and add the link
flags. This must be applied to bothsetuptools._distutils.compilers.C.msvcanddistutils.compilers.C.msvc; which one binds depends on SETUPTOOLS_USE_DISTUTILS.
extra_postargs is passed positionally, so the wrapper binds the real signature via
inspect.signature rather than assuming kwargs.
Do not set NVCC_PREPEND_FLAGS in the .bat. It bypasses these patches and silently
drops the optimizations.
Accuracy vs a float64 torch SDPA reference, 3 shapes × causal/non-causal:
backend
max_abs
rel_l2
cosine
sageattn (auto)
0.0099–0.189
0.0347–0.0385
≥ 0.99926
sageattn_qk_int8_pv_fp16_cuda
0.0026–0.044
0.0092–0.0112
≥ 0.99994
sageattn_qk_int8_pv_fp8_cuda
0.0113–0.189
0.0345–0.0381
≥ 0.99927
qk_int8_pv_fp16_cuda is consistently the most accurate — ~3.7× lower relative error
than the auto-selected backend. Worth pinning explicitly if quality matters.
Large-magnitude causal inputs produce no NaN and no Inf, confirming fast-math did not
poison masked positions.
With the wheel installed, optimized_attention resolves to attention_sage.
Two caveats worth knowing:
SAGE_ATTENTION_SUPPORTS_MASK is False for SageAttention 2.2.0, so ComfyUI falls
back to PyTorch attention on any masked call. Not all attention in a workflow is
actually accelerated.
The global flag is known to produce black or corrupted output with some models
(Qwen, Wan). For those, use KJNodes'
Patch Sage Attention node per workflow and select
sageattn_qk_int8_pv_fp16_cuda — also the most accurate backend per the table above.
A note on SageAttention 3 (FP4)
If you came here hoping to pair Sage3 with an NVFP4 model in ComfyUI, check your
ComfyUI version first. As of ComfyUI 0.28.0, comfy/ldm/modules/attention.py
registers a "sage3" backend, but:
get_attention_function has zero call sites — the registry is write-only
--use-sage-attention routes to Sage2 unconditionally
So sage3 is unreachable from stock ComfyUI regardless of whether sageattn3 is
installed; it needs a third-party node to set the override. Confirm with:
If the only hits are the definition and the one read site, building
sageattention3_blackwell (which also wants a ~600 MB CUTLASS clone) gets you nothing.
