# DF-GNN: Dynamic Fusion Framework for Attention Graph Neural Networks on GPUs

arXiv: https://arxiv.org/abs/2411.16127 · HTML: https://mlsanity.com/paper/2411.16127/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 22 of its 28 references that are in the map.
- First posted on arXiv 25 Nov 2024.

## Builds on (most cited first, 22 of 22)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [Graph Attention Networks](https://mlsanity.com/paper/1710.10903/) (2017)
- [How Powerful are Graph Neural Networks?](https://mlsanity.com/paper/1810.00826/) (2018)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Neural Message Passing for Quantum Chemistry](https://mlsanity.com/paper/1704.01212/) (2017)
- [Fast Graph Representation Learning with PyTorch Geometric](https://mlsanity.com/paper/1903.02428/) (2019)
- [How Attentive are Graph Attention Networks?](https://mlsanity.com/paper/2105.14491/) (2021)
- [Link Prediction Based on Graph Neural Networks](https://mlsanity.com/paper/1802.09691/) (2018)
- [Recipe for a General, Powerful, Scalable Graph Transformer](https://mlsanity.com/paper/2205.12454/) (2022)
- [Benchmarking Graph Neural Networks](https://mlsanity.com/paper/2003.00982/) (2020)
- [A Generalization of Transformer Networks to Graphs](https://mlsanity.com/paper/2012.09699/) (2020)
- [InfoGraph: Unsupervised and Semi-supervised Graph-Level Representation Learning via Mutual Information Maximization](https://mlsanity.com/paper/1908.01000/) (2019)
- [Neural Bellman-Ford Networks: A General Graph Neural Network Framework for Link Prediction](https://mlsanity.com/paper/2106.06935/) (2021)
- [A Fair Comparison of Graph Neural Networks for Graph Classification](https://mlsanity.com/paper/1912.09893/) (2019)
- [Deep Graph Library: A Graph-Centric, Highly-Performant Package for Graph Neural Networks](https://mlsanity.com/paper/1909.01315/) (2019)
- [Attention-based Graph Neural Network for Semi-supervised Learning](https://mlsanity.com/paper/1803.03735/) (2018)
- [SparseTIR: Composable Abstractions for Sparse Compilation in Deep Learning](https://mlsanity.com/paper/2207.04606/) (2022)
- [FeatGraph: A Flexible and Efficient Backend for Graph Neural Network Systems](https://mlsanity.com/paper/2008.11359/) (2020)
- [Understanding GNN Computational Graph: A Coordinated Computation, IO, and Memory Perspective](https://mlsanity.com/paper/2110.09524/) (2021)
- [FusedMM: A Unified SDDMM-SpMM Kernel for Graph Embedding and Graph Neural Networks](https://mlsanity.com/paper/2011.06391/) (2020)

## Built on by (most cited first, 2 of 2)

- [Fused3S: Fast Sparse Attention on Tensor Cores](https://mlsanity.com/paper/2505.08098/) (2025)
- [FlashSchNet: Fast and Accurate Coarse-Grained Neural Network Molecular Dynamics](https://mlsanity.com/paper/2602.13140/) (2026)

## Abstract (arXiv)

Attention Graph Neural Networks (AT-GNNs), such as GAT and Graph Transformer, have demonstrated superior performance compared to other GNNs. However, existing GNN systems struggle to efficiently train AT-GNNs on GPUs due to their intricate computation patterns. The execution of AT-GNN operations without kernel fusion results in heavy data movement and significant kernel launch overhead, while fixed thread scheduling in existing GNN kernel fusion strategies leads to sub-optimal performance, redundant computation and unbalanced workload. To address these challenges, we propose a dynamic kernel fusion framework, DF-GNN, for the AT-GNN family. DF-GNN introduces a dynamic bi-level thread scheduling strategy, enabling flexible adjustments to thread scheduling while retaining the benefits of shared memory within the fused kernel. DF-GNN tailors specific thread scheduling for operations in AT-GNNs and considers the performance bottleneck shift caused by the presence of super nodes. Additionally, DF-GNN is integrated with the PyTorch framework for high programmability. Evaluations across diverse GNN models and multiple datasets reveal that DF-GNN surpasses existing GNN kernel optimization works like cuGraph and dgNN, with speedups up to 7.0× over the state-of-the-art non-fusion DGL sparse library. Moreover, it achieves an average speedup of 2.16× in end-to-end training compared to the popular GNN computing framework DGL.