为模型推理准备输入

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

目的#

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

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

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

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

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

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

概述#

1. 获取输入#

获取输入的流程

获取 Token Positions：每个 Token 在其请求序列内的相对位置。
获取 Token Indices：每个已调度 Token 在 Token 表中的索引。
获取 Token IDs：使用 Token 索引从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：将每个块的逻辑地址（在其序列内）转换为设备内存中的全局物理地址。
Max Query Len：请求批次中最长的已调度 Token 长度。
Slot Mapping：输入 Token 将存储到其中的每个 Token 的索引。
Attention Mask：在 Softmax 之前应用于注意力分数的掩码矩阵，用于控制哪些 Token 可以相互关注（通常是因果注意力）。

开始之前#

主要有三种类型的变量。

Token 级别：表示与每个已调度 Token 对应的某个属性，因此此变量的长度是已调度 Token 的数量。
Request 级别：表示每个已调度请求的某个属性，其长度通常是已调度请求的数量。（Query Start Location 是一个特例，多一个元素）
System 级别
1. Token ID 表：存储每个请求的 Token ID（即模型的输入）。该表的形状为 (max num request, max model len)。其中，max num request 是在前向批次中允许的最大并发请求数，max model len 是模型中每个请求序列可以处理的最大 Token 数。
2. 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 个请求。它们的提示长度分别为 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 位置：#

首先，确定每个 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 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 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 表中，我们使用第一个块（即 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 的流程

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

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

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

目的：block offsets 指示 Token 在块内的偏移量。
使用 device block number 和 block offsets 构建 slot mapping。

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

细节

(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。


(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]
(Token 级别) 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]


(Request 级别) 因为我们知道已调度 Token 的数量为 [3, 2, 5]

(Request 级别) 使用前缀和计算 query start location：[0, 3, 5, 10]。
(Request 级别) 步骤 1 中的所有 Token 都处于预填充阶段，已计算 Token 的数量为 0；因此 sequence length = [3, 2, 5]。
(Request 级别) 如上所述，number of computed tokens 都为 0：[0, 0, 0]。
Number of Requests：3
(Request 级别) 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：对于所有开始预填充过程的请求，我们只创建一个掩码矩阵，供不同请求重用。此掩码矩阵的形状为 5 * 5。



步骤 2：分块预填充#
在步骤 2 中，我们不再提供解释或进行计算，而是直接给出最终结果。

获取输入#
每个请求的已调度 Token：{'0': 1, '1': 1, '2': 3}

Request Indices：[0, 1, 2, 2, 2]
Token Positions：[3, 2, 5, 6, 7]

当前的Token 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的新 Token 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，因为它们在 Token 生成或分块预填充后需要在设备中存储更多 KV 缓存空间。
当前的Block 表
| 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 | ...... 



(Token 级别) Block Table Indices：[1, 7, 14, 15, 15]
(Token 级别) Device Block Number：[2, 7, 6, 8, 8]
(Token 级别) Block Offsets：[1, 0, 1, 0, 1]
(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 如何为模型前向传播准备输入。如果您有任何好的想法，欢迎贡献给我们。


              

              
              
              
              
                
                  

       previous
 vLLM Ascend 中的 Patch
 
    next
 Disaggregated-prefill


            
            
              
                


  
  
     目录
  
    
目的
概述
1. 获取输入
2. 构建输入的注意力元数据


开始之前
提示


深入了解细节
步骤 1：所有请求都处于预填充阶段
获取输入
1. 获取 Token 位置
2. 获取 Token 索引
3. 检索 Token ID


构建输入的注意力元数据


步骤 2：分块预填充
获取输入
构建输入的注意力元数据




最后


          
            

  
  
    
由 vllm-ascend 团队提供

  
  
  
    

  © 版权所有 2025, vllm-ascend 团队。