Attention operator boundary protection method, apparatus, medium, and device

Abstract

The application discloses a boundary protection method and device of an attention operator, a medium and equipment. The method comprises the following steps: loading the data of each block in a key matrix according to a preset block size by using a tensor memory accelerator, and filling a special floating point number to an out-of-bound position of a tail block exceeding an effective sequence length to obtain a filled tail block; wherein the special floating point number comprises negative infinity, positive infinity or a non-number value; obtaining an attention score block corresponding to the filled tail block; and performing a bit operation conversion operation by using a specified thread group to convert all non-finite number values in the attention score block into negative infinity to obtain a converted attention score block, so that the out-of-bound position has no contribution to the calculation of a corresponding attention weight; wherein the bit operation conversion operation is realized based on element floating point number coding and does not include a conditional branch instruction. The application can avoid the thread bundle divergence problem caused by the conditional branch and improve the calculation efficiency of the attention.

Description

Technical Field

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

[0002] With the development of artificial intelligence technology, the attention mechanism has become a core component of various deep learning models. Attention calculation usually involves large matrix multiplication and softmax operations, which can easily create a bottleneck in accessing high-bandwidth memory (HBM).

[0003] Flash Attention is an attention mechanism that reduces HBM access through tiled computation. When the processor executes the Flash Attention kernel, the key matrix is ​​typically tiled along the sequence dimension. The actual sequence length of the key matrix may not be divisible by the target sequence length in the tile size, rendering some positions in the last tile (the tail tile) invalid. This is addressed by querying the matrix Q and the key matrix K, performing a Q×K operation. T When calculating partial blocks of the attention score matrix, these invalid positions need to be processed to avoid errors in the final attention output. Summary of the Invention

[0004] The purpose of this invention is to provide a boundary protection method, apparatus, computer-readable storage medium, and electronic device for attention operators. By using a tensor memory accelerator, special floating-point numbers are automatically filled into the out-of-bounds positions of the tail block. By using a unified bitwise operation to convert non-finite values ​​in the attention score block into negative infinity, the thread bundle divergence problem caused by conditional branches can be avoided, the computational efficiency of attention can be improved, and the accuracy of the attention output results can be guaranteed.

[0005] A first aspect of the present invention provides a boundary protection method for an attention operator, comprising: Using a tensor memory accelerator, data from each block in the key matrix is ​​loaded according to a preset block size. If the effective sequence length of the tail block is less than the target sequence length corresponding to the preset block size, special floating-point numbers are used to fill the out-of-bounds positions that exceed the effective sequence length, thus obtaining the filled tail block. The special floating-point numbers include negative infinity, positive infinity, or non-numerical values. The padded tail block and the corresponding query block are multiplied by a matrix to obtain the corresponding attention score block. By specifying a thread group, bitwise operation transformation is performed in parallel on the elements in the attention score block to convert all non-finite values ​​in the attention score block to negative infinity, resulting in a transformed attention score block, so that the out-of-bounds position does not contribute to the calculation of the corresponding attention weight; wherein, the bitwise operation transformation is implemented based on the floating-point encoding of the elements and does not include conditional branch instructions.

[0006] Optionally, the bitwise conversion operation is executed in parallel by each thread within a specified thread group, and each thread processes the corresponding element according to the same instruction sequence.

[0007] Optionally, the bitwise conversion operation includes: Extract the first exponent bit from the floating-point code corresponding to the current element; Based on the first exponent bit and using bitwise operation instructions, choose to retain the original finite value of the current element, or rewrite the non-finite value of the current element as negative infinity.

[0008] Optionally, the step of selecting to retain the original finite value of the current element or rewriting the non-finite value of the current element to negative infinity based on the first exponent bit and using bitwise operation instructions includes: A binary mask is generated based on the numerical characteristics of the first exponent bit; wherein the bit width of the binary mask is the same as the bit width of the floating-point code corresponding to the current element, and the value of the binary mask is either all one or all zero. For the current element, the binary mask is used to select and output either the original finite value or negative infinity.

[0009] Optionally, generating a binary mask based on the numerical characteristics of the first exponent bit includes: The first exponent bit is logically compared with the second exponent bit corresponding to the non-finite value to obtain the corresponding comparison result. Subtract one from the comparison result, and then arithmetically right-shift the result by N bits to obtain the binary mask with a bit width of N+1; where N≥1.

[0010] Optionally, selecting to output the original finite value or negative infinity using the binary mask includes: Perform a bitwise AND operation between the binary mask and the floating-point code for negative infinity to obtain the first intermediate result; After inverting the binary mask bit by bit, perform a bitwise AND operation with the floating-point code of the current element to obtain the second intermediate result; Perform a bitwise OR operation between the first intermediate result and the second intermediate result to obtain the floating-point code of the current element selection output.

[0011] Optionally, the method further includes: An online softmax operation is performed on the transformed attention score blocks to obtain the corresponding attention weight matrix.

[0012] A second aspect of the present invention provides a boundary protection device for an attention operator, comprising: The out-of-bounds filling module is used to load the data of each block in the key matrix according to the preset block size using a tensor memory accelerator, and fill the out-of-bounds positions that exceed the effective sequence length with special floating-point numbers when the effective sequence length of the tail block is less than the target sequence length corresponding to the preset block size, so as to obtain the filled tail block; wherein, the special floating-point numbers include: negative infinity, positive infinity or non-numerical. The matrix calculation module is used to perform matrix multiplication on the padded tail block and the corresponding query block to obtain the corresponding attention score block; The data conversion module is used to perform bitwise operation conversion on the elements in the attention score block in parallel through a specified thread group, converting all non-finite values ​​in the attention score block into negative infinity to obtain the converted attention score block, so that the out-of-bounds position does not contribute to the calculation of the corresponding attention weight; wherein, the bitwise operation conversion is implemented based on the floating-point encoding of the elements and does not include conditional branch instructions.

[0013] 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 boundary protection method of the attention operator described in any embodiment of the first aspect.

[0014] A fourth 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 boundary protection method of the attention operator described in any embodiment of the first aspect.

[0015] Compared with existing technologies, embodiments of the present invention provide a boundary protection method, apparatus, computer-readable storage medium, and electronic device for attention operators. Through a tensor memory accelerator, special floating-point numbers are automatically filled into the out-of-bounds positions of the tail block during the data loading stage. Furthermore, through bitwise operations, non-finite values ​​in the attention score blocks are uniformly converted to negative infinity, ensuring that out-of-bounds positions do not contribute to the calculation of the corresponding attention weights in the subsequent normalized exponent summation process, thus guaranteeing the accuracy of the final attention output. Simultaneously, since the bitwise operations are implemented through bitwise instructions without any branch conditional jump instructions (i.e., if-else branches), all threads within the specified thread group execute the exact same instruction sequence synchronously, avoiding thread bundle divergence and significantly improving hardware utilization and attention computation efficiency. Attached Figure Description

[0016] Figure 1 This is a flowchart illustrating an embodiment of the boundary protection method for the attention operator provided by the present invention; Figure 2 This is a schematic diagram of a structure of an embodiment of the boundary protection device for the attention operator provided by the present invention; Figure 3 This is a schematic diagram of the structure of an embodiment of the electronic device provided by the present invention; Figure 4 This is a schematic diagram of the structure of an embodiment of the artificial intelligence processor provided by the present invention. Detailed Implementation

[0017] 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.

[0018] In related calculations, independent explicit mask matrices are typically pre-constructed. For example, after calculating the attention score block (denoted as S_tile), the mask matrix is ​​multiplied by S_tile to mask invalid locations and perform boundary protection (i.e., data outside the boundary will not contaminate the final calculation result). However, this method requires additional storage of the mask matrix, resulting in a large memory overhead, and also requires additional element-wise multiplication operations (i.e., matrix multiplication), increasing the overall computational cost.

[0019] To address this, the inventors implemented boundary protection using a "coordinate calculation and branch judgment" method. Specifically, within each thread or warp, the coordinates (denoted as col) of the current attention score in the sequence dimension are calculated in real-time. A conditional judgment (if-else) determines whether to set this coordinate to -inf (negative infinity). If the current attention score's col is greater than the effective sequence length of the key matrix (denoted as seql_k), the coordinate is considered out of bounds, and the current attention score is assigned the value -inf. Otherwise, the coordinate is considered within bounds and its original value is retained. Subsequently, during the Softmax operation, since -inf becomes zero after exponential operation, it does not participate in subsequent normalization summation calculations, and the corresponding attention weight is 0, thus not affecting subsequent calculations.

[0020] After further research, the inventors discovered that although the "coordinate calculation and branch judgment" method can solve the problems of memory overhead and extra matrix multiplication operations in explicit mask matrices, it still has the following technical defects: (1) Heavy Single Instruction Multiple Threads (SIMT) overhead: Integer coordinate calculation and comparison operations need to be performed on each element in the block (such as if (col>= seql_k)), generating a large number of Arithmetic Logic Unit (ALU) instructions, thereby increasing the computational load; (2) There is a problem of thread bundle divergence: the conditional branch instruction of the out-of-bounds judgment causes the threads in the same thread bundle to go to different execution branches (some threads perform out-of-bounds assignment operations, and some threads skip the operation). The threads in different branches will be executed serially, thus wasting the parallel throughput of the SIMT architecture. (3) Not friendly to weak SIMT architecture: In some AI acceleration chips or processors, the computing power of their SIMT units is relatively weak and the branch prediction capability is limited. The above integer operations and branch divergence will become performance bottlenecks, increasing the execution latency of the FlashAttention kernel function.

[0021] See Figure 1 This is a flowchart illustrating an embodiment of the boundary protection method for the attention operator provided by the present invention.

[0022] To address the aforementioned technical problems, a first aspect of the present invention provides a boundary protection method for an attention operator, comprising steps S1 to S3, as detailed below: Step S1: Using a tensor memory accelerator, load the data of each block in the key matrix according to the preset block size, and when the effective sequence length of the tail block is less than the target sequence length corresponding to the preset block size, fill the out-of-bounds positions that exceed the effective sequence length with special floating-point numbers to obtain the filled tail block; wherein, the special floating-point numbers include: negative infinity, positive infinity or non-numerical numbers; Step S2: Perform matrix multiplication on the padded tail block and the corresponding query block to obtain the corresponding attention score block; Step S3: By specifying a thread group, perform bitwise operation conversion on the elements in the attention score block in parallel to convert all non-finite values ​​in the attention score block into negative infinity, so as to obtain the converted attention score block, so that the out-of-bounds position does not contribute to the calculation of the corresponding attention weight; wherein, the bitwise operation conversion is implemented based on the floating-point encoding of the elements and does not include conditional branch instructions.

[0023] It should be noted that in step S1, before starting the Flash Attention Kernel, the boundary padding parameters of the Tensor Memory Accelerator (TMA) are pre-configured. That is, for the key matrix to be loaded, its out-of-bounds padding value is set to a special floating-point number, such as negative infinity (-inf), positive infinity (+inf), or not-numerical (NaN); among which, -inf is the preferred out-of-bounds padding value.

[0024] After configuration, the TMA loads the data of each block in the key matrix into the on-chip cache according to the preset block size. The TMA has a hardware automatic boundary check function. When the effective sequence length of the tail block of the loaded key matrix is ​​less than the target sequence length corresponding to the preset block size, the TMA automatically fills the out-of-bounds positions in the tail block that exceed the effective sequence length with a pre-set special floating-point number (such as -inf) to obtain the padded tail block.

[0025] In step S2, the padded tail block obtained in step S1 is multiplied by the corresponding query block to obtain the corresponding attention score block. Specifically, let the size of the current query block (denoted as Q_tile) be [Br, d], the size of the key block (denoted as K_tile) be [chunk_K, d], and the size of the output attention score block (denoted as S_tile) be [Br, chunk_K]; where Br is the sequence length of the query block, d is the head dimension (headdim), and chunk_K is the target sequence length of the key block.

[0026] The element value S_tile[i][j] in the i-th row and j-th column of the attention score block is = Q[i][:]·K[j][:]. Since the out-of-bounds positions in the tail block have been filled with special floating-point numbers (such as -inf), according to the IEEE 754 floating-point arithmetic rules: if either operand of the dot product operation is -inf and the other operand is a finite value, then the dot product result at the corresponding position is -inf, inf, or NaN (-inf × 0 = NaN), which are all non-finite values. In summary, in this embodiment of the invention, the attention score corresponding to the out-of-bounds position must belong to a "non-finite value" (i.e., -inf, inf, or NaN).

[0027] In step S3, the attention score blocks calculated in step S2 are assigned to a designated thread group (such as a thread bundle). All threads within the thread group perform bitwise transformation operations on the elements in the attention score blocks in parallel to convert all non-finite values ​​in the attention score blocks to -inf, while finite values ​​remain unchanged, resulting in the transformed attention score blocks. The bitwise transformation operations are based on the floating-point encoding of the elements and are primarily implemented through bitwise operation instructions, without any conditional branch (if-else) instructions or integer coordinate checks.

[0028] After the transformation, the attention score corresponding to all out-of-bounds positions is -inf. An online softmax operation is then performed on the transformed attention scores in blocks to obtain the corresponding attention weight matrix. According to the rules of exponential operation, exp(-inf) = 0, and it does not participate in the normalization summation, ensuring that out-of-bounds positions do not contribute to the calculation of the corresponding attention weights (i.e., the attention weight corresponding to an out-of-bounds position is zero), thus not affecting subsequent calculations and completing the boundary protection of the attention operator.

[0029] In an optional embodiment, the bitwise transformation operation is executed in parallel by threads within a specified thread group, and each thread processes the corresponding element according to the same instruction sequence.

[0030] It should be noted that all threads within a specified thread group (such as a thread bundle) must execute the exact same instructions within the same clock cycle. If different threads within a thread bundle execute different instruction sequences, thread bundle divergence will occur, causing threads in different branches to be executed serially, resulting in a significant decrease in hardware utilization.

[0031] In this embodiment of the invention, each thread within a specified thread group executes the same instructions at the same time, only operating on different data; after each thread independently completes the transformation of the corresponding element, the transformed attention score block is obtained. Since bitwise operation transformation does not include any conditional branch instructions, but is implemented through bitwise operation instructions, there will be no control flow divergence within the thread group, and all threads always maintain a unified execution path (i.e., each thread executes the exact same instruction sequence), maximizing the utilization of the hardware computing unit and fundamentally eliminating the performance loss caused by thread bundle divergence.

[0032] In an optional embodiment, the bitwise conversion operation includes: Extract the first exponent bit from the floating-point code corresponding to the current element; Based on the first exponent bit and using bitwise operation instructions, choose to retain the original finite value of the current element, or rewrite the non-finite value of the current element as negative infinity.

[0033] It should be noted that, according to the IEEE 754 floating-point standard, the binary encoding of all floating-point numbers (i.e., floating-point encoding) consists of three parts: the sign bit, the exponent bit, and the mantissa bit. The exponent bit is used to quickly identify whether the current element is a non-finite value.

[0034] Based on the extracted numerical characteristics of the first exponent bit, and using bitwise operation instructions, a binary selection operation is performed: If the first index bit is not all 1, it means that the current element is a finite value, and its original value is retained unchanged. If all values ​​of the first exponent are 1 (e.g., 0xFF), it indicates that the current element is a non-finite value (NaN, -inf, or inf), and it is rewritten as -inf.

[0035] The above binary selection process does not include any conditional branch instructions and does not require jump operations. This invention is applicable to all standard-compliant floating-point formats, including 32-bit single-precision, 64-bit double-precision, and 16-bit half-precision. The choice of floating-point format depends on the specific product or technology in which this invention is applied.

[0036] In an optional embodiment, the step of selecting to retain the original finite value of the current element or rewriting the non-finite value of the current element to negative infinity based on the first exponent bit and using bitwise operation instructions includes: A binary mask is generated based on the numerical characteristics of the first exponent bit; wherein the bit width of the binary mask is the same as the bit width of the floating-point code corresponding to the current element, and the value of the binary mask is either all one or all zero. For the current element, the binary mask is used to select and output either the original finite value or negative infinity.

[0037] Further, the step of selecting to output the original finite value or negative infinity using the binary mask includes: Perform a bitwise AND operation between the binary mask and the floating-point code for negative infinity to obtain the first intermediate result; After inverting the binary mask bit by bit, perform a bitwise AND operation with the floating-point code of the current element to obtain the second intermediate result; Perform a bitwise OR operation between the first intermediate result and the second intermediate result to obtain the floating-point code of the current element selection output.

[0038] It should be noted that a binary mask (denoted as mask) can be generated without branches using bitwise operations (such as shift and XOR): when all bits of the first exponent are 1 (corresponding to a non-finite value), the generated binary mask is all one; otherwise, the generated binary mask is all zero. The bit width of this binary mask is exactly the same as the bit width of the floating-point code corresponding to the current element (32 bits).

[0039] For the current element, a branchless binary selection operation is performed using the generated binary mask, ultimately outputting either the original finite value or -inf. Specifically, this can be achieved using `result = (mask & neg_inf_bits) | (~mask & data)`; where `neg_inf_bits` is the floating-point encoding of -inf, `data` is the floating-point encoding of the current element, and `result` is the floating-point encoding of the selected output for the current element. When the mask is all ones, `result` is the floating-point encoding of -inf; when the mask is all zeros, `result` is the floating-point encoding of the current element. Finally, the selected floating-point encoding is reinterpreted as a floating-point number, thus completing the conversion of the current element.

[0040] In an optional embodiment, generating a binary mask based on the numerical characteristics of the first exponent bit includes: The first exponent bit is logically compared with the second exponent bit corresponding to the non-finite value to obtain the corresponding comparison result. Subtract one from the comparison result, and then arithmetically right-shift the result by N bits to obtain the binary mask with a bit width of N+1; where N≥1.

[0041] It should be noted that the second exponent bits are all 1, and logical comparison can be achieved through bitwise XOR operation: XOR the first exponent bits with the second exponent bits. If they are equal, the comparison result is zero; otherwise, the comparison result is a non-zero value.

[0042] By utilizing the sign bit extension property of arithmetic right shift (filling in the sign bit, 0 for positive numbers, and 1 for negative numbers), the comparison results of 0 and non-zero can be converted into a binary mask of all 0s or all 1s: when the comparison result is 0, 0 minus 1 equals -1, and negative numbers are arithmetically right-shifted, with the high-order bits filled with 1. If the arithmetic right shift is N=31 bits, a 32-bit binary mask of all 1s is obtained; when the comparison result is non-zero, any non-zero value minus 1 results in a non-negative number with its highest bit being 0. Arithmetic right shift is N=31 bits, resulting in a 32-bit binary mask of all zeros.

[0043] For example, the bitwise conversion operation in this embodiment of the invention can be performed by the following bitwise conversion function convert_to_neg_inf_if_inf_or_nan, which performs binary selection on each element (non-finite values ​​are converted to -inf, and finite values ​​retain their original values), as follows: __device__ float convert_to_neg_inf_if_inf_or_nan(float f) { / / Indicates that this function is executed on the processor device and called by the thread. / / 1. Define a union x, and interpret the floating-point number f bit by bit as a 32-bit unsigned integer xu union { float f; uint32_t u;} x = { .f = f}; / / 2. Extract the exponent bits (bits 30-23). ​​The exponent stores the 8-bit exponent value of the floating-point number. uint32_t exponent = (xu>>23)&0xFF; / / 3. Determine if it is inf or nan: all exponent bits are 1 (i.e., exponent == 0xFF) / / Tip: (exponent ^ 0xFF) is 0 if they are all equal; after subtracting 1, the highest bit becomes 1 (for the case where (exponent ^ 0xFF) equals 0, a borrow will occur). For the 0xFFFFFFFF (represented as -1 in floating-point numbers) generated by the borrow, after arithmetic right shift by 31 bits, it is still 0xFFFFFFFF (all 1s). / / Shift right by 31 bits to obtain the binary mask is_inf_nan: if it is ±inf / NaN, then is_inf_nan = 0xFFFFFFFF, otherwise it is 0. uint32_t is_inf_nan = ((exponent ^ 0xFF) - 1)>>31; / / 4. Bit selection: / / If is_inf_nan is all 1s, choose to output -inf (0xFF800000) / / Otherwise, retain the original bit mode xu = (is_inf_nan&0xFF800000) | (~is_inf_nan&x.u); return xf; / / Returns the floating-point member xf of union x. At this point, xu has been updated, and the corresponding floating-point number is the converted result. } It is worth noting that the implementation instructions in the above-mentioned conversion function convert_to_neg_inf_if_inf_or_nan are merely examples, and not the only implementation of this invention. Some of the implementation instructions in the above-mentioned conversion function can be implemented through the following embodiments: Example 1: Use the carry bit of the 8-bit first exponent bit to directly generate a binary mask, that is, first exponent bit + 1. If the first exponent bit = 0xFF, the carry bit generated by the overflow is 1; if the first exponent bit ≠ 0xFF, there is no carry, and the carry bit is 0; after negating the carry bit, it is expanded into 32 bits of all 1s / all 0s. Example 2: Perform an equality comparison operation between the extracted first exponent bit and 0xFF, return a boolean result, take the negative of the result and expand it to 32 bits of all 1s or all 0s; Example 3: The bit selection operation is xu = xu ^ ((xu ^ 0xFF800000)&is_inf_nan); if is_inf_nan is all one, then select the output -inf (0xFF800000); otherwise, retain the original bit mode.

[0044] In summary, compared with boundary protection methods based on coordinates and branch judgment, the boundary protection method using the attention operator provided in the first aspect of this invention has the following advantages: (1) Eliminate integer coordinate calculation: Related technologies require calculating row and column indices based on thread coordinates and performing integer comparisons (if (col>= seq_len)) within each thread, generating a large number of ALU instructions; the embodiments of the present invention do not require executing comparison instructions of the form col>= seql_k, nor do they require mapping to global coordinates through thread identifiers, thereby reducing the number of ALU instructions; (2) Eliminate thread bundle divergence: There are no if-else branches in the entire processing path, and all threads in the same warp execute the exact same instruction sequence (bit operations only), avoiding the serialization overhead caused by thread branches; (3) Make full use of hardware characteristics: transfer the computational burden of boundary detection and filling from software to TMA hardware unit; the boundary filling operation is automatically completed during the data transfer from HBM to on-chip cache, without occupying any resources of the computing unit, and the efficiency of TMA in processing boundary filling is much higher than that of general ALU. (4) Extremely low computational overhead: Bitwise operation conversion functions require only a very small number of pure integer instructions (such as shift, XOR, subtraction, bitwise AND, bitwise OR), and their computational overhead is far lower than the combined overhead of "coordinate calculation + branch judgment + mask writing"; (5) Particularly suitable for chips with weak SIMT computing power: For some processors or artificial intelligence accelerators with relatively weak integer operation and branch performance compared to floating point computing power, the embodiments of the present invention can significantly improve the execution efficiency of Flash Attention in scenarios with unaligned sequence lengths and greatly reduce the overhead of boundary processing.

[0045] See Figure 2 This is a schematic diagram of an embodiment of the boundary protection device for the attention operator provided by the present invention.

[0046] A second aspect of the present invention provides a boundary protection device for an attention operator, comprising: The out-of-bounds filling module 11 is used to load the data of each block in the key matrix according to the preset block size using a tensor memory accelerator, and fill the out-of-bounds positions that exceed the effective sequence length with special floating-point numbers when the effective sequence length of the tail block is less than the target sequence length corresponding to the preset block size, so as to obtain the filled tail block; wherein, the special floating-point numbers include: negative infinity, positive infinity or non-numerical. Matrix calculation module 12 is used to perform matrix multiplication calculation on the padded tail block and the corresponding query block to obtain the corresponding attention score block; The data conversion module 13 is used to perform bitwise operation conversion on the elements in the attention score block in parallel through a specified thread group, converting all non-finite values ​​in the attention score block into negative infinity, so as to obtain the converted attention score block, so that the out-of-bounds position does not contribute to the calculation of the corresponding attention weight; wherein, the bitwise operation conversion is implemented based on the floating-point encoding of the elements and does not include conditional branch instructions.

[0047] It should be noted that the boundary protection device for the attention operator provided in the second aspect embodiment of the present invention can realize all the processes of the boundary protection method for the attention operator 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 boundary protection method for the attention operator described in any of the first aspect embodiments, and will not be repeated here.

[0048] 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 boundary protection method of the attention operator described in any of the first aspects of the present invention.

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

[0050] A fourth 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, when executing the computer program, implements the boundary protection method of the attention operator described in any of the first aspects of the present invention.

[0051] 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.

[0052] 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.

[0053] 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.

[0054] 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 3 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.

[0055] See Figure 4 This is a schematic diagram of an embodiment of the artificial intelligence processor provided by the present invention.

[0056] The artificial intelligence processor provided in this embodiment of the invention includes multiple computing units (CUs). The CU is the core component for performing computing tasks and can run multiple threads in parallel. The threads can be divided into multiple consumer thread groups (such as consumer thread group 1 and consumer thread group 2) and one producer thread group.

[0057] The producer thread group is responsible for loading the key matrix block data from global memory (GLM) to group shared memory (GSM) according to the preset block size through the TMA unit. When loading the tail block, the TMA unit automatically fills the out-of-bounds positions that exceed the effective sequence length with special floating-point numbers. The consumer thread group is responsible for performing matrix multiplication calculations between query blocks and key blocks, bitwise operation conversion operations (i.e., converting all non-finite values ​​in the attention score block corresponding to the tail block to -inf through bitwise operations), and softmax operations.

[0058] GSM is used to cache query blocks, key blocks, and value blocks loaded by the producer thread group, as well as attention score blocks, transformed attention score blocks, and attention output blocks generated during the calculation process by the consumer thread group. CLM has the characteristics of large capacity and high bandwidth, and is used to store the original query matrix, key matrix, value matrix, and the final result of attention calculation.

[0059] 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 boundary protection method for an attention operator, characterized in that, include: Using a tensor memory accelerator, data from each block in the key matrix is ​​loaded according to a preset block size. If the effective sequence length of the tail block is less than the target sequence length corresponding to the preset block size, special floating-point numbers are used to fill the out-of-bounds positions that exceed the effective sequence length, thus obtaining the filled tail block. The special floating-point numbers include negative infinity, positive infinity, or non-numerical values. The padded tail block and the corresponding query block are multiplied by a matrix to obtain the corresponding attention score block. By specifying a thread group, bitwise operation transformation is performed in parallel on the elements in the attention score block to convert all non-finite values ​​in the attention score block to negative infinity, resulting in a transformed attention score block, so that the out-of-bounds position does not contribute to the calculation of the corresponding attention weight; wherein, the bitwise operation transformation is implemented based on the floating-point encoding of the elements and does not include conditional branch instructions.

2. The boundary protection method for the attention operator as described in claim 1, characterized in that, The bitwise transformation operations are executed in parallel by each thread within the specified thread group, and each thread processes the corresponding element according to the same instruction sequence.

3. The boundary protection method for the attention operator as described in claim 1, characterized in that, The bitwise conversion operations include: Extract the first exponent bit from the floating-point code corresponding to the current element; Based on the first exponent bit and using bitwise operation instructions, choose to retain the original finite value of the current element, or rewrite the non-finite value of the current element as negative infinity.

4. The boundary protection method for the attention operator as described in claim 3, characterized in that, The step of selecting to retain the original finite value of the current element or rewriting the non-finite value of the current element to negative infinity based on the first exponent bit and using bitwise operation instructions includes: A binary mask is generated based on the numerical characteristics of the first exponent bit; wherein the bit width of the binary mask is the same as the bit width of the floating-point code corresponding to the current element, and the value of the binary mask is either all one or all zero. For the current element, the binary mask is used to select and output either the original finite value or negative infinity.

5. The boundary protection method for the attention operator as described in claim 4, characterized in that, The step of generating a binary mask based on the numerical characteristics of the first exponent bit includes: The first exponent bit is logically compared with the second exponent bit corresponding to the non-finite value to obtain the corresponding comparison result. Subtract one from the comparison result, and then arithmetically right-shift the result by N bits to obtain the binary mask with a bit width of N+1; where N≥1.

6. The boundary protection method for the attention operator as described in claim 4, characterized in that, The step of selecting to output the original finite value or negative infinity using the binary mask includes: Perform a bitwise AND operation between the binary mask and the floating-point code for negative infinity to obtain the first intermediate result; After inverting the binary mask bit by bit, perform a bitwise AND operation with the floating-point code of the current element to obtain the second intermediate result; Perform a bitwise OR operation between the first intermediate result and the second intermediate result to obtain the floating-point code of the current element selection output.

7. The boundary protection method for the attention operator as described in claim 1, characterized in that, The method further includes: An online softmax operation is performed on the transformed attention score blocks to obtain the corresponding attention weight matrix.

8. A boundary protection device for an attention operator, characterized in that, include: The out-of-bounds filling module is used to load the data of each block in the key matrix according to the preset block size using a tensor memory accelerator, and fill the out-of-bounds positions that exceed the effective sequence length with special floating-point numbers when the effective sequence length of the tail block is less than the target sequence length corresponding to the preset block size, so as to obtain the filled tail block; wherein, the special floating-point numbers include: negative infinity, positive infinity or non-numerical. The matrix calculation module is used to perform matrix multiplication on the padded tail block and the corresponding query block to obtain the corresponding attention score block; The data conversion module is used to perform bitwise operation conversion on the elements in the attention score block in parallel through a specified thread group, converting all non-finite values ​​in the attention score block into negative infinity to obtain the converted attention score block, so that the out-of-bounds position does not contribute to the calculation of the corresponding attention weight; wherein, the bitwise operation conversion is implemented based on the floating-point encoding of the elements and does not include conditional branch instructions.

9. 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 is located to perform the boundary protection method of the attention operator as described in any one of claims 1 to 7.

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 boundary protection method of the attention operator as described in any one of claims 1 to 7.

Images