KV Cache 的追加与读取

示例从 3 个 prompt token 的 prefill 开始,展示后续 decode 中缓存的变化。

图示可左右滑动,保留矩阵比例与文字大小。

01

更新缓存

取一层中的一组 KV 头,K 和 V 分别沿位置轴追加当前 token 的一行。

K旧3 × 2追加K新1 × 2K4 × 2沿位置轴追加当前 K/VV旧3 × 2追加V新1 × 2V4 × 2沿位置轴追加当前 K/V
02

读取 Key 缓存

当前 Query 与缓存中的全部 Key 匹配,得到各位置的注意力概率。

q1 × 2×Kᵀ2 × 4缩放 → softmaxa1 × 4
03

汇总 Value 缓存

按注意力概率对缓存中的 Value 加权求和,K/V 继续保留供后续步骤使用。

a1 × 4×V4 × 2=H1 × 2各头输出拼接后乘 W_O得到当前 token 的注意力输出
维度
取 B = 1、d_h = 2,图中展示一组 KV 头。单层完整 K/V 的形状均为 B × n_kv × T × d_h。
对象
缓存包含应用 RoPE 后的 K,以及未经 RoPE 的 V。Q 仅用于当前步骤,色块深浅仅作数值示意。
容量
L 层缓存共包含 2 B L n_kv T d_h 个元素。图中展示位置的追加过程,实际实现可通过预分配等方式管理内存。