MagicPIG: LSH Sampling for Efficient LLM Generation

2024/10/21 by Zhuoming Chen, Ranajoy Sadhukhan, Chen, Zhuoming +19 · 42 citations
Computer Science · #Computation and Language (cs.CL) #Computer science #FOS: Computer and information sciences #Machine Learning (cs.LG) #Music and Audio Processing #Sampling (signal processing) #Speech Recognition and Synthesis #Telecommunications #Video Analysis and Summarization

paper · pdf · doi:10.48550/arxiv.2410.16179

published in arXiv (Cornell University) (Cornell University)

openalex publication_date 2024/10/21 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

Abstract

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× 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.

Cited by

• OrchANN: Hierarchical Orchestration for Skewed Out-of-Core Vector Search
• LOCKS: Page-Local Compact Key Summaries for Efficient Long-Context Decoding
• KV Admission: Learning What to Write for Efficient Long-Context Inference
• CTkvr: KV Cache Retrieval for Long-Context LLMs via Centroid then Token Indexing
• Hold Onto That Thought: Assessing KV Cache Compression On Reasoning
• ESS: An Offload-Centric Latent-Cache Management Architecture for DeepSeek-V3.2-Exp
• KVReviver: Reversible KV Cache Compression with Sketch-Based Token Reconstruction
• Test-Time Temporal Sampling for Efficient MLLM Video Understanding
• LiteCache: A Query Similarity-Driven, GPU-Centric KVCache Subsystem for Efficient LLM Inference
• Optimizing Native Sparse Attention with Latent Attention and Local Global Alternating Strategies
• MTraining: Distributed Dynamic Sparse Attention for Efficient Ultra-Long Context Training
• NOSA: Native and Offloadable Sparse Attention
• LouisKV: Efficient KV Cache Retrieval for Long Input-Output Sequences
• vAttention: Verified Sparse Attention
• SparseServe: Unlocking Parallelism for Dynamic Sparse Attention in Long-Context LLM Serving
• GraphKV: Breaking the Static Selection Paradigm with Graph-Based KV Cache Eviction
• Spotlight Attention: Towards Efficient LLM Generation via Non-linear Hashing-based KV Cache Retrieval
• Retrospective Sparse Attention for Efficient Long-Context Generation
• Trainable Dynamic Mask Sparse Attention
• LeanK: Learnable K Cache Channel Pruning for Efficient Decoding
• Efficient Attention Mechanisms for Large Language Models: A Survey
• SeerAttention-R: Sparse Attention Adaptation for Long Reasoning
• Memory Mosaics at scale
• AnTKV: Anchor Token-Aware Sub-Bit Vector Quantization for KV Cache in Large Language Models
• Kinetics: Rethinking Test-Time Scaling Laws
• Rectified Sparse Attention
• HATA: Trainable and Hardware-Efficient Hash-Aware Top-k Attention for Scalable Large Model Inference
• Learn from the Past: Fast Sparse Indexing for Large Language Model Decoding
• SALE : Low-bit Estimation for Efficient Sparse Attention in Long-context LLM Prefilling
• KVzip: Query-Agnostic KV Cache Compression with Context Reconstruction
• RetroInfer: A Vector-Storage Approach for Scalable Long-Context LLM Inference
• FIER: Fine-Grained and Efficient KV Cache Retrieval for Long-context LLM Inference
• Mustafar: Promoting Unstructured Sparsity for KV Cache Pruning in LLM Inference
• TailorKV: A Hybrid Framework for Long-Context Inference via Tailored KV Cache Optimization
• FreeKV: Boosting KV Cache Retrieval for Efficient LLM Inference
• RoPE-Aware Bit Allocation for KV-Cache Quantization
• Efficient Graph-Based Approximate Nearest Neighbor Search Achieving: Low Latency Without Throughput Loss
• Heterogeneous LLM Serving with General-Purpose Processing-Near-Memory for Retrieval-Based Sparse Attention
• HMI: Hierarchical Knowledge Management for Efficient Multi-Tenant Inference in Pretrained Language Models
• The Sparse Frontier: Sparse Attention Trade-offs in Transformer LLMs
• Training-Free Hashing-Based Attention via Binary Principal Components
• KeepKV: Achieving Periodic Lossless KV Cache Compression for Efficient LLM Inference
• AlayaDB: The Data Foundation for Efficient and Effective Long-context LLM Inference

Related

• Progressive Sparse Attention: Algorithm and System Co-design for Efficient Attention in LLM Serving
• MoSKA: Mixture of Shared KV Attention for Efficient Long-Sequence LLM Inference
• ALISA: Accelerating Large Language Model Inference via Sparsity-Aware KV Caching
• HGCA: Hybrid GPU-CPU Attention for Long Context LLM Inference
• Spotlight Attention: Towards Efficient LLM Generation via Non-linear Hashing-based KV Cache Retrieval