AI BriefingPyTorch→CUDA LLM 컴파일러
Writing an LLM compiler from scratch: PyTorch to CUDA in 5,000 lines of Python
핵심 내용
5,000줄 Python과 raw CUDA로 PyTorch 그래프를 CUDA까지 내리는 컴파일러를 구현했다.
자세히 보기
5,000줄 Python, 2주, 외부 라이브러리 없음이라는 제약으로 LLM 컴파일러를 처음부터 만들고, PyTorch 그래프를 Torch IR → Tensor IR → Loop IR → Tile IR → Kernel IR → CUDA로 내리는 파이프라인을 정리했다.
Tensor IR의 핵심 원시는 Elementwise, Reduction, IndexMap 세 가지다. add, mul, exp, rsqrt 같은 계산, sum, max, prod 같은 축 축소, reshape, transpose, slice, unsqueeze, concat 같은 레이아웃 변환을 모두 이 셋으로 환원한다.
예시로 RMSNorm은 PyTorch의 단일 op에서 15개 노드의 Tensor IR로 풀리고, Linear는 broadcast된 mul과 sum으로 분해된다. 중간의 M×K×N 형태는 실제로 만들지 않고, 이후 fusion 단계에서 하나의 커널로 합친다.
transpose와 slice처럼 연속된 IndexMap은 좌표 매핑을 합성해 하나로 접는다. 그다음 Loop IR로 lift한 뒤 인접한 루프를 fusion해 글로벌 메모리 왕복을 줄이고, 최종 CUDA 커널을 출력한다.
Part 1은 캡처, 분해, IndexMap 합성, 루프화, fusion까지를 다루며, 이어질 Part 2에서 Tile IR와 matmul 최적화, 코드 생성 세부를 설명한다.
이 한국어 요약은 AI가 자동으로 만들었습니다. 원문의 주장과 맥락은 원문에서 확인해 주세요. 저작권은 원저작자에게 있습니다.
