SageBwd: A Trainable Low-bit Attention

Jintao Zhang; Marco Chen; Haoxu Wang; Kai Jiang; Ion Stoica; Joseph E. Gonzalez; Jianfei Chen; Jun Zhu

cs.LGMar 2, 2026

SageBwd: A Trainable Low-bit Attention

Authors: Jintao Zhang, Marco Chen, Haoxu Wang, Kai Jiang, Ion Stoica, Joseph E. Gonzalez, Jianfei Chen, Jun Zhu

Organizations: Tsinghua University, UC Berkeley

Abstract

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.

Figures & tables

$\sigma_{Q},\sigma_{K}$	Output		dQ		dK		dV
$\sigma_{Q},\sigma_{K}$	CosSim	Rel- $\ell_{2}$	CosSim	Rel- $\ell_{2}$	CosSim	Rel- $\ell_{2}$	CosSim	Rel- $\ell_{2}$
1	0.9999	0.0160	0.9998	0.0184	0.9998	0.0220	0.9999	0.0159
3	0.9992	0.0389	0.9971	0.0758	0.9970	0.0777	0.9992	0.0387
5	0.9982	0.0603	0.9798	0.2014	0.9799	0.2007	0.9982	0.0605
8	0.9953	0.0972	0.8900	0.4666	0.8886	0.4699	0.9953	0.0973
10	0.9933	0.1161	0.7823	0.6648	0.7820	0.6684	0.9933	0.1157

Table 1: Sage vs. FPA across random QKV with varying

\sigma_{Q}

and

\sigma_{K}

Panel of Figure 1 — Figure 1: Pretraining loss over 78B tokens under a different number of tokens/step

Metric	$\boldsymbol{\delta}$	$\mathbf{P}$	$\mathbf{dP}$	$\mathbf{dS}$	$\mathbf{O}$	$\mathbf{dQ}$	$\mathbf{dK}$	$\mathbf{dV}$
CosSim	0.9973	0.9917	1.0000	0.9789	0.9969	0.9664	0.9537	0.9985
Rel-L2	0.0736	0.1293	0.0000	0.2045	0.0793	0.2579	0.3074	0.0540

Table 2: Cosine similarity and relative

\ell_{2}

error for intermediate tensors in SageBwd (vs. FPA).

Figure 2: Speed comparison between SageBwd and Baselines ( RTX4090 , headim=128).

Figure 3: Speed comparison between SageBwd and Baselines ( RTX4090 , headim=64).

Panel of Figure 4 — Figure 4: Ablation of Q-smoothing and K-smoothing pretraining loss over 78B tokens under a different number of tokens/step

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 5 — Figure 5: Cosine similarity between SageBwd and SDPA over layers on different settings

Panel of Figure 6 — Figure 6: Relative L2-Error between SageBwd and SDPA over layers on different settings

Attention Quantization

Attention quantization reduces the precision and memory cost of attention computations while preserving model quality.

See all 37 papers

From Attention Sensitivity to Layer Role: Revisiting Mixed-Precision Quantization of Transformers
Sep 28, 2026Nafiseh HosseinpourFardi, Negar Alihadi, Mahmoudreza Babaei +2Mixed-Precision Quantization LLM Quantization
Sharif University of Technology, Tehran, Iran · GISMA University of Applied Sciences, Potsdam, Germany · BRAINS, Brandenburg Research Center for Applied Intelligent Systems, Potsdam, Germany +3
Underscoring the Problem: Why Softpick Fails at Initialization
Oct 4, 2026Aryan Sood, Jaikaran Singh, Ishaan Bansal
IIT Roorkee
Pretraining Transformers with Quantized Softmax in Attention
Sep 27, 2026Shangzhen Zhu, Muyan Hu, Tomasz KozlowskiLanguage Model Pretraining Quantization-Aware Training
University of Illinois Urbana-Champaign

Sep 28, 2026J/K move · Enter open · S save

From Attention Sensitivity to Layer Role: Revisiting Mixed-Precision Quantization of Transformers

Nafiseh HosseinpourFardi, Negar Alihadi, Mahmoudreza Babaei, Milad Hosseini+1

Sharif University of Technology, Tehran, Iran · GISMA University of Applied Sciences, Potsdam, Germany · BRAINS, Brandenburg Research Center for Applied Intelligent Systems, Potsdam, Germany+3

Mixed-Precision Quantization LLM Quantization Attention Quantization Post-Training Quantization Neural Network Quantization Quantization

Open paper AskPDF arXiv

Most post-training quantization pipelines fit each weight matrix to its pretrained counterpart, one matrix at a time. Whether that proxy tracks what an attention block actually computes, or how errors in the Q, K and V projections compound inside the softmax, is rarely checked. We write the objective on the attention output instead, over all three projections at once, and reuse it throughout the pipeline. JAB defines one scalar loss over the joint Q, K, V weights of a block, evaluated against the block's real causally-masked attention output, and uses it twice: to fit the quantized weights (GPTQ warm start, then STE with learnable scales), and to score the block for a multiple-choice knapsack allocation. On attention-only quantization of Mistral-7B this works. At 3 bits JAB recovers 77-90% of the gap between uniform GPTQ and full precision, and its sensitivity estimate tracks an oracle costing 73 forward passes to within a fraction of a point. It stops working once MLP layers enter the allocation. A role-aware offset rule needing no sensitivity estimate at all beats JAB on GPT-2's MLP and on the full Mistral-7B model: with a 3-bit floor it quantizes 96.4% of the weights to 4.5 bits per parameter at 6.933 perplexity, within 4.4% of full precision (6.643) at 3.56x compression, against 7.158 for JAB at the same budget. Which matrix a weight sits in matters more than any sensitivity estimate we computed. Two things came out sideways. Block-local reconstruction is an unreliable proxy for end-to-end perplexity: one run improved a block's own objective 4.6x while perplexity rose 32x, which is why every allocation here is validated end-to-end. And on attention-only quantization, fine-tuning moved weights farther from their pretrained values while pulling attention outputs closer, with net gains. Post-training seems to recover attention behavior, not weights.

Abstract

Attention Quantization

Explore similar work

From Attention Sensitivity to Layer Role: Revisiting Mixed-Precision Quantization of Transformers

Underscoring the Problem: Why Softpick Fails at Initialization

Pretraining Transformers with Quantized Softmax in Attention