On Efficient Scaling of GNNs via IO-Aware Layers Implementations Графовые нейросети (GNN) плохо дружат с GPU: агрегация по соседям, основная операция в GNN, представляет собой чтения из неупорядоченных участков GPU-памяти с последующей агрегацией. Это плохо укладывается на современное железо, заточенное в первую очередь под упорядоченное чтение (например, в матричных умножениях/атеншене). В статье (спотлайт на ICML) сделали для GNN примерно то же, что FlashAttention сделал для трансформеров: переписали основные слои так, чтобы они меньше гоняли данные между памятью и вычислительными блоками. Для attention-слоёв используется IO-aware подход в духе FlashAttention, для агрегирующих - дополнительная параллелизация вершин с большим числом соседей. А ещё авторы показали, что для части свёрточных слоёв современные решения NVIDIA уже быстрее большинства специализированных реализаций. Результат - до 8.5× ускорения и до 76× меньше памяти в отдельных сценариях. Все реализации доступны как drop-in замена для популярных GNN-фреймворков. Кстати, статья - это по сути проект преподавателей и студентов ШАДа) Блогпост @ai_newz