5,000줄의 파이썬으로 구현한 원칙적인 ML 컴파일러 스택

실용적 조언

커스텀 커널 최적화 시 뱅크 충돌을 피하기 위해 공유 메모리 차원에 패딩을 추가하는 기법을 고려하라.
메모리 대역폭이 병목인 연산에서는 연산자 융합을 통해 HBM 읽기/쓰기 횟수를 줄이는 것이 가장 효과적이다.

섹션별 상세

작성자는 PyTorch FX 그래프에서 시작하여 최종 CUDA 코드에 이르기까지 총 6단계의 중간 표현(IR)을 정의했다. 각 단계는 Torch IR, Tensor IR, Loop IR, Tile IR, Kernel IR을 거치며 하드웨어 추상화 수준을 점진적으로 낮춘다. 특히 Tensor IR 단계에서는 모든 연산을 Elementwise, Reduction, IndexMap으로 분해하여 프론트엔드 확장성을 확보했다.

python

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

컴파일러 파이프라인을 통해 변환될 입력 PyTorch 코드 예시

연산자 융합(Fusion)을 통해 메모리 효율성을 극대화하는 메커니즘을 구현했다. Loop IR 단계에서 인접한 커널들을 하나의 루프 네스트로 통합하여 중간 버퍼 생성을 방지한다. 이를 통해 x + bias 연산과 matmul, relu가 하나의 커널 내에서 처리되며 HBM(High Bandwidth Memory) 접근 횟수를 최소화한다.

cpp

extern "C" __global__ __launch_bounds__(256) void k_relu_reduce(const float* bias, const float* x, const float* w, float* relu) {
  // ... (중략) ...
  #pragma unroll
  for (int a2 = 0; a2 < 2; a2++) {
    // Shared memory loading with async and padding
    // 2x2 register tile accumulation
    acc0 += v0 * in3; acc1 += v0 * in4; acc2 += v1 * in3; acc3 += v1 * in4;
  }
  relu[idx] = fmaxf(0.0f, acc0);
}

최종 단계에서 생성된 최적화된 CUDA 커널 코드

GPU 성능 최적화를 위해 Tile IR 단계에서 하드웨어 특화 최적화를 적용했다. 2x2 레지스터 타일링을 통해 스레드당 4개의 출력을 동시에 누적하고, 더블 버퍼링과 비동기 메모리 복사(cp.async)를 사용하여 연산과 통신을 오버랩했다. 또한 공유 메모리 뱅크 충돌을 방지하기 위해 패딩(Padding)을 추가하는 기법을 수치적으로 구현했다.

성능 측정 결과, Qwen 블록 크기의 어텐션 스코어 계산에서 torch.compile과 대등한 성능을 기록했다. 5,000줄의 파이썬 코드만으로도 TinyLlama 및 Qwen2.5-7B와 같은 실제 모델의 엔드투엔드 컴파일이 가능함을 입증했다. 이는 교육적 목적뿐만 아니라 해킹 가능한(hackable) 컴파일러로서의 실용성을 보여준다.

용어 해설

중간 표현(Intermediate Representation): — 소스 코드를 기계어로 번역하는 과정에서 거치는 추상적인 데이터 구조입니다. 컴파일러는 여러 단계의 IR을 거치며 최적화를 수행하고, 하드웨어에 종속되지 않는 분석을 가능하게 합니다.
연산자 융합(Operator Fusion): — 여러 개의 독립적인 연산을 하나의 커널로 결합하여 메모리 대역폭 사용을 줄이는 최적화 기법입니다. 중간 결과물을 메모리에 쓰고 다시 읽는 과정을 생략하여 실행 속도를 대폭 향상시킵니다.
공유 메모리(Shared Memory): — GPU의 스트리밍 멀티프로세서(SM) 내부에 위치한 고속 메모리 영역입니다. 전역 메모리보다 훨씬 빠르며, 동일한 블록 내의 스레드들이 데이터를 공유하고 재사용하는 데 사용됩니다.
타일링(Tiling): — 대규모 행렬 연산을 작은 블록(타일) 단위로 나누어 처리하는 기법입니다. 데이터의 지역성을 활용하여 고속 캐시나 공유 메모리에 데이터를 유지함으로써 메모리 병목 현상을 해결합니다.

언급된 도구

deplodock추천

5,000줄의 Python으로 작성된 참조 ML 컴파일러

triton중립

GPU 프로그래밍을 위한 언어 및 컴파일러 (비교 대상으로 언급)

언급된 리소스

문서A Principled ML Compiler Stack in 5,000 Lines of Python (Full Article)

GitHubdeplodock GitHub Repository