MagicDec: 推測デコーディングによる長文脈生成におけるレイテンシとスループットのトレードオフの打破

MagicDec: Breaking the Latency-Throughput Tradeoff for Long Context Generation with Speculative Decoding

August 20, 2024

著者: Jian Chen, Vashisth Tiwari, Ranajoy Sadhukhan, Zhuoming Chen, Jinyuan Shi, Ian En-Hsu Yen, Beidi Chen

cs.AI

要旨

大規模言語モデル（LLM）は、インタラクティブなチャットボット、ドキュメント分析、エージェントワークフローなどの長文脈アプリケーションでより一般的になってきているが、長文脈リクエストを低レイテンシと高スループットで提供することは困難である。推測的デコード（SD）は、パフォーマンスを犠牲にすることなくレイテンシを削減するために広く使用されている技術だが、従来の知見では、その効果は小規模なバッチサイズに限定されるとされていた。MagicDecでは、驚くべきことに、SDが中程度から長いシーケンスにおいても高スループット推論レジームで高速化を達成できることを示す。さらに興味深いことに、我々の厳密な分析に基づき、インテリジェントなドラフト戦略を用いることで、バッチサイズが増加するにつれてより良い高速化を実現できる。MagicDecはまず、バッチサイズとシーケンス長の増加に伴うボトルネックの変化を特定し、これらの洞察を活用して高スループット推論において推測的デコードをより効果的に展開する。次に、スパースKVキャッシュを備えたドラフトモデルを活用し、シーケンス長とバッチサイズの両方に比例してスケールするKVボトルネックに対処する。

English

Large Language Models (LLMs) have become more prevalent in long-context applications such as interactive chatbots, document analysis, and agent workflows, but it is challenging to serve long-context requests with low latency and high throughput. Speculative decoding (SD) is a widely used technique to reduce latency without sacrificing performance but the conventional wisdom suggests that its efficacy is limited to small batch sizes. In MagicDec, we show that surprisingly SD can achieve speedup even for a high throughput inference regime for moderate to long sequences. More interestingly, an intelligent drafting strategy can achieve better speedup with increasing batch size based on our rigorous analysis. MagicDec first identifies the bottleneck shifts with increasing batch size and sequence length, and uses these insights to deploy speculative decoding more effectively for high throughput inference. Then, it leverages draft models with sparse KV cache to address the KV bottleneck that scales with both sequence length and batch size.