TATAA: Programmable Mixed-Precision Transformer Acceleration with a Transformable Arithmetic Architecture

TL;DR

TATAA employs 8-bit integer (int8) arithmetic for quantized linear layer operations through post-training quantization, while it relies on bfloat16 floating-point arithmetic to approximate non-linear layers of a transformer model.

Abstract

Modern transformer-based deep neural networks present unique technical challenges for effective acceleration in real-world applications. Apart from the vast amount of linear operations needed due to their sizes, modern transformer models are increasingly reliance on precise non-linear computations that make traditional low-bitwidth quantization methods and fixed-dataflow matrix accelerators ineffective for end-to-end acceleration. To address this need to accelerate both linear and non-linear operations in a unified and programmable framework, this paper introduces TATAA. TATAA employs 8-bit integer (int8) arithmetic for quantized linear layer operations through post-training quantization, while it relies on bfloat16 floating-point arithmetic to approximate non-linear layers of a transformer model. TATAA hardware features a transformable arithmetic architecture that supports both formats during runtime with minimal overhead, enabling it to switch between a systolic array mode for int8 matrix multiplications and a SIMD mode for vectorized bfloat16 operations. An end-to-end compiler is presented to enable flexible mapping from emerging transformer models to the proposed hardware. Experimental results indicate that our mixed-precision design incurs only 0.14% to 1.16% accuracy drop when compared with the pre-trained single-precision transformer models across a range of vision, language, and generative text applications. Our prototype implementation on the Alveo U280 FPGA currently achieves 2935.2 GOPS throughput on linear layers and a maximum of 189.5 GFLOPS for non-linear operations, outperforming related works by up to 1.45x in end-to-end throughput and 2.29x in DSP efficiency, while achieving 2.19x higher power efficiency than modern NVIDIA RTX4090 GPU.

Figures (18)

• Figure 1: Illustration of a typical Transformer block, and how TATAA maps different operations in Transformers including linear MatMul and a variety of non-linear functions into transformable architecture, based on a general end-to-end framework design.
• Figure 2: TATAA hardware architecture and dual-mode processing unit (DMPU) design. The TATAA core can be configured for two kinds of workload during runtime, to support both linear MatMul and non-linear functions.
• Figure 3: Processing element (PE) design in the proposed DMPUs when deployed on FPGA devices. The converter from 2's complement to signed digit (SD) logic that locates in PE-S3's bottom logic is not depicted in this figure since it is similar to the SD to 2's complement circuit in PE-S1.
• Figure 4: Execution dataflow in int8MatMul mode, in which all the PEs are connected as a whole systolic array and deploy output stationary dataflow.
• Figure 5: bfloat16 mode dataflow and how to reuse the integer processing units. The Top-Logic, MUL & ADD, and Bottom-Logic processing are depicted in specific colors in this figure.