ICML2026
Understand and Accelerate Memory Processing Pipeline for Large Language Model Inference
Zifan He, Rui Ma, Yizhou Sun, Jason Cong
摘要
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 faster and 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.