Anyscale、Ray Serve LLM でトークン負荷意識型最適化を提案

大規模なLLMのサービス提供において、オーケストレーション層はシステム全体の効率性を決定づける重要な要素です。この層が、多様なリクエストストリームをLLMエンジンレプリカ群にどのように分散させるかが、TTFT（Time to First Token）、TPOT（Token Per Second Output Time）、スループットといった指標に直接影響を及ぼします。

vLLMのレプリカ群を運用する際、ルーティング層が下す最も重要な判断はシンプルです。「次のリクエストを誰に割り当てるか？」

本稿では、「KVキャッシュの再利用のみを最適化対象とする」という誤解を解き、トークン負荷を意識したルーティング（token-load-aware routing）の概念を紹介するとともに、それがなぜ効率的なLLMサービス提供にとって不可欠なのかを解説します。

LLMリクエストルーティングにおける課題

LLMのサービス提供は、推論が自己回帰的であることと、LLMリクエストの多様性が極めて高いという特性により、従来のマイクロサービス向けルーティングとは根本的に異なります。従来のマイクロサービスでは、リクエストを独立して処理することが一般的であり、関連するリクエストを同じレプリカにルーティングしても大きなメリットは得られません。

一方、LLMのサービス提供では、状態保持（statefulness）が必要となり、リクエスト間の多様性が高く、実行コストが予測困難であるという課題が生じます。

ステートフル実行では、エンジンが過去のリクエストで計算された KV キャッシュを再利用できます。リクエストのプレフィックス（先頭部分）を既に保持しているレプリカにルーティングすることで、事前計算（prefill）の負荷と TTFT（Time To First Token）を大幅に削減できます。

高い異質性も課題です。入力や出力シーケンスの長さがリクエスト間で大きく異なるため、GPU リソースの使用量や実行時間に大きな差が生じます。また、生成は非確定的な性質を持ちます。同じ入力であっても生成されるトークン数は変動するため、リクエスト全体のコストを事前に予測するのは困難です。

これらの特性が、オーケストレーション層にとって独自の課題を生み出しています。すなわち、KV キャッシュの再利用というメリットを活かしつつ、エンジンレプリカ間の異質性をどうバランスさせるかという問題です。

LinkNaive Optimization: KV キャッシュの再利用最大化

TTFT の短縮とサービング効率の最適化において、KV キャッシュの再利用を最大化する手法は最も素朴なアプローチと言えます。一般的には、セッションアフィニティ、プレフィックスアフィニティ、そして KV キャッシュアフィニティの 3 つの方法があります。

LinkSession Affinity（セッションアフィニティ）

Ray Serve LLM は、一貫性ハッシュ（consistent hashing）を通じてセッションアフィニティをサポートしています。これには、クライアントが同一セッションからのリクエストにセッション ID を関連付ける必要があります。マルチターンワークロードでは、この仕組みにより後続のターンを同じレプリカへルーティングし、先行するターンで計算済みの KV キャッシュを再利用可能になります。

一貫性ハッシュによるセッションアフィニティには、レベルごとの負荷分散という追加メリットもあります。仮想ノードを十分に用意すれば、セッションをレプリカ間で均等に分散させることが可能です。

図 1: 仮想ノードを備えた一貫性ハッシュはセッションレベルの負荷分散を実現します。

*仮想ノードを備えた一貫性ハッシュはセッションレベルの負荷分散を実現します*

session_affinity_config = LLMConfig(
    model_loading_config=...,
    deployment_config=(
        "request_router_config": RequestRouterConfig(
            request_router_class=(
                "ray.serve.experimental.",
                "consistent_hash_router.ConsistentHashRouter"
            )
        ),
    engine_kwargs={...},
)

LinkPrefix アフィニティ

Ray Serve LLM のプレフィックスアフィニティルーターは、受信したリクエストに基づいてプレフィックスツリーを構築し、各レプリカに送信されたプレフィックスを追跡します。後続のリクエストはこのツリーを探索し、最大のプレフィックス重複を持つレプリカを見つけます。

• 近似値の扱い: プレフィックスツリーは実際の KV キャッシュではなくリクエスト文字列を追跡しています。KV キャッシングはトークンレベルで動作するため、ルーターが把握する状態とエンジン内の実際のキャッシュ状態には差異が生じる可能性があります。
• _eviction_（キャッシュ削除）を考慮していない: ルーターは、以前に見たプレフィックスが常にキャッシュに残っていると仮定します。高同時実行環境や長時間稼働するサービスでは、キャッシュの削除（eviction）が発生し、この推測が古びてしまうリスクがあります。KV キャッシュオフローディング はキャッシュミス削減に寄与しますが、ルーターは依然としてすべてのキャッシュ階層を同等に扱い、オフロードされたブロックを GPU メモリへ再読み込みするコストを無視したままです。

prefix_affinity_config = LLMConfig(
    model_loading_config=...,
    deployment_config=(
        "request_router_config": RequestRouterConfig(
            request_router_class=(
                "ray.serve.llm.",
                "request_router.PrefixCacheAffinityRouter",
        ),
    engine_kwargs={...},
)

LinkKV Cache Affinity

Ray Serve LLM の KV 対応ルーター は、エンジンから KV キャッシュのイベントを受け取り、グローバルなラジックスツリーを構築します。vLLM はキャッシュブロックが作成されたり削除されたりするたびにイベントを出力するため、ルーターはレプリカ全体にわたる KV キャッシュの状態を常に把握できます。これにより、各リクエストに対して KV キャッシュブロック単位での重複計算が可能になります。

私たちは NVIDIA Dynamo チームと協力し、モジュラーインターフェースを通じて Dynamo の KV インデクサーを外部ライブラリから利用可能にしています。イベント、リクエスト、データ平面は Ray ネイティブのまま維持され、Dynamo の KV インデクサーが KV キャッシュの重複計算を行い、リクエストのスコアリングに貢献します。算出されたスコアはその後、Ray Serve のリクエスト平面で使用され、対象となるレプリカを選択するために活用されます。

図 2: KV イベント伝送とグローバル KV ラジックスツリーの構築。vLLM エンジンは、KV キャッシュブロックが作成または削除される際に KV イベントを出力します。Ray Serve LLM はこれらのイベントを KVAwareRouter に転送し、NVIDIA Dynamo の KV インデクサーがグローバルな KV ラジックスツリーを構築して、KV キャッシュの重複に基づいてリクエストにスコアを付けます。

*KV イベント伝送とグローバル KV ラジックスツリーの構築。*

kv_affinity_config = LLMConfig(
    model_loading_config=...,
    deployment_config=(
        "request_router_config": RequestRouterConfig(
            request_router_class=(
                "ray.serve.llm.",
                "request_router.KVAwareRouter",
        ),
    engine_kwargs={...},
)

KV キャッシュのオーバーラップだけでなく、Ray Serve LLM の KV キャッシュアフィニティルーターは、ルーティング決定時にトークン負荷も考慮します。これは以下のセクションで詳しく解説します。

なぜ KV キャッシュの再利用最大化だけでは不十分なのか？

KV キャッシュの再利用を最適化しても、必ずしもサービングパフォーマンスが最大化されるわけではありません。例えば、遅延するリクエスト（ストランガラー）が存在し、入力・出力シーケンス長が大きく異なる非同期マルチターン RL ロールアウトワークロードを考えてみましょう。

• 各ステップでは8つのロールアウトを実行し、ロールアウトあたり10ターンあります。
• 各ロールアウトは2K トークンの入力から開始され、各ターンで1K トークンを生成します。ただし、最終ターンで8K トークンを生成する2つのストランガラーを除きます。

図 3: 非同期 RL マルチターンロールアウトの可視化。

*Figure 3: Asynchronous RL multi-turn rollouts visualization.*

非同期 RL ロールアウトでは、通常はステップ時間の最小化、または固定並列度でのスループット最大化が目標となります。ここでは、ステップ時間の代わりとして p99 エンドツーエンドのロールアウトレイテンシを使用します。図 4 に示す通り、並列度が 16 の条件下で 3 つのルーター変種を比較しました。

• PureKVCacheAffinityRouter は、KV キャッシュの重なりを最適化するために KVAwareRouter を改変したものです。
• ConsistentHashRouter もまた、マルチターン会話における KV キャッシュの再利用を最大化しつつ、セッションレベルでの負荷分散を提供します。

バランスの取れた割り当てを行います。各マルチターンロールアウトにはセッション ID を付与します。

• KVAwareRouter：KV キャッシュの重複とトークン負荷の両方を考慮するルーティングアルゴリズムです。

図 4：非同期マルチターン RL ロールアウトにおけるルーターの比較。KVAwareRouter は KV キャッシュのヒット率をいくらか犠牲にする代わりに、トークン負荷のバランスを改善し、p99 のエンドツーエンドロールアウトレイテンシを低下させます。

*非同期マルチターン RL ロールアウト*

KVAwareRouter は、プレフィックスキャッシュのヒット率がやや低くても、p99 のロールアウトエンドツーエンドレイテンシにおいて最良の結果を示します。KV キャッシュのヒット率とトークン負荷（デコードリクエストがレプリカ上で占めるアクティブな KV キャッシュの推定規模）をより深く分析すると、KVAwareRouter は KV キャッシュのヒット率をある程度犠牲にしてでも、トークン負荷のバランスを取る戦略を採用していることがわかります。

LinkRequest Herding

PureKVCacheAffinityRouter が提供する純粋な KV キャッシュアフィニティは、リクエストの群れ（herding）を引き起こす可能性があります。一度レプリカがシステムプロンプトのような共有プレフィックスをキャッシュすると、その後のリクエストはキャッシュの重複を検知し、優先的にそのレプリカへルーティングされます。結果として、異なる会話からの多数のリクエストが同一のレプリカに集中し、トークン負荷の不均衡を生み出します。

リンクロードバランシング

一貫ハッシュでは仮想ノードを十分に用意することで、レプリカ間でセッション数を均等に分散できます。しかし、この手法はすべてのセッションを同質のものとして扱ってしまいます。LLM の推論サービスにおいては、セッションは均質ではありません。あるセッションには短時間のリクエストが含まれる一方、別のセッションは長時間実行される遅延要因（ストラングラー）になる可能性があります。複数のストラングラーが同じレプリカに集中してしまうと、リクエストの分布自体は均等であっても、実際の負荷バランスは崩れてしまいます。

セッションアフィニティは KV キャッシュの局所性を維持する点で優れた手法ですが、リクエストレベルでの均衡を保つだけでは、実際の負荷バランスを担保することはできません。

KV キャッシュ再利用を超えて：トークン負荷への意識が重要

エンジンが各リクエストに対して実際に実行する処理について考えてみましょう。1 つのリクエストは以下の 2 つのフェーズを経ます。

• プリフィル: エンジンはまず、重複している KV キャッシュを再利用します。その後、キャッシュされていない入力トークンに対する KV キャッシュの計算を行います。
• デコード: プリフィルが完了すると、リクエストは出力トークンの生成を開始します。各デコードステップにおいて、エンジンはこれまでに生成されたすべてのトークンの KV キャッシュにアクセスし、もう一度フォワードパスを実行します。

この仕組みから、KV キャッシュの重なりをルーティングの目的とする手法の限界が見えてきます。KV キャッシュの重なりは、プリフィル処理で節約できる量を示すだけで、エンジンが実際に行う必要がある作業量までは示しません。残りの負荷は以下の 2 つの要因から生じます。

• プリフィル負荷: アクティブなプリフィルトークンと、キャッシュされていない新規入力トークンの合計。
• デコード負荷: アクティブなリクエストによる進行中のデコード処理。

トークン負荷（token load）は、計算リソースに制約されるプレフィル負荷とメモリに制約されるデコード負荷の両方を一つの指標として統合し、各レプリカの状態を捉えます。

KV キャッシュのアフィニティも依然重要ですが、その役割は変化します。つまり、KV キャッシュの重なりは負荷推定の代理指標であり、ルーティング自体の目的ではないのです。キャッシュの重なりから、着信するプレフィル処理のうち何分をスキップできるかを把握し、それをもとに残りのプレフィル負荷を見積もります。さらに、現在進行中のデコード負荷と組み合わせることで、全体のトークン負荷を推定します。

KVAwareRouter は、Ray Serve LLM においてトークン負荷の意識を備えたルーティングポリシーです。これにより、一貫性ハッシュに基づくセッションアフィニティよりもさらに優れた機能が提供されます。

• 自動的な KV キャッシュの再利用: セッションアフィニティでは一貫性ハッシュを利用するため、クライアント側でセッション ID を指定する必要があり、アプリケーションがどのリクエストを同じセッションに属するものとして明示的に定義する必要があります。一方、KVAwareRouter はセッション ID の指定を不要とし、利用可能な KV キャッシュを自動的に発見して再利用します。
• 柔軟性: 負荷状況に応じて、スコアリング関数の 係数 を調整できます。

Ray Serve LLM は、リクエストが受付された瞬間から最初のトークン生成、デコードの進行、そして完了に至るまで、ライフサイクル全体を追跡します。各フェーズにおいて、ルーターは対応するトークン負荷を更新し、すべてのエンジンレプリカにわたる最新の負荷状況を維持しています。Ray Serve LLM はエンドツーエンドのリクエストおよび制御プレーンを提供しつつ、リクエストのスコアリングには Dynamo の選択サービスモジュールを統合しています。

図 5: トークン負荷の登録。各リクエストは、受付からプリフィル完了、デコード、そしてリクエスト完了へと進む過程で、KV Aware Router にトークン負荷を登録・更新します。ルーターは各エンジンのトークン負荷を常に最新の状態に保ち、これを基に後続のリクエストのスコアリングとルーティングを行います。

*KV aware router*

スケーラビリティの観点から、Ray Serve LLM はルーターをレプリケートし、ルーター間でのトークン負荷に関する最終的な整合性を保証したビューを維持します。KVAwareRouter はローカルでルーティング判断を行い、その負荷更新情報を非同期方式でピアのルーターと共有します。これにより、ルーティング処理やトークンのストリーミングが同期パスから外れつつも、すべてのルーターが同じ負荷ビューに収束できるようになります。

ただし、一貫性のあるハッシュによるセッションアフィニティは、以下のようなワークロードにおいて依然として有力な選択肢となります。

キャッシュミスは特にコストがかかります。長い多回話の会話において、すべての会話を同じレプリカで保持することは KV キャッシュの局所性を保証します。KVAwareRouter はスコアリング係数に応じて、キャッシュの重複をある程度犠牲にしてトークン負荷のバランスを改善する可能性があり、そのためにはチューニングが必要です。

セッション間の KV キャッシュ再利用には限界があります。共通のシステムプロンプト後に会話が分岐した場合、KVAwareRouter がセッション間で発見できる追加的なキャッシュ再利用はほとんどありません。

セッションには類似したワークロードが存在します。セッション間に入力・出力の長さや実行時間に大きな差異がない場合、一貫性ハッシュ法でもセッションレベルで十分な負荷分散が実現できます。

リンク：ケーススタディ Claude Code トレースの再生

コーディングエージェントは現在、最も人気のある LLM インフェレンスワークロードの一つです。そこで、実際の Claude Code のトレースからヒントを得てみましょう。gpt-oss-120b のコンテキストウィンドウに収まるようフィルタリングした Weka Claude Code トレースコーパス のサブセット上で、さまざまなルーティング戦略を評価します。

図 6：Claude Code トレースのサンプル。

*Figure 6: Sample Claude Code trace.*

まず、ワークロードの形状を見てみましょう。注目すべき点は二つあります。セッションが非常に多様であることと、各セッション内で KV キャッシュの再利用に大きな余地があることです。

個々のリクエストでも入力長や出力長のばらつきはすでに大きいですが、セッションレベルで見るとその差はさらに顕著になります。単一のターンで完了するセッションもあれば、数十回のターンを要し、数百万トークンもの入力を蓄積するセッションもあります。

Figure 7: Distribution of Weka Claude Code replay requests and reconstructed sessions: long individual contexts,

highly variable outputs, and heavy-tailed session lifetimes and turn counts.

*Figure 7: Distribution of Weka Claude Code replay requests and reconstructed sessions: long individual contexts, highly variable outputs, and heavy-tailed session lifetimes and turn counts.*

この状況は、興味深いルーティングのトレードオフを生み出します。セッションを同じレプリカに固定して KV キャッシュの再利用を最大化するか、それともトークン負荷をレプリカ間でより均等に分散させるか、どちらを選ぶべきかです。

Ray Serve LLM の KVAwareRouter は、レプリカを評価する際に KV キャッシュの重複率だけでなくトークン負荷も考慮します。その結果、一貫性ハッシュを用いたセッションアフィニティと比較して、TTFT（Time to First Token）、TPOT（Time Per Output Token）、スループットともに優れたパフォーマンスを発揮することが示されています。

図 8：トークン負荷の意識化は、多様なエージェント型マルチターンワークロードのパフォーマンスを向上させます。KVAwareRouter は、若干のプレフィックスキャッシュヒット率を犠牲にしてトークン負荷のバランスを改善し、TTFT、TPOT、スループットを向上させています。補足すると、並行度が高まるにつれてプレフィックスキャッシュヒット率が低下するのは、KV キャッシュの_eviction_（削除）の影響が顕著になることを示しています。

*トークン負荷の意識化は、多様なエージェント型マルチターンワークロードのパフォーマンスを向上させます。*

その理由を理解するために、アクティブなデコードブロックの係数分散（CV）とプレフィックスキャッシュヒット率を見てみましょう。一貫性ハッシュによるセッションアフィニティは、同じセッションからのリクエストが同一レプリカに留まるため、KV キャッシュの再利用を効果的に維持します。しかし、この手法で負荷分散が行われるのはセッションレベルのみです。あるセッションが他よりも大幅に大きかったり長かったりするケースでは、セッション数を均等に配分しても、トークン負荷まで均等になるわけではありません。

KVAwareRouter は、並行処理が増えるにつれて、一部のプレフィックスキャッシュヒット率を犠牲にしてでもトークン負荷のバランスを改善します。その結果、異種 LLM ワークロードにおいては、KV キャッシュの再利用だけを最大化するよりも、KV キャッシュの再利用とトークン負荷の両方をバランスさせることが、より優れた全体的なサービング性能をもたらすことが示されました。

注記: プレフィックスキャッシュヒット率は、vLLM が報告したキャッシュされたプロンプトトークンの数を総プロンプトトークン数で割った値をトークン重み付けした比率です。これは、両方の使用状況フィールドが利用可能なリクエストを対象に計算されます。平均的なアクティブデコードブロックの CV（変動係数）は、レプリカごとのデコーディング KV ブロックの不均衡をオフラインで再構築したものです。0.5 秒ごとのバインにおいて、各リクエストは最初のトークンから完了まで推定されるアクティブな KV ブロックのフットプリントを提供します。CV が低いほど、負荷がよりバランスしていることを意味します。

リンクと今後の課題

このブログ記事の冒頭で概説されたルーティングの課題を再検討すると、KVAwareRouter は状態実行とワークロードの多様性に対応していますが、非決定性の生成は依然として未解決の問題です。出力長は受付時点では不明であるため、当初は軽量なリクエストと見なされていたものが、長いデコードフェーズを経て同じレプリカで衝突し、予期せぬトークン負荷の不均衡を引き起こす可能性があります。この不確実性を正確に反映させることは、業界全体にとって依然として大きな課題です。一つの解決策として、エージェント型ハルネスが期待される生成挙動に関するヒントを提供することが考えられます。

近い将来、KVAwareRouter のサポートを、プリフィルとデコードを分離した展開環境やデータ並列展開、マルチモーダルワークロードにも拡大する予定です。ルーティングは依然として進化中の課題です。どこで、どのレベルでルーティングを行うかを決定するには、ワークロードと展開アーキテクチャに対する深い理解が必要です。今後のアップデートもぜひご注目ください。

LinkConclusion

最適なルーティングポリシーはワークロードに依存します。Ray Serve LLM では、すぐに組み合わせて使えるように設定可能なルーティングポリシーを提供しています。一貫したハッシュを用いたセッションアフィニティにより、多段対話における KV キャッシュの再利用を強力にサポートしつつ、セッションをレプリカ間でバランスよく分散できます。KVAwareRouter はさらに一歩進んで、KV キャッシュの重複とトークン負荷の両方をバランスさせることで、より細粒度な負荷分散を実現し、異種ワークロードにおける TTFT（Time to First Token）、TPOT（Time Per Output Token）、スループットを改善します。

最も重要な点は、KV キャッシュの重複だけをルーティングの目的にすべきではないということです。残りのプリフィル作業と進行中のデコード負荷の両方を捉える「トークン負荷」の方が、各エンジン上の実際の作業量をより正確に反映しています。この負荷をバランスさせつつ、KV キャッシュの再利用も活用することで、より効率的な LLM サービングが可能になります。

LinkAcknowledgment

NVIDIA Dynamo チームには、KV インデクサとリクエストスコアリングモジュールをモジュール化し、Ray Serve LLM との統合を可能にしたことに対し、特に感謝いたします。

LinkReproduction Notes

ベンチマーク コード。