support a2

Author: liziyu179

examples/disaggregate_prefill_v1/gen_ranktable.py

@@ -0,0 +1,90 @@
+import os
+import torch.distributed as dist
+import json
+import argparse
+from vllm_ascend.soc_info import NPUSocInfo
+
+parser = argparse.ArgumentParser(
+    description="Arguments of rank table generator",
+)
+parser.add_argument(
+    "--prefill-device-cnt", type=int, required=True, help="number of prefill devices"
+)
+parser.add_argument(
+    "--decode-device-cnt", type=int, required=True, help="number of decode devices"
+)
+args = parser.parse_args()
+prefill_device_cnt = args.prefill_device_cnt
+decode_device_cnt = args.decode_device_cnt
+
+print("enter py")
+
+master_addr = os.environ.get("MASTER_ADDR")
+master_port = os.environ.get("MASTER_PORT")
+rank = os.environ.get("RANK")
+# This variable is set by torchrun, 
+# and is different from WORLD_SIZE in gen_rank_table.sh.
+world_size = os.environ.get("WORLD_SIZE")
+soc_info = NPUSocInfo()
+
+def get_cmd_stdout(cmd):
+    import subprocess
+    return subprocess.run(
+        cmd,
+        capture_output=True,
+        shell=True
+    ).stdout.decode("utf-8").strip()
+
+local_host = get_cmd_stdout("hostname -I | awk -F \" \" \'{print$1}\'")
+print(f"local_host: {local_host}")
+print("gen ranktable.json")
+
+num_cards = get_cmd_stdout("npu-smi info -l | grep \"Total Count\"").split(":")[1].strip()
+num_cards = int(num_cards)
+chips_per_card = get_cmd_stdout("npu-smi info -l | grep \"Chip Count\"").split("\n")[0].split(":")[1].strip()
+chips_per_card = int(chips_per_card)
+
+local_device_list: list[dict[str, str]] = list()
+super_pod_id = "0"
+for card_id in range(num_cards):
+    for chip_id in range(chips_per_card):
+        device_id = card_id * chips_per_card + chip_id
+        if soc_info.is_a3:
+            device_ip = get_cmd_stdout(f"/usr/local/Ascend/driver/tools/hccn_tool -i {device_id} -vnic -g | grep ipaddr").split(":")[1].strip()
+            super_device_id = get_cmd_stdout(f"npu-smi info -t spod-info -i {card_id} -c {chip_id} | grep SDID").split(":")[1].strip()
+            super_pod_id = get_cmd_stdout(f"npu-smi info -t spod-info -i {card_id} -c {chip_id} | grep \"Super Pod ID\"").split(":")[1].strip()
+        else:
+            device_ip = get_cmd_stdout(f"/usr/local/Ascend/driver/tools/hccn_tool -i {device_id} -ip -g | grep ipaddr").split(":")[1].strip()
+
+        device_info = {
+                "server_id": local_host,
+                "device_id": str(device_id),
+                "device_ip": str(device_ip),
+            }
+        if soc_info.is_a3:
+            device_info.update({"super_pod_id": str(super_pod_id), "super_device_id": str(super_device_id)})
+        local_device_list.append(device_info)
+
+dist.init_process_group(backend=dist.Backend.GLOO)
+global_device_list = [None] * dist.get_world_size()
+dist.all_gather_object(global_device_list, local_device_list)
+global_device_list = [device_info for device_list in global_device_list for device_info in device_list]
+cnt = 1
+for device_info in global_device_list:
+    device_info["cluster_id"] = str(cnt)
+    cnt += 1
+assert (prefill_device_cnt + decode_device_cnt) <= len(global_device_list), \
+"prefill_device_cnt + decode_device_cnt must be less than or equal to number of all devices in cluster"
+ranktable = {
+    "version": "1.2",
+    "server_count": str(world_size),
+    "prefill_device_list": global_device_list[:prefill_device_cnt],
+    "decode_device_list": global_device_list[prefill_device_cnt:prefill_device_cnt+decode_device_cnt],
+    "status": "completed"
+}
+
+
+with open("ranktable.json", "w") as f:
+    json.dump(ranktable, f, indent=4)
+
+print("gen ranktable.json done")

examples/disaggregate_prefill_v1/gen_ranktable.sh

@@ -0,0 +1,46 @@
+gen_rank_table.sh
+#!/bin/bash
+
+source /usr/local/Ascend/ascend-toolkit/set_env.sh
+export LD_LIBRARY_PATH=/usr/local/Ascend/ascend-toolkit/latest/opp/vendors/customize/op_api/lib/:${LD_LIBRARY_PATH}
+
+IPs=('1.0.0.0' '1.0.0.1')
+LOCAL_HOST=`hostname -I|awk -F " " '{print$1}'`
+GPUS_PER_NODE=8
+MASTER_ADDR=${IPs[0]}
+MASTER_PORT=6657
+NNODES=${#IPs[@]}
+NODE_RANK="2"
+for i in "${!IPs[@]}";
+do
+    echo "${IPs[$i]}"
+    if [ "$LOCAL_HOST" == "${IPs[$i]}" ];
+    then
+        NODE_RANK=$i
+        break
+    fi
+done
+if [[ $NODE_RANK == "" ]];then
+    echo "[Error] para \"NODE_RANK\" must be confing"
+    exit 1
+fi
+
+WORLD_SIZE=$(($GPUS_PER_NODE * $NNODES))
+RANKSTART=`expr $GPUS_PER_NODE \* $NODE_RANK`
+
+echo "========>param:"
+echo "WORLD_SIZE: " $WORLD_SIZE
+echo "RANKSTART": $RANKSTART
+echo "NNODES": $NNODES
+echo "NODE_RANK": $NODE_RANK
+echo "==============="
+
+if [[ -n "${GEN_RANKTABLE}" || ! -e ${PWD}/ranktable.json ]]; then
+    GLOO_SOCKET_IFNAME=enp189s0f0 torchrun \
+        --nproc_per_node 1 \
+        --nnodes ${NNODES} \
+        --node_rank ${NODE_RANK} \
+        --master_addr ${MASTER_ADDR} \
+        --master_port ${MASTER_PORT} \
+        gen_ranktable.py --prefill-device-cnt $1 --decode-device-cnt $2
+fi

vllm_ascend/distributed/llmdatadist_connector_v1_a3.py

@@ -1,4 +1,5 @@
 import msgspec
+import os
 from dataclasses import dataclass
 
 from typing import Optional, Any, Tuple
@@ -28,6 +29,7 @@
 # from .llmdatadist_connector_v1 import TORCH_DTYPE_TO_NPU_DTYPE
 from vllm.v1.request import Request
 from vllm.utils import logger
+from vllm_ascend.soc_info import NPUSocInfo
 
 import llm_datadist
 from llm_datadist import LLMDataDist, LLMRole, CacheDesc, BlocksCacheKey, LLMConfig, LLMException
@@ -314,14 +316,15 @@ def __init__(
       self.llm_datadist = LLMDataDist(self.llm_datadist_role, self.local_agent_metadata.cluster_id)
       self.init_llm_datadist()
       self.finished_reqs = set()
+      self.soc_info = NPUSocInfo()
       # remote_ip, remote_rank = self.get_remote_ip_and_rank()
       # for idx in range(len(remote_ip)):
       #   remote_agent_meta = self.read_agent_metadata(global_rank_table, remote_ip[idx], remote_rank[idx], self.llm_datadist_remote_role)
       #   self.add_remote_agent(remote_agent_meta)
 
 
   def listen_for_agent_metadat_req(self, event: threading.Event):
-    port = envs.VLLM_LLMDD_CHANNEL_PORT + self.local_dp_rank * self.tp_size
+    port = envs.VLLM_LLMDD_CHANNEL_PORT + self.local_dp_rank * self.tp_size + self.tp_rank
     url = f"tcp://0.0.0.0:{port}"
     msg_encoder = msgspec.msgpack.Encoder()
     msg_decoder = msgspec.msgpack.Decoder()
@@ -394,11 +397,11 @@ def read_agent_metadata(self, global_rank_table, server_id, device_id, agent_rol
           continue
         if device_info["device_id"] != str(device_id):
           continue
-        super_pod_id_ = device_info["super_pod_id"]
+        super_pod_id_ = device_info.get("super_pod_id", None)
         server_id_ = device_info["server_id"]
         device_id_ = device_info["device_id"]
         device_ip_ = device_info["device_ip"]
-        super_device_id_ = device_info["super_device_id"]
+        super_device_id_ = device_info.get("super_device_id", None)
         cluster_id_ = int(device_info["cluster_id"])
         agent_metadata = LLMDataDistAgentMetadata(
           super_pod_id=super_pod_id_,
@@ -537,32 +540,38 @@ def add_remote_agent(self, metadata: LLMDataDistAgentMetadata) -> bool:
     decode_server_device_info = None
     prefill_server_device_info = {
       "device": [
-        {
-          "device_id": prefill_metadata.device_id,
-          "device_ip": prefill_metadata.device_ip,
-          "super_device_id": prefill_metadata.super_device_id,
-          "rank_id": "0"
+        { 
+          k: v for k, v in [
+          ("device_id", prefill_metadata.device_id),
+          ("device_ip", prefill_metadata.device_ip),
+          ("super_device_id", prefill_metadata.super_device_id),
+          ("rank_id", "0")]
+          if v is not None
         }
       ],
       "server_id": prefill_metadata.server_id
     }
     if is_same_server:
       prefill_server_device_info["device"].append(
-        {
-          "device_id": decode_metadata.device_id,
-          "device_ip": decode_metadata.device_ip,
-          "super_device_id": decode_metadata.super_device_id,
-          "rank_id": "1"
+        { 
+          k: v for k, v in [
+          ("device_id", decode_metadata.device_id),
+          ("device_ip", decode_metadata.device_ip),
+          ("super_device_id", decode_metadata.super_device_id),
+          ("rank_id", "1")]
+          if v is not None
         }
       )
     else:
       decode_server_device_info = {
         "device": [
-          {
-            "device_id": decode_metadata.device_id,
-            "device_ip": decode_metadata.device_ip,
-            "super_device_id": decode_metadata.super_device_id,
-            "rank_id": "1"
+          { 
+            k: v for k, v in [
+            ("device_id", decode_metadata.device_id),
+            ("device_ip", decode_metadata.device_ip),
+            ("super_device_id", decode_metadata.super_device_id),
+            ("rank_id", "1")]
+            if v is not None
           }
         ],
         "server_id": decode_metadata.server_id
@@ -571,28 +580,29 @@ def add_remote_agent(self, metadata: LLMDataDistAgentMetadata) -> bool:
     if decode_server_device_info is not None:
       rank_table["server_list"].append(decode_server_device_info)
 
-    # generate super_pod_list for rank table
-    super_pod_list = []
-    prefill_super_pod_info = {
-       "super_pod_id": prefill_metadata.super_pod_id,
-       "server_list": [
-          {"server_id": prefill_metadata.server_id}
-        ],
-    }
-    if is_same_pod and not is_same_server:
-      prefill_super_pod_info["server_list"].append(
-        {"server_id": decode_metadata.server_id}
-      )
-    super_pod_list.append(prefill_super_pod_info)
-    if not is_same_pod:
-      decode_super_pod_id = {
-        "super_pod_id": decode_metadata.super_pod_id,
+    if self.soc_info.is_a3:
+      # generate super_pod_list for rank table
+      super_pod_list = []
+      prefill_super_pod_info = {
+        "super_pod_id": prefill_metadata.super_pod_id,
         "server_list": [
-          {"server_id": decode_metadata.server_id}
-        ],
+            {"server_id": prefill_metadata.server_id}
+          ],
       }
-      super_pod_list.append(decode_super_pod_id)
-    rank_table["super_pod_list"] = super_pod_list
+      if is_same_pod and not is_same_server:
+        prefill_super_pod_info["server_list"].append(
+          {"server_id": decode_metadata.server_id}
+        )
+      super_pod_list.append(prefill_super_pod_info)
+      if not is_same_pod:
+        decode_super_pod_id = {
+          "super_pod_id": decode_metadata.super_pod_id,
+          "server_list": [
+            {"server_id": decode_metadata.server_id}
+          ],
+        }
+        super_pod_list.append(decode_super_pod_id)
+      rank_table["super_pod_list"] = super_pod_list
     logger.info(f"LLMDataDistConnectorWorker: try link with remote, comm id: {comm_name}")
     logger.info(f"rank table \n{rank_table}")
     logger.info(f"comm name: {comm_name}")

vllm_ascend/soc_info.py

@@ -0,0 +1,12 @@
+from dataclasses import dataclass
+import torch_npu
+
+@dataclass
+class NPUSocInfo:
+    is_a3: bool = False
+
+    def __post_init__(self):
+        torch_npu.npu._lazy_init()
+        self.soc_version = torch_npu._C._npu_get_soc_version()
+        if self.soc_version in (253, 254, 255):
+            self.is_a3 = True