BASS: Muestreo Especulativo Optimizado con Atención por Lotes

BASS: Batched Attention-optimized Speculative Sampling

April 24, 2024

Autores: Haifeng Qian, Sujan Kumar Gonugondla, Sungsoo Ha, Mingyue Shang, Sanjay Krishna Gouda, Ramesh Nallapati, Sudipta Sengupta, Xiaofei Ma, Anoop Deoras

cs.AI

Resumen

La decodificación especulativa ha surgido como un método poderoso para mejorar la latencia y el rendimiento al alojar modelos de lenguaje grandes. Sin embargo, la mayoría de las implementaciones existentes se centran en generar una única secuencia. Las aplicaciones de IA generativa del mundo real a menudo requieren múltiples respuestas, y cómo realizar la decodificación especulativa en un entorno por lotes mientras se preservan sus beneficios de latencia plantea desafíos no triviales. Este artículo describe un sistema de decodificación especulativa por lotes que establece un nuevo estado del arte en la latencia de generación de múltiples secuencias y que demuestra una utilización superior de la GPU, así como la calidad de las generaciones dentro de un límite de tiempo. Por ejemplo, para un modelo de 7.8B de tamaño en una sola GPU A100 y con un tamaño de lote de 8, cada secuencia se genera a una velocidad promedio de 5.8 ms por token, con un rendimiento general de 1.1K tokens por segundo. Estos resultados representan una latencia de vanguardia y una aceleración de 2.15X sobre la decodificación regular optimizada. Dentro de un límite de tiempo en el que la decodificación regular no finaliza, nuestro sistema es capaz de generar secuencias con un HumanEval Pass@First del 43% y un Pass@All del 61%, superando ampliamente lo que es factible con la decodificación especulativa de una sola secuencia. Nuestra utilización máxima de la GPU durante la decodificación alcanza hasta un 15.8%, más de 3X la más alta de la decodificación regular y alrededor de 10X la de la decodificación especulativa de una sola secuencia.

English

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.