PaperlayerView exact revision on arXiv

Source-linked AI summary

SageBwd: A Trainable Low-bit Attention

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

arXiv:2603.02170v1cs.LGcs.AI

TL;DR

Low-bit attention is effective for inference, but SageBwd previously showed a pre-training gap despite preserving fine-tuning performance. This work analyzes that gap and finds that SageBwd can match full-precision attention during pre-training under suitable conditions.

  • Problem

    SageBwd’s applicability to full pre-training was limited by a persistent performance gap relative to full-precision attention, despite successful fine-tuning performance.

  • Method

    The paper combines theoretical analysis, empirical observations, and targeted ablations to identify sources of SageBwd’s pre-training gap and its stability conditions.

  • Results

    SageBwd matches full-precision pre-training performance when tokens per step are reduced, while QK-norm stabilizes large-token training and dS is the main quantization bottleneck.

  • Takeaways & Limitations

    K-smoothing remains essential for training stability, whereas Q-smoothing provides limited pre-training benefit and can slightly degrade gradient accuracy.

  • Takeaways & Limitations

    SageBwd’s training stability degrades at very large batch sizes, motivating methods that reduce backward-pass quantization error without relying on smaller batches or increased gradient noise.

Abstract

from arXiv · show

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.

1 INTRODUCTION

Low-bit attention is promising for reducing attention costs, but its use in large-scale training remains less understood. This work analyzes SageBwd’s pre-training gap and identifies conditions that recover full-precision performance.

  • Low-bit attention reduces attention costs by enabling low-precision Tensor Cores, but training applicability remains less understood than inference effectiveness.
  • Backward computation is especially vulnerable because small-magnitude tensors and propagated forward-output errors amplify quantization deviations.
  • SageBwd quantizes six of seven attention matrix multiplications to INT8 while preserving fine-tuning performance, yet previously lagged full-precision attention during pre-training.
  • The dominant training deviation arises from the small-magnitude backward-pass dS tensor, while QK-norm stabilizes pre-training by constraining query–key outliers.
  • Reducing tokens per optimization step enables SageBwd to match full-precision pre-training performance, while K-smoothing remains necessary and Q-smoothing provides limited benefit.

2 RELATED WORK

Related work accelerates attention through hardware-efficient kernels and low-precision computation. Trainable low-bit attention extends these ideas beyond inference and fine-tuning toward full pre-training.

  • FlashAttention variants and xFormers accelerate attention through tiling, improved parallelism, warp partitioning, and custom CUDA kernels.
  • SageAttention methods use INT8 quantization and outlier-smoothing techniques, while FlashAttention3 also explores FP8 attention.
  • SageAttention3 introduces SageBwd, which quantizes most attention matrix multiplications while preserving fine-tuning performance.
  • This work builds on SageBwd by analyzing training instability and conditions for recovering full-precision performance during pre-training.

3 PRELIMINARIES

The preliminaries define attention, quantization, smoothing, and SageBwd’s mixed-precision backward design. SageBwd retains dP in FP16 while quantizing the other backward matrix multiplications with INT8.

  • Scaled dot-product attention computes S = QK⊤, P = softmax(S), and O = PV, with sequence length N and head dimension D.
  • INT8 quantization represents high-precision matrices using rounded integer values and floating-point scale factors, enabling tensor-core matrix multiplication.
  • Per-block quantization shares one scale factor across all elements in a block, such as a FlashAttention tile.
  • Q- and K-smoothing reduce channel-wise outlier effects before low-bit quantization by subtracting block-wise query and global key means.
  • SageBwd applies K-smoothing before INT8 QK⊤ quantization and uses mixed per-token or per-block quantization for the ˜PV product.
  • In the backward pass, SageBwd keeps dP = dOV⊤ in FP16 and quantizes the remaining four matrix multiplications with per-block INT8.

4 ANALYSIS OF SA G EBW D IN PRETRAINING

The analysis identifies QK-norm and tokens-per-step as key determinants of SageBwd’s pre-training behavior, with dS as the main backward-pass bottleneck. Smaller tokens-per-step improve agreement with full-precision attention, while larger activation scales worsen gradient accuracy.

  • QK-norm for logit stabilization and quantization: QK-norm controls query–key scale, stabilizes logits, and improves low-bit robustness by reducing the effective INT8 quantization step size.
  • Sensitivity of dS in the backward pass: The discrepancy between SageBwd and full-precision attention peaks at dS, with errors propagating further into dQ and dK.
  • Sensitivity of dS in the backward pass: dS becomes increasingly small for long sequences, reducing its effective signal-to-noise ratio under approximately fixed absolute INT8 noise.
  • Sensitivity of dS in the backward pass: RMS(P) ≈ 5 × 10^-3, RMS(dP) ≈ 5 × 10^-5, and RMS(dS) ≈ 1 × 10^-7 in a QK-normed checkpoint trained over 78B tokens.
  • Effect of tokens-per-step: At 2.1M tokens per step SageBwd underperforms full-precision attention, whereas at 260K it matches full-precision attention within noise.
  • Effect of tokens-per-step: Higher stochastic gradient noise at smaller tokens-per-step may make INT8 error a smaller perturbation, although other batch-size-dependent effects may also contribute.
  • Effect of tokens-per-step: The study varies tokens per step through batch size while holding sequence length fixed, leaving their interaction with quantization error for future work.
  • Effect of QK standard deviation on quantization error: At σQ,K = 10, dQ and dK show cosine similarity below 0.79 and relative ℓ2 error above 0.66, while O and dV remain relatively accurate.

5 EXPERIMENTS

Experiments show that SageBwd matches FPA at smaller tokens-per-step but requires QK-norm for stable large-TPS training. Its main backward-pass error concentrates in dS, while kernel benchmarks show substantial speedups over FlashAttention2.

  • QK-norm stability: QK-norm is necessary at 2.1M TPS because removing it causes training instability and eventual divergence.At 260K TPS, SageBwd can match FPA without QK-norm, although intermediate errors remain larger.
  • Backward-pass error: The dS tensor is the primary quantization bottleneck, with larger deviations than most intermediates and further errors propagating to dQ and dK.The pseudo-quantized FPA analysis finds O and dV close to FPA, while dS and downstream gradients diverge more substantially.
  • Kernel performance: SageBwd achieves up to a 1.67× speedup over FlashAttention2 across head dimensions D = 64 and D = 128 on RTX4090.It also exceeds Triton- and xFormers-based FlashAttention2 implementations.
  • Kernel performance: The current implementation prioritizes correctness and stability over aggressive kernel fusion, leaving further speed improvements possible through optimization.This bounds the reported kernel-performance results to the present implementation.

6 ABLATION STUDY

The ablation shows that K-smoothing is required for stable pre-training, whereas Q-smoothing offers no consistent benefit and may reduce gradient fidelity. Q-smoothing also requires a bias-gradient correction that introduces another quantization-noise pathway.

  • Experimental setup: Figure 4 compares FPA and SageBwd with no smoothing, K-smoothing, and QK-smoothing at 2.1M and 260K tokens per step.All runs use QK-norm and the Section 5 training hyperparameters; Q-smoothing alone is not evaluated.
  • K-smoothing: K-smoothing is necessary for stable pre-training, including at 260K tokens per step.It subtracts the token-wise mean of K before quantization and still achieves FPA-level performance in the more noise-tolerant regime.
  • K-smoothing: K-smoothing can be applied at kernel entry without modifying the backward pass or adding bias terms.The dQ computation remains valid because each row of dS sums to zero.
  • Q-smoothing: Q-smoothing provides no consistent improvement in pre-training loss or intermediate-tensor accuracy.With Q-smoothing enabled, dQ and dK can deviate marginally more from the FPA baseline.
  • Q-smoothing: Q-smoothing requires an additional bias branch so that the corrected dK equals dKcenter plus dKbias.The correction restores the full gradient but creates another pathway for quantization noise.

7 CONCLUSION AND FUTURE WORK

The paper concludes that SageBwd can match FPA during pre-training when QK-norm controls outliers and tokens per step are sufficiently small. Its main remaining boundary is degraded stability at very large batch sizes, motivating methods that reduce backward-pass quantization error without relying on smaller batches.

  • Conclusion: QK-norm is necessary for stability at large tokens per step, while dS is the dominant accuracy bottleneck affecting dQ and dK.Smaller tokens per step make training more tolerant to this noise; larger values expose a stable but suboptimal gap.
  • Conclusion: Smaller tokens per step enable SageBwd to match FPA performance, whereas larger tokens per step preserve a stable performance gap.The conclusion attributes the difference to training tolerance for backward-pass quantization noise.
  • Limitations and future work: SageBwd training stability degrades at very large batch sizes despite achieving FPA-level performance under moderate tokens per step.Future work targets mitigating dS-path quantization error without reduced batch size or increased gradient noise.

A.1 FORWARD PASS

The appendices describe tiled forward and backward INT8 attention algorithms that quantize inputs and selected intermediate products blockwise while retaining specified computations in higher precision. They return the attention output and its gradients after processing query, key, and value blocks.

  • Forward pass: The forward pass accepts FP16 Q, K, and V matrices, partitions them into query and key/value blocks, and applies per-block quantization.It computes blockwise scores, softmax statistics, and value products before returning O and log-sum-exp state L.
  • Forward pass: The forward algorithm computes Sij from quantized Q and K, performs per-token probability quantization, and accumulates Oij across blocks.Running maxima and row sums support the blockwise softmax normalization.
  • Backward pass: The backward pass receives quantized forward tensors, O, L, and dO, then iterates over key/value and query blocks.It computes dV, dP, dS, dQ, and dK before returning the three parameter gradients.
  • Backward pass: The backward algorithm keeps dP in FP16 while quantizing dS and the remaining backward matrix multiplications with INT8 blockwise scales.This design avoids quantizing dP before the sensitive dS computation.

B DS MAGNITUDE

The analysis bounds the magnitude of the softmax score gradient dS using the rowwise gradient dP and the probability matrix P. Because P is row-normalized, the average dS magnitude is suppressed relative to the largest per-row dP magnitude.

  • Rowwise analysis: For each row i, the analysis writes dS_i as P_i ◦ (dP_i − δ_i1) and evaluates its root-mean-square magnitude.P_i and dP_i are row vectors, and δ_i is the corresponding entry of δ.
  • Bound: Because each softmax row P_i is a probability vector with entries in [0, 1] summing to 1, its terms constrain the dS magnitude.The bound uses the infinity norm of dP_i − δ_i1.
  • Bound: The global RMS bound interprets average dS magnitude as no greater than the largest per-row dP gradient magnitude scaled by a factor involving the row length.This explains why dS can have small magnitude relative to upstream gradients.

C COSINE SIMILARITY AND REL-L2 ERROR

Figures 5 and 6 compare SageBwd with full-precision attention across layers, tokens-per-step settings, and architectural settings using cosine similarity and relative L2-error.

  • Figure 5 reports layerwise cosine similarity between SageBwd and FPA across different tokens-per-step and architectural settings.The comparisons use inputs and gradients extracted from a pretrained 325M Llama model.
  • Figure 6 reports layerwise relative L2-error between SageBwd and FPA across different tokens-per-step and architectural settings.The metric complements cosine similarity by measuring relative error magnitude.
Loading 2603.02170v1…