Hardware-Accelerated Large Language Model Inference Engine Using Processing-In-Memory: The PIM-Flux Adaptive Co-Scheduling Architecture

Abstract

Large language models (LLMs) are increasingly limited by the time and energy needed to move model weights and key-value (KV) cache data between DRAM and compute units. This memory bottleneck is especially serious during token-by-token decoding, where matrix-vector operations do not fully use GPU computing resources while memory bandwidth is heavily used (Horowitz, 2014; Wolters et al., 2024). Processing-in-memory (PIM) helps reduce this cost by performing computation close to DRAM banks, reducing the need to move data to separate compute units (Mutlu et al., 2022; Lee et al., 2021). This paper proposes PIM-Flux, a heterogeneous LLM inference engine that shares work dynamically between a GPU/NPU compute plane and a bank-level DRAM-PIM array. A runtime scheduler classifies each transformer sub-layer as compute-bound or memory-bound at each inference step and sends it to the more suitable compute resource. The design also stores the KV cache in mixed INT4/INT8 precision inside PIM memory and uses a lightweight instruction-queue interface without requiring changes to the host processor's instruction set. The design is evaluated analytically using seven published PIM and GPU baselines covering GPT-2, GPT-3-class, and LLaMA-2/Mixtral workloads. The projections indicate a 2.1-4.8x latency reduction and a 3.4-7.9x improvement in energy efficiency compared with an NVIDIA A100/H100-class GPU baseline for batched decoding, while preserving GPU/NPU compute capacity for compute-bound prefill. These results suggest that adaptive PIM-GPU co-scheduling is a practical direction for reducing the cost of LLM inference.