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.

错误:搜索内容不能为空,请输入英文关键词
错误:关键词超出字数限制,请精简
高级检索

GNNCacheInfer: Efficient Graph Neural Network Inference via Embedding Cache

  • Junkun Shen,
  • Yuezhi Che,
  • Qiao Zhan,
  • Dazhao Cheng

摘要

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.