为模型前向传播准备输入#
目的#
执行模型前向传播所需的信息：
输入
输入对应的注意力元数据
下图展示了我们需要为模型推理准备的内容。
+---------------+
inputs --> | |
| model | --> output
attn_meta --> | |
+---------------+
因此，只要我们拥有上述两方面的信息，就可以执行模型的前向传播。
本文将解释我们如何获取输入及其对应的注意力元数据。
概述#
1.获取输入#
获取输入的工作流程：
获取
token positions：每个 token 在其请求序列中的相对位置。获取
token indices：每个已调度 token 在 token 表中的索引。获取
Token IDs：使用 token indices 从 token id table 中检索 Token IDs。
最后，这些 Token IDs 需要输入到模型中，positions 也需要送入模型以创建 Rope（旋转位置编码）。两者共同构成模型的输入。
注意：Token IDs 是模型的输入，因此我们也称它们为 Input IDs。
2.构建输入注意力元数据#
模型在前向传播过程中需要以下注意力元数据：
query start location：每个请求对应的已调度 token 的起始和结束位置。sequence length：每个请求的长度，包括已计算 token 和新调度的 token。number of computed tokens：每个请求已计算 token 的数量。number of requests：本批次中的请求数量。number of tokens：本批次中已调度 token 的总数。block table：将每个块在其序列内的逻辑地址转换为其在设备内存中的全局物理地址。max query len：本请求批次中最长的已调度 token 长度。slot mapping：输入 token 将被存储到的每个 token 的索引。attention mask：在 softmax 之前应用于注意力分数的掩码矩阵，用于控制哪些 token 可以相互关注（通常是因果注意力）。
开始之前#
主要有三种类型的变量。
token 级别：代表每个已调度 token 对应的一个属性，因此该变量的长度等于已调度 token 的数量。
请求级别：代表每个已调度请求的一个属性，其长度通常等于已调度请求的数量。（
query start location是一个特例，它多一个元素。）系统级别：
Token IDs table：存储每个请求的 token IDs（即模型的输入）。此表的形状为
(max num request, max model len)。其中，max num request是前向批次中允许的最大并发请求数，max model len是该模型中单个请求序列可以处理的最大 token 数量。Block table：将每个块在其序列内的逻辑地址转换为其在设备内存中的全局物理地址。此表的形状为
(max num request, max model len / block size)
注意：这两个表都来自 准备输入 之前的 _update_states 方法。如果需要更多启发，可以查看一下。
提示#
简而言之，一个 token ID 是一个整数（通常是 int32），它代表一个 token。Token ID 示例：
| Token ID | Token |
|--------------|---------------|
| 0 | [PAD] |
| 1 | <|endoftext|> |
| 2 | <|start|> |
| 3 | [SEP] |
| 4 | I |
| 5 | the |
| 6 | be |
| 7 | of |
| 8 | and |
| ... | ... |
| ... | ... |
| vocab_size-1 | <|im_end|> |
深入细节#
假设：
一次可调度的最大 token 数：10
block size：2总共调度 3 个请求。它们的提示长度分别为 3、2 和 8。
max model length：12（模型中单个请求序列可以处理的最大 token 数量）。
这些假设是在启动 vLLM 时配置的。它们不是固定的，因此可以手动设置。
步骤 1：所有请求均处于预填充阶段#
获取输入#
由于一次可调度的最大 token 数为 10，每个请求的已调度 token 可以表示为 {'0': 3, '1': 2, '2': 5}。注意 request_2 使用了分块预填充，留下了 3 个提示 token 未调度。
1.获取 token positions#
首先，确定每个 token 属于哪个请求：token 0–2 分配给 request_0，token 3–4 分配给 request_1，token 5–9 分配给 request_2。为了表示这种映射，我们使用 request indices，例如，request indices：[0, 0, 0, 1, 1, 2, 2, 2, 2, 2]。
对于每个请求，使用 已计算 token 的数量 + 当前调度 token 的相对位置（request_0: [0 + 0, 0 + 1, 0 + 2]，request_1: [0 + 0, 0 + 1]，request_2: [0 + 0, 0 + 1,..., 0 + 4]），然后将它们连接在一起（[0, 1, 2, 0, 1, 0, 1, 2, 3, 4]）。
注意：在实际代码中，有一种更高效的方法（使用 request indices）来创建 positions。
最后，token positions 可以获取为 [0, 1, 2, 0, 1, 0, 1, 2, 3, 4]。此变量是 token 级别 的。
2.获取 token indices#
当前 Token IDs table 的形状为 (max num request, max model len)。
为什么表中的 T_3_5、T_3_6、T_3_7 没有被调度？
我们将一个请求序列中的所有 Token IDs 一次性填充到此表中，但我们只检索本次调度的 token。然后下次再检索剩余的 Token IDs。
| T_0_0 | T_0_1 | T_0_2 | ? | ? | ? | ? | ? | ? | ? | ? | ? |
| T_1_0 | T_1_1 | ? | ? | ? | ? | ? | ? | ? | ? | ? | ? |
| T_2_0 | T_2_1 | T_3_2 | T_3_3 | T_3_4 | T_3_5 | T_3_6 | T_3_7 | ? | ? | ? | ? |
| ? | ? | ? | ? | ? | ? | ? | ? | ? | ? | ? | ? |
......
......
......
注意 T_x_x 是一个 int32。
假设 M = max model len。那么我们可以使用 token positions 以及每个 token 的 request indices 来构造 token indices。
所以 token indices = [0 + 0 * M, 1 + 0 * M, 2 + 0 * M, 0 + 1 * M, 1 + 1 * M, 0 + 2 * M, 1 + 2 * M, 2 + 2 * M, 3 + 2 * M, 4 + 2 * M] = [0, 1, 2, 12, 13, 24, 25, 26, 27, 28]
3.检索 Token IDs#
我们使用 token indices 从 token 表中选择出对应的 Input IDs。伪代码如下：
input_ids = token_table[token_indices]
如前所述，我们将这些 Token IDs 称为 Input IDs。
Input IDs=[T_0_0, T_0_1, T_0_2, T_1_0, T_1_1, T_2_0, T_2_1, T_3_2, T_3_3, T_3_4]
构建输入注意力元数据#
在当前的块表中，我们使用第一个块（即 block_0）来标记未使用的块。块的形状为 (最大请求数, 最大模型长度 / 块大小)，其中 最大模型长度 / 块大小 = 12 / 2 = 6。
| 1 | 2 | 0 | 0 | 0 | 0 |
| 3 | 0 | 0 | 0 | 0 | 0 |
| 4 | 5 | 6 | 0 | 0 | 0 |
| 0 | 0 | 0 | 0 | 0 | 0 |
......
......
......
设备内存中的 KV 缓存块如下所示：
| 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | ......
假设 K = 最大模型长度 / 块大小 = 6，我们可以得到令牌的设备块编号。
实现槽映射的工作流程：
使用
K、positions和request indices获取块表索引。目的：对于每个令牌，它可用于从
块表中选择设备块编号。使用
块表索引获取设备块编号。目的：
设备块编号指示每个令牌属于哪个设备块。使用
positions和block size获取块内偏移。目的：
块内偏移指示每个令牌在块内的偏移量。使用
设备块编号和块内偏移构建槽映射。目的：我们可以使用
槽映射将令牌 ID 存储到令牌槽中。
详细信息：
(令牌级别) 使用一个简单的公式计算
块表索引：request indices * K + positions / block size。因此它等于[0 * 6 + 0 / 2, 0 * 6 + 1 / 2, 0 * 6 + 2 / 2, 1 * 6 + 0 / 2, 1 * 6 + 1 / 2, 2 * 6 + 0 / 2, 2 * 6 + 1 / 2, 2 * 6 + 2 / 2, 2 * 6 + 3 / 2, 2 * 6 + 4 / 2] = [0, 0, 1, 6, 6, 12, 12, 13, 13, 14]。这可用于从块表中选择设备块编号。(令牌级别) 使用
块表索引为每个已调度的令牌选择出设备块编号。伪代码为block_numbers = block_table[block_table_indices]。因此设备块编号=[1, 1, 2, 3, 3, 4, 4, 5, 5, 6](令牌级别)
块内偏移可以通过block offsets = positions % block size = [0, 1, 0, 0, 1, 0, 1, 0, 1, 0]计算得出。最后，使用
块内偏移和设备块编号创建槽映射：设备块编号 * 块大小 + 块内偏移 = [2, 3, 4, 6, 7, 8, 9, 10, 11, 12]
(请求级别) 已知已调度的令牌数量为 [3, 2, 5]：
(请求级别) 使用前缀和计算
查询起始位置：[0, 3, 5, 10]。(请求级别) 步骤 1 中的所有令牌都处于预填充阶段，已计算的令牌数量为 0；因此
序列长度=[3, 2, 5]。(请求级别) 如上所述，
已计算令牌数均为 0：[0, 0, 0]。请求数量：3(请求级别)
令牌数量：[3, 2, 5]最大查询长度：5(令牌级别)
槽映射：[2, 3, 4, 6, 7, 8, 9, 10, 11, 12]注意力掩码：对于所有发起预填充过程的请求，我们仅创建一个掩码矩阵，以便在不同请求间复用。该掩码矩阵的形状为5 * 5：
步骤 2：分块预填充#
在步骤 2 中，我们不再提供解释或进行计算；而是直接呈现最终结果。
获取输入#
每个请求的已调度令牌：{'0': 1, '1': 1, '2': 3}
请求索引：[0, 1, 2, 2, 2]令牌位置：[3, 2, 5, 6, 7]当前令牌 ID 表：
| T_0_0 | T_0_1 | T_0_2 | T_0_3 | ? | ? | ? | ? | ? | ? | ? | ? | | T_1_0 | T_1_1 | T_1_2 | ? | ? | ? | ? | ? | ? | ? | ? | ? | | T_2_0 | T_2_1 | T_3_2 | T_3_3 | T_3_4 | T_3_5 | T_3_6 | T_3_7 | ? | ? | ? | ? | | ? | ? | ? | ? | ? | ? | ? | ? | ? | ? | ? | ? | ...... ...... ......
注意：T_0_3、T_1_2 分别是 request_0 和 request_1 的新令牌 ID。它们是从模型输出中采样得到的。
令牌索引：[3, 14, 29, 30, 31]输入 ID：[T_0_3, T_1_2, T_3_5, T_3_6, T_3_7]
构建输入注意力元数据#
我们将块 7 和 8 分别分配给 request_1 和 request_2，因为它们在令牌生成或分块预填充后需要更多设备空间来存储 KV 缓存。
当前块表：
| 1 | 2 | 0 | 0 | 0 | 0 |
| 3 | 7 | 0 | 0 | 0 | 0 |
| 4 | 5 | 6 | 8 | 0 | 0 |
| 0 | 0 | 0 | 0 | 0 | 0 |
......
......
......
设备内存中的 KV 缓存块：
| 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | ......
(令牌级别)
块表索引：[1, 7, 14, 15, 15](令牌级别)
设备块编号：[2, 7, 6, 8, 8](令牌级别)
块内偏移：[1, 0, 1, 0, 1](令牌级别)
槽映射：[5, 14, 13, 16, 17]
已调度令牌数量：[1, 1, 3]
查询起始位置：[0, 1, 2, 5]序列长度：[4, 3, 8]已计算令牌数：[3, 2, 5]请求数量：3最大查询长度：3槽映射：[5, 14, 13, 16, 17]注意力掩码：5 * 8每个令牌有一个
1 * 8的向量，共有 5 个已调度的令牌。
最后#
如果您理解了步骤 1 和步骤 2，您就会知道所有后续步骤。
希望本文档能帮助您更好地理解 vLLM 如何为模型前向传播准备输入。如果您有任何好的想法，欢迎向我们贡献。