BASS: Campionamento Speculativo Ottimizzato con Attenzione in Batch

BASS: Batched Attention-optimized Speculative Sampling

April 24, 2024

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

cs.AI

Abstract

Il decoding speculativo è emerso come un metodo potente per migliorare la latenza e il throughput nell'hosting di grandi modelli linguistici. Tuttavia, la maggior parte delle implementazioni esistenti si concentra sulla generazione di una singola sequenza. Le applicazioni generative di intelligenza artificiale nel mondo reale spesso richiedono più risposte, e come eseguire il decoding speculativo in un contesto batch preservando i suoi vantaggi in termini di latenza pone sfide non banali. Questo articolo descrive un sistema di decoding speculativo in batch che stabilisce un nuovo stato dell'arte nella latenza di generazione di sequenze multiple e dimostra una superiore utilizzazione della GPU, nonché una qualità delle generazioni entro un budget di tempo. Ad esempio, per un modello di dimensioni 7.8B su una singola GPU A100 e con una dimensione del batch di 8, ogni sequenza viene generata a una velocità media di 5.8ms per token, con un throughput complessivo di 1.1K token al secondo. Questi risultati rappresentano una latenza all'avanguardia e un miglioramento di 2.15X rispetto al decoding regolare ottimizzato. Entro un budget di tempo in cui il decoding regolare non riesce a completare, il nostro sistema è in grado di generare sequenze con un HumanEval Pass@First del 43% e un Pass@All del 61%, superando di gran lunga quanto è fattibile con il decoding speculativo a singola sequenza. La nostra massima utilizzazione della GPU durante il decoding raggiunge il 15.8%, più di 3X rispetto al massimo del decoding regolare e circa 10X rispetto al decoding speculativo a singola sequenza.

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.