淺談 Prompt cache
什麼是 Prompt Cache ? 你每天在使用的 llm 如何 cache ? 如果失敗又會發生什麼事情?
Contents +
淺談 Prompt cache
你每天在等聊天機器人回話的那一兩秒,其實是在看一台「腦袋」努力翻頁。Prompt Cache 做的事,很像是幫這顆腦,把常翻的那幾頁預先摺角、貼標籤,甚至直接影印好放桌上,少翻幾次,回話就快很多。
多數人用 LLM 的方式,比自己以為的還「重複」。系統提示一樣、工具說明一樣、法條說明一樣、產品規格一樣,只換最後那幾行問題。人腦會習慣這些背景設定,機器卻每次都要從頭讀。Prompt Cache 把這些重複段落拆成一塊塊模組,先算完裡面的注意力狀態,放進一個「記憶櫃」。下次再遇到同一段,就直接把結果拿出來接上去,省掉整段重算。對使用者來說,差別就是:第一個字跑出來的時間,大幅縮短。
先把故事講白:現在 LLM 已經有一個叫 KV Cache 的東西,會記住自己剛剛算過的中間結果,下一個 token 就不用再從頭來過。Prompt Cache 多做一步,讓這個記憶不只在一條對話裡重複使用,而是可以在很多不同請求之間「模組化共享」。像是在伺服器旁邊放一櫃常用零件,組裝新機器時直接拿,不用每次重鑄。
讀到這裡,你可以先帶走一個判斷:如果你的應用很多「模板式長 prompt」,Prompt Cache 這種思路對你非常有利,其中最有名的大概是近期 claude code 因為各種 prompt cache 問題導致用戶使用量被亂刪減,所以到底什麼是cache為什麼會這樣 ? 我覺得可以寫一系列文章來導讀整理,這邊先用這篇論文帶大家看過 https://arxiv.org/abs/2311.04934 低延遲 LLM 推論:Prompt Cache 與模組化注意力重用
一個長 prompt 丟進模型前,那幾百毫秒到幾秒的空白,其實都耗在「重讀背景」。系統提示、模板、說明文件,對人來說只是理所當然的開場,對模型來說卻是每次都要重新掃過的負擔。Prompt Cache 這篇論文做的事,就是承認這些東西一直重複出現,把它們切成模組,提前算完注意力狀態,之後直接用結果,不再每次重來。
這篇文章主要根據論文本身的內容來寫。你另外開的參考網站目前頁面是載入訊息,沒有多出新的技術細節,所以實際可用資訊還是來自 arXiv 的全文。下面我會用比較工程實作的角度,把 paper 的想法重新整理成一篇可以當技術說明頁的長文,配合分層標題,方便之後掛在有 TOC 的網站上。
1. 從 KV Cache 到 Prompt Cache :問題是「重複」沒有被好好用
現在線上 LLM 幾乎都會用 KV Cache。概念很直接。模型在自回歸生成時,每出一個 token,本來都要重新計算整段輸入的 self-attention。有了 KV Cache 之後,前面 token 的 key / value 只算一次,後面每步只補新 token 的部分,算力從「每步看全部歷史」變成「每步只處理新增的那一格」。
這已經把生成階段的計算壓下來不少。不過 prompt 那一段背景文字,還是每個請求都重算一次。對很多實際應用來說,這塊才是真正大的成本。想一下你日常的長 prompt 場景。系統訊息同一份,工具說明同一份,RAG 撈回來的文件來來去去就那幾篇。真正每次不同的,只有使用者最後那一小段指令。
論文的觀察很樸素。既然有這麼多重複片段,為什麼只在一條對話裡重用中間狀態,不乾脆跨請求直接重用「某一段文字整條算完的注意力狀態」?這樣新請求來時,如果用到同樣那段文字,就不用再把文字餵進模型,只要把那段對應的 KV tensor 拿出來,像積木一樣接回去。
Prompt Cache 就是沿著這個思路往下做。它的目標很單純,減少 time-to-first-token,也就是第一個 token 出現前那段死寂時間。後面逐 token 的產生過程,還是交給原本的 KV Cache。
2. Prompt Cache 做了哪幾件事
從系統角度看,Prompt Cache 多出來的工作有三件。
第一,承認 prompt 裡有「可重用的模組」。把 system prompt、常用模板、RAG 文件、個人化描述這些重複段落拆出來,定義成獨立單元。論文叫這些東西 prompt module。
第二,定義一套標記語言,讓這些模組在文字層級變得顯式,而不是寫死在 string 裡。這就是 Prompt Markup Language。它有點像在 prompt 裡寫 HTML,用 <module>、<param> 這種 tag 把結構標出來,方便之後在服務端對應到一塊一塊的 KV cache。
第三,處理位置編號的問題。注意力狀態跟 token 的 position id 綁在一起。直接把一段 KV 結果搬到別的地方,理論上會破壞語意。作者做的事情,是把 schema 的整體布局固定住,給每個模組一段不重疊的位置區間,讓「同一個模組在不同 prompt 裡出現時,位置意義保持一致」。同時實驗發現,只要模組內的相對順序不變,位置編號之間留空不影響輸出品質。
這三件事疊在一起,Prompt Cache 得到一個能力。當有新請求進來,它先看 prompt 說自己是從哪一個 schema 衍生、用了哪些模組。這些模組之前已經在背景算過 KV 狀態,存進 cache。系統就直接把這幾塊 KV tensor 接起來,再針對這次新增或帶參數的文字部分去跑一次注意力。整段 prompt 的 prefill 等於被拆成「記憶拷貝」加上「少量新計算」。
3. Prompt Markup Language :把 prompt 寫成可被 cache 的結構
論文花不少篇幅在介紹 PML,因為這是整個系統的入口介面。對使用者來說,PML 有幾個核心概念。
第一個概念是 schema 跟 prompt 的區別。schema 是一份「模板描述檔」,用來定義有哪些 prompt module、它們在整體中的順序、層級關係。實際送進模型的一次請求,則是在 <prompt> 裡,宣告自己是基於哪個 schema,選擇要 import 哪些模組,再加上一些非 cache 的額外說明。
這樣做有一個直接效果。系統只需要對 schema 裡的模組做一次編碼,把每個模組對應的 KV 狀態算好存起來。之後所有從這個 schema 衍生的 prompt,只要有引用模組,就能拿現成的。prompt 裡新寫的文字,只算那一小段。
第二個概念是參數。很多模板彼此之間只有少數幾行不同。PML 用 <param> 這個 tag,在模組裡留了一塊長度受限的空間,後面 prompt 真的引用這個模組時,才能填入實際文字。編碼時,用 <unk> 之類的 placeholder 先佔位置,把那一段的位置 id 記錄下來。執行時,替換成真正的 token,再把對應位置上的注意力狀態更新。
這樣一來,一個很大的模組,裡面只要少數幾個位置會變。大部分 KV 結果可以重用。只要填入實際參數那一段重新計算。論文裡舉的是旅遊行程的例子,trip-plan 模組固定,只有天數這個參數在跑。
第三個概念是 union。現實裡常見一組候選片段之間互斥,比方說同一份說明文件的中、英文版。這種情況下,PML 用 <union> 把這幾個模組包起來,指示它們共用起始位置編號。實作上,這有助於節省 position id 的空間,未來也可以做一些預取上的小手腳。注意這跟參數不一樣,參數是在模組內留空間,union 則是多個完整模組共享一個位置區間。
最後還有 nested module 跟 system / user / assistant 這類跟具體模型模板對應的 tag。這些設計讓同一個 schema 可以對應到不同 LLM 的 chat 格式,而不用手動再多包一次 bracket 或特殊 token。
4. Schema 編碼與位置編號:怎麼讓「模組化」不破壞語意
要把一個模組的 KV 狀態事先算好,第一步是決定每個 token 的 position id。Prompt Cache 的做法是先根據 schema 的順序,依序累加每個模組的長度。比方說前兩個模組長度分別是 50 跟 60,第三個模組起始位置就從 110 開始。union 的情況比較特別,同一個 union 裡的模組共用一個起始位置,長度則取那組裡最大的那個做基準。
值得注意的一點是,這些位置編號不需要從零開始。中間有一段「像空白一樣的空窗」,對語意影響不大。這讓系統在日後組合模組時,有空間塞入新的非 cache 文字,或是留 buffer 給參數,卻不用重建整個編號表。
接著系統把這段 token 序列跟對應的 position id 丟進 LLM,讓模型算出 key / value 狀態。這時候會遇到一個實務細節。很多現成模型的實作假設 position id 是連續的,需要稍微改一下程式碼。論文以三種常見位置編碼為例。
對於直接用 lookup table 的位置 embedding,例如早期 BERT、GPT-2,問題比較小。因為 position id 本來就是拿來查表,改成非連續只是在 index 上跳來跳去。
對 RoPE 類的旋轉位置編碼,像 Llama2、Falcon,用法是在注意力計算時,根據位置產生旋轉矩陣。這裡他們會先把不同位置對應的旋轉結果預先算成表格,接著按 position id 查詢。換句話說,把「連續往前算」換成「按 id 查一次」。ALiBi 類似,抽象成給每個位置對一個 bias,再根據 id 查表套進 softmax 前的分數。
最後還有一個小但實務上很有用的修改。PyTorch 預設 tensor concat 每次都會配置新記憶體。Prompt Cache 在執行階段要頻繁把不同模組的 KV tensor 接起來,如果每次 concat 都重新配一塊,就會被 allocator 拖累。作者自己寫了一個有 buffer 的 concat 實作,盡量重用配置過的空間,減少記憶體壓力和分配時間。
5. Cached Inference :新請求進來時發生什麼事
當有一個新 prompt 進到系統,整個流程大致是這樣。先解析 PML,看這份 prompt 說自己是基於哪個 schema,再確認裡面 import 的模組名稱都存在。這一步類似靜態型別檢查,避免有人引用不存在的模組或參數。
確定合法之後,系統會找出每個被引用模組對應的 KV cache,把它們依照需要的順序接起來。這邊有一個有趣的小結論。理論上注意力對 token 順序很敏感,不過對於已經算好的 KV 狀態,只要相對順序跟位置編號一致,拼接時的實體記憶體排列順序不影響結果。論文引用了 permutation invariance 的分析,說明這點在 transformer 架構下是合理的。
接著輪到非 cache 的部分。包括 schema 之外的新文字、參數實際填入的 token。系統會根據它們在整個 prompt 裡的相對位置,安排對應的 position id。如果這段文字在兩個模組 A、B 中間,就沿用 A 結尾之後的那一段位置編號。對參數來說,則是用先前 <unk> 位置留下的那串 index。完成這一步,就得到「這次請求中需要新算的 token」和一組「它們對應的位置」。
然後把前面接好的 KV tensor 當作現成的 cache,連同剛產生的 token 和 position id 一起送進模型。模型看到的是一條完整的長序列,前半段 KV 來自 cache,後半段來自這次真正餵入的 token。注意 prefill 階段在這裡就結束。後面自回歸生成每一個新 token 的步驟,還是交給原本的 KV Cache pipeline 處理。Prompt Cache 只把「第一輪讀 prompt」這件事改寫成記憶拷貝加少量補算。
這個設計的另外一個效果,是在 batch 推論時有機會共享 KV。多個請求如果都是基於同一個 schema,自然會引用同一批模組。再配合像 paged attention 這種把 KV block 管理成頁面的技巧,可以讓同一個模組在一個 batch 裡只佔一份實體記憶體。對高併發的服務來說,這是額外的吞吐空間。
6. 實驗結果: TTFT 的加速幅度和準確率影響
論文用 Llama2、MPT、Falcon 等多個開源模型,加上 LongBench 這組長上下文基準,去量化 Prompt Cache 的效果。觀察重點有三個,時間、準確率、記憶體。
時間方面,重點放在 GPU 和 CPU 上的 time-to-first-token。GPU 那邊,他們把 prompt module 放在 GPU 記憶體,看到的加速落在大約 5 倍到 10 倍之間。模組放在 CPU 記憶體,推論時再搬進 GPU,速度提升也還有大約 1.5 倍到 3 倍。這兩個數字基本上可以當成實務上「全放 GPU」和「全放 CPU」的上下界。實際系統可能會介於中間。
CPU 那邊變化更大。Intel i9 這組,實驗結果最高達到 70 倍左右的 TTFT 縮短。AMD 那組則接近 20 倍。論文推估,這跟記憶體頻寬差異有關。整體來看,CPU 做長序列 self-attention 本來就比 GPU 慢得多,把這塊算力換成線性成長的 memcpy,是非常划算的交換。
準確率方面,他們在 LongBench 上用 deterministic 的取樣方式,比較 baseline KV Cache 和 Prompt Cache 的分數。大部分數據集差異落在微小範圍,偶爾有個別任務略高或略低。論文沒有看到明顯災難性的品質流失。這印證前面提到的那個前提,只要模組切在語意封閉的位置,且內部相對順序不變,模型對這種模組化注意力的調整並不敏感。
記憶體方面,他們把每個模型「每 token 需要的 KV 空間」列成表格。小模型像 Falcon 1B,一個 token 的 cache 大約 0.18 MB。代表 1K token 的模組大概 180 MB。更大的像 Llama 70B,每個 token 大約 2.5 MB,1K token 模組就要 2.5 GB。對應到實際系統,很明顯可以感覺出來,小模型可以大方把模組放在 GPU。大型模型如果要 cache 長文件,幾乎一定得退到 CPU,甚至再加一層壓縮。
7. 模組切分與注意力遮罩:品質安全帶怎麼綁
Prompt Cache 的一個設計副作用,是把注意力視窗縮在每個模組內。也就是說,一個模組裡的 token 看不到其他模組的內容。這點其實很像一種特殊的 attention mask。Longformer 那種長文模型也有用類似想法,只把注意力聚焦在鄰近區域,借此取得近似的表現。
這樣的遮罩對品質有兩面性。對語意上獨立的模組,比方說單純的系統說明、法條原文、個人背景描述,內部自成一個小上下文,關掉跨模組的注意力有機會反而減少雜訊。對需要跨段落推理的場景,情況複雜一些。如果你強迫模型在編碼階段把所有關聯都塞在同一個模組裡,不見得會符合現實需求。
論文為此提出 scaffolding 這個折衷。使用者可以指定一組模組在編碼時一起處理,讓它們共享注意力範圍。系統仍然會保留每個模組單獨的 KV 狀態,方便其他情況重用。不過當 prompt 一次引入了 scaffold 裡的全部成員,會用這個合併結果覆蓋單獨模組的版本。換句話說,把更多記憶體拿去換輸出的穩定度。
整個架構在品質上還有一個實務建議。每個 prompt module 盡量做到自洽,不依賴前後模組裡的細節。像你在 schema 裡常見的會是「完整 system prompt 一段」、「某一條業務規則一段」、「某個 RAG 文件自己一段」。這種切法一方面讓遮罩看起來合理,一方面也比較符合現有服務的配置方式。
8. 適合的應用場景:從程式碼輔助到個人化與 RAG
論文用三個例子來示範 Prompt Cache 在實際任務裡的用法,對應到現在很多團隊正在做的事情。
程式碼生成那個案例,把專案裡的幾個重要類別或檔案各自變成 prompt module。當開發者要針對某段程式碼提問時,在 prompt 裡「匯入」相關模組,例如 Unit、Map、Game、Player 這些。系統幫你把這些檔案對應的 KV 狀態接起來。實驗顯示,在 GPU 上 time-to-first-token 大約縮成原本的四分之一,輸出的程式碼內容一樣。
個人化的例子,則是假設有一組固定的使用者特徵。像年級、程度、學習歷史、喜好風格、評量方式等,把每個特徵做成獨立模組,還用 union 指出互斥關係。實際生成時,根據當下的 profile 選擇對應模組,組出一個完整的個人化 prompt。因為每個特徵片段在 cache 裡都準備好了,換使用者時只是在 KV 層換一批積木。
最後是帶參數的旅遊行程模板。整段行程安排寫在一個模組裡,中間留幾個 <param> 作為目的地、天數之類的插槽。使用者送出請求時,填好這些參數,Prompt Cache 只對這些位置附近重新算,其他文字完全照 cache 模式走。論文的實驗裡,這種參數化模板同樣達到明顯的 TTFT 降低。
這三個例子其實可以串成一個更大的圖。不管是程式碼上下文、個人化特徵、旅遊模板,本質上都是「有限集合裡反覆出現的長段文字」,差別只在於維度拆在檔案、使用者屬性,還是任務模板。Prompt Cache 把這種結構變成顯式的 cache 單位,而不是默默塞在字串常數裡。
9. 總結:把 prompt 當 layout 來設計,而不是一大串字
從工程角度看,Prompt Cache 的重要性不只在於一組漂亮的 TTFT 數字,更在於它逼大家承認一件事。長 prompt 不是一段線性的字串,而是一個有版面配置的東西,有模組、有層級、有布局。只要把這個結構標出來,很多原本看起來只能硬撐的成本,其實都可以搬到預先計算或 cache 層處理。
論文本身已經證明,這種模組化注意力重用在多種模型、多種任務上都能維持接近原本的準確率,同時大幅縮短 time-to-first-token,特別是長 prompt、長上下文的場景。代價在記憶體,要看模型大小和模組長度來權衡。對小模型,GPU cache 看起來非常實用。對大模型,CPU cache 搭配壓縮會變成現實選項。
如果把這篇論文放回實際產品脈絡,它比較像一個「下一代 LLM serving 架構的積木」而不是單一小技巧。未來各種 LLM 平台、gateway,很可能在 API 背後都有一層類似 PML 的 schema,幫你把重複段落拆開、預算記憶體、擺放在對應的 cache 層。你要做的事情,只是誠實把自己系統裡那些反覆拷貝的 prompt 原型畫出來,決定哪些該進模組、哪些還維持即時生成。
站在寫服務的人角度,最實際的下一步,就是回頭檢查現在線上的幾個核心 prompt。畫出它們的共同結構,統計真正會變動的部分有多少。這個過程本身就會讓架構更清楚。之後不管是要採用 Prompt Cache 這樣的實作,還是用其他方案改寫 prefill,心裡都會有一把尺,知道哪裡在燃燒 GPU,哪裡其實只是還沒把模組拆開。