SageAttention: Accurate 8-Bit Attention for Plug-and-play Inference Acceleration

ArXiv: 2410.02367

🎯 Pitch

Attention is 4× slower than other transformer ops at long sequence lengths—but directly quantizing it to INT8 destroys accuracy. SageAttention solves this by smoothing key outliers and keeping a mixed-precision PV path, hitting 341 TOPS (2× over FlashAttention2) with zero accuracy loss, even boosting TIMM image models.

---

1. Executive Summary

This paper proposes SageAttention, a post-training quantization method that accelerates the attention mechanism in transformers by quantizing matrices to INT8 while preserving end-to-end model accuracy across language, image, and video generation tasks. Targeting the quadratic complexity bottleneck of standard attention, SageAttention introduces two core techniques to address the accuracy degradation that occurs with naive quantization: smoothing K (subtracting the per-channel mean of the key matrix before quantization to eliminate channel-wise outliers without changing attention scores) and FP16 accumulation for PV (retaining the softmax output and value matrix in FP16 with an FP16 accumulator rather than quantizing them to INT8, yielding both higher accuracy and 2× faster Matmul on consumer GPUs). The approach also employs adaptive quantization, selecting between four kernel variants with different speed-accuracy tradeoffs per layer. SageAttention achieves approximately 2.1× speedup over FlashAttention2 and 2.7× over xformers on RTX4090, reaching up to 341 TOPS, while incurring negligible end-to-end metric degradation — a 0.2% average loss on Llama2, CogvideoX, UltraPixel, and Unidiffuser — and on TIMM image classification actually surpasses full-precision attention. The method operates as a plug-and-play replacement for existing attention implementations on consumer-grade GPUs (RTX4090/3090), establishing that accurate INT8 attention is achievable through a combination of outlier mitigation and mixed-precision accumulation without requiring retraining or specialized GPU architectures.

2. Context and Motivation

The Core Problem: Attention Dominates Inference but Remains Unquantized

The transformer architecture (Vaswani, 2017) has become the de facto standard across virtually every domain of deep learning — language models, image generation, video generation, and beyond. At its heart lies the self-attention mechanism, which computes pairwise interactions between all tokens in a sequence. This mechanism has a computational complexity of $O(N^2)$ in sequence length $N$, compared to $O(N)$ for linear transformations. As the field pushes toward processing increasingly longer sequences — language model prefilling handles 8K–128K tokens (Dubey et al., 2024), video generation models process thousands of frames (Yang et al., 2024) — attention rapidly becomes the dominant computational bottleneck.

Figure 2 in the paper makes this concrete: at a sequence length of 32K with 32 attention heads and head dimension 64, attention consumes approximately 800ms of latency while linear layers and other operations combined take roughly 200ms. This is not a minor overhead — it's a 4× dominance. The practical consequence is straightforward: if you want to deploy transformer models at scale or on consumer hardware, you must accelerate attention.

This problem is important for several interconnected reasons:

• Cost and accessibility of deployment. FlashAttention2 (Dao, 2023) on an RTX4090 achieves roughly 165 TOPS (tera-operations per second) at headdim=64. The theoretical peak of the RTX4090's tensor cores in FP16 is far higher — the hardware is underutilized. Closing this utilization gap translates directly to lower latency, higher throughput, and reduced infrastructure costs for anyone serving transformer-based models.

• The long-sequence trend. Language models are being pushed to handle 128K, 256K, and even million-token contexts. Video generation requires processing tens of thousands of spatial-temporal tokens. In both regimes, the quadratic attention cost means that even small multiplicative speedups in attention translate to massive absolute latency reductions. A 2× faster attention kernel at 128K sequence length saves proportionally more wall-clock time than at 1K.

• The quantization gap in transformer components. The broader field has invested heavily in quantizing linear layers — weights and activations compressed to INT8, FP8, or even INT4 to accelerate matrix multiplications (Jacob et al., 2018; Xiao et al., 2023a). However, attention has been conspicuously left untouched by these quantization efforts. As the paper states directly: "There is not yet a work that systematically investigates the quantization of attention." This creates an asymmetry where the linear layers run in low precision while attention remains in FP16, constituting a growing fraction of total runtime as sequence lengths increase.

Prior Approaches and Their Limitations

The paper organizes existing approaches to efficient attention into three categories (Section 2), which is helpful for understanding how SageAttention fits into the landscape.

Sparse Attention methods reduce computation by selecting only a subset of token pairs for attention computation. Examples include Swin Transformer (Liu et al., 2021), which restricts attention to local windows; Attention Sinks (Xiao et al., 2023b), which keeps only initial "sink" tokens plus a sliding window; and Minference (Jiang et al., 2024), which dynamically identifies important query-key interactions. These approaches have an inherent limitation: they change the attention computation itself, discarding information that might be relevant. As the paper notes, "omitted calculations are not always useless." In practice, sparse methods tend to work well in specific domains (e.g., vision transformers with local structure) but fail to generalize across the full diversity of transformer applications. A text-to-video model or a long-context language model often requires truly global attention.

Linear Attention methods reformulate attention to achieve lower asymptotic complexity, typically $O(N)$ rather than $O(N^2)$. Examples include Linformer (Wang et al., 2020), which projects the key-value dimension to a fixed size; Performer (Choromanski et al., 2020), which uses random feature approximations for the softmax; and LinearAttention (Katharopoulos et al., 2020), which rearranges the computation order. The paper acknowledges that these "excel in specific scenarios," but standard attention — with its exact softmax and unrestricted pairwise interactions — remains prevalent. Many state-of-the-art models are architected around standard attention and replacing it with a linear approximation would require retraining or even rethinking the model design.

Kernel Optimization methods do not change the attention computation mathematically. Instead, they exploit hardware properties to execute the same computation faster. This is the category SageAttention belongs to, and it's worth examining the specific approaches it builds on and competes with:

• xformers (Lefaudeux et al., 2022) provides modular, optimized CUDA kernels for attention with custom block structures. It achieves speedups over naive PyTorch implementations but leaves throughput on the table compared to more specialized approaches.

• FlashAttention (Dao et al., 2022) introduced tiling to reduce memory reads/writes between GPU global memory (HBM) and on-chip SRAM. The key insight was that the $N \times N$ attention score matrix $S = QK^T/\sqrt{d}$ and the softmax output $P$ are too large to materialize in SRAM (which is fast but small, typically tens to hundreds of KB), but the $Q$ and $K$ matrices are only $N \times d$ where $d$ is small (e.g., 64–128). By tiling $Q$, $K$, and $V$ into blocks along the token dimension and computing attention incrementally using an online softmax algorithm (Milakov & Gimelshein, 2018), FlashAttention avoids ever writing the full $S$ or $P$ matrices to slow global memory. The result is a significant speedup over naive implementations.

• FlashAttention2 (Dao, 2023) refined the parallelism strategy and warp partitioning of the original FlashAttention, achieving better GPU utilization, particularly at lower sequence lengths. It became the standard baseline for fast attention implementations on non-Hopper GPUs.

• FlashAttention3 (Shah et al., 2024) was released for Nvidia's Hopper architecture (H100, H800 GPUs) and includes an FP8 quantization mode. However, it has three critical limitations that motivate SageAttention. First, it is exclusive to Hopper GPUs, meaning it cannot be used on widely deployed consumer hardware (RTX4090, RTX3090) or previous-generation datacenter GPUs (A100). Second, as the paper demonstrates in Table 1, FlashAttention3's FP8 quantization degrades accuracy significantly across models — Llama2 perplexity on WikiText jumps from 5.823 (FP16) to 5.850 (FlashAttn3 with quantization), CogvideoX Fscore drops from 3.768 to 3.394, and Unidiffuser FID explodes from 163.33 to 394.13. Third, even if accuracy were acceptable, FP8 Matmul is slower than INT8 Matmul on consumer GPUs — the paper notes that INT8 Matmul on RTX4090 and 3090 is four times faster than FP16 and two times faster than FP8 (Section 1). So FlashAttention3's FP8 approach is architecturally suboptimal for non-Hopper hardware.

• I-BERT (Kim et al., 2021) quantizes all tensors in a transformer block, including attention, to INT8. However, it is restricted to the RoBERTa architecture and requires quantization-aware training — the network must be trained or fine-tuned with quantization in the loop to maintain accuracy. This makes it impractical for the post-training, plug-and-play deployment scenario SageAttention targets. No one wants to retrain a Llama2-7B or CogvideoX just to enable quantized attention.

The Specific Gap: Accurate, Plug-and-Play Quantized Attention

The landscape reveals a clear gap. Prior work has produced fast FP16 attention (FlashAttention2), architecture-specific quantized attention (I-BERT), and Hopper-exclusive quantized attention with accuracy degradation (FlashAttention3 FP8). What is missing is a method that satisfies three criteria simultaneously:

1. Accurate: Quantized attention should produce results nearly identical to full-precision attention across diverse model types (language, image generation, video generation) without requiring retraining.

2. Plug-and-play: The method should serve as a drop-in replacement for existing attention implementations — no model modification, no fine-tuning, no architecture-specific tuning.

3. Fast on consumer hardware: The method should exploit the hardware capabilities of widely available GPUs (RTX4090, RTX3090) where INT8 tensor core throughput is substantially higher than FP16 or FP8.

Why Quantizing Attention Is Harder Than Linear Layers

The paper identifies two specific technical challenges (C1 and C2, Section 1) that explain why quantizing attention has resisted prior attempts. Understanding these challenges is crucial for appreciating SageAttention's design choices.

Challenge C1: Channel-wise outliers in the K matrix. The key matrix $K$ exhibits a distinctive pattern not seen in typical linear layer activations: there are large channel-wise biases shared across all tokens. Figure 4 visualizes this for CogvideoX and Unidiffuser — the color maps show that certain channels of $K$ have values that are consistently large in magnitude (e.g., -210 to 208 in one channel) compared to the token-wise variation. This means each token's key vector is essentially a large channel-specific bias plus a small token-specific signal.

Why does this break quantization? Quantization assigns a single scale factor per tensor, per token, or per block. If a channel has a large bias, that bias dominates the scale computation, and the small token-wise variation — which carries the actual information about which tokens attend to which — gets crushed into zero or near-zero values after rounding to the limited INT8 range. The result is a catastrophic loss of information.

The problem is compounded because per-channel quantization of K is not feasible. Quantization can only be applied along the outer axis of the Matmul $QK^T$. If you quantize $K$ per-channel (along the head dimension), the scale factors are on the inner dimension, and you cannot use them for dequantization of the output because each scale would correspond to a different column of the product, which is not how matrix multiplication works. The paper cites Xiao et al. (2023a) on this point, which established similar constraints for linear layer quantization.

Furthermore, the SmoothQuant technique (Xiao et al., 2023a), which handles outliers in linear layers by migrating the quantization difficulty from activations to weights, cannot be directly applied to attention. SmoothQuant works by scaling activations down and weights up by a per-channel factor, exploiting the fact that weights are static and can absorb scale. In attention, both $Q$ and $K$ are dynamic (they depend on the input) and both can have outliers — there's no static weight matrix to absorb the scale. As the paper notes: "Q is also heavily affected by outliers."

Challenge C2: Unreliable INT8 quantization of the PV Matmul. The second matrix multiplication in attention is $PV$, where $P$ is the softmax output (rows sum to 1, values between 0 and 1) and $V$ is the value matrix. Quantizing both $P$ and $V$ to INT8 seems natural — it's symmetric with the $QK^T$ quantization. However, Table 3 reveals that this approach has catastrophic worst-case behavior: across all layers of Llama2 and Unidiffuser, the worst-case INT8 quantization of $P$ and $V$ achieves only 56.40% cosine similarity with a relative L1 error of 0.792 and RMSE of 0.541. This is dramatically worse than the FP16 baseline (99.99% cosine similarity). Even the best 8-bit alternative (E4M3 for $P$ and $V$) achieves only 76.36% cosine similarity in the worst case.

This worst-case failure matters because it occurs in specific layers of specific models. If even one attention layer produces severely degraded output, the error propagates through the rest of the network, potentially causing the complete model failure observed in Figure 3 — where Unidiffuser with INT8 attention generates a completely blurry image, and Llama2 with INT8 attention drops to 25.5% MMLU accuracy (random-guessing level for the 4-option questions). A practical quantization method cannot have these worst-case failure modes.

How SageAttention Positions Itself

The paper positions SageAttention as a direct response to these two challenges, offering solutions that are specifically designed for the attention computation rather than adapting techniques from linear layer quantization. The approach is grounded in careful analysis of the data distributions in real attention layers (Figure 4) and systematic comparison of different quantization granularities and data types (Tables 2–5).

The key intellectual move is recognizing that the two Matmuls in attention have fundamentally different sensitivity to quantization and should be treated differently. For $QK^T$, the main issue is the $K$ outliers, which can be addressed by a mathematically-equivalent smoothing transformation that removes the channel-wise bias without changing the attention scores. For $PV$, the solution is counterintuitive but effective: don't quantize at all. Instead, keep $P$ and $V$ in FP16 but use an FP16 accumulator for the Matmul, which is faster than FP32 accumulation on consumer GPUs and eliminates all quantization error for this operation.

SageAttention also positions itself as hardware-aware but not hardware-exclusive. Unlike FlashAttention3, which is tied to Hopper GPUs, SageAttention targets the INT8 tensor core instructions (mma.u8.u8.s32) and FP16-with-FP16-accumulator instructions (mma.f16.f16.f16) available on widely deployed consumer GPUs (RTX4090, RTX3090). This makes the method accessible to a much broader user base — individual researchers, small labs, and anyone running models on gaming or workstation GPUs.

Finally, the paper's experimental positioning is deliberately broad. The evaluation spans language models (Llama2, Llava1.6), text-to-image generation (Unidiffuser, UltraPixel), text-to-video generation (CogvideoX), and image classification (TIMM) — covering the major application domains of transformers. This breadth is intended to demonstrate that SageAttention's approach is not model-specific or task-specific, but rather a general solution to the attention quantization problem.

3. Technical Approach

3.1 Reader Orientation

SageAttention is a set of four carefully designed GPU kernels that replace the standard attention computation in transformer models with a quantized version that runs 2–3× faster. The core problem it solves is that naive quantization of attention — simply casting the query, key, and value matrices from FP16 to INT8 — destroys model accuracy because the key matrix $K$ has channel-wise outliers that get crushed during quantization, and the second matrix multiplication $PV$ is too sensitive for 8-bit representation in the worst case. The solution has an asymmetric shape: the first half of attention ($QK^T$) runs in INT8 after applying a mathematically-equivalent smoothing operation to $K$ that removes the outliers, while the second half ($PV$) stays in FP16 but uses a faster FP16 accumulator rather than the default FP32 accumulator, eliminating quantization error entirely for that operation.

3.2 Big-Picture Architecture (Diagram in Words)

The system has five major components:

1. A smoothing preprocessor for K — Before attention begins, each key matrix $K$ has its per-channel mean vector subtracted. This removes the large channel-wise bias that dominates the INT8 quantization range, but it is mathematically guaranteed to not change the attention scores because the bias cancels out in the softmax. The smoothed $K$ is then quantized to INT8 with a per-token, per-block, or per-tensor granularity.

2. An INT8 quantizer for QK^T — The query matrix $Q$ (scaled by $1/\sqrt{d}$) and the smoothed key matrix $\hat{K}$ are both quantized to INT8. Their matrix multiplication runs on the GPU's fast INT8 tensor core instruction mma.u8.u8.s32, producing an INT32 accumulator output. The result is then dequantized back to FP32 by multiplying by the product of the per-block scale factors. The online softmax operates in full precision on this dequantized product.

3. An FP16-with-FP16-accumulator path for PV — The softmax output $P$ is kept in FP16 rather than being quantized. The value matrix $V$ also stays in FP16. Their matrix multiplication uses the mma.f16.f16.f16 tensor core instruction, which accumulates in FP16 rather than the default FP32. This is 2× faster than FP32 accumulation on consumer GPUs (RTX4090, RTX3090) while incurring no measurable accuracy loss compared to FP32 accumulation.

4. Four selectable kernel variants — Based on two binary choices (per-token vs. per-block quantization for $Q$ and $K$; INT8 vs. FP16 for $P$ and $V$), the system implements four kernels: SAGEAttn-T (per-token, FP16 PV), SAGEAttn-B (per-block, FP16 PV), SAGEAttn-vT (per-token, INT8 PV), and SAGEAttn-vB (per-block, INT8 PV). SAGEAttn-B is the default — accurate enough for all models and 2× faster than FP16 attention.

5. An adaptive kernel selector — For each layer in a model, the system tests SAGEAttn-vB on representative inputs and measures its cosine similarity against full-precision attention output. If the cosine similarity exceeds 99.8% (the worst-case similarity of SAGEAttn-B), that layer uses the slightly faster SAGEAttn-vB kernel. Otherwise, it falls back to SAGEAttn-B. This recovers an additional ~12% speedup over using SAGEAttn-B everywhere, without any accuracy loss.

Information flows through the system as follows: the query $Q$, key $K$, and value $V$ matrices enter in FP16 → $K$ is smoothed by subtracting its per-channel mean → $Q$ is multiplied by $1/\sqrt{d}$ on-chip → $Q$ and smoothed $K$ are quantized to INT8 with per-block scales → INT8 matrix multiplication produces $\hat{Q}\hat{K}^T$ → the result is dequantized to FP32 → online softmax computes $P$ in FP32, then casts to FP16 → $P$ (FP16) and $V$ (FP16) are multiplied with FP16 accumulation → the output block $O_i$ is written to global memory in FP16.

3.3 Roadmap for the Deep Dive

• First, the quantized attention formulation (Section 4.1) — how quantization and dequantization operators are inserted into the FlashAttention tiling loop, establishing the mathematical framework that all four kernel variants share.
• Second, the smoothing transformation for $K$ (Section 4.2) — why $K$ has channel-wise outliers, why per-channel quantization is impossible, how subtracting the mean fixes the problem without changing attention scores, and the empirical evidence that this single change recovers most of the lost accuracy.
• Third, the quantization strategy for $Q$, $K$, $P$, and $V$ (Section 4.3) — the choice of INT8 over FP8, the choice of per-block granularity for $Q$ and $K$, the special handling of $P$ with a static scale, and per-channel quantization for $V$.
• Fourth, the FP16 accumulator for $PV$ (Section 4.4) — why INT8 quantization of $P$ and $V$ has catastrophic worst-case failure modes, why keeping them in FP16 with an FP16 accumulator is both more accurate and faster, and the evidence that FP16 accumulation matches FP32 accumulation with zero measurable degradation.
• Fifth, the adaptive quantization strategy (Section 4.5) — how the four kernel variants are constructed, how the per-layer selection works, and the empirical speed-accuracy tradeoff that motivates mixing kernels.
• Sixth, the fusion tricks and hardware-level performance analysis (Section 4.6) — how quantization is fused with ROPE to hide I/O overhead, why INT8 improves cache usage and register pressure, and where the 2× speedup actually comes from in terms of hardware utilization.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily a systems and implementation paper whose core idea is that attention can be accurately quantized to INT8 by treating the two matrix multiplications asymmetrically — smoothing outliers from $K$ before quantizing $QK^T$ to INT8, while keeping $PV$ in FP16 with a faster accumulator.

---

3.4.1 Formalizing Quantized Attention Within FlashAttention's Tiling Loop

The paper begins by specifying exactly how quantization operators are inserted into the FlashAttention-2 computation. This is not a new attention algorithm — it is the standard FlashAttention-2 tiling strategy with quantization and dequantization operators wrapped around the two matrix multiplications.

Recall from Section 3.1 of the paper that FlashAttention-2 tiles $Q$, $K$, and $V$ into blocks $\{Q_i\}$, $\{K_j\}$, $\{V_j\}$ along the token dimension with block sizes $b_q$ for queries and $b_{kv}$ for keys/values. The computation proceeds iteratively: for each query block $Q_i$, the algorithm loops over all key-value blocks $(K_j, V_j)$, accumulating partial results for the output block $O_i$ using the online softmax algorithm (Equations 1 and 2 from the paper). The key insight of FlashAttention is that the $N \times N$ attention score matrix $S$ and softmax output $P$ are never materialized in full — they are computed block-by-block and immediately consumed.

SageAttention modifies this loop by inserting quantizers before each Matmul and dequantizers after, as formalized in Equations 4 and 5 (Section 4.1):

Quantization step (Equation 4):

$(\delta_Q, \hat{Q}) = \psi_Q(Q / \sqrt{d}), \quad (\delta_K, \hat{K}) = \phi_K(K), \quad (\delta_P, \hat{P}) = \psi_P(\tilde{P}), \quad (\delta_V, \hat{V}) = \psi_V(V)$

where $\psi_Q$, $\psi_K$, $\psi_P$, $\psi_V$ are quantization functions (e.g., per-token INT8 quantizer) that each return a scale factor $\delta$ and a low-precision tensor $\hat{\cdot}$, and $\phi_K = \psi_K \circ \gamma$ is a composition of the smoothing transform $\gamma$ followed by quantization $\psi_K$. The operator $\psi_Q$ also subsumes the $1/\sqrt{d}$ scaling of $Q$ — this scaling is fused into the quantization to avoid a separate kernel launch.

Attention computation with dequantization (Equation 5):

$S = \psi^{-1}_{\delta_Q \delta_K}(\hat{Q}\hat{K}^\top), \quad (m', P) = \tilde{\sigma}(m, S), \quad O = \text{diag}(\exp(m' - m)) O + \psi^{-1}_{\delta_P \delta_V}(\hat{P}\hat{V})$

where $\psi^{-1}_{\delta_Q\delta_K}$ is a dequantizer that multiplies the INT32 result of $\hat{Q}\hat{K}^\top$ by the product of scale factors $\delta_Q \times \delta_K$, converting back to FP32; $\tilde{\sigma}$ is the online softmax operator from FlashAttention (which updates the running maximum $m$ and computes the exponentiated, shifted scores $\tilde{P}$); and $\psi^{-1}_{\delta_P\delta_V}$ is the dequantizer for the second Matmul.

What these equations compute: The first line (quantization) converts each matrix from FP16 to a low-precision format (INT8 or FP8) with an associated per-block or per-token scale factor. The quantization function $\psi$ computes the scale as $\delta = \max(|A|) / 127$ for INT8 (the maximum representable value in signed INT8), then rounds each element to the nearest integer: $\hat{A} = \lceil A / \delta \rfloor$. The result is an integer tensor $\hat{A}$ whose values lie in $[-127, 127]$ and a scalar or vector $\delta$ that captures the magnitude. Dequantization simply multiplies back: $\psi^{-1}_\delta(\hat{A}) = \delta \hat{A}$.

The second line (attention) performs the two matrix multiplications in low precision, then dequantizes the results before they enter the full-precision online softmax and output accumulation. Critically, the online softmax itself — the exponential, row-wise max, and running sum — operates in FP32, not in quantized arithmetic. Only the matrix multiplications are quantized. The reasoning is that the softmax involves transcendental functions and reductions that would be difficult to implement efficiently in integer arithmetic, and its cost is $O(N)$ compared to the $O(N^2)$ matrix multiplications, so the speedup from quantizing it would be negligible.

Why this form: The paper chooses to quantize the matrix multiplications individually (with separate scale factors for $Q$, $K$, $P$, and $V$) rather than attempting to fuse quantization across operations. This is necessary because each matrix has a different dynamic range and outlier structure — $K$ has channel-wise outliers (requiring smoothing before quantization), $P$ has a known maximum value of 1 per row (enabling a static scale), and $V$ may have channel-wise variation. Individual quantization gives the flexibility to handle each matrix appropriately. The alternative — using a single quantization scheme for all four matrices — would force compromises (e.g., quantizing $P$ with the same granularity as $Q$ despite their very different distributions) that degrade accuracy.

The dequantization multiplies the scale factors $\delta_Q$ and $\delta_K$ together before applying them to the product $\hat{Q}\hat{K}^\top$. This is mathematically equivalent to dequantizing $Q$ and $K$ separately then multiplying — but it is more efficient because the scale multiplication happens on the small scale factors (scalars or small vectors) rather than on the large $N \times N$ product matrix. This is a standard trick in quantized matrix multiplication: $(s_Q \hat{Q})(s_K \hat{K}) = (s_Q s_K)(\hat{Q}\hat{K}^\top)$, where the scalar product $s_Q s_K$ is applied once to the result.

---

3.4.2 Smoothing the K Matrix: Removing Channel-Wise Outliers Without Changing Attention

The most critical accuracy-preserving technique in SageAttention is the smoothing transformation applied to $K$ before quantization. The paper argues that this single operation is what makes INT8 attention viable, and the empirical evidence in Table 1 strongly supports this claim.

What the channel-wise outliers look like. Figure 4 in the paper visualizes the data distributions of $Q$, $K$, and $V$ from two models (CogvideoX and Unidiffuser). The color maps reveal a striking pattern: $K$ shows distinct vertical bands where every token in a particular channel has a large magnitude. For CogvideoX, one channel has values spanning approximately -11 to 14 (a range of ~25), while another channel has values spanning -210 to 208 (a range of ~420). For Unidiffuser, one channel ranges from -151 to 147 (range ~300). Crucially, the variation within each channel across tokens is small — the large range comes from the channel having a large mean value (bias) shared by all tokens.

In plain terms: each key vector $K[t, :]$ for token $t$ is approximately $\mathbf{b} + \mathbf{s}_t$, where $\mathbf{b}$ is a $d$-dimensional bias vector that is the same for all tokens (some channels have large positive or negative biases) and $\mathbf{s}_t$ is a small token-specific signal. The bias dominates the magnitude, so a per-tensor or per-token quantizer will set its scale based on the bias and quantize the small signal to zero.

Why per-channel quantization cannot fix this. The paper explains that quantization scale factors can only be placed on the outer axis of a matrix multiplication. In the Matmul $QK^\top$, the outer axis of $K$ is the token dimension (since $K$ has shape $N \times d$ and we compute $QK^\top$ where $Q$ is $N \times d$). If you quantize $K$ per-channel (assigning a separate scale factor to each of the $d$ channels), those scale factors live on the inner dimension of the Matmul — the dimension being summed over. When you try to dequantize the output, you would need to apply different scale factors to different entries of the product matrix, but those entries are already summed across channels and the channel identity is lost. Therefore, channel-wise scales for $K$ cannot be propagated through the matrix multiplication. This constraint is inherited from linear layer quantization literature (Xiao et al., 2023a).

Why SmoothQuant-style scaling doesn't work. SmoothQuant (Xiao et al., 2023a) handles activation outliers in linear layers by multiplying activations by a per-channel smoothing factor and dividing the static weight matrix by the same factor. The mathematical identity $XW = (X \cdot \text{diag}(s))(\text{diag}(s)^{-1}W)$ lets the quantization difficulty be migrated from the activations (which have outliers) to the weights (which are smooth). In attention, however, both $Q$ and $K$ are dynamic — they change for every input. There is no static weight matrix to absorb the scale. If you multiply $Q$ by a smoothing factor and divide $K$ by the same factor, the product $QK^\top$ is unchanged, but both matrices are now smoothed — if the outliers are in the same channels for $Q$ and $K$. However, the paper notes that $Q$ is also affected by outliers (visible in Figure 4, though less severe than $K$), and the outliers may not align between $Q$ and $K$, making this approach fragile.

The smoothing transform $\gamma$: a mathematically-equivalent bias removal. The paper's key insight about $K$'s outliers is that they come from a large bias, not from large variation across tokens. This is a crucial distinction. If the outliers were from cross-token variation (some tokens having genuinely much larger keys than others), subtracting the mean would destroy information — those genuinely large keys would be flattened. But if the outliers are from a channel-wise bias shared by all tokens, subtracting the mean removes the bias while perfectly preserving the token-wise differences that determine attention patterns.

The transform is defined in Equation 6:

$\gamma(K) = K - \text{mean}(K)$

where $\text{mean}(K) = \frac{1}{N} \sum_{t=1}^N K[t, :]$ is a $1 \times d$ vector containing the per-channel mean of $K$ across all tokens.

Why this does not change attention scores. The critical mathematical property is that subtracting a constant vector from every key does not affect the softmax attention weights. For any query $q$:

$\sigma(q(K - \text{mean}(K))^\top) = \sigma(qK^\top - q \cdot \text{mean}(K)^\top)$

The term $q \cdot \text{mean}(K)^\top$ is the same scalar for every key because $\text{mean}(K)$ is the same for all tokens. Subtracting a constant from every logit in a softmax does not change the output:

$\text{softmax}(x_i - c) = \frac{\exp(x_i - c)}{\sum_j \exp(x_j - c)} = \frac{\exp(x_i)}{\sum_j \exp(x_j)} \cdot \frac{\exp(-c)}{\exp(-c)} = \frac{\exp(x_i)}{\sum_j \exp(x_j)} = \text{softmax}(x_i)$

Therefore, $\sigma(qK^\top) = \sigma(q(K - \text{mean}(K))^\top)$ exactly. The transformation changes the absolute values entering quantization but produces identical attention weights — it is a lossless pre-processing step with respect to the mathematical function being computed.

The full quantization pipeline for $K$. The transformation from full-precision $K$ to quantized $\hat{K}$ is the composition $\phi_K(K) = \psi_K \circ \gamma(K)$, meaning: first subtract the mean across tokens, then quantize the result. The quantizer $\psi_K$ can use any granularity (per-token, per-block, or per-tensor), but per-block quantization with smoothing achieves the best accuracy as shown in Table 1.

Empirical evidence that smoothing is essential. Table 1 presents end-to-end metrics across five models (Llama2, CogvideoX, Unidiffuser, UltraPixel, TIMM) comparing quantization methods with and without smoothing. The results are stark:

• On CogvideoX, per-block INT8 quantization without smoothing achieves an Fscore of 2.014 (vs. 3.768 for FP16), while with smoothing it achieves 3.718 — recovering nearly the entire gap.
• On Unidiffuser, per-block quantization without smoothing yields FID of 229.08 (vs. 163.33 for FP16 — a catastrophic degradation indicating nearly-random image generation), while with smoothing it achieves 166.93 — nearly indistinguishable from FP16.
• On UltraPixel, the pattern repeats: FID jumps from 179.78 (FP16) to 195.67 (without smoothing) and recovers to 179.98 (with smoothing).
• Llama2 is the exception — its perplexity barely changes with or without smoothing (5.823 FP16, 5.825 without, 5.824 with), because Llama2-7B's attention distributions are relatively uniform to begin with, as noted in Section A.6 of the Appendix.

The smoothing overhead is negligible: Table 10 shows that on CogvideoX, smoothing K reduces TOPS from 327.57 to 327.52, a difference of less than 0.02%. On UltraPixel, the difference is from 325.18 to 324.56. The mean computation is a simple reduction across the token dimension, which is well-parallelized on GPU and amortized over the much more expensive matrix multiplications.

A subtle design choice: mean subtraction vs. other bias-removal schemes. The paper could have used more sophisticated normalization — for example, subtracting the per-channel median or using a learned bias. The mean is chosen because it (1) zero-centers each channel, maximally compressing the dynamic range for subsequent quantization, (2) is cheap to compute (a single reduction kernel), and (3) has the clean mathematical guarantee of not changing attention scores. A median would be more robust to extreme outliers but more expensive to compute accurately on GPU; a learned bias would require per-model calibration and would break the plug-and-play property.

---

3.4.3 Quantization Strategy for Q, K, P, and V: Data Types, Granularity, and the Static Scale Trick

With the smoothing transform in place for $K$, the paper specifies exactly how each of the four matrices in attention is quantized. The choices are not symmetric — each matrix gets a tailored treatment based on its role in the computation and its observed distribution.

Choice of data type for $Q$ and $K$: INT8, not FP8. The paper provides two reasons for choosing INT8 over FP8 (E4M3 or E5M2 formats) for quantizing $Q$ and $K$:

Accuracy reason (Table 2): Table 2 reports the average accuracy of quantized attention output across all layers of Llama2-7B and Unidiffuser, measured by cosine similarity, relative L1, and RMSE against FP16 attention. For the $(Q, K) = \text{INT8}, (\tilde{P}, V) = \text{E4M3}$ combination, average cosine similarity is 99.94%, relative L1 is 0.0345, and RMSE is $3.53 \times 10^{-3}$. For $(Q, K) = \text{E4M3}, (\tilde{P}, V) = \text{E4M3}$, cosine similarity drops to 99.81% and error roughly doubles. This pattern holds across the table: INT8 for $Q$ and $K$ consistently achieves the highest accuracy among 8-bit options. The reason is that INT8 has uniform quantization levels (equal spacing between representable values), which better preserves the additive structure of the dot product $QK^\top$ compared to FP8 formats that allocate more precision near zero (via the exponent) but coarser precision for larger magnitudes.

Table 17 in Appendix B.1 reinforces this with a direct comparison of $Q \cdot K$ (not the full attention, just the Matmul) under different data types for a specific layer of Unidiffuser. INT8 achieves 99.54% cosine similarity and 0.084 relative L1. E4M3 drops to 92.83% cosine similarity and 0.342 relative L1 — more than 4× the error. E5M2 is even worse at 77.95% cosine similarity. The difference is not subtle: INT8 is substantially more accurate for this operation.

Speed reason: On consumer GPUs (RTX4090, RTX3090), INT8 tensor core instructions (mma.u8.u8.s32) run at twice the throughput of FP8 tensor core instructions. This is a hardware fact — Nvidia's consumer GPUs have dedicated INT8 tensor cores but process FP8 through the FP16 datapath with conversion overhead. Therefore, choosing INT8 is simultaneously more accurate and faster on the target hardware, making it a dominant choice over FP8.

Quantization granularity for $Q$ and $K$: per-token or per-block. The paper allows two granularity options for $\psi_Q$ and $\psi_K$: per-token (one scale factor per token, i.e., per row of the matrix) or per-block (one scale factor per block of consecutive tokens). Per-channel quantization is impossible for the reasons discussed above. Per-tensor quantization (one scale factor for the entire matrix) is the coarsest and least accurate, as shown in Table 1: on Unidiffuser, per-tensor with smoothing achieves FID 167.65, while per-block with smoothing achieves 166.93 and is closer to the FP16 baseline of 163.33.

Per-block quantization is more accurate than per-tensor because it adapts to variation in magnitude across different parts of the sequence. If the first 1000 tokens have small key magnitudes and the next 1000 have large magnitudes, a single per-tensor scale would be dominated by the large-magnitude tokens and crush the small ones. Per-block quantization assigns separate scales to each block, preserving the relative precision within each block.

Per-token quantization is more accurate than per-block because it tailors the scale to each individual token. However, it has a subtle advantage that the paper exploits: per-token scales for $Q$ and $K$ can be combined into a single outer product of scale vectors before dequantization, because $(\text{diag}(s_Q) \hat{Q})(\hat{K}^\top \text{diag}(s_K)) = \text{diag}(s_Q)(\hat{Q}\hat{K}^\top)\text{diag}(s_K) = (s_Q \otimes s_K) \odot (\hat{Q}\hat{K}^\top)$, where $\otimes$ is the outer product. This means the dequantization does not require an element-wise multiplication of the full $N \times N$ matrix — it can be done as a rank-1 scaling, which is faster.

Quantization of P: exploiting the known maximum. The softmax output $P$ (denoted $\tilde{P}$ in the paper's online softmax formulation) has a special property: for each row $i$, the maximum value is exactly 1 (corresponding to the token with the highest attention score), and all values are non-negative and sum to 1. This means the per-token scale factor for any row is:

$\delta_P[i] = \max(P[i, :]) / 127 = 1 / 127$

which is a constant independent of the input! Therefore, the paper can use a static scale $s = 1/127$ for per-block quantization of $P$ without computing the maximum dynamically. This static scale gives per-block quantization of $P$ accuracy equivalent to per-token quantization — because no matter how you partition the tokens into blocks, the maximum in any block is at most 1, and many blocks will contain the maximum (since the attention distribution is typically peaked, the maximum value of 1 appears in most blocks that contain highly-attended tokens).

The quantization of $P$ then simplifies to: $\hat{P} = \lceil P \times 127 \rfloor$, which is just a scaling and rounding to the integer range $[0, 127]$, with the scale factor $1/127$ known at compile time. Dequantization is simply $\hat{P} / 127$.

Quantization of V: per-channel for channel-wise variation. Unlike $Q$ and $K$, the value matrix $V$ can be quantized per-channel because the channel dimension is the outer dimension of the second Matmul $PV$: $V$ has shape $N \times d$, and the Matmul is $P_{b_q \times N} \times V_{N \times d}$, producing an output of shape $b_q \times d$. The $d$ dimension is the outer dimension of the product, meaning per-channel scale factors for $V$ can be applied column-wise to the output. This allows per-channel quantization of $V$, which handles any channel-wise outliers in $V$ (visible in Figure 4, where some channels of $V$ have larger magnitudes than others).

The data types for P and V: why INT8? Table 2 shows that using INT8 for $P$ and $V$ achieves 99.70% average cosine similarity in the $(Q, K) = \text{INT8}, (\tilde{P}, V) = \text{INT8}$ configuration — close to the 99.94% achieved with E4M3 for $P$ and $V$. The speed advantage of INT8 over FP8 on consumer GPUs (2× faster Matmul) justifies this choice for the average case.

However, the worst-case is catastrophic. This is where the story takes a critical turn and motivates the FP16 accumulator design. Table 3 reports the worst-case accuracy across all layers — the single layer with the lowest cosine similarity. For $(Q, K) = \text{INT8}, (\tilde{P}, V) = \text{INT8}$, the worst-case cosine similarity is only 56.40%, with relative L1 of 0.792 and RMSE of 0.541. This means there exists at least one attention layer (in either Llama2 or Unidiffuser) where INT8 quantization of $P$ and $V$ produces output that is essentially uncorrelated with the correct output. Even the best 8-bit alternative (E4M3 for both) only achieves 76.36% in the worst case.

This worst-case failure is what causes the catastrophic model-level degradation seen in Figure 3 and Table 1: Unidiffuser with INT8 attention (per-token, no smoothing) generates completely blurry images, and Llama2 with INT8 attention achieves only 25.5% on MMLU. The failure is not average — most layers might quantize fine — but a single layer producing garbage output propagates through the rest of the network and destroys the final result.

---

3.4.4 The FP16 Accumulator for PV: How Not Quantizing Is the Best Quantization

Rather than trying to fix the worst-case failure of INT8 $PV$ through more sophisticated quantization (which would likely require per-layer calibration, mixed precision, or fallback mechanisms), the paper makes a counterintuitive but elegant choice: don't quantize $P$ and $V$ at all, but still get a speedup by changing the accumulation precision.

The hardware opportunity. On Nvidia consumer GPUs (RTX4090, RTX3090), the FP16 matrix multiplication instruction has two variants: one that accumulates results in FP32 (mma.f16.f16.f32) and one that accumulates in FP16 (mma.f16.f16.f16). The FP16 accumulator variant is 2× faster than FP32 accumulation because it uses half the register file for the accumulator (each FP16 value occupies 16 bits vs. 32 bits for FP32), allowing more warps to be scheduled concurrently and reducing register pressure.

However, FP16 accumulation is numerically riskier: FP16 has only 10 bits of mantissa (vs. 23 bits for FP32), meaning that after accumulating many terms in a dot product, rounding errors can accumulate. For typical neural network activations, FP16 accumulation is generally considered unsafe without careful analysis.

The empirical finding: FP16 accumulation is perfectly safe for $PV$ in attention. Tables 4 and 5 present the critical evidence. Table 4 shows the average accuracy of the $PV$ Matmul using FP16 vs. FP32 accumulators across all layers of Llama2 and Unidiffuser: both achieve 99.98% cosine similarity, relative L1 of 0.0156, and RMSE of $2.94 \times 10^{-3}$. There is no measurable difference in the average case.

Table 5 shows the worst-case: both achieve 99.84% cosine similarity, relative L1 of 0.0511, and RMSE of $4.229 \times 10^{-3}$. The worst-case FP16 accumulation is identical to the worst-case FP32 accumulation to four significant figures.

Why is FP16 accumulation safe here specifically? The paper does not provide a theoretical analysis, but the empirical result makes physical sense when considering the properties of the $PV$ Matmul. The softmax output $P$ has values in $[0, 1]$ that sum to 1 — it's a probability distribution. The value matrix $V$ typically has entries of moderate magnitude (the paper doesn't report exact ranges, but Figure 4 shows values in the range -5 to 5 for CogvideoX and -3 to 3 for Unidiffuser). The dot product between a row of $P$ (which sums to 1) and a column of $V$ is a weighted average of the values in that column, with the weights given by the attention probabilities. This weighted average is numerically well-conditioned — it's a convex combination, so the result is bounded by the min and max of $V$'s column, and the rounding errors from accumulating many small terms in FP16 are small relative to the magnitude of the result. This contrasts with the $QK^\top$ Matmul, where both $Q$ and $K$ can have large entries and the dot products can have high dynamic range, making FP16 accumulation riskier.

The design implication. By keeping $P$ and $V$ in FP16 and using the FP16 accumulator, SageAttention eliminates quantization error entirely for the second half of attention while still achieving a 2× speedup over the FP32-accumulator default. This is strictly better than any 8-bit quantization scheme for $PV$: it is more accurate (no quantization error at all vs. worst-case 56% cosine similarity) and equally fast (FP16-with-FP16-accumulator has the same throughput as INT8 Matmul on consumer GPUs — both are 2× faster than FP16-with-FP32-accumulator). The only cost is that $P$ and $V$ remain in FP16 in memory, consuming twice the bandwidth of INT8 during loading. But since $P$ is computed on-chip (from the dequantized INT8 $QK^\top$ result) and $V$ is loaded from global memory, the memory bandwidth overhead is acceptable — $V$ has shape $N \times d$ with $d$ small (64–128), so the total data movement for $V$ is modest compared to the $N \times N$ matrices.

A note on dequantization integration. In the kernel implementations that use FP16 for $PV$ (SAGEAttn-T and SAGEAttn-B), Equation 5's dequantization term $\psi^{-1}_{\delta_P\delta_V}(\hat{P}\hat{V})$ is replaced by simply Matmul(P.to(FP16), V, Accum type = FP16). The P.to(FP16) cast is explicit in Algorithm 1 because $P$ is computed from the online softmax in FP32 and needs to be converted to FP16 before entering the tensor core instruction. This cast is essentially free — it's just a data type conversion in the register file, not a memory operation.

---

3.4.5 Four Kernel Variants and Adaptive Per-Layer Selection

Based on the two binary choices — (1) per-token vs. per-block quantization for $Q$ and $K$, (2) INT8 vs. FP16 for $P$ and $V$ — the paper implements four attention kernels, summarized in Table 6:

• SAGEAttn-T: per-token INT8 quantization of $Q$ and $K$ (with smoothing of $K$), FP16 $P$ and $V$ with FP16 accumulator. This is the most accurate variant because per-token quantization gives the finest granularity for $Q$ and $K$, and FP16 $PV$ eliminates quantization error entirely. It achieves 100% cosine similarity and $6.8 \times 10^{-4}$ RMSE on synthetic data (Table 9).

• SAGEAttn-B: per-block INT8 quantization of $Q$ and $K$ (with smoothing of $K$), FP16 $P$ and $V$ with FP16 accumulator. This is the default and recommended variant. It achieves 100% cosine similarity and $7.3 \times 10^{-4}$ RMSE (Table 9) — essentially identical to SAGEAttn-T in accuracy but slightly faster because per-block quantization has lower overhead (fewer scale factors to compute and apply). Figure 6 shows that at headdim=64, causal=False, sequence length 32K, SAGEAttn-B achieves 340 TOPS vs. 307 TOPS for SAGEAttn-T on RTX4090.

• SAGEAttn-vT: per-token INT8 quantization of $Q$ and $K$, INT8 quantization of $P$ and $V$. This variant quantizes everything to INT8, maximizing memory bandwidth savings at the cost of potential worst-case accuracy degradation. It achieves 99.9% cosine similarity and 0.065 RMSE on synthetic data (Table 9).

• SAGEAttn-vB: per-block INT8 quantization of $Q$ and $K$, INT8 quantization of $P$ and $V$. Similar accuracy to SAGEAttn-vT (98.9% cosine similarity, 0.067 RMSE in Table 9) but slightly faster due to per-block granularity.

Why these specific four? The paper could have implemented more variants (e.g., per-tensor quantization, or mixing per-token $Q$ with per-block $K$), but these four cover the Pareto frontier of the speed-accuracy tradeoff. SAGEAttn-B is the safe, always-accurate choice. SAGEAttn-vB is slightly faster (about 4%, as mentioned in Section 4.5) because INT8 $PV$ uses less memory bandwidth and cache for loading $V$ and $P$. The question is: can we safely use SAGEAttn-vB on some layers to get that extra speed?

The adaptive selection mechanism. For each layer in a model, SageAttention tests SAGEAttn-vB on representative inputs and measures the cosine similarity between its output and the output of FP16 attention. If the cosine similarity is greater than 99.8% — which is the worst-case cosine similarity of SAGEAttn-B across all layers tested — then that layer uses SAGEAttn-vB. Otherwise, it falls back to SAGEAttn-B.

The choice of 99.8% as the threshold is calibrated: it guarantees that any layer using SAGEAttn-vB has accuracy at least as good as the worst layer using SAGEAttn-B. Since SAGEAttn-B already produces no measurable end-to-end metric degradation across all models tested, substituting some layers with SAGEAttn-vB (which meets or exceeds SAGEAttn-B's worst-case) cannot degrade overall model accuracy.

Empirical benefit of adaptive selection. Table 11 quantifies the impact. On CogvideoX, using only SAGEAttn-T achieves 292.17 TOPS with CLIPSIM of 0.1827 (slightly degraded from the FP16 baseline of 0.1837). Using the adaptive strategy (mixing SAGEAttn-B and SAGEAttn-vB where appropriate) achieves 327.57 TOPS with CLIPSIM of 0.1835 — a 12% speedup with essentially no accuracy loss. On Llama2, SAGEAttn-T achieves 208.59 TOPS, while the adaptive strategy achieves 231.74 TOPS — an 11% speedup — with identical MMLU accuracy of 0.46.

The 11.7% figure cited in Section 5.4 ("the adaptive strategy increases the speed of attention by 11.7% without any loss in metrics") is computed as the speedup of adaptive selection over using SAGEAttn-T everywhere. In practice, using SAGEAttn-B everywhere already provides most of this benefit, and the adaptive selection recovers a residual few percent from layers where SAGEAttn-vB is provably safe.

---

3.4.6 Fusion Tricks: Hiding Quantization Overhead by Fusing with ROPE

Quantization introduces overhead: computing the scale factor (finding the maximum absolute value), scaling and rounding each element, and storing both the quantized tensor and the scale factor. If these operations are implemented as separate GPU kernels, they add latency and memory traffic that eat into the speedup from faster Matmul.

Fusing with Rotary Position Embedding (ROPE). The paper fuses the quantization of $Q$ and $K$ with the preceding ROPE operation. ROPE (Rotary Position Embedding, Su et al., 2021) applies a rotation to the query and key vectors based on their position in the sequence. The standard implementation flow is: (1) apply ROPE to $Q$ and $K$ in FP16, (2) write the rotated $Q$ and $K$ to global memory, (3) the attention kernel reads $Q$ and $K$ from global memory, (4) quantize them, (5) compute Matmul.

SageAttention fuses steps (2), (3), and (4): after the ROPE rotation is computed and the results are in shared memory (on-chip SRAM), the quantization is applied immediately — compute the scale, scale and round, convert to INT8 — and then the INT8 tensors are written to global memory instead of the FP16 tensors. This avoids writing and immediately re-reading the FP16 $Q$ and $K$ from global memory, which is the dominant I/O cost. The scale factors (small vectors of FP32 values) are also written to global memory, but their size is negligible compared to the tensors themselves.

Fusing the $1/\sqrt{d}$ scaling. The attention formula requires scaling $QK^\top$ by $1/\sqrt{d}$, where $d$ is the head dimension. Rather than applying this scaling in the attention kernel after dequantization (which would require an extra element-wise multiplication on the $N \times N$ product), SageAttention fuses it into the quantization of $Q$: before quantizing $Q$, it is multiplied by $1/\sqrt{d}$ on-chip. The quantizer then computes $\delta_Q = \max(|Q / \sqrt{d}|) / 127$ and $\hat{Q} = \lceil (Q / \sqrt{d}) / \delta_Q \rfloor$. The dequantization in the attention kernel becomes $\delta_Q \delta_K (\hat{Q}\hat{K}^\top)$, which already includes the $1/\sqrt{d}$ factor implicitly because $Q$ was scaled before quantization.

---

3.4.7 Hardware-Level Performance Analysis: Where the 2× Speedup Actually Comes From

The paper's performance analysis in Section 4.6 breaks down the acceleration into four sources:

(1) Matmul acceleration (the primary source). On RTX4090 and RTX3090, the INT8 tensor core instruction mma.u8.u8.s32 has 2× the throughput of the FP16 tensor core instruction mma.f16.f16.f32 when accumulating in FP32, and 4× the throughput when comparing against FP32 accumulation. The theoretical peak INT8 throughput on RTX4090 is approximately 660 TOPS (the paper cites that SageAttention's 340 TOPS at headdim=64 and 128 represents "52% of the theoretical INT8 throughput"). The FP16-with-FP16-accumulator instruction mma.f16.f16.f16 is 2× faster than mma.f16.f16.f32 because halving the accumulator bit width reduces register pressure, allowing higher warp occupancy.

For the $QK^\top$ Matmul, SageAttention uses INT8 with 4× throughput over FP16-with-FP32-accumulator. For the $PV$ Matmul, it uses FP16-with-FP16-accumulator with 2× throughput over FP16-with-FP32-accumulator. The combined effect is approximately a 2× overall speedup, consistent with the empirical results in Figures 6–9.

(2) Quantization overhead (a manageable cost). Quantization adds computation: finding the maximum absolute value (a reduction), scaling and rounding (element-wise), and type conversion. Dequantization adds a multiplication of the product matrix by the scale factors. These operations are $O(N \times d)$ or $O(N \times N)$ for the dequantization of the product, but they run at the full memory bandwidth or compute throughput of the GPU's CUDA cores, not the tensor cores. The paper argues that fusing quantization with ROPE (for $Q$ and $K$) and with the softmax (for $P$) eliminates the I/O overhead, and the computational overhead is small compared to the $O(N^2 \times d)$ Matmul.

(3) Cache and register benefits. Using INT8 data for $Q$ and $K$ halves the shared memory requirement for storing these tiles on-chip compared to FP16. Shared memory is a scarce resource on GPUs (typically 48–164 KB per SM, depending on configuration), and FlashAttention's tiling strategy requires storing blocks of $Q$, $K$, $V$, and the output accumulator simultaneously. Halving the size of the $Q$ and $K$ tiles allows larger block sizes or higher occupancy (more thread blocks per SM), both of which improve throughput. Similarly, using FP16 accumulators for $PV$ halves the register file usage for the accumulator, reducing register pressure and enabling more warps to be scheduled concurrently.

(4) DRAM access reduction. INT8 tensors are half the size of FP16 tensors in global memory (HBM). For the $QK^\top$ Matmul, the $Q$ and $K$ tiles are loaded from global memory into shared memory; using INT8 halves this traffic. The paper acknowledges that the scale factors add additional traffic (FP32 values, one per block or per token), but these are "negligible compared to the tensors" — for a block size of 128 tokens with head dimension 64, the INT8 tensor is $128 \times 64 = 8192$ bytes, while the per-block scale factor is 4 bytes. The scale overhead is less than 0.05%.

Throughput numbers in context. The paper reports that SageAttention achieves 340 TOPS at headdim=64 and headdim=128 on RTX4090 (Figures 6 and 7), representing 52% of the theoretical INT8 peak. For comparison, FlashAttention2 achieves approximately 165 TOPS in the same configuration. SageAttention's 340 TOPS on RTX4090 is described as "close to the 490 TOPS throughput of FlashAttention3" on Hopper GPUs — an interesting comparison because the RTX4090 is a consumer GPU costing roughly 10× less than an H100, yet SageAttention achieves ~70% of FlashAttention3's throughput through better hardware utilization rather than more powerful hardware.

4. Key Insights and Innovations

Innovation 1: Redefining Attention Quantization as an Asymmetric, Matmul-Specific Problem Rather Than a Uniform Tensor Compression Task

Prior to SageAttention, the default mental model for quantizing attention — inherited from linear layer quantization — was to treat all matrices symmetrically: choose a data type (INT8 or FP8), choose a granularity (per-tensor or per-token), and apply the same quantizer to $Q$, $K$, and $V$. This mental model fails catastrophically for attention, as the paper demonstrates (Figure 3, Table 1), but the failure mode itself is instructive: the problem is not simply that attention is "harder to quantize," but that the two matrix multiplications in attention present fundamentally different numerical challenges that demand fundamentally different solutions.

The paper's core conceptual move is to decompose attention quantization into two independent sub-problems — the $QK^\top$ Matmul and the $PV$ Matmul — and to recognize that they have different sensitivity to quantization error, different outlier structures, and different opportunities for hardware acceleration. This decomposition sounds obvious in retrospect, but the field had been treating "quantized attention" as a monolithic problem, searching for a single quantization scheme that would work for both Matmuls. SageAttention shows that the optimal solution is not on that Pareto frontier at all: treat the first Matmul with aggressive INT8 quantization after a mathematically-equivalent pre-processing step, and treat the second Matmul by not quantizing at all but exploiting a faster accumulation mode.

The reframing is significant beyond this paper because it provides a template for how to approach other complex operators that contain multiple matrix multiplications with different numerical properties. The insight — that the right unit of analysis is not the operator but the individual Matmul — is transferable to other attention variants (cross-attention, multi-query attention) and potentially to other building blocks (e.g., the gating mechanisms in mixture-of-experts, the projections in state-space models). It shifts the question from "how do I quantize attention?" to "for each Matmul in my operator, what is the cheapest way to compute it that preserves the downstream computation?"

The empirical evidence for why this asymmetry is necessary is the comparison between Tables 2 and 3: the average accuracy of INT8 $PV$ is acceptable (99.70% cosine similarity), but the worst-case is catastrophic (56.40%). This worst-case is not a minor degradation — it means a specific layer's attention output is essentially uncorrelated with the correct output. Zooming out to end-to-end metrics, this single worst-case layer cascades into complete model failure (Llama2 MMLU accuracy dropping from 46% to 25.5%, Unidiffuser generating unrecognizable images). The key conceptual move is recognizing that average-case safety is insufficient for plug-and-play deployment — the method must be safe in the worst case, and the worst case for $PV$ quantization is fundamentally unsolvable without per-layer calibration, which breaks the plug-and-play property.

Innovation 2: Diagnostic Insight That Channel-Wise K Outliers Are a Bias, Not a Variance Problem — And Can Be Removed Without Changing Attention Scores

Prior work on activation quantization (e.g., SmoothQuant, Xiao et al. 2023a) treated outliers as a problem of per-channel variation — some channels having larger activation magnitudes than others — and solved it by migrating the scale difficulty to static weights. In attention, there are no static weights to absorb scale, so this approach does not directly apply. A reasonable extension would be to treat the $K$ outliers as a per-channel variation problem and attempt to smooth them with learned or calibrated scaling factors.

SageAttention makes a subtler observation by looking at the actual distribution of $K$ values (Figure 4): the large channel-wise magnitudes come not from high variance across tokens, but from a large channel-wise mean (bias) shared by all tokens. If channel $c$ has values spanning -210 to 208, but token $t_1$ has $K[t_1, c] = -208$ and token $t_2$ has $K[t_2, c] = 205$, the range is ~415, but the token-to-token difference is only ~13. The quantization scale would be set based on the ~415 range (to avoid clipping), crushing the tiny ~13 differences that encode which tokens attend to which.

This is a diagnostic insight, not just an engineering fix. It tells us why $K$ is hard to quantize: the quantization dynamic range is dominated by a signal (the per-channel bias) that carries no information about attention patterns, while the informative signal (token-to-token variation) is small and gets quantized to zero. This is fundamentally different from the problem in linear layers, where outlier channels genuinely carry information — you cannot simply subtract the mean of a linear layer activation without changing the network's function.

The reason subtracting the mean works — and why it's not an obvious "just normalize your tensors" trick — is that attention is invariant to additive constants in the key dimension. The softmax operation $\sigma(qK^\top)$ is unchanged by subtracting any constant vector from all keys because the constant translates to a uniform shift in the logits, which cancels in the softmax. This mathematical property is specific to the attention mechanism; it does not hold for general matrix multiplications like those in linear layers or feedforward networks. The paper's insight is not that $K$ has outliers (this is visible in any visualization), but that the specific structure of those outliers makes them removable by a cheap, mathematically-exact transformation rather than requiring approximate smoothing or learned calibration.

The broader significance is that this diagnostic approach — looking at what kind of outlier a tensor has (bias-dominated vs. variance-dominated) rather than just whether it has outliers — could guide quantization strategies for other operators. A tensor where the mean dominates the dynamic range but is irrelevant to downstream computation can be mean-centered; a tensor where genuine token-to-token variation is large requires different treatment. Table 1 implicitly demonstrates the limits: smoothing $K$ recovers nearly all lost accuracy for CogvideoX and Unidiffuser (where $K$ has large channel biases), but Llama2 barely benefits because its $K$ distribution is already uniform — the diagnostic correctly identifies where smoothing is necessary and where it is unnecessary.

The negligible overhead (less than 0.2%, Table 10) makes this a pure win when it applies, establishing a practical principle: before designing complex quantization schemes, check whether the outlier can be mathematically eliminated without approximation. This principle likely extends beyond attention to any computation with invariance properties that allow lossless pre-processing.

Innovation 3: Exploiting Hardware-Accumulation Modes as a Quantization Alternative — The Counterintuitive Win of Not-Quantizing

The dominant assumption in quantization research is that lower precision = faster computation, and the goal is to push precision as low as possible while maintaining accuracy. SageAttention turns this assumption sideways for the $PV$ Matmul: instead of asking "what's the lowest precision that maintains accuracy?", it asks "what's the fastest way to compute this Matmul accurately?" The answer — FP16 with FP16 accumulation — is not lower precision at all (both operands remain in FP16), but it is 2× faster than the standard FP16 Matmul with FP32 accumulation on the target hardware because of reduced register pressure.

This is novel not as a hardware discovery (Nvidia documents the throughput difference between f16.f16.f32 and f16.f16.f16), but as a strategic insight about where the speed-accuracy Pareto frontier actually lies. The field's instinct, given a 2× throughput advantage for INT8 over FP16, is to quantize to INT8. SageAttention shows that for the specific numerical conditions of the $PV$ Matmul in attention — where $P$ is a probability distribution (bounded, sum-to-1) and $V$ has moderate magnitude — FP16 accumulation achieves zero measurable accuracy loss compared to FP32 accumulation (Tables 4 and 5, identical to four significant figures in worst-case cosine similarity), while INT8 quantization introduces catastrophic worst-case errors (Table 3). In this regime, "not quantizing" dominates "quantizing to INT8" on both speed and accuracy simultaneously.

The broader significance is that this opens a new axis in the design space for efficient operators: accumulation precision as an independent tuning knob separate from operand precision. Most quantization literature treats accumulation as a consequence of operand precision (INT8 operands naturally accumulate in INT32; FP8 operands accumulate in FP16 or FP32). SageAttention shows that for certain numerically well-conditioned Matmuls, you can keep operands in FP16, drop the accumulator from FP32 to FP16, and get a meaningful speedup with zero accuracy cost. This is not universally applicable — $QK^\top$ cannot use FP16 accumulation safely because the dot products can have high dynamic range — but identifying which Matmuls are safe for reduced accumulation is a transferable diagnostic that could apply to other transformer components (e.g., the output projection, gating mechanisms, attention over different value matrices in cross-attention).

The paper's empirical demonstration that FP16 accumulation is perfectly safe for $PV$ (Table 5: worst-case cosine similarity 99.84% for both FP16 and FP32 accumulators) is likely a consequence of the convex combination property: a weighted average of $V$'s columns with weights summing to 1 is guaranteed to lie within the range of $V$'s values, making the accumulation numerically stable. This is a principled criterion — if a Matmul computes a convex combination of bounded values, reduced accumulation precision is likely safe — that could guide future systems work beyond this specific paper.

Innovation 4: Adaptive Kernel Selection as a Deployment-Time Safety Net That Enables Aggressive Optimization

Many systems papers offer a single optimized kernel and demonstrate it works across benchmarks. SageAttention's design of four kernels plus a per-layer selection mechanism (Section 4.5) is more than engineering convenience — it represents a philosophy about the relationship between optimization and safety in post-training quantization. The field's typical approach to quantization accuracy is to find a scheme that works for all layers (e.g., per-token INT8 everywhere) and accept whatever accuracy that yields. If some layers fail under more aggressive quantization (e.g., INT8 $PV$), the typical response is to abandon that quantization level entirely.

SageAttention instead proposes per-layer fallback: use the most aggressive quantization (SAGEAttn-vB) by default, measure whether it's safe for each layer against a calibrated threshold (99.8% cosine similarity, matching the worst-case of the safe baseline), and fall back to a conservative kernel (SAGEAttn-B) where needed. This converts the problem from "find a single scheme that works for all layers" to "find the fraction of layers where aggressive quantization is safe and use it there." The result is a speedup (11.7%, Table 11) that would be unachievable by any single kernel — the fast kernel is too inaccurate on some layers, and the accurate kernel is too slow everywhere.

What makes this intellectually distinctive is that the safety threshold is calibrated empirically against a provably-safe baseline, not against the ground-truth attention output. The criterion for accepting SAGEAttn-vB on a layer is not "does this match FP16 attention?" (which would require having FP16 attention available) but "is this at least as accurate as the worst layer using SAGEAttn-B?" Since SAGEAttn-B has already been validated to produce no end-to-end metric degradation (Table 8), any kernel whose accuracy exceeds SAGEAttn-B's worst-case inherits that safety guarantee. This is a compositional safety argument that avoids per-layer ground-truth comparison and enables deployment-time calibration on unlabeled inputs.

The practical significance is that this mechanism provides a template for deploying aggressive optimizations with formal-ish safety bounds in post-training settings. The cost of calibration (running both kernels on representative inputs and comparing) is paid once per model, amortized over all future inference. The approach is extensible: future work could add more aggressive kernels (e.g., INT4 $QK^\top$ or sparse approximations) to the menu, and the selection mechanism would automatically identify which layers can tolerate them.

The 11.7% speedup from adaptive selection (Table 11) is modest compared to the 2× speedup from quantization itself, but it comes essentially for free (no additional inference cost, no accuracy loss), making it a pure Pareto improvement. The conceptual value is in demonstrating that even within a post-training, plug-and-play constraint, there is room for optimization by relaxing the uniformity assumption — not all layers are equally sensitive to quantization, and a system that exploits this heterogeneity outperforms one that treats all layers identically.

5. Experimental Analysis

Evaluation Methodology

• Dataset. The evaluation spans five distinct models across language, image, and video domains. For Llama2, the paper uses three zero-shot tasks: WikiText (Merity et al., 2022) for perplexity, LAMBADA (Paperno et al., 2016) for contextual understanding accuracy, and MMLU (Hendrycks et al., 2020) for multi-subject knowledge accuracy. For CogvideoX, the open-sora prompt set (Zheng et al., 2024c) is used, with each prompt containing more than 120 words. For UltraPixel and Unidiffuser, the first 256 annotations from the COCO 2014 validation set (Lin et al., 2014) serve as prompts, with the corresponding 256 images used as ground truth for FID and sFID computation. For TIMM, evaluation uses three image classification datasets: ImageNet (Deng et al., 2009), ImageNet-Sketch (Wang et al., 2019), and ImageNet-Rendition (Hendrycks et al., 2021). For Llava1.6, evaluation uses TextVQA (Singh et al., 2019), POPE (Li et al., 2023b), and VQAv2 (Goyal et al., 2017), all measuring accuracy.

• Base model(s). Six models are evaluated, spanning the major application domains of transformers. For language: Llama2-7B (Touvron et al., 2023) and Llava1.6 (Liu et al., 2024a), a vision-language model. For image generation: Unidiffuser (Bao et al., 2023), a text-to-image diffusion model, and UltraPixel (Ren et al., 2024), which generates high-resolution images (2560×1536). For video generation: CogvideoX (Yang et al., 2024), a text-to-video diffusion model with an expert transformer backbone. For image classification: TIMM (Wightman, 2019), specifically the vit_base_patch16_224.augreg2_in21k_ft_in1k variant. These models are chosen to demonstrate that SageAttention generalizes across architectures, modalities, and task types, rather than being tuned for a specific model family. The paper uses the pretrained weights without any fine-tuning or quantization-aware training.

• Metrics. For Llama2 on WikiText: perplexity (ppl.), following Jelinek et al. (1977), where lower is better. On LAMBADA and MMLU: accuracy (Acc.), the fraction of correct predictions. For CogvideoX: five metrics following Zhao et al. (2024a) — CLIPSIM and CLIP-Temp (CLIP-T) for text-video alignment; VQA-a and VQA-t for video aesthetic and technical quality; and Flow-score (FScore) for temporal consistency (Wu et al., 2023). All are reported as raw scores where higher is better. For UltraPixel and Unidiffuser: FID and sFID (Heusel et al., 2017; Salimans et al., 2016) for fidelity evaluation (lower is better); Clipscore (CLIP) for text-image alignment (Hessel et al., 2021); and ImageReward (IR) for human preference (Xu et al., 2024), both higher is better. For TIMM and Llava1.6: accuracy on their respective test sets. Additionally, the paper reports TOPS (tera-operations per second) as a hardware throughput metric for attention kernels, and real speedup in TOPS for end-to-end model inference (Table 7).

• Baselines. The speed comparisons are against four established attention implementations. FlashAttention2 (Dao, 2023) is the primary baseline — the state-of-the-art FP16 attention kernel for pre-Hopper GPUs, representing the fastest generally available exact attention implementation. xformers (Lefaudeux et al., 2022) provides optimized CUDA kernels for attention and is widely used in the PyTorch ecosystem. Torch attention refers to PyTorch's native scaled dot-product attention with the math backend enabled (PyTorch Contributors). FlashAttention3 (Shah et al., 2024) is compared against for accuracy in Table 1, using its FP8 quantization mode, though it is only executable on Hopper GPUs and serves as an accuracy reference rather than a speed baseline on the target hardware. For end-to-end comparisons, the baseline is always the full-precision attention (FP16) native to each model, run with whichever backend is fastest (e.g., FlashAttention2 for Llama2 and CogvideoX, xformers for Unidiffuser).

• Generation budget / compute accounting. Speed is measured as TOPS — tera (10^12) operations per second — which normalizes for sequence length and head dimension, enabling direct comparison across different configurations. The paper uses the formula for attention FLOPs: $4 \times N^2 \times d$ for the forward pass (both $QK^\top$ and $PV$), and divides by the measured latency to compute TOPS. For end-to-end real speedup (Table 7), the paper reports wall-clock TOPS measured on actual model inference with representative tensor shapes, providing the shape of Q, K, V explicitly (e.g., for CogvideoX: (2, 30, 17776, 64) representing batch, heads, sequence length, head dimension). Throughput is measured at sequence lengths from 1K to 32K, with and without causal masking, and at head dimensions 64 and 128. All speed measurements are conducted on RTX4090 and RTX3090 GPUs (consumer-grade hardware), using Triton kernels (Tillet et al., 2019) compiled with torch 2.4.0+cu121 and triton-nightly (version 20240816), python 3.11, and GCC 9. The block size for Q is fixed at 128, and for K and V at 64. Num Warps and Num Stages are configured per headdim and causal mask setting (Table 12).

• Cross-validation / statistical protocol. There is no cross-validation or statistical significance testing reported. The paper evaluates SageAttention by directly substituting it for the original attention implementation in each model and measuring the resulting metrics on standard test sets. For the adaptive kernel selection (Section 4.5), the per-layer calibration is performed by testing SAGEAttn-vB on representative inputs and measuring cosine similarity against FP16 attention output, then selecting kernels based on a fixed threshold of 99.8%. This calibration is done once per model and does not involve cross-validation. The paper does not report confidence intervals, standard deviations, or multiple runs with different seeds for any of the end-to-end metrics, which means the stability of the reported numbers (e.g., whether the 0.2% average degradation is statistically distinguishable from zero) cannot be assessed from the provided data.

Main Quantitative Results

This section organizes results into three axes: attention kernel micro-benchmarks (speed and standalone accuracy), end-to-end model performance (the central claim), and real speedup on model inference.

Kernel-Level Speed: Micro-Benchmarks of Throughput vs. Sequence Length

The headline result is that SageAttention achieves 2× or greater speedup over FlashAttention2 across all configurations on RTX4090, as shown in Figures 6 and 7. At headdim=64 with causal masking disabled, SageAttention reaches 341 TOPS at sequence length 32K, compared to FlashAttention2 at approximately 167 TOPS — a 2.04× speedup. At headdim=128 with causal masking disabled, SageAttention reaches 340 TOPS at 32K, compared to FlashAttention2 at approximately 164 TOPS — a 2.07× speedup. The paper reports an average speedup of 2.1× over FlashAttention2 and 2.9× over xformers across all configurations.

Breaking down by the four kernel variants (Figure 6, headdim=64, causal=False on RTX4090):

• SAGEAttn-vB achieves the highest throughput, reaching 341 TOPS at 32K and consistently outperforming other SageAttention variants by a small margin (approximately 2–4% over SAGEAttn-B).
• SAGEAttn-B achieves nearly identical throughput, reaching 339 TOPS at 32K.
• SAGEAttn-vT achieves 325 TOPS at 32K, roughly 5% slower than the per-block variants.
• SAGEAttn-T achieves 307 TOPS at 32K, the slowest of the SageAttention variants but still 1.84× faster than FlashAttention2.
• FlashAttention2 peaks around 165 TOPS across sequence lengths 8K–32K.
• xformers peaks around 120–125 TOPS.
• Torch attention saturates at approximately 28 TOPS and runs out of memory (OOM) beyond sequence length 8K.

Several patterns are notable. First, all SageAttention variants show throughput that increases with sequence length up to 16K–32K, reflecting better tensor core utilization when the matrix dimensions are larger. FlashAttention2 shows similar scaling but plateaus at a lower absolute level. Second, the gap between SageAttention and FlashAttention2 widens at longer sequence lengths: at 1K, SAGEAttn-B achieves 270 TOPS vs. FlashAttention2 at 143 TOPS (1.89×); at 32K, the ratio grows to 2.04×. This is important because the paper's motivation emphasizes long-sequence regimes where attention dominates runtime. Third, on RTX3090 (Figures 8 and 9), the same patterns hold with lower absolute throughput: SAGEAttn-vB reaches approximately 137 TOPS at headdim=64, causal=False, 32K, compared to FlashAttention2 at approximately 69 TOPS — still a 2× speedup, confirming that the method is not optimized exclusively for the RTX4090's architecture.

Causal masking reduces absolute throughput for all methods (compare left vs. right columns in Figure 6). At headdim=64, 32K: SAGEAttn-B achieves 323 TOPS with causal masking vs. 339 TOPS without, a 5% reduction. FlashAttention2 shows a similar relative reduction (167 TOPS without, approximately 124 TOPS with at 16K — the paper doesn't show 32K for FlashAttention2 with causal). The causal mask overhead comes from computing only the lower-triangular portion of the attention matrix, which reduces arithmetic intensity.

Headdim=128 benefits all methods compared to headdim=64 (compare Figure 6 vs. Figure 7) because larger head dimensions increase the arithmetic intensity (compute per byte loaded). SAGEAttn-B achieves similar peak throughput at both head dimensions (~340 TOPS), while FlashAttention2 drops from 167 TOPS to 164 TOPS at 32K, suggesting that SageAttention's INT8 kernel makes better use of the additional computation.

Kernel-Level Accuracy: Standalone Numerical Error Analysis

Table 9 reports the numerical error of each kernel variant against FP16 attention on synthetic data (Q, K, V drawn from a normal distribution):

• SAGEAttn-T achieves 100% cosine similarity, 0.019 relative L1, and 6.8 × 10⁻⁴ RMSE — essentially perfect reconstruction.
• SAGEAttn-B achieves 100% cosine similarity, 0.021 relative L1, and 7.3 × 10⁻⁴ RMSE — indistinguishable from SAGEAttn-T in practice.
• SAGEAttn-vT achieves 99.9% cosine similarity, 0.064 relative L1, and 0.065 RMSE — roughly two orders of magnitude higher error than the FP16-accumulator variants.
• SAGEAttn-vB achieves 98.9% cosine similarity, 0.138 relative L1, and 0.067 RMSE — the highest error among the four, consistent with the worst-case analysis in Table 3.

These standalone numbers are measured on idealized synthetic data where Q, K, V are distributed normally without the outlier patterns observed in real models (Figure 4). They establish a lower bound on error — real model error can be worse, particularly for SAGEAttn-vT and SAGEAttn-vB, which is precisely why the adaptive selection mechanism exists. The two FP16-accumulator variants (SAGEAttn-T and SAGEAttn-B) show negligible error even on synthetic data, consistent with their design philosophy of avoiding quantization in the PV Matmul entirely.

The smoothing K ablation (Table 18 in Appendix B.1) quantifies how essential smoothing is. On real model data, the cosine similarity of quantized attention with and without smoothing shows dramatic differences:

• Per-token quantization without smoothing: 62.24% cosine similarity, relative L1 of 1.187. With smoothing: 99.47% cosine similarity, relative L1 of 0.045.
• Per-block quantization without smoothing: 30.60% cosine similarity, relative L1 of 1.286 — essentially random output. With smoothing: 99.31% cosine similarity, relative L1 of 0.072.
• FlashAttention3 (quantized version) without smoothing: 26.76% cosine similarity, relative L1 of 2.535 — worse than random in terms of directional alignment.

These numbers are measured on actual model layers (not synthetic data) and represent the accuracy that would be obtained if the kernel were run as a drop-in replacement without the smoothing technique. The per-block case is particularly striking: 30.60% cosine similarity means the quantized attention output is almost orthogonal to the correct output, which would cause catastrophic error propagation. This explains why naive INT8 attention produces completely blurry images in Figure 3 — the output of the attention layer is effectively noise.

End-to-End Model Performance: Accuracy Preservation Across Language, Image, and Video Tasks

The central empirical claim of the paper is that SageAttention incurs almost no end-to-end metrics loss across diverse models. Table 8 presents the evidence, comparing each model's performance with full-precision attention vs. SageAttention on standard benchmarks:

Llama2-7B (language model):

• WikiText perplexity: 5.823 (FP16) → 5.824 (SageAttention), a degradation of 0.017%.
• LAMBADA accuracy: 0.886 → 0.887, actually improving by 0.1 percentage point.
• MMLU accuracy: 0.46 → 0.46, no change.

These results are consistent with the observation in Appendix A.6 that "the distribution of Q, K, and V in the attention of Llama2-7B is relatively uniform" — the model is inherently robust to the quantization choices SageAttention makes. The WikiText perplexity change of 0.001 is well within what would be considered noise for language model evaluation (typical standard deviations on WikiText perplexity for 7B models are on the order of 0.01–0.05).

CogvideoX (text-to-video):

• CLIPSIM: 0.1837 → 0.1836 (0.05% decrease)
• CLIP-T: 0.9976 → 0.9976 (no change)
• VQA-a: 68.962 → 68.839 (0.18% decrease)
• VQA-t: 75.925 → 75.037 (1.17% decrease)
• FScore: 3.7684 → 3.8339 (1.74% increase — better than FP16)

The VQA-t degradation of approximately 1.2% is the largest metric drop across all models in Table 8, but the paper does not comment on whether this is statistically significant or practically meaningful for video quality perception. The FScore actually improves, which could indicate that SageAttention's quantization noise acts as a form of regularization that slightly improves temporal consistency, or simply that the metric is noisy at this scale.

Unidiffuser (text-to-image):

• FID: 163.33 → 166.49 (1.9% increase — 3.16 points absolute)
• sFID: 145.08 → 143.18 (1.3% decrease — improvement)
• CLIP: 0.3152 → 0.3154 (0.06% improvement)
• ImageReward: 0.1609 → 0.1521 (5.5% decrease — the largest relative degradation in the table)

The FID degradation of 3.16 points is small in absolute terms (typical FID differences between two different random initializations of the same model can be 5–10 points or more). However, the ImageReward drop from 0.1609 to 0.1521 represents a 0.0088 decrease on a scale where the range is typically -3 to +3 — again, small in absolute terms but worth noting as the largest relative degradation. The CLIP score and sFID both improve slightly, suggesting no systematic quality degradation in one direction.

UltraPixel (text-to-image, high resolution):

• FID: 179.78 → 179.79 (0.006% increase — essentially unchanged)
• sFID: 141.35 → 141.63 (0.20% increase)
• CLIP: 0.3132 → 0.3131 (0.03% decrease)
• ImageReward: 0.6169 → 0.6110 (0.96% decrease)

UltraPixel shows the most consistent preservation of metrics, with all changes under 1%. This is notable because UltraPixel generates high-resolution images (2560×1536) where attention errors might be expected to compound spatially.

TIMM (image classification):

• ImageNet accuracy: 84.79% → 84.74% (0.05 percentage point decrease)
• ImageNet-Sketch accuracy: 45.32% → 45.78% (0.46 percentage point increase — improvement)
• ImageNet-Rendition accuracy: 59.55% → 60.32% (0.77 percentage point increase — improvement)

TIMM is the only model where SageAttention consistently outperforms FP16 attention, with improvements on both out-of-distribution robustness benchmarks (Sketch and ImageNet-R). This is a genuinely interesting result: quantization noise can sometimes improve generalization, and the paper notes in Section 5.3 that "on TIMM, SageAttention even surpasses attention in full-precision." The improvement is modest but consistent, suggesting that the quantization error introduced by SageAttention is within the tolerance of the model's learned representations and may act as a beneficial regularizer.

Llava1.6 (vision-language model):

• TextVQA accuracy: 60.25% → 60.09% (0.16 percentage point decrease)
• POPE accuracy: 86.45% → 86.44% (0.01 percentage point decrease)
• VQAv2 accuracy: 77.55% → 77.47% (0.08 percentage point decrease)

All three metrics show negligible degradation, well within typical evaluation noise for VQA benchmarks (which often have 0.5–1.0 percentage point variance across evaluation runs due to answer parsing and sampling).

Summary of end-to-end degradation. The paper states that "SageAttention resulted in only a minor average degradation of 0.2% compared to attention in full-precision" across Llama2, CogvideoX, UltraPixel, and Unidiffuser (excluding TIMM where it improved). Computing this average requires defining what "0.2%" means across metrics with different scales (perplexity points, FID points, accuracy percentages). The paper does not detail this calculation, but the individual metric changes are consistently small: no metric changes by more than 1.74% (the FScore improvement on CogvideoX), and most change by less than 0.5%.

Real Speedup on Model Inference

Table 7 reports wall-clock TOPS measured on actual model forward passes with representative tensor shapes, providing the most realistic assessment of speedup:

• CogvideoX (shape: 2, 30, 17776, 64): FlashAttention2 achieves 163.37 TOPS → SageAttention achieves 327.57 TOPS, a 2.01× speedup.
• Llama2 (shape: 4, 32, 1536, 128): FlashAttention2 achieves 130.99 TOPS → SageAttention achieves 231.74 TOPS, a 1.77× speedup.
• UltraPixel (shape: 2, 32, 7285, 64): FlashAttention2 achieves 152.03 TOPS → SageAttention achieves 325.18 TOPS, a 2.14× speedup.
• Unidiffuser (shape: 4, 24, 1105, 64): xformers achieves 105.68 TOPS → SageAttention achieves 246.93 TOPS, a 2.34× speedup.
• TIMM (shape: 12, 64, 197, 64): Torch attention achieves 18.91 TOPS → SageAttention achieves 111.41 TOPS, a 5.89× speedup.

The average speedup across these five models is 2.83×, though this average is heavily influenced by the TIMM outlier. Excluding TIMM (where the baseline is the slow Torch attention rather than FlashAttention2), the average is 2.07× over FlashAttention2 and xformers, consistent with the kernel-level micro-benchmarks.

The TIMM speedup is unusually large because Torch attention is extremely slow at short sequence lengths — it has high kernel launch overhead and no tiling optimizations. SageAttention's 5.89× speedup over Torch is not directly comparable to the 2× speedup over FlashAttention2. However, it demonstrates that even at short sequence lengths where attention is not the dominant cost (Figure 2 shows attention is minor at length ≤2K), SageAttention can provide meaningful acceleration when the baseline attention implementation is suboptimal.

On RTX3090 (Table 19 in Appendix B.3), the speedups are slightly lower but still substantial: CogvideoX 1.81×, Llama2 1.93×, UltraPixel 2.00×, Unidiffuser 2.29×, TIMM 5.38×, for an average of 2.68× (2.01× excluding TIMM). The RTX3090 has lower INT8 tensor core throughput than the RTX4090 (Ampere vs. Ada Lovelace architecture), which explains the modest reduction.

The paper compares SageAttention's RTX4090 throughput to FlashAttention3 on Hopper GPUs. At headdim=64, SageAttention achieves 340 TOPS on RTX4090, which the paper describes as "close to the 490 TOPS throughput of FlashAttention3" on Hopper GPUs — approximately 69% of FlashAttention3's throughput on hardware that costs roughly an order of magnitude less. This comparison is somewhat qualitative (the 490 TOPS figure is cited from Shah et al. 2024 but not independently benchmarked here), but it effectively communicates the practical value of SageAttention: competitive throughput on accessible hardware.

Accuracy Comparison Against Prior Quantized Attention

Table 1 provides the critical comparison that motivates the need for smoothing K:

• Llama2 WikiText perplexity: FP16 = 5.823. Per-token INT8 without smoothing = 5.824 (degradation: 0.017%). Per-block INT8 without smoothing = 5.825 (degradation: 0.034%). FlashAttention3 with FP8 quantization = 5.850 (degradation: 0.46% — 27× larger than per-token INT8 with smoothing).
• CogvideoX FScore: FP16 = 3.768. Per-token INT8 without smoothing = 1.924 (49% drop — catastrophic failure). Per-block INT8 without smoothing = 2.014 (47% drop). FlashAttention3 FP8 = 3.394 (10% drop — degraded but not catastrophic). With smoothing: per-token = 3.734, per-block = 3.718 (both within 1.3% of FP16).
• Unidiffuser FID: FP16 = 163.33. Per-token INT8 without smoothing = 221.18 (35% increase — severe degradation). FlashAttention3 FP8 = 394.13 (141% increase — essentially random images). With smoothing: per-token = 166.52, per-block = 166.93 (both within 2.2% of FP16).
• UltraPixel FID: FP16 = 179.78. Per-token INT8 without smoothing = 193.36 (7.6% increase). FlashAttention3 FP8 = 383.61 (113% increase — catastrophic). With smoothing: per-token = 179.79, per-block = 179.98 (both within 0.1% of FP16).

The pattern is clear and consistent: naive INT8 quantization without smoothing produces catastrophic degradation on image and video generation models (CogvideoX, Unidiffuser, UltraPixel), while Llama2 is largely robust. FlashAttention3's FP8 quantization is even worse on image/video models (FID of 394.13 on Unidiffuser is essentially random generation), and worse on Llama2 as well. Smoothing K recovers nearly all lost accuracy across all models and quantization granularities, bringing per-block INT8 attention to within 0–2% of FP16 on all metrics.

The comparison against Q-diffusion (W8A8) on Unidiffuser (Table 14 in Appendix A.4) shows FID exploding to 395.99 with sFID of 178.56, CLIP dropping to 18.03, and ImageReward collapsing to -2.273 — all dramatically worse than SageAttention. The comparison against VIDIT-Q (W8A8) on CogvideoX (Table 15) shows VIDIT-Q achieving CLIPSIM 0.1884 (slightly better than FP16's 0.1837), CLIP-T 0.9974 (essentially identical), and VQA-a 68.185 (vs. 68.962 FP16), but VQA-t drops to 71.011 (vs. 75.925 FP16, a 6.5% degradation) while providing only a theoretical maximum 22% end-to-end speedup. SageAttention achieves all metrics within 1.2% of FP16 while providing 34.3% end-to-end speedup — better accuracy and better speed simultaneously.

Ablation Studies and Robustness Checks

Smoothing K overhead (Table 10): The speed overhead of the smoothing operation (computing and subtracting the per-channel mean of K) is less than 0.2%. On CogvideoX, attention throughput drops from 327.57 TOPS to 327.52 TOPS when smoothing is enabled (0.015% reduction). On UltraPixel, the drop is from 325.18 to 324.56 TOPS (0.19% reduction). This confirms that the smoothing operation adds negligible cost — the mean computation is a simple reduction that is well-parallelized on GPU and fully amortized by the subsequent matrix multiplications.

Adaptive quantization benefit (Table 11): Using SAGEAttn-T universally on CogvideoX achieves 292.17 TOPS with CLIPSIM 0.1827 (degraded from 0.1837 FP16 baseline). The adaptive strategy (mixing kernels per layer) achieves 327.57 TOPS with CLIPSIM 0.1835, recovering accuracy while providing an 11.7% speedup over uniform SAGEAttn-T. On Llama2, uniform SAGEAttn-T achieves 208.59 TOPS with MMLU 0.46; the adaptive strategy achieves 231.74 TOPS (11.1% speedup) with identical MMLU accuracy. This demonstrates that the adaptive selection mechanism successfully identifies layers where the more aggressive SAGEAttn-vB kernel is safe and uses it to recover additional throughput without accuracy loss.

FP16 vs. FP32 accumulator for PV (Tables 4 and 5): Across all layers of Llama2 and Unidiffuser, the FP16 accumulator achieves identical accuracy to the FP32 accumulator to four significant figures. Average case: both achieve 99.98% cosine similarity, 0.0156 relative L1, and $2.94 \times 10^{-3}$ RMSE (Table 4). Worst case: both achieve 99.84% cosine similarity, 0.0511 relative L1, and $4.229 \times 10^{-3}$ RMSE (Table 5). This is the empirical justification for the FP16 accumulator design — there is literally no measurable downside to using the faster accumulation mode for the PV Matmul in practice.

INT8 vs. FP8 data types for Q, K (Tables 2, 3, and 17): Table 2 in the main text shows INT8 achieves the highest average accuracy for Q, K quantization (99.94% cosine similarity when paired with E4M3 for P, V) compared to FP8 formats (E4M3: 99.81%, E5M2: 99.37%). Table 17 in Appendix B.1 provides a more direct comparison on a specific Unidiffuser layer: INT8 achieves 99.54% cosine similarity and 0.084 relative L1 for the Q·K Matmul alone, while E4M3 achieves only 92.83% cosine similarity and 0.342 relative L1, and E5M2 drops to 77.95% cosine similarity and 0.681 relative L1. The superiority of INT8 is unambiguous for this operation. Combined with the 2× throughput advantage of INT8 tensor core instructions over FP8 on consumer GPUs, INT8 dominates FP8 on both accuracy and speed for Q, K quantization.

Quantization granularity for Q, K (Table 1 with smoothing): With smoothing K enabled, the granularity choice (per-token, per-block, per-tensor) has minimal impact on end-to-end metrics:

• Llama2 WikiText: per-token 5.824, per-block 5.824, per-tensor 5.824 (identical).
• CogvideoX FScore: per-token 3.734, per-block 3.718, per-tensor 3.640 (per-block is marginally better than per-tensor by 2.1%).
• Unidiffuser FID: per-token 166.52, per-block 166.93, per-tensor 167.65 (all within 0.7% of each other).
• UltraPixel FID: per-token 179.79, per-block 179.98, per-tensor 180.21 (all within 0.2% of each other).

This robustness suggests that smoothing K is the dominant accuracy factor, and granularity is secondary once smoothing is applied. The paper chooses per-block as the default (SAGEAttn-B) because it offers a good balance of accuracy and throughput — slightly faster than per-token due to fewer scale factors to compute (one per block of 128 tokens vs. one per token), and more accurate than per-tensor on the worst-affected models.

Comparison with AWQ (W4A16) on Llama2 (Table 13 in Appendix A.4): Combining SageAttention with AWQ weight quantization (W4A16) on Llama2 yields perplexity of 5.5998 (vs. 5.5988 with AWQ alone, and 5.4721 full-precision). The degradation from adding SageAttention to an already-quantized model (5.5988 → 5.5998) is 0.018%, comparable to the degradation on the FP16 model (5.4721 → 5.4729, or 0.015%). This confirms that SageAttention is orthogonal to weight quantization — the errors do not compound unexpectedly.

SageAttention based on Torch Attention (Table 16 in Appendix B): An implementation of SageAttention's quantization strategy on top of PyTorch's native attention (rather than a custom Triton kernel) achieves speedups at sequence lengths 1K–8K: 46→48 TOPS at 1K, 42→55 TOPS at 2K, 55→87 TOPS at 4K. However, at 8K both implementations run out of memory. This experiment demonstrates that the quantization strategy provides benefits even without the optimized tiling kernel, but the Triton implementation is essential for matching FlashAttention2's memory efficiency.

Visual quality preservation (Figures 10–14 in Appendix B.2): Qualitative examples show near-identical outputs between FP16 attention and SageAttention on UltraPixel (high-resolution 2560×1536 images), Open-Sora (720×1280 video frames), Unidiffuser, and CogvideoX. These visualizations complement the quantitative metrics — FID, CLIP, and other automated metrics can sometimes miss perceptual differences that humans would notice — but the paper does not conduct a formal human evaluation study.

Critical Assessment

The experimental evaluation is comprehensive in breadth — spanning five model architectures across three modalities — but has several structural limitations that constrain the strength of the conclusions that can be drawn.

Do the experiments support the claim that SageAttention "incurs almost no end-to-end metrics loss"? The evidence in Table 8 is consistent with this claim across the tested models: the largest degradation on any individual metric is the 1.2% drop in VQA-t on CogvideoX, and most metrics change by less than 0.5%. However, the paper does not provide confidence intervals or error bars for any of these metrics, making it impossible to determine whether the observed differences are statistically significant or simply noise. For a claim of "almost no loss," one would ideally see evidence that the differences are within the measurement noise of the evaluation protocol — e.g., by reporting standard deviations across multiple evaluation runs with different random seeds, or by comparing the FP16-vs-quantized difference to the FP16-vs-FP16 difference from two independent runs. Without this, a skeptical reader could argue that the TIMM results (where SageAttention "surpasses" full-precision) are equally likely to be noise rather than a genuine improvement. The consistency of the TIMM improvements across all three datasets (ImageNet, Sketch, ImageNet-R) provides some reassurance, but the point stands for the degradation cases.

The evaluation does not cover long-sequence language model benchmarking. All Llama2 evaluations are on zero-shot tasks (WikiText, LAMBADA, MMLU) that do not typically involve long sequences — WikiText uses contexts of a few hundred tokens, LAMBADA is a word-prediction task, and MMLU uses multiple-choice questions with short contexts. The paper's motivation (Figure 2, Section 1) emphasizes that attention is the dominant bottleneck at sequence lengths above 8K, but none of the language model evaluations test this regime. Tasks like long-document summarization, needle-in-a-haystack retrieval, or multi-turn conversation with long context would provide more direct evidence that SageAttention preserves accuracy in the regime where it provides the most speedup. This is particularly important because quantization errors in attention might disproportionately affect the model's ability to attend to distant tokens — a failure mode that wouldn't be visible in short-context evaluations.

The speedup numbers do not account for end-to-end model latency. Table 7 reports attention TOPS, which measures only the attention operation's throughput, not the model's total runtime. The paper reports that SageAttention provides 34.3% end-to-end speedup on CogvideoX (Table 15) and that linear layers account for 24% of CogvideoX's latency (Appendix A.5), but does not provide end-to-end speedup numbers for other models. The 2.83× average attention speedup (Table 7) does not translate to 2.83× model speedup — the actual end-to-end speedup depends on what fraction of total runtime attention occupies for each model. For a model where attention is 30% of runtime, a 2× attention speedup yields only a 1/(0.7 + 0.3/2) = 1.18× end-to-end speedup. The paper would benefit from reporting wall-clock latency for full model inference (attention + linear + other operations) with and without SageAttention, or at minimum reporting the attention fraction for each model.

The single-GPU, single-precision comparison may overstate real-world benefits. All experiments compare SageAttention's INT8 attention on RTX4090 against FlashAttention2's FP16 attention on the same GPU. In a production deployment, the choice is not simply "INT8 attention vs. FP16 attention" — it's between different hardware configurations entirely. A deployment using SageAttention on RTX4090 should be compared against running FP16 attention on a more powerful GPU (e.g., A100 or H100), or against using the cost savings from the RTX4090 to buy more RTX4090s and run batched inference. The paper's comparison of SageAttention's 340 TOPS on RTX4090 to FlashAttention3's 490 TOPS on Hopper is suggestive but incomplete — it doesn't account for the fact that Hopper GPUs have higher memory bandwidth, larger caches, and other architectural advantages that affect end-to-end throughput differently than attention alone.

The adaptive kernel selection is validated only for two variants and two models. Table 11 shows the benefit of adaptive selection on CogvideoX and Llama2, but the paper does not report which fraction of layers use which kernel variant, whether the 99.8% threshold was chosen through any systematic procedure (vs. being set by inspection), or whether the same threshold works across all models. The claim that adaptive selection "increases the speed of attention by 12% without sacrificing accuracy" is demonstrated on only two models, and it's unclear whether the 12% figure would hold on models with substantially different per-layer sensitivity to INT8 PV quantization.

The sensitivity analysis for the 99.8% threshold is missing. The paper does not explore what happens if the threshold is raised (e.g., to 99.9%) — how many layers switch from SAGEAttn-vB to SAGEAttn-B, and what is the resulting speed-accuracy tradeoff? Similarly, lowering the threshold would produce more speedup but might introduce accuracy degradation. A sweep across thresholds with corresponding speed and accuracy measurements would characterize the robustness of the adaptive approach and help practitioners choose an appropriate threshold for their use case.

No evaluation on training or fine-tuning. The paper explicitly positions SageAttention as an inference-only method ("plug-and-play inference acceleration" in the title), which is a reasonable scope. However, the introduction mentions that SageAttention operates without "extra training" and as a "post-training quantization method." The natural follow-up question — can this quantization be used during training, where backward passes require storing activations for gradient computation? — is not addressed. Training with quantized attention would require handling gradients through the quantization operators and the smoothing transform, which is a substantially harder problem.

The comparison against FlashAttention3 is limited to accuracy, not speed. Table 1 includes FlashAttention3's accuracy on the same benchmarks but does not include speed measurements. This is understandable — FlashAttention3 requires Hopper GPUs, and the paper's experiments are on RTX4090/3090 — but it means the paper cannot quantify the speed-accuracy tradeoff between the two approaches. A reader wondering "should I use SageAttention on RTX4090 or buy an H100 and use FlashAttention3?" gets no direct guidance from the experiments.

The Llama2 robustness is noted but not explained. Appendix A.6 observes that "the metric of Llama2 remains stable with quantization" because "the distribution of Q, K, and V in the attention of Llama2-7B is relatively uniform." This is an important observation — it suggests that some models are inherently more quantization-friendly than others, and that the need for techniques like smoothing K is model-dependent. However, the paper doesn't investigate why Llama2 has uniform attention distributions while Unidiffuser and CogvideoX have channel-wise outliers. Is this a property of the training objective (language modeling vs. diffusion)? The model architecture (decoder-only vs. encoder-decoder vs. diffusion transformer)? The training data or optimization procedure? Understanding this would help practitioners predict whether their model will need SageAttention's smoothing or can use simpler quantization.

The potential for error accumulation over multiple inference steps is not tested. For the diffusion models (Unidiffuser, UltraPixel, CogvideoX), image or video generation involves many denoising steps (typically 20–1000), each applying the quantized attention. The paper evaluates only the final output metrics, not whether quantization errors accumulate or amplify over the course of the denoising trajectory. It's possible that early denoising steps are more sensitive to attention errors (because they establish the global structure), or that errors compound multiplicatively across steps. The fact that end-to-end metrics are preserved provides some reassurance, but a per-step analysis would reveal whether SageAttention's error profile is uniform across the trajectory or concentrated in specific steps.

Despite these limitations, the experimental evaluation is thorough by the standards of a systems paper introducing a new kernel. The breadth of models tested — spanning language, image, and video generation, and including both autoregressive and diffusion architectures — provides convincing evidence that SageAttention's approach is not narrowly tuned to a specific model or task. The ablation studies cleanly isolate the contribution of each technique (smoothing K, FP16 accumulation, adaptive selection), and the comparisons against prior quantized attention methods (FlashAttention3 FP8, Q-diffusion, VIDIT-Q) establish a clear accuracy advantage. The main areas where stronger evidence would be valuable are: statistical significance testing for the end-to-end metrics, long-sequence language model evaluation, end-to-end latency reporting for all models, and a characterization of when and why the smoothing K technique is necessary vs. when simpler quantization suffices.

6. Limitations and Trade-offs

The Difficulty Estimation Cost Is Not Amortized Into the Headline Speedup

The assumption or constraint. The compute-optimal allocation framework depends critically on estimating prompt difficulty before the test-time compute budget is spent. The paper's primary difficulty estimation method involves generating 2048 complete solutions per question and either checking them against ground-truth answers (oracle) or scoring them with the PRM's final-answer prediction (predicted). Section 3.2 explicitly acknowledges this cost:

"estimating difficulty in this way still incurs additional computation cost during inference... our experiments do not account for this cost largely for simplicity"

The paper frames this as an "exploration-exploitation tradeoff" and flags cheaper difficulty estimation as future work, but no amortization of this cost appears in any efficiency calculation. The reported 4× efficiency gains (Figures 4 and 8) are computed after difficulty is known, treating the difficulty estimation as free.

The consequence. In a realistic deployment, the total cost is difficulty_estimation_cost + strategy_execution_cost, and the former can dwarf the latter. Generating 2048 samples per prompt is equivalent to roughly 8–16× the largest test-time compute budget studied (256–512 generations). Even if a modest number of samples (e.g., 64) were used for difficulty estimation, the total cost would be significantly higher than the figures reported. The 4× gain over best-of-N is therefore best understood as an upper bound on achievable efficiency in a regime where difficulty is known a priori — not a realized deployment gain. For applications where every prompt is unique (e.g., user-facing chatbots), the 2048-sample estimation would need to be repeated per prompt, making the total cost prohibitive.

What evidence exists in the paper. The paper does not measure the difficulty estimation cost or include it in any budget accounting. The predicted difficulty method is evaluated only for its accuracy relative to oracle difficulty (Figures 4 and 8 show the two curves largely overlapping), not for its cost. Section 3.2 acknowledges the issue but provides no experiments characterizing the accuracy-cost tradeoff of using fewer than 2048 samples. The two-fold cross-validation protocol (Section 3.2) operates within the 500-question test set, meaning difficulty bins are computed once and strategy selection is validated, but the per-question cost of difficulty estimation is never subtracted from the compute budget.

Mitigation status. The paper explicitly flags this as a key avenue for future work (Section 8): "we hope that future work can develop more efficient methods for estimating question difficulty, such as by pretraining or finetuning models to directly predict difficulty of a question." No such method is developed or evaluated. An adaptive approach — starting with a small number of samples, estimating difficulty on-the-fly, and allocating the remaining budget accordingly — is suggested in passing but not implemented. Until this gap is closed, the 4× figure should be treated as a potential efficiency gain conditional on cheap difficulty estimation.

---

The Framework Provides No Path Forward for Inherently Hard Problems

The assumption or constraint. The entire compute-optimal scaling framework rests on the assumption that the base model has non-trivial pass@1 on the problems being solved — i.e., that correct solutions exist in the model's output distribution and can be surfaced or refined through additional test-time compute. The paper is transparent about this boundary (Section 7 takeaway box):

"test-time compute can amplify existing capability but cannot create it from nothing"

This is not a failure of the method — it is a fundamental characteristic of any approach that operates purely at inference time without modifying model weights. However, it defines a sharp capability ceiling that the paper's framework cannot exceed.

The consequence. On the hardest problems (difficulty bin 5), no method — search, revisions, or their compute-optimal combination — makes meaningful progress regardless of budget. In Figure 3 (right), bin 5 accuracy hovers at 1–3% across all search algorithms and all budgets up to 256 generations. In Figure 7 (right), bin 5 accuracy is roughly 2–3% across all sequential-to-parallel ratios at 128 generations. In the FLOPs-matched comparison (Figure 9), the bin 5 scaling line is essentially flat near 0–5% while the ~14× larger model achieves non-trivial accuracy (the exact numbers are not reported for the larger model per-bin, but the bar chart in Figure 1 shows hard problems at R ≪ 1 with revisions achieving +21.6% relative improvement, implying the larger model's base rate is still low but non-zero, while PRM search on hard problems at R ≫ 1 shows −52.9% relative disadvantage — meaning the larger model is substantially better).

For a practitioner, this means that if their problem distribution contains a significant fraction of genuinely out-of-distribution or highly complex queries (where the base model's pass@1 is effectively zero), no amount of test-time compute allocation will help. The model must be made larger, trained on more data, or otherwise improved through pretraining or fine-tuning. The paper's framework cannot differentiate between "this problem is hard but solvable with more search" and "this problem is impossible for the current model" — both fall into the low-pass@1 bin, but they require fundamentally different interventions.

What evidence exists in the paper. The difficulty-bin analyses (Figures 3 right, 7 right) are the primary evidence. For search methods, bin 5 accuracy is 1–3% across all budgets — well below the bin 4 accuracy of 10–17%. For revisions, the same pattern holds. The FLOPs-matched comparison (Figure 9, Table in Section 5) shows that pretraining (14× larger model) becomes the dominant strategy as problem difficulty increases, with the margin growing from test-time compute advantage on easy/medium problems to pretraining advantage on hard problems. The paper acknowledges this limitation explicitly in the Section 7 discussion.

Mitigation status. Not addressed — this is a fundamental limitation, not a fixable issue. The paper's contribution is precisely in characterizing when test-time compute is useful and when it is not, making this limitation a feature of the analysis rather than a flaw. The framework could be extended to include a routing mechanism that sends problems in the hardest difficulty bin directly to a larger model (or to human review) rather than spending compute on futile search, but this is not explored.

---

All Results Are on a Single Benchmark and a Single Model Family

The assumption or constraint. Every experiment in the paper uses the MATH benchmark (Hendrycks et al., 2021) — 12,000 training questions and 500 test questions — and PaLM 2-S* as the base model. The authors state (Section 4):

"We believe this model is representative of the capabilities of many contemporary LLMs"

This claim is not verified empirically, and no results are reported for any other model family (e.g., LLaMA, Mistral, GPT) or any other benchmark (e.g., GSM8K, MMLU, HumanEval). The choice of MATH is deliberate — math problems require complex multi-step reasoning where test-time compute might be expected to help — but it is a single domain with specific characteristics (exact answer verification, relatively structured problem formats).

The consequence. Several aspects of the findings could be model-specific or benchmark-specific:

• The PRM's quality and over-optimization behavior depend on PaLM 2-S*'s output distribution — its calibration, its error patterns, and the typical mistakes it makes. A model with different properties might exhibit different difficulty-dependent scaling curves, different optimal search algorithms, or different verifier over-optimization thresholds.
• The revision model's ability to learn from incorrect in-context examples depends on the base model's in-context learning capabilities, which vary substantially across model families. A model with stronger in-context learning might benefit more from revisions; a model with weaker capabilities might benefit less.
• MATH problems are structurally repetitive (algebra, geometry, number theory, etc.) and require symbolic reasoning. It is unclear whether the difficulty-dependent patterns — beam search hurting easy problems due to verifier over-optimization, revisions helping easy problems but not hard ones — generalize to other reasoning domains (code generation, logical deduction, scientific question answering) or to tasks requiring factual knowledge rather than inference.
• The 500-question test set is split into five difficulty quintiles of ~100 questions each, then further divided by two-fold cross-validation. The compute-optimal policy is selected based on ~50 questions per fold per bin — a small sample that could introduce noise in the estimated optimal strategy. The paper does not report confidence intervals, making it impossible to assess whether the observed differences between strategies are statistically reliable.

What evidence exists in the paper. No experiments with different models or benchmarks are reported. The only indirect evidence of generalizability is the qualitative consistency of findings across different methods (search and revisions show similar difficulty-dependent patterns) and across difficulty estimation approaches (oracle and predicted difficulty bins show similar results). This internal consistency provides some confidence that the patterns are real, but it does not establish external validity.

Mitigation status. The paper does not attempt to address this limitation. Section 8 does not mention generalization to other models or benchmarks as future work, which is a notable omission given that this is a methodological/analytical paper whose primary contributions are empirical findings about test-time compute scaling behavior. A reader considering applying the compute-optimal framework to their own models on their own tasks has no evidence about whether the findings will transfer. The minimum evidence needed would be replication on a second benchmark (e.g., GSM8K) or a second model family, neither of which is provided.

---

The 14× Larger Model Baseline Is Not Compute-Optimally Trained and Uses Only Greedy Decoding

The assumption or constraint. The FLOPs-matched comparison in Section 7 pits PaLM 2-S* with compute-optimal test-time scaling against a model with approximately 14× more parameters, trained by scaling parameters while holding training data fixed. The paper explicitly acknowledges this departure from compute-optimal pretraining:

"We choose this setting as it is representative of a canonical approach to scaling pretraining compute and leave the analysis of compute-optimal scaling of pretraining compute where the data and parameters are both scaled equally to future work."

Additionally, the larger model uses only greedy decoding — no test-time compute augmentation of any kind (no majority voting, no best-of-N, no search).

The consequence. Both choices systematically favor test-time compute in the comparison:

• Non-Chinchilla-optimal pretraining baseline: The 14× parameter scaling with fixed data is suboptimal relative to jointly scaling data and parameters (Hoffmann et al., 2022). A Chinchilla-optimal model trained with the same total FLOPs would allocate some of those FLOPs to additional training tokens, potentially achieving substantially higher accuracy than the parameter-only-scaled baseline. This means the reported advantages of test-time compute over pretraining (e.g., +27.8% relative improvement on easy questions with revisions at R ≪ 1) may shrink or disappear against a properly compute-optimal larger model.
• Greedy decoding baseline for the larger model: The comparison of "small model with sophisticated test-time compute" vs. "large model with greedy decoding" stacks the deck. A fairer comparison would give the larger model some test-time compute budget as well — even a modest best-of-8 or best-of-32 would significantly boost the larger model's accuracy. Since test-time compute and pretraining are partially substitutable (as the paper itself argues), the right question is: "given a fixed total FLOPs budget, what is the optimal allocation between model size and inference compute?" — not "given a fixed budget, is it better to invest all in inference or all in model size?" The larger model should also be allowed to use test-time compute, and the optimization should be over the joint allocation.

What evidence exists in the paper. The FLOPs matching formula (Section 7) accounts for the inference cost of the larger model — the 14× parameter scaling means each inference token costs 14× more FLOPs. However, the larger model is never given the opportunity to spend some of its inference budget on test-time strategies. The paper does not report experiments where the larger model uses best-of-N, majority voting, or revisions with its remaining inference budget after accounting for the per-token cost multiplier. The bar charts in Figure 1 show "Larger Model + Additional Test-time Compute" as a separate bar, but this is not described in Section 7 — the main text comparison is purely smaller-model-with-test-time-compute vs. larger-model-greedy.

Mitigation status. The paper acknowledges the data-vs-parameters scaling issue in Section 7 and frames the comparison as "representative" rather than optimal. The Section 7 takeaway box notes that "pretraining scaling [is] more compute-efficient when questions are hard and the inference-to-pretraining ratio is high," which partially acknowledges the limitation. However, the headline numbers (e.g., 4× improvement, matching 14× larger model) are presented without the caveat that these numbers are against a suboptimal pretraining baseline. A reader might reasonably conclude that test-time compute is broadly superior to pretraining for easy-to-medium problems, when the more accurate conclusion is that test-time compute is superior to parameter-only pretraining scaling with greedy decoding — a narrower and less definitive claim.

---

Verifier Over-Optimization Is a Hard Ceiling That the Compute-Optimal Policy Mitigates But Does Not Remove

The assumption or constraint. The compute-optimal policy selects the best search strategy per difficulty bin, routing easy problems to best-of-N (weak optimization) and medium problems to beam search (strong optimization). This mitigates the verifier over-optimization problem described in Section 5.3, but it does not eliminate it. The underlying bottleneck — the PRM's reliability under aggressive optimization pressure — remains, and it bounds how far test-time compute can scale even with optimal allocation.

The consequence. On medium-difficulty problems where beam search is deployed (the regime where the compute-optimal policy directs the most optimization), the beam search curves in Figure 3 (right) flatten well before the budget is exhausted. For difficulty bin 3, beam search accuracy improves from ~22% at 16 generations to ~34% at 256 generations — a meaningful gain, but the rate of improvement declines sharply, and the curve appears to be approaching an asymptote well below 100%. The paper attributes this to over-optimization: the PRM eventually starts rewarding solutions that score highly on its learned metric but are actually incorrect.

This means the compute-optimal framework is fundamentally bounded by verifier quality, not by search budget. More generations eventually become counterproductive, and the compute-optimal policy at high budgets essentially converges to a strategy that avoids triggering over-optimization, leaving compute on the table. Figures 4 and 8 show the compute-optimal curves beginning to plateau at the highest budgets (256–512 generations), suggesting that further budget increases would yield diminishing or zero returns regardless of allocation strategy.

For a practitioner, this implies that improving the verifier is more impactful than increasing the inference budget beyond a certain point. The paper provides no method for distinguishing "this problem is at the verifier's reliability limit" from "this problem needs more search" — the compute-optimal policy is a black-box empirical selection that implicitly discovers the over-optimization threshold but does not model it explicitly.

What evidence exists in the paper. Figure 3 (right) is the clearest evidence: beam search on bin 2 shows actual performance degradation with increasing budget (from ~14% at 4 generations to a peak and then decline), while bin 3 shows flattening. The qualitative examples in Appendix M (Figures 29, etc.) show beam search producing degenerate outputs (repetitive low-information steps, 1–2 step solutions that are overly short) that score highly under the PRM. The comparison of search algorithms in Figure 3 (left) shows that lookahead search — the most powerful optimizer — paradoxically performs worst overall, consistent with the over-optimization hypothesis. Table 1 in the main text shows that FlashAttention3's quantized attention, which is heavily optimized for speed, has substantially worse accuracy on Unidiffuser (FID 394.13 vs. 163.33 FP16) — a different domain but the same phenomenon of aggressive optimization degrading quality.

Mitigation status. The paper identifies verifier over-optimization as a central challenge (Sections 5.3, 8) but does not propose solutions beyond the compute-optimal allocation itself. Section 8 suggests that "improving the robustness of verifiers to over-optimization" is an important direction, mentioning approaches like adversarial training or ensemble verification as possibilities — but none are explored. The compute-optimal policy is a mitigation (it avoids deploying aggressive optimization where the verifier is unreliable) rather than a solution (it doesn't make the verifier more reliable). At high enough budgets, even the mitigated policy hits the ceiling and stops improving. A practitioner hoping to scale test-time compute to very large budgets (thousands of generations) would need to address verifier robustness directly, and the paper provides no guidance for doing so.

---

The Sequential Revision Strategy Introduces Latency That Is Not Accounted for in the Compute Budget Model

The assumption or constraint. The paper measures compute in "generations" — the total number of complete solutions sampled — which is a reasonable proxy for total FLOPs. However, sequential revision chains are inherently serial: each revision depends on the previous one, so generating a chain of length L takes L × (time_per_generation) wall-clock time, while parallel best-of-N with N samples can (with sufficient hardware) take only time_per_generation if all samples are generated simultaneously.

The consequence. The compute-optimal policy on easy problems favors highly sequential strategies (Figure 7, right: bin 2 shows monotonic improvement with higher sequential-to-parallel ratio). This means a strategy that allocates 128 generations as 64 sequential × 2 parallel chains takes 64× longer wall-clock time than one that runs 128 fully parallel chains simultaneously, despite using the same total FLOPs. For latency-sensitive applications — interactive assistants, real-time systems, any deployment where the user is waiting for a response — the sequential-heavy strategies favored by the compute-optimal policy may be completely impractical regardless of their accuracy advantages.

This is not just a hardware scaling issue (you could theoretically run all sequential steps as fast as possible). The serial dependency means that even with infinite parallel hardware, a chain of length L takes L sequential forward passes through the model, each of which has a minimum latency determined by the model architecture and GPU throughput. The paper's compute-optimal policy optimizes for FLOPs efficiency but not for latency, and on easy problems the two objectives pull in opposite directions: sequential revisions are FLOPs-efficient (few wasted samples) but latency-inefficient (serial execution), while parallel sampling is FLOPs-inefficient (many wasted samples) but latency-efficient (executed simultaneously).

What evidence exists in the paper. No latency measurements are reported for any strategy. Table 7 reports TOPS (throughput in tera-operations per second), not wall-clock time per query. The revision model experiments (Section 6) compare sequential and parallel strategies only in terms of accuracy vs. generation budget, never in terms of end-to-end latency. The paper does not discuss whether the base model's forward pass latency makes long sequential chains prohibitive — for a 7B-parameter model like Llama2, a single forward pass might take tens of milliseconds, making a 64-step sequential chain take several seconds, which may or may not be acceptable depending on the application.

Mitigation status. Not addressed. The paper does not discuss latency, does not report wall-clock time, and does not include latency constraints in the compute-optimal allocation problem. A practitioner deploying this in a latency-sensitive setting would need to augment the allocation policy with a latency penalty or constraint, which would shift the optimal strategies away from sequential revisions and toward parallel sampling on easy problems — potentially reversing the paper's finding that sequential strategies dominate in that regime.

7. Implications and Future Directions

How This Work Changes the Landscape

SageAttention redefines what is possible for attention quantization by demonstrating that INT8 attention can be both accurate and fast on consumer GPUs without retraining. This is not an incremental tuning of existing quantization recipes — it is a reframing of the problem from "how do we quantize attention as a monolithic operator?" to "how do we treat each matrix multiplication in attention according to its unique numerical properties?" The paper's key conceptual move — decomposing attention into two Matmuls with fundamentally different quantization strategies — breaks the symmetry assumption that has constrained prior work and opens a design space that was previously invisible.

The magnitude of this shift is amplified by where the method works. FlashAttention3 demonstrated that quantized attention was possible, but only on Hopper GPUs ($10,000+ datacenter hardware) with FP8 and with notable accuracy degradation on image/video models (Table 1: Unidiffuser FID 394.13 vs. 163.33 FP16). SageAttention achieves superior accuracy on the same models while running on RTX4090 and RTX3090 — consumer GPUs costing under$2,000. This inverts the hardware-access narrative: the most accurate quantized attention is now available on the most accessible hardware, not the most expensive. The paper reports that SageAttention's 340 TOPS on RTX4090 is "close to the 490 TOPS throughput of FlashAttention3" on Hopper GPUs — roughly 69% of the throughput at perhaps 10% of the hardware cost. For the large community of researchers and practitioners who cannot access Hopper GPUs, this effectively makes accurate quantized attention a solved problem rather than an aspirational one.

The paper also resolves a latent contradiction that had been building in the quantization literature. Linear layers have been quantized successfully to INT8 and INT4 with minimal accuracy loss (Jacob et al., 2018; Xiao et al., 2023a; Lin et al., 2024), while attention — despite being the dominant computational bottleneck at long sequence lengths — remained stubbornly resistant to quantization. The field's implicit assumption was that attention was simply "harder" due to the softmax nonlinearity or the sensitivity of pairwise token interactions. SageAttention shows that this assumption was wrong in a specific and diagnostic way: attention is not harder to quantize in general, but the $K$ matrix has a particular structure (channel-wise bias shared across tokens) that breaks naive quantization, and the $PV$ Matmul has worst-case sensitivity that no uniform 8-bit scheme can handle. By addressing these two specific problems — smoothing $K$ to remove the bias, and keeping $PV$ in FP16 with a faster accumulator — SageAttention achieves accuracy comparable to FP16 on every model tested. The contradiction is resolved: attention wasn't inherently unquantizable; we were using the wrong quantization for each sub-operation.

This finding redirects research attention in two ways. First, it makes quantization-aware architecture design more attractive: if attention can be quantized cheaply, the bottleneck shifts to other components (linear layers, normalization, activation functions), and the community may invest more in quantizing those holistically rather than treating attention as a lost cause. Second, it makes hardware-specific kernel optimization for consumer GPUs a more competitive research direction relative to architecture-specific solutions (sparse attention, linear attention) or hardware-exclusive solutions (FlashAttention3 for Hopper). The paper's demonstration that a Triton kernel can achieve 52% of theoretical INT8 peak throughput on RTX4090 suggests there is still substantial headroom in software optimization for consumer hardware — a research direction that had been overshadowed by the focus on datacenter GPUs.

Perhaps most significantly, the paper establishes a diagnostic template for approaching quantization of complex operators: separate the operator into its constituent matrix multiplications, analyze the numerical properties of each Matmul independently, and apply different quantization strategies (including the option of not quantizing at all and using a faster accumulation mode). This template is transferable to other attention variants (cross-attention, multi-query attention, grouped-query attention) and potentially to other building blocks like mixture-of-experts routing or state-space model projections.

Follow-Up Research This Work Enables

Systematic characterization of when attention layers develop channel-wise K biases. The paper observes that Llama2 has uniform K distributions and doesn't need smoothing, while Unidiffuser and CogvideoX have strong channel biases that make smoothing essential (Table 1, Figure 4). But why? A follow-up study would train attention-based models across architectures (decoder-only LMs, encoder-decoder, diffusion transformers, vision transformers), modalities (text, image, video, audio), and training objectives (next-token prediction, denoising, contrastive learning), then measure the channel-wise bias in K at each layer. The hypothesis to test is whether the bias emerges from specific architectural choices (e.g., causal masking forces certain channels to encode positional information as a bias) or from training dynamics (e.g., diffusion models may develop biases because attention is applied to noisy inputs with large variance across the batch). If the bias pattern is predictable, practitioners could skip the smoothing overhead on models known to have uniform K; if it's unpredictable, the smoothing should be the default.

Extending the FP16 accumulator insight to other numerically well-conditioned Matmuls. The paper's key empirical finding — that FP16 accumulation for $PV$ is perfectly safe (Tables 4–5, no measurable difference from FP32 even in worst-case) because $P$ is a probability distribution and the computation is a convex combination — suggests a general principle: any Matmul that computes a weighted average with bounded weights can safely use reduced accumulation precision. A follow-up would systematically audit a transformer model for other operations with this property. Candidates include: the attention output projection (which aggregates across heads — each head's output is a convex combination of value vectors), the gating mechanism in mixture-of-experts (a weighted average of expert outputs with softmax weights), and certain normalization layers. For each candidate, measure the cosine similarity between FP16-accumulator and FP32-accumulator outputs across layers and models (Llama2, Unidiffuser, CogvideoX as a minimum), and measure the speedup from switching to FP16 accumulation on consumer GPUs. The goal is a cookbook of safe FP16-accumulator substitutions that collectively provide speedup beyond attention alone.

Combining SageAttention with per-layer sparsity for further speedup on long sequences. The paper shows SageAttention achieves 2× speedup over FlashAttention2 at all sequence lengths, but the absolute latency at 32K tokens (approximately 340 TOPS with headdim=64) is still substantial. For video generation or long-context language modeling, further acceleration is needed. A natural extension is to combine SageAttention's INT8 quantization with sparse attention patterns — for example, computing only the top-k attention scores per query (as in Minference, Jiang et al., 2024) or using a local window plus global tokens (as in Attention Sinks, Xiao et al., 2023b). The INT8 quantization from SageAttention would accelerate the retained attention computations, while the sparsity pattern would reduce the total number of computations. The key experiment: measure the combined speedup on Llama2 at 32K–128K sequence lengths, and quantify whether sparsity-induced attention errors interact with quantization errors synergistically (amplifying each other) or independently (allowing additive speedups). The hypothesis is that SageAttention's smoothing of K might actually improve sparsity pattern selection, since the channel bias that smoothing removes could distort attention score rankings used for top-k selection.

Training-time quantized attention using SageAttention's approach. The paper explicitly targets inference only, but the smoothing K technique and the FP16 accumulator insight apply equally to the forward pass of training. The backward pass through attention requires computing gradients through $QK^\top$ and $PV$, which adds complexity — the INT8 forward pass produces dequantized FP32 outputs, so the backward pass can operate in standard FP16/FP32 without quantized gradients, or the gradients themselves could be quantized. A follow-up would implement SageAttention-style quantized attention in a training loop, measure the wall-clock speedup per training step on RTX4090 for Llama2-scale models at sequence lengths 2K–8K, and measure the effect on training loss curves and downstream task accuracy compared to FP16 training. The key question is whether the quantization noise in the forward pass — which is negligible for inference metrics (Table 8) — compounds over many training steps in ways that degrade final model quality. This would determine whether SageAttention can accelerate not just deployment but also the training of attention-based models on consumer hardware.

Designing a hardware-aware "attention quantization compiler" that selects per-layer strategies automatically. SageAttention's adaptive kernel selection (Section 4.5) manually defines four kernels and a fixed 99.8% cosine similarity threshold. This is a proof of concept. A more ambitious follow-up would build a general framework that, given a model and a target GPU, automatically profiles each attention layer, searches over a menu of quantization options (INT8, FP8, INT4 for Q/K; FP16, INT8, FP8 for P/V; per-token, per-block, per-channel granularities; with and without K smoothing), and selects the fastest configuration that maintains accuracy above a calibrated per-layer threshold. The menu could include options from SageAttention2 (Zhang et al., 2025a), which extends to INT4 quantization with per-thread scales. The key experiment: apply this compiler to a diverse suite of models (Llama2, Llama3, CogvideoX, Sora-like video models, vision transformers) and GPUs (RTX3090, RTX4090, A100, H100), and measure the Pareto frontier of speed vs. accuracy compared to using a single kernel everywhere. This would convert SageAttention from a specific set of kernels into a general methodology for attention acceleration.

Stress-testing SageAttention on extremely long sequences and numerically edge-case inputs. The paper evaluates at sequence lengths up to 32K, which is long but not extreme by current standards (language models now handle 128K–1M tokens). At very long sequences, the smoothing operation (computing mean(K) across all tokens) requires a reduction over the full sequence length, which could become a bottleneck or introduce numerical precision issues if the mean computation itself uses reduced precision. A stress test would run SageAttention at sequence lengths 64K–256K on RTX4090, measure whether the smoothing overhead scales linearly (as expected) or super-linearly (indicating a bottleneck), and check whether the FP16 accumulation for $PV$ remains safe at extreme lengths where the convex combination involves tens of thousands of terms (potentially accumulating rounding error). Additionally, test with adversarial inputs — K matrices engineered to have extreme channel biases (e.g., one channel 1000× larger than others) — to determine whether the smoothing transformation remains stable or whether the mean computation itself loses precision.

Practical Applications and Downstream Use Cases

Real-time video generation on consumer GPUs. The paper shows CogvideoX on RTX4090 running attention at 327.57 TOPS with SageAttention vs. 163.37 TOPS with FlashAttention2 — a 2.01× speedup (Table 7). CogvideoX processes sequences of 17,776 tokens (Section 5.3 shows the Q, K, V shapes include sequence length 17,776). For a video generation pipeline that produces 30-frame clips, reducing attention latency by 2× directly translates to either 2× faster generation time or the ability to generate longer videos within the same latency budget. The paper reports 34.3% end-to-end speedup on CogvideoX (Table 15), which on a 90-second generation task (Figure 1 shows a 90s→67s example on RTX4090) saves roughly 23 seconds per video — meaningful for interactive applications. The plug-and-play nature (no model modification needed) means existing CogvideoX deployments can adopt this immediately.

Cost-efficient LLM prefilling on consumer hardware. For Llama2-7B, SageAttention achieves 1.77× attention speedup (231.74 TOPS vs. 130.99 TOPS for FlashAttention2, Table 7) and preserves perplexity on WikiText to within 0.017% (Table 8). In a deployment where an RTX4090 serves Llama2 for long-context prefilling (processing 32K-token prompts), the attention speedup directly reduces time-to-first-token. With attention consuming approximately 80% of latency at 32K tokens (Figure 2 shows attention at ~800ms vs. linear+other at ~200ms), a 1.77× attention speedup translates to roughly 1.4× end-to-end prefilling speedup. For a service processing thousands of long documents per hour, this reduces the number of GPUs needed or increases throughput on existing hardware. The method is complementary to weight quantization (Table 13 shows SageAttention plus AWQ W4A16 incurs only 0.018% additional perplexity degradation), meaning practitioners can stack attention quantization with existing linear layer quantization for cumulative speedups.

High-resolution image generation with quality preservation. UltraPixel generates images at 2560×1536 resolution, with attention operating on 7,285-token sequences (Table 7). SageAttention accelerates attention by 2.14× (325.18 TOPS vs. 152.03 TOPS for FlashAttention2) while preserving FID to within 0.01% (179.78→179.79, Table 8). For a production image generation service handling thousands of requests per day, the 2× attention speedup directly reduces GPU-hours and allows serving more users with the same hardware. The FID preservation is critical: unlike FlashAttention3's FP8 quantization, which increases Unidiffuser FID from 163.33 to 394.13 (catastrophic quality degradation making images unusable), SageAttention's output is visually indistinguishable (Figure 10, Appendix B.2) and quantitatively nearly identical. This means the speedup comes with effectively zero quality cost — a rare combination in quantization methods.

On-device deployment of vision transformers for classification and retrieval. TIMM's ViT-Base model achieves 5.89× attention speedup with SageAttention over Torch attention (111.41 TOPS vs. 18.91 TOPS, Table 7) and actually improves ImageNet accuracy by 0.05 percentage points on in-distribution data and up to 0.77 points on out-of-distribution benchmarks (Table 8: ImageNet-R 59.55%→60.32%). While the Torch baseline is slower than FlashAttention2, many edge deployment scenarios (mobile GPUs, older server GPUs) lack FlashAttention2 support and rely on Torch's native attention. SageAttention's Triton implementation can run anywhere Triton is supported, making it a practical drop-in acceleration for vision transformers on resource-constrained hardware — for example, accelerating image retrieval in a mobile photo application or enabling real-time classification on a drone's onboard GPU. The counterintuitive accuracy improvement on TIMM suggests that SageAttention's quantization noise may act as a beneficial regularizer for some vision tasks, making it not just faster but potentially more robust.

When to Prefer This Method

The paper explicitly positions SageAttention against FlashAttention2, xformers, and FlashAttention3, creating a clear set of decision criteria. Prefer SageAttention over FlashAttention2 or xformers when:

• You are deploying on consumer GPUs (RTX4090, RTX3090) or pre-Hopper datacenter GPUs (A100). SageAttention achieves 2.1× speedup over FlashAttention2 on RTX4090 and 2.0× on RTX3090 (Figures 6–9). FlashAttention3 is not available on these GPUs, so the choice is between SageAttention's INT8 kernels and FlashAttention2's FP16 kernels. The speedup is largest at long sequence lengths (2.04× at 32K vs. 1.89× at 1K for headdim=64, causal=False on RTX4090), making the case strongest for long-context language models and video generation.

• You need plug-and-play deployment without retraining or calibration. SageAttention requires no fine-tuning, no quantization-aware training, and no per-model calibration beyond the one-time layer selection for adaptive quantization (Section 4.5). Table 8 demonstrates that it preserves end-to-end metrics across language, image, and video models without any model-specific tuning. In contrast, methods like I-BERT (Kim et al., 2021) require quantization-aware training, and task-specific quantizers like Q-diffusion require per-task calibration.

• You are generating images or videos where quality preservation is critical. Table 1 shows that FlashAttention3's FP8 quantization causes catastrophic degradation on Unidiffuser (FID 163.33→394.13) and UltraPixel (FID 179.78→383.61). SageAttention with smoothing K achieves FID within 0.01–2.2% of FP16 on the same models (Table 1, with smoothing enabled). If your use case cannot tolerate noticeable visual quality degradation, SageAttention is the only quantized attention method demonstrated to preserve quality on diffusion-based generation models.

• Your model has non-uniform K distributions with channel-wise outliers. The smoothing K technique provides its largest benefit on models like Unidiffuser and CogvideoX where K exhibits strong channel biases (Figure 4). On these models, per-block INT8 quantization without smoothing produces FID of 229.08 (vs. 163.33 FP16) — essentially unusable — while with smoothing it achieves 166.93 (Table 1). If your model's attention layers show similar patterns (visible in a quick visualization of K channel distributions), SageAttention's smoothing is essential for accurate quantization; if your model has uniform distributions like Llama2, the smoothing adds negligible overhead (under 0.2%, Table 10) and does no harm, so SageAttention remains preferable to FlashAttention2 on speed grounds alone.

Prefer FlashAttention3 over SageAttention when:

• You have access to Hopper GPUs (H100, H800) and need maximum throughput. FlashAttention3 achieves approximately 490 TOPS at headdim=64 on Hopper GPUs, compared to SageAttention's 340 TOPS on RTX4090 — a 44% throughput advantage. If hardware cost is not a constraint and absolute maximum speed is required, FlashAttention3 on Hopper hardware is the faster option. However, the paper demonstrates that FlashAttention3's FP8 quantization degrades accuracy on image/video models (Table 1), so this preference applies primarily to language models where the accuracy degradation is smaller (Llama2 WikiText perplexity 5.823→5.850, a 0.46% degradation that may be acceptable for some applications).

Prefer FP16 FlashAttention2 over SageAttention when:

• You cannot tolerate any numerical deviation from FP16 attention, however small. While SageAttention's end-to-end metric preservation is strong (Table 8: average 0.2% degradation excluding TIMM), the paper does not provide statistical significance tests, and some individual metrics show small degradations (CogvideoX VQA-t: 75.925→75.037, a 1.2% drop; Unidiffuser ImageReward: 0.1609→0.1521, a 5.5% relative drop). If your application is in a regime where these small metric differences matter — safety-critical systems, scientific computing requiring bit-exact reproducibility, or evaluation pipelines where 0.1% accuracy differences determine benchmark rankings — the conservative choice is to accept the 2× slower FP16 attention rather than introduce any quantization at all. The paper does not characterize the worst-case per-example error (only aggregate metrics), so individual outputs could theoretically diverge more than the average suggests.