GNNCacheInfer: Efficient Graph Neural Network Inference via Embedding Cache
摘要
Graph Neural Networks (GNNs) have emerged as powerful tools for a wide range of graph-based tasks. During mini-batch GNN inference, node-wise inference often results in significant redundant computation. To address this, existing methods propose layer-wise inference, which mitigates redundant computation of node embeddings. However, while layer-wise inference is well-suited for full-graph inference, it remains inefficient for partial-graph inference, as redundant computations persist, leading to suboptimal performance. In this work, we present GNNCacheInfer, a general-purpose framework designed to accelerate GNN inference by leveraging embedding cache. GNNCacheInfer is versatile, supporting both partial-graph and full-graph inference. Additionally, we introduce two key optimization techniques: GPU-accelerated embedding cache access and dynamic mini-batch size scheduling. GPU-accelerated embedding cache access employs Unified Virtual Addressing (UVA) to enable GPUs to directly access and update CPU memory, eliminating redundant memory copies. Dynamic mini-batch size scheduling adaptively adjusts the mini-batch size to improve GPU utilization and balance workload dynamics. Our evaluation demonstrates the effectiveness of GNNCacheInfer, establishing it as a practical and efficient solution for real-world GNN Inference.