Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

SOLE: Hardware-Software Co-design of Softmax and LayerNorm for Efficient Transformer Inference

Wenxun Wang, Shuchang Zhou, Wenyu Sun, Peiqin Sun, Yongpan Liu

TL;DR

SOLE addresses bottlenecks in transformer inference by targeting Softmax and LayerNorm with a hardware/software co-design. It introduces E2Softmax and AILayerNorm, leveraging $\mathrm{Softmax}(X_i) = \frac{\exp(X_i - X_{max})}{\sum_j \exp(X_j - X_{max})}$ and $\mathrm{LayerNorm}(X_i) = \frac{X_i - \mu}{\sigma} \cdot \gamma + \beta$ alongside log2 quantization, log-based division, dynamic compression, and Power-of-Two Factor quantization. The hardware realization includes E2Softmax Unit and AILayerNorm Unit, achieving negligible accuracy loss without retraining and substantial speedups (e.g., up to 36.2x Softmax and 61.3x LayerNorm over GPU) as well as energy and area efficiency gains over prior designs. These results demonstrate a viable path for memory-aware non-linear operation co-design in transformer inference and easy integration with existing quantization workflows.

Abstract

Transformers have shown remarkable performance in both natural language processing (NLP) and computer vision (CV) tasks. However, their real-time inference speed and efficiency are limited due to the inefficiency in Softmax and Layer Normalization (LayerNorm). Previous works based on function approximation suffer from inefficient implementation as they place emphasis on computation while disregarding memory overhead concerns. Moreover, such methods rely on retraining to compensate for approximation error which can be costly and inconvenient. In this paper, we present SOLE, a hardware-software co-design for Softmax and LayerNorm which is composed of E2Softmax and AILayerNorm. E2Softmax utilizes log2 quantization of exponent function and log-based division to approximate Softmax while AILayerNorm adopts low-precision statistic calculation. Compared with state-of-the-art designs, we achieve both low-precision calculation and low bit-width storage on Softmax and LayerNorm. Experiments show that SOLE maintains inference accuracy without retraining while offering orders of magnitude speedup and energy savings over GPU, achieving 3.04x, 3.86x energy-efficiency improvements and 2.82x, 3.32x area-efficiency improvements over prior state-of-the-art custom hardware for Softmax and LayerNorm, respectively.

SOLE: Hardware-Software Co-design of Softmax and LayerNorm for Efficient Transformer Inference

TL;DR

SOLE addresses bottlenecks in transformer inference by targeting Softmax and LayerNorm with a hardware/software co-design. It introduces E2Softmax and AILayerNorm, leveraging and alongside log2 quantization, log-based division, dynamic compression, and Power-of-Two Factor quantization. The hardware realization includes E2Softmax Unit and AILayerNorm Unit, achieving negligible accuracy loss without retraining and substantial speedups (e.g., up to 36.2x Softmax and 61.3x LayerNorm over GPU) as well as energy and area efficiency gains over prior designs. These results demonstrate a viable path for memory-aware non-linear operation co-design in transformer inference and easy integration with existing quantization workflows.

Abstract

Transformers have shown remarkable performance in both natural language processing (NLP) and computer vision (CV) tasks. However, their real-time inference speed and efficiency are limited due to the inefficiency in Softmax and Layer Normalization (LayerNorm). Previous works based on function approximation suffer from inefficient implementation as they place emphasis on computation while disregarding memory overhead concerns. Moreover, such methods rely on retraining to compensate for approximation error which can be costly and inconvenient. In this paper, we present SOLE, a hardware-software co-design for Softmax and LayerNorm which is composed of E2Softmax and AILayerNorm. E2Softmax utilizes log2 quantization of exponent function and log-based division to approximate Softmax while AILayerNorm adopts low-precision statistic calculation. Compared with state-of-the-art designs, we achieve both low-precision calculation and low bit-width storage on Softmax and LayerNorm. Experiments show that SOLE maintains inference accuracy without retraining while offering orders of magnitude speedup and energy savings over GPU, achieving 3.04x, 3.86x energy-efficiency improvements and 2.82x, 3.32x area-efficiency improvements over prior state-of-the-art custom hardware for Softmax and LayerNorm, respectively.
Paper Structure (17 sections, 17 equations, 6 figures, 3 tables, 2 algorithms)

This paper contains 17 sections, 17 equations, 6 figures, 3 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related work
  3. Prominent Non-Linear Operations in Transformers
  4. Accelerating Non-Linear Operations
  5. SOLE Software
  6. Preliminary
  7. E2Softmax
  8. AILayerNorm
  9. SOLE Hardware
  10. E2Softmax Unit
  11. AILayerNorm Unit
  12. Evaluation
  13. Experiment Setting
  14. Algorithm Performance
  15. Speedup
  16. ...and 2 more sections

Figures (6)

  • Figure 1: (a) Latency Breakdown of Deit-Tiny with the image size of 448 on a 2080Ti GPU. (b) Difference between MatMul operations with Softmax/LayerNorm operations.
  • Figure 2: Overview of SOLE
  • Figure 3: Distribution of $\mathbf{Exp}(X_i-X_{max})$ in logarithm.
  • Figure 4: Hardware design of E2Softmax Unit
  • Figure 5: Hardware design of AILayerNorm Unit
  • ...and 1 more figures