在 16.1 节中,我们指出显存瓶颈是大模型训练与推理的第一硬约束——单张 A100 的 80 GB HBM 看似充裕,但面对 70B 参数的全精度权重外加 KV 缓存、激活值、优化器状态时,依然捉襟见肘。16.3 节的 GPU 选型也反复印证:显存容量直接决定了你能跑多大的模型、多长的上下文。
本节从硬件显存管理视角出发,系统拆解四种最核心、最立竿见影的显存优化技术。它们不是互斥的,而是生产环境中常叠加使用的组合拳。理解每种技术的显存收益边界与延迟代价,是你做推理部署和训练调优时进行 trade-off 决策的基础。
一、量化(Quantization):用精度换容量
技术实质:将模型权重、激活值或 KV 缓存从高精度浮点数(FP16/BF16,各占 2 字节)映射到低精度整数(INT8 占 1 字节、INT4 占 0.5 字节),从而直接压缩显存占用。
显存收益测算:
- 权重量化:70B 模型 FP16 需约 140 GB;INT4 量化后降至约 35–40 GB,单张 A100 即可加载。
- KV 缓存量化:在 INT8 下,长上下文场景的 KV 缓存占用直接减半,这对 128K 上下文窗口的部署尤为关键。
工程分类与适用场景:
| 量化类型 | 代表方案 | 作用对象 | 特点 |
|---------|---------|---------|------|
| 训练后量化(PTQ) | GPTQ、AWQ、GGUF | 仅权重 | 无需重新训练,部署侧直接转换;适合推理 |
| 量化感知训练(QFT/QLoRA) | QLoRA、LLM-QAT | 权重+部分激活 | 在微调阶段引入低精度,精度损失更小 |
| KV Cache 量化 | KV Cache INT8/FP8 | KV 缓存 | 动态在线量化,缓解长序列显存压力 |
精度与速度的 trade-off:
- INT8:对大多数生成任务几乎无损,是性价比最高的甜点区。
- INT4/FP4:权重量化可运行更大模型,但可能触发精度悬崖(Perplexity 陡升),尤其在数学、代码等需要精确符号推理的任务中。
- 混合精度:常见做法是权重 INT4/INT8,激活 FP16,推理时动态反量化,兼顾容量与计算精度。
实用认知:量化省下的显存不是免费的午餐。低精度运算在部分 GPU(如消费级 RTX 4090 的 Tensor Core)上支持良好,但在某些架构上可能需要模拟运算,导致实际延迟不降反升。务必在目标硬件上做端到端 benchmark,而非只看模型文件体积。
二、KV 缓存(KV Cache):自回归生成的显存蓄水池
技术实质:在自回归生成中(1.1 节),Decoder 每生成一个新 Token,都需要计算当前 Token 与所有历史 Token 的注意力。为了避免重复计算前面已经算过的 Key 和 Value,模型在显存中把它们缓存起来。这构成了推理阶段最大的显存消耗项之一,常超过权重本身。
显存占用公式(以单 batch 为例):
KV Cache ≈ 2 × num_layers × hidden_size × seq_length × bytes_per_token
其中系数 2 代表 K 和 V 两组张量。以 Llama 3 70B(80 层,8192 hidden size)为例,FP16 下每 Token 的 KV 缓存约:
2 × 80 × 8192 × 2 bytes ≈ 2.5 MB/Token
若上下文长度为 8192,单条请求的 KV 缓存就高达 20 GB。当 batch size 提升到 16 时,仅 KV 缓存就吃掉 320 GB,远超单卡显存。
优化方向(硬件视角):
- 多查询注意力(MQA)与分组查询注意力(GQA):Llama 2/3、Mistral 等模型在架构层面减少 KV 头的数量,将缓存量压缩为原来的 1/8 或 1/4。
- KV Cache 量化:如前述,将 FP16 KV 压至 INT8,直接减半。
- 分页/碎片管理:vLLM 的 PagedAttention(将在 17.3 和 23.1 节详述)把 KV 缓存按块管理,避免传统实现中因预分配最大长度而产生的显存浪费。这在硬件层面等效于显存虚拟化。
实用认知:如果你在做 API 服务部署,KV 缓存管理比权重量化更能决定你的并发上限。一个 7B 模型权重只有 14 GB,但 64 路并发 × 8K 上下文的 KV 缓存可以轻松突破 100 GB。这也是为什么推理集群需要配备比训练集群更大单卡显存或更激进的缓存压缩策略。
三、注意力优化:从 O(N²) 显存到 O(N) 的硬件革命
技术实质:标准 Self-Attention 的计算复杂度是序列长度的平方 O(N²),而中间结果(注意力矩阵 S = QK^T/√d)的显存占用同样为 O(N²)。当上下文从 2K 扩展到 128K 时,注意力矩阵的显存占用会从 MB 级飙升到 GB 甚至十 GB 级,且对 HBM(高带宽显存)的读写带宽造成巨大压力。
FlashAttention 系列(17.3 节将深入原理)是目前最主流的硬件层优化方案。其核心不是减少计算量,而是通过分块计算(Tiling)和 Online Softmax,将中间注意力矩阵从 HBM 中剔除:
- 标准 Attention:Q、K、V 加载到 HBM → 计算 S → 存回 HBM → 计算 Softmax → 存回 HBM → 乘 V → 写回 HBM。中间矩阵反复读写,显存峰值高、带宽拥堵。
- FlashAttention:将 Q、K、V 切分成小块放在 SRAM(片上高速缓存)中计算,只输出最终的 O(注意力输出),不保存中间 S 矩阵。
显存收益:
- 标准 Attention 的显存峰值 ∝ N²;FlashAttention 的显存峰值 ∝ N(仅保留输出和少量累加器)。
- 在长序列场景(>4K)下,FlashAttention 不仅显著降低显存,还因减少了 HBM 读写而加速 2–4 倍。
变体与边界:
- FlashAttention-2/3:进一步优化了 Warp 级并行和异步数据拷贝,在 H100 上能更好利用 Tensor Core 和异步拷贝指令。
- Sparse Attention / Sliding Window:如 Mistral 的滑动窗口注意力,从算法层限制每个 Token 只能 attend 到局部窗口,从根本上把 N² 降到 N×W(窗口大小),但会牺牲部分长程依赖能力。
实用认知:FlashAttention 的收益与 GPU 架构强相关。在显存带宽受限的 A100/H100 上,它是“免费午餐”;但在某些早期架构或消费级卡上,若 SRAM 太小导致分块过细,收益会打折扣。此外,FlashAttention 目前对变长序列的 Padding 处理已有很好支持,但接入自定义 Attention 变体(如带有偏置项的 ALiBi)时,需确认实现版本是否兼容。
四、显存交换(Offloading / Swapping):用带宽换容量
技术实质:当 GPU HBM 装不下所有数据时,将部分张量(权重、优化器状态、KV 缓存或激活值)卸载到容量更大但速度更慢的存储层级——通常是 CPU 内存(DRAM),甚至 NVMe SSD。需要时再异步传输回 HBM。
典型方案:
| 方案 | 场景 | 卸载对象 | 技术要点 |
|------|------|---------|---------|
| ZeRO-Offload / ZeRO-Infinity | 训练 | 优化器状态、权重、梯度 | CPU 内存甚至 NVMe 作为显存扩展;计算时按需换入 |
| DeepSpeed Inference | 推理 | 不活跃的层权重 | 流水线式逐层加载,牺牲延迟换大模型支持 |
| vLLM PagedAttention | 推理 | KV Cache 块 | 将暂时不用的 KV 块换出到 CPU,需要时换入(类似 OS 虚拟内存) |
| FlexGen / Infinite-LLM | 推理 | KV Cache / 权重 | 磁盘 offload,支持单卡跑超大模型,但延迟极高 |
带宽瓶颈与延迟代价:
- PCIe 带宽:CPU DRAM ↔ GPU HBM 经 PCIe 4.0/5.0,带宽约 32–64 GB/s,仅为 H100 NVLink(900 GB/s)的几十分之一。频繁换入换出会导致 GPU 计算单元大量空等。
- NVMe 带宽:更低(数 GB/s 级别),仅适合极端离线场景。
- 异步流水线:成熟的 offload 方案会尝试用计算掩盖传输(Overlap),即 GPU 计算当前层时,预取下一层权重。但如果模型粒度太细或带宽太低,掩盖效率会急剧下降。
实用认知:显存交换是兜底方案,而非性能方案。在训练侧,ZeRO-Offload 让单卡跑大模型成为可能,但吞吐量会暴跌一个数量级;在推理侧,PagedAttention 的 swap 机制主要用于超长上下文 + 超大并发的极端情况,常规业务应优先通过量化、GQA、FlashAttention 把显存压进 HBM 内解决。
五、四种技术的组合策略与选型建议
在实际系统中,这四种技术通常叠加使用。以下是一张面向硬件/infra 工程师的决策速查表:
| 技术 | 显存节省力度 | 延迟影响 | 工程复杂度 | 推荐优先级 |
|------|-----------|---------|-----------|-----------|
| 权重量化(INT8/INT4) | ★★★★★(2–4×) | 轻微增加或降低 | 低(已有成熟工具链) | 首选 |
| KV Cache 量化 | ★★★★(2×) | 轻微增加 | 中(需推理框架支持) | 长上下文必选 |
| GQA / MQA 架构 | ★★★(4–8× 减少头数) | 无 | 极低(模型设计阶段决定) | 模型选型时优先 |
| FlashAttention | ★★★(长序列 O(N²)→O(N)) | 降低 | 低(主流框架已内置) | 长上下文必选 |
| 显存交换 Offload | ★★★★★(理论上无限扩展) | 显著增加 | 中–高 | 最后手段 |
典型组合示例:
- 单卡 4090(24 GB)跑 70B 模型:INT4 权重量化(35 GB→可压缩加载)+ CPU offload 剩余层 + 短上下文避免 KV 缓存爆炸。
- A100 80G 推理服务 8K 上下文:FP16 权重 + FlashAttention + GQA + INT8 KV Cache,追求并发与延迟平衡。
- 训练 65B 模型 on 8×A100:ZeRO-3 切分权重 + FlashAttention + BF16 混合精度 + Activation Checkpointing(用计算换显存,见 16.1 节)。
六、小结
显存优化的本质是在容量、带宽、精度、延迟的四维空间里寻找帕累托最优:
- 量化直接压缩数据位宽,是容量受限时的第一杠杆;
- KV 缓存管理决定了推理服务的并发天花板,其瓶颈往往比权重更早到来;
- 注意力优化(FlashAttention)通过重构计算流水减少 HBM 占用与带宽压力,是长上下文的刚需;
- 显存交换扩展了物理边界,但受限于 PCIe/NVMe 带宽,仅适合作为兜底。
在 17.3 节中,我们将进一步深入 CUDA 软件栈,剖析 FlashAttention 和 PagedAttention 的 Kernel 级实现原理;而在第 22 章(推理优化)和第 23 章(推理框架)中,这些技术将以更完整的产品形态(vLLM、TensorRT-LLM 等)出现,指导你完成从原理到落地的最后一公里。