제8장 — 차세대 KV 캐시 관리

LLM Primer VI: AI 시스템 확장하기를 챕터별로 걷는 워크스루의 여덟 번째 글입니다. 운영체제의 페이징 통찰을 추론 엔진으로 옮겨 와 KV 캐시를 예약된 바이트 슬래브에서 공유·축출·프리픽스 캐시가 가능한 자원으로 바꾸는 장입니다.

이 장이 존재하는 이유

제7장은 연속 배칭에 빚을 남긴 채 끝났습니다. 시퀀스가 매 이터레이션마다 들어오고 나가고, 소박한 배치는 각 슬롯에 최대치 크기의 슬래브를 주며, 그 대부분이 낭비되고, 배칭의 처리량 승리가 부분적으로 되돌아갑니다. 이 빚은 내부 단편화 — 1960년대에 운영체제가 페이징으로 해결한 바로 그 실패 모드 — 입니다. 제8장은 그 해법을 LLM 서빙에 적용하는 이야기입니다 — KV 캐시를 작은 물리 블록으로 나누고, 페이지 테이블로 논리 토큰 위치와 분리하며, 축출과 캐싱 정책이 블록을 회수하거나 시퀀스 사이에 공유하게 합니다. PagedAttention은 기초가 되는 수이고, H2O와 InfiniGen은 축출 정책이며, 프리픽스 캐시는 프로덕션 클러스터가 수백만 에이전트 요청을 열두 대 GPU로 서빙하게 해 주는 기법입니다.

8.1 PagedAttention은 KV 캐시의 가상 메모리이다

vLLM의 PagedAttention 논문(2023)은 OS의 통찰을 엔진으로 곧장 들어 올립니다. KV 캐시는 고정 크기의 블록 — 보통 각 16토큰 — 으로 나뉘어 평평한 물리 풀에 놓입니다. 시퀀스는 블록 테이블 — 논리 위치를 물리 블록 ID로 매핑하는 포인터 배열 — 로 표현됩니다. 어텐션 커널은 블록 테이블을 추가 입력으로 받아, 연속 슬라이싱이 아니라 간접 참조로 키와 값을 모읍니다. Hopper의 L2가 임의 접근 패턴을 충분히 잘 흡수해서, 커널은 슬래브 기반 버전의 몇 퍼센트 안에서 돕니다. 이득은 큽니다. 내부 단편화가 60~80퍼센트에서 대략 4퍼센트(시퀀스당 부분 꼬리 블록 하나)로 떨어져 사용 가능한 동시성이 2~4배 올라갑니다. 참조 카운트 블록 공유는 복잡한 샘플링을 거의 공짜로 만듭니다 — 2,000토큰 프롬프트에 대한 best-of-8이 16,800 토큰 블록에서 2,800 토큰 블록으로 떨어집니다 — 그리고 이것이 프리픽스 캐시가 짓는 기판입니다.

8.2 H2O와 InfiniGen은 중요하지 않은 토큰을 축출한다

PagedAttention은 단편화를 해결하지만, 어떤 배치로도 VRAM을 넘어가는 컨텍스트는 해결하지 못합니다. 200,000토큰 Llama-3-70B 컨텍스트는 가중치와 나란히 60GB의 KV를 요구합니다. H2O(Heavy Hitter Oracle)는 디코딩 중 어텐션 가중치가 작은 원천 위치 집합 — 최근성 토큰, 맨 처음의 어텐션 싱크 토큰, 콘텐츠에 관련된 희소한 히트 — 에 집중하며, 대부분의 이력 위치는 본질적으로 영에 가까운 가중치를 받는다고 관찰합니다. 엔진은 위치별로 누적 어텐션 점수를 추적하고, 시퀀스의 KV 예산이 상한에 다가가면 최근 토큰과 싱크 토큰의 보장된 유지 창을 제외한 가장 낮은 점수 위치를 축출합니다. 절감은 크고, 대가는 영구성입니다 — 이후 쿼리가 축출된 위치를 필요로 하면 엔진은 회복할 수 없습니다. InfiniGen은 동적이고 회복 가능한 선택으로 이 절충을 다듬습니다 — 토큰을 아예 버리는 대신 그 KV를 CPU 메모리로 오프로드하고, 어텐션이 다시 그 위치에 집중되기 시작하면 다시 페이징해 옵니다. 올바른 축출 정책은 워크로드가 자기 긴 이력을 얼마나 다시 조회하는지에 달려 있습니다 — 에이전트 워크로드는 영구 축출을 벌주고 InfiniGen 스타일의 회복을 보상합니다.

8.3 프리픽스 캐시는 PagedAttention이 잠금 해제하는 가장 강력한 레버

실제 트래픽에서 대부분 프롬프트의 첫 수천 토큰은 요청들 사이에서 동일합니다. 챗 서비스는 모든 대화에 같은 시스템 프롬프트를 재사용합니다. RAG 서비스는 같은 것을 검색한 모든 사용자에게 같은 검색 결과를 프롬프트에 붙입니다. 에이전트는 매 스텝마다 같은 도구 설명과 추론 스캐폴드를 주입합니다. PagedAttention은 공유를 기계적으로 만듭니다 — 프롬프트를 블록 크기 조각으로 해시하고, 그 해시가 연산으로 채워진 블록의 글로벌 캐시에 있으면 새 요청의 블록 테이블이 캐시된 블록을 가리키게 하고 그 프리픽스의 프리필을 완전히 건너뜁니다. 없으면 프리필이 돌고 결과 블록이 등록됩니다. 프로덕션 히트율은 극적입니다 — 챗 서비스 시스템 프롬프트에서 99퍼센트 이상, RAG 검색 의존 프리픽스에서 30~60퍼센트, 에이전트 스캐폴드에서 1.0에 가까움. SGLang의 RadixAttention은 블록 정렬되지 않은 어떤 길이의 공유 프리픽스도 키로 삼는 라덱스 트리로 이 아이디어를 더 밀고 갑니다. 프리픽스 캐시는 예산을 초과하던 서빙 클러스터를 예산 안으로 되돌려 놓는 데 가장 자주 성공하는 단일 기법입니다.

제8장이 준비하는 것

페이징, 축출, 프리픽스 캐시는 토큰당 KV 발자국을 줄였고 엔진의 메모리 행동을 높은 동시성에서 감당 가능하게 만들었습니다. 그것들이 다루지 않는 것은 디코딩의 근본적 순차 의존입니다 — 활성 슬롯이 몇 개든 시퀀스당 이터레이션당 한 토큰. 제9장은 그 구속 조건을 다룹니다 — 값싼 초안으로 여러 토큰을 미리 예측하고, 그 추측을 값비싼 한 번의 포워드 패스로 검증해, 사용자 지각 지연에 가장 중요한 시퀀스에 대해 이터레이션당 한 토큰의 바닥을 깨는 기법 계열, 투기적 디코딩입니다.

다음 → 제9장: 투기적 디코딩. 초안, 검증, 그리고 언제 투기가 이득인지의 산수.