大模型推理优化全景:从KV Cache到投机解码
推理瓶颈在哪里? 大模型推理的核心瓶颈是自回归生成的串行性:每生成一个token都需要将之前所有token的KV对参与注意力计算。随着序列变长,计算量和显存占用呈线性甚至二次增长。 推理优化的目标可以归纳为三个维度: 降低延迟(首token延迟TTFT + 每token延迟TPOT) 提升吞吐(每秒生成token总数) 减少显存(KV Cache占用) KV Cache:标准但不够 KV Cache是现代推理引擎的标配。在生成阶段,避免重复计算历史token的Key和Value向量,而是缓存复用。但KV Cache的显存占用非常大:以Llama-70B为例,单条请求的KV Cache在4K序列长度下约占2GB显存。 优化方向: PagedAttention(vLLM):将KV Cache分为固定大小的block,通过页表管理,显著减少碎片 量化KV Cache:将FP16的KV对压缩为INT8甚至INT4,显存减半甚至1/4 稀疏KV Cache:基于注意力分数剪枝,丢弃不重要的KV对 连续批处理(Continuous Batching) 传统static batching需要等待批次内所有请求完成才能释放资源,短请求被长请求拖累。连续批处理(也叫iteration-level batching)在每个forward pass级别动态加入新请求、移除已完成请求: 时间步1: [req_A(token5), req_B(token3), req_C(token1)] 时间步2: [req_A(token6), req_B(token4), req_D(token1)] # C完成,D加入 vLLM和TGI都实现了这一机制,实测吞吐量可提升3-5倍。 投机解码(Speculative Decoding) 投机解码利用一个小模型(draft model)快速生成候选token序列,再由大模型一次性验证。如果小模型的猜测正确率足够高,就能以接近小模型的速度输出大模型质量的token: 小模型生成N个候选token(并行,很快) 大模型对这N个token做一次forward(并行验证) 接受正确的前缀,拒绝处重新采样 Medusa在此基础上扩展,在单次forward中同时预测多个位置的token(多头预测),进一步提升加速比。 量化推理 训练时用FP16/BF16,推理时降低精度: 精度 方法 性能损失 INT8 W8A8 / SmoothQuant <1% INT4 W4A16 / GPTQ / AWQ 1-3% 2-bit BitNet / 1.58bit 训练时需特殊设计 AWQ和GPTQ是目前最流行的4bit权重量化方案,配合Marlin内核可达极高吞吐。需注意,4bit量化在70B以上模型上表现好,但7B级别可能有明显精度下降。 ...