An importance-aware large model prefix kv rearrangement method and system

Abstract

The application discloses an importance-aware large model prefix KV rearrangement method and system, comprising the following steps: judging whether all tokens in each node in the radix tree are important or not according to the calculation results of historical queries and prefixes through an important information acquisition module; reading the data blocks containing K or V of the Imp sequence into the CPU memory through a prefix KV rearrangement module, and then placing the K or V corresponding to the important tokens in a data block and placing the K or V corresponding to the unimportant tokens in other data blocks; and updating the sequence in each radix tree node according to the token order after rearrangement, and placing the important tokens in front and the unimportant tokens at the back through a metadata update module. The application alleviates the read amplification problem, further reduces the data reading amount of the disk, reduces the loading time during prefix KV multiplexing, and thus shortens the time for generating the first token during large model reasoning.

Description

Technical Field

[0001] This invention relates to the field of computer science and artificial intelligence systems, specifically to a method and system for importantness-aware large model prefix key-value rearrangement. Background Technology

[0002] Existing prefix-aware large-model inference systems often directly use the large model for inference, potentially generating suboptimal results. For example, when asked about recent events not included in the model's training data, the model might give an incorrect answer. Furthermore, due to issues such as hallucinations, the model's responses may contain inaccurate or misleading information. To improve response quality, applications typically add a prefix containing rich contextual information to the user query to form a complete request before feeding it into a large language model (LLM). For example, retrieval-enhanced generative techniques search for documents relevant to the user query from external knowledge bases. Advanced GPT plugins, such as Chameleon, include tool definitions in system prompts and use few-shot examples to guide the LLM in performing complex inference tasks.

[0003] During large model inference, a K tensor and a V tensor are computed and generated for each token at each layer of the model. However, researchers have found that different requests often have partially or completely identical prefixes (e.g., similar queries might use RAG to retrieve partially or entirely the same relevant documents; the same GPT plugin can be used multiple times, generating the same system hints in different requests). Therefore, existing prefix-aware large model inference systems typically store and reuse these K and V tensors with shared prefixes (called prefix KV) to avoid recompiling the KV for the prefix portion. When a new request with a duplicate prefix arrives, the system reuses its prefix KV and loads it into GPU memory, thus avoiding the regeneration of the prefix KV and only requiring computation for the query portion. Therefore, reusing prefix KV reduces the time to first token (TTFT) of the large model, improving the user experience. The following example illustrates this process:

[0004] For example, suppose a user uses a large model plugin specifically for translating Chinese to English, and the user wants to translate the words "apple" and "banana" in sequence. When the user enters the query "apple", the system will add a prefix in front: "Please translate the user's Chinese into English.", so 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 in the local disk (since the total amount of data is too large to fit in the CPU memory and needs to be stored on the local disk). Assume that 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" arrives, the request sent by the system to the large model is "Please translate the user's Chinese into English. Banana". At this time, in order to complete the inference calculation of each layer, the large model does not need to recalculate the KV of the prefix because it can directly reuse the previously stored 7 K and 7 V, and then only calculate the K and V of "banana", thus avoiding redundant calculations and shortening the TTFT.

[0005] In addition, existing prefix-aware large model inference systems also consider token importance and reduce 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 large model inference result as loading all the prefix KV "Please translate the user's Chinese into English" because the other tokens are actually not important and do not help with the inference result.

[0006] Explanation of terms:

[0007] Token: The words / phrases, etc. formed after the sentence input to the large model is tokenized, which is the basic unit encoded into the large model subsequently.

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

[0009] Prefix: The natural language sentence added in front of the query problem input by the user;

[0010] Key (K): When the large model makes inferences, it will calculate and generate a K tensor for each token in each layer of the model;

[0011] Value (V): When the large model makes inferences, it will calculate and generate a V tensor for each token in 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 transmitting 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 into Japanese. When the user inputs "pear", the request actually submitted by this 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 loading only some important prefix KV can reduce the data loading volume, existing systems have a 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 prior art, 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 includes the following steps:

[0021] Based on the calculation results of historical queries and prefixes, determine whether all tokens within each node in 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 , where S represents the natural language sentence input by the user, t represents a single token obtained after sentence tokenization, the subscript of t corresponds to the position index of each token, and n represents that the sentence contains n tokens after tokenization; i Let the important token sequence be Imp = [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 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 acquisition module: Acquire the importance of each K tensor and V tensor in all stored prefix KVs; the acquisition process is performed offline;

[0026] Prefix KV rearrangement module: According to whether the KV in each node in 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 incorrect matching and searching for reusable prefixes when new requests arrive.

[0028] Furthermore, the working method of the importance information acquisition module is specifically as follows: Based on the historical queries of multiple users, along with the prefix, it is sent to the large model for inference. The attention matrix generated during the analysis and calculation process is accumulated for its columns; the accumulated values of each column 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] Furthermore, 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] Furthermore, the working method of the metadata update module is specifically as follows: Write the rearranged tokens into the nodes of the radix tree, and add a mapping list inside each node of the radix tree. 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 determine 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 read amplification problem;

[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 disk data reading volume;

[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 large model inference. BRIEF DESCRIPTION OF THE DRAWINGS

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

[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. Specific implementation mode

[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 the important token are also important, and the corresponding K and V of the unimportant token are also unimportant. Suppose 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 indexes of the 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 rearranges 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 tokens 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 K and V, and the rest are regarded as unimportant K and V.

[0046] Prefix KV Rearrangement Module: Reads the existing 2 data blocks, and packs the important K or V scattered in the 2 data blocks into one data block, and packs 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 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 in each node of the radix tree, and the mapping list records the index list of the subscripts if the original token order is to be restored; 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 judge 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." 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 will store their K and V in the order of tokens, so their K and V are scattered in different data blocks (marked in black in Figure 3 (a) of the figure), and the K and V of other tokens are unimportant (represented by slash shadows). Figure 3In graph (b), each node contains only one S-sequence label for the original token ID sequence.

[0051] like Figure 2 As shown, the present invention first uses an importance information adoption module to find all important tokens, namely "please", "Chinese", "translate", and "English". Then, since 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. The K and V order in the disk is updated by reading and writing to the disk. Then, the system uses a metadata update module to update the radix tree, adding a new sequence S' and a mapping list m to each radix tree node.

[0052] After using the KV rearrangement technique of this invention, the K or V corresponding to important tokens are placed in a single 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 / of / into". The problem with this is that when a new request, "Please translate the user's Chinese into English", arrives, it can only match the token "Please", reducing the reusable KV quantity and thus limiting the reduction in TTFT time. Therefore, the metadata update module of this invention, in addition to writing a rearranged S' sequence, also writes an m sequence in each node of the cardinality tree to recover 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] This invention implements its technology using FlexGen code technology. Two different open-source OPT models (OPT-6.7B, OPT-13B, and OPT-30B) were used for testing. The experiments were conducted on a server equipped with two AMD EPYC 7763 CPUs (64 cores total), 128GB of DRAM, an NVIDIA A100 GPU with 80GB of HBM video memory, and a 2TB Intel SSD with a read throughput of approximately 5GB / s. The GPU and CPU are connected via a PCIe 4.0 x16 interface.

[0054] This invention was tested using the RTE dataset. Two to ten few-sample examples were appended as system hints before each query on this dataset. These system hints were shared across different queries, with reuse frequency following a normal distribution. The prefix of the OPT-30B model was extended to a maximum length of 4K, while other OPT models were extended to 10K. To prevent runtime memory overflow errors and ensure that only a portion of the prefix key-value pairs (KVs) were cached (the remaining KVs were stored on SSDs), this invention allocated 10GB of GPU cache and 32GB of CPU cache to the prefix KVs. The remaining GPU and CPU memory were 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 per request before and after using the present invention under different models. As can be seen from Table 1, using the present invention can accelerate the TTFT by up to 1.3 times.

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

[0057]

[0058]

Claims

1. A prefix key-value rearrangement method for a large-scale importance-aware model, characterized in that, It includes the following steps: Based on the calculation results of historical queries and prefixes, determine whether all tokens within each node in 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 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 back to the local disk; According to the reordered 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 key-value 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 carried out offline; Prefix KV rearrangement module: a process of reading the original data block, repacking, and writing according to whether the KV in each node of the radix tree is important; specifically: 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; 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 metadata update module is specifically as follows: write the reordered 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 index list of the subscripts if the original token order is to be restored; when a new request arrives, it will combine the reordered token sequence and the mapping list to restore the token sequence before rearrangement, so as to determine whether KV can be reused.

Images