Source-linked AI summary
POD-Attention: Unlocking Full Prefill-Decode Overlap for Faster LLM Inference
Aditya K Kamath, Ramya Prabhu, Jayashree Mohan, Simon Peter, Ramachandran Ramjee, Ashish Panwar
TL;DR
LLM inference needs both compute and memory bandwidth, but hybrid-batch attention kernels handle prefill and decode separately. POD-Attention runs both phases concurrently with SM-aware scheduling; attention speeds up and serving throughput rises while latency falls.
Problem
Hybrid batching combines compute-bound prefill and memory-bound decode, but attention kernels handle the phases separately and leave GPU resources underused.
Source: §1 · pp. 1–2Method
Results
Attention computation is up to 59% faster (mean 28%) than independently optimized kernels; serving throughput improves up to 22%, with reduced latency over Sarathi-Serve and vLLM.
Source: §7 · p. 12Takeaways & Limitations
The paper concludes POD-Attention enables hybrid-batching-based LLM inference with higher serving throughput and lower latency.
Source: §7 · p. 12Warp-parallel fusion is vulnerable to stragglers across varying batch sizes and context lengths, requiring extensive tuning.
Source: §3.1 · p. 5
Abstract
from arXiv · showhide
Each request in LLM inference goes through two phases: compute-bound prefill and memory-bandwidth-bound decode. To improve GPU utilization, recent systems use hybrid batching that combines the prefill and decode phases of different requests into the same batch. This approach optimizes linear operations but remains inefficient for attention computation because existing attention kernels specialize execution independently for the prefill and decode phases. In this paper, we present POD-Attention - the first GPU kernel that efficiently computes attention for hybrid batches. POD-Attention aims to maximize the utilization of both compute and memory bandwidth by carefully allocating the GPU's resources such that prefill and decode operations happen concurrently on the same multiprocessor. POD-Attention speeds up attention computation by up to $59\%$ (mean $28\%$), enabling higher throughput and lower latency LLM inference compared to the use of independently optimized prefill and decode attention kernels.
1 Introduction
LLM serving depends on GPU utilization, but prefill and decode demand compute and memory bandwidth at different times. POD-Attention targets underused resources in hybrid-batch attention by running both phases concurrently.
1 Introduction
- Prefill is compute-bound and decode is memory-bound, so efficient LLM inference requires using both compute and memory bandwidth despite the phases occurring at different times.GPU resource utilization is important to overall deployment efficiency.Source: §1 · p. 1
- Hybrid batching combines different requests’ prefill and decode inputs, avoiding separate model-weight fetches and allowing long prompts to be chunked to reduce tail latency.Schedulers can combine ongoing decodes with a new prefill chunk each iteration.Source: §1 · p. 2
- Existing attention kernels optimize prefill and decode separately, leaving prefill memory-bandwidth utilization often below 5% and decode compute utilization under 10%.The underutilization is especially consequential in hybrid batches, where the two kernels execute one after the other.Source: §1 · p. 2
- POD-Attention fuses prefill and decode attention with CTA-parallel execution and SM-aware software scheduling to use compute and memory resources together.The paper identifies straggler threads, synchronization barriers, and lack of guaranteed SM-level co-location as limitations of existing fusion techniques.Source: §1 · p. 2
- Attention computation is up to 59% faster (mean 28%) than FlashAttention and FlashInfer kernels; end-to-end throughput improves up to 22%, with lower latency metrics.The latency metrics named are time-to-first-token, time-between-tokens, and end-to-end request execution latency.Source: §1 · p. 2
2 Background and Motivation
LLM inference pairs compute-bound prefill with memory-bound decode, creating a throughput-latency tradeoff that hybrid batching addresses by combining requests’ phases. As context length grows, attention becomes a major bottleneck, while GPU execution offers several ways to run operations concurrently.
2.1 Large Language Model (LLM) Inference
- Prefill processes prompt tokens in parallel and is compute-bound, while autoregressive decode is memory-bound; their latency measures are TTFT and TBT, respectively.Prefill latency is generally higher than decode-iteration latency.Source: §2.1 · p. 2
- Chunked-prefill hybrid batching schedules one prompt chunk alongside ongoing decodes, avoiding generation stalls while supporting larger batches.A hybrid batch commonly contains one prefill chunk and multiple decodes, requiring both compute and memory bandwidth.Source: §2.1 · p. 3
- At 16K context, attention can exceed 60% of inference time, and separate prefill and decode kernels leave compute idle during decode.Prefill and decode attention execute immediately one after the other in hybrid batches.Source: §2.1 · p. 3
2.2 GPU Execution Model
- GPU execution is organized around SMs: warps form CTAs that share resources and are assigned to SMs when resources are available.Operations within one stream are serialized, while different streams can execute in parallel in any order.Source: §2.2 · p. 4
3 A Case Study on Concurrent Execution
Concurrent execution methods can improve attention runtime, but their gains are constrained by wave quantization, stragglers, barriers, and lack of guaranteed SM-level co-location. The case study finds that POD-Attention outperforms the tested stream- and warp-parallel alternatives.
3 A Case Study on Concurrent Execution
- Passing prefill and decode inputs together to an existing attention kernel is counter-productive and slower than serial execution.The case study evaluates concurrent GPU execution methods to identify limitations motivating a specialized attention kernel.Source: §3 · p. 4
3.1 Methods of Concurrent Execution
- Kernel-parallel and CTA-parallel execution do not guarantee SM-level co-location, while warp-parallel fusion guarantees it but can suffer from straggler delays.Intra-thread fusion can offer substantial overlap, but attention’s CTA-level synchronization barriers limit it.Source: §3.1 · pp. 4–5
3.2 Analysis of Readily Available Methods
- POD-Attention outperforms both FA_Streams and FA_HFuse, while each improves over serial execution in the tested attention workloads.FA_Streams gains up to 20% over FA_Serial when serial execution suffers from wave quantization; FA_HFuse degrades in later, prefill-dominated chunks due to stragglers.Source: §3.2 · p. 5
3.3 Analysis of Other Methods
- CTA- and kernel-parallel fusion improve average runtime over serial execution by 3% and 7%, respectively, while intra-thread fusion improves it by 13%.Intra-thread fusion’s benefit is limited because synchronization barriers hinder concurrent execution.Source: §3.3 · p. 6
4 POD-Attention
POD-Attention fuses prefill and decode attention in one GPU kernel, using SM-aware CTA scheduling to co-locate both operations. Its optimizations balance their different compute, bandwidth, and shared-memory demands.
4.1 SM-aware CTA Scheduling
- CTA-parallel fusion lets prefill and decode CTAs run independently, while SM-aware scheduling assigns their work after dispatch to ensure both operations share an SM.Each CTA checks its SM and the other CTAs running there before choosing an operation.Source: §4 · p. 6, §4.1 · p. 6
4.2.1 Tile Sizes.
- A QSL tile length of 16 reduces decode compute utilization to about 10%, freeing tensor cores for co-located prefills without harming decode performance at large batch sizes.Larger decode tiles create redundant computation that competes with prefill for tensor cores.Source: §4.2.1 · p. 7
4.2.2 Concurrent CTAs per SM.
- POD-Attention uses 2 CTAs per SM for prefill-dominant batches and 4 otherwise, selecting the configuration at runtime; 8 CTAs per SM usually underperforms.Fewer CTAs give prefill more shared memory for larger tiles, while more CTAs allow finer-grained scheduling for decode.Source: §4.2.2 · p. 7
4.2.3 Virtual Decode CTAs.
- Virtual decode CTAs divide a four-warp decode CTA into one-warp units, balancing shared-memory use against prefill CTAs.Each virtual CTA uses one quarter of the original decode CTA's shared memory.Source: §4.2.3 · p. 7
- Chunked-prefill splits are capped at two full waves to provide parallelism without excessive memory-bandwidth contention with concurrent decode CTAs.The limit is determined empirically.Source: §4.2.3 · pp. 7–8
4.3 Implementing CTA-parallel Fusion
- The fused wrapper remaps CTA IDs to call prefill or decode device functions, hand-tuning shared-memory use and adapting decode barriers for virtual CTAs.The kernel launches with enough shared memory for whichever operation requires more.Source: §4.3 · p. 8
4.4 Discussion on Alternative Implementations
- Persistent-thread execution with SM-aware scheduling performs on par with CTA-parallel fusion, while still requiring scheduling to guarantee same-SM co-location.Persistent CTAs pull prefill or decode work as needed.Source: §4.4 · p. 8
- MPS and MIG are inapplicable because hybrid batching combines prefill and decode operations within a single process.Source: §4.4 · p. 8
5 Evaluation
The evaluation measures attention performance, offline throughput, and online latency across three models and multiple workloads. POD-Attention improves attention and serving outcomes, while results vary with workload composition and configuration.
5.1 Evaluating Attention Computation
- 59% peak and 28% mean attention speedup exceed all alternatives; energy use falls up to 35% (mean 20.5%) versus FA_Serial.The speedup distribution is shown in Figure 11; energy savings are largely proportional to runtime reduction.Source: Figure 11 · p. 9, §5.1 · p. 9
5.2 Evaluating Throughput in Offline Inference
- Sarathi+POD achieves the highest offline throughput: 22%, 20%, and 19% above Sarathi, and 27%, 13%, and 12% above vLLM across the three models.Figure 12 presents serving throughput in offline inference.Source: Figure 12 · p. 9, §5.2 · p. 9
5.3 Evaluating Latency in Online Inference
- Sarathi+POD reduces Sarathi's median TTFT to 7.5 and 11.74 seconds at higher load, lowers tail TBT by 10–20%, and improves request latency over Sarathi across cases.It also outperforms vLLM on request latency in many cases; Sarathi+POD chunk size can trade lower TTFT for higher TBT.Source: §5.3.1 · p. 10, §5.3.2 · p. 10, §5.3.3 · pp. 10–11
5.4.1 CTAs per SM.
- For long contexts, 2 CTAs per SM performs better; as contexts shorten and decode cost grows, 4 CTAs per SM performs better.More CTAs allow more decode CTAs to run alongside fewer prefill CTAs.Source: §5.4.1 · p. 11
5.4.2 Scheduling Policy.
- Proportional CTA allocation performs up to 14% better than 50:50 allocation as load increases, spreading less frequent operations to improve overlap and reduce contention.Source: §5.4.2 · p. 11, Figure 14 · p. 11
5.4.3 Limiting Prefill Splits.
- Limiting prefill splits nearly doubles POD-Attention's speedup over FA_Serial in the last four chunks of a 16K prompt co-running with 64 same-context decode requests.The comparison is for Llama-3-8B with chunk size 512.Source: §5.4.3 · p. 11, Table 8 · p. 11
5.4.4 Sensitivity to Workload.
- For Llama-3-8B requests of about 16.5K tokens, Sarathi+POD's peak gains occur at P:D ratios of 12–18, where most batches are hybrid.Ratios below 12 yield many decode-only iterations; ratios above 18 yield many prefill-only iterations.Source: §5.4.4 · p. 11, Figure 15 · p. 12
6 Related Work
POD-Attention jointly fuses prefill and decode attention, unlike prior approaches that optimize the phases separately. Related work spans specialized attention kernels, kernel fusion, LLM serving systems, and complementary resource-utilization strategies.
Optimizing Attention Computation
- POD-Attention jointly optimizes and fuses prefill and decode attention, while prior attention kernels handle the phases separately.The section reviews FlashAttention, FA-2, FlashDecoding, FlashDecoding++, LeanAttention, and FlashInfer.Source: §6 · p. 12
Operation Fusion
- ISPA profiles fixed CTA counts for concurrent execution, while HFuse fuses operations warp-parallel; POD-Attention leverages SM-centric scheduling.Elastic kernels restrict resources but do not guarantee intra-SM co-location; fixed profiles can make load balancing difficult as attention inputs change.Source: §6 · p. 12
Optimizing LLM Inference
- Orca uses iteration-level scheduling, PagedAttention and vAttention address dynamic memory management, and Sarathi-Serve uses chunked prefills.Splitwise, DistServe, and TetriInfer disaggregate prefill and decode across GPU nodes; other work overlaps compute with communication.Source: §6 · p. 12
Comparison with NanoFlow
- NanoFlow overlaps operations through CUDA-stream nano-batches, whereas POD-Attention fuses prefill and decode attention within a batch.NanoFlow benefits from large batches; POD-Attention targets long-context scenarios where attention consumes significant time.Source: §6 · p. 12
7 Conclusion
POD-Attention computes prefill and decode attention in parallel to use GPU compute and memory bandwidth simultaneously. It improves attention performance and end-to-end serving outcomes against the stated comparisons.
7 Conclusion
- Up to 59% (mean 28%) faster attention computation than independently optimized prefill and decode attention kernels.The kernel computes prefill and decode attention in parallel to use compute and memory bandwidth simultaneously.Source: §7 · p. 12
- Up to 22% higher end-to-end serving throughput with POD-Attention.The comparison is against the use of independently optimized prefill and decode attention kernels.Source: §7 · p. 12
- Latency is significantly reduced over Sarathi-Serve and vLLM.The conclusion reports this outcome alongside improved end-to-end serving throughput.Source: §7 · p. 12
A.1 Abstract
POD-Attention overlaps prefill and decode attention operations for LLMs. It is built on FlashAttention kernels and integrated with Sarathi-Serve.
A.1 Abstract
- POD-Attention is a GPU kernel that overlaps prefill and decode attention operations for LLMs.It is intended for hybrid-batching-based LLM inference.Source: §A.1 · p. 13
- POD-Attention is built on FlashAttention kernels, version 2.6.1.Source: §A.1 · p. 13
- POD-Attention is integrated with Sarathi-Serve, a hybrid-batching-based LLM inference scheduler.Source: §A.1 · p. 13
A.2 Artifact check-list (meta-information)
The artifact checklist specifies the models, dataset, and runtime environment used, along with hardware and preparation and experiment-time estimates.
A.2 Artifact check-list (meta-information)
- The listed models are Llama-2-7B, Llama-3-8B, and Yi-6B, and the dataset is arXiv-Summarization.Source: §A.2 · p. 13
- Experiments use one to two NVIDIA A100 80 GB GPUs and an x86 machine, with Ubuntu 22.04, CUDA 12.4, Python 3.12, and PyTorch 2.4.Compilation lists CUDA 12.4 and GCC 11.4.Source: §A.2 · p. 13
- Workflow preparation is estimated at one minute with Docker or one to two hours from source; experiments take approximately 18 hours.The artifact is marked publicly available and lists DOI 10.5281/zenodo.14770841.Source: §A.2 · p. 13
A.3.2 Hardware dependencies.
The artifact documents hardware, software, data, and model requirements, installation options, experiment workflows, and expected evaluation outputs.
A.3.2 Hardware dependencies.
- The artifact requires an x86 machine with two 80GB NVIDIA A100 GPUs; with one GPU, all experiments except Table 6 and specified Figure 12 results remain runnable.The excluded Figure 12 results are for Llama-2-7B and Llama-3-8B.Source: §A.3.2 · p. 13
A.3.3 Software dependencies.
- POD-Attention was tested on Ubuntu 22.04, and installing the artifact resolves its other software dependencies.Source: §A.3.3 · p. 13
A.3.4 Data sets.
- Some experiments use a subset of the arXiv-Summarization dataset stored in the artifact's traces/ folder.Source: §A.3.4 · p. 13
A.3.5 Models.
- The artifact evaluates Yi-6B, Llama-2-7B, and Llama-3-8B; accessing Llama-3-8B requires a Hugging Face private token.Source: §A.3.5 · p. 13
A.4 Installation
- Installation is available through the recommended Docker option or a manual setup.The provided Docker image includes all dependencies pre-installed.Source: §A.4 · p. 13, §A.4.1 · p. 13
A.4.2 Manual installation.
- Manual installation uses Anaconda for compatible CUDA, Python, and PyTorch versions and can take up to two hours.The instructions specify CUDA Toolkit 12.4.0 and FlashInfer 0.1.5.Source: §A.4.2 · p. 13
A.5 Experiment workflow
- The artifact provides Makefile commands for attention-performance, end-to-end LLM, and kernel-fusion experiments; most need one GPU, while Table 6 and specified Figure 12 results need two.Those Figure 12 results are for Llama-2-7B and Llama-3-8B.Source: §A.5 · p. 13
A.6 Evaluation and expected results
- Evaluation outputs include raw numbers, logs, graphs, and CSV tables, and runtime variation may occur while general trends are expected to hold.Source: §A.6 · p. 13