KV Cache 的追加与读取
示例从 3 个 prompt token 的 prefill 开始,展示后续 decode 中缓存的变化。
图示可左右滑动,保留矩阵比例与文字大小。
01
更新缓存
取一层中的一组 KV 头,K 和 V 分别沿位置轴追加当前 token 的一行。
02
读取 Key 缓存
当前 Query 与缓存中的全部 Key 匹配,得到各位置的注意力概率。
03
汇总 Value 缓存
按注意力概率对缓存中的 Value 加权求和,K/V 继续保留供后续步骤使用。
- 维度
- 取 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 个元素。图中展示位置的追加过程,实际实现可通过预分配等方式管理内存。