BASS: Amostragem Especulativa com Atenção Otimizada em 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

Resumo

A decodificação especulativa surgiu como um método poderoso para melhorar a latência e a taxa de transferência na hospedagem de grandes modelos de linguagem. No entanto, a maioria das implementações existentes foca na geração de uma única sequência. Aplicações reais de IA generativa frequentemente exigem múltiplas respostas, e como realizar a decodificação especulativa em um cenário de processamento em lote, preservando seus benefícios de latência, apresenta desafios não triviais. Este artigo descreve um sistema de decodificação especulativa em lote que estabelece um novo estado da arte em latência de geração de múltiplas sequências e demonstra superior utilização da GPU, bem como qualidade de gerações dentro de um limite de tempo. Por exemplo, para um modelo de 7,8B em uma única GPU A100 e com um tamanho de lote de 8, cada sequência é gerada a uma velocidade média de 5,8ms por token, com uma taxa de transferência geral de 1,1K tokens por segundo. Esses resultados representam a latência mais avançada e uma aceleração de 2,15X em relação à decodificação regular otimizada. Dentro de um limite de tempo em que a decodificação regular não consegue terminar, nosso sistema é capaz de gerar sequências com HumanEval Pass@First de 43% e Pass@All de 61%, superando em muito o que é viável com a decodificação especulativa de sequência única. Nossa utilização máxima da GPU durante a decodificação chega a 15,8%, mais de 3X a maior utilização da decodificação regular e cerca de 10X a da decodificação especulativa de sequência única.

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.