ChunkAttention: 접두사 인식 KV 캐시와 이단계 분할을 통한 효율적인 자기 주의 메커니즘

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

February 23, 2024

저자: Lu Ye, Ze Tao, Yong Huang, Yang Li

cs.AI

초록

셀프 어텐션은 대규모 언어 모델(LLM)의 필수 구성 요소이지만, 긴 시퀀스에 대한 추론 지연의 주요 원인 중 하나입니다. 다중 테넌트 LLM 서비스 시나리오에서, 셀프 어텐션의 계산 및 메모리 연산 비용은 여러 LLM 요청이 접두사에서 공유 시스템 프롬프트를 가질 확률을 활용하여 최적화할 수 있습니다. 본 논문에서는 ChunkAttention을 소개합니다. 이는 접두사 인식 셀프 어텐션 모듈로, 여러 요청 간에 일치하는 프롬프트 접두사를 감지하고 런타임에 메모리에서 키/값 텐서를 공유하여 KV 캐시의 메모리 활용도를 향상시킬 수 있습니다. 이는 단일 키/값 텐서를 더 작은 청크로 분할하고 이를 보조 접두사 트리로 구조화함으로써 달성됩니다. 결과적으로, 접두사 트리 기반 KV 캐시 위에 효율적인 셀프 어텐션 커널을 설계하였으며, 여기서는 공유 시스템 프롬프트가 존재할 때 셀프 어텐션 계산 중 데이터 지역성을 개선하기 위해 2단계 분할 알고리즘이 구현되었습니다. 실험 결과, ChunkAttention은 최신 구현 대비 셀프 어텐션 커널을 3.2~4.8배 빠르게 수행할 수 있으며, 시스템 프롬프트의 길이는 1024에서 4096까지 다양합니다.

English

Self-attention is an essential component of large language models(LLMs) 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.8times compared to the start-of-the-art implementation, with the length of the system prompt ranging from 1024 to 4096.