Repository: pytorch
Issue: [RFC] PagedAttention Support
Labels: triaged,module: sdpa

### Feature request

PagedAttention has been a mainstream optimization technology for generation task based on LLMs.  It has been supported by a lot of server engines, e.g., [vllm](https://github.com/vllm-project/vllm), [tensorrt-llm](https://github.com/NVIDIA/TensorRT-LLM/blob/release/0.5.0/tensorrt_llm/runtime/kv_cache_manager.py) and [TGI](https://github.com/huggingface/text-generation-inference). To get OOB performance with PyTorch+Huggingface, in this RFC, we plan to propose the PageAttention KV Cache design which use compatible semantic as the HF KV Cache. 
  
### Motivation

KV cache is used to reduce computation for _**Decoder**_ layers, but it also brings memory overheads.  In general, the memory overheads are introduced by four parts. 
1. Firstly, in the attention module, the past key/value token states should be concated with the current token to get the entire context which will introduce additional memory read/write. 
2. Secondly, for the multiple sampling, e.g., beam search, the KV cache should be reordered according to latest beam index.  When the sequence is very long, the memory overheads of these two parts will be performance bottleneck. 
3. Thirdly, for batching inputs, all sequence will be generally padded to the longest sequence in this batch. The native KV cache in the transformers uses contiguous memory buffer to store KV cache and the shape is (batch, head_num, past_seq_len, head_dim), it means that the memory consumption of the KV cache is decided by the longest sequence. 
4. Lastly, for multiple sampling, there is only one prompts for different beams, and it can be shared across different beams, but the native implementation needs to store multiple copies of prompts tokens.  This will obviously waste a lot of memory and limits the max batch size for throughput mode.  

To address the above limitations, this [paper](https://arxiv.org/pdf/2309.06180.pdf) proposes PagedAtention, an attention algorithm inspired by the operating system’s (OS) solution to memory fragmentation and sharing: virtual memory with paging. PagedAttention divides KV cache into blocks, each of which can contain the attention keys and values of a fixed number of tokens. In PagedAttention, the blocks for the KV cache are not necessarily stored in contiguous space. Therefore, the KV cache can be managed in a more flexible way as in OS’s virtual memory: one can think of blocks as pages, tokens as bytes, and requests as processes. This design alleviates internal fragmentation by using relatively small blocks and allocating them on demand. Moreover, it eliminates external fragmentation as all blocks have the same size. Finally, it enables memory sharing at the granularity of a block, across the different sequences associated with the same request or even cross the different requests.

### Huggingface KV Cache

The HF already support 3 kinds of  [Cache ](https://github.com/huggingface/transformers/blob/main/src/transformers/cache_utils.py) implementation. DynamicCache, StaticCache and SinkCache.  Firstly, the modeling Attention layer uses the Cache.upate() function to cache the key/value of the curent step into KV Cache and get the key/value with entire context tokens states and then Cache.reorder_cache() is used to update the KV Cache if beam search is used.  In general, the attention layer will use _torch.nn.functional.scaled_dot_product_attention_ to do the the Scale Dot Product Attention (SDPA), e.g. the [LlamaSdpaAttention](https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L653).  The workflow of these KV Cache as following:  
![image](https://github.com/pytorch/pytorch/assets/46986936/bc0e68a7-657a-4b91-857a-02b7f7b12a3e)
![image](https://github.com/pytorch/pytorch/assets/46986936/81d34d8d-9ec3-45bc-8162-25deb4ffeab4)


### PagedAttention KV Cache Design

The native SDPA operators assume that the neighbor tokens are stored contiguously in the memory, but the PagedAttention partition the sequence into multiple blocks and these blocks are stored discretely. So the PagedAttention can not co-work with the current SPDA operator.  In this proposal, we want enable the PagedAttention in PyTorch/TorchAO. The PagedAttention KV Cache should be semantic compatible with HF KV cache abstraction and functionality compatible with _torch.compile_. To achieve above tow targets, we provide two components in this design. 

1. PagedAttentionCache 

The PagedAttentionCache is responsible to manage the allocate and free of physical space for every sequence. Due to the logical blocks of every sequence is discretely store in physical blocks, the PagedAttentionCache also need to maintain a mapping table from the logical block to every physical block. Besides, the PagedAttentionCache also provide the copy-on-write mechanism to enable the prompt sharing for beam search.  

2. PagedTensor 

PagedTensor is a tensor wrapper which include multiple pre-allocated cache tensors of every Decoder layers and other meta info(e.g., block mapping table) information. To get semantic compatible with _torch.nn.functional.scaled_dot_product_attention_ , we also rewrite this method to do SDPA based on the discrete paged memory and Attention layer do not need to aware these PagedAttention related meta information with this design. 

The PagedAttentionCache includes a lot python logic which is not friendly to _torch.compile_ and  there are many repeated Decoder layers for LLM models.   With the above design, The allocate can be only executed once before the model forward which can reduce the allocate overheads can avoid the graph compile issue. The workflow of our design as following figure. So, you can see that the workflow of the PagedAttentionCache is similar to other HF KV Caches except for the _Cache.allocate()_ before model forward and _allocate_ is used to allocate the free slots to store the key/value in pre-allocate memory 

![image](https://github.com/pytorch/pytorch/assets/46986936/b0678e9f-1d1f-4517-b405-86fe902917a8)

 #### Example 
```python 
# Initialize PagedAttentionCache
pagedcache = PagedAttentionCache(
    num_blocks, block_size, num_key_value_heads, head_dim, num_layers, device, dtype
)

key_states = torch.randn(batch_size, num_key_value_heads, prompt_len, head_dim)
value_states = torch.randn(batch_size, num_key_value_heads, prompt_len, head_dim)
query = torch.randn(batch_size, prompt_len, num_query_heads, head_dim)
# allocate slots before updating
pagedcache.allocate(batch_size, prompt_len)
# cache new tokens into kv cache
updated_key, updated_value = pagedcache.update(key_states, value_states, 1)
# Scaled dot product attention
output = torch.nn.functional.scaled_dot_product_attention(
    query.transpose(1, 2), updated_key, updated_value, scale=0.125
)
```