Faster Generalized Message Passing with Skew-Aware Parallelism and Work Partitioning

Graph neural networks (GNNs) are built from sparse operations that gather and aggregate features along the edges of a graph, and on GPUs their speed is limited by memory access and parallelism rather than by arithmetic. IO-aware kernels remove most of the data movement, yet still leave much of the GPU idle on real graphs and cover only a few fixed layers. We show that the remaining bottleneck is how work is partitioned: degree distributions are heavy-tailed, so a kernel that gives every node one unit of parallel work waits on its few highest-degree nodes while the rest of the device idles. We, therefore, make GNN kernels degree-aware. First, nodes are bucketed by degree, extending a technique of prior IO-aware kernels, then reordered within each bucket and processed concurrently, so that long jobs start first and short ones fill the idle capacity. Second, the heaviest neighborhoods are cut into fixed-size edge slices that are merged exactly, and per-edge computations are partitioned over edges entirely. Third, data is loaded asynchronously wherever there is computation to overlap it. We bring these techniques to the two generalized message-passing primitives, gSpMM and gSDDMM, and provide them as a portable library of $18$ and $32$ kernels with autograd, from which arbitrary GNN convolutions can be composed in pure PyTorch. Against DGL, a widely used and efficient graph learning library, we achieve $17\times$ speedup for GATv2, $2.1\times$ speedup for gSDDMM, and $1.4\times$ for gSpMM in the median over thousands of configurations, reaching $23\times$ on $\text{web-fraud}$ and $3.7\times$ on $\text{amazon-computers}$. Our implementations are also $16\times$ faster than PyG and can process graphs which cause other libraries run out of memory. Against prior IO-aware kernels, they are $2.1\times$ faster in the median and $12.7\times$ on the most skewed graph.