PaperlayerView exact revision on arXiv

Source-linked AI summary

Understand and Accelerate Memory Processing Pipeline for Large Language Model Inference

Zifan He, Rui Ma, Yizhou Sun, Jason Cong

arXiv:2603.29002v3cs.DCcs.AI

TL;DR

Long-context LLM optimizations lack a systematic understanding of their memory-processing costs and computational characteristics. This paper unifies them into a four-step pipeline and demonstrates a GPU-FPGA design that achieves up to 2.2× speedup and 4.7× energy reduction.

  • Problem

    Prior work treats LLM inference optimizations as isolated techniques, limiting systematic understanding of their computational characteristics and hardware-efficiency implications.

  • Method

    The paper unifies LLM inference optimizations into a four-step memory-processing pipeline and maps heterogeneous operations across GPUs and FPGAs.

  • Results

    1.04–2.2× speedup and 1.11–4.7× energy reduction are achieved across LLM inference optimizations with the GPU-FPGA system.

  • Takeaways & Limitations

    GPU-FPGA heterogeneous systems are a practical direction for accelerating LLM memory processing and informing future heterogeneous hardware designs.

  • Takeaways & Limitations

    The heterogeneous system does not deploy LaCT because its memory-processing bottleneck is dominated by compute-bound operations.

Abstract

from arXiv · show

Modern large language models (LLMs) increasingly depends on efficient long-context processing and generation mechanisms, including sparse attention, retrieval-augmented generation (RAG), and compressed contextual memory, to support complex reasoning. We show that these optimizations can be unified into a four-step 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 heterogeneity in its computational characteristics. Motivated by this insight, we argue that \textbf{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 achieves up to $4.7\times$ less energy across multiple LLM inference optimizations than the GPU baseline (similar results hold on NVIDIA A100). These results establish heterogeneous systems as a practical direction for efficient LLM memory processing and inform future heterogeneous hardware design.

1. Introduction

The introduction frames long-context LLM inference as a memory-processing problem and unifies existing optimizations into a four-step pipeline. It argues that heterogeneous computation and GPU-FPGA mapping can accelerate this pipeline, improving inference speed and energy efficiency.

  • Motivation: 128k to 1 million tokens per request make efficient long-input memorization and processing increasingly important, while KV caches impose substantial hardware costs.This motivates systematic acceleration of memory processing for long-context inference.
  • Problem: Existing LLM inference optimizations are largely treated as isolated techniques, limiting systematic understanding of their computational characteristics and hardware-efficiency implications.The paper presents three claims to establish a unified basis for further acceleration.
  • Claim 1: Modern LLM inference optimizations can be unified as a four-step memory-processing pipeline: Prepare Memory, Compute Relevancy, Retrieval, and Apply to Inference.The pipeline respectively preprocesses memory, scores importance, selects information, and applies the result during inference.
  • Claim 2: Memory-processing computations are heterogeneous in arithmetic intensity, resource utilization, memory access patterns, and data dependencies.For example, compressed key-embedding generation can be regular and compute-intensive, whereas score computation has different characteristics.
  • Claim 3: 1.5 ∼5.7× speedup for sparse attention produces up to 1.49× faster end-to-end inference, while 5.2 ∼7.7× speedup for RAG produces up to 2.2× end-to-end speedup.The heterogeneous system maps memory processing across off-the-shelf devices, using FPGAs for suitable operations and GPUs for compute-intensive work.

2. Background

Long-context LLM inference optimizations change how memory is accessed and managed, but share a common memory-processing sequence. This background covers sparse attention, retrieval-augmented generation, compressed contextual memory, and test-time training.

  • Sparse Attention: Sparse attention reduces long-context attention cost by selectively attending to a subset of past tokens.It addresses quadratic prefill and linear decoding complexity in standard transformers.
  • Retrieval-Augmented Generation: RAG retrieves relevant documents from an external corpus instead of storing static knowledge as per-layer KV-cache vectors.Dynamic systems include FLARE, which triggers retrieval when confidence drops, and DRAGIN, which uses attention statistics to detect uncertainty.
  • Compressed Contextual Memory: Compressed contextual memory stores information as embeddings or summarized text, enabling extremely long inputs and far-distance retrieval in latent space.Methods trade computation for runtime data efficiency by synthesizing textual memories or compressing segments into latent embeddings.
  • Test-time Training: Test-time training treats model parameters as internal memory and adapts them during inference through recurrent updates alternating backpropagation and generation.LaCT extends this approach with batched updates to improve GPU utilization.
  • Common Memory Processing: These methods share a four-step memory-processing pipeline: Prepare Memory, Compute Relevancy, Retrieval, and Apply to Inference.The steps respectively structure raw memory, score relevance, extract relevant content, and integrate it with the query for subsequent LLM operations.

3. Memory Processing in LLM Inference

LLM inference memory processing can be unified as a four-stage pipeline that transforms generated memory into inference outputs. Profiling shows that memory processing becomes a major, method-dependent latency bottleneck as memory size grows.

  • Memory-processing definition: Memory processing transforms generated memory and the current input into intermediate outputs used to produce the final language-model output.The model generates memory with g, processes it with f, and uses the resulting intermediate output to generate yt.
  • Four-stage pipeline: The pipeline comprises Prepare Memory, Compute Relevancy, Retrieval, and Apply to Inference.These stages respectively create memory indexes, compute scores, select or refine memory, and incorporate retrieved memory into subsequent computations.
  • Pipeline flexibility: Some methods skip pipeline stages, while the unified abstraction enables stage bypass, kernel reuse, and pipeline scheduling.MemAgent skips relevancy computation, TTT skips retrieval, and unnecessary stages add no overhead when data bypasses them.
  • Latency overhead: 22%–81% of decoding latency is spent on memory processing for 1M-token sparse-attention sequences, versus 1%–11% for 4K-token sequences.For RAG, memory processing reaches 40%–61% when processing 20M documents.
  • Method heterogeneity: Memory-processing costs vary substantially by method: sparse attention and RAG emphasize Compute Relevancy and Retrieval, while MemAgent reaches 97% latency in Prepare Memory.Because memory processing is a primary inference bottleneck, accelerating it can improve end-to-end latency and energy efficiency.

4. Computational Heterogeneity

The memory processing pipeline exhibits substantial computational heterogeneity across LLM inference optimizations, spanning memory- versus compute-bound operations and regular versus irregular access patterns. This heterogeneity motivates mapping different pipeline operations onto heterogeneous systems.

  • Computational characterization: The analysis characterizes each pipeline step by arithmetic intensity, data-access patterns, and dependencies, with detailed latency breakdowns provided in Appendix B.Arithmetic intensity distinguishes memory-bound from compute-bound operations.
  • Sparse Attention & RAG: Sparse attention and RAG have memory-bound Compute Relevancy and Retrieval steps involving irregular accesses and dependencies in BM25 scoring, top-k selection, and max reduction.BM25 accesses token-frequency histograms nondeterministically, while top-k maintains running maximum scores with dependencies.
  • Synthesized Memory: MemAgent combines memory-bound decoding during Prepare Memory with compute-bound prefilling during Apply to Inference.Its memory processing is effectively a sequence of LLM inferences.
  • Memory as Context: Memory as Context resembles sparse attention and RAG but adds query-generation calculations that are independent of the model forward pass and can be parallelized and fused.The extra calculations occur during Compute Relevancy.
  • TTT: LaCT is not deployed on the heterogeneous system because compute-bound forward and backward passes dominate memory-processing latency despite a more memory-intensive loss function.The loss function corresponds to Compute Relevancy.
  • Implication: This heterogeneity motivates mapping LLM memory-processing operations onto a heterogeneous system.The mapping is discussed in the following section.

5. GPU-FPGA Heterogeneous System

The GPU-FPGA system maps regular, compute-bound work to the GPU and irregular, data-dependent, memory-bound work to the FPGA, while prioritizing end-to-end communication costs. It provides distinct mappings for sparse attention and RAG, synthesized memory, and Memory as Context, supported by streaming FPGA kernels and reusable kernel composition.

  • Mapping Criteria: The system maps compute-bound, regular operations to the GPU and irregular, data-dependent, memory-bound operations to the FPGA, while prioritizing cross-device communication overhead over kernel speedup.For example, KV-cache extraction remains separate from FPGA top-k selection because PCIe overhead outweighs the fusion benefit.
  • General Setup: Sparse attention and RAG keep Prepare Memory and Apply Memory on the GPU while the FPGA executes a fused Compute Relevancy and Retrieval kernel.The GPU stores memory, target, and output data; the FPGA stores the query and processed memory.
  • Synthesized Memory: Synthesized Memory assigns LLM decoding to the FPGA and LLM prefilling to the GPU, exchanging KV caches and synthesized-memory token IDs for concatenation.For each segment, the GPU delivers the KV cache to the FPGA, which returns token IDs to the GPU.
  • Memory as Context: Memory as Context stores memory only on the FPGA and transfers retrieved embeddings to the GPU, enabling greater FPGA kernel fusion and sparing GPU memory for model storage.The mapping uses recurrent-loop communication for input segments, with GPU preparation requiring retrieved memory and the next input segment.
  • FPGA Kernel Architecture: The FPGA General Setup kernel uses streaming dataflow modules for Compute Relevancy and Retrieval, combining inner-product scoring, reduction, and a running top-k retriever.Reusable standalone or fused kernels form a library that users can recombine through the GPU-FPGA communication API.

6. Evaluation

Evaluation on an AMD MI210 GPU and Alveo U55C FPGA shows that heterogeneous acceleration improves memory processing and end-to-end inference through on-chip memory, pipelining, and flexible HBM use. Benefits vary across workloads, energy efficiency, and batch sizes.

  • Experimental Setup: The system uses an AMD Alveo U55C FPGA with an AMD Instinct MI210 GPU and AMD EPYC 7v13 host, comparing against the same GPU baseline.Estimations based on profiling address generalizability to systems pairing an NVIDIA A100 with U55C.
  • Experimental Setup: Evaluation reports end-to-end latency and memory-processing fractions across workload-specific datasets, sequence lengths, and per-token or total-request latency measures.Sparse attention uses varying past-token lengths; RAG, MemAgent, and Memory as Context use prefill-and-decode request latency under different input lengths.
  • Acceleration Results: On-chip URAM and BRAM provide about 5× more effective bandwidth than GPU SRAM, yielding 1.8–2.2× to 1.2–5.6× kernel speedups and 1.04–1.49× end-to-end gains.The kernel ranges are 1.8–2.2× for SeerAttention-R top-k, 2.6–4.9× with threshold, and 1.2–5.6× for LServe.
  • Acceleration Results: Pipelined streaming and flexible datapaths deliver 1.3–2.2× DeepSeek Attention memory-processing speedup, 5.1–6.6× single-stage RAG improvement, and 3.1–4.0× Memory as Context improvement.These translate to 1.1–1.2×, up to 1.47–1.84×, and 1.3–1.6× end-to-end speedups, respectively; two-stage RAG reaches only 1.1–2.1× memory-processing speedup because of reranker dominance.
  • Acceleration Results: 1.8× speedup is achieved consistently over a GPU-only baseline for MemAgent under prefill–decode disaggregation, while PCIe transfer overhead is roughly three orders of magnitude below end-to-end latency.MemAgent benefits because its memory generation relies on decoding, where FPGA-controlled HBM transactions sustain higher effective bandwidth.
  • Energy and Batch-Size Scaling: Energy cost decreases by 1.11−1.61× for sparse attention, 1.07−1.21× for RAG, 4.66× for MemAgent, and 1.65× for Memory as Context.Across batch sizes, speedup increases for sparse attention and RAG, decreases for Memory as Context, and becomes a slowdown for MemAgent.

7. Conclusion … C. Cross-device Communication

The paper unifies LLM memory processing into a four-step pipeline and accelerates its heterogeneous workloads with a GPU–FPGA system. Across inference optimizations, the design improves performance and energy efficiency while using PCIe communication whose overhead is outweighed by FPGA kernel gains.

  • 7. Conclusion: 1.04–2.2× speedup and 1.11–4.7× energy reduction demonstrate the GPU–FPGA system’s end-to-end benefits for LLM memory processing.The prototype uses off-the-shelf devices, while a heterogeneous ASIC could further improve energy efficiency.
  • A. Related Works: Prior work characterizes LLM memory through taxonomies and evaluates parametric, contextual, and external memory categories.The supplied surveys also motivate comparisons across memory types and usage patterns.
  • A. Related Works: Sparse-attention, advanced-RAG, and context-compression research develops selective retrieval, reranking, recurrent memory, and compressed long-context processing.Examples include Longformer, Reformer, MInference, JudgeRank, HLATR, QwenLong-L1.5, RMT, Mamba, and Gated DeltaNet.
  • B. Detail Computation Properties of Memory Processing Pipeline: Higher arithmetic intensity indicates more compute-bounded operations, while pipeline bottlenecks vary substantially across inference optimizations.The analysis uses arithmetic-intensity and latency-distribution profiling to expose these computational differences.
  • B. Detail Computation Properties of Memory Processing Pipeline: Sparse attention and RAG bottleneck on compute relevancy and retrieval, whereas MemAgent bottlenecks on prepare memory, essentially LLM decoding.For sparse attention and RAG, the bottleneck share increases with memory size.
  • B. Detail Computation Properties of Memory Processing Pipeline: Memory-centric, irregular operations such as relevance computation and top-k retrieval are offloaded to FPGA customized data paths to mitigate GPU inefficiencies.This mapping targets memory-bounded workloads and shortens inference’s critical path.
  • C. Cross-device Communication: PCIe peer-to-peer transfer avoids the conventional CPU-and-DRAM path, while measured PCIe overhead remains sufficiently low despite lower bandwidth than NVLink.The conventional method reaches only 1/20 of peak PCIe bandwidth; PCIe provides 32 GB/s versus NVLink 3.0’s 600 GB/s, but FPGA kernel gains compensate.

C.1. PCIe Latency

PCIe transfers add microsecond-to-millisecond latency, but this overhead remains small relative to GPU computation across the evaluated memory-processing methods. Even as transfer sizes increase, computation exceeds communication time by approximately 1000×.

  • C.1. PCIe Latency: 14 µs transfers occur for sparse attention, while corresponding GPU kernels take 128–2450 µs.Transfers include the query vector, new key indexing vectors, and retrieved indices.
  • C.1. PCIe Latency: 7 µs transfers occur for RAG, while corresponding GPU kernels take 23–1596 ms.RAG transfers include only the query document counter and retrieved indices.
  • C.1. PCIe Latency: 20–320 µs transfers occur for Memory as Context, while corresponding GPU kernels take 26–498 ms.Transfers include memory, query, and retrieved embeddings for each segment.
  • C.1. PCIe Latency: 14–218 ms transfers occur for MemAgent, while corresponding GPU kernels require 17–534 s.Transfers include KV cache and token IDs for each segment.
  • C.1. PCIe Latency: PCIe communication overhead remains small (∼1000× difference) relative to computation time across all methods, even when data transfer size increases.This comparison covers sparse attention, RAG, Memory as Context, and MemAgent.

D. Detail Experiment Settings · E. FPGA Kernel Design

The paper specifies experimental configurations across several memory-processing methods and implements heterogeneous kernels by assigning dense computation to GPUs and sparse or specialized processing to FPGAs. FPGA designs use dataflow architectures for memory retrieval and decoding-specialized parallelism for generation.

  • D. Detail Experiment Settings: DeepSeek Attention generates 64 query heads and a key-indexing vector, computes dot products, and uses weighted scores for top-k selection with k=2048 in DeepSeek V3.2 Exp.The GPU baseline uses vLLM, while the MI210 loads only the first layer and extrapolates latency across layers.
  • D. Detail Experiment Settings: DRAGIN, FLARE, and Fixed-sentence RAG use Llama 2 7B with BM25S retrieval, retrieving 64 documents and generating at most 32 tokens.BM25S replaces ElasticSearch as the BM25 backend.
  • D. Detail Experiment Settings: Two-stage RAG selects 64 documents with bge-large-en-v1.5 and BM25, then reranks them to 10 documents using bge-reranker-large.The maximum generated length is 32 tokens.
  • D. Detail Experiment Settings: Memory as Context uses 1024-token segments and 32-token outputs, with the GPU sending newly generated memory and next-segment embeddings to the FPGA.The method follows Titans and uses recurrent memory embeddings.
  • D. Detail Experiment Settings: MemAgent uses 5000-token segments, a 1024-token memory, and 32-token maximum outputs, sending the prefill KV cache to the FPGA and receiving memory token IDs.The GPU-FPGA configuration transfers generated memory information between devices.
  • E. FPGA Kernel Design: GPU kernels use cuBLAS/cuSparse or rocBLAS/rocSparse for linear operations and custom CUDA/HIP kernels for nonlinear operations.These libraries and kernels target state-of-the-art GPU performance.
  • E. FPGA Kernel Design: Memory as Context follows FlexLLM’s HMT plugin design, streaming segment embeddings from CPU through on-chip projection while cross-attending cached memory embeddings in HBM.Figure 17 describes a dataflow kernel whose modules compute individual operations as data arrives.
  • E. FPGA Kernel Design: MemAgent’s FPGA executes only LLM decoding, specializing prior designs for decoding by treating attention as GEMV sequences and increasing hidden-dimension parallelism.The architecture uses separate special-function units, including SwiGLU and LayerNorm.

F. FPGA Kernel Improvement Analysis

The analysis attributes FPGA gains to higher on-chip bandwidth, pipelined datapaths for dependent and irregular operations, and memory-aware decoding architectures. These mechanisms produce substantial kernel and end-to-end speedups across sparse attention and heterogeneous decoding.

  • Case 1: Large On-chip Memory: 1.8–2.2× speedup for SeerAttention-R with top-k, 2.6–4.9× with threshold, and 1.2–5.6× for LServe translate to end-to-end gains.The corresponding end-to-end speedups are 1.04–1.25× for SeerAttention-R and 1.15–1.49× for LServe; LServe becomes worse than the GPU at 1M tokens.
  • Case 2: Pipelined and Flexible Datapath: 1.3–2.2× speedup in memory processing for DeepSeek Attention yields a 1.1–1.2× end-to-end speedup.Fine-grained pipelining overlaps index score computation with top-k or threshold-based selection, while customized random access reduces communication and computation latency.
  • Case 3: Faster Decoding: 1.8× speedup over a GPU-only baseline is achieved consistently by the GPU-FPGA system under prefill-decode disaggregation.FPGAs exploit customized architectures and controlled off-chip memory transactions to better use HBM bandwidth during memory-bound decoding.

G. Expanded Results for Sparse Attention and RAG

This section compares end-to-end and memory-processing latency for sparse attention and RAG, alongside energy efficiency and kernel power consumption across the GPU-FPGA system and baselines.

  • Latency: Figures 19–22 report absolute end-to-end and memory-processing latency for sparse attention and RAG.Figure 19 specifically shows sparse-attention end-to-end latency versus sequence length for the baseline and GPU-FPGA system.
  • Energy efficiency: Figures 23–24 compare energy efficiency in joule per token for sparse attention and joule per request for RAG.
  • Kernel power: 26.4 W is DeepSeek Attention kernel power on U55C, compared with 55 W on MI210.
  • Kernel power: 24.9 W is SeerAttention-R threshold kernel power on U55C, compared with 45 W on MI210; top-k uses 25.3 W and 46 W, respectively.
  • Kernel power: 29.7 W is RAG kernel power on U55C, compared with 106 W on MI210 and 34 W on EPYC 7v13; MemAgent and HMT/Titans consume 44.2 W and 42.6 W on U55C.

H. Results with NVIDIA A100 · I. Practical Impact and Device Scaling

The NVIDIA A100 comparison aggregates FPGA, GPU, and PCIe latency components and shows heterogeneous systems can outperform the A100 in selected configurations. Device-scaling analysis highlights prototype process differences, newer FPGA capabilities, and persistent benefits for memory-bound or irregular operations.

  • H. Results with NVIDIA A100: The A100 is more widely adopted for LLM inference than the AMD MI210, but both GPUs could not be profiled with U55C in the same node.End-to-end latency is calculated by aggregating measured latency components.
  • I. Practical Impact and Device Scaling: Memory-processing latency is evaluated across sparse attention mechanisms and sequence lengths, while RAG latency is evaluated across document counts.The cited figures cover sparse-attention latency, RAG end-to-end latency, and single- and two-stage RAG processing.
  • H. Results with NVIDIA A100: The evaluation aggregates measured FPGA, GPU, and PCIe communication latency components, while profiling kernel execution latency separately.Figures 25 and 26 use DeepSeek Attention as a representative case study.
  • H. Results with NVIDIA A100: Even when paired with the MI210, the GPU–FPGA heterogeneous system can outperform the A100 in certain configurations.The A100 generally outperforms the MI210 under identical optimizations.
  • I. Practical Impact and Device Scaling: Architectural trends suggest heterogeneous execution benefits will persist because GPU improvements can enlarge non-GPU latency fractions and targeted operations are memory-bound or irregular.The passage notes that these operations do not scale linearly with GPU HBM bandwidth.
  • I. Practical Impact and Device Scaling: Energy efficiency is reported for sparse attention in Joule per token and for RAG systems in Joule per request.These measurements are presented for the evaluated sparse-attention mechanisms and RAG systems.
  • H. Results with NVIDIA A100: A100 is generally faster in LLM inference than MI210, yet integrating U55C with A100 can still speed up inference.This comparison uses DeepSeek V3.2 Exp with DeepSeek Attention across MI210, A100, MI210 + U55C, and A100 + U55C.
Loading 2603.29002v3…