LLM - GPT(Decoder Only) 类模型的 KV Cache 公式与原理 教程

2024-08-28 12:12

本文主要是介绍LLM - GPT(Decoder Only) 类模型的 KV Cache 公式与原理 教程,希望对大家解决编程问题提供一定的参考价值,需要的开发者们随着小编来一起学习吧!

欢迎关注我的CSDN:https://spike.blog.csdn.net/
本文地址:https://spike.blog.csdn.net/article/details/141605718

免责声明:本文来源于个人知识与公开资料,仅用于学术交流,欢迎讨论,不支持转载。


Img

在 GPT 类模型中,KV Cache (键值缓存) 是用于优化推理效率的重要技术,基本思想是通过缓存先前计算的 键(Key) 和 值(Value),避免在推理过程中,重复计算 Mask 的 注意力(Attention) 矩阵,从而加速生成过程。

1. 公式

矩阵乘法的基础性质:

A ⋅ B = [ A 1 A 2 … A n ] ⋅ [ B 1 B 2 ⋮ B n ] = A 1 B 1 + A 2 B 2 + ⋯ + A n B n A \cdot B = \begin{bmatrix} A_{1} & A_{2} & \dots & A_{n} \end{bmatrix} \cdot \begin{bmatrix} B_{1} \\ B_{2} \\ \vdots \\ B_{n} \end{bmatrix} = A_{1}B_{1} + A_{2}B_{2} + \dots + A_{n}B_{n} AB=[A1A2An] B1B2Bn =A1B1+A2B2++AnBn

其中 A i A_{i} Ai A A A 的列向量, B i B_{i} Bi B B B 的行向量,也就是说相同维度的向量相乘,可拆解成行向量乘以列向量,即 A A A n n n 列, B B B n n n 行。如图:

matrix

例如:基础的矩阵乘法:

A = [ 1 2 3 4 ] , B = [ 5 6 7 8 ] C = [ 1 ∗ 5 + 2 ∗ 7 1 ∗ 6 + 2 ∗ 8 3 ∗ 5 + 4 ∗ 7 3 ∗ 6 + 4 ∗ 8 ] = [ 19 22 43 50 ] A = \begin{bmatrix} 1 & 2 \\ 3 & 4 \end{bmatrix}, \quad B = \begin{bmatrix} 5 & 6 \\ 7 & 8 \end{bmatrix} \\ C = \begin{bmatrix} 1*5 + 2*7 & 1*6 + 2*8 \\ 3*5 + 4*7 & 3*6 + 4*8 \end{bmatrix} = \begin{bmatrix} 19 & 22 \\ 43 & 50 \end{bmatrix} A=[1324],B=[5768]C=[15+2735+4716+2836+48]=[19432250]

也可以写成,行列向量相乘的形式,即 A 拆分出多个行向量,B 拆分出多个列向量,即:

C = [ 1 3 ] ⋅ [ 5 6 ] + [ 2 4 ] ⋅ [ 7 8 ] = [ 1 ∗ 5 1 ∗ 6 3 ∗ 5 3 ∗ 6 ] + [ 2 ∗ 7 2 ∗ 8 4 ∗ 7 4 ∗ 8 ] C = \begin{bmatrix} 1 \\ 3 \end{bmatrix} \cdot \begin{bmatrix} 5 & 6 \end{bmatrix} + \begin{bmatrix} 2 \\ 4 \end{bmatrix} \cdot \begin{bmatrix} 7 & 8 \end{bmatrix}= \begin{bmatrix} 1*5 & 1*6 \\ 3*5 & 3*6 \end{bmatrix} + \begin{bmatrix} 2*7 & 2*8 \\ 4*7 & 4*8 \end{bmatrix} C=[13][56]+[24][78]=[15351636]+[27472848]
= [ 1 ∗ 5 + 2 ∗ 7 1 ∗ 6 + 2 ∗ 8 3 ∗ 5 + 4 ∗ 7 3 ∗ 6 + 4 ∗ 8 ] = [ 19 22 43 50 ] =\begin{bmatrix} 1*5 + 2*7 & 1*6 + 2*8 \\ 3*5 + 4*7 & 3*6 + 4*8 \end{bmatrix} = \begin{bmatrix} 19 & 22 \\ 43 & 50 \end{bmatrix} =[15+2735+4716+2836+48]=[19432250]

进一步拆解:

A ⋅ B = A 1 B 1 + A 2 B 2 + ⋯ + A n B n = [ a 1 , 1 B 1 a 2 , 1 B 2 ⋮ a m , 1 B n ] + [ a 1 , 2 B 1 a 2 , 2 B 2 ⋮ a m , 2 B n ] + ⋯ + [ a 1 , n B 1 a 2 , n B 2 ⋮ a m , n B n ] = [ a 1 , 1 B 1 + a 1 , 2 B 1 + ⋯ + a 1 , n B 1 a 2 , 1 B 2 + a 2 , 2 B 2 + ⋯ + a 2 , n B 2 ⋯ a m , 1 B n + a m , 2 B n + ⋯ + a m , n B n ] A \cdot B = A_{1}B_{1} + A_{2}B_{2} + \dots + A_{n}B_{n} \\ = \begin{bmatrix} a_{1,1}B_{1} \\ a_{2,1}B_{2} \\ \vdots \\ a_{m,1}B_{n} \end{bmatrix} + \begin{bmatrix} a_{1,2}B_{1} \\ a_{2,2}B_{2} \\ \vdots \\ a_{m,2}B_{n} \end{bmatrix} + \cdots + \begin{bmatrix} a_{1,n}B_{1} \\ a_{2,n}B_{2} \\ \vdots \\ a_{m,n}B_{n} \end{bmatrix} \\ = \begin{bmatrix} a_{1,1}B_{1} + a_{1,2}B_{1} + \cdots + a_{1,n}B_{1} \\ a_{2,1}B_{2} + a_{2,2}B_{2} + \cdots + a_{2,n}B_{2} \\ \cdots \\ a_{m,1}B_{n} + a_{m,2}B_{n} + \cdots + a_{m,n}B_{n} \end{bmatrix} AB=A1B1+A2B2++AnBn= a1,1B1a2,1B2am,1Bn + a1,2B1a2,2B2am,2Bn ++ a1,nB1a2,nB2am,nBn = a1,1B1+a1,2B1++a1,nB1a2,1B2+a2,2B2++a2,nB2am,1Bn+am,2Bn++am,nBn

基础的矩阵乘法的另一种形式:

C = [ 1 3 ] ⋅ [ 5 , 6 ] + [ 2 4 ] ⋅ [ 7 , 8 ] C=\begin{bmatrix} 1 \\ 3 \end{bmatrix} \cdot \begin{bmatrix} 5,6 \end{bmatrix} + \begin{bmatrix} 2 \\ 4 \end{bmatrix} \cdot \begin{bmatrix} 7,8 \end{bmatrix} C=[13][5,6]+[24][7,8]
[ 1 ∗ [ 5 6 ] 3 ∗ [ 5 6 ] ] + [ 2 ∗ [ 7 8 ] 4 ∗ [ 7 8 ] ] \begin{bmatrix} 1*[5&6] \\ 3*[5&6] \end{bmatrix} + \begin{bmatrix} 2*[7&8] \\ 4*[7&8] \end{bmatrix} [1[53[56]6]]+[2[74[78]8]]
[ 1 ∗ 5 1 ∗ 6 3 ∗ 5 3 ∗ 6 ] + [ 2 ∗ 7 2 ∗ 8 4 ∗ 7 4 ∗ 8 ] = [ 19 22 43 50 ] \begin{bmatrix} 1*5 & 1*6 \\ 3*5 & 3*6 \end{bmatrix} + \begin{bmatrix} 2*7 & 2*8 \\ 4*7 & 4*8 \end{bmatrix} = \begin{bmatrix} 19 & 22 \\ 43 & 50 \end{bmatrix} [15351636]+[27472848]=[19432250]

如果 A A A 是下三角矩阵,即包含 Mask 信息,Decoder 无法观察到之后的推理部分,则 A ⋅ B A \cdot B AB,输出:

A ⋅ B = [ a 1 , 1 B 1 a 2 , 1 B 2 + a 2 , 2 B 2 ⋯ a m , 1 B n + a m , 2 B n + ⋯ + a m , n B n ] A \cdot B = \left[ \begin{array}{llll} a_{1,1}B_{1}\\ a_{2,1}B_{2} + a_{2,2}B_{2}\\ \cdots \\ a_{m,1}B_{n} + a_{m,2}B_{n} + \cdots + a_{m,n}B_{n} \end{array} \right] AB= a1,1B1a2,1B2+a2,2B2am,1Bn+am,2Bn++am,nBn

2. 推理

第1步:

在 Decoder 解码过程中,只关注 Transformer 的 自注意力(Self-Attention),输入第 1 个 Token,将 Token 转换成 输入特征 I n p u t 1 = [ 1 , d e m b ] Input_{1}=[1,d_{emb}] Input1=[1,demb],暂时忽略 batch_size d e m b d_{emb} demb 表示 Embedding Size。

  1. 输入特征 I n p u t 0 = [ 1 , d e m b ] Input_{0}=[1,d_{emb}] Input0=[1,demb],乘以权重 W = [ d e m b , 3 ∗ d e m b ] W=[d_{emb}, 3*d_{emb}] W=[demb,3demb] (已训练完成,值是固定的),输出维度 [ 1 , 3 ∗ d e m b ] [1, 3*d_{emb}] [1,3demb],即作为 Q\K\V,每个向量 [ 1 , d e m b ] [1,d_{emb}] [1,demb]

    • Q 1 = [ 1 , d e m b ] Q_{1}=[1,d_{emb}] Q1=[1,demb] K 1 = [ 1 , d e m b ] K_{1}=[1,d_{emb}] K1=[1,demb] V 1 = [ 1 , d e m d ] V_{1}=[1,d_{emd}] V1=[1,demd],只与输入特征 I n p u t 0 Input_{0} Input0 的 Embedding 相关。
  2. 根据 Self-Attention 的公式,忽略 d \sqrt{d} d ,只有1维,mask 不起作用,即
    A t t ( Q , K , V ) = s o f t m a x ( Q K ⊤ + m a s k ) ∗ V A t t 1 ( Q , K , V ) = s o f t m a x ( Q 1 K 1 ⊤ ) V 1 其中  s o f t m a x ( x i ) = e x i ∑ j = 1 n e x j Att(Q,K,V)=softmax(QK^{\top}+mask)*V \\ Att_{1}(Q,K,V)=softmax(Q_{1}K_{1}^{\top})V_{1} \\ 其中 \ softmax(x_i) = \frac{e^{x_i}}{\sum_{j=1}^{n} e^{x_j}} Att(Q,K,V)=softmax(QK+mask)VAtt1(Q,K,V)=softmax(Q1K1)V1其中 softmax(xi)=j=1nexjexi

  3. A t t 0 Att_{0} Att0 ( [ 1 , d e m b ] [1,d_{emb}] [1,demb]) 经过一系列推理,最后输出 [ 1 , d v ] [1, d_{v}] [1,dv] d v d_{v} dv 是全部词元 Token 的数量,根据概率值即可获得最后的 Token。

第 2 步

将第 1 步输出的 Token 转换成 [ 1 , d e m b ] [1,d_{emb}] [1,demb],与第 1 步组合至一起,即 输入特征 I n p u t 1 = [ 2 , d e m b ] Input_{1}=[2,d_{emb}] Input1=[2,demb]

  1. 输入特征 I n p u t 1 = [ 2 , d e m b ] Input_{1}=[2,d_{emb}] Input1=[2,demb],乘以权重 W = [ d e m b , 3 ∗ d e m b ] W=[d_{emb}, 3*d_{emb}] W=[demb,3demb],权重是固定的,因此只需要计算第 2 个输入的特征 [ 1 , d e m b ] [1,d_{emb}] [1,demb],第 1 个不需要计算,也就是说 Q\K\V 的维度是 [ 2 , d e m b ] [2, d_{emb}] [2,demb],只需计算一次即可,剩余的可以直接 c o n c a t concat concat 到一起。

  2. 根据 Self-Attention 的公式,忽略 d \sqrt{d} d ,注意第1行,已经计算,第2行,需要使用 Q 2 Q_{2} Q2 K 2 K_{2} K2 V 2 V_{2} V2,进行计算,即:
    A t t 2 ( Q , K , V ) = s o f t m a x ( Q K ⊤ + m a s k ) ∗ V s o f t m a x ( [ Q 1 K 1 ⊤ Q 2 K 1 ⊤ + Q 2 K 2 ⊤ ] ) ⋅ [ V 1 V 2 ] = [ s o f t m a x ( Q 1 K 1 ⊤ ) V 1 s o f t m a x ( Q 2 K 1 ⊤ ) V 1 + s o f t m a x ( Q 2 K 2 ⊤ ) V 2 ] = [ A t t 1 ( Q , K , V ) s o f t m a x ( Q 2 K 1 ⊤ ) V 1 + s o f t m a x ( Q 2 K 2 ⊤ ) V 2 ] Att_{2}(Q,K,V) = softmax(QK^{\top}+mask)*V \\ softmax(\left[ \begin{array}{ll} Q_{1}K_{1}^{\top}\\ Q_{2}K_{1}^{\top} + Q_{2}K_{2}^{\top}\\ \end{array} \right]) \cdot \begin{bmatrix} V_{1} \\ V_{2} \\ \end{bmatrix} \\= \left[ \begin{array}{ll} softmax(Q_{1}K_{1}^{\top})V_{1}\\ softmax(Q_{2}K_{1}^{\top})V_{1} + softmax(Q_{2}K_{2}^{\top})V_{2}\\ \end{array} \right] \\ = \left[ \begin{array}{} Att_{1}(Q,K,V) \\ softmax(Q_{2}K_{1}^{\top})V_{1} + softmax(Q_{2}K_{2}^{\top})V_{2}\\ \end{array} \right] Att2(Q,K,V)=softmax(QK+mask)Vsoftmax([Q1K1Q2K1+Q2K2])[V1V2]=[softmax(Q1K1)V1softmax(Q2K1)V1+softmax(Q2K2)V2]=[Att1(Q,K,V)softmax(Q2K1)V1+softmax(Q2K2)V2]

  3. KV 都是成对出现的,如果 缓存 KV,则可以加快推理速度。

第 3 步:重复进行。

3. 缓存占用

关于 Llama3 的 KV Cache 源码,参考 model.py:

xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)
xk = xk.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)
xv = xv.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)self.cache_k = self.cache_k.to(xq)
self.cache_v = self.cache_v.to(xq)self.cache_k[:bsz, start_pos : start_pos + seqlen] = xk
self.cache_v[:bsz, start_pos : start_pos + seqlen] = xvkeys = self.cache_k[:bsz, : start_pos + seqlen]
values = self.cache_v[:bsz, : start_pos + seqlen]

关于 KV 的缓存内存占用:

相关参数 batch_size=32head=32layer=32dim_size=4096seq_length=2048,float32(4个字节)类,计算 KV cache 的缓存占用:
M = 2 ∗ N b s ∗ ( N d i m / N h e a d ∗ N h e a d ) ∗ N l a y e r ∗ N s e q ∗ 4 = 2 ∗ 32 ∗ 4096 ∗ 32 ∗ 2048 ∗ 4 / 1024 / 1024 / 1024 = 64 G M=2*N_{bs}*(N_{dim}/N_{head}*N_{head})*N_{layer}*N_{seq}*4 \\ =2*32*4096*32*2048*4/1024/1024/1024=64G M=2Nbs(Ndim/NheadNhead)NlayerNseq4=23240963220484/1024/1024/1024=64G
也就是说 head 数量无关,因为维度除以 Head 再乘以 Head。Llama3 使用 GQA (Grouped Query Attention) 分组查询注意力机制,降低 4 倍的 KV Cache,head=32,kv_head=8,即 scale=head/kv_head=4

参考:

  • CSDN - 从头开始实现 LLaMA3 的网络结构与推理流程 教程
  • Transformers KV Caching Explained

这篇关于LLM - GPT(Decoder Only) 类模型的 KV Cache 公式与原理 教程的文章就介绍到这儿,希望我们推荐的文章对编程师们有所帮助!



http://www.chinasem.cn/article/1114735

相关文章

Maven的使用和配置国内源的保姆级教程

《Maven的使用和配置国内源的保姆级教程》Maven是⼀个项目管理工具,基于POM(ProjectObjectModel,项目对象模型)的概念,Maven可以通过一小段描述信息来管理项目的构建,报告... 目录1. 什么是Maven?2.创建⼀个Maven项目3.Maven 核心功能4.使用Maven H

IDEA自动生成注释模板的配置教程

《IDEA自动生成注释模板的配置教程》本文介绍了如何在IntelliJIDEA中配置类和方法的注释模板,包括自动生成项目名称、包名、日期和时间等内容,以及如何定制参数和返回值的注释格式,需要的朋友可以... 目录项目场景配置方法类注释模板定义类开头的注释步骤类注释效果方法注释模板定义方法开头的注释步骤方法注

Spring Boot循环依赖原理、解决方案与最佳实践(全解析)

《SpringBoot循环依赖原理、解决方案与最佳实践(全解析)》循环依赖指两个或多个Bean相互直接或间接引用,形成闭环依赖关系,:本文主要介绍SpringBoot循环依赖原理、解决方案与最... 目录一、循环依赖的本质与危害1.1 什么是循环依赖?1.2 核心危害二、Spring的三级缓存机制2.1 三

C#中async await异步关键字用法和异步的底层原理全解析

《C#中asyncawait异步关键字用法和异步的底层原理全解析》:本文主要介绍C#中asyncawait异步关键字用法和异步的底层原理全解析,本文给大家介绍的非常详细,对大家的学习或工作具有一... 目录C#异步编程一、异步编程基础二、异步方法的工作原理三、代码示例四、编译后的底层实现五、总结C#异步编程

Python虚拟环境终极(含PyCharm的使用教程)

《Python虚拟环境终极(含PyCharm的使用教程)》:本文主要介绍Python虚拟环境终极(含PyCharm的使用教程),具有很好的参考价值,希望对大家有所帮助,如有错误或未考虑完全的地方,... 目录一、为什么需要虚拟环境?二、虚拟环境创建方式对比三、命令行创建虚拟环境(venv)3.1 基础命令3

使用Node.js制作图片上传服务的详细教程

《使用Node.js制作图片上传服务的详细教程》在现代Web应用开发中,图片上传是一项常见且重要的功能,借助Node.js强大的生态系统,我们可以轻松搭建高效的图片上传服务,本文将深入探讨如何使用No... 目录准备工作搭建 Express 服务器配置 multer 进行图片上传处理图片上传请求完整代码示例

Go 语言中的select语句详解及工作原理

《Go语言中的select语句详解及工作原理》在Go语言中,select语句是用于处理多个通道(channel)操作的一种控制结构,它类似于switch语句,本文给大家介绍Go语言中的select语... 目录Go 语言中的 select 是做什么的基本功能语法工作原理示例示例 1:监听多个通道示例 2:带

鸿蒙中@State的原理使用详解(HarmonyOS 5)

《鸿蒙中@State的原理使用详解(HarmonyOS5)》@State是HarmonyOSArkTS框架中用于管理组件状态的核心装饰器,其核心作用是实现数据驱动UI的响应式编程模式,本文给大家介绍... 目录一、@State在鸿蒙中是做什么的?二、@Spythontate的基本原理1. 依赖关系的收集2.

python连接本地SQL server详细图文教程

《python连接本地SQLserver详细图文教程》在数据分析领域,经常需要从数据库中获取数据进行分析和处理,下面:本文主要介绍python连接本地SQLserver的相关资料,文中通过代码... 目录一.设置本地账号1.新建用户2.开启双重验证3,开启TCP/IP本地服务二js.python连接实例1.

Python 安装和配置flask, flask_cors的图文教程

《Python安装和配置flask,flask_cors的图文教程》:本文主要介绍Python安装和配置flask,flask_cors的图文教程,本文通过图文并茂的形式给大家介绍的非常详细,... 目录一.python安装:二,配置环境变量,三:检查Python安装和环境变量,四:安装flask和flas