BASS: Batched Attention-optimized Speculative Sampling

要約

投機的デコードは、大規模な言語モデルをホストする際の待ち時間とスループットを改善する強力な方法として登場しました。
ただし、既存の実装のほとんどは、単一のシーケンスを生成することに重点を置いています。
現実世界の生成 AI アプリケーションでは複数の応答が必要になることが多く、レイテンシーの利点を維持しながらバッチ設定で投機的デコードを実行する方法は、簡単ではない課題を引き起こします。
この論文では、マルチシーケンス生成レイテンシーにおいて新しい最先端技術を確立し、優れた GPU 利用率と時間予算内での生成品質を実証する、バッチ投機的デコーディングのシステムについて説明します。
たとえば、単一の A100 GPU でバッチ サイズ 8 の 7.8B サイズのモデルの場合、各シーケンスはトークンあたり平均 5.8 ミリ秒の速度で生成され、全体のスループットは 1.1K トークン/秒になります。
これらの結果は、最先端のレイテンシと、最適化された通常のデコードと比べて 2.15 倍のスピードアップを示しています。
通常のデコードでは完了しない時間枠内で、私たちのシステムは HumanEval Pass@First が 43%、Pass@All が 61% のシーケンスを生成できます。これは、単一シーケンスの投機的デコードで実現可能なものをはるかに超えています。
デコード中のピーク GPU 使用率は 15.8% にも達し、通常のデコードの最高 3 倍以上、単一シーケンスの投機的デコードの約 10 倍に達します。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google