引言
随着大语言模型(LLM)参数的快速增长,推理效率成为制约其大规模部署的核心瓶颈。以GPT-3 175B为例,单是加载模型权重便需要数百GB显存,而自回归解码的序列生成特性更使得推理过程对内存带宽和计算效率提出了极高要求。本文将从底层原理出发,系统梳理LLM推理优化的核心技术路线。
1. KV Cache:自回归解码的核心加速机制
Transformer的自回归生成过程中,每一个新token的生成都需要计算所有历史token的注意力。若每次重新计算全部Key和Value矩阵,时间复杂度将达到O(N²)。KV Cache的核心思想是将已计算的K、V张量缓存复用,使得每步推理的计算量从O(N²)降至O(N)。
具体而言,在多头注意力层中,对于第i步推理,我们只需计算新输入token对应的k_i和v_i,并将其追加到已缓存的K_cache和V_cache中。注意力分数通过当前query与所有历史key的点积获得,有效利用了Transformer的并行计算优势。
2. PagedAttention:解决显存碎片化的关键创新
vLLM项目的核心贡献之一便是PagedAttention技术。传统KV Cache为每个序列预留一块连续的显存空间,这导致了严重的内存碎片问题——随着序列长度变化和并发请求的启停,大量不连续的显存空洞无法被有效利用。
PagedAttention借鉴操作系统中的虚拟内存分页机制,将KV Cache拆分成固定大小的block(通常为16个token),每个序列维护一个block table来映射逻辑block到物理block。这种设计使得:
- 显存按block粒度分配和回收,消除外部碎片
- 不同序列可共享相同前缀的KV Cache(prefix sharing)
- 显存利用率从传统方案的20-40%提升至接近100%
3. Continuous Batching:提升吞吐量的利器
传统静态批处理需等待组内所有序列完成推理,导致大量计算浪费在padding和空闲等待上。Continuous Batching(或称iteration-level scheduling)在每个解码迭代后动态调整批次组成:已完成的序列立即退出,新请求随时加入。
vLLM与TensorRT-LLM均实现了这一机制,相比静态批处理,吞吐量提升可达2-8倍。关键在于结合PagedAttention的灵活内存管理,实现计算与调度的高效协同。
4. 量化与精度压缩技术
模型量化是降低推理显存和带宽需求的主流手段:
- INT8量化:SmoothQuant、GPTQ等方法实现激活和权重的INT8映射,精度损失通常在1%以内,显存减半
- INT4/FP4量化:AWQ(Activation-aware Weight Quantization)利用激活分布保护重要权重列,在4-bit下保持可接受的生成质量
- 混合精度:对敏感层(如首尾层)保持FP16,其余层使用INT4,平衡精度与效率
5. Speculative Decoding:以小博大的加速艺术
Speculative Decoding(推测解码)的核心思路是使用一个小模型(draft model)快速生成k个候选token,然后由大模型(target model)一次性验证这些候选。若验证通过,相当于一次前向传播生成了k+1个token。
该方法之所以有效,是因为LLM推理在显存带宽受限时,大模型对多个token的验证几乎与单个token的前向传播耗时相当——读入权重的成本远大于计算本身的成本。当小模型准确率在80%以上时,可获得1.5-3倍加速。
6. FlashAttention与IO感知计算
FlashAttention通过IO-aware的精确注意力算法,将HBM(高带宽显存)与SRAM之间的数据传输最小化。在长序列场景下,标准注意力计算需要O(N²)的HBM访问,而FlashAttention将其降至O(N²/M)(M为SRAM大小)。
FlashAttention-2进一步优化了线程块和warp级别的并行策略,在A100 GPU上达到了理论峰值算力70%以上的利用率,成为现代LLM推理框架的标准组件。
7. 部署架构全景
现代LLM推理栈呈现分层架构:
- 硬件层:NVIDIA A100/H100、AMD MI300、AWS Inferentia等专用AI加速器
- 内核层:FlashAttention、fused CUDA kernel、TRITON算子
- 引擎层:vLLM、TensorRT-LLM、TGI(Text Generation Inference)、llama.cpp
- 服务层:OpenAI-compatible API、负载均衡、autoscaling、流式输出
总结
LLM推理优化是一个涵盖算法、系统、硬件的多层次工程问题。从KV Cache的基础加速到PagedAttention的显存革新,再到现在的多维优化技术,每一步突破都推动着大模型从实验室走向生产环境。随着MoE架构、长序列处理、多模态推理等新挑战的出现,推理优化领域将持续保持高速演进。

发表评论 取消回复