Large model prefix KV rearrangement method and system for importance perception

Abstract

The invention discloses an importance-aware large model prefix KV rearrangement method and system, and the method comprises the steps: judging whether all tokens in each node in a cardinal number tree are important or not through an important information collection module according to historical query and prefix calculation results; a prefix KV rearrangement module is used for reading K or V data blocks containing the Imp sequence into a CPU memory, then K or V corresponding to important tokens is placed in one data block again, and K or V corresponding to unimportant tokens is placed in other data blocks; and the metadata updating module updates the sequence in each cardinal number tree node according to the rearranged token sequence, the important tokens are placed in the front, and the unimportant tokens are placed in the back. According to the method, the read amplification problem is relieved, and the data reading amount of the disk is further reduced; and the loading time during prefix KV multiplexing is reduced, so that the time for generating a first token during large model reasoning is shortened.

Description

Technical Field

[0001] The present invention relates to the field of computer science artificial intelligence systems, and in particular to an importance-aware large model prefix KV rearrangement method and system. Background Art

[0002] Existing prefix-aware large model reasoning systems often use large models directly for reasoning, which may generate suboptimal results. For example, when asked about recent events that are not included in the model training data, the model may give incorrect answers. In addition, due to problems such as hallucinations, the model's answers may contain inaccurate or misleading information. To improve the quality of answers, applications usually add context containing rich information, namely prefixes, before the user query to form a complete request, which is then input into the large language model (LLM). For example, retrieval-augmented generation technology searches for documents related to the user query from an external knowledge base. Advanced GPT plugins, such as Chameleon, include tool definitions in system prompts and use few-sample examples to guide the LLM to perform complex reasoning tasks.

[0003] During the large model reasoning process, a K tensor and a V tensor are calculated and generated for each token in each layer of the model. However, researchers have found that the prefixes of different requests are often partially or completely identical (for example, similar queries may use RAG to retrieve some or all of the same related documents; the same GPT plug-in can be used multiple times, thereby generating the same system prompts in different requests). Therefore, existing prefix-aware large model reasoning systems usually store and reuse these shared prefix K and V tensors (called prefix KV) to avoid recalculating the KV of the prefix part. When a new request containing a repeated prefix arrives, the system reuses its prefix KV to load it into the GPU memory, thereby avoiding the regeneration of the prefix KV, and only the query part needs to be calculated. Therefore, reusing prefix KV will reduce the time to first token (TTFT) for large models to generate the first token and improve user experience. The following example illustrates this process:

[0004] For example, suppose a user uses a large model plug-in specifically for translating Chinese into English, and the user wants to translate the two words: apple and banana. When the user enters the query "apple", the system will add a prefix in front: "Please translate the user's Chinese into English." Therefore, the inference request sent to the large model is actually "Please translate the user's Chinese into English. Apple". The large model will calculate and generate K tensors and V tensors for each token in the entire request, and store the K and V corresponding to the entire prefix to the local disk (because the total amount of data is too large to be stored in the CPU memory, it needs to be stored on the local disk). Suppose the token division result of the prefix is ​​"Please translate / user / 's / Chinese / into / English." Since the prefix is ​​divided into 7 tokens, there are 7 K and 7 V tensors in each layer of the model. When the user's second query "banana" comes, the request sent to the big model by the system is "Please translate the user's Chinese into English. Banana". At this time, in order to complete the reasoning calculation of each layer, the big model does not need to recalculate the prefix KV, because it can directly reuse the 7 Ks and 7 Vs stored before, and then only calculate the K and V of "banana", thereby avoiding redundant calculations and shortening TTFT.

[0005] In addition, the existing prefix-aware big model reasoning system also considers the importance of tokens and reduces the loading of some tokens. For example, when reusing prefix KV, only loading the K and V of the four tokens "Please translate", "Chinese", "Translate", and "English" can obtain the same big model reasoning result as loading the entire prefix KV "Please translate the user's Chinese into English", because other tokens are actually not important and are not helpful for the reasoning result.

[0006] Explanation of terminology:

[0007] Token: The words / phrases formed after the sentences input to the big model are segmented, which are the basic units to be encoded in the big model later.

[0008] Query: The question that the user inputs to LLM (natural language description);

[0009] Prefix: A natural language sentence added to the front of the query question entered by the user;

[0010] Key (Key, K): When a large model is inferred, a K tensor is generated for each token in each layer of the model.

[0011] Value (V): When a large model is inferred, a V tensor is generated for each token at each layer of the model.

[0012] Prefix KV: The KV tensors generated by the prefix during inference.

[0013] Chunk Storage and Radix Tree:

[0014] To improve the efficiency of transferring K and V to the GPU, existing systems store data in chunks, that is, store the K or V corresponding to several consecutive tokens together. For example, in the above example, the K of "Please" and the K of "user" are stored as a chunk object.

[0015] In addition, to increase the proportion of reusable prefix KV as much as possible, existing systems use the data structure of a radix tree to record the prefixes of KV that have been calculated, which is convenient for new requests to match and search. In a radix tree, each node represents several tokens in the prefix, and the path from the root to any leaf node represents a complete prefix. If two prefixes have the same tokens, then the common tokens will only be represented once in the tree. Continuing with the above example to explain this process:

[0016] Suppose there is another large model plugin dedicated to translating Chinese to Japanese. When the user inputs "pear", the request actually submitted by the plugin to the large model is: "Please translate the user's Chinese into Japanese. Pear". Since the prefix of this plugin is not exactly the same as the prefix of the previous plugin, only the KV corresponding to the previous part of the tokens can be reused.

[0017] Problems Existing in Existing Prefix-Aware Large Model Inference Systems:

[0018] Although only loading some important prefix KV can reduce the data loading volume, existing systems have the problem of read amplification when loading important prefix KV. For example, in the above example, when loading the K or V of the important token "Please", since the K or V of the unimportant token "user" is stored in the same data block, it will also be loaded from the local disk into the CPU memory, resulting in the problem of read amplification. Generally speaking, existing systems will read unnecessary data, thus wasting the read bandwidth and increasing the read latency. Summary of the Invention

[0019] Aiming at the deficiencies of the existing technology, the purpose of the present invention is to provide an importance-aware large model prefix KV rearrangement method and system.

[0020] The purpose of the present invention is achieved by the following technical solutions: An importance-aware large model prefix KV rearrangement method, including the following steps:

[0021] Based on the calculation results of historical queries and prefixes, determine whether all tokens in each node of the radix tree are important. Among them, the K and V corresponding to important tokens are also important, and the K and V corresponding to unimportant tokens are also unimportant. Assume that the token sequence in a node containing n tokens is S = [t 0 ,t 1 ,…,t n-1 , where S represents the natural language sentence input by the user, t represents a single token obtained after sentence segmentation, the subscript of t corresponds to the position index of each token, and n represents the total number of tokens contained in the sentence after segmentation; Denote the sequence of important tokens as Imp = [t i ,t i+1 ,…,t j , 0 <= i < j < n, where i and j represent the position indices of important tokens in S, and n represents the total number of tokens in the S sentence;

[0022] Read the data block containing the K or V of the Imp sequence into the CPU memory, relocate the K or V corresponding to the important tokens in the data block into one data block, and place the K or V corresponding to the unimportant tokens in other data blocks; Rewrite the reorganized data block to the local disk;

[0023] According to the reordered token order, update the sequence S in each radix tree node to S'; Among them, place the important tokens in the front and the unimportant tokens in the back; And add a mapping list to record the index list for restoring S from S'.

[0024] The present invention also provides an importance-aware large model prefix KV rearrangement system, including:

[0025] Importance information collection module: Collect the importance of each K tensor and V tensor in all stored prefix KVs; The collection process is performed offline;

[0026] Prefix KV rearrangement module: According to whether the KV in each node of the radix tree is important, read the original data block and perform the process of repackaging and writing;

[0027] Metadata update module: Used to update the radix tree; Since rearrangement will disrupt the order of K and V, it will cause the inability to correctly match and find reusable prefixes when new requests arrive.

[0028] Further, the working method of the importance information acquisition module is specifically as follows: Based on the historical queries of multiple users, it is sent to the large model for inference together with the prefix, the attention matrix generated during the analysis and calculation process is analyzed, and the columns therein are accumulated; according to the accumulated values of each column, they are sorted in descending order from largest to smallest; assuming that the user designates x% of the tokens as important tokens, where 0 < x ≤ 100, then the KVs corresponding to the tokens in the first x% of the sorted columns are regarded as important Ks and Vs, and the rest are regarded as unimportant Ks and Vs.

[0029] Further, the working method of the prefix KV rearrangement module is specifically as follows: Read the existing 2 data blocks, and pack the important Ks or Vs scattered in the 2 data blocks into one data block, and pack the unimportant Ks or Vs into another data block; among them, the data block full of unimportant Ks or Vs will not be read, and the proportion of important Ks or Vs in the data block to be read is larger, alleviating the problem of read amplification.

[0030] Further, the working method of the metadata update module is specifically as follows: Write the rearranged tokens into the nodes of the radix tree, add a mapping list in each node of the radix tree, and the mapping list records the list of subscript indexes for restoring the original token order; when a new request comes, it will combine the rearranged token sequence and the mapping list to restore the token sequence before rearrangement, so as to judge whether the KV can be reused.

[0031] The beneficial effects of the present invention are as follows:

[0032] 1. By rearranging the KV, the present invention can reduce the KV data reading volume during the reuse of the large model prefix KV, alleviating the problem of read amplification;

[0033] 2. The present invention can reduce the amount of unimportant KV data entering the memory, so that more important KV can be accommodated under the same memory, thereby further reducing the data reading volume of the disk;

[0034] 3. The present invention can reduce the loading time during the reuse of the prefix KV, thereby shortening the time to generate the first token during the inference of the large model. BRIEF DESCRIPTION OF THE DRAWINGS

[0035] Figure 1 It is a schematic diagram of the structure of the data block storage and the radix tree;

[0036] Figure 2 It is a schematic diagram of the large model prefix KV rearrangement system with importance perception;

[0037] Figure 3 It is a schematic diagram of the composition of KV before rearrangement and the metadata structure, where (a) is the KV composition diagram and (b) is the metadata structure diagram;

[0038] Figure 4 It is a schematic diagram of the KV composition and metadata structure after rearrangement. Among them, (a) is the KV composition diagram, and (b) is the metadata structure diagram. Detailed implementation manners

[0039] The present invention will be described in detail below with reference to the accompanying drawings. Without conflict, the features in the following embodiments and implementation manners can be combined with each other.

[0040] As Figure 2 shown, an importance-aware large model prefix KV rearrangement method provided by an embodiment of the present invention includes the following steps:

[0041] 1. The important information acquisition module determines whether all tokens in each node of the radix tree are important according to the historical query and the calculation result of the prefix. Among them, the corresponding K and V of important tokens are also important, and the corresponding K and V of unimportant tokens are also unimportant. Assume that the token sequence in a node containing n tokens is S = [t 0 , t 1 , …, t n-1 (S represents the natural language sentence input by the user, t represents a single token obtained after sentence segmentation, the subscript of t corresponds to the position index of each token, and n represents that the sentence contains n tokens after segmentation); among them, the important token sequence is denoted as Imp = [t i , t i+1 , …, t j , 0 <= i < j < n (i and j represent the position indexes of important tokens in S, and n represents the total number of tokens in the S sentence);

[0042] 2. The prefix KV rearrangement module reads the data block containing the K or V of the Imp sequence into the CPU memory, and then relocates the K or V corresponding to the important tokens into one data block, and the K or V corresponding to the unimportant tokens into other data blocks; the rearranged data block will be rewritten to the local disk.

[0043] 3. The metadata update module updates the sequence S in each radix tree node to S' according to the rearranged token order, where the important tokens are placed in the front and the unimportant ones are placed in the back. In addition, a mapping list is added to record the index list for restoring S from S'.

[0044] An embodiment of the present invention also provides an importance-aware large model prefix KV rearrangement system, including:

[0045] Importance Information Collection Module: Based on the historical queries of multiple users, it is sent to the large model for inference together with the prefix, analyzes the attention matrix generated during the calculation process, and accumulates the columns therein; sorts them in descending order according to the accumulated values of each column; assuming that the user specifies x% of the tokens as important tokens, where 0 < x ≤ 100, then the KVs corresponding to the tokens in the first x% of the sorted columns are regarded as important Ks and Vs, and the rest are regarded as unimportant Ks and Vs.

[0046] Prefix KV Rearrangement Module: Reads the existing 2 data blocks, and packs the important Ks or Vs scattered in the 2 data blocks into one data block, and packs the unimportant Ks or Vs into another data block; among them, the data block full of unimportant Ks or Vs will not be read, and the proportion of important Ks or Vs in the data block to be read is larger, alleviating the problem of read amplification.

[0047] Metadata Update Module: Writes the rearranged tokens into the nodes of the radix tree, adds a mapping list inside each node of the radix tree, and the mapping list records the index list of subscripts for restoring the original token order; when a new request comes, it combines the rearranged token sequence and the mapping list to restore the token sequence before rearrangement, so as to determine whether the KV can be reused.

[0048] As Figure 1 shown, the following uses a specific embodiment to explain this process:

[0049] Suppose there is a user using a large model plugin dedicated to translating Chinese into English. After the user inputs the target vocabulary to be translated, the system will add a prefix (Prefix) in front of the user input: "Please translate the user's Chinese into English.", and assume that the token division result of the prefix is "Please translate / user's / Chinese / into / English."; in addition, there is another large model plugin dedicated to translating Chinese into Japanese. After the user inputs the target vocabulary to be translated, the system will add a prefix (Prefix) in front of the user input: "Please translate the user's Chinese into Japanese." Assume that the token division result of the prefix is "Please translate / user's / Chinese / into / Japanese". Then the system will use the radix tree to store the tokens corresponding to these two prefixes.

[0050] Since only the four tokens "Please translate", "Chinese", "translate", and "English" in the prefix "Please translate the user's Chinese into" are important tokens, the existing system stores their Ks and Vs in the order of tokens, so their Ks and Vs are scattered in different data blocks (marked in black in Figure 3 (a) of the figure), and the Ks and Vs of other tokens are unimportant (represented by diagonal shading). Figure 3In each node of the graph (b), there is only one S sequence marking the original token ID sequence.

[0051] like Figure 2 As shown, the present invention first uses the importance information adoption module to find all important tokens, namely "Please", "Chinese", "Translate", and "English"; then because the first three tokens are in the same node (root node) of the radix tree, the K and V corresponding to these three tokens are reordered and re-blocked, and the KV arrangement order in the disk is updated by reading and writing the disk; then the system uses the metadata update module to update the radix tree, and adds a new sequence S' and a mapping list m in each radix tree node.

[0052] After using the KV rearrangement technology of the present invention, the K or V corresponding to the important tokens is placed in one data block as much as possible, such as Figure 4 As shown in Figure (a); the original text corresponding to the rearranged root node is "Please translate / Chinese / into / user / ". The problem with this is that when a new request "Please translate the user's Chinese into English" arrives, only the token "Please translate" can be matched, which reduces the amount of reusable KVs, thereby limiting the reduction in TTFT time. To this end, the metadata update module of the present invention writes a rearranged S' sequence in each node of the radix tree, as well as an m sequence for restoring the original S sequence; for example, in Figure 4 In Figure (b), S0 can be recovered by vector operation of S'0[m0], that is, S0 = S'0[m0].

[0053] The present invention implements the technology of the present invention on the FlexGen code technology. Two open source OPT models of different sizes were used for testing (i.e., OPT-6.7B, OPT-13B, and OPT-30B). The experiment of the present invention was carried out on a server equipped with two AMD EPYC 7763 CPUs (total 64 cores), 128GB DRAM, an NVIDIA A100 GPU with 80GB HBM video memory, and a 2TB Intel SSD with a read throughput of about 5GB / s. The GPU and the CPU are connected via a PCIe 4.0×16 interface.

[0054] The present invention uses the RTE dataset for testing. Two to ten few-sample examples are appended before each query of the dataset as system prompts. These system prompts are shared between different queries, and the reuse frequency follows a normal distribution. The prefix of the OPT-30B model is extended to a maximum length of 4K, while other OPT models are extended to 10K. In order to prevent memory overflow errors at runtime and ensure that only some prefix key-value pairs (KVs) are cached (the remaining KVs are stored on the SSD), the present invention allocates 10GB of GPU cache and 32GB of CPU cache for prefix KVs, and the remaining GPU and CPU memory are used to store model weights, KV caches used in the decoding phase, and input data. Each block contains keys or values ​​from 64 tokens.

[0055] Table 1 compares the average TTFT of each request before and after using the present invention under different models. It can be seen from Table 1 that the TTFT can be accelerated by up to 1.3 times after using the present invention.

[0056] Table 1: Comparison of TTFT time (seconds) for large model inference under different models

[0057]

[0058]

Claims

1. An importance-aware large model prefix KV rearrangement method, characterized in that: It includes the following steps: Based on the calculation results of historical queries and prefixes, it is determined whether all tokens in each node of the radix tree are important; among them, the K and V corresponding to important tokens are also important, and the K and V corresponding to unimportant tokens are also unimportant; assume that the token sequence in a node containing n tokens is S = [t0, t1, …, t n-1 (S represents the natural language sentence input by the user, t represents a single token obtained after sentence segmentation, the subscript of t corresponds to the position index of each token, and n represents that the sentence contains n tokens after segmentation); among them, the important token sequence is denoted as Imp = [t i , t i+1 , …, t j , 0 <= i < j < n (i and j represent the position indices of important tokens in S, and n represents the total number of tokens in the S sentence); Read the data block containing K or V with Imp sequences into the CPU memory, relocate the K or V corresponding to the important tokens in the data block into one data block, and place the K or V corresponding to the unimportant tokens in other data blocks; rewrite the reorganized data block back to the local disk; According to the rearranged token order, update the sequence S in each radix tree node to S'; where the important tokens are placed in the front and the unimportant tokens are placed in the back; and add a mapping list to record the index list for restoring S from S'.

2. An importance-aware large-model prefix KV rearrangement system, characterized in that: It includes: Importance information acquisition module: Acquire the importance of each K tensor and V tensor in all stored prefix KVs; The acquisition process is performed offline; Prefix KV rearrangement module: Read the original data block and perform the process of repackaging and writing according to whether the KV in each node of the radix tree is important; Metadata update module: Used to update the radix tree; to solve the problem that the rearrangement disrupts the order of K and V, resulting in the inability to correctly match and find the prefix for reuse when a new request arrives.

3. The importance-aware large-model prefix KV rearrangement system according to claim 2, characterized in that: The working method of the importance information acquisition module is specifically as follows: According to the historical queries of multiple users, send them together with the prefix into the large model for inference, analyze the attention matrix generated during the calculation process, and accumulate the columns therein; sort them in descending order according to the accumulated values of each column; assume that the user designates x% of the tokens as important tokens, where 0 < x ≤ 100, then regard the KV corresponding to the tokens in the first x% columns after sorting as important K and V, and the rest as unimportant K and V.

4. The importance-aware large-model prefix KV rearrangement system according to claim 2, characterized in that: The working method of the prefix KV rearrangement module is specifically as follows: Read the existing 2 data blocks, and pack the important K or V scattered in the 2 data blocks into one data block, and pack the unimportant K or V into another data block; among them, the data block full of unimportant K or V will not be read, and the proportion of important K or V in the data block to be read increases.

5. The importance-aware large model prefix KV rearrangement system according to claim 2, characterized in that: The working method of the metadata update module is specifically as follows: Write the rearranged tokens into the nodes of the radix tree, add a mapping list in each node of the radix tree, and the mapping list records the subscript index list for restoring the original token order; when a new request arrives, it will combine the rearranged token sequence and the mapping list to restore the token sequence before rearrangement, so as to determine whether KV can be reused.

Images