KV 캐시 최적화로 긴 컨텍스트 추론 메모리 줄이기

긴 컨텍스트를 다루는 LLM 추론 서버를 운영하다 보면, 어느 순간 병목이 연산이 아니라 메모리라는 사실을 깨닫게 됩니다. 프롬프트가 수만 토큰에 이르는 RAG나 코드 어시스턴트 워크로드에서는 모델 가중치보다 KV 캐시가 GPU 메모리의 대부분을 잡아먹고, 결국 동시에 처리할 수 있는 요청 수(배치 크기)를 제한하는 실질적 상한이 됩니다.

이 글에서는 KV 캐시가 왜 그렇게 커지는지부터 페이지드 어텐션·GQA·양자화·프리픽스 공유·축출(eviction)까지, 긴 컨텍스트 추론에서 메모리를 실제로 줄이는 레버들을 정리합니다. 각 기법은 메모리·지연·품질 사이의 트레이드오프를 가지므로, 무엇을 언제 켜야 하는지를 함께 다룹니다.

KV 캐시가 무엇이고 왜 메모리를 먹는가

트랜스포머의 자기회귀 디코딩은 토큰을 하나씩 생성합니다. 각 단계에서 어텐션은 새 토큰의 쿼리를 이전 모든 토큰의 키·값과 곱해야 합니다. 매 스텝마다 과거 키·값을 다시 계산하면 O(n²)이 되므로, 한 번 계산한 키·값을 저장해 재사용하는 것이 KV 캐시입니다. 디코딩은 스텝당 O(n)으로 떨어지지만, 그 대가로 시퀀스 길이에 선형으로 증가하는 메모리를 지불합니다.

캐시 크기는 단순한 곱셈입니다. 2(K와 V) × 레이어 수 × KV 헤드 수 × 헤드 차원 × 시퀀스 길이 × dtype 바이트이고, 여기에 배치 크기를 곱하면 총량이 됩니다.

# KV 캐시 메모리 추정 (요청 1건 기준, GQA 미적용 가정)
def kv_cache_bytes(layers, kv_heads, head_dim, seq_len, dtype_bytes=2):
    # 2 = Key + Value 두 텐서
    return 2 * layers * kv_heads * head_dim * seq_len * dtype_bytes

# 예: 32 레이어, KV 헤드 32, 헤드차원 128, 32K 토큰, fp16(2바이트)
per_req = kv_cache_bytes(32, 32, 128, 32_768, 2)
print(per_req / 1024**3, "GiB")   # 약 16 GiB — 단 한 요청에!

# 배치 16이면 256 GiB → 단일 GPU로는 불가능
print(per_req * 16 / 1024**3, "GiB")

KV 캐시는 시퀀스 길이에 정비례하므로 긴 컨텍스트일수록 캐시가 지배적이 되고, 요청마다 독립적으로 쌓이므로 동시 처리량을 늘리려면 요청당 캐시를 줄여야 합니다. 아래 기법들은 모두 이 곱셈식의 어느 항을 깎느냐로 이해할 수 있습니다.

페이지드 어텐션: 단편화라는 숨은 낭비 제거

가장 먼저 잡아야 할 것은 알고리즘이 아니라 할당 방식입니다. 전통적 구현은 요청마다 “최대 길이”만큼 캐시 버퍼를 연속으로 예약했습니다. 생성 길이는 요청마다 다르므로, 예약하고 쓰지 않는 공간이 내부 단편화로 낭비되고, 크기가 다른 요청이 오가며 외부 단편화까지 생깁니다.

페이지드 어텐션(PagedAttention)은 운영체제의 가상 메모리 페이징을 KV 캐시에 적용합니다. 캐시를 고정 크기 블록(페이지)으로 쪼개고 논리 시퀀스를 물리 블록에 매핑하는 블록 테이블을 둡니다. 필요할 때만 블록을 할당하므로 내부 단편화는 마지막 블록 한 조각으로 줄고, 블록 크기가 균일해 외부 단편화는 사라집니다.

# vLLM: 페이지드 어텐션 + 블록 기반 캐시 (엔진이 자동 관리)
from vllm import LLM, SamplingParams

llm = LLM(
    model="my-org/long-context-7b",
    gpu_memory_utilization=0.90,  # KV 캐시가 쓸 수 있는 GPU 메모리 비율
    max_model_len=32_768,
    block_size=16,                # 블록당 토큰 수 — 단편화·오버헤드 균형
    enable_prefix_caching=True,   # 공유 프리픽스 블록 재사용(뒤에서 설명)
)
# 낭비가 줄면 같은 GPU에서 배치 크기가 몇 배로 늘어난다

페이지드 어텐션은 품질 손실 없이 순수하게 낭비만 걷어내므로 공짜 점심에 가깝습니다. 긴 컨텍스트 서빙이라면 이것을 기본값으로 깔고 나머지 기법을 그 위에 얹습니다.

GQA와 MQA: KV 헤드 수를 줄이는 구조적 절감

곱셈식에서 KV 헤드 수 항을 직접 깎는 구조가 그룹 쿼리 어텐션(GQA)입니다. 표준 멀티헤드 어텐션(MHA)은 쿼리 헤드마다 자기 키·값 헤드를 갖지만, GQA는 여러 쿼리 헤드가 하나의 KV 헤드를 공유합니다. 모든 쿼리가 단일 KV 헤드를 공유하는 극단이 MQA입니다.

• MHA: KV 헤드 = 쿼리 헤드(32/32). 캐시 최대, 품질 기준선.
• GQA: KV 헤드 = 쿼리 헤드의 1/4~1/8(예: 8/32). 캐시 4~8배 절감, 품질 손실 미미.
• MQA: KV 헤드 = 1. 캐시 최소, 큰 모델에서 품질 저하가 관찰되기도 함.

# GQA 캐시 절감 비율은 "쿼리 헤드 / KV 헤드"
def gqa_saving(q_heads, kv_heads):
    return q_heads / kv_heads   # 이 배수만큼 KV 캐시가 작아진다

print(gqa_saving(32, 8))   # 4.0배 절감 (GQA-8)
print(gqa_saving(32, 1))   # 32배 절감 (MQA) — 품질 리스크 있음

주의할 점은 GQA가 모델 아키텍처 결정이라는 것입니다. 이미 학습된 MHA 모델을 서빙 시점에 GQA로 바꿀 수는 없습니다. 따라서 새 모델을 고를 때 KV 헤드 수를 스펙 시트에서 확인하는 편이 좋습니다.

KV 캐시 양자화: dtype 바이트를 깎기

곱셈식의 마지막 항인 dtype 바이트를 줄이는 것이 KV 캐시 양자화입니다. 캐시를 fp16(2바이트) 대신 fp8·int8(1바이트)로 저장하면 절반, int4면 1/4이 됩니다. 가중치 양자화와 독립적으로 캐시만 낮은 정밀도로 두는 것이 핵심입니다.

# vLLM: KV 캐시를 fp8로 저장 → 캐시 메모리 절반
llm = LLM(
    model="my-org/long-context-7b",
    kv_cache_dtype="fp8",   # "auto"(fp16) 대신 fp8. int8 계열도 지원
    max_model_len=65_536,   # 절감분으로 컨텍스트를 두 배로 늘릴 수 있음
    gpu_memory_utilization=0.90,
)

품질 관점에서 키(K)와 값(V)은 민감도가 다릅니다. 어텐션 스코어는 키에 소프트맥스를 거치므로 키의 이상치 하나가 분포를 크게 흔들 수 있어, 키가 양자화에 더 취약합니다. 그래서 실무 구현은 종종 키는 채널별, 값은 토큰별 스케일을 적용하거나 키에만 더 높은 비트를 배정합니다.

• fp8: 대부분의 워크로드에서 품질 저하가 거의 없는 무난한 기본값.
• int8: 스케일 캘리브레이션이 필요하나 fp8이 없는 하드웨어에서 유효.
• int4: 절감폭은 크나 긴 생성에서 오차 누적으로 품질이 흔들려 검증 필수.

양자화는 반드시 실제 태스크로 회귀 검증해야 합니다. perplexity만 보면 놓치는 열화가 있으므로, 정답이 있는 대표 태스크로 fp16 대비 정확도를 확인합니다.

프리픽스 공유: 같은 컨텍스트를 한 번만 저장

RAG나 few-shot 프롬프트에서는 여러 요청이 동일한 앞부분(시스템 프롬프트, 검색된 문서, 예시)을 공유합니다. 순진하게 처리하면 요청마다 같은 프리픽스의 KV 캐시를 중복 저장합니다. 프리픽스 캐싱은 공통 프리픽스의 캐시 블록을 여러 요청이 공유하게 만들어 중복 저장과 중복 프리필 연산을 동시에 없앱니다. 페이지드 어텐션 위에서는 블록 테이블이 같은 물리 블록을 여러 논리 시퀀스가 가리키게 하고, 프리픽스를 넘어 분기하는 지점에서만 copy-on-write로 새 블록을 만드는 식으로 구현됩니다.

# 프리픽스가 길고 공유될수록 이득이 커진다
# 요청 N개가 같은 프리픽스 P를 공유할 때 절감되는 캐시:
def prefix_saving(prefix_tokens, n_requests, per_tok_bytes):
    # 공유 없으면 N번 저장, 공유하면 1번 → (N-1)번치 절약
    return prefix_tokens * (n_requests - 1) * per_tok_bytes

# 8K 프리픽스, 동시요청 32건, 토큰당 128KB 가정
saved = prefix_saving(8_192, 32, 128 * 1024)
print(saved / 1024**3, "GiB")   # 약 31 GiB 절약 + 프리필 연산도 절감

효과를 극대화하려면 공유 내용을 프롬프트 맨 앞에 두어야 합니다. 캐시는 접두부 매치이므로 앞쪽에 요청마다 달라지는 값(타임스탬프·사용자 ID)을 끼워 넣으면 그 뒤 프리픽스가 공유되지 못합니다. 고정 컨텍스트를 앞에, 가변 질의를 뒤에 두는 원칙이 적용됩니다.

축출과 압축: 무엇을 버릴 것인가

가장 공격적인 절감은 캐시의 일부를 아예 버리는 것입니다. 어텐션 패턴을 보면 많은 토큰은 뒤 토큰들이 거의 참조하지 않습니다. 이런 저중요 토큰의 KV를 축출(eviction)하면 캐시를 고정 예산에 가둘 수 있지만, 이는 근사이므로 명백한 품질 트레이드오프가 있습니다.

• 슬라이딩 윈도우: 최근 W개 토큰만 유지. 국소 의존이 강한 작업엔 효과적이나 먼 과거 참조엔 부적합.
• 어텐션 싱크: 맨 앞 몇 개 토큰(싱크)은 항상 유지하고 나머지는 윈도우로 관리. 초기 토큰이 어텐션을 대량 흡수하는 현상을 이용해 안정성 개선.
• 중요도 기반 축출: 누적 어텐션 스코어가 낮은 토큰을 우선 버림. 정확도 유지에 유리하나 추적 오버헤드가 있음.

# 개념: 어텐션 싱크 + 슬라이딩 윈도우로 캐시를 고정 예산에 가둔다
def evict_kv(cache, sink=4, window=4096):
    if len(cache) <= sink + window:
        return cache
    # 맨 앞 sink 토큰(어텐션 싱크)은 항상 보존
    keep_sink = cache[:sink]
    # 최근 window 토큰만 유지, 중간의 오래된 토큰은 축출
    keep_recent = cache[-window:]
    return keep_sink + keep_recent   # 캐시 크기가 sink+window 로 상한 고정

축출 전에 반드시 물어야 할 질문은 “이 워크로드가 먼 과거를 정확히 참조해야 하는가?”입니다. 긴 문서에서 특정 사실을 찾아내는 검색형 태스크는 중간 토큰 축출에 취약합니다. 반대로 최근 맥락이 주로 중요한 챗봇형이라면 싱크+윈도우가 큰 손실 없이 메모리를 크게 줄여 줍니다.

기법 조합과 검증 순서

이 레버들은 배타적이지 않으므로 계층으로 쌓는 것이 정석입니다. 리스크가 낮은 것부터 켜고 각 단계에서 메모리와 품질을 함께 측정합니다.

• 1단계(무손실): 페이지드 어텐션 + 프리픽스 공유. 품질 손실 없이 단편화·중복 제거. 항상 켠다.
• 2단계(구조적): 모델 선택 시 GQA 확인. 이미 정해졌으면 건너뜀.
• 3단계(저손실): KV 캐시 fp8 양자화. 대표 태스크로 회귀 검증 후 채택.
• 4단계(안전판): 메모리 압박 시 오래된 블록을 CPU RAM·NVMe로 오프로딩. 지연을 감내하는 대신 요청 거절을 피함.
• 5단계(유손실): 축출·압축. 먼 과거 참조에 둔감할 때만, 품질 게이트를 통과한 경우에 한해.

# 도입 전후 메모리·품질을 함께 계측하는 벤치 루프
for CFG in baseline fp8 fp8_prefix fp8_prefix_evict; do
  # 동일 프롬프트셋으로 최대 배치·P99 지연·정확도 수집
  python bench.py --config "$CFG" 
    --dataset long_ctx_eval.jsonl --report "results/${CFG}.json"
done
# 메모리만 보지 말고 accuracy 가 임계치 아래로 떨어지는 지점을 찾는다
python compare.py results/*.json --metric accuracy --min 0.97

마무리

긴 컨텍스트 추론에서 메모리 병목의 정체는 대부분 KV 캐시이고, 그 크기는 2 × 레이어 × KV헤드 × 헤드차원 × 길이 × dtype이라는 곱셈식으로 설명됩니다. 최적화란 결국 이 식의 어느 항을 깎느냐의 문제로, 단편화는 페이지드 어텐션과 프리픽스 공유로, KV 헤드 수는 GQA로, dtype 바이트는 양자화로, 길이 항은 축출·오프로딩으로 다룹니다.

실무의 핵심은 무손실 기법을 먼저, 유손실 기법을 나중에 쌓되 매 단계에서 대표 태스크의 정확도를 게이트로 두는 것입니다. 메모리 수치만 쫓다 보면 조용히 품질이 무너지기 쉽습니다. “이 워크로드가 먼 과거를 정확히 참조해야 하는가”라는 한 가지 질문이, 얼마나 공격적으로 캐시를 줄여도 되는지를 가장 빠르게 판단하게 합니다.