Статья представляет новое исследование под названием CODA: Rewriting Transformer Blocks as GEMM-Epilogue Programs, основная цель которого — оптимизировать эффективность обучения моделей Transformer, особенно путем решения «память-интенсивных» операций, которые кажутся разрозненными, но в совокупности занимают значительное время.

Автор статьи, источник: Machine Learning China

22 мая Три Дао перепостил в социальных сетях твит Хань Гуо. Он также написал: «После некоторых математических преобразований выяснилось, что всё в Transformer представляет собой серию GEMM + epilogue (матричное умножение плюс эпилог). При наличии оптимизированных примитивов LLM (и новички) смогут написать ядра для всех операций Transformer со скоростью света!»

Tri Dao — один из основных авторов серии FlashAttention, а этот твит ссылается на их статью, опубликованную в тот же день: CODA.

• Название статьи: CODA: Переопределение блоков Transformer как программ GEMM-Epilogue
• Адрес статьи: https://arxiv.org/abs/2605.19269
• Адрес кода: https://github.com/HanGuo97/coda-kernels

Это имя звучит как «финал» и произносится как «CUDA». Исследователи из MIT, Принстона, Together AI и Meta попытались систематически устранить мелкие, редко замечаемые, но постоянно потребляющие время вычисления в обучении Transformer с помощью новой программной абстракции.

Налог на лень при обучении больших моделей

Чтобы понять, какую проблему решает CODA, сначала нужно разобраться, куда уходит время при обучении крупных моделей.

На обучении модели LLaMA-3 стиля с 1 миллиардом параметров на одном GPU NVIDIA H100 большинство людей интуитивно считают, что всё время тратится на матричные умножения и вычисления внимания, ведь именно они — «настоящие вычисления». Эта интуиция в целом верна: матричные умножения (GEMM) и внимание действительно занимают основную часть вычислительных ресурсов.

Но если вы откроете профайлер и внимательно посмотрите, вы обнаружите, что еще есть ряд «малых операторов», тихо потребляющих время: нормализация (RMSNorm), функции активации (SwiGLU, RoPE), остаточное сложение, межуровневое сокращение… Их отдельные вычислительные затраты невелики, но они часто перемещают крупные промежуточные тензоры из видеопамяти и обратно.

Это то, что называется «узким местом пропускной способности памяти»: представьте себе повара с безупречным мастерством, которому для приготовления каждого блюда нужно таскать ингредиенты из удалённого склада и возвращать их обратно, вместо того чтобы держать их на рабочей поверхности рядом. Насколько бы быстро ни двигались руки повара, время ожидания транспортировки — это реальная трата времени.

Еще хуже то, что с появлением у NVIDIA низкопrecisionных форматов, таких как FP8 и FP4, матричные вычисления становятся все быстрее, относительная стоимость этих «перемещений» растет: матричное умножение ускоряется, но затраты на передачу тензоров туда и обратно не сокращаются пропорционально.

В статье представлена наглядная группа данных: при обучении модели с 1 млрд параметров на H100 с использованием TorchTitan нематричные операции занимают значительную долю общего времени выполнения, и эта доля еще больше возрастает при использовании точности FP8.

Существующие программные фреймворки практически бессильны перед этим. PyTorch представляет вычисления Transformer в виде последовательности операторов с четкими границами между ними. Эти границы отлично подходят для автоматического дифференцирования (autograd), но именно они препятствуют оптимизации объединения между операторами: каждая граница оператора часто означает ненужную запись в память.

CODA: В «Финале» спрятаны сокровища

Идея CODA основана на простом наблюдении.

На GPU высокопроизводительное ядро матричного умножения (GEMM) структурно разделено на две части: основной цикл (mainloop), отвечающий за основные вычисления умножения и сложения блоков матриц, и завершающая часть (epilogue), выполняющая дополнительные операции перед записью результата в видеопамять, такие как добавление смещения, преобразование типов и простое масштабирование.

Значение финального этапа заключается в том, что вывод матричного умножения пока «живет» в регистрах на чипе и еще не записан в глобальную видеопамять. Это короткое золотое окно: если на этом этапе выполнить дополнительные вычисления, можно полностью избежать лишнего цикла записи и чтения из видеопамяти.

Основной инсайт CODA заключается в том, что многие память-интенсивные операции в Transformer можно алгебраически переparameterизовать и выполнить в рамках этого «заключительного» окна.

Это требует некоторого математического мастерства. Возьмем наиболее распространенный шаблон GEMM-RMSNorm-GEMM: результат матричного умножения проходит через остаточное сложение, RMS-нормализацию, а затем еще одно матричное умножение. В традиционном подходе три отдельных оператора выполняются последовательно, а промежуточные результаты дважды записываются в виде памяти видеокарты.

Команда CODA обнаружила, что масштабирующий множитель строк r в нормализации RMS, будучи скалярной величиной, общей для каждой строки, коммутирует с последующим матричным умножением: применение r можно отложить с «периода перед вторым GEMM» на «конец второго GEMM». После откладывания на конце первого GEMM требуется вычислить только локальную «частичную среднеквадратичную норму» (partial RMS), которая объединяется с помощью крайне легковесного вспомогательного ядра сокращения, а полное вычисление RMSNorm исчезает.

Аналогичная переparameterизация применима и к операциям, таким как SwiGLU, RoPE (вращающееся позиционное кодирование), кросс-энтропийная потеря, и даже к обратному распространению. В статье приведена теорема, доказывающая: если прямой проход является «блочно-локальным», то обратное распространение автоматически наследует ту же структуру. Подробности см. в оригинальной статье.

Пять «кирпичиков» и набор «языка Лего»

CODA — это не конкретный фьюжн-ядро, а набор программных абстракций.

Он фиксирует оптимизированный экспертами основной цикл GEMM, а затем в завершающей части предоставляет пять типов комбинируемых базовых примитивов:

• Построчное преобразование (остаточное сложение, функция активации, RoPE)
• Загрузка и сохранение векторов (трансляция весов RMSNorm)
• Блочная загрузка и сохранение матриц (сохранение промежуточных активаций для обратного распространения)
• Блочное сокращение (локальное среднеквадратическое отклонение, блочное log-sum-exp)
• Состояние изменений (статистика max и sum-exp, необходимая для онлайн-нормализации)

С помощью этих пяти типов блоков можно охватить почти все операции, кроме внимания, в прямом и обратном проходах стандартного Transformer.

Еще более интересно, насколько эта абстракция терпима к вопросу «кто пишет код». В экспериментах статья оценивала два подхода: один — с написанием кода человеком-программистом, другой — с генерацией кода с помощью Claude Code: при заданных примитивах CODA, нескольких примерах и журналах реализации ИИ выполняет большую часть кода ядра, а человек обеспечивает легкий контроль.

Производительность обоих режимов достигла высокого уровня. Три Дайо написал в твиттере: «LLM и новички могут писать ядро со скоростью света» — это именно то, что демонстрируют экспериментальные результаты статьи в реальности.

Результаты эксперимента

Для бенчмаркинга CODA были выбраны сложные соперники: cuBLAS с torch.compile, а также Liger Kernel и FlashInfer, оптимизированные для LLM.

Для каждой ядерной функции оценивались две реализации: CODA (LLM), сгенерированная Claude Code, при этом исследователи предоставили описание примитивов, несколько примеров и постоянно обновляемый журнал техник реализации — ИИ написал основной код, а человек осуществлял легкий контроль; CODA (Human), написанная независимо человеческими программистами с использованием той же высокого уровня идеи переформулировки, но без использования самого набора примитивов CODA. Результаты обеих групп сравнивались с оптимизированными библиотеками, такими как cuBLAS + torch.compile, Liger Kernel и FlashInfer.

На уровне одиночного оператора, на примере типичного паттерна GEMM-RMSNorm-GEMM, CODA превзошла базовую реализацию cuBLAS + PyTorch при скрытых размерностях для трех моделей объемом 1B, 7B и 70B. Аналогичные результаты показаны и для других финальных комбинаций, таких как SwiGLU, RoPE и кросс-энтропия.

Ядра, сгенерированные с помощью LLM, показывают результаты, сопоставимые с ручными версиями на большинстве тестов, а в некоторых конфигурациях даже немного превосходят их. Это весьма необычный вывод в области оптимизации GPU-ядер, традиционно отличающейся крайне высоким порогом входа.

Обратное распространение особенно эффективно: обратное ядро GEMM-Residual-PartialRMS-GEMM демонстрирует ускорение в 1,6–1,8 раза по сравнению с базовой версией, а обратное вычисление SwiGLU улучшается на 1,4–1,6 раза. В этом направлении разрыв между LLM и ручной реализацией также минимальный. Это неудивительно: обратное распространение естественным образом предполагает больше доступов к промежуточным тензорам, что увеличивает выгоду от фьюжена; а первичные операции CODA достаточно четко спроектированы, чтобы позволить ИИ-моделям корректно выполнять комбинирование.

В энд-ту-энд бенчмарке полного трансформерного слоя ускорение прямого прохода CODA составляет от 5% до 20% в зависимости от размера, причем эффект более выражен при больших размерах модели (соответствующих скрытым измерениям размером 70B).

В плане численной точности переparameterизация CODA изменяет момент применения масштабирующего коэффициента RMSNorm, однако эксперименты показывают, что ее численная ошибка сопоставима с эталонной реализацией PyTorch и в некоторых конфигурациях даже меньше — благодаря более точному аккумулятору в основном цикле GEMM.

Что может CODA: справочная таблица, чтобы четко определить границы возможностей CODA, прежде чем перейти к более широкому контексту.

• Охватывает почти все вычисления в прямом и обратном проходах стандартных Transformer (например, архитектуры LLaMA), за исключением внимания и встраивания слов, включая RMSNorm, остаточное сложение, активацию SwiGLU, вращающееся позиционное кодирование RoPE, кросс-энтропийную потерю и обратные градиенты для указанных операций.
• Ускорение: при скрытых измерениях от 1B до 70B на уровне отдельных операторов наблюдается различная степень улучшения по сравнению с базовой линией cuBLAS + torch.compile, причем наибольшая выгода наблюдается при обратном распространении (некоторые ядра достигают ускорения более чем в 1,6 раза); конечное ускорение прямого прохода полного слоя Transformer составляет примерно 5–20%, причем эффект более выражен при больших размерах модели.
• Кто может использовать: CODA, реализованный на основе CuTeDSL (Python DSL для NVIDIA CUTLASS), поддерживающий два способа написания ядер — как человеком-программистом, так и AI-моделью, причем оба способа обеспечивают высокую производительность.
• Текущие ограничения: в настоящее время поддерживается только сценарий с одним GPU, распределённое обучение не предусмотрено; переparameterизация в основном ориентирована на стандартную архитектуру Transformer, применимость к другим архитектурам требует дополнительной проверки.

Заключение

CODA — это не изолированная работа. Это конкретная реализация одного из подходов: на GPU истинное пространство для оптимизации часто заключается не в «чём считать», а в «как переносить».

FlashAttention позволяет вычислениям внимания «переехать» в память на чипе, а CODA пытается сделать то же самое для нормализации и функций активации. Triton снижает барьер для написания пользовательских ядер, а ThunderKittens, TileLang и другие проекты продолжают исследовать эту область на разных уровнях. Вместе эти работы указывают на одну и ту же цель: объединить удобство выражения графа операторов PyTorch с эффективностью выполнения, близкой к ручному CUDA, в единой программируемой системе.

Последнее предложение твита Tri Dao стоит перечитать еще раз: «LLM и новички могут писать ядра с скоростью света для всех операций Transformer». За этим скрывается более глубокая логика: когда абстракции программирования спроектированы достаточно хорошо, сами модели ИИ могут участвовать в оптимизации своей собственной инфраструктуры обучения. Именно этот цикл делает CODA по-настоящему захватывающим.

С этой точки зрения, название «CODA», возможно, несет в себе глубокий смысл. В классической музыке Coda — это заключительная часть произведения, завершающая его. Здесь она является «финалом» ядра GEMM — и грамотное написание этого финала, возможно, станет следующим важным этапом в повышении эффективности систем обучения Transformer.