FlexInfer: Flexible LLM Inference with CPU Computations

FlexInfer: Flexible LLM Inference with CPU Computations

Speakers: Seonjin Na, Geonhwa Jeong, Byung Hoon Ahn, Tushar Krishna, Hyesoon Kim

Conference: MLSys 2025

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

Overview

The proliferation of Large Language Models (LLMs) has led to an explosion in demand for efficient inference, particularly in high-throughput applications like chatbots. However, a critical bottleneck in deploying these models is their immense memory footprint, encompassing both model weights and the dynamically growing Key-Value (KV) cache. Modern LLMs frequently exceed the memory capacity of even high-end accelerators like the NVIDIA H100 GPU. To address this, current solutions often resort to offloading model components to CPU memory, leveraging the CPU's larger capacity. While this enables the execution of larger models on memory-constrained GPUs, it introduces a significant performance penalty due to constant data transfers over the PCI Express (PCI-e) interconnect.

This talk, presented by Seonjin Na and a collaborative team from Georgia Tech, Meta, UCSD, and Intel Labs, introduces FlexInfer, a novel system designed to overcome the PCI-e bottleneck by intelligently and adaptively leveraging both GPU and CPU computation for LLM inference. FlexInfer re-evaluates the role of CPUs, particularly those equipped with dedicated matrix accelerators like Intel's Advanced Matrix Extensions (AMX), in the LLM inference pipeline. The core insight is that while GPUs excel in compute-intensive phases, CPUs can be highly effective in memory-bandwidth-bound phases, especially when equipped with modern acceleration features and abundant memory.

FlexInfer proposes a dynamic execution strategy that selects the optimal hardware (CPU, GPU, or a combination) for different phases of LLM inference – the pre-fill phase and the decode phase – based on runtime parameters and server configurations. By minimizing costly data transfers and maximizing the utilization of each processor's strengths, FlexInfer demonstrates significant latency reductions, making the deployment of very large LLMs more efficient and cost-effective on commodity hardware configurations consisting of a CPU and a single GPU.

Background

▶ Watch: Introduction to FlexInfer: Flexible LLM Inference with CPU (0:00)

Large Language Models perform inference in two distinct phases: the pre-fill phase (also known as the pre-computation phase) and the decode phase. During the pre-fill phase, the entire user input prompt is processed in parallel to generate the first output token and the initial KV cache. This phase is typically compute-bound, as it involves extensive matrix multiplications across the input sequence. Its performance is measured by Time to First Token (TFT). Following the pre-fill phase, the decode phase commences, where one output token is generated at a time, fed back into the model, and the process repeats until an end-of-sequence token is produced or the maximum sequence length is reached. The decode phase is largely memory bandwidth-bound, primarily due to the constant access and update of the KV cache. Its performance is measured by Time per Output Token (TPoPT).

A paramount challenge in LLM inference is the substantial memory requirement. Model sizes have grown exponentially; for instance, OPT 66 billion already exceeds the memory capacity of an H100 GPU, and models like Llama 3 405 billion continue to push these boundaries. Beyond model weights, the KV cache itself presents a significant memory burden. Its size scales linearly with factors such as batch size, sequence length, and the number of attention heads. For an OPT 66 billion model, the KV cache alone can exceed the model size when using a 2K sequence length and a batch size of 32. Consequently, it is common for the combined memory footprint of LLMs and their KV caches to surpass the capacity of contemporary GPU accelerators.

To mitigate these memory constraints, prior work introduced offloading-based inference techniques, such as Deep Speed Zero and FlexGen. These systems offload model weights, activations, and the KV cache to the more abundant CPU memory, transferring data to the GPU via the PCI-e interconnect only when needed for layer computation. While this approach effectively enables commodity GPUs to handle very large models that would otherwise exceed their memory, it introduces a major performance bottleneck: the overhead of PCI-e data transfer. Profiling Llama 2 70 billion on an H100 GPU with a 1K input sequence and 32 output tokens revealed that over 89% of the execution time was spent on PCI-e data transfer, even with aggressive overlapping of computation and data movement. This clearly highlighted that data transfer, not computation, was the dominant performance limiter in offloading scenarios.

In light of this bottleneck, FlexInfer revisits the opportunity for CPU computation during LLM inference. Recent advancements in CPU architecture have made them a more viable option. Modern server-class CPUs, including those from Intel (e.g., Sapphire Rapids with AMX), Arm, and AMD, are now equipped with dedicated on-chip accelerators similar to GPU tensor cores, supporting data types like BF16 and INT8. Furthermore, CPU servers inherently offer significantly larger memory capacities than GPUs, which can be further expanded using technologies like CXL (Compute Express Link). These two factors — enhanced computational capabilities and vast memory — position CPUs as potential contributors to efficient LLM inference, especially when model and KV cache sizes exceed GPU memory.

Key Findings

▶ Watch: Offloading bottleneck: 89% execution time on PCI transfer (3:40)

FlexInfer's research yielded several crucial findings that underpin its design and demonstrate its efficacy:

• CPUs excel in the decode phase: While CPUs generally exhibit lower computational throughput than GPUs, particularly for long sequences and large batch sizes, they significantly outperform GPUs in the memory-bandwidth-bound decode phase when offloading is involved. This is because CPUs can access the KV cache directly from their local memory without incurring PCI-e data transfer overheads, which plague GPU-based offloading solutions during token generation.
• CPUs are insufficient for the pre-fill phase with long sequences: Even with dedicated accelerators like AMX, CPU computation throughput remains insufficient to process longer input sequences and larger batch sizes efficiently during the compute-bound pre-fill phase. GPUs, despite data transfer overheads, still offer superior performance for this phase due to their parallel processing capabilities.
• Adaptive strategy is crucial: Given the complementary strengths and weaknesses of CPUs and GPUs across the pre-fill and decode phases, adaptively leveraging computation on both processors is essential for achieving optimal overall performance. A static approach or relying solely on one type of processor for all tasks in an offloading scenario leads to significant bottlenecks.
• FlexInfer significantly reduces latency: By dynamically selecting the optimal execution policy for each inference phase (pre-fill and decode) based on model architecture, runtime parameters, and server configuration, FlexInfer consistently achieves superior performance. On average, FlexInfer reduced latency by 75% on an Ice Lake CPU + A100 40GB GPU server and 76% on a Sapphire Rapids CPU + H100 80GB GPU server, compared to traditional FlexGen-like offloading baselines.
• Modern CPUs offer viable acceleration: CPUs equipped with on-chip accelerators like Intel's AMX are powerful enough to outperform GPU-based offloading solutions for certain workloads, particularly at smaller batch sizes and in the decode phase, validating the premise of revisiting CPU computation.

Technical Deep Dive

▶ Watch: Initial comparison: Sapphire Rapids CPU outperforms H100 with offloading (5:10)

FlexInfer is designed as a dynamic, adaptive system for LLM inference on a CPU + single GPU server architecture, aiming to minimize the PCI-e data transfer bottleneck inherent in offloading-based approaches. It comprises two primary components: the Execution Planner and the Inference Executor.

The Execution Planner is the brain of FlexInfer. Its core function is to determine the optimal execution policy for both the pre-fill and decode phases of LLM inference. This decision-making process relies heavily on a sophisticated performance estimator. This estimator predicts the expected Time to First Token (TFT) and Time per Output Token (TPoPT) for various candidate execution policies. It takes into account a comprehensive set of parameters:

• Model architecture: Specifics like layer count, hidden size, attention heads, etc.
• Runtime parameters: Input sequence length, output sequence length, batch size.
• Server configuration: Crucial hardware characteristics such as the compute throughput (e.g., FLOPS for GPU, AMX performance for CPU) and memory bandwidth (of both CPU and GPU) of the specific server setup.
• Profiling data: Empirical measurements gathered from the hardware to fine-tune predictions.

The performance estimator uses mathematical formulas to calculate the amount of computation and data load required for each phase under different policies. By meticulously modeling these factors, it can accurately predict which policy will yield the lowest latency for a given scenario.

Once the Execution Planner determines the optimal policies, the Inference Executor takes over. This component is responsible for executing the selected policies on the CPU and GPU hardware. The design of FlexInfer centers around intelligently leveraging a set of baseline execution policies, each with distinct advantages and limitations in a CPU + single GPU environment:

1. CPU-only Execution:

• Mechanism: The entire LLM layer execution, including model weights and KV cache, resides and operates solely on the CPU.
• Advantages: Completely eliminates PCI-e data transfer overhead, as all data is local to the CPU memory. Modern CPUs with accelerators (like AMX) can efficiently handle computations for certain workloads.
• Limitations: CPU computation throughput, even with accelerators, can be a bottleneck for the compute-bound pre-fill phase, especially with longer sequences and larger batch sizes.

1. GPU with Offloading (FlexGen-like):

• Mechanism: Model weights and the KV cache are stored in CPU memory. During computation, necessary data is transferred to the GPU via PCI-e, and computations are performed on the GPU. This policy often involves overlapping data transfers with GPU computation to maximize utilization.
• Advantages: Leverages the high computational throughput of GPUs, making it efficient for the compute-bound pre-fill phase.
• Limitations: Suffers from significant slowdown during the memory-bandwidth-bound decode phase due to the high and repetitive data transfer overheads between CPU and GPU for each generated token.

1. CPU-GPU Static Partitioning (SplitGen-like):

• Mechanism: This policy implements a form of pipeline parallelism, statically dividing the LLM layers between the CPU and GPU. For example, the first X layers might run on the CPU, and the remaining Y layers on the GPU. Data transfer between processors only occurs for the intermediate output of the partitioned layers.
• Advantages: Minimizes data transfer by only sending final layer outputs between processors, potentially leveraging both CPU and GPU concurrently.
• Limitations: The CPU portion can become a bottleneck, especially with longer sequences and larger batch sizes, if its computational throughput is insufficient to keep pace with the GPU or overall demand.

FlexInfer's strength lies in its ability to dynamically combine these policies for the pre-fill and decode phases. For instance:

• Scenario 1 (Longer input sequences, larger batch sizes): The Execution Planner might select GPU with offloading for the pre-fill phase to exploit the GPU's high computational throughput. For the decode phase, it might switch to CPU-GPU static partitioning to minimize data transfer overheads and leverage both processors more effectively during the token generation process. This combination takes advantage of GPU's raw compute for the initial heavy lift and then optimizes for memory access during sequential token generation.
• Scenario 2 (Shorter input lengths, smaller batch sizes): If the AMX-enabled CPU can efficiently handle the required computations, the planner might select CPU-only execution for both pre-fill and decode phases, completely eliminating transfer overheads and potentially offering a more cost-effective solution.

While FlexInfer currently focuses on these three core policies, the authors note that the approach is extensible. The framework can be adapted to consider additional execution strategies, such as tensor parallelism between CPU and GPU, as future needs arise. The implementation of FlexInfer involved extending existing frameworks, specifically Hugging Face Transformers and Intel Extension for PyTorch, to incorporate CPU-optimized kernels and support the deep execution strategy selection for both pre-fill and decode phases.

Experimental Setup & Results

▶ Watch: CPU strengths: decode phase; weakness: pre-phase for long sequences (6:05)

To thoroughly evaluate FlexInfer's performance, the researchers conducted extensive experiments, comparing it against several strong baselines on two distinct server configurations.

Baselines:

1. CPU-only: All operations performed on the CPU.
2. FlexGen: A representative GPU offloading solution where model weights and KV cache are stored in CPU memory and transferred to the GPU.
3. FlexGen Optimer: An optimized version of FlexGen that partially stores model weights and KV cache on the GPU to maximize available GPU memory, reducing some offloading overheads.
4. SplitGen: A baseline representing CPU-GPU static partitioning, where layers are divided between the CPU and GPU.

Server Configurations:

1. Server 1: An Ice Lake CPU (without AMX) paired with an NVIDIA A100 40GB GPU.
2. Server 2: A Sapphire Rapids CPU (with AMX) paired with an NVIDIA H100 80GB GPU.

Models and Workloads:

The evaluation used models from the OPT and Llama families, ranging in size from 30 billion to 70 billion parameters, ensuring that the models often exceeded the GPU's native memory capacity, necessitating offloading.

For the primary performance comparison, the input length was set to 512 tokens and the output length to 32 tokens. Initial exploratory analyses also used a 1K input length and 32 output tokens to understand the phase-specific performance characteristics.

Key Findings from Initial Analysis (CPU vs. GPU Offloading):

• When comparing CPU-only inference against GPU with offloading for a 128 input length and 32 output length, the Sapphire Rapids CPU (with AMX) consistently outperformed the H100 with offloading across all batch sizes. Even the Ice Lake CPU (without AMX) showed better performance than the GPU with offloading at smaller batch sizes.
• For longer input sequences (e.g., 1K input length), GPUs eventually outperformed CPUs despite the data transfer overhead. This was attributed to the CPU's computational throughput limitations for processing such long sequences and larger batch sizes in the pre-fill phase.
• A phase-wise breakdown (1K input, 32 output):
• Time to First Token (TFT - Pre-fill): CPUs exhibited significantly slower TFT compared to GPUs, even when GPUs suffered from data transfer overhead. This confirmed that CPU computation throughput, even with dedicated accelerators, was insufficient for compute-intensive pre-fill with long sequences.
• Time per Output Token (TPoPT - Decode): All CPU configurations clearly outperformed GPUs. This was a critical observation, as CPUs did not incur any data transfer overhead during the sequential generation of output tokens, unlike the GPU offloading approach which had to constantly move data.

FlexInfer's Overall Performance:

The results demonstrated that FlexInfer consistently achieved the best performance across both server configurations and various batch sizes. By adaptively selecting the optimal execution strategy for both the pre-fill and decode phases, FlexInfer effectively navigated the trade-offs between CPU computation throughput and PCI-e data transfer overhead.

• On average, FlexInfer reduced the end-to-end inference latency by a remarkable 75% on the Ice Lake CPU + A100 server.
• Similarly, it achieved a 76% latency reduction on the more powerful Sapphire Rapids CPU + H100 server.

While baselines like FlexGen Optimer and SplitGen did improve performance over vanilla FlexGen, they were still limited by either CPU computation throughput (SplitGen) or persistent data transfer overheads (FlexGen Optimer). FlexInfer's dynamic approach allowed it to always choose the most efficient path.

The paper also delves into further details such as the specific design of the performance estimator, sensitivity studies analyzing the impact of various parameters, and detailed PCI-e traffic analysis, providing a comprehensive understanding of the system's behavior. A minor but important detail confirmed in the Q&A was that the overhead for converting data layouts when transferring between CPU and GPU memory was found to be negligible.

Practical Implications

▶ Watch: FlexInfer's execution policies: CPU only, GPU offloading, static partitioning (7:25)

FlexInfer presents several significant practical implications for practitioners, infrastructure teams, model builders, and deployers working with large language models.

First and foremost, FlexInfer offers a cost-effective solution for deploying very large LLMs that exceed the memory capacity of a single high-end GPU. Instead of requiring multiple expensive GPUs for model parallelism or sharding, FlexInfer enables these models to run efficiently on more commodity hardware setups comprising a powerful CPU and a single GPU. This can drastically reduce the capital expenditure for inference infrastructure.

For infrastructure teams, FlexInfer provides a blueprint for optimizing resource utilization in heterogeneous environments. By dynamically allocating compute-intensive pre-fill tasks to the GPU and memory-bandwidth-sensitive decode tasks to the CPU (where data locality eliminates PCI-e overhead), FlexInfer ensures that each hardware component is leveraged for its strengths. This intelligent orchestration can lead to higher overall throughput and lower latency for LLM serving, directly impacting user experience in applications like chatbots. The insights on CPU capabilities, particularly with accelerators like AMX, highlight the importance of considering modern CPU architectures as active participants, not just passive memory banks, in the LLM inference pipeline.

Model builders and deployers will find FlexInfer particularly useful when dealing with memory-constrained deployment scenarios. It provides a robust strategy to overcome the critical PCI-e data transfer bottleneck that has historically plagued offloading-based solutions. The significant latency reductions (75-76%) directly translate to faster response times, which is crucial for real-time applications. Moreover, the flexibility to adapt to varying input lengths and batch sizes means that a single deployment can handle a wider range of user requests efficiently without needing to pre-provision for worst-case scenarios.

However, it's also important to consider the trade-offs and limitations. The current work focuses on a single CPU + single GPU system. While this is a common and cost-effective configuration, it may not be the optimal solution for environments with multiple GPUs (e.g., 8-GPU servers). In such setups, the CPU could potentially become a bottleneck if it's tasked with coordinating too many GPUs or handling disproportionately large portions of the workload. The authors acknowledge that using multiple GPUs would likely yield the best raw performance in terms of latency, but FlexInfer targets a different trade-off space focused on cost and efficiency for single-GPU systems.

Another practical consideration is the integration with continuous batching or other online serving optimizations. The current FlexInfer implementation assumes a more discrete batching scenario. For highly dynamic online serving environments where requests arrive asynchronously and are continuously batched, further research and integration would be required to ensure FlexInfer's dynamic policy selection remains optimal and doesn't introduce additional overheads. This is explicitly noted as a direction for future work.

In summary, FlexInfer offers a powerful, practical approach to unlock the full potential of large LLMs on existing and future CPU+GPU server infrastructure, making advanced AI more accessible and efficient for a broader range of applications and deployment environments.

Key Takeaways

• Large Language Models (LLMs) and their Key-Value (KV) caches often exceed the memory capacity of single high-performance GPUs.
• Offloading model components to CPU memory is a common strategy, but it introduces a major performance bottleneck: over 89% of execution time can be spent on PCI-e data transfers.
• Modern CPUs with dedicated matrix accelerators (e.g., Intel AMX) and larger memory capacities offer a viable opportunity for LLM inference.
• CPUs are highly beneficial for the memory-bandwidth-bound decode phase due to direct memory access, while GPUs remain superior for the compute-bound pre-fill phase.
• FlexInfer dynamically selects optimal execution policies (CPU-only, GPU with offloading, or CPU-GPU static partitioning) for the pre-fill and decode phases, based on runtime parameters and server hardware.
• By adaptively leveraging both CPU and GPU, FlexInfer significantly reduces inference latency, achieving average reductions of 75-76% compared to offloading-based baselines.

About the Speaker(s)

The talk was presented by Seonjin Na. The work on FlexInfer is a collaborative effort involving researchers from several prominent institutions and companies: Georgia Tech, Meta, UCSD, and Intel Labs. The list of co-authors includes Geonhwa Jeong, Byung Hoon Ahn, Tushar Krishna, and Hyesoon Kim. Their collective expertise spans various aspects of AI/ML systems, hardware architecture, and efficient inference techniques, contributing to the development of FlexInfer.