본문 바로가기
논문리딩

[논문 요약] Prompt Cache: Modular Attention Reuse for Low-Latency Inference

by MFDO 2024. 8. 19.

최근 클로드 모델에 Prompt Cache가 도입되며,

특정 조건 상황에서 API 비용이

이론상 최대 90%의 절약 효과를 보였다는 기사를 읽게되었다.

 

관련하여  예일대학교와 구글 연구진의 논문이 대두되었기에

읽어보는 시간을 가지게 되었다.

 

 

나는 인쇄해서 읽는게 제일 잘 읽힌다...

 

원문

https://arxiv.org/pdf/2311.04934

 


| 짧은 요약

LLM의 추론에서 자주 사용되는 텍스트 세그먼트의 어텐션 상태를 모듈화하여 재사용함으로써,

대기 시간을 크게 단축하는 기술을 제시하였다.

PML을 통해 프롬프트 구조를 명확히 정의하고,

이를 바탕으로 GPU와 CPU 기반의 성능을 향상시킬 수 있었다고 한다. 

 


 

 

| 아이디어 발단 상황

대규모 언어 모델(LLM)의 추론 과정에서 동일하거나 유사한 입력 프롬프트가 자주 재사용된다.

특히 시스템 메시지, 문서 템플릿 등과 같은 공통된 텍스트 세그먼트가

여러 프롬프트에 걸쳐 반복적으로 사용되면서, 이러한 반복적인 계산이 성능 병목 현상을 초래한다.

 

| 제안

Prompt Cache는 이러한 문제를 해결하기 위해 설계된 기술로,

LLM의 추론 과정에서 자주 사용되는 text segment의 *attention state를 미리 계산하고 저장하여 재사용

=> 추론 시 반복되는 계산을 줄여 대기 시간을 단축

 

Attention State?
"입력 시퀀스의 모든 부분이 동일한 중요성을 가지지 않는다."를 전제로 입력의 중요한 정보를 담고 있는 부분

왜 시간이 단축되는가?
"We note that the performance advantage becomes more pronounced as the size of cached segments grows since the computation overhead of attention states scales quadratically with input sequence size."
즉, 입력 시퀀스의 크기가 커질수록 Attenttion state를 계산하는 데 드는 오버헤드가 제곱의 비율로 증가하기 때문에
캐시된 segement의 크기가 커질수록 성능이 더 좋아진다.
논문에서는 입력 시퀀스가 크기에 따라, Attention sate를 계산시 필요 연산량이 제곱으로 증가하게 된다고 한다.

따라서, 계산 오버헤드가 커지면서 성능이 저하되는 것을 방지하기 위해 캐시를 사용하면, 이미 계산한 결과를 재사용하여 계산 부담을 줄이고 성능을 높일 수 있다.
캐시된 세그먼트가 크면 클수록, 더 많은 정보를 미리 저장하고 재사용할 수록, 성능 이점이 더 커진다는 의미다.

 

 


이후 내용 수정 및 다듬기 必

 

 

| 이전 사례

기존 Key-Value (KV) Cache는 single prompt  내에서 Attentions state를 재사용하여 계산량을 줄이는 방법으로 사용되어졌음

Prompt Cache는 이 multiple prompt 간에 어텐션 상태를 모듈화하여 재사용할 수 있도록 설계하였음

 

 - 각 도형의 의미 작성해주기

 - 파란 사각형 : prompt

 - 그냥 사각형 : 토큰

 - ...

 

 

| 시행 착오 및 극복

1) Transformer Positional encoding

2) efficiently recognize a text segment

 

Transformer의 Positional encoding 특징
Transformer 소개 여기 작성하기
입력 데이터 순서에 구애받지 않아 병렬처리가 가능해졌지만 이 입력데이터의 순서는 중요함.
이러한 데이터의 순서를 인코딩을 통해 인식하게 하는데 그것이 Transfomer의 Positional encoding 역할
 

 

| 분석 방법

Prompt Cache의 성능을 평가하기 위해 HuggingFace Transformer 라이브러리를 기반으로 한 프로토타입을 구현하고, 여러 LLM에서 이를 테스트

주 실험에서 GPU와 CPU에서의 시간 절약 효과를 측정하고, 다양한 long context 데이터셋을 사용하여 성능을 검증함.

 

 

| 도출 방안

Prompt Cache는 PML(Prompt Markup Language)이라는 언어를 도입하여 프롬프트의 구조를 명시적으로 정의하고,

자주 사용되는 텍스트 모듈의 위치 ID를 부여함으로써 어텐션 상태를 재사용할 수 있게 함.

 => 사용자가 프롬프트 모듈을 쉽게 선택하거나 업데이트할 수 있으며, 새로운 텍스트 세그먼트를 효율적인 처리가 가능해짐.

 

 

| 분석 결과

 - 대기 시간을 감소 성과  

  > GPU 기반 추론에서 1.5배에서 10배

  > CPU 기반 추론에서는 20배에서 70배

 

 - 정확도 손실 없이 달성, 다양한 LLM에 걸쳐 일관된 성능 향상을 보임.

 

| 의의 및 활용 가능성

Prompt Cache는 미래의 고성능 LLM 시스템을 위한 중요한 구성 요소로 확장 가능

 ex) 메모리 관리 및 캐시 교체 전략 등 최적화 기법 적용

연구 결과가 오픈 소스

 

핵심 기술 : 프롬프트 모듈화, 어텐션 상태의 재사용

  > PML(Prompt Markup Language)를 도입하여 프롬프트의 구조를 명확하게 정의

  > 캐시에서 결과를 가져오는 것이 아니라, 프롬프트의 구성 요소를 모듈화하여 필요한 부분만 선택적으로 재사용가능

 =>  템플릿 기반의 대규모 문서 생성, 자동 응답 시스템, 반복적인 데이터 처리 작업에서 큰 효과를 발휘

댓글