跳转至

为模型前向传播准备输入

目的

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

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

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

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

因此，只要我们拥有了上述这两类信息，就可以执行模型的前向传播。

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

概述

1. 获取输入

获取输入的工作流程：

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

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

  3. 获取 Token IDs：使用 token indices 从 token id table（token ID 表） 中检索 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：将每个 block 的逻辑地址（在其序列内）转换为设备内存中的全局物理地址。
  • max query len：本请求批次中最长的已调度 token 长度。
  • slot mapping：每个输入 token 将要存储到的槽位索引。
  • attention mask：在 softmax 之前应用于注意力分数的掩码矩阵，用于控制哪些 token 之间可以相互关注（通常是因果注意力）。

开始之前

变量主要分为三种类型。

  • token 级别：表示每个已调度 token 对应的一个属性，因此该变量的长度是已调度 token 的数量。
  • request 级别：表示每个已调度请求的一个属性，其长度通常是已调度请求的数量。（query start location 是一个特例，它多一个元素。）
  • 系统级别：
  • Token ID 表：存储每个请求的 token ID（即模型的输入）。该表的形状为 (max num request, max model len)。其中，max num request 是前向批次中允许的最大并发请求数，max model len 是在该模型的一个请求序列中能够处理的最大 token 数量。
  • Block table：将每个 block 的逻辑地址（在其序列内）转换为设备内存中的全局物理地址。该表的形状为 (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 个请求，它们的 prompt 长度分别为 3、2 和 8。
  • max model length：12（模型中的一个请求序列可以处理的最大 token 数量）。

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

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

获取输入

由于一次可调度的最大 token 数量为 10，因此每个请求的已调度 token 可以表示为 {'0': 3, '1': 2, '2': 5}。请注意，request_2 使用了 chunked prefill，因此还剩 3 个 prompt token 未被调度。

1. 获取 token 位置

首先，确定每个 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）。

最终，token positions 可以得到为 [0, 1, 2, 0, 1, 0, 1, 2, 3, 4]。该变量属于 token 级别。

2. 获取 token 索引

当前 Token ID 表 的形状为 (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 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）来标记未使用的块。该块的形状为 (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。

实现槽位映射的工作流程如下：

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

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

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

用途：device block number 表示每个 token 属于哪个设备块。

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

用途：block offsets 表示每个 token 在块内的偏移量。

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

用途：我们可以使用 slot mapping 将 Token IDs 存储到 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 过程的请求，我们仅创建一个 mask 矩阵，以便在不同请求之间复用。该 mask 矩阵的形状为 5 * 5。

步骤 2：分块 Prefill

在步骤 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 IDs。它们是从模型的输出中采样得到的。

  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 生成或分块 prefill 之后需要更多设备空间来存储 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. (Token 级别) block table indices：[1, 7, 14, 15, 15]
  2. (Token 级别) device block number：[2, 7, 6, 8, 8]
  3. (Token 级别) block offsets：[1, 0, 1, 0, 1]
  4. (Token 级别) slot mapping：[5, 14, 13, 16, 17]

调度的 token 数量：[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

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

最后

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

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

本页原文 Markdown：在 AtomGit 查看·内容源自开源项目 vl/vllm-ascend