PaperlayerView exact revision on arXiv

Source-linked AI summary

SparQ Attention: Bandwidth-Efficient LLM Inference

Luka Ribar, Ivan Chelombiev, Luke Hudlass-Galley, Charlie Blake, Carlo Luschi, Douglas Orr

arXiv:2312.04985v6cs.LG

TL;DR

Long-sequence and batched LLM generation can be limited by repeatedly fetching the growing KV cache. SparQ Attention approximates which cached tokens matter and selectively fetches them; across tested tasks and models, it achieves substantial compression with little loss in performance.

  • Problem

    For long sequences and batches, each generated token requires fetching the full, growing KV cache, making generation increasingly memory-bandwidth-limited.

    Source: §1 · p. 1

  • Method

    SparQ Attention approximates attention scores using subsets of query and key components to fetch only relevant tokens from the KV cache.

    Source: §1 · p. 1, §8 · p. 10

  • Results

    Across tested tasks and models, compression ratios of 1/2 to 1/8 are achievable with little to no task-performance loss; the paper reports up to 8× compression without substantial accuracy loss.

    Source: §1 · p. 2, §5.2 · p. 7

  • Takeaways & Limitations

    SparQ retains the full KV cache while selecting relevant tokens, and is reported as robust across tasks and models for reducing inference times in unseen settings.

    Source: §8 · p. 10

Abstract

from arXiv · show

The computational difficulties of large language model (LLM) inference remain a significant obstacle to their widespread deployment. The need for many applications to support long input sequences and process them in large batches typically causes token-generation to be bottlenecked by data transfer. For this reason, we introduce SparQ Attention, a technique for increasing the inference throughput of LLMs by utilising memory bandwidth more efficiently within the attention layers, through selective fetching of the cached history. Our proposed technique can be applied directly to off-the-shelf LLMs during inference, without requiring any modification to the pre-training setup or additional fine-tuning. We show that SparQ Attention brings up to 8x savings in attention data transfers without substantial drops in accuracy, by evaluating Llama 2 and 3, Mistral, Gemma and Pythia models on a wide range of downstream tasks.

1. Introduction

Long inputs and batched generation make fetching the growing KV cache a memory-bandwidth bottleneck. SparQ Attention estimates which tokens matter using subsets of query and key components, then selectively fetches cached history; the paper reports substantial transfer savings with little accuracy loss.

1. Introduction

  • Each generated token requires fetching the full KV cache, whose size grows with sequence length and batch size, making long batched generation bandwidth-limited.This transfer requirement limits token-generation speed.Source: §1 · p. 1
  • Because tokens generally attend to only a small part of the sequence, predicting high-attention tokens could reduce transfers by fetching only their key-value pairs.Source: §1 · p. 1
  • SparQ approximates attention scores from subsets of query and key components, then fetches only the most relevant tokens at each generation step.The introduction also reports evaluation on Llama 2 and 3, Mistral, Gemma, and Pythia, with practical benchmarks on IPU and GPU.Source: §1 · pp. 1–2
  • Up to 8× compression without substantial accuracy loss; Figure 1 shows matching SQuAD performance at 1/8–1/4 the dense model’s data transfers.The reported pattern is representative across various models and tasks.Source: Figure 1 · p. 1, §1 · p. 2

2. Background

Transformer generation can be bandwidth-bound when data transfers take longer than computation. As sequence length or batch size grows, attention’s KV-cache transfers can dominate, motivating transfer-efficient attention.

2. Background

  • When arithmetic intensity is below the compute-to-memory rate ratio, execution is limited by memory transfer rather than computation.Arithmetic intensity is defined as arithmetic operations divided by memory transfers.Source: §2 · p. 2
  • Typical transformer sequence-generation workloads are bandwidth-bound, so reducing data transfers is identified as the most effective way to accelerate generation.The text says a broader analysis also finds sequence generation typically bandwidth-bound.Source: §2 · p. 2
  • When ρ ≫ 6/B, such as with large sequence length or batch size, attention dominates data transfers because the full KV cache is transferred at each generation step.Here ρ = S/(g d_m), where g is grouped-query heads per key-value head.Source: §2 · p. 2
  • Because transfers limit performance and attention transfers dominate as sequence length increases, the paper motivates transfer-efficient alternatives to standard attention.Source: §2 · p. 2

3. Approximating Attention

The section develops attention approximations that reduce KV-cache reads by exploiting sparse attention scores and query vectors with heavy-tailed components. It also adds a running-mean correction and calibrates the approximate softmax temperature.

3. Approximating Attention

  • For each autoregressive step, an attention head uses query q with key and value caches K and V, which must be read while current-token k and v are appended.The transfer accounting counts scalar elements to separate cache compression from the number format used to represent the cache.Source: §3 · p. 3
  • Softmax attention scores are sparse enough to select a top-k mask and fetch only corresponding values, but computing those scores still requires reading the full key cache.Source: §3 · p. 3
  • Because value vectors show substantial sequence-axis autocorrelation, a running-mean value vector can correct the approximation with minimal additional overhead.The mean vector is updated and written back at each step.Source: §3 · p. 4, Table 1 · p. 4
  • Heavy-tailed query components motivate selecting the top-r components to approximate scores, so only corresponding key-cache components need to be fetched before selecting key-value pairs.Source: §3 · p. 4
  • With query sparsity, the temperature must balance the √r and √d_h extremes; the chosen temperature supports estimating top-128-score coverage and weighting the approximate output.The relative weight α is the estimated mass of the top-k terms.Source: §3 · p. 4

4. SparQ Attention

SparQ Attention estimates attention scores from selected query and key components, then fetches full key–value vectors only for the highest-scoring positions. A running mean value vector compensates for omitted values, while parameters r and k control the data-transfer and accuracy trade-off.

4. SparQ Attention

  • SparQ Attention selects r query components to estimate scores, fetches full key–value vectors for the top-k positions, then uses a running mean value vector to compensate for omitted values.The mean vector receives leftover attention mass based on approximate score weightings.Source: §4 · p. 4
  • The single-head transfer cost is MSparQ = S r + 2 k dh + 4 dh, accounting for selected key components, top-k keys and values, and additional current-vector transfers.The passage identifies the first term as reading r rows of K and the second as reading top-k columns of K and V.Source: §4 · p. 4
  • Varying r and k trades data transfer against approximation accuracy, with r the main compression control when sequence length S greatly exceeds head dimension dh.The passage states that typically S ≫ dh.Source: §4 · p. 5
  • For grouped-query attention, SparQ sums query magnitudes and approximate scores within each group before selection, and omits mean-value compensation for Llama 3 and Mistral.The authors report better performance for those models without Step 3.Source: §4 · p. 5

5. Experiments

Experiments test SparQ Attention across several model families and long-context tasks, comparing it with cache-selection baselines and examining sequence scaling and design choices. The reported results include robust task performance at compression ratios of 1/2 to 1/8 and sustained performance in extended-sequence tests.

5.1. Setup

  • The evaluation spans eight models from Llama 2 and 3, Mistral, Gemma, and Pythia, and tasks covering question answering, summarisation, language modelling, and text repetition.The main task examples were designed for roughly 1k–2k tokens; SQuAD examples add seven confusion contexts.Source: §5.1 · pp. 5–6, §5.2 · p. 6

5.2. Results

  • Compression ratios of 1/2 to 1/8 are achievable with little to no task-performance loss across tested tasks and models; H2O degrades on some tasks, while LM-Infinite degrades across all.The authors interpret LM-Infinite’s results as evidence that the tasks do not permit simply discarding the long input sequence.Source: §5.2 · p. 7

5.3. Sequence Length Scaling

  • On the extended SQuAD variation, SparQ maintains performance up to 128k sequence length; on needle-in-a-haystack tests, it remains close to dense attention even at high sparsity.The SQuAD variation increases sequence length and difficulty by adding confusion contexts; the needle results are averaged over depths and sequence lengths in Table 3.Source: §5.3 · p. 7, Table 3 · p. 7

5.4. Ablations

  • SparQ retains performance comparable to an exact top-k oracle across a wide range of compression ratios and outperforms random low-rank key-cache projection.The oracle supplies exact top-k keys without transferring data to calculate the selection.Source: §5.4 · p. 7
  • The proposed approximate-softmax temperature performs best among the tested settings, including τ = √dh, τ = √r, and α = 0.The α = 0 case corresponds to the limit as temperature tends to zero.Source: §5.4 · p. 7
  • The authors recommend fixing k = 128 and tuning r to balance data transfer and task performance across models and tasks.Reducing either hyperparameter improves bandwidth efficiency but can harm task performance.Source: §5.4 · p. 8

6. Benchmarking

Microbenchmarks and end-to-end tests assess whether SparQ Attention’s transfer savings translate into wall-clock speedups on IPU, CPU, and GPU. Reported results include a 7.41× IPU microbenchmark speedup in one configuration and CPU speedups at every evaluated sequence length.

6. Benchmarking

  • The benchmarking evaluates isolated attention microbenchmarks as well as end-to-end Transformer generation, complementing a hardware-independent theoretical memory-transfer cost model.The benchmarks include IPU and GPU implementations, while the end-to-end tests use CPU and GPU implementations.Source: §6 · p. 8, §6.1 · p. 8, §6.2 · p. 9

6.1. Microbenchmarks

  • On one IPU at batch size 1 and S = 16384, SparQ (r = 32, k = 128) takes 5.28 ms/query versus 40.4 ms/query for dense attention, a 7.41× speedup.The theoretical speedup is 7.53×; the authors attribute the near-match to attention being strongly memory-bound with remote memory.Source: §6.1 · p. 8
  • The authors report that microbenchmark speedups translate theoretical transfer benefits into substantial wall-clock gains, while noting that small-batch improvements and alternatives to storing K twice need further work.The implementation stores K in two layouts at a 50% extra-memory cost and uses fused gather-then-matmul to avoid writing gather results to memory.Source: §6 · p. 8, §6.1 · p. 9

6.2. End-to-End Performance

  • At the longest evaluated CPU sequence lengths, SparQ Attention achieves a 2.5× speedup over the dense baseline, with speedups at all evaluated lengths.The benchmark used AMD EPYC systems and Llama 2 7B; Figure 9 reports batch size 1 and a compression ratio of 1/8.Source: §6.2 · p. 9, Figure 9 · p. 9
  • GPU end-to-end inference on a single H100 PCIe achieves speedups over dense attention at modest sequence lengths, despite using high-bandwidth memory.Figure 10’s setup uses batch size 1, compression ratio 1/8, and 16-bit model weights and KV cache; its plotted theoretical line is an upper bound.Source: §6.2 · p. 9, Figure 10 · p. 9

7. Related Work

Related work includes efficient-attention schemes, inference-time selective or fixed cache lookup, cache eviction, and sparsity in model weights or activations. These approaches differ in whether they change pre-training architecture, retain selective access to the context, or permanently discard information.

7. Related Work

  • Sparse attention schemes such as Longformer (Beltagy et al., 2020) and BigBird use salient tokens or approximate dense attention maps, while MQA (Shazeer, 2019) and GQA (Ainslie et al., 2023) share KV heads.MQA and GQA are architectural methods that must be implemented during pre-training and can involve task-performance trade-offs or affect model quality and stability.Source: §7 · p. 9
  • FlexGen retrieves values for exact top-k scores but uses the full key cache, limiting transfer reduction to 50%; LM-Infinite (Han et al., 2023) and StreamingLLM (Xiao et al., 2023) use fixed, non-selective cache patterns.The fixed patterns retain recent tokens and a few initial tokens.Source: §7 · p. 9
  • Eviction methods including H2O (Zhang et al., 2023), Scissorhands (Liu et al., 2023a), and FastGen reduce cache size and transfer but permanently remove context tokens.The passage notes that this can cause mistakes on queries seeking less-attended parts of the sequence.Source: §7 · pp. 9–10
  • Weight and activation sparsity methods target parameter transfer in small-batch, short-sequence regimes and can complement sparse attention; separate 4-bit KV-cache compression is also described as complementary to reducing transferred elements.Some activation-sparsity approaches replace activation functions with ReLUs and require additional fine-tuning.Source: §7 · p. 10

8. Conclusion

SparQ Attention modifies attention to access only relevant KV-cache tokens at each generation step, reducing data transfers. The authors report robustness across tasks and models and highlight maintaining the full cache for task performance.

8. Conclusion

  • SparQ Attention accesses only relevant KV-cache tokens at each generation step, reducing data transfers, particularly in long-sequence regimes where inference is often transfer-bound.The method is presented as a technique for pre-trained LLM inference.Source: §8 · p. 10
  • Across numerous tasks and models, SparQ Attention is described as robust and viable for reducing inference times in unseen settings.The conclusion contrasts retaining the full KV cache with strategies that discard input information using heuristics or predefined policies.Source: §8 · p. 10

A. Detailed Results

The detailed-results appendix identifies compression/performance curves across evaluated models and tasks, and heatmaps for a needle-in-a-haystack evaluation. The code appendix shows SparQ’s approximate-score token selection and attention computation.

A. Detailed Results

  • Compression/performance trade-off curves cover all evaluated tasks for Llama 2, Llama 3, Mistral, Gemma, and Pythia model families.The curves include compression reference lines at 1/2×, 1/4×, and 1/8× versus dense; shaded lines show ±1 standard error.Source: A. Detailed Results · p. 14, Figure A1 · p. 14, Figure A2 · p. 15
  • Needle-in-a-haystack heatmaps encode correct answers in blue and incorrect answers in red across sequence lengths and needle depths.The evaluation compares SparQ Attention, H2O, LM-Infinite, and dense on togethercomputer/LLaMA-2-7B-32K at two compression ratios.Source: Figure A4 · p. 17

B. Code

  • The code approximates attention scores from the r largest query components, selects top-k positions, gathers their keys and values, then interpolates with V_mean.The interpolation weight alpha is the sum of estimated scores at selected positions.Source: B. Code · p. 18

C. Arithmetic Intensity

The appendix analyzes arithmetic intensity and the share of generation time spent in attention. It connects cache-to-model size with bandwidth limits and describes a sequence-length-based estimate using llama-bench measurements.

C. Arithmetic Intensity

  • As batch size grows, A/M approaches N/C + g, leaving the KV-cache size per item relative to model size as the limiting factor for large-batch inference.The full-layer analysis defines N as parameters, C as KV-cache elements per batch item, and g as grouped-query heads per key-value head.Source: C. Arithmetic Intensity · p. 18
  • With ρ = S/(g d_m), arithmetic intensity approaches g + 6/ρ as batch size increases; ρ determines when the model becomes memory-bandwidth bound.Figure C1 illustrates the relationship between ρ and arithmetic intensity for multi-head attention and grouped-query attention.Source: C. Arithmetic Intensity · pp. 18–19, Figure C1 · p. 19
  • Large-batch arithmetic intensity is just 7 for multi-head attention with S = d_m, and up to 100 for grouped-query attention with S ≪ d_m, while ML hardware can provide r_A/r_M > 200.These comparisons support the appendix’s conclusion that sequence generation is typically bandwidth-bound.Source: C. Arithmetic Intensity · p. 19

D. Measuring Time Spent in Attention

  • Generating a single token takes longer as the existing sequence length increases, and the measured time scales linearly with sequence length.The estimate uses llama-bench measurements on CPU and A100 GPU; the authors note that limited tooling and inference probing can complicate exact attention-time measurements.Source: D. Measuring Time Spent in Attention · pp. 19–20, Figure D1 · p. 20

E. Attention Sparsity Analysis

SparQ approximates attention rankings from sparse query components and partial key data, aiming to identify high-scoring tokens without fetching the entire cache. Query statistics and top-k comparisons suggest that this approximation can preserve useful attention rankings, though some later layers are harder to predict.

E. Attention Sparsity Analysis

  • Up to 20× more query elements than a normal distribution fall beyond absolute z-score 3.0, motivating magnitude-based sparsity in SparQ’s first stage.The authors theorise that pretrained queries encode information sparsely in their tails, so this sparsity does not significantly harm attention-mapping approximation.Source: E. Attention Sparsity Analysis · p. 20
  • Figure E1 reports heavy-tailed query-component kurtosis and outlier-heavy z-value mass across Llama 2 7B heads, measured on 40 SQuAD examples.The z-value mass is normalized against the Gaussian reference of 0.3%.Source: Figure E1 · p. 21
  • Query sparsity of 50% and 75% maintains high top-k correspondence across various k values.Top-k correspondence measures how much the approximate and exact attention-score top-k sets overlap.Source: E. Attention Sparsity Analysis · p. 21, Figure E2 · p. 22
  • Because dropping value positions alone saves at most half of transfers, SparQ approximates scores from partial keys to avoid fetching the full key matrix.It first transfers some key components for all positions, then all key components for selected positions.Source: E. Attention Sparsity Analysis · p. 21
  • Using the r largest-magnitude query components outperforms selecting the first r components, with reasonably high overlap possible at r = dh/8.Some later layers are harder to predict.Source: E. Attention Sparsity Analysis · p. 21

F. Benchmarking Detail

Benchmarking results show that SparQ speedups depend on device and input size: IPU speedups occur across a range of sizes, while GPU speedups require sufficient input size. The appendix also describes benchmark settings and provides illustrative task examples, including a WikiText-103 result.

F. Benchmarking Detail

  • On IPU, SparQ speedups occur across input sizes and depend principally on r; on GPU, sufficient input size is needed, with the bandwidth-limited A10G reaching speedups sooner.The authors suggest kernel fusion could improve small-input performance; Figure F1 spans batch sizes 1–64, sequence lengths 1024–65536, and k values 64–512.Source: F. Benchmarking Detail · p. 22, Figure F1 · p. 23

G. Methodology

  • The H2O baseline uses a fixed cache size k rather than a fixed sequence-length proportion, and the authors’ implementation often shared errors with the original.For Pythia-1.4B on SQuAD 1-shot at k = 256 and l = 64, the two implementations were correct on 60 and 57 of 200 examples, respectively.Source: G. Methodology · p. 23

G.1. Examples

  • The examples section presents one prompt and a cherry-picked output per task, using dense Llama 2 13B, SparQ, H2O, and LM-Infinite outputs.The displayed SparQ setting is r = 8, k = 128.Source: §G.1 · p. 25

G.1.1. QUESTION ANSWERING (SQUAD 1-SHOT)

  • In the displayed SQuAD example, H2O answers “1979” and LM-Infinite answers “1975” to the question about the theater’s opening year.The example uses a prompt augmented with unrelated contexts.Source: §G.1.1 · p. 25

G.1.2. QUESTION ANSWERING (TRIVIAQA 0-SHOT)

  • In the displayed TriviaQA example, dense attention and SparQ both answer “Dubonnet,” while LM-Infinite answers “Byrrh.”The example asks which aperitif was named for the Paris chemist who created it in 1846.Source: §G.1.2 · p. 25

G.1.3. SUMMARISATION (CNN/DAILYMAIL)

  • In the displayed CNN/DailyMail example, dense attention and SparQ begin with the same summary text, while LM-Infinite’s excerpt differs.The passages show truncated output excerpts rather than aggregate task scores.Source: §G.1.3 · p. 26

G.1.4. REPETITION (SHAKESPEARE)

  • In the displayed Shakespeare repetition example, dense attention and SparQ produce the same excerpt, while H2O and LM-Infinite excerpts differ.The excerpts are drawn from the repetition task’s Shakespeare text.Source: §G.1.4 · p. 26

G.1.5. LANGUAGE MODELLING (WIKITEXT-103)

  • Displayed results are 0.669 for dense attention and 0.673 for SparQ; H2O is 0.685, and LM-Infinite is 0.692.Source: §G.1.5 · p. 27
Loading 2312.04985v6…