# SageBwd: A Trainable Low-bit Attention

arXiv: https://arxiv.org/abs/2603.02170 · HTML: https://mlsanity.com/paper/2603.02170/

## Key facts

- Cited by 8 papers in the map (data through 14 Oct 2026).
- Builds on 28 of its 34 references that are in the map.
- First posted on arXiv 2 Mar 2026.

## Builds on (most cited first, 28 of 28)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference](https://mlsanity.com/paper/1712.05877/) (2017)
- [SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models](https://mlsanity.com/paper/2211.10438/) (2022)
- [On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima](https://mlsanity.com/paper/1609.04836/) (2016)
- [Scaling Vision Transformers to 22 Billion Parameters](https://mlsanity.com/paper/2302.05442/) (2023)
- [FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision](https://mlsanity.com/paper/2407.08608/) (2024)
- [MInference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attention](https://mlsanity.com/paper/2407.02490/) (2024)
- [Don't Decay the Learning Rate, Increase the Batch Size](https://mlsanity.com/paper/1711.00489/) (2017)
- [Query-Key Normalization for Transformers](https://mlsanity.com/paper/2010.04245/) (2020)
- [SageAttention: Accurate 8-Bit Attention for Plug-and-play Inference Acceleration](https://mlsanity.com/paper/2410.02367/) (2024)
- [SpargeAttention: Accurate and Training-free Sparse Attention Accelerating Any Model Inference](https://mlsanity.com/paper/2502.18137/) (2025)
- [SageAttention2: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantization](https://mlsanity.com/paper/2411.10958/) (2024)
- [Sparse VideoGen2: Accelerate Video Generation with Sparse Attention via Semantic-Aware Permutation](https://mlsanity.com/paper/2505.18875/) (2025)
- [SageAttention3: Microscaling FP4 Attention for Inference and An Exploration of 8-Bit Training](https://mlsanity.com/paper/2505.11594/) (2025)
- [SLA: Beyond Sparsity in Diffusion Transformers via Fine-Tunable Sparse-Linear Attention](https://mlsanity.com/paper/2509.24006/) (2025)
- [A Statistical Framework for Low-bitwidth Training of Deep Neural Networks](https://mlsanity.com/paper/2010.14298/) (2020)
- [TurboDiffusion: Accelerating Video Diffusion Models by 100-200 Times](https://mlsanity.com/paper/2512.16093/) (2025)
- [Quant VideoGen: Auto-Regressive Long Video Generation via 2-Bit KV-Cache Quantization](https://mlsanity.com/paper/2602.02958/) (2026)
- [SageAttention2++: A More Efficient Implementation of SageAttention2](https://mlsanity.com/paper/2505.21136/) (2025)
- [SLA2: Sparse-Linear Attention with Learnable Routing and QAT](https://mlsanity.com/paper/2602.12675/) (2026)
- [SpargeAttention2: Trainable Sparse Attention via Hybrid Top-k+Top-p Masking and Distillation Fine-Tuning](https://mlsanity.com/paper/2602.13515/) (2026)
- [Accurate INT8 Training Through Dynamic Block-Level Fallback](https://mlsanity.com/paper/2503.08040/) (2025)
- [Geometry-Aware Rotary Position Embedding for Consistent Video World Model](https://mlsanity.com/paper/2602.07854/) (2026)
- [Identifying Sensitive Weights via Post-quantization Integral](https://mlsanity.com/paper/2503.01901/) (2025)
- [Residual Context Diffusion Language Models](https://mlsanity.com/paper/2601.22954/) (2026)

## Built on by (most cited first, 8 of 8)

- [KernelBenchX: A Comprehensive Benchmark for Evaluating LLM-Generated GPU Kernels](https://mlsanity.com/paper/2605.04956/) (2026)
- [SVG-EAR: Parameter-Free Linear Compensation for Sparse Video Generation via Error-aware Routing](https://mlsanity.com/paper/2603.08982/) (2026)
- [Vidu S2: Real-Time Interactive, Editable, and Spatial Video Generation](https://mlsanity.com/paper/2609.11638/) (2026)
- [Vidu S1: A Real-Time Interactive Video Generation Model](https://mlsanity.com/paper/2607.03118/) (2026)
- [6Bit-Diffusion: Inference-Time Mixed-Precision Quantization for Video Diffusion Models](https://mlsanity.com/paper/2603.18742/) (2026)
- [Pretraining Transformers with Quantized Softmax in Attention](https://mlsanity.com/paper/2609.33591/) (2026)
- [Delta-Matching: Closing the Final Gap of Native 8-bit Training for LLMs](https://mlsanity.com/paper/2609.37852/) (2026)
- [Backward-State Policy Is Part of the Learning Algorithm](https://mlsanity.com/paper/2609.39813/) (2026)

## Abstract (arXiv)

Low-bit attention, such as SageAttention, has emerged as an effective approach for accelerating model inference, but its applicability to training remains poorly understood. In prior work, we introduced SageBwd, a trainable INT8 attention that quantizes six of seven attention matrix multiplications while preserving fine-tuning performance. However, SageBwd exhibited a persistent performance gap to full-precision attention (FPA) during pre-training. In this work, we investigate why this gap occurs and demonstrate that SageBwd matches full-precision attention during pretraining. Through experiments and theoretical analysis, we reach a few important insights and conclusions: (i) QK-norm is necessary for stable training at large tokens per step, (ii) quantization errors primarily arise from the backward-pass score gradient dS, (iii) reducing tokens per step enables SageBwd to match FPA performance in pre-training, and (iv) K-smoothing remains essential for training stability, while Q-smoothing provides limited benefit during pre-training.