POD-Attention: Unlocking Full Prefill-Decode Overlap for Faster LLM Inference
Fuente:
arXiv
Saved in:
| Main Authors: | , , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2024
|
| Subjects: | |
| Online Access: | |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866929716590542848 |
|---|---|
| author | Kamath, Aditya K Prabhu, Ramya Mohan, Jayashree Peter, Simon Ramjee, Ramachandran Panwar, Ashish |
| author_facet | Kamath, Aditya K Prabhu, Ramya Mohan, Jayashree Peter, Simon Ramjee, Ramachandran Panwar, Ashish |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2410_18038 |
| institution | arXiv |
| publishDate | 2024 |
| record_format | arxiv |
| spellingShingle | POD-Attention: Unlocking Full Prefill-Decode Overlap for Faster LLM Inference Kamath, Aditya K Prabhu, Ramya Mohan, Jayashree Peter, Simon Ramjee, Ramachandran Panwar, Ashish Machine Learning Distributed, Parallel, and Cluster Computing I.2.7; C.1.4 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. |
| title | POD-Attention: Unlocking Full Prefill-Decode Overlap for Faster LLM Inference |
| topic | Machine Learning Distributed, Parallel, and Cluster Computing I.2.7; C.1.4 |
| url | https://arxiv.org/abs/2410.18038 |