T-REX: Hardware–Software Co-Optimized Transformer Accelerator With Reduced External Memory Access and Enhanced Hardware Utilization

Transformers have emerged as the dominant neural network architecture thanks to their versatility- and human-level performance. However, a transformer model contains a very large number of parameters, which makes the overhead associated with external memory access (EMA), a critical bottleneck in computing systems. In light of this, we develop T-REX, a hardware accelerator for transformers with the goal of minimizing EMA. We develop a novel factorizing training method to reduce the parameter count substantially. We also combine a range of data compression techniques to effectively compress dense and sparse parameter matrices. Furthermore, we develop a dynamic batching technique, which can process multiple short inputs in parallel to increase parameter reuse. Last but not least, we develop a two-directional accessible register file (TRF) and the optimized dataflow to improve throughput and hardware utilization in many back-to-back matrix multiplications (MMs) commonly in a transformer model. The T-REX test chip is prototyped in Intel 16-nm FinFET. It achieves 68–<inline-formula xmlns:mml="http://www.w3.org/1998/Math/MathML" xmlns:xlink="http://www.w3.org/1999/xlink"> <tex-math notation="LaTeX">$567~\mu $</tex-math> </inline-formula>s/token and 0.41–<inline-formula xmlns:mml="http://www.w3.org/1998/Math/MathML" xmlns:xlink="http://www.w3.org/1999/xlink"> <tex-math notation="LaTeX">$3.95~\mu $</tex-math> </inline-formula>J/token, demonstrating 1–2 orders of magnitude improvement over the prior work.