背景

参考笔记

• (万字长文)说说大模型中的推理加速技术

推理背景

预填充和解码阶段

• 预填充阶段

  • 输入prompt做前向计算，生成第一个token

  • TTFT指标 (Time To First Token)：生成第1个token所需的时间

• 解码阶段

  • 逐个生成剩余token，直到达到EOS或最大长度
  • TPOT(Time Per Output Token)

提高效率方法

• 批处理Batching

  • 大批数据传到GPU，做一次处理，提高GPU利用率。
    • 静态批处理：等一批中最长的生成完成才结束，差异会有问题。
  • 太大了会导致内存溢出，可以考虑KV Cache/LLM内存要求。

• KV 缓存

  • 生成新token，依赖之前所有token的key/value，但这部分不用重新计算，可以缓存在GPU
  • 如llama2-7B模型，fp16精度，b=1，seq长度4096，32层，hsize4096
    • KV缓存大小：1*4096*2*32*4096*2=2GB，其中
  $$\text{总KV缓存大小}=b\ast s\ast 2\ast \text{n\_layers}\ast \text{h\_size}\ast \mathrm{精}\mathrm{度}$$

• LLM 内存需求

  • 权重：fp16，2m, 7B，14GB显存
  • KV缓存

模型并行化

具体：模型并行

模型并行化

Pipeline 并行

• 按层并行

张量并行

• 层内并行

序列并行

• 按长度切分为多块

Attention 优化技术

结构优化 MHA/MQA/GQA/MLA

注意力机制优化

MHA

MQA

• 所有Query共用一对Key和Value，减少KV数量来降低缓存的内存

GQA

• 按Query分组，每组共用一对Key和Value，减少KV数量来降低缓存的内存

MLA

FlashAttention

FlashAttention

FlashAttention

• 分块计算softmax tiling算法
• 反向不使用中间注意力矩阵，通过只存softmax归一化系数来降低HBM内存消耗

KV Cache技术

参考文章

• 大模型推理加速：看图学KV Cache
• 图解 KV Cache 过程
• Attention解码详细推导

KV Cache

KV Cache

背景

• 每一步的Attention计算存在大量冗余，通过推导公式可知，第k步的完整$Att{n}_{stepk}$矩阵：
  • 前k-1行：已由前面的[1, k-1]步骤计算过
  • 第k行：$Att{n}_k$只和$Q_k$、前面的每一对$K_i,V_i$有关，与其余$Q$均无关
  • 每一行，只和对应步骤的Q有关

方法

• 解码新token时，缓存前面每一步的$K_i、V_i、Att{n}_i$ 结果
• 只需用当前Q和前面每一步的K、V来计算$Att{n}_k$ 即可，属于增量QKV计算

优点

• 大大减少Time Per Output Token，减少了GPU在一次推理中的运算量
• 注意：对预填充TTFT没有影响

缺点

• 做推理占用更多显存，会限制序列长度和BatchSize，影响吞吐量
• 花费更多时间在读取缓存数据上。
  • 时间：读1MB数据到GPU >> 在GPU上算1MB数据
• 具有较多碎片
  • 内部碎片：达到最大长度前，结束了，实际长度小于最大长度
  • 外部碎片：对每个请求预先分配的显存空间大小不一样，易产生外部碎片
• 仅在1个请求内复用，无法做到跨请求，不适用于多轮对话场景

对比示意图，一目了然，每次只用当前Q和前的KV做计算即可，因为前面的行和$Q_k$无关，对应行只和对应步骤的Q有关。

KV Cache 显存分析

显存分析

参数定义

• 每个batch：b个序列
• 每个序列：t个token，包括prompt和completion
• 每个token：需缓存key和value 2类张量
• 每个decoder：包含n_layers注意力层
• 每个多头注意力：n_heads个头
• 注意力向量：d_head向量维度
• 精度：1参数p_a字节，fp32-4，fp16/bf16-2，fp8/int8-1

KV Cache大小 (单token)

• MHA

$$2\ast n_{layers}\ast n_{head}\ast d_{head}\ast p_a$$

• MQA：所有query共用1个head，减少了注意力头数$n_{head}$
  • 效果不行，无法做到有效多卡并行

$$2\ast n_{layers}\ast 1\ast d_{head}\ast p_a$$

• GQA：query分为$G$组，每组共用1个head

$$2\ast n_{layers}\ast G\ast d_{head}\ast p_a$$

KV Cache大小 (整个batch) MHA

$$b\ast t\ast 2\ast n_{layers}\ast n_{head}\ast d_{head}\ast p_a$$

Paged Attention 显存管理

Paged Attention

背景

• 显存存在较多碎片
  • 内部碎片：按固定最大长度去缓存KV，空间连续存储，但实际长度小于最大长度
  • 外部碎片：对每个请求预先分配的显存空间大小不一样，易产生外部碎片

方法

• 参考OS的分页机制
• 把每个生成序列的KV Cache划分为多个块，不连续存储，每块包含固定数量的key和value
• 注意力计算时，使用block table来找到块，进而取出k和v
• 新token产生时，进行新的区块分配，块固定大小

优点

• 降低了内存浪费；提升更大Batch

碎片很多：

PageAttention 运行机制

Copy-On-Write 解码技术

Copy-On-Write

背景

• 并行采样：1个prompt一次性生成多个不同的输出序列
• Beam Search的多个输出可能存在相同的前缀，可以共享KV Cache来降低显存占用

方法

• 当多个请求/序列需要共享相同的KV Cache块时，这些块会被映射到相同的物理块
• vLLM写时复制：当一个序列会改变KV Cache时，vLLM会先创建副本，这样修改就不会影响其他序列

Prefix Caching 技术

Prefix Caching

背景

• PagedAttention的KV Cache仅在1个请求内复用，无法做到跨请求，不适用于多轮对话场景

核心思想

• 跨请求KV Cache
  • 新一轮prompt的KV Cache复用上一轮结果，提升prefix性能，降低新一轮的TTFT
  • 注意：无法提升TPOP，仅能提升TTFT
• 缓存系统提示和历史对话中的KV缓存，以便在后续请求中重用

SGLang Prefix Caching

• 每次请求结束后KV Cache不直接丢弃，而是由RadixAttention算法保留在GPU显存中
• 新请求到来时使用RadixTree做前缀匹配，命中缓存则复用
• 通过LRU(Least Recenty Used)来逐出未被使用的KV Cache

vLLM Prefix Caching

• 思路估计和SGLang一致，只是采用Hash RadixAttention的方法
• hash(prefix tokens + block tokens) <--> Logical KV blocks -> Physical KV blocks

DeepSeek Prefix Caching / Context Caching

• 采用上下文硬盘缓存技术
  • 把预计未来会重复使用的内容，缓存在分布式硬盘阵列，而非GPU中
  • 64tokens为一个存储单元，不足不会被缓存
  • 命中缓存：0.1元；未命中缓存：1元 (百万tokens)
• MLA 结构：大大压缩了KV Cache的大小
  • 传输和存储量均大大减少

模型优化技术

量化

量化

核心思想

• 降低模型权重和激活精度，模型训练32/16 -> 存储8/4等。
• 舍入会丢失一些精度，剪裁丢失一些动态范围

优点

• 内存占用减少，可以运行更大模型

稀疏

稀疏

核心思想

• 使用0来代替一些接近0的值
• 对很多元素为0的稀疏矩阵，可以用压缩形式表示

蒸馏

蒸馏

核心思想

• 把大模型知识蒸馏到小模型

模型服务技术

模型执行通常受内存带宽限制(特别是权重)，期望加载权重时尽可能多的处理他们。

动态批处理

动态批处理

思想

• 同时执行多个不同的请求。
• 会立即剔除批处理中已完成的序列，而不是等待整批完成后再进行下一组请求。

预测推理

预测推理/推测采样/辅助生成

思想

• 使用多个便宜的、小的临时模型，并行生成多个token的长序列
• 主模型并行验证生成的token，决定是否接受
• 并行执行序列的多个不同步骤以尝试节省时间。