TileLang — это высокоуровневый Python-язык предметной области (DSL) для проектирования и компиляции ориентированных на производительность GPU-ядер через TVM. По мере того как большие языковые модели и другие transformer-архитектуры повышают вычислительные требования, возможность создавать пользовательские GPU-ядра, полностью использующие tensor cores — специализированные ускорители матричного умножения NVIDIA, — становится все более критичной для эффективности обучения и инференса. Однако написание таких ядер традиционно требует глубоких знаний CUDA C++, программирования на уровне warp и аппаратно-зависимых иерархий памяти. TileLang закрывает этот пробел, позволяя разработчикам описывать вычисления на уровне тайлов, тогда как компилятор берет на себя низкоуровневые детали. Это руководство дает всесторонний обзор возможностей TileLang, начиная с проверки окружения и создания переиспользуемых утилит для бенчмаркинга и численной верификации. Затем в нем последовательно реализуются сложение векторов, тайловое матричное умножение на tensor cores, исследование расписаний, объединенные эпилоги GEMM, построчный softmax и FlashAttention.

На протяжении руководства разработчики напрямую работают с тайлами TileLang в shared memory, регистровыми фрагментами, конвейеризованными циклами, примитивами параллельной итерации, редукциями и операторами tensor-core GEMM. Компилятор автоматически управляет сопоставлением потоков, схемами размещения памяти, синхронизацией, векторизацией и генерацией низкоуровневых CUDA-инструкций. Ядра сравниваются с базовыми реализациями PyTorch и cuBLAS; также выполняются инспекция сгенерированного CUDA-кода, оценка пропускной способности памяти и вычислений, а также автотюнинг для выявления архитектурно-зависимых конфигураций ядер. TileLang доступен на GitHub.

Руководство начинается с настройки CUDA-окружения Google Colab, установки TileLang с резервным вариантом nightly-сборки и импорта необходимых модулей PyTorch и TileLang. Определяются переиспользуемые утилиты для бенчмаркинга, валидации и отчетности, которые измеряют задержку ядра и сравнивают численные результаты с использованием относительной ошибки. Затем реализуется ядро сложения векторов на TileLang, запускается на GPU, сравнивается по пропускной способности с PyTorch и анализируется по CUDA-коду, сгенерированному компилятором.

Далее реализуется тайловое ядро матричного умножения на tensor cores, которое перемещает входные тайлы через global memory, shared memory и регистровые фрагменты. Размеры тайлов, стадии конвейера, число потоков и L2 swizzling задаются вручную, тогда как TileLang генерирует инструкции tensor cores, синхронизацию и логику передачи данных. Несколько конфигураций расписаний проходят бенчмаркинг, проверку численной точности, после чего определяется наиболее производительная архитектурно-зависимая конфигурация ядра. Полный код руководства доступен здесь.

Затем ядро матричного умножения расширяется за счет объединения добавления bias и активации GELU непосредственно в аккумуляторе, находящемся в регистрах. Такой подход снижает промежуточный трафик global memory, завершая эпилог до записи итогового выходного тензора. Подобное объединение ядер является хорошо известным методом сокращения узких мест, связанных с пропускной способностью памяти, которые часто доминируют в задержке крупномасштабных нейросетевых нагрузок. Объединенная реализация сравнивается с eager-выполнением PyTorch. Кроме того, реализуется ядро построчного softmax с использованием редукций максимума и суммы на уровне фрагментов, при этом процесс нормализации в значительной степени остается в регистрах.

Реализуется объединенное forward-ядро FlashAttention, которое обрабатывает тайлы query, key и value без материализации полной матрицы attention-score в global memory. Онлайн-обновления softmax применяются с использованием текущих максимумов, сумм нормализации, коэффициентов масштабирования и тайловых матричных умножений на tensor cores. FlashAttention, изначально предложенный исследователями из Stanford, стал широко применяемым методом сокращения объема памяти self-attention с квадратичного пространства до линейного и теперь интегрирован в основные фреймворки, включая API scaled dot-product attention в PyTorch. Как causal, так и non-causal attention валидируются относительно scaled dot-product attention в PyTorch, с сопоставлением задержки и вычислительной пропускной способности.

Определяется пространство поиска автотюнинга по размерам матричных тайлов, размерностям K-блоков, глубинам конвейера и числу потоков, при этом конфигурации, превышающие бюджет shared memory, отфильтровываются. Декоратор автотюнинга TileLang используется для компиляции, бенчмаркинга, валидации и кэширования нескольких расписаний ядра для одной и той же нагрузки матричного умножения. Выбранное ядро выполняется, проверяется относительно PyTorch и оценивается по достигнутой задержке и пропускной способности tensor cores. Этот автоматизированный поиск решает практическую проблему разработки GPU-ядер: оптимальные размеры тайлов и глубины конвейера различаются между GPU-архитектурами, такими как Ampere и Hopper, что делает ручную настройку хрупкой при ориентации на несколько поколений оборудования.

Рабочий процесс отладки и интроспекции TileLang представлен через вывод с устройства, инспекцию сгенерированного CUDA-кода и встроенный профилировщик ядер. Рассматриваются маркеры, сгенерированные компилятором, включая операции tensor cores, асинхронные копирования, барьеры синхронизации и инструкции загрузки матриц. Все разделы руководства организованы в отказоустойчивый runner, который записывает статус выполнения, сообщает информацию о времени и выводит компактный справочник по программированию на TileLang.

Руководство демонстрирует, как TileLang переводит Python-программы на уровне тайлов в оптимизированные GPU-ядра без необходимости вручную управлять индексами потоков, схемами данных на уровне warp, инструкциями tensor cores или асинхронными барьерами памяти. Реализованные и проверенные ядра охватывают элементные операции, ограниченные пропускной способностью памяти, вычислительно интенсивные нагрузки GEMM, объединенные эпилоги нейронных сетей, редукции, размещенные в регистрах, и attention с online-softmax. Исследование также показывает, как размеры блоков, потребление shared memory, глубина конвейера, число потоков, форма тайла и L2 swizzling влияют на производительность в разных GPU-архитектурах. Инспекция сгенерированного исходного кода, отладка на стороне устройства, профилирование и автоматизированный поиск расписаний вместе формируют полный рабочий процесс для разработки, проверки, бенчмаркинга и улучшения пользовательских ядер TileLang.