首页 › LLM大语言模型 › 从0-1构建LLM大语言模型 › 大语言模型四大组件:KV、前缀、提示与语义缓存

大语言模型四大组件:KV、前缀、提示与语义缓存

在大型语言模型的架构中,有四个组件分别存储着四种不同的对象,而它们都被统称为缓存。​

​

飞书文档 - 图片

​

•

KV缓存为单个请求存储注意力张量。​

•

前缀缓存则在服务器上以基于标记ID的哈希链作为键,存储相同的张量。​

•

提示缓存是提供商对该相同查找操作的计费版本,其读取费用为基本输入费率的0.1倍,而写入费用则为基本费率的1.25倍。​

•

语义缓存存储已完成的响应字符串,并以嵌入空间中的余弦相似度作为键。​

​

前三个是精确匹配且与正确性无关,因此一次未命中会带来成本和延迟。第四个是模糊匹配,它会在返回200错误码的同时给出一个错误答案。​

​

那么今天,我们就来逐一介绍这四种类型,看看它们各自存储什么,以及哪些因素会悄悄地破坏它们。​

此处的所有内容均运行于一台机器上,包括 CPU 在内,采用的是一个拥有 3.6 亿参数的模型。此外,还提供了一个 Anthropic API 示例,以及一个基于 sentence-transformers 构建的小型语义缓存。对于那些仅存在于推理引擎内部的机制,我们以伪代码形式呈现其逻辑,而非假定它能在普通笔记本电脑上复现。​

​

代码块​

Bash

pip install "transformers>=5.0" torch​

​

# only for the quantized cache example​

pip install optimum-quanto​

​

# only for the semantic cache example​

pip install sentence-transformers ​

​

# only for the prompt caching example​

pip install anthropic​

​

​

KV缓存​

在预填充过程中,模型会为每一层中的每个提示标记计算一个键向量和值向量,并将其存储起来。​

解码阶段会遍历已存储的向量,并为每个生成的标记追加一对新的向量,而不是在每一步都重新计算整个序列。​

​

飞书文档 - 图片

​

查询不会被缓存,原因在于因果掩码机制。一个标记的查询向量仅在该标记被处理时使用一次,此后不再被读取。而它的键和值则会被其后的每一个标记所读取,因此这二者是最值得保存的两个向量。​

•

如果不存储它们,每一步解码都需要对迄今为止生成的完整序列进行一次矩阵乘法。​

•

而有了它,这一步就只需对一个新标记进行一次矩阵与向量的乘法,所需的浮点运算次数要少得多。​

​

飞书文档 - 图片

​

​

🧠

尽管这降低了每个标记的计算量,但你必须在每一步都将整个缓存从 HBM 加载进来,因此解码不再受计算能力的限制,而是转变为内存带宽的瓶颈。注意力核的计算速度比缓存数据的加载速度更快,因此在解码的一个步骤中,GPU大部分时间都在等待内存访问。​

​

​

飞书文档 - 图片

​

​

<b每个token都会导致KV缓存增长​

transformers库将缓存作为一等对象公开,因此您可以保存、检查并将其重新传递回去。​

<b以下是它的最小代码示例:​

​

代码块​

Python

import torch​

from transformers import AutoTokenizer, AutoModelForCausalLM, DynamicCache​

model_id = "HuggingFaceTB/SmolLM2-360M-Instruct"​

tokenizer = AutoTokenizer.from_pretrained(model_id)​

model = AutoModelForCausalLM.from_pretrained(​

​

​

​

​

​

50%

​

​

​

​

50%

​

​

•

通常情况下,您只需调用 generate 方法,缓存就会在内部自动创建并销毁,对您而言是透明的。而在这里,我们自行构建了一个 DynamicCache 并将其传入,这意味着在生成完成后,我们仍然持有对该缓存的引用。​

•

get_seq_length() 会报告缓存中存储的标记位置数。运行该函数时,输出结果为提示长度加上生成的标记数,再减去1。​

•

最后一个标记的键和值已被计算,但从未被任何东西访问过。​

​

缓存随请求一起被释放​

​

🧠

以上所有操作均在一次调用中完成。引擎会在请求结束后释放这些内存块,因此在第20回合时,20轮的聊天预填充会再次从第1轮到第19轮依次执行,且需支付全额费用。​

​

​

​

您可以通过在各回合之间自行保持缓存的存活来查看备选方案。​

​

代码块​

Python

​

•

past_key_values 对象在循环外部被创建一次,并在每次生成调用时传递进去。这样一来,缓存不会在第一轮结束时被释放,而在第二轮开始时仍保持已填充状态。​

•

在每一轮中,我们都会重新构建完整的消息列表,并通过 apply_chat_template 重新渲染。第二轮发送的提示信息包含了第一轮的所有内容,以及新的问题。​

•

由于缓存中已保存了第一轮的标记,模型只需对新的后缀进行前填充。每轮打印的input_length都会增加,而实际的前填充工作量却保持不变。​

•

完成内容会从生成的ID中截取,并重新附加到消息中。这正是使下一轮的提示严格延续上一轮的原因。​

重用之所以有效,是因为第二轮的标记序列从第一轮的标记序列开始,逐位完全相同。如果你修改了历史记录中更早的任何内容,缓存就会失效。​

在本代码演示中,缓存属于单个进程中的一个 Python 变量。而在服务引擎中,它则归属于一个由数千个请求共同访问的共享池。接下来我们就来了解一下这一点。​

​

前缀缓存​

上述提到的共享池源于一种行为上的变化。​

当一个请求执行完毕后,引擎会将其 KV 块保留在内存中,而不是将其释放,并为其建立索引,以便后续的请求能够快速定位。这就是前缀缓存。​

该索引必须遵循聊天循环中所涉及的同一规则,即仅当先前的标记完全相同时,重用才有效。​

vLLM 通过默认存储 16 个标记的缓存,并利用父块哈希值与其中标记 ID 的组合哈希值来标识每个区块。​

​

​

将父哈希链入子块,可将区块查找转化为前缀查找,因为只有在其之前的所有内容均匹配时,该区块才会匹配。​

调度器按顺序遍历传入的块,并在遇到首次未命中时停止。若发生命中,则该块的引用计数加一,同时在有请求正在使用该块期间将其固定,防止被驱逐。​

从miss开始的所有内容都会获得新的分配和新的预填充。​

查找代码​

vLLM在其调度器内部运行这一过程,并由负责管理实际张量的内存管理模块进行封装。​

下面的代码仅保留了决定是否重用的两个部分,即把标记序列转换为块键的函数,以及遍历这些键以确定前缀中可以跳过预填充的部分的函数。​

​

代码块​

Python

BLOCK_SIZE = 16​

def block_hashes(token_ids, salt=None):​

"""Chain-hash a token sequence into per-block keys."""​

hashes, parent = [], hash(salt)​

# Only complete blocks are hashed. A partial tail block is skipped.​

for start in range(0, len(token_ids) - BLOCK_SIZE + 1, BLOCK_SIZE):​

block = tuple(token_ids[start : start + BLOCK_SIZE])​

parent = hash((parent, block))​

hashes.append(parent)​

​

•

block_hashes 方法将代币序列切分为固定长度的16个代币区块。每个区块的密钥通过哈希运算(parent, block)与前一个区块的密钥进行折叠,因此第五个密钥编码的是从第一个区块到第五个区块的信息,而非仅第五个区块本身。​

​

​

•

范围止于 len(token_ids) - BLOCK_SIZE + 1,这会丢弃尾部的任何不完整块。这些标记永远不会被索引,并且会在每次以该位置结束的请求中重新计算。​

•

schedule 方法按顺序遍历键,并在遇到第一个缺失的键时停止。它不会尝试在序列中后续位置继续匹配,因为后续块的键已经依赖于先前失败的那个键。​

•

ref_count += 1 会将该块标记为正在使用中。驱逐操作仅作用于计数为零的块,这正是为了避免正在运行的请求被突然移除其缓存。​

•

凡是匹配到的内容都会被放入 reused_tokens,而其后的部分则会被重新预填充。​

在我们刚刚讨论的代码中,还有一件重要的事情:​

​

代码块​

Python

BLOCK_SIZE = 16​

def block_hashes(token_ids, salt=None):​

"""Chain-hash a token sequence into per-block keys."""​

hashes, parent = [], hash(salt)​

# Only complete blocks are hashed. A partial tail block is skipped.​

for start in range(0, len(token_ids) - BLOCK_SIZE + 1, BLOCK_SIZE):​

block = tuple(token_ids[start : start + BLOCK_SIZE])​

parent = hash((parent, block))​

hashes.append(parent)​

return hashes​

​

请注意上述函数中的 salt 参数。​

当两条请求发送的内容完全相同的时候,它们会生成相同的块密钥,因此最终指向 GPU 显存中的同一物理 KV 块。这些张量在显存中仅有一份副本,而两条请求都会读取这份副本。​

当两个请求都来自同一应用程序时,这正是您期望的行为。​

但当它们来自不同的客户时,可能需要作出区分。因此,将每个租户的值作为盐值进行传递,会改变首个父哈希值,从而使相同的文本为每个租户生成不同的密钥,进而确保各租户的请求不会落在同一区块上。​

这样一来,每个租户都会获得自己的副本,这会占用内存并降低缓存命中率,但能够实现隔离。​

​

在 transformers 中的实现​

transformers 允许你预先填充一次提示,并在多个不同的续写中复用生成的缓存。​

​

代码块​

Python

import copy​

import torch​

from transformers import AutoModelForCausalLM, AutoTokenizer, StaticCache​

model_id = "HuggingFaceTB/SmolLM2-360M-Instruct"​

tokenizer = AutoTokenizer.from_pretrained(model_id)​

model = AutoModelForCausalLM.from_pretrained(​

model_id, dtype=torch.bfloat16, device_map="auto"​

)​

SHARED_PREFIX = """You are a careful assistant. ​

​

•

由于我们需要一种可以四处复制的固定分配,因此使用了StaticCache而非DynamicCache。​

•

模型(...) 调用属于预填充。此处不会采样任何 token。我们仅将共享前缀输入模型,以填充缓存,随后保留返回的 past_key_values。​

•

在循环内部,每个问题都被拼接到同一个前缀上。整个字符串会被分词,因此前缀部分的标记 ID 每次都相同,而这正是引擎的哈希链所检查的条件。​

•

copy.deepcopy 为每个请求提供一份预先填充的缓存副本。生成过程通过追加操作就地修改缓存,因此如果没有复制,第一个问题就会污染第二个问题的前缀。生产环境中的引擎并不会复制张量。相反,它共享物理块并跟踪引用计数,这使得重用几乎无需额外开销,而非与前缀长度成正比。​

​

逐出对命中率的影响​

如上所述,只有完整的区块才会被索引,因此每次都会重新计算尾部的不完整区块。​

这意味着块大小应当适当调整​

•

更大的块意味着更少的表查找和更好的内存局部性​

•

更小的块意味着更细粒度的共享以及尾部更少的浪费。​

​

​

正如预期,驱逐操作会降低命中率。​

缓存与正在运行的批处理共享同一块GPU显存池,因此更大的缓存会导致并发序列数量减少;在压力较大时,vLLM会按最近最少使用算法回收未被引用的块。​

混合流量使情况更加恶化,因为较长的共享前缀占用最多的地址块,而这些前缀一旦丢失确实会造成严重损失。​

在启用此功能之前,您应当了解两件事​

•

它仅保存预填充数据,因此解码时间保持不变,将全部加速效果归因于缓存会高估其作用。​

•

而且哈希运算本身也需要消耗一定的开销,因此在处理真正唯一的查询时,基准测试显示吞吐量反而出现了下降,而非提升。​

还有一个第三方面的问题,它与工作负载相关,并且对RAG的影响最为显著。​

​

​

A RAG 提示包含一条系统指令、若干检索得到的文本块,以及查询;这些文本块在每次请求中都会变化,且不同请求之间的顺序也可能不同。即使两次请求检索到的文档完全相同,但顺序不同,它们在链式哈希下也毫无关联。​

​

​

单独预填充每个数据块并将其缓存拼接起来是行不通的。​

拼接后的张量携带了错误的位置编码。每个分块从未关注过其他任何分块,而且每个分块都在模型认为的位置零处贡献了自己的注意力汇聚点。要使其正常工作,需要在边界处进行部分重计算,而非简单的拼接。​

​

​

顺便说一句,该解决方案在开源社区中早已存在。​

(开源)实现了 CacheBlend,其工作原理是:不将分块缓存首尾相连,而是将其在任意位置复用,并仅重新计算一小部分标记,这些标记由预计算值与全注意力机制所生成结果之间的最大偏差来确定。LM缓存​

​

​

该子集恢复了跨分块注意力机制,并修正了位置编码,因此输出能够保持全预填充质量。​

这使得首次生成标记的耗时相比重新计算全部内容提升了约两到三倍,且重新计算的开销与从较慢的存储介质中获取已缓存的数据块的开销实现了流水线化。​

它可与 vLLM 集成,并从您的提示中读取块边界,因此即使每次检索到的文档顺序不同,检索流量也能得到复用。​

​

​

这是仓库地址: https://github.com/LMCache/LMCache​

​

提示缓存​

在托管模型中,你无法获得任何块表或逐出策略。相反,你会得到一份基于提供商自身前缀复用的定价表,以及两个用于控制的参数。​

缓存的对象仍然是KV张量,而非您的提示文本,并且仍然需要对完整渲染后的上下文进行精确的前缀匹配。​

​

​

渲染后的上下文包含了您从未编写过的提供方系统内容,这也是为什么从外部看最小长度和失效规则显得随意的原因。​

以下是通过代码演示的提示缓存版本:​

​

代码块​

Python

import anthropic​

client = anthropic.Anthropic() # reads ANTHROPIC_API_KEY from the environment​

# Must clear the model's minimum cacheable length or nothing is cached at all.​

LONG_INSTRUCTIONS = "You are a precise technical editor. " * 400​

def ask(question: str):​

return client.messages.create(​

model="claude-sonnet-4-6",​

max_tokens=512,​

system=[​

{​

"type": "text",​

"text": LONG_INSTRUCTIONS,​

"cache_control": {"type": "ephemeral"}, # everything above is cacheable​

}​

],​

messages=[{"role": "user", "content": question}],​

​

那段代码中只有一行会触及缓存。​

您指定 cache_control 的位置决定了请求的哪一部分会被写入条目,而使用计数器则能告知您后续的调用是否已读取该条目。​

•

标记被附加到您希望覆盖的最后一个块,而非某个范围。它会写入一条缓存条目,涵盖从请求开始直至该块(含该块)的所有内容。​

•

用户消息位于标记下方,因此它始终处于缓存区域之外,因为每次调用时都会发生变化,所以不能放在缓存区域内。​

•

使用计数器可以帮您了解底层的运行情况。第一次调用报告了非零的 cache_creation_input_tokens 以及零次读取;第二次则相反,且指令费用按输入速率的十分之一计费。​

•

如果两个计数器都返回零,则该前缀的长度低于模型的最小可缓存长度,请求将完全不进行缓存处理。对此不会产生任何错误。​

从直觉上来看(正如上文所述),如果我们将 cache_control 下移到用户消息层面,那么读取计数器将始终为零,因为在每次调用时,被标记的块都会发生变化。​

提示缓存的经济性​

Anthropic 对写入条目收取基础输入费率的1.25倍,读取条目则收取0.1倍;若希望条目保留更长时间,则适用更高的写入倍率。OpenAI 在其现有模型上也采用了同样的两种倍率。​

由于在TTL范围内被重复使用的任何内容都会避免重新计算,因此在后续请求中可收回溢价成本。​

读取操作只能找到由先前某个请求写入的条目,而写入操作仅会在您设置的断点处发生。​

​

​

每次调用都会检查您的断点;若未命中,则会向前回溯有限数量的块,以查找较早的写操作。​

由于Anthropic的上限为20个区块,因此在两次调用之间添加超过20个区块的对话内容,会导致最后一次写入超出范围,从而停止命中。​

​

语义缓存​

语义缓存会嵌入传入的提示,在已存储的提示中执行最近邻搜索,并在相似度超过阈值时直接返回一条已存储的响应。​

​

​

这就是为什么它既会保存输出标记,也会保存输入标记。这也正是为什么每次请求都必须进行一次嵌入的往返操作,包括每一次未命中。​

以下是一个用几行代码实现的语义缓存演示:​

​

代码块​

Python

# requires: pip install sentence-transformers​

import numpy as np​

from sentence_transformers import SentenceTransformer​

encoder = SentenceTransformer("all-MiniLM-L6-v2")​

class SemanticCache:​

def __init__(self, threshold=0.95):​

self.threshold = threshold​

self.vectors = np.empty((0, encoder.get_sentence_embedding_dimension()))​

​

上述类中的每一种方法都对应着你在生产环境中必须做出的一个决策:​

•

normalize_embeddings=True会将每个向量归一化为单位长度,从而使 self.vectors @ vec 计算的余弦相似度等同于普通的点积。如果跳过归一化,则无法比较不同长度的提示之间的得分。​

•

查找操作会将嵌入向量与结果一同返回,因此答案可以稍后直接使用该嵌入向量,而无需再次进行嵌入。这一点很重要,因为无论命中还是未命中,每次请求都需要支付嵌入费用;如果对同一查询重复计算嵌入,就会使缓存的总体成本翻倍。​

•

在演示场景下,暴力求解的 argmax 方法尚可;但在大规模场景下则不适用。一旦数据条目超过数千条,这种方法实际上就变成了近似最近邻索引,而这会在阈值之外额外引入召回率方面的权衡。​

•

存储操作仅在未命中路径上、模型作出回答之后才会被调用。在该回答成为未来所有得分高于阈值的提示的响应之前,没有任何机制对其正确性进行验证。这凸显了该方法的最大风险:缓存并不知道所存储的响应是否正确,它只关心新的提示与旧提示是否相似。​

下面的代码展示了最后一点:​

​

代码块​

Python

​

这是我们得到的输出:​

​

代码块​

Python

0.961 'How do I reset my password?' vs 'How can I reset my password?'​

0.952 'Is the API rate limited?' vs 'Is the API not rate limited?'​

0.887 'Refund policy for annual plans' vs 'Refund policy for monthly plans'​

​

•

第一组是真正的释义,应给出相同的答案。​

•

第二组仅有一个否定词的差异,需给出相反的答案。​

•

第三组仅有一个运算值的差异,需给出不同的答案。​

尽管存在一些偏差,但三者的得分非常接近。改写与否定之间的差距还不到0.01分,这一差距在实际流量中根本无法维持。​

•

如果将其调高,命中率会大幅下降,而每次调用仍需为嵌入付费。​

•

如果将其调低,命中率会上升,但同时也会导致大量自信的错误答案。​

•

据不同人士的说法,已发布的默认值范围在0.75至0.97之间,这表明该参数取决于您的流量特征,而非可供照搬的固定数值。​

这本身并不是一种完全可靠的技术,因为某些失败情况(如上文所示)能够绕过任何阈值,它们源于嵌入所表达的内容。​

​

四种技术的总结​

​

​

上述四种技术中有三种与正确性无关,因此其未命中仅体现在开销和延迟上,而不会在其他方面体现。​

语义缓存的工作方式不同,因此在此处不宜以命中率作为衡量指标。​

​

此外,还存在第五层——一种使用较少的缓存机制。​

它是一种精确匹配响应缓存,当请求与存储的答案在字节层面完全一致时,便会直接返回该存储答案。​

它像语义缓存一样同时保存输入与输出,且不存在误报风险,因为它完全不进行任何相似度匹配。​

在使用嵌入之前,你只需测量一下字节级的重复率即可。当然,其中也存在一些问题,相信你现在已经发现了。​

​

​

生产方面的要点​

每种技术在投入生产使用前,都有一些需要注意的失效点:​

•

如果提示词的开头包含任何变量,例如系统提示中的时间戳、请求 ID 或用户名,则其后的所有块都将失效。请始终将稳定内容置于最前,可变内容置于最后,并在边界处设置标记。​

​

​

•

工具模式通常置于系统提示之前,因此重新排序可能会使整个缓存失效。​

•

请检查将被渲染到提示中的各项设置。在 Anthropic 平台上,切换网页搜索、引用、思维配置或工具选择选项都会重写提示文本,并使下游模块失效。对两种推理方案进行 A/B 测试时,您的缓存将被一分为二。​

•

对历史进行汇总会重写前缀,因此下一次调用将按冷令牌的全价计费。就地截断工具输出可保持前缀字节完全一致,并使缓存继续有效。​

​

​

•

缓存条目的键值基于模型,因此即使路由到更便宜的模型,仍会以冷启动费率预填充整个累积的历史记录。​

​

​

要准确判断两个提示在何处停止匹配,请直接比较它们的标记 ID,而非您记录的文本。以下是演示:​

​

代码块​

Python

​

在您的日志中看似相同的两个提示,可能因开头序列标记(BOS)、尾部换行符或重新序列化的工具模式而有所不同。​

通过比较标记 ID 而非渲染后的文本,可以精确定位复用停止的索引位置;而对两侧少量 ID 进行解码,通常也能还原出准确的文本。​

​

上面的运行展示了一个常见的例子。将一条指定的无系统消息进行转换,使聊天模板填充为模型的默认值,且在索引3处两个提示看起来不同,因此无法复用。​

​

前三层涵盖一个理念,并应用于三个范围。​

•

KV缓存会在单个请求的持续时间内保存注意力状态。​

•

前缀缓存会在请求结束后保留该状态,以便后续请求可以进行查找。​

•

提示缓存是指提供商在其自有硬件上运行前缀缓存,并对你重复使用的部分另行收费。​

语义缓存的工作方式有所不同。它以嵌入相似度为键来存储响应文本,因此在发生命中时,会完全跳过模型,并将输出标记与输入标记一同保存。不过,命中也可能错误,此时它仍会返回正常的成功状态。​