개념
프리필 대 디코드(Prefill vs Decode)
모든 LLM 추론 요청에는 두 단계가 있습니다.
프리필(Prefill) 은 전체 입력 프롬프트를 처리합니다. 모든 토큰이 이미 알려져 있으므로 전체 시퀀스에 대한 어텐션(attention)을 병렬로 계산할 수 있습니다. 이는 큰 행렬곱(matrix multiplication)이며 GPU 코어가 바쁘게 일합니다. 병목은 연산입니다. 하드웨어가 초당 몇 FLOPS를 낼 수 있는지가 중요합니다. A100은 BF16 기준 312 TFLOPS입니다. 단일 A100에서 70B 모델의 4,096토큰 프롬프트 프리필은 약 400ms가 걸립니다.
디코드(Decode) 는 출력 토큰을 하나씩 생성합니다. 새 토큰은 모든 이전 토큰을 참조(attend)하지만, 한 번의 순전파(forward pass)당 토큰 하나만 생성합니다. 가중치 행렬의 크기는 프리필 때와 같지만, 행렬이 아니라 벡터 하나와 곱합니다. GPU 코어는 마이크로초 단위로 계산을 끝낸 뒤, 다음 가중치 묶음이 메모리에서 도착하기를 기다립니다. 병목은 메모리 대역폭(memory bandwidth)입니다. A100은 2 TB/s 대역폭을 가집니다. FP16의 70B 모델은 140GB입니다. 모델 전체를 한 번 읽는 데 70ms가 걸립니다. 이것이 단일 디코드 스텝의 바닥값입니다.
graph LR
subgraph "Prefill (compute-bound)"
P1["All prompt tokens"] --> P2["Parallel attention"]
P2 --> P3["Full matmul utilization"]
end
subgraph "Decode (memory-bound)"
D1["One token at a time"] --> D2["Sequential generation"]
D2 --> D3["Waiting on memory reads"]
end
P3 --> D1
ops:byte 비율(ops:byte ratio), 즉 산술 강도(arithmetic intensity)는 이 절충을 포착합니다. 메모리에서 읽은 바이트당 몇 연산을 수행하는지 측정하는 지표입니다.
ops:byte ratio = FLOPs per token / bytes read from memory
4,096토큰 배치의 프리필에서는 가중치를 한 번 읽을 때 약 4,096번의 곱셈-누산(multiply-accumulate)을 수행합니다. 비율이 높고 계산 병목(compute-bound)입니다. 배치 크기 1의 디코드에서는 가중치를 한 번 읽어 거의 1번의 유용한 연산만 수행합니다. 비율이 낮고 메모리 병목(memory-bound)입니다.
핵심 통찰은 이것입니다. 디코드는 단일 토큰을 만들기 위해 모델 전체를 읽기 때문에 메모리 병목입니다. 아래의 모든 최적화는 읽을 양을 줄이거나, 한 번 읽을 때 처리하는 토큰 배치를 늘리거나, 읽기를 아예 피하는 방식입니다.
KV 캐시(KV Cache)
어텐션 중에는 각 토큰의 쿼리(query)가 이전 모든 토큰의 키(key)와 값(value) 벡터를 참조합니다. 캐싱이 없으면 토큰 N을 생성할 때 이전 N-1개 토큰의 키/값 사영(projection)을 다시 계산해야 합니다. 토큰 1은 토큰 2를 생성할 때 사영되고, 토큰 3을 생성할 때 다시 사영되고, 토큰 4에서도 다시 사영됩니다. 토큰 1,000까지 가면 토큰 1은 총 999번 사영됩니다.
KV 캐시는 이전 모든 토큰의 키/값 사영을 저장합니다. 토큰 N을 생성할 때는 토큰 N의 키/값만 계산하고, 토큰 1부터 N-1까지의 캐시된 K/V와 이어 붙입니다.
graph TD
subgraph "Without KV Cache"
A1["Token 5: recompute K,V for tokens 1-4"]
A2["Token 6: recompute K,V for tokens 1-5"]
A3["Token 7: recompute K,V for tokens 1-6"]
end
subgraph "With KV Cache"
B1["Token 5: compute K5,V5, read K1-4,V1-4 from cache"]
B2["Token 6: compute K6,V6, read K1-5,V1-5 from cache"]
B3["Token 7: compute K7,V7, read K1-6,V1-6 from cache"]
end
KV 캐시 메모리 공식:
KV cache size = 2 * num_layers * num_kv_heads * head_dim * seq_len * bytes_per_param
Llama 3 70B, 즉 80개 레이어, 그룹 쿼리 어텐션(Grouped-Query Attention; GQA)을 쓰는 8개 KV 헤드, head_dim=128, BF16 기준:
per token: 2 * 80 * 8 * 128 * 2 bytes = 327,680 bytes = 320 KB
at 4,096 tokens: 320 KB * 4,096 = 1.28 GB
at 128K tokens: 320 KB * 131,072 = 40 GB
Llama 3 70B의 단일 128K 컨텍스트 대화 하나가 KV 캐시만 40GB를 소비합니다. A100 메모리의 절반입니다. 4K 토큰 사용자 100명이 동시에 있다면 KV 캐시만 128GB가 필요합니다. 그래서 KV 캐시 관리는 추론 최적화의 핵심 과제입니다.
연속 배칭(Continuous Batching)
정적 배칭(static batching)은 N개 요청 묶음(batch)이 도착할 때까지 기다리고, 함께 처리한 뒤, 모든 요청이 끝날 때까지 새 요청을 받지 않습니다. 어떤 요청은 500토큰이 필요하고 다른 요청은 10토큰만 필요하다면, 짧은 요청은 끝난 뒤에도 490 디코드 스텝 동안 유휴 상태(idle)로 배치 슬롯을 차지합니다.
연속 배칭(continuous batching), 또는 반복 단위 배칭(iteration-level batching)은 어떤 요청이 끝나는 즉시 새 요청을 배치에 넣습니다. 배치는 매 디코드 스텝마다 다시 평가됩니다. 10토큰 뒤 끝난 요청은 즉시 대기 중인 요청으로 교체됩니다.
sequenceDiagram
participant GPU
participant R1 as Request 1 (50 tokens)
participant R2 as Request 2 (10 tokens)
participant R3 as Request 3 (30 tokens)
participant R4 as Request 4 (waiting)
Note over GPU: Static batching
GPU->>R1: Process batch [R1, R2, R3]
Note over R2: R2 done at step 10
Note over R2: Wasting 40 steps...
Note over R3: R3 done at step 30
Note over R3: Wasting 20 steps...
GPU->>R4: Finally start R4 at step 50
Note over GPU: Continuous batching
GPU->>R1: Process batch [R1, R2, R3]
Note over R2: R2 done at step 10
GPU->>R4: Insert R4 at step 11
Note over R3: R3 done at step 30
처리량 향상은 출력 길이의 변동이 얼마나 큰지에 달려 있습니다. 길이가 모두 같으면 연속 배칭은 정적 배칭과 비슷합니다. 길이가 다양한 일반적인 작업 부하(workload)에서는 GPU 슬롯이 비지 않으므로 처리량이 2-5배 높아질 수 있습니다.
PagedAttention
각 요청의 KV 캐시는 연속된(contiguous) 메모리 블록입니다. 요청이 들어오고 나가면서 메모리는 운영체제의 RAM처럼 단편화(fragmentation)됩니다. 4K 토큰 요청은 1.28GB의 연속된 블록이 필요합니다. 전체 여유 메모리가 2GB라도 1.28GB가 연속되어 있지 않으면 요청을 받을 수 없습니다. 메모리를 낭비하거나 요청을 거절해야 합니다.
PagedAttention(vLLM)은 운영체제 스타일의 가상 메모리(virtual memory)를 KV 캐시에 적용합니다. 요청마다 하나의 연속 블록을 할당하는 대신, 고정 크기 페이지(page)를 할당합니다. 보통 16토큰 단위입니다. 페이지는 GPU 물리 메모리 어디에나 있을 수 있습니다. 페이지 테이블(page table)이 요청의 논리적(logical) 시퀀스 위치를 물리적(physical) 페이지 위치로 매핑합니다.
graph TD
subgraph "Contiguous allocation"
C1["Request A: 2GB block"]
C2["[free: 0.5GB]"]
C3["Request B: 1GB block"]
C4["[free: 1.5GB -- but fragmented]"]
end
subgraph "PagedAttention"
P1["Page pool: 256 pages of 16 tokens each"]
P2["Request A: pages 3,7,12,45,88..."]
P3["Request B: pages 1,4,9,22,67..."]
P4["No fragmentation, no waste"]
end
PagedAttention은 공유 프리픽스(shared prefix)에 대한 쓰기 시 복사(copy-on-write)도 가능하게 합니다. 50개 요청이 같은 시스템 프롬프트(system prompt)를 공유하면 해당 시스템 프롬프트의 KV 캐시 페이지는 한 번만 저장되고 50개 요청이 참조합니다. 요청이 서로 다른 사용자 메시지로 갈라질 때만 자기 페이지를 얻습니다. 같은 시스템 프롬프트가 많은 애플리케이션에서 메모리를 크게 절약합니다.
vLLM은 PagedAttention으로 메모리 낭비를 거의 0에 가깝게 줄였다고 보고합니다. 단순(naive) 할당 방식의 낭비가 60-80%인 데 비해 약 4% 수준입니다.
추측 디코딩(Speculative Decoding)
디코드는 순차적이기 때문에 느립니다. 토큰 하나를 만들고, 다시 입력하고, 다음 토큰을 만듭니다. 하지만 다음 5개 토큰을 싸게 추측하고 한 번에 검증할 수 있다면 어떨까요?
추측 디코딩은 작고 빠른 드래프트 모델(draft model)이 K개의 후보 토큰(candidate token)을 생성하게 합니다. 큰 타깃 모델(target model)은 K개 후보를 단일 순전파에서 처리합니다. 이는 프리필처럼 병렬이고 계산 병목이며 효율적입니다. 타깃 모델이 드래프트 모델의 예측에 동의하면 타깃 순전파 한 번의 시간으로 K개 토큰을 모두 채택(accept)합니다. 위치 j에서 불일치가 나오면 1부터 j-1까지의 토큰을 채택하고 나머지는 버립니다.
graph LR
D["Draft model (1B)"] -->|"Generate 5 tokens<br/>~5ms"| C["Candidates: the cat sat on the"]
C --> T["Target model (70B)"]
T -->|"Verify all 5 in one pass<br/>~70ms"| V{"Match?"}
V -->|"4 of 5 match"| A["Accept 4 tokens in 75ms<br/>vs 280ms sequential"]
V -->|"Mismatch at pos 5"| R["Reject token 5<br/>Resample from target"]
속도 향상은 수락률(acceptance rate), 즉 드래프트 모델의 예측이 타깃과 얼마나 자주 일치하는지에 달려 있습니다. Llama 3 8B가 Llama 3 70B의 드래프트 역할을 하면 자연어에서 70-85%의 수락률이 흔합니다. 이는 디코드 속도 2-3배 향상으로 이어집니다.
추측 디코딩에는 세 가지 접근이 있습니다.
| 방법 | 드래프트 출처(Draft source) | 수락률(Acceptance rate) | 오버헤드(Overhead) |
|---|
| Draft-target(Leviathan et al.) | 별도의 작은 모델 | 70-85% | 드래프트 모델 메모리 |
| EAGLE(Li et al.) | 타깃 위의 경량 헤드(lightweight head) | 75-90% | 약 1% 추가 파라미터 |
| N-gram lookup | 토큰 n-gram 테이블 | 40-60% | 거의 없음 |
EAGLE은 타깃 모델의 은닉 상태(hidden state) 위에 작은 자기회귀 헤드를 학습합니다. 타깃 모델의 끝에서 두 번째 레이어의 특징(feature)을 사용해 다음 토큰의 임베딩(embedding)을 예측합니다. 별도 모델의 표현이 아니라 타깃 모델 자체의 표현 위에서 동작하므로, 매우 적은 추가 메모리로 더 높은 수락률을 얻습니다. EAGLE-2는 문맥에 따라 후보 수를 조정하는 동적 드래프트 트리(dynamic draft tree)를 추가합니다.
N-gram 추측 디코딩(N-gram speculative decoding)은 현재 문맥이나 사전 구축된 코퍼스에서 n-gram 연속(continuation) 테이블을 유지합니다. 드래프트가 같은 대화 안에서 이전에 나타난 반복 패턴, 코드, 구조화된 출력과 맞으면 신경망 비용 없이 작동합니다. 평균 수락률은 낮지만 추측 비용은 사실상 무료입니다.
추측 디코딩은 수학적으로 정확합니다. 출력 분포는 타깃 모델의 분포와 동일합니다. 근사가 아닙니다. 검증 단계가 채택된 토큰이 타깃 모델이 부여했을 확률을 정확히 갖도록 보장합니다.
프리픽스 캐싱(Prefix Caching)
많은 요청은 같은 프리픽스(prefix)를 공유합니다. 챗봇 시스템 프롬프트, RAG 컨텍스트 블록, few-shot 예시 묶음 같은 것들입니다. 프리픽스 캐싱이 없으면 모든 요청은 이 공유 토큰들의 KV 캐시를 처음부터 다시 계산합니다.
프리픽스 캐싱은 공통 프리픽스의 KV 캐시를 저장하고 요청 간에 재사용합니다. 새 요청이 알려진 프리픽스로 들어오면 캐시된 KV 항목을 복사하거나 참조하고, 고유한 접미부(unique suffix)의 KV만 계산합니다.
2,000토큰짜리 시스템 프롬프트가 모든 요청에 공유된다면 프리픽스 캐싱은 요청당 약 400ms의 프리필을 제거합니다. 초당 100요청이면 초당 40초의 GPU 연산을 절약합니다. 이는 GPU 한 장 이상의 일에 해당합니다.
SGLang의 RadixAttention은 프리픽스 캐싱을 라딕스 트리(radix tree), 즉 트라이(trie)로 구현합니다. 토큰 내용으로 프리픽스를 색인합니다. 저장된 프리픽스와 일치하는 요청은 KV 캐시를 무료로 얻습니다. 부분 프리픽스 매칭도 가능합니다. 2,000 프리픽스 토큰 중 1,500개가 캐시 항목과 같다면 1,500개는 재사용하고 500개만 다시 계산합니다.
추론 엔진(Inference Engines)
프로덕션 LLM 서빙은 세 엔진이 주도합니다.
| 엔진 | 핵심 혁신 | 가장 적합한 용도 |
|---|
| vLLM | PagedAttention, continuous batching | 범용 서빙, 가장 높은 호환성 |
| SGLang | RadixAttention(prefix caching), structured generation | 멀티턴(multi-turn) 챗봇, 제약 디코딩(constrained decoding) |
| TensorRT-LLM | NVIDIA kernel fusion, FP8 quantization | NVIDIA 하드웨어에서 단일 GPU 최대 처리량 |
vLLM은 기본 출발점입니다. 가장 넓은 모델 범위를 지원하고, NVIDIA, AMD, Intel 등 다양한 GPU 벤더에서 동작하며, PagedAttention과 연속 배칭으로 강한 처리량을 냅니다. OpenAI 호환 API(OpenAI-compatible API) 덕분에 기존 OpenAI API 호출을 쉽게 대체할 수 있습니다.
SGLang은 vLLM과 같은 기반 위에 RadixAttention 프리픽스 캐싱과 구조화된 LLM 프로그램을 위한 도메인 특화 언어(domain-specific language; DSL)를 추가합니다. 작업 부하가 멀티턴 대화, 도구 사용(tool use), 제약 디코딩(JSON 출력, 정규식 기반 생성(regex-guided generation))을 포함한다면, 프리픽스 재사용을 통해 SGLang이 vLLM보다 2-5배 빠를 수 있습니다.
TensorRT-LLM은 모델을 최적화된 NVIDIA GPU 커널로 컴파일합니다. 어텐션, 선형 연산, 활성화 함수(activation)를 하나의 커널로 융합(fuse)하고, H100에서 FP8을 사용하며, NVIDIA Triton Inference Server와 통합됩니다. NVIDIA 하드웨어에서 가장 높은 단일 GPU 처리량을 내지만 설정이 더 복잡하고 NVIDIA GPU에서만 동작합니다.
Llama 3 70B, 4xA100-80GB, BF16 기준의 실측 수치입니다.
| 지표(Metric) | vLLM | SGLang | TensorRT-LLM |
|---|
| 처리량(1 user) | 약 50 TPS | 약 55 TPS | 약 65 TPS |
| 처리량(100 users) | 총 약 2,500 TPS | 총 약 3,200 TPS | 총 약 3,000 TPS |
| 첫 토큰까지 시간(Time to first token) | 약 400ms | prefix hit 시 약 300ms | 약 350ms |
| 최대 컨텍스트(Max context) | 128K | 128K | 128K |
ops:byte 프레임워크(The Ops:Byte Framework)
측정하지 않는 것은 최적화할 수 없습니다. ops:byte 비율은 작업 부하가 계산 병목인지 메모리 병목인지 알려주고, 어떤 최적화가 중요한지 결정합니다.
Compute roof: peak FLOPS of the GPU
Memory roof: peak bandwidth * ops:byte ratio
ops:byte가 낮으면 디코드와 작은 배치처럼 메모리 대역폭 상한(memory bandwidth roof)에 부딪힙니다. 더 많은 연산, 높은 클럭(clock), 많은 코어는 도움이 되지 않습니다. 메모리 읽기를 줄이거나, 양자화(quantization)와 KV 캐시 압축(KV cache compression)을 쓰거나, 배치 크기를 늘려 한 번 읽은 가중치를 더 많은 유용한 작업에 분산(amortize)시켜야 합니다.
ops:byte가 높으면 프리필과 큰 배치처럼 연산 상한(compute roof)에 부딪힙니다. 메모리 대역폭 최적화는 큰 도움이 되지 않습니다. 더 빠른 GPU, 커널 융합(kernel fusion), 정밀도 축소(reduced precision)로 FLOPS를 더 뽑아야 합니다.
| 시나리오(Scenario) | ops:byte | 병목(Bound) | 최적화 수단(Optimize with) |
|---|
| Prefill, batch=1 | 약 4,096 | Compute | Kernel fusion, FP8 |
| Decode, batch=1 | 약 1 | Memory | Quantization, KV compression |
| Decode, batch=32 | 약 32 | Memory | Larger batch, continuous batching |
| Decode, batch=256 | 약 256 | Transitioning | Both matter |
| Decode, batch=1024 | 약 1,024 | Compute | Kernel fusion, tensor parallelism |
A100의 교차점(crossover point)은 ops:byte 약 156입니다. 312 TFLOPS / 2 TB/s로 계산됩니다. 156보다 낮으면 메모리 병목이고, 높으면 계산 병목입니다. 연속 배칭은 반복(iteration)당 더 많은 토큰을 묶어 디코드를 이 교차점 쪽으로 밀어 올립니다.
직접 만들기
Step 1: KV 캐시를 처음부터 만들기
레이어, 헤드, 헤드 차원(head dimension)별로 키/값 사영을 저장하는 멀티헤드 KV 캐시(multi-head KV cache)를 만들고, 메모리 증가 패턴을 확인합니다.
import numpy as np
class KVCache:
def __init__(self, num_layers, num_heads, head_dim, max_seq_len, dtype=np.float16):
self.num_layers = num_layers
self.num_heads = num_heads
self.head_dim = head_dim
self.max_seq_len = max_seq_len
self.dtype = dtype
self.k_cache = np.zeros(
(num_layers, num_heads, max_seq_len, head_dim), dtype=dtype
)
self.v_cache = np.zeros(
(num_layers, num_heads, max_seq_len, head_dim), dtype=dtype
)
self.seq_len = 0
def update(self, layer_idx, new_keys, new_values):
num_new = new_keys.shape[1]
end = self.seq_len + num_new
self.k_cache[layer_idx, :, self.seq_len:end, :] = new_keys
self.v_cache[layer_idx, :, self.seq_len:end, :] = new_values
return (
self.k_cache[layer_idx, :, :end, :],
self.v_cache[layer_idx, :, :end, :]
)
def advance(self, num_tokens):
self.seq_len += num_tokens
def memory_bytes(self):
return self.k_cache.nbytes + self.v_cache.nbytes
def used_bytes(self):
per_token = 2 * self.num_layers * self.num_heads * self.head_dim * np.dtype(self.dtype).itemsize
return per_token * self.seq_len
Step 2: KV 캐시를 사용하는 어텐션
디코드 스텝에서 KV 캐시를 사용하는 단순화된 멀티헤드 어텐션을 구현합니다.
def scaled_dot_product_attention(query, keys, values):
head_dim = query.shape[-1]
scores = np.matmul(query, keys.transpose(0, 1, 3, 2)) / np.sqrt(head_dim)
seq_len_q = scores.shape[-2]
seq_len_k = scores.shape[-1]
if seq_len_q > 1:
mask = np.triu(np.ones((seq_len_q, seq_len_k), dtype=np.float32), k=seq_len_k - seq_len_q + 1)
scores = scores + mask * (-1e9)
max_scores = np.max(scores, axis=-1, keepdims=True)
exp_scores = np.exp(scores - max_scores)
attn_weights = exp_scores / np.sum(exp_scores, axis=-1, keepdims=True)
return np.matmul(attn_weights, values)
class MultiHeadAttention:
def __init__(self, d_model, num_heads):
self.num_heads = num_heads
self.head_dim = d_model // num_heads
scale = np.sqrt(2.0 / d_model)
self.W_q = np.random.randn(d_model, d_model).astype(np.float32) * scale
self.W_k = np.random.randn(d_model, d_model).astype(np.float32) * scale
self.W_v = np.random.randn(d_model, d_model).astype(np.float32) * scale
self.W_o = np.random.randn(d_model, d_model).astype(np.float32) * scale
def forward(self, x, kv_cache=None, layer_idx=0):
batch, seq_len, d_model = x.shape
Q = np.matmul(x, self.W_q).reshape(batch, seq_len, self.num_heads, self.head_dim).transpose(0, 2, 1, 3)
K = np.matmul(x, self.W_k).reshape(batch, seq_len, self.num_heads, self.head_dim).transpose(0, 2, 1, 3)
V = np.matmul(x, self.W_v).reshape(batch, seq_len, self.num_heads, self.head_dim).transpose(0, 2, 1, 3)
if kv_cache is not None:
K_full, V_full = kv_cache.update(layer_idx, K[0], V[0])
K = K_full[np.newaxis, :, :, :]
V = V_full[np.newaxis, :, :, :]
if seq_len == 1:
kv_cache.advance(1)
attn_out = scaled_dot_product_attention(Q, K, V)
attn_out = attn_out.transpose(0, 2, 1, 3).reshape(batch, -1, d_model)
return np.matmul(attn_out, self.W_o)
Step 3: 연속 배칭 시뮬레이터
정적 배칭과 연속 배칭의 스케줄링 차이를 시뮬레이션합니다.
import heapq
class Request:
def __init__(self, request_id, prompt_tokens, output_tokens, arrival_step):
self.request_id = request_id
self.prompt_tokens = prompt_tokens
self.output_tokens = output_tokens
self.arrival_step = arrival_step
self.tokens_generated = 0
self.start_step = None
self.end_step = None
def is_done(self):
return self.tokens_generated >= self.output_tokens
def simulate_static_batching(requests, batch_size):
step = 0
completed = []
queue = list(requests)
queue.sort(key=lambda r: r.arrival_step)
while queue:
batch = []
while queue and len(batch) < batch_size:
r = queue.pop(0)
r.start_step = max(step, r.arrival_step)
batch.append(r)
if batch:
step = max(step, max(r.start_step for r in batch))
max_output = max(r.output_tokens for r in batch)
for r in batch:
r.tokens_generated = r.output_tokens
r.end_step = step + max_output
step += max_output
completed.extend(batch)
return completed
def simulate_continuous_batching(requests, batch_size):
step = 0
completed = []
queue = sorted(requests, key=lambda r: r.arrival_step)
queue_idx = 0
active = []
waiting = []
while queue_idx < len(queue) or active or waiting:
while queue_idx < len(queue) and queue[queue_idx].arrival_step <= step:
waiting.append(queue[queue_idx])
queue_idx += 1
while waiting and len(active) < batch_size:
r = waiting.pop(0)
r.start_step = step
active.append(r)
if not active:
if waiting:
step += 1
continue
elif queue_idx < len(queue):
step = queue[queue_idx].arrival_step
continue
else:
break
for r in active:
r.tokens_generated += 1
done = [r for r in active if r.is_done()]
for r in done:
r.end_step = step + 1
completed.append(r)
active = [r for r in active if not r.is_done()]
step += 1
return completed
def batching_stats(completed):
latencies = [r.end_step - r.arrival_step for r in completed]
total_time = max(r.end_step for r in completed) - min(r.arrival_step for r in completed)
total_tokens = sum(r.output_tokens for r in completed)
return {
"avg_latency": np.mean(latencies),
"p50_latency": np.median(latencies),
"p99_latency": np.percentile(latencies, 99),
"total_time": total_time,
"throughput": total_tokens / total_time if total_time > 0 else 0,
}
Step 4: 프리픽스 캐시
공유 프리픽스의 KV 항목을 저장하는 트라이(trie) 기반 프리픽스 캐시를 구현합니다.
class TrieNode:
def __init__(self):
self.children = {}
self.kv_data = None
self.hit_count = 0
class PrefixCache:
def __init__(self, max_entries=1000):
self.root = TrieNode()
self.max_entries = max_entries
self.total_entries = 0
self.hits = 0
self.misses = 0
def _walk(self, token_ids):
node = self.root
depth = 0
for tid in token_ids:
if tid not in node.children:
break
node = node.children[tid]
depth += 1
return node, depth
def lookup(self, token_ids):
node, depth = self._walk(token_ids)
if depth > 0:
self.hits += 1
current = self.root
for tid in token_ids[:depth]:
current = current.children[tid]
current.hit_count += 1
kv_entries = []
current = self.root
for tid in token_ids[:depth]:
current = current.children[tid]
if current.kv_data is not None:
kv_entries.append(current.kv_data)
return depth, kv_entries
self.misses += 1
return 0, []
def insert(self, token_ids, kv_per_token):
node = self.root
for i, tid in enumerate(token_ids):
if tid not in node.children:
if self.total_entries >= self.max_entries:
return i
node.children[tid] = TrieNode()
self.total_entries += 1
node = node.children[tid]
if i < len(kv_per_token):
node.kv_data = kv_per_token[i]
return len(token_ids)
def hit_rate(self):
total = self.hits + self.misses
return self.hits / total if total > 0 else 0.0
Step 5: 추측 디코딩 시뮬레이터
수락률을 조절할 수 있는 드래프트-타깃 추측 디코딩을 시뮬레이션합니다.
class DraftModel:
def __init__(self, vocab_size, acceptance_rate=0.8):
self.vocab_size = vocab_size
self.acceptance_rate = acceptance_rate
def generate(self, context, num_tokens):
tokens = np.random.randint(0, self.vocab_size, size=num_tokens)
return tokens
def get_probs(self, context, token):
probs = np.random.dirichlet(np.ones(self.vocab_size))
return probs
class TargetModel:
def __init__(self, vocab_size):
self.vocab_size = vocab_size
def get_probs(self, context, tokens=None):
if tokens is not None:
return [np.random.dirichlet(np.ones(self.vocab_size)) for _ in tokens]
return np.random.dirichlet(np.ones(self.vocab_size))
def speculative_decode(draft_model, target_model, context, num_speculative=5,
draft_cost=1.0, target_cost=10.0, verify_cost=12.0):
total_tokens = 0
total_cost = 0.0
accepted_counts = []
context = list(context)
max_tokens = 100
while total_tokens < max_tokens:
draft_tokens = draft_model.generate(context, num_speculative)
total_cost += draft_cost * num_speculative
target_probs = target_model.get_probs(context, draft_tokens)
total_cost += verify_cost
accepted = 0
for i, token in enumerate(draft_tokens):
draft_p = draft_model.get_probs(context + list(draft_tokens[:i]), token)
target_p = target_probs[i]
r = np.random.random()
acceptance_prob = min(1.0, target_p[token] / (draft_p[token] + 1e-10))
if r < draft_model.acceptance_rate:
accepted += 1
context.append(token)
total_tokens += 1
else:
new_token = np.random.choice(draft_model.vocab_size, p=target_p)
context.append(new_token)
total_tokens += 1
break
accepted_counts.append(accepted)
if accepted == num_speculative:
bonus_probs = target_model.get_probs(context)
bonus_token = np.random.choice(draft_model.vocab_size, p=bonus_probs)
context.append(bonus_token)
total_tokens += 1
sequential_cost = total_tokens * target_cost
return {
"total_tokens": total_tokens,
"speculative_cost": total_cost,
"sequential_cost": sequential_cost,
"speedup": sequential_cost / total_cost if total_cost > 0 else 1.0,
"avg_accepted": np.mean(accepted_counts),
"acceptance_rate": np.mean(accepted_counts) / num_speculative,
}
def compare_speculation_strategies(vocab_size=1000, num_trials=20):
results = {}
for name, acceptance_rate, spec_tokens in [
("Draft-target (8B->70B)", 0.78, 5),
("EAGLE", 0.85, 6),
("N-gram", 0.50, 4),
("No speculation", 0.0, 0),
]:
if spec_tokens == 0:
results[name] = {
"speedup": 1.0,
"acceptance_rate": 0.0,
"avg_accepted": 0.0,
}
continue
trial_results = []
for _ in range(num_trials):
draft = DraftModel(vocab_size, acceptance_rate=acceptance_rate)
target = TargetModel(vocab_size)
context = list(np.random.randint(0, vocab_size, size=10))
result = speculative_decode(draft, target, context, num_speculative=spec_tokens)
trial_results.append(result)
results[name] = {
"speedup": np.mean([r["speedup"] for r in trial_results]),
"acceptance_rate": np.mean([r["acceptance_rate"] for r in trial_results]),
"avg_accepted": np.mean([r["avg_accepted"] for r in trial_results]),
}
return results
Step 6: KV 캐시 메모리 프로파일러
실제 모델 설정에 대한 KV 캐시 메모리 요구량을 계산합니다.
MODEL_CONFIGS = {
"Llama-3-8B": {
"num_layers": 32, "num_kv_heads": 8, "head_dim": 128,
"model_params_b": 8, "gqa": True,
},
"Llama-3-70B": {
"num_layers": 80, "num_kv_heads": 8, "head_dim": 128,
"model_params_b": 70, "gqa": True,
},
"Llama-3-405B": {
"num_layers": 126, "num_kv_heads": 8, "head_dim": 128,
"model_params_b": 405, "gqa": True,
},
"Mistral-7B": {
"num_layers": 32, "num_kv_heads": 8, "head_dim": 128,
"model_params_b": 7, "gqa": True,
},
"GPT-4-est": {
"num_layers": 120, "num_kv_heads": 96, "head_dim": 128,
"model_params_b": 1800, "gqa": False,
},
}
def kv_cache_memory(config, seq_len, dtype_bytes=2):
per_token = 2 * config["num_layers"] * config["num_kv_heads"] * config["head_dim"] * dtype_bytes
total = per_token * seq_len
return {
"per_token_bytes": per_token,
"per_token_kb": per_token / 1024,
"total_bytes": total,
"total_mb": total / (1024 ** 2),
"total_gb": total / (1024 ** 3),
}
def memory_budget(config, gpu_memory_gb, model_dtype_bytes=2, kv_dtype_bytes=2):
model_memory_gb = config["model_params_b"] * 1e9 * model_dtype_bytes / (1024 ** 3)
overhead_gb = gpu_memory_gb * 0.1
available_for_kv = gpu_memory_gb - model_memory_gb - overhead_gb
if available_for_kv <= 0:
return {"error": "Model does not fit in GPU memory", "model_memory_gb": model_memory_gb}
per_token = 2 * config["num_layers"] * config["num_kv_heads"] * config["head_dim"] * kv_dtype_bytes
max_tokens = int(available_for_kv * (1024 ** 3) / per_token)
return {
"gpu_memory_gb": gpu_memory_gb,
"model_memory_gb": round(model_memory_gb, 1),
"overhead_gb": round(overhead_gb, 1),
"available_for_kv_gb": round(available_for_kv, 1),
"max_total_tokens": max_tokens,
"max_users_at_2k": max_tokens // 2048,
"max_users_at_4k": max_tokens // 4096,
"max_users_at_32k": max_tokens // 32768,
}
사용해보기
vLLM 사용:
from vllm import LLM, SamplingParams
llm = LLM(
model="meta-llama/Llama-3-70B-Instruct",
tensor_parallel_size=4,
enable_prefix_caching=True,
max_model_len=8192,
gpu_memory_utilization=0.9,
)
params = SamplingParams(temperature=0.7, max_tokens=256)
outputs = llm.generate(["Explain inference optimization in one paragraph."], params)
프리픽스 캐싱과 구조화된 출력을 위한 SGLang 사용:
import sglang as sgl
@sgl.function
def classify(s, text):
s += sgl.system("You are a classifier. Output JSON only.")
s += sgl.user(f"Classify this text: {text}")
s += sgl.assistant(sgl.gen("result", regex=r'\{"label": "(positive|negative|neutral)"\}'))
runtime = sgl.Runtime(model_path="meta-llama/Llama-3-70B-Instruct", tp_size=4)
sgl.set_default_backend(runtime)
results = classify.run_batch([
{"text": "This product is amazing!"},
{"text": "Terrible experience."},
{"text": "It was okay I guess."},
])
TensorRT-LLM 사용:
import tensorrt_llm
from tensorrt_llm.runtime import ModelRunner
runner = ModelRunner.from_dir("./llama-70b-trt-engine/", rank=0)
outputs = runner.generate(
batch_input_ids=[tokenizer.encode("Explain KV caching.")],
max_new_tokens=256,
temperature=0.7,
)
산출물 만들기
이 강의는 다음을 만듭니다.
outputs/skill-inference-optimization.md — LLM 추론 서빙의 처리량, 지연 시간, 비용을 진단하고 최적화하기 위한 스킬(skill).
연습문제
-
(쉬움) KV 캐시 프로파일러를 수정해 FP16, FP8, INT4 KV 캐시 양자화를 비교합니다. Llama 3 70B, 4K 컨텍스트에서 4xA100-80GB 기준 각 정밀도(precision)의 최대 동시 사용자 수를 계산합니다. KV 양자화를 INT4로 낮추면 사용자 수용량이 대략 4배가 되어야 합니다.
-
(중간) 연속 배칭 시뮬레이터를 확장해 GPU 사용률(스텝마다 채워진 배치 슬롯 비율)을 추적합니다. 출력 길이가 파레토 분포(Pareto distribution; shape=1.5, scale=20)를 따르는 50개 요청에 대해 정적 배칭과 연속 배칭의 사용률을 시간에 따라 그립니다. 연속 배칭은 80% 이상의 사용률을 유지해야 합니다.
-
(중간) num_kv_heads < num_query_heads인 그룹 쿼리 어텐션(GQA) 버전의 KV 캐시를 구현합니다. Llama 3 70B는 64개의 쿼리 헤드와 8개의 KV 헤드를 사용합니다. 전체 멀티헤드 어텐션 대비 KV 캐시 크기의 메모리 절감량(8배 감소)을 계산합니다.
-
(중간) LRU(Least Recently Used) 축출(eviction)을 사용하는 프리픽스 캐시를 만듭니다. max_entries=500으로 두고, 1,000개 요청 중 60%가 5개 공통 프리픽스 중 하나를 공유하도록 생성합니다. 적중률(hit rate)을 측정하고 무제한 캐시와 비교합니다. 좋은 축출 정책이라면 적중률이 55% 이상 유지되어야 합니다.
-
(어려움) 추측 디코딩 시뮬레이터를 트리 기반 추측(EAGLE-2 스타일)으로 확장합니다. K개 토큰의 단일 체인 대신 후보 트리(예: 3개 레벨에서 각 2 분기 = 8개 리프 후보)를 생성합니다. 검증 라운드당 채택된 토큰 수를 선형 추측과 비교합니다.
핵심 용어
| 용어 | 흔한 설명 | 실제 의미 |
|---|
| 프리필(Prefill) | "프롬프트 처리" | 모든 입력 토큰에 대해 병렬 어텐션을 계산하는 단계입니다. 큰 행렬곱이 GPU 코어를 바쁘게 하므로 계산 병목(compute-bound)입니다. |
| 디코드(Decode) | "토큰 생성" | 한 번의 순전파당 토큰 하나를 생성하며, 매번 전체 모델 가중치를 읽습니다. 다음 가중치가 도착하기 전에 연산이 끝나므로 메모리 병목(memory-bound)입니다. |
| KV 캐시(KV cache) | "어텐션 상태 캐싱" | 이전 모든 토큰의 키/값 사영을 저장해 각 디코드 스텝에서 다시 계산하지 않도록 합니다. 연산을 줄이기 위해 메모리를 사용합니다. |
| 연속 배칭(Continuous batching) | "동적 배칭" | 전체 배치 완료를 기다리지 않고, 각 디코드 반복마다 완료된 요청 자리에 새 요청을 넣는 스케줄링입니다. |
| PagedAttention | "KV 캐시를 위한 가상 메모리" | KV 캐시를 연속 블록이 아닌 고정 크기 페이지로 할당해 단편화를 없애고, 공유 프리픽스에 대한 쓰기 시 복사를 가능하게 합니다. |
| 추측 디코딩(Speculative decoding) | "초안 작성 후 검증" | 빠른 드래프트 모델이 여러 토큰을 제안하고 타깃 모델이 한 번에 검증합니다. 수학적으로 정확하며 보통 2-3배 빨라집니다. |
| EAGLE | "자기 추측 디코딩(self-speculative decoding)" | 타깃 모델의 은닉 상태 위에 경량 헤드를 학습해 별도 드래프트 모델보다 높은 수락률을 얻는 추측 디코딩 변형입니다. |
| 프리픽스 캐싱(Prefix caching) | "시스템 프롬프트 KV 재사용" | 시스템 프롬프트, few-shot 예시 같은 공통 프리픽스의 KV 캐시를 저장해 요청 간 재사용하고 중복된 프리필을 건너뜁니다. |
| ops:byte 비율(Ops:byte ratio) | "산술 강도(arithmetic intensity)" | 메모리 바이트당 연산 횟수의 비율입니다. 작업 부하가 계산 병목인지 메모리 병목인지를 결정합니다. |
| 첫 토큰까지 시간(Time to first token; TTFT) | "TTFT" | 요청을 받은 뒤 첫 출력 토큰이 나올 때까지의 지연 시간입니다. 긴 프롬프트에서는 프리필 시간이 지배합니다. |
더 읽을거리
- Kwon et al., "Efficient Memory Management for Large Language Model Serving with PagedAttention" (2023) — 페이지 단위 KV 캐시 관리를 도입한 vLLM 논문이며, 이제는 추론 서빙의 업계 표준입니다.
- Leviathan et al., "Fast Inference from Transformers via Speculative Decoding" (2023) — 드래프트-검증 추측이 타깃 모델의 정확한 분포를 유지하면서 2-3배의 속도 향상을 얻을 수 있음을 증명한 기초 논문입니다.
- Li et al., "EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty" (2024) — 별도의 드래프트 모델 대신 타깃 모델 자체의 특징(feature) 위에 헤드를 학습해 더 높은 수락률을 얻는 방법입니다.
- Zheng et al., "SGLang: Efficient Execution of Structured Language Model Programs" (2024) — 프리픽스 캐싱을 위한 RadixAttention과 다중 호출 LLM 프로그램을 위한 프로그래밍 모델을 소개합니다.
- Williams et al., "Roofline: An Insightful Visual Performance Model for Multicore Architectures" (2009) — 연산과 메모리 병목을 추론하기 위한 ops:byte 프레임워크를 정식화한 원 논문입니다.