AI TECHNOLOGY

AI Technology ClassroomLesson 2 / 3

From Tokens and Attention to the KV Cache

A precise, example-led walkthrough of tokenization, position information, Query-Key matching, Value mixing, residual connections, and what the KV cache does and does not save during autoregressive decoding.

11 min read

先試著不要想 Transformer、矩陣或 cache。假設你只看到一句話的前半段,現在要把最後一格填起來。人讀過一次,大概不需要停頓;對模型來說,這一格卻是一連串小動作的結果。

這篇就跟著那一格往回走。我們會看文字怎麼被拆開、模型怎麼在前文裡找線索,以及答案出現後,哪些筆記可以留到下一輪繼續用。工程名詞還是會出現,但都等到它真的派上用場時再介紹。

全文只用這一句 小美養貓,小強養狗。貓吃魚,狗吃肉。小美的寵物吃__。
目前在看
小美 小強 小美寵物

為什麼用這一句

先自己讀一次:「小美的寵物吃__。」要填出「魚」,第一步不是回想貓喜歡吃什麼,而是回到前文確認小美養的是哪一種動物。找到「小美養貓」之後,還得再往後找「貓吃魚」。答案其實經過兩次連結才到手。

句子裡還故意放了「小強養狗、狗吃肉」。它不是多餘的背景,而是一個干擾項。如果只看到「寵物」就隨便抓附近的動物,答案很可能會跑到狗或肉。模型必須分辨哪些字跟眼前的問題比較有關。

所有線索也都放在空格左邊。這很重要,因為自回歸語言模型在產生下一個 token 時,只能使用已經出現的上下文。這個小例子同時具備左側線索、干擾項與兩跳關係,後面的每個步驟都能回到同一句話檢查。

01|先把文字切成 token

人看到「小美養貓」會直接讀出意思,模型不能把整句中文原封不動拿去做矩陣運算。它得先把文字拆成一塊一塊的積木,每一塊再換成編號。

這些文字積木叫做 token;負責切分的工具叫 tokenizer;查到的編號則叫 token id。你可以把 token id 想成圖書館裡的索書號:它方便系統找到資料,但編號本身不等於書的內容。

小美#4821#1930#7715#11小強#5602#1930#7208#12

圖上的編號只是教學示意,不是真實 tokenizer 的輸出。Token 也不一定剛好是一個中文字或一個完整英文單字;它可能是一個字、一段常見字串或標點。換一個模型,同一句話甚至可能被拆成不同數量的 token。

例句裡有兩個「小美」。在同一個 tokenizer 下,它們會拿到相同的 token id,但一個出現在句首,另一個靠近空格,功能顯然不同。模型還需要位置資訊,才知道誰先出現、誰離眼前的問題比較近。這些位置資訊可以透過 positional embedding、RoPE 或其他方法加入,因此「字相同」不代表兩個位置會被當成同一件事。

02|Token id 先變成向量

有了 token id,模型只是知道「該去哪一格拿資料」,還沒有真正可以運算的內容。下一步,它會依編號查出一長串學習而來的數字,這串數字叫 embedding vector。

如果 token id 像索書號,embedding 就比較像那本書的內容特徵:談動物、飲食還是人物關係,都藏在許多維度的數值裡。我們常把它畫成平面座標,方便想像「語境相近的 token 在表示空間中可能呈現相近結構」;真實向量的維度則高得多。

同一個 token id 會查到同一份 token embedding。

但完整表示還會受到位置與前文影響。「小美」第一次和第二次出現,後面的 hidden state 不會一樣。

到這裡仍然比較像查表:給定編號,讀出一份 token embedding。進入 Transformer layers 之後才開始大量計算。模型會讓每個位置一邊保留自己的內容,一邊吸收前文中與自己有關的資訊;那份一路被更新的表示,就是後面常看到的 hidden state。

03|同一份輸入投影成 Query、Key、Value

現在回到句尾的「吃」。它想補出下一個字,得先從前文找出跟自己最有關的線索。為了完成這件事,每個位置的 hidden state 會準備三種不同用途的向量。

Query 是目前想找什麼,Key 是自己可以怎麼被找到,Value 則是被找到之後真正交出去的內容。句尾的「吃」帶著 Query 出發;前面的「小美」「貓」「魚」各自亮出 Key;配對完成後,模型再取回它們的 Values。

Q

Query/我要找什麼

目前位置帶著 Query 去比對其他位置。

K

Key/我能被怎麼找到

Key 決定這個位置與 Query 的匹配程度。

V

Value/被找到後交出什麼

Value 是加權混合時真正被取回的內容。

工程上,Q、K、V 都是由同一份輸入分別乘上不同矩陣得到的,可以平行計算,沒有「先算 Q,再用 Q 算 K」的順序。這三組矩陣則是在訓練過程中一起學出來的。

「Key 是標籤、Value 是內容」只是幫助入門的比喻。Key 不是資料庫裡非黑即白的分類欄位,而是連續向量;它跟 Query 的配對也不是「找到或沒找到」兩種結果,而是一連串高低不同的分數。

04|Query 對照 Keys,再混合 Values

把注意力放回例句,就比較容易看懂。句尾正在處理「小美的寵物吃什麼」,所以「小美」「貓」「魚」應該得到較高比重;「小強」「狗」「肉」仍在上下文裡,但和眼前問題的關係比較弱,比重就低一些。

這不是把不相關的字整個丟掉。標準 dense attention 比較像調音量:有些線索開得很大,有些只留一點聲音。模型再依照這些音量,把各位置的 Value 混合成新的資訊。

小美
最高
小強
低,但不是零

正式寫法是:Query 與可見位置的 Keys 計算相似度,經過縮放、causal mask 與 softmax 後得到 attention weights,再用這些權重加權混合 Values。上圖的數字只是示意;沒有被 mask 掉的位置通常都會分到權重,只是大小可能差很多。未來尚未出現的 token 會被 causal mask 擋住,不能偷看。

05|Attention 輸出不等於最後答案

Value 混合完成後,模型還沒有直接得到「魚」。比較好的想像方式是:句尾的「吃」原本有一張自己的卡,attention 從前文帶回一些線索,再把這些線索寫回卡片。卡片沒有被換掉,只是內容變得更完整。

進入這一層之前,它主要帶著「吃」這個位置已經累積的資訊;經過 attention 之後,又多知道了哪些人物、動物和食物與自己相關。往後每經過一層,這張卡還會繼續更新。

工程上,加權混合的結果會通過 output projection,並透過 residual connection 加回原本的表示。Transformer block 通常還包含 normalization 與 MLP,各模型的排列也可能不同。這些運算共同更新 hidden state,所以 attention output 只是中間結果,不是模型最後吐出的答案。

06|為什麼需要多層

「小美的寵物吃什麼」不只要找一次。可以把它想成讀兩遍:第一遍先整理人物關係,讓「貓」這個位置吸收「我是小美養的」;第二遍再從句尾出發,透過已經帶有飼主資訊的「貓」,找到它對應的食物「魚」。

這也是多層 Transformer 的直覺價值。每一層都讓各位置重新整理一次上下文,前一層剛吸收的關係,到了後一層就能成為新的查找線索。

淺層第一跳:小美 → 貓

「貓」的位置吸收「小美養的」這層關係。

較深層第二跳:貓 → 魚

句尾再利用已經帶有飼主資訊的「貓」,連到它吃的「魚」。

這是理解層數的概念圖,不是模型內部的逐層錄影。

真實模型有多個 attention heads、MLP 與分散式表示,不一定整齊地由某一層做第一跳、下一層做第二跳。

07|真正拿去猜下一個 token 的是哪一份資料

走完最後一層後,句尾那張卡已經帶著模型整理過的上下文。接著,系統會把所有可能的下一個 token 排成一張排行榜。在這個例句裡,我們期待「魚」排在最上面,於是它被選出來,接回原句末尾。

這張排行榜的工程名稱是 logits。模型取最後位置的 final hidden state,經過輸出層轉成整個詞彙表的 logits,再換算成下一個 token 的機率。

那前面那些位置呢?它們不是算完就全部丟掉。每一層產生的表示會參與後續 attention,也會留下下一輪需要的 K/V。只是在這個 decoding step,真正拿來挑選下一個 token 的,是最後位置那一列 logits;其他位置不需要再各選一次答案。

08|KV cache:省下重算,沒有省下重讀

「魚」接回句尾後,模型還要繼續猜下一個 token。這時最直覺的笨方法,是把整句話從頭讀一遍,重新替每個舊 token 計算每一層的 Key 與 Value;答案越長,重複工作就越多。

KV cache 像把已經做好的閱讀筆記留在桌上。舊 token 的 K/V 已經算過,就不用每次重新抄一遍。新 token 只要做好自己的 Q/K/V,再拿新的 Query 去翻閱那些舊筆記。

不過,「不用重抄」不等於「不用翻閱」。上下文裡保存的 K/V 還是得被讀出來,才能判斷哪些內容跟新 Query 有關。這就是 KV cache 最容易被誤解的地方。

下一個 decoding step 的工作是否仍要做
重算舊 token 的 K/V不用,直接讀 cache
新 token 計算自己的 Q/K/V
新 token 走完所有 Transformer layers
新 Query 對照所有 cached Keys
依權重讀取所有需要的 cached Values

表格可以拿來做最後核對:舊 K/V 不必重算,新 token 卻仍要走完所有 Transformer layers,也要用新 Query 對照 cached Keys、讀取需要的 cached Values。

因此,上下文越長,桌上的筆記就越厚。KV cache 會占用更多容量,每生成一個新 token,也可能需要搬動更多資料。長上下文的代價不只來自運算,還包括記憶體容量與頻寬。

為什麼存 K/V,不存舊 Query

Query 像一次性的提問。舊 token 當時問過什麼,完成那一輪之後就不再需要;下一輪真正要發問的是新 token。相反地,舊 Keys 仍要讓後面的 Query 查找,舊 Values 也仍是被取回的內容,所以它們必須留下。

換成工程說法:每一步只需要新 token 的 Query,拿它去對照所有舊 Keys,再依權重混合 Values。舊 Query 不會參與新 token 的這一列 attention。

一句話記住:KV cache 省了重寫,沒省重讀。

這也是 LLM decoding 常受記憶體頻寬影響的原因之一。

09|MQA 與 GQA 為什麼從 K/V 下手

如果 KV cache 是倉庫,真正會一輪一輪留在裡面的是 K/V。Query 用完就離開,不會長期占著貨架。因此,想縮小倉庫、減少每一步要搬的資料,最直接的方向就是減少 K/V heads。

傳統 multi-head attention 讓每個 Query head 都有自己的 K/V heads。Multi-Query Attention(MQA)讓多個 Query heads 共用一組 K/V;Grouped-Query Attention(GQA)則折衷處理,讓一群 Query heads 共用較少量的 K/V heads。

減少 K/V heads 可以縮小 KV cache,也降低自回歸 decoding 時反覆讀取 K/V 的記憶體頻寬。不過這不是免費改動,設計時仍要在模型品質、訓練方式、吞吐量與延遲之間取捨。

三個容易講錯的地方

讀到這裡,可以回頭檢查三個常見誤會。Key 不是「貓、狗、人物」這種離散類別,而是拿來和 Query 做軟性匹配的連續向量,所以干擾項通常不是直接變成零。

同一個 token 的 K/V 也不會從第一層一路固定到最後。每一層都會根據該層輸入重新投影;hidden state 改變,下一層算出的 Q/K/V 也跟著改變。

最後,attention output 不是模型已經寫好的答案。它還要經過 projection、residual、MLP 與後續 layers。走到最後一層後,最後位置的 hidden state 才會被轉成 token logits。

本課整理

整段流程其實可以收回一句話:文字先被切成 token,位置與語意被放進向量;Query 回頭對照 Keys,再把相關的 Values 帶回來,逐層更新 hidden state。最後位置的表示經過輸出層,才變成下一個 token 的候選排行榜。

答案接回句尾後,下一輪又開始。KV cache 讓舊 K/V 不必重新計算,MQA/GQA 則進一步減少需要保存的 K/V heads。但新 token 仍要計算,也仍要讀取前文留下的資料。

KV cache 省下重算,卻沒有省下重讀。

References

  1. Vaswani et al., “Attention Is All You Need”:Transformer、scaled dot-product attention、multi-head attention 與 residual 架構。
  2. Su et al., “RoFormer: Enhanced Transformer with Rotary Position Embedding”:RoPE 如何把位置資訊帶進 self-attention。
  3. Shazeer, “Fast Transformer Decoding: One Write-Head is All You Need”:MQA 與 incremental decoding 的 K/V 記憶體頻寬問題。
  4. Ainslie et al., “GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints”:GQA 的 K/V head 分組與品質、速度取捨。
  5. Kwon et al., “Efficient Memory Management for Large Language Model Serving with PagedAttention”:KV cache 記憶體管理如何影響 LLM serving 吞吐量。

例句、token id 與 attention weights 都是教學示意。不同模型的 tokenizer、位置編碼、Transformer block 與推論實作會有所不同。

Learning guide

AI Technology Classroom

0 / 3

Prerequisites

  • Tokens, vectors, and basic matrix multiplication

What I learned

  • Explain how position information distinguishes repeated tokens
  • Trace Query-Key matching and Value mixing
  • Separate attention output from the final hidden state
  • Explain what KV cache saves and what it still needs to read

Key terms

Open glossary →

Further reading

Knowledge check

1. How can a model distinguish the same token at two positions?
2. What can KV cache reuse from old tokens?
3. What does KV cache not eliminate?
4. Why do MQA and GQA reduce decoding memory pressure?

Thanks for reading.

Take the concept with you, not just the terminology.

#AI Technology#LLM#Token#Embedding#Position Encoding#RoPE#Transformer#Attention#Query#Key#Value#KV Cache#MQA#GQA#Inference#Memory Bandwidth