Hybrid PIM-oriented Architecture-Dataflow Co-Design for Heterogeneous LLM Inference

Recent studies show that incorporating Processing-in-Memory (PIM) into neural processing units (NPUs) can significantly improve the efficiency of large language model (LLM) inference on edge devices. However, the limited compute capability of DRAM-PIM and the use of fixed dataflow strategies still constrain overall performance gains. To address these limitations, we propose a hybrid PIM architecture that integrates non-volatile spin-orbit torque magnetic random-access memory (SOT-MRAM)-based PIM into existing NPU-PIM systems, fully exploiting the high compute capability and bandwidth of SOT-MRAM together with the large storage capacity of DRAM. On top of this architecture, we further design an adaptive compute scheduling and dataflow optimization framework. Using NSGA-II-based multi-objective dataflow space exploration, our framework identifies Pareto-optimal hardware resource allocations and dataflow configurations under different deployment objectives. Experimental results show that our approach reduces inference latency by up to 8.72× and energy consumption by 11.74× compared to NPUs, and further achieves 6.24× lower latency and 7.76× higher energy efficiency than NPU-PIM baselines.