Triton от OpenAI: язык для GPU-ядер без боли CUDA

OpenAI выложила в открытый доступ Triton — язык и компилятор для написания GPU-ядер на Python-подобном синтаксисе. Заявленная цель амбициозна: дать исследователю без опыта в CUDA писать код, который по скорости конкурирует с рукописными библиотеками NVIDIA. По внутренним замерам OpenAI, ядро умножения матриц FP16 на Triton достигает производительности cuBLAS в пределах нескольких процентов — при объёме кода примерно в 25 строк.

Почему это вообще нужно

Писать быстрые ядра под GPU — работа для узкой касты. CUDA требует ручного управления разделяемой памятью, коалесингом обращений, синхронизацией варпов и планированием инструкций. Один неверный шаг с шаблоном доступа к памяти — и ядро в разы медленнее теоретического предела.

Отсюда парадокс: исследователи придумывают новую архитектуру нейросети, но упираются в отсутствие эффективного ядра. Приходится либо мириться с медленной наивной реализацией на PyTorch, либо ждать, пока кто-то напишет CUDA-ядро вручную. Triton пытается закрыть этот разрыв.

Наша цель — сделать так, чтобы разработчик без опыта в CUDA мог писать код, сопоставимый по скорости с работой экспертов.

Как это устроено

Ключевая идея — уровень абстракции между CUDA и высокоуровневыми фреймворками. Программист на Triton работает с блоками данных, а не с отдельными потоками. Компилятор сам решает задачи, которые в CUDA лежат на человеке:

• распределение работы по потокам внутри блока;
• размещение промежуточных значений в разделяемой памяти;
• коалесинг обращений к глобальной памяти;
• планирование инструкций и использование тензорных ядер.

Ядро описывается декоратором @triton.jit и выглядит как обычная Python-функция. Вот минимальный пример поэлементного сложения векторов:

import triton
import triton.language as tl

@triton.jit
def add_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr):
    pid = tl.program_id(axis=0)
    offsets = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offsets < n
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)
    tl.store(out_ptr + offsets, x + y, mask=mask)

Никакого threadIdx, никакого __syncthreads(), никакой ручной работы с shared memory. Программист описывает, что делает один блок над куском данных; распределение по потокам компилятор берёт на себя.

Что решает компилятор

Внутри Triton опирается на LLVM и собственный промежуточный уровень представления. На нём выполняются оптимизации, характерные именно для GPU: анализ шаблонов доступа, автоматическая векторизация, размещение тайлов в регистрах и разделяемой памяти. Для матричного умножения это даёт то, ради чего всё затевалось, — использование тензорных ядер без единой строчки о них в исходнике.

Triton против альтернатив

Главное отличие от TVM и Halide — Triton не заставляет отдельно описывать «что вычислить» и «как распланировать». Пишете как обычный код с блочными операциями, а планирование остаётся компилятору.

Где предел

Публичный релиз работает на GPU NVIDIA, начиная с архитектуры Volta и новее. Никакой поддержки AMD, Apple Silicon или встраиваемых ускорителей в первой версии нет. Компилятор ориентирован на тензорные ядра и вычисления с плотными матрицами; на нерегулярных задачах, вроде разреженных графов или сложной динамики управления, преимущества над CUDA становятся менее очевидными.

Ещё одна оговорка — сравнения производительности OpenAI приводит на своих сценариях. Independent-бенчмарки на нестандартных задачах могут выглядеть иначе. Точных цифр по разным моделям GPU в исходном анонсе не приведено.

Кому это пригодится прямо сейчас

1. Исследователям, которым нужен нестандартный слой — attention с маской особой формы, кастомная нормализация, экзотическая функция активации.
2. Инженерам, оптимизирующим инференс: слияние операций в одно ядро экономит проходы по памяти и заметно снижает латентность.
3. Авторам библиотек, которым важна поддержка широкого набора форм тензоров без ручного тюнинга под каждый случай.

Для большинства прикладных задач в PyTorch Triton не заменяет ничего — типовые операции уже упираются в cuBLAS и cuDNN. Смысл появляется там, где готовых ядер нет или где нужно соединить несколько операций в одну.

Проект открыт под лицензией MIT и живёт на GitHub. Установка сводится к pip install triton, зависимость — CUDA-совместимая видеокарта и драйверы NVIDIA. Экосистема уже впитала язык: Triton используется внутри torch.compile для автоматической генерации ядер, а также в ряде библиотек ускорения LLM.