Can I read SqueezeAttention: 2D Management of KV-Cache in LLM Inference via Layer-wise Optimal Budget on EtoBox?

SqueezeAttention: 2D Management of KV-Cache in LLM Inference via Layer-wise Optimal Budget by Wang, Zihao; Cui, Bin; Gan, Shaoduo is a scholarly article available to read on EtoBox.

What is SqueezeAttention: 2D Management of KV-Cache in LLM Inference via Layer-wise Optimal Budget about?

Optimizing the Key-Value (KV) cache of the Large Language Model (LLM) has been considered critical to saving the cost of inference. Most of the existing KV-cache compression algorithms attempted to sparsify the sequence of tokens by taking advantage of the different importance of tokens. However, most of these methods treat all layers equally, allocating the same KV budget to each layer. This approach is suboptimal, as some layers may be less sensitive to input tokens yet still receive the same budget as others. In this work, we found that by identifying the importance of attention layers, we could optimize the KV-cache jointly from two dimensions, i.e., sequence-wise and layer-wise. Based on our observations regarding layer-wise importance in inference, we propose SqueezeAttention to precisely optimize the allocation of KV-cache budget among layers on-the-fly and then incorporate three representative sequence-wise algorithms to compress the KV-cache for each layer with its very own budget. Specifically, we first measure each layer's importance by calculating the cosine similarity of the input prompt differences before and after the self-attention layers. Based on this similarity,

Author

Wang, Zihao; Cui, Bin; Gan, Shaoduo

Published

2024

Language

EN