Large model reasoning length and performance optimization method

By dynamically filtering and calculating the aggregated key matrix, combined with CPU offloading and GPU reloading techniques, the inference length and performance of large models are optimized, solving the problems of large KV cache memory consumption and high computational load, and achieving computational efficiency and storage optimization.

CN121882287AActive Publication Date: 2026-04-17KYLIN CORP
View PDF 4 Cites 0 Cited by

Patent Information

Authority / Receiving Office
CN · China
Patent Type
Applications(China)
Current Assignee / Owner
KYLIN CORP
Filing Date
2026-03-20
Publication Date
2026-04-17

AI Technical Summary

Technical Problem

Existing large models have high inference costs, especially due to the large GPU memory space occupied by the KV cache and the failure to effectively reduce the computational cost of the attention distribution matrix, which leads to increased computational and storage pressure.

Method used

By presetting that each sequence uses a maximum of N1 logical memory blocks, and storing a maximum of N2 and N3 logical memory blocks in the GPU and CPU respectively, the key matrix is ​​dynamically filtered and aggregated, attention scores are calculated and sorted, only logical memory blocks with high scores are retained for computation, and the KV cache is managed through CPU offloading and GPU reloading techniques.

Benefits of technology

It reduces computational load and GPU storage pressure while maintaining model prediction accuracy. The complexity of calculating the attention distribution is reduced from O(T) to O(MN2+BN1), and the KV cache size is reduced, thus reducing GPU memory usage.

✦ Generated by Eureka AI based on patent content.

Smart Images

  • Figure CN121882287A_ABST
    Figure CN121882287A_ABST
Patent Text Reader

Abstract

The invention provides a large model reasoning length and performance optimization method, which belongs to the technical field of computers, and comprises the following steps: presetting that each sequence uses N1 logic memory blocks at most, N2 logic memory blocks are saved in a GPU at most, and N3 logic memory blocks are saved in a CPU at most; aiming at key caches corresponding to the B tokens in each logic memory block, sequentially dividing the key caches into M non-overlapped blocks according to a sequence, and respectively calculating aggregation characteristics to obtain an aggregation key matrix; if T / Blt is in the current calculation process; = N1, and a conventional attention calculation mode is used; if T / Bgt is in the current calculation process; n1, after obtaining a query vector, splicing the aggregation key matrix corresponding to each logic memory block to obtain the query vector; calculating attention scores on the basis of the query vector and the aggregation key matrix, and sorting the attention scores to obtain N1 logic memory blocks with higher attention scores; and performing conventional attention mechanism calculation. According to the invention, the calculation amount and the GPU storage pressure are reduced.
Need to check novelty before this filing date? Find Prior Art

Description

Technical Field

[0001] This invention relates to large model optimization, and more particularly to a method for optimizing the inference length and performance of large models, belonging to the field of computer technology. Background Technology

[0002] With the rapid development of large model technology, the number of model parameters and the maximum supported inference length of large models have increased rapidly, leading to a corresponding increase in model inference costs. To accelerate inference speed and reduce the computational load required for inference, mainstream inference frameworks, such as VLLM, generally use a key-value cache mechanism. This mechanism saves the historical information of the key and value matrices already computed in the attention mechanism, reusing them in subsequent decoding processes to reduce redundant computation. The KV cache is updated during each decoding process. The overall process is as follows: Figure 1 As shown.

[0003] vLLM uses a paged attention mechanism to divide the key-value cache and store it in multiple logical key-value blocks. Each logical block has a fixed size; in vLLM, the default size is 16, which can hold the key / value values ​​of 16 tokens. This allows for dynamic allocation of key-value cache memory for requests, improving memory utilization. Assuming the sequence length (number of tokens) is T in the current computation, reusing the key-value cache reduces the computational cost of calculating the key and value matrices in the attention mechanism from O(T) to O(T). 2) becomes O(2).

[0004] While the KV cache reduces inference time by preventing redundant calculations of the key and value matrices in the attention mechanism, it doesn't effectively reduce the computational cost of subsequently calculating the memory distribution matrix. Furthermore, storing the KV cache requires additional GPU memory. Assuming the current large model has L layers, D hidden layer dimensions, and T sequence length in the current computation, the computational cost of the memory distribution matrix is ​​O(T), and the KV cache requires an additional L units of storage space. D T 2. As the length of the processed feature sequence increases, the space occupied by the KV cache increases rapidly.

[0005] Previous studies have found that when calculating the attention distribution matrix, attention often focuses on specific tokens. Therefore, sparse attention mechanisms are often used to actively filter and retain the tokens that need attention while discarding irrelevant tokens, achieving KV cache sparsification, shortening the sequence length of the stored KV cache, and reducing storage pressure and computational cost. "Longformer: The Long-Document Transformer" uses a sliding window attention mechanism to force the query vector to focus only on multiple historical tokens at specific positions directly adjacent to the current token. However, this method filters and retains historical tokens in a static way, failing to fully consider the differences and correlations between tokens, thus resulting in poor performance. The paper "H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models" proposes a key-value (KV) cache maintenance / eviction strategy. By setting a maximum number of cached tokens, whenever the number of cached tokens exceeds the maximum, a scoring function is used to calculate the score of each token, and the token with the lowest score is discarded. This achieves relatively good results. However, the method of discarding tokens relies on the average of all historical attention scores. Since different decoding stages focus on different tokens, the average attention score cannot fully represent the importance of the token. In addition, the method of discarding tokens is discrete and irreversible. Once a token is discarded, it cannot be reused. Furthermore, the sparse discarding of tokens at specific positions cannot be fully integrated with the Paged Attention technique of VLLM. Summary of the Invention

[0006] To address the aforementioned issues, this invention provides a method for optimizing the inference length and performance of large models, thereby reducing computational load and GPU storage pressure.

[0007] To achieve the above objectives, the technical solution of the present invention is: a method for optimizing the inference length and performance of large models, comprising the following steps: S1, presupposes that each sequence uses a maximum of N1 logical memory blocks, stores a maximum of N2 logical memory blocks in the GPU, and stores a maximum of N3 logical memory blocks in the CPU; where N1 <= N2; S2, cache the keys corresponding to the B tokens in each logical memory block. The data is sequentially divided into M non-overlapping blocks of size E, and the aggregation features are calculated for each block to obtain the aggregation bond matrix. ; S3, if the currently calculated and saved logical memory blocks N≤N1, use the conventional attention calculation method; if the currently calculated and saved logical memory blocks N>N1, then obtain the query vector. Then, the aggregated key matrix corresponding to each logical memory block is... Piecing them together, we get Based on query vector and aggregate bond matrix Calculate and sort the attention scores to obtain N1 logical memory blocks with higher attention scores. Perform conventional attention mechanism calculations on the N1 logical memory blocks with higher attention scores. S4, the number of times the logical memory blocks containing the N1 logical memory blocks with higher attention scores are activated. Increase the attribute by 1.

[0008] Furthermore, each logical memory block in step S1 includes two additional attributes: (1) the number of times the logical memory block is activated. , indicates the number of times the current logical memory block has been called, with an initial value of 0; (2) the number of times the logical memory block has not been activated. This represents the number of times the current logical memory block has not been invoked since it was unloaded onto the CPU; its initial value is 0. In step S2, the aggregation feature is calculated by averaging along the time dimension or by applying max pooling along the time dimension. In step S3, based on the query vector and aggregate bond matrix Calculate and sort the attention scores to obtain N1 logical memory blocks with higher attention scores, including the following calculation process:

[0009]

[0010]

[0011]

[0012]

[0013]

[0014]

[0015]

[0016] In the formula, Let d be the feature vector corresponding to the T-th token in the autoregressive calculation process, where d is the feature dimension. and The weights and biases used by the linear layer in the self-attention layer to transform the input into a query vector. for The corresponding query vector; and The weights and biases used by the linear layer in the self-attention layer to transform the input into a key vector. for The corresponding key vector; and The weights and biases used by the linear layer in the self-attention layer to transform the input into a value vector. for The corresponding value vector; and This represents the key vector and value vector corresponding to the first T-1 tokens; and This is the result of concatenating the KV cache and the newly computed key / value vectors, where T is the total sequence length; @ represents the matrix multiplication operation. This is the square root of the feature dimension, used to scale the attention score to reduce variance. Rate attention Used for attention scoring The attention distribution matrix obtained after normalization, The result of the attention mechanism calculation; Among them, attention score It is a scalar representing the key vector. and query vector The correlation between them; The aggregate bond matrix corresponding to logical memory block i The j-th vector The calculation method is as follows

[0017] Based on query vector and aggregate bond matrix The Middle vectors The calculation of attention score is corresponding key vector The mean of attention scores, i.e.

[0018]

[0019] Then query vector With the aggregate bond matrix Attention score It is a scalar, that is, the query vector. The maximum value of the product with each vector in the aggregation key matrix, that is Among

[0020] By sorting in descending order and finding the N1 logical memory blocks with the largest scores to participate in the calculation.

[0021] Furthermore, in step S3, when N > N1, the following situations are also included: If N1 < N ≤ N2, only select N1 logical memory blocks with higher attention scores to participate in the conventional attention mechanism calculation; If N2 < N ≤ N2 + N3, on the basis of selecting N1 logical memory blocks to participate in the conventional attention mechanism calculation, sequentially transfer the logical memory block with the lowest attention score from the GPU to the CPU until only N2 logical memory blocks remain on the GPU; If N2 + N3 < N, on the basis of selecting N1 logical memory blocks to participate in the conventional attention mechanism calculation and sequentially transferring the logical memory block with the lowest attention score from the GPU to the CPU until only N2 logical memory blocks remain on the GPU, according to the logical memory blocks saved on the CPU, based on Sorting is performed, and the highest logical memory blocks saved on the CPU are sequentially deleted until only N3 logical memory blocks remain on the CPU.

[0022] Furthermore, in step S4, if the corresponding logical memory block among the N1 logical memory blocks is on the CPU, it is re - loaded from the CPU to the GPU.

[0023] Furthermore, in step S3, for the cases of N2 < N ≤ N2 + N3 and N2 + N3 < N, only retain N2 logical memory blocks with higher attention scores on the GPU, unload the remaining logical memory blocks to the CPU, and increment the number of times the logical memory block of all logical memory blocks located on the CPU is not activated by 1.

[0024] Furthermore, in step S3, for the case of N2 + N3 < N, only retain the N3 logical memory blocks with the lowest number of times the logical memory block is not activated on the CPU, and delete the remaining logical memory blocks on the CPU.

[0025] The beneficial effects of a method for optimizing the inference length and performance of a large model according to the present invention are: ​This invention extracts the aggregation key matrix from the logical memory blocks in Paged Attention and dynamically filters the logical memory blocks used in the current stage to prevent the omission of effective information, thus reducing the computational cost of calculating the attention distribution from O(T) to O(M). N2+B The computational complexity of calculating the attention aggregation matrix changes from O(T) to O(B) for N1). N1).

[0026] This invention uses a CPU cache to temporarily store logical memory blocks, and employs CPU offload / GPU reload technology to prevent the premature loss of important information. The size of the KV cache is determined by... Become This reduces the storage pressure on the GPU. Attached Figure Description

[0027] The present invention will now be described in further detail with reference to the accompanying drawings and specific embodiments.

[0028] Figure 1 This is a schematic diagram of the attention mechanism applied by KV cache; Figure 2 This is a schematic diagram of logical memory block i and its corresponding aggregate key vector; Figure 3 This is a flowchart of the two-stage attention mechanism calculation. Figure 4 This is a flowchart of the logical memory block unloading / reloading process. Detailed Implementation

[0029] The technical solutions of the embodiments of the present invention will be clearly and completely described below with reference to the accompanying drawings. Many specific details are set forth in the following description to provide a thorough understanding of the present invention. However, the present invention may also be implemented in other ways different from those described herein. Those skilled in the art can make similar extensions without departing from the spirit of the present invention. Therefore, the present invention is not limited to the specific embodiments disclosed below.

[0030] The present invention provides a method for optimizing the inference length and performance of large models without training. It can effectively combine the Paged Attention technology of VLLM to dynamically filter and utilize important tokens, and reuse deleted tokens through CPU-offload / GPU reload. This method can reduce the amount of KV cache memory and the computational cost of the attention mechanism without compromising the model's prediction accuracy.

[0031] Because text information has semantic continuity, previous studies have shown that the attention distributions of adjacent tokens are similar. Furthermore, considering that the Paged Attention technique in VLLM stores adjacent tokens in the same logical memory block, the overall idea of ​​this invention includes the following: 1. It is assumed that each sequence uses a maximum of N1 logical memory blocks, and a maximum of N2 logical memory blocks are stored in the GPU and a maximum of N3 logical memory blocks are stored in the CPU, where N1≤N2.

[0032] 2. For each logical memory block containing B tokens, calculate M aggregate key vectors and store them on the GPU. After each decoding, count the number of times the logical memory block is accessed and the number of times it is not accessed. Figure 2 As shown.

[0033] It should be noted that in the Page attention mechanism, multiple logical memory blocks are created. Each logical memory block contains key vectors and value vectors corresponding to B tokens. Generally, the value of B is 16 or 32. Figure 2 middle This represents the i-th logical memory block. This represents the j-th key vector stored in the i-th logical memory block. This represents the j-th value vector stored in the i-th logical memory block. Taking B=16 as an example, at this point, the logical memory block... It contains B tokens, and then adjacent key vectors are aggregated to obtain an aggregated key vector. For example, if M is 2, and the aggregation method is to take the average, then... , Based on aggregated key vectors , Obtain the aggregated bond matrix for each logic block. .

[0034] 3. A two-stage approach is used to filter and calculate the attention distribution matrix. First, the query vector is used... For aggregate bond matrix Calculate the attention score, select N1 logical memory blocks for this stage, and then apply the full attention mechanism to the selected N1 logical memory blocks to calculate the output. This changes the computational cost of calculating the attention distribution from O(T) to O(M). N2+B N1), calculate the attention aggregation matrix through the attention distribution. The computational complexity changes from O(T) to O(B). N1), where B N1< <M N2+B N1<< T, its calling flow is as follows: Figure 3 As shown.

[0035] 4. A CPU cache is used to temporarily store KV cache logical memory blocks that have not been accessed for a long time. Only N² logical memory blocks are stored on the GPU, and a maximum of N³ logical memory blocks are stored on the CPU. When the number of memory blocks on the GPU exceeds N², the lower-rated memory blocks are unloaded to the CPU, but their corresponding aggregate key matrices are preserved. It is still stored on the GPU and participates in the selection of logical memory blocks, thereby reducing the storage pressure on the GPU. In actual use, the context length T is relatively large, T / B >> N2, which causes the size of the KV cache to be reduced from... Become ( (This is a rounding-up function). When a logical memory block on the CPU corresponds to an aggregated key matrix with a high attention score, it will be reloaded onto the GPU (GPU reload) to prevent important information from being lost prematurely. Its call flow is as follows: Figure 4 As shown.

[0036] It should be noted that the GPU's video memory on a server is limited. Following the logic of current technology, assuming there are T tokens, the memory required will be... There are 3 logical memory blocks, each containing 1 B key vectors and 1 value vectors corresponding to each attention layer, with each key vector and 1 value vector having a dimension of D. Therefore, each logical memory block occupies 1 L of GPU memory. D B 2, therefore, a total of The GPU memory space. After applying the logic of this invention, each B tokens will generate an aggregated key vector of length M, thus occupying additional GPU memory space. The video memory space. And when When the number of logical memory blocks exceeds N2, the extra logical memory blocks will be offloaded to the CPU, and only N2 logical memory blocks will be stored on the GPU. Therefore, the GPU memory occupied by the logical memory blocks will be reduced by... Change to L D N2 B 2. Therefore, the size of the KV cache stored on the GPU is determined by... Become .

[0037] Furthermore, "when a logical memory block on the CPU corresponds to a high attention score in the aggregate bond matrix, it will be reloaded onto the GPU (GPU reload)," which means "when a memory block on the CPU has an attention score that ranks in the top N1, it will be loaded onto the GPU." "High attention score" actually refers to "attention score ranking in the top N1."

[0038] Example 1 The present invention provides a method for optimizing the inference length and performance of large models, which specifically includes the following steps: S1. Preset that each sequence uses a maximum of N1 logical memory blocks, a maximum of N2 logical memory blocks are stored in the GPU, and a maximum of N3 logical memory blocks are stored in the CPU. Each logical memory block will be assigned two additional attributes: (1) the number of times the logical memory block is activated. , indicating the number of times the current logical memory block has been accessed, with an initial value of 0. (2) Number of times the logical memory block has not been activated. This represents the number of times the current logical memory block has not been accessed since it was unloaded from the CPU; its initial value is 0. and Conjugate; N1 <= N2, but the value of N3 is unconstrained and unrestricted; N1 is used to control the amount of computation, N2 is used to control the GPU memory usage, and N3 is used to control the CPU usage.

[0039] S2. Cache the keys corresponding to the B tokens in each logical memory block i. The data is sequentially divided into M non-overlapping blocks of size E, and aggregate features are calculated for each block (aggregate features can be calculated by averaging along the time dimension or by applying max pooling along the time dimension), thus obtaining aggregated key vectors. , And obtain the aggregated key matrix of each logic block based on the aggregated key vector. , where D refers to the dimension of the hidden layer.

[0040] S3. If T / B <= N1 during the current calculation process, meaning the number of saved logical memory blocks N is less than N1, use the conventional attention calculation method. If T / B > N1 during the current calculation process, meaning the number of saved logical memory blocks N is greater than N1, obtain the query vector. Then, the aggregated key matrix corresponding to each logical memory block is... Piecing them together, we get .

[0041] It should be noted that this embodiment assumes... , there are the following situations: (1) If N <= N1, only the conventional attention mechanism needs to be used for calculation at this time; (2) If N1 < N ≤ N2, it means that the logical memory blocks participating in the calculation are redundant. Based on the query vector and the aggregation key matrix calculate the attention scores and sort them, and only select N1 logical memory blocks to participate in the conventional attention mechanism calculation; (3) If N2 < N ≤ N2 + N3, it means that the logical memory blocks saved on the GPU have reached the upper limit. At this time, the extra logical memory blocks need to be unloaded to the CPU. Therefore, based on the query vector and the aggregation key matrix calculate the attention scores and sort them, and only select N1 logical memory blocks to participate in the conventional attention mechanism calculation (if the corresponding logical memory block is on the CPU, it will be reloaded from the CPU to the GPU). Then, transfer the logical memory block with the lowest attention score from the GPU to the CPU in turn until only N2 logical memory blocks remain on the GPU.

[0042] (4) If N2 + N3 < N, it means that the logical memory blocks saved on the GPU and the CPU have reached the upper limit. At this time, the extra logical memory blocks on the GPU need to be unloaded to the CPU, and then the extra logical memory blocks on the CPU will be directly deleted. Therefore, based on the query vector and the aggregation key matrix calculate the attention scores and sort them, and only select N1 logical memory blocks to participate in the conventional attention mechanism calculation (if the corresponding logical memory block is on the CPU, it will be reloaded from the CPU to the GPU). Then, transfer the logical memory block with the lowest attention score from the GPU to the CPU in turn until only N2 logical memory blocks remain on the GPU. After that, sort according to the logical memory blocks saved on the CPU based on and delete the logical memory block with the highest on the CPU in turn (the corresponding aggregation key matrix will also be deleted) until only N3 logical memory blocks remain on the CPU.

[0043] In addition, it should also be noted that the memory blocks must be created on the GPU and will be temporarily transferred to the CPU in specific scenarios. And the logical memory blocks are unique and will not exist simultaneously on the GPU and the CPU. And the logical memory blocks can only participate in the "conventional attention mechanism calculation" on the GPU, which is also the reason for "if the corresponding logical memory block is on the CPU, it will be reloaded from the CPU to the GPU". If the logical memory block does not exist on the CPU, it must exist on the GPU and can directly participate in the calculation without transfer.

[0044] The above is based on query vectors and aggregate bond matrix The calculation and sorting of attention scores to obtain the N1 logical memory blocks with higher attention scores specifically includes the following calculation process: For the input of attention layer i in the autoregressive process The corresponding calculation results, and the most basic attention mechanism calculation method are as follows:

[0045]

[0046]

[0047]

[0048]

[0049]

[0050]

[0051]

[0052] In the formula, Let d be the feature vector corresponding to the T-th token in the autoregressive calculation process, where d is the feature dimension. and The weights and biases used by the linear layer in the self-attention layer to transform the input into a query vector. for The corresponding query vector; and The weights and biases used by the linear layer in the self-attention layer to transform the input into a key vector. for The corresponding key vector; and The weights and biases used by the linear layer in the self-attention layer to transform the input into a value vector. for The corresponding value vector; and This represents the key vector and value vector corresponding to the first T-1 tokens, also known as the KV cache. This part of the calculation results will be retained and reused until the decoding is finished. and This is the result of concatenating the KV cache and the newly computed key / value vector, where T is the total sequence length, or the number of tokens. @ represents a matrix multiplication operation. This is the square root of the feature dimension, used to scale the attention score to reduce variance. Rate attention Used for attention scoring The attention distribution matrix obtained after normalization, This represents the computational result of the attention mechanism.

[0053] Among them, attention score It is a scalar representing the key vector. and query vector The higher the value, the stronger the correlation between the two.

[0054] The aggregate bond matrix corresponding to logical memory block i The j-th vector The calculation method is as follows

[0055] Based on query vector and aggregate bond matrix The Middle vectors The same calculation method was used to calculate the attention score, which can also be understood as... corresponding key vector The mean of attention scores, i.e.

[0056]

[0057] Then query vector With the aggregate bond matrix Attention score It is a scalar, that is, the query vector. With the aggregate bond matrix The maximum value of the product of each vector in the vector, i.e.

[0058] The larger this value, the stronger the input for this operation. The stronger the correlation between the information and the historical information contained in the corresponding logical memory block.

[0059] Therefore, through [ Sort the data from largest to smallest and select the N1 logical memory blocks with the highest scores to participate in the calculation. This way, the computational complexity can be reduced as much as possible while retaining the most relevant historical information.

[0060] In addition, it should be noted that the query vector and aggregate bond matrix Attention score The larger the value, the greater the input value. The stronger the correlation between the information and the historical information contained in the corresponding logical memory block.

[0061] If the attention score of a logical memory block ranks in the top N1 after a single sort, it can be considered that the logical memory block is more relevant to the current inference task, and therefore its activation count should be increased. Increment the attribute by 1. This attribute can be understood as the number of times the logical memory block is used, thus representing its relevance to the current inference task. The larger the value, the stronger the relevance of the logical memory block to the current inference task, and the higher the probability of it being used again. Therefore, it should be kept on the GPU as much as possible to prevent the IO burden caused by the logical memory block being unloaded to the CPU and then read back to the GPU. Similarly, if the value is smaller, it indicates that the logical memory block is less relevant to the current inference task, and the probability of it being used again is also lower. Therefore, under preset conditions, it can be unloaded to the CPU and loaded onto the GPU only when needed.

[0062] S4, the number of times the logical memory blocks containing the N1 logical memory blocks with higher attention scores are activated. Increase the attribute by 1.

[0063] The above steps S1-S4 are repeated.

[0064] In summary, unlike existing schemes that filter independent tokens through dynamic / static sparsity, this invention fully considers the semantic relevance of adjacent tokens and the role mechanism of Paged Attention in VLLM, and proposes a method for optimizing the inference length and performance of large models. The key technical points are: (1) By extracting the aggregation key matrix from the logical memory blocks in Paged Attention, the logical memory blocks used in the current stage are dynamically filtered to prevent the omission of effective information. (2) The logical memory blocks are temporarily stored by CPU caching, and important information is prevented from being lost prematurely by CPU offloading / GPU reloading technology.

[0065] Obviously, the described embodiments are only some, not all, of the embodiments of the present invention. All other embodiments obtained by those skilled in the art based on the embodiments of the present invention without inventive effort are within the scope of protection of the present invention.

Claims

1. A large model inference length and performance optimization method, characterized in that, It includes the following steps: S1. Preset that each sequence uses at most N1 logical memory blocks, at most N2 logical memory blocks are saved in the GPU, and at most N3 logical memory blocks are saved in the CPU; where N1 <= N2; S2, cache the keys corresponding to the B tokens in each logical memory block , sequentially divide into M non-overlapping blocks of size E and calculate the aggregated features respectively, and obtain the aggregated key matrix ; S3, if the currently calculated and saved logical memory blocks N≤N1, use the conventional attention calculation method; if the currently calculated and saved logical memory blocks N>N1, then obtain the query vector. Then, the aggregated key matrix corresponding to each logical memory block is... Piecing them together, we get Based on query vector and aggregate bond matrix Calculate and sort the attention scores to obtain N1 logical memory blocks with higher attention scores. Perform conventional attention mechanism calculations on the N1 logical memory blocks with higher attention scores. S4, the number of times of activating the logical memory block containing N1 logical memory blocks with higher attention score attribute plus 1; Each logical memory block in step S1 includes two additional attributes: (1) the number of times the logical memory block is activated , which is initially 0; and (2) the number of times the logical memory block is not activated , which is initially 0. Based on the query vector in step S3 and the aggregation bond matrix Calculate the attention score and sort to get N1 higher attention score logical memory blocks, including the following calculation process: , , , , , , , , In the formula, Let d be the feature vector corresponding to the T-th token in the autoregressive calculation process, where d is the feature dimension. and The weights and biases used by the linear layer in the self-attention layer to transform the input into a query vector. for The corresponding query vector; and The weights and biases used by the linear layer in the self-attention layer to transform the input into a key vector. for The corresponding key vector; and The weights and biases used by the linear layer in the self-attention layer to transform the input into a value vector. for The corresponding value vector; and This represents the key vector and value vector corresponding to the first T-1 tokens; and This is the result of concatenating the KV cache and the newly computed key / value vectors, where T is the total sequence length; @ represents the matrix multiplication operation. This is the square root of the feature dimension, used to scale the attention score to reduce variance. Rate attention Used for attention scoring The attention distribution matrix obtained after normalization, The result of the attention mechanism calculation; Among them, attention score It is a scalar representing the key vector. and query vector The correlation between them; The aggregate bond matrix corresponding to logical memory block i The j-th vector The calculation method is as follows , Based on query vector and aggregate bond matrix The Middle vectors The calculation of attention score is corresponding key vector The mean of attention scores, i.e. , , Then query vector With the aggregate bond matrix Attention score It is a scalar, that is, the query vector. With the aggregate bond matrix The maximum value of the product of each vector in the vector, i.e. , Through [ Sort the data from largest to smallest and find the N1 logical memory blocks with the highest scores to participate in the calculation.

2. The method according to claim 1, characterized in that, In step S2, the calculation of the aggregated feature uses the average along the time dimension or applies max pooling along the time dimension.

3. The method according to claim 2, characterized in that, In step S3, when N > N1, the following situations are further included: If N1 < N ≤ N2, only select N1 logical memory blocks with higher attention scores to participate in the calculation of the conventional attention mechanism; If N2 < N ≤ N2 + N3, on the basis of selecting N1 logical memory blocks to participate in the calculation of the conventional attention mechanism, sequentially transfer the logical memory block with the lowest attention score from the GPU to the CPU until only N2 logical memory blocks remain on the GPU; If N2 + N3 < N, when selecting N1 logical memory blocks to participate in the calculation of the conventional attention mechanism, sequentially transfer the logical memory blocks with the lowest attention scores from the GPU to the CPU until only N2 logical memory blocks remain on the GPU. Based on the logical memory blocks saved on the CPU, sort them according to and sequentially delete the logical memory blocks with the highest scores on the CPU until only N3 logical memory blocks remain on the CPU.

4. The method according to claim 3, characterized in that, In step S4, if the corresponding logical memory block among the N1 logical memory blocks is located on the CPU, reload it from the CPU to the GPU.

5. The method according to claim 4, characterized in that, In step S3, for the cases where N2 < N ≤ N2 + N3 and N2 + N3 < N, only N2 logical memory blocks with higher attention scores are retained on the GPU, and the remaining logical memory blocks are unloaded to the CPU, and the number of times the logical memory blocks of all the logical memory blocks located on the CPU are not activated The attribute is incremented by 1.

6. The method according to claim 4, characterized in that, In step S3, for the case where N2 + N3 < N, only the number of times the logical memory blocks are not activated is retained on the CPU. The lowest N3 logical memory blocks are retained, and the remaining logical memory blocks on the CPU are deleted.

Citation Information

Patent Citations

  • Method and system for improving reasoning length and performance of large model and medium

    CN120197699A

  • Key value cache compression and sparse attention calculation method and system for large language model reasoning

    CN120952055A

  • Large language model key value cache optimization method based on cooperative game

    CN121143950A

  • Algorithm method for large-model long-context reasoning

    CN121365738A