Оптимизация операций в GNN: ускорение до 10x

🛠 Метод Работа фокусируется на ускорении трех ключевых типов операций в 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. Хоть она и довольно нишевая и на любителя, тем не менее импакт более чем серьезный. Интересно, где еще есть низковисящие фрукты для оптимизации?