スループットが5倍に向上、バックエンドシステムとフロントエンド言語を共同設計するLLMインターフェースが登場

大規模言語モデル (LLM) は、複数の連鎖生成呼び出し、高度なプロンプト技術、制御フロー、および外部環境との対話を必要とする複雑なタスクにますます使用されるようになっています。しかし、これらのアプリケーションをプログラミングおよび実行するための既存の効率的なシステムには、重大な欠点があります。

現在、オープンソース コミュニティの研究者は、LLM 用の構造化生成言語 (SGLang) を提案しています。 SGLang は、バックエンドのランタイム システムとフロントエンド言語を共同で設計することで、LLM との相互作用を強化し、LLM をより高速かつ制御可能にします。機械学習分野の著名な学者であり、CMU の助教授である Chen Tianqi 氏もこの研究を転送しました。

一般的に、SGLang の貢献には主に次のものが含まれます。

• バックエンドでは、研究チームは、複数の LLM にわたって呼び出しを自動的かつ効率的に生成する KV キャッシュ再利用テクノロジーである RadixAttention を提案しました。
• フロントエンドでは、研究チームは生成プロセスを制御するために Python に組み込まれた柔軟なドメイン固有言語を開発しました。言語は、インタープリタ モードまたはコンパイラ モードで実行できます。

バックエンド コンポーネントとフロントエンド コンポーネントが連携して、複雑な LLM プログラムの実行とプログラミングの効率を向上させます。

この研究では、SGLang を使用してエージェント、推論、抽出、対話、および少数ショット学習タスクを含む一般的な LLM ワークロードを実装し、NVIDIA A10G GPU で Llama-7B および Mixtral-8x7B モデルを採用しました。下の図 1 と 2 に示すように、SGLang のスループットは既存のシステム (Guidance および vLLM) と比較して 5 倍向上しています。

図 1: LLM タスクにおけるさまざまなシステムのスループット (A10G、FP16 上の Llama-7B、テンソル並列性 = 1)

図 2: LLM タスクにおけるさまざまなシステムのスループット (A10G 上の Mixtral-8x7B、FP16、テンソル並列性 = 8)

バックエンド: RadixAttention を使用した KV キャッシュの自動再利用

SGLang ランタイムの開発中に、複雑な LLM プログラムを最適化するための鍵となる KV キャッシュの再利用が研究者によって発見されましたが、これは現在のシステムでは適切に処理されていません。 KV キャッシュの再利用とは、同じプレフィックスを持つ異なるプロンプトが中間 KV キャッシュを共有し、冗長なメモリと計算を回避できることを意味します。複数の LLM 呼び出しを含む複雑なプログラムでは、さまざまな KV キャッシュ再利用パターンが存在する可能性があります。下の図 3 は、LLM ワークロードでよく見られる 4 つのパターンを示しています。一部のシステムでは特定のシナリオで KV キャッシュの再利用を処理できますが、通常は手動の構成とアドホックな調整が必要になります。さらに、再利用パターンの多様性のため、既存のシステムは手動で構成してもすべてのシナリオに自動的に適応することはできません。

図 3: KV キャッシュ共有の例。青いボックスは共有可能なプロンプト部分、緑のボックスは共有不可能な部分、黄色のボックスは共有不可能なモデル出力です。共有可能な部分には、少数の学習例、自己一貫性の質問、複数ターンの対話における会話履歴、思考ツリーにおける検索履歴が含まれます。

これらの再利用の機会を体系的に活用するために、本研究では実行時に KV キャッシュを自動的に再利用する新しい方法である RadixAttention を提案しました。このメソッドは、ビルド要求の完了後に KV キャッシュを破棄する代わりに、プロンプトとビルド結果の KV キャッシュを基数ツリーで保持します。このデータ構造により、効率的なプレフィックス検索、挿入、および削除が可能になります。この研究では、キャッシュ ヒット率を向上させるために、キャッシュを考慮したスケジューリング戦略を補完した、最近最も使用されていない (LRU) 削除戦略を実装しました。

基数木は、トライ (プレフィックス木) のスペースを節約する代替手段として使用できます。一般的なツリーとは異なり、基数ツリーのエッジは単一の要素だけでなく、異なる長さの要素のシーケンスでもマークできるため、基数ツリーの効率が向上します。

この研究では、キーとして機能するトークン シーケンスと値として機能する対応する KV キャッシュ テンソル間のマッピングを管理するために基数ツリーを利用します。これらの KV キャッシュ テンソルはページ レイアウトで GPU に保存され、各ページのサイズはトークンに相当します。

GPU のメモリ容量が限られていることを考慮すると、無限の KV キャッシュ テンソルを再トレーニングすることは不可能であり、排除戦略が必要になります。この研究では、LRU 削除戦略を採用して、リーフ ノードを再帰的に削除します。さらに、RadixAttention は、連続バッチ処理やページング アテンションなどの既存の手法と互換性があります。マルチモーダル モデルの場合、RadixAttention は画像トークンを処理するように簡単に拡張できます。

次の図は、複数の受信要求を処理しながら基数ツリーがどのように維持されるかを示しています。フロントエンドは常に完全なプロンプトをランタイムに送信し、ランタイムはプレフィックスのマッチング、再利用、およびキャッシュを自動的に実行します。ツリー構造は CPU 上に保存され、メンテナンスのオーバーヘッドが低くなります。

図 4. LRU 削除ポリシーを使用した RadixAttention 操作の例 (9 つのステップで説明)。

図 4 は、さまざまな要求に応じて基数ツリーが動的に進化する様子を示しています。これらのリクエストには、2 つのチャット セッション、小規模サンプルの学習クエリのバッチ、および自己一貫性サンプリングが含まれます。各ツリー エッジには、サブ文字列またはトークン シーケンスを表すラベルが付けられます。ノードはさまざまな状態を反映するために色分けされています。緑は新しく追加されたノード、青はその時点でアクセスされているキャッシュ ノード、赤は削除されたノードを表します。

フロントエンド: SGLang による簡単な LLM プログラミング

フロントエンドでは、この研究では、高度なプロンプト技術、制御フロー、マルチモーダル性、デコード制約、および外部インタラクションの表現を可能にする、Python に埋め込まれたドメイン固有言語である SGLang を提案しています。 SGLang 関数は、OpenAI、Anthropic、Gemini、ローカル モデルなどのさまざまなバックエンドで実行できます。

図 5. SGLang を使用した多次元記事スコアリングの実装。

図5に具体的な例を示します。多次元の記事スコアリングを実現するために、ブランチ・解決・マージ プロンプト手法を使用します。この機能は、LLM を使用して複数の側面から記事の品質を評価し、判断を統合し、要約を生成し、最終的な成績を割り当てます。強調表示された領域は、SGLang API の使用を示しています。 （１）forkはpromptの複数の並列コピーを作成します。 （２）genはLLM生成を呼び出し、その結果を変数に格納する。この呼び出しは非ブロッキングなので、複数の生成呼び出しをバックグラウンドで同時に実行できます。 （３）[variable_name]生成された結果を検索します。 （４）必要に応じて生成に制約を課す。 （５）runは引数を使ってSGLang関数を実行します。

このような SGLang プログラムが与えられた場合、それをインタープリター経由で実行するか、データフロー グラフとしてトレースし、グラフ エグゼキューターを使用して実行することができます。後者の場合、コードの移動、命令の選択、自動チューニングなど、コンパイラの最適化の余地が生まれます。

SGLang の構文は Guidance に大きく影響を受けており、新しいプリミティブを導入し、プログラム内並列処理とバッチ処理も処理します。これらすべての新機能が、SGLang の優れたパフォーマンスに貢献します。

ベンチマーク

研究チームは、一般的な LLM ワークロードでシステムをテストし、達成されたスループットを報告しました。

具体的には、この研究では、1 つの NVIDIA A10G GPU (24GB) 上で Llama-7B を、8 つの NVIDIA A10G GPU 上で Mixtral-8x7B を FP16 精度を使用したテンソル並列処理でテストし、vllm v0.2.5、guidance v0.1.8、Hugging Face TGI v1.3.0 をベースライン システムとして使用しました。

図 1 および 2 に示すように、SGLang はすべてのベンチマークでベースライン システムを上回り、スループットが 5 倍向上しています。また、レイテンシの点でもパフォーマンスが優れており、特にプレフィックス キャッシュ ヒットが大きなメリットをもたらす最初のトークン レイテンシでは優れています。これらの改善は、RadixAttention の自動 KV キャッシュの再利用、インタープリターによって実現されるプログラム内並列性、およびフロントエンド システムとバックエンド システムの共同設計によるものです。さらに、アブレーション研究では、キャッシュ ヒットがない場合でも、実行時に RadixAttention を常に有効にすることで、目立ったオーバーヘッドが発生しないことが示されています。