sageattention

What it is and what it does

SageAttention is a specialized attention kernel library that replaces the standard transformer attention computation with quantized variants using 8-bit (or 4-bit in SageAttention2) precision. It integrates smoothing techniques, per-block quantization for queries and keys, and FP16 accumulators to maintain accuracy while reducing memory bandwidth and computation cost during inference.

The package is designed as a drop-in replacement for scaled_dot_product_attention, allowing users to accelerate existing models with minimal code changes. It currently supports head dimensions of 64, 96, and 128, variable sequence lengths between queries and key-values, and group-query attention patterns. Performance is optimized for RTX4090 and RTX3090 GPUs; benefits on other architectures are not guaranteed.

Use it for:

• Accelerate video generation models like CogVideoX by replacing their attention layers with quantized kernels
• Speed up large language model inference on supported GPUs by patching scaled_dot_product_attention globally
• Reduce memory bandwidth requirements in transformer-based inference pipelines without retraining
• Deploy transformer models with lower latency on resource-constrained inference servers using RTX-series GPUs
• Benchmark quantized attention accuracy on models with variable sequence lengths per batch

Worth the install?

AI-flagged interpretation of the facts on this page — verify before relying

SageAttention provides quantized attention kernels for transformer inference that replace the standard scaled dot-product attention with 8-bit or 4-bit quantized variants, designed to accelerate inference on specific GPU architectures.

Yes, if you run transformer inference on RTX4090 or RTX3090 GPUs and can verify that the supported head dimensions and attention patterns match your model. The low install friction, permissive license, and active maintenance make it a reasonable experiment. No, if you use other GPU architectures or require head dimensions outside 64, 96, 128—the stated optimization scope is narrow and performance gains are not guaranteed elsewhere.

Install

sageattention on PyPI

pip install sageattention

uv add sageattention

poetry add sageattention

Installing sageattention

Before you install

Low install friction with a pure-Python wheel distribution. The package is in beta status and has been actively maintained since its first release on 2024-10-05, with the latest release on 2024-11-20. Requires Python 3.9+ and torch 2.3.0+, with triton 2.3.0+ as a build dependency.

License in practice

BSD 3-Clause License is permissive and poses no significant restrictions on use, modification, or distribution in commercial or private projects.

Quickstart

pip install sageattention

from sageattention import sageattn

attn_output = sageattn(q, k, v, tensor_layout="HND", is_causal=False, smooth_k=True)

Requires Python>=3.9, torch>=2.3.0, and triton>=2.3.0. Performance is optimized for RTX4090 and RTX3090 GPUs; other architectures may not see significant speedup. Head dimension must be one of 64, 96, or 128.

Verify before relying

• Whether performance gains materialize on GPU architectures other than RTX4090 and RTX3090, given the stated optimization focus
• Actual accuracy impact of disabling smooth_k in production workloads with irregular q, k, v distributions
• Compatibility with transformer models that do not use F.scaled_dot_product_attention directly

Package facts

Evidence: sageattention-1.0.6-py3-none-any.whl

Tags

More Python Modules packages

Further reading