SPIN: Shadow Predictive Indexer for Sparse Attention
Yao Fu, Jiahan Chang, Yueying Li, Ritchie Zhao, Bryce Long, Mohammad Mahdi Kamani, Samkit Jain, Rahul Raman 等 14 位作者
原文摘要
Indexer-based sparse attention reduces the cost of core attention by passing only a fixed, small number of important tokens to it. However, the indexer must still score the entire KV cache at every decoding step. This scoring overhead becomes a major bottleneck as the context length grows. We propose SPIN (Shadow Predictive Indexer) to reduce this indexer overhead. SPIN uses lightweight, history-based prediction to identify important KV blocks, avoiding the need to score the full KV cache at every decoding step. SPIN keeps the score history fresh with random exploration and allocates the block budget across layers by their measured miss rates. SPIN treats KV blocks and speculative decoding as first-class design and implementation considerations. Across extensive evaluations on long-context and agentic benchmarks, SPIN achieves 30-40% sparsity while preserving task quality. With larger exploration shares and per-layer block budgets, SPIN keeps task quality on the especially difficult MRCRv2 benchmark close to baseline at 50-57% effective block sparsity. In end-to-end vLLM serving, SPIN improves output throughput by up to 13.9% and reduces median inter-token latency by up to 13.6%.
cs.LG