ChunkAttention: Efficient Self-Attention with Prefix-Aware KV Cache and Two-Phase Partition

Lu YeZe TaoYong HuangYang Li

article2024ACL78 citations

Proposes ChunkAttention, a self-attention module that uses a prefix-tree structure and a two-phase partition kernel to share key/value caches across requests with matching system prompts, accelerating attention computation by up to 4.8× in multi-tenant LLM serving.

Listen

Serving large language models in multi-tenant environments requires substantial memory and computational resources. The self-attention mechanism, a core computational component of these models, suffers from high latency and memory bottlenecks during real-time token generation because it repeatedly processes large key-value (KV) data caches. At the same time, many user requests share identical system instructions, examples, or tool definitions at the beginning of their prompts. Storing duplicate KV caches for these shared prefixes wastes significant hardware memory and limits overall serving capacity.

The article evaluates ChunkAttention, a self-attention module designed to automatically identify shared prompt prefixes across requests and share their KV cache representations in memory during inference. The authors develop and assess this system to reduce memory consumption and accelerate decoding speeds without requiring manual prompt pre-configuration.

The approach introduces a prefix-tree structure that divides monolithic KV caches into smaller chunks and links shared prefix paths dynamically across active requests. On top of this tree structure, the authors implement a two-phase partition self-attention kernel that batches shared computations across sequences before processing sequence-specific tokens. The authors benchmarked this system against industry-standard baselines, including PagedAttention and FlashAttention, using microkernel tests and end-to-end evaluations on a 7-billion-parameter language model hosted on NVIDIA A100 GPUs under varying batch sizes and arrival rates.

The evaluation yields several key findings. First, ChunkAttention accelerates the self-attention microkernel by 3.2× to 4.8× compared to standard PagedAttention when shared system prompts range from 1,024 to 4,096 tokens. Second, end-to-end serving evaluations demonstrate a 70% to 90% reduction in peak KV cache memory usage when handling long shared prefixes. Third, the system achieves 1.6× to 2.3× higher request throughput while maintaining an average generation latency under 40 milliseconds per token. Finally, when no prompt tokens are shared, ChunkAttention exhibits no performance regression compared to existing baseline systems.

These findings indicate that dynamic prefix-aware caching significantly lowers hardware serving costs and enhances system capacity for applications with repetitive context, such as customer service chatbots, tool-augmented models, and batch benchmark evaluations. The automated runtime detection eliminates operational overhead for infrastructure teams, as developers do not need to manually pre-register static system prompts.

Organizations operating large language model infrastructure should consider adopting dynamic prefix-aware KV caching to improve throughput and reduce memory footprints in shared workloads. Engineering teams should keep shared system instructions at the very beginning of prompts to maximize cache sharing benefits. Further community development is recommended to generalize these low-level GPU kernel optimizations across different model architectures and diverse hardware accelerators.

The reported performance gains depend heavily on shared text appearing at the exact beginning of prompt sequences; any modifications or variations in the initial tokens eliminate prefix matching benefits. Additionally, because the current implementation is heavily tuned for specific GPU architectures and standard attention dimensions, additional testing and optimization are necessary before deploying the solution across non-standard hardware environments.

Microsoft cache inference language model ChunkAttention prefix-aware KV cache efficient self-attention

arXiv: 2402.15220 microsoft/chunk-attention

Paper: Efficient Memory Management for Large Language Model Serving with PagedAttention, Woosuk Kwon et al. (2023). PagedAttention introduces the foundational non-contiguous KV-cache management framework and serving architecture that ChunkAttention directly builds upon and benchmarks against.
Paper: FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning, Tri Dao (2023). FlashAttention-2 provides the core hardware-aware IO-tiling and attention work-partitioning techniques foundational to developing ChunkAttention's two-phase GPU kernel.
Paper: SARATHI: Efficient LLM Inference by Piggybacking Decodes with Chunked Prefills, Amey Agrawal et al. (2023). SARATHI introduces chunked prefill strategies for LLM inference that ChunkAttention builds upon for prefix-aware chunking and computation partitioning.
Paper: Orca: A Distributed Serving System for Transformer-Based Generative Models, Gyeong-In Yu et al. (2022). Orca establishes iteration-level scheduling for generative transformer serving, providing the architectural foundation for modern multi-tenant batch serving systems.

Paper: ThinK: Thinner Key Cache by Query-Driven Pruning, Yuhui Xu et al. (2025). ThinK explores complementary memory reduction by pruning redundancy in the KV cache channel dimension rather than sharing across prompt sequence prefixes.
Paper: PolarQuant: Quantizing KV Caches with Polar Transformation, Insu Han et al. (2025). PolarQuant investigates coordinate-transformation-based KV cache quantization to achieve orthogonal memory compression alongside prefix-sharing serving systems.
Paper: Kascade: A Practical Sparse Attention Method for Long-Context LLM Inference, Dhruv Deshmukh et al. (2025). Kascade develops cross-layer sparse attention to cut attention computation during long-context prefill and decode stages.
Paper: QJL: 1-Bit Quantized JL Transform for KV Cache Quantization with Zero Overhead, Amir Zandieh et al. (2025). QJL presents extreme 1-bit JL-transform quantization for KV caches to tackle inference memory bottlenecks from a numeric precision perspective.
Paper: Language Models Can Control Their Own Attention, Namgyu Ho et al. (2026). Declarative Attention enables models to dynamically scope and mask their own KV cache reads at decode time within serving engines like vLLM.
Paper: XAttention: Block Sparse Attention with Antidiagonal Scoring, Ruyi Xu et al. (2025). XAttention extends long-context inference acceleration through block-sparse attention scoring along antidiagonal patterns.
Paper: Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention, Jingyang Yuan et al. (2025). Native Sparse Attention introduces hardware-aligned, trainable sparse attention kernels to natively alleviate the quadratic scaling of standard KV caches.

Cover for ChunkAttention: Efficient Self-Attention with Prefix-Aware KV Cache and Two-Phase Partition

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 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 run-time 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× compared to the start-of-the-art implementation, with the length of the system prompt ranging from 1024 to 4096.

Topics

cache inference language model large language models multi tenants random access memory ranging probabilistic inference

Keywords

ChunkAttention prefix-aware KV cache efficient self-attention two-phase partition KV cache memory management iteration-based batching large language models data locality

1 Introduction
2 Preliminaries
2.1 Shared System Prompt
2.2 LLMInferencing
3 Our Approach
3.1 Prefix Aware KV Cache (PAKV)
3.2 Two-phase Partition (TPP)
3.3 Further Optimizations
4 Experiments
4.1 Microkernel Evaluation
4.2 End-to-end Evaluation
5 Related Work
6 Conclusion
7 Limitations
References
A System Prompt for Chatbot Applications with Plugins

Knowls

Knowl 1 — Prefix-Aware Key-Value Cache (PAKV) Structure and Dynamic Management
model/method
The Prefix-Aware Key-Value Cache (PAKV) organizes the cached key and value tensors of active inference sequences into a prefix tree (or forest of trees for distinct system prompts). Instead of storing monolithic contiguous tensors of shape $b \times h \times n \times d$ (batch size $b$ , attention heads $h$ , sequence length $n$ , head dimension $d$ ), key and value tensors along the sequence length dimension are sliced into fixed-size chunks of length $c$ .

Each node in the prefix tree corresponds to a chunk $C$ containing three elements:
1. A segment of $c$ context tokens shared by all sequences passing through the node.
2. A physical key tensor slice of size $h \times c \times d$ .
3. A physical value tensor slice of size $h \times c \times d$ .
Every path from the root node to a leaf node uniquely defines an individual sequence. A structural invariant of this tree organization is that all sequences covered by any given node $C$ are contiguous in the batch sequence index dimension $[i, j]$ .

Dynamic memory management is handled via a pool-based memory allocator tracking free and allocated chunks of fixed size $c$ :
- Sequence arrival: A prefix lookup traverses the tree. For matched prefix chunks, existing cached tensors are shared directly; newly computed suffix tokens are allocated chunks and inserted into the tree as a branching path.
- Decoding iteration: Newly generated completion tokens are appended to the active leaf chunk until full, at which point a new leaf chunk is allocated.
- Sequence departure: When generation reaches the end-of-sequence token or maximum completion length, the sequence's path is pruned, and non-shared chunks are returned to the free list without releasing physical OS memory allocations.
Knowl 2 — Two-Phase Partition (TPP) Self-Attention Decoding Algorithm
algorithm
The Two-Phase Partition (TPP) algorithm computes self-attention during autoregressive decoding across a prefix-aware KV cache. For a decoding batch of $b$ sequences, the query matrix is $Q \in \mathbb{R}^{b \times d}$ (representing one query vector per sequence per head). The algorithm partitions the workload into two phases: a chunk-first phase that batches multi-sequence operations over shared prefix chunks using Tensor Cores, and a sequence-first phase that merges running attention statistics and processes sequence-specific unshared suffix chunks.
Input: Query matrix $Q \in \mathbb{R}^{b \times d}$ , prefix tree $T$
Output: Attention output matrix $O \in \mathbb{R}^{b \times d}$
// Phase 1: Chunk-First (Batched Shared Chunks)
Identify all shared chunks $C_1, \dots, C_k$ in $T$ covered by multiple sequences
for each chunk $C \in \{C_1, \dots, C_k\}$ do
    Retrieve key and value cache slices $K^{(C)}, V^{(C)} \in \mathbb{R}^{c \times d}$
    Determine contiguous sequence index range $[i, j]$ covered by chunk $C$
     $(O^{(C)}, m^{(C)}, n^{(C)}) \leftarrow \text{partial\_attn}(Q_{i:j,:}, K^{(C)}, V^{(C)})$
    Save partial attention outputs $(O^{(C)}, m^{(C)}, n^{(C)})$ to global memory
end for
// Phase 2: Sequence-First (Per-Sequence Accumulation and Suffix Chunks)
for each sequence query $q_u = Q_{u,:} \in \mathbb{R}^d$ ( $u \in \{1, \dots, b\}$ ) do
    Initialize sequence accumulators $o_u \leftarrow \mathbf{0} \in \mathbb{R}^d$ , $m_u \leftarrow -\infty$ , $n_u \leftarrow 0$
    for each shared chunk $C \in \{C_1, \dots, C_k\}$ covering sequence $u$ do
        Slice sequence-specific entry: $(o^{(C)}, m^{(C)}, n^{(C)}) \leftarrow (O^{(C)}_{u-i,:}, m^{(C)}_{u-i}, n^{(C)}_{u-i})$
         $(o_u, m_u, n_u) \leftarrow \text{attn\_reduce}(o^{(C)}, m^{(C)}, n^{(C)}, o_u, m_u, n_u)$
    end for
    Retrieve unshared sequence-specific chunks $C_{k+1}, \dots, C_l$ for sequence $u$
    for each chunk $C \in \{C_{k+1}, \dots, C_l\}$ do
        Retrieve key and value cache slices $K^{(C)}, V^{(C)} \in \mathbb{R}^{c \times d}$
         $(o^{(C)}, m^{(C)}, n^{(C)}) \leftarrow \text{partial\_attn}(q_u, K^{(C)}, V^{(C)})$
         $(o_u, m_u, n_u) \leftarrow \text{attn\_reduce}(o^{(C)}, m^{(C)}, n^{(C)}, o_u, m_u, n_u)$
    end for
     $O_{u,:} \leftarrow o_u / n_u$
end for
return $O$
Knowl 3 — Partial Attention and Online Softmax Reduction Equations
equation
In the Two-Phase Partition attention kernel, partial_attn computes independent attention statistics over a key-value chunk $C$ of length $c$ with key slice $K^{(C)} \in \mathbb{R}^{c \times d}$ and value slice $V^{(C)} \in \mathbb{R}^{c \times d}$ . For a query slice $Q_{i:j,:} \in \mathbb{R}^{(j-i) \times d}$ spanning sequence indices $i$ to $j$ :
$\begin{aligned} W^{(C)} &= Q_{i:j,:} (K^{(C)})^T \in \mathbb{R}^{(j-i) \times c} \\ m^{(C)} &= \max_{\text{row}}(W^{(C)}) \in \mathbb{R}^{j-i} \\ E^{(C)} &= \exp\left(W^{(C)} - m^{(C)} \mathbf{1}^T\right) \in \mathbb{R}^{(j-i) \times c} \\ n^{(C)} &= \sum_{\text{row}} E^{(C)} \in \mathbb{R}^{j-i} \\ O^{(C)} &= E^{(C)} V^{(C)} \in \mathbb{R}^{(j-i) \times d} \end{aligned}$
To combine the partial attention result of chunk $C$ for sequence $u$ (denoted $o^{(C)} \in \mathbb{R}^d, m^{(C)} \in \mathbb{R}, n^{(C)} \in \mathbb{R}$ ) with running accumulators $o_u \in \mathbb{R}^d, m_u \in \mathbb{R}, n_u \in \mathbb{R}$ , attn_reduce applies online softmax scaling:
$\begin{aligned} x^{(C)} &= \exp\left(m^{(C)} - \max(m^{(C)}, m_u)\right) \in \mathbb{R} \\ y^{(C)} &= \exp\left(m_u - \max(m^{(C)}, m_u)\right) \in \mathbb{R} \\ o_u &\leftarrow x^{(C)} o^{(C)} + y^{(C)} o_u \in \mathbb{R}^d \\ n_u &\leftarrow x^{(C)} n^{(C)} + y^{(C)} n_u \in \mathbb{R} \\ m_u &\leftarrow \max(m^{(C)}, m_u) \in \mathbb{R} \end{aligned}$
After all covering chunks are processed, the final attention output vector for sequence $u$ is given by element-wise division: $O_{u,:} = o_u / n_u$ .
Knowl 4 — Theoretical Bounds on Batch Capacity Gain and KV Cache Fragmentation
theoretical result
Let a batch of requests have sequence prompt lengths $n_p$ , generated completion lengths $n_c$ , and shared prefix token lengths $n_s$ . The sharing ratio $r$ is defined as:

$r = \frac{n_s}{n_p + n_c}$

Under memory-capacity-limited inference scenarios, sharing the key-value tensors of the common prefix across sequences increases the maximum number of concurrent sequences that can fit into memory by a factor of approximately:

$\text{Capacity Multiplier} \approx \frac{1}{1 - r}$

For fixed chunk size $c$ in the prefix tree and total sequence context length $n$ , memory loss resulting from alignment internal fragmentation in leaf chunks is strictly upper-bounded by:

$\text{Memory Loss} \le \frac{c - 1}{n}$
Knowl 5 — CPU-GPU Context Management and Latency Hiding for Prefix Trees
model/method
Because the prefix tree structure is maintained dynamically in host CPU memory, execution of the two-phase partition kernel on GPU requires dispatching execution context tuples $(C, i, j)$ , where $C$ is the GPU chunk memory address and $[i, j]$ is the sequence index range covered by $C$ . ChunkAttention mitigates host-device synchronization and copy overheads through two mechanisms:
1. Latency hiding: The CPU traversal and context generation routines are overlapped asynchronously with earlier GPU kernel executions in the transformer decoder layer (such as RMSNorm/LayerNorm and QKV projection) prior to the self-attention kernel.
2. Lazy context copy: The prefix tree topology remains constant across consecutive decoding steps that do not modify tree branches. Context metadata is cached in GPU memory and transfers are triggered only upon structural changes: when a sequence's active leaf chunk becomes full (occurring once every $c$ iterations), when a new sequence joins, or when a finished sequence terminates.

Knowl 6 — Self-Attention Microkernel Latency Across Context and Shared Prefix Lengths

data/table

The microkernel latency of self-attention was evaluated on an NVIDIA A100 GPU (80GB, CUDA 11.8) decoding one token across batch size $b=32$ , number of heads $h=32$ , head dimension $d=128$ , chunk size $c=64$ , and FP16 precision. Baselines include naive PyTorch attention, xformers memory-efficient attention, PyTorch FlashAttention, standard vLLM PagedAttention, and PagedAttention* (where non-shared virtual pages are manually mapped to shared physical pages to isolate memory savings from kernel algorithm design).

$n_p$	$n_s$	Naive ( $\mu$ )s)	xformers ( $\mu$ )s)	FlashAttn ( $\mu$ )s)	PagedAttn ( $\mu$ )s)	PagedAttn* ( $\mu$ )s)	ChunkAttn ( $\mu$ )s)
1024	0	363.35	378.19	1586.73	356.17	355.82	332.50
1024	512	364.73	385.79	1587.14	355.88	257.74	198.87
1024	768	362.43	378.50	1591.61	356.02	215.18	131.21
1024	1024	361.76	379.36	1586.90	355.44	154.46	56.00
2048	0	686.40	816.44	3175.25	702.98	703.50	655.44
2048	1024	687.52	828.76	3173.53	703.35	505.32	384.37
2048	1536	685.78	820.19	3174.96	702.90	421.25	247.14
2048	2048	688.41	823.60	3152.25	703.72	338.41	110.48
4096	0	1369.52	1720.00	6289.55	1400.61	1400.17	1301.78
4096	2048	1370.47	1722.42	6303.21	1400.99	998.78	747.56
4096	3072	1369.74	1725.57	6301.41	1400.30	828.98	477.66
4096	4096	1370.41	1713.13	6300.65	1399.51	663.84	206.22

When prefixes are fully shared ( $n_s = n_p$ ), ChunkAttention achieves a $3.2\times$ to $4.8\times$ speedup over standard PagedAttention (e.g., $56.00\,\mu\text{s}$ vs $355.44\,\mu\text{s}$ at $n_p=1024$ ; $206.22\,\mu\text{s}$ vs $1399.51\,\mu\text{s}$ at $n_p=4096$ ). Comparing ChunkAttention with PagedAttn* shows a $2.8\times$ to $3.2\times$ speedup solely attributable to the Two-Phase Partition algorithm. When no prefix is shared ( $n_s=0$ ), ChunkAttention incurs no performance regression over baseline kernels.

Knowl 7 — End-to-End LLM Serving Performance with Prefix Sharing (ChunkLlama)

data/table

ChunkLlama incorporates ChunkAttention into LLaMA-2 7B (FP16) and was evaluated against vLLM 0.2.7 and Text Generation Inference (TGI 1.3.4) on an NVIDIA A100 GPU (80GB). Requests arrive according to a Poisson process with arrival rate $\lambda$ (requests per second, RPS). Maximum batch size is capped at 32, with completion token target $n_c=512$ .

$n_p$	$n_s$	$n_c$	RPS	Latency (ms/tok)		Peak KV Cache (GB)		Peak Batch Size
				vLLM	ChunkLlama	vLLM	ChunkLlama	vLLM	ChunkLlama
1024	0	512	1.0	19.92	19.11	14.73	11.90	23	18
1024	1024	512	1.0	20.80	14.07	14.79	3.28	23	14
2048	0	512	0.6	21.90	19.43	21.70	22.41	19	20
2048	2048	512	0.6	21.61	15.20	21.09	3.40	19	12
4096	0	512	0.4	26.23	26.88	34.59	35.13	16	16
4096	4096	512	0.4	27.62	17.16	35.42	4.00	16	11

Under shared prefix settings ( $n_s=n_p$ ), ChunkLlama reduces peak KV cache memory consumption by $70\%$ to $90\%$ (e.g., from $35.42\,\text{GB}$ to $4.00\,\text{GB}$ at $n_p=4096$ ). Serving throughput under a normalized latency bound of $40\,\text{ms/token}$ improves by $1.6\times$ (2.9 RPS vs 1.8 RPS) at $n_p=n_s=1024$ and by $2.3\times$ (2.3 RPS vs 1.0 RPS) at $n_p=n_s=2048$ . Peak batch size drops by $20\%$ to $40\%$ due to faster per-token decode latency.

Knowl 8 — Throughput Dynamics Under Sequence Divergence During Generation
empirical result
As autoregressive decoding proceeds, generated completion tokens ( $n_c$ ) diverge across sequences, lowering the effective prefix sharing ratio $r = n_s / (n_p + n_c)$ . In microkernel throughput experiments ( $b=32$ , $c=64$ , FP16, A100 GPU):
- At $n_s = 1024$ shared prompt tokens:
  - $n_c = 256$ : ChunkAttention reaches $241.93\times 10^3$ tokens/s vs PagedAttention's $76.35\times 10^3$ tokens/s ( $3.2\times$ speedup).
  - $n_c = 512$ : ChunkAttention reaches $186.44\times 10^3$ tokens/s vs $69.15\times 10^3$ tokens/s ( $2.7\times$ speedup).
  - $n_c = 1024$ : ChunkAttention reaches $127.85\times 10^3$ tokens/s vs $58.12\times 10^3$ tokens/s ( $2.2\times$ speedup).
- At $n_s = 2048$ shared prompt tokens:
  - $n_c = 512$ : ChunkAttention reaches $145.41\times 10^3$ tokens/s vs PagedAttention's $39.85\times 10^3$ tokens/s ( $3.6\times$ speedup; $2.0\times$ over PagedAttention*'s $73\times 10^3$ tokens/s).
  - $n_c = 1024$ : ChunkAttention reaches $107.37\times 10^3$ tokens/s vs $36.18\times 10^3$ tokens/s ( $3.0\times$ speedup).
  - $n_c = 2048$ : ChunkAttention reaches $70.33\times 10^3$ tokens/s vs $30.17\times 10^3$ tokens/s ( $2.3\times$ speedup; $1.5\times$ over PagedAttention*'s $46\times 10^3$ tokens/s).
- At $n_s = 4096$ shared prompt tokens:
  - $n_c = 512$ : ChunkAttention reaches $101.69\times 10^3$ tokens/s vs PagedAttention's $21.04\times 10^3$ tokens/s ( $4.8\times$ speedup).
  - $n_c = 4096$ : ChunkAttention reaches $37.05\times 10^3$ tokens/s vs $15.12\times 10^3$ tokens/s ( $2.4\times$ speedup).
Knowl 9 — Batch Size Scaling and Arithmetic Intensity Behavior in Prefix Attention
empirical result
Standard self-attention kernels during iterative decoding have an arithmetic intensity of approximately $0.99\,\text{FLOPs/byte}$ because query vectors ( $1 \times d$ ) are multiplied against stored KV tensors independently per sequence, making them strictly memory-bound. As a result, baseline kernels (Naive PyTorch, xformers, FlashAttention, and PagedAttention) plateau in throughput at batch size $b=16$ .

In contrast, ChunkAttention batches query vectors across sequences sharing common chunks in Phase 1 into a query matrix $Q_{i:j,:} \in \mathbb{R}^{(j-i) \times d}$ . When evaluated at $n_s = 2048$ shared tokens and $n_c = 64$ generated completion tokens on an A100 GPU:
- ChunkAttention throughput increases monotonically from $155\times 10^3$ tokens/s at $b=16$ to $224\times 10^3$ tokens/s at $b=96$ .
- This continuous throughput scaling is driven by improved temporal data locality on shared KV chunks in GPU SRAM and the conversion of matrix-vector operations into matrix-matrix operations executable on GPU Tensor Cores.
Knowl 10 — Structural and Deployment Limitations of ChunkAttention
limitation
ChunkAttention has three principal operational limitations:
1. Strict Prefix Location Constraint: Key-value tensor sharing requires identical tokens to be aligned at the exact beginning of the prompt sequence (index 0). If shared text, system instructions, or documents appear in the middle or end of user prompts, the prefix tree treats them as diverging paths, preventing KV cache deduplication.
2. Incompatibility with Custom Fine-Tuning Paradigms: The architectural benefit relies on multi-tenant serving workloads with long shared prompts (e.g., tool schemas or few-shot examples). If applications replace prompting with parameter-efficient fine-tuning or model fine-tuning, system prompt lengths and sharing opportunities decrease.
3. Hardware and Configuration Portability: The Two-Phase Partition kernel is implemented in low-level CUDA and specialized for common head dimensions (e.g., $d=128$ ) and specific GPU architectures (NVIDIA A100, RTX 4090). Adapting TPP to arbitrary head dimensions, novel accelerators, or specialized CPU environments requires dedicated manual tuning and implementation.

Coverage note — None was omitted; all core contributions, data structures, algorithms, mathematical formulations, optimization techniques, empirical microkernel/end-to-end results, and limitations are fully covered.

References

1.Reza Yazdani Aminabadi, Samyam Rajbhandari, Ammar Ahmad Awan, Cheng Li, Du Li, Elton Zheng, Olatunji Ruwase, Shaden Smith, Minjia Zhang, Jeff Rasley, and Yuxiong He. 2022. Deepspeed inference: Enabling efficient inference of transformer models at unprecedented scale. In SC22: International Conference for High Performance Computing, Networking, Storage and Analysis, pages 1–15.
2.Rohan Anil, Andrew M Dai, Orhan Firat, Melvin Johnson, Dmitry Lepikhin, Alexandre Passos, Siamak Shakeri, Emanuel Taropa, Paige Bailey, Zhifeng Chen, et al. 2023. Palm 2 technical report. arXiv e-prints, pages arXiv–2305.
3.Anthropic. 2023. How to use system prompts. https://docs.anthropic.com/claude/docs/how-to-use-system-prompts.
4.Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. 2020. Language models are few-shot learners. Advances in neural information processing systems, 33:1877–1901.
5.Yupeng Chang, Xu Wang, Jindong Wang, Yuan Wu, Linyi Yang, Kaijie Zhu, Hao Chen, Xiaoyuan Yi, Cunxiang Wang, Yidong Wang, Wei Ye, Yue Zhang, Yi Chang, Philip S. Yu, Qiang Yang, and Xing Xie. 2023. A survey on evaluation of large language models.
6.Jack Choquette, Wishwesh Gandhi, Olivier Giroux, Nick Stam, and Ronny Krashinsky. 2021. Nvidia a100 tensor core gpu: Performance and innovation. IEEE Micro, 41(2):29–35.
7.Zheng Chu, Jingchang Chen, Qianglong Chen, Weijiang Yu, Tao He, Haotian Wang, Weihua Peng, Ming Liu, Bing Qin, and Ting Liu. 2023. A survey of chain of thought reasoning: Advances, frontiers and future.
8.Together Computer. 2023. Redpajama-data: An open source recipe to reproduce llama training dataset.
9.Tri Dao. 2023. Flashattention-2: Faster attention with better parallelism and work partitioning. arXiv preprint arXiv:2307.08691.
10.Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. 2022. Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances in Neural Information Processing Systems, 35:16344–16359.
11.Qingxiu Dong, Lei Li, Damai Dai, Ce Zheng, Zhiyong Wu, Baobao Chang, Xu Sun, Jingjing Xu, Lei Li, and Zhifang Sui. 2023. A survey on in-context learning.
12.Pin Gao, Lingfan Yu, Yongwei Wu, and Jinyang Li. 2018. Low latency rnn inference with cellular batching. In Proceedings of the Thirteenth EuroSys Conference, pages 1–15.
13.Gemini. 2023. Gemini: A family of highly capable multimodal models.
14.Xinyang Geng and Hao Liu. 2023. Openllama: An open reproduction of llama.
15.gyudoza. 2023. jujumilk3/leaked-system-prompts: Collection of leaked system prompts. https://github.com/jujumilk3/leaked-system-prompts.
16.Steve Hill. 1992. A simple fast memory allocator. In DAVID KIRK, editor, Graphics Gems III (IBM Version), pages 49–50. Morgan Kaufmann, San Francisco.
17.Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin De Laroussilhe, Andrea Gesmundo, Mona Attariyan, and Sylvain Gelly. 2019. Parameter-efficient transfer learning for nlp. In International Conference on Machine Learning, pages 2790–2799. PMLR.
18.Zhiqiang Hu, Lei Wang, Yihuai Lan, Wanyu Xu, Ee-Peng Lim, Lidong Bing, Xing Xu, Soujanya Poria, and Roy Ka-Wei Lee. 2023. Llm-adapters: An adapter family for parameter-efficient fine-tuning of large language models.
19.HuggingFace. 2023. huggingface/text-generation-inference: Large language model text generation inference. https://github.com/huggingface/text-generation-inference.
20.Yunho Jin, Chun-Feng Wu, David Brooks, and Gu-Yeon Wei. 2023. S3: Increasing gpu utilization during generative inference for higher throughput. arXiv preprint arXiv:2306.06000.
21.Sehoon Kim, Coleman Hooper, Thanakul Wattanawong, Minwoo Kang, Ruohan Yan, Hasan Genc, Grace Dinh, Qijing Huang, Kurt Keutzer, Michael W. Mahoney, Yakun Sophia Shao, and Amir Gholami. 2023. Full stack optimization of transformer inference: a survey.
22.Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph E. Gonzalez, Hao Zhang, and Ion Stoica. 2023. Efficient memory management for large language model serving with pagedattention.
23.Benjamin Lefaudeux, Francisco Massa, Diana Liskovich, Wenhan Xiong, Vittorio Caggiano, Sean Naren, Min Xu, Jieru Hu, Marta Tintore, Susan Zhang, Patrick Labatut, and Daniel Haziza. 2022. xformers: A modular and hackable transformer modelling library.
24.Minghao Li, Yingxiu Zhao, Bowen Yu, Feifan Song, Hangyu Li, Haiyang Yu, Zhoujun Li, Fei Huang, and Yongbin Li. 2023. Api-bank: A comprehensive benchmark for tool-augmented llms.
25.Nelson F. Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang. 2023. Lost in the middle: How language models use long contexts.
26.Pan Lu, Swaroop Mishra, Tanglin Xia, Liang Qiu, Kai-Wei Chang, Song-Chun Zhu, Oyvind Tafjord, Peter Clark, and Ashwin Kalyan. 2022. Learn to explain: Multimodal reasoning via thought chains for science question answering. Advances in Neural Information Processing Systems, 35:2507–2521.
27.Pan Lu, Baolin Peng, Hao Cheng, Michel Galley, Kai-Wei Chang, Ying Nian Wu, Song-Chun Zhu, and Jianfeng Gao. 2023a. Chameleon: Plug-and-play compositional reasoning with large language models. In Advances in Neural Information Processing Systems, volume 36, pages 43447–43478. Curran Associates, Inc.
28.Pan Lu, Liang Qiu, Kai-Wei Chang, Ying Nian Wu, Song-Chun Zhu, Tanmay Rajpurohit, Peter Clark, and Ashwin Kalyan. 2023b. Dynamic prompt learning via policy gradient for semi-structured mathematical reasoning. In International Conference on Learning Representations (ICLR).
29.Maxim Milakov and Natalia Gimelshein. 2018. Online normalizer calculation for softmax. arXiv preprint arXiv:1805.02867.
30.oneDNN Contributors. 2023. oneapi deep neural network library (onednn). https://github.com/oneapi-src/oneDNN.
31.OpenAI. 2023a. Chatgpt plugins. https://platform.openai.com/docs/plugins/introduction.
32.OpenAI. 2023b. Function calling - openai api. https://platform.openai.com/docs/guides/function-calling.
33.OpenAI. 2023c. Gpt-4 technical report. arXiv preprint arXiv:2303.08774.
34.OpenAI. 2023d. How to call functions with chat models. https://cookbook.openai.com/examples/how_to_call_functions_with_chat_models.
35.OpenAI. 2023e. openai/tiktoken: tiktoken is a fast bpe tokeniser for use with openai’s models. https://github.com/openai/tiktoken.
36.Cheng Qian, Chi Han, Yi R. Fung, Yujia Qin, Zhiyuan Liu, and Heng Ji. 2023. Creator: Tool creation for disentangling abstract and concrete reasoning of large language models.
37.Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. 2018. Improving language understanding by generative pre-training.
38.Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. 2019. Language models are unsupervised multitask learners. OpenAI blog, 1(8):9.
39.Jon Saad-Falcon, Joe Barrow, Alexa Siu, Ani Nenkova, David Seunghyun Yoon, Ryan A. Rossi, and Franck Dernoncourt. 2023. Pdftriage: Question answering over long, structured documents.
40.Timo Schick, Jane Dwivedi-Yu, Roberto Dessì, Roberta Raileanu, Maria Lomeli, Luke Zettlemoyer, Nicola Cancedda, and Thomas Scialom. 2023. Toolformer: Language models can teach themselves to use tools.
41.Ying Sheng, Lianmin Zheng, Binhang Yuan, Zhuohan Li, Max Ryabinin, Daniel Y Fu, Zhiqiang Xie, Beidi Chen, Clark Barrett, Joseph E Gonzalez, et al. 2023. High-throughput generative inference of large language models with a single gpu. arXiv preprint arXiv:2303.06865.
42.Franyell Silfa, Jose Maria Arnau, and Antonio González. 2022. E-batch: Energy-efficient and high-throughput rnn batching. ACM Trans. Archit. Code Optim., 19(1).
43.Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. 2023a. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971.
44.Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. 2023b. Llama: Open and efficient foundation language models.
45.Mariano Trebino. 2016. mtrebi/memory-allocators: Custom memory allocators in c++ to improve the performance of dynamic memory allocation. https://github.com/mtrebi/memory-allocators#pool-allocator.
46.Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, et al. 2022. Chain-of-thought prompting elicits reasoning in large language models. Advances in Neural Information Processing Systems, 35:24824–24837.
47.Jules White, Quchen Fu, Sam Hays, Michael Sandborn, Carlos Olea, Henry Gilbert, Ashraf Elnashar, Jesse Spencer-Smith, and Douglas C. Schmidt. 2023. A prompt pattern catalog to enhance prompt engineering with chatgpt.
48.Samuel Williams, Andrew Waterman, and David Patterson. 2009. Roofline: An insightful visual performance model for multicore architectures. Commun. ACM, 52(4):65–76.
49.Gyeong-In Yu, Joo Seong Jeong, Geon-Woo Kim, Soojeong Kim, and Byung-Gon Chun. 2022. Orca: A distributed serving system for {Transformer-Based} generative models. In 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22), pages 521–538.
50.Yongchao Zhou, Andrei Ioan Muresanu, Ziwen Han, Keiran Paster, Silviu Pitis, Harris Chan, and Jimmy Ba. 2023. Large language models are human-level prompt engineers.
51.Yuchen Zhuang, Yue Yu, Kuan Wang, Haotian Sun, and Chao Zhang. 2023. Toolqa: A dataset for llm question answering with external tools.

Citation

MLA

Ye, L., et al. “ChunkAttention: Efficient Self-Attention with Prefix-Aware KV Cache and Two-Phase Partition”. Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), 2024, pp. 11608–20, https://doi.org/10.18653/v1/2024.acl-long.623.

APA

Ye, L., Tao, Z., Huang, Y., & Li, Y. (2024). ChunkAttention: Efficient Self-Attention with Prefix-Aware KV Cache and Two-Phase Partition. Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), 11608–11620. https://doi.org/10.18653/v1/2024.acl-long.623

Chicago

Ye, L., Z. Tao, Y. Huang, and Y. Li. 2024. “ChunkAttention: Efficient Self-Attention with Prefix-Aware KV Cache and Two-Phase Partition”. Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), 11608–20. https://doi.org/10.18653/v1/2024.acl-long.623.

Harvard

Ye, L. et al. (2024) “ChunkAttention: Efficient Self-Attention with Prefix-Aware KV Cache and Two-Phase Partition”, Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). Association for Computational Linguistics, pp. 11608–11620. Available at: https://doi.org/10.18653/v1/2024.acl-long.623.

Vancouver

1. Ye L, Tao Z, Huang Y, Li Y (2024) ChunkAttention: Efficient Self-Attention with Prefix-Aware KV Cache and Two-Phase Partition. In: Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). Association for Computational Linguistics, pp 11608–11620

BibTeX

@inproceedings{ye-etal-2024-chunkattention,
    title = "{C}hunk{A}ttention: Efficient Self-Attention with Prefix-Aware {KV} Cache and Two-Phase Partition",
    author = "Ye, Lu  and
      Tao, Ze  and
      Huang, Yong  and
      Li, Yang",
    editor = "Ku, Lun-Wei  and
      Martins, Andre  and
      Srikumar, Vivek",
    booktitle = "Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)",
    month = aug,
    year = "2024",
    address = "Bangkok, Thailand",
    publisher = "Association for Computational Linguistics",
    url = "https://aclanthology.org/2024.acl-long.623/",
    doi = "10.18653/v1/2024.acl-long.623",
    pages = "11608--11620"
}

Metadata:ACL Anthology

Source Code

This paper has an official code repository available. Click below to access the source code.

View Repository

Access the Paper

This paper is available from its original source. Click below to access the PDF.

Open PDF

License: https://creativecommons.org/licenses/by/4.0/

ChunkAttention: Efficient Self-Attention with Prefix-Aware KV Cache and Two-Phase Partition

Abstract

Topics

Keywords

Table of Contents

Knowls

Knowl 1 — Prefix-Aware Key-Value Cache (PAKV) Structure and Dynamic Management

Knowl 2 — Two-Phase Partition (TPP) Self-Attention Decoding Algorithm

Knowl 3 — Partial Attention and Online Softmax Reduction Equations

Knowl 4 — Theoretical Bounds on Batch Capacity Gain and KV Cache Fragmentation

Knowl 5 — CPU-GPU Context Management and Latency Hiding for Prefix Trees