Investigate PagedAttention KV-cache memory management for faster inference

[Azeirah]

New research just came out on using a technique inspired by kernel virtual memory and pages to manage the KV cache.

Results? Way faster inference!

https://vllm.ai/

They claim up to 24x the throughput (measured in requests handled per second) compared to huggingface's transformers library

How?

Inference is bottlenecked by memory, most notably the KV cache. They say the KV cache's most notable features are

• That it's very large
• That it's dynamic, size depends on sequence length which is variable. Existing systems waste 60-80% of memory due to fragmentation and over-reservation

PagedAttention is an alternative approach to managing the KV cache which is inspired by virtual memory, pages and blocks. By allocating the space dynamically with this approach, only up to about 4% of memory will be wasted, instead of the aforementioned 60-80.

For further details, better refer to their website and Github.

[JohannesGaessler]

llama.cpp currently only ever serves one user at a time so this optimization is not applicable.

[nivibilla]

I assume it would be useful if we want to host the models and have a interface like chat.openai.com?

[JohannesGaessler]

Yes, for enterprise use where you have one server generating responses for many users in parallel the optimization would be useful.

[Azeirah]

llama.cpp currently only ever serves one user at a time so this optimization is not applicable.

Oh I wasn't aware this was exclusively for a client-server application, that explains why they measure performance in requests/sec 🥲

[howard0su]

this optimization is still applicable as it can save vram usage of kv tensor.

[nivibilla]

If we do end up building this for server use and I think that would be a good idea. Then this paging system would be very useful.

[howard0su]

Read through the blog and the code. It turns out the paged attention is a way to manage the memory so that the compute kernel doesn't require kv have to be continues. This make it possible that you can have one prompt's kv append by multi output's KVs. like the following

Prompt KV Block ------ Output 1 KV Block
                            ------ Output 2 KV block
                              ....

This is super helpful if your prompt is long and you need to output multi results. This is a purely engineering trick. The change is mainly around the how we manage the KV in VRAM. If we are using CPU, this is even simpler to implement. (simple as list v.s. vector)

[slaren]

We allocate all the KV memory required for the maximum context length on startup in one block, so we shouldn't have any fragmentation either.

[randxie]

@JohannesGaessler Is serving multiple users concurrently or batch inference on the roadmap of llama.cpp?

[JohannesGaessler]

I don't have any plans for it because I don't care about commercial use but I can't speak for the other devs.

[okpatil4u]

Should it not be on the list ? Today we are talking about chatbots, in 6 months or so, people will start looking for autonomous agents. Would it not make sense to build a system that can process multiple requests simultaneously and efficiently ?

…

On Sun, 25 Jun 2023 at 4:20 PM, Rand Xie ***@***.***> wrote: @JohannesGaessler <https://github.com/JohannesGaessler> Is serving multiple users concurrently or batch inference on the roadmap of llama.cpp? — Reply to this email directly, view it on GitHub <#1955 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAXGU4HR4DKJ5ZMVHMQ2T7DXNAJWLANCNFSM6AAAAAAZN5MVXY> . You are receiving this because you are subscribed to this thread.Message ID: ***@***.***>

[nivibilla]

Yeah, I think first we need to solve batch inference. It's implemented in babyllama but I'm haven't tried to port it over to the main llama yet