[PaperReview] Prompt Cache: Modular Attention Reuse For Low-Latency Inference

https://arxiv.org/pdf/2311.04934

 

1. 프롬프트 캐시(Prompt Cache) 개발: 자주 사용하는 텍스트 세그먼트의 주의 상태를 미리 계산하여 저장, 재사용으로 지연 시간 감소

2. 프롬프트 마크업 언어(PML) 활용: 재사용 가능한 텍스트 세그먼트를 프롬프트 모듈로 정의, 이용자가 쉽게 접근 및 수정 가능하게 함.

3. 지연시간 개선 평가: 다양한 LLM에서 평균 8배에서  60배까지 지연 시간 감소 확인, 정확성 유지

 

요약

본 논문에서는 Prompt Cache라는 접근법을 제안하고 있음. 이 방법은 대형 언어모델 (LLM)의 추론을 가속화하는데  중점을 두며, 서로 다른 LLM 프롬프트 간에 주의 상태를 재사용하는 방식이다. 많은 입력 프롬프트는 시스템 메시지, 프롬프트 템플릿 및 문서와 같은 중첩된 텍스트 세그먼트를 포함하고 있는데, 이러한 세그먼트를 사전에 계산하고 저장함으로써 각각의 사용자 프롬프트에 대해 효율적으로 재사용할 수 있다.

 

핵심 방법론

1. 프폼프트 모듈 정의: Prompt Cache 는 재사용 가능한 텍스트 세그먼트를 명확히 정의하는 스키마인 프롬프트 마크업 언어(PML)를 사용한다. 이는 주의 상태 재사용  시의 위치 정확성을 보장하고, 사용자에게 캐시된 상태에 접근할 수있는 인터페이스를 제공한다.
2. 프롬프트 모듈 인코딩: 프롬프트 모듈은 스키마로부터 추출된 토큰 시퀀스를 기반으로 위치 ID가 할당 된다. 이 과정을 통해 모든 프롬프트 모듈에 대해 (k,v) 주의 상태를 계산하고 캐시한다.
3. 주의 상태 재사용: 사용자가 프롬프트를 작성하면, Prompt Cache는 스키마를 처리하고 캐시된 모듈의 주의 상태를 조회한 후, 새로운 텍스트 세그먼트에 대해 주의 상태를 계산한다. 사용자가 제공한 매개 변수에 따라서 유연한 구조의 프롬프트를 제공할 수 있다.
4. 프롬프트 캐시 구조: 프롬프트 모듈은 서로 독립적인 구성을 가지면서도, 필요시 모든 모듈의 주의 상태를 함께 계산하고 캐시를 통해 재사용  할 수 있다. 이는 여러 프롬프트 간의 중복을 줄이고, 전체 시스템의 효율성을 높인다.

 

결론

Prompt Cache는 LLM의 추론을 가속화하기 위한 혁신적인 접근법으로, 입력 프롬프트의 구조적 중복을 효과적으로  활용한다. 이 연구는 향후 LLM 제공 시스템의 기본 구성 요소로 사용할 수 있느 가능성을 열어두며, GPU 캐시 대체 전력 등이 포함된 더 발전된 시스템 설계가 필요함을 강조한다.

이 논문은 또한 향후 RAG와  같이 레이턴시가 중요한 Application에서도 유용하게 작용할 것으로 기대함. 이와 같은 방식으로 연구자들은 보다 복잡한 프롬프트 구조와 고급 기능을 필요로 하는 작업에 이 접근 법을 활용 할 수 있을 것이다. 

 

핵심 기술 요약

1) 프롬프트 모듈화 및 캐싱

• *프롬프트 모듈: 자주 사용되는 텍스트 세그먼트를 프롬프트 모듈로 정의하고, 해당 모듈의 어텐션 상태를 미리 계산하여 저장합다.
• *재사용: 프롬프트에 이러한 모듈이 등장할 때, 미리 저장된 어텐션 상태를 불러와 재사용함으로써 중복 계산을 방지합다.

2) Prompt Markup Language (PML)

• *구조 명시화: 프롬프트의 구조를 명시적으로 정의하기 위해 PML을 도입하여, 프롬프트 모듈의 위치와 내용을 체계적으로 관리합다.
• *위치 정확성 보장: PML을 통해 각 모듈의 위치 ID를 부여하여, 어텐션 상태 재사용 시 위치 정확성을 유지합다. 

---

📊 실험 결과

• *성능 향상: Prompt Cache를 활용하여 GPU 기반 추론에서 최대 8배, CPU 기반 추론에서 최대 60배의 Time-to-First-Token(TTFT) 개선을 달성하였다.
• *정확도 유지: 어텐션 상태 재사용에도 불구하고, 출력 정확도에는 손실이 없음을 확인

---

✅ 장점 & ❌ 단점

✅ 장점
추론 지연 시간 단축	자주 사용되는 프롬프트 세그먼트의 어텐션 상태를 재사용하여 추론 속도를 크게 향상
계산 효율성 증대	중복 계산을 방지하여 계산 자원을 절약하고, 시스템의 전반적인 효율성을높임
모델 수정 불요	기존 모델의 파라미터를 변경하지 않고도 성능 개선을 달성

 

❌ 단점
초기 설정 복잡성	프롬프트 모듈의 정의와 PML 적용 등 초기 설정 과정이 복잡
메모리 사용 증가	어텐션 상태를 저장하기 위한 추가적인 메모리 공간이필요
동적 프롬프트 리 제한	새로운 프롬프트 세그먼트가 등장할 경우, 해당 세그먼트의 어텐션 상태를 새롭게 계산하고 저장

 

Prompt Cache는 대규모 언어 모델 추론 시 자주 사용되는 프롬프트 세그먼트의 어텐션 상태를 미리 계산하여 저장하고 재사용함으로써, 추론 지연 시간을 획기적으로 단축하는 기법입니다

 

 

 

---

RAGCache 구현

 

7. RAGCache 전체적인 구조

Cache에서는 Replacement 기법이 제일 핵심. 

Cache메모리에 자주 사용되는 프로세스가 오래 상주하고있어야만 메인 메모리를 거치지 않고 Cache에서 빠르게 엑세스하여 이득을 볼 수있기 때문.

해당 논문에서는 Prefix-awareGreedy-Dual-Size-Frequency(PGDSF) replacement policy가 적용됨.

만약 D1, D2, D3의 3개의 추출된 문서들이 있다고 가장해보자, 이때, D1-> D3순으로 입력된 컨텍스트가 하나, D2-> D3 순으로 입력된 컨텍스트가 하나있다고 해보자. D3가 두번 반복해서 사용되었지만, D3는 앞에 입력된 내용이 달라짐ㅂ에 따라서 똑같은 D3이더라도 내용이 달라지기 때문에 Cache를 적용할 수 없다. 

BUT, D1-> D2, D1->D3라면? D1이 앞에서 반복해서 사용되고 있기 때문에 D1을 cache에 저장해서 재사용하는 것이 가능

PGDSF는 이러한 입력의 순서에 따라서 Cache를 재사용할 수있는여부가 달라지는것을 반영한 Replacement방법이다.

 

8.Knwledge Tree 구조

 

위의 그림에서 트리의 루트 노드인 S는 공통적으로 계속 사용되는 시스템 프롬프트의 States를 나타낸다. 순서상 Parent위치에 있는 자주 호출될 확률이 높은 Cache들은 전송속도가 훨씬 빠른 VRAM 메모리에 Cache를 저장하는것을 볼 수 있고, 비교적 우선순위가 낮은 노드들은 RAM메모리에 Cache가 저장되는 것을 확인할 수 있다. 

-> 우선순위를 정하는 기준은 무엇인가?

 

Cache의 우선순위를 계산하는 방법

 

→ Clock은 해당 Cache가 최근에 호출된 시간을 나타내며, 더 최근에 호출이 되었을수록 높은 우선순위를 가짐.

     Frequency는 얼마나 자주 호출되는지를 나타내는 호출빈도이며, Size는 토큰의 개수

 

Cost는 토큰당 얼만큼의 연산 시간이 드는지 나타내는 수식

 

→ Eviction(축출) 되는 토큰은 위의 수식의 Priority가 가장 낮은 Cache를 우선적으로 쳐내는데, GPU Memory에 상주하는 Cache가 꽉차서 교체를 해야될 때는 GPU 메모리에 상주하는 Leaf Node중 가장 우선 순위가 낮은 Node를 메인 메모리로 옮기고 새로운 노드를 VRAM에 추가한다. 그리고 메인 메모리에서도 가장 우선순위가 낮은 Leaf 노드를 탈락시키고 새로운 노드를 채워 넣는 방식으로 설계된다.

 

10. Cache-Aware Reordering 방법 예시

 

다음은 Cache-Aware Reordering 방식이다. 이부분은 입력들이 여러개가 동시에 입력되었을 때, Cache를 적용하기에 가장 유리한 방식으로 입력의 실행 순서를 바꾸는 방법이다.

위의 그림(10)에서 왼쪽은 이미 Cache에 저장된 States의 토큰 길이가 하나는 2이고 하나는1이고, 새롭게 컴퓨팅 해야한 되는 토큰의 길이가 각각 2인 상황이다. 

 

Q1을 먼저 실행하면 2의 State를 꺼내와서 쓰고, 2를 연산한 다음에 Cache에 저장하려고 하는데 공간이 꽉찼을 때, 1을 Cache에 Eviction하고 채워 넣었다고 해보자.

→ 1은 Cache에서 탈락되었기 때문에 새로운 연산을 하고 뒤따라오는 2의 길이 입력도 연산을 하면 총 5의 길이만큼 compute가 수행된것이다. 

 

Q2를 먼저 실행하고 Q1을 실행하면 ?? 

→ 1을 Cache에서 꺼내쓰고 2를 연산한 다음에 2를 저장하기 위해서 저장되어 있는 Q1의 Cache를 Eviction하고 저장한다. 그러면 Q1의 Cache에 있던 2의 입력을 다시 Compute하고 뒤따라오는 2를 Compute하면 총 6의 연산을 수행하므로 순서만 바뀌었을 뿐인데 연산량이 증가함

 

오른쪽은 저장할 수 있는 Cache의 용량이 총 5일 때를 가정함.

Q2를 먼저 실행하면 길이 1의 입력을 Compute하고 Eviction 없이 Cache에 저장할 수 있다. 하지만 Q1이 먼저 실행되면 Q2의 길이 2짜리 Cache를 Eviction 해야하기 떄문에 2의 연산이 추가로 적용됨

 

이처럼 Cache-Aware Reordering은 컴퓨팅 해야하는 길이가 같을때는 캐싱된 입력의 길이가 긴 입력을 우선적으로, 캐싱된 입력의 길이가 같을 때는, 컴퓨팅 해야하는 입력의 길이가 짧은 것을 우선적으로 실행하도록 순서를 조정하는 방법

 

11. Speculative Pipelining 구조

→  해당 방법 (11)은 Retrieval 과정에 필요한 시간이 많이 길어질 때, Generation을 위한 대기 시간이 길어지면서 생기는 병목 현상을 최소화하기 위한 방법이다. vector DB를 이용한 Retrieval 과정은 일반적으로 IVF나 HNSW와 같은 알고리즘을 통해서 이루어짐.

이때, vector 탐색을 한번에 처리하게 하지 않고 작은 타임 슬라이드로 나누어서 단계별로 진행하도록 하면 vector 탐색의 중간 결과를 얻을 수 있다. →  최종 단계의 결과값이 정확하겠지만, 때로는 중간 값에서 구해진 결과랑 최종 결과값이 일치할 수도 있다.

이에 Speculative pipelining은 중간 결과가 나오면 미리 generation을 수행하고 있다가, 최종 결과값과 수행하고 있던 입력이 일치하면 계속해서 진행하고, 만약 다른 최종 결과값이 나오면 달라진 입력으로 다시 생성과정을 수행하도록 함

 

 

 

 

 

 

 

 

 

 

RAG 구현 참고 : https://delosycho.co.kr/%EB%85%BC%EB%AC%B8-%EC%A0%95%EB%A6%AC-ragcache-efficient-knowledge-caching-for-retrieval-augmented-generation/ 

논문 정리: RAGCache: Efficient Knowledge Caching for Retrieval-Augmented Generation - delosycho