譯者 | 核子可樂

審校 | 重樓

是不是總感覺ChatGPT響應(yīng)緩慢？

大家可能沒有留意，大模型應(yīng)用程序的提示詞重復(fù)率高達(dá)70%，問天氣、問翻譯和問節(jié)日安排的內(nèi)容大量出現(xiàn)，且每次都要消耗算力進(jìn)行處理。這樣的情況在分布式集群的各節(jié)點(diǎn)上被無數(shù)次放大，白白燒掉寶貴的能源和金錢。

為此，Anthropic日前詳細(xì)介紹了如何利用提示詞緩存技術(shù)將推理成本降低90%。其實(shí)不少開源大模型運(yùn)行時(shí)（包括vLLM、TRT-LLM和SGLang等）都擁有自動(dòng)前綴緩存（也稱上下文緩存）功能，負(fù)責(zé)將相同前綴請(qǐng)求中的輸入提示詞自動(dòng)緩存起來。

前綴緩存的工作原理

為了更好地理解前綴緩存，我們先來聊聊大模型推理的工作原理。

推理過程在宏觀上分為兩個(gè)步驟：

• 通過正向傳遞處理給定的輸入標(biāo)記序列，即預(yù)填充階段。
•  解碼階段，從首個(gè)token連續(xù)生成至最后一個(gè)token，且當(dāng)前token依賴于上一token。

圖一

由于此過程的自回歸屬性（即新token依賴于前一token），因此有效的內(nèi)存管理非常重要。多數(shù)大模型會(huì)采取為中間狀態(tài)保留KV緩存的做法。與簡(jiǎn)單提示詞或語義緩存的不同之處在于，其不會(huì)將全文輸入和輸出保存在數(shù)據(jù)庫內(nèi)，因?yàn)檫@樣就只有完全匹配（或者幾乎完全相同的查詢）才能立即命中緩存并收到響應(yīng)。

在預(yù)填充階段，在大模型處理token時(shí)會(huì)計(jì)算“注意力”，即每個(gè)token與其他token的關(guān)系。計(jì)算過程會(huì)為每個(gè)token生成鍵-值矩陣。如果不經(jīng)任何KV緩存，那么模型每次回顧此前token時(shí)都需要重新計(jì)算這些矩陣。KV緩存在設(shè)計(jì)上只支持一次生成，即只會(huì)在生成一條輸出的過程中捕捉中間狀態(tài)。

如果有兩條具有相同前綴的請(qǐng)求，該怎么處理？

KV緩存的基本思路啟發(fā)并衍生出了前綴緩存，確保在生成包含相同前綴的提示詞時(shí)給出不同的響應(yīng)。簡(jiǎn)單類比一下，假設(shè)已經(jīng)計(jì)算過2 * 6的結(jié)果，那么對(duì)于2 * 6 * 3 * 5 這個(gè)新問題，可以直接復(fù)用之前的答案，避免在序列中重復(fù)計(jì)算。

這對(duì)應(yīng)用程序有何幫助？

我們可以使用以下最佳實(shí)踐來充分發(fā)揮前綴緩存的優(yōu)勢(shì)：

為提示詞結(jié)構(gòu)設(shè)計(jì)策略

可以將系統(tǒng)提示詞、基礎(chǔ)指令或者共享上下文等常量元素放在提示詞的開頭（圖二），從而為多條查詢建立可復(fù)用基礎(chǔ)。其他動(dòng)態(tài)或特殊內(nèi)容則可放在末尾。

圖二

對(duì)請(qǐng)求進(jìn)行重新分組

將共享通用結(jié)構(gòu)/前綴的請(qǐng)求捆綁在一起（圖三）。例如，在處理以常見問候語或稱呼開頭的多條客戶查詢時(shí)，可以嘗試將它們分為一組，盡可能提高計(jì)算過程的緩存和復(fù)用率。

圖三

監(jiān)控緩存利用率

另外，需要注意跟蹤緩存利用率。

包括命中率與未命中率：

• 找出哪些前綴比其他一般前綴更重要
• 識(shí)別緩存未命中的模式

依托這些見解，就能優(yōu)化提示詞結(jié)構(gòu)以獲得最佳性能。

簡(jiǎn)單示例

以下示例為當(dāng)多條查詢共享相同的上下文時(shí)，前綴緩存如何優(yōu)化大模型推理。我們使用一份簡(jiǎn)單的員工數(shù)據(jù)庫表，并對(duì)其中內(nèi)容進(jìn)行不同查詢。

Python
import time
from vllm import LLM, SamplingParams

# A small table containing employee information
LONG_PROMPT = """You are a helpful assistant that recognizes content in markdown tables. Here is the table:

| ID | Name          | Department | Salary  | Location    | Email                |
|----|---------------|------------|---------|-------------|---------------------|
| 1  | Alice Smith   | Engineering| 85000   | New York    | alice@company.com   |
| 2  | Bob Johnson   | Marketing  | 65000   | Chicago     | bob@company.com     |
| 3  | Carol White   | Sales      | 75000   | Boston      | carol@company.com   |
| 4  | David Brown   | Engineering| 90000   | Seattle     | david@company.com   |
| 5  | Eve Wilson    | Marketing  | 70000   | Austin      | eve@company.com     |
"""

def get_generation_time(llm, sampling_params, prompts):
    start_time = time.time()
    output = llm.generate(prompts, sampling_params=sampling_params)
    end_time = time.time()
    print(f"Output: {output[0].outputs[0].text}")
    print(f"Generation time: {end_time - start_time:.2f} seconds")

# Initialize LLM with prefix caching enabled
llm = LLM(
    model='lmsys/longchat-13b-16k',
    enable_prefix_caching=True
)

sampling_params = SamplingParams(temperature=0, max_tokens=50)

# First query - will compute and cache the table
get_generation_time(
    llm,
    sampling_params,
    LONG_PROMPT + "Question: What is Alice Smith's salary? Your answer: Alice Smith's salary is "
)

# Second query - will reuse the cached table computation
get_generation_time(
    llm,
    sampling_params,
    LONG_PROMPT + "Question: What is Eve Wilson's salary? Your answer: Eve Wilson's salary is "
)

運(yùn)行以上代碼，即可查詢不同查詢間的實(shí)際時(shí)間差異。第二條查詢明顯更快，因?yàn)槠鋸?fù)用了緩存中的表上下文。具體時(shí)間將根據(jù)硬件和設(shè)置而有所浮動(dòng)。

總結(jié)

前綴緩存是一項(xiàng)強(qiáng)大的大模型應(yīng)用優(yōu)化技術(shù)。實(shí)施上述最佳實(shí)踐將幫助開發(fā)人員顯著降低推理成本，且不致影響響應(yīng)質(zhì)量。參考示例也表明其操作難度極低，推薦大家馬上在自己的應(yīng)用程序中試一試。

原文標(biāo)題：90% Cost Reduction With Prefix Caching for LLMs，作者：Mahak Shah