AIモデルとプラットフォーム

SGLang:構造化言語モデルプログラムの効率的な実行

mm
Unite.AI を Google の優先ソースに追加

大規模言語モデル(LLM)は、複数の生成呼び出し、先進的なプロンプティング技術、制御フロー、構造化された入出力が必要な複雑なタスクにますます利用されています。しかし、これらのアプリケーションをプログラミングし、実行するための効率的なシステムは不足しています。SGLangは、この問題に対処するために、構造化言語モデルプログラムの効率的な実行を提供することを目的とした新しいシステムです。SGLangは、フロントエンド言語とランタイムで構成されています。フロントエンドは、生成と並列性制御のプリミティブを提供してプログラミングを簡素化し、ランタイムは、RadixAttentionによるKVキャッシュの再利用や圧縮された有限状態マシンによる構造化出力の高速化などの新しい最適化によって実行を高速化します。実験では、SGLangは、さまざまな大規模言語モデルと多モーダルモデルで、最先端の推論システムと比較して、最大6.4倍のスループットを達成することが示されています。これらのモデルは、エージェント制御、論理的推論、少샷ラーニングベンチマーク、JSONデコーディング、リトリーバ増強生成パイプライン、多ターンチャットなどのタスクを処理します。

最近のLLMの能力の向上により、そのユーティリティが拡大し、より広範な一般タスクを処理し、自律エージェントとして機能できるようになりました。これらのアプリケーションでは、LLMは、ツールの使用、複数の入力モダリティ、さまざまなプロンプティング技術(例:少샷ラーニング、自己一貫性、スケルトンオブソシエーション、ツリーオブソシエーション)を介して、複数ラウンドの計画、推論、外部環境との相互作用に従事します。これらの新しいユースケースでは、複雑なタスクを完了するために、多くの場合、依存関係のあるLLM生成呼び出しが必要になります。

これは、単純なチャットから、LLMのプログラム的な使用に移行することを示しています。ここで、プログラムはLLMの生成プロセスをスケジュールして制御します。これらのプログラムは「言語モデルプログラム」(LMプログラム)と呼ばれます。先進的なプロンプティング技術とエージェントワークフローは、LMプログラムの範囲内にあります。LMプログラムには、2つの共通の特性があります。(1)LMプログラムは、通常、複雑なタスクを完了し、全体の品質を向上させるために、制御フローを伴う複数のLLM呼び出しを含みます。(2)LMプログラムは、構造化された入力を受け取り、構造化された出力を生成します。これにより、LMプログラムを組み合わせて既存のソフトウェアシステムに統合することができます。

この記事では、SGLangフレームワークについて深く掘り下げ、其のアーキテクチャを探り、其のパフォーマンスを分析し、最先端のフレームワークと比較します。では、始めましょう。

SGLangの紹介

LMプログラムの広範な使用にもかかわらず、現在のシステムは、それらを効率的に表現し、実行することができません。SGLangは、LMプログラムの効率的な使用に関連する2つの主要な課題を特定しています。

  • プログラミングの複雑さ:LLMの非決定的な性質により、LMプログラムの開発は面倒で難しいです。これには、広範な文字列操作、プロンプトの実験的な調整、脆い出力の解析、複数の入力モダリティの処理、並列性メカニズムの実装が含まれます。これらの複雑さは、さらにもう一つのプログラムの読みやすさを大幅に低下させます。
  • 実行の非効率性:LMプログラムの実行は、冗長な計算とメモリ使用により非効率的です。最先端の推論エンジンは、待ち時間を低減し、スループットを向上させるために最適化されていますが、ワークロードに関する直接の知識が欠けているため、重大な非効率性が生じます。KVキャッシュの再利用は、生成推論に不可欠な再利用可能な中間テンソルで構成されており、LLM呼び出しの共通プレフィックスを共有するLLM呼び出しの間でKVキャッシュを再利用するための効果的なメカニズムが不足しています。また、構造化出力(例:JSONモード)の制約付きデコーディングは、既存のシステムが1回の呼び出しから1トークンのみをデコーディングするため、非最適化されています。

これらの課題に対処するために、SGLangは、LLMの構造化言語を導入します。基本的な考え方は、LMプログラムのマルチコール構造を効率的な実行のために体系的に利用することです。次の図に示すように、SGLangは2つの部分で構成されています。フロントエンド言語とバックエンドランタイムです。

フロントエンドは、LMプログラムのプログラミングを簡素化し、ランタイムは、実行を高速化します。これらの部分は、パフォーマンスを向上させるために共同で動作するか、独立して機能することができます。

SGLangは、Pythonに埋め込まれたドメイン固有言語であり、生成(例:extend、gen、select)と並列性制御(例:fork、join)に対するプリミティブを提供します。Pythonの制御フローとライブラリと互換性があり、ユーザーがネイティブのPython構文で高度なプロンプティングワークフローを簡単に開発できるようにします。SGLangには、インタープリターとコンパイラーが含まれています。インタープリターは、プロンプト状態をストリームとして管理し、プリミティブ操作をストリームに非同期的に実行することで、同期とプログラム内並列性の適切な制御を保証します。さらに、SGLangプログラムはトレースされ、コンパイルされてさらに最適化できます。SGLangのランタイムは、LMプログラムの実行を高速化するためのいくつかの新しい最適化を提案します。

  • RadixAttention:この技術は、複数の生成呼び出しの間でKVキャッシュを自動的に再利用することを可能にします。既存の推論エンジンでは、リクエストごとにKVキャッシュが破棄され、複数の呼び出しの間で再利用できず、実行を遅くします。SGLangは、KVキャッシュをラディックス木内にLRUキャッシュとして保持し、ラディックス木を使用して効率的なマッチング、挿入、追い出しを実行します。これにより、ランタイムはさまざまな再利用パターンを効率的に処理できます。
  • 圧縮された有限状態マシン:この技術は、構造化出力の制約付きデコーディングを高速化します。既存のシステムは、次のトークンのみに従って制約を適用し、1回の呼び出しから1トークンのみをデコーディングできます。代わりに、SGLangは、制約を分析し、それらを表すために圧縮された有限状態マシンを構築し、可能な場合はマルチトークンパスを1ステップのパスに圧縮し、複数のトークンを一度にデコーディングして高速化します。
  • API推測的実行:APIのみのモデル(例:OpenAIのGPT-4)に対して、SGLangは、API推測的実行を導入して、マルチコールプログラムを最適化します。

SGLangを使用して、エージェント制御、論理的推論、少샷ラーニングベンチマーク、JSONデコーディング、リトリーバ増強生成パイプライン、多ターンチャット、多モーダリティ処理などのさまざまなLLMアプリケーションが実装されました。パフォーマンスは、Llama-7B/70B、Mistral-8x7B、LLaVA-v1.5-7B(画像)、LLaVA-NeXT-34B(ビデオ)などのモデルで、NVIDIA (NVDA ) A10GおよびA100 GPUでテストされました。実験結果は、SGLangが、さまざまなワークロード、モデル、ハードウェア設定で、Guidance、vLLM、LMQLを含む既存のプログラミングおよび推論システムと比較して、最大6.4倍のスループットを達成することを示しています。

SGLang:プログラミングモデルと方法論

SGLangのプログラミングモデルは、ランニングエクサンプルを介して導入され、言語プリミティブと実行モードを説明し、ランタイム最適化の機会を概説します。このモデルは、マルチコールワークフロー(例:文字列操作、API呼び出し、制約指定、並列性)における面倒な操作を、柔軟で組み合わせ可能なプリミティブを提供することで簡素化します。SGLangは、Pythonに埋め込まれたドメイン固有言語です。次の図は、ブランチソルブマージプロンプティング方法を使用して画像についてのエッセイを評価するプログラムを示しています。

関数multi_dimensional_judgeは3つの引数を取ります:`s`、`path`、`essay`。sはプロンプト状態を管理し、pathは画像ファイルパス、essayはエッセイのテキストです。新しい文字列とSGLangプリミティブは、+=演算子を使用して状態sに追加され、実行されます。まず、関数は画像とエッセイをプロンプトに追加します。次に、エッセイが画像に関連しているかどうかをselectを使用して確認し、結果をs[“related”]に保存します。関連している場合、プロンプトは3つのコピーに分岐され、genを使用して結果をf[“judgment”]に保存します。次に、判断を結合し、要約を生成し、レター評価を割り当てます。最後に、結果をJSON形式で返します。JSON形式は、正規表現制約regexによって定義されます。SGLangは、このプログラムを大幅に簡素化します。OpenAI APIのようなインターフェイスを使用する場合、同等のプログラムは2.1倍のコード行数が必要になります。これは、手動の文字列操作と並列性制御のためです。

SGLangは、プロンプト状態、生成、並列性制御に対するプリミティブを提供します。これらは、Python構文とライブラリと共に使用できます。以下はプリミティブの一覧です。

gen:モデルを呼び出して生成し、結果を指定された名前の変数に保存します。結果を特定のグラマーに従うように制約するために、正規表現を指定できます(例:JSONスキーマ)。

  • select:モデルを呼び出して、リストからの最も可能性の高いオプションを選択します。
  • += または extend:文字列をプロンプトに追加します。
  • [変数名]:生成の結果を取得します。
  • fork:プロンプト状態の並列フォークを作成します。
  • join:プロンプト状態の並列フォークを再結合します。
  • imageとvideo:画像とビデオ入力を取り込みます。

SGLangプログラムを実行する最も簡単な方法は、インタープリターを介して行うことです。ここで、プロンプトは非同期ストリームとして扱われ、プリミティブ操作(extendgenselectなど)はストリームに非同期的に提出されます。これらの非ブロッキング呼び出しにより、Pythonコードは生成が完了するのを待たずに実行を続行できます。各プロンプトは、バックグラウンドスレッドのストリームエクスキューターによって管理され、プログラム内並列性が可能になります。生成結果の取得は、結果が準備できたときにブロックします。正しい同期を保証するために。代わりに、SGLangプログラムは計算グラフとしてコンパイルされ、グラフエクスキューターで実行できます。これにより、さらに最適化が可能になります。この論文では、デフォルトでインタープリターモードを使用し、コンパイラモードの結果については付録Dで説明します。SGLangは、独自のSGLangランタイム(SRT)を使用してオープンウェイトモデルをサポートし、OpenAIやAnthropicなどのAPIモデルもサポートします。

LLM用のプログラミングシステムは、高レベル(例:LangChain、DSPy)と低レベル(例:LMQL、Guidance、SGLang)に分類できます。高レベルシステムは、事前定義されたまたは自動生成されたプロンプト(例:DSPyのプロンプトオプティマイザー)を提供します。低レベルシステムは、通常、プロンプトを変更せずに、プロンプトとプリミティブの直接操作を許可します。SGLangは、LMQLやGuidanceと同様の低レベルシステムです。次の表は、これらのシステムの機能を比較しています。

SGLangは、ランタイム効率に重点を置いており、独自の共同設計ランタイムを備えています。これにより、新しい最適化が可能になります。高レベル言語(例:DSPy)は、低レベル言語(例:SGLang)にコンパイルできます。DSPyのバックエンドとしてSGLangを統合して、ランタイム効率を向上させることが示されています。

上記の例は、9つの時間点におけるRadixAttention操作とLRU追い出しポリシーを示しています。ここでは、さまざまなリクエスト(チャットセッション、少샷ラーニングの質問、自己一貫性サンプリング)に対するラディックスツリーの動的割り当てと追い出しの動作を示しています。各ツリーのエッジは、サブ文字列またはトークンのシーケンスを表すラベルを持ち、ノードは、追加された新しいノード(緑)、キャッシュされたノード(青)、追い出されたノード(赤)を示すために色分けされています。

ステップ1:ラディックスツリーは最初から空です。

ステップ2:サーバーは、ユーザーからの「Hello」というメッセージを受け取り、「Hi」というLLM出力を返します。システムプロンプト「You are a helpful assistant」、ユーザーメッセージ「Hello!」、LLMの応答「Hi!」は、ツリー内の新しいノードに接続された単一のエッジとして統合されます。

ステップ3:新しいプロンプトが到着し、サーバーはツリー内でプロンプトのプレフィックス(会話の最初のターン)を見つけ、KVキャッシュを再利用します。新しいターンはツリーに新しいノードとして追加されます。

ステップ4:新しいチャットセッションが開始されます。ステップ3のノードは、2つのチャットセッションがシステムプロンプトを共有できるように2つのノードに分割されます。

ステップ5:2番目のチャットセッションが続きます。しかし、メモリ制限のため、ステップ4のノードが追い出される必要があります。新しいターンは、ステップ4の残りのノードの後に追加されます。

ステップ6:サーバーは、少샷ラーニングのクエリを受け取り、処理し、ツリーに挿入します。ルートノードは、新しいクエリが既存のノードとプレフィックスを共有しないため、分割されます。

ステップ7:サーバーは、さらに少샷ラーニングのクエリのバッチを受け取ります。これらのクエリは、同じセットの少샷例を共有するため、ステップ6のノードが分割されて共有が可能になります。

ステップ8:サーバーは、最初のチャットセッションからの新しいメッセージを受け取り、2番目のチャットセッションのすべてのノードを最も最近使用されていないノードとして追い出します。

ステップ9:サーバーは、ステップ8のノードの質問に対してさらに回答をサンプリングするリクエストを受け取り、スペースを確保するために複数のノードを追い出します。可能な場合は、ラディックスツリーのエッジは、サブ文字列またはトークンのシーケンスを表すラベルを持ち、ノードは、追加された新しいノード(緑)、キャッシュされたノード(青)、追い出されたノード(赤)を示すために色分けされています。

この例は、RadixAttentionがさまざまなリクエストに応じてノードの動的割り当てと追い出しの動作をどのように処理するかを示しています。KVキャッシュの再利用とメモリ管理を効率的に行う方法を示しています。

SGLang:評価と結果

オープンウェイトモデルでの結果

待ち時間とスループットの結果は、次の図に示されています。SGLangは、最大6.4倍のスループットを向上させ、最大3.7倍の待ち時間の短縮を実現します。これらの改善は、KVキャッシュの再利用、プログラム内並列性の活用、制約付きデコーディングの高速化によるものです。

これらのベンチマークでは、キャッシュヒット率は50%から99%の範囲でした。付録13に、すべてのベンチマークの達成されたキャッシュヒット率と最適なキャッシュヒット率がリストされています。SGLangのキャッシュを意識したスケジューリングアプローチは、平均で96%の最適なヒット率に近づいています。

テンソル並列性を持つ大規模モデルでの結果

テンソル並列性を使用した大規模モデル(Mixtral-8x7B、Llama-70B)も、同じベンチマークでテストされ、結果は次の図に示されています。大規模モデルのスピードアップは、小規模モデルで観察されたのと同様の傾向を示しています。これは、SGLangの最適化が大規模モデルにも一般化できることを示しています。GuidanceとLMQLは、テンソル並列性の効率的な実装が不足しているため、省略されました。

多モーダルモデルでの結果

SGLangには、画像とビデオのプリミティブを使用した多モーダルモデルのネイティブサポートがあります。この論文の最適化は、多モーダルモデルと互換性があります。RadixAttentionの場合、入力画像のハッシュが計算され、ラディックスツリーのキーとして使用され、同じ画像からの画像トークンのKVキャッシュの再利用が可能になります。LLaVA-v1.5-7B(画像)がllava-bench-in-the-wildで、LLaVA-NeXT-34B(ビデオ)がActivityNetで実行されました。これらのモデルは他のベースラインシステムで十分にサポートされていないため、Hugging Face Transformersのモデル著者のオリジナル実装がベースラインとして使用されました。次の表に示すように、SGLangはこれらのベンチマークで最大6倍のスループットを提供します。llava-bench-in-the-wildでは、同じ画像について複数の質問が処理され、SGLangランタイムはKVキャッシュを再利用しました。

本稼働へのデプロイ

SGLangは、オープンウェイトモデルを提供するChatbot Arenaにデプロイされています。トラフィックが低いモデルがあるため、各モデルにはSGLangワーカーが1つしかありません。一ヶ月後、LLaVA-Next-34BのRadixAttentionキャッシュヒット率は52.4%、Vicuna-33Bのキャッシュヒット率は74.1%でした。キャッシュヒットは、共通のシステムメッセージ、頻繁に再利用される画像の例、多ターンチャット履歴から来ました。これにより、Vicuna-33Bの最初のトークンの待ち時間が平均1.7倍短縮されました。

最終的な考察

この記事では、SGLangについて説明しました。SGLangは、構造化言語モデルプログラムの効率的な実行を提供することを目的とした新しいシステムです。フロントエンド言語とランタイムで構成されており、ランタイムは、RadixAttentionによるKVキャッシュの再利用や圧縮された有限状態マシンによる構造化出力の高速化などの新しい最適化を提供します。実験では、SGLangは、さまざまな大規模言語モデルと多モーダルモデルで、最大6.4倍のスループットを達成することが示されています。これらのモデルは、エージェント制御、論理的推論、少샷ラーニングベンチマーク、JSONデコーディング、リトリーバ増強生成パイプライン、多ターンチャットなどのタスクを処理します。

職業はエンジニア、心は作家。クナルは、AIとMLを深く愛し理解しているテクニカルライターで、これらの分野の複雑な概念を魅力的で情報の多いドキュメンテーションを通じて簡素化することに尽力しています。