TileLang: высокоуровневый DSL для GPU-ядер — от сложения векторов до FlashAttention

TileLang — это высокоуровневый предметно-ориентированный язык на Python для проектирования и компиляции GPU-ядер, построенный на базе TVM. Он позволяет разработчикам работать с shared-memory tiles, регистровыми фрагментами, конвейерными циклами и тензорными инструкциями, оставляя компилятору управление потоками, синхронизацией и генерацией низкоуровневого CUDA-кода.

В туториале последовательно реализуются сложение векторов, блочное матричное умножение с тензорными ядрами, исследование расписаний, fused GEMM epilogues, построчный softmax и FlashAttention. Для каждого шага проводится сравнение с эталонными реализациями PyTorch и cuBLAS, а также инспекция сгенерированного CUDA-кода.

Отдельное внимание уделено автоподбору параметров ядра (autotuning), который позволяет находить оптимальные конфигурации для конкретной архитектуры GPU. В результате разработчик получает производительность, близкую к ручной оптимизации, при минимальных усилиях.