Команда Hazy Research представила прорывной мегакернел для ускорения работы Llama*-70B на графических процессорах H100. Новое решение превосходит популярный SGLang на 22% по общей пропускной способности при обработке 65 536 промптов из бенчмарка ShareGPT. Секрет кроется в агрессивном перекрытии вычислений, работы с памятью и коммуникационных операций — как будто вы научились одновременно жарить яичницу, варить кофе и читать новости, вместо того чтобы делать это по очереди.
В отличие от предыдущего мегакернела для Llama*-1B, который был заточен под низкую задержку на одном GPU, новый вариант оптимизирован для высокопроизводительной обработки больших батчей на нескольких устройствах. Если раньше проблема была в простом ожидании загрузки весов из памяти, то теперь исследователи столкнулись с гетерогенной нагрузой: одни операции ограничены вычислениями (матричные умножения), другие — пропускной способностью памяти (attention decode), а третьи — межпроцессорными связями NVLink.
Ключевая идея заключается в использовании всех ресурсов GPU одновременно, а не по частям. Традиционные системы разбивают прямой проход модели на десятки или сотни отдельных ядер, каждое из которых имеет периоды настройки и завершения, когда полезная работа не выполняется. Мегакернел объединяет весь процесс в единое ядро с интерпретатором инструкций, работающим на каждом потоковом мультипроцессоре (SM). Это позволяет начинать загрузку данных для следующей операции, пока текущая еще выполняется.
Исследователи добились перекрытия ресурсов на трех уровнях иерархии GPU. Внутри отдельных SM они используют конвейерную обработку инструкций, что дает прирост 2-6% производительности. Между SM применяется глобальная очередь работ, обеспечивающая прирост до 14.2% при размере батча 8192. А между GPU коммуникационные операции выполняются в фоновых потоках, позволяя другим потокам продолжать вычисления.
Итоговые результаты впечатляют: мегакернел показал 14 425 токенов/с на входе против 11 783 у SGLang и 9 043 токена/с на выходе против 7 387. При этом авторы честно предупреждают, что их код — это исследовательский прототип, чувствительный к версиям компилятора, настройкам GPU и «иногда даже к неправильному взгляду». Но идея настолько многообещающая, что они планируют создать более универсальный набор инструкций для мегакернелов, возможно, даже для обучения моделей.
*Продукт принадлежит компании Meta, которая признана экстремистской и запрещена в РФ.
Источник новости и обложки: hazyresearch.stanford.edu