롱컨텍스트 LLM 추론 최적화 완전 정리: KV 캐시 병목부터 MagicPIG 샘플링 기법까지

핵심 메시지

• 긴 문맥을 다루는 LLM의 속도 병목은 계산량이 아니라, 토큰 하나를 만들 때마다 GPU 메모리에서 KV 캐시를 읽어와야 하는 데이터 이동에 있다.
• 중요도가 낮은 토큰을 버리는 정적 KV 압축은 요약이나 단순 질의응답에서는 통하지만, 문맥 전체를 엮어야 답이 나오는 문제에서는 무너진다.
• 압축한 KV 캐시로 초안을 만들고 원본 모델이 검증하는 추측 디코딩은 출력을 바꾸지 않으면서 99%에 가까운 수용률을 보였다.
• MagicPIG는 상위 k개를 고르는 대신 해시 기반 샘플링으로 어텐션을 추정하고 무거운 조회를 CPU로 넘겨 처리량을 2~4배 끌어올렸다.
• 논리 깊이를 조절해 만든 새 평가 데이터에서, 8K에서 128K로 컨텍스트를 늘린 모델들은 사실상 암기 수준 능력만 유지했다.

쉽게 이해하기

카네기멜런대 고급 자연어처리 수업(CS 11-711, 2024년 가을)의 초청 강의에서 베이디 첸(Beidi Chen) 교수는 긴 문맥을 처리하는 언어 모델의 효율 문제를 다뤘다. 불과 2년 전만 해도 기본 컨텍스트 길이는 4K~8K 토큰 수준이었지만, 지금은 오픈소스 모델도 128K가 기본이고 제미나이 프로 1.5처럼 100만 토큰을 내세우는 모델도 있다. 문서를 통째로 넣는 비서, 코드베이스 전체를 참조하는 개발 도구, AI 에이전트가 모두 이 긴 문맥에 기대고 있다.

문제는 문맥이 길어질수록 KV 캐시가 시퀀스 길이와 배치 크기에 비례해 커진다는 점이다. 강의에서는 모델 가중치와 캐시를 합치면 수백 기가바이트에 이르는 사례를 들며, 80GB짜리 최신 GPU 한 장에 올리는 것조차 만만치 않다고 지적했다. 게다가 이 캐시는 토큰을 하나 생성할 때마다 매번 다시 읽어야 해서, 성능이 연산 속도가 아니라 메모리 대역폭에 묶인다. GPU가 자랑하는 높은 연산 성능이 놀고 있는 셈이라, 새로운 하드웨어에게는 오히려 기회가 된다는 설명이 이어졌다.

그래서 나온 첫 번째 접근이 정적 압축이다. 지금까지 받은 어텐션 점수를 누적해 두고, 메모리가 부족해지면 점수가 낮은 토큰을 버리는 방식이다. 이 과정에서 연구진은 문장 맨 앞 한두 토큰이 의미와 무관하게 항상 큰 어텐션을 받는 '어텐션 싱크' 현상을 발견했다. 첫 토큰을 줄바꿈 문자로 바꿔도 현상이 유지되는 것으로 보아, 소프트맥스 구조상 어디도 보지 않을 수 없는 상황에서 모델이 첫 토큰을 임시 대피소처럼 쓰는 편향으로 해석됐다.

그러나 버린 정보는 되돌아오지 않는다. 여러 변수의 관계를 단계적으로 따라가야 답이 나오는 문제에서는 압축 모델이 어긋나기 시작했고, 심지어 어텐션을 모두 계산한 뒤 상위 k개만 고르는 '이상적인' 상한선조차 실패했다. 대안으로 제시된 것이 추측 디코딩이다. 작은 모델 대신 1% 수준으로 압축한 KV 캐시가 초안을 만들고 전체 캐시를 쓴 모델이 검증하는 구조인데, 출력 분포가 보존되면서도 수용률이 99%에 달했다. 배치가 커지면 이점이 사라지는 기존 추측 디코딩과 달리, KV 캐시는 예시마다 따로 존재해 항상 메모리 병목 상태라 배치가 클수록 오히려 유리했다.

강의 후반의 MagicPIG는 여기서 한 걸음 더 나간다. 상위 k개 선택이 사실은 편향된 추정이라는 점을 동물원 비유로 설명한 뒤, 어텐션 점수에 비례한 확률로 표본을 뽑으면 편향 없는 추정이 된다고 정리했다. 이 샘플링을 1990년대 해싱 기법으로 구현해, 해시 계산은 GPU가 맡고 메모리를 많이 쓰는 해시 테이블 조회는 CPU가 맡는 분업 구조를 만들었다. CPU는 연산이 100배 느리지만 메모리 대역폭은 10배 차이에 그친다는 점을 파고든 설계로, 처리량은 2~4배 개선됐다.

주요 인사이트

• 병목이 연산이 아니라 데이터 이동이라면, 최적화의 승부처는 커널 튜닝이 아니라 무엇을 어디에 두고 언제 옮길지를 정하는 설계로 옮겨간다.
• 손실 압축은 되돌릴 수 없다. 지금 필요 없어 보이는 토큰이 나중에 질문의 핵심이 될 수 있으므로, 출력이 보존되는 무손실 경로를 함께 두는 편이 안전하다.
• 상위 k개 선택은 직관적이지만 통계적으로는 편향된 추정이다. 확률에 비례한 표본 추출이 이론적으로 더 정확하고, 실제 벤치마크에서도 상위 k 방식을 앞섰다.
• 검색으로 뚫리는 벤치마크는 긴 문맥 능력을 측정하지 못한다. 문맥 전체를 엮어야만 풀리는 관계형 문제여야 롱컨텍스트의 값어치가 드러난다.
• 컨텍스트 창에 적힌 숫자와 실제로 추론할 수 있는 범위는 다르다. 사전학습 길이를 넘겨 확장한 구간에서는 암기 이상의 능력이 잘 따라오지 않았다.

자주 묻는 질문

관련 AI 소식