Week 6: Beyond Pytorch: Custom Kernel과 vLLM

모두연 Pytorch + NPU 온라인 모임 #6 | 2025-02-05

소개

이번 강의에서는 PyTorch를 넘어서(Beyond PyTorch) LLM 추론에 필요한 기술들을 다룹니다. 크게 네 가지 주제로 구성됩니다:

LLM inference의 성능 문제 — Prefill과 Decode의 차이, Roofline analysis, Memory wall
Memory overhead를 줄이기 위한 기술들 — Continuous batching, Speculative decoding, Flash Attention, Paged Attention, Quantization
Kernel Programming — CUDA, CUTLASS, OpenAI Triton
Serving 최적화 Framework — vLLM과 기타 상용 프레임워크

강의 진행 현황

원래 계획에서는 PyTorch internal 기초 3회와 심화 8회로 구성되어 있었으나, 현재 계획은 다음과 같이 조정되었습니다:

Week 1: 기술적인 배경
Week 2: eager mode
Week 3: graph mode
Week 4: automatic differentiation
Week 5: distributed programming
Week 6: beyond Pytorch: custom kernel과 vLLM (이번 강의)
Week 7: CPU / GPU / NPU
Week 8: 리벨리온 NPU (마지막 강의)

오늘의 핵심 질문

이번 강의를 관통하는 핵심 질문은 세 가지입니다. 첫째, AI 반도체 관점에서 PyTorch만 지원하면 되는가? 특히 추론에 특화된 AI 반도체라면 어떤 추가 기술이 필요한지를 살펴봅니다. 둘째, 성능 측면에서 LLM 추론의 핵심 이슈는 무엇이며 그 해결책은 무엇인가? 셋째, 개발자들이 이러한 해결책을 구현하기 위해 PyTorch만으로 충분한가?

결론부터 말하면, LLM 추론 관점에서 PyTorch는 큰 그림의 일부입니다. Open Pretrained LLM(Llama, DeepSeek 등)이 Hugging Face 같은 플랫폼을 통해 배포되고, PyTorch 같은 ML Framework 위에서 실행되지만, 실제 성능을 끌어내려면 Kernel Programming과 서빙 최적화라는 PyTorch를 넘어서는 기술이 필수적입니다.

LLM Inference의 성능 문제

Self Attention과 Scaled Dot-Product

LLM의 핵심 연산인 Self Attention은 모든 단어의 Key-Value 값과 특정 단어의 Query를 사용하여 계산됩니다. 예를 들어 “it”이라는 단어를 처리할 때, “it”의 query가 문장 내 모든 단어의 key에 반영되고, 각 단어별로 query-key 내적값과 value를 곱한 뒤, 모든 단어의 영향을 합산하여 self-attention이 반영된 새로운 embedding을 생성합니다. 하단 레이어의 연산은 병렬 처리가 가능하지만, softmax 등 상위 연산은 병렬화가 까다롭습니다.

모든 단어의 kv값과 특정 단어의 query를 사용

Non-Causal vs. Causal Attention

Attention에는 크게 두 가지 유형이 있습니다. Non-Causal Attention은 앞뒤 모든 단어를 고려하는 양방향 방식이고, Causal Attention은 이전 단어만 고려하는 단방향 방식입니다. GPT 등장 이후 대부분의 LLM은 Causal Attention(Auto-regressive) 방식을 사용합니다.

Non-Causal Attention

앞/뒤 모든 단어를 고려 (양방향)

Causal Attention

이전 단어만 고려 (단방향)

Prefill과 Decode

Decode-only LLM은 다음 단어를 반복적으로 예측하는 구조입니다. 모델과 데이터로 다음 단어를 예측하도록 학습하여 Trained Weights를 만들고, 이 weights를 기반으로 Prefill과 Decode 두 단계를 거쳐 추론을 수행합니다.

Prefill은 프롬프트를 이해하는 과정으로, 첫 번째 토큰 생성까지의 단계입니다. Decode는 문장이 끝날 때까지 토큰을 반복 생성하며, 실제 답변을 만들어내는 과정입니다. Prefill과 Decode는 동일한 weight를 공유하지만, 구조가 다른 모델이라고 생각할 수 있습니다.

Prefill 단계의 성능 지표는 TTFT(Time to First Token)이고, Decode 단계의 성능 지표는 TPS(Tokens Per Second)입니다. Self-Attention 계산 시 고려해야 하는 이전 문맥이 점점 길어지므로, 이론적으로 스텝이 진행될수록 TPS가 감소합니다.

Prefill vs. Decode: 병렬성의 차이

Prefill과 Decode의 가장 중요한 차이는 병렬성에 있습니다.

Prefill 단계에서는 프롬프트의 모든 토큰에 대한 KV 값을 생성하면서 첫 번째 단어를 생성합니다. 모든 프롬프트 토큰을 동시에 처리할 수 있는데, 다음 단어가 이미 정해져 있으므로 이전 토큰 결과를 기다릴 필요가 없기 때문입니다. 배치 사이즈가 큰 계산과 유사합니다.

Decode 단계에서는 이전 토큰이 완전히 처리되어야 다음 토큰을 처리할 수 있으며, 매 스텝마다 하나의 토큰만 생성합니다. 따라서 순차 처리가 불가피합니다.

병렬성 비교를 다이어그램으로 보면, Prefill은 수직 방향(레이어 간)과 우상향(KV 전달) 의존성만 존재하여 횡적으로 병렬 처리가 가능한 반면, Decode는 토큰 간 순차적 의존성이 있어 시퀀셜 처리가 필요합니다.

Prefill (Parallel)              Decode (Sequential)
┌───┐ ┌───┐ ┌───┐ ┌───┐       ┌───┐ ┌───┐
│ I │ │like│ │my │ │cat│       │ a │ │lot│
└─┬─┘ └─┬─┘ └─┬─┘ └─┬─┘       └─┬─┘ └─┬─┘
  │    │    │    │         │    │
  ▼    ▼    ▼    ▼         ▼ ──→ ▼
┌───┐ ┌───┐ ┌───┐ ┌───┐       ┌───┐ ┌───┐
│L1 │ │L1 │ │L1 │ │L1 │       │L1 │ │L1 │
└─┬─┘ └─┬─┘ └─┬─┘ └─┬─┘       └─┬─┘ └─┬─┘
  │    │    │    │         │    │
  ▼    ▼    ▼    ▼         ▼ ──→ ▼
┌───┐ ┌───┐ ┌───┐ ┌───┐       ┌───┐ ┌───┐
│L2 │ │L2 │ │L2 │ │L2 │       │L2 │ │L2 │
└───┘ └───┘ └───┘ └───┘       └───┘ └───┘

Roofline Analysis

Roofline analysis를 통해 Prefill과 Decode의 성능 특성을 분석하면 명확한 차이가 드러납니다. Decode는 Arithmetic Intensity가 낮아 연산기가 남아있는데 데이터 전송 속도가 병목이 되는 Memory Bound 상태입니다. 반면 Prefill은 Arithmetic Intensity가 높아 연산기를 최대한 활용하는 Compute Bound 상태입니다.

Memory Wall

메모리 성능 문제는 오래된 이슈입니다. 1998년 “Memory Wall” 논문 이후, 프로세서 성능은 연산 속도보다 데이터 전송 속도에 의해 결정된다는 사실이 널리 알려졌습니다.

메모리 인터페이스 기술 자체는 꾸준히 발전하고 있지만, 하드웨어의 연산 밀도(Computing Density) 증가 속도가 메모리 인터페이스 발전 속도보다 훨씬 빠르기 때문에 상대적인 메모리 성능 문제는 오히려 점점 더 심해지고 있습니다. 따라서 메모리 대역폭(Memory Bandwidth)을 효율적으로 사용하는 것이 매우 중요합니다.

Memory Overhead를 줄이기 위한 기술들

Continuous Batching

여러 Query를 병렬로 처리하는 가장 기본적인 방법은 배칭입니다. Static/Dynamic Batching은 여러 요청을 묶어 Lock-Step으로 처리하지만, 가장 느린 요청까지 대기해야 하는 문제가 있습니다.

이를 해결하기 위해 등장한 것이 Continuous Batching(Orca, 서울대 정병권 교수)입니다. 각 레이어 단계에서 동적으로 요청을 끼워 넣어 병렬 처리하는 방식으로, 시퀀스 길이가 달라도 효과적으로 처리할 수 있습니다. 현재 LLM 서빙의 필수 기술로 자리잡았습니다.

Speculative Decoding

Speculative Decoding은 여러 토큰을 병렬로(Speculative하게) 생성하는 기법입니다. 동작 방식은 다음과 같습니다:

작은 모델(Draft Model)로 여러 토큰의 초안을 생성합니다
큰 모델(Main Model)로 초안을 병렬 검증합니다 (Prefill과 유사한 방식)
오류 발생 지점 이전 토큰만 Accept하고, 이후는 폐기한 뒤 재시도합니다

이 방식을 통해 한 요청 내에서 검증 과정을 활용하여 여러 토큰을 동시에 처리할 수 있습니다.

Attention의 비용 문제

Attention 연산은 매우 비용이 높습니다. Attention의 복잡도는 $O(n^2)$ 으로, 여기서 $n$ 은 sequence length입니다. 마지막 토큰을 처리할 때 이전 모든 토큰으로부터 Attention을 계산해야 하고, 이를 $N$ 번 반복하면 $O(N^2)$ 복잡도가 됩니다. 시퀀스가 길어질수록 메모리와 시간이 기하급수적으로 증가합니다.

Flash Attention

Flash Attention은 $O(N^2)$ 복잡도를 근본적으로 해결할 수는 없지만, 훨씬 효율적으로 처리할 수 있게 합니다. 핵심 기법은 Tiling + Fusion입니다. 기존 방식에서는 반복적으로 DRAM에서 읽기/쓰기를 수행하여 메모리 접근 오버헤드가 컸지만, Flash Attention은 데이터를 한 번 올리고 계산을 완료한 후 한 번에 내보내는 방식으로 메모리 접근을 최소화합니다.

Paged Attention

기존 방식에서는 Max sequence length에 맞춰 KV Cache를 할당하면 메모리 낭비가 발생합니다. Paged Attention은 Logical Cache Block과 Physical Cache Block을 분리(Decouple)하여, 필요한 부분만 물리 메모리에 할당(On-Demand Allocation)하는 방식으로 이 문제를 해결합니다. 이는 OS의 Virtual Memory / Paging 기법과 매우 유사한 접근입니다.

모델 경량화

Quantization(양자화)과 Sparsity(희소성) 등의 모델 압축 기술도 메모리 overhead를 줄이는 데 활용됩니다.

Kernel Programming

Flash Attention과 GPU-Aware 알고리즘

Flash Attention은 단순히 알고리즘 수준의 최적화가 아니라, GPU 하드웨어 특성을 깊이 이해한 GPU-Aware Attention 알고리즘입니다. PyTorch의 Op 추상화에는 한계가 있습니다. 각 커널이 DRAM에서 읽고 쓰는 I/O 모델을 따르기 때문에 타일링과 퓨전이 불가능합니다. 이를 해결하려면 저수준 API로 커널을 직접 작성하고, 이를 Custom Op으로 래핑하여 PyTorch에서 사용해야 합니다.

GPU 메모리 계층에서 DRAM은 용량이 크지만 느리고, Shared Memory는 용량이 작지만 빠릅니다. Flash Attention은 데이터를 쪼개서 Shared Memory에 올려놓고 고대역폭으로 연산하는 전략을 취합니다.

Tiling과 Fusion의 적용

Attention 연산 $\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right) V$ 에서 각 연산의 Tiling 적용 가능성은 다릅니다:

곱셈 (QK, …V): 자연스럽게 tiling 적용 가능
Masking: 마찬가지로 tiling 가능
Softmax: 전체 분포를 봐야 하므로 tiling이 까다로움 — 핵심 기술적 챌린지

모든 연산이 같은 방식으로 tiling된다면 자연스럽게 fusion 적용이 가능합니다. 하지만 이러한 tiling과 fusion은 PyTorch 수준에서는 표현이 불가능하며, Kernel Programming이 꼭 필요합니다.

Softmax의 도전: Online Softmax

일반적인 softmax는 Sum을 구하는 pass가 필요한 2-pass 방식이고, 수치 안정성을 위해 Max를 구하는 pass를 추가하면 3-pass(Safe softmax)가 됩니다:

$y_i = \frac{e^{x_i - \max(x)}}{\sum_j e^{x_j - \max(x)}}$

Safe softmax는 수치 안정성은 보장하지만, 메모리를 3번 읽어야 하는 오버헤드가 있습니다.

일반적인 softmax (2-pass)

Safe softmax (3-pass)

FlashAttention은 Online Algorithm을 사용하여 이 문제를 해결합니다. 최댓값을 미리 계산하지 않고, 현재까지 확인된 최댓값만 유지하면서 매 tile마다 max값을 update하고 보정합니다. Q는 모든 tile에 동일하게 적용되고, K는 i번째 tile만 읽어들이며, 최종 결과를 DRAM에 저장합니다. 이렇게 FlashAttention은 softmax의 numerical stability에 대한 영향을 최소화한 single-pass, tiled/fused 알고리듬을 구현합니다.

FlashAttention의 진화: v1 → v2 → v3

FlashAttention은 알고리즘을 재구성하여 읽기/쓰기 횟수를 줄이고 효율성을 높인 기술로, Stanford PhD 학생 Tri Dao가 개발했습니다. 버전별 진화 과정은 다음과 같습니다:

Version	주요 특징	GPU Utilization
Standard	PyTorch ops only	Low
v1	Tiling + Fusion (CUDA)	~25%
v2	CUTLASS 활용, Better warp partitioning	50-70%
v3	Hopper 최적화 (async copy, low-precision)	~75%

GPU 아키텍처의 진화와 CUTLASS

GPU 아키텍처는 CUDA가 가정했던 깨끗한 SIMT 아키텍처에서, 다양한 최적화 기능이 추가되면서 현저히 복잡해졌습니다.

CUDA가 가정한, 깨끗한 SIMT 아키텍쳐

최적화 기능이 추가되면서 현저히 복잡해진 최근 GPU

CUTLASS는 이렇게 복잡해진 GPU의 계층적 병렬 수행을 개발자가 직접 control할 수 있도록 C++ template을 제공하는 라이브러리입니다.

하지만 최적화된 CUDA Kernel을 작성하는 것은 여전히 만만치 않습니다. GPU Peak 성능 대비 달성 가능한 수준을 비교하면 그 난이도가 명확히 드러납니다:

접근 방식	GPU Peak 성능 대비
CUDA 프로그래밍 가이드 기반	~10%
CUTLASS 활용	80-90%
NVIDIA 비공개 소스 (cuBLAS)	~100%

출처: https://dyninst.github.io/scalable_tools_workshop/petascale2023/assets/slides/kerenz.pdf

Triton: CUDA 대안 오픈소스

Triton은 Harvard PhD Philippe Tillet이 개발하고, OpenAI 지원 하에 오픈소스로 발전한 kernel programming language입니다. CUDA의 복잡성은 감추면서 타일링과 퓨전을 표현할 수 있으며, PyTorch 2.0의 Inductor 백엔드에 통합되었습니다.

“OpenAI’s Triton is very disruptive angle to Nvidia’s closed-source software moat for machine learning.”

Triton의 핵심은 Block Programming Model입니다. CUDA와 Block Programming 간에는 Semantic Gap이 크지만, Triton은 개발자가 의도를 손쉽게 표현할 수 있는 block programming model을 제공하고, GPU 최적화를 위한 세부 사항은 컴파일러 기술로 해결합니다.

CUDA와 Block Programming 간의 Semantic Gap이 큼
Triton은 개발자가 의도를 손쉽게 표현할 수 있는 **“block programming model”**을 제공
GPU 최적화를 위한 detail은 컴파일러 기술로 해결

Triton은 다양한 백엔드를 지원합니다: Nvidia GPU, AMD GPU, Intel GPU, AWS Trainium, Qualcomm Hexagon NPU, Azure MAIA, ARM CPU, x86 CPU 등 폭넓은 하드웨어에서 동작합니다.

Serving 최적화 Framework

vLLM: LLM Serving의 업계 표준

vLLM은 PagedAttention을 개발한 PhD 학생들이 만든 프로젝트에서 시작했습니다. Continuous Batching, Speculative Decoding 등 앞서 다룬 고수준 최적화 기법들이 서빙 프레임워크 내부에 통합되어 있으며, 오픈 소스 생태계에서 성공적으로 자리 잡아 업계 표준으로 인정받고 있습니다.

vLLM Architecture

vLLM의 성능

앞서 다룬 중요한 최적화 기술들은 대부분 vLLM에 반영되어 있습니다.

vLLM의 PyTorch 중심 전략

vLLM은 PyTorch 기반으로, 하드웨어 호환성을 PyTorch를 통해 확보하는 전략을 취하고 있습니다. 이는 AI 반도체 회사 입장에서 매우 중요한 의미를 가집니다. PyTorch에 하드웨어를 잘 붙이면 vLLM에서도 자연스럽게 사용할 수 있기 때문입니다.

기타 상용 프레임워크

오픈 소스로는 vLLM이 널리 사용되고 있으며, 상용 솔루션으로는 Firework AI, Together AI, Friendly AI 등이 독점적인 최적화 기술을 가진 서빙 프레임워크를 제공하며 활발한 펀딩과 사업화를 진행하고 있습니다.

정리

이번 강의에서 다룬 내용을 정리하면 다음과 같습니다:

LLM inference의 성능 문제 — Prefill은 Compute Bound, Decode는 Memory Bound라는 근본적으로 다른 특성을 가지며, Roofline analysis와 Memory wall 분석을 통해 이를 확인할 수 있습니다.
Memory overhead를 줄이기 위한 기술들 — Continuous batching은 다양한 시퀀스 길이에서도 높은 활용도를 제공하고, Speculative decoding은 검증 과정으로 문제를 전환하여 병렬 처리를 가능하게 합니다. Flash Attention은 타일링과 퓨전으로 메모리 접근을 최소화하고, Paged Attention은 논리/물리 캐시 블록을 분리하여 메모리 낭비를 줄입니다.
Kernel Programming — PyTorch 추상화의 한계를 넘어서기 위해 CUDA, CUTLASS, 그리고 특히 Triton이 중요한 역할을 합니다.
Serving 최적화 Framework — vLLM이 업계 표준으로 자리잡았으며, 기타 상용 프레임워크들도 활발히 발전하고 있습니다.

결론적으로 PyTorch는 LLM 추론 생태계의 핵심이지만, 커널 프로그래밍과 고수준 서빙 프레임워크가 함께 필요합니다.