NEO: Saving GPU Memory Crisis with CPU Offloading for Online LLM Inference

NEO: Saving GPU Memory Crisis with CPU Offloading for Online LLM Inference

Speakers: Xuanlin Jiang, Yang Zhou, Shiyi Cao, Ion Stoica, Minlan Yu

Conference: MLSys 2025

YouTube: https://www.youtube.com/watch?v=None

Overview

In the rapidly evolving landscape of large language models (LLMs), online inference has become a cornerstone for numerous cutting-edge applications. However, the relentless growth in LLM size has precipitated a significant challenge: the GPU memory crisis. This talk, "NEO: Saving GPU Memory Crisis with CPU Offloading for Online LLM Inference," presented by Yang Zhou and co-authored with Xuanlin Jiang, Shiyi Cao, Ion Stoica, and Minlan Yu from a collaboration spanning Peking University, UC Berkeley, Davis, and Harvard, introduces a novel system designed to mitigate this crisis. NEO proposes an intelligent CPU offloading strategy that leverages the often-underutilized CPU resources to enhance GPU efficiency and overall inference throughput for LLMs.

The core problem addressed by NEO is the disproportionate growth of GPU compute power versus its memory capacity. While GPU compute capabilities have expanded dramatically, memory capacity has lagged, leading to bottlenecks in online LLM inference. Specifically, limited GPU memory restricts the size of the KV cache and, consequently, the effective batch size, resulting in low GPU utilization and reduced throughput. NEO tackles this by intelligently offloading the memory-intensive decoding attention computation and its associated KV cache to the CPU, thereby freeing up valuable GPU memory and compute cycles. The significance of this work lies in its ability to achieve substantial performance gains—up to 6.6 times higher throughput on memory-constrained GPUs like the T4 and 14% on H100s—without compromising model accuracy, offering a practical and cost-effective solution for deploying and scaling LLM services.

Background

▶ Watch: Introduction: GPU memory crisis in online LLM inference (0:00)

The advent of large language models has revolutionized AI applications, but their immense scale presents significant infrastructure challenges, particularly in the context of online inference. Online inference, where individual requests are processed with low latency requirements, is crucial for interactive applications. As LLMs continue to grow in parameter count, the memory footprint required for their operation, especially the KV cache (key-value cache), escalates dramatically. The KV cache stores intermediate attention states for each token in a sequence, allowing for efficient generation of subsequent tokens without recomputing past states. This cache, however, consumes a substantial amount of GPU memory.

The fundamental issue, termed the "GPU memory crisis" by the researchers, stems from a divergence in hardware evolution: while GPU compute capabilities have seen "vast compute improvement," there has been "no significant memory expansion" in recent years. This imbalance means that even powerful GPUs can become memory-bound, unable to fully utilize their compute potential. A smaller KV cache directly translates to smaller maximum batch sizes, which in turn leads to under-saturated GPUs, low utilization, and ultimately, reduced inference throughput and higher latency.

Existing solutions to this memory bottleneck have their own set of limitations. Quantization and sparsification techniques can reduce the model's memory footprint by representing parameters with fewer bits or pruning less important connections. For instance, quantization might halve model parameters, saving considerable GPU memory for the KV cache. However, these methods often come at the cost of "output quality," a compromise many applications cannot afford. Another approach is memory offloading, where excessive GPU data is moved to the CPU's much larger memory and swapped back on demand. While seemingly straightforward, the primary hurdle here is the "relatively pretty low" PCI bandwidth between the CPU and GPU. For an operation like LLM inference, where the entire KV cache might need frequent swapping, this low bandwidth can introduce severe overheads, making naive offloading "even make inference throughput and latency worse than the GPU only approach" when the GPU already has enough memory for model weights. This highlights the critical need for a more intelligent, bandwidth-aware offloading strategy, which NEO aims to provide.

Key Findings

▶ Watch: Key Observation 1: KV cache accessed only by decoding attention (4:00)

NEO's design is predicated on two critical observations about LLM inference and hardware characteristics, which collectively form the foundation for its high-performance CPU offloading strategy. These observations address the shortcomings of prior offloading attempts and pave the way for an efficient hybrid CPU-GPU execution model.

The first key observation is that "only the decoding attention operation accesses the KV cache." The LLM inference process typically involves an initial pre-fill stage where the input prompt is processed to generate the first set of KV cache entries, followed by a decoding stage where one token is generated at a time. Within each transformer layer, operations include pre-projection, storing KV cache, computing attention, and post-projection. Crucially, the researchers found that "all the KV cache accesses happen only during the decoding attention phases." This insight is profound because it implies that if the decoding attention computation can be offloaded to the CPU, then the need for constant KV cache swapping between GPU and CPU via the slow PCI bus can be "possibly eliminate[d]." Instead of swapping the KV cache, the computation that uses it is moved to where the KV cache resides.

The second key observation concerns the relative performance characteristics of GPUs and CPUs. While GPUs are renowned for their massive parallel compute capabilities, their advantage over CPUs in terms of memory bandwidth is significantly smaller. The talk highlights that "GPUs are actually much closer than to CPU in terms of memory bandwidth, rather than the compute." Citing prior work like Fast Decode, a modern X86 CPU's memory bandwidth is "only about three times higher" than an A10G GPU, whereas the GPU is "over two orders of magnitude faster than the CPU in terms of compute." Crucially, the decoding attention operation—the very operation identified in the first observation—is "exactly memory bound." This confluence of factors makes offloading the memory-bound decoding attention computation to the CPU not just a "possible solution" but "even a promising option," as the CPU's memory bandwidth is relatively competitive for this specific workload, and offloading it frees up the GPU for its compute-intensive tasks.

Based on these two observations, NEO's central insight is to offload both the KV cache and the decoding attention computation to the CPU. However, realizing this insight for online LLM serving with strict latency Service Level Objectives (SLOs) and achieving performance gains necessitates addressing two primary challenges: first, "how do we efficiently overlap the CPU and GPU within each inference iteration," given their fundamental differences in compute capacity; and second, "how do we efficiently schedule requests across inference iterations," especially considering the dynamic and varied nature of real-world LLM workloads with diverse input/output lengths. NEO's subsequent technical contributions are built upon overcoming these challenges.

Technical Deep Dive

▶ Watch: NEO's insight: Offload KV cache and decoding attention to CPU (6:00)

NEO's innovative approach to CPU offloading for LLM inference revolves around an asymmetric pipelining strategy, meticulously designed to overcome the limitations of prior attempts and maximize CPU-GPU overlap. The researchers first analyze two "strawman" designs to illustrate the problems NEO addresses.

The first strawman, "simple offloading," extracts the decoding attention computation and its associated KV cache and offloads them entirely to the CPU, while the GPU handles the rest, including the pre-fill attention and linear operations. The main drawback of this design is a "lack of any overlapping." There's no concurrent execution between CPU and GPU computations, nor between computation and PCI communication, leading to inefficient resource utilization.

The second strawman, "asymmetric pipelining," attempts to introduce overlap. In this design, the pre-fill stage remains on the GPU. During the decoding stage, a single decoding batch is split into two sub-batches. The GPU performs linear operations for one sub-batch while the CPU concurrently performs attention operations for the other. This creates a symmetric data flow where GPU linear operations of one batch overlap with CPU attention operations of the other. However, this design also suffers from several issues:

1. Insufficient Overlap: It "entirely overlook[s] the pre-fill stage and KV cache swap out time," during which the CPU can remain idle.
2. Bottleneck Imbalance: The "linear stage on the GPU is typically much slower than the attention stage on the CPU." This imbalance prevents effective overlap, leading to a "CPU bottleneck and also wasting GPU cycles."
3. Wasted GPU Memory: This design only stores model weights and runtime activations on the GPU, leaving "significant utilization of GPU memory" unused as it "does not store any KV cache."

Motivated by these deficiencies, NEO proposes its refined asymmetric pipelining with asymmetric batch division to achieve full GPU and CPU overlapping. Instead of a symmetric split, NEO consolidates operations to balance the workload across the heterogeneous devices:

• Sub-batch Zero: This batch is designed to be "slightly longer to more job on the GPU." It includes "all the GPU pre-fill operation including linear and pre-fill attention," "GPU decoding of all requests of some request," and "CPU decoding of some request." By mixing pre-fill computation and a portion of GPU decode computation, NEO "largely extend[s] the duration of the GPU compute and utilize their GPU memory," enabling effective CPU-GPU overlap.
• Sub-batch One: In contrast, this batch is "pretty simple," primarily containing "the CPU decoding attention of most request." This keeps the CPU busy with its specialized, memory-bound task.

This "mixed batching" strategy, combined with layer-wise swapping, is critical for high overlapping. NEO initiates PCI transmission "immediately after each layer's KV cache is computed." This fine-grained, layer-by-layer data transfer overlaps the PCI communication with ongoing computation, minimizing idle times.

Furthermore, NEO addresses the issue of unused GPU memory. Unlike the second strawman that keeps no KV cache on the GPU, NEO performs partial offloading. This means "some request decoding is on the CPU" (with its KV cache stored in CPU memory) while "some is on the CPU [sic - likely meant GPU here, as it's partial offloading with KV cache in both places]" (with its KV cache stored in GPU memory). This allows NEO to store KV cache in both CPU and GPU memory, thereby "fully utilized GPU memory" for requests that can benefit from GPU-resident KV cache and computation.

To manage the complexities of "real-world inference workload[s]" which are "complex, irregular, and sometimes dynamic changing," NEO employs an intelligent dynamic scheduling policy. This scheduler is guided by several high-level principles:

1. Greedy: At the start of each iteration, the scheduler estimates throughput for two options—a purely GPU-only inference schedule or the two-batch asymmetric pipeline schedule—and selects the one with the higher estimated throughput.
2. Balancing: For the asymmetric pipelining, the scheduler's goal is to "minimize the pipeline bubbles," ensuring continuous data flow and computation.
3. Hide CPU: A key objective is to keep the GPU busy when the CPU is busy, ensuring that the GPU, being the more powerful compute resource, is never idled waiting for the CPU.
4. Maximize GPU: The scheduler ensures that the batch size is "sufficient enough to fully utilize the GPU," preventing under-saturation.

While the talk outlines these principles, the detailed scheduling algorithms are referred to the paper, indicating a sophisticated mechanism for adaptive resource allocation in dynamic environments.

Experimental Setup & Results

▶ Watch: Analyzing Strawman 1: Simple offloading's lack of overlapping (8:00)

To validate the efficacy of NEO, the researchers conducted extensive evaluations using a diverse set of real-world and synthetic inference datasets. The experimental setup was designed to assess NEO's performance across various LLM sizes and different GPU architectures, specifically focusing on scenarios where model weights could fit into GPU memory and remain resident there.

Datasets: The evaluation employed "several real-world inference data set and various synthetic ones" to simulate different inference workloads, covering a spectrum of input and output sequence lengths and request arrival patterns characteristic of online LLM serving.

Models: NEO was evaluated using "various size of models." While specific model names or parameter counts were not explicitly stated in the transcript, the context implies modern LLMs that strain GPU memory. A key assumption in the evaluation was that "The model weights in our setting can fit into the GPU memory and start always stay there," meaning NEO primarily focuses on offloading the KV cache and associated decoding computation, not the base model weights themselves.

Hardware: The experiments utilized a range of NVIDIA GPUs to demonstrate NEO's versatility and impact across different memory and compute profiles:

• NVIDIA H100: A high-end, recent GPU representing state-of-the-art compute.
• NVIDIA T4: A more memory-constrained, older generation GPU, often found in cloud environments.
• NVIDIA A10G: Another cloud-oriented GPU, used to specifically study the impact of CPU capacity. For these experiments, "four kinds of EC2 instances" were used, each with "one single A10G GPU" but varying numbers of CPU cores, allowing for an analysis of how different CPU capacities affect NEO's performance gains.

Metrics: The primary metrics for evaluation included:

• Online Latency: The time taken to process individual requests, crucial for online inference SLOs.
• Throughput: The number of requests processed per unit of time, indicating overall system capacity.
• Throughput Gains: The percentage or factor by which NEO improves throughput compared to baseline systems.

Baselines: NEO's performance was compared against a "GPU only inference schedule," referred to as the "VM" or "baseline." This baseline represents a standard LLM serving setup without CPU offloading.

Headline Results:

• Overall Throughput: NEO consistently sustained higher workloads than the GPU-only baseline while maintaining "comparable latency at lower rate."
• H100 Performance: On the powerful H100 GPU, NEO achieved "around 14% higher throughput." This demonstrates that even with high-end GPUs, intelligent offloading can yield measurable improvements.
• T4 Performance: The gains were particularly dramatic on the memory-constrained T4 GPU, where NEO achieved "6.6 times throughput" at certain latency SLOs. The speaker later clarified this, stating "we achieve nearly nine times throughput gains" in the T4 settings. This significant improvement is attributed to the T4's "extremely constrained memory budget for the KV cache," which severely limits batch size for the GPU-only baseline. By offloading the KV cache to the CPU, NEO enabled "much higher batch size," unlocking substantial throughput gains.
• Impact of CPU Capacity: The experiments on A10G GPUs with varying CPU cores showed that NEO's performance scales with CPU power. With a "more powerful CPU," NEO achieved "up to 79% higher throughput over the baseline." This confirms that the CPU is an active participant in the pipeline, and its capabilities directly influence the benefits derived from offloading. The researchers noted that these EC2 instances are "all available on the AWS cloud," indicating the practical applicability of their findings.

In summary, NEO demonstrated considerable throughput gains across different GPU types and CPU configurations, especially excelling in scenarios where GPU memory is a bottleneck. The results confirm that intelligent CPU offloading, as implemented by NEO, is a viable and effective strategy for mitigating the GPU memory crisis in online LLM inference without sacrificing inference accuracy.

Practical Implications

▶ Watch: NEO's Asymmetric Pipelining for full CPU/GPU overlapping (12:00)

NEO's innovative approach to CPU offloading for online LLM inference carries significant practical implications for a wide range of stakeholders in the AI/ML ecosystem, including practitioners, infrastructure teams, model builders, and deployers.

For practitioners and infrastructure teams, NEO offers a compelling solution to the omnipresent challenge of scaling LLM services. The ability to achieve "considerable super gains on different GPUs"—up to 6.6x or even 9x on memory-constrained T4 GPUs, and a solid 14% on high-end H100s—translates directly into higher request throughput and improved utilization of existing hardware. This means infra teams can serve more users or larger models with their current GPU clusters, potentially delaying costly hardware upgrades. The fact that NEO achieves this "with the same hardware cost," by leveraging local host CPU resources, makes it an economically attractive proposition, especially in cloud environments where CPU capacity is often abundant and underutilized relative to GPU compute.

For model builders and deployers, NEO opens up possibilities for deploying larger and more complex LLMs in latency-sensitive online inference scenarios. The critical advantage is that NEO achieves these performance gains "if inference accuracy, we are not sacrificing any accuracy part." This is a crucial distinction from techniques like quantization, which often trade off accuracy for memory savings. Deployers can now confidently use NEO to serve models closer to their full fidelity, maintaining output quality while enhancing throughput. It also means that models previously deemed too large for certain GPU configurations due to KV cache limitations can now be deployed more effectively.

However, realizing these benefits involves certain tradeoffs and limitations. The primary tradeoff is the increased complexity in system design and orchestration. NEO requires careful "restructur[ing] the pipeline of inference" and sophisticated "adaptive scheduling policy" to balance workloads across heterogeneous CPU and GPU resources. Infrastructure teams would need to implement or integrate NEO's asymmetric pipelining and dynamic scheduling mechanisms, which are more intricate than a purely GPU-only setup. The performance gains are also tied to CPU capacity, as evidenced by the "up to 79% higher throughput over the baseline with a more powerful CPU." This implies that while local CPUs are leveraged, their specifications (e.g., core count, memory bandwidth) will influence the maximum achievable benefits. Teams would need to consider the CPU-to-GPU ratio when designing their serving infrastructure.

Looking ahead, the discussion about future hardware like NVIDIA's GB200 series highlights potential shifts. The GB200's "really unified interconnect between the CPU and GPU" with "much higher PCI bandwidth" could alter the optimal offloading strategy. With higher bandwidth, "you can probably just do offloading but not do any compute" on the CPU and still get performance gains. However, the speaker notes that "the KV cache size you can offload is still limited because compared to the GPU memory bandwidth, the PCI bandwidth is still bottleneck." This suggests that even with improved interconnects, offloading computation to powerful ARM cores on future platforms might still be beneficial for "less compute intensive but memory intensive operations," especially given that these ARM cores "actually have pretty high memory bandwidth." This indicates that while hardware evolves, the fundamental principles of balancing memory-bound tasks to appropriate hardware, as demonstrated by NEO, will likely remain relevant.

In essence, NEO provides a robust, accuracy-preserving method to extract more performance from existing hardware for online LLM inference, especially on memory-constrained GPUs. Its practical implication is a more efficient, scalable, and cost-effective deployment of large language models in production environments, albeit with an increased demand for intelligent system design and resource orchestration.

Key Takeaways

• GPU Memory Crisis: The talk highlights a critical bottleneck in online LLM inference, where limited GPU memory for the KV cache restricts batch sizes and leads to low GPU utilization and throughput, despite vast compute improvements.
• Targeted Offloading: NEO's core insight is to offload the KV cache and the memory-bound decoding attention computation to the CPU, based on observations that only decoding attention accesses the KV cache and CPUs are relatively competitive in memory bandwidth for this task.
• Asymmetric Pipelining: The system introduces an advanced asymmetric pipelining strategy with asymmetric batch division and layer-wise swapping to achieve high CPU and GPU overlap, significantly improving efficiency over naive offloading or symmetric pipelining.
• Dynamic Scheduling: NEO employs a dynamic scheduler guided by principles of greediness, balancing, hiding CPU idle time, and maximizing GPU utilization to adapt to irregular and dynamic real-world LLM workloads.
• Significant Performance Gains: NEO demonstrates substantial throughput improvements, achieving up to 6.6x (or nearly 9x) higher throughput on memory-constrained T4 GPUs, 14% on H100s, and up to 79% with more powerful CPUs, all without compromising inference accuracy.
• Cost-Effective Scaling: By intelligently leveraging local CPU resources, NEO offers a practical and cost-effective solution for scaling LLM inference, enabling higher utilization of existing hardware and facilitating the deployment of larger models.

About the Speaker(s)

The talk "NEO: Saving GPU Memory Crisis with CPU Offloading for Online LLM Inference" was presented by Yang Zhou. The work is a collaborative effort, with Xuanlin Jiang identified as the first author, who unfortunately could not attend the conference due to visa issues, but whose "hard work goes to him." The research team also includes Shiyi Cao, Ion Stoica, and Minlan Yu. The authors represent a joint effort across several prominent academic institutions: Peking University, UC Berkeley, UC Davis, and Harvard. This diverse institutional collaboration underscores the interdisciplinary nature of the research, combining expertise in systems, machine learning, and hardware optimization.