POD-Attention: Unlocking Full Prefill-Decode Overlap for Faster LLM Inference

TL;DR

POD-Attention introduces a novel GPU kernel that efficiently overlaps prefill and decode phases in LLM inference, significantly boosting GPU utilization and reducing latency for faster, more efficient large language model deployment.

Contribution

This paper presents the first GPU kernel capable of concurrently computing attention for hybrid batches, optimizing resource utilization during LLM inference.

Findings

Speeds up attention computation by up to 59%.

Achieves mean speedup of 28% over traditional kernels.

Enables higher throughput and lower latency in LLM inference.

Abstract

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.