FlexGen: High-Throughput Generative Inference of Large Language Models with a Single GPU

Most LLM serving talk assumes you care about interactive latency: first token soon, stream smoothly. FlexGen asks a different question. What if the job is latency-insensitive — benchmark suites, offline grading, batch rewriting, HELM-style evals — and the scarce resource is a single commodity GPU with CPU RAM and an SSD behind it?

Then the bottleneck is not “make one request snappy.” It is “keep the GPU fed while tensors live across a memory hierarchy.” FlexGen’s punchline: treat GPU / CPU / disk as one pool, place weights–activations–KV with a searchable policy, schedule layer compute in a zig-zag block pattern that reuses expensive loads, and optionally compress weights and KV to 4 bits so effective batch size can explode.

What this paper explains

Generative inference has two phases that both need memory:

1. Prefill — process the prompt; materialize KV for the prefix.
2. Decode — emit tokens one by one; append KV each step.

On a cluster with enough HBM, you keep weights and a large KV cache on-GPU and chase low latency. On a 16GB T4 trying to run OPT-175B, that fantasy collapses. Weights alone are hundreds of GB in FP16; the KV cache for a large batch can be multiple× the weight footprint. Prior offloaders often push weights to CPU/disk but still park KV on the GPU, which caps batch size at 1–2 and kills throughput.

FlexGen’s unit of success is generation throughput (tokens/s counting prefill + decode) under a chosen latency budget — a Pareto curve, not a single “fastest chat” point.

Prior limits

Systems in the DeepSpeed ZeRO-Inference and Hugging Face Accelerate lineage already offload. FlexGen’s critique is that those strategies still leave throughput on the table for huge models:

• KV stuck on GPU → tiny batches when sequence × batch grows.
• Naive layer-by-layer I/O → reload the same weight blocks too often for large batched decode.
• No unified search over where weights, activations, and KV live → hand-tuned policies that miss the high-batch frontier.
• No co-designed compression → cannot enlarge the feasible batch region when HBM and even DRAM are saturated.

Interactive engines (continuous batching, paged KV) optimize a different objective. FlexGen is explicitly about high-throughput / limited-resource generation.

The mechanism

1. Three-tier memory. Every tensor class — weights $W$, activations, KV cache — can be placed partially on GPU, CPU, or disk. Percentages $(w_g, w_c, \ldots)$ define a policy. More GPU residency means less I/O but less room for batch; more disk means huge models fit and huge batches become possible if you can hide or amortize transfers.

2. Zig-zag block schedule. Instead of finishing one full sequence through all layers before starting the next micro-batch’s I/O pattern, FlexGen processes a block of GPU micro-batches with a zig-zag over layers: load a layer’s weights once, run that layer for many micro-batches, then move on. Theorem 4.1 in the paper states the I/O cost of this schedule is within $2\times$ of an idealized optimal schedule they analyze but do not fully implement.

3. Overlap. Weight prefetch for the next layer, cache/activation traffic for neighboring batches, and compute for the current batch are overlapped where the hardware allows.

4. Policy search. Given hardware sizes and bandwidths, FlexGen casts throughput maximization as a linear program over placement variables and schedule parameters, then materializes a concrete policy (GPU batch size × number of GPU-batches in a block, plus residency percentages).

5. Optional 4-bit compression. Group-wise compression of weights and attention cache (no retraining / calibration in their reported setup) shrinks footprints enough that OPT-175B can keep weights in CPU and avoid disk for some high-throughput points — unlocking effective batch sizes like $144$ in their abstract’s headline setting.

Conceptually, throughput is:

and the lever FlexGen turns hardest is effective batch size under the I/O schedule, not speculative decoding or phase disaggregation.

Algorithm / figure walkthrough

1. Profile the machine — GPU HBM, CPU DRAM, disk bandwidth (their Table 1: T4 16GB, ~208GB DRAM, NVMe SSD).
2. Search a policy — LP over where $W$, activations, and KV live and how to block micro-batches.
3. Run zig-zag blocks — for each layer block, load weights, sweep micro-batches through that layer, overlap next I/O.
4. Prefill then decode under the same placement philosophy; decode’s growing KV is why unified KV placement matters.
5. Optionally compress to 4 bits to expand the feasible batch frontier.
6. Pick a point on the latency–throughput curve — small batch ≈ lower latency; huge batch ≈ max tokens/s for offline jobs.

The opening visual shows a request’s tensors moving through GPU → CPU → disk tiers while a zig-zag path reuses a loaded weight slab across several micro-batches.

What to notice when reading

• FlexGen is not claiming to beat vLLM on chat TTFT. It targets throughput under limited accelerators.
• Unified placement of KV (not only weights) is the quiet systems move that enables large batches.
• Zig-zag is an I/O amortization trick: pay for a weight load, harvest it across micro-batches.
• The LP is a search tool over a discrete-ish policy space — the paper’s contribution is the formulation + schedule, not “ML magic.”
• Compression here is a throughput enabler (fit more batch), adjacent to but distinct from AWQ’s activation-aware accuracy story.
• Related: ZeRO-Inference / Accelerate (baselines), Petals (decentralized collective), DistServe (phase split for SLOs), H₂O (KV eviction) — different levers.

Results and evidence

From the authors’ evaluation (NVIDIA T4 16GB, their Table 1 hardware, OPT models — not universal constants):

• On OPT-175B, FlexGen reports a new Pareto frontier versus DeepSpeed ZeRO-Inference and Hugging Face Accelerate, including up to about $100\times$ higher maximum throughput when 4-bit compression and large effective batches are allowed (abstract / Fig. 1 narrative: effective batch $144$, ~$1$ token/s generation throughput class result on one 16GB GPU).
• Without compression, they still report large gains from better placement + zig-zag; e.g. on the order of $40\times$ higher throughput than DeepSpeed/Accelerate at a ~5000s latency regime with effective batch $64$ (2048 tokens) in their cited comparison.
• At still higher latency budgets they report up to about $69\times$ higher maximum throughput by enlarging effective batch to $256$ (8192 tokens) in the discussed setting.
• Table-style throughput snapshots (e.g. prompt $512$, generate $32$ on one T4) list FlexGen around $0.69$ token/s on OPT-175B and $7.32$ token/s on OPT-30B for concrete policies; compressed variants rise further (e.g. ~$1.12$ token/s on OPT-175B in the same table family).
• HELM: they report benchmarking a 30B model on a 16GB GPU across 7 representative sub-scenarios in about 21 hours.
• Accuracy: they report 4-bit weight/KV compression with negligible loss on their OPT perplexity / task checks (see paper §6.2 tables).

Limitations

• Optimized for throughput, not TTFT/TPOT chat SLOs; huge batches mean huge per-batch latency.
• Disk offload is sensitive to SSD bandwidth; slow disks flatten the curve (paper discusses SSD speed sensitivity).
• LP policies assume relatively stable hardware and sequence shapes; highly dynamic interactive traffic needs a different stack.
• Padding variable-length prompts to a max length is a simple but wasteful batching choice.
• Collective / multi-GPU systems (e.g. Petals-style) can win under good networks; FlexGen’s headline is the one GPU + DRAM + SSD regime.
• 4-bit results are for their compression recipe on OPT; do not conflate with every modern weight-only quantizer.

How to read the paper

1. Abstract + §1 — latency-insensitive throughput; single-GPU motivation.
2. Background — memory math for weights vs KV at large batch.
3. Design — placement space, zig-zag schedule, overlap, LP search.
4. Compression + approximations — 4-bit and sparse attention as optional enlargeers.
5. Evaluation — vs DeepSpeed / Accelerate / Petals; latency–throughput curves; HELM.
6. Appendix — I/O optimality sketch (Theorem 4.1), ablations, breakdowns.

Knowledge check

Keep reading / Sources

• Paper: arXiv:2303.06865 · PDF
• Code: FMInference/FlexGen
• Related Daily issues: ZeRO, PagedAttention, Orca, DistServe, AWQ