Transform decoder calculation optimization method, system and application

Abstract

The invention discloses a calculation optimization method of a transform decoder, which is characterized in that by changing a flash attention / feed forward network, a linear layer calculation process and a segmentation mode, the number of unnecessary ddr memory access times is reduced, the memory occupation of an SRAM (Static Random Access Memory) is reduced, the calculation time consumption of the whole process of the transform decoder is optimized, and the calculation optimization of the transform decoder is realized. The invention further discloses application of the calculation optimization method in machine translation, text generation, speech recognition and the like. The invention further discloses a hardware system and a computer readable storage medium for implementing the calculation optimization method.

Description

Technical Field

[0001] The present invention belongs to the field of large language models and artificial intelligence technology, and relates to a transformer decoder calculation optimization method, system and application. Background Art

[0002] Transformer decoder has two main time-consuming aspects. One aspect is memory access time. The traditional calculation process is to read input data and weights from DDR memory, complete the calculation in the on-chip SRAM, and then write the calculation results to the DDR memory. Due to the limited on-chip SRAM memory, there will be a large number of repeated accesses to the DDR memory, which increases the transformer decoder time consumption. The other aspect is the calculation time consumption. Transformer decoder has many matrices with large weight parameters. In the autoregressive mode, there will also be a large number of repeated calculations.

[0003] The current mainstream optimization methods are as follows:

[0004] 1. Since transformer attention has a lot of repeated calculations in the autoregressive mode, kv cache is proposed to cache k and v to reduce repeated calculations and thus reduce calculation time.

[0005] 2. Flash attention is introduced. By dividing the input into blocks and performing multiple passes on the input blocks, the calculation of small blocks is performed on SRAM. Each small block calculates softmax separately, reducing the read and write operations of DDR. Finally, the overall softmax calculation is compensated and the result is written to DDR.

[0006] It can be seen that flash attention can reduce a lot of unnecessary DDR memory access times and reduce calculation time. It is currently a better optimization method. However, since flash attention only considers the optimization of the input and the block calculation of q, k, and v, it does not consider the combination of kv cache / rotation encoding in attention, nor the feedforward module in the decoder and the operations of the remaining linear layers. Therefore, there is still a lot of room for optimization in the entire transformer decoder calculation. Summary of the invention

[0007] In order to solve the deficiencies in the prior art, an object of the present invention is to provide a calculation optimization method for a transformer decoder.

[0008] The computational optimization method of the present invention reduces the SRAM memory usage and optimizes the computational time of the entire transformer decoder process by changing the computational process and segmentation method of the flash attention / feedforward network and other linear layers while reducing the number of unnecessary DDR accesses, thereby achieving computational optimization of the transformer decoder.

[0009] The flash attention is an optimized attention mechanism, which is usually used to accelerate the calculation process of the Transformer model; the feedforward network refers to a feedforward network in which the connections do not form a loop and the data is transmitted in only one direction.

[0010] Specifically, the method of the present invention aims at solving the problem of redundant memory accesses and calculation operations in the current decoder, integrates the kv cache, optimizes the relevant calculation processes in flash attention and feedforward, and performs multiple effective segmentation calculations to reduce unnecessary memory accesses and calculation time.

[0011] Specifically, the calculation optimization method of the present invention mainly includes the following steps:

[0012] Step 1: Optimize the attention module: First, divide the last dimension d of the input, weights wq, wk, wv, key cache k_cache, and value cache v_cache into i parts, read one part from the DDR each time, reuse it i times, and reduce the memory usage by i times; by reducing the amount of data read from the DDR memory each time, the memory usage is reduced;

[0013] Then the input sequence length seq_len is cut into seq_len parts. In the present invention, the input sequence length is cut into seq_len parts. According to the number of cut parts, each part is calculated by matrix multiplication to obtain Q, K, V, and the rotation code of Q and K is calculated. Then the key cache k_cache and the value cache v_cache are updated, and the local softmax after the cut and block is calculated. After all the cut seq_len parts are calculated, the overall softmax value is updated to obtain the output result of softmax; finally, the first dimension of the weight wo is cut into i parts, and the matrix multiplication operation is performed on the output result of softmax and each part wo, which is executed i times, and the output results of the matrix multiplication of all parts are added to obtain the output result of attention;

[0014] wq: query weight, used to convert the input into a query vector.

[0015] wk: key weight, used to convert the input into a key vector.

[0016] wv: value weight, used to convert the input into a value vector.

[0017] k_cache: key cache, storing key vectors from past time steps, used to speed up consecutive decoding steps.

[0018] v_cache: value cache, which stores value vectors in past time steps, also used to improve decoding efficiency.

[0019] The block operation divides the weight matrix and input sequence into multiple small segments, so that the results can be calculated and cached separately at different time steps, reducing the load of each calculation and memory access.

[0020] Step 2: Optimize the feedforward module: First, divide the last dimension of the weights W1 and W3 of the feedforward network into j parts to reduce the amount of data processed in one calculation, thereby reducing memory usage and possible calculation time; then perform matrix multiplication on the output result of attention and W1, and then obtain O1 through the silu activation function. The above nonlinear transformation process can increase the expressiveness of the model; at the same time, perform matrix multiplication on the output result of attention and W3 to obtain O3, and perform element-by-element multiplication on O1 and O3 to obtain O2. By merging two intermediate results generated by different weights, the feature expression can be enhanced; finally, divide the first dimension of the weight W2 into j parts, perform matrix multiplication on O2 and W2, execute j times, and concatenate the output results of all parts to obtain the output result of the feedforward module;

[0021] Step 3: Repeat the above steps to optimize each layer of the transformer encoder;

[0022] Step 4: Linear layer optimization: Since the last dimension of the linear layer is very large, the last dimension is divided into k parts, and the output result of the feedforward module is matrix multiplied with the linear after division. This is performed k times, and the output results of all parts are concatenated to obtain the output result of the last linear layer.

[0023] For the above four parts, the calculation process is executed completely serially, so after each part is executed, the on-chip SRAM can be fully reused.

[0024] The present invention also provides applications of the above-mentioned computational optimization method in machine translation, text generation, speech recognition, etc.

[0025] The present invention also provides a hardware system for implementing the above method, the hardware system comprising: a memory and a processor; a computer program is stored in the memory, and when the computer program is executed by the processor, the above calculation optimization method is implemented.

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

[0027] The beneficial effects of the present invention include: compared with the prior art, the method of the present invention further divides the data calculation on the basis of flash attention and kvcache caching technology, first saving the kv calculation in the autoregressive mode, and in addition, by dividing the data for calculation, the space of the on-chip SRAM is saved, so that each calculation can be performed on the on-chip SRAM, and there is no need to write the intermediate results into the DDR, and frequently read and write. The calculation results obtained by the attention module will be directly placed on the SRAM for the next module feedforward calculation, and each part is serialized to save the time of writing the output result of each part to the DDR. In addition, since the space of the SRAM is saved as much as possible, the DDR memory access times can reach the theoretical memory access times, and no additional repeated read and write operations are required. Through data segmentation, data is read from the DDR in batches for calculation, so that the memory access time can cover the calculation time of the on-chip SRAM data. Therefore, compared with the existing scheme, the method of the present invention can realize that the transformer decoder calculation time is similar to the time used for the theoretical DDR memory access times. BRIEF DESCRIPTION OF THE DRAWINGS

[0028] In order to more clearly illustrate the embodiments of the present invention or the technical solutions in the prior art, the drawings required for use in the embodiments or the description of the prior art will be briefly introduced below. Obviously, the drawings described below are only some embodiments of the present invention. For those skilled in the art, other drawings can be obtained based on these drawings without paying any creative work.

[0029] Figure 1 Flowchart for optimizing transformer decoder calculations in the present invention. DETAILED DESCRIPTION

[0030] The present invention is further described in detail with reference to the following specific examples and drawings. The process, conditions, experimental methods, etc. for implementing the present invention, except for the contents specifically mentioned below, are all common knowledge and common common sense in the art and are not particularly limited by the present invention.

[0031] Some of the nouns involved in the present invention are described as follows:

[0032] Transformer: A Large Language Model Architecture

[0033] decoder: decoder

[0034] Attention: Attention mechanism module

[0035] ddr:Double Data Rate Synchronous Dynamic Random Access Memory, double data rate synchronous dynamic random access memory

[0036] SRAM: Static random-access memory

[0037] kv cache: kv cache mechanism

[0038] Softmax: Softmax is an activation function that normalizes a numerical vector into a probability distribution vector, and the sum of each probability is 1

[0039] Feedforward network: Feedforward neural network

[0040] linear layer

[0041] seq_len sequence length

[0042] The present invention provides a calculation optimization method for a transformer decoder. The method reduces unnecessary DDR memory access times and SRAM memory occupancy by changing the flashattention / feedforward network and linear layer calculation process and segmentation method, optimizes the calculation time of the entire process of the transformer decoder, and realizes the calculation optimization of the transformer decoder.

[0043] Specifically, the method comprises the following steps:

[0044] Step 1: Optimize the attention module: split the input, weight and cache, split the input sequence length and calculate the Q, K, V vectors and the rotation encoding of the Q, K vectors, update the cache, calculate the local and global softmax values, split the weight wo and perform matrix multiplication and sum of the softmax output to obtain the final attention output;

[0045] Step 2: Optimize the feedforward module: Split weights W1 and W3, perform matrix multiplication with the attention output obtained in step 1 and apply the activation function to generate O1 and O3, and multiply O1 and O3 element by element to obtain O2; split weight W2, perform matrix multiplication with O2, and concatenate the results of each matrix multiplication of weight W2 to form the final feedforward network feedforward output;

[0046] Step 3: Optimize each layer in the transformer decoder using the optimization methods in steps 1 and 2.

[0047] Step 4: Linear layer optimization: Split the last dimension of the linear layer, perform matrix multiplication with the feedforward output, concatenate the results, and complete the calculation optimization.

[0048] Step 1 includes the following sub-steps:

[0049] Step 1.1, divide the input, weight and the last dimension of the cache into i parts, read one part from the DDR memory each time, and reuse it i times; the weight includes the query weight wq, the key weight wk, and the value weight wv; the cache includes the key cache k_cache and the value cache v_cache;

[0050] Step 1.2, divide the input sequence length seq_len into seq_len parts, and calculate the query Q, key K and value V vectors for each part, and calculate the rotation encoding of Q and K;

[0051] Step 1.3, update the key cache K_cache and the value cache V_cache, and calculate the local softmax value for each block;

[0052] Step 1.4: After all blocks are calculated, update the overall softmax value to get the complete softmax output;

[0053] Step 1.5: Divide the first dimension of the weight wo into i parts, perform matrix multiplication on each part with the output of softmax, and add the results of all matrix multiplications to get the final attention output result.

[0054] Step 2 includes the following sub-steps:

[0055] Step 2.1, divide the last dimension of the weights W1 and W3 of the feedforward network into j parts;

[0056] Step 2.2, use the split weight W1 to perform matrix multiplication with the output of the attention module, and obtain O1 through the silu activation function;

[0057] Use the split weight W3 to perform matrix multiplication with the output of the attention module to obtain O3;

[0058] Step 2.3, perform element-by-element multiplication of O1 and O3 to obtain O2;

[0059] Step 2.4, divide the first dimension of weight W2 into j parts, and then perform matrix multiplication of O2 with each part of W2;

[0060] Step 2.5: Concatenate the matrix multiplication output results of each part in step 2.4 to obtain the complete feedforward network output.

[0061] Step 4 includes the following sub-steps:

[0062] Step 4.1, divide the last dimension of the linear layer into k parts;

[0063] Step 4.2: Perform matrix multiplication on each portion obtained in step 4.1 and the feedforward output obtained in step 2, and perform this operation k times in total. Concatenate the k output results to obtain the final linear layer output result.

[0064] Figure 1 The optimized flowchart for transformer decoder calculation is as follows:

[0065] Input load: first load the input data;

[0066] Weight slicing (wa / wk / wv slice): divide the weights of query (Q), key (K), and value (V) into multiple small segments respectively;

[0067] Cache slice (kv cache slice): splits the cache of keys (K) and values ​​(V);

[0068] Input slice: splits the input data into multiple parts for subsequent calculations;

[0069] Q / K / V matmul: Performs matrix multiplication on the split input and weights to compute the query (Q), key (K), and value (V) matrices.

[0070] Rotation encoding (Q / K rotary_encode): Applies rotation encoding to the calculated query (Q) and key (K) matrices;

[0071] Cache update (kv cache update): updates the cache of keys (K) and values ​​(V);

[0072] Flash Attention: performs optimized attention mechanism calculations;

[0073] Weight slicing (wo slice): cut the output weight (wo) into multiple small segments;

[0074] Attention output (attention output matmul / add): perform matrix multiplication on the calculated attention result and the split wo weight, and add the results to get the final attention module output;

[0075] Feedforward network weight W1 slice (ffn w1 slice): split the first set of weights W1 of the feedforward network (ffn) into multiple segments;

[0076] Feedforward network weight W3 slice (ffn w3 slice): Similarly, the third group of weights W3 of the feedforward network is sliced;

[0077] Matrix multiplication and SiLU (matmul / silu): Matrix multiplication of the split W1 and W3 weights with the output of the attention module, and W1 also needs to apply the SiLU activation function;

[0078] Element-wise multiplication (element-wise mul): multiply the results after activation function processing element by element;

[0079] Feedforward network weight W2 slice (ffn w2 slice): Slice the second set of weights W2 of the feedforward network;

[0080] Matrix multiplication (matmul / add): perform matrix multiplication on the output of the previous step and the split W2 weight and add them together;

[0081] Last linear w slice: slice the last linear layer weights;

[0082] Output concatenation (matmul / concat): concatenate each copy of the last dimension of the linear layer with the output of the feedforward network;

[0083] Output: Get the final output of the entire Transformer decoder.

[0084] The following takes the large model llama 7B as an example to actually analyze the full-process calculation optimization of the transformerdecoder (decoder) calculation optimization method of the present invention:

[0085] 1. Linear+attention

[0086] The specific implementation pseudo code process is as follows:

[0087]

[0088] head represents the number of shares wq, wk, wv, k_cache, v_cache are divided into; Tr represents the number of shares Q is divided into, each share is Br*d / 32, Br is the size of each input seqlen share; Tc represents the number of shares K, V are divided into, each share is Bc*d / 32, Bc is the size of each kv_cache seqlen share; Oi, li, mi, Pij are all intermediate results of flash attention calculation; after loading the various divided data, each share Q, K, V is calculated, and then after rotation encoding, the updated Q, K are obtained, and then the kv cache is updated. The last j loop is to calculate the local softmax. After the calculation, the global softmax value is updated by the coefficients of li and mi to obtain the output of attention.

[0089] The calculation formula of Attention is as follows:

[0090]

[0091] Specifically, load input[seq_len,d]: loads an input matrix with a sequence length of seq_len and a dimension of d.

[0092] for head = 1:32: For 32 attention heads, loop to perform the following operations. Here 32 represents the number of attention heads in the model.

[0093] loadwq[d,d / 32],wk[d,d / 32],wv[d,d / 32]: loads weight matrices wq,wk,wv, whose dimensions are [d,d / 32].

[0094] loadk_cache[prev_seq_len,d / 32],v_cache[prev_seq_len,d / 32]: loads the key cache k_cache and value cache v_cache of the previous sequence length, their dimensions are [prev_seq_len,d / 32].

[0095] for i=1:Tr: the loop variable i ranges from 1 to Tr, where Tr represents the length of the target sequence.

[0096] Q = matmul(input, wq): Calculate the query matrix Q by performing matrix multiplication using the input matrix and weight wq.

[0097] Q=rotary_emb(Q,freqs_cis): applies rotary embedding coding to Q.

[0098] K = matmul(input, wk): Perform matrix multiplication using the input matrix and weight wk to calculate the key matrix K.

[0099] K=rotary_emb(K,freqs_cis): applies rotary embedding coding to K.

[0100] k_cache=concat(K,k_cache): concatenates the calculated K with the key cache k_cache.

[0101] v_cache = concat(matmul(input,wv),v_cache) : Performs matrix multiplication on the input matrix and weights wv to calculate the value matrix V and concatenates it with the value cache v_cache.

[0102] write k_cache,v_cache: write the updated key cache k_cache and value cache v_cache.

[0103] init Oi,li,mi: Initialize the intermediate variables Oi,li,mi, which will be used for softmax calculation.

[0104] forj=1:Tc: The loop variable j ranges from 1 to Tc, where Tc represents the number of blocks.

[0105] sij = Q*K^T / sqrt(d / 32): Calculate the dot product of the query matrix Q and the transposed key matrix K, and then divide by the square root of d / 32.

[0106] mij=max(mi,sij): find the maximum value of mi and sij.

[0107] Pij=exp(sij-mij): calculates the exponential of the difference between sij and mij.

[0108] lij=exp(mi-mij)*li+Pij: Calculate the exponential of the difference between mi and mij, multiply by li, and then add Pij.

[0109] Oij=exp(mi-mij)*Oi+Pij*Vj: Calculate the exponential of the difference between mi and mij, multiply by Oi, and add the product of Pij and Vj.

[0110] mi=mij,li=lij,Oi=Oij:Update mi,li,Oi for the calculation of the next time step.

[0111] o=Oi / li: Calculate the final output o, which is Oi divided by li.

[0112] Analysis of SRAM occupation (taking fp16 as an example):

[0113] ax→2*seqlen*d

[0114] bO→2*seqlen*d

[0115] a.

[0116] a.

[0117] a.

[0118] f.Sij: Br *Bc

[0119] g.Pij→Reuse Sij

[0120] Specifically,

[0121] aX->2*seq_len*d: This represents the size of the input matrix X, where seq_len is the length of the sequence and d is the dimension of the model. The factor 2 represents the conversion of bits to bytes. The examples are all in fp16 format, and 16 bits is converted to 2 bytes, so multiply by 2.

[0122] bO->2*seq_len*d: The size of the output matrix O. Similar to the input matrix, seq_len is the length of the sequence, d is the dimension of the model, and the output dimension is the same as the input. The factor 2 represents the conversion of bits to bytes. The examples are all in fp16 format. 16 bits is converted to 2 bytes, so multiply by 2.

[0123] c.wq,wk,wv->2*d^2 / 32*3: This is the size of the query (wq), key (wk), value (wv) weight matrix. The dimension of each weight matrix is ​​d*d. Since the weight is divided into 32 parts (there are 32 attention heads), the size of each part is d^2 / 32; the factor 2 represents the conversion of bits to bytes. The examples are all in fp16 format. 16 bits is converted to bytes, which is 2 bytes, so multiply by 2. The factor 3 represents the total memory occupied by wq, wk, and wv, so multiply by 3.

[0124] d.k_cache,v_cache->2*d / 32*Bc: the size of the key cache (k_cache) and the value cache (v_cache). The factor 2 represents the conversion of bits to bytes. The examples are all in fp16 format. 16 bits is converted to 2 bytes, so multiply by 2.

[0125] eQ,K,V->2*d / 32*Br+2*2*d / 32*Bc: the size of query (Q), key (K), and value (V). 2*d / 32*Br represents the matrix size of each header, the factor 2 represents the conversion of bits to bytes, the examples are all in fp16 format, 16 bits converted to bytes is 2 bytes, so multiply by 2, and 2*2*d / 32*Bc represents twice the size of the cache batch, the two factors 2 represent the first 2 represents the conversion of bits to bytes, the examples are all in fp16 format, 16 bits converted to bytes is 2 bytes, so multiply by 2, the second 2 represents the memory occupied by K and V together.

[0126] f.Sij:Br*Bc: The size of the matrix Sij whose size is Br and Bc.

[0127] g.Pij->Reuse Sij: The size of the probability matrix Pij is the same as the calculated result of Sij.

[0128] Worst case: At this time, seqlen = 1024, prev_seqlen = 1024, and the occupied SRAM is about 22M;

[0129] Autoregressive case: At this time, seqlen = 1, pre_seqlen = 1024, and the occupied SRAM is about 5.5M.

[0130] The worst case refers to the maximum load in sequence processing. Here, seq_len and prev_seqlen are both set to 1024, which means that the length of the currently processed sequence and the length of the previously processed sequence are both the maximum. This will result in the highest SRAM usage because the information of the entire sequence and all previous sequence information must be stored. In this case, the SRAM usage is about 22 megabytes.

[0131] The autoregressive case refers to the case where the model generates only one output at a time in autoregressive generation, and then uses the output as the input for the next time step. Here, seq_len=1 means that the model processes only one output (usually a word or a token) at the current time step, and prev_seqlen=1024 means that the model has generated 1024 outputs so far. In the autoregressive mode, since only one new output is processed each time, the SRAM occupancy is significantly reduced to about 5.5 megabytes compared to the case of processing the entire sequence.

[0132] 2. Generate attention dimension transformation

[0133] Note: i represents the number of parts into which wo is cut

[0134] The specific implementation pseudo code process is as follows:

[0135] for i=1:32

[0136] load wo[d / 32,d]

[0137] attention_o+=matmul(O,wO)

[0138] Here we mainly do the dimension conversion of the attention output, split wo, and update the overall attention output value by accumulating the result of each multiplication of the wo matrix after the split;

[0139] Specifically,

[0140] for i=1:32: defines a loop that iterates from 1 to 32. This loop corresponds to the 32 attention heads in the transformer model;

[0141] load wo[d / 32,d]: In each iteration, load a weight matrix wo of dimension [d / 32,d]. Here d is the dimension of the model.

[0142] attention_o+=matmul(O,wo): Performs matrix multiplication, multiplies the output of softmax by the weight matrix wo, and then accumulates the result to the attention_o variable.

[0143] Analysis of SRAM occupation (taking fp16 as an example):

[0144] aO→2*seqlen*d

[0145] b.attention_o→2*seqlen*d

[0146] c.

[0147] Specifically,

[0148] aO->2*seq_len*d: This represents a matrix O whose size is twice the sequence length (seq_len) multiplied by the model dimension (d). The factor 2 represents the conversion of bits to bytes. The examples are all in fp16 format. 16 bits is converted to 2 bytes, so it is multiplied by 2.

[0149] b.attention_o->2*seq_len*d: This represents the size of the matrix attention_o that stores the output of the attention module. It is the same as O, which is twice the sequence length multiplied by the model dimension. The factor 2 represents the conversion of bits to bytes. The examples are all in fp16 format. 16 bits is converted to 2 bytes, so multiply by 2.

[0150] c.wo->2*d^2 / 32: This represents the size of the weight matrix wo, which is twice the square of the model dimension (d^2) divided by 32. Due to the multi-head attention mechanism in the model, each head corresponds to a part of the weight. The factor 2 represents the conversion of bits to bytes. The examples are all in fp16 format. 16 bits is converted to bytes, which is 2 bytes, so multiply by 2.

[0151] Worst case: At this time, seqlen = 1024, prev_seqlen = 1024, and the occupied SRAM is about 17M;

[0152] Autoregressive case: At this time, seqlen = 1, prev_seqlen = 1024, and the occupied SRAM is about 1M.

[0153] 3. Feedforward network

[0154] Note: i represents the number of parts w1,w3,w2 are divided into

[0155] The specific implementation pseudo code process is as follows:

[0156] for i=1:128

[0157] load w1[d,11008 / 128]

[0158] load w3[d,11008 / 128]

[0159] o1=silu(matmul(attention_o,w1))

[0160] o3=matmul(attention_o,w3)

[0161] o2=o1*o3#shares memory with o1, the operation is dot multiplication

[0162] load w2[11008 / 128,d]#share memory with w1

[0163] ffn_o+=matmul(o2,w2)

[0164] Similarly, w1 and w3 are split first. After splitting, matrix multiplication operation is performed with the output result of the previous attention step, attention_o, to obtain o1, o3, where o1 also needs to be updated to the new o1 through the silu activation function, and then o1 and o3 are multiplied element by element to obtain o2 to obtain deeper information. Finally, by splitting w2 and accumulating the results of each matrix multiplication with o2, the update of the overall ffn output value is completed.

[0165] Specifically,

[0166] for i=1:128:loop 128 times. It means the weight matrix is ​​divided into 128 segments for processing.

[0167] load w1[d,11088 / 128]: loads the weight matrix w1, whose dimensions are the dimensions of the model d and 11088 divided by 128.

[0168] load w3[d,11088 / 128]: Also loads the weight matrix w3, with the same dimension as w1.

[0169] o1 = silu(matmul(attention_o, w1)) : Performs matrix multiplication, multiplies the attention output attention_o by the weight w1, and then processes the result through the SiLU activation function to obtain the output o1.

[0170] o3 = matmul(attention_o, w3) : performs another matrix multiplication, multiplying attention_o by weight w3 to get output o3.

[0171] o2=o1*o3: multiply o1 (the result after the activation function) and o3.

[0172] load w2[11088 / 128,d]: Loads another weight matrix w2 whose dimensions are 11088 divided by 128 and the dimension d of the model.

[0173] ffn_o+=matmul(o2,w2): Perform the last matrix multiplication, multiply o2 by weight w2, and add the result to ffn_o. Get the final output of the feedforward network.

[0174] Analysis of SRAM occupation (taking fp16 as an example):

[0175] a.attention_o→2*seqlen*d

[0176] b.ffn_o-2*seqlen*d

[0177] a.

[0178] a.

[0179] Specifically,

[0180] a.attention_o->2*seq_len*d: indicates the storage space required for the output of the attention module (attention_o), where seq_len is the length of the sequence and d is the dimension. The factor 2 indicates the conversion of bits to bytes. The examples are all in fp16 format, and 16 bits is converted to 2 bytes, so multiply by 2.

[0181] b.ffn_o->2*seq_len*d: indicates the storage space required for the feedforward network (FFN) output (ffn_o), which is calculated in the same way as the attention module output. The factor 2 indicates the conversion of bits to bytes. The examples are all in fp16 format, and 16 bits is converted to 2 bytes, so multiply by 2.

[0182] c.w1,w3->2*2*11088 / 128*d: represents the storage requirements of weight matrices w1 and w3. 11088 / 128 means that each weight matrix is ​​divided into 128 parts, each part has 11088 / 128 elements. The first factor 2 represents the conversion of bits to bytes. The examples are all in fp16 format. 16 bits is converted to bytes, which is 2 bytes, so multiply by 2. The second factor 2 represents the memory occupied by w1 and w3 together.

[0183] d.o1,o3->2*2*seq_len*11088 / 128*d: represents the storage requirement of the intermediate results o1 and o3. The calculation method is similar to the weight matrix, but it is related to the sequence length seq_len. The first factor 2 represents the conversion of bits to bytes. The examples are all in fp16 format. 16 bits are converted to bytes, which is 2 bytes, so multiply by 2. The second factor 2 represents the memory occupied by o1 and o3 together.

[0184] Worst case: At this time, seqlen = 1024, prev_seqlen = 1024, and the occupied SRAM is about 19.75M;

[0185] Autoregressive case: At this time, seqlen = 1, prev_seqlen = 1024, and the occupied SRAM is about 1.25M.

[0186] 4.Last linear

[0187] Note: i represents the number of parts of the last linear w.

[0188] The specific implementation pseudo code process is as follows:

[0189] for i=1:256

[0190] load w[d,32000 / 256]

[0191] output=concat(op,matmul(ffn_o,w))

[0192] After executing all the above steps, the last layer of linear operation needs to be performed. This layer of linear dimension is large and needs to be split separately. The w small blocks are matrix multiplied with the ffn result of the previous step each time. Finally, the overall data is updated through each splicing.

[0193] Specifically,

[0194] for i=1:256: defines a loop that iterates 256 times.

[0195] load w[d,32000 / 256]: In each iteration, load a weight matrix w whose dimensions are the dimensions of the model d multiplied by 32000 divided by 256. This means that the weight matrix is ​​divided into 256 small pieces, and one piece is loaded in each iteration.

[0196] output = concat(op, matmul(ffn_o, w)): Perform matrix multiplication on the feedforward network output ffn_o and the loaded weight w, and then concatenate the result of this multiplication with a previous output op to update the final output.

[0197] Analysis of SRAM occupation (taking fp16 as an example):

[0198] οw->2*d*32000 / 256

[0199] οoutput->2*seqlen*32000

[0200] Specifically,

[0201] w->2*d*32000 / 256: represents the storage requirement of the weight matrix w. d is the dimension of the model, 32000 is the total number of elements in the weight matrix, and 256 means that the weight matrix is ​​divided into 256 segments. The factor 2 represents the conversion of bits to bytes. The examples are all in fp16 format. 16 bits is converted to 2 bytes, so multiply by 2.

[0202] output->2*seq_len*32000: indicates the storage requirement of the output matrix output. seq_len is the length of the sequence, and 32000 indicates the number of features of each sequence element in the output matrix. The factor 2 indicates the conversion of bits to bytes. The examples are all in fp16 format. 16 bits is converted to 2 bytes, so multiply by 2.

[0203] Whether it is the worst case seqlen = 1024 or the autoregressive seqlen = 1, the number of tokens generated is always 1. Therefore, the calculation of SRAM occupancy only considers seqlen = 1, which is about 1.04M.

[0204] Since the above steps are executed completely serially, SRAM can be fully reused. Therefore, in the worst case, the maximum SRAM occupancy of llama 7B is about 22M. In this solution, through data segmentation, data is read from DDR in batches for calculation, so that each memory access time can cover the calculation time of the on-chip SRAM data. After llama is optimized in this solution, the calculation time of the entire process is close to the time of the theoretical DDR memory access times. Finally, the theoretical DDR memory access times are analyzed, and the actual time of the entire llama process is the time it takes for the actual memory bandwidth of the hardware platform to handle these DDR memory access times.

[0205] Analysis of DDR theoretical memory access times:

[0206] Transformer decoder 1 layer memory access times (taking fp16 as an example)

[0207] a.loadwq,wk,wv→2d 2 *3

[0208] b.loadk_cache,v_cache→2*d*prevseqen (only used in autoregressive mode)

[0209] c.writek_cache,v_cache→2*d*seqlen

[0210] d.loadwo→2*d 2

[0211] e.load w1,w2,w3→2d*11008*3

[0212] Specifically,

[0213] a.load wq,wk,wv->2*d^2*3: loads the weight matrices of query (wq), key (wk), and value (wv). The size of each matrix is ​​the square of the model dimension d. The factor 2 represents the conversion of bits to bytes. The examples are all in fp16 format. 16 bits is converted to 2 bytes, so multiply by 2. The factor 3 represents the memory occupied by wq, wk, and wv together.

[0214] b.load k_cache,v_cache->2*d*prev_seq_len (only used in autoregressive mode): load key cache (k_cache) and value cache (v_cache), the size is the model dimension d multiplied by the length of the previous sequence prev_seq_len, the factor 2 represents the conversion of bits to bytes, the examples are all in fp16 format, 16 bits converted to bytes is 2 bytes, so multiply by 2. At the same time, these caches are only used in the autoregressive decoding process of the model.

[0215] c.write k_cache,v_cache->2*d*seq_len: writes the updated key and value cache. The size is the same as the cache loaded in b, but the current sequence length seq_len is used here. The factor 2 represents the conversion of bits to bytes. The examples are all in fp16 format. 16 bits converted to bytes is 2 bytes, so multiply by 2.

[0216] d.load wo->2*d^2: loads the output weight matrix wo, whose size is the square of the model dimension d. The factor 2 represents the conversion of bits to bytes. The examples are all in fp16 format. 16 bits is converted to 2 bytes, so multiply by 2.

[0217] e.loadw1,w2,w3->2*d*11008*3: loads the weight matrices w1, w2, and w3 of the feedforward network. The size of each weight matrix is ​​the model dimension d multiplied by 11008. The factor 2 represents the conversion of bits to bytes. The examples are all in fp16 format. 16 bits is converted to 2 bytes, so multiply by 2. The factor 3 represents the memory occupied by w1, w2, and w3 together.

[0218] Total memory access count (taking llama 7B as an example):

[0219] Autoregressive mode: (a+b+c+d+e)*layers=13488881664;

[0220] Non-autoregressive mode: (a+c+d+e)*layers=13220446208.

[0221] The protection content of the present invention is not limited to the above embodiments. Without departing from the spirit and scope of the present invention, changes and advantages that can be thought of by those skilled in the art are included in the present invention and are protected by the attached claims.

Claims

1. A computational optimization method for a transformer decoder, characterized in that: The method reduces unnecessary DDR memory access times and SRAM memory occupancy by changing the flashattention / feedforward network and linear layer calculation process and segmentation method, optimizes the calculation time of the entire process of the transformer decoder, and realizes the calculation optimization of the transformer decoder.

2. The calculation optimization method according to claim 1, characterized in that: The method comprises the following steps: Step 1: Optimize the attention module: split the input, weight and cache, split the input sequence length and calculate the Q, K, V vectors and the rotation encoding of the Q, K vectors, update the cache, calculate the local and global softmax values, split the weight wo and perform matrix multiplication and sum of the softmax output to obtain the final attention output; Step 2: Optimize the feedforward module: Split weights W1 and W3, perform matrix multiplication with the attention output obtained in step 1 and apply the activation function to generate O1 and O3, and multiply O1 and O3 element by element to obtain O2; split weight W2, perform matrix multiplication with O2, and concatenate the results of each matrix multiplication of weight W2 to form the final feedforward network feedforward output; Step 3: Optimize each layer in the transformer decoder using the optimization methods in steps 1 and 2. Step 4: Linear layer optimization: Split the last dimension of the linear layer, perform matrix multiplication with the feedforward output, concatenate the results, and complete the calculation optimization.

3. The calculation optimization method according to claim 2, characterized in that: Step 1 includes the following sub-steps: Step 1.1, divide the input, weight and the last dimension of the cache into i parts, read one part from the DDR memory each time, and reuse it i times; the weight includes the query weight wq, the key weight wk, and the value weight wv; the cache includes the key cache k_cache and the value cache v_cache; Step 1.2, divide the input sequence length seq_len into seq_len parts, and calculate the query Q, key K and value V vectors for each part, and calculate the rotation encoding of Q and K; Step 1.3, update the key cache K_cache and the value cache V_cache, and calculate the local softmax value for each block; Step 1.4: After all blocks are calculated, update the overall softmax value to get the complete softmax output; Step 1.5: Divide the first dimension of the weight wo into i parts, perform matrix multiplication on each part with the output of softmax, and add the results of all matrix multiplications to get the final attention output result.

4. The calculation optimization method according to claim 2, characterized in that: Step 2 includes the following sub-steps: Step 2.1, divide the last dimension of the weights W1 and W3 of the feedforward network into j parts; Step 2.2, use the split weight W1 to perform matrix multiplication with the output of the attention module, and obtain O1 through the silu activation function; Use the split weight W3 to perform matrix multiplication with the output of the attention module to obtain O3; Step 2.3, perform element-by-element multiplication of O1 and O3 to obtain O2; Step 2.4, divide the first dimension of weight W2 into j parts, and then perform matrix multiplication of O2 with each part of W2; Step 2.5: Concatenate the matrix multiplication output results of each part in step 2.4 to obtain the complete feedforward network output.

5. The calculation optimization method according to claim 2, characterized in that: Step 4 includes the following sub-steps: Step 4.1, divide the last dimension of the linear layer into k parts; Step 4.2: Perform matrix multiplication on each portion obtained in step 4.1 and the feedforward output obtained in step 2, and perform this operation k times in total. Concatenate the k output results to obtain the final linear layer output result.

6. Application of the computational optimization method according to any one of claims 1 to 5 in machine translation, text generation, and speech recognition.

7. A hardware system for implementing the computational optimization method according to any one of claims 1 to 5, 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 computational optimization method according to any one of claims 1 to 5 is implemented.

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

Images