BASS: Batched Attention-optimized Speculative Sampling

TL;DR

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 is described.

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.

Figures (6)

• Figure 1: Comparing latency and GPU utilization of auto-regressive regular decoding (RD), single-sequence speculative decoding (SD) and our BASS method on two models. RD and BASS are measured with exponentially increasing batch sizes (BS).
• Figure 2: (a) Inference steps in regular decoding of an LLM. (b) Operations in multi-head attention.
• Figure 3: Standard speculative decoding. The draft model (Draft M) generates $k$ draft tokens auto-regressively, which are then processed by the main model (M) in parallel to verify correctness.
• Figure 4: Attention calculation in BASS: (a) Attention compute flow, (b) BASS-PAD launches one kernel for QK GEMM and one kernel for PV GEMM by padding the $K$, $V$ and $P$ tensors to the maximum sequence length across the batch, and (c) BASS-SPLIT launches one kernel per sequence and thereby accommodates variable sequence lengths.
• Figure 5: A 7.8B code model's accuracy on HumanEval with BASS, within a time budget of 2.5 seconds. $t$ is temperature.