KV缓存的管理方法、调度器、介质、设备及产品与流程

本公开涉及计算机，具体地，涉及一种kv缓存的管理方法、调度器、介质、设备及产品。

背景技术：

1、大语言模型的推理过程，是指利用已经训练好的大语言模型进行推理，其中，用户输入问题信息，大语言模型可根据问题信息进行推理，并输出对应的答案信息，输入的问题信息和输出的答案信息合称为上下文（context）。

2、transformer模型在大语言模型（llm，large language model）中扮演着重要角色，其核心模块是注意力机制。大语言模型的推理过程可分为预填充（prefill）和解码（decoding）两个阶段，在prefill阶段，大语言模型处理输入文本prompt，生成第一个token和kv缓存，在decoding阶段，大语言模型根据上一个生成的token和历史kv缓存，逐个生成输出token。随着上下文的增长，大语言模型推理所需的计算量会呈线性增长，因此，目前采用kv缓存来保存已经得到的kv矩阵，以避免重复计算。然而，kv缓存占用的内存空间也会随着上下文长度的增长而呈线性增加，在长序列推理、多并行推理任务的场景下，很容易出现kv缓存爆炸问题。

技术实现思路

1、本公开的目的是提供一种kv缓存的管理方法、调度器、介质、设备及产品，可以降低推理系统的总成本，避免主节点在进行长序列推理时kv缓存的爆炸问题。

2、为了实现上述目的，第一方面，本公开提供一种transformer模型中kv缓存的管理方法，应用于基于transformer模型的推理系统中的调度器，所述推理系统包括所述调度器、m个主节点和n个从节点，所述m个主节点和所述n个从节点用于执行推理任务，所述m个主节点用于加载所述transformer模型，所述主节点的计算能力高于所述从节点的计算能力，m≥1，n≥1；所述方法包括：

3、若监测到所述主节点在执行所述推理任务的过程中，存在满足迁移条件的推理序列，则从所述n个从节点中确定目标从节点，其中，所述迁移条件为所述推理序列包括的token的数量达到第一数量；

4、分别向各个所述主节点下发迁移指令，所述迁移指令用于指示所述主节点将第一kv缓存发送至所述目标从节点，所述第一kv缓存为所述主节点生成的第二数量个连续的token对应的kv缓存，所述第二数量个连续的token为所述推理序列包括的token，所述第二数量小于或等于所述第一数量；

5、向所述目标从节点下发接收指令，所述接收指令用于指示所述目标从节点接收并存储各个所述主节点分别发送的所述第一kv缓存。

6、第二方面，本公开提供一种调度器，其特征在于，应用于基于transformer模型的推理系统，所述推理系统包括所述调度器、m个主节点和n个从节点，所述m个主节点和所述n个从节点用于执行推理任务，所述m个主节点用于加载所述transformer模型，所述主节点的计算能力高于所述从节点的计算能力，m≥1，n≥1；所述调度器包括：

7、确定模块，用于若监测到所述主节点在执行所述推理任务的过程中，存在满足迁移条件的推理序列，则从所述n个从节点中确定目标从节点，其中，所述迁移条件为所述推理序列包括的token的数量达到第一数量；

8、第一指令发送模块，用于分别向各个所述主节点下发迁移指令，所述迁移指令用于指示所述主节点将第一kv缓存发送至所述目标从节点，所述第一kv缓存为所述主节点生成的第二数量个连续的token对应的kv缓存，所述第二数量个连续的token为所述推理序列包括的token，所述第二数量小于或等于所述第一数量；

9、第二指令发送模块，用于向所述目标从节点下发接收指令，所述接收指令用于指示所述目标从节点接收并存储各个所述主节点分别发送的所述第一kv缓存。

10、第三方面，本公开提供一种计算机可读存储介质，其上存储有计算机程序，该程序被处理器执行时实现本公开第一方面提供的所述方法的步骤。

11、第四方面，本公开提供一种电子设备，包括：

12、存储器，其上存储有计算机程序；

13、处理器，用于执行所述存储器中的所述计算机程序，以实现本公开第一方面提供的所述方法的步骤。

14、第五方面，本公开提供一种计算机程序产品，包括计算机程序，该计算机程序被处理器执行时实现本公开第一方面提供的所述方法的步骤。

15、通过上述技术方案，基于transformer模型的推理系统可包括调度器、m个主节点和n个从节点，主节点的计算能力高于从节点的计算能力，m个主节点用于加载transformer模型，通过采用主节点和从节点异构的配置方式，可以降低推理系统的总成本，并且提高资源利用率。调度器若监测到主节点在执行推理任务的过程中，存在满足迁移条件的推理序列，可分别向各个主节点下发迁移指令，其中，迁移条件为推理序列包括的token的数量达到第一数量，这样，可以在主节点执行的推理任务中存在长序列的情况下，将主节点生成的kv缓存迁移至从节点中，并且，从主节点中迁移出的第一kv缓存，为推理序列包括的第二数量个连续的token对应的kv缓存，可以将连续token对应的完整的kv缓存迁移至目标从节点，便于在目标从节点中基于token进行kv缓存的寻址。如此，降低主节点中kv缓存的内存占用量，避免主节点在进行长序列推理时kv缓存的爆炸问题，保证推理任务的正常执行。

16、本公开的其他特征和优点将在随后的具体实施方式部分予以详细说明。

技术特征：

1.一种transformer模型中kv缓存的管理方法，其特征在于，应用于基于transformer模型的推理系统中的调度器，所述推理系统包括所述调度器、m个主节点和n个从节点，所述m个主节点和所述n个从节点用于执行推理任务，所述m个主节点用于加载所述transformer模型，所述主节点的计算能力高于所述从节点的计算能力，m≥1，n≥1；所述方法包括：

2.根据权利要求1所述的方法，其特征在于，所述m个主节点执行所述推理任务采用的并行策略包括张量并行和流水线并行，m为张量并行度ntp与流水线并行度npp的乘积；

3.根据权利要求2所述的方法，其特征在于，所述目标从节点包括ntp个第二存储单元组，所述第二存储单元组包括所述transformer模型中每一层分别对应的第二存储单元；

4.根据权利要求2所述的方法，其特征在于，所述目标从节点包括每一所述主节点分别对应的第三存储单元组，第i个所述主节点对应的所述第三存储单元组中，包括第m层至第n层中每一层分别对应的第三存储单元；

5.根据权利要求1所述的方法，其特征在于，所述从所述n个从节点中确定目标从节点，包括：

6.根据权利要求1所述的方法，其特征在于，所述方法还包括：

7.一种调度器，其特征在于，应用于基于transformer模型的推理系统，所述推理系统包括所述调度器、m个主节点和n个从节点，所述m个主节点和所述n个从节点用于执行推理任务，所述m个主节点用于加载所述transformer模型，所述主节点的计算能力高于所述从节点的计算能力，m≥1，n≥1；所述调度器包括：

8.一种计算机可读存储介质，其上存储有计算机程序，其特征在于，该程序被处理器执行时实现权利要求1-6中任一项所述方法的步骤。

9.一种电子设备，其特征在于，包括：

10.一种计算机程序产品，包括计算机程序，其特征在于，该计算机程序被处理器执行时实现权利要求1-6中任一项所述方法的步骤。

技术总结
本公开涉及一种KV缓存的管理方法、调度器、介质、设备及产品，所述方法包括：若监测到主节点在执行推理任务的过程中，存在满足迁移条件的推理序列，则从N个从节点中确定目标从节点；分别向各个主节点下发迁移指令，迁移指令用于指示主节点将第一KV缓存发送至目标从节点，第一KV缓存为主节点生成的第二数量个连续的token对应的KV缓存；向目标从节点下发接收指令，接收指令用于指示目标从节点接收并存储各个主节点分别发送的第一KV缓存。通过上述技术方案，采用主节点和从节点异构的配置方式，可以降低推理系统的总成本，将主节点生成的KV缓存迁移至从节点中，避免主节点在进行长序列推理时KV缓存的爆炸问题。

技术研发人员：陈杰,张楚文,阚宏伟
受保护的技术使用者：新紫光集团有限公司
技术研发日：
技术公布日：2025/3/18