mlsanityData to 7 Oct 2026

SageBwd: A Trainable Low-bit Attention

arXiv 2603.02170 ↗First posted 2 Mar 2026cs.LG cs.AIcited by 7

Not on the map. Why?

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.

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

Builds on

28 of 34 references.

References · 28

  1. 2017Attention Is All You Needcited by 59,932
  2. 2024The Llama 3 Herd of Modelscited by 14,192
  3. 2022FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesscited by 3,611
  4. 2023FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioningcited by 2,551
  5. 2017Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inferencecited by 1,548
  6. 2022SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Modelscited by 1,338
  7. 2016On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minimacited by 1,337
  8. 2023Scaling Vision Transformers to 22 Billion Parameterscited by 690
  9. 2024FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precisioncited by 507
  10. 2024MInference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attentioncited by 375
  11. 2017Don't Decay the Learning Rate, Increase the Batch Sizecited by 313
  12. 2020Query-Key Normalization for Transformerscited by 266
16 more
  1. 2024SageAttention: Accurate 8-Bit Attention for Plug-and-play Inference Accelerationcited by 158
  2. 2024SageAttention2: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantizationcited by 113
  3. 2025SpargeAttention: Accurate and Training-free Sparse Attention Accelerating Any Model Inferencecited by 113
  4. 2025Sparse VideoGen2: Accelerate Video Generation with Sparse Attention via Semantic-Aware Permutationcited by 100
  5. 2025SageAttention3: Microscaling FP4 Attention for Inference and An Exploration of 8-Bit Trainingcited by 63
  6. 2025SLA: Beyond Sparsity in Diffusion Transformers via Fine-Tunable Sparse-Linear Attentioncited by 55
  7. 2020A Statistical Framework for Low-bitwidth Training of Deep Neural Networkscited by 51
  8. 2025TurboDiffusion: Accelerating Video Diffusion Models by 100-200 Timescited by 31
  9. 2025SageAttention2++: A More Efficient Implementation of SageAttention2cited by 21
  10. 2026SLA2: Sparse-Linear Attention with Learnable Routing and QATcited by 21
  11. 2026Quant VideoGen: Auto-Regressive Long Video Generation via 2-Bit KV-Cache Quantizationcited by 18
  12. 2026SpargeAttention2: Trainable Sparse Attention via Hybrid Top-k+Top-p Masking and Distillation Fine-Tuningcited by 16
  13. 2025Accurate INT8 Training Through Dynamic Block-Level Fallbackcited by 14
  14. 2026Geometry-Aware Rotary Position Embedding for Consistent Video World Modelcited by 12
  15. 2025Identifying Sensitive Weights via Post-quantization Integralcited by 10
  16. 2026Residual Context Diffusion Language Modelscited by 9

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2026: 7.

7 papers.

  1. 2026KernelBenchX: A Comprehensive Benchmark for Evaluating LLM-Generated GPU Kernelscited by 10
  2. 2026SVG-EAR: Parameter-Free Linear Compensation for Sparse Video Generation via Error-aware Routingcited by 9
  3. 2026Vidu S1: A Real-Time Interactive Video Generation Modelcited by 4
  4. 2026Vidu S2: Real-Time Interactive, Editable, and Spatial Video Generationcited by 4
  5. 20266Bit-Diffusion: Inference-Time Mixed-Precision Quantization for Video Diffusion Modelscited by 0
  6. 2026Pretraining Transformers with Quantized Softmax in Attentioncited by 0
  7. 2026Backward-State Policy Is Part of the Learning Algorithmcited by 0