[RFC]: Sparse KV cache management framework

[ShawnD200]

Motivation.

Efficient use of KV cache not only saves scarce GPU memory but also accelerates speed. Evicting potentially many low-utilization tokens, especially in long context, is the way to do it. To this end, this RFC proposes a flexible and performant KV cache management framework.

Proposed Change.

The design is as the following figure shows:

(1) CachePolicy is a general interface that manages KV cache allocations. It generalizes from basic cache behavior that always allocates new slots/blocks to append new tokens to the more sophisticated that dictates how cache space is used should there be new tokens, such as sliding window, H2O, FastGen, etc.

(2) CachePolicy (modified from previous BlockTable) has two straightforward methods: add_tokens_prefill and add_tokens_decode to give slots/blocks to new tokens. The given slots/blocks may be new allocations, and if so they will be requested from allocator. Or they can be old ones, which are from the previous blocks at its disposal. It is the policy that selects victims, evicts them, and replaces them with the new ones.

(3) To do so, two data structures are employed to keep track of blocks and tokens: a PhysicalBlockTable and a VirtualBlockTable, and maintains these two to manage the cache space. The PhysicalBlockTable (previously BlockList) holds the physical block IDs, while the VirtualBlockTable maps tokens to blocks and backward with slot mappings and token mappings respectively.

(4) Implemented first two concrete policies which are used today, the defaultCachePolicyBase and CachePolicySlidingWindow, which keeps context of sliding window. Specifically, it is reimplemented such that the sliding window takes effect in prefill and is adhered to in decode, in which never a new block will be allocated should a window of blocks have been had.

(5) Per layer cache management capability.

(6) Attention weights from attention backend.

This work has a draft PR #11928. It was worked on V0 and will be revamped on V1 in the following weeks.

Tasks:

1. Unify add_tokens_prefill and add_tokens_decode (which was suggested by @comaniac)
2. Abstract new cache policy interface and develop the supporting data structures.
3. Wait for and coordinate with Hybrid Memory Allocator [RFC]: Hybrid Memory Allocator #11382
4. Getting attention weights from backend.

Feedback Period.

2 weeks.

There are efforts already there, like #10646 and #10942. Welcome to discuss and this work is open and can be adjusted by existing work.

CC List.

@comaniac @simon-mo @WoosukKwon @youkaichao @heheda12345 @YaoJiayi @Zachary-ai-engineer

Any Other Things.

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[heheda12345]

Agree that we need some better abstraction on the block allocation logic of different types of KV cache. I'm trying to make the abstraction right now for the hybrid allocator and we can have some discussion about it. (My initial attempt is here in #11938 , and I'm thinking of whether there is a better interface). But I'm not sure whether we need to introduce a token-slot mapping.

For sliding window, we have the opportunity to support prefix caching, so I'm implementing allocating new blocks for new tokens and free old blocks that go out of sliding window simutaneously. The block_table will be [-1, -1, -1, 4, 5] where -1 are freed blocks, 4&5 are blocks in use. And the token-slot mapping is exactly the same with standard attention.

For token drop like H2O, I don't have a good idea for prefix caching. Without prefix caching, I think we can just allocate min(num_tokens, num_tokens_to_keep) slots. We do not need to maintain the exact mapping between slot and tokens. It is correct to put the kv cache of a new token to any of the free slots.

I haven't investigate how to support FastGen, and I'm happy to discuss about that.

[ShawnD200]

Thanks for your reply.

For token drop like H2O, I don't have a good idea for prefix caching. Without prefix caching, I think we can just allocate min(num_tokens, num_tokens_to_keep) slots. We do not need to maintain the exact mapping between slot and tokens. It is correct to put the kv cache of a new token to any of the free slots.

For sliding window, mapping between slots and tokens are not required due to that mapping can be reconstructed at any time. For H2O you mentioned, I could be terribly wrong, but I don't see how it works without mapping, and you could have it either by explicit mapping or making block_size=1 so that block ids can be arranged to be a mapping as well.

Prefix caching adds a layer of complexity, for "full attention" and "sliding window" it does not affect much, but for "random" selective eviction policy, either have block_size=1, or the block after token in it evicted can't be shared.

[heheda12345]

Based on my understanding, BlockManager only needs to know the number of slot needed by the model if we don't plan to implement prefix caching of these token drop models. The exact mapping between slot and token, and the intermediate state needed for selecting which token to drop can be maintained by some buffer in the model side. These state can be treated as another KV cache layer with very small hidden size in the BlockManager side. I think it would be more flexible than modifying block manager for each token drop method.

[ShawnD200]

Based on my understanding, BlockManager only needs to know the number of slot needed by the model if we don't plan to implement prefix caching of these token drop models.

BlockManager (V1 does not have it, perhaps the backend) assumes the kv is a list of consecutive blocks as a sequence (of seq_len), it is correct that mapping is not needed in this regard. But it must be a "sequence", meaning no holes (if you may drop any), albeit there might be other ways to allow holes.

In my abstraction, we don't make holes after eviction by replacing the eviction with new tokens, and then give the attn backend the mapping to copy new kv to cache.

The exact mapping between slot and token, and the intermediate state needed for selecting which token to drop can be maintained by some buffer in the model side. These state can be treated as another KV cache layer with very small hidden size in the BlockManager side. I think it would be more flexible than modifying block manager for each token drop method.

Model side? what change?

[heheda12345]

I feel that filling the holes with new tokens can be achieved by a special Attention layer called H2OAttention instead of modifying the block manager.

[ShawnD200]

I feel that filling the holes with new tokens can be achieved by a special Attention layer called H2OAttention instead of modifying the block manager.

Never heard such an attention layer can ever exist or be created.

[heheda12345]

What is the blocker of supporting hole-filling by creating a H2OAttention layer?

[ShawnD200]

What is the blocker of supporting hole-filling by creating a H2OAttention layer?

Correct me if I am wrong:

1. an attention layer in the model, which manipulates cache slots for itself and/or for the other layers?
2. Even if you can do it, it does not mean it should be done without the kv manager knowing it (how the kv manager and the 'H2OAttention' layer cooperate with each other?)

Admittedly, if cache compression is done in attention layer themselves, it has performance benefit. This would be a totally different story.

Is this what you really meant?

[heheda12345]

It seems that each H2O attention layer has different slot mapping (correct me if it is wrong), so I think it is natural to ask each layer to maintain its own slot mapping. block manager does not need to know the exact slot mapping, and I think we can add an interface to let the model to pass the number of slot it needs to the block manager. In this way, we can minimize the change to the block manager and still keep the flexibility of supporting different compression algorithms.

[ShawnD200]

It seems that each H2O attention layer has different slot mapping (correct me if it is wrong), so I think it is natural to ask each layer to maintain its own slot mapping. block manager does not need to know the exact slot mapping, and I think we can add an interface to let the model to pass the number of slot it needs to the block manager. In this way, we can minimize the change to the block manager and still keep the flexibility of supporting different compression algorithms.

Then it comes down to: change block manager vs. attention module, this seems not a question, from single responsibility (and single place -- there is only one block manager but a few backends) point of view.