BaKlaVa -- Budgeted Allocation of KV cache for Long-context Inference

2025-02-18 · Ahmed Burak Gulhan, Krishna Teja Chitty-Venkata, Murali Emani, Mahmut Kandemir, Venkatram Vishwanath

In Large Language Model (LLM) inference, Key-Value (KV) caches (KV-caches) are essential for reducing time complexity. However, they result in a linear increase in GPU memory as the context length grows. While recent work explores KV-cache eviction and compression policies to reduce memory usage, they often consider uniform KV-caches across all attention heads, leading to suboptimal performance. We introduce BaKlaVa, a method to allocate optimal memory for individual KV-caches across the model by estimating the importance of each KV-cache. Our empirical analysis demonstrates that not all KV-caches are equally critical for LLM performance. Using a one-time profiling approach, BaKlaVa assigns optimal memory budgets to each KV-cache. We evaluated our method on LLaMA-3-8B, and Qwen2.5-7B models, achieving up to a 70\% compression ratio while keeping baseline performance and delivering up to an order-of-magnitude accuracy improvement at higher compression levels.

📄 PDF Abstract BibTeX arXiv:2502.13176

Code (0)

등록된 구현이 없습니다.

Tasks

GPULanguage ModelingLanguage ModellingLarge Language Model

Methods 이 논문이 사용한 방법론

Softmax The Softmax output function transforms a previous layer's output into a vector of probabilities. It is commonly used for multiclass classification. Given an input vector $x$…

Attention 설명 없음

MEDA: Dynamic KV Cache Allocation for Efficient Multimodal Long-Context Inference

2025-02-24 · Zhongwei Wan, Hui Shen, Xin Wang, Che Liu 외

Long-context Multimodal Large Language Models (MLLMs) that incorporate long text-image and text-video modalities, demand substantial resources as their multimodal Key-Value (KV) caches grow with increasing input lengths,…

DepthKV: Layer-Dependent KV Cache Pruning for Long-Context LLM Inference

2026-04-27 · Zahra Dehghanighobadi, Asja Fischer arxiv

Long-context reasoning is a critical capability of large language models (LLMs), enabling applications such as long-document understanding, summarization, and code generation. However, efficient autoregressive inference …

Code Generation

Paged Attention Meets FlexAttention: Unlocking Long-Context Efficiency in Deployed Inference

2025-06-08 · Thomas Joshi, Herman Saini, Neil Dhillon, Antoni Viros i Martin 외

Large Language Models (LLMs) encounter severe memory inefficiencies during long-context inference due to conventional handling of key-value (KV) caches. In this work, we introduce a novel integration of PagedAttention wi…

GPU

EvolKV: Evolutionary KV Cache Compression for LLM Inference

2025-09-10 · Bohan Yu, Yekun Chai arxiv

Existing key-value (KV) cache compression methods typically rely on heuristics, such as uniform cache allocation across layers or static eviction policies, however, they ignore the critical interplays among layer-specifi…

Code Completion

DSA: More Efficient Budgeted Pruning via Differentiable Sparsity Allocation

2020-04-05 · ECCV 2020 8 · Xuefei Ning, Tianchen Zhao, Wenshuo Li, Peng Lei 외

Budgeted pruning is the problem of pruning under resource constraints. In budgeted pruning, how to distribute the resources across layers (i.e., sparsity allocation) is the key problem. Traditional methods solve it by di…