LLM系列：LLM推理加速技术

Albert_Liyan
807 阅读6分钟

本文梳理主要LLM推理加速技术，讨论加速原理及发展方向。

LLM推理关注的性能指标有哪些？

  • 吞吐量 吞吐量的定义：单位时间内处理的token数量。（输入序列长度 + 输出序列长度）/ 耗时。

  • 时延 时延定义：平均收到每个token的时间。用户发出请求到收到响应所需时间除以生成序列长度。

    吞吐量和时延是两个不同性能指标。吞吐量越高表示系统单位时间处理请求越多，系统利用率高；时延越高表示推理越快，用户体验越好。通常提高吞吐的做法是提升batch，单请求串行改为并行；但是batch增大计算量同时增大，时延也会相应增大。

LLM主流模型结构

Attention mask：Prefix mask，Casual mask
Attetion type: multi-head => multi group（节省kv cache缓存）

Decode Only特点

自回归：下一个token生成依赖上一个token
两阶段： prefill，decode
Kv Cache：空间换时间

1 主流推理优化技术

image.png

1.1 算子融合

算子融合即是多个算子OP融合成一个算子OP，减少kernel调用和显存读写。
当前业界主流LLM推理框架，都实现针对Transformer结构的算子融合。例如 DeepSpeed Inference，对Transformer结构支持如下几类算子融合：
（1）归一化层和QKV融合
（2）Self-Attention计算融合，大名鼎鼎的FlashAttention
（3）全连接层、归一化层、Bias-Add、残差连接和GELU
（4）Bias-Add和残差连接融合

捕获.JPG

1.2 高性能算子

对运行热点函数实现高性能算子。单算子优化即根据单个算子，结构计算模式和硬件架构特点，选择性能最高的kernel函数执行。典型的是通用矩阵-矩阵乘GEMM、通用矩阵-向量乘GEMV。
（1）GEMM优化
LLM推理的prefill阶段，self-attention和MLP层存在多个GEMM操作，占据推理时延80%以上。Nvidia提供cuBLAS、CUDA、CUTLASS多层级优化方案。
（2）GEMV优化
LLM推理的decoding阶段，热点函数变为GEMV。GEMV的计算强度相对更低，优化点转向访存开销。

Ascend加速库（AscendC CCE算子）。

2 显存优化

2.1 KV Cache优化

Decode-Only结构的大模型，由于在生成每个新token时都存在重复计算，例如当生成第N个token时，第N-1个token作为输入，需要关注第（N-2）、第（N-3）直到第1个token。KV Cache即是缓存K、V值，避免重复计算。 KV-Cache优化参考：LLM系列：KVCache及优化方法

2.2 Page Attention

研究发现由于缺乏显存有效管理，显存浪费严重（60%-80%），显存浪费主要是kv cache。在PageAttention出现前，业界主流LLM推理框架对kv cache管理低效。 FasterTransformer中会预先分配一个足够长的显存空间用于存储上下文，例如LLama-7B的上下文长度是2k，kv cache需要支持2k个token的显存空间，当实际使用上下文长度小于2k时，发生显存浪费。而且由于序列长度是动态的，kv缓存也是动态变化的，kv cache的缓存管理挑战非常大。过度保留和碎片化导致显存浪费，而且显存动态申请会造成较大的时延开销。

PagedAttention是一种借鉴操作系统中虚拟内存和分页思想的注意力算法，PagedAttention允许在非连续的内存空间中存储连续的 key 和 value。PagedAttention 将每个序列的 KV cache划分为块，每个块包含固定数量token的键和值。 在推理过程中，维护一个逻辑块到物理块的映射表，多个逻辑块可以对应一个物理块，通过引用计数表示物理块被引用的次数（引用计数大于1，表示物理块被使用，引用计数等于零，表示物理块被释放），实现将地址不连续的物理块串联在一起统一管理。

低显存推理优化技术

FlashAttention

  • V1 重排Attention的计算，通过利用tiling、recompute来显著加快速度并将内存使用量从二次降低到线性。

  • V2 a) 减少非矩阵乘计算（FLOPs）提升吞吐量; b) 优化Attention部分thread blocks计算，使其在序列长度维度上并行化，让SM的利用率尽量打满 (CUDA GEMM优化) c) 优化thread blocks内部warp层级的工作模式，尽量减少warp间的通讯和读取shared memory的次数 (CUDA GEMM优化)

  • V3 利用Hopper架构新硬件特性 Tensor Core WGMMA 计算和 Tensor Memory Accelerator（TMA）的异步性来重叠计算与通信，利用FP8低精度计算

FlashDecoding
是在FlashAttention 的基础上针对 inference 的优化，主要分为三步： 1 长文本下将KV分成更小且方便并行的chunk 2 对每个chunk的KV，Q和他们进行之前一样的FlashAttention获取这个chunk的结果 3 对每个chunk结果reduce

PageAttension

  • 背景：KV cache显存占用大、碎片化严重
  • 解决：KV Cache划分成block，block不需要再内存空间连续。实现高效的内存共享。

Offload技术（混合部署或混合推理）
offload对象：权重、KV cache 原理：offload技术将对象卸载到外部资源或将其分布在多种设备上以减少主设备的内存负载，例如CPU内存、Disk等。

Offload技术在 AI框架中应用

HuggingFace Transformers  
DeepSpeed ZeRO-Inference  
llama.cpp
vLLM

FlashAttention
https://zhuanlan.zhihu.com/p/638468472
https://zhuanlan.zhihu.com/p/664061672

Flash Decoding
https://zhuanlan.zhihu.com/p/664264445

投机解码
Lookahead Decoding
https://zhuanlan.zhihu.com/p/701015670

推测解码（Speculative Decoding）
https://blog.csdn.net/qq_27590277/article/details/135812738

4 模型压缩

  • A16W8/A16W4 权重压缩，计算时插入反量化算子，精度无损，GPU收益比较明显，Ascend带宽比较弱

  • W8A8 SmoothQuant量化

  • KV Cache量化 插入Quant、De-quant 算法：xxx

5 服务

  • continuous batching
  • dynamic batching

Dynamic Batching与Continuous Batching类似，都是为了提升推理的吞吐。两者区别是Continuous Batching是在时间轴方向动态插入新队列，Dynamic Batching则是在batch维度方向动态插入维度。

6 分布式

DP、TP、PP，推理场景主要是TP

  • 数据并行
  • 张量并行
  • 流水线并行

embedding
权重竖切，[word_size, hidden_size/2]，最后AllGather拼接

attention
按照head个数切，每个卡跑 head/n 个头 Multi-query，需要复制head
卡间通信：hccs通信

7 模型结构优化

FlashAttension QKV分块计算，中间结果缓存在cache，减少HBM 来回搬运 以增量方式执行Softmax归一化

MQA、GQA K,V 只有一个头，Q是多头，共享kv，减少mem和带宽，但有精度损失。

参考：
www.microsoft.com/en-us/resea…

avatar
文章
阅读
粉丝