当你满怀信心地启动一个 70B 参数的大模型推理任务，看着进度条稳步前进，突然——屏幕一闪，一行冰冷的红字弹出：CUDA out of memory。你检查显存，明明还剩 20GB 空间，为什么还是 OOM？

这不是个例。在深度学习领域，尤其是大语言模型（LLM）时代，GPU 显存不足早已成为最常见的部署瓶颈。但问题的根源往往比"买更好的显卡"复杂得多：有时候是硬件层面的内存带宽限制，有时候是软件层面的碎片化问题，更多时候是两者的叠加效应。

显存到底被谁吃掉了？

要理解显存瓶颈，首先需要拆解一个模型的显存占用结构。以一个 13B 参数的模型在 NVIDIA A100 40GB 上推理为例，显存分配大致如下：

  • 模型权重：约占 65%（约 26GB），这部分是静态的，只要模型加载就会一直占用
  • KV Cache：约 30%（可变），这是动态分配的核心区域
  • 激活值和其他临时数据：约 5%，转瞬即逝

看起来权重是大头，但问题恰恰出在 KV Cache 上——它会随着序列长度和批量大小线性增长，而且增长方式极其"混乱"。

KV Cache：显存消耗的隐形杀手

KV Cache（Key-Value Cache）是 Transformer 自回归推理的核心优化技术。简单说，每生成一个新 token，模型需要计算它与之前所有 token 的注意力关系。如果不缓存之前的 Key 和 Value 矩阵，每次生成都要重新计算所有历史 token，计算量会呈二次方增长。

缓存确实解决了计算效率问题，但带来了新的显存压力。以 OPT-13B 模型为例，单个 token 的 KV Cache 需要约 800KB：

KV Cache per token = 2 (K和V) × 5120 (隐藏层维度) × 40 (层数) × 2 (FP16字节数)
                   ≈ 800 KB

如果一个序列最长 2048 个 token，单个请求的 KV Cache 就要消耗 1.6GB。当批量增加时，这个数字会迅速膨胀。

但更关键的问题是：现有的内存管理方式极其低效。

碎片化：显存利用率的隐形杀手

传统推理框架（如 FasterTransformer、HuggingFace Transformers）为每个请求预分配一块连续内存，大小按 max_seq_len 计算。这种看似简单的策略带来了三种严重的浪费：

预留碎片（Reservation Fragment）：因为不知道实际输出多长，系统必须预留最大可能长度。一个 prompt 可能只生成 50 个 token，但系统已经预留了 2048 个位置的空间。

内部碎片（Internal Fragment）：当序列提前结束时，预留空间的后半部分彻底浪费。比如 max_seq_len 设为 2048，实际只输出了 100 个 token，剩下 1948 个位置白白占用。

外部碎片（External Fragment）：不同请求的内存块大小不同，频繁分配释放后，显存空间变得支离破碎。可能显存统计上还有 10GB 空闲，但找不到一块足够大的连续区域来容纳新请求。

根据 vLLM 论文的实测数据，在现有系统中，只有 20.4% 到 38.2% 的 KV Cache 内存实际用于存储 token 状态，其余全部被各种形式的碎片浪费掉了。

pie title 传统系统KV Cache内存分布
    "实际token状态 (20-38%)" : 30
    "预留碎片" : 25
    "内部碎片" : 25
    "外部碎片" : 20

这种低效的内存利用直接限制了批处理能力。批处理是提升吞吐量的关键，但受限于显存，很多请求只能在队列中等待。

内存墙：硬件层面的结构性困境

碎片化是软件层面的问题，但即使解决了碎片化，还有一个更深层的硬件瓶颈在等着我们——内存墙（Memory Wall）。

2024 年，加州大学伯克利分校的研究团队在《AI and Memory Wall》论文中揭示了一个令人不安的趋势：过去 20 年间，服务器级硬件的峰值算力（FLOPS）以每两年 3.0 倍的速度增长，而 DRAM 带宽和互连带宽仅以每两年 1.6 倍和 1.4 倍的速度增长。

xychart-beta
    title "硬件性能增长对比（20年跨度）"
    x-axis ["峰值算力", "DRAM带宽", "互连带宽"]
    y-axis "增长倍数" 0 --> 60000
    bar [60000, 100, 30]

这意味着什么？算力增长了 60,000 倍，但内存带宽只增长了 100 倍。对于 LLM 推理来说，这是致命的。

为什么解码阶段受影响最大？

Transformer 模型的推理分为两个阶段：

Prefill 阶段：处理整个输入 prompt，可以使用矩阵-矩阵乘法，计算密集度高，算术强度（FLOPs per byte）可达 200 以上。

Decode 阶段：逐个生成 token，每次只能使用矩阵-向量乘法，算术强度骤降到个位数。

算术强度低意味着什么？意味着 GPU 大部分时间在"等数据"，而不是"算数据"。计算单元极其强大，但数据传输速度跟不上，就像用一根吸管给消防车加水。

对于 OPT-13B 模型，单个 token 的 KV Cache 约 800KB。假设 GPU 带宽 2TB/s，读取这些数据需要约 0.4 微秒。但如果是批量推理，每个请求都要读取自己的 KV Cache，带宽压力成倍增加。

PagedAttention：借鉴操作系统的智慧

面对碎片化困境，2023 年 UC Berkeley 团队提出了 PagedAttention，核心思想借鉴自操作系统的虚拟内存分页管理。

从分段到分页：一场内存管理的革命

早期的操作系统使用分段式内存管理，为每个进程分配一段连续的物理内存。问题显而易见：进程结束后留下不规则的空洞，新进程可能因为找不到足够大的连续空间而无法加载。

分页式管理彻底改变了这一切。物理内存被划分为固定大小的"页"（page），进程的虚拟地址空间也按同样大小划分。一个进程的虚拟页可以映射到任意物理页，不需要连续。这样，无论物理内存多么碎片化，总能找到可用的页。

PagedAttention 将这一思想引入 KV Cache 管理：

  • 逻辑块（Logical Block）：每个请求视角下的连续 KV Cache，按固定大小（如 16 个 token）划分
  • 物理块（Physical Block）：GPU 显存中实际存储的数据块
  • 块表（Block Table）：记录逻辑块到物理块的映射关系

当生成新 token 时，系统只需找到任意一个空闲物理块，将其映射到请求的逻辑块序列末尾。不需要预留，不需要连续，内存利用率接近 100%。

Copy-on-Write：支持复杂解码策略

很多 LLM 应用需要复杂的解码策略，如并行采样（为同一 prompt 生成多个回答）和束搜索（Beam Search）。这些场景下，多个序列会共享部分 KV Cache。

PagedAttention 引入了操作系统的写时复制（Copy-on-Write）机制。当多个序列共享同一物理块时，系统只维护一个引用计数。当某个序列需要修改该块时，才真正复制数据。对于并行采样，prompt 部分的 KV Cache 完全共享，只有最后一个块可能需要复制。

vLLM 论文的实验数据显示，PagedAttention 在并行采样场景可节省 6%-30% 的内存，在束搜索场景可节省高达 55%-66% 的内存。

性能提升的代价

天下没有免费的午餐。PagedAttention 的动态映射引入了额外的开销：

  1. 块表查询：每次注意力计算都需要先查询块表，增加了内存访问
  2. 非连续访问：物理块分散在显存各处，可能影响缓存命中率
  3. 内核复杂度：需要专门的 CUDA 内核处理分块读取和计算

实测显示，PagedAttention 的注意力内核延迟比高度优化的连续内存实现高 20-26%。但这个代价完全值得——更高效的内存利用意味着更大的批量，而大批量带来的吞吐提升远超内核开销。

vLLM 论文报告，相比当时最先进的 Orca 系统，PagedAttention 带来了 2-4 倍的吞吐提升，在某些场景下甚至更高。

FlashAttention：从 IO 视角重新设计

PagedAttention 解决了内存管理问题，但注意力计算本身的开销还没解决。2022 年，斯坦福大学团队提出的 FlashAttention 从一个全新的角度切入——IO 感知（IO-awareness）。

GPU 内存层级：速度与容量的权衡

GPU 的内存并非铁板一块，而是分层的：

  • HBM（高带宽内存）：几十 GB 容量，带宽约 2TB/s
  • SRAM（片上缓存）：几 MB 容量，带宽约 19TB/s，延迟低一个数量级

传统注意力算法将所有中间结果写入 HBM：计算 Q×K 得到注意力分数矩阵，写入 HBM；读取后做 Softmax，再写入 HBM；读取后乘 V，最后写入 HBM。对于序列长度 N，这需要 O(N²) 次 HBM 读写。

FlashAttention 的核心洞察是：与其频繁读写 HBM，不如把计算"搬"到 SRAM 里完成。具体策略是分块计算（Tiling）：

  1. 将 Q、K、V 矩阵分成小块，每块小到可以放入 SRAM
  2. 在 SRAM 内完成注意力计算的所有步骤
  3. 只将最终结果写回 HBM

这样，HBM 读写次数从 O(N²) 降到 O(N²/M)，其中 M 是 SRAM 大小。更重要的是，注意力矩阵本身不需要存储，显著降低了峰值内存占用。

分块 Softmax 的数学技巧

分块计算的最大挑战是 Softmax——它需要知道整个序列的信息才能归一化。FlashAttention 采用了一种巧妙的在线 Softmax 算法：

  1. 分块计算时，维护两个累加器：最大值和指数和
  2. 每处理一个新块，更新这两个值并调整之前的结果
  3. 最终得到与全局 Softmax 数学等价的结果

这需要额外的计算（重归一化），但相比于减少的 HBM 读写，整体效率反而更高。

性能数据

FlashAttention 论文报告了显著的性能提升：

  • BERT-large 训练：相比 MLPerf 1.1 记录，端到端加速 15%
  • GPT-2 训练（序列长度 1K）：加速 3 倍
  • 长序列任务（序列长度 1K-4K）：加速 2.4 倍

更重要的是，FlashAttention 使超长上下文成为可能。在此之前，处理 16K 甚至 64K 的序列几乎不可想象；有了 FlashAttention，这些变成了现实。

量化：用精度换空间

除了优化内存管理，另一个直接减少显存占用的方法是降低数据精度。这就是量化技术。

从 FP32 到 INT4：显存的瘦身之旅

传统模型使用 FP32（32 位浮点）存储权重，一个 13B 参数模型需要约 52GB。改用 FP16 或 BF16，显存需求减半；进一步量化到 INT8，再减半；INT4 甚至可以将显存需求降到原来的 1/8。

量化方法主要分两类：

训练后量化（Post-Training Quantization, PTQ）：在已训练好的模型上直接转换精度，无需重新训练。主流方法如 GPTQ、AWQ 可以在几乎不损失精度的情况下将模型量化到 INT4。

量化感知训练（Quantization-Aware Training, QAT）：在训练过程中模拟量化效果，让模型适应低精度表示。精度更高，但需要额外的训练开销。

精度损失的实际影响

量化的代价是精度损失。但研究表明，对于 LLM，INT8 甚至 INT4 量化通常只带来可接受的精度下降。

根据 Hivenet 的量化指南，GPTQ-INT8 量化可以实现：

  • 1.6 倍加速
  • 7.5 倍容量增加
  • 统计上零精度损失

但并非所有场景都适合激进量化。对于需要高精度的任务（如数学推理、代码生成），INT8 可能比 INT4 更稳妥。实际部署中，需要针对具体任务进行基准测试。

其他优化技术的权衡

梯度检查点（Gradient Checkpointing）

也叫激活重计算（Activation Recomputation），核心思想是：在前向传播时不保存所有中间激活值，只保存部分"检查点"；反向传播时需要时重新计算。

这可以将激活值的内存占用减少 5 倍，但代价是增加 20%-30% 的计算量。对于训练阶段，这是一个经典的"空间换时间"到"时间换空间"的反转。

模型并行（Model Parallelism）

当单个 GPU 放不下模型时，需要将模型拆分到多个 GPU：

张量并行（Tensor Parallelism）：将单个层（如注意力头、线性层）拆分到不同 GPU。需要高带宽互连，延迟敏感。

流水线并行（Pipeline Parallelism）：将不同层分配给不同 GPU，形成流水线。对互连带宽要求较低，但存在"气泡"问题。

混合并行是当前大规模模型训练和推理的主流方案，但部署复杂度显著增加。

CPU Offloading

将部分模型权重或 KV Cache 卸载到 CPU 内存甚至 SSD，需要时再加载回 GPU。微软的 DeepSpeed ZeRO-Offload 是典型代表。

这可以突破 GPU 显存限制，但代价是高昂的数据传输开销。仅适合对延迟不敏感的离线推理场景。

工程实践：如何诊断和优化显存问题？

显存监控工具

nvidia-smi：最基础的工具，显示总体显存使用情况。但注意，它显示的是已分配内存，不反映碎片化程度。

PyTorch Memory Snapshot：提供细粒度的显存分配可视化，可以追踪每个张量的分配和释放时间，帮助定位内存泄漏和碎片化热点。

import torch
torch.cuda.memory._record_memory_history(max_entries=100000)
# ... 运行你的模型 ...
snapshot = torch.cuda.memory._snapshot()
# 保存快照用于分析

常见 OOM 场景与对策

场景一：显存看似充足但 OOM

  • 原因：碎片化
  • 对策：使用 PagedAttention（vLLM）、减小 max_seq_len、重启进程清理碎片

场景二：长序列推理 OOM

  • 原因：KV Cache 随序列长度线性增长
  • 对策：量化（INT8/INT4）、减少批量大小、使用 FlashAttention

场景三：微调训练 OOM

  • 原因：优化器状态和梯度占用大量显存
  • 对策：ZeRO Offload、梯度检查点、LoRA（只微调部分参数）

硬件演进：HBM 的崛起

从硬件角度看，解决内存墙的根本在于提升带宽。近年来，HBM（高带宽内存）技术快速演进：

内存类型 典型带宽 容量 延迟
GDDR6 约 1TB/s 较大 较高
HBM2 约 2TB/s 较小 较低
HBM2E 约 3TB/s 更大 低
HBM3 约 3TB/s+ 更大 更低

HBM 采用 3D 堆叠设计，将多个 DRAM 芯片垂直堆叠，与 GPU 裸片近距离连接。相比 GDDR 的 PCB 布线，HBM 的信号路径更短，带宽更高，延迟更低。

但 HBM 也有限制：容量增长较慢，成本更高。从 A100 到 H100，FLOPS 增长超过 2 倍，但显存容量仍停留在 80GB 上限。这意味着即使硬件不断升级，内存墙问题仍将持续存在。

未来展望

显存瓶颈不会在短期内消失，但技术正在快速演进：

算法层面：更高效的注意力机制（如线性注意力、稀疏注意力）、更激进的量化技术（如 1.58-bit 量化）

系统层面：更智能的内存管理、更高效的内核实现、更灵活的调度策略

硬件层面：HBM 容量和带宽持续提升、片上存储增大、专用 AI 加速器

终极解决方案可能是打破 Transformer 的框架，探索线性时间复杂度的新架构。但在那之前，PagedAttention、FlashAttention 和量化技术的组合，已经为我们争取到了宝贵的时间。

显存优化不是单一技术的问题，而是从硬件架构到算法设计的系统工程。理解内存墙、碎片化、带宽瓶颈的本质，才能在有限的硬件资源中榨取每一滴性能。

参考文献

  1. Kwon, W., et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. SOSP ‘23.
  2. Dao, T., et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022.
  3. Gholami, A., et al. (2024). AI and Memory Wall. IEEE Micro, 2024.
  4. NVIDIA. (2023). Mixed Precision Training Documentation.
  5. Dettmers, T., et al. (2023). GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers.
  6. NVIDIA. (2023). FasterTransformer: Transformer-related optimization.
  7. Yu, G., et al. (2022). Orca: A Distributed Serving System for Transformer-Based Generative Models. OSDI ‘22.
  8. Ren, J., et al. (2021). ZeRO-Offload: Democratizing Billion-Scale Model Training. USENIX ATC ‘21.
  9. Williams, S., et al. (2009). Roofline: An Insightful Visual Performance Model for Multicore Architectures. Communications of the ACM.
  10. Hugging Face. (2024). Model Memory Anatomy Documentation.