【LLM 專欄】進階 context engineering:KV cache centric LLM 應用設計
如何藉由 prompt design 來優化系統效能、降低成本並提升用戶體驗
兩週前 Manus 的團隊發布一篇 Blog <Context Engineering for AI Agents: Lessons from Building Manus>,詳細描述了他們如何基於 KV cache 的特性來設計整個 Manus 的 Agentic system。
整篇都非常值得參考,從 KV cache centric 的角度談清楚了很多 LLM agents 的設計細節。
KV cache centric 的應用設計筆者認為其實是業界內一個很常見的進階技巧,設計的好的話可以節省運算、成本,同時還能提升用戶體驗,但因為對 senior LLM engineers 太直覺了,反而很少大規模公開討論,屬於容易被忽視的部分。
類似的技巧筆者在過去一兩年中,也已經實踐在許多內部專案的架構設計中。看到這篇文章後,不禁技癢難耐,也準備撰寫一篇文章來和大家分享一下心得與經驗 XD
也希望本文可以幫助大家更好理解 KV cache centric design 到底是甚麼意思,並且更好的理解 Manus 的 blog,推廣 KV cache centric design 的思考方法。
同時本文會提出 2 個常見但容易犯錯的設計做改良:
- System prompt 的內容先後順序安排。
- Chat History Management 的邏輯。
同時提出如果自有 GPU host LLM 能夠對 RAG 做的重大體驗改進,以及針對 workflow, agent 的優化設計。
也希望藉由上述這些實務場景,提出一個這三年筆者 LLM 工作,以及近十年的 DL 研究經驗帶來的一個最重要的觀念,Model-system-application co-design。
0. Preliminary: KV cache, prefix caching
如果最近有在追蹤筆者的幾篇新文章,應該不難預期我們又要從 KV cache 的概念開始談起,畢竟建立完整且正確的基礎觀念,是進一步理解後續內容的重要前提 XD。
.
1️⃣ KV cache
要理解 KV cache 之前必須先重新來看 Attention 的運作方式,在前文中我們已經非常非常詳細講解了整個 LLM 正向傳導(Forward pass)的過程,有興趣深入學習的可以再去複習一下。
這裡我們將 Attention 抽離出來,並為了簡化討論,假設 Multi-head attention 的 head 為 1,也就是最基本的 self attention。如此一來,我們的 Attention 計算流程便可歸納成以下三個步驟:
其中第二步 softmax self attention 就是整個 Transformer 的核心(畢竟原論文稱為的 Attention is all you need XD)。
為什麼說 Attention 是核心?因為 Attention 掌管了 Transformer 架構跟 context 的主要互動,而一個 sequence model 最重要的就是處理 context,如果能正確理解 context,就更有可能建立語言認知。
如果我們把 softmax self attention 這個步驟圖像化,則會如下圖。
其中 query q_t 會跟歷史所有的 key k_1, k_2, …, k_t 內積,找出最相關的 context,並且經過 softmax 後變成 attention score,再藉由這個 attention score 來 weighted sum 我們的所有歷史資訊的 Value v_1, v_2, …, v_t。
簡單來說,Attention 機制的本質就是判斷哪些 context 最重要,並依此進行加權總和。
.
假設我們今天把 inference 的時候的 t = 3 跟 t = 4 時刻的操作畫出來,那就會如下圖。
細心的朋友應該馬上就看出來了,在 t = 3 跟 t = 4 的這兩個時間點,其實有很多內容是共同的,key 的前三個 [k_1, k_2, k_3] 跟 value 的前三個 [v_1, v_2, v_3] 都一樣,也就是如下圖,因此我們其實不需要重新計算這些東西。
⭐ 這就是 KV cache!
當我們在計算第 t 個時刻的 Attention 時, [k_1, k_2, …, k_{t-1}] 跟 [v_1, v_2, …, v_{t-1}] 都已經在 t-1 輪已經被計算出來了,所以我們只要前一輪有把這些東西記起來,我們就不用重複計算他們。
.
反過來說,如果我們沒有存 t-1 輪以及以前的所有 Key 跟 Value(KV cache),那在算第 t 輪時,就要同時把 1~t-1 的 KV projection 也重做一次,如下圖,
而且筆者上述的計算其實也還是要存 h_t,不然我們甚至要從 input 開始把所有運算都重算一次,等於是重新 prefill,因此既然都要存了,就乾脆存 Key Value,可以連 KV projection 一起省掉。
.
總結來說,KV cache 的邏輯是以下三者。
- softmax attention 每一次需要使用到過去的所有 key 跟 value。
- 我們可以把每一輪計算出的 key, value 存下來(KV Cache),提供未來使用。
- 這樣每一步就只需要計算當下 token 的 qkv projection 即可,需要用的 context 資料都已經在 KV cache 中。(等於節省了對過去 token 的 KV cache)
事實上,KV cache 早在 2019 年 GPT-2 時就已經提出(若有更早的案例,歡迎指正)。值得一提的是,KV cache 的概念在 BERT 時代未被廣泛應用的主要原因,是因為 BERT 使用的是 Bidirectional attention,每個 token 不僅要使用過去還需要「未來」的 Key 和 Value,cache 的使用沒有單向依賴,因此難以直接套用 KV cache。而以 GPT 為代表的 Causal attention,正是 KV cache 的最佳應用場景。
.
2️⃣ 什麼是 Prefix Caching?
前面講了 KV Cache,接著我們就來提一提 Prefix Caching 的概念。KV Cache 教了我們在做 LLM 的正向傳導(Forward Pass)時,每一個 token 都可以沿用過去所有 tokens 的 Key Value,因此我們應該存出 KV cache,節省運算量跟時間。
實際上這不只是在 LLM decoding 的時候有這個性質,只要兩次 LLM completion 的 prompt 有「相同的前綴(Prefix)」,即可共用 KV cache。
這種場景其實在我們大多數 LLM 應用場景中非常常見,最常見的不外乎是以下兩種場景:
- 多輪對話,第 N 輪對話可以跟 N-1 輪共用 N-1 輪的 KV cache。
- 不同對話但使用相同的 system prompt,可以共用 system prompt 的 KV cache。
因此關鍵是有沒有相同的 prefix,只要有相同的 prefix 即可以共用 KV cache,省去計算這段 context 每一層 Key Value 的 computation。
⭐因此我們在設計 application 的時候,能最大化共用 Prefix 共用 KV cache 的設計,就是最 KV cache friendly 或是說 KV cache centric 的設計。
.
🌰 舉個簡單的例子,方便大家先對 KV cache centric design 建立一個簡單的認識:prompt 設計
假設我們利用 LLM 來解決某種任務(像是翻譯、檢查錯字、寫文章 ... 等),一般而言這個任務的 prompt 會包含幾種訊息:
- 這個助理的自我介紹、功能介紹(以防用戶詢問)
- 工作要求
- 限制
- User input
- Few shot examples
而其中 User input 是幾乎保證每一次都不一樣,因此我們就要盡量把 User input 放在整個 prompt 的最後面,才能營造出最長的共用 Prefix,進而最大化 KV cache 的使用。🔥
你是一個翻譯助理,請根據以下指引,將使用者提供的文字從來源語言準確、流暢地譯成目標語言,同時保留原文格式(如 Markdown、程式碼區塊與 LaTeX 公式)。
# 翻譯風格
1. 保持原文語氣與文體特徵。
2. 關鍵詞、專有名詞以**粗體**標示,數學公式使用 LaTeX。
3. 如遇術語,可在譯文後以括號補充原文(選擇性)。
# 系統限制
- 不可自行加入與原文無關的內容。
- 遇到不確定的詞彙或文化參考,請在括號中簡短說明。
# Few shot examples
example1
example2
example3
使用者輸入文字: {User_input}如果今天我們把 User input 提前放,則會造成較短的共用 Prefix,進而影響到我們可以共用的 KV cache。👿
如果我們的 few shot examples 是隨著 user input 而變動的話(Ex: Retrieval Few shot examples),那也需要盡量放在整個 prompt 的最後面。
⭐ KV centric prompt design:把不變、共通的內容放在前面,不同、會變動的內容放在後面,營造最長相同 Prefix。
.
對於自己 host LLM 的企業來說,更多的 prefix caching 代表能夠更節省 GPU computation,進而更高效使用 GPU,並且提升用戶體驗(更低的 TTFT time to first token)。
如果是使用第三方 LLM API(如 OpenAI, Anthropic, Gemini, …),則除了用戶體驗以外,主要則是省錢。💰
以 OpenAI 為例,主要模型如果使用到 prefix caching,則大多是只收取 1/2~1/4 的費用。(畢竟對他們幾乎是零代價)💰
而 Anthropic 更激進,如果你啟用 Prompt caching,第一次需要多付錢,但後續每一次使用到同樣的 prefix,都只需要 1/10 的價格。(這個 KV cache 只會預設存在 5 分鐘,只要 5 分鐘內有 refresh 就可以持續使用)💰
因此可以看到,如果應用設計的好,最大幅度去使用 Prefix Caching,我們是可以常態達到 50%~90% 的成本節省。🔥
而這可能是對「套皮 GPT/Claude 的企業」最重要的一個生存問題:如何用更低成本達成一樣的事情,因此 Manus 才會特別寫一篇文章來討論相關技術。
.
3️⃣ Prefix Caching 系統設計
接下來我們稍微延伸一點,來討論一下如果今天你是 LLM inference service 的負責人,應該怎麼樣設計 inference 系統讓 Prefix Caching 發生機率最大化。
在筆者前一篇深入探討 LLM inference efficiency 的文章中,筆者有提到假設在單張 H100 上 host LLaMA 3.1 8b (bf16),理論上我們可以同時存放最多 488k tokens 的 KV cache。
也就是說對於系統設計而言,我們問題就變成「如何最有效利用這 488k tokens KV cache,讓最多 request 可以使用到 KV cache?」
.
這邊撇除掉比較細節的 data structure,筆者想特別提兩個值得關注設計:
- KV cache eviction:當我們 GPU memory 滿時,要丟掉(或是 offload)哪些 KV cache。
- KV cache aware scheduling:如何基於 KV cache 來 schedule 我們的 concurrent requests。(如何安排 KV cache 的使用)
.
首先來看 KV cache eviction,其實就跟所有 cache system 一樣,當 GPU Memory 滿的時候,我們要讓未來最有機會被使用的 KV cache 留下來,而把使用機率最低的 KV cache 丟掉(或是 offload 到 CPU Memory)。
在各種 cache system 最常見的 eviction 機制不外乎就是兩個:
- 淘汰使用率最低的(LFU:Least Frequently Used)
- 淘汰最久未用的(LRU:Least Recently Used)
而大名鼎鼎的 Radix Attention [1] 的主要就是採取 LRU 的 KV cache eviction 策略,也被使用在 vllm 跟 sglang 中。
.
我們來稍微細看 vllm 的 KV cache eviction 具體會發生甚麼事。參考 vllm document 以及 source code,可以知道 vllm 的 KV cache eviction 會遵循以下邏輯:
當需要 allocate 新的 KV cache,但 free list 已空 時,會依序套用下列規則:
- 只從
ref_count == 0的 blocks 挑選可淘汰的候選 KV block(保證不會清掉正在被使用的 KV cache)。 - 若候選很多,採取 LRU 策略淘汰(
last_access_ts最舊者先 evict)。 - 若還有 tie,優先淘汰位於最長前綴「尾端」的 block(可視為「前面堆了最多個 blocks 的尾巴」)。這個選擇能最大化保留公共前綴,不至於從中段挖洞。
其中前兩部都是常見的 LRU 策略,第三步除了能優先保留最長前綴以外,另一個主要的理由其實是 LLM 的使用特性,LLM 使用中容易出現相同公共前綴,然後多種不同回答的應用場景,像是 self-consistency,使用同樣的前綴 decode 5 種答案,再做 Majority Vote。
.
接著來簡單看 KV cache aware scheduling,假設我們有 20 個 requests 在我們的 queue 裡面,我們其實可以採取藉由對這個 queue 重新排隊,來最高效使用 KV cache。
正常來說我們可以做以下兩個策略:
- 優先發送 KV cache 已經在 GPU Memory 的 requests。
- 把相同 prefix 的 requests 排在一起發送。
這兩個策略都能減少我們 offload/reload KV cache 的次數,避免 cache thrashing 造成額外 latency。
.
從 KV cache eviction 跟 KV cache aware scheduling 這兩點應該可以發現,系統設計者也可以基於 LLM 的特性,來設計更好的系統邏輯,來優化整體系統指標。
1. chat history management
接下來就是這篇文章的重點,基於 prefix caching 的觀點我們要怎麼樣設計整個 application 的各個方面。
前面我們以 prompt 做過一次說明,在設計 prompt 的時候把共通、不會變動的內容放在前面,而把不同、會變動的內容放在後面,就可以顯著更高效的使用 prefix caching,大幅節省運算與成本,並且提供更好的用戶體驗(更低的 TTFT)。
Radix Attention [1] 的論文就體現出使用 prefix caching 帶來顯著更低的 Latency 跟更好的 throughput。
接下來我們來討論另一個看似簡單但其實不容易的話題:如何優化 chat history 的 KV cache 使用。
有開發過 LLM chatbot 的朋友應該都知道,chat history 包含了所有的對話紀錄,通常全部送給 LLM 可以讓用戶得到最好的體驗(LLM 可以理解更多上下文,進而更好回答下一輪),一般我們送入 LLM api 的 chat history 會長的如下圖。
而因為前面說的 KV cache 的特性,雖然整體 chat history 很長,但我們真正需要計算的其實只有最新一輪的 User input,前面的對話輪次全部都已經計算過了 KV cache,因此運算代價其實比想像中低。
看似很美好,但 chat 的場景卻有一個無可避免的缺陷:chat 可以無限延長,但是 chat history 卻有一個硬性限制,不能超過 llm context length。
而目前最主流的 chat history management 方法,是當我們的 chat history 即將超過 context length 時我們正常的做法是把最早以前的對話紀錄丟掉。
這種方法在目前超級常見,就筆者所知,主流的開源 LLM 框架、軟體專案、甚至台灣主要企業內的 LLM chat history management 都是這樣設計的。
但這種 chat history management 的方案卻有一個顯著的缺點,就是當我們的 chat history 只要一超過 context length,就會觸發 truncate 把最早的對話清理掉,而導致我們整串 chat history 的 prefix 性質從最早的地方被破壞掉,而無法繼續使用 KV cache。
而在當下,很多 LLM 的 context length 動輒就是 32k 甚至 128k,假設我們 system prompt 只寫 4k,留 4k 給 output,那 chat history 長度就大約是 24k 甚至是 120k。原本開開心心使用 124k 的 KV cache(system prompt + chat history),但出現一次 chat history truncate 後,直接 120k context 全部都要重算。
而且更慘的是接下來幾乎每一輪對話,因為 context length 已經很滿了,因此幾乎都要 truncate 並且重新計算一次整個 chat history 的 KV cache造成大規模的運算跟成本浪費。👿👿👿
.
⭐ 解決方案其實很簡單,就是我們在 truncate chat history 的時候採取更激進 truncate 策略,把最早的 50% chat history 都丟掉。用一次的對整體 chat history 的重新計算,換取後面更多輪次不用計算。
假設我們在 H100 80GB 上 host llama3.1–8b,現在 concurrent user 只有一個用戶,並且他連續多輪對話每一輪都是 4k input 然後 LLM response 4k output,那我們採取兩種策略的 theoretical TTFT 會顯現出巨大的差異。
熟悉資料結構的朋友應該也很快就發現,這跟我們以前大學學的 dynamic table 是同一個邏輯,算法設計上非常簡單,只是我們基於對 LLM inference 的知識,來反哺應用上的設計。
我們目前只是基於更改 truncate chat history 的邏輯,就得到了大幅度的 TTFT 改善,節省 computation 也節省成本。🔥
.
⁉️ 一定會有朋友現在心中有一個問題:「一次丟一半的 chat history 難道不會影響回答品質很大嗎?」
實際上在 32k 甚至 128k 的對話場景下,大部分情況不會。
不過如果我們真的擔心,我們也可以在每一次 truncate 的時候對被丟棄的 chat history 做 summarize,把過去重要的對話資訊壓成 memory。
不過這樣就會帶來一個問題,就是 summarize 本身可能要生成 2k~4k tokens,以 LLaMA-3.1 8b 而言,因為要壓縮 64k 的 context,所以相當於一個 context length = 64k 的 decoding 問題, TPS 大概就 150 tokens/s ,因此大約也要花 15~30s 才能做完一次 summarize。
而這些時間都會被加在用戶的 TTFT 裡面,也就是說如果在原有 truncate half chat history 的邏輯上加入 summarize 的操作,則會讓那偶發的幾次 truncate + summarize + recomputation 的 TTFT 變得無法忍受。
.
這就提到 KV cache centric chat history management 的第二個重點,chat history summary/truncate 的觸發邏輯並不應該基於 user requests,而應該盡量找機會在 GPU resource idle 的時候偷偷做掉,或是盡量 overlay 掉。
我們先來看 summarize 這步,首先我們要知道 summarize 的發生時機其實非常多,並不是只有到要撞到 context length 的時候才可以做 summarize,因為我們是 summarize 前半的 chat history,因此只要 chat history 夠長就隨時都可以同步進行 summarize,並把 summarize 的結果先存起來,等到要用到的時候再取出。
同時我們還可以在某一次回復 user input 的時候,同時把 summarize 做完,只要把 summary 的 prompt 放在最後(當成另一種 user input),兩個任務幾乎可以共用所有的 KV cache,主要的時間花在 decoding。
而我們又知道 decoding 主要是 Memory bandwidth bound,上面兩個操作既使用相同模型、又共用大部分 KV cache,所以連需要搬運的內容都高度相同,幾乎不會有多餘代價。
⭐ 因此我們 summarize chat history 這步可以很簡單的幾乎零成本在 background 做完,而不會讓用戶感覺到任何 Latency。
.
接下來看 truncate + recomputation 這步,其實如果理解 summarize 可以在背景做完,truncate + recomputation 也幾乎可以套用相同邏輯。
我們可以設定適當的 threshold,例如當我們的 chat history 已經超過 80% 的 context length 時,不等 User input 進來,我們就在 backend 提前把 truncate 跟 recomputation 做完。
而因為 truncate + recomputation 等於要重新 prefill 半個 chat history,是嚴格的 computation bound,因此更好的作法是要在確認 GPU 有閒置的 computation resource 的時候就偷偷做完。
.
細心的朋友可能馬上也會想到,這些技巧也可用在 truncate 1 chat message 的場景,也就是原本常見的 chat history management 方法。
但因為 truncate 1 chat message 所需要的 recomputation 太頻繁加上每一次的運算量又大,所以更難遮蓋掉,因此仍然不是好辦法。即便遮蓋掉,花費的計算資源跟成本也降不下去,因此筆者還是在大多數 chatbot 應用中採取 truncate half chat history 的做法。
而這就是了解 LLM inference 後才能想清楚的邏輯,也就是利用 inference 的特性回頭來大幅度優化應用。🔥🔥🔥
上述做法乍看之下需要自有 GPU host LLM 才可以操作,實際上在有支援 prompt caching 的 API 上也都可以做相似/相同的設計,並且獲得一定的收益。只是第三方 API 的收益通常要更精準去設計,因為自由度較低。
2. KV cache center design for RAG, agent
接下來我們來更進一步,如果我們自有 GPU host LLM,能做什麼 KV cache centric design 來優化我們的 RAG/agent system ,降低計算量/成本或是提供更好的用戶體驗。
1️⃣ Reuse KV cache in agentic RAG system
2025 年我想大部分的 RAG system 應該都是某種 workflow based 或是 agentic based RAG system,而這種 multi-step LLM system 的最大好處就是有很多機會可以共用 KV cache。
只要我們 prompt/workflow 設計的好,很多步驟可以有效的 share KV cache,來降低 multi-step 帶來的 latency。
因為不太能公開之前工作開發的 RAG system,因此我們來看 Langchain 的 Agentic RAG pipeline。
其中我們來看 Check Relevance 的 node,一般來說 check relevance 做的事情就是讓 LLM 判斷某些 retrieval chunks 是否相關,把 irrelevant chunks 丟掉,避免太多 noisy context 影響到最後生成品質。
而 Check Relevance node 就有兩種常見的 share KV cache 方法。
第一、不同 chunks 可以直接平行 decode
以往大家可能習慣把所有 chunks 並在一起,一次性讓 LLM 對所有 chunks 判斷是否有關,希望藉此來節省 system prompt 的 token cost,畢竟 system prompt。最知名的論點就是 Stanford 的 FrugalGPT [2]。
但這裡筆者要 argue 相反的觀點,如果我們自有 GPU host LLM 的話,全部分開來平行做判斷才是顯著更好的方法。(如下圖)
因為不同 request 中重複的地方全部都可以藉由 share KV cache 來進行,所以我們並不會引入多的計算需求,同時每一個 request 的 context 更短,所以不論是 prefilling 還是 decoding 的 attention 運算總量都是更小的。
而且原本假設要 output 1k tokens 解決的問題,現在可能被拆成 20 組 50 tokens 的問題,固定 TPS 底下解決效率會大幅度提升。
因為我們熟知 Decoding 是 Memory Bandwidth Bound,我們知道拉大 Request 的 Batch size 還可以讓整體 arithmetic intensity 更往右靠,進而更好利用 GPU 資源。
再加上分開判斷對於 LLM 也是更容易的任務。
⭐ 因此對於 Check Relevance node,只要自有 GPU host LLM,分開處理的不論是效率、效果、成本都是顯著更低的。
.
第二、Check Relevance Node 跟後續 Generate Node 可以共用 KV cache
不過把所有 chunks 放在一個 prompt 做 check relevance 有另一個隱藏的好處,就是很大概率這組 KV cache 可以被 Generate Node 無縫使用。
也就是說我們犧牲 Check Relevance Node 的效率,來換 Generate Node 可以直接 Reuse 前面 retrieval context 的 KV cache。
⁉️ 不過前面我們不是說 Check Relevance Node 的目的是篩選掉不相關的 Context 嗎?如果 Generate Node 直接使用相同的 KV cache,這些 noisy context 不就跟著進去了嗎?
沒錯這是 reuse KV cache 最重要考量的一個點,因此我們需要對 Generate Node 做對應的設計:
最基本的做法是在 Generate Node Prompt 裡面把 Check Relevance 的結論加進去,也就是說我們沒有顯式移除掉不相關 context,而是用補充說明的方式希望 LLM 理解(Ex: 告訴他 context 1, 2, 5, 9 無關)。
但我個人更喜歡進階一點的作法,我們在 calculate attention score 的時候去修改我們的 Causal Mask,把被判定不相關的 chunks 的 tokens 也跟著 mask 掉,如下圖。
這樣我們在最大幅度 share KV cache 的前提下,一樣讓 Generate Node 盡量忽視掉不相關的 Chunks。(Attention score 不會去 weight 這些 chunks)
不過這個做法其實並不完全等價於我們一開始就把 context 拿掉,因為其他 Chunks 還是有可能偷偷帶著一些被丟掉 chunks 的資訊,因此還是要實驗確定在自己的場景是否穩定有效。
實際上比起修改 Causal Mask ,我們還可以再更進階直接物理 Drop Tokens,真正把這些 tokens 的 KV cache 丟掉才能真正省 computation,但是這在大多框架中需要更精緻去修改底層 attention 的計算,修改成本比較高 XD
.
這裡我們隨便拿一個 Node 出來討論就有好幾種不一樣的設計方式,而這些 KV centric 的設計方法能有效優化我們的 KV cache 使用率,藉此節省運算、降低成本並減低 TTFT。
不過敏感的朋友應該發現了,不同 node 之間 share retrieval context 跟不同 requests 之間 share system prompt 這兩種邏輯,其實是衝突的。
- 如果我們把 retrieval context 放在 prompt 的前面 ==> 不同 node 可以共享 retrieval context 的 KV cache,但不同 requests 進來則無法 share。
- 如果把 node system prompt 放在前面 ==> 不同 requests 可以共享 system prompt 的 KV cache,但不同 workflow node 則無法。
也就是呈現一個有一好沒兩好的情況,如下圖。
這其實就是真正考驗 application design 的地方,要回答現在需要採取哪一種設計需要全面的思考以下幾種問題:
- 我們應用的 SLO 是什麼?
- 我們比較常遇到高並發事件,同時大量 requests 一起進來,還是比較常是低並發的環境,大多時候都只有 1 個 requests?
- 我們是專用 LLM inference server 還是通用 LLM inference server?
- 我們的 KV cache 平均會在 GPU Memory 中存活多久?
- 我們是 retrieval context 比較長還是 node system prompt 比較長?哪一個 recompute 的代價比較高。
- ...
當然我們也需要採取更進階的 prompt 設計來盡量結合兩者優點,舉例而言我們可以再把 workflow 的 prompt 進行拆解,把共通的需求寫成 share system prompt 並放在最前面,並把不一樣的需求寫在最後面。
.
2️⃣ RAG with document KV cache
前面我們都還是在調整 prompt 的位置、調整 request 打法,來盡量 reuse KV cache,是較通用的方法,但卻不見得真正利用到 RAG 的特性。
針對 RAG 其實我們還有一個更大的夢想,因為 RAG 可能被 retrieved 的 chunks 是有限的,能不能把整個 vector 所有 context 的 KV cache 都先預先算出來,這樣 retrieval 到後直接 reuse 現成的 KV cache,大幅節省 prefilling 的時間。
但現實沒有這麼美好,現在的 LLM 都是採取 prefix caching,只有相同 prefix 才可以 reuse KV cache。
也就是說如果兩個 query 分別 retrieve 到 [1, 3, 4, 5] 跟 [1, 2, 4, 6] 四個 chunks,我們最多就 reuse [1, 4] 這兩個 chunks 的 KV cache,還需要把這兩個 chunks 放在 context 的最前面。
而如果加入排列順序,chunks 由相似度高排到低,則 reuse 的機率還會顯著降低,同樣 retrieve 5 個 chunks,如果排列順序不同,也有效無法 reuse。
假設 db 裡面總共 1,000 個 chunks,每次 retrieve 5 個出來照相似度排列,我們總共就有 990 Trillions 的排列可能,不可能也不值得做出所有的 KV cache。
實際上這個問題在 2020 年就已經被廣泛討論 [3, 4, 5, 6],大部分結論是加以訓練可以得到一定的成效,但當時的 T5 或是 BERT based results 跟現在 LLM 是天差地遠。
.
而在 LLM RAG 時代,如何讓 RAG system 最高效使用 context cache,則主要發展成兩脈:
第一、設計系統來最大化 Prefix Caching 命中率的 RAGCache [7],每次新的 requests 都先檢查有沒有過去某一次用到相同的 chunk prefix,如果有的話就盡量 reload KV cache 回來,如果沒有則建立新的 KV cache。
雖然整個 DB 的 prefix 排列是天文數字,但大多系統短時間內會出現某種規律 Pattern(Ex: 事件觸發大家都在問相似的問題、同一個人一直續問),因此只要把 prefix caching 發揮好也能得到有效加速,以 RAGCache 而言可以得到平均 4x 加速 TTFT,以及 2.1x throughput。
RAGCache 的主要邏輯是,有 prefix caching 總比沒有好,即便這些 KV cache 要從其他地方 reload 進來。因此怎麼樣設計一個完整的 cache system,讓最有機會被用到的 chunks prefix 對應的 KV cache 留在最近的地方,變成主要的設計原則。
.
第二、研究讓 LLM 直接使用預先算出的 chunks KV cache 的方法
如果我們從更底層來重新思考問題,假設我們先對每一個 chunk 都先獨立用 LLM 算一次 KV cache,在 RAG inference 的時候,被 retrieved 到 chunk 直接把這些 KV cache 拿去給 LLM 用到底會有甚麼問題?(如下圖)
主要問題有兩個:
- Position Misalignment:這些 chunks 裡面每一個 token 在 precompute 時候的 position 跟真實使用時的 position 大概不一樣。
- Lack of cross attention:不同 chunks 之間從來沒有互相被 attention 過,缺乏 chunks 之間的聯繫。
而這兩個問題又進一步導致算出來的 KV 以及 Attention 結果跟真正一次性 prefill 算出來的不一樣,進而導致 performance drop。
Position 的問題還稍微好解決一點,Gim In, et al. [8] 發現 LLM 的 position 不連續也可以有好表現,也就是說我們只要維持幾種主要的 position 區間,然後把我們 inference 的 context 重組變成對的 position 即可。
Cross Attention 則較難解決,因為這就是我們 precomputed KV cache 裡面真實缺乏的資訊,重算這些 cross attention 幾乎就等於整個都重算了,也不值得。而解決這件事的邏輯主要就聚焦在「能不能用少量的 recomputation 來重新對齊整個 attention score?」
CacheBlend [9] 首先發現了如果我們無視 Cross attention 的問題,我們 chunks 越多反而會效果越差,而直接整個重新 prefill 則通常越多 chunks 效果還可以提升,顯現出明顯的效果差異。
接著 CacheBlend [9] 提出了我們可能只需要重新計算少部分 tokens 的 cross attention 就可以解決這個 performance drop,關鍵變成怎麼樣挑選這些需要被重新計算的 tokens。
並且提出只要看第一層的 KV cache 差異就可以精準篩選出這些要被重新計算的 tokens。等於是我們只需要 prefill 一層就可以做挑選,然後挑選出一定比例的 tokens 重算所有 cross attention,就可以得到非常好的效果,有效減低運算量。
而 CacheBlend 在大量場景下都可以接近 Full KV recompute 的效果,但只需要小量的運算。(Prefix caching 因為比較嚴格,要相同 prefix 才能使用,因此通常也需要大量的運算量來處理 prompt)
而 Epic [10] 則延伸 CacheBlend 的 idea,不滿足於用第一層的 KV 差異來挑選 tokens recompute,直接對每一個 chunks 的前幾個 tokens 重新 compute 就可以得到非常好的效果。
在大量模型、問題下 Epic(LegoLink)都可以接近全部重新計算(FR: Full Recompute)的效果,並且接近 Naive(Full reuse)的計算量。
相似的 idea 還有 KV link [11],同時採用了 Prompt cache 的 position 修改方法,對齊 cache 的 position,也採用類似 Epic 的方式,主要運算新增的 link tokens(其實非常類似 Epic 的固定 recompute 某幾個 tokens)。
.
⭐ CacheBlend, Epic 跟 KV link 都展現了我們可能可以藉由少量的 recomputation 來達到跟全部重新 prefill 相似的效果,因此 RAG with document KV cache 這個夢想是可以實現的,如果我們系統 TTFT 今天是關鍵的 bottleneck,這一系列是非常值得參考的。
如果對 attention sink 熟悉的朋友應該看的出來 Epic 其實就是針對 attention sink 在設計,只要修正每一個 chunks 的 sink tokens 就可以得到好效果,但這個議題筆者暫時不打算深入,之後有機會可以討論 XD 筆者在極端的 computed limited 的機器上嘗試過 Epic 跟 KVlink,品質還不錯 XD
.
3️⃣ Workflow aware KV cache scheduling
接下來談第三個部分,基於 Workflow 的特性我們其實也可以更好 schedule KV cache 的 reload/offload 或是 evict KV cache。
這應該很直觀,假設我們 workflow 就是固定的 5 步驟,在進行第一步驟的 decoding 的時候,就可以先行把 2 3 步驟的 KV cache reload 回來,也就是 OS 熟悉的 Prefetching,而 KVflow [12] 則針對 Workflow 特性提出整套 KV cache scheduling 的 solution。
前文我們提到目前 KV cache eviction 的邏輯普遍都是 LRU eviction,而這種策略其實對於 workflow 是相對不友善的,因為大部分時候 workflow 都是 loop/sequence of nodes,最近使用的 node 剛好是接下來一大段時間最不可能再使用的 node。(如下圖)
同理,對於相較固定的 workflow 以及大部分的 Multi Agent system,我們大多都可以設計更有效的 prefetching 方法,在 node 發生前就 prefetch 回對應的 KV cache,把 KV cache transmission 的時間全部 overlap 掉。
.
看到這裡大家應該也有感覺了,如果我們自有 GPU 來 host LLM,我們能玩的花樣就多了,一個懂得 KV cache 跟 LLM inference 的 LLM 工程師,光是重新設計 LLM pipeline,就可能可以讓企業內的 GPU 利用效率大幅提升,節省大量成本並且提升用戶體驗。
結論
本文受到 Manus 的 Blog 所啟發,比較全面的講解了 KV cache centric design 的關鍵。
KV cache centric design 的好處很明顯,最常講的有以下三個:
- 成本更低
- 運算量更少
- 用戶體驗更好(TTFT 顯著降低)
對於 LLM 應用/套皮的公司,這幾乎是服務與企業的命脈。
而要最大化 KV cache 的使用,則需要盡量營造出相同的 prefix,因此我們從 prompt 到 workflow 都可以做很多相對應的設計。
如果我們是自己 GPU host LLM,則可以更細緻的去設計 KV cache 如何在我們的 GPU 中存活,設計更適合我們應用的 eviction、scheduling 邏輯,甚至像是 RAG 有機會直接使用 precomuted 的 chunks KV cache,大幅節省時間。
這些設計方案的本質都一樣,就是從系統分析(prefilling / decoding 的資源分析、KV cache 分析)來反哺我們的應用設計,很多方案都很簡單,但如果缺乏完整對 LLM inference 系統的認識,就無法找到這些簡單有效的方案。
因此筆者還是很推薦 LLM application engineer 可以多去讀 LLM system 相關的論文,用更全局的視角思考問題。
Reference
- Zheng, Lianmin, et al. “Sglang: Efficient execution of structured language model programs.” Advances in neural information processing systems 37 (2024): 62557–62583.
- Chen, Lingjiao, Matei Zaharia, and James Zou. “Frugalgpt: How to use large language models while reducing cost and improving performance.” arXiv preprint arXiv:2305.05176 (2023).
- Burtsev, Mikhail S., et al. “Memory transformer.” arXiv preprint arXiv:2006.11527 (2020).
- De Jong, Michiel, et al. “Mention memory: incorporating textual knowledge into transformers through entity mention attention.” arXiv preprint arXiv:2110.06176 (2021).
- Févry, Thibault, et al. “Entities as experts: Sparse memory access with entity supervision.” arXiv preprint arXiv:2004.07202 (2020).
- Chen, Wenhu, et al. “Augmenting pre-trained language models with qa-memory for open-domain question answering.” arXiv preprint arXiv:2204.04581 (2022).
- Jin, Chao, et al. “Ragcache: Efficient knowledge caching for retrieval-augmented generation.” arXiv preprint arXiv:2404.12457 (2024).
- Gim, In, et al. “Prompt cache: Modular attention reuse for low-latency inference.” Proceedings of Machine Learning and Systems 6 (2024): 325–338.
- Yao, Jiayi, et al. “CacheBlend: Fast large language model serving for RAG with cached knowledge fusion.” Proceedings of the Twentieth European Conference on Computer Systems. 2025.
- Hu, Junhao, et al. “EPIC: Efficient Position-Independent Caching for Serving Large Language Models.” arXiv preprint arXiv:2410.15332 (2024).
- Yang, Jingbo, et al. “Kvlink: Accelerating large language models via efficient kv cache reuse.” arXiv preprint arXiv:2502.16002 (2025).
- Pan, Zaifeng, et al. “KVFlow: Efficient Prefix Caching for Accelerating LLM-Based Multi-Agent Workflows.” arXiv preprint arXiv:2507.07400 (2025).
