[RFC]: KV cache offloading

[orozery]

Motivation.

Currently, in vLLM v1 there is no in-house solution for offloading KV cache data from the GPU memory to other medium (in particular, CPU memory).
There is a proposed RFC (#16144) and respective PRs (#13377 and #17653) that try to address that.
The approach they take is somewhat similar to the way offloading was implemented in V0:

1. On the scheduler side, extend the core GPU allocator (KVCacheManager) to support CPU offloading
2. On the worker side, add a synchronous call to handle the actual CPU<->GPU transfer in the execute_model function.

In this RFC I propose an alternative approach which supports the following requirements:

• Async saving of new KV data from GPU to cache. GPU memory will not be freed until save is completed (similar to NixlConnector)
• Async loading of KV data from the cache to GPU. Requests waiting for cache load won't be scheduled until load is completed (similar to NixlConnector)
• Support pluggable backends for cache (CPU backend for start)
• Allow pulling of cache events (cache insert, evict, access), in the same way as the GPU cache. This will allow a unified KVCacheEvent stream.
• Enable LRU eviction of offloaded data

Proposed Change.

I suggest we will enable offloading using a new OffloadingConnector, with minimal changes to vLLM's core.

On the scheduler side, this connector will delegate to an abstract OffloadingManager.
The OffloadingManager will be responsible for bookkeeping allocation and evictions of offloaded data. Its output will be encoded over KVConnectorMetadata and sent to workers.

On the worker side, we will also have an OffloadingConnector which parses load/store requests sent by the OffloadingManager and executed asynchronously using a set of worker threads, one per transfer type (e.g. 1 thread doing GPU->CPU, 1 for CPU->GPU, and in the future also 1 CPU->Disk, etc.).

Each transfer request submitted by the OffloadingManager will be responded by a unique job_id, which will be used to track completions.

To enable this design, we need 3 changes in vLLM's core:

1. PR [V1] - Enable worker -> scheduler connector metadata #19555 Introduce connector-metadata also in the direction of worker->scheduler (currently, only the scheduler connector can pass-on metadata to workers, but not the other way around).
2. PR v1: Add Request.block_hashes #19728 Introduce Request.block_hashes to allow the OffloadingConnector to re-use the block-hashes computed by the KVCacheManager.
3. PR v1: Support KV events from connectors #19737 Add a connector API for collecting KV cache events (cache insertion, deletion).

Aside from the above PRs with changes to vLLM's core, I already opened PR #19848 including the basis of pluggable offloading implementation.
On-top of this PR, there will be PRs for a concrete CPU offloading implementation (probably one PR for the scheduler side, and one PR for the worker side).
The last step will be to introduce the actual OffloadingConnector that will enable the e2e use of offloading in v1.

Feedback Period.

One week.

CC List.

@WoosukKwon @simon-mo @robertgshaw2-redhat @njhill

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.

[josephrocca]

Exciting! Do you have a branch with all the required changes merged to be able to test it? I'm particularly interested in testing it with DeepSeek on 8xH100 - i.e. VRAM constrained, hence prefix cache size/hits reduced when no offloading is available.

FWIW, I did test @chunxiaozheng's #17653 and got a 1.3x throughput (on real world "chat" use case, many concurrent users) when running RedHatAI/DeepSeek-R1-0528-quantized.w4a16 on 8xH100.

Script for #17653 reproduction, just for reference

cd /root
git clone --branch cpu-offload-v1-new2 https://github.com/chunxiaozheng/vllm.git
cd vllm
git reset --hard 12f1a3fbf5c715ca1d7154e9a833b6a290792912
VLLM_USE_PRECOMPILED=1 pip install --editable .
VLLM_WORKER_MULTIPROC_METHOD=spawn VLLM_MARLIN_USE_ATOMIC_ADD=1 python -m vllm.entrypoints.openai.api_server --host 0.0.0.0 --port 3000 --max-model-len 8192 --max-seq-len-to-capture 8192 --enable-chunked-prefill --enable-prefix-caching --trust-remote-code --disable-log-requests --tensor-parallel-size 8 --gpu-memory-utilization 0.95 --served-model-name deepseek-chat --model RedHatAI/DeepSeek-R1-0528-quantized.w4a16 --enable-kvcache-cpu-offloading --swap-space 60

[orozery]

Exciting! Do you have a branch with all the required changes merged to be able to test it? I'm particularly interested in testing it with DeepSeek on 8xH100 - i.e. VRAM constrained, hence prefix cache size/hits reduced when no offloading is available.

Not yet, but I will sure keep you posted.
I would be also very interested to get performance numbers.

[njhill]

Thanks @orozery. So I understand better, this would be a concrete implementation of a KV connector using the existing interface (with the other additions to that which you're proposing)? I like the "minimal changes to vllm core" in principle.

A few general thoughts:

• For a hierarchical cache we likely want the GPU->CPU saving to happen at the time of eviction rather than initial computation or else we're unnecessarily mirroring the GPU cache in CPU mem. I think this would require extending the connector API.
• Cached CPU blocks would't necessarily have all preceding blocks also in the CPU cache, some prefix might only reside in the GPU cache. For a given cache tier, we may some length threshold below which it's not worth retrieving (and therefore possibly also not worth storing) since it would be faster to recompute.
• An important question I think is whether these other cache tiers might be shared between DP GPU instances. If so it might change the first bullet assumption.

[orozery]

• For a hierarchical cache we likely want the GPU->CPU saving to happen at the time of eviction rather than initial computation or else we're unnecessarily mirroring the GPU cache in CPU mem. I think this would require extending the connector API.

My initial thoughts were that num_cpu_blocks would be much larger than num_gpu_blocks, and in that case it would be OK to have the GPU cache duplicated to CPU.
Looking it up, I understand this may not be the case, and CPU size could be roughly the same size as the GPU, or larger by a small factor.
So you make up a good point.
I think that design-wise, this does not effect the OffloadingManager API and downwards to the workers.
Instead, this effects only the OffloadingConnector (which is responsible for invoking prepare_store).
But more importantly, since we want this to be async, this will require changes to vllm's core, allowing the connector to be notified on "soon-to-be-evicted" GPU blocks.
Do you agree?
Concretely, we can extend KVCacheManager.allocate_slots to return a list of block IDs that are not yet evicted, but are next-in-line.
Then, passing this information to the connector, it can kick-off an offloading op.

• Cached CPU blocks would't necessarily have all preceding blocks also in the CPU cache, some prefix might only reside in the GPU cache. For a given cache tier, we may some length threshold below which it's not worth retrieving (and therefore possibly also not worth storing) since it would be faster to recompute.

What about having only full sequences offloaded to CPU, not just the tails that spill off the GPU cache?
I agree that in any case we should optimize to avoid inefficient IOs.

• An important question I think is whether these other cache tiers might be shared between DP GPU instances. If so it might change the first bullet assumption.

In DP you have a scheduler per DP-rank. Our cache state will thus be disjoint per each rank.
How would you suggest to share cache-state between DP-ranks?
Don't you think this actually falls outside of N/S scope? (E/W)