A 52.03TOPS/W DCIM-Based Accelerator with FlashAttention and Sparsity-Aware Alignment for LLMs

Decoder-based large language models (LLMs) are designed specifically for text generation, consisting primarily of multiple decoder layers that utilize unidirectional attention mechanism [1], as shown in Fig. 1. Due to the enormous number of parameters and the complexity of the attention mechanism, LLMs require substantial computational and memory resources in deployment for resource-constrained edge devices. Algorithms and model compression methods such as KV Cache and pruning have been proposed to reduce the workload, achieving remarkable success [2], [3]. However, high-efficient hardware implementation such as sparsity exploitation [14] still remains unexplored. Meanwhile, Digital Computing-in-Memory (DCIM) can reduce the computation and memory demands by minimizing data movement and enabling more efficient parallel processing directly within the memory, making it suitable for memory-bound LLMs [6–9, 11–13]. Designing DCIM-based LLM accelerators faces the following challenges: 1) The KV Cache size scales linearly with the context length, which induces a mismatch between high storage demand and on-chip memory size, as well as frequent external memory access. 2) Scaling the CIM Macro size can reduce computation complexity from the FlashAttention (FA) mechanism [4], [5], but tiling size in FA introduces imbalanced workload distribution in CIM Macro computation stages, exacerbating inefficiencies in resource utilization. 3) Unstructured sparsity brought by model pruning causes low spatial and temporal utilization of CIM Macro, which compromises hardware efficiency.