Как проект из ШАДа попал в Spotlight статей на конференции ICML 2026

Граф из миллионов вершин не загружает современную GPU на все 100%: видеокарта почти всё время не вычисляет, а ждёт загрузки данных из памяти ^[1]. Графовые нейросети, или GNN, упираются в это давно: сами операции достаточно простые, но доступ к памяти нерегулярный и разреженный. И чем мощнее GPU, тем заметнее недостаточная её утилизация.

Идея выросла из проектного курса в ШАДе. Толчком стало то, что один из самых популярных фреймворков для работы с графами, Deep Graph Library, на момент начала работы не обновлялся уже около года — это знак того, что в области что‑то застряло.

Меня зовут Федя Великонивцев, я старший исследователь Yandex Research, руковожу группой, которая занимается эффективными вычислениями на GPU. На том курсе мы с коллегами — Дарьей Фоминой из команды ML‑инфраструктуры Яндекса, Вячеславом Ждановским из команды разработки инференса — и студентами Даниилом Красильниковым, Алексеем Бойковым и Андреем Долговязовым взялись выяснить, почему графовые нейросети тормозят на современных GPU.

Так появился проект, который мы оформили в отдельную статью — On Efficient Scaling of GNNs via IO‑Aware Layer Implementations ^[2]. Её приняли на ICML-2026 со статусом Spotlight. Для контекста: из 23 918 поданных работ приняли 6 352 (26,6%), а Spotlight достался только 536 работам — это 2,2% заявок с самыми высокими оценками программного комитета.

Дальше расскажу, как мы прошли путь от этого вопроса до трёх семейств специализированных GPU‑кернелов — с парой неожиданных находок по дороге.

Почему графовые нейросети до сих пор медленные

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

Но есть проблема: графы очень неудобны для GPU.

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

Сами же графовые агрегации — это в большинстве своём вычислительно несложные операции (например, усреднение признаков по соседям), но есть и исключения в виде порёберного внимания ^[3], когда признаки соседей агрегируются с обучаемыми коэффициентами по аналогии с трансформерами.

Всё это приводит к тому, что GPU нередко проводят больше времени за ожиданием данных из памяти, чем за выполнением вычислений.

По мере роста производительности GPU эта проблема становится только заметнее: вычислительные блоки становятся быстрее, а пропускная способность памяти растёт значительно медленнее.

FlashAttention показал путь

Похожая проблема несколько лет назад возникла у трансформеров. Работа FlashAttention и её последующие продолжения показали, что ускорение может достигаться в том числе за счёт уменьшения количества обращений к памяти.

Мы решили посмотреть на графовые нейросети с той же точки зрения ^[4]. Соответственно, один из вопросов, который мы задали, был: «Что, если оптимизировать не сами вычисления GNN, а то, как данные перемещаются между HBM (High Bandwidth Memory) и быстрой памятью GPU?»

Три семейства графовых операций

Большинство популярных GNN можно свести к трём типам вычислительных паттернов:

1. Свертки, представляемые в виде матричного умножения, — простые арифметические операции взвешенной суммы признаков соседей, например усреднение. Такие операции представляются в виде умножения матрицы признаков на разреженную матрицу смежности графа — этот паттерн вычисления называется Sparse‑Dense Matrix Multiplication (SpMM).

2. Свертки на основе редукции — различные варианты min/max‑агрегации и pooling’ов, которые нельзя представить в виде матричного умножения.

3. Attention‑свертки — свертки с обучаемыми коэффициентами для агрегации признаков. Например, GATv2 и Graph Transformer.

Для каждого семейства боттлнеки оказались разными, поэтому и решения пришлось разрабатывать отдельно.

Attention для графов по мотивам FlashAttention

Расскажу, как мы получили заметный выигрыш для attention‑based‑моделей.

Обычные реализации GATv2 и Graph Transformer создают большие промежуточные тензоры для каждого ребра графа. Эти данные записываются в глобальную память и затем снова считываются для вычисления softmax и агрегации.

Мы убрали этот этап.

Вместо материализации промежуточных значений вычисления выполняются непосредственно во время прохода по соседям вершины. Для этого используется online softmax — тот же базовый принцип, который лежит в основе FlashAttention.

В результате:

• уменьшается объём передаваемых данных;

• снижается нагрузка на HBM;

• сокращается пиковое потребление памяти.

Для backward pass мы используем идею из FlashAttention-2: веса внимания не сохраняются целиком, а пересчитываются на лету из компактной статистики.

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

Кроме того, мы реализовали вариант с блочно‑разреженным представлением графа, позволяющим использовать тензорные ядра — специализированные модули на GPU, отвечающие за выполнение инструкций матричного умножения. Такой подход особенно эффективен для графов с высокой плотностью.

Что делать с вершинами, у которых тысячи соседей

Для операций на основе редукции проблема оказалась другой. У реальных графов обычно тяжёлый хвост распределения степеней: большинство вершин имеют немного соседей, но существует небольшое количество очень крупных вершин. Именно такие вершины часто становятся главным источником замедления.

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

Вычисления мы распределяем по‑разному. Для тяжёлых вершин мы параллельно считаем редукцию по чанкам и затем параллельно агрегируем результат. Мы можем себе это позволить, так как такая операция ассоциативна (то есть не зависит от порядка вычисления).

Это позволяет лучше загрузить GPU и убрать дисбаланс вычислений.

Не всегда нужно писать собственный кернел

Один из интересных результатов работы связан с GNN, представляемыми в виде SpMM. Мы ожидали, что специализированные реализации смогут существенно обогнать существующие решения.

На практике оказалось иначе.

Для многих SpMM‑based‑операций современные библиотеки NVIDIA (например, взятая нами cuSPARSE ^[5]) уже настолько хорошо оптимизированы, что оказываются быстрее большинства опубликованных специализированных реализаций. После правильного кеширования и настройки cuSPARSE зачастую показывает лучший результат.

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

Результаты

На наборе реальных графов и популярных архитектур нам удалось получить:

• до 8,5× ускорения для attention‑сверток;

• до 10× ускорения для некоторых сверток‑редукций операций;

• до 8× ускорения при использовании модифицированного cuSPARSE для SpMM‑сверток;

• до 76× снижения пикового потребления памяти в отдельных конфигурациях.

Особенно важно, что все реализации доступны как drop‑in replacement для популярных GNN‑фреймворков. Для использования не требуется переписывать модель или менять архитектуру.

Помогает ли переупорядочивание вершин в графе

Граф — структура без порядка: вершины можно пронумеровать как угодно, и это не повлияет на результат вычислений. Но то, в каком порядке признаки вершин лежат в памяти GPU, очень влияет на производительность.

Дело в том, что GPU подгружает данные из HBM не по одному числу, а блоками. Если нужный байт лежит рядом с уже загруженными, он достаётся «бесплатно», прямо из кеша. При этом связи в графе — кто с кем сосед — изначально никак не привязаны к тому, как признаки вершин разложены в памяти. Когда кернел проходит по соседям одной вершины, он может прыгать по совершенно случайным адресам. Каждый такой прыжок — отдельный «холодный» поход в HBM: кеш ограничен, и данные, которые мы доставали секунду назад, могли уже быть вытеснены другими.

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

Вот так, например, может выглядеть матрица смежности графа после переупорядочивания:

Цвет показывает плотность ненулевых значений в блоке 128 × 128. Слева — до переупорядочивания: соседи разбросаны по всей матрице. Справа — после. Видно, что вершины по большей части сконцентрированы рядом с теми, с которыми они связаны

Кстати, этим вопросом занимались и до нас. Например, работа Can Graph Reordering Speed Up Graph Neural Network Training ^[6] показывает, что для обучения ^[7] GNN переупорядочивание может давать честный выигрыш. Мы решили разобраться, при каких условиях это работает, а при каких — нет.

Оказалось, что эффект зависит от того, как именно кернел читает данные. Мы сравнили два паттерна доступа: один параллелит вычисления по признакам (feature‑parallel, данные читаются предсказуемо, подряд), другой — по соседям (neighbor‑parallel, каждый поток читает своего соседа из случайного места). Переупорядочивание помогает в основном второму паттерну — там, где без него память читается совсем вразнобой.

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

А вот на графах с низкой степенью — скажем, дорожных сетях, где у вершины в среднем два соседа, — выигрыша почти нет. Соседей просто слишком мало, чтобы локальность вообще на что‑то влияла, и время съедают накладные расходы на каждую вершину, а не промахи кеша.

Получается, переупорядочивание — это не универсальный буст, а инструмент, который нужно применять осознанно: смотреть на структуру графа и на то, как написан кернел, который будет по нему бегать.

Вместо заключения

Последние несколько лет машинное обучение училось простой вещи: быстрый алгоритм — это не только хорошая математика ^[8], но и реализация, которая дружит с железом. FlashAttention показал это на трансформерах. Мы убедились, что с графовыми нейросетями всё ровно так же.

И чем больше становятся графы, чем быстрее GPU, тем сильнее перекос: считать мы научились быстро, а вот подвозить данные из памяти — нет. Так что работа с памятью потихоньку превращается в отдельную инженерную задачу, не менее важную, чем сами вычисления. Похоже, писать алгоритмы под конкретное железо придётся всё чаще.

Подробнее — в нашей статье ^[2]. Исходный код и библиотека с GPU‑кернелами доступны на GitHub ^[9] и могут использоваться как drop‑in replacement для существующих архитектур.

Если тема для вас актуальна, подписывайтесь на YR Reading Group ^[10]. А всё самое интересное с конференции, как всегда, расскажем в MLunderhood ^[11].