KV Cache Compression, But What Must We Give in Return? A Comprehensive Benchmark of Long Context Capable Approaches (2024)

Abstract: Long context capability is a crucial competency for large language models (LLMs) as it mitigates the human struggle to digest long-form texts. This capability enables complex task-solving scenarios such as book summarization, code assistance, and many more tasks that are traditionally manpower-intensive. However, transformer-based LLMs face significant challenges with long context input due to the growing size of the KV cache and the intrinsic complexity of attending to extended inputs; where multiple schools of efficiency-driven approaches — such as KV cache quantization, token dropping, prompt compression, linear-time sequence models, and hybrid architectures — have been proposed to produce efficient yet long context-capable models. Despite these advancements, no existing work has comprehensively benchmarked these methods in a reasonably aligned environment. In this work, we fill this gap by providing a taxonomy of current methods and evaluating 10+ state-of-the-art approaches across seven categories of long context tasks. Our work reveals numerous previously unknown phenomena and offers insights — as well as a friendly workbench — for the future development of long context-capable LLMs. The source code is available at https://github.com/henryzhongsc/longctx_bench.

## introduction large language models (llms) have gained significant popularity and recognition due to their exceptional generalizability across a wide range of intellectual tasks. like any other tool, their most precious utility is demonstrated when they enable us to accomplish tasks beyond our innate capabilities @xcite @xcite . for instance, while driving nails with bare hands is impractical, a hammer makes it feasible. similarly, humans struggle with digesting and retaining long information, making it essential for llms to bridge this gap. the need for longcontext capable llms is almost universally agreed upon, with different llm service providers racing to launch models with even greater context lengths. for example, google's gemini 1.5 supports a context length of 128k tokens @xcite , and anthropic's claude 3 offers a context length of 200k tokens. however, this powerful long context capability comes with significantly higher costs. in long context scenarios, the key-value cache (kv cache)which stores attention keys and values during generation to prevent re-computation -becomes the new memory and speed bottlenecks, as its size grows linearly with the number of tokens in the batch. for instance, a 500b model with a batch size of 128 and a context length of 8,192 typically requires a 3tb kv cache, imposing a substantial processing burden even on the most advanced hardware solutions @xcite . similarly, in open-source models like qwen @xcite , the kv cache size for a 4k context is 0.91 gb, whereas, for a 100k context, it is 22.8 gb @xcite -which is a non-negligible growth regardless of the serving scenario. given the limited memory space available for serving the model, supporting longer contexts usually requires reducing the number of requests that can be processed, leading to higher inference costs. naturally, many efficiency-driven approaches have been proposed to enable llms to handle long contexts with reduced resource burdens, with a healthy selection of them featured in ## reviewing different schools of efficient long context handling approcahes before going into the details of the experiment, we will briefly introduce different schools of long context-capable approaches and their corresponding exemplary methods. in ## benchmarking benchmarking such a variety of methods in a reasonable manner requires significant effort in terms of experiment design, execution, and computational resources. we first introduce the datasets and methods covered, along with the justifications for their selection. then, we detail the experiment setup and explain how to interpret our experiment reports. finally, we analyze the reported results by highlighting some interesting phenomena and providing insights for future scholars. all experiments are conducted on one or more 80g nvidia a100 gpus under dgxa100 servers. ## challenges and opportunities in this section, we share our insights regarding different long context challenges and highlight several opportunities derived from our benchmarking observations. how to effectively reduce prefill time and footprint? based on our empirical observations, most kv cache compression methods struggle to make the prefill stage efficient without compromising performance (ob ❶), which calls for investments in more performant prefill-time compression methods. however, other than the performance requirement on accuracy-like metrics, prefill-time compression methods are entangled with non-trivial technical comparability challenges. recall that flashattention (fa) @xcite is inevitable during the prefill stage to improve hardware utility, with the key spirit of fa being to avoid the generation of a full attention matrix. thus, methods that rely on the availability of a full attention matrix cannot be easily integrated. therefore, we advocate future research on prefilltime compression methods with fa compatibility in mind. how to build efficient yet long context-capable architectures? we empirically observe that pure linear-time sequence models that mix input tokens together struggle with information retrieval (ob ❸), where some sort of attention mechanism provides visible improvements (ob ❷). therefore, an important future direction is to explore how to efficiently combine attention layers with linear-time sequence model layers and determine the optimal number of attention layers needed to achieve an ideal performance-efficiency balance. how to cash-in real-world efficiency? different methods often have varying levels of optimization while being comparable in theoretical efficiency, meaning whether a method is practically efficient in real-world application is highly related to factors like the ease of optimization (e.g., quantization is well-studied and easy to optimize, while some unstructured methods will involve extra challenges @xcite ) and compatibility with established software or hardware frameworks (e.g., compatibility with flashattention, as mentioned above). based on these factors, it is challenging to provide a fair apple-to-apple comparison regarding efficiency. researchers should keep this challenge in mind and develop efficient yet long context-capable methods. ## conclusion our benchmark fills a critical gap by evaluating 10+ methods across 65 settings, uncovering new insights on long context-capable approaches. we also provide a minimalistic and extensible package for reproducible research. 9 4631