Source-linked AI summary
SageAttention3: Microscaling FP4 Attention for Inference and An Exploration of 8-Bit Training
Jintao Zhang, Jia Wei, Pengle Zhang, Xiaoming Xu, Haofeng Huang, Haoxu Wang, Kai Jiang, Jianfei Chen, Jun Zhu
TL;DR
Attention’s quadratic complexity motivates more efficient inference and training, but prior low-bit attention primarily targeted inference. The paper introduces microscaling FP4 attention and trainable 8-bit attention, achieving strong inference speedups and lossless fine-tuning performance while exposing pretraining limitations.
Problem
Prior low-bit attention work focused on inference, leaving the feasibility of low-bit attention for training large models unexplored despite the importance of training efficiency.
Method
The paper applies microscaling FP4 quantization to attention for inference and designs trainable 8-bit attention for forward and backward propagation.
Results
1038 TOPS on RTX5090 delivers a 5× speedup over the fastest FlashAttention, while 8-bit attention achieves lossless fine-tuning performance but is unsuitable for pretraining tasks.
Takeaways & Limitations
The approach provides plug-and-play inference acceleration and supports lossless fine-tuning performance, but low-bit attention remains limited for pretraining.
Takeaways & Limitations
The current 8-bit attention has limitations in pretraining tasks, and its speed remains below theoretical upper bounds.
Abstract
from arXiv · showhide
The efficiency of attention is important due to its quadratic time complexity. We enhance the efficiency of attention through two key contributions: First, we leverage the new FP4 Tensor Cores in Blackwell GPUs to accelerate attention computation. Our implementation achieves 1038 TOPS on RTX5090, which is a 5x speedup over the fastest FlashAttention on RTX5090. Experiments show that our FP4 attention can accelerate inference of various models in a plug-and-play way. Second, we pioneer low-bit attention to training tasks. Existing low-bit attention works like FlashAttention3 and SageAttention focus only on inference. However, the efficiency of training large models is also important. To explore whether low-bit attention can be effectively applied to training tasks, we design an accurate and efficient 8-bit attention for both forward and backward propagation. Experiments indicate that 8-bit attention achieves lossless performance in fine-tuning tasks but exhibits slower convergence in pretraining tasks. The code is available at https://github.com/thu-ml/SageAttention.
1 Introduction
The paper targets attention’s quadratic cost with microscaling FP4 inference and trainable 8-bit attention. It reports major inference acceleration and lossless fine-tuning performance, while identifying limitations for pretraining.
- Attention efficiency becomes increasingly important for generation models because attention has quadratic time complexity with sequence length.
- The paper introduces FP4 attention for inference and trainable 8-bit attention for training, addressing a gap in prior low-bit attention work.
- FP4 attention uses microscaling for QK^T and PV, with 1x16 quantization groups and two-level quantization for the attention map.
- 1038 TOPS on RTX5090 represents a 5× speedup over FlashAttention for the proposed FP4 attention.
2 Preliminary
The preliminary section describes attention as two matrix multiplications separated by softmax and explains how FlashAttention reduces memory input/output through online softmax. It also introduces quantization as low-bit matrix multiplication with scale factors.
- Attention computes S = QK^T, P = Softmax(S), and O = PV, with Q, K, and V shaped by sequence length and head dimension.
- FlashAttention partitions Q, K, and V into blocks and uses online softmax to avoid large memory input/output for S and P.
- Quantization converts high-bit matrices and their scale factors into low-bit representations so Tensor Cores can accelerate matrix multiplication.
3 FP4 Attention for Inference Acceleration
This section applies microscaling FP4 quantization to both attention matrix multiplications and introduces two-level quantization for the attention map. The resulting design addresses FP4 accuracy challenges through blockwise quantization, scale-range utilization, and hardware-aware implementation choices.
- 3.1 Microscaling FP4 Attention: FP4 microscaling partitions matrices into 1 × n blocks, assigning each block an FP8 scale factor and FP4 representation.Quantization and dequantization operate through paired low-bit values and scale factors.
- 3.1 Microscaling FP4 Attention: SageAttention3 applies FP4 microscaling quantization to both QK⊤ and PV matrix multiplications.The attention pipeline quantizes Q, K, the online-softmax output, and V before the two FP4 matrix multiplications.
- 3.3 Hardware Implementation Optimization: The implementation incorporates smoothing, accumulator-layout permutation, fused online-softmax quantization, and producer-warp epilogues to improve accuracy or kernel performance.These optimizations adapt the FP4 attention kernel to hardware layout, reduction, and register constraints.
- 3.1 Microscaling FP4 Attention: NVFP4 is selected over MXFP4 because empirical CogVideoX results show higher attention accuracy across layers.NVFP4 uses 1 × 16 quantization blocks with E4M3 scale factors, whereas MXFP4 uses 1 × 32 blocks with E8M0 scale factors.
- 3.2 Two-Level Quantization: Two-level quantization addresses direct FP4 quantization errors in eP by improving utilization of the E4M3 scale-factor range.The method normalizes eP before quantizing it and then reconstructs it using two scale factors.
4 INT8 Attention for Training
This section introduces SageBwd, an INT8 attention implementation covering forward and backward propagation for training. It preserves selected FP16 computation where gradient accuracy is especially sensitive while quantizing the remaining matrix multiplications.
- 4 INT8 Attention for Training: SageBwd extends low-bit attention from inference to training by quantizing six of seven attention matrix multiplications to INT8.The method is designed for both forward and backward propagation.
- 4.1 Forward: The forward pass uses smoothed K, per-block INT8 quantization for Q, K, and V, and per-token quantization for the attention map.Per-token quantization for eP improves accuracy compared with a static per-block scale factor.
- 4.2 Backward: The backward pass quantizes the gradient-related matrix multiplications with INT8 while retaining dOV⊤ in FP16.Keeping dOV⊤ in FP16 improves dQ accuracy because errors there propagate through dP and dS into dQ and dK.
- 4.2 Backward: The backward computation contains five attention matrix multiplications, with four accelerated using INT8 per-block quantization.The design avoids quantizing dOV⊤ because longer sequences can accumulate its downstream gradient errors.
5 Experiments
Experiments evaluate SageAttention3 and SageBwd across language, image, and video generation models, measuring kernel speed, end-to-end latency, quality, and training behavior. SageAttention3 preserves end-to-end quality while accelerating inference, whereas SageBwd provides lossless fine-tuning performance but slower pretraining convergence.
- Kernel speed: 5× and 11×: SageAttention3 is faster than FlashAttention and xformers on RTX5090 while maintaining end-to-end metrics across models.The paper reports these comparisons as main results.
- Experimental setup: SageAttention3 is evaluated across representative language, image, and video-generation models against FlashAttention2, xformers, SageAttention, and SageAttention2.The evaluated models include Qwen2.5, Llama3.2, CogVideoX, HunyuanVideo, Mochi, Flux, and Stable-Diffusion3.5.
- Kernel speed: 4–5× over FlashAttention2 and 8–11× over xformers: SageAttention3 accelerates RTX5090 kernel execution.SageBwd reaches up to 1.67× speedup over FlashAttention2 on RTX4090 for forward-plus-backward execution.
- Quality: SageAttention3 incurs almost no end-to-end quality loss across the evaluated models and maintains full generation quality in visible examples.The visible examples cover HunyuanVideo and Stable-Diffusion3.5, with additional examples in the appendix.
- Training effectiveness: SageBwd matches BF16 in fine-tuning but converges relatively slowly during FineWeb-Edu pretraining, limiting its applicability to pretraining tasks.Fine-tuning evaluations cover Qwen2.5 and Llama3.2 on GSM8K, DROP, MMLU, and HELLASWAG.
- End-to-end speed: About 3× and 2.4×: SageAttention3 improves end-to-end inference generation speed for HunyuanVideo and CogVideoX on RTX5090.SageBwd accelerates Llama (1B) training by about 1.15× with 8K/16K-token micro-batches on RTX4090.
- Combined use: INT8 SageBwd followed by FP4 SageAttention3 inference achieves higher GSM8K and MMLU accuracy than BF16 fine-tuning followed by FP4 inference.The paper attributes this improvement to more similar representable data distributions between INT8 and FP4, reducing mismatch error.
6 Related Work
Efficient attention methods primarily exploit GPU hardware and memory-access strategies to accelerate attention computation. FlashAttention variants improve tiling, parallelism, warp partitioning, or GPU-specific kernels, while xformers uses dedicated CUDA kernels.
- FlashAttention: FlashAttention reduces GPU memory I/O between global memory and on-chip SRAM through tiling.FlashAttention2 further improves parallelism and warp partition strategies.
- Other efficient attention methods: FlashAttention3 targets kernel speed specifically on Hopper GPUs, while xformers accelerates attention with dedicated CUDA kernels.
7 Conclusions
The paper contributes SageAttention3, an FP4 inference method, and SageBwd, a trainable 8-bit attention method. The results show strong inference and fine-tuning performance, but pretraining remains limited and SageBwd has room to approach its theoretical speed.
- Contributions: SageAttention3 is presented as the first microscaling FP4 attention for inference, reaching 1038 TOPS on RTX5090 and a 5× speedup over the fastest FlashAttention.It accelerates various models without end-to-end quality-metric degradation.
- Contributions: SageBwd is presented as the first trainable 8-bit attention, achieving lossless fine-tuning performance but limitations in pretraining tasks.
- Limitations and future work: SageBwd remains below its theoretical speed potential, possibly because of suboptimal Triton kernel implementations.The authors identify improving low-bit attention for pretraining as an important direction for future work.
A.1 Visible Comparison Examples
The appendix provides visual examples and comparisons for image and video generation across several models. These figures are intended to show qualitative outputs and scale-strategy differences.
- Image generation: Figures 10 and 11 show visible image-generation examples from Stable-Diffusion3.5 and Flux.
- Scale strategies: Figure 12 compares different scale strategies for eP in CogVideoX.
- Video generation: Figures 13 and 14 show visible video-generation examples from CogVideoX and HunyuanVideo.
A.2 Additional Kernel Speed Comparison
SageBwd accelerates both forward and backward attention kernels relative to FlashAttention on RTX4090, with the strongest gains in forward propagation.
- Experimental setup: Figures 15–18 compare SageBwd with baselines on RTX4090 for forward and backward kernels at head dimensions 128 and 64.The forward comparisons are shown in Figures 15–16, while backward comparisons are shown in Figures 17–18.
A.3 Datasets, Metrics, and Hyperparameters
The experiments evaluate language, video, and image-generation models using task-specific datasets and metrics, with separate hyperparameter settings for pretraining and fine-tuning.
- Datasets: Language models are evaluated on GSM8K, DROP, MMLU, and HELLASWAG, while text-to-video models use open-sora prompt sets and text-to-image models use COCO annotations.These datasets cover the language, video-generation, and image-generation evaluations.
- Metrics: Evaluation uses Accuracy and F1 for text-to-text, five video-quality metrics, and FID and sFID for text-to-image outputs.The video metrics measure alignment, aesthetic and technical quality, and temporal consistency.
- Hyperparameters: Pretraining uses a 400M model with 1e-3 learning rate, whereas fine-tuning uses 3e-5 learning rate for 700 steps with task-specific batch sizes.Both schedules use linear decay; pretraining has 1000 warmup steps and fine-tuning has 100 warmup steps.
- Implementation details: The appendix includes register layouts for FP4 operands and FP32 accumulators used in the implementation.Figures 19–21 show operand, accumulator, and permuted accumulator layouts.
A.4 Additional Experiments of using SageBwd
Additional experiments examine SageBwd’s speed, accuracy, quantization behavior, implementation trade-offs, and theoretical throughput across attention settings.
- Fine-tuning performance: SageBwd remains highly consistent with BF16 across multiple models, datasets, and random seeds in fine-tuning experiments.Tables 9–14 compare SageBwd and BF16 performance for Qwen2.5 and Llama3.2 across four datasets.
- Layer-wise accuracy: SageAttention3’s accumulated L1 error generally increases with layer depth, but retaining three high-error layers in FP16 is used to mitigate this drift.The analysis also observes occasional decreases in deeper layers, suggesting partial error cancellation.
- Theoretical speed comparison: SageAttention3’s theoretical throughput is compared with FlashAttention3 across GPUs because direct empirical comparison is unavailable when FlashAttention3 runs only on Hopper GPUs.The comparison uses NVIDIA’s official throughput documentation.
- Two-level quantization: Two-level quantization refines the quantization intervals for attention probabilities and uses an FP8 scale-factor range distinct from direct quantization.The appendix analyzes the resulting relative-error comparison and representable output count.
- Backward quantization: Keeping dOV⊤ in FP16 avoids error accumulation through dP, dS, dQ, and dK during the backward sequence recurrence.The other four matrix multiplications are accelerated using INT8 per-block quantization.