BASS: Batched Attention-optimized Speculative Sampling
Fuente:
arXiv
Salvato in:
| Autori principali: | , , , , , , , , |
|---|---|
| Natura: | Preprint |
| Pubblicazione: |
2024
|
| Soggetti: | |
| Accesso online: | |
| Tags: |
Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
|
| _version_ | 1866911933253287936 |
|---|---|
| author | Qian, Haifeng Gonugondla, Sujan Kumar Ha, Sungsoo Shang, Mingyue Gouda, Sanjay Krishna Nallapati, Ramesh Sengupta, Sudipta Ma, Xiaofei Deoras, Anoop |
| author_facet | Qian, Haifeng Gonugondla, Sujan Kumar Ha, Sungsoo Shang, Mingyue Gouda, Sanjay Krishna Nallapati, Ramesh Sengupta, Sudipta Ma, Xiaofei Deoras, Anoop |
| contents | Speculative decoding has emerged as a powerful method to improve latency and throughput in hosting large language models. However, most existing implementations focus on generating a single sequence. Real-world generative AI applications often require multiple responses and how to perform speculative decoding in a batched setting while preserving its latency benefits poses non-trivial challenges. This paper describes a system of batched speculative decoding that sets a new state of the art in multi-sequence generation latency and that demonstrates superior GPU utilization as well as quality of generations within a time budget. For example, for a 7.8B-size model on a single A100 GPU and with a batch size of 8, each sequence is generated at an average speed of 5.8ms per token, the overall throughput being 1.1K tokens per second. These results represent state-of-the-art latency and a 2.15X speed-up over optimized regular decoding. Within a time budget that regular decoding does not finish, our system is able to generate sequences with HumanEval Pass@First of 43% and Pass@All of 61%, far exceeding what's feasible with single-sequence speculative decoding. Our peak GPU utilization during decoding reaches as high as 15.8%, more than 3X the highest of that of regular decoding and around 10X of single-sequence speculative decoding. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2404_15778 |
| institution | arXiv |
| publishDate | 2024 |
| record_format | arxiv |
| spellingShingle | BASS: Batched Attention-optimized Speculative Sampling Qian, Haifeng Gonugondla, Sujan Kumar Ha, Sungsoo Shang, Mingyue Gouda, Sanjay Krishna Nallapati, Ramesh Sengupta, Sudipta Ma, Xiaofei Deoras, Anoop Machine Learning Computation and Language Speculative decoding has emerged as a powerful method to improve latency and throughput in hosting large language models. However, most existing implementations focus on generating a single sequence. Real-world generative AI applications often require multiple responses and how to perform speculative decoding in a batched setting while preserving its latency benefits poses non-trivial challenges. This paper describes a system of batched speculative decoding that sets a new state of the art in multi-sequence generation latency and that demonstrates superior GPU utilization as well as quality of generations within a time budget. For example, for a 7.8B-size model on a single A100 GPU and with a batch size of 8, each sequence is generated at an average speed of 5.8ms per token, the overall throughput being 1.1K tokens per second. These results represent state-of-the-art latency and a 2.15X speed-up over optimized regular decoding. Within a time budget that regular decoding does not finish, our system is able to generate sequences with HumanEval Pass@First of 43% and Pass@All of 61%, far exceeding what's feasible with single-sequence speculative decoding. Our peak GPU utilization during decoding reaches as high as 15.8%, more than 3X the highest of that of regular decoding and around 10X of single-sequence speculative decoding. |
| title | BASS: Batched Attention-optimized Speculative Sampling |
| topic | Machine Learning Computation and Language |
| url | https://arxiv.org/abs/2404.15778 |