READ MORE

ArXiv: 2303.06865

🎯 Pitch

FlexGen achieves 1 token/s throughput for OPT-175B on a single 16GB GPU—a 100× improvement over prior offloading systems—by treating generative inference as a block-level schedule that reuses loaded weights across a massive effective batch size (up to 144). This throughput is unlocked only when weights, activations, and the attention cache are aggressively quantized to 4 bits and offloaded across GPU, CPU, and disk.

1. Executive Summary

This paper introduces FlexGen, a high-throughput generation engine for running large language model inference under severely limited GPU memory — specifically, a single commodity GPU — by aggregating memory across GPU, CPU, and disk and using a linear programming-based search algorithm to find efficient patterns for storing and accessing tensors (weights, activations, and key-value cache). The system also compresses weights and the attention cache to 4 bits with group-wise quantization, enabling a dramatically larger effective batch size that amortizes expensive I/O operations. On a single 16GB NVIDIA T4 GPU, FlexGen running OPT-175B achieves up to 1.12 token/s generation throughput and a 100× improvement in maximum throughput over state-of-the-art offloading systems (DeepSpeed ZeRO-Inference and Hugging Face Accelerate), reaching an effective batch size of 144 with compression enabled — the first time token/s throughput is demonstrated for OPT-175B on such constrained hardware — establishing that throughput-oriented generative inference on limited resources is viable only when the block schedule reuses loaded weights across a large batch and the CPU and disk are aggressively leveraged as part of the tensor placement strategy.

2. Context and Motivation

The Core Problem: LLM Inference Is Prohibitively Resource-Intensive

The fundamental problem this paper tackles is that running inference on large language models requires enormous GPU memory that most practitioners simply don't have. The numbers are stark. To load the OPT-175B model in FP16 precision requires approximately 325 GB of GPU memory just for the weights (Section 3). This means that inference — even for a single query — demands at least five A100 (80GB) GPUs and complex parallelism strategies to split the model across them (Pope et al., 2022; Aminabadi et al., 2022). When you account for the key-value (KV) cache — an additional memory structure that stores intermediate attention states for each token in the sequence — the memory demands explode further. The paper calculates that with a batch size of 512, an input sequence length of 512, and an output sequence length of 32 on OPT-175B, the KV cache alone consumes 1.2 TB, which is 3.8× the model weights themselves (Section 3). This makes the KV cache "a new bottleneck of large-batch high-throughput inference."

This resource barrier creates a fundamental inequity in who can use these models. Only organizations with access to datacenter-scale GPU clusters can deploy the largest LLMs, while individual researchers, small labs, and companies without massive infrastructure budgets are locked out. The paper frames this democratization angle explicitly: lowering resource requirements enables broader access to LLM capabilities, which matters for both research reproducibility and practical deployment.

The Throughput-Oriented Use Case: Why Latency Can Be Sacrificed

The paper distinguishes between two fundamentally different inference scenarios (Section 1). Latency-oriented inference is what most people think of: interactive chatbots where each query needs a response in milliseconds. This is the domain of systems like FasterTransformer (NVIDIA, 2022) and Orca (Yu et al., 2022), which optimize for wall-clock time per query on high-end hardware.

But the paper argues there exists an equally important — and under-explored — class of throughput-oriented inference tasks: batch processing of large datasets where individual query latency matters far less than the total number of tokens processed per second. The authors cite several concrete examples:

• Benchmarking: The HELM benchmark (Liang et al., 2022) requires evaluating models across thousands of prompts. Running this on a single GPU — rather than a cluster — would dramatically lower the barrier to model evaluation.
• Information extraction: Processing large document corpora to extract structured data (Narayan et al., 2018).
• Data wrangling: Using LLMs to clean, transform, and integrate datasets (Narayan et al., 2022).
• Form processing: Automatic extraction and structuring of information from forms (Chen et al., 2021).

In these tasks, the workflow is batched: you have hundreds or thousands of prompts to process, and nobody is waiting for any individual response. If processing the entire batch takes hours instead of minutes, that's acceptable — what matters is the total throughput in tokens per second. The paper's key insight is that this latency-tolerant setting opens up optimization opportunities that are fundamentally impossible in interactive settings. Specifically, you can use extremely large batch sizes to amortize the cost of moving data between slow storage and the GPU, overlapping I/O with computation across many queries simultaneously.

Prior Approaches and Why They Fall Short

The paper identifies three categories of existing work, each with critical limitations for the throughput-oriented, resource-constrained setting (Section 1, Section 2):

1. Model Compression

Techniques like INT8 quantization (Dettmers et al., 2022), GPTQ 4-bit weight quantization (Frantar et al., 2022), and SmoothQuant (Xiao et al., 2022) reduce model memory footprint, but as the paper notes, these methods "often assume that the model fits into the GPU memory." For a 175B parameter model on a 16GB GPU, even 4-bit weights alone consume roughly 87.5 GB — still far exceeding GPU capacity. Compression alone cannot bridge a ~20× memory gap. It must be combined with something else, yet prior compression work does not systematically study this combination with offloading.

2. Collaborative Inference

Systems like Petals (Borzunov et al., 2022) distribute model layers across a decentralized network of volunteer GPUs. This is a creative approach that can work, but it introduces network communication as a bottleneck: every forward pass requires transferring activations between machines. The paper demonstrates (Section 6.3, Figure 4) that under realistic network conditions (100ms delay, 100Mbps bandwidth), Petals' per-GPU throughput drops dramatically — from roughly 7 token/s under ideal conditions to under 1 token/s — because activation transfer dominates latency. Furthermore, Petals does not use offloading, so individual GPUs are limited to small batch sizes constrained by their VRAM capacity. The paper argues this makes Petals better suited for latency-sensitive scenarios where users want interactive responses, rather than throughput-oriented batch processing.

3. Offloading-Based Systems

This is the most directly comparable category, and where the paper's primary critique lands. Two systems support offloading LLM inference when GPU memory is insufficient:

• DeepSpeed ZeRO-Inference (Aminabadi et al., 2022): Can offload entire model weights to CPU or disk, and uses ZeRO data parallelism across GPUs.
• Hugging Face Accelerate (HuggingFace, 2022): Supports offloading a fraction of weights to CPU or disk.

The paper's central critique of these systems is that they inherit offloading strategies from training systems (Rajbhandari et al., 2021; Ren et al., 2021) without adapting to the unique structure of generative inference. This adaptation failure manifests in three specific ways:

First, the row-by-row (batch-by-batch) computation schedule is I/O-inefficient. Figure 3(a) illustrates the schedule used by DeepSpeed and Accelerate: process all layers for one batch (one prompt) completely, then move to the next batch. This means that for every batch, the system must load all model weights from slow storage into GPU, compute, then discard them. Since consecutive batches don't share weights in this schedule, the weights are loaded from scratch for every single batch — even though the same weights are needed repeatedly across batches. The paper quantifies this inefficiency indirectly through Table 2: with OPT-175B, Accelerate achieves only 0.01 token/s throughput (batch size 2) and DeepSpeed achieves 0.01 token/s (batch size 1), while FlexGen reaches 0.69 token/s — a 69× improvement from the schedule alone, before compression.

Second, existing systems restrict KV cache and activations to GPU memory only. The paper states that DeepSpeed and Accelerate "can only put cache/activations on GPU" (Section 6.1). This is a critical limitation because the KV cache — not the weights — becomes the memory bottleneck at large batch sizes. By forcing the KV cache to stay on GPU, these systems hit out-of-memory errors at tiny batch sizes (1–2 for OPT-175B), preventing them from amortizing I/O costs across a large batch. The paper makes this explicit: "DeepSpeed Zero-Inference and Hugging Face Accelerate cannot use a batch size larger than 2 due to out-of-memory issues" (Section 1).

Third, these systems do not exploit the column-by-column (layer-by-layer) weight reuse opportunity. In generative inference, the same layer weights are needed for every token generation step and every prompt in a batch. The row-by-row schedule treats each batch independently, missing the opportunity to load a layer's weights once, compute that layer for all tokens and all prompts, then move to the next layer. The zig-zag block schedule (Figure 3(b)) is FlexGen's solution to this: compute a block of prompts through a block of layers, reusing loaded weights across multiple prompts before swapping them out.

4. The Gap in the Literature: No Systematic Study of Throughput-Optimal Offloading

The paper identifies a specific gap: no existing system provides a principled framework for optimizing the complex design space of offloading strategies. The problem involves:

• Three types of tensors (weights, activations, KV cache)
• Three levels of the memory hierarchy (GPU, CPU, disk)
• Multiple possible computation orders (row-by-row, column-by-column, hybrid)
• Trade-offs between batch size, memory capacity, and I/O bandwidth at each level

Prior systems made these choices heuristically, inheriting patterns from training where the computation graph is fundamentally different (training processes one batch at a time through all layers, while inference generates tokens sequentially, requiring the KV cache to persist across generation steps). The paper argues this heuristic approach leaves enormous performance on the table — the 69–100× throughput gap between FlexGen and baselines (Table 2) quantifies precisely how much.

How This Paper Positions Itself

The paper positions FlexGen not as an improvement on existing offloading systems, but as a fundamentally different approach to the problem. Rather than adapting training offloading strategies to inference, FlexGen starts from the computational structure of generative inference itself and builds an optimization framework around it.

The key positioning moves are:

From heuristic to principled optimization. Instead of hard-coding a schedule and placement strategy, FlexGen defines a search space of possible strategies (Section 4.2), builds an analytical cost model to estimate throughput (Section 4.3), and uses linear programming to find the optimal configuration for a given hardware setup (Equation 1). This means the same FlexGen code can automatically adapt to different GPUs, CPU memory sizes, disk speeds, and model architectures — the paper emphasizes this is a "general framework" that can work across diverse hardware.

From latency-oriented to throughput-oriented thinking. The paper explicitly reframes the problem away from minimizing per-query latency toward maximizing total tokens processed per second. This reframing is what makes the large-batch, column-by-column schedule viable — low per-query latency would demand row-by-row processing to deliver results quickly, but throughput optimization rewards batching to amortize I/O.

From GPU-centric to whole-machine thinking. Rather than treating CPU and disk as inferior fallbacks when GPU memory is insufficient, FlexGen treats the entire memory hierarchy as a unified resource. The linear programming optimizer decides what fraction of each tensor type lives at each level, including the possibility of splitting tensors across levels (e.g., 50% of weights on GPU, 50% on CPU). The paper even finds that CPU computation — normally considered useless for deep learning — can be beneficial for attention score computation when the KV cache is stored on CPU, because moving the entire KV cache to GPU would be more expensive than simply computing attention on the CPU (Section 4.2, "Computation delegation").

From uncompressed to aggressively compressed. Unlike prior work that compresses weights only, FlexGen compresses both weights and the KV cache to 4 bits. The paper argues this is particularly important for the offloading setting because compression reduces I/O volume — the dominant bottleneck — not just memory footprint. The group-wise quantization method (Section 5) is chosen specifically because it allows fine-grained dequantization back to FP16 before computation on GPU, making it suitable for the offloading pipeline where compressed tensors are stored on CPU/disk and decompressed when loaded to GPU.

The paper's ambition is clear from the introduction: to demonstrate that throughput-oriented LLM inference on a single commodity GPU is not just possible, but can achieve performance that makes it practically useful — processing hundreds of tokens per second — by systematically solving the joint optimization problem of tensor placement, computation scheduling, and compression under memory and bandwidth constraints.

3. Technical Approach

3.1 Reader Orientation

FlexGen is a generation engine — a piece of software that takes a large language model (like OPT-175B) and a batch of prompts, and produces completions by orchestrating computation across a single GPU, CPU, and disk, even when the model is hundreds of gigabytes and the GPU has only 16 GB of memory. The problem it solves is that existing offloading systems (DeepSpeed, Accelerate) are forced to use tiny batch sizes on large models because they don't reuse loaded weights across prompts and they restrict the KV cache to GPU memory, which means their throughput collapses. FlexGen solves this by recasting the joint problem of what tensors to put where and in what order to compute as a linear programming optimization, finding configurations that enable batch sizes hundreds of times larger than baselines and amortizing the I/O cost of moving tensors across the slow memory hierarchy over many parallel queries.

3.2 Big-Picture Architecture (Diagram in Words)

The FlexGen system has five conceptual components working together:

1. Hardware Model (three-level memory hierarchy). The machine has a GPU (small, fast memory), CPU (medium, medium-speed memory), and disk (large, slow memory). Each device has known capacities and bandwidths between them. FlexGen models these as constraints in an optimization problem.

2. Computational Graph of LLM Inference. Figure 2 defines the abstract computation as a 3D grid: rows are layers (transformer layers 1 through l), columns are tokens (generation steps 0 through n), and the depth (colored squares) represents batches (different prompts in the batch). Each colored square is a "compute this layer for this batch at this generation step" operation, and squares of the same color share the same layer weights.

3. Offloading Strategy Search Space. A strategy has three parts: a compute schedule (the order in which squares of the grid are traversed), a tensor placement (what fraction of weights, activations, and KV cache live on GPU vs. CPU vs. disk), and a computation delegation choice (whether attention scores are computed on GPU or CPU). These variables define the entire set of possible execution plans.

4. Analytical Cost Model and Linear Programming Solver. For any given (schedule, placement, delegation) configuration, the cost model analytically estimates total execution time and peak memory usage as a function of hardware specs (bandwidths, FLOPs, capacities) and workload parameters (batch size, sequence lengths, model dimensions). The solver enumerates schedule choices and solves a linear program to find the placement percentages that minimize latency per token subject to memory constraints.

5. Runtime Execution Engine. Once a policy is selected, a PyTorch-based runtime executes Algorithm 1's block schedule with overlapping I/O and computation across multiple CUDA streams and CPU threads, managing pre-allocated buffers for tensors and using memory-mapped files for disk-resident data.

Information flows as follows: the user specifies the model, hardware specs, and workload parameters (prompt length, generation length) → the cost model is profiled on the hardware to fit bandwidth and FLOPs parameters → the optimizer searches over (block size, GPU batch size) pairs and solves LP for optimal placement → the selected policy configures the runtime → prompts are processed in blocks, with weights loaded layer-by-layer and reused across the GPU batches in a block, activations and KV cache streamed between devices, and the final generated tokens collected.

3.3 Roadmap for the Deep Dive

• First, the formal problem formulation (Section 4.1): the generative inference with offloading as a graph traversal problem on the 3D computation grid, with explicit constraints on what must be computed before what and where tensors can live. This defines what a valid strategy is.

• Second, the search space construction (Section 4.2): the three axes of strategy design — compute schedule (row-by-row vs. zig-zag block), tensor placement (percentages at each memory tier), and computation delegation (GPU vs. CPU for attention). This defines what strategies FlexGen can express, including the zig-zag block schedule that is provably within 2× of optimal I/O complexity.

• Third, the cost model and policy search (Section 4.3): how execution time and peak memory are estimated analytically, how overlapping I/O and compute is modeled, and how the optimization problem decomposes into an outer enumeration over schedule parameters and an inner linear program over placement percentages. This defines how FlexGen finds a good strategy.

• Fourth, the multi-GPU extension (Section 4.4): how pipeline parallelism across multiple GPUs reduces per-GPU memory pressure, potentially enabling super-linear scaling in decoding throughput by allowing larger batch sizes when offloading to CPU instead of disk.

• Fifth, the approximate methods (Section 5): the 4-bit group-wise quantization for weights and KV cache, and the Top-K sparse attention mechanism, both designed to reduce I/O volume (not just memory footprint) in the offloading setting. This defines how FlexGen trades negligible accuracy for substantial throughput gains.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily a systems paper whose core idea is that the joint optimization of computation scheduling, tensor placement across a three-level memory hierarchy, and compression enables LLM inference with high throughput on hardware that would otherwise be considered impossibly constrained. The intellectual heavy lifting is in formalizing the design space as a tractable optimization problem and proving that the chosen schedule class is near-optimal.

Problem Formulation: Generative Inference as Constrained Graph Traversal

The paper models the problem of LLM inference with offloading as finding a valid path through a computational graph (Section 4.1, Figure 2). This formalization is the foundation for everything that follows, because it makes precise what "valid" means and what degrees of freedom exist.

The graph structure. Consider a machine with three devices: a GPU, a CPU, and a disk. The GPU and CPU can both perform computation (at very different speeds), while the disk cannot. These three devices form a three-level memory hierarchy: the GPU has the smallest but fastest memory, the disk has the largest but slowest memory, and the CPU sits in between. Since the LLM cannot fit entirely in GPU memory, computation must proceed part-by-part by partially loading the model.

The computation to be performed is represented as a 3D grid where:

• Each row corresponds to a transformer layer (1 through l).
• Each column corresponds to a generation step (token 0, the prefill step, through token n, the final generated token).
• Each square in a given (row, column) position actually represents multiple squares stacked in the "batch" dimension — one per prompt in the dataset. Squares with the same color share the same layer weights (because they correspond to the same layer number).

Figure 2 shows this for a 4-layer model generating 3 tokens per prompt. The prefill step (Token 0) processes all prompt tokens at once, while subsequent steps (Tokens 1, 2) process one new token each, generating the KV cache from previous steps.

What a valid path means. A "path" is a sequence that visits — i.e., computes — every square exactly once. The paper defines "a valid path as a path that traverses (i.e., computes) all squares" subject to four types of constraints:

1. Dependency constraint: "A square can only be computed if all squares to its left on the same row were computed." This encodes the autoregressive dependency: you cannot generate token t for a given batch and layer until you have generated token t-1 for that same batch and layer, because the attention computation needs the KV cache from previous steps.

2. Input availability constraint: "To compute a square on a device, all its inputs (weights, activations, cache) must be loaded to the same device." If you want to compute layer i on the GPU, the weight matrix for that layer, the activation from the previous layer, and the relevant KV cache entries must all reside in GPU memory at computation time.

3. Output persistence constraint: "After being computed, a square produces two outputs: activations and KV cache. The activations should be stored until its right sibling is computed. The KV cache should be stored until the rightmost square on the same row is computed." The activation from layer i at step t is needed as input to layer i+1 at step t (the "right sibling"). The KV cache from layer i at step t is needed for attention computation in all subsequent steps (up to the rightmost square on that row). Neither can be discarded prematurely.

4. Memory capacity constraint: "At any time, the total size of tensors stored on a device cannot exceed its memory capacity." Each device has a fixed memory budget, and the tensors residing there at any moment must fit.

The optimization objective. "The goal is to find a valid path that minimizes the total execution time, which includes the compute cost and I/O cost when moving tensors between devices." This is not simply a scheduling problem — it's a joint scheduling, placement, and routing problem where decisions about when to compute a square interact with decisions about where to store its inputs and outputs.

Why this formulation matters. By formalizing the problem this way, the paper makes two critical observations that motivate the rest of the design: (1) the dependency structure is not fully sequential — there is parallelism across batches and across tokens within a batch, limited only by the KV cache dependency; and (2) the I/O cost of moving weights dominates if you reload them for every batch, so reusing weights across multiple batches before swapping them out is the key to high throughput. These observations directly lead to the zig-zag block schedule.

Search Space Construction: Schedule, Placement, and Delegation

Given the problem formulation, FlexGen defines a search space of possible strategies along three orthogonal axes. The search space is designed to be expressive enough to capture near-optimal strategies but structured enough that optimization over it is tractable.

Compute Schedule: From Row-by-Row to Zig-Zag Block

The baseline: row-by-row schedule. Figure 3(a) illustrates the schedule used by DeepSpeed and Accelerate: process each batch (row in the "batch" dimension) completely before moving to the next. Specifically, for the first prompt: compute all tokens through all layers, top to bottom. Then move to the second prompt. This is the intuitive order because "it is the fastest way to finish the generation for one batch and the KV cache can be freed immediately after a row." But it has a fatal flaw: "because every two contiguous squares do not share weights, this schedule has to repeatedly load the weights and incurs huge I/O costs." Every time you switch to a new batch, you need to reload all layer weights for that new batch, even though you just had them loaded for the previous batch.

The insight for weight reuse: column-by-column traversal. The paper observes that "all squares in a column share weights." If you traverse the graph column-by-column (computing one layer across all tokens and all batches before moving to the next layer), you can load that layer's weights once, keep them on GPU, and compute all dependent squares without ever swapping them out. The I/O per byte of weights loaded is amortized across all tokens and all batches that use those weights.

However, pure column-by-column traversal is impossible: "we cannot traverse a column all the way to the end because the activations and KV cache still need to be stored. Hence, we have to stop when they fill the CPU and disk memory." The activations from computing layer i for all batches must be held until layer i+1 finishes its computation. The KV cache must be held much longer — from its creation until the final generation step. Both grow with batch size and sequence length, limited by total system memory.

The compromise: zig-zag block schedule. Figure 3(b) illustrates the resulting compromise. Instead of processing all batches for all layers column-by-column (which would overflow memory), or one batch at a time row-by-row (which wastes I/O), the system processes a block of GPU batches through a column of layers, moves to the next column (next generation step), and repeats. Within a block, the weights for each layer are loaded once and reused across all GPU batches in that block. The block size is the product of the GPU batch size (how many prompts are processed together at the GPU level) and the number of GPU batches in a block (how many such groups are chained together before moving to the next token).

Algorithm 1: Block Schedule with Overlapping. The paper presents the pseudocode:

for i = 1 to generation length do
    for j = 1 to num layers do
        // Compute a block with multiple GPU batches
        for k = 1 to num GPU batches do
            load weight(i, j + 1, k)
            store activation(i, j, k - 1)
            store cache(i, j, k - 1)
            load cache(i, j, k + 1)
            load activation(i, j, k + 1)
            compute(i, j, k)
            synchronize()
        end for
    end for
end for

The paper notes that "the first six functions in the innermost loop can be seen as launched in parallel with six logical threads because there are no dependencies." This is the overlapping optimization: while batch k is computing, the system simultaneously loads weights for the next layer (j+1), loads the cache and activation for the next batch (k+1), and stores the cache and activation for the previous batch (k-1). The synchronize() call waits for all six operations to complete before advancing the loop, ensuring correctness.

Theoretical guarantee: within 2× of optimal I/O complexity. Theorem 4.1 states: "The I/O complexity of the zig-zag block schedule is within 2× of the optimal solution." The proof in Appendix A.2 compares the zig-zag block schedule to a more sophisticated "diagonal block schedule" (Figure 5) that is proven I/O-optimal asymptotically, and shows that the zig-zag variant achieves at least half the throughput of the optimal schedule. The core reason: in the zig-zag schedule, peak memory is highest when processing later generation steps (because the KV cache has grown), so early steps underutilize available memory. The diagonal schedule rebalances this by interleaving computation across tokens so memory usage is closer to constant. The paper acknowledges that the diagonal schedule is not implemented due to "practical implementation difficulty" (specifically, managing non-contiguous KV cache buffers for efficient attention computation), but uses it to provide a theoretical worst-case bound on the simpler block schedule's suboptimality.

The two schedule parameters. The block schedule introduces exactly two integer parameters into the search space:

• GPU batch size (gbs): the number of prompts processed simultaneously on the GPU within one iteration of the innermost loop. This determines the granularity of parallelism: larger values increase GPU utilization but consume more GPU memory for intermediate tensors.
• Number of GPU batches in a block (#gb): how many innermost-loop iterations are chained before moving to the next generation step. The product bls = gbs × #gb is the block size (or effective batch size), which determines how many prompts share a single load of each layer's weights.

Tensor Placement: Fractional Allocation Across the Memory Hierarchy

Beyond the computation order, a strategy must specify where each of the three tensor types (weights, activations, KV cache) resides within the three-level memory hierarchy. The paper uses nine continuous variables to capture this:

• Weights: wg, wc, wd — the fraction of model weights stored on GPU, CPU, and disk respectively. These sum to 1 (wg + wc + wd = 1).
• KV cache: cg, cc, cd — the fraction of the KV cache stored on GPU, CPU, and disk. These also sum to 1 (cg + cc + cd = 1).
• Activations: hg, hc, hd — the fraction of activations stored on GPU, CPU, and disk. These also sum to 1 (hg + hc + hd = 1).

In principle, these percentages are continuous between 0 and 1, meaning the optimizer can decide, for example, to keep 20% of weights on GPU and 80% on CPU. The paper acknowledges that "the percentage cannot be an arbitrary real number between 0 and 1, because the tensor cannot be split arbitrarily," but relaxes the variables to be continuous "since it is changing gradually." In practice, the runtime implements tensor splitting at layer granularity for weights (e.g., keep layers 1–19 on GPU, layers 20–96 on CPU), and at tensor granularity for activations and KV cache (splitting individual tensors across devices). The paper justifies this choice: "Coarser granularity leads to lower runtime overhead but it is less flexible and its cost is difficult to analyze. Considering both the runtime overhead and desired flexibility, we use layer granularity for weights, and tensor granularity for activations and the KV cache."

Why fractional placement matters. The search space allows, for instance, the KV cache to be split across GPU and CPU. This is a departure from baselines (DeepSpeed, Accelerate) that "can only put cache/activations on GPU." Allowing the KV cache to spill to CPU enables much larger effective batch sizes, because the KV cache — not the weights — is the memory bottleneck at scale. The paper calculates that with b = 512, s = 512, n = 32, and OPT-175B dimensions, the KV cache alone is 1.2 TB, which is 3.8× the model weights and far exceeds any single GPU's memory.

Computation Delegation: When to Use CPU Compute

A counterintuitive design choice in FlexGen is the option to perform attention score computation on the CPU rather than the GPU. The paper explains: "This is because the computation of attention scores during decoding is I/O-bounded. Consider a case where the KV cache is stored on the CPU. Computing the attention scores on the GPU requires moving the entire KV cache to the GPU, which incurs a substantial I/O cost as the KV cache is huge. In contrast, computing the attention score on the CPU does not require moving the KV cache. It only requires moving the activations from the GPU to the CPU."

The quantitative analysis is striking. Let b be the GPU batch size, s be the sequence length, and h1 be the hidden size. The size of the KV cache that must move to GPU for attention computation is b × s × h1 × 4 bytes (in FP16, each element is 2 bytes, and the KV cache for one layer contains both keys and values, hence 4×). The size of the activation that must move from GPU to CPU for CPU-side attention is b × h1 × 4 bytes. The ratio is s: CPU attention computation reduces I/O volume by a factor equal to the sequence length. "For long sequences (e.g., s ≥ 512), it is better to compute the attention scores on the CPU if the associated KV cache is not stored on the GPU." This is a design choice that follows directly from the observation that in the offloading regime, I/O bandwidth dominates compute throughput — moving fewer bytes wins, even if the computation itself is slower on CPU.

Cost Model and Policy Search: From Search Space to Optimal Policy

The search space defines what strategies FlexGen can express. The cost model and policy search define how FlexGen picks a particular strategy given hardware specs and workload parameters.

Analytical Cost Model

The cost model predicts the total execution time for processing a block of prompts, and also estimates peak memory usage to enforce capacity constraints. It operates at the granularity of one layer during one phase (prefill or decoding), and then scales up to the full model and full generation.

Latency estimation for one layer, one phase. The total latency for processing one layer during prefill is denoted Tpre, and the averaged latency for one layer during decoding is Tgen. The total latency for a block is then estimated as:

$T = T_{\text{pre}} \cdot l + T_{\text{gen}} \cdot (n - 1) \cdot l$

where l is the number of layers and n is the number of tokens to generate. The (n-1) factor accounts for the fact that the first token (prefill) is handled separately, and the remaining n-1 decoding steps each process one new token per prompt.

The max-over-I/O-and-compute assumption. The paper assumes perfect overlapping of I/O and computation: the six logical threads in Algorithm 1 run in parallel, and the latency of each phase is bounded by the slowest among them. Specifically:

$T_{\text{pre}} = \max(\text{ctog}_p, \text{gtoc}_p, \text{dtoc}_p, \text{ctod}_p, \text{comp}_p)$

where:

• ctog_p is the latency of reading data from CPU to GPU during prefill for one layer,
• gtoc_p is the latency of writing data from GPU to CPU during prefill,
• dtoc_p is the latency of reading from disk to CPU during prefill,
• ctod_p is the latency of writing from CPU to disk during prefill,
• comp_p is the computation latency during prefill.

Tgen is defined analogously with subscript g for the decoding phase:

$T_{\text{gen}} = \max(\text{ctog}_g, \text{gtoc}_g, \text{dtoc}_g, \text{ctod}_g, \text{comp}_g)$

What the I/O terms look like. Consider dtoc_g (disk-to-CPU read during decoding) as an example. The paper estimates it by summing the I/O volume for each tensor type, divided by the disk-to-CPU bandwidth:

$\text{dtoc}_g = \frac{1}{\text{disk bandwidth}} \cdot \left[(8h_1^2 + 4h_1 \cdot h_2) \cdot wd + 4 \cdot \text{bls} \cdot (s + \frac{n}{2}) \cdot h_1 \cdot cd + 2 \cdot \text{bls} \cdot h_1 \cdot hd\right]$

where:

• h1 is the hidden size of the transformer (12288 for OPT-175B),
• h2 is the hidden size of the second MLP layer (49152 for OPT-175B),
• 8h1² + 4h1·h2 is the size of FP16 weights for one transformer layer (attention weights: 4 matrices of shape h1 × h1 = 4h1² elements × 2 bytes = 8h1² bytes; MLP weights: two matrices of shapes h1 × h2 and h2 × h1 = 2h1·h2 elements × 2 bytes = 4h1·h2 bytes; total: 8h1² + 4h1·h2 bytes),
• wd is the fraction of weights on disk,
• bls is the block size,
• s is the prompt length,
• n is the output sequence length,
• 4·bls·(s + n/2)·h1 is the average size of the KV cache for one layer (4 accounts for keys + values in FP16; s + n/2 is the average sequence length during decoding — the KV cache starts at size s after prefill and grows to s + n at the last step, averaging to approximately s + n/2),
• cd is the fraction of KV cache on disk,
• 2·bls·h1 is the size of activations for one layer (one activation tensor in FP16),
• hd is the fraction of activations on disk.

What this equation computes: it estimates the total time the disk read channel is busy during one layer's decoding phase, by dividing the total bytes that must be read from disk (weights fraction on disk + KV cache fraction on disk + activation fraction on disk) by the disk bandwidth. This time, taken as a maximum over all I/O paths and computation, determines the layer latency under the perfect-overlap assumption.

Why this form: the max-over-paths formulation captures the bottleneck principle in overlapping systems — the slowest operation determines overall throughput. Summing I/O volumes before dividing by bandwidth correctly aggregates multiple independent data transfers on the same channel. The s + n/2 approximation for average KV cache size is an analytical simplification of the true sum Σ_{t=1}^{n}(s + t) which grows quadratically; the linear approximation is justified because it makes the cost model a linear function of the placement variables, enabling linear programming.

Computation terms. The comp_g term breaks into GPU and CPU computation:

$\text{comp}_g = \text{gpu\_comp}_g + \text{cpu\_comp}_g$

where GPU computation includes matrix multiplications for the linear layers and batched matrix multiplications for attention on the GPU-resident portion of the KV cache, and CPU computation includes attention on the CPU/disk-resident KV cache:

$\text{cpu\_comp}_g = \frac{4(cc + cd) \cdot \text{bls} \cdot (s + \frac{n}{2}) \cdot h_1}{\text{cpu flops}}$

This models the CPU attention computation as the number of floating-point operations needed for computing attention scores over the CPU-resident cache portion, divided by CPU FLOPs/sec.

Memory constraints. In addition to latency, the cost model estimates peak memory usage for GPU, CPU, and disk at two critical points: during prefill and during decoding. These estimates include both "home" memory (tensors that permanently reside on the device) and "working" memory (scratch space needed during computation). The full model in Appendix A.3 tracks detailed memory for attention intermediates (QKV projections, attention scores, attention outputs), MLP intermediates, and embedding lookups. These constraints ensure the optimizer does not select policies that would cause out-of-memory errors.

Policy Search as a Two-Level Optimization

The full policy has 11 variables: block size (bls), GPU batch size (gbs), and the nine placement percentages (wg, wc, wd, cg, cc, cd, hg, hc, hd). The paper decomposes the optimization into two levels:

Outer level: enumerate schedule parameters. "We first enumerate a few choices of (bls, gbs) tuple. Typically, gbs is a multiple of 4, and bls is less than 20 so there are not too many choices." This restriction is practical: (bls, gbs) are discrete variables in a limited range (the paper shows policies with bls up to 256 in Table 2, but the enumeration handles far fewer candidate pairs because bls = gbs × #gb, and #gb is the variable being enumerated).

Inner level: linear programming over placement percentages. For a fixed (bls, gbs) pair, finding the optimal placement p = (wg, wc, wd, cg, cc, cd, hg, hc, hd) becomes a linear programming problem:

$\min_p \quad T / \text{bls}$

subject to: