vLLM入门

vLLM是伯克利大学组织开源的大语言模型高速推理框架，使用 PagedAttention 高效管理注意力键和值内存，支持连续批处理和快速模型执行，通过引入操作系统的虚拟内存分页思想，提高语言模型服务在实时场景下的吞吐与内存使用效率。

KV Cache

背景介绍

Transformer 模型是一种广泛应用于自然语言处理（NLP）任务的深度学习模型。它通过自注意力机制（Self-Attention）来理解和生成文本。Transformer 模型在生成文本时，通常采用自回归推理的方式。

自回归推理的特点

自回归推理指的是模型在生成文本时，每次只预测一个token（即一个基本的文本单位，通常是一个字或一个词）。具体过程如下：

1. 初始输入：模型接收一段初始的输入文本（历史输入 tokens）。
2. 预测下一个 token：模型根据当前的输入，预测下一个 token。
3. 更新输入：将预测出的 token 添加到历史输入中，形成新的输入序列。
4. 重复执行：使用更新后的输入序列，继续预测下一个 token。

这种方法会反复执行多次，每次只增加一个 token。这意味着在每一轮推理中，当前的输出 token 与之前的所有输入 tokens 拼接，作为下一轮的输入。

重复计算的问题

在自回归推理过程中，每一轮的输入序列与前一轮相比，只增加了一个新的 token，但大部分输入 tokens 是相同的。因此，前后两轮的输入只相差一个 token，这导致大量的计算是重复的。例如，每次都需要重新计算整个输入序列的表示（如 Wq(X), Wk(X), Wv(X)），这在计算上非常耗时，尤其是在处理长文本时。

KV Cache 技术的解决方案

为了避免这种重复计算，KV Cache（键值缓存）技术被引入。KV Cache 的核心思想是将模型在前一轮推理中计算得到的键（Key）和值（Value）向量保存下来，以便在后续的推理过程中复用这些向量，而无需重新计算。

具体流程如下：

1. 没有 KV Cache 的流程（Without KV Cache）：

   • 每次推理时，模型需要重新计算所有输入 tokens 的 Wq(X), Wk(X), Wv(X)。
   • 计算出全量的注意力（Attention）机制。

2. 使用 KV Cache 的流程（With KV Cache）：

   • 第一步：在第一次推理时，模型计算完整的注意力机制，并将生成的键值对（KV）保存到 KV_cache 中。
   • 第二步：在后续的推理中，模型只需计算新增加的 token 的 Q（查询）、K（键）、V（值）向量。
   • 联合使用缓存：将新计算的 KV 向量与之前缓存的 KV_cache 结合起来，形成完整的 QKV（查询、键、值）集合。
   • 计算注意力：使用联合后的 QKV 来计算注意力，从而得到下一个 token 的预测结果。

KV Cache 的优势

通过使用 KV Cache 技术，模型在每次推理时只需计算新增的 token 的相关向量，而无需重新计算之前所有的输入。这大大减少了计算量，提高了推理的效率，尤其在处理长文本或需要快速响应的应用场景中，效果尤为显著。

总结

这段话主要介绍了 Transformer 模型在自回归推理过程中存在的重复计算问题，以及 KV Cache 技术如何通过缓存键值向量来避免这些重复计算，从而提升推理效率。通过这种优化，模型能够更高效地处理连续生成的文本，减少计算资源的消耗。

希望这个解释能帮助你更好地理解原文内容！

注意力机制

缩放点积注意力的公式

缩放点积注意力的核心公式如下：

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

其中：

• $Q$（Query）：查询矩阵。用于查询信息，表示当前需要关注的内容。
• $K$（Key）：键矩阵。用于键入信息，表示输入序列中各个位置的特征。
• $V$（Value）：值矩阵。用于提供实际的信息，表示与输入序列中各个位置相关的内容。
• $d_k$：键向量的维度。
• $\text{softmax}$：归一化函数，将相似度转换为概率分布。

公式解析

3.1 输入矩阵的构建

在注意力机制中，输入序列首先通过三个不同的线性变换（即三个不同的权重矩阵）生成查询（$Q$）、键（$K$）和值（$V$）矩阵：

$Q=X{W}_Q$
$K=X{W}_K$
$V=X{W}_V$

其中：

• $X$：输入序列的表示，形状为 $(n,d_{model})$，$n$ 为序列长度，$d_{model}$ 为模型的隐藏维度。
• $W_Q,W_K,W_V$：线性变换的权重矩阵，形状分别为$(d_{model},d_k)$、$(d_{model},d_k)$ 和 $(d_{model},d_v)$。

3.2 点积计算相似度

查询矩阵 $Q$ 和键矩阵$K$通过点积操作计算相似度：

$$Q{K}^T$$

这一步的结果是一个形状为$(n,n)$ 的相似度矩阵，每个元素$(i,j)$ 表示输入序列中第 $i$ 个位置的查询与第$j$个位置的键的相似度。

3.3 缩放（Scaling）

为了防止在高维空间中点积结果过大，导致梯度消失或梯度爆炸的问题，引入缩放因子 $\sqrt{d_k}$：

$$\frac{Q{K}^T}{\sqrt{d_k}}$$

3.4 归一化（Softmax）

将缩放后的相似度通过 softmax 函数进行归一化，得到权重矩阵：

$$\text{softmax}\left(\frac{Q{K}^T}{\sqrt{d_k}}\right)$$

这一步将每一行的相似度值转换为概率分布，确保权重之和为1。

3.5 加权求和（Weighted Sum）

将归一化后的权重矩阵与值矩阵 ( V ) 相乘，得到最终的注意力输出：

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

输出矩阵的形状为 $ (n, d_v) $，每个位置的表示是根据输入序列中所有位置的值向量加权求和得到的。

缩放点积简单示例

理解缩放点积注意力（Scaled Dot-Product Attention）通过具体示例

缩放点积注意力（Scaled Dot-Product Attention）是Transformer模型中的核心机制之一。为了更好地理解这个公式，我们将通过一个简单的、具体的例子一步步演示其工作原理。

公式回顾

首先，让我们再次回顾缩放点积注意力的公式：

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

其中：

• $Q$ 是查询（Query）矩阵。
• $K$ 是键（Key）矩阵。
• $V$ 是值（Value）矩阵。
• $d_k$ 是键向量的维度。
• $\text{softmax}$是将数值转换为概率分布的函数。

具体示例

假设我们有一个非常简单的场景：

• 序列长度 $n=2$（即输入有两个词）
• 模型的隐藏维度 $ d_{model} = 4 $
• 键（Key）和查询（Query）的维度$ d_k = 2 $
• 值（Value）的维度 $d_v=2$

步骤 1：定义输入矩阵 ( X )

假设输入序列 ( X ) 有两个词，每个词用一个4维向量表示：

$$X=\left[\begin{array}{cccc}1& 0& 1& 0\\ 0& 1& 0& 1\end{array}\right]$$

步骤 2：定义权重矩阵$ W_Q, W_K, W_V $

为了简化计算，我们手动定义三个权重矩阵：

$$W_Q=\left[\begin{array}{cc}1& 0\\ 0& 1\\ 1& 0\\ 0& 1\end{array}\right],W_K=\left[\begin{array}{cc}1& 0\\ 0& 1\\ 0& 1\\ 1& 0\end{array}\right],W_V=\left[\begin{array}{cc}1& 0\\ 0& 1\\ 1& 0\\ 0& 1\end{array}\right]$$

步骤 3：计算 ( Q, K, V ) 矩阵

通过将输入 ( X ) 与权重矩阵相乘，得到查询 ( Q )、键 ( K ) 和值 ( V ) 矩阵。

$$Q=X{W}_Q=\left[\begin{array}{cc}1& 0\\ 0& 1\end{array}\right]$$

$$K=X{W}_K=\left[\begin{array}{cc}1& 0\\ 0& 1\end{array}\right]$$

$$V=X{W}_V=\left[\begin{array}{cc}1& 0\\ 0& 1\end{array}\right]$$

注：在这个简化的例子中，权重矩阵被设计成使得 ( Q, K, V ) 都等于 ( X ) 的前两列。

步骤 4：计算 $QK^T $

计算查询 ( Q ) 与键 ( K ) 的点积：

$$Q{K}^T=\left[\begin{array}{cc}1& 0\\ 0& 1\end{array}\right]{\left[\begin{array}{cc}1& 0\\ 0& 1\end{array}\right]}^T=\left[\begin{array}{cc}1& 0\\ 0& 1\end{array}\right]\left[\begin{array}{cc}1& 0\\ 0& 1\end{array}\right]=\left[\begin{array}{cc}1& 0\\ 0& 1\end{array}\right]$$

步骤 5：缩放

将点积结果除以$\sqrt{d_k}$ 进行缩放。这里 $d_k = 2 $，所以 $\sqrt{d_k}=\sqrt{2}\approx 1.414$。

$$\frac{Q{K}^T}{\sqrt{d_k}}=\frac{1}{1.414}\left[\begin{array}{cc}1& 0\\ 0& 1\end{array}\right]=\left[\begin{array}{cc}0.707& 0\\ 0& 0.707\end{array}\right]$$

步骤 6：应用 Softmax

对缩放后的点积结果应用 softmax 函数。softmax 是按行进行归一化，使每行的元素和为1。

$$\text{softmax}\left(\left[\begin{array}{cc}0.707& 0\\ 0& 0.707\end{array}\right]\right)=\left[\begin{array}{c}\text{softmax}(0.707,0)\\ \text{softmax}(0,0.707)\end{array}\right]$$

计算每个元素的 softmax：

• 第一行：
  $$\text{softmax}(0.707,0)=\left[\frac{e^{0.707}}{e^{0.707}+e^0},\frac{e^0}{e^{0.707}+e^0}\right]\approx \left[\frac{2.028}{2.028+1},\frac{1}{2.028+1}\right]\approx [0.669,0.331]$$

• 第二行：
  $$\text{softmax}(0,0.707)=\left[\frac{e^0}{e^0+e^{0.707}},\frac{e^{0.707}}{e^0+e^{0.707}}\right]\approx \left[\frac{1}{1+2.028},\frac{2.028}{1+2.028}\right]\approx [0.331,0.669]$$

所以，归一化后的权重矩阵为：

$$\text{softmax}\left(\frac{Q{K}^T}{\sqrt{d_k}}\right)\approx \left[\begin{array}{cc}0.669& 0.331\\ 0.331& 0.669\end{array}\right]$$

步骤 7：加权求和得到注意力输出

将归一化后的权重矩阵与值 ( V ) 矩阵相乘，得到最终的注意力输出：

$$\text{Attention}(Q,K,V)=\left[\begin{array}{cc}0.669& 0.331\\ 0.331& 0.669\end{array}\right]\left[\begin{array}{cc}1& 0\\ 0& 1\end{array}\right]=\left[\begin{array}{cc}0.669\times 1+0.331\times 0& 0.669\times 0+0.331\times 1\\ 0.331\times 1+0.669\times 0& 0.331\times 0+0.669\times 1\end{array}\right]=\left[\begin{array}{cc}0.669& 0.331\\ 0.331& 0.669\end{array}\right]$$

结果解释

每一行的输出表示输入序列中对应位置的新表示：

• 第一行：[0.669, 0.331] 表示第一个词的新表示是第一个值向量的 66.9% 加上第二个值向量的 33.1%。
• 第二行：[0.331, 0.669] 表示第二个词的新表示是第一个值向量的 33.1% 加上第二个值向量的 66.9%。

这意味着模型通过注意力机制，根据查询和键的相似度动态地调整每个值向量的权重，从而生成新的表示。

更直观的图示

为了更直观地理解，我们可以将上述步骤图示化：

1. 输入序列 ( X )：

$$\begin{array}{cc}\text{词1}& [1,0,1,0]\\ \text{词2}& [0,1,0,1]\end{array}$$

2. 计算 ( Q, K, V )：

$$Q=K=V=\left[\begin{array}{cc}1& 0\\ 0& 1\end{array}\right]$$

3. 计算 ( QK^T ) 并缩放：

$$Q{K}^T=\left[\begin{array}{cc}1& 0\\ 0& 1\end{array}\right],\frac{Q{K}^T}{\sqrt{2}}\approx \left[\begin{array}{cc}0.707& 0\\ 0& 0.707\end{array}\right]$$

4. 应用 Softmax：

$$\text{softmax}\left(\frac{Q{K}^T}{\sqrt{2}}\right)\approx \left[\begin{array}{cc}0.669& 0.331\\ 0.331& 0.669\end{array}\right]$$

5. 加权求和得到输出：

$$\text{Attention}(Q,K,V)\approx \left[\begin{array}{cc}0.669& 0.331\\ 0.331& 0.669\end{array}\right]$$

总结

通过这个具体的例子，我们可以看到缩放点积注意力是如何通过查询、键和值之间的相似度计算权重，并用这些权重对值向量进行加权求和，生成新的表示的。这种机制允许模型动态地关注输入序列中的不同部分，从而更好地理解和生成语言。

多头注意力机制

多头注意力机制（Multi-Head Attention）详解

在深入理解了 自注意力（Self-Attention） 机制后，接下来让我们探讨 多头注意力机制（Multi-Head Attention）。多头注意力机制是Transformer架构中的关键组成部分，它通过并行地应用多个注意力头，增强了模型捕捉不同子空间信息的能力，从而提升了模型的表达能力和性能。

1. 多头注意力机制的基本概念

多头注意力机制 通过同时使用多个不同的注意力头（Attention Heads）来处理信息。每个注意力头在不同的子空间中学习不同的表示，从而使模型能够从多个角度捕捉输入数据的相关信息。这种机制不仅丰富了模型的表达能力，还提高了其对复杂模式的捕捉能力。

2. 多头注意力机制的目的和优势

2.1 目的

• 捕捉多样化的关系：不同的注意力头可以学习到输入数据中不同类型的关系或特征。例如，一个头可能专注于捕捉语法关系，另一个头可能专注于语义关系。
• 增强表示能力：通过多个头的并行计算，多头注意力机制可以生成更加丰富和多样化的表示，提升模型的整体性能。

2.2 优势

• 并行处理：多个头可以并行计算，提高计算效率。
• 灵活性：每个头可以专注于不同的子空间，增强模型的灵活性和适应性。
• 稳定性：通过分散学习，多头注意力机制可以减少单个头可能出现的过拟合或偏差问题。

3. 多头注意力机制的工作原理

多头注意力机制的核心思想是将查询（Query）、键（Key）和值（Value）向量分别通过多个线性变换，形成多个不同的子空间表示，然后分别计算每个子空间中的注意力，最后将所有头的输出拼接并通过一个线性变换得到最终的输出。

3.1 公式表示

假设有 ( h ) 个注意力头，多头注意力机制的计算过程如下：

1. 线性变换：

   对输入的查询 ( Q )、键 ( K )、值 ( V ) 分别应用 ( h ) 个不同的线性变换，得到 ( h ) 个不同的子空间表示：

   $$Q_i=Q{W}_i^Q,K_i=K{W}_i^K,V_i=V{W}_i^V\text{for}i=1,2,\dots ,h$$

其中，$ W_i^Q, W_i^K, W_i^V $ 是每个头的线性变换矩阵。

2. 计算注意力：

   对每个头分别计算缩放点积注意力：

   $${\text{head}}_i=\text{Attention}(Q_i,K_i,V_i)=\text{softmax}\left(\frac{Q_i{K}_i^T}{\sqrt{d_k}}\right)V_i$$

   其中，$d_k $ 是每个头的键向量维度。

3. 拼接和线性变换：

   将所有头的输出拼接起来，并通过一个线性变换得到最终的输出：

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

   其中，$ W^O $ 是输出的线性变换矩阵。

3.2 示例解析

为了更直观地理解多头注意力机制，我们通过一个具体的例子进行解析。

示例参数

• 序列长度（Sequence Length）：$n = 2 $
• 模型隐藏维度（Model Dimension）：$d_{model} = 4 $
• 注意力头数（Number of Heads）：$h=2$
• 每个头的维度（Dimension per Head）：$d_k = d_v = 2 $

步骤 1：定义输入矩阵 ( Q, K, V )

假设我们有以下输入：

$$Q=K=V=\left[\begin{array}{cccc}1& 0& 1& 0\\ 0& 1& 0& 1\end{array}\right]$$

步骤 2：定义权重矩阵

为简化计算，我们手动定义权重矩阵：

$$W_i^Q=W_i^K=W_i^V=\left[\begin{array}{cccc}1& 0& 0& 0\\ 0& 1& 0& 0\end{array}\right]\text{for}i=1,2$$

$$W^O=\left[\begin{array}{cccc}1& 0& 1& 0\\ 0& 1& 0& 1\end{array}\right]$$

步骤 3：线性变换

对每个头分别应用线性变换：

$$Q_1=Q{W}_1^Q=\left[\begin{array}{cc}1& 0\\ 0& 1\end{array}\right]$$

$$K_1=K{W}_1^K=\left[\begin{array}{cc}1& 0\\ 0& 1\end{array}\right]$$

$$V_1=V{W}_1^V=\left[\begin{array}{cc}1& 0\\ 0& 1\end{array}\right]$$

同理：

$$Q_2=Q{W}_2^Q=\left[\begin{array}{cc}1& 0\\ 0& 1\end{array}\right]$$

$$K_2=K{W}_2^K=\left[\begin{array}{cc}1& 0\\ 0& 1\end{array}\right]$$

$$V_2=V{W}_2^V=\left[\begin{array}{cc}1& 0\\ 0& 1\end{array}\right]$$

步骤 4：计算每个头的注意力

对于每个头，计算缩放点积注意力：

$${\text{head}}_1=\text{Attention}(Q_1,K_1,V_1)=\text{softmax}\left(\frac{Q_1{K}_1^T}{\sqrt{2}}\right)V_1$$

$${\text{head}}_2=\text{Attention}(Q_2,K_2,V_2)=\text{softmax}\left(\frac{Q_2{K}_2^T}{\sqrt{2}}\right)V_2$$

与之前的单头注意力示例相同，每个头的输出为：

$${\text{head}}_1={\text{head}}_2=\left[\begin{array}{cc}0.669& 0.331\\ 0.331& 0.669\end{array}\right]$$

步骤 5：拼接和线性变换

将所有头的输出拼接起来：

$$\text{Concat}({\text{head}}_1,{\text{head}}_2)=\left[\begin{array}{cccc}0.669& 0.331& 0.669& 0.331\\ 0.331& 0.669& 0.331& 0.669\end{array}\right]$$

通过线性变换 ( W^O )：

$$\text{MultiHead}(Q,K,V)=\left[\begin{array}{cccc}0.669& 0.331& 0.669& 0.331\\ 0.331& 0.669& 0.331& 0.669\end{array}\right]\left[\begin{array}{cccc}1& 0& 1& 0\\ 0& 1& 0& 1\end{array}\right]=\left[\begin{array}{cc}0.669\times 1+0.331\times 0+0.669\times 1+0.331\times 0& 0.669\times 0+0.331\times 1+0.669\times 0+0.331\times 1\\ 0.331\times 1+0.669\times 0+0.331\times 1+0.669\times 0& 0.331\times 0+0.669\times 1+0.331\times 0+0.669\times 1\end{array}\right]=\left[\begin{array}{cc}1.338& 0.662\\ 0.662& 1.338\end{array}\right]$$

结果解释

最终的多头注意力输出是通过将所有头的注意力输出拼接并线性变换得到的矩阵：

$$\text{MultiHead}(Q,K,V)=\left[\begin{array}{cc}1.338& 0.662\\ 0.662& 1.338\end{array}\right]$$

这表示每个位置的最终表示是多个注意力头综合考虑不同子空间信息的结果。

我的理解

你的理解基本上是正确的！多头注意力机制（Multi-Head Attention）确实可以被看作是在自注意力机制（Self-Attention）的基础上，通过定义多套权重矩阵来实现的。具体来说，你的理解步骤如下：

1. 定义多套权重矩阵：

   • 对于每一个注意力头（head），定义一套独立的权重矩阵 ( W_Q )、( W_K ) 和 ( W_V )。

2. 不同变换：

   • 对输入 $X$ 使用不同的权重矩阵，得到多套查询（$ Q$）、键（$K $）和值（$V）矩阵。
   • 每个注意力头都有自己独立的 $Q$、$K$、$V$ 矩阵。

3. 缩放点积注意力：

   • 对每一套 $Q$、$K$、$V$ 矩阵，计算缩放点积注意力：
     $${\text{Attention}}_i(Q_i,K_i,V_i)=\text{softmax}\left(\frac{Q_i{K}_i^T}{\sqrt{d_k}}\right)V_i$$
   • 这里，$d_k$ 是每个注意力头中键向量的维度，通常 $d_k=\frac{d_{\text{model}}}{h}$，其中 $h$ 是注意力头的数量。

4. 拼接输出：

   • 将所有注意力头的输出拼接在一起，形成一个大的输出向量：
     $$\text{Concat}({\text{head}}_1,{\text{head}}_2,\dots ,{\text{head}}_h)$$

5. 线性变换：

   • 对拼接后的向量应用一个线性变换（通常是一个全连接层），将其映射回原始的模型维度 ( d_{\text{model}} )：
     $$\text{MultiHead}(Q,K,V)=\text{Concat}({\text{head}}_1,{\text{head}}_2,\dots ,{\text{head}}_h)W^O$$
   • 其中，$W^O$ 是输出的线性变换矩阵，形状为 $(h\cdot d_k,d_{\text{model}})$。

更详细的解释

为了更清晰地理解多头注意力机制，下面我们通过一个具体的例子来详细解释。

示例参数

• 模型维度（Model Dimension）：$d_{\text{model}}=4$
• 注意力头数（Number of Heads）：$h=2$
• 每个头的维度（Dimension per Head）：$d_k=d_v=\frac{d_{\text{model}}}{h}=2$
• 输入序列长度（Sequence Length）：$n=2$（即有两个词）

步骤 1：定义输入矩阵 ( X )

假设输入序列 ( X ) 有两个词，每个词用一个4维向量表示：

$$X=\left[\begin{array}{cccc}1& 0& 1& 0\\ 0& 1& 0& 1\end{array}\right]$$

步骤 2：定义权重矩阵 $W_Q$、$W_K$、$W_V$ 和$W^O$

为了简化计算，我们手动定义这些权重矩阵。假设每个头共享相同的权重矩阵（实际中通常每个头有独立的权重）：

$$W_Q=W_K=W_V=\left[\begin{array}{cccc}1& 0& 0& 0\\ 0& 1& 0& 0\end{array}\right]$$

输出线性变换矩阵 ( W^O ) 为：

$$W^O=\left[\begin{array}{cccc}1& 0& 1& 0\\ 0& 1& 0& 1\end{array}\right]$$

步骤 3：线性变换生成 ( Q )、( K )、( V ) 矩阵

对于每个头，计算：

$$Q_i=X{W}_Q=\left[\begin{array}{cc}1& 0\\ 0& 1\end{array}\right]$$

$$K_i=X{W}_K=\left[\begin{array}{cc}1& 0\\ 0& 1\end{array}\right]$$

$$V_i=X{W}_V=\left[\begin{array}{cc}1& 0\\ 0& 1\end{array}\right]$$

步骤 4：计算每个头的注意力

对于每个头，计算缩放点积注意力：

$${\text{head}}_i=\text{Attention}(Q_i,K_i,V_i)=\text{softmax}\left(\frac{Q_i{K}_i^T}{\sqrt{d_k}}\right)V_i$$

计算过程如下：

1. 点积：

$$Q_i{K}_i^T=\left[\begin{array}{cc}1& 0\\ 0& 1\end{array}\right]\left[\begin{array}{cc}1& 0\\ 0& 1\end{array}\right]=\left[\begin{array}{cc}1& 0\\ 0& 1\end{array}\right]$$

2. 缩放：

$$\frac{Q_i{K}_i^T}{\sqrt{2}}=\frac{1}{\sqrt{2}}\left[\begin{array}{cc}1& 0\\ 0& 1\end{array}\right]\approx \left[\begin{array}{cc}0.707& 0\\ 0& 0.707\end{array}\right]$$

3. 应用 Softmax：

$$\text{softmax}\left(\left[\begin{array}{cc}0.707& 0\\ 0& 0.707\end{array}\right]\right)\approx \left[\begin{array}{cc}0.669& 0.331\\ 0.331& 0.669\end{array}\right]$$

4. 加权求和：

$${\text{head}}_i=\left[\begin{array}{cc}0.669& 0.331\\ 0.331& 0.669\end{array}\right]\left[\begin{array}{cc}1& 0\\ 0& 1\end{array}\right]=\left[\begin{array}{cc}0.669& 0.331\\ 0.331& 0.669\end{array}\right]$$

步骤 5：拼接和线性变换

将所有头的输出拼接起来：

$$\text{Concat}({\text{head}}_1,{\text{head}}_2)=\left[\begin{array}{cccc}0.669& 0.331& 0.669& 0.331\\ 0.331& 0.669& 0.331& 0.669\end{array}\right]$$

然后通过线性变换 ( W^O )：

$$\text{MultiHead}(Q,K,V)=\left[\begin{array}{cccc}0.669& 0.331& 0.669& 0.331\\ 0.331& 0.669& 0.331& 0.669\end{array}\right]\left[\begin{array}{cccc}1& 0& 1& 0\\ 0& 1& 0& 1\end{array}\right]=\left[\begin{array}{cc}1.338& 0.662\\ 0.662& 1.338\end{array}\right]$$

结果解释

最终的多头注意力输出是通过将所有头的注意力输出拼接并线性变换得到的矩阵：

$$\text{MultiHead}(Q,K,V)=\left[\begin{array}{cc}1.338& 0.662\\ 0.662& 1.338\end{array}\right]$$
这表示每个位置的最终表示是多个注意力头综合考虑不同子空间信息的结果。

多头注意力机制的意义

你的理解准确地捕捉了多头注意力机制的核心思想。进一步来说，多头注意力机制的意义和优势体现在以下几个方面：

1. 多样性：

   • 每个注意力头可以学习到不同类型的关系或特征。例如，一个头可能专注于捕捉句子的语法结构，另一个头可能专注于语义信息。

2. 丰富的表示能力：

   • 通过并行地应用多个注意力头，模型能够生成更加丰富和多样化的表示，从而提升整体性能。

3. 增强的学习能力：

   • 多头注意力机制使得模型能够在不同的子空间中独立学习，这有助于捕捉更复杂的模式和依赖关系。

4. 高效的并行计算：

   • 多头注意力机制能够充分利用现代硬件（如GPU）的并行计算能力，提高计算效率。

总结

你对多头注意力机制的理解是正确且全面的。通过定义多套权重矩阵，分别对输入进行不同的线性变换，生成多个 ( Q )、( K )、( V ) 矩阵，然后分别计算缩放点积注意力，最后将这些头的输出拼接并通过线性变换得到最终的输出，多头注意力机制能够显著增强模型的表达能力和性能。

解码器

KV Cache优化解码过程

PagedAttention

vLLM比直接调用模型的generate方法快14-24倍

以13B的模型为例，在推理时模型大约占用26GB显存约65%，KV Cache占用约12GB显存

KV Cache的利用率只用20%-40%。

在大模型推理时，按照可生成最大序列长度分配显存。

造成三种类型的浪费：

1. 预分配，但不会用到。
2. 预分配，但尚未用到。
3. 显存之间的间隔碎片，不足以分配给下一个文本生成。

vLLM服务部署

python -m vllm.entrypoints.openai.api_server --model xx --host 0.0.0.0 --tensor-parallel-size=2 --dtype half &

这是一个使用 Python 和 FastAPI 部署 vLLM 模型的基本示例，包含你的要求：

from fastapi import FastAPI, HTTPException
import logging
import vllm
import torch

app = FastAPI()

# 配置日志记录
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 加载模型（可选 LoRA 权重）
def load_model(model_name: str, lora_weights: str = None, quantized: bool = False):
    if quantized:
        model = vllm.Model(model_name, dtype='half', tensor_parallel_size=2)
    else:
        model = vllm.Model(model_name)
    if lora_weights:
        model.load_lora_weights(lora_weights)
    return model

model = load_model("your_model_name", lora_weights=None, quantized=True)

@app.post("/predict/")
async def predict(prompt: str, lora_weights: str = None):
    try:
        if lora_weights:
            model.load_lora_weights(lora_weights)
        response = await model.generate(prompt)
        return {"response": response}
    except Exception as e:
        logger.error(f"预测时出错: {e}")
        raise HTTPException(status_code=500, detail="预测错误")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)

关键点：

1. 模型加载：load_model 函数允许加载量化模型和可选的 LoRA 权重。
2. 日志记录：配置了基本的日志记录，用于跟踪错误。
3. 并发处理：FastAPI 默认处理多个请求并发。
4. 接口：/predict/ 接口接受输入提示和可选的 LoRA 权重进行推理。

根据需要调整日志记录和错误处理！

一、LLM推理的两阶段

一个常规的LLM推理过程通常分为两个阶段：prefill和decode。通常会使用KV cache技术加速推理。

1. Prefill预填充阶段

• 输入处理：在这个阶段，整个prompt（输入文本）被输入到模型中进行前向计算（forward pass）。
• KV Cache生成：
  • 模型通过参数 $W_k$和 $W_v$ 将输入转化为键（Key）和值（Value），得到 $X_k$和$X_v$。
  • 这些键值对会被存储在 KV Cache 中（分别命名为 cache_k 和 cache_v）。
• 效率提升：存储在 KV Cache 中的键值对可以在后续的注意力计算中复用，避免对之前的 token 进行重复计算，从而节省推理时间。

2. Decode生成响应阶段

• 逐步生成：在这个阶段，模型逐个生成响应的 token。每次生成一个 token，都会使用前面所有已生成的 token 的 KV 值进行计算。
• KV Cache更新：
  • 每生成一个 token，例如 $ t_4$，模型会使用t0t3的‘cach{e}_k‘和‘cach{e}_v‘中的值计算注意力，然后将$t_4$的 KV 值存入 KV Cache 中。
  • 同理，对于下一个生成的 token$t_6$，也会更新 KV Cache。

3. 效率和限制

• 计算时间：Decode 阶段的计算通常比 Prefill 阶段更耗时，因为它是逐个生成 token，而不是并行处理。
• 内存压力：
  • 随着输入（prompt）的长度和输出（response）序列的增加，KV Cache 的大小也会增长，这对 GPU 的显存造成压力。
  • 由于输出序列的长度通常无法预先确定，因此很难为 KV Cache 提前分配固定的存储空间。

下图展示了一个13B的模型在A100 40GB的gpu上做推理时的显存占用分配（others表示forward过程中产生的activation的大小，这些activation你可以认为是转瞬即逝的，即用完则废，因此它们占据的显存不大），从这张图中我们可以直观感受到推理中KV cache对显存的占用。因此，如何优化KV cache，节省显存，提高推理吞吐量，就成了LLM推理框架需要解决的重点问题。

activation指的是在模型前向推理过程中产生的中间激活值。这些激活值是在神经网络的每一层中计算出来的，用于存储每个节点的输出（即神经元的激活状态），以供后续计算和反向传播使用。

二、为KV Cacha分配存储空间的常规方式

对于训练好的模型，一种常用的部署方式是将其打包成一个推理服务（server），它接收客户端发送来的请求（request），读取请求中的数据（prompt）来做推理。一个请求中可以只有1个prompt，也可以包含多个prompt。

在常规的推理框架中，当我们的服务接收到一条请求时，它会为这条请求中的prompts分配gpu显存空间，其中就包括对KV cache的分配。**由于推理所生成的序列长度大小是无法事先预知的，所以大部分框架会按照(batch_size, max_seq_len)这样的固定尺寸，在gpu显存上预先为一条请求开辟一块连续的矩形存储空间。然而，这样的分配方法很容易引起“gpu显存利用不足”的问题，进而影响模型推理时的吞吐量。**你可能觉得这个描述有点抽象，别着急，我们来具体看一个例子。

下图展示了一个常规的推理框架是如何为请求中的prompt在gpu显存上分配KV cache的。在本例中，我们假设一个请求只发送1条prompt（本例中共有3条请求）：

我们假设max_seq_len = 8，所以当第1条请求(prompt1)过来时，我们的推理框架为它安排了(1, 8)大小的连续存储空间。

当第2条请求（prompt2）过来时，同样也需要1块(1, 8)大小的存储空间。但此时prompt1所在的位置上，只剩3个空格子了，所以它只能另起一行做存储。对prompt3也是同理。

仔细观察这3条prompt的KV cache排布，你是不是隐约觉得这种排布似乎没有充分利用起gpu的显存？：

• **浅色块：**观察图中的浅色块，它是prefill阶段prompt的KV cache，是无论如何都会被使用的空间，它不存在浪费。
• 中色块：观察图中的中色块，它是decode阶段的KV cache，其中<eos>表示序列生成的截止符。虽然这些中色块最终都会被我们用上，但是在decode阶段一个个token生成时，我们并不能预知哪些块会被最终用上。例如对于prompt2，当你生成when的时候，你无法知道下一个会生成<eos>，还是会生成别的词。所以这些中色块都是一种“潜在的浪费”，我们称中色块的部分为预留碎片（reservation fragment）。
• **深色块：**观察图中的深色块，它也是decode阶段的KV cache，但直到序列生成完毕，它都没有被用上。由于这些深色块是预留的KV cache的一部分，所以我们称其为内部碎片（internal fragment）。
• **灰色块：**观察图中的灰色块，它不是我们预留的KV cache的一部分，且最终也没有被用上，**我们称这些灰色块为外部碎片（external fragment）。**想象一下，此时新来了一条prompt4，它也要求显存中的8个格子作为KV cache。此时你的显存上明明有9个空格子，但因为它们是不连续的碎片，所以无法被prompt4所使用。这时prompt4的这条请求只好在队列中等待，直到gpu上有足够显存资源时再进行推理，这不就对模型推理的吞吐量造成显著影响了吗？

观察整个KV cache排布，你会发现它们的毛病在于太过“静态化”。当你无法预知序列大小时，你为什么一定要死板地为每个序列预留KV cache空间呢？**为什么不能做得更动态化一些，即“用多少占多少”呢？**这样我们就能减少上述这些存储碎片，使得每一时刻推理服务能处理的请求更多，提高吞吐量，这就是vLLM在做的核心事情，我们先通过一张实验图来感受下vLLM在显存利用上的改进效果（VS 其它推理框架）：

不难发现，相比于别的推理框架，vLLM几乎能做到将显存完全打满。

读到这里，你可能会有以下疑问：

• vLLM是通过什么技术，动态地为请求分配KV cache显存，提升显存利用率的？
• 当采用动态分配显存的办法时，虽然明面上同一时刻能处理更多的prompt了，但因为没有为每个prompt预留充足的显存空间，如果在某一时刻整个显存被打满了，而此时所有的prompt都没做完推理，那该怎么办？

在后文的第三～四章，我们将回答问题1。第五章回答问题2。

三、PagedAttention原理

在本节中，我们先来回答问题1：vLLM通过一种名为PagedAttention的技术，动态地为请求分配KV cache显存，提升显存利用率。

**整体上来说，PagedAttention的设计灵感来自操作系统中虚拟内存的分页管理技术。**所以本节会先通过通俗易懂的方式，和大家一起快速回顾操作系统的虚拟内存技术，在这个过程中和大家一起具象化感受PagedAttention的设计思想。然后再来详细介绍PagedAttention的运作流程。

3.1 操作系统的虚拟内存

（1）不使用虚拟内存

我们知道程序运行时，会将代码、数据等内容存放在物理内存上。在最原始的做法中（没有操作系统，例如单片机），程序直接对物理内存进行操作，决定使用它的哪些存储地址。

如果你只跑一个进程，那还好说。但如果需要运行多个进程时，麻烦就来了：由于我直接操作了物理内存地址，所以我在为自己的进程分配物理内存时，还要考虑别的进程是如何分配物理内存的（别人已经占用的我不能用）。这样不同进程间的耦合性太高了，给开发带来难度。

有没有一种办法，让各个进程间的开发能够相互独立呢？一种直觉的做法是：

• 给每个进程分配一个虚拟内存。每个进程在开发和运行时，可以假设这个虚拟内存上只有自己在跑，这样它就能大胆操作。
• **虚拟内存负责统一规划代码、数据等如何在物理内存上最终落盘。**这个过程对所有进程来说都是透明的，进程无需操心

虚拟内存的核心思想可简化成下图:

（2）虚拟内存的分段管理

在分段式内存管理中，虚拟内存会尽量为每个进程在物理内存上找到一块连续的存储空间，让进程加载自己的全部代码、数据等内容。我们来看一个具体的例子：

在这个例子中，3个进程的虚拟内存各自为它们在物理内存上映射了一块连续的存储空间。在某一时刻，我释放了进程2，同时想运行进程4。这时我尴尬地发现，虽然物理内存上有640M的空间剩余，但因为是碎片化的，我的进程4无法加载进去，因此它只能等待（回想一下本文第二部分对传统KV cache显存分配的分析）。

在这个情况下，如果我硬要运行进程4，也是有办法的：我可以先把进程3从物理内存上交换（swap）到磁盘上，然后把进程4装进来，然后再把进程3从磁盘上加载回来。通过这种方法我重新整合了碎片，让进程4能够运行。

但这种办法的显著缺点是：如果进程3过大，同时内存到磁盘的带宽又不够，整个交换的过程就会非常卡顿。这就是分段式内存管理的缺陷。

这时，我自然而然会想到：**我为什么要给所有进程都预分配一个固定的存储块（段）呢？**假设这个进程是一个浏览器，我难道会一下就用到这个进程里所有的功能吗？就不能进程运行到哪里，或者我想用哪个具体功能时，再加载这部分相关的内容去内存，以此让整个内存分配更加动态？

（3）虚拟内存的分页管理

我们可以将进程1、进程2想成是两本书。代码分布在书的不同page上。我们希望想读哪一页，就加载哪一页，而不是一下把两本书都加载进来。同时，当我们不想读某些页的时候，我们也能根据页码将其清空。

现在，我们希望读进程1和进程2的page1，我们就将其加载到物理内存上。虚拟内存会帮我们做好映射，把来自不同进程的这两页分别加载到物理内存对应位置。

虚拟内存的分业管理技术总结起来就是：

• 将物理内存划分为固定大小的块，我们称每一块为页（page）。从物理内存中模拟出来的虚拟内存也按相同的方式做划分
• 对于1个进程，我们不需要静态加载它的全部代码、数据等内容。我们想用哪部分，或者它当前跑到哪部分，我们就动态加载这部分到虚拟内存上，然后由虚拟内存帮我们做物理内存的映射。
• 对于1个进程，虽然它在物理内存上的存储不连续（可能分布在不同的page中），但它在自己的虚拟内存上是连续的。通过模拟连续内存的方式，既解决了物理内存上的碎片问题，也方便了进程的开发和运行。

3.2 PagedAttention

（1）处理单个请求
现在，你已经知道虚拟内存分页管理的基本原理和优势，趁热打铁，我们马上来看以其为灵感的PagedAttention技术是如何操作的。我们还是从具体的例子讲起。

假设现在你向模型server发送一条请求，prompt为Four score and seven years ago our，你希望模型能做续写。PagedAttention的运作流程如下图：

在图中：

• 请求（request）可理解为操作系统中的一个进程
• 逻辑内存（logical KV blocks）可理解为操作系统中的虚拟内存，每个block类比于虚拟内存中的一个page。每个block的大小是固定的，在vLLM中默认大小为16，即可装16个token的K/V值
• 块表（block table）可理解为操作系统中的虚拟内存到物理内存的映射表
• 物理内存（physical KV blocks）可理解为操作系统中的物理内存，物理块在gpu显存上，每个block类比于虚拟内存中的一个page

图中带圈的序号表示操作步骤，我们就按这个顺序来看看。

(i) Prefill阶段

• 划分逻辑块：vLLM拿到这条prompt，先按照设定好的block大小B（本例中B=4），为prompt划分逻辑块（Logical KV blocks）。由于prompt中有7个token，所以vLLM用2个逻辑块（block 0， block 1）来装它们的KV值。其中，在逻辑块1中目前只装了"years", “ago”, "hour"这3个token的KV值，有1个位置是空余的。这个位置就被称为保留位（reservation）

• 划分物理块：划分好逻辑块后，我们就可以将其映射到物理块中去了。物理块是实际存放KV值的地方。我们通过一张block table来记录逻辑块和物理块的映射关系，block table的主要内容包括：

• • 逻辑块和物理块的映射关系（physical block number）：例如逻辑块0对应物理块7
  • 每个物理块上被填满的槽位（# filled）：例如在prefill阶段，对物理块7，其4个槽位都被填满；对物理块1，其3个槽位被填满。

• 正常计算prompt的KV值，并通过划分好的关系填入物理块中。

(ii) Decode阶段-生成第1个词。

• 使用KV cache计算attention，生成第1个词fathers。不难发现，当我们计算时，我们使用的是逻辑块，即形式上这些token都是连续的。与此同时，vLLM后台会通过block table这个映射关系，帮我们从物理块上获取数据做实际计算。通过这种方式，每个request都会认为自己在一个连续且充足的存储空间上操作，尽管物理上这些数据的存储并不是连续的。
• 基于新生成的词，更新逻辑块、物理块和block table。对于block table，vLLM将它filled字段由3更新至4。
• 分配新的逻辑块和物理块。当fathers更新进去后，逻辑块已装满。所以vLLM将开辟新的逻辑块2，并同时更新对应的block table和物理块。

(ii) Deocde阶段-生成第2个词

类比步骤（2）来进行。

（2）处理多个请求

有了（1）的解释，大家看懂这张图应该不难。通过多个请求（prompt）同时做推理的例子，大家可以更好感受到PagedAttention是如何通过动态存储KV cache的方式，来更充分利用gpu显存的。

。

3.2 PagedAttention

（1）处理单个请求
现在，你已经知道虚拟内存分页管理的基本原理和优势，趁热打铁，我们马上来看以其为灵感的PagedAttention技术是如何操作的。我们还是从具体的例子讲起。

假设现在你向模型server发送一条请求，prompt为Four score and seven years ago our，你希望模型能做续写。PagedAttention的运作流程如下图：

在图中：

• 请求（request）可理解为操作系统中的一个进程
• 逻辑内存（logical KV blocks）可理解为操作系统中的虚拟内存，每个block类比于虚拟内存中的一个page。每个block的大小是固定的，在vLLM中默认大小为16，即可装16个token的K/V值
• 块表（block table）可理解为操作系统中的虚拟内存到物理内存的映射表
• 物理内存（physical KV blocks）可理解为操作系统中的物理内存，物理块在gpu显存上，每个block类比于虚拟内存中的一个page

图中带圈的序号表示操作步骤，我们就按这个顺序来看看。

(i) Prefill阶段

• 划分逻辑块：vLLM拿到这条prompt，先按照设定好的block大小B（本例中B=4），为prompt划分逻辑块（Logical KV blocks）。由于prompt中有7个token，所以vLLM用2个逻辑块（block 0， block 1）来装它们的KV值。其中，在逻辑块1中目前只装了"years", “ago”, "hour"这3个token的KV值，有1个位置是空余的。这个位置就被称为保留位（reservation）

• 划分物理块：划分好逻辑块后，我们就可以将其映射到物理块中去了。物理块是实际存放KV值的地方。我们通过一张block table来记录逻辑块和物理块的映射关系，block table的主要内容包括：

• • 逻辑块和物理块的映射关系（physical block number）：例如逻辑块0对应物理块7
  • 每个物理块上被填满的槽位（# filled）：例如在prefill阶段，对物理块7，其4个槽位都被填满；对物理块1，其3个槽位被填满。

• 正常计算prompt的KV值，并通过划分好的关系填入物理块中。

(ii) Decode阶段-生成第1个词。

• 使用KV cache计算attention，生成第1个词fathers。不难发现，当我们计算时，我们使用的是逻辑块，即形式上这些token都是连续的。与此同时，vLLM后台会通过block table这个映射关系，帮我们从物理块上获取数据做实际计算。通过这种方式，每个request都会认为自己在一个连续且充足的存储空间上操作，尽管物理上这些数据的存储并不是连续的。
• 基于新生成的词，更新逻辑块、物理块和block table。对于block table，vLLM将它filled字段由3更新至4。
• 分配新的逻辑块和物理块。当fathers更新进去后，逻辑块已装满。所以vLLM将开辟新的逻辑块2，并同时更新对应的block table和物理块。

(ii) Deocde阶段-生成第2个词

类比步骤（2）来进行。

（2）处理多个请求

有了（1）的解释，大家看懂这张图应该不难。通过多个请求（prompt）同时做推理的例子，大家可以更好感受到PagedAttention是如何通过动态存储KV cache的方式，来更充分利用gpu显存的。

PagedAttention 概述这个好懂

PagedAttention 是一种借鉴虚拟内存分页管理的技术，旨在优化大模型推理中的KV Cache管理，以更有效地利用GPU显存。

1. 处理单个请求

(i) Prefill阶段

• 逻辑块划分：

  • 输入的prompt（如“Four score and seven years ago our”）被划分为固定大小的逻辑KV块（每个块默认大小为16，示例中设为4）。
  • 由于prompt中有7个token，划分为2个逻辑块（block 0和block 1）。block 1中有一个空余位置，被称为“保留位”。

• 物理块划分：

  • 将逻辑块映射到物理块，记录在块表中。每个物理块对应某个逻辑块，并记录已填满的槽位数。

• KV值计算与存储：

  • 在prefill阶段，模型计算每个token的KV值，并将其存储到相应的物理块中。

(ii) Decode阶段 - 生成第1个词

• 计算注意力：

  • 使用逻辑块中的KV值计算生成第一个token（如“fathers”）。虽然这些token在物理存储上不连续，但逻辑上视为连续的存储空间。

• 更新状态：

  • 根据生成的token，更新逻辑块、物理块和块表。例如，将block table中filled字段从3更新至4。

• 分配新块：

  • 由于逻辑块已装满，vLLM会开辟新的逻辑块，并更新对应的块表和物理块。

2. 处理多个请求

在处理多个请求时，PagedAttention展示了其动态管理KV Cache的能力：

• 多个请求的输入被同时划分和处理，确保在GPU显存中高效存储和利用。
• 每个请求仍然保持对逻辑连续存储的“感知”，即使物理存储并不连续。

优势

• 显存优化：通过动态管理KV Cache和块表，PagedAttention能够更有效地利用GPU显存，减少资源浪费。
• 灵活性：允许处理多个请求，同时保持操作的高效性。

总结

PagedAttention通过引入虚拟内存分页的概念，优化了大模型推理中KV Cache的管理，能够在处理单个或多个请求时，显著提高显存的利用效率和计算的灵活性。

如果你对PagedAttention的具体实现或其他方面有疑问，欢迎继续探讨！

四、PagedAttention在不同解码策略下的运作

下面说明在parallel sampling的场景下，vLLM（PagedAttention）是怎么做到节省显存的。

传统KV cache怎么做：
假设模型的max_seq_len = 2048。传统KV cache可能在显存中分配两块长度是2048的空间。由于prompt一致，这两块2048的空间中存在大量重复的KV cache。

vLLM怎么做：

假定我们发给模型1个request，这个request中包含2个prompt/sample，记为Sample A1和Sample A2，这两个prompt完全一致，都为Four score and seven years ago our，我们希望模型对这两个prompt分别做续写任务。

（1）首先，Prefill阶段，vLLM拿到Sample A1和Sample A2，根据其中的文字内容，为其分配逻辑块和物理块。

• **分配逻辑块：**对于A1，vLLM为其分配逻辑块block0和block1；对于A2，vLLM为其分配逻辑块block0和block1。需要注意的是，A1的逻辑块和A2的逻辑块是独立的（尽管它们都叫block0和block1），你可以将A1和A2视作操作系统中两个独立运行的进程。
• 分配物理块：对于A1和A2，虽然逻辑块独立，但因为它们的文字完全相同，所以可以在物理内存上共享相同的空间。所以A1的逻辑块block0/1分别指向物理块block7/1；A2的逻辑块block0/1分别指向物理块block7/1。我们设每个物理块下映射的逻辑块数量为ref count，所以对物理块block7/1来说，它们的ref count都为2。

（2）然后，进入decode阶段，A1和A2各自做推理，得到第一个token，分别为fathers和mothers。

• 将生成的token装入逻辑块：对于A1和A2来说，将其生成的token装入各自的逻辑块block1。
• 触发物理块copy-on-write机制：由于fathers/mothers是两个完全不同的token，因此对物理块block1触发复制机制，即在物理内存上新开辟一块空间。此时物理块block1只和A2的逻辑块block1映射，将其ref count减去1；物理块block3只和A1的逻辑块block1映射，将其ref count设为1。

总结起来，vLLM节省KV cache显存的核心思想是，对于相同数据对应的KV cache，能复用则尽量复用；无法复用时，再考虑开辟新的物理空间。

4.2 Beam Search

**我们从右往左来看这张图。**虚线位置表示“当前decoding时刻”，beam width = 4。图中所有的block皆为逻辑块。

因为beam width = 4，这意味着根据beam search算法，在当前阶段我们生成了top 4个概率最大的token（我们记这4个token为beam candidate 0/1/2/3），它们分别装在block5，block6，block7和block8中。

现在我们继续使用beam search算法做decoding，继续找出top 4个最可能的next token。经过我们的计算，这top 4 next token，有2个来自beam candidate 1，有2个来自beam candidate 2。因此我们在block6中引出block9和block10，用于装其中两个top 2 next token；对block7也是同理。

现在，block9/10/11/12中装的top 4 next token，就成为新的beam candidates，可以按照和上述一样的方式继续做beam search算法。而对于block5和block8，它们已经在beam search的搜索算法中被淘汰了，后续生成的token也不会和它们产生关系，所以可以清除掉这两个逻辑块，并释放它们对应的物理块的内存空间。

好，我们继续往左边来看这幅图。block3引出block5/6/7，block4引出block8，这意味着当前这4个top4 token，是上一个timestep下candidate1和candidate3相关序列生成的（candidate0和2的block没有画出，是因为它们所在的序列被beam search算法淘汰了，因此没有画出的必要）。由于block8已经被淘汰，所以block4也相继被淘汰，并释放对应的物理内存空间。

由此往左一路推，直到block0为止（block0代表着prompt，因此被beam seach中所有的序列共享）。这一路上，我们都根据最新时刻的beam search decoding结果，释放掉不再被需要的逻辑块和对应的物理内存空间，达到节省显存的目的。

---

好的，让我们更详细地探讨PagedAttention在不同解码策略下的运作，特别是在parallel sampling和beam search场景中如何节省显存。

1. Parallel Sampling

传统KV Cache的缺陷

• 显存占用：
  • 假设模型的最大序列长度（max_seq_len）为2048。在传统的KV缓存方案中，为每个请求（request）分配两块2048长度的显存。这意味着即使多个请求具有相同的内容（例如相同的prompt），它们仍然会在显存中各自占用一块2048的空间，从而导致重复数据的显著浪费。

vLLM的优化策略

• 逻辑块和物理块的分配：

  • 当vLLM接收到Sample A1和Sample A2（两个完全相同的prompt）时，它为每个样本分配逻辑块（如block0和block1）。需要注意的是，尽管A1和A2的逻辑块命名相同，但它们在逻辑上是独立的，可以看作操作系统中两个独立的进程。

• 共享物理内存：

  • 由于这两个请求的内容相同，vLLM能够在物理内存中共享相同的空间。例如，A1的逻辑块block0和block1都指向物理块block7和block1，而A2的逻辑块也指向这两个物理块。通过这种方式，物理块的引用计数（ref count）将为2，表示有两个逻辑块正在使用这块物理空间。

Decode阶段的操作

• 生成token的存储：

  • 在decode阶段，当A1和A2分别生成第一个token（如fathers和mothers）时，vLLM将这些token存入各自的逻辑块（block1）。

• 触发物理块的复制机制：

  • 由于fathers和mothers是不同的token，这会触发物理块的“写时复制”机制。vLLM会为新的token分配新的物理空间。例如，物理块block1现在只与A2的逻辑块block1关联，引用计数减为1；而物理块block3则仅与A1的逻辑块block1关联，引用计数设为1。

2. Beam Search

解码过程

• 逻辑块的运用：
  • 在beam search策略中，假设束宽（beam width）为4，这意味着在当前时刻，模型会生成概率最高的4个token，分别存放在不同的逻辑块（如block5, block6, block7, block8）中。

清理不必要的块

• 淘汰逻辑块：

  • 随着解码过程的推进，某些逻辑块（如block5和block8）可能会因未被选中而被淘汰。在后续的解码步骤中，这些块会被清除，释放对应的物理内存空间。

• 更新逻辑块的映射：

  • 当继续进行beam search时，新的top token会从先前的候选中生成。例如，block3引出block5/6/7，block4引出block8，这表示新的top token是来自上一个时刻候选序列的结果。淘汰的块不会被后续的token引用，从而可以有效释放内存。

总结

vLLM通过以下方式优化显存使用：

• 复用机制：对于相同的KV cache，尽可能复用已有的物理空间，避免重复分配内存。
• 动态管理：在解码过程中，根据需要动态释放不再使用的逻辑块及其对应的物理内存，确保显存利用率最大化。

这种方法尤其适合处理重复请求和复杂生成任务，通过高效的显存管理提升了大模型的推理性能和效率。如果还有其他具体问题或细节需要探讨，请随时告诉我！

五、调度和抢占

到目前为止，我们已经回答了“vLLM是如何优化KV cache显存分配”的问题，现在我们来回答另一个重要的问题：

• 当采用动态分配显存的办法时，虽然明面上同一时刻能处理更多的prompt了，但因为没有为每个prompt预留充足的显存空间，如果在某一时刻整个显存被打满了，而此时所有的prompt都没做完推理，那该怎么办？

5.1 总原则

当有一堆请求来到vLLM服务器上时，vLLM需要一个调度原则来安排如何执行这些请求，这个调度原则概括如下：

• 先来的请求先被服务（First-Come-First-Serve, FCFS）
• 如有抢占的需要，后来的请求先被抢占（preemption）

（1）先来的请求先被服务
这个很好理解，当有一堆请求到达vLLM服务器时，vLLM肯定优先处理来得早的请求

（2）后来的请求先被抢占
想象一下，当一堆请求来到vLLM服务器做推理，导致gpu显存不足时，vLLM会怎么做呢？

最直接的办法，就是暂停这堆请求中最后到达的那些请求的推理，同时将它们相关的KV cache从gpu上释放掉，以便为更早到达的请求留出足够的gpu空间，让它们完成推理任务。如果不这样做的话，各个请求间相互争夺gpu资源，最终将导致没有任何一个请求能完成推理任务。等到先来的请求做完了推理，vLLM调度器认为gpu上有足够的空间了，就能恢复那些被中断的请求的执行了。

在资源不足的情况下，暂时中断一些任务的执行，这样的举动就被称为“抢占（preemption）”。

5.2 终止和恢复被抢占的请求

对于这些因gpu资源不足而被抢占的任务，vLLM要完成两件事：

• 暂停它们的执行，同时将与之相关的KV cache从gpu上释放掉
• 等gpu资源充足时，重新恢复它们的执行

针对这两件事，vLLM分别设计了**Swapping（交换策略）和Recomputation（重计算策略）**来解决。我们来细看这两个策略。

（1）Swapping

对于被抢占的请求，vLLM要将其KV cache从gpu上释放掉，那么：

• 问题1：该释放哪些KV cache？
• 问题2：要把这些KV cache释放到哪里去？

先看问题1。由前文PagedAttention原理可知，一个请求可能对应多个block。我们既可以选择释放掉部分block，也可以选择释放掉全部block，或者更科学地，我们可以预测一下哪些block被使用的频率最低，然后释放掉这些低频block（但这种方式实现起来难度较大，性价比不是很高）。在vLLM中，采取的是all-or-nothing策略，即释放被抢占请求的所有block。

再来看问题2。对于这些被选中要释放的KV block，如果将它们直接丢掉，那未免过于浪费。vLLM采用的做法是将其从gpu上交换（Swap）到cpu上。这样等到gpu显存充份时，再把这些block从cpu上重载回来。

（2）Recomputation

知道了Swapping机制，重计算的过程也很好理解了：对于有些任务（比如parallel sampling中并行采样数n=1的任务），当它们因为资源不足而被抢占时，可以不做swap，而是直接释放它们的物理块，把它们重新放入等待处理的队列中，等后续资源充足时再重新从prefill阶段开始做推理

好，到这里，我们总结一下vLLM对请求的调度处理流程：

• 当一堆请求来到vLLM服务器上时，按照**First-Come-First-Serve（FCFS）**原则，优先处理那些最早到来的请求。
• 当gpu资源不足时，为了让先来的请求能尽快做完推理，vLLM会对那些后到来的请求执行“抢占”，即暂时终止它们的执行。
• 一旦vLLM决定执行抢占操作，它会暂停处理新到来的请求。在此期间，它会将被抢占的请求相关的KV block全部交换（swap）至cpu上。等交换完成后，vLLM才会继续处理新到来的请求。
• 当vLLM认为gpu有足够资源时，它会将cpu上的KV block重新加载回gpu，恢复被抢占请求的执行（recomputation）

让我们逐步解析这段关于vLLM如何处理显存分配和调度请求的内容。

1. 关键问题

在采用动态分配显存的策略时，尽管可以在同一时刻处理更多的prompt，但如果GPU显存被填满，而某些prompt的推理尚未完成，就会面临资源不足的问题。那么，vLLM是如何应对这一挑战的呢？

2. 调度原则

vLLM采用了一套明确的调度原则，主要包括：

• 先来先服务（FCFS）：

  • 当请求同时到达时，优先处理最早到达的请求。

• 抢占（Preemption）：

  • 在显存资源不足的情况下，后来的请求可以暂时中断，释放其相关的KV cache，以确保优先请求可以顺利完成。

3. 抢占机制

处理请求的具体步骤

• 暂停新到来的请求：一旦检测到显存不足，vLLM会暂停处理新的请求。

• 释放被抢占请求的KV cache：为了为优先请求留出足够的显存，vLLM会暂停后来的请求并释放其占用的KV cache。

抢占的操作细节

• 暂停和释放：

  • vLLM会将被抢占请求的所有相关block从GPU上释放，确保先来的请求能够使用这些显存资源。

• 资源恢复：

  • 在优先请求完成后，vLLM会检查显存状态，如果资源充足，就会将之前被交换到CPU的KV cache重新加载到GPU，恢复被抢占的请求的执行。

4. Swapping（交换策略）

释放KV cache的方式

• 释放哪些KV cache：

  • vLLM采用“全有或全无”的策略，释放被抢占请求的所有block，而不是部分释放。这种方法简单且直接。

• 交换到哪里：

  • 被释放的KV cache不会直接丢弃，而是通过交换（swap）机制转移到CPU上，待GPU资源充足后再从CPU加载回来。

5. Recomputation（重计算策略）

• 重计算的使用场景：
  • 对于某些任务（例如parallel sampling中n=1的情况），在资源不足时，可以选择直接释放物理块，而不执行swap。这些任务会被重新放入等待处理的队列中，待资源充足时从prefill阶段重新开始推理。

6. 总结

vLLM的请求调度处理流程总结如下：

• 依据FCFS原则，优先处理早到的请求。
• 在GPU资源不足时，通过抢占机制暂停后来的请求，并释放其相关的KV block至CPU。
• 等待资源恢复后，将之前的KV block重新加载回GPU，并恢复被抢占请求的执行。

这种调度策略确保了即使在资源有限的情况下，vLLM仍然能高效地完成任务，提高了整体的推理性能和资源利用率。如果需要进一步深入某个方面，请告诉我！

模型权重

这两种模型权重文件格式（*.safetensor 和 *.bin）各有特点，以下是它们的详细技术解读：

1. *.safetensor 格式

• 背景和发展：safetensor 是一种新型的文件格式，旨在提高模型权重的安全性和读取效率。由 Hugging Face 社区开发，主要用来替代传统的 .bin 格式。
• 安全性：
  • safetensor 通过固定的文件格式防止潜在的文件结构被篡改或错误读写，减少了恶意攻击的风险。传统的 .bin 文件可能包含不规范的序列化数据，容易造成解析错误或安全漏洞。
• 性能：
  • safetensor 文件的读取速度通常比 .bin 更快，因为它的格式更简洁高效，优化了 IO 性能，特别是在大规模分布式系统中。
  • 文件中包含元数据，并且可以直接加载特定的张量（tensor）而不必全部加载整个文件，这对于存储非常大的模型权重时特别有用。
• 跨平台和兼容性：与 PyTorch 和 TensorFlow 兼容性较好，尤其是在 Hugging Face 的库中已经被广泛应用。

2. *.bin 格式

• 背景和历史：.bin 格式是深度学习社区最早采用的标准之一，通常用于保存模型的权重和参数。很多框架（例如 PyTorch、TensorFlow）默认以这种格式保存模型。
• 文件结构：
  • .bin 文件本质上是二进制文件，存储经过序列化后的权重。它可以包含任意张量（tensor）的数据，包括浮点数和整数。由于格式较为通用，它在多种场景下得到广泛使用。
  • 这种格式没有严格的文件安全性限制，容易出现由于解析器或文件损坏导致的读取错误。
• 读取性能：
  • .bin 文件的读取性能相对较慢，因为读取时通常需要一次性加载全部权重。对于非常大的模型，特别是在有限内存的设备上，这可能会导致性能瓶颈。
  • .bin 文件不包含内置的张量元数据或安全校验功能，这意味着在一些大规模的分布式训练环境中，可能会出现一致性或安全性问题。

总结

• safetensor：更安全、读取速度更快，适合现代化、大规模、分布式的模型部署场景。它提供了更好的元数据支持，允许部分加载张量，提升了灵活性。
• bin：历史悠久、应用广泛，但在现代大规模部署场景中逐渐表现出一些局限性，比如加载速度慢和安全性不足。

选择哪种格式取决于你对安全性、性能和兼容性的需求。如果追求更高效和安全的模型部署，safetensor 是优选；而 .bin 格式则在更传统或通用的场景中仍然非常有用。

内存共享

张量并行与数据并行

博客

感觉张量并行重要一点

张量并行（Tensor Parallelism）和数据并行（Data Parallelism）是两种常见的并行化技术，广泛应用于 深度学习 任务，尤其是在训练和推理大规模神经网络时。这两种方法通过不同的方式利用多张 GPU 来加速计算，减轻单个 GPU 的负担。下面我们来详细比较这两者。

1. 数据并行（Data Parallelism）

概念

数据并行是指将同一个模型的多个副本复制到不同的 GPU 上，使用不同的数据批次（mini-batches）进行并行计算。每个 GPU 处理不同部分的数据，独立地计算损失和梯度。然后，这些 GPU 将它们的梯度汇总起来，以更新整个模型的参数。

工作流程

• 模型复制：在多张 GPU 上复制同一个模型，每张 GPU 运行相同的模型副本。
• 数据分片：将数据集划分为多个小批次，每个 GPU 处理一个批次的数据。举例来说，如果你有 4 张 GPU 和一个大小为 1024 的 mini-batch，那么每张 GPU 将处理大小为 256 的 mini-batch。
• 并行计算：每个 GPU 独立计算损失和梯度。
• 梯度同步：每个 GPU 计算的梯度会被同步和汇总（通常通过通信操作，如 all-reduce），主 GPU 或 CPU 收集所有梯度，并更新模型参数。
• 更新模型：所有 GPU 都使用同一个更新后的模型参数进行下一轮计算。

优点

• 易实现：数据并行的实现较为简单，适合大多数深度学习框架，如 TensorFlow、PyTorch。
• 可扩展性：适合模型参数较小、数据量较大的情况，因为模型副本不会消耗太多显存，GPU 间的计算是相对独立的。

缺点

• 通信开销：每一轮训练后，各 GPU 都要同步它们的梯度，这会产生较大的通信开销，尤其是当 GPU 数量增多时。
• 显存浪费：因为每个 GPU 都保留了一份完整的模型副本，所以当模型非常大时（如 GPT-3），单张 GPU 的显存可能不足以容纳整个模型。

场景

数据并行通常用于模型较小但数据量很大的场景。例如在图像分类任务中，每张 GPU 处理不同的图像批次，模型的更新是基于所有 GPU 计算结果的汇总。

2. 张量并行（Tensor Parallelism）

概念

张量并行是指将同一个模型的不同部分（例如一个层的参数或计算）分割到多个 GPU 上进行计算，而不是在每个 GPU 上保留完整的模型副本。通过这种方式，超大规模模型可以分布到多个 GPU 上进行训练或推理。

工作流程

• 模型切分：模型的张量（如权重矩阵、神经网络层）按某一维度被切分，切分后的部分分布在不同的 GPU 上。例如一个 1000x1000 的矩阵可以切分成 4 个 1000x250 的子矩阵，每个 GPU 只计算其中的一部分。
• 并行计算：各 GPU 处理自己负责的部分计算，并通过通信进行必要的数据交换。
• 结果汇总：计算完成后，GPU 间交换中间结果以确保整个模型的计算一致性。

在推理过程中，GPU 之间的通信相对较少，因为模型参数是固定的。而在训练过程中，GPU 间需要频繁的同步和通信。

优点

• 处理大模型：张量并行适合处理超大模型，尤其是单个 GPU 无法容纳整个模型时。通过将模型分割，可以让每张 GPU 只处理它可以承载的那部分。
• 显存节省：每个 GPU 只存储部分模型的权重和中间激活值，因此可以处理更大规模的模型。

缺点

• 实现复杂：张量并行比数据并行复杂得多，需要对模型的结构进行拆解和重组。尤其是需要管理 GPU 间的通信，这增加了实现和调试的难度。
• 通信开销：虽然数据并行的通信开销主要在梯度同步上，但张量并行的通信开销在前向传播和反向传播中都会出现，因为 GPU 之间必须共享中间计算结果。

场景

张量并行多用于处理 超大规模模型，如 GPT-3、BERT 等大语言模型。因为这些模型的参数量非常大，单个 GPU 的显存难以完全加载整个模型，必须通过张量并行进行拆分和处理。

3. 数据并行 vs 张量并行：关键对比

特性	数据并行	张量并行
模型存储	每个 GPU 保存完整的模型副本	模型的不同部分分布在不同 GPU 上
数据处理方式	数据划分为不同的批次，分配到不同的 GPU	各 GPU 处理模型的不同部分
计算模式	每个 GPU 独立计算并汇总梯度	每个 GPU 处理部分张量，需频繁通信
适用模型规模	适用于模型较小，数据量大的情况	适用于模型较大，单个 GPU 无法容纳时
通信开销	通信主要在每轮计算后的梯度汇总中	GPU 间频繁通信，通信开销更大
实现难度	相对简单，易于在现有框架中实现	复杂，需要更多的模型和计算切分
典型应用	图像分类、大型数据集的小模型	GPT-3、BERT 等大规模语言模型

4. 混合并行（Hybrid Parallelism）

在实际应用中，尤其是处理超大规模深度学习模型时，混合并行 是常见的策略，即结合数据并行和张量并行的优势。比如，某些大型语言模型（如 GPT-3）采用了张量并行来分割单个模型层，同时使用数据并行来加速整体计算。

• 混合实现：模型层的计算可以通过张量并行分配到多个 GPU 上，而整个数据集的批次处理可以通过数据并行的方式进行分片。这样，既能处理大模型，又能最大化利用多卡并行的优势。

混合并行的一个典型例子是在一些大型训练平台上使用模型并行和数据并行结合的方式，来加速大规模模型的训练和推理。

总结

• 数据并行：通过将数据分割并分配到不同 GPU 上运行相同的模型副本来加速计算。适合模型相对较小、但数据量大的场景。
• 张量并行：通过将模型的张量（如神经网络层、矩阵）切分并分配到多个 GPU 上来分担计算。适合模型特别大时，单个 GPU 无法加载整个模型的情况。

这两种方法各有优势，选择哪种方法取决于模型的规模、数据量以及硬件资源配置等因素。

流水线并行

核心思想

流水线并行（Pipeline Parallelism）

• 核心思想：**模型按层（或模块）**划分为多个部分，每个GPU负责处理模型的某一部分。这就像在工厂的生成流水线上，每个工人（GPU）只负责产品的一部分加工步骤。
• 工作方式：输入数据在流水线中按顺序传递，各GPU按顺序处理自己负责的模型部分。多个输入样本可以同时在流水线中流动，以提高效率。
• 适用场景：适合模型规模非常大，单张GPU无法容纳整个模型的情况，例如具有多个层的深度模型（如Transformer、GPT系列）。

张量并行（Tensor Parallelism）

• 核心思想：将模型的**单个层的张量（如权重矩阵）**按某个维度切分，分配到多个GPU上。 每个GPU负责处理该层张量的一部分，多个GPU协同计算一个层的输出。
• 工作方式：在单层的计算中，多个GPU同时处理不同的张量部分，并通过通信共享计算结果，最终合成整个层的输出。
• 适用场景：适合单层的计算非常大，无法由单个GPU处理的情况。例如，具有极大参数矩阵的层，如GPT-3中的自注意力机制。

并行方式差异

特性	流水线并行（Pipeline Parallelism）	张量并行（Tensor Parallelism）
模型切分方式	将模型按层或模块切分，不同 GPU 负责不同的层	将单个层的张量按维度切分，多个 GPU 共同计算
工作流	输入数据流经不同的 GPU，各 GPU 处理模型的不同部分	单个层的计算被多个 GPU 并行处理，GPU 之间共享计算结果
通信开销	GPU 间需传递中间层的输出结果，前后层通信较少	每个前向/反向传播步骤都需 GPU 之间大量通信
计算效率	需要通过批量处理和流水线技术来提升并行度（存在流水线泡沫）	多个 GPU 同时处理单层的计算，提升并行效率
实现难度	实现相对简单，但需精心设计以减少流水线延迟	实现复杂，需要处理张量切分和大量 GPU 间通信
适合的模型类型	模型层数较多，且单层可以容纳在 GPU 内	单个层的计算非常大，无法在单个 GPU 内完成

适用场景

流水线并行适用场景

• 超大规模深层模型：例如 GPT-3、BERT 等深度学习模型。流水线并行可以通过将模型的不同层分配到不同的 GPU 上，解决模型参数量过大、单个 GPU 无法容纳的问题。
• 显存限制：当模型参数量巨大时，单个 GPU 无法存储整个模型的权重和中间激活值，流水线并行通过分层计算降低每个 GPU 的显存需求。
• 多层 Transformer：Transformer 模型常常具有几十层到上百层的结构，这样的模型非常适合流水线并行。

张量并行适用场景

• 单层参数非常大：例如 GPT-3 中的多头注意力机制，每个注意力头的计算涉及非常大的张量。张量并行可以将这种大矩阵按维度切分，分配到多个 GPU 上同时处理。
• 并行化计算量巨大：张量并行可以利用多个 GPU 共同处理一个层的计算任务，适合那些单层计算负担极重的模型。
• 显存较大但计算需求较高的模型：如果单个层的计算非常复杂，需要分布到多个 GPU 上同时进行时，张量并行是理想的选择。

注意力头数和tensor_parallel_size

1.多头数与tensor_parallel_size不兼容

多头注意力机制通常会将输入张量按照注意力头的数量切分，并在不同头之间并行执行。**多头数（number of heads）**和 tensor_parallel_size 的设置需要能够很好地匹配。特别是在张量并行中，注意力头的数量需要能够被 tensor_parallel_size 整除，否则会导致切分时出现维度不匹配的错误。

输入张量：输入数据经过embeding后的序列向量，形状为（batch_size,seq_len,hidden_size）

batch_size：输入批次的大小，即一次处理的数据样本数量。

seq_len：输入序列的长度，通常是一个自然语言的词数或子词数。

hidden_size：隐藏层的维度，即每个单词或序列元素

每个头处理的事不同的维度，比如embeding嵌入14维，2头的话，可能就是第一个头处理前7维，第二个头处理后7维

2.如何在3张卡上运行

vLLM参数

博客

Qwen2转成ONNX

博客