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\times$ faster and $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.
Lay Summary
Large language models can now read long documents and answer complex questions by reasoning step by step, but this ability is expensive. As the document gets longer and the model produces longer answers, many LLMs require optimizations to efficiently find useful information and generate answers wisely, which involves additional overhead that slows down the generation speedup. Our paper shows that many long-document LLM methods follow the same basic pattern: first, organize the information, then decide what is relevant, then retrieve the useful parts, and finally use them to generate an answer. These steps are quite different from each other. Some look like regular math that GPUs are good at, while others involve searching, filtering, and moving data around, which GPUs handle less efficiently. This observation motivates us to split the work across different kinds of hardware. In our prototype GPU-FPGA system, this design improves LLM inference speed by up to $2.2\times$ and reduces energy use by up to $4.7\times$ compared with a GPU-centered system. These results suggest that future LLM systems should not rely on one type of hardware to do everything. Instead, AI may benefit from systems where different devices handle the parts of the workload they are best suited for. Our work provides a framework for understanding these workloads and offers guidance for building faster, cheaper, and more energy-efficient AI inference systems.
Video
Chat is not available.
Successful Page Load