Key-value cache compression method and system for large language model reasoning

Abstract

The application discloses a key-value cache compression method and system for large language model reasoning, and belongs to the technical field of cache optimization, comprising: in the pre-filling stage of large language model reasoning, performing: using a semantic scoring model to divide input text into word units and predict the semantic scores of each word unit, obtaining a word unit set and the semantic scores of each word unit in the word unit set; mapping the semantic scores of each word unit to the word unit set obtained by the large language model dividing the input text into word units; according to the pruning ratio, removing the word unit features (key matrix and value matrix) of part of the word units with the lowest semantic scores from the KV cache in the video memory; translating the word unit features forward, so that the word unit features are stored continuously in the KV cache; and updating the word unit features that have been translated using the translated positions. The application can effectively balance the calculation efficiency and cache hit rate in the large language model reasoning process under the condition that the video memory capacity is limited and the cross-layer I / O communication overhead is high.

Description

Technical Field

[0001] This invention belongs to the field of cache optimization technology, and more specifically, relates to a key-value cache compression method and system for reasoning in large language models. Background Technology

[0002] With the continuous evolution of deep learning and artificial intelligence technologies, generative large language models have demonstrated excellent performance in various natural language processing tasks. In recent years, application scenarios have gradually evolved from single-round question answering and text translation to highly autonomous agent workflows. In agent applications, models typically adopt an "observation-reasoning-action" execution paradigm, requiring frequent invocation of external tools, reading long documents, and accumulating multi-round contextual memory.

[0003] In the underlying mechanism of large language models, the inference process is divided into a pre-filling stage and a decoding stage. In the pre-filling stage, the model inputs all tokens of the input sequence in parallel, calculates the key and value matrices for each token (i.e., the initial establishment of the KV cache), and ensures autoregressive properties through causal attention masks, ultimately generating the first output token. The decoding stage is the recursive stage of LLM inference, used to generate subsequent output tokens one by one. In this stage, the model only inputs the previously generated token, calculates the attention of the current token using the cached KV cache, generates the next token, and appends the KV of the current token to the cache.

[0004] To avoid redundant computations of historically calculated context, modern inference systems commonly employ prefix caching. Due to the nature of the causal attention mechanism in the Transformer architecture, identical prefix inputs will inevitably generate identical key-value caches (KV caches). Therefore, storing historically generated key-value caches in storage and reusing them in subsequent requests with the same prefix can significantly reduce the initial token response time.

[0005] However, the extremely long context accumulation under agent workloads presents an unprecedented "memory bloat" challenge to existing prefix caching mechanisms. This memory bloat severely constrains existing inference architectures in the following dimensions: First, the severely limited GPU memory capacity leads to a precipitous drop in cache hit rate. Modern GPU memory is expensive and has limited capacity. When the key-value cache expands rapidly due to agent tasks, the working set quickly exceeds the capacity limit of HBM. This forces the system to frequently evict hot data, causing subsequent requests with the same prefix to miss the cache, thus forcing the model to recalculate massive amounts of historical context during the pre-filling phase, resulting in a significant waste of computing power.

[0006] Second, the cross-tier I / O communication overhead of tiered storage becomes a performance bottleneck. To overcome the capacity bottleneck of a single-card HBM, the industry has attempted to introduce a heterogeneous tiered storage architecture of VRAM-RAM-SSD. However, the migration of terabytes of prefix cache data between storage tiers (especially via the PCIe bus) causes extremely high I / O latency and bandwidth congestion. In many cases, this data movement overhead even offsets the benefits of avoiding recomputation.

[0007] Overall, given the limited video memory capacity and high overhead of cross-layer I / O communication, effectively balancing computational efficiency and cache hit rate in the inference process of large language models is a technical problem that urgently needs to be solved. Summary of the Invention

[0008] To address the shortcomings and improvement needs of existing technologies, this invention provides a key-value cache compression method and system for large language model inference. Its purpose is to effectively balance computational efficiency and cache hit rate in the large language model inference process, under the conditions of limited video memory capacity and high cross-layer I / O communication overhead.

[0009] To achieve the above objectives, according to one aspect of the present invention, a key-value cache compression method for large language model inference is provided, comprising the following steps performed during the pre-filling stage of large language model inference: S1: Use a semantic scoring model to segment the input text of the current round into lexical units and predict the semantic score of each lexical unit to obtain a lexical set. and the semantic scores of each word element therein; S2: Set of words The semantic scores of each word are mapped to the word set obtained by the large language model from word segmentation of the input text. Above, we obtain the word set. Semantic scores for each word element in the text; S3: Prune the word set according to the preset pruning ratio. The lexical features of the lexical units with the lowest semantic scores are removed from the KV cache in the video memory; the lexical features include the key matrix and value matrix of the lexical unit; S4: Shift the word features after removing them forward so that the remaining word features are stored in a contiguous space in the KV cache; update the shifted word features using the updated positions; Among them, the semantic scoring model is smaller in scale than the large language model.

[0010] Furthermore, the lexical set The semantic scores of each word are mapped to the word set obtained by the large language model from word segmentation of the input text in the current round. Above, including: Record lexicon set and word set The offset range of each word in the input text; With lexical set and word set Using the lexical units as vertices, construct an undirected graph; if the lexical units offset interval and word If the offset intervals overlap, then in the word units... and word elements Add an edge between them; Breadth-first search is used to traverse the undirected graph to extract mutually independent connected components; The semantic score is mapped within the connected components according to the following rules: If within the same connected component, the word set One word element with lexical set If a word has a one-to-one relationship with another word, then the word will be... The semantic score is assigned to the corresponding word unit; if the word unit set One word element with lexical set If multiple lexical units in a word have a one-to-many relationship, then the lexical units will be... The semantic score is assigned to each corresponding lexical unit; if the lexical set Multiple lexical units and lexical sets in If a word element has a many-to-one relationship, then the word element set... The maximum semantic score of these multiple lexical units is assigned to the corresponding lexical unit; if the lexical set Multiple lexical units and lexical sets in If multiple lexical units in a word have many-to-many relationships, then the lexical set... The maximum semantic score of these multiple lexical units is assigned to the corresponding lexical unit.

[0011] Furthermore, the semantic scoring model is the BERT model.

[0012] Furthermore, when updating the translated lexical features using the updated positions, the following formula is used: ; in, and These represent the position indices before and after the translation, respectively. Indicates the lexical features before translation. Indicates the lexical features after translation. Indicates rotation parameters, This represents the rotation matrix that encodes the rotation position.

[0013] Furthermore, the key-value caching compression method for large language model inference provided by this invention also includes: Maintain the prefix tree corresponding to the KV cache; For nodes in the prefix tree Generate its child nodes hour: If node For nodes The first child node, and used for storing nodes. The last physical page of the corresponding lexical feature If the page boundary is not aligned, then the node will be... Lexical features are appended to the physical page. For the unwritten lexical feature parts, allocate free physical pages for writing, and simultaneously, at the node... Inject split markers into the metadata to indicate physical pages Some of the data stored in the middle belongs to the node. child nodes; If node For nodes The first child node, and used for storing nodes. The last physical page of the corresponding lexical feature If page boundaries are aligned, then free physical pages are allocated and written to the node. Lexical features; If node Not a node The first child node, and the node If a split marker exists in the metadata, then it is a node. The lexical features are assigned to physical pages, and the physical pages are... After the data in the node is copied to the starting position in the allocated physical page, the node will be... Lexical features are appended to the text; If node Not a node The first child node, and the node If no split marker is found in the metadata, a free physical page is allocated and written to the node. Lexical features.

[0014] Furthermore, when searching the prefix tree to perform prefix path matching, if the currently searched node If a split marker exists in the metadata, it will be used to store the node. The last physical page of the corresponding lexical feature After popping from the currently matched prefix path, perform subsequent matching.

[0015] According to another aspect of the present invention, a key-value cache compression system for large language model inference is provided, comprising: a three-level heterogeneous storage system and a controller; The controller includes a compression module; the compression module is used to execute the key-value cache compression method for large language model inference provided by the present invention, so as to compress lexical features in the pre-filling stage of large language model inference; The three-tier heterogeneous storage system includes: video memory, RAM, and solid-state drives (SSDs). The video memory is used to perform large language model inference tasks. The video memory, RAM, and SSDs are all equipped with key-value (KV) caches to cache compressed blocks of word features after compression.

[0016] Furthermore, the controller also includes: a management module and an execution module; The management module is used to manage the storage space allocation at each level of the three-tier heterogeneous storage system and the indexing of the KV cache; The execution module is used to handle read and write operations on memory and solid-state drives, as well as data transfer operations between memory and video memory.

[0017] Furthermore, the management module is also used to perform asynchronous lifecycle scheduling for the three-tier heterogeneous storage system; asynchronous lifecycle scheduling includes: Loading: When the inference task enters the scheduling pool and the model forward propagation has not yet started, check the prefix hit. If the prefix hits but the GPU memory does not hit, allocate space in the GPU memory and send an instruction to the execution module so that the execution module loads the corresponding compressed block from memory into the GPU memory, or loads the corresponding compressed block from the solid-state drive into memory and then loads the compressed block into the GPU memory. Prefetch: When the total size of the accumulated prompt words reaches a preset safety threshold since the last prefetch, an instruction is sent to the execution module to prefetch the compressed block that exists only on the solid-state drive and not in memory into memory; Storage: During the decoding stage of large language model inference, each time the decoding generates new lexical units that cause the current physical page in the KV cache to be filled, space is allocated in memory and solid-state drive, and instructions are sent to the execution module so that the execution module can move the filled physical page from the KV cache to memory and solid-state drive in sequence to achieve persistent data storage. Evicting: When the KV cache space of any level in the three-level heterogeneous storage system reaches the corresponding exhaustion index, compressed blocks are selected from the KV cache space for eviction according to the preset eviction policy.

[0018] Furthermore, the communication channel between the management module and the execution module is a hybrid inter-process communication channel based on ZeroMQ and shared memory.

[0019] In summary, the above-described technical solutions conceived in this invention can achieve the following beneficial effects: (1) This invention uses a lightweight semantic scoring model to perform word segmentation on the input text and performs semantic scoring on each word segmentation. By mapping the semantic scores of words to each word segmented by the large language model, and based on the mapped word semantic scores, the word features of the words with the lowest scores are removed from the KV cache. Thus, the semantic and syntactic importance of each word in the current round input context can be evaluated before the query Q arrives, and KV cache compression is achieved. Therefore, this invention implements a query-independent key-value cache compression method. Based on this query-independent key-value cache compression method, this invention can actively reduce the size of the key-value cache without affecting the model inference accuracy, so that the limited GPU memory can accommodate a multiplied increase in the request sequence, thereby significantly improving the global cache hit rate and reducing cross-layer I / O transmission. The KV cache effectively improves the computational efficiency of large model inference tasks.

[0020] (2) In a preferred embodiment of the present invention, an alignment algorithm based on a bipartite graph is proposed to map the semantic scores of the word units obtained by the semantic scoring model to the word units obtained by the large language model. Specifically, word units in two word unit sets are used as nodes. Edges between nodes are constructed according to the overlap of the offset intervals of word units in different sets. After extracting mutually independent connected components, semantic scores are transferred within the connected components. For 1-to-1 and 1-to-N, the scores of the scoring model are directly assigned to the corresponding word units of the large language model. For N-to-1 and N-to-M, the maximum value of the scores of the associated word units in the scoring model is extracted and uniformly assigned to all word units of the large model within the component. Thus, the mapping of word unit semantic scores can be completed without omitting important semantic information.

[0021] (3) In the preferred embodiment of the present invention, when the node of the prefix tree corresponding to the KV cache splits at a point where the page boundary alignment is not satisfied, a split marker is injected into the metadata of the parent node to indicate that part of the data stored in the last physical page used to store the lexical features corresponding to the parent node belongs to its first child node. For other child nodes, the parent node is not modified in place, but a completely new free physical page is requested, and the valid common lexicals in the original shared page are copied to the starting position of the new page. Then, the lexicals of the new branch are safely appended to the new page. This realizes a copy-on-write (CoW) and tail redundancy strategy. Through the copy-on-write and tail redundancy strategies, the present invention effectively solves the problem of "misalignment between logical sequence and physical page boundary" caused by fine-grained lexical compression. It can effectively prevent concurrent write conflicts in the multi-branch search and reuse common in intelligent agent scenarios, and ensure the data security and system robustness of the underlying video memory pages.

[0022] (4) The present invention splits the three-level storage scheduling into a central management module responsible for overall planning and a distributed execution module responsible for transporting, so that the uplink and downlink operations of the large key-value cache between video memory, memory and solid-state drive are completely asynchronous and separated from the inference main path, thereby reducing the access latency of the key-value cache. Attached Figure Description

[0023] Figure 1 This is a schematic diagram of a bipartite graph scoring alignment method for query-irrelevant cache compression provided in an embodiment of the present invention.

[0024] Figure 2 This is a schematic diagram of rotational position encoding remapping and memory compaction provided in an embodiment of the present invention.

[0025] Figure 3 A flowchart of a key-value caching compression method for large language model inference provided in an embodiment of the present invention.

[0026] Figure 4 The flowchart for the execution of a large language model inference task provided in this embodiment of the invention.

[0027] Figure 5 A flowchart of a key-value caching compression system for large language model inference provided in an embodiment of the present invention. Detailed Implementation

[0028] To make the objectives, technical solutions, and advantages of this invention clearer, the invention will be further described in detail below with reference to the accompanying drawings and embodiments. It should be understood that the specific embodiments described herein are merely illustrative and not intended to limit the invention. Furthermore, the technical features involved in the various embodiments of this invention described below can be combined with each other as long as they do not conflict with each other.

[0029] In this invention, the terms "first," "second," etc. (if present) in the invention and the accompanying drawings are used to distinguish similar objects and are not necessarily used to describe a specific order or sequence.

[0030] Before discussing the exemplary embodiments in more detail, it should be mentioned that the core of large language models is the multi-head self-attention mechanism. For a given input text, its latent state is projected to generate query (Q), key (K), and value (V) matrices. The K and V matrices generated in the pre-filling step during the inference phase are cached (i.e., key-value cache) for direct use in subsequent decoding steps. The extremely long context accumulation under the agent's workload generates huge key and value matrices. Due to the severely limited GPU memory capacity and the huge communication overhead caused by cross-layer I / O communication, how to effectively balance the computational efficiency and cache hit rate in the inference process of large language models is a technical problem that urgently needs to be solved. To solve this technical problem, this invention provides a key-value cache compression method and system for large language model inference. The overall concept is to achieve KV cache compression by pruning the volume of the key and value matrices, thereby storing more effective data in a limited storage space, improving the global hit rate, and reducing cross-layer I / O overhead.

[0031] Based on the above technical concept, this invention considers that in large language model inference, if the bloated original key-value cache is allowed to flood into the storage system, it will lead to exhaustion of video memory and cross-layer I / O bandwidth. Therefore, this invention proposes a query-independent compression method for asynchronous offline compression of the key-value cache.

[0032] Most existing lexical filtering algorithms rely on query Q (i.e., query awareness), which requires the input to be available. However, this necessitates reloading the entire cache back into memory before compression, failing to fundamentally save storage. Therefore, this invention employs a lightweight semantic scoring model that runs in a bypass mode. Without relying on specific queries, it pre-evaluates the semantic and syntactic importance of each lexical in the previous round's input and output context after the previous round's request is completed, and saves the compressed state for future use.

[0033] The additional semantic scoring model suffers from a "segmenter mismatch" problem with the original large language model, meaning that the word set obtained by the two models is inconsistent. For example, when using BERT as the semantic scoring model, the BERT model is based on WordPiece, while mainstream generative large language models use byte-pair encoding. The word segmentation results for the same input text are quite different. Therefore, semantic scoring cannot be mapped one-to-one between the two word sets.

[0034] Therefore, this invention proposes a semantic alignment algorithm based on bipartite graphs, which mainly includes the following steps: (1) Obtain the word set of the scoring model and large model lexical set Record the start and end offset range of each word in the original text string.

[0035] (2) Construct an undirected graph with vertices as follows: and .like offset range and If the offset intervals of the two elements intersect, then an edge is added between them.

[0036] (3) Use breadth-first search to traverse the graph and extract mutually independent connected components. Each connected component represents a shared text cluster across the word segmenters.

[0037] (4) Semantic scoring is passed within connected components, and the passing rules are as follows: Figure 1 As shown: For 1-to-1 and 1-to-N, the score of the scoring model is directly assigned to the corresponding large language model lexical unit; for N-to-1 and N-to-M, in order to ensure that important semantics are not missed, the maximum value of the semantic score of the associated lexical unit of the scoring model is extracted and uniformly assigned to all large model lexical units within the connected component.

[0038] After the semantic scoring mapping is completed, according to the preset pruning ratio, the lexical features (key matrix and value matrix) of the lexical units with the lowest semantic scores are removed from the cache queue of the video memory. This allows the size of the key-value cache to be actively reduced without affecting the model's inference accuracy, thus achieving key-value cache compression. This enables the limited video memory and system memory to accommodate a multiplied number of request sequences, thereby significantly improving the global cache hit rate.

[0039] In the query-independent compression method proposed in this embodiment, some lexical features are forcibly removed, resulting in "honeycomb-like" memory holes in the KV cache and causing the attention mechanism to crash. To address this issue, this invention further proposes methods for maintaining data continuity from both physical and logical levels, such as... Figure 2 As shown.

[0040] Modern large-scale model inference frameworks widely use PagedAttention, which requires key-value caches to be stored contiguously, such as with 16 terms per physical page. Except for the last page of the current sequence, which can be partially filled, all other pages must be full. To maintain data continuity at the physical level, this invention scans the discarded term feature sequence and compactly assembles the discrete valid terms in sequence, recombining fragments that originally spanned several physical pages into a few contiguous physical pages, thereby achieving memory compaction. Memory compaction completely eliminates memory fragmentation, significantly reducing GPU memory usage and cross-layer I / O communication volume. Optionally, in this embodiment, terms are scanned by fusing CUDA operators.

[0041] Research has found that the early attention layers of large models are highly dependent on the continuity of absolute positions. Ordinary compression can cause gaps, significantly reducing the model's accuracy. Large language models are based on rotational positional encoding, where the encoding results are closely related to the storage location.

[0042] Based on the mathematical properties of vector rotation in rotational position encoding compared to relative position encoding, for the original position... Features That is, the key matrix and value matrix before encoding, whose encoded representation is as follows: It can be reversibly restored by multiplying by the inverse rotation matrix, where, Indicates rotation parameters, This represents the rotation matrix.

[0043] In this invention, the memory compaction operation ensures that the remaining lexical features are stored contiguously in physical pages, while also changing the position of the lexical features. The original position after memory compaction and reorganization... The valid tokens were shifted to the new continuous physical index. At this point, the corresponding lexical features are updated according to the following formula: ; in, and These represent the position indices before and after the translation, respectively. Indicates the lexical features before translation. This indicates the lexical features after translation.

[0044] Based on the above formula, this invention achieves rotational positional encoding remapping. This operation forcibly eliminates the positional encoding jumps caused by the original elimination, disguising the new sequence as a logically completely continuous sequence, thereby maintaining the continuity of data at the logical level. This eliminates the confusion caused by positional discontinuities, achieving inference accuracy comparable to the full attention model in large language models.

[0045] Based on the above query-independent compression method combined with memory compaction operations and rotation position encoding remapping, lexical-level key-value cache compression is achieved. However, this lexical-level key-value cache compression implementation presents an irreconcilable contradiction in multi-branch search and reuse common in agent scenarios: the underlying storage and retrieval system is page-level. When the agent performs multiple rounds of divergent thinking, different rounds will share the same prefix sequence (i.e., the branch points of the prefix tree).

[0046] For example, after a shared prefix sequence is compressed, its end may not fill the last physical page. At this point, branches A and B each generate new terms, and both branches will attempt to append to the remaining space of the physical page. This lack of page alignment will directly lead to multi-threaded write conflicts and data pollution. To address this problem, this invention further proposes a sparse radix tree indexing mechanism, mainly including: (1) Split mark injection: When a tree node splits at a point where page boundary alignment is not satisfied, the system injects a split mark (last_page_padding) into the parent node's metadata, indicating how many slots in the last physical page contain data but logically belong to the separated child nodes. For the first branch that is joined, the data of the first child node of the corresponding split node will be appended to the last physical page.

[0047] (2) Copy-on-Write (CoW): For a newly accessed branch sequence, when it is detected that a write to an unaligned physical page is necessary, if the corresponding parent node has been injected with a split marker, no in-place modification is performed. Instead, after freeing up the physical page, the valid common terms in the original shared page are copied to the beginning of the new page, and then the terms of the new branch are safely appended to the new page. This achieves tail redundancy and the copy-on-write strategy.

[0048] By employing tail redundancy and copy-on-write strategies, this invention cleverly avoids data overwriting in the memory management framework while fully preserving the extremely high efficiency of token compression, thus ensuring the absolute security of the system in high-concurrency scenarios.

[0049] When executing the prefix path matching algorithm, if a split marker (last_page_padding is non-zero) is encountered, the system performs a tail pop operation, actively ignoring the last ambiguous physical page.

[0050] The following is an example.

[0051] Example 1: A key-value caching compression method for reasoning in large language models, such as Figure 3 As shown, the following steps are performed during the pre-filling stage of large language model inference: S1: Use a semantic scoring model to segment the input text of the current round into lexical units and predict the semantic score of each lexical unit to obtain a lexical set. and the semantic scores of each word element therein; S2: Set of words The semantic scores of each word are mapped to the word set obtained by the large language model from word segmentation of the input text. Above, we obtain the word set. Semantic scores for each word element in the text; S3: Prune the word set according to the preset pruning ratio. The lexical features of the lexical units with the lowest semantic scores are removed from the KV cache in the video memory; the lexical features include the key matrix and value matrix of the lexical unit; S4: Shift the word features after removing them forward so that the remaining word features are stored in a contiguous space in the KV cache; update the shifted word features using the updated positions; Among them, the semantic scoring model is smaller in scale than the large language model.

[0052] Optionally, in this embodiment, the lightweight semantic evaluation model selected is the BERT (Bidirectional Encoder Representations from Transformers) model. The BERT model is a Transformer architecture that only contains an encoder. Its core advantage lies in its bidirectional attention mechanism, which can effectively capture and extract deep semantic information. At the same time, its parameter count is only about 500M. Due to the bidirectional global perception capability of the BERT model, this invention can complete semantic scoring in parallel with extremely high throughput (processing more than 33k words per second) without relying on specific queries.

[0053] In this embodiment, the lexical set The semantic scores of each word are mapped to the word set obtained by the large language model from word segmentation of the input text in the current round. Above, including: Record lexicon set and word set The offset range of each word in the input text; With lexical set and word set Using the lexical units as vertices, construct an undirected graph; if the lexical units offset interval and word If the offset intervals overlap, then in the word units... and word elements Add an edge between them; Breadth-first search is used to traverse the undirected graph to extract mutually independent connected components; The semantic score is mapped within the connected components according to the following rules: If within the same connected component, the word set One word element with lexical set If a word has a one-to-one relationship with another word, then the word will be... The semantic score is assigned to the corresponding word unit; if the word unit set One word element with lexical set If multiple lexical units in a word have a one-to-many relationship, then the lexical units will be... The semantic score is assigned to each corresponding lexical unit; if the lexical set Multiple lexical units and lexical sets in If a word element has a many-to-one relationship, then the word element set... The maximum semantic score of these multiple lexical units is assigned to the corresponding lexical unit; if the lexical set Multiple lexical units and lexical sets in If multiple lexical units in a word have many-to-many relationships, then the lexical set... The maximum semantic score of these multiple lexical units is assigned to the corresponding lexical unit.

[0054] This embodiment further includes: Maintain the prefix tree corresponding to the KV cache; For nodes in the prefix tree Generate its child nodes hour: If node For nodes The first child node, and used for storing nodes. The last physical page of the corresponding lexical feature If the page boundary is not aligned, then the node will be... Lexical features are appended to the physical page. For the unwritten lexical feature parts, allocate free physical pages for writing, and simultaneously, at the node... Inject split markers into the metadata to indicate physical pages Some of the data stored in the middle belongs to the node. child nodes; If node For nodes The first child node, and used for storing nodes. The last physical page of the corresponding lexical feature If page boundaries are aligned, then free physical pages are allocated and written to the node. Lexical features; If node Not a node The first child node, and the node If a split marker exists in the metadata, then it is a node. The lexical features are assigned to physical pages, and the physical pages are... After the data in the node is copied to the starting position in the allocated physical page, the node will be... Lexical features are appended to the text; If node Not a node The first child node, and the node If no split marker is found in the metadata, a free physical page is allocated and written to the node. Lexical features.

[0055] In this embodiment, when searching the prefix tree to perform prefix path matching, if the currently searched node If a split marker exists in the metadata, it will be used to store the node. The last physical page of the corresponding lexical feature After popping from the currently matched prefix path, perform subsequent matching.

[0056] Based on the key-value caching compression method for large language model inference provided in this embodiment, the complete execution process of a large language model inference task is as follows: Figure 4 As shown.

[0057] In summary, this embodiment achieves word-level key-value cache compression. Under limited constraints, without affecting model inference accuracy, it proactively reduces the size of the key-value cache, allowing limited GPU memory to accommodate a significantly increased number of request sequences. This significantly improves the global cache hit rate and reduces cross-layer I / O transmission, effectively enhancing the computational efficiency of large model inference tasks. Furthermore, addressing the "physical fragmentation" and "logical position discontinuity" issues caused by word-level compression, this embodiment's memory compaction and rotation position encoding remapping operations are adapted to the memory allocation rules of the underlying PagedAttention framework. It also reshapes the continuous relative positions of discrete words, eliminating the accuracy degradation caused by positional discontinuities in the early attention layers. In addition, this embodiment employs "tail redundancy" and "copy-on-write" strategies to solve the "misalignment between logical sequences and physical page boundaries" problem caused by fine-grained word compression. In multi-branch search and reuse common in agent scenarios, this effectively prevents concurrent write conflicts, ensuring data security and system robustness of the underlying GPU memory pages.

[0058] Example 2: A key-value caching compression system for large language model inference, such as Figure 5As shown.

[0059] This embodiment aims to apply the key-value cache compression method for large language model inference provided in Embodiment 1 above. After completing the compression and index construction of the prefix cache using the key-value cache compression method for large language model inference provided in Embodiment 1, the traditional single-level memory allocation mode of the LLM inference framework can no longer meet the needs of high-concurrency applications in order to efficiently manage and schedule these data. Therefore, as... Figure 5 As shown in this embodiment: a three-tier heterogeneous storage system and a controller; The controller includes a compression module; the compression module is used to execute the key-value cache compression method for large language model inference provided by the present invention, so as to compress lexical features in the pre-filling stage of large language model inference; The three-tier heterogeneous storage system includes: video memory, RAM (DRAM), and solid-state drives (SSDs). The video memory is used to perform large language model inference tasks. The video memory, RAM, and SSDs are all equipped with key-value (KV) caches to cache compressed blocks of word features after compression.

[0060] In order to decouple the entire key-value caching compression system from the physical and logical levels, this embodiment adopts the design concept of "management-execution separation" at the software module level, in addition to the three-level heterogeneous storage system. Accordingly, the controller also includes a management module and an execution module.

[0061] The Management Module, an extension of the core scheduler, resides in the main process. It is responsible for request interception, batch processing, and the formulation of uplink and downlink routing schemes for the key-value cache. It also manages the storage space allocation for each level of the three-tier heterogeneous storage system and the indexing of the KV cache. Internally, the Management Module maintains the tree structure using a prefix block prefix tree and a hash algorithm. The algorithm is expressed as follows: for non-root nodes, any prefix sequence has a unique index ID globally.

[0062] The Execution Module creates an independent Worker process for each compute card (GPU / NPU). This Worker process lacks a global scheduling view and is solely responsible for executing commands issued by the management module. Specifically, the Execution Module handles read / write operations on memory and SSDs, as well as data transfer operations between memory and video memory. In this embodiment, the Execution Module is configured with a high-performance asynchronous I / O framework. For example, it includes an io_uring instance to handle large-scale concurrent read / write operations between memory and SSDs, and several independent threads that use cudaMemcpyAsync or aclrtMemcpyAsync to handle data transfer between DRAM and video memory.

[0063] To ensure real-time performance in high-concurrency scenarios with tens of thousands of requests per second, a hybrid inter-process communication channel based on ZeroMQ and shared memory was established between the management module and the execution module. Once the management module has formulated the data transfer plan, it sends a fixed-length header of only a few dozen bytes via the REQ-REP socket in ZeroMQ, while thousands of Block IDs and memory offset information are directly written to shared memory. The execution module, after obtaining the task list with zero-copy processing, immediately returns a response and asynchronously hands it over to a background lock-free queue for the actual data copying. When the copying is complete, the execution module reports a completion signal to the management module's PULL socket via the PUSH socket, and the management module updates the task counter. This mechanism ensures virtually zero latency in network transmission.

[0064] In this embodiment, the management module is also used to perform asynchronous lifecycle scheduling of the three-tier heterogeneous storage system; asynchronous lifecycle scheduling includes: Loading: When an inference task enters the scheduling pool but before model forward propagation begins, the management module calls an interface to check for prefix hits. For compressed blocks that are hit but not in GPU memory, the management module first allocates storage space for them in GPU memory. At this time, the uplink task counter is incremented. Subsequently, instructions are sent to the execution module via shared memory. The lock-free queue scheduling CUDA / CANN asynchronous DMA engine in the execution module performs the actual fetching, and the block can only be unblocked and the next computation can proceed after all associated batches of compressed blocks have arrived in GPU memory (uplink task counter reaches zero).

[0065] Prefetching: To overcome the limitation of extremely slow SSD bandwidth on first character generation time, the management module actively polls the background request queue. The system sets a safety threshold (e.g., the total size of accumulated prompt words). (Total memory storage capacity - Total GPU memory storage capacity). When the total number of prompts since the last prefetch reaches this safety threshold, a prefetch task is issued for compressed blocks that exist only on the SSD and not in memory. Using the asynchronous io_uring mechanism, cold data on the SSD is silently promoted to hot data in memory. This operation does not consume any critical inference resources (GPU resources).

[0066] Storage: Each time a new token is generated during decoding, causing the physical page to fill up, the inference framework notifies the management module of the full physical page, initiating a two-step confirmation mechanism: First, during the pre-registration phase, the management module allocates storage space in memory and SSD and temporarily stores the information; second, in the distribution and protection phase, before initiating the downlink write, the management module forcibly increments the reference count of the compressed block in the video memory. The execution module then initiates a cascading downlink transfer from HBM to DRAM to SSD until the data is fully persisted to disk. After this, the management module decrements the reference count via an IPC callback, ensuring that the underlying framework does not mistakenly release video memory during this period.

[0067] Eviction: When key-value caches at each level face space exhaustion, the system performs eviction accordingly. To prevent tree-structured index topology breakage, eviction for solid-state drives (SSDs) relies on a specific leaf node eviction list. This list is based on an LRU policy and only includes leaf node blocks with zero upstream and downstream task counts and zero reference counts. For memory caches, since they inherently possess a safe copy from the SSD, the system maintains an independent memory eviction list. This allows for the free eviction of any block in memory based on the LRU policy, without considering leaf node limitations, thus maximizing the efficiency of limited memory space utilization.

[0068] In summary, this embodiment compresses the originally extremely large and redundant key-value cache both physically and logically using an asynchronous and query-independent algorithm. The compressed high-value cache is then asynchronously stored, retrieved, and scheduled by an architecture-decoupled hierarchical storage engine. Compared to existing approximation algorithms, quantization, or online sparsity strategies, this solution fundamentally reshapes the caching ecosystem disaster caused by the accumulation of agent workflows. It effectively solves the memory utilization problem in complex long-text inference and the I / O communication bottlenecks in PCIe / NVMe storage tier interconnection, significantly improving the computational efficiency of large language model inference tasks.

[0069] Those skilled in the art will readily understand that the above description is merely a preferred embodiment of the present invention and is not intended to limit the present invention. Any modifications, equivalent substitutions, and improvements made within the spirit and principles of the present invention should be included within the scope of protection of the present invention.

Claims

1. A key-value caching compression method for reasoning in large language models, characterized in that, This includes the following steps performed during the pre-filling stage of large language model inference: S1: Use a semantic scoring model to segment the input text of the current round into lexical units and predict the semantic score of each lexical unit to obtain a lexical set. and the semantic scores of each word element therein; S2: The word set The semantic score of each word in the input text is mapped to the word set obtained by the large language model through word segmentation. The above yields the set of lexical terms. Semantic scores for each word element in the text; S3: Prune the word set according to the preset pruning ratio. The lexical features of the lexical units with the lowest semantic scores are removed from the KV cache in the video memory; Lexical features include the key matrix and value matrix of the lexical unit; S4: Shift the lexical features after removing them forward so that the remaining lexical features are stored in a contiguous space in the KV cache; update the shifted lexical features using the updated positions; The semantic scoring model is smaller in scale than the large language model.

2. The key-value caching compression method for large language model inference as described in claim 1, characterized in that, The word set The semantic score of each word is mapped to the word set obtained by the large language model from word segmentation of the input text in the current round. Above, including: Record the set of lexical terms and the set of terms The offset interval of each word in the input text; With the aforementioned lexical set and the set of terms Using the lexical units as vertices, construct an undirected graph; if the lexical units offset interval and word If the offset intervals overlap, then in the word units... and word elements Add an edge between them; The undirected graph is traversed using a breadth-first search to extract mutually independent connected components; The semantic score is mapped within the connected components according to the following rules: If within the same connected component, the term set One word element With the set of lexical terms If a word has a one-to-one relationship with another word, then the word will be... The semantic score is assigned to the corresponding word unit; if the word unit set One word element With the set of lexical terms If multiple lexical units in a word have a one-to-many relationship, then the lexical units will be... The semantic score is assigned to each corresponding lexical unit; if the lexical set Multiple lexical elements and the set of lexical elements If a word element has a many-to-one relationship, then the word element set... The maximum semantic score of these multiple lexical units is assigned to the corresponding lexical unit; if the lexical set Multiple lexical elements and the set of lexical elements If multiple lexical units in a word have many-to-many relationships, then the lexical set... The maximum semantic score of these multiple lexical units is assigned to the corresponding lexical unit.

3. The key-value caching compression method for large language model inference as described in claim 1 or 2, characterized in that, The semantic scoring model is the BERT model.

4. The key-value caching compression method for large language model inference as described in claim 1 or 2, characterized in that, When updating the translated lexical features using the updated positions, the following formula is used: ； in, and These represent the position indices before and after the translation, respectively. Indicates the lexical features before translation. Indicates the lexical features after translation. Indicates rotation parameters, The rotation matrix represents the rotation position encoding.

5. The key-value caching compression method for large language model inference as described in claim 1 or 2, characterized in that, Also includes: Maintain the prefix tree corresponding to the KV cache; For the nodes in the prefix tree Generate its child nodes hour: If node For nodes The first child node, and used for storing nodes. The last physical page of the corresponding lexical feature If the page boundary is not aligned, then the node will be... Lexical features are appended to the physical page. For the unwritten lexical feature parts, allocate free physical pages for writing, and simultaneously, at the node... Inject split markers into the metadata to indicate physical pages Some of the data stored in the middle belongs to the node. child nodes; If node For nodes The first child node, and used for storing nodes. The last physical page of the corresponding lexical feature If page boundaries are aligned, then free physical pages are allocated and written to the node. Lexical features; If node Not a node The first child node, and the node If a split marker exists in the metadata, then it is a node. The lexical features are assigned to physical pages, and the physical pages are... After the data in the node is copied to the starting position in the allocated physical page, the node will be... Lexical features are appended to the text; If node Not a node The first child node, and the node If no split marker is found in the metadata, a free physical page is allocated and written to the node. Lexical features.

6. The key-value caching compression method for large language model inference as described in claim 5, characterized in that, When searching the prefix tree to perform prefix path matching, if the currently searched node If a split marker exists in the metadata, it will be used to store the node. The last physical page of the corresponding lexical feature After popping from the currently matched prefix path, perform subsequent matching.

7. A key-value caching and compression system for large language model inference, characterized in that, include: Three-tier heterogeneous storage system and controller; The controller includes a compression module; the compression module is used to execute the key-value cache compression method for large language model inference as described in any one of claims 1 to 6, so as to compress lexical features in the pre-filling stage of large language model inference; The three-tier heterogeneous storage system includes: video memory, RAM, and solid-state drives; The video memory is used to perform large language model inference tasks. The video memory, the main memory, and the solid-state drive are all equipped with KV caches for caching compressed blocks of word features after compression.

8. The key-value caching and compression system for large language model inference as described in claim 7, characterized in that, The controller further includes: a management module and an execution module; The management module is used to manage the storage space allocation at each level of the three-level heterogeneous storage system and the index of the KV cache; The execution module is used to handle read and write operations of the memory and the solid-state drive, as well as data transfer operations between the memory and the video memory.

9. The key-value caching and compression system for large language model inference as described in claim 8, characterized in that, The management module is also used to perform asynchronous lifecycle scheduling of the three-tier heterogeneous storage system; the asynchronous lifecycle scheduling includes: Loading: When the inference task enters the scheduling pool but has not yet started model forward propagation, the prefix hit is checked. If the prefix is ​​hit but the video memory is not hit, space is allocated in the video memory and an instruction is sent to the execution module so that the execution module loads the corresponding compressed block from the memory into the video memory, or loads the corresponding compressed block from the solid-state drive into the memory and then loads the compressed block into the video memory. Prefetch: When the total size of the accumulated prompt words reaches a preset safety threshold since the last prefetch, an instruction is sent to the execution module so that the execution module prefetches the compressed block that exists only in the solid-state drive and not in the memory into the memory; Storage: During the decoding stage of the large language model inference, each time the decoding generates new lexical units that cause the current physical page in the KV cache to be filled, space is allocated in the memory and the solid-state drive, and an instruction is sent to the execution module so that the execution module moves the filled physical page from the KV cache to the memory and the solid-state drive in sequence to achieve persistent data storage; Evicting: When the KV cache space of any level in the three-level heterogeneous storage system reaches the corresponding exhaustion index, a compressed block is selected from the KV cache space for eviction according to a preset eviction strategy.

10. The key-value caching compression system for large language model inference as described in claim 8 or 9, characterized in that, The communication channel between the management module and the execution module is a hybrid inter-process communication channel based on ZeroMQ and shared memory.

Images