[Feature] Question about kvcache offloading to disk or CPU memory

[Yang-x-Zhao]

Motivation

In sglang's paper, it mentions that sglang plans to implement a kvcache offloading mechanism. I notice that there is a disable_disk_cache=False option in sglang's Server args, and wonder if this option controls the kvcache offloading mechanism.

Related resources

No response

[merrymercy]

disable_disk_cache is for caching other things like the regex conversion. The kvcache offloading mechanism has not been implemented and we plan to support it in the next few months.

[AllenHW]

Hi @merrymercy - I recently benchmarked sglang and was impressed with its throughput for offline generations. We're a startup called Outerport (https://www.outerport.com/), exploring large scale data processing use cases of LLMs

I'm pretty invested in exploring offloading KV cache to CPU. In fact we built a rust process that maintains a highspeed CPU cache to store arbitrary dictionary of tensors (which can be used to keep model weights, or kv cache). Being able to do that in sglang sounds like a promising direction. Is there a thread around this topic I can participate?

[xiezhq-hermann]

Hi @AllenHW, we currently have an implementation of hierarchical caching implemented on Python and we will keep you posted about the release and there will be an open thread for it soon. Would definitively be interested in learning more from you on high performance I/O management. Is there any code / document that can serve as a reference already for the work mentioned?

[AllenHW]

Awesome, yes would love to.

we have not open sourced the code yet, but on a high level:

• the rust process acquires RAM memory from the OS and builds a page table. this becomes the cache.
• when offloading to cpu, client process (Python) sends a list of cuda allocations (managed by pytorch's memory allocator, can be non-contiguous) and the tensors contained in each block
• the rust process takes out the needed pages and begins a series of cudaMemcpy calls.
• when loading from CPU to GPU, the client allocates the gpu memory and sends ipc handle to the rust process
• rust process opens the ipc memory and begins transferring data

For kv cache, a separate process probably does not offer much? We designed the system to be able to send and receive models from a remote registry.

I'm still ramping up with sglang and llm inference in general but a few thoughts:

• cpu <-> gpu bandwidth is around 11-12GB/s on PCIe 4.0 and 2x that on PCIe 5.0. there is probably a point (for short sequences) where recomputing kv cache is faster than the over head loading from CPU
• but the nice thing about loading kv cache from cpu is that it can happen concurrently with kv cache computation (https://arxiv.org/pdf/2406.19707)
• it gets intricate, but being able to schedule kv cache loading from cpu and recomputing kv cache concurrently could be interesting for the offline cases.