PaperlayerView exact revision on arXiv

Source-linked AI summary

MagicPIG: LSH Sampling for Efficient LLM Generation

Zhuoming Chen, Ranajoy Sadhukhan, Zihao Ye, Yang Zhou, Jianyu Zhang, Niklas Nolte, Yuandong Tian, Matthijs Douze, Leon Bottou, Zhihao Jia, Beidi Chen

arXiv:2410.16179v4cs.CLcs.LG

TL;DR

Long-context LLM generation is constrained by the growing KV cache, and TopK attention can degrade on tasks that use information across the full context. MagicPIG approximates attention with LSH sampling and CPU offloading; it maintains high accuracy across diverse tasks while improving throughput and decoding latency.

  • Problem

    Long-context generation is memory-bound because the KV cache grows with batch size and sequence length, while TopK attention can degrade on tasks using the full context.

    Source: §1 · pp. 1–2

  • Method

    MagicPIG uses LSH sampling to approximate attention-score distributions and estimate attention outputs, with hash tables and attention computation offloaded to the CPU.

    Source: §1 · p. 3

  • Results

    MagicPIG maintains high accuracy across diverse tasks and improves serving throughput by 1.5 ∼5× across A100, L20, and RTX 4090, with 54ms decoding latency on RTX 4090 for Llama-3.1-8B-Instruct at 96K context.

    Source: §1 · p. 3

  • Takeaways & Limitations

    The authors report that the LSH and CPU system co-design reduces attention-computation workload while preserving high accuracy across diverse tasks.

    Source: §6 · p. 12

    Oracle sampling cannot exceed 2× wall-clock speedup because obtaining the attention distribution requires full computation of all qkT.

    Source: §4.1 · p. 7

Abstract

from arXiv · show

Large language models (LLMs) with long context windows have gained significant attention. However, the KV cache, stored to avoid re-computation, becomes a bottleneck. Various dynamic sparse or TopK-based attention approximation methods have been proposed to leverage the common insight that attention is sparse. In this paper, we first show that TopK attention itself suffers from quality degradation in certain downstream tasks because attention is not always as sparse as expected. Rather than selecting the keys and values with the highest attention scores, sampling with theoretical guarantees can provide a better estimation for attention output. To make the sampling-based approximation practical in LLM generation, we propose MagicPIG, a heterogeneous system based on Locality Sensitive Hashing (LSH). MagicPIG significantly reduces the workload of attention computation while preserving high accuracy for diverse tasks. MagicPIG stores the LSH hash tables and runs the attention computation on the CPU, which allows it to serve longer contexts and larger batch sizes with high approximation accuracy. MagicPIG can improve decoding throughput by up to $5\times$ across various GPU hardware and achieve 54ms decoding latency on a single RTX 4090 for Llama-3.1-8B-Instruct model with a context of 96k tokens. The code is available at https://github.com/Infini-AI-Lab/MagicPIG.

1 Introduction

Long-context generation is constrained by KV-cache memory and decoding costs, while TopK attention can lose accuracy on tasks that aggregate information across the context. MagicPIG uses LSH sampling and CPU offloading to make attention estimation more practical, reporting higher throughput while maintaining high accuracy across diverse tasks.

1 Introduction

  • Exact TopK attention preserves accuracy on retrieval tasks but degrades severely on aggregated word-extraction tasks, where useful information is distributed more broadly.The paper attributes the contrast to aggregation tasks having less-peaked attention distributions.Source: Figure 1 · p. 1, §1 · p. 1
  • In some layers, the top 20% of attention covers only 70% of total attention scores, showing that attention can be long-tailed rather than sparse.This is reported for tasks that leverage the full context.Source: §1 · p. 2
  • Attention sinks make scores appear sparser, although scores are more uniformly distributed across tokens apart from the sink, whose orientation is nearly input-invariant.The paper’s geometric account places queries near k0 and keys in a narrow cone far from queries.Source: §1 · p. 2, Figure 2 · p. 2
  • Query and key distributions lying in narrow, nearly opposite cones make nearest-neighbor methods perform poorly when searching for TopK keys.The paper describes this mismatch as a reason TopK-key search is intrinsically costly.Source: §1 · p. 2
  • Oracle sampling achieves up to 4× lower estimation error than naive TopK, and MagicPIG approximates it with LSH while offloading hash tables and attention computation to the CPU.The system is designed to support larger batches or longer contexts and larger hashing tables.Source: §1 · pp. 2–3
  • MagicPIG maintains high accuracy across diverse tasks, improves serving throughput by 1.5 ∼5× across A100, L20, and RTX 4090, and reaches 54ms decoding latency on RTX 4090 for Llama-3.1-8B-Instruct at 96K context.The authors also report that MagicPIG outperforms TopK attention on two aggregation tasks.Source: §1 · p. 3

2 Background

The background formulates attention approximation over the KV cache and situates the work among hashing-based attention methods and KV-cache reduction approaches.

2.1 Problem formulation

  • Self-attention forms a weighted average of cached values using softmax-normalized query-key scores, and the formulation seeks a lower-cost approximation using m ≪ n sampled entries.The sampling and diagonal matrices are chosen to minimize the stated approximation objective; TopK instead selects the highest-weight entries.Source: §2.1 · p. 3

2.2 Related works

  • Prior LSH attention methods target training or encoder models, while improved theoretical bounds have not established practical wall-clock acceleration for LLM decoding.FlashAttention, flash-decoding, and SlimAttention losslessly accelerate the attention operator and are described as orthogonal to this approach.Source: §2.2 · p. 4
  • LSH retrieves vectors sharing a hash code with the query in at least one of L independently built tables, with retrieval probability determined by K and L.The background describes SimHash as the cosine-similarity LSH family and gives its retrieval-probability expression.Source: §2.2 · p. 4
  • KV-cache methods reduce memory or loading through token selection, dynamic sparsity, or quantization, addressing distinct dimensions of the cache bottleneck.Examples include H2O, SnapKV, and Keyformer for token preservation; Quest and Loki for dynamic sparsity; and KIVI and QServe for quantization.Source: §2.2 · p. 4

3 Rethinking attention sparsity

The section argues that TopK attention is suboptimal for long-context tasks and examines sampling-based attention estimation as an alternative. Oracle sampling improves estimation and downstream accuracy, but assumes the exact attention distribution is known.

3.1 Achilles’ heel of TopK attention

  • TopK degrades aggregation-task performance despite preserving retrieval accuracy; when attention is long-tailed, the Top20% of tokens covers only 70–80% of attention scores, leaving 15–20% estimation error.The cited Figure 4 reports TopK estimation error for a 16k-token KV cache.Source: §3.1 · pp. 4–5, Figure 4 · p. 4
  • The attention sink is nearly static across inputs while the center of key states is stable and nearly opposite, so TopK can over-weight the sink and lose contextual information.The reported cosine similarities are above 0.99 for sink keys, above 0.9 for center keys, and −0.9 to −0.8 between them.Source: §3.1 · p. 5, Figure 5 · p. 5

3.2 Estimate attention with sampling

  • Up to 4× lower relative error; oracle sampling also preserves high accuracy at a computation cost of 0.002% of a 16k-token context, approximately 32 tokens.Figure 6 compares oracle sampling with TopK estimation and downstream performance.Source: Figure 6 · p. 6, §3.2 · p. 6
  • Oracle sampling is not the lowest-variance estimator and assumes the exact attention vector w is known, which is not true for sparse attention.Source: §3.2 · p. 6
  • Oracle sampling is unbiased, and its covariance trace decreases with budget; its expected unique-sample computation is bounded by 1 + Bϵ, where ϵ = 1 − max_i w_i.Repeated draws need to be computed or loaded only once.Source: §3.2 · p. 6

4 MagicPIG

MagicPIG uses sampling-based estimation to address limitations of TopK attention, then approximates oracle sampling with LSH and a GPU–CPU system design. Its algorithm uses cosine-similarity hashing and partitions projections, sampling, and attention across hardware.

4.1 Self-normalized importance sampling for attention estimation

  • Oracle sampling estimates attention output by sampling from a proposed distribution, but obtaining the true attention distribution requires full query–key computation and caps its speedup at 2×.Self-normalized importance sampling converges to the attention output, while its variance depends on how well the proposal matches the target distribution.Source: §4.1 · p. 7

4.2 Variance reduction with LSH

  • LSH approximates the sampling distribution by transforming query–key inner products into cosine similarities and using a similarity-preserving hash function.The variance-reduction derivation treats the log value-norm term as relatively stable compared with the query–key term.Source: §4.2 · p. 7

4.3 Algorithm Design

  • MagicPIG adapts self-normalized importance sampling to unnormalized hash probabilities, uses SimHash, and samples a key only when it collides with the query in at least two hash tables.It centers keys before hashing because softmax is translation invariant; without centering, keys are poorly distinguished by random projections.Source: §4.3 · p. 8

4.4 System co-design

  • The system runs linear projections and random projections on the GPU while placing hash-table sampling and attention on the CPU.CPU DRAM provides capacity for large hash tables, while sink and local tokens remain on-device to reduce CPU workload.Source: §4.4 · pp. 8–9, Figure 7 · p. 9

5 Evaluation

Across the evaluated tasks, MagicPIG preserves accuracy with low computation cost and improves decoding throughput across three hardware settings. Ablations show that centering is important and sampling can outperform exact TopK attention on two aggregated tasks.

5.1 MagicPIG Preserves Accuracy

  • Across all task types, MagicPIG reports less than 2% accuracy degradation at 2%–5% computation cost; with LSH sampling, it matches or exceeds Quest accuracy at half the computation cost.MagicPIG also demonstrates good performance on moderate-context tasks, indicating more robust performance in general serving.Source: §5.1 · p. 10

5.2 MagicPIG Shows Impressive Efficiency across Various Hardware Settings

  • 1.5× throughput improvement on A100, 5.0× on L20, and 3.3× on RTX 4090 are reported; RTX 4090 reaches 54ms single-request latency at 96K context.With KV-cache offloading, MagicPIG fits over 12× larger batches than GPU full-attention baselines.Source: Figure 8 · p. 11, §5.2 · p. 12

5.3 Ablation Study

  • Without centering, retrieval accuracy falls almost to zero and FWE accuracy to 65%, while fewer than 0.1% of keys are sampled.The authors attribute the sampling failure to the keys’ orientation relative to queries.Source: §5.3 · p. 12
  • On CWE and FWE, MagicPIG beats exact TopK attention by margins of up to 3% and 8%, respectively.The comparison concerns two aggregated tasks where TopK attention experiences significant performance degradation.Source: §5.3 · p. 12, Figure 9 · p. 12

6 Conclusion

The conclusion presents MagicPIG as an LSH-sampling approach paired with CPU offloading to reduce attention workload while retaining accuracy across diverse tasks.

6 Conclusion

  • MagicPIG combines LSH sampling with CPU offloading of hash tables and reduced attention computation to improve throughput and latency across hardware configurations.The authors report high accuracy across diverse tasks and characterize the approach as a co-design of attention approximation and system architecture.Source: §6 · p. 12

A Proofs for theorems

The appendix proof excerpts give a covariance argument, an inequality derivation, and a discussion of the oracle sampling distribution’s trade-offs.

A.1 Proof for Theorem 3.2

  • The proof states that the trace of the covariance matrix decreases monotonically with sample count B.It uses that E[||vX||²] − ||o||² is constant.Source: §A.1 · p. 18

A.2 Proof for Theorem 3.3

  • The proof invokes convexity and Jensen’s inequality, then establishes g(x) ≥ 0 for x in (0, 1) and B ≥ 1 to derive a bound.The displayed bound is expressed in terms of n, ϵ, and B.Source: §A.2 · p. 18

B Oracle sampling

  • The variance-minimizing oracle sampling probability is proportional to attention weight times value norm, but the paper says it is not optimal for downstream accuracy or efficiency.The authors attribute this trade-off to the sink token’s smaller value norm and the flatter resulting distribution, which increases computation cost.Source: B Oracle sampling · p. 18

C Supplementary analysis

The supplementary analysis reports that log |vi −o| is stable during a decoding step, while the attention sink’s value norm is smaller than those of other tokens.

C Supplementary analysis

D Additional evaluation

The additional evaluation covers longer contexts, larger models, reasoning tasks, and system performance; the hyperparameter discussion describes how LSH settings affect sampling, cost, accuracy, and memory.

D.1 Longer Contexts

D.2 Scaling up to larger models

  • MagicPIG preserves high accuracy with low computation for 70B-level models on RULER, with four layers preserved.Source: Table 5 · p. 20, §D.2 · p. 20

D.3 Reasoning

  • MagicPIG consistently outperforms Quest across infini igsm reasoning complexities, while TopK degrades and Oracle Sampling maintains high accuracy.Source: §D.3 · p. 20, Table 6 · p. 21

D.4 System performance

  • For Llama-3.1-8B-Instruct at 96K context, the system evaluation reports decoding latency, maximum throughput, and throughput under a 200 ms latency constraint.The table marks configurations that hit the CPU memory limit.Source: Table 7 · p. 21, §D.4 · p. 21

E.1 (K, L) in LSH

  • In MagicPIG, a key is sampled when its hash collides with the query in at least two tables; K random projections form each hash code across L tables.Source: §E.1 · p. 22
  • Small K can retrieve too many keys, whereas large K lowers collision probability and may require more tables; increasing K and L also raises CPU memory overhead.Source: §E.1 · p. 22

E.2 (K, L) and memory overhead

  • Hash-table memory grows linearly with context length and the model’s number of key-value heads; larger tables also constrain increases to K and L.Source: §E.2 · p. 22

E.3 (K, L) and computation cost/budget

  • Increasing K reduces the estimated sampling budget, while increasing L raises it; measured retrieved-key ratios can differ from estimates because keys are not perfectly distributed.Source: §E.3 · pp. 22–23, Table 9 · p. 23

E.4 (K, L) and accuracy

The section examines accuracy trade-offs from LSH settings and gives a practical, explicitly limited example of how to choose them.

E.4 (K, L) and accuracy

  • Increasing (K, L) can significantly improve accuracy in relatively longer contexts.Source: §E.4 · p. 23, Table 10 · p. 23
  • At ∼2% computation cost, (11, 300) achieves higher RULER accuracy than (10, 150) on MegaBeam-Mistral-7B-512K.Source: Table 10 · p. 23
  • The same (K, L) settings can generalize to larger LLMs; Table 11 covers 8B and 70B models on RULER at 64K.Source: §E.4 · p. 23, Table 11 · p. 23

E.5 How to select (K, L)

  • The authors manually select K=8–10 from small-task ablations, then adjust L for the desired computation budget; under cost below 5% and L below 200, they use (8, 75), (9, 120), and (10, 150).Source: §E.5 · pp. 23–24
  • The proposed tuning procedure is only an example for MagicPIG; advanced hashing methods that may improve the memory–accuracy trade-off remain future work.Source: §E.5 · p. 24

F TopK vs. Sampling

The example contrasts Top-K’s ranking-based estimate with sampling from the full distribution: sampling is unbiased in expectation, while Top-K remains biased.

F TopK vs. Sampling

  • Top-K uses ranking information, whereas sampling accounts for the full attention-score distribution after Softmax.Source: F TopK vs. Sampling · p. 24
  • In the zoo example, the true mean is 8.7 lb; sampling with replacement proportional to animal counts has expected estimate 8.7 lb and standard deviation 4.7 lb at a 10-sample budget.Source: F TopK vs. Sampling · pp. 24–25
  • With 20 samples, the sampling standard deviation falls to 3.4 lb, while Top-K with the same budget yields a biased 17 lb estimate.Source: F TopK vs. Sampling · p. 25

G Limitations and future work

MagicPIG’s CPU-DRAM storage is unsuitable when serving systems lack sufficient DRAM, and the system has not yet been implemented for prefilling. Future work includes improving LSH, CPU-GPU efficiency, GPU attention support, and hyperparameter tuning.

G Limitations and future work

  • MagicPIG’s offloaded KV cache and hash tables require CPU DRAM, making the system unsuitable for serving scenarios with insufficient DRAM.The authors identify this as a limitation of the current system.Source: G Limitations and future work · p. 25
  • MagicPIG has not been implemented in the prefilling stage, which the authors identify as an important direction for long-context LLM serving.Source: G Limitations and future work · p. 25
  • Cross-polytope hashing is proposed as a future direction to reduce hash-table size while improving estimation accuracy.Source: G Limitations and future work · p. 25
  • CPU-GPU pipelines and CPUs’ AVX512 bf16 features are proposed as ways to improve efficiency.Source: G Limitations and future work · p. 25
  • For high-end GPUs with sufficient HBM, the authors identify GPU-friendly LSH algorithms and efficient GPU kernels as requirements for LSH-based attention acceleration.Automatically tuning LSH hyperparameters K and L is also identified as future work.Source: G Limitations and future work · p. 25

Related papers

Loading 2410.16179v4…