Perplexity переводит ROSE на CuTeDSL для ускорения инференса на GPU Hopper и Blackwell

Perplexity встроила язык CuTeDSL в собственный inference‑движок ROSE, чтобы генерировать JIT‑компилируемый PTX‑код и точно контролировать исполнение пропускных по нагрузке операций на GPU Hopper и Blackwell. Решение ориентировано не на переписывание математических формул, а на то, как именно GPU выполняет простые по сути, но требовательные по пропускной способности задачи — от нормализаций до маршрутизации MoE и операций эмбеддингов.

ROSE (Runtime — Optimized Serving Engine) выступает интерфейсом для моделей: он принимает запросы, формирует планы исполнения, бэтчит входы и возвращает декодированные токены, эмбеддинги или оценки. Движок управляет инициализацией устройств, межпроцессным и межузловым взаимодействием и загрузкой весов; для LLM‑задач он отвечает за батчинг, чанкинг, сэмплинг, аллокацию KV‑памяти как для полного, так и для линейного внимания, а также за префикс‑матчинг.

ROSE изначально создавался для удобного развёртывания кастомных Llama‑моделей и имел интерфейс, совместимый с Triton Inference Server; со временем он вырос в полноценный движок для ранжирования, классификации, скоринга и других трансформерных задач. Экосистема ядер ROSE включает реализации на CuTeDSL, Triton, CUDA, CUTLASS и cuBLAS; для базовых операций матричных умножений и ядров внимания преимущественно используются готовые реализации CUTLASS и cuBLAS. Техническое преимущество CuTeDSL — сочетание высокоуровневых абстракций и контроля низкоуровневых примитивов. Язык, основанный на CuTe layout algebra и MLIR, компилирует код JIT в оптимизированный PTX, позволяя агрессивно специализировать ядра по скрытым размерностям и конкретным устройствам. Такая специализация критична для вспомогательных ядер — нормализаций, MoE‑роутинга, функций активации и эмбеддингов — которые требуют тонкой настройки, чтобы достичь пикового быстродействия.

Практические эффекты для инженеров и команд: сочетание готовых NVIDIA‑ядр для тяжёлых линейных алгебраических операций и кастомных CuTeDSL‑ядр для остальной части пайплайна упрощает экспериментирование с архитектурами и ускоряет выкат оптимальных вариантов в продакшен. ROSE также реализует межустройственные примитивы для MoE‑маршрутизации, диспетчинга и комбинирования, что помогает масштабировать специализированные реализации под архитектуры Hopper и Blackwell. В результате переход на CuTeDSL команда получила более гибкий стек для создания специализированных GPU‑ядер и упрощённый процесс оптимизации производительности: от генерации JIT‑PTX до развертывания в распределённых средах с учётом требований к памяти и общению между узлами.