KV 缓存

主线位置：推理主线的第四步「带 KV 缓存生成」。prefill 一次送进整段 prompt，之后每步只送一个新 token，前面位置的键和值从缓存里取。本页对照 transformers 的 DynamicCache 和 candle llama.rs 的 Cache：谁管位置、缓存放在哪、带缓存和不带缓存的 logits 差多少，以及几种用错的样子。

版本：PyTorch 2.15.0 / transformers 5.22.0（CPU）、candle 0.11.0（CPU，accelerate；Metal）

sequenceDiagram
  participant U as 调用方
  participant M as Llama
  participant C as Cache
  U->>M: forward 5 个 token, index_pos 0
  M->>C: 每层存 k v, 长 5
  M-->>U: 最后一个位置的 logits
  U->>M: forward 1 个 token, index_pos 5
  C-->>M: 取出旧的 k v
  M->>C: cat 之后存回, 长 6
  M-->>U: logits

实验代码：Python 侧 pt_kv.py、pt_tf.py；candle 侧 crate llm-exp 的 bin tf、kvpit。release 构建，平台 M2 Max、macOS 27、rustc 1.99.0。模型是本地的 SmolLM2-135M-Instruct 权重，f32；PyTorch 固定 torch.set_num_threads(8)，注意力用 eager（另注明的除外）。缓存在 candle 内部怎么存、candle-nn 里另外几种缓存，见源码卷的 KV 缓存与注意力。

两边的缓存对象

transformers 5.22.0 里，缓存是一个独立的对象，model(...) 把它作为 past_key_values 接进去、更新、再放进输出里交回来。手写一个带缓存的贪心循环：

def greedy(prompt, cache, n=20):
  out = model(prompt, past_key_values=cache, use_cache=True)
  nxt = out.logits[0, -1].argmax().view(1, 1)
  toks = [nxt.item()]
  for _ in range(n - 1):
    out = model(nxt, past_key_values=out.past_key_values, use_cache=True)
    nxt = out.logits[0, -1].argmax().view(1, 1)
    toks.append(nxt.item())
  return toks

with torch.inference_mode():
  cache = DynamicCache(config=model.config)
  first = greedy(ids, cache)
  print("cache type", type(cache).__name__, "layers", len(cache.layers), "seq_length", cache.get_seq_length())
  print("layer 0 keys", tuple(cache.layers[0].keys.shape), cache.layers[0].keys.dtype)
  print("first:", repr(tok.decode(first)))

cache type DynamicCache layers 30 seq_length 24
layer 0 keys (1, 3, 24, 64) torch.float32
first: ' Paris. Paris is the largest city in France and the capital of the French department of the Espace'

30 层，每层的键是 (batch, kv_heads, seq, head_dim)。长度是 24 而不是 25：最后生成的那个 token 还没送进去。注意 3 个 KV 头：缓存里存的是 GQA 复制之前的键和值（对齐 logits），只有查询头数的三分之一。

candle 的 Cache 同样是模型之外的一个对象，每次 forward 以 &mut 传进去：

let mut cache = Cache::new(true, DType::F32, &cfg, &dev)?;
let x = Tensor::new(ids, &dev)?.unsqueeze(0)?;
let mut next = model.forward(&x, 0, &mut cache)?.squeeze(0)?.argmax(0)?.to_scalar::<u32>()?;
let mut out = vec![next];
let mut pos = ids.len();
for _ in 1..n {
  let x = Tensor::new(&[next], &dev)?.unsqueeze(0)?;
  next = model.forward(&x, pos, &mut cache)?.squeeze(0)?.argmax(0)?.to_scalar::<u32>()?;
  out.push(next);
  pos += 1;
}

Cache 里放的东西比 DynamicCache 多：

pub struct Cache {
    masks: HashMap<(usize, usize), Tensor>,
    pub use_kv_cache: bool,
    kvs: Vec<Option<(Tensor, Tensor)>>,
    cos: Tensor,
    sin: Tensor,
    device: Device,
}

• kvs：每层一对 (k, v)，形状和 transformers 一样是 (1, 3, seq, 64)。每步 CausalSelfAttention::forward 用 Tensor::cat 在序列维上接一格再存回去。
• cos / sin：RoPE 表。transformers 把它放在模型的 rotary_emb 里，candle 放在缓存里，所以 Cache::new 要 dtype 参数，而且必须和模型权重的 dtype 一致（半精度）。
• masks：按 (seq_len, kv_len) 缓存的因果 mask。
• use_kv_cache 是唯一的公开字段。Cache::new(false, ..) 时每次 forward 都不读也不存 kvs，调用方每步要送整段前缀。

除了 use_kv_cache，字段都是私有的，也没有读取它们的方法。想像 cache.layers[0].keys.shape 那样看一眼缓存：

error[E0616]: field `kvs` of struct `candle_transformers::models::llama::Cache` is private
  --> src/bin/kvpriv.rs:10:24
   |
10 |   println!("{}", cache.kvs.len());
   |                        ^^^ private field

Cache 也没有 get_seq_length、reset、crop 这类方法。缓存有多长，只有调用方记着；要重新开始，就再 Cache::new 一个。

位置：transformers 自己推，candle 要你传

上面两段循环最大的差别是 candle 的 forward 多了第二个参数 index_pos：这一段输入的第一个 token 在整个序列里的位置。transformers 的 Python 循环里没有位置，因为 LlamaModel.forward 从缓存的长度推出来（transformers 5.22.0 的代码）：

if position_ids is None:
  past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
  position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
  position_ids = position_ids.unsqueeze(0)

candle 不核对 index_pos 和缓存的长度是否一致：index_pos 只用来从 RoPE 表里取行、决定 mask 的形状，键和值总是接在缓存后面。照 Python 的习惯不管位置、每步都传 0（kvpit pos0）：

pos0: [7042, 28, 260, 3575, 282, 4649, 28, 260, 1797, 1918, 28, 260, 1797, 1918, 28, 284, 260, 3575, 282, 4649]

解码出来是 Paris, the capital of France, the United States, the United States, and the capital of France，正确的是 Paris. Paris is the largest city in France and the capital of the French department of the Espace。没有报错：decode 一步只有一个 token，不建 mask（对齐 logits），每个新 token 都按位置 0 旋转，和缓存里的键之间的相对位置全错了。prefill 之后的第一个 token Paris 还是对的，因为它来自 prefill 的输出，位置是从 0 正确递增的。quantized_llama 上同一类错误（计数器按调用次数前进）的现象和修法见源码卷的坑 cannot broadcast [13, 45] to [1, 9, 13, 13]。

带缓存的 logits 等于不带缓存的吗

缓存只是把算过的键和值存下来，数学上带不带缓存结果相同。沿参考序列（prompt 加参考的 19 个生成 token，共 24 个）逐步取 logits，每一步都喂参考序列里的 token，而不是自己生成的，这样一处分叉不会传下去（teacher forcing）。transformers 一侧，不带缓存是一次前向 24 个 token，取后 20 个位置；带缓存是先 prefill 5 个，再一个一个送：

if mode == "full":
  rows = model(ids).logits[0, 4:]
else:
  cache = DynamicCache(config=model.config)
  out = model(ids[:, :5], past_key_values=cache, use_cache=True)
  rows = [out.logits[0, -1]]
  for p in range(5, len(seq)):
    out = model(ids[:, p:p + 1], past_key_values=out.past_key_values, use_cache=True)
    rows.append(out.logits[0, -1])

candle 的 forward 只返回最后一个位置，不带缓存就每一步新建一个 Cache::new(false, ..)，从头算整段前缀：

match mode {
  "nocache" => {
    for end in 5..=seq.len() {
      let mut cache = Cache::new(false, dt, &cfg, &dev)?;
      rows.push(row(&seq[..end], 0, &mut cache)?);
    }
  }
  _ => {
    let mut cache = Cache::new(true, dt, &cfg, &dev)?;
    rows.push(row(&seq[..5], 0, &mut cache)?);
    for pos in 5..seq.len() {
      rows.push(row(&seq[pos..pos + 1], pos, &mut cache)?);
    }
  }
}

20 步里每一步 49152 个 logits，20 步合起来最大 |ref| 32.73：

都在 1e-4 以内，但没有一行是 0，包括 transformers 自己和自己比。带缓存时，decode 一步的查询只有一行，q @ k.t()、lm_head 都是 1 行的矩阵乘；不带缓存时同样的数是一个多行矩阵乘里的一行，BLAS 的分块和求和顺序不同（对齐 logits 的动脑筋里，transformers 只换 lm_head 的行数就差 2.48e-5）。candle 的第 0 步是 0，因为两边的第 0 步是同一个 5 token 的 prefill；transformers 的不带缓存是一次 24 个 token，第 0 步就已经是不同形状的矩阵乘。

candle 带缓存对 transformers 的误差（8.20e-5）和 transformers 自己带不带缓存的差（8.01e-5）一样大：缓存这一步没有引入额外的偏差。

分段 prefill

prompt 不一定一次送完：多轮对话里，先 prefill 系统提示和历史，再把新一轮用户输入作为一段多 token 输入接上去。这时 seq_len > 1 且 index_pos > 0，mask 是长方形的：前 index_pos 列全能看，后 seq_len 列是因果三角。

build_causal_mask(2, 3):
[[0, 0, 0, 0, 1],
 [0, 0, 0, 0, 0]]
Tensor[[2, 5], u8]

tf chunk 先送 3 个 token（index_pos 0），再送 2 个（index_pos 3），之后一个一个 decode，和一次 prefill 5 个比：

row(&seq[..3], 0, &mut cache)?;
rows.push(row(&seq[3..5], 3, &mut cache)?);

tf_c_cpu_f32_chunk.json        steps=20 max|diff| all=2.909e-05 first=2.909e-05 last=1.335e-05 top1 20/20 (max|ref| 32.73)
tf_c_metal_f32_chunk.json      steps=20 max|diff| all=2.861e-05 first=2.861e-05 last=1.478e-05 top1 20/20 (max|ref| 32.73)

CPU 差 2.91e-5，Metal 差 2.86e-5，20 步的 argmax 全同。transformers 一侧同样先送 3 个、再接着同一个 DynamicCache 送 2 个，位置由缓存长度推出，最后一个位置的 logits 和一次送 5 个差 4.67e-5：

cache = DynamicCache(config=model.config)
model(ids[:, :3], past_key_values=cache, use_cache=True)
two = model(ids[:, 3:], past_key_values=cache, use_cache=True).logits[0, -1]

chunk 3+2 vs 5: max|diff|=4.673e-05

缓存用错的样子

换一个 prompt，没换缓存

transformers 里把一个已经用过的 DynamicCache 接着传给新 prompt，不报错，新 prompt 接在旧序列后面，位置从 24 开始：

again = greedy(second, cache)  # 不清缓存，接着用
print("reuse: seq_length", cache.get_seq_length(), repr(tok.decode(again)))
fresh = greedy(second, DynamicCache(config=model.config))
print("fresh:", repr(tok.decode(fresh)))

reuse: seq_length 48 ' Berlin. Berlin is the largest city in Germany and the capital of the German state of the EThe'
fresh: ' Berlin. Berlin is the largest city in Germany and the capital of the Federal Republic of Germany. Berlin'

（second 是 The capital of Germany is。）模型同时看得见关于法国的那一段，前 14 个 token 和新缓存相同，从第 15 个起不同（ German state 对 Federal Republic）。

candle 里同样不新建 Cache、把新 prompt 从 index_pos 0 送进去（kvpit reuse）：

Error: cannot broadcast [5, 5] to [1, 9, 5, 29]

index_pos 是 0，mask 按 5 × 5 建；缓存里已有 24 个位置，cat 之后键有 29 个，分数是 (1, 9, 5, 29)。和 quantized_llama 不同，llama.rs 在 index_pos == 0 时不丢掉旧缓存（quantized_llama 的行为见源码卷 KV 缓存与注意力），它也没有 clear_kv_cache。修法只有一个：每段新对话 Cache::new 一个新的。

位置超过 max_position_embeddings

candle 的 RoPE 表在 Cache::new 里按 max_position_embeddings（SmolLM2 是 8192）建好，位置超过就取不出来（kvpit far，在 index_pos 8192 送一个 token）：

max_position_embeddings=8192
Error: narrow invalid args start + len > dim_len: [8192, 32], dim: 0, start: 8192, len:1

transformers 的 LlamaRotaryEmbedding.forward 每次按 position_ids 现算 cos / sin，没有上限：

far = model(ids[:, :1], position_ids=torch.tensor([[8192]])).logits
print("position 8192:", tuple(far.shape), "finite", bool(far.isfinite().all()))

position 8192: (1, 1, 49152) finite True

能算不等于算得对，超过训练长度的位置模型没见过。candle 用报错把这条线画了出来，transformers 不画。

一段前缀，两条分支

Cache 实现了 Clone，kvs 里是 Tensor（克隆只复制 Arc）。llama.rs 每步都 cat 出新张量再存回，从不原地写，所以克隆出来的两份缓存之后各走各的。共同前缀 The capital of 只 prefill 一次，克隆后分别接 France is 和 Germany is（kvpit clone）：

let x = Tensor::new(&PROMPT[..3], &dev)?.unsqueeze(0)?;
model.forward(&x, 0, &mut cache)?;
let mut other = cache.clone();
let a = generate(&model, &mut cache, &dev, &PROMPT[3..], 3, 20, false)?;
let b = generate(&model, &mut other, &dev, &second[3..], 3, 20, false)?;

branch France == fresh France: true
clone: [11452, 30, 11452, 314, 260, 3995, 2240, 281, 4521, 284, 260, 3575, 282, 260, 7030, 5261, 282, 4521, 30, 11452]

France 那一支和从头生成的 20 个 token 完全相同；Germany 那一支解码出来是 Berlin. Berlin is the largest city in Germany and the capital of the Federal Republic of Germany. Berlin，和 transformers 用新缓存生成的一样。注意第二段输入是两个 token，index_pos 3，走的正是上一节的长方形 mask。candle-nn 的 KvCache 克隆后共享缓冲区，不能这样用（KV 缓存与注意力）。

动脑筋答案

不报错，生成的是接着法国那一段往下说的内容（kvpit reuse1）：

reuse1: [3575, 282, 4649, 30, 657, 314, 260, 3995, 2240, 281, 1910, 411, 2155, 284, 260, 1796, 768, 35488, 6151, 282]

解码是 capital of France. It is the largest city in Europe by population and the second most populous department of。同一个 token 用新 Cache 跑（kvpit fresh1）是：

fresh1: [198, 198, 504, 1646, 288, 253, 3424, 2329, 1740, 314, 288, 1464, 253, 2437, 284, 4798, 1380, 1786, 338, 27784]

解码是 \n\nThe key to a successful online course is to create a clear and engaging learning experience that resonates。

• 一个 token 时 seq_len == 1，不建 mask，也就没有形状要对，cat 之后这个 The 能看见缓存里关于法国的 24 个位置。
• 它按位置 0 旋转，而缓存里的键在位置 0 到 23：相对位置是负的，模型没见过这种排列，但上下文仍然在，于是续写出了 capital of France。
• 多 token 输入会报错，单 token 输入静默串到上一段对话里。所以「新对话新建 Cache」不能靠报错来提醒。

小结

• transformers 的 DynamicCache 和 candle 的 Cache 都在模型外面，每层存 (1, 3, seq, 64) 的键和值（GQA 复制之前）。candle 的 Cache 还装着 RoPE 表和 mask，字段私有，没有长度、清空、裁剪的方法。
• transformers 从缓存长度推位置；candle 的 index_pos 由调用方传、不核对。decode 时一直传 0 不报错，生成的内容变了。
• 带缓存和不带缓存的 logits 不逐位相同：transformers 自己差 8.01e-5，candle CPU 6.25e-5、Metal 6.87e-5，最大 |ref| 32.73，argmax 20/20；分段 prefill（3 + 2）和一次 prefill 差 2.9e-5。
• 换 prompt 不换缓存：transformers 静默接上；candle 多 token 报 cannot broadcast [5, 5] to [1, 9, 5, 29]，单 token 静默串题。llama.rs 在 index_pos == 0 时不丢缓存，新对话要新建 Cache。
• 位置到 8192（max_position_embeddings）candle 报 narrow 错误，transformers 照算。Cache 可以克隆，做共同前缀的分支。