准备模型前向传播的输入

准备模型前向传播的输入#

目的#

执行模型前向传播所需的信息：

  • 输入

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

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

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

因此，只要拥有上述两部分信息，即可执行模型的前向传播。

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

概述#

1.获取输入#

获取输入的工作流程：

  1. 获取token positions：每个token在其请求序列中的相对位置。

  2. 获取token indices：每个已调度token在token表中的索引。

  3. 获取Token IDs：使用token索引从token id table中检索Token ID。

最后，这些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将被存储到的槽位索引。

  • attention mask：在softmax之前应用于注意力分数的掩码矩阵，用于控制哪些token可以相互关注（通常是因果注意力）。

开始之前#

主要有三种类型的变量。

  • token级别：表示每个已调度token对应的一个属性，因此该变量的长度等于已调度token的数量。

  • 请求级别：表示每个已调度请求的一个属性，其长度通常等于已调度请求的数量。（query start location是特殊情况，多一个元素。）

  • 系统级别：

    1. Token IDs table：存储每个请求的token ID（即模型的输入）。该表的形状为(max num request, max model len)。其中，max num request是前向批次中允许的最大并发请求数，max model len是该模型单个请求序列能处理的最大token数。

    2. 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 表 的形状为 (max num request, max model len)。

为什么 T_3_5、T_3_6、T_3_7 出现在此表中但未被调度？

  • 我们将一个请求序列中的所有 Token ID 一次性填入此表，但仅检索本次调度的 token。剩余的 Token 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 的 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.检索 Token ID#

我们使用 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 Table 中，我们使用第一个 block（即 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，我们可以得到 token 的 device block number。

实现 slot mapping 的工作流程：

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

    目的：对于每个 token，它可用于从 block table 中选择 device block number。

  2. 使用 block table indices 获取 device block number。

    目的：device block number 指示每个 token 属于哪个设备块。

  3. 使用 positions 和 block size 获取 block offsets。

    目的：block offsets 指示每个 token 在块内的偏移量。

  4. 使用 device block number 和 block offsets 构建 slot mapping。

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

详细信息：

  1. (Token 级别) 使用简单公式计算 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。

  2. (Token 级别) 使用 block table indices 为每个调度的 token 选出 device block number。伪代码为 block_numbers = block_table[block_table_indices]。因此 device block number=[1, 1, 2, 3, 3, 4, 4, 5, 5, 6]

  3. (Token 级别) block offsets 可通过 block offsets = positions % block size = [0, 1, 0, 0, 1, 0, 1, 0, 1, 0] 计算。

  4. 最后，使用 block offsets 和 device block number 创建 slot mapping：device block number * block size + block_offsets = [2, 3, 4, 6, 7, 8, 9, 10, 11, 12]

(请求级别) 如我们已知，调度的 token 数量为 [3, 2, 5]：

  • (请求级别) 使用前缀和计算 query start location：[0, 3, 5, 10]。

  • (请求级别) 步骤 1 中的所有 token 均处于 prefill 阶段，已计算的 token 数量为 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

  • (Token 级别) slot mapping：[2, 3, 4, 6, 7, 8, 9, 10, 11, 12]

  • attention mask：对于所有发起 prefill 过程的请求，我们仅创建一个掩码矩阵以便在不同请求间复用。该掩码矩阵的形状为 5 * 5：

步骤 2：分块预填充#

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

获取输入#

每个请求的调度 token：{'0': 1, '1': 1, '2': 3}

  1. request indices：[0, 1, 2, 2, 2]

  2. token positions：[3, 2, 5, 6, 7]

    当前 Token IDs 表：

    | 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 的新 Token ID。它们是从模型输出中采样得到的。

  3. token indices：[3, 14, 29, 30, 31]

  4. Input IDs：[T_0_3, T_1_2, T_3_5, T_3_6, T_3_7]

构建输入注意力元数据#

我们将块 7 和 8 分别分配给 request_1 和 request_2，因为它们在 token 生成或分块预填充后需要在设备上占用更多空间来存储 KV 缓存。

当前 Block Table：

| 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. (Token 级别) 块表索引: [1, 7, 14, 15, 15]

  2. (Token 级别) 设备块编号: [2, 7, 6, 8, 8]

  3. (Token 级别) 块偏移量: [1, 0, 1, 0, 1]

  4. (Token 级别) 槽位映射: [5, 14, 13, 16, 17]

已调度 token 数量: [1, 1, 3]

  • 查询起始位置: [0, 1, 2, 5]

  • 序列长度: [4, 3, 8]

  • 已计算 token 数量: [3, 2, 5]

  • number of requests：3

  • 最大查询长度: 3

  • 槽位映射: [5, 14, 13, 16, 17]

  • 注意力掩码: 5 * 8

    每个 token 有一个 1 * 8 的向量，共有 5 个已调度 token。

最后#

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

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