BASS: Batched Attention-optimized Speculative Sampling
Authors: name: Haifeng Qian; name: Sujan Kumar Gonugondla; name: Sungsoo Ha; name: Mingyue Shang; name: Sanjay Krishna Gouda; name: Ramesh Nallapati; name: Sudipta Sengupta; name: Xiaofei Ma; name: Anoop Deoras
Published: 2024-04-24 09:57:11

Abstract: 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.