Large model construction method and device capable of presetting KV cache capacity

Abstract

The invention provides a large model construction method and device capable of presetting KV cache capacity, and is applied to an attention layer in a Transform structure of a large model, and the attention layer comprises a key value vector sequence MV composed of M key value vectors and a key vector sequence MK composed of M key vectors. Wherein M is equal to a preset value; the method comprises the following steps: aiming at an Nth input vector, mapping the Nth input vector into a write query vector wq and a first write key value vector wv; calculating by using the write query vector wq and the M key vectors to obtain a write weight vector ww; and updating M key value vectors in the historical key value vector sequence MV'by using the write-in weight vector ww and the first write-in key value vector wv, and writing the updated key value vector sequence MV into a cache. In this way, a KV cache capacity scheme with the preset length can be achieved to replace a KV cache capacity scheme infinitely increasing along with the context length.

Description

Technical Field

[0001] The present application relates to the field of artificial intelligence technology, and in particular to a method and device for constructing a large model capable of presetting a KV cache capacity. Background Art

[0002] KV cache technology is an efficient caching method. Based on KV cache technology, the model can store information in the form of KV cache sequence during reasoning. In this way, in the subsequent reasoning process, the large model can quickly access historical information, thereby improving the overall reasoning efficiency and speed. However, the sequence length of the KV cache gradually increases with the reasoning of the large model. As the sequence length increases, the cache demand also continues to grow, which makes the reasoning of the large model a memory constraint problem, greatly limiting the throughput of the system.

[0003] Using a limited-capacity KV cache to cope with infinite context reasoning is the key to solving this problem. Streaming LLM and Low-rank Embedding Sidekick with Sparse policy (LESS) are two common methods to implement limited-capacity KV cache.

[0004] However, both StreamingLLM and LESS discard some KV information and do not cache it, which may easily lead to a decrease in the reasoning effect of large models. Summary of the invention

[0005] The embodiments of the present application provide a large model construction method and device that can preset KV cache capacity to solve the problem of poor reasoning effect caused by discarding KV information in traditional caching methods.

[0006] In a first aspect, an embodiment of the present application provides a large model construction method with preset KV cache capacity, which is applied to an attention layer in a Transformer structure of a large model, wherein the attention layer includes a key-value vector sequence MV consisting of M key-value vectors, and a key vector sequence MK consisting of M key vectors; wherein M is equal to a preset value, and the key-value vector sequence MV is composed of initial values;

[0007] The method comprises: for the Nth input vector of the input attention layer, using the first matrix Wwq to map it to a write query vector wq, and using the second matrix Wwv to map it to a first write key value vector wv; wherein N≥1, the input vector is determined based on the embedding of the text token of the input large model or the output of the previous layer Transformer block; using the write query vector wq and M key vectors to calculate, and obtaining the write weight vector ww; when N=1, obtaining the key value vector sequence MV composed of the initial value from the cache, and confirming it Defined as the historical key vector sequence MV', when N>1, obtain the key vector sequence MV corresponding to the N-1th input vector from the cache and determine it as the historical key vector sequence MV'; use the write weight vector ww and the first write key vector wv to update the M key vectors in the historical key vector sequence MV' to obtain the updated key vector sequence MV; write the updated key vector sequence MV into the cache to replace the historical key vector sequence MV', the updated key vector sequence MV corresponds to the Nth input vector, that is, MV(N). MV(0) is the initial form of MV, MV(0) is composed of initial values, as the historical key vector sequence when N=1, MV(N-1) is the historical key vector sequence when N>1.

[0008] Generally, MK, MV(0), Wwq, and Wwv are obtained by training a large model, where MV(0) can also be preset to all zero values ​​or a random value and not participate in training updates.

[0009] In one achievable manner, the dimension of the key-value vector is D, and the dimension of the key vector is E.

[0010] In one implementable manner, the dimension of the first matrix Wwq is [D, E], so that the dimension of the write query vector wq is the same as the dimension of the key vector, and both are E; the dimension of the second matrix Wwv is [D, D], so that the dimension of the first write key value vector wv is the same as the dimension of the key value vector, and both are D.

[0011] In one feasible method, the steps of calculating using the write query vector wq and M key vectors to obtain the write weight vector ww include: respectively calculating the similarity between the write query vector wq and the M key vectors to obtain M weight scores; normalizing the M weight scores to obtain the write weight vector ww, and the dimension of the write weight vector ww is M.

[0012] In one implementable method, the write weight vector ww and the first write key value vector wv are used to update the M key value vectors in the historical key value vector sequence MV', and the step of obtaining the updated key value vector sequence MV includes: transposing the write weight vector ww to convert the write weight vector ww into a write weight matrix ww' of dimension [M, 1], so that the number of columns of the write weight matrix ww' matches the number of rows of the first write key value vector wv; performing matrix multiplication calculation on the write weight matrix ww' and the first write key value vector wv to obtain M second write key value vectors wMV, and the dimension of the second write key value vector wMV is D; adding the M second write key value vectors wMV and the M key value vectors in the historical key value vector sequence MV' to obtain the updated key value vector sequence MV, i.e., MV(N).

[0013] In one achievable manner, M≥1; and the historical key-value vector sequence MV(0) corresponding to the first input vector is obtained through training or is preset.

[0014] In one achievable manner, the process of using Wwq or Wwv for mapping can be replaced by a single-layer or multi-layer perceptron (MLP), which is not specifically limited in the embodiments of the present application.

[0015] In one feasible way, during the inference process, when processing token by token, MV can use the same memory block / video memory block for different N; during the training process, it can be N independent memory blocks / video memory blocks to facilitate parallel training.

[0016] In one feasible manner, the updated MV is used as a KV cache for use in subsequent calculation processes. The subsequent calculation process can be the same as the traditional large model method, and the embodiments of the present application do not make specific limitations on this.

[0017] In a second aspect, an embodiment of the present application further provides a large model construction method capable of presetting a KV cache capacity, which is applied to an attention layer in a Transformer structure of a large model, wherein the attention layer includes a key-value vector sequence MV consisting of M key-value vectors, and a key vector sequence MK consisting of M key vectors; wherein M is equal to a preset value, and the M key-value vectors are updated based on the large model construction method provided by the first aspect and its various implementations, and before the update, the key-value vector sequence MV is composed of an initial value;

[0018] The method includes: for the Nth input vector of the input attention layer, using the third matrix Wq to map the input vector to the current query vector q; wherein N≥1, the input vector is the embedding vector corresponding to the text tokenization of the input large model or the process vector output by the previous layer Transformer block of the current Transformer block; using the current query vector q and M key vectors to calculate and obtain the current weight vector w; when N=1, obtaining the key value vector sequence MV composed of initial values ​​from the cache, and when N>1, obtaining the key value vector sequence MV corresponding to the N-1th input vector from the cache; using the current weight vector w to perform weighted summation on the obtained key value vector sequence MV to obtain the output vector.

[0019] In one implementable manner, the dimension of the third matrix Wq is [F, E], where F is the dimension of the input vector and the dimension of the current query vector q is E; the dimension of the current weight vector w is M; the attention layer also includes a fourth matrix Wk with a dimension [F, E], and the fourth matrix Wk is used to map the input vector to the current key vector k with a dimension E; and the dimension of the key vector is the same as the dimension of the current key vector k; the attention layer also includes a fifth matrix Wv with a dimension [F, D], and the fifth matrix Wv is used to map the input vector to the current key value vector v with a dimension D; and the dimension of the key value vector is the same as the dimension of the current key value vector v.

[0020] In a third aspect, an embodiment of the present application further provides a large model construction device capable of presetting a KV cache capacity, which is applied to an attention layer in a Transformer structure of a large model, wherein the attention layer includes a key-value vector sequence MV consisting of M key-value vectors, and a key vector sequence MK consisting of M key vectors; wherein M is equal to a preset value, and the key-value vector sequence MV is composed of an initial value;

[0021] The device includes: a first mapping module, for mapping the Nth input vector of the input attention layer into a write query vector wq using a first matrix Wwq, and mapping it into a first write key value vector wv using a second matrix Wwv; wherein N≥1, the input vector is an embedding vector corresponding to the text tokenization of the input large model or a process vector output by the previous layer Transformer block of the current Transformer block; a first calculation module, for calculating using the write query vector wq and M key vectors to obtain a write weight vector ww; a first acquisition module, when N=1, for acquiring the weight vector ww from the cache. A key value vector sequence MV composed of initial values ​​is determined as the historical key value vector sequence MV'. When N>1, it is used to obtain the key value vector sequence MV corresponding to the N-1th input vector from the cache, and determine it as the historical key value vector sequence MV'; an update module is used to use the write weight vector ww and the first write key value vector wv to update the M key value vectors in the historical key value vector sequence MV' to obtain an updated key value vector sequence MV; a cache module is used to write the updated key value vector sequence MV into the cache to replace the historical key value vector sequence MV', and the updated key value vector sequence MV corresponds to the Nth input vector.

[0022] In a fourth aspect, an embodiment of the present application further provides a large model construction device capable of presetting a KV cache capacity, which is applied to an attention layer in a Transformer structure of a large model, wherein the attention layer includes a key-value vector sequence MV consisting of M key-value vectors, and a key vector sequence MK consisting of M key vectors; wherein M is equal to a preset value, and the M key-value vectors are updated based on the large model construction method provided by the first aspect and its various implementations, and before the update, the key-value vector sequence MV is composed of initial values;

[0023] The device includes: a second mapping module, which is used to map the Nth input vector of the input attention layer to the current query vector q using a third matrix Wq; wherein N≥1, and the input vector is the embedding vector corresponding to the text tokenization of the input large model or the process vector output by the previous layer Transformer block of the current Transformer block; a second calculation module, which is used to calculate using the current query vector q and M key vectors to obtain the current weight vector w; a second acquisition module, which is used to obtain a key value vector sequence MV composed of initial values ​​from the cache when N=1, and to obtain a key value vector sequence MV corresponding to the N-1th input vector from the cache when N>1; an output module, which is used to perform weighted summation on the obtained key value vector sequence MV using the current weight vector w to obtain an output vector.

[0024] From the above content, it can be seen that the embodiment of the present application provides a large model construction method and device with preset KV cache capacity, which is applied to the attention layer in the Transformer structure of the large model, and the attention layer includes a key-value vector sequence MV composed of M key-value vectors, and a key vector sequence MK composed of M key vectors; wherein M is equal to a preset value, and the key-value vector sequence MV is composed of initial values; the method includes: for the Nth input vector of the input attention layer, using the first matrix Wwq to map it to a write query vector wq, and using the second matrix Wwv to map it to a first write key-value vector wv; wherein N≥1, the input vector is the corresponding embedding vector after the text of the input large model is tokenized or the previous layer Tr of the current Transformer block. The process vector output by the ansformer block; the write query vector wq is used to calculate with the M key vectors to obtain the write weight vector ww; when N=1, the key value vector sequence MV composed of the initial value is obtained from the cache, and it is determined as the historical key value vector sequence MV', when N>1, the key value vector sequence MV corresponding to the N-1th input vector is obtained from the cache, and it is determined as the historical key value vector sequence MV'; using the write weight vector ww and the first write key value vector wv, the M key value vectors in the historical key value vector sequence MV' are updated to obtain the updated key value vector sequence MV; the updated key value vector sequence MV is written into the cache to replace the historical key value vector sequence MV', and the updated key value vector sequence MV corresponds to the Nth input vector. The embodiment of the present application can use the MKV cache to replace the KV cache, and by continuously updating the key value vector sequence MV, the limited KV cache capacity can be used to cope with the infinite context length. A KV (ie, MKV) cache capacity solution with a preset length can be implemented to replace the KV cache capacity solution that grows infinitely with the context length, relieve cache pressure, solve memory constraints, and does not involve the step of discarding KV information. BRIEF DESCRIPTION OF THE DRAWINGS

[0025] Figure 1 A schematic diagram of a large model construction method provided in an embodiment of the present application;

[0026] Figure 2 A flowchart of a first large model construction method capable of presetting KV cache capacity provided in an embodiment of the present application;

[0027] Figure 3 A schematic diagram of a large model construction process provided in an embodiment of the present application;

[0028] Figure 4 A flowchart of a second large model construction method capable of presetting KV cache capacity provided in an embodiment of the present application;

[0029] Figure 5A schematic diagram of the structure of a first large model building device capable of presetting a KV cache capacity provided in an embodiment of the present application;

[0030] Figure 6 A schematic diagram of the structure of a second large model building device with preset KV cache capacity provided in an embodiment of the present application. DETAILED DESCRIPTION

[0031] In order to enable those skilled in the art to better understand the technical solutions in the present application, the technical solutions in the embodiments of the present application will be clearly and completely described below in conjunction with the drawings in the embodiments of the present application. Obviously, the described embodiments are only part of the embodiments of the present application, not all of the embodiments. Based on the embodiments in the present application, all other embodiments obtained by ordinary technicians in the field without creative work should fall within the scope of protection of the present application.

[0032] Before introducing the technical solutions of the embodiments of the present application, the terms involved in the embodiments of the present application are first introduced by way of example.

[0033] KV caching technology: A data storage technology. KV caching specifically refers to storing data in the form of key-value pairs in memory or video memory. The key is the unique identifier of the data, and the value is the data associated with the key. Since the access speed of memory or video memory is much faster than the hard disk, storing frequently accessed data in the cache can significantly increase the speed of data retrieval.

[0034] The construction of a large model includes two main stages: training and inference.

[0035] Training phase: In this phase, the large model is learned through a large amount of data. The training process usually involves data preparation, parameter tuning and other processes. Once the model training is completed, it can be used for reasoning in actual applications. The large model reasoning process: refers to the process of using a trained model (such as a language model) to generate a predicted output given data (such as input text).

[0036] The following is an illustrative introduction to the model reasoning process using a large-scale language model (LLM) as an example.

[0037] The reasoning process includes at least the following steps:

[0038] 1. Input processing:

[0039] (1) Text preprocessing: Convert the input text into a format that the model can understand. This process usually includes tokenization of text data. Tokenization is the process of dividing text data into smaller units (such as words, subwords, or characters). The resulting units are called tokens. Tokenization is a key step in text preprocessing, which helps the model better understand the underlying structure of the text and process it more efficiently.

[0040] (2) Embedding: Convert the tokenized text into an embedding vector (or word vector).

[0041] 2. Model Reasoning:

[0042] (1) Input embedding: The processed input text (embedding vector) is input into the neural network of the language model.

[0043] (2) Forward propagation: The language model processes input embedding through multiple neural network structures (such as the Transformer structure). Forward propagation mainly consists of two important components: the self-attention mechanism and the feed-forward neural network.

[0044] Self-attention mechanism: In the attention layer, the model uses the self-attention mechanism to calculate the attention weight of each token to other tokens, thereby capturing the dependencies and contextual information in the input sequence. It includes at least the following three steps:

[0045] ① Key-value pair conversion: The embedding vector corresponding to each token is converted into three different vectors: query vector, key vector, and key-value vector. These vectors are generated from the embedding through matrix multiplication.

[0046] ② Calculate attention weights: Calculate the dot product of the query vector and all key vectors, and then calculate the attention weight of each tag through the Softmax function.

[0047] ③ Weighted average: Use the attention weights to perform a weighted average on all value vectors to generate a contextual representation of the tag.

[0048] (3) Generate output: After multiple layers of processing, the model generates prediction distributions for each tag, which represent the model’s probability estimates for each possible output tag.

[0049] 3. Output generation: Select the output text generated by the model based on the generated probability distribution and predefined strategy.

[0050] It is understandable that the self-attention mechanism involves generating key and value vectors for each input word, that is, generating key-value pairs. As the vectors are continuously generated, the number of key-value pairs increases. In addition, the self-attention mechanism also involves obtaining previously generated key-value pairs for subsequent reasoning. Therefore, an efficient cache method is needed to store key-value pairs to ensure that previously calculated results can be reused during reasoning.

[0051] KV cache technology is an efficient caching method. Based on KV cache technology, the model can quickly access historical information when generating new content, thereby improving the overall reasoning efficiency and speed. In actual application, using KV cache technology as a caching method has become a common means to speed up the reasoning generation speed of large models (including large language models LLM).

[0052] When KV caching is applied to large model reasoning, the following two terms are specifically involved.

[0053] Key / value (KV): The collective name for two pairs of vector data (key and value) used in the attention layer in the neural network model transformer. In each attention layer of the large model network structure, one token generally corresponds to one KV.

[0054] Key-value cache (KV Cache): A sequence of one or more KV data, generally referring to the KV sequence of all tokens before the token currently being processed (it can also include the KV of the current token, without limitation).

[0055] KV cache example: The sentence "I love you very much" can be tokenized into [I, very, love, you]. When processing [love], the KV cache refers to the sequence of [I, very]; when processing [you], the KV cache refers to the sequence of [I, very, love]. It can be seen that the sequence length of the transformer's KV cache gradually increases as the token processing progresses.

[0056] However, as the sequence length increases, the cache requirements are also growing, which makes large model reasoning a memory-constrained problem, greatly limiting the throughput of the system. Using a limited-capacity KV cache to cope with infinite context reasoning becomes the key to solving this problem. Streaming LLM and Low-rank Embedding Sidekick with Sparse policy (LESS) are two common methods to implement limited-capacity KV cache.

[0057] The key idea of ​​StreamingLLM is to use window attention and change global attention to local attention, so that the KV cache that needs to be saved during reasoning can be kept at a fixed size, that is, kept at the window length. However, this solution will lead to incomplete utilization of historical KV information. When the context length is long, most of the KV information is discarded and cannot provide all the information for subsequent reasoning, which can easily lead to a decline in the effect of large models.

[0058] LESS refers to learning the residual between the original attention output and the attention output approximated by the sparse strategy, and achieves this by accumulating the information discarded by the sparse strategy into a low-level cache or state of constant size. The sparse strategy refers to retaining the KV information that is considered critical and discarding other unimportant KV information. However, this scheme adds additional low-level cache overhead and implementation complexity, and the weak importance of KV information does not mean that the KV information is completely useless. Directly discarding unimportant KV information will lead to a decrease in the inference effect of large models.

[0059] To solve the above problems, the embodiment of the present application first provides a large model construction method, which can be applied to the attention layer in the Transformer structure of the large model. The attention layer may include a third matrix Wq, a fourth matrix Wk and a fifth matrix Wv.

[0060] Figure 1 A schematic diagram of the process of constructing a large model provided in an embodiment of the present application.

[0061] like Figure 1 As shown, the large model construction method provided in the embodiment of the present application includes the following steps S101-S103.

[0062] S101: For the Nth input vector of the input attention layer, use the fourth matrix Wk to map the input vector to the current key vector k.

[0063] In the large model, the text input therein can first be converted into an embedding vector, each vector representing a word or subword unit in the text, i.e., a token. Afterwards, the embedding vector can be input as an input vector into the attention layer in the Transformer structure for the model building process. It is understandable that a piece of text can be tokenized into multiple embedding vectors (input vectors) to form a vector sequence, which is then input into the attention layer. That is, in step S100, the input vector is determined based on the text input into the large model, the Nth input vector is one of the vector sequence, and N ≥ 1.

[0064] It should also be noted that the Transformer structure is generally a multi-layer structure, which may include multiple attention layers. The attention layer in step S100 may be any layer of the multiple attention layers, not limited to the first layer. Accordingly, the input vector may be an embedding vector directly from the input text (the input of the first attention layer), or an intermediate vector processed by a previous attention layer.

[0065] The fourth matrix Wk is a weight matrix, which is used to calculate the key key of the input vector, that is, to map the input vector to the current key vector k. The dimension of the fourth matrix Wk can be [F, E], where F is the dimension of the input vector. Then, the dimension of the current key vector k is E. The input vector is mapped from the original F-dimensional space to a new E-dimensional space. The specific values ​​of F and E can be designed based on actual conditions, and the embodiments of the present application do not make specific limitations on this.

[0066] For the input vector x, the mapping formula is k=X·Wk; wherein Wk represents the fourth matrix, and k represents the current key vector.

[0067] S102: For the Nth input vector of the input attention layer, use the fifth matrix Wv to map the input vector to the current key-value vector v.

[0068] The fifth matrix Wv is also a weight matrix, which is used to calculate the key value of the input vector, that is, to map the input vector to the current key value vector v. The dimension of the fifth matrix can be [F, D]. Then, the dimension of the current key value vector v is D. The input vector is mapped from the original F-dimensional space to a new D-dimensional space. The specific value of D can be designed based on the actual situation, and the embodiment of the present application does not make specific limitations on this.

[0069] For the input vector x, the mapping formula is V=X·Wv; wherein Wv represents the fifth matrix, and V represents the current key-value vector v.

[0070] S103: For the Nth input vector of the input attention layer, use the third matrix Wq to map the input vector to the current query vector q.

[0071] The third matrix Wq is also a weight matrix, which is used to calculate the query query of the input vector, that is, to map the input vector to the current query vector q. The dimension of the third matrix Wq can be [F, E], then the dimension of the current query vector q is E. The input vector is mapped from the original F-dimensional space to a new E-dimensional space.

[0072] It is worth noting that for the current query vector q, the current key vector k and the current key value vector v, operations such as calculating similarity and normalization can also be performed. The specific steps will be described in detail below and will not be repeated here.

[0073] In some implementations, steps S101, S102, and S103 may be performed in other sequences, which is not specifically limited in the embodiments of the present application.

[0074] The embodiment of the present application can also provide a large model construction method with a preset KV cache capacity. The method can use a KV cache of a preset length (i.e., MKV) to cache KV information. As the construction proceeds, each time a new token is generated, the newly added KV information can be written (memorized) into the MKV in a target manner, and the target manner will not change the length of the MKV. It can be understood that the traditional KV cache step is to append the KV information corresponding to the current token to the corresponding position of the KV cache, and the sequence length of the KV cache will increase by 1 accordingly. The method provided in the embodiment of the present application can write the KV information corresponding to the current token into a KV cache of a preset length, i.e., MKV, without discarding the KV information, and can maintain the integrity of the KV information. It will not be limited by memory capacity, and can increase system throughput.

[0075] Specifically, the large model construction method with preset KV cache capacity provided in the embodiment of the present application can be applied to the attention layer in the Transformer structure of the large model, and the attention layer includes a key-value vector sequence MV composed of M key-value vectors, and a key vector sequence MK composed of M key vectors; wherein M is equal to a preset value, M≥1, and exemplarily, M can be equal to 1, 2, 3 or other, and can also be dynamically expanded during actual use, which is not specifically limited in the embodiment of the present application. The key-value vector sequence MV is composed of initial values, which can be obtained by training the large model, or can be all zero values ​​or random values, which is not specifically limited in the embodiment of the present application.

[0076] In an embodiment of the present application, the key-value vector sequence MV and the key vector sequence MK can be stored in a cache, and each element of the vector in the key-value vector sequence MV and the key vector sequence MK can have a preset value, which can be obtained by initializing the vector. The preset value can be, for example, 0 or 1, or other values, which is not specifically limited in the embodiment of the present application.

[0077] It should be noted that the dimension of the key value vector can be the same as the dimension of the current key value vector v, that is, both are D. Then, the key value vector sequence MV is equivalent to a matrix with a dimension of [M, D]. The dimension of the key vector can be the same as the dimension of the current key vector k, that is, both are E. Then, the key vector sequence MK is equivalent to a matrix with a dimension of [M, E] to ensure the feasibility of matrix calculation.

[0078] Figure 2A flowchart of the first large model construction method with preset KV cache capacity provided in an embodiment of the present application.

[0079] Figure 3 A schematic diagram of the large model construction process provided in an embodiment of the present application.

[0080] like Figure 2 and Figure 3 As shown, the large model construction method with preset KV cache capacity provided in the embodiment of the present application may include the following steps S201-S205.

[0081] S201: For the Nth input vector of the input attention layer, use the first matrix Wwq to map it to a write query vector wq, and use the second matrix Wwv to map it to a first write key-value vector wv; wherein N ≥ 1, and the input vector is determined based on the embedding of the text token of the input large model or the output of the previous layer Transformer block.

[0082] Among them, the Nth input vector of the input attention layer corresponds to the current token. The first matrix Wwq is a weight matrix, which is used to calculate the query query of the input vector, that is, to map the input vector to the write query vector wq (wquery). The dimension of the first matrix Wwq can be [D, E], so that the dimension of the write query vector wq is the same as the dimension of the key vector, and both are E.

[0083] The second matrix Wwv is also a weight matrix, which is used to calculate the key value of the input vector, that is, to map the input vector to the write key value vector wv(wvalue). The dimension of the second matrix Wwv can be [D, D], so that the dimension of the first write key value vector wv is the same as the dimension of the key value vector, and both are D.

[0084] S202: Calculate using the write query vector wq and the M key vectors to obtain a write weight vector ww.

[0085] Step S202 is specifically an attention calculation step, which at least includes the process of calculating similarity and normalization. Specifically, step S202 includes the following steps S2021-S2022.

[0086] S2021: Calculate the similarity between the written query vector wq and the M key vectors respectively to obtain M weight scores.

[0087] In the embodiment of the present application, the similarity may be calculated by a dot product or cosine similarity method, or by a method such as a Euclidean distance calculation, which is not specifically limited in the embodiment of the present application.

[0088] Exemplarily, M may be equal to 3, and the M key vectors mkey may be, for example, mkey1, mkey2, and mkey3, and the M weight scores may be score1, score2, and score3.

[0089] S2022: Normalize the M weight scores to obtain a write weight vector ww, and the dimension of the write weight vector ww is M.

[0090] In the embodiment of the present application, the Softmax function can be used for normalization. The sum of the normalized weight scores is equal to 1, and the normalized weight scores can form a write weight vector ww. Since the number of weight scores is M, the dimension of the write weight vector ww is M.

[0091] Exemplarily, when M is equal to 3, the normalized weight scores may be w1, w2, and w3.

[0092] It is worth noting that the above steps S2021-S2022 are exemplary introductions to step S202. In practical applications, other calculation steps may be used to calculate the write query vector wq and the M key vectors to obtain the write weight vector ww.

[0093] S203: When N=1, obtain the key value vector sequence MV composed of the initial value from the cache, and determine it as the historical key value vector sequence MV', that is, MV(0). When N>1, obtain the key value vector sequence MV corresponding to the N-1th input vector from the cache, and determine it as the historical key value vector sequence MV', that is, MV(N-1).

[0094] Among them, the N-1th input vector corresponds to the previous token of the current token. For example, the sentence "I love you very much" can be tokenized as [I, very, love, you]. When processing [love], [love] is the current token and [very] is the previous token.

[0095] It is worth noting that MV(0) is the initial form of MV, which is the historical key value vector sequence when N=1, and MV(N-1) is the historical key value vector sequence when N>1. MK, MV(0), Wwq, and Wwv are obtained by training a large model, where MV(0) can also be preset to all 0 values ​​or random values ​​and not participate in training updates.

[0096] As the construction proceeds, the key-value vector sequence MV can be continuously updated, and each input vector corresponds to a different key-value vector sequence MV. The details are as follows:

[0097] S204: Using the write weight vector ww and the first write key value vector wv, the M key value vectors in the historical key value vector sequence MV' are updated to obtain an updated key value vector sequence MV, namely MV(N).

[0098] The updating step may include the following steps S2041-S2043.

[0099] S2041: Transpose the write weight vector ww to convert the write weight vector ww into a write weight matrix ww' of dimension [M, 1], so that the number of columns of the write weight matrix ww' matches the number of rows of the first write key value vector wv.

[0100] It can be understood that the write weight vector ww is composed of M normalized weight scores, so the write weight vector ww can be converted into a write weight matrix ww' with a dimension of [M, 1] by transposing.

[0101] S2042: Perform matrix multiplication calculation on the write weight matrix ww′ and the first write key value vector wv to obtain M second write key value vectors wMV, and the dimension of the second write key value vector wMV is D.

[0102] According to the rules of matrix multiplication, the number of columns of the write weight matrix ww' matches the number of rows of the first write key value vector wv, both equal to 1, so the two can perform matrix multiplication, and since the dimension of the write weight matrix ww' is [M, 1], the dimension of the first write key value vector wv is [1, D], the dimension of the second write key value vector wMV is D, and M second write key value vectors wMV can form a matrix with a dimension of [M, D]. Figure 3 As shown, when M is equal to 3, the M second write key value vectors wMV can be wvalue1, wvalue2, and wvalue3 respectively.

[0103] S2043: Add the M second write-in key-value vectors wMV and the M key-value vectors in the historical key-value vector sequence MV' to obtain an updated key-value vector sequence MV, namely MV(N).

[0104] It can be understood that the M second written key value vectors wMV are added to the M key value vectors in the historical key value vector sequence MV' to obtain the updated key value vector sequence MV (new MV), which is equivalent to completing the memory process of the KV information corresponding to the current token. Exemplarily, when M is equal to 3, the updated key value vector sequence MV includes three key value vectors, namely new mvalue1, new mvalue2 and new mvalue3.

[0105] It is worth noting that when N is equal to 2, the key-value vector sequence MV corresponding to the previous input vector, ie, the first input vector, may be preset. The first input vector corresponds to the first token of the text input into the large model.

[0106] S205: Write the updated key-value vector sequence MV into the cache to replace the historical key-value vector sequence MV', wherein the updated key-value vector sequence MV corresponds to the Nth input vector.

[0107] In this way, the KV information can be cached. Further, when processing the next token of the current token, the key-value vector sequence MV of the Nth input vector can be obtained, and the key-value vector sequence MV can be updated using the N+1th input vector, which will not be described here.

[0108] It should be noted that steps S201-S205 can also be referred to as the process of writing, memorizing or updating MKV.

[0109] It should be noted that, like the traditional transformer, MKV can be fused with position coding, multi-head attention, group or group convolution, or any other form, and the embodiments of the present application do not impose any limitations on this.

[0110] It should be noted that the construction method provided in the embodiment of the present application only modifies the KV cache formation mechanism of the attention layer of the large model, and does not limit other parts of the large model.

[0111] It can be understood that for an inference process with an expected context length (including prompts and generated content) of L, using a KV cache with a preset length of M instead of a KV cache with a maximum length of L can compress the KV cache when M is less than L, thereby further reducing the video memory usage of large model inference and improving the speed of large model inference.

[0112] It should be noted that the process of using Wwq or Wwv for mapping can be replaced by a single-layer or multi-layer perceptron (MLP), and the embodiments of the present application do not specifically limit this.

[0113] It should also be noted that during the inference process, when processing token by token, MV can use the same memory block / video memory block for different N; during the training process, it can be N independent memory blocks / video memory blocks to facilitate parallel training.

[0114] It should also be noted that the updated MV is used as a KV cache for subsequent calculation processes, and the subsequent calculation process can be the same as the traditional large model method, and the embodiments of the present application do not make specific limitations on this.

[0115] From the above content, it can be seen that the embodiment of the present application provides a large model construction method with a preset KV cache capacity, which is applied to the attention layer in the Transformer structure of the large model, and the attention layer includes a key-value vector sequence MV composed of M key-value vectors, and a key vector sequence MK composed of M key vectors; wherein M is equal to a preset value, and the key-value vector sequence MV is composed of initial values; the method includes: for the Nth input vector of the input attention layer, using the first matrix Wwq to map it to a write query vector wq, and using the second matrix Wwv to map it to a first write key-value vector wv; wherein N≥1, the input vector is the corresponding embedding vector after the text of the input large model is tokenized or the previous layer Tra of the current Transformer block. The process vector output by the nsformer block; the write query vector wq is used to calculate with the M key vectors to obtain the write weight vector ww; when N=1, the key value vector sequence MV composed of the initial value is obtained from the cache, and it is determined as the historical key value vector sequence MV', when N>1, the key value vector sequence MV corresponding to the N-1th input vector is obtained from the cache, and it is determined as the historical key value vector sequence MV'; using the write weight vector ww and the first write key value vector wv, the M key value vectors in the historical key value vector sequence MV' are updated to obtain the updated key value vector sequence MV; the updated key value vector sequence MV is written into the cache to replace the historical key value vector sequence MV', and the updated key value vector sequence MV corresponds to the Nth input vector. The embodiment of the present application can use the MKV cache to replace the KV cache, and by continuously updating the key value vector sequence MV, the limited KV cache capacity can be used to cope with the infinite context length. A KV (ie, MKV) cache capacity solution with a preset length can be implemented to replace the KV cache capacity solution that grows infinitely with the context length, relieve cache pressure, solve memory constraints, and does not involve the step of discarding KV information.

[0116] Figure 4A flowchart of a second large model construction method with preset KV cache capacity provided in an embodiment of the present application.

[0117] like Figure 4 As shown, an embodiment of the present application also provides a large model construction method with a preset KV cache capacity, which can be combined with the two large model construction methods in the aforementioned embodiments, and applied to the attention layer in the Transformer structure of the large model to realize model construction. The attention layer may include a key-value vector sequence MV composed of M key-value vectors, and a key vector sequence MK composed of M key vectors; wherein M is equal to a preset value, and the M key-value vectors are updated based on the aforementioned first large model construction method with a preset KV cache capacity, and before the update, the key-value vector sequence MV is composed of initial values. For more introduction to the key-value vector sequence MV and the key vector sequence MK, please refer to the aforementioned content, which will not be repeated here.

[0118] Continue to see Figure 4 The large model construction method with preset KV cache capacity provided in the embodiment of the present application may include the following steps S301-S304. Steps S301-S304 are steps of performing attention operation on MKV.

[0119] S301: For the Nth input vector of the input attention layer, use the third matrix Wq to map the input vector to the current query vector q; wherein N ≥ 1, the input vector is the embedding vector corresponding to the tokenized text of the input large model or the process vector output by the previous layer of the Transformer block of the current Transformer block.

[0120] For the introduction of the third matrix Wq, please refer to the aforementioned S103, which will not be repeated here. The dimension of the third matrix Wq can be [F, E], where F is the dimension of the input vector, and the dimension of the current query vector q is E. The embodiment of the present application does not specifically limit the values ​​of F and E.

[0121] In some implementations, step S301 and step S103 may be performed alternatively.

[0122] S302: Calculate using the current query vector q and the M key vectors to obtain a current weight vector w.

[0123] In this step, the dot product calculation can be performed between the current query vector q and each key vector to obtain the attention scores. These scores represent the relevance of the current query vector q to each key vector. After that, the normalized SoftMax operation is performed on the attention scores to obtain the attention weights, which represent the degree of attention of the current query vector q to each key vector.

[0124] It can be understood that the current weight vector w is composed of M normalized attention scores, and the dimension of the current weight vector w is M.

[0125] S303: when N=1, obtain a key value vector sequence MV consisting of initial values ​​from the cache; when N>1, obtain a key value vector sequence MV corresponding to the N-1th input vector from the cache.

[0126] The introduction of the key-value vector sequence MV composed of initial values ​​and the key-value vector sequence MV corresponding to the N-1th input vector can refer to the aforementioned step S203, which will not be repeated here.

[0127] S304: Using the current weight vector w, perform weighted summation on the acquired key-value vector sequence MV to obtain an output vector.

[0128] Step S304 refers to using the attention weights to perform weighted averaging on the M key vectors in the key vector sequence MV to generate the final attention weighted output. The contribution of each key vector is determined by its corresponding attention weight. It can be understood that the key vector in the key vector sequence MV refers to the key vector of MKV, not the current key vector v in step S102.

[0129] The weighted sum calculation formula can be:

[0130]

[0131] Among them, Output represents the output vector, V i Represents the i-th key-value vector, 0<i≤M, Attention Represents the attention weight of the i-th key-value vector.

[0132] It should be noted that steps S201-S205 may also be referred to as the process of recalling or reading MKV.

[0133] In some implementations, the attention layer also includes a fourth matrix Wk of dimension [F, E], which is used to map the input vector to a current key vector k of dimension E; and the dimension of the key vector is the same as the dimension of the current key vector k.

[0134] In some implementations, the attention layer also includes a fifth matrix Wv of dimension [F, D], which is used to map the input vector to a current key-value vector v of dimension D; and the dimension of the key-value vector is the same as the dimension of the current key-value vector v.

[0135] Figure 5 A schematic diagram of the structure of a first large model building device with preset KV cache capacity provided in an embodiment of the present application.

[0136] like Figure 5 As shown, the embodiment of the present application provides a large model construction device capable of presetting the KV cache capacity, which is applied to the attention layer in the Transformer structure of the large model, and the attention layer includes a key-value vector sequence MV composed of M key-value vectors, and a key vector sequence MK composed of M key vectors; wherein M is equal to a preset value, and the key-value vector sequence MV is composed of initial values;

[0137] The device includes:

[0138] The first mapping module 1001 is used to map the Nth input vector of the input attention layer to a write query vector wq using the first matrix Wwq, and to map it to a first write key value vector wv using the second matrix Wwv; wherein N≥1, the input vector is the embedding vector corresponding to the tokenized text of the input large model or the process vector output by the previous layer Transformer block of the current Transformer block;

[0139] A first calculation module 1002 is used to calculate using the write query vector wq and the M key vectors to obtain a write weight vector ww;

[0140] The first acquisition module 1003, when N=1, is used to acquire a key value vector sequence MV composed of initial values ​​from the cache, and determine it as a historical key value vector sequence MV'; when N>1, is used to acquire a key value vector sequence MV corresponding to the N-1th input vector from the cache, and determine it as a historical key value vector sequence MV';

[0141] An updating module 1004 is used to update the M key value vectors in the historical key value vector sequence MV' by using the write weight vector ww and the first write key value vector wv to obtain an updated key value vector sequence MV;

[0142] The cache module 1005 is used to write the updated key value vector sequence MV into the cache to replace the historical key value vector sequence MV', and the updated key value vector sequence MV corresponds to the Nth input vector.

[0143] In some implementations, the dimension of the key-value vector is D, and the dimension of the key vector is E.

[0144] In some implementations, the dimension of the first matrix Wwq is [D, E], so that the dimension of the write query vector wq is the same as the dimension of the key vector, and both are E; the dimension of the second matrix Wwv is [D, D], so that the dimension of the first write key value vector wv is the same as the dimension of the key value vector, and both are D.

[0145] In some implementations, the first calculation module 1002 is specifically used to: respectively calculate the similarity between the write query vector wq and M key vectors to obtain M weight scores; normalize the M weight scores to obtain a write weight vector ww, and the dimension of the write weight vector ww is M.

[0146] In some implementations, the update module 1004 is specifically used to: transpose the write weight vector ww to convert the write weight vector ww into a write weight matrix ww' of dimension [M, 1], so that the number of columns of the write weight matrix ww' matches the number of rows of the first write key value vector wv; perform matrix multiplication calculation on the write weight matrix ww' and the first write key value vector wv to obtain M second write key value vectors wMV, and the dimension of the second write key value vector wMV is D; add the M second write key value vectors wMV and the M key value vectors in the historical key value vector sequence MV' to obtain an updated key value vector sequence MV.

[0147] In some implementations, M≥1; and, the key vector sequence MK, the first matrix Wwq and the second matrix Wwv are obtained by training a large model; the key value vector sequence MV is obtained by training a large model, or, the initial value of the key value vector sequence MV is all zero values ​​or random values.

[0148] Figure 6 A schematic diagram of the structure of a second large model building device with preset KV cache capacity provided in an embodiment of the present application.

[0149] like Figure 6 As shown, the embodiment of the present application also provides a large model construction device with preset KV cache capacity, which is applied to the attention layer in the Transformer structure of the large model, and the attention layer includes a key value vector sequence MV composed of M key value vectors, and a key vector sequence MK composed of M key vectors; wherein M is equal to a preset value, and the M key value vectors are updated based on the aforementioned first large model construction method with preset KV cache capacity, and before the update, the key value vector sequence MV is composed of initial values;

[0150] The device includes:

[0151] The second mapping module 2001 is used to map the Nth input vector of the input attention layer to the current query vector q using the third matrix Wq; wherein N ≥ 1, and the input vector is the embedding vector corresponding to the tokenized text of the input large model or the process vector output by the previous layer Transformer block of the current Transformer block;

[0152] The second calculation module 2002 is used to calculate using the current query vector q and the M key vectors to obtain the current weight vector w;

[0153] The second acquisition module 2003 is used to acquire a key value vector sequence MV composed of initial values ​​from the cache when N=1, and is used to acquire a key value vector sequence MV corresponding to the N-1th input vector from the cache when N>1;

[0154] The output module 2004 is used to perform weighted summation on the acquired key-value vector sequence MV using the current weight vector w to obtain an output vector.

[0155] In some implementations, the dimension of the third matrix Wq is [F, E], where F is the dimension of the input vector and the dimension of the current query vector q is E; the dimension of the current weight vector w is M; the attention layer also includes a fourth matrix Wk with a dimension [F, E], and the fourth matrix Wk is used to map the input vector to the current key vector k with a dimension E; and the dimension of the key vector is the same as the dimension of the current key vector k; the attention layer also includes a fifth matrix Wv with a dimension [F, D], and the fifth matrix Wv is used to map the input vector to the current key value vector v with a dimension D; and the dimension of the key value vector is the same as the dimension of the current key value vector v.

[0156] In a specific implementation, the present invention further provides a computer storage medium, wherein the computer storage medium may store a program, and when the program is executed, the program may include some or all of the steps in each embodiment of the large model construction method with preset KV cache capacity provided by the present invention. The storage medium may be a disk, an optical disk, a read-only memory (ROM) or a random access memory (RAM).

[0157] It is easy to understand that those skilled in the art can combine, split, reorganize, etc. the embodiments of the present application to obtain other embodiments based on the several embodiments provided in the present application, and these embodiments do not exceed the protection scope of the present application.

[0158] The above specific implementation methods further explain in detail the purpose, technical solutions and beneficial effects of the embodiments of the present application. It should be understood that the above are only specific implementation methods of the embodiments of the present application and are not used to limit the protection scope of the embodiments of the present application. Any modifications, equivalent substitutions, improvements, etc. made on the basis of the technical solutions of the embodiments of the present application should be included in the protection scope of the embodiments of the present application.

Claims

1. A method for constructing a large model capable of presetting KV cache capacity, characterized in that: An attention layer in the Transformer structure of the large model, the attention layer comprising a key-value vector sequence MV consisting of M key-value vectors, and a key-value vector sequence MK consisting of M key vectors; wherein M is equal to a preset value, and the key-value vector sequence MV is composed of initial values; The method comprises: For the Nth input vector input to the attention layer, use the first matrix Wwq to map it to a write query vector wq, and use the second matrix Wwv to map it to a first write key value vector wv; wherein N≥1, the input vector is the embedding vector corresponding to the tokenized text input to the large model or the process vector output by the previous layer of the Transformer block of the current Transformer block; Calculate using the write query vector wq and the M key vectors to obtain a write weight vector ww; When N=1, the key value vector sequence MV composed of the initial value is obtained from the cache, and it is determined as the historical key value vector sequence MV'; when N>1, the key value vector sequence MV corresponding to the N-1th input vector is obtained from the cache, and it is determined as the historical key value vector sequence MV'; Using the write weight vector ww and the first write key value vector wv, the M key value vectors in the historical key value vector sequence MV' are updated to obtain the updated key value vector sequence MV; The updated key-value vector sequence MV is written into the cache to replace the historical key-value vector sequence MV', and the updated key-value vector sequence MV corresponds to the Nth input vector.

2. The method for constructing a large model with preset KV cache capacity according to claim 1, characterized in that: The dimension of the key-value vector is D, and the dimension of the key vector is E.

3. The method for constructing a large model with preset KV cache capacity according to claim 2, characterized in that: The dimension of the first matrix Wwq is [D, E], so that the dimension of the write query vector wq is the same as the dimension of the key vector, and both are E; the dimension of the second matrix Wwv is [D, D], so that the dimension of the first write key value vector wv is the same as the dimension of the key value vector, and both are D.

4. The method for constructing a large model with preset KV cache capacity according to claim 3, characterized in that: The step of calculating using the write query vector wq and the M key vectors to obtain a write weight vector ww includes: Calculate the similarity between the write query vector wq and the M key vectors respectively to obtain M weight scores; The M weight scores are normalized to obtain the write weight vector ww, and the dimension of the write weight vector ww is M.

5. The method for constructing a large model capable of presetting KV cache capacity according to claim 4, characterized in that: The step of updating the M key value vectors in the historical key value vector sequence MV' by using the write weight vector ww and the first write key value vector wv to obtain the updated key value vector sequence MV includes: Transpose the write weight vector ww to convert the write weight vector ww into a write weight matrix ww' of dimension [M, 1], so that the number of columns of the write weight matrix ww' matches the number of rows of the first write key value vector wv; Performing matrix multiplication calculation on the write weight matrix ww′ and the first write key value vector wv to obtain M second write key value vectors wMV, and the dimension of the second write key value vector wMV is D; The M second write-in key-value vectors wMV and the M key-value vectors in the historical key-value vector sequence MV' are added together to obtain the updated key-value vector sequence MV.

6. The method for constructing a large model with preset KV cache capacity according to claim 1, characterized in that: M≥1； And, the key vector sequence MK, the first matrix Wwq and the second matrix Wwv are obtained by training the large model; The key-value vector sequence MV is obtained by training the large model, or the initial value of the key-value vector sequence MV is all zero values ​​or random values.

7. A method for constructing a large model capable of presetting KV cache capacity, characterized in that: An attention layer in the Transformer structure applied to the large model, the attention layer comprising a key-value vector sequence MV consisting of M key-value vectors, and a key-value vector sequence MK consisting of M key-value vectors; wherein M is equal to a preset value, the M key-value vectors are updated based on the large model construction method according to any one of claims 1 to 6, and before the update, the key-value vector sequence MV is composed of initial values; The method comprises: For the Nth input vector input to the attention layer, the input vector is mapped to the current query vector q using a third matrix Wq; wherein N ≥ 1, and the input vector is the embedding vector corresponding to the tokenized text input to the large model or the process vector output by the previous layer of the Transformer block of the current Transformer block; Calculate using the current query vector q and the M key vectors to obtain a current weight vector w; When N=1, the key value vector sequence MV composed of the initial values ​​is obtained from the cache; when N>1, the key value vector sequence MV corresponding to the N-1th input vector is obtained from the cache; The current weight vector w is used to perform weighted summation on the acquired key-value vector sequence MV to obtain an output vector.

8. The method for constructing a large model capable of presetting KV cache capacity according to claim 7, characterized in that: The dimension of the third matrix Wq is [F, E], where F is the dimension of the input vector and the dimension of the current query vector q is E; The dimension of the current weight vector w is M; The attention layer further includes a fourth matrix Wk of dimension [F, E], wherein the fourth matrix Wk is used to map the input vector to a current key vector k of dimension E; and the dimension of the key vector is the same as the dimension of the current key vector k; The attention layer also includes a fifth matrix Wv of dimension [F, D], which is used to map the input vector to a current key-value vector v of dimension D; and the dimension of the key-value vector is the same as that of the current key-value vector v.

9. A large model construction device capable of presetting KV cache capacity, characterized in that: An attention layer in the Transformer structure of the large model, the attention layer comprising a key-value vector sequence MV consisting of M key-value vectors, and a key-value vector sequence MK consisting of M key vectors; wherein M is equal to a preset value, and the key-value vector sequence MV is composed of initial values; The device comprises: A first mapping module is used to map the Nth input vector input to the attention layer into a write query vector wq using a first matrix Wwq, and to map it into a first write key value vector wv using a second matrix Wwv; wherein N≥1, the input vector is an embedding vector corresponding to the tokenized text input to the large model or a process vector output by a previous layer of the Transformer block of the current Transformer block; A first calculation module, configured to calculate using the write query vector wq and the M key vectors to obtain a write weight vector ww; A first acquisition module, when N=1, is used to acquire the key value vector sequence MV composed of the initial value from the cache, and determine it as the historical key value vector sequence MV'; when N>1, is used to acquire the key value vector sequence MV corresponding to the N-1th input vector from the cache, and determine it as the historical key value vector sequence MV'; An updating module, configured to update the M key value vectors in the historical key value vector sequence MV' by using the write weight vector ww and the first write key value vector wv to obtain the updated key value vector sequence MV; The cache module is used to write the updated key value vector sequence MV into the cache to replace the historical key value vector sequence MV', and the updated key value vector sequence MV corresponds to the Nth input vector.

10. A large model construction device capable of presetting KV cache capacity, characterized in that: An attention layer in the Transformer structure applied to the large model, the attention layer comprising a key-value vector sequence MV consisting of M key-value vectors, and a key-value vector sequence MK consisting of M key-value vectors; wherein M is equal to a preset value, the M key-value vectors are updated based on the large model construction method according to any one of claims 1 to 6, and before the update, the key-value vector sequence MV is composed of initial values; The device comprises: A second mapping module is used to map the Nth input vector input to the attention layer to a current query vector q using a third matrix Wq; wherein N ≥ 1, and the input vector is an embedding vector corresponding to the tokenized text input to the large model or a process vector output by a previous layer of the Transformer block of the current Transformer block; A second calculation module is used to calculate using the current query vector q and the M key vectors to obtain a current weight vector w; A second acquisition module, when N=1, used to acquire the key value vector sequence MV composed of the initial value from the cache; when N>1, used to acquire the key value vector sequence MV corresponding to the N-1th input vector from the cache; The output module is used to perform weighted summation on the acquired key-value vector sequence MV using the current weight vector w to obtain an output vector.

Images