Large model chip pre-filling stage reasoning optimization method and system under long context and application

Abstract

The invention discloses a large model chip pre-filling stage reasoning optimization method under a long context. The reasoning optimization method comprises the following steps: step 1, dynamically slicing input data and model parameters in a pre-filling stage; step 2, a QKV module, a self-attention module, an attention probe module and an FFN module are loaded and calculated in parallel in sequence, and residual error connection processing is carried out after loading calculation of the attention probe module and the FFN module is carried out; and step 3, completing the whole reasoning calculation. The invention further discloses a reasoning optimization system for implementing the reasoning optimization method and corresponding application, and the reasoning optimization system and the corresponding application have wide application scenarios.

Description

Technical Field

[0001] The present invention belongs to the technical field of chip optimization, and relates to an inference optimization method, system and application for the pre-fill stage of large model chips in long contexts. Background Art

[0002] In recent years, fields such as large language models have developed rapidly and achieved excellent results. However, due to the huge number of parameters and computing requirements in large language models, it poses a great challenge to deploy them on resource-constrained edge devices. Especially, more and more models support ultra-long contexts of 8K, 32K, 200K, or even 2M to support the long text processing ability of large models, which usually requires high-end GPUs with high costs to complete.

[0003] For edge devices, limited by the shortage of computing resources, scarce on-chip resources, and low-power requirements, it is currently difficult for large language models to run on their edge chips. According to the computing characteristics, the inference of large language models is usually divided into a pre-fill stage and a decoding stage. The former mainly outputs the first token after inputting an ultra-long text, with extremely high computing and storage requirements. The latter is the process of generating a new token using context information for each input token during the inference process, with relatively lower computational complexity and storage requirements. The computational complexity and resource requirements of the two stages vary greatly. In terms of storage, the storage in the pre-fill stage is 200,000 times that of the decoding stage. In terms of computational complexity, the computation in the pre-fill stage is 40 billion times that of the decoding stage. The optimization methods also vary greatly.

[0004] When deploying large language models on inference chips, limited by the shortage of computing resources and on-chip resources of the inference chips, a large amount of time is wasted on transferring input data and models from slow DDR storage devices to high-speed SRAM, and writing intermediate results or final results to low-speed storage devices. Due to the latency of DDR access, it greatly affects the overall running speed and increases its power consumption. Some slicing methods attempt to convert the problem into on-chip computing after tiling, such as flash-attention, which is based on local small modules, such as only accelerating the attention computing module, and is optimized for the characteristics of GPUs. It is not suitable for the low-resource and low-power information of the inference side, and the optimization algorithm design is unreasonable, so it cannot achieve a good acceleration effect. Summary of the Invention

[0005] In order to solve the deficiencies of the prior art, the purpose of the present invention is to provide an inference optimization method, system and application for the pre-fill stage of large model chips in long contexts, which are related to artificial intelligence, AI chip optimization and microarchitecture.

[0006] The present invention provides an inference optimization method for the pre-filling stage of a large model chip in a long context. The inference optimization method includes the following steps:

[0007] Step 1: Dynamically slice the input data and model parameters in the pre-filling stage;

[0008] Step 2: Sequentially and parallelly load and calculate the QKV, self-attention, attention_project, and FFN modules; and / or,

[0009] Perform residual connection processing after the loading and calculation of the attention_project and FFN modules;

[0010] Step 3: Complete the entire inference calculation.

[0011] In Step 1, the input data of each transformer layer is stored in a double data rate synchronous dynamic random access memory (DDR), with a shape represented as [s, d], and sliced into a shape of [s / n, d];

[0012] Dynamically slicing the model parameters includes: in the QKV layer processing, slicing the weight matrices Wq, Wk, and Wk according to the multi-head attention mechanism, slicing the shape of the weights from [d, d] to [d, d / n]; in the self-attention layer processing, slicing the matrices Q, K, and V; in the attention_project part processing, slicing the weight Wo, slicing the shape of the weights from [d, d] to [d, d / n]; in the FFN module part processing, slicing the parallel branch weights from [d, ffn_dim] to [d, ffn_dim / n], and slicing the serial branch weights from [ffn_dim, d] to [ffn_dim / n, d].

[0013] In Step 2, when loading and calculating QKV, it includes the following steps:

[0014] Step 2.1.1: Sequentially load the slices of Wq, Wk, and Wv from the double data rate synchronous dynamic random access memory (DDR) into the static random access memory (SRAM), and read the slice Xi of the input data from the DDR into the SRAM;

[0015] Step 2.1.2: Calculate the query vector Qi = Xi * Wq and the position embedding Rope(Qi) of the query vector Q respectively, calculate the key vector Ki = Xi * Wk and the position embedding Rope(Ki) of the key vector K, calculate the value vector Vi = Xi * Wv, and write the results into the Q, K, and V queues in the DDR for transfer;

[0016] Step 2.1.3: Reuse the SRAM storage space of the weight slice and the input data slice, and sequentially calculate the other input data slices in the DDR until all input data slices are executed;

[0017] Step 2.1.4: Release the SRAM space occupied by the weight slice, the calculated intermediate result, and the input data slice, complete the calculation of Q, K, and V, and store them in the DDR.

[0018] In Step 2, when loading and calculating self-attention, the following steps are included:

[0019] Step 2.2.1: Load Qi from the Q queue of the double data rate synchronous dynamic random access memory DDR into the SRAM, and perform the initialization operations of O, init M, and init L; where O is used to save the output of self-attention, M is used to save the local maximum value, and L is used to calculate the accumulated sum of the softmax denominator;

[0020] Step 2.2.2: While performing the O initialization operation, load K from the K queue of the DDR j into the SRAM;

[0021] Step 2.2.3: After K j is loaded, calculate S ij =Q i K j T , and while calculating, load V from the V queue j into the SRAM;

[0022] Step 2.2.4: Update m ij =max(m i，j-1 , rowmax(S ij ))), Pij = exp(S ij - m ij ), l ij = exp(m i，j-1 - m ij )l i，j-1 + rowsum(P ij ), O ij = exp(m i，j-1 - m ij )O i，j-1 + P ij V j , complete the single-layer inner-layer attention calculation; and while calculating, load K from the K queue of the DDR j+1 into the SRAM, which reuses the SRAM space occupied by K j ;

[0023] Step 2.2.5: Perform an inner loop according to the number of input slices;

[0024] Step 2.2.6: Execute the outer update of Oi and Li, and store the results in the DDR;

[0025] Step 2.2.7: Repeat the above operations, synchronize and load the next Q from the Q queue in the DDR i+1 into the SRAM, and reuse the SRAM space used by Q i and perform the same calculations as the above steps, and store all results in the DDR;

[0026] Step 2.2.8: Loop through the results of all heads in the multi-head attention mechanism and store them in the DDR.

[0027] In Step 2, when loading and calculating the attention projection, slice the input data by rows and slice the weight matrix by columns; for each sliced input data, sequentially load the weight slices and perform matrix multiplication o_out = att_o * Wo, and concatenate all results to obtain the output; among them, perform an outer loop on the sliced input data and an inner loop on the sliced weight matrix;

[0028] Load the input of the current layer from the DDR into the SRAM, perform element-wise addition, perform residual connection, process all sliced input data, complete the calculation of the attention projection, and write the output here to the DDR.

[0029] In Step 2, when loading and calculating the FFN module, perform an outer loop by slicing the input data by rows and an inner loop by slicing the weight matrix by columns;

[0030] The FFN includes two parallel weight calculation branches W1, W3 and one serial weight calculation branch W2. Slice the weights W1, W3 by columns and slice the weight W2 by rows. After writing the attention projection calculation results to the DDR, ffn_in reads them back from the DDR to the SRAM;

[0031] For each iteration of the outer loop, a piece of SRAM is pre-allocated for the memory ffn_out of the full reduction operation and initialized to 0; in the inner loop, the weight W1 is first loaded, and then in the next time slice, O1 = silu(ffn_in * W1) is calculated, and the memory of W1 is released; during the calculation and release process, W3 is loaded in parallel; in the next time slice, O3 = ffn_in * W3 is calculated, and after the calculation is completed, W3 is released; during the calculation of O3 and the release of W3, W2 is loaded; in the next time slice, O2 = Hadamard_product(O1, O3) is calculated, and then out = O2 * W2 is executed. During the execution of out = O2 * W2, O1 and O3 are released. After out is executed, W2 and O2 are released. Thus, a complete FFN calculation is completed; finally, through the full reduction operation of element-wise addition, the output result of the final inner loop is obtained;

[0032] In the outer loop, a residual connection process of element-wise addition is performed on the inner loop output result and the input data ffn_in of the FFN module, ffn_out = ffn_out + ffn_in, and then the result is written into the DDR, and the ffn_in is released. The outer loop is repeated until the entire FFN module calculation is completed and the result is written into the DDR.

[0033] The present invention also provides an inference optimization system for the pre-filling stage of a large model chip in a long context. The inference optimization system includes: an input processing module, a QKV calculation module, a self-attention calculation module, an attention projection module, a feed-forward neural network FFN module, and a result output module;

[0034] The input processing module is used to load the input data and perform preprocessing, and store the data in the SRAM;

[0035] The QKV calculation module is used to calculate Query, Key, and Value vectors, perform positional encoding on Query and Key, and store the results;

[0036] The self-attention calculation module is used to calculate attention scores and weights, and initially obtain the attention output;

[0037] The attention projection module is used to perform projection calculation on the output of the self-attention calculation module and merge the results to obtain the final attention output;

[0038] The feed-forward neural network FFN module is used to perform calculations of dimension elevation, element-wise multiplication, and dimension reduction;

[0039] The result output module is used to output the final calculation result in the decoding stage and release all intermediate calculation resources.

[0040] The present invention also provides the above-mentioned inference optimization method, or the application of the above-mentioned inference optimization system in the optimized scheduling of system resources for long-context large model inference.

[0041] The present invention also provides a hardware system for implementing the above-mentioned inference optimization method, the hardware system includes: a memory and a processor; a computer program is stored on the memory, and when the computer program is executed by the processor, the above-mentioned inference optimization method is implemented.

[0042] The present invention also provides a computer-readable storage medium, on which a computer program is stored, and when the computer program is executed by a processor, the above-mentioned inference optimization method is implemented.

[0043] The beneficial effects of the present invention include: Compared with the original 2048 context, the storage condition of the long context of 200k has increased by 100 times, the amount of calculation has increased by 10,000 times, the storage in the prefill stage is 200,000 times that of the decoding stage, and the amount of calculation, the amount of calculation in the prefill stage is 40 billion times that of the decoding stage. Its calculation mode is completely different from that of the decoding stage, and the design of the pipeline and calculation mode is completely different. It cannot be simply confused. Under the challenge of such a ten-order-of-magnitude difference, the present invention uses fine slicing, draws on the tiling idea of flash-attention2, globally processes the entire network, slices the calculation into sub-blocks, and simultaneously performs model parallel slicing and data parallel slicing on the weights and inputs respectively. And at the right time, use DDR for data transfer. Without precision loss, utilize the characteristics of the storage structure of the processor to complete the inference of long texts. With relatively small SRAM occupancy, optimize the implementation of long text inference on edge devices, which is of great help to the design of the processor chip microarchitecture and the optimization of the deployment algorithm. Description of the Drawings

[0044] In order to more clearly illustrate the technical solutions in the embodiments of the present invention or the prior art, the following will briefly introduce the drawings required for the description of the embodiments or the prior art. Obviously, the drawings in the following description are only some embodiments of the present invention. For those skilled in the art, without creative efforts, other drawings can be obtained based on these drawings.

[0045] Figure 1 It is the overall prefill structure diagram of the present invention

[0046] Figure 2 It is the schematic diagram of calculating Q, K, V of the present invention and caching them to DDR.

[0047] Figure 3 It is the pipeline diagram of calculating Q, K, V of the present invention and caching them to DDR.

[0048] Figure 4 This is the basic schematic diagram of flash attention in the present invention.

[0049] Figure 5 This is the schematic diagram of calculating self-attention in the present invention.

[0050] Figure 6 This is the pipeline schematic diagram of calculating self-attention in the present invention.

[0051] Figure 7 This is the schematic diagram of attention_project in the present invention.

[0052] Figure 8 This is the pipeline diagram of attention_project in the present invention.

[0053] Figure 9 This is the schematic diagram of the FFN module in the present invention.

[0054] Figure 10 This is the pipeline diagram of the FFN part in the present invention. Detailed implementation manners

[0055] Combined with the following specific embodiments and drawings, the present invention will be further described in detail. The processes, conditions, experimental methods, etc. for implementing the present invention, except for the specifically mentioned content below, are all common knowledge and well-known common sense in the art, and the present invention has no special limiting content.

[0056] The explanations of relevant English nouns in the present invention are as follows:

[0057] long context: Long context, referring to data input with a long context length.

[0058] prefill stage: Prefill stage, referring to the stage in the inference process of the model where, for the first input prompt (usually composed of multiple tokens, which may be dozens or millions), the first new token is generated.

[0059] SRAM: Static Random Access Memory, a type of memory with fast read and write speeds.

[0060] DDR: Double Data Rate Synchronous Dynamic Random Access Memory, a type of memory, usually used for main memory, with slower speeds.

[0061] tiling: Tiling, a data slicing processing method.

[0062] flash-attention2: An efficient attention mechanism method that uses tiling and parallel processing.

[0063] QKV: Query, Key, Value, which are three important matrices in the attention mechanism.

[0064] self-attention: Self-attention, a computational mechanism used to calculate the correlation between positions in an input sequence.

[0065] attention_project: Attention projection, which performs projection calculations on the output of the attention mechanism.

[0066] FFN: Feed-forward neural network, a neural network structure commonly used to transform input features.

[0067] Element-Wise Add: Element-wise addition, which performs an addition operation on the corresponding elements in two matrices.

[0068] residual: Residual, which refers to the structure in a neural network that directly adds the input to the output.

[0069] Hadamard_product: It performs an element-wise product of two matrices, and the two input matrices need to have the same shape.

[0070] fuse kernel: Kernel fusion, which means in a neural network, combining several layers that are related before and after into one layer to reduce the data reading and writing operations between layers.

[0071] memory-bound: Memory-bound, where the performance of a computing task is mainly limited by the memory access speed rather than the computing power of the processor.

[0072] compute-bound: Compute-bound, where the performance of a computing task is mainly limited by the computing power of the processor rather than the memory bandwidth.

[0073] In the prefill stage of the present invention, the complete input tokens need to be used, and the computational complexity has a quadratic relationship with the token length, which requires a large amount of storage and computing. This belongs to the compute-bound mode. Various methods such as parallelization, pre-computation, space-for-time, and pipeline parallelism need to be designed to minimize the access to slow storage during computing. The present invention designs a corresponding optimization and acceleration algorithm. By borrowing the tiling and model parallelization data slicing ideas of flash-attention2 and reasonably using DDR as a storage transfer, an optimized on-chip computing scheme is designed. The input and the model are reasonably sliced, and the advantages of high-efficiency memory access of SRAM are fully utilized. All computations are converted into in-SRAM (on-chip) computations, which can greatly reduce the latency caused by memory access and reduce the number of times of accessing DDR sequentially for successive computations. Through a carefully designed pipeline, the storage and computing are maximally interleaved to reduce the total computing time, reduce power consumption, and reduce latency, converting the huge computing and storage requirements into an acceptable feasible solution.

[0074] The specific data involved in the present invention are special cases of data applied in a specific implementation scenario and are applicable in the corresponding implementation scenario.

[0075] Through the processing of this method, all computations can be converted into on-chip computations. Taking the llama2 7b model as an example, for a 200K context length, the SRAM occupancy in the prefill stage is controlled within 12.4MB, and all computations can be completely converted into on-chip computations. (If lower requirements for SRAM are needed, finer-grained slicing processing can be performed on the basis of this method to reduce the demand for SRAM), thereby enabling the realization of ultra-long context model inference on low-power edge chips, which is of great help for the low-cost use of large models.

[0076] The present invention provides an inference optimization method for the prefill stage of large model chips in an extremely long context (greater than or equal to 200K) environment. Without increasing computing resources and storage resources, it can be extended to support contexts of any length. For a given inference chip, its computing model is finely sliced and scheduled, making full use of the fast read and write speed of on-chip storage resources (usually on-chip SRAM) (usually 1-2 orders of magnitude higher than DDR). This reduces the I / O time in computing, decreases the access to DDR, and improves the computing efficiency. Through carefully designed software scheduling and resource allocation, it can support inference of any length with extremely low resource occupancy, making the inference time of the context length only affected by computing resources and independent of the context inference length and storage. Thus, it can support efficient inference of any length context on low-power and low-cost AI inference chips. In the prefill stage, compared with the original 2048 context for a long context of 200k, the storage condition increases by 100 times and the computation amount increases by 10,000 times. Only 12.4M of SRAM is required to complete the processing of 200K context (if lower requirements for SRAM are needed, finer-grained slicing processing can be performed based on the method of the present invention to reduce the SRAM requirement). All computations are completed in on-chip SRAM, reducing DDR access, which provides favorable assistance for low-power chip design and microarchitecture.

[0077] In practical application scenarios, the inference of large language models is usually divided into a prefill stage and a decode stage. The difference in the computation amount and resource demand between the two stages reaches 40 billion. Therefore, their computing modes and pipeline designs are very different. The present invention designs a corresponding optimization and acceleration method for the characteristics of large computation amount and large resource demand in the prefill stage.

[0078] For the pre-filling stage, the input shape is [b, s, d], where b is the batch size, s is the sequence length, and d is the latent variable dimension. For the sake of convenience, assume b = 1 and ignore this dimension. For pre-filling, take s = 200,000. According to requirements, s can be larger. The length of the sequence in the pre-filling stage is 200,000 times that of the sequence in the decoding stage. The theoretical computational volume in the pre-filling stage is 40 billion times that of the decoding stage. Compared with the decoding stage, the design scheme of the inference optimization method in the pre-filling stage will be completely different. The pre-filling stage is computationally intensive, and both storage and calculation are relatively large challenges. It is necessary to slice the data, adopt a suitable slicing scheme, on-chip computing strategy, and use DDR to complete data transfer at the appropriate time to achieve inference optimization in the pre-filling stage. In the pre-filling stage, 200,000 tokens are processed at a time, that is, s = 200,000, and the input shape is [200,000, d]. The model is divided into L layers. Since the model structure of each layer is the same (transformer structure) and the data size is also the same, the optimization scheme can be fully reused. Therefore, one of the layers is taken as an example for illustration, as Figure 1 is the overall structure diagram of the pre-filling stage; the input data of each transformer layer is represented as [s, d] and sliced into [s / n, d], where n is the number of slices.

[0079] 1. In the present invention, for the current Transformer layer, its input is the output of the previous Transformer layer, with a shape of [200000, d]. Compared with the decoding stage, its value is 200,000 times larger. Taking Llama 2 7B as an example, d = 4096. Stored in 16-bit floating-point numbers, the space it occupies reaches 1.6G. For other models, d > 4096, and the corresponding storage requirements are even greater. Its value is extremely large. Therefore, when performing calculations, the output of the previous layer, output, can only be written to the DDR. Since the output of the previous layer serves as the input of the current layer, that is, the input of the current layer is stored in the DDR. For the Transformer layer input [200000, d] located in the DDR, following the storage idea similar to the pages in X86 in memory, the input shape is divided into multiple blocks; in a specific embodiment, taking the input size of Llama 2 as an example, the input of the Transformer layer is selected to be divided into 200 blocks. The reason for choosing 200 is that the shape of each block after division is [1000, 4096], and the storage space it requires is 8M. For most acceleration chips, 8M SRAM is an acceptable setting. For some larger models, due to the larger hidden variable dimension d, the corresponding input occupied space will be larger, and the number of slices will be adjusted. So, 200 here is only for illustrative purposes and is not a fixed value. In fact, if the available SRAM is allowed to be larger, the number of slices can be reduced. On the contrary, if the available SRAM is allowed to be smaller, the number of slices needs to be increased accordingly. For the processing of weights, the three weights Wq, Wk, and Wv need to be loaded from the DDR into the SRAM. The shapes of the three weights are all [d, d], and slicing is required. Since the Transformer defaults to using the multiple-header attention mechanism, although the multi-head processing is in the inner-product stage, in the present invention, it can be introduced in the projection stage. After projecting Wq, Wk, and Wv, Q, K, and V are obtained. Q, K, and V are sliced on the last dimension, that is, the hidden variable dimension d. Here, it is chosen to perform multi-head slicing on the three weights Wq, Wk, and Wk. By performing projection calculations after slicing the weights, the purpose of dividing the multi-heads is achieved. The number of headers is n, and the weight shape of each head is [d, d / n]. At this time, the weights can be directly placed into the limited SRAM space. Therefore, they can be directly loaded from the DDR onto the SRAM and the on-chip calculations can be completed on the SRAM. In order to control the SRAM usage amount of the weights, the weights are split into a multi-head form. The original shape of the weights [d, d] is split into the shape of the weights after multi-heads [d, d / n], where n is the number of multi-heads, and it is used as the number of slices here;

[0080] In a specific implementation process, since the input will be sliced into 200 pieces, Wq, Wk, and Wv need to be repeatedly used for different input slices (for different input slices, calculations need to be performed with Wq, Wk, and Wv). To minimize the repeated loading of Wq, Wk, and Wv as much as possible, the Wq, Wk, and Wv of the multi-head slices (with a shape of [d, d / n]) are loaded from the DDR into the SRAM in sequence. Then, one slice of input data is read from the DDR, denoted as Xi, with a shape of [1000, d], and Q is calculated respectively i = X i * W q , Rope(Q i ), K i = X i * W k , Rope(K i ), V i = X i * W v , and their output shapes are all [1000, d / n]. The results are written into the Q, K, and V queues for transfer in the DDR. Then, the next slice of input data is read, denoted as X j , X j reuses the storage space occupied by X i , with a shape of [1000, d]. For X j , Q is calculated respectively j = X j * W q , Rope(Q j ), K j = X j * W k , Rope(K j ), V j = X j * Wv, where Q j , K j , V j reuses the storage space occupied by Q i , K i , V i , with shapes all being [1000, d / n]. The results are written into the Q, K, and V queues for transfer in the DDR until all slices of the input data X have been processed, which is performed 200 times and the calculations of Q, K, and V are completed. At this time, the SRAM space occupied by X 200 , Q 200 , K 200 , V 200 , W K , W V , W Q is released. That is, in this step, the calculations of Q, K, and V are completed and stored in the DDR. The schematic diagram is as shown inFigure 2 As shown, its pipeline diagram is as Figure 3 shown below.

[0081] In a specific implementation process, the execution steps of the Q, K, and V calculation parts are as follows:

[0082] In the first time slice, load X from the DDR 1 into the SRAM, and its shape is [1000, d];

[0083] In the second time slice, execute Q 1 = X 1 * Wq, Q 1 = Rope(Q 1 ), and its output shape is [1000, d / n];

[0084] In the third time slice, execute K 1 = X 1 * Wk, K 1 = Rope(K 1 ), and its output shape is [1000, d / n]. At the same time, execute in parallel to write the previous Q 1 to the DDR;

[0085] In the fourth time slice, execute V 1 = X 1 * Wv, and its output shape is [1000, d / n]. At the same time, execute in parallel to write the previous K 1 to the DDR;

[0086] In the fifth time slice, load X from the DDR 2 into the SRAM, and reuse the memory of X 1 , and its shape is [1000, d]. At the same time, execute in parallel to write the previous V 1 to the DDR;

[0087] In the sixth time slice, execute Q 2 = X 2 * Wq, Q 2 = Rope(Q 2 ), and its output shape is [1000, d / n]; ...

[0089] In the 801st time slice, write the previous V 200 to the DDR, and then release X 200 , Q 200 , K 200 , V 200 , Wq, Wk, Wv, and complete the calculation of Q, K, and V.

[0090] During the above execution process, every 4 time slices complete one round of calculation of Q, K, and V.

[0091] 2. After completing the calculation of Q, K, and V corresponding to one head and storing them in the DDR, start the calculation of self-attention. For the calculation of self-attention, Q, K, and V were split into sub-blocks and stored in the DDR in the previous step. Here, referring to the idea of Flash-attention2 tiling, the outer layer executes 200 loops. The 200 here is consistent with the previous setting. Taking the input size of llama2 as an example, choose to execute 200 loops in the outer layer. The shape of each block after splitting is [1000, 4096], and the storage space it requires is 8M. For most acceleration chips, 8M SRAM is an acceptable setting. For some larger models, since the hidden dimension d is larger, the corresponding input occupied space will be larger, and the number of slices will be adjusted. So the 200 here is just for illustration and is not a fixed value. In fact, if the available SRAM is allowed to be larger, the number of slices can be reduced. On the contrary, if the available SRAM is allowed to be smaller, the number of slices needs to be increased accordingly. Load Qi from the Q queue in the DDR into the SRAM, and then sequentially enter some initializations of O (O is used for the output of self-attention, with a shape of [1000, d / n]), init M (M is used to save the local maximum value, with a shape of

[1000] ), init L (L is used to calculate the cumulative sum of the denominator in softmax, with a shape of

[1000] ). The selected input sequence length is 200000, and the number of slices is 200, so the sequence length in a single calculation is 1000; the numbers 200, 1000, 200000, etc. in the above description are just special cases of a specific implementation. These initialization processes are parallel to the process of loading Kj from the K queue in the DDR into the SRAM in the inner loop. Here, for the convenience of understanding, the original schematic diagram of flash-attention2 is given, as Figure 4 shown. Figure 4 It is only used to illustrate the original Flash-attention2 algorithm, just for a better understanding of the method of the present invention. The method of the present invention is not an implementation of the flash-attention algorithm and there is no inclusion relationship.

[0092] Figure 4 In, the meaning of the key parts in the figure:

[0093] Q, K, V matrices:

[0094] Q (Query): Represents the query vector of the input sequence, with a size of Q: N×d, where N is the sequence length and d is the dimension of the hidden layer.

[0095] K (Key) and V (Value): respectively represent the key and value vectors. In the figure, K T represents the transpose matrix of K, with dimensions K T : d × N; the dimensions of V are V: N × d.

[0096] SRAM operations:

[0097] The orange squares in the figure represent the process of loading data from external memory (such as DDR or HBM) into on-chip memory (SRAM).

[0098] Copy Block to SRAM: refers to loading the block data of the K and V matrices into SRAM.

[0099] Compute Block on SRAM: perform calculations on SRAM. The calculation block shown in the figure is QK T · M × N, to calculate the attention scores.

[0100] Inner Loop and Outer Loop: These loop operations refer to the block processing of the input data, processing each block in SRAM, and finally outputting the results to HBM (High Bandwidth Memory).

[0101] Pipeline operations:

[0102] Figure 4 In [the figure], it shows how to block-process the data in the pipeline operation through Inner Loop and Outer Loop, and gradually store the results into HBM after processing each block.

[0103] sm(QK T )V: represents performing the Softmax operation on the result of QK T calculation, and then multiplying by the V vector, with the output dimensions of N × d.

[0104] Data blocking and loading:

[0105] Figure 4 It shows how to block (slice by column) and load the large-scale K and V matrices into SRAM. Since the SRAM space is limited and only a small part of the data can be processed at a time, the K and V need to be sliced into smaller blocks and loaded and processed step by step.

[0106] Matrix multiplication and Softmax operation:

[0107] On SRAM, first calculate QK TThe result is then applied with the Softmax operation, and finally multiplied by the V matrix to obtain the final attention output sm(QK T )V.

[0108] To efficiently utilize SRAM, the result of each calculation is immediately processed and output to HBM to prevent excessive occupation of SRAM space.

[0109] Pipeline and loop operations:

[0110] The pipelining of the Inner Loop and Outer Loop continuously loads new data blocks during the calculation and outputs the processed data. In this way, the utilization rate of hardware resources is maximized and the overhead of data transmission is reduced.

[0111] In a specific implementation process, such as Figure 5 , the outer loop runs 200 times. 200 is just a special case of a specific implementation. Each time the loop reads a slice Qi from the Q queue in DDR into SRAM. The inner loop runs 200 times. 200 is just a special case of a specific implementation. Each time the loop first loads Kj from the K queue in DDR into SRAM. After Kj is loaded, the calculation is performed in sequence. Among them, S ij = Q i K j T , because it is a computational operator and the calculation takes the main time, data loading can be performed in parallel while calculating, that is, using the calculation time to cover the data loading time. Here, S ij = Q i K j T and loading Vj from the V queue in DDR into SRAM are performed in parallel. After S ij = Q i K j T , the following calculations are performed serially in sequence, m ij = max(m i，j-1 , rowmax(S ij ), Pij = exp(S ij - m ij ), l ij = exp(m i，j-1 - m ij )l i，j-1 + rowsum(P ij ), O ij = exp(m i，j-1 - m ij )O i，j-1 + P ij V j, namely Update Max, Update L ij , Update O ij , that is, the calculation of the inner-layer attention can be completed, and then the next V is loaded from the K sequence in the DDR j+1 data block. After the inner loop completes 200 calculations, the outer-layer Update Qi and Update Li are executed, and the result [1000, d / n] is written back to the DDR. Then the outer loop reads the next Q from the Q sequence in the DDR i+1 , execute the same process, finally write all the results [200000, d / n] back to the DDR, and then process the calculation of the next header. After all n heads are calculated, the total shape after being written back to the DDR is [200000, d], and its schematic diagram is as Figure 5 shown

[0112] In a specific implementation process, the execution of the self-attention calculation part is as Figure 6 shown, and the steps are as follows

[0113] In the first time slice, Q is read from the Q bank (the memory unit composed of Q slices located on the DDR) 1 to the SRAM, and its shape is [1000, d / n]. 1000 is just a special case of a specific implementation

[0114] In the second time slice, K is read from the K bank (the memory unit composed of K slices located on the DDR) 1 to the SRAM, and its shape is [1000, d / n]. While reading K 1 , some initializations in the outer loop are completed in parallel, O (O is used as the output of self-attention, with a shape of [1000, d / n]), init M (M is used to save the local maximum value, with a shape of

[1000] ), initL (L is used to calculate the cumulative sum of the denominator in the softmax, with a shape of

[1000] );

[0115] In the third time slice, calculate S 1，1 = Q 1 K 1 T , and its output shape is [1000, 1000]. While calculating S 1,1 , V is read from the V bank (the memory unit composed of V slices located on the DDR) 1 to the SRAM, and its shape is [1000, d / n];

[0116] In the fourth time slice, update the local maximum value max_value and update P1 , update L 1 , complete the calculation of softmax;

[0117] In the 5th time slice, calculate O 1 O = 1 O 1 + P 1 V 1 whose output shape is [1000, d / n]. While calculating O 1 , read K from Kbank (a memory unit composed of K slices located on the DDR) to SRAM, whose shape is [1000, d / n]; 2

[0118] Thereafter, complete a set of loop Si,j calculations every three times, update Max, P, L, O = O + PV;

[0119] In the 6th time slice, calculate S 1,2 S = 1 Q 2 T K T whose output shape is [1000, 1000]. While calculating S 1,2 , read V from Vbank (a memory unit composed of V slices located on the DDR) to SRAM, whose shape is [1000, d / n]; 2

[0120] In the 7th time slice, update the local maximum value max_value, update P 1 , update L 1 , complete the calculation of softmax;

[0121] In the 8th time slice, calculate O 1 O = 1 O 1 + P 2 V 2 whose output shape is [1000, d / n]. While calculating O 1 , read K from Kbank (a memory unit composed of K slices located on the DDR) to SRAM, whose shape is [1000, d / n]; 3

[0122] Thereafter, complete a set of loop Si,j calculations every three times, update Max, P, L, O = O + PV;

[0123] In the 9th time slice, calculate S 1,3 S = 1 Q 3 T K T whose output shape is [1000, 1000]. While calculating S 1,3Meanwhile, read V from Vbank (a memory cell composed of V slices located on the DDR). 3 into SRAM, with a shape of [1000, d / n]; ...

[0125] For 200 inner loops, since the first data loading cannot be parallelized and pipelined, it occupies two time slices. In total, 2 + 3 * 200 = 602. That is, at the 602nd time slice, the last O = O + PV operation of the inner loop is executed.

[0126] At the 602nd time slice, calculate O 1 = O 1 + P 1 V 200 , with an output shape of [1000, d / n];

[0127] At the 603rd time slice, complete one outer loop and write O 1 to the O bank, which is located on the DDR and has a shape of [1000, d / n]; ...

[0129] That is, after 603 time slices, complete a set of complete inner loops, which is one outer loop. After 603 * 200 = 120600 time slices, that is, at the 120600th time slice, complete the last execution process.

[0130] At the 120600th time slice; complete the 200th outer loop and write O 200 to the O bank; it is located on the DDR and has a shape of [1000, d / n]. Thus, all outer loops are completed.

[0131] The total output written to the DDR is [200000, d / n].

[0132] Repeat the process of the first step, read in the next headers Q, K, V, and then execute the calculation process of this step to complete the calculation of attention for all headers. The results are written to the DDR, and the total output written to the DDR is [200000, d].

[0133] 3. After completing the calculation of the self-attention part in the previous step, next, perform a projection calculation on Wo (also known as attention_project in some literature). The output result of the second step is denoted as attn, and its shape is [200000, d]. 200000 is just a special case of a specific implementation. Taking llama2 7B as an example, d = 4096, stored in 16-bit floating-point numbers, and the space it occupies reaches 1.6G. For other models, d > 4096, and the corresponding storage requirements are even greater. Therefore, it is also necessary to slice the input. For the weights of this layer, the shape of Wo is [d, d]. To save SRAM, it is also necessary to slice the weight Wo. For the input data with a shape of [200000, d], similar to the second step, it is sliced into 200 blocks, each with a shape of [1000, d]; for the weight Wo, it is sliced into 8 blocks by columns, and the size of each block is [d, d / 8]. Here, choosing to slice into 8 blocks takes the input size of llama2 as an example. The reason for choosing 8 is that the shape of each block after splitting is [4096, 512], and the storage space it requires is 4M. Plus, the input also needs to occupy SRAM. Here, 4M SRAM for the weight is an acceptable setting. For some larger models, due to the larger hidden dimension d, the corresponding input occupied space will be larger, and the number of slices will be adjusted. So, 8 here is just for illustrative purposes and is not a fixed value. In fact, if more available SRAM is allowed, the number of slices can be reduced. On the contrary, if less available SRAM is allowed, the number of slices needs to be increased accordingly. The above 8 is just a special case of a specific implementation. In this step, execute o_out = attn * Wo, which is attention_project. Here, the shape of the input attn is [200000, d], the shape of the weight Wo is [d, d], and the shape of the output o_out is [200000, d]. Because the memory required for the input attn is much larger than that required for the weight Wo, the input is placed in the outer loop to reduce the number of times the inner loop is loaded. In this specific implementation, the input attn needs to be sliced into 200 data blocks, and the weight weight needs to be sliced 8 times. Therefore, the number of outer loop iterations is 200 times to complete the loading of attn, and the inner loop iterates 8 times. The inner loop respectively completes the loading of the sliced weights and the calculation of o_out ij = attn i * Wo j calculation, and the concatenation operation of the last dimension, and reuse the sub-block weight Wo in the inner loop j corresponding SRAM. This calculation process is computationally sensitive. Through pipeline arrangement, the calculation and data transfer are parallelized, and the total time is determined by the calculation time. After the inner loop ends, the result o_out is obtainedi , with a shape of [1000, d], release attn i 's memory and Wo j 's SRAM, and then load the corresponding block X of the current block from the input bank of DDR i , with a shape of [1000, d], for the output o_out i Execute o_out i＝ o_out i+ X i , that is, the corresponding residual link operation is completed. The operation performed is Element - WiseAdd, and the result is written into the DDR queue of o_out, with a shape of [1000, d]. Then, for the next entity in the outer loop, read the next attn from the DDR queue of attn i+1 , execute the same steps as above until all outer loops are completed. The output is written into the DDR queue of o_out, with a shape of [200000, d], until the entire attention_project is completed. The schematic diagram for calculating attention_project is as Figure 7 shown.

[0134] In a specific implementation process, the execution of the attention_project part is as Figure 8 shown, and the steps are as follows:

[0135] In the first time slice, read attn from the attn bank (a memory unit composed of attn slices located on DDR) to SRAM, with a shape of [1000, d]; 1 to SRAM, with a shape of [1000, d];

[0136] In the second time slice, read Wo_ from Wo (located on DDR) to SRAM, with a shape of [d, d / 8]; 1 to SRAM, with a shape of [d, d / 8];

[0137] In the third time slice, the calculation is performed in parallel with data reading and writing. During the calculation, o_out is executed serially 1,1 = attn 1 * Wo_ 1 , o_out 1,1 = o_out 1,1 + x 1,1 , and the output shape is [1000, d / 8]. The calculation is in parallel with data reading and writing, and for data reading and writing, it is executed serially. Read X from DDR 1,1 to SRAM, with a shape of [1000, d / 8], read Wo_ from Wo 2 to SRAM, with a shape of [d, d / 8];

[0138] In the above implementation process, it can be regarded as completing the output projection calculation and the residual calculation within a time slice;

[0139] In the 4th time slice, the calculation and data reading / writing are performed in parallel, and o_out is executed serially during the calculation 1,2 = attn 1 * Wo_ 2 , o_out 1,2 = o_out 1,2 + x 1,2 , and the output shape is [1000, d / 8]. The calculation and data reading / writing are parallel, while for the data reading / writing, it is executed serially, and written into o_out 1,1 to the DDR, with the shape of [1000, d / 8], and read X from the DDR 1,2 , with the shape of [1000, d / 8] to the SRAM, and read Wo_ from Wo 3 to the SRAM, with the shape of [d, d / 8];

[0140] In the 5th time slice, the calculation and data reading / writing are performed in parallel, and o_out is executed serially during the calculation 1,3 = attn 1 * Wo_ 3 , o_out 1,3 = o_out 1,3 + x 1,3 , and the output shape is [1000, d / 8]. The calculation and data reading / writing are parallel, while for the data reading / writing, it is executed serially, and written into o_out 1,2 to the DDR, with the shape of [1000, d / 8], and read X from the DDR 1,3 to the SRAM, with the shape of [1000, d / 8], and read Wo_ from Wo 4 to the SRAM, with the shape of [d, d / 8]; ...

[0142] In the 10th time slice, the calculation and data reading / writing are performed in parallel, and o_out is executed serially during the calculation 1,8 = attn 1 * Wo_ 8 , o_out 1,8 = o_out 1,8 + x 1,8 , and the output shape is [1000, d / 8]. The data reading / writing is executed serially, and written into o_out 1,7 to the DDR, with the shape of [1000, d / 8], and read X from the DDR 1,8 to the SRAM, with the shape of [1000, d / 8];

[0143] In the 11th time slice, write into o_out1,8 To the DDR, with the shape of [1000, d / 8];

[0144] As can be seen from the above steps, after 11 time slices, an operation of one outer loop is completed;

[0145] In the 12th time slice, read attn from the attn bank (a memory unit composed of attn slices located on the DDR) 2 to the SRAM, with the shape of [1000, d];

[0146] In the 13th time slice, read Wo_ from Wo (located on the DDR) 1 to the SRAM, with the shape of [d, d / 8];

[0147] In the 14th time slice, the calculation is carried out in parallel with data reading and writing, and o_out is serially executed during the calculation 2,1 = attn 2 * Wo_ 1 , o_out 2,1 = o_out 2,1 + x 2,1 , with the output shape of [1000, d / 8], data reading and writing are serially executed, and read X from the DDR 2,1 to the SRAM, with the shape of [1000, d / 8], read Wo_ from Wo 2 to the SRAM, with the shape of [d, d / 8]; ...

[0149] One outer loop requires 11 time slices, and 200 loops in total require 2200 time slices;

[0150] In the 2199th time slice, the calculation is carried out in parallel with data reading and writing, and o_out is serially executed during the calculation 200,8 = attn 200 * Wo_ 8 , o_out 200,8 = o_out 200,8 + x 200,8 , with the output shape of [1000, d / 8], data reading and writing are serially executed, and write o_out 200,7 to the DDR, with the shape of [1000, d / 8], read X from the DDR 200,8 to the SRAM, with the shape of [1000, d / 8];

[0151] In the 2200th time slice, write o_out 200,8 to the DDR, with the shape of [1000, d / 8];

[0152] At this point, all calculations of attention_project are completed, and the results are stored in DDR with a shape of [200000, d].

[0153] 4. Next, we enter the FFN module, which has a parallel two-way weight calculation branch for dimension enhancement, followed by element-by-element multiplication, followed by a serial weight branch for dimension reduction. The FFN input shape is [200000, d], where 200000 is just a special case of a specific implementation, with llama2 7B is an example, where d=4096, stored as a 16-bit floating point number, and the space occupied reaches 1.6G. For other models, d>4096, and the corresponding storage requirements are even greater. Therefore, it is also necessary to slice the input and the weights. The input is divided into 200 blocks to execute the outer loop, and the weights are divided into 16 blocks to execute the inner loop. The two parallel weight matrices of the former can be sliced ​​in columns, and the serial weight matrices of the latter can be sliced ​​in rows. Here, 16 blocks are selected. Taking the input size of llama2 as an example, the hidden dimension ffn_dim of its FFN is 11088. The reason for choosing 16 is that the shape of each block after segmentation is [4096, 688], and the storage space required is 5.64M, plus the input also needs to occupy SRAM, here the weight is 5.64M SRAM is an acceptable setting. For some larger models, since the hidden dimension ffn_dim is larger, the corresponding input space will be larger, and the number of slices will be adjusted. So the 16 here is just for illustration, not a fixed value. In fact, if the available SRAM is allowed to be larger, the number of slices can be reduced. On the contrary, if the available SRAM is allowed to be smaller, the number of slices needs to be increased accordingly.

[0154] The input is denoted as ffn_in, and its shape is [200000, d]. Slice it and divide it into 200 slices. 200 is just a special case of a specific implementation. The shape of each slice is [1000, d]. Execute the outer layer 200 times. In each outer loop, input ffn_in iThe shape is [1000, d]. There is an inner loop nested 16 times in each outer loop. 16 is just a special case of a specific implementation. In the inner loop, the weights will be sliced. W1 and W3 are the weights of a set of parallel branches, and column slicing is used for both of them. After slicing, the sliced shapes are both [d, ffn_dim / 16]. W2 is sliced by rows, and the sliced shape is [ffn_dim / 16, d]. With this slicing, the output shape of each iteration of the inner loop is [1000, d]. Therefore, for each iteration of the outer loop, a piece of SRAM is pre-allocated for the memory ffn_out for the all-reduce operation of the 16 iterations of the inner loop, and its shape is [1000, d], which is initialized to 0. In the 16 inner loops, first, the sliced W1 is loaded, which is sliced by columns and has a shape of [d, ffn_dim / 16]. Then, in the next time slice, the sliced W3 is loaded, which is sliced by columns and has a shape of [d, ffn_dim / 16]. Here, the loading of W3 and O1 = silu(ffn_in * W1) are executed in parallel. After this step is completed, the SRAM occupied by W1 is immediately released. In the next time slice, load W2, which is sliced by rows and has a shape of [ffn_dim / 16, d]. This time, the loading of W2 and O3 = ffn_in * W3 are executed in parallel. After the calculation, W3 is released. After execution is completed, in the next time slice, O2 = Hadamard_product(O1, O3) is executed serially (Hadamard_product is also known as the Hadamard product, which is the element-wise product of two matrices and requires the two input matrices to have the same shape). After O2 = Hadamard_product(O1, O3) is executed, out = O2 * W2 is executed serially. During the execution of out, the SRAM occupied by O1 and O3 is released simultaneously. After out = O2 * W2 is executed, W2 and O2 are released. Thus, one inner loop is completed, and the result shape is [1000, d]. Since column slicing is performed on W1 and W3 and row slicing is performed on W2 during the calculation, the result needs to perform an all-reduce operation, that is, ffn_out = ffn_out + out, and the calculation result shape is [1000, d]. Thus, one inner loop is completed. In the outer loop, the input ffn_in performs an Element Wise-Add residual short-cut process ffn_out = ffn_out + ffn_in, and then the result is written to the DDR. The outer loop is repeated 200 times to complete the entire FFN calculation, and the result is written to the DDR, with a shape of [200000, d]. The principle is as Figure 9 shown.

[0155] In a specific implementation process, the execution of the FFN module part is as follows: Figure 10 as shown below:

[0156] In the first time slice, a slice ffn_in of the input ffn_in is loaded from the DDR 1 into the SRAM, with a shape of [1000, d];

[0157] In the second time slice, a slice W1 of the weight W1 is loaded from the DDR 1 into the SRAM, with a shape of [d, ffn_din / 16]. At the same time, a section of SRAM is allocated and named ffn_out 1 for use in performing eletwiseadd for the inner loop iteration and initialized to 0;

[0158] In the third time slice, data calculation and data loading are performed in parallel. The calculation part is O1 1,1 = silu(ffn_in 1 * w1 1 ), with an output shape of [1000, ffn_dim / 16]. For the data loading part, a slice W3 of the weight W3 is loaded from the DDR 1 into the SRAM, with a shape of [d, ffn_din / 16];

[0159] In the fourth time slice, data calculation and data loading are performed in parallel. The calculation part is divided into two parts, which are executed serially, namely O3 1,1 = ffn_in 1 * w3 1 , with an output shape of [1000, ffn_dim / 16], and O2 1,1 = H(O1 1,1 , O3 1,1 ), where H represents Hadamard_product, which is element-wise multiplication of corresponding matrices. For the data loading part, W11 is released (only indicating reusability, this step does not take time), and a slice W2 of the weight W2 is loaded from the DDR 1 into the SRAM, with a shape of [ffn_din / 16, d];

[0160] In the fifth time slice, data calculation and data release are performed in parallel. The calculation part is Out 1,1 = O2 1,1 * w2 1 , with an output shape of [1000, d]. For the data loading part, W3 1 , O1 1,1 , O3 1,1 are released (only indicating reusability, this step does not take time);

[0161] The 6th time slice, perform data calculation and data loading in parallel. The calculation part is ffn_out 1 = ffn_out 1 + Out 1,1 , this step completes eletwiseadd. For the data loading part, release O2 1,1 , W2 1 (only indicate that it can be reused, this step does not take time), load the slice W1 of weight W1 from DDR 2 to SRAM, with the shape of [d, ffn_din / 16];

[0162] It can be seen from the above description that 4 time slices are required to complete the ffn operation;

[0163] The 7th time slice, perform data calculation and data loading in parallel. The calculation part is O1 1,2 = silu(ffn_in 1 * w1 2 ), with the output shape of [1000, ffn_dim / 16]. For the data loading part, load the slice W3 of weight W3 from DDR 2 to SRAM, with the shape of [d, ffn_din / 16];

[0164] The 8th time slice, perform data calculation and data loading in parallel. The calculation part is divided into two parts, which are executed serially, namely O3 1,2 = ffn_in 1 * w3 2 , with the output shape of [1000, ffn_dim / 16], O2 1,2 = H(O1 1,2 , O3 1,2 ), where H represents Hadamard_product, which is element-wise multiplication of corresponding matrices. For the data loading part, release W1 2 (only indicate that it can be reused, this step does not take time), load the slice W2 of weight W2 from DDR 2 to SRAM, with the shape of [ffn_din / 16, d];

[0165] The 9th time slice, perform data calculation and data release in parallel. The calculation part is Out 1,2 = O2 1,2 * w2 2 , with the output shape of [1000, d]. For the data loading part, release W3 2 , O1 1,2 , O3 1,2(Only indicate that it can be reused. This step does not take time.)

[0166] In the 10th time slice, data calculation and data loading are performed in parallel. The calculation part is ffn_out 1 = ffn_out 1 + Out 1,2 , this step completes eletwiseadd. For the data loading part, O2 is released 1,2 , W2 2 (Only indicate that it can be reused. This step does not take time.), load the slice W1 of weight W1 from DDR 3 to SRAM, and its shape is [d, ffn_din / 16];

[0167] It can be seen from the above description that 4 time slices are required to complete the ffn operation;

[0168] Since the initialization pipeline requires 2 time slices and 16 inner loops of the wide area need to be executed, the required number of time slices is 2 + 4 * 16 = 66; ...

[0170] In the 66th time slice, data calculation and data loading are performed in parallel. The calculation part is ffn_out 1 = ffn_out 1 + Out 1,16 , this step completes eletwiseadd. For the data loading part, O2 is released 1,16 , W2 16 (Only indicate that it can be reused. This step does not take time.);

[0171] In the 67th time slice, data calculation is performed, ffn_out 1 = ffn_out 1 + ffn_in 1 , perform a residual connection operation, that is, element-wise addition (Element-WiseAdd) is required;

[0172] In the 68th time slice, write ffn_out 1 to DDR, and its shape is [1000, d];

[0173] At this point, one execution of the outer loop is completed, and a total of 68 * 20 = 13600 time slices are required to complete the total execution;

[0174] In the 69th time slice, load the slice ffn_in of the input ffn_in from DDR 2 to SRAM (reuse the previous ffn_in 1of the SRAM memory space), with a shape of [1000, d]; ...

[0176] At the 13,598th time slice, data calculation and data loading are performed in parallel. The calculation part is ffn_out 200 = ffn_out 200 + Out 200,16 , this step completes eletwiseadd, the data loading part, and releases O2 200,16 , W2 16 (Only indicate that it can be reused, and this step will not take time);

[0177] At the 13,599th time slice, data calculation is performed, ffn_out 200 = ffn_out 200 + ffn_in 200 , perform a residual connection operation, that is, element-wise addition (Element-WiseAdd) needs to be performed;

[0178] At the 13,600th time slice, write ffn_out 200 to the DDR, with a shape of [1000, d];

[0179] So far, the entire FFN calculation process has been completed, and all results are written to the DDR, with a shape of [200,000, d].

[0180] Currently, there is no report on the optimization work of long text, especially extremely long text (greater than or equal to 200K context) inference on low-power chips at the edge. Through careful design of data slicing, pipelining, and resource utilization, this article converts all complex calculations to on-chip computing, greatly reducing the characteristics of limited bandwidth, limited storage resources, and power consumption sensitivity (DDR access accounts for half of the chip's power consumption) of edge devices. The inference model originally requiring expensive GPUs is optimized and implemented on edge chips. For long text processing, such as 200K, the entire storage requirement has increased by 100 times compared to the original model (2048), and the computational complexity has increased by 10,000 times. Under this huge complexity, the optimization algorithm in this article overcomes this problem. With only a linear increase in computational complexity and no increase in storage resources, it becomes feasible to implement it, which is a very effective method.

[0181] In a specific embodiment, the statistical usage amount of SRAM is 12,386,048 Byte. If there are lower requirements for SRAM, more fine-grained slicing can be performed on the basis of the specific implementation of the present invention.

[0182] Those skilled in the art should understand that the embodiments of the present application can be provided as a method, a system, or a computer program product. Therefore, the present application can take the form of a completely hardware embodiment, a completely software embodiment, or an embodiment combining software and hardware aspects. Moreover, the present application can take the form of a computer program product implemented on one or more computer-usable storage media (including but not limited to disk memory, CD-ROM, optical memory, etc.) containing computer-usable program code. The solutions in the embodiments of the present application can be implemented in various computer languages. For example, object-oriented programming languages such as Java and interpreted scripting languages such as JavaScript, etc.

[0183] The present application is described with reference to the flowcharts and / or block diagrams of methods, apparatuses (systems), and computer program products according to the embodiments of the present application. It should be understood that each flow and / or block in the flowchart and / or block diagram, as well as the combination of flows and / or blocks in the flowchart and / or block diagram, can be implemented by computer program instructions. These computer program instructions can be provided to the processor of a general-purpose computer, a special-purpose computer, an embedded processor, or other programmable data processing devices to generate a machine, so that the instructions executed by the processor of the computer or other programmable data processing devices generate means for implementing the functions specified in Figure 1 one or more flows and / or Figure 1 blocks.

[0184] These computer program instructions can also be stored in a computer-readable memory that can direct a computer or other programmable data processing device to work in a specific manner, so that the instructions stored in the computer-readable memory generate a manufactured article including instruction means, and the instruction means implement the functions specified in Figure 1 one or more flows and / or Figure 1 blocks.

[0185] These computer program instructions can also be loaded onto a computer or other programmable data processing device, so that a series of operation steps are executed on the computer or other programmable device to generate a computer-implemented process, and thus the instructions executed on the computer or other programmable device provide steps for implementing the functions specified in Figure 1 one or more flows and / or Figure 1 blocks.

[0186] Although the preferred embodiments of the present application have been described, those skilled in the art can make additional changes and modifications once they learn the basic creative concepts. Therefore, the appended claims are intended to be construed as including the preferred embodiments and all changes and modifications falling within the scope of the present application.

[0187] Obviously, those skilled in the art can make various changes and modifications to this application without departing from the spirit and scope of this application. Thus, if these modifications and variations of this application fall within the scope of the claims of this application and their equivalent technologies, this application is also intended to include these modifications and variations.

[0188] The protection scope of the present invention is not limited to the above embodiments. Without departing from the spirit and scope of the inventive concept of the present invention, changes and advantages that can be conceived by those skilled in the art are included in the present invention, and the scope of protection is defined by the appended claims.

Claims

1. A method for optimizing the inference of a large model chip pre-filling phase under a long context, characterized in that: The reasoning optimization method comprises the following steps: Step 1: Dynamically slice the input data and model parameters in the pre-filling phase; Step 2: Load and calculate QKV, self-attention, attention_project, and FFN modules in parallel; perform residual connection processing after loading and calculating the attention_project and FFN modules; Step 3: Complete the entire inference calculation.

2. The inference optimization method according to claim 1, characterized in that: In step 1, the input data of each transformer layer is stored in a double data rate synchronous dynamic random access memory DDR, the shape is represented as [s, d], and the slice is of shape [s / n, d]; Dynamically slicing the model parameters includes: in QKV layer processing, multi-head slicing the weight matrices Wq, Wk, Wk according to the multi-head attention mechanism, slicing the shape of the weight from the shape [d, d] to the shape [d, d / n]; in the self-attention layer processing, slicing the matrices Q, K, V; In the attention_project part of the processing, the weight Wo is sliced, and the shape of the weight is sliced ​​from shape [d, d] to shape [d, d / n]; in the FFN module part of the processing, the parallel branch weight is sliced ​​from shape [d, ffn_dim] to shape [d, ffn_dim / n], and the serial branch weight is sliced ​​from shape [ffn_dim, d] to shape [ffn_dim / n, d].

3. The reasoning optimization method according to claim 1, characterized in that: In step 2, when loading and calculating QKV, the following steps are included: Step 2.1.1, sequentially load slices of Wq, Wk, and Wv from the double data rate synchronous dynamic random access memory DDR into the static random access memory SRAM, and read the slice Xi of input data from the DDR into the SRAM; Step 2.1.2, respectively calculate the query vector Qi=Xi*Wq and the position embedding Rope(Qi) of the query vector Q, calculate the key vector Ki=Xi*Wk and the position embedding Rope(Ki) of the key vector K, calculate the value vector Vi=Xi*Wv, and write the results into the Q, K, V queues of DDR for transfer; Step 2.1.3, reuse the SRAM storage space of the weight slices and input data slices, and calculate the other input data slices in DDR in sequence until all input data slices are executed; Step 2.1.4, release the SRAM space occupied by weight slices, intermediate calculation results and input data slices, complete the calculation of Q, K, V and store them in DDR.

4. The reasoning optimization method according to claim 1, characterized in that: In step 2, when loading and calculating self-attention, the following steps are included: Step 2.2.1, load Qi from the Q queue of the double data rate synchronous dynamic random access memory DDR into the SRAM, and perform initialization operations of O, init M, and init L; where O is used to save the output of self-attention, M is used to save the local maximum value, and L is used to calculate the cumulative sum of the softmax denominator; Step 2.2.2: Load K from the K queue of DDR during the O initialization operation j to SRAM; Step 2.2.3, K j After loading is completed, calculate S ij =Q i K j T , and load V from the V queue of DDR while calculating j to SRAM; Step 2.2.4, Update m ij =max(m i，j-1 , rowmax(S ij )), Pij = exp(S ij -m ij ), l ij =exp(m i，j-1 -m ij ) i，j-1 +rowsum(P ij ), O ij =exp(m i，j-1 -m ij ) i，j-1 +P ij V j , complete the single-layer inner attention calculation; and load K from the K queue of DDR while calculating j+1 to SRAM, which reuses K j The SRAM space occupied; Step 2.2.5, perform inner loop according to the number of input slices; Step 2.2.6: Execute outer layer update O i and L i , the result is stored in DDR; Step 2.2.

7. Repeat the above steps to synchronize and load the next Q from the DDR Q queue. i+1 to SRAM and reuse Q i The SRAM space used and the same calculations as above are performed, and all results are stored in DDR; Step 2.2.8: Loop through all the multi-head results in the multi-head attention mechanism and store them in DDR.

5. The reasoning optimization method according to claim 1, characterized in that: In step 2, when loading and calculating the attention projection, the input data is sliced ​​by row and the weight matrix is ​​sliced ​​by column; for each sliced ​​input data, the weight slices are loaded in turn and the matrix multiplication o_out = att_o * Wo is performed, and all the results are concatenated to obtain the output; wherein, the sliced ​​input data is looped in the outer layer and the sliced ​​weight matrix is ​​looped in the inner layer; The input of the current layer is loaded from DDR to SRAM, element-by-element addition is performed, residual connection is performed, the input data after all slices are processed, the attention projection calculation is completed, and the output is written to DDR.

6. The reasoning optimization method according to claim 1, characterized in that: In step 2, when loading and calculating the FFN module, the input data is sliced ​​by row to execute the outer loop, and the weight matrix is ​​sliced ​​by column to execute the inner loop; The FFN includes two parallel weight calculation branches W1 and W3 and one serial weight calculation branch W2. The weights W1 and W3 are sliced ​​by columns, and the weight W2 is sliced ​​by rows. After the attention projection calculation result is written to DDR, ffn_in is used to read it out from DDR to SRAM again. For each iteration of the outer loop, a piece of SRAM is pre-allocated for the memory ffn_out of the full reduction operation and initialized to 0; the inner loop first loads the weight W1, and then calculates O1=silu(ffn_in*W1) in the next time slice, and releases the W1 memory; during the calculation and release process, W3 is loaded in parallel; O3=ffn_in*W3 is calculated in the next time slice, and W3 is released after the calculation is completed; W2 is loaded during the calculation of O3 and the release of W3; O2=Hadamard_product(O1,O3) is calculated in the next time slice, and then out=O2*W2 is executed. During the execution of out=O2*W2, O1 and O3 are released. After out is executed, W2 and O2 are released, and a complete FFN calculation is completed; finally, the output result of the last inner loop is obtained through the full reduction operation of element-by-element addition; In the outer loop, the inner output result and the input data ffn_in of the FFN module are subjected to an element-by-element addition residual connection process ffn_out=ffn_out+ffn_in, and then the result is written to the DDR, and ffn_in is released. The outer loop is repeated until the entire FFN module calculation is completed and the result is written to the DDR.

7. A system for inference optimization of a large model chip pre-filling phase under a long context, characterized in that: The reasoning optimization system includes: an input processing module, a QKV calculation module, a self-attention calculation module, an attention projection module, a feedforward neural network FFN module, and a result output module; The input processing module is used to load input data and perform preprocessing, and store the data in SRAM; The QKV calculation module is used to calculate the Query, Key, and Value vectors, perform position encoding on the Query and Key, and store the results; The self-attention calculation module is used to calculate the attention score and weight to initially obtain the attention output; The attention projection module is used to perform projection calculation on the output of the self-attention calculation module and merge the results to obtain the final attention output; The feed-forward neural network FFN module is used to perform dimension enhancement, element-by-element multiplication and dimension reduction calculations; The result output module is used to output the final calculation result of the decoding stage and release all intermediate calculation resources.

8. Application of the reasoning optimization method as described in any one of claims 1 to 6, or the reasoning optimization system as described in claim 7 in system resource optimization scheduling for long-context large-model reasoning.

9. A hardware system for implementing the reasoning optimization method according to any one of claims 1 to 6, characterized in that: The hardware system includes: a memory and a processor; a computer program is stored in the memory, and when the computer program is executed by the processor, the reasoning optimization method according to any one of claims 1 to 6 is implemented.

10. A computer-readable storage medium having a computer program stored thereon, characterized in that: When the computer program is executed by a processor, the reasoning optimization method according to any one of claims 1 to 6 is implemented.

Images