一、显存基础与瓶颈分析

大家好，我是你们的老朋友。今天咱们聊聊显存——这个在TensorRT-LLM推理中绕不开的话题。

说实话，我刚开始接触大模型推理时，也踩过不少坑。有一次部署一个70B的模型，明明算力够，但就是跑不起来。后来一查，显存爆了。嗯，从那以后，我养成了一个习惯：先搞清楚显存怎么用，再动手优化。

1.1 GPU显存架构：HBM与L2 Cache

先说说GPU的显存架构。你想想看，GPU处理数据，就像工厂流水线。HBM（高带宽内存）是原材料仓库，L2 Cache是车间里的临时货架。

HBM的特点：

  • 容量大——通常几十GB，能装下整个模型
  • 带宽高——A100的HBM带宽约2TB/s，H100能到3.35TB/s
  • 延迟高——访问一次要几百个时钟周期

L2 Cache的特点：

  • 容量小——A100的L2 Cache是40MB，H100是50MB
  • 带宽极高——是HBM的几倍
  • 延迟低——几十个时钟周期就能拿到数据

我个人习惯把L2 Cache看作「热数据缓冲区」。那些频繁访问的权重、激活值，如果能留在L2里，推理速度能快不少。我在项目中遇到过，同样的模型，L2命中率从60%优化到85%，端到端延迟直接降了30%。

核心观点：显存优化的本质，就是让数据尽量待在L2 Cache里，少去HBM「搬砖」。

下面这张图，是我自己总结的GPU显存层次结构：

GPU显存层次结构（以A100为例） HBM2e（高带宽内存） 容量：80GB | 带宽：2TB/s | 延迟：~400 cycles L2 Cache 容量：40MB | 带宽：~10TB/s | 延迟：~50 cycles SM（流式多处理器） + L1/Shared Memory 每个SM 192KB L1/Shared Memory | 延迟：~5 cycles 数据流动方向：HBM → L2 Cache → SM L1/Shared Memory

1.2 TensorRT-LLM推理中的显存占用模型

搞清楚了硬件，咱们再看看软件层面。TensorRT-LLM推理时，显存主要被三样东西吃掉：模型权重、KV Cache、激活内存。

模型权重：

  • 就是模型参数本身，比如70B模型，FP16精度下约140GB
  • 量化后能降到70GB（INT4）甚至35GB（INT8）
  • 这部分是「死内存」，加载后基本不变

KV Cache：

  • 存储Attention层的Key和Value
  • 随序列长度线性增长——序列越长，吃得越多
  • 我见过最夸张的情况：一个8K长度的对话，KV Cache占了总显存的60%

激活内存：

  • 中间计算结果，比如LayerNorm的输出、FFN的中间激活
  • 跟batch size强相关——batch越大，激活内存越大
  • 这部分是「临时内存」，用完就释放

个人经验：我曾经优化过一个对话系统，发现激活内存占了40%的显存。后来通过「激活重计算」技术，把峰值降到了15%。说白了，就是用时间换空间——多算一次，少存一点。

下面这个表格，是我整理的不同配置下的显存占用估算：

模型规模 精度 权重显存 KV Cache（4K长度） 激活内存（batch=1） 总计
7B FP16 14 GB 2 GB 1 GB 17 GB
13B FP16 26 GB 4 GB 2 GB 32 GB
70B INT4 35 GB 8 GB 4 GB 47 GB
70B FP16 140 GB 8 GB 4 GB 152 GB

注意：上面的估算只是「裸模型」的占用。实际部署时，还要加上TensorRT-LLM的运行时内存、CUDA context、以及各种临时缓冲区。我建议至少预留20%的余量。

1.3 显存瓶颈的Profiling方法

光知道理论没用，得能「看见」显存到底怎么用的。Profiling就是干这个的。

我常用的工具有三个：

  1. nvidia-smi——最基础，看个大概
  2. Nsight Systems——时间线视图，看显存分配和释放
  3. Nsight Compute——深入kernel级别，看L2命中率、带宽利用率

举个例子，用Nsight Systems抓显存分配：

# 命令行启动
nsys profile -o trace_output -t cuda,nvtx \
  --gpu-metrics-device=0 \
  python run_inference.py

# 然后打开生成的 .nsys-rep 文件
# 在「CUDA Memory」视图里，你能看到：
# - 每个alloc/free的时间点
# - 峰值显存占用
# - 哪些操作在「吃」显存

我个人习惯先跑一次nvidia-smi看看整体情况：

# 每0.5秒刷新一次，显示显存使用
watch -n 0.5 nvidia-smi --query-gpu=memory.used,memory.total --format=csv

# 输出示例：
# memory.used [MiB], memory.total [MiB]
# 45230, 81920
# 45310, 81920
# ...

如果发现显存占用异常高，我会用Nsight Compute深入看kernel的L2命中率：

# 对特定kernel做profiling
ncu --set full -o kernel_profile \
  --target-processes all \
  python run_inference.py

# 关键指标：
# L1命中率：越高越好，目标>80%
# L2命中率：目标>60%
# HBM带宽利用率：目标>70%

避坑指南：我曾经遇到过一个情况，nvidia-smi显示显存只用了60%，但程序报OOM。后来发现是「显存碎片化」导致的——虽然总量够，但连续的大块内存不够。这种情况，用Nsight Systems的「Memory Fragmentation」视图才能看出来。

最后，我总结一下Profiling的「三步走」策略：

  • 第一步：nvidia-smi看总量，确认有没有爆显存
  • 第二步：Nsight Systems看时间线，找到显存峰值点
  • 第三步：Nsight Compute看细节，定位L2命中率、带宽瓶颈

嗯，这一章的内容就到这里。显存基础打牢了，后面咱们才能谈优化。记住一句话：你没法优化你看不见的东西。先学会Profiling，再谈优化策略。

专注资料整理
← 上一章 📖 返回目录 下一章 →