Effective Activation Scratchpad Management in NPUs for LLMs via Inter- and Intra-Decoder Allocation

To accelerate emerging Large Language Models (LLMs), modern Neural Processing Units (NPUs) typically adopt a scratchpad memory for activations as on-chip memory. However, managing an activation scratchpad in NPUs becomes challenge, as the size of LLMs and their context sequence lengths continue to grow. Suboptimal activation memory management may cause low on-chip memory usage and excessive off-chip memory accesses due to large activation data, leading to significant performance degradation during LLM inference on NPUs. This challenge necessitates careful methods to maximize on-chip memory utilization. To address the problem, this work proposes an effective activation scratchpad management in NPUs for LLMs via inter-and intra-decoder allocation. Specifically, inter-decoder optimization facilitates the sharing of on-chip memory across stacked decoders, and intra-decoder one interleaves on-chip memory usage among operations in a decoder. We evaluate our approach with OPT models. The experimental results demonstrate that the proposed approach achieves 56–79 × reduction in memory usage to naïve on-chip memory management methods.