paperbot · PL 论文追踪

全部论文 待读 收藏

P2 POD-A <scp>ttention</scp> : Unlocking Full Prefill-Decode Overlap for Faster LLM Inference

POD-Attention:解锁完整的 Prefill-Decode 重叠以实现更快的 LLM 推理
ASPLOS Proceedings of the 30th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 22025匹配度 2/10 · 该论文关注 LLM 推理的 GPU kernel 优化,属于机器学习系统性能优化,与用户研究方向(静态分析、形式化验证、浮点/贝叶斯程序分析及其与 LLM 的交叉)无直接关联。虽然涉及 LLM 系统,但未涉及程序分析或形式化方法。引用 96 inbox
Aditya K. Kamath, Ramya Prabhu, Jayashree Mohan, Simon Peter, Ramachandran Ramjee, Ashish Panwar
crossref: 2026-09-12dblp: 2026-09-12arxiv: 2026-09-12

AI 速览

一句话:首个能在混合批次上并行执行 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%),带来更高的吞吐与更低的延迟。

问题
混合批处理中的注意力计算仍沿用分别针对 prefill 与 decode 特化的 kernel,无法同时利用计算与内存带宽,成为 LLM 推理效率瓶颈。
方法
设计 POD-Attention,通过在单个流多处理器上协同调度 prefill 与 decode 的注意力计算,使两者并发执行并平衡计算与内存资源占用。
评估
与独立优化的 prefill 与 decode 注意力 kernel 做对比,测量注意力计算加速比以及端到端吞吐与延迟。
原文摘要(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 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

BibTeX
@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}
}