Large model chip decoding stage reasoning optimization method and system under long context and application

Abstract

The invention discloses a large model chip decoding stage reasoning optimization method under a long context, and the method comprises the following steps: 1, carrying out the dynamic slicing of KV cache data and model parameters in a decoding stage, and enabling the KV cache data and model parameters to be called by a static random access memory (SRAM); step 2, a QKV module, a self-attention module, an attention project module and an FFN module are loaded and calculated in sequence in parallel, and the QKV module, the self-attention module, the attention project module and the FFN module are loaded and calculated in sequence; carrying out residual error connection processing after the loading calculation of the attention project and the FFN module, and carrying out the residual error connection processing after the loading calculation of the attention project and the FFN module; 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 decoding stage of a large model chip under long context. Background Art

[0002] In recent years, the fields of large language models and the like have developed rapidly and achieved excellent results. However, due to the huge number of parameters and computing requirements in large language models, it brings great challenges 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 the requirement of low power consumption, 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 generates the first token after inputting the text. In the generation process, all inputs need to participate in the calculation, and the computing and storage requirements are extremely high. The amount of calculation is quadratic with the input length. The latter is the process of generating a new token by using context information for each newly input token during the inference process. Its computing amount and storage requirements are relatively low. Since the computing amounts and resource demands of the two stages are quite different, the optimization methods are also quite different.

[0004] When deploying large language models on inference chips currently, limited by the shortage of computing resources and on-chip storage 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 and final results to low-speed storage devices. Due to the latency of DDR access, it greatly affects its overall running speed and increases its power consumption at the same time. Some slicing methods attempt to convert data or models into on-chip calculations after tiling, such as flash-attention, which is based on local small modules and only accelerates the attention calculation module. Moreover, it is optimized for the characteristics of GPUs and is not suitable for low-resource and low-power chips on the inference side, and cannot achieve good acceleration effects. 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 decoding stage of a large model chip under long context, which is related to artificial intelligence, AI chip optimization and microarchitecture.

[0006] The optimization method of the present invention mainly targets the decoding stage. In the decoding stage, although the computational complexity is smaller than that in the pre-filling stage, the data access does not decrease significantly. Therefore, it is memory access sensitive, specifically manifested as follows: when a new token is input, projection calculation for this token is performed. Subsequently, the K and V values will be combined with the historical KV cache for calculation. Essentially, it belongs to the multiplication of a vector and a matrix, which is a memory-bound situation. The intermediate calculation results can be stored in the cache to complete the fuse kernel operation, reducing memory access and the corresponding IO overhead. The present invention designs an optimization acceleration and resource scheduling method for the decoding stage. By referring to the tiling idea in flash-attention2 and the data and model slicing parallelism in model parallelization, an optimization scheme for on-chip computing is designed. The input and the model are reasonably sliced to make full use of the efficient memory access advantage of SRAM, converting the main calculations into on-chip calculations, which can greatly reduce the latency caused by memory access, reduce the number of times to access DDR during sequential calculations, and through a carefully designed pipeline, maximize the parallel interleaving coverage of storage and calculation to reduce the total time, power consumption, and latency.

[0007] The present invention provides an inference optimization method for the decoding stage of a large model chip under long context. The inference optimization method includes the following steps:

[0008] Step 1: Dynamically slice the KV cache data and model parameters in the decoding stage for static random access memory (SRAM) call;

[0009] Step 2: Sequentially and parallelly load and calculate the QKV, self-attention, attention_project, and FFN modules; 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, slice the KV cache data according to the sequence length and load the sliced KV cache data into the static random access memory (SRAM);

[0012] Dynamically slicing the model parameters includes: in the QKV processing, slicing the weights Wq, Wk, and Wv according to the multi-head attention mechanism, slicing the shape of the weights from [d, d] to [d, d / n]; in the self-attention part processing, slicing the shape of the KV cache from [T, d / n] to [T / m, d / n]; in the projection calculation processing, slicing the shape of the Wo matrix from [d, d] to [d / n, d]; in the FFN module processing, slicing the parallel weight matrix column-wise, slicing the shape from [d, d] to [d, d / n], and slicing the serial weight matrix row-wise, slicing the shape from [d, d] to [d / n, d].

[0013] In step two, when loading and calculating QKV, the weights Wq, Wk, and Wv in the transformer layer are loaded from the double data rate synchronous dynamic random access memory DDR into the static random access memory SRAM; it includes the following steps:

[0014] Step 2.1.1: Load Wq from DDR into SRAM;

[0015] Step 2.1.2: Load Wk from DDR into SRAM, and simultaneously execute Q = X * Wq serially, RopeQKcode, Qrope, and after Q = X * Wq is executed, release Wq;

[0016] Step 2.1.3: Load Wv from DDR into SRAM, and simultaneously execute K = X * Wk, K rope, releaseQK code, and after K = X * Wk is executed, release Wk;

[0017] Step 2.1.4: Calculate the value vector V = X * Wv, and release Wv after the calculation is completed.

[0018] When loading and calculating self-attention, load the sliced KV cache from the double data rate synchronous dynamic random access memory DDR into the static random access memory SRAM; it includes the following steps:

[0019] Step 2.2.1: Load K 0 Cache, with the shape of [T / m, d / n];

[0020] Step 2.2.2: Load V 0 Cache, with the shape of [T / m, d / n], and execute S 0 in parallel with loading V 0 = QK 0 ’, P 0 = softmax(S0 ); S 0 = QK 0 ’, P 0 = softmax(S 0 ) is executed serially and releases S 0 and K 0 ;

[0021] Step 2.2.3, Load K 1 Cache, with shape [T / m, d / n], executes O 1 in parallel with the loaded K 0 = P 0 V 0 , attn += O 0 ; O 0 = P 0 V 0 , attn += O 0 is executed serially and releases V 0 and O 0 ;

[0022] Step 2.2.4, Repeat the above steps until the self-attention part is completed.

[0023] When loading and calculating attention_project, the sliced Wo matrix is loaded from the double data rate synchronous dynamic random access memory DDR into the static random access memory SRAM; the following steps are included:

[0024] Step 2.3.1, Multiply the first head attn 0 in the self-attention step with the sliced Wo read from DDR, o 0 = attn 0 * Wo, perform the operation O = O + o 0 ;

[0025] Step 2.3.2, Multiply the second head attn 1 in the self-attention step with the sliced Wo read from DDR, o 1 = attn 1 * Wo, perform the operation O = O + o 1 ;

[0026] Step 2.3.3, Repeat the above operation n times to obtain the final result O, initialize it to the value of input, with shape [1, d];

[0027] Step 2.3.4: Add the output result and the input data element by element to complete the residual connection, and finally obtain the calculation result of the attention_project link.

[0028] When loading and calculating the FFN module, load the weights after column slicing and / or row slicing from the double data rate synchronous dynamic random access memory DDR into the static random access memory SRAM; the steps are as follows:

[0029] Step 2.4.1: Load W1 from DDR.

[0030] Step 2.4.2: Load W3 from DDR, calculate o1 = (ffn_in * W1) in parallel, o1 = silu(o1), and release W1 after the calculation is completed.

[0031] Step 2.4.3: Load W2 from DDR, calculate o3 = (ffn_in * W3) in parallel, o2 = Hadamard(o1, o3), and release W3, o1, and o3 after the calculation is completed.

[0032] Step 2.4.4: Execute ffn_out = o2 * W2, perform a full reduction operation of element-by-element addition, and merge the results.

[0033] Step 2.4.5: Add the output result and the input data element by element to complete the residual connection, and finally obtain the calculation result of the FFN link.

[0034] The present invention also discloses an inference optimization system for the decoding stage of a large model 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;

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

[0036] The QKV calculation module is used to calculate Query, Key, and Value vectors, perform position encoding on Query and Key, and store the results.

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

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

[0039] The feed-forward neural network FFN module is used to perform calculations of dimension elevation, element-by-element multiplication, and dimension reduction.

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

[0041] The present invention also discloses an application of the above-mentioned inference optimization method or the above-mentioned inference optimization system in the system resource optimization scheduling of long-context large model inference.

[0042] The present invention also discloses 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.

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

[0044] The beneficial effects of the present invention include: through fine slicing and drawing on the tiling idea of flash-attention2, the entire network is globally processed, and the calculation is sliced into sub-blocks. At the same time, model parallel slicing and data parallel slicing are performed on the weight and input respectively. Without precision loss, the characteristics of the storage structure of the processor are utilized to complete the inference of long texts. With a relatively small SRAM occupancy, the long-text inference is optimized and implemented on edge devices, which is of great help to the design of the processor chip microarchitecture and optimized deployment.

[0045] The present invention can be extended to support contexts of any length without increasing computing resources and storage resources. For a given inference chip, fine slicing scheduling of its model is performed, making full use of the characteristics of the on-chip storage resource (usually on-chip SRAM) with fast read and write speeds (the SRAM read and write speed is usually 1-2 orders of magnitude higher than that of DDR). This reduces the I / O time in the calculation, reduces the access to DDR, improves the computing efficiency, and performs on-chip calculations. Through carefully designed software scheduling and resource allocation, it is possible to support the inference of contexts of any length with extremely low resource occupancy, and can support the efficient inference of contexts of any length on low-power and low-cost AI inference chips. BRIEF DESCRIPTION OF THE DRAWINGS

[0046] 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 following drawings 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.

[0047] Figure 1 This is the complete pipeline diagram of the decoding stage of the present invention.

[0048] Figure 2 This is the pipeline diagram of the Q, K, V calculation of the present invention.

[0049] Figure 3 This is the basic schematic diagram of the flash attention of the present invention.

[0050] Figure 4 This is the pipeline diagram of the self-attention calculation of the present invention.

[0051] Figure 5 This is the pipeline diagram of the attention_project of the present invention.

[0052] Figure 6 This is the pipeline diagram of the FFN module of the present invention. Detailed implementation manners

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

[0054] The explanations of relevant English terms in the present invention are as follows:

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

[0056] decode stage: Decoding stage, referring to the stage where the model generates a new token using context information for each input token during the inference process.

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

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

[0059] Tiling: Tiling, a data slicing processing method.

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

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

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

[0063] attention_project: Attention projection, which performs a projection calculation on the output of the attention mechanism.

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

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

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

[0067] Hadamard_product: It performs an element-by-element multiplication of two matrices, and the two input matrices need to have the same shape.

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

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

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

[0071] The present invention provides a large model inference optimization method 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 model is finely sliced and scheduled to make full use of the fast read and write speed (the SRAM read and write speed is usually 1 - 2 orders of magnitude higher than that of DDR) of on-chip storage resources (usually on-chip SRAM), so as to reduce the I / O time in computing, reduce the access to DDR, improve the computing efficiency, and perform on-chip computing. Through carefully designed software scheduling and resource allocation, it can support inferences of contexts of any length with extremely low resource occupancy, and can support efficient inferences of contexts of any length on low-power and low-cost AI inference chips. In the decoding stage, taking llama2 7B as an example, only 1.06MB of SRAM is required to complete the processing of 200K contexts, and all calculations are completed on on-chip SRAM, reducing the access to DDR, which provides favorable help for low-power chip design and microarchitecture.

[0072] In a specific implementation, the usage statistics of the SRAM is 1,057,280 Bytes.

[0073] The inference of large language models is usually divided into a prefill stage and a decode stage. The computational requirements and resource demands of the two stages vary significantly. In view of the decode stage and its characteristics of relatively small computational and resource demands, the present invention designs a corresponding optimization and acceleration method:

[0074] For the decode stage, the input shape is [b, s, d], where b is the batch number, s is the sequence length, and d is the latent variable dimension. For the sake of convenience, it is assumed that b = 1 and this dimension is ignored. For the decode stage, one token is processed at a time, that is, s = 1, and the input shape is [1, d]. The model has L layers. Since the model structures of each layer are the same (transformer structure), the input, output, and weight data shapes are also the same. Therefore, the optimization scheme can be fully reused. For the sake of convenience, one of the layers is taken as an example for illustration here.

[0075] 1. For the current Transformer layer, its input is the output of the previous Transformer layer, with a shape of [1, d] and relatively small values. During the calculation, the output "output" of the previous layer is fully written to the SRAM and used as the input "input" of the current layer, which is stored in the SRAM. For the three weights Wq, Wk, and Wv, they need to be loaded from the DDR into the SRAM. Each weight has a shape of [d, d], and the weights need to be sliced. The Transformer by default uses the mechanism of multiple - header attention. Although the multiple - head processing is in the inner - product stage, in the present invention, it is introduced in the projection stage. After projecting Wq, Wk, and Wv to obtain Q, K, and V, slice Q, K, and V along the last dimension, i.e., the latent variable dimension. Here, it is chosen to perform multiple - head slicing on the three weights Wq, Wk, and Wk. By performing projection calculations after slicing the weights, the purpose of dividing multiple heads is achieved. The number of heads is n, and the weight shape of each head becomes [d, d / n]. At this time, the weights can be directly placed into the limited SRAM space, so they can be directly loaded from the DDR into the SRAM and the on - chip calculation is completed on the SRAM. In order to control the SRAM usage of the weights in the present invention, the weights are split into a multiple - head form. The original shape of the weights is [d, d], and after splitting into multiple heads, the weight shape is [d, d / n], where n is the number of multiple heads, and it is used as the slicing quantity here. To efficiently reuse the SRAM, the design principle followed is: after completing one load, immediately use the obtained data for calculation. After the calculation is completed, if there is no other calculation depending on it, release the storage resources, and the process of this calculation is parallel to the next data transfer. For the decoding stage, one token is processed at a time, i.e., s = 1. The number of activations generated in the middle is relatively small compared to the number of weights, and the computational amount is also small. This process is a memory - sensitive operation (i.e., memory - bound type, mainly time is spent on data transfer), that is, the time for data transfer can be used to parallelly cover the time for calculation. Since the SRAM resources are limited, a tight pipeline arrangement is made for it. The pipeline for this part is as Figure 1 shown as:

[0076] 1) The input is X. Sequentially load Wq, Wk, and Wv from the DDR into the SRAM, and parallelly complete the calculations of Q = X * Wq, K = X * Wk, and V = X * Wv;

[0077] 2) Execute slicing and loading KVcache from the DDR, and complete S = QK' (where K' is the transpose of K), P = softmax(S), and O = PV calculations;

[0078] 3) Load Wo to complete the projection calculation and residual calculation, O+ = att * Wo;

[0079] 4) Load W1, W2, W3 to complete the calculation of the FFN.

[0080] Specifically, as Figure 2 shown, the execution process of Q, K, V calculation in the present invention is as follows:

[0081] Input X, and the shape of X is [1, d];

[0082] In the first time slice, load Wq from the DDR to the SRAM, and its shape is [d, d / n];

[0083] In the second time slice, load Wk from the DDR to the SRAM, and its shape is [d, d / n]. At the same time, it is parallel to the following three steps: Q = X * Wq (project X using Wq to obtain Q), RopeQKcode (obtain the matrix for rotational encoding), and Q rope (perform Rope encoding on Q). And these three steps (Q = X * Wq, RopeQKcode, Q rope) are executed serially; after Q = X * Wq is executed, release Wq.

[0084] In the third time slice, load Wv from the DDR to the SRAM, and its shape is [d, d / n]. At the same time, parallelly execute the following three steps: K = X * Wk (project X using Wk to obtain K), K rope (perform Rope encoding on K), and release QKcode (release the matrix for rotational encoding); among them, loading Wv from the DDR to the SRAM and (K = X * Wk, K rope, release QK code) are executed in parallel, and (K = X * Wk, K rope, release QK code) are executed serially inside; after K = X * Wk is executed, release Wk.

[0085] In the fourth time slice, execute V = X * Wv, and release Wv after the calculation is completed.

[0086] 2. After the calculation of Q, K, V is completed, if Wq, Wk, Wv are no longer needed, then immediately release Wq, Wk, Wv, as Figure 2As shown in the middle pipeline, for the release on SRAM, only need to mark this section of storage as available, and the time consumption can be almost ignored. In the inference dialogue of the context up to 200K, the maximum length of KVcache is 199999, that is, the loaded KVcache length is 199999. Here, consider the most extreme case, that is, the length that has entered the KVcache before is 199999, the current input seqlen = 1, and the total context length is 200000. For convenience, the length of KVcache is denoted as 200000, denoted as T. Due to the capacity limitation of SRAM, it is also necessary to slice and load KVcache multiple times here. For a header, the total shape of the Kcache and Vcache to be loaded is [T, d / n]. After using the sliced loading method, considering the full utilization of SRAM resources, that is, to reuse the SRAM space used in the previous Q, K, V calculations as little wastefully as possible, the KVCache with a maximum length of 200000 is sliced into 50 parts here, denoted as m, then the shape of each part is [T / m, d / n], that is, [4000, d / n]. Therefore, here will loop and execute 50 small loops through the pipeline method to process the loading of KVcache and related calculations. The introduction of the pipeline can parallelize the GEMM and softmax operation time and data loading in the attention, so as to hide the calculation time. To save the usage amount of SRAM, considering the different data dependencies of the two GEMMs in the attention, the first S = QK', only depends on Q and K, and the second O = PV only depends on V. Therefore, for KVcache, load Kcache and Vcache separately. The process of loading Vcache is parallel with the current two operations of S = QK' and softmax. When loading Kcache, it will be parallel with O = PV in the previous group of sliced data. That is, first load the sliced K_cache from DDR, with the shape of [T / m, d / n]. After its loading is completed, because curretK is appended at the last position of the total K_cache sequence, so the concat (a data append method) will not be performed in the previous 49 loops, and directly use the K segment from Kcache, that is, K = K_cache. After the loading of K_cache is completed and the K data is available, in the next time slice, parallelly execute the loading of the slices of V_cache and part of the calculations in the attention (that is, the first GEMM and softmax calculations. When performing this calculation, Q, K j has been prepared, and the self-correlation operation is S j = QK j T , and the local maximum value is updated m j = max(mj-1 , rowmax(S j ), local softmax denominator correction P j = exp(S j - m j ); l j = exp(m j-1 - m j )l j-1 + rowsum(P j )), These above operations refer to the calculation part of the flash - attention2 method. Different from flash - attention2, in the present invention, Q is not partitioned, and its flash - attention2 principle is as Figure 3 shown.

[0087] Figure 3 In, the meaning of the key parts in the figure:

[0088] Q, K, V matrices:

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

[0090] K (Key) and V (Value): Represent the key and value vectors respectively. 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.

[0091] SRAM operations:

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

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

[0094] Compute Block on SRAM: Perform calculations on SRAM. The calculation block shown in the figure is QK T · M × N, calculating the attention scores.

[0095] 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).

[0096] Pipeline operations:

[0097] Figure 3In it, through the Inner Loop and Outer Loop, it shows how to process data in chunks during the pipelined operation and gradually store the results in the HBM after processing each chunk.

[0098] sm(QK T )V: It means to perform a Softmax operation on the calculation result of QK T and then multiply by the V vector, with the output size of N×d.

[0099] Data chunking and loading:

[0100] Figure 3 It shows how to load large-scale K and V matrices in chunks (sliced by columns) into the SRAM. Since the SRAM space is limited and only a small part of the data can be processed at a time, it is necessary to cut K and V into smaller chunks and gradually load and process them.

[0101] Matrix multiplication and Softmax operation:

[0102] On the SRAM, first calculate the result of QK T , then apply the Softmax operation to this result, and finally multiply by the V matrix to obtain the final attention output sm(QK T )V.

[0103] To efficiently utilize the SRAM, the result of each calculation will be immediately processed and output to the HBM to prevent occupying too much SRAM space.

[0104] Pipelining and loop operations:

[0105] The pipelined processing of the Inner Loop and Outer Loop continuously loads new data chunks while calculating and outputs the processed data. In this way, the utilization rate of hardware resources is maximized and the overhead of data transmission is reduced.

[0106] After the loading of the V_cache slice is completed, the same applies to the K_cache. curretV is concatenated (conca) at the end of the total V_cache sequence. Therefore, the concat operation will not be performed in the previous 49 loops, that is, V = V_cache. If the V data is available, in the next time slice, the data loading will be the next batch of K_cache, and the partial calculation of the previous batch will be performed in parallel at the same time. In this way, the calculation of one batch of the attention module is completed through the pipelined arrangement.

[0107] When the loop executes for the 50th time, it is necessary to perform concat on currentK and currentV calculated using the total tokens in the first step, that is, K = Concat(K_cache, currentK), V = Concat(V_cache, currentV), and then complete the calculation of attention according to the pipeline similar to the above description, and write currentK and currentV into K-cache and V-cache respectively. Thus, the self-attention calculation for the 200k context cache is completed.

[0108] As Figure 4 shown, the execution process of the entire self-attention calculation is as follows:

[0109] In the first time period, load K 0 Cache, with a shape of [T / m, d / n];

[0110] In the second time period, load V 0 Cache, with a shape of [T / m, d / n], and execute S 0 in parallel with the loaded V 0 Cache = QK 0 ’, P 0 = softmax(S 0 ); S 0 = QK 0 ’, P 0 = softmax(S 0 ) is executed serially, and after execution, release S 0 and K 0 ;

[0111] In the third time period, load K 1 Cache, with a shape of [T / m, d / n], and execute O 1 in parallel with the loaded K 0 Cache = P 0 V 0 , attn += O 0 ; O 0 = P 0 V 0 , attn += O 0 is executed serially, and after execution, release V 0 and O 0 ;

[0112] In the fourth time period, load V 1 Cache, with a shape of [T / m, d / n], and execute S 1 in parallel with the loaded V 1= QK 1 ’, P 1 = softmax(S 1 )); S 1 = QK 1 ’, P 1 = softmax(S 1 ) is executed serially, and K 1 and S 1 are released after execution; ....

[0114] Load K into the 99th time period 49 Cache, with shape [T / m - 1, d / n], the current currentK, with shape [1, d / n], after concat operation, K 49 = concat(K 49 Cache, currentK), with shape [T / m, d / n], and execute O 49 in parallel with the loaded K 48 = P 48 V 48 , attn += O 48 ; O 48 = P 48 V 48 , attn += O 48 is executed serially, and V 48 and O 48 are released after execution;

[0115] Load V into the 100th time period 49 Cache, with shape [T / m - 1, d / n], the current currentV, with shape [1, d / n], after concat operation, V 49 = concat(V 49 Cache, currentV), with shape [T / m, d / n], and execute S 49 in parallel with the loaded V 49 = QK 49 ’, P 49 = softmax(S 49 )); S 49 = QK 49 ’, P 49 = softmax(S 49 ) is executed serially, and K 49 and S 49 are released after execution;

[0116] In the 101st time period, execute O serially 49 = P49 V 49 , attn+=O 49 , after the execution is completed, release V 49 and O 49 .

[0117] 3. After completing the calculation of the self-attention part in the previous step, the next step is to perform projection calculation on Wo (some literature also calls it attention_project). In step 2, the output corresponding to each head is recorded as attn, and its shape is [1, d / n], while the shape of the weight Wo is [d, d]. In order to efficiently use SRAM and keep the calculation shape consistent with the previous result, Wo is sliced ​​into n pieces. After slicing, the shape of each piece of Wo is [d / n, d]. Then the projection calculation C_o=attent*Wo is calculated as [1, d / n]*[d / n, d]. After slicing, the data is small and can be calculated on-chip (i.e., on SRAM). The output shape is [1, d]. The multi-head processing in step 2 is equivalent to slicing the input by column, while Wo performs row slicing. Therefore, for this pair of slices, After the slice processing, the matrix calculation result needs to perform an element-wise addition (Element-WiseAdd) reduction (all-reduce) operation, and the number of executions is the number of multiple heads (that is, the number of slices), here is n. At the same time, there is a residual short-cut operation after attention_project, because the input data Xi (whose shape is [1, d]) needs to be element-wise added (Element-WiseAdd) to complete the residual short-cut processing. At this point, the entire calculation of the attention module is completed, and only the final output output is retained, with a shape of [1, d], which will be used as the input of the subsequent FFN module. At this time, other resources have been released, and only the results after the attentionproject projection are saved in the storage, such as Figure 5 As shown, the calculation process of attention_project is as follows:

[0118] The total output of this layer is recorded as O, which is initialized to the input value with a shape of [1, d];

[0119] The input comes from the output of the first head in step 2, attn 0 , whose shape is [1, d / n], read the slice of Wo from DDR, whose shape is [d / n, d], and perform matrix multiplication operation, o 0 =attn 0 *Wo, whose shape is [1,d], executes O=O+o 0Operation;

[0120] The input comes from the output of the second head in step 2, attn 1 , which has a shape of [1, d / n], reads a slice of Wo from the DDR, which has a shape of [d / n, d], and performs a matrix multiplication operation, o 1 = attn 1 * Wo, which has a shape of [1, d], and executes O = O + o 1 Operation;

[0121] Repeat the above operation n times to obtain the final result O, which has a shape of [1, d].

[0122] 4. Next, enter the processing of the FFN module. The FFN module is characterized by a large amount of weight data. Its structure has a parallel two-way weight calculation branch for dimensionality increase, then performs an element-wise product, followed by a serial weight branch for dimensionality decrease, and two matrix module groups. The former parallel two weight matrices can perform column slicing, and the latter serial weight matrix can perform row slicing. After the attention_project module ends in step 3, its output is on the SRAM, with a shape of [1, d], which is used as the input of the FFN module, denoted as ffn_in. At this time, except for ffn_in on the SRAM, the SRAM has been completely released. The shapes of the parallel branches W1 and W3 in the FFN are [d, ffn_hid_d], so W1 and W3 need to be sliced (by column). ffn_hid_d is sliced into fn slices, and each slice has a corresponding shape of [d, d / n]. The number of fn is ffn_hid_d / (d / n), so fn small loops need to be executed. After slicing, the shapes of W1 and W3 are both [d, d / n] and can be put into the SRAM.

[0123] Figure 6 The execution steps for the FFN module are as follows:

[0124] In the first time slice, load the data of W1 (sliced by column), which has a shape of [d, d / n]

[0125] In the second time slice, load the data of W3 (sliced by column), which has a shape of [d, d / n], and simultaneously execute o1 = (ffn_in * W1) in parallel, o1 = silu(o1). After execution, release W1. Here, loading W3 and the following three operations (o1 = (ffn_in * W1), o1 = silu(o1), release W1) are executed in parallel, while the three operations (o1 = (ffn_in * W1), o1 = silu(o1), release W1) are executed serially internally;

[0126] In the third time slice, load the data of W2 (sliced by row) with the shape of [d / n, d]. At the same time, execute o3 = (ffn_in * W3) in parallel, o2 = Hadamard(o1, o3), release W3, release o1 and o3. Here, loading W2 and the following three operations (o3 = (ffn_in * W3), o2 = Hadamard(o1, o3), release W3, release o1 and o3) are executed in parallel, while the three operations of (o3 = (ffn_in * W3), o2 = Hadamard_product(o1, o3), release W3) are executed serially inside; releasing o1 and o3 is parallel to releasing W3;

[0127] In the fourth time slice, execute ffn_out = o2 * W2 with the shapes of [1, d], [1, d / n], and [d / n, d] respectively, and then perform a reduction (all-reduce) operation of element-wise addition (Element-WiseAdd).

[0128] Repeat the above 4 time slices. The number of repetitions is the number of chunks, which is fn here. At the same time, there is also a residual short-cut operation. Since ffn_in (with the shape of [1, d]) also needs to be element-wise added (Element-WiseAdd) to complete the residual short-cut processing. Thus, the entire calculation of FFN is completed, and only the final output with the shape of [1, d] is retained. Thus, the entire calculation of the ffn module is completed, and thus all the pipelining operations in the decode stage are completed.

[0129] 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 sharding, pipeline arrangement, and resource utilization, the present invention transforms all complex calculations into on-chip calculations, 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 is increased by 100 times compared to the original model (2048), and the computational complexity is increased by 100 times. Under this huge complexity, the optimization method of the present invention overcomes this problem. On the premise that the computational complexity only increases linearly and the storage resources do not increase, it becomes feasible to implement it, which is a very effective method.

[0130] Through the processing of this method, all calculations can be converted into on-chip calculations. Taking the llama2 7b model as an example, for a context length of 200K, the SRAM occupancy in the decode stage is controlled at 1.05MB, and only a very small amount of cache is required, which can be completely converted into on-chip calculations, thus enabling the implementation of ultra-long context model inference on low-cost and low-power edge chips, which is of great help for the low-cost use of large models.

[0131] 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 complete hardware embodiment, a complete 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 storage, CD-ROM, optical storage, etc.) that contain 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.

[0132] The present application is described with reference to the flowcharts and / or block diagrams of methods, devices (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, and 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 specified functions in Figure 1 one process or multiple processes and / or blocks Figure 1 one block or multiple blocks.

[0133] 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 implements the specified functions in Figure 1 one process or multiple processes and / or blocks Figure 1 one block or multiple blocks.

[0134] 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 means for implementing the specified functions in Figure 1One or more processes and / or blocks Figure 1 Steps of the functions specified in one or more blocks.

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

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

[0137] 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 appended claims are used as the protection scope.

Claims

1. A method for optimizing the decoding phase inference of a large model chip under a long context, characterized in that: The reasoning optimization method comprises the following steps: Step 1: Dynamically slice the KV cache data and model parameters in the decoding 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 KV cache data is sliced ​​according to the sequence length, and the sliced ​​KV cache data is loaded into a static random access memory SRAM; Dynamically slicing the model parameters includes: in QKV processing, slicing the weights Wq, Wk, and Wv according to the multi-head attention mechanism, and slicing the shape of the weights from shape [d, d] to shape [d, d / n]; in the self-attention part processing, slicing the shape of the KV cache from shape [T, d / n] to shape [T / m, d / n]; in the projection calculation processing, slicing the shape of the Wo matrix from shape [d, d] to shape [d / n, d]; in the FFN module processing, slicing the parallel weight matrix by columns, and the shape is sliced ​​from shape [d, d] to shape [d, d / n], and the serial weight matrix is ​​sliced ​​by rows, and the shape is sliced ​​from shape [d, d] to shape [d / n, d].

3. The reasoning optimization method according to claim 1, characterized in that: In step 2, when loading and calculating QKV, the sliced ​​weights Wq, Wk, and Wv are loaded from the double data rate synchronous dynamic random access memory DDR into the static random access memory SRAM; The steps include: Step 2.1.1, load Wq from DDR into SRAM; Step 2.1.2, load Wk from DDR into SRAM, and serially execute Q=X*Wq, RopeQKcode, Q rope, Q=X*Wq, and then release Wq; Step 2.1.3, load Wv from DDR into SRAM, and serially execute K=X*Wk, K rope, release QKcode, K=X*Wk, and then release Wk; Step 2.1.4, calculate the value vector V = X*Wv, and release Wv after the calculation is completed.

4. The reasoning optimization method according to claim 1, characterized in that: In step 2, when loading and calculating self-attention, the sliced ​​KV cache is loaded from the double data rate synchronous dynamic random access memory DDR into the static random access memory SRAM; including the following steps: Step 2.2.1, load K0Cache, the shape is [T / m,d / n]; Step 2.2.2, load V0Cache, the shape is [T / m, d / n], and execute S0=QK0', P0=softmax(S0) in parallel with loading V0Cache; S0=QK0', P0=softmax(S0) is executed serially, and S0 and K0 are released after execution; Step 2.2.3, load K1Cache, the shape is [T / m, d / n], and execute O0=P0V0, attn+=O0 in parallel with loading K1Cache; O0=P0V0, attn+=O0 is executed serially, and V0 and O0 are released after execution; Step 2.2.4: Repeat the above steps until the self-attention part is completed.

5. The reasoning optimization method according to claim 1, characterized in that: In step 2, when loading and calculating attention_project, the sliced ​​Wo matrix is ​​loaded from the double data rate synchronous dynamic random access memory DDR into the static random access memory SRAM; including the following steps: Step 2.3.1, perform matrix multiplication operation on the first header attn0 in the self-attention step and the slice read into Wo in DDR, o0=attn0*Wo, and execute O=O+o0 operation; Step 2.3.2, perform matrix multiplication operation on the second header attn1 in the self-attention step and the slice of Wo read into DDR, o1=attn1*Wo, and execute O=O+o1 operation; Step 2.3.3, repeat the above operation n times to get the final result O, initialize it to the input value, and the shape is [1, d]; Step 2.3.4: Add the output result to the input data element by element to complete the residual connection, and finally obtain the calculation result of the attention_project link.

6. The reasoning optimization method according to claim 1, characterized in that: In step 2, when loading and calculating the FFN module, the weights after column slicing and / or row slicing are loaded from the double data rate synchronous dynamic random access memory DDR into the static random access memory SRAM; the steps include: Step 2.4.1, load W1 from DDR; Step 2.4.2, load W3 from DDR, calculate o1=(ffn_in*W1) in parallel, o1=silu(o1), and release W1 after calculation; Step 2.4.3, load W2 from DDR, calculate o3 = (ffn_in*W3) and o2 = Hadamard (o1, o3) in parallel, and release W3 and o1, o3 after the calculation is completed; Step 2.4.4, execute ffn_out = o2*W2, perform full reduction operation of element-by-element addition, and merge the results; Step 2.4.5: Add the output result to the input data element by element to complete the residual connection, and finally obtain the calculation result of the FFN link.

7. A large model decoding stage reasoning optimization system under 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 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