PaperlayerView exact revision on arXiv

Source-linked AI summary

SageAttention: Accurate 8-Bit Attention for Plug-and-play Inference Acceleration

Jintao Zhang, Jia Wei, Haofeng Huang, Pengle Zhang, Jun Zhu, Jianfei Chen

arXiv:2410.02367v9cs.LG

TL;DR

Long sequences make quadratic attention a major inference bottleneck, while prior quantization mainly targets linear layers and direct attention quantization can harm accuracy. The paper proposes SageAttention, a plug-and-play INT8 attention method combining quantization strategies for efficient computation. It outperforms FlashAttention2 and xformers by about 2.1× and 2.7×, respectively, while preserving end-to-end metrics across language, image, and video models.

  • Problem

    Attention becomes the primary cost for long sequences, but existing quantization primarily optimizes linear layers and direct attention quantization can severely degrade model performance.

  • Method

    SageAttention is a plug-and-play post-training INT8 quantization method that adds quantization around attention’s matrix multiplications while keeping online softmax in full precision.

  • Results

    2.1× and 2.7× faster than FlashAttention2 and xformers, respectively, SageAttention maintains end-to-end metrics across language, image, and video generation models.

  • Takeaways & Limitations

    SageAttention provides plug-and-play attention acceleration with negligible end-to-end metric loss across diverse models.

Abstract

from arXiv · show

The transformer architecture predominates across various models. As the heart of the transformer, attention has a computational complexity of $O(N^2)$, compared to $O(N)$ for linear transformations. When handling large sequence lengths, attention becomes the primary time-consuming component. Although quantization has proven to be an effective method for accelerating model inference, existing quantization methods primarily focus on optimizing the linear layer. In response, we first analyze the feasibility of quantization in attention detailedly. Following that, we propose SageAttention, a highly efficient and accurate quantization method for attention. The OPS (operations per second) of our approach outperforms FlashAttention2 and xformers by about 2.1 times and 2.7 times, respectively. SageAttention also achieves superior accuracy performance over FlashAttention3. Comprehensive experiments confirm that our approach incurs almost no end-to-end metrics loss across diverse models, including those for large language processing, image generation, and video generation. The codes are available at https://github.com/thu-ml/SageAttention.

1 INTRODUCTION

Long sequences make attention a dominant transformer cost, while existing quantization mainly targets linear layers and direct attention quantization can severely damage accuracy. SageAttention addresses this gap with plug-and-play 8-bit attention acceleration.

  • Motivation: O(N^2) attention cost dominates other operations when sequence lengths reach 8K∼128K.This motivates faster attention for video generation and language-model prefilling.
  • Research gap: Existing quantization methods primarily optimize linear layers, leaving attention in high precision and lacking systematic attention-quantization studies.FlashAttention3 offers FP8 attention but is restricted to Nvidia Hopper architecture.
  • Challenge: Direct INT8 quantization of attention can severely degrade performance, producing blurry Unidiffuser images and 25.5% MMLU accuracy for Llama2.The difficulty arises because attention combines softmax with QK⊤ and PV matrix multiplications.
  • Proposed approach: SageAttention is a plug-and-play post-training quantization method that replaces the original high-precision attention implementation at inference time.It targets efficient attention computation without requiring additional training.
  • Efficiency: 2.1× and 2.7× faster than FlashAttention2 and xformers, respectively, SageAttention accelerates attention on RTX4090 and 3090 GPUs.The implementation uses Triton and fused attention kernels with INT8 and FP16 Tensor Core operations.

2 RELATED WORK

Efficient attention methods reduce computation in different ways, but sparse attention is limited because omitted calculations make it applicable only in some scenarios.

  • Sparse Attention: Sparse attention selects only parts of a sequence for processing with standard attention.Examples include Swin Transformer, Twins, UniFormer, Attentionsinks, InfLLM, LongLoRA, Minference, and SkipAttention.

3 PRELIMINARY

FlashAttention reduces attention memory traffic through tiling and online softmax, while quantization accelerates matrix multiplication by converting operands to low precision and restoring scale afterward.

  • Attention computation: Attention computes S = QK⊤/d, applies softmax to obtain P, and produces O = PV, with S and P sized N × N.Because N can be very large, these intermediate matrices create substantial memory-I/O costs.
  • FlashAttention: FlashAttention tiles Q, K, and V across tokens and uses online softmax to compute output blocks without reading or writing full S and P matrices.This reduces global-memory I/O for the large N × N intermediates.
  • FlashAttention: FlashAttention iteratively processes K and V blocks to update each output block until all token blocks are consumed.The supplied formulation initializes running vectors and produces O_i after the final block iteration.
  • Quantization: Quantized matrix multiplication converts high-precision operands into low-precision matrices with scales, multiplies them, and dequantizes the result.Modern GPUs generally execute low-precision matrix multiplication faster than higher-precision multiplication.
  • Dynamic quantization: Dynamic quantizers choose scales from maximum absolute values, with per-tensor, per-token, per-channel, and per-block granularity options.Per-block quantization assigns one scale to each token block, matching FlashAttention’s tiling structure.

4 SAGE ATTENTION

SageAttention accelerates attention with plug-and-play 8-bit quantization while preserving accuracy through targeted quantization choices, smoothing, and mixed-precision computation. Its formulation quantizes Q, K, P, and V around FlashAttention-style tiling, retains full-precision online softmax, and selects kernels balancing accuracy and speed.

  • 4.1 FORMULATION: SageAttention adds quantizers to Q, K, P, and V and dequantizers to both QK^T and PV Matmuls, while keeping online softmax in full precision.The computation remains organized with FlashAttention-style tiles.
  • 4.2 SMOOTH MATRIX K: Smoothing K significantly improves end-to-end accuracy, with attention speed overhead below 0.2%.The method subtracts the mean key before quantization without changing attention scores.
  • 4.3 QUANTIZATION FOR Q, K, P, V: INT8 is selected for Q and K because it provides higher accuracy than tested FP8 formats and is faster on commonly used GPUs.The reported GPU comparison states that INT8 Matmul is two times faster than FP8 on RTX4090 and 3090.
  • 4.4 FP16 ACCUMULATOR: FP16 Matmul with an FP16 accumulator is used for P and V because INT8 produces unacceptable worst-layer accuracy, while FP16 incurs no accuracy loss versus FP32 accumulation.The FP16-accumulator choice also reduces computation and register costs on relevant GPUs.
  • KERNEL IMPLEMENTATIONS: SageAttention offers four kernels that vary Q,K granularity and P,V precision, with SAGEAttn-B accurate across all models and SAGEAttn-vB about 4% faster for selected layers.Layer-wise cosine similarity is used to select the faster variant where appropriate.
  • IMPLEMENTATION: The implementation fuses quantization with preceding operations such as ROPE and uses INT8 and FP16 Tensor Core instructions to reduce overhead.The design also fuses the QK^T scaling coefficient into Q quantization and reduces data-transfer and register demands.

5 EXPERIMENTS

Experiments evaluate SageAttention’s speed, kernel accuracy, end-to-end metrics, and overhead across language, image, and video models on RTX4090 and RTX3090.

  • Models: SageAttention was evaluated on Llama2, CogvideoX, Unidiffuser, UltraPixel, TIMM, and Llava1.6 across language, image, video, classification, and visual-question-answering tasks.
  • Speed: 341 TOPS and approximately 2x faster than FlashAttention2 and 2.9x faster than xformers were achieved on RTX4090 speed tests.Tests covered headdim=64 or 128, varying sequence lengths, and causal or non-causal attention; RTX3090 showed similar speedup performance.
  • Accuracy: 100% cosine similarity and RMSE in the e-4 level were achieved by SAGEAttn-T and SAGEAttn-B against full-precision attention.
  • End-to-end performance: 2.83x average real speedup over original attention was measured across Unidiffuser, UltraPixel, CogvideoX, Llama2, and TIMM.
  • End-to-end performance: SageAttention matched full-precision end-to-end performance across models, with 0.2% average degradation on four models and higher performance on TIMM.
  • Ablations: Adaptive quantization increased attention speed by 11.7% without metric loss, while smoothing K reduced attention speed by less than 0.2%.

6 CONCLUSION AND FUTURE WORK

The paper introduces SageAttention as an efficient, precise INT8 attention quantization method that preserves end-to-end metrics across language, image, and video generation models.

  • SageAttention uses INT8 quantization for attention and maintains end-to-end metrics across language, image, and video generation models.
  • Smoothing K adds under 0.2% speed overhead, FP16 accumulators improve the P,V Matmul, and adaptive quantization improves OPS by 12% without sacrificing accuracy.
  • SageAttention is approximately 2.1x faster than FlashAttention2 and 2.7x faster than xformers.

ACKNKOWLEGEMENT

The paper acknowledges implementation assistance, institutional support, computing resources, and experimental environment details.

  • The authors thank Haofeng Huang for implementation assistance and acknowledge support from NSFC, Tsinghua Institute for Guo Qiang, and the High Performance Computing Center.
  • Experiments used OpenAI Triton on Ubuntu 22.04 servers with RTX4090 and RTX3090 hardware configurations.
  • The implementation specifies block sizes of 128 for Q and 64 for K and V, with warp and stage parameters detailed in Table 12.
  • UltraPixel and Unidiffuser used 256 COCO prompts and corresponding images, while CogvideoX used open-sora prompts exceeding 120 words.

A.4 ADDITIONAL EXPERIMENTS

Additional experiments compare SageAttention with weight-and-activation quantization baselines on Llama2 and Unidiffuser.

  • Table 13 compares SageAttention with AWQ (W4A16) on Llama2.
  • Table 14 compares SageAttention with Q-diffusion (W8A8) on Unidiffuser.

A.5 COMPARISON WITH OTHER METHODS

SageAttention is compared with task-specific quantization methods and alternative attention implementations, with experiments covering language and image-generation settings. The comparisons emphasize its distinct focus on attention quantization and its stable accuracy and speed benefits.

  • SageAttention is orthogonal to AWQ, Q-diffusion, and ViDiT-Q because those methods primarily quantize linear layers.
  • AWQ compresses LLM parameters without accelerating computation, while Q-diffusion and ViDiT-Q have limited reported or released acceleration support.
  • SageAttention is evaluated against other methods through Llama2 WikiText perplexity and prefilling speed, plus Unidiffuser comparisons with Q-diffusion.
  • Llama2 metrics remain stable after quantization because its attention Q, K, and V distributions are relatively uniform.

B IMPLEMENTATION BASED ON TORCH ATTENTION

This implementation section documents Torch Attention comparisons and quantization-error analyses across numerical formats and smoothing choices. The supplied captions also identify related image and video examples and an RTX3090 speedup comparison.

  • FlashAttention and Torch attention are identified as commonly used attention implementations, with implementation speeds reported using Torch.
  • Table 17 compares Q · K numerical error across quantization types, using per-token quantization on Q and K from Unidiffuser’s 24th layer.
  • The section includes captions for UltraPixel and Open-Sora generation examples, additional examples for UltraPixel, Unidiffuser, and CogvideoX, and RTX3090 real-speedup results.
  • INT8 quantization achieves higher Q · K precision than E4M3 and E5M2 in the Unidiffuser experiment.
  • Table 18 compares quantized-attention error with and without smoothed K across models, showing significant precision benefits from smoothing K.

B.2 VISUALIZED RESULTS

Visual comparisons show SageAttention preserving the apparent quality of full-precision image and video generation, while RTX3090 measurements report faster execution across evaluated models.

  • SageAttention matches full precision in high-resolution UltraPixel images and produces an identical Open-Sora video across precisions.
  • Additional visual comparisons are reported for UltraPixel, Unidiffuser, and CogvideoX.
  • 2.7× average speedup over original attention is reported on RTX3090 across Unidiffuser, UltraPixel, CogvideoX, Llama2, and TIMM.
Loading 2410.02367v9…