概要
Tokasaurus は、高スループットLLM推論に特化した新しいエンジン 小型モデル では低CPUオーバーヘッドとHydragenによる動的グルーピングで高速化 大型モデル ではNVLink有無に応じたパラレル手法で高効率を実現 ベンチマークで vLLMやSGLangを最大3倍以上 上回る性能を発揮 GitHubとPyPIで公開、 Llama-3やQwen-2 などをサポート
Tokasaurus：高スループットLLM推論エンジン
- Tokasaurus は、スループット重視のLLM推論ワークロード向けに最適化されたエンジン
- 小型モデル と 大型モデル の両方に対応し、従来エンジンを上回る性能を実現
- FlexGen、vLLM、SGLang などのOSSエンジンを参考に独自最適化を追加
- 主な用途は大量シーケンスのバッチ処理や合成データ生成、RLパイプラインなど
- 単一生成のレイテンシではなく、 全体の処理時間とコスト を重視
小型モデル最適化
- ShareGPTデータセットによるチャットボット推論 と GSM8Kの大量サンプリング で検証
- 2つの主要最適化 で他エンジンに対し2倍以上のスループットを達成
- CPUオーバーヘッド最小化
- Webリクエスト処理、トークナイズ、KVキャッシュ管理などを非同期化
- マネージャー が入力キューを監視し、GPUスタールを防止
- 必要に応じてオプション処理（停止文字列チェック等）を自動スキップ
- 動的なプレフィックス識別と活用
- 共有プレフィックスを深さ優先探索で検出し、 Hydragen で効率的に利用
- 小型モデルほど注意機構の計算負荷が大きく、効果が顕著
- CPUオーバーヘッド最小化
大型モデル最適化
- 複数GPU 環境での効率的な推論を実現
- Pipeline Parallelism（PP） ：NVLink非搭載GPU（例：L40S）向け
- 通常より少ない通信コストで大規模バッチを分割・並列実行
- vLLMやSGLangのPP実装と比較し、 Llama-3.1-70B×8GPU で3倍以上のスループット
- Async Tensor Parallelism（Async-TP） ：NVLink搭載GPU向け
- PyTorchのAsync-TPに対応し、通信と計算のオーバーラップが可能
- 大規模バッチ時に自動でAsync-TPを有効化し、最適なスループットを確保
導入方法と対応モデル
- GitHub でコード公開、 PyPI からpip install tokasaurusでインストール可能
- Llama-3、Qwen-2 ファミリー等に対応
- 1ノード内でデータ並列・テンソル並列・パイプライン並列を任意に組み合わせ可能
- 純Python実装 で、FlashInferのAttention/Sampling演算を利用
- GPT-fast のようなフォーク・カスタマイズも容易
ベンチマーク詳細
- 全エンジン共通 のKVキャッシュサイズ・リクエスト数で公平に比較
- 各種パラメータも最適化し、ウォームアップ後の平均スループットを計測
- ShareGPT ベンチはSGLangのスクリプト、 Large Language Monkeys ベンチは独自スクリプトを使用
- 全実験はOpenAI API経由で統一実施
- vLLM Python API（LLM.generate()）利用時、Llama-1Bで約5%のスループット増加も確認
謝辞・引用
- Prime Intellect および Together AI の計算資源提供に感謝
- Dan Biderman、Simon Guo、Manat Kaur、Avanika Narayan によるβテスト協力
- 利用時は以下の引用を推奨
- @misc{juravsky2025tokasaurus, author = {Jordan Juravsky and Ayush Chakravarthy and Ryan Ehrlich and Sabri Eyuboglu and Bradley Brown and Joseph Shetaye and Christopher Ré and Azalia Mirhoseini}, title = {Tokasaurus: An LLM Inference Engine for High-Throughput Workloads}, year = {2025}, howpublished = {\url{https://scalingintelligence.stanford.edu/blogs/tokasaurus/}} }