Data processing method, data processing chip and computer device

Abstract

The application provides a data processing method, a data processing chip and a computer device. The method is applied to any core included in the data processing chip, and the core is used for processing a corresponding data block in to-be-processed data. The method comprises the following steps: dividing the data block corresponding to the core into a plurality of sub-data blocks; multiplying each sub-data block with a first mask matrix to obtain a first prefix-sum matrix corresponding to each sub-data block, wherein the first prefix-sum matrix represents a prefix-sum result in each sub-data block; determining a second prefix-sum matrix corresponding to the data block based on a second mask matrix and the first prefix-sum matrix; and the second prefix-sum matrix represents a prefix-sum result in the core. According to the application, the mask storage cost, the calculation cost and the calculation complexity can be reduced.

Description

Technical Field

[0001] This application relates to, but is not limited to, data processing technology, and in particular to a data processing method, a data processing chip, and a computer device. Background Technology

[0002] Prefix sums are a fundamental operator widely used in various deep learning models. Among related technologies, the Blech algorithm can compute prefix sums in parallel. It constructs a bottom-up binary tree, where nodes with even indices (starting from 1) aggregate the values ​​of preceding nodes (upward reduction), and then propagates the prefix sum to child nodes through a top-down binary tree (downward distribution). However, when the Blech algorithm performs interval addition on n elements, the size of the first-level mask matrix is ​​n×n (n is on the order of tens of thousands in large language models), making the storage cost of the mask matrix extremely high. Furthermore, the intervals in each level of the Blech algorithm are different: the first level adds adjacent elements (interval of 1), the second level adds elements with an interval of 2, and the third level adds elements with an interval of 4. Therefore, a large number of mask matrices with different patterns are needed to achieve selective addition at different levels, significantly increasing the number and management cost of the mask matrices. Summary of the Invention

[0003] This application provides a data processing method, a data processing chip, and a computer device.

[0004] The technical solution of this application embodiment is implemented as follows:

[0005] This application provides a data processing method applied to any core in a data processing chip. The core is used to process a corresponding data block in the data to be processed, including: dividing the data block corresponding to the core into multiple sub-data blocks; multiplying each sub-data block by a first mask matrix to obtain a first prefix sum matrix corresponding to each sub-data block, the first prefix sum matrix representing the prefix sum result within each sub-data block; determining a second prefix sum matrix corresponding to the data block based on a second mask matrix and the first prefix sum matrix; the second prefix sum matrix representing the prefix sum result within the core.

[0006] This application provides a data processing chip, in which any core is used to: process a corresponding data block in the data to be processed; divide the data block corresponding to the core into multiple sub-data blocks; multiply each sub-data block by a first mask matrix to obtain a first prefix sum matrix corresponding to each sub-data block, wherein the first prefix sum matrix represents the prefix sum result within each sub-data block; and determine a second prefix sum matrix corresponding to the data block based on a second mask matrix and the first prefix sum matrix; wherein the second prefix sum matrix represents the prefix sum result within the core.

[0007] This application provides a computer device including the data processing chip described above.

[0008] The technical solutions provided by the embodiments of this application may include the following beneficial effects:

[0009] In this embodiment, during the processing of a corresponding data block, any core in the data processing chip divides the data block into multiple sub-data blocks and multiplies each sub-data block by a first mask matrix to obtain a first prefix sum matrix that represents the prefix sum result within each sub-data block. Then, based on the second mask matrix and the first prefix sum matrix, the second prefix sum matrix corresponding to the data block of that core is determined, and the second prefix sum matrix represents the prefix sum result within that core. In this way, on the one hand, by dividing the data block corresponding to each core into multiple sub-data blocks, the size of the mask matrix required for calculating the prefix sum is significantly reduced, and the prefix sum calculation corresponding to the data block is decomposed into multiple simple matrix operations, thereby effectively reducing both mask storage costs and computational costs and complexity. On the other hand, the prefix sum calculation process within the core can be completed using only the first and second mask matrices, eliminating the need to construct multiple sets of mask matrices based on the different summation intervals of each layer to select elements for calculation, reducing the number of mask matrices, thereby reducing the complexity of the prefix sum calculation and mask management overhead, making the prefix sum calculation process more stable and controllable.

[0010] It should be understood that the above general description and the following detailed description are exemplary and explanatory only, and do not limit this application. Attached Figure Description

[0011] The accompanying drawings, which are incorporated in and form part of this specification, illustrate embodiments consistent with this application and, together with the description, serve to explain the principles of this application.

[0012] Figure 1 This is a schematic diagram of the composition structure of a data processing chip provided in an embodiment of this application.

[0013] Figure 2 This is a schematic diagram illustrating the implementation flow of a data processing method provided in an embodiment of this application.

[0014] Figure 3 This is a schematic diagram of a sub-data block prefix sum calculation provided in an embodiment of this application.

[0015] Figure 4 This is a schematic diagram illustrating the calculation of inter-block prefix sums between sub-data blocks provided in an embodiment of this application.

[0016] Figure 5 This is a schematic diagram of prefix sum calculation within a core provided in an embodiment of this application.

[0017] Figure 6 This is a schematic diagram of an inter-core prefix and sum provided in an embodiment of this application.

[0018] Figure 7 This is a schematic diagram of the total prefix of each core provided in an embodiment of this application.

[0019] Figure 8 This is a schematic diagram of the composition structure of a computer device provided in an embodiment of this application. Detailed Implementation

[0020] To make the objectives, technical solutions, and advantages of this application clearer, the application will be further described in detail below with reference to the accompanying drawings. The described embodiments should not be regarded as limitations on this application. All other embodiments obtained by those skilled in the art without creative effort are within the scope of protection of this application.

[0021] In the following description, references are made to “some embodiments,” which describe a subset of all possible embodiments. However, it is understood that “some embodiments” may be the same subset or different subsets of all possible embodiments and may be combined with each other without conflict.

[0022] In the following description, the terms "first, second, third" are used merely to distinguish similar objects and do not represent a specific ordering of objects. It is understood that "first, second, third" may be interchanged in a specific order or sequence where permitted, so that the embodiments of this application described herein can be implemented in an order other than that illustrated or described herein.

[0023] In the embodiments of this application, the terms "module" or "unit" refer to a computer program or part of a computer program that has a predetermined function and works with other related parts to achieve a predetermined goal, and can be implemented wholly or partially using software, hardware (such as processing circuitry or memory), or a combination thereof. Similarly, a processor (or multiple processors or memory) can be used to implement one or more modules or units. Furthermore, each module or unit can be part of an overall module or unit that includes the functionality of that module or unit.

[0024] Unless otherwise defined, all technical and scientific terms used in the embodiments of this application have the same meaning as commonly understood by one of ordinary skill in the art to which this application belongs. The terminology used in the embodiments of this application is for the purpose of describing the embodiments of this application only and is not intended to limit this application.

[0025] Before providing a further detailed description of the embodiments of this application, the nouns and terms involved in the embodiments of this application will be explained, and the nouns and terms involved in the embodiments of this application shall be interpreted as follows.

[0026] 1) Prefix sum (Cumsum) is used to efficiently calculate the sum of elements from the start position to the current position in a sequence. It is a basic operator widely used in various deep learning models, especially large language models such as qwen3next and mamba2.

[0027] 2) A mask matrix is ​​a special matrix used in matrix operations or data processing to mark, filter, or control specific positions. It is used to indicate which positions need to be retained, ignored, or specially processed by means of element values ​​(usually 0 or 1, or other conventional values).

[0028] 3) The Blech algorithm is a Blech tree-based prefix sum calculation method. It constructs a bottom-up binary tree, where nodes with even indices (starting from 1) aggregate the values ​​of their preceding nodes (upward reduction), and then propagate the prefix sum to child nodes through a top-down binary tree (downward distribution). The Blech algorithm reduces the time complexity of element-wise prefix sum calculation from O(n) to O(log n), making it a widely adopted fundamental parallel primitive in parallel computing hardware in related technologies.

[0029] 4) Single Instruction Multiple Threads (SIMT): Multiple threads execute the same instructions simultaneously but process different data. Threads can access memory independently and support branching and dynamic scheduling.

[0030] This application provides a data processing chip. Figure 1 This is a schematic diagram of the composition structure of a data processing chip provided in an embodiment of this application, as shown below. Figure 1 As shown, the data processing chip 100 includes multiple cores 101, each core 101 being used to process the corresponding data block in the data to be processed.

[0031] Data processing chip 100 is an integrated circuit chip used to efficiently perform data calculation, processing, analysis and conversion tasks.

[0032] In some implementations, the data processing chip 100 may be a chip based on a Single Instruction Multiple Data (SIMD) architecture. For example, the data processing chip may be at least one of a Tensor Processing Unit (TPU), a Neural Processing Unit (NPU), etc.

[0033] Single Instruction Multiple Data (SID) refers to performing the same operation on multiple data simultaneously with a single instruction. The data must be stored contiguously, the operation type is fixed, and the branching capability is relatively weak (it needs to be optimized through masking or loop expansion).

[0034] In some implementations, the data processing chip 100 may be a hardware computing unit based on the Tensor Core architecture with ultra-large-scale matrix multiplication capabilities, which has a performance advantage over the Bllelloch algorithm in related technologies in the context of prefix sum computation of ultra-long sequences.

[0035] In some implementations, the data processing chip 100 can perform prefix sum, cumulative sum, and other calculations on the data to be processed, and can divide the data to be processed into multiple data blocks, with multiple cores in the data processing chip processing multiple data blocks respectively.

[0036] For example, a data processing chip includes a first core, a second core, and a third core, which divides the data to be processed into a first data block, a second data block, and a third data block; the first data block is assigned to the first core for processing, the second data block is assigned to the second core for computation, and the third data block is assigned to the third core for computation.

[0037] This application provides a data processing method. Figure 2 This is a schematic diagram illustrating the implementation flow of a data processing method provided in an embodiment of this application, such as... Figure 2 As shown, the method may include the following steps S201 to S203:

[0038] Step S201: Divide the data block corresponding to the core into multiple sub-data blocks;

[0039] Here, the data block corresponding to the core refers to the data block allocated by the core among the multiple data blocks included in the data to be processed. The core is used to process the data to be processed. For example, the core performs calculations such as prefix sum and cumulative sum on the corresponding data block.

[0040] In some implementations, a data block is divided into multiple sub-data blocks of the same size.

[0041] In some implementations, the data block can be arbitrarily divided into multiple sub-data blocks according to the size of the data block corresponding to the core. For example, if the data block corresponding to the core includes 12 elements, the 12 elements can be divided into 2 sub-data blocks, each of which includes 6 elements. Alternatively, the 12 elements can be divided into 3 sub-data blocks, each of which includes 4 elements.

[0042] In some implementations, the data block can be divided into multiple sub-data blocks based on the size of the data block corresponding to the core, with the goal of minimizing the number of sub-data blocks and the number of elements in each sub-data block. This minimizes the mask matrix required for subsequent prefix sum calculations. For example, if the data block corresponding to the core contains 12 elements, dividing it into 2 sub-data blocks would result in each sub-data block containing 6 elements, making the mask matrix at least 6*6. Dividing it into 3 sub-data blocks would result in each sub-data block containing 4 elements, making the mask matrix at least 4*4. Therefore, dividing the 12 elements into 3 sub-data blocks is a better solution.

[0043] In some implementations, each sub-data block has the same block size. For example, the sub-data block size is set to 128. When calculating the prefix sum for data that is not a multiple of 128, the end of the data to be processed is first padded with 0. 0 does not affect the accumulation result.

[0044] Step S202: Multiply each sub-data block by the first mask matrix to obtain the first prefix sum matrix corresponding to each sub-data block. The first prefix sum matrix represents the prefix sum result in each sub-data block.

[0045] Here, the first mask matrix is ​​a matrix used to calculate the prefix sum of each element in each sub-data block.

[0046] In some implementations, the size of the first mask matrix is ​​related to the size of the sub-data blocks. The first mask matrix can be an upper triangular matrix. Here, an upper triangular matrix refers to a square matrix (a matrix with the same number of rows and columns), where all elements below the main diagonal are zero, while the elements on and above the main diagonal can be any value (including zero). In the embodiments of this application, the elements on and above the main diagonal of the first mask matrix can all be 1.

[0047] In some implementations, for each sub-data block, each element in the sub-data block is multiplied by the first mask matrix to obtain the prefix sum result corresponding to the sub-data block; the prefix sum results corresponding to each sub-data block form the first prefix sum matrix.

[0048] For example, a sub-data block includes a first sub-data block, a second sub-data block, and a third sub-data block. The prefix sum result corresponding to the first sub-data block is the first row of the first prefix sum matrix, the prefix sum result corresponding to the second sub-data block is the second row of the first prefix sum matrix, and the prefix sum result corresponding to the third sub-data block is the third row of the first prefix sum matrix.

[0049] In some implementations, the prefix sum result corresponding to each sub-data block includes the cumulative sum from the first element to the last element of the sub-data block. This can be understood as the prefix sum result corresponding to the sub-data block including the first element, the sum between the first and second elements, the sum between the first, second, and third elements, and so on until the last element is accumulated. For example, if the sub-data block includes 4 elements (0.1, 0.1, 0.1, 0.1), and the first mask matrix is ​​an upper triangular matrix, the prefix sum result corresponding to the sub-data block (0.1, 0.1, 0.1, 0.1) is obtained by multiplying (0.1, 0.1, 0.1, 0.1) with the first mask matrix.

[0050] In some implementations, the elements in each sub-data block can be summed sequentially to obtain the prefix sum result within that sub-data block, thereby obtaining the first prefix sum matrix corresponding to each sub-data block.

[0051] In related technologies, the Blelloch parallel prefix sum algorithm requires the construction of a mask matrix of size n×n (n is on the order of tens of thousands in the case of large language models), which consumes a large amount of on-chip storage for TPU and significantly increases the matrix multiplication overhead. In the embodiments of this application, the data to be processed (such as long sequences) is divided into fixed-length data blocks in each core, and the prefix calculation within the block is completed through the first mask matrix to obtain the first prefix sum matrix. This makes the maximum size of the mask matrix only "number of blocks × number of blocks", which effectively reduces the mask storage and calculation cost, reduces the pressure of constructing and storing giant masks, and makes the calculation volume only related to the block dimension, thus solving the problem of excessively large mask matrix size in related technologies.

[0052] Step S203: Based on the second mask matrix and the first prefix sum matrix, determine the second prefix sum matrix corresponding to the data block; the second prefix sum matrix represents the prefix sum result within the core.

[0053] Here, the second mask matrix is ​​a matrix used to calculate the prefix sum between each sub-data block within the core.

[0054] In some implementations, the size of the second mask matrix is ​​related to the number of sub-data blocks. The second mask matrix can be a lower triangular matrix with zeros on the main diagonal. Here, a lower triangular matrix refers to a square matrix (a matrix with the same number of rows and columns), where all elements above the main diagonal are zero, while elements on and below the main diagonal can have any value (including zero). In this embodiment, the main diagonal of the second mask matrix is ​​also set to 0, and all elements below the main diagonal are zero.

[0055] In some implementations, the second prefix sum matrix includes the sum of each element in the data block and all the elements preceding that element.

[0056] In some implementations, the inter-block prefix sum results between each sub-data block are first determined using the last column of the second mask matrix and the first prefix sum matrix, and then the second prefix sum matrix is ​​determined based on the first prefix sum matrix and the inter-block prefix sum results between each sub-data block.

[0057] In some implementations, the second prefix sum matrix is ​​obtained by adding each element in the first prefix sum matrix to the last prefix sum value of the previous sub-data block corresponding to each element.

[0058] In related technologies, the Blech algorithm calculates prefix sums at different summation intervals (1, 2, 4, 8…) at each layer, requiring the construction of multiple sets of mask matrices to select elements for computation. This application's embodiment uses a fixed-form first and second mask matrix to handle intra-block, inter-block, and inter-core prefix sums, eliminating the need for additional masks and replacing the complex multi-layer mask design of the Blech algorithm. This reduces implementation complexity and mask management overhead, making the prefix sum calculation process more stable and controllable, and better suited to the matrix operation characteristics of TPUs.

[0059] In some implementations, the size and number of sub-data blocks in each core are consistent. Assuming the size of the sub-data blocks is 4 and the number is 3, the core can process the accumulation of 12 elements at a time. The three sub-data blocks first determine the prefix sum value (i.e., the first prefix sum matrix) corresponding to each sub-data block in parallel. However, the subsequent sub-data blocks also need to add the accumulated value of the previous sub-data blocks. Therefore, the last number after the accumulation within each sub-data block is taken to calculate the accumulated value between sub-data blocks.

[0060] In related technologies, when the core data block is not divided into sub-data blocks, the prefix sum calculation of 12 elements requires a triangular matrix of shape (12, 12), totaling 144 numbers. In this embodiment, the data block is divided into 3 sub-data blocks of size 4, so the prefix sum can be calculated using a matrix of shape (4, 4) and a matrix of shape (3, 3), requiring only 25 numbers. It can be seen that compared with related technologies, this embodiment can save a lot of on-chip memory.

[0061] In this embodiment, during the processing of a corresponding data block, any core in the data processing chip divides the data block into multiple sub-data blocks and multiplies each sub-data block by a first mask matrix to obtain a first prefix sum matrix that represents the prefix sum result within each sub-data block. Then, based on the second mask matrix and the first prefix sum matrix, the second prefix sum matrix corresponding to the data block of that core is determined, and the second prefix sum matrix represents the prefix sum result within that core. In this way, on the one hand, by dividing the data block corresponding to each core into multiple sub-data blocks, the size of the mask matrix required for calculating the prefix sum is significantly reduced, and the prefix sum calculation corresponding to the data block is decomposed into multiple simple matrix operations, thereby effectively reducing both mask storage costs and computational costs and complexity. On the other hand, the prefix sum calculation process within the core can be completed using only the first and second mask matrices, eliminating the need to construct multiple sets of mask matrices based on the different summation intervals of each layer to select elements for calculation, reducing the number of mask matrices, thereby reducing the complexity of the prefix sum calculation and mask management overhead, making the prefix sum calculation process more stable and controllable.

[0062] In some embodiments, step S203 may include steps S211 to S212:

[0063] Step S211: Multiply the second mask matrix by the last column of the first prefix sum matrix to obtain the third prefix sum matrix; the third prefix sum matrix represents the inter-block prefix sum results between each sub-data block;

[0064] In some implementations, the second mask matrix is ​​a lower triangular matrix with 0s above and 1s below the diagonal. The last column of the first prefix sum matrix is ​​the last prefix sum value of each sub-data block, that is, the sum of all elements of each sub-data block. The second mask matrix is ​​multiplied by the last column of the first prefix sum matrix to obtain the inter-block prefix sum result between each sub-data block, that is, the third prefix sum matrix.

[0065] In some implementations, when the data block corresponding to the core is divided into 3 sub-data blocks, the first row of the third prefix sum matrix is ​​0, the second row is the sum of the elements of the first sub-data block, and the third row is the sum of the elements of the first and second sub-data blocks.

[0066] In some implementations, the last column of the first prefix sum matrix is ​​the last prefix sum value of each sub-data block. By simply passing the last prefix sum value of each sub-data block and combining it with the second mask matrix, the inter-block prefix sum result between the sub-data blocks, i.e., the third prefix sum matrix, can be calculated. For example, the last column of the first prefix sum matrix might be (0.4, 1.0, 0.8). T Where 0.4 is the sum of the elements in the first sub-data block, 1.0 is the sum of the elements in the second sub-data block, and 0.8 is the sum of the elements in the third sub-data block; multiplying the second mask matrix by the last column of the first prefix sum matrix yields (0.0, 0.4, 1.4). T , where 0.4 is the sum of all elements of the first sub-data block and all elements of the 0th sub-data block (i.e., 0.0), and 1.4 is the sum of all elements of the first sub-data block and all elements of the 2nd sub-data block.

[0067] In some implementations, the third prefix sum matrix includes the prefix sum results between the sub-data blocks before the last sub-data block. For example, when the data block is divided into 3 sub-data blocks, the third prefix sum matrix includes the sum of the elements in the first sub-data block, and the sum of the elements in the first sub-data block and the elements in the second sub-data block.

[0068] In some implementations, a third prefix sum matrix can be constructed using the last column of the first prefix sum matrix.

[0069] Step S212: Add the first prefix sum matrix and the third prefix sum matrix to obtain the second prefix sum matrix corresponding to the data block.

[0070] In some implementations, adding the first and third prefix sum matrices is equivalent to adding the sum of the elements of the previous sub-data block to each element in the first prefix sum matrix. For example, the first row of the first prefix sum matrix is ​​the prefix sum result corresponding to the first sub-data block, and adding 0 to the first row of the first prefix sum matrix represents the sum of each element in the first sub-data block with all elements preceding it; the second row of the first prefix sum matrix is ​​the prefix sum result corresponding to the second sub-data block, and adding the second row of the first prefix sum matrix to the sum of the elements in the first sub-data block represents the sum of each element in the second sub-data block with all elements preceding it; the third row of the first prefix sum matrix is ​​the prefix sum result corresponding to the third sub-data block, and adding the third row of the first prefix sum matrix to the sum of the elements in the first and second sub-data blocks represents the sum of each element in the third sub-data block with all elements preceding it.

[0071] In some implementations, the second prefix sum matrix includes the sum of each element in the data to be processed and all the elements preceding it. For example, the second prefix sum matrix includes the first element, the sum of the first element and the second element, the sum of the first element, the second element and the third element, and so on, until the last element of the data to be processed is added.

[0072] In this embodiment, the second mask matrix is ​​multiplied by the last column of the first prefix sum matrix to obtain a third prefix sum matrix representing the inter-block prefix sum results between each sub-data block; the first and third prefix sum matrices are then added to obtain the second prefix sum matrix corresponding to each data block. Thus, the inter-block prefix sum results between each sub-data block can be determined using only the last column of the first prefix sum matrix and the second mask matrix, thereby simplifying the calculation of the prefix sum results between elements of the data block within the core.

[0073] In some embodiments, the above method may further include the following step S221:

[0074] Step S221: After determining the second prefix sum matrix corresponding to the data block, write the last prefix sum value in the second prefix sum matrix as the kernel prefix sum value corresponding to the kernel into the data buffer.

[0075] Here, the last prefix sum value in the second prefix sum matrix represents the sum of all elements in the data block corresponding to that core. For example, if a data block contains 12 elements, the last prefix sum value in the second prefix sum matrix is ​​the sum of all 12 elements.

[0076] The kernel prefix sum represents the sum of all elements in the data block corresponding to the kernel.

[0077] In some implementations, each core writes the last prefix sum value from the second prefix sum matrix calculated by that core into a data buffer to pass the core's intra-core prefix sum value to other cores for calculating inter-core prefix sums. In some implementations, the data buffer can be a secondary cache (L2 memory).

[0078] In related technologies, the doubling interval summation method of the Blech algorithm will perform cross-core communication multiple times (log n level) in the later stage, which increases the communication cost between multiple cores. In the embodiment of this application, the prefix sum value of each core is written into the data buffer, which can realize the synchronization of prefix sum values ​​between multiple cores. By using a single inter-core synchronization, the communication cost is reduced to a one-time fixed delay.

[0079] In this embodiment, after determining the second prefix sum matrix corresponding to the data block, the last prefix sum value in the second prefix sum matrix is ​​written into the data buffer as the core's intra-core prefix sum value. Thus, by writing the intra-core prefix sum value corresponding to each core into the data buffer, prefix sum value synchronization between multiple cores is achieved using only a single inter-core synchronization, reducing the communication cost between multiple cores to a one-time fixed latency.

[0080] In some embodiments, the first mask matrix is ​​an upper triangular matrix, and the size of the first mask matrix is ​​determined based on the size of the sub-data blocks; the second mask matrix is ​​a lower triangular matrix, and the size of the second mask matrix is ​​determined based on the number of sub-data blocks.

[0081] In some implementations, the first mask matrix is ​​set as an upper triangular matrix with 1s at and above the diagonal and 0s below the diagonal, so that when each sub-data block is multiplied by the first mask matrix, the sum of each element in each sub-data block can be obtained, that is, the first prefix sum matrix corresponding to each sub-data block is determined.

[0082] In some implementations, the size of the first mask matrix can be the same as the maximum size of the sub-data block. For example, if the sub-data block is a 4*1 matrix, the size of the first mask matrix can be set to 4*4.

[0083] In some implementations, the size of the first mask matrix is ​​at least equal to, or may be greater than, the maximum size of the sub-data block.

[0084] In some implementations, the second mask matrix is ​​set as a lower triangular matrix with 0s at and above the diagonal and 1s below the diagonal. This allows the second prefix sum matrix corresponding to the data blocks within the core to be determined based on the second mask matrix and the first prefix sum matrix.

[0085] In some implementations, the size of the second mask matrix can be determined based on the number of sub-data blocks. For example, if the number of sub-data blocks is 3, the size of the second mask matrix can be set to 3*3.

[0086] In some implementations, the size of the second mask matrix is ​​at least equal to, or may be greater than, the number of sub-data blocks.

[0087] In some implementations, the first mask matrix and the second mask matrix in each core are the same.

[0088] In this embodiment, a fixed mask matrix is ​​set in each core based on the divided sub-data blocks, so that the prefix sum in each core can be calculated. This not only reduces the number and size of the mask matrix, but also reduces the complexity of the prefix sum calculation.

[0089] In some embodiments, the above method may further include the following steps S231 to S232:

[0090] Step S231: Obtain the intra-core prefix sum value corresponding to each core in the data processing chip from the data buffer;

[0091] In some implementations, each core reads the Hanoi prefix sum value corresponding to each core from the data buffer to synchronize the prefix sum calculation results of other cores.

[0092] In some implementations, the intra-core prefix sum values ​​read by each core can be a vector with a length equal to the number of cores. For example, in a data processing chip, there are 8 cores participating in the prefix sum calculation. Each core reads the intra-core prefix sum values ​​corresponding to the 8 cores from the data buffer to form a vector with a length of 8.

[0093] Step S232: Based on the second prefix sum matrix corresponding to the data block, the prefix sum value within each core, and the third mask matrix, determine the fourth prefix sum matrix. The fourth prefix sum matrix represents the total prefix sum result between cores before the core.

[0094] Here, the third mask matrix is ​​a matrix used to calculate the prefix sum among multiple cores.

[0095] In some implementations, the size of the third mask matrix is ​​related to the number of cores involved in the prefix sum calculation. For example, when there are 3 cores involved in the prefix sum calculation, the size of the third mask matrix can be 3*3.

[0096] In some implementations, the third mask matrix can be a lower triangular matrix in which the second mask matrix is ​​a matrix with 0s at and above the main diagonal and 1s below the main diagonal.

[0097] In some implementations, inter-core prefix propagation is completed by multiplying the third mask matrix with the vector formed by the prefix sums of each core to obtain the inter-core prefix sum result. The inter-core prefix sum result is then passed to each row of the second prefix sum matrix to obtain the total prefix sum result among multiple cores, i.e., the fourth prefix sum matrix.

[0098] In some implementations, the fourth prefix sum matrix includes the sum of each element in the data block corresponding to the core and all elements preceding that element.

[0099] In some implementations, the intra-kernel prefix sums can be accumulated to determine the inter-kernel prefix sum result, and then the fourth prefix sum matrix can be determined based on the inter-kernel prefix sum result.

[0100] In related technologies, when using the Blellach algorithm to calculate the total prefix sum between cores, distant elements are accessed in the later stages of the summation, resulting in a large number of cross-core synchronizations in a multi-core environment, affecting throughput. In this embodiment, the intra-core prefix sum value of each core is obtained from the data buffer, and the inter-core prefix sum value between cores is obtained using a third mask matrix, thereby simplifying the number of communication steps between cores and improving the calculation efficiency of the prefix sum.

[0101] In this embodiment, the intra-core prefix sum value corresponding to each core is obtained from the data buffer; based on the second prefix sum matrix corresponding to the data block, the intra-core prefix sum value, and the third mask matrix, a fourth prefix sum matrix is ​​determined. The fourth prefix sum matrix represents the total prefix sum result between cores preceding each other. Thus, by obtaining the intra-core prefix sum value corresponding to each core from the data buffer, prefix sum synchronization between cores is achieved by only transmitting the intra-core prefix sum value of each core. The total prefix sum result between cores preceding each other is then determined through the third mask matrix, simplifying cross-core communication from multi-stage synchronization in related technologies to a one-time matrix operation, reducing the communication cost between cores.

[0102] In some embodiments, step S232 may include steps S241 to S243:

[0103] Step S241: Multiply the third mask matrix by the prefix sum value within each core to obtain the fifth prefix sum matrix; the fifth prefix sum matrix represents the inter-core prefix sum result between each core.

[0104] Here, the fifth prefix sum matrix can include the prefix sum results between each core. For example, when there are 3 cores involved in the prefix sum calculation, the first row of the fifth prefix sum matrix is ​​the intra-core prefix sum value (0) corresponding to the 0th core, the second row is the result of adding the intra-core prefix sum value of the 0th core to the intra-core prefix sum value of the first core, and the third row is the result of accumulating the prefix sum values ​​from the 0th core to the second core.

[0105] In some implementations, the third mask matrix can be a lower triangular matrix with 0s at and above the main diagonal and 1s below the main diagonal. Multiplying the third mask matrix by the vector formed by the prefix sums of each core yields the inter-core prefix sums.

[0106] In some implementations, the third mask matrix in each core is the same.

[0107] Step S242: Based on the position of the data block in the data to be processed, determine the total prefix sum value of each core before the core from the fifth prefix sum matrix;

[0108] In some implementations, each core can calculate the prefix sum result among all cores. Since each core processes data blocks in the order of the data to be processed, the core corresponding to the data block before the data block processed by that core can be determined based on the position of the data block in the data to be processed, thereby determining the total prefix sum value corresponding to the cores before that core.

[0109] For example, in a data processing chip, there are 3 cores involved in prefix sum calculation. The data block being processed by the current core is the second data block in the data to be processed. After the current core calculates the prefix sum results between all cores, it selects the inter-core prefix sum results of all cores before the current core (i.e., the sum of the intra-core prefix sum values ​​of the 0th core and the 1st core) from the prefix sum results between all cores based on the position of the data block being processed by the current core.

[0110] For example, if the data block being processed by the current core is the third data block in the data to be processed, after the current core calculates the prefix sum result between all cores, it selects the inter-core prefix sum result of all cores before the current core (i.e., the result of accumulating the intra-core prefix sum values ​​from the 0th core to the 2nd core) based on the position of the data block being processed by the current core.

[0111] In some implementations, the current core may only read the kernel prefix sum of all cores preceding it, but each core needs to be configured with a triangular matrix of a different size. For example, the first core does not need to know the accumulated result of the subsequent cores and therefore does not need a triangular matrix that can accumulate the kernel prefix sum of the subsequent cores. The second core only needs to know the accumulated result of the first core and can read only one value, so it does not need a triangular matrix. The third core needs the accumulated result of the first two cores and needs to be configured with a (2,2) shaped triangular matrix. The fourth core needs a (3,3) shaped triangular matrix, which increases the complexity of configuring the third mask matrix.

[0112] Step S243: Add the second prefix sum matrix corresponding to the data block to the total prefix sum value to obtain the fourth prefix sum matrix.

[0113] In some implementations, the fourth prefix sum matrix, which includes the prefix sums calculated by the current core, can be determined by comparing each element in the second prefix sum matrix with the total prefix sum value corresponding to the core.

[0114] In this embodiment, the third mask matrix is ​​multiplied by the prefix sum value within each core to obtain the fifth prefix sum matrix. This fifth prefix sum matrix represents the inter-core prefix sum result between cores. Based on the position of the data block in the data to be processed, the total prefix sum value of each core preceding the current core is determined from the fifth prefix sum matrix. The second prefix sum matrix corresponding to the data block is added to the total prefix sum value to obtain the fourth prefix sum matrix. Thus, by calculating the inter-core prefix sum result between cores using the third mask matrix and selecting an appropriate total prefix sum value to add to the second prefix sum value corresponding to the current core, the total prefix sum result between the current core and the cores preceding it is obtained, reducing the computational complexity of the prefix sum.

[0115] In some embodiments, the data blocks corresponding to the core are allocated sequentially based on the location of the core in the data processing chip, and the size of the data blocks corresponding to the core is determined based on the number of cores included in the data processing chip and the size of the data to be processed.

[0116] This application provides a data processing chip, such as... Figure 1 As shown, any core included in the data processing chip is used for: processing the corresponding data block in the data to be processed; dividing the data block corresponding to the core into multiple sub-data blocks; multiplying each sub-data block by a first mask matrix to obtain a first prefix sum matrix corresponding to each sub-data block, the first prefix sum matrix representing the prefix sum result within each sub-data block; determining the second prefix sum matrix corresponding to the data block based on the second mask matrix and the first prefix sum matrix; the second prefix sum matrix representing the prefix sum result within the core.

[0117] In this embodiment, during the processing of a corresponding data block, any core in the data processing chip divides the data block into multiple sub-data blocks and multiplies each sub-data block by a first mask matrix to obtain a first prefix sum matrix that represents the prefix sum result within each sub-data block. Then, based on the second mask matrix and the first prefix sum matrix, the second prefix sum matrix corresponding to the data block of that core is determined, and the second prefix sum matrix represents the prefix sum result within that core. In this way, on the one hand, by dividing the data block corresponding to each core into multiple sub-data blocks, the size of the mask matrix required for calculating the prefix sum is significantly reduced, and the prefix sum calculation corresponding to the data block is decomposed into multiple simple matrix operations, thereby effectively reducing both mask storage costs and computational costs and complexity. On the other hand, the prefix sum calculation process within the core can be completed using only the first and second mask matrices, eliminating the need to construct multiple sets of mask matrices based on the different summation intervals of each layer to select elements for calculation, reducing the number of mask matrices, thereby reducing the complexity of the prefix sum calculation and mask management overhead, making the prefix sum calculation process more stable and controllable.

[0118] In some embodiments, the core is further configured to: multiply the second mask matrix by the last column of the first prefix sum matrix to obtain a third prefix sum matrix; the third prefix sum matrix represents the inter-block prefix sum result between each sub-data block; and add the first prefix sum matrix and the third prefix sum matrix to obtain the second prefix sum matrix corresponding to the data block.

[0119] In this embodiment, the second mask matrix is ​​multiplied by the last column of the first prefix sum matrix to obtain a third prefix sum matrix representing the inter-block prefix sum results between each sub-data block; the first and third prefix sum matrices are then added to obtain the second prefix sum matrix corresponding to each data block. Thus, the inter-block prefix sum results between each sub-data block can be determined using only the last column of the first prefix sum matrix and the second mask matrix, thereby simplifying the calculation of the prefix sum results between elements of the data block within the core.

[0120] In some embodiments, the core is further configured to: after determining the second prefix sum matrix corresponding to the data block, write the last prefix sum value in the second prefix sum matrix as the kernel-internal prefix sum value corresponding to the core into the data buffer.

[0121] In this embodiment, after determining the second prefix sum matrix corresponding to the data block, the last prefix sum value in the second prefix sum matrix is ​​written into the data buffer as the core's intra-core prefix sum value. Thus, by writing the intra-core prefix sum value corresponding to each core into the data buffer, prefix sum value synchronization between multiple cores is achieved using only a single inter-core synchronization, reducing the communication cost between multiple cores to a one-time fixed latency.

[0122] In some embodiments, the first mask matrix is ​​an upper triangular matrix, and the size of the first mask matrix is ​​determined based on the size of the sub-data blocks; the second mask matrix is ​​a lower triangular matrix, and the size of the second mask matrix is ​​determined based on the number of sub-data blocks.

[0123] In this embodiment, a fixed mask matrix is ​​set in each core based on the divided sub-data blocks, so that the prefix sum in each core can be calculated. This not only reduces the number and size of the mask matrix, but also reduces the complexity of the prefix sum calculation.

[0124] In some embodiments, the core is further configured to: obtain the intra-core prefix sum value corresponding to each core in the data processing chip from the data buffer; and determine a fourth prefix sum matrix based on the second prefix sum matrix corresponding to the data block, the intra-core prefix sum value, and the third mask matrix, wherein the fourth prefix sum matrix represents the total prefix sum result between the cores preceding the core.

[0125] In this embodiment, the intra-core prefix sum value corresponding to each core is obtained from the data buffer; based on the second prefix sum matrix corresponding to the data block, the intra-core prefix sum value, and the third mask matrix, a fourth prefix sum matrix is ​​determined. The fourth prefix sum matrix represents the total prefix sum result between cores preceding each other. Thus, by obtaining the intra-core prefix sum value corresponding to each core from the data buffer, prefix sum synchronization between cores is achieved by only transmitting the intra-core prefix sum value of each core. The total prefix sum result between cores preceding each other is then determined through the third mask matrix, simplifying cross-core communication from multi-stage synchronization in related technologies to a one-time matrix operation, reducing the communication cost between cores.

[0126] In some embodiments, the core is further configured to: multiply the third mask matrix by the prefix sum value within each core to obtain a fifth prefix sum matrix; the fifth prefix sum matrix represents the inter-core prefix sum result between cores; determine the total prefix sum value of each core before the core from the fifth prefix sum matrix based on the position of the data block in the data to be processed; and add the second prefix sum matrix corresponding to the data block to the total prefix sum value to obtain a fourth prefix sum matrix.

[0127] In this embodiment, the third mask matrix is ​​multiplied by the prefix sum value within each core to obtain the fifth prefix sum matrix. This fifth prefix sum matrix represents the inter-core prefix sum result between cores. Based on the position of the data block in the data to be processed, the total prefix sum value of each core preceding the current core is determined from the fifth prefix sum matrix. The second prefix sum matrix corresponding to the data block is added to the total prefix sum value to obtain the fourth prefix sum matrix. Thus, by calculating the inter-core prefix sum result between cores using the third mask matrix and selecting an appropriate total prefix sum value to add to the second prefix sum value corresponding to the current core, the total prefix sum result between the current core and the cores preceding it is obtained, reducing the computational complexity of the prefix sum.

[0128] In some embodiments, the data blocks corresponding to the core are allocated sequentially based on the location of the core in the data processing chip, and the size of the data blocks corresponding to the core is determined based on the number of cores included in the data processing chip and the size of the data to be processed.

[0129] The following will describe an exemplary application of the embodiments of this application in a real-world application scenario.

[0130] Prefix sums are a fundamental operator widely used in various deep learning models, especially large language models such as qwen3next and mamba2. Serial prefix sums in related technologies require adding n elements sequentially, resulting in a time complexity of O(n), which is difficult to achieve ideal performance in large-scale data processing scenarios. To explore more efficient prefix sum calculation methods under parallel computing, Bllelloch proposed the classic Bllelloch tree-based prefix sum calculation method. By constructing a bottom-up binary tree, nodes with even indices (starting from 1) converge the values ​​of preceding nodes (upward reduction), and then the prefix sum is propagated to child nodes through a top-down binary tree (downward distribution). This reduces the time complexity of element-wise prefix sum calculation from O(n) to O(log n), becoming a fundamental parallel primitive widely adopted in subsequent parallel computing hardware.

[0131] The Blelloch algorithm is essentially a prefix sum calculation algorithm for SIMT (Single Instruction Multiple Threads, where multiple threads execute the same instruction simultaneously but process different data, threads can access memory independently, and support branching and dynamic scheduling) hardware. It relies on hardware with a large number of independent threads, and multiple threads can simultaneously and independently retrieve numbers by index and perform binary addition. Since the tensor core architectures such as TPU and NPU are based on SIMD (Single Instruction Multiple Data, where multiple data are processed by a single instruction, data must be stored continuously, operation types are fixed, and branching capabilities are weak (requiring optimization through masking or loop expansion)), compared to SIMT architecture hardware, the data processed by a single instruction in SIMD hardware is generally continuous data, making it difficult to meet the requirement of processing interval data by a single instruction in the Blelloch algorithm. Forcing implementation will lead to: (1) memory access bottlenecks caused by interval data retrieval; (2) the need for multiple instruction combinations to achieve the corresponding function (due to the SIMD architecture, the instruction granularity is large, and the configuration and issuance overhead of a single instruction is large), resulting in unsatisfactory performance.

[0132] In related technologies, the specific problems encountered when programming with a TPU according to the Blelloch algorithm are as follows:

[0133] (1) The size of the mask matrix is ​​too large:

[0134] TPUs compute by executing the same instruction in parallel across multiple lanes. However, the doubling interval (binary tree) summation of the Blech algorithm (adding the first and second numbers, then the third and fourth numbers, and then the second and fourth numbers again) does not conform to the instruction characteristics of TPUs. Generally, a mask matrix is ​​used to implement the addition of specified elements. When a TPU core needs to perform interval addition of n elements, the size of the first-level mask matrix becomes n×n (n is on the order of tens of thousands in large language model scenarios), making the storage cost of the mask matrix extremely high. In addition, the increase in the size of the mask matrix also leads to a significant increase in the computational overhead of matrix multiplication.

[0135] (2) The number of masks increases continuously with the number of elements:

[0136] The Blelloch algorithm uses different intervals at each layer. The first layer adds adjacent elements (with a gap of 1), the second layer adds elements with a gap of 2, the third layer adds elements with a gap of 4, and so on. In large language models, the number of elements is often related to the vocabulary size, frequently reaching tens of thousands, requiring multiple layers of interval-based summation. This means constructing a large number of mask matrices with different patterns to achieve selective addition at different levels, significantly increasing the number and management costs of the mask matrices, placing a heavy burden on both developers and users.

[0137] (3) The parallel computing process of the Blelloch algorithm is incompatible with the TPU communication method:

[0138] To fully utilize the computing power of the TPU, when processing prefix sums of a large number of elements, the data is typically split across different cores for parallel computation. However, the Bllelloch algorithm's computation rules stipulate that the element summation interval increases exponentially, and the elements to be summed in pairs may be distributed across different cores, leading to a large number of cross-core synchronizations during the computation process. For example, as the summation interval increases, cross-core communication may occur between core 7 and core 5, or between core 7 and core 3, impacting the overall performance of the TPU when calculating prefix sums.

[0139] Based on the above description, this application provides a prefix sum calculation method based on the Tile (block) concept. This method solves the affinity problem between the algorithm and SIMD architecture hardware by dividing the entire input sequence into blocks, performing local prefix sum calculations first, and then obtaining the global prefix sum. The instructions used in the calculation method are mainly large-granularity matrix multiplication instructions and continuous data transport instructions (SIMD-compatible instructions). Prefix sum calculation is efficiently achieved with a small number of SIMD instructions.

[0140] In this embodiment, a block-based prefix sum calculation method is designed. The size of the sub-data blocks in each TPU core is set (e.g., block is 128). The data in each core is then divided into multiple sub-data blocks according to the sub-data block size, so that the maximum size of the mask matrix is ​​reduced from the original n×n to (n / block)×(n / block). In the scenario of large language models, n is on the order of tens of thousands, while n / block can be reduced to the order of hundreds, effectively reducing the mask storage and calculation costs.

[0141] In this embodiment, a mask matrix for calculating the prefix sum within a sub-data block, a mask matrix for calculating the prefix sum between sub-data blocks, and a mask matrix for calculating the prefix sum between cores are designed. Each core normally retains these three mask matrices, eliminating the need to generate multiple mask matrices as the layer increases.

[0142] In this embodiment, all TPU cores synchronize results between cores by calculating the prefix sum matrix between cores, requiring only one write and one read operation, without the need for multiple cross-core synchronizations using the Blech algorithm.

[0143] In this embodiment, a new prefix sum calculation method is proposed for the TPU architecture. This method achieves the O(log n) time complexity of the Blelloch algorithm while reducing the size and number of mask matrices and the number of inter-core communications. Specifically:

[0144] (1) A framework for block-based prefix sum computation using upper / lower triangular matrices (mask matrices) is proposed;

[0145] The Blelloch algorithm's parallel prefix sum algorithm requires constructing a mask matrix of size n×n (n is on the order of tens of thousands in large language model scenarios), which consumes a large amount of on-chip storage for TPUs and significantly increases matrix multiplication overhead. This application's embodiment divides long sequences into fixed-length data blocks within each TPU core, and uses an upper triangular matrix (mask matrix) to complete prefix calculation within each sub-block. This ensures that the maximum mask matrix size is always only "number of blocks × number of blocks" or "block size × block size," reducing the pressure of constructing and storing giant masks. The computational load is only related to the dimension of the sub-blocks, solving the problem of excessively large mask matrix sizes in related technologies.

[0146] (2) Use three triangular mask matrices instead of multiple mask matrices;

[0147] The Blelloch algorithm calculates prefix sums at different intervals (1, 2, 4, 8…) at each layer, requiring the construction of multiple sets of mask matrices to select elements for computation. This application proposes using fixed-form upper / lower triangular matrices to handle intra-block, inter-block, and inter-kernel prefix sums, eliminating the need for additional mask matrices. This replaces the complex multi-layer mask design of the Blelloch algorithm, reducing implementation complexity and mask management overhead. It makes the prefix sum calculation process more stable and controllable, and is better suited to the matrix operation characteristics of TPUs.

[0148] (3) A mechanism is proposed for independent computing within the core and one-time synchronization between cores.

[0149] The doubling interval summation characteristic of the Blelloch algorithm leads to accessing distant elements in the later stages of summation, resulting in a large amount of cross-core synchronization in a multi-core TPU environment, severely impacting throughput. This application's embodiment constructs a lightweight inter-core merging mechanism: "completely calculate the block-level prefix sum within the core → transmit only the block tail result → complete inter-core prefix propagation with a one-time lower triangular matrix calculation." This simplifies cross-core communication from multi-stage synchronization to a single matrix operation, significantly reducing multi-core communication costs and improving the computational efficiency of prefix summation on the TPU.

[0150] The overall architecture of the prefix sum calculation method based on the block-based approach provided in this application embodiment may include:

[0151] (1) Data Shard: Divide the elements to be accumulated into 8 parts according to the number of TPU cores and distribute them to the cores of each TPU; within each core, divide the data into several blocks according to the set sub-data block size to obtain an element matrix of (num_blocks_per_core, block_size).

[0152] (2) Intra-Block Prefix Sum Calculation: Within each TPU core, the element matrix of (num_blocks_per_core, block_size) is multiplied by the upper triangular matrix of (block_size, block_size) to obtain the cumulative result of (num_blocks_per_core, block_size). Each row of the result is the prefix sum of a block.

[0153] (3) Calculation and propagation of inter-block prefix sums: Take the last value of the prefix sum corresponding to each sub-data block to form the prefix sum vector within the sub-data block (num_blocks_per_core, 1), use the lower triangular matrix of (num_blocks_per_core, num_blocks_per_core) (with the diagonal set to 0) to calculate the inter-sub-data block prefix sum, and then add the prefix sum within the sub-data block and the prefix sum between the sub-data blocks to obtain the prefix sum within the core.

[0154] (4) Inter-Core Prefix Accumulation: Each TPU core writes its last prefix value to L2 Memory. Assuming there are 8 TPU cores participating in the prefix accumulation calculation, each TPU core reads 8 intra-core prefix sum values ​​from L2 Memory, forming a vector of length 8. This vector is then multiplied by an 8×8 lower triangular matrix (with the diagonals set to 0) to complete the inter-core prefix accumulation. Finally, the intra-core prefix sum is added to the inter-core prefix sum to obtain the final prefix sum result for all cores of all TPUs.

[0155] In this embodiment of the application, the prefix sum calculation method based on the block-based approach provided may include the following steps S1 to S3:

[0156] Step S1: Calculate the prefix sum within the block in parallel;

[0157] Within the core of each TPU, the prefix sum of each block within the TPU core is directly calculated using an upper triangular matrix. For example... Figure 3 As shown, assuming the 12 elements in the core are divided into 3 sub-data blocks, each sub-data block includes 4 elements, multiplying each sub-data block by the upper triangular matrix, we get the prefix sum of (0.1, 0.1, 0.1, 0.1) as (0.1, 0.2, 0.3, 0.4), the prefix sum of (0.2, 0.2, 0.3, 0.3) as (0.2, 0.4, 0.7, 1.0), and the prefix sum of (0.2, 0.3, 0.1, 0.2) as (0.2, 0.5, 0.6, 0.8).

[0158] Step S2: Calculation and propagation of inter-block prefix sums;

[0159] Within each TPU core, such as Figure 4 As shown, the last column of the prefix sum within the sub-block is taken, and then the vector is multiplied by the lower triangular matrix (with the diagonal set to 0) to obtain the inter-block prefix sum.

[0160] like Figure 5As shown, the prefix sum result matrix within the block is added to the prefix sum result matrix between blocks to obtain the prefix sum result within the entire TPU core.

[0161] Step S3: Calculation and propagation of inter-kernel prefix sums.

[0162] Each TPU core writes the last prefix sum value within its own core to L2 Memory. All cores then read the prefix sum value of each core from L2 Memory and use a lower triangular matrix (with the diagonals set to 0) to calculate the inter-core prefix sum.

[0163] Assume three TPU cores participate in the prefix sum calculation, and that the last prefix sums of the three TPU cores are 1.0, 0.8, and 2.2, respectively. The calculation process for the inter-core prefix sum is as follows: Figure 6 As shown.

[0164] Add the prefix sum within the core to the prefix sum of all preceding cores to obtain the final prefix sum result within this TPU core (the final prefix sum result within this TPU core refers to the prefix sum within the current core plus the prefix sum of all preceding cores), such as... Figure 7 As shown, assuming the prefix sum of the first two TPU cores is 1.8, then add 1.8 to all elements of the third TPU core to obtain the final prefix sum of the third TPU core.

[0165] In this embodiment, the block-based prefix sum calculation method is applicable to TPU architecture, especially suitable for scenarios where large language models are deployed on TPUs and probability accumulation (Top-P) needs to be calculated. If a user only wants to select the next token output by the large language model from all (lexical) tokens whose probability accumulation reaches Top-P, the probability values ​​of the vocabulary size are first sorted in ascending order. Then, the prefix sum calculation process proposed in this embodiment is used to obtain the probability accumulation value. Finally, tokens corresponding to indices with probability accumulation values ​​less than (1-Top-P) are removed from the candidate tokens to complete the Top-P post-processing logic.

[0166] The prefix sum calculation algorithm provided in this application has a strong affinity with the SIMD architecture and is applicable to hardware computing platforms with different SIMD architectures, giving full play to the computing power of the SIMD computing platform.

[0167] In summary, in related technologies, the Bllelloch algorithm requires the construction of an n×n mask matrix. The prefix sum calculation method provided in this application reduces the size of the maximum mask matrix from O(n²) to O(block_num²). Taking the 129280-word vocabulary of the large DeepSeek-V3 model as an example, assuming each TPU core is allocated 16160 words, the mask matrix size is reduced from 16160*16160 to 127*127. This reduces the size of the mask matrix and decreases storage and computational costs.

[0168] In related technologies, the prefix sum calculation in the Blelloch algorithm requires generating multiple sets of mask matrices as the interval (1, 2, 4, 8...) increases. However, the embodiments of this application replace the multi-level mask matrices with fixed triangular matrices, reducing the number of mask matrices to a constant level.

[0169] In related technologies, the doubling interval summation method of the Blelloch algorithm involves multiple (log n level) cross-core communications in the later stages, which increases the cost of inter-core communication. However, the embodiment of this application uses a single inter-core synchronization, reducing the communication cost to a one-time fixed delay.

[0170] This application provides a computer device, such as... Figure 8 As shown, the computer device 800 includes a data processing chip 100.

[0171] It should be understood that the phrase "one embodiment" or "an embodiment" throughout the specification means that a specific feature, structure, or characteristic related to the embodiment is included in at least one embodiment of this application. Therefore, "in one embodiment" or "in an embodiment" appearing throughout the specification does not necessarily refer to the same embodiment. Furthermore, these specific features, structures, or characteristics can be combined in any suitable manner in one or more embodiments. It should be understood that in the various embodiments of this application, the sequence numbers of the above steps / processes do not imply a sequential order of execution; the execution order of each step / process should be determined by its function and internal logic, and should not constitute any limitation on the implementation process of the embodiments of this application. The sequence numbers of the above embodiments of this application are merely descriptive and do not represent the superiority or inferiority of the embodiments.

[0172] It should be noted that, in this document, the terms "comprising," "including," or any other variations thereof are intended to cover non-exclusive inclusion, such that a process, method, article, or apparatus that comprises a list of elements includes not only those elements but also other elements not expressly listed, or elements inherent to such a process, method, article, or apparatus. Unless otherwise specified, an element defined by the phrase "comprising one..." does not exclude the presence of other identical elements in the process, method, article, or apparatus that includes that element.

[0173] In the several embodiments provided in this application, it should be understood that the disclosed devices and methods can be implemented in other ways. The device embodiments described above are merely illustrative. For example, the division of units is only a logical functional division, and in actual implementation, there may be other division methods, such as: multiple units or components can be combined, or integrated into another system, or some features can be ignored or not executed. In addition, the coupling, direct coupling, or communication connection between the various components shown or discussed can be through some interfaces, and the indirect coupling or communication connection between devices or units can be electrical, mechanical, or other forms.

[0174] The units described above as separate components may or may not be physically separate, and the components shown as units may or may not be physical units; they may be located in one place or distributed across multiple network units; some or all of the units may be selected to achieve the purpose of this embodiment according to actual needs. Furthermore, the functional units in the embodiments of this application may all be integrated into one processing unit, or each unit may be a separate unit, or two or more units may be integrated into one unit; the integrated unit may be implemented in hardware or in a combination of hardware and software functional units.

[0175] Those skilled in the art will understand that all or part of the steps of the above method embodiments can be implemented by hardware related to program instructions. The aforementioned program can be stored in a computer-readable storage medium. When the program is executed, it performs the steps of the above method embodiments. The aforementioned storage medium includes various media that can store program code, such as mobile storage devices, read-only memory (ROM), magnetic disks, or optical disks.

[0176] Alternatively, if the integrated units described above are implemented as software functional modules and sold or used as independent products, they can also be stored in a computer-readable storage medium. Based on this understanding, the technical solution of this application, in essence or the part that contributes to related technologies, can be embodied in the form of a software product. This computer software product is stored in a storage medium and includes several instructions to cause a computer device (which may be a personal computer, server, or network device, etc.) to execute all or part of the methods of the various embodiments of this application. The aforementioned storage medium includes various media capable of storing program code, such as mobile storage devices, ROMs, magnetic disks, or optical disks.

[0177] The above are merely embodiments of this application, but the scope of protection of this application is not limited thereto. Any variations or substitutions that can be easily conceived by those skilled in the art within the scope of the technology disclosed in this application should be included within the scope of protection of this application.

Claims

1. A data processing method, characterized in that, The method is applied to any core included in a data processing chip, the core being used to process corresponding data blocks in the data to be processed, the method comprising: The data block corresponding to the core is divided into multiple sub-data blocks; Each of the sub-data blocks is multiplied by the first mask matrix to obtain the first prefix sum matrix corresponding to each sub-data block. The first prefix sum matrix represents the prefix sum result in each sub-data block. Based on the second mask matrix and the first prefix sum matrix, the second prefix sum matrix corresponding to the data block is determined; the second prefix sum matrix represents the prefix sum result within the core.

2. The method according to claim 1, characterized in that, The step of determining the second prefix sum matrix corresponding to the data block based on the second mask matrix and the first prefix sum matrix includes: Multiplying the second mask matrix by the last column of the first prefix sum matrix yields a third prefix sum matrix; the third prefix sum matrix represents the inter-block prefix sum results between each of the sub-data blocks; The first prefix sum matrix and the third prefix sum matrix are added together to obtain the second prefix sum matrix corresponding to the data block.

3. The method according to claim 1, characterized in that, The method further includes: After determining the second prefix sum matrix corresponding to the data block, the last prefix sum value in the second prefix sum matrix is ​​written into the data buffer as the kernel prefix sum value corresponding to the core.

4. The method according to claim 1, characterized in that, The first mask matrix is ​​an upper triangular matrix, and the size of the first mask matrix is ​​determined based on the size of the sub-data block; The second mask matrix is ​​a lower triangular matrix, and the size of the second mask matrix is ​​determined based on the number of sub-data blocks.

5. The method according to any one of claims 1 to 4, characterized in that, The method further includes: Obtain the core prefix sum and value corresponding to each core in the data processing chip from the data buffer; Based on the second prefix sum matrix corresponding to the data block, the prefix sum values ​​within each core, and the third mask matrix, a fourth prefix sum matrix is ​​determined, which represents the total prefix sum result between the core and the cores preceding the core.

6. The method according to claim 5, characterized in that, The step of determining the fourth prefix sum matrix based on the second prefix sum matrix corresponding to the data block, each of the kernel prefix sum values, and the third mask matrix includes: The third mask matrix is ​​multiplied by the intra-core prefix sum value to obtain the fifth prefix sum matrix; the fifth prefix sum matrix represents the inter-core prefix sum result between the cores. Based on the position of the data block in the data to be processed, the total prefix sum value of each core preceding the core is determined from the fifth prefix sum matrix; The fourth prefix sum matrix is ​​obtained by adding the second prefix sum matrix corresponding to the data block to the total prefix sum value.

7. The method according to any one of claims 1 to 4, characterized in that, The data block corresponding to the core is allocated sequentially based on the position of the core in the data processing chip, and the size of the data block corresponding to the core is determined based on the number of cores in the data processing chip and the size of the data to be processed.

8. A data processing chip, characterized in that, Any core included in the data processing chip is used for: Process the corresponding data blocks in the data to be processed; The data block corresponding to the core is divided into multiple sub-data blocks; each sub-data block is multiplied by the first mask matrix to obtain the first prefix sum matrix corresponding to each sub-data block, and the first prefix sum matrix represents the prefix sum result in each sub-data block; Based on the second mask matrix and the first prefix sum matrix, the second prefix sum matrix corresponding to the data block is determined; The second prefix sum matrix represents the prefix sum results within the core.

9. The data processing chip according to claim 8, characterized in that, The core is also used to: multiply the second mask matrix by the last column of the first prefix sum matrix to obtain a third prefix sum matrix; the third prefix sum matrix represents the inter-block prefix sum result between each of the sub-data blocks; and add the first prefix sum matrix and the third prefix sum matrix to obtain the second prefix sum matrix corresponding to the data block.

10. The data processing chip according to claim 8, characterized in that, The core is also used to: after determining the second prefix sum matrix corresponding to the data block, write the last prefix sum value in the second prefix sum matrix as the kernel-internal prefix sum value corresponding to the core into the data buffer.

11. The data processing chip according to claim 8, characterized in that, The first mask matrix is ​​an upper triangular matrix, and the size of the first mask matrix is ​​determined based on the size of the sub-data block; The second mask matrix is ​​a lower triangular matrix, and the size of the second mask matrix is ​​determined based on the number of sub-data blocks.

12. The data processing chip according to any one of claims 8 to 11, characterized in that, The core is also used to: obtain the intra-core prefix sum value corresponding to each core in the data processing chip from the data buffer; and determine a fourth prefix sum matrix based on the second prefix sum matrix corresponding to the data block, the intra-core prefix sum value, and the third mask matrix, wherein the fourth prefix sum matrix represents the total prefix sum result between the core and the cores preceding the core.

13. The data processing chip according to any one of claims 8 to 11, characterized in that, The data block corresponding to the core is allocated sequentially based on the position of the core in the data processing chip, and the size of the data block corresponding to the core is determined based on the number of cores in the data processing chip and the size of the data to be processed.

14. A computer device, characterized in that, Includes the data processing chip as described in any one of claims 8 to 13.

Images