ICML2026

Understand and Accelerate Memory Processing Pipeline for Large Language Model Inference

Zifan He, Rui Ma, Yizhou Sun, Jason Cong

Abstract

Modern large language models (LLMs) increasingly depend on efficient long-context processing and generation mechanisms, including sparse attention, retrieval-augmented generation (RAG), and compressed contextual memory, to solve complex tasks. We show that these optimizations can be unified into a four-stage memory processing pipeline: Prepare Memory, Compute Relevancy, Retrieval, and Apply to Inference. Through systematic profiling, we identify a 22%-97% memory processing overhead in LLM inference and strong computational heterogeneity across stages in memory processing. Motivated by this insight, we argue that heterogeneous systems are well-suited to accelerate memory processing and thus end-to-end inference. We demonstrate this approach on a GPU-FPGA system by offloading sparse, irregular, and memory-bounded operations to FPGAs while retaining compute-intensive operations on GPUs. Evaluated on an AMD MI210 GPU and an Alveo U55C FPGA, our system is up to 2.2×2.2\times faster and 4.7×4.7\times energy reduction across multiple LLM optimizations than the GPU baseline (with similar results on NVIDIA A100), establishing heterogeneous systems as a practical direction for efficient LLM inference and informing future heterogeneous hardware design.