BASS: Gebatchte Attention-geoptimaliseerde Speculatieve Sampling

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

Samenvatting

Speculatief decoderen is naar voren gekomen als een krachtige methode om de latentie en doorvoer te verbeteren bij het hosten van grote taalmodellen. De meeste bestaande implementaties richten zich echter op het genereren van een enkele reeks. In real-world generatieve AI-toepassingen zijn vaak meerdere reacties vereist, en het uitvoeren van speculatief decoderen in een batchomgeving terwijl de latentievoordelen behouden blijven, vormt niet-triviale uitdagingen. Dit artikel beschrijft een systeem van batchgewijs speculatief decoderen dat een nieuwe standaard zet op het gebied van latentie bij het genereren van meerdere reeksen en dat superieure GPU-gebruik en kwaliteit van generaties binnen een tijdsbudget aantoont. Bijvoorbeeld, voor een model van 7,8B grootte op een enkele A100 GPU en met een batchgrootte van 8, wordt elke reeks gegenereerd met een gemiddelde snelheid van 5,8ms per token, waarbij de totale doorvoer 1,1K tokens per seconde bedraagt. Deze resultaten vertegenwoordigen state-of-the-art latentie en een 2,15X versnelling ten opzichte van geoptimaliseerd regulier decoderen. Binnen een tijdsbudget waarin regulier decoderen niet voltooid wordt, is ons systeem in staat om reeksen te genereren met een HumanEval Pass@First van 43% en een Pass@All van 61%, wat ver uitstijgt boven wat haalbaar is met enkelvoudig speculatief decoderen. Ons piek-GPU-gebruik tijdens het decoderen bereikt wel 15,8%, meer dan 3X het hoogste niveau van regulier decoderen en ongeveer 10X dat van enkelvoudig speculatief decoderen.

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.