[Core] Support offloading KV cache to CPU #10874

ApostaC · 2024-12-03T20:52:05Z

An implementation for CPU KV cache offloading (#7697)

TL; DR: CPU offloading is better than prefix caching in our benchmark, we also found that the evictor can be optimized to save 10-30% of the runtime.

This PR is for fixing the DCO issue for the Kuntai's original CPU offloading PR #9682 . It also contains new CUDA kernels to improve the KV cache offloading performance.

End-to-end benchmarking results:

A long document QA workload (see benchmarks/benchmark_long_document_qa.py) running on A100-40G-SXM GPU. The GPU can cache 8 documents and the CPU can cache 30 documents.

(Following are the original data for the above figure)

Num documents	vLLM	vLLM w/ prefix caching	vLLM w/ prefix caching + CPU offloading
8	13.66	0.49	0.5
16	27.28	7.22	2.3
32	54.54	49.96	17.26
64	109.27	126.08	110.96

New kernel implementation microbenchmark

The numbers are collected on A100-40GB-SXM GPUs

# of pages	New implementation	Old implementation
Swap in 1 page	1.21 GB / second	1.21 GB / second
Swap in 1250 pages	12.6 GB / second	3.11 GB / second
Swap out 1 page	1.22 GB / second	1.21 GB / second
Swap out 1250 pages	12.5 GB / second	3.12 GB / seocond

The new kernel can achieve 4x better throughput than the old swap_block implementation.
Also, it won't decrease the performance when the number of pages are small.

Potential improvement:

Currently, the swap_block is invoked once per layer. If we can aggregate the copy of all the layers into one kernel, the throughput of copying 1 page will also achieve >10GB/s.

Implementation

This PR has much less features compared to #8694, but it is really minimum and creates very little core change. So I guess we can use this PR to enable CPU KV cache offloading first, and then focus on disk.

The key idea of this implementation is to maintain those allocated blocks that didn't hit the cache, and constantly copy them into CPU after each scheduler step.

Here is the flow diagram

This idea is borrowed from ConServe (paper link: https://arxiv.org/abs/2410.01228), based on the assumption that the CPU-GPU bandwidth is much higher than GPU KV cache generation throughput. Thanks Yifan for this idea.

github-actions · 2024-12-03T20:52:20Z

👋 Hi! Thank you for contributing to the vLLM project.
Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can do one of these:

Add ready label to the PR
Enable auto-merge.

🚀

Signed-off-by: ApostaC <[email protected]> Co-authored-by: KuntaiDu <[email protected]>

Signed-off-by: ApostaC <[email protected]>

…ssues Signed-off-by: ApostaC <[email protected]>

Signed-off-by: ApostaC <[email protected]>

KuntaiDu · 2024-12-04T20:50:54Z

tests/kernels/test_cache.py

    block_mapping_tensor = torch.tensor(block_mapping,
                                        dtype=torch.int64,
-                                        device="cpu").view(-1, 2)
+                                        device=device).view(-1, 2)


Is this because that this tensor need to be accessed by the new CUDA memcpy kernel?

Yes. The new paged_copy kernel need to access the block mapping from GPU.

lixiaobai09 · 2024-12-06T08:54:16Z

vllm/core/block_manager.py

        """
        return self._computed_blocks_tracker.get_num_cached_tokens(seq)
+
+    def get_and_reset_swaps(self,


This function seems like not get the real physical block ID from get_physical_block_id? Especially for CPU PrefixCachingBlockAllocator, whose start ID is not zero.

If I understand correctly, this function should not return the physical block IDs because the get_physical_block_id will be called in block_manager.swap_in()/block_manager.swap_out() later.

Call chain is: scheduler._swap_in() --> block_manager.swap_in() --> block_allocator.get_physical_block_id() (similar for swapping out).

(Let me know if my understanding is incorrect and I will fix it asap, thanks!)

lixiaobai09 · 2024-12-07T08:41:44Z

vllm/core/scheduler.py

+
+        # NOTE(Kuntai): extend the swapping list for CPU offloading
+        new_swap_out, new_swap_in = \
+                self.block_manager.get_and_reset_swaps(time.time())


However, the get_and_reset_swaps is called directly here without get_physical_block_id. I think these block IDs are sent to the cache engine directly later, so not the real physical block IDs.

Got it! I double-checked the logic and you are right. Just pushed another commit to fix the issue and update the docstring. Thanks for the catch!

Signed-off-by: ApostaC <[email protected]>

mergify · 2024-12-09T01:29:31Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @ApostaC.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

Signed-off-by: ApostaC <[email protected]>

DearPlanet · 2024-12-13T02:57:33Z

vllm/core/block/cpu_offloading_block_allocator.py

+          uncached: allocated blocks that didn't hit any cache
+          cached: allocated blocks that are cached, either in GPU or in CPU
+          free: the blocks are not allocated by block allocator
+        This implementation aims to transform uncacherd blocks to cached blocks


uncacherd blocks -> uncached blocks?

Thanks for the catch! Fixed!

Signed-off-by: ApostaC <[email protected]>

comaniac · 2024-12-17T19:29:59Z

vllm/engine/arg_utils.py

+        parser.add_argument(
+            '--block-allocator',
+            type=str,
+            default='CpuGpuBlockAllocator',
+            choices=['CpuGpuBlockAllocator', 'CpuOffloadingBlockAllocator'],
+            help='The block allocator that vLLM uses. Currently'
+            ' can be CpuGpuBlockAllocator (the default) and '
+            'CpuOffloadingBlockAllocator (experimental) that '
+            'supports offloading the KV cache to CPU . '
+            'When using CpuOffloadingBlockAllocator, the '
+            'preemption mode must be recompute.')


Can we avoid to expose block allocator? Instead we should provide something like --kv-cache-offloading.

After a second thought, for the sake of modularity and extensibility, I think it's fine to expose a block allocator argument, but we should keep the default name concise, and also keep the possibility to accept third-party allocators.

I recommend adding --allocator, which can take several default values as above, and also a class qualname like mod.cls .

How does this interface work with v1?

comaniac · 2024-12-17T19:30:29Z

vllm/worker/worker.py

        # they contain parameters to launch cudamemcpyasync.
        blocks_to_swap_in = torch.tensor(execute_model_req.blocks_to_swap_in,
-                                         device="cpu",
+                                         device="cuda",


This won't be compatible with other devices I think?

As suggested by Kaichao, we will use a static pinned CPU memory buffer to host this array.

comaniac · 2024-12-17T19:32:17Z

benchmarks/benchmark_long_document_qa.py

Can we separate the benchmark scripts to another PR to reduce the size of this one?

mergify · 2024-12-17T19:48:54Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @ApostaC.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

ApostaC · 2024-12-18T00:00:14Z

csrc/cache_kernels.cu

+    return static_cast<T*>(tensor.data_ptr());
+  } else if (device.is_cpu() && tensor.is_pinned()) {
+    T* ptr;
+    cudaHostGetDevicePointer((void**)&ptr, static_cast<T*>(tensor.data_ptr()),


Save this pointer at the creation of the CPU memory allocation -- making it CUDA graph compatible.

Potentially in a new PR

lixiaobai09 · 2024-12-26T02:12:23Z

vllm/core/block/cpu_offloading_block_allocator.py

+            block_id = block.block_id
+
+            # check if this block is freed
+            if block_id is None:


Because the block structure in _uncached_blocks is from a block memory pool, the block can be freed and allocated as a new block_id. However, while the block_id is allocated as a new block_id, this may cause some mistakes (or hurt performance). So, maybe we can remove the block in _uncached_blocks while block_manager calls the free function in cpu_offloading_block_allocator?

I see what you mean. This makes sense. Let me think what's the best way to fix it. Iterating over the _uncache_blocks during free() will be too expensive IMO.

@lixiaobai09 Hey, I've fixed this. Please see #11532

ApostaC · 2024-12-26T23:25:26Z

As suggested by @simon-mo and @youkaichao , I split this PR into 3 smaller PRs to make it easier to review.

CUDA kernel: [Core] Performance optimization for swap_blocks by cuda kernels #11531
CPU offloading block allocator: [Core] Block Allocator to support KV cache CPU offloading #11532
Benchmark script: [Benchmark] Add benchmark script for CPU offloading #11533

chunxiaozheng · 2025-03-10T09:11:12Z

hi~ @ApostaC,I have some questions to ask. When will the offloading be triggered from GPU to CPU？
From the code, it depends on _uncached_blocks, but the content in _uncached_blocks will be free after the sequence inference is completed. Therefore, only blocks are hit (computed is true) will trigger the offloading from GPU to CPU?

hmellor · 2025-04-02T20:37:37Z

After speaking to @ApostaC I'm closing this as it only optimises the V0 engine. There is work in progress for a similar PR in V1.

ApostaC requested review from tlrmchlsmth, WoosukKwon, zhuohan123, youkaichao, alexm-redhat, comaniac and njhill as code owners December 3, 2024 20:52

mergify bot added the frontend label Dec 3, 2024

KuntaiDu mentioned this pull request Dec 3, 2024

[Core] Support offloading KV cache to CPU #9682

Closed

ApostaC and others added 2 commits December 3, 2024 20:59

Move to a new branch to fix the DCO issues.

f60a8fa

Signed-off-by: ApostaC <[email protected]> Co-authored-by: KuntaiDu <[email protected]>

[Fix] the failed unit tests

e6654f2

Signed-off-by: ApostaC <[email protected]>

ApostaC force-pushed the yihua-cpu-offloading2 branch from da35ed9 to e6654f2 Compare December 3, 2024 21:00

ApostaC added 2 commits December 4, 2024 06:07

[Fix] CPU offloading not working bug and [fix] unit test and format i…

ba6c9e3

…ssues Signed-off-by: ApostaC <[email protected]>

[fix] broken tests for cpu offloading allocator

1c94985

Signed-off-by: ApostaC <[email protected]>

KuntaiDu reviewed Dec 4, 2024

View reviewed changes

KuntaiDu added ready ONLY add when PR is ready to merge/full CI is needed and removed ready ONLY add when PR is ready to merge/full CI is needed labels Dec 5, 2024

lixiaobai09 reviewed Dec 6, 2024

View reviewed changes

lixiaobai09 reviewed Dec 7, 2024

View reviewed changes

ApostaC added 2 commits December 8, 2024 21:15

[Fix] add the call to get_physical_block_ids

daab0d6

Signed-off-by: ApostaC <[email protected]>

[Add] faster unsafe implementation for get_physical_block_id

919e5e3

Signed-off-by: ApostaC <[email protected]>

mergify bot added the needs-rebase label Dec 9, 2024

Merge branch 'main' into yihua-cpu-offloading2

52185bf

mergify bot removed the needs-rebase label Dec 9, 2024

ApostaC added 3 commits December 10, 2024 17:37

Updating the benchmark script with correct usage instructions

505e60c

Signed-off-by: ApostaC <[email protected]>

make yapf happy

a517a29

Signed-off-by: ApostaC <[email protected]>

fix format checker issues

789b00e

Signed-off-by: ApostaC <[email protected]>

lixiaobai09 mentioned this pull request Dec 11, 2024

[Core] Efficient transmission for CPU prefix caching, based on PR#10874 #11099

Open

KuntaiDu added the ready ONLY add when PR is ready to merge/full CI is needed label Dec 11, 2024

DearPlanet reviewed Dec 13, 2024

View reviewed changes

ApostaC force-pushed the yihua-cpu-offloading2 branch 2 times, most recently from 583984f to 789b00e Compare December 13, 2024 16:41

ApostaC added 2 commits December 13, 2024 16:50

Merge remote-tracking branch 'upstream/main' into yihua-cpu-offloading2

a0b5061

Fix the compatibility witht the latest main

2648fa5

Signed-off-by: ApostaC <[email protected]>

ApostaC force-pushed the yihua-cpu-offloading2 branch from 9fcf23e to 2648fa5 Compare December 13, 2024 17:12

Fix the typo uncacherd -> uncached

278e166

Signed-off-by: ApostaC <[email protected]>

comaniac reviewed Dec 17, 2024

View reviewed changes

mergify bot added the needs-rebase label Dec 17, 2024

ApostaC commented Dec 18, 2024

View reviewed changes

lixiaobai09 reviewed Dec 26, 2024

View reviewed changes

This was referenced Dec 26, 2024

[Core] Performance optimization for swap_blocks by cuda kernels #11531

Closed

[Core] Block Allocator to support KV cache CPU offloading #11532

Closed

[Benchmark] Add benchmark script for CPU offloading #11533

Merged

zeroorhero mentioned this pull request Mar 1, 2025

[Bug]: Runtime error when running MLA models with "prefix caching enabled" and "chunked prefill disabled" #14069

Closed

1 task

hmellor closed this Apr 2, 2025

Uh oh!

[Core] Support offloading KV cache to CPU #10874

[Core] Support offloading KV cache to CPU #10874

Uh oh!

Conversation

ApostaC commented Dec 3, 2024 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

An implementation for CPU KV cache offloading (#7697)

End-to-end benchmarking results:

New kernel implementation microbenchmark

Potential improvement:

Implementation

Uh oh!

github-actions bot commented Dec 3, 2024

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mergify bot commented Dec 9, 2024

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mergify bot commented Dec 17, 2024

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ApostaC commented Dec 26, 2024

Uh oh!

chunxiaozheng commented Mar 10, 2025

Uh oh!

hmellor commented Apr 2, 2025

Uh oh!

Uh oh!

ApostaC commented Dec 3, 2024 •

edited by github-actions bot

Loading