ホームへ戻る

Inference systems

FluxBin、二値基底と LUT CUDA kernel により、単一 A100 で最大 5.92 倍の推論高速化を報告

FluxBin は超低ビット重みを FP16 に再展開せず、二値パターンを使って事前計算済みの部分和を直接参照する。論文では単一 A100 で最大 5.92 倍の高速化と 10.19 倍のエネルギー効率改善を報告しているが、現時点では公開コードへのリンクにアクセスできない。

André Karwath aka Aka · CC BY-SA 2.5 · Image source
zh-Hant

香港大学とハルビン工業大学の研究チームは、二値量子化でよく見られる「モデルは大幅に小さくなるものの、実際には高速化されない」という問題の解決を目指す FluxBin を発表した。多くの 1~4-bit 手法では、実行時に依然として重みを FP16 に逆量子化する必要があるほか、非構造化された高精度の例外値によって不規則なメモリアクセスが発生する。そのため、理論上のメモリ効率やビット演算の利点が、変換処理と kernel のオーバーヘッドによって相殺されてしまう。

FluxBin はまず、複数組の ±1 の二値基底で重みを近似し、各組のスケーリング係数を行ベクトルと列ベクトルに分解することで、方向ごとに異なる重み分布を捉える。さらに Hessian 近似を用いて敏感な列を特定し、それらの列にのみ追加の二値基底を割り当てる。疎に分布する敏感な列によって GPU が低速化するのを防ぐため、Virtual Columnar Mapping は論理的に分散した列を物理的に連続する行列へ並べ替え、グローバル基底と敏感列の分岐で同種の密な GEMV kernel を再利用できるようにする。

実行レイヤーの鍵となるのは lookup table(LUT)だ。kernel は、入力の部分ベクトルで生じ得る二値の正負の組み合わせを部分和として事前計算し、圧縮された重みのビットパターンをそのままインデックスとして使用する。列方向のスケーリングは LUT の構築時に融合し、行方向のスケーリングはレジスタでの累積後に適用する。これにより、重みは Uint32 圧縮形式のままストリーミングされ、メインループ内で浮動小数点行列へ復元する必要がない。測定では CUDA Graph によって KV cache を固定し、prefill と decode の kernel をリプレイすることで、Python のスケジューリングと launch overhead による影響も抑えている。

単一の 80GB A100 上で、著者らは最大 5.92 倍の高速化、10.19 倍のエネルギー効率改善、約 4 分の 1へのメモリ使用量削減を報告し、70B 規模のモデルを収容できるとしている。LLaMA-2 7B の場合、2.75-bit の混合設定では 250.87 token/s に達し、FP16 の 42.94 token/s を上回った。一方、ゼロショット評価の平均スコアは 67.07 から 59.15 に低下しており、スループット向上を精度劣化のない代替手段と解釈することはできない。また、すべての結果は A100、カスタム Transformers パス、静的な CUDA Graph に集中している。論文に記載された GitHub repository は確認時点で 404 を返しており、エンジニアは依然として kernel、対応モデル、エンドツーエンドのメモリ使用量を独立に検証できない。

出典

  1. FluxBin: Flexible LUT-based Ultra-low-bit LLM Inference by Algorithm-Kernel Synergy
  2. CUDA Graphs — CUDA Programming Guide
  3. Transformers text-generation API documentation