Large model reasoning acceleration method and device, storage medium and terminal
By using KV Cache and the division of select propagation layer sets in large language models, the problems of slow model inference speed and large resource overhead are solved, and a more efficient inference process is achieved.
Patent Information
- Application Number
- CN202510228210.4
- Authority / Receiving Office
- CN · China
- Patent Type
- Applications(China)
- Current Assignee / Owner
- Filing Date
- 2025-02-27
- Publication Date
- 2025-06-17
AI Technical Summary
Large language models have high computational overhead when performing inference and slow in reasoning, especially when generating long text.
By obtaining input text data from the pre-trained large model, the control model calculates keys and values of all pre-filled tokens based on multiple attention layers, caches these keys and values to KV Cache. At least two consecutive sets of select propagation layers are divided among the multiple attention layers, and the selection propagation layer set includes a selection layer and a corresponding subsequent layer. Use KV Cache in the selection layer for calculations and determine the target token for the importance condition. Only the KV Cache subset corresponding to these target tokens is used for calculation in the subsequent layer.
By reducing the overhead of the model's use of tokens, the inference speed and efficiency are improved, while ensuring that the performance of the model is not affected.
Smart Images

Figure CN120163246A_ABST
Abstract
Description
Technical Field
[0001] The embodiments of this specification relate to the field of computer technology, and in particular, to a method, device, storage medium, and terminal for accelerating large model inference. Background Art
[0002] Large language models (LLMs) are neural networks designed to understand and generate human language. They are trained on large amounts of text data to learn patterns, context, and nuances in language. LLMs can perform various tasks, such as answering questions, writing articles, translating, and participating in conversations, which makes large language models widely used in many scenarios, such as chatbots, agents, and physical robots. However, since large language models perform inference token by token in a recursive manner, that is, if the model needs to output hundreds of tokens, it has to perform hundreds of inferences. This results in a very large computational overhead and greatly affects the inference speed when the model generates text, especially long text. Summary of the Invention
[0003] The embodiments of this specification provide a method, device, storage medium, and terminal for accelerating large model inference, which can solve the technical problems of slow model inference speed and large resource overhead in related technologies.
[0004] In a first aspect, the embodiments of this specification provide a method for accelerating large model inference, the method includes:
[0005] Obtain the input text data of the pre-trained large model, control the pre-trained large model to calculate the keys and values of all pre-filled tokens in the input text data based on multiple attention layers, and cache all keys and values to obtain the KV Cache;
[0006] Divide at least two consecutive sets of selection propagation layers from the multiple attention layers, where the set of selection propagation layers includes a selection layer and the subsequent layer corresponding to the selection layer, and the selection layer is the first layer in the set of selection propagation layers;
[0007] When calculating in each selection layer, obtain the KV Cache, use the KV Cache as the attention context of the current layer and perform calculations to obtain the output token, and determine the target tokens among all pre-filled tokens that meet the importance condition;
[0008] When calculating in the subsequent layer corresponding to each selection layer, obtain the KV Cache subset corresponding to the target token in the KV Cache, use the KV Cache subset as the attention context of the current layer and perform calculations to obtain the output token.
[0009] In a possible implementation, dividing at least two consecutive sets of selection and propagation layers from the multiple attention layers includes: determining the first M layers among the multiple attention layers as look-ahead layers, and dividing at least two consecutive sets of selection and propagation layers from the remaining attention layers, where M is a positive integer; when performing calculations in each look-ahead layer, obtaining the KV Cache, using the KV Cache as the attention context of the current layer and performing calculations to obtain an output token.
[0010] In a possible implementation, using the KV Cache as the attention context of the current layer and performing calculations to obtain an output token includes: receiving the output token of the previous layer as the current token, calculating the key and value of the current token; performing attention calculations based on the KV Cache and the key and value of the current token to obtain an output token; updating the key and value of the current token to the KV Cache.
[0011] In a possible implementation, determining target tokens that meet the importance condition among all pre-filled tokens includes: calculating the attention scores of each pre-filled token, and determining target tokens that meet the importance condition among all pre-filled tokens according to the attention scores of each pre-filled token.
[0012] In a possible implementation, calculating the attention scores of each pre-filled token and determining target tokens that meet the importance condition among all pre-filled tokens according to the attention scores of each pre-filled token includes: constructing a context hidden state window of a preset length, calculating the attention scores of each pre-filled token based on the context hidden state window; performing weighted summation on the attention scores of pre-filled tokens within each context hidden state window based on a weighted vector corresponding to a preset selection mode to obtain the context scores of each pre-filled token; selecting target tokens that meet the importance condition according to the context scores of each pre-filled token.
[0013] In a possible implementation, the preset selection mode is at least one of uniform selection, exponential selection, and last token selection, and the calculation method of the weighted vector is different in different preset selection modes.
[0014] In a possible implementation manner, caching all keys and values to obtain a KV Cache includes: caching all keys and values in the CPU memory to obtain a KV Cache; obtaining the KV Cache includes: loading the KV Cache in the CPU memory into the GPU memory for use in the calculation of the current layer; obtaining the KV Cache subset corresponding to the target token in the KV Cache includes: obtaining the KV Cache subset corresponding to the target token in the KV Cache from the CPU memory and loading the KV Cache subset into the GPU memory for use in the calculation of the current layer.
[0015] In a possible implementation manner, obtaining the KV Cache subset corresponding to the target token in the KV Cache, using the KV Cache subset as the attention context of the current layer and performing calculations to obtain an output token includes: during the process of loading the KV Cache subset into the GPU memory, when the complete KV Cache subset is not obtained, performing calculations using the loaded part of the KV Cache subset as the attention context of the current layer; when the complete KV Cache subset is obtained, using the KV Cache subset as the attention context of the current layer and performing calculations to obtain an output token.
[0016] In a second aspect, an embodiment of this specification provides a large model inference acceleration device, and the device includes:
[0017] A KV calculation cache module, configured to obtain the input text data of the pre-trained large model, control the pre-trained large model to calculate the keys and values of all pre-filled tokens in the input text data based on multiple attention layers, and cache all keys and values to obtain a KV Cache;
[0018] A multi-attention layer division module, configured to divide at least two consecutive selected propagation layer sets from the multiple attention layers, where the selected propagation layer set includes a selected layer and the subsequent layer corresponding to the selected layer, and the selected layer is the first layer in the selected propagation layer set;
[0019] A selected layer calculation module, configured to, when performing calculations in each selected layer, obtain the KV Cache, use the KV Cache as the attention context of the current layer and perform calculations to obtain an output token, and determine the target tokens among all pre-filled tokens that meet the importance condition;
[0020] A subsequent layer calculation module, configured to obtain a KV Cache subset corresponding to the target token in the above KV Cache when performing calculations in subsequent layers corresponding to each selected layer, use the KV Cache subset as the attention context of the current layer, and perform calculations to obtain an output token.
[0021] In a third aspect, an embodiment of this specification provides a computer program product containing instructions. When the computer program product runs on a computer or a processor, the computer or the processor is caused to execute the steps of the above method.
[0022] In a fourth aspect, an embodiment of this specification provides a computer storage medium. The computer storage medium stores multiple instructions, and the instructions are adapted to be loaded and executed by a processor to perform the steps of the above method.
[0023] In a fifth aspect, an embodiment of this specification provides a terminal, including a memory, a processor, and a computer program stored on the memory and executable on the processor. The computer program is adapted to be loaded and executed by the processor to perform the steps of the above method.
[0024] The beneficial effects brought by the technical solutions provided in some embodiments of this specification at least include:
[0025] An embodiment of this specification provides a large model inference acceleration method. The input text data of a pre-trained large model is obtained, and the pre-trained large model is controlled to calculate the keys and values of all pre-filled tokens in the input text data based on multiple attention layers, and all the keys and values are cached to obtain a KV Cache. At least two consecutive selected propagation layer sets are divided in the multiple attention layers. A selected propagation layer set includes a selected layer and the subsequent layer corresponding to the selected layer, and the selected layer is the first layer in the selected propagation layer set. When calculating in each selected layer, the KV Cache is obtained, the KV Cache is used as the attention context of the current layer and calculations are performed to obtain an output token, and target tokens that meet the importance condition among all the pre-filled tokens are determined. When calculating in the subsequent layer corresponding to each selected layer, the KV Cache subset corresponding to the target tokens in the KV Cache is obtained, and the KV Cache subset is used as the attention context of the current layer and calculations are performed to obtain an output token. In the embodiment of this specification, multiple attention layers in the large model are divided into several selected propagation layer sets. The first layer in each set is the selected layer, and the rest are the subsequent layers. In each selected layer, the KV of all tokens is used for calculation, and at the same time, important tokens among the pre-filled tokens are also selected. In the subsequent layer corresponding to it, only the KV subset corresponding to these important tokens needs to be used for calculation, and there is no need to use all the pre-filled tokens. Since there is a sparse pattern similarity between adjacent layers, the important tokens are also similar for these layers. Therefore, after the important tokens are selected in the previous selected layer, the selected important tokens can be directly used in multiple subsequent layers, thereby reducing the token usage overhead of the model while ensuring that the performance of the model is not affected. In addition, for the large model, the important tokens are selected multiple times based on multiple selected layers, which is beneficial for the large model to adapt to the importance changes of each token generated during the iteration process. BRIEF DESCRIPTION OF THE DRAWINGS
[0026] In order to more clearly illustrate the technical solutions in the embodiments of this specification or the prior art, the following will briefly introduce the drawings required for the description of the embodiments or the prior art. Obviously, the following drawings are only some embodiments of this specification. For those skilled in the art, without creative efforts, other drawings can also be obtained based on these drawings.
[0027] Figure 1 It is an exemplary system architecture diagram of a large model inference acceleration method provided by an embodiment of this specification;
[0028] Figure 2 It is a flowchart of a large model inference acceleration method provided by an embodiment of this specification;
[0029] Figure 3 A flowchart showing the process of a large model inference acceleration method provided by an embodiment of this specification;
[0030] Figure 4 A flowchart showing the usage process of KV Cache in a large model inference acceleration method provided by an embodiment of this specification;
[0031] Figure 5 A structural block diagram of a large model inference acceleration device provided by an embodiment of this specification;
[0032] Figure 6 A structural diagram of a terminal provided by an embodiment of this specification. Detailed implementation manners
[0033] To make the features and advantages of the embodiments of this specification more obvious and understandable, the technical solutions in the embodiments of this specification will be clearly and completely described below in conjunction with the accompanying drawings in the embodiments of this specification. Obviously, the described embodiments are only a part of the embodiments of this specification, rather than all embodiments. Based on the embodiments in this specification, all other embodiments obtained by those skilled in the art without creative efforts belong to the scope protected by the embodiments of this specification.
[0034] When the following description involves the accompanying drawings, unless otherwise indicated, the same numbers in different drawings represent the same or similar elements. The implementation manners described in the following exemplary embodiments do not represent all implementation manners consistent with the embodiments of this specification. On the contrary, they are only examples of devices and methods consistent with some aspects of the embodiments of this specification as detailed in the appended claims. And in the description of the embodiments of this specification, unless otherwise stated, " / " means "or". For example, A / B can represent A or B. The "and / or" in the text is only a description of the association relationship of the associated objects, indicating that there can be three relationships. For example, A and / or B can represent: A exists alone, A and B exist simultaneously, and B exists alone. In addition, in the description of the embodiments of this specification, "a plurality of" means two or more than two.
[0035] Hereinafter, the terms "first" and "second" are only used for descriptive purposes and cannot be construed as implying or suggesting relative importance or implicitly indicating the quantity of the indicated technical features. Thus, the features defined with "first" and "second" may explicitly or implicitly include one or more of such features.
[0036] Almost all current large language models are built based on the Transformer architecture. Therefore, based on the working method of the Transformer, the model needs to generate each token (the smallest text unit) of the text one by one, and each token depends on the previous context. This means that to predict the 1000th token, the information of the previous 999 tokens is required, which usually involves a series of matrix multiplication Q operations on the representations of these tokens. However, after obtaining a token, the previously calculated context information cannot be remembered. When inferring the next token, the previous context information needs to be recalculated again. That is, when predicting the 1001st token, not only the information of the previous 999 tokens needs to be recalculated, but also the information of the 1000th token needs to be added. Therefore, considering that when the model calculates the attention score Attention, the K (Key) and V (Value) calculated for the context tokens will be repeatedly calculated and used during the calculation process, by introducing KVcache, the calculation results of the previous K and V can be stored and reused during the subsequent token generation, thus avoiding a large amount of repeated calculations.
[0037] However, with the surge in the demand for long text processing in large language models, the increase in the number of tokens has brought about a large-scale KV Cache, which results in the model needing to use a large amount of GPU memory to store the KV Cache and occupying more memory as the sequence length increases. For example, in the case of some models with a batch size of 8 and a context length of 128k, the KV Cache alone occupies more than 134GB of GPU memory.
[0038] To alleviate the GPU memory usage of the KV Cache, a common approach is to compress the KV Cache to reduce video memory consumption and improve efficiency. Among them, one method is to directly compress the input text at the language level. An additional language model is used to select and discard the relatively less important text in the input, reducing the length of the input text and thus the KV Cache. Although the model tries to retain key information as much as possible, some details and context are inevitably lost during the compression process, resulting in poor performance in some scenarios that require precise details or in certain specific fields, especially larger deviations in highly specialized or technical content. Moreover, since an additional language model is introduced for compression, it also increases the additional resource consumption and the complexity of the overall system. Another method is to directly compress the KV Cache, that is, to select and use the KV Cache of some important tokens while discarding the others. Since only the importance of each token's current state can be considered during discarding, it may lead to the discarding of the KV of tokens that are currently less important but may be important in the future. The loss of these key information will further lead to a decline in the inference performance of the model.
[0039] Therefore, the embodiments of this specification provide a large model inference acceleration method to solve the above technical problems of slow model inference speed and large resource overhead.
[0040] Please refer to Figure 1 , Figure 1 which is an exemplary system architecture diagram of a large model inference acceleration method provided by the embodiments of this specification.
[0041] As Figure 1 shown, the system architecture may include a terminal 101, a network 102, and a server 103. The network 102 is used to provide a medium for the communication link between the terminal 101 and the server 103. The network 102 may include various types of wired communication links or wireless communication links. For example, wired communication links include optical fibers, twisted pairs, or coaxial cables, and wireless communication links include Bluetooth communication links, Wireless-Fidelity (Wi-Fi) communication links, or microwave communication links, etc.
[0042] The terminal 101 can interact with the server 103 via the network 102 to receive messages from the server 103 or send messages to the server 103. Alternatively, the terminal 101 can interact with the server 103 via the network 102 to receive messages or data sent by other users to the server 103. The terminal 101 can be various electronic devices, including but not limited to tablet computers, laptop portable computers, and desktop computers, etc. And the terminal 101 can be a single device or a group of devices composed of multiple devices. Additionally, in the embodiments of this specification, a pre-trained large model for text understanding and generation is installed in the terminal 101.
[0043] In the embodiments of this specification, the terminal 101 first obtains the input text data of the pre-trained large model, controls the pre-trained large model to calculate the keys and values of all pre-filled tokens in the input text data based on multiple attention layers, and caches all the keys and values to obtain the KV Cache. During the inference process of the model, the terminal 101 divides at least two consecutive sets of selection propagation layers in the multiple attention layers. Each set of selection propagation layers includes a selection layer and the subsequent layer corresponding to the selection layer, and the selection layer is the first layer in the set of selection propagation layers. When calculating in each selection layer, the pre-trained large model in the terminal 101 obtains the KV Cache, uses the KV Cache as the attention context of the current layer and performs calculations to obtain the output token, and determines the target tokens among all the pre-filled tokens that meet the importance condition. When calculating in the subsequent layers corresponding to each selection layer, the pre-trained large model in the terminal 101 obtains the KV Cache subset corresponding to the target tokens in the KV Cache, uses the KV Cache subset as the attention context of the current layer and performs calculations to obtain the output token.
[0044] In the embodiments of this specification, the server 103 can be a business server that provides various services. It should be noted that the server 103 can be hardware or software. When the server 103 is hardware, it can be implemented as a distributed server cluster composed of multiple servers or as a single server. When the server 103 is software, it can be implemented as multiple software or software modules (such as those used to provide distributed services) or as a single software or software module, and specific limitations are not made here.
[0045] Alternatively, the system architecture may not include the server 103. In other words, the server 103 can be an optional device in the embodiments of this specification. That is, the method provided in the embodiments of this specification can be applied to a system structure that only includes the terminal 101, and specific limitations are not made in the embodiments of this specification.
[0046] It should be understood that Figure 1The numbers of terminals, networks, and servers in [it] are only illustrative. According to implementation requirements, there can be any number of terminals, networks, and servers.
[0047] Please refer to Figure 2 , Figure 2 which is a schematic flowchart of a large model inference acceleration method provided by an embodiment of this specification. The execution subject of the embodiment of this specification can be a terminal that executes large model inference acceleration, or a processor in the terminal that executes the large model inference acceleration method, or the large model inference acceleration service in the terminal that executes the large model inference acceleration method. For ease of description, below, taking the execution subject as the processor in the terminal as an example, the specific execution process of the large model inference acceleration method is introduced.
[0048] As Figure 2 shown, the large model inference acceleration method can at least include:
[0049] S202. Obtain the input text data of the pre-trained large model, control the pre-trained large model to calculate the keys and values of all pre-filled tokens in the input text data based on multiple attention layers, and cache all the keys and values to obtain the KV Cache.
[0050] Optionally, in the pre-filling stage of the model, when the model receives the input text data, except for the output of the last token of the input text data, the outputs of other tokens are not needed. That is, in the token sequence of the entire question, only the output of the last token of the question is needed because this is the first token of the answer. Therefore, when the model performs attention calculations on each pre-filled token in the input text data through multiple attention layers, the keys and values of all pre-filled tokens can be obtained, and then all the keys and values are cached to obtain the KV Cache.
[0051] Optionally, assume the model has a total of P layers. Then the model creates the initialized KV Cache by using the hidden state h i of each pre-filled token at the i-th layer and the attention projection matrices W i,k and W i,v : Generate tensors of K and V with a dimension of R H×N×d . Wherein, N represents the length of the pre-filled context, that is, the length of the input text data, H represents the number of attention heads in the network, and d represents the hidden size of each attention head. In the subsequent decoding stage of the model, that is, when outputting new tokens, the KV Cache of the previously cached pre-filled context can be directly read to avoid a large amount of repeated calculations.
[0052] S204. Divide at least two consecutive sets of selection propagation layers from multiple attention layers. Each set of selection propagation layers includes a selection layer and the subsequent layer corresponding to the selection layer, and the selection layer is the first layer in the set of selection propagation layers.
[0053] Optionally, as the sequence length increases, on the one hand, the problem of OOD (out-of-distribution) is likely to occur, and on the other hand, the attention calculation amount will also increase rapidly. To reduce the computational and storage requirements during the model calculation process, considering the sparsity of the attention matrix and the fact that for long sequences, the sparsity degree will further expand, the embodiments of this specification solve the problem based on the sparsity of the attention matrix.
[0054] Specifically, the model can utilize this property of sparsity and perform attention calculation on some important tokens through Selective Attention, thereby significantly reducing the calculation amount. However, when the model selects important tokens, the more accurate the token selection method is, the more flexible and precise it can locate key information. However, the greater the computational overhead introduced by the more accurate token selection method. Therefore, how to balance the selection efficiency of important tokens becomes particularly important.
[0055] Furthermore, most of the inputs and calculation contents between adjacent layers in the model are the same. Therefore, the sparse patterns between these adjacent layers are also similar, that is, the important tokens considered by each layer from the second layer to the eighth layer of the model are almost the same as those considered by the first layer. From the similarity of the sparse patterns between adjacent layers, it can be seen that after the model selects important tokens at a certain layer, in some subsequent layers, the selection results of these important tokens can be directly used without further selection. In this way, only a few selections are required in the entire model system to propagate the selection results to most layers that do not require selection, significantly reducing unnecessary computational overhead. For example, for most models, only less than or equal to three selections are required.
[0056] In the embodiments of this specification, at least two consecutive sets of selection propagation layers can be divided from multiple attention layers of the model through means such as experiments. In each set of selection propagation layers, it includes a selection layer and the subsequent layer corresponding to the selection layer. The selection layer is the first layer in the set of selection propagation layers, and the similarity of the sparse patterns between the selection layer and the subsequent layer can reach a preset threshold, thereby ensuring that the propagation of important tokens does not affect the original performance of the model.
[0057] S206. When performing calculations in each selection layer, obtain the KV Cache, use the KV Cache as the attention context for the current layer and perform calculations to obtain the output token, and determine the target tokens among all pre-filled tokens that meet the importance condition.
[0058] Optionally, when performing calculations in each selection layer, it still needs to load the full KV Cache of the pre-filled tokens by itself, use it as the attention context for the current layer and perform calculations to obtain the output token. At the same time, this layer also needs to select important tokens, that is, determine the target tokens among all pre-filled tokens that meet the importance condition.
[0059] S208. When performing calculations in the subsequent layers corresponding to each selection layer, obtain the KV Cache subset corresponding to the target tokens in the KV Cache, use the KV Cache subset as the attention context for the current layer and perform calculations to obtain the output token.
[0060] Correspondingly, when performing calculations in the subsequent layers corresponding to each selection layer, according to the target tokens determined by the corresponding selection layer, load the KV Cache subset corresponding to these target tokens in the KV Cache, and use the KV Cache subset as the attention context for the current layer and perform calculations to obtain the output token. During the calculation process of these subsequent layers, there is no need to calculate the right KV Cache, nor to select important tokens, which reduces the usage overhead of tokens by the model while ensuring that the performance of the model is not affected. And the embodiments of this specification also dynamically select multiple sets of selection propagation layers for the large model, and perform multiple selections of important tokens based on multiple selection layers, which is beneficial for the large model to adapt to the importance changes of each token generated during the iteration process.
[0061] In the embodiments of this specification, a method for accelerating large model inference is provided. The input text data of the pre-trained large model is obtained, and the pre-trained large model is controlled to calculate the keys and values of all pre-filled tokens in the input text data based on multiple attention layers, and all the keys and values are cached to obtain the KV Cache. At least two consecutive sets of selective propagation layers are divided in the multiple attention layers. A set of selective propagation layers includes a selection layer and the subsequent layer corresponding to the selection layer, and the selection layer is the first layer in the set of selective propagation layers. When calculating in each selection layer, the KV Cache is obtained, and the KV Cache is used as the attention context of the current layer to perform calculations to obtain the output token, and the target tokens that meet the importance condition among all the pre-filled tokens are determined. When calculating in the subsequent layers corresponding to each selection layer, the KVCache subset corresponding to the target tokens in the KV Cache is obtained, and the KV Cache subset is used as the attention context of the current layer to perform calculations to obtain the output token. In the embodiments of this specification, multiple attention layers in the large model are divided into several sets of selective propagation layers. The first layer in each set is the selection layer, and the rest are the subsequent layers. In each selection layer, the KV of all tokens is used for calculation, and at the same time, the important tokens among the pre-filled tokens are selected. In the subsequent layers corresponding to it, only the KV subset corresponding to these important tokens needs to be used for calculation, and there is no need to use all the pre-filled tokens. Since there is a sparse pattern similarity between adjacent layers, the important tokens are also similar for these layers. Therefore, when the important tokens are selected in the previous selection layer, the selected important tokens can be directly used in multiple subsequent layers, thus reducing the token usage overhead of the model while ensuring that the performance of the model is not affected. In addition, for the large model, the important tokens are selected multiple times based on multiple selection layers, which is beneficial for the large model to adapt to the importance changes of each token generated during the iteration process.
[0062] Please refer to Figure 3 , Figure 3 which is a schematic flowchart of a method for accelerating large model inference provided by the embodiments of this specification.
[0063] As Figure 3 shown, the method for accelerating large model inference may at least include:
[0064] S302. Obtain the input text data of the pre-trained large model, control the pre-trained large model to calculate the keys and values of all pre-filled tokens in the input text data based on multiple attention layers, and cache all the keys and values in the CPU memory to obtain the KV Cache.
[0065] Optionally, for the process of calculating the keys and values of all pre-filled tokens, please refer to the detailed description in step S202, which will not be elaborated here. Usually, the KV Cache is directly cached in the GPU memory, but this leads to excessive occupancy of the video memory. Then, please refer to Figure 4 , Figure 4 , which is a schematic diagram of the usage process of the KV Cache in a large model inference acceleration method provided by an embodiment of this specification. As Figure 4 shown, when caching the KV Cache, the KV Cache cached in the GPU memory during generation can be unloaded to the CPU memory. When the model needs to use the KV Cache, the required part can be loaded from the CPU memory to the GPU memory. This significantly reduces the video memory usage of the KV Cache. At the same time, since there is no need to compress the tokens, it will not cause performance degradation, thus supporting the model to maintain or even accelerate the original inference efficiency in the long text scenario.
[0066] S304. Determine the first M layers in multiple attention layers as look-ahead layers, and divide at least two consecutive sets of selection propagation layers in the remaining attention layers, where M is a positive integer. When performing calculations in each look-ahead layer, obtain the KV Cache, use the KV Cache as the attention context of the current layer and perform calculations to obtain the output tokens.
[0067] Optionally, the sparsity of the first few layers in the model is usually low. Then, in order to maintain the original performance as much as possible while maintaining efficiency, the first M layers in multiple attention layers can be determined as look-ahead layers, where M is a positive integer, and then a set of selection propagation layers L is divided in the remaining layers. These look-ahead layers also need to perform full attention calculations on all tokens and the complete KV Cache. However, different from the selection layers, there are no subsequent layers that need to propagate tokens, so there is no need to select important tokens. Please continue to refer to Figure 4 . The first M layers of the model need to load the complete KV Cache from the CPU memory and do not perform selection. The value of M can be determined according to actual needs and can be 3, 4, etc.
[0068] S306. When performing calculations in each selection layer, load the KV Cache in the CPU memory to the GPU memory for use in the current layer's calculations.
[0069] Please continue to refer to Figure 4 . When performing calculations in each selection layer, load the KV Cache in the CPU memory to the GPU memory for use in the current layer's calculations, and these KV Caches will also be used to select important tokens.
[0070] S308. Receive the output token from the previous layer as the current token, calculate the key and value of the current token; perform attention calculation based on the KV Cache and the key and value of the current token to obtain the output token; update the key and value of the current token to the KV Cache.
[0071] Further, please continue to refer to Figure 4 , when the model performs calculations, it receives the output token from the previous layer as the current token, calculates the key and value for the current token, that is, K and V, and then performs attention calculation based on the KV Cache and the key and value of the current token to obtain the output token. After obtaining the new output token, update the key and value of the current token calculated in this round to the KV Cache for subsequent calculations.
[0072] S310. Calculate the attention scores of each pre-filled token, and determine the target tokens that meet the importance condition among all pre-filled tokens according to the attention scores of each pre-filled token.
[0073] Optionally, when selecting important tokens, as Figure 4 shown, it can be implemented through a ContextSelector. When selecting the propagation layer set L, where the first layer is denoted as L0, and the target tokens that meet the importance condition are denoted as T i , i ∈ L, then the target tokens of each layer in the selected propagation layer set L can be obtained as:
[0074]
[0075] Specifically, within the ContextSelector, for the context hidden state h ∈ R (N×d) , maintain a context hidden state window h ω with a preset length, and initialize a context score S ∈ R N . Use the local context window as Q, the full context as K, and according to the pre-filled context length h c and the current window length h ω h ω calculate the corresponding attention scores A for each pre-filled token:
[0076] Q = W q h ω , K = W k h c ,
[0077] Further, to select important target tokens, first use the max pooling algorithm to obtain the maximum attention scores across attention heads, and then weight and sum the attention scores of the pre-filled tokens in each context hidden state window based on the weighted vector α corresponding to the preset selection mode to obtain the context scores of each pre-filled token. Finally, select the top_k tokens as the target tokens that meet the importance condition:
[0078]
[0079] T = argtop k (S), T ∈ R k ;
[0080] Among them, in different preset selection modes, the calculation method of the weighted vector α is different. Different weighted vector α values can assign different weights to the tokens in the window, so as to achieve different selection effects. Three types are provided in the embodiments of this specification, including uniform selection, exponential selection, and last token selection. In the uniform selection method, This means that each token in the window h ω makes the same contribution when weighting and summing the attention scores. This method is relatively conventional. In the exponential selection method, This method means that the tokens closer to the end of the window h ω make higher contributions. This method has good comprehensiveness and can be applied to most scenarios, especially code task scenarios. In the last token selection method, This method means that only the attention score of the last token in the window h ω is considered. This method is especially suitable for question-and-answer task scenarios. In practical applications, the weighted vector α can also have other calculation methods and calculation effects, which are not limited in the embodiments of this specification.
[0081] S312. When calculating in the subsequent layers corresponding to each selection layer, obtain the KV Cache subset corresponding to the target token in the KV Cache from the CPU memory and load the KV Cache subset into the GPU memory for use in the current layer calculation.
[0082] Optionally, as Figure 4 shown, after identifying the target token: T i in each selection layer, its corresponding subsequent layer l (where l < L i+1 ) can directly load the KV Cache subset corresponding to T i from the CPU memory: K i [T i and Vi [T i into the GPU memory.
[0083] S314. During the process of loading the KV Cache subset into the GPU memory, when the complete KV Cache subset is not obtained, calculations are performed based on the already loaded part of the KV Cache subset as the attention context for the current layer; when the complete KV Cache subset is obtained, the KV Cache subset is used as the attention context for the current layer and calculations are performed to obtain the output token.
[0084] Optionally, during the calculation process of subsequent layers, considering that it takes time to transfer the KV Cache subset of the target token from the CPU content to the GPU memory, in order to avoid unnecessary waiting when loading the KV Cache subset K i [T i and V i [T i , for the layers that have not received the complete KV Cache subset K i [T i and V i [T i , the already loaded part is used for complete attention calculations, so that meaningful calculations are also performed during the process of loading the complete data, achieving the effect of covering the data transfer time with effective calculations:
[0085]
[0086] From the above formula, it can be known that when the complete KV Cache subset K i [T i and V i [T i is not obtained, complete attention calculations are performed first, and when the complete KV Cache subset K i [T i and V i [T i is obtained, calculations are performed using the complete KV Cache subset to obtain the output token.
[0087] In the embodiments of this specification, a method for accelerating large model inference is provided. The generated KV Cache is unloaded to the CPU memory. When the model needs to use the KV Cache, the required part is loaded from the CPU memory to the GPU memory. This significantly reduces the video memory usage of the KV Cache. At the same time, since there is no need to compress the tokens, there will be no performance degradation. Thus, it supports the model to maintain or even accelerate the original inference efficiency in the long text scenario. Considering that the first few layers in the model usually have low sparsity, then complete attention calculations are performed on these low-sparsity look-ahead layers using all tokens, so as to maintain the original performance to the greatest extent while maintaining efficiency. Through the context selector, important target tokens are selected based on the attention scores of the tokens. When selecting, the embodiments of this specification provide three possible weighted summation calculation methods, so that the context selector can select appropriate target tokens in various different scenarios. In addition, when performing calculations, for the layers that have not received the complete KV Cache subset, complete attention calculations are performed using the loaded part, so as to cover the data transmission time with effective calculations and prevent the model from entering a meaningless waiting state.
[0088] Please refer to Figure 5 , Figure 5 which is a structural block diagram of a large model inference acceleration device provided by the embodiments of this specification. As Figure 5 shown, the large model inference acceleration device 500 includes:
[0089] The KV calculation cache module 510 is used to obtain the input text data of the pre-trained large model, control the pre-trained large model to calculate the keys and values of all pre-filled tokens in the input text data based on multiple attention layers, and cache all the keys and values to obtain the KV Cache;
[0090] The multi-attention layer division module 520 is used to divide at least two consecutive selection propagation layer sets in the multiple attention layers. The selection propagation layer set includes a selection layer and the subsequent layers corresponding to the selection layer, and the selection layer is the first layer in the selection propagation layer set;
[0091] The selection layer calculation module 530 is used to, when performing calculations in each selection layer, obtain the KV Cache, use the KV Cache as the attention context of the current layer and perform calculations to obtain the output tokens, and determine the target tokens that meet the importance conditions among all the pre-filled tokens;
[0092] The subsequent layer calculation module 540 is used to obtain the KV Cache subset corresponding to the target token in the KV Cache when calculating in the subsequent layers corresponding to each selected layer, use the KV Cache subset as the attention context of the current layer, and perform calculations to obtain the output token.
[0093] Optionally, the multi-attention layer division module 520 is further used to determine the first M layers in the multiple attention layers as the look-ahead layers, and divide at least two consecutive selected propagation layer sets in the remaining attention layers, where M is a positive integer; when calculating in each look-ahead layer, obtain the KV Cache, use the KV Cache as the attention context of the current layer, and perform calculations to obtain the output token.
[0094] Optionally, the selected layer calculation module 530 is further used to receive the output token of the previous layer as the current token, calculate the key and value of the current token; perform attention calculation based on the KV Cache and the key and value of the current token to obtain the output token; update the key and value of the current token to the KV Cache.
[0095] Optionally, the selected layer calculation module 530 is further used to calculate the attention scores of each pre-filled token, and determine the target token that meets the importance condition among all pre-filled tokens according to the attention scores of each pre-filled token.
[0096] Optionally, the selected layer calculation module 530 is further used to construct a context hidden state window with a preset length, calculate the attention scores of each pre-filled token based on the context hidden state window; perform weighted summation on the attention scores of the pre-filled tokens within each context hidden state window based on the weighted vector corresponding to the preset selection mode to obtain the context scores of each pre-filled token; select the target token that meets the importance condition according to the context scores of each pre-filled token.
[0097] Optionally, the preset selection mode is at least one of uniform selection, exponential selection, and last token selection, and the calculation method of the weighted vector is different in different preset selection modes.
[0098] Optionally, the KV calculation cache module 510 is further configured to cache all keys and values into the CPU memory to obtain a KV Cache; the selection layer calculation module 530 is further configured to load the KV Cache in the CPU memory into the GPU memory for use in the current layer calculation; the subsequent layer calculation module 540 is further configured to obtain the KV Cache subset corresponding to the target token in the KV Cache from the CPU memory and load the KV Cache subset into the GPU memory for use in the current layer calculation.
[0099] Optionally, in the process of loading the KV Cache subset into the GPU memory, when the complete KV Cache subset is not obtained, the subsequent layer calculation module 540 is further configured to perform calculations based on the already loaded part of the KV Cache subset as the attention context of the current layer; when the complete KV Cache subset is obtained, the KV Cache subset is used as the attention context of the current layer and calculations are performed to obtain the output token.
[0100] In the embodiments of this specification, a large model inference acceleration device is provided. Among them, the KV calculation cache module is configured to obtain the input text data of the pre-trained large model, control the pre-trained large model to calculate the keys and values of all pre-filled tokens in the input text data based on multiple attention layers, and cache all keys and values to obtain a KV Cache;
[0101] The multi-attention layer division module is configured to divide at least two consecutive selection propagation layer sets from the multiple attention layers. The selection propagation layer set includes a selection layer and the subsequent layer corresponding to the selection layer, and the selection layer is the first layer in the selection propagation layer set;
[0102] The selection layer calculation module is configured to, when performing calculations in each selection layer, obtain the KV Cache, use the KV Cache as the attention context of the current layer and perform calculations to obtain the output token, and determine the target tokens among all pre-filled tokens that meet the importance condition;
[0103] Subsequent layer calculation module, which is used to obtain the corresponding KV Cache subset of the target token in the KV Cache when performing calculations in the subsequent layers corresponding to each selected layer, use the KV Cache subset as the attention context of the current layer, and perform calculations to obtain the output token. In the embodiments of this specification, multiple attention layers in the large model are divided into several selected propagation layer sets. The first layer in each set is the selected layer, and the remaining layers are the subsequent layers; in each selected layer, the KV of all tokens is used for calculation, and at the same time, important tokens in the pre-filled tokens are selected; in the corresponding subsequent layers, only the KV subset corresponding to these important tokens needs to be used for calculation, and there is no need to use all the pre-filled tokens. Since there is sparse pattern similarity between adjacent layers, the important tokens are also similar for these layers. Therefore, when important tokens are selected in the previous selected layer, the subsequent multiple layers can directly use the selected important tokens, thus reducing the token usage overhead of the model while ensuring that the performance of the model is not affected. In addition, for the large model, important tokens are selected multiple times based on multiple selected layers, which is beneficial for the large model to adapt to the importance changes of each token generated during the iteration process.
[0104] The embodiments of this specification provide a computer program product including instructions, which, when the computer program product runs on a computer or a processor, cause the computer or the processor to execute the steps of the method in any one of the above embodiments.
[0105] The embodiments of this specification also provide a computer storage medium, which can store multiple instructions, and the instructions are suitable for being loaded and executed by a processor to execute the steps of the method in any one of the above embodiments.
[0106] Please refer to Figure 6 , Figure 6 which is a schematic structural diagram of a terminal provided by the embodiments of this specification. As Figure 6 shown, the terminal 600 may include: at least one terminal processor 601, at least one network interface 604, a user interface 603, a memory 605, and at least one communication bus 602.
[0107] Among them, the communication bus 602 is used to realize the connection and communication between these components.
[0108] Among them, the user interface 603 may include a display screen (Display), a camera (Camera), and optionally the user interface 603 may further include a standard wired interface and a wireless interface.
[0109] Among them, the network interface 604 may optionally include a standard wired interface, a wireless interface (such as a WI-FI interface).
[0110] Among them, the terminal processor 601 may include one or more processing cores. The terminal processor 601 connects various parts within the entire terminal 600 through various interfaces and lines. By running or executing instructions, programs, code sets, or instruction sets stored in the memory 605, and by invoking data stored in the memory 605, it executes various functions of the terminal 600 and processes data. Optionally, the terminal processor 601 may be implemented in at least one of the following hardware forms: Digital Signal Processing (DSP), Field-Programmable Gate Array (FPGA), and Programmable Logic Array (PLA). The terminal processor 601 may integrate one or a combination of several of the following: Central Processing Unit (CPU), Graphics Processing Unit (GPU), and modem. Among them, the CPU mainly processes the operating system, user interface, and application programs, etc.; the GPU is responsible for rendering and drawing the content to be displayed on the display screen; the modem is used to process wireless communication. It can be understood that the above-mentioned modem may not be integrated into the terminal processor 601 and may be implemented separately through a single chip.
[0111] Among them, the memory 605 may include Random Access Memory (RAM), and may also include Read-Only Memory (ROM). Optionally, the memory 605 includes a non-transitory computer-readable storage medium. The memory 605 can be used to store instructions, programs, code, code sets, or instruction sets. The memory 605 may include a program storage area and a data storage area. Among them, the program storage area may store instructions for implementing the operating system, instructions for at least one function (such as touch function, sound playback function, image playback function, etc.), instructions for implementing the above-mentioned various method embodiments, etc.; the data storage area may store the data involved in the above-mentioned various method embodiments. Optionally, the memory 605 may also be at least one storage device located far from the aforementioned terminal processor 601. As Figure 6 shown, the memory 605, as a computer storage medium, may include an operating system, a network communication module, a user interface module, and a large model inference acceleration program.
[0112] In Figure 6 In the terminal 600 shown in Figure 6 , the user interface 603 is mainly used to provide an interface for the user to input data and obtain the data input by the user; while the terminal processor 601 can be used to call the large model inference acceleration program stored in the memory 605 and specifically perform the following operations:
[0113] Obtain the input text data of the pre-trained large model, control the pre-trained large model to calculate the keys and values of all pre-filled tokens in the input text data based on multiple attention layers, and cache all the keys and values to obtain the KV Cache;
[0114] Divide at least two consecutive sets of selection propagation layers in multiple attention layers. The set of selection propagation layers includes a selection layer and the subsequent layer corresponding to the selection layer, and the selection layer is the first layer in the set of selection propagation layers;
[0115] When calculating in each selection layer, obtain the KV Cache, use the KV Cache as the attention context of the current layer and perform calculations to obtain the output token, and determine the target tokens among all pre-filled tokens that meet the importance condition;
[0116] When calculating in the subsequent layer corresponding to each selection layer, obtain the KV Cache subset corresponding to the target token in the KV Cache, use the KV Cache subset as the attention context of the current layer and perform calculations to obtain the output token.
[0117] In some embodiments, when the terminal processor 601 executes dividing at least two consecutive sets of selection propagation layers in multiple attention layers, it specifically performs the following steps: determine the first M layers in the multiple attention layers as the look-ahead layers, and divide at least two consecutive sets of selection propagation layers in the remaining attention layers, where M is a positive integer; when calculating in each look-ahead layer, obtain the KV Cache, use the KV Cache as the attention context of the current layer and perform calculations to obtain the output token.
[0118] In some embodiments, when the terminal processor 601 executes using the KV Cache as the attention context of the current layer and performing calculations to obtain the output token, it specifically performs the following steps: receive the output token of the previous layer as the current token, calculate the key and value of the current token; perform attention calculations based on the KV Cache and the key and value of the current token to obtain the output token; update the key and value of the current token to the KV Cache.
[0119] In some embodiments, when the terminal processor 601 executes to determine the target tokens that meet the importance condition among all pre-filled tokens, the following steps are specifically executed: calculate the attention scores of each pre-filled token, and determine the target tokens that meet the importance condition among all pre-filled tokens according to the attention scores of each pre-filled token.
[0120] In some embodiments, when the terminal processor 601 executes to calculate the attention scores of each pre-filled token and determine the target tokens that meet the importance condition among all pre-filled tokens according to the attention scores of each pre-filled token, the following steps are specifically executed: construct a context hidden state window with a preset length, and calculate the attention scores of each pre-filled token based on the context hidden state window; perform weighted summation on the attention scores of the pre-filled tokens within each context hidden state window based on the weighted vector corresponding to the preset selection mode to obtain the context scores of each pre-filled token; select the target tokens that meet the importance condition according to the context scores of each pre-filled token.
[0121] In some embodiments, the preset selection mode is at least one of uniform selection, exponential selection, and last token selection, and the calculation methods of the weighted vectors are different in different preset selection modes.
[0122] In some embodiments, when the terminal processor 601 executes to cache all keys and values to obtain the KV Cache, the following steps are specifically executed: cache all keys and values into the CPU memory to obtain the KV Cache; when the terminal processor 601 executes to obtain the KV Cache, the following steps are specifically executed: load the KV Cache in the CPU memory into the GPU memory for use in the current layer calculation; when the terminal processor 601 executes to obtain the KV Cache subset corresponding to the target token in the KV Cache, the following steps are specifically executed: obtain the KV Cache subset corresponding to the target token in the KV Cache from the CPU memory, and load the KV Cache subset into the GPU memory for use in the current layer calculation.
[0123] In some embodiments, when the terminal processor 601 executes the steps of obtaining the KV Cache subset corresponding to the target token in the KV Cache, using the KV Cache subset as the attention context of the current layer, and performing calculations to obtain the output token, the following steps are specifically executed: During the process of loading the KV Cache subset into the GPU memory, when the complete KV Cache subset has not been obtained, calculations are performed based on the already loaded part of the KV Cache subset as the attention context of the current layer; when the complete KV Cache subset has been obtained, the KV Cache subset is used as the attention context of the current layer and calculations are performed to obtain the output token.
[0124] In several embodiments provided in this specification, it should be understood that the disclosed devices and methods can be implemented in other ways. For example, the device embodiments described above are merely illustrative. For example, the division of modules is only a logical function division. In actual implementation, there may be other division methods. For example, multiple modules or components can be combined or integrated into another system, or some features can be ignored or not executed. Another point is that the displayed or discussed coupling or direct coupling or communication connection between each other can be through some interfaces. The indirect coupling or communication connection of devices or modules can be in electrical, mechanical or other forms.
[0125] The modules described as separate components may or may not be physically separated. The components shown as modules may or may not be physical modules, that is, they may be located in one place, or they may be distributed to multiple network modules. Some or all of the modules can be selected according to actual needs to achieve the purpose of the solution of this embodiment.
[0126] In the above embodiments, it can be implemented in whole or in part by software, hardware, firmware, or any combination thereof. When implemented using software, it can be implemented in whole or in part in the form of a computer program product. The above computer program product includes one or more computer instructions. When the above computer program instructions are loaded and executed on a computer, the processes or functions described above in accordance with the embodiments of this specification are generated in whole or in part. The above computer may be a general-purpose computer, a special-purpose computer, a computer network, or other programmable devices. The above computer instructions may be stored in a computer-readable storage medium or transmitted through the above computer-readable storage medium. The above computer instructions may be transmitted from one website, computer, server, or data center to another website, computer, server, or data center in a wired manner (such as coaxial cable, optical fiber, Digital Subscriber Line (DSL)) or wirelessly (such as infrared, wireless, microwave, etc.). The above computer-readable storage medium may be any available medium that a computer can access or a data storage device such as a server or data center that includes one or more integrated available media. The above available medium may be a magnetic medium (for example, a floppy disk, a hard disk, a magnetic tape), an optical medium (for example, a Digital Versatile Disc (DVD)), or a semiconductor medium (for example, a Solid State Disk (SSD)), etc.
[0127] It should be noted that for the foregoing method embodiments, for the sake of simplicity of description, they are all expressed as a series of action combinations. However, those skilled in the art should know that the embodiments of this specification are not limited by the described order of actions, because according to the embodiments of this specification, certain steps may be performed in other orders or simultaneously. Secondly, those skilled in the art should also know that the embodiments described in the specification are all preferred embodiments, and the actions and modules involved are not necessarily essential for the embodiments of this specification.
[0128] In addition, it should also be noted that the information (including but not limited to user equipment information, user personal information, etc.), data (including but not limited to data for analysis, stored data, displayed data, etc.), and signals involved in the embodiments of this specification are all authorized by the user or fully authorized by all parties, and the collection, use, and processing of relevant data need to comply with the relevant laws, regulations, and standards of relevant countries and regions.
[0129] The above describes specific embodiments of this specification. Other embodiments are within the scope of the appended claims. In some cases, the actions or steps recited in the claims may be performed in a different order than in the embodiments and still achieve the desired result. Additionally, the processes depicted in the drawings do not necessarily require the particular order or sequential order shown to achieve the desired result. In certain embodiments, multitasking and parallel processing are also possible or may be advantageous.
[0130] In the above embodiments, the descriptions of the respective embodiments have their own emphases. For parts not detailed in a certain embodiment, reference may be made to the relevant descriptions of other embodiments.
[0131] The above is the description of a large model inference acceleration method, device, storage medium, and terminal provided by the embodiments of this specification. For those skilled in the art, based on the ideas of the embodiments of this specification, there will be changes in the specific implementation manners and application scopes. In summary, the content of this specification should not be construed as a limitation on the embodiments of this specification.
Claims
1. A large model reasoning acceleration method, the method comprising: Obtain input text data of a pre-trained large model, control the pre-trained large model to calculate keys and values of all pre-filled tokens in the input text data based on multiple attention layers, and cache all keys and values to obtain a KV Cache; Dividing at least two consecutive selection propagation layer sets from the multiple attention layers, the selection propagation layer set includes a selection layer and a subsequent layer corresponding to the selection layer, and the selection layer is the first layer in the selection propagation layer set; When performing calculations in each selection layer, the KV Cache is obtained, the KV Cache is used as the attention context of the current layer and calculations are performed to obtain an output token, and a target token that meets the importance condition is determined among all pre-filled tokens; When calculations are performed in subsequent layers corresponding to each selection layer, the KV Cache subset corresponding to the target token in the KV Cache is obtained, the KV Cache subset is used as the attention context of the current layer and calculations are performed to obtain the output token.
2. The method according to claim 1, wherein dividing at least two consecutive sets of selective propagation layers from the multiple attention layers comprises: Determine the first M layers of the multiple attention layers as forward-looking layers, and divide at least two consecutive sets of selective propagation layers from the remaining attention layers, where M is a positive integer; When performing calculations in each forward-looking layer, the KV Cache is obtained, the KV Cache is used as the attention context of the current layer and calculations are performed to obtain an output token.
3. According to the method of claim 1, the step of using the KV Cache as the attention context of the current layer and performing calculations to obtain an output token comprises: Receive the output token of the previous layer as the current token, and calculate the key and value of the current token; Perform attention calculation based on the key and value of the KV Cache and the current token to obtain an output token; Update the key and value of the current token into the KV Cache.
4. According to the method of claim 1, the step of determining the target token that satisfies the importance condition among all pre-filled tokens comprises: Calculate the attention score of each pre-filled token, and determine the target token that meets the importance condition among all pre-filled tokens according to the attention score of each pre-filled token.
5. According to the method of claim 4, the calculating of the attention score of each pre-filled token and determining the target token that meets the importance condition among all pre-filled tokens according to the attention score of each pre-filled token comprises: Construct a context hidden state window of a preset length, and calculate the attention score of each pre-filled token based on the context hidden state window; Based on the weighted vector corresponding to the preset selection mode, the attention scores of the pre-filled tokens in each context hidden state window are weighted summed to obtain the context score of each pre-filled token; The target token that meets the importance condition is selected based on the context score of each pre-filled token.
6. According to the method of claim 1, the preset selection mode is at least one of uniform selection, exponential selection, and last token selection, and the calculation method of the weighted vector is different in different preset selection modes.
7. According to the method of claim 1, the caching of all keys and values to obtain KV Cache comprises: Cache all keys and values into CPU memory to get KV Cache; The obtaining of the KV Cache includes: Load the KV Cache in the CPU memory into the GPU memory for use in current layer calculation; The obtaining of the KV Cache subset corresponding to the target token in the KV Cache includes: The KV Cache subset corresponding to the target token in the KV Cache is obtained from the CPU memory, and the KV Cache subset is loaded into the GPU memory for use in current layer calculation.
8. According to the method of claim 7, obtaining the KVCache subset corresponding to the target token in the KV Cache, using the KV Cache subset as the attention context of the current layer and performing calculation to obtain the output token, comprises: In the process of loading the KV Cache subset into the GPU memory, when the complete KV Cache subset is not obtained, performing calculation based on the loaded portion of the KV Cache subset as the attention context of the current layer; When the complete KV Cache subset is obtained, the KV Cache subset is used as the attention context of the current layer and calculation is performed to obtain the output token.
9. A large model reasoning acceleration device, the device comprising: A KV calculation cache module is used to obtain input text data of a pre-trained large model, control the pre-trained large model to calculate the keys and values of all pre-filled tokens in the input text data based on multiple attention layers, and cache all keys and values to obtain a KV Cache; A multi-attention layer division module, used to divide at least two consecutive selection propagation layer sets from the multiple attention layers, wherein the selection propagation layer set includes a selection layer and a subsequent layer corresponding to the selection layer, and the selection layer is the first layer in the selection propagation layer set; A selection layer calculation module, used for obtaining the KV Cache when performing calculations in each selection layer, using the KV Cache as the attention context of the current layer and performing calculations to obtain an output token, and determining a target token that meets an importance condition among all pre-filled tokens; The subsequent layer calculation module is used to obtain the KV Cache subset corresponding to the target token in the KV Cache when performing calculations in the subsequent layers corresponding to each selection layer, use the KV Cache subset as the attention context of the current layer and perform calculations to obtain the output token.
10. A computer program product comprising instructions, which, when executed on a computer or a processor, enables the computer or the processor to execute the steps of the method according to any one of claims 1 to 8.
11. A computer storage medium storing a plurality of instructions, wherein the instructions are suitable for being loaded by a processor and executing the steps of the method according to any one of claims 1 to 8.
12. A terminal comprising a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein the processor implements the steps of the method according to any one of claims 1 to 8 when executing the computer program.
Citation Information
Cited By
Large model reasoning method and device, related equipment and computer program product
CN120430418A