2026-04-17 Mari Galicer Ivan Nikulin Chris Branch 12分読了
はじめに
世界のインターネット接続人口の95%の50ms以内で推論を実行するには、GPUメモリを徹底的に効率化する必要があります。昨年、Rust ベースの推論エンジンである Infire でメモリ利用率を改善し、モデルスケジューリングプラットフォームである Omni でコールドスタートを排除しました。現在、推論プラットフォームの次の大きなボトルネックであるモデルの重みに取り組んでいます。
LLMから単一のトークンを生成するには、GPU メモリからすべてのモデルの重みを読み込む必要があります。多くのデータセンターで使用している NVIDIA H100 GPU では、テンソルコアはメモリが供給できるよりも約600倍高速でデータを処理できるため、ボトルネックは計算ではなくメモリ帯域幅にあります。メモリバスを通過するすべてのバイトは、重みがより小さければ回避できたはずのバイトです。
Unweight の紹介
この問題を解決するために、Unweight を構築しました。これは、特別なハードウェアに依存せずに、ビット単位で正確な出力を保持しながら、モデルの重みを最大15～22%小さくできるロスレス圧縮システムです。
ここでの核となるブレークスルーは、高速なオンチップメモリで重みを解凍し、それらをテンソルコアに直接供給することで、遅いメインメモリへの余分なラウンドトリップを回避することです。ワークロードに応じて、Unweight のランタイムは複数の実行戦略から選択します。一部は単純性を優先し、他は メモリトラフィックを最小化します。オートチューナーが重み行列とバッチサイズごとに最適なものを選択します。
この投稿では Unweight の仕組みについて詳しく説明しますが、透明性の向上とこの急速に発展する分野でのイノベーション促進の精神に基づいて、技術論文も公開し、GPU カーネルもオープンソース化しています。
初期結果
Llama-3.1-8B での初期結果は、多層パーセプトロン (MLP) の重みだけで約30%の圧縮を示しています。Unweight はデコード用のパラメータに選択的に機能するため、これはモデルサイズの15～22%削減と約3 GB の VRAM 節約につながります。
下のグラフィックに示すように、これにより GPU からより多くを絞り出すことができ、Cloudflare のネットワークでより多くの場所でより多くのモデルを実行できます。
Unweight のおかげで、単一の GPU により多くのモデルを適合させることができます
圧縮がそれほど簡単ではない理由
推論を高速化したり、より小さな GPU で実行したりするために、モデルの重みを創造的な方法で圧縮する方法を探索する研究が増えています。最も一般的なのは量子化です。これは、大きな32ビットまたは16ビットの浮動小数点数をより小さな8ビットまたは4ビットの整数に変換することで、モデルの重みとアクティベーションのサイズを削減する手法です。
これはロッシー圧縮の一形態です。異なる16ビット浮動小数点値が同じ4ビット整数に変換される可能性があります。この精度の低下は、予測不可能な方法で応答の品質に影響します。
多様なユースケースに対応する本番推論サービスでは、正確なモデルの動作を保持するロスレスなものが必要であることがわかっていました。
既存のアプローチとの比較
最近のいくつかのシステム (Huff-LLM、ZipNN、ZipServ) は、LLM の重みを大幅に圧縮できることを示していますが、これらのアプローチは私たちのものとは異なる問題を対象としています。
- ZipNN: 配布と保存のために重みを圧縮し、CPU で解凍が行われます
- Huff-LLM: デコード用のカスタム FPGA ハードウェアを提案します
- ZipServ: GPU 推論で解凍を融合させますが、コンシューマーグレード GPU を対象としており、H100 GPU では機能しません
これらのどれも、Hopper GPU でのロスレス推論時間解凍で、Rust ベースの推論エンジンと統合できるものを提供していませんでした。
核となる課題
核となる課題は、バニラ圧縮ではありません。BF16 の重みの指数バイトは非常に冗長であるため、エントロピーコーディングはそれらに対してうまく機能します。課題は、推論を遅くしないほど十分に高速に解凍することです。
H100 では、テンソルコアはほとんどの時間メモリを待機しています。しかし、このアイドル容量を単純に解凍に転用することはできません。各 GPU コンピュートユニットは、共有メモリの制約により、解凍カーネルまたは行列乗算カーネルのいずれかを実行できますが、両方を同時に実行することはできません。
行列乗算と完全にオーバーラップしていない解凍レイテンシは、トークンレイテンシに直接加算されます。
Unweight の答えは、高速なオンチップ共有メモリで重みを解凍し、結果をテンソルコアに直接供給することです。ただし、異なるバッチサイズと重み形状全体でそれを効率的に機能させることが、実際のエンジニアリングが存在する場所です。
モデルの重みを効果的に圧縮する方法
AI モデルのすべての数値は、16ビットの「ブレインフロート」(BF16) として保存されます。各 BF16 値には3つの部分があります。
- 符号 (1ビット): 正または負
- 指数 (8ビット): 大きさ
- 仮数 (7ビット): その大きさ内の正確な値
指数の冗長性
符号と仮数は重み全体で予測不可能に変動します。ランダムデータのように見え、意味のある方法で圧縮することはできません。しかし、指数は異なるストーリーを語ります。
指数は驚くほど予測可能です
先行研究により、訓練済み LLM 全体で、256 の可能な指数値のうち、ほんの一握りが支配的であることが確立されています。最も一般的な上位16の指数は、典型的なレイヤーのすべての重みの99%以上をカバーしています。情報理論によれば、この分布を表すのに必要なのは約2.6ビットだけです。割り当てられた8ビットよりもはるかに少ないです。
典型的な LLM レイヤーの指数値分布を見ると、上位16の指数がすべてのモデルの重みの99%を占めていることがわかります。
典型的な LLM レイヤーの指数値分布
これが Unweight が利用する冗長性です。符号と仮数はそのままにして、Huffman コーディングを使用して指数バイトのみを圧縮します。これは、一般的な値に短いコードを割り当て、稀な値に長いコードを割り当てる古典的な手法です。
指数分布が非常に歪んでいるため、これは指数ストリームで約30%の圧縮を達成します。
選択的な適用
これを MLP 重み行列 (ゲート、アップ、ダウン投影) に選択的に適用します。これらはモデルのパラメータの約3分の2を構成し、トークン生成中のメモリトラフィックを支配します。注意の重み、埋め込み、レイヤー正規化は圧縮されません。
全体として、最適化は多層パーセプトロン (MLP) の重みサイズの約20%削減に変換されます。詳細は技術レポートで完全に説明されています。
エッジケースの処理
稀な指数を持つ少数の重みは別々に処理されます。64 の行内の任意の重みが上位16パレットの外の指数を持つ場合、行全体は逐語的に保存されます。このアプローチは、ホットパスでの要素ごとの分岐を排除します。すべての単一の重みをエッジケースについてチェックする代わりに、事前に行ごとに1つの決定を下します。
GPU メモリボトルネック
NVIDIA H100 GPU には、2つの関連するメモリの種類があります。
- 高帯域幅メモリ (HBM): 大きいですが、アクセスが比較的遅い。ここにモデルの重みが存在します
- 共有メモリ (SMEM): 小さいですが、非常に高速。ここで GPU は数学を行う直前にデータをステージングします
推論中、各トークンを生成するには、HBM から完全な重み行列を読み込む必要があります。HBM と SMEM 間のメモリバスがパフォーマンスボトルネックです。数学自体ではありません。
バスを通過するバイト数が少ないほど、トークン生成が高速になります。
H100 のテンソルコアは、HBM がデータを供給できるよりもはるかに高速に数値をクランチできます。圧縮は、より少ないバイトがバスを通過する必要があるため、役立ちます。
しかし、落とし穴があります。GPU は圧縮されたデータで数学を行うことはできません。重みは最初に解凍する必要があります。
従来のアプローチの問題
先行研究のほとんどは、重み行列全体を HBM に解凍してから、標準的な行列乗算を実行します。これはストレージ容量に役立ちますが、すべてのトークンに対して HBM から完全な非圧縮行列を読み込むため、帯域幅には役立ちません。
圧縮された重みを使用する4つの方法
推論中に圧縮された重みを使用する単一の最良の方法はありません。正しいアプローチはワークロードに依存します。バッチサイズ、重み行列の形状、解凍に利用可能な GPU 時間です。
Unweight は4つの圧縮実行パイプラインを提供し、それぞれ解凍の努力と計算の複雑さの間で異なるバランスを持っています。
- 完全な Huffman デコード
- 指数のみのデコード
- パレットトランスコード
- 前処理を完全にスキップ
4つの異なる実行パイプライン
パイプラインのスペクトラム
4つのパイプラインはスペクトラムを形成します。
一方の端: 完全デコード
完全デコードは、元の BF16 重みを完全に再構成し、標準的な行列乗算のために NVIDIA の cuBLAS ライブラリに渡します。これは cuBLAS が通常のデータで全速で実行される最も単純なパスですが、前処理ステップはメインメモリに最も多くのバイトを書き込みます。行列乗算が小さく、カスタムカーネルのオーバーヘッドが支配的な小さなバッチサイズで機能します。
もう一方の端: ダイレクトパレット
ダイレクトパレットは前処理を完全にスキップします。重みはモデルロード時にコンパクトな4ビット形式に事前トランスコードされ、行列乗算カーネルはこれらのインデックスからオンザフライで BF16 値を再構成します。前処理コストはゼロですが、カーネルは要素ごとにより多くの作業を行います。
中間: 2つの独立したパス
1つは指数バイトのみをデコードします (前処理トラフィックを半減)。もう1つは実行時に4ビットパレットインデックスにトランスコードします (4分の1に削減)。どちらも再構成的な行列乗算を使用します。カスタムカーネルは圧縮データを読み込み、高速共有メモリで BF16 を再構成し、メインメモリへのラウンドトリップなしでテンソルコアに直接供給します。
単一のパイプラインが勝たない理由
前処理が少ないほど、HBM に書き込まれるデータが少なくなり、メモリバスがより早く解放されます。しかし、matmul カーネルにより多くの再構成作業をシフトさせます。そのトレードオフが支払われるかどうかは、状況によって異なります。
小さなバッチサイズ (1～64トークン)
小さなバッチサイズでは、matmul は小さいため、オーバーラップする計算がそれほど多くなく、カスタムカーネルの固定コストが支配的です。完全デコード + cuBLAS は、cuBLAS のオーバーヘッドが低いため、単純に勝つことがよくあります。
大きなバッチサイズ (256+トークン)
大きなバッチサイズでは、matmul は十分に長く実行され、余分な再構成作業を吸収できます。軽い前処理がより早く終了し、解放されたバス帯域幅と計算オーバーラップが支払われます。パレットまたは指数パイプラインが先に進みます。
同じレイヤー内の異なる重み行列
同じレイヤー内の異なる重み行列は、異なるパイプラインを優先できます。「ゲート」と「アップ」投影は「ダウン」投影とは異なる次元を持ち、matmul 内で実行される操作の順序を変更し、異なるパフォーマンストレードオフが必要です。
スループット対パイプライン戦略
これが Unweight が単一の戦略をハードコードしない理由です。ランタイムは、オートチューニングプロセスによって通知される各バッチサイズで各重み行列に最適なパイプラインを選択します。このプロセスは、ターゲットハードウェアで実際のエンドツーエンドスループットを測定します。
再構成的な matmul の仕組み
4つのパイプラインのうち3つは、解凍と計算を融合させるカスタム行列乗算カーネルを使用します。このカーネルは HBM から圧縮データを読み込み、共有メモリで元の BF16 値を再構成し、テンソルコアに直接供給します。すべて1つの操作で。
再構成された重みはメインメモリに存在することはありません。
従来の解凍対 Unweight
Unweight を使用すると、MLP 重み行列のメモリバスを通過するバイト数は約30%少なくなります。
カーネル内の役割分担
このカーネル内では、GPU のスレッドグループは2つの役割に分割されます。
プロデューサーグループ
プロデューサーグループは、専用のメモリコピーハードウェア (TMA) を使用して、HBM から圧縮入力を共有メモリに読み込みます。符号+仮数バイト、指数データ (またはパレットインデックス)、および稀な指数を持つ行の場合は逐語的な指数行をステージングします。コンシューマーより先に実行され、循環バッファを満たすため、データが必要になる前に準備ができています。
コンシューマーグループ
コンシューマーグループは、指数を符号+仮数バイトと組み合わせることで BF16 値を再構成し、すぐに Hopper の WGMMA テンソルコア命令に結果を供給します。再構成された重みは、共有メモリを離れることなく、アセンブリから計算に直接進みます。