高帶寬內(nèi)存的激增需求和供應(yīng)鏈約束正在推高AI和HPC基礎(chǔ)設(shè)施成本，也推動了系統(tǒng)設(shè)計的創(chuàng)新需求。內(nèi)存已成為當前AI周期的決定性約束之一。這種約束不僅限于硬件的可用性，還影響著現(xiàn)代AI模型的運行方式。

現(xiàn)在，谷歌的研究人員正在解決大語言模型推理中最消耗內(nèi)存的組件之一：鍵值（KV）緩存。在最近的一篇研究論文中，該公司介紹了TurboQuant，這是一種新的壓縮方法，旨在在不犧牲模型準確性的情況下，減少推理期間KV緩存的內(nèi)存占用。該方法旨在實現(xiàn)作者描述的壓縮和失真之間的近最優(yōu)權(quán)衡，接近模型數(shù)據(jù)在不破壞其結(jié)構(gòu)的情況下可以壓縮多少的理論極限。

KV緩存存儲先前token的中間向量表示，以便模型能夠生成響應(yīng)而無需從頭開始重新計算先前的token。這些向量捕獲了token之間的關(guān)系，是注意力機制工作原理的核心，使緩存在長對話或文檔的速度和響應(yīng)性方面變得至關(guān)重要。但這同時也使其成為內(nèi)存消耗的主要來源。隨著上下文窗口增長到數(shù)萬或數(shù)十萬token，緩存相應(yīng)擴展，即使是配置良好的系統(tǒng)也可能會被內(nèi)存需求迅速壓倒。

谷歌的方法是以非常低的精度壓縮此緩存，同時保持使注意力機制起作用的數(shù)學特性。根據(jù)該公司的說法，TurboQuant可以將KV緩存內(nèi)存使用減少大約六倍，在某些情況下將數(shù)據(jù)表示降低到每個值僅幾位。重要的是，這種壓縮不需要重新訓練模型或在校準數(shù)據(jù)上進行微調(diào)。作者聲稱，它可以在推理時直接應(yīng)用，對準確性的影響最小。

高壓縮、最小準確性損失和無重訓練要求的結(jié)合，是使這項研究值得注意的原因。壓縮技術(shù)并不新鮮，量化已被廣泛用于縮小模型權(quán)重。但KV緩存壓縮已被證明更加困難。涉及的數(shù)據(jù)結(jié)構(gòu)是高維的，對失真敏感。小誤差可以通過注意力計算傳遞并降低輸出質(zhì)量。

TurboQuant通過兩部分方法解決了這個問題。第一步稱為PolarQuant，將向量表示轉(zhuǎn)換為一種可以以非常低精度更高效壓縮的形式。第二步應(yīng)用基于Johnson-Lindenstrauss引理的輕量級校正機制，這是一種在保持高維空間中距離的數(shù)學技術(shù)。這一步驟補償了在壓縮期間引入的失真，有助于保持向量之間的相互關(guān)系，注意力機制使用這些關(guān)系來確定哪些token最重要。

結(jié)果是這個系統(tǒng)實現(xiàn)了比典型方法更高水平的壓縮，而沒有引入通常伴隨低精度的偏差或不穩(wěn)定性。實際上，這允許更多token同時存儲在內(nèi)存中，并使相同的工作負載能夠在更少的硬件上運行。

其影響可能不僅限于模型效率。TurboQuant方法可以在推理期間直接應(yīng)用，不需要重新訓練、校準數(shù)據(jù)或更改模型架構(gòu)，這一事實可能使其更容易集成到現(xiàn)有系統(tǒng)中，而無需重新設(shè)計模型的構(gòu)建或服務(wù)方式。

該工作還達到了歷史上難以實現(xiàn)的壓縮水平。先前KV緩存壓縮的方法在非常低的比特表示下難以保持穩(wěn)定性，而TurboQuant在每值約3至3.5比特的比特寬度下保持穩(wěn)定的性能。

這里仍然存在明顯的限制。結(jié)果基于基準評估而非生產(chǎn)規(guī)模系統(tǒng)，且該方法僅針對整體內(nèi)存占用的一部分。模型權(quán)重、激活和其他系統(tǒng)開銷保持不變，對高帶寬內(nèi)存的需求不太可能消失。即便如此，該研究反映了推理時效率的不斷擴大關(guān)注。隨著模型規(guī)模的持續(xù)擴大，像TurboQuant這樣的技術(shù)表明，不僅可以新硬件來解決內(nèi)存的一些壓力，還可以通過更高效地使用推理期間計算的中間表示。