Обучение современных трансформеров завязано на плотные линейные алгебраические операции — в первую очередь на GEMM. Но немалая часть времени уходит вовсе не на матричные умножения, а на окружающие их операции с памятью: нормализацию, функции активации, обновление residual connection, редукции и прочие «склейки». Эти маленькие, но многочисленные шаги каждый раз гоняют огромные промежуточные тензоры через глобальную память GPU, хотя арифметики там почти нет. В итоге data movement становится главным тормозом в хорошо оптимизированных пайплайнах.
Исследователи представили CODA — новую абстракцию GPU-ядер, которая переписывает эти операции как GEMM-plus-epilogue программы. Ключевая идея: многие операторы, которые обычно выносятся в отдельные framework-ядра, можно алгебраически перепараметризовать так, чтобы они выполнялись прямо на чипе, пока тайл результата GEMM ещё лежит в регистрах — до того, как его запишут в память. Вместо того чтобы писать новый сложный mainloop, CODA фиксирует уже готовый «основной цикл» GEMM и даёт разработчику небольшой набор композовых epilogue-примитивов: скейлинг, редукции, попарные преобразования и накопление.
Этого ограниченного интерфейса хватает, чтобы покрыть почти все не-attention вычисления в forward и backward pass стандартного Transformer-блока. При этом он сохраняет производительность экспертных GEMM-ядер. На практике и человеческие, и сгенерированные LLM-моделями CODA-ядра работают так же быстро, как рукописный оптимизированный код. Авторы считают, что подход GEMM-plus-epilogue — это практичный путь, чтобы совместить продуктивность разработки на уровне фреймворка с аппаратной эффективностью.