大模型推理加速
大模型推理
基于Transformer解码器的大模型,推理阶段计算过程。
Transformer Decoder推理计算过程
标准Transformer Decoder推理时间复杂度是O(N^2)
推理加速方法
量化
16bit、8bit量化,通常会损失一定的效果。
模型结构改进
Muti-Query Attention
每个Head有一个独立的Query,所有Head共享一组Key、Value。减少Key、Value的内存占用和带宽开销。
Grouped Query Attention
每个Head有一个独立的Query,所有Head划分为G组,每组共享一组Key、Value。GQA 是介于 MHA 与 MQA 之间的 中间折中方案。
KVCache及其优化
缓存LLM Decode阶段的Key、value,避免重复计算Key、Value。
PageAttention
通过将传统连续缓存(KV Cache)拆分成可管理的小块(pages / blocks),显著提升了内存利用率与推理效率。
FlashAttention
分配计算、充分利用GPU。
FlashAttention 1
采用tiling技术,分块计算softmax。减少内存读写操作。
FlashAttention2相比FlashAttention1加速2x,计算效率达到GEMM的50%~70%。
改进1:
改进2:
改进3:
因果掩码,跳过无效计算,降低构造[L,L]完整掩码矩阵带来的显存开销,带来1.7~1.8倍计算加速。
2.1 上三角块跳过计算
2.2 下三角块无需掩码
2.3 仅针对处在对角线的块进行局部因果掩码。
softmax
safe softmax
block softmax
内存墙(Memory Wall)
硬件层面
高速带宽内存
软件层面
FlashAttention、PageAttention、MultiQueryAttention可以理解为从软件层面缓解内存墙问题。
Dynamic Batch
Speculative Decoding
使用小模型一次生成多个token,大模型对生成的token序列进行并行验证。
节省数据传输
SparQ Attention修改注意力
ChatGPT总结
大模型(如 GPT-3、LLaMA、Qwen 等)的推理速度,是影响其在线服务成本和响应时延的核心瓶颈。除了 FlashAttention,还有多种有效的推理加速技术,涵盖 算子级优化、模型结构压缩、系统并行、KV 缓存改进等方面。
下面我系统总结 当前主流 LLM 推理加速方法,并说明它们背后的原理与适用场景:
🧠 一、算子级优化(Operator-Level)
| 方法 | 原理 | 效果 |
|---|---|---|
| ✅ FlashAttention | Kernel 融合 + 分块 softmax | 内存↓、速度↑,适合长上下文 |
| xformers / Triton kernel | 自定义 fused kernel | 替代标准 PyTorch 算子,更高效 |
| TensorRT / ONNX 加速 | 转为图执行,优化图结构 | 大幅降低推理延迟,适合部署 |
| FlashDecoding(来自 FlashInfer) | 加速 KV cache 使用 | 对解码阶段大幅提速 |
💾 二、KV Cache 优化(Decoder Only 模型特有)
| 方法 | 原理 | 效果 |
|---|---|---|
| ✅ KV Cache | 缓存历史 token 的 K/V,避免重复计算 | 将 self-attn 从 $O(n^2)$ 降为 $O(n)$ |
| Paged KV Cache(vLLM) | 将 KV Cache 分页管理,按需释放、重用 | 支持多请求高效复用显存 |
| MQA / GQA(多组共享 key) | 多个 head 共享 key,减少 cache 大小 | 比标准 KV 缓存减少一半显存 |
| Speculative Decoding | 猜测多个 token、再校验 | 多 token 并行预测,加速解码阶段 |
📉 三、模型结构压缩与低精度推理
| 方法 | 原理 | 效果 |
|---|---|---|
| ✅ 量化(Int8/Int4) | 减少参数存储和矩阵乘计算精度 | 显存 ↓ 4-8x,速度↑1.5-3x |
| BitsAndBytes (bnb) | 支持训练和推理的量化引擎 | Hugging Face 提供接口 |
| LoRA + 推理合并 | 推理时合并 adapter 和主模型 | 多任务共享权重,适合多模型部署 |
| 模型蒸馏 | 训练小模型拟合大模型输出 | 能大幅缩小模型体积,提升速度 |
🚀 四、并行策略与系统优化
| 方法 | 原理 | 效果 |
|---|---|---|
| Tensor Parallel | 拆分每层的参数矩阵,在多个 GPU 上并行 | 多卡推理/训练 |
| Pipeline Parallel | 拆分不同层,流水线执行 | 减少 GPU 空闲时间 |
| Sequence Parallel(如 DeepSpeed) | 对单序列进行细粒度并行 | 减少单卡内存瓶颈 |
| Prefill + Decode 分离 | 预填充阶段并行,decode 阶段增量计算 | 大幅提升多请求吞吐量(vLLM) |
💡 五、推理策略层面优化
| 方法 | 原理 | 效果 |
|---|---|---|
| Top-k, Top-p Sampling | 限制候选词数量 | 加速生成阶段 |
| Early Exit | 模型前几层就做输出判断 | 降低平均计算量(需训练) |
| Dynamic Prompt Pruning | 对 prompt 进行裁剪和重写 | 降低 context 长度 |
🔬 六、业界代表系统(你可以参考/复用)
| 系统/框架 | 核心优势 | 来源 |
|---|---|---|
| vLLM | Paged KV Cache + 高并发 + FlashAttention | UC Berkeley |
| Triton / TensorRT-LLM | 高性能推理 + 模型优化 | NVIDIA |
| DeepSpeed-Inference | 分布式推理 + INT8/FP16 支持 | Microsoft |
| TGI (Text Generation Inference) | Hugging Face 部署方案 | HuggingFace |
| FasterTransformer | 编译优化 + 分布式支持 | NVIDIA |
参考
论文
Fast Transformer Decoding
Grouped Query Attention
DeepSeek Innovative Techniques
SparQ Attention: Bandwidth-Efficient LLM Inference