nanoVLM 에서부터 KV Cache 구현하기

우리는 nanoVLM 저장소 (순수 PyTorch 를 사용하여 자체 Vision Language Model 을 훈련하는 작은 코드베이스) 에서 KV Caching 을 처음부터 구현했습니다. 이를 통해 생성 단계에서 **38%**의 속도 향상을 얻었습니다. 이 블로그 포스트에서는 KV Caching 과 이를 구현하면서 얻은 모든 경험을 다룹니다. 배운 교훈은 일반적이며 모든 autoregressive language model generation 에 적용할 수 있습니다. 작은 코드베이스에서 처음부터 구현하는 것은 훌륭한 학습 경험입니다. 함께 따라오세요!

Autoregressive language models 는 한 번에 한 토큰을 샘플링하여 텍스트를 생성합니다. 추론 (inference) 단계에서는 모델이 주어진 입력 시퀀스를 처리하고 다음 토큰을 예측하며, 이를 시퀀스에 추가하고 이 과정을 특정 종료 기준 (stopping criterion) 에 도달할 때까지 반복합니다.

이 단계별 생성은 본질적으로 순차적입니다:

• 토큰 을 생성하려면 모델은 에서부터 전체 시퀀스를 고려해야 합니다. 위의 예에서 첫 번째 인스턴스인 the 는 , 반면 모든 이전 토큰들은 [What, is, in] 입니다.
• Transformer 가 내부적으로 병렬 처리되지만, 각 새로운 예측은 모든 Transformer 레이어를 통한 완전한 forward pass 를 필요로 하며, 이는 시퀀스 길이에 대해 제곱의 메모리/컴퓨트 비용을 발생시킵니다.

이 반복은 계산적 중복성 (redundancy) 을 초래합니다. 이 포스트에서는 이러한 비효율성을 완화하는 최적화 기법인 KV Caching을 탐구합니다.

목차:

• Transformer 아키텍처 재검토
• 중복이 발생하는 곳
• KV Caching 이 이를 어떻게 해결하는지
• nanoVLM 의 KV Caching: 이론에서 실용까지
• 요약: KV Caching 이 중요한 이유

캐싱에 돌입하기 전에, 먼저 Transformer 모델에서 attention 이 어떻게 작동하는지 다시 검토해봅시다. Transformer language model 은 다음으로 구성된 스택된 레이어로 구성됩니다:

• Multi-head self-attention
• Feed-forward network (MLP)
• Residual connections 과 layer normalisation

KV Caching이 도움이 되는 곳을 이해하려면, self-attention 메커니즘에 집중합니다. 특히 단일 attention head 내부에서.

핵심 계산을 시각화하기 위해 간단한 PyTorch 구현을 살펴보겠습니다.

import torch
input_seq_length = 5
dim_model = 10
...

입력 임베딩으로 표현된 시퀀스 에서, self-attention 은 다음과 같이 계산됩니다:

• , 와 함께
• , 와 함께
• , 와 함께
• 미래 토큰 접근을 방지하기 위한 Causal mask

최종 출력은 다음과 같습니다:

다음과 같이 Causal mask 를 사용한 최소한의 PyTorch 동등체입니다:

import torch.nn.functional as F
import math
d_k = K.shape[-1]
...

Autoregressive generation 에서 모델은 한 번에 한 토큰을 생성합니다. 각 단계에서, 전체 시퀀스를 위해 , , 를 다시 계산합니다.

new_token_emb = torch.randn(1, dim_model)
extended_input = torch.cat([input_ids_emb, new_token_emb], dim=0)
Q_ext = extended_input @ W_q
...

중복성을 확인하기 위해:

torch.testing.assert_close(K, K_ext[:input_seq_length]) # test pass
torch.testing.assert_close(V, V_ext[:input_seq_length]) # test pass

이러한 체크는 모든 최신 토큰을 제외한 경우, 와 는 이전에 계산된 값과 동일함을 보여줍니다.

Original (5×5): Extended (6×6):
■ ■ ■ ■ ■ ■ ■ ■ ■ ■ □
■ ■ ■ ■ ■ ■ ■ ■ ■ ■ □
...

■= 이미 계산되고 재사용됨 □= 불필요하게 다시 계산됨

대부분의 attention 계산은 불필요하게 반복됩니다. 시퀀스가 커질수록 더 비싸집니다.

이러한 비효율성을 제거하기 위해, 우리는 KV Caching을 사용합니다:

초기 프롬프트를 처리한 후, 각 레이어에 대해 계산된 키 (key) 와 값 (value) 를 캐시합니다.

생성 과정 중에는 새로운 토큰을 위해만 계산하고 추가하며, 이를 캐시에 붙여넣기합니다.

현재 토큰을 위해 계산한 값을 사용하여, 캐시된 키와 값과 함께 출력을 얻습니다.

이것은 전체 시퀀스 재계산을 기반으로 한 생성 방식을 경량화되고 점진적인 업데이트 방식으로 변화시킵니다.

✅ 실제로 이 캐시는 각 레이어마다 "key" 와 "value"라는 키를 가진 사전 (dictionary) 으로, 각각의 형태는 (
batch_size
,
num_heads
,
seq_len_cached
,
head_dim
) 입니다.

이것이 현대적인 LLM 이 긴 출력을 효율적으로 생성할 수 있는 기반입니다.

이제 KV 캐싱에 대한 이론을 이해했으니, nanoVLM 저장소 내부에서 실제로 어떻게 구현되는지 살펴보겠습니다. 이는 초단축적이고 자체적으로 완성된 코드베이스이기 때문에 이상적인 테스트베드입니다.

KV 캐시는 우리 모델의 세 가지 핵심 구성 요소 전반에 걸쳐 활성화됩니다:

• Attention block - 키와 값을 사용하여 KV 캐시를 업데이트합니다.
• Language model - 레이어별로 캐시를 추적합니다.
• Generation loop - 초기 패스 (입력 프롬프트) 와 순차적 decode 단계를 분리합니다.

LanguageModelGroupedAttention 클래스에서, 우리는 forward 함수를 키와 값의 캐시 (block_kv_cache) 를 받아들이고 업데이트하도록 수정했습니다.

과거에는 모델이 각 생성 단계마다 키와 값을 재계산했습니다. 이제 우리는 현재 토큰을 위해만 계산하고, 이를 캐시된 값에 추가합니다.

def forward(self, x, cos, sin, attention_mask=None, block_kv_cache=None):
is_prefill = block_kv_cache is None
B, T_curr, C = x.size()
...

LanguageModel 클래스에서 우리는 레이어별 캐시 추적을 도입했습니다. start_pos 인수는 모델이 새로 생성된 토큰에 대해 올바른 로터리 위치 인코딩 (rotary positional encodings) 을 계산하는 데 도움을 줍니다.

def forward(self, x, kv_cache=None, start_pos=0):
T_curr = x.size(1)
position_ids = torch.arange(start_pos, start_pos + T_curr, device=x.device)
...

kv_cache: 각 트랜스포머 레이어 하나에 대한 사전 목록으로, 이전 키와 값을 보유합니다.

start_pos: 로터리 임베딩이 현재 생성 인덱스와 정렬되도록 보장합니다.

가장 큰 구조적 변화는 VisionLanguageModel 의 generate() 방법입니다.

우리는 생성을 두 단계로 분리했습니다:

PREFILL PHASE: 전체 프롬프트를 인코딩하고 초기 캐시를 구축합니다.

DECODE PHASE: 캐시된 키/값을 사용하여 토큰 하나씩 생성합니다.

# PREFILL: Process the input prompt, fill the cache
prompt_output, kv_cache_list = self.forward(
inputs,
...

이러한 단계를 분리함으로써 우리는 중복 계산을 피하고 추론 속도를 크게 향상시킵니다. 특히 긴 프롬프트의 경우 더욱 그렇습니다.

모듈	원래 동작	새로운 동작
LanguageModelGroupedAttention.forward	각 단계마다 키, 값을 재계산합니다	키와 값을 사용하여 KV 캐시를 업데이트합니다
LanguageModel.forward	이전 상태에 대한 기억 없음	레이어별 KV 캐시를 추적하고 start_pos 를 처리합니다
VisionLanguageModel.generate	한 단계 생성 루프	PREFILL 과 DECODE 단계로 분리됩니다

이점	설명
Incremental growth	새로운 토큰마다 캐시가 한 행씩 성장합니다
Position-aware decoding	start_pos 는 위치 인코딩 계산의 정확성을 보장합니다
Efficiency	토큰당 추론을 O(seq len) 으로 줄여, 이차적 (quadratic) 로 변경합니다

KV 캐싱은 자기회귀 생성 (autoregressive generation) 과정에서 불필요한 계산을 제거하여 더 빠르고 효율적인 추론을 가능하게 하며, 특히 긴 시퀀스 및 실시간 애플리케이션에서 더욱 효과적입니다. 이는 속도와 메모리 간의 트레이드오프이며, 그 단점은 더 복잡한 코드와 빔 검색 (beam-search) 등 다양한 추론 스키마를 제한할 수 있습니다. KV 캐싱은 LLM 추론 속도를 높이는 데 널리 사용되는 방법으로, 이를 통해 소비자 하드웨어에서도 실행이 가능해졌으며, 이제 작동 원리도 알고 계십니다!