一种张量显存映射管理方法、设备、介质及产品与流程

本技术实施例涉及人工智能芯片，尤其涉及一种张量显存映射管理方法、设备、介质及产品。

背景技术：

1、近些年来深度学习领域技术飞速发展，大语言模型层出不穷，随着模型能力的不断突破和ai智能体（agent）的发展，大模型的推理应用场景得到快速发展。目前主流的推理框架包括：vllm、sglang、lightllm、triton server、tensorrt-llm等，这些推理框架通常运行在图形处理器（graphics processing unit，简称gpu）等主流芯片上。

2、对于主流芯片来说，其显存架构是连续排列的方式，即数据的逻辑形状与物理形状是一致的。因此，主流的推理框架都是以这种连续排列方式为基础进行显存管理。

3、然而，对于追求高算力高性能的芯片架构来说，其显存架构的排布方式并非连续且紧密的，即数据的逻辑形状与物理形状存在差异，这样，在对主流的推理框架适配新芯片时，针对显存架构采用非连续排列方式的新芯片，往往需要修改推理框架的显存管理模块，其工作量比较大，且效率较低。

技术实现思路

1、本技术实施例提供了一种张量显存映射管理方法、设备、介质及产品，用于在避免推理框架大量改动的情况下，解决显存架构差异性带来的适配问题。

2、一方面，本技术实施例提供了一种张量显存映射管理方法，该方法包括：

3、接收推理框架针对逻辑张量发起的显存分配请求；

4、基于人工智能芯片的硬件维度约束，对所述逻辑张量的逻辑形状进行维度折叠，获得物理形状和折叠系数；

5、基于所述物理形状和所述逻辑张量的数据类型，获得显存需求大小，并基于所述显存需求大小申请相应的物理显存和虚拟地址区间；

6、基于所述折叠系数、所述物理显存的物理地址和所述物理显存的硬件显存属性，构建张量对象；

7、建立并保存所述张量对象与所述虚拟地址区间的映射关系；

8、将所述虚拟地址区间的虚拟首地址返回给所述推理框架，以使所述推理框架基于所述虚拟首地址操作所述物理显存。

9、一方面，本技术实施例提供了一种张量显存映射管理装置，该装置包括：

10、接收模块，用于接收推理框架针对逻辑张量发起的显存分配请求；

11、处理模块，用于基于人工智能芯片的硬件维度约束，对所述逻辑张量的逻辑形状进行维度折叠，获得物理形状和折叠系数；基于所述物理形状和所述逻辑张量的数据类型，获得显存需求大小，并基于所述显存需求大小申请相应的物理显存和虚拟地址区间；

12、构建模块，用于基于所述折叠系数、所述物理显存的物理地址和所述物理显存的硬件显存属性，构建张量对象；建立并保存所述张量对象与所述虚拟地址区间的映射关系；

13、返回模块，用于将所述虚拟地址区间的虚拟首地址返回给所述推理框架，以使所述推理框架基于所述虚拟首地址操作所述物理显存。

14、可选地，所述显存分配请求包括：显存架构标识；

15、所述处理模块具体用于：

16、选取与所述显存架构标识匹配的显存管理实例；

17、从所述显存管理实例对应的显存池中，选取满足所述显存需求大小且处于空闲状态的物理显存，并将所述物理显存的状态更新为已占用。

18、可选地，所述处理模块还用于：

19、接收所述推理框架发起的携带第一虚拟地址的显存释放请求；

20、当所述第一虚拟地址命中所述虚拟地址区间时，从所述张量对象中获取所述物理地址和所述硬件显存属性，所述硬件显存属性包括：显存架构标识；

21、选取与所述显存架构标识匹配的显存管理实例，并基于所述物理地址从所述显存管理实例对应的显存池中确定所述物理显存；

22、将所述物理显存的状态更新为空闲状态。

23、可选地，所述第一虚拟地址为所述虚拟首地址；或者，所述虚拟首地址与预设的虚拟字节偏移的叠加值。

24、可选地，所述处理模块还用于：

25、接收所述推理框架发起的携带第二虚拟地址的算子调用请求；

26、当所述第二虚拟地址命中所述虚拟地址区间时，通过硬件算子基于所述张量对象，从所述物理显存中读取数据并处理，获得处理结果。

27、可选地， 所述硬件显存属性包括：数据排布类型；

28、所述处理模块具体用于：

29、针对所述逻辑张量中的每个元素，通过所述硬件算子执行以下操作：

30、获取所述元素在所述逻辑张量中的局部字节偏移；

31、基于所述物理形状、所述折叠系数和所述数据排布类型，将所述局部字节偏移转换为物理字节偏移；

32、基于所述物理地址和所述物理字节偏移，获得目标访存地址；

33、按照所述目标访存地址从所述物理显存中读取所述元素并处理，获得所述元素对应的子处理结果。

34、可选地，所述处理模块具体用于：

35、当所述第二虚拟地址为所述虚拟首地址与预设的虚拟字节偏移的叠加值时，将所述局部字节偏移与所述虚拟字节偏移叠加，获得总字节偏移；

36、基于所述物理形状、所述折叠系数和所述数据排布类型，将所述总字节偏移转换为物理字节偏移。

37、可选地，所述处理模块还用于：

38、接收所述推理框架发起的拷贝请求，所述拷贝请求包括：源虚拟地址和目的虚拟地址，所述目的虚拟地址为主机内存地址；

39、从映射表中获取所述源虚拟地址命中的第一张量对象，所述第一张量对象用于寻址第一物理显存；

40、基于所述第一张量对象从所述第一物理显存中读取源数据，并对所述源数据进行反向布局转换，获得目标数据；

41、将所述目标数据拷贝到所述目的虚拟地址。

42、可选地，所述处理模块还用于：

43、接收所述推理框架发起的拷贝请求，所述拷贝请求包括：源虚拟地址和目的虚拟地址，所述源虚拟地址为主机内存地址；

44、从映射表中获取所述目的虚拟地址命中的第二张量对象，所述第二张量对象用于寻址第二物理显存；

45、从所述源虚拟地址读取源数据；

46、基于所述第二张量对象，将所述源数据进行布局转换后写入所述第二物理显存。

47、可选地，所述处理模块还用于：

48、接收所述推理框架发起的拷贝请求，所述拷贝请求包括：源虚拟地址和目的虚拟地址；

49、从映射表中获取所述源虚拟地址命中的第一张量对象和所述目的虚拟地址命中的第二张量对象，所述第一张量对象用于寻址第一物理显存，所述第二张量对象用于寻址第二物理显存；

50、基于所述第一张量对象，从所述第一物理显存中读取源数据；

51、基于所述第二张量对象，将所述源数据写入所述第二物理显存。

52、一方面，本技术实施例提供了一种计算机设备，包括存储器、人工智能芯片、中央处理器及存储在存储器上并在中央处理器上运行的计算机程序，所述中央处理器执行所述计算机程序时实现上述张量显存映射管理方法的步骤。

53、一方面，本技术实施例提供了一种计算机可读存储介质，其存储有可由计算机设备执行的计算机程序，当所述计算机程序在计算机设备上运行时，使得所述计算机设备执行上述张量显存映射管理方法的步骤。

54、一方面，本技术实施例提供了一种计算机程序产品，所述计算机程序产品包括存储在计算机可读存储介质上的计算机程序，所述计算机程序包括程序指令，当所述程序指令被计算机设备执行时，使所述计算机设备执行上述张量显存映射管理方法的步骤。

55、本技术实施例中，张量显存映射层接收显存分配请求时，基于人工智能芯片的硬件维度约束，对逻辑张量的逻辑形状进行维度折叠，获得物理形状和折叠系数；然后基于物理形状和逻辑张量的数据类型，获得显存需求大小，并基于显存需求大小申请相应的物理显存和虚拟地址区间；同时，基于折叠系数、物理显存的物理地址和物理显存的硬件显存属性构建张量对象，并保存张量对象与虚拟地址区间的映射关系，即由张量显存映射层维护虚拟地址与物理显存之间的映射关系，该映射关系可覆盖逻辑形状与物理形状一致或不一致的情况，解决了显存架构差异性带来的问题，而推理框架基于张量显存映射层返回的虚拟地址即可完成物理显存操作，有效避免推理框架的大量改动适配，从而大幅度减少工作量和难度，实现高效的硬件适配。