KV cache：记住已经算过的历史

先写出张量形状，再讨论缓存

B 是请求数，T 是本次输入长度，D 是 hidden size；Hq / Hkv 是 query / KV 头数，d 是每头维度。

布局可以换轴；存储的元素数不变。这里假设 K 与 V 维度相同。

一行 attention，真的算出一个数

已缩放的 logits 为 [0, ln 2, ln 3]；exp 后是 [1, 2, 3]，除以总和 6。V 为 [2, 4, 8]。

所有数值为手算教学示例。

为什么缓存 K、V，却通常不缓存 Q？

后续位置需要历史 K 来匹配、历史 V 来取内容。历史 Q 只服务历史位置的输出，不会被未来查询读取。

这里讨论因果解码器。双向 encoder 的旧表示可能因新输入而变化。

最小增量实现：追加后再 attention

q, k_new, v_new = project(x_new)
q, k_new = rope(q, k_new, positions)
k = concat(k_cache, k_new, axis="time")
v = concat(v_cache, v_new, axis="time")
scores = q @ k.transpose(-1, -2) / sqrt(d)
scores = where(allowed, scores, -inf)
y = softmax(scores, axis=-1) @ v
return y, (k, v)

追加当前 token 后再读，才能关注自己。每一层有独立 cache；上一层的输出才是下一层输入。

concat 用于理解；服务端通常预分配或分页写入，避免每步复制整个历史。

缓存长度、绝对位置、物理槽位要分开

已有 5 个 token，窗口物理容量为 4。新 token 的 RoPE 位置继续增长，写入槽位按容量取模。

所有数值为手算教学示例。

一次追加两个 token，mask 不是 2×2

cache 已有 P = 3 个位置；本次 chunk 长 T = 2。Q 的长度是 2，K 的长度是 5。

若直接用左上对齐的 2×5 下三角，Q3 会只看到 K0，结果不会报错，却已经错了。

不同 API 的非方形 causal mask 对齐规则不同；先用显式布尔 mask 与参考结果比对。

GQA 的头映射是一条索引关系

例：Hq = 32、Hkv = 8。每 4 个 query heads 共用一个 KV head，仍各自计算自己的 attention 权重。

参考代码可以 repeat_interleave 验证；高效 kernel 应复用原始 KV，不把它永久复制 4 遍。

每卡 KV 不是总量除以卡数就结束

假想模型：32 层、8192 token、8 个 KV 头、d = 128、BF16。完整逻辑 KV 为 1 GiB。

第二行假设每个 KV head 被 2 张卡持有。真实实现可用其他布局；从每个 rank 的实际缓存形状核算。

请求长短不同，要把各自长度加起来

L = 32、Hkv = 8、d = 128、b = 2 bytes：每个 token 的所有层 KV 是 131072 bytes = 128 KiB。

若全部填充到 8192，会占 3 GiB。分页按块取整，还要加尾块、元数据与预留空间。

把字节数变成延迟下界

沿用 1 GiB / 序列的 KV。假设有效带宽为 1.5 TB/s，单步把所有层 KV 各读一次。

TB 用十进制，GiB 用二进制；这里的带宽是教学假设，不是某张卡的测量值。

为什么单 query 的算术强度很低？

一个 head 读 T 个 key / value，每个向量 d 维。一次乘加按 2 FLOP 计，不计 softmax。

BF16 的 b = 2，所以约 1 FLOP/byte。把历史加长，算术和读取一起增多，强度不会自动提升。

GQA 若多个 Q 共享一次 KV 读取，强度可提高；重复读取、cache 命中和融合方式也会改变实际结果。

用长度扫描定位瓶颈

假想剖析：固定 batch = 1、输出 64 token，预热后测稳定 decode 步骤。

attention 随历史增长，32K 已占约一半。优先检查 KV 流量、GQA 复用和压缩；2K 场景的最大收益则被固定部分限制。

对照时保持输出策略、batch、精度与频率一致；同时记录 TTFT、TPOT 和显存。

先证明增量输出等于完整前向

若首个 token 对、第二个开始错：先查追加顺序与 position offset。若只有 chunk decode 错：先查非方形 causal mask。

推理用 eval 模式。不要只比较最终采样文本：采样会放大微小差异，也会掩盖一些数值错误。

推理缓存与训练激活解决不同问题

训练通常一次并行处理整段序列，并保留反向所需信息；推理则重复追加 token，复用历史前向结果。

固定权重的评估、流式模型的截断反传有专门定义；先确定需要的梯度语义。

48 GiB 还能同时接几条请求？

先写出假设和单位，再代入数字。

三行 query，应该看到哪些列？

先写出假设和单位，再代入数字。

压缩 KV 一半，端到端能快多少？

先写出假设和单位，再代入数字。

亲手验证：完整前向 = 增量前向

用 NumPy 在 CPU 上跑一个可检查的参考实现。小尺寸便于把中间张量打印出来。

python3 -m pip install numpy
python3 lab.py
python3 lab.py --seed 7

full vs token: PASS
full vs chunks: PASS
wrong mask detected: PASS
128 KiB/token; 8192 tokens = 1.000 GiB

把 chunk 长度改成 [1, 3, 2, 5]，再把 offset 故意写成 0；观察第一处不一致。增加 Hq，保持 Hkv 不变，确认缓存元素数不变。

这是正确性与成本的教学实验；耗时不代表 GPU kernel 性能。