PIM-SHERPA: Software Method for On-device LLM Inference by Resolving PIM Memory Attribute and Layout Inconsistencies

TL;DR

This work introduces PIM-SHERPA, a software-only method for efficient on-device LLM inference by resolving PIM memory attribute and layout inconsistencies and proposes effective solutions on product-level PIM-enabled systems.

Abstract

On-device deployments of large language models (LLMs) are rapidly proliferating across mobile and edge platforms. LLM inference comprises a compute-intensive prefill phase and a memory bandwidth-intensive decode phase, and the decode phase has been widely recognized as well-suited to processing-in-memory (PIM) in both academia and industry. However, practical PIM-enabled systems face two obstacles between these phases, a memory attribute inconsistency in which prefill favors placing weights in a cacheable region for reuse whereas decode requires weights in a non-cacheable region to reliably trigger PIM, and a weight layout inconsistency between host-friendly and PIM-aware layouts. To address these problems, we introduce \textit{PIM-SHERPA}, a software-only method for efficient on-device LLM inference by resolving PIM memory attribute and layout inconsistencies. PIM-SHERPA provides two approaches, DRAM double buffering (DDB), which keeps a single PIM-aware weights in the non-cacheable region while prefetching the swizzled weights of the next layer into small cacheable buffers, and online weight rearrangement with swizzled memory copy (OWR), which performs the on-demand swizzled memory copy immediately before GEMM. Compared to a baseline PIM emulation system, PIM-SHERPA achieves approximately 47.8 - 49.7\% memory capacity savings while maintaining comparable performance to the theoretical maximum on the Llama 3.2 model. To the best of our knowledge, this is the first work to identify the memory attribute inconsistency and propose effective solutions on product-level PIM-enabled systems.

Figures (13)

• Figure 1: Model structure and inference process of LLM
• Figure 2: (a) LPDDR-PIM architecture, and (b) GEMV operating sequence of LPDDR-PIM
• Figure 3: Mapping of host-friendly and PIM-aware weights inside DRAM. (a) Example address map. (b) Host and DRAM view of host-friendly weight and (c) PIM-aware weight where each bank holds a complete matrix row. Each rectangular block represents 16 weight elements in FP16/BF16 data format, which are stored in column-major order.
• Figure 4: Estimated DRAM usage before and after weight duplication across representative LLMs. Y-axes use a logarithmic scale.
• Figure 5: Illustration of DRAM double buffering. Two small DRAM buffers are allocated in the cacheable region to prefetch weights of the next layer from the non-cacheable region while GEMM for the current layer executes. This double-buffering mechanism effectively hides the latency of online rearrangement behind compute time. Each decoder layer (Q, K, V, O, FF0 (gate), FF1 (up), and FF2 (down)) corresponds to the Llama 3.2 structure, where H denotes the hidden size.