[Fixbug] Fix soc_version for 310p

Signed-off-by: hfadzxy <starmoon_zhang@163.com>

Author: zhangxinyuehfad

Dockerfile.a3

@@ -50,6 +50,7 @@ RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi
     source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
     source /usr/local/Ascend/nnal/atb/set_env.sh && \
     export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
+    export SOC_VERSION=Ascend910_9392 && \
     python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
     python3 -m pip cache purge
 

Dockerfile.a3.openEuler

@@ -48,6 +48,7 @@ RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi
     source /usr/local/Ascend/nnal/atb/set_env.sh && \
     export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
     export CPLUS_INCLUDE_PATH=$CPLUS_INCLUDE_PATH:/usr/include/c++/12:/usr/include/c++/12/`uname -i`-openEuler-linux && \
+    export SOC_VERSION=Ascend910_9392 && \
     python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
     python3 -m pip cache purge
 

examples/disaggregated_prefill_v1/gen_ranktable.py

@@ -4,7 +4,7 @@
 
 import torch.distributed as dist
 
-from vllm_ascend.utils import AscendSocVersion, init_ascend_soc_version, get_ascend_soc_version
+from vllm_ascend import _build_info  # type: ignore
 
 parser = argparse.ArgumentParser(
     description="Arguments of rank table generator", )
@@ -38,8 +38,7 @@
 # and is different from WORLD_SIZE in gen_rank_table.sh.
 world_size = os.environ.get("WORLD_SIZE")
 
-init_ascend_soc_version()
-soc_info = get_ascend_soc_version()
+soc_info = _build_info.__ascend_soc_version__
 
 
 def get_cmd_stdout(cmd):
@@ -75,7 +74,7 @@ def get_cmd_stdout(cmd):
         device_id = local_device_ids[idx]
         chip_id = device_id % chips_per_card
         card_id = device_id // chips_per_card
-        if soc_info == AscendSocVersion.A3:
+        if soc_info == "A3":
             device_ip = get_cmd_stdout(
                 f"{hccn_tool_path} -i {device_id} -vnic -g | grep ipaddr"
             ).split(":")[1].strip()
@@ -87,15 +86,15 @@ def get_cmd_stdout(cmd):
             ).split(":")[1].strip()
         else:
             device_ip = get_cmd_stdout(
-                f"{hccn_tool_path} -i {device_id} -ip -g | grep ipaddr"
-            ).split(":")[1].strip()
+                f"{hccn_tool_path} -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 == AscendSocVersion.A3:
+        if soc_info == "A3":
             device_info.update({
                 "super_pod_id": str(super_pod_id),
                 "super_device_id": str(super_device_id)

setup.py

@@ -32,6 +32,11 @@
 from setuptools.command.install import install
 from setuptools_scm import get_version
 
+# Supported SOC_VERSION codes
+ASCEND_A2_SOC_VERSION = ["ASCEND910B1"]
+ASCEND_A3_SOC_VERSION = ["ASCEND910_9392"]
+ASCEND_310P_SOC_VERSION = ["ASCEND310P3"]
+
 
 def load_module_from_path(module_name, path):
     spec = importlib.util.spec_from_file_location(module_name, path)
@@ -91,11 +96,20 @@ def run(self):
             raise ValueError(
                 "SOC version 310 only supports custom kernels. Please set COMPILE_CUSTOM_KERNELS=1 to enable custom kernels."
             )
+        if soc_version in ASCEND_A2_SOC_VERSION:
+            ascend_soc_version = "A2"
+        elif soc_version in ASCEND_A3_SOC_VERSION:
+            ascend_soc_version = "A3"
+        elif soc_version in ASCEND_310P_SOC_VERSION:
+            ascend_soc_version = "310P"
+        else:
+            ascend_soc_version = "UNDEFINED"
 
         package_dir = os.path.join(ROOT_DIR, "vllm_ascend", "_build_info.py")
         with open(package_dir, "w+") as f:
             f.write('# Auto-generated file\n')
             f.write(f"__soc_version__ = '{soc_version}'\n")
+            f.write(f"__ascend_soc_version__ = '{ascend_soc_version}'\n")
             f.write(
                 f"__sleep_mode_enabled__ = {envs.COMPILE_CUSTOM_KERNELS}\n")
         logging.info(

tests/ut/conftest.py

@@ -23,4 +23,4 @@
 adapt_patch(True)
 
 # register Ascend CustomOp here because uts will use this
-register_ascend_customop()
+register_ascend_customop()

tests/ut/models/conftest.py

@@ -98,7 +98,7 @@ def mock_distributed():
                   return_value=Mock(is_first_rank=False, is_last_rank=False)), \
             patch("vllm_ascend.ops.fused_moe.get_current_vllm_config", return_value=mock_vllm_config), \
             patch("vllm_ascend.ops.moe.token_dispatcher.torch.distributed.get_rank", return_value=0), \
-            patch("vllm_ascend.ops.moe.token_dispatcher.get_ascend_soc_version", return_value=None), \
+            patch("vllm_ascend._build_info.__ascend_soc_version__", return_value=""), \
             patch.dict("vllm.distributed.parallel_state.__dict__", _TP=tp_group, _EP=ep_group, _DP=dp_group,
                        _PP=pp_group), \
             patch.dict("vllm_ascend.distributed.parallel_state.__dict__", _MC2=ep_group), \

tests/ut/ops/test_fused_ops.py

@@ -28,7 +28,7 @@
                                        AscendUnquantizedFusedMoEMethod)
 from vllm_ascend.ops.moe.experts_selector import select_experts
 from vllm_ascend.ops.moe.moe_mlp import cumsum_group_list, unified_apply_mlp
-from vllm_ascend.utils import AscendSocVersion, adapt_patch
+from vllm_ascend.utils import adapt_patch
 
 adapt_patch(True)
 
@@ -125,7 +125,13 @@ def mock_finalize(hidden_states, **kwargs):
             return_value=mock_forward_context_obj), \
         patch('vllm_ascend.ops.moe.fused_moe_prepare_and_finalize.get_forward_context',
             return_value=mock_forward_context_obj), \
-        patch("vllm_ascend.utils.get_ascend_soc_version", return_value=AscendSocVersion.A3), \
+        patch('vllm_ascend.ops.fused_moe.get_current_vllm_config',
+                return_value=MagicMock(
+                    parallel_config=MagicMock(tensor_parallel_size=2),
+                    scheduler_config=MagicMock(max_num_seqs=4),
+                    model_config=MagicMock(max_model_len=2048)
+                )), \
+        patch("vllm_ascend._build_info.__ascend_soc_version__", return_value="A3"), \
         patch('vllm_ascend.ops.moe.moe_mlp.get_forward_context',
                 return_value=mock_forward_context_obj), \
         patch('vllm_ascend.ops.moe.moe_comm_method.MC2CommImpl._get_token_dispatcher',
@@ -406,7 +412,8 @@ def test_apply_with_expert_map(self, moe_method, mock_dist_env,
         forward_context = mock_dist_env['mock_forward_context_obj']
 
         with patch("vllm_ascend.ops.fused_moe.get_forward_context", return_value=forward_context), \
-             patch("vllm_ascend.utils.get_ascend_soc_version", return_value=AscendSocVersion.A3):
+             patch("vllm_ascend._build_info.__ascend_soc_version__", return_value="A3"):
+
             expert_map = torch.tensor([0, 1, 2, -1, -1, -1, -1, -1])
             moe_method.ep_size = ep_size
             x = torch.randn(8, 2, 2)

tests/ut/ops/test_token_dispatcher.py

@@ -22,8 +22,8 @@
 from tests.ut.base import TestBase
 
 from vllm_ascend.ops.moe.token_dispatcher import (  # isort: skip
-    AscendSocVersion, TokenDispatcherWithAll2AllV,
-    TokenDispatcherWithAllGather, TokenDispatcherWithMC2)
+    TokenDispatcherWithAll2AllV, TokenDispatcherWithAllGather,
+    TokenDispatcherWithMC2)
 
 
 class TestTokenDispatcherWithMC2(TestBase):
@@ -50,10 +50,9 @@ def setUp(self):
             return_value=self.forward_context)
         self.forward_context_patch.start()
 
-        # Mock get_ascend_soc_version()
+        # Mock __ascend_soc_version__
         self.ascend_soc_version_patch = patch(
-            "vllm_ascend.ops.moe.token_dispatcher.get_ascend_soc_version",
-            return_value=AscendSocVersion.A3)
+            "vllm_ascend._build_info.__ascend_soc_version__", new="A3")
         self.ascend_soc_version_patch.start()
 
         kwargs = {"with_quant": False, "top_k": 8, "num_experts": 128}

tests/ut/torchair/ops/test_torchair_fused_moe.py

@@ -26,7 +26,7 @@
 from vllm_ascend.quantization.quant_config import AscendFusedMoEMethod
 from vllm_ascend.torchair.ops.torchair_fused_moe import (
     TorchairAscendFusedMoE, TorchairAscendUnquantizedFusedMoEMethod)
-from vllm_ascend.utils import AscendSocVersion, adapt_patch  # noqa E402
+from vllm_ascend.utils import adapt_patch  # noqa E402
 
 adapt_patch(True)
 
@@ -383,7 +383,7 @@ def test_apply_with_expert_map(self, moe_method, mock_dist_env,
         forward_context = MagicMock(
             fused_moe_state=_get_fused_moe_state(ep_size, is_prefill, True))
         with patch("vllm_ascend.torchair.ops.torchair_fused_moe.get_forward_context", return_value=forward_context), \
-             patch("vllm_ascend.torchair.ops.torchair_fused_moe.get_ascend_soc_version", return_value=AscendSocVersion.A3):
+             patch("vllm_ascend._build_info.__ascend_soc_version__", return_value="A3"):
             expert_map = torch.tensor([0, 1, 2, -1, -1, -1, -1, -1])
             moe_method.ep_size = ep_size
             x = torch.randn(8, 2, 2)

tests/ut/worker/test_model_runner_v1.py

@@ -16,7 +16,6 @@
 import pytest
 
 from vllm_ascend.ascend_forward_context import MoECommType
-from vllm_ascend.utils import AscendSocVersion
 from vllm_ascend.worker.model_runner_v1 import NPUModelRunner
 
 
@@ -25,21 +24,25 @@
     "soc_version, enable_expert_parallel, world_size, num_tokens, mc2_tokens_capacity, quant_type, expected_method",
     [
         # Case 1: Expert parallel is disabled, should always be 'allgather'
-        (AscendSocVersion.A2, False, 8, 100, 256, None, MoECommType.ALLGATHER),
-        (AscendSocVersion.A3, False, 16, 500, 256, None, MoECommType.ALLGATHER),
+        ("A2", False, 8, 100, 256, None, "allgather"),
+        ("A3", False, 16, 500, 256, None, "allgather"),
 
         # Case 2: A2 SOC with w4a8_dynamic -> use alltoall when not mc2
-        (AscendSocVersion.A2, True, 8, 100, 256, "w4a8_dynamic", MoECommType.ALLTOALL),
-        (AscendSocVersion.A2, True, 16, 257, 256, "w4a8_dynamic", MoECommType.ALLTOALL),
-        (AscendSocVersion.A2, True, 16, 100, 256, "w4a8_dynamic", MoECommType.MC2),  # meets mc2 condition
+        ("A2", True, 8, 100, 256, "w4a8_dynamic", "alltoall"),
+        ("A2", True, 16, 257, 256, "w4a8_dynamic", "alltoall"),
+        ("A2", True, 16, 100, 256, "w4a8_dynamic", "mc2"),  # meets mc2 condition
 
         # Case 3: A2 SOC without w4a8_dynamic -> fallback to allgather
-        (AscendSocVersion.A2, True, 8, 100, 256, None, MoECommType.ALLGATHER),
-        (AscendSocVersion.A2, True, 16, 257, 256, None, MoECommType.ALLGATHER),
+        ("A2", True, 8, 100, 256, None, "allgather"),
+        ("A2", True, 16, 257, 256, None, "allgather"),
 
         # Case 4: A3 SOC
-        (AscendSocVersion.A3, True, 8, 100, 256, None, MoECommType.MC2),
-        (AscendSocVersion.A3, True, 8, 257, 256, None, MoECommType.ALLTOALL),
+        ("A3", True, 8, 100, 256, None, "mc2"),
+        ("A3", True, 8, 257, 256, None, "alltoall"),
+
+        # Case 5: P3 SOC
+        ("310P", True, 8, 100, 256, None, "allgather"),
+        ("310P", True, 8, 257, 256, None, "allgather"),
     ])
 # yapf: enable
 def test_select_moe_comm_method(soc_version, enable_expert_parallel,
@@ -65,8 +68,8 @@ def test_select_moe_comm_method(soc_version, enable_expert_parallel,
     mock_runner.vllm_config = mock_vllm_config
 
     # Patch the helper functions
-    with patch('vllm_ascend.worker.model_runner_v1.get_ascend_soc_version',
-               return_value=soc_version), \
+    with patch('vllm_ascend._build_info.__ascend_soc_version__',
+               new=soc_version), \
          patch('vllm_ascend.worker.model_runner_v1.is_global_first_rank',
                return_value=True):
 
@@ -98,8 +101,8 @@ def test_select_moe_comm_method_unsupported_soc():
 
     unsupported_soc = "UnsupportedSOC"
 
-    with patch('vllm_ascend.worker.model_runner_v1.get_ascend_soc_version',
-               return_value=unsupported_soc), \
+    with patch('vllm_ascend._build_info.__ascend_soc_version__',
+               new=unsupported_soc), \
          patch('vllm_ascend.worker.model_runner_v1.is_global_first_rank',
                return_value=True), \
          pytest.raises(ValueError, match=f"Unsupported soc_version: {unsupported_soc}"):