[WIP] [Core][P/D] CPU connector for PD disagg #18332

ApostaC · 2025-05-19T05:44:48Z

TL;DR:
This PR implements CPU-based connector for PD disaggregation, with the following features

Async layerwise D2H copies at prefiller side
(WIP) Async layerwise H2D copies at decoder side
(WIP) Pure async CPU-CPU data transmission (plan: NIXL with Mooncake transfer engine or UCX)

(Update May 22) Now the code is functional and can produce correct results in offline test scripts. There are a few more minor things to do.

Backend perf bench: Mooncake vs UCX
- See https://github.com/ApostaC/nixl_benchmark
test TP compatibility
add online examples
remove the hard code
add potential performance improvements
fix the linter issues

High-level designs

Prefiller node

When a layer finishes the prefill at the prefiller node, the following things will happen:

Prefill worker process launch async d2h copies
Prefill worker sends a request to the decode worker (via a side-channel)
Prefill save this request to the "processing queue"

Upon the finish of each layer's prefill, the prefill worker will also do a quick "status check" on the "processing queue":

If the request finished both d2h copy and got a "ready" response from the decode worker, the prefiller worker will start sending the KV cache
If the request has finished sending, the prefiller worker will free the resources of the request and remove it from the "processing queue"

The following screenshot demonstrates the layer-wise async KV cache storing:

Decoder node

Every time the decoder worker has been called, it will do the following things:

Get new pending requests from the scheduler
Check the pending requests from the prefiller side-channel, and allocate the memory to receive the data and send the response back to the prefiller
check if there are new data receiving tasks finished and updates its internal states
Load the KV cache to GPU layer-by-layer
- Start the async H2D copy of layer i+1
- Wait for the async H2D copy of layer i to be finished
- Return to the model execution of layer i
Return the ready request IDs back to the scheduler

Immediate TODO items

Finish the functionality
Fix linter issues
Unit tests

Future works

Load-balancer
Separate the CPU connector into a different process
Async parallelism support

Signed-off-by: ApostaC <yihua98@uchicago.edu>

github-actions · 2025-05-19T05:44:58Z

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

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 either: Add ready label to the PR or enable auto-merge.

🚀

Signed-off-by: ApostaC <yihua98@uchicago.edu>

robertgshaw2-redhat · 2025-06-03T13:51:47Z

vllm/distributed/kv_transfer/kv_connector/v1/cpu_connector.py

+    dst_v.copy_(src_v, non_blocking=True)
+
+
+def h2d_page_copy(src_buffer: torch.Tensor, dst_layer: torch.Tensor,


QQ - do you think we can tweak the code to have an invariant that start_token_idx and stop_token_idx are multiples of block_size? This would really simplify the implementation I think

There will be some annoying corner cases with chunked prefill.

For example, assume the max_num_batched_tokens is 8192, and we have 2 requests with 8200 tokens.
In the first scheduling step, the first 8192 tokens of the first request will be scheduled as a new request.
In the second scheduling step, the last 8 tokens of the first request will be scheduled as "running request", and then the first 8192 - 8 = 8184 tokens of the second request will be scheduled as a new request.

In the above example, the second request will have the stop_token_idx not multiples of block_size.

Sure, but this does not matter for p/d since the offloading tokens are not "scheduled", they are treated like a prefix cache hit from the POV of the scheduler

oh i see, its because we are doing the sending chunk by chunk

njhill

Thanks @ApostaC! Still reviewing but posting comments so far...

vllm/v1/core/sched/scheduler.py

vllm/distributed/kv_transfer/kv_connector/v1/base.py

vllm/distributed/kv_transfer/kv_connector/v1/cpu_connector.py

njhill · 2025-06-02T22:01:32Z

vllm/distributed/kv_transfer/kv_connector/v1/cpu_connector.py

+        # NOTE(ApostaC): For a single request, this function will be called
+        # two times if the first time we returned async_load flag as True.
+        # The second time will be the "real schedule" time


Actually this is no longer the case, it will now typically only be called the first time in this case. It may still be called a second time if kv cache allocation fails the first time (without any call to update_state_after_alloc in-between).

vllm/distributed/kv_transfer/kv_connector/v1/cpu_connector.py

njhill · 2025-06-02T22:09:24Z

vllm/distributed/kv_transfer/kv_connector/v1/cpu_connector.py

+        prefill_request_id = kv_transfer_params["prefill_request_id"]
+        self._connect_request_ids(prefill_request_id, request_id)
+        self._should_be_ready_reqs.add(request_id)


I think we may want to move this to update_state_after_alloc. We should assume that get_num_new_matched_tokens is side-effect free. The multi-connector implementation will assume this.

See discussion and MultiConnector updates in #18632.

vllm/distributed/kv_transfer/kv_connector/v1/cpu_connector.py

robertgshaw2-redhat · 2025-06-03T16:07:50Z

vllm/distributed/kv_transfer/kv_connector/v1/cpu_connector.py

+            )
+
+            # Create a destination spec
+            dest_spec = DestinationSpec(


How will we support heterogeneous TP?

Good question. I'm envisioning the support in the following PRs.
When everything is offloaded to CPU, it would be easier to rearrange the memory there.

one idea is to have a separate ring buffer for each KV head

mergify · 2025-06-03T18:22:30Z

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

vllm/distributed/kv_transfer/kv_connector/v1/cpu_connector.py

njhill · 2025-06-04T17:48:20Z

vllm/distributed/kv_transfer/kv_connector/v1/cpu_connector.py

+        layer_id = self._get_layer_id(layer_name)
+        event = self._decoder_cuda_events.pop(layer_id, None)
+        if event is not None:
+            event.synchronize()
+
+        if layer_id == len(self._gpu_kv_caches) - 1:
+            # Free the memory for the whole request
+            for p_req_id in self._inflight_h2d_requests:
+                logger.info("Freeing request %s, current watermark: [%d, %d]",
+                            p_req_id,
+                            self._kv_receiver._allocator.low_watermark,
+                            self._kv_receiver._allocator.high_watermark)
+                self._kv_receiver.free_request(p_req_id)
+            self._inflight_h2d_requests.clear()


Why do we need to synchronize/block here? The idea of the async parts of the connector API is that these loads can span forward passes. If we are blocking in-line then there's no need to return load_kv_async=True from the scheduler side request_finished method, and no need to return anything from the worker side get_finished method...

Sorry if I've misunderstood something here

There are two "ready" steps:

first is the tcp transfer

second if the cpu -> gpu load

only the tcp transfer happens async in this PR

robertgshaw2-redhat · 2025-06-06T22:23:17Z

justfile

# Setting this allows creating a symlink to Justfile from another dir
set working-directory := "/home/rshaw/vllm/pd_examples/"

# Needed for the proxy server
vllm-directory := "/home/rshaw/vllm/" 

MODEL := "Qwen/Qwen3-0.6B"
P_PORT := "8100"
D_PORT := "8200"
PROXY_PORT := "8192"

prefill:
    CUDA_VISIBLE_DEVICES=0 \
    vllm serve {{MODEL}} \
      --port {{P_PORT}} \
      --disable-log-requests \
      --enforce-eager \
      --kv-transfer-config '{"kv_connector":"CPUConnector","kv_role":"kv_producer","kv_connector_extra_config": {"host": "localhost", "port": "54321", "size": 8}}'

decode:
    CUDA_VISIBLE_DEVICES=1 \
    vllm serve {{MODEL}} \
      --port {{D_PORT}} \
      --disable-log-requests \
      --enforce-eager \
      --kv-transfer-config '{"kv_connector":"CPUConnector","kv_role":"kv_consumer","kv_connector_extra_config": {"host": "localhost", "port": "54321", "size": 8}}'


proxy:
    python "{{vllm-directory}}tests/v1/kv_connector/cpu_kv_integration/toy_proxy_server.py" \
      --port {{PROXY_PORT}} \
      --prefiller-port {{P_PORT}} \
      --decoder-port {{D_PORT}}


send_request:
  curl -X POST http://localhost:{{PROXY_PORT}}/v1/completions \
    -H "Content-Type: application/json" \
    -d '{ \
      "model": "{{MODEL}}", \
      "prompt": "Red Hat is the best open source company by far across Linux, K8s, and AI, and vLLM has the greatest community in open source AI software infrastructure. I love vLLM because", \
      "max_tokens": 150, \
      "temperature": 0.7 \
    }'

benchmark:
  python {{vllm-directory}}/benchmarks/benchmark_serving.py --port {{PROXY_PORT}} --model {{MODEL}} --dataset-name random --random-input-len 1000 --random-output-len 100

benchmark_one INPUT_LEN:
  python {{vllm-directory}}benchmarks/benchmark_one_concurrent_req.py \
    --model {{MODEL}} \
    --input-len {{INPUT_LEN}} \
    --output-len 1 \
    --num-requests 5 \
    --seed $(date +%s) \
    --port {{PROXY_PORT}}
  

eval:
  lm_eval --model local-completions --tasks gsm8k \
    --model_args model={{MODEL}},base_url=http://127.0.0.1:{{PROXY_PORT}}/v1/completions,num_concurrent=50,max_retries=3,tokenized_requests=False \
    --limit 300

robertgshaw2-redhat · 2025-06-06T22:24:22Z

Some Notes:

I ran LM Eval with the Justfile above. The PR crashes under load.
I ran the following benchmark_one on a local H100 machine ... so its going over localhost. We see about 3s to send 10k tokens

Scenario	100 Tokens	1000 Tokens	10000 Tokens
P/D	168ms	406ms	3039ms
No P/D	78ms	168ms	470ms

ApostaC · 2025-06-09T16:42:25Z

I'm traveling these days. Will come back to this PR after this Wednesday.

Signed-off-by: ApostaC <yihua98@uchicago.edu>

ApostaC · 2025-06-24T04:56:22Z

Fixed the crash problem. Now the lm_eval runs with the following output on llama-3.1-8B model:

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.7933|±  |0.0234|
|     |       |strict-match    |     5|exact_match|↑  |0.7400|±  |0.0254|

Will clean the debug code and push commits soon.

Signed-off-by: ApostaC <yihua98@uchicago.edu>

ApostaC · 2025-06-26T04:01:14Z

@njhill @robertgshaw2-redhat Now the crashing & hanging issue should be fixed.

ApostaC added 3 commits May 19, 2025 05:26

[Add] cpu kv sender interfaces

6a9eee3

Signed-off-by: ApostaC <yihua98@uchicago.edu>

[WIP] Rewrite the whole implementation for layer-wise pipeline

ad7d852

Signed-off-by: ApostaC <yihua98@uchicago.edu>

[Add] adding tests for cpu kv pd

b1e003e

Signed-off-by: ApostaC <yihua98@uchicago.edu>

mergify bot added the v1 label May 19, 2025

ApostaC added 16 commits May 20, 2025 00:28

[Add] adding nixl transfer impl WIP

ccc8c1d

Signed-off-by: ApostaC <yihua98@uchicago.edu>

[Add] tests for nixl

050cfe6

Signed-off-by: ApostaC <yihua98@uchicago.edu>

Passed the nixl protocol unit tests

e5034b0

Signed-off-by: ApostaC <yihua98@uchicago.edu>

[Add] correct nixl data plane functionality

81e31b8

Signed-off-by: ApostaC <yihua98@uchicago.edu>

[Add] sender to receiver data plane finished

a6ffb26

Signed-off-by: ApostaC <yihua98@uchicago.edu>

[add] NixlPrefillManager and NixlDecodeManager and test example

9b0c66b

Signed-off-by: ApostaC <yihua98@uchicago.edu>

[Add] unit tests for more functionalities

2647ce5

Signed-off-by: ApostaC <yihua98@uchicago.edu>

[Ckpt] everything is functional

96cb2b5

Signed-off-by: ApostaC <yihua98@uchicago.edu>

[Add] remove the hard-coded host and port

005d5c1

Signed-off-by: ApostaC <yihua98@uchicago.edu>

[Add] precommit fixes

44b36be

Signed-off-by: ApostaC <yihua98@uchicago.edu>

[fix] ruff errors

2e2937f

Signed-off-by: ApostaC <yihua98@uchicago.edu>

[fix] format checker issue for tests

242098b

Signed-off-by: ApostaC <yihua98@uchicago.edu>

[remove] outdated tests

76e1473

Signed-off-by: ApostaC <yihua98@uchicago.edu>

[remove] hardcodes and fix precommit issues

01de06a

Signed-off-by: ApostaC <yihua98@uchicago.edu>

[remove] previous debug codes

104418e

Signed-off-by: ApostaC <yihua98@uchicago.edu>

[Add] bug fix for TP and correctly shutdown

b4994f0

Signed-off-by: ApostaC <yihua98@uchicago.edu>

ApostaC marked this pull request as ready for review May 28, 2025 15:39

ApostaC added 5 commits May 28, 2025 17:02

[fix] concurrency bug in TP > 1

1f01921

Signed-off-by: ApostaC <yihua98@uchicago.edu>

[Add] nsys analysis and add potential optimizations

af03fd5

Signed-off-by: ApostaC <yihua98@uchicago.edu>

[Add] small fixes for corner cases

483ed5a

Signed-off-by: ApostaC <yihua98@uchicago.edu>

temp fix for pending request ids

b27c101

Signed-off-by: ApostaC <yihua98@uchicago.edu>

[fix] online problems

78e6c02

Signed-off-by: ApostaC <yihua98@uchicago.edu>

ApostaC requested review from WoosukKwon, robertgshaw2-redhat and njhill as code owners May 30, 2025 21:48

ApostaC requested review from ywang96, comaniac and alexm-redhat as code owners May 30, 2025 21:48

[Add] passed the initial benchmark test

32dc419

Signed-off-by: ApostaC <yihua98@uchicago.edu>

robertgshaw2-redhat reviewed Jun 3, 2025

View reviewed changes

njhill reviewed Jun 3, 2025

View reviewed changes

robertgshaw2-redhat reviewed Jun 3, 2025

View reviewed changes

vllm/distributed/kv_transfer/kv_connector/v1/cpu_connector.py Outdated Show resolved Hide resolved

robertgshaw2-redhat reviewed Jun 3, 2025

View reviewed changes

mergify bot added the needs-rebase label Jun 3, 2025

vllm-project deleted a comment from gemini-code-assist bot Jun 3, 2025

njhill reviewed Jun 4, 2025

View reviewed changes

ApostaC added 2 commits June 15, 2025 16:13

Address the review comments

9683b48

Signed-off-by: ApostaC <yihua98@uchicago.edu>

Merge branch 'main' into local-dev/cpu-kv

60c43be

Signed-off-by: ApostaC <yihua98@uchicago.edu>

mergify bot removed the needs-rebase label Jun 15, 2025

ApostaC added 2 commits June 25, 2025 20:16

[fix] crash problem

3720cf8

Signed-off-by: ApostaC <yihua98@uchicago.edu>

fix the hang problem

7d40a1d

Signed-off-by: ApostaC <yihua98@uchicago.edu>

tlrmchlsmth mentioned this pull request Jun 27, 2025

[RFC]: Data Parallel Attention and Expert Parallel MoEs #16037

Open

22 tasks

		dst_v.copy_(src_v, non_blocking=True)


		def h2d_page_copy(src_buffer: torch.Tensor, dst_layer: torch.Tensor,

Uh oh!

[WIP] [Core][P/D] CPU connector for PD disagg #18332

Are you sure you want to change the base?

[WIP] [Core][P/D] CPU connector for PD disagg #18332

Uh oh!

Conversation

ApostaC commented May 19, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

High-level designs

Prefiller node

Decoder node

Immediate TODO items

Future works

Uh oh!

github-actions bot commented May 19, 2025

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!

njhill left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

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 Jun 3, 2025

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

robertgshaw2-redhat commented Jun 6, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

robertgshaw2-redhat commented Jun 6, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ApostaC commented Jun 9, 2025

Uh oh!

ApostaC commented Jun 24, 2025

Uh oh!

ApostaC commented Jun 26, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

ApostaC commented May 19, 2025 •

edited

Loading

robertgshaw2-redhat commented Jun 6, 2025 •

edited

Loading

robertgshaw2-redhat commented Jun 6, 2025 •

edited

Loading

ApostaC commented Jun 26, 2025 •

edited

Loading