Optimized Spatial Architecture Mapping Flow for Transformer Accelerators

Abstract

Recent innovations in Transformer-based large language models have significantly advanced the field of general-purpose neural language understanding and generation. With billions of trainable parameters, deployment of these large models relies on high-performance hardware accelerators to efficiently deliver the required computation. Spatial architectures, such as TPUs, offer a promising solution to accelerating computation-intensive workloads. However, the design process for existing spatial architectures is predominantly manual, and it often involves time-consuming redesigns for new applications and new problem dimensions, which greatly limits the development of optimally designed accelerators for Transformer models. To address these challenges, we propose SAMT (Spatial Architecture Mapping for Transformers), a comprehensive framework designed to optimize the dataflow mapping of Transformer inference workloads onto spatial accelerators. We demonstrate the effectiveness of SAMT in improving the performance of spatial accelerators for Transformer models. We propose and leverage the dynamic operator fusion schemes for the Transformer models and co-search the optimal dataflow mapping strategies for spatial accelerators. SAMT significantly reduces inference latency by 12% to 91% and energy consumption by 3% to 23% for evaluated Transformer models compared to traditional spatial accelerator designs among edge, mobile and cloud settings.

Figures (13)

• Figure 1: Lack of fusion exploration in spatial accelerators. Both traditional fixed dataflow accelerators (like ShiDianNaoShiDianNao, NVDLANVDLA, Eyerisseyeriss, and TPUjouppi2023tpu) and flexible dataflow accelerators (like MAERIMAERI and GAMMA gamma) did not consider operator fusion opportunities. *FLAT kao2023flat only considers one fusion while SAMT (ours) considers 64 fusion schemes. Details of accelerators without fusion can be found in Fig. \ref{['fig:acc_type']}.
• Figure 2: Major Computation Steps in a Transformer layer
• Figure 3: (a) Arithmetic intensity with different sequence length $l$ for BERT-Base model (embedding size $d$ = 768 and number of heads $n_h$ = 12) and GPT3-Medium model's prefilling stage ($d$ = 768 and number of heads $n_h$ = 12) (b) The percentage of the memory operations for the operations (for example, $A$ and $S$ in Fig. \ref{['fig:fusion_diagram']}) that scale with the sequence length. (c) Arithmetic intensity for different operators within the same transformer model.
• Figure 4: Abstraction of spatial accelerator architecture
• Figure 5: Example of mapping in MAESTRO_FUSION