BASS: バッチ化された注意機構最適化型推測サンプリング

BASS: Batched Attention-optimized Speculative Sampling

April 24, 2024

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

cs.AI

要旨

推測デコードは、大規模言語モデルのホスティングにおけるレイテンシとスループットを改善する強力な手法として登場しました。しかし、既存の実装のほとんどは単一のシーケンス生成に焦点を当てています。現実世界の生成AIアプリケーションでは、複数の応答が必要となることが多く、バッチ設定で推測デコードを実行しつつそのレイテンシの利点を維持することは、自明ではない課題を提起します。本論文では、バッチ推測デコードのシステムを説明し、複数シーケンス生成におけるレイテンシの新たな最先端を確立するとともに、時間予算内での生成品質とGPU利用率の優位性を実証します。例えば、7.8Bサイズのモデルを単一のA100 GPUでバッチサイズ8で実行した場合、各シーケンスは平均5.8ms/tokenの速度で生成され、全体のスループットは1.1K tokens/秒となります。これらの結果は、最適化された通常のデコードと比較して、最先端のレイテンシと2.15倍の高速化を表しています。通常のデコードでは完了しない時間予算内で、本システムはHumanEval Pass@Firstで43%、Pass@Allで61%のシーケンスを生成することができ、単一シーケンス推測デコードで実現可能な範囲を大幅に上回ります。デコード中のピークGPU利用率は15.8%に達し、通常のデコードの最高値の3倍以上、単一シーケンス推測デコードの約10倍となります。

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.