KV Cache 优化与显存管理：长上下文推理的算账方法

说明：本篇为综述/无第一手实验数据，采用公开架构资料与技术解析路线。

实验目标

长上下文推理（128K、1M token）把推理的瓶颈从"算力"挪到了"显存里的 KV Cache"。一个 70B 模型，8K 上下文的 KV Cache 约 10-20GB；到 128K 就轻松破 200GB——远超模型权重本身。本篇目标：建立 KV Cache 的算账框架（体积怎么算、长在哪、怎么压），并梳理四类主流优化（PagedAttention、连续批处理、量化与剪枝、投机解码）各自的收益与边界，读完能对"这条推理链路上 KV 占了多少显存、还能压到多少"给出量化回答。

代码逐段拆解

第一段：KV Cache 体积核算

def kv_cache_bytes(n_layers, n_kv_heads, head_dim, seq_len, batch, dtype_bytes=2):
    """单次生成时的 KV Cache 占用（仅当前 seq_len，不含预分配余量）"""
    per_token_per_layer = 2 * n_kv_heads * head_dim * dtype_bytes  # K 一份 + V 一份
    return n_layers * per_token_per_layer * seq_len * batch

# 例：70B 类模型常见配置（估算口径）
print(kv_cache_bytes(80, 8, 128, 8192, 1) / 2**30, "GB @ 8K, bs=1")
print(kv_cache_bytes(80, 8, 128, 131072, 1) / 2**30, "GB @ 128K, bs=1")
# 输出约：12.6 GB / 8K；201.3 GB / 128K

要点：

1. KV Cache 随序列长线性增长，且与 batch 相乘：公式里 seq_len × batch 是乘积关系——并发越高、上下文越长，缓存膨胀越快。这就是为什么"长上下文 + 高并发"是推理系统最烧显存的组合。
2. 2 * n_kv_heads * head_dim：K 和 V 各一份。注意是 KV heads 而不是 Q heads——GQA（Grouped Query Attention）能砍缓存的关键就在这：70B 类模型常见 8 个 KV head，相比 MHA 的 64 个 Q head，KV 只有 Q 的 1/8。
3. dtype_bytes=2：FP16/BF16 每元素 2 字节。量化到 8bit 直接砍半，4bit 再砍半——KV 量化是"乘法"级的显存收益。

第二段：显存水位预算

def memory_budget(model_gb, kv_gb_per_req, max_concurrent, total_gb, safety=0.05):
    usable = total_gb * (1 - safety)
    reserved = model_gb + 2.0            # 权重 + CUDA 上下文/临时缓冲
    budget_for_kv = usable - reserved
    max_requests = int(budget_for_kv // kv_gb_per_req)
    return max_requests, budget_for_kv

# 例：H100 80G 单卡，70B 模型 BF16 140GB 装不下 → 量化到 4bit 70GB
max_req, budget = memory_budget(70, 12.6, 0, 80)
print(f"4bit 70B 权重 + 8K KV：单卡可并 {max_req} 个 8K 请求")

关键理解：

• 权重量化是前提：70B 的 BF16 权重 140GB 放不进 80GB 卡，先量化到 4bit（约 70GB），KV Cache 才有预算。推理侧的显存账 = 权重 + KV + 临时缓冲 + CUDA 开销，四笔都要算。
• safety=0.05 是防 OOM 的底线：显存碎片、动态形状、CUDA graph 的额外占用都要留 5% 余量，硬塞满的推理服务会在长序列请求下 OOM 崩掉。

原理讲解

KV Cache 为什么必须存在

Transformer 的注意力要"看"之前所有 token 的 K/V。把历史 K/V 重算一次（无缓存），复杂度随序列长度平方增长，128K 上下文根本跑不动；所以推理框架把每个已生成 token 的 K/V 存下来（KV Cache），后续每一步只计算新 token 的 K/V 并追加。用显存换时间——这就是 KV Cache 的本质。

四类主流优化与边界

优化	原理	典型收益	边界/代价
PagedAttention（vLLM）	把 KV 按页分配，消除碎片，按需扩容	并发提升 2-4 倍	需框架支持；页调度有 CPU 开销
连续批处理（Continuous Batching）	请求级调度，不等整批完成	吞吐提升 5-10 倍	调度器复杂度；长尾请求占显存
KV 量化 / GQA / 剪枝	8bit/4bit 存 KV、减少 KV head、稀疏剪枝	缓存体积 1/2-1/4	精度损失需评测；剪枝要重训练
投机解码（Speculative）	小模型起草，大模型并行验证	减少生成步数→少写 KV	接受率决定收益（见 12 章投机解码篇）

设计要点：

• PagedAttention 是"减少浪费"，量化是"减少总量"：前者让现有显存装更多请求，后者让每个请求更省——两者正交，生产环境通常叠加使用。
• 连续批处理解决"吞吐天花板"：传统静态 batching 里最慢的请求拖累整批；连续批处理让每个请求独立进出，GPU 永远在处理"有活干"的请求，吞吐和显存利用率同时上升。
• 长上下文是量化最划算的场景：8K 上下文 KV 只有 12GB，量化省 6GB 意义有限；128K 时省 100GB，量化从"锦上添花"变"必需品"。

显存账本（核算口径）

以 80GB 单卡推理 70B 为例（估算口径）：

项	占用	说明
权重（4bit 量化）	~70GB	必须量化才装得下
CUDA 上下文/临时	~2GB	框架固定开销
KV 预算	~4GB	剩余空间很小 → 单卡并发受限
结论	—	70B 长上下文推理基本需要多卡 TP 或更大显存卡

实验结果

（估算口径，非实测）

场景	无优化	+PagedAttention	+KV 8bit	+连续批处理
单卡并发 8K 请求（70B 4bit）	0-1	3-4	6-8	与上叠加
128K 请求可行性	否（显存爆）	勉强	可	可

结论：长上下文推理的组合拳 = 权重量化（让模型装得下）+ KV 量化（让缓存省一半）+ PagedAttention（让显存无碎片）+ 连续批处理（让吞吐拉满），四者缺一，128K 场景都会撞显存墙。

工程启示

1. 先算账再优化：用上面的公式算出"权重 + KV + 固定开销"三笔账，才知道瓶颈在哪——很多"推理慢"其实是 KV 挤爆显存导致换页/排队。
2. 量化优先做 KV 而非只做权重：权重量化常被优先做，但在长上下文场景，KV 量化收益可能更大（缓存体积减半直接翻并发）。
3. 监控要加 KV 维度：DCGM 看到的是显存总量，看不到"其中多少是 KV"——vLLM 的 metrics（vllm:num_preemptions、gpu_cache_usage_perc）要进监控面板，KV 打满预占水位就是扩容信号。
4. 把接受率纳入 KV 优化决策：投机解码（12 章已述）通过减少生成步数间接减少 KV 写入，与量化叠加时收益相乘，值得一起评估。

下一篇预告：本篇完成了"推理显存账本"，下一篇回到智算×超算双栈的应用侧——"科学计算代理模型：从高保真仿真到 AI 代理模型的降阶之路"。