Practical overview
When it fits
Use SageAttention when optimizing attention computation in an existing CUDA-based inference pipeline and the target GPU, software stack and tensor formats match the implementation's requirements.
Who it is for
- Developers building GPU inference pipelines for language models.
- Teams integrating attention kernels into image-generation or video-generation inference systems.
- Researchers and infrastructure engineers comparing attention implementations across GPU hardware.
Documented uses
- Replace or integrate attention computation with quantized GPU kernels.
- Run attention using FP16 or BF16 query, key and value tensors in HND or NHD layout.
- Select attention kernels automatically by GPU or use APIs for specific quantization and backend options.
- Use attention in distributed inference or with torch.compile in non-cudagraphs mode.
Requirements & limits
The source code is publicly available under the Apache-2.0 license. The documented package installation uses pip; the repository also provides source-compilation instructions.
- The documented base environment requires Python 3.9+, PyTorch 2.3.0+ and Triton 3.0.0+.
- CUDA version requirements differ by GPU generation and feature; Blackwell or SageAttention2++ requires CUDA 12.8 or later.
- The primary API expects FP16 or BF16 Q, K and V tensors in HND or NHD layout.
- Different APIs expose different quantization and backend combinations; the general
sageattnAPI automatically chooses a kernel based on GPU. - The repository describes optimized kernels for Ampere, Ada and Hopper, and separately documents SageAttention3 code for Blackwell.
Check current requirements in the official documentation. Source-based guidance does not mean OpenViST has run or independently tested this project.
Install SageAttention 2.2.0
Install the documented SageAttention 2.2.0 package in an environment with the listed Python, PyTorch, Triton and CUDA prerequisites.
Before you start
- Python 3.9 or later
- PyTorch 2.3.0 or later
- Triton 3.0 or later
- CUDA version appropriate to the target GPU and selected features: 12.0 or later for Ampere, 12.3 or later for Hopper FP8, 12.4 or later for Ada FP8, or 12.8 or later for Blackwell or SageAttention2++
Steps
Prepare an environment meeting the documented Python, PyTorch, Triton and CUDA requirements.
Install the package with the command below.
pip install sageattention==2.2.0 --no-build-isolation
Expected outcome
The SageAttention 2.2.0 package, which contains SageAttention2++, is installed in the prepared environment.
Guide limitations
- CUDA requirements vary by GPU generation and feature; the listed version thresholds do not establish compatibility for every configuration.
- The README also documents source compilation, but this recipe covers the package-install route only.
Call the SageAttention API on attention tensors
Pass documented FP16 or BF16 query, key and value tensors to the `sageattn` function with an explicitly selected supported tensor layout.
Before you start
- SageAttention installed in a suitable CUDA environment
- Query, key and value tensors in FP16 or BF16
- Inputs arranged as HND
(batch_size, head_num, seq_len, head_dim)or NHD(batch_size, seq_len, head_num, head_dim)
Steps
Import
sageattnfromsageattention.Call
sageattn(q, k, v, tensor_layout="HND", is_causal=False)for HND inputs.For NHD inputs, set
tensor_layout="NHD"; setis_causalaccording to whether a causal mask is needed.
Expected outcome
The function returns an attention output, with the API selecting a kernel based on the GPU.
Guide limitations
- The documented API requires FP16 or BF16 inputs in one of the stated tensor layouts.
- This recipe describes the documented call pattern, not a tested model integration or performance result.
Sources & review
Source record details
- Source availability
- Active
- Sources checked
- 7 October 2026
- Original website title
- thu-ml/SageAttention · Official homepage title · retained in its source language