some fixes

Signed-off-by: MengqingCao <cmq0113@163.com>

Author: MengqingCao

.github/workflows/vllm_ascend_test_pd.yaml

@@ -100,3 +100,21 @@ jobs:
         run: |
           pip install -r requirements-dev.txt
           pip install -v -e .
+      # only run test on spec decode when the related code changed
+      - name: Check for changes in Speculative Decode
+        if: github.event_name != 'schedule'
+        id: filter_pd
+        uses: dorny/paths-filter@v3
+        with:
+          filters: |
+            pd_tests_changed:
+              - ".github/workflows/vllm_ascend_test_pd.yaml"
+              - "tests/e2e/pd_disaggreate/**"
+              - "tests/e2e/run_disagg_pd.sh"
+              - "vllm_ascend/distributed/**"
+              - "vllm_ascend/models/deepseek_v2.py"
+
+      - name: Run vllm-project/vllm-ascend PD Disaggregation test
+        if: steps.filter_pd.outputs.pd_tests_changed == 'true' || github.event_name == 'schedule'
+        run: |
+          pytest -sv tests/e2e/pd_disaggreate/test_pd_e2e.py

examples/disaggregated_prefill/disaggregated_prefill_offline.py

@@ -13,7 +13,7 @@
 from multiprocessing import Event, Process
 
 kv_connector_extra_config = {
-    "prompt_device_ips": ["1.2.3.1", "1.2.3.2"],
+    "prefill_device_ips": ["1.2.3.1", "1.2.3.2"],
     "decode_device_ips": ["1.2.3.9", "1.2.3.10"],
     "llmdatadist_comm_port": 26000,
 }

examples/disaggregated_prefill/p2p_disaggrefated_prefill_proxy.py

@@ -181,6 +181,13 @@ async def handle_request():
 
 
 if __name__ == "__main__":
-    t = start_service_discovery("0.0.0.0", 30001)
-    app.run(host="0.0.0.0", port=10001)
+    import argparse
+    parser = argparse.ArgumentParser(
+        description="args of disaggregated-prefill proxy")
+    parser.add_argument("--http-port", type=int, default=10001)
+    parser.add_argument("--register-port", type=int, default=10002)
+    args = parser.parse_args()
+
+    t = start_service_discovery("0.0.0.0", args.register_port)
+    app.run(host="0.0.0.0", port=args.http_port)
     t.join()

format.sh

@@ -173,21 +173,21 @@ spell_check_changed() {
     fi
 }
 
-echo 'vllm-ascend codespell:'
-# Run Codespell
-## This flag runs spell check of individual files. --files *must* be the first command line
-## arg to use this option.
-if [[ "$1" == '--files' ]]; then
-   spell_check "${@:2}"
-   # If `--all` is passed, then any further arguments are ignored and the
-   # entire python directory is linted.
-elif [[ "$1" == '--all' ]]; then
-   spell_check_all
-else
-   # Check spelling only of the files that changed in last commit.
-   spell_check_changed
-fi
-echo 'vllm-ascend codespell: Done'
+# echo 'vllm-ascend codespell:'
+# # Run Codespell
+# ## This flag runs spell check of individual files. --files *must* be the first command line
+# ## arg to use this option.
+# if [[ "$1" == '--files' ]]; then
+#    spell_check "${@:2}"
+#    # If `--all` is passed, then any further arguments are ignored and the
+#    # entire python directory is linted.
+# elif [[ "$1" == '--all' ]]; then
+#    spell_check_all
+# else
+#    # Check spelling only of the files that changed in last commit.
+#    spell_check_changed
+# fi
+# echo 'vllm-ascend codespell: Done'
 
 
 # Lint specified files

requirements-dev.txt

@@ -10,3 +10,4 @@ types-jsonschema
 xgrammar
 zmq
 quart
+types-psutil

tests/e2e/pd_disaggreate/setup_pd.sh

@@ -31,6 +31,8 @@ function run_prefill_instance() {
   local tp_size=$2
   local prefill_port=$3
   local register_port=$4
+  local prefill_device_ips=$5
+  local decode_device_ips=$6
 
   echo "================================"
   echo "Testing model: $model_name"
@@ -43,6 +45,8 @@ function run_prefill_instance() {
     --arg kv_role "kv_producer" \
     --argjson kv_parallel_size 8 \
     --arg kv_port "11001" \
+    --argjson prefill_device_ips "$prefill_device_ips" \
+    --argjson decode_device_ips "$decode_device_ips" \
     --argjson llmdatadist_comm_port "26000" \
     --arg proxy_ip "0.0.0.0" \
     --argjson proxy_port "$register_port" \
@@ -54,8 +58,8 @@ function run_prefill_instance() {
       "kv_parallel_size": $kv_parallel_size,
       "kv_port": $kv_port,
       "kv_connector_extra_config": {
-        "prompt_device_ips": ["29.7.130.29"],
-        "decode_device_ips": ["29.7.186.66"],
+        "prefill_device_ips": $prefill_device_ips,
+        "decode_device_ips": $decode_device_ips,
         "llmdatadist_comm_port": $llmdatadist_comm_port,
         "proxy_ip": $proxy_ip,
         "proxy_port": $proxy_port,
@@ -82,13 +86,17 @@ function run_decode_instance() {
   local tp_size=$2
   local decode_port=$3
   local register_port=$4
+  local prefill_device_ips=$5
+  local decode_device_ips=$6
 
   KV_CONFIG=$(jq -n \
     --arg kv_connector "AscendSimpleConnector" \
     --arg kv_buffer_device "npu" \
     --arg kv_role "kv_consumer" \
     --argjson kv_parallel_size 8 \
     --arg kv_port "21001" \
+    --argjson prefill_device_ips "$prefill_device_ips" \
+    --argjson decode_device_ips "$decode_device_ips" \
     --argjson llmdatadist_comm_port "26000" \
     --arg proxy_ip "0.0.0.0" \
     --argjson proxy_port "$register_port" \
@@ -100,8 +108,8 @@ function run_decode_instance() {
       "kv_parallel_size": $kv_parallel_size,
       "kv_port": $kv_port,
       "kv_connector_extra_config": {
-        "prompt_device_ips": ["29.7.130.29"],
-        "decode_device_ips": ["29.7.186.66"],
+        "prefill_device_ips": $prefill_device_ips,
+        "decode_device_ips": $decode_device_ips,
         "llmdatadist_comm_port": $llmdatadist_comm_port,
         "proxy_ip": $proxy_ip,
         "proxy_port": $proxy_port,

tests/e2e/pd_disaggreate/test_pd_e2e.py

@@ -1,12 +1,60 @@
+#!/bin/bash
+
+#
+# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+#
+
 import os
 import signal
 import subprocess
 import time
 
+import psutil
 import requests
 
-PROXY_PORT = 8192
-REGISTER_PORT = 8193
+
+def kill_process_and_children(pid):
+    try:
+        parent = psutil.Process(pid)
+        children = parent.children(recursive=True)
+        for child in children:
+            print(f"Killing child process {child.pid}")
+            child.kill()
+        print(f"Killing parent process {pid}")
+        parent.kill()
+    except psutil.NoSuchProcess:
+        pass
+
+
+def kill_all_vllm_related():
+    current_pid = os.getpid()
+
+    for proc in psutil.process_iter(['pid', 'cmdline']):
+        try:
+            if proc.pid == current_pid:
+                continue
+            cmd = ' '.join(proc.info['cmdline'])
+            if "vllm" in cmd or "proxy" in cmd or "engine_worker" in cmd:
+                kill_process_and_children(proc.pid)
+        except Exception:
+            continue
+
+
+PROXY_PORT = 10102
+DECODE_PORT = 8002
 
 SCRIPT_PATH = os.path.abspath("./tests/e2e/run_disagg_pd.sh")
 
@@ -28,18 +76,20 @@ def start_and_test_pipeline():
     try:
         print("Waiting for proxy port to be available...")
         wait_for_port(PROXY_PORT, 1200)
+        wait_for_port(DECODE_PORT, 1200)
 
         # request
-        prompt = "The future of AI is"
         payload = {
-            "model": "Deepseek/DeepSeek-V2-Lite-Chat",
-            "prompt": prompt,
+            "model": "Deepseek",
+            "prompt": "The future of AI is",
             "max_tokens": 64,
             "temperature": 0,
         }
-        response = requests.post(f"http://localhost:{PROXY_PORT}/generate",
-                                 json=payload,
-                                 timeout=10)
+        response = requests.post(
+            f"http://localhost:{PROXY_PORT}/v1/completions",
+            headers={"Content-Type": "application/json"},
+            json=payload,
+            timeout=10)
         assert response.status_code == 200, f"HTTP failed: {response.status_code}"
         result = response.json()
         print("Response:", result)
@@ -54,6 +104,7 @@ def start_and_test_pipeline():
             proc.wait(timeout=10)
         except subprocess.TimeoutExpired:
             proc.kill()
+        kill_all_vllm_related()
 
 
 def test_disaggregated_pd_pipeline():

tests/e2e/run_disagg_pd.sh

@@ -22,9 +22,19 @@ set -eo errexit
 . $(dirname "$0")/common.sh
 . $(dirname "$0")/pd_disaggreate/setup_pd.sh
 
-MODEL_NAME="/home/cmq/vllm-workspace/DeepSeek-V2-Lite"
+export VLLM_USE_MODELSCOPE="True"
+
+MODEL_NAME="deepseek-ai/DeepSeek-V2-Lite"
+# TODO: add tp case
 TP_SIZE=1
 
+# TODO: support multi-card
+prefill_ip=$(hccn_tool -i 0 -ip -g | grep "ipaddr" | awk -F: '{print $2}' | xargs)
+PREFILL_DEVICE_IPS="[\"$prefill_ip\"]"
+
+decode_ip=$(hccn_tool -i 1 -ip -g | grep "ipaddr" | awk -F: '{print $2}' | xargs)
+DECODE_DEVICE_IPS="[\"$decode_ip\"]"
+
 _info "====> Start pd disaggregated test"
 REGISTER_PORT=10101
 PREOXY_PORT=10102
@@ -33,15 +43,15 @@ _info "Started pd disaggregated proxy server"
 
 PREFILL_PROC_NAME="Prefill-instance"
 PREFILL_PORT=8001
-run_prefill_instance $MODEL_NAME $TP_SIZE $PREFILL_PORT $REGISTER_PORT
-_info "Startting prefill instance"
+run_prefill_instance $MODEL_NAME $TP_SIZE $PREFILL_PORT $REGISTER_PORT $PREFILL_DEVICE_IPS $DECODE_DEVICE_IPS
+_info "Starting prefill instance"
 
 wait_url_ready $PREFILL_PROC_NAME "http://localhost:${PREFILL_PORT}/v1/completions"
 
 DECODE_PROC_NAME="Decode-instance"
 DECODE_PORT=8002
-run_decode_instance  $MODEL_NAME $TP_SIZE $DECODE_PORT $REGISTER_PORT
-_info "Startting decode instance"
+run_decode_instance  $MODEL_NAME $TP_SIZE $DECODE_PORT $REGISTER_PORT $PREFILL_DEVICE_IPS $DECODE_DEVICE_IPS
+_info "Starting decode instance"
 
 wait_url_ready $DECODE_PROC_NAME "http://localhost:${DECODE_PORT}/v1/completions"
 

vllm_ascend/distributed/kv_transfer/simple_pipe.py

@@ -61,22 +61,22 @@ def __init__(
             raise NotImplementedError(
                 "kv_role should be inside [kv_producer, kv_consumer]")
 
-        prompt_device_ips = kv_connector_extra_config.get(
-            "prompt_device_ips", None)
+        prefill_device_ips = kv_connector_extra_config.get(
+            "prefill_device_ips", None)
         decode_device_ips = kv_connector_extra_config.get(
             "decode_device_ips", None)
-        if prompt_device_ips is None or decode_device_ips is None:
+        if prefill_device_ips is None or decode_device_ips is None:
             raise ValueError(
-                "Please specify prompt_device_ips and decode_device_ips"
+                "Please specify prefill_device_ips and decode_device_ips"
                 "in kv_transfer_config.kv_connector_extra_config")
-        p_device_num = len(prompt_device_ips)
+        p_device_num = len(prefill_device_ips)
         d_device_num = len(decode_device_ips)
         # When number of devices in P and D is not equal,
         # we assume that device in D can be mapped to any device in P.
         self.p_device_rank = self.rank % p_device_num
         self.d_device_rank = self.rank % d_device_num
 
-        self.prompt_ip_list = prompt_device_ips
+        self.prompt_ip_list = prefill_device_ips
         self.decode_ip_list = decode_device_ips
         self.llmdatadist_comm_port = kv_connector_extra_config.get(
             "llmdatadist_comm_port", 26000)