AN の日記

優化 RAG 檢索的數個方法

Tue, 12 May 2026 00:00:00 +0000

RAG (Retrieval Augmented Generation) 透過外部來源取得資訊，並將資訊提供給 LLM 產生回答，成功地解決了 LLM 回答受限於其訓練資料的問題。但是我們希望檢索到的答案是我們想要的，否則只是一坨垃圾話。

因此這邊整理了一些或許有用的方法，當檢索成效不如預期時，不妨參考看看嘿

我分成三個階段做說明，分別是前檢索階段、檢索階段、後檢索階段

前檢索階段：資料前處理，優化資料索引與使用者查詢

檢索階段：優化向量搜尋

後檢索階段：過濾檢索到的資訊

前檢索階段

此階段重點是如何將資料建立索引 (index)，以及檢索向量前可以優化的方法

建立索引

滑動視窗 (Sliding Window)：相鄰的文字分段 (chunk) 邊緣之間產生重疊部分，避免重要資訊發生在邊緣時被切開來，然後遺失
注意切割邊緣：切割時請避免斷在句子一半。以段落結尾為斷點會是相對好的方法。像是 LangChain 的 RecursiveCharacterTextSplitter 通常比 TokenTextSplitter 來得好
添加 Metadata：附上日期、URL、章節、頁面等資訊，有助於後續檢索能過濾結果
優化索引結構：使用不同索引技巧，像是使用多種 chunk size 來切分。像是 PDF 內的資料有文字圖片穿插，怎麼切特別重要
資料清洗：去除垃圾話、驗證事實正確性、更新過時資訊等
由小到大檢索 (Small-to-Big Retrieval)：用一小段中最重要的部分轉化為嵌入作為索引，並將周圍的上下文保留在 metadata 中。生成回應時再提供給 LLM。解決了小區塊不含雜訊但資訊量不足，大區塊資訊充足但雜訊太多的問題。LangChain 可以參考 ParentDocumentRetriever

檢索向量前

查詢路由 (query routing)：當資料爆多時，可以考慮將資料分類。例如資料有文章、程式碼兩種，可以將資料庫分兩個；在使用者問程式碼相關的問題時，將後續檢索範圍導到程式碼資料庫中。增加檢索效率的同時也減少檢索錯誤資料的機會
查詢改寫 (query rewriting)：將使用者的查詢換句話說，但保留原本意思，增加與向量資料庫匹配的機會。例如將較少見的詞彙換成較常見的，擴大搜尋範圍；對於較長的 query 也可以考慮拆成多個 sub-query，分開檢索
HyDE (Hypothetical Document Embeddings)：透過 LLM 先對 query 生成一個假設性回應，然後用原始 query 與上述假設回應進行檢索
查詢擴展 (query expansion)：在使用者的 query 加入額外字詞讓查詢更全面。例如有人搜尋『沙發』，可以加上同義詞像是『長椅』或『Couch』
關鍵字過濾：先以 LLM 偵測查詢中如果出現特殊關鍵字例如人名地名，可以利用關鍵字過濾向量搜尋的範圍

檢索階段

檢索主要有兩種：基於語意的搜索 and 基於關鍵字的搜索

語意搜索其實就是比對 query vector 與 chunk vector 的向量相似度，相似度越高代表語意越接近

關鍵字搜索會利用出現文字與出現頻率去計算分數，方法像是 TF-IDF, BM25

微調 Embedding 模型：向量資料庫是從 embedding 來的，如果應用場合含有大量罕見詞彙，可以考慮微調模型
利用 LLM 模型引導：如果不想微調模型，也可以考慮用 Instructor Embedding Model (例如這個 hkunlp/instructor-large · Hugging Face) 告訴目前任務背景，在 embedding 時去調整 query 到符合資料庫的資料
善用 metadata 過濾與搜尋：如果資料庫有附帶 metadata，可以用它幫忙過濾
混合搜尋 (hybrid search)：語意搜索+關鍵字搜索。兩種搜索各自有優缺點，因此利用兩種搜尋得到各自分數，然後利用 RRF (Reciprocal rank fusion) 根據權重算出最後的排名。權重可以根據領域做調整

後檢索階段

檢索出來的資料會餵給 LLM 協助生成答案。為了避免受到 LLM 的 context window size 等限制，這部分也是可以考慮優化的一個階段

提示詞壓縮：如果訊息太長，在保留核心訊息的前提下，將廢話去除。著名技術有 LLMLingua
重新排序 (Re-ranking)：將檢索後的資料排名重新排序。利用『使用者輸入』與『每一個檢索片段』的匹配分數去排序。匹配分數可以由 Bi-Encoder (兩個片段的向量相似度) 得到，或是利用 Cross-Encoder (Reranker) 得到。後者效果較好但速度較慢。此外HuggingFace 有提供許多 Reranker。要注意有些 Reranker 不支援中文。針對 Reranker 的研究可以看這一篇使用繁體中文評測各家 Reranker 模型的重排能力 – ihower { blogging }

另外現在也有基於圖的方法 GraphRAG，又名知識圖譜，著名工具有像是 Neo4j

不同的應用場景適合不同方法，沒有絕對好的方法，只有多實驗並評估才是真理

【筆記】llama.cpp 的 KV Cache

Mon, 09 Feb 2026 00:00:00 +0000

前陣子在研究 llama.cpp 框架支援的 KV Cache 操作，順手將一些常用或者有關 KV Cache 的參數記錄一下。預設值用粗體表示，隨時可能會變所以僅供參考。

llama-server 參數

https://github.com/ggml-org/llama.cpp/tree/master/tools/server

1

llama-server -m [模型.gguf] -np 1 -c 2048

-m [gguf模型]
-np [N]：Slots 數量（可同時處理的 requests）
-c [N]：Context window size。預設 0 表示從模型 metadata 讀取

KV 量化

-ctk [f32, f16, bf16, q8_0, q4_0, q4_1, iq4_nl, q5_0, q5_1]：Key cache 量化格式
-ctv [f32, f16, bf16, q8_0, q4_0, q4_1, iq4_nl, q5_0, q5_1]：Value cache 量化格式

Offloading

-ngl [N|auto|all]：放在 GPU 的模型層數。CPU + GPU 混合推論是 llama.cpp 的特色。auto是2025年底這個 PR 新加的…
-kvo/-nkvo：KV Cache 是否要 offload 到 GPU
–mlock：把模型檔案 lock 在 RAM，避免被 OS swapping 到 disk
–mmap/–no-mmap：是否 mmap 模型，預設開啟（甚麼時候會想關掉？例如特定情況想減少 pageout 時）
-ts [N0,N1,N2]：模型 offload 到 GPU 的比例，例如雙 GPU 給 [3,1]

Context Shift

Prompt + generated tokens 超過 context window size 就不會報 error 了；反之會自動丟棄最前面的訊息

–context-shift/–no-context-shift (2025/08 預設變 disabled)：解決長文問題；context 超過 window size 時會 truncate
-keep [N]：位移時想保留的 initial prompt tokens 數量。預設 0。

system prompt 通常不想被丟棄，這時就會用到 keep

KV Shifting

讓 prompt 前綴不用完全一樣也可以複用 KV Cache（模型需要是 RoPE）

–cache-reuse [N]：透過 KV Shifting 複用 KV 時的最小 chunk size。預設 0 表示不啟用

Chunk size 太小有個缺點：復用的 KV Cache 意義上可能會不一樣

Slots 相關

–slot-save-path [路徑]：存放 KV Cache of slots 的位置

常見 KV Cache 優化 (預設都開啟)

-fa [on|off|auto]：FlashAttention
-cb/-nocb：Continous Batching
–cache-prompt/–no-cache-prompt：啟用 Prompt caching（複用同 slot 上一筆 prompt 的 KV Cache）

Server API Endpoints

POST /completion (not OAI-compatible)

1
2
3
4


curl --request POST \
 --url http://localhost:8080/completion \
 --header "Content-Type: application/json" \
 --data '{"prompt": "1+1=?", "n_predict": 128}'

1
2
3
4
5
6
7


{
 "prompt": "1+1=?",
 "n_predict": 128,
 "top_k": 40,
 "grammar": "root::=[0-9]+",
 "cache_prompt": true
}

“cache_prompt”：啟用 Prompt caching（複用同 slot 上一筆 prompt 的 KV Cache）。預設打開
“id_slot”：指定的 slot。預設 -1，Server 會選擇 idle slot 中相似度最高的 (common prefix/input tokens)

KV slots 操作

POST /slots/{id_slot}?action=save (將 slot N 的 KV cache 寫到指定 filename)
POST /slots/{id_slot}?action=restore (將 filename 內的 KV cache 讀到 slot N)
POST /slots/{id_slot}?action=erase (刪掉 slot N 的 KV Cache)

GET /slots

查看所有 slots 狀態

Tokenizer演算法詳解：BPE, WordPiece, Unigram

Fri, 09 Jan 2026 00:00:00 +0000

本文內容難度：★ ★ ☆ ☆ ☆

建議閱讀對象：想知道 Tokenizer 是如何建立詞彙表（vocabulary），以及如何將句子切割成 subwords 的人。

大型語言模型 (LLM) 在推論時的第一步為 Tokenization，Tokenizer 會將句子拆解成最小的語意單位（Tokens）。可以是一個單字（word）、一個字母（character）或一個子詞（subword）。來看看有哪些方式吧！

Word-based

把句子中的每個單字（word）當成不同 token。使用空白字元與標點符號隔開即可達成。

優點：簡單
缺點：如果有 50 萬個英文單字，輸出的 dimension 就高達 50 萬導致參數過大；另外也無法處理不在詞彙表（Out of vocabulary, OOV）的單字

Character-based

把句子中的每個字母（character）切割成不同 token。

優點：簡單、字彙量只要 26 個英文字母加特殊符號、不會有 OOV 的問題
缺點：token 太多導致序列太長，可能會超過上下文上限；單一字母代表的意義不足導致 tokenizer 很難訓練

上面兩種方法都有明顯的缺點，因此普遍不被使用。接下來要介紹的是目前常見的三種 Subword Tokenizer 方法，會分為訓練階段以及切割階段去說明

訓練階段： 指定 Vocabulary Size（例如希望最後有 100000 個 token），利用訓練資料 (corpus) 建立出 Vocabulary

切割階段： 將句子根據 Vocabulary 去切割成 tokens

Subword Tokenizer：結合 Word-based 機制只保留常見的單字避免參數過大，同時也保留了 Character-based 的機制，可以把單字拆成更小的單元。

例如 “unfortunately” 可能被拆成 “un” + “fortunate” + “ly” 共 3 個 tokens。其中字首 un- 與字尾 -ly 都是常見且有意義的。

Subword Tokenizer 的比較

BPE (Byte-Pair Encoding)

A. 訓練階段

反覆合併出現頻率最高的 pair 並加入詞彙裡。

直接舉個例子～假設今天訓練資料來源如下：“hug” 出現 10 次，“pug” 出現 5 次 … 以此類推

1
2


# 訓練資料 (Corpus)
("hug", 10), ("pug", 5), ("pun", 12), ("bun", 4), ("hugs", 5)

一開始先把出現過的字母都加進詞彙中

1

Vocabulary = ["b", "g", "h", "n", "p", "s", "u"]

接著我們會計算詞彙中每個 pair 在訓練資料裡出現的頻率，並挑選頻率最高的合併並加入 vocabulary，並合併 corpus

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13


# 訓練資料 (Corpus)
("h" "u" "g", 10), ("p" "u" "g", 5), ("p" "u" "n", 12), ("b" "u" "n", 4), ("h" "u" "g" "s", 5)

# 每個 pair 的出現頻率 (由高到低排序)
("u", "g") --> 20 次
("p", "u") --> 17 次
("u", "n") --> 16 次
("h", "u") --> 15 次
...

# 將出現頻率最高的 ("ug") 加入 vocabulary，同時更新 corpus
Vocabulary = ["b", "g", "h", "n", "p", "s", "u", "ug"]
Corpus: ("h" "ug", 10), ("p" "ug", 5), ("p" "u" "n", 12), ("b" "u" "n", 4), ("h" "ug" "s", 5)

接著就反覆重複這個動作：計算每個 pair 的出現頻率 ➜ 挑選頻率最高的合併並加入 vocabulary 同時合併 corpus …

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13


# 訓練資料 (Corpus)
("h" "ug", 10), ("p" "ug", 5), ("p" "u" "n", 12), ("b" "u" "n", 4), ("h" "ug" "s", 5)

# 出現頻率 (由高到低排序)
("u", "n") --> 16 次
("h", "ug") --> 15 次
("p", "u") --> 12 次
("p", "ug") --> 5 次
...

# 將出現頻率最高的 "un" 加入 vocabulary，同時更新 corpus
Vocabulary = ["b", "g", "h", "n", "p", "s", "u", "ug", "un"]
Corpus: ("h" "ug", 10), ("p" "ug", 5), ("p" "un", 12), ("b" "un", 4), ("h" "ug" "s", 5)

直到 vocabulary size 達到我們設定的目標 (10) 就停止，這就是我們目前 Tokenizer 所有的詞彙了

1
2


# 這是最終結果 (假設設定的目標為 vocabulary size = 10)
Vocabulary = ["b", "g", "h", "n", "p", "s", "u", "ug", "un", "hug"]

到目前為止 Tokenizer 已經訓練好 vocabulary了，現在來看看它是怎麼將句子切割成 subwords

B. 切割階段

我們現在給一個句子 "bugs" ， BPE 用剛剛建立的詞彙表與合併規則去切割 tokens，最後得到結果 ["b", "ug", "s"] ，步驟如下：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25


"""
筆者備註：在這個例子中沒有特別去處理 out of vocabulary 的問題；
 不過實際應用上不會有 OOV 的問題 (因為訓練資料一定會出現 a-z 與所有符號)
"""

# 剛剛建立的詞彙表
Vocabulary = ["b", "g", "h", "n", "p", "s", "u", "ug", "un", "hug"]

# 訓練階段建立的合併規則 (按照產生的先後順序)
1. (u, g) -> ug
2. (u, n) -> un
3. (h, ug) -> hug

# -------------------------------------------- #
# 給一個輸入句子 s = "bugs"
# 一開始先拆成最小單位 ["b", "u", "g", "s"]
s.tokens = ["b", "u", "g", "s"]

# 根據上面的合併規則(1~3)順序，依序檢查是否有出現在句子裡
# 1. 檢查規則1 (u,g) --> 中間有出現，合併它們！
s.tokens = ["b", "ug", "s"]
# 2. 檢查規則2 (u,n) --> 沒出現，跳過
# 3. 檢查規則3 (h,ug) --> 沒出現，跳過

最後得到結果 ["b", "ug", "s"]

在實務上，BPE 會使用 Ġ 表示單字前面的空格。例如 "Hello world" 在最一開始會被解析成兩個單字 (Hello, Ġworld) ，目的是為了讓 Tokenizer 能夠「看見」空格，以及具有可逆性：如果沒有 Ġ 來標記空格，當你把 Token 轉換回文字時，會不知道哪裡該加空格。

WordPiece

A. 訓練階段

由 Google 發明，用在 BERT 系列模型上。

使用 ##來代表該 subword 是前面一個 subword 的延續，它們在原始文本中是連在一起的。例如 “word” 一開始會被切成 w ##o ##r ##d ；("play", "##ing") 在原始文本表示單字 playing，具有可逆性。

跟 BPE 一樣，WordPiece 學習合併規則。

跟 BPE 的差別差在 WordPiece 不是選擇最頻繁的 pair 去合併，而是使用以下公式計算 pair 的分數

1
2
3
4


# 假設 pair 為 "hap","py"

# 分數 = "happy"出現頻率 / ("hap"出現頻率 × "py"出現頻率)
score = (freq_of_pair) / (freq_of_first_element × freq_of_second_element)

該算法優先合併單個部分在詞彙表中頻率較低的 pair。例如，它不一定會合併 ("un", "##able") 即使它在詞彙表中出現的頻率很高，因為 "un" 和 "##able" 很可能頻繁地出現在其他詞中。

接下來我們以前面 BPE 舉的相同例子進行示範

1
2


# 訓練資料 (Corpus)
("hug", 10), ("pug", 5), ("pun", 12), ("bun", 4), ("hugs", 5)

一樣先把出現過的字母都加進詞彙中

1

Vocabulary = ["b", "h", "p", "##g", "##n", "##s", "##u"]

接著利用上面的公式計算每個 pair 的分數，並挑選分數最高的合併並加入 vocabulary，並合併 corpus

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12


# 訓練資料 (Corpus)
("h" "##u" "##g", 10), ("p" "##u" "##g", 5), ("p" "##u" "##n", 12), ("b" "##u" "##n", 4), ("h" "##u" "##g" "##s", 5)

# 每個 pair 的分數 (由高到低排序)
("##g", "##s") --> 5 次 / (20×5) = 1/20
("h", "##u") --> 15 次 / (15×36) = 1/36
("##u", "##g") --> 20 次 / (36×20) = 1/36
...

# 將分數最高的 "##gs" 加入 vocabulary，同時更新 corpus
Vocabulary = ["b", "h", "p", "##g", "##n", "##s", "##u", "##gs"]
Corpus: ("h" "##u" "##g", 10), ("p" "##u" "##g", 5), ("p" "##u" "##n", 12), ("b" "##u" "##n", 4), ("h" "##u" "##gs", 5)

接著就反覆重複這個動作：計算每個 pair 的分數 ➜ 挑選分數最高的合併並加入 vocabulary 同時合併 corpus …

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11


# 訓練資料 (Corpus)
("h" "##u" "##g", 10), ("p" "##u" "##g", 5), ("p" "##u" "##n", 12), ("b" "##u" "##n", 4), ("h" "##u" "##gs", 5)

# 每個 pair 的分數 (由高到低排序)
("h", "##u") --> 15 次 / (15×36) = 1/36
("##u", "##g") --> 20 次 / (36×20) = 1/36
...

# 將分數最高的 "hu" 加入 vocabulary，同時更新 corpus
Vocabulary = ["b", "h", "p", "##g", "##n", "##s", "##u", "##gs", "hu"]
Corpus: ("hu" "##g", 10), ("p" "##u" "##g", 5), ("p" "##u" "##n", 12), ("b" "##u" "##n", 4), ("hu" "##gs", 5)

直到 vocabulary size 達到我們設定的目標 (10) 就停止，這就是我們目前 Tokenizer 所有的詞彙了

1
2
3
4
5
6
7


# 下一輪 pair 的分數 (由高到低排序)，因此第三輪將加入 "hugs"
("hu", "##gs") --> 5 次 / (15×5) = 1/15
("p", "##u") --> 17 次 / (17×21) = 1/21
...

# 這是最終結果 (假設設定的目標為 vocabulary size = 10)
Vocabulary = ["b", "h", "p", "##g", "##n", "##s", "##u", "##gs", "hu", "hugs"]

B. 切割階段

現在來看看 WordPiece 是怎麼將句子切割成 subwords。前面提到的 BPE 在 tokenization 過程會根據學習到的合併規則去合併，但是 WordPiece 不是！

WordPiece 從句子開頭開始，找到 vocabulary 中與句子前綴相同且最長的 subword 然後拆分，反覆直到句尾。

一樣舉相同句子 "bugs"， vocabulary 中 b開頭只有它自己，所以得到 ["b", "##ugs"]。(如果 vocabulary 有"##bu"則得到的結果不同)

接著看"##ugs"，"##u" 是 vocabulary 中與 "##ugs" 前綴相同且最長的subword，所以拆分並得到 ["b", "##u", "##gs"] 。最後，"##gs" 在 vocabulary 中所以不拆，最後結果為 ["b", "##u", "##gs"] 。

Unigram

A. 訓練階段

跟 BPE 和 WordPiece 不同，Unigram 從一個較大的 vocabulary 開始，然後從中刪除 tokens，直到達到所需的 vocabulary size。

直接使用先前的範例

1
2


# 訓練資料 (Corpus)
("hug", 10), ("pug", 5), ("pun", 12), ("bun", 4), ("hugs", 5)

我們先列出所有的 substrings 並加入 vocabulary。為了簡化問題所以只保留 strict substrings

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11


# 為了簡化問題，只保留 strict substrings
# 排除以下 strings ("hug" 不排除因為它是 "hugs" 的 strict substring)
# 'pug', 'pun', 'ugs', 'hugs'

Vocabulary = ["h","u","g","hu","ug","p","pu","n","un","b","bu","s",
 "hug","gs","ugs"]

# 順便算出它們的出現頻率，出現頻率的總和是 210
# "p" 出現頻率是 17 所以出現機率是 17/210
("h", 15) ("u", 36) ("g", 20) ("hu", 15) ("ug", 20) ("p", 17) ("pu", 17) ("n", 16)
("un", 16) ("b", 4) ("bu", 4) ("s", 5) ("hug", 15) ("gs", 5) ("ugs", 5)

我們先來為 corpus 的每個 word 定義它的「分數」。我們會對 word 列舉所有的分割方式，並選出其中最高的機率乘積當作「分數」，以 "pug" 舉例

1
2
3
4
5
6
7


# "pug" 可以切割成三種方式，以下是每種方式的分數
# 「分數」為每個 subword 出現機率的乘積
["p", "u", "g"] : (17/210)×(36/210)×(20/210) = 0.001321 (score)
["p", "ug"] : (17/210)×(20/210) = 0.007710 (score)
["pu", "g"] : (17/210)×(20/210) = 0.007710 (score)

# 我們挑分數最高的

我們得到 "pug" 的「分數」是 0.007709，且我們會傾向將它拆成 ("p", "ug")或("pu", "g")而不是 ("p", "u", "g") 。這是每個 word 的分數與拆法：

1
2
3
4
5


"hug": ["hug"] (score 0.071428)
"pug": ["pu", "g"] (score 0.007710)
"pun": ["pu", "n"] (score 0.006168)
"bun": ["bu", "n"] (score 0.001451)
"hugs": ["hug", "s"] (score 0.001701)

回到目前例子

1
2


# 訓練資料 (Corpus)
("hug", 10), ("pug", 5), ("pun", 12), ("bun", 4), ("hugs", 5)

現在要來決定我們要從 vocabulary 中踢掉哪一個 token，我們會檢查刪除哪個 token 會讓所有 word 的 Loss (NLL, negative log likelihood) 增加最少

Loss Function (NLL) 為每個 word 的 出現頻率 * -log(該word分數) 的總和，值為 169.8

1

10 * (-log(0.071428)) + 5 * (-log(0.007710)) + 12 * (-log(0.006168)) + 4 * (-log(0.001451)) + 5 * (-log(0.001701)) = 169.8

假設我們踢掉 "pu" ，那 "pug" 仍然可以拆成 ("p", "ug")並維持相同的分數 (0.7710)。因此踢掉"pu"將給出完全相同的 Loss

但如果我們選擇踢掉"hug" ，那麼 Loss 會增加很多，因為 "hug" 和 "hugs" 的標記化會變成：

1
2


"hug": ["hu", "g"] (score 0.006802)
"hugs": ["hu", "gs"] (score 0.001701)

這些變化將導致 Loss 上升 23.5

1
2
3
4


# "hug" 增加的Loss
- 10 * (-log(0.071428)) + 10 * (-log(0.006802)) = 23.5

# "hugs" 則沒有增加Loss，因為 ["hug", "s"] 與 ["hu", "gs"] 的 score 都是 0.001701

因此我們選擇踢掉 "pu" 而不是 "hug"

Unigram 希望減少 vocabulary 中相似冗餘的 token，確保詞表裡的每一個 token 都是不可或缺的，不希望有「換了別人也能拼出來，且效果差不多」的 Token 存在（在本例"pu"與其它 token 功能太接近，不需要它）

接著就重複上面流程，反覆挑選 token 踢掉，直到 vocabulary size 達到我們設定的目標就可以停止啦～

B. 切割階段

Unigram 在切割句子中的 word 時，會使用 Viterbi 演算法，以 NLL 最低的方式去切割，有興趣的讀者可以再深入研究

以上就是 BPE, WordPiece, Unigram 三種演算法的介紹了，希望對大家有幫助～

參考資料

Hugging Face LLM Course — chapter 6 (the tokenizers)

【論文】PagedAttention — 高吞吐量LLM推論框架 vLLM 的設計

Tue, 09 Dec 2025 00:00:00 +0000

論文重點

提出了 PagedAttention 機制管理 KV Cache，大幅提升模型推論時的吞吐量

論文原文：Efficient Memory Management for Large Language Model Serving with PagedAttention

在大型語言模型（Large Language Models, LLM）爆炸式發展的這兩年，大家談論的焦點多半放在模型規模、能力與訓練資料上：參數越來越多、上下文長度越來越長、能做的任務也越來越多。然而，當這些模型真正走出論文、被放進產品裡，例如對話式助理、程式碼自動補全、雲端的文字生成 API，工程團隊很快會遇到另一個現實問題：推論時 GPU 記憶體的不足，變成推論（Inference）階段真正的瓶頸。

vLLM 是目前熱門的 AI 推論框架之一，專門針對高吞吐伺服器環境；最初由 UC Berkeley 開發，目前已成為 GitHub 上社群導向的開源專案，至截稿前有 64.9k 的星星與 11.8k 的 Fork，獲得高度關注與採用。

本文將會介紹 vLLM 創始團隊於 2023 發表的論文。其提出了受到作業系統虛擬記憶體啟發的設計：PagedAttention，解決 LLM 進行推論時，大量 KV Cache 於 GPU 記憶體（VRAM）碎片化而浪費的問題。在開始閱讀之前，建議先了解 Transformer 模型架構中的自注意力機制（self-attention）。

這邊簡單介紹一下 KV Cache 是什麼；在 Transformer Decoder 模型推論時會逐輪產生 token 與計算句子的 self-attention ，而 Attention 需要句子每個 token 的 Key & Value 向量；圖中上半部描繪了在沒有 Cache 的情況下， Transformer 計算整個句子 Attention 的過程

而 Key & Value 向量是由模型花時間 forward 計算得出的，我們可以把先前計算過的 Key & Value 向量（上圖紫色部分）存在記憶體中，只需要計算新 token 的 Key & Value 即可，這樣推論時就可以省下很多時間，是目前 LLM 推論階段關鍵的加速技術。

模型推論時用到的記憶體為了講求快速，都會選擇放在 GPU VRAM上（這邊不考慮 Offloading 到 DRAM/SSD，因為很慢）。

回到論文，來看看在 vLLM 出現之前，推論系統 Orca （不公開，由論文作者實現）在 KV Cache 利用的表現如何

作者分別實現了三種 Orca 版本，分別是 Max（總是分配模型最大句子長度的記憶體）、Pow2（分配output句子長度兩倍的記憶體）、 Oracle（分配output句子剛好的記憶體）。可以看到不管是哪種，KV Cache 都會出現碎片化的問題導致記憶體浪費。來看看 GPU 記憶體有多珍貴？

論文中以一個 40 GB VRAM 的 NVIDIA A100 GPU 為例子，從下圖左半邊可以看到，目前 LLM 推論時的記憶體用量，模型參數占了65%，KV Cache 占了 >30%，剩下為保留給 activation 的空間。模型參數與 activation 是固定大小，因此可優化的部分就剩 KV Cache 了。

從上圖右下方綠線可以看到，Batch size 越大，單位時間內能處理的 requests 就越多，Throughput 就越高。在 40 GB 記憶體的限制下，可以看到藍線（vLLM）相較於橘線（Orca/others），從 0.3k 提升到 0.9k，總共 3 倍的 Throughput 提升！由此可見避免 KV Cache 白白浪費記憶體有多重要。

因此作者用了作業系統管理虛擬記憶體的 Paging 技巧，將 KV Cache 分成 “Blocks” 來儲存，稱之為 “PagedAttention”。如下圖，每個 KV Block 包含了 4 個 Tokens （實際上預設是 16，論文 Ablation Study 章節指出 16 在大資料集與小資料集上都達到相對低的 latency），搭配 Block Table（映射表），一個 Logical KV Block 會對應到一個 Physical KV Block （如同虛擬記憶體以 Page 為單位對應到實體記憶體）

透過 PagedAttention 在分配 KV Cache 時就能以 Block 為單位分配，因為單位大小固定，可以簡單避免碎片化的問題，不再浪費 GPU 記憶體了！

從上圖還可以看到，如果不同 sequence 有著完全相同的 KV Block（左方 logical KV block #0 與右方 logical KV block #7），PagedAttention 可以映射到同一個 physical KV block #7，只需要存一份即可！

⚠️注意⚠️兩個 KV Blocks 前面整段 Tokens 要長度、內容完全一樣才算相同。

Why? 還記得 Transformer 論文提到的 Positional Encoding 嗎？Block 在句子中的絕對位置必須一樣；另外 KV Cache 存在於模型每一層，每層輸入是上一層經過 Attention 機制後的輸出。因此也會受前面 Tokens 影響。

P.S. 目前普遍的 Encoding 方式是 RoPE

在許多場景下，相同前綴 tokens 的多個 sequences 會經常出現。論文中用 Beam search 舉例，如下圖，在搜尋（width=4）時，4 個 sequences 的前綴經常有部分重疊，產生大量相同的 Logical KV Blocks，vLLM 只需要存一份實體的，可以大大減少 GPU 記憶體用量！

除了 PagedAttention 架構，論文也針對 Request Scheduling 與 Preemption 兩個策略去說明，基本上目前 LLM 推論框架對這兩項都有參數可以調整

Request Scheduling：多個 requests 同時正在排隊時，哪一個先做？

論文採用了 FCFS (First-Come, First-Served)，先抵達的先做

1

vLLM 參數為 "--scheduling-policy"

Preemption：當產生新 token 且發現 VRAM 不夠時，怎麼辦？

論文提出兩種方向：Swapping 與 Re-computation

Swapping (offloading) 就是把暫時用不到的 sequence 的 KV Cache 放到比 VRAM 慢的 DRAM 上，等 VRAM 有空間時再搬回來（當 kv block size 設定較小時表現較差，因為 CPU/GPU 頻繁進行少量數據傳輸，使 PCIE 頻寬無法有效利用；反之較好）
Re-computation 就是把暫時用不到的 sequence 的 KV Cache 扔了，之後再重新計算（因為一個完整句子的 KV 可以平行計算，在 GPU 算力無限的假設下，只需花費一次 prefilling 的時間）。

1
2


vLLM V0 兩個方式都能用，參數為 "--preemption-mode"
vLLM V1 後只支援 Re-computation

以上是論文針對 KV Cache 提出的管理方式。此時 KV Cache 只是在一個 request 內複用。

論文發表時並未實現跨 requests 的 KV Cache 共享。如果新一輪的 request 能夠直接復用上一輪計算好的 KV Cache（例如多輪對話），就能夠大幅降低新一輪 request 的 TTFT，俗稱 Prefix Caching。Prefix Caching 在 SGLang 推論框架論文中提出，透過 RadixAttention 來實現。

在 SGLang 論文發表後，vLLM 也透過 Hash RadixAttention 實現了 Prefix Caching（Issue #2614），官方稱作 Automatic Prefix Caching。每個 KV Blocks 會以目前+前面所有的 Tokens 經過 Hash 後，用 Hash Key 來檢索；Hash Key 一樣才是相同的 KV Blocks。有興趣的讀者可以看官網詳細的說明。

以上就是 vLLM 論文的重點介紹了～

透過論文，我們可以學到 vLLM 最重要的 PagedAttention 架構，它是如何改善 KV Cache 記憶體分配的，這是未來 LLM 推論框架勢必要面對的一個課題。

vLLM 從 2023 六月的第一版，到 2024 三月實現了 Prefix Caching，再到 2025 今年陸續把整個架構從 V0 重新翻新變成 V1，程式碼上經歷了許多改變，也多了超多其他 LLM 加速推論的技術，變化很快。

在 vLLM 2023 推出後同年底，有另一個開源推論框架 SGLang 也登場了，使用了 RadixAttention（類似 trie）去管理 KV Cache。這兩個應該算是現在最熱門的 LLM 推論框架了，對最新大模型的支援速度也很快。