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

ApostaC · 2024-12-26T23:15:38Z

Note: This PR is part of the big CPU offloading PR #10874 -- this PR contains the CPU-offloading block allocator implementation as well as the changes in the scheduler.

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.

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

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.

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

github-actions · 2024-12-26T23:15:52Z

👋 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]>

youkaichao · 2024-12-30T05:41:37Z

I will hand it over to @comaniac for final reviews.

comaniac · 2024-12-30T20:29:35Z

@ApostaC could you take a look at #11385 and see if it's related?

ApostaC · 2024-12-31T20:53:52Z

@ApostaC could you take a look at #11385 and see if it's related?

I've gone through that PR. Seems like that PR implements a similar functionality of CPU offloading, but not sure what the performance will be.

By the way, is the implementation (where offloading the KV during the execution of model runner) of that PR duplicated with Kuntai's previous disaggregated prefill PR (#10502)?
@KuntaiDu Please chime in if you have more understanding about this, thanks.

DearPlanet · 2025-01-03T09:59:46Z

@ApostaC I tried this pr with flashinfer backend, but get wrong decoding results after running serveral requests (maybe 100 to 200), I have no idea about how to trace the error.
It works well on xformers/default attention backend.

KuntaiDu · 2025-01-05T13:22:50Z

@ApostaC could you take a look at #11385 and see if it's related?

I've gone through that PR. Seems like that PR implements a similar functionality of CPU offloading, but not sure what the performance will be.

By the way, is the implementation (where offloading the KV during the execution of model runner) of that PR duplicated with Kuntai's previous disaggregated prefill PR (#10502)? @KuntaiDu Please chime in if you have more understanding about this, thanks.

For PR #11385, it is essentially "sending" KV cache to the CPU pool after prefill and "receiving" KV cache from the CPU pool before prefill, so the abstractions exposed by disaggregated prefill can help that PR handle all the control-plane stuff.

ApostaC · 2025-01-06T15:20:20Z

@ApostaC I tried this pr with flashinfer backend, but get wrong decoding results after running serveral requests (maybe 100 to 200), I have no idea about how to trace the error. It works well on xformers/default attention backend.

Hey @DearPlanet , can you give some basic scripts to help reproduce the problem? This will be very helpful to debug.

DearPlanet · 2025-01-07T08:18:46Z

@ApostaC , here is a simple reproducing process, the commands below executed on RTX3090x2:

Start service:

VLLM_ATTENTION_BACKEND=FLASHINFER CUDA_VISIBLE_DEVICES=0,1 vllm serve /host/models/Qwen2.5-32B-Instruct-AWQ/ --served-model-name qwen2.5-32b --enable-prefix-caching --block-allocator CpuOffloadingBlockAllocator --preemption_mode recomputation --swap-space 25  --tensor-parallel-size 2 --host 0.0.0.0 --port 8080 --gpu-memory-utilization 0.65 --max-model-len 3000

Run benchmark script at vllm/benchmark/ :

python3 benchmark_serving.py --base-url http://0.0.0.0:8080 --dataset-path ./sonnet.txt --model qwen2.5-32b --tokenizer /mnt/root/models/Qwen2.5-32B-Instruct-AWQ/ --request-rate 3 --backend openai-chat --endpoint /v1/chat/completions --dataset-name sonnet

Print output content of responses, then you can see the abnormal decoding results.

I tried with default/xformers/flashinfer backends:
3090x2: default✅xformers✅ flashinfer❌
L20x2(with fp8 kv cache): default❓xformers❌flashinfer❌

The correct output log file:
test_out_sonnet_default.log

The error output log file:
test_out_sonnet_flashinfer.log

boposki · 2025-01-17T07:59:23Z

I think there is a bug in _uncached_blocks. If a block is stored in _uncached_blocks, the block will be released before it is saved to the CPU after the inference is complete. Other requests will reuse the block, causing the block ID to be rewritten.

boposki · 2025-01-17T08:35:55Z

I think there is a bug in _uncached_blocks. If a block is stored in _uncached_blocks, the block will be released before it is saved to the CPU after the inference is complete. Other requests will reuse the block, causing the block ID to be rewritten.
sorry， i found the problem is solved。 code： self._uncached_blocks.remove(block_id)

mergify · 2025-02-28T08:54:39Z

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

maobaolong

@ApostaC Thanks for this great feature, we will test it in our test env.

BTW, i left a bundle of minor comment related to simplify code.

maobaolong · 2025-03-11T00:25:37Z

vllm/core/block/cpu_offloading_block_allocator.py

+            "handles CPU offloading internally."\
+        # mark this block as uncached
+
+        block = self._allocators[device].allocate_mutable_block(


We can call super implemented methods?

block = super().allocate_mutable_block( prev_block, device, extra_hash=extra_hash )

maobaolong · 2025-03-11T00:33:55Z

vllm/core/block/cpu_offloading_block_allocator.py

+        ), "cpu and gpu block allocators can't have intersection of block ids"
+
+        super().__init__(cpu_block_allocator, gpu_block_allocator)
+        self._allocators: Dict[Device,


This seems defined in super claas, can you use the it from super class?

maobaolong · 2025-03-11T00:35:19Z

vllm/core/block/cpu_offloading_block_allocator.py

+            " handles CPU offloading internally."
+
+        # allocate a GPU block
+        block = self._allocators[device].allocate_immutable_block(


this can be

super().allocate_immutable_block(.....)

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

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.

kyet · 2025-04-03T00:21:25Z

..snip.. There is work in progress for a similar PR in V1.

Which PR are you referring to? Is it possibly #13377?

hmellor · 2025-04-03T14:04:34Z

I don't think the PR has been created yet, but I believe it's coming from the LMCache team

[Add] CPU offloading block allocator and corresponding changes in the…

7401673

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

ApostaC requested review from zhuohan123, youkaichao, alexm-redhat, comaniac and njhill as code owners December 26, 2024 23:15

mergify bot added the frontend label Dec 26, 2024

ApostaC mentioned this pull request Dec 26, 2024

[Core] Support offloading KV cache to CPU #10874

Closed

Use a better data structure for uncached blocks

b083364

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

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

lyppg mentioned this pull request Jan 3, 2025

[Core] Support global prefix caching #11385

Closed

mergify bot added the needs-rebase label Feb 28, 2025

maobaolong reviewed Mar 11, 2025

View reviewed changes

ShangmingCai mentioned this pull request Mar 21, 2025

KV cache SSD Offloading kvcache-ai/Mooncake#159

Closed

hmellor closed this Apr 2, 2025

Uh oh!

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

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

Uh oh!

Conversation

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

Uh oh!

End-to-end benchmarking results:

Implementation

Uh oh!

github-actions bot commented Dec 26, 2024

Uh oh!

youkaichao commented Dec 30, 2024

Uh oh!

comaniac commented Dec 30, 2024

Uh oh!

ApostaC commented Dec 31, 2024

Uh oh!

DearPlanet commented Jan 3, 2025

Uh oh!

KuntaiDu commented Jan 5, 2025

Uh oh!

ApostaC commented Jan 6, 2025

Uh oh!

DearPlanet commented Jan 7, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

boposki commented Jan 17, 2025

Uh oh!

boposki commented Jan 17, 2025

Uh oh!

mergify bot commented Feb 28, 2025

Uh oh!

maobaolong left a comment

Choose a reason for hiding this comment

Uh oh!

maobaolong Mar 11, 2025

Choose a reason for hiding this comment

Uh oh!

maobaolong Mar 11, 2025

Choose a reason for hiding this comment

Uh oh!

maobaolong Mar 11, 2025

Choose a reason for hiding this comment

Uh oh!

hmellor commented Apr 2, 2025

Uh oh!

kyet commented Apr 3, 2025

Uh oh!

hmellor commented Apr 3, 2025

Uh oh!

Uh oh!

ApostaC commented Dec 26, 2024 •

edited by github-actions bot

Loading

DearPlanet commented Jan 7, 2025 •

edited

Loading