LLM系列：KVCache及优化方法

前言

Transformer encode-base模型，推理和训练过程高度统一（差异仅仅是否存在反向），而decoder-base模型（如GPT、LLama2），推理与训练差异性比较大：

自回归推理
全量prompt+增量token
KV Cache

KV Cache

如下以GPT2结构第i层推理过程为例，分析KV Cahche计算过程，其中WQi，WKi，WVi，WOi，Wupi，WdowniW_Q^i，W_Ki，W_V^i，W_Oi，W_{up}^i，W_{down}iWQi，WKi，WVi，WOi，Wupi，Wdowni表示第i层decoder权重矩阵。

全量prompt阶段

第i层输入xix^ixi（[b,s,h]），self-attention块 xKi，xVi，xQix_K^i，x_Vi，x_Q^ixKi，xVi，xQi（[b,s,h]）

KV Cahche计算：

xKi=xi⋅WKi（[b,s,h]）x_K^i = x^i \cdot W_K^i（[b,s,h]）xKi=xi⋅WKi（[b,s,h]） xVi=xi⋅WVi（[b,s,h]）x_V^i = x^i \cdot W_V^i（[b,s,h]）xVi=xi⋅WVi（[b,s,h]）

第i层attention、mlp计算：

xQi=xi⋅WQi（[b,s,h]）x_Q^i = x^i \cdot W_Q^i（[b,s,h]）xQi=xi⋅WQi（[b,s,h]）

xouti=softmax(xQixKiTh)⋅xVi⋅WOi+xi([b,s,h])x_{out}^{i=softmax(\frac{x_Q}i{x_K^i}T}{\sqrt{h}})\cdot x_V^i \cdot W_O^i + x^i ([b,s,h])xouti=softmax(hxQixKiT)⋅xVi⋅WOi+xi([b,s,h])

xi+1=fgelu(xouti⋅Wupi)⋅Wdowni+xouti([b,s,h])x^{{i+1}=f_{gelu}(x_{out}}i \cdot W_{up}^i) \cdot W_{down}^i + x_{out}^i ([b,s,h])xi+1=fgelu(xouti⋅Wupi)⋅Wdowni+xouti([b,s,h])

增量token阶段

增量推理时，当前生成词在第i层表示ti([b,1,h])t^i ([b,1,h])ti([b,1,h])，推理时执行：更新KV Cache和计算第i层输出。

更新KV Cahce：

xKi←Concat(xKi,ti⋅WKi)([b,s+1,h])x_K^i \leftarrow Concat(x_K^i, t^i \cdot W_K^i) ([b,s+1, h])xKi←Concat(xKi,ti⋅WKi)([b,s+1,h]) xVi←Concat(xVi,ti⋅WVi)([b,s+1,h])x_V^i \leftarrow Concat(x_V^i, t^i \cdot W_V^i) ([b,s+1, h])xVi←Concat(xVi,ti⋅WVi)([b,s+1,h])

第i层计算过程：

tQi=ti⋅WQi([b,1,h])t_Q^i = t^i \cdot W_Q^i ([b,1,h])tQi=ti⋅WQi([b,1,h]) touti=softmax(tQixKiTh)⋅xVi⋅WOi+ti([b,1,h])t_{out}^i = softmax(\frac{t_Q^i{x_Ki}^T}{\sqrt{h}})\cdot x_V^i \cdot W_O^i + t^i ([b,1,h])touti=softmax(htQixKiT)⋅xVi⋅WOi+ti([b,1,h])

ti+1=fgelu(touti⋅Wupi)⋅Wdowni+touti([b,1,h])t^{{i+1}=f_{gelu}(t_{out}}i \cdot W_{up}^i) \cdot W_{down}^i + t_{out}^i ([b,1,h])ti+1=fgelu(touti⋅Wupi)⋅Wdowni+touti([b,1,h])

KV Cache缓存机制如图：

KV Cache显存分析

KV cache的峰值显存占用大小: b(s+n)∗h∗l∗2∗2=4blh(s+n)b(s+n)∗h∗l∗2∗2=4blh(s+n)b(s+n)∗h∗l∗2∗2=4blh(s+n)，输入序列长度s，输出序列长度n，第一个2表示k/v cache，第二个2表示fp16占用2个字节，transformer模型的层数为l，隐藏层维度为h。

以GPT3(175B)为例分析KV Cache与模型参数大小，GPT3模型weight占用350GB(FP16)，层数l=96，维度h=12888。

bs	s+n	kv cache(GB)	kv cache/weight
4	4096	75.5	0.22
16	4096	302	0.86
64	4096	1208	3.45

根据上述数据，随着batch增大和长度增大，KV Cahche开销快速增大，甚至超过模型参数本身。

LLM的窗口长度不断增大，KV Cahche开销随之不断增大，优化KV Cahche非常必要。

KV Cache优化方法

总结典型KV Cache优化手段如下。

1 共用KV Cache（MQA和GQA）

MQA（Multi Query Attention）多查询注意力是MHA多头注意力的变体。两者主要区别是MQA中不同头共享一组KV，每个头只保留查询参数Q。KV矩阵只有一份，大幅减少内存。由于MQA改变注意力机制结构，模型需要从训练开始就支持MQA，或通过对已训练好的模型微调支持MQA，仅需约5%的原始数据量即可达到不过效果。Falcon、SantaCoder、StarCoder 等在内很多模型都采用了MQA机制

# Multi Head Attention
self.Wqkv = nn.Linear(     # Multi-Head Attention 的创建方法
    self.d_model,
    3 * self.d_model,     # Q、K和V 3 个矩阵, 所以是 3 * d_model
    device=device
)
query, key, value = qkv.chunk(3, dim=2)      # 每个 tensor 都是 (1, 512, 768)

# Multi Query Attention
self.Wqkv = nn.Linear(       # Multi-Query Attention 的创建方法
    d_model,
    d_model + 2 * self.head_dim,    # 只创建Q的头向量，所以是 1* d_model, 而K和V不再具备单独的头向量, 所以是 2 * self.head_dim
    device=device,
)
query, key, value = qkv.split(
    [self.d_model, self.head_dim, self.head_dim],    # query -> (1, 512, 768), key   -> (1, 512, 96), value -> (1, 512, 96)
    dim=2
)

GQA (Grouped Query Attention，分组查询注意力)，介于MHQ和MQA之间的折中方案。按查询头Q分组，每个组共享一个K和V。表达能力与推理性能兼顾。

MHA、MQA与GQA原理:

MQA与GQA性能对比：

2 窗口优化

3 量化与稀疏

通过量化与稀疏压缩 KV cache的显存消耗。

稀疏方法典型稀疏方式：

结果显示，KV cache稀疏到只有原来20%时仍然可以保持很高精度。

4 存储与计算优化

FlashDecoding 是在FlashAttention基础上对inference的优化，主要分三步：

（1）长文本下将KV分成更小且方便并行的chunk
（2）对每个chunk的KV，Q和他们进行之前一样的FlashAttention获取这个chunk的结果
（3）对每个chunk的结果进行reduce gif图如下：

如何系统的去学习大模型LLM ？

作为一名热心肠的互联网老兵，我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。

但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的 AI大模型资料 包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

😝有需要的小伙伴，可以V扫描下方二维码免费领取🆓

一、全套AGI大模型学习路线

AI大模型时代的学习之旅：从基础到前沿，掌握人工智能的核心技能！

二、640套AI大模型报告合集

这套包含640份报告的合集，涵盖了AI大模型的理论研究、技术实现、行业应用等多个方面。无论您是科研人员、工程师，还是对AI大模型感兴趣的爱好者，这套报告合集都将为您提供宝贵的信息和启示。

三、AI大模型经典PDF籍

随着人工智能技术的飞速发展，AI大模型已经成为了当今科技领域的一大热点。这些大型预训练模型，如GPT-3、BERT、XLNet等，以其强大的语言理解和生成能力，正在改变我们对人工智能的认识。那以下这些PDF籍就是非常不错的学习资源。

在这里插入图片描述

四、AI大模型商业化落地方案

阶段1：AI大模型时代的基础理解

目标：了解AI大模型的基本概念、发展历程和核心原理。
内容：
- L1.1 人工智能简述与大模型起源
- L1.2 大模型与通用人工智能
- L1.3 GPT模型的发展历程
- L1.4 模型工程
- L1.4.1 知识大模型
- L1.4.2 生产大模型
- L1.4.3 模型工程方法论
- L1.4.4 模型工程实践
- L1.5 GPT应用案例

阶段2：AI大模型API应用开发工程

目标：掌握AI大模型API的使用和开发，以及相关的编程技能。
内容：
- L2.1 API接口
- L2.1.1 OpenAI API接口
- L2.1.2 Python接口接入
- L2.1.3 BOT工具类框架
- L2.1.4 代码示例
- L2.2 Prompt框架
- L2.2.1 什么是Prompt
- L2.2.2 Prompt框架应用现状
- L2.2.3 基于GPTAS的Prompt框架
- L2.2.4 Prompt框架与Thought
- L2.2.5 Prompt框架与提示词
- L2.3 流水线工程
- L2.3.1 流水线工程的概念
- L2.3.2 流水线工程的优点
- L2.3.3 流水线工程的应用
- L2.4 总结与展望

阶段3：AI大模型应用架构实践

目标：深入理解AI大模型的应用架构，并能够进行私有化部署。
内容：
- L3.1 Agent模型框架
- L3.1.1 Agent模型框架的设计理念
- L3.1.2 Agent模型框架的核心组件
- L3.1.3 Agent模型框架的实现细节
- L3.2 MetaGPT
- L3.2.1 MetaGPT的基本概念
- L3.2.2 MetaGPT的工作原理
- L3.2.3 MetaGPT的应用场景
- L3.3 ChatGLM
- L3.3.1 ChatGLM的特点
- L3.3.2 ChatGLM的开发环境
- L3.3.3 ChatGLM的使用示例
- L3.4 LLAMA
- L3.4.1 LLAMA的特点
- L3.4.2 LLAMA的开发环境
- L3.4.3 LLAMA的使用示例
- L3.5 其他大模型介绍