Cache optimization method and device in large model reasoning

Abstract

The embodiment of the invention provides a cache optimization method in large model reasoning, which comprises the following steps: in a pre-filling stage, cache operation is carried out layer by layer for a plurality of attention layers in a large model, and the cache operation for any ith layer comprises the step of acquiring a target attention matrix of the ith layer. And determining a first index value and a second index value according to the distribution of the row data and the column data of the target attention matrix. And determining an ith preference score corresponding to the ith layer according to the first index value and the second index value. And according to the ith preference score, determining a target cache region allocated for the ith layer in the total cache region, and storing attention data of a target character in the input text in the target cache region. And according to the ith preference score, updating a previous cache region of each layer before the ith layer, and updating attention data of characters stored in the previous cache region.

Description

Technical Field

[0001] One or more embodiments of the present specification relate to the field of computer technology, and in particular, to a cache optimization method and device in large model reasoning. Background Art

[0002] Large Language Model (LLM), or LLM for short, is an AI model with a large number of parameters that specializes in processing and generating natural language text. They are trained on massive amounts of text data, usually based on the Transformers architecture, have multiple self-attention layers, and contain billions or even hundreds of billions of parameters. The more these layers, the better the model is usually able to understand and generate complex text structures.

[0003] At present, the ability of large models to process long texts has been greatly improved. For example, some large models can process more than 128k characters (tokens). However, as the length of the text increases, the amount of computation of the attention layer of the large model will increase significantly. A common strategy is to store the intermediate results of the keys and values ​​calculated for the attention layer of the large model during the reasoning process of the large model, also known as the key value cache (KV Cache). This can significantly reduce the computational complexity because the large model does not have to recalculate the keys and values ​​of all previous characters when generating new characters, but can efficiently generate new characters by referencing the information in the cache. However, due to the memory limitations of hardware devices such as GPUs, it is important to optimize the cache in large model reasoning. Summary of the invention

[0004] One or more embodiments of the present specification describe a cache optimization method and device in large model reasoning, which can effectively perform key-value caching in limited memory resources, thereby ensuring the generation quality of the large model.

[0005] In a first aspect, a cache optimization method in large model reasoning is provided, comprising: in a pre-filling stage, for multiple attention layers in the large model, cache operations are performed layer by layer, wherein the cache operation for any i-th layer includes:

[0006] Get the target attention matrix of the i-th layer, which contains the attention coefficients between characters in the input text calculated according to the parameters of the i-th layer;

[0007] Determine a first index value representing the spatial dispersion according to the distribution of the row data of the target attention matrix; and determine a second index value representing the temporal deviation according to the distribution of the column data of the target attention matrix;

[0008] Determine an i-th preference score corresponding to the i-th layer according to the first index value and the second index value;

[0009] According to the i-th preference score, determine a target cache area allocated for the i-th layer in the total cache area, and store the attention data of the target character in the input text therein;

[0010] According to the i-th preference score, the previous cache areas of each layer before the i-th layer are updated, and the attention data of the characters stored therein are updated.

[0011] In a second aspect, a cache optimization device is provided, comprising:

[0012] An acquisition unit, used for acquiring a target attention matrix of the i-th layer, which includes attention coefficients between characters in the input text calculated according to the parameters of the i-th layer;

[0013] A determination unit, configured to determine a first index value representing a spatial dispersion according to a distribution of row data of the target attention matrix; and to determine a second index value representing a temporal deviation according to a distribution of column data of the target attention matrix;

[0014] The determining unit is further configured to determine an i-th preference score corresponding to the i-th layer according to the first index value and the second index value;

[0015] a storage unit, for determining a target cache area allocated for the i-th layer in the total cache area according to the i-th preference score, and storing the attention data of the target character in the input text therein;

[0016] An updating unit is used to update the previous cache areas of each layer before the i-th layer according to the i-th preference score, and to update the attention data of the characters stored therein.

[0017] According to a third aspect, a computer-readable storage medium is provided, on which a computer program is stored. When the computer program is executed in a computer, the computer is caused to execute the method of the first aspect.

[0018] In a fourth aspect, a computing device is provided, including a memory and a processor, wherein the memory stores executable code, and when the processor executes the executable code, the method of the first aspect is implemented.

[0019] One or more embodiments of the present specification provide a cache optimization method and device for large model inference. In the pre-filling stage, cache operations are performed layer by layer for multiple attention layers in the large model. Specifically, for the current attention layer, the attention layer's preference for cache is dynamically evaluated by considering the changes in the attention coefficient of the attention layer in the spatial and temporal dimensions, and the corresponding cache is allocated to the attention layer based on the preference, and the caches of the previous attention layers are updated in a cascade manner. In this way, limited memory resources can be better and more fully utilized, while hardly affecting the generation quality of the large model. BRIEF DESCRIPTION OF THE DRAWINGS

[0020] In order to more clearly illustrate the technical solutions of the embodiments of this specification, the drawings required for use in the description of the embodiments will be briefly introduced below. Obviously, the drawings described below are only some embodiments of this specification. For ordinary technicians in this field, other drawings can be obtained based on these drawings without paying any creative work.

[0021] Figure 1 A comparison chart showing the effects of the traditional key-value caching method and this solution is shown;

[0022] Figure 2 A schematic diagram of an implementation scenario of an embodiment disclosed in this specification;

[0023] Figure 3 A flow chart of a cache optimization method in large model reasoning according to an embodiment of the present specification is shown;

[0024] Figure 4 A schematic diagram of a target attention matrix in one example of this specification is shown;

[0025] Figure 5 A schematic diagram of a cache optimization method in an example of this specification is shown;

[0026] Figure 6 A schematic diagram of a cache optimization device according to an embodiment of the present specification is shown. DETAILED DESCRIPTION

[0027] The solution provided in this specification is described below in conjunction with the accompanying drawings.

[0028] As mentioned earlier, in order to improve the efficiency of character generation for large models, key-value caching is required.

[0029] Currently, key-value caching is mainly performed in the following two ways:

[0030] First, for each attention layer in the large model, a cache area of ​​the same size is allocated, and the key-value pairs of characters are retained according to a certain predetermined size. However, this solution does not take into account the heterogeneity of the model and each attention layer within the model, and thus cannot effectively utilize cache resources.

[0031] Second, for each attention layer of the large model, the corresponding cache area is allocated in a fixed ratio according to the pre-observation rules. This method is only applicable to specific models and cannot be generalized to multiple models.

[0032] To this end, this solution proposes that, in the pre-filling stage, cache operations are performed layer by layer for multiple attention layers in the large model. Specifically, for the current attention layer, the preference of the attention layer for cache is dynamically evaluated by considering the changes in the attention coefficient of the attention layer in the spatial and temporal dimensions, and the corresponding cache is allocated to the attention layer based on the preference, and the caches of the previous attention layers are updated in a cascade manner. In this way, limited memory resources can be better and more fully utilized, while hardly affecting the generation quality of the large model.

[0033] Figure 1 A comparison chart showing the effects of the traditional key-value caching method and this solution is shown. Figure 1 In the figure, two traditional key-value caching methods are shown above. The left part is a key-value caching method based on a uniform allocation ratio, that is, for any input text, the cache allocation ratio of the four attention layers of the large model is 0.25. The right part is a key-value caching method based on a fixed allocation mode, that is, for any input text, the cache allocation ratio of the four attention layers of the large model is fixed to: 0.50, 0.34, 0.18 and 0.02.

[0034] In this solution, for input text A or input text B, two index values ​​corresponding to each of the four attention layers of the large model will be determined based on it, wherein one index value represents the degree of dispersion of the attention coefficient of the attention layer in space, and the other index value represents the degree of deviation of the attention coefficient of the attention layer in time. Then, based on the two index values ​​corresponding to each attention layer, the corresponding layer preference can be determined, and the corresponding cache can be allocated to each attention layer of the large model based on the layer preference. Among them, for input text A, the cache allocation ratios of the four attention layers of the large model can be: 0.25, 0.19, 0.27 and 0.29. For input text B, the cache allocation ratios of the four attention layers of the large model can be: 0.31, 0.11, 0.35 and 0.23.

[0035] It should be understood that Figure 1This is just an exemplary description. In practice, the number of attention layers included in the large model may be more than 4 layers, for example, 10 layers or even more, and this specification does not limit this. In addition, the cache allocation ratio corresponding to each attention layer is not limited to the above values.

[0036] It can be seen that in this scheme, the cache allocation ratio between different layers can be adjusted according to the layer preference to adapt to various contexts and a given memory budget, thereby making better and more effective use of limited memory resources while hardly affecting the generation quality.

[0037] Figure 2 A schematic diagram of an implementation scenario of an embodiment disclosed in this specification. Figure 2 In , the reasoning process of the large model includes a pre-filling phase and a decoding phase.

[0038] In the pre-filling stage, the cache operation is performed layer by layer for multiple attention layers in the large model. For any i-th layer, firstly, based on the distribution of the attention coefficient in the target attention matrix corresponding to the i-th layer, two index values ​​representing the spatial dispersion and the temporal deviation are determined respectively. Then, based on the two index values, the preference score corresponding to the i-th layer is determined. Finally, based on the preference score, the corresponding cache area is determined to store the attention data of the target characters in the input text, and based on the preference score, the previous cache areas of each layer before the i-th layer are updated, and the attention data of the characters stored therein are updated.

[0039] In the decoding stage, the generated characters for the input text are output based on the attention data of the characters stored in the cache areas corresponding to each of the multiple attention layers.

[0040] It should be understood that, in practice, the above decoding stage will be iteratively executed multiple times to output multiple generated characters.

[0041] As mentioned above, in this solution, during the pre-filling phase, cache operations are performed layer by layer for multiple attention layers in the large model. Since the cache operations for each layer are similar, the corresponding cache operations are described below using any i-th layer as an example.

[0042] Figure 3 The flowchart of the cache optimization method in large model reasoning according to one embodiment of the present specification is shown. The method can be executed by any device, equipment, platform, or device cluster with computing and processing capabilities. Figure 3 As shown, the method may include the following steps:

[0043] Step S302, obtaining the target attention matrix of the i-th layer, which contains the attention coefficients between characters in the input text calculated according to the parameters of the i-th layer.

[0044] Generally speaking, for each attention layer of the large model, based on the character representation output by the previous attention layer for the characters in the input text, and the parameters of the attention layer (i.e., the key matrix, vector matrix, and value matrix), the attention matrix corresponding to the attention layer can be calculated, which contains the attention coefficients between the characters in the input text.

[0045] Taking the input text length as S (i.e. containing S characters) as an example, the size of the above attention matrix is ​​S×S, which can be expressed as follows:

[0046]

[0047] In formula 1, D is the dimension of character representation, Q and K are query vector and key vector respectively, where Q = XW Q , K = XW K , where W is the character representation output by the previous attention layer for the characters in the input text, W Q and W K are the query matrix and the key matrix respectively.

[0048] It should be understood that the rows from left to right (or the columns from top to bottom) in the above-mentioned attention matrix of size S×S (hereinafter also referred to as the original attention matrix) correspond to the characters from front to back in the input text respectively.

[0049] In this solution, the target attention matrix of the i-th layer mentioned above may refer to a submatrix of the original attention matrix corresponding to the i-th layer within a preset observation window, wherein the size of the observation window may be S W ×S, can also be S W ×(SS W ).

[0050] Among them, the size of the observation window is S W ×S, the target attention matrix can be expressed as: i [-S w :,:], which can be Figure 4 As shown on the left (filled with slashes), it contains the back S in the original attention matrix W rows and a total of S columns. The size of the observation window is S w ×(SS W ), the target attention matrix can be expressed as: A i [-S W :,:-S W ], which can be Figure 4 As shown on the right (filled with slashes), it contains the back S in the original attention matrix W Rows and 1st to SSW List.

[0051] In summary, the above target attention matrix will contain the following S W It should be understood that the following W The rows correspond to the S starting from the end of the input text. W characters (also called the most recent characters), similarly, the last S in the original attention matrix W The columns also correspond to the most recent characters. That is, in the target attention matrix: A i [-S W :,:-S W ], the target attention matrix does not contain the columns corresponding to the most recent characters.

[0052] Of course, in practice, the original attention matrix corresponding to the i-th layer can also be directly used as the target attention matrix corresponding to the i-th layer, and this specification does not limit this.

[0053] In addition, in practice, since a multi-head attention mechanism is usually adopted, when a multi-head attention mechanism is adopted, the same observation window can be used to select the corresponding target attention matrix from the original attention matrix corresponding to each attention head, and then the selected multiple target attention matrices are integrated (for example, averaged) to obtain the final target attention matrix, and then the following steps are performed based on the final target attention matrix.

[0054] Step S304, determining the index value E representing the spatial dispersion according to the distribution of the row data of the target attention matrix; and determining the index value V representing the temporal deviation according to the distribution of the column data of the target attention matrix.

[0055] Specifically, for each target row in the target attention matrix, the corresponding information entropy is calculated, and then the information entropy is summed to obtain the above index value E. For each target column in the target attention matrix, the corresponding variance is calculated, and then the variance is summed to obtain the index value V.

[0056] It should be understood that the above is only an example of how to calculate the index value E and the index value V. In practice, other operations can be performed on each target row or each target column in the target attention matrix to obtain the above two index values. For example, when calculating the index value V, the variance can be replaced by the covariance or standard deviation.

[0057] Step S306: Determine the i-th preference score corresponding to the i-th layer according to the index value E and the index value V.

[0058] In one example, the product of the index value E to the power of T1 and the index value V to the power of T2 can be determined as the i-th preference score, where T1 and T2 are temperature parameters, which are respectively used to control the influence of the index value E and the index value V on the i-th preference score.

[0059] For example, the i-th preference score can be calculated according to the following formula:

[0060]

[0061] Among them, P i is the i-th preference score, A i [-S W :,:-S W ] is the target attention matrix, the E() function represents first finding the information entropy of each target row and then summing them up, the V() function represents first finding the variance of each column and then summing them up, T1 and T2 are temperature parameters.

[0062] It should be noted that the importance of attention distraction and offset may vary under different models and cache limitations. This solution can dynamically adjust the importance of attention distraction and offset by using the above-mentioned temperature parameters, so that this solution can flexibly adapt to various scenarios.

[0063] Of course, in practice, the product of the index value E and the index value V can be directly determined as the i-th preference score corresponding to the i-th layer, or other operations except power operation can be performed on the index value E and the index value V respectively to determine the i-th preference score. As long as the influence of the two index values ​​on the i-th preference score can be flexibly adjusted, this manual does not limit this.

[0064] It should be understood that, based on the above-mentioned index value E and index value V, when calculating the i-th attention score, the cache requirements of each layer are actually evaluated dynamically by considering the attention in the spatial and temporal dimensions.

[0065] It should also be understood that since the above-mentioned index value E and index value V are calculated based on the target attention matrix, and the target attention matrix will change with different input texts, this scheme will calculate different preference scores for different input texts, thereby ensuring adaptability.

[0066] Step S308, determining the target cache area allocated to the i-th layer in the total cache area according to the i-th preference score, and storing the attention data of the target characters in the input text therein.

[0067] Specifically, the preference scores corresponding to the layers before the i-th layer can be obtained, and the obtained preference scores and the i-th preference scores can be summed to obtain the sum result. That is, the i preference scores corresponding to the first i layers are summed. Afterwards, the ratio of the i-th preference score to the above sum result can be determined as the cache allocation ratio corresponding to the i-th layer, and the product of the cache allocation ratio and the preset total number of characters can be determined as the target number of characters corresponding to the i-th layer. Finally, according to the target number of characters corresponding to the i-th layer, the target cache area allocated to the i-th layer in the total cache area is determined.

[0068] In one example, the target number of characters may be determined according to the following formula:

[0069]

[0070] Among them, B i is the number of target characters corresponding to the i-th layer, P i is the i-th preference score corresponding to the i-th layer, P k is the kth layer before the i-th layer, is the cache allocation ratio corresponding to the i-th layer, B total The total number of characters is preset.

[0071] More specifically, the cache occupancy corresponding to a single character (i.e., the occupancy of a key-value pair of a single character) can be preset, and then the target number of characters B i By multiplying the cache occupancy corresponding to a single character, the target occupancy corresponding to the i-th layer can be obtained. Afterwards, an area matching the target occupancy can be allocated in the total cache area as the above target cache area.

[0072] It should be noted that since the above-mentioned target number of characters is calculated based on the i-th preference score, it can also be considered that the cache area is allocated to the i-th layer based on the i-th preference score. Since the i-th preference score will be adjusted accordingly with the different input texts, this solution actually implements a preference-priority adaptive allocation strategy.

[0073] Of course, in practice, the above-mentioned preset total number of characters can also be replaced by a preset total cache occupancy, so that after determining the cache allocation ratio corresponding to the i-th layer, the determined cache allocation ratio can be directly multiplied by the preset total cache occupancy to obtain the target occupancy corresponding to the i-th layer.

[0074] Finally, since only the attention data of the target characters are stored in the target cache area corresponding to the i-th layer, a removal operation is actually performed on the i-th layer (i.e., the attention data of non-target characters in the input text are removed).

[0075] The following describes the process of determining the target character stored in the target buffer area.

[0076] This scheme determines the target character based on the above target attention matrix, and as mentioned above, the target attention matrix can be A i [-S W :,:], can also be A i [-S W :,:-S W ], therefore, the method for determining the target character is described in two cases as follows.

[0077] First, the target attention matrix is: A i [-S W :,:], that is, when the target column corresponding to the most recent character in the input text is included, for the most recent character, an arbitrary target number is determined as the importance score of the most recent character. And for the remaining characters in the input text except the most recent character, the mean and variance are calculated for the corresponding target columns in the target attention matrix, and the mean and variance are weighted and integrated to obtain the importance score corresponding to the remaining characters. In this way, the importance score of each character in the time dimension can be obtained.

[0078] After obtaining the importance score of each character in the input text in the time dimension, select the target number of characters B from each character i For example, the characters can be sorted from high to low according to their importance scores, and then the target characters with the highest order are selected. i as the target character.

[0079] It should be understood that since the importance score of the most recent character is a target large number, it will definitely be selected as the target character, so the most recent character mentioned above can also be called a reserved character.

[0080] In one example, the importance score of each character can be calculated according to the following formula:

[0081]

[0082] Among them, I i [n] is the importance score of the nth character, A i [-S W :,n] is the nth column of the target attention matrix, Mean() is the mean function, which is used to measure the continuous importance of attention, Var() is the variance function, which is used to measure the variability of attention, and Ω is the target large number.

[0083] It can be seen that this scheme will consider the importance of characters based on multiple factors.

[0084] It should be understood that in practice, various modifications may be made to Formula 4, for example, γ may be removed and the mean and variance may be summed directly. For another example, Var() may be replaced with a function for calculating the standard deviation or covariance, and so on.

[0085] Secondly, the target attention matrix is: A i [-S W :,:-S W ], that is, when the target column corresponding to the most recent character in the input text is not included, the most recent character is taken as the reserved character, and the characters corresponding to each target column in the target attention matrix are taken as candidate characters. Then, the target candidate characters are selected from the candidate characters so that the target candidate characters plus the reserved characters are exactly equal to the number of target characters.

[0086] The process of selecting target candidate characters can be as follows: for each candidate character, the mean and variance are calculated for the target column corresponding to the candidate character in the target attention matrix, and the mean and variance are weighted and integrated to obtain the importance score corresponding to the candidate character. Afterwards, the candidate characters can be sorted from high to low according to the importance score, and the target candidate characters with the highest number of rankings are selected, and the target characters are formed based on the target candidate characters and the reserved characters. The target number here is obtained by multiplying the above target character number B i The result is obtained by subtracting the number of reserved characters.

[0087] Considering that the importance scores of each target character will be needed in the subsequent update of the target cache area corresponding to the i-th layer, the importance scores of the reserved characters not included in the target attention matrix can be assigned to any target large number to ensure that the corresponding attention data will not be removed.

[0088] Finally, the attention data of the target character described in this scheme refers to the character representation determined for the target character using the i-1th layer and the key vector and value vector determined by the key matrix and value matrix corresponding to the i-th layer.

[0089] Specifically, the key vector can be expressed as: K = XW K , the value vector can be expressed as: V = XW V , where X is the character representation output by the i-1th layer for the target character, W K and W V are the key matrix and value matrix (i.e., the attention parameters of the i-th layer), respectively.

[0090] In practice, the above attention data is also called key-value pairs or KV pairs.

[0091] Step S310, based on the i-th preference score, updating the previous cache areas of each layer before the i-th layer, and updating the attention data of the characters stored therein.

[0092] As mentioned above, this solution performs cache operations layer by layer. Therefore, when the i-th layer is reached, the previous layers have all been cached, that is, the corresponding cache areas have been allocated to the previous layers. However, since the cache requirements of the i-th layer are not considered when allocating cache areas to the previous layers, the previous cache areas of the previous layers need to be updated.

[0093] It should be noted that in this solution, the updating method for the previous cache area of ​​each layer before the i-th layer is similar. Therefore, the corresponding updating method is explained below by taking any target layer in each layer as an example.

[0094] Specifically, for the target layer, the sum of the i-th preference score corresponding to the i-th layer and the preference scores corresponding to the layers before the i-th layer is calculated, that is, the sum of the i preference scores corresponding to the first i layers is calculated. According to the ratio of the target attention score corresponding to the target layer to the sum result and the preset total number of characters, the number of updated characters corresponding to the target layer is determined. Based on the number of updated characters, the target prior cache area corresponding to the target layer is updated.

[0095] It should be understood that the above target attention score is calculated when the target layer is cached.

[0096] Secondly, the ratio of the target attention score corresponding to the target layer to the summed result can be understood as the update allocation ratio corresponding to the target layer, that is, the cache allocation ratio re-determined for the target layer after considering the cache requirements of the i-th layer.

[0097] Finally, the number of updated characters can be calculated based on Formula 3. We only need to replace P in Formula 3 with i Just replace it with the target attention score.

[0098] It should be noted that after calculating the above-mentioned number of updated characters, the above-mentioned number of updated characters can be multiplied by the cache occupancy corresponding to a single character to obtain the corresponding update occupancy of the target layer. Afterwards, the previous cache area of ​​the target layer can be updated to a cache area matching the update occupancy.

[0099] It should be understood that for the target layer, the updated cache area will be smaller than the previously allocated cache area because the denominator in Formula 3 for calculating the number of updated characters becomes larger (the i-th preference score is newly added). Therefore, after executing the operation of updating the cache area, some of the data stored therein needs to be removed, and the removal process is described below.

[0100] Specifically, the characters stored in the previous cache area corresponding to the target layer can be sorted from high to low according to the corresponding importance scores, and the remaining characters except the characters with the updated number of characters ranked first are determined as removed characters, and the attention data of the removed characters are removed.

[0101] Assuming that the target layer is represented as layer l, the corresponding removal operation can be expressed as: EVICT({K l ,V l},B l ,I l ), where K l is a matrix consisting of the key vectors of each character stored in the previous cache area of ​​the lth layer, V l is a matrix consisting of the value vectors of each character stored in the previous cache area of ​​the lth layer, B l To update the number of characters, I l is a vector composed of the importance scores of each character stored in the previous cache area of ​​the lth layer. The meaning of this removal operation is: retain the corresponding characters of I l Middle front B l The KV pairs of the characters with the highest importance scores are retained and can be expressed as:

[0102] in, D l =TopK(B l ,I l ).

[0103] Among them, D l Indicates that it corresponds to I l Middle front B l The index of the KV pair of the character with the highest importance score.

[0104] Similar to the above-mentioned target layer update method, the previous cache areas of other layers before the i-th layer can be updated, and the attention data of the characters stored therein can be updated, which will not be repeated in this manual.

[0105] In practice, the update operations for the layers preceding the i-th layer can be executed in parallel, thereby reducing the time complexity of the update operations for each layer to the level of a single-layer update operation, thereby improving the efficiency of cache operations.

[0106] It should be noted that this solution actually considers the importance of characters changing over time when removing corresponding key-value pairs based on their importance scores, thus solving the limitation of existing methods that often ignore time dynamics. In addition, since the importance score of characters is determined by comprehensively considering multiple factors such as the continuous importance and changes of attention, this solution is a robust removal strategy that can tolerate attention shifts.

[0107] At this point, the caching operation for the i-th layer is completed, and then the caching operation can be performed on the i+1-th layer until the last layer is reached.

[0108] Among them, when the caching operation is performed on the i+1th layer, the previous cache area of ​​the previous i layers will be updated (for example, for the i-th layer, the corresponding target cache area will be updated; for example, for the target layer, the cache area corresponding to the target layer will be updated again), and the attention data of the characters stored therein will be updated. The specific updating method can refer to the above-mentioned target layer, and this manual will not repeat it here.

[0109] In summary, this scheme is a dynamic cache management method. After the cache operation is performed on the last layer of all attention layers, the cache allocation ratio corresponding to each attention layer is finally determined, as well as the cache area corresponding to each attention layer and the attention data of the characters stored therein are also finally determined.

[0110] It should be noted that the dynamic cache management method provided by this solution can maintain effective attention data in limited memory resources. In other words, this solution can ensure the generation quality of large models under limited memory resources.

[0111] The cache optimization method in large model reasoning provided in the embodiments of this specification is described below with examples.

[0112] Figure 5 A schematic diagram of a cache optimization method in an example of this specification is shown. Figure 5 In the large model, there are 4 attention layers, and the caching operations are performed on the 4 attention layers in turn:

[0113] First enter the first layer (not shown in the figure). Since there is only a preference score P0 for the first layer, the corresponding cache allocation ratio is 1, so the total cache area is allocated to the first layer, which is specifically expressed as: And store the KV pairs of all characters in the input text.

[0114] Entering the second layer, first calculate the corresponding preference score P1, and based on the preference scores P1 and P0, determine the cache allocation ratio of the second layer: 0.26 and the matching cache area: Next, the KV pairs of the second layer are removed and the remaining KV pairs are stored in the cache area: Finally, based on the preference scores P1 and P0, the cache allocation ratio of the first layer is determined: 0.74 and the matching cache area: And for the cache area of ​​layer 1: The KV pairs in the , perform the removal operation, and the cache area of ​​the first layer after the removal operation is expressed as:

[0115] Next, we enter the third layer and calculate the corresponding preference score P2. Based on the preference scores P2, P1, and P0, we determine the cache allocation ratio of the third layer: 0.46 and the matching cache area: Next, the KV pairs at the third layer are removed and the remaining KV pairs are stored in the cache area: Finally, according to the preference scores P2, P1, and P0, the cache allocation ratios of the second layer and the first layer are determined as 0.14 and 0.40, respectively, and the corresponding cache areas are determined as follows: And for the cache areas of layer 2 and layer 1 respectively: The KV pairs in the detach operation are removed, and the cache areas of the second and first layers after the removal operation are expressed as:

[0116] Finally, we enter the 4th layer and first calculate the corresponding preference score P3. Based on the preference scores P3, P2, P1, and P0, we determine the cache allocation ratio of the 4th layer: 0.23 and the matching cache area: Next, the KV pairs at the 4th layer are removed and the remaining KV pairs are stored in the cache area: Finally, according to the preference scores P3, P2, P1, and P0, the cache allocation ratios of the third layer, the second layer, and the first layer are determined as 0.35, 0.11, and 0.31, respectively, as well as the matching cache areas: And for the cache areas of layer 3, layer 2, and layer 1 respectively: The KV pairs in the table are removed. The cache areas of the third, second, and first layers after the removal operation are expressed as follows:

[0117] At this point, the cache operation has been completed for all four attention layers, so the cache allocation ratios corresponding to the 1st to 4th layers are fixed to: {0.31, 0.11, 0.35, 0.23}, and their respective cache areas are also fixed to:

[0118] from Figure 5As can be seen from the figure, at each layer, the KV cache budget is reallocated to the previous layers according to the preference score obtained by the current layer to always maintain a given cache budget size.

[0119] It should be understood that Figure 5 This is just an exemplary description. In practice, the number of attention layers included in the large model may be more than 4 layers, for example, 10 layers or even more, and this specification does not limit this. In addition, the cache allocation ratio corresponding to each attention layer is not limited to the above values.

[0120] In summary, the cache optimization method in large model reasoning provided by the embodiments of this specification has the following advantages:

[0121] First, this scheme allows a global view of cache size allocation, so that memory resources can be optimally allocated according to the different attention mechanisms of each layer. That is, this scheme can use a global attention model to consider the unique characteristics of each layer when allocating memory resources.

[0122] Second, our scheme analyzes layer-specific key-value cache preferences during the pre-population phase and uses these preferences to manage cache budgets in a cascading manner, which can effectively reduce peak memory usage to the target level.

[0123] Furthermore, this scheme proposes a robust removal strategy that can tolerate attention shift. When attention shift occurs in a long context, this scheme can avoid the information loss problem caused by overly aggressive removal strategies.

[0124] Finally, our scheme consistently outperforms current baseline methods under a variety of models and memory constraints, and performs particularly well in low-memory scenarios.

[0125] Corresponding to the cache optimization method in the above large model reasoning, an embodiment of this specification also provides a cache optimization device, such as Figure 6 As shown, the device may include:

[0126] The acquisition unit 602 is used to obtain the target attention matrix of the i-th layer, which includes the attention coefficients between characters in the input text calculated according to the parameters of the i-th layer.

[0127] Determination unit 604 is used to determine a first index value representing the spatial dispersion according to the distribution of row data of the target attention matrix; and to determine a second index value representing the temporal deviation according to the distribution of column data of the target attention matrix.

[0128] The determination unit 604 is further configured to determine an i-th preference score corresponding to the i-th layer according to the first indicator value and the second indicator value.

[0129] The storage unit 606 is used to determine the target cache area allocated to the i-th layer in the total cache area according to the i-th preference score, and store the attention data of the target character in the input text therein.

[0130] The updating unit 608 is used to update the previous cache areas of each layer before the i-th layer according to the i-th preference score, and update the attention data of the characters stored therein.

[0131] In one embodiment, the determining unit 604 is specifically configured to:

[0132] For each target row in the target attention matrix, the corresponding information entropy is calculated respectively, and then the information entropy is summed to obtain the first index value;

[0133] For each target column in the target attention matrix, the corresponding variance is calculated respectively, and then the variances are summed to obtain the second index value.

[0134] In one embodiment, the determining unit 604 is further specifically configured to:

[0135] The product of the first index value to the power of T1 and the second index value to the power of T2 is determined as the i-th preference score, wherein T1 and T2 are respectively used to control the influence of the first index value and the second index value on the i-th preference score.

[0136] In one embodiment, the storage unit 606 includes:

[0137] The acquisition submodule 6062 is used to acquire the preference scores corresponding to the layers before the i-th layer, and sum the preference scores with the i-th preference score to obtain a sum result;

[0138] The calculation submodule 6064 is used to determine the target number of characters corresponding to the i-th layer according to the ratio of the i-th preference score to the sum result and the preset total number of characters;

[0139] The determination submodule 6066 is used to determine the target cache area according to the target number of characters.

[0140] In one embodiment, each target column in the target attention matrix corresponds to each character in the input text, and the apparatus further comprises:

[0141] A calculation unit 610, used to calculate the importance score of each character in the time dimension;

[0142] The selection unit 612 is used to select a target number of characters from each character as target characters according to the importance scores.

[0143] In one embodiment, the calculation unit 610 is specifically configured to:

[0144] For reserved characters in the input text, an arbitrary target large number is determined as the importance score of the reserved characters, the reserved characters including a number of characters starting from the end of the input text;

[0145] For the remaining characters except the reserved characters, the mean and variance are calculated for their corresponding target columns in the target attention matrix, and the mean and variance are weighted combined to obtain the importance score corresponding to the remaining characters.

[0146] In one embodiment, each target column in the target attention matrix corresponds to each candidate character, each candidate character is a remaining character in the input text except for the reserved characters, and the reserved characters include a number of characters starting from the end of the input text; the above device also includes:

[0147] A calculation unit 610 is used to calculate the importance score of each candidate character in the time dimension based on each target column of the target attention matrix;

[0148] A selection unit 612 is used to select a target number of target candidate characters from each candidate character based on the importance score, where the target number is obtained by subtracting the target number of characters from the number of reserved characters;

[0149] The forming unit 614 is used to form a target character based on the target candidate character and the reserved characters.

[0150] In one embodiment, the attention data of the target character includes a key vector and a value vector determined by using a character representation determined for the target character at the i-1th layer and a key matrix and a value matrix corresponding to the i-th layer.

[0151] In one embodiment, the updating unit 608 is specifically configured to:

[0152] For any target layer, calculate the sum of the i-th preference score and the preference scores corresponding to each layer;

[0153] Determine the number of updated characters corresponding to the target layer according to the ratio of the target attention score corresponding to the target layer to the sum result and the preset total number of characters;

[0154] Based on the number of updated characters, the target previous cache area corresponding to the target layer is updated.

[0155] In one embodiment, the updating unit 608 includes:

[0156] The sorting submodule 6082 is used to sort the characters stored in the target previous cache area according to the corresponding importance scores from high to low, and determine the remaining characters except the characters with the updated character number that are ranked first as the removal characters;

[0157] The removal submodule 6084 is used to remove the attention data of the removed characters from the target previous cache area.

[0158] The functions of the functional modules of the device in the above-mentioned embodiment of this specification can be implemented through the steps of the above-mentioned method embodiment. Therefore, the specific working process of the device provided by one embodiment of this specification will not be repeated here.

[0159] A cache optimization device provided in an embodiment of the present specification can effectively perform key-value caching in limited memory resources, thereby ensuring the generation quality of large models.

[0160] According to another embodiment, there is also provided a computer-readable storage medium having a computer program stored thereon, which, when executed in a computer, causes the computer to execute a combination of Figure 3 The method described.

[0161] According to another embodiment of the present invention, a computing device is provided, including a memory and a processor, wherein the memory stores executable code, and when the processor executes the executable code, the Figure 3 The method described.

[0162] Each embodiment in this specification is described in a progressive manner, and the same or similar parts between the embodiments can be referred to each other, and each embodiment focuses on the differences from other embodiments. In particular, for the medium or device embodiment, since it is basically similar to the method embodiment, the description is relatively simple, and the relevant parts can be referred to the partial description of the method embodiment.

[0163] The above is a description of a specific embodiment of the specification. Other embodiments are within the scope of the appended claims. In some cases, the actions or steps recorded in the claims can be performed in an order different from that in the embodiments and still achieve the desired results. In addition, the processes depicted in the drawings do not necessarily require the specific order or continuous order shown to achieve the desired results. In some embodiments, multitasking and parallel processing are also possible or may be advantageous.

[0164] The specific implementation methods described above further illustrate the purpose, technical solutions and beneficial effects of this specification. It should be understood that the above description is only the specific implementation method of this specification and is not intended to limit the protection scope of this specification. Any modifications, equivalent substitutions, improvements, etc. made on the basis of the technical solutions of this specification should be included in the protection scope of this specification.

Claims

1. A cache optimization method in large model reasoning, comprising: In the pre-filling stage, a cache operation is performed layer by layer for multiple attention layers in the large model, wherein the cache operation for any i-th layer includes: Get the target attention matrix of the i-th layer, which contains the attention coefficients between characters in the input text calculated according to the parameters of the i-th layer; Determine a first index value representing the spatial dispersion according to the distribution of the row data of the target attention matrix; and determine a second index value representing the temporal deviation according to the distribution of the column data of the target attention matrix; Determine an i-th preference score corresponding to the i-th layer according to the first index value and the second index value; According to the i-th preference score, determine a target cache area allocated for the i-th layer in the total cache area, and store the attention data of the target character in the input text therein; According to the i-th preference score, the previous cache areas of each layer before the i-th layer are updated, and the attention data of the characters stored therein are updated.

2. The method according to claim 1, wherein: The determining of a first index value characterizing the spatial dispersion comprises: For each target row in the target attention matrix, calculate the corresponding information entropy respectively, and then sum the information entropies to obtain the first index value; For each target column in the target attention matrix, the corresponding variance is calculated respectively, and then the variances are summed to obtain the second index value.

3. The method according to claim 1, wherein: The determining the i-th preference score corresponding to the i-th layer includes: The product of the first index value to the power of T1 and the second index value to the power of T2 is determined as the i-th preference score, wherein T1 and T2 are respectively used to control the influence of the first index value and the second index value on the i-th preference score.

4. The method according to claim 1, wherein: The determining of the target cache area allocated to the i-th layer in the total cache area includes: Obtaining each preference score corresponding to each layer before the i-th layer, and summing each preference score with the i-th preference score to obtain a sum result; Determine the target number of characters corresponding to the i-th layer according to the ratio of the i-th preference score to the sum result and the preset total number of characters; The target cache area is determined according to the target number of characters.

5. The method according to claim 4, wherein: Each target column in the target attention matrix corresponds to each character in the input text, and the method further includes: Calculating the importance score of each character in the time dimension; According to the importance score, characters of the target number of characters are selected from the characters as the target characters.

6. The method according to claim 5, wherein: The calculating the importance score of each character in the time dimension includes: For the reserved characters in the input text, an arbitrary target large number is determined as the importance score of the reserved characters; the reserved characters include a number of characters starting from the end of the input text; For the remaining characters except the reserved characters, the mean and variance are calculated for their corresponding target columns in the target attention matrix, and the mean and variance are weightedly combined to obtain the importance score corresponding to the remaining characters.

7. The method according to claim 4, wherein: Each target column in the target attention matrix corresponds to each candidate character; each candidate character is the remaining characters in the input text except the reserved characters; the reserved characters include a number of characters starting from the end of the input text; the method further includes: Calculating the importance score of each candidate character in the time dimension based on each target column of the target attention matrix; Based on the importance scores, a target number of target candidate characters is selected from the candidate characters; the target number is obtained by subtracting the target number of characters from the number of reserved characters; The target character is formed based on the target candidate character and the reserved character.

8. The method according to claim 1, wherein: The attention data of the target character includes a character representation determined for the target character using the i-1th layer and a key vector and a value vector determined by a key matrix and a value matrix corresponding to the i-th layer.

9. The method according to claim 1, wherein: The updating of the previous cache areas of each layer before the i-th layer includes: For any target layer, calculate the sum of the i-th preference score and the preference scores corresponding to the layers; Determine the number of updated characters corresponding to the target layer according to the ratio of the target attention score corresponding to the target layer to the summation result and the preset total number of characters; Based on the updated character number, the target previous cache area corresponding to the target layer is updated.

10. The method according to claim 9, wherein: The updating of the attention data of the characters stored therein comprises: Sort the characters stored in the target cache area according to the corresponding importance scores from high to low, and determine the remaining characters except the characters with the updated number of characters at the top of the sort as the removal characters; The attention data of the removed character is removed from the target previous cache area.

11. A cache optimization device, comprising: An acquisition unit, used for acquiring a target attention matrix of the i-th layer, which includes attention coefficients between characters in the input text calculated according to the parameters of the i-th layer; A determination unit, configured to determine a first index value representing a spatial dispersion according to a distribution of row data of the target attention matrix; and to determine a second index value representing a temporal deviation according to a distribution of column data of the target attention matrix; The determining unit is further configured to determine an i-th preference score corresponding to the i-th layer according to the first index value and the second index value; a storage unit, for determining a target cache area allocated for the i-th layer in the total cache area according to the i-th preference score, and storing the attention data of the target character in the input text therein; An updating unit is used to update the previous cache areas of each layer before the i-th layer according to the i-th preference score, and to update the attention data of the characters stored therein.

12. The device according to claim 11, wherein The determining unit is specifically used for: For each target row in the target attention matrix, calculate the corresponding information entropy respectively, and then sum the information entropies to obtain the first index value; For each target column in the target attention matrix, the corresponding variance is calculated respectively, and then the variances are summed to obtain the second index value.

13. The device according to claim 11, wherein: The determining unit is further specifically configured to: The product of the first index value to the power of T1 and the second index value to the power of T2 is determined as the i-th preference score, wherein T1 and T2 are respectively used to control the influence of the first index value and the second index value on the i-th preference score.

14. The device according to claim 11, wherein: The storage unit comprises: An acquisition submodule, used to acquire the preference scores corresponding to the layers before the i-th layer, and sum the preference scores with the i-th preference score to obtain a sum result; A calculation submodule, used for determining the target number of characters corresponding to the i-th layer according to the ratio of the i-th preference score to the sum result and a preset total number of characters; The determination submodule is used to determine the target cache area according to the target number of characters.

15. The device according to claim 14, wherein: Each target column in the target attention matrix corresponds to each character in the input text, and the device further includes: A calculation unit, used for calculating the importance score of each character in the time dimension; A selection unit is used to select the target number of characters from the characters as the target characters according to the importance scores.

16. The device according to claim 15, wherein: The computing unit is specifically used for: For the reserved characters in the input text, an arbitrary target large number is determined as the importance score of the reserved characters; the reserved characters include a number of characters starting from the end of the input text; For the remaining characters except the reserved characters, the mean and variance are calculated for their corresponding target columns in the target attention matrix, and the mean and variance are weightedly combined to obtain the importance score corresponding to the remaining characters.

17. The device according to claim 14, wherein: Each target column in the target attention matrix corresponds to each candidate character; each candidate character is the remaining characters in the input text except the reserved characters; the reserved characters include a number of characters starting from the end of the input text; the device also includes: A calculation unit, used for calculating the importance score of each candidate character in the time dimension based on each target column of the target attention matrix; A selection unit, configured to select a target number of target candidate characters from the candidate characters based on the importance scores; the target number is obtained by subtracting the target number of characters from the number of reserved characters; A forming unit is used to form the target character based on the target candidate character and the reserved character.

18. The device according to claim 11, wherein The attention data of the target character includes a character representation determined for the target character using the i-1th layer and a key vector and a value vector determined by a key matrix and a value matrix corresponding to the i-th layer.

19. The device according to claim 11, wherein: The updating unit is specifically used for: For any target layer, calculate the sum of the i-th preference score and the preference scores corresponding to the layers; Determine the number of updated characters corresponding to the target layer according to the ratio of the target attention score corresponding to the target layer to the summation result and the preset total number of characters; Based on the updated character number, the target previous cache area corresponding to the target layer is updated.

20. The device according to claim 19, wherein The updating unit comprises: A sorting submodule, used for sorting the characters stored in the target previous cache area according to the corresponding importance scores from high to low, and determining the remaining characters except the characters of the updated character number that are ranked first as the removal characters; The removal submodule is used to remove the attention data of the removed character from the target previous cache area.

21. A computer-readable storage medium having a computer program stored thereon, wherein: When the computer program is executed in a computer, the computer is caused to execute the method according to any one of claims 1 to 10.

22. A computing device comprising a memory and a processor, wherein: The memory stores executable codes, and when the processor executes the executable codes, the method according to any one of claims 1 to 10 is implemented.

Images