本篇工作已被電子設(shè)計(jì)自動(dòng)化領(lǐng)域頂級(jí)會(huì)議 DAC 2025 接收，由上海交大計(jì)算機(jī)學(xué)院蔣力教授與劉方鑫助理教授帶領(lǐng)的 IMPACT 課題組完成，同時(shí)也獲得了華為 2012 實(shí)驗(yàn)室和上海期智研究院的支持。第一作者是博士生汪宗武與碩士生許鵬。
在通用人工智能的黎明時(shí)刻，大語(yǔ)言模型被越來(lái)越多地應(yīng)用到復(fù)雜任務(wù)中，雖然展現(xiàn)出了巨大的潛力和價(jià)值，但對(duì)計(jì)算和存儲(chǔ)資源也提出了前所未有的挑戰(zhàn)。在以 transformer 模型為基礎(chǔ)的大模型中，鍵值緩存雖然用以存代算的思想顯著加速了推理速度，但在長(zhǎng)上下文場(chǎng)景中成為了存儲(chǔ)瓶頸。例如，半精度的 LLaMA-2-7B 模型權(quán)重約 14GB，在上下文長(zhǎng)度為 128K 時(shí)鍵值緩存占據(jù) 64GB，總和已經(jīng)接近高端卡 NVIDIA A100 的 80GB 顯存容量上限。鍵值量化可被用于壓縮緩存，但往往受到異常值的干擾，導(dǎo)致模型性能的顯著下降。為此，本文的研究者提出了 MILLION，一種基于乘積量化的鍵值緩存壓縮和推理加速設(shè)計(jì)。
- arxiv 鏈接：https://arxiv.org/abs/2504.03661
- 開(kāi)源鏈接：https://github.com/ZongwuWang/MILLION
整型量化的軟肋：異常值
圖 1：矩陣量化可視化。紅色代表的異常值顯著大于其他值，導(dǎo)致均勻量化后高位編碼被浪費(fèi)。
量化中受到廣泛使用的整型均勻量化受到異常值的影響較為顯著。圖 1 展示了矩陣中的量化。在一組分布較為集中的數(shù)據(jù)中，一個(gè)顯著偏離其他值的異常值會(huì)導(dǎo)致其他值的量化結(jié)果全部落在較低區(qū)間，浪費(fèi)了高位編碼的表示能力。
圖 2：圖中使用 「通道熵」 定量比較不同方案的量化效果，越大表明越有效地利用了通道容量，即整型的寬度。沿通道量化只能解決沿該方向分布的異常值，而在面對(duì)另一方向異常值時(shí)效果不佳。
在實(shí)際的鍵值量化中，為了更好的表示能力，通常對(duì)于每個(gè)通道（即鍵值向量的維度）或每個(gè) token 采取不同的量化參數(shù)，這種方法被稱為沿通道量化（channel-wise quantization）或沿詞元量化（token-wise quantization）。然而，如圖 2 所示，沿特定方向量化只能解決沿該方向分布的異常值。
圖 3：實(shí)際采樣獲得的鍵值緩存分布。在 llama-2-7b-wikitext-103-v1-layer10-value 中，異常值并不遵循簡(jiǎn)單的沿通道分布，而是呈現(xiàn)為較復(fù)雜的點(diǎn)狀和團(tuán)狀。
研究團(tuán)隊(duì)通過(guò)實(shí)際采樣數(shù)據(jù)發(fā)現(xiàn)，在鍵值緩存中，沿通道方向分布的異常值占多數(shù)，但也存在并不明顯的情況，如圖 3 所示。這表明，上述量化方案并不是一勞永逸的解決方式，仍然存在優(yōu)化空間。
異常值的解決方案：乘積量化
圖 4：數(shù)軸上的均勻和非均勻量化對(duì)比。在對(duì) 8 個(gè)數(shù)據(jù)點(diǎn)進(jìn)行 2 比特量化過(guò)程中，均勻量化浪費(fèi)了 10 編碼。而基于聚類的非均勻量化則編碼更合理。
如圖 4 所示，非均勻量化通過(guò)聚類的方式允許量化區(qū)間不等長(zhǎng)，從而更合理地分配編碼，提升量化效率。研究團(tuán)隊(duì)觀察到，由于通道間的數(shù)據(jù)分布可能存在關(guān)聯(lián)（即互信息非負(fù)），將通道融合后在向量空間中聚類，效果一定不亞于獨(dú)立通道的量化，如圖 5 所示。
圖 5：左圖為兩個(gè)通道獨(dú)立進(jìn)行 1 比特量化，右圖為在通道融合后進(jìn)行 4 分類的 KMeans 聚類。融合通道量化的通道熵更加接近 2 比特的容量極限，展示出更好的量化效果。
由于高維空間中聚類較為困難，因此將整個(gè)向量空間劃分為多個(gè)低維子空間的笛卡爾積，可以平衡聚類復(fù)雜度和量化效果。這與最近鄰搜索中使用的乘積量化思想一致。研究團(tuán)隊(duì)通過(guò)實(shí)驗(yàn)發(fā)現(xiàn)，子空間維度為 2 或 4 是較好的平衡點(diǎn)。
推理加速手段：高效的系統(tǒng)和算子實(shí)現(xiàn)
圖 6：三階段的推理系統(tǒng)設(shè)計(jì)
圖 7：分塊注意力機(jī)制使得批量延遲量化成為可能
圖 6 展示了離線訓(xùn)練、在線預(yù)填充、在線解碼三階段的量化推理系統(tǒng)設(shè)計(jì)。其中，碼本訓(xùn)練（量化校準(zhǔn)）屬于秒級(jí)輕量化任務(wù)，并且離線進(jìn)行，不影響運(yùn)行時(shí)開(kāi)銷；在線預(yù)填充階段使用訓(xùn)練好的碼本對(duì)鍵值緩存進(jìn)行量化壓縮，達(dá)到節(jié)省顯存的目的；在線解碼階段采用分塊注意力機(jī)制的方法，將預(yù)填充階段的歷史注意力和生成 token 的自注意力分開(kāi)計(jì)算（如圖 7 所示），達(dá)成批量延遲量化的目的，掩藏了在線量化的開(kāi)銷，確保模型輸出的高速性。并且，在歷史注意力階段，由于歷史鍵值對(duì)數(shù)遠(yuǎn)大于碼本長(zhǎng)度，因此先用查詢向量與碼本計(jì)算好非對(duì)稱距離查找表（ad-LUT），可以大大減少內(nèi)積距離計(jì)算量，達(dá)到加速計(jì)算的目的。
圖 8：向量化加載可有效使帶寬飽和
在算子優(yōu)化方面，研究團(tuán)隊(duì)在 flash decoding 的基礎(chǔ)上使用了寬數(shù)據(jù)（如 float4）向量化加載的方式，將多個(gè)乘積量化編碼打包為寬數(shù)據(jù)，有效使帶寬飽和（如圖 8 所示）。同時(shí)，在表查找階段，子空間之間的表具有獨(dú)立性，并且可以被放入少量緩存行中，研究團(tuán)隊(duì)利用這一空間局部性極大提高了表查找的 L2 緩存命中率。此外，研究團(tuán)隊(duì)還仔細(xì)掃描了不同上下文長(zhǎng)度下可能的內(nèi)核參數(shù)，找到最優(yōu)配置，形成了細(xì)粒度的預(yù)設(shè)，在實(shí)際運(yùn)行時(shí)動(dòng)態(tài)調(diào)整，充分利用 GPU 的計(jì)算資源。具體實(shí)現(xiàn)可以在開(kāi)源倉(cāng)庫(kù)中找到。
實(shí)驗(yàn)結(jié)果
實(shí)驗(yàn)設(shè)置
圖 9：實(shí)驗(yàn)設(shè)置
實(shí)驗(yàn)采用了不同位置編碼、不同上下文長(zhǎng)度的多種模型進(jìn)行了詳細(xì)的評(píng)估。在模型性能方面，采用困惑度（Perplexity，PPL）和 Longbench 兩種指標(biāo)；在系統(tǒng)性能方面，采用每詞元輸出間隔（Time Per Output Token, TPOT）定量分析，并給出了注意力層詳細(xì)的剖析。對(duì)比采用方案和乘積量化參數(shù)如圖 9 所示。
模型性能
圖 10：困惑度指標(biāo)。其中 「-1%」 表示該方法額外存儲(chǔ) 1% 的異常值不參與量化。
困惑度越小表明模型輸出質(zhì)量越高。實(shí)驗(yàn)結(jié)果表明，MILLION 與額外處理了異常值的 SOTA 方案輸出質(zhì)量保持一致，展現(xiàn)出對(duì)異常值良好的魯棒性。而 SOTA 方案在不處理異常值的情況下可能會(huì)遭遇嚴(yán)重的輸出質(zhì)量損失。
圖 11：Longbench 問(wèn)答數(shù)據(jù)集得分展示
在長(zhǎng)對(duì)話問(wèn)答任務(wù)中，不同模型在各種數(shù)據(jù)集上的得分均表明，MILLION 方案能夠在 4 倍鍵值緩存壓縮效率下保持幾乎無(wú)損的表現(xiàn)。
系統(tǒng)性能
圖 12：每詞元輸出時(shí)間。對(duì)比其他方案，MILLION 的優(yōu)勢(shì)持續(xù)增長(zhǎng)，在 32K 上下文時(shí)達(dá)到 2 倍加速比。
圖 13：注意力層時(shí)間剖析
在 TPOT 評(píng)估中，MILLION 能夠在 32K 上下文語(yǔ)境下同時(shí)達(dá)成 4 倍鍵值緩存壓縮比和 2 倍端到端加速比。注意力層的深入分析表明，MILLION 在訪存和內(nèi)核函數(shù)方面對(duì)比 baseline 取得顯著優(yōu)勢(shì)。
總結(jié)
MILLION 的主要貢獻(xiàn)在于：（1）深入分析鍵值緩存分布；（2）提出基于乘積量化的非均勻量化算法；（3）設(shè)計(jì)高效的推理系統(tǒng)及內(nèi)核。研究團(tuán)隊(duì)首先證實(shí)了鍵值緩存中異常值存在的普遍性，并指出異常值的不同分布是當(dāng)前主流的量化方案精度不足的根本原因；然后提出通過(guò)將高維向量空間分解為多個(gè)子空間，并在每個(gè)子空間內(nèi)獨(dú)立進(jìn)行向量量化的方法，更有效地利用了通道間的互信息，并且對(duì)異常值展現(xiàn)出極強(qiáng)的魯棒性；接著通過(guò) CUDA 異步流和高效的算子設(shè)計(jì)，充分利用了 GPU 的并行計(jì)算能力和內(nèi)存層次結(jié)構(gòu)，以支持乘積量化的高效執(zhí)行。實(shí)驗(yàn)表明，對(duì)比主流框架 transformers 的半精度實(shí)現(xiàn)，MILLION 在 32K 上下文場(chǎng)景中同時(shí)達(dá)成 4 倍壓縮率和 2 倍加速比，并且在多種語(yǔ)言任務(wù)中精度表現(xiàn)幾乎無(wú)損。
熱門跟貼