"""第 37 章：为已有 TinyTransformer 添加显式 KV 缓存；不训练新模型。""" from dataclasses import dataclass from math import sqrt import time import torch from tiny_transformer import TinyTransformer @dataclass class Cache: model: TinyTransformer layers: list length: int @property def nbytes(self): return sum(t.numel()*t.element_size() for pair in self.layers for t in pair) @torch.inference_mode() def cached_forward(model, ids, cache=None, use_sdpa=False): if model.training: raise ValueError("缓存接口只用于 eval 模式的推理") if ids.ndim != 2 or ids.shape[0] < 1 or ids.shape[1] < 1: raise ValueError("输入须为非空 [batch, length]") batch, length = ids.shape start = 0 if cache is None else cache.length config = model.config if cache is not None: if cache.model is not model or len(cache.layers) != config["layers"]: raise ValueError("缓存不属于当前模型") expected = (batch, config["heads"], start, config["width"]//config["heads"]) for key, value in cache.layers: if key.shape != expected or value.shape != expected: raise ValueError("缓存形状或长度不一致") if key.device != ids.device or value.device != ids.device: raise ValueError("缓存与输入设备不同") if key.dtype != model.tokens.weight.dtype or value.dtype != key.dtype: raise ValueError("缓存精度不同") if start+length > config["max_context"]: raise ValueError("超过位置表长度；本例不实现滑动窗口") position = torch.arange(start, start+length, device=ids.device) x = model.tokens(ids)+model.positions(position) new_layers = [] for index, block in enumerate(model.blocks): attention = block.attention normalized = block.norm1(x) projected = attention.qkv(normalized).chunk(3, dim=-1) query, key, value = [ part.reshape(batch, length, attention.heads, -1).transpose(1, 2) for part in projected] if cache is not None: old_key, old_value = cache.layers[index] key = torch.cat([old_key, key], dim=-2) value = torch.cat([old_value, value], dim=-2) key_position = torch.arange(start+length, device=ids.device) allowed = key_position[None, :] <= position[:, None] if use_sdpa: # 非方形注意力显式传绝对位置掩码，不直接套 is_causal=True。 mixed = torch.nn.functional.scaled_dot_product_attention( query, key, value, attn_mask=allowed, dropout_p=0.0, is_causal=False) else: scores = query@key.transpose(-2, -1)/sqrt(query.shape[-1]) weights = scores.masked_fill(~allowed, float("-inf")).softmax(-1) mixed = weights@value mixed = mixed.transpose(1, 2).contiguous().reshape(batch, length, config["width"]) x = x+attention.output(mixed) x = x+block.feedforward(block.norm2(x)) new_layers.append((key, value)) logits = model.output(model.final_norm(x)) return logits, Cache(model, new_layers, start+length) @torch.inference_mode() def generate(model, prompt, new_tokens, cached=True): if new_tokens < 1 or prompt.ndim != 2: raise ValueError("至少生成一个词元，提示为二维 ID") if prompt.shape[1]+new_tokens-1 > model.config["max_context"]: raise ValueError("生成所需输入超过位置表长度") model.eval() sequence = prompt.clone() cache, positions, pairs = None, 0, 0 stamps = [] started = time.perf_counter() for index in range(new_tokens): inputs = sequence if not cached or cache is None else sequence[:, -1:] if cached: logits, cache = cached_forward(model, inputs, cache) else: logits = model(inputs) # 每序列、每层的投影位置数；分数元素按每层、每个注意力头计数。 positions += inputs.shape[1] pairs += inputs.shape[1]*sequence.shape[1] token = logits[:, -1].argmax(-1, keepdim=True) sequence = torch.cat([sequence, token], dim=1) stamps.append(time.perf_counter()-started) return dict(ids=sequence, projected_positions=positions, score_elements=pairs, cache_bytes=cache.nbytes if cache else 0, first_seconds=stamps[0], total_seconds=stamps[-1], token_seconds=stamps) def corrected_distribution(target, draft): """枚举一次草稿采样、接受或残差重采样，验证分布；不是多词元加速内核。""" p, q = torch.tensor(target, dtype=torch.float64), torch.tensor(draft, dtype=torch.float64) if (p.shape != q.shape or p.ndim != 1 or torch.any(p < 0) or torch.any(q < 0) or not torch.isfinite(p).all() or not torch.isfinite(q).all() or not torch.isclose(p.sum(), torch.tensor(1., dtype=p.dtype)) or not torch.isclose(q.sum(), torch.tensor(1., dtype=q.dtype))): raise ValueError("需要两个等长的归一化非负概率向量") accepted_mass = torch.minimum(p, q) rejection = 1-accepted_mass.sum() residual = (p-q).clamp_min(0) correction = residual/residual.sum() if residual.sum() > 0 else torch.zeros_like(p) return accepted_mass+rejection*correction, float(accepted_mass.sum()) def main(): torch.set_num_threads(1) torch.manual_seed(2026) model = TinyTransformer(35, max_context=32, width=32, heads=4, layers=2).double().eval() ids = torch.randint(0, 35, (2, 11)) with torch.inference_mode(): whole = model(ids) _, cache = cached_forward(model, ids[:, :8]) chunk, _ = cached_forward(model, ids[:, 8:], cache) optimized, _ = cached_forward(model, ids[:, 8:], cache, use_sdpa=True) print("分块缓存与完整重算分数一致：", torch.allclose(chunk, whole[:, 8:], atol=1e-12)) print("显式掩码的 SDPA 与手写注意力一致：", torch.allclose(chunk, optimized, atol=1e-12)) prompt = ids[:1, :8] without, with_cache = [generate(model, prompt, 12, cached=flag) for flag in [False, True]] print("连续生成 12 个词元的 ID 一致：", torch.equal(without["ids"], with_cache["ids"])) print("每层投影位置数：", without["projected_positions"], "→", with_cache["projected_positions"]) print("每层、每个头的密集分数元素数：", without["score_elements"], "→", with_cache["score_elements"]) print("最终保留的 KV 张量：", with_cache["cache_bytes"], "字节") restored, accepted = corrected_distribution([.6, .3, .1], [.2, .5, .3]) print("推测采样接受概率：", f"{accepted:.3f}") print("接受加残差重采样后的分布：", [round(x, 3) for x in restored.tolist()]) print("本例随机权重用于核对计算；默认不报告随机波动的运行时间。") if __name__ == "__main__": main()