12.1 KV Cache Management

자기회귀 serving에는 성격이 다른 두 hardware 구간이 있습니다. Prefill 은 prompt를 병렬 처리하므로 행렬 곱셈 처리량이 주된 병목인 경우가 많습니다. Decode 는 sequence마다 한 번에 한 token을 생성하며, 매 단계 weight와 attention state를 읽습니다. 따라서 memory bandwidth, cache 용량, scheduler가 성능을 좌우하는 경우가 많습니다. 이 두 구간을 연결하는 상태가 KV cache입니다.

중요한 질문은 “cache가 들어가는가?”만이 아닙니다. 모델 품질을 지키면서 길이가 다른 요청이 섞였을 때 inter-token latency를 예측 가능하게 유지하고, 끝난 요청의 용량을 즉시 회수할 수 있어야 합니다.

1. Cache가 절약하는 연산

Layer $l$ 의 attention은 token state $x_t$ 를 투영합니다.

q_t=x_tW_Q,\qquad k_t=x_tW_K,\qquad v_t=x_tW_V.

Causal decode에서 이전 token의 key와 value는 변하지 않습니다. $\{k_i,v_i\}_{i\le t}$ 를 보관하면 $t+1$ 단계에서 새 projection만 계산하고 저장된 history를 조회합니다.

o_{t+1}=\operatorname{softmax}\left( \frac{q_{t+1}K_{\le t+1}^{\top}}{\sqrt{d_h}}+M \right)V_{\le t+1}.

$d_h$ 는 head dimension, $M$ 은 causal 또는 padding mask입니다.

복잡도는 정확히 구분해야 합니다. Cache가 없으면 decode $t$ 단계마다 길이 $t$ 인 전체 forward를 다시 실행하므로 self-attention score 연산은 $O(t^2)$ 입니다. 새 token $N$ 개에 대해 합하면 attention은 $O(N^3)$ , projection과 MLP는 $O(N^2)$ 입니다. Cache가 있으면 $t$ 단계 attention이 $O(t)$ , 전체는 $O(N^2)$ 가 됩니다. 이 표기는 model width와 batch dimension을 생략합니다. 실제 시간은 점근 복잡도보다 memory traffic과 kernel utilization에 좌우되는 경우가 많습니다.

2. 용량 모델

길이가 같은 sequence로 구성된 dense cache는

\text{bytes}_{KV}=2BSLH_{KV}d_hp

byte를 사용합니다. $B$ 는 동시 sequence 수, $S$ 는 sequence당 cache token, $L$ 은 layer 수, $H_{KV}$ 는 KV head 수, $d_h$ 는 head당 값의 수, $p$ 는 값 하나의 byte 수입니다. 앞의 2는 key와 value입니다.

실제 요청 길이가 다르면 $BS$ 대신 할당 token 수 $\sum_i S_i^{\text{allocated}}$ 를 써야 합니다. Padding, block metadata, allocator 여유 공간, 임시 workspace, speculative branch, beam hypothesis, tensor-parallel 복제 또는 sharding 비용은 별도입니다.

동시 sequence 32개, sequence당 8,192 token, 80개 layer, 8개 KV head, $d_h=128$ , BF16이라면:

2\cdot32\cdot8192\cdot80\cdot8\cdot128\cdot2 =85{,}899{,}345{,}920\text{ bytes}=80\text{ GiB}.

요청 하나가 아니라 32개 요청의 합계입니다. KV cache가 항상 weight보다 큰 것도 아닙니다. 동시성과 context가 충분히 커지면 더 커질 수 있으므로 capacity plan에는 weight와 최대 live-cache 점유량을 모두 포함해야 합니다.

KV cache 용량 계산기

조밀한 cache payload를 추정합니다. allocator metadata, padding, 임시 buffer, tensor-parallel 복제 비용은 포함하지 않습니다.

동시 sequence 수 (B): 32

Sequence당 cache token (S): 8,192

Transformer layer (L): 80

KV head (H_KV): 8

Head dimension (d_h): 128

Cache precision (p): 2 B/value

조밀한 KV payload:

80.00 GiB

2 × 32 × 8192 × 80 × 8 × 128 × 2 bytes

용량 추정치일 뿐입니다. 실제 모델과 serving engine에서 latency와 품질을 benchmark해야 합니다.

3. KV 폭 줄이기

$H_Q$ 를 query head 수라고 하겠습니다.

Multi-Head Attention (MHA): $H_{KV}=H_Q$
Multi-Query Attention (MQA): 모든 query가 하나의 KV head를 공유
Grouped-Query Attention (GQA): $1<H_{KV}<H_Q$ , 여러 query head가 하나의 KV head를 공유

Payload는 $H_{KV}$ 에 선형 비례합니다. GQA는 품질·용량·bandwidth trade-off를 제공합니다. 원 논문은 평가한 변환 및 학습 설정에서 MHA에 가까운 품질과 MQA에 가까운 속도를 보고했습니다 [1]. 임의의 checkpoint가 재학습 없이 head 구조를 바꿔도 품질이 유지된다는 보장은 아닙니다.

Multi-Head Latent Attention

DeepSeek-V2는 Multi-Head Latent Attention (MLA) 를 제안했습니다 [2]. MLA는 전통적인 head별 KV payload 대신 low-rank latent $c_t^{KV}$ 와 decoupled RoPE에 필요한 위치 성분을 cache합니다. Decode에서는 선형 key/value up-projection을 인접 projection에 대수적으로 흡수할 수 있습니다. Position-dependent rotation인 RoPE는 static absorption이 불가능해 별도로 처리합니다.

DeepSeek-V2 report는 자사의 DeepSeek 67B 비교 시스템 대비 KV cache 93.3% 감소와 최대 generation throughput 5.76배를 보고합니다. 해당 architecture와 구현의 system result이지 보편적인 MLA 배수가 아닙니다. Kernel 지원, latent width, positional encoding, parallel layout, traffic 형태가 실제 이득을 결정합니다.

4. Cache Precision은 시스템 Trade-off다

FP8 storage는 FP16/BF16 대비 dense payload를 절반으로, 4-bit storage는 4분의 1로 줄입니다. 용량 이득은 계산대로지만 latency와 품질 이득은 그렇지 않습니다. Serving engine은 값 변환, scale 계산, metadata 이동에 추가 명령을 쓸 수 있습니다. KV bandwidth가 병목이 아닌 workload는 속도 이득이 작을 수 있습니다.

Scale granularity, calibration data, E4M3/E5M2 같은 format, outlier, accumulation precision을 함께 정해야 합니다. vLLM 문서는 가능하면 calibrated scale을 사용하고 실제 workload에서 정확도를 평가하라고 명시합니다 [3].

건전한 rollout은 다음을 비교합니다.

Context-length slice별 품질
Prefill latency, p50/p95 time to first token, p50/p95 inter-token latency
고정 latency SLO에서의 throughput
Cache 점유율, allocation failure, eviction 또는 recomputation 비율

“Cache token이 두 배 들어간다”와 “token 생성 속도가 두 배다”는 다른 주장입니다.

5. 교육용 GQA Cache

다음 kernel은 cache shape와 causal masking을 보여줍니다. Cache된 token 뒤에 여러 새 token을 붙이는 chunked decode도 처리합니다.

의도적으로 단순화한 구현입니다. RoPE, padding mask, sliding window, quantized storage, tensor parallelism, paged allocation, prefix sharing, in-place writer를 생략했습니다. torch.cat을 사용하므로 production에 적합하지 않습니다. enable_gqa 지원도 설치된 PyTorch backend에 따라 다릅니다.

import torch
import torch.nn as nn
import torch.nn.functional as F

KVCache = tuple[torch.Tensor, torch.Tensor]  # 각각 [B, H_kv, S, d_h]


class CausalGQAAttention(nn.Module):
    def __init__(self, d_model: int = 4096, n_heads: int = 32, n_kv_heads: int = 8):
        super().__init__()
        if d_model % n_heads or n_heads % n_kv_heads:
            raise ValueError("d_model % n_heads와 n_heads % n_kv_heads는 0이어야 합니다")
        self.n_heads = n_heads
        self.n_kv_heads = n_kv_heads
        self.d_head = d_model // n_heads
        self.q_proj = nn.Linear(d_model, n_heads * self.d_head, bias=False)
        self.k_proj = nn.Linear(d_model, n_kv_heads * self.d_head, bias=False)
        self.v_proj = nn.Linear(d_model, n_kv_heads * self.d_head, bias=False)
        self.o_proj = nn.Linear(d_model, d_model, bias=False)

    @torch.no_grad()
    def forward(self, x: torch.Tensor, cache: KVCache | None = None):
        batch, new_tokens, width = x.shape
        q = self.q_proj(x).view(batch, new_tokens, self.n_heads, self.d_head)
        k = self.k_proj(x).view(batch, new_tokens, self.n_kv_heads, self.d_head)
        v = self.v_proj(x).view(batch, new_tokens, self.n_kv_heads, self.d_head)
        q, k, v = (z.transpose(1, 2) for z in (q, k, v))

        past_tokens = 0
        if cache is not None:
            past_k, past_v = cache
            past_tokens = past_k.size(2)
            k = torch.cat((past_k, k), dim=2)
            v = torch.cat((past_v, v), dim=2)

        # PyTorch SDPA boolean mask에서 True는 attention 허용입니다.
        q_pos = torch.arange(past_tokens, past_tokens + new_tokens, device=x.device)
        k_pos = torch.arange(k.size(2), device=x.device)
        allowed = (k_pos[None, :] <= q_pos[:, None])[None, None, :, :]

        y = F.scaled_dot_product_attention(
            q, k, v, attn_mask=allowed, is_causal=False, enable_gqa=True
        )
        y = y.transpose(1, 2).contiguous().view(batch, new_tokens, width)
        return self.o_proj(y), (k, v)

명시적인 position mask가 중요합니다. Cache가 비어 있지 않은 multi-token chunk에 is_causal=True를 순진하게 적용하면 kernel 계약에 따라 triangular mask 원점이 잘못 정렬될 수 있습니다.

6. Production Cache 계약

Production engine은 block을 미리 할당하고 새 KV row를 제자리에 기록합니다. PagedAttention은 virtual-memory-style block mapping을 사용해 sequence마다 하나의 연속 공간을 요구하지 않으며, 공통 prefix block을 공유할 수 있습니다 [4]. 다음 절에서 allocator를 다룹니다.

출시 전에는 다음을 기록합니다.

Artifact: model, tokenizer, positional encoding, cache dtype/scale, kernel/engine version, block size, parallel layout
Capacity: weight, cache, activation, workspace, CUDA graph, fragmentation, 장애 여유분을 분리
Telemetry: logical/allocated token, block utilization, prefix hit rate, eviction/preemption, recomputed token, HBM 사용량, batch token, TTFT, inter-token latency, output tokens/s
Failure policy: allocation failure 전에 admission을 중단하거나 낮은 우선순위 요청을 줄이고, 품질 slice 또는 tail latency가 기준을 넘으면 cache format을 rollback
Recovery: preemption/migration 뒤 token continuation을 검증하고, 직렬화한 KV tensor가 model 또는 kernel revision 사이에 호환된다고 가정하지 않기

7. 핵심 정리

KV caching은 중복 연산을 persistent memory traffic으로 바꿉니다. MQA/GQA/MLA는 폭을, 저정밀도는 값당 byte를, paged allocation은 낭비를 줄입니다. 어느 기법도 독립적인 switch가 아닙니다. 품질, kernel, scheduler, 요청 길이 분포가 작은 cache를 실제로 더 빠르고 안정적인 서비스로 만들 수 있는지 결정합니다.

Quizzes

Quiz 1: “KV cache가 generation을 quadratic에서 linear로 바꾼다”는 설명이 불완전한 이유는 무엇인가요?

Decode $t$ 단계에서도 $O(t)$ history를 읽으므로 새 token $N$ 개의 전체 attention은 $O(N^2)$ 입니다. Cache는 단계마다 전체 prefix projection과 $O(t^2)$ attention을 재계산하는 일을 제거합니다.

Quiz 2: Logical token 수가 같아도 cache 용량이 달라질 수 있는 이유는 무엇인가요?

Block rounding, 길이 분산, speculative branch, padding, prefix sharing, 복제, fragmentation이 allocated token과 overhead를 바꾸기 때문입니다.

Quiz 3: FP8 KV cache의 throughput 이득을 주장하기 전에 무엇이 필요한가요?

Model, scheduler, traffic, latency 목표, output policy를 고정하고 bandwidth, 점유율, 변환 overhead, tail latency, throughput, 품질 slice를 측정해야 합니다.

Quiz 4: MLA가 position 관련 cache 성분을 별도로 유지하는 이유는 무엇인가요?

RoPE는 위치에 따른 rotation이므로 순수 선형 up-projection처럼 static projection matrix에 흡수할 수 없습니다.

Quiz 5: torch.cat으로 cache tensor를 계속 늘리는 방식의 문제는 무엇인가요?

증가하는 history를 매번 새로 할당하고 복사합니다. Production engine은 block을 예약하고 row를 제자리에 씁니다.

References

Ainslie, J., et al. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. arXiv:2305.13245.
DeepSeek-AI. (2024). DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model. arXiv:2405.04434.
vLLM Project. Quantized KV Cache. 공식 문서.
Kwon, W., et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. arXiv:2309.06180.