一句话:首个能在混合批次上并行执行 prefill 与 decode 注意力计算的 GPU kernel,最高提速 59%。
LLM 推理中每个请求都要经过计算密集的 prefill 阶段与访存密集的 decode 阶段。为提升 GPU 利用率,近期系统采用混合批处理,把不同请求的 prefill 与 decode 放进同一批次,但既有的注意力 kernel 仍分别为两个阶段独立特化执行,导致注意力计算在混合批次下效率低下。本文提出 POD-Attention:第一个能高效计算混合批次注意力的 GPU kernel。它通过精细调配 GPU 资源,让 prefill 与 decode 操作在同一个流多处理器上并发执行,从而同时最大化计算与内存带宽的利用率。评估表明,相比分别优化的 prefill 与 decode 注意力 kernel,POD-Attention 将注意力计算加速最高 59%(平均 28%),带来更高的吞吐与更低的延迟。
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 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.
DOI 原文 · arXiv · PDF(开放获取) · DBLP
@inproceedings{abs-2410-18038,
title = {POD-A
<scp>ttention</scp>
: Unlocking Full Prefill-Decode Overlap for Faster LLM Inference},
author = {Aditya K. Kamath and Ramya Prabhu and Jayashree Mohan and Simon Peter and Ramachandran Ramjee and Ashish Panwar},
booktitle = {ASPLOS},
year = {2025},
doi = {10.1145/3676641.3715996}
}