Source-linked AI summary
Exploiting Sparsity for Long Context Inference: Million Token Contexts on Commodity GPUs
Ryan Synk, Monte Hoover, John Kirchenbauer, Neel Jain, Alex Stein, Manli Shu, Josue Melendez Sanchez, Ramani Duraiswami, Tom Goldstein
TL;DR
Long-context inference incurs O(N^2) prefill cost, repeated O(N) decoding cost, and large KV caches that strain commodity hardware. The method stores keys and values in a CPU vector database and uses approximate nearest-neighbor search to retrieve the top-k keys for each decoding query. Over 95% of dense-attention performance is achieved while attending to only 2% of context tokens on average, with inference demonstrated at the million-token scale.
Problem
Long-context inference incurs O(N^2) prefill cost, repeated O(N) decoding cost, and large KV caches that strain commodity hardware.
Method
The method stores keys and values in a CPU vector database and uses approximate nearest-neighbor search to retrieve the top-k keys for each decoding query.
Results
Over 95% of dense-attention performance is achieved while attending to only 2% of context tokens on average, with inference demonstrated at the million-token scale.
Takeaways & Limitations
Top-k attention provides a tunable compute-performance trade-off for long-context inference on a single commodity GPU.
Takeaways & Limitations
The approach relies on the assumption that modern language models have sparse attention patterns with a small number of tokens carrying most attention mass.
Abstract
from arXiv · showhide
There is growing demand for performing inference with hundreds of thousands of input tokens on trained transformer models. Inference at this extreme scale demands significant computational resources, hindering the application of transformers at long contexts on commodity (i.e not data center scale) hardware. To address the inference time costs associated with running self-attention based transformer language models on long contexts and enable their adoption on widely available hardware, we propose a tunable mechanism that reduces the cost of the forward pass by attending to only the most relevant tokens at every generation step using a top-k selection mechanism. We showcase the efficiency gains afforded by our method by performing inference on context windows up to 1M tokens using approximately 16GB of GPU RAM. Our experiments reveal that models are capable of handling the sparsity induced by the reduced number of keys and values. By attending to less than 2% of input tokens, we achieve over 95% of model performance on common benchmarks (RULER, AlpacaEval, and Open LLM Leaderboard).
1. Introduction
Long-context inference becomes impractical because attention and KV-cache costs grow with context length, especially during repeated decoding. The paper proposes top-k attention over CPU-resident cached states to retrieve only influential keys, enabling sublinear generation on commodity hardware.
- Problem: O(N^2) prefill cost, O(N) decoding cost, and a 52GB KV cache at N = 100,000 make long-context inference prohibitively expensive.The cache estimate assumes D = 4096 and L = 32 for Llama-3 8B.
- Problem: CPU offloading adds prohibitive data movement, while cache eviction can hurt performance by removing vectors needed later.A single layer’s cache can require 1.6GB to transfer repeatedly during generation.
- Approach: Top-k attention stores keys and values in a CPU vector database and retrieves the k highest-scoring keys using approximate nearest-neighbor search.The design separates retrieval from the remaining feed-forward computation, allowing smaller GPU matrix operations while cached states remain in CPU memory.
- Evidence: On roughly 1000-token questions, Figure 2 reports that 10 keys can match full-attention performance on selected OpenLLM Leaderboard tasks.The paper also identifies broader evaluation across models in Figure 6.
- Approach: The method targets sublinear long-context generation by exploiting the observation that only a handful of tokens contribute substantially to attention.The figure describes omitting insignificant query-key interactions to reduce wasted computation.
2. Motivation
The motivation is that modern language models may concentrate most attention mass on a small, relevant subset of tokens. Controlled experiments test this sparsity assumption and find that attention scores identify useful keys for near-accurate inference.
- Attention concentration: Only a small number of keys are needed to collect 75% of attention mass in 4000-token Wikipedia contexts, especially in deeper layers.The analysis uses 50 samples encoded with Llama-3-8B and examines the final token’s attention.
- Attention concentration: Attention focuses more strongly on tokens from the document selected by a copying prompt in multi-article contexts.The comparison averages attention scores across all heads and layers.
- Attention structure: Attention entropy is low across layers and decreases significantly after the first layer, indicating increasingly concentrated score distributions.Entropy is used to measure concentration: low entropy places more attention on fewer tokens, whereas high entropy spreads it more uniformly.
- Performance test: Across Llama-family benchmark tests, performance saturates by the 15th key when each layer uses the same top-k budget.The experiments average performance across selected OpenLLM Leaderboard tasks, and layer-1 tail mass appears unnecessary for good benchmark performance.
- Conclusion: The experiments support the assumption that attention scores reliably identify critical key-value pairs and that accurate inference needs very few keys.The next section uses a vector database to retrieve the most influential tokens.
3. Methodology
The method reduces decoding costs by selecting only the top-k keys and values relevant to each new query, while storing the large context cache on the CPU. It targets O(k) GPU memory instead of O(N), with k around 1% of N for near-equivalent performance.
- Top-k attention: At each decoding step, a nearest-neighbor search selects the k largest query-key scores and returns their corresponding value vectors.The search uses dot-product distance to mirror the attention score mechanism and produces score values plus index positions.
- Index construction: The method builds nearest-neighbor indices from cached keys across layers, using separate indices for each key head in multi-head attention.The index may support exact or approximate nearest-neighbor search.
- Cache placement: After the first generated token, context keys and values remain in a large CPU cache, while generated-token keys and values stay on the GPU for direct attention.The CPU search covers the original context, and GPU attention directly handles previously generated tokens, resembling windowed attention.
- Top-k attention: Top-k decoding reduces peak GPU memory to O(k) instead of O(N) and can use k around 1% of N while recovering near-equivalent performance.The selected context values and scores are moved to the GPU for the final attention computation.
- Cache construction: Prefilling can require O(N^2) memory, so the cache may be constructed using parallelized, approximate, or chunked approaches before top-k decoding.The experiments prefilled 1M-token caches with Flash Attention on an H100 GPU and used chunking to accommodate the memory requirements.
4. Evaluating Top-k Attention at Scale
Across RULER, OpenLLM Leaderboard, and AlpacaEval, top-k attention preserves most dense-attention performance with a small fraction of tokens, including at million-token context length. Performance varies by task, and layer-wise k allocation offers additional flexibility.
- RULER: 95% of RULER baseline performance is reached with k at or below 1% of context length, while approximately 98% is reached using 12.5% at 131k tokens.At k = 2, performance exceeds 60% at every evaluated context length, and scores improve as k increases.
- Open LLM Leaderboard Tasks: Top-k performance saturates below k = 10 on OpenLLM Leaderboard tasks, with similar behavior across pretrained and instruction-tuned models and independence from model size.The evaluation averages MMLU, ARC, HellaSwag, Winogrande, OpenbookQA, BoolQ, and PiQA performance.
- AlpacaEval: 95% of dense-attention performance is achieved with k equal to 2% of context length on AlpacaEval 2.0, regardless of model size.The same small-k behavior holds across model sizes on this generation-intensive benchmark.
- 1M Token Generation with Top-k: 100% success is achieved on one-million-token Needle In A Haystack retrieval with top-k attention, including when the needle appears anywhere in the context.The experiment uses a single GPU and compares top-k attention with cache eviction methods.
- Optimal k Across Task Types: Task requirements differ: most RULER tasks need under 1% of attention scores for 95% performance, but Word Counting requires nearly 9%.Needle In A Haystack needs only a tiny fraction, whereas word-counting tasks are most affected by k.
- Layer-Wise Settings for k: Adaptive layer-wise k allocation can improve RULER performance over uniform allocation while keeping the total k budget fixed.The adaptive strategy increases k linearly from the first to the last layer.
5. Related Work
Prior long-context inference methods reduce memory or distribute computation, but typically assume GPU-resident caches or datacenter-scale resources. This method instead combines sparse top-k retrieval with a nearest-neighbor data structure and extends it to million-token contexts.
- Systems Approaches: Flash Attention and Paged Attention reduce memory through blockwise computation, but Paged Attention assumes the cache fits on the GPU.Flash Attention provides theoretical linear memory complexity, while vLLM’s Paged Attention targets throughput across requests.
- Systems Approaches: Ring Attention scales long-sequence computation across devices, but assumes access to datacenter-level compute.It distributes sequences across devices and overlaps key-value communication with attention computation.
- Cache Eviction: Cache-eviction methods such as sliding-window attention and StreamingLLM save GPU memory by retaining selected tokens, but can fail when relevant information lies outside retained regions.StreamingLLM addresses attention sinks within a modified sliding window.
- Sparse Attention: Earlier top-k methods either compute all attention scores before filtering or use sampling and proxy approximations, limiting their suitability for resource-constrained long-context decoding.The cited nearest-token methods differ in their search structures and approximations.
- Position of This Work: This work uses a nearest-neighbor data structure for top-k retrieval, generalizes across databases, and extends inference to one-million-token contexts and beyond.The authors describe this as the first million-token inference result on a single commodity GPU.
6. Conclusion
The paper demonstrates top-k attention at million-token scale on a single GPU, recovering over 95% of dense-attention accuracy while using about 2% of context tokens on average. Layer- and task-dependent sparsity supports tunable compute–performance trade-offs.
- Conclusion: Over 95% of dense-attention accuracy is achieved on common benchmarks while using only 2% of the context length on average.The method operates at million-token scale on a single GPU with sublinear complexity.
- Conclusion: Attention distributions across layers suggest that top-k budgets can be adapted across tasks and layers to target deployment-specific compute–performance trade-offs.The conclusion frames this flexibility as a direction for future variations of the method.
A.1. Distribution of Attention Scores
Attention sparsity is quantified through entropy, which tracks how many scores are needed to retain dense-attention performance and varies across task categories and layers.
- 1% of attention scores generally preserves 95% of dense-attention performance across Open LLM Leaderboard, AlpacaEval, and RULER.The required threshold varies among subtasks within each benchmark.
- Attention entropy measures the concentration of softmaxed attention distributions, with lower entropy indicating sparser attention.It is computed from attention rows and averaged across generated tokens and text samples.
- Task-level attention entropy correlates strongly with the k required for 95% performance, reaching a Pearson correlation coefficient of 0.85.The comparison uses the first ten generated tokens from fifty samples per task category.
- Attention distributions are analyzed across model layers and task categories, with all 32 layers plotted in order.Figure 10 reports entropy by layer and task category, while Figure 11 organizes layers into rows.
A.2. Additional RULER and AlpacaEval Results
Additional results show that very small top-k budgets approach dense-attention performance across context lengths, models, and benchmark settings.
- Very low k values achieve near-dense performance on RULER across context lengths, matching the behavior observed on Open LLM Leaderboard and AlpacaEval.
- AlpacaEval 2.0 compares Llama instruction-tuned generations and examines how model size affects performance at small k values.