안녕하세요 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

...

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

• Part 1. IR Hierarchy — From PyTorch to Emitted CUDA
• Part 2. Tile IR — Scheduling Loops onto a GPU
• Part 3. Autotuning — A Search Loop Over Tile-IR Rewrites
• Repo: https://github.com/cloudrift-ai/deplodock