第八章运行时状态与显存

参数显存回答“模型权重放在哪里”，KV Cache 回答“历史 token 的信息如何复用”。普通推理不改变权重，却会不断追加缓存；长上下文和高并发因此可能耗尽显存。

1 为什么只缓存 K 和 V

在因果 decoder 中，旧位置的 hidden states 不依赖后来的 token，所以旧位置在各层产生的 K/V 可复用。新位置的 Query 需要与历史 K 匹配，再聚合历史 Value。

旧 Query 不参与新位置的输出计算，因此标准自回归解码不需要长期保存它。KV Cache 也不是只缓存最后一层：每一层有自己的 K/V，下一步都会用到。

2 缓存大小从形状推导

每层的 K 与 V 各有形状 $[B,h_{kv},T,d_h]$。同构全注意力模型的缓存字节数为：

$$ \boxed{M_{\rm KV}=2LBTh_{kv}d_hb_{kv}=2LBTd_{kv}b_{kv}} $$

这里的 2 来自 K 和 V，不是 FLOPs 的乘加系数。不同请求长度为 $T_i$ 时：

$$ M_{\rm KV}=2Ld_{kv}b_{kv}\sum_{i=1}^{B}T_i $$

单条序列追加一个已处理 token：

$$ \boxed{\Delta M_{\rm KV,one}=2Ld_{kv}b_{kv}} $$

整个 batch 每条序列各追加一个 token，则再乘 $B$。若题目只问元素数，不能乘字节数。

例如 $L=32,h_{kv}=8,d_h=128,b_{kv}=2$：每个 token 增加 $131072$ bytes，即 128 KiB；一条 4096-token 序列占 512 MiB；32 条占 16 GiB。这是未计分页和元数据的逻辑缓存大小。

3 总显存不能只算权重

一个实用预算式是：

$$ M_{\rm device}\ge M_{\rm weights}+M_{\rm KV}+M_{\rm temporary}+M_{\rm runtime}+M_{\rm reserve} $$

其中 $M_{\rm weights}\approx P_{\rm total}b_w$；临时空间包括中间激活、算子 workspace、logits；运行时还可能包含通信缓冲区、图捕获内存和分配器保留空间。各项峰值有生命周期重叠，不能把所有张量大小随意相加就当成精确峰值。

若扣除权重及其他预留后，可用于 KV 的空间为 $M_{\rm budget}$，则逻辑 token 容量上限为：

$$ N_{\rm cached,max}=\left\lfloor\frac{M_{\rm budget}}{2Ld_{kv}b_{kv}}\right\rfloor $$

等长请求最大并发还要除以每条请求预期缓存长度。请求会持续增长，因此只按当前 prompt 长度准入可能在生成中途达到容量限制。

4 激活显存为什么没有一个通用常数

一个 hidden tensor 为 $BSdb_a$ bytes；一个 FFN 中间 tensor 为 $BSmb_a$ bytes。SwiGLU 可存在 gate、up、乘积等中间结果，但是否同时落到 HBM、能否复用缓冲区，依赖融合实现。

朴素 Attention 如果显式保存分数或概率矩阵，单个矩阵大小为：

$$ B h_q S^2 b_a $$

例如 $B=1,h_q=32,S=4096,b_a=2$，一个矩阵就是 1 GiB。若多个中间量同时存在，显存还会更高。

FlashAttention 通过分块与在线 Softmax，避免将完整注意力矩阵写入 HBM，从而显著降低中间存储与 IO；它仍是精确注意力，数学上的全注意力主项仍随 $S^2$ 增长。FlashAttention 论文

纯推理中，处理完一层后可释放或复用不少临时激活，因此不能自动把一层临时激活峰值乘以 $L$。训练为了反向传播通常要保留或重计算更多信息，显存模型完全不同。

5 分页和前缀共享

PagedAttention 将 KV 按块管理，减少大段连续预分配造成的浪费，并支持缓存共享。PagedAttention 论文

若块大小为 $q$ 个 token，不共享时物理分配的 token 槽位约为：

$$ \sum_i q\left\lceil\frac{T_i}{q}\right\rceil $$

再乘每 token KV 大小，并计入页表等元数据。逻辑缓存每增加一个 token 都增加有效内容，物理分配则可能在跨页时一次增加一整页。

前缀共享使多个请求的相同前缀只保留一份物理 KV，改变的是实际存储的唯一缓存块数。它不意味着所有请求的注意力读取工作都消失；不同 Query 仍需访问前缀内容。

6 量化与多卡

权重量化改变 $b_w$，KV 量化改变 $b_{kv}$，两者是独立选择。INT4 权重裸数据约 0.5 byte/parameter，但 scale、zero-point、分组、padding 与高精度保留层会增加实际占用。不能因为权重是 4-bit，就把 KV 自动按 4-bit 计算。

Tensor Parallel 的 KV 能否均分，取决于 head 分配。若 $h_{kv}$ 可被并行度整除且实现确实切分，可以按对应份额计算；若并行度大于 KV head 数，可能复制 KV heads。Pipeline Parallel 则主要按每卡持有的层数计算 KV。单卡显存预算必须按真实分片与复制关系逐项列出。

标准 KV 公式还需要针对以下结构修改：

滑动窗口层只保留有效窗口时，用受窗口限制的缓存长度；混合全局层和窗口层逐层相加。
MLA 使用压缩的 latent 等缓存，应按实际缓存张量形状重算。
Encoder-decoder 可能额外缓存交叉注意力的 K/V。
Beam search、前缀共享和推测解码会改变分支、共享和临时缓存需求。

7 训练显存的延伸

训练还要考虑梯度、优化器状态、可能存在的 FP32 master weights、保存的激活。一个常见但非普适的混合精度 Adam 预算是：2-byte 权重、2-byte 梯度、4-byte master weights、两个 4-byte 动量，共 16 bytes/parameter；若梯度采用 FP32，则变成 18 bytes/parameter。优化器实现与分片策略不同，数值会变化。

MoE 的“激活参数较少”并不自动让全部训练状态只为这部分分配。是否分片、按需分配或卸载，要看具体训练系统。

下一章：Prefill 与 Decode 计算量。

第八章 运行时状态与显存