SageAttention: Quantized Attention Library

What is SageAttention?

SageAttention is a three-generation CUDA/Triton library for quantized attention that accelerates transformer inference by 2-5x over FlashAttention on modern GPUs. It applies INT8 smoothing to QK^T products and FP8 or FP16 accumulation to PV products, with kernels tuned for Ampere, Ada, Hopper, and Blackwell silicon. The API is a single `sageattn` function callable from PyTorch, with per-version variants exposing different precision/accuracy trade-offs.

Core features and supported GPU architectures

• ✓INT8 quantization for QK^T with per-token, per-head, or per-thread granularity (SageAttention2)
• ✓FP8 quantization for PV with FP16 accumulator — avoids the accuracy collapse naive FP8 accumulation causes on Hopper
• ✓Two-level accumulation strategy for PV on Hopper WGMMA, preserving accuracy at FP8 precision
• ✓Smooths outliers in attention logits before quantization, preventing the large errors that standard INT8 rounding produces
• ✓Triton backend (v1 pip) and CUDA backend (v2++), with automatic kernel selection via `sageattn`
• ✓Sparse attention variant SpargeAttn extends the quantized kernel to arbitrary block-sparse patterns without training
• ✓Variable-length sequences via `sageattn_varlen` and native group-query attention support
• ✓RTX 5090 delivers 560 TOPS on attention kernels — 2.7x faster than FlashAttention2 on the same hardware

Using SageAttention in your code

import torch from sageattention import sageattn # Standard HND layout (batch, heads, seq, dim) q = torch.randn(1, 16, 512, 64, dtype=torch.float16, device='cuda') k = torch.randn(1, 16, 512, 64, dtype=torch.float16, device='cuda') v = torch.randn(1, 16, 512, 64, dtype=torch.float16, device='cuda') output = sageattn(q, k, v, tensor_layout='HND', is_causal=False) # Drop-in SDPA replacement for compatible models: # torch.nn.functional.scaled_dot_product_attention = sageattn For CogVideoX-style video models, the README shows a direct SDPA swap works. For image and video DiT models, the recommended path is modifying the Attention class directly (example in example/modify_mochi.py).

Installing SageAttention

pip install sageattention==2.2.0 --no-build-isolation For the latest v2++ features including Blackwell support: compile from source by cloning the repo, then run python setup.py install. Set EXT_PARALLEL=4, NVCC_APPEND_FLAGS="--threads 8", and MAX_JOBS=32 for faster compilation. Note that pip's sageattention v1 uses slower Triton kernels — source compilation enables the full CUDA performance.

SageAttention vs FlashAttention

Who should use SageAttention?

Use SageAttention if you're serving or fine-tuning LLMs, video generation models (CogVideoX, Mochi), or DiT-based image models on Ampere, Ada, or Hopper GPUs and the attention kernel is your bottleneck. It shines when you have long contexts or high batch sizes. Skip it if you're on older hardware, cannot meet CUDA 12+, or need to preserve bit-exact results — and for precision-sensitive production pipelines, default to SageAttention2 rather than v3.

Frequently asked questions

The problem it solves

Attention computation is the dominant bottleneck in long-context transformers, and naive FP16/FP32 attention scales poorly with sequence length. FlashAttention addresses the memory wall but leaves accuracy-pressure quantization on the table. SageAttention closes that gap by quantizing QK^T to INT8 and PV to FP8, squeezing more throughput from tensor cores without the accuracy degradation naive low-precision math would cause.

Best use cases

• •Long-context LLM inference on H100/H800/H20 with FP8 PV accumulation
• •Video generation (CogVideoX, Mochi) inference acceleration via DiT attention replacement
• •Image transformers (DiT-based) with INT8 QK^T and FP16 PV for lossless speedup
• •Batch inference with variable-length sequences in the same batch
• •Serving quantized models where attention computation is the bottleneck
• •Blackwell GPU deployment with SageAttention2++ targeting RTX 5090

Related repositories