大模型推理加速手段整理

总览

读完A Survey from Algorithms to Systems后，根据自己的理解进行的总结。列出了一些MLSys的优化，以及推理加速工程师应该了解的attention算法的一些变种。


flowchart LR
InferOpt(LLM推理优化) --> Sys(MLSys优化) & Alg(ML算法优化) 
Alg --> KVS(KV sharing) & AS(Attention compression) & SD(投机采样 Speculative decoding) & MLS(...)
AS --> SW(Sliding Window Attention) & ASA(Adaptively Sparse Attention) & ASS(...)
KVS --> MQA & GQA & MLA
Sys --> quant(量化/低精度) & parallel(并行) & fd(kernel优化:Flash Decoding++) & MM(memory管理) & cb(Request 调度：Continous Batching)
parallel --> TP & PP & SP
MM --> kvc(KV Cache) & PA(Paged Attenton) & PC(Prefix Cache)

style InferOpt fill:#c699f0
style Alg fill:#f5dd05
style Sys fill:#f5dd05
style KVS fill:#99c1f2

style AS fill:#99c1f2
style SD fill:#99c1f2
style MLS fill:#99c1f2

style quant fill:#99c1f2
style parallel fill:#99c1f2
style fd fill:#99c1f2
style MM fill:#99c1f2
style cb fill:#99c1f2

推理优化可分为：

AI算法优化：优化模型，在不明显降低质量的情况下，加速推理。如：GQA, MQA, Sliding Windwo Attention …

系统优化：不改变模型语义，优化MLSyS，加速推理。如：KV Cache, Flash Decoding++…

LLM 推理指标：

FTT(First Token Time): 一个request从输入到输出的第一个token的延迟

TPS(Tokens Per Second):一个request的平均每秒的输出tokens数

Throughput：吞吐量，即服务在一秒内，对所有request生成的token的数量之和。

Quality: 推理结果的质量

Sys优化对指标的影响：❤️表示有提升，🖤表示有降低

ㅤ	FTT	TPS	ThroughPut	Quality
量化/低精度	❤️	❤️	❤️	🖤
并行/FlashDecoding	❤️	❤️	❤️	ㅤ
KV Cache	ㅤ	❤️	❤️	ㅤ
Paged Attention	ㅤ	ㅤ	❤️	ㅤ
Continous Batching	🖤	🖤	❤️	ㅤ

System Optimization

系统优化，不改变模型语义的情况下，加速运行

Low-bit / Quantization

Quantization-Aware Training

Post-Training Quantization

fp8/bf8…

Model parallelism

SP (ring Attention， seq维度并行)

Memory Management (KV Cache)

KV Cache：cache之前tokens的kv，下一轮直接使用。kv cache所需要的memory大小无法确定，可能会很大，因此需要一种管理KV Cache的策略。

naive approach: FasterTransformer根据max seq_len预分配一块连续的memory。内存浪费比较大。

Page-Attention: vLLM 将KV cache分为多个不连续的物理page，decoding时，按需分配，使用block table记录需要的page的位置，与操作系统的页表相似。

由于page attention天然的对kv cache进行分块，所以在decoding计算attention时，可以自然地使用FlashDecoding。

Preifx Cache：将KV Cache以前缀树的形式，组织起来，在不同对话中共享。

Request Scheduling

continous batching：动态删除batch中提前结束的request，并插入新到来的request

Kernel Optimization

flash decoding (++): 利用online softmax实现分块attention，在kv维度可以进行并行

Algorithmic Innovations

网络算法优化

Decoding Algorithm

非自回归

非自回归

半自会归

blockwise parallel decoding：一次前向，生成多个token

当前在quality上，多数还是不如自会归

Speculative decoding

小模型推理后续r个token，大模型一次前向进行验证并推理下一个token。验证通过则采用，不通过则丢弃后续token，然后继续前向生成一个正确的token。

优点是保证了quality，因为所有token都会经过大模型验证。

challenge在于如何小而准的小模型来做预测。

拓展阅读：Online Speculative Decoding

Tree-based Speculative Inference，使用多个小模型，每个小模型都向后预测多个token，然后组织成一棵树，通过精心设计的mask，可交由大模型一次验证。

Early exiting

idea: 中间层有足够的信息得到目标分布。在中间层插入classifier，动态退出。针对简单的request，进行的计算会更少。

Cascade inference

首先分析request的难度，简单的问题使用规模较小的模型，复杂的问题使用规模较大的模型。

Architecture Design，主要针对attention 架构的改动

Config Downsizing

缩减参数规模，但是影响quality

Attention compression

压缩self-attentiond的context规模，减少计算量以及KV Cache的大小。

直接保留特定位置的token：

sliding window attention

dilated attention

通过某种评估方式，丢弃掉不重要的token：

adaptively sparse attention

Activation sharing

多个self-attention head共享一组KV

MQA (Multi-query Attention)

GQA (Group-query Attention)

最新的MLA(Multi-head Latent Attention)

Conditional computing

Moe , 混合专家模型。 Router + 多个专家网络。 router动态决定采用哪些专家,而不是全量计算。