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

背景知识

大模型推理阶段

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

相关推荐
万岳科技程序员小金7 小时前
多商户商城APP源码开发的未来方向:云原生、电商中台与智能客服
人工智能·云原生·开源·软件开发·app开发·多商户商城系统源码·多商户商城app开发
CoookeCola8 小时前
MovieNet (paper) :推动电影理解研究的综合数据集与基准
数据库·论文阅读·人工智能·计算机视觉·视觉检测·database
火星资讯8 小时前
多形态机器人协同发力优艾智合引领核电运维智能化升级
人工智能
qq_420362038 小时前
AI在前端工作中的应用
前端·人工智能·sse
亚马逊云开发者8 小时前
Agentic AI基础设施实践经验系列(一):Agent应用开发与落地实践思考
人工智能
6v6-博客8 小时前
【效率工具】EXCEL批注提取工具
人工智能
晨非辰8 小时前
《数据结构风云》:二叉树遍历的底层思维>递归与迭代的双重视角
数据结构·c++·人工智能·算法·链表·面试
JJJJ_iii8 小时前
【机器学习12】无监督学习:K-均值聚类与异常检测
人工智能·笔记·python·学习·机器学习·均值算法·聚类
DogDaoDao9 小时前
OpenCV音视频编解码器详解
人工智能·opencv·音视频·视频编解码·h264·h265·音视频编解码