C hunk A ttention: Efficient Self-Attention with Prefix-Aware KV Cache and Two-Phase Partition (2024)

Abstract: Self-attention is an essential component of large language models (LLM) but a significant source of inference latency for long sequences. In multi-tenant LLMs serving scenarios, the compute and memory operation cost of self-attention can be optimized by using the probability that multiple LLM requests have shared system prompts in prefixes. In this paper, we introduce ChunkAttention, a prefix-aware self-attention module that can detect matching prompt prefixes across multiple requests and share their key/value tensors in memory at runtime to improve the memory utilization of KV cache. This is achieved by breaking monolithic key/value tensors into smaller chunks and structuring them into the auxiliary prefix tree. Consequently, on top of the prefix-tree based KV cache, we design an efficient self-attention kernel, where a two-phase partition algorithm is implemented to improve the data locality during self-attention computation in the presence of shared system prompts. Experiments show that ChunkAttention can speed up the self-attention kernel by 3.2-4.8 \times compared to the start-of-the-art implementation, with the length of the system prompt ranging from 1024 to 4096.

## introduction over the last few years, large language models (llm) have developed various capabilities, from in-context learning @xcite to chainof-thought reasoning @xcite , and achieved remarkable success in a wide range of natural language processing related tasks @xcite . representive models are the gpt @xcite @xcite , llama @xcite , palm @xcite and @xcite series. following the success of chat-gpt and gpt store, llm-based applications start to surge, and the demand to optimize llm's inference cost has been a new area of research interest @xcite @xcite . the self-attention module, as one of the critical components in llms, has poor performance during inference ( on the other hand, the shared system prompt in llm-based applications leads to redundancy in kv cache @xcite . typically, llms are pre-trained and deployed in a multi-tenant architecture for multiple applications to share. due to the in-context learning abilities of llms, using the system prompt to guide llms with instructions and few-shot examples is a common practice in designing llm-based applications @xcite . the system prompt is shared between multiple requests and can be very long. this can be observed in various llm-based applications, from online chatbots to offline experiments ( § 2.1). for instance, gyudoza (2023) shows that system prompts of various llm-based applications have more than 1k tokens. for a chatgpt-like online chatbot, the system prompt can be as long as 1766 tokens with only 6 plugins activated, and all requests sent to the chatbot share one single system prompt (appendix a). in the chameleon @xcite ) work, 4 system prompts are shared by 4241 queries to run the scienceqa benchmark @xcite , and 7 system prompts are shared by 7685 queries to run the tabmwp benchmark @xcite . an important question is whether we can leverage the sharing characteristic of system prompts to make the self-attention module faster and more memory efficient. to our knowledge, the only related work is a proposal by @xcite , in which the service provider reserves memory for key/value tensors of a set of predefined system prompts from application developers. the proposal has limitations: i) predefined system prompts are static and inflexible in frequent refreshes for largescale deployments since both application developers and the service provider are involved in the operation loop; ii) there is memory waste in case of long system prompts and low hit rate; iii) no work has been done to optimize the self-attention kernel in the presence of shared system prompts. to fill the gap, we propose chunkattention, a novel self-attention module featuring the prefixaware kv cache (pakv) and two-phase partition (tpp). kv cache in chunkattention is a prefix tree built with chunked context tokens and key/value tensors. thus, the kv cache is prefixaware and can dynamically detect and remove redundancy at runtime without human involvement. the kv cache only stores key/value tensors of sequences currently in decoding and has zero memory waste. in addition, the prefix-tree structure provides context for chunkattention to redesign a highly-optimized self-attention kernel with twophase partition: chunk-first phase and sequencefirst phase. query tensors from sequences with matching prompt prefixes are batched together to perform attention with key/value tensors. the main contributions of this paper are as follows: i) we reveal that system prompts can be long ( §2.1), providing opportunities for optimizing selfattention; ii) we propose to use prefix tree to implement kv cache, which is out-of-the-box, scalable and robust in terms of redundancy removal; iii) we implement a two-phase partition algorithm to speed up self-attention kernel on prefix-aware kv cache; iv) we prove the feasibility and empirically quantify the gain self-attention can achieve from shared system prompts under various system configurations. our experiments show that chunkattention can be significantly faster as the length of shared system prompts grows and has no performance degrada- ## experiments the evaluations are conducted at both the selfattention microkernel level and the end-to-end gptstyle model level. the microkernel level evaluations only capture time spent in the self-attention cuda kernel. the side effects of pakv and tpp, e.g., prefix tree operations, are captured in endto-end evaluations. we run all experiments with nvidia a100 gpu (80g) and cuda 11.8. ## conclusion in this paper, we propose chunkattention, a novel self-attention module, to efficiently manage kv cache and accelerate the self-attention kernel for llms inference. we successfully adopt the prefix tree to create a prefix-aware kv cache. it addresses the challenge of detecting and removing redundant kv cache at runtime. we evaluate chunkattention in various configurations and at different levels, proving its feasibility and the side effects can be managed. experiments show that the chunkattention kernel can achieve comparable throughput with sota pagedattention kernel without shared system prompts and can outperform it by 3.2-4.8× with a shared system prompt of 1024 to 4096 tokens on a100 (80g) by applying prefix-aware kv cache and two-phase partition.