CacheBlend: Fast Large Language Model Serving with Cached Knowledge Fusion | AI Research Paper Details

Large language models (LLMs) often incorporate multiple text chunks in their inputs to provide the necessary contexts. To speed up the prefill of the long LLM inputs, one can pre-compute the KV cache of a text and re-use the KV cache when the context is reused as the prefix of another LLM input. However, the reused text chunks are not always the input prefix, and when they are not, their precomputed KV caches cannot be directly used since they ignore the text's cross-attention with the preceding text in the LLM input. Thus, the benefits of reusing KV caches remain largely unrealized. This paper tackles just one question: when an LLM input contains multiple text chunks, how to quickly combine their precomputed KV caches in order to achieve the same generation quality as the expensive full prefill (i.e., without reusing KV cache)? We present CacheBlend, a scheme that reuses the pre-computed KV caches, regardless prefix or not, and selectively recomputes the KV values of a small subset of tokens to partially update each reused KV cache. In the meantime,the small extra delay for recomputing some tokens can be pipelined with the retrieval of KV caches within the same job,allowing CacheBlend to store KV caches in slower devices with more storage capacity while retrieving them without increasing the inference delay. By comparing CacheBlend with the state-of-the-art KV cache reusing schemes on three open-source LLMs of various sizes and four popular benchmark datasets of different tasks, we show that CacheBlend reduces time-to-first-token (TTFT) by 2.2-3.3X and increases the inference throughput by 2.8-5X, compared with full KV recompute, without compromising generation quality or incurring more storage cost.

https://www.aimodels.fyi/papers/arxiv/cacheblend-fast-large-language-model-serving-rag

CacheBlend: Fast Large Language Model Serving for RAG with Cached...

Large language models (LLMs) often incorporate multiple text chunks in their inputs to provide the necessary contexts. To speed up the prefill of the long LLM inputs, one can pre-compute the KV...

https://arxiv.org/abs/2405.16444

边读边想

1.Introduction

To ensure high-quality and consistent responses, applications often supplement the user query with additional texts to provide the necessary context of domain knowledge or user-specific information.

A typical example is Retrieval-Augmented Generation (RAG) where a user query will be prepended by multiple text chunks retrieved from a database to form the LLM input.

before generating any token, an LLM first uses prefill to go through the entire LLM input to produce the KV cache—concatenation of tensors associated with each input token that embeds the token’s “attention” with its preceding tokens.

Despite many optimizations, the delay and computation of prefill grow super-linearly with the input length, and can easily slow down the service, especially on long LLM inputs

Cross Attention in Transformer

Cross attention is a key component in transformers, where a sequence can attend to another sequence’s information, making it essential for…

https://medium.com/@sachinsoni600517/cross-attention-in-transformer-f37ce7129d78

selective KV recompute performs prefill on the input text in a traditional layer-by-layer fashion; however, in each layer, it updates the KV of a small fraction of tokens while reusing the KV of other tokens.

Comparing with full KV reuse, CacheBlend achieves much better generation quality with a small amount of extra KV update.

Fortunately, this small amount of extra computation does not affect the end-to-end inference latency, as it is hidden through pipeline parallelism by CacheBlend. Specifically, CacheBlend parallelizes partial KV update on one layer with the fetching of the KV cache on the next layer into GPU memory

Though pipelining KV compute and fetching is not new, we note that pipelining enables CacheBlend to store KV caches in slower non-volatile devices (e.g., disk or a separate storage server) without incurring extra delay as putting them in CPU memory.

2.Background

The prefill phase computes the KV cache layer by layer. The input tokens embeddings on each layer are first transformed into query (Q), key (K), and value (V) vectors, of which the K and V vectors form one layer of the KV cache. The LLM then multiplies Q and K vectors to obtain the attention matrix—the attention between each token and its preceding tokens—and does another dot product between the (normalized and masked) attention matrix with the V vector. The resulting vector will go through multiple neural layers to obtain the tokens’ embeddings on the next layer.

When the KV cache of a prefix is available, the prefill phase will only need to compute the forward attention matrix (between the suffix tokens and the prefix tokens) on each layer which directly affects the generated token.

3. Motivation

3.1 Opportunities of reusing KV caches

Since the reused contexts typically contain more information than the user queries, the prefill on the “context” part of the input accounts for the bulk of prefill overhead. Thus, it would be ideal to store and reuse the KV caches of reused texts, in order to avoid the prefill overhead when these texts are used again in different LLM inputs.

3.2 Why is prefix caching insufficient?

in prefix caching, the KV cache of a reusable text chunk is precomputed once, and if the text chunk is at the prefix of an LLM input, then the precomputed KV cache can be reused to avoid prefill on the prefix. The advantage of prefix caching is that the KV cache of a prefix is not affected by the succeeding text, so the generation result will be identical to full KV recompute (without the KV cache).

The disadvantage of prefix caching is also clear. To answer one query, applications, such as RAG, often prepend multiple text chunks in the LLM input to provide different contexts necessary for answering the query. As a result, except the first chunk, all other chunks’ KV caches are not reused since they are not the prefix of the LLM input.

In short, prefix caching can only save the prefill of the first text chunk, so the saving will be marginal when the LLM input includes more text chunks, even if they are reused

3.3 Why is full KV reuse insufficient?

It concatenates independently precomputed KV caches of recurring text chunks with the help of buffers to maintain the positional accuracy of each text chunk.

For instance, to concatenate the KV caches of chunks 𝐶1 and 𝐶2, PromptCache first needs to precompute the KV cache of 𝐶2 by running prefill on a hypothetical input that prepends 𝐶2 with a dummy prefix of length greater or equal to 𝐶1 This way, even if 𝐶2 is not the prefix, we still correctly preserve the positional information of 𝐶2’s KV cache, though each chunk’s KV cache will have to be precomputed multiple times.

a more fundamental problem is that the KV cache of nonprefix text chunk (e.g., 𝐶2) ignores the cross-attention between the chunk and the preceding text (e.g., 𝐶1). This is because the preceding text is not known when precomputing the KV cache.Ignoring cross-attention can lead to a wrong response.

With full prefill or prefix caching, the result is clear and correct. With full KV reuse the KV caches of the two text chunks are precomputed, with each chunk having the right positional embedding, and then concatenated to form the KV cache. However, if the LLM uses this KV cache to generate the answer, it will start to ramble and not produce the right answer

The absence of cross-attention in full KV reuse causes significant discrepancies in the forward attention matrix (explained in §2), which contains the attention between context tokens and the last few tokens, and directly affects the generated tokens.

4 Fast KV Cache Fusing

Goal. When an LLM input includes multiple re-used text chunks, how to quickly update the pre-computed KV cache, such that the forward attention matrix (and subsequently the output text) has minimum difference with the one produced by full KV recompute.

To achieve our goal, we present CacheBlend, which recomputes the KV of a selective subset of tokens on each layer while reusing other tokens’ KV

We begin with the notations (§4.1), and then describe how to recompute the KV of only a small subset of tokens (§4.2), and finally explain how to select the tokens on each layer whose KV will be recomputed (§4.3).

4.1 Terminology

Table 1 summarizes the notations used in this section. For a given list of 𝑁 text chunks, we use to denote the KV cache from full KV recompute, to denote the precomputed KV cache, and to denote the CacheBlend-updated KV cache. Here, each of these KV caches is a concatenation of KV caches associated with different text chunks. Each layer 𝑖 of the KV cache, , produces the forward attention matrix .

KV deviation: We define the KV deviation of a KV cache 𝐾𝑉 on layer 𝑖 of token 𝑗 as the absolute difference between and , denoted as . It measures how much different the given KV is on a particular token and layer compared to the full-prefilled KV cache. We will later use the KV deviation to identify which tokens’ KV has higher deviation and thus need to be updated.

Attention deviation: Similarly, for the forward attention matrix on layer 𝑖, we define the attention deviation, denoted as , to be the L-2 norm of its difference with . Recall from §3.3 that full KV reuse suffers from deviation in the forward attention matrix (illustrated in Figure 4) due to the absence of cross-attention.

Using these notations, our goal can be formulated as how to quickly update the precomputed KV cache to the new KV cache , such that the attention deviation on any layer 𝑖, is minimized.

4.2 Selectively recomputing KV cache

Workflow: The default implementation of prefill (depicted in Figure 5(a)) does not “skip” tokens while only computes the KV of a subset of tokens. Instead, CacheBlend runs the following steps (depicted in Figure 5(b)): • It first applies a mask on the input of each layer 𝑖 to reduce it to a subset of selected tokens. • It then transforms the reduced input into the vectors will also be restricted to the selected tokens. • It then expands the vector and vector by reusing the KV cache entries associated with the un-selected tokens on layer 𝑖, so that the attention matrix includes attention between selected tokens and all other tokens. • Finally, it runs the same attention module to produce the input of the next layer

These changes make little assumption on the exact transformer process and can be integrated with many popular transformers (more details in §6). It is important to notice that the compute overhead is proportional to the number of selected tokens. This is because it only runs computation associated with the selected tokens. If we recompute 𝑟% of tokens per layer, the total compute overhead will be 𝑟% of full prefill.

4.3 Selecting which tokens to recompute

Our intuition, therefore, is to prioritize recomputing the KV of tokens who have high KV deviations. Of course, this intuitive scheme is not feasible as it needs to know fullprefilled KV cache, and we will make it practical shortly.

To show the effectiveness of selecting tokens with high KV deviations, Figure 6 uses three models on the dataset of Musique (please see §7.1 for details). It shows the change of average attention deviation across all layers 𝑖, , after we use the aforementioned scheme (§4.2) to select and recompute the 𝑟% of tokens 𝑗 who have the highest KV deviation . As the recompute ratio (𝑟) increases, we can see that the attention deviation gradually reduces, and the biggest drops happen when the top few tokens with the highest KV deviations are recomputed.

Empirically, it suggests the following insight. Insight 1. On layer 𝑖, recomputing the KV of token 𝑗 who has a higher KV deviation (i.e. ) reduces the attention deviation (i.e., ) by a greater amount.

Insight 2. Tokens with the highest KV deviations on one layer are likely to have the highest KV deviations on the next layer.

Thus, if we recompute the KV of, say 10%, of tokens on a layer 𝑖, we should choose the 10% of tokens which have the highest KV deviations. We refer to these tokens as the High-KV-Deviation (or HKVD) tokens on layer 𝑖 .

Do we need to recompute KV for most tokens? In §7, we empirically show that choosing 10-20% tokens as HKVD tokens and recomputing their KV suffices to greatly reduce the attention deviation and preserve generation quality. This can be intuitively explained by attention sparsity, a well-studied property observed in many transformer models by prior research. It says that in an attention matrix, high attention typically only occurs between a small number of tokens and their preceding tokens.

We can see that a small fraction, about 10-15%, of tokens have much higher KV deviation than others, which corroborates the sparsity of cross-attention. If a token has very low attention with other chunks’ tokens (i.e., low cross-attention with other chunks), the KV deviation between and will be low and thus do not need to be recomputed. Only when a token has a high attention with other chunks (high KV deviation compared with ground truth), should its KV be recomputed.

How to identify the HKVD tokens without knowing the true KV values or attention matrix? Naively, to identify the HKVD tokens, one must know the fully recomputed 𝐾𝑉 full 𝑖 of each layer 𝑖 in the first place, but doing so is too expensive and defeats the purpose of selective KV recompute. Instead, we observe that the HKVD tokens on different layers are not independent: Insight 2. Tokens with the highest KV deviations on one layer are likely to have the highest KV deviations on the next layer.

For instance, if the HKVD tokens on the first layer are tokens 2, 3, and 5, these three tokens will likely also have higher KV deviations than most other tokens on the second layer. Figure 8 uses the same setting as Figure 7 and shows Spearman’s rank correlation score between the KV deviation of tokens between two neighboring layers. The figure shows a consistently high similarity of HKVD tokens between different layers

The intuition behind this correlation lies in the previous observation that the input embedding of each token changes slowly between layers in transformer models. Therefore, KV cache between layers should also bear similarity as KV cache is generated from the input embedding with a linear transformation.(We should clarify that although the HKVD tokens are similar across layers, the attention matrices between layers can still be quite different.)

Given the substantial correlation between the HKVD tokens, a straightforward solution is that we can perform prefill on the first layer first, pick the HKVD tokens of the first layer, and only update their KV on all other layers. Since an LLM usually has over 30 layers, this process can save most of the compute compared to full KV recompute. That said, using only the attention deviation of different tokens on the first layer may not be statistically reliable to pick HKVD tokens of all layers, especially deeper layers.

Thus, we opt for a gradual filtering scheme (depicted in Figure 9). If on average we want to pick 𝑟% HKVD tokens per layer, we will pick tokens based on the token-wise attention deviation on the first layer, with being slightly higher than 𝑟, and use them as the HKVD tokens on the second layer. Then we recompute the KV of these HKVD tokens on the second layer and pick tokens that have the highest token-wise attention deviation, with slightly less than , as the HKVD tokens on the next layer, and so forth. Intuitively, this gradual-filtering scheme eventually picks the HKVD tokens who have high attention deviation, not only on the first layer but also on multiple layers, which empirically is statistically more reliable to identify the HKVD tokens on each layer.

5 CacheBlend System Design

We present a concrete system design for CacheBlend, which reduces the impact of the selective KV recompute using the following basic insight.

Basic insight: If the delay for selective KV recompute (§4.3) is faster than the loading of KV into GPU memory, then properly pipelining the selective KV recompute and KV loading makes the extra delay of KV recompute negligible.

Pipelining KV loading and recompute: In CacheBlend, the selective recompute of one layer can start immediately after pre-computed the KV cache of the previous layer is loaded into the GPU. This is because which tokens’ KV to recompute on one layer only depends on the KV deviation of the previous layer’s tokens. As a result, if loading the pre-computed KV for one layer is faster（? slower） or equal to selective KV recompute of one layer, the KV-loading delay should be able to hide the selective recompute delay, i.e., without incurring any extra delay on time-to-first-token (TTFT).

Take the Llama-7B model and a 4K-long context, recomputing 15% of the tokens (the default recompute ratio) only takes 3 ms per layer, while loading one layer’s KV cache takes 16 ms from an NVME SSD (§7). In this case, KV loading can hide the delay for KV recompute on 15% of the tokens, i.e., KV recompute incurs no extra delay. Recomputing more tokens, which can slightly improve generation quality, may not incur extra delay either, as long as the delay is below 16 ms. On the contrary, with another model, Llama-70B, recomputing 15% of tokens takes 7 ms, but it only takes 4 ms to load one layer’s KV from an NVME SSD. Here KV loading does not completely hide the recompute delay. In short, a controller is needed to intelligently pick the recompute ratio as well as where to store the KV cache (if applicable).

5.1 Key Components

To realize the benefit of pipelining KV loading and recompute, our system has three major components. Loading Controller: We face two design questions in practice: First, given a storage device, how to choose a recompute ratio (what fraction of tokens to recompute KV per layer) without incurring extra delay to time-to-first-token (TTFT)? Figure 10(a) illustrates an example that, if we select a recompute ratio wisely, the recompute will not cause any extra delay to loading.

For this, the controller uses two delay estimators to find an idealized recompute ratio, such that the recompute delay is close to the loading delay. Given the recompute ratio 𝑟, length of context to be loaded 𝐿, and LLM, the recompute delay estimator calculates the expected delay (𝑟%, 𝐿𝐿𝑀, 𝐿).The loading delay estimator estimates the loading delay of the KV cache of one layer, (𝐿𝐿𝑀, 𝐿, 𝑠𝑡𝑜𝑟𝑎𝑔𝑒_𝑑𝑒𝑣𝑖𝑐𝑒) , based on the LLM, the storage device’s speed (which is measured offline), and the length of context 𝐿.

The controller calculates an idealized recomputation ratio such that the loading delay can hide the recompute delay, without degrading the inference quality. It first picks the recompute ratio 𝑟% such that (𝑟%, 𝐿𝐿𝑀, 𝐿) is equal to (𝐿𝐿𝑀, 𝐿, 𝑠𝑡𝑜𝑟𝑎𝑔𝑒_𝑑𝑒𝑣𝑖𝑐𝑒), and then takes the max of 𝑟% and %, where % is the minimal recompute ratio that empirically has low negligible quality drop from full KV recompute. In practice, we found % to be 15% from Figure 16.

Another design question facing the loading controller is: If we only want to do KV recompute of a fixed selective recompute ratio 15%, how can we choose the right storage device to store KVs such that no extra delay is caused? As shown in Figure 10(b), under a fixed recompute ratio, the controller should pick the cheapest storage device among all devices that do not increase the delay. The system developers can provide a list of potential storage device(s), and the controller uses a storage cost estimator which estimates the cost of storing KVs for each device, namely (𝐿𝐿𝑀, 𝐿,𝑇 , 𝑠𝑡𝑜𝑟𝑎𝑔𝑒_𝑑𝑒𝑣𝑖𝑐𝑒), based on the LLM, length of context 𝐿 and time duration 𝑇 needed to store it (if it is cloud storage). Then it uses (15%, 𝐿𝐿𝑀, 𝐿) and (𝐿𝐿𝑀, 𝐿, 𝑠𝑡𝑜𝑟𝑎𝑔𝑒_𝑑𝑒𝑣𝑖𝑐𝑒) to estimate the recompute and loading delays for all devices. Lastly, it finds out which storage device is the cheapest where ≥ .

KV cache store (mapping LLM input to KV caches): The KV cache store splits an LLM input into multiple text chunks, each of which can be reused or new. For instance, a RAG input typically consists of multiple retrieved context chunks (likely of a fixed length) and the user input. The splitting of LLM inputs is specific to the application, and we implement the same strategy as described in recent work. Once the input is split into text chunks, each chunk is hashed to find their corresponding KV cache, in the same way as the block hashing is implemented in vLLM. The KV caches of new chunks by the fusor (explained soon) are added to the devices. When the storage devices are full, the least recently used KV caches are evicted.

Fusor: The cache fusor (§4) merges pre-computed KV caches via selective recompute. Recall from §4.3, the decision of which tokens need to be recomputed for one layer depends on the recompute of the previous layer. Thus, the fusor waits until the recompute for the previous layer is done, and the KV caches for layer 𝐿 are loaded into the queue on GPU memory and then perform selective recompute using the recompute ratio 𝑟% calculated by the loading controller. The fusor repeats this process until all the layers are recomputed.

We put the key components together in an LLM inference workflow in Figure 11.

5.2 Putting them together

When a user of an LLM application submits a question（1）, a list of relevant text chunks will be queried. The loading controller then queries the KV cache manager on whether the KV caches for those text chunks exist, and where they are stored.（2）（3） Next, the KV cache manager returns this information back to the loading controller and the controller computes the idealized selective recomputation ratio, sends it to the fusor（4）, and loads the KV caches into a queue in GPU memory.（4） The KV cache fusor continuously recomputes the KV caches in the queue, until all layers are recomputed（5）. Lastly, the fused KV cache is input into the LLM inference engine, which generates the answer to the user question based on the KV cache.

6 Implementation

• fetch_kv(text, layer_id) -> KVCache: given a piece of text and a layer id, CacheBlend fetches the corresponding KV cache from KV store into the GPU. Returns -1 if the KV cache is not in the system. • prefill_layer(input_dict, KVCache) -> output_dict: CacheBlend takes in the input and KV cache of this layer and performs the partial prefill process for this particular layer. The output is used as the input for the next layer. • synchronize(): CacheBlend requires synchronization before prefilling every layer to make sure the KV cache of this layer has already been loaded into the GPU.

We implement these three interfaces inside vLLMs. For fetch_kv, we first calculate the hash of the text and search if it is inside the KV store system. If it is present, we call torch.load() to load it into GPU memory if KV cache is on disk or use torch.cuda() if the KV cache is inside CPU memory.

For prefill_layer, we implement this interface on top of the original layer function in vLLM that performs one layer of prefill. Three key-value pairs are recorded in the input_dict: (1) the original input data input_org required for prefilling an LLM layer (e.g., input_tensor, input_metadata), (2) a check_flag indicating whether HKVD tokens will be selected in this layer, and (3) HKVD_indices that track the indices of HKVD tokens. If check_flag is True, the input tokens with the largest deviation between the newly computed KV cache and the loaded KV cache will be selected as the HKVD tokens. If check_flag is False, partial prefill will only be performed on the current HKVD tokens indicated by HKVD_indices. Only the KV cache of the HKVD tokens will be computed and updated at each layer.

In the partial prefill for layer 𝑖, two threads are used to pipeline the computation (prefill_layer) of layer 𝑖 and the KV cache loading (fetch_kv) of the next layer 𝑖 + 1. synchronize is called before prefill_layer to assure the KV cache needed for prefill has been loaded into GPU.

Managing KV cache: CacheBlend manages the KV caches such that: If KV cache is not inside the system and is recomputed by the LLM engine in the runtime, we will move the KV cache into CPU by torch.cpu() and open a thread to write it back to disk in the background with torch.save(). During fetch_kv, we go through through the hash tables to fetch KV cache for the fusor. The hash tables are kept in CPU for their relatively small size (16MB for one million chunks).

7 Evaluation

Our key takeaways from the evaluation are: • TTFT reduction: Compared to full KV recompute, CacheBlend reduces TTFT by 2.2-3.3× over several models and tasks. • High quality: Compared with full KV reuse, CacheBlend improves quality from 0.15 to 0.35 in F1-score and Rouge-L score, while having no more than 0.01-0.03 quality drop compared to full KV recompute and prefix caching. • Higher throughput: At the same TTFT, CacheBlend can increase throughput by up to 5× compared with full KV recompute and 3.3× compared with prefix caching.

复现

关于推理过程（特别是hidden states）

How a Transformer works at inference vs training time

I made this video to illustrate the difference between how a Transformer is used at inference time (i.e. when generating text) vs. how a Transformer is trained. Disclaimer: this video assumes that you are familiar with the basics of deep learning, and that you've used HuggingFace Transformers at least once. If that's not the case, I highly recommend this course: http://cs231n.stanford.edu/ which will teach you the basics of deep learning. To learn HuggingFace, I recommend our free course: https://huggingface.co/course. The video goes in detail explaining the difference between input_ids, decoder_input_ids and labels: - the input_ids are the inputs to the encoder - the decoder_input_ids are the inputs to the decoder - the labels are the targets for the decoder. Resources: - Transformer paper: https://arxiv.org/abs/1706.03762 - Jay Allamar's The Illustrated Transformer blog post: https://jalammar.github.io/illustrated-transformer/ - HuggingFace Transformers: https://github.com/huggingface/transformers - Transformers-Tutorials, a repository containing several demos for Transformer-based models: https://github.com/NielsRogge/Transformers-Tutorials.

https://www.youtube.com/watch?v=IGu7ivuy1Ag

tokenizer

python相关

python的装饰器

Python 函数装饰器 | 菜鸟教程

装饰器(Decorators)是 Python 的一个重要部分。简单地说：他们是修改其他函数的功能的函数。他们有助于让我们的代码更简短，也更Pythonic（Python范儿）。大多数初学者不知道在哪儿使用它们，所以我将要分享下，哪些区域里装饰器可以让你的代码更简洁。 首先，让我们讨论下如何写你自己的装饰器。 这可能是最难掌握的概念之一。我们会每次只讨论一个步骤，这样你能完全理解它。 一切皆对象 首先我们来理解下 Python 中的函..

https://www.runoob.com/w3cnote/python-func-decorators.html

python的类方法

python 中 cls 指的是什么

在Python中，cls是指向类本身的引用、用于指代当前类或构造类的实例。它主要用于类方法的第一个参数，这与实例方法的第一个参数self相对应，self指向类的实例对象。使用cls的好处是，它可以让类的某个方法被类本身调用，而不是由类的某个具体实例调用。这种机制对于创建类的实例或实现基于当前类属性的逻辑非常有用。

https://docs.pingcode.com/ask/294997.html

 

深入代码：

问题

Q1：

Q2：

Q3:

When the KV cache of a prefix is available, the prefill phase will only need to compute the forward attention matrix (between the suffix tokens and the prefix tokens) on each layer which directly affects the generated token.

复现问题Q4:

q1

q2

Q3：