추론 시스템·커널 최적화 — FlashAttention, PagedAttention, Continuous Batching, Prefix Caching

한 줄 요약: 같은 모델·같은 GPU에서 더 빨리, 더 많은 사용자를, 더 싸게 굴리는 시스템 레이어 기법들이에요.

연관 문서:

• 파라미터 착시: KV 캐시의 진실 — 본 부록이 줄이려는 그 KV 캐시의 정체
• LLM 구동의 주요 스펙: 메모리 용량과 대역폭 — HBM 대역폭이 왜 처리량 상한이 되는가
• DFlash — 블록 확산 드래프터 — 같은 시스템 레이어의 다른 갈래(투기적 디코딩)

0. 출발점 — 어디가 병목인가

같은 70B 모델을 RTX 4090 한 장으로 12 tok/s 굴릴 때 H100 SXM에서 60 tok/s 가 나오는 이유는, 메모리 대역폭 부록에서 본 것처럼 거의 전적으로 HBM 대역폭 한 줄로 설명돼요. 그런데 같은 GPU 안에서도 추론 스택을 어떻게 짜느냐에 따라 처리량이 또 한 자릿수 이상 갈립니다. 같은 H100에 같은 Llama-13B 가 올라가 있어도, naive 한 PyTorch 루프로 굴릴 때와 vLLM 으로 굴릴 때의 시간당 처리 가능한 사용자 수가 14~24배 갈려요(vLLM 공식 보고).

이 차이가 어디서 오느냐를 한 줄로 정리하면 네 군데예요.

각 항목은 서로 독립적이라서 함께 적용할 수 있고, 실제로 vLLM·SGLang·TGI 같은 현대 서빙 엔진은 네 가지를 모두 켜놓고 돌아갑니다. 본문 §2에서 "GPU를 두 배 비싼 걸로 바꾸면 응답이 두 배 빨라지냐"라고 물었을 때의 소프트웨어 쪽 답이 이 네 기법이에요.

1. FlashAttention — HBM 왕복을 없애는 커널

1.1 표준 attention 의 메모리 패턴

표준 attention 은 길이 N 시퀀스에 대해 N×N 행렬 두 개(어텐션 점수 S, 소프트맥스 출력 P)를 만들어서 HBM 에 적었다가 다시 읽어요. 시퀀스가 길어지면 이 N×N 행렬이 메모리도, 대역폭도 다 잡아먹습니다.

문제는 GPU 의 SRAM 은 작고 빠르고, HBM 은 크고 느린 계층 구조라는 점이에요. SRAM 대역폭은 HBM 대역폭의 보통 10배 이상인데, 표준 attention 은 중간 결과를 굳이 HBM 에 쓰고 다시 읽으면서 그 빠른 SRAM 을 활용하지 못하고 있었어요. 어텐션이 메모리 바운드(memory-bound — 연산보다 메모리 이동이 더 오래 걸리는 상태) 라는 말의 구체적 의미가 이거예요.

HBM (High Bandwidth Memory): GPU 본체 옆에 붙은 큰 메모리. 대역폭은 H100 SXM 기준 ~3.35 TB/s (상세). SRAM: GPU 코어 바로 옆 캐시 같은 메모리. 한 SM(streaming multiprocessor) 당 100 KB 안팎으로 작지만 대역폭은 수십 TB/s.

1.2 핵심 트릭 — Tiling + Online Softmax

FlashAttention(Dao et al., 2022)이 한 일은 두 가지예요.

1. Tiling — Q, K, V 를 작은 블록으로 잘라서 SRAM 에 한 블록씩 올리고, 그 안에서 어텐션 한 조각을 끝까지 계산.
2. Online softmax — softmax 는 보통 "전체 점수를 다 본 다음 정규화" 가 필요한데, 블록을 순회하면서 러닝 통계(최댓값·합) 를 갱신해 가며 정규화. 즉 N×N 행렬을 통째로 만들지 않고도 정확한 softmax 를 얻습니다.

핵심은 N×N 어텐션 행렬을 HBM 에 올리지 않고도 동일한 결과 O 를 얻는다는 점이에요. 메모리 복잡도는 O(N²) 에서 O(N) 으로 떨어지고, HBM 트래픽은 SRAM 크기 M 에 반비례해서 줄어들어요.

한 줄 비유: 표준 attention 이 큰 노트(N×N) 한 장에 다 적었다가 다시 읽는 학생이라면, FlashAttention 은 작은 메모지(SRAM 블록) 한 장씩 돌려가며 머릿속(러닝 통계)에서 합산해서 최종 답안만 적는 학생이에요.

1.3 FA-v1 — 메모리는 O(N), 속도는 ×2~4

FlashAttention v1 의 보고 수치(논문 Abstract·§3):

부수 효과로 컨텍스트를 더 길게 쓸 수 있게 됐어요. Path-X(seq 16K) 에서 Transformer 가 처음으로 무작위 이상의 정확도(61.4%) 를 기록했고, Path-256(seq 64K) 에서도 63.1% 를 찍었습니다. 메모리 O(N²) → O(N) 으로 떨어지면서 시퀀스를 한 자릿수 더 늘릴 여력이 생긴 결과예요.

저자: Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré (Stanford, 2022-05).

1.4 FA-v2 — A100 에서 v1 대비 ×2

FlashAttention-2 (Dao, 2023-07) 가 v1 위에 얹은 세 가지 변경:

1. non-matmul FLOPs 감소 — softmax 의 rescale 같은 비-행렬곱 연산 횟수를 줄여요. GPU 텐서 코어는 행렬곱에서만 진가를 발휘해서, non-matmul 은 곱셈기 입장에서 빈손이에요.
2. 시퀀스 축 병렬화 — 한 헤드 안에서도 시퀀스 길이 방향으로 thread block 을 나눠서 GPU 점유율(occupancy) 을 높여요.
3. warp 단위 작업 분배 — 한 thread block 안의 warp 들 사이 공유 메모리 통신을 줄이는 방식으로 work 을 재분배.

수치(논문 Abstract·§3):

MFU (Model FLOPs Utilization): GPU 가 이론 최대 FLOPs/s 대비 실제로 모델 학습/추론에 쓴 비율. 72% 는 행렬곱(GEMM) 에 가까운 수치예요.

FA-v2 가 등장한 시점부터 PyTorch 의 scaled_dot_product_attention (SDPA), vLLM, TGI 같은 메이저 스택이 기본 백엔드로 깔기 시작했습니다.

1.5 FA-v3 — Hopper 비동기 + FP8

FlashAttention-3 (Shah, Bikshandi, Zhang, Thakkar, Ramani, Dao, 2024-07) 는 H100 의 새 하드웨어 기능을 정면으로 활용해요.

• Tensor Core ↔ TMA 비동기 실행 — 데이터 이동(TMA) 과 행렬곱(WGMMA) 을 warp specialization 으로 겹쳐서 하나가 멈춰도 다른 하나가 진행.
• block-wise matmul 과 softmax 인터리빙 — softmax 가 끝날 때까지 행렬곱이 노는 것을 막아요.
• FP8 양자화 + incoherent processing — H100 의 FP8 텐서 코어를 활용하면서 수치 안정성도 확보. incoherent processing: FP8 양자화에서 outlier(튀는 큰 값) 가 분산되도록 random orthogonal 변환을 가해 분포를 평탄화하는 기법. 큰 값 하나가 양자화 오차를 좌우하는 걸 막아요.

수치(논문 Abstract):

H100 위에서 FA-2 를 그대로 돌리면 35% 활용률에 그쳤는데, FA-3 는 75% 까지 끌어올렸다는 점이 핵심이에요. 같은 GPU 의 비동기 기능을 안 쓰면 그만큼이 그냥 버려지고 있었다는 뜻입니다.

1.6 어디서 쓰이나

• PyTorch SDPA — 2.0 부터 기본 백엔드로 FA 계열 커널 선택.
• vLLM — 기본 attention 백엔드.
• TGI (Text Generation Inference) — Hugging Face 의 서빙 엔진.
• flash-attn 패키지 — Tri Dao 가 직접 관리하는 CUDA 구현.

1.7 한계·전제

• 실용 헤드 차원 d ≤ 256 에서 가장 큰 효과. 헤드 차원이 작으면 SRAM 블록이 더 커서 효율이 좋아져요.
• FA-v3 는 Hopper(H100/H200) 전용. Ampere(A100) 에는 v2 까지만 의미. (주의) Blackwell(B100/B200) 대응 커널은 시점에 따라 별도 구현이 필요.
• 학습용·추론용 코드 경로가 부분적으로 다릅니다(backward 가 있느냐 없느냐).

2. PagedAttention — KV 캐시 페이징

2.1 문제 — KV 캐시 단편화로 60–80% 가 낭비

표준 서빙 시스템은 요청 하나가 들어오면 그 요청의 최대 길이만큼 KV 캐시 메모리를 미리 한 덩어리로 잡아둡니다. 문제는 두 가지예요.

1. 내부 단편화(internal fragmentation) — 실제로는 짧게 끝나는 요청이 많은데도 최대 길이만큼 잡혀 있음. 안 쓴 자리가 그대로 죽은 메모리.
2. 외부 단편화(external fragmentation) — 요청마다 잡아둔 덩어리 사이에 자투리가 생겨, 새 요청을 위한 큰 연속 공간이 없음.

vLLM 공식 보고(Kwon et al., SOSP 2023; vLLM 런치 블로그 2023-06)에 따르면 이 두 가지로 KV 캐시 메모리의 60–80% 가 낭비되고 있었어요. 그 결과 한 GPU 에 동시에 올릴 수 있는 요청 수(batch size) 가 작아지고, 처리량 상한이 낮아집니다.

2.2 OS 가상메모리 차용

한 줄 비유: OS 가 프로세스 하나에 메모리를 줄 때 큰 연속 공간을 통째로 주는 게 아니라 4 KB 페이지 단위로 잘게 잘라서 비연속 물리 위치에 분산 배치하고, 페이지 테이블로 연결해두잖아요. PagedAttention 은 그걸 KV 캐시에 그대로 가져온 거예요.

2.3 메커니즘 — 블록 테이블 + 비연속 KV 블록

KV 캐시를 고정 크기 블록(보통 16~32 토큰 단위)으로 잘라요. 각 요청은 자신의 블록 테이블 만 가지고 있고, 실제 KV 블록은 GPU 메모리 어디에 흩어져 있어도 무방.

새 요청이 오면 빈 블록만 골라서 할당하면 되고, 요청이 끝나면 그 블록만 풀에 반납. 큰 연속 공간을 미리 잡아둘 필요가 없어서 단편화가 사라져요.

부가 효과로 블록 공유가 가능해집니다. 두 요청이 같은 시스템 프롬프트로 시작하면, 그 prefix 에 해당하는 KV 블록을 둘 다 같은 물리 블록을 가리키게 만들 수 있어요(§4 Prefix Caching 으로 이어집니다).

2.4 효과 — 메모리 낭비 4% 미만, 처리량 2~24×

vLLM 공식 보고:

중요한 단서: vLLM 의 14~24× 수치는 PagedAttention 단독이 아니라 continuous batching(§3) 과의 결합 효과예요. 둘이 시너지로 동작합니다. 단편화를 없애면 batch size 를 키울 수 있고, batch size 를 키우면 continuous batching 이 더 큰 효과를 내요.

2.5 구현

• vLLM — PagedAttention 의 원조 구현(2023).
• TGI — 자체 paged KV cache 로 따라옴.
• SGLang — paged 레이아웃 기반 위에 RadixAttention(§4) 을 얹음.

3. Continuous Batching — 동적 배칭

3.1 문제 — static batching 의 head-of-line blocking

전통적 배칭은 요청 N개를 모아서 한 번에 forward 하고, N개가 모두 끝날 때까지 다음 배치를 시작하지 않아요. 그런데 LLM 디코드는 요청마다 출력 길이가 천차만별입니다 — 어떤 요청은 50 토큰, 어떤 요청은 2,000 토큰이에요.

A, C, D 가 일찍 끝나도 B 가 끝날 때까지 빈 자리에서 GPU 가 놀아요. 이걸 head-of-line blocking 이라고 부릅니다.

3.2 메커니즘 — iteration-level scheduling

ORCA(Yu et al., OSDI 2022) 가 처음 제안한 방식이에요. 배칭 단위를 "요청" 이 아니라 "한 디코딩 스텝" 으로 바꿉니다.

매 토큰을 뽑을 때마다 스케줄러가:

• 끝난 요청은 배치에서 빼냄.
• 대기 중이던 새 요청을 빈 자리에 끼워 넣음.
• 다음 한 스텝을 forward.

ORCA 는 또 selective batching 을 도입했는데, attention 처럼 시퀀스 길이가 다르면 함께 묶기 곤란한 연산만 따로 처리하고, 나머지(linear, layernorm 등)는 통째로 배칭해요.

3.3 효과 — 처리량 ×8~24

수치는 출처마다 비교 대상이 다르므로 분리해서 봐야 해요(Anyscale 블로그 정리, 2023; vLLM 런치 블로그):

표의 "naive batching" 베이스라인은 HuggingFace Pipelines 같은 단일 시퀀스 처리 루프를 가리켜요. 입력 길이 변동이 큰 워크로드에서는 1× 가 81 tok/s 수준의 약한 베이스라인이라, 8× / 23× 같은 배수가 그만큼 절대 속도는 아니에요(Anyscale 정리). 의미 있는 비교는 같은 표의 FasterTransformer(4×) 대비 비율이에요.

요지는 두 가지예요. (1) static 을 잘 최적화하는 것보다 iteration-level 로 패러다임을 바꾸는 게 더 큼. (2) PagedAttention 과 결합했을 때 batch size 한계가 또 풀려서 합산 효과가 한 자릿수 더 늘어남.

3.4 PagedAttention 과의 시너지

continuous batching 은 batch size 를 키울 때 진가를 발휘해요. 그런데 batch size 를 키우려면 KV 캐시 메모리가 더 필요하고, 단편화가 심하면 못 키워요. PagedAttention 이 단편화를 4% 이하로 잡아주니 batch 가 커지고, batch 가 커지니 continuous batching 의 처리량이 또 올라갑니다. 이게 vLLM 의 14~24× 라는 숫자의 정체예요.

3.5 구현

• vLLM, TGI, SGLang, NVIDIA TensorRT-LLM — 모두 기본으로 켬.
• DeepSpeed-Inference, llama.cpp 의 continuous batching 모드 — 따라옴.

4. Prefix Caching — 공통 prefix 재사용

클라이언트 관점(API 캐시 TTL·요청 측 매칭)은 transformer-inference.md 의 Prompt Caching 절에서 다뤄요. 여기서는 서버측 KV 재사용 구현(vLLM APC·SGLang RadixAttention) 을 정리해요.

4.1 동기 — 같은 시작을 N 번 처리할 이유가 있나

서빙 트래픽을 들여다보면 prefix 가 겹치는 경우가 압도적으로 많아요.

• 시스템 프롬프트 — 한 서비스의 모든 요청이 같은 system message 로 시작.
• few-shot 예시 — 분류·추출 작업에서 5~10개 예시를 매 요청 앞에 붙임.
• 멀티턴 채팅 — 한 세션에서 대화가 누적되면, 매 턴마다 그 전 턴 전체가 컨텍스트로 다시 들어옴.
• 에이전트 루프 — 같은 도구 정의·시스템 지침이 매 step 마다 반복.
• 긴 문서 QA — 같은 문서에 여러 질문을 던지는 패턴.

이 경우 prefill 단계(컨텍스트를 한 번에 처리해서 KV 캐시를 만드는 단계) 를 매번 처음부터 다시 하면, 사실상 같은 계산을 N 번 반복하고 있어요.

4.2 vLLM Automatic Prefix Caching (해시 기반)

vLLM 의 APC 는 KV 캐시 블록(§2 의 paged 블록 단위) 마다 해시 키를 만들어요. 새 요청이 들어오면 prefix 의 토큰 시퀀스를 같은 방식으로 해시해서, 일치하는 KV 블록이 이미 풀에 있으면 그대로 가리키게 합니다(공식 docs, enable_prefix_caching=True).

핵심 특징:

• 블록 단위라서 정확히 같은 토큰이 블록 경계까지 일치해야 hit. 한 토큰만 달라도 그 블록부터는 새로 계산.
• prefill 단계만 가속. 디코딩(토큰 생성) 자체는 어차피 새로 해야 함.
• 출력이 길고 prefix 가 짧은 워크로드에서는 효과 작음. 반대로 prefix 가 길고 출력이 짧으면 효과 큼.

4.3 SGLang RadixAttention — 트라이 구조

SGLang(Zheng et al., 2023-12; lmsys 블로그 2024-01) 은 한 단계 더 나가서 radix tree(트라이) 로 prefix 를 관리해요.

각 노드는 KV 캐시 슬라이스 하나를 들고 있고, LRU(least-recently-used)로 차가운 leaf 부터 evict 해요. 캐시 인식 스케줄러는 가장 따뜻한(=많이 hit 되는) 가지를 공유하는 요청을 먼저 골라 잡습니다.

세 가지 부속 메커니즘이 함께 동작해요.

1. Radix tree (trie) 구조 — token sequence 를 키로, KV 캐시 텐서를 값으로 가진 트리. 같은 prefix 는 한 노드만 차지.
2. LRU eviction — GPU 메모리가 차면 가장 오래 안 쓴 leaf 부터 재귀적으로 evict.
3. Cache-aware scheduling — 대기 중 요청 중에서 이미 캐시된 prefix 를 가장 많이 공유하는 요청을 우선 골라 hit rate 를 끌어올림.

vLLM APC 가 "들어온 순서대로 처리하면서 hit 가 나면 좋고" 의 수동적 캐시라면, RadixAttention 은 "스케줄러가 hit 를 적극적으로 만들어내는" 능동적 캐시예요.

4.4 정량 효과

(주의) 두 SGLang 수치(5× vs 6.4×) 는 모델·워크로드 조합에 따라 달라지는 범위예요. 어느 쪽이든 prefix 공유가 많은 워크로드에서만 그 차이가 납니다.

가장 큰 효과가 나는 워크로드들(SGLang 논문·블로그):

• 멀티턴 채팅
• few-shot learning 벤치마크
• 에이전트 루프 (같은 도구 정의·지침 반복)
• tree-of-thought 추론 (여러 분기가 같은 줄기를 공유)
• JSON 추출, RAG 파이프라인

4.5 함의 — 에이전트 워크플로 시대

도구를 호출하는 LLM 에이전트는 매 step 마다 같은 시스템 프롬프트·도구 정의·이전 turn 전체를 다시 보내요. prefix 가 5,000 토큰이고 step 이 20번이면, prefix caching 없이는 5,000 × 20 = 10만 토큰 prefill 을 하지만 캐싱이 있으면 사실상 5,000 토큰 한 번 + 추가분만 prefill. 에이전트가 빈번해질수록 prefix caching 의 비중이 커진다는 뜻이에요.

이 흐름은 DFlash (블록 확산 드래프터) 에서 다룬 투기적 디코딩과는 결이 다릅니다 — 그쪽은 decode 단계의 가속이고, prefix caching 은 prefill 단계의 가속이에요. 두 기법은 서로 상호보완적이라 한 스택 안에서 같이 켤 수 있습니다.

5. 통합 효과 — 같은 GPU, 같은 모델

행마다 비교 대상이 다르므로 수치를 곱해서 합치면 안 됩니다. 다만 vLLM/SGLang 같은 현대 서빙 엔진은 위 7행 중 다섯 행(FA-v2/v3, PagedAttention, Continuous Batching, Prefix Caching) 을 동시에 켜고 돌아가요. naive PyTorch 루프 대비 한 자릿수 후반 ~ 두 자릿수의 처리량 차이가 나는 이유가 이 합산 효과예요.

6. 어떤 환경에서 무엇이 효과 있나

본문 §2 의 결론 — "같은 모델·같은 GPU 에서도 스택 선택에 따라 처리량이 한 자릿수 이상 갈린다" — 의 구체적 근거가 이 표예요. GPU 를 새로 사기 전에, 지금 스택이 위 다섯 행을 다 켜고 있는지부터 확인하는 게 보통 더 큰 절감입니다.

출처 메모

• FlashAttention v1: Dao, Fu, Ermon, Rudra, Ré. "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness." arXiv:2205.14135 (2022-05).
• FlashAttention v2: Dao. "FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning." arXiv:2307.08691 (2023-07).
• FlashAttention v3: Shah, Bikshandi, Zhang, Thakkar, Ramani, Dao. "FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision." arXiv:2407.08608 (2024-07).
• PagedAttention / vLLM: Kwon, Li, Zhuang, Sheng, Zheng, Yu, Gonzalez, Zhang, Stoica. "Efficient Memory Management for Large Language Model Serving with PagedAttention." arXiv:2309.06180, SOSP 2023.
• vLLM 런치 블로그(2023-06): https://vllm.ai/blog/vllm
• ORCA (continuous batching, selective batching): Yu, Jeong, Kim, Chun. "Orca: A Distributed Serving System for Transformer-Based Generative Models." OSDI 2022. (주의) USENIX 페이지 직접 접근은 본 작업 시점에서 403 응답이라, 수치는 Anyscale 정리 블로그(2023) 와 vLLM 런치 블로그를 통해 확인.
• Anyscale "Continuous batching to increase LLM inference throughput" (2023): https://www.anyscale.com/blog/continuous-batching-llm-inference
• SGLang / RadixAttention: Zheng, Yin, Xie 외. "SGLang: Efficient Execution of Structured Language Model Programs." arXiv:2312.07104 (2023-12, rev 2024-06).
• SGLang 런치 블로그(2024-01): https://lmsys.org/blog/2024-01-17-sglang/
• vLLM Automatic Prefix Caching docs: https://docs.vllm.ai/en/latest/features/automatic_prefix_caching.html