模型前向计算输入准备

模型前向计算输入准备#

目的#

执行模型前向计算所需的信息：

输入数据
输入对应的注意力元数据

下图展示了模型推理需要准备的内容。

              +---------------+
  inputs  --> |               |
              |     model     |  --> output
attn_meta --> |               |
              +---------------+  

因此，只要拥有上述两部分信息，我们就能执行模型的前向计算。

本文档将解释如何获取输入及其对应的注意力元数据。

概述#

1. Obtain inputs#

获取输入的工作流程：

获取 token positions（词元位置）：每个词元在其请求序列中的相对位置。
获取 token indices（词元索引）：每个被调度词元在词元表中的索引。
获取 Token IDs：使用词元索引从词元ID表中检索出对应的词元ID。

最后，这些 Token IDs 需要输入到模型中，同时 positions 也需要送入模型以生成 Rope（旋转位置编码）。两者共同构成模型的输入。

注意：Token IDs 是模型的输入，因此我们也称其为 Input IDs。

2. Build inputs attention metadata#

模型在前向计算过程中需要以下注意力元数据：

query start location：每个请求对应的被调度词元的起始和结束位置。
sequence length：每个请求的长度，包括已计算的词元和新调度的词元。
number of computed tokens：每个请求已计算的词元数量。
number of requests：本轮次中的请求数量。
number of tokens：本轮次中被调度词元的总数。
block table：将每个块的逻辑地址（在其序列内）转换为其在设备内存中的全局物理地址。
max query len：本轮次请求中最长的被调度词元长度。
slot mapping：每个输入词元将要存储到的位置索引。
attention mask：在softmax之前应用于注意力分数的掩码矩阵，用于控制哪些词元可以相互关注（通常是因果注意力）。

开始之前#

主要有三种类型的变量。

词元级别：表示每个被调度词元对应的一个属性，因此该变量的长度等于被调度词元的数量
请求级别：表示每个被调度请求的一个属性，其长度通常等于被调度请求的数量。（query start location 是一个特例，它多一个元素）
系统级别：
1. 词元ID表：存储每个请求的词元ID（即模型的输入）。该表的形状为 (max num request, max model len)。其中，max num request 是前向批次中允许的最大并发请求数，max model len 是该模型中单个请求序列能处理的最大词元数量。
2. 块表：将每个块的逻辑地址（在其序列内）转换为其在设备内存中的全局物理地址。该表的形状为 (max num request, max model len / block size)

注意：这两个表都来自准备输入之前的 _update_states 方法。如果需要更多灵感，可以查看该方法。

提示#

简单来说，token ID 是一个整数（通常是 int32），代表一个词元。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|>    |

深入细节#

假设条件：

一次可调度的最大词元数：10
block size：2
总共调度3个请求。它们的提示词长度分别为3、2和8。
max model length：12（模型中单个请求序列能处理的最大词元数量）。

这些假设在启动vLLM时配置。它们不是固定的，因此您可以手动设置。

步骤1：所有请求都处于预填充阶段#

获取输入#

由于可调度的最大词元数为10，各请求的被调度词元数可表示为 {'0': 3, '1': 2, '2': 5}。请注意 request_2 使用了分块预填充，还有3个提示词元未被调度。

1. Get token positions:#

首先，确定每个词元属于哪个请求：词元0-2分配给 request_0，词元3-4分配给 request_1，词元5-9分配给 request_2。我们用 request indices 表示这种映射关系，例如：request indices：[0, 0, 0, 1, 1, 2, 2, 2, 2, 2]。

对于每个请求，使用已计算的词元数 + 当前调度词元的相对位置（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）来创建位置。

最终，token positions 为 [0, 1, 2, 0, 1, 0, 1, 2, 3, 4]。此变量为词元级别。

2. Get token indices:#

当前词元ID表的形状为 (max num request, max model len)。

为什么 T_3_5、T_3_6、T_3_7 在表中却未被调度？

我们将一个请求序列中的所有词元ID一次性填入此表，但本次只检索被调度的词元。剩余的词元ID将在下次检索。

| 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 和每个词元的 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. Retrieve the Token IDs#

我们使用 token indices 从词元表中选出对应的 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）标记未使用的块。块的形状为 (max num request, max model len / block size)，其中 max model len / block size = 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 = max model len / block size = 6，我们可以得到词元的 device block number。

获取槽位映射的工作流程：

使用 K、positions 和 request indices 获取 block table indices。

目的：对于每个词元，它可用于从 block table 中选择 device block number。
使用 block table indices 获取 device block number。

目的：device block number 指示每个词元属于哪个设备块。
使用 positions 和 block size 获取 block offsets。

目的：block offsets 指示每个词元在块内的偏移量。
使用 device block number 和 block offsets 构造 slot mapping。

目的：我们可以使用 slot mapping 将词元ID存储到词元槽中。

详细说明：

（词元级别）使用简单公式计算 block table indices：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 table 中选择 device block number。
（词元级别）使用 block table indices 为每个被调度词元选择 device block number。伪代码为 block_numbers = block_table[block_table_indices]。因此 device block number=[1, 1, 2, 3, 3, 4, 4, 5, 5, 6]
（词元级别）block offsets 可通过 block offsets = positions % block size = [0, 1, 0, 0, 1, 0, 1, 0, 1, 0] 计算。
最后，使用 block offsets 和 device block number 创建 slot mapping：device block number * block size + block_offsets = [2, 3, 4, 6, 7, 8, 9, 10, 11, 12]

（请求级别）已知被调度词元数为 [3, 2, 5]：

（请求级别）使用前缀和计算 query start location：[0, 3, 5, 10]。
（请求级别）步骤1中的所有词元都处于预填充阶段，已计算词元数为0；因此 sequence length = [3, 2, 5]。
（请求级别）如上所述，number of computed tokens 全为0：[0, 0, 0]。
number of requests：3
（请求级别）number of tokens：[3, 2, 5]
max query len：5
（词元级别）slot mapping：[2, 3, 4, 6, 7, 8, 9, 10, 11, 12]
attention mask：对于所有发起预填充过程的请求，我们简单地创建一个掩码矩阵供不同请求复用。该掩码矩阵的形状为 5 * 5：

步骤2：分块预填充#

在步骤2中，我们不再提供解释或进行计算；而是直接呈现最终结果。

获取输入#

各请求的被调度词元数：{'0': 1, '1': 1, '2': 3}

request indices：[0, 1, 2, 2, 2]
token positions：[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。它们是从模型输出中采样得到的。

token indices：[3, 14, 29, 30, 31]
Input IDs：[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 | ...... 

（词元级别）block table indices：[1, 7, 14, 15, 15]
（词元级别）device block number：[2, 7, 6, 8, 8]
（词元级别）block offsets：[1, 0, 1, 0, 1]
（词元级别）slot mapping：[5, 14, 13, 16, 17]

被调度词元数：[1, 1, 3]

query start location：[0, 1, 2, 5]
sequence length：[4, 3, 8]
number of computed tokens：[3, 2, 5]
number of requests：3
max query len：3
slot mapping：[5, 14, 13, 16, 17]
attention mask：5 * 8

每个词元有一个 1 * 8 的向量，总共有5个被调度词元。

最后#

如果你理解了步骤1和步骤2，就会知道所有后续步骤。

希望本文档能帮助你更好地理解vLLM如何为模型前向计算准备输入。如果你有任何好想法，欢迎提供给我们。

模型前向计算输入准备

目录

模型前向计算输入准备#

目的#

概述#

1. Obtain inputs#

2. Build inputs attention metadata#

开始之前#

提示#

深入细节#

步骤1：所有请求都处于预填充阶段#

获取输入#

1. Get token positions:#

2. Get token indices:#

3. Retrieve the Token IDs#

构建输入注意力元数据#

步骤2：分块预填充#

获取输入#

构建输入注意力元数据#

最后#