Ai model inference method and apparatus, and related device

Abstract

An artificial intelligence (AI) model inference method and apparatus, and a related device, relating to the technical field of AI. Iterative inference is performed by using an AI model and acquired input data, to obtain inference results. A processing module in the AI model corresponds to a first physical storage region indicated by a first virtual address space, and the first physical storage region is used for storing KVs generated by the processing module during iterative inference processes. When the first physical storage region is insufficient to continue storing a KV generated by the processing module during execution of data processing in an N-th inference process, a second physical storage region indicated by a second virtual address space is allocated to the processing module, so as to store the KV. The second virtual address space is contiguous with the first virtual address space. Thus, the KVs generated by the processing module are stored by means of allocating physical storage regions according to requirements, such that the utilization rate of the physical storage regions can be improved, and an operator in the processing module does not need to have a capability of being merged with a KV in a virtual address space.

Description

AI model inference method, device and related equipment

[0001] This application claims priority to the Chinese patent application filed with the State Intellectual Property Office on February 2, 2024, with application number 202410156724.9 and application name “AI model reasoning method, device and related equipment”, the entire contents of which are incorporated by reference into this application. Technical Field

[0002] The present application relates to the field of artificial intelligence technology, and in particular to an AI model reasoning method, apparatus, and related equipment. Background Art

[0003] With the development of artificial intelligence (AI) technology, the parameter scale of AI models is getting larger and larger. For example, the parameter scale of AI models can reach hundreds of billions, which makes the reasoning performance of AI models an important factor restricting the promotion of AI models. In some application scenarios, AI models usually perform multiple reasoning processes when performing an inference task. After each reasoning, the AI ​​model outputs an element, which can also be called a token. In addition, the element output by the AI ​​model in each reasoning process will be spliced ​​with the element that serves as the input of the AI ​​model in this reasoning process, and the spliced ​​result will be used as the input for the next reasoning of the AI ​​model, and so on, until the reasoning task is completed.

[0004] Typically, key-value cache (KV cache) technology can be used to optimize the reasoning performance of AI models. KV cache technology improves the reasoning performance of AI models by trading space for time without affecting the calculation accuracy of the AI ​​model. Specifically, it caches the key-value (KV) generated during the AI ​​model reasoning process, allowing the KV to be directly read for subsequent calculations without having to recalculate the KV.

[0005] At present, when an AI model performs an inference task, a physical storage area of ​​a fixed size is usually allocated to each network layer in the AI ​​model. The physical storage area is used to cache the KV generated by the network layer during each inference. Since the number of inferences performed by the AI ​​model for different inference tasks may vary, the size of the physical storage area allocated to each network layer is usually the maximum size, which is determined based on the maximum number of inferences that the AI ​​model may perform. However, in most cases, the number of inferences performed by the AI ​​model during the execution of an inference task is usually less than the maximum number of inferences, which makes the physical storage area allocated to each network layer usually too large, and there are often many areas where key values ​​are not stored, resulting in a more serious waste of storage resources and a large amount of resource fragmentation.

[0006] Summary of the Invention

[0007] This application provides an AI model inference method to alleviate the problem of physical storage resource waste and reduce resource fragmentation. In addition, this application also provides an AI model inference device, a computing device, a computer-readable storage medium, and a computer program product.

[0008] In a first aspect, the present application provides an AI model inference method, which can be performed by a corresponding AI model inference device. Specifically, the AI ​​model inference device obtains input data, which can be text data or image data, and the AI ​​model inference device uses the AI ​​model and the input data to perform iterative inference to obtain an inference result, wherein the AI ​​model can output a token each time it infers, and the inference result can be the concatenation result of multiple tokens output by the AI ​​model during iterative inference. The AI ​​model includes one or more processing modules, each of which can be, for example, a network layer in the AI ​​model or an attention head in a network layer. The processing module in the AI ​​model corresponds to a first physical storage area indicated by a first virtual address space, and the first physical storage area is used to store the KV (key value) generated by the processing module when the AI ​​model performs iterative inference. In the Nth inference process of the iterative inference (taking the Nth time as an example, N is a positive integer), the processing module in the AI ​​model will perform data processing, such as performing calculations based on the input data, and outputting the calculated data to the next processing module for processing. In addition, when performing data processing, the processing module will generate the KV for the Nth inference process. Since the processing module has been allocated a first physical storage area, the processing module can store the newly generated KV in the first physical storage area. However, when the allocated first physical storage area is insufficient to continue to store the KV generated by the processing module during the Nth reasoning process, a second physical storage area indicated by a second virtual address space is allocated to the processing module, and the KV generated during the Nth reasoning process is stored in the second physical storage area according to the second virtual address space. The second virtual address space is continuous with the first virtual address space.

[0009] During the iterative reasoning process of the AI ​​model, the AI ​​model reasoning device does not pre-allocate a fixed-length physical storage area for the processing module to store the KV generated during each reasoning process. Instead, when the allocated storage area (i.e., the first physical storage area) is insufficient to continue to store the newly generated KV, it dynamically allocates a new physical storage area (i.e., the second physical storage area) for the processing module, and uses the dynamically allocated new physical storage area to store the newly generated KV. In this way, the AI ​​model reasoning device stores the KV generated during multiple reasoning processes by allocating physical storage areas on demand, which can effectively alleviate the problem of storage resource waste caused by allocating fixed-length physical storage areas to the processing module, thereby improving the utilization rate of the physical storage area and solving the problem of resource fragmentation in the physical storage area. Moreover, since the first virtual address space for storing KVs is continuous with the second virtual address space, the operators in the processing module (such as the attention operator and its variants) do not need to splice the read KVs in the virtual address space (that is, splice the key values ​​in the order in which the key values ​​are generated) when reading KVs in different physical storage areas. Therefore, the operators in the processing module are not required to have the ability to splice KVs in the virtual address space, so that they can adapt to various variants of operators in the AI ​​model and improve the generalization of the operators.

[0010] In one possible implementation, the AI ​​model inference device can also allocate a fixed-size target virtual address space to the processing module before using the AI ​​model for iterative inference. The target virtual address space includes a first virtual address space and a second virtual address space with continuous virtual addresses. Then, when the AI ​​model inference device allocates the second physical storage area to the processing module, it can specifically first determine the second physical storage area to be allocated, and then establish a mapping between the second virtual address space and the second physical storage area. In this way, the AI ​​model inference device dynamically allocates a physical storage area to the processing module by mapping the second physical storage area to the allocated second virtual address space, thereby improving the utilization rate of the physical storage area and solving the resource fragmentation problem of the physical storage area.

[0011] In one possible implementation, when determining the second physical storage area, the AI ​​model inference device may specifically select the second physical storage area from a plurality of available physical storage areas included in the storage resource pool, where each of the plurality of physical storage areas has the same size, for example, each physical storage area has a size of 2 megabits. In this way, the AI ​​model inference device can dynamically select and allocate physical storage areas for the processing module from the storage resource pool, thereby improving the utilization rate of the physical storage areas in the storage resource pool.

[0012] In one possible implementation, after the AI ​​model completes iterative reasoning, the AI ​​model inference device releases multiple physical storage areas allocated to the processing module during the iterative reasoning process of the AI ​​model, wherein the multiple physical storage areas include a first physical storage area and a second physical storage area. In this way, after the iterative reasoning of the AI ​​model is completed, by releasing the allocated physical storage areas, the physical storage areas can be recycled so that the recycled physical storage areas can be used to meet the needs of the next iterative reasoning process of the AI ​​model, thereby realizing the recycling of physical storage resources.

[0013] In one possible embodiment, when the AI ​​model inference device determines whether the first physical storage area is insufficient to continue storing KV, it may specifically determine whether the virtual address used to store the KV generated by the processing module in the Nth inference process is mapped to a physical address, and when the virtual address is not mapped to a physical address, it is determined that the first physical storage area is insufficient to continue storing the KV generated in the Nth inference process; or, the AI ​​model inference device may detect whether there is remaining storage space in the first physical storage area, and when there is no remaining storage space in the first physical storage area, it is determined that the first physical storage area is insufficient to continue storing the KV generated in the Nth inference process; or, the AI ​​model inference device may detect whether the number of KVs that can be stored in the first physical storage area is zero, and when the number of KVs that can be stored is zero, it is determined that the first physical storage area is insufficient to continue storing the KV generated in the Nth inference process.

[0014] In one possible embodiment, the processing module includes M (M is a positive integer greater than 1) sub-processing modules, and the first virtual address space includes M first sub-address spaces, and the physical storage area indicated by each first sub-address space in the M first sub-address spaces is used to store the KV generated by a sub-processing module, and the physical storage areas indicated by different first sub-address spaces are used to store the KVs generated by different sub-processing modules. In this way, when the AI ​​model inference device calls the interface to read the KVs respectively generated by multiple sub-processing modules, since multiple sub-processing modules can share the same physical storage area to store KVs, the AI ​​model inference device only needs to call the interface once to obtain the KVs respectively generated by multiple sub-processing modules in the physical storage area, thereby effectively reducing the number of times the AI ​​model inference device calls the interface, that is, reducing the overhead generated by the AI ​​model inference device calling the interface.

[0015] In one possible embodiment, after allocating a second physical storage area to the processing module, the AI ​​model inference device can also divide the first virtual address space and the second virtual address space into M second sub-address spaces, each of the M second sub-address spaces is used to store the KV generated by a sub-processing module, and different second sub-address spaces are used to store the KV generated by different sub-processing modules. In this way, after allocating a new physical storage area to the processing module, the virtual address of the KV generated by each sub-processing module can be made continuous by redistributing the allocated physical storage area. This makes it possible for the operator in the sub-processing module to read the KV in different physical storage areas without splicing the read KV in the virtual address space, thereby adapting to various variant operators in the AI ​​model and improving the generalization of the operator.

[0016] In one possible embodiment, the processing module includes multiple network layers in the AI ​​model; or, the processing module includes a network layer in the AI ​​model; or, the processing module includes multiple attention heads in a network in the AI ​​model; or, the processing module includes an attention head in the AI ​​model.

[0017] In the second aspect, the present application provides an AI model inference device, which includes: an inference unit, used to obtain input data, and use the AI ​​model and the input data to perform iterative inference to obtain an inference result, wherein the AI ​​model includes a processing module, and the processing module corresponds to a first physical storage area indicated by a first virtual address space, and the first physical storage area is used to store the key value KV generated by the processing module when the AI ​​model performs iterative inference; wherein, in the Nth inference process, the processing module performs data processing to generate the KV in the Nth inference process; an allocation unit, used to allocate a second physical storage area indicated by a second virtual address space to the processing module when the first physical storage area is insufficient to continue to store the KV generated by the processing module in the Nth inference process, wherein the second virtual address space is continuous with the first virtual address space, and N is a positive integer; the inference unit is also used to store the KV generated in the Nth inference process in the second physical storage area according to the second virtual address space.

[0018] In one possible embodiment, the allocation unit is further used to allocate a fixed-size target virtual address space to the processing module before using the AI ​​model for iterative reasoning, and the target virtual address space includes the first virtual address space and the second virtual address space with continuous virtual addresses; the allocation unit is used to: determine the second physical storage area; and establish a mapping between the second virtual address space and the second physical storage area.

[0019] In a possible implementation, the allocation unit is configured to select the second physical storage area from a plurality of available physical storage areas included in a storage resource pool, and each of the plurality of physical storage areas has the same size.

[0020] In a possible embodiment, the allocation unit is also used to release multiple physical storage areas allocated to the processing module during the iterative reasoning process of the AI ​​model after the AI ​​model completes the iterative reasoning, and the multiple physical storage areas include the first physical storage area and the second physical storage area.

[0021] In one possible embodiment, the allocation unit is also used to: determine whether the virtual address used to store the KV generated by the processing module during the Nth reasoning process is mapped to a physical address, and when the virtual address is not mapped to a physical address, determine that the first physical storage area is insufficient to continue storing the KV generated during the Nth reasoning process; or, detect whether there is remaining storage space in the first physical storage area, and when there is no remaining storage space in the first physical storage area, determine that the first physical storage area is insufficient to continue storing the KV generated during the Nth reasoning process; or, detect whether the number of KVs that can be stored in the first physical storage area is zero, and when the number of KVs that can be stored is zero, determine that the first physical storage area is insufficient to continue storing the KV generated during the Nth reasoning process.

[0022] In one possible embodiment, the processing module includes M sub-processing modules, the first virtual address space includes M first sub-address spaces, the physical storage area indicated by each first sub-address space in the M first sub-address spaces is used to store the KV generated by a sub-processing module, and the physical storage areas indicated by different first sub-address spaces are used to store the KVs generated by different sub-processing modules, and M is a positive integer greater than 1.

[0023] In one possible embodiment, the allocation unit is further used to: after allocating the second physical storage area to the processing module, divide the first virtual address space and the second virtual address space equally into M second sub-address spaces, each of the M second sub-address spaces is used to store the KV generated by a sub-processing module, and different second sub-address spaces are used to store KVs generated by different sub-processing modules.

[0024] In one possible implementation, the processing module includes multiple network layers in the AI ​​model; or, the processing module includes one network layer in the AI ​​model; or, the processing module includes multiple attention heads in a network in the AI ​​model; or, the processing module includes one attention head in the AI ​​model.

[0025] It is worth noting that the AI ​​model inference device provided in the second aspect corresponds to the AI ​​model inference method provided in the first aspect. Therefore, the technical effects of the second aspect and any implementation method of the second aspect can refer to the technical effects of the first aspect or the corresponding implementation method of the first aspect.

[0026] In a third aspect, the present application provides a computing device comprising a processor and a memory, wherein the memory stores program instructions, and the processor runs the program instructions to execute the AI ​​model inference device method in the above-mentioned first aspect or any possible implementation of the first aspect.

[0027] In a fourth aspect, the present application provides a computer-readable storage medium, which stores instructions. When the computer-readable storage medium is run on a computing device, the computing device executes the operating steps of the AI ​​model reasoning method described in the first aspect or any implementation of the first aspect.

[0028] In a fifth aspect, the present application provides a computer program product comprising instructions, which, when executed on a computing device, enables the computing device to execute the operational steps of the AI ​​model reasoning method described in the first aspect or any one of the implementations of the first aspect.

[0029] Based on the implementation methods provided in the above aspects, this application can also be further combined to provide more implementation methods. BRIEF DESCRIPTION OF THE DRAWINGS

[0030] FIG1 is a schematic diagram of an exemplary AI model performing an inference task;

[0031] FIG2 is a flow chart of an AI model reasoning method provided by this application;

[0032] FIG3 is a schematic diagram of dynamically allocating physical storage areas to store KVs provided by the present application;

[0033] FIG4 is a flow chart of another AI model reasoning method provided by this application;

[0034] FIG5 is a schematic diagram of the aggregate storage of KVs generated by multiple sub-processing modules provided by the present application;

[0035] FIG6 is a schematic diagram of dynamically allocating physical storage areas and adjusting sub-address spaces provided by the present application;

[0036] FIG7 is a schematic diagram of the structure of an exemplary AI model inference device provided by this application;

[0037] FIG8 is a schematic diagram of the hardware structure of a computing device provided in this application. DETAILED DESCRIPTION

[0038] To make the above-mentioned objects, features, and advantages of the present application more clearly understood, various non-limiting embodiments of the embodiments of the present application will be exemplified below with reference to the accompanying drawings. Obviously, the described embodiments are only a part of the embodiments of the present application, not all of them. Based on the embodiments in this application, all other embodiments obtained based on the above content are within the scope of protection of this application.

[0039] Refer to Figure 1, which is a schematic diagram of an exemplary AI model performing an inference task. As shown in Figure 1, when the AI ​​model performs an inference task, it will perform iterative inference based on the input data. Figure 1 takes the execution of P inference processes as an example for illustration, where P is a positive integer greater than 1. In actual application, the AI ​​model can be deployed on an AI model inference device, so that the AI ​​model can be run by the AI ​​model inference device. Exemplarily, the AI ​​model inference device can include a computing device, such as the AI ​​model inference device can be implemented by a server, etc.; alternatively, the AI ​​model inference device can include multiple computing devices, such as the AI ​​model inference device can be implemented by a server cluster, etc.

[0040] During the first inference, the AI ​​model can reason based on the input data and output token1.

[0041] During the second inference, the AI ​​model inference device may use token 1 as the input of the AI ​​model, so that the AI ​​model can perform inference based on token 1 and output token 2. During the second inference process, each network layer in the AI ​​model can save the key 1 and value 1 generated during the second inference process.

[0042] During the third inference, the AI ​​model inference device can concatenate token2 and token1 to obtain the concatenation result "token1+token2", and use the concatenation result as the input of the AI ​​model, so that the AI ​​model can perform inference based on the input "token1+token2" and output token3. During the third inference process, each network layer in the AI ​​model no longer needs to repeatedly calculate the key 1 and value 1 corresponding to token1 (which have been saved during the second inference process), and can simply save the key 2 and value 2 newly generated for token2 during the third inference process.

[0043] And so on, until after the Pth inference, the AI ​​model can use the input "token1+token2+...+token P-1 "Infer and output the token as the end character P , the token P For example, it can be the character "EOS" (end of sentence). At this time, the AI ​​model can P Confirm that the reasoning task is completed, and the AI ​​model can output the final reasoning result: "token1+token2+...+token P-1 ”.

[0044] Before the AI ​​model performs the first inference, if the AI ​​model inference device pre-allocates a physical storage area of ​​the maximum size for each network layer based on the maximum number of inferences Q (Q is a positive integer greater than 1) that the AI ​​model can perform, since the number of inferences performed by the AI ​​model when performing a single inference task is usually less than Q, this causes a large number of areas in the allocated physical storage area to be idle, resulting in serious waste of resources and a large amount of resource fragmentation. If the size of the physical storage area allocated to each network layer is small, the physical storage area may only support the storage of the KV (key value) generated during part of the inference process, and cannot continue to store the newly generated KV, resulting in the failure of the AI ​​model to perform the inference task due to insufficient allocated physical storage area. Exemplarily, the physical storage area can be, for example, a storage area in a memory, or a storage area in other types of storage media, and this is not limited.

[0045] Based on this, an embodiment of the present application provides a model reasoning method to alleviate the resource waste problem of physical storage resources and reduce the resource fragmentation generated when storing KV based on sufficient storage resources.

[0046] In specific implementation, during the iterative reasoning process of the AI ​​model, the network layer in the AI ​​model will perform data processing and generate KV during each reasoning process. Among them, the AI ​​model reasoning device will use the physical storage area 1 indicated by the virtual address space 1 allocated to the network layer to store the KV generated by the network layer. During the Nth reasoning process of the AI ​​model (N is a positive integer), the network layer will generate a new KV during this reasoning process, and for the newly generated KV of the network layer, if the remaining available area in the physical storage area 1 can continue to store the newly generated KV, the AI ​​model reasoning device can use the physical storage area 1 to store the newly generated KV. If the physical storage area 1 cannot continue to store the newly generated KV, the AI ​​model reasoning device will allocate a new physical storage area for the network layer, specifically, allocate the physical storage area 2 indicated by the virtual address space 2, and use the newly allocated physical storage area 2 to store the newly generated KV, wherein the virtual address space 1 is continuous with the virtual address space 2, that is, the next address of the tail address of the virtual address space 1 is the first address of the virtual address space 2. By analogy, the AI ​​model inference device can refer to the above method and store the KV generated by the network layer in multiple inference processes by dynamically allocating physical storage areas until the iterative inference is completed.

[0047] Because during the AI ​​model inference process, the AI ​​model inference device does not pre-allocate a fixed-length physical storage area for the network layer to store the KV generated during each inference process, but when the existing physical storage area (i.e., physical storage area 1) is insufficient to continue to store the newly generated KV, it dynamically allocates a new physical storage area (i.e., physical storage area 2) to the network layer, and uses the dynamically allocated new physical storage area to store the newly generated KV. In this way, by allocating physical storage areas on demand to store the KV generated during multiple inference processes, the problem of storage resource waste caused by allocating fixed-length physical storage areas to the network layer can be effectively alleviated, thereby improving the utilization rate of the physical storage area and reducing the storage resource fragmentation generated during the model inference process.

[0048] Moreover, since the first virtual address space for storing KVs is continuous with the second virtual address space, the operators in the network layer (such as the attention operator and its variants) do not need to splice the read KVs in the virtual address space when reading KVs in different physical storage areas (that is, splicing them in the order in which the KVs are generated so that subsequent calculations can be performed based on the sequentially generated KVs). Therefore, the operators in the network layer are not required to have the ability to splice KVs in the virtual address space, which can adapt to various variants of operators in the AI ​​model and improve the generalization of the operators.

[0049] For ease of understanding, an embodiment of the model reasoning method provided in this application is described below in conjunction with the accompanying drawings.

[0050] See Figure 2, which is a flow chart of a model inference method provided in an embodiment of the present application. For ease of explanation, this embodiment uses the example of an AI model performing P inferences when performing an inference task to introduce the process of allocating physical storage areas to cache key values ​​by the AI ​​model inference device.

[0051] As shown in Figure 2, the model inference method may specifically include:

[0052] S201: The AI ​​model inference device obtains input data of the AI ​​model.

[0053] Among them, the AI ​​model, for example, can be a generative reasoning model such as a large language model (LLM), which can include large-scale model parameters, such as hundreds of millions of model parameters. Alternatively, the AI ​​model can be other types of reasoning models, which are not limited to this. The AI ​​model can include multiple network layers, and in a reasoning process, the multiple network layers can perform corresponding calculation processes in sequence, wherein the calculation results output by the previous network layer can be used as input to the next network layer. Each network layer can include one attention head, or can include multiple attention heads, each attention head is used to perform corresponding attention calculations. Among them, when a network layer includes multiple attention heads, the network layer can aggregate the calculation results obtained by multiple attention heads respectively, so as to improve the spatial resolution of the attention mechanism and reduce the possibility of overfitting.

[0054] For example, the input data obtained by the AI ​​model inference device can be any type of data. For example, the input data can be text data, such as "Please introduce the characteristics of apple ripening" input by the user; or the input data can be image data, such as an image input by the user; or the input data can be audio or video data.

[0055] In one possible implementation, the AI ​​model inference device can provide services when running the AI ​​model, such as cloud services that provide intelligent dialogue. In addition, the AI ​​model inference device can provide a client to the outside world. The client can, for example, be an application running on the user-side device, or a web browser, so that when the user uses the service, he can interact with the AI ​​model inference device through the client. The client can create an inference task based on the user's operation and feed back the inference task and the data provided by the user to the AI ​​model inference device. Thus, the AI ​​model inference device can use the data sent by the client as the input data of the AI ​​model and start executing the inference task. Accordingly, after the AI ​​model completes the inference task, the AI ​​model inference device can feed back the final inference result to the user through the client.

[0056] S202: The AI ​​model reasoning device uses the AI ​​model and the acquired input data to perform the first reasoning process and obtains the reasoning result 1.

[0057] In this embodiment, during the execution of the inference task, the AI ​​model inference device can use the AI ​​model to perform P inferences, and can integrate the inference results output by the AI ​​model after each inference as the final inference result. For example, assuming that the inference result output by the AI ​​model each time is a token, the final inference result can be, for example, "token1+token2+...+token" as shown in Figure 1 above. P-1 ”.

[0058] When performing the first inference, after obtaining the input data, the AI ​​model inference device can use the AI ​​model to infer the obtained input data, and the AI ​​model outputs the corresponding inference result, which is referred to as inference result 1 in this embodiment for easy distinction.

[0059] Exemplarily, the inference result 1 can be a token in the complete inference result finally fed back by the AI ​​model inference device, such as token1 shown in Figure 1, etc. The token1 can be, for example, a word, letter, number, symbol, etc., and there is no limitation on this.

[0060] S203: The AI ​​model inference device uses the inference result 1 as a new input of the AI ​​model, and uses the AI ​​model to perform a second inference process on the inference result 1 to obtain the inference result 2.

[0061] For example, in the process of the AI ​​model reasoning based on the inference result 1, each network layer / attention head in the AI ​​model will perform data processing, specifically generating corresponding KV for the inference result 1, and the K and V generated by each network layer / attention head can be represented by corresponding vector matrices. For ease of distinction, K1 and V1 are used for description below. In actual application, each network layer / attention head can also generate a query value (query), which can also be represented by a corresponding vector matrix.

[0062] In this embodiment, the AI ​​model inference device can store the K1 and V1 generated by each network layer / attention head, so that in each subsequent inference process of the AI ​​model (i.e., the 3rd to the Pth inference), it can directly read the K1 and V1 corresponding to the inference result 1 (for performing the corresponding calculation) without having to repeatedly calculate the K1 and V1. For ease of description, the processing module is used below to refer to the network layer or attention head that can generate KV.

[0063] Specifically, as shown in Figure 2, before the AI ​​model performs a second inference, the AI ​​model inference device can allocate a fixed-size target virtual address space for each processing module in the AI ​​model. Furthermore, the target virtual address space allocated by the AI ​​model inference device for each processing module may not be mapped to a specific physical storage area. The allocated target virtual address space may include multiple contiguous virtual address spaces. For ease of understanding and explanation, the following example uses the target virtual address space comprising a first virtual address space and a second virtual address space. The first virtual address space is contiguous with the second virtual address space, meaning that the next address after the last address in the first virtual address space is the first address in the second virtual address space. Furthermore, the size of the target virtual address space can be the maximum size supported by the AI ​​model inference device. For example, the AI ​​model inference device can determine the size of the target virtual address space as ((Q-1) * the storage space required to store the KV generated during a single inference process) based on the maximum number of inferences Q that the AI ​​model can perform. Alternatively, the size of the target virtual address space can be configured by a technician.

[0064] Taking the allocation of a physical storage area for one of the processing modules in the AI ​​model as an example, as shown in Figure 2, before performing the second inference, the AI ​​model inference device can determine an available first physical storage area for the processing module and establish a mapping between the first physical storage area and the first virtual address space. In this way, during the second inference of the AI ​​model, the AI ​​model inference device can store K1 and V1 calculated by the processing module in the first physical storage area according to the first virtual address space, as shown in Figure 2.

[0065] It is worth noting that if the size of the first physical storage area allocated to the processing module is too large, this may easily lead to a large number of areas in the first physical storage area not storing data during the entire process of executing the inference task, thereby causing resource waste and a large amount of resource fragmentation. Therefore, the first physical storage area allocated by the AI ​​model inference device to the processing module can be a smaller physical storage area, such as a physical storage area in units of pages. The size of the physical storage area can be 1 megabyte (M) or 2M, etc., and can be set according to the needs of the actual application, and there is no limitation on this.

[0066] In actual application, the AI ​​model inference device can pre-create a storage resource pool, which can include multiple available physical storage areas, and the sizes of the multiple available physical storage areas can be the same, such as 2M storage blocks. In this way, when the AI ​​model inference device allocates the first physical storage area to the processing module, it can select an available physical storage area from the multiple physical storage areas included in the storage resource pool, such as by selecting through a random algorithm, so that the selected physical storage area can be used as the first physical storage area and mapped to the first virtual address space. Furthermore, the AI ​​model inference device can also remove the first physical storage area that has been allocated to the processing module from the storage resource pool to avoid the subsequent repeated allocation of the first physical storage area to other processing modules.

[0067] Furthermore, when the AI ​​model inference device allocates a target virtual address space to the processing module, it can also create metadata for the target virtual address space. The metadata can include the starting address of the target virtual address space, as shown in FIG3 , and can also include information such as the size of the target virtual address space. After the AI ​​model inference device establishes a mapping between the first virtual address space and the first physical storage area, the AI ​​model inference device can update the metadata, wherein the updated metadata can also include a list of physical storage areas that have completed the mapping, the number of remaining KVs that can be stored, etc. Among them, the list of physical storage areas can include information such as the starting address and identification of the physical storage area. Moreover, before storing K1 and V1, the number of remaining KVs that can be stored can be n (n is a positive integer greater than 1), and after the storage of K1 and V1, the AI ​​model inference device can update the number of remaining KVs that can be stored in the metadata to (n-1).

[0068] The above process is explained by taking the example of allocating target virtual address space, allocating physical storage area and creating metadata for the processing module. For other processing modules in the AI ​​model, the AI ​​model inference device can refer to the above method to perform similar processing for other processing modules, which will not be repeated here.

[0069] S204: The AI ​​model inference device splices the inference result 1 and the inference result 2 to obtain a splicing result, and uses the splicing result as a new input of the AI ​​model, and uses the AI ​​model to perform a third inference process on the splicing result to obtain an inference result 3.

[0070] Similar to the second inference process of the AI ​​model, the AI ​​model inference device can continue to use the unused area in the allocated first physical storage area to store the newly generated K2 and V2 for inference result 2 when the processing module performs calculations, and update the number of remaining KVs that can be stored in the metadata to (n-2). The order in which the processing module generates KVs can be indicated by the order of the virtual addresses used when storing the KVs.

[0071] By analogy, the AI ​​model inference device can refer to the above method to complete N inferences of the AI ​​model (N is a positive integer) and obtain the inference result N. At this time, the KV stored in the first physical storage area includes the new KVs generated in the second inference to the Nth inference process, namely K1 to K N-1 , and V1 to V N-1 , as shown in Figure 3.

[0072] In this embodiment, the size of the first physical storage area allocated to the processing module is relatively small and may not be able to store all the KVs generated by the processing module during the execution of the inference task. Therefore, it is assumed that after the first physical storage area stores the new KVs generated in the second to Nth inference processes, the number of KVs that can be stored is 0. The following describes the process in which the AI ​​model inference device uses the AI ​​model to continue to perform the (N+1)th inference and subsequent inferences after the first physical storage area is fully used.

[0073] S205: The AI ​​model inference device splices the inference result 1 and the inference result N to obtain a splicing result, and uses the splicing result as a new input of the AI ​​model, and uses the AI ​​model to perform the (N+1)th inference process on the splicing result to obtain the inference result (N+1).

[0074] Since all the storage space in the first physical storage area has been used to store the KVs generated during the first N inference processes, when the AI ​​model performs the (N+1)th inference, the AI ​​model inference device can allocate a new physical storage area to the processing module, which is hereinafter referred to as the second physical storage area.

[0075] In specific implementation, as shown in Figure 2, before the AI ​​model performs the (N+1)th reasoning, the AI ​​model reasoning device can determine whether the remaining number of KVs that can be stored is 0 based on the metadata corresponding to the target virtual address space. If not, the AI ​​model reasoning device can continue to use the first physical storage area that has been mapped to the target virtual address space to store the newly generated KV. If so, it indicates that the first physical storage area has been used up. At this time, as shown in Figure 2, the AI ​​model reasoning device can select a new available physical storage area from the storage resource pool, and map the physical storage area (as the second physical storage area) to the second virtual address space, so as to allocate a new physical storage area to the processing module. Then, the AI ​​model reasoning device can update the metadata of the target virtual address space, including updating the list of physical storage areas that have been mapped and the remaining number of KVs that can be stored (that is, the number of KVs that can be stored in the second physical storage area).

[0076] In this way, during the AI ​​model's (N+1)th reasoning process, the AI ​​model reasoning device can generate the K address generated by the processing module for the reasoning result (N+1) according to the second virtual address space. N and V N , stored in the newly allocated second physical storage area, as shown in Figure 2. In addition, the AI ​​model inference device can also update the number of KVs that can be stored in the metadata (specifically, the number - 1). Then, the AI ​​model inference device can continue to perform the (N+2)th and subsequent inference processes on the AI ​​model.

[0077] Similarly, during the process of performing the (N+2)th to the Pth reasoning for the AI ​​model, the AI ​​model reasoning device can continue to use the second physical storage area to store the newly generated KV of the processing module, as shown in Figure 3. Moreover, when the second physical storage area is fully used, the AI ​​model reasoning device can refer to the above-mentioned method of dynamically allocating the second physical storage area to the processing module, continue to allocate a new physical storage area to the processing module, and use the newly allocated physical storage area to store the newly generated KV until the AI ​​model completes the reasoning task. The following continues to introduce the process of the AI ​​model reasoning device determining whether the AI ​​model has completed the reasoning task. It is assumed that the AI ​​model has currently completed P reasonings and output the reasoning result P.

[0078] It should be noted that in the above embodiment, the AI ​​model inference device determines whether the allocated first physical storage area supports the storage of a newly generated KV based on whether the number of KVs that can be stored in the remaining storage of the operation and maintenance is 0. In other embodiments, the AI ​​model inference device can also determine whether the allocated first physical storage area supports the storage of a new KV in other ways. Some exemplary explanations are given below.

[0079] Example 1: After the processing module generates a new KV, the AI ​​model reasoning device can determine the virtual address used to store the new KV, and determine whether the virtual address is mapped to a physical address. If the virtual address is mapped to a physical address (the physical address is a physical address in the first physical storage area), the AI ​​model reasoning device can store the newly generated KV to the storage location indicated by the physical address. If the virtual address is not mapped to a physical address), indicating that the first physical storage area has been used up, the AI ​​model reasoning device can determine that the first physical storage area is insufficient to store the newly generated KV of the processing module, and can further dynamically allocate a new physical storage area (that is, the above-mentioned second physical storage area) to the processing module to store the newly generated KV.

[0080] Example 2: The AI ​​model inference device can detect whether there is any remaining storage space in the first physical storage area after the processing module generates a new KV. When there is any remaining storage space in the first physical storage area, the AI ​​model inference device can continue to use the remaining storage space in the first physical storage area to store the newly generated KV. When there is no remaining storage space in the first physical storage area, indicating that the first physical storage area has been used up, the AI ​​model inference device can determine that the first physical storage area is insufficient to store the newly generated KV of the processing module, and can further dynamically allocate a new physical storage area to the processing module to store the newly generated KV.

[0081] S206: The AI ​​model reasoning device determines whether the reasoning result P is a terminator, and when it is determined that the reasoning result P is a terminator, determines that the reasoning task is completed, and outputs the reasoning result 1 to the reasoning result (P-1).

[0082] In actual applications, after each inference process of the AI ​​model, the AI ​​model inference device can determine whether the inference result currently output by the AI ​​model is a terminator. If not, the AI ​​model inference device can continue to use the AI ​​model to perform the next inference process; if so, the AI ​​model inference device can output each inference result obtained by the AI ​​model inference.

[0083] In this way, during the AI ​​model's execution of the reasoning task, the AI ​​model reasoning device dynamically establishes a mapping between the physical storage area and the virtual address space based on the need to store the newly generated KV of the processing module, thereby dynamically allocating the physical storage area to the processing module and avoiding the physical storage area allocated to the processing module being in a state of not storing data during the entire process, thereby effectively improving the utilization rate of physical storage resources and reducing the number of fragments of physical storage resources. At the same time, the AI ​​model reasoning device dynamically allocates the virtual address space mapped to the physical storage area for the processing module to be continuous, so that when reading the KV generated by the processing module, the KV generated in sequence by the processing module can be determined based on the continuous virtual address space, so that the second processing module can execute the corresponding operator based on the KV. Among them, the KV calculated by the processing module can be used as the input of the operator in the second processing module. In this way, during the AI ​​model reasoning process, the AI ​​model reasoning device can determine the order of the various KVs generated by the processing module without modifying the operator in the second processing module, thereby not only improving the generalization of the operator, but also improving the efficiency of the second processing module in calculating based on the KV, thereby improving the reasoning efficiency of the AI ​​model.

[0084] Furthermore, after the AI ​​model completes the reasoning task, the AI ​​model reasoning device can release the physical storage area allocated to each processing module. Taking the processing module as an example, the AI ​​model reasoning device can release the mapping between each physical storage area and the virtual address space based on the metadata of the target virtual address space, and release each physical storage area to the storage resource pool, so that the physical storage area becomes an available physical storage area in the storage resource pool. In addition, the AI ​​model reasoning device can also update the metadata, including clearing the list of physical storage areas in the metadata (to indicate that no physical storage area is currently mapped) and clearing the number of remaining KVs that can be stored in the metadata (to indicate that KV cannot be stored currently). When the AI ​​model performs the next reasoning task, the AI ​​model reasoning device can reallocate the physical storage area to the processing module and update the metadata.

[0085] It should be noted that in this embodiment, the AI ​​model performs an inference task on a set of input data as an example. In actual application, the AI ​​model can perform the inference process in parallel on multiple sets of input data, that is, the batch size is greater than 1. For example, the AI ​​model inference device can provide services to multiple users at the same time based on the AI ​​model, so that multiple users can provide input data to the AI ​​model inference device through their respective clients. At this time, the AI ​​model inference device can create multiple inference tasks for the input data provided by different users, and can use the AI ​​model to execute the multiple inference tasks in parallel. In this process, the AI ​​model inference device can create multiple target virtual address spaces for each processing module in the AI ​​model, and each target virtual address space corresponds to an inference task. Therefore, the AI ​​model inference device can use the process of the embodiment shown in Figure 2 to dynamically allocate physical storage areas to the processing module during the execution of each inference task by the AI ​​model, thereby improving the utilization rate of physical storage resources and improving the generalization of operators in the processing module.

[0086] In the embodiment shown in FIG2 above, the processing module is taken as an example of a network layer or an attention head in the AI ​​model. In other embodiments, the processing module may also include multiple network layers in the AI ​​model, or may be multiple attention heads included in a network layer. In this case, the KVs generated by multiple network layers or multiple attention heads can be stored in the same physical storage area, that is, the KVs are aggregated. For example, in a multi-query attention (MQA) scenario, multiple attention heads in the same network layer share the same KV, so the AI ​​model inference device can aggregate the KVs of multiple network layers for storage.

[0087] Below, in conjunction with Figure 4, the process of dynamically allocating physical storage areas for processing modules by the AI ​​model inference device in the scenario of aggregated storage KV is introduced.

[0088] Refer to Figure 4, which shows a flow chart of another model inference method. For ease of understanding, this embodiment uses the example of aggregating and storing the KV generated by M network layers or M attention heads. As shown in Figure 4, the method may specifically include:

[0089] S401: The AI ​​model inference device obtains input data of the AI ​​model.

[0090] S402: The AI ​​model reasoning device uses the AI ​​model to perform the first reasoning process on the acquired input data to obtain the reasoning result 1.

[0091] The implementation of steps S401 to S402 can refer to the relevant description of steps S201 to S202 in the embodiment shown in FIG2 , and will not be described in detail here.

[0092] S403: The AI ​​model inference device uses the inference result 1 as a new input of the AI ​​model, and uses the AI ​​model to perform a second inference process on the inference result 1 to obtain the inference result 2.

[0093] In this embodiment, as shown in FIG4 , before the AI ​​model performs the second inference, the AI ​​model inference device may allocate a fixed-size target virtual address space for each processing module in the AI ​​model, and the target virtual address space allocated by the AI ​​model inference device to each processing module may not be mapped to a specific physical storage area. The allocated target virtual address space may include multiple continuous virtual address spaces. In addition, the AI ​​model inference device will also create metadata for the target virtual address space, which may include the first address of the target virtual address space and other information.

[0094] Each processing module may include multiple network layers or multiple attention heads in a network layer. For ease of description, the following example illustrates a processing module including M sub-processing modules, where M is a positive integer greater than 1, and each sub-processing module may be a network layer or an attention head. The KVs generated by the M sub-processing modules may be stored in the same physical storage area.

[0095] Moreover, before the second reasoning is performed, the AI ​​model reasoning device can determine an available first physical storage area for the processing module from the storage resource pool, and establish a mapping between the first physical storage area and the first virtual address space in the target virtual address space. Exemplarily, the size of the first physical storage area can be a smaller physical storage area, such as a physical storage area in units of pages, and the size of the physical storage area can be 2M, etc. Since the first physical storage area will store the KVs generated by M sub-processing modules respectively, the AI ​​model reasoning device can further divide the first virtual address space into M sub-address spaces, each sub-address space being used to indicate a portion of the physical storage area in the first physical storage area, as shown in FIG4 .

[0096] For example, assuming that V of M is 2, that is, the processing module includes sub-processing module 1 and sub-processing module 2, then the AI ​​model inference device can divide the first virtual address space into 2 sub-address spaces, as shown in Figure 5, wherein the physical storage area indicated by sub-address space 1 is used to store the KV generated by sub-processing module 1, and the physical storage area indicated by sub-address space 2 is used to store the KV generated by sub-processing module 2.

[0097] In the process of the AI ​​model reasoning based on the reasoning result 1, each sub-processing module in the M sub-processing modules can generate a corresponding KV for the reasoning result 1, and the KV generated by each sub-processing module can be represented by a corresponding vector matrix. For the sake of distinction, the KV generated by each sub-processing module in the first reasoning process is described below as K1 and V1, and the KVs generated by different sub-processing modules are different. In this way, after each sub-processing module generates K1 and V1, the AI ​​model reasoning device can store the K1 and V1 generated by the M sub-processing modules in the first physical storage area according to the first virtual address space, as shown in Figure 4. Specifically, the K1 and V1 generated by each sub-processing module are stored in the physical storage area indicated by a sub-address space. Then, the AI ​​model reasoning device can update the metadata of the target virtual address space, such as updating the number of remaining KVs that can be stored.

[0098] S404: The AI ​​model inference device splices the inference result 1 and the inference result 2 to obtain a splicing result, and uses the splicing result as a new input of the AI ​​model, and uses the AI ​​model to perform a third inference process on the splicing result to obtain an inference result 3.

[0099] Among them, similar to the process of the AI ​​model performing the second reasoning, the AI ​​model reasoning device can continue to use the allocated first physical storage area to store the K2 and V2 newly generated by the N sub-processing modules for the reasoning result 2, and update the remaining number of KVs that can be stored in the metadata.

[0100] By analogy, the AI ​​model reasoning device can refer to the above method to complete N times of reasoning of the AI ​​model (N is a positive integer) and obtain the reasoning result N, which will not be described in detail here. At this time, the KV stored in the first physical storage area includes the new KVs generated in the second reasoning to the Nth reasoning process, namely K1 to K N-1 , and V1 to V N-1 .

[0101] Since the size of the first physical storage area is small, the K newly generated by the M sub-processing modules is stored. N-1 and V N-1 After the first physical storage area is fully used, it is impossible to continue to store new KVs. At this time, the AI ​​model inference device can dynamically allocate new physical storage areas to the M sub-processing modules to store the KVs generated by the M sub-processing modules in subsequent inference processes. The following continues to describe the process of the AI ​​model inference device using the AI ​​model to continue to perform the (N+1)th inference and subsequent inferences after the first physical storage area is fully used.

[0102] S405: The AI ​​model inference device splices the inference result 1 and the inference result N to obtain a splicing result, and uses the splicing result as a new input of the AI ​​model, and uses the AI ​​model to perform the (N+1)th inference process on the splicing result to obtain the inference result (N+1).

[0103] In a specific implementation, before the AI ​​model performs the (N+1)th reasoning, the AI ​​model reasoning device can determine whether the number of remaining KVs that can be stored is 0 based on the metadata corresponding to the target virtual address space. When it is determined that the number of remaining KVs that can be stored is 0, it indicates that the first physical storage area has been used up. At this time, as shown in Figure 4, the AI ​​model reasoning device can select an available second physical storage area from the storage resource pool, and map the second physical storage area to the second virtual address space in the target virtual address space, so as to allocate new physical storage areas to the M sub-processing modules. Among them, the first virtual address space is continuous with the second virtual address space.

[0104] In this embodiment, in the target virtual address space, a continuous section of virtual address space can be used to store multiple KVs generated by each sub-processing module. Therefore, after dynamically allocating the second physical storage area to the M sub-processing modules, the AI ​​model inference device can redivide the virtual address space that has been mapped to the physical storage area (i.e., the continuous first virtual address space and the second virtual address space) into M sub-address spaces, and use the physical storage area indicated by the redivided M sub-address spaces to store the KVs generated by the M sub-processing modules.

[0105] For example, assuming that the value of M is 2, before allocating the second physical storage area, the sub-address space used to store the KV generated by each sub-processing module is shown in the upper part of Figure 6; and after allocating the second physical storage area, the AI ​​model inference device can re-divide the virtual address space of the mapped physical storage area into 2 sub-address spaces, as shown in the lower part of Figure 6, and still use continuous virtual addresses to store the KV generated by the same sub-processing module. During this process, the AI ​​model inference device can migrate the KV generated by the sub-processing module 2 originally stored in the sub-address space 1 to the physical storage area indicated by the sub-address space 2. Then, the AI ​​model inference device can use the unused physical storage area indicated by the sub-address space 1 to store the newly generated KV of the sub-processing module 1, and use the physical storage area indicated by the sub-address space 2 to store the newly generated KV of the sub-processing module 2, as shown in the KV in Figure 6. N and V N .

[0106] After completing the storage of the newly generated KVs of the M sub-processing modules, the AI ​​model inference device can update the metadata of the target virtual address space, including updating the list of physical storage areas that have been mapped, the number of remaining KVs that can be stored, etc.

[0107] Similarly, during the (N+2)th to Pth reasoning process for the AI ​​model, the AI ​​model reasoning device can continue to use the physical storage area indicated by the re-divided sub-address space to store the newly generated KVs of each sub-processing module. Moreover, when the physical storage area indicated by each sub-address space is fully used, the AI ​​model reasoning device can refer to the above-mentioned method of dynamically allocating a second physical storage area to the M sub-processing modules, continue to allocate new physical storage areas to the M sub-processing modules, and re-divide the virtual address space of the mapped physical storage area, and use the physical storage area indicated by the equally divided sub-address space to store the newly generated KV until the AI ​​model completes the reasoning task.

[0108] S406: The AI ​​model reasoning device determines whether the reasoning result P is a terminator, and when it is determined that the reasoning result P is a terminator, determines that the reasoning task is completed, and outputs the reasoning result 1 to the reasoning result (P-1).

[0109] In this way, during the process of the AI ​​model performing reasoning tasks, the AI ​​model reasoning device can dynamically allocate physical storage areas to M sub-processing modules, which not only improves the utilization rate of physical storage resources and reduces the number of fragments of physical storage resources, but also the KV generated by each sub-processing module can be stored in the physical storage area indicated by a continuous virtual address, which can improve the generalization of operators in the sub-processing modules and improve the computing efficiency of the sub-processing modules.

[0110] When the AI ​​model inference device calls the interface to read the KVs generated by multiple sub-processing modules, since multiple sub-processing modules can share the same physical storage area to store KVs, the AI ​​model inference device only needs to call the interface once to obtain the KVs generated by multiple sub-processing modules in the physical storage area. This can effectively reduce the number of times the AI ​​model inference device calls the interface, that is, reduce the overhead generated by the AI ​​model inference device calling the interface.

[0111] It is worth noting that other reasonable step combinations that can be thought of by those skilled in the art based on the above description also fall within the scope of protection of this application. Secondly, those skilled in the art should also be familiar with that the embodiments described in the specification are all preferred embodiments, and the actions involved are not necessarily required by this application.

[0112] The above introduces the AI ​​model reasoning method provided in the embodiment of the present application in combination with Figures 1 to 6. Next, the structure of the AI ​​model reasoning device and computing equipment provided in the embodiment of the present application is introduced in combination with the accompanying drawings.

[0113] 7 , which shows a schematic diagram of the structure of a processor, the AI ​​model inference device 700 includes:

[0114] An inference unit 701 is configured to obtain input data and perform iterative inference using an AI model and the input data to obtain an inference result, wherein the AI ​​model includes a processing module, the processing module corresponds to a first physical storage area indicated by a first virtual address space, and the first physical storage area is configured to store a key value (KV) generated by the processing module when the AI ​​model performs iterative inference. During an Nth inference process, the processing module performs data processing to generate a KV for the Nth inference process.

[0115] an allocating unit 702 configured to allocate a second physical storage area indicated by a second virtual address space to the processing module when the first physical storage area is insufficient to continue storing the KV generated by the processing module in the Nth inference process, wherein the second virtual address space is continuous with the first virtual address space, and N is a positive integer;

[0116] The inference unit 701 is further configured to store the KV generated in the Nth inference process in the second physical storage area according to the second virtual address space.

[0117] In one possible implementation, the allocating unit 702 is further configured to allocate a target virtual address space of a fixed size to the processing module before performing iterative inference using the AI ​​model, where the target virtual address space includes a first virtual address space and a second virtual address space with continuous virtual addresses;

[0118] The allocation unit 702 is specifically configured to:

[0119] determining a second physical storage area;

[0120] A mapping between the second virtual address space and the second physical storage area is established.

[0121] In a possible implementation, the allocating unit 702 is specifically configured to select a second physical storage area from a plurality of available physical storage areas included in the storage resource pool, where each of the plurality of physical storage areas has the same size.

[0122] In one possible embodiment, the allocation unit 702 is also used to release multiple physical storage areas allocated to the processing module during the iterative reasoning process of the AI ​​model after the AI ​​model completes the iterative reasoning, and the multiple physical storage areas include a first physical storage area and a second physical storage area.

[0123] In a possible implementation, the allocating unit 702 is further configured to:

[0124] determining whether a virtual address for storing a KV generated by the processing module during the Nth inference process is mapped to a physical address, and determining that the first physical storage area is insufficient to continue storing the KV generated during the Nth inference process when the virtual address is not mapped to the physical address;

[0125] Alternatively, detecting whether there is any remaining storage space in the first physical storage area, and when there is no remaining storage space in the first physical storage area, determining that the first physical storage area is insufficient to continue storing the KV generated in the Nth inference process;

[0126] Alternatively, it is detected whether the remaining number of KVs that can be stored in the first physical storage area is zero, and when the remaining number of KVs that can be stored is zero, it is determined that the first physical storage area is insufficient to continue storing the KVs generated in the Nth reasoning process.

[0127] In one possible embodiment, the processing module includes M sub-processing modules, the first virtual address space includes M first sub-address spaces, the physical storage area indicated by each first sub-address space in the M first sub-address spaces is used to store the KV generated by a sub-processing module, and the physical storage areas indicated by different first sub-address spaces are used to store the KVs generated by different sub-processing modules, and M is a positive integer greater than 1.

[0128] In one possible embodiment, the allocation unit 702 is also used to: after allocating a second physical storage area to the processing module, divide the first virtual address space and the second virtual address space equally into M second sub-address spaces, each of the M second sub-address spaces is used to store the KV generated by a sub-processing module, and different second sub-address spaces are used to store the KVs generated by different sub-processing modules.

[0129] In one possible implementation, the processing module includes multiple network layers in the AI ​​model;

[0130] Alternatively, the processing module includes a network layer in the AI ​​model;

[0131] Alternatively, the processing module includes multiple attention heads in a network in the AI ​​model;

[0132] Alternatively, the processing module includes an attention head in the AI ​​model.

[0133] Since the AI ​​model reasoning device 700 shown in Figure 7 corresponds to the AI ​​model reasoning device in the embodiments shown in Figures 2 and 4 above, the specific implementation method of the AI ​​model reasoning device 700 shown in Figure 7 and its technical effects can be found in the relevant descriptions in the embodiments shown in Figures 2 and 4 above, and will not be repeated here.

[0134] Figure 8 is a schematic diagram of the hardware structure of a computing device 800 provided in the present application. The computing device 800 can, for example, implement the AI ​​model inference device in the embodiments shown in Figures 2 and 4 above. As shown in Figure 8, the computing device 800 includes a processor 801, a memory 802, and a communication interface 803. Among them, the processor 801, the memory 802, and the communication interface 803 communicate through a bus 804, and can also communicate through other means such as wireless transmission. The memory 802 is used to store instructions, and the processor 801 is used to execute the instructions stored in the memory 802. Furthermore, the computing device 800 may also include a memory unit 805, and the memory unit 805 may be connected to the processor 801, the storage medium 802, and the communication interface 803 through a bus 804. Among them, the memory 802 stores program code, and the processor 801 can call the program code stored in the memory 802 to perform the following operations:

[0135] Get input data;

[0136] Performing iterative reasoning using the AI ​​model and the input data to obtain a reasoning result, wherein the AI ​​model includes a processing module, the processing module corresponds to a first physical storage area indicated by a first virtual address space, and the first physical storage area is used to store a key value KV generated by the processing module when the AI ​​model performs iterative reasoning;

[0137] During the Nth reasoning process, the processing module performs data processing to generate the KV of the Nth reasoning process;

[0138] When the first physical storage area is insufficient to continue storing the KV generated by the processing module during the Nth inference process, allocating a second physical storage area indicated by a second virtual address space to the processing module, wherein the second virtual address space is continuous with the first virtual address space, and N is a positive integer;

[0139] The KV generated in the Nth inference process is stored in the second physical storage area according to the second virtual address space.

[0140] It should be understood that in this embodiment, the processor 801 may be a CPU, or may be other general-purpose processors, digital signal processors (DSPs), application-specific integrated circuits (ASICs), field-programmable gate arrays (FPGAs) or other programmable logic devices, discrete gate or transistor logic devices, discrete device components, etc. A general-purpose processor may be a microprocessor or any conventional processor, etc.

[0141] The memory 802 may include a read-only memory and a random access memory, and provides instructions and data to the processor 801. The memory 802 may also include a nonvolatile random access memory.

[0142] The memory 802 may be a volatile memory or a non-volatile memory, or may include both volatile and non-volatile memories. The non-volatile memory may be a read-only memory (ROM), a programmable read-only memory (PROM), an erasable programmable read-only memory (EPROM), an electrically erasable programmable read-only memory (EEPROM), or a flash memory. The volatile memory may be a random access memory (RAM), which is used as an external cache. By way of example and not limitation, many forms of RAM are available, such as static RAM (SRAM), dynamic random access memory (DRAM), synchronous DRAM (SDRAM), double data rate synchronous dynamic random access memory (DDR SDRAM), enhanced synchronous dynamic random access memory (ESDRAM), synchronous link DRAM (SLDRAM), and direct rambus RAM (DR RAM).

[0143] The communication interface 803 is used to communicate with other devices connected to the computing device 800. The bus 804 includes not only a data bus but also a power bus, a control bus, and a status signal bus. However, for the sake of clarity, all buses are labeled as bus 804 in the figure.

[0144] It should be understood that the computing device 800 according to the embodiment of the present application may correspond to the AI ​​model inference device 700 in the embodiment of the present application, and may correspond to the method executed by the AI ​​model inference device in the method shown in Figures 2 and 4 in the embodiment of the present application. The above-mentioned and other operations and / or functions implemented by the computing device 800 are respectively for implementing the processes of the corresponding methods in Figures 2 and 4. For the sake of brevity, they will not be repeated here.

[0145] The embodiment of the present application also provides a computer-readable storage medium. The computer-readable storage medium can be any available medium that can be stored by a computing device or a data storage device such as a data center containing one or more available media. The available medium can be a magnetic medium (for example, a floppy disk, a hard disk, a tape), an optical medium (for example, a DVD), or a semiconductor medium (for example, a solid-state drive). The computer-readable storage medium includes instructions that instruct the computing device to execute the above-mentioned AI model reasoning method.

[0146] The present application also provides a computer program product. The computer program product includes one or more computer instructions. When the computer instructions are loaded and executed on a computing device, the computer program product fully or partially generates the process or function described in the present application.

[0147] The computer instructions may be stored in a computer-readable storage medium or transmitted from one computer-readable storage medium to another computer-readable storage medium. For example, the computer instructions may be transmitted from one website, computer, or data center to another website, computer, or data center via wired (e.g., coaxial cable, optical fiber, digital subscriber line (DSL)) or wireless (e.g., infrared, wireless, microwave, etc.) means.

[0148] The computer program product may be a software installation package, which may be downloaded and executed on a computing device when any of the aforementioned AI model inference methods is required.

[0149] The above embodiments can be implemented in whole or in part by software, hardware, firmware or any other combination. When implemented using software, the above embodiments can be implemented in whole or in part in the form of a computer program product. The computer program product includes one or more computer instructions. When the computer program instructions are loaded or executed on a computer, the process or function described in the embodiment of the present application is generated in whole or in part. The computer can be a general-purpose computer, a special-purpose computer, a computer network, or other programmable device. The computer instructions can be stored in a computer-readable storage medium or transmitted from one computer-readable storage medium to another computer-readable storage medium. For example, the computer instructions can be transmitted from one website, computer, server or data center to another website, computer, server or data center via a wired (e.g., coaxial cable, optical fiber, digital subscriber line (DSL)) or wireless (e.g., infrared, wireless, microwave, etc.) method. The computer-readable storage medium can be any available medium that a computer can access or a data storage device such as a server or data center that contains one or more available media sets. The available medium can be a magnetic medium (e.g., a floppy disk, a hard disk, a tape), an optical medium (e.g., a DVD), or a semiconductor medium. The semiconductor medium can be a solid-state drive.

[0150] The terms used in the above embodiments are only for the purpose of describing specific embodiments and are not intended to limit the present application. As used in the specification and claims of this application, the singular expressions "one", "a kind of", "said", "above", "the" and "this" are intended to also include expressions such as "one or more", unless the context clearly indicates otherwise. It should also be understood that in the embodiments of the present application, "one or more" refers to one, two or more; the character " / " generally indicates that the objects associated with each other are in an "or" relationship. In the embodiments of the present application. "Simultaneously" means within the same time period, including situations at the same time. The terms "first", "second", etc. in the specification, claims and drawings of this application are used to distinguish similar objects, and are not necessarily used to describe a specific order or sequence. It should be understood that the terms used in this way can be interchangeable where appropriate, and this is merely a way of distinguishing objects with the same properties when describing them in the embodiments of the present application.

[0151] References to "one embodiment" or "some embodiments" in this specification mean that a particular feature, structure, or characteristic described in conjunction with that embodiment is included in one or more embodiments of the present application. Thus, phrases such as "in one embodiment," "in some embodiments," "in other embodiments," and "in yet other embodiments" appearing in various places in this specification do not necessarily refer to the same embodiment, but rather mean "one or more but not all embodiments," unless otherwise specifically emphasized. The terms "including," "comprising," "having," and variations thereof mean "including but not limited to," unless otherwise specifically emphasized.

[0152] The above description is merely a specific embodiment of the present application, but the scope of protection of the present application is not limited thereto. Any person skilled in the art can easily conceive of various equivalent modifications or substitutions within the technical scope disclosed in the present application, and such modifications or substitutions should be included in the scope of protection of the present application. Therefore, the scope of protection of the present application should be based on the scope of protection of the claims.

Claims

1. An artificial intelligence (AI) model reasoning method, characterized in that: The method comprises: Get input data; Performing iterative reasoning using the AI model and the input data to obtain a reasoning result, wherein the AI model includes a processing module, the processing module corresponds to a first physical storage area indicated by a first virtual address space, and the first physical storage area is used to store a key value KV generated by the processing module when the AI model performs iterative reasoning; During the Nth reasoning process, the processing module performs data processing to generate the KV of the Nth reasoning process; When the first physical storage area is insufficient to continue storing the KV generated by the processing module during the Nth inference process, allocating a second physical storage area indicated by a second virtual address space to the processing module, wherein the second virtual address space is continuous with the first virtual address space, and N is a positive integer; The KV generated by the processing module during the Nth reasoning process is stored in the second physical storage area according to the second virtual address space.

2. The method according to claim 1, characterized in that The method further comprises: Before performing iterative reasoning using the AI model, allocating a fixed-size target virtual address space to the processing module, the target virtual address space including the first virtual address space and the second virtual address space having continuous virtual addresses; The allocating the second physical storage area indicated by the second virtual address space to the processing module includes: determining the second physical storage area; A mapping between the second virtual address space and the second physical storage area is established.

3. The method according to claim 2, characterized in that The determining the second physical storage area includes: The second physical storage area is selected from a plurality of available physical storage areas included in a storage resource pool, and each of the plurality of physical storage areas has the same size.

4. The method according to claim 2 or 3, characterized in that The method further comprises: After the AI model completes iterative reasoning, multiple physical storage areas allocated to the processing module during the iterative reasoning process of the AI model are released, and the multiple physical storage areas include the first physical storage area and the second physical storage area.

5. The method according to any one of claims 1 to 4, characterized in that The method further comprises: determining whether a virtual address for storing the KV generated by the processing module during the Nth inference process is mapped to a physical address, and determining that the first physical storage area is insufficient to continue storing the KV generated during the Nth inference process when the virtual address is not mapped to a physical address; Alternatively, detecting whether there is any remaining storage space in the first physical storage area, and when there is no remaining storage space in the first physical storage area, determining that the first physical storage area is insufficient to continue storing the KV generated in the Nth inference process; Alternatively, it is detected whether the remaining number of KVs that can be stored in the first physical storage area is zero, and when the remaining number of KVs that can be stored is zero, it is determined that the first physical storage area is insufficient to continue storing the KVs generated in the Nth reasoning process.

6. The method according to any one of claims 1 to 4, characterized in that The processing module includes M sub-processing modules, the first virtual address space includes M first sub-address spaces, the physical storage area indicated by each first sub-address space in the M first sub-address spaces is used to store the KV generated by a sub-processing module, and the physical storage areas indicated by different first sub-address spaces are used to store the KVs generated by different sub-processing modules, and M is a positive integer greater than 1.

7. The method according to claim 6, characterized in that After allocating the second physical storage area to the processing module, the method further includes: The first virtual address space and the second virtual address space are equally divided into M second sub-address spaces, each of the M second sub-address spaces is used to store the KV generated by a sub-processing module, and different second sub-address spaces are used to store the KV generated by different sub-processing modules.

8. The method according to any one of claims 1 to 7, characterized in that The processing module includes multiple network layers in the AI model; Alternatively, the processing module includes a network layer in the AI model; Alternatively, the processing module includes multiple attention heads in a network in the AI model; Alternatively, the processing module includes an attention head in the AI model.

9. An artificial intelligence (AI) model inference device, characterized in that: The device comprises: an inference unit, configured to obtain input data and perform iterative inference using the AI model and the input data to obtain an inference result, wherein the AI model includes a processing module, the processing module corresponds to a first physical storage area indicated by a first virtual address space, and the first physical storage area is used to store a key value KV generated by the processing module when the AI model performs iterative inference; wherein, during an Nth inference process, the processing module performs data processing to generate the KV for the Nth inference process; an allocating unit, configured to allocate a second physical storage area indicated by a second virtual address space to the processing module when the first physical storage area is insufficient to continue storing the KV generated by the processing module during the Nth inference process, wherein the second virtual address space is continuous with the first virtual address space, and N is a positive integer; The inference unit is further configured to store the KV generated by the processing module during the Nth inference process in the second physical storage area according to the second virtual address space.

10. The device according to claim 9, characterized in that The allocation unit is further configured to allocate a target virtual address space of a fixed size to the processing module before performing iterative inference using the AI model, wherein the target virtual address space includes the first virtual address space and the second virtual address space having continuous virtual addresses; The distribution unit is used to: determining the second physical storage area; A mapping between the second virtual address space and the second physical storage area is established.

11. The device according to claim 10, characterized in that The allocation unit is configured to select the second physical storage area from a plurality of available physical storage areas included in a storage resource pool, wherein each of the plurality of physical storage areas has the same size.

12. The device according to claim 10 or 11, characterized in that The allocation unit is also used to release multiple physical storage areas allocated to the processing module during the iterative reasoning of the AI model after the AI model completes the iterative reasoning, and the multiple physical storage areas include the first physical storage area and the second physical storage area.

13. The device according to any one of claims 9 to 12, characterized in that Allocation unit, also used for: determining whether a virtual address for storing the KV generated by the processing module during the Nth inference process is mapped to a physical address, and determining that the first physical storage area is insufficient to continue storing the KV generated during the Nth inference process when the virtual address is not mapped to a physical address; Alternatively, detecting whether there is any remaining storage space in the first physical storage area, and when there is no remaining storage space in the first physical storage area, determining that the first physical storage area is insufficient to continue storing the KV generated in the Nth inference process; Alternatively, it is detected whether the remaining number of KVs that can be stored in the first physical storage area is zero, and when the remaining number of KVs that can be stored is zero, it is determined that the first physical storage area is insufficient to continue storing the KVs generated in the Nth reasoning process.

14. The device according to any one of claims 9 to 12, characterized in that The processing module includes M sub-processing modules, the first virtual address space includes M first sub-address spaces, the physical storage area indicated by each first sub-address space in the M first sub-address spaces is used to store the KV generated by a sub-processing module, and the physical storage areas indicated by different first sub-address spaces are used to store the KVs generated by different sub-processing modules, and M is a positive integer greater than 1.

15. The device according to claim 14, characterized in that The allocation unit is further configured to: After allocating the second physical storage area to the processing module, the first virtual address space and the second virtual address space are equally divided into M second sub-address spaces, each of the M second sub-address spaces is used to store the KV generated by a sub-processing module, and different second sub-address spaces are used to store the KVs generated by different sub-processing modules.

16. The device according to any one of claims 9 to 15, characterized in that The processing module includes multiple network layers in the AI model; Alternatively, the processing module includes a network layer in the AI model; Alternatively, the processing module includes multiple attention heads in a network in the AI model; Alternatively, the processing module includes an attention head in the AI model.

17. A computing device, characterized in that The computing device includes a processor and a memory; The memory is used to store instructions, and the processor executes the instructions stored in the memory to enable the computing device to perform the method according to any one of claims 1 to 8.

18. A computer-readable storage medium, characterized in that The method comprises instructions which, when executed on a computing device, cause the computing device to perform the method according to any one of claims 1 to 8.

19. A computer program product comprising instructions, characterized in that When the method is executed on at least one computing device, the method enables the at least one computing device to execute the method according to any one of claims 1 to 8.

Images