HEAT: NPU-NDP HE terogeneous A rchitecture for T ransformer-Empowered Graph Neural Networks

Transformer-empowered Graph Neural Networks (TF-GNNs) are gaining significant attention in AI research because they leverage the front-end Transformer's ability to process textual data while also harnessing the back-end GNN's capacity to analyze graph structures. Typically, TF-GNNs follow the sequential execution mode, where the front-end Transformer first encodes vertex features, followed by subgraph sampling and subsequent processing by the back-end GNN. However, due to the massive computation workloads of Transformers and the irregular memory access patterns of GNNs, achieving efficient inference for TF-GNNs remains a challenge. Although architectures like FACT and MEGA have been proposed to separately accelerate the Transformer and GNN, they overlook the new opportunities arising from the coupling of the Transformer and GNN.To enable efficient TF-GNNs, we propose HEAT, a heterogeneous architecture with a Neural Processing Unit (NPU) and a DIMM-based Near-Data Processing (NDP). Such a heterogeneous architecture can utilize both the high computational power of NPU and the high internal bandwidth of NDP. To fully unleash the potential of the NPU-NDP architecture, HEAT makes the following three contributions: First, HEAT leverages graph topology to identify the importance of vertices and encodes their features in the Transformer using varying precision accordingly. Second, HEAT gives more flexibility to the execution granularity and execution order of subgraphs in GNN, allowing their DRAM accesses to scatter across different banks and enhancing locality between subgraphs. Third, HEAT decouples the inherent dependency between the front-end Transformer and the back-end GNN by allowing the Transformer to selectively encode only the vertex features required by the current subgraph. This enables concurrent execution of the NPU and NDP, thereby fully exploiting the parallelism of the heterogeneous architecture. Comprehensive evaluations on a wide range of datasets demonstrate that HEAT outperforms the high-performance A100 GPU, the state-of-the-art (SOTA) Transformer accelerator FACT, and GNN accelerator MEGA in terms of execution time and energy efficiency.