HEIRS-PS: Hybrid 3-D RRAM- and SRAM-CIM for Pipelined Sparse Multitask Transformer Acceleration

Abstract

Multitask transformers have garnered widespread attention for fast task adaptation at edge. However, in addition to matrix–vector multiplication (MVM) with static and dynamic weights in the data-intensive transformer, task adaptation introduces new challenge of weight updating. Besides, while sparsity can help mitigate the rapid growth in parameter size and computational demand, it requires great hardware flexibility. Compute-in-memory (CIM) is widely explored to overcome the memory bottleneck in neural network acceleration, mainly with single type of memory devices, but struggles to balance the storage density with computing flexibility across multiple tasks. Meanwhile, the conventional pipeline parallel dataflow for high throughput suffers from the underutilization with bubbles caused by variable input sequence length and model sparsity. In this article, HEIRS-PS, a hybrid three-dimension RRAM- and SRAM-CIM architecture, is proposed for the pipelined multitask transformer acceleration with sparsity. With the 3-D integrated RRAM-CIM and SRAM-CIM arrays, a novel hybrid-CIM paradigm is proposed with the collaborative addition to achieve both high storage density and computing flexibility for multitask inference. Moreover, by fully leveraging the hybrid-CIM structure in HEIRS-PS, a new dynamic workload scheduling method is proposed with a layer regrouping method to balance the computing latency across pipeline stages. By eliminating the bubbles, the model inference latency is reduced with higher throughput. Experiment shows that the proposed HEIRS-PS architecture improves the energy efficiency by up to $13.3\times $ and $8.68\times $ , compared with state-of-the-art (SOTA) RRAM-CIM and SRAM-CIM works on Llama-1.3B. Moreover, HEIRS-PS also reduces the computing latency by up to 29.5% with the balanced latency across pipeline stages.