一种相似性感知的重要前缀KV投机预取方法及系统

本发明涉及计算机科学人工智能系统领域，具体涉及一种相似性感知的重要前缀kv投机预取方法及系统。

背景技术：

1、基于transformer的大语言模型在推理时通常包括预填充阶段和解码阶段。对于包含长上下文的请求，预填充阶段需要为输入前缀计算每一层的键向量和值向量，即kv缓存。检索增强生成、多轮对话、工具调用和少样本提示等应用中，不同请求之间经常共享相同或部分相同的前缀，因此将共享前缀的kv缓存保存并在后续请求中复用，可以减少重复计算并降低首个token生成延迟。

2、随着前缀长度、并发请求数和模型规模增加，gpu显存和cpu内存难以容纳全部可复用前缀kv缓存。为扩大缓存容量，现有系统可将前缀kv缓存存放在外存中，并在请求到达时将相应层的kv缓存加载到gpu显存。然而，外存到gpu显存的数据传输带宽和延迟成为新的瓶颈，尤其当长前缀kv缓存需要逐层加载时，i/o等待可能占据首个token生成延迟的主要部分。

3、为了减少加载量，现有技术可只加载对当前请求较重要的前缀kv缓存。该类方法通常根据当前层查询向量与前缀键向量之间的注意力权重确定重要token，再读取重要token对应的kv。该策略能够降低数据传输量，但也引入新的依赖关系：第i+1层的重要令牌集合需要在第i层计算完成并得到第i+1层输入后才能确定，因此下一层kv缓存无法在当前层计算期间提前预取，导致选择性加载与预取难以同时发挥作用。并且，重要kv的选择性加载暴露在gpu推理的关键路径上，导致gpu空闲，利用率低。

4、因此，需要一种能够利用模型层间特征，在不破坏推理准确性的情况下提前预取下一层重要前缀kv缓存的技术方案，以减少关键路径上的i/o等待并改善大语言模型长前缀推理的延迟。

技术实现思路

1、针对现有技术的不足，本发明的目的在于提供一种相似性感知的重要前缀kv投机预取方法及系统。

2、本发明的目的是通过以下技术方案来实现的：一种相似性感知的重要前缀kv投机预取方法，包括如下步骤：

3、步骤一，接收当前请求，在前缀索引结构中匹配当前请求对应的可复用前缀，并获得所述可复用前缀在各transformer层中的kv元数据及存储位置；

4、步骤二，对于transformer层中的第i层，获取当前请求在该层对应的重要前缀token索引集合s_{i}，并加载所述第i层计算所需的重要kv；

5、步骤三，在第i层开始计算后，基于相邻transformer层的重要前缀token索引集合具有相似性的特征，根据所述步骤二中得到的重要前缀token索引集合s_{i}生成第i+1层的候选重要前缀token索引集合s’{i+1}，并在第i层计算期间异步预取s’_{i+1}对应的第i+1层候选重要kv；

6、步骤四，在第i层计算结束时，停止发起新的第i+1层候选重要kv预取请求，以避免在当前层计算窗口之外继续加载无效数据；

7、步骤五，在第i+1层计算开始前，获取第i+1层的实际重要前缀token索引集合s_{i+1}，将s_{i+1}与已预取到的候选重要kv进行比对校验；其中，对于属于s_{i+1}的候选重要kv，作为重要kv直接使用；对于不属于s_{i+1}的候选重要kv，直接丢弃；对于属于s_{i+1}但未预取到的重要kv，进行同步补齐加载；

8、步骤六，使用校验后的重要kv缓存执行第i+1层计算，并将层索引递增，重复执行步骤三至步骤五，直至预填充阶段结束或到达最后一个transformer层；

9、其中，通过上述步骤，利用当前层计算时间隐藏下一层重要kv的i/o加载延迟，并通过校验与补齐机制修正投机预取误差，从而在保持推理结果稳定的同时降低首字延迟。

10、进一步地，所述kv元数据包括层索引、token索引、注意力头索引、缓存块索引以及对应kv数据所在的存储层级。

11、进一步地，所述步骤二中，所述重要前缀token索引集合s_{i}通过以下方式获取：针对第i层的每个注意力头，计算该注意力头的查询向量与可复用前缀键向量之间的注意力权重矩阵，并对所述注意力权重矩阵按列求和，得到该注意力头下各前缀token的重要性分数；将各注意力头的重要性分数按前缀token对应相加，得到由所有注意力头共用的各前缀token的汇总重要性分数；按照所述汇总重要性分数从高到低排序，选择预设比例或预设数量的前缀token，形成第i层的重要token索引集合s_{i}。

12、进一步地，所述步骤三中，所述生成第i+1层的候选重要前缀token索引集合s’{i+1}，包括：根据token的重要性分数排序、历史预取命中率、预取缓冲区大小或当前i/o带宽，对s_{i}进行截取或调整后生成s’_{i+1}。

13、进一步地，所述步骤三中，所述在第i层计算期间异步预取s’_{i+1}对应的第i+1层候选重要kv，包括：按重要性分数由高到低的顺序预取候选重要kv；其中，若所述候选重要kv位于gpu显存，则直接记录其地址或写入预取缓冲区；若位于cpu内存，则将其传输至gpu显存；若位于外存或磁盘，则先读取至cpu内存，再传输至gpu显存。

14、本发明还提供了一种相似性感知的重要前缀kv投机预取系统，包括：

15、重要前缀kv识别模块，用于在前缀索引结构中匹配当前请求对应的可复用前缀，并获得所述可复用前缀在各transformer层中的kv元数据及存储位置；对于transformer层中的第i层，获取当前请求在该层对应的重要前缀token索引集合s_{i}，并加载第i层计算所需的重要kv；

16、相似性感知投机预取模块，被配置为在第i层开始计算后，根据所述重要前缀token索引集合s_{i}生成第i+1层的候选重要前缀token索引集合s’{i+1}，并在第i层计算期间异步预取s’_{i+1}对应的第i+1层候选重要kv；持续监控gpu状态，在检测到gpu完成第i层计算时，停止发起新的第i+1层候选重要kv预取请求；

17、校验恢复模块，被配置为在第i+1层计算开始前，获取第i+1层的实际重要前缀token索引集合s_{i+1}，将s_{i+1}与已预取到的候选重要kv进行比对校验；其中，对于属于s_{i+1}的候选重要kv，作为重要kv直接使用；对于不属于s_{i+1}的候选重要kv，直接丢弃；对于属于s_{i+1}但未预取到的重要kv，进行同步补齐加载；

18、迭代执行模块，被配置为在校验恢复模块完成第i+1层候选重要kv的比对校验后，触发使用校验后的重要kv缓存执行第i+1层计算，并在第i+1层计算完成后将层索引递增，继续调用相似性感知投机预取模块和校验恢复模块，直至预填充阶段结束或到达最后一个transformer层。

19、进一步地，所述kv元数据包括层索引、token索引、注意力头索引、缓存块索引以及对应kv数据所在的存储层级。

20、进一步地，所述重要前缀kv识别模块包括：针对第i层的每个注意力头，计算该注意力头的查询向量与可复用前缀键向量之间的注意力权重矩阵，并对所述注意力权重矩阵按列求和，得到该注意力头下各前缀token的重要性分数；将各注意力头的重要性分数按前缀token对应相加，得到由所有注意力头共用的各前缀token的汇总重要性分数；按照所述汇总重要性分数从高到低排序，选择预设比例或预设数量的前缀token，形成第i层的重要token索引集合s_{i}。

21、进一步地，所述相似性感知投机预取模块包括：根据token的重要性分数排序、历史预取命中率、预取缓冲区大小或当前i/o带宽，对s_{i}进行截取或调整后生成s’_{i+1}。

22、进一步地，所述相似性感知投机预取模块还包括：按重要性分数由高到低的顺序预取候选重要kv；其中，若所述候选重要kv位于gpu显存，则直接记录其地址或写入预取缓冲区；若位于cpu内存，则将其传输至gpu显存；若位于外存或磁盘，则先读取至cpu内存，再传输至gpu显存。

23、本发明的有益效果在于：

24、第一，本发明将下一层重要kv缓存的加载提前到当前层计算期间执行，利用gpu计算时间隐藏外存或cpu内存到gpu显存的数据传输延迟，从而减少首个token生成延迟中的关键路径i/o等待。

25、第二，本发明采用推测后核验的方式处理跨层预测误差，实际参与下一层注意力计算的kv缓存仍由该层真实重要前缀令牌集合决定，因此可以降低预取误差对推理结果的影响。

26、第三，本发明只预取候选重要前缀令牌对应的kv缓存，不需要预取完整前缀kv缓存，因而能够同时利用重要性稀疏性和跨层预取，适用于gpu显存、cpu内存和外存构成的多级kv缓存系统。

27、第四，本发明不要求修改模型权重，也不依赖额外训练，可作为推理系统中的运行时调度机制与现有前缀kv复用、kv压缩、缓存替换和分层存储策略结合使用。