MATA: A Memory-Efficient Attention Accelerator for LLMs Exploiting Look-Back KV Cache Pruning

Authors:
Publisher website
Google Scholar
Abstract
Transformer-based Large Language Models (LLMs) have sparked a new wave of AI applications. However, their large computational complexity and memory footprint pose significant challenges for real-world deployment. Although dedicated transformer accelerators have been widely explored, we observe that they are unefficient for decoder-only LLMs that feature autoregressive computations with KV Cache. Our in-depth analysis reveals that DRAM accesses induced by the KV Cache dominate the overall attention process. To address this issue, we propose a Memory-efficient ATtention Accelerator (MATA) for LLMs through algorithm and hardware co-design. Specifically, at the algorithm level, to mitigate the overhead caused by the linear increase of KV Cache, we propose a post-training Look-Back pruning method. It dynamically discards unimportant tokens through a comprehensive scoring scheme, thereby restricting KV Cache to a constant volume. At the hardware level, to identify important tokens with low latency, we design a Slice Top-K (STK) engine that can complete top-k-based sorting with O(N) time complexity. Moreover, we present the Adaptive Dataflow, which adaptively performs different inference phases of LLMs, thus significantly enhancing the PE array utilization. On average, our MATA can achieve speedups of 3.56×, 2.23× and 2.04×, 1.56× energy savings over two state-of-the-art transformer accelerators SpAtten and FACT, respectively.
Funding Information
  • Natural Science Foundation of Jiangsu Province (BK20243051)
  • Beijing Natural Science Foundation (4254088)
  • Science and Technology Major Special Program of Jiangsu (BG2024028)

This publication has 26 references indexed in Scilit: