ADiP: Adaptive Precision Systolic Array for Matrix Multiplication Acceleration

TL;DR

Transformers demand intensive matrix multiplications, driving a need for efficient hardware under quantization. The authors introduce ADiP, an adaptive-precision systolic array with NxN PEs, shared shifters and accumulators, and a DiP-like dataflow to support symmetric and asymmetric matrix multiplications across multiple bit-precisions. Analytical latency/throughput models, a 22nm design-space study, and attention-workload evaluations (GPT-2 Medium, BERT Large, BitNet-1.58B) show up to 4× throughput gains and notable latency/energy improvements, especially for highly quantized workloads. ADiP demonstrates high compute density and scalable performance for quantized Transformer inference, making it suitable for edge and energy-constrained deployments.

Abstract

Transformers are at the core of modern AI nowadays. They rely heavily on matrix multiplication and require efficient acceleration due to their substantial memory and computational requirements. Quantization plays a vital role in reducing memory usage, and can be exploited for computations by designing reconfigurable architectures that enhance matrix multiplication by dynamically adjusting the precision. This paper proposes ADiP, a novel adaptive-precision systolic array architecture designed for efficient matrix multiplication acceleration.The proposed architecture consists of NxN adaptive-precision processing elements (PEs) and shared accumulators. ADiP supports multiple computation modes, including symmetric single-matrix multiplication as well as asymmetric multi-matrix multiplication with a shared input matrix, thereby improving data-reuse and PE utilization. In addition, ADiP maximizes the computational density by adapting to different precisions, such as 8bitx8bit, 8bitx4bit, and 8bitx2bit. Analytical models are developed for ADiP architecture, including latency and throughput for versatile architecture configurations. A comprehensive hardware design space exploration is demonstrated using 22nm commercial technology, achieving up to a 4x higher computational throughput. Furthermore, ADiP is evaluated on different transformer workloads from GPT-2 Medium, BERT Large, and BitNet-1.58B models, delivering latency improvement up to 53.6%, and energy improvement up to 24.4% for BitNet-1.58B MHA workloads. At a 64x64 size with 4096 PEs, ADiP achieves a peak throughput of 8.192 TOPS, 16.384 TOPS, and 32.768 TOPS for 8bitx8bit, 8bitx4bit, and 8bitx2bit operations, respectively.

Figures (10)

• Figure 1: Matrix multiplication stages for multi-head attention layers.
• Figure 2: Analytical modeling results for ADiP architecture in terms of D-MUL latency, ADiP latency and ADiP throughput across varying numbers of 2-bit multipliers ($M$ = 2, 4, 8, 16) at ADiP size of 64×64. Each group of bars represents different operand configurations: 8b×8b, 8b×4b, and 8b×2b. (a) D-MUL latency demonstrates a reduction as $M$ increases and saturates at one cycle. (b) ADiP latency representing the combined latency of D-MUL and ADiP systolic array. (c) Throughput for the ADiP architecture improves significantly as the number of 2-bit multipliers $M$ increases and precision decreases, highlighting the throughput gain.
• Figure 3: (a) $N$×$N$ ADiP systolic array architecture. Inputs (I) move diagonally across PE rows, transitioning from one row to the next. The boundary PEs are diagonally connected, so that the registered inputs from the leftmost PE column feed into the inputs of the rightmost PE column in the subsequent row. Weights are loaded vertically, and psums are accumulated vertically along the columns as well. (b) PE block diagram, consisting of 16× 2-bit MULs, arranged into four groups with their internal accumulators. In addition, four dedicated partial summations (psums) accumulators, and enabled registers for weights, inputs and psums. (c) Shows the unit of shared shifters and accumulators per each PE column adapting to different precisions (8b×8b, 8b×4b, and 8b×2b).
• Figure 4: An example of interleaving process in ADiP for different matrix multiplications: (a) 8b×8b multiplication using a single tile without interleaving, (b) 8b×4b multiplication with two interleaved weight tiles, where every two adjacent elements are converted to 4-bit precision, (c) 8b×2b multiplication with four interleaved weight tiles, where every four adjacent elements are converted to 2-bit, (d) Matrix multiplication (8b×2b) for Q, K, and V workloads in attention layers, where three weight tiles (one from each matrix) are interleaved, with every three adjacent elements converted to 2-bit.
• Figure 5: An example of weights preparation for 8b√ó2b matrix multiplication on the ADiP architecture. The preprocessing involves permutation and interleaving four weight tiles, with every group of four adjacent elements converted to 2-bit and concatenated.