Autoregressive next token prediction and KV Cache in transformers

개요

LLM의 토큰 생성 속도를 높이는 최적화 기법인 KV 캐시는 모델의 컨텍스트 윈도우 역할을 하며, 이전 토큰의 재계산 없이 생성 속도를 향상시키고 운영 비용을 절감합니다.

주요 내용

* Autoregressive 모델의 기본 구조: 텍스트 프롬프트는 토크나이저를 통해 정수 ID로 분할되고, 각 ID는 임베딩 벡터로 변환됩니다. 이 벡터들은 여러 개의 디코더 블록을 거치면서 변환되고, 최종적으로 마지막 디코더 블록의 출력이 다음 토큰을 예측하는 데 사용됩니다.
* Prefill 단계: 언어 모델이 첫 번째 새 토큰을 생성하기 전에 전체 입력 시퀀스를 한 번의 병렬 순방향 패스로 처리합니다. 이 단계는 첫 번째 예측 토큰을 생성하고, 이후 디코드 단계를 효율적으로 만들기 위해 KV 캐시를 채우는 역할을 합니다.
* KV 캐시의 작동 방식 (Decode 단계): Prefill 단계가 완료되면 모델은 디코드 모드로 전환됩니다. 이후 생성되는 각 토큰은 이전 토큰들을 재처리하는 대신, 저장된 KV 캐시를 참조하여 단일 토큰 단위로 처리됩니다. 새로운 토큰의 K와 V 값은 캐시에 추가되며, 어텐션 메커니즘은 이 캐시를 활용하여 컨텍스트를 파악합니다.
* KV 캐시의 이점: KV 캐시를 사용하면 각 토큰 생성 시 전체 시퀀스를 다시 처리할 필요 없이, 이전 토큰 정보를 효율적으로 재사용할 수 있습니다. 이로 인해 단순 순방향 패스 대비 계산량이 크게 줄어들어, O(n³)에서 O(n²)의 전체 생성 시간 복잡도로 개선됩니다.
* KV 캐시와 컨텍스트 윈도우: Prefill 후 KV 캐시는 모델의 컨텍스트 윈도우 역할을 하며, 이전 토큰들의 정보는 KV 캐시에 저장된 K 및 V 텐서로만 유지됩니다. 모델이 광고하는 컨텍스트 윈도우 크기는 GPU 메모리에 저장할 수 있는 KV 캐시의 최대 행 수를 의미합니다.
* KV 캐시의 메모리 비용: KV 캐시는 각 레이어, 각 토큰마다 K와 V 벡터를 저장하므로 메모리 사용량이 컨텍스트 길이에 따라 선형적으로 증가합니다. 이는 GPU 용량, 메모리 대역폭, 추론 처리량에 병목 현상을 유발하며, 캐시 압축 및 하이브리드 메모리 솔루션과 같은 효율적인 메모리 관리 기법에 대한 연구가 진행 중입니다.

시사점

KV 캐시는 LLM에서 장문의 텍스트를 효율적으로 생성하기 위한 핵심 최적화 기법으로, 생성 속도 향상 및 비용 절감에 기여하지만, 동시에 메모리 사용량 증가라는 현실적인 제약을 수반합니다.