PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Lu Ye, Ze Tao, Yong Huang, Yang Li

arXiv:2402.15220v4cs.LGcs.CL

TL;DR

Self-attention is a memory-bound inference bottleneck, and shared system prompts create redundant KV-cache storage across requests. ChunkAttention uses a prefix tree to share KV tensors at runtime and a two-phase partition kernel. On A100 (80G), it outperforms the state-of-the-art PagedAttention kernel by 3.2-4.8× for shared prompts of 1024 to 4096 tokens, while achieving comparable throughput without shared prompts.

  • Problem

    Self-attention is memory-bound during inference, while shared system prompts create redundant KV-cache storage across requests.

    Source: §1 · p. 1

  • Method

    ChunkAttention stores chunked key/value tensors in a prefix tree for runtime prefix-aware sharing and applies two-phase partitioning to its self-attention kernel.

    Source: §1 · p. 2

  • Results

    3.2-4.8× faster than the SOTA PagedAttention kernel with shared system prompts of 1024 to 4096 tokens on A100 (80G), and comparable throughput without shared prompts.

    Source: §6 · p. 8

  • Takeaways & Limitations

    Prefix-aware KV caching and two-phase partitioning improve attention throughput in the reported shared-prompt setting while retaining comparable throughput without shared prompts.

    Source: §6 · p. 8

    Sharing requires the common system prompt to appear at the beginning of the sequence, and the authors call for generalizing two-phase partitioning to more model configurations and hardware.

    Source: §7 · p. 9

Abstract

from arXiv · show

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 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.

1 Introduction

Self-attention slows long-context inference through memory-bound KV-cache operations, while shared system prompts create redundant cache across requests. ChunkAttention proposes runtime prefix-aware KV sharing and a two-phase attention kernel; experiments report greater speed as shared prompts grow without regression when none are shared.

1 Introduction

  • Self-attention is memory-bound: KV-cache operations drive latency, cache size grows linearly with context, and the cache restricts batch size and throughput.The introduction illustrates the capacity constraint with an FP16 example: an 8×A100 (80G) server holds 35 sequences of 2K context tokens.Source: §1 · p. 1
  • Shared system prompts create redundant KV cache across requests and can exceed 1K tokens in LLM applications.A ChatGPT-like chatbot’s shared prompt can reach 1,766 tokens with six plugins activated.Source: §1 · p. 1
  • Prior reservation of predefined prompt tensors is inflexible to refreshes, can waste memory at low hit rates, and does not optimize the attention kernel.The proposal requires application developers and the service provider to participate in the operation loop.Source: §1 · p. 2
  • ChunkAttention organizes chunked KV tensors in a prefix tree for runtime redundancy removal and uses two-phase partitioning to batch queries with matching prefixes.The prefix-aware cache detects and removes redundancy without human involvement; the kernel uses chunk-first and sequence-first phases.Source: §1 · p. 2
  • ChunkAttention is significantly faster as shared system prompts grow and has no performance degradation without shared prompts, compared with existing optimized implementations.The introduction reports the overall trend but gives no numerical speedup.Source: §1 · p. 2

2 Preliminaries

System prompts supply instructions, examples, or external knowledge and are reused across requests in online applications and templated research workloads. LLM inference prefills and caches prompt tokens before autoregressive decoding, where iteration-based batching combines one token from each active sequence.

2.1 Shared System Prompt

  • System prompts supply instructions, few-shot examples, or external knowledge and are shared across requests; with six plugins, one chatbot prompt can reach 1,766 tokens.The final prompt concatenates the system prompt with task-specific input.Source: §2.1 · p. 2
  • Offline LLM experiments issue many templated requests with identical instructions, examples, or external knowledge, reusing system prompts across tasks.Examples include compositional reasoning, question generation, and document question-answering workloads.Source: §2.1 · pp. 2–3

2.2 LLM Inferencing

  • Prefilling processes prompt tokens and caches their key/value tensors; decoding then generates tokens autoregressively until an end condition.Decoding stops at the end-of-sequence token or the maximum completion length.Source: §2.2 · p. 3
  • Iteration-based batching combines one current token from each sequence during decoding, while new sequences can join and completed ones leave; ChunkAttention assumes this batching is enabled.The batched tokens undergo QKV projection before self-attention, while attention is computed independently for each sequence.Source: §2.2 · p. 3

3 Our Approach

ChunkAttention organizes chunked key/value tensors in a prefix tree so sequences with common prefixes can share cache memory. Its two-phase kernel batches work on shared chunks before processing sequence-specific chunks.

3.1 Prefix Aware KV Cache (PAKV)

  • Chunked key/value tensors for common prefixes are shared in a prefix tree, with each path representing a sequence.Each tree node stores a shared context-token segment and its key/value slices.Source: §3.1 · p. 3, Figure 1 · p. 4
  • New sequences add paths to the prefix tree, completed sequences remove paths, and decoding appends tokens to leaf chunks or grows new chunks.Source: §3.1 · p. 4
  • Sharing common prefixes increases simultaneously processed sequences by approximately 1/(1-r), where r=ns/(np+nc).In memory-limited inference, this can help increase batch size and throughput.Source: §3.1 · p. 4

3.2 Two-phase Partition (TPP)

  • The chunk-first phase batches queries for shared chunks and saves partial attention results; the sequence-first phase restores them and processes remaining chunks per query.Source: §3.2 · p. 4, Figure 2 · p. 5
  • Batching queries for sequences sharing a chunk improves data locality and turns query vectors into matrices for tensor-core matrix multiplication.Source: §3.2 · p. 5
  • Without chunk-first partial results, sequence-first processing reloads shared chunks from RAM b times; the two phases balance parallelization and cache locality.Source: §3.2 · p. 6

3.3 Further Optimizations

  • GPU context copies are cached until the tree changes, while CPU devices can merge partial results using spin locks instead of GPU atomics.CPU context generation can overlap GPU work; copies are triggered by a full chunk, a joining sequence, or a leaving sequence.Source: §3.3 · p. 6

4 Experiments

The experiments evaluate ChunkAttention at both the self-attention microkernel and end-to-end model levels, with prefix sharing, completion length, and batch size varied. Results show higher throughput with shared prefixes and lower KV-cache memory use in end-to-end serving.

4 Experiments

  • Microkernel tests measure self-attention CUDA-kernel time, while end-to-end tests capture prefix-tree overhead; experiments use an A100 80G with CUDA 11.8.The microkernel workload varies prompt and shared-token counts, completion length, and batch size.Source: §4 · p. 6, §4.1 · p. 6

4.1 Microkernel Evaluation

  • At ns=4096, Naive is 6.6× slower than ChunkAttn and 2.1× slower than PagedAttn*; across ns=1024–4096, ChunkAttn outperforms PagedAttn* by 2.8–3.2×.PagedAttn* simulates physically shared KV memory; it does not implement PAKV or TPP.Source: §4.1 · pp. 6–7, Table 3 · p. 7
  • At ns=4096, physical KV sharing reduces PagedAttn* latency by up to 52% versus PagedAttn, and TPP causes no regression when ns=0.The latency comparison uses chunk size 64 and batch size 32.Source: §4.1 · p. 7, Table 3 · p. 7
  • With ns=2048, ChunkAttn remains 2.0× faster than PagedAttn* at nc=512 and 1.5× faster at nc=2048, although its advantage over PagedAttn narrows as decoding proceeds.The reported rates against PagedAttn* are 145K versus 73K at nc=512 and 70K versus 46K at nc=2048.Source: §4.1 · p. 7, Figure 3 · p. 7
  • With ns=2048, ChunkAttn throughput rises from 155K to 224K toks/s as batch size increases from 16 to 96.The authors attribute the trend to better data locality and improved arithmetic intensity.Source: §4.1 · p. 7, Figure 4 · p. 7

4.2 End-to-end Evaluation

  • 1.6× (2.9 against 1.8) and 2.3× (2.3 against 1.0) higher throughput than vLLM with 1024 and 2048 shared prefix tokens, respectively, while normalized latency stays below 40 ms/token.ChunkLlama's KV-cache memory use falls by 70%–90% with long shared prefixes.Source: §4.2 · pp. 7–8, Figure 5 · p. 8, Table 4 · p. 8

5 Related Work

Related work addresses KV-cache memory management and self-attention tiling, while ChunkAttention combines runtime prefix sharing with a partition strategy for decoding.

5 Related Work

  • PagedAttention addresses memory waste from dynamic sequence lengths, whereas ChunkAttention uses a prefix tree to discover redundant KV cache across requests at runtime (Kwon et al., 2023).The cited shared-prompt proposal was not implemented in vLLM releases up to 0.2.7.Source: §5 · p. 8
  • ChunkAttention targets multi-tenant deployments with runtime discovery rather than precompiling shared prompts, and uses prefix-tree context to guide its two-phase partition algorithm.The paper describes the approach as suited to centrally hosted models with scalability requirements.Source: §5 · p. 8
  • FlashAttention's tiling accelerates self-attention, but its inflexibility with non-contiguous memory and variable sequence lengths limits its fit for inference; ChunkAttention handles variable-length decoding and is described as complementary.ChunkAttention batches attention operations across sequences to reduce memory operations.Source: §5 · p. 8

6 Conclusion

ChunkAttention uses a prefix tree for runtime KV-cache deduplication and a two-phase partition kernel for faster inference. The kernel matches PagedAttention throughput without shared prompts and outperforms it by 3.2–4.8× with 1024–4096 shared system-prompt tokens on A100 80G.

6 Conclusion

  • ChunkAttention matches PagedAttention throughput without shared prompts and outperforms it by 3.2–4.8× with 1024–4096 shared system-prompt tokens on A100 80G.The reported result uses both prefix-aware KV cache and two-phase partition.Source: §6 · p. 8

7 Limitations

ChunkAttention’s prefix-aware KV sharing depends on matching prompts at the start of sequences, and its kernel’s tuning is limited to common model and hardware configurations. Fine-tuning could also reduce the opportunity to share system prompts.

The Position of System Prompt

  • System-prompt sharing requires the prompt at the sequence start; matching tokens elsewhere cannot save KV-cache memory when the full sequence prefixes differ.The paper notes that placing the prompt first is common but not mandatory.Source: §7 · p. 9

Fine Tuning

  • As fine-tuning becomes more practical, applications may need shorter system prompts, reducing the opportunities for prompt sharing.The paper says it has not seen more promising and cost-efficient fine-tuning and hosting solutions than system prompts as of its writing.Source: §7 · p. 9

Model and Hardware Compatibility

  • The CUDA two-phase partition kernel is tuned for common configurations and hardware; other combinations require case-by-case tuning and verification, adding development costs.Examples of tuned settings include a 128-dimensional head and NVIDIA A100, GeForce RTX 4090, and Intel Xeon CPU.Source: §7 · p. 9

A System Prompt for Chatbot Applications with Plugins

The appendix gives a 1,766-token system prompt for selecting plugin APIs from user requests, parsing parameters, and formatting calls.

API selection

  • The prompt instructs the chatbot to choose an API matching the user’s intent, parse its parameters, and return not_found when no API matches or a required parameter is absent.It specifies a strict api_chosen(param1=PARSED_PARAM1, ...) response syntax.Source: A System Prompt for Chatbot… · p. 12

API specifications

  • Available endpoints cover Bing web and image search, Expedia hotel and flight search, OpenTable restaurant search, and Spotify catalog search.The prompt describes these API functions and their parameters.Source: A System Prompt for Chatbot… · pp. 12–13
  • For Bing web search, the optional count parameter specifies how many search results to return; the default is 10, and the cited passage also notes a maximum value.Source: A System Prompt for Chatbot… · p. 12

Examples

Loading 2402.15220v4…