ChunkAttention: Efficient Self-Attention with Prefix-Aware KV Cache and Two-Phase Partition (2402.15220v4)

Abstract: Self-attention is an essential component of LLMs (LLM) but a significant source of inference latency for long sequences. In multi-tenant LLM 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 state-of-the-art implementation, with the length of the system prompt ranging from 1024 to 4096.

• The paper introduces ChunkAttention, a novel self-attention module that uses a prefix-aware KV cache to reduce redundancy and improve inference efficiency.
• It employs a two-phase partition algorithm that segments computations into chunk-first and sequence-first phases for enhanced data locality.
• Empirical results demonstrate speedups between 3.2 and 4.8 times, underscoring significant efficiency gains for large language models in memory-constrained scenarios.

Efficient Optimization of Self-Attention in LLMs with ChunkAttention

Introduction

With the proliferation of LLMs in multi-tenant serving scenarios, optimizing the inference cost, particularly for self-attention mechanisms dealing with long sequences, has emerged as a critical area of focus. The introduction of ChunkAttention—a novel, prefix-aware self-attention module—marks a significant step towards addressing this challenge. This methodology capitalizes on the observation that many LLM requests share system prompts, thereby allowing for the shared use of key/value (KV) tensors in memory, leading to improved memory utilization and inference efficiency.

Shared System Prompt: A Catalyst for Efficiency

A pivotal observation serving as the foundation for ChunkAttention is the presence of shared system prompts in LLM-based applications, leading to a considerable overlap in context tokens. This redundancy, while traditionally overlooked, represents a valuable opportunity for optimization. By dissecting monolithic KV tensors into smaller chunks organized in an auxiliary prefix tree, ChunkAttention ensures dynamic redundancy removal at runtime without manual intervention. This approach not only saves memory but also allows for the processing of a larger number of sequences simultaneously, thus enhancing throughput in memory-constrained scenarios.

Implementation: The Two-Phase Partition Algorithm

The implementation of ChunkAttention features two core components: a Prefix Aware KV Cache (PAKV) and a Two-phase Partition (TPP) algorithm. The former introduces a scalable, out-of-the-box solution for redundancy elimination through the use of a prefix tree structure for KV cache management. The latter, TPP, seeks to optimize the self-attention computation by dividing it into chunk-first and sequence-first phases. This division allows for the batching of query tensors from sequences with matching prompt prefixes, thus improving data locality and reducing necessary computational operations.

Empirical Validation and Implications

The experimentation with ChunkAttention across varied settings underlines its effectiveness in speeding up the self-attention computation significantly—by factors of 3.2 to 4.8—compared to state-of-the-art implementations. These findings underscore the importance of system prompt design in leveraging shared KV caches for computational efficiency. Moreover, the scalability and adaptive nature of the prefix-tree-based KV cache emerge as potent tools against the exponential growth in context lengths, offering a sustainable path forward as demands for more extensive context understanding grow.

Future Directions in AI and LLM Development

The introduction of ChunkAttention sets the stage for further explorations into memory and compute optimizations in the field of AI and LLMs. As the field evolves, the integration of such efficient algorithms could become standard, pushing the boundaries of what is computationally feasible. Looking ahead, the adoption of ChunkAttention-like methodologies could also spur the development of more sophisticated, context-aware models capable of handling increasingly complex tasks with greater efficiency. Moreover, the foundational principles laid out could inspire novel approaches to tackle the inherent challenges of scaling LLMs, both from a performance and an environmental sustainability perspective.

Indeed, the journey of refining and optimizing the performance of LLMs is far from over. The continued exploration of solutions like ChunkAttention will be pivotal in navigating the complexities of future AI applications. The potential for further optimizations—whether through algorithmic refinement, architectural changes, or hardware advancements—remains vast, promising exciting developments in the quest for ever-more capable and efficient LLMs.