Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 47 additions & 0 deletions docs/source/automatic_prefix_caching/apc.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,47 @@

# Introduction

## What is automatic prefix caching

Automatic prefix caching (APC in short) caches the intermediate computation results of existing queries (in the format of KV cache), so that a new query can directly reuse the intermediate computation results if it shares the same prefix with one of existing queries, allowing the new query to be served much quicker.
Comment thread
KuntaiDu marked this conversation as resolved.
Outdated

Comment thread
KuntaiDu marked this conversation as resolved.
Outdated
## Enabling APC in vllm

Set `enable_prefix_caching=True` in vllm engine to enable APC. Here is a minimum example (this example is extracted from [vllm/benchmarks/benchmark_prefix_caching.py](https://github.com/vllm-project/vllm/blob/main/benchmarks/benchmark_prefix_caching.py)):

```python

from vllm import LLM, SamplingParams

LONG_PROMPT = "You are a helpful assistant in recognizes the content of tables in markdown format. Here is a table as fellows. You need to answer my question about the table.\n# Table\n|Opening|Opening|Sl. No.|Film|Cast|Director|Music Director|Notes|\n|----|----|----|----|----|----|----|----|\n|J A N|9|1|Agni Pushpam|Jayabharathi, Kamalahasan|Jeassy|M. K. Arjunan||\n|J A N|16|2|Priyamvada|Mohan Sharma, Lakshmi, KPAC Lalitha|K. S. Sethumadhavan|V. Dakshinamoorthy||\n|J A N|23|3|Yakshagaanam|Madhu, Sheela|Sheela|M. S. Viswanathan||\n|J A N|30|4|Paalkkadal|Sheela, Sharada|T. K. Prasad|A. T. Ummer||\n|F E B|5|5|Amma|Madhu, Srividya|M. Krishnan Nair|M. K. Arjunan||\n|F E B|13|6|Appooppan|Thikkurissi Sukumaran Nair, Kamal Haasan|P. Bhaskaran|M. S. Baburaj||\n|F E B|20|7|Srishti|Chowalloor Krishnankutty, Ravi Alummoodu|K. T. Muhammad|M. S. Baburaj||\n|F E B|20|8|Vanadevatha|Prem Nazir, Madhubala|Yusufali Kechery|G. Devarajan||\n|F E B|27|9|Samasya|Madhu, Kamalahaasan|K. Thankappan|Shyam||\n|F E B|27|10|Yudhabhoomi|K. P. Ummer, Vidhubala|Crossbelt Mani|R. K. Shekhar||\n|M A R|5|11|Seemantha Puthran|Prem Nazir, Jayabharathi|A. B. Raj|M. K. Arjunan||\n|M A R|12|12|Swapnadanam|Rani Chandra, Dr. Mohandas|K. G. George|Bhaskar Chandavarkar||\n|M A R|19|13|Thulavarsham|Prem Nazir, sreedevi, Sudheer|N. Sankaran Nair|V. Dakshinamoorthy||\n|M A R|20|14|Aruthu|Kaviyoor Ponnamma, Kamalahasan|Ravi|G. Devarajan||\n|M A R|26|15|Swimming Pool|Kamal Haasan, M. G. Soman|J. Sasikumar|M. K. Arjunan||\n\n" # noqa: E501

# let enable_prefix_caching=True to enable APC
llm = LLM(model='baichuan-inc/Baichuan2-13B-Chat',
enable_prefix_caching=True)

sampling_params = SamplingParams(temperature=0, max_tokens=10)

# Asking the content of cell (1,1)
Comment thread
KuntaiDu marked this conversation as resolved.
Outdated
llm.generate(
LONG_PROMPT + "# Question\nWhat' s the content in the (1,1) cells\n",
sampling_params=sampling_params
)

# Asking the content of cell (2,2)
# Will be faster than the last query
# as the computation of LONG_PROMPT can be reused
llm.generate(
LONG_PROMPT + "# Question\nWhat' s the content in the (2,2) cells\n",
sampling_params=sampling_params
)

```

## Example workloads

We describe two example workloads, where APC can provide huge performance benefit:
- Long document query, where the user repeatedly query the same long document (e.g. software mannual or annual report) with different queries. In this case, instead of processing the long document again and again, APC allows vllm to process this long document (and generate its corresponding KV cache) *only once*, and all future requests can avoid recomputing this long document by reusing its KV cache. This allows vllm to serve future requests with much higher throughput and much lower latency.
- Multi-round conversation, where the user may chat with the application multiple times in the same chatting session. In this case, instead of processing the whole chatting history again and again, APC allows vllm to reuse the processing results of the chat history across all future rounds of conversation, allowing vllm to serve future requests with much higher throughput and much lower latency.


APC in general does not reduce the performance of vllm. With that being said, APC only reduces the time of processing the queries (the prefilling phase) and does not reduce the time of generating new tokens (the decoding phase). So APC does not bring performance gain when vllm spends most of the time generating answers to the queries (e.g. when the length of the answer is long), or new queries do not share the same prefix with any of existing queries (so that the computation cannot be reused).
48 changes: 48 additions & 0 deletions docs/source/automatic_prefix_caching/details.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,48 @@

# How vllm implements APC

(This text is for vllm developers and it is a bit long, feel free to skip it.)
Comment thread
KuntaiDu marked this conversation as resolved.
Outdated

The core idea of PagedAttention is to partition the KV cache of each request into KV Blocks. Each block contains the attention keys and values for a fixed number of tokens. The PagedAttention algorithm allows these blocks to be stored in non-contiguous physical memory so that we can eliminate memory fragmentation by allocating the memory on demand.

To automatically cache the KV cache, we utilize the following key observation: Each KV block can be uniquely identified by the tokens within the block and the tokens in the prefix before the block.

Block 1 Block 2 Block 3
[A gentle breeze stirred] [the leaves as children] [laughed in the distance]
Block 1: |<--- block tokens ---->|
Block 2: |<------- prefix ------>| |<--- block tokens --->|
Block 3: |<------------------ prefix -------------------->| |<--- block tokens ---->|
Comment thread
KuntaiDu marked this conversation as resolved.


In the example above, the KV cache in the first block can be uniquely identified with the tokens “A gentle breeze stirred”. The third block can be uniquely identified with the tokens in the block “laughed in the distance”, along with the prefix tokens “A gentle breeze stirred the leaves as children”. Therefore, we can build the following one-to-one mapping:

hash(prefix tokens + block tokens) <--> KV Block

With this mapping, we can add another indirection in vLLM’s KV cache management. Previously, each sequence in vLLM maintained a mapping from their logical KV blocks to physical blocks. To achieve automatic caching of KV blocks, we map the logical KV blocks to their hash value and maintain a global hash table of all the physical blocks. In this way, all the KV blocks sharing the same hash value (e.g., shared prefix blocks across two requests) can be mapped to the same physical block and share the memory space.


This design achieves automatic prefix caching without the need of maintaining a tree structure among the KV blocks. More specifically, all of the blocks are independent of each other and can be allocated and freed by itself, which enables us to manages the KV cache as ordinary caches in operating system.
Generalizedable Caching Policy

Keeping all the KV blocks in a hash table enables vLLM to cache KV blocks from earlier requests to save memory and accelerate the computation of future requests. For example, if a new request shares the system prompt with the previous request, the KV cache of the shared prompt can directly be used for the new request without recomputation. However, the total KV cache space is limited and we have to decide which KV blocks to keep or evict when the cache is full.


# Cache eviction policy

Managing KV cache with a hash table allows us to implement flexible caching policies. As an example, in current vLLM, we implement the following eviction policy:

When there are no free blocks left, we will evict a KV block with reference count (i.e., number of current requests using the block) equals 0.
If there are multiple blocks with reference count equals to 0, we prioritize to evict the least recently used block (LRU).
If there are multiple blocks whose last access time are the same, we prioritize the evictionto evict of the block that is at the end of the longest prefixwith the maximum number of prefix tokens. (i.e., has the maximum number of blocks before it)

Note that this eviction policy effectively implements the exact policy as in RadixAttention when applied to models with full attention, which prioritizesprioritzes to evict reference count zero and least recent used leaf nodes in the prefix tree.

However, the hash-based KV cache management gives us the flexibility to handle more complicated serving scenarios and implement more complicated eviction policies beyond the policy above:

Multi-LoRA serving. When serving requests for multiple LoRA adapters, we can simply let the hash of each KV block to also include the LoRA ID the request is querying for to enable caching for all adapters. In this way, we can jointly manage the KV blocks for different adapters, which simplifies the system implementation and improves the global cache hit rate and efficiency.

Multi-modal models. When the user input includes more than just discrete tokens, we can use different hashing methods to handle the caching of inputs of different modalities. For example, perceptual hashing for images to cache similar input images.

Non-LRU caching policy. Since we can manage the KV cache as a standard cache in operating systems, most existing cache eviction policies can be used to manage the KV cache in vLLM. As an example, instead of LRU, vLLM can additionally record the total access count of each cache block and implement a least frequently used (LFU) eviction policy.

Application-specific caching policy. For some applications, we can use some application-specific information to further improve the eviction policy. For example, suppose we know one specific request is the start of a multi-round conversation. In that case, we can prioritize caching the KV blocks for the specific request for future conversation rounds.
Comment thread
KuntaiDu marked this conversation as resolved.
Outdated
7 changes: 7 additions & 0 deletions docs/source/index.rst
Original file line number Diff line number Diff line change
Expand Up @@ -105,6 +105,13 @@ Documentation
quantization/fp8_e5m2_kvcache
quantization/fp8_e4m3_kvcache

.. toctree::
:maxdepth: 1
:caption: Automatic prefix caching

automatic_prefix_caching/apc
automatic_prefix_caching/details

.. toctree::
:maxdepth: 2
:caption: Developer Documentation
Expand Down