LLM 推理显存计算法 — 比权重先爆的是 KV cache

引言 — “这个模型，能装进我们的 GPU 吗”

如果对这个问题的回答是“得跑跑看才知道”，那在基础设施预算会议上就没什么可说的了。好在推理内存和训练不同，几乎全部可以用算术预测，因为既没有优化器状态，也没有梯度。

问题在于，大多数计算只做了一半。“7B 模型 fp16 是 14GB，所以 24GB 的卡装得下”这一步是对的。可接下来把批大小提到 32、上下文拉到四千 token 的那一刻，进程就因显存不足而死掉了。缺的那一项是 KV cache，而这一项和权重不同，会随用户流量的形状一起长大。

本文把推理显存拆成四块(权重、KV cache、激活、框架开销)，并整理每一块的计算方法。所有数字都以可自行验算的公式给出，建议您打开自己模型的 config 文件对照着读。

权重内存 — 一次乘法就结束的部分

权重是诚实的。参数量乘以每参数字节数，就完了。

BYTES_PER_PARAM = {
    "fp32": 4.0,
    "fp16": 2.0,   # bf16 相同
    "int8": 1.0,
    "fp8": 1.0,
    "int4": 0.5,   # 4bit，实际上还会多出缩放系数/零点
}

GIB = 1024 ** 3

def weight_gib(params_billions: float, dtype: str = "fp16") -> float:
    return params_billions * 1e9 * BYTES_PER_PARAM[dtype] / GIB

for n in (7, 13, 70):
    print(n, [round(weight_gib(n, d), 1) for d in ("fp16", "int8", "int4")])
# 7  [13.0, 6.5, 3.3]
# 13 [24.2, 12.1, 6.1]
# 70 [130.4, 65.2, 32.6]

这里先把单位说清楚。70 亿乘以 2 字节是 140 亿字节，除以 10 的 9 次方得 14GB，除以 2 的 30 次方得 13.0GiB。显卡规格书上的“80GB”多半是十进制表示，而 nvidia-smi 显示的数字是二进制表示，这 7% 的差距会让本来算得刚刚好的结果对不上。下面的计算全部统一以 GiB 为准。

4bit 那一项的数字是理想值。真实的 4bit 检查点会按分组保存 fp16 的缩放系数和零点，分组大小为 128 时，每个参数大约多出 0.5 比特。按每参数 0.53 字节而不是 0.5 字节来估算更安全。而且不少实现会把嵌入层和输出层留着不量化，词表越大的模型，这个差距越明显。

KV cache — 真正爆掉的地方

自回归解码每生成一个 token，都要对之前的全部 token 计算注意力。如果丢掉已经算过的键和值，每一步就得重算整条序列，所以要全部留着。这就是 KV cache。

一个 token 占用的字节数如下。

def kv_bytes_per_token(layers: int, kv_heads: int, head_dim: int,
                       bytes_per_elem: float = 2.0) -> float:
    # 2 是因为要同时保存 K 和 V
    return 2 * layers * kv_heads * head_dim * bytes_per_elem

# 7B 级 MHA 配置: 32 层，32 个 KV 头，头维度 128
print(kv_bytes_per_token(32, 32, 128) / 1024, "KiB")   # 512.0 KiB
# 70B 级 GQA 配置: 80 层，8 个 KV 头，头维度 128
print(kv_bytes_per_token(80, 8, 128) / 1024, "KiB")    # 320.0 KiB

一个 token 512KiB。如果这个数字没有直观感受，可以这样看。在 7B 级 MHA 模型上，一条 4096 token 上下文的请求，光 KV cache 就要 2GiB。32 条并发就是 64GiB，是 13GiB 权重的五倍。

这张表要横着读，而不是竖着读。权重是固定值，最后一列与流量成正比。把上下文从 4K 拉到 32K，只有最后一列变成 8 倍。那些宣传支持长上下文的服务实际上把并发压得很低运营，原因就在这里。

有一点要注意。KV cache 与存活 token 的总量成正比，而不是与活跃序列数成正比。100 条请求各 500 token，和 5 条请求各 10,000 token 占用的内存是一样的。容量规划要按并发 token 数来做，而不是按每秒请求数。

GQA 与 MQA — 只减少 KV 头的结构改动

上表中 7B 的 MHA 行和 GQA 行正好差 4 倍，是因为 KV 头从 32 个减到了 8 个。查询头仍然是 32 个，每 4 个一组共享一个 KV 头。

# 查询头不变，只减少 KV 头
n_heads, head_dim, layers = 32, 128, 32

for kv_heads, name in [(32, "MHA"), (8, "GQA-8"), (1, "MQA")]:
    per_token = kv_bytes_per_token(layers, kv_heads, head_dim)
    ratio = n_heads // kv_heads
    print(f"{name:6} kv_heads={kv_heads:2}  {per_token/1024:6.1f} KiB/token  ({ratio}배 절감)")
# MHA    kv_heads=32   512.0 KiB/token  (1배 절감)
# GQA-8  kv_heads= 8   128.0 KiB/token  (4배 절감)
# MQA    kv_heads= 1    16.0 KiB/token  (32배 절감)

节省比例正好等于查询头数除以 KV 头数。这不是近似，而是定义上就如此。

在质量方面，GQA 通常被视为一笔便宜的交易。原论文报告称，即使只有少数几个 KV 头，质量也接近 MHA，此后大量公开模型都采用了 8 个上下的 KV 头。相反，MQA 只有一个 KV 头，观测到质量下降的报告更多。不过这是预训练阶段的架构选择，不是已经部署的模型还能改的值。您能选的只到“选一个用 GQA 的模型”为止。

如果模型 config.json 里的 num_key_value_heads 小于 num_attention_heads，那就是 GQA。计算时把这两个搞混，结果会差 4 倍。

激活、开销，以及完整的计算器

剩下两块的大小随情况变化，很难给出精确公式。不过知道它们与哪些变量成正比就够了。

激活内存在解码阶段几乎不成问题，因为每一步每个批元素只流过一个 token。问题出在预填充。整段提示一次通过，中间张量与批大小乘以提示长度成正比。

def prefill_activation_gib(batch, prompt_len, hidden, intermediate,
                           live_buffers=4, bytes_per_elem=2.0):
    """预填充峰值激活的粗略下界。live_buffers 因框架而异，必须实测。"""
    per_token = (hidden + intermediate) * bytes_per_elem
    return batch * prompt_len * per_token * live_buffers / GIB

# 7B 级: hidden 4096, intermediate 11008
print(round(prefill_activation_gib(8, 4096, 4096, 11008), 2), "GiB")  # 3.69 GiB

live_buffers=4 是我随手定的常数。实际同时存活的缓冲区数量取决于算子融合程度和内存规划器的实现，所以这个值必须用实测来校正。唯一确定的只是它与批大小和提示长度的乘积成正比，也正因如此，分块预填充(把长提示切成若干段送入)能把这一项压成常数。

框架开销是 CUDA 上下文、cuBLAS 工作区、通信缓冲区和分配器余量的总和。每进程按 1GiB 上下预留通常就对了。在此之上，再考虑分配器碎片，给总量乘上 10~15% 的余量是务实的做法。

现在把它们合起来做成一个计算器。

def total_vram_gib(params_b, weight_dtype, layers, kv_heads, head_dim,
                   seq_len, batch, kv_bytes=2.0, overhead=1.15, fixed_gib=1.0):
    w = weight_gib(params_b, weight_dtype)
    kv = kv_bytes_per_token(layers, kv_heads, head_dim, kv_bytes) * seq_len * batch / GIB
    return (w + kv) * overhead + fixed_gib

# 70B，4bit 权重，GQA-8，8K 上下文，并发 16
print(round(total_vram_gib(70, "int4", 80, 8, 128, 8192, 16), 1), "GiB")
# 84.6 GiB  → 单张 80GiB 卡装不下

权重是 32.6GiB，而 KV cache 是 40GiB。就算把权重压到 4bit，KV cache 还是更大。忽略这一点，就会卡在“都量化了为什么还是不行”的状态里。

实务中用得更多的是反向计算：卡已经定了，能同时接多少条。

def max_concurrent(vram_gib, params_b, weight_dtype, layers, kv_heads, head_dim,
                   seq_len, util=0.9, kv_bytes=2.0, fixed_gib=1.0):
    usable = vram_gib * util - fixed_gib - weight_gib(params_b, weight_dtype)
    if usable <= 0:
        return 0
    per_seq_gib = kv_bytes_per_token(layers, kv_heads, head_dim, kv_bytes) * seq_len / GIB
    return int(usable / per_seq_gib)

print(max_concurrent(80, 70, "int4", 80, 8, 128, seq_len=8192))   # 15
print(max_concurrent(80, 70, "int4", 80, 8, 128, seq_len=2048))   # 61
print(max_concurrent(80, 70, "int4", 80, 8, 128, seq_len=8192, kv_bytes=1.0))  # 30

这三行各自对应一个运营决策。把上下文上限从 8K 降到 2K，吞吐变成 4 倍；把 KV cache 用 fp8 存，变成 2 倍。两个决策都会影响质量，但至少能看清得到了多少、又押上了什么。

PagedAttention — 算得对却用不上的内存

到目前为止的计算都假设“只用真正需要的那么多”。早期的服务实现做不到这一点。请求到来时，它们会按这条请求可能达到的最大长度，预先在连续空间里占好 KV cache。最大长度是 4096，而实际输出只有 200 token，剩下的就白白扔掉了。

vLLM 论文把这种浪费分成两类来解释：预留却不使用的内部碎片，以及夹在块之间谁也用不上的外部碎片。在论文测量的工作负载上，真正被有效使用的 KV 内存比例相当低，其余都消失在这两种碎片里。具体数值依赖于工作负载，所以与其照搬引用，不如记住这个结构：输出长度的方差越大，浪费越多。

PagedAttention 用操作系统的方式解决这个问题。它把 KV cache 切成固定大小的块(通常 16 个 token)，把逻辑上连续的序列映射到物理上分散的块上。结果是浪费被限制在每条序列最后一块的空位，也就是最多 15 个 token。

还有一个副作用，实际上它更重要。按块管理之后，多条序列可以共享同一个块。使用同一段系统提示的 100 条请求，只保留该前缀对应的一份块，通过引用计数共享。在前缀很长、请求量很大的服务里，这种前缀共享带来的节省比消除碎片还大。

# 在 vLLM 中实际会调整的参数
vllm serve <모델경로> \
  --max-model-len 8192 \            # 上下文上限。直接决定 KV cache 的上限
  --gpu-memory-utilization 0.90 \   # 全部显存中的使用比例。其余留给开销
  --kv-cache-dtype fp8 \            # 只把 KV cache 降到 8bit。容量减半
  --max-num-seqs 64 \               # 并发序列上限
  --enable-prefix-caching           # 复用公共前缀块

# 请在启动日志里确认这几行。它们应当与上面的计算吻合。
#   "GPU KV cache size: 129,024 tokens"
#   "Maximum concurrency for 8192 tokens per request: 15.75x"

如果启动日志打印的 KV cache token 数与前面 max_concurrent 的计算相差很大，那两者中必有一个是错的。多数情况下是把 KV 头数当成了查询头数。

量化不是免费的

把权重降到 4bit，内存变成四分之一。这是事实。问题在于这句话通常就到此为止了。

第一，质量损失很难被平均指标捕捉到。一个被概括为“WikiText 困惑度上升了 0.1”的变化，实际上常常在长输出的后半段、代码生成、多语言、格式遵守这些尾部区域表现得大得多。困惑度是平均 token 预测难度，而不是任务成功率。决定是否量化时，必须用您自己任务的评测集来衡量。这个评测集怎么造，在不靠感觉做 LLM 评测那一篇里单独讲过。

第二，4bit 并不总是更快。仅权重量化只在存储上是 4bit，运算仍用 fp16，所以每次矩阵乘法都要做一次反量化。批很小、内存带宽是瓶颈时，要读的字节数只有四分之一，确实会更快。反过来，批很大、算力成为瓶颈时，剩下的只有反量化开销，可能比 fp16 还慢。“量化了就会变快”是批大小为 1 的基准测试里的说法。

第三，KV cache 量化是与权重量化彼此独立的决策，而且通常更安全。fp8 的 KV cache 有很多报告说损失很小，而 4bit 的 KV 则有在长上下文下出现劣化的报告。正如前面的计算所示，长上下文服务里主导内存的是 KV 这一侧，所以把顺序反过来——先把 KV cache 降到 fp8，权重留在 fp16 或 8bit——往往是更好的组合。

归纳起来，选择顺序是这样：先看 KV cache 的 dtype，再看上下文上限，权重量化放到最后看。大多数团队都以完全相反的顺序切入，因而浪费时间。

结语 — 权重是常量，KV cache 是变量

要记住的只有一句话。权重内存是选定模型那一刻就确定的常量，KV cache 是由您决定接收的流量形状所决定的变量。容量不足几乎总是来自变量那一侧。

所以部署前需要算的值只有两个：每 token 的 KV 字节数(2 × 层数 × KV 头数 × 头维度 × 每元素字节数)，以及用剩余显存除以该值得到的总 token 数。有了这两个数字，上下文上限与并发上限之间的取舍就能在会议室里当场算出来。之后再谈量化也不迟。