KV Cache by hand

개요

KV Cache는 LLM 추론 시 발생하는 반복적인 계산을 최적화하여 성능을 향상시키는 핵심 기법이며, 계산량을 크게 줄이는 대신 메모리 사용량 증가라는 트레이드오프를 동반합니다.

주요 내용

* KV Cache의 필요성: LLM은 이전 토큰들의 정보를 바탕으로 다음 토큰을 예측하는 자기회귀적 방식으로 텍스트를 생성합니다. Transformer의 어텐션 메커니즘에서는 각 토큰마다 Query(Q), Key(K), Value(V) 벡터를 계산하는데, KV Cache가 없다면 새로운 토큰을 생성할 때마다 이전 모든 토큰의 K와 V를 재계산해야 합니다.
* KV Cache의 작동 방식: KV Cache는 이전 토큰들의 K와 V 벡터를 메모리에 저장해두고, 새로운 토큰의 Q와만 연산하여 어텐션 점수를 계산합니다. Q 벡터는 캐싱되지 않는데, 이는 각 토큰의 Q가 해당 토큰 처리 시점에만 사용되고 다시 필요하지 않기 때문입니다.
* KV Cache 미사용 시 계산 복잡성: KV Cache 없이 텍스트를 생성할 때, N번째 토큰을 예측하는 데 O(N²)의 연산량이 필요합니다. 이는 이전 모든 토큰을 처음부터 다시 처리해야 하기 때문이며, 다층 모델에서는 모든 레이어에서 각 위치마다 어텐션을 재계산해야 하는 더 큰 비효율을 초래합니다.
* KV Cache 사용 시 계산 복잡성: KV Cache를 사용하면 새로운 토큰에 대해서만 Q, K, V를 계산하고, 이 Q를 캐싱된 K와 연산하므로 N번째 토큰 예측에 O(N)의 연산량만 필요합니다. 이는 이전 토큰에 대한 K, V 재계산을 완전히 제거하여 계산량을 획기적으로 줄입니다.
* Compute vs. Memory Trade-off: KV Cache는 계산량을 크게 줄여주지만, 시퀀스 길이가 길어질수록 K와 V 벡터의 크기가 선형적으로 증가하여 GPU 메모리(VRAM) 사용량이 늘어나는 단점이 있습니다. 이는 특히 생성되는 출력 길이가 예측 불가능할 때 메모리 관리 문제를 야기하며, PagedAttention과 같은 메모리 관리 기법의 필요성을 증대시킵니다.
* 다층 모델에서의 KV Cache: 다층 모델에서 KV Cache가 없을 경우, 특정 레이어에서 이전 토큰들의 출력을 재계산하지 않으면 이후 레이어에서 필요한 K, V 값을 얻지 못하게 됩니다. 따라서 모든 레이어에서 모든 토큰에 대한 어텐션을 재계산해야 하므로, 문맥 길이가 길수록 N² 복잡성이 모델의 레이어 수만큼 곱해져 매우 비효율적입니다.

시사점

KV Cache는 LLM 추론 속도를 비약적으로 향상시키는 필수 최적화 기법이지만, 증가하는 메모리 요구량을 관리하기 위한 PagedAttention과 같은 새로운 메모리 관리 기술 및 Grouped-Query Attention (GQA), Multi-Head Latent Attention (MLA) 등의 모델 아키텍처 혁신이 병행되어야 함을 시사합니다.