Mastering Long Contexts in LLMs with KVPress

🖼️ 인포그래픽

🖼️ 4컷 인포그래픽

💡 한 줄 요약

KVPress는 긴 문맥에서 선형으로 증가하는 키·값 캐시를 중요도 기반으로 압축해 메모리 사용량을 줄이고 생성 속도를 높이는 모듈형 도구다.

📌 핵심 요약

• 긴 문맥은 문맥 내 검색·학습과 장기 추론을 가능하게 하지만, 문맥 토큰 수에 비례해 커지는 키·값 캐시 때문에 막대한 메모리를 요구한다.
• Llama 3-70B가 100만 토큰을 bfloat16으로 처리할 때 키·값 캐시는 약 327.6GB, 모델 가중치까지 포함한 전체 메모리는 약 470GB에 이른다.
• KVPress는 각 어텐션 헤드에서 중요도가 낮은 키·값 쌍을 제거하는 여러 압축 알고리즘을 제공하며, 순방향 훅을 통해 모델의 어텐션 계층에 결합된다.
• Llama 3.1 8B의 12만 8천 토큰 문맥에서 50% 압축을 적용한 사례는 최대 메모리를 45GB에서 37GB로 줄이고 A100의 디코딩 속도를 초당 11토큰에서 17토큰으로 높였다.
• 압축률이 높아질수록 정확도가 저하될 수 있으므로, KVPress는 여러 기법을 표준 장문 문맥 데이터셋에서 비교해 메모리 효율과 성능 사이의 균형을 평가하도록 지원한다.

🧩 주요 포인트

1. 긴 문맥은 문맥 내 검색·학습과 장기 추론을 가능하게 하지만, 문맥 토큰 수에 비례해 커지는 키·값 캐시 때문에 막대한 메모리를 요구한다.
2. Llama 3-70B가 100만 토큰을 bfloat16으로 처리할 때 키·값 캐시는 약 327.6GB, 모델 가중치까지 포함한 전체 메모리는 약 470GB에 이른다.
3. KVPress는 각 어텐션 헤드에서 중요도가 낮은 키·값 쌍을 제거하는 여러 압축 알고리즘을 제공하며, 순방향 훅을 통해 모델의 어텐션 계층에 결합된다.
4. Llama 3.1 8B의 12만 8천 토큰 문맥에서 50% 압축을 적용한 사례는 최대 메모리를 45GB에서 37GB로 줄이고 A100의 디코딩 속도를 초당 11토큰에서 17토큰으로 높였다.
5. 압축률이 높아질수록 정확도가 저하될 수 있으므로, KVPress는 여러 기법을 표준 장문 문맥 데이터셋에서 비교해 메모리 효율과 성능 사이의 균형을 평가하도록 지원한다.

🧠 상세 정리

1. 긴 문맥이 제공하는 가능성과 메모리 문제

대규모 언어 모델의 문맥 창은 한 번의 요청에서 처리할 수 있는 최대 토큰 수를 뜻하며, 모델이 발전하면서 그 길이도 빠르게 늘고 있다. 긴 문맥은 하나의 질의 안에서 방대한 텍스트를 참조하는 문맥 내 검색, 같은 세션에 제시된 사례에 맞춰 행동을 조정하는 문맥 내 학습, 매우 긴 추론 과정을 끊김 없이 처리하는 장기 추론을 가능하게 한다. 그러나 문맥이 길어질수록 이를 보존하는 키·값 캐시가 차지하는 메모리도 함께 증가해 실제 배포의 핵심 병목이 된다. 원문은 Llama 3-70B로 100만 토큰을 처리할 때 키·값 캐시에만 약 330GB가 필요하다는 사례를 들며, 긴 문맥의 이점을 활용하려면 캐시 자체를 더 효율적으로 관리해야 한다고 문제를 설정한다.

2. 키·값 캐시가 자동회귀 생성을 가속하는 원리

자동회귀 모델은 텍스트를 한 토큰씩 생성하며, 새로운 토큰을 예측할 때 앞서 등장한 모든 토큰의 정보를 이용한다. 예를 들어 1000번째 토큰을 생성할 때는 1번째부터 999번째 토큰까지의 표현이 필요하고, 1001번째 토큰을 생성할 때도 기존 999개 토큰과 새로 생성된 1000번째 토큰을 다시 고려해야 한다. 매 단계마다 이전 토큰 전체의 어텐션 계산을 반복하면 문장이 길어질수록 중복 연산이 커지므로, 키·값 캐시는 각 어텐션 계층에서 계산된 키와 값의 중간 결과를 저장한다. 이후 토큰 생성에서는 저장된 결과를 재사용함으로써 같은 정보를 계속 다시 계산하지 않아도 되며, 이 때문에 키·값 캐시는 효율적인 생성에 필수적인 장치가 된다.

3. 문맥 길이에 선형 비례하는 캐시 규모

키·값 캐시의 크기는 키와 값을 모두 저장하는 계수 2, 수치 정밀도, 계층 수, 어텐션 헤드 수, 헤드 차원, 문맥 토큰 수의 곱으로 결정된다. 따라서 모델 구조와 정밀도가 같다면 캐시 메모리는 토큰 수에 선형으로 비례하며, 긴 문맥에서는 모델 가중치보다도 큰 비중을 차지할 수 있다. 원문이 제시한 Llama 3-70B 사례에서는 bfloat16 2바이트, 80개 계층, 8개 키·값 헤드, 헤드 차원 128, 문맥 100만 토큰을 적용했을 때 캐시 크기가 327.6GB로 계산된다. 700억 개 매개변수의 모델 가중치도 bfloat16에서 약 140GB가 필요하므로 전체 요구량은 약 470GB이며, 이 가운데 키·값 캐시가 약 70%를 차지한다.

4. 압축 기법을 모은 KVPress의 역할

KVPress는 긴 문맥에서 발생하는 키·값 캐시의 메모리 문제를 해결하기 위해 엔비디아가 개발한 파이썬 도구 모음이다. 최신 캐시 압축 기법을 하나의 모듈형 체계로 제공하며, 저장 정밀도를 낮춰 공식의 정밀도 항을 줄이는 트랜스포머 라이브러리의 키·값 캐시 양자화와도 함께 사용할 수 있다. 압축 연구자는 기존 방법의 구조를 이해하고 새로운 방법을 추가할 수 있고, 개발자는 개별 논문의 구현을 처음부터 재현하지 않고도 실제 생성 파이프라인에 여러 기법을 빠르게 적용할 수 있다. 즉 KVPress의 핵심 가치는 단일 압축 알고리즘 하나가 아니라, 다양한 방법을 동일한 인터페이스에서 실험·확장·배포할 수 있도록 만든 통합 기반에 있다.

5. 프레스의 중요도 평가와 캐시 가지치기

KVPress에서 실제 압축을 수행하는 알고리즘은 프레스라고 불리며, 다수의 프레스는 각 어텐션 헤드 안에서 키·값 쌍의 중요도 점수를 계산한 뒤 점수가 낮은 항목을 제거한다. KnormPress는 키 벡터의 노름이 낮은 쌍을 가지치기하고, SnapKVPress는 최근 질의가 부여한 어텐션 가중치가 낮은 쌍을 제거한다. ExpectedAttentionPress는 향후 질의에서 기대되는 어텐션 가중치가 가장 낮은 키·값 쌍을 제거하도록 설계됐으며, 각 프레스의 압축 비율 속성이 제거 강도를 결정한다. 이러한 프레스는 순방향 훅을 통해 모델의 어텐션 계층에 연결되므로 모델 전체 구조를 직접 다시 작성하지 않고도 생성 도중 캐시를 동적으로 압축할 수 있다.

6. 사전 채움 단계 적용과 실제 자원 절감

KVPress는 사용자 문맥을 처음 처리해 캐시를 만드는 사전 채움 단계에 압축을 집중하며, 이 시점은 긴 입력으로 인해 캐시가 가장 크게 형성되는 구간이다. 사용자는 전용 텍스트 생성 파이프라인에 모델과 프레스를 지정하고, 긴 문맥과 그 문맥에 대한 질문을 전달해 압축된 캐시를 기반으로 답을 생성할 수 있다. Llama 3.1 8B의 짧은 입력에서는 약 15GB의 모델 가중치가 메모리 대부분을 차지하지만, 입력이 길어지면 키·값 캐시가 전체 사용량의 주요 부분으로 커진다. 12만 8천 토큰 문맥에 50% 압축을 적용한 사례에서는 최대 메모리가 45GB에서 37GB로 줄었고, 더 작아진 캐시 덕분에 A100에서 디코딩 속도도 초당 11토큰에서 17토큰으로 향상됐다.

7. 표준 데이터셋을 활용한 압축 성능 비교

키·값 캐시 압축 분야에는 서로 다른 중요도 기준과 제거 전략을 사용하는 여러 연구가 존재하며, KVPress는 이미 열두 개가 넘는 프레스를 제공하면서 외부 연구자의 기여도 장려한다. 도구에 포함된 명령줄 인터페이스를 이용하면 RULER, InfiniteBench, Loogle과 같은 표준 장문 문맥 데이터셋에서 각 프레스의 성능을 비교할 수 있다. 원문은 4천 토큰 길이의 RULER 데이터셋에서 아홉 가지 프레스를 여러 압축 비율로 평가한 결과를 제시한다. 해당 실험에서는 AdaKVPress와 KVPress 저자들이 만든 미공개 가지치기 기법인 ExpectedAttentionPress를 결합한 방식이 가장 높은 성능을 기록했으며, 이는 압축률뿐 아니라 어떤 중요도 평가 방식을 선택하고 조합하는지도 결과에 영향을 준다는 점을 보여준다.

8. 정확도와 효율 사이의 균형

KVPress는 긴 문맥의 키·값 캐시를 줄여 메모리 효율과 디코딩 속도를 개선하지만, 압축 비율을 높일수록 모델 정확도가 떨어질 수 있다는 한계가 벤치마크에 나타난다. 따라서 최대한 많이 제거하는 것보다 주어진 데이터와 작업에서 성능 저하를 허용 가능한 수준으로 유지하면서 적절한 압축률과 프레스를 선택하는 것이 중요하다. 원문의 댓글에서는 캐시 축소가 어텐션 연산량에 미치는 영향도 질문됐지만, 저자는 부동소수점 연산 횟수를 직접 측정하지 않았다고 답했으며 대신 전체 생성 시간을 측정한 자료를 안내했다. 저자에 따르면 대부분의 프레스에서 압축 계산은 긴 문맥의 순방향 처리보다 매우 가벼우며, 향후 과제는 정확도 손실을 더 줄이는 압축 알고리즘을 개발하는 것이다.

🧾 핵심 주장 / 시사점

• 긴 문맥 배포의 메모리 병목은 모델 가중치만으로 설명되지 않으며, 100만 토큰 사례에서는 키·값 캐시가 전체 메모리의 약 70%를 차지한다.
• 캐시 압축은 저장 공간만 줄이는 기법이 아니라 이후 어텐션이 다루는 캐시 규모도 줄이므로, 제시된 실험에서는 메모리 절감과 디코딩 속도 향상이 함께 나타났다.
• 압축 효과는 비율 하나로 결정되지 않으며, 키 노름·최근 어텐션·기대 어텐션 등 중요도 평가 기준과 기법의 조합을 표준 데이터셋에서 함께 검증해야 한다.

✅ 액션 아이템

• 긴 문맥에서 키·값 캐시가 문맥 토큰 수에 선형으로 증가해 메모리 부담이 급증하므로, KVPress 적용 대상 토큰 구간을 명확히 정한다.
• Llama 3-70B 100만 토큰의 KV 327.6GB, 전체 470GB 수치를 메모리 상한 기준으로 두고, Llama 3.1 8B 50% 압축의 45GB→37GB와 11→17 토큰/초 개선을 성능 지표로 맞춘다.
• 저중요도 KV 제거를 포함한 KVPress 압축 기법들을 표준 장문 문맥 데이터셋에서 병렬로 시험해 압축률·정확도 저하 트레이드오프를 정량 비교한다.

❓ 열린 질문

• 긴 문맥 처리가 필요한 실제 워크로드에서 KVPress 적용을 시작할 토큰 임계치와 적용 범위를 어떻게 정할 것인가?
• 압축률을 높였을 때 정확도 저하가 허용 가능한지 판단할 임계 지표는 어떤 값으로 설정할 것인가?
• Llama 3.1 8B의 45GB→37GB, 11→17 토큰/초 개선이 우리 하드웨어 환경에서도 동일하게 재현되는지 어떤 검증 절차로 확인할 것인가?