BASS : Échantillonnage Spéculatif Optimisé par Attention par Lots

BASS: Batched Attention-optimized Speculative Sampling

April 24, 2024

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

cs.AI

Résumé

Le décodage spéculatif s'est imposé comme une méthode puissante pour améliorer la latence et le débit lors de l'hébergement de grands modèles de langage. Cependant, la plupart des implémentations existantes se concentrent sur la génération d'une seule séquence. Les applications réelles d'IA générative nécessitent souvent plusieurs réponses, et la manière d'effectuer un décodage spéculatif dans un contexte par lots tout en préservant ses avantages en termes de latence pose des défis non triviaux. Cet article décrit un système de décodage spéculatif par lots qui établit un nouvel état de l'art en matière de latence pour la génération de séquences multiples et démontre une utilisation supérieure des GPU ainsi qu'une qualité de génération optimale dans un budget de temps donné. Par exemple, pour un modèle de 7,8 milliards de paramètres sur un seul GPU A100 et avec une taille de lot de 8, chaque séquence est générée à une vitesse moyenne de 5,8 ms par token, avec un débit global de 1,1K tokens par seconde. Ces résultats représentent une latence de pointe et une accélération de 2,15X par rapport à un décodage régulier optimisé. Dans un budget de temps où le décodage régulier ne parvient pas à terminer, notre système est capable de générer des séquences avec un HumanEval Pass@First de 43 % et un Pass@All de 61 %, dépassant largement ce qui est réalisable avec un décodage spéculatif mono-séquence. Notre utilisation maximale des GPU pendant le décodage atteint jusqu'à 15,8 %, soit plus de 3X le maximum du décodage régulier et environ 10X celui du décodage spéculatif mono-séquence.

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.