TokenRouter：トークン単位のLLMルーティングを支える推論基盤

導入

一般的なLLMルーターは、セッションまたはリクエスト全体に対して一度だけモデルを選びます。入力の難易度や目的を判断し、生成全体を適切なモデルへ送る方式です。実装しやすい一方、生成途中で必要となる品質、コスト、速度の変化を細かく反映することはできません。

これに対してトークン単位のルーティングでは、生成中の状態に応じて各トークンを異なるモデルへ振り分けられます。より柔軟な設計ですが、サービング基盤に大きな負担を与えます。単一LLMを前提とする推論システムでは、モデルごとに処理速度が異なるため、デコードステップの同期が崩れやすくなります。また、リクエストが複数モデルへ動的に分散されることで、効率的なバッチを作るまで待たされるケースも増えます。

TokenRouterの設計

TokenRouterは、新しいルーティングアルゴリズムそのものではなく、トークン単位のルーティングを実行するためのサービングシステムです。中心となる考え方は、「リクエスト中心のプログラミング、モデル中心の実行」です。開発者は一つのリクエストから見たルーティングロジックを記述し、ランタイムがLLMごとにサブサーバーを起動して、処理を非同期に振り分けます。

この構成では、ルーティングロジックとモデル内部のスケジューリングを分離できます。開発者がモデルごとの複雑な制御コードを個別に実装する必要が減るだけでなく、各サブサーバーが自分の負荷と速度に合わせて処理を進められます。全モデルを一つの同期ループに合わせる必要がない点が特徴です。

さらに各サブサーバーには遅延バッチ処理のスケジューラーを配置します。新しい処理を常に即時実行するのではなく、追加の待ち時間と、より大きなバッチによるスループット向上を比較します。論文では、スループットの数学モデルからスケジューラーの主要なハイパーパラメーターを導出しています。

主なポイント

• リクエスト単位ではなくトークン単位でモデルを選択。
• LLMごとに独立したサブサーバーを配置。
• モデル間の処理を非同期にディスパッチ。
• 動的な分散で生じる小規模バッチの問題を遅延バッチで緩和。
• 数学的なスループットモデルで調整値を決定。
• 複数のルーティング方式、負荷、モデル組み合わせで評価。
• 既存システムに対して2.01〜64.15倍のデコードスループットを報告。

意義と留意点

TokenRouterが示すのは、トークン単位のルーティングがアルゴリズムだけで完結しないという点です。ルーティングによる品質やコストの利点を実サービスへ反映するには、同期処理とスケジューリングのオーバーヘッドを抑えるランタイムが必要です。リクエスト側の記述とモデル側の実行を分ける設計は、能力や速度の異なるモデルを組み合わせる際の抽象化としても分かりやすいものです。

ただし、提供された素材には、完全なハードウェア構成、各ベースラインの詳細、実験ごとの絶対スループットは含まれていません。そのため、報告された倍率があらゆる本番環境で再現されるとは限らず、モデル構成や負荷、ルーティング分布によって結果は変わる可能性があります。それでも、モデルを一つ選ぶ従来型の推論基盤から、トークンごとにモデルを切り替える基盤へ移行する際の重要な設計方向を示しています。

出典：Hugging Face Daily Papers