[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

@@ -362,7 +362,7 @@ def test_swap_blocks(
    block_mapping = list(zip(src_blocks, dst_blocks))
    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

@@ -508,3 +523,19 @@ def get_num_cached_tokens(self, seq: Sequence) -> int:
        cached in the block manager for the sequence.
        """
        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

@@ -322,10 +322,10 @@ def prepare_worker_input(
        # `blocks_to_swap_in` and `blocks_to_swap_out` are cpu tensors.
        # 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

ApostaC requested review from tlrmchlsmth, WoosukKwon, zhuohan123, youkaichao, alexm-neuralmagic, 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

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Core] Support offloading KV cache to CPU #10874

[Core] Support offloading KV cache to CPU #10874

ApostaC commented Dec 3, 2024 •

edited by github-actions bot

Loading

github-actions bot commented Dec 3, 2024

KuntaiDu Dec 4, 2024

ApostaC Dec 4, 2024

lixiaobai09 Dec 6, 2024

ApostaC Dec 6, 2024

lixiaobai09 Dec 7, 2024

ApostaC Dec 8, 2024

mergify bot commented Dec 9, 2024

DearPlanet Dec 13, 2024

ApostaC Dec 13, 2024

comaniac Dec 17, 2024

youkaichao Dec 18, 2024

comaniac Dec 18, 2024

comaniac Dec 17, 2024

ApostaC Dec 18, 2024

comaniac Dec 17, 2024

KuntaiDu Dec 18, 2024

mergify bot commented Dec 17, 2024

ApostaC Dec 18, 2024

ApostaC Dec 18, 2024

[Core] Support offloading KV cache to CPU #10874

Are you sure you want to change the base?

[Core] Support offloading KV cache to CPU #10874

Conversation

ApostaC commented Dec 3, 2024 • edited by github-actions bot Loading

An implementation for CPU KV cache offloading (#7697)

End-to-end benchmarking results:

New kernel implementation microbenchmark

Potential improvement:

Implementation

github-actions bot commented Dec 3, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mergify bot commented Dec 9, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mergify bot commented Dec 17, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ApostaC commented Dec 3, 2024 •

edited by github-actions bot

Loading