生成与缓存：让同一模型逐步续写

从给定真实前缀，走到自己产生前缀

本章使用第八章保存、第九章评价过的同一个模型。前面评分时，每一步都能读取真实过去；现在只有用户提供的开头，之后每一步只能读取自己刚选出的词元。目标是实现一条可停止、可复现、可核验的生成循环，再减少重复计算。先修为因果注意力、绝对位置、softmax 和保存恢复。缓存不改变模型参数，也不应该替换它给同一前缀的分数；选择策略则会改变输出分布，两件事须分开检验。

先把输入开头按训练时保存的词表编码，在前面放一个 BOS。把当前输入送进模型，最后一个位置的 logits 表示紧接这个前缀的下一词元分数。选择一个新编号，追加到序列，再计算下一步。遇到 EOS 即停止；若始终没有 EOS，就在预先设置的新词元数上限停止。这个上限是资源约束，不是模型预测的“自然结束”。输出记录应区分两种停止原因。

令前缀编号序列为 $s_{1:t}$，下一词元编号为 $v\in\{0,\ldots,V-1\}$。模型输出 $z(s_{1:t})\in\mathbb R^V$，默认概率为 $p(v\mid s_{1:t})=\operatorname{softmax}(z)_v$。循环的状态至少包括序列、缓存、已经生成的数量与采样随机状态。若输入开头含训练词表外的字，编码器只能用 UNK；界面能显示原始字并不意味着模型内部读取了那个字的独立表示。

选词规则本身也是一个明确的算法

贪心选择返回最大分数的编号。相同分数可能出现并列，需要固定规则，例如选择编号最小者；否则两台实现即使 logits 相同，也可能续写出不同的序列。贪心只最大化眼前这一步的条件概率，不能保证整段序列的联合概率最大，更不能保证事实正确。后续条件概率取决于已经选出的词元，所以逐步局部最大与整段搜索是不同问题。

温度 $\tau$ 是严格正的标量。将所有 logits 除以它，再做 softmax，定义采样分布 $q_\tau$。当 $\tau=1$，分布与模型原分布相同；当温度降低，已有分数差被放大。两种编号的概率比变成 $\exp((z_i-z_j)/\tau)$，这说明变化来自相对分数差，而不是给某个词元额外添加了事实证据。项目把 temperature=0 特别定义为贪心分支，不能真的执行除以零。

设原概率为二分之一、十分之三、五分之一。温度等于二时，未归一化权重相当于对这三个概率开平方，再除以平方根之和；分布会更平。温度等于二分之一时，权重相当于把概率平方：四分之一、零点零九、零点零四，相加零点三八，所以新概率约为零点六五七九、零点二三六八、零点一零五三。这个手算只改变采样，原检查点及其教师强制损失仍是原来的对象。

top-k 再保留分数最高的 $k$ 个编号，组成集合 $A$。在被保留集合上重新归一化，其他编号概率置零。若原分布为零点五、零点三、零点二，取前两个后得到 $0.5/0.8=0.625$ 和 $0.3/0.8=0.375$。忘记重新归一化就得不到和为一的离散分布。若正确答案恰好位于集合外，它的生成机会就变为零；截断可以控制随机性，不能作为真实性保证。

import math
def truncated_distribution(logits, temperature, k):
    assert temperature > 0 and 1 <= k <= len(logits)
    keep = sorted(range(len(logits)), key=lambda i:(-logits[i], i))[:k]
    shifted = [v/temperature for v in logits]
    maximum = max(shifted[i] for i in keep)
    weights = [math.exp(shifted[i]-maximum) if i in keep else 0.0 for i in range(len(logits))]
    total = sum(weights)
    return [v/total for v in weights]
logits = [math.log(.5), math.log(.3), math.log(.2)]
q = truncated_distribution(logits, 1.0, 2)
assert all(abs(a-b)<1e-12 for a,b in zip(q,[.625,.375,0.0]))
assert abs(sum(q)-1.0)<1e-12
print(q, truncated_distribution(logits,.5,3))

一条因果性质，怎样变成可以复用的中间结果

不使用缓存时，追加一个词元后，我们会重新计算整个窗口每层的 Q、K、V 和中间状态。但一个过去位置的计算不允许读取未来：它的输入词元、位置编号和合法前缀都没有变化。在评估模式、固定参数、固定同一前缀的前提下，过去位置在第一层的表示不变；第二层读取的过去表示也不变，逐层可以重复这个论证。这就是复用过去 K、V 的理由，不是凭经验猜测它们“大概差不多”。

对第 $l$ 层，缓存已经处理的键和值，形状分别为 $[B,h,t,d_h]$。新进入的长度为 $r$ 的片段先通过该层归一化与投影，得到新查询、键、值。把旧键与新键沿时间轴拼接，值也同样拼接。新查询与长度 $t+r$ 的全部键计算分数，并使用绝对位置比较构造因果遮罩。输出仍只计算新片段的 $r$ 个位置，再进入下一层。最后才把每层更新后的 K、V 返回给调用者。

位置编号是这条等价论证的前提。若旧缓存长度是五，新输入一个词元，应使用位置五，而不是又从零开始。若新输入两个词元，它们的位置是五和六：第一个新查询不能读第二个新键，第二个才能读两者。简单给一个局部二乘二下三角遮罩，无法同时正确表示五个旧键。项目用 past_length 加新片段索引构建查询位置，再与所有键位置比较。

还有两个必须保持的条件。第一，所有层缓存的时间长度相同，否则第几步的含义已经冲突；批量、头数、每头宽度和设备也须一致。第二，数值对照采用评估模式且关闭 dropout。训练时直接复用带不同随机掩码的中间值，既破坏前向等价，也改变梯度图。本课只把缓存作为生成优化，不把它暗中插进教师强制的训练循环。

第一次调用把整个提示送进模型，返回的缓存已经覆盖这些提示位置。选出新编号后，下次只把这个新编号送入，并携带旧缓存；不要再次传入整个提示，否则历史会被重复拼接。选择一个编号，与把它作为下一次输入真正处理，是前后两个时刻。若选到EOS便停止，则没有必要再计算以EOS为输入的隐藏状态。用每一步“输入长度、旧缓存长度、新缓存长度”的记录，可以独立发现重复喂入造成的长度错误。

项目返回stop_reason字段：仅在允许EOS停止且本轮选到EOS时为eos，其他由新词元上限结束的情况为budget。零预算也属于budget；若显式禁止遇EOS停止，即使最后一个编号恰好是EOS，也不能把原因记成eos。这让停止规则成为可测试的状态转移，而不是事后从一段解码文本猜测程序为什么结束。

import math
vectors = [[1.,0.],[0.,1.],[1.,1.],[-1.,.5]]
def attend(query, keys, values):
    score = [sum(a*b for a,b in zip(query,key))/math.sqrt(2) for key in keys]
    maximum = max(score)
    weights = [math.exp(s-maximum) for s in score]
    weights = [w/sum(weights) for w in weights]
    return [sum(w*value[d] for w,value in zip(weights,values)) for d in range(2)]
cached_keys, cached_values = [], []
for end, vector in enumerate(vectors, 1):
    cached_keys.append(vector[:]); cached_values.append(vector[:])
    cached = attend(vector,cached_keys,cached_values)
    reference = attend(vectors[end-1],vectors[:end],vectors[:end])
    assert all(abs(a-b)<1e-12 for a,b in zip(cached,reference))
print('single-head identity-projection cache check passed')

缓存花多少内存，省了哪一部分计算

定义每个标量存储 $b$ 字节，层数 $N$，批量 $B$，当前缓存长度 $t$，注意力总宽度 $d=h d_h$。每层一个K和一个V，共有 $2Bhtd_h$ 个标量；乘层数和字节数，得到 $2NBtdb$。它不包括模型权重、当前查询、临时分数、分配器余量或框架对象。把这个公式当作整个进程内存，会系统性低估需求。

用课程小配置，层数二、批量一、窗口十六、宽度三十二、float32每标量四字节，缓存为八千一百九十二字节，即八KiB。如果并行生成八个样本，就是六十四KiB。若改成三十二层、宽度四千零九十六、长度四千零九十六、批量一、每标量两字节，普通多头缓存为两GiB。这里KiB与GiB分别按二的十次方和三十次方换算，不能把单位不一致的数字直接相减。

计算量也要限定模型。以下只统计矩阵乘加次数，固定批量，讨论前缀从一增长到 $m$、始终未超过位置上限 $C$ 的情形；softmax、归一化、非线性与访存另有代价。令前馈宽度为 $f$，对长度为 $t$ 的完整前向，投影和前馈的量级为 $Nt(d^2+df)$，注意力两个矩阵乘的量级为 $Nt^2d$。只有固定 $f$ 与 $d$ 的比例，例如本项目常用四倍宽度，前一项才可简写为 $Ntd^2$。不缓存且每步从头计算，需要累加所有前缀长度：注意力项的总和为 $O(Nm^3d)$。缓存后每步只增加一个查询，但仍需读取全部过去键和值，所以该步注意力是 $O(Ntd)$，总和为 $O(Nm^2d)$。

这不是“缓存后注意力变成常数时间”。上下文越长，每一步仍需处理越多K、V；硬件上还可能被读取带宽限制。首次处理提示的预填充依然要完成整段计算，不能把预填充耗时混入或移出统计后，就声称某种缓存策略全面更快。项目的数值测试回答输出是否等价；实际首词元延迟、逐词元间隔和总耗时属于需要另行测量的性能问题。

窗口满了：最容易破坏无损优化的边界

项目使用有限的学习式绝对位置表。当缓存已满，再追加词元会超出可用位置。不能直接截掉最老的K、V，再把剩余向左搬一个位置，假装它们仍有效：位置嵌入改变了，过去隐藏状态也可能包含被删词元的信息。即使把键的某个位置标签改掉，也没有撤销前面层已经发生的混合。这个错误常常仍会生成文字，因此必须用与完整重算的分数对照发现。

本项目采用明确的有限窗口协议：窗口内可以增量追加；需要淘汰最老词元时，清空缓存，把最后C个编号从位置零开始完整重算，然后继续。无缓存参考也采用同样取窗规则。两者比较的才是同一个算法目标。它无法保留窗口外的记忆，也不应宣称通过缓存实现了无限上下文。若想采用不同位置体系或循环记忆，需要另定义模型与训练方法，不能只改缓存容器。

这也限制了前面复杂度结论的应用：本项目一旦满窗，后续每一步都须重算长度 $C$ 的窗口，注意力约为 $O(NC^2d)$，线性与前馈约为 $O(NC(d^2+df))$。因此不能把未满窗时的缓存收益延伸到任意长生成。若把 $C$ 固定、只让生成步数增加，满窗后的总成本对步数近似线性，但每步带着完整窗口重算的常数；它与“每步只算一个新查询”是两个不同区间。性能报告应把提示长度、满窗时刻和重算次数同时列出。

生成中还要分清概率相近与编号相同。浮点矩阵乘的计算分组可能造成很小的logits误差；如果两个最大值几乎并列，这个误差就可能改变贪心结果。项目先比较每一步同一前缀的logits最大误差，再比较固定条件下的生成序列；不能只因最终文本偶然相同就认为所有中间值等价。数值容差应预先指定，并保留所用精度、设备和模式。

采样复现还有一个额外状态：随机数生成器。仅记录整数种子，而中间又多抽一次数，后续序列仍会不同。项目为生成使用本地generator，不让别的代码偷偷消耗同一随机流；保存或重放采样时，还需保留具体策略与状态。固定种子不修复错误的位置offset，固定输出字符串也不证明分布正确，二者应接受不同的检查。

停止判断使用特殊词元编号，而不是在已经解码的普通文本里搜索字母EOS。输入的字面字符串、保留编号和可读显示是三个层次。项目的生成接口一次处理一个提示；模型本身虽支持批量张量，同时生成多条不同长度的请求还需要逐行维护存活状态、停止原因与缓存位置，不能把某一行结束广播成全部结束。并发服务的排队和连续批处理留给高效推理课程，不混进这里的单请求正确性结论。

本章交付：先建立参考，再启用优化

提交时保留一个每步全窗口重算的参考分支，一个增量缓存分支；两者读取同一检查点、词表、输入与选择规则。依次检查空提示的BOS输入、一个新词元、多个新词元、窗口边界、超长提示和EOS停止。记录cache张量形状、重算次数与分数差，先通过这些检查，再谈速度。下一章会分析减少KV头数和避免显式存储大注意力矩阵的机制；它们不能代替本章的位置与因果契约。

# 沿用同一starter.py；第3/7章方法已经完成，现增加以下分支。
# causal_attention: query_positions = arange(Q) + query_offset
# TransformerLM.forward: 验证past各层形状/长度，positions从L开始
# 保留原无缓存路径与有效目标损失。

def causal_attention(query, key, value, query_offset=0, dropout=None):
    """[B,H,Q,D] × [B,H,K,D]; row i can see keys j <= offset+i.

    This intentionally uses matmul/softmax rather than a fused attention API.
    Probabilities are returned before any dropout so rows sum to one.
    """
    scores = query @ key.transpose(-2, -1) / math.sqrt(query.shape[-1])
    query_positions = torch.arange(query.shape[-2], device=query.device) + query_offset
    key_positions = torch.arange(key.shape[-2], device=key.device)
    future = key_positions.unsqueeze(0) > query_positions.unsqueeze(1)
    probabilities = torch.softmax(scores.masked_fill(future, float("-inf")), dim=-1)
    mixed_probabilities = probabilities if dropout is None else dropout(probabilities)
    return mixed_probabilities @ value, probabilities

# TransformerLM.forward（替换第7章的无缓存方法）
def forward(self, tokens, targets=None, past_key_values=None, use_cache=False):
        if tokens.ndim != 2 or tokens.dtype != torch.long or tokens.shape[1] == 0:
            raise ValueError("tokens must be a nonempty [batch,time] int64 tensor")
        batch, length = tokens.shape
        offset = 0
        if past_key_values is not None:
            if len(past_key_values) != self.config.n_layers:
                raise ValueError("past_key_values must contain one (K,V) pair per layer")
            lengths = []
            for key, value in past_key_values:
                expected = (batch, self.config.n_heads, key.shape[-2], self.config.d_model // self.config.n_heads)
                if tuple(key.shape) != expected or tuple(value.shape) != expected:
                    raise ValueError("invalid cached K/V shape")
                if key.device != tokens.device or value.device != tokens.device:
                    raise ValueError("cached K/V and tokens must be on the same device")
                lengths.append(key.shape[-2])
            if len(set(lengths)) != 1:
                raise ValueError("all layers must have the same cache length")
            offset = lengths[0]
        if offset + length > self.config.block_size:
            raise ValueError(
                f"cached length {offset} + input length {length} exceeds block_size "
                f"{self.config.block_size}; recompute a sliding window from position zero"
            )
        positions = torch.arange(offset, offset + length, device=tokens.device)
        hidden = self.embedding_dropout(
            self.token_embedding(tokens) + self.position_embedding(positions)
        )
        presents = []
        for index, block in enumerate(self.blocks):
            past = None if past_key_values is None else past_key_values[index]
            hidden, present = block(hidden, past, use_cache)
            if use_cache:
                presents.append(present)
        logits = self.head(self.final_norm(hidden))
        loss = None
        if targets is not None:
            if targets.shape != tokens.shape or targets.dtype != torch.long:
                raise ValueError("targets must match tokens shape and int64 dtype")
            valid_tokens = (targets != -100).sum()
            if valid_tokens == 0:
                raise ValueError("targets must contain at least one scored token")
            loss = F.cross_entropy(
                logits.reshape(-1, self.config.vocab_size), targets.reshape(-1),
                ignore_index=-100, reduction="sum",
            ) / valid_tokens
        return {"logits": logits, "loss": loss,
                "past_key_values": tuple(presents) if use_cache else None}