当社がCloudflareのネットワーク向けに最も効率的な推論エンジンを構築した方法

推論は、チャットボット応答、AIエージェント、自動運転車の判断、不正検出などの、今日の最も強力なAI製品のいくつかを強化しています。問題は、これらの製品の1つをハイパースケーラーの上に構築する場合、推論タスクを実行するために、大規模な集中型データセンターから高価なGPUをレンタルする必要が出てくる可能性が高いことです。このモデルは、Cloudflareでは機能しません。Cloudflareのグローバル分散ネットワークと、大規模なマルチGPUノードを使用した一般的な集中型AIのデプロイメントとの間には矛盾があるのです。世界のインターネット人口の95%から50ミリ秒以内にある、無駄のない、高速で広く分散したネットワーク上で自社のコンピューティングを運用する企業として、推論タスクを他社よりも効率的に実行する必要があります。

AIモデルが大型化し、複雑化していることで、この状況はさらに悪化しています。Llama 4herdやgpt-ossといったこれらのモデルのサポートをし始めた時、私たちはGPUを追加購入することでスケーリングの問題にお金を注ぎ込むことはできないと考えました。アイドル状態のあらゆる容量を利用し、各モデルが展開される場所に俊敏に対応する必要がありました。

広く使用されているオープンソースの推論とサービスエンジンであるvLLM上でモデルのほとんどを実行した後、私たちはエッジのGPUを十分に活用できないことが判明しました。個人のデバイスからデータセンターまで、非常に幅広いハードウェア上で実行できますが、大規模なデータセンター向けに最適化されています。vLLMは、特定のモデルに対応する強力なハードウェア上で専用の推論サーバーとして実行することで、真に活躍します。しかし、動的なワークロードや分散ネットワーク、さらには他のサービスと並行してエッジで推論を実行する際の固有のセキュリティ制約への対応には、あまり最適化されていません。

そこで私たちは、Cloudflare推論ワークロードのニーズを今後何年にも満たせるものを構築しようと決めました。InfireはRustで書かれたLLM推論エンジンで、さまざまなテクニックを用いてメモリ、ネットワークI/O、GPU使用率を最大化します。より少ないGPUでより多くのリクエストに対応でき、CPUオーバーヘッドを大幅に削減できるため、ネットワーク全体の時間、リソース、エネルギーを節約できます。

最初のベンチマーキングでは、H100 NLP GPUを搭載した無負荷のマシンで、InfireがvLLM 0.10.0よりも最大7%速く推論タスクを完了することが示されました。実際の負荷がかかるインフラでは、パフォーマンスは大幅に向上します。

現在、InfireはWorkers AIのLlama 3.1 8Bモデルを強化しており、今すぐ@cf/meta/llama-3.1-8b-instructでテストできます。

業界の努力により、推論はここ数年で大幅に改善されました。vLLMは最近、KVキャッシュの最適化、バッチ処理の改善、Flash Attention 3の実装などの機能を備えたvLLM V1エンジンの最近のリリースにより、ここでの先端を行くことになりました。vLLMは、ほとんどの推論ワークロードに最適であり、現在、いくつかのサービスに使用しています。Workers AIカタログ。AIのワークロードとカタログが拡大するにつれて、まさにハードウェアとパフォーマンスの要件に合わせて推論を最適化する必要性が高まっています。

Cloudflareは新しいインフラストラクチャの多くをRustで書いており、vLLMはPythonで書かれています。PythonはMLワークロードのプロトタイプとして優れた言語であることが証明されていますが、効率性を最大化するには低レベルの実装の詳細を制御する必要があります。複数の抽象化レイヤーとPythonライブラリを使って低レベルの最適化を実装すると、不必要な複雑さが加わり、CPUパフォーマンスがかなり低下します。これは単に解釈済み言語としてのPythonの非効率性が原因です。

私たちは、自分たちが使用しているオープンソースプロジェクトに貢献したいと思っていますが、今回は私たちの優先順位がvLLMプロジェクトの目的と合致しない可能性があるので、私たちのニーズに合わせてサーバーを書くことにしました。例えば、vLLMは、マルチインスタンスGPU（MIG）を使用せずに同じGPU上で複数のモデルを共同ホストすることをサポートしていず、ダウンタイムを最小限に抑えるために、同じGPU上で複数のモデルを動的にスケジュールできる必要があります。また、社内のAI研究チームは、vLLMへのアップストリームは不可能ではないにしろ、困難なユニーク機能を模索しています。

最後に、コードを安全に実行することはプラットフォーム全体における最優先事項であり、Workers AIも例外ではありません。強力なサンドボックス化なしに他のサービスと並行してエッジノードで実行するサードパーティのPythonプロセスは信頼できません。したがって、gvisorを経由してvLLMを実行せざるを得ません。追加の仮想化レイヤーがあることで、vLLMに別のパフォーマンスのオーバーヘッドが加わります。さらに重要なのは、ただでさえかなり長いvLLMインスタンスの起動と取り下げの時間を増加させることです。エッジノードに全負荷がかかると、gvisor経由で実行されるvLLMは2.5ものCPUコアを消費し、他の重要なサービスとCPU時間を競合させられます。その結果、vLLMが遅くなり、結果としてGPU使用率が低下します。

Infireを開発する際、推論効率に関する最新の研究を取り入れてきました。実際に構築したものを詳しく見てみましょう。

Infireは、OpenAI互換HTTPサーバー、バッチ処理、Infireエンジンという3つの主要コンポーネントで構成されています。

^{Infireのアーキテクチャの概要}

Cloudflareの自動スケーリングサービスによって、あるデータセンター内の特定ノードでモデルの実行がスケジュールされた場合、まずモデルの重みをR2オブジェクトストレージから取得しなければなりません。重みがダウンロードされると、将来の再利用のためにエッジノードにキャッシュされます。

キャッシュまたはR2のいずれかで重みが利用可能になると、InfireはGPUへのモデルの読み込みを開始できます。

モデルのサイズは大きく異なりますが、ほとんどが大きいため、Infireの起動プロセスにおいてGPUメモリへの転送は時間がかかる可能性があります。たとえば、ほとんどの非量子化モデルは、重みをBF16浮動小数点形式で保存します。このフォーマットは、32ビットの浮動小数点数と同じ動的範囲を持ちますが、精度は低下します。サイズ、パフォーマンス、精度のスイートスポットを提供する推論に最適です。名前が示すように、BF16形式は16ビット、つまりウェイトあたり2バイトを必要とします。したがって、与えられたモデルのおよそメモリ内サイズは、そのパラメータのサイズの2倍になります。例えば、LLama3.1 8Bには約8Bのパラメーターがあり、そのメモリフットプリントは約16GBです。LLama4 Scoutのような大規模なモデルは、1090億のパラメーターを持ち、約218GBのメモリを必要とします。Infireは、Page LockedメモリとCUDA非同期コピーメカニズムの組み合わせを複数のストリームで利用し、GPUメモリへのモデル転送を高速化します。

モデルの重みをロードしている間に、Infireはモデルのパラメータに基づいて必要なカーネルの実行を開始し、デバイスにロードします。コンパイルをモデルの読み込みと並列化することで、両方のプロセスの遅延を考慮します。Llama-3-8B-Instructモデルをディスクから読み込むときのInfireの起動時間は4秒弱です。

Infireサーバーは、高性能なHTTPクレートであるhyperの上に構築されており、最低限のCPU時間を消費しながら、何百もの接続を並行して処理することが可能です。ChatGPTのユビキタス性により、vLLMやその他の多くのサービスは、OpenAIと互換性のあるエンドポイントをすぐに提供することができます。Infireもその点では変わりません。サーバーは、クライアントとの通信を処理する役割を担います。接続の受け入れ、プロンプトの処理、応答を返すことです。プロンプトは通常、いくつかのテキスト、またはチャットセッションの「トランスクリプト」と、応答の生成方法に影響する追加のパラメータで構成されます。プロンプトに付属するパラメーターには、応答のランダム性に影響する温度や、可能な応答のランダム性と長さに影響する他のパラメータが含まれます。

リクエストが有効と判断されると、Infireはそれをトークンマイザーに渡し、トークンにより、生のテキストが一連のトークン、つまりモデルが消費できる数字に変換されます。モデルによってトークンサイダーの種類は異なりますが、最も一般的なものはバイトペアエンコーディングを使用しています。トークン化には、HuggingFaceのtokenizers createを使用します。トークン化されたプロンプトとパラムは、その後バッチに送信され、GPUでの処理をスケジュールされます。そこで、エンベッディングと呼ばれる数字のベクトルとして処理されます。

Infireで最も重要な部分は、バッチ処理の方法です。複数のリクエストを並行して実行することです。これにより、メモリ帯域幅とキャッシュをよりよく利用できるようになります。

バッチ処理がなぜ重要なのかを理解するには、推論アルゴリズムがどのように機能するかを理解する必要があります。モデルの重みは、基本的に二次元行列（テンソルとも呼ばれる）の集まりです。ベクトルとして表現されるプロンプトは、主に1つの操作であるベクトル対行列の乗法で支配される一連の変換を通過します。モデルの重みが非常に大きいため、乗算のコストは、メモリからフェッチするのにかかる時間によって支配されます。さらに、現代のGPUには、行列の乗法専用のハードウェアユニット（Nvidia GPUではTensor Coreと呼ばれる）があります。メモリアクセスのコストを増幅し、テンソルコアを活用するためには、複数の操作をより大きな行列乗法に集約する必要があります。

Infireは、2つのテクニックを利用して、これらの行列操作のサイズを大きくします。1つ目はプリフィルと呼ばれるもので、この技術はプロンプトトークンに適用されます。プロンプトトークンはすべて事前に利用可能で、デコーディングを必要としないため、すべて並行して処理することができます。これが、入力トークンが出力トークンよりも安価（かつ高速）である理由の1つです。

^{Infireは、バッチ処理によりより大きな行列乗法を可能にする方法}

もう1つの手法はバッチ処理と呼ばれ、複数のプロンプトを1つのデコード操作に集約します。

Infireは両方のテクニックを組み合わせたものです。できるだけ多くのプロンプトを並行して処理し、残りのスロットを着信プロンプトからのプリフィルトークンでバッチ処理します。これは、チャンク化されたプリフィルを使った連続バッチとも呼ばれます。

トークンがInfireエンジンによってデコードされると、バッチ処理者はStreamのエンドトークンに到達したプロンプトを処理し、トークンをデコーダーに送り返し、テキストに変換します。

バッチ処理用のもう一つの仕事は、KVキャッシュの処理です。推論プロセスの中で要求の厳しい操作の1つがAttentionと呼ばれます。注意は、現在のトークンまでのすべてのトークンについて計算されたKV値を調べる必要があります。デコードする新しいトークンごとに、以前に発生したKV値を再計算しなければならないとしたら、コンテキストサイズが長くなり、プロセスの実行時間が爆発的に増えます。しかし、キャッシュを使用すると、以前の値をすべて保存し、連続するトークンごとに再度読み込むことができます。プロンプトのKVキャッシュは、コンテキストウィンドウが許可する数のトークンのKV値を保存できる可能性があります。LLama 3では、最大コンテキストウィンドウは128Kトークンです。各プロンプトにKVキャッシュを事前に割り当てていれば、H100 GPU上で4つのプロンプトを並行して実行するのに十分なメモリを確保できません。これに対するソリューションは、ページKVキャッシュです。ページKVキャッシングでは、キャッシュはページと呼ばれる小さな塊に分割されます。バッチ処理者が、プロンプトがKVキャッシュを超えることを検出すると、そのプロンプトに別のページを割り当てるだけです。ほとんどのプロンプトが最大コンテキストウィンドウにヒットすることは稀なため、この手法では通常の負荷の下で基本的に無制限の並列処理が可能になります。

最後に、バッチ処理者はGPU上で実行する必要のあるカーネルをスケジュールすることで、Infireのフォワードパスを実行します。

Infireを開発することで、当社が使用しているハードウェア（現在はNvidia Hopper GPU）に焦点を当てることができます。これにより、この特定のアーキテクチャに低レベルのPTX命令を使用して、特定のコンピュートカーネルのパフォーマンスを向上させることができました。

Infire Just-in-timeは、実行している特定のモデル用にカーネルをコンパイルし、隠れ状態のサイズ、辞書サイズ、実行しているGPUなど、モデルのパラメータを最適化します。大規模な行列乗法などの一部の操作では、Infireは、より高速とみなされると、高性能なcuBLASltライブラリを利用します。

Infireはまた、非常に粒度の細かいCUDAグラフを使用し、基本的に、可能な限りのバッチサイズに対して専用のCUDAグラフをオンデマンドで作成します。その後、将来の発表のために保存します。概念的には、CUDAグラフはジャストインタイムコンパイラーの別の形態です。CUDAドライバは、一連のカーネルのローンチを、分割されたカーネルのローンチコストが大幅に低い単一のコンテキスト（グラフ）に置き換えるため、カーネルは次に実行されますが、個別のローンチよりも、1つのグラフとしてローンチされた方が速く実行できます。

Cloudflareでは、H100 NML GPUを使用して、エッジノードの1つで合成ベンチマークを実行しました。

私たちが実行したベンチマークは、広く使用されているShareGPT v3データセットでした。4,000プロンプトのセットでベンチマークを実行し、コンカレンシーは200でした。そして、ベアメタル上で動作するInfireとvLLM、そしてgvisorの下で動作するvLLMを比較しました。これは、現在本番環境で動作している方法です。本番トラフィックのシナリオでは、エッジノードは他のトラフィックとリソースを競合することになります。これをシミュレートするために、利用可能なCPUが1つだけのgvisorで実行するvLLMをベンチマークしました。

ベンチマークからわかるように、私たちは当初の目標であるvLLMのパフォーマンスに匹敵するレベルをわずかに上回るという目標を達成しましたが、さらに重要なのは、CPU使用率が大幅に低いことで実現できたことです。これは信頼できるベアメタルプロセスとしてInfireを実行できることが主な理由です。 .推論が他のサービスから貴重なリソースを浪費することはなくなり、GPUの使用率は80%以上となり、運用コストの削減につながりました。

これは始まりにすぎません。まだInfireに実装されていないことが証明されたパフォーマンス最適化が複数あります。たとえば、Flash Attention 3を統合していますが、当社のカーネルのほとんどはカーネル融合を利用しません。これらの最適化やその他の最適化によって、近い将来、さらに高速な推論を実現できるでしょう。

AI推論の実行は、インフラストラクチャに新たな課題と要求をもたらします。Infireは、世界中のユーザーの近くでAIを効率的に実行する方法です。継続的バッチ、ページKVキャッシュ、当社のハードウェアに合わせた低レベル最適化などのテクニックを構築することで、Infireはオーバーヘッドを最小限に抑えつつ、GPUの使用率を最大化します。Infireは、特に当社が必要とする厳格なセキュリティ制約の下で、以前のvLLMベースのセットアップの何分の1かに比べて、推論タスクをより速く完了します。これにより、より少ないリソースでより多くのリクエストに対応できるようになり、Workers AI経由で提供されるリクエストの処理速度と効率性が向上します。

しかし、これは最初のイテレーションに過ぎません。より大規模なモデルのマルチGPUサポート、量子化、真のマルチテナンシーをInfireの次のバージョンに組み込むことに興奮しています。これは、Cloudflareを、開発者がAIアプリを構築するための最高のプラットフォームにするという目標の一部です。

CloudflareでAIワークロードが高速化されているかどうかを確認したいですか？Workers AI をいますぐはじめましょう