Abstract
Each request in LLM inference goes through two phases: compute-bound prefilland memory-bandwidth-bound decode. To improve GPU utilization, recent systemsuse hybrid batching that combines the prefill and decode phases of differentrequests into the same batch. Hybrid batching works well for linear operationsas it amortizes the cost of loading model weights from HBM. However, attentioncomputation in hybrid batches remains inefficient because existing attentionkernels are optimized for either prefill or decode. In this paper, we present POD-Attention -- the first GPU kernel thatefficiently computes attention for hybrid batches. POD-Attention aims tomaximize the utilization of both compute and memory bandwidth by carefullyallocating the GPU's resources such that prefill and decode operations happenconcurrently on the same multiprocessor. We integrate POD-Attention in astate-of-the-art LLM inference scheduler Sarathi-Serve. POD-Attention speeds upattention computation by up to 75% (mean 28%) and increases LLM servingthroughput by up to 22% in offline inference. In online inference,POD-Attention enables lower time-to-first-token (TTFT), time-between-tokens(TBT), and request execution latency versus Sarathi-Serve.