提出新型注意力核POD-Attention，实现预填充与解码并行计算，显著提升大模型推理速度。
POD-Attention: Unlocking Full Prefill-Decode Overlap for Faster LLM Inference
- 设计统一注意力核，让预填充和解码在同一个计算单元上并发执行
- 在混合批处理下，注意力计算速度最高提升59%（平均28%）
- 适合追求高吞吐、低延迟的LLM部署场景
大模型推理中每个请求都经历计算密集型预填充阶段和内存带宽受限的解码阶段。为提升GPU利用率，现有系统采用混合批处理，将不同请求的预填充和解码合并到同一批次。该方法虽优化了线性运算，但对注意力计算仍效率低下，因现有注意力内核分别独立优化预填充和解码阶段。本文提出POD-Attention——首个可高效计算混合批次注意力的GPU内核。通过精细分配GPU资源，使预填充与解码操作在同个流多处理器上并发运行，最大化利用计算与内存带宽。实验表明，相较于独立优化的预填充与解码注意力内核，POD-Attention在注意力计算上最快提速59%（平均28%），显著提升大模型推理的吞吐量与降低延迟。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。