초록

현대의 AI 워크로드는 학습 및 추론 모두에서 최적화된 계산 커널에 크게 의존한다. 이러한 AI 커널은 DRAM과 SRAM 간에 타일(tile)을 이동시키고, 해당 타일들에 대해 일련의 계산을 수행하는 등 명확한 데이터 흐름 패턴을 따르고 있다. 그러나 이러한 패턴이 명확함에도 불구하고, 고성능 커널을 작성하는 것은 여전히 복잡한 과제이다. 최고 성능을 달성하기 위해서는 현대의 가속기들을 최대한 활용하기 위해 철저한 하드웨어 중심 최적화가 필요하다. 도메인 특화 컴파일러들은 고성능 커널 작성의 부담을 줄이려는 시도를 하지만, 사용성과 표현력 측면에서 여전히 한계를 겪는다. 본 논문에서는 보다 효율적인 AI 커널 프로그래밍을 위한 일반화된 타일 기반 프로그래밍 모델인 TileLang을 제안한다. TileLang은 스케줄링 공간(스레드 바인딩, 레이아웃, 텐서화, 파이프라인 등)을 데이터 흐름에서 분리하여, 사용자 정의 어노테이션과 기초(primitive)로 캡슐화한다. 이 접근 방식을 통해 사용자는 커널의 데이터 흐름 자체에 집중할 수 있고, 나머지 대부분의 최적화는 컴파일러가 담당하게 된다. 다양한 일반적으로 사용되는 장치에서 광범위한 실험을 수행한 결과, TileLang은 주요 커널에서 최첨단 성능을 달성함을 입증하였다. 이는 통합된 블록-스레드 패러다임과 투명한 스케줄링 기능이 현대 AI 시스템 개발에서 요구하는 강력함과 유연성을 동시에 제공함을 보여준다.