本篇工作已被電子設計自動化領域頂級會議 DAC 2025 接收，由上海交大計算機學院蔣力教授與劉方鑫助理教授帶領的 IMPACT 課題組完成，同時也獲得了華為 2012 實驗室和上海期智研究院的支持。第一作者是博士生汪宗武與碩士生許鵬。

在通用人工智能的黎明時刻，大語言模型被越來越多地應用到復雜任務中，雖然展現(xiàn)出了巨大的潛力和價值，但對計算和存儲資源也提出了前所未有的挑戰(zhàn)。在以 transformer 模型為基礎的大模型中，鍵值緩存雖然用以存代算的思想顯著加速了推理速度，但在長上下文場景中成為了存儲瓶頸。例如，半精度的 LLaMA-2-7B 模型權重約 14GB，在上下文長度為 128K 時鍵值緩存占據(jù) 64GB，總和已經(jīng)接近高端卡 NVIDIA A100 的 80GB 顯存容量上限。鍵值量化可被用于壓縮緩存，但往往受到異常值的干擾，導致模型性能的顯著下降。為此，本文的研究者提出了 MILLION，一種基于乘積量化的鍵值緩存壓縮和推理加速設計。

• arxiv 鏈接：https://arxiv.org/abs/2504.03661
• 開源鏈接：https://github.com/ZongwuWang/MILLION

整型量化的軟肋：異常值

圖 1：矩陣量化可視化。紅色代表的異常值顯著大于其他值，導致均勻量化后高位編碼被浪費。

量化中受到廣泛使用的整型均勻量化受到異常值的影響較為顯著。圖 1 展示了矩陣中的量化。在一組分布較為集中的數(shù)據(jù)中，一個顯著偏離其他值的異常值會導致其他值的量化結果全部落在較低區(qū)間，浪費了高位編碼的表示能力。

圖 2：圖中使用 「通道熵」 定量比較不同方案的量化效果，越大表明越有效地利用了通道容量，即整型的寬度。沿通道量化只能解決沿該方向分布的異常值，而在面對另一方向異常值時效果不佳。

在實際的鍵值量化中，為了更好的表示能力，通常對于每個通道（即鍵值向量的維度）或每個 token 采取不同的量化參數(shù)，這種方法被稱為沿通道量化（channel-wise quantization）或沿詞元量化（token-wise quantization）。然而，如圖 2 所示，沿特定方向量化只能解決沿該方向分布的異常值。

圖 3：實際采樣獲得的鍵值緩存分布。在 llama-2-7b-wikitext-103-v1-layer10-value 中，異常值并不遵循簡單的沿通道分布，而是呈現(xiàn)為較復雜的點狀和團狀。

研究團隊通過實際采樣數(shù)據(jù)發(fā)現(xiàn)，在鍵值緩存中，沿通道方向分布的異常值占多數(shù)，但也存在并不明顯的情況，如圖 3 所示。這表明，上述量化方案并不是一勞永逸的解決方式，仍然存在優(yōu)化空間。

異常值的解決方案：乘積量化

圖 4：數(shù)軸上的均勻和非均勻量化對比。在對 8 個數(shù)據(jù)點進行 2 比特量化過程中，均勻量化浪費了 10 編碼。而基于聚類的非均勻量化則編碼更合理。

如圖 4 所示，非均勻量化通過聚類的方式允許量化區(qū)間不等長，從而更合理地分配編碼，提升量化效率。研究團隊觀察到，由于通道間的數(shù)據(jù)分布可能存在關聯(lián)（即互信息非負），將通道融合后在向量空間中聚類，效果一定不亞于獨立通道的量化，如圖 5 所示。

圖 5：左圖為兩個通道獨立進行 1 比特量化，右圖為在通道融合后進行 4 分類的 KMeans 聚類。融合通道量化的通道熵更加接近 2 比特的容量極限，展示出更好的量化效果。

由于高維空間中聚類較為困難，因此將整個向量空間劃分為多個低維子空間的笛卡爾積，可以平衡聚類復雜度和量化效果。這與最近鄰搜索中使用的乘積量化思想一致。研究團隊通過實驗發(fā)現(xiàn)，子空間維度為 2 或 4 是較好的平衡點。

推理加速手段：高效的系統(tǒng)和算子實現(xiàn)

圖 6：三階段的推理系統(tǒng)設計

圖 7：分塊注意力機制使得批量延遲量化成為可能

圖 6 展示了離線訓練、在線預填充、在線解碼三階段的量化推理系統(tǒng)設計。其中，碼本訓練（量化校準）屬于秒級輕量化任務，并且離線進行，不影響運行時開銷；在線預填充階段使用訓練好的碼本對鍵值緩存進行量化壓縮，達到節(jié)省顯存的目的；在線解碼階段采用分塊注意力機制的方法，將預填充階段的歷史注意力和生成 token 的自注意力分開計算（如圖 7 所示），達成批量延遲量化的目的，掩藏了在線量化的開銷，確保模型輸出的高速性。并且，在歷史注意力階段，由于歷史鍵值對數(shù)遠大于碼本長度，因此先用查詢向量與碼本計算好非對稱距離查找表（ad-LUT），可以大大減少內積距離計算量，達到加速計算的目的。

圖 8：向量化加載可有效使帶寬飽和

在算子優(yōu)化方面，研究團隊在 flash decoding 的基礎上使用了寬數(shù)據(jù)（如 float4）向量化加載的方式，將多個乘積量化編碼打包為寬數(shù)據(jù)，有效使帶寬飽和（如圖 8 所示）。同時，在表查找階段，子空間之間的表具有獨立性，并且可以被放入少量緩存行中，研究團隊利用這一空間局部性極大提高了表查找的 L2 緩存命中率。此外，研究團隊還仔細掃描了不同上下文長度下可能的內核參數(shù)，找到最優(yōu)配置，形成了細粒度的預設，在實際運行時動態(tài)調整，充分利用 GPU 的計算資源。具體實現(xiàn)可以在開源倉庫中找到。

實驗結果

實驗設置

圖 9：實驗設置

實驗采用了不同位置編碼、不同上下文長度的多種模型進行了詳細的評估。在模型性能方面，采用困惑度（Perplexity，PPL）和 Longbench 兩種指標；在系統(tǒng)性能方面，采用每詞元輸出間隔（Time Per Output Token, TPOT）定量分析，并給出了注意力層詳細的剖析。對比采用方案和乘積量化參數(shù)如圖 9 所示。

模型性能

圖 10：困惑度指標。其中 「-1%」 表示該方法額外存儲 1% 的異常值不參與量化。

困惑度越小表明模型輸出質量越高。實驗結果表明，MILLION 與額外處理了異常值的 SOTA 方案輸出質量保持一致，展現(xiàn)出對異常值良好的魯棒性。而 SOTA 方案在不處理異常值的情況下可能會遭遇嚴重的輸出質量損失。

圖 11：Longbench 問答數(shù)據(jù)集得分展示

在長對話問答任務中，不同模型在各種數(shù)據(jù)集上的得分均表明，MILLION 方案能夠在 4 倍鍵值緩存壓縮效率下保持幾乎無損的表現(xiàn)。

系統(tǒng)性能

圖 12：每詞元輸出時間。對比其他方案，MILLION 的優(yōu)勢持續(xù)增長，在 32K 上下文時達到 2 倍加速比。

圖 13：注意力層時間剖析

在 TPOT 評估中，MILLION 能夠在 32K 上下文語境下同時達成 4 倍鍵值緩存壓縮比和 2 倍端到端加速比。注意力層的深入分析表明，MILLION 在訪存和內核函數(shù)方面對比 baseline 取得顯著優(yōu)勢。

總結

MILLION 的主要貢獻在于：（1）深入分析鍵值緩存分布；（2）提出基于乘積量化的非均勻量化算法；（3）設計高效的推理系統(tǒng)及內核。研究團隊首先證實了鍵值緩存中異常值存在的普遍性，并指出異常值的不同分布是當前主流的量化方案精度不足的根本原因；然后提出通過將高維向量空間分解為多個子空間，并在每個子空間內獨立進行向量量化的方法，更有效地利用了通道間的互信息，并且對異常值展現(xiàn)出極強的魯棒性；接著通過 CUDA 異步流和高效的算子設計，充分利用了 GPU 的并行計算能力和內存層次結構，以支持乘積量化的高效執(zhí)行。實驗表明，對比主流框架 transformers 的半精度實現(xiàn)，MILLION 在 32K 上下文場景中同時達成 4 倍壓縮率和 2 倍加速比，并且在多種語言任務中精度表現(xiàn)幾乎無損。