EQUIP: EQUivariant preserving In-Place updates for Efficient Token Pruning
Arun Ramachandran, R Govindarajan, Murali Annavaram University of Southern California, Prakash Raghavendra
Token-pruning has emerged as a primary focus in large language models (LLMs) to enhance model efficiency while preserving accuracy, especially for large sequence lengths. However, the eviction operation of token-pruning methods causes “holes” in KV tensors, posing two major challenges: (1) The shift operation, required to make the KV tensor contiguous, results in significant copy overheads; (2) The changes in position indices due to token eviction lead to increased computational requirements for Rotary Positional Encoding (RoPE). To address these issues, we introduce EQUIP, an EQUivariant preserving in-place token update mechanism that ensures the equivariance property of the operations performed in the attention computation. EQUIP offers two fundamental advantages: First, it combines eviction and a subsequent token insertion into an in-place replacement operation, which reduces the KV cache copy overheads significantly. Second, EQUIP reduces recomputation of rotation operations through a combination of in-place update, caching and a re-indexing strategy. Together, these optimizations enable EQUIP to achieve geomean speedups of 1.62× (or 1.47×) on CPU (GPU) over StreamingLLM, and 3.45× (or 1.86×) on CPU (GPU) over Heavy Hitters (H2O). EQUIP with Paged Attention achieves speedups of 4.18×(2.61×) on CPU (GPU) over auto-regressive baselines. EQUIP matches the model accuracy of baseline pruning methods while delivering superior performance.
尚无 AI 解读，展示原文摘要
Proceedings of the 64th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) · 2026 · 正式发表