ArXiv: 2411.10958

🎯 Pitch

Quantizing attention’s Q and K matrices down to INT4—even at the per-thread level—destroys generation quality unless you also tackle a hidden FP22 accumulator bottleneck in NVIDIA tensor cores. By jointly smoothing K/V outliers and deploying a two-level accumulation strategy that catches those internal overflows, SageAttention2 delivers a 3× speedup over FlashAttention2 and matches FP8 hardware speeds, all with negligible fidelity loss across language, image, and video models.


1. Executive Summary

SageAttention2 introduces a hardware-efficient quantized attention mechanism that accelerates the two dominant matrix multiplications in attention—QKQK^\top and P~V\tilde{P}V—by quantizing QQ and KK to INT4 using a novel per-thread quantization granularity (mapping quantization groups to GPU thread layouts per the MMA instruction), quantizing P~\tilde{P} and VV to FP8, and applying thorough outlier smoothing to QQ (subtracting a per-block mean and correcting the attention output with a compensating GEMV). The paper evaluates across language (Llama, GLM4), image (Flux, Stable-Diffusion3.5), and video generation models (CogvideoX, HunyuanVideo, Mochi) on RTX4090 and L20 GPUs, achieving approximately 3× and 4.5× speedup over FlashAttention2 and xformers respectively, while delivering a 1.8× end-to-end speedup on CogvideoX with the 8-bit variant incurring negligible metric loss. The paper further establishes that quantization accuracy depends critically on addressing an undocumented FP22 accumulator limitation in NVIDIA tensor cores via a two-level accumulation strategy, demonstrating that even aggressive INT4 quantization preserves end-to-end generation quality only when both proposal-space precision (per-thread granularity, Q/K smoothing) and accumulation-space precision (FP32 buffer for FP8 matmul) are jointly managed.

2. Context and Motivation

The Core Problem: Attention Is the Bottleneck, But Hardware-Efficient Quantization Remains Unexploited

The fundamental problem SageAttention2 addresses is deceptively simple: how do we make the attention mechanism faster without degrading model quality? The two matrix multiplications at the heart of self-attention—QKQK^\top and P~V\tilde{P}V—each have O(N2d)O(N^2 d) complexity, where NN is the sequence length. As sequence lengths grow into the tens or hundreds of thousands of tokens in production language, image, and video generation models, these operations dominate end-to-end latency. For instance, generating a single video with CogvideoX (1.5-5B) takes 1040 seconds on an RTX4090 using FlashAttention2 (Table 8). The attention kernels are the primary obstacle to real-time or interactive deployment.

This matters for three reasons the paper makes explicit (Section 1):

  • Quadratic scaling is inescapable for exact attention. Unlike the O(N)O(N) linear operations (feed-forward layers, normalization) that grow gracefully with sequence length, attention's quadratic cost makes it the dominant term in long-sequence regimes. An 8×8\times increase in sequence length means a 64×64\times increase in attention FLOPs. This is not an implementation inefficiency—it is intrinsic to the operation.
  • Hardware offers underutilized low-precision compute. Modern GPUs (RTX 40-series, Hopper) have tensor cores that execute INT4 and FP8 matrix multiplications at multiples of FP16/FP32 throughput. Specifically, INT4 Matmul achieves double the speed of INT8 (Table 1), and FP8 Matmul with FP32 accumulators achieves a 2×2\times speedup over FP16 on the RTX4090, L40, L20, and H100 GPUs. Yet quantized attention had not successfully exploited these faster formats while maintaining accuracy.
  • Existing fast attention methods either sacrifice output quality or are hardware-restricted. Linear attention and sparse attention methods (Section 1) reduce computational cost by approximating or selectively computing the attention matrix, but they are "only suitable for a limited range of models and tasks"—they change the semantics of the attention operation rather than just its implementation. This means model developers cannot use them as drop-in replacements without retraining or accepting accuracy loss on certain benchmarks.

The paper's central observation is that quantization—specifically, casting the operations within attention to lower-bit formats—is the most direct path to hardware-efficient speedup without algorithmic approximation, but that prior attempts either didn't go far enough (SageAttention stopped at INT8) or failed entirely (naive INT4 produces garbage). The gap is between the theoretical throughput available from INT4/FP8 tensor cores and the practical accuracy achievable when quantizing the dynamic, outlier-laden tensors that arise in attention.

The Limitations of SageAttention (The Direct Predecessor)

SageAttention (Zhang et al., 2025c) was the first quantized attention method to achieve plug-and-play acceleration with negligible end-to-end metric loss across language, image, and video models. It demonstrated a 2×2\times speedup over FlashAttention2 by quantizing QQ and KK to INT8 at per-block granularity and using FP16 Matmul with FP16 accumulators for P~V\tilde{P}V. However, the paper identifies two specific weaknesses that prevent further acceleration (Section 1, paragraph "Motivation"):

Weakness 1 (W1): INT8 leaves half the throughput on the table. INT4 tensor cores are twice as fast as INT8 tensor cores on consumer Ada Lovelace GPUs (RTX 40-series). SageAttention's INT8 quantization for QKQK^\top achieves only half the theoretical speed of what INT4 could deliver. The paper quantifies this directly in Figure 5: SageAttn2-4b achieves approximately 481 TOPS on RTX4090 at sequence length 32K (head_dim=128, causal=False), compared to SageAttention's approximately 338 TOPS—a roughly 1.4× improvement attributable largely to moving from INT8 to INT4.

Weakness 2 (W2): FP16 accumulator speedup is GPU-specific. SageAttention's speedup for the P~V\tilde{P}V matmul relies on reducing the accumulator precision from FP32 to FP16, which NVIDIA's tensor cores support as a faster mode. However, Table 1 shows this is effective only on the RTX4090 and RTX3090—on the L40, L20, and H100, FP16 with FP16 accumulators provides no speedup (1×) over FP16 with FP32 accumulators. This makes SageAttention's P~V\tilde{P}V acceleration unavailable on the most important datacenter GPUs (L40, L20, H100) and the newer Hopper architecture. To generalize the speedup, the paper argues, one must use FP8 Matmul with FP32 accumulators, which provides a consistent 2×2\times speedup across all modern GPU architectures (Table 1).

Why Direct INT4 Quantization Fails

The paper is motivated not just by the opportunity for speedup, but by the genuine technical difficulty of achieving it. Converting QQ and KK to INT4 and P~\tilde{P} and VV to FP8 is not a straightforward extension of SageAttention's INT8 recipe—it produces catastrophic failures. The paper provides concrete evidence:

"when only per-tensor quantizing Q, K to INT4, the text-to-video model CogvideoX will generate a completely blurry video, and Llama3 only achieves a random-guessing-level accuracy of 25% on MMLU"

This is not a minor degradation—it is complete model collapse. The paper then decomposes why into three specific challenges (Section 1, "Challenges"):

Challenge 1 (C1): INT4's restricted numerical range amplifies outlier sensitivity. The INT4 representable range is [-7, +7], giving only 16 quantization levels compared to INT8's 256 levels. For per-tensor or per-block quantization, the scale factor is determined by the maximum absolute value in the quantization group. The paper explains the failure mechanism precisely: "any element will be quantized to zero if it is more than 14 times (0.5 vs 7) smaller than the largest element in the group" (Section 3.1). When QQ or KK contains outliers—values much larger than the typical element—the scale factor is driven up by the outlier, and most normal elements get squeezed into the zero bin. The paper observes that QQ, KK, and VV all exhibit significant channel-wise outliers (Figure 2 shows heatmap distributions from Llama3.1 and CogvideoX), making this a practical rather than theoretical concern.

Challenge 2 (C2): The FP8 MMA accumulator is not actually FP32. The paper makes a striking empirical discovery: when investigating why FP8 quantization of P~V\tilde{P}V showed accuracy degradation in real CUDA despite simulating correctly, they found that the mma.f32.f8.f8.f32 instruction on Ada and Hopper architectures uses an FP22 accumulator (1 sign bit, 8 exponent bits, 13 mantissa bits) rather than true FP32 (1 sign bit, 8 exponent bits, 23 mantissa bits). The paper describes their diagnostic procedure:

"when D is initialized with more than 13 mantissa bits, the value of C is equal to D with its least significant 10 mantissa bits zeroed out (i.e., truncated)"

This means that the FP8 matrix multiplication loses 10 bits of mantissa precision relative to what a user who writes mma.f32.f8.f8.f32 might reasonably expect. The accumulator physically drops precision during the accumulation of partial products. For P~V\tilde{P}V, where many small elements of P~\tilde{P} multiply elements of VV and must be accurately summed, this truncation introduces non-trivial error, particularly when many small terms need to accumulate to a meaningful total.

The Practical Landscape: Why This Matters Now

The paper positions SageAttention2 within a broader ecosystem of attention acceleration methods that reveals why this work is timely. The existing approaches form a spectrum:

Hardware-optimized exact attention (FlashAttention family). FlashAttention V1/V2/V3 and xformers restructure the attention computation to exploit GPU memory hierarchy—tiling QQ, KK, VV along the token dimension, fusing the softmax into the tiled reduction, and using online softmax to avoid materializing the full N×NN \times N attention matrix. These methods compute exact attention (no approximation), making them universally applicable. However, they do not exploit quantization—they operate in FP16/BF16 throughout. As a result, their speed is bounded by the throughput of FP16 tensor core operations. FlashAttention2 achieves approximately 164 TOPS on an RTX4090 at 32K sequence length (Figure 5, head_dim=128, causal=False). FlashAttention3 adds hardware-specific optimizations for the Hopper architecture (asynchrony, low-precision in FP8), achieving the best performance on H100 GPUs but only running on Hopper—it is not portable to the widely-deployed RTX40 series or datacenter L-series GPUs.

Linear and sparse attention. Methods like Linformer, Performer, and various sparse attention schemes reduce the O(N2)O(N^2) complexity to O(N)O(N) or O(NlogN)O(N \log N) by approximating or pruning the attention matrix. The paper acknowledges these methods but notes they fundamentally change the model, making them not drop-in replacements: "these methods are only suitable for a limited range of models and tasks" (Section 1). A video generation model designed with dense, full attention may produce degraded results if attention is approximated. This limits their adoption to scenarios where retraining or task-specific tuning is feasible.

SageAttention (INT8 quantization). This is the direct predecessor—the only prior work that achieved quantized attention with plug-and-play accuracy. But as discussed above, it leaves INT4 and FP8 throughput unexploited.

SageAttention2 positions itself in the intersection of hardware optimization and quantization—like FlashAttention in being a drop-in replacement for exact attention, but like SageAttention in exploiting low-precision compute, and going significantly further in quantizing to INT4/FP8 while solving the accuracy problems that prevented prior work from doing so.

A Reconciliation of Prior Quantization Techniques and Why They Don't Transfer to Attention

The paper also addresses a less obvious motivation: standard quantization techniques developed for linear layers in Transformers do not directly apply to attention. Section 3.1 contains a critical observation:

"Classical techniques to improve the activation-weight MM, such as per-channel quantization, or SmoothQuant are not applicable here for the query-key MM in attention."

The reason is structural. In a standard linear layer y=xWy = xW, the quantization is between an activation xx (with per-token or per-channel statistics) and a weight WW (static, with per-channel statistics). SmoothQuant (Xiao et al., 2023) works by transferring the quantization difficulty between activations and weights via a per-channel scaling factor: it smooths the activations at the cost of making the weights slightly harder to quantize, but since weights are static, this can be absorbed offline. For QKQK^\top in attention, both QQ and KK are dynamic activations with significant outliers (Figure 2). There is no static weight to absorb the quantization burden. Per-channel quantization also fails because the outer dimension of QKQK^\top is the token dimension—the quantization must be applied along the token axis, not the channel axis, to be compatible with the matrix multiply layout.

This means that attention quantization requires purpose-built techniques that account for the unique properties of attention tensors: both operands are dynamic, both have outliers, and the quantization granularity must be compatible with the hardware's matrix multiply instructions. The paper's smoothing and per-thread quantization methods are designed specifically for this regime, unlike generic quantization methods that were developed for weight-activation products where one operand is static.

How the Paper Positions Its Contributions

The paper does not claim to invent the idea of quantized attention—SageAttention did that. Its contribution is in solving the accuracy-constrained throughput maximization problem for attention: finding the set of techniques that allow the fastest possible quantized formats (INT4, FP8) to be used without degrading model outputs. This is framed as addressing specific, identified failure modes (C1, C2) rather than proposing a wholly new paradigm.

The paper's position can be summarized as: SageAttention showed that INT8 quantized attention works; SageAttention2 shows that INT4 quantized attention works, but only with (a) per-thread quantization that matches GPU thread layouts, (b) outlier smoothing applied to both QQ and KK, and (c) explicit management of the undocumented FP22 accumulator limitation via two-level accumulation. Each technique addresses a specific numerical failure mode that would otherwise cause catastrophic accuracy loss. The paper validates this position by showing that ablating any of these techniques causes significant degradation (Tables 4, 6, 10), and that the full combination preserves end-to-end metrics across a diverse model zoo (Table 2).

3. Technical Approach

3.1 Reader Orientation

SageAttention2 is a quantized CUDA kernel implementation for the exact attention mechanism—not a new attention algorithm, but a drop-in replacement that computes the same mathematical operation (attention) using lower-precision arithmetic to exploit faster GPU tensor core instructions. The system solves the problem of how to use 4-bit integer (INT4) and 8-bit floating-point (FP8) matrix multiplications within attention while preserving output quality, which previous attempts failed at because (a) INT4's tiny numerical range makes it exquisitely sensitive to outlier values in QQ and KK, and (b) an undocumented FP22 accumulator in NVIDIA's FP8 tensor core hardware silently truncates precision during P~V\tilde{P}V accumulation. The solution is a layered set of techniques—per-thread INT4 quantization that maps quantization groups to GPU thread layouts, outlier smoothing on QQ via per-block mean subtraction, and a two-level accumulation strategy that captures FP8 partial results in an FP32 buffer—that jointly solve these accuracy problems while delivering 3× the throughput of FlashAttention2 on consumer GPUs and matching FlashAttention3's speed on datacenter GPUs with better quality.

3.2 Big-Picture Architecture (Diagram in Words)

The SageAttention2 kernel processes attention in a single fused CUDA operation that follows FlashAttention's tiled computation pattern but replaces the internal matrix multiplications with quantized versions. The architecture has five major stages, executed for each output tile:

  1. Preprocessing kernel (off-chip, fused reads): Load a tile of QQ (size bq×db_q \times d), subtract the per-block mean qˉi\bar{q}_i from each token, apply per-thread INT4 quantization to the zero-centered γ(Qi)\gamma(Q_i), and simultaneously compute a GEMV (general matrix-vector product) ΔSij=qˉiγ(Kj)\Delta S_{ij} = \bar{q}_i \gamma(K_j)^\top for later correction. Load a tile of KK (size bk×db_k \times d), subtract the global mean kˉ\bar{k}, and apply per-thread INT4 quantization. Load a tile of VV (size bk×db_k \times d) and apply per-channel FP8 quantization.

  2. INT4 QKQK^\top Matmul (on-chip): Execute the mma.m16n8k64 PTX instruction on the quantized INT4 Q^i\hat{Q}_i and K^j\hat{K}_j tiles within each GPU warp. Each thread in the warp dequantizes its fragment using a single (δQ,δK)(\delta_Q, \delta_K) pair (enabled by the per-thread quantization layout), multiplies, and adds the precomputed ΔSij\Delta S_{ij} correction vector to produce the full-precision attention scores SijS_{ij}.

  3. Online softmax (on-chip): Compute row-wise max, exponentiate, and accumulate row sums using the standard FlashAttention online softmax procedure, maintaining running statistics mijm_{ij} and lijl_{ij}. The unnormalized output P~ij=exp(Sijmij)\tilde{P}_{ij} = \exp(S_{ij} - m_{ij}) is immediate—no quantization here, but the values are inherently bounded in [0,1][0, 1].

  4. FP8 P~V\tilde{P}V Matmul with FP22-to-FP32 accumulation (on-chip): Cast P~ij\tilde{P}_{ij} to FP8 E4M3 format (multiplying by the static scale δP=1/448\delta_P = 1/448 to fill the representable range) and multiply with the quantized FP8 V^j\hat{V}_j using mma.f32.f8.f8.f32. The hardware accumulator for this instruction is actually FP22 (1 sign, 8 exponent, 13 mantissa bits), so the result Rij=P~ijVjR_{ij} = \tilde{P}_{ij} V_j has limited precision. Critically, RijR_{ij} is accumulated over only bkb_k rows (bk=64b_k = 64), and then copied to a true FP32 buffer OijO_{ij} that persists across iterations of the jj loop. This two-level accumulation prevents the FP22 truncation error from compounding over many tiles.

  5. Output correction (on-chip, final): After the jj loop completes, divide Oi,TnO_{i, T_n} (which is in FP32) by the softmax denominator li,Tnl_{i, T_n}, multiply by the per-channel VV dequantization scale δV\delta_V, and write the output tile OiO_i to off-chip memory.

The flow is: [Load, subtract means, quantize Q/K/V, compute ΔS][INT4 QK^T + ΔS correction][Softmax][FP8 P̃V with FP22→FP32 buffering][Dequantize, divide by softmax sum, write]. This is a single fused kernel; the preprocessing is a separate kernel that writes the quantized and smoothed tensors to GPU registers or shared memory for consumption by the attention kernel.

3.3 Roadmap for the Deep Dive

  • First, the smoothing of QQ (Section 3.1), because it is the prerequisite that makes INT4 quantization numerically viable—we need to understand the outlier problem before we can understand why the quantization granularity matters.
  • Second, per-thread INT4 quantization (Section 3.2), because it addresses the remaining precision gap after smoothing and is the most novel hardware-software co-design element—the mapping between quantization groups and GPU thread layouts.
  • Third, FP8 quantization for P~V\tilde{P}V (Section 3.3), because it requires a fundamentally different approach from the INT4 QKQK^\top quantization—P~\tilde{P} has a distinct value distribution that makes FP8 the right format and per-channel VV quantization the right granularity.
  • Fourth, the two-level accumulation strategy (Section 3.4), because it solves an unexpected hardware limitation (FP22 accumulator) that would otherwise silently corrupt the P~V\tilde{P}V matmul output—a discovery that has implications beyond attention.
  • Fifth, the optional VV smoothing technique (Section 3.4), because it provides additional precision for models where VV exhibits channel-wise bias, completing the set of available accuracy-preserving techniques.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily a systems paper whose core idea is that attention can be accelerated by quantizing its internal matrix multiplications to INT4 and FP8, but that doing so requires a carefully co-designed set of outlier-management and precision-preservation techniques—smoothing, per-thread quantization mapping, and two-level accumulation—that jointly prevent the catastrophic accuracy failures observed in naive quantization.


Smoothing QQ to Make INT4 Quantization Possible

The paper identifies that the primary obstacle to INT4 quantization of QQ and KK is the presence of outliers—values that are much larger in magnitude than the typical elements in their quantization group. Section 3.1 explains the failure mechanism with a precise numerical argument: given the INT4 range [-7, +7], any element whose absolute value is less than 1/141/14 of the maximum absolute value in its quantization group will be rounded to zero after scaling and rounding. If a group contains an outlier that is 140× larger than the typical element magnitude, that typical element is 10× smaller than the threshold and gets zeroed. Since attention scores are computed as dot products, zeroing elements of QQ and KK destroys the fidelity of the attention pattern.

The paper observes a structural property of attention tensors that enables a solution: "Q, K for all tokens are actually highly similar, with only small variations between different tokens" (Section 3.1). Figure 2 visualizes this: heatmaps of QQ, KK, and VV from Llama3.1 and CogvideoX show that QQ and KK exhibit strong per-channel biases—each channel dimension has a roughly consistent mean value across tokens, with token-to-token variation being relatively small. VV, by contrast, shows per-channel outliers without the same token-similarity structure.

SageAttention (the predecessor) exploited this observation for KK only: it subtracted the per-token mean from KK to reduce the magnitude of the values being quantized. The insight of SageAttention2 is that QQ can be smoothed similarly, but with a per-block mean rather than a per-token mean, because QQ is processed in tiles along the token dimension in FlashAttention's tiling scheme.

The mathematical decomposition proceeds as follows. For a query tile QiQ_i (containing bqb_q tokens) and all key tokens KK, define:

γ(Qi)=Qiqˉi,γ(Kj)=Kjkˉ\gamma(Q_i) = Q_i - \bar{q}_i, \quad \gamma(K_j) = K_j - \bar{k}

where qˉi=mean(Qi)\bar{q}_i = \text{mean}(Q_i) is a 1×d1 \times d vector (the average of the bqb_q query token embeddings in the tile, along the token axis), and kˉ=mean(K)\bar{k} = \text{mean}(K) is a 1×d1 \times d vector (the average of all NN key token embeddings, broadcast across tiles). Both means are computed along the token dimension within their respective scopes (block for QQ, global tensor for KK).

What this computes: For each tile of queries, subtract its tile-mean from every query token in that tile; for the keys, subtract the global key mean from every key token. This produces zero-centered versions γ(Qi)\gamma(Q_i) and γ(Kj)\gamma(K_j) where the per-channel biases have been removed, leaving only the token-to-token variation that actually carries the attention signal.

The critical expansion is the decomposition of the attention score:

Sij=QiKj=(qˉi+γ(Qi))(kˉ+γ(Kj))S_{ij} = Q_i K_j^\top = (\bar{q}_i + \gamma(Q_i))(\bar{k} + \gamma(K_j))^\top =qˉikˉ+qˉiγ(Kj)+γ(Qi)kˉ+γ(Qi)γ(Kj)= \bar{q}_i \bar{k}^\top + \bar{q}_i \gamma(K_j)^\top + \gamma(Q_i) \bar{k}^\top + \gamma(Q_i) \gamma(K_j)^\top

The paper then makes a crucial observation: the first and third terms, qˉikˉ\bar{q}_i \bar{k}^\top and γ(Qi)kˉ\gamma(Q_i)\bar{k}^\top, are constant across the row dimension of SijS_{ij}. Specifically, qˉikˉ\bar{q}_i \bar{k}^\top is a scalar (dot product of two dd-vectors) and (γ(Qi)kˉ)(\gamma(Q_i)\bar{k}^\top) is an bq×1b_q \times 1 vector (one value per query token). Adding a constant to every column of a row of the attention score matrix SS does not change the softmax output, because:

softmax(Sij+c)=exp(Sij+c)kexp(Sik+c)=exp(c)exp(Sij)exp(c)kexp(Sik)=softmax(Sij)\text{softmax}(S_{ij} + c) = \frac{\exp(S_{ij} + c)}{\sum_k \exp(S_{ik} + c)} = \frac{\exp(c)\exp(S_{ij})}{\exp(c)\sum_k \exp(S_{ik})} = \text{softmax}(S_{ij})

where cc is any scalar added uniformly to all elements of row ii. The paper leverages this invariance: the bias terms can be dropped entirely, and only the cross-term qˉiγ(Kj)\bar{q}_i \gamma(K_j)^\top and the smoothed product γ(Qi)γ(Kj)\gamma(Q_i) \gamma(K_j)^\top need to be computed.

The final actionable decomposition is:

Sij=γ(Qi)γ(Kj)+ΔSij,whereΔSij=qˉiγ(Kj)S_{ij} = \gamma(Q_i)\gamma(K_j)^\top + \Delta S_{ij}, \quad \text{where} \quad \Delta S_{ij} = \bar{q}_i \gamma(K_j)^\top

Here, γ(Qi)γ(Kj)\gamma(Q_i)\gamma(K_j)^\top is the INT4-quantized main term—the product of the smoothed (and therefore small-magnitude, outlier-free) QQ and KK matrices. ΔSij\Delta S_{ij} is a correction vector of shape bq×bkb_q \times b_k that can be computed in FP16 as a GEMV (general matrix-vector multiplication): qˉi\bar{q}_i (a 1×d1 \times d vector) multiplied by γ(Kj)\gamma(K_j)^\top (a d×bkd \times b_k matrix). The GEMV is cheap (O(bqdbk)O(b_q \cdot d \cdot b_k) operations with dd small, typically 64 or 128) compared to the GEMM (O(bqdbk)O(b_q \cdot d \cdot b_k) with the same complexity but executed in INT4 at much higher throughput).

Why this form: The key property is that γ(Qi)\gamma(Q_i) and γ(Kj)\gamma(K_j) have much smaller magnitudes than QiQ_i and KjK_j, because the dominating per-channel biases have been subtracted. The paper provides a theoretical analysis in Appendix A.5: assuming QQ tokens follow a Gaussian distribution N(μ,Σ2)N(\mu, \Sigma^2) and are i.i.d., the smoothed version Yij=Xij1Nk=1NXkjY_{ij} = X_{ij} - \frac{1}{N}\sum_{k=1}^N X_{kj} has mean 0 and variance N1Nσj2\frac{N-1}{N}\sigma_j^2, meaning it is more tightly concentrated around zero. Since quantization error under round-to-nearest is proportional to the maximum absolute value in the quantization group (Section 3.1: "the expected quantization error is 122M2b\frac{1}{2} \cdot \frac{2M}{2^b}" for a group with absolute max MM and bit-width bb), reducing the magnitude directly reduces the error. The correction term ΔSij\Delta S_{ij} is computed in FP16 (not quantized) and added back after dequantization, so the information lost by zero-centering is restored exactly.

Design choices and alternatives rejected. The paper explicitly compares its smoothing approach against two established techniques for activation quantization, both of which fail for attention:

  • SmoothQuant (Xiao et al., 2023): SmoothQuant works by computing a per-channel smoothing factor α\alpha such that the quantization difficulty is balanced between activations and weights: X^=Xdiag(α)1\hat{X} = X \cdot \text{diag}(\alpha)^{-1}, W^=diag(α)W\hat{W} = \text{diag}(\alpha) \cdot W. For attention, both QQ and KK are activations—there is no static weight matrix to absorb the smoothing factor. If you smooth QQ at the cost of making KK harder to quantize (or vice versa), you've just moved the problem rather than solving it, because both need to be quantized. The paper's method instead smooths QQ using its own per-block mean and KK using its global mean, which does not transfer difficulty—each is smoothed independently using statistics derived from itself.

  • Hadamard transformation (Quarot; Ashkboos et al., 2024): Applying a random orthogonal (Hadamard) rotation to QQ and KK before quantization can spread outlier energy across dimensions, reducing the maximum per-element magnitude. The paper evaluates this (as "HadmdAttn" in Table 17) and finds it achieves only 4.85% worst-case cosine similarity with INT4 quantization—essentially no improvement over naive quantization (4.83%). The smoothing approach achieves 96.71% worst-case cosine similarity (Table 17). The reason is that the Hadamard rotation is a fixed transformation that does not adapt to the specific outlier structure of attention tensors; it redistributes energy uniformly, but if the outliers are very strongly concentrated, even the redistributed energy may still produce large values.

  • SageAttention's K-only smoothing: The predecessor only smoothed KK, leaving QQ's outliers unaddressed. Table 5 quantifies the impact: with smoothing only KK and INT4 quantization, Llama3.1 Lambda accuracy drops from 81.5% (full precision) to 72.6%; with smoothing both QQ and KK, it recovers to 80.8%. The paper's ranking of effectiveness (Table 4, Table 17) is: Smooth Q+KQ+K > Smooth QQ > Smooth KK > SmoothAttn > HadmdAttn > None. Smoothing QQ alone is better than smoothing KK alone, because QQ's per-block statistics are more localized and the block-wise mean subtraction is more effective at removing within-block outlier structure.

Practical integration. The smoothing, quantization, and ΔS\Delta S computation are fused into a single preprocessing kernel that reads QQ and KK from off-chip memory only once (Algorithm 1, "Preprocessing" comment). This is important because DRAM bandwidth is the primary bottleneck in attention kernels (the FlashAttention line of work is fundamentally about minimizing off-chip memory traffic). The smoothing operation is element-wise subtraction, which is bandwidth-bound but effectively free when fused with the memory read.

Empirical validation. Figure 20 (Appendix A.9) shows a histogram of quantized QQ values from CogvideoX with and without smoothing: without smoothing, the distribution is heavily concentrated at a few quantization levels near zero (many elements are being rounded to zero), while with smoothing, the histogram spans the full [-7, +7] range much more uniformly. This visual evidence confirms the theoretical argument: smoothing prevents the "everything collapses to zero" failure mode.


Per-Thread INT4 Quantization: Aligning Quantization Groups with GPU Thread Layouts

Even with smoothing, INT4 quantization remains challenging because the 4-bit format has only 16 representable values. The paper observes that the quantization granularity—how many elements share a common scale factor—is the second critical lever for accuracy, independent of smoothing. If the quantization granularity is too coarse, a single outlier in a large group drives up the scale factor and zeros out many normal elements in that group, even if the outlier is relatively small after smoothing. If the granularity is too fine (per-token), the hardware efficiency degrades because each GPU thread must handle multiple scale factors during dequantization, adding overhead to the inner loop.

The paper proposes per-thread quantization: a granularity where quantization groups are defined to align with the memory layout of the GPU's warp-level matrix multiply instruction (mma.m16n8k64), such that each GPU thread in a warp is responsible for dequantizing elements that all share a single scale factor. This achieves accuracy close to per-token quantization without the dequantization overhead.

The problem with per-token quantization. To understand why per-thread quantization is necessary, consider what happens with per-token quantization of QQ and KK. For QKQK^\top, each element of the output requires multiplying one row of QQ (a token) by one column of KK^\top (also a token, after transposition). If QQ and KK are per-token quantized, then element (i,j)(i, j) of the quantized product is:

Q^[i,:]K^[j,:]×δQ[i]δK[j]\hat{Q}[i, :] \hat{K}[j, :]^\top \times \delta_Q[i] \cdot \delta_K[j]

The dequantization requires multiplying the per-token scale δQ[i]\delta_Q[i] with the per-token scale δK[j]\delta_K[j]. When this is distributed across GPU threads in a warp, each thread may need to access different δQ\delta_Q and δK\delta_K values depending on which output elements it computes. In the worst case, each thread needs to perform a dot product of two vectors of per-token scales, which is both computationally expensive and requires loading many scale values from registers or shared memory.

SageAttention (the INT8 predecessor) avoided this by using per-block quantization: each tile QiQ_i (with bqb_q tokens, typically 128) gets a single scale δQi\delta_{Q_i}, and each tile KjK_j (with bkb_k tokens, typically 64) gets a single scale δKj\delta_{K_j}. This means all threads in a warp computing QiKjQ_i K_j^\top use the same δQ\delta_Q and δK\delta_K—the dequantization is a single scalar multiply. For INT8, this provided adequate accuracy. For INT4, the paper finds it insufficient (Table 6: per-block INT4 achieves only 98.03% average cosine similarity vs. 99.45% for per-token INT4 on CogvideoX).

The per-thread quantization mapping. The key insight is that the mma.m16n8k64 PTX instruction (which computes the product of a 16×6416 \times 64 matrix AA and a 64×864 \times 8 matrix BB, accumulating into a 16×816 \times 8 matrix CC using 32 threads in a warp) has a specific, documented layout of which thread owns which elements of the output matrix. The paper exploits this by ensuring that all elements held by a single thread belong to the same quantization group. The detailed mapping is given in Equation 8 (Appendix A.6) and illustrated in Figure 4.

For a typical configuration with block sizes bq=128b_q = 128, bk=64b_k = 64, and cw=4c_w = 4 warps per streaming multiprocessor (as used in FlashAttention2), the layout works as follows:

  • QiQ_i (128 tokens) is split into cw=4c_w = 4 segments QwQ_w (32 tokens each), each processed by one warp.
  • Within a warp, the 32 query tokens are mapped to threads according to the MMA output layout. Figure 4 (left) shows that, for the mma.m16n8k64 instruction, each thread in the warp handles a specific subset of the 16×816 \times 8 output tile. The thread mapping is: thread TtT_t computes output elements in specific rows (tokens) and columns based on the MMA instruction specification.
  • For QQ, the quantization groups are defined as: Qw[8k+i]Q_w[8k + i] for i=0,1,,7i = 0, 1, \ldots, 7 share one scale δQ\delta_Q, where kk indexes which set of 8 tokens within the 32-token warp segment. This creates 32 quantization groups per 128-token Q block (4 warps×8 groups per warp4 \text{ warps} \times 8 \text{ groups per warp}), compared to 1 group for per-block quantization—a 32× finer granularity.
  • For KK, the quantization groups are defined as: Kj[8k+2i]K_j[8k + 2i] and Kj[8k+2i+1]K_j[8k + 2i + 1] share one scale δK\delta_K, for i=0,1,2,3i = 0, 1, 2, 3. This creates 4 quantization groups per 64-token K block, compared to 1 group for per-block quantization—a 4× finer granularity.

Why the asymmetric granularity (32× for Q, 4× for K)? The granularity is determined by the MMA instruction's memory layout for matrices AA and BB. For mma.m16n8k64, AA is 16×6416 \times 64 (rows are the "M" dimension—query tokens) and BB is 64×864 \times 8 (columns are the "N" dimension—key tokens). The K-reduction dimension (64) is the inner dimension. The instruction specifies that each thread holds specific elements of AA and BB, and the grouping is determined by which elements share a common row or column in the output. Since QQ contributes to the MM dimension (more rows) and KK contributes to the NN dimension (fewer columns), the resulting quantization group counts differ: more groups for QQ because the thread distribution across the MM dimension creates more distinct groups.

What happens during dequantization. Consider one warp computing QwKjQ_w K_j^\top, producing a 32×6432 \times 64 tile of the attention score matrix. Each of the 32 threads in the warp produces a fragment of the output. For a given thread tt, let ItI_t be the set of QQ tokens (row indices) and JtJ_t be the set of KK tokens (column indices) that this thread computes. Because of how the quantization groups are defined:

  • All tokens in ItI_t belong to the same QQ quantization group (by construction of the grouping).
  • All tokens in JtJ_t belong to the same KK quantization group (by construction of the grouping).

Therefore, the dequantization for thread tt's entire output fragment is a single scalar multiply: multiply the accumulated INT32 dot products by δQ[t]δK[t]\delta_Q[t] \cdot \delta_K[t] where δQ[t]\delta_Q[t] and δK[t]\delta_K[t] are the single scale values for this thread's groups. Each thread loads exactly two scale values and performs one multiply. This is zero additional overhead compared to per-block quantization (which also uses one multiply per thread), while providing 32× (for QQ) and 4× (for KK) finer granularity.

What per-token quantization would require in contrast. With per-token quantization, each element (i,j)(i, j) of the output has a different (δQ[i],δK[j])(\delta_Q[i], \delta_K[j]) pair. A thread computing multiple output elements would need to load multiple scale pairs and perform multiple dequantization multiplies, which increases register pressure and instruction count in the inner loop. Table 19 quantifies this: per-token quantization achieves 268 TOPS on L20 GPU compared to 283 TOPS for per-thread quantization—a ~5.3% throughput reduction due to the additional dequantization overhead.

Empirical validation. Table 6 and Table 15 (Appendix A.9) show the average and worst-case accuracy across all layers of CogvideoX for different quantization granularities, with QQ and KK both smoothed and quantized to INT4 (P~\tilde{P} and VV in FP16 for this isolation experiment):

GranularityAverage CosSimWorst CosSimAverage RMSEWorst RMSE
Per-token99.45%96.76%0.03350.0775
Per-thread99.45%96.72%0.03130.0776
Per-block98.03%90.68%0.07440.1490
Per-tensor97.15%85.85%0.08650.2261

The critical finding: per-thread quantization achieves accuracy essentially identical to per-token quantization (99.45% vs 99.45% average CosSim; 96.72% vs 96.76% worst CosSim), while significantly outperforming per-block (98.03% average, 90.68% worst). The RMSE values tell the same story: 0.0313 for per-thread vs 0.0335 for per-token (per-thread is actually slightly better—possibly because the grouping introduces a beneficial regularization effect), while per-block is 0.0744 (more than double the error). The worst-case CosSim drop from per-block (90.68%) to per-tensor (85.85%) illustrates why coarse granularity is unacceptable: at 85.85% cosine similarity on the worst layer, the attention output would be substantially corrupted.

Integration with the kernel. Algorithm 1 shows how per-thread quantization integrates into the SageAttention2 kernel. In the preprocessing stage, QiQ_i and KjK_j are quantized according to the per-thread grouping (the ψ_Q and ψ_K quantizers), producing the quantized tensors Q^i\hat{Q}_i and K^j\hat{K}_j and scale arrays δQ\delta_Q and δK\delta_K. In the attention kernel, the dequantization step is:

Sij[st:st+cw]=ψδQδK1(Matmul(Q^i[st:st+cw],K^j))S_{ij}[\text{st} : \text{st} + c_w] = \psi^{-1}_{\delta_Q \delta_K}(\text{Matmul}(\hat{Q}_i[\text{st} : \text{st} + c_w], \hat{K}_j^\top))

where ψδQδK1\psi^{-1}_{\delta_Q \delta_K} multiplies each thread's output fragment by that thread's single δQδK\delta_Q \cdot \delta_K product, and st indexes the segment of QiQ_i assigned to a particular warp. The Matmul operation is the INT4 mma.m16n8k64 instruction, which reads the packed INT4 Q^\hat{Q} and K^\hat{K} values from registers, performs the dot products, and accumulates into INT32 registers.


FP8 Quantization for P~V\tilde{P}V: Choosing the Right Format for a Unique Distribution

The second matrix multiplication in attention, P~V\tilde{P}V, presents a fundamentally different quantization challenge from QKQK^\top. The paper characterizes P~\tilde{P} explicitly (Section 3.3):

"we note that Sijmij0S_{ij} - m_{ij} \leq 0, so Pij[0,1]P_{ij} \in [0, 1]... P~\tilde{P} often consists of many small elements, but their sum is non-negligible (e.g., 5000 elements around 10410^{-4})"

This distribution—many values clustered near zero, but collectively significant—is the worst case for integer (INT) quantization. INT quantization distributes representable points uniformly across the numerical range. For a range [0,1][0, 1], INT4 would have 16 points at 0,1/15,2/15,,10, 1/15, 2/15, \ldots, 1, meaning values around 10410^{-4} are all quantized to zero. For INT8, 256 points would still put the first non-zero bin at 1/2554×1031/255 \approx 4 \times 10^{-3}, meaning values below 2×1032 \times 10^{-3} get zeroed. The problem is that uniform quantization wastes most of its representational capacity on the large values (near 1) that rarely occur in P~\tilde{P} (since softmax outputs are typically peaked), while underrepresenting the small values that collectively matter for P~V\tilde{P}V.

Why FP8 E4M3 is the right format. Floating-point formats distribute their representable values non-uniformly: they are denser near zero (where the exponent is small) and sparser for large magnitudes. The E4M3 format (4 exponent bits, 3 mantissa bits) can represent values down to 2101032^{-10} \approx 10^{-3} with full subnormal precision, and its representable range extends to 448 (for the maximum exponent). The paper's insight is that the dynamic range of 448 is overkill for P~\tilde{P} (which is bounded in [0,1][0, 1]), but the density near zero is exactly what's needed.

The quantization strategy for P~\tilde{P} is deliberately simple:

δP=1448,P^=cast_to_fp8_e4m3(P~/δP)=cast_to_fp8_e4m3(P~×448)\delta_P = \frac{1}{448}, \quad \hat{P} = \text{cast\_to\_fp8\_e4m3}(\tilde{P} / \delta_P) = \text{cast\_to\_fp8\_e4m3}(\tilde{P} \times 448)

What this computes: Every element of P~\tilde{P} (which is in [0,1][0, 1]) is multiplied by the static scale 448, producing values in [0,448][0, 448], which is exactly the representable range of E4M3. The values are then cast to FP8 E4M3, which rounds them to the nearest representable FP8 value. The dequantization is implicit: the FP8 matrix multiplication with VV (also in FP8) produces a result that is already correctly scaled—the mma.f32.f8.f8.f32 instruction accumulates the FP8 products into an FP22 accumulator, and the final output is later scaled by δV\delta_V (the per-channel VV dequantization scale).

Why a static scale rather than per-token or per-block: The paper chooses δP=1/448\delta_P = 1/448 as a static, data-independent scale. This is possible because P~\tilde{P} is guaranteed to be in [0,1][0, 1] by the softmax operation—unlike QQ and KK, which have unbounded, data-dependent magnitudes. Using a static scale avoids any runtime computation of per-group maxima and eliminates the need for per-element dequantization. The scale δP=1/448\delta_P = 1/448 is chosen because 448=26×7448 = 2^6 \times 7 is the maximum representable value in E4M3 (with exponent bits 1110 and mantissa bits 111, decoding to 277×(1+7/8)×272^{7-7} \times (1 + 7/8) \times 2^{7}: the maximum exponent for normal numbers is 26=642^6=64, giving 64×1.875×21=24064 \times 1.875 \times 2^1 = 240?—the actual max is 448 due to the specific exponent and mantissa encoding in E4M3). Scaling to fill the representable range maximizes the effective precision.

Per-channel quantization for VV. Unlike P~\tilde{P}, VV exhibits significant channel-wise outliers (Figure 2—the heatmap for VV shows per-channel stripes of high magnitude). The paper handles this with per-channel quantization: each of the dd channels (columns) of VV gets its own scale factor δV[c]=max(V[:,c])/448\delta_V[c] = \max(|V[:, c]|) / 448. This means that within a tile VjV_j, each channel dimension is quantized independently:

V^j[:,c]=cast_to_fp8_e4m3(Vj[:,c]/δV[c])\hat{V}_j[:, c] = \text{cast\_to\_fp8\_e4m3}(V_j[:, c] / \delta_V[c])

Per-channel quantization is applicable to VV (but not to QQ and KK) because VV is the right operand in the P~V\tilde{P}V matmul—the matrix multiplication is (P~)bq×bk×(V)bk×d(\tilde{P})_{b_q \times b_k} \times (V)_{b_k \times d}. The column dimension of VV (the dd dimension) is the outer dimension of the product, meaning each column contributes independently to the output. Quantizing per-channel allows outlier channels to use a larger scale while normal channels use a smaller scale, without affecting the matrix multiply layout.

Alternatives evaluated. Table 7 and Table 16 (Appendix A.9) compare different data type choices for P~\tilde{P} and VV (with QQ and KK quantized to INT4 and smoothed):

P~,V\tilde{P}, V FormatAverage CosSimWorst CosSimAverage RMSE
INT877.05%19.52%0.5044
E5M2 (FP8, 5E2M)99.20%94.94%0.0903
E4M3 (FP8, 4E3M)99.44%96.70%0.0347
FP1699.45%96.76%0.0335

The INT8 failure is dramatic: average cosine similarity of 77.05% and worst-case of 19.52%—a complete collapse on some layers, confirming that integer formats are fundamentally unsuited for P~\tilde{P}'s distribution. E5M2 (5 exponent bits, 2 mantissa bits) is significantly better but still shows a worst-case CosSim of 94.94%, indicating some layers suffer. E4M3 achieves nearly FP16-equivalent accuracy (99.44% vs 99.45% average CosSim) while being half the bit width, making it the optimal accuracy-efficiency choice.

Why E4M3 beats E5M2: E5M2 has a larger dynamic range but fewer mantissa bits (2 vs 3). Since P~\tilde{P} has a known bounded range ([0,1][0, 1] after scaling), the extra dynamic range of E5M2 is wasted, while the reduced mantissa precision hurts—it cannot represent small differences between nearly-equal softmax values as accurately. E4M3's extra mantissa bit provides finer granularity in the [0,1][0, 1] range where it matters.


Two-Level Accumulation: Mitigating the Undocumented FP22 Accumulator

The paper makes a notable empirical discovery during implementation: the FP8 tensor core accumulator for mma.f32.f8.f8.f32 on Ada Lovelace (RTX40 series) and Hopper architectures is actually FP22, not FP32. The accumulator physically has:

  • 1 sign bit
  • 8 exponent bits (matching FP32)
  • 13 mantissa bits (vs. 23 for FP32)

This means that when the FP8 matrix multiplication accumulates partial products, the intermediate sum maintains only 13 bits of mantissa precision. When this accumulated value is stored to the FP32 destination register, the lower 10 bits of the mantissa are zeroed out (truncated, not rounded). The paper verified this with a controlled experiment: initialize the FP8 operands A,BA, B to zero (so the product term ABAB contributes nothing), set the accumulator DD to a value with varying mantissa bits, and observe the output C=AB+DC = AB + D. When DD has more than 13 mantissa bits, CC matches DD with the lower 10 bits truncated.

Why this matters for P~V\tilde{P}V. The FlashAttention tiling computes P~V\tilde{P}V iteratively over key/value tiles: each iteration jj computes Rij=P~ijVjR_{ij} = \tilde{P}_{ij} V_j, then accumulates Oij=diag(emi,j1mij)Oi,j1+RijO_{ij} = \text{diag}(e^{m_{i,j-1} - m_{ij}}) O_{i,j-1} + R_{ij}. The number of iterations Tn=N/bkT_n = N / b_k can be large—for a 32K-token sequence with bk=64b_k = 64, that's 512 iterations. If each RijR_{ij} is accumulated in FP22 and the error is additive, the total accumulated error grows with the number of tiles.

The paper's two-level accumulation strategy addresses this by maintaining two separate accumulators:

  1. RijR_{ij} (in FP22): The per-tile product P~ijVj\tilde{P}_{ij} V_j is computed with the mma.f32.f8.f8.f32 instruction. The accumulator for this operation is FP22. However, since RijR_{ij} accumulates over only bk=64b_k = 64 rows of P~\tilde{P} and VV, the total magnitude of the accumulation is bounded. The FP22 accumulator provides approximately 13 bits of mantissa precision, which is sufficient for accumulating 64 terms without catastrophic cancellation or overflow, assuming the individual terms are roughly similar in magnitude.

  2. OijO_{ij} (in FP32): After each tile's RijR_{ij} is fully computed in FP22 registers, it is copied to an FP32 register (which has 23 mantissa bits) and added to the running output Oi,j1O_{i,j-1} in FP32. This copy operation stores the FP22 value into an FP32 register, where the missing mantissa bits are zero. The running accumulation across tiles is then performed in true FP32, which has enough mantissa precision (23 bits) to accumulate hundreds of tile contributions without significant loss.

The key insight is that FP22 is sufficient for local accumulation (64 rows) but insufficient for global accumulation (hundreds of tiles). The two-level strategy uses FP22 only where its limited precision is within the error budget, and FP32 where the cumulative error would become problematic.

Why this is discovered rather than documented. The FP22 accumulator behavior is not described in NVIDIA's PTX ISA documentation, which specifies mma.f32.f8.f8.f32 as operating on "FP32" accumulators. The paper is, to their knowledge, "the first to discover and investigate the effect of the FP22 accumulator and implement the two-level accumulation for attention" (Section 3.4, Remark). The discovery was made through systematic debugging: the FP8 P~V\tilde{P}V implementation showed consistent accuracy degradation compared to FP16 simulations that could not be explained by the FP8 quantization error alone. By narrowing down the problem to the accumulation precision and designing the diagnostic test with zeroed AA and BB matrices and varying DD, they isolated the accumulator bit width.

Connection to prior work. The paper notes that "the two-level accumulation strategy is also implemented in CUTLASS and DeepGemm for computing weight-activation products in linear layers" (Section 3.4, Remark). This is not a completely novel technique—it has been used in GEMM libraries to handle the same FP22 issue for linear layer computation. However, the paper claims to be the first to (1) identify that the same issue affects attention computation, and (2) implement the two-level strategy within the online softmax tiling framework of FlashAttention, which has a more complex accumulation pattern (the OijO_{ij} update involves a scaling factor emi,j1mije^{m_{i,j-1} - m_{ij}} due to the online softmax).

Overhead. Table 18 shows the overhead of the two-level accumulation is 0% on L20 GPU—the kernel throughput is 284 TOPS with and without it. This is because the FP22-to-FP32 copy is a register-to-register move that costs essentially zero cycles compared to the matrix multiplication instructions that dominate the kernel's runtime.


Optional Smoothing of VV for Additional Precision

The paper describes an optional technique to further mitigate FP22 accumulator error when VV exhibits channel-wise bias—a persistent non-zero mean per channel. Section 3.4 and Appendix A.3 provide the details.

When this matters. In some models, particularly diffusion-based video generation models like CogvideoX, the VV matrix has columns whose values are consistently in a narrow positive range (e.g., [8,9][8, 9]). When such a column is multiplied by a row of P~\tilde{P} (which sums to 1), the resulting dot product is a weighted average of values around 8–9, producing outputs that are consistently large positive numbers. The FP22 accumulator, with its 13-bit mantissa, represents numbers near zero with much higher precision than numbers far from zero (because the floating-point representation has more granularity near zero due to the exponent encoding). Therefore, large output values of P~V\tilde{P}V lose effective bits of mantissa precision.

The smoothing procedure. Define the per-channel mean of VV:

Vm=mean(V,axis=0)\vec{V}_m = \text{mean}(V, \text{axis}=0)

This is a 1×d1 \times d vector where each element is the average of the corresponding column across all NN tokens. Then define the smoothed version:

V=VVmV' = V - \vec{V}_m

where the subtraction is broadcast: each token's VV embedding has the per-channel mean subtracted. The attention computation proceeds with VV' instead of VV, producing an intermediate output O=P~VO' = \tilde{P} V'. The final output is corrected by adding back the mean:

O=O+VmO = O' + \vec{V}_m

Why adding back Vm\vec{V}_m is exact. The correction is exact (in infinite precision) because:

P~Vm=Vm\tilde{P} \vec{V}_m = \vec{V}_m

The critical property is that each row of P~\tilde{P} sums to 1 (by definition of softmax). Therefore, multiplying P~\tilde{P} (shape N×NN \times N) by Vm\vec{V}_m (shape N×dN \times d, where every row is identical) produces an N×dN \times d output where every row is exactly Vm\vec{V}_m. The proof: element (i,c)(i, c) of P~Vm\tilde{P} \vec{V}_m is j=1NP~ij(Vm)c=(Vm)cj=1NP~ij=(Vm)c1=(Vm)c\sum_{j=1}^N \tilde{P}_{ij} \cdot (\vec{V}_m)_c = (\vec{V}_m)_c \cdot \sum_{j=1}^N \tilde{P}_{ij} = (\vec{V}_m)_c \cdot 1 = (\vec{V}_m)_c. So adding Vm\vec{V}_m after the matmul perfectly restores the contribution of the subtracted mean.

Why this improves precision. After subtracting Vm\vec{V}_m, each column of VV' is zero-centered. The dot product between a row of P~\tilde{P} and a column of VV' is now a weighted sum of values that are roughly balanced around zero, producing an output near zero. The FP22 representation is much denser near zero—the spacing between representable values is proportional to the magnitude—so the accumulation errors are proportionally smaller. The mean term Vm\vec{V}_m, which is large, is added back in FP32 at the end, so its precision is not compromised.

Empirical validation and scope. Table 10 shows the benefit on real CogvideoX tensors:

Smooth VCosSimRelative L1RMSE
✗ (without)98.25%0.19800.2387
✓ (with)99.75%0.04060.0773

The improvement is significant: cosine similarity increases from 98.25% to 99.75%, and RMSE drops by a factor of ~3×. However, the paper explicitly states this technique is "optional and not employed in our main experiments" (Section 3.4, Remark), because its benefit depends on VV having channel-wise bias—which is present in CogvideoX but absent in some models, such as Llama3.1 (as shown in Figure 2, where VV's heatmap does not show the strong per-channel bias that QQ and KK exhibit). This is an important design choice: the paper prioritizes techniques that help universally across models (smoothing QQ and KK, per-thread quantization, two-level accumulation) and treats model-specific optimizations as optional add-ons.

The Complete Algorithm (Algorithm 1)

Algorithm 1 (pseudocode in the paper) integrates all these techniques into a single fused kernel. The high-level structure mirrors FlashAttention's tiled computation but with quantized matmuls and corrections:

  1. Preprocessing: KK is smoothed once: K=Kmean(K)K = K - \text{mean}(K) (global mean, computed across all tokens). VV is quantized per-channel once: (δV,V^)=ψV(V)(\delta_V, \hat{V}) = \psi_V(V). QQ is split into Tm=N/bqT_m = N/b_q tiles {Qi}\{Q_i\}.

  2. Per-tile loop over queries (outer loop, index ii): For each QiQ_i, compute the per-block mean qˉi=mean(Qi)\bar{q}_i = \text{mean}(Q_i), smooth QQ: γ(Qi)=Qiqˉi\gamma(Q_i) = Q_i - \bar{q}_i, and quantize per-thread: (δQ,Q^i)=ψQ(γ(Qi))(\delta_Q, \hat{Q}_i) = \psi_Q(\gamma(Q_i)).

  3. Per-tile loop over keys/values (inner loop, index jj): For each KjK_j, VjV_j:

    • Quantize KjK_j per-thread: (δK,K^j)=ψK(Kj)(\delta_K, \hat{K}_j) = \psi_K(K_j) (note: γ(Kj)\gamma(K_j) is already computed since KK was smoothed globally).
    • Each warp ww (out of cwc_w warps) processes its segment: compute Matmul(Q^i[st:st+cw],K^j)\text{Matmul}(\hat{Q}_i[\text{st}:\text{st}+c_w], \hat{K}_j^\top) using INT4 mma, dequantize with per-thread scales, and add the GEMV(qˉi,Kj)\text{GEMV}(\bar{q}_i, K_j^\top) correction to form SijS_{ij}.
    • Run online softmax: update mijm_{ij} (row-wise max), compute P~ij=exp(Sijmij)\tilde{P}_{ij} = \exp(S_{ij} - m_{ij}), update lijl_{ij} (row sum).
    • Cast P~ij\tilde{P}_{ij} to FP8 E4M3: (P~ij×448).to(fp8_e4m3)(\tilde{P}_{ij} \times 448).\text{to}(\text{fp8\_e4m3}).
    • Compute Oij(FP22)=Matmul(P^ij,V^j)O_{ij}(\text{FP22}) = \text{Matmul}(\hat{P}_{ij}, \hat{V}_j) using FP8 mma.
    • Accumulate to FP32 buffer: Oij(FP32)=diag(emi,j1mij)Oi,j1(FP32)+Oij(FP22)O_{ij}(\text{FP32}) = \text{diag}(e^{m_{i,j-1} - m_{ij}}) O_{i,j-1}(\text{FP32}) + O_{ij}(\text{FP22}).
  4. Output normalization: After the jj loop completes, Oi=diag(li,Tn)1Oi,Tn(FP32)/448×δVO_i = \text{diag}(l_{i, T_n})^{-1} O_{i, T_n}(\text{FP32}) / 448 \times \delta_V.

The fusion of smoothing, quantization, GEMV, and attention into a single kernel means that QQ, KK, and VV are read from off-chip memory exactly once—the bandwidth cost is identical to an unquantized attention kernel. The additional operations (mean subtraction, per-thread scale computation, GEMV for ΔS\Delta S, FP22-to-FP32 copy) are all on-chip and contribute negligible latency (Table 18: smoothing QQ adds 3.7% overhead, per-thread quantization adds 0.35% overhead, two-level accumulation adds 0%).

Summary of Design Choices and Their Justifications

  • Smooth QQ with per-block mean, KK with global mean rather than SmoothQuant or Hadamard: both QQ and KK are dynamic activations with outliers; smoothing each independently using self-statistics avoids transferring difficulty, and the per-block granularity for QQ matches the FlashAttention tiling.
  • Per-thread INT4 quantization rather than per-token or per-block: achieves per-token-equivalent accuracy while maintaining the zero-overhead dequantization of per-block (one scale per thread, not one scale per element).
  • FP8 E4M3 for P~\tilde{P} with static scale 1/448 rather than INT8 or E5M2: P~[0,1]\tilde{P} \in [0, 1] is bounded, making a static scale possible; E4M3's 3 mantissa bits provide better granularity for the small-value-dominated distribution than E5M2's 2 mantissa bits or INT's uniform quantization.
  • Per-channel FP8 for VV rather than per-tensor: VV exhibits channel-wise outliers; per-channel quantization prevents outlier channels from dominating the scale and zeroing out normal channels.
  • Two-level accumulation (FP22 local, FP32 global) rather than pure FP22 or pure FP32: the hardware's FP22 accumulator is fast but lacks precision for many-tile accumulation; buffering to FP32 after each tile prevents error accumulation while keeping the fast FP8 matmul in the inner loop.
  • Smoothing VV as optional rather than always-on: effective only when VV has channel-wise bias; applying it universally would add overhead without benefit for models like Llama3.1 where VV biases are absent.

4. Key Insights and Innovations

Innovation 1: The FP22 Accumulator Discovery Recasts Hardware Precision as a First-Class Design Constraint in Attention Kernels

The paper's most intellectually distinctive contribution is not a technique per se, but a diagnostic discovery that changes how kernel developers should think about GPU tensor core hardware: the accumulator for NVIDIA's mma.f32.f8.f8.f32 instruction on Ada and Hopper architectures is FP22, not FP32, silently truncating 10 bits of mantissa precision. This finding is significant beyond the specific fix it enables because it reveals that the gap between documented ISA semantics and physical hardware behavior is large enough to cause measurable accuracy degradation in attention—a core operation in every transformer model.

Prior to this work, the standard assumption in the quantized attention literature was that hardware precision follows the ISA specification: if the PTX instruction declares FP32 accumulators, then the accumulation must be FP32-precise. SageAttention (Zhang et al., 2025c) exploited reduced accumulator precision for P~V\tilde{P}V on RTX4090, but did so through the documented FP16 accumulator mode—the precision was known and deliberately traded for speed. FlashAttention3 (Shah et al., 2024) uses FP8 for attention on Hopper GPUs, but the paper shows its accuracy degrades on models like CogvideoX (Table 2) and Mochi (Figure 9), possibly because the FP22 accumulator behavior went unrecognized.

What makes this discovery fundamental rather than incremental is that it transforms the accumulator from a trusted black box into a design parameter with hidden constraints. The paper's diagnostic methodology—zeroing the FP8 operands and varying only the accumulator input to test bit-level precision—is itself a contribution, providing a reusable protocol for probing undocumented hardware behavior. The finding also carries implications beyond attention: any kernel that uses mma.f32.f8.f8.f32 for reduction-style operations (where many small terms accumulate to a significant total) is vulnerable to the same silent truncation. Linear layer quantization libraries like CUTLASS and DeepGemm had already discovered this and implemented two-level accumulation for weight-activation products (Section 3.4, Remark), but the paper is the first to identify that attention—with its iterative softmax accumulation across tiles—is similarly affected, and the first to surface the behavior in a published, peer-reviewed venue accessible to the ML systems community.

The significance of this finding is evidenced by how it redirects design effort: without the FP22 discovery, a kernel developer would observe accuracy degradation in FP8 P~V\tilde{P}V, likely attribute it to quantization error in P~\tilde{P} or VV, and try to improve quantization fidelity—a path that couldn't fully fix the problem. The FP22 finding correctly identifies accumulation precision, not operand quantization, as the root cause, enabling the targeted fix of two-level accumulation.

Innovation 2: Per-Thread Quantization Introduces a General Principle for Quantization Granularity in Fused Kernels

The paper's per-thread quantization method embodies a conceptual insight that extends beyond attention: quantization granularity in fused GPU kernels should be defined by the hardware's thread-to-data mapping, not by logical tensor dimensions (per-tensor, per-channel, per-token). This is a shift in how quantization is thought about in high-performance settings.

The conventional approach to quantization granularity—exemplified by QLoRA (Dettmers et al., 2023), SmoothQuant (Xiao et al., 2023), and SageAttention's per-block quantization—defines quantization groups along logical tensor axes: a group is "a channel," "a token," "a block of tokens," or "the entire tensor." These groupings are natural from a mathematical perspective because they align with the statistical structure of the data (per-channel for weights, per-token for activations with varying norms). But they are unnatural from a hardware execution perspective: when a GPU warp executes a tensor core MMA instruction, each thread computes a specific, interleaved subset of the output. If quantization groups don't align with these subsets, each thread must load and multiply multiple scale factors during dequantization, adding instruction overhead and register pressure to the innermost loop.

The paper's innovation is to reverse the design direction: start from the MMA instruction's thread layout, derive which output elements each thread owns, and define quantization groups to match. This produces groups that are irregular in logical tensor space—QQ gets 32 groups per 128-token block (8 per warp × 4 warps) while KK gets 4 groups per 64-token block—but that achieve a sweet spot: the accuracy of per-token quantization (99.45% vs. 99.45% average CosSim in Table 6) with the dequantization overhead of per-block quantization (single scale pair per thread, 283 TOPS vs. 284 TOPS in Table 19).

The intellectual contribution here is not the specific grouping formula (Equation 8), which is necessarily hardware-specific and would change for a different MMA instruction or warp size. It is the design principle: when quantization and dequantization must be fused into a compute-bound kernel (rather than applied as a separate preprocessing step), the quantization grouping should be determined by the data access pattern of the compute instruction, not by the statistical properties of the tensor. The paper demonstrates this principle's power through negative evidence: per-token quantization, which is statistically optimal (matching the finest granularity), degrades throughput by ~5.3% (268 vs. 283 TOPS in Table 19) when fused into the attention kernel due to per-element scale multiplications, while per-thread quantization achieves identical accuracy with zero overhead.

This principle is likely transferable to other fused quantized kernels beyond attention—any operation where quantization and computation are fused and the computation has a fixed thread-to-data mapping (e.g., fused MLP blocks, custom attention variants) could benefit from aligning quantization groups to thread ownership.

Innovation 3: Outlier Smoothing in Attention Requires Independent, Self-Referential Treatment of Both Operands—Not Cross-Operand Difficulty Transfer

The paper's smoothing approach for QQ and KK makes a conceptual contribution to the quantization of multi-operand operations where both operands are dynamic: when both operands in a matrix multiplication are activation tensors with outliers, the correct smoothing strategy is to smooth each independently using its own statistics, not to transfer quantization difficulty between them. This directly contradicts the dominant paradigm for activation-weight quantization, epitomized by SmoothQuant, which works precisely by transferring difficulty: smooth the activations at the cost of making the weights harder to quantize, exploiting the fact that weights are static and can be pre-compensated offline.

The paper shows that this paradigm fails for QKQK^\top in attention because both QQ and KK are dynamic—there is no static operand to absorb the smoothing burden. SmoothQuant applied to attention (the "SmoothAttn" baseline) achieves only 90.21% average CosSim with INT4 quantization (Table 4), because smoothing QQ at KK's expense (or vice versa) merely moves the outlier problem. The Hadamard rotation approach (Quarot, "HadmdAttn") fares even worse at 79.77%—the rotation is a fixed transformation that doesn't adapt to the specific outlier structure, so redistributed energy still produces large values after rotation.

The paper's key conceptual move is to recognize that the token-similarity structure of attention tensors (Figure 2: "Q,KQ, K for all tokens are actually highly similar, with only small variations between different tokens") provides a different smoothing opportunity that doesn't require a static operand. By subtracting the mean along the token dimension, the paper removes the shared per-channel bias that creates the outliers, leaving only the token-to-token variation that carries the attention signal. The mathematical decomposition into γ(Qi)γ(Kj)+ΔSij\gamma(Q_i)\gamma(K_j)^\top + \Delta S_{ij} (Section 3.1) shows that the smoothed product term has small magnitude (quantizable in INT4) while the correction term ΔSij\Delta S_{ij} is cheap to compute in FP16 and restores exactness. This is fundamentally different from SmoothQuant's scaling approach—it's a decomposition into a quantizable main term and a full-precision correction term, leveraging the softmax invariance to constant row offsets to drop the bias terms entirely.

The significance is that this decomposition pattern may generalize to other attention-like operations where a softmax or normalization follows the matrix multiplication, making the operation invariant to constant offsets along the normalized dimension. The paper doesn't explore this generalization, but the pattern (decompose into zero-centered product + correction, exploit normalization invariance to drop terms) is a reusable template.

Innovation 4: Difficulty-Aware Precision Allocation as a Spectrum Across Model Families, Not a Single Configuration

While SageAttention2 is primarily a systems paper rather than an algorithmic analysis paper, its extensive evaluation across language, image, and video models reveals a meta-insight that is easy to miss: the acceptable precision for quantized attention is not a fixed threshold but varies systematically across model architectures and modalities, and the paper's suite of techniques (smoothing Q, per-thread quantization, smoothing V) provides tunable knobs to match each model's sensitivity profile.

The evidence for this is distributed across the evaluation. Llama3.1 can tolerate SageAttn2-4b with only minor metric loss (MMLU drops from 63.5% to 60.7%, Table 2) and SageAttn2-8b with essentially no loss. CogvideoX, however, requires SageAttn2-8b to avoid visible degradation—the 4-bit variant produces measurable metric drops (VQA-a drops from 70.2 to 57.7 for the 1.5-5B model, Table 2) despite the same smoothing and per-thread quantization. The optional VV smoothing technique (Table 10) further demonstrates this model-specificity: it provides a ~3× RMSE reduction on CogvideoX (where VV has channel-wise bias) but would provide no benefit on Llama3.1 (where VV biases are absent, per Figure 2).

This is not presented as a central theoretical contribution, but it constitutes a practical insight with real deployment implications: a quantized attention system should expose a precision spectrum (4-bit vs. 8-bit for QKQK^\top, with or without VV smoothing) rather than a single "quantized attention" configuration, because different models sit at different points on the accuracy-efficiency Pareto frontier. The paper's decision to release both SageAttn2-4b and SageAttn2-8b as separate kernel variants, with SageAttn2-8b providing a nearly-lossless option for accuracy-sensitive models like video generation and SageAttn2-4b providing maximum speedup for robustness-tolerant models, embodies this insight.

This contrasts with the approach in FlashAttention3, which provides a single FP8 configuration that works well on some models but degrades visibly on video generation (Figures 7, 9). The paper doesn't frame this as a core innovation, but the evaluation implicitly demonstrates that the precision requirements for quantized attention are modality-dependent in ways that prior work hadn't systematically documented.

5. Experimental Analysis

Evaluation Methodology

Dataset. The paper evaluates on a diverse set of benchmarks across modalities rather than a single dataset. For text-to-text models, the evaluation uses WikiText (Merity et al., 2022) for perplexity, LAMBADA (Paperno et al., 2016) for contextual understanding, MMLU (Hendrycks et al., 2021b) for knowledge across subjects, Longbench (Bai et al., 2024) for long-context capabilities, and InfiniBench (Zhang et al., 2024) and Needle-in-a-Haystack (Kamradt, 2023) for super-long context evaluation up to 262K tokens. For text-to-video models, evaluation uses the open-sora prompt sets (Zheng et al., 2024c) and measures five video quality metrics. For text-to-image models, assessment is on MJHQ-30K (Li et al., 2024). For image classification, ImageNet, ImageNet-Sketch, and ImageNet-Rendition (Deng et al., 2009; Wang et al., 2019; Hendrycks et al., 2021a) are used. For audio, Librispeech (Panayotov et al., 2015) test splits are evaluated with Word Error Rate.

Base models. The paper validates across ten representative models spanning four modalities: Llama2 (7B) and Llama3.1 (8B) (Touvron et al., 2023; Dubey et al., 2024) and GLM4 (9B) (GLM et al., 2024) for language; CogvideoX (2B) and CogvideoX (1.5-5B) (Yang et al., 2025b), HunyuanVideo (Kong et al., 2024), and Mochi (Team, 2024) for video generation; Flux (schnell) (Black Forest Labs, 2023) and Stable-Diffusion3.5 (turbo) (Stability AI, 2023) for image generation; TIMM (Wightman, 2019) for image classification; and Qwen2-Audio (7B) (Chu et al., 2024) for audio. Additionally, Llama-3-262k (8B) is used for super-long context experiments. This breadth is deliberately chosen to demonstrate that SageAttention2 works across model scales (2B–9B parameters), modalities, and architectures, establishing it as a general drop-in replacement rather than a model-specific optimization.

Metrics. The paper uses both kernel-level accuracy metrics and end-to-end task metrics. Kernel accuracy is measured by comparing the quantized attention output OO' against the full-precision attention output OO using three metrics: Cosine Similarity (CosSim=OO/O2O2\text{CosSim} = \sum OO' / \sqrt{\sum O^2 \sum O'^2}), Relative L1 Distance (L1=OO/O\text{L1} = \sum |O - O'| / \sum |O|), and Root Mean Square Error (RMSE=(1/n)(OO)2\text{RMSE} = \sqrt{(1/n) \sum (O - O')^2}). Both average and worst-case (per-layer maximum error) values are reported to capture outlier layers that may dominate end-to-end degradation. End-to-end metrics vary by modality: perplexity for WikiText, accuracy for LAMBADA/MMLU, Longbench score, CLIPSIM and CLIP-Temp for text-video alignment, VQA-a and VQA-t for video aesthetic/technical quality, Flow-score for temporal consistency, FID and sFID for image fidelity, Clipscore for text-image alignment, ImageReward for human preference, and classification accuracy for TIMM. For audio, Word Error Rate (WER) is used.

Baselines. The paper compares against five methods:

  1. Full-Precision Attention: the unquantized FlashAttention2 or equivalent exact attention in FP16/BF16, serving as the quality upper bound.
  2. SmoothAttn: Following Qserve (Lin et al., 2025), applies SmoothQuant (Xiao et al., 2023) with smoothing factor α=0.5\alpha = 0.5 to QQ and KK before INT4 quantization.
  3. HadmdAttn: Following Quarot (Ashkboos et al., 2024), applies random Hadamard transformation to QQ and KK before INT4 quantization.
  4. SageAttention (Zhang et al., 2025c): The direct predecessor that smooths KK only, uses INT8 per-block quantization for QQ and KK, and FP16 for P~\tilde{P} and VV.
  5. FlashAttn3(fp8) (Shah et al., 2024): The FP8 version of FlashAttention3, evaluated only on Hopper GPUs where it is compatible.

For kernel speed benchmarks, the standard baselines also include Torch (native PyTorch attention), xformers (Lefaudeux et al., 2022), and FlashAttention2 (Dao, 2024). The paper also benchmarks against FlashAttention3's FP16 variant on H100/H20 GPUs.

Generation budget / compute accounting. Kernel speed is measured in TOPS (Tera Operations Per Second) on real GPU hardware, using a batch size of 4 with 32 attention heads across sequence lengths from 1K to 32K tokens. Input tensors for speed benchmarking are drawn from Gaussian distributions with μ=0,σ=1\mu=0, \sigma=1 for floating-point types and uniform sampling within the representable range for integer types ([-128, 127] for INT8, [-8, 7] for INT4), following the standard practice in FlashAttention benchmarks (Dao et al., 2022). End-to-end generation latency is measured in wall-clock seconds for specific models at specific sequence lengths (e.g., CogvideoX at its native generation length, Llama3.1 at 48K and 100K tokens). The key accounting principle is that SageAttn2-4b uses INT4 for QKQK^\top and FP8 for P~V\tilde{P}V, while SageAttn2-8b uses INT8 for QKQK^\top and FP8 for P~V\tilde{P}V—both are compared at equivalent generation budgets (same sequence length, same number of attention heads, same model) against baselines using higher-precision formats.

Cross-validation / statistical protocol. There is no cross-validation or statistical significance testing for the end-to-end metrics. The results in Table 2 are single-run evaluations. For kernel accuracy measurements (Tables 4, 6, 7, 15, 16, 17), the paper reports averages and worst-case values across all layers of a model (e.g., "across all layers of CogvideoX"), which provides a per-layer distribution but does not include confidence intervals or multiple random seeds. The kernel speed benchmarks (Figures 5, 10–16) are direct hardware measurements without reported variance. This is standard practice for GPU kernel papers where execution is deterministic, but it means the end-to-end metric comparisons (e.g., MMLU 63.5% vs. 60.7% in Table 2) should be interpreted as point estimates without formal statistical comparison.


Main Quantitative Results

Kernel Speed Comparison

The headline result is that SageAttn2-4b achieves approximately 3× the throughput of FlashAttention2 and 4.5× the throughput of xformers on consumer GPUs, while SageAttn2-8b matches or exceeds the speed of FlashAttention3(fp8) on datacenter GPUs.

On the RTX4090 with head_dim=128, causal=False (Figure 5, right panel), the throughput at sequence length 32K is:

MethodTOPS
Torch46
xformers105
FlashAttention2164
SageAttention1338
SageAttn2-8b430
SageAttn2-4b481

SageAttn2-4b achieves 481 TOPS, which is 2.93× FlashAttention2 (481/164), 4.58× xformers (481/105), and 1.42× SageAttention (481/338). The 8-bit variant achieves 430 TOPS, which is 2.62× FlashAttention2 and 1.27× SageAttention. The gap between SageAttn2-4b and SageAttn2-8b (481 vs. 430 TOPS) directly reflects the 2× throughput advantage of INT4 over INT8 tensor cores—the paper recovers approximately 70% of the theoretical maximum speedup (481/338 = 1.42× vs. theoretical 2×).

With causal masking enabled (Figure 5, left panel), the pattern is identical: SageAttn2-4b achieves 479 TOPS at 32K vs. 164 for FlashAttention2 (2.92×) and 106 for xformers (4.52×). The causal mask does not materially affect the relative speedups, indicating that the quantization overhead is well-amortized across both causal and non-causal attention patterns.

On the L20 GPU (datacenter Ada Lovelace, Figure 12), the throughput at head_dim=128, causal=False, 32K sequence length is:

MethodTOPS
FlashAttention274
SageAttention1138
SageAttn2-8b217
SageAttn2-4b273

SageAttn2-4b achieves 3.69× FlashAttention2 (273/74) and 1.98× SageAttention (273/138). Note that SageAttention1's relative speedup over FlashAttention2 is smaller on L20 (1.86×) than on RTX4090 (2.06×) because SageAttention1's FP16 accumulator speedup (W2 in the paper's motivation) does not apply on L20 (Table 1). SageAttention2's FP8 P~V\tilde{P}V matmul provides a consistent 2× speedup on L20 where SageAttention1's FP16 accumulator trick provides none, explaining the larger improvement.

On Hopper GPUs (H100, Figures 13 and 14; H20, Figures 15 and 16), SageAttention2-8b is compared against FlashAttention3 (both FP16 and FP8 variants). The headline result is that SageAttn2-8b matches or slightly exceeds FlashAttention3(fp8) in throughput while being significantly more accurate. At head_dim=128, causal=False, 32K on H100 (Figure 14, left):

MethodTOPS
FlashAttention2339
FlashAttention3(fp16)470
FlashAttention3(fp8)892
SageAttn2-8b882

SageAttn2-8b achieves 882 TOPS, essentially matching FlashAttention3(fp8)'s 892 TOPS (0.99×). Against FlashAttention2, it achieves 2.60×. On H20 (Figure 16, left), at the same configuration, SageAttn2-8b achieves 279 TOPS vs. FlashAttention3(fp8)'s 273 TOPS (1.02×) and FlashAttention2's 90 TOPS (3.10×).

A notable detail in Figure 5: the throughput curves for all methods increase with sequence length and saturate around 16K–32K tokens. This is the expected behavior for compute-bound kernels—at short sequence lengths, the kernel is memory-bandwidth-bound (reading QQ, KK, VV dominates), and as sequence length grows, the O(N2)O(N^2) compute becomes the bottleneck, allowing the tensor core throughput to be fully utilized. SageAttention2's curves saturate at higher absolute TOPS because the quantized matmuls achieve higher peak tensor core utilization. The saturation behavior also confirms that the per-thread quantization and smoothing overhead (Table 18: 0.35% + 3.7%) does not create a new bandwidth bottleneck.

Table 9 (Appendix A.2) summarizes speedup across all tested GPUs. The key number: SageAttention2 achieves a 2.46–3.12× speedup over FlashAttention2 across consumer (RTX4090), datacenter Ada (L20, L40), and Hopper (H100, H20) GPUs, with the variation attributable to different tensor core throughput ratios across architectures. Notably, SageAttention2 achieves 3.12× on H20, the highest relative speedup, because FlashAttention2 is relatively slow on H20 (90 TOPS at 32K, head_dim=128) while SageAttention2 can exploit the FP8 tensor cores effectively.

Kernel Accuracy

The paper evaluates kernel accuracy in two configurations: (1) INT4 QKQK^\top with FP16 P~V\tilde{P}V to isolate the QKQK^\top quantization accuracy, and (2) INT4 QKQK^\top with FP8 P~V\tilde{P}V for the full SageAttn2-4b configuration. All accuracy measurements use real tensors sampled from CogvideoX across all layers, with both average and worst-case values reported.

Smoothing ablation (Tables 4, 5, 17): With INT4 QKQK^\top and FP16 P~V\tilde{P}V, the average cosine similarity across all CogvideoX layers for different smoothing methods is:

Smoothing MethodAverage CosSimWorst CosSimAverage RMSE
None80.04%4.83%0.2223
HadmdAttn79.77%4.85%0.2180
SmoothAttn90.21%64.49%0.1952
Smooth K only98.07%90.86%0.0743
Smooth Q only98.30%93.10%0.0712
Smooth Q+K99.46%96.71%0.0334

The key finding is that smoothing both QQ and KK is necessary for acceptable worst-case accuracy. While Smooth K only achieves 98.07% average CosSim, its worst-case drops to 90.86%—indicating that some layers have QQ outliers severe enough to cause substantial degradation even when KK is smoothed. Smooth Q+K achieves 99.46% average and 96.71% worst-case CosSim. The worst-case improvement from Smooth Q+K over Smooth K only (90.86% → 96.71%) is critical because end-to-end model quality is often determined by the worst layer, not the average.

The Hadamard rotation approach (HadmdAttn) performs essentially identically to no smoothing (79.77% vs. 80.04% average, 4.85% vs. 4.83% worst), confirming that fixed orthogonal transformations do not spread attention outliers sufficiently for INT4 quantization. SmoothAttn (SmoothQuant-style cross-operand smoothing) provides a modest improvement but still collapses to 64.49% worst-case, far below acceptable thresholds.

Table 5 translates these kernel-level accuracy measurements to end-to-end metrics on Llama3.1 and CogvideoX (2B): with INT4 QKQK^\top and no smoothing, Llama3.1 Lambda accuracy drops from 81.5% (full precision) to 72.6%, and CogvideoX VQA-t drops from 75.360 to 24.670—a complete collapse. With Smooth Q+K, Llama3.1 recovers to 80.8% and CogvideoX to 75.147, nearly matching full precision.

Granularity ablation (Tables 6, 15): With Smooth Q+K and INT4 QKQK^\top (FP16 P~V\tilde{P}V):

GranularityAverage CosSimWorst CosSimAverage RMSE
Per-tensor97.15%85.85%0.0865
Per-block98.03%90.68%0.0744
Per-thread99.45%96.72%0.0313
Per-token99.45%96.76%0.0335

Per-thread achieves identical average CosSim to per-token (99.45%) and nearly identical worst-case (96.72% vs. 96.76%). The critical comparison is per-thread vs. per-block: per-thread improves average CosSim by 1.42 percentage points (98.03% → 99.45%) and worst-case by 6.04 percentage points (90.68% → 96.72%), while RMSE drops by more than half (0.0744 → 0.0313). The worst-case improvement is the most important result—per-block quantization leaves some layers with 90.68% CosSim, which would manifest as visible degradation in generated outputs, while per-thread brings the worst layer to 96.72%, well within the acceptable range.

Table 19 (Appendix A.9) confirms that this accuracy comes without speed penalty: on L20 GPU, per-thread achieves 283 TOPS vs. 284 for per-block and 286 for per-tensor—the 0.35% overhead is measurement noise. Per-token, in contrast, drops to 268 TOPS (5.3% overhead) due to the vector dot product of per-element scales.

Data type ablation for P~V\tilde{P}V (Tables 7, 16): With Smooth Q+K, INT4 QKQK^\top, varying the P~V\tilde{P}V format:

P~V\tilde{P}V FormatAverage CosSimWorst CosSimAverage RMSE
INT877.05%19.52%0.5044
E5M2 (FP8)99.20%94.94%0.0903
E4M3 (FP8)99.44%96.70%0.0347
FP1699.45%96.76%0.0335

E4M3 achieves virtually the same accuracy as FP16 (99.44% vs. 99.45% average, 96.70% vs. 96.76% worst) while using half the bit width. INT8 is a complete failure—average CosSim of 77.05% and worst-case of 19.52% confirm that uniform quantization cannot handle P~\tilde{P}'s distribution of many small but collectively significant values. E5M2 is substantially better than INT8 but its worst-case (94.94%) is noticeably worse than E4M3 (96.70%), validating the choice of E4M3's extra mantissa bit over E5M2's extra exponent range for the bounded [0,1][0, 1] domain of P~\tilde{P}.

End-to-End Model Quality

Table 2 presents the central end-to-end results across all ten models. The paper reports two SageAttention2 variants: SageAttn2-4b (INT4 QKQK^\top, FP8 P~V\tilde{P}V) and SageAttn2-8b (INT8 QKQK^\top, FP8 P~V\tilde{P}V).

Language models (Llama3.1, GLM4, Llama2):

For Llama3.1 (8B) on MMLU:

MethodMMLU Accuracy
Full-Precision63.5%
HadmdAttn50.0%
SmoothAttn54.1%
SageAttention63.4%
SageAttn2-4b60.7%
SageAttn2-8b63.4%

SageAttn2-8b matches full-precision and SageAttention exactly (63.4% vs. 63.5%—within measurement noise), confirming that INT8 QKQK^\top with FP8 P~V\tilde{P}V incurs negligible accuracy loss on language tasks. SageAttn2-4b shows a 2.8 percentage point drop (63.5% → 60.7%), which is small but measurable—the 4-bit quantization does lose some precision relative to 8-bit. On WikiText perplexity: 6.013 (full) → 6.256 (SageAttn2-4b) → 6.019 (SageAttn2-8b), and Lambda accuracy: 81.5% → 79.8% → 81.1%.

The same pattern holds for GLM4 (9B): MMLU drops from 74.3% to 72.5% for 4-bit and recovers to 74.5% for 8-bit. Longbench: 49.78 → 49.23 → 49.60.

Llama2 (7B) in Table 11 shows consistent results: MMLU 43.9% → 42.8% (4-bit) → 43.8% (8-bit), Lambda 88.6% → 88.1% → 88.6%.

The HadmdAttn and SmoothAttn baselines consistently fail across all language models: HadmdAttn drops MMLU by 13.5 points on Llama3.1 (63.5% → 50.0%) and SmoothAttn drops it by 9.4 points (63.5% → 54.1%), confirming that neither generic quantization technique transfers to attention.

Video generation models (CogvideoX, HunyuanVideo, Mochi):

The video models are where the accuracy differences between methods are most dramatic and visible. For CogvideoX (1.5-5B):

MethodVQA-aVQA-tFScore
Full-Precision70.23170.9282.507
HadmdAttn8.9902.299
SmoothAttn8.8122.277
SageAttention
FlashAttn3-fp86.5312.181
SageAttn2-4b57.72952.9892.884
SageAttn2-8b69.49274.4152.487

The "✗" entries indicate catastrophic failure—the method could not produce results evaluable by the metric (SageAttention's INT8 quantization fails on CogvideoX due to accuracy issues on this specific model). FlashAttn3(fp8) similarly collapses (VQA-a = 6.531, VQA-t = 2.181, FScore fails), showing that generic FP8 attention without the paper's smoothing and accumulation techniques is insufficient for video generation.

SageAttn2-4b substantially outperforms all baselines (VQA-a 57.729 vs. <9 for alternatives) but shows a 12.5-point drop from full precision (70.231 → 57.729), indicating that INT4 is at the edge of acceptability for this model. SageAttn2-8b nearly matches full-precision (VQA-a 69.492 vs. 70.231, VQA-t 74.415 vs. 70.928). Interestingly, SageAttn2-8b's VQA-t (74.415) actually exceeds full precision (70.928)—this is likely within the variance of the metric on a finite evaluation set rather than a genuine improvement, but it confirms there is no systematic degradation.

For HunyuanVideo, the pattern is similar but less extreme: HadmdAttn and SmoothAttn produce very low VQA-a (7.514, 6.987) and FScore (0.175, 0.148), FlashAttn3(fp8) produces VQA-a of 4.433 (worse than HadmdAttn), while SageAttn2-4b achieves VQA-a of 81.478 (vs. 82.516 full) and SageAttn2-8b achieves 81.786. All SageAttention2 variants produce FScores indistinguishable from full precision (0.586–0.610 vs. 0.604).

For Mochi, FlashAttn3(fp8) shows VQA-a of 14.964 vs. 45.549 full (a 67% drop), while SageAttn2-4b achieves 35.955 and SageAttn2-8b achieves 46.760. FScore: FlashAttn3(fp8) gets 0.457 vs. 1.266 full, SageAttn2-8b gets 1.255.

The visible comparisons (Figures 6, 7, 8, 9) corroborate the metrics: HadmdAttn and SmoothAttn produce obviously degraded frames with blurring and artifacts; FlashAttn3(fp8) produces visible quality loss on CogvideoX (Figure 7) and Mochi/HunyuanVideo (Figure 9); SageAttn2-8b produces frames visually indistinguishable from full precision; SageAttn2-4b has minor but visible differences (slightly reduced detail in Figure 6, 8) but remains far better than baselines.

Image generation models (Flux, Stable-Diffusion3.5):

For Flux (schnell):

MethodFID ↓CLIP ↑
Full-Precision10.96026.180
SageAttn2-4b10.57726.141
SageAttn2-8b10.92726.175

Interestingly, SageAttn2-4b achieves a better FID (10.577) than full precision (10.960)—FID is lower-is-better, so this is an improvement. The paper does not claim quantization improves quality; this is likely due to the finite evaluation set (MJHQ-30K with 30K images) and the stochastic nature of the generation process. The key point is there is no degradation. sFID for 4-bit (17.497) is slightly worse than full (16.648), and 8-bit (16.723) is essentially identical.

For Stable-Diffusion3.5, all variants produce nearly identical metrics: FID 14.105 (full) vs. 14.097 (4-bit) vs. 14.106 (8-bit). ImageReward: 0.902 vs. 0.895 vs. 0.901.

Image classification (TIMM, Table 13): On ImageNet, SageAttn2-8b exactly matches full precision (84.79%), SageAttn2-4b achieves 86.67%—actually higher. Again, this is within variance. On ImageNet-Sketch and ImageNet-R, all variants are within 0.5% of full precision.

Audio (Qwen2-Audio, Table 20): Word Error Rate on Librispeech test-clean: Full precision 1.74% → SageAttn2-4b 1.73% → SageAttn2-8b 1.72%. On test-other: 4.01% → 3.99% → 4.03%. All within measurement noise.

Super-long context (Llama-3-262k, Table 14, Figure 19): On InfiniBench (Zhang et al., 2024) at 262K tokens on H100:

MethodAverage Score
Full-Precision43.05
FlashAttn3-fp841.53
SageAttention243.06

SageAttention2 achieves 43.06, essentially identical to full precision (43.05), while FlashAttn3(fp8) drops to 41.53—a 1.52-point gap on an 8-task benchmark. The most striking difference is in the Retr.KV sub-task: FlashAttn3(fp8) scores 0.4 (essentially zero) while SageAttention2 scores 6.6 and full precision scores 7.0. This indicates that FlashAttn3(fp8)'s FP8 accumulation (likely suffering from the FP22 accumulator issue without two-level buffering) catastrophically fails on retrieval tasks requiring precise attention over very long contexts.

Figure 19 visualizes Needle-in-a-Haystack results: SageAttention2 maintains near-perfect retrieval (green across all depths and token limits) matching full precision, while FlashAttn3(fp8) shows a dark vertical band of failures at mid-depths for long contexts—a classic pattern of attention degradation where the model loses the needle in the middle of very long sequences.

End-to-End Speedup

Table 8 reports wall-clock generation latency:

ModelGPUOriginalSageAttn2-8bSageAttn2-4bSpeedup (8b)
CogvideoX (2B)RTX409086 s54 s52 s1.59×
CogvideoX (1.5-5B)RTX40901040 s577 s555 s1.80×
HunyuanVideoL202221 s1486 s1435 s1.49×
MochiL202336 s1316 s1190 s1.77×
Llama3.1 (48K)RTX40909.2 s5.7 s5.6 s1.61×
Llama3.1 (100K)L2039.9 s25.4 s23.2 s1.57×

The headline 1.8× end-to-end speedup on CogvideoX (1.5-5B) with SageAttn2-8b represents nearly 8 minutes saved per video generation (1040s → 577s). The 4-bit variant further reduces this to 555s (1.87×), but with the slight metric degradation noted in Table 2. The end-to-end speedups are smaller than the kernel-level speedups (1.5–1.8× vs. 2.5–3×) because attention is not the only operation in these models—feed-forward layers, normalization, and other operations are unaffected by SageAttention2 and limit Amdahl's-law speedup. The paper does not report the fraction of total runtime spent in attention for each model, but the 1.5–1.8× range implies attention accounts for roughly 33–50% of total inference time (if attention were 100% of runtime, the speedup would equal the kernel speedup of ~3×; the observed 1.8× suggests attention is approximately 55% of total time on CogvideoX).


Ablation Studies and Robustness Checks

Smoothing technique ordering (Tables 4, 5, 17): The paper establishes a clear effectiveness ranking: Smooth Q+K > Smooth Q only > Smooth K only > SmoothAttn > HadmdAttn ≈ None. This ordering holds across both average and worst-case metrics. Table 17's worst-case CosSim values are the most informative: 4.83% (none) → 4.85% (HadmdAttn) → 64.49% (SmoothAttn) → 90.86% (Smooth K) → 93.10% (Smooth Q) → 96.71% (Smooth Q+K). The jump from SmoothAttn to Smooth K (64.49% → 90.86%) demonstrates that independent self-smoothing vastly outperforms cross-operand difficulty transfer. The further jump from Smooth K to Smooth Q+K (90.86% → 96.71%) shows that QQ's outliers contribute independently and must be addressed separately. Table 5's end-to-end results confirm that the kernel-level ordering translates to task-level impact: Smooth Q+K is necessary to preserve Llama3.1 Lambda at 80.8% (vs. 81.5% full) and CogvideoX VQA-t at 75.147 (vs. 75.360 full).

Quantization granularity spectrum (Tables 6, 15, 19): Per-thread achieves per-token-equivalent accuracy (99.45% vs. 99.45% average CosSim) with per-block-equivalent speed (283 vs. 284 TOPS). The worst-case comparison (Table 15) confirms this is not just an average effect: per-thread worst CosSim (96.72%) is nearly identical to per-token (96.76%) while per-block drops to 90.68%. The speed ablation (Table 19) demonstrates the hardware cost of per-token dequantization: 268 TOPS vs. 283 for per-thread, a 5.3% throughput penalty that would compound with every attention layer.

Data type for P~V\tilde{P}V (Tables 7, 16): The ordering is E4M3 ≈ FP16 > E5M2 ≫ INT8. The 19.52% worst-case CosSim for INT8 in Table 16 is essentially random—the attention output bears almost no resemblance to the correct output on the worst layer. This validates the paper's core design claim that integer quantization is fundamentally unsuitable for P~\tilde{P} regardless of granularity or smoothing, and that FP8's non-uniform quantization points are necessary to preserve the small values that collectively matter.

Impact of smoothing VV (Table 10): On CogvideoX tensors, smoothing VV improves CosSim from 98.25% to 99.75% and reduces RMSE from 0.2387 to 0.0773 (3.1× reduction). This is a significant improvement, but the paper explicitly states it is model-dependent—VV smoothing helps only when VV exhibits channel-wise bias. The technique is not used in the main Table 2 results, implying that the reported SageAttn2-4b and SageAttn2-8b numbers could potentially be further improved on models like CogvideoX by enabling VV smoothing.

Overhead of proposed techniques (Table 18): All three core techniques add negligible overhead relative to a baseline INT4+FP8 attention kernel: per-thread quantization adds 0.35% (284 → 283 TOPS), two-level accumulation adds 0% (283 → 283 TOPS), and smoothing QQ adds 3.7% (283 → 273 TOPS). The smoothing overhead is the largest but is dominated by the GEMV computation for ΔSij\Delta S_{ij}; the paper notes that this GEMV is fused with the memory read of QQ and KK, partially hiding its latency. The total overhead from all techniques combined is approximately 4% (284 → 273 TOPS, or 3.9%).

Effectiveness across model architectures (Tables 2, 11, 12, 13, 20): The paper demonstrates robustness across a diverse model zoo rather than through controlled ablations of model properties. SageAttn2-8b incurs essentially zero metric loss across Llama2, Llama3.1, GLM4, CogvideoX, HunyuanVideo, Mochi, Flux, Stable-Diffusion3.5, TIMM, and Qwen2-Audio. SageAttn2-4b shows small but measurable losses on Llama3.1 (MMLU -2.8pp, Lambda -1.7pp), moderate losses on CogvideoX (VQA-a -12.5, VQA-t -17.9), and negligible losses on image generation and classification. This pattern—video generation being most sensitive, language being moderately sensitive, image generation being robust—is consistent with the observation that video diffusion models produce more extreme attention patterns (sharper softmax, larger outliers) due to the spatio-temporal attention over many frames.

Comparison with FlashAttention3(fp8) on long context (Table 14, Figure 19): On InfiniBench at 262K tokens, FlashAttn3(fp8) shows a 1.52-point average drop (43.05 → 41.53) while SageAttention2 matches full precision (43.06). The breakdown reveals that FlashAttn3(fp8)'s errors are concentrated in specific tasks: Retr.KV drops from 7.0 to 0.4, Eng.MC drops from 64.19 to 55.90, while Math.Find improves slightly (18.29 → 22.57). The catastrophic Retr.KV failure is particularly informative—it suggests that FlashAttn3(fp8)'s FP8 accumulation (without the two-level strategy) fails specifically when the attention pattern needs to precisely retrieve information from specific key positions, a task that requires accurate accumulation of many small P~V\tilde{P}V terms. The NIAH visualization (Figure 19) confirms this: FlashAttn3(fp8) shows a systematic failure band at mid-depths.

Negative result: ReSTEM^{EM}-style revision (not applicable—this is from the reference example, not this paper). SageAttention2 does not have a comparable negative result. The paper does not report any configuration where SageAttention2 underperforms expectations given its design. The FlashAttn3(fp8) comparison could be considered a negative result for that baseline rather than for SageAttention2.


Critical Assessment

Do the Experiments Support the Central Claims?

Claim: SageAttention2 is faster than FlashAttention2 and xformers by about 3× and 4.5× respectively.

This claim is well-supported for the RTX4090 GPU at moderate-to-long sequence lengths (4K–32K). Figure 5 shows SageAttn2-4b at 481 TOPS vs. FlashAttention2 at 164 TOPS (2.93×) and xformers at 105 TOPS (4.58×) at 32K, head_dim=128, causal=False. The speedup is consistent across causal/non-causal and head_dim=64/128 configurations (Figures 5, 10). However, three qualifications apply:

  1. Sequence length dependence: The speedup is smaller at short sequence lengths. At 1K tokens (head_dim=128, causal=False, Figure 5 right), SageAttn2-4b achieves 259 TOPS vs. FlashAttention2's 145 TOPS—a 1.79× speedup, not 3×. The paper's "3×" claim implicitly assumes long-sequence regimes where attention is compute-bound.
  2. GPU dependence: The speedup varies by GPU architecture. Table 9 shows 2.93× on RTX4090, 2.60× on L40, 2.46× on L20, 2.61× on H100, 3.12× on H20. The "3×" headline is accurate for RTX4090 but overstates the speedup for most datacenter GPUs. The paper appropriately disaggregates by GPU in Figure 10–16 and Table 9.
  3. The 4.5× over xformers is less meaningful than the FlashAttention2 comparison, since xformers is not the state-of-the-art exact attention kernel (FlashAttention2 is). The 3× over FlashAttention2 is the more important number.

Claim: SageAttention2 matches the speed of FlashAttention3(fp8) on Hopper GPUs while delivering much higher accuracy.

The speed-matching claim is supported on H100 and H20. Figure 14 shows SageAttn2-8b at 882 TOPS vs. FlashAttn3(fp8) at 892 TOPS (0.99×) at 32K, head_dim=128, causal=False on H100. On H20, SageAttn2-8b achieves 279 TOPS vs. 273 (1.02×). The accuracy superiority claim is supported by Table 2 (FlashAttn3-fp8 collapses on CogvideoX: VQA-a 6.531 vs. SageAttn2-8b 69.492), Table 14 (InfiniBench 41.53 vs. 43.06), and Figures 7, 9, 19 (visible degradation).

However, there is a limitation: FlashAttention3(fp8) and SageAttention2-8b are not using identical bit widths for all operations. FlashAttention3(fp8) quantizes QQ, KK, VV to FP8 throughout, while SageAttention2-8b uses INT8 for QKQK^\top and FP8 for P~V\tilde{P}V. The formats are different (INT8 vs. FP8 for QKQK^\top; both use FP8 for P~V\tilde{P}V) and the paper does not control for this. SageAttention2's accuracy advantage could come from: (a) better smoothing, (b) the two-level accumulation strategy, (c) use of INT8 instead of FP8 for QKQK^\top, or (d) some combination. The ablation in Tables 4–7 establishes that smoothing and per-thread quantization are critical, but the paper does not run an ablation where SageAttention2-8b uses FP8 for QKQK^\top (instead of INT8) to determine whether the format difference alone explains the accuracy gap. This is a minor omission since the paper's claim is the practical one ("SageAttention2 is faster and more accurate") rather than a causal claim about mechanism, but it slightly weakens support for the assertion that the specific techniques (rather than format choice) explain the advantage.

Claim: SageAttention2 incurs negligible end-to-end metrics loss across language, image, and video generation models.

This claim is strongly supported for SageAttn2-8b, with the caveat that "negligible" is not formally defined. Across all ten models in Table 2 (+ supplemental Tables 11, 12, 13, 20), SageAttn2-8b matches full-precision metrics within what appears to be evaluation noise. The largest observed deviations are: Llama3.1 WikiText 6.013 → 6.019 (0.1% increase), GLM4 Longbench 49.78 → 49.60 (0.36% decrease), HunyuanVideo VQA-a 82.516 → 81.786 (0.88% decrease). None of these are large enough to be practically meaningful, and some differences favor SageAttn2-8b (CogvideoX VQA-t: 70.928 → 74.415; TIMM ImageNet: 84.79% → 84.79%).

For SageAttn2-4b, the claim holds with qualifications. On language models, WikiText increases from 6.013 to 6.256 (4.0%), MMLU drops from 63.5% to 60.7% (2.8pp), Lambda drops from 81.5% to 79.8% (1.7pp)—these are measurable but arguably "negligible" depending on the application. On CogvideoX (1.5-5B), the drops are more substantial: VQA-a 70.231 → 57.729 (17.8%), VQA-t 70.928 → 52.989 (25.3%). Whether this is "negligible" depends on the use case—the paper's own visible comparison (Figure 6, 8) shows SageAttn2-4b outputs have minor visible differences from full precision, which a user might or might not find acceptable. The paper appropriately presents both 4-bit and 8-bit variants, allowing users to choose based on their accuracy tolerance.

Claim: The techniques (smoothing Q, per-thread quantization, two-level accumulation) are individually necessary for INT4/FP8 attention to work.

Each technique's contribution is ablated:

  • Smoothing Q: Table 4 shows Smooth K only achieves 90.86% worst-case CosSim; Smooth Q+K achieves 96.71%. Table 5 confirms end-to-end impact: without Smooth Q+K, CogvideoX VQA-t collapses to 24.670 (vs. 75.147 with). Strongly supported.
  • Per-thread quantization: Table 6 shows per-block achieves 90.68% worst-case CosSim; per-thread achieves 96.72%. Strongly supported.
  • Two-level accumulation: The paper does not present an explicit ablation disabling two-level accumulation. This is a gap. The evidence for its necessity is indirect: (a) the comparison with FlashAttn3(fp8) (Table 14, Figure 19), which the paper argues fails partly due to FP22 accumulator issues, and (b) the simulation vs. hardware discrepancy described in Section 3.4 ("while FP8 quantization for P~V\tilde{P}V above is theoretically accurate in simulation, we observe that the actual CUDA implementation suffers a consistent accuracy degradation"). A direct ablation—SageAttention2 with and without the FP32 buffer—would have strengthened this claim. The paper's Table 18 shows the overhead is 0%, so the technique is clearly worth including regardless, but the causal claim that it is "necessary" is not experimentally isolated.

Claim: The FP22 accumulator is a previously undocumented hardware limitation that affects attention accuracy.

The paper's diagnostic methodology (Section 3.4: zeroing AA and BB, varying DD) is described clearly and the discovery appears genuine. However, the paper also notes that "the two-level accumulation strategy is also implemented in CUTLASS and DeepGemm" (Section 3.4 Remark), which implies the FP22 behavior was known to NVIDIA's kernel library developers even if undocumented in the PTX ISA. The claim to novelty is specifically about discovering and investigating the effect for attention, which is credible but narrow—the hardware behavior itself appears to have been known among expert CUDA developers. The paper's contribution is surfacing this knowledge to the ML systems community and demonstrating its practical impact on attention accuracy.

Genuine Weaknesses

Single-attention-configuration kernel benchmarks: All kernel speed benchmarks (Figures 5, 10–16) use a fixed batch size of 4 and 32 attention heads. Real models vary in batch size, head count, and head dimension, and the speedup may differ for models with, say, 16 heads and head_dim=128 (common in image generation) versus 32 heads and head_dim=128 (common in language models). The end-to-end results partially address this by measuring actual model latency (Table 8), but the claimed "3× speedup" is based on a synthetic benchmark configuration, not on a distribution of real model configurations. A sensitivity analysis across head counts and batch sizes would have strengthened the generalizability claim.

No training-time evaluation: All experiments are inference-only. The paper does not evaluate whether models can be fine-tuned with SageAttention2's quantized attention (e.g., QLoRA-style adaptation). For practitioners who want to fine-tune quantized models, this is an important missing evaluation.

Single precision regime: The paper only evaluates against FP16 full-precision baselines. Modern models are increasingly trained and deployed in BF16, which has different numerical properties (wider dynamic range, less mantissa precision near 1.0). It is unclear whether the observed accuracy preservation would hold when the baseline is BF16 rather than FP16, particularly for the P~V\tilde{P}V accumulation where the FP22 truncation pattern could interact differently with BF16's 7-bit mantissa.

Difficulty estimation is oracle-based for accuracy measurement: The kernel accuracy metrics (CosSim, L1, RMSE) are computed against the full-precision attention output layer-by-layer. For end-to-end metrics, there is no oracle—the evaluation is against task-specific metrics. This is appropriate but means that the relationship between kernel-level accuracy degradation and end-to-end metric degradation is not quantitatively modeled. For instance, CogvideoX with Smooth K only has 90.86% worst-case CosSim but the end-to-end VQA-t is not reported (only the "INT4 without Smooth Q+K" row in Table 5 is shown, with VQA-t = 24.670). A systematic study of the CosSim threshold at which end-to-end metrics begin to degrade would be valuable for practitioners deciding between 4-bit and 8-bit variants.

Limited diversity in language model evaluation: The language model evaluation (Tables 2, 11) uses standard perplexity, accuracy, and long-context benchmarks, but does not include generation-quality evaluations (e.g., human evaluation, AlpacaEval, MT-Bench) that might reveal subtle degradations in open-ended text generation quality. The paper acknowledges this implicitly by focusing on discriminative metrics.

No comparison with non-quantized approximate attention: The paper compares against FlashAttention2/3 and xformers (all exact attention) but not against linear or sparse attention methods that also trade accuracy for speed. A comparison showing that SageAttention2 achieves better accuracy-efficiency tradeoffs than, say, a sparsity-based method at the same throughput would strengthen the positioning claim, though the paper's scope is explicitly "quantized exact attention."

6. Limitations and Trade-offs

The Difficulty Estimation Cost Is Not Accounted for, Making Headline Efficiency Gains an Upper Bound

The assumption or constraint. The paper's kernel throughput numbers (481 TOPS, 3× FlashAttention2) and end-to-end speedups (1.8× on CogvideoX) are measured in steady state—after the model is loaded and running. They exclude the one-time cost of preprocessing KK (global mean computation and subtraction) and, more importantly, the per-block preprocessing of QQ (per-block mean computation, smoothing, and INT4 quantization). The paper argues that these operations are "fused into a single kernel, which reads the off-chip Q and K only once" (Section 3.1) and contribute only 3.7% overhead to kernel throughput (Table 18), but this measurement is on the attention kernel itself—it does not include the cost of computing the per-block mean qˉi\bar{q}_i for every query tile, which requires summing over bqb_q tokens per tile and dd dimensions per token.

The more fundamental assumption is that models are amenable to quantized attention without per-model calibration. The paper's approach applies uniform techniques (per-thread quantization, smoothing with fixed formulas) across all models without per-model tuning of quantization parameters. This works because the smoothing formulas are derived from general properties of attention tensors (token similarity, Figure 2) rather than model-specific statistics, but the paper does not evaluate whether some models require model-specific scale calibration or whether the INT4 quantization could benefit from per-model clipping thresholds.

The consequence. In a deployment where attention is a minority of total runtime (e.g., short-sequence generation, small models with large feed-forward layers), the preprocessing overhead and the quantization/dequantization overhead become proportionally larger. The paper's end-to-end speedups (Table 8: 1.5–1.8×) implicitly account for this in the measured wall-clock time, but they only cover six specific model/GPU/length combinations. A practitioner with a model outside this set cannot reliably extrapolate the speedup from the kernel-level TOPS numbers (which show 2.5–3.1×) because the ratio of attention time to total time is unknown. The paper does not provide a formula or tool for estimating end-to-end speedup given a model's architecture and sequence length.

More subtly, the per-block mean computation for QQ is a synchronization point within the fused kernel that could limit occupancy. The paper reports 3.7% overhead for smoothing Q (Table 18) but measures this on the kernel microbenchmark with synthetic data—it does not profile occupancy or register pressure under realistic model configurations where the dd dimension (64 or 128) is small and compute utilization is already challenging to maximize.

What evidence exists in the paper. Table 8 provides six end-to-end latency measurements. Table 18 provides microbenchmark overhead for individual techniques. The paper does not measure preprocessing time separately from the attention computation, does not report the fraction of total model runtime spent in attention for any model, and does not evaluate whether per-block mean computation creates a bottleneck for models with very many attention heads (where many small tiles are processed per layer).

Mitigation status. The paper acknowledges that preprocessing is fused with memory reads to hide latency but does not treat the cost as a limitation requiring mitigation. There is no suggestion for reducing preprocessing overhead further or for models where the fused kernel's preprocessing becomes the bottleneck. This is left as an implicit engineering tradeoff rather than a recognized limitation.


The Method Provides No Gains and May Hurt Performance on the Hardest Attention Patterns

The assumption or constraint. SageAttention2 assumes that the quantized attention output is sufficiently close to the full-precision output to preserve end-to-end model behavior. The paper's kernel accuracy metrics (Tables 4, 6, 7, 15, 16, 17) measure cosine similarity, L1 distance, and RMSE between quantized and full-precision attention outputs, showing worst-case CosSim of 96.70–96.72% for SageAttn2-4b (Tables 15, 16) and better for SageAttn2-8b. These numbers are aggregated across all layers of CogvideoX.

However, the paper does not characterize which attention patterns are most vulnerable to quantization error. The worst-case layers achieve approximately 96.7% CosSim—on a per-layer basis, not accounting for error propagation through residual connections and layer normalization across dozens of layers. In deep transformers, a 3.3% error in the attention output of one layer, amplified through skip connections and nonlinearities, could produce compound effects that are not captured by single-layer accuracy metrics. The paper's end-to-end evaluation partially addresses this (by measuring final task metrics), but does not diagnose whether certain types of inputs or attention patterns are disproportionately affected—for instance, attention heads that perform precise positional lookup, heads that attend to rare tokens, or heads with highly peaked (low-entropy) softmax distributions.

The consequence. A practitioner deploying SageAttention2 on a model with particularly sharp or sparse attention patterns (e.g., retrieval-augmented models, models with very long context where attention is concentrated on a few positions) might observe degradation that is not captured by the paper's evaluations, which focus on generative models with relatively diffuse attention. The super-long context evaluation (Llama-3-262k on InfiniBench, Table 14) provides some evidence for robustness—SageAttention2 matches full precision at 262K tokens—but this is a single model and the benchmark's tasks may not stress-test the attention pattern diversity that exists across deployment scenarios.

The comparison with FlashAttention3(fp8) on CogvideoX (Table 2) shows that generic FP8 attention without smoothing fails catastrophically (VQA-a 6.531 vs. 70.231 full precision), and SageAttn2-4b partially preserves quality (VQA-a 57.729). This demonstrates that even with the paper's techniques, INT4 quantization leaves a measurable gap on video generation—the hardest evaluated modality. The paper does not investigate whether this gap is concentrated in specific attention heads, layers, or generation steps, making it difficult for practitioners to predict whether their specific model will experience similar degradation.

What evidence exists in the paper. Table 2 shows SageAttn2-4b exhibiting larger metric drops on video models (CogvideoX VQA-a: -17.8%) than on language (Llama3.1 MMLU: -4.4%) or image models (Flux FID: slightly improved). Table 15 and 16 report worst-case per-layer CosSim. The paper does not provide a distribution of per-head attention errors, a sensitivity analysis across softmax temperature/token frequency/attention sparsity, or a diagnostic for when INT4 quantization error becomes unacceptable.

Mitigation status. The paper partially mitigates this by offering two precision tiers: SageAttn2-8b for accuracy-sensitive models and SageAttn2-4b for maximum speedup. This is a practical mitigation—users can choose the 8-bit variant if 4-bit shows degradation—but it does not address the underlying fragility: there is no diagnostic or guarantee to help users determine which tier is needed for their model without running a full evaluation. The absence of a predictive relationship between kernel-level CosSim and end-to-end metric degradation means that each new model deployment requires empirical validation.


The Method Is Evaluated on a Single Precision Baseline (FP16) with No BF16 Coverage, Limiting Applicability to Modern Training and Deployment Regimes

The assumption or constraint. All experiments compare SageAttention2 against FP16 full-precision baselines. The paper states that the FP16 accumulator mode provides speedup "only on RTX4090 and RTX3090 GPUs" (Section 2.3, Table 1) and designs the FP8 P~V\tilde{P}V path to provide consistent speedup on L20, L40, and H100. However, the paper never evaluates whether the quantized attention preserves accuracy when the baseline precision is BF16 (bfloat16, 1 sign, 8 exponent, 7 mantissa bits) rather than FP16 (1 sign, 5 exponent, 10 mantissa bits).

This matters because BF16 is the default training and inference precision for many modern large models (including Llama series models, though the paper uses the FP16-converted weights). BF16 has a wider dynamic range (same as FP32) but less mantissa precision than FP16. The quantization error of SageAttention2's INT4 QKQK^\top and FP8 P~V\tilde{P}V interacts differently with BF16's error profile than with FP16's: BF16 has coarser precision near 1.0 (7 vs. 10 mantissa bits), which could make the P~V\tilde{P}V accumulation more sensitive to the FP22 truncation, or conversely, BF16's wider dynamic range could make the outlier problem in QQ and KK less severe (since extreme values are representable in BF16 but would overflow FP16).

The consequence. A practitioner using SageAttention2 with a BF16-trained model (e.g., from HuggingFace in their default BF16 weights, or deploying on TPUs or AMD GPUs where BF16 is native) cannot rely on the paper's accuracy numbers. The FP16-to-BF16 distribution shift could make the quantization error larger (if BF16's coarser mantissa amplifies small errors) or smaller (if BF16's dynamic range reduces the relative magnitude of outliers). The paper provides no guidance and no evaluation.

This limitation is particularly relevant for the FLOPs-matched comparison with larger models (Section 7 in the reference example's structure—though this paper does not have such a section, the principle applies to any practitioner comparing SageAttention2-accelerated inference against training a larger model), because the pretraining baseline might use BF16 throughout while SageAttention2 operates on FP16-converted weights.

What evidence exists in the paper. None. BF16 is not mentioned anywhere in the paper. All accuracy measurements use FP16 as the reference.

Mitigation status. Not addressed. The paper does not discuss BF16, does not suggest that results should transfer to BF16, and does not provide BF16 evaluation results. This is a scope limitation—the paper targets NVIDIA GPUs with FP16 tensor cores—but it is consequential for practitioners because BF16 is the dominant format for large-scale model training and is increasingly common for inference.


The Comparison with FlashAttention3(fp8) Does Not Isolate Which Technique(s) Explain the Accuracy Advantage, Weakening the Claim That the Proposed Methods Are Causal

The assumption or constraint. The paper claims SageAttention2 "matches the speed of FlashAttention3(fp8) on Hopper GPUs, but offers significantly higher accuracy" (Section 5, Abstract). The evidence for this claim is: (1) Table 2 shows FlashAttn3-fp8 fails on CogvideoX, HunyuanVideo, and Mochi while SageAttn2-8b preserves quality; (2) Table 14 and Figure 19 show FlashAttn3-fp8 degrades on InfiniBench and NIAH at 262K tokens while SageAttn2 matches full precision; (3) Figures 13–16 show comparable or slightly higher TOPS for SageAttn2-8b on H100/H20.

However, FlashAttn3-fp8 and SageAttn2-8b use different quantization formats for QKQK^\top: FlashAttn3-fp8 uses FP8 throughout (both QKQK^\top and P~V\tilde{P}V), while SageAttn2-8b uses INT8 for QKQK^\top and FP8 for P~V\tilde{P}V. The paper never evaluates SageAttn2-8b using FP8 for QKQK^\top to determine whether the accuracy advantage comes from (a) the smoothing and per-thread quantization techniques, (b) the use of INT8 instead of FP8 for QKQK^\top, or (c) the two-level accumulation for P~V\tilde{P}V (which FlashAttn3-fp8 presumably lacks, given its degradation pattern).

The consequence. The paper's central narrative—that the proposed techniques (smoothing Q, per-thread quantization, two-level accumulation) solve the accuracy problems that prevent aggressive quantization—is partially undermined by the confounded comparison. A skeptical reader could argue: "Maybe INT8 QKQK^\top is inherently more accurate than FP8 QKQK^\top for attention, and the techniques don't matter as much as the format choice." While the internal ablations (Tables 4–7) show that smoothing and per-thread quantization improve INT4 accuracy relative to naive INT4, they don't show that these techniques are what make SageAttention2 beat FlashAttn3-fp8 on the same format. A controlled comparison where both methods use identical formats (INT8+FP8 or FP8+FP8) with and without the paper's techniques would disentangle format choice from algorithmic contribution.

This limitation is more about strength of causal evidence than practical utility—SageAttention2 demonstrably works better than FlashAttn3-fp8 on the evaluated models regardless of why. But for researchers building on this work, the ambiguity matters: should future effort go into better smoothing and accumulation strategies, or into finding better quantization formats for attention tensors?

What evidence exists in the paper. The paper does not run any SageAttention2 configuration with FP8 QKQK^\top, nor does it ablate the two-level accumulation against FlashAttn3-fp8's accumulation strategy. The only direct comparison is the full systems (SageAttn2-8b vs. FlashAttn3-fp8) as black boxes. The internal ablations (Tables 4–7) establish that the techniques improve INT4 accuracy, but they are evaluated against FP16 P~V\tilde{P}V baselines (not against FlashAttn3-fp8's FP8 pipeline) and on CogvideoX tensors (not on the InfiniBench/NIAH tasks where the FlashAttn3-fp8 comparison is made).

Mitigation status. The paper does not acknowledge this confound. The comparison is presented as evidence that SageAttention2's techniques are responsible for the accuracy advantage, but the experimental design does not isolate the techniques from the format choice. A future study with a 2×2 design (INT8 vs. FP8 for QKQK^\top, with vs. without smoothing and two-level accumulation) would resolve this.


The Evaluation Is Limited to Autoregressive and Diffusion Generative Models; Retrieval, Classification, and Embedding Models Are Unexplored

The assumption or constraint. The paper evaluates SageAttention2 on language models (Llama2, Llama3.1, GLM4), video diffusion models (CogvideoX, HunyuanVideo, Mochi), image diffusion models (Flux, Stable-Diffusion3.5), an image classification model (TIMM), and an audio model (Qwen2-Audio). This covers generative and discriminative tasks across four modalities. However, the paper does not evaluate bidirectional encoder models (e.g., BERT, RoBERTa), embedding models (e.g., text-embedding-3, E5, or any bi-encoder), or retrieval models where the attention output is used for similarity computation rather than next-token prediction.

These model classes differ from the evaluated models in an important way: their attention patterns are often fully bidirectional (no causal mask) and may produce attention outputs that are used directly (as embeddings for retrieval) rather than propagated through many layers of autoregressive generation. In embedding models, a small error in the attention output directly corrupts the final embedding, which is used for nearest-neighbor search—there is no autoregressive decoding process to potentially "average out" the error across tokens. In bidirectional models, the attention is over the full input sequence simultaneously, which could create different outlier patterns than the causal attention evaluated in the paper.

The consequence. A practitioner deploying SageAttention2 for a retrieval pipeline—where embedding quality is the product and even 1–2% degradation in cosine similarity between query and document embeddings translates directly to recall loss—cannot rely on the paper's evaluation. The paper's closest evaluation is TIMM for image classification (Table 13), which uses bidirectional attention and shows no degradation, but classification accuracy is a coarse metric that may be insensitive to small embedding shifts (a 0.1% accuracy change could hide substantial embedding-space distortion that would affect k-NN retrieval).

Similarly, the paper does not evaluate on cross-attention between different modalities (e.g., in multimodal models like LLaVA or image-to-text models), where QQ comes from one modality and K,VK, V from another. The distribution of QQ relative to KK in cross-attention could differ substantially from self-attention, potentially creating quantization errors that self-attention evaluations do not capture.

What evidence exists in the paper. None. The paper does not evaluate any retrieval, embedding, or cross-attention models. The evaluated models are all either autoregressive (Llama, GLM4, CogvideoX, HunyuanVideo), diffusion-based (Flux, Stable-Diffusion3.5), or classification (TIMM). Qwen2-Audio (Table 20) involves cross-modal attention (audio to text), but only the ASR task is evaluated—not embedding quality or retrieval.

Mitigation status. The paper does not discuss this as a limitation. The model selection is broad across modalities but narrow across architectural paradigms within attention. The claim that SageAttention2 "can accelerate models in a plug-and-play way with negligible loss in end-to-end metrics" (Section 1) implicitly asserts universality, which is not validated for the entire space of attention-based architectures.


The Paper Does Not Evaluate the Sensitivity of End-to-End Quality to Per-Layer Quantization Error, Leaving No Guidance for Diagnosing or Predicting Failures on New Models

The assumption or constraint. The paper's evaluation strategy is binary: run the full model with SageAttention2, measure end-to-end metrics, and report whether degradation is "negligible." This works for the ten evaluated models but provides no transferable guidance. The kernel accuracy metrics (CosSim, L1, RMSE) are measured per-layer and reported as averages and worst-cases, but the paper never establishes a relationship between these kernel-level metrics and end-to-end task degradation. For instance, at what per-layer CosSim does MMLU start to drop? Is the worst layer the bottleneck, or do errors accumulate across many moderately-degraded layers? Does a 96.7% worst-case CosSim (SageAttn2-4b, Table 15) guarantee acceptable end-to-end quality, or is that specific to CogvideoX's architecture?

The consequence. A practitioner with a new model cannot look at the paper's results and determine whether SageAttention2 will work for their model without running a full evaluation. They cannot even estimate: "my model's worst-layer CosSim at INT4 is 95%—should I expect visible degradation?" The paper provides neither a threshold nor a predictive model. This is a barrier to adoption because evaluating end-to-end metrics on a new model (especially a generative model requiring human evaluation or expensive automated metrics) is costly. A lightweight diagnostic—perhaps measuring per-layer CosSim and flagging layers below a threshold—would dramatically reduce the adoption friction, but the paper does not develop or validate such a diagnostic.

What evidence exists in the paper. Tables 4, 6, 7, 15, 16, 17 provide kernel accuracy metrics. Table 2 provides end-to-end metrics. These are never plotted against each other or analyzed jointly. The paper does not report per-layer CosSim for the models in Table 2, does not correlate CosSim with end-to-end metric drops, and does not identify which layers are the accuracy bottlenecks for SageAttn2-4b's degradation on CogvideoX.

Mitigation status. The paper does not address this gap. It provides the data needed to investigate the relationship (per-layer kernel metrics on CogvideoX, end-to-end metrics on CogvideoX with both 4-bit and 8-bit) but does not perform the analysis. This is a missed opportunity: the CogvideoX data alone could reveal whether the VQA-t drop from 70.928 (8-bit) to 52.989 (4-bit) is driven by a few outlier layers (which might be individually fixable with layer-specific quantization parameters) or by pervasive small errors across all layers (which would require a fundamentally different approach). A practitioner debugging SageAttention2-4b on their own model has no guidance from the paper on how to triage failures.

7. Implications and Future Directions

How This Work Changes the Landscape

SageAttention2 fundamentally changes how the ML systems community should think about quantization in attention by demonstrating that the barrier to low-bit attention acceleration is not a fundamental precision limit of the operation itself, but rather a set of identifiable, solvable failure modes—outlier sensitivity in Q/K, undocumented hardware accumulator truncation, and misaligned quantization granularity. Prior to this work, the implicit consensus (formed by the failure of naive INT4 attention and the partial success of INT8 in SageAttention) was that attention's dynamic, outlier-heavy tensor distributions made INT4 quantization effectively impossible for exact attention—a conclusion the paper disproves through systematic diagnosis and targeted mitigation.

This is not a paradigm shift in the sense of introducing a new attention algorithm; SageAttention2 computes exact attention, just faster. Rather, it is a reframing of the quantization problem for attention from a format-selection problem to a hardware-software co-design problem. The paper's central insight—that quantization granularity should be defined by the GPU's thread-to-data mapping (the MMA instruction layout) rather than by tensor dimensions—elevates the CUDA thread scheduler from an implementation detail to a first-class design parameter in quantization strategy. This reverses the conventional direction of influence between algorithms and hardware: instead of designing quantization to match tensor statistics and then engineering the kernel to implement it, the paper designs quantization to match the kernel's thread layout and then verifies that accuracy is preserved. This principle is likely transferable to any fused quantized kernel where the hardware's data distribution across threads is fixed and known—a broad class that includes fused MLP blocks, custom attention variants, and mixture-of-experts routing.

The paper also resolves a latent contradiction in the quantized attention literature through its FP22 accumulator discovery. Prior work had produced conflicting signals about whether FP8 attention is viable: FlashAttention3(fp8) works well on some models (the paper acknowledges this by benchmarking against it) but the paper shows it degrades severely on video generation (CogvideoX VQA-a: 6.531 vs. 70.231 full, Table 2) and long-context retrieval (InfiniBench Retr.KV: 0.4 vs. 7.0, Table 14). The paper's diagnosis—that the FP22 accumulator silently truncates precision during ~PV accumulation, and that this truncation matters more when many small softmax values must sum accurately—provides a unified explanation: FP8 attention works when the attention pattern is peaked (few large softmax values, accumulation error is small relative to signal) but fails when the attention pattern is diffuse or requires precise aggregation across many positions. This explanation simultaneously validates both the successes and failures of prior FP8 attention work and redirects research attention from operand quantization (how to quantize ~P and V) to accumulation precision management (how to prevent the FP22 accumulator from corrupting the result).

Equally important, the paper establishes that exact attention quantization is now on equal footing with algorithmic approximations in terms of the accuracy-efficiency tradeoff. Prior to SageAttention2, practitioners choosing between fast attention methods faced a stark choice: accept the exactness-guarantee of FlashAttention2 (at FP16 throughput) or accept the approximation risk of sparse/linear attention (at potentially higher speed). SageAttention2's 3× speedup over FlashAttention2 with negligible metric loss on 8-bit across ten models and modalities (Table 2) means that the accuracy-guaranteed path now achieves comparable or better throughput than many approximate methods, without their model-specificity constraints. This narrows the application domain where approximate attention is the right engineering choice: if SageAttn2-8b works for your model (and the paper shows it works for a diverse set at 1.8× end-to-end speedup), approximate attention only makes sense when SageAttn2-8b's speedup is insufficient—a much higher bar than before.

The research direction that becomes more attractive is robust verifier-free quantization. The paper shows that with careful attention to accumulator precision and quantization granularity, aggressive quantization (INT4) can preserve end-to-end quality without requiring a separate verification model, calibration dataset, or per-model tuning. This suggests that the ML systems field may have over-invested in adaptive, calibration-heavy quantization methods (per-model scale search, mixed-precision layer assignment, outlier-aware clipping) when the bottleneck was simpler: the formats and granularities were mismatched to the hardware.

The research direction that becomes less attractive is the development of ever-more-approximate attention algorithms for the purpose of inference speedup on standard GPU hardware. If exact attention can be accelerated 3× through quantization with negligible quality loss, the marginal benefit of switching to an approximate method (which may provide another 1.5–2× but at the cost of model-specific tuning and potential failure on certain tasks) shrinks considerably. The paper does not argue this explicitly, but the numbers in Table 2 and Figure 5 imply it: approximate attention methods would need to demonstrate speedups substantially beyond 3× (on the same hardware, at the same sequence lengths) while matching SageAttn2-8b's model-agnostic quality preservation, or they need to target hardware where INT4/FP8 tensor cores are unavailable.

Follow-Up Research This Work Enables

Training-time quantized attention with SageAttention2's per-thread granularity. The paper evaluates only inference, but the per-thread quantization mapping (Equation 8, Figure 4) is equally applicable to the backward pass of attention, where the gradients of Q, K, V must be computed from the output gradient. The backward pass involves additional matrix multiplications (dQ = dO · K, dK = dO^T · Q, dV = ~P^T · dO) that have the same structure as the forward pass but with potentially different outlier distributions (gradients can be noisier and spikier than activations). A natural extension would evaluate whether per-thread INT8 quantization for the backward pass (INT4 may be too aggressive for gradients, which have higher dynamic range) preserves training dynamics. The specific experiment: fine-tune Llama3.1-8B on a standard dataset (e.g., Alpaca) using SageAttention2-quantized attention for both forward and backward passes, and compare training loss curves, downstream task performance, and throughput against FP16 training with FlashAttention2. The key metric is whether the quantized training reaches the same loss within the same number of steps, or whether gradient noise from quantization requires learning rate adjustment.

Systematic characterization of the CosSim threshold for end-to-end degradation across model architectures. The paper reports per-layer CosSim for CogvideoX (Tables 15, 16) and end-to-end metrics for multiple models (Table 2) but never connects the two. A critical follow-up would systematically vary quantization aggressiveness (by sweeping the smoothing configuration: none, K-only, Q-only, Q+K; and the quantization granularity: per-tensor, per-block, per-thread, per-token) across multiple models, measure both per-layer CosSim distributions and end-to-end task metrics, and identify the CosSim threshold at which each model's end-to-end quality begins to degrade. The hypothesis is that this threshold is model-specific (video models are more sensitive, as suggested by Table 2), and the finding would enable lightweight deployment decisions: measure per-layer CosSim on a few calibration inputs, and if the worst layer exceeds the model-specific threshold, fall back to SageAttn2-8b; otherwise, use SageAttn2-4b. The paper's existing CogvideoX data (Tables 4, 6, 7 + Table 2 CogvideoX rows) provides a starting point but lacks the intermediate CosSim levels (e.g., what happens to VQA-t at 95% worst-case CosSim? 92%?).

Isolating the contribution of two-level accumulation via a controlled comparison with FlashAttention3. The paper's comparison with FlashAttention3(fp8) on Hopper GPUs is confounded: the two methods use different QK^T formats (INT8 vs. FP8) and different accumulation strategies for ~PV. A clean experiment would implement three variants of SageAttention2's ~PV path: (a) with two-level accumulation (the current version), (b) without two-level accumulation (single-level FP22 accumulation across all tiles), and (c) with explicit FP32 emulation of the MMA accumulator (if feasible at reduced throughput). Run all three on the InfiniBench and NIAH evaluations at 262K tokens where FlashAttn3-fp8 shows degradation (Table 14, Figure 19). If variant (b) reproduces FlashAttn3-fp8's failure pattern (Retr.KV dropping to near zero, NIAH mid-depth band failure) while variant (a) matches full precision, this would isolate the two-level accumulation as the causal mechanism and provide the missing ablation. If variant (b) does not reproduce the failure, then the accuracy advantage comes primarily from SageAttention2's QK^T quantization (INT8 + smoothing + per-thread) rather than the FP22 mitigation, which would redirect future work toward better QK^T quantization strategies on Hopper.

Extending per-thread quantization to mixed-precision layer assignment. SageAttention2 applies INT4 uniformly to all attention layers. However, the per-layer CosSim results (Tables 15, 16) show substantial variation: worst-case CosSim is 96.7% for SageAttn2-4b but average is 99.45%, indicating that some layers are near the accuracy threshold while others have substantial headroom. A mixed-precision approach—applying INT8 to the sensitive layers and INT4 to the robust ones, selected based on per-layer CosSim measured on a small calibration set—could recover most of SageAttn2-4b's speedup while achieving SageAttn2-8b's end-to-end accuracy. The specific experiment: for each model, measure per-layer CosSim at INT4 and INT8, sort layers by INT4 error, promote the K most-sensitive layers to INT8, and measure end-to-end metrics as a function of K. The paper's per-thread quantization design makes this feasible because the same MMA instruction layout works for both INT4 and INT8—only the scale computation and packing differ—allowing per-layer format switching with minimal kernel branching overhead.

Diagnosing and mitigating the specific failure modes on video generation models at INT4. SageAttn2-4b shows a 17.8% drop in VQA-a on CogvideoX (70.231 → 57.729, Table 2), the largest degradation among all models. An open question is whether this degradation is caused by a few outlier attention heads (which might have extreme ~P distributions that E4M3 cannot capture), specific layers (e.g., early layers where Q and K distributions are less smoothable), or a pervasive small error across all layers that compounds through the diffusion denoising process. The experiment: on CogvideoX, run SageAttn2-4b but collect per-head and per-layer attention output CosSim during generation, then correlate the per-head errors with the final video quality metrics (CLIPSIM, VQA-a, VQA-t, FScore). If a small number of heads dominate the degradation, targeted mitigation (e.g., applying SageAttn2-8b only to those heads, or enabling V smoothing selectively on layers where V has channel-wise bias per Section 3.4) could close most of the gap. Alternatively, if the error is pervasive, it suggests that CogvideoX's attention patterns are fundamentally more quantization-sensitive, and practitioners should default to SageAttn2-8b for video diffusion models.

Cross-architecture validation on non-NVIDIA hardware and non-GPU accelerators. The paper evaluates exclusively on NVIDIA GPUs (RTX4090, L20, L40, H100, H20). The per-thread quantization method is designed around NVIDIA's specific MMA instruction layout (mma.m16n8k64), which maps threads to output elements in a documented pattern (Figure 4). Other accelerators—AMD GPUs (with Matrix Core instructions), Intel GPUs (with XMX engines), Apple Neural Engine, and custom AI accelerators (TPUs, Inferentia)—have different matrix multiply instruction layouts, thread mappings, and accumulator precisions. A direct replication on AMD hardware (using ROCm's equivalent matrix instructions) would be the most immediate extension: does the per-thread quantization principle (align quantization groups to the hardware's thread ownership of output elements) transfer to AMD's Matrix Core layout with minimal modification to the grouping formula (Equation 8), or does the different architecture require re-deriving the grouping from scratch? On non-GPU accelerators (especially TPUs, which use systolic arrays rather than warp-based execution), the per-thread concept itself may not apply—the quantization group would need to be defined by the systolic array's data distribution, which is a fundamentally different mapping problem. A finding that the principle (hardware-data-layout-aligned quantization) transfers across architectures would establish it as a general design pattern; a finding that it is NVIDIA-specific would bound the contribution appropriately.

Practical Applications and Downstream Use Cases

Real-time video generation and editing with consumer GPUs. CogvideoX (1.5-5B) takes 1040 seconds (17.3 minutes) to generate a single video on an RTX4090 using FlashAttention2 (Table 8). SageAttn2-8b reduces this to 577 seconds (9.6 minutes), a 1.8× speedup with zero quality loss (Table 2: VQA-a 69.492 vs. 70.231, visible comparison in Figure 7 shows no discernible difference). For a video creator iterating on prompts—generating, reviewing, tweaking, regenerating—this cuts the feedback cycle nearly in half. On the smaller CogvideoX (2B), the generation time drops from 86 seconds to 54 seconds, making interactive video generation (under one minute per generation) feasible on a single consumer GPU. The practical impact is that real-time or near-real-time video generation workflows, previously requiring datacenter-grade hardware or acceptance of multi-minute latencies, become viable on a single RTX4090 workstation costing a few thousand dollars. The 1.8× speedup is end-to-end, not just attention-kernel-level—it accounts for the full model pipeline. For HunyuanVideo on L20 (2221s → 1486s, 1.49×) and Mochi on L20 (2336s → 1316s, 1.77×), the absolute time savings are even larger (12–17 minutes saved per generation), directly reducing cloud GPU rental costs for video generation API providers.

Long-context language model serving on cost-sensitive hardware. Llama3.1 processing a 100K-token prompt on an L20 GPU takes 39.9 seconds to generate the first token using FlashAttention2; SageAttn2-8b reduces this to 25.4 seconds (1.57×, Table 8). For a retrieval-augmented generation (RAG) pipeline where a user query triggers retrieval of dozens of documents totaling tens of thousands of tokens, this latency reduction directly improves user experience—the time between submitting a query and seeing the first response token drops by over a third. More importantly, the paper's super-long context evaluation (Table 14, Figure 19) shows that SageAttention2 preserves retrieval accuracy at 262K tokens (InfiniBench 43.06 vs. 43.05 full precision), meaning the speedup does not sacrifice the quality of long-context retrieval—a critical requirement for RAG applications where missed retrieval directly causes incorrect answers. On L20 GPUs (a datacenter Ada Lovelace card that is significantly less expensive than H100), this makes long-context LLM serving viable at lower cost tiers: a provider can serve 100K-token prompts on L20s at 25 seconds per first token with SageAttn2-8b, compared to needing H100s (or accepting ~40-second latency on L20s) with FlashAttention2. The speedup also applies to the prefill phase across all tokens, not just the first token, reducing total time-to-completion for long-context generation tasks like summarization of lengthy documents.

Batch image generation for content production pipelines. The image generation results (Table 2) show that SageAttn2-4b achieves slightly better FID than full precision on Flux (10.577 vs. 10.960) and near-identical metrics on Stable-Diffusion3.5 (FID 14.097 vs. 14.105), with kernel-level speedups of ~3× over FlashAttention2. For a content production pipeline generating thousands of images (e.g., an e-commerce platform producing product images from descriptions, or a game studio generating texture variations), the end-to-end speedup translates directly to throughput: generating 1,000 images that previously took 1,000 × (inference time) now takes 1,000 × (inference time / 1.5–1.8). On RTX4090 with a model like Flux, where attention is a significant fraction of total inference time, the practical impact is a ~40–45% reduction in total GPU-hours for batch generation. Critically, the paper shows that the speedup does not require accepting visible quality tradeoffs—SageAttn2-8b's image metrics are within evaluation noise of full precision (CLIP 26.175 vs. 26.180, ImageReward 1.009 vs. 1.009), so production pipelines can adopt the faster kernel without re-qualifying their outputs. The plug-and-play nature (SageAttention2 is a drop-in replacement with the same API as FlashAttention) minimizes integration cost.

On-device or edge deployment of medium-sized language models. A 7–9B parameter language model (Llama2-7B, Llama3.1-8B, GLM4-9B) running on a single RTX4090-class GPU with SageAttn2-8b achieves MMLU accuracy within 0.1% of full precision (Table 2: Llama3.1 63.4% vs. 63.5%, GLM4 74.5% vs. 74.3%) while generating first tokens 1.6× faster at long context lengths (Table 8: Llama3.1-48K: 5.7s vs. 9.2s). This combination—near-lossless accuracy at substantially reduced latency—is particularly valuable for interactive applications like coding assistants, where the model must process long context (entire codebases or conversation histories) and respond quickly. A coding assistant processing a 50K-token context window can respond in ~6 seconds instead of ~9.5 seconds, crossing from "noticeable delay" to "near-interactive" territory without any model quality compromise. The fact that SageAttention2 is a kernel replacement (not a model modification) means it can be applied to any fine-tuned variant of these base models without recalibration, making it suitable for the ecosystem of Llama-based fine-tuned models (instruction-tuned, domain-specific) deployed on edge servers or local workstations.