Research map / Architectures and efficiency / Quantization
SageAttention3: Microscaling FP4 Attention for Inference and An Exploration of 8-Bit Training
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 SageAttention3, an FP4 microscaling attention kernel for Blackwell GPUs with two-level P quantization for inference, and SageBwd, an INT8 trainable attention that quantizes six of seven attention matmuls for fine-tuning.
Problem
Attention's quadratic cost makes inference and training slow; prior low-bit attention (FlashAttention3, SageAttention) either lacks FP4 support for plug-and-play inference or is inference-only, with no low-bit trainable attention.
Approach
Apply NVFP4 microscaling (1x16 blocks, E4M3 scales) to both attention matmuls, with a two-level scaling for the softmax matrix P that renormalizes each row to [0, 448x6] so the FP8 scale factors use their full range; separately, an INT8 trainable attention that keeps dO V^T in FP16 to avoid gradient error accumulation.
Key results
- SageAttention3 reaches 1038 TOPS on RTX5090, about 5x faster than FlashAttention2 on the same GPU, with CosSim 99.55% against full precision.
- SageBwd fine-tuning matches BF16 on Qwen2.5 (1.5B/3B) and Llama3.2-1B across GSM8K, DROP, MMLU and HellaSwag, e.g. Qwen2.5-3B GSM8K 0.607 vs 0.601.
- SageBwd speeds up Llama (1B) training by about 1.15x at 8K/16K tokens on RTX4090, while pretraining a 400M model converges noticeably slower than BF16.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use FP4 microscaling attention instead of the fastest FlashAttention | Attention inference on RTX5090 · RTX5090 | TOPS: 1038 TOPS; 5× speedup | Helped |
| Use FP4 microscaling attention | CogVideoX inference · CogVideoX | inference latency: 64 s -> 27 s | Helped |
| Use INT8 SageBwd | Llama training · Llama (8K) | one iteration training latency: 2.1 s -> 1.9 s | Helped |
| Replace MXFP4 with NVFP4 (ablation) | FP4 attention on CogVideoX · CogVideoX layers | CosSim: 98.37% -> 99.52% | Helped |
| Use two-level scaling for attention map P instead of direct quantization (ablation) | FP4 attention on CogVideoX · CogVideoX layers | CosSim: 93.32% -> 99.52% | Helped |
| Keep dO Vᵀ in FP16 instead of quantizing it to INT8 (ablation) | SageBwd · CogVideoX gradients | dQ CosSim: 97.47% -> 99.77% | Helped |
| Retain the three most sensitive layers in FP16 (ablation) | SageAttention3 · CogVideoX-2B | Layer10 L1 error: 0.0922 -> 0.0447 | Helped |
| Add Q smoothing (ablation) | SageAttention3 · CogVideoX-2B layers | CosSim: 0.915642 -> 0.982848 | Helped |
| Add K smoothing (ablation) | SageAttention3 · CogVideoX-2B layers | CosSim: 0.915642 -> 0.991176 | Helped |
| Use INT8 instead of FP8 SageBwd (ablation) | CogVideoX-2B backward gradients · CogVideoX-2B gradients | dQ L1 error: 0.0696 -> 0.0290 | Helped |
| Use INT8 instead of FP8 SageBwd (ablation) | Qwen2.5-1.5B fine-tuning · Qwen2.5-1.5B fine-tuning | GSM8K: 0.5031 -> 0.5232 | Helped |
| Use SageBwd instead of BF16 attention | Qwen2.5-1.5B fine-tuning · GSM8K | GSM8K: 0.521 -> 0.520 | No clear effect |
Surprising findings
- 8-bit attention training converges slowly in pretraining even though it matches BF16 after fine-tuning.
- Keeping only the dO V^T matmul in FP16 is critical for dQ accuracy (CosSim 99.77% vs 97.47% for INT8).
Limitations
- SageBwd kernels are implemented in Triton and still lag theoretical speed bounds.
- 8-bit attention converges slowly in pretraining, so it is not suitable for pretraining yet.
- FP8 attention in FlashAttention3 is Hopper-only, and the FP4 kernel targets Blackwell GPUs, limiting hardware coverage.
Experimental setup
Datasets: GSM8K DROP MMLU HellaSwag FineWeb-Edu open-sora prompts COCO
Benchmarks: GSM8K DROP MMLU HellaSwag VBench-style video metrics (CLIPSIM, CLIP-T, VQA-a, VQA-t, FScore) COCO image metrics (FID, sFID, CLIP, ImageReward)
Metrics: TOPS CosSim L1 RMSE accuracy FID CLIPSIM
Baselines: FlashAttention2 xformers SageAttention SageAttention2 BF16 full-precision attention
Builds on: FlashAttention SageAttention2 SageAttention
Closest papers by meaning
- SageAttention: Accurate 8-Bit Attention for Plug-and-play Inference Acceleration (ICLR 2025 · similarity 0.90)
- SageAttention2: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantization (ICML 2025 · similarity 0.89)
- FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision (NeurIPS 2024 · similarity 0.87)
- Delta Attention: Fast and Accurate Sparse Attention Inference by Delta Correction (NeurIPS 2025 · similarity 0.87)
- SpargeAttention: Accurate and Training-free Sparse Attention Accelerating Any Model Inference (ICML 2025 · similarity 0.86)
- FlashBias: Fast Computation of Attention with Bias (NeurIPS 2025 · similarity 0.86)
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.