Cache optimization method and device in large model inference

Abstract

The embodiment of the specification provides a cache optimization method and device in large model reasoning, comprising: in a pre-filling stage, performing a cache operation layer by layer for a plurality of attention layers in a large model, wherein the cache operation of any i-th layer comprises: obtaining a target attention matrix of the i-th layer. According to the distribution of the target attention matrix row data and column data, a first index value and a second index value are determined. According to the first index value and the second index value, an i-th preference score corresponding to the i-th layer is determined. According to the i-th preference score, a target cache area for the i-th layer is allocated in a total cache area, and attention data of a target character in an input text is stored therein. According to the i-th preference score, the previous cache area of each layer before the i-th layer is updated, and the attention data of the character stored therein is updated.

Description

TECHNICAL FIELD

[0001] One or more embodiments of the present specification relate to the technical field of computer technology, and more particularly to a cache optimization method and device in large model inference. BACKGROUND

[0002] A large language model (LLM), simply referred to as a large model, is an artificial intelligence model with a large number of parameters, which is specifically designed to process and generate natural language text. They are trained on massive amounts of text data, usually based on the Transformers architecture, with multiple attention (Self-Attention) layers and containing tens of billions or even hundreds of billions of parameters. The more layers, the better the model can understand and generate complex text structures.

[0003] Currently, the ability of large models to handle long texts has been greatly improved, for example, some large models can handle more than 128k characters (tokens). However, as the length of the text increases, the computational complexity of the attention layers of the large model will increase significantly. A common strategy is to store the intermediate results of the key and value of the attention layer calculation of the large model during the inference process of the large model, also known as Key Value Cache (KV Cache). This can significantly reduce the computational complexity, because the large model does not have to recalculate all the keys and values of the 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 inference. SUMMARY

[0004] One or more embodiments of the present specification describe a cache optimization method and device in large model inference, 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 inference is provided, comprising: in a pre-filling stage, performing cache operations layer by layer for a plurality of attention layers in the large model, wherein the cache operation for any ith layer includes:

[0006] obtaining a target attention matrix of the ith layer, which contains attention coefficients between characters in the input text calculated according to the parameters of the ith layer;

[0007] determining a first index value representing spatial dispersion from the distribution of row data of the target attention matrix; and determining a second index value representing time offset from the distribution of 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] determine a target cache area for the i-th layer in a total cache area according to the i-th preference score, and store attention data of a target character in the input text in the target cache area;

[0010] update a previous cache area of each layer before the i-th layer and attention data of a character stored in the previous cache area according to the i-th preference score.

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

[0012] an obtaining unit configured to obtain a target attention matrix of an i-th layer, which contains attention coefficients between characters in an input text calculated according to parameters of the i-th layer;

[0013] a determining unit configured to determine a first index value representing spatial dispersion according to a distribution of row data of the target attention matrix, and determine a second index value representing time offset 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 storing unit configured to determine a target cache area for the i-th layer in a total cache area according to the i-th preference score, and store attention data of a target character in the input text in the target cache area;

[0016] an updating unit configured to update a previous cache area of each layer before the i-th layer and attention data of a character stored in the previous cache area according to the i-th preference score.

[0017] In a third aspect, a computer readable storage medium is provided, which stores a computer program, and when the computer program is executed in a computer, the computer executes the method of the first aspect.

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

[0019] The cache optimization method and device in large model inference provided by one or more embodiments of the specification, in the pre-filling stage, for a plurality of attention layers in the large model, the cache operation is performed layer by layer, specifically, for the current attention layer, by considering the change of the attention coefficient of the attention layer in the spatial and temporal dimensions, the preference of the attention layer for caching is dynamically evaluated, and the corresponding cache is allocated for the attention layer based on the preference, and the caches of the previous attention layers are updated in a cascading manner, thereby, the limited memory resources can be better and more fully utilized, while the generation quality of the large model is almost not affected. BRIEF DESCRIPTION OF DRAWINGS

[0020] In order to more clearly illustrate the technical solutions of the embodiments of the specification, the drawings needed in the embodiment description will be briefly introduced below. Obviously, the drawings in the following description are only some embodiments of the specification, and other drawings can be obtained by those skilled in the art without creative labor on the basis of these drawings.

[0021] Figure 1 A comparison chart of the effect of the traditional key-value cache method and the present scheme is shown.

[0022] Figure 2 A schematic diagram of an implementation scenario of an embodiment disclosed in the specification is shown.

[0023] Figure 3 A flowchart of the cache optimization method in large model inference according to an embodiment of the specification is shown.

[0024] Figure 4 A schematic diagram of the target attention matrix in an example of the specification is shown.

[0025] Figure 5 A schematic diagram of the cache optimization method in an example of the specification is shown.

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

[0027] The scheme provided by the specification will be described below in conjunction with the drawings.

[0028] As described above, in order to improve the efficiency of large model character generation, key-value caching is needed.

[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 a key-value pair of a certain predetermined size is reserved. However, this scheme does not take into account the heterogeneity of the model and each attention layer inside the model, so it cannot effectively utilize cache resources.

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

[0032] Therefore, in the pre-filling stage, for multiple attention layers in the large model, the cache operation is performed layer by layer. Specifically, for the current attention layer, by considering the changes of the attention coefficients of the attention layer in the spatial and temporal dimensions, the preference of the attention layer for caching is dynamically evaluated, 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 cascading manner. In this way, the limited memory resources can be better and more fully utilized, while the generation quality of the large model is almost not affected.

[0033] Figure 1 The effect comparison chart of the traditional key-value cache method and the present scheme is shown. Figure 1 In the above, the upper part shows two traditional key-value cache methods, wherein the left part is a key-value cache method based on uniform allocation proportion, that is, for any input text, the cache allocation proportion of each of the four attention layers of the large model is 0.25. The right part is a key-value cache method based on fixed allocation mode, that is, for any input text, the cache allocation proportion of the four attention layers of the large model is fixed as 0.50, 0.34, 0.18 and 0.02.

[0034] In the present scheme, for input text A or input text B, first, the corresponding two index values of each of the four attention layers of the large model are determined based thereon, wherein one index value represents the dispersion degree of the attention coefficients of the attention layer in the spatial dimension, and the other index value represents the offset degree of the attention coefficients of the attention layer in the time dimension. Then, the corresponding layer preference can be determined based on the two index values corresponding to each attention layer, and the corresponding cache is allocated to each attention layer of the large model based on the layer preference. For input text A, the cache allocation proportion 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 proportion 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 merely an illustrative example. In practice, large models may include more than four attention layers, such as ten or even more. This specification does not limit this. Furthermore, the cache allocation ratio for each attention layer is not limited to the values ​​mentioned above.

[0036] As can be seen, in this scheme, the cache allocation ratio between different layers can be adjusted according to layer preferences to adapt to various contexts and given memory budgets. This allows for better and more efficient use of limited memory resources, while having almost no impact on the generation quality.

[0037] Figure 2 This is a schematic diagram illustrating an implementation scenario of one of the embodiments disclosed in this specification. Figure 2 In the large model, the reasoning process includes a pre-filling stage and a decoding stage.

[0038] In the pre-filling stage, caching is performed layer by layer for multiple attention layers in the large model. Specifically, for any i-th layer, firstly, based on the distribution of attention coefficients in the target attention matrix corresponding to the i-th layer, two indicators representing spatial dispersion and temporal offset are determined. Next, based on these two indicators, a preference score is determined for the i-th layer. Finally, based on this preference score, a corresponding cache region is determined to store the attention data of the target characters in the input text, and the prior cache regions of all layers before the i-th layer are updated based on the preference score, along with the attention data of the characters stored therein.

[0039] During the decoding phase, based on the attention data of the characters stored in the corresponding buffer areas of multiple attention layers, the generated characters for the input text are output.

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

[0041] As mentioned earlier, in this scheme, during the pre-filling stage, caching operations are performed layer by layer for multiple attention layers in the large model. Since the caching operations for each layer are similar, the following explanation uses any i-th layer as an example to illustrate the corresponding caching operations.

[0042] Figure 3 This diagram illustrates a caching optimization method for large model inference according to one embodiment of this specification. This method can be executed by any device, apparatus, platform, or cluster of devices with computing and processing capabilities. Figure 3 As shown, the method may include the following steps:

[0043] Step S302: Obtain the target attention matrix of the i-th layer, which contains the attention coefficients between characters in the input text calculated based on the parameters of the i-th layer.

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

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

[0046]

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

[0048] It should be understood that in the attention matrix of size S×S (hereinafter also called the original attention matrix), each row from left to right (or each column from top to bottom) corresponds to each character from front to back in the input text.

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

[0050] Where the size of the observation window is S W In the case of ×S, the target attention matrix can be represented as: A i [-S w :,:], which can be like Figure 4 As shown on the left (filled with diagonal lines), this contains the post-S from the original attention matrix. W Rows and all S columns. The size of the observation window is S. w ×(SS W In the case of A), the target attention matrix can be represented as: A i [-S W :,:-S W ], which can be like Figure 4 As shown on the right (filled with slashes), this contains the post-S from the original attention matrix. W Line and 1 to SSW column.

[0051] In summary, the target attention matrix contains the last S W rows, it should be understood that the last S W rows correspond to the S W characters (also referred to as the most recent characters) from the end of the input text, and similarly, the last S W columns in the original attention matrix also correspond to the most recent characters. That is, in the case of 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 the present specification does not limit this.

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

[0054] In step S304, an index value E representing spatial dispersion is determined according to the distribution of the target attention matrix row data, and an index value V representing time offset is determined according to the distribution of the target attention matrix column data.

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

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

[0057] In step S306, the i-th preference score corresponding to the i-th layer is determined according to the index value E and the index value V.

[0058] In one example, the product of the T1 power of the indicator value E and the T2 power of the indicator value V can be determined as the i-th preference score, where T1 and T2 are temperature parameters respectively used to control the influence of the indicator value E and the indicator 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] where P i is the i-th preference score, A i [-S W :,:-S W ] is the target attention matrix, the E() function represents summing the information entropy of each target row first, and the V() function represents summing the variance of each column first, and T1 and T2 are temperature parameters.

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

[0063] Of course, in practice, the product of the indicator value E and the indicator value V can also be directly determined as the i-th preference score corresponding to the i-th layer, or other operations other than power operation can also be performed on the indicator value E and the indicator value V to determine the i-th preference score, as long as the influence of the two indicator values on the i-th preference score can be flexibly adjusted. The present specification does not limit this.

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

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

[0066] Step S308, according to the i-th preference score, determine the 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.

[0067] Specifically, the preference scores corresponding to each of the layers before the i-th layer can be obtained, and the obtained preference scores and the i-th preference score are summed to obtain a summation result. That is, the i preference scores corresponding to the first i layers are summed. Then, the ratio of the i-th preference score to the summation 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 character number can be determined as the target character number corresponding to the i-th layer. Finally, the target cache region allocated for the i-th layer in the total cache region is determined according to the target character number corresponding to the i-th layer.

[0068] In one example, the target character number can be determined according to the following formula:

[0069]

[0070] wherein B i is the target character number corresponding to the i-th layer, P i is the i-th preference score corresponding to the i-th layer, P k is the k-th layer before the i-th layer, is the cache allocation ratio corresponding to the i-th layer, B total is the preset total character number.

[0071] More specifically, the cache occupancy corresponding to a single character (i.e., the occupancy of the key-value pair of a single character) can be preset in advance, and then the target character number B i is multiplied by the cache occupancy corresponding to a single character to obtain the target occupancy corresponding to the i-th layer. Then, a region matching the target occupancy can be allocated in the total cache region as the target cache region.

[0072] It should be noted that since the target character number is calculated based on the i-th preference score, it can also be considered that the cache region is allocated for the i-th layer based on the i-th preference score, and since the i-th preference score is adjusted accordingly as the input text changes, the scheme actually implements a preference-based adaptive allocation strategy.

[0073] Of course, in practice, the preset total character number can also be replaced by a preset total cache occupancy, so that after the cache allocation ratio corresponding to the i-th layer is determined, 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 is stored in the target cache region corresponding to the i-th layer, the removal operation (i.e., removing the attention data of the non-target characters in the input text) is actually performed for the i-th layer.

[0075] The determination process of the target character stored in the target cache area is described as follows.

[0076] The present solution is based on the target attention matrix to determine the target character, and as previously described, the target attention matrix can be A i [-S W :,:], or A i [-S W :,:-S W Therefore, the determination method of the target character is described as follows in two cases.

[0077] First, in the case where the target attention matrix is A i [-S W :,:], that is, the case where it contains the target column corresponding to the most recent character in the input text, for the most recent character, an arbitrary target large number is determined as the importance score of the most recent character. And for the remaining characters in the input text, the mean and variance of the target column corresponding to each character in the target attention matrix are calculated respectively, 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 time dimension of the input text, the target character number B i characters are selected from the characters as target characters. For example, the characters can be sorted in descending order of importance score, and then the top B i characters are selected as target characters.

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

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

[0081]

[0082] Where 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 a mean function for measuring the continuous importance of attention, Var() is a variance function for measuring the variability of attention, and Ω is a target large number.

[0083] It can be seen that the importance of the character is considered based on multiple factors.

[0084] It should be understood that in practice, Formula 4 can also be modified in various ways, such as removing γ and directly summing the mean and variance. For example, Var() can be replaced by a function that calculates the standard deviation or covariance, and so on.

[0085] Secondly, in the case of the target attention matrix being A i [-S W :,:-S W ], that is, not containing the target column corresponding to the most recent character in the input text, 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, and then a target candidate character is selected from the candidate characters, so that the target candidate character plus the reserved character is exactly equal to the target character number.

[0086] The process of selecting the target candidate character can be: for each candidate character, the mean and variance of the target column corresponding to the candidate character in the target attention matrix are calculated respectively, and the mean and variance are weighted and integrated to obtain the importance score of the candidate character. Then, the candidate characters can be sorted in descending order of importance score, and the target candidate characters with the top target number are selected, and the target characters are formed based on the target candidate characters and the reserved character. The target number here is obtained by subtracting the reserved character number from the target character number B i .

[0087] Considering that the importance score of each target character is needed in the subsequent process of updating the target cache area corresponding to the i-th layer, the importance score of the reserved character not included in the target attention matrix can be assigned an arbitrary large target number to ensure that the attention data corresponding to it can not be removed.

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

[0089] Specifically, the key vector can be represented as: K = XW K , and the value vector can be represented as: V = XW V , where X is the character representation output by the i-1-th 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 pair or KV pair.

[0091] At step S310, according to the i-th preference score, the previously cached area of each layer before the i-th layer is updated, and the attention data of the characters stored therein is updated.

[0092] As mentioned above, the present scheme is to perform the caching operation layer by layer, thus, when the i-th layer is reached, the caching operation has been performed on each of the previous layers, i.e., the corresponding cached area has been allocated for each of the previous layers. However, since the caching requirement of the i-th layer is not considered when the cached area is allocated for each of the previous layers, the previously cached area of each of the previous layers needs to be updated.

[0093] It should be noted that in the present scheme, the updating method of the previously cached area of each of the layers before the i-th layer is similar, thus, the corresponding updating method is described below by taking any target layer in each of the layers as an example.

[0094] Specifically, for the target layer, the sum of the i-th preference score corresponding to the i-th layer and each of the preference scores corresponding to each of the layers before the i-th layer is calculated, i.e., 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, and the preset total number of characters, the updated number of characters corresponding to the target layer is determined. Based on the updated number of characters, the target previously cached area corresponding to the target layer is updated.

[0095] It should be understood that the target attention score described above is calculated when the target layer is subjected to the caching operation.

[0096] Secondly, the ratio of the target attention score corresponding to the target layer to the sum can be understood as the updated allocation ratio corresponding to the target layer, i.e., the caching allocation ratio of the target layer determined again after considering the caching requirement of the i-th layer.

[0097] Finally, the updated number of characters described above can be calculated based on formula 3, which only needs to replace P i in formula 3 with the target attention score.

[0098] It should be noted that after the updated number of characters described above is calculated, the updated number of characters can be multiplied by the cache occupancy corresponding to a single character to obtain the updated occupancy corresponding to the target layer. Then, the previously cached area of the target layer can be updated to a cached area matching the updated occupancy.

[0099] It should be understood that for the target layer, the updated cached area will be smaller than the previously cached area allocated initially, because the denominator in formula 3 for calculating the updated number of characters becomes larger (the i-th preference score is added). Thus, after the operation of updating the cached area is performed, part of the data stored therein needs to be removed, which is described below.

[0100] Specifically, the characters stored in the prior cache area corresponding to the target layer can be sorted from high to low according to their corresponding importance scores, and the remaining characters except for the characters whose number of updated characters is at the top of the sort can be identified as characters to be removed, and the attention data of the removed characters can be removed.

[0101] Assuming the target layer is represented as layer l, the corresponding removal operation can be represented as: EVICT({K l V l},B l ,I l ), where K l V is a matrix consisting of the key vectors of each character stored in the prior cache region of layer l. l B is a matrix consisting of the value vectors of each character stored in the prior cache region of layer l. l To update the number of characters, I l This is a vector consisting of the importance scores of each character stored in the prior cache region of layer l. The removal operation means: retain the character corresponding to I... l Middle front B l The key-value pairs of the characters with the highest importance scores, the retained key-value pairs can be represented as:

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

[0103] Among them, D l Indicates corresponding to I l Middle front B l Index of the key-value pairs of the characters with the highest importance scores.

[0104] Similar to the update method for the target layer described above, the prior cache areas of other layers preceding the i-th layer can be updated, as well as the attention data of the characters stored therein. These details will not be elaborated upon here.

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

[0106] It should be noted that the scheme actually considers the importance of characters changing over time when removing the corresponding key-value pair based on the importance score of the characters, solving the limitation that the existing method often ignores the time dynamics. In addition, since the importance score of the characters is determined by comprehensively considering the continuous importance and changes of attention and other factors, the scheme is a robust removal strategy that can tolerate attention shifts.

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

[0108] Among them, when performing cache operation for the i+1-th layer, the previous cache area of the first i layers (for example, for the i-th layer, the corresponding target cache area is updated; for example, for the target layer, the cache area corresponding to the target layer is updated again) is updated, and the attention data of the characters stored therein is updated. The specific updating method can refer to the above target layer, and the description will not be repeated here.

[0109] In summary, the scheme is a dynamic cache management method, and the cache allocation ratio of each attention layer is finally determined after the cache operation is performed for the last layer in all attention layers, and 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 the scheme can maintain effective attention data in limited memory resources. In other words, the scheme can ensure the generation quality of large models in the case of limited memory resources.

[0111] The following describes the cache optimization method provided by the large model inference method of the embodiment of the present application in conjunction with an example.

[0112] Figure 5 The cache optimization method in one example of the present application is shown in the following schematic diagram, Figure 5 In which the large model includes 4 attention layers, and the cache operation is performed for the 4 attention layers in turn:

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

[0114] Then enter 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 remove operation is performed on the KV pairs in the cache area of the 2nd layer, and the remaining KV pairs are stored to the cache area of the 2nd layer. Finally, according to the preference scores P1 and P0, the cache allocation ratio of the 1st layer is determined as 0.74, and the matching cache area is: and the KV pairs in the cache area of the 1st layer are removed, and the cache area of the 1st layer after the remove operation is represented as:

[0115] Next, the 3rd layer is entered, the corresponding preference score P2 is calculated first, and based on the preference scores P2, P1 and P0, the cache allocation ratio of the 3rd layer is determined as 0.46, and the matching cache area is: Next, the remove operation is performed on the KV pairs in the cache area of the 3rd layer, and the remaining KV pairs are stored to the cache area of the 3rd layer. Finally, according to the preference scores P2, P1 and P0, the cache allocation ratios of the 2nd layer and the 1st layer are determined as 0.14 and 0.40 respectively, and the matching cache areas are: and the KV pairs in the cache area of the 2nd layer and the 1st layer are removed respectively, and the cache areas of the 2nd layer and the 1st layer after the remove operation are represented as:

[0116] Finally, the 4th layer is entered, the corresponding preference score P3 is calculated first, and based on the preference scores P3, P2, P1 and P0, the cache allocation ratio of the 4th layer is determined as 0.23, and the matching cache area is: Next, the remove operation is performed on the KV pairs in the cache area of the 4th layer, and the remaining KV pairs are stored to the cache area of the 4th layer. Finally, according to the preference scores P3, P2, P1 and P0, the cache allocation ratios of the 3rd layer, the 2nd layer and the 1st layer are determined as 0.35, 0.11 and 0.31 respectively, and the matching cache areas are: and the KV pairs in the cache area of the 3rd layer, the 2nd layer and the 1st layer are removed respectively, and the cache areas of the 3rd layer, the 2nd layer and the 1st layer after the remove operation are represented as:

[0117] By now, the cache operation has been performed for the 4 attention layers, so the cache allocation ratios corresponding to the 1st-4th layers are fixed as {0.31, 0.11, 0.35, 0.23}, and the respective cache areas are also fixed as:

[0118] From Figure 5 ​​​​​​As can be seen, at each layer, the KV cache budget is re-allocated to the previous layers according to the preference score obtained by the current layer, so as to always maintain a given cache budget size.

[0119] It should be understood that, Figure 5 It is only an example, in practice, the number of attention layers included in the large model can be more than 4, such as 10 or even more, and the 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 the large model inference provided by the embodiments of the specification has the following advantages:

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

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

[0123] Third, the scheme proposes a robust removal strategy that can tolerate attention shifts. In the case of attention shifts in long contexts, the scheme can avoid the information loss problem caused by overly aggressive removal strategies.

[0124] Finally, the scheme is always superior to the current benchmark method under various models and memory constraints, and performs better in low memory conditions.

[0125] Corresponding to the above cache optimization method in the large model inference, one embodiment of the specification also provides a cache optimization device, as shown in Figure 6 The device can include:

[0126] The acquisition unit 602 is configured to acquire a target attention matrix of the i th layer, which contains attention coefficients between characters in the input text calculated according to parameters of the i th layer.

[0127] The determination unit 604 is configured to determine a first index value representing spatial dispersion according to the distribution of the target attention matrix row data, and determine a second index value representing time offset according to the distribution of the target attention matrix column data.

[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 index value and the second index value.

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

[0130] The updating unit 608 is configured to update, according to the i-th preference score, the prior cache area of each layer before the i-th layer, and update the attention data of characters stored in the prior cache area.

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

[0132] For each target row in the target attention matrix, calculate a corresponding information entropy, and sum the information entropies to obtain a first index value;

[0133] For each target column in the target attention matrix, calculate a corresponding variance, and sum the variances to obtain a second index value.

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

[0135] Determine, as the i-th preference score, a product of T1 power of the first index value and T2 power of the second index value, where 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 an embodiment, the storage unit 606 includes:

[0137] The obtaining sub-module 6062 is configured to obtain each preference score corresponding to each layer before the i-th layer, and sum each preference score and the i-th preference score to obtain a sum result;

[0138] The calculating sub-module 6064 is configured to determine, according to a ratio of the i-th preference score and the sum result and a preset total number of characters, a target number of characters corresponding to the i-th layer;

[0139] The determining sub-module 6066 is configured to determine, according to the target number of characters, the target cache area.

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

[0141] The calculating unit 610 is configured to calculate an importance score of each character in a time dimension;

[0142] The selecting unit 612 is configured to select, according to the importance score, a number of characters as target characters from the characters.

[0143] In an embodiment, the calculating unit 610 is specifically configured to:

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

[0145] For the remaining characters other than the reserved characters, a mean value and a variance are respectively calculated for corresponding target columns of the target attention matrix, and the mean value and the variance are weighted and integrated to obtain an importance score corresponding to the remaining characters.

[0146] In an embodiment, each target column in the target attention matrix corresponds to a candidate character, and each candidate character is a remaining character in the input text other than the reserved characters, and the reserved characters include a number of characters of the input text from the end to the front; the apparatus further includes:

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

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

[0149] The formation unit 614 is configured to form the target character based on the target candidate characters and the reserved characters.

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

[0151] In an embodiment, the update unit 608 is specifically configured to:

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

[0153] According to the ratio of the target attention score corresponding to the target layer to the sum result, and the preset total character number, determine the update character number corresponding to the target layer;

[0154] Based on the update character number, update the target previous cache area corresponding to the target layer.

[0155] In an embodiment, the update unit 608 includes:

[0156] The sorting sub-module 6082 is configured to sort the characters stored in the target previous cache area in descending order of the corresponding importance scores, and determine the remaining characters other than the update character number of characters ranked first as removal characters;

[0157] The removing submodule 6084 is configured to remove attention data of the removing character from the target previous cache area.

[0158] The functions of each functional module of the device in the foregoing embodiments of the present specification can be realized by each step of the foregoing method embodiments, and thus the specific working process of the device provided by one embodiment of the present specification is not described herein.

[0159] The cache optimization device provided by one embodiment of the present specification can effectively perform key-value caching in limited memory resources, thereby ensuring the generation quality of a large model.

[0160] According to another aspect, an embodiment also provides a computer-readable storage medium having stored thereon a computer program which, when executed in a computer, causes the computer to perform the method described above. Figure 3 The method described above.

[0161] According to another aspect, an embodiment also provides a computer-readable storage medium having stored thereon a computer program which, when executed in a computer, causes the computer to perform the method described above. Figure 3 The method described above.

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

[0163] The above describes specific embodiments of the present specification. Other embodiments are within the scope of the appended claims. In some cases, the actions or steps recited in the claims can be performed in an order different than the order in which they are recited and still achieve desirable results. In addition, the processes depicted in the figures do not necessarily require the particular order shown, or sequential order, to achieve the desired results. In some implementations, multitasking and parallel processing can be advantageous.

[0164] The above specific embodiments describe the purpose, technical solutions and beneficial effects of the present specification in further detail. It should be understood that the above description is only a specific embodiment of the present specification and is not used to limit the protection scope of the present specification. Any modification, equivalent replacement, improvement, etc. made on the basis of the technical solutions of the present specification shall be included in the protection scope of the present specification.

Claims

1. A method for cache optimization in large model inference, comprising: In the pre-filling stage, a caching operation is performed layer by layer for multiple attention layers in the large model, wherein the caching operation for any ith layer includes: obtaining a target attention matrix of the ith layer, which contains attention coefficients between characters in the input text calculated according to parameters of the ith layer; determining a first index value representing spatial dispersion degree according to the distribution of row data of the target attention matrix, and a second index value representing time offset degree according to the distribution of column data of the target attention matrix; determining an ith preference score corresponding to the ith layer according to the first index value and the second index value; determining a target cache area allocated for the ith layer in a total cache area according to the ith preference score, and storing attention data of target characters in the input text in the target cache area; updating the previously cached areas of each layer before the ith layer and the attention data of characters stored therein according to the ith preference score.

2. The method of claim 1, wherein the determination of the first index value representing the spatial dispersion degree includes: calculating respective information entropies for each target row in the target attention matrix, and summing the information entropies to obtain the first index value; calculating respective variances for each target column in the target attention matrix, and summing the variances to obtain the second index value.

3. The method of claim 1, wherein, the determination of the ith preference score corresponding to the ith layer includes: determining the product of the first index value raised to the power of T1 and the second index value raised to the power of T2 as the ith preference score, wherein T1 and T2 are used to control the influence of the first index value and the second index value on the ith preference score, respectively.

4. The method of claim 1, wherein, the determination of the target cache area allocated for the ith layer in the total cache area includes: obtaining respective preference scores corresponding to each layer before the ith layer, and summing each preference score with the ith preference score to obtain a sum result; determining a target character number corresponding to the ith layer according to the ratio of the ith preference score to the sum result and a preset total character number; determining the target cache area according to the target character number.

5. The method of claim 4, wherein, each target column in the target attention matrix corresponds to a character in the input text, and the method further includes: calculating importance scores of the characters in the time dimension; selecting characters of the target character number from the characters as the target characters according to the importance scores.

6. The method of claim 5, wherein, the calculation of the importance scores of the characters in the time dimension includes: for a reserved character in the input text, determining any target large number as the importance score of the reserved character; the reserved character includes a number of characters from the end of the input text forward; for the remaining characters other than the reserved characters, calculating a mean value and a variance for the target column corresponding to each remaining character in the target attention matrix, and performing weighted integration on the mean value and the variance to obtain an importance score corresponding to the remaining character.

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

8. The method of claim 1, wherein, The attention data of the target characters includes character representations determined by the i-1th layer for the target characters and key vectors and value vectors determined by the i-th layer corresponding key matrix and value matrix.

9. The method of claim 1, wherein, The updating of the prior cache area of each layer before the i-th layer includes: For any target layer, the sum of the i-th preference score and the preference scores corresponding to each layer is calculated; According to the ratio of the target attention score corresponding to the target layer to the sum and a preset total character number, an updated character number corresponding to the target layer is determined; Based on the updated character number, the target prior cache area corresponding to the target layer is updated.

10. The method of claim 9, wherein, The updating of the attention data of the characters stored therein includes: The characters stored in the target prior cache area are sorted in descending order of corresponding importance scores, and the remaining characters except for the top updated character number are determined as removal characters; The attention data of the removal characters is removed from the target prior cache area.

11. A cache optimization apparatus, comprising: An acquisition unit configured to acquire a target attention matrix of an i-th layer, which contains attention coefficients between characters in an input text calculated according to parameters of the i-th layer; A determination unit configured to determine a first index value representing space dispersion according to the distribution of row data of the target attention matrix, and determine a second index value representing time offset according to the distribution of column data of the target attention matrix; The determination 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 configured to determine a target cache area allocated for the i-th layer in a total cache area according to the i-th preference score, and store attention data of target characters in the input text in the target cache area; An update unit configured to update prior cache areas of each layer before the i-th layer according to the i-th preference score, and update attention data of characters stored therein.

12. The apparatus of claim 11, wherein, The determination unit is specifically configured to: Calculate each information entropy corresponding to each target row in the target attention matrix, and then sum the information entropies to obtain the first index value; Calculate each variance corresponding to each target column in the target attention matrix, and then sum the variances to obtain the second index value.

13. The apparatus of claim 11, wherein, The determination unit is further specifically configured to: A product of T1th power of the first indicator value and T2th power of the second indicator value is determined as the i-th preference score, where T1 and T2 are used to control influences of the first indicator value and the second indicator value on the i-th preference score respectively.

14. The apparatus of claim 11, wherein, The storage unit comprises: An obtaining submodule is configured to obtain respective preference scores corresponding to respective layers before the i-th layer, and sum the respective preference scores and the i-th preference score to obtain a sum result; A calculating submodule is configured to determine a target character number corresponding to the i-th layer according to a ratio of the i-th preference score and the sum result and a preset total character number; A determining submodule is configured to determine the target cache region according to the target character number.

15. The apparatus of claim 14, wherein, Each target column in the target attention matrix corresponds to a character in the input text, and the apparatus further comprises: A calculating unit is configured to calculate an importance score of each character in a time dimension; A selecting unit is configured to select characters of the target character number from the characters as the target characters according to the importance scores.

16. The apparatus of claim 15, wherein, The calculating unit is specifically configured to: For a reserved character in the input text, any target large number is determined as an importance score of the reserved character; the reserved character includes a number of characters of the input text from the end to the front; For remaining characters other than the reserved character, a mean value and a variance are calculated for a target column corresponding to each of the remaining characters in the target attention matrix respectively, and the mean value and the variance are weighted and integrated to obtain an importance score corresponding to each of the remaining characters.

17. The apparatus of claim 14, wherein, Each target column in the target attention matrix corresponds to a candidate character; the candidate characters are the remaining characters other than the reserved characters in the input text; the reserved characters include a number of characters of the input text from the end to the front; and the apparatus further comprises: A calculating unit is configured to calculate an importance score of each candidate character in a time dimension based on each target column of the target attention matrix respectively; A selecting unit is configured to select target candidate characters of a target number from the candidate characters based on the importance scores; the target number is obtained by subtracting a reserved character number from the target character number; A forming unit is configured to form the target characters based on the target candidate characters and the reserved characters.

18. The apparatus of claim 11, wherein, The attention data of the target characters comprises a character representation determined by an i-1-th layer for the target characters and a key vector and a value vector determined by a key matrix and a value matrix corresponding to an i-th layer.

19. The apparatus of claim 11, wherein, The updating unit is specifically configured to: For any target layer, calculate a sum result of the i-th preference score and respective preference scores corresponding to respective layers; Determine an updating character number corresponding to the target layer according to a ratio of a target attention score corresponding to the target layer and the sum result and a preset total character number; Update a target previous cache region corresponding to the target layer based on the updating character number.

20. The apparatus of claim 19, wherein, The updating unit comprises: The sorting submodule is configured 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 top sorted characters as removal characters. The removal submodule is configured to remove the attention data of the removal characters from the target previous cache area.

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

22. A computing device comprising a memory and a processor, wherein, The memory stores executable code, and the processor executes the executable code to implement the method in any one of claims 1-10. The memory stores executable code, and the processor executes the executable code to implement the method in any one of claims 1-10.

Images