突破内存瓶颈：llama.cpp 项目中 KV 缓存优化策略全解析

在初始化过程中，缓存系统会根据模型的层数和每层的维度，计算所需的内存空间，并为不同的硬件设备（如 CPU、GPU）分配相应的缓存区域。例如，对于支持硬件加速的层，缓存会被分配到 GPU 内存中，以提高访问速度。

缓存大小计算

KV 缓存的大小直接影响模型的内存占用和推理性能。llama.cpp 在初始化时会打印缓存的详细信息，包括总大小、K 和 V 缓存的分别占用等：

LLAMA_LOG_INFO("%s: size = %7.2f MiB (%6u cells, %3d layers, %2u/%u seqs), K (%s): %7.2f MiB, V (%s): %7.2f MiB\n",
    __func__, (float)(memory_size_k + memory_size_v) / (1024.0f * 1024.0f),
    kv_size, (int) layers.size(), n_seq_max, n_stream,
    ggml_type_name(type_k), (float)memory_size_k / (1024.0f * 1024.0f),
    ggml_type_name(type_v), (float)memory_size_v / (1024.0f * 1024.0f));

这段代码计算并打印了 KV 缓存的总大小、K 缓存和 V 缓存的大小，以及缓存的层数和序列数。这些信息对于评估模型的内存需求和性能优化至关重要。

创新优化：llama.cpp 的 KV 缓存策略

llama.cpp 在 KV 缓存优化方面引入了多项创新技术，旨在平衡推理速度、内存占用和模型性能。

1. 动态内存管理与缓存复用

llama.cpp 的 KV 缓存系统采用了动态内存管理策略，能够根据输入序列的长度和数量，灵活调整缓存的分配和使用。核心函数 llama_kv_cache::seq_rm 用于从缓存中移除指定序列的数据，释放内存空间：

bool llama_kv_cache::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
    GGML_ASSERT(seq_id == -1 || (seq_id >= 0 && (size_t) seq_id < seq_to_stream.size()));
    if (p0 < 0) { p0 = 0; }
    if (p1 < 0) { p1 = std::numeric_limits<llama_pos>::max(); }
    if (seq_id >= 0) {
        auto & cells = v_cells[seq_to_stream[seq_id]];
        auto & head = v_heads[seq_to_stream[seq_id]];
        uint32_t new_head = cells.size();
        for (uint32_t i = 0; i < cells.size(); ++i) {
            if (!cells.pos_in(i, p0, p1)) { continue; }
            if (cells.seq_has(i, seq_id) && cells.seq_rm(i, seq_id)) {
                if (new_head == cells.size()) { new_head = i; }
            }
        }
        // 如果释放了插槽，更新 head 以便下次从这里开始搜索
        if (new_head != cells.size() && new_head < head) { head = new_head; }
    } else {
        // 匹配所有序列
        // ... 代码省略 ...
    }
    return true;
}

seq_rm 函数通过遍历缓存中的单元格，移除与指定序列相关的数据，并更新缓存的头部指针（head），以便下次分配时从释放的位置开始搜索，提高缓存利用率。

2. 分层 KV 缓存与设备卸载

llama.cpp 支持将不同层的 KV 缓存分配到不同的计算设备上，实现'设备卸载'（offload）。例如，可以将计算密集型的层缓存分配到 GPU，而将其他层缓存保留在 CPU 内存中。这一功能在 llama_kv_cache 的构造函数中实现：

ggml_backend_buffer_type_t buft = ggml_backend_cpu_buffer_type();
if (offload) {
    auto * dev = model.dev_layer(il);
    buft = ggml_backend_dev_buffer_type(dev);
    dev_name = ggml_backend_dev_name(dev);
}
LLAMA_LOG_DEBUG("%s: layer %3d: dev = %s\n", __func__, il, dev_name);

通过这种方式，llama.cpp 可以充分利用异构计算资源，平衡内存占用和计算效率。相关的设备管理逻辑在 llama_model 类中实现。

3. 滑动窗口注意力（SWA）支持

为了处理更长的输入序列，llama.cpp 集成了滑动窗口注意力（Sliding Window Attention, SWA）机制。SWA 通过限制注意力计算的窗口大小，只关注最近的 k 个 token，从而降低内存占用和计算复杂度。这一功能在 llama_kv_cache_iswa 类中实现：

llama_kv_cache_iswa::llama_kv_cache_iswa(
    const llama_model & model, ggml_type type_k, ggml_type type_v,
    bool v_trans, bool offload, bool swa_full, bool unified,
    uint32_t kv_size, uint32_t n_seq_max, uint32_t n_ubatch,
    uint32_t n_pad, const layer_filter_cb & filter,
    const layer_reuse_cb & reuse)
: hparams(model.hparams), unified(unified) {
    // 创建非 SWA 层的 KV 缓存
    kv_base = std::make_unique<llama_kv_cache>(
        model, type_k, type_v, v_trans, offload, unified,
        size_base, n_seq_max, n_pad, 0, LLAMA_SWA_TYPE_NONE, filter_base, reuse);
    // 创建 SWA 层的 KV 缓存
    kv_swa = std::make_unique<llama_kv_cache>(
        model, type_k, type_v, v_trans, offload, unified,
        size_swa, n_seq_max, n_pad, hparams.n_swa, hparams.swa_type, filter_swa, reuse);
}

llama_kv_cache_iswa 类维护了两个 KV 缓存实例：kv_base 用于非 SWA 层，kv_swa 用于 SWA 层。这种分离设计允许模型对不同层采用不同的注意力策略，在性能和效率之间取得平衡。

高级特性：KV 缓存的动态管理

llama.cpp 的 KV 缓存系统还提供了多项高级特性，支持复杂场景下的缓存管理和优化。

序列复制与状态迁移

在多轮对话或批处理场景中，经常需要复制或迁移序列的 KV 缓存状态。llama.cpp 提供了 seq_cp 函数，用于复制序列的缓存数据：

void llama_kv_cache::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) {
    GGML_ASSERT(seq_id_src >= 0 && (size_t) seq_id_src < seq_to_stream.size());
    GGML_ASSERT(seq_id_dst >= 0 && (size_t) seq_id_dst < seq_to_stream.size());
    const auto s0 = seq_to_stream[seq_id_src];
    const auto s1 = seq_to_stream[seq_id_dst];
    if (s0 == s1) {
        // 同一流内的复制，只需更新元数据
        // ... 代码省略 ...
    } else {
        // 跨流复制，需要复制实际数据
        // ... 代码省略 ...
    }
}

当源序列和目标序列位于同一流（stream）时，seq_cp 只需更新缓存的元数据，无需复制实际的 Key 和 Value 数据，从而提高效率。当跨流复制时，函数会安排实际的数据复制操作，确保目标序列能够正确复用源序列的缓存结果。

K-shift：缓存的高效更新

在处理长序列时，当缓存空间不足，llama.cpp 会使用 K-shift 技术来更新缓存内容。K-shift 通过移动缓存中的 Key 和 Value 矩阵，为新的 token 腾出空间，同时保持注意力计算的正确性。这一过程在 llama_kv_cache::update 函数中实现：

bool llama_kv_cache::update(llama_context * lctx, bool do_shift, const stream_copy_info & sc_info) {
    bool updated = false;
    // 处理流复制操作
    // ... 代码省略 ...
    if (do_shift) {
        if (!get_can_shift()) {
            GGML_ABORT("The current KV cache / model configuration does not support K-shift");
        }
        LLAMA_LOG_DEBUG("%s: applying K-shift\n", __func__);
        // 应用 K-shift
        if (hparams.rope_type != LLAMA_ROPE_TYPE_NONE) {
            ggml_backend_sched_reset(sched);
            auto * res = lctx->get_gf_res_reserve();
            res->reset();
            auto * gf = build_graph_shift(res, lctx);
            if (!ggml_backend_sched_alloc_graph(sched, gf)) {
                LLAMA_LOG_ERROR("%s: failed to allocate compute graph for K-shift\n", __func__);
                return updated;
            }
            res->set_inputs(nullptr);
            if (lctx->graph_compute(gf, false) != GGML_STATUS_SUCCESS) {
                LLAMA_LOG_ERROR("%s: failed to compute K-shift\n", __func__);
                return updated;
            }
            updated = true;
        }
        // 重置 shift 状态
        for (uint32_t s = 0; s < n_stream; ++s) {
            auto & cells = v_cells[s];
            cells.reset_shift();
        }
    }
    return updated;
}

K-shift 技术的核心在于通过旋转位置编码（RoPE）来调整缓存中的 Key 矩阵，使得旧的 token 数据能够被新的 token 覆盖，同时保持相对位置信息的正确性。这一技术大大提高了缓存的利用率，允许模型处理更长的序列。

实践指南：KV 缓存的调优与应用

了解了 llama.cpp 中 KV 缓存的原理和实现后，我们来看看如何在实际应用中优化和配置 KV 缓存，以获得最佳的性能。

缓存大小配置

KV 缓存的大小是影响模型性能的关键参数。在 llama.cpp 中，可以通过命令行参数 --kvsize 来指定 KV 缓存的大小（以 token 数为单位）。例如：

./main -m models/7B/ggml-model-q4_0.bin -p "Hello world" --kvsize 2048

这个命令将 KV 缓存大小设置为 2048 tokens。选择合适的缓存大小需要平衡内存限制和模型性能：

• 缓存过小时：模型需要频繁地进行 K-shift 或重新计算，导致推理速度下降。
• 缓存过大时：会占用过多的内存资源，可能导致内存溢出或其他应用程序性能受影响。

llama.cpp 在启动时会打印 KV 缓存的详细信息，例如：

llama_kv_cache_init: size = 256.00 MiB ( 4096 cells, 32 layers, 1/1 seqs), K (f16): 128.00 MiB, V (f16): 128.00 MiB

通过这些信息，你可以判断当前的缓存配置是否合理，并进行相应调整。

SWA 参数调优

对于支持滑动窗口注意力（SWA）的模型，可以通过调整 SWA 相关参数来优化长序列处理性能。在 llama.cpp 中，SWA 的配置主要通过模型的超参数（hparams）实现。

关键的 SWA 参数包括：

• n_swa：滑动窗口的大小。
• swa_type：SWA 的类型，如 LLAMA_SWA_TYPE_NONE（禁用 SWA）、LLAMA_SWA_TYPE_SLIDING（滑动窗口）等。

你可以通过修改模型的超参数或在推理时指定相关参数来调整 SWA 行为。例如，在加载模型时指定滑动窗口大小：

./main -m models/7B/ggml-model-q4_0.bin --swa-window 512

监控与调试

llama.cpp 提供了多种工具和选项，帮助你监控和调试 KV 缓存的行为：

1. 环境变量 LLAMA_KV_CACHE_DEBUG：设置该变量可以启用 KV 缓存的调试日志。例如：

   export LLAMA_KV_CACHE_DEBUG=1
   ./main -m models/7B/ggml-model-q4_0.bin -p "Hello world"
   

   启用调试后，llama.cpp 会打印详细的 KV 缓存操作日志，包括缓存的分配、更新和释放等信息。

2. 缓存使用统计：通过 llama_kv_cache::memory_breakdown 函数可以获取不同设备上 KV 缓存的内存占用情况：

   std::map<ggml_backend_buffer_type_t, size_t> llama_kv_cache::memory_breakdown() const {
       std::map<ggml_backend_buffer_type_t, size_t> ret;
       for (const ggml_backend_buffer_ptr & buf_ptr : bufs) {
           ret[ggml_backend_buffer_get_type(buf_ptr.get())] += ggml_backend_buffer_get_size(buf_ptr.get());
       }
       return ret;
   }
   

   这一功能可以帮助你识别内存瓶颈，优化缓存的设备分配策略。

总结与展望

KV 缓存在大模型推理中扮演着至关重要的角色，直接影响模型的推理速度和内存占用。llama.cpp 作为一个高效的 LLaMA 模型实现，在 KV 缓存优化方面提供了丰富的功能和灵活的配置选项，包括动态内存管理、分层设备卸载和滑动窗口注意力支持等。

通过深入理解 llama.cpp 中 KV 缓存的实现，如核心算法和数据结构，以及 SWA 优化，你可以更好地调优模型性能，满足不同场景的需求。

未来，随着模型规模的不断增大和硬件技术的发展，KV 缓存技术仍有很大的优化空间。llama.cpp 社区也在持续探索新的缓存优化策略，如更智能的缓存淘汰算法、自适应的窗口大小调整等。我们期待看到 llama.cpp 在 KV 缓存优化方面带来更多创新，为大模型的高效部署做出更大贡献。