Llama.cpp Attention kv Cache

Transformer 的注意力机制 Attention 是核心组件，用于捕捉输入序列中 token 之间的关系，这个关系是通过 Q、K、V 结构建模的, QKV 是 Attention 的核心矩阵

KV-cache 是 Transformer 模型推理中的关键优化技术，用于存储注意力机制 Attention 中的键 Key 和值 Value 张量，以加速自回归生成

在自回归生成中，每个新 token 的生成需要计算当前 token 的 Query (Q) 与之前所有 token 的 Key 和 Value 进行注意力计算。KV-cache 保存之前的 K 和 V，避免重复计算

KV-cache 占用额外内存，但它避免重复计算整个序列的 K 和 V，整体效率更高

Transformer 中的 Attention

绝大多数 LLM（如 Qwen3、LLaMA、GPT）使用 Scaled Dot-Product Attention 或其变体（比如 GQA），Attention 都是 Q、K、V 的函数

Q（Query）：当前 token 的查询向量。
K（Key）：所有 token 的键向量。
V（Value）：所有 token 的值向量。

自回归生成：每个新 token 的 Q 需要与之前所有 token 的 K 和 V 计算注意力

Attention的输入 $x$ 其实是

$ Q_i = W_q x_i $, $ K_i = W_k x_i $, $ V_i = W_v x_i $ w 和 x 分别是什么？

$ x_i \in \mathbb{R}^{d_{\text{model}}} $ 是序列中第 $ i $ 个 token 的嵌入向量，来源于embedding层。$ d_{\text{model}} $ 是模型的隐藏维度。
$ W_q, W_k, W_v \in \mathbb{R}^{d_{\text{model}} \times d_{\text{model}}} $ 是注意力层的线性变换矩阵，用于将输入 $ x_i $ 投影到查询、键、值空间。

$ W_q $, $ W_k $, $ W_v $ 是 Transformer 模型参数的一部分，在预训练或微调阶段通过优化算法（如 Adam）学习。最小化损失函数（如交叉熵），通过梯度下降更新这些权重，使模型更好地捕捉 token 间的关系。

在训练开始时，$ W_q $, $ W_k $, $ W_v $ 通常通过某种方式初始化（如 Xavier 或 He 初始化）设置初始值。训练后优化好的值已经存在于GGUF文件中。

KAQ：隐藏层和隐藏维度

指 Transformer 的编码器（Encoder）和解码器（Decoder）中多层的整体结构。每一个隐藏层包括 Attention层、FFN层、残差连接层。隐藏层的数量具体看模型的配置。隐藏层的输入输出维度就是隐藏维度：$ d_{\text{model}} $

也就是说，隐藏层的第一个模块是 aattention,第二个模块是FNN，残差连接和归一化是隐藏层的最后一个模块。

所有子模块（注意力、FFN、残差、归一化）输入输出均为 $d_{\text{model}}$ 维（如 512），确保隐藏层间数据传递一致。

KV-cache 优化

初始生成：计算第一个 token 时，生成并存储所有 token 的 K 和 V 到 KV-cache 空间
后续生成：每个新 token 只计算当前 Q、K、V，其中 K 和 V 追加到 KV-cache，注意力计算仅涉及当前 Q 和缓存的 K、V
内存结构：KV-cache 按层存储：每 layer 一个 K 和 V 张量。形状：[n_layers, n_heads, seq_len, head_dim]，其中 seq_len 随生成增长。

对于 llama-cli，prompt 的长度通过 -c, --ctx-size设置，默认值是 4096

KV-cache 的工作流程

初始化：./llama-cli -m ./models/Qwen3-1.7B-Q4_K_M.gguf -c 4096 设置最大序列长度 4096

在推理过程中

第一次 FFN 时，计算所有 token （即prompt）的 Q、K 和 V，将 K 和 v 存储到 KV-cache
当前 token 先计算 Q、K、V，追加 K、V 到 KV-cache，再计算 Attention。然后用当前 Q 与整个 KV-cache（包括新 K 和 V）计算注意力。
推理完成或上下文重置时，清空 KV-cache

注：

上述的 “计算所有 token 的”，“所有”指的是输入prompt 经过tokenize 后的token 和在自回归过程中已生成的 token
一般会对 kv-cache 进行量化。

我的 case ：

qwen3.attention.head_count = 16    注意力头数（Query 头），即多头注意力（MHA）的总头数。 
qwen3.attention.head_count_kv = 8  键值头数，表明使用 Grouped Query Attention GQA，K 和 V 的头数少于 Q 的头数。
qwen3.rope.freq_base = 1000000.000000  旋转位置编码（RoPE）基频，影响 KV-cache 的位置嵌入计算
qwen3.attention.layer_norm_rms_epsilon = 0.000001  层归一化的 epsilon，间接影响注意力计算稳定性
qwen3.attention.key_length = 128    每个 Key 向量的维度
qwen3.attention.value_length = 128  每个 Key 向量的维度

表明 Qwen3 使用 GQA，K 和 V 的头数减少一半（8 vs 16）。这将 KV-cache 的内存占用减半，因为 KV-cache 只存储 K 和 V 的张量

如果层数是28层，上下文长度是 4096，存储精度是 FP16，则 kv-cache 占用内存为：

2 (K+V) × 28 (layers) × 8 (heads_kv) × 4096 (seq_len) × 128 (head_dim) × 2 (FP16 字节) = 469,762,048 bytes = 0.469 GB

key_length = value_length = head_dim 是标准设计, 用户输入上下文长度 n_ctx 就直接影响了 kv-cache 的内存占用。

Scaled Dot-Product Attention

公式：

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}} + M\right)V$$

输入：

查询（Query）：$ Q \in \mathbb{R}^{n \times d_k} $，$ n $ 是序列长度（你的 n_ctx = 4096），$ d_k $ 是每个头的维度（例如 qwen3.embedding_length = 2048 除以 head_count = 16，约 $ d_k $ = 128 与log 中 qwen3.attention.key_length 和 qwen3.attention.value_length 相同）。
键（Key）：$ K \in \mathbb{R}^{n \times d_k} $。
值（Value）：$ V \in \mathbb{R}^{n \times d_v} $，通常 $ d_v = d_k $。
可选掩码（Mask）：$ M \in \mathbb{R}^{n \times n} $，用于因果注意力（防止未来 token 影响当前 token）。

步骤：

点积：计算 $ QK^T \in \mathbb{R}^{n \times n} $，表示 q 和 k 的相似度。
缩放：除以 $ \sqrt{d_k} $ 防止点积过大导致 softmax 饱和（数值稳定性）
掩码（可选）：添加掩码 $ M $（例如因果掩码，未来 token 设为 $-\infty$）。
Softmax：归一化得到注意力权重： $$A = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}} + M\right), \quad A \in \mathbb{R}^{n \times n}$$
加权值：用注意力权重 $ A $ 加权值矩阵 $ V $： $$\text{Output} = AV \in \mathbb{R}^{n \times d_v}$$

伪代码：

// Scaled dot-product attention
ggml_tensor * qk = ggml_matmul(ctx, q, ggml_transpose(ctx, k)); // QK^T
qk = ggml_scale(ctx, qk, 1.0f / sqrtf(d_k)); // 缩放
qk = ggml_add(ctx, qk, mask); // 掩码
qk = ggml_softmax(ctx, qk); // Softmax
output = ggml_matmul(ctx, qk, v); // AV

Multi-Head Attention

公式： $$\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \text{head}_2, \dots, \text{head}_h)W^O$$

其中：

每个头：$ \text{head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V) $。
$ W_i^Q, W_i^K, W_i^V \in \mathbb{R}^{d_{\text{model}} \times d_k} $ 是Q、K、V的投影矩阵。
$ W^O \in \mathbb{R}^{h \cdot d_v \times d_{\text{model}}} $ 是输出投影矩阵。
$ h $ 是头数（我的 28），$ d_{\text{model}} = 2048 $。

GQA 变体：

键和值共享 $ h_{\text{kv}} = 8 $ 个头，查询有 $ h = 28 $ 个头。
每 $ \frac{h}{h_{\text{kv}}} = \frac{28}{8} = 3.5 $ 个查询头共享一组键/值头，减少 KV 缓存内存。

Rotary Positional Encoding (RoPE)

llama.cpp 使用 RoPE（qwen3.rope.freq_base = 1000000）为注意力机制添加位置信息，取代绝对位置编码。

公式首先：

对查询和键应用旋转矩阵： $$Q’ = Q \cdot R(\theta), \quad K’ = K \cdot R(\theta)$$
$ R(\theta) $ 是旋转矩阵，基于位置 $ m $ 和频率 $ \theta_i = \text{freq_base}^{-2i/d_k} $。
例如，对于我的case $d_k$ = 128， $ \theta_i = (10^6)^{-2i/128} $

然后计算注意力： $$\text{Attention}(Q’, K’, V) = \text{softmax}\left(\frac{Q’ K’^T}{\sqrt{d_k}} + M\right)V$$

源码 self-attention 的计算步骤

我的实例是 qwen3 架构，所以在 build_graph() 中的 struct llm_build_qwen3 : public llm_graph_context 中实现了 self-attention 的计算：

struct llm_build_qwen3 : public llm_graph_context {
...
// self-attention
{
    // compute Q and K and RoPE them
    // W^Q，W^K，W^V 已经是学习过的权矩阵了
    // X 是embedding 层的输出
    // Q = X * W^Q
    ggml_tensor * Qcur = build_lora_mm(model.layers[il].wq, cur);
    cb(Qcur, "Qcur", il);
    // K = X * W^K
    ggml_tensor * Kcur = build_lora_mm(model.layers[il].wk, cur);
    cb(Kcur, "Kcur", il);
    // V = X * W^V
    ggml_tensor * Vcur = build_lora_mm(model.layers[il].wv, cur);
    cb(Vcur, "Vcur", il);

    Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head,    n_tokens);
    Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens);
    Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens);

    Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, NULL, LLM_NORM_RMS, il);
    cb(Qcur, "Qcur_normed", il);

    Qcur = ggml_rope_ext(
            ctx0, Qcur, inp_pos, nullptr,
            n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
            ext_factor, attn_factor, beta_fast, beta_slow
            );

    Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, NULL, LLM_NORM_RMS, il);
    cb(Kcur, "Kcur_normed", il);

    Kcur = ggml_rope_ext(
            ctx0, Kcur, inp_pos, nullptr,
            n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
            ext_factor, attn_factor, beta_fast, beta_slow
            );

    cb(Qcur, "Qcur", il);
    cb(Kcur, "Kcur", il);
    cb(Vcur, "Vcur", il);

    cur = build_attn(inp_attn,
            model.layers[il].wo, model.layers[il].bo,
            Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, 1.0f/sqrtf(float(n_embd_head)), il);
}}

上述实现了应用 RoPE 到 MHA 的 Q 和 K，然后计算注意力。在最后步骤 build_attn() 时，已将发生了存储 kv 的动作。

KAQ：attention 和 kv-cache 工作流程

初始阶段：

输入 prompt 分词后（如 100 个 token），计算所有 token 的 Q、K、V。K 和 V 存储到 KV-cache，seq_len=100。

初始阶段应该发生在 build_graph()

生成阶段：

对于每个新 token 计算当前 Q、K、V。
当前 Q 与 KV-cache 中之前所有 token 的 K 和 V（包括 prompt 和已生成 token）进行注意力计算：Attention(Q, K, V)
新 K 和 V 追加到 KV-cache，seq_len 增 1。

自然地，生成阶段发生在 compute_graph()，如何做的步骤已经在graph中了。

KAQ：build_graph 时计算和存储的是什么

build_attn() 函数中有 store kv-cache。但是不更新 kv-cache？【然】只有在 computer_graph() 中一个接一个 token 生成时才会更新 kv-cache。？

ggml_tensor * llm_graph_context::build_attn(){

    // 这是存储 kv-cahe 对象的结构，是该函数的输入
    llm_graph_input_attn_kv * inp;

    // 这是 kv-cache 的 context
    const auto * mctx_cur = inp->mctx;
    {
    // 获取k和v在 kv-cache 中的索引位置
    const auto & k_idxs = inp->get_k_idxs();
    const auto & v_idxs = inp->get_v_idxs();

    // 拷贝 k 和 v 到 kv-cache 中
    // 其中这两个cpy操作动作是，存储 k_cur and v_cur 在cache中，基于头位置信息
    // cpy_k() 实际上执行 kv->cpy_k(),  其中 kv 的类型是 llama_kv_cache pointer， 
    // 它属于 mctx_cur 对象的一个成员，是实际上存储kv-cache的 内存
    ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il));
    ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, v_cur, v_idxs, il));
    }
    // ggml_build_forward_expand 将复制操作添加到 GGML 计算图（gf）
}

所以是的，build_graph 阶段只是计算初始的 qkv，并将 kv-cache 存储在 graph中。

KAQ：kv-cahe 如何更新的

KV-cache 更新不发生在 build_graph 阶段，而是运行时状态，在 llama_decode 阶段动态更新？【非也】

构建 graph 时 store kv-cache 和访问 kv-cache，然后在 decode 时，即执行推理时，kv-cache 的更新就是根据前面定义好的 graph 进行的，因为 graph中已经有了更新的节点了，所以执行到那里，自然就更新了 kv-cache 了。

KAQ：code 中 kv-cache 存在什么对象中

存储在 llama_kv_cache 对象 kv 中，which 在 llama_graph 对象中（属于 llm_graph_context 类的 ggml_cgraph 对象中）。即在 build_graph 时，在 llama_graph 对象中创建了 kv-cache。

位置见 llm_graph_context::build_attn()

llama_kv_cache_context 对象中有成员 llama_kv_cache* kv 而 llama_kv_cache 对象有两个方法 cpy_k() 和 cpy_v() 用于拷贝 K, V 到 kv-cache。


class llm_graph_input_attn_kv {
    const llama_kv_cache_context * mctx;
};

class llama_kv_cache_context {
    llama_kv_cache * kv;
};

class llama_kv_cache {
    struct kv_layer {
        uint32_t il; // layer index in the model
        ggml_tensor * k;
        ggml_tensor * v;
        ...
    };
    // 每个层对应一个 kv_layer对象，即对应每层中存储的 k 和 v
    std::vector<kv_layer> layers
    // 存储k_cur and v_cur 在对应位置上
    ggml_tensor * cpy_k();
    ggml_tensor * cpy_v();
};

KAQ：Q 不需要缓存？

K, V 需要存储在 KV-cache，目的是复用历史 token 的信息。而 Q 不需要缓存，因为每次生成新 token，重新计算 Q，因为它依赖当前输入。

这正是推理过程的逻辑：使当前输入token的 Q 与之前所有token的 K、V 进行注意力计算。如果我需要生成 5 个新的token，就需要5次前向计算，即 5 次计算 Q，5 次计算 K 和 V，然后 5 次注意力计算。这是冗余的（为什么是冗余的，这涉及计算细节），在生成第一个 token 时计算 k 和 v ，之后将其缓存起来，在计算第二个token时，只要将第二个token的 k 和 v 追加到缓存中即可得到完整的用于第二个token的 k 和 v。

$ Q_i = W_q x_i $, $ K_i = W_k x_i $, $ V_i = W_v x_i $

在计算时的 shape 是什么样的？

$ x_i $: 当前 token 的嵌入，shape 为 [batch_size, hidden_size]
$ W_q, W_k, W_v $: 权重矩阵，shape 为 [hidden_size, n_head * head_dim]
$ Q_i, K_i, V_i $ 的 shape 相同，为 [batch_size, n_head, head_dim]
在Multihead Attention 中，Q, K, V 会被 reshape 为 [batch_size, n_head, seq_len, head_dim]

QKV shape在计算过程中的变化：给出一实例。假设：Qwen3-1.7B，hidden_size=2048, n_head=16, head_dim=128, 序列 “我爱中国”。

Token “[CLS]”：
- 预测 “我”。
- 输入：[CLS]，seq_len=1。
- $ Q_0, K_0, V_0 $: [1, 16, 1, 128]。
- KV Cache：[1, 16, 1, 128]（seq_len=1）。
Token “我”：
- 预测 “爱”。
- 输入：[CLS], 我，seq_len=2。
- $ Q_1 $: [1, 16, 1, 128]。
- $ K_1, V_1 $: [1, 16, 1, 128]，追加到 KV Cache。
- KV Cache：[1, 16, 2, 128]（seq_len=2）。
Token “爱”：
- 预测 “中国”。
- 输入：[CLS], 我, 爱，seq_len=3。
- $ Q_2 $: [1, 16, 1, 128]。
- KV Cache：[1, 16, 3, 128]（seq_len=3）。

其中的 seq_len 表示当前处理的序列长度，即输入序列中 token 的数量，每预测一个 Token，seq_len 都会 +1。上述例子中：

$ Q $ 只与当前Token有关: [batch_size, n_head, 1, head_dim]（当前 token）。
$ K, V $ 与当前token和历史token有关: [batch_size, n_head, seq_len, head_dim]（历史 + 当前 token）。

不使用 kv-cache 的过程

对于已有 100 个 token 的序列，生成 5 个新 token（$ t_{101} $, …, $t_ {105} $）：

第 $ n $ 次前向传递（$ n=1,…,5 $）：

输入序列：[$t_1$, …, $t_{100+n-1}$]（长度 $ 100+n-1 $）。
嵌入：生成 [$x_1$, …, $x_{100+n-1}$]，形状 [100+n-1, 2048]。
对于每层（共 28 层）：
- 计算所有 token 的 Q、K、V： $$Q_i = W_q x_i, \quad K_i = W_k x_i, \quad V_i = W_v x_i \quad (i=1,…,100+n-1)$$
- 应用 RoPE, 其中 $ \theta_i = 1000000^{-2i/128} $.
- Attention：对最后一个 token 的 Q： $$\text{Attention}(Q_{100+n-1}, [K_1, …, K_{100+n-1}], [V_1, …, V_{100+n-1}])$$

输出：预测 $ t_{100+n} $。

重复计算：每次前向传递重新计算整个序列（包括 prompt 和已生成 token）的 Q、K、V，导致计算量为 $O(n²)$。

Transformer 中的 Attention#

Attention的输入 $x$ 其实是#

KAQ：隐藏层 和 隐藏维度#

KV-cache 优化#

KV-cache 的工作流程#

Scaled Dot-Product Attention#

Multi-Head Attention#

Rotary Positional Encoding (RoPE)#

源码 self-attention 的计算步骤#

KAQ：attention 和 kv-cache 工作流程#

初始阶段：#

生成阶段：#

KAQ：build_graph 时计算和存储的是什么#

KAQ：kv-cahe 如何更新的#

KAQ：code 中 kv-cache 存在什么对象中#

KAQ：Q 不需要缓存 ？#

不使用 kv-cache 的过程#

更多内容#