PaperlayerView exact revision on arXiv

Source-linked AI summary

Post-Training Sparse Attention with Double Sparsity

Shuo Yang, Ying Sheng, Joseph E. Gonzalez, Ion Stoica, Lianmin Zheng

arXiv:2408.07092v2cs.LGcs.AIcs.CL

TL;DR

LLM inference is memory-intensive, with KV-cache access a key bottleneck, and post-training sparse attention has faced accuracy and runtime challenges. Double Sparsity combines token and channel sparsity, using offline calibration to identify important tokens efficiently. It reports near-lossless performance at 1/16 sparsity, with attention and end-to-end acceleration, while its offloading variant reduces KV-cache memory use.

  • Problem

    KV-cache access is a major inference bottleneck, while post-training sparse attention has struggled to identify important tokens accurately and efficiently at runtime.

    Source: §1 · p. 1

  • Method

    Double Sparsity combines token sparsity with channel sparsity, using offline calibration of relatively static channel patterns to identify important tokens at runtime.

    Source: §1 · p. 2

  • Results

    At 1/16 sparsity, Double Sparsity achieves nearly lossless performance, up to 16× faster attention computations, and 1.9× end-to-end speedup.

    Source: §8 · p. 10

  • Takeaways & Limitations

    Double Sparsity-Offload reduces KV-cache memory usage to 1/16 and outperforms previous state-of-the-art offloading techniques in throughput.

    Source: §8 · p. 10

    Perfectly overlapping communication with computation remains challenging; improving asynchronous capabilities to mask communication overhead is a future direction.

    Source: §8 · p. 10

Abstract

from arXiv · show

The inference process for large language models is slow and memory-intensive, with one of the most critical bottlenecks being excessive Key-Value (KV) cache accesses. This paper introduces "Double Sparsity," a novel post-training sparse attention technique designed to alleviate this bottleneck by reducing KV cache access. Double Sparsity combines token sparsity, which focuses on utilizing only the important tokens for computing self-attention, with channel sparsity, an approach that uses important feature channels for identifying important tokens. Our key insight is that the pattern of channel sparsity is relatively static, allowing us to use offline calibration to make it efficient at runtime, thereby enabling accurate and efficient identification of important tokens. Moreover, this method can be combined with offloading to achieve significant memory usage reduction. Experimental results demonstrate that Double Sparsity can achieve $\frac{1}{16}$ token and channel sparsity with minimal impact on accuracy across various tasks, including wiki-2 perplexity, key-value retrieval, and long context benchmarks with models including Llama-2-7B, Llama-2-70B, and Mixtral-8x7B. It brings up to a 14.1$\times$ acceleration in attention operations and a 1.9$\times$ improvement in end-to-end inference on GPUs. With offloading, it achieves a decoding speed acceleration of 16.3$\times$ compared to state-of-the-art solutions at a sequence length of 256K. Our code is publicly available at https://github.com/andy-yang-1/DoubleSparse.

1 Introduction

The paper targets costly KV-cache access during memory-bound, token-by-token LLM decoding, where post-training sparse attention has struggled to combine accuracy with runtime efficiency. Double Sparsity combines token and channel sparsity, using offline-calibrated channels to identify important tokens at runtime.

1 Introduction

  • KV-cache access is an under-addressed inference bottleneck, and for large batches or long sequences the cache can exceed model-weight size.Token-by-token decoding is memory-bound, with both weights and the KV cache contributing to memory access.Source: §1 · p. 1
  • Post-training sparse attention has struggled to identify important tokens accurately and efficiently without sacrificing practical wall-clock acceleration.Computing and sorting full attention weights identifies important tokens but does not speed up the computation being avoided.Source: §1 · p. 2
  • Double Sparsity pairs important-token attention with channel-based token identification, leveraging relatively static channel patterns and offline calibration.The method also predicts important tokens for the next layer using adjacent-layer embedding similarity, enabling host-memory KV-cache offloading and selective GPU prefetch.Source: §1 · p. 2
  • 1/16 token and channel sparsity incurs negligible accuracy loss across language modeling, question answering, and retrieval; attention speedup reaches 14.1× and end-to-end speedup reaches 1.9×.With offloading, decoding throughput is 16.3× higher than state-of-the-art offloading solutions at a sequence length of 256K.Source: §1 · p. 2

2 Background

The background introduces self-attention notation and frames post-training sparse attention as a way to exploit inherent model sparsity without additional training.

2.1 Preliminaries on Self-Attention and Notations

  • For a decoding step, the query represents the input token, while keys and values represent all tokens; the attention calculation is a major bottleneck at long sequence lengths.The passage attributes the bottleneck to quadratic computational complexity.Source: §2.1 · p. 2

2.2 Post-training Sparse Attention

  • Post-training sparse attention exploits inherent model sparsity, such as token-level sparsity, to accelerate attention calculations without additional training.The passage notes prior examples including H2O, StreamingLLM, and SparQ.Source: §2.2 · p. 3

3 Challenges in Post-Training Sparse Attention

The section identifies three obstacles for post-training sparse attention: preserving retrieval accuracy, achieving hardware-friendly speedups, and reducing memory use. The proposed techniques address these challenges with sparse attention and KV-cache offloading.

3.1 Retrieval Accuracy

  • Discarding tokens can lose critical information and compromise retrieval accuracy; preserving the complete KV cache is an alternative explored by sparse-attention methods.StreamingLLM discards earlier tokens, while H2O drops tokens based on previous attention scores.Source: §3.1 · p. 3

3.2 Hardware Friendliness

  • Non-contiguous channel and token selection can waste GPU bandwidth through cache misses, limiting SparQ to a reported 1.3× attention-computation speed increase.The passage points to continuous memory access and avoiding dynamic channel selection as important hardware-friendliness goals.Source: §3.2 · p. 3

3.3 Memory Usage

  • Preserving the complete KV cache consumes substantial GPU memory, while offloading the whole cache can incur substantial communication overhead.FlexGen offloads each layer's full KV cache to the GPU during computation; transferring only selected tokens is described as requiring less time.Source: §3.3 · p. 3
  • The paper proposes Double Sparsity for attention acceleration and Double Sparsity-Offload for reducing memory use.The passage states up to 16× attention acceleration with minimal additional memory for Double Sparsity, and memory use reduced to 1/16 without increased latency for Double Sparsity-Offload.Source: §3.3 · p. 3

4 Double Sparsity

Double Sparsity retains the full KV cache while using offline-calibrated channels and a compact label cache to make decoding more hardware-friendly. Its channel-selection approach is calibrated offline, and the label cache supports contiguous access to heavy-channel values.

4 Double Sparsity

  • Double Sparsity retains the entire KV cache and uses offline calibration to preselect outlier channels, avoiding runtime sorting and its associated cache misses.A compact label cache stores outlier Key-cache channel values to favor contiguous GPU memory access.Source: §4 · pp. 3–4

4.1 Offline Calibration

  • Offline calibration identifies critical channels by computing their contributions to attention scores on a small validation set.The method builds on the observation that only a few channel contributions have significant impact.Source: §4.1 · p. 4
  • When the channel ratio exceeds 0.25, offline-calibrated outlier indices overlap with online-decoding outliers by 0.95.The passage presents this overlap as evidence of offline calibration's reliability.Source: §4.1 · p. 4

4.2 Forwarding with Label Cache

  • A 4-bit label cache stores heavy-channel Key values contiguously, has 1/16 the size of the K cache, and accelerates decoding 2–4× versus no label cache.All heavy-channel values are stored during prefilling; decoding adds only heavy-channel values from new tokens.Source: §4.2 · p. 5

5 Reducing GPU Memory Usage with Double Sparsity-Offload

Double Sparsity-Offload reduces GPU memory needs by keeping the complete KV cache on the CPU and prefetching selected tokens with double buffering. Its feasibility rests on embedding similarity across most consecutive layers, while the complexity analysis finds IO chiefly depends on attention over selected tokens.

5 Reducing GPU Memory Usage with Double Sparsity-Offload

  • Double Sparsity-Offload reduces GPU memory requirements to 1/16 of the original KV caches, supporting more efficient decoding with limited GPU memory.The section describes this as an extension of Double Sparsity.Source: §5 · p. 5

5.1 Prefetching Tokens with Double Buffer

  • The GPU retains the label cache and a double buffer while the complete KV cache stays on the CPU; approximate next-layer queries drive token prefetching.Prefetch transfers overlap with current-layer attention and feed-forward computation.Source: §5.1 · p. 5

5.2 Empirical Analysis: Embedding Similarity Between Layers

  • On Llama-2-7B Pile validation, consecutive-layer embedding cosine similarity exceeded 90% for all but the second–third and final layer pairs, with most pairs above 95%.The exceptions were the first two layers, the second and third layers, and layers 30 and 31.Source: §5.2 · p. 5

5.3 Complexity Analysis

  • Double Sparsity’s overall IO complexity is approximated by O(2 × β × S × d), so time complexity depends linearly on β and extra memory overhead linearly on α.The approximate-attention phase omits softmax and allows high parallelism; the table compares sparsity methods by overhead, complexity, and speedup.Source: §5.3 · p. 5, Table 1 · p. 6
  • Perplexity changes minimally from sparsity levels 1 to 1/16, while a significant performance gap emerges between 1/16 and 1/32.The table covers models at various sparsity levels.Source: Table 2 · p. 6

6 Experiment

The experiments report robust performance at 1/16 sparsity across benchmarks, with gains in attention and end-to-end inference speed. They evaluate long-context and retrieval accuracy, operator speed, and end-to-end performance under stated hardware and workload settings.

6 Experiment

  • Double Sparsity maintains robust performance at 1/16 sparsity across Wiki-2 perplexity, long-context, retrieval, and other benchmarks.The experiments also include ablations across model sizes, attention mechanisms, and mixture-of-experts configurations; Figure 3 covers long-context benchmarks across sparsity levels.Source: §6 · p. 6, §6.1.1 · p. 6, Figure 3 · p. 7

6.1.1 Wiki-2 Perplexity

  • The Wiki-2 evaluation examines perplexity across sparsity levels, with a noticeable shift as sparsity goes beyond 1/16.Table 3 presents ablations across model configurations at 1/16 sparsity; GQA models are incompatible with K outlier channels.Source: §6.1.1 · p. 6, Table 3 · p. 7

6.1.2 Long Context Benchmarks

  • At 1/16 sparsity, Double Sparsity maintains nearly unchanged accuracy on long-context benchmarks and outperforms StreamingLLM and H2O.The evaluation uses Llama-2-7B across multiple long-context benchmarks.Source: §6.1.2 · p. 6, Figure 3 · p. 7

6.1.3 Key-Value Retrieval

  • Double Sparsity significantly outperforms H2O, StreamingLLM, and RTN quantization on key-value retrieval, while Double Sparsity-Offload shows equivalent performance.The benchmark evaluates in-context retrieval, including accuracy as context length increases.Source: §6.1.3 · p. 7, Figure 4 · p. 8

6.2.1 Setups

  • The setup tests attention and end-to-end inference on A10G and A100 GPUs, including Double Sparsity-Offload at sequence lengths from 64K to 256K.Attention comparisons use scaled_dot_product_attention; end-to-end baselines are gpt-fast and FlexGen Offloading.Source: §6.2.1 · pp. 7–8

6.2.2 Attention Operator Speedup

  • On A10G, every tested case achieves at least 5× speedup over scaled_dot_product_attention, with more than half exceeding 9×.The reported A100 cases are nearly all at least 4× faster, with larger batches reaching up to 10× speedup.Source: §6.2.2 · p. 8, Figure 5 · p. 8

6.2.3 End-to-End Inference Speedup

  • Double Sparsity-Offload reaches 16× speedup over FlexGen at sequence lengths of 64K–256K, while Double Sparsity yields at least 1.3× speedup across tested end-to-end conditions.Under regular workloads, Double Sparsity-Offload achieves 4–8× speedup over FlexGen; Double Sparsity approaches twofold speedup in some scenarios.Source: §6.2.3 · p. 9

7 Related Work

Related work covers post-training sparse attention, attention sparsity methods developed through training, and other attention or inference optimizations. These approaches vary in their sparse patterns, token retention, and targeted inference stage.

Sparse Attention Inference

  • Post-training sparse attention methods include static and dynamic sparse patterns, token eviction, and dynamic decoding methods that preserve all tokens.StreamingLLM and LM-Infinite use static patterns with token eviction; H2O and Scissorhands use dynamic patterns with token eviction, while SparQ and Quest preserve all tokens.Source: §7 · p. 9

Sparse Attention Training

  • Training-based sparse attention reduces attention complexity through sparse factorization, locality-sensitive hashing, or architectures with linear complexity.The cited examples include Sparse Transformer, Reformer, Longformer, BigBard, and Linformer.Source: §7 · p. 9

Other Attention and Inference Optimizations

  • Other attention and inference optimizations include quantization, compression, efficient attention architectures, memory-efficient algorithms, and state space models.The passage also lists batching, memory optimization, parameter sharing, speculative decoding, scheduling, and sparsification.Source: §7 · p. 9

8 Future Directions and Conclusion

The authors identify communication-computation overlap as a remaining challenge and present asynchronous communication masking as a future direction. Their conclusion reports near-lossless performance at 1/16 sparsity and speedups for Double Sparsity and Double Sparsity-Offload.

Future Directions

  • Communication and computation are difficult to overlap perfectly; enhancing asynchronous capabilities to mask communication overhead is a proposed future direction.The authors connect this direction with significant acceleration at minimal memory footprint.Source: §8 · p. 10

Conclusion

  • At 1/16 sparsity, Double Sparsity achieves nearly lossless benchmark performance, up to 16× faster attention, and 1.9× end-to-end speedup; Double Sparsity-Offload reduces KV-cache use to 1/16 and reports 16× throughput over prior offloading techniques.Source: §8 · p. 10

A Offline Calibration Illustration

Figure 7b relates the selected-channel fraction to the overlap between offline and online outliers.

A Offline Calibration Illustration

B.1 Forward without Label Cache

The label cache substantially speeds the forward pass; the accompanying figures report latency comparisons and cross-layer embedding similarity.

B.1 Forward without Label Cache

  • The latency comparison reports a 2–4× forward-pass speedup with label caches.Table 4 compares performance with and without the label cache.Source: Table 4 · p. 14, §B.1 · p. 14

B.2 Embedding Similarity Across Layers

  • Figure 8 plots average cosine similarity across attention heads between adjacent layers on Pile embeddings for Llama-2-7B.It covers layer pairs 0–1, 1–2, 2–3, and 3–4.Source: Figure 8 · p. 14

C Perplexity Selection Illustration

The perplexity illustration supports selecting 1/16 sparsity for both tokens and channels; checklist responses describe the paper’s empirical verification and reporting practices.

C Perplexity Selection Illustration

  • At 1/16 token and channel sparsity, performance remains largely consistent with the original model, while perplexity shifts suddenly beyond 1/16.Figure 9 uses token-level sparsity on the x-axis, channel-level sparsity on the y-axis, and 10-perplexity on the z-axis; higher bars indicate better performance.Source: Figure 9 · p. 15, C Perplexity Selection Illustration · p. 15

NeurIPS Paper Checklist

  • The checklist says the introduction’s claims are empirically verified and that limitations and future directions are discussed in Section 8.The authors identify Sections 6.1 and 6.2 as describing experimental setups and state that code is available.Source: NeurIPS Paper Checklist · pp. 16–17
Loading 2408.07092v2…