こんにちは。コロンビア大学の博士課程でAI・セキュリティの研究をしているKoukyosyumeiです。
最近、OpenAIやAnthropicなどのLLM APIを統一的なインターフェースで利用するためのAI Gatewayが数多く登場しており、例えばPython製のLiteLLMなどのOSSが注目を集めています。LiteLLMは、モデルの振り分けだけでなく、APIキーの管理、レート制限、利用料金の記録なども担当できる部分が強みになっています。
今回、これらの機能をRustで再実装し、オリジナルのLiteLLMと比較したところ、通常のチャットAPIで約41倍、ストリーミングでも約40倍のスループットが得られ、さらにメモリ使用量も約4.4GBから68MBに削減できました。
AI Gatewayとは？
AI Gatewayは、アプリケーションとLLMプロバイダーの間に配置するプロキシです。
Application
|
v
AI Gateway
|
+------------+------------+
| | |
v v v
OpenAI Anthropic Gemini
例えば、認証、利用可能なモデルの制限、予算管理、複数モデルへのルーティング、ストリーミングレスポンスの中継などを行います。
しかし、AIエージェントが大量のLLMリクエストを発行するようになると、Gateway自身のCPU使用率やメモリ使用量も無視できなくなります。そこで今回は、LiteLLMのPython製プロキシをRustで再実装した場合、どのくらい高速化できるか試してみました。
Rustで再実装してみた
今回の実装では、Rust Webフレームワーク h5i-app を使用しました。h5i-appはAxumベースのフレームワークで、アプリケーションのビジネスロジックを独立したkernelとして実装し、Lean 4で仕様を形式検証したり、GPUによる高速化などを図れることが特徴となっています。
AI Gatewayでは、HTTP通信と、認証・予算管理・ルーティングなどのロジックを分離しています。
HTTP Request
|
v
Gateway Server
|
v
Rust Kernel
|
+-------+-------+
| |
Admit Settle
| ^
v |
Provider API ------+
|
v
HTTP Response
基本的な流れとしては、各リクエストにおいて、まずAdmitでAPIキー、モデルへのアクセス権限、レート制限、予算などを確認し、利用予定額を予約します。その後、OpenAIやAnthropicなどのAPIを呼び出し、Settleで実際のトークン使用量に基づいて料金を確定します。また、チャット、embeddings、ストリーミング、複数プロバイダーへのルーティングなどの機能もサポートされています。
ベンチマーク結果
それでは、実際のパフォーマンスを比較してみます。LLMの推論時間が結果に影響しないように、両方のGatewayから同じモックLLMサーバーを呼び出す構成で測定しました。
比較対象は、LiteLLM Python 1.106.2と、今回実装したh5i-gatewayです。
8 CPU、64並列接続での結果は以下のとおりです。
| 処理 | LiteLLM Python | Rust版 | 高速化 |
|---|---|---|---|
| Chat API | 801 req/s | 33,200 req/s | 41.4倍 |
| Streaming | 277 req/s | 11,208 req/s | 40.5倍 |
| Embeddings | 572 req/s | 17,200 req/s | 30.1倍 |
| PostgreSQL + APIキー | 480 req/s | 1,389 req/s | 2.9倍 |
メモリ使用量も大幅に削減され、約65分の1のメモリで動作しています。
| 実装 | メモリ使用量 |
|---|---|
| LiteLLM Python（8 workers） | 4.4 GB |
| Rust版 | 68 MB |
さらに、チャットAPIのレイテンシを比較すると、中央値（p50）は44.4msから1.91msに、p99は346msから3.08msに短縮されました。また、実際のLLMに近い50msの上流遅延を設定したベンチマークでも、256並列接続でRust版は4,887 req/s、LiteLLM Python版は657 req/sを記録しました。
詳細なベンチマーク条件、結果、再現スクリプトはこちらで公開しています。
最適化
今回の開発では、単純にPythonをRustに置き換えただけでなく、いくつかの最適化を行いました。
1. HTTPとストリーミング処理の最適化
LLMのストリーミングでは、SSE（Server-Sent Events）を使って小さなデータを連続的に送信します。このような小さな書き込みが重なると、TCPのdelayed ACKなどの影響もあり、余計な待ち時間が発生していきます。
そこで、受け付けたTCPソケットでTCP_NODELAYを有効にしたところ、64並列接続のストリーミング性能が約1,550 req/sから11,600 req/sに改善しました。SSEのように小さなメッセージを大量に送信する処理では、ネットワーク設定も重要だと分かります。
2. PostgreSQLへの書き込みをバッチ化
また、AI Gatewayでは、リクエストごとに予算やトークン使用量を更新する必要があります。今回のRust版では、予算を正確に管理するため、リクエストの受付時と完了時にPostgreSQLのトランザクションを使用しましたが、毎回個別にトランザクションを実行すると、データベースとの通信がボトルネックになります。
そこで、同じテナントから届いた複数の更新をまとめて処理するGroup Commitを実装し、さらに、同じ行に対する複数の変更をまとめて書き込むようにしました。この最適化によって、64並列接続でのPostgreSQLを使った処理が、約70 req/sから877 req/sまで改善しました。
3. 必要なデータだけを読み込む
初期実装では、リクエストを処理する際に、テナントに属する大量のデータをまとめて読み込んでいたため、APIキーの数が増えるほど、読み込みコストも増大してしまいます。
そこで、各コマンドが必要とする行だけを取得するRead Planを実装しました。例えば、あるAPIキーからのリクエストを認可する場合、そのキーや所属チーム、利用可能なモデル、レート制限など、判断に必要なデータだけを読み込みます。この変更によって、1万個のAPIキーを持つテナントでのスループットが、約336 req/sから1,260 req/sまで改善しました。
Lean 4による形式検証
今回使用したh5i-appでは、Rustで実装したkernelをLean 4に変換し、形式的に検証することもできます。例えば、予算を超えるリクエストが許可されないことや、管理者権限のないAPIキーが制限値を変更できないことなどを検証しています。特に今回は、バッチ処理による最適化が、逐次実行と同じ結果になることなどを証明しています。
まとめ
今回は、LiteLLMのAI GatewayをRustで再実装し、パフォーマンスを比較してみました。結果として、通常のチャットAPIで約41倍、ストリーミングで約40倍のスループットを実現し、メモリ使用量も約65分の1に削減できました。一方で、PostgreSQLを使った場合の高速化は約2.9倍にとどまっており、データベースへのアクセス方法やトランザクション設計が重要なことも分かりました。
実装とベンチマークはOSSとして公開しています。興味があればぜひ試してみてください！
- h5i-ai-gateway: https://github.com/h5i-dev/h5i-ai-gateway
- h5i-app: https://github.com/h5i-dev/h5i
- ベンチマーク: https://github.com/h5i-dev/h5i-ai-gateway/tree/main/bench