JOVANA
Explore Library Glossary Getting Started Three Levels Fields How it works Mission
Join the mission
All guides

塞得進記憶體的注意力

服務時真正的瓶頸是 KV 快取。多查詢、分組查詢與滑動視窗注意力,是架構上的解答。

沒人警告你的那個快取

模型生成文字時,一次產生一個 token,而每產生一個新 token,都必須回頭注意每一個先前的 token。每步重算全部會無望地慢,所以模型保留一個 KV 快取:每個過去 token 的鍵與值存在記憶體裡重複使用。這讓生成變快——但快取隨每個 token 成長,在長對話裡,塞不進 GPU 的東西變成是它、而非模型權重。

M_{\text{KV}} = 2 \cdot L \cdot H \cdot d_h \cdot S \cdot B

KV 缓存占用的内存随每个因子同时增长——层数、注意力头数、每头维度、序列长度和批量——所以削减头数才是关键。

多查詢:所有頭共用一組鍵

標準的多頭注意力給每個頭自己的查詢、鍵、值——比方 32 個頭,就有 32 組獨立的鍵/值要快取。多查詢注意力(MQA)保留 32 個查詢頭,但讓它們全部共用單一鍵頭與單一值頭。快取一夜縮小 32 倍,生成也變快,因為每步要從記憶體讀的少得多。

标准多头注意力让每个头都有自己的键和值——32 个头就要缓存 32 套键/值,这正是多查询注意力要削减的开销。

多头注意力示意图:多个并行的头分别产生输出,再拼接并混合。

代價是:把所有鍵與值塌縮成一個頭,丟掉了表達空間,純 MQA 模型品質會掉一點,訓練也可能變得不穩。這是個正確、但稍微走過頭的想法。

分組查詢:甜蜜點

分組查詢注意力(GQA)是如今人人都採用的折衷。不是 32 個鍵/值頭(完整多頭)也不是 1 個(MQA),而是用少數幾個——比方 8 個——讓每個鍵/值頭由一 4 個查詢頭共用。你幾乎拿回完整注意力的全部品質,同時把快取砍 4 倍。它是 Llama、Mistral、Qwen 與多數現行開放模型的預設,因為它正好坐在品質對記憶體的甜蜜點上。

分组查询注意力保留许多查询头,但让一组查询头共享同一套键/值——这正是 GQA 所节省的查询-键-值软查找。

注意力示意图:查询通过 softmax 与各个键匹配,从而对值加权。

# config knobs you will actually see
num_attention_heads: 32   # query heads
num_key_value_heads: 8    # GQA: 32/8 = 4 query heads share each KV head
# num_key_value_heads == 32 -> plain multi-head
# num_key_value_heads == 1  -> multi-query
一個數字——鍵/值頭的數量——讓你在完整注意力、GQA 與 MQA 之間滑動。

滑動視窗:別什麼都看

GQA 縮小每個 token 的快取,但快取仍隨序列長度成長。滑動視窗注意力直接攻擊長度:不讓每個 token 回頭注意整段歷史,而是限制在最近的 W 個 token——一個固定視窗,比方 4096。一旦超過視窗,快取就停止成長,因為舊的鍵可以丟掉。

点击一个词,看它关注哪些词——滑动窗口注意力只是把这种关注范围限制在最近的固定窗口内。

交互式自注意力:选中一个词元会高亮它所关注的其他词元。

這不會丟掉長距離資訊嗎?比你想的少。疊很多視窗層,有效觸及範圍會累積:一個 token 可以透過一連串重疊視窗,影響到 4096 個之外的 token,就像深層卷積網路用小濾波器看見大圖。許多模型把視窗層與完整注意力層交錯,以取得便宜的局部性偶爾的全域觸及——這是最後一篇長上下文架構的初嘗。