158 lines
6.6 KiB
Python
158 lines
6.6 KiB
Python
#
|
|
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
|
|
# Copyright 2023 The vLLM team.
|
|
#
|
|
# 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.
|
|
#
|
|
"""Shared UT setup.
|
|
|
|
NPU vs CPU routing is determined by directory convention, not decorators.
|
|
See ``.github/workflows/scripts/select_tests.py`` and
|
|
``.github/workflows/scripts/test_config.yaml`` for the routing rules.
|
|
|
|
Conventions for UT directories:
|
|
tests/ut/<module>/ -> CPU runner (default)
|
|
tests/ut/<module>/a2/ -> A2 NPU x1
|
|
tests/ut/<module>/a2_2/ -> A2 NPU x2
|
|
tests/ut/<module>/a3_2/ -> A3 NPU x2
|
|
tests/ut/<module>/a3_4/ -> A3 NPU x4
|
|
tests/ut/<module>/310p/ -> 310P NPU x1
|
|
"""
|
|
|
|
import importlib.util
|
|
import subprocess
|
|
import sys
|
|
import types
|
|
from unittest.mock import MagicMock
|
|
|
|
try:
|
|
# Note: do not import torch here for cpu env, which will lead to circle import error.
|
|
subprocess.run(["npu-smi", "info"], capture_output=True, check=True)
|
|
_npu_available = True
|
|
except (subprocess.CalledProcessError, FileNotFoundError):
|
|
_npu_available = False
|
|
|
|
if not _npu_available:
|
|
triton_runtime = MagicMock()
|
|
triton_runtime.driver.active.utils.get_device_properties.return_value = {
|
|
"num_aic": 8,
|
|
"num_vectorcore": 8,
|
|
}
|
|
sys.modules["triton.runtime"] = triton_runtime
|
|
torch_npu = types.ModuleType("torch_npu")
|
|
torch_npu.__spec__ = importlib.util.spec_from_loader("torch_npu", loader=None)
|
|
torch_npu.__path__ = []
|
|
torch_npu.npu = MagicMock() # type: ignore[attr-defined]
|
|
torch_npu.profiler = MagicMock() # type: ignore[attr-defined]
|
|
torch_npu.npu_fusion_attention = MagicMock() # type: ignore[attr-defined]
|
|
torch_npu.npu_format_cast = MagicMock(side_effect=lambda weight, fmt: weight) # type: ignore[attr-defined]
|
|
torch_npu._C = MagicMock() # type: ignore[attr-defined]
|
|
torch_npu._C._NPUTaskGroupHandle = MagicMock
|
|
sys.modules["torch_npu"] = torch_npu
|
|
sys.modules["torch_npu._C"] = torch_npu._C
|
|
sys.modules["torch_npu._C._distributed_c10d"] = torch_npu._C._distributed_c10d
|
|
acl_rt = types.ModuleType("acl.rt")
|
|
acl_rt.__spec__ = importlib.util.spec_from_loader("acl.rt", loader=None)
|
|
acl_rt.memcpy = MagicMock() # type: ignore[attr-defined]
|
|
acl_mod = types.ModuleType("acl")
|
|
acl_mod.__spec__ = importlib.util.spec_from_loader("acl", loader=None)
|
|
acl_mod.rt = acl_rt # type: ignore[attr-defined]
|
|
sys.modules["acl"] = acl_mod
|
|
sys.modules["acl.rt"] = acl_rt
|
|
mooncake_engine = types.ModuleType("mooncake.engine")
|
|
mooncake_engine.__spec__ = importlib.util.spec_from_loader("mooncake.engine", loader=None)
|
|
mooncake_engine.TransferEngine = MagicMock() # type: ignore[attr-defined]
|
|
sys.modules["mooncake.engine"] = mooncake_engine
|
|
import torch
|
|
|
|
try: # noqa: SIM105
|
|
torch.utils.rename_privateuse1_backend("npu")
|
|
except RuntimeError:
|
|
pass
|
|
torch.npu = MagicMock()
|
|
torch.npu.Stream = MagicMock
|
|
torch.version.cann = None
|
|
torch.distributed.is_hccl_available = MagicMock(return_value=True)
|
|
|
|
import pytest
|
|
|
|
mooncake_engine = types.ModuleType("mooncake.engine")
|
|
mooncake_engine.__spec__ = importlib.util.spec_from_loader("mooncake.engine", loader=None)
|
|
mooncake_engine.TransferEngine = MagicMock() # type: ignore[attr-defined]
|
|
sys.modules.setdefault("mooncake.engine", mooncake_engine)
|
|
|
|
from vllm_ascend.utils import ( # noqa: E402
|
|
adapt_patch,
|
|
clear_enable_sp,
|
|
register_ascend_customop,
|
|
)
|
|
|
|
# Mock torch_npu AFTER vllm_ascend import to avoid circular import in accelerate
|
|
if not _npu_available:
|
|
sys.modules["torch_npu"].npu.current_device = MagicMock(return_value=0)
|
|
sys.modules["torch_npu._inductor"] = MagicMock()
|
|
sys.modules["torch_npu"]._npu_flash_attention = MagicMock() # type: ignore[attr-defined]
|
|
sys.modules["torch_npu"]._npu_paged_attention_splitfuse = MagicMock() # type: ignore[attr-defined]
|
|
sys.modules["torch_npu"]._npu_reshape_and_cache = MagicMock() # type: ignore[attr-defined]
|
|
sys.modules["torch_npu"].npu_fused_infer_attention_score = MagicMock() # type: ignore[attr-defined]
|
|
sys.modules["torch_npu"]._npu_fused_infer_attention_score_get_max_workspace = MagicMock() # type: ignore[attr-defined]
|
|
sys.modules["torch_npu"].npu_moe_gating_top_k_softmax = MagicMock() # type: ignore[attr-defined]
|
|
sys.modules["torch_npu"].npu_quant_matmul = MagicMock() # type: ignore[attr-defined]
|
|
sys.modules["torch_npu"].npu_rms_norm = MagicMock() # type: ignore[attr-defined]
|
|
sys.modules["torch_npu"].npu_swiglu = MagicMock() # type: ignore[attr-defined]
|
|
sys.modules["torch_npu"].npu_convert_weight_to_int4pack = MagicMock() # type: ignore[attr-defined]
|
|
|
|
adapt_patch()
|
|
adapt_patch(True)
|
|
|
|
# register Ascend CustomOp here because uts will use this
|
|
register_ascend_customop()
|
|
|
|
# Clean up any stale mock modules that may have been installed by
|
|
# other test files (e.g., ascend_store/_mock_deps.py) which replace
|
|
# real subpackages with MagicMock, breaking later imports.
|
|
_stale_modules = [
|
|
k
|
|
for k in sys.modules
|
|
if k.startswith("vllm_ascend.distributed.kv_transfer.") and not isinstance(sys.modules[k], types.ModuleType)
|
|
]
|
|
for _m in _stale_modules:
|
|
del sys.modules[_m]
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _clear_enable_sp_before_test():
|
|
clear_enable_sp()
|
|
yield
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _mock_ascend_store_deps(request):
|
|
# ascend_store code imports vllm_ascend helpers (AttentionComputeStartGate,
|
|
# get/reset_attention_compute_start_gate, ...) which _mock_deps.py no longer
|
|
# mocks globally (mutating the real modules leaked into other UTs). Mock them
|
|
# per-test, scoped to the ascend_store tests only.
|
|
if "distributed/ascend_store/" not in request.node.nodeid:
|
|
yield
|
|
return
|
|
from unittest.mock import patch
|
|
|
|
_pfx = "vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store"
|
|
with (
|
|
patch(f"{_pfx}.pool_worker.get_attention_compute_start_gate"),
|
|
patch(f"{_pfx}.pool_worker.reset_attention_compute_start_gate"),
|
|
patch(f"{_pfx}.config_data.AttentionComputeStartGate", type("AttentionComputeStartGate", (), {})),
|
|
):
|
|
yield
|