生成文字和 KV 缓存

第 01 模块第 3 课用 DeepSeek 做过温度实验，那时模型是一个黑箱。现在模型是我们自己训练的，可以打开看看它每一步到底在做什么。

python generate.py

模型给出的只是概率

模型读完一段文字，在最后一个位置给词表里的 6289 个字各打一个分，softmax 之后变成概率：

== 1. 读完一段文字，模型给下一个字的概率（前 8 名）
  「白日依山尽，黄河入海」→ 无 6.9%  间 4.5%  多 4.3%  深 3.7%  边 3.0%  空 2.6%  难 2.3%  遥 2.1%
  「床前明月光，疑是地上」→ 天 3.2%  花 2.9%  楼 2.4%  清 1.8%  看 1.7%  游 1.7%  人 1.6%  枝 1.3%
  「春风」→ 。 14.2%  ， 5.1%  吹 3.3%  满 2.9%  雨 1.9%  起 1.9%  尽 1.4%  落 1.3%

"黄河入海"后面，原诗是"流"，可"流"不在前 8 名里。"疑是地上"后面，"霜"也不在。上一课检查过，这个模型不背诵，它没有记住这两首名诗，只是按学到的规律给出"合理"的字。

顺便说一个有意思的发现：《全唐诗》里收录的《静夜思》是"床前看月光，疑是地上霜。举头望山月，低头思故乡"，和课本上的"床前明月光……举头望明月"不一样。一般认为，我们熟悉的版本出自明清时期的选本。所以就算模型背了诗，它背的也是"看月光"。

"春风"后面，概率最高的是句号（14.2%）。我数了一下训练数据："春风"一共出现 797 次，其中 111 次后面紧跟着句号，占 13.9%，和模型给出的概率几乎一样。模型学到的，就是这些字在唐诗里的出现规律。

生成文字，就是反复做同一件事：拿到概率，选一个字，接到后面，再算下一个字的概率。怎么"选"，决定了写出来的是什么。

温度

第 01 模块讲过温度：把分数除以温度，再 softmax。温度小于 1，高分和低分的差距被拉大，模型更"保守"；大于 1，差距被缩小，模型更"大胆"。

logits = logits[:, -1, :] / temperature
...
next_id = torch.multinomial(F.softmax(logits, dim=-1), num_samples=1)

同样的随机种子，三个温度各写 3 首：

  温度 0.3：
    白云无限意，白发不如何。何事无人在，无人不可知。
    一望东山路，千行万里愁。夜深人不见，秋色月无声。
    江南春水上，江上月明时。独有清风起，何人别有期。
  温度 1.0：
    白露起春光，知君得舞酲。霜时飏先没，桐死势悠扬。
    白发摩牛放，真关万趾兴。夜寒休绕郡，秋杀杜娘宫。玉镜收新镜，金尊坐晚麕。何言吊鱼道，双锡空余年。
    戍楼映江岸，淼危棹度斜。欲分凝炯娩，归盖乱精冰。
  温度 1.5：
    淡露依春似映兹，襞舞凋年雪路揉筝衬情。死势悠扬双觉处，摩牛放将溉噭胞。……

温度 0.3 写得最通顺，"一望东山路，千行万里愁。夜深人不见，秋色月无声"几乎可以乱真。代价是用字很常见，重复也多："无人"出现了两次，"江"出现了两次，读多了会觉得都差不多。

温度 1.0 就是按模型给出的原始概率选字。用字丰富了，但开始出现"麕""酲""炯娩"这些生僻字，意思也更难接上。

温度 1.5 彻底乱了：格式崩了（"襞舞凋年雪路揉筝衬情"一句九个字），满篇生僻字。原因是温度高时，概率分布被压平，6000 多个字里那些本来概率极低的字，加起来也分走了可观的概率。一旦选中一个奇怪的字，后面的上下文就变得奇怪，错误越积越多。

top-k：只在靠前的字里选

温度高时的问题是"长尾"：几千个低概率的字加起来占了太多份额。top-k 的办法很直接：只保留概率最高的 k 个字，其余的全部去掉，在这 k 个里按概率选。

if top_k is not None:
    kth = torch.topk(logits, top_k).values[:, -1:]
    logits = logits.masked_fill(logits < kth, float("-inf"))

温度同样是 1.5，加上 top-k=20：

== 3. 温度 1.5，但只从概率最高的 20 个字里选（top-k=20）
    白头白马在西陵，自是西林下路尘。今日不辞青史士，可堪回首白苹州。
    万国风尘里，千金雨色深。不知归未到，应似到人中。
    东北山川尽不愁，三千春尽一枝愁。春寒未入千门色，雨尽空闻九夜寒。……

格式恢复了，生僻字也不见了，但比温度 0.3 更有变化。高温度负责"大胆"，top-k 负责"别太离谱"，两个一起用效果最好。

另一个常用的办法叫 top-p（核采样）：不固定个数，而是从最高的开始往下加，加到累计概率达到 p（比如 0.9）为止。概率集中时只剩几个字，分散时多留一些。第 01 模块用 API 时见过这个参数。

续写和藏头诗

生成的开头不一定是空的。给出前半句，模型就接着写：

== 4. 给出开头，让它续写
    床前明月光，已有四邻情。知君得舞袖，何奈在山期。
    大漠孤烟直，高楼楚水清。雁雕摩露下，鸿劒鼓歌中。夜雨初分郡，秋风又过巴。此才应不遇，空与钓鱼舟。
    人工智能好，言蠹不见人。徒思一岁月，独在岘生春。

"大漠孤烟直"续成了一首完整的五言律诗，八句。"人工智能"这个唐朝没有的词，它也照样接得下去，因为这四个字各自都在唐诗里出现过。

更进一步，可以在生成过程中插手。藏头诗要求每句的第一个字是指定的：先把指定的字放进去，再让模型一个字一个字地写完这一句，写的时候不许它输出标点和换行，最后由我们加上标点：

@torch.no_grad()
def acrostic(heads, n_char=7, temperature=0.8):
    ids = encode("\n")
    for i, head in enumerate(heads):
        ids = torch.cat([ids, encode(head)], dim=1)  # 先把指定的字放进去
        for _ in range(n_char - 1):  # 再让模型一个字一个字地写完这一句
            logits = model(ids)[0][0, -1] / temperature
            logits[BANNED] = float("-inf")
            next_id = torch.multinomial(F.softmax(logits, dim=-1), 1)
            ids = torch.cat([ids, next_id[None]], dim=1)
        ids = torch.cat([ids, encode("，。"[i % 2])], dim=1)  # 标点由我们加，保证格式
    return decode(ids[0].tolist()).strip()

== 5. 藏头诗：每句开头的字由我们指定，其余由模型写
    春色欲依春似春，眠花先得舞花年。不知诗句先酬情，觉后还随楚水人。
    学得重雕出，无情忘却回。止行空向兴，境静更闻闻。

每句的开头连起来是"春眠不觉"和"学无止境"。

这个技巧有一个正式的名字：约束解码。在选字之前，把不符合要求的选项的分数设成负无穷。第 02 模块第 4 课让 DeepSeek 输出 JSON 时，服务端保证输出格式正确，用的就是同一个原理：在每一步，把会让 JSON 不合法的词元都排除掉。

重复的计算

现在看生成的效率。回到 generate 方法，最简单的写法是每一步都把整段文字送进模型：

logits, _ = self(idx)  # 不用缓存：每一步都把整段序列从头算一遍

写第 50 个字时，模型要处理前面全部 49 个字；写第 51 个字时，又要把这 49 个字重新处理一遍，再加上第 50 个字。

但仔细想想，前 49 个字的计算结果并没有变。因果掩码保证每个位置只看前面，后面加了新字，前面位置的计算不受任何影响。真正新的，只有最后一个位置。

最后一个位置要算注意力，需要什么？它自己的 q，以及前面所有位置的 k 和 v。所以只要把每一层的 k 和 v 存起来，每一步只算新字的 q、k、v，把新的 k、v 接到存好的后面，就够了。这就是 KV 缓存。

实现

gpt.py 的注意力里，缓存只多了几行：

if cache is not None:  # KV 缓存：把之前算过的 K、V 接在前面，不用重算
    if "k" in cache:
        k = torch.cat([cache["k"], k], dim=2)
        v = torch.cat([cache["v"], v], dim=2)
    cache["k"], cache["v"] = k, v

有一个细节要注意：因果掩码。不用缓存时，q 和 k 的长度一样，掩码是一个正方形的下三角。用缓存时，q 只有新来的 1 个字，k 却有全部的字，而且新字排在最后，它可以看到所有的 k。所以掩码要错开：

total = k.size(2)
mask = torch.ones(T, total, dtype=torch.bool, device=x.device).tril(diagonal=total - T)

位置嵌入也要跟着调整：新字不是第 0 个位置，而是第 start 个。这就是 forward 里 start 参数的作用。

生成时，第一步把整段提示词送进去，把所有位置的 k、v 存进缓存；之后每一步只送最新的一个字：

if use_cache:
    # 第一步把整段提示词送进去；之后每步只送最新的一个词元，前面的 K、V 从缓存里取
    logits, _ = self(idx[:, start:], caches=caches, start=start)
    start = idx.size(1)

量一量

== 6. KV 缓存：结果一样吗？快了多少？
  同样的随机种子，写 120 个字：用缓存和不用缓存的结果完全一样
  写  32 个字：不用缓存     23 毫秒，用缓存    12 毫秒，快了 1.8 倍
  写  64 个字：不用缓存     52 毫秒，用缓存    26 毫秒，快了 2.0 倍
  写 120 个字：不用缓存    130 毫秒，用缓存    49 毫秒，快了 2.6 倍
  这个模型每个词元的缓存：2 × 4 层 × 128 维 × 4 字节 = 4 KB，128 个词元共 512 KB

先确认正确性：同样的随机种子，用不用缓存，写出的 120 个字完全一样。缓存只是省掉了重复计算，结果不变。

再看速度。写的字越多，省下的越多：32 个字快 1.8 倍，120 个字快 2.6 倍。不用缓存时，写第 n 个字要处理 n 个字，总的计算量随长度按平方增长；用缓存后，每一步只处理 1 个字。

在我们这个模型上，加速只有两倍多，并不惊人。原因是模型太小、文字太短：每一步真正的计算只要零点几毫秒，Python 本身调用函数、拼接张量的开销占了很大比例，这部分开销缓存省不掉。在大模型上，每一步的计算量大得多，文字动辄几千上万个词元，KV 缓存省下的就是数量级的差距。

缓存的代价：内存

KV 缓存用内存换时间。每个词元，每一层都要存一个 k 和一个 v。我们的模型每个词元只要 4 KB，128 个词元 512 KB，微不足道。

大模型就不一样了。一个几十层、每层几千维的模型，每个词元的缓存可能要几百 KB 到 1 MB 以上。一个十几万词元的长对话，缓存就要占到几十 GB 的显存，常常比模型本身还大。同时服务很多用户时，这就是最主要的显存开销。

所以上一课末尾提到的那些改进，很多都是冲着 KV 缓存来的：多个头共用 k 和 v（分组查询注意力），或者把 k 和 v 压缩成更小的向量再存（DeepSeek 的多头潜在注意力，MLA）。第 10 模块讲 vLLM 时，还会看到服务端怎么管理这些缓存。

第 06 模块第 4 课讲的"缓存命中的输入更便宜"，是另一层的缓存：服务商把常用前缀（比如固定的系统提示词）的 KV 缓存保存下来，下一个请求用到相同的前缀时，直接取出来，不用重新计算，所以能便宜很多。

练习

1. 实现 top-p 采样：把概率从大到小排序，保留累计概率达到 p 的那些字。用温度 1.5、p=0.9 写几首，和 top-k=20 比较。
2. 写一个"固定格式"的约束解码：只允许生成七言绝句，第 8、16、24 个字符只能是逗号或句号，其他位置不许出现标点。
3. 把 gpt.py 的 block_size 改成 512 并重新训练（或者只是随机初始化），测一下写 500 个字时，用缓存和不用缓存各要多久。

自测

1. 温度和 top-k 分别改变了什么？为什么常常一起用？

温度在 softmax 之前把分数整体缩放，改变概率分布的"尖锐"程度：温度低时集中在少数高分的字上，温度高时更平均。top-k 只保留概率最高的 k 个字，把其余的去掉。温度高能带来多样性，但会让大量低概率的字分走份额，top-k 把这些长尾去掉，两者一起用既有变化又不至于太离谱。

2. KV 缓存为什么不会改变生成的结果？

因为有因果掩码，每个位置的计算只依赖它自己和前面的位置，后面加上新字不会改变前面位置的 k 和 v。缓存只是把这些不会变的结果存起来重复使用，省掉了重复计算，数学上完全等价。

3. KV 缓存的代价是什么？为什么在大模型上这个代价很重要？

代价是内存：每个词元在每一层都要存一份 k 和 v。大模型层数多、维度大，每个词元的缓存很大，长上下文和多用户同时使用时，缓存占的显存可能超过模型本身，所以很多模型结构上的改进（如分组查询注意力、多头潜在注意力）都是为了减小它。