A sliding window attention acceleration method and device based on circular cache
By adopting a sliding window attention acceleration method based on ring cache in the large language model, the problem of low memory storage efficiency in the prior art is solved, and more efficient memory utilization and inference speed are achieved.
Patent Information
- Application Number
- CN202411709463.5
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- Filing Date
- 2024-11-27
- Publication Date
- 2025-02-07
- Estimated Expiration
- 2044-11-27
AI Technical Summary
The prior art has low memory storage efficiency in sliding window attention mechanism, resulting in slow processing speed of large models for sequences.
Using a sliding window attention acceleration method based on ring cache, the cache position index of the K matrix and V matrix is cyclically cached and reused by assigning two ring cache blocks K Rcache and V Rcache of length N to each decoding layer.
It improves the memory utilization efficiency of K Rcache and V Rcache in SWA, improves the speed of K Rcache and V Rcache data storage, thereby accelerating the inference process of large language models.
Smart Images

Figure CN119204230B_ABST
Abstract
Description
Technical Field
[0001] The present invention relates to the field of computer technology, and in particular to a sliding window attention acceleration method and device based on circular cache. Background Art
[0002] With the development of artificial intelligence technology, especially the advancement of deep learning technology, the attention mechanism has gradually become a core component in natural language processing, speech recognition and other technical fields due to its superior performance in processing long dependencies. Traditional neural network models have the problem of gradient vanishing or gradient exploding when processing long sequences. At the same time, as the length of the sequence increases, the complexity of the model calculation is linearly related to the square of the sequence length, which makes the processing of large models on long sequence data very time-consuming.
[0003] In recent years, the Transformer architecture has significantly improved training efficiency by introducing the sliding window self-attention mechanism (SWA), which allows the model to process data of fixed window size in parallel.
[0004] The calculation process of SWA is as follows:
[0005]
[0006] in, is a two-dimensional matrix, N1 is the text length, H is the hidden layer length; d k is the length of the second dimension of the matrix H, K T is the transpose of matrix K. Multiply by a mask before softmax processing ( ), SA's Mask is a lower triangular matrix, so that the token at each text position only pays attention to the text content at that position and before. SWA's Mask is different from SA. There is a window size, so that Attention only focuses on the text content within the window size.
[0007] In order to reduce repeated calculations of the model during the large language model (LLM) reasoning process, KV Cache technology is generally used to cache the K and V matrices calculated in Attention, avoiding repeated calculations of historical KV matrices when new query reasoning calculations are requested, reducing the amount of model reasoning calculations and improving the model reasoning speed.
[0008] However, some existing KV Cache technologies have problems of low memory utilization and low storage efficiency for SWA calculations. Summary of the invention
[0009] The present invention provides a sliding window attention acceleration method and device based on circular cache to solve the problem that the memory storage efficiency in the sliding window attention mechanism in the prior art is low, which leads to slow sequence processing speed for large models.
[0010] In a first aspect, the present invention provides a sliding window attention acceleration method based on circular cache, which specifically comprises the following steps:
[0011] Step S1, obtaining parameters of a large language model (LLM);
[0012] Step S2: According to the parameters of the LLM, two ring buffer cache blocks (Ring Buffer Cache Block, RBCB) KRcache and VRcache with a length of N are respectively allocated to each decoding layer (Decoder Layer) in the LLM before LLM inference calculation;
[0013] Step S3: In the LLM input data preparation stage, according to the input data token0, token1, ..., token k and corresponding subscripts (in this application, the "corresponding subscripts" means that for token i , whose corresponding subscript is i), assigns the ring cache location index of K Rcache and V Rcache to each token;
[0014] Step S4: According to the ring cache location index of all tokens, the K matrix and V matrix of each token in the LLM reasoning process are stored in K Rcache and V Rcache respectively (in this application, the "corresponding" means that for a token i , the ring buffer location index assigned to it by step S3 is index i , whose K matrix is stored in the index of K Rcache i The V matrix is stored in the index of V Rcache. i locations) of the ring buffer.
[0015] Where N represents the size of the Sliding Window Attention (SWA) window, and N is a positive integer.
[0016] Among them, Rcache stands for Ring Buffer Cache.
[0017] Wherein, i and k are both positive integers, i is greater than or equal to 1 and i is less than or equal to k.
[0018] Preferably, in step S2, the annular cache block represents a cache block composed of N blocks and used to store K matrix and V matrix cache.
[0019] Preferably, in step S2, each K matrix and V matrix of each decoding layer is allocated a ring cache block, namely [KC0, KC1, ..., KC N-1 ] and [VC0,VC1,...,VC N-1 ].
[0020] In this application, because LLM generation is a cyclic reasoning process, input data preparation or processing is required before each reasoning, so it is described as the LLM input data preparation stage. In the LLM cyclic reasoning process, the first reasoning is called the prefill stage, and the second reasoning until the end of the reasoning is called the decode stage.
[0021] Preferably, in step S3, in the LLM input data preparation stage, according to the input data token0, token1, ..., token k And the corresponding subscript, assign K Rcache and V Rcache index positions to each token, specifically including the following steps:
[0022] Step S301: During the first inference of LLM (or the data preparation phase before the prefill phase of the first inference), k tokens are input. When j is less than or equal to N-1, token i The ring buffer position index is assigned to j, and the ring buffer position index is input into the index queue Q j ; When j is greater than N-1, the K matrix and V matrix of the last N tokens are cached, and the ring cache position index is assigned to [0,1,...,N-1];
[0023] Step S302: In the LLM autoregressive reasoning process (or can be expressed as the data preparation stage before each reasoning in the decoding stage), a token is obtained each time according to the subscript from small to large, and the last index n of the index queue Q at this time is obtained, and the ring buffer position index index of the current token in the ring buffer block is reallocated according to the index n and the ring buffer position index index is added to the index queue Q;
[0024] Step S303, repeating step S302 until the LLM autoregression generation is completed;
[0025] Wherein, j is a positive integer, j is greater than or equal to 1 and j is less than or equal to k.
[0026] Preferably, in step S302, the length of the index queue Q is the same as the length of K Rcache and V Rcache, and the specific calculation of the ring cache location index index is as follows:
[0027] index=(n+1)%N;
[0028] Where n is a positive integer, % represents modulo, and index is a positive integer. As can be seen from the above formula, the maximum index stored in the index queue Q is N-1, which corresponds to the last position of the circular cache. As n increases, the circular cache position index index will cycle from 0 to N-1, that is, the K matrix and V matrix of the newly generated token are stored from the starting position of the circular cache block to the last position of the cache block, and then stored from the starting position, forming a circular reuse of the cache space.
[0029] Preferably, a sliding window attention acceleration method based on circular cache further includes:
[0030] Step S5: In the decoding calculation phase, SWA obtains the K matrix and V matrix stored in K Rcache and V Rcache according to the ring cache position index of all tokens, and calculates through the K matrix and V matrix to form the attention calculation result of LLM.
[0031] Preferably, in step S5, during each Attention calculation process, the indexes of the K matrix and the V matrix correspond one to one (for example, , , );
[0032] From the above calculation process, we can see that the subscripts of the K matrix and the V matrix are consistent each time they are calculated. Attn3(Q,K,V) is the sum of Q3 and each pair of K matrix and V matrix. The subscript has nothing to do with the order of calculation of each pair of matrix K and matrix V. Therefore, in the Attention calculation process, there is no need to calculate according to the order of input tokens, but only need to calculate from the starting position of the ring buffer block.
[0033] Among them, in the present application, the index queue Q is used to store the circular cache position index; the circular cache position index represents the position number of the K matrix stored in the K Rcache and the position number of the V matrix stored in the V Rcache; the last index of the index queue Q represents the last value of the index queue Q.
[0034] In a second aspect, the present invention further provides a sliding window attention acceleration device based on circular cache, which specifically includes the following modules:
[0035] Large language model parameter acquisition module, used to obtain the parameters of the large language model (LLM);
[0036] A cache allocation module, configured to allocate two ring cache blocks K Rcache and V Rcache of length N to each decoding layer in the LLM before LLM inference calculation according to the parameters of the LLM;
[0037] The cache index allocation module is used to allocate the indexes of token0, token1, ..., token according to the input data in the LLM input data preparation stage. k And the corresponding subscript, assign the ring cache location index of K Rcache and V Rcache to each token;
[0038] The cache storage module is used to store the K matrix and V matrix of each token in the LLM reasoning process in the corresponding ring cache positions of K Rcache and V Rcache respectively according to the ring cache position indexes of all tokens.
[0039] Where N represents the size of the Sliding Window Attention (SWA) window, and N is a positive integer.
[0040] Among them, Rcache stands for Ring Buffer Cache.
[0041] Wherein, i and k are both positive integers, i is greater than or equal to 1 and i is less than or equal to k.
[0042] Preferably, in the cache allocation module, the annular cache block represents a cache block composed of N blocks, which is used to store K matrix and V matrix caches.
[0043] Preferably, in the cache allocation module, each K matrix and V matrix of each decoding layer is allocated a ring cache block, namely [KC0, KC1, ..., KC N-1 ] and [VC0,VC1,...,VC N-1 ].
[0044] Wherein, N is a positive integer.
[0045] Preferably, the cache index allocation module specifically includes the following submodules:
[0046] The first cache index allocation submodule is used to input k tokens during the first LLM reasoning, and when j is less than or equal to N-1, the token iThe ring buffer position index is assigned to j, and the ring buffer position index is input into the index queue Q j ; When j is greater than N-1, the K matrix and V matrix of the last N tokens are cached, and the ring cache position index is assigned to [0,1,...,N-1];
[0047] The second cache index allocation submodule is used to obtain a token each time according to the subscript from small to large during the LLM autoregressive inference process, obtain the last index n of the index queue Q at this time, reallocate the ring cache position index index of the current token in the ring cache block according to the index n, and add the ring cache position index index to the index queue Q;
[0048] A third cache index allocation submodule, used for repeatedly executing the second cache index allocation submodule until the LLM autoregression generation is completed;
[0049] Wherein, j is a positive integer, j is greater than or equal to 1 and j is less than or equal to k.
[0050] Preferably, in the second cache index allocation submodule, the length of the index queue Q is the same as the lengths of K Rcache and V Rcache, and the specific calculation of the ring cache location index index is as follows:
[0051] index=(n+1)%N;
[0052] Where n is a positive integer, % represents modulo, and index is a positive integer. As can be seen from the above formula, the maximum index stored in the index queue Q is N-1, which corresponds to the last position of the circular cache. As n increases, the circular cache position index index will cycle from 0 to N-1, that is, the K matrix and V matrix of the newly generated token are stored from the starting position of the circular cache block to the last position of the cache block, and then stored from the starting position, forming a circular reuse of the cache space.
[0053] Preferably, a sliding window attention acceleration device based on circular cache further includes:
[0054] The attention calculation module is used in the decoding calculation stage of SWA. According to the ring cache position index of all tokens, it obtains the K matrix and V matrix stored in K Rcache and V Rcache, and calculates through the K matrix and V matrix to form the attention calculation result of LLM.
[0055] Preferably, in the attention calculation module, in each Attention calculation process, the indexes of the K matrix and the V matrix correspond one to one.
[0056] In a third aspect, the present invention further provides a computer-readable storage medium having a computer program stored thereon, which, when executed by a processor, implements a sliding window attention acceleration method based on a circular cache as described in any one of the first aspects of the present application.
[0057] In a fourth aspect, the present invention further provides an electronic device, comprising: a memory storing a computer program; and a processor communicatively connected to the memory, for executing a sliding window attention acceleration method based on circular cache as described in any one of the first aspects of the present application when the computer program is called.
[0058] Compared with the prior art, the present invention has the following obvious outstanding substantive features and significant advantages:
[0059] The present invention provides a sliding window attention acceleration method and device based on circular cache, which solves the problem that the memory storage efficiency in the sliding window attention mechanism in the prior art is low, which leads to the slow sequence processing speed of the large model. Based on the circular cache mechanism, the large language model with sliding window attention is optimized, which improves the memory utilization efficiency of KRcache and VRcache in SWA, and also improves the data storage speed of KRcache and VRcache, thereby improving the reasoning speed of the entire large language model. BRIEF DESCRIPTION OF THE DRAWINGS
[0060] The accompanying drawings, which constitute part of the present invention, are used to provide a further understanding of the present invention. The exemplary embodiments of the present invention and their descriptions are used to explain the present invention and do not constitute improper limitations on the present invention. In the accompanying drawings:
[0061] Figure 1 (a) is an Attention Mask map of SA.
[0062] Figure 1 (b) is an Attention Mask image of SWA.
[0063] Figure 2 It is a K / V Cache cache graph.
[0064] Figure 3 This is a SWA K / V Cache caching process diagram.
[0065] Figure 4 It is a flow chart of a sliding window attention acceleration method based on circular cache according to a preferred embodiment of the present invention.
[0066] Figure 5 (a) is a schematic diagram of a Ring Buffer cache block according to a preferred embodiment of the present invention.
[0067] Figure 5 (b) is a schematic diagram of the storage direction of a Ring Buffer cache block in a preferred embodiment of the present invention.
[0068] Figure 6 It is a structural schematic diagram of a sliding window attention acceleration device based on circular cache according to a preferred embodiment of the present invention. DETAILED DESCRIPTION
[0069] The present invention provides a method and device for accelerating attention of a sliding window based on a circular cache. In order to make the purpose, technical solution and effect of the present invention clearer and more specific, the present invention is further described in detail with reference to the accompanying drawings and examples. It should be understood that the specific embodiments described herein are only used to explain the present invention and are not used to limit the present invention.
[0070] It should be noted that the terms "first", "second", etc. in the specification and claims of the present invention and the above-mentioned drawings are used to distinguish similar objects, and are not necessarily used to describe a specific order or sequence. It should be understood that the data used in this way can be interchanged under appropriate circumstances. In addition, the terms "including" and "having" and any variations thereof are intended to cover non-exclusive inclusions. For example, a process, method, system, product or device that includes a series of steps or units is not necessarily limited to those steps or units clearly listed, but may include other steps or units that are not clearly listed or inherent to these processes, methods, products or devices.
[0071] Sliding Window Attention (SWA) is a variant of the Self-Attention (SA) mechanism in Transformer. The calculation process of SWA and SA is as follows:
[0072]
[0073] in, is a two-dimensional matrix, N1 is the text length, H is the hidden layer length, d k is the length of the second dimension of the matrix H, K T is the transpose of matrix K. Multiply by a mask before softmax , the Mask of SA is a lower triangular matrix, such as Figure 1 Part (a) makes the token at each text position only pay attention to the text content at that position and before. The Mask of SWA is different from that of SA, such as Figure 1In part (b), there is a window size, so that Attention only focuses on the text content within the window size. In the large language model (LLM) reasoning process, in order to reduce the repeated calculation of the model, KV Cache technology is generally used to cache the K matrix and V matrix calculated in Attention, avoiding the repetition of historical K matrix and V matrix calculation when a new query reasoning calculation is requested, reducing the amount of model reasoning calculation and improving the model reasoning speed.
[0074] Existing K / V Cache (K Cache or V Cache) technologies are generally as follows: Figure 2 As shown in the figure, with the input of new query, K / V Cache will add a new cache, and the new cache will generally be spliced at the end of the original caches block. This method will take a long time to splice the new cache with the original caches block each time, involving new data memory allocation and data copying.
[0075] There is also a Page Cache mechanism designed based on the Page Cache principle, which will pre-cache multiple pages of Caches Block and continuously add new Cache data to it. These two methods are mainly suitable for the calculation of the SA method, which will continuously add new caches and gradually increase the memory usage. They are not suitable for the K / VCache caching method with a fixed window size in the SWA scenario, and the cache utilization efficiency is not high.
[0076] Common K / V Cache technologies in SWA scenarios include: Figure 3 As shown in the figure, as a new K / V Cache is added, the data of a cache is removed from the original Caches Block data header, and a new cache data is spliced at the end. The length of the entire Caches Block is always maintained at the size of the window length. This method also involves new data memory allocation and data copying during the caching process, which is time-consuming.
[0077] Embodiment 1:
[0078] like Figure 1-Figure 5 As shown, the sliding window attention acceleration method based on circular cache described in this embodiment specifically includes the following steps:
[0079] Step S1, obtaining parameters of a large language model (LLM);
[0080] Step S2: According to the parameters of the LLM, two ring cache blocks K Rcache and V Rcache of length N are respectively allocated to each decoding layer in the LLM before LLM inference calculation; wherein N represents the size of a sliding window attention (SWA) window, and N is a positive integer.
[0081] The ring buffer block is composed of N blocks and is used to store the K matrix and V matrix cache. The Ring Buffer Cache mechanism will allocate a window size video memory or memory space according to each layer of SWA before the large model inference, which is called a ring buffer cache block. Figure 5 As shown in (a), new K / V Caches are continuously generated during the SWA calculation process. K / V Caches are stored in the cache in order from the starting address of the cache block. When the loop is completed, the data is stored again from the starting address to overwrite the original data, and is always stored in a fixed memory space.
[0082] Since the memory address is gradually increasing and cannot be cyclic, it is actually implemented as a continuous memory address, such as Figure 5 As shown in (b). Assuming that the window size of SWA is N, the memory address can be divided into N blocks, which can store N K / V Caches with storage indexes of [0,1,2,…,N-1]. When the stored K / V Cache index exceeds N-1, it is re-stored from address 0. Among them, each K matrix and V matrix of each decoding layer is allocated a ring cache block, which are [KC0,KC1,…,KC N-1 ] and [VC0,VC1,...,VC N-1 ].
[0083] Step S3: In the LLM input data preparation stage, according to the input data token0, token1, ..., token k And the corresponding subscript, assign K Rcache and V Rcache ring cache location index to each token;
[0084] Optionally, step S3 specifically includes the following steps:
[0085] Step S301: During the first LLM reasoning, k tokens are input. When j (where j is a positive integer, j is greater than or equal to 1 and j is less than or equal to k) is less than or equal to N-1, the token i The ring buffer position index is assigned to j, and the ring buffer position index is input into the index queue Q j; When j is greater than N-1, the K matrix and V matrix of the last N tokens are cached, and the ring cache position index is assigned to [0,1,...,N-1].
[0086] Step S302: In the LLM autoregressive reasoning process, a token is obtained each time according to the subscript from small to large, and the last index n of the index queue Q at this time is obtained. The ring cache position index index of the current token in the ring cache block is reallocated according to the index n, and the ring cache position index index is added to the index queue Q; wherein the length of the index queue Q is the same as the length of K Rcache and V Rcache, and the specific calculation of the ring cache position index index is as follows:
[0087] index=(n+1)%N;
[0088] Where n is a positive integer, % represents modulo, and index is a positive integer. As can be seen from the above formula, the maximum index stored in the index queue Q is N-1, which corresponds to the last position of the circular cache. As n increases, the circular cache position index index will cycle from 0 to N-1, that is, the K matrix and V matrix of the newly generated token are stored from the starting position of the circular cache block to the last position of the cache block, and then stored from the starting position, forming a circular reuse of the cache space.
[0089] For example, a sequence includes token0, token1, token2, token3, and N=3; the first time token0 is obtained, the last index n=0 of the index queue Q, the index index=1 of token0 is calculated, and the ring buffer position index=1 is stored in the first position of the index queue Q; the second time token1 is obtained, the last index n=1 of the index queue Q, the shape buffer position index=2 is calculated, and the shape buffer position index=2 is stored in the second position of the index queue Q; the third time token2 is obtained, the last index n=2 of the index queue Q, the shape buffer position index=0 is calculated, and the shape buffer position index=0 is stored in the third position of the index queue Q; the fourth time token3 is obtained, the last index n=3 of the index queue Q, the shape buffer position index=1 is calculated, and the shape buffer position index=1 is stored in the fourth position of the index queue Q. The size of the shape buffer position index index is the position index number of the K matrix and V matrix cache stored in the ring buffer block. Step S303: Repeat step S302 until the LLM autoregression generation is completed.
[0090] Preferably, a sliding window attention acceleration method based on circular cache further includes:
[0091] Step S5: In the decoding calculation phase, SWA obtains the K matrix and V matrix stored in K Rcache and V Rcache according to the ring cache position index of all tokens, and calculates through the K matrix and V matrix to form the attention calculation result of LLM.
[0092] Preferably, in step S5, in each Attention calculation process, the indexes of the K matrix and the V matrix correspond one to one. Therefore, when calculating Attention, it is not necessary to calculate in the order of the K matrix and the V matrix index queues, but it is sufficient to calculate from the starting position of the K Rcache and the V Rcache, and it is not necessary to change the original calculation method of Attention or FlashAttention to calculate in the order of the K matrix and the V matrix index.
[0093] Embodiment 2:
[0094] like Figure 6 As shown, the sliding window attention acceleration device based on circular cache described in this embodiment specifically includes the following modules:
[0095] The large language model parameter acquisition module is used to obtain the parameters of the large language model (LLM).
[0096] A cache allocation module is used to allocate two ring cache blocks K Rcache and V Rcache of length N to each decoding layer in the LLM before LLM inference calculation according to the parameters of the LLM; wherein N represents the size of the sliding window attention (SWA) window, and N is a positive integer.
[0097] In the cache allocation module, each K matrix and V matrix of each decoding layer is allocated a ring cache block, namely [KC0, KC1, ..., KC N-1 ] and [VC0,VC1,...,VC N-1 ].
[0098] The cache index allocation module is used to allocate the indexes of token0, token1, ..., token according to the input data in the LLM input data preparation stage. k (where k is a positive integer) and the corresponding subscript, assigning each token the ring cache location index of K Rcache and VRcache.
[0099] The cache index allocation module specifically includes the following submodules:
[0100] The first cache index allocation submodule is used to input k tokens during the first LLM reasoning, and when j (where j is a positive integer, j is greater than or equal to 1 and j is less than or equal to k) is less than or equal to N-1, the token i The ring buffer position index is assigned to j, and the ring buffer index position is input into the index queue Q j ; When j is greater than N-1, the K matrix and V matrix of the last N tokens are cached, and the ring cache position index is assigned to [0,1,...,N-1].
[0101] The second cache index allocation submodule is used to obtain a token each time according to the subscript from small to large during the LLM autoregressive inference process, obtain the last index n of the index queue Q at this time, reallocate the ring cache position index index of the current token in the ring cache block according to the index n, and add the ring cache position index index to the index queue Q.
[0102] The length of the index queue Q is the same as that of K Rcache and V Rcache. The specific calculation of the ring cache location index index is as follows:
[0103] index=(n+1)%N;
[0104] Where n is a positive integer, % represents modulo, and index is a positive integer. As can be seen from the above formula, the maximum index stored in the index queue Q is N-1, which corresponds to the last position of the circular cache. As n increases, the circular cache position index index will cycle from 0 to N-1, that is, the K matrix and V matrix of the newly generated token are stored from the starting position of the circular cache block to the last position of the cache block, and then stored from the starting position, forming a circular reuse of the cache space.
[0105] The third cache index allocation submodule is used to repeatedly execute the second cache index allocation submodule until the LLM autoregression generation is completed.
[0106] The cache storage module is used to store the K matrix and V matrix of each token in the LLM reasoning process in the corresponding ring cache positions of K Rcache and V Rcache respectively according to the ring cache position indexes of all tokens.
[0107] Optionally, a sliding window attention acceleration device based on circular cache further includes:
[0108] The attention calculation module is used in the decoding calculation stage of SWA. According to the ring cache position index of all tokens, the K matrix and V matrix stored in K Rcache and V Rcache are obtained, and the K matrix and V matrix are used for calculation to form the attention calculation result of LLM. In each Attention calculation process, the indexes of K matrix and V matrix correspond one to one.
[0109] The present invention designs a Ring Buffer Cache mechanism based on the circular buffer concept for SWA reasoning acceleration in large model reasoning scenarios. Ring Buffer can cyclically utilize a fixed memory space, which coincides with the window cache concept of SWA. A window-sized memory is pre-allocated through the RingBuffer Cache mechanism to cache SWA's KVCache cyclically, avoiding repeated memory allocation and time consumed by large amounts of data copying when caching a new cache, while also avoiding the growth of cache memory during the reasoning process, improving memory usage efficiency, and thus improving the reasoning speed of large models when using SWA calculations.
[0110] The specific embodiments of the present invention are described in detail above, but they are only examples, and the present invention is not limited to the specific embodiments described above. For those skilled in the art, any equivalent modifications and substitutions made to the present invention are also within the scope of the present invention. Therefore, the equalization changes and modifications made without departing from the spirit and scope of the present invention should be included in the scope of the present invention.
Claims
1. A sliding window attention acceleration method based on circular cache, characterized in that: The specific steps include: Step S1, obtaining parameters of a large language model; Step S2: according to the parameters of the large language model, before LLM inference calculation, two ring cache blocks K Rcache and V Rcache of length N are respectively allocated to each decoding layer in the LLM; Step S3: In the LLM input data preparation stage, according to the input data token0, token1, ..., token k And the corresponding subscript, assign the ring cache location index of K Rcache and V Rcache to each token; Step S4: According to the ring cache index positions of all tokens, the K matrix and V matrix of each token in the LLM reasoning process are stored in the ring cache positions corresponding to K Rcache and V Rcache respectively; Among them, N represents the size of the sliding window attention window, and N is a positive integer; Among them, Rcache represents a ring cache; Wherein, k is a positive integer.
2. According to the sliding window attention acceleration method based on circular cache in claim 1, it is characterized in that: In step S2, the ring cache block represents a cache block composed of N blocks and used to store K matrix and V matrix cache.
3. According to the sliding window attention acceleration method based on circular cache in claim 1, it is characterized in that: In step S3, in the LLM input data preparation stage, according to the input data token0, token1, ..., token k And the corresponding subscript, assign the ring cache location index of K Rcache and V Rcache to each token, specifically including the following steps: Step S301: During the first LLM reasoning, input k tokens. When j is less than or equal to N-1, the token i The ring buffer position index is assigned to j, and the ring buffer position index is input into the index queue Q j ; When j is greater than N-1, the K matrix and V matrix of the last N tokens are cached, and the ring cache position index is assigned to [0,1,...,N-1]; Step S302: In the LLM autoregressive reasoning process, a token is obtained each time according to the subscript from small to large, and the last index n of the index queue Q at this time is obtained, and the ring buffer position index index of the current token in the ring buffer block is reallocated according to the index n and the ring buffer position index index is added to the index queue Q; Step S303, repeating step S302, and repeating step S302 until the LLM autoregression generation is completed; Wherein, j is a positive integer, j is greater than or equal to 1 and j is less than or equal to k.
4. A sliding window attention acceleration method based on circular cache according to claim 3, characterized in that: In step S302, the length of the index queue Q is the same as the length of K Rcache and V Rcache, and the specific calculation of the ring cache location index index is as follows: index=(n+1)%N; Where n is a positive integer, % means remainder, and index is a positive integer.
5. The method for accelerating the sliding window attention based on circular buffer according to claim 1, characterized in that: Also includes: Step S5: In the decoding calculation phase, SWA obtains the K matrix and V matrix stored in K Rcache and VRcache according to the ring cache position index of all tokens, and calculates through the K matrix and V matrix to form the attention calculation result of LLM.
6. A sliding window attention acceleration device based on circular cache, characterized in that: The modules include: A large language model parameter acquisition module, used to obtain parameters of the large language model; A cache allocation module, configured to allocate two ring cache blocks K Rcache and V Rcache of length N to each decoding layer in the LLM before LLM inference calculation according to the parameters of the LLM; The cache index allocation module is used to allocate the indexes of token0, token1, ..., token according to the input data in the LLM input data preparation stage. k And the corresponding subscript, assign the ring cache location index of K Rcache and V Rcache to each token; The cache storage module is used to store the K matrix and V matrix of each token in the LLM reasoning process in the corresponding ring cache positions of K Rcache and V Rcache respectively according to the ring cache position indexes of all tokens; Among them, N represents the size of the sliding window attention window, and N is a positive integer; Among them, Rcache represents a ring cache; Wherein, k is a positive integer.
7. A computer-readable storage medium, characterized in that: The computer-readable storage medium stores a computer program, and when the computer program is executed by a processor, it implements a sliding window attention acceleration method based on circular cache as described in any one of claims 1-5.
8. An electronic device, characterized in that: It includes a memory, a processor, and a computer program stored in the memory and executable on the processor. When the processor executes the computer program, it implements a sliding window attention acceleration method based on circular cache as described in any one of claims 1 to 5.
Citation Information
Patent Citations
Data caching method and device based on attention mechanism, equipment and medium
CN118113630A
Large language model reasoning acceleration method and device based on sparse sliding window
CN118132682A
Cited By
Fast attention calculation
GB2704436A