引言:LLM推理的计算鸿沟
大语言模型的参数量从7B扩展到70B乃至405B,推理成本已成为AI应用落地的核心瓶颈。以Llama-3-70B为例, FP16推理需要140GB显存,单张A100 80GB无法加载完整的模型。在与之前LLM推理优化主题不重复的前提下,本文聚焦量化压缩与推理系统架构两大技术路径。
1. 量化技术谱系
量化通过降低模型权重和激活值的数值精度来减少显存占用和加速计算:
1.1 训练后量化(PTQ):GPTQ / AWQ / SmoothQuant
三类主流PTQ方法各有适用场景:
- GPTQ:基于Optimal Brain Quantization的逐列校准,利用Hessian矩阵的逆自动确定量化参数,实现3-bit量化。对于OPT-66B,3-bit GPTQ的困惑度增长小于0.5%
- AWQ(Activation-Aware Weight Quantization):通过逐通道缩放(per-channel scaling)保护重要权重通道。相比GPTQ,同等bit位数下perplexity更低,且对硬件更友好(避免了随机顺序的列更新)
- SmoothQuant:发现在LLM中激活值比权重更难量化(存在离群通道)。该方法将激活的量化难度通过数学等价变换迁移到权重上,实现W8A8全量化,配合INT8 GEMM kernel加速
1.2 量化感知训练(QAT)与校准策略
- QAT在训练中模拟量化噪声,典型loss增长低于PTQ但需要GPU训练资源
- 校准数据集的选择至关重要:GPTQ论文建议使用C4数据集的128个2048-token的样本作为校准集
- 领域适配问题:医疗/法律等垂直领域模型通用校准集效果差,需使用领域数据重新校准
1.3 FP8与新一代数值格式
随着NVIDIA H100/B200原生支持FP8,硬件级低精度推理成为标配:
- E4M3 FP8:4位指数+3位尾数,适合推理(动态范围更大)
- E5M2 FP8:5位指数+2位尾数,适合训练梯度
- FP8混合精度:用FP8做GEMM,FP16/BF16做SoftMax和LayerNorm,典型加速1.5-2x
2. KV Cache与内存管理
自回归推理中,每生成一个token需要计算与所有历史token的注意力。KV Cache技术避免重复计算,但带来了显存管理的挑战:
2.1 vLLM的PagedAttention:虚拟内存管理革命
vLLM的核心创新是借鉴操作系统虚拟内存的思想,将KV Cache分页管理:
- 分块机制:将序列的KV Cache固定大小的Block(如每Block存储16个token的KV向量)
- Block Table:每序列维护Block到物理显存Block的映射表,支持非连续存储
- 消除碎片:相比传统连续分配消除了内部碎片和外部碎片,内存利用率从20-40%提升至接近100%
- Copy-on-Write:Beam Search场景下多个序列共享前缀KV Cache,写时复制避免物理拷贝,内存节省数倍
2.2 KV Cache的内存占用计算
对于Transformer模型,KV Cache的内存占用为:
Memory = 2 × n_layers × n_head × d_head × seq_len × batch_size × sizeof(dtype)
以Llama-3-70B为例(n_layers=80, n_head=64, d_head=128),FP16下单个2048-token序列的KV Cache为:2×80×64×128×2048×2Byte = 512MB。当并发batch_size=64时,KV Cache占用高达32GB!这解释了为何内存管理对推理吞吐至关重要。
3. 推理引擎架构对比
| 引擎 | 核心技术 | 核心优势 | 适用场景 |
|---|---|---|---|
| vLLM | PagedAttention + Continuous Batching | 高吞吐+内存效率 | 通用API服务 |
| TensorRT-LLM | KV Cache + CUDA Graph + Custom Kernel | 极致GPU利用率 | NVIDIA GPU环境 |
| SGLang | RadixAttention + 压缩FSM | 结构化生成优化 | 工具调用/Agent场景 |
| llama.cpp | GGUF量化 + Metal/CUDA混合 | 消费级硬件 | 个人开发者 |
3.1 TensorRT-LLM的极致优化
NVIDIA官方推理引擎将多个GPU操作融合为单个CUDA Graph:
- Custom Attention Kernel:FlashAttention-2的TensorRT定制版,减少memory-bound瓶颈
- In-flight Batching:类似Continuous Batching,在不同batch slot间动态调度token级别执行
- Multi-GPU Tensor Parallel:通过NCCL高效实现Attention和MLP的层内切分
- FP8/INT4 Weight-Only Quantization:利用H100的FP8 Tensor Core实现极致吞吐
3.2 SGLang的RadixAttention
SGLang(Structured Language Model)提出的RadixAttention优化了共享前缀的系统提示词:
- 使用Radix Tree管理KV Cache前缀,将前缀共享从线性查找优化为对数级
- 在多轮对话和少样本系统提示词场景下,提升缓存命中率,减少prefill时间
- 对于结构化输出(JSON Schema/Regex),基于DFA约束的token masking,减少无效采样
4. 关键量化指标的权衡
生产部署时需要平衡延迟、吞吐、显存和精度四个维度的指标:
- TTFT(Time to First Token):Prefill阶段延迟,受序列长度和GPU显存带宽主导
- TPS(Tokens Per Second):生成阶段吞吐,受解码步长和数据类型影响
- Memory Footprint:FP16/INT8/INT4的显存占用比例大致为4:2:1(含KV Cache)
- Quality Degradation:4-bit量化在大多数任务中保持FP16性能99%+,3-bit以下质量急剧下降
5. 未来趋势
- Speculative Decoding:用小型模型快速生成多个token,大模型并行验证,实现2-3x加速
- Mamba/SSM架构:线性复杂度的状态空间模型替代Transformer,推理吞吐提升5x
- Disaggregated Serving:将Prefill(计算密集)和Decode(内存密集)分离到不同GPU集群
- FP4/BitNet 1-bit推理:极致量化方向,开拓端侧AI推理的可能性
结语
LLM推理系统的优化已进入软硬件协同设计的新阶段:量化压缩降低理论算力需求,PagedAttention和调度策略释放系统效率,FlashAttention和CUDA Graph榨取GPU极限性能。对于AI基础设施团队而言,理解全栈优化技术是支撑大规模AI应用的必备能力。

发表评论 取消回复