- Published on
Transformers 로컬 LLM 속도 2배 - KV캐시·PagedAttention
- Authors
- Name
- 스타차일드
- https://x.com/ETFBITX
로컬에서 LLM을 돌리면 생각보다 쉽게 “GPU는 남는데 토큰 생성이 느린” 상황을 만나게 됩니다. 대부분의 병목은 모델 파라미터 자체보다 디코딩 단계의 KV 캐시 메모리 트래픽과 배치/동시성에서의 캐시 관리 방식에서 터집니다. 이 글은 Transformers 기반 로컬 서빙에서 체감 성능을 2배 가까이 끌어올리기 위해, KV 캐시를 어떻게 줄이고(또는 효율화하고), PagedAttention 같은 페이지드 KV 관리로 단편화와 OOM을 줄이며, FlashAttention 및 컴파일 옵션으로 디코드 효율을 높이는지를 실전 관점에서 정리합니다.
메모리 최적화 전반이 필요하다면 이미지 생성 쪽이지만 접근이 유사한 글인 Stable Diffusion VRAM OOM 없애는 7가지 최적화도 함께 보면 좋습니다. “VRAM은 있는데 왜 느리지?”라는 질문의 답이 결국 메모리 트래픽과 할당 전략인 경우가 많습니다.
1) 로컬 LLM이 느려지는 진짜 이유: 프리필 vs 디코드
LLM 추론은 크게 두 구간으로 나뉩니다.
- Prefill(프리필): 입력 프롬프트 전체를 한 번에 통과시키며 KV 캐시를 채우는 단계
- Decode(디코드): 한 토큰씩 생성하면서 매 스텝 KV 캐시를 읽고 새 KV를 추가하는 단계
체감 속도(TPS, tokens per second)는 대개 디코드가 지배합니다. 디코드에서는 매 토큰마다 모든 레이어의 KV 캐시를 참조하므로, 연산량보다 메모리 대역폭과 캐시 접근 패턴이 성능을 좌우합니다.
KV 캐시가 왜 병목인가
KV 캐시는 레이어마다, 토큰마다 Key/Value 텐서를 저장합니다. 일반적인 추정식은 아래처럼 볼 수 있습니다.
- KV 메모리(바이트)
≈ batch * seq_len * n_layers * 2(K,V) * n_kv_heads * head_dim * dtype_size
여기서 중요한 포인트는 seq_len(컨텍스트 길이) 과 batch(동시 요청 수) 가 곱으로 들어간다는 점입니다. 즉,
- 컨텍스트를 2배로 늘리면 KV도 2배
- 동시 요청을 2배로 늘리면 KV도 2배
그리고 디코드에서는 이 거대한 KV를 매 스텝 읽습니다. 연산 최적화만으로는 한계가 생깁니다.
2) “속도 2배”의 현실적인 목표 설정
2배라는 수치는 환경에 따라 다르지만, 아래 조건 중 하나라도 해당하면 2배 개선이 현실적입니다.
- Transformers 기본
generate()로만 돌리고 있고, FlashAttention 미사용 - 긴 컨텍스트(예: 8k 이상)에서 TPS가 급락
- 동시성 올리면 OOM 또는 TPS 급락(단편화, 비효율적 KV 할당)
torch.compile또는 BetterTransformer, fused kernel 미적용
반대로 이미 vLLM + PagedAttention + FlashAttention2/3 조합을 잘 쓰고 있다면 2배는 어렵고, 미세 튜닝(10~30%)이 목표가 됩니다.
3) 측정부터: TPS를 프리필/디코드로 쪼개서 보자
최적화는 “측정 단위”가 있어야 합니다. 최소한 아래 3가지는 분리해서 보세요.
- 프리필 처리량(프롬프트 길이에 비례)
- 디코드 TPS(생성 토큰 수에 비례)
- VRAM 사용량(피크 및 안정 상태)
간단 벤치 코드(Transformers)
아래 코드는 단일 GPU에서 대략적인 TPS를 측정합니다. 실제로는 워밍업과 여러 번 반복이 필요합니다.
import time
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
model_id = "meta-llama/Llama-3.1-8B-Instruct" # 예시
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
model_id,
torch_dtype=torch.bfloat16,
device_map="cuda",
)
model.eval()
prompt = "다음 요구사항을 만족하는 API 설계를 제안해줘: " + "A" * 4000
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
# 워밍업
with torch.no_grad():
_ = model.generate(**inputs, max_new_tokens=32, do_sample=False)
torch.cuda.synchronize()
start = time.time()
with torch.no_grad():
out = model.generate(**inputs, max_new_tokens=256, do_sample=False)
torch.cuda.synchronize()
end = time.time()
new_tokens = out.shape[1] - inputs["input_ids"].shape[1]
print("new_tokens=", new_tokens)
print("sec=", end - start)
print("TPS=", new_tokens / (end - start))
이 수치가 “전체 TPS”라면, 이후 튜닝의 목표는 디코드 TPS를 올리는 것입니다.
4) KV 캐시 튜닝 1: dtype와 캐시 정밀도부터
가장 간단하면서 효과가 큰 축이 KV 캐시 dtype입니다.
- FP16 또는 BF16: 기본적으로 많이 사용
- FP8 KV 캐시: 지원되는 스택에서 메모리/대역폭 이점을 크게 봄(환경 의존)
Transformers 단독으로는 KV 캐시를 FP8로 “항상” 깔끔하게 내리기 어렵고, 보통은 vLLM, TensorRT-LLM, 일부 커스텀 커널에서 더 잘 지원합니다.
실전 팁
- GPU가 BF16에 강하면
torch_dtype=torch.bfloat16가 안정적입니다. - FP16에서 오버플로/언더플로 문제가 있거나 품질이 흔들리면 BF16을 우선 고려합니다.
- KV 캐시가 병목이면 “파라미터 양자화”만으로는 TPS가 잘 안 오를 수 있습니다. 파라미터는 줄었는데 KV는 그대로라서 디코드가 여전히 메모리 바운드인 경우가 많습니다.
5) KV 캐시 튜닝 2: 컨텍스트 길이와 슬라이딩 윈도우
속도와 비용을 동시에 잡는 가장 확실한 방법은 실제 필요한 컨텍스트만 남기는 것입니다.
- RAG라면 “검색 결과를 더 정확히” 해서 프롬프트 자체를 줄이기
- 대화형이라면 “요약/메모리” 전략으로 히스토리를 압축하기
- 모델이 지원하면 Sliding Window Attention 사용(일부 모델/서빙에서 가능)
컨텍스트 8k를 4k로 줄이면, 디코드 단계에서 KV 접근량이 거의 절반으로 내려가 TPS가 크게 개선됩니다.
6) PagedAttention: KV 캐시를 페이지로 쪼개서 관리하는 이유
Transformers 기본 방식은 요청마다 KV 캐시를 연속 텐서로 크게 잡는 형태가 많습니다. 동시 요청이 늘고, 요청마다 길이가 제각각이면 다음 문제가 터집니다.
- 단편화: 빈 공간이 생기는데 재사용이 어려움
- 할당/해제 오버헤드: 요청이 많을수록 allocator 부담 증가
- OOM의 비결정성: 총량은 남아 보여도 연속 블록이 부족해 OOM
PagedAttention은 KV 캐시를 “페이지(블록)” 단위로 관리해 이런 문제를 완화합니다.
- 요청의 KV는 여러 페이지로 구성
- 길이가 늘면 페이지를 추가로 붙임
- 짧은 요청이 끝나면 페이지를 풀에 반환
결과적으로 동시성에서 안정적으로 VRAM을 쓰고, 불필요한 낭비를 줄이며, 스케줄링 효율을 높여 TPS를 끌어올리는 효과가 납니다.
7) vLLM로 PagedAttention 적용: 가장 빠른 지름길
Transformers만으로 “PagedAttention급”의 이점을 얻기는 어렵고, 현실적으로는 vLLM이 가장 접근성이 좋습니다.
vLLM 실행 예시
아래는 OpenAI 호환 서버로 띄우는 예시입니다.
python -m vllm.entrypoints.openai.api_server \
--model meta-llama/Llama-3.1-8B-Instruct \
--dtype bfloat16 \
--gpu-memory-utilization 0.90 \
--max-model-len 8192 \
--max-num-batched-tokens 8192
핵심 옵션은 아래입니다.
--gpu-memory-utilization: KV 캐시 풀을 얼마나 공격적으로 잡을지 결정--max-model-len: 최악의 컨텍스트 길이에 맞춰 KV 풀/스케줄링이 잡히므로 필요 이상 크게 잡지 말 것--max-num-batched-tokens: 동적 배칭 상한. 너무 낮으면 GPU가 놀고, 너무 높으면 지연이 늘거나 메모리 압박
튜닝 순서(권장)
--max-model-len을 실제 요구사항에 맞게 최소화--gpu-memory-utilization을 0.85에서 시작해 OOM 없는 선까지 증가- 동시 요청 패턴에 맞춰
--max-num-batched-tokens조정
이 조합만으로도 “Transformers 단독 generate” 대비 체감 2배는 흔히 가능합니다(특히 동시성에서).
8) FlashAttention(또는 SDPA)로 디코드 커널 효율 올리기
Attention 계산은 커널 구현에 따라 성능이 크게 갈립니다.
- PyTorch SDPA(Scaled Dot-Product Attention): 환경에 따라 Flash 계열을 자동 선택
- FlashAttention2/3: 더 최적화된 커널(지원 GPU, 설치 상태에 따라 다름)
Transformers에서는 보통 아래 옵션/환경으로 유도합니다.
PyTorch SDPA 활성화 체크
import torch
print(torch.backends.cuda.flash_sdp_enabled())
print(torch.backends.cuda.mem_efficient_sdp_enabled())
print(torch.backends.cuda.math_sdp_enabled())
가능하면 flash 또는 mem_efficient가 켜지도록 맞추는 게 좋습니다.
Transformers에서 attention 구현 지정(모델/버전 의존)
일부 모델은 attn_implementation 인자를 받습니다.
model = AutoModelForCausalLM.from_pretrained(
model_id,
device_map="cuda",
torch_dtype=torch.bfloat16,
attn_implementation="sdpa",
)
주의할 점은 모델 아키텍처/Transformers 버전에 따라 지원 여부가 달라서, 에러가 나면 해당 모델 문서를 확인해야 합니다.
9) torch.compile: 잘 먹히면 공짜 성능, 안 먹히면 시간 낭비
torch.compile 은 모델/환경에 따라 효과 편차가 큽니다.
- 프리필에서 이득이 크고, 디코드는 제한적인 경우가 많음
- shape가 자주 바뀌는 워크로드(가변 길이)에서는 컴파일 이점이 줄어듦
그래도 시도 가치는 있습니다.
import torch
model = torch.compile(model, mode="reduce-overhead")
권장 접근은 “벤치로 확인 후 채택”입니다. 컴파일 시간이 길 수 있으니 서버 부팅 시 워밍업 전략도 함께 설계하세요.
10) 동시성에서의 속도: 배칭, 스트리밍, 큐잉
로컬 서빙에서 “속도 2배”는 단일 요청 TPS보다 동시 요청 처리량에서 더 잘 나옵니다. 핵심은 배칭입니다.
- 동시 요청을 그냥 스레드로 때리면 컨텍스트 길이가 제각각이라 GPU 활용이 들쭉날쭉
- 동적 배칭(continuous batching)이 가능한 런타임(vLLM 등)을 쓰면 디코드 스텝을 묶어서 커널 효율이 올라감
다만 배칭을 올리면 지연(latency)이 늘 수 있어, 다음을 분리해 운영하는 게 좋습니다.
- 인터랙티브(채팅): 낮은 지연 우선, 배칭 상한을 낮게
- 배치 작업(요약/분석): 처리량 우선, 배칭 상한을 높게
스트리밍을 켰을 때 429 폭주나 thundering herd 같은 문제가 생기면 요청 큐/리트라이 설계가 필요합니다. API 기반 스트리밍 운영 관점은 LangChain OpenAI 스트리밍 중 429 폭주 해결법의 “폭주 제어” 아이디어가 로컬 서빙에도 그대로 적용됩니다.
11) OOM이 아니라 “느림”도 메모리 문제다: 단편화와 allocator
로컬 LLM에서 자주 보는 증상은 이렇습니다.
- 처음엔 빠른데 시간이 갈수록 느려짐
- 간헐적으로만 OOM
nvidia-smi상 VRAM은 남아 보이는데 커널이 느림
이 경우는 KV 캐시뿐 아니라 CUDA allocator 단편화와 빈번한 할당/해제가 원인일 수 있습니다. 해결 방향은 다음과 같습니다.
- 가능한 한 “요청마다 큰 텐서를 새로 만들지 않기” (런타임이 풀링해주면 베스트)
- max length, batch 상한을 운영 요구에 맞게 고정해 shape 변동을 줄이기
- vLLM처럼 KV를 풀 기반으로 관리하는 엔진 사용
운영 중 장애/재시작 이슈까지 포함하면 쿠버네티스에서 OOM 이후 재기동, 프로브 실패로 CrashLoopBackOff 가 나기도 합니다. 인프라 레벨에서 빠르게 원인 분리하는 방법은 Kubernetes CrashLoopBackOff 원인 7가지·즉시복구도 참고할 만합니다.
12) 실전 체크리스트: 2배를 만드는 조합
아래는 “가장 흔하게” 성능이 뛰는 조합의 우선순위입니다.
(1) Transformers 단독에서 할 수 있는 것
torch_dtype=bfloat16또는float16로 통일- SDPA/Flash 계열 attention 커널이 실제로 활성화됐는지 확인
- 불필요하게 큰
max_new_tokens,max_length제한 - 프롬프트 길이 절감(요약, RAG 품질 개선)
- 가능하면
torch.compile벤치 후 적용
(2) 동시성/서빙에서 체감이 큰 것
- vLLM로 전환해 PagedAttention + continuous batching 사용
--max-model-len을 운영 요구에 맞게 낮추기--gpu-memory-utilization로 KV 풀을 안정적으로 확보- 배칭 상한(
--max-num-batched-tokens)을 트래픽 패턴에 맞게 튜닝
13) 마무리: “KV를 지배하는 자가 TPS를 지배한다”
로컬 LLM 최적화는 결국 디코드 단계의 KV 캐시 접근을 얼마나 효율적으로 만들고, 동시성에서 KV를 얼마나 똑똑하게 관리하느냐로 수렴합니다. 단일 요청만 빠르게 만드는 것보다, PagedAttention 같은 페이지 기반 KV 관리와 동적 배칭을 도입하면 “서비스로서의 처리량”이 크게 뛰는 경우가 많습니다.
다음 액션만 실행해도 성능 개선 체감이 큽니다.
- Transformers 벤치로 현재 TPS를 측정하고(프롬프트 길이별로)
- SDPA/Flash 커널 활성 여부를 확인한 뒤
- 동시 요청이 있다면 vLLM로 옮겨
--max-model-len과 KV 풀을 튜닝
이 3단계가 로컬 LLM 속도 2배에 가장 가까운, 가장 재현성 높은 루트입니다.