从头编写 LLM 编译器：PyTorch 到 CUDA 的 5,000 行 Python

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

저는 현대 ML 컴파일러 스택의 간단한 개요를 작성하고 싶었습니다. 즉, model.generate()에서 GPU 가 커널을 실행하기까지 어떤 일이 일어나는지입니다. 그러나 이 스택은 읽기 매우 어렵습니다. TVM 는 50 만 줄 이상의 C++ 코드를 가지고 있습니다. PyTorch 는 Dynamo, Inductor, Triton 을 서로 위에 쌓아놓았습니다. 그리고 XLA, MLIR, Halide, Mojo 가 있습니다.

대신, 저는 다른 접근 방식을 선택하고从头부터 하나를 구축하기로 결정했습니다. 순수 Python 과 raw CUDA 만 사용합니다. 작은 모델 (Qwen2.5-7B, TinyLlama) 을 취하여 CUDA 커널의 시퀀스로 컴파일합니다. 목표는 오늘 Triton 을 이기는 것이 아니라, PhD 가 필요하지 않고 수정하기 쉬운 해커블 (hackable) 컴파일러를 만들거나, 적어도 더 쉽게 따라가게 하는 것입니다.

최종 성능은 PyTorch Eager 와 torch.compile 를 기준으로 생산 스택의 약 50-90% 입니다.

저는 원칙적인 방식으로 구축했습니다. 레이어화된 파이프라인과 명확히 분리된 관심사를 가지고 있습니다:

1. Torch IR — FX 그래프 (rmsnorm, linear, softmax, ...) 를 포착합니다
2. Tensor IR — 모든 연산이 Elementwise / Reduction / IndexMap 으로 분해됩니다
3. Loop IR — 다른 커널과 융합된 루프 네스트로 작성된 커널입니다
4. Tile IR — GPU 에 스케줄링되는 커널 (스레드, 블록, 공유 메모리)
5. Kernel IR — 하드웨어 원시물로 스케줄링 물질화됩니다
6. CUDA — nvcc 를 위한 소스 방출 준비 완료

Tensor IR 은 ONNX 와 Jax 와 같은 미래 프론트엔드를 지원하기 위해 도입되었습니다. Loop fusion 은 긴 점근적 연산과 reduction 체인의 융합을 처리합니다. Lowering 단계를 통해 타일링된 matmul, smem 스테이징, 더블 버퍼링과 같은 최적화가 소개됩니다.

각 단계는 독립적으로 검사하고 디버깅할 수 있습니다 (repository link). GPU 는 필요 없습니다:

deplodock compile -c "nn.RMSNorm(2048)(torch.randn(1,32,2048))" --ir tensor|loop|tile|kernel|cuda

벤치마킹:

deplodock run --bench --profile -c "torch.nn.Softmax(dim=-1)(torch.randn(1,28,2048,2048))"

엔드 투 엔드 컴파일레이션:

deplodock compile Qwen/Qwen2.5-7B

생성된 RMSNorm 의 CUDA 커널은 다음과 같습니다:

extern "C" global
launch_bounds(256) void k_rms_norm_reduce(const float* x, const float* p_weight, float* rms_norm) {
// ... (코드 번역)