A cache method in an attention mechanism, device and medium

Abstract

The application provides a caching method in an attention mechanism, an apparatus and a medium, relates to the technical field of artificial intelligence, and aims to solve the technical problems of excessive storage occupation and large computing cost of KV caching of a large language model attention mechanism. The method comprises the following steps: for all word units in a vocabulary, a set of query Q, key K and value V vectors not containing position information is generated by pre-computation; the set of vectors is stored in a storage device, and a mapping relationship between the word units and the vectors is established; during model inference, for a target word unit in an input sequence, the corresponding Q, K and V vectors are retrieved and acquired from the storage device based on the mapping relationship; finally, the attention result is calculated based on the acquired vectors and the position information of the target word unit. Through pre-computation and global sharing, the application converts real-time computation into efficient retrieval, significantly reduces the computing cost and storage occupation, and improves the long text processing capability and concurrent service capability of the model.

Description

Technical Field

[0001] This application relates to the field of artificial intelligence technology, and in particular to an attention mechanism caching optimization method, apparatus and computer-readable storage medium for large language models. Background Technology

[0002] In recent years, large language models based on the Transformer architecture have achieved great success in the field of artificial intelligence, and their core attention mechanism is key to their powerful performance. During the model's inference process, to improve processing speed, it is typically necessary to cache a large number of key-value pairs generated during attention computation, i.e., a KV cache. However, this approach comes at the cost of expensive and scarce storage resources in computing devices, such as the high-bandwidth memory and cache in graphics processing units (GPUs). As the length of input sequences increases and the number of concurrently processed sequences rises, the KV cache quickly becomes a system bottleneck, severely limiting the model's ability to process long texts, its inference speed, and its service capacity.

[0003] To alleviate the aforementioned problems, existing technologies have proposed several optimization schemes. For example, some schemes attempt to pre-compute and cache the key-value (KV) states of reusable static text portions (or independent modules) in the input, and then directly query and assemble these caches during subsequent inference to reduce redundant computations. However, this method based on "text block" reuse has limited applicability, does not reduce the computational and storage overhead for non-repeating parts of the sequence, and requires complex corrections or shifts of the positional information contained within different cache blocks when assembling them, increasing processing complexity. Other optimization methods, such as window attention, sparse attention, or quantization compression of the cache, can reduce overhead to some extent, but they mainly optimize within the framework of existing KV caching, and the fundamental problems of storage and computational overhead remain unresolved. Therefore, how to fundamentally reduce the attention computation and storage overhead in the inference process of large language models is a technical problem that urgently needs to be solved by those skilled in the art. Summary of the Invention

[0004] The purpose of this application is to provide a global caching method, apparatus and computer-readable storage medium for attention mechanisms, which aims to solve the technical problems in the prior art where the KV caching of the attention mechanism of large language models will excessively occupy storage resources, real-time generation of KV vectors will have a large computational overhead, and additional storage of projection matrices will be required, thus limiting the performance of model inference.

[0005] To achieve the above objectives, this application provides a caching method in an attention mechanism, comprising: for all lexical units in the vocabulary, based on the projection matrix of at least one specified layer in the model, pre-calculating and generating a set of query Q, key K, and value V vectors that correspond one-to-one with all lexical units and do not contain positional information; storing the set of Q, K, and V vectors in at least one storage device and establishing a mapping relationship between all lexical units and the set of Q, K, and V vectors; during model inference, for a target lexical unit in the input sequence, retrieving the Q, K, and V vectors corresponding to the target lexical unit from the at least one storage device based on the mapping relationship; and finally, calculating an attention result based on the obtained Q, K, and V vectors and the positional information of the target lexical unit in the input sequence.

[0006] Optionally, storing the set of Q, K, V vectors in at least one storage device includes:

[0007] The set of Q, K, V vectors are stored hierarchically in a first storage device and a second storage device, wherein the access speed of the first storage device is higher than that of the second storage device;

[0008] Furthermore, the hierarchical storage includes: based on the preset usage frequency of each word in all the words, prioritizing the storage of Q, K, and V vectors corresponding to high-frequency words in the first storage device, and storing the Q, K, and V vectors corresponding to low-frequency words in the second storage device.

[0009] In a preferred embodiment of this application, the first storage device is a cache and the second storage device is high-bandwidth memory.

[0010] Furthermore, the method further includes: monitoring the available space of the first storage device; when the available space is less than a preset threshold, determining the amount of storage space to be released, and deleting or moving at least one Q, K, V vector corresponding to the low-frequency word from the first storage device based on the amount of storage space to be released.

[0011] Optionally, the number m of Q, K, and V vectors deleted or moved is based on the required storage space size n, calculated using the formula... It is determined, where λ is the proportionality coefficient.

[0012] Furthermore, during model inference, the projection matrix used for pre-calculation is not stored in the video memory or cache of the computing device.

[0013] In a preferred embodiment of this application, the designated layer is the first attention layer of the model.

[0014] This application also provides a caching device in an attention mechanism, comprising: a pre-computation unit, configured to pre-compute, based on the projection matrix of at least one specified layer in the model, a set of query Q, key K, and value V vectors corresponding one-to-one with all lexical units in the vocabulary and without positional information; a storage management unit, configured to store the set of Q, K, and V vectors in at least one storage device and establish a mapping relationship between all lexical units and the set of Q, K, and V vectors; a retrieval unit, configured to, during model inference, retrieve, based on the mapping relationship, the Q, K, and V vectors corresponding to the target lexical unit in the input sequence from the at least one storage device; and a calculation unit, configured to calculate the attention result based on the obtained Q, K, and V vectors and the positional information of the target lexical unit in the input sequence.

[0015] Furthermore, the set of Q, K, V vectors are stored hierarchically in a first storage device and a second storage device, wherein the access speed of the first storage device is higher than that of the second storage device; and, according to the preset usage frequency of each word in all the words, the Q, K, V vectors corresponding to high-frequency words are preferentially stored in the first storage device, and the Q, K, V vectors corresponding to low-frequency words are stored in the second storage device; the available space of the first storage device is monitored; when the available space is less than a preset threshold, the amount of storage space to be released is determined, and at least one Q, K, V vector corresponding to the low-frequency word is deleted or moved from the first storage device based on the amount of storage space to be released.

[0016] This application also provides a computer-readable storage medium having a computer program stored thereon, which, when executed by a processor, implements the method described in any of the preceding claims.

[0017] Compared with existing technologies, this application has the following beneficial effects: 1. Significantly reduced computational overhead: By pre-compiling and directly retrieving Q, K, and V vectors, this application avoids performing real-time matrix multiplication operations for specified layers during model inference, moving the computation to the offline stage, thereby significantly reducing the load on computing devices and improving inference speed. 2. Significantly reduced storage footprint: The pre-computed Q, K, and V vectors are globally shared, and all sequences reference the same set of vectors, completely changing the traditional method of generating and storing KV caches separately for each sequence. When processing long sequences or high-concurrency requests, this sharing mechanism can greatly reduce the occupation of storage resources. 3. Saved weight matrix storage: Since the Q, K, and V vectors for specified layers are directly obtained from storage, there is no need to retain their corresponding projection matrices in valuable storage resources such as GPU memory or cache during inference, further saving storage space. 4. Improved model processing capabilities: Combining the above advantages, this application effectively alleviates the storage bottleneck, enabling the model to process longer input sequences and support a higher number of concurrent users, thus improving the overall service capability of the model. Attached Figure Description

[0018] To more clearly illustrate the technical solutions of the embodiments of this application, the accompanying drawings used in the description of the embodiments will be briefly introduced below.

[0019] Figure 1 This is a flowchart illustrating a pre-calculation method provided in an embodiment of this application. Figure 2 This is a flowchart illustrating a reasoning method provided in an embodiment of this application. Figure 3 This is a schematic diagram of a hierarchical storage strategy provided in an embodiment of this application. Figure 4 This is a flowchart illustrating a dynamic cache management method provided in an embodiment of this application. Figure 5 This is a schematic diagram of the inference system architecture according to an embodiment of this application. Figure 6 This is a signaling interaction timing diagram of one embodiment of this application.

[0020] The main reference numerals in the attached figures are explained as follows: Step 10: Obtain the correspondence step; Step 20: Obtain the Q / K / V vector step; Step 30: Calculate the attention step by combining the position information P; 100: Inference system; 110: Inference engine; 120: Storage management unit; 130: Cache; 140: High-bandwidth memory; 150: Solid state drive. Detailed Implementation

[0021] To make the objectives, technical solutions, and advantages of this application clearer, the application will be further described in detail below with reference to the accompanying drawings and specific embodiments. It should be noted that the specific embodiments described herein are for illustrative purposes only and are not intended to limit the scope of this application.

[0022] Example 1

[0023] This embodiment provides a caching method in the attention mechanism, which optimizes the computation and storage overhead during the inference process of large language models through pre-computation and global sharing. As an optional implementation, this method can be executed within a computing system, such as a server configured with a graphics processor.

[0024] Please see Figure 5 This illustration shows a schematic diagram of an inference system architecture according to an embodiment of this application. An inference system 100 may include an inference engine 110, a storage management unit 120, and a tiered storage architecture. The storage architecture may include storage devices of different speeds and capacities, such as a cache 130, high-bandwidth memory 140, and a larger-capacity but relatively slower solid-state drive 150. The basic method flow described in this embodiment can be implemented on such a system architecture.

[0025] In one embodiment of this application, the method is mainly divided into two stages: an offline pre-computation stage and an online inference stage.

[0026] First, refer to Figure 1 The offline pre-computation stage is described below. The goal of this stage is to generate a globally shared set of query, key, and value vectors that do not contain positional information for all lexical units in the model vocabulary.

[0027] Specifically, in the offline pre-computation phase, the complete vocabulary of the large language model used is first obtained, along with the word embedding vector corresponding to each lexical unit in the vocabulary. Assuming the model vocabulary contains 100,000 unique lexical units, the system loads the word embedding vectors corresponding to these 100,000 lexical units, forming a large vector set, denoted as X. It can be understood that these word embedding vectors are the initial semantic representations of lexical units in the model.

[0028] Subsequently, the attention mechanism of at least one specified layer in the model is optimized. As a preferred implementation, the first attention layer of the model can be optimized because it processes the longest sequence and yields the most significant optimization benefits. The system loads the projection matrices used to generate the query, key, and value vectors from the attention mechanism of this specified layer (e.g., the first layer), denoted as follows: , and Subsequently, the system performs large-scale matrix multiplication operations to pre-calculate the corresponding query, key, and value vectors for all 100,000 lexical units in the vocabulary. This calculation process can be represented as: in, It is a collection containing all 100,000 word embedding vectors. , and This is the calculated set of global query, key, and value vectors, each corresponding one-to-one with all tokens. It's important to note that the query, key, and value vectors generated in this step are purely semantic representations of the tokens themselves, completely lacking any positional information (such as positional encoding) that will appear later in the sequence. This step can be performed by a pre-computation unit, which can be a dedicated computation module or an offline program.

[0029] To facilitate rapid lookup in subsequent inference stages, a mapping relationship needs to be established between lexical units and pre-computed vectors. The system establishes a one-to-one mapping between all lexical units in the vocabulary and the query, key, and value vectors generated in the above steps. This mapping relationship can be implemented using various data structures. For example, a hash table can be created, using the unique identifier of each lexical unit as the key and the address of its corresponding query, key, and value vector in the storage device as the value. Under a multi-head attention mechanism, this mapping relationship can also record the attention head information to which each vector belongs.

[0030] Finally, all generated query, key, and value vectors are stored in one or more designated storage devices according to the established mapping relationship. In this basic embodiment, these 100,000 sets of query, key, and value vectors can be stored as a contiguous data block in the high-speed storage of a computing device, such as the high-bandwidth memory 140 of a graphics processor. A storage management unit can be responsible for performing the mapping relationship establishment and vector storage, for example... Figure 5 The storage management unit 120 shown is shown.

[0031] Through the offline processing described above, the system prepares a global QKV vector "dictionary" that can be shared by all inference tasks. Its advantage lies in the fact that, after pre-computation and storage, the original projection matrix used to generate these vectors is readily available. , For this specific layer, it is no longer necessary during inference. Therefore, when the model is inferred online, there is no need to load these projection matrices into the valuable video memory or cache of the computing device, thus further saving storage space.

[0032] The following will refer to Figure 2 This section explains the online inference phase. This phase occurs when the model receives user input and needs to perform tasks such as text generation or understanding.

[0033] Suppose the model receives an input sequence, such as "The weather is nice today". The inference engine 110 processes the words in the sequence one by one. For any target word in the sequence, such as "today", the system no longer performs traditional real-time computation (i.e., obtaining the word embedding vector of "today" and then...). , Instead of matrix multiplication, the following retrieval and calculation process is used:

[0034] First, in step 10, the mapping relationship established in the offline phase is used to find the storage information of the query, key, and value vector corresponding to the target term "today". This is equivalent to performing a query in the hash table previously created using the term identifier "today".

[0035] Secondly, in step 20, based on the address information retrieved in step 10, the query, key, and value vector corresponding to the term "today," without location information, is directly retrieved from the specified storage device (e.g., high-bandwidth memory 140). This process is completed by a retrieval unit, which may be part of the storage management unit 120. This operation is essentially an efficient memory read operation, with computational overhead far lower than large-scale matrix multiplication.

[0036] Finally, in step 30, the query, key, and value vector obtained in the previous step is combined with the specific position information P of the target word "today" in the current input sequence "Today the weather is very good" (e.g., the position vector obtained through rotational position encoding, absolute position encoding, etc.), and then fed into the attention calculation module. A typical attention calculation function f can be expressed as: Here, the function f can be a standard attention calculation formula. For example, after incorporating location information P, the following calculation is performed: Where Q', K', and V' are vectors fused with positional information P. In this way, the inherent semantic representation of a word is decoupled from its contextual positional information in the sequence, and fused dynamically only when computation is required. This step can be performed by a single computational unit.

[0037] For other lexical units in the input sequence ("weather", "very", "good"), the system repeats steps 10, 20, and 30 above. In this way, this embodiment pre-processes most of the computation for the specified attention layer offline, transforming it into an efficient "lookup-read-fusion computation" mode during online inference, significantly reducing computational latency. More importantly, since all concurrently processed sequences share the same global QKV vector table, storage usage no longer increases linearly with the number or length of sequences, thus greatly saving valuable storage resources and solving the bottleneck problem of existing technologies.

[0038] Example 2

[0039] Based on Embodiment 1, this embodiment proposes a better storage scheme, namely a hierarchical storage strategy, aiming to further improve the access speed of vectors and thus enhance the overall inference performance. The technical solution of this embodiment can be understood in combination with Figure 3 and Figure 6 for understanding.

[0040] In Embodiment 1, we stored all one hundred thousand groups of query, key, and value vectors in the high-bandwidth memory 140. However, in practical applications, there are huge differences in the usage frequencies of different tokens in the vocabulary, which generally follow a long-tailed distribution. That is to say, a small number of high-frequency tokens (such as "的", "是", "我", and common punctuation marks) account for the vast majority of the text composition, while a large number of low-frequency tokens (such as rare characters, technical terms) rarely appear. This embodiment precisely utilizes this characteristic to perform an intelligent layout of the pre-computed query, key, and value vectors.

[0041] Referring to Figure 3 , this embodiment constructs a two-level storage system. Specifically, the cache 130 with extremely high access speed but small capacity (such as the secondary cache L2 Cache) is used as the first storage device, and the high-bandwidth memory 140 with slightly slower access speed but larger capacity is used as the second storage device. Obviously, the access speed of the first storage device is higher than that of the second storage device.

[0042] The hierarchical storage strategy can be executed according to the following steps: First, in the offline stage, in addition to performing the pre-computation described in Embodiment 1, it is also necessary to analyze the preset usage frequencies of all tokens in the vocabulary. This frequency can be obtained through statistical analysis of a large-scale and representative corpus (such as Wikipedia, general web texts, etc.), thereby obtaining a token frequency sorted list.

[0043] Secondly, according to the token frequency sorting and the capacity of the storage device, a storage strategy is formulated. For example, if it is statistically found that the tokens ranked in the top 5% of the frequency in the vocabulary (i.e., 5000 high-frequency tokens) contribute 80% of the access volume in daily texts, at this time, the storage management unit 120 will preferentially load the query, key, and value vectors corresponding to these 5000 high-frequency tokens into the fastest first storage device, that is, the cache 130.

[0044] For the remaining 95,000 low-frequency tokens, or as a more robust strategy, for the query, key, and value vectors of all one hundred thousand tokens, they are stored in the second storage device with a larger capacity, that is, the high-bandwidth memory 140. Storing all the data in the high-bandwidth memory 140 can serve as a backup for the data in the cache 130 to ensure that any token can be accessed.

[0045] When the model performs online inference, if the inference engine 110 needs to obtain the query, key, and value vector of a target term, the retrieval logic of the storage management unit 120 is adjusted accordingly: 1. First, check whether the target term is a high-frequency term and attempt to obtain its query, key, and value vector from the first storage device (cache 130). 2. If the acquisition is successful, i.e., a "cache hit", the data is returned to the computing unit with extremely low latency, and inference continues. Since most accesses are for high-frequency terms, a cache hit is achieved in most cases, significantly improving the average access speed. 3. If the data cannot be obtained from the cache 130, i.e., a "cache miss", the storage management unit 120 will turn to the second storage device (high-bandwidth memory 140) for searching and reading. Since the high-bandwidth memory 140 contains full or low-frequency data, this operation can successfully obtain the data. After obtaining the data, it is returned to the computing unit.

[0046] Please see Figure 6 This diagram illustrates the typical processing flow of a "cache miss" event through a signaling interaction sequence diagram. 1. The inference engine 110 sends a request to the storage management unit 120 to retrieve the vector of word A. 2. The storage management unit 120 first sends a query request to the cache 130. 3. The cache 130 does not contain data for word A, so it returns a miss response to the storage management unit 120. 4. Upon receiving the miss response, the storage management unit 120 immediately sends a read request to the high-bandwidth memory 140, requesting the vector data of word A. 5. The high-bandwidth memory 140 finds the data and returns it to the storage management unit 120. 6. The storage management unit 120 finally returns the vector data retrieved from the high-bandwidth memory 140 to the inference engine 110 for subsequent calculations. 7. As an optional optimization strategy, while returning data to the inference engine, the storage management unit 120 can write the vector data of word A, which has just been read from high-bandwidth memory 140, into the cache 130. This "write-back" strategy is based on the principle of locality of access, meaning that recently accessed items are likely to be accessed again in the near future. Thus, if the vector of word A is needed again soon, it can be directly retrieved from the cache.

[0047] Through this hierarchical storage and "miss-origin-write-back" mechanism, this embodiment fully utilizes the performance differences of different storage hardware without sacrificing functional completeness, enabling the system to serve the vast majority of vector requests at near-caching speed. This further reduces inference latency and improves the system's response speed compared to Embodiment 1.

[0048] Example 3

[0049] This embodiment, based on the hierarchical storage strategy described in Embodiment 2, further introduces a dynamic cache management mechanism. This aims to address the problem that static storage layouts cannot adapt to dynamically changing system loads and input data distributions, thereby enabling the cache system to have adaptive capabilities and higher resource utilization efficiency. The technical solution of this embodiment can be found in [reference needed]. Figure 4 To understand.

[0050] In Example 2, we place a fixed vector of 5000 high-frequency terms into cache 130. However, in actual operation, the following situations may occur: First, a sudden surge in requests may require caching many terms not currently in cache 130, causing cache 130 to quickly run out of space; second, user input may be concentrated in a specific domain for a period of time, causing some originally low-frequency terms (such as domain-specific terms) to become high-frequency in a short period. Static strategies cannot efficiently cope with these changes.

[0051] To this end, this embodiment adds a dynamic cache management module to the storage management unit 120, which executes a continuous "monitor-decision-execution" cycle.

[0052] Please see Figure 4 This flowchart details the process of a typical passive cleanup (or "eviction") strategy in dynamic cache management. 1. Monitor storage space: The dynamic cache management module continuously monitors the available space of the first storage device (i.e., cache 130). This monitoring can be periodic or event-triggered (e.g., whenever a new write request is received).

[0053] 2. Judgment and Triggering: The module compares the monitored available space with a preset threshold. For example, this threshold can be set to 15% of the total cache capacity. When the available space is less than this threshold, it indicates that the cache is under pressure, and the system will trigger the cache cleanup process.

[0054] 3. Decision-Making: After the cleanup process starts, the first step is to determine the size of the space to be released, denoted as n. This value n can be a fixed value or dynamically calculated based on the current system load. Then, it needs to determine which data should be evicted. To minimize the impact on performance, the cleanup process should prioritize evicting the lowest-value data. In the scenario of this application, the value of data can be measured based on the frequency of term usage. The cleanup process will select the query, key, and value vectors corresponding to the few least frequently used terms in the current cache 130 as the evictory targets. Furthermore, to ensure data consistency and the correctness of inference, an important constraint is that any term vectors currently being used by the active inference sequence cannot be evicted.

[0055] 4. Execution: After identifying m low-frequency and inactive terms to be eliminated, the cleanup program performs deletion or relocation operations. Specifically, it deletes the query, key, and value vectors corresponding to these m terms from the cache 130. Since a full data backup is stored in the high-bandwidth memory 140, this deletion operation is safe and will not result in data loss.

[0056] In one specific implementation of this embodiment, the process of determining the number of terms to be eliminated, m, can be more precise and quantifiable. For example, m can be calculated using a formula based on the required storage space size n to be released. Assuming that the storage space occupied by the query, key, and value vectors of each term is fixed, a proportionality coefficient λ can be defined, whose value is approximately equal to 1 / (the size of a single term's QKV vector). In this case, the number of terms to be eliminated, m, can be determined using the following formula: Here, ceil() is a floor function that ensures the space released is not less than the required space n. This formula provides a clear and executable calculation basis for dynamic management.

[0057] In addition to the passive cleanup strategies mentioned above, the dynamic cache management module can also include proactive loading and update strategies. For example, when the system load decreases and the available space in cache 130 becomes sufficient, the system can proactively preload some term vectors whose usage frequency has significantly increased and which are not currently in cache 130 from high-bandwidth memory 140 into cache 130 based on recent access frequency statistics. The system can maintain a dynamic term frequency counter and periodically update the contents of cache 130 to ensure that it always contains the set of terms with the highest current usage frequency.

[0058] By introducing this dynamic cache management mechanism, this embodiment transforms the entire cache system from a static layout into a dynamic system capable of adapting to load changes and learning. Whether under sudden high-concurrency scenarios or in a stable operating state, this mechanism ensures that cache resources are always used to store the most valuable data, thereby guaranteeing system robustness while continuously providing high-performance inference services.

[0059] Example 4

[0060] This embodiment aims to demonstrate the flexibility and scalability of the global caching framework proposed in this application, especially its compatibility with different types of attention mechanisms and support for dynamic switching between them to adapt to different business needs and hardware resource constraints.

[0061] There are several variations of attention mechanisms in existing technologies, such as standard multi-head attention, multi-query attention, and grouped query attention. Multi-head attention equips each query head with an independent key and value header, providing optimal performance but incurring the highest overhead. Multi-query attention allows all query heads to share the same set of key and value headers, greatly saving storage and bandwidth, but potentially sacrificing some performance. Grouped query attention is a compromise between the two, grouping query heads and sharing key and value headers within each group.

[0062] In practical inference services, different attention mechanisms may need to be flexibly selected based on different business scenarios (e.g., online dialogues with extremely high latency requirements versus offline translations with high throughput requirements) or the current system load. The framework in this application can well support this requirement.

[0063] The specific implementation process is as follows:

[0064] 1. Offline preparation: In the offline pre-computation stage, the system no longer generates query, key, and value vectors for only one attention mechanism, but performs pre-computation separately for multiple mechanisms that need to be supported.

[0065] For multi-head attention patterns: Following the method in Example 1, a complete set of query, key, and value vector tables is pre-computed and stored. In this table, for a multi-head attention layer, each attention head for each word has its own independent query, key, and value vector.

[0066] For grouped query attention mode: Pre-computation is also performed, but multiple query heads share the same set of key and value vectors. Therefore, the number of generated key and value vectors will be much smaller than in multi-head attention mode.

[0067] For multi-query attention mode: similarly, fewer key and value vectors are generated, and all query headers share a single set of keys and values.

[0068] All these pre-computed query, key, and value vector tables for different schemas can be viewed as different "vector libraries." These vector libraries can be stored entirely in a large-capacity third-level storage, for example... Figure 5 The solid-state drive shown is 150.

[0069] 2. Mode loading and switching: When the inference system 100 starts up, or when the mode needs to be changed according to the operation and maintenance instructions, the storage management unit 120 will perform the corresponding loading operation.

[0070] Loading multi-head attention mode: If the system needs to run at the highest performance, the control module will instruct the storage management unit 120 to read the QKV vector library prepared for multi-head attention mode from the solid-state drive 150 and load it into the graphics processor's storage system (i.e., high-bandwidth memory 140 and cache 130, which can also adopt the hierarchical and dynamic management strategies of embodiments 2 and 3).

[0071] Switching to Grouped Query Attention Mode: Suppose that after the system has been running for a period of time, the increased number of concurrent users causes the video memory bandwidth to become a bottleneck. The operations and maintenance personnel can issue an instruction to switch to the grouped query attention mode, which has lower video memory usage. At this time, the storage management unit 120 will first clear the vector data related to multi-head attention currently in the high-bandwidth memory 140 and cache 130, then read the QKV vector library corresponding to grouped query attention from the solid-state drive 150 and load it into the graphics processor.

[0072] From the model's perspective, the entire switching process requires no code recompilation or modification of the model structure; it is simply a data loading and replacement process. This makes the switching of attention mechanisms very fast and flexible.

[0073] The beneficial effects of this embodiment are that it provides inference systems with significant operational flexibility and cost-effectiveness. Service providers can dynamically adjust the implementation of the underlying attention algorithm based on real-time performance requirements (using multi-head attention when high performance is needed), cost considerations (using grouped query attention or multi-query attention when resource conservation is needed), or specific service level agreements requested by users, without downtime or complex software updates. The global pre-computation framework of this application, by separating computation from data, simplifies algorithm switching to data file switching, greatly enhancing the system's adaptability and maintainability.

[0074] In summary, this application provides a complete and flexible attention mechanism caching optimization scheme through the above embodiments. From basic pre-computation and global sharing to advanced hierarchical storage and dynamic management, and then to flexible multi-mode compatibility, the various technical points of this application are interconnected and work together to solve the core bottlenecks in large language model inference, thereby significantly improving the model's processing power and overall service efficiency.

[0075] Accordingly, this application also provides an apparatus embodiment, such as an inference system 100 configured with the aforementioned method. The pre-computation unit, storage management unit (e.g., storage management unit 120), retrieval unit, and computation unit (which may be integrated into the inference engine 110) in this system work together to execute the aforementioned method flow. Similarly, storing a computer program implementing the aforementioned method in a computer-readable storage medium allows the same technical effect to be achieved when the processor executes the program.

[0076] The above description is merely a preferred embodiment of this application and is not intended to limit the scope of this application. Various modifications and variations can be made to this application by those skilled in the art. Any modifications, equivalent substitutions, improvements, etc., made within the spirit and principles of this application should be included within the protection scope of this application.

Claims

1. A caching method in an attention mechanism, characterized in that, include: For all lexical units in the vocabulary, based on the projection matrix of at least one specified layer in the model, a set of query Q, key K, and value V vectors that correspond one-to-one with all lexical units and do not contain position information are pre-calculated and generated. The set of Q, K, V vectors is stored in at least one storage device, and a mapping relationship is established between all the word elements and the set of Q, K, V vectors. During the model inference process, for the target word in the input sequence, based on the mapping relationship, the Q, K, and V vectors corresponding to the target word are retrieved from the at least one storage device; The attention result is calculated based on the obtained Q, K, V vectors and the position information of the target word in the input sequence.

2. The method according to claim 1, characterized in that, Storing the set of Q, K, V vectors in at least one storage device includes: The set of Q, K, V vectors are stored hierarchically in a first storage device and a second storage device, wherein the access speed of the first storage device is higher than that of the second storage device; Furthermore, the hierarchical storage includes: based on the preset usage frequency of each word in all the words, prioritizing the storage of Q, K, and V vectors corresponding to high-frequency words in the first storage device, and storing the Q, K, and V vectors corresponding to low-frequency words in the second storage device.

3. The method according to claim 2, characterized in that, The first storage device is a cache, and the second storage device is high-bandwidth memory.

4. The method according to claim 2 or 3, characterized in that, Also includes: Monitor the available space of the first storage device; When the available space is less than a preset threshold, the required amount of storage space to be released is determined, and at least one Q, K, V vector corresponding to the low-frequency word is deleted or moved from the first storage device based on the required amount of storage space to be released.

5. The method according to claim 4, characterized in that, The number m of Q, K, and V vectors to be deleted or moved is based on the required storage space size n, calculated using the formula... It is determined, where λ is the proportionality coefficient.

6. The method according to claim 1, characterized in that, During model inference, the projection matrix used for pre-calculation is not stored in the video memory or cache of the computing device.

7. The method according to claim 1, characterized in that, The designated layer is the model's first attention layer.

8. A caching device in an attention mechanism, characterized in that, include: The pre-computation unit is used to pre-compute a set of query Q, key K, and value V vectors that correspond one-to-one with all the lexical units in the vocabulary and do not contain position information, based on the projection matrix of at least one specified layer in the model. A storage management unit is used to store the set of Q, K, V vectors in at least one storage device and establish a mapping relationship between all the word elements and the set of Q, K, V vectors; The retrieval unit is used, during the model inference process, to retrieve, based on the mapping relationship, the Q, K, and V vectors corresponding to the target word in the input sequence from the at least one storage device. The computing unit is used to calculate the attention result based on the acquired Q, K, V vectors and the position information of the target word in the input sequence.

9. The apparatus according to claim 8, characterized in that, The storage management unit is further configured to: store the set of Q, K, V vectors hierarchically in a first storage device and a second storage device, wherein the access speed of the first storage device is higher than that of the second storage device; and, based on the preset usage frequency of each word in all the words, prioritize storing the Q, K, V vectors corresponding to high-frequency words in the first storage device and the Q, K, V vectors corresponding to low-frequency words in the second storage device; monitor the available space of the first storage device; when the available space is less than a preset threshold, determine the amount of storage space to be released, and delete or move at least one Q, K, V vector corresponding to the low-frequency word from the first storage device based on the amount of storage space to be released.

10. A computer-readable storage medium having a computer program stored thereon, which, when executed by a processor, implements the method of any one of claims 1 to 7.

Images