migrate P2pNcclConnector to P2pHcclConnector #2686

IuoIun · 2025-09-02T00:53:36Z

What this PR does / why we need it?

Related issue: #2594

The PR migrate the P2pNcclConnector to vllm-ascend project. It supports running disaggregated prefill in the similar way to P2pNcclConnector in vllm.

The only dependency for P2pHcclConnector is hccl, context and stream api in cann. Cann may possibly be installed when using with ascend npu. Therefore, it can be used with fewer pre-dependencies. This is helpful when:

Want to quickly run some test case for disaggregated prefill in users environments.
Run disaggregated prefill and don't want to introduce other dependency.
Quickly run disaggregated prefill in single node with hccl and simple config.

Does this PR introduce any user-facing change?

Yes, this PR introduces a new type of KVConnector(P2pHcclConnector) that can be enabled via configuration.

How was this patch tested?

cann version: 8.2.RC1

test case:

# setup variable and environmentss
export ASCEND_VISIBLE_DEVICES=5,4 # modify according to your environment
export ASCEND_RT_VISIBLE_DEVICES=0,1 # modify according to your environment
modelName=xx # use your own model
servedModelName=xxx # use your own model name

# for prefill instance
vllm serve ${modelName} --host 0.0.0.0 --port 20005 --seed 42 --served-model-name ${servedModelName} --max-model-len 10000 --max-num-seqs 256 --gpu-memory-utilization 0.7 --disable-log-request --kv-transfer-config \{\"kv_rank\":0,\"kv_connector\":\"P2pHcclConnector\",\"kv_role\":\"kv_producer\",\"kv_buffer_size\":\"8e9\",\"kv_port\":\"21001\",\"kv_connector_extra_config\":\{\"proxy_ip\":\"127.0.0.1\",\"proxy_port\":\"30001\",\"http_port\":\"20005\",\"send_type\":\"PUT_ASYNC\"}\} > prefill.log &

# for decode instance
vllm serve ${modelName} --host 0.0.0.0 --port 20009 --seed 1024 --served-model-name  ${servedModelName} --max-model-len 10000 --max-num-seqs 256 --gpu-memory-utilization 0.7 --disable-log-request --kv-transfer-config \{\"kv_rank\":1,\"kv_connector\":\"P2pHcclConnector\",\"kv_role\":\"kv_consumer\",\"kv_buffer_size\":\"8e9\",\"kv_port\":\"22001\",\"kv_connector_extra_config\":\{\"mem_pool_size_gb\":256,\"proxy_ip\":\"127.0.0.1\",\"proxy_port\":\"30001\",\"http_port\":\"20009\",\"send_type\":\"PUT_ASYNC\"\}\} > decode.log &

# start the proxy in vllm project
python vllm/examples/online_serving/disaggregated_serving_p2p_nccl_xpyd/disagg_proxy_p2p_nccl_xpyd.py > proxy.log &

# test request
curl  http://localhost:10001/v1/completions     -H "Content-Type: application/json"     -d '{
    "model": "'$servedModelName'",
    "prompt": "Why is the sky blue?",
    "max_tokens": 100,
    "temperature": 0
}'

vLLM version: v0.10.2
vLLM main: vllm-project/vllm@releases/v0.11.0

Signed-off-by: luolun <luolun1995@cmbchina.com>

github-actions · 2025-09-02T00:53:43Z

👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:‌‌

A PR should do only one thing, smaller PRs enable faster reviews.
Every PR should include unit tests and end-to-end tests ‌to ensure it works and is not broken by other future PRs.
Write the commit message by fulfilling the PR description to help reviewer and future developers understand.

If CI fails, you can run linting and testing checks locally according Contributing and Testing.

gemini-code-assist

Code Review

This pull request migrates the P2pNcclConnector to a P2pHcclConnector for Ascend NPUs, enabling disaggregated prefill. The changes are extensive, introducing a new KVConnector, its engine, and HCCL wrappers. While the overall structure is sound, I've identified several critical issues, many of which appear to be copy-paste errors from the original NCCL implementation. These include incorrect API usage for HCCL streams, wrong exception types for error handling, and logging errors that could lead to runtime crashes. Addressing these issues is crucial for the stability and correctness of the feature on Ascend hardware.

vllm_ascend/distributed/device_communicators/pyhccl.py

vllm_ascend/distributed/kv_transfer/kv_connector/v1/p2p/p2p_hccl_engine.py

vllm_ascend/distributed/device_communicators/pyhccl_wrapper.py

vllm_ascend/distributed/kv_transfer/kv_connector/v1/p2p/p2p_hccl_connector.py

gemini-code-assist · 2025-09-02T00:56:01Z

vllm_ascend/distributed/kv_transfer/kv_connector/v1/p2p/p2p_hccl_engine.py

+    def get_finished(
+            self, finished_req_ids: set[str], no_compile_layers
+    ) -> tuple[Optional[set[str]], Optional[set[str]]]:
+        """
+        Notifies worker-side connector ids of requests that have
+        finished generating tokens.
+
+        Returns:
+            ids of requests that have finished asynchronous transfer,
+            tuple of (sending/saving ids, recving/loading ids).
+            The finished saves/sends req ids must belong to a set provided in a
+            call to this method (this call or a prior one).
+        """
+
+        # Clear the buffer upon request completion.
+        for request_id in finished_req_ids:
+            for layer_name in no_compile_layers:
+                tensor_id = request_id + "#" + layer_name
+                if tensor_id in self.recv_store:
+                    with self.recv_store_cv:
+                        tensor = self.recv_store.pop(tensor_id, None)
+                        self.send_request_id_to_tensor_ids.pop(
+                            request_id, None)
+                        self.recv_request_id_to_tensor_ids.pop(
+                            request_id, None)
+                    if isinstance(tensor, tuple):
+                        addr, _, _ = tensor
+                        self.pool.free(addr)
+
+        # TODO:Retrieve requests that have already sent the KV cache.
+        finished_sending: set[str] = set()
+
+        # TODO:Retrieve requests that have already received the KV cache.
+        finished_recving: set[str] = set()
+
+        return finished_sending or None, finished_recving or None


The get_finished method has an incomplete implementation, with TODO comments indicating that the logic for retrieving finished sending and receiving requests is missing. It currently returns (None, None), which might prevent the scheduler from correctly freeing resources for completed requests. This could lead to resource leaks or deadlocks. Please complete the implementation to correctly track and return the sets of finished request IDs.

Signed-off-by: luolun <luolun1995@cmbchina.com>

into p2pHcclConnector

Signed-off-by: luolun <luolun1995@cmbchina.com>

xueliangyang-oeuler

good job !

wangxiyuan · 2025-09-08T02:01:46Z

vllm_ascend/__init__.py

    register_model()
+
+
+def register_kvconnector():


No need to add a new plugins, I think just register the connector in https://github.com/vllm-project/vllm-ascend/blob/main/vllm_ascend/distributed/__init__.py is enough.

Agreed, fixed

wangxiyuan · 2025-09-08T02:03:37Z

vllm_ascend/distributed/kv_transfer/kv_connector/v1/p2p/__init__.py

I think the module is too complex, vllm_ascend/distributed/kv_connector is OK IMO

wangxiyuan · 2025-09-08T02:04:31Z

vllm_ascend/attention/attention_v1.py

                                              AttentionLayer, AttentionType)
 from vllm.attention.backends.utils import CommonAttentionState
 from vllm.config import VllmConfig
+from vllm.distributed.kv_transfer import (get_kv_transfer_group,


the main branch now support this change. please rebase to main.

wangxiyuan · 2025-09-08T02:11:05Z

vllm_ascend/distributed/kv_transfer/kv_connector/v1/p2p/p2p_hccl_connector.py

+            scheduler_output (SchedulerOutput): the scheduler output object.
+        """
+
+        meta = P2pNcclConnectorMetadata()


so here we reuse nccl metadata totally?

Yes, it can be reused at the moment

github-actions · 2025-09-11T08:47:31Z

This pull request has conflicts, please resolve those before we can evaluate the pull request.

Directly register connector instead of using plugin. Signed-off-by: luolun <luolun1995@cmbchina.com>

Signed-off-by: luolun <luolun1995@cmbchina.com>

Directly register connector instead of using plugin. Signed-off-by: luolun <luolun1995@cmbchina.com>

Signed-off-by: luolun <luolun1995@cmbchina.com>

2. add P2pHcclConnectorMetadata 3. fix connector register Signed-off-by: luolun <luolun1995@cmbchina.com>

jianzs · 2025-09-27T14:52:39Z

vllm_ascend/distributed/kv_transfer/kv_connector/p2p_hccl_connector.py

+            ReqMeta.make_meta(request_id, token_ids, block_ids, block_size))
+
+
+class P2pHcclConnector(KVConnectorBase_V1):


Can we inherit from the upstream class?

migrate P2pNcclConnector to P2pHcclConnector

e47a592

Signed-off-by: luolun <luolun1995@cmbchina.com>

github-actions bot added the module:core label Sep 2, 2025

gemini-code-assist bot reviewed Sep 2, 2025

View reviewed changes

IuoIun added 5 commits September 2, 2025 08:56

Merge branch 'main' into p2pHcclConnector

ad13ed5

rename p2p_nccl_engine to p2p_hccl_engine and improve stream handling

408a19a

Signed-off-by: luolun <luolun1995@cmbchina.com>

Merge branch 'p2pHcclConnector' of https://github.com/luolun/vllm-ascend

8d8b154

into p2pHcclConnector

remove unnesscery log

ec28f86

Signed-off-by: luolun <luolun1995@cmbchina.com>

improve handing of oom

e1fc641

Signed-off-by: luolun <luolun1995@cmbchina.com>

luolun force-pushed the p2pHcclConnector branch from 7e0fadd to e1fc641 Compare September 2, 2025 06:56

fix oom problem

0d5daa3

Signed-off-by: luolun <luolun1995@cmbchina.com>

luolun force-pushed the p2pHcclConnector branch from 2f0c150 to 0d5daa3 Compare September 2, 2025 15:00

xueliangyang-oeuler reviewed Sep 3, 2025

View reviewed changes

wangxiyuan reviewed Sep 8, 2025

View reviewed changes

github-actions bot added the merge-conflicts label Sep 11, 2025

IuoIun added 2 commits September 12, 2025 18:20

Fix: only send blocks in use.

4331bc8

Directly register connector instead of using plugin. Signed-off-by: luolun <luolun1995@cmbchina.com>

Fix: initialize kv_rank and kv_parallel_size based on kv_transfer_config

82a3644

Signed-off-by: luolun <luolun1995@cmbchina.com>

luolun force-pushed the p2pHcclConnector branch from b913b5d to 82a3644 Compare September 12, 2025 10:22

IuoIun added 11 commits September 12, 2025 18:25

Remove unused kvconnector registration from setup and init files

691756d

Signed-off-by: luolun <luolun1995@cmbchina.com>

Clean up code formatting in p2p_hccl_connector and p2p_hccl_engine

0251156

Signed-off-by: luolun <luolun1995@cmbchina.com>

rename p2p_nccl_engine to p2p_hccl_engine and improve stream handling

b5fcb03

Signed-off-by: luolun <luolun1995@cmbchina.com>

remove unnesscery log

583b972

Signed-off-by: luolun <luolun1995@cmbchina.com>

improve handing of oom

af2ee8c

Signed-off-by: luolun <luolun1995@cmbchina.com>

fix oom problem

264ed42

Signed-off-by: luolun <luolun1995@cmbchina.com>

Fix: only send blocks in use.

a1c1649

Directly register connector instead of using plugin. Signed-off-by: luolun <luolun1995@cmbchina.com>

Fix: initialize kv_rank and kv_parallel_size based on kv_transfer_config

eef8adb

Signed-off-by: luolun <luolun1995@cmbchina.com>

Remove unused kvconnector registration from setup and init files

27f239c

Signed-off-by: luolun <luolun1995@cmbchina.com>

Clean up code formatting in p2p_hccl_connector and p2p_hccl_engine

3c291a0

Signed-off-by: luolun <luolun1995@cmbchina.com>

1. load acllib by cann version

ff98f26

2. add P2pHcclConnectorMetadata 3. fix connector register Signed-off-by: luolun <luolun1995@cmbchina.com>

IuoIun added 7 commits September 19, 2025 22:29

reserve kv_cache_memory to avoid oom

2fe9694

reserve kv_cache_memory to avoid oom

25742d3

Merge branch 'p2pHcclConnector_3e330fcb' into p2pHcclConnector

f2c3564

store tensor to pool if kv_buffer is full

1e76c9f

Merge branch 'p2pHcclConnector_3e330fcb' into p2pHcclConnector

6c756d4

Merge remote-tracking branch 'origin/main' into p2pHcclConnector

c9213e5

Fix tensor size calculation in recv_tensor method for improved clarity

0735ab6

github-actions bot removed the merge-conflicts label Sep 27, 2025

jianzs reviewed Sep 27, 2025

View reviewed changes

		ReqMeta.make_meta(request_id, token_ids, block_ids, block_size))


		class P2pHcclConnector(KVConnectorBase_V1):

migrate P2pNcclConnector to P2pHcclConnector #2686

Are you sure you want to change the base?

migrate P2pNcclConnector to P2pHcclConnector #2686

Conversation

IuoIun commented Sep 2, 2025 • edited by jianzs Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What this PR does / why we need it?

Does this PR introduce any user-facing change?

How was this patch tested?

Uh oh!

github-actions bot commented Sep 2, 2025

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

gemini-code-assist bot Sep 2, 2025

Choose a reason for hiding this comment

Uh oh!

xueliangyang-oeuler left a comment

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!

github-actions bot commented Sep 11, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

IuoIun commented Sep 2, 2025 •

edited by jianzs

Loading