模型推理如何利用非前缀缓存

背景知识

大模型推理阶段

transformer

先看下transformer架构图,来了解到大模型推理在该架构的大概定位

该图的左侧,我们可以这样来理解他的作用:为了让模型能够更好的了解prompt(即对于模型来说,prompt的语义怎么能更好的表示)

该图的右侧,就是模型推理阶段,该阶段每次产生token的类型是自回归生成的,即模型生成一个token会基于输入prompt的tokens 和 已生成的tokens来推理出一个新的token

由此我们知道了,该transformer架构图的右侧,是大模型推理阶段

推理流程

  1. Input:用户输入的prompt(文本)
  2. Tokenization:将文本转换为模型可处理的token
  3. Prefill:一次性处理所有输入token,计算每个token KV,并存入KV Cache中(注意:这一步利用GPU并行性来加速计算过程)
  4. KV Cache:空间换时间的思想,通过缓存每个token的kv(transformer注意力机制的中间结果),来避免重复计算,从而节省时间
  5. Decode:逐一生成新的token,每次用当前的tokens(历史tokens + 已生成)来预测下一个token,同时将新的token的kv加入缓存中(供下一轮用),一直循环,直到满足停止条件(生成 或 达到长度限制)
  6. De-tokenization:将生成的token序列转换为人类可读的文本

至此,我们知道了模型推理的大概流程。此外,由于我们的主题是和KV Cache有关的,所以在这里我们需要先详细了解一下KV Cache(由于KV Cache是用于加速自注意力机制的计算,所以会结合自注意力机制计算来讲)

自注意力机制 & KV Cache & 前缀缓存

自注意力机制

在注意力机制计算之前,我们要先得到该token的q、k、v

以上的

  • X代表该token经过位置编码后的词嵌入向量
  • W_Q、W_K、W_V,是模型在训练过程中学到的参数矩阵,它们决定了如何通过X得到QKV
  • Q、K、V:
    • Q代表:代表当前 token 在"提问"时所关注的信息,它会去匹配所有历史 token 的 Key,寻找与自己最相关的内容。可以理解为:"我现在在想什么?我需要哪些上下文来帮助我?"
    • K代表:代表每个已生成 token 对自身内容的"摘要"或"索引"。
    • V代表:代表每个 token 所承载的实际语义信息,当 Query 找到匹配的 Key 后,就会从对应的 Value 中提取内容,作为最终输出的一部分。可以理解为:"我不仅告诉你'我在',还告诉你'我是什么意思'。"

至此,我们知道了一个token的Q、K、V三个矩阵是如何得到的,接下来,我们来看,这三个矩阵是如何在自注意力机制中发挥自身作用的。

上图是自注意力机制的计算的主要逻辑:

我们从Score这行开始看(对于score之前行,就是得到每个token的Q、K、V):

q1 . k1、q2 . k2:这个就是该token的q,需要 去和所有历史token的k做匹配,得到该token与历史token的相似度,也可以这样理解。首先,Q 代表当前 token 在当前上下文中的"信息需求",然后它通过匹配历史 token 的 Key,来决定从对应的 Value 中提取多少语义内容,从而完成上下文感知的表示更新。

Divide by (d_k) ^1/2:缩放点积,避免softmax梯度接近0

softmax函数:

当输入x_i时候,todo

softmax(归一化):将每个token的相似度,变成概率

softmax*value:有了概率列表,就对历史token的v,加权求和

z:这个就是我们的注意力的结果啦,注意它是一个融合了上下文信息的新表示

此处,小插曲,为什么要用点积计算呢?(todo)

KV Cache

至此,我们知道了一个token的qkv是如何产生的,也知道了qkv在自注意力机制中发挥了重要作用,经过自注意力机制的计算,使每个token都有一个融合上下文信息的新的表示,现在我们来思考一下,以上步骤有可优化之处吗?

我们通过详细看下注意力机制计算的公式找出优化点:

当模型生成第一个"遥"字时,input="", ""是起始字符。Attention的计算如下:

Attention的计算公式如下:

当模型生成第二个"遥"字时,input="遥", Attention的计算如下:

Attention的计算公式如下:

当模型生成第三个"领"字时,input="遥遥"Attention的计算如下:

Attention的计算公式如下:

我们发现,当生成一个新toke时,历史token的k,v也是要参与的,那么我们想要用历史token的k、v,每次用它们时候都要算出来吗?所以我们可以用空间换时间的思想,将历史token的k、v存起来,我们需要用的时候直接取就行,这样一来,矩阵相乘的次数就少了,那么整体推理时间也会降低,那么这就是KV Cache

因此,我们可以选择缓存历史token的kv,如下图:


至此,我们知道了一个token的qkv是如何产生的,也知道了各token的qkv通过参与自注意力机制的计算,使每个token都有一个融合上下文信息的新的表示,以及加速注意力机制计算的KV Cache

前缀缓存

为避免重复计算token的kv,所以在 prefill 阶段主要作用就是给迭代的生成阶段准备 KV Cache。但这些 KV Cache 仅仅是为单次生成式请求服务的,那很自然的一种想法就是,KV Cache 能不能跨请求复用?

在某些场景下,多次请求的 Prompt 可能会共享同一个前缀(Prefix),比如system prompt,这种情况下,请求的前缀的 KV Cache 计算的结果是相同的,所以可以被缓存起来,给下一个请求复用。

但 KV Cache 跟其它服务缓存不一样的地方是,它太大了,以至于(目前)很难通过 Redis/Memcache 这种分布式缓存服务存取。比如对 13B LLM 模型来说,在 FP16 精度下单 token 的 KV Cache 大约是 1MB,假设要缓存的前缀有 500 个 token(大约800多个汉字),那就是 500MB。一般来说,我们不会每次请求去从分布式系统里读取/传输 500MB 的缓存,甚至都不会每次请求从内存往显存中拷贝 500MB 的缓存,所以大部分情况下,prefix cache 都会放在显存里。

这也就意味着,如果你想命中 prefix cache,必须把相同 prefix 的请求发到同一张 GPU卡上才行。

具体实现

以sglang中的实现介绍

  • 在每次请求结束后,这次请求里的KV Cache不直接丢弃,而是借由RadixAttention算法将KV Cache仍然保留在GPU的显存中。RadixAttention算法还将序列token到KV Cache的映射关系,保存在RadixTree这种数据结构中。
  • 当一个新的请求到来的时候,SGLang的调度器通过RadixTree做前缀匹配,查找新的请求是否命中缓存。如果命中了缓存,SGLang的调度器会复用KV Cache来完成这次请求。
  • 由于GPU的显存是非常有限的,缓存在内存中的KV Cache不能永久被保留,需要设计相应的逐出策略。RadixAttention引入的逐出策略是LRU(least recently used)

此时,我们知道了用前缀缓存可以加速推理,那非前缀可以吗?我们先来分析一下:

第一层 KV 的生成公式

h是词嵌入 + 位置编码;k是该token的k,v是该token的v

注意:这一步是纯局部的,每个 token 独立计算自己的 QKV,不与其他 token 交互,因此,第一层的 K 和 V 完全由 (token ID, 位置 t) 决定。

自注意力机制计算

自注意力机制大概过程:当前 token 的 q 会 与所有已生成 token(包括自己)的 k 矩阵相乘计算,再经过softmax得到多个系数,再用这些系数对所有已生成的v加权求和,最终得到注意力分数。

自注意力机制是 Transformer 的核心,但单层自注意力只能捕捉 "扁平的关联",无法处理语言中的复杂依赖。

  • 每一层都对前一层的输出进行非线性变换 + 上下文整合
  • 信息被逐步提炼:
    • Layer 1: 词性、局部语法
    • Layer 5: 句法结构、依存关系
    • Layer 10: 语义角色、意图理解
    • Layer 20+: 世界知识、逻辑推理

第一层的自注意力机制计算后,得到h1(可以理解为 对于第二层来说,未变化成k、v的一个矩阵)。

第二层的自注意力机制的k 和 v是有h1 分别和 w_k 和 w_v矩阵得来,但此时的kv是已经有上下文信息的了!

第三层,类推...

所以,即使两个徐请求中,即使相同(词嵌入 + 位置编码都相同)的token也不能利用kv cache,因为会有注意力缺失问题。

那么我们有办法利用非前缀缓存吗?来看Cache blend技术

CacheBlend

...todo

相关推荐
亚马逊云开发者41 分钟前
Amazon OpenSearch 助力高效 RAG 系统落地
人工智能
suke43 分钟前
听说前端又死了?
前端·人工智能·程序员
小小小怪兽1 小时前
关于氛围编程(Vibe Coding)的一切
人工智能·全栈
付玉祥1 小时前
从谷歌白皮书看 Prompt 工程
人工智能
松岛雾奈.2301 小时前
机器学习--数据集的标准化和归一化算法;随机森林
人工智能·算法·机器学习
阿明Drift1 小时前
用 RAG 搭建一个 AI 小说问答系统
前端·人工智能
朱龙凯1 小时前
LangChain学习笔记
人工智能
飞哥数智坊1 小时前
Cursor 2.1 发布实测:计划能点了,审查能用了,CR 花多少?我也替你试了
人工智能·ai编程·cursor
凯子坚持 c1 小时前
Doubao-Seed-Code模型深度剖析:Agentic Coding在Obsidian插件开发中的应用实践
网络·人工智能
iFlow_AI2 小时前
iFlow CLI快速搭建Flutter应用记录
开发语言·前端·人工智能·flutter·ai·iflow·iflow cli