BASS: Stapelverarbeitungsoptimierte spekulative Stichprobenziehung

BASS: Batched Attention-optimized Speculative Sampling

April 24, 2024

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

cs.AI

Zusammenfassung

Spekulatives Decodieren hat sich als eine leistungsstarke Methode zur Verbesserung von Latenz und Durchsatz bei der Bereitstellung großer Sprachmodelle etabliert. Die meisten bestehenden Implementierungen konzentrieren sich jedoch darauf, eine einzelne Sequenz zu generieren. Generative KI-Anwendungen in der realen Welt erfordern oft mehrere Antworten, und die Durchführung des spekulativen Decodierens in einem Stapelmodus unter Beibehaltung seiner Latenzvorteile birgt nicht-triviale Herausforderungen. Dieser Artikel beschreibt ein System des stapelweisen spekulativen Decodierens, das einen neuen Stand der Technik in Bezug auf die Latenz bei der Mehrfachsequenzgenerierung setzt und eine überlegene GPU-Auslastung sowie Qualität der Generierungen innerhalb eines Zeitbudgets aufzeigt. Zum Beispiel wird für ein Modell der Größe 7,8 Milliarden auf einer einzelnen A100 GPU und mit einer Stapelgröße von 8 jede Sequenz mit einer durchschnittlichen Geschwindigkeit von 5,8 ms pro Token generiert, wobei der Gesamtdurchsatz bei 1,1K Tokens pro Sekunde liegt. Diese Ergebnisse repräsentieren den Stand der Technik in Bezug auf die Latenz und eine 2,15-fache Beschleunigung gegenüber optimiertem regulärem Decodieren. Innerhalb eines Zeitbudgets, das reguläres Decodieren nicht abschließt, ist unser System in der Lage, Sequenzen mit einer HumanEval Pass@First-Rate von 43% und Pass@All-Rate von 61% zu generieren, was bei weitem über das hinausgeht, was mit spekulativem Decodieren von Einzelsequenzen möglich ist. Unsere Spitzen-GPU-Auslastung während des Decodierens erreicht bis zu 15,8%, mehr als das 3-fache des höchsten Wertes des regulären Decodierens und etwa 10-fach des spekulativen Decodierens von Einzelsequenzen.

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.