利用开源Ollama快速熟悉LLM设计方法（8. KV-cache）

KV-Cache这玩意儿，本质上就是一种典型的空间换时间策略。说白了，就是把注意力机制中的 K 和 V 缓存下来，等到后续生成新 token 时直接拿来用，省得每次都要从头再算一遍。

看上面那张图就很好理解了：每次计算产生的 K 和 V 都先存着，当新序列进来时，只需要从缓存池子里把之前的 K 和 V 读出来就行。这样一来，那些已经被计算过的部分就没必要重复劳动了。

不过嘛，原理听着简单，可一摸到 llma.cpp 的源代码，事情就变得有意思起来了。先来看它的核心数据结构：

从代码注释能看出来，这是一个循环存储的 buffer，也就是俗称的 ring buffer。里面的几个关键字段干了什么，可以这么理解：

has_shift——标记格子里的位置是否被整体挪动过，偏移量就存在 cell[i].delta 里。

do_defrag——一个内存整理标记，表示是否需要执行碎片整理。

do_copy——也是控制标记，告诉系统要不要做一次复制操作。

recurrent——这个有点特殊，它主要是为循环状态模型准备的，比如 Mamba 模型就需要它。

head——表示当前需要参与计算的 KV 头位置。在上面的示意图里，那些橙色的格子就是它。实际上，它的数值等于已经存了数据的单元数量。

size——缓存的容量上限，也就是能容纳多少键值对。

used——已经被占用的单元格数量。只要某个格子里至少有一个序列 ID（也常叫 slot ID），就算是被用了。

n——真正的有效细胞数量，标准是非负位置且有关联的序列 ID。

type_k、type_v——分别存储键和值的具体数据类型。

cells——用来记录缓存中每个格子状态的容器。

k_l、v_l——分别对应每一层的键和值。

ctxs、bufs——这两个是搞内存管理的，ctxs 管上下文，bufs 是 buffer 数组。

还有一个叫 llama_kv_cell 的结构体，主要职责是记录 token 位置的变动：

pos 就是 token 的原始位置，delta 记录位置偏移量，src 用在循环状态模型里做状态复制，seq_id 则标明这个 token 归属于哪个序列 ID（可以是 task ID 或 slot ID）。

所有的 KV-cache 操作都在 cells 里完成。比如 rm，就是在指定区域里把某个序列 ID 移除；add 则是往区域里塞入新的序列 ID；cp 负责在指定区域插入序列 ID；clear 就是清空，把所有单元格的 pos 重置为 -1，同时擦掉所有序列 ID。

那 KV-cache 是怎么初始化的呢？llama_kv_cache_init 大致走了这么几步：

第一步，设好最基本的参数，像 head、size、used、type_k、type_v 这些。

第二步，清理 cells 数组并重新调整大小到 kv_size。

第三步，算一下每种 buffer 要覆盖多少层 decode。buffer 的类型挺多：GPU、SYCL（一种异构计算框架）、VULKAN（图像计算框架）、HBM（高内存带宽）以及普通的 CPU。不过 ggml 定义的 backend 其实就三种——CPU、GPU、GPU_SPLIT。

第四步，给每种 buffer 分配一个 ggml context，也就是 ggml 的内存管理对象。

第五步，为 decode 的每一层创建一个初始为空的 KV 张量。

第六步，最后给每种 buffer 实际分配一块 buffer。