Venue: NAACL 2022

# Easy and Efficient Transformer: Scalable Inference Solution For Large NLP Model

Gongzheng Li\({}^{1}\), Yadong Xi\({}^{1}\), Jingzhen Ding\({}^{1}\), Duan Wang\({}^{1}\), Ziyang Luo\({}^{2}\),

**Rongsheng Zhang\({}^{1}\),Bai Liu\({}^{1}\), Changjie Fan\({}^{1}\), Xiaoxi Mao\({}^{1}\), Zeng Zhao\({}^{1}\)\({}^{\dagger}\)**

\({}^{1}\) Fuxi AI Lab, NetEase Inc., Hangzhou, China

\({}^{2}\) Department of Computer Science, Hong Kong Baptist University, Hong Kong SAR, China {ligongzheng,xiyadong,maoxia;hzzhaozeng}@corp.netease.com

###### Abstract

Recently, large-scale transformer-based models have been proven to be effective over various tasks across many domains. Nevertheless, applying them in industrial production requires tedious and heavy works to reduce inference costs. To fill such a gap, we introduce a scalable inference solution: **Easy and Efficient Transformer (EET)**, including a series of transformer inference optimization at the algorithm and implementation levels. First, we design highly optimized kernels for long inputs and large hidden sizes. Second, we propose a flexible CUDA memory manager to reduce the memory footprint when deploying a large model. Compared with the state-of-the-art transformer inference library (Faster Transformer v4.0), EET can achieve an average of 1.40-4.20x speedup on the transformer decoder layer with an A100 GPU.

+
Footnote †: Corresponding Author

+
Footnote †: Corresponding Author

+
Footnote †: Corresponding Author

## 1 Introduction

In recent years, transformer-based models have achieved impressive performance across variant domains, such as natural language processing [35, 87, 69, 1], computer vision [16, 17] and speech processing [14, 15]. The scaling law proposed by Kaplan et al. [15] indicates that the validation PPL of a neural language model scales as a power-law with model sizes, dataset sizes, and the amount of training computation. Such law is corroborated empirically by many following works [17, 18].

However, the mega-sized models are notoriously expensive for deployment in the industry. For example, GPT-2 medium model (700M parameters [23]) spends up to 10s to generate 512 tokens given a prompt with the length of 512 on an RTX 2080ti GPU, which is not allowed in the industrial application. Multiple approaches have been proposed to solve such problems, including knowledge distillation [14, 15], model pruning [16], and quantization [20]. Apart from these works, much attention has also been paid to optimizing CUDA implementation of a transformer layer for better hardware utilization. Previous works (e.g.: TensorRT [21], LightSeq [22] and Faster Transformer (FT) [23]) have implemented many optimization techniques, including kernels fusion, gemm optimization, quantization, etc. However, these works still have several limitations. TensorRT only contains the multi-head attention(MHA) operation, lacking a complete transformer model. LightSeq cannot support the model hidden size and input sequence length over 1024. FT contains some performance flaws which need to be improved.

In this paper, we propose a novel transformer inference acceleration library, **Easy and Efficient Transformer (EET)** . First, we implement custom CUDA kernels to avoid explicit matrix addition of attention and padding masks with attention weights. As a result, the attention mask matrix is no longer required, while FT spends overhead to initialize an attention mask on the CPU and push it to CUDA. In addition, compared with FT, padding masks are no longer needed in computation, leading to additional performance improvement. Second, we propose a new method, _thread block folding_, to extend all kernels to support a larger model size up to 12288 and a longer sequence up to 4096. For FT, it directly assigns the thread number in a CUDA block, which may hurt the parallel efficiency. Finally, we design a dynamic CUDA memory management mechanism to reduce the CUDA memory occupation for the same model size, while FT needs to manually allocate memory usage.

We have conducted comprehensive experiments

to compare EET with Fairseq,1 LightSeq and FT. In our experiments, EET achieves about 4.48-20.27x and 4.30-27.43x speedup over Fairseq on 2080ti and A100 respectively. When we set the model size to 768 and 1024 on 2080Ti, EET makes 0.82-2.46x speedup over LightSeq. Compared to FT(v3.1), EET achieves about 1.21-6.30x and 1.62-8.16x speedup on 2080ti and A100 respectively. Compared to FT(v4.0), EET achieves about 1.40-4.20x speedup on A100. The remarkable experimental results corroborate the effectiveness of our EET.

Footnote 1: [https://github.com/pytorch/fairseq](https://github.com/pytorch/fairseq)

## 2 Custom Kernels

FT NVIDIA (2021a) has implemented highly optimized CUDA kernels for transformer inference. To make further optimization, we design our custom kernels with the considerations below:

\(\bullet\) Because padding tokens do not affect the final results, preventing padding tokens from participating in MHA instead of simply applying padding masks can significantly reduce the computational overhead.

\(\bullet\) Although an attention mask is essential for MHA in text generation, constructing a mask that varies with the input length is time-consuming.

\(\bullet\) The hidden sizes and input lengths of the large-scale pre-trained models can easily exceed 1024. It is necessary to extend these kernels to support large hidden sizes and input lengths elegantly and efficiently.

To remove previously mentioned masks in computation, we redesign the kernels and call the mechanism _mask fusion_. To extend all the kernels to support the model size or sequence length greater than 1024, we improve the CUDA thread structure and call the method as _thread block folding_. Next, we describe these two methods in detail.

### Mask Fusion

The attention mask indicates the attention boundary for each token to prevent the attention from looking forward. The padding mask indicates where the padding tokens are. Thus they both characterize the position information of the tokens in a sequence. Meanwhile, each CUDA thread also has a unique positional index. So we can map each token in the MHA to a thread or block in the CUDA kernels. The function of the attention mask is achieved by comparing whether the CUDA position of the query token being processed is larger than the CUDA position of the key token. The function of the padding mask is achieved by starting the valid calculations from the padding offset when sequentially processing each token. Therefore, we transform the mask computation to logical operation with CUDA thread index comparison. Thus there is no need to store any explicit functional parameters of the masks and the computation overhead of masking operation is saved. The algorithm pseudo-code is shown in Algorithm 1.

```
Input:\(qk,paddingLen,seqLen,batch,headNum\) Output: the attention weights back to \(qk\) CUDA Initialize \(grid\leftarrow(batch*headNum)\) CUDA Initialize \(block\leftarrow(seqLen)\)\(batchId+blockIdx.r/headNorm\)\(padLen\gets paddingLen[batchId]\)\(qQOffset\leftarrow blockIdx.*seqLen*seqLen\)\(qKOffset\gets qKOffset+paddLen*seqLen\)\(s\gets padLen\)\(\triangleright\) start at first non-pad \(e\gets seqLen\)\(\triangleright\) end at last token \(reduceMax\leftarrow-inf\)\(reduceSum\gets 0\)for\(i=s\)to\(c\)do \(position\gets qkOffset+threadIdx.x\)\(data\gets qkPosition\) \(u\gets pddLen\)\(\triangleright\) upper boundary \(l\gets i\) if\(l<threadIdx.z<u\)then\(reduceMax\gets blockReduceMax(data)\)\(reduceSum\gets blockReduceSum(data)\)\(data\gets softmax(reduceMax,reduceSum)\)endif\(qk[position]\gets data\) endfor
```

**Algorithm 1** MHA with _mask fusion_ 

the model. We call this method _thread block folding_, which allows us to extend any kernel to any model size and any sequence length with minimum changes and non-degraded performance. For instance, assuming the model size is 1280, we fold it once and create two half-size blocks, then the data can be assigned into two separate blocks with 640 threads in each.

We introduce a folding coefficient to characterize the number of folding. Given the model size \(h\), the folding coefficient \(t\) and the number of threads \(n\) in one block is defined as:

\[t=2^{\lceil\frac{h}{1024}\rceil-1};~{}~{}~{}n=\frac{h}{2^{t}}\]

As for simplicity, _thread block folding_ only adds a new dimension for the block, which slightly impacts the basic CUDA thread grid structure. As for efficiency, the minimum thread number is 512 when the model size or input length is larger than 1024 and makes full use of thread parallelism. The sequence expansion process is similar to the model expansion process. Finally, we support the model size no larger than 16384 and sequence length no longer than 4096.

## 3 Dynamic Memory Manager

The inference is much more sensitive to latency compared to training. As a result, model parallelism Shoeybi et al. (2020) and pipeline parallelism Huang et al. (2019) are undesirable for inference. Their communication overhead introduced by tensor slicing or layer split is significant even with the support of NVLink and GPUDirect. To reduce the latency and hardware requirements for online service, minimizing the memory footprint is of exceptional value when loading very large models. Thus we propose a dynamic memory management strategy for this issue.

Except for the model weights, the memory footprint includes the caches and the buffers. It is hard to reduce the memory footprint of weights because they are inherent to the model. Similarly, The _K/V_ caches for MHA are also hard to compress because they are pre-allocated to avoid runtime memory requests, the size of which depends on the model size, maximum batch size, and maximum sequence length. Whereas the activation cache and the buffers used to store the operator's results are compressible. Hence our dynamic memory management strategy mainly focuses on the activation caches and the operation result buffers.

### Cache Reuse

The caches include _K/V_ caches and activation caches. In incremental decoding, the keys and the values for every step are stored for the next step's attention computation. The maximum size of _K/V_ caches is predictable because we can determine the maximum batch size and decoding steps at the start of the running instance. We allocate the maximum required memory in advance to reduce the forward latency, avoiding malloc overhead and memory corruption.

Different from _K/V_ caches, the activation results are useless after we have calculated and passed them to the next layer. The memory for these activations can be reused across different layers and different operators. We could reuse the activation caches in the following cases.

\(\bullet\) The embedding operator shares the cache with the feed-forward operator and the final output. Yet the attention operator holds another cache because of the residual connection.

\(\bullet\) The cache for input sequences can be reused by the decoded tokens. The maximum size is determined by the maximum input length.

\(\bullet\) The cache can be reused across different layers.

We use the following notations: \(b\), the maximum batch size; \(s\), the maximum sequence length; \(p\), the maximum prompt length; \(h\), the hidden units; \(l\), the layer number. The total activation cache size is:

Figure 1: The schematic diagram of _thread block folding_.

 

The total K/V cache size is :

\[2*b*h*s*l\]

### Buffer Reuse

The continuous CUDA kernels are not always fused, especially when it comes to Cublas GEMM calls. So we need some buffers to store the returns for those non-fused kernels. Managing the buffers manually like FT is complicated and inefficient. We develop a dynamic buffer manager to avoid the tedium of manual design and achieve a highly efficient memory allocation.

We maintain a list of buffers and use different strategies within and across modules to improve memory utilization. When within modules, we reuse the buffer only when the request size is identical to an idle buffer in the list, preventing memory fragmentation. When across modules, we reuse the buffer when the request size is smaller than any idle buffer in the list, avoiding duplicated malloc. The decision process is demonstrated in Figure 2. In our design, the developer only needs to request a buffer of a specified size and mark it as idle when it is useless, without concerning how to reuse memory exactly. The total buffer size is:

\[b*p*(6*h+n*p)\]

where \(b\) is the batch size, \(p\) is the input length, \(h\) is the hidden size, and \(n\) is the head number.

## 4 Experiments

During inference, many factors can affect the actual performance, including model size, prompt length, sequence length, padding ratio in a batch, and hardware feature. Completely traversing all combinations requires a huge amount of works. Because the dataset has no effect on the experiment results, we adopt the fake inputs for convenience. To compare fairly and reduce our works, we define some typical experiment settings. If there is no special instruction, the experiment is conducted based on Configuration A in Table 1. Fairseq is an intuitional baseline because it is implemented using pure PyTorch code.

### Speedup for GPT-2 Layer with Different Sequence Lengths

We first apply EET over GPT-2 on NVIDIA 2080ti and A100. Figure 3 and 4 reveal that EET achieves about 4.48-20.27x and 4.30-27.43x speedup than Fairseq and about 1.21-6.30x and 1.62-8.16x speedup than FT(v3.1), on 2080ti and A100 respectively. For Fairseq and FT(v3.1), the incremental decoding processes the input tokens one by one, while EET improves the tokens parallelism by processing input tokens all at once. As a result, the speedup grows with the increase of the input length.

Figure 3: Inference speedup of EET with different prompt lengths on 2080ti compared to Fairseq and FT(v3.1).

\begin{table}
\begin{tabular}{l c c} \hline \hline  & Config A & Config B \\ \hline Batch size & 4 & 8 \\ Model size & 1024 & 2048 \\ Max prompt & 1024 & 1024 \\ Max sequence & 1024 & 1024 \\ Datatype & fp16 & fpf16 \\ \hline \hline \end{tabular}
\end{table}
Table 1: Configuration A and B

Figure 2: The schematic diagram of buffer decision strategy.

 

The recent version of FT(v4.0) also introduces the parallel decoding of the input sequences for text generation as we did, so the performance of EET and FT(v4.0) is getting closer with the input length increasing. However, EET still has some performance advantages, which are attributed to our operation kernel optimization. Figure 5 shows that EET achieves about 1.40-2.54x speedup compared to FT(v4.0) with the configuration B in Table 1.

When processing a batch of inputs, the length of them may be uneven. The FT(v4.0) uses the minimum length of the prompts for full decoding, while the EET uses the maximum length. For example, if there is a batch containing sequences of different length like [5; 2; 4; 10], the final prompt length used for parallelism is 2 in the FT. In contrast, it is 10 in the EET. Figure 6 shows that we make 2.74-4.42x speedup with the prompt fixed to 512 and other configurations keeping the same as the configuration B in Table 1.

Unlike Fairseq and FT(v4.0), LightSeq only supports model sizes that are smaller than 1024, we also make a comparison here as a supplement. Figure 7 shows that we make 0.82-2.46x speedup when we set the model size to 768 and 1024.

### Speedup for Transformer Decoder Layer with Different Model Sizes

To prove the scalability of our EET, we evaluate the performance on different model sizes with configuration C in Table 2. Figure 8 and Figure 9 reveal that EET achieves about 2.25-7.50x speedup than Fairseq and about 1.71-4.61x speedup than FT(v4.0). The acceleration ratio decreases as the model size increases due to the increased ratio of matrix multiplication in the inference. Nevertheless, with the help of _thread block folding_, EET can still deliver significant speedup with very large model sizes, compared to Fairseq and FT(v4.0).

 

half of the total tokens. The result is shown in Figure 10. Deprecation of the padding masks with the _mask fusion_ trick brings in 0.99-1.27x speedup. As for Bert, its hidden size is fixed to 1024 and it has no sequence mask, which kicks off the optimization of thread-block folding and sequence mask fusion, then the speedup is not as significant as GPT2.

### Memory distribution

Given the batch size 16, the maximum sequence length 1024, the vocab size 13672, we plot the memory distribution of the hidden size of 1024 and 4096 with layer numbers 24 and 40 respectively, as shown in Figure 11. Regardless of the hidden size, we can find that model weights and \(K/V\) caches occupy most memory. The activation caches and the buffers only take up a small part, which shows the effectiveness of our dynamic memory management strategy.

Given the batch size 4, the maximum sequence length 1024, we plot the memory occupancy of different model parameter sizes, see Figure 12. Compared with the 10 billion of PyTorch's maximum model parameter sizes, it is up to 18 billion for our EET, which proves that we can place much larger models onto one GPU, thus avoiding unnecessary waste of GPU resources and inter-GPU communication overhead on multiple cards.

## 5 Conclusion

This paper comprehensively describes a series of optimization techniques for transformer inference acceleration exploiting both algorithmic and GPU hardware features. These techniques are packed into the EET, a library dedicated to inference acceleration for large transformer-based models and long input lengths. EET has a 1.40-4.42x speedup for the GPT-2 layer and a 0.99-1.27x speedup for the Bert layer compared to the state-of-the-art transformer inference library FT. To make EET easier to

Figure 11: Memory distribution for 1024/4096 hidden sizes.

Figure 8: Speedup with different model sizes on 2080ti and A100 compared to Fairseq.

Figure 12: Memory occupancy for different model sizes.

Figure 10: Performance speedup for Bert layer on 2080ti compared to FT(v4.0).

Figure 9: Speedup with different hidden sizes compared to FT(v4.0).

 

apply to a specific task, we provide operation level and model level API, meanwhile integrating web service with dynamic batching. We will continue to improve and keep it up-to-date.

AcknowledgmentsWe would like to thank all the users of EET and the anonymous reviewers for their excellent feedback. This work is supported by the Key Research and Development Program of Zhejiang Province (No. 2022C01011).

## References

* A. Baevski, W. Hsu, A. Conneau, and M. Auli (2021)Unsupervised speech recognition. arXiv preprint arXiv:2105.11084. Cited by: SS3.1.
* A. Baevski, H. Zhou, A. Mohamed, and M. Auli (2020)Wav2vec 2.0: a framework for self-supervised learning of speech representations. arXiv preprint arXiv:2006.11477. Cited by: SS1.
* T. B. Brown, B. Mann, N. Ryder, M. Subbiah, J. Kaplan, P. Dhariwal, Arvind Neelakantan, P. Shyam, G. Sastry, A. Askell, S. Agarwal, A. Herbert-Voss, G. Krueger, T. Henighan, R. Child, A. Ramesh, D. M. Ziegler, J. Wu, C. Winter, C. Hesse, M. Chen, E. Sigler, M. Litwin, S. Gray, B. Chess, J. Clark, C. Berner, S. McCandlish, A. Radford, I. Sutskever, and D. Amodei (2020)Language models are few-shot learners. External Links: Link Cited by: SS1.
* J. Devlin, M. Chang, K. Lee, and K. Toutanova (2019)BERT: pre-training of deep bidirectional transformers for language understanding. External Links: Link Cited by: SS1.
* A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, J. Uszkoreit, and N. Houlsby (2020)An image is worth 16x16 words: transformers for image recognition at scale. External Links: Link Cited by: SS1.
* G. Hinton, O. Vinyals, and J. Dean (2015)Distilling the knowledge in a neural network. External Links: Link Cited by: SS1.
* Y. Huang, Y. Cheng, A. Bapna, O. Firat, M. Xu Chen, D. Chen, H. Lee, J. Ngiam, Q. V. Le, Y. Wu, and Z. Chen (2019)Gpipe: efficient training of giant neural networks using pipeline parallelism. External Links: Link Cited by: SS1.
* Y. Jiang, S. Chang, and Z. Wang (2021)Transgan: two transformers can make one strong gan. External Links: Link Cited by: SS1.
* X. Zhao, Y. Yin, L. Shang, X. Jiang, X. Chen, L. Li, F. Wang, and Q. Liu (2020)Tinybert: distilling bert for natural language understanding. External Links: Link Cited by: SS1.
* J. Kaplan, S. McCandlish, T. Henighan, T. B. Brown, B. Chess, R. Child, S. Gray, A. Radford, J. Wu, and D. Amodei (2020)Scaling laws for neural language models. External Links: Link Cited by: SS1.
* NVIDIA (2021a)"faster transformer". External Links: Link Cited by: SS1.
* A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, I. Sutskever, et al. (2019)Language models are unsupervised multitask learners. OpenAI blog1 (8), pp. 9. External Links: Link Cited by: SS1.
* C. Raffel, N. Shazeer, A. Roberts, K. Lee, S. Narang, M. Matena, Y. Zhou, W. Li, and P. J. Liu (2020)Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of Machine Learning Research21 (140), pp. 1-67. External Links: ISSN 1574-469X, Document, Link Cited by: SS1.
* S. Shen, Z. Dong, J. Ye, L. Ma, Z. Yao, A. Gholami, M. W. Mahoney, and K. Keutzer (2019)Q-bert: hessian based ultra low precision quantization of bert. External Links: Link Cited by: SS1.
* M. Shoeybi, M. Patwary, R. Puri, P. LeGresley, J. Casper, and B. Catanzaro (2020)Megatron-lm: training multi-billion parameter language models using model parallelism. External Links: Link Cited by: SS1.
* A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin (2017)Attention is all you need. External Links: Link Cited by: SS1.
* E. Voita, D. Talbot, F. Moiseev, R. Sennrich, and I. Titov (2019)Analyzing multi-head self-attention: specialized heads do the heavy lifting, the rest can be pruned. External Links: Link Cited by: SS1.
* X. Wang, Y. Xiong, Y. Wei, M. Wang, and L. Li (2021)Lightseq: a high performance inference library for transformers. External Links: Link Cited by: SS1.