LLM Serving Optimization with Variable Prefill and Decode Lengths

技術要約：可変プリフィルおよびデコード長を伴うLLMサービングの最適化

問題定式化
本論文は、固定されたKey-Value（KV）キャッシュメモリ予算を持つ単一の計算ワーカーにおける、大規模言語モデル（LLM）のオフラインスケジューリング問題に対処するものである。システムは、$t=0$ において同時に到着する $n$ 個のリクエストのバックログを処理する。各リクエスト $i$ は、不均一なプロンプト長（プリフィルサイズ、$s_i$）とレスポンス長（デコードサイズ、$o_i$）によって特徴付けられる。

核心となる課題は、自己回帰的なデコーディング中のメモリ消費の動的な性質から生じる。プリフィルフェーズは $s_i$ に比例するメモリを消費する一方で、デコードフェーズは生成されるトークンごとにメモリ消費量が増加する（$j$ 番目のトークンに対して $s_i + j$）。その結果、ワーカーが同時に処理できるリクエスト数は、アクティブなリクエストの生成進行状況に基づいて動的に変動する。目的は、任意のタイムステップにおける集約メモリ使用量が容量 $M$ を超えないという制約の下で、総エンドツーエンドレイテンシ（Total End-to-End Latency, TEL）、すなわち全リクエストの完了時刻の総和を最小化することである。

著者らは、先行研究（Jaillet et al., 2025）に見られる一様入力サイズの仮定を緩和し、短文の対話プロンプトと長文ドキュメントの要約タスクが混在する現実的なシナリオを考慮している。

理論的知見と課題
本論文は、入力サイズの一様性を仮定を緩和することが、スケジューリング問題の複雑さを根本的に変えることを確立している：

1. NP困難性: 不均一な $s_i$ と $o_i$ を持つ一般化されたスケジューリング問題は、制限付き3-Partition問題への帰着を通じて、NP困難であることが証明されている。
2. 標準的なヒューリスティックの失敗: FCFS（先着順）、MC-SF（最短出力優先）、およびTotal-Size-Based（総サイズベース）の優先順位付けを含む標準的なポリシーは、この不均一な設定においては近似比が非有界になることが示されている。具体的には、入力は大きいが出力が小さいリクエストは一時的に大量のメモリを消費する可能性があり、一方で入力は小さいが出力が長いリクエストは長期間メモリを占有する可能性があり、単純なヒューリスティックでは解決できないトレードオフが生じる。

手法：Sorted-F アルゴリズム
これらの課題に対処するため、著者らは定数倍近似保証を実現するように設計されたスケジューリングアルゴリズムである Sorted-F を提案する。このアルゴリズムは2つのフェーズで動作する：

1. バッチ構築（フェーズ1）: アルゴリズムは、順序付けられたリストを形成するために、実行可能なリクエストのバッチを反復的に選択する。選択は、リクエスト集合 $X$ に対して定義される新しい品質指標 $F(X)$ によって導かれる：
   $$F(X) = \frac{\sum_{r_i \in X} o_i}{|X|^2}$$
   この指標は、バッチサイズ（$|X|$）と集約されたダウンストリーム・デコードコスト（$\sum o_i$）のバランスを取る。$F(X)$ を最小化することは、「スループット密度」（平均デコード時間あたりの完了リクエスト数）を最大化することと同等である。選択された各バッチ内では、リクエストは非減少の出力長に従ってソートされる。
2. 実行（フェーズ2）: システムは離散的なタイムステップにわたってスケジュールを実行し、連続バッチングおよびプリフィル/デコード混合を利用しながら、メモリの実現可能性制約が満たされる限り、順序付けられたリストからリクエストを現在のバッチに受け入れる。

近似保証
著者らは、Sorted-F が最適なオフラインスケジュールの少なくとも 48 倍の定数近似比を達成することを証明している。この証明は、Sorted-F のスケジュールと最適スケジュールの両方を比較可能な形式に変換する、一連の構造的変換（分離、グルーピング、およびアライメント）を含み、各ステップでのレイテンシの増加が定数倍の範囲内に収まることを示している。

実践的な実装
大規模な $n$ に対してフェーズ1の組合せ最適化が計算コストの高いことを認識し、論文では3つの実用的なソルバーを提案している：

• 厳密な動的計画法（Exact Dynamic Programming）: $O(n^2 M)$ の計算量で、小規模なインスタンス（$n \le 100$）に対して厳密解を提供する。
• ローカル・スワップ探索（Local Swap Search）: 中規模なスケール（$n \le 500$）のためのヒューリスティックであり、ペアワイズ交換（$O(n^2)$）を通じて初期の貪欲解を反復的に改善する。
• 分位点貪欲選択（Quantile Greedy Selection）: 大規模システム（$n \ge 1000$）のための線形時間ヒューリスティックであり、統計的サンプリングを使用して外れ値をフィルタリングし、容量を貪欲に埋める。

さらに、リクエストを順序付けるための線形計画法（LP）緩和の解を活用したヒューリスティックである Sorted-LP、および推論中の出力長の予測不確実性を扱うために動的に推定値を精緻化する適応型バリアントである Sorted-F+Amin を導入している。

実験結果
混合ワークロード（LMSYS-ChatデータセットとArxiv長文ドキュメント要約の組み合わせ）を用いた数値実験により、以下のことが示された：

• レイテンシ削減: Sorted-F（Local Swap Searchを使用）は、FCFS、MC-SF、およびLP誘導ベースラインを一貫して上回った。2000個のリクエストを含む80/20の短文/長文混合において、Sorted-F は FCFS に対して 4.87倍、MC-SF に対して 2.09倍、平均レイテンシを削減した。
• 最適性ギャップ: 取り扱い可能なインスタンスにおいて、Sorted-F と LP 緩和の下限との経験的なギャップは小さく（約1.03–1.09）、これは、理論的な境界は緩いものの、実際には Sorted-F がほぼ最適に動作していることを示唆している。
• ワークロードへの感度: Sorted-F の利点は、メモリが拘束条件となる不均一かつ過負荷のレジームで最も顕著である。均一または低負荷のレジームでは、その優位性は減少する。
• オンライン性能: ポアソン到着下でテストされた後退ホライゾン・バリアント（Sorted-F-Online）は、過負荷のオンラインレジームにおいて MC-SF や FCFS に対して大幅な改善を示したが、安定した低レイテンシのレジームでは効果は限定的であった。

意義と主張
本論文は、可変プリフィル長は単なる些細なモデリングの詳細ではなく、標準的なスケジューリング・ヒューリスティックを無効にし、問題をNP困難にする根本的な要因であると主張している。主要な貢献は Sorted-F アルゴリズム であり、これは不均一な入力/出力制約の下でのLLMスケジューリングに対する最初の定数倍近似保証を提供するものである。本研究は、理論的なスケジューリングの限界と実用的なLLMサービングの間の溝を埋め、出力長の完全な予測を必要とせずに、混合ワークロードのシナリオでレイテンシを大幅に削減する、スケーラブルでメモリを意識したフレームワークを提供している。著者らは、理論的な境界は48であるが、経験的な結果は、現実的な設定においてアルゴリズムが最適解に非常に近い形で動作していることを示唆していると強調している。