Research map / Architectures and efficiency / Attention and transformers
SageAttention: Accurate 8-Bit Attention for Plug-and-play Inference Acceleration
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes SageAttention, a plug-and-play post-training 8-bit quantization method for attention (INT8 QK^T with K mean-smoothing, FP16 PV with FP16 accumulator) that speeds up attention about 2x over FlashAttention2 and xformers with negligible end-to-end metric loss.
Problem
Attention has O(N^2) cost and dominates latency for long sequences, but prior quantization work focuses on linear layers; directly quantizing attention to 8-bit degrades accuracy substantially across models.
Approach
Quantize Q and K to INT8 with per-token or per-block scales after subtracting the mean of K across tokens (which leaves softmax attention unchanged), keep P and V in FP16 with an FP16 accumulator for the PV matmul, and adaptively pick the fastest kernel per layer that preserves cosine similarity above 99.8%.
Key results
- SageAttention kernels reach about 340-341 TOPS on RTX4090, roughly 2x faster than FlashAttention2 and 2.7-2.9x faster than xformers on average.
- End-to-end speedup of attention averages 2.83x on RTX4090 across Unidiffuser, UltraPixel, CogvideoX, Llama2 and TIMM, with metrics essentially matching full precision (e.g., Llama2 WikiText ppl 5.823 vs 5.824).
- Smoothing K adds under 0.2% overhead in TOPS and substantially improves accuracy; e.g., per-block CogVideoX Fscore rises from 2.014 without smoothing to 3.718 with it (full precision 3.768).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add adaptive selection of SAGEAttn-vB for suitable layers, using SAGEAttn-B otherwise (ablation) | SageAttention · CogvideoX attention | TOPS: 292.17 -> 327.57 | Helped |
| Add adaptive selection of SAGEAttn-vB for suitable layers, using SAGEAttn-B otherwise (ablation) | SageAttention · Llama2 attention | TOPS: 208.59 -> 231.74 | Helped |
| Add K smoothing by subtracting the mean key across tokens (ablation) | SAGEAttn-T per-token quantization · Attention output accuracy | cosine similarity: 62.24% -> 99.47% | Helped |
| Add K smoothing by subtracting the mean key across tokens (ablation) | SAGEAttn-B per-block quantization · Attention output accuracy | cosine similarity: 30.60% -> 99.31% | Helped |
| Add K smoothing by subtracting the mean key across tokens (ablation) | Per-tensor Q,K quantization · Attention output accuracy | cosine similarity: 41.40% -> 98.06% | Helped |
| Enable K smoothing (ablation) | SAGEAttn-B · CogvideoX attention | TOPS: 327.57 -> 327.52 | No clear effect |
| Enable K smoothing (ablation) | SAGEAttn-B · UltraPixel attention | TOPS: 325.18 -> 324.56 | No clear effect |
| Use FP16 instead of FP32 accumulator (ablation) | FP16 P,V attention matmul · Attention output accuracy across real-model layers | cosine similarity: 99.98% -> 99.98% | No clear effect |
| Quantize P,V to INT8 rather than retaining FP16 (ablation) | SageAttention · Attention output accuracy across real-model layers | cosine similarity: 99.94% -> 99.70% | Hurt |
| Use per-block rather than per-token Q,K quantization (ablation) | SageAttention · Attention output accuracy across real-model layers | cosine similarity | No clear effect |
| Use per-block rather than per-token Q,K quantization (ablation) | SAGEAttn-vT · Attention output accuracy across real-model layers | cosine similarity: 99.9% -> 98.9% | Hurt |
| Replace full-precision attention with SageAttention | Llama2 · WikiText | perplexity: 5.823 -> 5.824 | No clear effect |
Surprising findings
- Direct INT8 quantization of Q and K yields completely blurry images and video, and Llama2 falls to random-guess-level MMLU accuracy (25.5%) with INT8 attention.
- The quantized FlashAttention3 (FP8) baseline performs worse than the proposed INT8 approach in accuracy metrics, e.g., Unidiffuser FID 394.13 vs 166.52 for per-token smoothed K.
Limitations
- Kernels are implemented for RTX4090 and RTX3090 consumer GPUs using Triton; speed gains on Hopper or other hardware are not shown.
- Adaptive selection of kernels per layer requires a per-layer cosine similarity calibration pass over inputs.
- Some accuracy degradation remains in certain layers for the INT8 variants, so the FP16-accumulator kernel is needed for worst-case accuracy.
Experimental setup
Datasets: WikiText LAMBADA MMLU open-sora prompt set COCO 2014 val ImageNet ImageNet-Sketch ImageNet-Rendition TextVQA POPE VQAv2
Benchmarks: WikiText LAMBADA MMLU VBench-style CogVideoX metrics (CLIPSIM, CLIP-T, VQA-a, VQA-t, FScore) COCO FID/sFID/CLIP/ImageReward ImageNet ImageNet-Sketch ImageNet-Rendition TextVQA POPE VQAv2
Metrics: perplexity accuracy cosine similarity relative L1 RMSE TOPS FID sFID CLIPSIM CLIP-T
Baselines: FlashAttention2 xformers FlashAttention3 (FP8) Torch attention AWQ Q-diffusion ViDiT-Q
Builds on: FlashAttention-2 dynamic quantization online softmax Triton
Closest papers by meaning
- SageAttention2: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantization (ICML 2025 · similarity 0.92)
- SageAttention3: Microscaling FP4 Attention for Inference and An Exploration of 8-Bit Training (NeurIPS 2025 · similarity 0.90)
- Delta Attention: Fast and Accurate Sparse Attention Inference by Delta Correction (NeurIPS 2025 · similarity 0.88)
- XAttention: Block Sparse Attention with Antidiagonal Scoring (ICML 2025 · similarity 0.88)
- SpargeAttention: Accurate and Training-free Sparse Attention Accelerating Any Model Inference (ICML 2025 · similarity 0.88)
- FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning (ICLR 2024 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.