Transformer 추론(Inference)이 내부적으로 실제로 어떻게 작동하는지, 그리고 왜 KV 캐시(KV Cache)가 LLM의 속도가 느려지는 것을 방지하는 가장 중요한 단일 최적화 기술인지 알아봅니다. 긴 프롬프트(Prompt)에서도 LLM이 왜 빠르게 응답하는지 궁금했다면, 그 답은 바로 KV 캐시(KV Cache)에 있습니다. 하지만 대부분의 설명은 "Key와 Value를 저장한다"는 수준에서 그칩니다. 이 글은 더 깊이 들어갑니다.
이 글을 통해 배우게 될 내용:
이 글을 마칠 때쯤 여러분은 다음을 이해하게 될 것입니다:
- 왜 자기회귀(Autoregressive) LLM 생성은 설계상 비용이 많이 드는가
- 어텐션(Attention)이 실제로 무엇을 계산하는지 — 그리고 왜 이를 재계산하는 것이 낭비인지
- 프리필(Prefill) 단계와 디코드(Decode) 단계의 차이
- KV 캐시(KV Cache)가 어떻게 성장하며, 언제 GPU의 가장 큰 적이 되는지
- vLLM의 PagedAttention이 메모리 파편화(Memory Fragmentation) 문제를 어떻게 해결했는지
- 다음에 올 기술들: 양자화된 캐시(Quantized Cache), 슬라이딩 윈도우(Sliding Window), 투기적 디코딩(Speculative Decoding)
서론: LLM 추론(Inference)이 비싼 이유
불편한 진실부터 시작해 봅시다. 여러분이 GPT-4나 Claude에 프롬프트(Prompt)를 보내고 첫 번째 토큰(Token)이 나타나는 것을 지켜볼 때, 여러분의 GPU는 단 하나의 문자를 생성하기 전 이미 수백만 번의 부동 소수점 연산(Floating-point operations)을 수행했습니다. 그리고 응답의 그다음 토큰이 나올 때마다 — 똑같은 과정을 반복합니다. 작은 버전이 아닙니다. 전체 계산을 수행합니다. 전체 시퀀스(Sequence)에 대한 어텐션(Attention)을 매번 수행합니다.
최적화가 없다면, 200개의 토큰 응답을 생성하는 7B 파라미터(Parameter) 모델은 계속 커지는 전체 시퀀스에 대해 어텐션(Attention)을 200번 재계산해야 합니다. 4,096개의 컨텍스트(Context)를 가진 70B 모델의 경우, 이는 단순히 느린 수준이 아니라 실제 서비스(Production)에서 거의 사용 불가능한 수준입니다. 이것이 LLM 추론(Inference)의 핵심적인 경제적 문제입니다. 자기회귀(Autoregressive) 생성은 본질적으로 순차적이며 비용이 많이 듭니다. 배치(Batch) 단위로 학습(Training)을 병렬화하는 것처럼 출력 토큰(Output Token) 간의 생성을 병렬화할 수 없습니다. 각각의 새로운 토큰은 그 이전에 나온 모든 토큰에 의존합니다.
KV 캐시(KV Cache)는 현대적인 LLM 추론(Inference)을 경제적으로 실행 가능하게 만드는 엔지니어링 솔루션입니다. 이것은 마법이 아니라, 의도적인 메모리-연산 트레이드오프(Memory-compute tradeoff)입니다. 이를 깊이 이해하는 것은 모델을 단순히 배포하는 ML 엔지니어와 모델을 최적화하는 엔지니어 사이의 차이를 만듭니다.
토큰 생성 과정에서 일어나는 일
캐싱(Caching)에 대해 이야기하기 전에, 우리가 무엇을 캐싱하는지부터 이해해 봅시다. LLM은 텍스트를 왼쪽에서 오른쪽으로 한 번에 하나의 토큰씩 생성합니다. 각 생성 단계는 다음과 같습니다:
- 전체 입력 프롬프트(Input Prompt) + 이전에 생성된 모든 토큰을 가져옵니다.
- 트랜스포머(Transformer)를 통해 완전한 순전파(Forward Pass)를 실행합니다.
- 어휘 사전(Vocabulary)에 대한 확률 분포(Probability Distribution)를 생성합니다.
- 해당 분포에서 다음 토큰을 샘플링(Sample)합니다.
그다음 새로운 토큰이 시퀀스(Sequence)에 추가되며, 이 과정이 반복됩니다. 다음은 이 루프를 의사코드(Pseudocode)로 나타낸 것입니다:
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")
input_ids = tokenizer.encode("The capital of France is", return_tensors="pt")
generated = input_ids
for _ in range(50):
with torch.no_grad():
outputs = model(generated)
logits = outputs.logits[:, -1, :]
next_token = torch.argmax(logits, dim=-1)
generated = torch.cat([generated, next_token.unsqueeze(0)], dim=-1)
print(tokenizer.decode(generated[0]))
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
여기서 발생하는 문제를 주목하십시오: 50번째 단계에서 generated에는 50개 이상의 토큰이 들어 있습니다. 우리는 단지 51번째 토큰을 예측하기 위해 이 모든 토큰에 대해 전체 트랜스포머 순전파(Forward Pass)를 수행하고 있습니다. 연산량은 시퀀스 길이(Sequence Length)에 따라 선형적으로 계속 증가합니다. 질문은 이것입니다: 왜 매번 이전의 모든 토큰을 다시 처리해야 할까요?
트랜스포머 어텐션(Transformer Attention): 빠른 복습
KV 캐시(KV Cache)가 왜 존재하는지 이해하려면, 어텐션(Attention)이 실제로 무엇을 계산하는지 이해해야 합니다. 트랜스포머 추론(Inference)의 핵심은 멀티 헤드 셀프 어텐션(Multi-head Self-attention)입니다.
각 레이어(Layer)와 각 토큰 위치(Token position)에 대해, 어텐션(Attention)은 세 가지 투영(Projection)을 계산합니다:
- Q (Query): "내가 찾고 있는 것은 무엇인가?"
- K (Key): "내가 무엇을 포함하고 있는가?"
- V (Value): "내가 실제로 담고 있는 정보는 무엇인가?"
위치 $i$에 대한 어텐션 출력은 다음과 같습니다:
$\text{Attention}(Q_i, K, V) = \text{softmax}(Q_i \cdot K^T / \sqrt{d_k}) \cdot V$
이를 말로 풀어서 설명하면 다음과 같습니다: 토큰 $i$가 질문($Q$)을 던지면, 이를 모든 토큰의 키($K$)에 전파하여 어텐션 가중치(Attention weights)를 얻고, 그 가중치를 사용하여 모든 값($V$)의 가중 합(Weighted sum)을 구합니다. $n$개의 토큰으로 이루어진 시퀀스의 경우, 각 토큰은 $n$개의 모든 토큰을 참조합니다. 이는 시간과 메모리 측면 모두에서 $O(n^2)$의 복잡도를 가지며, 이것이 긴 컨텍스트(Long context)가 성능에 큰 악영향을 미치는 이유입니다.
여기 핵심적인 통찰이 있습니다: $Q$는 매 단계마다 변합니다. 하지만 이미 처리된 토큰들의 $K$와 $V$는 변하지 않습니다. 일단 토큰이 트랜스포머(Transformer)에 의해 처리되면, 해당 토큰의 Key와 Value 투영값은 고정됩니다. 이 값들은 오직 토큰의 내용과 위치에만 의존하며, 미래의 토큰에는 의존하지 않습니다. 이것이 바로 KV 캐시(KV cache)를 사용하는 수학적 근거입니다.
어텐션을 재계산하는 것이 비효율적인 이유
숫자를 통해 구체적으로 살펴보겠습니다. 표준적인 LLaMA-2 7B 모델의 사양은 다음과 같습니다:
- 32개의 트랜스포머 레이어 (Transformer layers)
- 32개의 어텐션 헤드 (Attention heads)
- 은닉 차원 (Hidden dimension): 4096
- KV 헤드 차원 (KV head dimension): 128
단일 토큰에 대해, 한 레이어에서의 $K$와 $V$ 투영값은 각각 크기가 128인 벡터입니다. 32개의 헤드와 32개의 레이어 전체에 걸쳐 한 토큰의 KV 상태(KV state)를 저장하는 비용은 다음과 같습니다:
$2 (K와 V) \times 32 (레이어) \times 32 (헤드) \times 128 (head_dim) \times 2 \text{ bytes (fp16)} = 524,288 \text{ bytes} \approx \text{토큰당 } 0.5 \text{ MB}$
이제 2,048개의 토큰으로 구성된 프롬프트(Prompt)를 가정해 봅시다:
$2,048 \text{ 토큰} \times 0.5 \text{ MB} = 1 \text{ GB의 KV 상태}$
캐싱(Caching)을 하지 않는다면, 매 디코딩 단계(Decode step)마다 이 1 GB의 KV 상태 전체를 처음부터 다시 계산해야 합니다. 200번의 디코딩 단계를 거친다면, 단 몇 백 단어를 생성하기 위해 200 GB에 달하는 계산량을 반복하게 되는 셈입니다.
이러한 재계산 경로는 메모리 대역폭(Memory bandwidth) 또한 포화시킵니다. A100(대역폭 2 TB/s)에서 13B 모델을 사용할 때, 매 단계 모델 가중치(Model weights)를 단 한 번씩만 읽는 데에도 다음과 같은 시간이 소요됩니다:
$13 ext{B 파라미터} \times 2 \text{ bytes/param} = \text{순전파(Forward pass)당 } 26 \text{ GB}$
$26 \text{ GB} / 2000 \text{ GB/s} = \text{토큰당 약 } 13 \text{ms} \rightarrow \text{초당 약 } 77 \text{개 토큰의 한계치(Ceiling)}$
어떠한 불필요한 재계산도 이 예산(Budget)을 직접적으로 깎아먹게 됩니다.
KV 캐시(KV Cache)가 저장하는 것
KV 캐시는 이전에 처리된 모든 토큰에 대한 Key 및 Value 텐서(tensors)를 저장하여, 이후의 디코딩 단계(decode steps)에서 이를 재계산할 필요가 없도록 합니다. 개념적인 레이아웃은 다음과 같습니다:
각 디코딩 단계에서 모델은:
- 새로운 토큰에 대해서만 Q(Query)를 계산합니다.
- 새로운 토큰에 대한 K(Key)와 V(Value)를 계산합니다.
- 새로운 K/V를 캐시에 추가(Append)합니다.
- 새로운 Q를 캐시된 모든 K/V에 대해 어텐션(Attention) 연산을 수행합니다.
- 새로운 토큰에 대한 출력 로짓(output logit)을 반환합니다.
KV 캐시를 사용한 어텐션:
def attention_with_cache(query_new, key_cache, value_cache, key_new, value_new):
keys = torch.cat([key_cache, key_new], dim=1)
values = torch.cat([value_cache, value_new], dim=1)
scores = torch.einsum("bhd,bshd->bhs", query_new, keys) / math.sqrt(d_k)
weights = torch.softmax(scores, dim=-1)
output = torch.einsum("bhs,bshd->bhd", weights, values)
return output, keys, values
결과: 단계당 $O(n)$의 재계산 대신, $O(1)$의 쿼리 계산을 수행합니다. 이는 시퀀스(sequence)가 얼마나 길어지든 상관없이 일정한 비용(constant cost)이 든다는 것을 의미합니다.
Prefill 단계 vs. Decode 단계
실제 운영 환경의 LLM 추론(inference)에는 근본적으로 다른 두 가지 동작 모드가 있습니다.
Prefill 단계 (Prefill Phase)
프롬프트(prompt)를 처음 제출하면, 모델은 모든 프롬프트 토큰을 병렬로 처리합니다:
프롬프트: $[T_0, T_1, T_2, ..., T_{2047}]$
$\downarrow \downarrow \downarrow \downarrow$
$[모든 토큰이 배치 행렬 연산(batch matrix ops)을 통해 동시에 처리됨]$\downarrow$
$[2048개 위치 모두에 대해 전체 KV 캐시가 채워짐]$\downarrow$
[첫 번째 출력 토큰 생성]
Prefill은 연산 제한적(compute-bound)입니다. 즉, 모든 위치에 대해 대규모 행렬 곱셈(matrix multiplications)을 한꺼번에 수행합니다. GPU가 이 작업에 탁월한 성능을 발휘합니다.
Prefill 시간은 첫 번째 토큰 생성 시간(TTFT, Time to First Token)이 됩니다. 사용자는 이를 첫 글자가 나타나기 전까지의 지연 시간으로 느끼게 됩니다.
디코딩 단계 (Decode Phase)
Prefill 이후, 각 새로운 토큰은 한 번에 하나씩 생성됩니다:
단계 1: 새로운 토큰 $T_{2048}$ $\rightarrow$ $T_{2048}$에 대해서만 $Q, K, V$ 계산 $\rightarrow$ $K_{2048}, V_{2048}$을 캐시에 추가 $\rightarrow$ 0부터 2048까지의 위치에 대해 어텐션 (Attend) 수행 $\rightarrow$ $T_{2049}$ 샘플링
단계 2: 새로운 토큰 $T_{2049}$ $\rightarrow$ 캐시에 이제 2050개의 항목이 있음 $\rightarrow$ 0부터 2049까지의 위치에 대해 어텐션 (Attend) 수행 $\rightarrow$ $T_{2050}$ 샘플링
디코딩은 메모리 대역폭 제한 (Memory-bandwidth-bound) 상태입니다. 즉, GPU의 HBM으로부터 계속해서 커지는 KV 캐시를 읽어와야 합니다. 연산 코어 (Compute cores)는 메모리 읽기를 기다리며 대부분 유휴 (Idle) 상태로 있게 됩니다.
디코딩 단계 메모리 읽기 (LLaMA 7B, 2048 캐시 기준):
- 모델 가중치 (Model weights): ~14 GB
- KV 캐시 (KV cache): ~1 GB (계속 증가 중)
- 총합: ~15 GB
2 TB/s 대역폭에서: ~7.5ms/token $\rightarrow$ 이론적 최대치 약 133 tokens/sec
GPU 메모리와 KV 캐시의 증가
여기서부터 비용이 급격히 비싸집니다.
KV 캐시 크기 = batch_size $\times$ seq_len $\times$ num_layers $\times$ num_heads $\times$ head_dim $\times$ 2 $\times$ dtype_bytes
실제 예시 — LLaMA-2 13B, 배치 크기 (batch size) 8, 4K 컨텍스트 (context):
batch_size = 8
seq_len = 4096
num_layers = 40
num_kv_heads = 40
head_dim = 128
dtype_bytes = 2
kv_cache_bytes = (
batch_size * seq_len * num_layers * num_kv_heads * head_dim * 2 * dtype_bytes
)
print(f"KV Cache: {kv_cache_bytes / 1e9:.2f} GB")
이제 컨텍스트 8K, 배치 16으로 확장하면:
batch_size = 16
seq_len = 8192
A100-80GB 메모리 예산 (LLaMA-2 13B, fp16)
┌─────────────────────────────────────────────┐
│ 모델 가중치 (Model Weights) ~26 GB │
├─────────────────────────────────────────────┤
│ KV 캐시 (KV Cache) ~30–40 GB │ ← 격전지 (the battleground)
├─────────────────────────────────────────────┤
│ 활성화 값 / 기타 (Activations / Other) ~5–10 GB │
├─────────────────────────────────────────────┤
│ CUDA 런타임 / 기타 (CUDA Runtime / Misc) ~2–3 GB │
└─────────────────────────────────────────────┘
KV 캐시는 유일한 동적 구성 요소 (dynamic component)입니다. 시퀀스가 길어짐에 따라 증가하고, 시퀀스가 끝나면 줄어들며, 주의 깊게 관리하지 않으면 파편화 (fragment)됩니다. 다른 모든 구성 요소는 로드 시점에 고정됩니다.
Continuous Batching 및 PagedAttention
정적 배치 (Static Batching)의 문제점
기존의 추론 서버들은 정적 배치 (Static Batching)를 사용했습니다. 즉, N개의 요청이 모일 때까지 기다렸다가 모든 요청이 완료될 때까지 실행하는 방식입니다. 하지만 요청들은 서로 다른 시점에 종료됩니다. 예를 들어, 8개의 요청 중 7개가 50단계에서 종료되고 나머지 하나가 500단계까지 실행된다면, 긴 꼬리 (long tail) 구간 동안 GPU는 87.5%의 유휴 (idle) 상태가 됩니다.
연속 배치 (Continuous Batching)
최신 엔진들은 요청 (request) 단위가 아닌 반복 (iteration) 단위로 배치를 구성합니다:
시간 → T0 T1 T2 T3 T4 T5 T6 T7
Slot 0: [A A A A ←종료→ E E E ]
Slot 1: [B B ←종료→ D D D D ]
Slot 2: [C C C C C C ←종료→ F ]
요청 A가 종료되는 즉시, Slot 0은 곧바로 요청 E를 시작합니다. 이를 통해 GPU 활용률 (utilization)을 높게 유지할 수 있습니다. 하지만 이는 새로운 문제, 즉 KV 캐시 (KV cache) 메모리 파편화 (fragmentation)를 야기합니다. 각 요청마다 길이가 다르기 때문에, 엄청난 양의 메모리를 낭비하지 않고는 연속적인 메모리 블록을 미리 할당할 수 없습니다.
PagedAttention (vLLM)
vLLM의 PagedAttention (2023)은 운영체제 (OS)의 페이징 (paging)에서 빌려온 가상 메모리 (virtual memory) 비유를 통해 이 문제를 해결했습니다. 요청당 연속적인 KV 블록을 할당하는 대신, 메모리를 고정된 크기의 페이지 (pages, 일반적으로 각 16개 토큰)로 나눕니다. 블록 테이블 (block table)은 논리적 페이지 (logical pages)를 물리적 GPU 메모리 (physical GPU memory)에 매핑합니다:
논리적 뷰 (Sequence A): [Page 0 → Page 1 → Page 2 → Page 3]
물리적 GPU 메모리:
Block 7: Sequence A의 0–15번 토큰
Block 2: Sequence A의 16–31번 토큰
Block 15: Sequence A의 32–47번 토큰
Block 3: Sequence A의 48–63번 토큰
블록 테이블 (Block table): Seq A: { logical 0 → physical 7, logical 1 → physical 2, logical 2 → physical 15, logical 3 → physical 3 }
이점:
- 내부 파편화 (internal fragmentation) 없음 — 시퀀스가 성장함에 따라 한 번에 하나의 페이지씩 할당
- 접두사 공유 (Prefix sharing) — 동일한 시스템 프롬프트 (system prompt) KV 페이지를 수천 개의 요청 간에 공유 (copy-on-write)
- 스와핑 (Swapping) — 메모리 압박이 있을 때 페이지를 CPU 메모리로 축출 (evict) 가능
- 메모리 낭비가 약 30–40% (파편화로 인해)에서 4% 미만으로 감소
간소화된 블록 테이블 개념 class BlockTable:
def init(self, block_size=16, num_blocks=1000):
self.block_size = block_size
self.free_blocks = list(range(num_blocks))
self.block_data = {} # physical_block_id → KV tensor
self.tables = {}
seq_id → {logical → physical}
def allocate_block(self):
return self.free_blocks.pop()
def get_kv(self, seq_id, logical_block_idx):
physical = self.tables[seq_id][logical_block_idx]
return self.block_data[physical]
How vLLM and Inference Engines Optimize KV Cache
vLLM in Practice
from vllm import LLM, SamplingParams
llm = LLM(
model="meta-llama/Llama-2-7b-hf",
tensor_parallel_size=2, # split across 2 GPUs
gpu_memory_utilization=0.90, # leave 10% headroom
max_model_len=4096
)
outputs = llm.generate([
"Explain KV cache to an ML engineer"
], SamplingParams(temperature=0.7, max_tokens=512))
vLLM ships: PagedAttention, continuous batching, Flash Attention v2, tensor parallelism, and prefix caching out of the box.
Flash Attention
Flash Attention avoids the expensive O(n²) memory allocation by tiling computations to fit in SRAM:
Standard attention: Q, K, V → HBM → compute N×N matrix → HBM → output
Memory: O(n²)
Flash Attention: Q, K, V → SRAM tiles → fused kernel → output
Never materializes full N×N matrix in HBM
Memory: O(n)
Speed: 2–4× faster
Prefix Caching
If thousands of users share the same system prompt, their KV caches for that prefix are identical. Compute it once, share it everywhere:
Without prefix caching:
Request 1: compute KV for [SYSTEM: 2048 tokens] + [query A]
Request 2: compute KV for [SYSTEM: 2048 tokens] + [query B]
Request 3: