より小さく、より速く、より安全に：KimiとGLMを大規模で実行する

Workers AIは、世界中の優れたオープンソースモデルの推論を、ユーザーに近いCloudflareのデータセンター内のGPUで実行しています。最も高性能で要求の高いモデルの2つは、MoonshotのKimi KシリーズとZ.aiのGLMです。これらは大規模で、長文コンテキストに対応し、Mixture-of-Expertsモデルであり、非常に使いやすいものです。しかし、メモリの制約から効率的に提供するのは非常に困難です。 以前、Workers AIで大規模モデルをどのように提供するか、また推論のプリフィル（prefill）とデコード（decode）のフェーズを分離して各GPUを最大限に活用する方法について記述しました。この記事では、これらのモデルをメモリに収め、高速に保つために、その上に重ねる3つの技術に焦点を当てます。それは、KVキャッシュの量子化、モデルウェイトの圧縮、そして、これら2つの技術が共有ハードウェアにより多くのリクエストを詰め込むことを可能にするため、それらのリクエストが共有するキャッシュの保護です。これらの最適化により、モデルの精度を変更することなく、より低いコストでより多くの顧客をサポートできるようになります。 すべての実験と本番トラフィックは、オープンソースの推論サービングフレームワークであるSGLangで実行およびベンチマークされています。SGLangは市場で最高のパフォーマンスを提供すると判断し、SGLangチームと緊密に連携してパッチや新機能をアップストリームし、私たちの作業をオープンソースコミュニティで利用できるようにしています。 KVキャッシュの量子化 モデルがテキストを生成する際、すでに処理したすべてのトークンのアテンションキー（K）とバリュー（V）をKVキャッシュと呼ばれる構造に保存します。このキャッシュにより、モデルは新しいトークンごとにコンテキスト全体を再読み込みすることなく、長い会話を続けることができます。長文コンテキストモデルの場合、KVキャッシュは急速に増加し、通常はモデルのウェイトではなく、KVキャッシュが最初にGPUメモリを使い果たします。 デフォルトでは、キャッシュは16ビット精度（BF16）で保存されます。代わりに8ビット浮動小数点（FP8、e4m3）で保存することで、サイズを半分にします。Kimi K2.6では、メモリに保持できるコンテキスト量が約68万6千トークンから約137万トークンに増加し、2倍になります。 その利点がどこから来るのかを正確に述べる価値があります。それは生の速度ではありません。キャッシュの量子化は、FP8アテンションカーネルが値を読み取る際に変換する必要があるため、トークンごとにわずかな作業が増加します。しかし、それが変更するのは、一度に保持できるリクエストの数です。以下の測定値は、非集約型H200デプロイメントでのKimi K2.6デコードのもので、アテンションカーネルを直接比較しています。 | 同時リクエスト数 | BF16 KVキャッシュ (トークン/秒) | FP8 KVキャッシュ (トークン/秒) | |---|---|---| | 1 | 137 | 125 | | 8 | 731 | 689 | | 16 | 1,061 | 1,028 | | 32 | 1,558 | 1,489 | | 64 | 2,192 | 1,616 | | 64 (BF16) | アウトオブメモリ | 1,028 | 任意の同時リクエストレベルで、BF16はトークンあたり数パーセント高速です。しかし、BF16は32件の同時リクエストでキャッシュを使い果たし、33件目を受け入れることができません。一方、FP8は64件まで継続でき、毎秒2,192トークンに達し、BF16のピークよりも約41%高く、トークンあたりのコストは約30%低くなります。プリフィルとデコードを別々のプールとして実行するため、最も役立つ場所に適用できます。プリフィルは計算バウンドでありメモリバウンドではないため、キャッシュはBF16のままにして、わずかに高いスループットを維持します。 モデルの回答が変わらなければ、これらは問題になりません。そのため、確認しました。評価スイート全体で、FP8とBF16キャッシュは区別がつきません。 | ベンチマーク | BF16 KV | FP8 KV | |---|---|---| | GSM8K | 94.24 | 94.09 | | ARC-Easy | 89.06 | 89.14 | | ARC-Challenge | 66.72 | 67.49 | | MMLU | 89.11 | 89.04 | | MMLU-Pro | 80.29 | 79.29 | | mcxams (内部ベンチマーク) | 61 / 63 | 61 / 63 | | Tool-call validity | 92.2% | 92.6% | モデルウェイトの圧縮 KVキャッシュはGPUメモリへの要求の1つですが、モデルのウェイトももう1つです。GLM 5.2では、ウェイトを8ビット浮動小数点から4ビット整数（INT4）に圧縮し、精度は失われません。チェックポイントは705 GBから421 GBに約40%縮小され、8ウェイテンソル並列デプロイメント全体でのGPUあたりのメモリ使用量は約88 GBから52 GBに低下します。これにより、同じハードウェアで約118万トークンのKVキャッシュのためのスペースが生まれます。 評価スイート全体で、INT4とFP8ウェイトは区別がつきません。 | ベンチマーク / 機能 | メトリック | FP8 | INT4 | |---|---|---|---| | GSM8K | Exact match | 94.39% | 93.56% | | GSM8K | Flexible | 94.24% | 93.48% | | ARC-Easy | Accuracy | 86.62% | 86.15% | | ARC-Easy | Acc (norm) | 84.51% | 85.19% | | ARC-Challenge | Accuracy | 64.93% | 64.85% | | ARC-Challenge | Acc (norm) | 67.24% | 66.64% | | MMLU | Average | 86.60% | 86.54% | | MMLU-Pro | Exact | 80.80% | 80.47% | | mcxams (内部ベンチマーク) | Passed | 62 / 63 | 62 / 63 | ウェイトが小さいとデコードフェーズが速くなります。その理由は明らかです。各トークンを生成するということは、モデルのウェイトをGPUメモリからストリーミングすることを意味するため、デコード速度はメモリ帯域幅によって制限されます。より少ないデータを移動すれば、各トークンはより早く到着します。 | 同時リクエスト数 | GLM FP8 (トークン/秒) | GLM INT4 (トークン/秒) | INT4 ゲイン | |---|---|---|---| | 1 | 609 | 92 | +55% | | 8 | 425 | 551 | +21% | | 16 | 683 | 825 | +21% | | 32 | 994 | 1,267 | +27% | | 64 | 1,672 | 1,933 | +16% | プリフィルは異なる動作をします。これは計算バウンドであり、INT4ウェイトはモデルがそれらと乗算する前に展開し直す必要があるため、その追加ステップによりプリフィルは速くなるのではなく遅くなります。GLMは、FP8では毎秒約10,160トークン、INT4では8,660トークンのプリフィルを維持します。KVキャッシュと同様に、非集約型設計はこれを妥協ではなく選択に変えます。デコードではINT4を使用し、そこで勝利します。プリフィルではFP8を使用し、そこで勝利します。モデルの精度は、実行するすべてのベンチマークでFP8モデルから0.8ポイント以内に収まり、品質は区別がつきません。 共有KVキャッシュの保護 上記の2つの技術は同じ効果をもたらします。それは、より多くのリクエストが1つのGPUのメモリを同時に共有できるようにすることです。その効率性が目的ですが、数百のリクエストが同じ物理KVキャッシュのページを読み書きしていることを意味します。これを高速にするメカニズム（ページングアテンション、連続バッチ処理、キャッシュ再利用）はすべて、ブックキーピングを正確に行うことに依存しています。リクエスト量が多い場合、10億分の1の間違いでも定期的に発生します。 そこで、KVキャッシュ整合性チェックを防御層として構築しました。アイデアは単純です。各物理キャッシュページには、ページが再割り当てされるたびに変化するタグが付けられ、サーバーは各リクエストが使用を期待するページとタグを記録します。サポートされているデコード操作がキャッシュから読み取る前に、それらのマッピングがチェックされます。一致しない場合は、影響を受けたリクエストは、間違ったページからのデータを返すことを許可されるのではなく、中止されます。 安全チェックが出荷されるかどうかを決定する質問は、そのコストです。8,192トークンの入力と1,000トークンの出力を持つ、2つのプリフィル、2つのデコード構成の中規模本番モデルで測定しました。 | 同時実行数 | スループット変化 | p95 レイテンシ変化 | |---|---|---| | 1 | -0.53% | +0.42% | | 2 | -0.38% | +0.54% | | 4 | -0.79% | +0.63% | | 8 | -0.43% | +0.80% | コストはスループットとテールレイテンシの両方で1%未満であり、95%信頼区間の上限でさえ1%近くに留まります。検証をアテンションカーネルに統合するのではなく、別のバッチチェックとして実行することで計算コストを低く抑えました。これは、GPUスレッドグループ間の競合を導入する可能性があります。これはデプロイメントごとに有効化され、デフォルトパスは測定可能なオーバーヘッドのないno-opトラッカーを使用するため、必要としないデプロイメントは何も負担しません。 次は何ですか 最先端モデルを効率的に提供することは、常に変化する目標であり、これはその背後にある継続的な作業です。FP8 KVキャッシュをフリート全体に拡大し、Blackwell（NVIDIAのGPUアーキテクチャ）でNVFP4ウェイトを検証し、整合性チェックを negligible cost で常に有効にできるように取り組んでいます。これらの最適化により、より低いコストで、同じ精度で、より多くのお客様をサポートし続けることができます。 最高のオープンモデルをGPUに搭載し、数百万人の開発者に提供するという問題が、あなたの興味を引くものであれば、ぜひ私たちと一緒に働きましょう。