[PYTHON] LLM 모델 서빙 시 KV Cache가 추론 속도에 미치는 3가지 영향과 성능 해결 방법

KV Cache (Key-Value Cache)

생성형 AI(Generative AI) 시대의 핵심인 거대언어모델(LLM)을 효율적으로 서빙하기 위해서는 단순한 하드웨어 가속을 넘어 소프트웨어적인 아키텍처 최적화가 필수적입니다. 특히 KV Cache(Key-Value Cache)는 오토리그레시브(Autoregressive) 모델의 추론 지연 시간(Latency)을 획기적으로 단축시키는 마법 같은 기술입니다. 본 가이드에서는 KV Cache의 메커니즘과 이것이 Python 기반 서빙 환경에서 성능을 어떻게 좌우하는지 심층 분석합니다.

---

1. KV Cache의 본질: 왜 매번 다시 계산하지 않는가?

LLM은 이전 토큰들을 바탕으로 다음 토큰을 하나씩 예측하는 방식으로 작동합니다. 이때 매 단계마다 전체 문맥(Context)을 다시 어텐션(Attention) 연산하면, 문장이 길어질수록 계산 복잡도가 $O(N^2)$으로 증가하여 속도가 기하급수적으로 느려집니다.

KV Cache는 이미 계산된 이전 토큰들의 Key와 Value 행렬을 GPU 메모리에 저장해 두었다가 다음 토큰 예측 시 재사용하는 기법입니다. 이를 통해 중복 계산을 제거하고 $O(N)$의 선형적인 복잡도로 추론 속도를 해결할 수 있습니다.

2. KV Cache 유무에 따른 성능 및 리소스 차이 비교

KV Cache를 적용했을 때와 그렇지 않았을 때, 시스템 리소스와 추론 효율성이 어떻게 변하는지 분석한 결과입니다.

비교 항목	KV Cache 미적용 (Re-computation)	KV Cache 적용 (Caching)
추론 단계별 계산량	누적되는 토큰 수에 따라 증가	새로운 토큰 1개에 대해서만 계산
추론 속도 (TTFT vs TPOT)	문맥이 길어질수록 지연 시간 급증	일정한 수준의 토큰당 생성 속도 유지
GPU 메모리 점유율	상대적으로 낮음 (연산 중심)	매우 높음 (캐시 저장 공간 필요)
처리량 (Throughput)	낮음 (연산 병목 발생)	높음 (메모리 대역폭 병목으로 전이)
복잡도	단순함	캐시 관리 및 페이징 로직 필요

---

3. 실무 적용을 위한 7가지 KV Cache 최적화 및 서빙 예제

Python 환경에서 Hugging Face Transformers나 vLLM 스타일의 최적화 기법을 직접 구현하고 활용하는 방법입니다.

Example 1: Transformers 라이브러리에서의 use_cache 활성화

가장 기본적으로 모델 추론 시 캐시 기능을 사용하여 속도를 높이는 방법입니다.

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_id = "meta-llama/Llama-2-7b-hf"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16).cuda()

inputs = tokenizer("KV Cache is essential for", return_tensors="pt").to("cuda")

# use_cache=True가 기본이지만 명시적으로 확인
output = model.generate(
    **inputs, 
    max_new_tokens=50, 
    use_cache=True, 
    return_dict_in_generate=True
)
print(tokenizer.decode(output.sequences[0]))
        

Example 2: Manual KV Cache 객체 다루기 (Past Key Values)

반복 루프 내에서 직접 캐시를 전달하여 추론 단계를 세밀하게 제어하는 해결책입니다.

# 초기 단계 (Prefill)
past_key_values = None
input_ids = inputs.input_ids

for _ in range(5):
    outputs = model(input_ids, past_key_values=past_key_values, use_cache=True)
    next_token_logits = outputs.logits[:, -1, :]
    next_token_id = torch.argmax(next_token_logits, dim=-1).unsqueeze(-1)
    
    # 캐시 갱신 및 입력 업데이트
    past_key_values = outputs.past_key_values
    input_ids = next_token_id
    print(f"Generated Token ID: {next_token_id.item()}")
        

Example 3: PagedAttention 스타일의 메모리 레이아웃 시뮬레이션

vLLM에서 사용하는 메모리 파편화 해결 기법인 PagedAttention의 개념적 구조입니다.

class VirtualBlockManager:
    def __init__(self, num_blocks, block_size):
        self.free_blocks = list(range(num_blocks))
        self.block_size = block_size
        self.mapping = {} # Logical to Physical

    def allocate(self, seq_id, num_tokens):
        needed_blocks = (num_tokens + self.block_size - 1) // self.block_size
        allocated = [self.free_blocks.pop(0) for _ in range(needed_blocks)]
        self.mapping[seq_id] = allocated
        return allocated

manager = VirtualBlockManager(num_blocks=100, block_size=16)
print(f"Allocated Blocks for Seq 1: {manager.allocate('seq_1', 35)}")
        

Example 4: KV Cache 양자화(8-bit/4-bit)를 통한 메모리 절약

캐시 용량이 GPU 메모리를 초과할 때 발생하는 병목을 해결하기 위해 캐시 자체를 압축합니다.

# BitsAndBytes 등의 라이브러리를 활용한 설정 예시
# 가상의 로직으로 표현
def quantize_kv_cache(key_states, value_states):
    # 8-bit 양자화 적용
    q_key = (key_states / key_states.max()).to(torch.int8)
    q_value = (value_states / value_states.max()).to(torch.int8)
    return q_key, q_value
        

Example 5: Multi-Query Attention (MQA) 적용 모델 활용

KV Head 수를 줄여 캐시 크기를 획기적으로 줄인 모델(예: Falcon, Mistral 일부) 활용법입니다.

# MQA/GQA 모델은 동일한 메모리에서 더 큰 배치 사이즈를 가질 수 있음
from transformers import AutoConfig

config = AutoConfig.from_pretrained("tiiuae/falcon-7b")
print(f"Multi-Query Attention: {config.multi_query}") # True 인지 확인
# 이는 KV Cache 메모리 사용량을 헤드 수만큼 감소시킴
        

Example 6: Flash Attention 연산과 캐시의 상호작용

메모리 대역폭 병목을 해결하기 위해 Flash Attention을 연동하는 최적화 코드입니다.

# 최신 버전의 Transformers와 PyTorch 2.0+ 사용 시
with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False):
    # 이 컨텍스트 안에서 모델 추론 실행 시 캐시 접근 속도가 최적화됨
    model.generate(**inputs, max_new_tokens=20)
        

Example 7: 캐시 용량 계산기를 통한 인프라 사이징

실제 서빙 시 필요한 GPU 메모리 요구량을 계산하는 Python 함수입니다.

def calculate_kv_cache_size(batch_size, seq_len, num_layers, num_heads, head_dim, precision=2):
    # precision=2 (FP16), 4 (FP32)
    # Key + Value 이므로 2를 곱함
    total_bytes = batch_size * seq_len * num_layers * num_heads * head_dim * 2 * precision
    return total_bytes / (1024**3) # GB 단위

# Llama-7B, 2048 Context, Batch 1 기준
gb_needed = calculate_kv_cache_size(1, 2048, 32, 32, 128)
print(f"Required KV Cache Memory: {gb_needed:.2f} GB")
        

---

4. 결론: KV Cache 최적화가 미래다

LLM 추론 속도의 병목은 이제 연산(Compute)이 아닌 메모리 대역폭(Memory Bandwidth)에 있습니다. KV Cache는 이 문제를 해결하는 핵심 열쇠이지만, 동시에 막대한 메모리를 점유하는 양날의 검이기도 합니다. 따라서 PagedAttention, 양자화, MQA/GQA와 같은 고도화된 기술을 통해 캐시 효율을 극대화하는 것이 실무 MLOps의 정수입니다. Python 개발자라면 단순한 라이브러리 호출을 넘어, 이러한 캐시 메커니즘이 인프라 비용과 사용자 경험(Latency)에 미치는 영향을 데이터로 증명하고 최적화할 줄 알아야 합니다.

내용 출처 및 참고 문헌

• vLLM Project: "PagedAttention: Software-defined memory management for LLM serving"
• Hugging Face Blog: "Optimizing LLM Inference with KV Cache"
• Dao, T., et al. (2022). "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness"
• NVIDIA Technical Blog: "Mastering LLM Serving Efficiency"