PaperlayerView exact revision on arXiv

Source-linked AI summary

SARATHI: Efficient LLM Inference by Piggybacking Decodes with Chunked Prefills

Amey Agrawal, Ashish Panwar, Jayashree Mohan, Nipun Kwatra, Bhargav S. Gulavani, Ramachandran Ramjee

arXiv:2308.16369v1cs.LGcs.DC

TL;DR

LLM inference suffers from low-utilization decodes and pipeline bubbles caused by unequal prefill and decode costs. SARATHI combines chunked-prefills with decode-maximal batching to create uniform hybrid work units, improving throughput across models, hardware, and parallelism settings.

  • Problem

    LLM inference underutilizes GPUs during autoregressive decode and experiences pipeline bubbles when prefill and decode costs imbalance micro-batches.

  • Method

    SARATHI splits prefill requests into equal compute-sized chunks and fills hybrid batches with one prefill chunk plus decodes that piggyback on it.

  • Results

    SARATHI improves throughput across evaluated models and hardware, including up to 1.91× end-to-end throughput improvement and up to 10× decode-throughput improvement.

  • Takeaways & Limitations

    Uniform hybrid batches improve GPU utilization and significantly reduce pipeline bubbles by balancing micro-batch compute requirements.

  • Takeaways & Limitations

    Selecting an optimal chunk size remains open because it depends on hardware, model characteristics, sequence length, token composition, and potentially unknown P:D ratios.

Abstract

from arXiv · show

Large Language Model (LLM) inference consists of two distinct phases - prefill phase which processes the input prompt and decode phase which generates output tokens autoregressively. While the prefill phase effectively saturates GPU compute at small batch sizes, the decode phase results in low compute utilization as it generates one token at a time per request. The varying prefill and decode times also lead to imbalance across micro-batches when using pipeline parallelism, resulting in further inefficiency due to bubbles. We present SARATHI to address these challenges. SARATHI employs chunked-prefills, which splits a prefill request into equal sized chunks, and decode-maximal batching, which constructs a batch using a single prefill chunk and populates the remaining slots with decodes. During inference, the prefill chunk saturates GPU compute, while the decode requests 'piggyback' and cost up to an order of magnitude less compared to a decode-only batch. Chunked-prefills allows constructing multiple decode-maximal batches from a single prefill request, maximizing coverage of decodes that can piggyback. Furthermore, the uniform compute design of these batches ameliorates the imbalance between micro-batches, significantly reducing pipeline bubbles. Our techniques yield significant improvements in inference performance across models and hardware. For the LLaMA-13B model on A6000 GPU, SARATHI improves decode throughput by up to 10x, and accelerates end-to-end throughput by up to 1.33x. For LLaMa-33B on A100 GPU, we achieve 1.25x higher end-to-end-throughput and up to 4.25x higher decode throughput. When used with pipeline parallelism on GPT-3, SARATHI reduces bubbles by 6.29x, resulting in an end-to-end throughput improvement of 1.91x.

1 Introduction

LLM inference is inefficient because compute-rich prefills are followed by low-utilization decodes, while varying phase costs create pipeline bubbles. SARATHI combines chunked-prefills and decode-maximal batching to improve utilization and balance pipeline work.

  • LLM inference has become a dominant GPU workload as large models see widespread use across applications.
  • Prefill processes prompt tokens in parallel and saturates GPU compute at small batch sizes, whereas decode generates one token per autoregressive pass with low utilization.
  • At small batch sizes, decode cost per token can reach approximately 200 times the prefill cost, and repeated decode passes substantially affect overall efficiency.
  • Pipeline parallelism can suffer bubbles because prefill and decode times vary across micro-batches, despite micro-batching intended to mitigate stalls.
  • SARATHI splits prefills into equal compute-sized chunks and fills each hybrid batch’s remaining slots with decodes that piggyback on the prefill chunk.
  • 1.91× end-to-end throughput improvement is achieved with SARATHI in the reported pipeline-parallel GPT-3 evaluation, alongside gains across models and hardware.

2 Background

Transformer inference applies decoder-block operations differently in prefill and decode: prefills process sequences in parallel, while decodes process one newly generated token using cached prior keys and values. Multi-GPU deployment uses tensor or pipeline parallelism to scale models and batch capacity.

  • Transformer decoder block: A transformer decoder block contains self-attention and a feed-forward network, comprising preproj, attn, postproj, ffn_ln1, ffn_ln2, and other operations.
  • Prefill phase: During prefill, the decoder-block input tensor has shape [B,L,H], where B is batch size, L is sequence length, and H is embedding size.
  • Transformer decoder block: Self-attention transforms X into Q, K, and V through preproj, computes Y through attention, and returns Z through postproj.
  • Transformer decoder block: The FFN applies two batched matrix-matrix multiplications, producing an intermediate shape [B,L,H2] before returning to shape [B,L,H].
  • Decode phase: During decode, the input shape is [B,1,H], and prior tokens’ K and V tensors are stored in GPU memory as the KV cache.
  • Multi-GPU inference: Tensor parallelism shards layers across GPUs, whereas pipeline parallelism assigns subsets of layers to GPUs and uses micro-batches moving between stages.

3 Motivation

LLM inference combines a compute-saturating prefill phase with a memory-bound decode phase, while pipeline parallelism introduces bubbles when micro-batch runtimes differ. SARATHI addresses both through uniform hybrid batches.

  • Prefill and decode efficiency: Prefill has nearly constant per-token cost across batch sizes, indicating GPU saturation even at batch size one, while decode cost falls as batch size increases.
  • Prefill and decode efficiency: 200×, 100×, and 16.7× are the decode-to-prefill per-token cost ratios at batch sizes 1, 2, and 18, respectively.
  • Arithmetic intensity: Decode arithmetic intensity drops by more than two orders of magnitude relative to prefill and becomes compute-intensive only at batch size 256, which is infeasible because of KV-cache footprint.
  • Pipeline bubbles: Pipeline parallelism creates GPU-inactivity bubbles because later stages wait for preceding micro-batches, and runtime variation persists under iteration-level scheduling.
  • Insights: SARATHI slices prefills into compute-efficient chunks and piggybacks decodes alongside them, producing uniformly compute-intensive batches with reduced runtime variance across stages.

4 SARATHI: Design and Implementation

SARATHI combines chunked-prefills with decode-maximal batching to create compute-saturating, uniform hybrid batches for efficient LLM inference. The design addresses decode inefficiency, pipeline imbalance, and chunking overhead through batch construction, fused operations, and tile-aware chunk sizing.

  • Overview: SARATHI uses chunked-prefills and decode-maximal batching to improve LLM inference performance.Chunked-prefills splits prefills, while decode-maximal batching combines one prefill chunk with piggybacked decodes.
  • Chunked-prefills: Chunked-prefills splits large prefill requests into compute-sized chunks that can saturate GPU execution.Prefill throughput shows diminishing returns beyond a model- and GPU-dependent token count, enabling carefully sized chunks.
  • Chunked-prefills: Chunked-prefills incur repeated KV-cache memory accesses, while smaller chunks can reduce arithmetic intensity and prefill efficiency.The authors state that attention overhead is limited because attention is a small fraction of the forward-pass time; profiling can select an effective chunk size.
  • Decode-maximal batching: Decode-maximal batching combines one prefill chunk with decode tokens, producing compute-saturating and uniform work units.The remaining batch slots are filled with decodes, which piggyback on the prefill computation.
  • Identifying the ideal chunk size: Tile-aware sizing ensures the fused-operation matrix dimension remains a multiple of the tile size.The method selects a chunk for prefill efficiency and adjusts it with piggybacked decode tokens to align with the tile size.

5 Evaluation

The evaluation examines SARATHI's throughput effects, comparison with iteration-level scheduling, and impact on pipeline bubbles. It spans multiple models, GPU settings, and large-scale parallel deployments.

  • Evaluation questions: The evaluation measures decode and end-to-end throughput across sequence lengths, batch sizes, and prefill-to-decode ratios.These measurements address SARATHI's throughput impact under varying workload conditions.
  • Evaluation questions: The evaluation compares SARATHI with iteration-level scheduling mechanisms such as Orca.This comparison is listed as a distinct evaluation question.
  • Evaluation questions: The evaluation measures SARATHI's effect on GPU bubbles and throughput for pipeline-parallel models.This tests whether the techniques improve pipeline-parallel execution efficiency.

5.1 Evaluation on a Single GPU

On a single GPU, SARATHI combines chunked-prefills with decode-maximal batching to improve decode and end-to-end throughput across models, sequence lengths, batch sizes, and P:D ratios.

  • Decode speedup: Up to 10× decode-throughput improvement is achieved for LLaMA-13B on an A6000 GPU, with gains varying by batch size and sequence length.The evaluation compares SARATHI with separate prefill-only and decode-only baseline batches.
  • Varying P:D ratio: Peak throughput occurs when P:D = C/(B−1), so decode iterations align with the number of prefill chunks and decodes perfectly piggyback.For chunk size 256 and batch size 18, the peak improvement is 1.27× at P:D = 14 for sequence length 1K.
  • Varying P:D ratio: Smaller chunks increase piggybacking opportunities, but very small chunks reduce arithmetic intensity and add KV-cache read overhead.Chunk size 256 therefore outperforms chunk sizes such as 128 in the reported experiments.
  • Varying P:D ratio: SARATHI’s gains remain around 10% across a broad range of P:D ratios, although performance peaks only when prefills and decodes are balanced.At low P:D, prefill tokens can run out; at high P:D, decode tokens can run out.
  • Varying batch and chunk sizes: Up to 1.6× runtime reduction occurs in linear operations, with FFN reductions of 1.3×–1.6× and preproj/postproj reductions of 1.05×–1.38×.At small batch sizes, most throughput improvement comes from more efficient FFN computation in decode-maximal batches.

5.2 Comparison to Iteration-level Scheduling

Compared with Orca’s iteration-level scheduling, SARATHI provides more consistent throughput by deliberately constructing hybrid batches and using smaller prefill chunks.

  • Scheduling design: Orca’s overlap between prefills and decodes is incidental and varies with request arrivals and departures, whereas SARATHI constructs hybrid batches explicitly.Best-case Orca can be viewed as a special SARATHI case with chunk size equal to the maximum sequence length.
  • Throughput comparison: 1.11× higher throughput is achieved by best-case Orca scheduling for 1K sequences, while worst-case Orca performs similarly to the baseline.As sequence length increases, best-case Orca approaches baseline performance because its full prefill exhausts available prefill work.
  • Throughput comparison: 1.27×, 1.25×, and 1.23× overall throughput gains are achieved by SARATHI across the three evaluated sequence lengths.These results are reported under the optimal P:D setting with chunk size 256.
  • Request latency: SARATHI avoids the latency increase caused when Orca adds a longer prefill sequence to a running batch by using smaller prefill chunks.Longer prefills can delay ongoing decodes in iteration-level scheduling.

5.3 Pipeline Parallelism with SARATHI

In simulated GPT-3 pipeline-parallel inference across 64 A100 GPUs, SARATHI reduces pipeline bubbles and accelerates request completion compared with Orca-style pipeline scheduling.

  • Evaluation setup: The evaluation profiles GPT-3 operation and communication runtimes, then uses regression to estimate missing values in a simulated serving system.The simulated deployment spans 64 A100 GPUs across eight servers with InfiniBand.
  • Pipeline bubbles: 6.29× median bubble-time reduction per request is achieved by creating equal-compute units of work.Bubble time sums bubble time across all micro-batches and iterations for a request.
  • Request completion: 1.91× faster PP execution is achieved with SARATHI than baseline TP-PP execution with Orca scheduling.The comparison uses simulated request completion times for 10K requests.
  • Request completion: 1.48× faster execution is achieved by SARATHI-enabled PP compared with TP-only execution.TP-PP supports a 2.45× higher batch size than TP-only because it leaves more room for the KV cache.

5.4 Ablation Study of Chunked-prefills

The ablation shows that smaller prefill chunks increase attention and prefill overhead, while larger chunks preserve prefill efficiency; smaller chunks can nevertheless improve throughput by piggybacking more decodes.

  • Overhead of chunking: Chunk size 64 incurs 3× attention overhead and about 5× overall prefill-time overhead.Smaller chunks reduce GPU utilization and require more KV-cache reloads.
  • Overhead of chunking: Chunk sizes 256 and 512 limit end-to-end prefill computation loss to within 20% and 10%, respectively.Larger chunks benefit from higher GPU utilization and fewer KV-cache reloads.
  • Throughput trade-offs: Chunk size 64 almost matches baseline throughput despite being 5× slower in prefill.Decode-throughput gains compensate for the prefill slowdown.
  • Throughput trade-offs: Chunk size 128 yields up to 1.16× higher throughput despite prefill being more than 2× slower than baseline.The improvement is mainly due to piggybacking more decodes.
  • Throughput trade-offs: Throughput improvement is higher when chunk size is a multiple of 128; chunk size 256 outperforms 320.The figure shows this tile-quantization effect across chunk sizes from 64 to 512.

6 Discussion

The discussion identifies unresolved deployment and configuration challenges despite demonstrating SARATHI across multiple models and hardware configurations.

  • Deployment challenges: SARATHI optimizes throughput scheduling, while real-world serving also requires latency, queuing-delay, and fairness optimization.Meeting these goals requires revisiting scheduling policies.
  • Configuration challenges: Optimal chunk size remains future work because it depends on hardware, model characteristics, sequence length, token composition, and potentially unknown P:D ratios.The paper reports appropriate chunk sizes for given P:D ratios but does not solve general selection.
  • Scope boundaries: The evaluation assumes equal prefill and decode token counts within a batch except in simulation experiments, whereas real requests can vary significantly.This assumption limits direct coverage of heterogeneous request lengths.
  • Scope boundaries: The study focuses on sequence lengths up to 3K and P:D ratios from 1-200.The authors consider these ranges representative of many real-world deployments.

7 Related Work

Related work spans systems optimizations and model innovations, including memory management, attention efficiency, quantization, and mixture-of-experts architectures.

  • Organization: Prior systems work addresses optimization along the dimensions of systems design and model innovation.The paper organizes related work into these two categories.
  • Systems optimizations: vLLM replaces conservative KV-cache preallocation with incremental allocation motivated by virtual memory.This improves batch-size support for autoregressive decoding with unknown output lengths.
  • Systems optimizations: Prior optimizations can complement SARATHI by enabling longer sequences, larger batch sizes, or both.The paper cites optimized attention implementations and dynamic memory allocation as examples.
  • Model innovations: Model innovations include multi-query attention, quantization, and mixture-of-experts models.These approaches respectively reduce KV-cache size, compress model sizes, or reduce activated parameters.

8 Conclusion

SARATHI targets decode inefficiency and pipeline bubbles by combining chunked-prefills with decode-maximal batching, improving end-to-end throughput across models and hardware configurations.

  • Conclusion: SARATHI addresses low GPU utilization in memory-bound decodes and pipeline bubbles from inconsistent prefill and decode times.These problems produce micro-batch imbalance during pipeline-parallel inference.
  • Conclusion: Decode-maximal batching piggybacks decodes with prefills, making the combined workload compute bound.Chunked-prefills supplies more prefill chunks and creates uniform work units that reduce pipeline bubbles.
  • Conclusion: SARATHI improves end-to-end throughput across multiple models and hardware configurations.The conclusion states this as the paper’s overall demonstrated outcome.
Loading 2308.16369v1…