KV 缓存与注意力

主线位置：推理主线「注意力 + KV 缓存」那一步。生成第 t 个 token 时，注意力要用到前面所有位置的键（K）和值（V）；它们在前面的步骤里已经算过，存下来就不用每步重算整段前缀。本页讲这些张量存在哪、怎么追加，以及 candle-nn 里不经过 matmul 的 CPU 注意力实现。模型怎么在 forward 里调用它们，见量化 llama 的一次前向。

stateDiagram-v2
  Empty --> Allocated: 第一次 append 建全零缓冲
  Allocated --> Allocated: slice_set 写入新位置
  Allocated --> Grown: 放不下时 cat 一块全零
  Grown --> Allocated: 继续 slice_set
  Allocated --> Empty: reset
  Grown --> Empty: reset

图里是 KvCache 的状态。本页的实验是 nn-exp 的 kv、kvprobe 和 flash（都是 release），平台 M2 Max、macOS 27、rustc 1.99.0。

quantized_llama：自己拼

推理主线上的模型没有用 candle-nn 的缓存。quantized_llama 每层存一个 kv_cache: Option<(Tensor, Tensor)>，forward_attn 里：

        let (k, v) = match &self.kv_cache {
            None => (k, v),
            Some((k_cache, v_cache)) => {
                if index_pos == 0 {
                    (k, v)
                } else {
                    let k = Tensor::cat(&[k_cache, &k], 2)?;
                    let v = Tensor::cat(&[v_cache, &v], 2)?;
                    (k, v)
                }
            }
        };
        self.kv_cache = Some((k.clone(), v.clone()));

形状是 (batch, kv_heads, seq, head_dim)，在第 2 维（序列）上拼。每一步 cat 都分配一个长一格的新张量，把旧的整段复制过去，旧张量随即释放：第 t 步复制 t 行，生成 n 个 token 一共复制约 n²/2 行。index_pos == 0 时直接丢掉旧缓存，所以从头开始一段新对话不调 clear_kv_cache 也行；clear_kv_cache 把每层设成 None，用来提前释放内存。

KvCache：预分配，再 slice_set

KvCache 是 K、V 两个 Cache。Cache::append：

    pub fn append(&mut self, src: &Tensor) -> Result<()> {
        let seq_len = src.dim(self.dim)?;
        // ...
        if self.all_data.is_none() {
            let mut shape = src.dims().to_vec();
            shape[self.dim] = self.max_seq_len;
            let ad = Tensor::zeros(shape, src.dtype(), src.device())?;
            self.all_data = Some(ad)
        };
        let ad = self.all_data.as_mut().unwrap();
        while self.current_seq_len + seq_len > self.max_seq_len {
            let mut shape = src.dims().to_vec();
            shape[self.dim] = self.grow_by;
            let next_ad = Tensor::zeros(shape, src.dtype(), src.device())?;
            *ad = Tensor::cat(&[&*ad, &next_ad], self.dim)?;
            self.max_seq_len += self.grow_by;
        }
        ad.slice_set(src, self.dim, self.current_seq_len)?;
        self.current_seq_len += seq_len;
        Ok(())
    }

• 缓冲区在第一次 append 时才建：那时才知道批大小、头数、dtype 和设备。形状照抄 src，只把序列维换成 max_seq_len，填零。
• 写入用 Tensor::slice_set：它取 &self，却经 storage_mut() 拿到存储的写锁，用 copy2d 把 src 原地复制到偏移 current_seq_len 处。每步只复制新来的那几行，和缓存已有多长无关。它要求两边都连续、不共享存储、dtype、设备、维数都一致。
• KvCache::append 返回 current_data()，即 all_data.narrow(dim, 0, current_seq_len)：一个和缓冲区共享存储的视图，序列维没填满时不连续。
• 放不下时按 grow_by 增长。grow_by 在 new 里设成初始的 max_seq_len，之后不变，所以是线性增长（4、8、12、16……），每次增长用 cat 把整个缓冲区复制一遍。

slice_set 要求连续，KvCache::append 不替你调 contiguous()。注意力里最常见的写法是 reshape((b, seq, heads, d))?.transpose(1, 2)?，转置后不连续，直接 append 会报：

non-contiguous append: Err(slice-set only supports contiguous tensors)
seq 1 transposed: contiguous true, append Ok(())

第二行是序列长度为 1 的同样写法：(1, 1, 3, 64) 转置成 (1, 3, 1, 64) 后，长度为 1 的维度不参与连续性判断，仍算连续，append 成功。所以 decode 一步一个 token 时不会出错，只有多 token 的 prefill 才报；只用单 token 测过的代码，换成一段 prompt 才撞上这个错误。

Cache 的源码注释说 all_data 用 Option 的另一个理由是「makes it safe to clone a KvCache that has been reset」。反过来说，没有 reset 的 KvCache 克隆之后，两份共享同一个缓冲区（Tensor 克隆的是 Arc），各自的 current_seq_len 却是分开的。实验（kv clone，序列维长 8，每次追加一个值）：

c2 after its append: [1.0, 2.0]
c1 after its append: [1.0, 3.0]
c2 read again:       [1.0, 3.0]
ConcatKvCache: d1 [1.0, 3.0], d2 [1.0, 2.0]

c2 写在位置 1 的 2，被 c1 后来写在同一位置的 3 覆盖了。想从同一段前缀分出两条生成（例如束搜索），克隆 KvCache 是错的；ConcatKvCache 每次 cat 都产生新张量，克隆后互不影响。

ConcatKvCache：每步 cat

ConcatKvCache 就是把 quantized_llama 的做法包了一层：

    pub fn append(&mut self, k: &Tensor, v: &Tensor) -> Result<(Tensor, Tensor)> {
        // Detach inputs to break BackpropOp chain - KV caches are inference-only.
        let k = k.contiguous()?.detach();
        let v = v.contiguous()?.detach();

        self.k = Some(match &self.k {
            None => k,
            Some(k_cache) => Tensor::cat(&[k_cache, &k], self.dim)?.detach(),
        });
        // ...

它替你调 contiguous()，还 detach()，返回的总是连续张量。文档注释写着它「providing significant GPU performance improvements for autoregressive generation」，推荐 GPU 用它、CPU 用 KvCache；引入它的 PR #3143 标题是「ConcatKvCache for 2-5x GPU speedup on autoregressive generation」，提交信息里没有测量数据。

在 CPU 和 Metal 上量一下。模拟 SmolLM2 的 30 层、3 个 KV 头、64 维，f32，先 prefill 128 个位置，再 decode 200 步；每种缓存先整轮跑一遍预热，第二轮计时。「append」只追加，「append + attn」每层再做一次单 token 注意力，注意力有三种写法：GQA 时把 KV 头复制成 9 个再 matmul（像 repeat_kv），不复制、只用 3 个头 matmul，以及 Metal 的 ops::sdpa。单位是每个 decode 步（30 层）的毫秒数：

• 追加本身：KvCache 每步只复制一行，CPU 上 30 层 10 μs；cat 每步复制整段（这里 128 到 328 行），慢 24 到 35 倍。Metal 上两者都被内核调用的固定开销拉平到零点几毫秒，cat 仍慢一倍。
• 加上注意力，KvCache 返回的不连续视图并没有拖慢后面的 matmul 和 sdpa，三种写法下它都是最快的。单独测一个 (1, 3, 329, 64) 的视图（kvprobe），q @ k.t() 在视图上和在连续张量上都是 6 到 13 μs，缓冲区开到 4096 也一样。
• 这台机器上没有复现 PR 标题说的「2 到 5 倍」；它说的 GPU 也许是 CUDA，本卷没有实测 CUDA。ConcatKvCache 实际的好处是：不用预先定上限、克隆安全、返回的张量总是连续（下游如果自己要求连续，就省一次复制）。
• 「复制头」那一列 Metal 比 CPU 还慢：expand 之后的 reshape 要复制，单 token 的小张量上每多一次内核调用都很显眼。quantized_llama 在 Metal 上单 token 时不复制头，直接调 sdpa，源码注释是「SDPA will do MQA for us」。

早先没有预热轮的版本里，Metal 上进程里第一个测的缓存出现过 55 ms/步（不复制头）和 23 ms/步（sdpa）的「append + attn」，同一进程后面的缓存都在 1 ms 左右；加了预热轮之后，预热轮里也见过一次 12 ms/步。这类首次运行的抖动属于 Metal 后端的话题，表里只用预热后的第二轮。

其他几种缓存

kv_cache.rs 一共 1142 行，除了上面两种还有：

KvCache 的用户有 Quantized Phi-3、GLM-4、SmolLM3（量化版不开 use_flash_attn 时也是）、Qwen3-VL，以及 Gemma 3 / 4 的非滑动窗口层；ConcatKvCache 的用户是 Qwen3、Quantized Qwen3（GPU）、Quantized Qwen3-MoE。Llama（llama.rs）和 quantized_llama 都自己 cat。

RawInterleavedKvCache::write_kv 满了时 self.buf.resize(self.buf.len() * 2, 0.0)，容量为 0 时翻倍还是 0。RawInterleavedKvCache::new(3, 64, 0) 之后写一个位置：

thread 'main' (3697912) panicked at ~/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/candle-nn-0.11.0/src/kv_cache.rs:1090:21:
range end index 64 out of range for slice of length 0

模型里传的都是 4096 之类的正数，碰不到；但这是一个只靠调用方避开的 panic。

attention/：CPU 上的 flash attention

candle-nn/src/attention/ 是 0.11.0 新加的一块（PR #3254，首次进入 0.11.0，之后 #3513、#3584、#3634、#3643 又改过），不经过 matmul 和 softmax_last_dim，直接在 CPU 存储上算注意力。入口 flash_attn：输入 (B, S, H, D)，输出 (B, H, S, D)，掩码是 AttnMask（无、带偏移的因果、显式张量）。

• B = 1：因果掩码走 causal（按行只循环到对角线，被掩住的位置不算），其余走 standard；两者再按「单 token decode / 多 token prefill」和「有没有掩码、ALiBi、softcap」分出精简版。
• B > 1：只支持 f32 / f16、无 softcap、无掩码或因果，打包成变长格式（varlen）算；不满足就报错，不回退。
• 只认 CPU：取存储时不是 Storage::Cpu 就报「Expected CPU storage for q」。

核心是在线 softmax（online_softmax_step）：

pub(crate) fn online_softmax_step(
    score: f32,
    m: &mut f32,
    ssum: &mut f32,
    acc: &mut [f32],
    v_apply: impl FnOnce(&mut [f32], f32),
) {
    if score > *m {
        let scale_old = (*m - score).exp();
        for a in acc.iter_mut() {
            *a *= scale_old;
        }
        *ssum = *ssum * scale_old + 1.0;
        *m = score;
        v_apply(acc, 1.0);
    } else {
        let w = (score - *m).exp();
        v_apply(acc, w);
        *ssum += w;
    }
}

每个查询行顺着键走一遍：算出一个分数，就用它更新当前最大值 m、指数和 ssum 和输出累加器 acc（D 个 f32）；最大值变大时把之前累加的结果按 exp(旧 m − 新 m) 缩小。走完一遍除以 ssum 就是输出。S × KV 的分数矩阵从来不出现，内存只和 D 有关。点积用 f32 累加，BF16 / F16 输入也是（dot_f32 的注释专门解释了为什么不用 VecOps::vec_dot：它会把累加结果先舍入回半精度）。

谁在用：qwen3.rs 在 CPU 上（且没开 flash-attn feature 时）走 forward_cpu_flash_attn，quantized_qwen3.rs（CPU 上）、smol/quantized_smollm3.rs（use_flash_attn 时）也用它，decode 时调的是读交错缓存的 causal_decode_f32_interleaved。Llama、quantized_llama、BERT 都不用。旧入口 cpu_flash_attention::run_flash_attn_cpu 从 0.9.2 起标成 deprecated，转发到新的 flash_attn。

和「matmul + 因果掩码 + softmax_last_dim + matmul」比一下（flash，CPU，f32，9 个查询头、3 个 KV 头、64 维，各 20 次取平均）：

q len  512, kv len  512: flash_attn 12.494 ms, matmul+softmax 12.741 ms, out [1, 9, 512, 64], max diff 1.0728836e-6
q len    1, kv len 1024: flash_attn 0.288 ms, matmul+softmax 0.433 ms, out [1, 9, 1, 64], max diff 2.5331974e-7
q len    1, kv len 4096: flash_attn 0.582 ms, matmul+softmax 1.981 ms, out [1, 9, 1, 64], max diff 1.7508864e-7

prefill 时两者差不多（参照实现里还包括了在主机上建 512 × 512 掩码的时间）；decode 时 flash 版本快 1.5 到 3.4 倍，上下文越长差得越多：参照实现要先把 3 个 KV 头复制成 9 个、再写出整行分数，flash 版本直接按 h / k_h 找到对应的 KV 头读，只走一遍。

动脑筋答案

  let mut c = KvCache::new(2, 4);
  for (step, n) in [3usize, 1, 1, 10].into_iter().enumerate() {
    let k = Tensor::ones((1, 3, n, 64), DType::F32, &dev)?;
    let (kk, _) = c.append(&k, &k)?;
    let all = c.k_cache().all_data().as_ref().unwrap();
    // print kk.dims(), kk.is_contiguous(), all.dims(), max_seq_len
  }

step 0: append 3 -> returned k [1, 3, 3, 64] contiguous false, buffer [1, 3, 4, 64], max_seq_len 4
step 1: append 1 -> returned k [1, 3, 4, 64] contiguous true, buffer [1, 3, 4, 64], max_seq_len 4
step 2: append 1 -> returned k [1, 3, 5, 64] contiguous false, buffer [1, 3, 8, 64], max_seq_len 8
step 3: append 10 -> returned k [1, 3, 15, 64] contiguous false, buffer [1, 3, 16, 64], max_seq_len 16

• 第一次建出 4 格，放 3 个，返回的视图是 4 格里的前 3 格，不连续（头之间的步长按 4 格算）。
• 第二次正好填满，narrow 出整个缓冲区，连续。
• 第三次放不下，长一块 grow_by = 4，变成 8 格。
• 第四次要放到第 15 格：while 循环长了两次（8 → 12 → 16）。在 #3035 之前的 if 版本里，这一步只会长到 12，然后 slice_set 报错。
• 增长量一直是 4，不会翻倍。max_seq_len 给小了，长对话里每长一块就把整个缓冲区复制一次；给大了，第一次 append 就按上限分配全零内存。

小结

• quantized_llama 和 llama 都不用 candle-nn 的缓存，每层自己 Tensor::cat，每步复制整段；index_pos == 0 时丢掉旧缓存。
• KvCache 第一次 append 时按上限分配全零缓冲，之后用 slice_set 原地写入，每步只复制新行；放不下时按初始大小线性增长；输入必须连续；克隆后共享缓冲区。
• ConcatKvCache 每步 cat、替你 contiguous 和 detach，克隆安全。实测 M2 Max 上它在 CPU 和 Metal 都比 KvCache 慢，没有复现「GPU 快 2 到 5 倍」。
• 另有环形（滑动窗口）、散列（批内不同位置）、交错（给 CPU flash attention）几种缓存；RawInterleavedKvCache 容量为 0 时写入会 panic。
• attention/ 的 CPU flash attention 用在线 softmax 一遍算完，不建分数矩阵；decode 时比 matmul 写法快 1.5 到 3.4 倍；只有 Qwen3 和 SmolLM3 在用。