Accelerating text generation by minimizing memory transfer in attention mechanism

Abstract

Text generation models based on autoregressive transformer models have been instrumental in advancing applications such as chatbot systems and virtual assistants. When the model generates text with multiple batching, the key/value pairs used in the attention mechanism cannot be shared, thus leading to prolonged execution time. As the attention mechanism is memory bounded, off-chip memory accesses should be minimized for faster execution. Although previous methods reduced the off-chip memory accesses regarding unimportant tokens, they fall short in selectively removing the negligible tokens in each instance. Rather, this dissertation estimates the weight using bit chunks of K vectors, effectively removing the memory accesses for low weight tokens and achieving an $12.1x$ pruning ratio without fine-tuning. Additionally, this dissertation present consecutive bit chunk request that prevents the underutilization of Processing Elements (PEs) induced by on-demand DRAM access. Finally, a dedicated hardware equipped with PEs and auxiliary modules is designed, which supports the proposed methods. As a result, it shows $2.6x$ reduced memory accesses, leading to an average $2.3x$ speedup and a $2.4x$ energy efficiency.