Как сообщает PyTorch, команда разработчиков представила Helion — новый язык программирования предметно-ориентированной области (DSL), который компилируется в оптимизированный код для Triton. Этот инструмент решает фундаментальную проблему разработки высокопроизводительных ядер машинного обучения: компромисс между производительностью и портируемостью.
Проблема специализированных ядер
Современное машинное обучение требует всё большей вычислительной мощности, что приводит к созданию множества специализированных ядер. Хотя такие ядра могут демонстрировать впечатляющую производительность, они обычно написаны на низкоуровневых языках, специфичных для конкретного оборудования. Это создает долгосрочную проблему поддержки: ядро, тщательно оптимизированное для одной архитектуры, быстро становится техническим долгом, который сложно и дорого переносить на другое оборудование.
Helion решает эту проблему, компилируя высокоуровневый Python-встроенный DSL в автоматически настраиваемый код Triton. Он устанавливает новый уровень абстракции, который соединяет удобство PyTorch с производительностью низкоуровневых языков.
Сравнение подходов к разработке ядер
- CUDA/Gluon/TLX: Прямое написание ядер на языках вроде CUDA дает максимальный контроль, но требует значительных усилий для достижения высокой производительности
- Triton: Хотя представляет большой шаг вперед, всё еще требует ручного управления индексированием тензоров и определения пространства поиска для автотюнинга
- PyTorch: Предлагает исключительную простоту использования, но ограничивает детальный контроль над стратегиями слияния операций
Модель программирования: PyTorch с тайлами
Цель модели программирования Helion — минимизировать шаблонный код и использовать существующие знания разработчиков о PyTorch. Типичное ядро Helion состоит из двух частей:
import torch, helion, helion.language as hl @helion.kernel() def matmul(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: # --- Код хоста (выполняется на CPU) --- m, k = x.size() k, n = y.size() out = torch.empty([m, n], dtype=x.dtype, device=x.device) # --- Код устройства (компилируется в ядро Triton) --- for tile_m, tile_n in hl.tile([m, n]): acc = hl.zeros([tile_m, tile_n], dtype=torch.float32) for tile_k in hl.tile(k): acc = torch.addmm(acc, x[tile_m, tile_k], y[tile_k, tile_n]) out[tile_m, tile_n] = acc return out
Ключевая языковая конструкция hl.tile разделяет пространство итераций ядра на тайлы. Программист только указывает необходимость тайлинга, а конкретные детали реализации, такие как размеры тайлов, порядок итераций и оптимизации размещения памяти, обрабатываются автотюнером Helion.
Helion выглядит как попытка решить проблему, которая десятилетиями мучает разработчиков системного ПО: как совместить производительность ассемблера с читабельностью Python. Интересно, удастся ли команде PyTorch избежать судьбы многих DSL, которые начинались как элегантные абстракции, а заканчивались как очередной слой legacy-кода, который нужно поддерживать. Автотюнинг — это мощно, но 10 минут на поиск оптимальной конфигурации могут оказаться слишком дорогими для разработки в режиме реального времени.
Автотюнинг: неявные пространства поиска
Ключевое отличие Helion — его автоматизированный автотюнинг на этапе компиляции. В Triton разработчики должны вручную определять пространство поиска для оптимизаций, что требует явного перечисления каждой тестируемой конфигурации.
Helion меняет эту динамику с помощью неявных пространств поиска. Высокоуровневый язык автоматически строит обширное многомерное пространство поиска по вариантам реализации. Например, один вызов hl.tile неявно инструктирует автотюнер исследовать различные размеры блоков, порядки циклов и способы преобразования пространства итераций.
Рабочий процесс автотюнинга
Когда ядро запускается впервые без указанной конфигурации, автотюнер инициирует автоматизированный поиск. Этот процесс, который обычно занимает около 10 минут, оценивает тысячи кандидатов конфигураций ядер Triton, используя стратегии поиска вроде дифференциальной эволюции или поиска по шаблонам.
По завершении автотюнер выводит единственную лучшую найденную конфигурацию, которую можно зафиксировать для последующих запусков. Это позволяет значительно сократить время разработки и обеспечить производительность, близкую к ручной оптимизации.
Helion также включает возможности шаблонизации, позволяющие передавать лямбда-функции в качестве аргументов ядра. Это особенно полезно для реализации универсальных ядер с настраиваемыми эпилогами.
Практические преимущества
Реализация ядра Attention занимает всего 30 строк в Helion по сравнению со 120 строками в Triton и тысячами строк в CUDA. Эта высокоуровневая декларативная модель программирования становится производительной благодаря механизму автотюнинга.
Важное преимущество Helion заключается в том, что знание PyTorch означает знание большей части Helion. Разработчики могут использовать стандартные операторы PyTorch внутри ядер, а Helion автоматически отображает их на соответствующие низкоуровневые реализации Triton.
Новый инструмент обещает значительно сократить время разработки высокопроизводительных ядер машинного обучения, одновременно обеспечивая их портируемость между различными аппаратными архитектурами. Это может ускорить инновации в области ML, снизив барьер для создания оптимизированных вычислительных ядер.