Multi-head attention (MHA) is a core component of the transformer architecture, with each transformer model consisting of $L$ stacked layers. Each layer comprises two submodules: a multi-head attention mechanism (MHA) and a feed-forward network (FFN). Given input token embeddings $\boldsymbol{X}\in\mathbb{R}^{n\times d}$, the MHA computes attention in parallel across $H$ attention heads, defined as:

where $\boldsymbol{Q}^{(h)}=\boldsymbol{X}\boldsymbol{W}_{q_{h}},\boldsymbol{K}^{(h)}=\boldsymbol{X}\boldsymbol{W}_{k_{h}},\boldsymbol{V}^{(h)}=\boldsymbol{X}\boldsymbol{W}_{v_{h}}$ are the Query, Key, and Value matrices for head $h$. The projection matrices $\boldsymbol{W}_{q_{h}},\boldsymbol{W}_{k_{h}},\boldsymbol{W}_{v_{h}}\in\mathbb{R}^{d\times d_{H}}$ map the input embedding $\boldsymbol{X}$ to the respective attention components, and $d_{H}$ is typically set to $d/H$ to balance the dimensionality across heads.

Prefill and decoding. During generation, let the model generate $n_{g}$ tokens. In the initial step, the input tokens $\boldsymbol{X}_{0}\in\mathbb{R}^{n\times d}$ are prefilled, and the Keys and Values ($\boldsymbol{K},\boldsymbol{V}$) for each head and each layer are cached for future use. This prefill stage results in the KV caches: $\boldsymbol{K}_{0}=\textrm{Concat}(\boldsymbol{K}^{(1)},\dots,\boldsymbol{K}^{(H)})$ and $\boldsymbol{V}_{0}=\textrm{Concat}(\boldsymbol{V}^{(1)},\dots,\boldsymbol{V}^{(H)})$, where $\boldsymbol{K}_{0},\boldsymbol{V}_{0}\in\mathbb{R}^{n\times d}$.

At each step $t$ ($1\leq t\leq n_{g}$) of autoregressive decoding, the model generates a new token $\boldsymbol{x}_{t}$, conditioned on the input and previously generated tokens. For the subsequent steps, multi-head attention (MHA) only computes the Query/Key/Value vectors ($\boldsymbol{q}_{t},\boldsymbol{k}_{t},\boldsymbol{v}_{t}\in\mathbb{R}^{d}$) for the newly generated token $\boldsymbol{x}_{t}$. These vectors are then appended to the KV cache: $\boldsymbol{K}_{t}=\boldsymbol{K}_{t-1}|\boldsymbol{k}_{t}$ and $\boldsymbol{V}_{t}=\boldsymbol{V}_{t-1}|\boldsymbol{v}_{t}$. The attention mechanism is then performed between $\boldsymbol{q}_{t}$ and the updated KV cache, i.e., $\boldsymbol{K}_{t}$ and $\boldsymbol{V}_{t}$.

The query matrix $Q\in R^{N_{q}\times d}$ and key and value matrices $K$, $V$ $\in R^{N_{k}\times d}$ are inputs to the following equation which is computed independently for each head and batch instance. The output matrix $H\in R^{N_{q}\times d}$ is obtained in essentially three steps as shown in Equation 2.

These intermediate activations $S$ and $P$ are large and thus place a significant demand on memory bandwidth. FlashAttention (Dao et al., 2022) addresses this by chunking the Query, Key, and Value matrices along the token dimension and tiling the attention mechanism to accelerate the process as described below.

Flash Attention. FlashAttention integrates the online softmax algorithm Milakov & Gimelshein (2018) to fuse the three operations illustrated in Equation 2. It requires only a single pass over an entire row of tokens to compute attention. This method partitions the attention output matrix $H$ into independent output tiles, with each tile’s computation being independent of the others. This approach eliminates the need for intermediate global memory reads and writes of $S$ and $P$ tensors. FlashAttention employs a tiling strategy, where small chunks of $query$, $key$, and $value$ are transferred to the compute cores to first perform “partial attention” for each pair of chunks. This is followed by normalization at the end with the help of additional values (sum and max).

Additionally, optimization techniques such as Ring Attention Liu et al. (2023a), Striped Attention Brandon et al. (2023), and Lean Attention Sanovar et al. (2024) aim to balance computation across GPU resources while using FlashAttention’s core algorithm. These techniques significantly reduce the memory bandwidth overhead during inference, especially for long-context generations. However, the tiled dataflow of FlashAttention’s core algorithm, as shown in Figure 2, involving three activation tensors, makes it unsuitable for direct application of state-of-the-art quantization techniques. The quantization techniques do not address the tiled nature of $query$, $key$, and $value$ tensors, thus, techniques which utilize token-wise or channel-wise quantization scale and min cannot be implemented directly without loss of efficiency. On the other hand, these quantization techniques can be applied ot vanilla attention algorithm, but they have the overhead of storing the large intermediate tensors resulting in high memory bandwidth demand.

Further, to avoid overflow, the exponentiation (in partial-attention calculation) in FlashAttention is performed in FP32 data format (exponentiation in Figure 2(a)), while the matrix multiplications are performed in FP16 using GPU tensor cores. Current generation GPUs are inefficient in handling FP32 and do not make use of tensor cores, resulting in higher latency and inefficiency in the attention operation (FP32 CUDA cores deliver only $\sim$3% of the performance of FP16 Tensor Cores). The usage of FP16 and FP32 also adds pressure to the scratchpad and register file on the compute units of GPU.

Progressive quantization (PQ) uses a combination of INT8 and INT4 representations progressively, combining the computational efficiency of symmetric quantization with the ability of asymmetric quantization to represent values accurately at INT4. An overview of PQ follows.

Quantization maps high-precision floating-point values ($x$) into discrete levels ($Q(x)$), expressed as:

where $s$ is a scaling factor and $z$ is a zero-point. These depend on the quantization type, symmetric (sym.) or asymmetric (asym.):

For two quantized matrices $\hat{A}$ and $\hat{B}$, the low-bit integer matrix multiplication is:

We see that asymmetric quantization introduces computation (the last three terms) that are not necessary in symmetric computation (where $z=0$). In general, asymetric quantization leads to lower error rates but results in higher computational overhead.

Progressive quantization proposes two stage of quantization to solve this problem. First, symmetric quantization is used to build INT8 representations for computational efficiency, avoiding the extra overhead of asymmetric quantization (the three non-zero terms above). Then, for better memory efficiency, 8-bit integers are further compressed to INT4 using asymmetric quantization.

The integer inference formula for PQ is:

where $\hat{Q}(A)$ and $\hat{Q}(B)$ denote:

Here, $s_{a}$ and $s_{b}$ are the (FP16) scales from the first INT8 quantization step. Equations 7 and 8 represent dequantization from asymmetric INT4 quantization into INT8. The scales $s^{\operatorname{int}}$ and zero points $z^{\operatorname{int}}$ are stored in INT8, but the majority of the data $Q(Q(x))$ are stored in INT4.

This approach, proposed in Qserve (Lin et al., 2024) for weight quantization, is more hardware-friendly compared to decompression or other low-bit integer inference methods, such as those in Atom (Zhao et al., 2024), LLM.int8 (Dettmers et al., 2022), or TensorRT-LLM . In our work, we extend the design of PQ to represent KV cache, with the objective of making the dequantization into the attention computations faster while keeping it compatible to attention acceleration mechanisms (such as FlashAttention).

As we saw in 1(a), the attention operation can constitute up to 80% of the overall inference execution time when combined with state-of-the-art quantization techniques for other parts of transformer. This is majorly due to the high-precision execution of attention operation in state-of-the-art techniques, such as FlashAttention. As we saw in subsection 2.2, the usage of FP16 and FP32 not only adds pressure to the scratchpad and register file on the compute units of GPU, but slows down the matrix multiplications and softmax operation inside the attention operation because of the higher precision. For example, the peak performance of FP32 CUDA cores used in FlashAttention’s exponentiation operation is only $\sim$3% of the FP16 tensor performance delivered by Nvidia’s A100-SXM. Thus, there is a need to enable quantized execution of the attention operation to leverage the faster low-precision tensor cores in modern GPUs. However, a naive quantization of the attention operation could lead to a loss of the model’s generative performance (accuracy).

On the other hand, compression techniques such as KIVI (Zirui Liu et al., 2023), GEAR (Kang et al., 2024), and Qserve (Lin et al., 2024) have proposed techniques to compress the KV cache to 4-bit or even 2-bit formats. However, these techniques aim to reduce the memory footprint of KV cache, but rely on time-intensive decompression to FP16 before executing the attention mechanism using acceleration mechanisms such as FlashAttention.

Additionally, these methods frequently employ dynamic reordering of outlier channels or smooth attention Lin et al. (2024) techniques to adjust key-query interactions, aiming to mitigate outliers Dettmers et al. (2022). However, these strategies introduce additional latency, hindering the adoption of integer inference within the attention mechanism. Specifically, the smoothing factor(Xiao et al., 2024b) that transfers outliers from key to query is not applicable in cases where both tensors are quantized activations. Similarly, dynamic reordering requires synchronized reordering of query, key, and value components, reducing its overall effectiveness. Thus, there is a huge need to develop a quantization technique that will enable a quantized execution of the attention operation in a FlashAttention compatible manner and reduce the memory footprint of KV cache while not losing the accuracy of LLMs.

TurboAttention thus presents a novel unified technique for enabling quantized execution of attention along with a cooperative KV cache compression mechanism which reduces latency, memory footprint, incurs negligible accuracy loss. We present the in two parts; FlashQ (section 3), which enables quantization of Key, Value, and Query in a FlashAttention compatible manner and SAS, which enables efficient GPU tensor core-friendly, softmax computation (section 4).

As we saw in subsection 2.1, multi-head attention generates large intermediate activations, which place a significant demand on memory bandwidth. FlashAttention (Dao et al., 2022) addresses this by chunking the Query, Key, and Value matrices along the token dimension and tiling the attention mechanism to accelerate the process. FlashAttention divides the $h$-th head’s query, key, and value matrices ($\boldsymbol{Q}^{(h)}$, $\boldsymbol{K}^{(h)}$, and $\boldsymbol{V}^{(h)}$) into sub-blocks of size $T_{c}$ and $T_{r}$, denoted as $\boldsymbol{Q}^{(h)}_{c}$, $\boldsymbol{K}^{(h)}_{r}$, and $\boldsymbol{V}^{(h)}_{r}$, for efficient tiled computation.

To enable a FlashAttention-compatible quantized execution of attention we propose Blockwise Progressive Quantization (BPQ). In BPQ, each of the sub-block is first compressed using 8-bit symmetric quantization. Unlike Qserve, which applies per-channel PQ to weights, we apply BPQ at a sub-block granularity via the function:

where $\mathbf{X}\in{\boldsymbol{Q}^{(h)}_{c},\boldsymbol{K}^{(h)}_{r},\boldsymbol{V}^{(h)}_{r}}$ represents a sub-block, and $\text{Quant}_{\text{sym}}^{8}$ denotes symmetric 8-bit quantization. Given the critical memory bottleneck posed by the KV cache, as described in subsection 2.4, we apply progressive quantization after computation to further compress the Key and Value tensors for storage.

Further, building on the techniques from KIVI (Liu et al., 2024) and informed by our own analysis (Figure 4), we further reduce quantization errors by compressing the Key and Value 8-bit tensors in a channel-wise manner using asymmetric quantization:

where $\hat{K}^{q1}_{g}$ and $\hat{V}^{q1}_{g}$ are groups in each channel of the key and value 8-bit tensors. This minimizes the errors caused by outliers in certain channels while keeping the tiled nature needed to make it compatible with FlashAttention’s core algorithm. Since a new query vector is generated at each decoding step, compressing it further beyond the initial 8-bit quantization is unnecessary.

To further enhance progressive quantization and achieve substantial memory savings in the KV cache, we explore reducing the bit-width to 2-bit. Although 4-bit KV cache quantization has shown near-lossless performance across various models, as demonstrated by Atom (Zhao et al., 2024), QuaRot (Ashkboos et al., 2024), and Qserve (Lin et al., 2024), uniformly applying 2-bit compression to all attention heads often leads to significant model performance degradation. Approaches like smooth factors (Lin et al., 2024; Xiao et al., 2024a) and offline reordering (e.g., Qserve and Atom) can improve compression accuracy but introduce additional latency overhead, making them challenging to integrate with dynamically expanding KV caches. While token-wise mixed precision could further enhance accuracy, it incurs dynamic execution overhead in CUDA kernels and is not FlashAttention-compatible, thus severely reducing hardware efficiency.

Inspired by recent work on head pruning in multi-head attention (Ge et al., 2024; Liu et al., 2023b; Rajput et al., 2024), we propose a headwise mixed-precision approach. This approach combines 2-bit and 4-bit compression, and thus achieves a favorable trade-off between accuracy and memory savings. It applies 2-bit and 4-bit BPQ for different heads based on their priority. The priority is calculated by:

Here, $\text{gap}^{(h)}$ represents the difference between the maximum and minimum values across all channels for head $h$, which captures the range of values in the attention mechanism. A larger gap indicates that compressing this head may introduce a larger quantization error, thus making it more important to preserve precision. $\text{std}^{(h)}$ is the standard deviation of the channel-wise gaps within each head. This measures the variability of the gaps, where higher variability (i.e., a larger standard deviation) implies that the head’s distribution is more uneven. Heads with a higher $\text{std}^{(h)}$ are more sensitive to quantization, and thus require higher bit precision to minimize performance degradation.

Using this metric, we rank all heads in terms of their priority scores. Instead of using a fixed threshold, we compress the lowest-priority $n_{h}$ heads in each layer to 2-bit, while the remaining heads are quantized to 4-bit. The quantization strategy is thus defined as:

Here, $n_{h}$ is the number of heads to be compressed to 2-bit, determined per layer. By ranking the heads based on their priority and applying lower-bit quantization to the least critical heads (i.e., those with the smallest priority), we achieve a balanced trade-off between memory savings and model performance. Figure 3 illustrates the head-wise quantization dataflow of FlashQ, where the second head is block-wise progressively quantized to INT2. As we will see later, this technique presents a robust way to compress the KV-cache and can be employed in a majority of models easily.

FlashQ introduces an efficient decoding buffer to accelerate inference during long-context generation. To support integer inference in the attention mechanism, newly generated tokens are temporarily stored in a buffer $\mathcal{B}$ of size $n_{b}$ (e.g., $n_{b}=64$) and quantized using 8-bit symmetric quantization (Equation 9). This avoids the need for compression at every decoding iteration. Once the buffer reaches its capacity, FlashQ applies progressive quantization (Equation 10) to further compress the keys and values from 8 bits to lower bit-widths for improved memory efficiency.

To ensure stability, we employ a universal scale for 8-bit quantization and clamp outliers that exceed this range, preventing the need to re-compress previously stored tokens in the buffer when new tokens with larger values are added. This approach contrasts with previous methods like KIVI (Liu et al., 2024) and GEAR (Kang et al., 2024), which keep buffers in full precision, introducing significant memory overhead and preventing the use of integer inference in the attention mechanism. We summarize the detailed algorithm of TurboAttention in Appendix A.

Models and dataset We conducted an extensive evaluation of TurboAttention for generative inference using a variety of open-source pre-trained and instruction-tuned models, including LLaMA3-8B-inst (Dubey et al., 2024), Phi3-mini-inst (Abdin et al., 2024), and Qwen2.5-7B (Yang et al., 2024), on multiple generative tasks. These tasks included mathematical reasoning datasets such as GSM8k (Cobbe et al., 2021) and AQuA (Ling et al., 2017), symbolic reasoning tasks from BigBench Hard (BBH) (Suzgun et al., 2022). Given the complexity of these tasks, we use the chain-of-thought prompts created by (Fu et al., 2023) to improve reasoning, which contains 8-shot demonstrations of multi-step reasoning. With the CoT demonstrations, we have the average prefill length of GSM8k, AQuA, and BBH as 900, 1304, 1021 respectively. We then prompt the model to generate 256 tokens and extract answers from them.

To minimize overheads, we implemented and optimized TurboAttention with OpenAI Triton(Tillet et al., 2019) kernel support as follows. First, to address the memory bottlenecks caused by frequent quantization and dequantization, we fused the QKV projection with quantization (Equation 9). Additionally, efficient dequantization from progressive quantization was integrated directly into the attention mechanism to reduce decoding latency (Equation 10).

Second, we implement a KV cache buffer (subsection 3.3 for the newly generated Key and Value vectors during decode. These Key/Value vectors are further compressed to INT4 and INT2 (depending on the head) every $n_{b}$ steps. To avoid recompression when newly generated KV tokens exceed the previous cache value range maximum, we clamp outliers to ensure only newly generated tokens are compressed to INT8. This further improves decoding efficiency by reducing unnecessary recompression.

Third, we leveraged the Triton framework to implement INT8 matrix multiplications for the attention mechanism, thereby reducing both memory and computation bottlenecks. Finally, we implemented an efficient polynomial approximation and lookup table for SAS, significantly enhancing the overall performance of the softmax operation. For TurboAttention, we fix block size $B_{c}$, $B_{r}$, and $n_{b}$ to 64, SAS threshold $n_{r}$ to -6, and set half of the heads’ KV cache to 2-bit quantization for head-wise mixed precision. Further optimization can be achieved by implementing CUDA kernels, as they can deliver higher efficiency and substantial speedup.

Framework. For our experiments, we applied TurboAttention and the baseline methods to models available on Huggingface Wolf et al. (2019), with our inference framework implemented in PyTorch Paszke et al. (2019). Since our focus is on attention acceleration, all other parts of the model are maintained in FP16 unless otherwise stated.

Precision. In our experiments, we explore ultra-low precision quantization techniques, reporting performance results for 4-bit, 3-bit, and mixed headwise 2/4-bit quantization of the KV cache. Our findings indicate that TurboAttention achieves near-lossless performance with 4-bit KV cache compression across various tasks and models. However, consistent with findings in other KV cache compression methods, compressing to 2-bit can led to accuracy degradation. To balance accuracy and memory savings, we applied 2-bit precision to $50\%$ of model heads, achieving a practical trade-off. Consequently, we benchmark TurboAttention with mixed precision quantization against other 3-bit baseline methods for a comprehensive comparison in 2.

Additional ablation studies examining the impact of the number of 2-bit heads and block size variations are provided in subsection 5.6, while the results for pure 2-bit KV cache quantization are included in the Appendix.

We compare TurboAttention against the following state-of-the-art techniques:

$\bullet$ FP16 A dense model where both activations and weights are represented as 16-bit floating-point numbers. This serves as our base dense model without any attention optimizations.

$\bullet$ FlashAttention An attention algorithm applying FlashAttention (Shah et al., 2024) with FP16 computation (and FP32 for exponentiation). This method is "exact" and does not alter accuracy but improves system efficiency by optimizing the attention mechanism.

$\bullet$ KIVI (Liu et al., 2024) A near-lossless KV cache quantization method that compresses the key cache per-channel and the value cache per-token using fine-grained grouping. It stores residual tokens of length $n_{b}$ in full precision, and achieving 4-bit KV cache compression at its best accuracy mode.

$\bullet$ GEAR(Kang et al., 2024) A concurrent KV cache quantization method that achieves 2-bit near-lossless compression by leveraging low-rank approximations to reduce quantization errors. Residual tokens of length $n_{b}$ are stored in full precision for error compensation. Here we use the efficiency version GEAR-L, that deploy quantiations and low rank approximation together to compress KV cache.

It is important to note that TurboAttention is not just a KV cache compression algorithm but rather a unified approximation algorithm for the attention mechanism. However, we compare it with KV cache compression algorithms to highlight its ability to preserve both performance and efficiency in large language models (LLMs).

Generative performance on reasoning tasks. Table 2 demonstrate that TurboAttention achieves comparable or superior performance to recent KV cache compression algorithms across different models and the challenging reasoning tasks, in both 4-bit and lower-bit compression settings. Despite further compressing the query and approximating the softmax operation, TurboAttention maintains near lossless accuracy, achieving an average accuracy of 60.27% across three models and datasets, which is closely aligned with the FP16 baseline average accuracy of 61.89%. Notably, TurboAttention also delivers outstanding 3-bit performance, surpassing baseline methods while providing reduced KV cache size, and enhanced inference efficiency.

In this section, we evaluate latency(wall clock time) and maximum throughputs of TurboAttention on a single A100 GPU(80GB). We measure the attention mechanism of the model under varying batch sizes with a fixed context length of 1k, as well as under different context lengths (ranging from 4k to 32k) with a fixed batch size of 4. We separately assess the prefill and decoding phases and report the speedup ratio compared to the FlashAttention FP16/32 baseline. Figure 6 shows the latency comparison across various input settings and methods, all using the same attention mechanism configuration.

Our results indicate that, across batch sizes from 1 to 64 and context lengths from 4k to 32k, TurboAttention consistently outperforms both FlashAttention FP16 and other KV cache compression techniques combined with Flash Attention baselines. Specifically, TurboAttention achieves a speedup of up to 1.8× for prefill and up to 1.7× for decoding, highlighting significant improvements in softmax, matrix multiplication, and dequantization efficiency. Additionally, by compressing the KV cache to int-4/2, TurboAttention enables long-context generation up to 32k with a batch size of 4, whereas the FP16 baseline encounters out-of-memory (OOM) issues beyond a 4k context length. We also observe that due to the dequantization overhead in other KV cache compression methods, such as KIVI, these methods exhibit worse latency relative to the FP16 baseline.

7(a)shows the throughputs comparison across various batch sizes with fixed context length 1k and generation length 125 tokens. It demonstrats that, compared to the Flash-FP16 baseline, TurboAttention significantly improves maximum throughputs by up to $2.37\times$. Meanwhile, TurboAttention achieves better throughputs compared to KIVI or GEAR due to our efficient progressive dequantization, compressed KV cache, and SAS which improves the latency.

Ablation study on head wise selection methodology In this study, we aim to validate the effectiveness of our priority-based head selection strategy as detailed in subsection 3.2. We compare our approach against several baseline methods that use simpler, less adaptive metrics to drive head selection. As explored in subsection 3.2 and illustrated in Figure 4, channelwise outliers critically impact quantization error. To address this, we examine head selection strategies that incorporate channelwise value ranges as the basis for selection.

Baseline methods include: Entropy: Selecting 2-bit heads based on the entropy of each head. Min-Max: Using the range between minimum and maximum values within heads. Variation: Evaluating variation in channelwise value ranges. The benchmarks are run on the LLaMA-3-8B-inst model with 8-shot COT prompts on the AQua dataset. Given LLaMA-3-8B-inst’s 8-head configuration per key/value cache, we perform head selection across a range of 0 to 16 heads. Results, as shown in 7(b), indicate that our prioritized strategy significantly surpasses other selection approaches in minimizing quantization error, supporting its efficiency and robustness in managing outliers.

Ablation study on block size We evaluate the sensitivity of TurboAttention to varying block sizes ($B_{c}$ and $B_{r}$) for the query, key, and value matrices. These block sizes are closely related to the device’s SRAM capacity and have a substantial impact on system efficiency. Consequently, assessing TurboAttention’s performance across different block sizes is essential to confirm its robustness and adaptability across configurations. As shown in Table 3 , we tested Phi3-mini with various block sizes on the GSM8K-COT dataset using an 8-shot setup. The results demonstrate that TurboAttention maintains robustness across different block sizes, underscoring its ability to integrate seamlessly with Flash Attention across various systems while preserving model accuracy.