sgl-project/sglang-jax

解決する課題

SGL-JAXは、大規模言語モデル（LLM）をGoogle TPU上で高スループットかつ低レイテンシでデプロイするという課題を解決します。要求の厳しいプロダクションワークロードに対して、ハードウェアの利用率を最大化する特化した推論エンジンを提供します。

仕組み

このエンジンは、OpenAI互換のHTTPサーバー、継続的バッチングのためのスケジューラ、およびJAXを使用してモデルを実行するTensor Parallel Workersからなる分散アーキテクチャを採用しています。以下の主要なメカニズムを通じてパフォーマンスを最適化します：

• Radix Tree KV Cache: メモリを効率的に管理し、共通の接頭辞を持つリクエスト間でキャッシュデータを共有できるようにすることで、冗長な計算を削減します。
• Continuous Batching: 入力リクエストを動的にグループ化し、TPUの稼働率を最大限に維持します。
• FlashAttention: 高性能なカーネルを統合し、長いシーケンスのAttention計算を高速化します。
• Tensor Parallelism: 大規模なモデルを複数のTPUデバイスに分割し、単一チップのメモリを超えるモデルの処理を可能にします。

対象ユーザー

Google TPUハードウェア上で、大規模なLLMやマルチモーダルモデル（テキストからビデオへの生成や視覚言語モデルなど）をデプロイする開発者および組織向けに設計されています。

ハイライト

• OpenAI互換API: 既存のOpenAIベースのツールやSDKの代替としてそのまま利用可能です。
• 幅広いモデルサポート: Qwen（MoEバリアントを含む）、Llama、Gemma 2、DeepSeekなどに最適化されています。
• マルチモーダル機能: テキストからビデオへの生成（Wan 2.1/2.2）および視覚言語（Qwen2.5-VL）モデルをサポートしています。
• JAXベース: TPU上での高性能な実行のためにゼロから構築されています。