GQLA:面向硬體自適應的大語言模型解碼的分組查詢潛在注意力機制
研究人員提出了分組查詢潛在注意力(GQLA),這是對DeepSeek多頭潛在注意力(MLA)的改進,在不重新訓練的情況下提供兩種硬體自適應的解碼路徑。該方法能在H100和H20 GPU上實現高效推理,幷包含TransGQLA用於轉換預訓練的GQA模型。
文章情報
要點
- GQLA在DeepSeek的MLA基礎上擴充套件出雙解碼路徑(MQA吸收和GQA),以適應不同硬體的roofline模型。
- 同一組GQLA權重既可用於H100(MQA路徑),也可用於H20(GQA路徑配合多token預測)。
- TransGQLA可將預訓練的GQA檢查點轉換為GQLA,在LLaMA-3-8B上將KV快取壓縮至28.125%。
- 在GQA路徑上支援多達8路零冗餘張量並行。
為什麼重要
這條新聞值得關注,因為GQLA在DeepSeek的MLA基礎上擴充套件出雙解碼路徑(MQA吸收和GQA),以適應不同硬體的roofline模型。
技術影響
可能影響模型選型、推理成本、產品能力和評測基準。
在大型語言模型推理中,注意力機制是決定效能和效率的關鍵元件。DeepSeek-V2/V3採用的多頭潛在注意力(MLA)透過將鍵和值聯合壓縮到低秩潛在空間,幾乎完美匹配了H100 GPU的計算與儲存roofline。然而,MLA的訓練權重僅暴露一種解碼路徑——吸收式多查詢注意力(MQA)形式。這意味著高效推理嚴重依賴於H100級別的計算頻寬比,無法利用沿頭維度的張量並行,並且在像H20這樣受出口限制的商用GPU上無法獲得多token預測(MTP)帶來的加速優勢。
針對這一侷限,研究者提出了分組查詢潛在注意力(GQLA),作為MLA的最小修改方案。GQLA的訓練權重在同一引數集上暴露兩種代數等價的解碼路徑:一條是與MLA相同的MQA吸收路徑,另一條是帶有每組擴充套件快取的分組查詢注意力(GQA)路徑。在執行時,系統會根據目標硬體的特性自動選擇最合適的路徑,無需重新訓練或自定義核心。因此,同一組GQLA權重能夠在H100上使用MQA吸收路徑(s_q=1)達到roofline效能,同時在H20上切換到GQA路徑(s_q=2)並利用MTP加速,並在GQA路徑上支援最多8路的零冗餘張量並行。
為了降低部署成本,研究者還將TransMLA擴充套件為TransGQLA,用於將預訓練的GQA檢查點轉換為GQLA模型。在LLaMA-3-8B上的實驗顯示,在MQA吸收路徑上,每個token的KV快取大小壓縮至GQA基線的28.125%,同時在每組路徑上保留了GQA級別的流量結構。這一方法不僅提升了推理效率,還增強了模型在不同硬體環境下的部署靈活性。GQLA的提出為大型語言模型在多樣化硬體上的高效推理提供了一種簡潔實用的解決方案,尤其對受出口限制的GPU部署具有重要意義。