서비스로 들어오는 프롬프트는 서로 많이 겹친다. 모든 요청 앞에 같은 시스템 프롬프트가 붙고, 여러 턴짜리 대화는 매 턴마다 이전 대화 전체를 다시 보내며, 에이전트는 같은 도구 설명과 지시문을 호출마다 되풀이한다. 겹치는 앞부분의 KV 캐시는 언제 계산해도 같으므로 한 번 계산해 두면 다음 요청은 그 부분의 프리필을 건너뛸 수 있다. 이렇게 앞부분의 KV 캐시를 재사용하는 방법을 프리픽스 캐시(prefix caching)라고 부른다.
재사용할 수 있는 조건
KV 캐시를 재사용하려면 앞부분이 토큰 단위로 정확히 같아야 한다. 어텐션은 각 토큰이 자기보다 앞선 토큰만 보도록 계산하므로 위치 i까지의 KV 값을 정하는 것은 0부터 i까지의 토큰뿐이다. 두 요청의 첫 1,000토큰이 같으면 그 1,000토큰의 KV 값도 같다. 반대로 1,001번째 토큰부터 달라지면 그 뒤의 KV 값은 앞부분이 같아도 재사용할 수 없다.
그래서 재사용은 항상 앞에서부터, 처음 달라지는 토큰 직전까지만 된다. 가운데 같은 문단이 있어도 그 앞이 다르면 재사용할 수 없다. 같은 모델, 같은 가중치, 같은 KV 정밀도라는 조건도 함께 붙는다.
얻는 것
프리필을 건너뛴 만큼 TTFT가 줄고 그 계산에 쓰였을 연산기가 다른 요청에 돌아간다. 효과는 공유하는 앞부분이 길고 요청마다 새로 붙는 부분이 짧을수록 크다.
| 부하 | 공유되는 앞부분 | 효과 |
|---|---|---|
| 긴 시스템 프롬프트를 쓰는 챗봇 | 시스템 프롬프트 | 시스템 프롬프트가 수천 토큰이면 매 요청의 프리필 대부분을 건너뛴다 |
| 여러 턴 대화 | 이전 턴까지의 대화 | 턴이 쌓일수록 공유 부분이 길어져 새 턴의 프리필이 새 메시지만큼으로 줄어든다 |
| 같은 문서에 여러 질문 | 문서 본문 | 문서를 한 번만 프리필하고 질문마다 짧은 프리필만 한다 |
| 에이전트 루프 | 도구 정의, 지시문, 이전 단계 기록 | 단계마다 앞부분이 그대로 이어지므로 적중률이 높다 |
메모리 쪽 이득도 있다. 같은 앞부분을 여러 요청이 동시에 쓰고 있으면 그 블록은 한 벌만 저장한다. 시스템 프롬프트 2,000토큰을 요청 100개가 공유하면 블록 테이블은 100개지만 KV 블록은 한 벌, 약 260MB다(토큰당 128KiB 기준).
해시 블록 방식
vLLM의 자동 프리픽스 캐시는 페이지 단위 KV 관리 위에 올라간다. 각 KV 블록을 그 블록 안의 토큰과 그 블록 앞의 모든 토큰으로 해시해 식별한다. 앞부분이 다르면 같은 토큰을 담은 블록이라도 해시가 다르므로 해시가 같다는 것은 그 블록까지의 앞부분 전체가 같다는 뜻이다.
새 요청이 들어오면 엔진은 프롬프트를 블록 크기로 나눠 앞에서부터 해시를 계산하고, 해시 테이블에 같은 블록이 있는지 찾는다. 처음으로 못 찾는 블록 직전까지는 기존 블록을 가리키도록 블록 테이블을 채우고, 그 뒤부터 프리필한다. 블록 단위로 맞추므로 마지막 블록을 다 채우지 못한 부분은 재사용되지 않는다.
요청이 끝나도 블록을 바로 지우지 않고 해시 테이블에 남겨 둔다. 아무 요청도 쓰지 않는 블록은 빈 블록처럼 다른 요청에 내줄 수 있는 후보가 되고, 메모리가 모자랄 때 가장 오래 쓰이지 않은 것부터 내보낸다. 그 전에 같은 앞부분을 가진 요청이 오면 그대로 재사용한다.
RadixAttention
SGLang(NeurIPS 2024)의 RadixAttention은 캐시된 시퀀스들을 기수 트리(radix tree)로 관리한다. 트리의 간선이 토큰 시퀀스이고 루트에서 어떤 노드까지의 경로가 하나의 앞부분이다. 새 요청이 오면 트리를 따라 내려가며 가장 길게 일치하는 앞부분을 찾고, 그 지점부터 프리필한 뒤 새 부분을 가지로 붙인다.
트리 구조는 앞부분끼리의 포함 관계를 그대로 드러내므로 축출과 스케줄링에 쓰기 좋다. 축출은 지금 쓰이지 않는 잎 노드부터 오래된 순으로 한다. 공통 앞부분은 여러 잎이 공유하는 안쪽 노드라서 잎들이 다 빠지기 전에는 남는다. 스케줄러는 대기 중인 요청 가운데 캐시와 일치하는 앞부분이 긴 요청을 먼저 처리해 적중률을 높일 수 있다.
해시 블록 방식과 기수 트리 방식은 같은 문제를 다른 자료구조로 푼다. 사용자 입장에서 차이는 일치 단위(블록 대 토큰), 축출 정책, 스케줄링이 캐시를 얼마나 고려하는지에서 나온다.
적중률을 높이는 프롬프트 구성
프리픽스 캐시의 효과는 엔진보다 프롬프트를 만드는 쪽이 더 크게 좌우한다. 재사용은 앞에서부터만 되므로 요청마다 바뀌는 내용이 앞에 오면 그 뒤 전체가 캐시를 놓친다.
- 바뀌지 않는 것을 앞에 둔다. 시스템 프롬프트, 도구 정의, 고정 지시문, 예시를 맨 앞에 두고 사용자 입력과 검색 결과는 뒤에 둔다
- 앞부분에 요청마다 다른 값을 넣지 않는다. 시스템 프롬프트 첫 줄의 현재 시각, 요청 ID, 사용자 이름은 그 뒤 전체의 재사용을 막는다. 꼭 필요하면 뒤로 옮긴다
- 토큰화까지 같게 만든다. 공백 하나, 줄바꿈 하나, 채팅 템플릿의 미세한 차이도 토큰을 바꾼다. 같은 내용을 같은 코드 경로로 직렬화한다
- 대화 이력은 덧붙인다. 이전 턴을 요약하거나 잘라 내면 앞부분이 바뀌어 재사용이 끊긴다. 문맥 상한 때문에 잘라야 한다면 자르는 시점을 드물게 만든다
라우팅과의 관계
캐시는 GPU마다 따로 있다. 서버가 여러 대이고 로드 밸런서가 요청을 무작위로 나눠 주면 같은 대화의 다음 턴이 이전 턴을 처리한 서버로 갈 확률은 서버 수에 반비례한다. 서버 8대면 7/8의 확률로 캐시를 놓친다.
그래서 프리픽스 캐시를 쓰는 서비스는 라우팅도 캐시를 보고 한다. 대화 ID나 사용자 ID로 같은 서버에 붙이는 방법이 가장 단순하고, 프롬프트 앞부분의 해시로 서버를 고르거나 각 서버의 캐시 상태를 알고 있는 라우터를 두는 방법도 있다. 캐시만 보고 보내면 특정 서버에 부하가 몰릴 수 있으므로 서버별 대기 중인 요청 수와 함께 저울질한다.
주의할 점
프리픽스 캐시는 출력을 바꾸지 않는 최적화다. 그래도 효과를 재는 방법과 보안, 테스트 조건에서 챙길 것이 생긴다.
- 적중률을 지표로 남긴다. 캐시가 켜져 있어도 프롬프트 구성 때문에 적중률이 낮으면 이득이 없다. 엔진이 내보내는 프리픽스 캐시 적중 지표를 TTFT와 함께 본다
- 다른 사용자 간 공유를 검토한다. 같은 앞부분을 가진 요청의 TTFT가 짧아지는 것을 관찰하면, 다른 사용자가 어떤 프롬프트를 보냈는지 짐작할 여지가 생긴다. 사용자나 테넌트별로 캐시를 나눠야 하는지 보안 요구사항을 확인한다. vLLM은 해시에 별도의 값(cache salt)을 섞어 캐시 공간을 나누는 방법을 제공한다
- 부하 테스트에 반영한다. 테스트 요청이 모두 같은 프롬프트면 적중률이 비현실적으로 높게 나오고, 모두 무작위면 반대다. 운영 로그의 앞부분 공유 정도를 재현한다