SageAttention 2.2.0 — Blackwell (sm_120a) wheel + reproducible optimized build

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:

pip install sageattention-2.2.0-cp313-cp313-win_amd64.whl

sha256: b82441d7e8beb2b6c0fac69e6f6967c89d000785855d9db19f85e2ba7bdde184

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.

nvcc's host pass needs -Xcompiler /Zc:preprocessor.

2.3 clang thinLTO emits LLVM bitcode, which MSVC's linker cannot read

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

$ nvcc -ccbin clang-cl.exe ...
nvcc fatal : Host compiler targets unsupported OS.

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.

pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu130
pip install ninja packaging

Steps

git clone https://github.com/thu-ml/SageAttention
cd SageAttention
# copy this repo's build/ files into the checkout root

Then either edit the four paths at the top of build_optimized.bat, or set them as environment variables:

set "VCVARS=C:\Program Files\Microsoft Visual Studio\2022\BuildTools\VC\Auxiliary\Build\vcvars64.bat"
set "SAGE_PYTHON=C:\path\to\venv\Scripts\python.exe"
set "SAGE_CUDA_HOME=C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.3"
set "SAGE_CLANG_DIR=C:\path\to\clang+llvm\bin"
set "SAGE_MARCH=native"

build_optimized.bat

Produce a redistributable wheel instead of installing:

python -m pip wheel . --no-build-isolation --no-deps -w dist

Tuning knobs

SAGE_ARCH and SAGE_TORCH_ARCH must agree. For an RTX 4090: SAGE_ARCH=89, SAGE_TORCH_ARCH=8.9.

---

4. Flags applied

Device (nvcc):

-gencode arch=compute_120a,code=sm_120a
-O3 --use_fast_math --threads 0
-Xptxas -O3,-dlcm=ca,--allow-expensive-optimizations=true
--extra-device-vectorization
-Xcompiler /fp:fast -Xcompiler /MD -Xcompiler /Zc:preprocessor

Host (clang-cl, pybind glue only):

/std:c++20 /clang:-O3 /clang:-march=znver3 /clang:-mtune=znver3 /clang:-flto=thin
/fp:fast /clang:-fno-finite-math-only /clang:-fhonor-infinities /clang:-fhonor-nans

Link (lld-link): /OPT:REF /OPT:ICF=10

-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.

1. _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.
2. Compiler.link / Compiler.initialize — repoint at lld-link and add the link flags. This must be applied to both setuptools._distutils.compilers.C.msvc and distutils.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.

---

6. Verification

cuobjdump confirms the arch claim:

$ cuobjdump --list-elf _qattn_sm89.cp313-win_amd64.pyd
ELF file 1: _qattn_sm89.cp313-win_amd64.1.sm_120a.cubin
...
$ cuobjdump --list-ptx _qattn_sm89.cp313-win_amd64.pyd
(empty — SASS only)

Accuracy vs a float64 torch SDPA reference, 3 shapes × causal/non-causal:

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.

Speed, B=2 H=24 L=4096 D=128 fp16 non-causal, RTX 5070:

---

7. ComfyUI usage

python main.py --use-sage-attention

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
• nothing ever writes transformer_options["optimized_attention_override"]
• --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:

grep -rn "get_attention_function\|optimized_attention_override" comfy/

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.

---

Files

Credits

• SageAttention by THU-ML — Apache 2.0
• woct0rdho/SageAttention for Windows wheel groundwork

The wheel is a build of upstream SageAttention and inherits its Apache 2.0 license. The build scripts in build/ are released under the same terms.