A prefix caching method and apparatus suitable for a hybrid attention model architecture
By using a unified prefix indexing mechanism and a dual verification mechanism based on chained hashing, the full attention key-value cache and linear attention state in the hybrid attention model are managed collaboratively, which solves the problem of high computational complexity in the hybrid attention model and improves inference efficiency and system stability.
Patent Information
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- SHANGHAI SUIYUAN TECH CO LTD
- Filing Date
- 2026-04-27
- Publication Date
- 2026-07-07
Smart Images
Figure CN122088718B_ABST
Abstract
Description
Technical Field
[0001] This invention relates to the fields of artificial intelligence and natural language processing, and in particular to a prefix caching method and apparatus suitable for hybrid attention model architectures. Background Technology
[0002] Traditional Transformer architectures employ a full attention mechanism, resulting in a computational complexity of O(n²). During inference, key-value pairs are typically stored in a key-value cache (KV Cache) to avoid redundant computation. Existing prefix caching techniques usually divide the token sequence into fixed-size blocks and construct a prefix hash tree based on chained hash values, reusing physical blocks of the KV Cache through longest common prefix matching.
[0003] To address the computational complexity bottleneck of full attention mechanisms in scenarios with extremely long contexts, linear attention mechanisms (such as GatedDeltaNet or Mamba architectures) are proposed, reducing computational complexity to O(n). Linear attention layers maintain the continuity of sequence information through convolutional and recurrent states, rather than relying on a key-value cache. The recurrent state, which compresses and stores historical information using a fixed-size matrix, is key to achieving efficient inference with linear attention. In next-generation models, full attention and linear attention layers are typically alternated within the same network structure. For example, some layers employ full attention to handle precise long-range dependencies, while others use linear attention to efficiently compress historical information. This hybrid attention architecture improves performance but also introduces complexity in cache management.
[0004] Current technologies only manage the KV cache, neglecting to cache the linear attention states in the hybrid attention model. Even if the KV cache prefix is fully hit, the linear attention layer still needs to recalculate from the initial state, thus reducing the overall effectiveness of the prefix cache. Furthermore, the KV cache stores data in blocks, supporting reference counting and sharing, while the linear attention state is a compressed representation of the entire prefix sequence and cannot be managed in blocks. These two technologies differ fundamentally in their management methods and reuse strategies. Therefore, a unified prefix caching scheme that simultaneously supports full attention caching and linear attention state caching is urgently needed to improve the inference efficiency of the hybrid attention model. Summary of the Invention
[0005] This invention provides a prefix caching method and apparatus suitable for hybrid attention model architectures, so as to achieve collaborative caching and reuse of full attention key-value caching and linear attention states in hybrid attention models.
[0006] According to a first aspect of the present invention, a prefix caching method suitable for hybrid attention model architectures is provided, the method comprising:
[0007] The request sequence received based on the hybrid attention model is divided into blocks and a chained hash value is calculated. The chained hash value is then matched with the longest common prefix of the prefix hash tree constructed from the historical request sequence to obtain the matching cache block. The hybrid attention model includes a linear attention layer and a full attention layer.
[0008] Perform linear attention snapshot verification on the matching cache blocks, and determine the matching cache blocks that pass the verification as valid cache blocks;
[0009] Based on the valid cache block, determine the first sequence corresponding to the valid cache block and the remaining second sequence in the request sequence;
[0010] For the first sequence, the key-value cache information and linear attention snapshot corresponding to the effective cache block are reused;
[0011] For the second sequence, perform pre-filled inference.
[0012] Optionally, the linear attention snapshot verification of the matching cache block includes:
[0013] Starting from the last matching cache block, traverse backwards one by one to obtain the hash value of the current matching cache block, and query the local valid snapshot hash set based on the hash value. The local valid snapshot hash set includes the correspondence between hash values and linear attention snapshots, and the linear attention snapshot includes convolution state and loop state.
[0014] When the hash value exists in the local valid snapshot hash set, local verification is determined to be successful, and a query command containing the hash value is sent to the GPU inference process via inter-process communication for remote verification.
[0015] When the hash value does not exist in the local valid snapshot hash set, the current matching cache block is determined to be an invalid cache block, and the previous matching cache block is traversed and verified.
[0016] Optionally, the remote verification includes:
[0017] When a confirmation instruction is received from the GPU inference process, the remote verification is confirmed to be successful, and the currently matching cache block is determined to be a valid cache block.
[0018] If the confirmation instruction is not received, the currently matching cache block will be determined as an invalid cache block.
[0019] Optionally, before performing pre-padded inference on the second sequence, the method further includes:
[0020] The linear attention snapshot is then validated a second time.
[0021] When the secondary verification passes, pre-filled inference is performed on the second sequence based on the key-value cache information and the linear attention snapshot;
[0022] After the pre-filled inference is completed, the request sequence is marked as recovered.
[0023] Optionally, the method further includes: immediately capturing a linear attention snapshot and key-value cache information at the current sequence position after completing pre-filled inference for the second sequence;
[0024] During subsequent decoding and inference, when the number of newly added sequence word elements reaches the block boundary, a linear attention snapshot and key-value cache information are captured at the current sequence position;
[0025] The key-value cache information is saved to the key-value cache storage area in the form of cache blocks, and a corresponding prefix hash tree is constructed based on the hash value of each cache block;
[0026] The linear attention snapshot is saved as a state tensor to the corresponding slot in the state cache storage area, wherein the linear attention snapshot and key-value cache information corresponding to the same capture position use the hash value of the same cache block as the index.
[0027] Optionally, after saving the linear attention snapshot as a state tensor to the corresponding slot in the state cache storage area, the method further includes:
[0028] The set of valid snapshot hash values is determined by the dual-caching coordinator based on the hash values of cache blocks;
[0029] Establish a forward mapping between cache block identifiers and snapshot hash value sets, and a reverse mapping between snapshot hash values and cache block identifiers.
[0030] Optionally, after performing pre-filled inference, the following is also included:
[0031] When the eviction triggering conditions are met, the target key-value cache information in the key-value cache storage area is evicted, and an eviction cache block list is generated according to the cache block identifier corresponding to the evicted key-value cache information. The eviction triggering conditions include prefix cache capacity exceeding the limit, GPU memory pressure, insufficient swap space, or global release.
[0032] Based on the list of evicted cache blocks, remove the hash values associated with the list of evicted cache blocks from the set of valid snapshot hash values;
[0033] Remove the entry associated with the snapshot hash from the reverse mapping from the snapshot hash to the cache block identifier;
[0034] Remove entries associated with the list of evicted cache blocks from the positive mapping from the cache block identifier to the set of snapshot hashes.
[0035] According to another aspect of the present invention, a prefix caching device suitable for hybrid attention model architectures is provided, the device comprising:
[0036] The matching module is used to divide the request sequence received based on the hybrid attention model into blocks and calculate the chained hash value, and perform the longest common prefix matching between the chained hash value and the prefix hash tree of the historical request sequence to obtain the matching cache block. The hybrid attention model includes a linear attention layer and a full attention layer.
[0037] The verification module is used to perform linear attention snapshot verification on the matching cache block and determine the matching cache block that passes the verification as a valid cache block;
[0038] The sequence determination module is used to determine, based on the valid cache block, a first sequence corresponding to the valid cache block and the remaining second sequence in the request sequence;
[0039] The reuse module is used to reuse the key-value cache information and linear attention snapshot corresponding to the effective cache block for the first sequence;
[0040] A pre-filled inference module is used to perform pre-filled inference for the second sequence.
[0041] According to another aspect of the present invention, an electronic device is provided, the electronic device comprising: one or more processors; and a storage device for storing one or more programs;
[0042] When the one or more programs are executed by the one or more processors, the one or more processors are made to implement the method described in any embodiment of the present invention.
[0043] According to another aspect of the present invention, a computer-readable storage medium is provided having a computer program stored thereon, which, when executed by a processor, implements any of the methods described in the embodiments of the present invention.
[0044] The technical solution of this invention adopts a unified prefix index mechanism based on chain hashing, which enables two types of caches to share a hash space, performs the longest prefix matching and traverses backward from the initial matching position to verify and lock the consistency boundary that is valid for both types of caches, and collaboratively restores and reuses key-value caches and linear attention snapshots, significantly reducing the risk of heterogeneous state out-of-sync.
[0045] It should be understood that the description in this section is not intended to identify key or important features of the embodiments of the present invention, nor is it intended to limit the scope of protection of the present invention. Other features of the present invention will become clearer from the following description. Attached Figure Description
[0046] To more clearly illustrate the technical solutions in the embodiments of the present invention, the accompanying drawings used in the description of the embodiments will be briefly introduced below. Obviously, the accompanying drawings described below are only some embodiments of the present invention. For those skilled in the art, other drawings can be obtained based on these drawings without creative effort.
[0047] Figure 1 This is a flowchart of a prefix caching method suitable for a hybrid attention model architecture according to Embodiment 1 of the present invention;
[0048] Figure 2 This is a diagram of a prefix caching system architecture for a hybrid attention model according to Embodiment 1 of the present invention;
[0049] Figure 3 This is a synchronous capture timing diagram summarizing the inference process provided in Embodiment 1 of the present invention;
[0050] Figure 4 This is a flowchart of another prefix caching method suitable for hybrid attention model architectures provided in Embodiment 2 of the present invention;
[0051] Figure 5 This is a schematic diagram of a prefix caching device suitable for a hybrid attention model architecture according to Embodiment 3 of the present invention;
[0052] Figure 6 This is a structural block diagram of an electronic device provided in Embodiment 4 of the present invention. Detailed Implementation
[0053] To enable those skilled in the art to better understand the present invention, the technical solutions of the present invention will be clearly and completely described below with reference to the accompanying drawings of the embodiments of the present invention. Obviously, the described embodiments are only some embodiments of the present invention, and not all embodiments. Based on the embodiments of the present invention, all other embodiments obtained by those skilled in the art without creative effort should fall within the scope of protection of the present invention.
[0054] It should be noted that the terms "first," "second," etc., in the specification, claims, and accompanying drawings of this invention are used to distinguish similar objects and are not necessarily used to describe a specific order or sequence. It should be understood that such data can be interchanged where appropriate so that the embodiments of the invention described herein can be implemented in orders other than those illustrated or described herein. Furthermore, "comprising" and "having," and any variations thereof, are intended to cover non-exclusive inclusion; for example, a process, method, apparatus, product, or terminal device that includes a series of steps or units is not necessarily limited to those steps or units explicitly listed, but may include other steps or units not explicitly listed or inherent to such processes, methods, products, or terminal devices.
[0055] Figure 1 This is a flowchart illustrating a prefix caching method suitable for a hybrid attention model architecture, as provided in Embodiment 1 of the present invention. This embodiment is applicable to scenarios where prefixes are cached based on a hybrid attention model. The method can be executed by a prefix caching device suitable for a hybrid attention model architecture. This device can be implemented in hardware and / or software and can be integrated into an electronic device with data processing capabilities.
[0056] like Figure 1 As shown, the method includes:
[0057] S101, the request sequence received based on the hybrid attention model is divided into blocks and chained hash values are calculated. The chained hash values are then matched with the longest common prefix of the prefix hash tree of the historical request sequence to obtain the matching cache block.
[0058] Among them, such as Figure 2 The diagram shown illustrates the prefix caching system architecture of the hybrid attention model in this embodiment. In this architecture, the decoder layers are stacked as follows: "Layer 0 (Full Attention) → Layer 1 (Linear Attention) → Layer 2 (Full Attention) → Layer 3 (Linear Attention) → ...". It should be noted that the above structure is merely an example and does not limit the specific stacking method of the decoder layers.
[0059] The full attention layer connects to the KV Cache storage area, while the linear attention layer connects to the MambaCache storage area (state cache storage area). Both are managed uniformly by a dual-cache coordinator and connected by a unified PrefixCache hash index.
[0060] The dual-buffer coordination manager incorporates a triple mapping structure:
[0061] Forward mapping: cache block identifier A set of snapshot hashes, used to record the linear attention snapshot hashes associated with each cache block;
[0062] Reverse mapping: Snapshot hash Cache block identifier, used to record the cache block corresponding to each snapshot hash value;
[0063] Valid snapshot hash set: used to record all currently valid snapshot hash values and supports fast verification with O(1) complexity.
[0064] Through the above triple mapping relationship, the dual-cache coordinator binds the lifecycle of cache blocks in the KV Cache with the logical lifecycle of linear attention snapshots, thereby achieving coordinated management.
[0065] In addition, the system includes a scheduler whose cooperative operation instructions are executed in the following order during inference: matching, verification, allocation, capture, and eviction. When the hybrid attention model receives a new request sequence, it will execute the corresponding inference process based on the above system architecture. It should be noted that the above structure is only an example and does not limit the specific implementation of the system architecture.
[0066] In a specific implementation, when a request sequence is received, it can be parsed based on the above system architecture to identify which parts of the request sequence have already undergone inference and cached key-value cache information and linear attention snapshot information. For the identified parts, redundant calculations can be skipped, and historical inference results can be directly reused, thereby significantly improving inference efficiency.
[0067] Specifically, in this implementation, the request sequence is divided into blocks of a fixed size (e.g., 16 tokens), and a hash value is calculated for each block. The hash value of each block contains the hash information of the preceding block, thus forming a chain dependency relationship. That is, the hash value of the i-th block depends not only on the tokens within that block, but also on the hash value of the previous block.
[0068] For example, the request sequence can be divided into blocks such as Block1, Block2, ..., Blockn, and their corresponding chain hashes are H1, H2, H3, ..., Hn.
[0069] Furthermore, the prefix hash tree stores the KV cache block data structure that has been computed in the historical request sequence. Therefore, during matching, the chained hashes H1, H2, H3, ... are searched sequentially in the prefix hash tree. As long as a certain Hi exists in the cache, it means that the first i blocks have been computed in the historical request and can be directly reused.
[0070] For example, if a match is found in the prefix hash tree and `matched_blocks = N`, it means that there are N matching cache blocks. It should be noted that this example is for illustrative purposes only and does not limit the specific form of the matching cache blocks.
[0071] Furthermore, after obtaining the matching cache block, it is necessary to determine whether the last block in the request sequence is a complete block. If the last block is not full (for example, the block size is 16, but it only contains 15 tokens), the number of matching cache blocks is reduced by 1, that is, the final number of matching cache blocks is N-1, and the incomplete block is discarded; if the last block is a complete block, no adjustment is needed, and the number of matching cache blocks remains N. It should be noted that the above processing method is only an example and does not constitute a limitation on specific adjustment strategies.
[0072] S102, perform linear attention snapshot verification on the matching cache blocks, and determine the verified matching cache blocks as valid cache blocks.
[0073] Optionally, linear attention snapshot verification is performed on the matching cache blocks, including: traversing backwards from the last matching cache block to obtain the hash value of the current matching cache block, and querying the local valid snapshot hash set based on the hash value, wherein the local valid snapshot hash set includes the correspondence between hash values and linear attention snapshots, and the linear attention snapshots include convolutional states and loop states; when the hash value exists in the local valid snapshot hash set, the local verification is determined to be successful, and a query instruction containing the hash value is sent to the GPU inference process via inter-process communication for remote verification; when the hash value does not exist in the local valid snapshot hash set, the current matching cache block is determined to be an invalid cache block, and the traversal and verification of the previous matching cache block continues.
[0074] Optionally, remote verification includes: when a confirmation instruction is received from the GPU inference process, determining that the remote verification is successful and identifying the currently matching cache block as a valid cache block; when no confirmation instruction is received, identifying the currently matching cache block as an invalid cache block.
[0075] In this implementation, a hybrid attention model incorporating both linear and full attention layers is used. In the pure Transformer model, key-value caching is universally applicable as long as the token sequence is the same. However, in the hybrid attention model with linear attention, the linear attention snapshot maintains a compressed state that depends not only on the current token but also strictly on the historical state. Therefore, even after confirming a successful key-value cache match, further linear attention snapshot verification of the matching cache block is required. Only when the linear attention snapshot verification passes can the cached information in the matching cache block be considered correct and directly reused during the inference process of the current request sequence, thus ensuring the integrity and accuracy of the reused information.
[0076] Specifically, in this implementation, verification is performed by starting from the last matched cache block and proceeding backwards. For example, when N matching cache blocks are obtained by matching with the prefix hash tree, the system does not directly trust the matching result, but instead checks and verifies each cache block in reverse order starting from the Nth cache block.
[0077] In the verification of the Nth cache block, a dual verification mechanism is adopted, namely local verification and remote verification. Local verification refers to maintaining a set of local valid snapshot hashes in CPU memory. This set records the correspondence between hash values and linear attention snapshots (including convolutional and cyclic states), and is used to identify all cache blocks known to have complete linear attention states. It can be regarded as a whitelist of valid snapshot hashes. When the hash value of the Nth matching cache block is not in this set, it means that it does not have a valid linear attention state and can be directly judged as a verification failure without accessing the GPU. This process only requires a query operation with an O(1) time complexity.
[0078] When local verification fails, the Nth matching cache block is directly identified as invalid without remote verification, and the verification process continues for the (N-1)th matching cache block, thus avoiding unnecessary computation and communication overhead. Only when local verification succeeds is a query command containing the hash value sent to the GPU inference process via inter-process communication for remote verification. The reason for remote verification is that the hash table maintained by the CPU may lag behind the actual state of the GPU memory. For example, the GPU may have reclaimed the corresponding state data due to memory pressure, while the CPU has not yet updated it. Furthermore, the storage structure of linear attention states is relatively complex, and only the GPU process can accurately determine whether the state is still valid.
[0079] When the GPU receives a query command containing a hash value, it will perform a retrieval based on that hash value and return either a confirmation command or a non-confirmation command based on the retrieval result. A confirmation command indicates that the linear attention state exists and is valid, while a non-confirmation command indicates that the linear attention state does not exist or is invalid. Only when a confirmation command is received is the remote verification considered successful, and the Nth matching cache block is determined to be a valid cache block; otherwise, it is determined to be an invalid cache block, and the verification process continues for the (N-1)th matching cache block.
[0080] Through the aforementioned dual verification mechanism of "local verification + remote verification + failure fallback," a set of valid cache blocks that have passed verification can be obtained. The number of valid cache blocks does not exceed the number of matching cache blocks; if all verifications pass, the two numbers are the same; however, if any matching cache block fails verification, the number of valid cache blocks will be less than the number of matching cache blocks. This implementation does not limit the number of valid cache blocks, where valid cache blocks are truly secure and reusable cache blocks.
[0081] S103, based on the valid cache blocks, determine the first sequence in the request sequence that corresponds to the valid cache blocks and the remaining second sequence.
[0082] Specifically, when the request sequence is divided into blocks (Block1, Block2, Block3, Block4, and Block5), if Block1 to Block4 have already obtained corresponding valid cached blocks during the historical request sequence reasoning process, then the sequence consisting of Block1 to Block4 corresponding to valid cached blocks in the request sequence (e.g., token1…token64) is determined as the first sequence, and the sequence corresponding to Block5 (e.g., token65…token80) is determined as the second sequence. It should be noted that this implementation only uses a scenario where each block consists of 16 tokens as an example. For scenarios with other block granularities, the determination of the first and second sequences is similar to the above process and will not be repeated here.
[0083] S104, for the first sequence, reuse the key-value cache information and linear attention snapshot corresponding to the valid cache block.
[0084] Specifically, for the first sequence, the inference engine does not need to allocate new GPU memory; instead, it directly reuses the key-value cache information and linear attention snapshot of the corresponding valid cache block. Simultaneously, a reference counter is set for each valid cache block, incrementing by 1 each time it is referenced. This reference counting mechanism prevents valid cache blocks from being erroneously reclaimed in subsequent processes; a valid cache block is only allowed to be released or overwritten when all requests referencing it have ended (i.e., the reference count reaches zero).
[0085] Furthermore, in this embodiment, the hash value (denoted as last_hash) corresponding to the end of the last valid cache block is recorded in the context data of the request sequence, and this hash value is used as the anchor point for state recovery. For example, when the GPU starts processing the subsequent Block 5, the corresponding linear attention state vector can be extracted from the video memory according to last_hash and loaded into a register or shared memory as the initial state for subsequent calculations, thereby ensuring that the calculation of the linear attention layer can be seamlessly connected from this "breakpoint" without having to recalculate from the initial state.
[0086] S105, Perform pre-filling inference for the second sequence.
[0087] Optionally, before performing pre-filling inference on the second sequence, the method further includes: performing secondary verification on the linear attention snapshot; when the secondary verification passes, performing pre-filling inference on the second sequence based on the key-value cache information and the linear attention snapshot; and marking the request sequence as recovered after the pre-filling inference is completed.
[0088] Specifically, in this embodiment, a new key-value cache block is allocated for the second sequence. During the inference process for the second sequence, the key-value cache block is gradually populated with data. Before the key-value cache block allocation is completed and inference officially begins, the cache state may change due to time delays or resource contention caused by concurrent requests. Therefore, a fast secondary verification is performed before the scheduler starts the GPU kernel.
[0089] For example, this fast verification may include checking whether the reference count of the previously reusable valid cache block is still non-zero, and whether the corresponding linear attention state snapshot still exists in GPU memory and has not been marked as invalid. The secondary verification is considered successful only if the reference count is non-zero and the linear attention snapshot is not invalid. This mechanism effectively avoids the problem of erroneous reuse caused by cache blocks being reclaimed or overwritten in a very short time under extreme concurrency conditions. If the secondary verification fails, it must immediately revert to the recalculation mode and stop executing the reuse path, thereby ensuring the correctness of the inference result.
[0090] If the secondary verification passes, pre-filling inference is performed on the second sequence based on the key-value cache information corresponding to the valid cache block and the linear attention snapshot.
[0091] Optionally, the method further includes: immediately capturing a linear attention snapshot and key-value cache information at the current sequence position after completing the pre-filling inference of the second sequence; capturing a linear attention snapshot and key-value cache information again at the current sequence position during subsequent decoding inference when the number of newly added sequence tokens reaches the block boundary; storing the key-value cache information in the form of cache blocks in the key-value cache storage area, and constructing a corresponding prefix hash tree based on the hash value of each cache block; storing the linear attention snapshot in the form of a state tensor in the corresponding slot of the state cache storage area, wherein the linear attention snapshot and key-value cache information corresponding to the same capture position use the same hash value as the index.
[0092] Optionally, after saving the linear attention snapshot as a state tensor to the corresponding slot in the state cache storage area, the method further includes: determining the set of valid snapshot hash values based on the hash values of the cache blocks through the dual cache coordinator; establishing a forward mapping relationship between the cache block identifier and the set of snapshot hash values, and a reverse mapping relationship between the snapshot hash values and the cache block identifier.
[0093] Among them, such as Figure 3 The diagram illustrates the synchronous capture timing during inference, demonstrating the synchronous execution relationship between key-value cache writing and linear attention snapshot capture at the same time. The horizontal axis represents time or the token sequence (Token1→2→…→Token16 (block boundary)→Token17→…→Token32 (block boundary)→…→EOS). At the key-value cache operation level, KV data is continuously written to the current physical cache block during inference. When the cache block is full or the sequence ends, the hash value of the complete cache block is inserted into the prefix index tree. At the state cache operation level, the convolutional state and the loop state are continuously updated at each token, and a snapshot of the current state is captured at the block boundary. The final state is captured at the end of the sequence. At each "capture point," the above two types of operations are executed synchronously, using the same hash value as a unified index to achieve consistency identification.
[0094] Furthermore, in this embodiment, key-value cache information is stored in the key-value cache storage area in the form of cache blocks, and linear attention snapshots are stored in the corresponding slots of the state cache storage area in the form of state tensors. After completing the cache write, the dual-cache coordination manager synchronously maintains three types of mapping relationships: a bidirectional mapping between cache block identifiers and snapshot hash values, and the inclusion of hash values into the set of valid snapshot hash values. It should be noted that this embodiment only provides an illustrative example of the reasoning process of the second sequence and does not constitute a limitation on the specific implementation method.
[0095] The technical solution of this invention adopts a unified prefix index mechanism based on chain hashing, which enables two types of caches to share a hash space, performs the longest prefix matching and traverses backward from the initial matching position to verify and lock the consistency boundary that is valid for both types of caches, and collaboratively restores and reuses key-value caches and linear attention snapshots, significantly reducing the risk of heterogeneous state out-of-sync.
[0096] Figure 4 This is a flowchart illustrating another prefix caching method suitable for a hybrid attention model architecture, provided by an embodiment of the present invention. Based on the above embodiments, this embodiment further includes performing a linked eviction when the eviction triggering condition is met after performing pre-filled inference. Figure 4 As shown, the method includes:
[0097] S201, the request sequence received based on the hybrid attention model is divided into blocks and chained hash values are calculated. The chained hash values are then matched with the longest common prefix of the prefix hash tree constructed from the historical request sequences to obtain matching cache blocks.
[0098] S202, perform linear attention snapshot verification on the matching cache blocks, and determine the verified matching cache blocks as valid cache blocks.
[0099] Optionally, linear attention snapshot verification is performed on the matching cache blocks, including: traversing backwards from the last matching cache block to obtain the hash value of the current matching cache block, and querying the local valid snapshot hash set based on the hash value, wherein the local valid snapshot hash set includes the correspondence between hash values and linear attention snapshots, and the linear attention snapshots include convolutional states and loop states; when the hash value exists in the local valid snapshot hash set, the local verification is determined to be successful, and a query instruction containing the hash value is sent to the GPU inference process via inter-process communication for remote verification; when the hash value does not exist in the local valid snapshot hash set, the current matching cache block is determined to be an invalid cache block, and the traversal and verification of the previous matching cache block continues.
[0100] Optionally, remote verification includes: when a confirmation instruction is received from the GPU inference process, determining that the remote verification is successful and identifying the currently matching cache block as a valid cache block; when no confirmation instruction is received, identifying the currently matching cache block as an invalid cache block.
[0101] S203, based on the valid cache blocks, determine the first sequence in the request sequence that corresponds to the valid cache blocks and the remaining second sequence.
[0102] S204, for the first sequence, reuse the key-value cache information and linear attention snapshot corresponding to the valid cache block.
[0103] S205, Perform pre-filling inference for the second sequence.
[0104] Optionally, before performing pre-filling inference on the second sequence, the method further includes: performing secondary verification on the linear attention snapshot; when the secondary verification passes, performing pre-filling inference on the second sequence based on the key-value cache information and the linear attention snapshot; and marking the request sequence as recovered after the pre-filling inference is completed.
[0105] Optionally, the method further includes: immediately capturing a linear attention snapshot and key-value cache information at the current sequence position after completing the pre-filling inference of the second sequence; capturing a linear attention snapshot and key-value cache information at the current sequence position when the number of newly added sequence tokens reaches the block boundary during subsequent decoding inference; saving the key-value cache information to the key-value cache storage area in the form of cache blocks, and constructing a corresponding prefix hash tree based on the hash value of each cache block; saving the linear attention snapshot to the corresponding slot in the state cache storage area in the form of a state tensor, wherein the linear attention snapshot and key-value cache information corresponding to the same capture position use the same hash value as the index.
[0106] Optionally, after saving the linear attention snapshot as a state tensor to the corresponding slot in the state cache storage area, the method further includes: determining the set of valid snapshot hash values based on the hash values of the cache blocks through the dual cache coordinator; establishing a forward mapping relationship between the cache block identifier and the set of snapshot hash values, and a reverse mapping relationship between the snapshot hash values and the cache block identifier.
[0107] S206, when the eviction triggering condition is met, the target key-value cache information in the key-value cache storage area is evicted, and an eviction cache block list is generated according to the cache block identifier corresponding to the evicted key-value cache information.
[0108] Specifically, in the hybrid attention model, key-value cache information and linear attention snapshots are tightly bound together based on a unified hash value. The key-value cache (KV Cache) stores the key / value matrix of the full attention layer, while the linear attention snapshot stores the compressed hidden state of the linear attention layer. Both must exist simultaneously or be deleted simultaneously. Deleting only one will lead to serious errors in subsequent cache reuse. For example, if the system triggers an eviction operation due to insufficient video memory, and only the KV Cache is deleted without simultaneously deleting the corresponding linear attention snapshot, a new request may be misjudged as a hit during the prefix matching stage and attempt to reuse the cache. However, when actually reading the KV Cache, data loss or errors may occur, leading to computational anomalies or output errors. Therefore, this implementation avoids the above problems through a linked eviction mechanism.
[0109] The eviction triggering conditions include, but are not limited to: prefix cache capacity exceeding the limit, excessive GPU memory pressure, insufficient swap space, or global release. When at least one eviction triggering condition is met, an eviction operation is performed on the target key-value cache information in the key-value cache storage area, and a list of evicted cache blocks is generated, which records the identifiers of the deleted cache blocks. It should be noted that key-value cache information is managed using cache block identifiers, while linear attention snapshots are indexed by hash values.
[0110] S207, based on the list of evicted cache blocks, remove the hash values associated with the list of evicted cache blocks from the set of valid snapshot hash values.
[0111] After obtaining the list of evicted cache blocks, the system finds the hash value associated with each evicted cache block and removes the hash value from the set of valid snapshot hash values, thereby deleting the corresponding invalid hash value from the whitelist. Thus, once a hash value is removed, subsequent requests, even if they calculate the same hash value, will be judged as a miss during the verification phase, forcing a recalculation, serving as the first line of defense against erroneous reuse.
[0112] S208, Remove the entry associated with the hash value from the reverse mapping from the snapshot hash value to the cache block identifier.
[0113] Furthermore, after deleting the hash value from the set of valid snapshot hash values, the corresponding entry must also be deleted from the reverse mapping relationship of "snapshot hash value → cache block identifier". If this operation is not performed, the system may still retain the mapping relationship of "hash value → old cache block identifier", which, although the hash value is no longer valid, may cause subsequent logical errors or conflicts. Therefore, it is necessary to thoroughly clean up the mapping relationships pointing to freed memory.
[0114] S209, Remove entries associated with the list of evicted cache blocks from the forward mapping from cache block identifiers to snapshot hash value sets.
[0115] In addition, the corresponding entries need to be removed from the forward mapping relationship of "cache block identifier → snapshot hash set". Since the cache block has been evicted and no longer carries valid data, all mapping information with the cache block identifier as the key should be deleted to release related metadata resources and prevent the CPU-side index data from continuing to expand.
[0116] It is worth mentioning that this embodiment addresses the challenge of heterogeneous management of key-value caches and state snapshots in a hybrid architecture by employing a unified prefix index mechanism based on chained hashing, enabling both types of caches to share a hash space. This method performs longest prefix matching and verifies by traversing backward from the initial matching position. It locks the consistency boundary where both types of caches are valid through local snapshot queries and remote confirmations, thereby collaboratively restoring and reusing key-value caches and linear attention snapshots, significantly reducing the risk of heterogeneous state synchronization failures. During inference execution, both types of states are captured synchronously and use the same hash index. A dual-cache coordinator is used to establish a forward and reverse mapping relationship between cache block identifiers and snapshot hashes. When eviction conditions are met, the two types of caches are cleaned up in conjunction based on the mapping relationship. This invention solves the problem of heterogeneous cache storage and lifecycle collaborative management in a hybrid attention model architecture, achieving consistency maintenance throughout the entire process and ensuring system robustness while improving inference performance.
[0117] The technical solution of this invention adopts a unified prefix indexing mechanism based on chained hashing, enabling two types of caches to share a hash space. It performs longest prefix matching and traverses backward from the initial matching position to verify and lock the consistency boundary where both types of caches are valid. This collaboratively restores and reuses key-value caches and linear attention snapshots, significantly reducing the risk of heterogeneous state desynchronization. Simultaneously, through a linked eviction mechanism, when either type of cache is evicted due to resource pressure, the other type of cache and its index relationships are automatically and synchronously cleaned up, thereby ensuring the overall consistency and stability of the system.
[0118] Figure 5 This is a schematic diagram of a prefix caching device suitable for a hybrid attention model architecture, provided as an embodiment of the present invention. Figure 5 As shown, the device includes: a matching module 310, a verification module 320, a sequence determination module 330, a reuse module 340, and a pre-filled inference module 350.
[0119] in:
[0120] The matching module 310 is used to divide the request sequence received based on the hybrid attention model into blocks and calculate the chain hash value, and perform the longest common prefix matching between the chain hash value and the prefix hash tree constructed by the historical request sequence to obtain the matching cache block. The hybrid attention model includes a linear attention layer and a full attention layer.
[0121] The verification module 320 is used to perform linear attention snapshot verification on the matching cache blocks and determine the verified matching cache blocks as valid cache blocks.
[0122] The sequence determination module 330 is used to determine the first sequence corresponding to the valid cache block and the remaining second sequence in the request sequence based on the valid cache block;
[0123] The reuse module 340 is used to reuse the key-value cache information and linear attention snapshot corresponding to the effective cache block for the first sequence;
[0124] Pre-filled inference module 350 is used to perform pre-filled inference for the second sequence.
[0125] Optionally, the verification module 320 includes:
[0126] The effective snapshot hash set query unit is used to traverse backwards from the last matching cache block to obtain the hash value of the current matching cache block, and to query the local effective snapshot hash set based on the hash value. The local effective snapshot hash set includes the correspondence between hash values and linear attention snapshots, and the linear attention snapshots include convolution state and loop state.
[0127] The verification unit is used to determine that local verification is successful when the hash value exists in the local valid snapshot hash set, and to send a query instruction containing the hash value to the GPU inference process via inter-process communication for remote verification; when the hash value does not exist in the local valid snapshot hash set, the current matching cache block is determined to be an invalid cache block, and the previous matching cache block is traversed and verified. Specifically, when a confirmation instruction is received from the GPU inference process, the remote verification is determined to be successful, and the current matching cache block is determined to be a valid cache block; when no confirmation instruction is received, the current matching cache block is determined to be an invalid cache block.
[0128] Optionally, the device further includes a secondary verification module for performing secondary verification on the linear attention snapshot; when the secondary verification passes, pre-filling inference is performed on the second sequence based on the key-value cache information and the linear attention snapshot; after the pre-filling inference is completed, the request sequence is marked as recovered.
[0129] Optionally, the device further includes a capture module, configured to: immediately capture a linear attention snapshot and key-value cache information at the current sequence position after completing the pre-filling inference of the second sequence; capture a linear attention snapshot and key-value cache information at the current sequence position when the number of newly added sequence tokens reaches the block boundary during subsequent decoding inference; save the key-value cache information in the form of cache blocks to the key-value cache storage area, and construct a corresponding prefix hash tree based on the hash value of each cache block; save the linear attention snapshot in the form of a state tensor to the corresponding slot in the state cache storage area, wherein the linear attention snapshot and key-value cache information corresponding to the same capture position use the same hash value as the index.
[0130] Optionally, the apparatus further includes a mapping relationship establishment module, which is used to determine the set of valid snapshot hash values based on the hash values of cache blocks through the dual cache coordination manager; establish a forward mapping relationship between cache block identifiers and the set of snapshot hash values, and a reverse mapping relationship between snapshot hash values and cache block identifiers.
[0131] Optionally, the device further includes a linkage eviction module, configured to: evict target key-value cache information in the key-value cache storage area when eviction trigger conditions are met, and generate an eviction cache block list based on the cache block identifier corresponding to the eviction key-value cache information, wherein the eviction trigger conditions include prefix cache capacity exceeding the limit, excessive GPU memory pressure, insufficient swap space, or global release; remove the hash value associated with it from the set of valid snapshot hash values based on the eviction cache block list; remove the entry associated with the hash value from the reverse mapping relationship of "snapshot hash value to cache block identifier"; and remove the entry associated with the eviction cache block list from the forward mapping relationship of cache block identifier to snapshot hash value set.
[0132] The prefix caching device for hybrid attention model architecture provided in this embodiment of the invention can execute the method of any of the above embodiments, and has functional modules corresponding to the method and corresponding beneficial effects.
[0133] Figure 6 A schematic diagram of an electronic device 10 that can be used to implement embodiments of the present invention is shown. The electronic device is intended to represent various forms of digital computers, such as laptop computers, desktop computers, workstations, personal digital assistants (PDAs), servers, blade servers, mainframe computers, and other suitable computing devices. The electronic device can also represent various forms of mobile terminal devices. The components shown herein, their connections, and their respective functions are illustrative only and are not intended to limit the implementation of the present invention.
[0134] like Figure 6 As shown, the electronic device 10 includes at least one processor 11 and a memory, such as a read-only memory (ROM) 12 and a random access memory (RAM) 13, communicatively connected to the at least one processor 11. The memory stores computer programs executable by the at least one processor 11. The processor 11 can perform various operations and processes based on the computer program stored in the read-only memory 12 or loaded from storage unit 18 into the random access memory 13. The RAM 13 can also store various programs and data required for the operation of the electronic device 10. The processor 11, ROM 12, and RAM 13 are interconnected via a bus 14, and an input / output (I / O) interface 15 is also connected to the bus 14.
[0135] Multiple components in electronic device 10 are connected to I / O interface 15, including: input unit 16, such as keyboard, mouse, etc.; output unit 17, such as various types of displays, speakers, etc.; storage unit 18, such as disk, optical disk, etc.; and communication unit 19, such as network card, modem, wireless transceiver, etc. Communication unit 19 allows electronic device 10 to exchange information / data with other electronic devices through computer networks such as the Internet and / or various telecommunications networks.
[0136] Processor 11 can be a general-purpose and / or special-purpose processing component with processing and computing capabilities. Examples of processor 11 include, but are not limited to, a central processing unit (CPU), a graphics processing unit (GPU), a dedicated artificial intelligence (AI) computing chip, a processor for running machine learning models, a digital signal processor (DSP), and other suitable processors, controllers, or microcontrollers. Processor 11 is used to perform the methods described in the foregoing embodiments, such as a prefix caching method based on a hybrid attention model.
[0137] Specifically, the operations performed by processor 11 include:
[0138] The request sequence received based on the hybrid attention model is divided into blocks, chained hashes are calculated, and the chained hashes are matched with the longest common prefix of the prefix hash tree of the historical request sequence to obtain matching cache blocks. The hybrid attention model includes a linear attention layer and a full attention layer.
[0139] Perform linear attention snapshot verification on the matching cache blocks, and determine the matching cache blocks that pass the verification as valid cache blocks;
[0140] Based on the valid cache blocks, determine the first sequence in the request sequence that corresponds to the valid cache blocks and the remaining second sequence;
[0141] For the first sequence, reuse the key-value cache information and linear attention snapshots corresponding to the valid cache blocks;
[0142] For the second sequence, perform pre-filled inference.
[0143] Computer programs for implementing the prefix caching method of the present invention applicable to hybrid attention model architectures can be written in one or more programming languages. The computer program can be provided to a processor of a general-purpose computer, a special-purpose computer, or other computing device, such that the processor, when executing the program, performs the functions or operations described in the flowcharts and / or block diagrams. The computer program can be executed entirely on a local device, partially on a local device, partially on a remote device, or entirely on a remote device or server.
[0144] In the context of this invention, a computer-readable storage medium can be a tangible medium for containing or storing a computer program for use by, or in connection with, an instruction execution apparatus, device, or electronic device. Computer-readable storage media can include, but are not limited to, electronic media, magnetic media, optical media, electromagnetic media, infrared media, or semiconductor media, or any combination thereof. Specific examples include, but are not limited to, electrical connections (such as wires), portable storage disks, hard disks, random access memory (RAM), read-only memory (ROM), erasable programmable read-only memory (EPROM or flash memory), optical fibers, compact optical disc read-only memory (CD-ROM), optical storage devices, magnetic storage devices, etc.
[0145] It should be understood that the above process can be adjusted, such as reordering the steps, adding or deleting steps. For example, the steps described in this invention can be executed in parallel, sequentially, or in other orders, as long as the desired effect of the technical solution of this invention can be achieved, they all fall within the protection scope of this invention.
[0146] The above embodiments are only used to illustrate the technical solutions of the present invention, and are not intended to limit the scope of protection of the present invention. Those skilled in the art should understand that various modifications, equivalent substitutions, or improvements can be made to the above embodiments without departing from the spirit and principles of the present invention, and all such modifications, substitutions, or improvements should be included within the scope of protection of the present invention.
Claims
1. A prefix caching method suitable for hybrid attention model architectures, characterized in that, The method includes: The request sequence received based on the hybrid attention model is divided into blocks and a chained hash value is calculated. The chained hash value is then matched with the longest common prefix of the prefix hash tree constructed from the historical request sequence to obtain the matching cache block. The hybrid attention model includes a linear attention layer and a full attention layer. Perform linear attention snapshot verification on the matching cache blocks, and determine the matching cache blocks that pass the verification as valid cache blocks; Based on the valid cache block, determine the first sequence corresponding to the valid cache block and the remaining second sequence in the request sequence; For the first sequence, the key-value cache information and linear attention snapshot corresponding to the effective cache block are reused; For the second sequence, perform pre-filled inference.
2. The method according to claim 1, characterized in that, The linear attention snapshot verification of the matching cache block includes: Starting from the last matching cache block, it iterates backwards to obtain the hash value of the current matching cache block, and queries the local valid snapshot hash set based on the hash value. The local valid snapshot hash set includes the correspondence between hash values and linear attention snapshots, and the linear attention snapshot includes convolution state and loop state. When the hash value exists in the local valid snapshot hash set, local verification is determined to be successful, and a query command containing the hash value is sent to the GPU inference process via inter-process communication for remote verification. When the hash value does not exist in the local valid snapshot hash set, the current matching cache block is determined to be an invalid cache block, and the previous matching cache block is traversed and verified.
3. The method according to claim 2, characterized in that, The remote verification process includes: When a confirmation instruction is received from the GPU inference process, the remote verification is confirmed to be successful, and the currently matching cache block is determined to be a valid cache block. If the confirmation instruction is not received, the currently matching cache block will be determined as an invalid cache block.
4. The method according to claim 1, characterized in that, Before performing pre-padded inference on the second sequence, the following is also included: The linear attention snapshot is then validated a second time. When the secondary verification passes, pre-filled inference is performed on the second sequence based on the key-value cache information and the linear attention snapshot; After the pre-filled inference is completed, the request sequence is marked as recovered.
5. The method according to claim 1, characterized in that, The method also includes: After completing the pre-filled inference for the second sequence, immediately capture a linear attention snapshot and key-value cache information at the current sequence position; During subsequent decoding and inference, when the number of newly added sequence word elements reaches the block boundary, a linear attention snapshot and key-value cache information are captured at the current sequence position; The key-value cache information is saved to the key-value cache storage area in the form of cache blocks, and a corresponding prefix hash tree is constructed based on the hash value of each cache block; The linear attention snapshot is saved as a state tensor to the corresponding slot in the state cache storage area, wherein the linear attention snapshot and key-value cache information corresponding to the same capture position use the hash value of the same cache block as the index.
6. The method according to claim 5, characterized in that, After saving the linear attention snapshot as a state tensor to the corresponding slot in the state cache storage area, the method further includes: The set of valid snapshot hash values is determined by the dual-caching coordinator based on the hash values of cache blocks; Establish a forward mapping between cache block identifiers and snapshot hash value sets, and a reverse mapping between snapshot hash values and cache block identifiers.
7. The method according to claim 6, characterized in that, After performing pre-filled inference, it also includes: When the eviction triggering conditions are met, the target key-value cache information in the key-value cache storage area is evicted, and an eviction cache block list is generated according to the cache block identifier corresponding to the evicted key-value cache information. The eviction triggering conditions include prefix cache capacity exceeding the limit, GPU memory pressure, insufficient swap space, or global release. Based on the list of evicted cache blocks, remove the hash values associated with the list of evicted cache blocks from the set of valid snapshot hash values; Remove the entry associated with the snapshot hash from the reverse mapping from the snapshot hash to the cache block identifier; Remove entries associated with the list of evicted cache blocks from the positive mapping from the cache block identifier to the set of snapshot hashes.
8. A prefix caching device suitable for hybrid attention model architectures, characterized in that, The device includes: The matching module is used to divide the request sequence received based on the hybrid attention model into blocks and calculate the chained hash value, and perform the longest common prefix matching between the chained hash value and the prefix hash tree constructed from the historical request sequence to obtain the matching cache block. The hybrid attention model includes a linear attention layer and a full attention layer. The verification module is used to perform linear attention snapshot verification on the matching cache block and determine the matching cache block that passes the verification as a valid cache block; The sequence determination module is used to determine, based on the valid cache block, a first sequence corresponding to the valid cache block and the remaining second sequence in the request sequence; The reuse module is used to reuse the key-value cache information and linear attention snapshot corresponding to the effective cache block for the first sequence; A pre-filled inference module is used to perform pre-filled inference for the second sequence.
9. An electronic device, characterized in that, The electronic device includes: One or more processors; and a storage device for storing one or more programs. When the one or more programs are executed by the one or more processors, the one or more processors cause the one or more processors to implement the method as described in any one of claims 1-7.
10. A computer-readable storage medium having a computer program stored thereon, characterized in that, The computer program, when executed by a processor, implements the method as described in any one of claims 1-7.