🧸 ELI5 · Transformers 核心内功

别再从头重读了!KV Cache 记忆魔法

10 个前后咬合的卡通两格手绘漫画,带你彻底看透 Hugging Face Transformers 的缓存机制与显存演进历程。

小绿:正在计算的当前 Token 与自注意力头
小橙:新生成的 Query 探针与新吐出的词
小蓝:缓存里存着的老朋友(Key 与 Value 记忆条)
红虚线:因果掩码遮蔽 / 滑动窗口过期被遗忘淘汰

⚡ 核心主线:为什么大模型需要“记性”?

一句话真相: 大模型吐字是“一个字一个字蹦”的(自回归)。如果它是个严重的健忘症患者,每次吐第 1000 个字都要把前面 999 个字全部重新乘以一遍矩阵,计算量就会发生 $O(N^2)$ 算力大雪崩!KV Cache 就是它的小账本——算过的 Key 和 Value 记在账上,以后只算当前最新这 1 个字!

历史已存:K_past & V_past (1 ~ t-1 词免重算) 新词 q_t Softmax(Q K^T) V 仅做 1 次点积更新 词 t+1
已学完
1

自回归重复计算与 KV 缓存起源

💡 一句话解释: 不用缓存时,生成 1000 个字总共要算 50 万次矩阵;有了缓存,只要算 1000 次!
# 没有缓存:O(N^2) 耗时爆炸 | 启用缓存:O(N) 线性平稳 outputs = model(**inputs, past_key_values=past_key_values, use_cache=True)
烦恼:吐第100个字要重读前99个
O(N^2) 重复计算 每前进一步,全盘重新推演

每写一个字,从头背一遍全书

办法:老词做成书签,只算最新一个
历史已存好 K, V Pairs 只算最新 Token t 计算复杂度直降为 O(N)

记在小账本,只写新一笔

解决了什么:彻底消灭了生成过程中对历史 Token 的全量重复投影运算,推理速度飙升几十倍。
代价是什么:显存开销随生成长度线性膨胀,长文本生成极容易遭遇 GPU 显存爆炸(OOM)。
已学完
2

因果掩码的魔法 (Causal Mask)

💡 一句话解释: 为什么老词的 K 和 V 能够原封不动复用?因为时间不能倒流,后说的话绝不能影响前面的字!
$$ ext{Attention}(Q, K, V) = ext{softmax}\left( rac{Q K^ op}{\sqrt{d_{ ext{head}}}} imes ext{mask} ight) V$$
困惑:以后产生的新词会不会改写历史?
❌ 老词 未来词 双向模型 BERT 必须全部重算

若是未来改过去,缓存全报废

办法:下三角单向因果掩码遮蔽未来
下三角 🔒恒定 K_past / V_past 终身有效!

单向时间流,历史永不变

解决了什么:因果掩码阻断了未来信息对过去的篡改,奠定了自回归生成可以安全复用历史向量的数学基石。
代价是什么:仅适用于自回归单向语言模型(如 GPT、Llama),掩码语言模型(如 BERT)无法使用此缓存机制。
已学完
3

Transformers Cache 抽象基类

💡 一句话解释: 模型有几十个 Transformer 层,不能到处乱塞,官方抽象出标准 `Cache` 基类统一对账!
new_K, new_V = cache.update(k_t, v_t, layer_idx) # 自动按层(layer_idx)管理 [batch, num_heads, seq_len, head_dim]
烦恼:以前几十层元组到处传参
32层 past_key_values 元组乱飞

元组满天飞,代码极难维护

办法:分层抽屉柜,统一 update 存取
Layer 0 [B, H, T, D] Layer 1 [B, H, T, D] Layer 31 [B, H, T, D] 面向对象管理,清爽直观

标准抽屉柜,按层归档

解决了什么:将以往脆弱杂乱的 `tuple(tuple(tensor))` 重构成可扩展的 `Cache` 对象,解耦底层存储。
代价是什么:打破了早期的元组签名习惯,自定义解码循环时需要主动维护注意力掩码形状。
已学完
4

动态按需扩容 (DynamicCache)

💡 一句话解释: 每蹦出一个字,就把新的 K 和 V 往张量尾巴上一贴,像不断拉长的小火车。
from transformers import DynamicCache past_key_values = DynamicCache() # 随着 token 增加在 seq_len 维度动态 torch.cat
烦恼:每次加车厢都要重搭轨道
原内存 重开新空间 频繁重新分配显存与拷贝

反复重新找地盘,效率有损

办法:最灵活开箱即用,自动拼接
t=1 t=2 t=3 new 零配置默认启用,开箱即用

自动贴尾巴,最省心简单

解决了什么:无需预先指定生成长度,根据实际生成字数动态弹性分配,内存浪费最小。
代价是什么:张量形状不断变化(Dynamic Shape),完全无法使用 CUDA Graph 和 torch.compile 加速。
已学完
5

编译加速神装:静态缓存 (StaticCache)

💡 一句话解释: 一开局就预先包下一整排 1024 个固定座位,新字直接按号入座,形状再也不变!
from transformers import StaticCache # 预先固定最大长度,完美解锁 torch.compile 与 CUDA Graph cache = StaticCache(config=model.config, max_batch_size=1, max_cache_len=1024, dtype=torch.bfloat16)
烦恼:形状一变,编译器当场崩溃
Len: 12 Len: 13 💥 torch.compile 触发反复重编译

忽长忽短,硬件图无法固化

办法:开局圈好固定包厢,指针就位
Shape 永远不变,跑满 CUDA 图

固定包厢,流水线全速狂飙

解决了什么:张量尺寸在生命周期内恒定,完美支持 `torch.compile` 和 CUDA 图录制,吞吐翻倍。
代价是什么:无论实际说几个字,显存必须按 `max_cache_len` 顶格预先占满,短文本会浪费显存。
已学完
6

滑动窗口遗忘法 (SlidingWindowCache)

💡 一句话解释: 只保留最近 4096 个字的清晰记忆,来了新词就把最古老的老词踢出窗口(金鱼的记忆法)!
# 针对 Mistral / Gemma 等支持局部注意力模型,窗口满了自动左移剔除 cache = StaticSlidingWindowCache(config, max_cache_len=4096, sliding_window=4096)
烦恼:聊了上万句,显存彻底爆仓
无限堆积的历史 💥 显存随着对话长度无限线性上升

什么都不忘,迟早要撑死

办法:固定窗框移动,新来旧退
踢出 最近 W 个词 新入 显存彻底封顶,无限畅聊

窗口滑动,显存永远恒定

解决了什么:显存开销固定为 $O(W)$(常数级),再长上下文生成也不会显存溢出。
代价是什么:模型彻底遗忘窗口之外的前文细节,需搭配注意力汇(Attention Sink)防止文本崩溃。
已学完
7

记忆瘦身减脂:量化缓存 (QuantizedCache)

💡 一句话解释: 原来存一个数要用 16 位的胖箱子(FP16/BF16),压缩成 8 位甚至 4 位的小便签存(INT8/INT4)!
from transformers import QuantoQuantizedCache # 或 HqqQuantizedCache # 内存直降 50% ~ 75%,原本只能跑 8K 的卡能跑 32K!
烦恼:高精度浮点数塞满显卡
FP16 (2B) FP16 (2B) 70B 模型千字缓存吞噬数十 GB

箱子太大,卡车很快装满

办法:无损压缩为 INT8 / INT4 压缩包
INT4 INT4 INT4 INT4 容量瞬间缩水 4 倍!

真空打包,容量轻松翻几倍

解决了什么:在显存极其紧张的端侧或单卡上,以几乎可忽略的精度损失将 KV Cache 显存砍半甚至变为 1/4。
代价是什么:每次读取与写入时伴随轻微的反量化/量化运算开销。
已学完
8

架构瘦身:分组查询注意力 (GQA / MQA)

💡 一句话解释: 以前 32 个侦探(Q 头)每人都要配一套完整的卷宗(KV),现在每 4 个侦探共享 1 份卷宗!
# Llama-3 / Qwen 等现代模型默认架构:num_key_value_heads = 8, num_attention_heads = 32 (GQA 4:1)
烦恼:32 个头各自存一份纯属浪费
MHA: 32 个头配 32 份 KV

人手一份,严重冗余

办法:一组人共享一份核心档案
Shared 1×KV KV 缓存直接缩减到原来的 1/4 ~ 1/8

合资拼单,效果完全不减

解决了什么:从模型骨架层面将每层需缓存的 KV 头数由 32/64 缩减为 8/4,大幅减轻显存带宽瓶颈。
代价是什么:模型需在预训练阶段按 GQA 结构训练,无法对旧版 MHA 模型无缝直接启用。
已学完
9

前缀命中复用 (Prompt Caching)

💡 一句话解释: 当你多次对同一个 PDF 提问,或者共用长达数千字的 System Prompt 时,这段前缀只算一次并锁进库里!
# 提示词完全相同的前缀块生成哈希,下次请求直接从缓存提取,首字等待时间(TTFT)暴降 80%
烦恼:换个问题,整篇论文又重读一次
问答 1 重算 Prefill 问答 2 又从头算一遍 算力大部分浪费在已知前文

反复重读同一本书,极慢极贵

办法:前缀打上哈希戳,直接秒读快照
公共前缀KV Hash #A8F2 秒回 TTFT 暴跌至数十毫秒

熟读成诵,换个问题秒接

解决了什么:跨多轮对话或同文档并发请求时,跳过庞大 System Prompt 的 Prefill 算力消耗。
代价是什么:需要缓存调度器维护前缀树(Radix Tree)并在请求稍有改动时及时置换。
下一站
10

生产终极形态:分页内存池 (PagedAttention)

💡 一句话解释: 当几百个人同时在线聊天时,像操作系统的虚拟内存一样,把显存切成小分页(Page),零碎片极致利用!
# vLLM / SGLang 等现代化推理服务器的底层支柱:彻底告别连续显存预分配 Logical KV Block ➔ Page Table ➔ Physical GPU Pages (显存利用率达 96%+)
烦恼:连续预留显存,到处是空闲碎片
实际用 预留未用虚胖显存 60% 以上的显存死于碎片与虚占

虚占茅坑不拉屎,显存极度浪费

办法:按需分配 16-Token 物理小切片
物理空间 100% 紧密拼贴

小块拼图,并发翻三倍

解决了什么:将非连续显存物理块通过页表映射串联,彻底消灭显存碎片,支持超大规模高并发服务。
代价是什么:需要重写底层 GPU 算子(非标准 PyTorch 连续张量原生算子),工程实现极其复杂。

📏 性能衡量尺:显存占用 vs 算力加速的核心权衡

在 Transformers 推理落地中,KV Cache 的选型永远在以下四个维度相互博弈:

推理吞吐 Throughput 显存消耗 VRAM (GB) 编译优化 Graph Fit 首字延迟 TTFT (ms) 支点:单卡硬件物理显存上限 (VRAM Wall)

🚀 静态缓存 (StaticCache)

大幅改善: 支持 torch.compile,吐字极速

牺牲妥协: 预分配固定显存,短序列有空间闲置

🪟 滑动窗口 (SlidingWindow)

大幅改善: 显存彻底封顶,无限长度不 OOM

牺牲妥协: 丧失对超长远古背景细节的全局记忆

🗜️ 量化缓存 (Quantized)

大幅改善: 显存直接缩水 50%~75%

牺牲妥协: 轻微反量化计算开销与极端长尾精度损失

👶 零基础版:手把手算一次你的显存够不够?

很多同学经常问:“为什么我买了 24GB 显卡,加载完 14GB 的模型,聊两句长天就爆显存了?” 让我们用公式算一次 70 亿参数模型(Llama-3-8B) 吐 4096 个字需要吃掉多少缓存:

经典计算公式:

$$ ext{KV Cache 显存} = 2 imes ( ext{层数 } L) imes ( ext{键值头数 } H_{kv}) imes ( ext{每个头维度 } D) imes ( ext{序列长度 } S) imes ( ext{精度字节数 } B)$$
  • 系数 2:因为既要存 Key,又要存 Value;
  • Llama-3-8B 参数:层数 $L=32$、GQA 键值头数 $H_{kv}=8$、头维度 $D=128$、半精度 $B=2$ 字节(FP16/BF16);
  • 每增加 1 个 Token:$2 imes 32 imes 8 imes 128 imes 2 = 131,072$ 字节 $pprox 128 ext{ KB}$;
  • 当生成到 4096 个字:$128 ext{ KB} imes 4096 pprox \mathbf{0.5 ext{ GB}}$ 单请求显存!如果是老一代 MHA(32个头),将暴增至 2 GB!如果并发 10 个用户就是 20 GB 显存!
缓存类型 存储结构特点 最适合的应用场景 一句话大白话比喻
DynamicCache 随生成动态拼接张量 单机本地调试、快速验证、短问答 随时往后贴便利贴,贴多少用多少
StaticCache 预先开辟最大长度张量 高并发生产服务、开启 torch.compile 预先包下整排连座,对号入座不挪窝
SlidingWindowCache 循环滚动的滑动窗口 长篇角色扮演、长程闲聊、无限流交互 金鱼的记忆,只看最近一屏内容
QuantizedCache 将 FP16 量化为 INT8/INT4 显存极其紧缺的边缘端、消费级显卡 把厚重档案袋抽成真空压缩包携带
PagedAttention 虚拟内存分页与页表映射 大型推理集群(vLLM / SGLang)生产部署 紧密拼积木,每一寸显存都塞满

提示:在 Hugging Face 最新主分支中,`past_key_values` 已经全面支持传入各种 `Cache` 实例对象以替换老式元组。

🗺️ 一遍过完:Hugging Face 官方进阶路线

所有资料已全部通过真实网络连接核验,从官方讲解到源码级 Cache 类层层剖析:

🚩 第一站:Hugging Face Cache 官方图解总纲
为什么去: 官方文档第一视角,搞清为什么训练不用缓存、推理必须用缓存。
预计耗时: 10 分钟
立即前往 ↗
🚩 第二站:文本生成全流程剖析 (Text Generation Guide)
为什么去: 掌握 GenerationMixin、贪心采样、温度调节与内部缓存传递全过程。
预计耗时: 15 分钟
立即前往 ↗
🚩 第三站:源码级 Cache 类家族详解 (Internal Generation Utils)
为什么去: 查看 DynamicCache、StaticCache、QuantizedCache 的 Python 源码与继承体系。
预计耗时: 20 分钟
立即前往 ↗