SageAttention Documentation

For more granular control, you can replace specific attention processors in your model source code. For example, in Mochi models, you can replace the MochiAttnProcessor2_0 from diffusers with a custom attention class.

HunyuanVideo Workaround: Because HunyuanVideo uses attention_mask (which sageattn does not support), you must modify the official attention implementation to split text tokens from image tokens. Apply SageAttention only to the large, mask-free image-token self-attention, while keeping the masked text part on SDPA or FlashAttention.

Base Environment Requirements

• python>=3.9
• torch>=2.3.0
• triton>=3.0.0
• flash-attn (for benchmarking)

CUDA Version Requirements

• >=12.8 for Blackwell or SageAttention2++
• >=12.4 for fp8 support on Ada
• >=12.3 for fp8 support on Hopper
• >=12.0 for Ampere

Installation Methods

Via pip (Recommended for v2.2.0)

To use SageAttention 2.2.0 (which includes SageAttention2++):

pip install sageattention==2.2.0 --no-build-isolation

To use SageAttention V1 (Triton-based, slower):

pip install sageattention==1.0.6

From Source

git clone https://github.com/thu-ml/SageAttention.git
cd SageAttention 
# Optional: set build flags for faster compilation
export EXT_PARALLEL=4 NVCC_APPEND_FLAGS="--threads 8" MAX_JOBS=32 
python setup.py install

To run parallel SageAttention inference, you must install the latest xfuser (xDiT) and a specific development version of diffusers from source.

Follow these steps:

1. Install xfuser with flash_attn support.
2. Clone and install the latest diffusers from GitHub.
3. Run the parallel inference script.

Note: diffusers must be version >=0.32.0.dev0 to be compatible with the latest xDiT.

# install latest xDiT(xfuser).
pip install "xfuser[flash_attn]"

# install latest diffusers (>=0.32.0.dev0), need by latest xDiT.
git clone https://github.com/huggingface/diffusers.git
cd diffusers && python3 setup.py bdist_wheel && cd dist && python3 -m pip install *.whl

# then run parallel sage attention inference.
./run_parallel.sh

SageAttention3 is optimized for specific model types.

Supported Models:

• Video generation: CogVideoX-2B, HunyuanVideo, Mochi.
• Image generation: Flux, Stable-Diffusion3.5, and almost all other image generation models.

Hybrid Approach for Lossless Acceleration: SageAttention3 does not guarantee lossless acceleration for all models. For other video generation models, you can achieve lossless acceleration by using a hybrid approach:

1. Apply SageAttention2++ only at the first and last timesteps.
2. Use SageAttention3 for all other timesteps.

SageAttention3 requires a specific base environment and must be compiled from source.

Base Environment Requirements:

• python >= 3.13
• torch >= 2.8.0
• CUDA >= 12.8

Installation Steps: Clone the repository, navigate to the sageattention3_blackwell directory, and run the setup script.

git clone https://github.com/thu-ml/SageAttention
cd SageAttention/sageattention3_blackwell 
python setup.py install

To run inference for CogVideoX using SageAttention, navigate to the example directory and use the cogvideox_infer.py script with the --attention_type sage flag.

Example command:

cd example
python cogvideox_infer.py --model cogvideox-2b --compile --attention_type sage

Output videos will be saved to ./example/videos/<model>/<attention_type>/.

You can achieve plug-and-play acceleration by monkey-patching torch.nn.functional.scaled_dot_product_attention with sageattn.

Note: Not all models support this direct replacement. For complex models (like image/video DiT models), you may need to modify the specific Attention Class of the target model instead.

import torch.nn.functional as F
from sageattention import sageattn

F.scaled_dot_product_attention = sageattn

The sageattn function automatically selects the optimal kernel based on your GPU to balance performance and accuracy.

Input Requirements:

• q, k, v must be FP16/BF16 dtype.
• Default shape: (batch_size, head_num, seq_len, head_dim) with tensor_layout="HND".
• For shape (batch_size, seq_len, head_num, head_dim), set tensor_layout="NHD".

from sageattention import sageattn
attn_output = sageattn(q, k, v, tensor_layout="HND", is_causal=False)

The sageattn3_blackwell function provides microscaling FP4 attention.

Arguments:

• q, k, v: Tensors with FP16/BF16 dtype and shape (batch_size, head_num, seq_len, head_dim).
• is_causal: Boolean determining whether to use a causal mask.

Returns:

• attn_output: The result of the attention operation.

from sageattn3 import sageattn3_blackwell
attn_output = sageattn3_blackwell(q, k, v, is_causal=False)

SageAttention provides several specialized kernels for different quantization and backend requirements: