핵심 개념
대규모 언어 모델의 자기 주의 메커니즘은 추론 지연의 주요 원인이지만, 시스템 프롬프트의 공유 특성을 활용하여 메모리 사용과 계산 비용을 크게 줄일 수 있다.
초록
이 논문은 ChunkAttention이라는 새로운 자기 주의 모듈을 제안한다. ChunkAttention은 접두사 인식 KV 캐시(PAKV)와 두 단계 분할(TPP) 기법을 통해 대규모 언어 모델의 자기 주의 메커니즘을 최적화한다.
PAKV는 KV 캐시를 접두사 트리 구조로 구성하여 다중 요청 간 공유되는 시스템 프롬프트 접두사를 동적으로 탐지하고 제거한다. 이를 통해 메모리 사용량을 크게 줄일 수 있다.
TPP는 PAKV 기반의 KV 캐시에 최적화된 자기 주의 커널을 구현한다. 청크 단위와 시퀀스 단위로 병렬화를 수행하여 데이터 지역성을 높이고 메모리 연산을 최소화한다.
실험 결과, ChunkAttention은 기존 최적화 기법 대비 3.2-4.8배 빠른 자기 주의 커널 성능을 보였으며, 전체 추론 속도에서도 1.6-2.3배 향상된 처리량을 달성했다. 또한 KV 캐시 메모리 사용량을 70-90% 줄일 수 있었다.
통계
시스템 프롬프트 길이가 1024에서 4096 토큰 범위일 때, ChunkAttention의 자기 주의 커널 성능이 기존 구현 대비 3.2-4.8배 향상되었다.
ChunkLlama의 전체 추론 속도가 기존 구현 대비 1.6-2.3배 향상되었다.
ChunkLlama의 KV 캐시 메모리 사용량이 70-90% 감소했다.
인용구
"시스템 프롬프트는 대규모 언어 모델 기반 애플리케이션 설계에 있어 필수적인 요소이지만, 자기 주의 메커니즘의 성능과 메모리 사용에 부정적인 영향을 미친다."
"ChunkAttention은 접두사 인식 KV 캐시와 두 단계 분할 기법을 통해 자기 주의 메커니즘의 성능과 메모리 효율성을 크게 향상시킬 수 있다."
"실험 결과, ChunkAttention은 기존 최적화 기법 대비 3.2-4.8배 빠른 자기 주의 커널 성능과 1.6-2.3배 향상된 전체 추론 속도를 보였다."