Граф из миллионов вершин не загружает современную GPU на все 100%: видеокарта почти всё время не вычисляет, а ждёт загрузки данных из памяти. Графовые нейросети, или GNN, упираются в это давно: сами операции достаточно простые, но доступ к памяти нерегулярный и разреженный. И чем мощнее GPU, тем заметнее недостаточная её утилизация.
Идея выросла из проектного курса в ШАДе. Толчком стало то, что один из самых популярных фреймворков для работы с графами, Deep Graph Library, на момент начала работы не обновлялся уже около года — это знак того, что в области что‑то застряло.
Меня зовут Федя Великонивцев, я старший исследователь Yandex Research, руковожу группой, которая занимается эффективными вычислениями на GPU. На том курсе мы с коллегами — Дарьей Фоминой из команды ML‑инфраструктуры Яндекса, Вячеславом Ждановским из команды разработки инференса — и студентами Даниилом Красильниковым, Алексеем Бойковым и Андреем Долговязовым взялись выяснить, почему графовые нейросети тормозят на современных GPU.
Так появился проект, который мы оформили в отдельную статью — On Efficient Scaling of GNNs via IO‑Aware Layer Implementations. Её приняли на ICML-2026 со статусом Spotlight. Для контекста: из 23 918 поданных работ приняли 6 352 (26,6%), а Spotlight достался только 536 работам — это 2,2% заявок с самыми высокими оценками программного комитета.
Дальше расскажу, как мы прошли путь от этого вопроса до трёх семейств специализированных GPU‑кернелов — с парой неожиданных находок по дороге.
Почему графовые нейросети до сих пор медленные
Графовые нейросети используются во множестве прикладных задач: рекомендательные системы, антифрод, анализ социальных сетей, работа с молекулами, прогнозирование временных рядов.
Но есть проблема: графы очень неудобны для GPU.
В отличие от трансформеров, где данные обычно лежат в памяти компактно и регулярно, для выполнения агрегации признаков соседей графовые модели постоянно обращаются к случайным участкам памяти. Дело в том, что на вершинах нет порядка и индексы вершин могут быть выданы произвольно. У одной вершины может быть три соседа, у другой — три тысячи, а их признаки могут находиться в совершенно разных местах памяти.
Сами же графовые агрегации — это в большинстве своём вычислительно несложные операции (например, усреднение признаков по соседям), но есть и исключения в виде порёберного внимания, когда признаки соседей агрегируются с обучаемыми коэффициентами по аналогии с трансформерами.
Всё это приводит к тому, что GPU нередко проводят больше времени за ожиданием данных из памяти, чем за выполнением вычислений.
По мере роста производительности GPU эта проблема становится только заметнее: вычислительные блоки становятся быстрее, а пропускная способность памяти растёт значительно медленнее.
FlashAttention показал путь
Похожая проблема несколько лет назад возникла у трансформеров. Работа FlashAttention и её последующие продолжения показали, что ускорение может достигаться в том числе за счёт уменьшения количества обращений к памяти.
Мы решили посмотреть на графовые нейросети с той же точки зрения. Соответственно, один из вопросов, который мы задали, был: «Что, если оптимизировать не сами вычисления GNN, а то, как данные перемещаются между HBM (High Bandwidth Memory) и быстрой памятью GPU?»
Три семейства графовых операций
Большинство популярных GNN можно свести к трём типам вычислительных паттернов:
Свертки, представляемые в виде матричного умножения, — простые арифметические операции взвешенной суммы признаков соседей, например усреднение. Такие операции представляются в виде умножения матрицы признаков на разреженную матрицу смежности графа — этот паттерн вычисления называется Sparse‑Dense Matrix Multiplication (SpMM).
Свертки на основе редукции — различные варианты min/max‑агрегации и pooling’ов, которые нельзя представить в виде матричного умножения.
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) уже настолько хорошо оптимизированы, что оказываются быстрее большинства опубликованных специализированных реализаций. После правильного кеширования и настройки 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 показывает, что для обучения GNN переупорядочивание может давать честный выигрыш. Мы решили разобраться, при каких условиях это работает, а при каких — нет.
Оказалось, что эффект зависит от того, как именно кернел читает данные. Мы сравнили два паттерна доступа: один параллелит вычисления по признакам (feature‑parallel, данные читаются предсказуемо, подряд), другой — по соседям (neighbor‑parallel, каждый поток читает своего соседа из случайного места). Переупорядочивание помогает в основном второму паттерну — там, где без него память читается совсем вразнобой.
Дальше — интереснее: эффект зависит ещё и от степени вершин. На графах с высокой степенью — например, в графах покупок или социальных сетях, где у вершины могут быть десятки и сотни соседей, — переупорядочивание даёт стабильный и заметный прирост: там действительно есть что кешировать.
А вот на графах с низкой степенью — скажем, дорожных сетях, где у вершины в среднем два соседа, — выигрыша почти нет. Соседей просто слишком мало, чтобы локальность вообще на что‑то влияла, и время съедают накладные расходы на каждую вершину, а не промахи кеша.
Получается, переупорядочивание — это не универсальный буст, а инструмент, который нужно применять осознанно: смотреть на структуру графа и на то, как написан кернел, который будет по нему бегать.
Вместо заключения
Последние несколько лет машинное обучение училось простой вещи: быстрый алгоритм — это не только хорошая математика, но и реализация, которая дружит с железом. FlashAttention показал это на трансформерах. Мы убедились, что с графовыми нейросетями всё ровно так же.
И чем больше становятся графы, чем быстрее GPU, тем сильнее перекос: считать мы научились быстро, а вот подвозить данные из памяти — нет. Так что работа с памятью потихоньку превращается в отдельную инженерную задачу, не менее важную, чем сами вычисления. Похоже, писать алгоритмы под конкретное железо придётся всё чаще.
Подробнее — в нашей статье. Исходный код и библиотека с GPU‑кернелами доступны на GitHub и могут использоваться как drop‑in replacement для существующих архитектур.
Если тема для вас актуальна, подписывайтесь на YR Reading Group. А всё самое интересное с конференции, как всегда, расскажем в MLunderhood.