基于循环神经网络的键值缓存管理方法及系统与流程

本技术涉及大语言模型，尤其是涉及一种基于循环神经网络的键值缓存管理方法及系统。

背景技术：

1、大语言模型基于transformer架构，在推理过程中采用自回归方式逐token生成输出。设模型共l层，第l层的多头注意力计算中，在每一解码步，模型需要访问序列中所有前序token的键向量和值向量。kvcache机制将这些键值向量缓存以避免重复计算。然而，kvcache的总显存占用随序列长度、隐藏维度、模型层数和参数精度线性增长。以llama-2-70b模型为例，在长序列推理场景下，kvcache的显存需求远超单卡gpu显存容量，导致推理服务的吞吐量和并发能力受到严重制约。

2、目前，为解决kvcache的显存溢出问题，普遍采用基于静态启发式规则的缓存管理策略，如最近最少使用（lru）策略或先进先出（fifo）策略。以lru策略为例，当gpu显存容量不足时，系统选择访问时间戳最早的缓存块进行淘汰，将其从gpu显存移除；当后续解码步需要访问已被淘汰的缓存块时，则触发重计算重新生成其键值向量。部分方案进一步引入cpu内存作为二级存储，将淘汰的缓存块暂存于cpu内存中以替代直接释放，从而避免部分重计算开销。

3、针对上述技术，淘汰策略仅依据历史访问时间等静态信息进行决策，无法预判缓存块在未来解码步中被访问的可能性，导致高访问概率的缓存块可能被误淘汰，短期内被重新访问时需执行重计算，产生额外计算代价。

4、基于此，本技术提供一种基于循环神经网络的键值缓存管理方法及系统。

技术实现思路

1、为了改善淘汰策略仅依据历史访问时间等静态信息进行决策，无法预判缓存块在未来解码步中被访问的可能性，导致高访问概率的缓存块可能被误淘汰，短期内被重新访问时需执行重计算，产生额外计算代价的问题，本技术提供一种基于循环神经网络的键值缓存管理方法及系统。

2、第一方面，本技术提供的一种基于循环神经网络的键值缓存管理方法，采用如下的技术方案：包括：

3、构建基于循环神经网络的访问模式预测器；所述访问模式预测器以连续多个解码步的kv缓存块访问状态序列为输入，输出各缓存块在未来预测窗口内被访问的概率分数；所述访问状态序列基于大语言模型各解码步的注意力权重矩阵构造得到；

4、在生成新的kv缓存块时，基于所述访问模式预测器对已有缓存块输出的概率分数，采用基于位置邻域的初始化策略计算新块的初始概率分数；根据所述初始概率分数与自适应准入阈值的比较结果，执行新块的准入决策，将新块分配至gpu显存或cpu内存；自适应准入阈值随当前gpu显存占用率动态调整；

5、当gpu显存中缓存块数量超过容量上限时，利用所述访问模式预测器对各缓存块输出的概率分数，按概率分数从低到高选择若干个缓存块作为淘汰对象，将其kv数据从gpu显存迁移至cpu内存；

6、对于被淘汰至cpu内存的缓存块，基于缓存块被淘汰时的概率分数，通过单调非递减的保留时长映射函数确定其在cpu内存中的保留步数；在保留期内，若目标缓存块被再次访问则换入gpu显存并取消保留计时；超出保留期后，释放目标缓存块占用的cpu内存空间。

7、优选的，所述访问模式预测器以连续多个解码步的kv缓存块访问状态序列为输入，输出各缓存块在未来预测窗口内被访问的概率分数，包括：

8、初始化隐藏状态为零向量，维度为预设的隐藏状态维度；

9、对于历史窗口内的每个时间步，依次执行以下计算：

10、将上一时间步的隐藏状态与当前时间步的访问状态向量进行拼接，经第一权重矩阵线性变换和第一偏置向量偏移后，通过sigmoid激活函数得到重置门向量；

11、将上一时间步的隐藏状态与当前时间步的访问状态向量进行拼接，经第二权重矩阵线性变换和第二偏置向量偏移后，通过sigmoid激活函数得到更新门向量；

12、将所述重置门向量与上一时间步的隐藏状态进行逐元素相乘，并与当前时间步的访问状态向量拼接，经第三权重矩阵线性变换和第三偏置向量偏移，通过双曲正切激活函数得到候选隐藏状态；所述重置门向量用于控制上一时间步隐藏状态的保留程度；

13、以所述更新门向量为插值系数，对上一时间步的隐藏状态与所述候选隐藏状态进行线性插值，得到当前时间步的隐藏状态；其中，当所述更新门向量趋近于零时，保留旧状态，趋近于1时采用候选状态；

14、经过历史窗口长度个时间步的迭代后，获取最终隐藏状态，所述最终隐藏状态编码了历史访问序列中的时序依赖关系；

15、将所述最终隐藏状态经输出权重矩阵线性变换和输出偏置向量偏移，通过sigmoid激活函数获取各缓存块的概率分数向量；

16、其中，所述访问模式预测器的可训练参数集合包括所有权重矩阵和偏置向量，其总参数量相对于kv缓存总规模为低阶项，单步推理的计算复杂度与缓存块总数和隐藏状态维度呈线性关系。

17、优选的，所述访问模式预测器以连续多个解码步的kv缓存块访问状态序列为输入，输出各缓存块在未来预测窗口内被访问的概率分数，包括：

18、初始化隐藏状态为零向量，细胞状态为零向量；

19、对于历史窗口内的每个时间步，依次执行以下计算：

20、将上一时间步的隐藏状态与当前时间步的访问状态向量进行拼接，分别经遗忘门权重矩阵、输入门权重矩阵、候选状态权重矩阵和输出门权重矩阵线性变换，并分别经对应的偏置向量偏移后，依次通过激活函数得到遗忘门向量、输入门向量、候选细胞状态向量和输出门向量；其中，遗忘门向量和输入门向量经sigmoid激活，候选细胞状态向量经双曲正切激活，输出门向量经sigmoid激活；

21、将上一时间步的细胞状态与遗忘门向量逐元素相乘，加上输入门向量与候选细胞状态向量的逐元素相乘结果，得到当前时间步的细胞状态；

22、将输出门向量与经双曲正切激活后的当前时间步细胞状态进行逐元素相乘，得到当前时间步的隐藏状态；

23、经过历史窗口长度个时间步的迭代后，获取迭代隐藏状态，并将所述迭代隐藏状态经输出权重矩阵线性变换和输出偏置向量偏移，通过sigmoid激活函数，得到各缓存块的概率分数向量；

24、其中，lstm的可训练参数集合包括四个门控权重矩阵、输出权重矩阵以及所有偏置向量。

25、优选的，所述访问模式预测器的训练过程，包括：

26、定义真实访问标签：对于每个解码步和每个缓存块，若目标块在之后未来预测窗口长度个解码步内至少被访问一次，则标签为正样本，否则为负样本；

27、收集大语言模型在真实推理工作负载下的kv缓存块访问轨迹数据，构成轨迹数据集；采用滑动窗口方式从所述轨迹数据集中构造训练样本集，每个训练样本包括由连续历史窗口长度个解码步的访问状态向量构成的输入序列以及对应的标签向量；

28、采用二值交叉熵作为单样本损失函数，计算所述访问模式预测器对每个缓存块输出的概率分数与对应真实访问标签之间的损失；定义经验风险函数为所有训练样本的损失均值加上正则化项，正则化系数用于控制过拟合程度；

29、以最小化所述经验风险函数为训练目标，采用adam优化器进行参数更新，其中，一阶矩和二阶矩估计的衰减率以及数值稳定项采用预设值；

30、学习率采用余弦退火调度策略，随训练步数增加从学习率上界按余弦曲线递减至学习率下界；

31、在验证集上监控损失值，若连续预设耐心参数个训练轮次验证损失未下降，则提前终止训练。

32、优选的，所述自适应准入阈值的定义及性质，包括：

33、自适应准入阈值函数以当前gpu常驻块数量为自变量，以基准准入阈值为基础值，以调节系数为增幅系数，以显存利用警告水位线和gpu可容纳块数上限为分段边界；

34、若当前gpu常驻块数量低于所述显存利用警告水位线，则准入阈值恒等于所述基准准入阈值，准入条件处于最宽松状态；

35、若当前gpu常驻块数量介于所述显存利用警告水位线与所述gpu可容纳块数上限之间，则准入阈值随当前gpu常驻块数量线性递增，准入条件逐步趋严，使得只有预测访问概率更高的新块才能进入gpu显存；

36、若当前gpu常驻块数量趋近于所述gpu可容纳块数上限，则准入阈值达到最大值，为所述基准准入阈值与所述调节系数之和，准入条件处于最严格状态；

37、在准入决策执行后，以gpu常驻块集合中各缓存块的概率分数的算术均值作为gpu显存有效利用率指标，有效利用率指标用于衡量gpu显存中缓存块的平均预测访问概率。

38、优选的，所述当gpu显存中缓存块数量超过容量上限时，利用所述访问模式预测器对各缓存块输出的概率分数，按概率分数从低到高选择若干个缓存块作为淘汰对象，将其kv数据从gpu显存迁移至cpu内存，包括：

39、按概率分数从低到高选择若干缓存块作为淘汰对象，等价于求解优化问题：在满足释放指定数量缓存块的约束下，使被淘汰缓存块的概率分数之和最小化；优化问题的最优解满足如下性质：淘汰后保留在gpu显存中的缓存块集合具有最大的概率分数总和，即gpu显存中保留的是最有可能被访问的缓存块集合；

40、定义被淘汰缓存块的预期重计算代价为目标块的概率分数与单块重计算开销的乘积，其中，单块重计算开销包括重新计算key向量和value向量所需的浮点运算量和时间延迟；淘汰策略等价于最小化被淘汰缓存块的预期重计算代价总和；

41、淘汰执行过程中，将每个被淘汰缓存块的kv数据从gpu显存复制至cpu内存预留区域，其数据传输代价为被淘汰块数与单块传输延迟的乘积，单块传输延迟包括pcie带宽限制和内存拷贝时间；

42、对于每个被淘汰块，将映射条目从gpu显存地址更新为cpu内存地址，并记录淘汰发生的时间步。

43、优选的，所述对于被淘汰至cpu内存的缓存块，基于缓存块被淘汰时的概率分数，通过单调非递减的保留时长映射函数确定其在cpu内存中的保留步数，包括：

44、述保留时长映射函数以被淘汰缓存块的概率分数为输入，以缩放系数与概率分数的指数调节参数次幂的乘积为基础值，加上最小保留步数后，经向下取整得到该缓存块在cpu内存中的保留步数；其中，所述缩放系数用于控制保留时长的整体幅度，所述指数调节参数用于控制概率分数到保留时长的非线性映射程度，所述最小保留步数用于保证所有被淘汰缓存块至少在cpu内存中保留预设步数；

45、所述保留时长映射函数满足单调性约束：对于任意两个概率分数，若前者大于后者，则前者的保留步数大于或等于后者的保留步数，使得概率分数越高的缓存块在cpu内存中保留时间越长，即高概率分数的缓存块即使被淘汰也更有可能在保留期内被重新访问；

46、当所述指数调节参数取值为一时，所述保留时长映射函数退化为线性映射，保留步数与概率分数成正比；当所述指数调节参数大于一时，高概率分数区域的保留时长差异被放大，低概率分数区域的保留时长差异被压缩，使得概率分数的微小差异在高分区域导致更大的保留时长差异。

47、第二方面，本技术一种基于循环神经网络的键值缓存管理装置，采用如下技术方案，包括：

48、访问预测模块，用于构建基于循环神经网络的访问模式预测器；所述访问模式预测器以连续多个解码步的kv缓存块访问状态序列为输入，输出各缓存块在未来预测窗口内被访问的概率分数；所述访问状态序列基于大语言模型各解码步的注意力权重矩阵构造得到；

49、缓存调整模块，用于在生成新的kv缓存块时，基于所述访问模式预测器对已有缓存块输出的概率分数，采用基于位置邻域的初始化策略计算新块的初始概率分数；根据所述初始概率分数与自适应准入阈值的比较结果，执行新块的准入决策，将新块分配至gpu显存或cpu内存；自适应准入阈值随当前gpu显存占用率动态调整；

50、缓存淘汰模块，用于当gpu显存中缓存块数量超过容量上限时，利用所述访问模式预测器对各缓存块输出的概率分数，按概率分数从低到高选择若干个缓存块作为淘汰对象，将其kv数据从gpu显存迁移至cpu内存；

51、缓存保留模块，用于对于被淘汰至cpu内存的缓存块，基于缓存块被淘汰时的概率分数，通过单调非递减的保留时长映射函数确定其在cpu内存中的保留步数；在保留期内，若目标缓存块被再次访问则换入gpu显存并取消保留计时；超出保留期后，释放目标缓存块占用的cpu内存空间。

52、第三方面，本技术还提供一种控制设备，所述设备包括：

53、包括存储器和处理器，所述存储器上存储有能够被处理器加载并执行如上述基于循环神经网络的键值缓存管理方法的计算机程序。

54、第四方面，本技术还提供一种计算机可读存储介质，存储有能够被处理器加载并执行如上述基于循环神经网络的键值缓存管理方法的计算机程序。

55、综上所述，本技术中系统构建基于门控循环单元或长短期记忆网络的访问模式预测器，以连续多个解码步的注意力权重矩阵构造访问状态序列作为输入，输出各缓存块在未来窗口内被访问的概率分数，从而建模kv缓存访问的时序依赖关系。在新缓存块生成时，利用已有块的预测概率通过位置邻域加权初始化新块概率，并与随gpu显存占用率动态调整的自适应准入阈值比较，决定将新块分配至gpu显存或cpu内存，实现预测驱动的准入控制。当gpu显存容量超限时，基于概率分数从低到高选取淘汰对象，将其迁移至cpu内存，并以被淘汰块的概率分数通过单调保留时长映射函数确定其在cpu内存中的差异化保留步数，超出保留期则释放内存。在推理服务运行过程中，持续采集访问轨迹数据对预测器进行在线增量更新，使其适应工作负载分布的变化。从而将传统被动式缓存管理转变为主动预测驱动，通过rnn建模访问时序依赖使淘汰决策趋向最小化期望重计算代价，准入策略将低价值块直接分配至cpu以提高gpu显存有效利用率，保留策略避免过早释放高概率块的同时及时清理低概率块以优化cpu内存使用效率，且预测器参数量和计算量相对于kv缓存总规模为低阶项，对推理引擎核心逻辑的侵入性修改仅需拦截块管理器的四个接口。