Allspark: Workload Orchestration for Visual Transformers on Processing In-Memory Systems

Abstract

The advent of Transformers has revolutionized computer vision, offering a powerful alternative to convolutional neural networks (CNNs), especially with the local attention mechanism that excels at capturing local structures within the input and achieve state-of-the-art performance. Processing in-memory (PIM) architecture offers extensive parallelism, low data movement costs, and scalable memory bandwidth, making it a promising solution to accelerate Transformer with memory-intensive operations. However, the crucial issue lies in efficiently deploying an entire model onto resource-limited PIM system while parallelizing each transformer block with potentially many computational branches based on local-attention mechanisms. We present Allspark, which focuses on workload orchestration for visual Transformers on PIM systems, aiming at minimizing inference latency. Firstly, to fully utilize the massive parallelism of PIM, Allspark employs a fine-grained partitioning scheme for computational branches, and formats a systematic layout and interleaved dataflow with maximized data locality and reduced data movement. Secondly, Allspark formulates the scheduling of the complete model on a resource-limited distributed PIM system as an integer linear programming (ILP) problem. Thirdly, as local-global data interactions exhibit complex yet regular dependencies, Allspark provides a two-stage placement method, which simplifies the challenging placement of computational branches on the PIM system into the structured layout and greedy-based binding, to minimize NoC communication costs. Extensive experiments on 3D-stacked DRAM-based PIM systems show that Allspark brings $1.2\times$ $\sim$ $24.0\times$ inference speedup for various visual Transformers over baselines. Compared to Nvidia V100 GPU, Allspark-enriched PIM system yields average speedups of $2.3\times$ and energy savings of $20\times$ $\sim$ $55\times$.