shiichan

Kimi と GLM のコンテキストがまさかの 2 倍に!Cloudflare の軽量化テクニックがすごすぎる

みんな、しぃちゃんだよ!今日は Cloudflare のエンジニアリングブログから、AI モデルをどうやって「軽く・速く・安全に」動かすかというガチ技術寄りの話を見つけてきたよ!

Cloudflare Blog blog.cloudflare.com

なにが発表されたの?

Cloudflare の公式ブログで、Moonshot AI の Kimi や Z.ai の GLM のような大規模なオープンウェイトモデルを自社インフラで大量にサーブするための最適化手法が詳しく紹介されたよ。これらのモデルは巨大な MoE(Mixture of Experts)構成で、しかも長いコンテキストに対応しているから、ボトルネックになるのは計算力じゃなくて GPU のメモリなんだって。

Cloudflare が取った対策は大きく 3 つ。

  • KV キャッシュの量子化(16 bit から 8 bit に圧縮してメモリを節約)
  • モデル重みの圧縮(8 bit から 4 bit にしてチェックポイントを小さく)
  • KV キャッシュの整合性チェック(大規模並列実行でも壊れたデータを返さないようにする安全対策)

今までどうだったの?

今まで KV キャッシュは BF16(16 bit)で保持するのが基本で、たとえば Kimi K2.6 だとこの方式でメモリに保持できるコンテキスト量はおよそ 686,000 トークンが上限だったの。長いコンテキストを扱うモデルだと、モデルの重みよりも先に KV キャッシュが GPU メモリを埋め尽くしちゃうのが課題だったんだって。

そこで Cloudflare は KV キャッシュを FP8(e4m3 形式の 8 bit 浮動小数点)に変換。これだけでサイズがほぼ半分になって、保持できるコンテキストがおよそ 1.37 百万トークンまで、約 2 倍に伸びたよ。しかも面白いのは、メモリが厳しくなるのは「デコード」の場面だけだから、計算が重い「プリフィル」は精度重視でこれまで通り BF16 のままにしているところ。使い分けがうまいの!

単発リクエストのスループットは BF16 の 137 トークン/秒に対して FP8 は 125 トークン/秒とわずかに下がるものの、同時実行数を増やすと真価を発揮するよ。64 同時リクエストで比べると、BF16 はすでに 32 リクエストの時点でメモリ不足になっちゃうのに対して、FP8 は 2,192 トークン/秒を叩き出したんだって。全体で見るとスループットは 41% 向上、トークンあたりのコストは約 30% 下がる計算になるみたい。

精度への影響も気になるところだけど、GSM8K(94.24% から 94.09%)、ARC-Challenge(66.72% から 67.49%)、MMLU(89.11% から 89.04%)と、ベンチマークではほとんど差が出ていないよ。

これで何が変わるの?

この最適化によって、Cloudflare のインフラ上で Kimi や GLM のような巨大モデルを動かすときに、同じ GPU でより多くの同時リクエストをさばけるようになるよ。長文コンテキストの推論を使う開発者にとっては、レイテンシが安定しやすくなったり、コストパフォーマンスが上がったりするという形で恩恵が返ってくるはず。

さらに後述する整合性チェックのおかげで、大規模に並列処理していても「壊れたデータをそのまま返してしまう」というリスクを抑えられるから、安全性の面でも一歩前進だね。

深く潜ってみよう

GLM 5.2 については KV キャッシュとは別に、モデル重みそのものの圧縮にも挑戦しているよ。もともと FP8(8 bit)で保存されていた重みを INT4(4 bit の整数)に変換したところ、こんな結果になったんだって。

  • チェックポイントのサイズ: 705 GB から 421 GB に(約 40% 削減)
  • 8-way テンソル並列時の GPU 1 枚あたりのメモリ使用量: 88 GB から 52 GB に
  • 浮いたメモリ分、同じハードウェアで約 1.18 百万トークン分の KV キャッシュを追加確保できるように

デコードのスループットも改善していて、単発リクエストなら 60 トークン/秒 から 92 トークン/秒 と 55% アップ。ただし 64 同時リクエストまで負荷を上げると差は 16%(1,672 から 1,933 トークン/秒)まで縮まるんだって。

一方でプリフィル(最初にプロンプト全体を処理する工程)は逆に INT4 だと遅くなる弱点もあるみたい。INT4 の重みは計算の前に一度展開(デコンプレス)する必要があるから、FP8 の 10,160 トークン/秒に対して INT4 は 8,660 トークン/秒に留まるの。そこで Cloudflare は「プリフィルは FP8、デコードは INT4」という使い分け(disaggregated アーキテクチャ)で両方のいいとこ取りをしているよ。精度面でも、全ベンチマークを通じて FP8 との差は 0.8 ポイント以内(MMLU は FP8 が 86.60%、INT4 が 86.54%)に収まっているんだって。

もうひとつの柱が KV キャッシュの整合性チェック。数百件のリクエストが物理的なキャッシュページを共有する規模になると、めったに起きないはずのバグでも高頻度で顕在化しちゃうから、Cloudflare は各物理ページに「再割り当てされるたびに変わるタグ」を付けて、サーバー側でリクエストごとに「使うはずのページとタグ」を記録する仕組みを導入したよ。タグが一致しなければ、壊れたデータをそのまま返さずにリクエストを中断する設計なの。

この整合性チェックのオーバーヘッドは、同時実行数 1 から 8 の範囲でスループット・p95 レイテンシともに 1% 未満に抑えられているとのこと。しかも Attention のカーネルに直接組み込まず、あえて別のバッチ処理として実行することで、GPU のスレッド間で起きがちな同期のレースコンディションを避けているんだって。

ちなみにこれらのベンチマークはオープンソースの推論エンジン SGLang を使って計測されていて、Cloudflare はそのパッチを本家にもアップストリームしているそう。今後の展望としては、こんな方向性が挙げられているよ。

  • FP8 の KV キャッシュをフリート全体に展開する
  • Blackwell 世代の GPU 向けに NVFP4 形式の重みを検証する
  • 整合性チェックのオーバーヘッドをさらに減らす

まとめ

  • Kimi や GLM のような巨大な MoE モデルは GPU メモリがボトルネックになりやすい
  • KV キャッシュを BF16 から FP8(e4m3)にすることで保持できるコンテキストが約 2 倍(686,000 から 1.37 百万トークン)に伸びた
  • GLM 5.2 の重みを FP8 から INT4 に圧縮してチェックポイントを 40% 削減、精度低下は 0.8 ポイント以内に抑えた
  • プリフィルは FP8・デコードは INT4 という使い分けで速度と精度を両立させている
  • KV キャッシュに整合性チェックを追加し、壊れたデータを返す前にリクエストを中断できるようにした

自分で LLM 推論基盤を運用していたり、GPU メモリのやりくりに頭を悩ませているインフラ・MLOps エンジニアには、たまらなく刺さる内容だよ!