[Core] Implementing disaggregated prefilling, and caching KV cache in CPU/disk/database. #8498

KuntaiDu · 2024-09-16T01:36:23Z

TL; DR: implemented disaggregated prefill with <100 core line change (and most of them are comments)

This PR is a continuation of PR #6170 , with a new design that allows future extension.

Current supported applications:

Disaggregated prefilling. Check examples/disagg_prefill/disagg_prefill_example.sh for an example, and benchmarks/disagg_prefill for various benchmarks. Benchmarking script are all one-click runnable (after setting HF_TOKEN)
Connecting to a KV cache storage service LMCache. Examples TBD.

Two roles: KV provider (e.g. prefill vLLM instance) and KV consumer (e.g. decode vLLM instance)

Provider side: insert: insert a KV cache to a buffer, so that it can be transferred upon request
Consumer side: drop_select: select a KV cache based on tokens, transfer the selected KV, and drop this KV out from the buffer

Example workflow (the buffer in the following figure is the same as insert)

PR Checklist (Click to Expand)

Thank you for your contribution to vLLM! Before submitting the pull request, please ensure the PR meets the following criteria. This helps vLLM maintain the code quality and improve the efficiency of the review process.

PR Title and Classification

Only specific types of PRs will be reviewed. The PR title is prefixed appropriately to indicate the type of change. Please use one of the following:

[Bugfix] for bug fixes.
[CI/Build] for build or continuous integration improvements.
[Doc] for documentation fixes and improvements.
[Model] for adding a new model or improving an existing model. Model name should appear in the title.
[Frontend] For changes on the vLLM frontend (e.g., OpenAI API server, LLM class, etc.)
[Kernel] for changes affecting CUDA kernels or other compute kernels.
[Core] for changes in the core vLLM logic (e.g., LLMEngine, AsyncLLMEngine, Scheduler, etc.)
[Hardware][Vendor] for hardware-specific changes. Vendor name should appear in the prefix (e.g., [Hardware][AMD]).
[Misc] for PRs that do not fit the above categories. Please use this sparingly.

Note: If the PR spans more than one category, please include all relevant prefixes.

Code Quality

The PR need to meet the following code quality standards:

We adhere to Google Python style guide and Google C++ style guide.
Pass all linter checks. Please use format.sh to format your code.
The code need to be well-documented to ensure future contributors can easily understand the code.
Include sufficient tests to ensure the project to stay correct and robust. This includes both unit tests and integration tests.
Please add documentation to docs/source/ if the PR modifies the user-facing behaviors of vLLM. It helps vLLM user understand and utilize the new features or changes.

Adding or changing kernels

Each custom kernel needs a schema and one or more implementations to be registered with PyTorch.

Make sure custom ops are registered following PyTorch guidelines: Custom C++ and CUDA Operators and The Custom Operators Manual
Custom operations that return Tensors require meta-functions. Meta-functions should be implemented and registered in python so that dynamic dims can be handled automatically. See above documents for a description of meta-functions.
Use torch.libary.opcheck() to test the function registration and meta-function for any registered ops. See tests/kernels for examples.
When changing the C++ signature of an existing op, the schema must be updated to reflect the changes.
If a new custom type is needed, see the following document: Custom Class Support in PT2.

Notes for Large Changes

Please keep the changes as concise as possible. For major architectural changes (>500 LOC excluding kernel/data/config/test), we would expect a GitHub issue (RFC) discussing the technical design and justification. Otherwise, we will tag it with rfc-required and might not go through the PR.

What to Expect for the Reviews

The goal of the vLLM team is to be a transparent reviewing machine. We would like to make the review process transparent and efficient and make sure no contributor feel confused or frustrated. However, the vLLM team is small, so we need to prioritize some PRs over others. Here is what you can expect from the review process:

After the PR is submitted, the PR will be assigned to a reviewer. Every reviewer will pick up the PRs based on their expertise and availability.
After the PR is assigned, the reviewer will provide status update every 2-3 days. If the PR is not reviewed within 7 days, please feel free to ping the reviewer or the vLLM team.
After the review, the reviewer will put an action-required label on the PR if there are changes required. The contributor should address the comments and ping the reviewer to re-review the PR.
Please respond to all comments within a reasonable time frame. If a comment isn't clear or you disagree with a suggestion, feel free to ask for clarification or discuss the suggestion.

Thank You

Finally, thank you for taking the time to read these guidelines and for your interest in contributing to vLLM. Your contributions make vLLM a great tool for everyone!

github-actions · 2024-09-16T01:36:35Z

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

🚀

ahg-g · 2024-09-18T16:19:43Z

/cc

rkooo567 · 2024-09-18T16:48:57Z

@KuntaiDu this is ready to review?

KuntaiDu · 2024-09-18T22:32:29Z

@KuntaiDu this is ready to review?

Final bug fix on TP case when connecting to LMCache. Will ping you when it's ready for review (most likely tomorrow)

zeroorhero · 2024-09-19T02:13:52Z

It‘s a good job!

KuntaiDu · 2024-09-20T01:15:03Z

@rkooo567 Now this code is ready for review

KuntaiDu · 2024-09-20T05:51:34Z

Lol huggingface is down and I will rerun the test tomorrow.

kuangdao · 2024-09-20T08:37:36Z

mark

lixiaolx · 2024-09-20T11:53:11Z

mark

KuntaiDu · 2024-09-22T18:02:16Z

All tests are now passed!

zeroorhero · 2024-09-23T03:25:29Z

@KuntaiDu hi ！According to your commit, I implemented valkey (supporting TCP and rdma) as the kv cache storage pool in the prefill and decode nodes(#8724). Due to license reasons, I implemented valkey instead of redis, and valkey also supports rdma after version 8.0. Can you give me some suggestions?

comaniac

First review the changes in the vLLM core.

vllm/utils.py

vllm/worker/worker.py

vllm/worker/worker_base.py

vllm/worker/model_runner.py

comaniac · 2024-09-23T23:41:56Z

vllm/distributed/parallel_state.py

+            - 2 * tp * pp
+            - Why: both prefill vLLM and decode vLLM is in the world


Does that mean we now only support 1-prefill + 1-decode?

Yes for this PR. Support on multi prefil and multi decode vllm instances with different tp/pp is on the roadmap but not in this PR.

I guess you are wondering if this architecture is scalable enough to support XpYd scenario. It can. Here is a planned architecture:

(prefill worker --- database frontend) --- database backend --- (database frontend --- decode worker)

The instances in the same bracket share the same world in torch.distributed.

Basically for XpYd & multiple prefill multiple decode worker scenario, I am planning to not connect prefill worker to decoder worker. Instead, I will connect them to a database and let the database handle all the network topology related optimizations and let it handle different tp/pp conversion.

So in summary if we want to achieve XpYd, then we should separately configure Xp workers and Yd workers. Can you elaborate how to configure Xp workers as an example? I suppose we will only configure KV producer?

Example architecture:

Each gray box is inside the same torch.distributed world.

(And LMCache can be one example solution for the database. It implements this frontend and backend, and the frontend placement can be in various places (CPU, GPU, disk)).

Right now, to configure Xp workers and Yd workers will be quite messy in parallel state. For example, when we have 4 prefill workers(TP=4) and 2 decode workers(TP=2), world size will not be 4 * 2 here, and _DISAGG tuples will not be one to one. Another issue will occur in torch_distributed_pipe. Since send operations and recv operations will be doubled when prefill has twice worker as decode.

Even with a database, the database frontend will need to manage the difference between prefill workers and decode workers, and perform concat or scatter operations when necessary.

Example architecture:

Each gray box is inside the same torch.distributed world.

Will the latency (connecting to a db) be a concern?

Yes, but the overhead is engineeringly optimizable by, e.g., pipelining the send/recv of the KV cache with model forwarding, or by using RDMA database, etc.

Rest assured that we are planning to support both one-hop connection (i.e. prefiller directly connected to decoder) and two-hop connection (i.e. prefiller --- database --- decoder).

KuntaiDu · 2024-09-24T00:16:33Z

@KuntaiDu hi ！According to your commit, I implemented valkey (supporting TCP and rdma) as the kv cache storage pool in the prefill and decode nodes(#8724). Due to license reasons, I implemented valkey instead of redis, and valkey also supports rdma after version 8.0. Can you give me some suggestions?

Thank you (rdma yes!). I left some comments to your PR.

zeroorhero · 2024-09-24T01:40:13Z

@KuntaiDu hi ！According to your commit, I implemented valkey (supporting TCP and rdma) as the kv cache storage pool in the prefill and decode nodes(#8724). Due to license reasons, I implemented valkey instead of redis, and valkey also supports rdma after version 8.0. Can you give me some suggestions?

Thank you (rdma yes!). I left some comments to your PR.

Thanks!

pizhenwei · 2024-09-24T01:49:04Z

comments

Hi @KuntaiDu

I'm the author of Valkey Over RDMA, please feel free to contact me [email protected] on any issue or suggestion.

KuntaiDu · 2024-09-24T07:14:06Z

@comaniac I have fixed the comments. Thank you for your review 😄. As for the hidden states send/recv and sampler output send/recv, I need to do some extra model runner condition checks, and I will make a follow-up PR to do this (this PR is already huge).

Yang-x-Zhao · 2024-09-24T12:12:53Z

vllm/worker/model_runner.py

+        bypass_model_exec = False
+        if self.need_recv_kv(model_input, kv_caches):
+            hidden_or_intermediate_states, bypass_model_exec, model_input = \
+                get_disagg_group().recv_kv_caches_and_hidden_states(


Just curious. Is this recv function blocking?

It seems to me that adding an independent thread that constantly waiting for kvcache might hide the communication cost, but I am not sure if this is feasible under current block manager logic.

Yeah currently it is blocking (and it is an intentional design choice).
Reasons behind this design:

Transfer delay is not significant (~5ms). Most of the overheads in the implementation is tokenizing the whole request twice and de-tokenizing the first token twice (~20ms). Optimizing other overheads will give us more perf improvement.

Non-blocking KV transfer requires adding GPU buffer at decoding vLLM side. Since decoding is memory-bounded, this is not ideal (as buffering reduces # of available GPU memory for decoding KV caches).

Transfer delay is not significant (~5ms). Most of the overheads in the implementation is tokenizing the whole request twice and de-tokenizing the first token twice (~20ms).

Thanks for this helpful information! I agree with your observation. Skipping sampling & output processor in prefill side and sending kvcache immediately after hidden states will be more significant.

sjnaj · 2024-09-24T12:51:06Z

It would be nice if the decode part could support non-blocking reception and continious batching.
For example, when executing the step function on a batch, according to the maximum capacity obtained by the profile, the required requests can be obtained from a queue that stores some transfer completed (prefetched)kvcache and other infos, then add them to the batch.

artetaout · 2024-09-24T14:29:40Z

when the consumer receives a complete new request , should it work or reject it？in this case, should we add some warnings？

KuntaiDu · 2024-09-24T16:24:21Z

It would be nice if the decode part could support non-blocking reception and continious batching. For example, when executing the step function on a batch, according to the maximum capacity obtained by the profile, the required requests can be obtained from a queue that stores some transfer completed (prefetched)kvcache and other infos, then add them to the batch.

The goal of this PR is to create minimum core change so that people can easily understand the code when review it. We can add this optimization in future PRs.

KuntaiDu · 2024-09-24T16:26:47Z

when the consumer receives a complete new request , should it work or reject it？in this case, should we add some warnings？

It will redo the prefill. Currently we won't show a warning, but we do expect some logging and warning mechanism in the future.

comaniac

Took a pass to the kv transfer implementation. Overall is clean is easy to follow. Will take a deeper look at the performance later.

Meanwhile, @youkaichao could you take a look at the parallel states? The current way of handling torch.distributed is a bit messy, but I don't have a better way to deal with it in my mind.

@rkooo567 please also take a look at the parallel states. I believe the current implementation would have some issues with Ray integration.

vllm/utils.py

comaniac · 2024-09-24T17:19:11Z

vllm/distributed/parallel_state.py

+            - 2 * tp * pp
+            - Why: both prefill vLLM and decode vLLM is in the world


So in summary if we want to achieve XpYd, then we should separately configure Xp workers and Yd workers. Can you elaborate how to configure Xp workers as an example? I suppose we will only configure KV producer?

benchmarks/benchmark_serving.py

vllm/distributed/kv_transfer/kv_lookup_buffer/base.py

vllm/distributed/kv_transfer/kv_lookup_buffer/simple_buffer.py

vllm/distributed/kv_transfer/kv_pipe/torch_distributed_pipe.py

comaniac · 2024-09-24T17:38:22Z

vllm/distributed/kv_transfer/vllm_adapter.py

@@ -0,0 +1,444 @@
+"""vLLM distributed KV cache transfer API.


Is there any planned adaptor in the future?

Personally I feel vLLM will only have one adapter (not sure if other guys will implement new adapters in vLLM tho). But I do expect that other applications (like other serving engines, or a KV database application) implement their own adapters to integrate the transferred KV caches to their application workflow.

I see that's interesting. So you're expecting someone to just from vllm.distributed import kv_transfer, and only use the kv_transfer for other frameworks?

Yes. I am expecting different applications to follow the same set of kv_transfer protocol so that they can communicate with vLLM (by, for example, "pretending" to be a vLLM decode instance).

Also kv_transfer library is intentionally designed in a way that it has no code dependency on vLLM so that other applications can use it.

KuntaiDu · 2024-09-24T22:32:36Z

@comaniac Thanks for reviewing! Code changes done. Also pinged Kaichao for further review on parallel_state.py.

Performance-wise, here are some bottlenecks I know (these bottlenecks are sorted by how significant they are).

In current workflow (shown in the Figure), the decode vLLM will start to run after first token being generated by prefill instance, and decode vLLM instance will also generate the first token.
- Issue: tokenize for the full sequence twice, detokenize the first token twice
- Overhead: ~20ms.
- Solution: detokenization overhead can be fixed easily (just send SamplerOutput to the decode vLLM), tokenization is also fixable but not easy.
Current recv is blocking
- Overhead: ~5ms.
- Not easy to fix:
  - A strawmen solution: add a buffer at the decode vLLM side and let it pre-fetch the KV cache
  - Not good: adding the buffer in GPU reduces the KV cache space for decode vLLM (and decoding phase is GPU-memory-bound). Adding the buffer in CPU will slow down the distributed communication (cannot go through NVLink).
Current KV cache lookup operation is O(n), n is # of cached KV caches.
- Overhead: not sure but won't be large (the buffer size won't be (and shouldn't be) large in order to leave more GPU memory space for inference).
- Easy to fix but needs an efficient hashing logic.

Signed-off-by: Russell Bryant <[email protected]>

…lm-project#9766) Signed-off-by: ElizaWszola <[email protected]>

…lls (vllm-project#10452) Signed-off-by: Max de Bayser <[email protected]>

Signed-off-by: Alexei V. Ivanov <[email protected]>

Signed-off-by: Lucas Wilkinson <[email protected]>

Signed-off-by: Yanyi Liu <[email protected]>

…ading (vllm-project#10456) Signed-off-by: Isotr0py <[email protected]>

…t#10470) Signed-off-by: DarkLight1337 <[email protected]>

…10423) (vllm-project#10424)

Signed-off-by: skylee-01 <[email protected]>

…ect#10471) Signed-off-by: DarkLight1337 <[email protected]>

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

Signed-off-by: Kuntai Du <[email protected]>

mergify · 2024-11-20T07:38:55Z

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

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

KuntaiDu · 2024-11-20T07:47:22Z

lol something happened when I tried to sign off. I'll fix it tmrw

KuntaiDu requested review from rkooo567, simon-mo and WoosukKwon September 16, 2024 01:39

KuntaiDu mentioned this pull request Sep 16, 2024

[Core] implement disaggregated prefilling via KV cache transfer #6170

Closed

7 tasks

KuntaiDu requested a review from youkaichao September 16, 2024 01:47

KuntaiDu force-pushed the kuntai-disagg-refactor branch from eb751d6 to 60ede08 Compare September 16, 2024 02:11

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

comaniac reviewed Sep 23, 2024

View reviewed changes

Yang-x-Zhao reviewed Sep 24, 2024

View reviewed changes

comaniac reviewed Sep 24, 2024

View reviewed changes

mergify bot added the needs-rebase label Nov 19, 2024

russellb and others added 16 commits November 19, 2024 21:05

[Core] Avoid metrics log noise when idle (vllm-project#8868)

704bacd

Signed-off-by: Russell Bryant <[email protected]>

[Model][Quantization] HQQ support through Marlin kernel expansion (vl…

278b10a

…lm-project#9766) Signed-off-by: ElizaWszola <[email protected]>

Change granite chat template to keep json list formatting for tool ca…

b3daa75

…lls (vllm-project#10452) Signed-off-by: Max de Bayser <[email protected]>

[CI/Build] Update Dockerfile.rocm (vllm-project#10434)

b961015

Signed-off-by: Alexei V. Ivanov <[email protected]>

[Bugfix] Marlin 2:4 temp fix for large M dim (>256) (vllm-project#10464)

87b0a41

Signed-off-by: Lucas Wilkinson <[email protected]>

[Misc] Add __setitem__ for LazyDict (vllm-project#10469)

4e187ef

Signed-off-by: Yanyi Liu <[email protected]>

[Bugfix] Fix Mamba model initialization and MLP Speculator weights lo…

0e9d8d9

…ading (vllm-project#10456) Signed-off-by: Isotr0py <[email protected]>

[Bugfix] Enforce no chunked prefill for embedding models (vllm-projec…

c41c605

…t#10470) Signed-off-by: DarkLight1337 <[email protected]>

[CI/Build] Add sphinx/rst linter for docs (vllm-project#10366)

0b38eef

[CI/Build] Support compilation with local cutlass path (vllm-project#…

de2483a

…10423) (vllm-project#10424)

[ci/build] Combine nightly and optional (vllm-project#10465)

8f5691c

[model] Reduce medusa weight (vllm-project#10454)

173c5aa

Signed-off-by: skylee-01 <[email protected]>

[Bugfix] Handle conflicts between modern and legacy fields (vllm-proj…

da0d34c

…ect#10471) Signed-off-by: DarkLight1337 <[email protected]>

[Platforms] Refactor xpu code (vllm-project#10468)

acc7c54

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

resolve merge conflicts

bae609a

Signed-off-by: Kuntai Du <[email protected]>

fall back to full prefill when any of the KV cache receive fails.

1780820

Signed-off-by: Kuntai Du <[email protected]>

mergify bot removed the needs-rebase label Nov 20, 2024

KuntaiDu force-pushed the kuntai-disagg-refactor branch from 529c425 to 1780820 Compare November 20, 2024 07:38

KuntaiDu requested review from mgoin, robertgshaw2-neuralmagic, tlrmchlsmth, LiuXiaoxuanPKU, DarkLight1337 and ywang96 as code owners November 20, 2024 07:38

mergify bot added documentation Improvements or additions to documentation frontend labels Nov 20, 2024

mergify bot added the needs-rebase label Nov 20, 2024

		- 2 * tp * pp
		- Why: both prefill vLLM and decode vLLM is in the world

[Core] Implementing disaggregated prefilling, and caching KV cache in CPU/disk/database. #8498

Are you sure you want to change the base?

[Core] Implementing disaggregated prefilling, and caching KV cache in CPU/disk/database. #8498

Conversation

KuntaiDu commented Sep 16, 2024 • edited Loading

PR Title and Classification

Code Quality

Adding or changing kernels

Notes for Large Changes

What to Expect for the Reviews

Thank You

github-actions bot commented Sep 16, 2024

ahg-g commented Sep 18, 2024

rkooo567 commented Sep 18, 2024

KuntaiDu commented Sep 18, 2024

zeroorhero commented Sep 19, 2024

KuntaiDu commented Sep 20, 2024

KuntaiDu commented Sep 20, 2024

kuangdao commented Sep 20, 2024

lixiaolx commented Sep 20, 2024

KuntaiDu commented Sep 22, 2024 • edited Loading

zeroorhero commented Sep 23, 2024

comaniac left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

KuntaiDu Sep 24, 2024 • edited Loading

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

KuntaiDu commented Sep 24, 2024

zeroorhero commented Sep 24, 2024

pizhenwei commented Sep 24, 2024

KuntaiDu commented Sep 24, 2024

Choose a reason for hiding this comment

KuntaiDu Sep 24, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sjnaj commented Sep 24, 2024

artetaout commented Sep 24, 2024

KuntaiDu commented Sep 24, 2024

KuntaiDu commented Sep 24, 2024

comaniac left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

KuntaiDu Sep 24, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

KuntaiDu Sep 25, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

KuntaiDu commented Sep 24, 2024

mergify bot commented Nov 20, 2024

KuntaiDu commented Nov 20, 2024

KuntaiDu commented Sep 16, 2024 •

edited

Loading

KuntaiDu commented Sep 22, 2024 •

edited

Loading

KuntaiDu Sep 24, 2024 •

edited

Loading

KuntaiDu Sep 24, 2024 •

edited

Loading

KuntaiDu Sep 24, 2024 •

edited

Loading

KuntaiDu Sep 25, 2024 •

edited

Loading