Enabling High-Throughput Inference of Transformers on Near-Data Processing Architectures

The rapid adoption of Transformer models in AI has exposed critical inefficiencies in conventional computing architectures, particularly due to their large memory footprint and low data reuse. Near-Data Processing (NDP) architectures have emerged as a promising solution to mitigate the memory wall problem, especially for memory-intensive neural network (NN) workloads. However, existing AI compilation frameworks, optimized for von Neumann architectures, fail to exploit the distributed nature and fine-grained parallelism inherent in the NDP architectures, especially for memory-intensive Transformer models. This work proposes a mapping framework that leverages novel pipeline parallelism to maximize the inference throughput of Transformers on NDP architectures. Firstly, we propose a pipeline layout strategy called Rect-zigzag, which offers superior flexibility to existing pipeline layout strategies (such as 1D static layout and zigzag layout) while adapting to the fine-grained partitioning required by Transformer models. Secondly, we propose a dynamic programming-based mapping algorithm capable of achieving joint optimization of partitioning and pipeline layout for Transformer models. Experiments demonstrate that our method enables the inference throughput of Transformers to reach 1.14∼3.92× that of the baseline methods.