Year: 2026
Type: conference-paper
Abstract: 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.
Language: English
Subfield: Artificial Intelligence
Field: Computer Science
Domain: Physical Sciences
Open Access status: hybrid
Funders: Indian Institute of Science, VMware
OpenAlex ID: W7166798098