TileLang — 고성능 GPU·CPU·가속기 커널 개발을 위한 도메인 특화 언어

AI 요약

TileLang은 고성능 GPU·CPU·NPU 커널을 작성하는 도메인 특화 언어(DSL)입니다. GEMM, 양자화 GEMM, FlashAttention, LinearAttention처럼 모델 학습과 추론에서 쓰이는 연산을 대상으로 합니다. Python과 비슷한 문법을 쓰되 TVM 기반 컴파일러가 하드웨어에 맞는 코드를 생성합니다. 개발자는 커널의 타일 구성과 데이터 이동, 연산을 표현하면서 저수준 최적화도 적용할 수 있습니다.

커널 작성과 컴파일

README의 예제는 FP16 행렬 곱에 FP32 누산을 적용한 뒤 ReLU를 융합하는 커널입니다. @tilelang.jit 함수 안에서 입력 텐서의 크기와 자료형을 지정하고, T.Kernel로 처리할 타일과 스레드를 정합니다. T.alloc_shared는 공유 메모리 타일을 만들고 T.alloc_fragment는 누산용 타일을 준비합니다. T.Pipelined는 전역 메모리에서 공유 메모리로 데이터를 옮기는 단계와 연산을 파이프라인으로 구성합니다. 마지막으로 T.gemm이 타일 행렬 곱을 표현하고 T.Parallel이 원소별 ReLU를 처리합니다.

첫 호출 때 @tilelang.jit은 입력 크기와 컴파일 시점 인자를 바탕으로 커널을 특수화합니다. 예제는 생성한 결과를 PyTorch의 행렬 곱과 ReLU 결과에 대조해 검증합니다. TileLang은 현재 환경의 장치를 감지해 기본 컴파일 대상을 선택하며, 다른 백엔드나 아키텍처를 지정하려면 Target 객체를 사용할 수 있습니다.

백엔드와 지원 범위

CUDA가 주 백엔드이며 SM70부터 SM120까지의 코드 경로를 제공합니다. TMA, WGMMA, TMEM 같은 기능은 해당 기능을 지원하는 GPU가 필요합니다. AMD ROCm/HIP과 Apple Metal도 지원 백엔드로 안내하며, ROCm은 Linux와 CDNA·RDNA GPU를 대상으로 합니다. Metal은 Apple 실리콘에서 동작하고, 지원되는 M5 시스템에서는 Metal 4 cooperative tensor 기능을 사용할 수 있습니다.

LLVM CPU, NVIDIA CuTe DSL, WebGPU는 실험 단계입니다. CPU 백엔드는 LLVM 15 이상으로 소스 빌드해야 하며, CuTe DSL은 별도 패키지가 필요합니다. Huawei Ascend와 MetaX, Moore Threads, HYGON, Sunrise-AI 장치용 구현도 생태계 어댑터로 나뉘어 있습니다. 이 어댑터는 TileLang 본체에 포함되지 않고 별도 저장소에서 개발되므로 릴리스와 호환성 일정도 독립적입니다. 공식 휠은 Linux x86-64·AArch64, Windows x86-64, macOS arm64용으로 제공됩니다.

컴파일러 개발과 디버깅

저장소에는 커널 작성뿐 아니라 컴파일 과정을 살피는 기능도 포함됩니다. 컴파일 오류에 Python 소스 위치를 표시하고, 중간 표현(IR)이 각 컴파일 단계를 거치며 어떻게 바뀌는지 확인하는 추적·비교 도구를 제공합니다. 레이아웃 시각화와 컴파일 단계별 시간 측정 기능도 안내합니다. 자동 튜닝은 커널 후보를 벤치마크하며, README에는 파이프라인·그룹 컴파일과 멀티 GPU 벤치마킹 지원도 나와 있습니다.

설치와 활용

PyPI의 안정판은 pip install tilelang으로 설치합니다. 최신 기능을 먼저 시험하려면 별도 nightly 휠을 쓸 수 있지만 안정판보다 불안정할 수 있습니다. AMD GPU에서는 ROCm 빌드 PyTorch와 호스트 ROCm 런타임이 필요합니다. 프로젝트는 GEMM·양자화와 어텐션 커널, 시퀀스 모델 연산, 특정 GPU 아키텍처용 예제와 컴파일러 도구를 각각 안내합니다. 성능 비교 스크립트와 설정은 별도 tilelang-benchmark 저장소에서 확인하도록 연결합니다.

원문: GitHub / 번역·요약: Trawling