BASS : Batched Attention-optimized Speculative Sampling (2024)

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.15 \times 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 3 \times the highest of that of regular decoding and around 10 \times of single-sequence speculative decoding.

## introduction in recent years, generative large language models (llms) have rapidly gained popularity due to their ability to generalize across a wide variety of tasks. these models are increasingly deployed commercially for applications such as coding assistants, writing aids, conversational agents, search and summarization tools and more. the accuracy ## batched attention-optimized speculative sampling batched attention-optimized speculative sampling (bass) extends speculative decoding by enabling batch processing across multiple sequences. while speculative sampling improves gpu utilization for a single sequence, parallelism is limited to the small number of draft tokens per sequence (typically 5-10). batching sequences with speculative decoding can further maximize gpu usage. to fully realize these benefits, specialized tensor manipulations and cuda kernels are required. ## experiments in this section we demonstrate the benefits of bass over batched auto-regressive regular decoding (rd) and single-sequence speculative decoding. ## related work efficient inference of llms has been a popular research topic in recent years. model quantization techniques @xcite @xcite sparse attention techniques @xcite limit the number of tokens to attend to in order to reduce the complexity of attention layers, and thereby extend the maximum allowable sequence length. since its introduction, speculative decoding has seen numerous variations and improvements. some proposals take a draft-model-free approach, by using an n-gram model to predict draft tokens @xcite , or by using the model embedding to predict drafts @xcite . specinfer @xcite uses a draft tree to generate and organize multiple drafts for the mainmodel sequence in order to maximize the number of tokens accepted per step. @xcite study the relation between batch size and the optimal fixed draft length for max throughput; it is however based on a primitive prototype implementation: rejected draft tokens are masked rather than discarded, which achieve sequence lengths that are uniform across a batch yet are unnecessarily large and inefficient. the above works on speculative decoding are orthogonal to the discussions and ideas in this paper and can be combined. the conclusion presented in @xcite may change with the kernel implementations in this paper. ## conclusion this paper presents batched attention-optimized speculative sampling (bass), a system that ad-vances the state of the art in fast multi-sequence generation by llms. by addressing the unique challenges of extending speculative decoding to batched inference without sacrificing latency, we demonstrate superior latency, gpu utilization as well as accuracy of generations within a time limit.