スタンフォード大がTokasaurus公開：高スループット特化でvLLMを最大3倍凌駕

Tokasaurusは、スタンフォード大学の研究チームが公開した、バッチ処理や高スループット推論に特化したオープンソースのLLM推論エンジンだ。

何ができるのか

合成データの生成や強化学習のロールアウト、大量のコードベース解析といった作業では、チャットボットのような1リクエストごとの応答速度よりも、一定時間内にどれだけ多くのトークンを処理できるかが重要になる。

Tokasaurusはまさにこのスループット最大化を狙って設計された。一般的なShareGPTベンチマークや、共通プロンプトを大量に再利用するLarge Language MonkeysテストにおいてvLLMやSGLangを上回り、特に分岐生成が多いタスクでは2倍以上のスループットを達成している。

さらにNVLink非搭載のL40Sを8基使ったLlama-3.1-70Bのテストでは、既存エンジンのパイプライン並列実装と比較して3倍以上の処理能力を叩き出した。

仕組みと利用環境

この性能を支えているのは主に3つの最適化だ。

1つ目はCPUオーバーヘッドの最小化だ。GPUが計算待ちになるのを防ぐため、内部のマネージャーがキューを常時監視し、枯渇しそうになると停止文字列の検出など優先度の低い処理を自動で省いてGPUの稼働率を維持する。

2つ目は共有プレフィックスの動的活用だ。同一の長文やシステムプロンプトを参照するリクエストに対し、モデル実行直前に貪欲法で最長の一致部分を特定し、Hydragenの仕組みを使って重複するAttention計算を削ぎ落とす。

3つ目はハードウェア構成に応じた並列化の使い分けだ。安価な非NVLink環境向けには高効率なパイプライン並列を用意し、H100やB200のようなNVLink搭載機ではバッチサイズが6000トークンを超えた段階でPyTorchの非同期テンソル並列へ自動で切り替え、通信遅延を計算の裏に隠蔽する。

現在は単一ノード内でのLlama-3およびQwen-2ファミリーに対応しており、実装はFlashInferを利用した純粋なPythonで書かれている。コードはGitHubで公開されており、pip install tokasaurusですぐに導入可能だ。

なぜ重要か

強化学習用のデータ生成やバッチ評価を日常的に回しているエンジニアなら、Tokasaurusを検証環境で試す価値は十分にある。チャットサービス向けの低遅延エンジンではないが、大量のトークンを少しでも安く速く処理したい現場では確実なコスト削減につながる。

特に評価できるのは、安価な非NVLink環境への最適化を本気で突き詰めた点だ。多くの先端実装が高価な相互接続を前提とする中、パイプラインのスケジューリングを磨くことで普及帯のGPUから圧倒的なスループットを引き出してみせた。対応モデルの広さではまだvLLMに及ばないものの、目的が合致するバッチ用途なら極めて実用的な選択肢になるはずだ。