현대적 LLM 컴파일러 스택 개요: 대화형 및 해킹 가능한 컴파일러 작성하기

요약

본 글은 라이브러리 없이 순수 Python과 raw CUDA만을 사용하여 바닥부터 구축한 소규모 ML 컴파일러 스택의 설계 과정을 다룹니다. Transformer 모델을 6단계의 중간 표현(IR)을 통해 CUDA 커널로 변환하는 과정을 설명하며, 실험 결과 PyTorch 프로덕션 스택에 근접하는 성능과 일부 커널에서의 압도적인 속도 향상을 입증했습니다.

핵심 포인트

순수 Python과 raw CUDA를 사용하여 6단계의 중간 표현(IR)을 갖춘 ML 컴파일러를 직접 구현함
Torch IR부터 CUDA 소스 코드 생성까지 이어지는 단계별 컴파일 파이프라인(Tensor, Loop, Tile, Kernel IR)을 제시함
RTX 5090 환경에서 PyTorch 스택 대비 기하평균 0.96배의 성능을 기록하며 실용성을 증명함
특정 커널의 경우 PyTorch의 수동 최적화 커널보다 최대 5.6배 빠른 성능을 달성함

안녕하세요 r/LocalLLaMA 여러분,

실제 서비스용 (Production) ML 컴파일러 스택은 매우 가혹합니다. TVM은 50만 줄 이상의 C++ 코드로 이루어져 있습니다. PyTorch는 Dynamo, Inductor, 그리고 Triton을 서로 쌓아 올린 구조입니다. XLA, MLIR, Halide, Mojo 등도 존재합니다.

이는 단연코 현대 컴퓨팅 인프라에서 가장 중요한 부분 중 하나임에도 불구하고, 핵심 개념에 대해 이용 가능한 정보가 거의 없습니다. 이 간극을 메우기 위해, 저는 라이브러리를 전혀 사용하지 않고 순수 Python과 로우 레벨 (raw) CUDA만을 사용하여 바닥부터 작은 ML 컴파일러를 구축했습니다. 이 컴파일러는 작은 트랜스포머 (Transformer) 모델 (TinyLlama, Qwen2.5-7B)을 입력받아 6개의 중간 표현 (IR, Intermediate Representation)을 거쳐 일련의 CUDA 커널 (Kernel)로 낮춥니다 (lowering). RTX 5090에서 오토튜닝 (Autotuned)된 스택은 PyTorch 프로덕션 스택 대비 기하평균 0.96배의 성능을 기록했으며, 84개의 커널 형태 중 32개가 PyTorch의 수동 최적화 (hand-optimized) 커널을 능가했습니다 (최대 5.6배 속도 향상).

한 달간의 작업 끝에, 3부작 시리즈가 마침내 완성되었습니다:

Part 1에서는 RMSNorm 레이어를 파이프라인의 상단 절반을 통해 엔드 투 엔드 (end-to-end)로 살펴봅니다:

Torch IR — 캡처된 FX 그래프 (rmsnorm, linear, softmax, ...)
Tensor IR — 모든 연산 (op)이 Elementwise / Reduction / IndexMap으로 분해됨
Loop IR — 다른 커널과 융합 (fused)된 루프 중첩 (loop nest) 형태로 작성된 커널
Tile IR — GPU에 스케줄링된 커널 (threads, blocks, shared memory)
Kernel IR — 하드웨어 프리미티브 (primitives)로 구체화된 스케줄
CUDA — nvcc를 위한 준비가 된 소스 코드 생성

예를 들어, PyTorch 표현식이 IR 레벨을 통과하는 과정은 다음과 같습니다:

torch.relu(torch.matmul(x + bias, w))   # x: (16, 64), bias: (64,), w: (64, 16)

Torch IR:

bias_bc =  bias[j]                         -> (16, 64) float32
add     =  add(x, bias_bc)                 -> (16, 64) float32
matmul  =  matmul(add, w, has_bias=False)  -> (16, 16) float32
...

Tensor IR:

bias_bc  =  bias[j]                 -> (16, 64) float32
w_bc     =  w[j, k]                 -> (16, 64, 16) float32
add      =  add(x, bias_bc)         -> (16, 64) float32
...

Loop IR

=== merged_relu -> relu ===
for a0 in 0..16:  # free (M)
    for a1 in 0..16:  # free (N)
...

Part 2에서는 하반부, 즉 루프 중첩 (loop nest)이 어떻게 GPU 스케줄 (GPU schedule)로 변환되는지를 설명합니다. 계산을 블록 (blocks)으로 분할하고, 스레드 (threads)에 매핑하며, 입력을 공유 메모리 (smem)로 스테이징 (stage)하는 등의 작업을 수행하는 16개의 기계적인 Tile-IR 패스 (passes)가 포함됩니다. 각 패스는 CLI 상에서 하나의 diff로 나타납니다. 이는 CUDA 엔지니어가 수행하는 최적화 시퀀스를 모방합니다.

예를 들어, 입력을 smem으로 스테이징하는 패스는 다음과 같습니다:

deplodock compile \
  -c "torch.nn.RMSNorm(2048)(torch.randn(1,32,2048))" \
  --ir tile -vv \
...

>>> t:007_stage_inputs
@@ matched at rms_norm (in-place) @@
@@ -2,6 +2,7 @@
...

Part 3는 오토튜닝 (autotuning)으로 시리즈를 마무리합니다. Part 2의 모든 파라미터 (블록 크기 (block size), 레지스터 타일 (register tile), K-청크 (K-chunk), 스테이징 여부, 더블 버퍼링 (double-buffer) 여부 등)는 휴리스틱 (heuristics)을 사용하여 수동으로 선택되었습니다. 이러한 휴리스틱은 제가 맞춘 형태 (shapes)에는 작동했지만, 다른 곳에서는 실패했습니다. Part 3는 휴리스틱을 탐색 루프 (search loop)로 대체합니다: 규칙 파라미터들의 크로스 프로덕트 (cross-product)에 대한 SP-MCTS를 사용합니다.

전체 파이프라인은 하나의 CLI로 구성됩니다:

# 모든 IR 단계 검사
deplodock compile -c "nn.RMSNorm(2048)(torch.randn(1,32,2048))" --ir tensor|loop|tile|kernel|cuda

...

각 파트는 충분히 독립적이어서 특정 레이어에만 관심이 있다면 건너뛸 수 있습니다:

AI 자동 생성 콘텐츠

원문 바로가기

현대적 LLM 컴파일러 스택 개요: 대화형 및 해킹 가능한 컴파일러 작성하기

요약

핵심 포인트

댓글