KV 캐시(KV Cache), Flash Attention과 추론 최적화(Inference Optimization)

훈련(training)은 병렬적이고 연산량(FLOP) 한계(bound)에 부딪힙니다. 추론(inference)은 직렬적이고 메모리 대역폭 한계에 부딪힙니다. 병목(bottleneck)이 다르면 기법(trick)도 달라집니다.

유형: Build 언어: Python 선수 지식: Phase 7 · 02 (Self-Attention), Phase 7 · 05 (Full Transformer), Phase 7 · 07 (GPT) 예상 시간: 약 75분

문제

순진한(naive) 자기회귀 디코더(autoregressive decoder)는 N개의 토큰(token)을 생성할 때 O(N²)만큼의 일을 합니다. 매 스텝(step)마다 앞쪽 전체 접두부(prefix)에 대한 어텐션(attention)을 다시 계산하기 때문입니다. 4K 토큰짜리 응답이라면 1,600만 번의 어텐션 연산(operation)이 일어나고, 그 대부분은 중복(redundant)된 계산입니다. 접두부 토큰의 은닉 상태(hidden state)는 한 번 계산되면 결정적(deterministic)으로 정해지므로, 사실은 새 토큰의 질의(query)만 앞쪽에 캐싱된 키(key)와 값(value)에 적용하면 충분합니다.

여기에 더해 어텐션 자체도 데이터 이동(data movement)이 많은 연산입니다. 표준 어텐션은 N×N 점수(score) 행렬, N×d 소프트맥스(softmax) 출력, N×d 최종 출력을 모두 실체화(materialize)하기 때문에 HBM에 대한 읽기/쓰기가 너무 많습니다. N이 2K 이상이 되면 어텐션은 연산량 한계에 도달하기 전에 메모리 대역폭 한계에 먼저 도달합니다. 고전적인 어텐션 커널(kernel)은 최신 GPU의 성능을 4~10배 정도 충분히 활용하지 못합니다.

Dao 외 연구진이 제안한 두 가지 최적화가 최전선(frontier) 추론을 "느림"에서 "빠름"으로 끌어올렸습니다.

1. KV 캐시(KV cache). 접두부 토큰의 K와 V 벡터(vector)를 저장합니다. 새 토큰의 어텐션은 캐싱된 키 전체에 질의 하나를 던지는 형태로 바뀝니다. 추론은 생성 스텝당 O(N²) 재계산에서 O(N)으로 줄어듭니다.
2. Flash Attention. 어텐션 계산을 타일(tile) 단위로 쪼개서 전체 N×N 행렬이 HBM에 올라가지 않게 합니다. 소프트맥스와 행렬곱(matmul)이 모두 SRAM 안에서 끝납니다. A100에서는 24배, H100의 FP8에서는 510배의 실제 시간(wall-clock) 속도 향상을 줍니다.

2026년 기준 두 기법 모두 보편적입니다. vLLM, TensorRT-LLM, SGLang, llama.cpp 같은 운영(production) 추론 스택은 이 둘을 기본 전제로 둡니다. 최전선 모델도 Flash Attention을 켠 상태로 배포됩니다.

개념

KV 캐시 수식(KV cache math)

디코더 계층(decoder layer), 토큰, 헤드(head) 하나당 다음과 같이 계산합니다.

bytes_per_token_per_layer = 2 * d_head * dtype_size
                          ^
                          K and V

32 계층, 32 헤드, d_head=128, fp16인 7B 모델 기준입니다.

per token per layer = 2 * 128 * 2 = 512 bytes
per token (32 layers) = 16 KB
per 32K context = 512 MB

80 계층, d_head=128, KV 헤드 8개의 그룹 질의 어텐션(Grouped-Query Attention; GQA)을 쓰는 Llama 3 70B 기준입니다.

per token per layer = 2 * 8 * 128 * 2 = 4096 bytes (4 KB)
per 32K context = 10.4 GB

이 10 GB 때문에 Llama 3 70B를 128K 문맥(context)으로 돌리면 배치 크기(batch size) 1에서도 40 GB짜리 A100의 상당 부분을 KV 캐시만으로 차지합니다.

GQA가 곧 KV 캐시 측면의 승부수입니다. 64 헤드짜리 표준 다중 헤드 어텐션(Multi-Head Attention; MHA)이었다면 32 GB가 필요했을 것입니다. 잠재(Multi-head Latent) 어텐션(MLA)은 이를 더 압축합니다.

Flash Attention — 타일링 기법(tiling trick)

표준 어텐션은 다음과 같이 동작합니다.

S = Q @ K^T          (HBM read, N×N, HBM write)
P = softmax(S)       (HBM read, HBM write)
O = P @ V            (HBM read, HBM write)

HBM 왕복이 세 번 일어납니다. H100 기준 HBM 대역폭(bandwidth)은 약 3 TB/s이고 SRAM은 약 30 TB/s입니다. HBM을 한 번 거치는 비용은 칩(on-chip) 안에서 처리하는 것보다 약 10배 느립니다.

Flash Attention은 다음과 같이 동작합니다.

for each block of Q (tile size ~128 × 128):
    load Q_tile into SRAM
    for each block of K, V:
        load K_tile, V_tile into SRAM
        compute S_tile = Q_tile @ K_tile^T     (SRAM)
        running softmax aggregation             (SRAM)
        accumulate into O_tile                  (SRAM)
    write O_tile to HBM

타일 하나당 HBM 왕복은 한 번입니다. 전체 메모리 사용량(footprint)은 O(N²)에서 O(N)으로 줄어듭니다. 역방향 전파(backward pass)는 순방향 전파(forward pass)에서 값을 저장해 두는 대신 일부 값을 다시 계산(recompute)해 메모리를 추가로 절약합니다.

수치적 기법(Numerical trick). 누적(running) 소프트맥스는 타일마다 (max, sum) 쌍을 유지하기 때문에 최종 정규화(normalization)가 정확히 맞아 떨어집니다. 근사(approximation)가 아닙니다. Flash Attention은 표준 어텐션과 비트(bit) 단위로 동일한 출력을 계산합니다(fp16 비결합성으로 인한 차이 정도만 있을 수 있습니다).

버전 진화(Version evolution):

Flash 4는 출시 시점에는 순방향 전파만 지원합니다. 훈련에는 여전히 Flash 3을 사용합니다. Flash 4의 GQA와 가변 길이(varlen) 지원은 2026년 중반 기준 아직 미적용 상태입니다.

추측 디코딩(Speculative decoding) — 또 다른 지연(latency) 절감 기법

값싼 모델이 N개의 토큰을 제안하고, 큰 모델이 그 N개를 한꺼번에 검증(verify)합니다. 검증에서 k개가 받아들여지면 큰 모델의 순방향 전파 1번으로 k개의 생성을 마친 셈입니다. 코드와 산문 영역에서 일반적인 k는 3~5 정도입니다.

2026년 기본값은 다음과 같습니다.

• EAGLE 2 / Medusa. 검증 모델의 은닉 상태를 공유하는 통합 초안 헤드(integrated draft head)입니다. 품질 손실 없이 2~3배 속도 향상을 냅니다.
• 초안 모델(draft model)을 사용하는 추측 디코딩. 소비자(consumer)용 하드웨어에서도 2~4배 속도 향상을 보입니다.
• 선행 디코딩(Lookahead decoding). 자코비(Jacobi) 반복(iteration)을 사용하는 방식입니다. 초안 모델이 필요 없지만 활용 범위가 좁고(niche) 무료로 적용할 수 있다는 장점이 있습니다.

연속 배칭(Continuous batching)

고전적인 배칭(batched) 추론은 가장 느린 시퀀스(sequence)가 끝나기를 기다린 뒤에야 새 배치를 시작합니다. 짧은 응답이 먼저 끝나면 GPU 자원을 낭비합니다.

연속 배칭(처음에는 Orca에서 출시되었고 현재는 vLLM, TensorRT-LLM, SGLang에 들어가 있습니다)은 오래된 시퀀스가 끝나는 즉시 새 요청을 배치에 끼워 넣습니다. 일반적인 채팅(chat) 작업 부하(workload)에서 처리량(throughput)을 5~10배 끌어올립니다.

PagedAttention — 가상 메모리(virtual memory)로서의 KV 캐시

vLLM의 대표 기능입니다. KV 캐시를 16-토큰 블록(block) 단위로 할당하고, 페이지 테이블(page table)이 논리적(logical) 위치를 물리적(physical) 블록에 대응시킵니다. 이를 통해 빔 탐색(beam search)이나 병렬 샘플링(parallel sampling) 같은 병렬 분기 사이에서 KV를 공유하고, 프롬프트 캐싱(prompt caching)을 위해 접두부를 빠르게 교체(hot-swap)할 수 있으며, 메모리 단편화(fragmentation)도 줄일 수 있습니다. 단순한 연속(contiguous) 할당 대비 처리량을 약 4배까지 끌어올릴 수 있습니다.

직접 만들기

code/main.py에서는 다음 세 가지를 구현합니다.

1. 순진한 O(N²) 점진적(incremental) 디코더.
2. O(N) KV 캐시 기반(cached) 디코더.
3. Flash Attention의 누적 최댓값(running-max) 알고리즘을 모사(simulate)하는 타일 소프트맥스(tiled softmax).

Step 1: KV 캐시

class KVCache:
    def __init__(self, n_layers, n_heads, d_head):
        self.K = [[[] for _ in range(n_heads)] for _ in range(n_layers)]
        self.V = [[[] for _ in range(n_heads)] for _ in range(n_layers)]

    def append(self, layer, head, k, v):
        self.K[layer][head].append(k)
        self.V[layer][head].append(v)

    def read(self, layer, head):
        return self.K[layer][head], self.V[layer][head]

간단합니다. 계층별, 헤드별 리스트에 토큰 단위로 K, V 벡터를 계속 추가(append)합니다.

Step 2: 타일 소프트맥스(tiled softmax)

def tiled_softmax_dot(q, K, V, tile=4):
    """Flash-attention-style softmax(qK^T)V with running max/sum."""
    m = float("-inf")
    s = 0.0
    out = [0.0] * len(V[0])
    for start in range(0, len(K), tile):
        k_block = K[start:start + tile]
        v_block = V[start:start + tile]
        scores = [sum(qi * ki for qi, ki in zip(q, k)) for k in k_block]
        new_m = max(m, *scores)
        exp_old = math.exp(m - new_m) if m != float("-inf") else 0.0
        exp_new = [math.exp(sc - new_m) for sc in scores]
        s = s * exp_old + sum(exp_new)
        for j in range(len(out)):
            out[j] = out[j] * exp_old + sum(e * v[j] for e, v in zip(exp_new, v_block))
        m = new_m
    return [o / s for o in out]

softmax(qK) V를 한 번에 계산한 것과 동일한 출력을 내지만, 어느 순간에도 작업 공간(working set)은 전체 N × d_head가 아니라 tile × d_head 블록만큼만 차지합니다.

Step 3: 100 토큰 생성에서 순진한 방식과 캐싱 방식 비교

어텐션 연산 수를 셉니다. 순진한 방식은 O(N²) = 5050, 캐싱 방식은 O(N) = 100입니다. 코드가 두 값을 모두 출력합니다.

사용해보기

# HuggingFace transformers는 디코더 전용 generate()에서 KV 캐시를 자동으로 켭니다.
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.2-3B",
    attn_implementation="flash_attention_2",  # Hopper라면 FA3 사용
    torch_dtype="bfloat16",
)
# generate()는 KV 캐시를 자동으로 사용합니다.

vLLM 운영 예시는 다음과 같습니다.

pip install vllm
vllm serve meta-llama/Llama-3.1-70B-Instruct \
    --tensor-parallel-size 4 \
    --max-model-len 32768 \
    --enable-prefix-caching \
    --kv-cache-dtype fp8

요청 사이의 접두부 캐싱(prefix caching)은 2026년의 큰 이득 중 하나입니다. 동일한 시스템 프롬프트(system prompt), 퓨샷(few-shot) 예시, 긴 문맥 문서의 KV를 여러 호출 사이에서 재사용합니다. 도구(tool) 프롬프트가 반복되는 에이전트(agent) 작업 부하에서는 접두부 캐싱만으로도 처리량이 흔히 5배 정도 늘어납니다.

산출물 만들기

outputs/skill-inference-optimizer.md를 살펴봅니다. 이 스킬(skill)은 새 추론 배포(inference deployment)에 맞는 어텐션 구현, KV 캐시 전략, 양자화(quantization), 추측 디코딩 조합을 골라줍니다.

연습문제

1. 쉬움. code/main.py를 실행합니다. 순진한 디코더와 캐싱 디코더가 같은 출력을 내는지 확인하고, 연산 횟수(op-count) 차이를 기록합니다.
2. 중간. 접두부 캐싱을 구현합니다. 프롬프트 P와 여러 완성(completion)이 있을 때, P에 대해 순방향 전파를 한 번 돌려 KV 캐시를 채우고 완성별로 분기합니다. 매번 P를 다시 인코딩(encoding)하는 방식과 비교해 속도 향상을 측정합니다.
3. 어려움. 장난감(toy) 수준의 PagedAttention을 구현합니다. KV 캐시를 고정된 16-토큰 블록 단위로 나누고 빈 블록 목록(free-list)을 둡니다. 시퀀스가 끝나면 블록을 풀(pool)에 반납합니다. 길이가 다양한 채팅 완성 1,000건을 모사하고, 연속 할당 방식과 비교해 메모리 단편화 정도를 비교합니다.

핵심 용어

더 읽을거리

• Dao et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness — Flash 1 논문입니다.
• Dao (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning — Flash 2 논문입니다. 더 나은 병렬성과 작업 분할을 다룹니다.
• Shah et al. (2024). FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision — Flash 3 논문입니다. Hopper 비동기 처리와 저정밀도(low-precision)를 설명합니다.
• FlashAttention-4 release notes (Dao-AILab, 2026) — Blackwell 5단 파이프라인과 소프트웨어 exp2 트릭을 다룹니다. 본 강의에서 언급한 순방향 전용(forward-only) 출시 주의 사항은 저장소 README에서 확인할 수 있습니다.
• Kwon et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention — vLLM 논문입니다.
• Leviathan et al. (2023). Fast Inference from Transformers via Speculative Decoding — 추측 디코딩 논문입니다.
• Li et al. (2024). EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty — 강의에서 언급한 통합 초안(integrated-draft) 방식의 EAGLE-1/2 논문입니다.
• Cai et al. (2024). Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads — EAGLE과 함께 언급한 Medusa 방식의 논문입니다.
• vLLM docs — PagedAttention — 16-토큰 블록과 페이지 테이블 설계를 정공법으로 다룬 공식 문서입니다.