Methods, computing devices, media, and program products for performing attention computations

Abstract

The present application relates to a method, a computing device, a medium and a program product for performing attention calculation. The method comprises: obtaining key-value indexes corresponding to a plurality of selected key-value entries participating in attention operation via TopK calculation; calculating addresses of a plurality of data blocks to be read based on the key-value indexes, each data block in the plurality of data blocks comprising a corresponding selected key-value entry; reading a plurality of redundant matrix data based on the addresses of the plurality of data blocks using a matrix loading instruction for attention operation to generate a preliminary calculation result; and obtaining a target calculation result based on the preliminary calculation result and a target mask within the data block. The present application makes full use of the high-bandwidth matrix loading capability of the chip, and can significantly improve the overall efficiency of attention calculation and the utilization rate of hardware resources with the introduction of a small amount of calculation redundancy.

Description

Technical Field

[0001] The embodiments of the present invention generally relate to the field of artificial intelligence, and more specifically to a method, computing device, computer-readable storage medium, and computer program product for performing attention computation. Background Technology

[0002] In the inference and training process of Large Language Models (LLMs), the attention mechanism is one of the core computational steps. For LLM inference with long contexts, standard attention computation can account for more than 50% of the time. To reduce computational cost, artificial intelligence models have proposed sparse attention algorithms, which significantly reduce computational complexity by selecting only the top K most important key-value pairs for computation. However, sparse attention algorithms have significant hardware efficiency issues when performing discrete reads of the KV cache. Specifically, the original KV cache is usually logically stored contiguously, but sparse attention needs to randomly read key-value entries distributed in different locations in the cache based on the top K indices.

[0003] In related technologies, one attention calculation scheme is to read all KV cache data and use a mask to block invalid positions to obtain the attention calculation result. However, this method leads to huge data reading and matrix multiplication redundancy. Another scheme is to use single instruction multithreading (SIMT) instructions to load discrete key-value entries one by one. However, this method is inefficient on architectures that focus on dense data reading and requires additional data concatenation operations, wasting bandwidth and computing resources.

[0004] In summary, the shortcomings of traditional methods for performing attention calculations are as follows: when it is necessary to read key-value cache entries discretely according to the sparse attention algorithm, the full read method leads to excessive data redundancy and computational redundancy, wasting bandwidth and computing power; while the discrete loading method is limited by the memory access granularity of the hardware, resulting in low loading efficiency and requiring additional data concatenation operations, which cannot fully utilize the matrix loading capability of the hardware. Summary of the Invention

[0005] This invention provides a method, computing device, computer-readable storage medium, and computer program product for performing attention computation. This invention fully utilizes the dense memory access and matrix computation capabilities of hardware, significantly reducing the overhead of discrete data reading while introducing only a small amount of computational redundancy, improving bandwidth utilization and instruction efficiency, thereby achieving higher overall computing performance and resource utilization.

[0006] According to one aspect of the present invention, a method for performing attention computation is provided. The method includes: obtaining key-value indices corresponding to a plurality of selected key-value entries participating in the attention computation via TopK computation; calculating addresses of a plurality of data blocks to be read based on the key-value indices, each of the plurality of data blocks including a corresponding selected key-value entry; reading a plurality of redundant matrix data using matrix loading instructions based on the addresses of the plurality of data blocks for use in the attention computation, thereby generating a preliminary computation result; and obtaining a target computation result based on the preliminary computation result and a target mask within the data blocks.

[0007] In some embodiments, the method further includes: generating a target mask within a data block, at least based on a key-value index.

[0008] In some embodiments, generating a target mask within a data block, at least based on a key-value index, includes: generating an initial mask based on the key-value index, the initial mask corresponding to the positions of multiple selected key-value entries in a key-value cache; and determining a target mask within the data block based on the initial mask and multiple redundant matrix data read.

[0009] In some embodiments, the target mask contains fewer mask bits than the initial mask.

[0010] In some embodiments, determining the target mask within a data block based on an initial mask and multiple redundant matrix data read includes: removing the mask value corresponding to the position of an unread key-value entry in the initial mask, wherein the unread key-value entry is the key-value entry in the key-value cache other than the multiple redundant matrix data.

[0011] In some embodiments, calculating the address of multiple data blocks to be read based on the key-value index includes: calculating the address of multiple data blocks to be read based on the key-value index and the data reading granularity of the matrix loading instruction, such that each data block in the multiple data blocks includes a corresponding selected key-value entry, and the data reading granularity is an integer multiple of the hardware alignment granularity of the artificial intelligence chip, which is used to perform attention calculation.

[0012] In some embodiments, the method further includes: determining the data granularity of the matrix loading instruction based on the size of each selected key-value entry, such that the data granularity of the matrix loading instruction is larger than the size of each selected key-value entry.

[0013] In some embodiments, the size of the redundant matrix data is larger than the size of the corresponding selected key-value entry.

[0014] According to a second aspect of the present invention, a computing device is also provided. The computing device includes: at least one processor; and a memory communicatively connected to the at least one processor; wherein the memory stores instructions executable by the at least one processor to enable the computing device to perform the method of the first aspect of the present invention.

[0015] According to a third aspect of the present invention, a computer-readable storage medium is also provided. The computer-readable storage medium stores a computer program that, when executed by a machine, performs the method of the first aspect of the present invention.

[0016] According to a fourth aspect of the present invention, a computer program product is also provided, comprising a computer program that, when executed by a machine, performs the method of the first aspect of the present invention.

[0017] This invention calculates data block addresses based on TopK indexes, reads matrix data containing necessary redundancy using matrix loading instructions at the block level, and performs post-processing with a simplified target mask. This avoids the significant waste of fully reading the KV cache and circumvents inefficient discrete SIMT loading. Furthermore, it fully leverages the high-bandwidth matrix loading capabilities of AI chips (especially those with block-level access granularity), significantly improving the overall efficiency and hardware resource utilization of sparse attention computation with minimal computational redundancy.

[0018] It should be understood that the description in this section is not intended to identify key or essential features of the embodiments of the present invention, nor is it intended to limit the scope of the invention. Other features of the invention will become readily apparent from the following description. Attached Figure Description

[0019] The above and other features, advantages, and aspects of the various embodiments of the present invention will become more apparent from the accompanying drawings and the following detailed description. In the drawings, the same or similar reference numerals denote the same or similar elements.

[0020] Figure 1 The illustration shows a schematic diagram of a general-purpose graphics processor implementing a method for performing attention computation according to an embodiment of the present invention.

[0021] Figure 2 The flowchart of a method for performing attention calculation according to some embodiments of the present invention is shown.

[0022] Figure 3 A schematic diagram of key-value caching and data block partitioning provided according to some embodiments of the present invention is shown.

[0023] Figure 4A schematic diagram of a simplified mask provided according to some embodiments of the present invention is shown.

[0024] Figure 5 A flowchart illustrating the determination of target attention results according to some embodiments of the present invention is shown.

[0025] Figure 6 A schematic diagram of a computing device for performing attention calculation according to an embodiment of the present invention is shown.

[0026] Figure 7 A schematic diagram of a conventional method for reading multiple key-value entries is shown. In the various figures, the same or corresponding reference numerals indicate the same or corresponding parts. Detailed Implementation

[0027] Preferred embodiments of the invention will now be described in more detail with reference to the accompanying drawings. While preferred embodiments of the invention are shown in the drawings, it should be understood that the invention can be implemented in various forms and should not be limited to the embodiments set forth herein. Rather, these embodiments are provided so that the invention will be thorough and complete, and will fully convey the scope of the invention to those skilled in the art.

[0028] The term "comprising" and its variations as used herein signify open inclusion, i.e., "including but not limited to". Unless otherwise stated, the term "or" means "and / or". The term "based on" means "at least partially based on". The terms "one example embodiment" and "one embodiment" mean "at least one example embodiment". The term "another embodiment" means "at least one additional embodiment". The terms "first", "second", etc., may refer to different or the same objects.

[0029] The attention operator is a crucial operator in artificial intelligence (AI) models, serving as the core computational module for implementing the attention mechanism. By calculating the correlation (attention weights) between elements at different positions within the input tensor, the attention operator dynamically adjusts the representation of each element, enabling the model to focus on information more relevant to the current task. For input data, the attention operator generates query, key, and value matrices. The computation of attention utilizes simple matrix operations on the query, key, and value matrices.

[0030] In this invention, a key-value cache (KV cache) refers to a cache structure that stores historical key and value vectors within an attention mechanism. A key-value entry refers to a combination of a key and a value at a single location in the cache. A data block refers to the granularity of data loaded by a matrix load instruction, which typically contains multiple consecutive key-value entries and is an integer multiple of the hardware alignment granularity. A matrix load instruction is an instruction capable of loading matrix data in batches at the data block granularity. Redundant matrix data refers to unselected data read in addition to the actually selected key-value entries due to block-based reading.

[0031] As described earlier, various sparse attention algorithms can be employed to reduce the computational overhead of attention. Taking DeepSeek's Sparse Attention (DSA) algorithm as an example, its core idea is to select only the K most important (K being a natural number) key-value entries (TopK) from all historical key-value entries for each query, ignoring the majority of other key-value entries. This method can reduce computational complexity while maintaining model accuracy, thereby significantly improving inference speed.

[0032] However, sparse attention algorithms face a challenge in practical hardware implementation: the problem of discrete data retrieval. That is, to support efficient dense access, the original key-value cache is typically stored logically contiguously and arranged sequentially. However, the indexes output by the TopK algorithm are discrete and non-contiguous, meaning that the required data needs to be retrieved separately from various random locations in the KV cache.

[0033] One related technique ignores the discreteness of the TopK index and directly reads all data from the entire KV cache, performing a complete dense matrix multiplication. Then, a global mask is generated based on the TopK index, and after Softmax, the results at invalid locations are set to negative infinity or zero, thus retaining only the valid computation results. However, this technique requires an extremely large amount of data to be read (proportional to the size of the entire KV cache), and the matrix multiplication computation is also extremely large (containing a large number of redundant operations), severely wasting the chip's bandwidth and computing power resources. Although the mask can guarantee the correctness of the final result, a large number of invalid calculations occupy valuable computing units, resulting in actual performance far below the hardware peak.

[0034] In another related technology, such as Figure 7As shown, the single-instruction, multi-threaded architecture of a general-purpose computing on graphics processing units (GPGPU) can be utilized to execute an independent load instruction (e.g., an ld instruction) for each selected KV index, retrieving the corresponding key-value entry 704 from the KV Cache 702. For example... Figure 7 As shown, because these loading operations are scattered and fine-grained, they trigger a large number of memory access requests, leading to increased cache miss rates. Furthermore, the loaded data blocks are logically discontinuous, requiring additional concatenation instructions (e.g., packaging the scattered data into a continuous tensor via a vector processing unit) before being fed into the matrix multiplication unit for computation. In other words, SIMT discrete loading is inefficient and cannot fully utilize the chip's memory bandwidth. The additional data concatenation introduces numerous auxiliary instructions, further reducing overall computational efficiency.

[0035] To at least partially address one or more of the aforementioned problems and other potential issues, exemplary embodiments of the present invention provide at least one method for performing attention computation. This method includes: obtaining key-value indices corresponding to multiple selected key-value entries participating in the attention computation via TopK computation; calculating addresses of multiple data blocks to be read based on the key-value indices, each data block including a corresponding selected key-value entry; reading multiple redundant matrix data using matrix loading instructions based on the addresses of the multiple data blocks for attention computation, thereby generating a preliminary computation result; and obtaining a target computation result based on the preliminary computation result and a target mask within the data blocks.

[0036] Specifically, compared to the aforementioned techniques based on full read and full computation, the example embodiment of this invention reads only the data block containing the selected key-value entry, rather than the entire KV cache, thus significantly reducing the amount of data read and matrix multiplication computation. Compared to the aforementioned techniques based on SIMT discrete loading and matrix concatenation, the example embodiment of this invention utilizes matrix loading instructions to read data in blocks in batches, avoiding the large number of fine-grained memory access requests and data concatenation instructions caused by discrete loading, thereby significantly improving bandwidth utilization and instruction pipeline efficiency. In addition, the example embodiment of this invention also generates a simplified target mask within the data block, performing masking operations only on the small number of data blocks actually read, further reducing the computational overhead of post-processing. Furthermore, since the attention calculations between different data blocks are naturally independent, the example embodiment of this invention can be easily extended to multi-core parallel processing, further improving throughput.

[0037] Figure 1A schematic diagram illustrates a general-purpose graphics processor 100 implementing a method for performing attention calculation according to an embodiment of the present invention. Figure 1 As shown, the general-purpose graphics processor 100 includes multiple streaming processor clusters (SPCs), such as SPC 1, ..., SPC M (where M is a natural number, M>1). Each streaming processor cluster can independently process a computational task, or multiple clusters can collaboratively process a computational task. Data is shared between the SPCs through a global cache / global memory.

[0038] Taking streaming processor cluster 1 as an example, it includes: multiple computation units (such as vector operation units) and tensor operation units. The multiple vector operation units are, for example, computation unit 1, computation unit 2, ..., computation unit N (N is a positive integer). The tensor operation units are specifically used to perform tensor computations, such as matrix multiplication (MMA), convolution, etc.

[0039] In some embodiments of the present invention, each vector operation unit includes multiple cores, a register file, and a shared cache for performing arithmetic and logical operations, such as accumulation, reduction, and regular addition, subtraction, multiplication, and division. The vector operation unit integrates the following functional modules: Key-value index acquisition module 102: Utilizes the core in the vector operation unit to perform TopK calculation and obtain the key-value index corresponding to the multiple selected key values ​​participating in the attention operation.

[0040] Address calculation module 104: Calculates the physical address of the data block to which each key-value index belongs. This module can be implemented by address generation logic within the vector operation unit or by loading the memory unit.

[0041] Target calculation result acquisition module 108: In the preliminary calculation result returned by the tensor operation unit, the position corresponding to the mask 0 is invalid (-INF or 0), and the Softmax and AND-value matrix multiplication are completed.

[0042] In some embodiments, the tensor operation unit includes: a cache block (e.g., two source operand caches) for temporarily storing input data read from memory or a higher-level cache; a computation core for performing the actual matrix multiplication operations; a matrix loading unit (which can be extended or reused from an existing loading unit to support instruction decoding logic), supporting matrix loading instructions, capable of reading consecutive data blocks from the KV cache at a data block granularity with high bandwidth, and directly filling the cache block; and an intermediate-level cache for sharing data between the vector operation unit and the tensor operation unit within the streaming processor cluster. Optionally, the intermediate-level cache can be omitted, and interaction can be performed directly through a global cache / memory. The tensor operation unit integrates a matrix data reading module 106, which supports matrix loading instructions and is capable of reading consecutive data blocks from the KV cache at a data block granularity with high bandwidth, and directly filling the cache block.

[0043] It is understandable that the above Figure 1 The general-purpose graphics processor 100 shown and its ability to execute the attention calculation method of the present invention can be widely used in artificial intelligence scenarios, especially suitable for processing long sequence data. For example, it can be applied to real-time dialogue systems with large language models (e.g., intelligent customer service, virtual personal assistants, etc.), multimodal large model video processing (e.g., analyzing one hour of surveillance video or movie clips, etc.), autonomous driving (e.g., vehicle trajectory prediction, pedestrian intent recognition, etc.), recommendation systems, code generation and completion tools (e.g., AI programming assistants, etc.), text processing (e.g., document summarization and question answering systems, machine translation, sentiment analysis, etc.), audio processing (e.g., automatic speech recognition, music generation, speaker separation, etc.), bioinformatics and scientific computing, etc.

[0044] In the above scenario, Figure 1 The general-purpose graphics processor shown executes the attention calculation method described in this invention through the collaboration of vector operation units and tensor operation units, achieving sparse attention calculation with high throughput, low latency, and high bandwidth utilization.

[0045] The following is combined Figure 2 A flowchart is provided to describe a method 200 for performing attention calculation according to an embodiment of the present invention. It should be noted that the method 200 according to an embodiment of the present invention can, for example, be used in… Figure 1 The method is implemented at the general-purpose graphics processor 100 shown. It should be understood that method 200 may also include additional actions not shown, and the scope of the invention is not limited in this respect.

[0046] At step 202, the key-value indices corresponding to the multiple selected key-value entries participating in the attention operation are obtained via TopK calculation.

[0047] During the inference process of an artificial intelligence model, for the currently processed query, the model calculates the relevance score between the query and the keys at all historical locations. The relevance score can be obtained through dot product, cosine similarity, or other metrics. Specifically, TopK calculation refers to selecting a number of key-value pairs from historical key-value pairs for the current query, such as a user-input request, to participate in subsequent attention calculations. For example, using a sparse attention algorithm, importance scores are calculated for all historical key-value entries, and the K key-value entries with the highest scores are selected to obtain K key-value indexes.

[0048] For example, in long-context reasoning, the original KV cache contains 200K key-value entries. After TopK selection, 2048 key-value indices with high relevance are retained. The key-value index refers to the position information of the key-value entry in the original KV cache (e.g., the 0th, 1st, etc.). It should be noted that the key-value index is discrete and not necessarily continuous; for example, it might be [5, 23, 87, 1024, …].

[0049] At step 204, the addresses of multiple data blocks to be read are calculated based on the key-value index, and each data block includes the corresponding selected key-value entry.

[0050] Hardware such as chips typically uses data blocks as the smallest unit of access, and each data block can contain multiple consecutive key-value entries. That is, when reading memory data, it is not read as individual key-value entries, but rather in fixed-size blocks. Each block can contain multiple consecutive key-value entries (e.g., a data block contains n KVentries, where n is an integer). Therefore, to avoid loading discrete KV entries one by one, data can be read in units of data blocks.

[0051] In some embodiments, the starting address of the data block to which each key-value index belongs (i.e., the one to which it falls) is calculated based on each selected key-value index. For example, the starting address of the data block is calculated based on the key-value index and hardware alignment granularity to ensure that the address of each data block conforms to the cache line alignment rules of the AI ​​chip, thus guaranteeing the correct execution of the matrix loading instructions.

[0052] It should be noted that multiple indexes may fall within the same data block, therefore, key-value indexes need to be deduplicated to obtain a list of data block addresses to be read. The length of this list (i.e., the number of data blocks to be read) is less than or equal to K. Each selected data block contains at least one valid KV entry selected by TopK (and may also contain other unselected redundant KV entries).

[0053] In step 206, based on the addresses of the multiple data blocks, multiple redundant matrix data are read using matrix loading instructions for attention operations, thereby generating preliminary calculation results.

[0054] The addresses of multiple data blocks are the set of starting addresses of the data blocks to be read, calculated based on key-value indices. The matrix loading instruction is a hardware-supported instruction that loads matrix data in batches at the data block level. The matrix loading instruction reads the entire content of a data block from the KV cache according to this address set, block by block. Because a data block may contain key-value entries not selected by TopK, the read data is redundant relative to the actually selected entries. The read data blocks are concatenated into a continuous matrix, and the tensor computation unit performs attention operations (such as matrix multiplication, scaling, etc.) to obtain a preliminary calculation result. That is, this preliminary calculation result includes the operation results corresponding to both valid and redundant key-value entries.

[0055] In step 208, the target calculation result is obtained based on the preliminary calculation result and the target mask in the data block.

[0056] The target mask is used to identify whether key-value entries in each data block are valid, that is, which positions in the data block are truly selected valid key-value entries. For example, the mask can be a binary mask, with mask 1 used to identify valid key-value entries and mask 0 used to identify redundant key-value entries. This mask can be generated based on the TopK index and the actual data blocks read, and it only covers the positions within the actually read data blocks, thus being more concise than a global mask that covers the entire KV cache. In other words, based on the actual data blocks read, only the mask bits corresponding to the read data blocks are retained, and the mask bits corresponding to the unread data blocks are removed, thus obtaining a concise target mask. Using this target mask, the values ​​of the corresponding redundant entries in the initial calculation result are set to negative infinity or zero, and then subsequent operations such as Softmax are performed to determine the corresponding target calculation result.

[0057] For example, multiple data blocks are concatenated into input matrices Q and K, and dense matrix multiplication is performed to obtain a preliminary result matrix S (containing the product of effective and redundant positions). A target mask M is applied, replacing positions in S corresponding to mask 0s with negative infinity or zero, resulting in S'. After Softmax and multiplication with the Value matrix, the desired attention result is output.

[0058] Figure 3 A schematic diagram illustrating key-value caching and data block partitioning according to some embodiments of the present invention is shown. For example... Figure 3As shown, the KV cache 302 is logically contiguous and consists of multiple key-value entries. In physical storage logic, these key-value entries are divided into multiple data blocks, such as data block 304, data block 306, data block 308, and data block 310. Each data block contains a fixed number of contiguous key-value entries. The starting address of each data block is hardware aligned. Assume that the selected key-value indices for TopK calculation are indices A, B, and C (which may be scattered in different locations, corresponding to key-value entries A, B, and C, respectively). Each index maps to a data block (for example, indices A, B, and C map to data blocks 304, 306, and 308, respectively; that is, data block 304 contains key-value entry A, data block 306 contains key-value entry B, and data block 308 contains key-value entry C). In addition to containing the key-value entries corresponding to indices A, B, and C (black parts), each data block also contains multiple adjacent redundant unselected key-value entries (white parts). It should be noted that these redundant key-value entries, although unnecessary, are still read in and participate in subsequent matrix multiplications.

[0059] In some embodiments, the number of mask bits in the target mask (i.e., the data length of the mask) is less than the number of mask bits in the initial mask. Specifically, the initial mask covers all locations throughout the key-value cache, and its length is equal to the total number of key-value entries in the key-value cache. The target mask, however, only covers locations within the actual data blocks being read, and its length is equal to the total number of key-value entries contained in all read data blocks. Since the number of data blocks actually read is much smaller than the total number of blocks in the key-value cache, and the size of each data block is fixed, the length of the target mask is significantly shorter than the length of the initial mask, making it more concise. This conciseness reduces the amount of data that subsequent masking operations need to process, thereby reducing computational and storage overhead.

[0060] In some embodiments, based on the actual range of data blocks read, only valid bits within the read blocks can be retained from the initial mask, while all mask bits corresponding to unread data blocks can be deleted, resulting in a simplified target mask that only covers the read data. For example, an initial mask can be generated based on the key-value index of the acquired key-value entries. The initial mask is a binary vector with a length equal to the total number of KV cache entries, where the selected index positions are 1, and the rest are 0. For example, if the KV cache has 200K key-value entries, the initial mask will have 200K bits. If K is 2048, then 2048 positions in the KV cache will be 1. Based on this, according to the address of the data block actually read, only the mask bits belonging to the read data block are retained, and the remaining mask bits are discarded, thus obtaining a simplified target mask. For example, based on a determined list of addresses of data blocks to be read, mask bits belonging only to the positions within these data blocks are extracted from the initial mask, and all other mask bits are discarded to obtain the target mask.

[0061] Assuming there are M data blocks to be read, each containing B key-value entries, the target mask has a length of M×B, much smaller than the initial mask. The target mask retains only the valid positions (1s) and the rest (0s). Since M×B is typically only slightly larger than K (because the valid entries are concentrated in the M blocks), the target mask is very concise. This significantly reduces the amount of data processed in subsequent masking operations, and the mask can directly correspond to the read data blocks in memory layout, eliminating the need for additional index mapping. Furthermore, the overhead of subsequent masking operations is proportional to the amount of data read, rather than the size of the entire KV cache, thus saving costs.

[0062] Figure 4 A schematic diagram of a simplified mask provided according to some embodiments of the present invention is shown. For example... Figure 4 As shown, the left side is the initial mask, whose length is equal to the total number of entries in the entire KV cache. Only the positions selected by TopK are 1 (black parts), and the rest are 0 (white parts). The right side is the target mask, whose length is equal to the total number of entries contained in all read data blocks. Only the mask bits corresponding to the read data blocks are retained in the target mask; the rest are discarded. It can be seen that the target mask is shorter and more concise than the initial mask.

[0063] In some embodiments, the granularity of data read by the matrix load instruction is determined based on the size of each selected key-value entry, such that the granularity of data read by the matrix load instruction is larger than the size of each selected key-value entry. The size of a key-value entry refers to the number of bytes occupied by a single key-value entry in memory. For example, if the key and value use BF16 format, each BF16 occupies 2 bytes, and a key-value pair is typically 4 bytes. The granularity of data read by the matrix load instruction refers to the minimum fixed size of the data block read continuously from memory in a single matrix load operation. This granularity is typically an alignment value determined during hardware design, such as 64 bytes, 128 bytes, 256 bytes, or 512 bytes.

[0064] The size of the data block is a configurable parameter. Optionally, the data block size can be matched to the cache line size or DMA burst transfer length of hardware such as an AI chip. For example, the data read granularity of a matrix load instruction is an integer multiple of the hardware alignment granularity. By setting an appropriate data read granularity, a balance can be struck between redundancy and hardware efficiency. Larger granularity results in higher efficiency for the matrix load instruction but increases the amount of redundant data; smaller granularity reduces redundancy but may not fully utilize the block load instruction. This application can select an appropriate granularity based on specific hardware characteristics and application scenarios, without imposing any restrictions.

[0065] By setting a granularity larger than that of a single entry, batch reading in blocks can be achieved, thereby leveraging the high bandwidth of the hardware. In other words, setting the data read size to be larger than that of a single key-value entry can fully utilize the burst transfer mode of the memory controller and bus, reduce address handshaking and command overhead, and improve bandwidth utilization.

[0066] It should be noted that the data reading granularity can be dynamically adjusted based on the redundancy ratio and the size of the data blocks. For example, when the estimated redundancy ratio exceeds a preset threshold (e.g., 80%), a smaller data block granularity can be temporarily switched (e.g., using subblocks), or a rollback SIMT load can be used for the few sparse indexes.

[0067] In some embodiments, since the attention calculations corresponding to different data blocks are independent of each other (i.e., key-value entries between different blocks do not interact in the actual attention), multiple computing cores can be used to process different sets of data blocks in parallel. Specifically, a task scheduler can be used to divide the list of data block addresses to be read into multiple sublists and assign them to different computing cores. Each computing core independently performs matrix loading, attention calculations, and masking. Finally, the outputs of each core are concatenated in their original order to obtain the complete attention result. This parallel approach can fully utilize multi-core resources and significantly reduce end-to-end latency.

[0068] In some embodiments, the standard attention mechanism requires reading not only the key matrix K but also the value matrix V. Reading the value matrix can be done in exactly the same way: based on the same TopK index, calculate the address of the data block containing the valid value entries, use matrix loading instructions to read the redundant value data, and mask it using the same target mask within the data block. Since the TopK indices of Q, K, and V are consistent, the mask generation logic can be shared, further improving efficiency.

[0069] Figure 5 A flowchart illustrating the process of determining target attention results according to some embodiments of the present invention is shown. Figure 5 As shown, in step 502, K KV indices for participating in the attention operation are determined via TopK calculation.

[0070] In step 504, each index is mapped to the complete data block it belongs to in physical storage, and the address or address index of each data block is determined. Each data block contains at least one valid key-value entry selected by TopK.

[0071] In step 506, a target mask is determined based on the key-value index of the acquired key-value entry and the data block to be read. The target mask can be used for subsequent masking operations, such as masking the preliminary calculation results to determine the target calculation result.

[0072] In step 508, a matrix loading instruction is invoked to read data from the data block. This involves reading continuous data from memory in blocks with high bandwidth and in batches, and directly loading it into the register of the tensor computation unit. The matrix loading instruction can be, for example, a Matrix Load instruction (e.g., tensor_ld_block) or another matrix loading instruction; the present invention does not limit the form of the matrix loading instruction. It should be noted that the data blocks read from memory are organized into one or more continuous matrices.

[0073] This matrix contains both selected valid data and potentially unselected redundant data. The degree of redundancy depends on the clustering of valid entries within a block. If valid entries are scattered, more blocks need to be read, potentially resulting in a higher redundancy ratio; however, the hardware computing power is sufficient to handle this redundancy, and the resulting high bandwidth gains outweigh the losses. Therefore, the redundancy is appropriate, rather than uncontrollable.

[0074] In step 510, standard dense matrix multiplication (Multi-head Latent Attention, MLA) is performed on the entire read matrix to obtain preliminary calculation results. Since redundant data also participates in the multiplication, the preliminary results may contain both valid and invalid results. Based on this, invalid results can be masked using a target mask (set to -INF or 0). The target mask only covers the read blocks, making it very concise.

[0075] In step 512, the attention calculation result is output. After the target masking operation described above, Softmax and multiplication with the Value matrix are performed to obtain the attention calculation result.

[0076] In this way, starting from the discrete TopK index, sparse and discrete accesses are transformed into dense and continuous block accesses through block mapping and matrix loading instructions, and the redundancy effect is eliminated by using intra-block masks, thereby improving the overall computing performance and resource utilization.

[0077] Figure 6 A schematic diagram of a computing device for performing attention calculation according to an embodiment of the present invention is shown. Figure 6As shown, the computing device 600 includes, for example, an artificial intelligence chip for performing attention calculations. The computing device 600 may have one or more processors and a memory communicatively connected to the one or more processors; wherein the memory stores instructions executable by at least one processor. The processors include dedicated processing units such as graphics processing units (GPUs), field-programmable gate arrays (FPGAs), application-specific integrated circuits (ASICs), general-purpose graphics processors, and general-purpose processors such as CPUs.

[0078] In some embodiments, the computing device 600 further includes at least: a key value index acquisition module 602, an address calculation module 604, a matrix data reading module 606, and a target calculation result acquisition module 608.

[0079] Regarding the key-value index acquisition module 602, it is used to obtain the key-value index corresponding to the multiple selected key-value entries participating in the attention operation via TopK calculation.

[0080] Regarding the address calculation module 604, it is used to calculate the address of multiple data blocks to be read based on the key-value index, where each data block includes the corresponding selected key-value entry. Regarding the matrix data reading module 606, it is used to read multiple redundant matrix data based on the addresses of multiple data blocks using matrix loading instructions, for use in attention operations, thereby generating preliminary calculation results; and Regarding the target calculation result acquisition module 608, it is used to acquire the target calculation result based on the preliminary calculation result and the target mask in the data block.

[0081] The various processes and procedures described above, such as method 200, can be executed at a computing device. This computing device may include, for example, at least one processor (at least one graphics processor and at least one central processing unit); and a memory communicatively connected to the at least one processor; wherein the memory stores instructions executable by the at least one processor. In some embodiments, method 200 may be implemented as a computer software program or program product tangibly contained in a machine-readable medium. In some embodiments, part or all of the computer program may be loaded and / or installed on the computing device via read-only memory (ROM) and / or a communication unit. When the computer program is loaded into random-access memory (RAM) and executed by the GPU and CPU, one or more actions of method 200 described above can be performed.

[0082] This invention can be a method, apparatus, system, and / or computer program product. A computer program product may include a computer-readable storage medium having computer-readable program instructions loaded thereon for performing various aspects of the invention. The computer-readable storage medium may be a tangible device capable of holding and storing instructions used by an instruction execution device. The computer-readable storage medium may be, for example, but not limited to, an electrical storage device, a magnetic storage device, an optical storage device, an electromagnetic storage device, a semiconductor storage device, or any suitable combination thereof.

[0083] The computer-readable program instructions described herein can be downloaded from a computer-readable storage medium to various computing / processing devices, or downloaded via a network, such as the Internet, a local area network, a wide area network, and / or a wireless network, to an external computer or external storage device. Various aspects of the invention are described herein with reference to flowchart illustrations and / or block diagrams of methods, apparatus (systems), and computer program products according to embodiments of the invention. It should be understood that each block of the flowchart illustrations and / or block diagrams, and combinations of blocks in the flowchart illustrations and / or block diagrams, can be implemented by computer-readable program instructions.

[0084] These computer-readable program instructions can be provided to the central processing unit of a general-purpose computer, a special-purpose computer, or other programmable data processing apparatus to produce a machine such that, when executed by the central processing unit of the computer or other programmable data processing apparatus, they create means for implementing the functions / actions specified in one or more blocks of a flowchart and / or block diagram. These computer-readable program instructions can also be stored in a computer-readable storage medium that causes a computer, programmable data processing apparatus, and / or other device to operate in a particular manner. Thus, the computer-readable medium storing the instructions comprises an article of manufacture that includes instructions for implementing aspects of the functions / actions specified in one or more blocks of a flowchart and / or block diagram.

[0085] The flowcharts and block diagrams in the accompanying drawings illustrate the architecture, functionality, and operation of possible implementations of systems, methods, and computer program products according to various embodiments of the present invention. In this regard, each block in a flowchart or block diagram may represent a module, segment, or portion of an instruction, which contains one or more executable instructions for implementing a specified logical function. In some alternative implementations, the functions marked in the blocks may occur in a different order than those shown in the drawings. For example, two consecutive blocks may actually be executed substantially in parallel, and they may sometimes be executed in reverse order, depending on the functions involved. It should also be noted that each block in the block diagrams and / or flowcharts, and combinations of blocks in the block diagrams and / or flowcharts, may be implemented using a dedicated hardware-based system that performs the specified function or action, or using a combination of dedicated hardware and computer instructions.

[0086] It should be understood that the various forms of processes shown above can be used to rearrange, add, or delete steps. For example, the steps described in this application can be executed in parallel, sequentially, or in different orders, as long as the desired result of the technical solution disclosed in this application can be achieved, and this is not limited herein.

[0087] The specific embodiments described above do not constitute a limitation on the scope of protection of this application. Those skilled in the art should understand that various modifications, combinations, sub-combinations, and substitutions can be made according to design requirements and other factors.

Claims

1. A method for performing attention computation, the method comprising: include: By using TopK calculation, obtain the key-value indices corresponding to the multiple selected key-value entries participating in the attention operation; Based on the key-value index, the addresses of multiple data blocks to be read are calculated, and each data block includes the corresponding selected key-value entry; Based on the addresses of the multiple data blocks, a matrix loading instruction is used to read multiple redundant matrix data for attention operations, thereby generating preliminary calculation results. as well as Based on the preliminary calculation results and the target mask within the data block, the target calculation results are obtained.

2. The method of claim 1, wherein, Also includes: At least based on the key-value index, generate the target mask within the data block.

3. The method of claim 2, wherein, At least based on the key-value index, generating the target mask within the data block includes: Based on the key-value index, an initial mask is generated, which is relative to the position of the selected key-value entries in the key-value cache; Based on the initial mask and the multiple redundant matrix data read, the target mask within the data block is determined.

4. The method according to claim 3, characterized in that, The target mask contains fewer mask bits than the initial mask contains.

5. The method according to claim 3, characterized in that, Based on the initial mask and the multiple redundant matrix data read, the target mask within the data block is determined as follows: In the initial mask, the mask values ​​corresponding to the positions of unread key-value entries are removed. The unread key-value entries are the key-value entries in the key-value cache other than multiple redundant matrix data.

6. The method according to claim 1, characterized in that, Based on the key-value index, calculating the addresses of the multiple data blocks to be read includes: Based on the key-value index and the data reading granularity of the matrix loading instruction, the addresses of multiple data blocks to be read are calculated, such that each data block includes a corresponding selected key-value entry. The data reading granularity is an integer multiple of the hardware alignment granularity of the artificial intelligence chip, and the artificial intelligence chip is used to perform the attention calculation.

7. The method according to claim 6, characterized in that, Also includes: The granularity of data read for the matrix loading instruction is determined based on the size of each selected key-value entry, such that the granularity of data read for the matrix loading instruction is larger than the size of each selected key-value entry.

8. The method according to claim 1, characterized in that, The size of the redundant matrix data is larger than the size of the corresponding selected key-value entry.

9. A computing device, characterized in that, include: At least one processor; as well as A memory that is communicatively connected to the at least one processor; in The memory stores instructions that can be executed by the at least one processor to enable the at least one processor to perform the method of any one of claims 1-8.

10. A computer-readable storage medium, characterized in that, The computer-readable storage medium stores a computer program that, when executed by a machine, performs the method according to any one of claims 1-8.

11. A computer program product, characterized in that, Includes a computer program, which, when executed by a machine, performs the method according to any one of claims 1-8.

Images