llama.cpp KV Cache 与注意力机制：环形缓冲到 Flash Attention

1. 背景与目标

自回归 transformer 推理时，每生成一个新 token 都要重新计算整个序列的注意力，但经过 softmax 之后的键（K）和值（V）矩阵只依赖于已解码的位置，不必反复重算。这就是 KV cache 的基本动机：将每一层的 K、V 保存下来，后续步骤只需追加最新的 KV 向量，大大削减计算量。

然而，朴素的 KV cache 存在两个突出矛盾：内存爆炸（长上下文或大批次时 cache 可能超出显存）和高效管理（多个并发序列如何复用、淘汰，如何在不复制大量数据的情况下实现位置偏移）。llama.cpp 在纯 C/C++ 框架内给出了一套完整的 KV cache 管理方案，并结合 hadamard 旋转、flash attention 等技巧将注意力计算的精度与速度推到实用水平。

本文从内存布局、分配算法、KQ mask 生成、逻辑移位，到 flash attention 的 CUDA 实现，深入剖析其底层机制。

2. 核心设计与实现

2.1 KV 缓存的 3D 张量与单元元数据

llama.cpp 的 KV 缓存并非一个简单的二维数组，而是按层、流（stream）组织的三维结构。每个 transformer 层存储两个张量：键 k 与值 v，形状为 [n_embd_gqa, kv_size, n_stream]（参见 src/llama-kv-cache.cpp:80-250）。其中 n_embd_gqa 是分组查询注意力中的头维度，kv_size 是缓存槽位数量，n_stream 则用于支持并行序列解码。

为什么三维？ 单独一个序列占用一个流，当 kv_unified=false 时，每个序列拥有独立的 KV 缓冲区，避免了不同序列的键值互相覆盖，大幅简化隔离逻辑。开关统一模式后所有序列共享一个二维缓冲区以减少内存开销。

除了原始数据，缓存还维护了一套单元元数据（llama_kv_cells），每个槽位对应一个 cell，记录：位置 pos、逻辑移位 delta、归属序列 bitset、是否为空等信息。这样数据本身与“谁在用哪个槽”解耦，后续的分配、重用、分片查找都基于 cell 元数据快速进行，无需扫描整个 key/value 数据区。

2.2 槽位分配算法：基于窗口的 MRU 扫描

当需要为新 token 分配缓存空间时，核心函数 llama_kv_cache::find_slot（src/llama-kv-cache.cpp:500-650）执行一次智能扫描。其目标是在有限的 kv_size 内找到可复用的 cell，同时尽量保证连续分配（contiguous）以利于矩阵运算。

算法流程简化如下（非连续模式）：

// src/llama-kv-cache.cpp: find_slot 内环
for (uint32_t i = 0; i < n_test; i++) {
    const auto idx = head_cur;
    head_cur++; n_tested++;
    bool can_use = cells.is_empty(idx);
    if (!can_use && cells.seq_count(idx) == 1) {
        const llama_pos pos_cell = cells.pos_get(idx);
        if (!can_use) {
            const llama_seq_id seq_id_cell = cells.seq_get(idx);
            if (llama_hparams::is_masked_swa(n_swa, swa_type, 
                    pos_cell, cells.seq_pos_max(seq_id_cell) + 1)) {
                can_use = true;
            }
        }
    }
    if (can_use) { res.idxs[s].push_back(idx); } else if (cont) { break; }
}

关键点：扫描起点是上次分配的 head_cur（最近用过的槽之后），这样能跳过大概率还在活跃的区域（warm start）。如果空闲空间足够多，甚至会直接重置到 0 号槽。cell 可用的两个条件：

• 空 cell：从未被占用。
• 单序列占用且被滑动窗口注意力（SWA）遮蔽：当模型启用 SWA 时，检查该 cell 的位置是否已经超出序列当前窗口的最大位置。若被遮蔽，意味着注意力计算中这个位置的权重已经被强制置零，因此 cell 可被安全覆盖，达到类似 LRU 的效果，但无需维护全局时间戳。

这种 SWA-aware slot reuse 使得在开启滑动窗口后，固定大小的 KV 缓存能够近似支持无限上下文（超出窗口的历史自动被淘汰），而无需物理驱逐和移动数据。

2.3 逻辑位置移位：避免物理内存搬动

序列长度超过缓存大小时，需要“丢弃”最前面的 token。传统做法是物理地将整个 tensor 数据向前移动，代价极大。llama.cpp 引入了逻辑移位（K-shift）机制。

在 seq_add（src/llama-kv-cache.cpp）中，当上下文溢出时并不立即移动数据，而是给对应的 cell 记录一个 shift 增量。当后续需要执行注意力时，再构建一个计算图 build_graph_shift（src/llama-kv-cache.cpp:1100-1200）对所有受影响的键张量原地进行 RoPE 旋转调整。具体步骤（量化键的情形）：

// src/llama-kv-cache.cpp: build_rope_shift 量化分支
if (ggml_is_quantized(cur->type)) {
    tmp = ggml_cast(ctx, cur, GGML_TYPE_F32);
    tmp = ggml_mul_mat_aux(ctx, tmp, rot);  // 撤销 Hadamard
    tmp = ggml_rope_ext(ctx, tmp, shift, factors, n_rot, ...);
    tmp = ggml_mul_mat_aux(ctx, tmp, rot);  // 重新施加 Hadamard
    tmp = ggml_cpy(ctx, tmp, cur);
}

对于量化键（比如 Q8_0），不能直接应用 RoPE（旋转位置编码），因为量化表示与旋转操作不兼容。代码中先反量化到 FP32，然后通过 Hadamard 矩阵（rot）解回原始空间，施用 RoPE 旋转，再重新 Hadamard 变换、重新量化写回。Hadamard 矩阵（满足 H^T H = I）用于将量化噪声去相关，这个技术在量化 KV cache 注意力中尤为关键，详见后文优化部分。

整个过程仅在调用 get_has_shift() 为 true 时触发，且为批量图执行，无需每次解码都运行，将 O(N) 的内存搬动降为 O(1) 的 delta 记录加上偶尔的 O(layers * len) 旋转计算。

2.4 KQ Mask 的高效生成与复用

计算注意力前需要生成一个 [n_tokens, n_kv] 的 mask 矩阵，指示哪些 key 对 query 有效（因果遮蔽、SWA、二维位置等）。直接对每个 token 遍历所有 KV cell 开销巨大。llama.cpp 在 set_input_kq_mask_impl 模板中进行了精细优化（src/llama-kv-cache.cpp:750-950）：

• 模板特化：通过 causal, swa, is_2d, alibi 等布尔模板参数，编译器能在编译期消除分支，生成高效的内部循环。
• Mask 复用：对于同一序列的多个 token，如果它们的遮蔽条件大部分相同（除了临近位置可能会因为新 token 加入而变化），则复用上一次计算的 mask，只重新计算可能被新 token 影响的小范围区域，大幅减少计算量。
• 数据类型分发：最终写入的目标 buffer 可能是 fp16 或 fp32，不同数据类型的写入采用不同函数模板，避免转换开销。

复用逻辑利用了每个序列在缓存中的 seq_srct（上一次的位置范围），通过差分更新将渐进复杂度从 O(n_tokens * n_kv) 降低到接近 O(n_kv)。

2.5 多流批量与统一缓存

n_stream 维度允许在一个 layer 的 KV tensor 中放置多个序列的 cache，每个流独立分配 slot。当请求统一缓存时（kv_unified=true），所有序列共享同一个二维平面，但 cell 元数据仍然记录各自的所有权，支持多个序列同时使用同一个 cell（通过 bitset 记录）。这种设计用位掩码管理多序列 scene，避免了为每个序列维护完整数据块的冗余。

跨流复制则通过 seq_cp 实现（如 speculative decoding 中将 draft 模型的序列复制到 target 模型流中）。数据移动在 backend 上执行，性能取决于后端拷贝能力。

3. 优化技术细节

3.1 Hadamard 旋转：量化 KV 下的精确注意力

当 KV 缓存被量化存储时（例如 Q8_0、Q4_0），直接计算 Q·K^T 会引入显著的量化误差，因为点积对噪声敏感。llama.cpp 借鉴了 QuaRot 的思想，但在 KV cache 层次应用了 Walsh-Hadamard 旋转。

实现位于 src/llama-kv-cache.cpp 中的 ggml_gen_hadamard：

// src/llama-kv-cache.cpp: ggml_gen_hadamard 生成
data[0*n + 0] = 1.0 / sqrtf(n);
for (int s = 1; s < n; s *= 2) {
    for (int i = 0; i < s; i++) {
        for (int j = 0; j < s; j++) {
            const float val = data[i*n + j];
            data[(i + s)*n + (j    )] =  val;
            data[(i    )*n + (j + s)] =  val;
            data[(i + s)*n + (j + s)] = -val;
        }
    }
}

Hadamard 矩阵是正交矩阵，即 H^T H = I。将 K 事先用 H 旋转，存储 K' = K·H，则注意力计算变为 Q·(K')^T = Q·H^T·K^T = (Q·H^T)·K^T。如果同时对 Q 也施加相同的旋转（通常放在投影层之后），则效果等价于未经旋转的标准计算。关键在于，量化操作施加在旋转后的 K' 上，其量化噪声因为 Hadamard 的变换而分散到所有维度，不再集中在某些大值维度上，使得最终点积对量化误差不敏感。

在 ggml 中，这个旋转是通过 GGML_HINT_SRC0_IS_HADAMARD 提示给 matmul 内核的，内核可以利用该信息做进一步的优化。

3.2 SWA 感知的槽位重用：无限上下文的无额外代价

前面已经描述了 SWA-aware reuse 的机理，这里再强调其作为隐式 LRU 的巧妙之处：它不需要维护任何访问时间戳或引用计数，只需借力于模型本身的注意力掩码。一个 cell 是否可用，仅取决于它对应的位置是否还落在最新序列 token 的滑动窗口内。窗口之外的 cell 自动变为可被覆盖，这就是以注意力窗口性质作为驱逐策略。

这种设计使得在支持 SWA 的模型（如 Mistral、最新的 Llama 3.1）上，即使物理缓存只有 4096 个槽，也能支持远远超过 4096 的上下文，因为旧位置会被透明地淘汰。对于非 SWA 模型，这个条件不会触发，cell 必须等到序列结束或显式释放才能被复用，策略自然退化。

3.3 Flash Attention 集成

llama.cpp 的 CUDA 后端内建了 FlashAttention 的实现，代码位于 ggml/src/ggml-cuda/fattn-flash.cu。这部分代码实现了 IO-aware 的 tiling 算法：

• 将 Q 矩阵按块加载到共享内存 / 寄存器，与分块的 K、V 做在线 softmax 规约。
• 利用 warp 级操作和 shuffle 指令在块内并行计算。
• 支持 causal masking，并且与 llama.cpp 的 KQ mask 生成对接，确保 mask 的正确应用。

虽然因篇幅无法展开 CUDA kernel 的全部细节，但其设计遵循原始 FlashAttention-2 的骨架构：将 Q 作为外循环，K、V 分块流入，通过维护一个运行中的 softmax 统计量避免存储中间结果，从而将 HBM 访问量从 O(N^2) 降至 O(N)。

4. 性能与权衡

上述各种技术并非零成本，各自牺牲了某些方面：

总体来看，llama.cpp 在设计时充分考虑了缓存命中率与内存限制的平衡，利用逻辑移位和 SWA 策略避免物理移动，用 Hadamard 旋转打通了量化存储与精确注意力之间的障碍，这些设计非常适合消费级硬件上的长文本推理。

5. 相关论文

• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness — Dao et al., 2022. (NeurIPS 2022, arXiv:2205.14135) — 本篇为 FlashAttention 集成提供了理论基础。
• QuaRot: Outlier-Free 4-Bit Inference in Rotated LLMs — Ashkboos et al., 2024. — 启发了 Hadamard 旋转应用于量化 KV 缓存的方法。

6. 总结

llama.cpp 的 KV cache 是一套精密的工程系统：以 3D tensor 承载多流，以 cell 元数据解耦分配，以 SWA 窗口为隐式 LRU 实现无限上下文，以逻辑移位规避物理搬移，以 Hadamard 旋转在低位宽 KV 存储下保持注意力质量，并最终通过 Flash Attention 将计算压力压在更靠近芯片的位置。这些技术的组合使得 llama.cpp 能够在哪怕一张消费级显卡上，流畅地运行数十万 token 的上下文推理，而不至于被内存和计算所困。