gpumode · 강의 아카이브
《GPU Mode》 L079 2025 High priority transcript · failed

Mirage (MPK) — Compiling LLMs into Mega Kernels

LLM inference 의 latency 곡선에서 보이지 않는 두 도둑 — kernel launch overheadHBM 왕복. CMU 의 Mirage Persistent Kernel 은 LLM forward 전체를 한 개의 대형 fused 커널로 컴파일해서 두 비용을 동시에 절감한다. 1.2-6.7× latency 개선이라는 숫자가 의미하는 것 — search space, super-optimizer, mega-kernel 의 인지적 모델까지 학습 노트.

Mirage MPK mega kernel super-optimizer CMU OSDI 2025 LLM inference launch overhead
M
Speaker
Mengdi Wu
CMU · Mirage 공동저자
X
Speaker
Xinhao Cheng
CMU · Mirage 공동저자
강의 번호
L079
학습 우선순위
High
자막
failed
출처
OSDI 2025
§ 01강의가 풀려는 문제· why this lecture exists

"하나의 inference token = 수백 개 kernel launch" 의 자리를 다시 본다

LLM 의 inference 를 표준 PyTorch / vLLM 으로 돌리면 한 token 의 forward 가 보통 200-500 개의 GPU kernel 로 쪼개진다. attention, GEMM, layer norm, residual, RMSNorm, rotary embedding... 각자가 별도 launch. small batch (특히 batch=1) 에서는 kernel launch overhead 가 wallclock 의 30-50% 가 되는 자리.

강의의 frame.

  1. 각 kernel 사이의 HBM 왕복 — 한 layer 의 output 이 HBM 에 써지고 다음 kernel 이 다시 읽는다. 같은 데이터를 두 번 옮긴다. memory bandwidth 의 낭비.
  2. launch overhead — 각 launch 가 ~5μs 의 host-device 통신. 200 launch 면 1ms. small-batch decode latency 에 직접 박힌다.
  3. 이걸 한 커널로 묶을 수 있다 — Mirage Persistent Kernel (MPK). "한 번 launch 된 mega-kernel 이 layer 들을 inside 에서 sequencing".
강의의 인지적 frame

vLLM / SGLang 같은 inference engine 의 다음 자리는 커널 갯수 자체를 줄이는 것. fusion 한두 개가 아니라 전체 forward 를 단일 커널로 컴파일. 그것이 Mirage 의 출발점.

"우리가 본 inference 의 진짜 도둑은 compute 가 아니라 launch 와 HBM 왕복이다 — 둘 다 한 커널로 묶어서 죽인다."Mirage 저자 · 확인 필요
§ 02mega-kernel 의 동기· launch overhead 절감

many-kernel vs mega-kernel — timeline 의 차이

왜 mega-kernel 이 의미가 있는지를 timeline 한 장으로 본다.

FIG · LLM 한 token forward 의 timeline 비교vLLM many-kernel vs Mirage mega-kernel
vLLM
L
k1
L
k2
L
k3
L
k4
L
k5
L
attn
L
k6
L
k7
L
k8
Mirage MPK
L
single mega-kernel · all layers in-kernel scheduled
절감
latency saving (launch + HBM)
L = launch overhead bar, k = kernel. vLLM 의 갭마다 launch 가 박혀 있고, kernel 사이에 HBM round-trip 의 implicit cost. Mirage 는 한 launch + on-chip scheduling.

mega-kernel 이 왜 단순히 fusion 의 한 사례 가 아닌가. 일반 fusion (예: torch.compile) 은 인접한 op 만 묶는다 (matmul + bias + relu). Mirage 는 layer 의 경계를 넘어서 attention 과 다음 layer 의 GEMM 까지 한 커널 안에 들어간다. 이게 가능한 이유 — § 03 의 super-optimizer 가 dependency 를 분석해 inside-kernel scheduling 으로 표현하기 때문.

"persistent" 의 의미

"Persistent Kernel" 은 한 번 launch 된 후 모든 SM 위에 머무르며 작업을 처리하는 패턴. work-queue 에서 task 를 fetch 해서 실행. forward 의 모든 layer 가 task 로 분해되어 한 persistent kernel 안에서 도는 것이 MPK.

§ 03super-optimizer· why "super"

"여러 단계의 IR 을 동시에 search" 가 핵심 차별점

Mirage 가 부른 표어 — multi-level superoptimizer for tensor programs. 일반 컴파일러가 한 IR level 에서만 최적화한다면, Mirage 는 graph-level / tile-level / SM-level 을 동시에 search.

L0 graph IR model 의 computational graph — matmul, attention, norm 같은 op 들의 DAG
L1 block / tile IR graph 의 op 를 tile 단위로 분해. 어떤 op 끼리 fuse 할지, tile 크기는 얼마인지
L2 thread / warp IR tile 안의 thread 분배. shared memory layout. WGMMA 의 매핑
L3 PTX → SASS 최종 코드 생성. nvcc 가 담당

"super" 의 진짜 의미 — L0/L1/L2 의 결정을 동시에 search. graph 단계에서 fusion 후보를 먼저 정하고 → tile 단계에서 최적 size 를 정하고 → thread 단계에서 layout 을 정하는 sequential pipeline 이 아니다. 세 단계의 결정이 맞물려 있어 분리하면 local optimum 에 갇힌다.

예시 — attention 의 KV cache load 와 GEMM 의 weight load 를 동시에 보면.

  • graph 단계 — KV load 와 weight load 가 같은 tile group 에 들어가면 fusion 가능.
  • tile 단계 — 두 load 가 같은 SM 위에서 돌면 SMEM 공유로 bandwidth 절감.
  • thread 단계 — thread 분배가 두 load 의 latency 를 가리도록 hand-tuned.

이 셋이 분리되면 — graph 단계에서 fusion 결정 시 SMEM 비용을 모르고, tile 단계에서 thread 를 모르고. super-optimizer 는 세 결정을 joint search.

왜 "super-optimization" 이라 부르는가

Massalin 의 1987 superoptimizer 의 정신 — 모든 가능한 program 을 enumerate 하고 검증한다. Mirage 는 LLM-scale 에서 같은 정신을 살리되 search space 를 학습한 heuristic 으로 줄이고, 검증을 자동화. 결과 = "사람이 짤 수 없는 형태의 mega-kernel".

"우리는 graph + tile + thread 를 함께 search 한다 — 분리하면 좋은 mega-kernel 이 나오지 않는다."Mirage 논문 · 확인 필요
§ 04search space· what is searched

scheduler 가 결정해야 하는 4 차원

Mirage 의 search space 를 명시적으로 풀어본다. 4 개 축이 있다.

TAB · search space 의 4 차원모든 결정의 joint optimization
총 search space 는 천만 단위 후보. cost model 로 빠르게 prune, 살아남은 후보만 실측. 한 model 당 컴파일 시간은 분-시간 단위.
FIG · search space sketchgraph × tile × thread
tile size large tile fusion depth best vLLM (many-kernel) Mirage 의 search candidate ★ = 측정 후 선택 ○ = pruned candidates · = full search space
vLLM 의 default 는 좌하단 (작은 tile, 얕은 fusion). Mirage 의 best 는 우상단 (큰 tile, 깊은 fusion). 그 영역에 mega-kernel 이 산다.
§ 05검증· verification

"super-optimizer 가 짠 코드가 진짜 맞는가" 의 답

super-optimizer 의 위험 — 컴파일러가 짜낸 코드가 사람이 검토하기 어려운 형태. silent wrong result 가 production 에 풀려 나갈 risk. Mirage 의 답.

이 검증 layer 가 산업적 채택 가능성의 핵심. 빠른 mega-kernel 이 silent wrong 이면 production 에 풀어둘 수 없다.

test-time vs build-time

Mirage 의 "컴파일" 은 작은 sample input 으로 build-time 검증을 강제. 그러나 production 의 모든 input shape 을 cover 하지 못한다 — 그래서 inference-time 에 shape 가 새로 등장하면 fallback 또는 재컴파일. dynamic shape 의 대응이 § 08 의 한계.

§ 06LLM 추론 결과· latency numbers

1.2× ~ 6.7× — 어떤 자리에서 큰 속도가 나는가

Mirage 가 보고하는 latency 개선 범위는 1.2× 부터 6.7× 까지 넓다. 이 변동의 원인을 보면 어떤 자리에 mega-kernel 이 효과적인지가 보인다.

TAB · model × scenario 별 speedup (개념적)baseline = vLLM many-kernel
큰 speedup 은 batch=1 decode 에 집중 — 즉 latency-critical inference, 즉 chat 응답 / coding assistance 같은 자리. throughput-critical 자리 (큰 batch prefill) 에서는 효과 약함.

모델별로도 변동이 있다. 강의에서 보고한 패턴 (확인 필요).

"Mirage 의 speedup 곡선은 batch size 와 prefill/decode 의 함수 — chat application 의 latency-critical 자리에서 가장 빛난다."학습 노트
§ 07hardware 활용· SM occupation

persistent kernel 이 SM 위에 오래 머무르는 의미

mega-kernel 의 hardware 측면 효과. 단순히 launch 절감 외에 SM 활용도 자체가 변한다.

이 모델의 trade-off — register / SMEM pressure 가 매우 커진다. mega-kernel 안에 모든 데이터가 들어와야 하므로 occupancy 감소 risk. Mirage 의 search 가 이를 cost 로 반영해 균형점을 찾는다.

CUDA Graphs 와의 관계

CUDA Graphs (cudaGraph) 도 launch overhead 를 절감한다. 차이는 — Graphs 는 여러 kernel 을 미리 묶어 한 번에 launch, kernel 자체는 분리. Mirage 는 코드 자체를 한 커널로 합성. fusion 효과가 추가됨.

§ 08한계· caveats

모든 자리에서 mega-kernel 이 답은 아니다

§ 09다음 단계· future work

강의에서 명시적으로 던진 다음 자리

§ 10기억할 메모· key takeaways

다시 열었을 때 손에 잡혀야 할 것

launch overhead 의 진짜 비용
batch=1 decode 에서 wallclock 의 30-50%. 200-500 launch 가 한 token 마다.
mega-kernel
forward 전체를 한 커널로 합성. persistent + on-chip scheduling. fusion 의 한 단계 위.
super-optimizer
graph + tile + thread 를 joint search. multi-level. cost model + 실측 검증.
search space 4 차원
graph fusion / tile shape / SM mapping / memory placement. 천만 단위 candidate, pruning 후 실측.
검증 layer
numerical (atol/rtol) + differential testing + timing profile. silent wrong 차단.
speedup 곡선
batch=1 decode 에서 ~6.7×, large batch 에서 1.2-2×. latency-critical 자리에 best.
CUDA Graph 와의 차이
Graph = 여러 kernel 묶어 launch. Mirage = 코드 자체를 한 kernel 로 합성. fusion 효과 추가.
한계
compile time 큼, dynamic shape 약점, training 미지원 (강의 시점). multi-GPU / Blackwell 은 future.
YouTube강의 영상 (확인 필요)
PaperMirage — A Multi-Level Superoptimizer for Tensor Programs · OSDI 2025
RelatedCUDA Graphs · TVM · Halide super-optimization

손에 새기기 — 실습 시퀀스

  1. Mirage 의 hello-world — repo 의 simple model (Llama 1B 또는 GPT-2) 를 컴파일. compile time 측정. 출력 mega-kernel 의 PTX 살펴보기.
  2. baseline 비교 — 같은 model 을 vLLM / Mirage / torch.compile 세 가지로 batch=1, batch=64 decode latency 측정.
  3. search space 직접 보기 — Mirage 의 search log 를 dump. 어떤 후보가 prune 됐고 어떤 게 살아남았는지.
  4. numerical 검증 — Mirage output 을 PyTorch eager 와 atol 1e-3 비교. 어디서 차이가 나는지 layer 별로 추적.
  5. shape sensitivity — sequence length 를 바꿔가며 재컴파일 비용 측정. cache 효과.
  6. nsys 로 timeline — vLLM 의 launch overhead 와 Mirage 의 mega-kernel 을 같은 timeline 에 본다.
§ 11다른 강의로의 연결· connections

이 강의가 시리즈 안에서 어디로 이어지는가

§ 12열린 질문· open questions

다음에 다시 들었을 때 직접 검증해야 할 것들

검증 메모

이 노트의 모든 latency 수치는 Mirage repo / OSDI 2025 paper / 강의 metadata 의 재구성. 강의 시점 이후 활발한 개발이 이뤄졌을 가능성. 정확한 API 와 성능 수치는 repo README + paper 직접 참조.

← Lecture 078Iris — Multi-GPU Programming in Triton Lecture 080 →How FlashAttention 4 Works