Research map / Architectures and efficiency / Attention and transformers
MagicPIG: LSH Sampling for Efficient LLM Generation
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes MagicPIG, which replaces TopK attention approximation with LSH-based importance sampling of KV cache entries, offloading hash tables and sparse attention to CPU to speed up long-context LLM decoding.
Problem
Long-context LLM decoding is bottlenecked by the KV cache, and TopK-based sparse attention approximations both degrade accuracy on aggregation tasks and incur high search overhead and no memory savings.
Approach
Treat attention output as an expectation and estimate it by self-normalized importance sampling with SimHash-based LSH sampling probabilities (keys sampled if they collide with the query in at least two of L hash tables), with hash tables and attention computation placed on CPU and projections on GPU.
Key results
- On RULER 96K with Llama-3.1-8B-Instruct, MagicPIG (10,150) reaches 80.0 vs 83.0 full attention at about 2% cost, while Quest (16,0.04) reaches 74.9 at about 10% cost.
- MagicPIG improves throughput by 1.5x (A100, 34B, 16K), 5.0x (L20, 13B, 16K), and 3.3x (simulated RTX 4090, 8B, 96K), achieving 54ms decoding latency for a single 96K request on RTX 4090.
- MagicPIG beats exact TopK attention on CWE and FWE aggregation tasks in RULER 16K by up to 3% and 8%.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use MagicPIG (10,220) instead of full attention | Llama-2-7b-chat · lm-eval-harness comprehensive tasks | Avg. accuracy: 49.1 -> 47.4 | Hurt |
| Use MagicPIG (10,220) instead of Quest (32,0.1) | Llama-3.1-8B-Instruct · lm-eval-harness comprehensive tasks | Avg. accuracy: 59.2 -> 71.2 | Helped |
| Use MagicPIG (8,75) instead of full attention | Llama-3.1-8B-Instruct · LongBench long-context tasks | Avg. accuracy: 71.2 -> 70.7 | Hurt |
| Use MagicPIG (9,120) instead of Quest (16,0.04) | Llama-3.1-8B-Instruct · RULER | Avg. accuracy: 82.1 -> 87.6 | Helped |
| Offload attention and KV cache work to CPU with MagicPIG | LLM decoding · L20, CodeLlama-13B, 16K context | Throughput: 5.0× | Helped |
| Add centering before LSH sampling (ablation) | MagicPIG · RULER retrieval and FWE tasks | Accuracy: Without centering, the accuracy drops to almost zero in retrieval (NIAH) and degrades to 65% in FWE. | Helped |
| Use sampling-based MagicPIG instead of exact TopK attention (ablation) | MagicPIG · RULER CWE | Accuracy: up to 3% | Helped |
| Use sampling-based MagicPIG instead of exact TopK attention (ablation) | MagicPIG · RULER FWE | Accuracy: up to 8% | Helped |
| Increase L from 35 to 240 at 4% computation cost while varying K from 7 to 10 (ablation) | MagicPIG · RULER 16K | Accuracy: 88.5 -> 94.2 | Helped |
| Set K to 8 instead of 7 with L fixed at 120 (ablation) | MagicPIG · RULER 16K | Accuracy: 94.3 -> 94.1 | No clear effect |
Surprising findings
- Exact TopK attention significantly degrades accuracy on aggregation tasks (cwe, fwe) even at moderate density, while sampling outperforms it.
- Attention is not always sparse: in some layers the top 20% of tokens covers only 70% of attention mass, and the attention sink dominates apparent sparsity.
Limitations
- Hash tables and KV cache stored in CPU DRAM, unsuitable where DRAM is limited.
- Not implemented for the prefilling stage.
- LSH hyperparameters (K, L) must be tuned manually and hash tables grow with context length and KV heads.
Experimental setup
Datasets: GSM8K MMLU CoQA QASPER LCC RepoBench-P TriviaQA PassageRetrieval TREC RULER infini-igsm LongBench
Benchmarks: GSM8K MMLU CoQA LongBench RULER infini-igsm
Metrics: accuracy computation cost ratio (FLOPs vs full attention) decoding throughput (tokens/s) time between tokens latency
Baselines: Quest TopK attention Full attention Loki Oracle sampling
Builds on: Quest Reformer ScatterBrain SimHash recursive attention
Closest papers by meaning
- Spotlight Attention: Towards Efficient LLM Generation via Non-linear Hashing-based KV Cache Retrieval (NeurIPS 2025 · similarity 0.91)
- ShadowKV: KV Cache in Shadows for High-Throughput Long-Context LLM Inference (ICML 2025 · similarity 0.90)
- RocketKV: Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression (ICML 2025 · similarity 0.90)
- Efficient Low Rank Attention for Long-Context Inference in Large Language Models (NeurIPS 2025 · similarity 0.89)
- Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs (ICLR 2024 · similarity 0.89)
- DeFT: Decoding with Flash Tree-attention for Efficient Tree-structured LLM Inference (ICLR 2025 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.