Memory management method, device, medium, equipment and product for attention calculation

Abstract

The application discloses a memory management method and device for attention calculation, a medium, equipment and products. The method comprises the following steps: allocating value data blocks and key data blocks to a first address interval and a second address interval of a first storage body region in a shared memory respectively; allocating a query data block to a third address interval with a first offset from the first address interval; wherein the first offset is determined by the memory occupation of the value data block; the combined interval of the first address interval and the third address interval has the same first boundary address as the shared memory; allocating an output data block to a fourth address interval with a second offset from the second address interval; wherein the second offset is determined by the memory occupation of the key data block; the combined interval of the second address interval and the fourth address interval has the same second boundary address as the shared memory. The application can maximize the memory access bandwidth, thereby effectively improving the efficiency of attention calculation.

Description

Technical Field

[0001] This invention relates to the field of artificial intelligence technology, and in particular to a memory management method, apparatus, computer-readable storage medium, electronic device, and computer program product for attention computing. Background Technology

[0002] With the rapid development of artificial intelligence technology, the attention mechanism has become a core operator in many fields such as natural language processing and computer vision. It achieves dynamic weighted aggregation of input information by calculating the interaction relationship between queries, keys, and values. As model complexity and data scale continue to increase, the efficiency bottleneck of attention computation becomes increasingly prominent, among which memory access performance is one of the factors restricting computational efficiency. Summary of the Invention

[0003] The purpose of this invention is to provide a memory management method, apparatus, computer-readable storage medium, electronic device, and computer program product for attention computing, which maximizes memory access bandwidth by optimizing the layout of data in shared memory, thereby effectively improving the efficiency of attention computing.

[0004] A first aspect of the present invention provides a memory management method for attention computation, comprising:

[0005] Allocate the value data block to the first address range of the first storage area in the shared memory;

[0006] Allocate the key data block to the second address range of the second storage area in the shared memory;

[0007] The query data block is allocated to a third address interval with a first offset from the first address interval; wherein the first offset is determined by the memory usage of the value data block; the merged interval of the first address interval and the third address interval has the same first boundary address as the shared memory;

[0008] The output data block is allocated to a fourth address range with a second offset from the second address range; wherein the second offset is determined by the memory usage of the key data block; the merged range of the second address range and the fourth address range has the same second boundary address as the shared memory.

[0009] Optionally, the method further includes:

[0010] After allocating the first address range, the second address range, the third address range, and the fourth address range, the storage capacity of the remaining space in the shared memory is calculated.

[0011] When the storage capacity is not less than the memory usage of the exponent and the logarithm result, the exponent and the logarithm result are allocated to the fifth address range of the remaining space;

[0012] When the storage capacity is smaller than the storage size of the exponent and the logarithm result, the exponent and the logarithm result are stored in global memory.

[0013] Optionally, the first address range may contain the first boundary address.

[0014] Optionally, the second address range may contain the second boundary address.

[0015] Optionally, the memory usage of the first address range, the second address range, the third address range, and the fourth address range is determined by the number of blocks cached in shared memory in each batch, the size of a single block, and the data type.

[0016] Optionally, cross-region data storage is supported between the first storage region and the second storage region.

[0017] A second aspect of the present invention provides a memory management apparatus for attention computation, comprising:

[0018] The value data allocation module is used to allocate value data blocks to the first address range of the first storage area in the shared memory.

[0019] The key data allocation module is used to allocate key data blocks to the second address range of the second storage area in the shared memory.

[0020] A query data allocation module is used to allocate query data blocks to a third address range with a first offset from the first address range; wherein, the first offset is determined by the memory usage of the value data block; the merged range of the first address range and the third address range has the same first boundary address as the shared memory;

[0021] An output data allocation module is used to allocate output data blocks to a fourth address range with a second offset from the second address range; wherein the second offset is determined by the memory usage of the key data block; and the merged range of the second address range and the fourth address range has the same second boundary address as the shared memory.

[0022] A third aspect of the present invention provides a computer-readable storage medium comprising a stored computer program; wherein, when the computer program is executed, it controls the device in which the computer-readable storage medium is located to perform the memory management method for attention computation as described in any embodiment of the first aspect.

[0023] A fourth aspect of the present invention provides a computer program product including computer instructions that, when executed by a processor, implement the memory management method for attention computation described in any embodiment of the first aspect.

[0024] A fifth aspect of the present invention provides an electronic device including a processor, a memory, and a computer program stored in the memory and configured to be executed by the processor, wherein the processor, when executing the computer program, implements the memory management method for attention computation as described in any embodiment of the first aspect.

[0025] Compared with existing technologies, embodiments of the present invention provide a memory management method, apparatus, computer-readable storage medium, electronic device, and computer program product for attention computation, which has the following beneficial effects: Embodiments of the present invention allocate key data blocks and value data blocks to a first storage area and a second storage area of ​​shared memory, respectively, and store query data blocks and value data blocks adjacently, and output data blocks and key data blocks adjacently. Simultaneously, the merged areas of value data blocks and query data blocks, and the merged areas of key data blocks and output data blocks are deployed at opposite ends of the shared memory. This storage layout supports parallel loading and access of data sets within query data blocks and key data blocks, key data blocks and value data blocks, and output data blocks and query data blocks, thereby significantly reducing data waiting time and effectively improving memory access bandwidth and the overall efficiency of attention computation. Attached Figure Description

[0026] Figure 1 This is a flowchart illustrating an embodiment of the memory management method for attention computation provided by the present invention;

[0027] Figure 2 This is a schematic diagram of an embodiment of the storage region division provided by the present invention;

[0028] Figure 3 This is a schematic diagram of an embodiment of the shared memory allocation address range provided by the present invention;

[0029] Figure 4 This is a schematic diagram of another embodiment of the shared memory allocation address range provided by the present invention;

[0030] Figure 5 This is a schematic diagram of a structure of an embodiment of the memory management device for attention computing provided by the present invention;

[0031] Figure 6 This is a schematic diagram of the structure of an embodiment of the electronic device provided by the present invention. Detailed Implementation

[0032] The technical solutions of the embodiments of the present invention will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only some embodiments of the present invention, and not all embodiments. Based on the embodiments of the present invention, all other embodiments obtained by those skilled in the art without creative effort are within the scope of protection of the present invention.

[0033] The artificial intelligence processor involved in this invention can be any one of CPU (Central Processing Unit), GPU (Graphics Processing Unit), TPU (Tensor Processing Unit), NPU (Neural Network Processing Unit), DPU (Deep Learning Processing Unit), APU (Accelerated Processing Unit), and GPGPU (General-Purpose computing on Graphics Processing Unit), depending on its application to a specific product or technology in the embodiments of this invention.

[0034] See Figure 1 This is a flowchart illustrating an embodiment of the memory management method for attention computation provided by the present invention.

[0035] A first aspect of the present invention provides a memory management method for attention computation, including steps S11 to S14, as follows:

[0036] Step S11: Allocate the value data block to the first address range of the first storage area in the shared memory;

[0037] Step S12: Allocate the key data block to the second address range of the second storage bank region in the shared memory;

[0038] Step S13: Allocate the query data block to a third address interval with a first offset from the first address interval; wherein, the first offset is determined by the memory usage of the value data block; the merged interval of the first address interval and the third address interval has the same first boundary address as the shared memory;

[0039] Step S14: Allocate the output data block to a fourth address range with a second offset from the second address range; wherein the second offset is determined by the memory usage of the key data block; the merged range of the second address range and the fourth address range has the same second boundary address as the shared memory.

[0040] It should be noted that Group Shared Memory (GSM) includes a first storage area and a second storage area, which are used to cache the query data blocks, value data blocks, key data blocks, and output data blocks (i.e. blocks of attention calculation results) required for pipelined attention computation.

[0041] Specifically, for a given AI chip, the GSM memory is divided into two parallel regions based on address: a low-address memory region (lowbank, capacity 128KB) and a high-address memory region (highbank, capacity 96KB). This allows simultaneous access to data in both regions during attention calculations, improving memory access efficiency. For example... Figure 2 The diagram shown is a schematic representation of an embodiment of the storage region partitioning provided by the present invention. Figure 2 In this embodiment, GSM has two boundary addresses: a start address (first boundary address) and an end address (second boundary address). The first storage area is configured as the low bank (corresponding to "LowAddress"), and the second storage area is configured as the high bank (corresponding to "High Address"). Memory access operations 1 and 2 can access both the low bank and the high bank simultaneously. Alternatively, the first storage area can be configured as the high bank, and the second storage area as the corresponding low bank. In this case, the first boundary address becomes the end address of the GSM, and the second boundary address becomes the start address. This embodiment of the invention does not limit this specific configuration.

[0042] During the calculation of attention (such as flash attention), the inventors discovered that the data required differs at different calculation stages. For example, in the first stage of calculating the attention score, it is necessary to query the data block (denoted as Q) and the key data block (denoted as K) and execute "QK". TThe first stage involves matrix multiplication of P and V. The second stage, calculating the attention of the current block, requires an exponentially increased attention score (P) and a value data block (V) to perform matrix multiplication of P and V. The third stage is the data unloading and prefetching stage, where the output data block (O) is unloaded from the thread register to the GSM, and the query data required for the next round of computation is prefetched from the Global Memory (GLM) to the GSM. Based on this, this embodiment of the invention stores Q and K separately in the first and second storage areas; similarly, K and V are placed in the two storage areas; and O and Q are also placed in the two storage areas. This enables parallel loading / access of each set of data, thereby improving memory access bandwidth.

[0043] As can be deduced from the above, K and V are cached in the first and second storage areas of the GSM, respectively, Q and V are stored adjacently, and O and K are stored adjacently. Specifically, the merged areas of Q and V, and the merged areas of O and K, are deployed at opposite ends of the GSM to minimize fragmentation caused by data block allocation across low and high banks, while fully leveraging the parallel access capabilities of memory access operations 1 and 2 to the two storage areas, thereby effectively improving memory access bandwidth. To provide a clearer description of the technical solution provided by the embodiments of the present invention, some specific embodiments are provided below for reference:

[0044] Example 1: Figure 3 The diagram shown is a schematic representation of an embodiment of the shared memory allocation address range provided by the present invention. Figure 3 In the configuration, the first memory bank region is configured as a low bank, and the second memory bank region is configured as a high bank. The starting address of the merged region of Q and V coincides with the first boundary address (starting address) of GSM, and the ending address of the merged region of O and K coincides with the second boundary address (ending address) of GSM. In the low bank, the merged region of Q and V has two layout options: ① The first offset between the third address interval corresponding to Q and the first address interval corresponding to V is the memory footprint of V (i.e., +V_Size), in which case the starting address of Q and the ending address of V are continuous; ② The first offset is -V_Size, in which case the ending address of Q and the starting address of V are continuous. Similarly, in the high bank, the merged region of O and K also has two layout options: ① The second offset between the fourth address interval corresponding to O and the second address interval corresponding to K is the memory footprint of K (i.e., +K_Size), in which case the starting address of O and the ending address of K are continuous; ② The second offset is -K_Size, in which case the ending address of O and the starting address of K are continuous.

[0045] Example 2: Figure 4The diagram shown is a schematic representation of another embodiment of the shared memory allocation address range provided by the present invention. Figure 4 In the configuration, the first storage region is configured as a high bank, and the second storage region is configured as a low bank. The end address of the merged region of Q and V coincides with the first boundary address (end address) of GSM, and the start address of the merged region of O and K coincides with the second boundary address (start address) of GSM. In the high bank, the merged region of Q and V has two layout options: ① the first offset is +V_Size, in which case the start address of Q and the end address of V are continuous; ② the first offset is -V_Size, in which case the end address of Q and the start address of V are continuous. Similarly, in the low bank, the merged region of O and K also has two layout options: ① the second offset is +K_Size, in which case the start address of O and the end address of K are continuous; ② the second offset is -K_Size, in which case the end address of O and the start address of K are continuous.

[0046] On the one hand, this invention addresses the multi-stage nature of attention computation by allocating query data blocks (Q) and key data blocks (K), key data blocks (K) and value data blocks (V), and output data blocks (O) and query data blocks (Q) to the first and second storage areas of the GSM memory, respectively. This enables parallel loading and access of each set of data, significantly reducing data waiting time and effectively improving memory access bandwidth and the overall efficiency of attention computation. On the other hand, by deploying the merged areas of value data blocks (V) and query data blocks (Q), and the merged areas of key data blocks (K) and output data blocks (O) at opposite ends of the shared memory, fragmentation caused by data block allocation across low and high banks can be minimized, thereby improving the memory access efficiency of data blocks.

[0047] In an optional embodiment, the method further includes: calculating the storage capacity of the remaining space in the shared memory after allocating the first address range, the second address range, the third address range, and the fourth address range;

[0048] When the storage capacity is not less than the memory usage of the exponent and the logarithm result, the exponent and the logarithm result are allocated to the fifth address range of the remaining space;

[0049] When the storage capacity is smaller than the storage size of the exponent and the logarithm result, the exponent and the logarithm result are stored in global memory.

[0050] It should be noted that after allocating the first address range (value data block V), the second address range (key data block K), the third address range (query data block Q), and the fourth address range (output data block O), this embodiment of the invention calculates the storage capacity of the unoccupied remaining space in the GSM. When the storage capacity of the remaining space is greater than or equal to the result of the exponent and the logarithm (LogSumExp, LSE), the LSE is allocated to the fifth address range within the remaining space / GSM; when the storage capacity of the remaining space is less than the LSE, the LSE is directly stored in the GLM to avoid crowding out the storage space of Q, K, V, and O; wherein, the LSE is calculated by "first taking the logarithm, then summing, and finally exponentiating" based on the attention score.

[0051] In an optional embodiment, the memory usage of the first address range, the second address range, the third address range, and the fourth address range is determined by the number of blocks cached in shared memory in each batch, the size of a single block, and the data type.

[0052] It should be noted that the memory usage of the first, second, third, and fourth address ranges is determined by the following parameters of their respective stored data blocks (value data block V, key data block K, query data block Q, and output data block O): the number of blocks cached in GSM per batch, the size of a single block, and the data type used by the data blocks. The specific calculations are as follows:

[0053] The memory usage of the first address range (V) is: V_SIZE = kStages × n × k × sizeof(Element);

[0054] The memory usage of the second address range (K) is K_SIZE = kStages × n × k × sizeof(Element);

[0055] The memory usage of the third address range (Q) is Q_SIZE = kConsumers × m × k × sizeof(Element);

[0056] The memory usage of the fourth address range (O) is O_SIZE = kConsumers × m × k × sizeof(Element);

[0057] The memory footprint of the fifth address range (LSE) is LSE_SIZE = kConsumers × m × sizeof(ElementLSE).

[0058] In the above formula, `kConsumers` represents the number of consumer thread groups enabled, used to execute the "outer loop calculation of value blocks" and the "inner loop calculation of key blocks - value blocks" in parallel, obtaining `kConsumers` output result blocks in parallel. Therefore, the number of value blocks is equal to `kConsumers`. `kStages` represents the number of blocks cached in shared memory per batch during the inner loop calculation; the number of value blocks and key blocks is the same. The size of the query block is `m×k`, where `m` is the sequence length of the query block and `k` is the attention head dimension (`headdim`). The size of both the value block and the key block is `n×k`, where `n` is the sequence length of the value block / key block. `Element` represents the data types Q, K, V, and O (e.g., FP16, FP8); `ElementLSE` is a data type specific to LSE, typically FP32. `sizeof(Element)` and `sizeof(ElementLSE)` are used to calculate the number of bytes of memory occupied by a single element of the corresponding data type.

[0059] In an optional embodiment, cross-region data storage is supported between the first storage region and the second storage region.

[0060] It should be noted that although GSM is divided into two storage banks supporting parallel access—low bank (128KB capacity) and high bank (96KB capacity)—data blocks can still be stored across these two storage banks according to actual storage needs. In attention-based computation scenarios, value blocks and key blocks are typically similar in size, and the number of blocks cached in shared memory per batch is the same; similarly, query blocks and output result blocks are also similar in size, and the number of blocks cached in shared memory per batch is the same. Therefore, the overall size of the merged regions for Q and V, and the merged regions for O and K, are similar. In actual computation, the overall size of these two merged regions may exceed 96KB, requiring one of the merged regions to be stored across two storage banks (i.e., across banks), thus leading to data storage fragmentation.

[0061] In an optional embodiment, the first address range contains the first boundary address.

[0062] Furthermore, the second address range contains the second boundary address.

[0063] It should be noted that, regarding the potential fragmentation problem in data storage, the inventors further investigated and found that key data blocks and value data blocks are used for inner loop calculations, both of which are frequently accessed data; while query data blocks are used for outer loop calculations, and output data blocks are the results obtained after the outer loop ends, with relatively low access frequencies. Therefore, this embodiment of the invention deploys key data blocks and value data blocks at opposite ends of the GSM to avoid fragmentation, i.e., the first address interval contains the first boundary address of the GSM, and the second address interval contains the second boundary address; simultaneously, query data blocks and value data blocks are stored adjacently, and output data blocks and key data blocks are stored adjacently. Under this storage layout, even if query data blocks or output data blocks need to be stored across banks due to actual storage requirements, their impact on overall attention computing performance can be minimized due to the low access frequency of such data. In summary, this embodiment of the invention allocates memory regions more likely to be stored across banks to data blocks with lower access frequencies, thereby minimizing the negative impact of cross-bank storage on memory access performance.

[0064] See Figure 5 This is a schematic diagram of an embodiment of the memory management device for attention computing provided by the present invention.

[0065] A second aspect of the present invention provides a memory management apparatus for attention computation, comprising:

[0066] Value data allocation module 11 is used to allocate value data blocks to the first address range of the first storage area in the shared memory;

[0067] The key data allocation module 12 is used to allocate key data blocks to the second address range of the second storage area in the shared memory;

[0068] The query data allocation module 13 is used to allocate the query data block to a third address interval with a first offset from the first address interval; wherein, the first offset is determined by the memory usage of the value data block; the merged interval of the first address interval and the third address interval has the same first boundary address as the shared memory;

[0069] The output data allocation module 14 is used to allocate the output data block to a fourth address interval with a second offset from the second address interval; wherein the second offset is determined by the memory usage of the key data block; the merged interval of the second address interval and the fourth address interval has the same second boundary address as the shared memory.

[0070] It should be noted that the memory management device for attention computing provided in the second aspect embodiment of the present invention can implement all the processes of the memory management method for attention computing described in any of the first aspect embodiments. The functions and technical effects of each module and unit in the device are the same as the functions and technical effects of the memory management method for attention computing described in any of the first aspect embodiments, and will not be repeated here.

[0071] A third aspect of the present invention provides a computer-readable storage medium comprising a stored computer program; wherein, when the computer program is executed, it controls the device in which the computer-readable storage medium is located to perform the memory management method for attention computation described in any of the first aspects of the present invention.

[0072] A fourth aspect of the present invention provides a computer program product including computer instructions that, when executed by a processor, implement the memory management method for attention computation described in any of the first aspects of the present invention.

[0073] See Figure 6 This is a schematic diagram of an embodiment of the electronic device provided by the present invention.

[0074] A fifth aspect of the present invention provides an electronic device including a processor 21, a memory 22, and a computer program stored in the memory 22 and configured to be executed by the processor 21, wherein the processor 21, when executing the computer program, implements the memory management method for attention computation described in any of the first aspects of the present invention.

[0075] Preferably, the computer program can be divided into one or more modules / units (such as computer program one, computer program two, ...), and the one or more modules / units are stored in the memory 22 and executed by the processor 21 to complete the present invention. The one or more modules / units can be a series of computer program instruction segments capable of performing specific functions, and the instruction segments are used to describe the execution process of the computer program in the electronic device.

[0076] The processor 21 can be any one of a CPU (Central Processing Unit), GPU (Graphics Processing Unit), TPU (Tensor Processing Unit), NPU (Neural Network Processing Unit), DPU (Deep Learning Processing Unit), APU (Accelerated Processing Unit), and GPGPU (General-Purpose Computing on Graphics Processing Unit). The processor 21 is the control center of the electronic device, connecting various parts of the electronic device via various interfaces and lines.

[0077] The memory 22 mainly includes a program storage area and a data storage area. The program storage area can store the operating system, applications required for at least one function, etc., and the data storage area can store related data, etc. In addition, the memory 22 can be a high-speed random access memory, or a non-volatile memory, such as a plug-in hard disk, a smart media card (SMC), a secure digital (SD) card, and a flash card, etc., or the memory 22 can also be other volatile solid-state storage devices.

[0078] It should be noted that the aforementioned electronic devices may include, but are not limited to, processors and memory, as will be understood by those skilled in the art. Figure 6 The structural block diagram shown is merely a structural example of the above-described electronic device and does not constitute a limitation on the structure of the above-described electronic device. The above-described electronic device may include more or fewer components than shown, or combine certain components, or different components.

[0079] The above description is only a preferred embodiment of the present invention. It should be noted that for those skilled in the art, several improvements and modifications can be made without departing from the technical principles of the present invention, and these improvements and modifications should also be considered within the scope of protection of the present invention.

Claims

1. A memory management method for attention computation, characterized in that, include: Allocate the value data block to the first address range of the first storage area in the shared memory; Allocate the key data block to the second address range of the second storage area in the shared memory; The query data block is allocated to a third address interval with a first offset from the first address interval; wherein the first offset is determined by the memory usage of the value data block; the merged interval of the first address interval and the third address interval has the same first boundary address as the shared memory; The output data block is allocated to a fourth address range with a second offset from the second address range; wherein the second offset is determined by the memory usage of the key data block; the merged range of the second address range and the fourth address range has the same second boundary address as the shared memory.

2. The memory management method for attention computation as described in claim 1, characterized in that, The method further includes: After allocating the first address range, the second address range, the third address range, and the fourth address range, the storage capacity of the remaining space in the shared memory is calculated. When the storage capacity is not less than the memory usage of the exponent and the logarithm result, the exponent and the logarithm result are allocated to the fifth address range of the remaining space; When the storage capacity is smaller than the storage size of the exponent and the logarithm result, the exponent and the logarithm result are stored in global memory.

3. The memory management method for attention computation as described in claim 1, characterized in that, The first boundary address exists within the first address range.

4. The memory management method for attention computation as described in claim 1, characterized in that, The second address range contains the second boundary address.

5. The memory management method for attention computation as described in claim 1, characterized in that, The memory usage of the first address range, the second address range, the third address range, and the fourth address range is determined by the number of blocks cached in shared memory in each batch, the size of a single block, and the data type.

6. The memory management method for attention computation as described in claim 1, characterized in that, Data storage across regions is supported between the first storage region and the second storage region.

7. A memory management device for attention computation, characterized in that, include: The value data allocation module is used to allocate value data blocks to the first address range of the first storage area in the shared memory. The key data allocation module is used to allocate key data blocks to the second address range of the second storage area in the shared memory. A query data allocation module is used to allocate query data blocks to a third address range with a first offset from the first address range; wherein, the first offset is determined by the memory usage of the value data block; the merged range of the first address range and the third address range has the same first boundary address as the shared memory; An output data allocation module is used to allocate output data blocks to a fourth address range with a second offset from the second address range; wherein the second offset is determined by the memory usage of the key data block; and the merged range of the second address range and the fourth address range has the same second boundary address as the shared memory.

8. A computer-readable storage medium, characterized in that, The computer-readable storage medium includes a stored computer program; wherein, when the computer program is executed, it controls the device on which the computer-readable storage medium resides to perform the memory management method for attention computation as described in any one of claims 1 to 6.

9. A computer program product, characterized in that, It includes computer instructions that, when executed by a processor, implement the memory management method for attention computation as described in any one of claims 1 to 6.

10. An electronic device, characterized in that, It includes a processor, a memory, and a computer program stored in the memory and configured to be executed by the processor, wherein the processor, when executing the computer program, implements the memory management method for attention computation as described in any one of claims 1 to 6.

Images