理解llama.cpp如何进行LLM推理

目录

• • 前言
  • [0. 简述](#0. 简述)
  • [1. 从prompt输入到输出的整体流程](#1. 从prompt输入到输出的整体流程)
  • [2. 理解ggml的张量](#2. 理解ggml的张量)
  • • [2.1 张量的基本结构](#2.1 张量的基本结构)
    • [2.2 张量操作和视图](#2.2 张量操作和视图)
    • [2.3 张量计算](#2.3 张量计算)
    • [2.4 将计算放到GPU上](#2.4 将计算放到GPU上)
  • [3. Tokenization](#3. Tokenization)
  • [4. Embeddings](#4. Embeddings)
  • [5. Transformer](#5. Transformer)
  • • [5.1 Self-attention](#5.1 Self-attention)
    • [5.2 Transformer的各层](#5.2 Transformer的各层)
    • [5.3 计算logits](#5.3 计算logits)
  • [6. Sampling](#6. Sampling)
  • • [6.1 Greedy sampling](#6.1 Greedy sampling)
    • [6.2 Temperature sampling](#6.2 Temperature sampling)
  • [7. 推理优化](#7. 推理优化)
  • • [7.1 KV cache](#7.1 KV cache)
    • [7.2 持续优化](#7.2 持续优化)
    • [7.3 KV cache的实际应用](#7.3 KV cache的实际应用)
  • 结语
  • 参考

前言

看到一篇非常不错的文章和大家分享下，主要是解决了博主之前的很多困惑，记录下个人学习笔记，仅供自己参考😄

本文分享 llama.cpp 中 LLM 推理的整体流程，对 llama.cpp 框架代码的解析并不是很多，后面有机会的话再和大家分享

refer1：Understanding how LLM inference works with llama.cpp

refer2：理解llama.cpp怎么完成大模型推理的

refer3：https://chatgpt.com/

0. 简述

在这篇文章中，我们将深入研究 LLM （L arge L anguage Model） 大语言模型的内部原理，以切实了解它们是如何工作的。我们将使用 llama.cpp 框架，它是 meta 开源的一个纯 c/c++ 的 llm 推理框架，其代码较简洁，也没有过多的抽象概念，是一个非常值得学习的推理框架

Note ：我们将使用 这个提交版本 来进行讲解

这里我们主要关注 LLM 的推理部分，更加侧重从工程角度出发来理解 llm 的内部原理。在本文中，我们将从头到尾介绍整个推理过程，包括如下内容：

• Tensors：简要概述如何使用 tensor 进行数学运算，部分运算可能在 GPU 上
• Tokenization：将用户输入的文本提示词（prompt）拆分成一个个 token 的过程，llm 会将这些 token 用作其输入
• Embedding：将 token 转换为向量表示的过程
• The Transformer：llm 架构的核心部分，负责实际推理过程，我们将重点关注 self-attention 自注意力机制部分
• Sampling：选择下一个预测 token 的过程，我们将探讨两种采样技术
• The KV cache：这是一种常用的优化技术，用于加快 llm 的推理速度，我们将探讨一个基本的 kv cache 实现

1. 从prompt输入到输出的整体流程

llm 例如 llama 的工作原理是接收提示词输入文本即 "prompt" 来预测下一个 token 或 word 是什么。为了说明这一过程，我们将以 Wikipedia 中关于描述量子力学文章的第一句话为例，我们的 prompt 是：

Quantum mechanics is a fundamental theory in physics that

llm 会尝试按照训练后认为最有可能的续写内容来续写这个句子，使用 llama.cpp，我们可以得到如下的续写：

provides insights into how matter and energy behave at the atomic scale.

让我们先来看看这一过程的整体流程，llm 的核心是每次只预测一个 token，因此要生成一个完整的句子，需要对同一 prompt 重复应用 llm 模型，并将之前的输出 token 添加到 prompt 中。这类模型也被称为自回归模型，这里我们将主要关注单个 token 的生成过程，如下图所示：

Note：从用户输入 prompt 到生成单个 token 的完整流程包括 tokenizaton、embedding、transformer 以及 sampling 等多个阶段，本篇文章将重点介绍这些内容

根据图示，整个流程如下：

• 1. 分词：tokenizer 分词器会将 prompt 分割成一系列 token，根据模型的 vocabulary 词汇表，有些单词会被拆分成多个 token，每个 token 由一个唯一的数字表示
• 2.嵌入 embedding 转换：每个数字 token 都被转换成一个 embedding 嵌入。embedding 是一个固定大小的向量，它能以更有效的方式表示 token，以便 llm 处理。所有的嵌入向量共同组成一个 embedding matrix 嵌入矩阵
• 3.输入 Transformer：嵌入矩阵是 transformer 的输入，transformer 是一个神经网络，是 llm 的核心。transformer 由多个 layer 组成，每一个 layer 接收一个输入矩阵，并利用模型参数对其进行各种数学运算，其中最核心的是 self-attention layer。该层的输出被用作下一层的输入
• 4. logits 生成：最后一个神经网络将 Transformer 的输出转换为 logits，每个可能的下一个 token 都有一个对应的 logit，它代表了该 token 是句子 "正确" 延续的概率
• 5. 采样：从 logits 列表中选择下一个 token 时，会使用几种采样技术中的一种
• 6. 生成输出 ：被选中的 token 作为输出返回。若要继续生成 token，则将所选 token 添加到 步骤 1 的 token 列表中，然后重复该过程。这人过程可以一直持续到生成所需的 token 数量，或者 llm 发出一个特殊的结束流（EOS）token

接下来的内容就是围绕这整个流程，我们将详细介绍其中的每一个步骤，但在此之前，我们需要先熟悉一下 tensor

2. 理解ggml的张量

tensor（张量）是神经网络中用于执行数学运算的主要数据结构。llama.cpp 使用 ggml，它是一个纯 c++ 实现的张量库，相当于 python 生态系统中的 pytorch 或 tensorflow。我们将使用 ggml 来了解 tensor 是如何操作的

Note ：ggml 是一个用 c/c++ 编写、专注于 transformer 架构模型推理的机器学习库，它和 pytorch、tensorflow 等机器学习库比较类似。关于 ggml 更多的介绍大家可以参考 Georgi Gerganov 在 hugging face 上发布的一篇文章：Introduction to ggml

tensor 可以表示为一个多维数组，它可能包含一个单一的数字（标量）、一个向量（一维数组）、一个矩阵（二维数组）甚至三维或四维数组。

tensor 可以分为两类：

• 数据张量：这些张量保存实际数据，包含一个多维数组的数值
• 运算张量：这些张量只表示一个或多个其它张量之间的计算结果，只有在实际计算时才会保存数据

我们下面就来探讨下它们之间的区别

2.1 张量的基本结构

在 ggml 中，张量由 ggml_tensor 结构体表示，为方便起见，我们将其略作简化，其定义如下：

// ggml/include/ggml.h
struct ggml_tensor {
    enum ggml_type    type;
    enum ggml_backend backend;

    int     n_dims;
    // number of elements
    int64_t ne[GGML_MAX_DIMS];
    // stride in bytes
    size_t  nb[GGML_MAX_DIMS];

    enum ggml_op op;

    struct ggml_tensor * src[GGML_MAX_SRC];

    void * data;

    char name[GGML_MAX_NAME];
};

其中：

• type 字段表示 tensor 元素的原始类型。例如，GGML_TYPE_F32 表示每个元素都是 32 位浮点数
• backend 字段用来表示该 tensor 是在 CPU 上还是在 GPU 上，即数据后端是 CPU 后端还是 GPU 后端
• n_dims 表示 tensor 的维度数量，取值范围从 1 到 4
• ne 表示每个维度中的元素数量。ggml 采用行优先顺序，即 ne[0] 表示每一行的大小，ne[1] 表示每一列的大小，以此类推
• nb 稍微复杂一些，它表示步长（stride），即不同维度相邻元素之间的字节数。在第一个维度中，这将是原始元素的大小，在第二个维度中，它是行大小乘以元素大小，以此类推。例如，对于一个 4x3x2 的 tensor 来说，其步长 nb 计算如下：

其中：

• nb[0] = 4：表示沿着第 0 维（最内层）的移动，每移动 1 个元素需要 4 个字节的大小
• nb[1] = 16：表示沿着第 1 维移动 1 步，跳过 4 个元素，所以需要 4x4=16 个字节的大小
• nb[2] = 48：表示沿着第 2 维移动 1 步，跳过 4x3=12 个元素，所以需要 12x4=48 个字节的大小

使用步长的目的是为了在不复制任何数据的情况下执行某些张量运算。例如，在一个二维 tensor 上执行转置操作时（行变为列），只需翻转 ne 和 nb 并指向相同的底层数据即可完成：

// ggml/src/ggml.c (the function was slightly simplified).
struct ggml_tensor * ggml_transpose(
        struct ggml_context * ctx,
        struct ggml_tensor  * a) {
    // Initialize `result` to point to the same data as `a`
    struct ggml_tensor * result = ggml_view_tensor(ctx, a);

    result->ne[0] = a->ne[1];
    result->ne[1] = a->ne[0];

    result->nb[0] = a->nb[1];
    result->nb[1] = a->nb[0];

    result->op   = GGML_OP_TRANSPOSE;
    result->src[0] = a;

    return result;
}

下面是该函数的具体分析：

1. 创建 View：

函数首先调用 ggml_view_tensor(ctx, a)，创建了一个与原始张量 a 共用同一块数据的 "视图" 张量 result。这样做的好处是避免了数据的拷贝，只是在元数据（比如维度和步长）上做了调整

2. 交换维度（ne）：

• result->ne[0] = a->ne[1]
• result->ne[1] = a->ne[0]

对于一个形状为 4x3x2 的张量，原始的 a->ne[0]=4, a->ne[1]=3, a->ne[2]=2，经过交换后，result->ne[0]=3, result->ne[1]=4, result->ne[2]=2，即张量的形状变为 3x4x2

3. 交换步长（nb）：

同时交换了第一个和第二个维度对应的字节步长，这保证了在计算元素偏移时能够正确地访问数据

4. 操作标识与来源：

设置 result->op = GGML_OP_TRANSPOSE; 表明这是一个转置操作，同时将原始张量 a 地指针赋给 result->src[0]

总的来说，该函数仅仅是交换了 tensor 的前两个维度以及对应的步长，并没有对数据进行重新排列，生成的是一个共享原数据的视图。

这里博主对于步长交换有一个困惑，那就是转置后的 tensor 维度是 3x4x2，其步长为 nb[0]=16, nb[1]=4, nb[2]=48，但是按照 ggml_tensor 结构体的定义来说，对于一个 3x4x2 维度的 ggml_tensor 而言其步长应该为 nb[0]=4, nb[1]=12, nb[2]=48，为什么同样是 3x4x2 维度的 ggml_tensor，其步长却不同呢？🤔

这两者之所以不同，根本原因是：ggml_transpose() 函数并没有拷贝或重新布局内存，而是仅仅创建了一个"视图"（view） ，这个视图依然指向原来那片内存，但通过修改元数据（维度 ne 和步长 nb）来让 "逻辑上的索引" 满足转置关系。它只是告诉后续的计算："当我说 ( i 0 , i 1 , i 2 ) (i_0,i_1,i_2) (i0,i1,i2) 时，其实要去取原始张量里 ( i 1 , i 0 , i 2 ) (i_1,i_0,i_2) (i1,i0,i2) 位置的元素"

我们举个具体的索引例子来说明下：

• 在转置后新张量中，假设我们要访问索引为 ( i 0 , i 1 , i 2 ) = ( 2 , 1 , 1 ) (i_0,i_1,i_2)=(2,1,1) (i0,i1,i2)=(2,1,1) 位置的元素
• 逻辑上，这个原始应该等于原始张量索引为 ( x , y , z ) = ( 1 , 2 , 1 ) (x,y,z)=(1,2,1) (x,y,z)=(1,2,1) 处的值（因为转置交换了前两个维度）

下面让我们分别计算内存偏移：

1. 对于新张量（转置视图）

o f f s e t n e w ( 2 , 1 , 1 ) = 2 × 16 + 1 × 4 + 1 × 48 = 32 + 4 + 48 = 84 \mathrm{offset}_{\mathrm{new}}(2,1,1)=2\times 16+1\times 4+1\times 48=32+4+48=84 offsetnew(2,1,1)=2×16+1×4+1×48=32+4+48=84

2. 对于原始张量

o f f s e t o r i g ( 1 , 2 , 1 ) = 1 × 4 + 2 × 16 + 1 × 48 = 4 + 32 + 48 = 84 \mathrm{offset}_{\mathrm{orig}}(1,2,1)=1\times 4+2\times 16+1\times 48=4+32+48=84 offsetorig(1,2,1)=1×4+2×16+1×48=4+32+48=84

结果都是 84，说明这两个 "逻辑上转置" 的索引确实映射到同一个 物理地址。如果我们没有交换 nb[0] 和 nb[1]，而保持原来的 [4,16,48]，那么对 ( i 0 , i 1 , i 2 ) = ( 2 , 1 , 1 ) (i_0,i_1,i_2)=(2,1,1) (i0,i1,i2)=(2,1,1) 的偏移就会算成 2 × 4 + 1 × 16 + 1 × 48 = 8 + 16 + 48 = 72 2\times 4+1\times 16+1\times 48=8+16+48=72 2×4+1×16+1×48=8+16+48=72，显然不是我们想要的 84，也就拿错了数据

因此，经过 ggml_transpose() 函数之后维度为 3x4x2 的张量表示如下：

而不是下面的一个新的 3x4x2 维度的 ggml_tensor 表示：

2.2 张量操作和视图

如前所述，一些 tensor 保存数据，而另一些则表示其它 tensor 之间操作的理论结果，回到 ggml_tensor 结构体：

• op 可以是 tensor 之间任何受支持的操作，将其设置为 GGML_OP_NONE 则表示该 tensor 持有数据。其它值表示不同的操作，例如，GGML_OP_MUL_MAT 表示该 tensor 不保存数据，只表示其它两个 tensor 之间的矩阵乘法结果
• src 是一个指针数组，指向要在其间进行运算的 tensor，例如，如果 op == GGML_OP_MUL_MAT，那么 src 将包含指向两个要相乘的 tensor 的指针，如果 op == GGML_OP_NONE，那么 src 将是空的
• data 指向实际的 tensor 数据，如果该 tensor 是一个操作，则指向 NULL。它也可以指向另一个 tensor 的数据，即视图（view）。例如，在上面的 ggml_transpose() 函数中，生成的 tensor 是原始 tensor 的视图，只是翻转了维度（ne）和步长（nb），data 指向内存中的同一个位置

下面的矩阵乘法函数很好地诠释了这些概念：

// ggml/src/ggml.c (simplified and commented)
struct ggml_tensor * ggml_mul_mat(
        struct ggml_context * ctx,
        struct ggml_tensor  * a,
        struct ggml_tensor  * b) {
    // Check that the tensors' dimensions permit matrix multiplication.
    GGML_ASSERT(ggml_can_mul_mat(a, b));

    // Set the new tensor's dimensions
    // according to matrix multiplication rules.
    const int64_t ne[4] = { a->ne[1], b->ne[1], b->ne[2], b->ne[3] };
    // Allocate a new ggml_tensor.
    // No data is actually allocated except the wrapper struct.
    struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, MAX(a->n_dims, b->n_dims), ne);

    // Set the operation and sources.
    result->op   = GGML_OP_MUL_MAT;
    result->src[0] = a;
    result->src[1] = b;

    return result;
}

在上述函数中，result 不包含任何数据，它只是表示 a 和 b 相乘的理论结果

2.3 张量计算

上述 ggml_mul_mat() 函数或其它 tensor 运算并不计算任何东西，而只是为计算准备 tensor。换一种方式理解，它建立了一个 graph 计算图，其中每个 tensor 运算都是一个节点，而运算的来源是该节点的子节点。在矩阵乘法场景中，graph 中有一个父节点，其操作是 GGML_OP_MUL_MAT，同时还有两个子节点。

以 llama.cpp 为例，下面的代码实现了 self-attention，它是 Transformer layer 的一个组件，稍后我们将对其进行更深入的探讨：

// llama.cpp
static struct ggml_cgraph * llm_build_llama(/* ... */) {
    // ...

    // K,Q,V are tensors initialized earlier
    struct ggml_tensor * KQ = ggml_mul_mat(ctx0, K, Q);
    // KQ_scale is a single-number tensor initialized earlier.
    struct ggml_tensor * KQ_scaled = ggml_scale_inplace(ctx0, KQ, KQ_scale);
    struct ggml_tensor * KQ_masked = ggml_diag_mask_inf_inplace(ctx0, KQ_scaled, n_past);
    struct ggml_tensor * KQ_soft_max = ggml_soft_max_inplace(ctx0, KQ_masked);
    struct ggml_tensor * KQV = ggml_mul_mat(ctx0, V, KQ_soft_max);

    // ...
}

上述代码是一系列 tensor 运算，并构建了一个与 原始 Transformer 论文 中描述的计算图完全相同的计算图：

为了计算得到最终的 tensor 结果（这里是 QKV），需要采取以下步骤：

• 1. 加载数据 ：数据被加载到每个叶子张量的 data 指针中，示例中的叶子张量是 Q，K，V
• 2. 构建计算图 ：输出张量（QKV）通过 ggml_build_forward() 转换为计算图，该函数相对简单，按深度优先顺序排列节点
• 3. 执行计算图 ：计算图使用 ggml_graph_compute()，它以深度优先的顺序在每个节点上执行 ggml_compute_forward()。它负责主要的数学运算，并将结果填入到张量的 data 指针中
• 4. 结果输出 在这一过程结束时，输出张量的 data 指针指向最终的计算结果

2.4 将计算放到GPU上

由于 GPU 的高度并行性，许多张量运算（如矩阵加法和乘法）在 GPU 上的计算效率要高得多。当 GPU 可用时，可以用 tensor->backend = GGML_BACKEND_GPU 来标记 tensor。在这种情况下，ggml_compute_forward() 会尝试将计算放到 GPU 上，GPU 将执行张量运算，运算结果将存储在 GPU 内存中（而不是 data 指针中）

参考之前的 self-attention 计算图，假设 Q，K，V 是固定的张量，整个计算可以放到 GPU 上执行，如下图所示：

这一过程首先是将 Q，K，V 复制到 GPU 内存中，然后，CPU 逐个张量驱动计算，但实际的数学运算会放到 GPU 上执行。当计算图中的最后一个操作结束时，结果张量的数据将从 GPU 内存复制回 CPU 内存

Note ：在实际的 Transformer 中，Q，K，V 并非固定不变，QKV 也不是最终的输出，这个我们稍后再详细解释

有了对张量的理解之后，我们就可以回到 LLaMA 的流程

3. Tokenization

推理的第一步是 tokenization（分词），它是将 prompt 提示词拆分成一系列较短的字符串（称为 token）的过程。这些 token 必须是模型词汇表中的一部分，也就是 LLM 所训练的 token 列表。例如，LLaMA 的词汇表由 32k 个 token 组成，并作为模型的一部分发布

对于我们的 prompt 示例，tokenization 会将这个 prompt 分成 11 个 token（空格用特殊元符号 '__'（U+2581） 代替）

|Quant|um|▁mechan|ics|▁is|▁a|▁fundamental|▁theory|▁in|▁physics|▁that|

Note ：U+2581 是一个 Unicode 编码点，代表字符 '__'。U+2581 中的 U+ 前缀是标准的 Unicode 表示法，后面的 2581 是该字符的编码

在分词器方面，LLaMA 使用的是都带有字节对编码（byte-pair-encoding，BPE）算法的 SentencePiece tokenizer。这种分词器非常有意思，因为它是基于子词的，这意味着单词可以由多个 token 来表示。例如，在我们的 prompt 中，"Quantum" 被拆分成了 "Quant" 和 "um"

在训练过程中，当词汇表被导出时，BPE 算法会确保常用词以单个 token 的形式包含在词汇表中，而罕见词则被分解成子词。在上面的例子中，"Quantum" 一词不在词汇表中，但 "Quant" 和 "um" 作为两个单独的 token 出现在词汇表中。此外，空格不会被特殊处理，如果足够常见，空格会作为元字符包含在 token 中

基于子词的分词方法之所以强大，有以下几个原因：

• 基于子词的分词方法通过将常见的后缀和前缀表示为单独的 token，使得 LLM 可以学习 "Quantum" 等罕见词汇的含义，同时保持相对较小的词汇量
• 它可以学习特定语言的特征，而无需采用特定语言的分词方案。引用 BPE 编码论文中的例子：考虑德语的复合词如 Abwasser|behandlungs|anlange（污水处理厂），分段的、可变长度的表示形式比将该词编码为固定长度的向量更加直观。
• 同样，它在解析代码时也很有用。例如，一个 model_size 的变量将被分词为 model|_|size，这样 LLM 就能 "理解" 该变量的用途（这也是为变量赋予有意义名称的另一个原因！）

在 llama.cpp 中，分词是通过 llama_tokenize() 函数执行的，该函数将 prompt 提示词字符串作为输入，并返回一个 token 列表，其中每个 token 用一个整数表示：

// llama.h
typedef int llama_token;

// common.h
std::vector<llama_token> llama_tokenize(
        struct llama_context * ctx,
        // the prompt
        const std::string & text,
        bool   add_bos);

分词过程首先将 prompt 提示词分解为单字符 token，然后，它会反复尝试将两个相应的 token 合并成一个更大的 token，只要合并后的 token 在词汇表中。这样可以确保生成的 token 尽可能大，对于我们的 prompt 例子，分词步骤如下：

Q|u|a|n|t|u|m|▁|m|e|c|h|a|n|i|c|s|▁|i|s|▁a|▁|f|u|n|d|a|m|e|n|t|a|l|

Qu|an|t|um|▁m|e|ch|an|ic|s|▁|is|▁a|▁f|u|nd|am|en|t|al|

Qu|ant|um|▁me|chan|ics|▁is|▁a|▁f|und|am|ent|al|

Quant|um|▁mechan|ics|▁is|▁a|▁fund|ament|al|

Quant|um|▁mechan|ics|▁is|▁a|▁fund|amental|

Quant|um|▁mechan|ics|▁is|▁a|▁fundamental|

值得注意的是，每个中间步骤都包括根据模型词汇表进行有效的分词，不过，只有最后一步才会作为 LLM 的输入

至此，流程中的 prompt 到 tokens 的转换讲解完成

4. Embeddings

分词得到的这些 token 将作为 LLaMA 的输入，用于预测下一个 token，这里的核心函数是 llm_build_llama()：

// llama.cpp (simplified)
static struct ggml_cgraph * llm_build_llama(
         llama_context & lctx,
     const llama_token * tokens,
                   int   n_tokens,
                   int   n_past);

该函数由 tokens 和 n_tokens 参数表示的 token 列表作为输入，然后，它将构建 LLaMA 的完整张量计算图，并以 struct ggml_cgraph 的形式返回。这个阶段实际上不会进行计算，n_past 参数设置为 0，暂时可以忽略，我们将在后面讨论 kv cache 时再来讨论它

除了 token，该函数还使用了模型权重参数，这些权重是 LLM 在训练过程中学习到的固定 tensor，是模型的一部分。这些模型参数会在推理开始前预加载到 lctx 中

现在，我们来开始探索计算图结构，计算图的第一部分涉及将 token 转换为 embedding。embedding 是每个 token 的固定向量表示，它比纯整数更适合深度学习，因为它能捕捉词语的语义。该向量的大小就是模型维度，不同模型的 embedding 维度不同，例如，在 LLaMA-7B 中，嵌入维度为 n_embed=4096

模型参数包括一个将 token 转换为 embedding 的嵌入矩阵，由于我们的词汇量 n_vocab=32000，因此这是一个 32000 x 4096 的矩阵，每行表示一个 token 的嵌入向量，如下图所示：

计算图的第一部分是从嵌入矩阵中提取每个 token 的相关行：

// llama.cpp (simplified)
static struct ggml_cgraph * llm_build_llama(/* ... */) {
    // ...

    struct ggml_tensor * inp_tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
    memcpy(
        inp_tokens->data,
        tokens,
        n_tokens * ggml_element_size(inp_tokens));

    inpL = ggml_get_rows(ctx0, model.tok_embeddings, inp_tokens);
}

代码首先创建了一个新的一维整数张量，称为 inp_tokens，用来保存数字 token。然后，将 token 复制到该张量的 data 指针中。最后，创建一个新的 GGML_OP_GET_ROWS 张量操作，将嵌入矩阵 model.tok_embeddings 与我们的 token 结合起来

如下图所示，这一操作将从嵌入矩阵中抽取行，创建一个新的 n_tokens x n_embed 矩阵，其中只包含按原始顺序排列的 token 的嵌入向量

至此，流程中的 tokens 到 embeddings 的转换讲解完成

5. Transformer

计算图的主要部分是 Transformer，Transformer 是一个神经网络架构，是 LLM 的核心，执行主要的逻辑推理。在下面的章节中，我们将从工程的角度探讨 transformer 的一些关键技术，重点是 self-attention 自注意力机制。如果大家想更详细的了解 Transformer 架构，可以参考：The Illustrated Transformer

5.1 Self-attention

我们先来看看 self-attention 是什么，再看看它如何与整个 Transformer 架构相匹配

self-attention 是一种机制，它采用 token 序列并生成该序列的紧凑向量表示，同时考虑了 token 之间的关系。它是 LLM 架构中唯一计算 token 之间关系的地方，因此它构成了语言理解的核心，涵盖了对词汇关系的理解。由于它涉及跨 token 计算，因此从工程角度来看，它也是最有趣的地方，因为计算量可能会相当大，尤其是对于一些较长的序列

self-attention 的输入是 n_tokens x n_embd 的嵌入矩阵，每一行或向量代表一个独立的 token，然后，每个向量都被转换成三个不同的向量，分别称为查询（query）、键（key）、值（value）向量。这种转换是通过将每个 token 的嵌入向量与固定的 wq，wk，wv 矩阵相乘来实现的，如下图所示，这些矩阵是模型参数的一部分

每个 token 都要重复这一过程，即 n_tokens 次。从理论上讲，这个过程可以在一个循环中完成，但为了提高效率，所有行的转换都是通过矩阵乘法一次操作来完成的。相关代码如下：

// llama.cpp (simplified to remove use of cache)

// `cur` contains the input to the self-attention mechanism
struct ggml_tensor * K = ggml_mul_mat(ctx0,
    model.layers[il].wk, cur);
struct ggml_tensor * Q = ggml_mul_mat(ctx0,
    model.layers[il].wq, cur);
struct ggml_tensor * V = ggml_mul_mat(ctx0,
    model.layers[il].wv, cur);

我们最终得到 K，Q，V 三个矩阵，其大小也是 n_tokens x n_embd，每个 token 的 key、query 和 value 向量堆叠在一起

self-attention 的下一步是将查询矩阵 Q 与键矩阵 K 的转置相乘，这一运算主要是计算每一对查询向量和键向量的联合得分，我们将使用 S(i,j) 来表示查询 i 和键 j 的得分

这个过程会产生 n_tokens^2 个得分，每个 query-key 对一个分数，放在一个名为 QK^T 的矩阵中，随后会对该矩阵进行 mask 掩码，以去除对角线上方的元素，如下图所示：

Note ：通过将 Q 与 K 的转置相乘，计算出每个 query-key 对的联合得分 S(i,j)，上图显示的是前四个 token 的结果，以及每个得分所对应的 token。mask 掩码步骤可确保只保留当前 token 与其前面 token 之间的分数。为简单起见，这里省略了中间的 scaling 缩放操作

mask 掩码操作是一个关键步骤，对于每个 token，它只保留其前面 token 的分数。在训练阶段，这一限制确保 LLM 只根据过去的 token 而不是未来的 token 来学习预测 token。此外，它还允许在预测未来 token 时进行重大优化，稍后我们再来详细探讨

self-attention 的最后一步是将掩码得分 QK^T_masked 与之前的值向量（V）相乘。这样的矩阵乘法运算将前面所有 token 的值向量加权相加，其中的权重就是得分 S(i,j)

例如，对于第四个 token ics，它会生成 Quant、um、_mechan 和 ics 的这几个 token 的值向量的加权和，权重为 S(3,0) 至 S(3,3)，这些得分是根据 ics 的查询向量和之前所有 token 的键向量计算出来的，如下图所示：

Note ：QK^TV 矩阵包含值向量的加权和，例如，上图高亮显示的最后一行是前四个值向量的加权和，权重就是高亮显示的分数

QK^TV 矩阵标志着自注意力机制的结束，之前在张量计算中已经介绍过实现自注意力的相关代码，现在我们应该可以更好的来理解它

5.2 Transformer的各层

self-attention 是 transformer 各层的组成部分之一，除了 self-attention 外，每一层还包含其他多种张量运算，主要是矩阵加法、乘法和激活，这些都是前馈神经网络（feed-forward network, ffn）的一部分。我们将不再详细探讨这些操作，只需要注意以下内容：

• 前馈网络中使用了较大的固定参数矩阵，在 LLaMA-7B 中，它们的大小为 n_embd x n_ff = 4096 x 11008
• 除了自注意力外，所有其他操作都可以认为是逐行或逐个 token 进行的，只有自注意力包含跨 token 计算，这一点在下面讨论 kv cache 时非常重要
• 输入和输出的大小总是 n_tokens x n_embd，每个 token 各一行，每行的大小与模型的维度相同

下面是 LLaMA-7B 中单个 Transformer 架构的示意图：

Note：未来的模型架构可能会略有不同

LLaMA-7B 中的 Transformer 的完整计算图如上所示，它包含自注意力层和前馈网络层，每一层的输出都是下一层的输入。自注意力阶段和前馈阶段都使用了大量的参数矩阵，这些参数构成了该模型 70 亿各参数中的大部分

在 Transformer 架构中有多个层，例如，在 LLaMA-7B 中有 n_layers=32 层，这些层完全相同，只是每个层都有自己的一组参数矩阵（例如，用于自注意力机制的 wq，wk，wv 矩阵）。第一层的输入是上述嵌入的矩阵，第一层的输出作为第二层的输入，以此类推。我们可以认为，每一层都会产生一个嵌入列表，但每个嵌入不再直接与单个 token 相关联，而是与某种更复杂的 token 关系理解相关联

5.3 计算logits

Transformer 的最后一步是计算 logits，logit 是一个浮点数，表示某个 token 是下一个 "正确" token 的概率，其数值越高，相应 token 是 "正确" token 的可能性就越大

logits 的计算方法是将最后一层 Transformer 的输出与固定的 n_embd x n_vocab 参数矩阵（在 llama.cpp 中也称为 output）相乘，如下图所示，这一运算的结果是我们词汇表中每个 token 的 logit，例如，在 LLaMA 中，它的结果是 n_vocab=32000 个 logits

Note ：Transformer 的最后一步是将最后一层的输出与一个固定参数矩阵（也称为 output）相乘，从而计算处 logits，从上图来看只有结果的最后一行（高亮显示）才值得关注，它包含词汇表中每个可能的下一个 token 的 logits

logits 是 Transformer 的输出结果，它告诉我们最有可能的下一个 token 是什么。至此，所有张量计算结束，下面这个经过简化和注释的 llm_build_llama() 函数总结了本节中描述的所有步骤：

// llama.cpp (simplified and commented)

static struct ggml_cgraph * llm_build_llama(
         llama_context & lctx,
     const llama_token * tokens,
                   int   n_tokens,
                   int   n_past) {
    ggml_cgraph * gf = ggml_new_graph(ctx0);
    struct ggml_tensor * cur;
    struct ggml_tensor * inpL;

    // Create a tensor to hold the tokens.
    struct ggml_tensor * inp_tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, N);
    // Copy the tokens into the tensor
    memcpy(
        inp_tokens->data,
        tokens,
        n_tokens * ggml_element_size(inp_tokens));

    // Create the embedding matrix.
    inpL = ggml_get_rows(ctx0,
        model.tok_embeddings,
        inp_tokens);

    // Iteratively apply all layers.
    for (int il = 0; il < n_layer; ++il) {
        struct ggml_tensor * K = ggml_mul_mat(ctx0, model.layers[il].wk, cur);
        struct ggml_tensor * Q = ggml_mul_mat(ctx0, model.layers[il].wq, cur);
        struct ggml_tensor * V = ggml_mul_mat(ctx0, model.layers[il].wv, cur);

        struct ggml_tensor * KQ = ggml_mul_mat(ctx0, K, Q);
        struct ggml_tensor * KQ_scaled = ggml_scale_inplace(ctx0, KQ, KQ_scale);
        struct ggml_tensor * KQ_masked = ggml_diag_mask_inf_inplace(ctx0,
            KQ_scaled, n_past);
        struct ggml_tensor * KQ_soft_max = ggml_soft_max_inplace(ctx0, KQ_masked);
        struct ggml_tensor * KQV = ggml_mul_mat(ctx0, V, KQ_soft_max);

        // Run feed-forward network.
        // Produces `cur`.
        // ...

        // input for next layer
        inpL = cur;
    }

    cur = inpL;

    // Calculate logits from last layer's output.
    cur = ggml_mul_mat(ctx0, model.output, cur);

    // Build and return the computation graph.
    ggml_build_forward_expand(gf, cur);
    return gf;
}

要执行实际推理，需要使用 ggml_graph_compute() 计算该函数返回的计算图，然后将 logits 从最后一个张量的 data 指针复制到浮点数组中，为下一步的采样做好准备

至此，流程中的 embeddings 到 intermediate matrices 再到 logits 的转换讲解完成

6. Sampling

有了 logits 列表，下一步就是根据该列表选择下一个 token，这个过程称为 采样（Sampling）。这里有多种采样方法可供选择，适用于不同的情况，在本节中，我们将介绍两种基本的采样方法，grammar sampling 语法采样等更高级的采样方法将留在以后的文章中介绍

6.1 Greedy sampling

贪婪采样（Greedy sampling）是一种直接的方法，它选择与之相关的 logits 最高的 token

在我们的 prompt 例子中，以下 token 的 logits 最高：

token	logits
__describes	18.990
__provides	17.781
__explains	17.403
__de	16.361
__gives	15.007

因此，贪婪采样将确定性地选择 __describes 作为下一个 token，当重新评估相同 prompt 需要确定性输出时，贪婪采样最有用

6.2 Temperature sampling

温度采样是概率性的，这意味着同一 prompt 在重新评估时可能会产生不同的输出结果。它使用一个名为 temperature 的参数，该参数是一个介于 0~1 之间的浮点数值，会影响结果的随机性，过程如下：

• 1. 将 logits 从高到低排序，并使用 softmax 函数进行归一化处理，以确保它们的总和为 1
• 2. 应用阈值（默认 0.95），只保留累积概率低于阈值的 token，这一步骤可有效去除低概率的 token，防止 "坏的" 或 "不正确的" token 很少被采样
• 3. 将剩余 token 的 logits 除以温度参数并再次归一化，使其总和为 1 并代表概率
• 4. token 会根据这些概率随机抽样，例如，在我们的 prompt 中，token __describes 的概率为 p=0.6，这意味着大约有 60% 的概率会选择它，一旦重新评估时，可能会选择不同的 token

步骤 3 中的温度参数可增加或减少随机性，较低的温度值会抑制较低概率的 token，使相同的 token 更有可能在重新评估时被选中，因此，较低的温度值会降低随机性。相反，温度值越高，概率分布越 "扁平化"，越强调概率较低的 token，这就增加了每次重新评估产生不同 token 的可能性，从而增加了随机性

下图展示了不同温度下的下一个 token 概率：

我们的示例 prompt 经过归一化后的下一个 token 概率与温度相关，温度越低，低概率的 token 就越少，而温度越高，低概率的 token 就越多，temp=0 时与贪婪采样基本相同

现在我们来通过一个例子来说明，假设我们示例 prompt 经过 LLM 后输出的下一个 token 和对应的 logits 如之前的表格所示，温度采样（temp=0.8）的流程如下：

步骤 1： softmax 转换

为了计算概率，同时保证数值稳定性，我们使用 safe softmax，其公式如下：

s o f t m a x ( { x 1 , ... , x N } ) = { e x i ∑ j = 1 N e x j } i = 1 N \mathrm{softmax}(\{x_1,\ldots,x_N\})=\left\{\frac{e^{x_i}}{\sum_{j=1}^Ne^{x_j}}\right\}_{i=1}^N softmax({x1,...,xN})={∑j=1Nexjexi}i=1N

s a f e s o f t m a x ( { x 1 , ... , x N } ) = { e x i − m ∑ j = 1 N e x j − m } i = 1 N \mathrm{safe\ softmax}(\{x_1,\ldots,x_N\})=\left\{\frac{e^{x_{i}-m}}{\sum_{j=1}^{N}e^{x_{j}-m}}\right\}_{i=1}^N safe softmax({x1,...,xN})={∑j=1Nexj−mexi−m}i=1N

其中 m = max ⁡ j = 1 N ( x j ) m=\max_{j=1}^{N}(x_j) m=maxj=1N(xj)

在当前示例中最大值 m = 18.990 m=18.990 m=18.990，每个 token 的差值为：

• __describes： x 0 − m = 18.990 − 18.990 = 0 x_{0}-m=18.990-18.990=0 x0−m=18.990−18.990=0， e x 0 − m = e 0 = 1 e^{x_{0}-m}=e^{0}=1 ex0−m=e0=1
• __provides： x 1 − m = 17.781 − 18.990 = − 1.209 x_{1}-m=17.781-18.990=-1.209 x1−m=17.781−18.990=−1.209， e x 1 − m = e − 1.209 ≈ 0.298 e^{x_{1}-m}=e^{-1.209} \approx 0.298 ex1−m=e−1.209≈0.298
• __explains： x 2 − m = 17.403 − 18.990 = − 1.587 x_{2}-m=17.403-18.990=-1.587 x2−m=17.403−18.990=−1.587， e x 2 − m = e − 1.587 ≈ 0.204 e^{x_{2}-m}=e^{-1.587} \approx 0.204 ex2−m=e−1.587≈0.204
• __de： x 3 − m = 16.361 − 18.990 = − 2.629 x_{3}-m=16.361-18.990=-2.629 x3−m=16.361−18.990=−2.629， e x 3 − m = e − 2.629 ≈ 0.072 e^{x_{3}-m}=e^{-2.629} \approx 0.072 ex3−m=e−2.629≈0.072
• __gives： x 4 − m = 15.007 − 18.990 = − 3.983 x_{4}-m=15.007-18.990=-3.983 x4−m=15.007−18.990=−3.983， e x 4 − m = e − 3.983 ≈ 0.019 e^{x_{4}-m}=e^{-3.983} \approx 0.019 ex4−m=e−3.983≈0.019

总和 ∑ j = 1 N e x j − m ≈ 1 + 0.298 + 0.204 + 0.072 + 0.019 ≈ 1.593 {\sum_{j=1}^{N}e^{x_{j}-m}}\approx 1+0.298+0.204+0.072+0.019 \approx 1.593 ∑j=1Nexj−m≈1+0.298+0.204+0.072+0.019≈1.593

因此各 token 的初步概率为：

• __describes： 1 1.593 ≈ 0.628 \frac{1}{1.593} \approx 0.628 1.5931≈0.628
• __provides： 0.298 1.593 ≈ 0.187 \frac{0.298}{1.593} \approx 0.187 1.5930.298≈0.187
• __explains： 0.204 1.593 ≈ 0.128 \frac{0.204}{1.593} \approx 0.128 1.5930.204≈0.128
• __de： 0.072 1.593 ≈ 0.045 \frac{0.072}{1.593} \approx 0.045 1.5930.072≈0.045
• __gives： 0.019 1.593 ≈ 0.012 \frac{0.019}{1.593} \approx 0.012 1.5930.019≈0.012

步骤 2：应用 Top-p 阈值

按照概率从大到小累加：

• __describes：0.628（累积 0.628）
• __provides：0.187（累积 0.628 + 0.187 = 0.815）
• __explains：0.128（累积 0.815 + 0.128 = 0.943）
• __de：0.045（累积 0.943 + 0.045 = 0.988，超过 0.95）
• __gives：不考虑

根据设定的阈值（0.95），只保留累积概率在 0.95 以下的 token，即：

• 保留的 token ：__describes、__provides、__explains
• 剔除的 token ：__de、__gives

步骤 3：温度调节与重新归一化

对于保留的 token，我们取它们原来的 logits 重新进行温度调节，假设温度参数 T = 0.8 T=0.8 T=0.8，计算新的 logits：

n e w _ l o g i t i = l o g i t i T \mathrm{new\_{logit}}_i = \frac{\mathrm{logit}_i}{T} new_logiti=Tlogiti

计算得到：

• __describes： 18.990 0.8 = 23.7375 \frac{18.990}{0.8}=23.7375 0.818.990=23.7375
• __provides： 17.781 0.8 = 22.22625 \frac{17.781}{0.8}=22.22625 0.817.781=22.22625
• __explains： 17.403 0.8 = 21.75375 \frac{17.403}{0.8}=21.75375 0.817.403=21.75375

继续使用 safe softmax 进行归一化（此时最大值 m = 23.7375 m=23.7375 m=23.7375），得到：

• __describes： x 0 − m = 23.7375 − 23.7375 = 0 x_{0}-m=23.7375-23.7375=0 x0−m=23.7375−23.7375=0， e x 0 − m = e 0 = 1 e^{x_{0}-m}=e^{0}=1 ex0−m=e0=1
• __provides： x 1 − m = 22.22625 − 23.7375 = − 1.51125 x_{1}-m=22.22625-23.7375=-1.51125 x1−m=22.22625−23.7375=−1.51125， e x 1 − m = e − 1.51125 ≈ 0.220 e^{x_{1}-m}=e^{-1.51125} \approx 0.220 ex1−m=e−1.51125≈0.220
• __explains： x 2 − m = 21.75375 − 23.7375 = − 1.98375 x_{2}-m=21.75375-23.7375=-1.98375 x2−m=21.75375−23.7375=−1.98375， e x 2 − m = e − 1.98375 ≈ 0.137 e^{x_{2}-m}=e^{-1.98375} \approx 0.137 ex2−m=e−1.98375≈0.137

总和 ∑ j = 1 N e x j − m ≈ 1 + 0.220 + 0.137 ≈ 1.357 {\sum_{j=1}^{N}e^{x_{j}-m}}\approx 1+0.220+0.137 \approx 1.357 ∑j=1Nexj−m≈1+0.220+0.137≈1.357

因此，各 token 的最终采样概率为：

• __describes： 1 1.357 ≈ 0.737 \frac{1}{1.357} \approx 0.737 1.3571≈0.737
• __provides： 0.220 1.357 ≈ 0.162 \frac{0.220}{1.357} \approx 0.162 1.3570.220≈0.162
• __explains： 0.137 1.357 ≈ 0.101 \frac{0.137}{1.357} \approx 0.101 1.3570.137≈0.101

温度参数的影响：

• 低温：放大 logits 差异，生成结果更确定，低概率 token 的机会更少
• 高温：减小 logits 差异，使分布更平坦，低概率 token 的采样概率提高，从而增加随机性

步骤 4：最终采样

根据步骤 3 得到的最终概率分布，通过随机抽样选择一个 token

总的来说，温度采样结合了两个策略：

• Top-p（核采样）：用于剔除那些累积概率超过一定阈值的低概率 token，从而保证生成内容不会因为极低概率 token 而出现不合理情况
• 温度缩放：调节了生成的随机性，低温使输出更确定（更 "保守"），高温则引入更多随机性（更 "发散"）

对于步骤 4 中的随机抽样博主有些困惑，怎么确保采样概率越高的 token 最终选择的机会越大呢？🤔

随机抽样通常使用一种叫做逆变换抽样 （Inverse Transform Sampling ）的方法来完成，假设我们已经得到了每个 token 的最终采样概率 p 1 , p 2 , ... , p n p_1,p_2,\ldots,p_n p1,p2,...,pn，其总和为 1，具体抽样步骤如下：

1. 计算累积概率分布

c i = ∑ j = 1 i p i c_i= \sum_{j=1}^{i}p_{i} ci=j=1∑ipi

例如，如果有三个 token，其累计概率可能为：

• token 1： c 1 = p 1 c_1=p_1 c1=p1
• token 2： c 2 = p 1 + p 2 c_2=p_1+p_2 c2=p1+p2
• token 3： c 3 = p 1 + p 2 + p 3 c_3=p_1+p_2+p_3 c3=p1+p2+p3

2. 生成随机数

从均匀分布 [0,1) 中生成一个随机数 r r r

3. 选择 Token

找到第一个累计概率 c i c_i ci 满足 r ≤ c i r \leq c_i r≤ci 的 token，然后选择这个 token。这样一来，概率越高的 token 占据的累计区间越大，被选中的可能性也就越高

这种方法确保了如果某个 token 的概率 p i p_i pi 较高，它在累计分布中所占的区间就更大，因此生成的随机数 r r r 落在这个区间内的概率也更大，从而更有可能被采样到

对 token 采样完之后就实现了 LLM 的一次完整迭代，在对初始 token 采样后，它将被添加到 token 列表中，整个过程再次运行。之后输出会不断成为 LLM 的输入，每次迭代增加一个 token

至此，流程中的 logits 到 output 的转换讲解完成

理论上，之后的迭代完全可以按照相同的步骤进行，不过，为了解决随着 token 列表的增加而出现的性能下降问题，我们需要采用某些优化方法。接下来我们就来介绍这些优化措施

7. 推理优化

随着 LLM 输入 token 列表的增加，Transformer 的 self-attention 操作可能会成为性能瓶颈。更长的 token 列表意味着需要将更大的矩阵相乘。每个矩阵乘法都是由许多较小的数字运算（即浮点运算）组成，这些运算受到 GPU 每秒浮点运算容量（flops）的限制。在 Transformer 推理运算中，根据计算，对于 52B 参数模型在 A100 GPU 上，由于过多的浮点运算，性能在 208 个 token 时开始下降。为了解决这一瓶颈，最常用的优化技术是 kv cache

7.1 KV cache

回顾一下，每个 token 都有一个相关的嵌入向量，通过与参数矩阵 wk 和 wv 相乘，进一步转换为 key 向量和 value 向量。kv cache 就是这些 key 向量和 value 向量的缓存，通过缓存，我们可以节省每次迭代时重新计算所需的浮点运算

kv cache 的工作原理如下：

• 如前所述，在初始迭代过程中，会计算所有 token 的 key 和 value 向量，然后保存到 kv cache 中
• 在之后的迭代中，我们只需要计算最新 token 的 key 和 value 向量，缓存的 k-v 向量和新 token 的 k-v 向量被拼接在一起，形成 K，V 矩阵。这样就省去了重新计算之前所有 token 的 k-v 向量的过程，这点非常重要

Note ：在随后的迭代中，只计算最新 token 的 key 向量即可，其余的从缓存中提取，共同组成 K 矩阵，此外，新计算的 key 向量也会保存到缓存中，同样的过程也适用于 value 向量

我们之所以能够利用缓存来处理 key 和 value 向量，主要是因为这些向量在两次迭代之间保持相同。例如，如果我们先处理前四个 token，然后处理第五个 token，最初的四个 token 保持不变，那么前四个 token 的 key 和 value 向量在第一次迭代和第二次迭代之间将保持相同，因此，我们其实根本不需要在第二次迭代中重新计算前四个 token 的 key 和 value 向量

这一原则适用于 Transformer 中的所有层，而不仅仅是第一层，在所有层中，每个 token 的 key 和 value 向量仅依赖于先前的 token。因此，在后续迭代中添加新 token 时，现有 token 的 key 和 value 向量将保持不变

对于第一层，这一概念的验证相对简单：token 的 key 向量是通过将 token 的固定嵌入向量与固定的 wk 参数矩阵相乘来确定的。因此，在后续的迭代中，无论是否引入额外的 token，它都不会发生变化，同样的道理也适用于 value 向量。

对于第二层及之后的层，这一原则就不那么明显了，但仍然适用。要理解原因，请看第一层的 QKV 矩阵，它是 self-attention 步骤的输出，QKV 矩阵中每一行都是一个加权和，它取决于：

• 之前 token 的 value 向量
• 根据之前 token 的 key 向量计算出的分数

因此，QKV 中的每一行都完全依赖之前的 token，这个矩阵在经过一些额外的基于行的运算后（softmax），将作为第二层的输入，这意味着第二层的输入在未来的迭代中将保持不变，只是增加了新的行。归纳起来，同样的逻辑也适用于其他层

Note ：再来看看 QKV 矩阵是如何计算的，第三行（高亮显示的）仅根据第三个查询向量 '_mechan' 和前三个 key 和 value 向量（也是高亮显示）确定，随后的 token 不会对其产生影响，因此，它在未来的迭代中将保持不变

7.2 持续优化

你可能会问，既然我们缓存了 key 和 value 向量，为什么不同时缓存 query 向量呢？答案是，事实上，除了当前 token 的 query 向量外，之前 token 的 query 向量在后续迭代中都是不必要的

有了 kv cache 之后，我们实际上就可以只向 self-attention 中提供最新 token 的 query 向量，该 query 向量与缓存的 K 矩阵相乘，只计算 QKV 矩阵的最新一行。事实上，在所有层中，我们现在传递的是 1xn_embd 大小的向量，而不是第一次迭代中计算的 n_token x n_embd 矩阵，为了说明这一点，我们可以将下图中的后一次迭代与前一次迭代进行比较：

Note ：在上图的例子中，第一次迭代有四个 token，第二次迭代增加了第五个 token 即 _is，最新 token 的 key、query 和 value 向量以及缓存的 key 和 value 向量被用来计算 QKV 的最后一行，这就是预测下一个 token 所需要的全部内容

这一过程在所有层中重复进行，利用每一层的 kv cahce。因此，在这种情况下，Transformer 的输出是由一个 n_vocab logits 组成的单一向量，用来预测下一个 token

通过这种优化，我们节省了在 QK 和 QKV 中计算不必要的浮点运算，随着 token 列表的增大，这种运算会变得非常重要

7.3 KV cache的实际应用

我们可以来深入研究下 llama.cpp 的代码，看看 kv cache 是如何实现的：

// llama.cpp (simplified)

struct llama_kv_cache {
    // cache of key vectors
    struct ggml_tensor * k = NULL;

    // cache of value vectors
    struct ggml_tensor * v = NULL;

    int n; // number of tokens currently in the cache
};

我们可以看到它是使用 tensor 来构建的，一个用于 key 向量，一个用于 value 向量。缓存初始化时，会为每一层分配足够的空间来保存 512 个 key 和 value 向量：

// llama.cpp (simplified)
// n_ctx = 512 by default
static bool llama_kv_cache_init(
    struct llama_kv_cache & cache,
    ggml_type   wtype,
    int   n_ctx) {
    // Allocate enough elements to hold n_ctx vectors for each layer.
    const int64_t n_elements = n_embd*n_layer*n_ctx;

    cache.k = ggml_new_tensor_1d(cache.ctx, wtype, n_elements);
    cache.v = ggml_new_tensor_1d(cache.ctx, wtype, n_elements);

    // ...
}

回想一下，在推理过程中，计算图是使用 llm_build_llama() 这个函数来构建的，这个函数有一个名为 n_past 的参数，我们之前忽略了这个参数。在第一次迭代中，n_tokens 参数表示 token 的数据，而 n_past 被设置为 0。在随后的迭代中，n_tokens 被设置为 1，因为只有最新的 token 才会被处理，而 n_past 则包含过去 token 的数量

下面我们展示了该函数中利用缓存计算 K 矩阵的相关部分代码，代码稍作了简化，忽略了多头注意力，并为每个步骤添加了注释：

// llama.cpp (simplified and commented)

static struct ggml_cgraph * llm_build_llama(
    llama_context & lctx,
    const llama_token * tokens,
    int   n_tokens,
    int   n_past) {
    // ...

    // Iteratively apply all layers.
    for (int il = 0; il < n_layer; ++il) {
         // Compute the key vector of the latest token.
         struct ggml_tensor * Kcur = ggml_mul_mat(ctx0, model.layers[il].wk, cur);
         // Build a view of size n_embd into an empty slot in the cache.
         struct ggml_tensor * k = ggml_view_1d(
            ctx0,
            kv_cache.k,
            // size
            n_tokens*n_embd,
            // offset
            (ggml_element_size(kv_cache.k)*n_embd) * (il*n_ctx + n_past)
         );

         // Copy latest token's k vector into the empty cache slot.
         ggml_cpy(ctx0, Kcur, k);

         // Form the K matrix by taking a view of the cache.
         struct ggml_tensor * K =
             ggml_view_2d(ctx0,
                 kv_self.k,
                 // row size
                 n_embd,
                 // number of rows
                 n_past + n_tokens,
                 // stride
                 ggml_element_size(kv_self.k) * n_embd,
                 // cache offset
                 ggml_element_size(kv_self.k) * n_embd * n_ctx * il);
    }
}

首先，计算出新的 key 向量，然后，用 n_past 在缓存中找到下一个 xxx，并把新的 key 向量复制到那里。最后，通过在缓存中提取具有正确 token 数（n_past + n_tokens）的视图，形成矩阵 K

kv cache 是 LLM 推理优化的基础，值得注意的是，llama.cpp 中实现的版本并不是最优的（截至原文撰写时），例如，它提前分配了大量内存，以保存所支持的最大 key 和 value 向量数（本例中为 512）。更先进的实现方法如 vLLM 旨在提高内存使用效率，并可能进一步改善性能，有机会的话讲留待今后的文章中介绍

OK，以上就是这篇文章的全部内容了

结语

这篇文章主要介绍了 LLM 的推理流程，整个推理过程梳理得非常清楚，从 prompt➡tokens➡embeddings➡transformer➡logits➡sampling➡output，每个过程都讲解得非常详细

输入的文本提示词经过分词器拆分成一个个 token，接着将 token 转换为嵌入向量作为 transformer 的输入，经过 transformer 中 self-attention 等组件的处理后得到 logits，然后利用 logits 通过采样方法得到预测的下一个 token，这便是 LLM 推理的一次完整过程。接着我们只要将新生成的 token 添加到之前的 token 列表，然后重复这个过程即可

这是一篇非常不错的文章，大家感兴趣的可以看看🤗

参考

• Understanding how LLM inference works with llama.cpp
• 理解llama.cpp怎么完成大模型推理的
• https://chatgpt.com/
• https://github.com/ggml-org/ggml
• https://github.com/ggml-org/llama.cpp