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

Автор статті, джерело: Machine Learning News

22 травня Tri Dao у соціальних мережах перепостив твіт Han Guo. Він також написав: «Після деяких математичних перетворень виявилось, що все в 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, спочатку потрібно зрозуміти, куди витрачається час на навчання великих моделей.

Під час навчання моделі з 1 мільярдом параметрів стилю LLaMA-3 на одній GPU NVIDIA H100 більшість людей інтуїтивно вважають, що весь час витрачається на матричне множення та обчислення уваги, адже саме вони є «справжніми обчисленнями». Ця інтуїція в цілому правильна: матричне множення (GEMM) та увага дійсно займають основну частину обчислювальних ресурсів.

Але якщо ви уважно розглянете профілювальник, ви помітите, що ще є серія «малих операторів», які тихо споживають час: нормалізація (RMSNorm), функції активації (SwiGLU, RoPE), залишкова додавання, зведення між шарами… Вони окремо мають невеликий обчислювальний навантаження, але часто переміщують великі проміжні тензори з пам’яті GPU.

Це так званий «обмеження пропускної здатності пам’яті»: як якщо б кухар з ідеальним майстерством мусив кожного разу переносити інгредієнти з віддаленого складу, використовувати їх і повертати назад, а не тримати на робочій поверхні поруч. Незалежно від швидкості руху рук кухаря, час очікування перенесення — це реальна втрата.

Ще гірше те, що зі зростанням швидкості матричних обчислень завдяки низькоточним форматам, таким як FP8 і FP4 від NVIDIA, відносна вартість цих «переміщень» зростає: матричне множення прискорюється, а витрати на перенесення тензорів не скорочуються пропорційно.

У статті наведено наочні дані: під час навчання моделі з 1 мільярдом параметрів за допомогою TorchTitan на H100, нематричні операції займають значну частину загального часу виконання, і цей відсоток ще більше зростає з введенням точності FP8.

Існуючі програмні фреймворки майже безсили перед цим. PyTorch представляє обчислення Transformer як послідовність операторів із чіткими межами між ними. Ці межі дуже добре підходять для автоматичного диференціювання (autograd), але саме вони перешкоджають оптимізації об’єднання між операторами: кожна межа оператора часто є непотрібним записом у пам’ять.

CODA: У «фіналі» сховано скарби

Вихідною точкою CODA є просте спостереження.

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

Значення фіналу полягає в тому, що наразі вихід матричного множення ще «живе» у регістрах на кристалі, а ще не був записаний у глобальну пам’ять GPU. Це короткий золотий вікно: якщо в цей момент виконати додаткові обчислення, можна повністю уникнути однієї подвійної операції запису та зчитування з пам’яті GPU.

Основний інсайт CODA полягає в тому, що багато пам’ятємістких операцій у Transformer можна алгебраїчно переформулювати і виконати всередині цього «фінального» вікна.

Це вимагає певних математичних навичок. Наприклад, у найпоширенішому шаблоні GEMM-RMSNorm-GEMM: результат множення матриць проходить через додавання залишкової зв’язки, RMS-нормалізацію, а потім ще одне множення матриць. У традиційному підході три окремі оператори виконуються послідовно, а проміжні результати двічі зберігаються у відеопам’яті.

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

Подібні переформулювання застосовні й до операцій, таких як SwiGLU, RoPE (обертальні позиційні кодування), функція втрат перехресної ентропії, а також для зворотного поширення. У статті наведено теорему, що доводить: якщо прямий прохід є «блочним локальним», то зворотне поширення автоматично успадковує ту саму структуру. Докладніше дивіться у оригінальній статті.

П’ять «кірпичиків» і набір «мови LEGO»

CODA — це не конкретний фузійний ядро, а набір програмних абстракцій.

Він фіксує оптимізований експертами основний цикл GEMM, а потім у кінці відкриває п’ять типів комбінованих базових примітивів:

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

За допомогою цих п’яти типів блоків можна охопити майже всі операції, крім уваги, в прямому та зворотному поширенні стандартного Transformer.

Ще цікавіше, наскільки ця абстракція відкрита для того, хто пише код. У експерименті стаття оцінила два підходи: один — коли код пишуть людські програмісти, інший — коли код генерується Claude Code: за допомогою початкових описів CODA, кількох прикладів та журналів реалізації AI виконує більшу частину ядерного коду, а людина здійснює легкий нагляд.

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

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

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

У роботі для кожного ядра оцінено два варіанти реалізації: CODA (LLM), згенеровану Claude Code, яка отримала від дослідників опис примітивів, кілька прикладів та щоденно оновлюваний журнал технік реалізації — AI написав основний код, а людина здійснювала легкий контроль; 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 достатньо чіткі, щоб дозволити моделям ШІ правильно виконати комбінування.

У повному енд-ту-енд бенчмарку для цілої шару Transformer, прискорення прямого проходу 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, без використання розподіленого навчання; репараметризація в першу чергу призначена для стандартної архітектури Transformer, а застосовність до інших архітектур потребує додаткової перевірки.

Висновок

CODA — це не ізольована робота. Це конкретна реалізація певного підходу: на GPU справжній простір для оптимізації часто полягає не в тому, «що обчислювати», а в тому, «як переносити».

FlashAttention дозволяє обчислення уваги «переїхати» до пам’яті на кристалі, а CODA намагається зробити те саме з нормалізацією та функціями активації. Triton знижує бар’єри для написання власних ядер, а ThunderKittens, TileLang та інші досліджують цю область на різних рівнях. Ці роботи разом вказують на одну й ту ж мету: поєднати зручність виразу графів операторів PyTorch із ефективністю виконання, близькою до ручного CUDA, у єдиній програмованій фреймворці.

Останній рядок твіту Tri Dao варто перечитати ще раз: «LLM та новачки можуть писати ядра зі швидкістю світла для всіх операцій Transformer». За цим стоїть глибша логіка: коли абстракції програмування розроблені достатньо добре, самі AI-моделі можуть брати участь у оптимізації своєї власної інфраструктури навчання. Саме цей цикл — найбільш захоплива частина CODA.

З цієї точки зору, назва «CODA» може мати глибший зміст. У класичній музиці Coda — це фінальна частина твору, що завершує його. Тут вона є «фіналом» ядра GEMM — і правильно написаний цей фінал, можливо, стане наступним важливим етапом у підвищенні ефективності систем навчання Transformer.