🛠 Метод

Работа фокусируется на ускорении трех ключевых типов операций в GNN:

- 🔹 Разреженные свертки.
- 🔹 Редукции на графах (min/max).
- 🔹 Внимание на графах.

Замечают, что распределение вершин графов по степеням обыкновенно имеет тяжелые хвосты. Потому разделяют вершины на две группы по степени — легкие и тяжелые вершины.

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

Объем подгружаемых данных не меняется, но ускорение достигается за счет того, что один блок обрабатывает несколько узлов, тем самым повышая memory-level параллелизм.

Для слоев внимания реализуют кастомный вариант FlashAttention для CSR-матриц. В наивной реализации материализовали матрицы внимания, а тут оно и не нужно.

CSR-кернелы для произвольной sparsity уже неплохи, но не используют ценный ресурс современных GPU — тензорные ядра. Если запаковать окрестности вершин в тайлы 16×16, то можно воспользоваться ими и, даже выполняя некоторые лишние вычисления, выжать ускорение.

Кроме того, существующие либы почему-то не используют cuSPARSE — рабочий и мощный инструмент. Одно его накатывание + кастомные адаптации уже дает заметный профит.

📊 Эксперименты

Эффективность реализации проверяют на бенчмарке GraphLand — наборе задач с большим разнообразием структур и свойств графов.

Удается выжать медианное ускорение порядка 2× (вплоть до 8.5×) против DGL на GATv2, медианное ускорение 2.6× (вплоть до 10×) на операциях редукции. Операции разреженных сверток на cuSPARSE дают вплоть до 8× ускорения. Кроме того, пиковое потребление памяти оказывается обычно в разы меньше (а то и в десятки раз) по сравнению с DGL. PyG тоже в среднем ест меньше GPU VRAM, но при этом он куда медленнее.

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

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

💡 Выводы

Классная и полезная работа, дающая существенный толчок всей области GraphML. Хоть она и довольно нишевая и на любителя, тем не менее импакт более чем серьезный. Интересно, где еще есть низковисящие фрукты для оптимизации?