POD-Attention: 解锁完整的预填充-解码重叠以实现更快的LLM推理
机器学习
2025-02-18 v2 分布式、并行与集群计算
摘要
LLM推理中的每个请求都经历两个阶段：计算受限的预填充和内存带宽受限的解码。为了提高GPU利用率，最近的系统使用混合批处理，将不同请求的预填充和解码阶段合并到同一批次中。这种方法优化了线性操作，但对于注意力计算仍然效率低下，因为现有的注意力内核专门独立执行预填充和解码阶段。在本文中，我们提出了POD-Attention——第一个能高效计算混合批次注意力的GPU内核。POD-Attention旨在通过精心分配GPU资源，使预填充和解码操作在同一多处理器上并发执行，从而最大化计算和内存带宽的利用率。与使用独立优化的预填充和解码注意力内核相比，POD-Attention将注意力计算速度提升高达（平均），从而实现更高吞吐量和更低延迟的LLM推理。
引用
@article{arxiv.2410.18038,
title = {POD-Attention: 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},
journal= {arXiv preprint arXiv:2410.18038},
year = {2025}
}
备注
Proceedings of the 30th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2 (ASPLOS '25), March 30 - April 3, 2025, Rotterdam, Netherlands