返回首頁

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,嘗試解決二元量化常見的「模型壓得很小,實際卻沒有變快」問題。許多一至四位元方法在執行時仍須把權重反量化成 FP16,或因非結構化高精度例外產生不規則存取;因此理論上的記憶體與位元運算優勢,會被轉碼和 kernel 開銷吃掉。

FluxBin 先以多組正負一二元基底近似權重,並把每組縮放因子拆成行向量與列向量,以捕捉不同方向的權重分布。它再利用 Hessian 近似找出敏感欄位,只為這些欄位配置額外二元基底。為避免稀疏敏感欄位拖慢 GPU,Virtual Columnar Mapping 會把邏輯上分散的欄位重排成實體連續矩陣,讓全域基底與敏感分支重用同類密集 GEMV kernel。

執行層的關鍵是 lookup table。kernel 先把輸入子向量可能遇到的二元正負組合預算成部分和,壓縮權重的位元圖樣直接充當索引;列縮放在建立 LUT 時融合,行縮放則於暫存器累加後套用。權重因此以 Uint32 壓縮格式串流,不必在主迴圈還原成浮點矩陣。測試亦以 CUDA Graph 固定 KV cache 並重播 prefill、decode kernel,降低 Python 排程與 launch overhead 對測量的干擾。

在單張 80GB A100 上,作者報告最高 5.92 倍速度、10.19 倍能源節省及約四倍記憶體縮減,並稱可容納 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